Repository navigation
Conversation
) * refactor(qwen35): adopt the shared step contract Signed-off-by: RicardoMin <17879681016@163.com> * refactor(qwen35): address step contract review Signed-off-by: RicardoMin <17879681016@163.com> * test(qwen35): cancel harness requests on drop Signed-off-by: RicardoMin <17879681016@163.com> --------- Signed-off-by: RicardoMin <17879681016@163.com>
pegainfer-project#1125) Signed-off-by: Feathbow <feathbow@gmail.com>
pegainfer-project#1126) Signed-off-by: Feathbow <feathbow@gmail.com>
sysuTobo
force-pushed
the
feat/gemma4-tp2-concurrent
branch
from
October 4, 2026 12:03
776f10c to
aca34ae
Compare
…r-project#1129) * feat(kernels): make the gated RMSNorm activation selectable Qwen3.8-Flash-Next declares `output_gate_type: "sigmoid"`, and its reference implementation passes that field straight to the gated RMSNorm on the Gated DeltaNet output. The in-tree kernel hardcoded SiLU with no way to select, so serving this checkpoint would have computed all 36 linear-attention layers with the wrong gate. Template the kernel on the activation and add a second entry point rather than changing the existing signature: `rms_norm_gated_cuda` keeps its exact semantics and no Qwen3.5 call site moves. The SiLU branch also keeps its original `g / (1 + exp(-g))` expression instead of the equivalent `g * sigmoid(g)`, because the two are not bit-identical in fp32 and Qwen3.5's gates are bit-exact against it. Verified by diffing the PTX of this file before and after at `-arch=compute_89`: the SiLU instantiation is instruction-for-instruction identical (147 instructions), and the sigmoid instantiation differs in exactly one arithmetic slot (`div.rn.f32` -> `rcp.rn.f32`), the rest being basic-block label names. Gated on an L20 (sm_89) against an f64 CPU reference at the Flash-Next GDN geometry (48 value heads at head_dim 128), which also pins `silu_out == gate * sigmoid_out`; the four pre-existing norm parity tests still pass bit-exactly. Three things the change touches that are not the kernel: - The new parity gate ranks the worst element by its margin against the actual tolerance and checks every element. The bound is affine in the reference magnitude, so ranking by raw deviation and then asserting that affine bound orders differently, which lets a large-magnitude element past its own bound while a small-magnitude element, ranked higher, passes. - No `#[allow(clippy::too_many_arguments)]` on the new entry point: the workspace already allows that lint, so the attribute is a no-op. The one this change removes from `rms_norm_gated_batch_into` is the same attribute. Two more elsewhere in the file are left alone as unrelated to this change. - `docs/models/qwen35/support-qwen38.md` cites this file to argue the 27B line must not implement `output_gate_type`. That citation pointed at line 58, which is the SiLU expression on the base revision and moves to line 70 here, so the pointer is repaired and the claim scoped to `qwen3_5` — the one architecture that never reaches the new entry point. Scope note: `pegainfer-core::ops` re-exports what model crates currently call, not every kernel op, and this one has no caller yet because Flash-Next's GDN forward is slice B (pegainfer-project#1127). `pegainfer-kernels` is a direct dependency of every model crate, so the new entry point is reachable; the `pegainfer-core` re-export lands with its first caller rather than as unused surface here. Note this is the inverse of the Qwen3.8-27B lesson, where a field of the same name is inert and must not be implemented. Read the field's consumer, not its name. Refs pegainfer-project#1104 Signed-off-by: bohaotu <bohaotu@pony.ai> * fix(kernels): pass the gated norm activation as a parameter, per review Review asked for one public entry point rather than a second function plus a wrapper that nothing calls until pegainfer-project#1127: `rms_norm_gated_batch_into` now takes `GatedNormActivation`, Qwen3.5's two call sites pass `Silu`, and the enum crosses `pegainfer-core::ops` so they can name it. Both `extern "C"` entries are unchanged, so the PTX evidence for the SiLU instantiation still holds. Also from the review: - `assert_close` rejects non-finite values on either side. Ranking by margin silently skipped a NaN, so an all-NaN output passed; a CPU test guards it. - The two redundant gate-relation assertions are gone, along with the note justifying them (wrong at g = 0, where silu is 0 and sigmoid is 0.5): the two distinct references already fail if the template collapses. - `norm.cu`'s comment gives the real reason for keeping `g / (1 + exp(-g))` (bit-for-bit for Qwen3.5) and says its golden gates would not catch the swap. - `support-qwen38.md` names the SiLU branch of `rms_norm_gated_kernel` instead of citing a line number that the change moved. Signed-off-by: bohaotu <bohaotu@pony.ai> --------- Signed-off-by: bohaotu <bohaotu@pony.ai>
…project#1115) * feat(gemma4): serve Gemma 4 31B bf16 at tensor parallel 2 One rank per device: per-layer-kind shards of the text tower, one all_reduce after o_proj and one after down, and a per-rank GemmaKv, all driven lock-step from the existing single scheduler thread. The 31B bf16 checkpoint then serves on the 48 GB-class sm_89 devices (L20/L40S/RTX 6000 Ada), which cannot hold its ~59.8 GiB of weights at TP1. Invariant: a request is now a sum of P sharded partial products instead of one whole pass; both reductions run before the residual they feed is formed. Two environment hazards found and fixed, in docs/models/gemma4/tp.md: a cold-process cublasLt lazy load deadlocks against NCCL's module loader, so each rank's kernels are warmed before the communicator exists; ncclCommAbort wedges with live captured graphs, so EngineState::drop releases them first. Decode graphs also work under TP now. Verified on 2xL20 (sm_89, CUDA 12.9, NCCL_PROTO=LL128): fmt and clippy clean, gemma4 lib suite 63 passed, TP2 parity 48/48 with worst picked-token gap 0.0000 eager and captured, real 31B TP2 cold start 30 s, TTFT 57 ms, 49.4 ms per step, 143 tok/s at c8. Snapshot: docs/benchmarks/gemma4-31b-tp2-l20.md. Fixes pegainfer-project#1114. Signed-off-by: bohaotu <bohaotu@pony.ai> * fix(gemma4): keep every tensor-parallel rank on one collective stream Rank 0's step and the extra ranks' steps share one collective sequence, so a rank-0 failure must never skip the peer: `launch_staged` and `step_plain_segment` propagated rank 0's error straight out, which left the peer's matching calls unpaired and shifted the pairing for every later step. Both now hold rank 0's verdict until the extra ranks have run, and one shared helper (`step_extra_ranks`) owns the drain and the frontier check. Also from the TP2 review audit: - cross-rank invariants are inspected now: the extra ranks' frontiers after every prefill (always) and the free-page counts after every admission (debug builds) - `validate_for`'s global GQA check runs only for the sharding branch, where the per-rank group can actually split; the replicate branch's group of one made it vacuous - a multi-rank launch checks the visible CUDA device count first, and startup logs the NCCL version and the effective NCCL_PROTO, warning when it is not LL128 — the requirement this container needs and nothing enforced - `RankKv::new` starts at a sentinel id rather than 0, so a state built outside `GemmaKv::multi` cannot collide with a real request's steady-decode fingerprint - the stale `Send`/`Sync` rationale is corrected and `Sync` is gone: the decode-graph sweep walks the ranks on one thread - the TP2 gate serves its first prompt alone before the batch, so the solo admission path is covered on every run Verified on 2xL20 (sm_89, CUDA 12.9, NCCL_PROTO=LL128): fmt and clippy clean, gemma4 lib suite 63 passed, TP2 gate 48/48 with worst picked-token gap 0.0000. Signed-off-by: bohaotu <bohaotu@pony.ai> * fix(gemma4): drive rank 0 before the extra ranks in the pipeline collect The audit's path-B fix covered `launch_staged` and `step_plain_segment`, but `decode_round_collect` — the batched decode the staged pipeline drains — had the extra ranks added *ahead* of rank 0 with a `?`, so a failure on an extra rank returned before rank 0 issued its half of the step. That is the mirror of path B: the ranks share one collective sequence, so whichever half is skipped shifts the pairing for every later step. Rank 0 now goes first, its verdict is held, the extra ranks are driven unconditionally, and rank 0's failure is reported through the same `fail_active_batch` path as before. Verified on 2xL20 (sm_89, CUDA 12.9, NCCL_PROTO=LL128): fmt and clippy clean, gemma4 lib suite 63 passed, TP2 gate 48/48 with worst picked-token gap 0.0000. Signed-off-by: bohaotu <bohaotu@pony.ai> * refactor(gemma4): let the weights own their shard, and name the KV rank Three hardening items from the PR review. No behavior changes: - `Gemma4Weights` now carries the `TensorParallelConfig` it was cut as, and `GemmaServe::new` reads its geometry and pools from there instead of taking a second `tp` argument: the weights and the serve can no longer disagree about which rank they belong to. - `upload_experts` and `upload_w4a16` state in their docs that they stage whole tensors, so they are only reachable at world size 1 — where `validate_for` refuses routed and W4A16 checkpoints under TP. - `GemmaKv::core` / `core_mut` name the rank that was asked for when it is past the twins, and `multi` documents the one-twin-per-rank contract. Verified on 2xL20 (sm_89, CUDA 12.9, NCCL_PROTO=LL128): fmt and clippy clean, gemma4 lib suite 63 passed, TP2 gate 48/48 with worst picked-token gap 0.0000. Signed-off-by: bohaotu <bohaotu@pony.ai> * docs(gemma4): say which checkpoint the TP2 gate is about The gate reads whatever `PEGAINFER_TEST_MODEL_PATH` points at; only its `#[ignore]` reason and the tp.md paragraph named a pinned 12B, which reads as if the 12B were the point of the check. Say instead that the checkpoint is what picks the branch — the shipped 12B takes the replicate branch, a 31B geometry takes the shard branch production serves — and why a full-depth 31B cannot be gated here at all: the gate needs a single-rank baseline, and 57 GiB of weights do not fit one card. The six-layer synthetic is that shape at a depth that fits. Signed-off-by: bohaotu <bohaotu@pony.ai> * fix(gemma4): drive every TP rank before any one is waited on The two-rank scored path read its logprobs back on rank 0 inside the step, before the peer's tower was launched, so the readback waited on a collective that had no peer and the scheduler hung. Launch every rank's tower first and read back afterwards; the same ordering now governs the extra ranks (launch all, then drain), a rank-0 tower failure aborts the communicators instead of driving a peer that would wait forever, every arena resets its bucket floor, and a failed load releases each rank's graphs before its communicator drops. The reduction sends only the live rows, the gate compares the two runs' top logprobs (not just the picked token) against both one-rank runs, the TP2 arms carry the same envelope overrides as the rest, and the gate runner claims one device reservation for both gate types. Gated on two L20s: the one-rank comparison 48/48 picks at a 0.0000 gap, the scored probe on the real 31B, and the real 60-layer 31B against its Hugging Face dump (7/8 rows, worst shared-token logprob gap 0.6753). A long prompt still stalls the prefill -- a host sync mid-tower before the peer launches -- which docs/models/gemma4/tp.md records as a known bound. Signed-off-by: bohaotu <bohaotu@pony.ai> * fix(gemma4): answer the second review round Fail-closed: a TP failure now aborts the comms (releasing each rank's graphs first, as teardown does), marks the engine broken, and every later step refuses — so a comm-less reduction can never return partial sums. The mixed step checks rank 0's verdict before driving the peers, the stale "drive the peer either way" comments and the tp.md bound say what the code does, the gate stops comparing a request at its first differing pick, the redundant repeat comparison is gone, and the probe submits both requests before draining either. The single-threaded driver has a known limit: a prefill past ~100 tokens reaches a host-side device sync on rank 0 before the peers have launched, and hangs. The TP prompt ceiling refuses such a prompt (`PEGAINFER_TP_MAX_PROMPT`, default the `TP_MAX_PROMPT` constant, `off` to lift) rather than hang, and the scored probe runs at that same ceiling so guard and gate agree. Verified on two L20s: the synthetic one-rank comparison 48/48 at gap 0.0000, the scored probe at the ceiling, the ceiling refusing 128, and the real 31B against its HF dump (7/8 rows, worst shared-token gap 0.6753). Signed-off-by: bohaotu <bohaotu@pony.ai> * fix(gemma4): answer the third review round A walk tail after a broken pair no longer scores its first token from partial sums: `finish_plain_walker`'s swallow-and-continue arm returns the error when `tp_broken` is set, as the solo path does. Under tensor parallelism the mixed step's gathered rows are a prefill, so the ceiling bounds their **sum** — four 64-token prompts beside the decode rows were never measured complete. A single **named** device now takes the drop path: the selector is still never substituted, but naming fewer than a two-rank run needs drops its two-rank gates (as one-GPU hosts do) instead of dying — only naming none at all is an error. The scored probe's first request asks for four tokens, so the second is admitted beside a live batch rather than beside one that finished at admission. The stale `abort_comms` doc above the current one is gone, and the docs state the prompt limit the ceiling imposes — `tp.md`, `serving.md`, the benchmark TL;DRs and the root env list all said 8192 while the ceiling refuses 64. Verified on two L20s: the synthetic one-rank comparison 48/48 at gap 0.0000, the scored probe at the ceiling, the ceiling refusing 128, and the real 31B against its HF dump at 7/8 rows and 0.6753. Signed-off-by: bohaotu <bohaotu@pony.ai> * fix(gemma4): answer the fourth review round A chunked walk's rounds gather across prompts, which the TP gates never covered, so `PEGAINFER_MIX_CHUNK_TOKENS` is now refused under tensor parallelism instead of being left to build an unmeasured mixed step. The TP ceiling bounds a whole step, not just the prompt. A newcomer rides beside the live decode rows only while its remaining prompt rows plus those rows fit, and the gather budget leaves room for the decode rows too — before, the budget counted the gathered prompts alone, so four 64-token prompts beside the decode rows could exceed the ceiling that was never measured complete. The scored-prompt ceiling test now sits on the boundary — a prompt exactly at the ceiling passes, one token past it is refused — instead of `Some(0)` / `Some(usize::MAX)`, which never touched either side. The 31B HF fixture check matches the artifact's exact name rather than a `*31b*` substring that would also accept a fixture dumped for a different 31B. Signed-off-by: bohaotu <bohaotu@pony.ai> * fix(gemma4): answer the fifth review round `abort_comms` returns early when there are no extra ranks. At world size 1 there is no collective to abort and no peer frontier to keep in step, so this is the ordinary single-rank prefill failure the caller answers by failing that request — but it set `broken`, which the solo path reads to stop the whole engine, so a request-local error such as a scratch allocation took the scheduler down with it. A pick flip in the two-rank gate now measures the two picks' own gaps. Before, `worst_pick` accumulated only while the runs picked the same token, so a token that traded places (one run `(-0.1, -5)`, the other `(-5, -0.1)`) left every bound at zero while the same token moved 4.9 nat; only the printed `worst_shared` saw it. The flip now folds both picks' gaps in before its `break`, `worst_top`'s doc no longer claims it catches a flip (it does not), and the one-rank control compares the two drained runs exactly — ids and logprobs — instead of through the same tolerant metrics, which can pass while tokens trade places. The TP-ceiling asserts move out of `only_a_scored_prompt_is_bound_to_the_whole _prompt_ceiling` into a test named for them, off a shared `ceiling_probe`. The gate runner drops the HF fixture's file-name check: `golden_bytes` already holds the fixture's manifest against the checkpoint's digests, so the name could not catch a mismatched dump — only refuse a valid one named otherwise. The `tp.md` "How a step runs" line loses "and drains it" (only the prefill path drains the extra ranks; decode and the mixed step do not), and Known bounds loses the second long-prompt entry the ceiling entry supersedes. Signed-off-by: bohaotu <bohaotu@pony.ai> * fix(gemma4): answer the sixth review round A self-audit of this PR at its fifth-round head turned up one regression on the single-rank path and a set of documentation claims that the code does not make. This commit answers both. The scored-prompt path returned a rank-0 failure straight out of `admit_and_prefill`, past the `tp_broken` gate that prices every other prefill failure. `Err` from `Scheduler::step` means "the engine is beyond use", so at world size 1 — where the fifth round made `abort_comms` a no-op precisely to keep a request-local fault request-local — a `prompt_logprobs` request that failed in `launch_prompt_tower` or in the readback's own scratch allocation took the whole engine down with it. That is what `main` answers by failing the one request. Both now land in `stepped`, so the gate decides: a broken pair stops, a single rank keeps serving. A readback failure is request-scoped at any world size, because by then every rank's tower has completed and its collectives are paired. `ready_rows_pinned` was the one admission site whose extra-rank failure left the row in the batch: rank 0 had already reserved, and the caller's fallback reserves again, so the pair would carry a two-token frontier skew into a decode step that nothing inspects (the frontier check only runs on prefill). It returns `Result` now and an extra rank that cannot match rank 0's reservation stops the engine, which is what `admit_extra_ranks`' own contract already said nothing could repair. `PEGAINFER_TP_MAX_PROMPT=off` lifts the guard on a hang this build cannot recover from, and said nothing at startup; it warns now, naming the wedge. A gate covers the single-rank half of the first fix on both prefill paths. It needs no injection hook: a token id outside the embedding fails in `prepare_single`'s `validate_tokens` before the tower allocates anything, so the fault is deterministic, and the property asserted is the one that regressed — the engine still serves the next request. The documentation changes are all claims the code does not make. The pool table converted MiB to decimal GB and labelled it GiB, overstating the default point by 1.06 GiB and contradicting the binary rows two lines below it; the corrected 44.29 GiB plus the ~0.90 GiB of context and workspace the same run reports still misses an L20's 44.99 GiB usable, but by ~0.2 GiB rather than by the implied margin. The world-size rule block stated the per-rank global GQA check unconditionally, which the third round made conditional in the code, and omitted `P > 0` and `rank < P`. "Bit-identical" described a printed measurement as an assertion; the two-rank comparison asserts against 0.5, so the line is now labelled the first cut it is, beside `LOGBROB_LINE`'s 1.3x of margin. "A prefill drains every non-primary rank's stream" was true of the solo, scored and walker prefills and not of a mixed step. The refusal list in `serving.md` was missing the chunked walk and the prompt cap, and "single GPU" survived in `CLAUDE.md`, `README.md` and the index's serving row. Signed-off-by: bohaotu <bohaotu@pony.ai> * refactor(gemma4): one helper drives the extra ranks of a decode-shaped step Signed-off-by: bohaotu <bohaotu@pony.ai> * test(gemma4): calibrate the TP parity line from the drift NCCL actually produces The two-rank parity gate failed a correct run. Measured on two L20s with `NCCL_PROTO=LL128`, the 12B replicate branch's worst picked-token logprob gap is 0.5738 on this container's stock NCCL 2.18.3 and 0.379 on the 2.32.3 the recorded run used — and 0.5738 is bit-identical at every commit of this branch, `14017842` included, so the spread is the environment's and not the code's. NCCL picks the reduction order, so the drift it leaves behind moves with the build; a 0.5 line sat inside that spread. `LOGBROB_LINE` is 1.0 now: ~1.7x over the worst measurement, and the same value `DRIFT_LINE` uses for the Hugging Face comparison. That is the shape of the one calibration precedent this repo has for a quantity of this kind — `serve_oracle`'s `neutral_scale`, two algorithms over one context, 0.31..5.75 observed and a line at 12.0, roughly 2x the worst drift seen. The line's job is a structural error, a wrong shard or a missing reduction or a rank out of step, which moves a logprob by many nats and flips most picks rather than 2 of 48; reduction noise is what the near-tie rule beside it polices. Both measurements go in the doc as two columns naming their NCCL build, and the gate prints the version it ran on, so the next number can be attributed instead of compared blind. The formatter already existed inline in the start-up log; it is `nccl_version()` now and both call sites use it. Signed-off-by: bohaotu <bohaotu@pony.ai> * test(gemma4): set the parity line from this gate's own floor and a fault The line's two numbers came from two different gates: the older reading batched all three prompts, kept comparing after a flip, left flipped picks out, and recorded no NCCL build, so nothing tied the 0.5 it produced to the 1.0 that replaced it. Re-measured at one head with one procedure on two L20s (`NCCL_PROTO=LL128`, NCCL 2.18.3): four runs of the 12B replicate branch each read `22/24` picks kept and `worst_pick` 0.5738, the same to four decimals every run, and one run with the `o_proj` `all_reduce` skipped — the structural error the line is for — puts the one-rank pick outside the two-rank top-8 at request 0 step 0, and reads `worst_top` 2.1895 with the near-tie rule relaxed. `worst_pick` reads 0 there because the two lists share no token, which is what the near-tie rule exists for. `LOGBROB_LINE` stays 1.0: ~1.7x over the floor and ~2.2x under the fault, and the value `DRIFT_LINE` already uses. `0.5` sits below the floor and so fails a correct run. The comment and the `tp.md` table carry these readings now and drop the claim that the spread between readings is the NCCL build's: the gate prints its build and pair because a reading needs them, not because they explain it. Two doc fixes from the same round: the 31B benchmark names `LOGBROB_LINE` instead of the literal 0.5, and the replicate-branch reading of the global head check is corrected in `tp.md` and `config.rs` — each KV head's `Q / G` query heads are split across the `P / G` ranks holding that head, so `Q / P` is `(Q / G) / (P / G)` (8 at 12B with P = 2), not `P / G` copies of `Q / G`. Signed-off-by: bohaotu <bohaotu@pony.ai> * test(gemma4): assert the compared logprobs are finite in the TP2 gates A NaN row is invisible to every comparison beside it: `worst_pick`, `worst_top` and `worst_shared` accumulate through `f32::max`, and the Hugging Face gate through `gap > worst`, and a NaN moves neither bound — while the top-k *ids* stay right, so the near-tie rules stay quiet too. One NaN logit folds its row's logsumexp into NaN, and the gate passes on a row it exists to fail. `serve_oracle::compare_row` already asserts finiteness over both arms. `assert_finite` checks the row each comparison reads from: the one-rank and the two-rank side in `distribution_gap`, and the echo row plus the reference fixture's values in the HF gate. Reproduced on two L20s by injecting a NaN row into both runs at the comparison's input (scratch patch, reverted, never committed): the pre-fix gate passes on it — `22/24 steps keep the one-rank pick; worst picked-token logprob gap 0.5738`, the NaN step silently dropped from every bound — and the fixed gate fails at `tp2: request 0 step 0 (one rank) scored a non-finite logprob: NaN with top [(236770, NaN), ...]`. The other logprob gates needed nothing: they compare per value (`delta <= LOGPROB_TOLERANCE`), where a NaN fails `<=` instead of slipping past an accumulated bound. Signed-off-by: bohaotu <bohaotu@pony.ai> --------- Signed-off-by: bohaotu <bohaotu@pony.ai>
…tall The single-threaded driver ran rank 0's whole segment before any peer's, so a host-blocking operation inside it — rank 0's per-row readback, and a GEMM workspace allocation behind it — waited on a collective whose peer call had not been issued yet, and the scheduler hung. Admission refused a prompt past the length the gates had been verified at (`PEGAINFER_TP_MAX_PROMPT`) to keep that from being reached. `EngineState::drive_ranks` runs rank 0's segment on the scheduler thread and each extra rank's on its own scoped thread, then joins, so every peer's collectives are in flight before rank 0 can block. The payload is transposed once up front — `rank_payloads`/`prefill_payloads`, with `GemmaKv::split` lending each rank one mutable family per request — because the threads need disjoint borrows. The six rank-driving paths (solo, scored and chunked prefill, the mixed step, both decode paths) go through it, and the ceiling, its parser, its knob and its unit assertions are deleted rather than left as scaffolding. A rank-0 failure releases rank 0's graphs and detaches its communicator *before* the join, so a peer already waiting on a call that will never come is unblocked; the extras' teardown follows the join, and the engine is marked broken only when there is a peer to keep in step — so a world-size-1 failure stays request-scoped, which `a_failed_prefill_costs_that_request_not_the_engine` pins. Verified on 2×L20 (`NCCL_PROTO=LL128`): the stall is gone — a scored prefill of 64, 128, 256 and 1024 tokens all pass on the 60-layer 31B, where 128 and up used to hang. The 12B replicate parity keeps `22/24` picks at gap `0.5738`; the 31B-geometry synthetic keeps `48/48` at `0.0000`; the real 31B against its HF dump keeps `7/8` rows, worst shared-token gap `0.6753`; the new prefix gate reads `4.0538` against its `12.0` line. `fmt`, `clippy --all-targets -D warnings` and the single-rank suite (63 passed) are clean. Signed-off-by: bohaotu <bohaotu@pony.ai>
sysuTobo
force-pushed
the
feat/gemma4-tp2-concurrent
branch
from
October 7, 2026 02:03
aca34ae to
9d000cb
Compare
Owner
Author
|
Superseded. |
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
4 participants
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.
Stacked on pegainfer-project#1115 (base is its head branch). This is the "B" half of the split: pegainfer-project#1115 carries the TP2 feature and refuses a prompt past the length the single-threaded driver has been verified at; this PR removes that limit by driving the ranks concurrently. Retarget to
mainonce pegainfer-project#1115 merges — the diff below is B only.What changed
One thread per rank.
EngineState::drive_ranksruns rank 0's segment on the scheduler thread and each extra rank's on its own scoped thread, then joins — so by the time rank 0 reaches a host-blocking operation (a GEMM workspace allocation, first reached at 128 prompt tokens) every peer's collectives are already in flight. The six rank-driving paths (solo, scored and chunked prefill, the mixed step, both decode paths) go through it;kv.rsgrowsGemmaKv::split/worldfor the per-rank payloads, andlayer.rsreduces only the live rows.The prompt ceiling is gone: admission no longer refuses long prompts, and
PEGAINFER_TP_MAX_PROMPT— with its parser, its knob and its unit assertions — is deleted rather than left as dead scaffolding.A rank-0 failure inside
drive_ranksaborts rank 0's communicator (releasing its graphs first, as teardown does) to unblock the peers before the join, sets the engine broken, and every later step refuses — so a comm-less reduction cannot return partial sums.Evidence (2×L20,
NCCL_PROTO=LL128)the_two_rank_engine_is_prefix_consistent(no reference needed): the same 512-token prefix scored in a 1024-token run differs by 4.05 on the 31B and 1.69 on the four-layer synthetic, from row 0 on. That is the shape-dependenceserve_oracle'sneutral_scalealready measures at 0.31..5.75 raw logits against a 12.0 line, so the gate's own line is on that scale; the 1024-token HF case needs that gate's tolerance-plus-top-1-share discipline, which is the remaining follow-up.