Repository navigation
Merged
Conversation
`kern cut <manifest> --at <calls> --module pp_mailbox.cubin --out <dir>`
splits a verified manifest into S stage manifests. Liveness across each
cut point finds the boundary workspaces; they go through a send/recv op
pair over a mailbox on the downstream GPU (tools/kernels-src/pp_mailbox.cu:
full/empty counters, system-scope release/acquire, a device-side spin that
traps after 60 s). Each stage carries `cut {id, stage, stages}` and a
`clock` output fill the mailbox kernels stamp per item with globaltimer.
A stage's group `pp.<e>` joins two different manifests that agree only on
the peer's `of` buffer; import_peers already matches by name and size, so
the runtime is unchanged. Protocol requires a token fill on stage 0 only.
Gate: a property test (400 random manifests, symbolic interpreter) shows
every cut that is made computes the whole; qwen3-4b cut at l18 into two
stages, run in sequence (examples/pp_chain), gives bitwise identical
logits and tokens to the whole manifest on 7 prompts, chunk 512 and 64.
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Stage 0 (`--pp-listen`) is the head: HTTP, scheduler and the only pool,
sized to the fewest pages any stage has. Other stages (`--pp-head`) send
Hello {stage, cut, pages, mailbox}, get Table {downstream_mailbox, pages},
then replay the head's Item {program, vars, rows} in order; the last
stage answers tokens. Activations go GPU to GPU through the mailbox. Any
error exits the process; a closed connection ends the other side. Each
stage logs one `pp item` line per item from its clock: rows, busy,
recv_wait, send_wait.
New dependencies: tokio-util (codec) for length-delimited framing,
futures-util (sink/stream) to drive it, bytes for its frames; all three
were already in the lock through pegainfer.
Gate (tray07, qwen3-4b, 2 stages): first tokens identical to single-GPU
kern-serve on 7 prompts; 8193-token prompt TTFT 313 -> 174 ms at conc1.
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
gen_qwen3_decode.py emits `kv.0`..`kv.35` (vLLM's per-layer pool layout, 4096 bytes per token each) instead of one layer-interleaved `kv`; the pool hands every paged state the same page ids, so nothing else changes. The draft's `draft_kv` stays interleaved. The generator had not run since the registry move (dead `handwritten` import, `variant` unimported, launch `label` dropped); it runs again, so the regenerated examples also take the registry builds of argmax, embedding and weight_prep. Gate (tray07): `kern test qwen3-4b` bit-identical at every span; new vs old qwen3-4b.json bit-identical at every span; `kern run` 32 steps text identical; dspark round and --rows 1, 128 steps, text identical to before. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
… much of each it uses A stage keeps a state whole or drops it, so a whole-model state is carried by every stage that touches a piece of it. `kern cut` now prints, per stage, the regions (distinct call offsets) of each kept state its calls touch against the whole manifest's; a share below one is memory held for other stages, reported, not refused. A page or slot table whose state a stage dropped indexes a kept state of the same kind instead of losing its domain (and with it its place in the protocol). Gate (tray07, qwen3-4b per-layer kv cut before l18): each stage keeps 18 states, all used whole (the single-state manifest: `kv` 36 of 72 regions on both stages); pages per stage without --capacity 122993 -> 245760; the chain gate (7 prompts, chunk 512 and 64) and the 2-process serve gate are bitwise identical as before; sweep unchanged within noise. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
The regenerated qwen3-4b examples carry one kv state per layer and the registry builds of argmax, embedding and weight_prep; kern test and the pipeline gates on PR #42 cover the change. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
…y load A checkpoint tensor that must be transformed before any kernel can read it (K3's mxfp4 experts shuffled into the batched GEMM's layout) is declared a `source` and the transform a `derive` program. The runtime runs it call by call inside load_weights, each source staged only between its first and last reader in one shared allocation, so the peak is the derived weights plus the widest set of live sources, not the checkpoint twice. Carried over from the derive-weights work that the full-K3 stage path was measured with. Signed-off-by: JinYan Su <751080330@qq.com> Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
…age's chunk hand back the token A stage derives only the weights its own calls read, like the tables of a once program. A chunk program that writes the first token (a model whose chunked kernels must see every prompt token, K3) is no longer refused: only the last stage writes an output, so the token comes from it either way. Signed-off-by: JinYan Su <751080330@qq.com> Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
A pipeline stage holds the head's slot ids for its own layers' recurrent state, so it zeroes a fresh sequence's slot on the head's word rather than through a lease of its own. kern-run writes a line table by name, whole. Signed-off-by: JinYan Su <751080330@qq.com> Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
…that emits the token An item carries the sequence's line tables and the slot to zero on its first call, and says whether it finishes the prompt; the last stage answers only that one. Without a one-row step the whole prompt goes through the chunk program and the last chunk's token is the answer (full K3 PP8). Signed-off-by: JinYan Su <751080330@qq.com> Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
…ups from the paged table Stage k>0 of a cut chunk program gets the sequence count from the mailbox, not a seq_len input, so the protocol takes the groups axis from the first dim of the unfilled input that indexes a paged state. Signed-off-by: JinYan Su <751080330@qq.com> Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
…ompt as chunks The head runs stage 0's manifest, which cannot tell whether the last stage's chunk emits the token, so the step is optional rather than inferred from the head's own chunk. Fallible setup also runs before the stage listeners start, so its error is not masked by their exit. Signed-off-by: JinYan Su <751080330@qq.com> Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
…eqlens fill A call whose rows are set per call may now span several groups when a `cu_seqlens` fill says where each sequence's rows start; without one it is still an error. The context var becomes the call's sequences' lengths summed (a packed prefill expands every one of them), equal to the old longest-sequence value wherever groups is 1, the only place one is named today. `kern_run::fills` takes per-sequence lengths and lays the sequences back to back; `vars` takes the call's total rows. Gate: K3 PP8 packed prefill (16 sequences per 8192-row item), 192 frozen prompts conc16 on two trays: 47.0k tok/s vs 42.6k unpacked; single-sequence first tokens identical, 187/192 under load with the 5 others near-ties. Signed-off-by: JinYan Su <751080330@qq.com> Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
… prefill item A pipeline fed one prompt per item pays a full weight pass on every stage for each prompt's last partial chunk. The head now keeps the admitted prompts active and fills each item with `pack()`: oldest first, partial chunks allowed, up to the call's rows, sequences and context; a pure function with a property test that every small input drains within its bounds, in order, every row once. A full item goes out at once, a partial one only while fewer items than stages are in flight. The last stage answers every item, so the head counts what is in the pipeline and finishes each prompt from its entry in the reply; an item zeroes the slot of each sequence it starts, and a line table is sent only when it changes. Gate: K3 PP8, 192 prompts of 4.5-9.5k tokens, conc16: 47.0k tok/s (5.88k/GPU) vs 42.6k unpacked, 161 items instead of 219, slowest stage 50.1k tok/s while busy. Signed-off-by: JinYan Su <751080330@qq.com> Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
![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: 1589fb9559
ℹ️ 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".
Comment on lines
+34
to
+36
| let room = b.context.saturating_sub(context + s.pos); | ||
| let n = s.left.min(b.rows - rows).min(room); | ||
| if n == 0 || out.len() == b.seqs { |
There was a problem hiding this comment.
When the manifest's batch.context maximum is smaller than the runtime's sequence-token limit, a prompt can be admitted successfully and then become permanently stuck: after its first chunk advances pos to the context maximum, room is zero, pack returns an empty plan forever, and the request is never completed or rejected. Admission needs to account for the context bound (or packing needs to surface this as a context-length rejection) so oversized prompts cannot leave the scheduler wedged.
Useful? React with 👍 / 👎.
…er runs on Runtime::mark records a blocking-sync event on the compute stream; Mark::wait binds the context on whatever thread holds it and returns the fault of any launch before it. A caller learns when the device got somewhere without the issuing thread ever waiting. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Signed-off-by: JinYan Su <751080330@qq.com>
… host Every stage marks each item it issues and a thread beside it waits on the mark: one `pp item` line per item with its number (its place in the TCP order, the same on every stage), rows, issued_us and done_us. Nothing on the issuing thread waits any more; the device clock's half-ring read was a stream sync every 128 items. A device fault (a mailbox wait that gave up) surfaces at the item's mark and ends the process, which is how a dead stage is noticed. tools/pp_timeline.py joins the stages' lines by item into per-stage period, held and wait, and a gantt. Gate (qwen3-4b, 2 stages, tray07): gate prompts' first tokens equal to unsplit; head and stage logged 498 items each; last stage killed -> head exits; last stage frozen -> head exits at 60 s on its item's mark, the resumed stage 62 s later the same way. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Signed-off-by: JinYan Su <751080330@qq.com>
The `clock` fill made a stage's output carry globaltimer stamps, so a forward was no longer a pure function of its inputs, and its [ring x 5] layout was spread over the mailbox kernel, the cut, the protocol and kern-serve. The host now knows when each item left a device (previous commit), so the fill, `pp.clock`, the stamp kernels and the waits' clock argument are deleted; a first stage has no receive calls, a last stage no send calls. Gate: kern-manifest tests (cut property test included) pass; pp_chain on qwen3-4b cut at 218, 7 prompts, chunk 512 x 8 steps and chunk 64 x 16 steps: logits and tokens bitwise identical to the unsplit manifest. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Signed-off-by: JinYan Su <751080330@qq.com>
xiaguan
changed the title
This was referenced Oct 8, 2026
xiaguan
added a commit
to pegainfer-project/kern-k3-pp-prefill
that referenced
this pull request
Oct 8, 2026
kern's Protocol reads a call's axes off the vars (`axis: rows | groups | tray`) instead of inferring them from the slot and seq_len fills, so a downstream stage that reads no length still has its sequence axis. Needs pegainfer-project/kern#44. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: JinYan Su <751080330@qq.com>
…ger imply the axes `vars.<name>.axis` is `rows`, `groups` or `tray`, at most one var each. Protocol reads a call's axes off the vars. Before, it reverse-engineered them from the roles: the rows var was whichever var the `slot` fill's buffer was shaped over, the groups var the `seq_len` fill's, the tray var whatever third var a fill or line table used. So every forward had to carry `slot`, `seq_len` and `token`, which is a serving loop's habit written into the function's type: a pipeline stage past the first takes activations, not tokens, and reads no length, yet has both axes. The cut-stage special cases in Protocol (`cut.stage > 0`: no token required, groups guessed from the page table) are gone with the inference; Protocol does not read `cut` at all. `token_rows()` is an Option now and the drivers that feed tokens (kern run, kern test) say so when a manifest takes none; `slots()` / `seq_lens()` go, `any(Fill::...)` covers them. Schema golden regenerated; every generator and example manifest with a batch names its axes (k3: `rows` is the tray axis, as the old inference found it). Gate: kern-manifest / kern-pool / kern-run / kern-test / kern-serve tests pass; qwen3-4b cut at 218 with the axis manifest: pp_chain chunk 512 x 8 and 64 x 16 bitwise identical to unsplit, serve gate first tokens equal, killed stage -> head exits, frozen stage -> head exits at 61 s. The K3 PP8 downstream stages (no token, no seq_len) pass Protocol without a special case. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: JinYan Su <751080330@qq.com>
The five hashed examples gained `axis` on their vars with the manifest change; the gate re-pins them. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: JinYan Su <751080330@qq.com>
Wiring is deployment, not manifest; the kind carries only direction, dtype and shape so EP all-to-all can be its second user. Decode over the pipeline waits for it rather than building on `pp.*` ops. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: JinYan Su <751080330@qq.com>
The slot zeroing and line tables an item carried for a per-sequence state had no model through the pipeline (qwen3 and K3 have none); they come back with the model that needs them and the gate that proves it. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: JinYan Su <751080330@qq.com>
A buffer crossing a cut is an output on the stage before and an input or an inout (a new general buffer kind: an input the program rewrites in place) on the stage after, exported. `kern cut` only changes signatures; the manifest carries nothing pipeline-specific: the mailbox kernel, the six `pp.*` ops and the per-edge peer group are gone. Transport belongs to the serving shell: the next stage's arrivals are mapped into this one, `Runtime::push` copies on the compute stream, `Runtime::mark` says when it landed, and item numbers go over TCP. The next stage's host does not wait for the data: it issues item k behind `Runtime::hold`, a kernel spinning on a pinned count that `Gate::open` raises when "k arrived" comes in. Every signal the waiting thread handles stays off the driver: marks are a kernel writing a sequence number into pinned memory, outputs come back through `Runtime::fetch` into pooled pinned buffers, faults are the issuing thread's `Runtime::fault`. The four deadlocks that forced each of these are in docs/lessons.md. Gate (2026-10-08): qwen3-4b pp2 chain and serve first tokens identical to the unsplit manifest, faults as before; K3 PP8 c16 44.3k tok/s (mailbox 44.9-46.8k); qwen3 conc1 TTFT unchanged, conc4 5-13 ms behind the mailbox (docs/serve.md). Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: JinYan Su <751080330@qq.com>
… stop policy Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: JinYan Su <751080330@qq.com>
xiaguan
added a commit
to pegainfer-project/kern-k3-pp-prefill
that referenced
this pull request
Oct 9, 2026
… and --state-per-layer * k3: the stage kernels split a packed call's rows by cu_seqlens Varlen siblings of the prefill's per-sequence kernels, each next to the one it generalises: the KDA short-conv gather and window update and the recurrent-state load/store per sequence (k3_span_gather, k3_span_state), the latent gather placing each sequence's cached rows back to back (k3_mla_v2), the batched FMHA's length tables (k3_prefill) and each sequence's last row for the head (copy_rows). kern_k3_fmha_lens_varlen keeps cum_seq_lens_kv unpadded. The TRT-LLM gen context FMHA aligns its causal mask to the end of each cum_kv span, so a span padded to 128 rows made a 24-token sequence see 104 zero keys; with the prefix sums unpadded, packed prefill matches unpacked on every single-sequence length probed (6..1400 tokens) and on a 25k-token prompt over four chunks. The existing kernels are unchanged; the four PP4 manifests only follow the five rebuilt modules' hashes. Signed-off-by: JinYan Su <751080330@qq.com> Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * k3: flash-kda builds a varlen variant for the packed prefill With -DKERN_VARLEN the trimmed instantiation list keeps upstream's varlen launch (<128, true, true, true, true>) instead of the fixed-span one: cu_seqlens splits the rows into sequences, each with its own state block, and upstream's build_tile_prefix kernel lays out the tiles. build.sh passes KERN_DEFINES through. The cubin is bundled under prebuilt/ like flash_kda_d128, since its CUTLASS headers are not pinned here. Signed-off-by: JinYan Su <751080330@qq.com> Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * k3: the generator packs prefill sequences and keeps one state per layer --state-per-layer gives each layer its own KV / KDA state, so kern cut can split one whole-model prefill into stages that carry only their own layers' states, and drops the topology when every group is size one. --pack N makes the prefill ragged: up to N sequences per call split by a cu_seqlens fill, per-sequence KDA line tables, FlashKDA's varlen build, the context FMHA batched over the sequences (varlen_abi.py), and the head on each sequence's last row. The full model's raw experts are bound as sources and shuffled by a derive program while the weights load, so a PP8 stage never holds both copies. Gate: full K3 PP8 on two trays, 192 prompts of 4.5-9.5k tokens at conc16: packed 47.0k tok/s vs 42.6k unpacked; 187/192 first tokens identical, the 5 others near-ties (top-1/top-2 margin 0.10-0.26). Signed-off-by: JinYan Su <751080330@qq.com> Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * gen_stage: the vars say which axis of a call they count kern's Protocol reads a call's axes off the vars (`axis: rows | groups | tray`) instead of inferring them from the slot and seq_len fills, so a downstream stage that reads no length still has its sequence axis. Needs pegainfer-project/kern#44. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: JinYan Su <751080330@qq.com> --------- Signed-off-by: JinYan Su <751080330@qq.com> Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.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.
Prefill-only pipeline parallelism, from
kern cutto full Kimi-K3 (896 experts, 93 layers) on 8 GB300s. This PR folds in #42 (cut + per-stage serving) and #43 (derive weights at load), so the whole feature reviews in one place.A forward is a pure function, so cutting at a call index changes no arithmetic. The stages run the same calls on the same values, and the pipeline's outputs are the unsplit manifest's byte for byte.
What lands
Cutting (
kern-manifest)kern cutsplits a manifest at call indices. Liveness finds the workspaces that are live across a cut; each becomes anoutputon the stage before and aninput(read only) orinout(rewritten in place, a new general buffer kind) on the stage after, exported. A cut only changes signatures: the manifest carries nothing pipeline-specific, andkern-manifestknows no kernel.cu_seqlensfill (packed prefill). A downstream stage that has noseq_lenfill takes its groups from the paged table.cut { id, stage, stages }, used only to check the handshake. Thetokenfill is required only on stage 0.Derive weights at load (
kern-manifest,kern-runtime)sourcebuffer, turned into a weight by aderive: trueprogram insideload_weights.Runtime
Runtime::zero_slot: a stage zeroes the sequence slot that the head's pool handed out.Runtime::push: copy a buffer into another runtime's mapped buffer on the compute stream.Runtime::hold/Gate::open: a kernel spinning on a count in pinned memory holds the compute stream until the host raises the count, so a caller can issue work before its inputs are there.Runtime::mark/Mark::wait: a kernel writes a sequence number into pinned memory; waiting is a memory read, never a driver call.Runtime::fetchcopies an output into pooled pinned memory for a waiting thread to read;Runtime::faultis how the issuing thread hears a fault.Serving (
kern-serve)--pp-listen/--pp-head).Hello/Tableover TCP, carrying each stage's arriving buffers as fabric handles and the minimum page count across stages.Items downstream (vars, fill values, page and line tables, slots to zero) and tokens back. Page ids are shared by every stage.pp itemline (item, rows,issued_us,done_us), tells the neighbours, and on the last stage answers the head with the tokens. That thread never calls the driver: the four deadlocks that rule out each kind of call are indocs/lessons.md(2026-10-08).pack(): oldest prompts first, partial chunks allowed, stops at the call's rows, sequences and context. A full item goes out at once; a partial one only while fewer items than stages are in flight.tools/pp_timeline.pyjoins the stages' lines into per-stage period, held time and wait time, plus a gantt.Generators
tools/once.py/tools/kern_manifest.pycan emit derive programs.crates/.Results
Full K3 PP8, one GB300 per stage, two trays, chunk 8192, 16 sequences per item.
Against SGLang TP8 / DCP8 / EP8 on the same 8 GPUs. Random 2k–32k cold prompts (mean 17.9k),
max_tokens=1:The capacity rows are from the device-mailbox build (44.9–46.8k tok/s over its runs). With the host transport of this PR's head the same run measures 44.3k tok/s / 2.47 req/s (tray12 + tray17, 2026-10-08); the stage period (175–179 ms per 8192-row item) is compute, so the transport is not on its critical path.
Packed vs unpacked. 192 frozen prompts of 4.5–9.5k tokens, concurrency 16: 42.6k → 47.0k tok/s (+10.4%), 219 → 161 items, TTFT p50 2.4 → 2.1 s.
qwen3-4b, 2 stages, chunk 512 (median TTFT,
max_tokens=1, 1 / 4 / 16 chunks). conc1: 8.7 / 25.3 / 172.4 ms against 313 ms unsplit for 16 chunks. conc4: 22.7 / 90.1 / 677.4 ms, 5–13 ms behind the mailbox build (17.6 / 79.6 / 664.7): the push waits one TCP round trip for "k−1 taken", and a cuBLAS launch queued behind a closed hold blocks inside the call until the hold opens. Pages per stage: 122,993 → 245,760 with per-layer KV.Gates
kern test qwen3-4b: PASS, bitwise at every span. The per-layer and single-state manifests are bitwise identical, andkern runtext is unchanged.pp_chain(stages in sequence vs the whole manifest), qwen3-4b, 7 prompts of 1–2049 tokens, at chunk 512 × 8 steps and at chunk 64 × 16 steps: logits and tokens bitwise identical.cargo fmt --check, clippy-D warnings, the crate tests (the cut property test, thepack()property test, the derive staging plan), the schema golden, hawk's restricted-visibility class, andcargo build/cargo test -p kern-serve.Limits (roadmap PP2)
max_tokens1), with no prefix cache.🤖 Generated with Claude Code