feat(k3): cross-process CP whale gang — fleet-wide long-prompt prefill - #970
Open
xiaguan wants to merge 25 commits into
Open
feat(k3): cross-process CP whale gang — fleet-wide long-prompt prefill#970xiaguan wants to merge 25 commits into
xiaguan wants to merge 25 commits into
Conversation
The four invariants from docs/models/glm52/free-running-dp.md keep getting relearned the hard way when designing EP coordination (K3's gang lane included). Put them where every session loads them. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01LMbYEexVS2Qe6KBLoqjBrW Signed-off-by: xiaguan <751080330@qq.com>
…core CP8/CP16 gangs span processes, so the M0.5 in-process leveling board (one Mutex) no longer reaches the members, and the free-running discipline forbids replacing it with anything that waits. The rendezvous rides the two facts that stay true while everyone keeps stepping: the launch count is a global clock (mega pairing pins the world to ±1), and whales are sparse events. Consensus degenerates to one line — "whale w runs at absolute launch L" — sequenced by rank0 over a host TCP side channel. scheduler/whale.rs: pure state machines. Gather (descriptor broadcast; members reply with their launch count and arm — no multi-launch operations, count advances exactly one per step) then Commit (L = max(replies) + slack 4, strictly above the previous whale's L); every member free-runs to L and enters unanimously. Cancels only disarm — committed whales are uncancellable, unwanted results are dropped. Policy: widest admitting power-of-two gang (a superstep stalls the whole fleet, wider = shorter), tray-aligned with the poster rotated last, and quadratic-model segment leveling (t ∝ len + Q·depth·len, Q = 2.7e-6 from the CP4 16k profile) so the deep-prefix ranks get shorter segments; K3_CP_SEGMENT_FLOOR = 2048 is the FlashKDA saturation floor from the segment sweep. scheduler/whale_hub.rs: transports. LocalWhaleHub (in-process degenerate case) and TcpWhaleHub (rank0 hosts the sequencer; ep.rs-bootstrap-style framing, magic/version-tagged, gather deadline driven by the transport). Verified without a GPU, as designed: 14 deterministic protocol/policy tests (mid-op late replies, concurrent-whale serialization, timeout + stale ready, corrupted descriptors and slack breaches all die loudly) plus a 48-seed fleet fuzzer (±1 launch pinning, 0–1 round message latency over per-destination FIFOs, random multi-launch operations) asserting unanimous entry, seq-ordered launches, and quiescence. The fuzzer caught a real wedge — a width-refused whale stranded the queue behind it — now fixed in pump() and pinned by a regression test. TCP hub proven over loopback. Design and theory notes: docs/models/k3/cp-lane-design.md (whale rendezvous section: slack bound, leveling model, O(CP²) KDA package traffic and the Blelloch prefix-scan fallback). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01LMbYEexVS2Qe6KBLoqjBrW Signed-off-by: xiaguan <751080330@qq.com>
The whale lane joins the K3 scheduler as the fleet counterpart of the in-process CP gang (the two are mutually exclusive). Every step consults the rendezvous at its launch boundary: the member answers gathers with its exact launch count and arms, an armed or committed whale restricts the step to single-launch work (decode only) so the committed launch is hit dead on, and Enter calls the new StepExecutor::prefill_whale seam on every gang member at that one count. Posters keep exactly one whale post outstanding — the slot is reserved at post time so a commit can never land on a full batch — and a sequencer cancel routes the post to a plain local prefill at the next unrestricted launch. Mock-fleet tests drive real schedulers over a LocalWhaleHub and pin the serving contract: unanimous entry at one launch count with the poster as the last CP rank, non-members never hearing of a partial-width whale, concurrent and same-rank whales serializing onto distinct supersteps, and short prompts staying on the local prefill path. The engine assembly keeps whale: None for now — arming comes with the fleet data plane (fabric scratch + doorbells) in the next change. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01LMbYEexVS2Qe6KBLoqjBrW Signed-off-by: xiaguan <751080330@qq.com>
…l_whale The cross-machine half of the whale lane. The in-process CP exchange rests on two things a process boundary kills — peer-access pool pointers and CUDA events — and the fleet replaces both with the substrate MegaMoE already runs on. Every rank's CP publish surface (conv halo tail, KDA (M,D) packages, MLA latent/rope rows) moves into one fabric-exportable slab; the whale hub grows a startup allgather that moves the world's 64-byte handles once (doubling as the fleet's startup barrier, like the EP bootstrap) and each process imports the table once. Ordering becomes doorbells: u64 flags in the same slabs, remote cuStreamWriteValue64 to announce, local cuStreamWaitValue64/GEQ to wait — every wait on the rank's own memory, every remote touch a plain NVLink store, no host sync anywhere in a superstep. Doorbell values derive entirely from the rendezvous ((seq+1)*stride + window), so rotating gang membership never needs a host negotiation and the values stay strictly monotonic across whales a rank sits out. The forward path is untouched by the split: K3CpScratch now carries a K3CpSyncHandle (in-process group or fleet gang) and the three exchange windows dispatch through it. K3Executor::prefill_whale enters a committed superstep as the descriptor's seated member — same segment walk, same chunk step, leveled whale segments — with the shared body factored out of prefill_cp. PEGAINFER_K3_WHALE=<addr> arms the lane at launch: slabs allocate, the hub exchanges, executors import, and the scheduler's rendezvous serves whales with no further setup. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01LMbYEexVS2Qe6KBLoqjBrW Signed-off-by: xiaguan <751080330@qq.com>
…ne-design Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01LMbYEexVS2Qe6KBLoqjBrW Signed-off-by: xiaguan <751080330@qq.com>
…hostname /etc/hosts maps a tray's own name to 127.0.1.1, which would put the sequencer's listener on loopback while peers dial the fabric address — the same silent connection-refused trap the EP bootstrap already guards against. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01LMbYEexVS2Qe6KBLoqjBrW Signed-off-by: xiaguan <751080330@qq.com>
The stream memops engine rejects fabric-imported mappings outright (CUDA_ERROR_INVALID_VALUE on GB300, two-process probe), so every cross-process doorbell write died on the fleet's first whale superstep. SM stores go through the same mappings on every MegaMoE step, so the ring becomes a one-thread kernel batching all of a beat's flags; waits stay cuStreamWaitValue64 on the rank's own locally allocated slab, which memops accept. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01LMbYEexVS2Qe6KBLoqjBrW Signed-off-by: xiaguan <751080330@qq.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01LMbYEexVS2Qe6KBLoqjBrW Signed-off-by: xiaguan <751080330@qq.com>
The hub acceptor ran the slab exchange inline per accepted socket, and serve_slab_exchange blocks in wait_complete until the whole world's slabs are in — so with more than one peer process the first accepted connection wedged the accept loop and the remaining peers sat unserved in the TCP backlog until the connect timeout. CP8 (one peer process) worked; the CP16 4-tray fleet deadlocked at arming. Each accepted connection now gets its own thread running hello → slab exchange → writer registration → reader loop, so wait_complete only blocks that connection. The ordering constraint stands: the writer is registered only after the world table frame is written on that socket, so the sequencer can never interleave a commit with the table. Also breaks the whale arming log into phase timings (slab alloc / slab exchange barrier / import+install) to account for startup time, and adds a three-process slab-exchange regression test — the single-peer loopback test could not catch a serial-acceptor deadlock. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01LMbYEexVS2Qe6KBLoqjBrW Signed-off-by: xiaguan <751080330@qq.com>
The prefill wrapper refused any call where t_kv * k_stride_tok overflowed int32 — with 96 heads that is ~116k tokens, so the whale's 256k single-superstep prefill fail-stopped the whole EP16 fleet (CUDA_ERROR_INVALID_VALUE from the guard, not the kernel). The overflowing product was only ever the batch stride, and with b = 1 the batch coordinate is always 0: the value is never used to offset anything. Real addressing runs on the per-token and per-head strides, and TMA's in-descriptor address math is 64-bit. Pass 0 batch strides and drop the guard. Adds an ignored GPU depth test driving the kernel at 260k tokens x 96 heads: identical K rows make every causal softmax uniform and V rows varying only per head make the expected output exactly the per-head value, so any high-offset addressing error surfaces as a wrong element. Passes on GB300 in 2.7s. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01LMbYEexVS2Qe6KBLoqjBrW Signed-off-by: xiaguan <751080330@qq.com>
The padded per-row families run at the covering chunk bucket (4224/8448/16896) — only attention is varlen. Pure leveling stretches the head segment past the bucket the mean sits in whenever the mean runs near a boundary: at 65k over 8 ranks (mean 8,140) the leveled head crossed 8,448, landing that rank in the 16,896 bucket and doubling its padded rows while the lockstep superstep waited — a measured ~900ms TTFT step (2,502ms at 57k, 3,638ms at 65k, then back to the smooth line at 73k). Cap segments at the mean's bucket: every rank stays in the same bucket and leveling still balances the attention triangle inside it. Measured on the CP8 2-tray fleet: 65,116 tokens 3,638 -> 2,747ms (5.07x over CP1 local), neighbors unchanged (57k 2,500ms, 73k 3,954ms, 128k 6,382ms), 64k output unchanged. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01LMbYEexVS2Qe6KBLoqjBrW Signed-off-by: xiaguan <751080330@qq.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01LMbYEexVS2Qe6KBLoqjBrW Signed-off-by: xiaguan <751080330@qq.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01LMbYEexVS2Qe6KBLoqjBrW Signed-off-by: xiaguan <751080330@qq.com>
…ks levers Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01LMbYEexVS2Qe6KBLoqjBrW Signed-off-by: xiaguan <751080330@qq.com>
Free framing functions replace the zero-state FrameWriter/FrameReader newtypes; Hello and the peer table carry a rank Range instead of re-derived first+count pairs; the slab body codec is written once; the four-arm exchange match collapses to a let-else; the host-only addr field existed for one test that now uses the bind-and-drop pattern its siblings use. A peer send that fails mid-frame now poisons the hub so a sibling rank cannot append into the abandoned frame. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01LMbYEexVS2Qe6KBLoqjBrW Signed-off-by: xiaguan <751080330@qq.com>
…ases Doorbell: flag_addrs/flag_count name what the parameters hold, the arg struct zero-inits, the single-call error mapper inlines, and the header carries the same-stream ordering requirement instead of a fence argument. FlashMLA: the zero-batch-stride comment shrinks to the b=1 precondition; the deep-context test doc drops its changelog half. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01LMbYEexVS2Qe6KBLoqjBrW Signed-off-by: xiaguan <751080330@qq.com>
Scheduler side: WhaleDuty::Free is a unit variant (is_quiet carries the single-launch contract), the posted whale rides its RequestId through the descriptor instead of a parallel key counter, and the merged fake executor serves both the CP and whale integration worlds. Data plane: the scratch constructors share one body, arm_fleet refuses a non-increasing whale seq (doorbell values would alias), the per-window gang re-checks move to arm time, and the whale admission floor is derived from K3_WHALE_SEGMENT_FLOOR instead of a bare 4096. The executor re-checks k3_whale_admits on the cross-process descriptor before deriving segments. Docs: cp-lane-design's doorbell paragraph now describes the SM-kernel ring (not the pre-probe memops design) and records the fleet-wide chunk_tokens agreement the segment derivation assumes. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01LMbYEexVS2Qe6KBLoqjBrW Signed-off-by: xiaguan <751080330@qq.com>
k3_whale_width/admits/gang/segments and their tests move from scheduler/whale.rs to executor/cp.rs, beside the CP admission and split they parameterize — the rendezvous module keeps only the protocol state machines. The executor's whale startup sequence (arm slab, import world, install gang, re-arm scratch) moves from executor/mod.rs into whale_gang.rs impl blocks, so the fleet data plane is one file. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01LMbYEexVS2Qe6KBLoqjBrW Signed-off-by: xiaguan <751080330@qq.com>
k3_whale_segments drops the total == 0 early return (admission rejects before any caller reaches it, and the general loop already yields width segments), making its always-width doc true. The doorbell comment stops claiming a kernel fence orders the publication (stream order does), the PostedWhale doc stops claiming cancels pair by id (they pair by seq), and k3_whale_prompt_hash goes private. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01LMbYEexVS2Qe6KBLoqjBrW Signed-off-by: xiaguan <751080330@qq.com>
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 0a4cdf6277
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| self.model.rank | ||
| ); | ||
| let layout = K3WhaleSlabLayout::new(world, k3_chunk_bucket(self.chunk_tokens)?); | ||
| let (base, wire) = k3_whale_slab_alloc(self.ctx.device_ordinal, &layout)?; |
There was a problem hiding this comment.
When --k3-weight-staging loads multiple local ranks, load_rank_executors performs every load on scoped worker threads, leaving the launch thread without a current CUDA context after they join. This call then reaches cuMemsetD8 and cuCtxSynchronize in k3_mega_fabric_slab_alloc without calling bind_thread/gpu.set_current, so the supported staging-plus-whale configuration can fail startup with an invalid-context CUDA error. Bind this executor's context before allocating the slab.
Useful? React with 👍 / 👎.
Comment on lines
+872
to
+874
| pub fn drain(&self, rank: GlobalRank) -> Vec<WhaleToMember> { | ||
| drain_mailbox(&self.mailboxes, self.first_local, rank) | ||
| } |
There was a problem hiding this comment.
If the TCP reader detects a disconnection after a helper has sent Ready but before it receives Commit or Cancel, it records failed, yet this boundary check silently returns an empty mailbox and never reports that failure unless a later outbound send happens. The member therefore remains armed indefinitely and stops admitting prefills; worse, a connection failure during commit broadcast can leave only part of the gang entering the CP exchange. Make every scheduler-boundary drain observe the stored failure so the fleet fails loudly instead of hanging or partially committing.
AGENTS.md reference: AGENTS.md:L110-L114
Useful? React with 👍 / 👎.
…thread Two full-model fleet failures the pruned gates never reached: - The stream-ordered pool pins its release threshold at MAX, so after a 1.5T weight load it holds ~66 GiB beyond the model's accounting and cuMemCreate(FABRIC) cannot see any of it. The slab allocator now trims the default pool and retries once when the first attempt reports OOM, logging the pool's reserved/used footprint as it does. - Under staged (parallel) loading the engine-launch thread never binds a CUDA context, so the whale slab's VMM calls fail with INVALID_CONTEXT. arm_whale_slab binds the executor's context first. Serial loading masked this by leaving the last rank's context current. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: xiaguan <751080330@qq.com>
…edger Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: xiaguan <751080330@qq.com>
…x_ctx scale The chunked-prefill MLA attention materialized expanded K/V for the whole context in one shot: 21 GiB of scratch at ctx=262144 (the 74x latent inflation), which is what pushed the full 1.5T model past the 276.6 GiB HBM ceiling and forced the acceptance run down to ctx=135168. The expansion scratch now covers min(max_ctx, 16896) rows. A context that fits takes the exact single causal FMHA it always did — bitwise unchanged, existing gates untouched. A deeper one walks the past in dense windows and closes with a causal t_q x t_q call over the chunk's own keys, folding each window's output through the FMHA's log-sum-exp into an f32 accumulator: - The SM100 FMHA is templated over mask x scheduler: the causal entry keeps CausalMask + its TileH=8 scheduler; a new dense entry pairs ResidualMask with the individual tile scheduler (t_q and t_kv unrelated, so a ragged context tail narrower than the chunk is fine). Both write the f32 [heads, t_q] LSE when asked. - k3_mla_prefill_lse_merge folds one window into the accumulator (one block per (q, head) row, no cross-block coordination); o_finalize converts the merged f32 back to bf16. Only the merge order differs from a single call, so the paths agree to bf16 rounding. - gemm_rows_span_into_checked expands kv_b over a token span of the gathered latent; expand_k gains the matching rope row offset. At 262144 the expansion buffers drop 21 GiB to 1.4 GiB against ~210 MB of new accumulators; only the 576 B/token latent/rope gather still scales with max_ctx. That puts the full model's 256k serving budget back under the ceiling. Verified on GB300: new windowed-vs-single equivalence GPU test (ragged tail included), the full k3 gate suite 6/6, and the cp_prefill CP4-vs- CP1 gate at a 65536-token prompt, which drives the window walk on both lanes. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: xiaguan <751080330@qq.com>
…ared Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: xiaguan <751080330@qq.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: xiaguan <751080330@qq.com>
…padding The buckets are already geometric (256..16896, worst-case <=2x) and the mega dispatch sends only live rows above one rank, so MoE never pays for padding. The short-prompt gap is the ~270ms coordination intercept plus fixed costs; the lever is de-bucketing the local dense families to live_rows and profiling the intercept, not adding bucket steps. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: xiaguan <751080330@qq.com>
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.
feat(k3): cross-process CP whale gang — fleet-wide long-prompt prefill
What / invariant changed
K3's context-parallel prefill lane previously stopped at process boundaries
(
PEGAINFER_K3_CP= local ranks only). This PR extends it across the whole EPfleet: a long prompt admitted on any endpoint is split across up to 16 ranks on
4 machines ("whale gang"), each rank prefills its context segment in lockstep
supersteps, and KV/activations flow between processes over CUDA fabric (MNVL)
slabs with SM-kernel doorbells. The EP free-running discipline (CLAUDE.md) is
preserved: no rank ever waits in the engine loop; the gang is booked as agreed
future launch steps; padding ranks step constructively.
New pieces:
scheduler/whale.rs— gang planning: depth-leveled segment split (quadraticattention model, capped at the mean's chunk bucket), rendezvous state machine.
scheduler/whale_hub.rs— TCP hub: per-connection threads, slab-tableexchange, sequenced plan broadcast.
executor/whale_gang.rs+executor/cp.rs— fabric slab data plane;doorbells ring via a one-thread SM kernel (GB300 stream memops reject
fabric-imported VAs; waits stay memops on local slabs).
never used to offset; pass 0, delete the INT_MAX guard). Includes an
#[ignore]260k-token GPU test.Behavior tightening:
PEGAINFER_K3_WHALE_MINnow refuses values below 4096(2× the 2048 segment floor) — a narrower floor could not split into two
legal segments anyway.
Evidence (all on GB300 trays, pruned-224 K3 checkpoint)
Correctness gates
prompts byte-identical across every fleet endpoint vs the local lane;
16k + 48 greedy tokens coherent and identical on all endpoints.
(post-ceiling-fix); FlashMLA
deep_context_attention_addresses_past_int32passes on hardware (2.7s).
Same-fleet A/B — TTFT min ms (CP1 local lane vs whale, identical fleet/binary)
CP16 matches CP8 at these lengths (intercept-bound; 28.5k: 1,228ms, 4.79×).
nsys anatomy (whale CP8@128k, deep rank): superstep wall 6,233ms =
FMHA 35.8% + mega MoE 19.8% (pairing-wait-inflated) + dense GEMM 15.5% +
elementwise 10.4% + KDA 4.2% + gaps 6.7%. Measured Amdahl cap 6.4× at 128k;
we serve 93% of it. Full account and lever ranking in
docs/models/k3/cp-lane-design.md.
Fixed en route (each hardware-reproduced)
accept loop (masked at CP8, deadlocked CP16). Per-connection threads.
(
cuStreamWriteValue64→CUDA_ERROR_INVALID_VALUE); doorbell writes movedinto an SM kernel.
0.0.0.0.
sat near it → padded-row families doubled fleet-wide (~+900ms at 65k). Segments
now cap at the mean's bucket. 65k TTFT 3,638 → 2,747ms.
Not in this PR: 256k single-superstep gate and the full-896-expert EP16 vs
vLLM TP16 rematch (blocked on free trays); FMHA striping / superstep graphing /
finer bucket ladder (profile-ranked follow-ups, see cp-lane-design.md).
Tests
cargo test --release --workspace --libgreen (includes new no-GPU tests:gang consensus, slab-table exchange with multi-peer processes, leveling
bucket-cap, hub loopback).
🤖 Generated with Claude Code