Repository navigation
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>
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. |
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>
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 59590c7821
ℹ️ 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".
| let (vars, values) = kern_run::fills(&self.p, std::slice::from_ref(lease), &[pos], ids.len(), ids); | ||
| let mut rows: BTreeMap<String, Vec<i64>> = values.into_iter().map(|(f, v)| (f.name.clone(), v)).collect(); | ||
| for t in &self.p.page_tables { | ||
| let mut v = Vec::with_capacity(t.width); | ||
| lease.extend_row(&t.name, &mut v)?; | ||
| rows.insert(t.name.clone(), v.into_iter().map(i64::from).collect()); | ||
| } |
There was a problem hiding this comment.
Forward inputs required only by downstream stages
When a later stage has a fill or page-table buffer that is not used by stage 0, send_item never includes it in Item.rows because it derives rows only from the head's protocol. The downstream stage still retains and writes the corresponding input before issuing its program, but receives no value, so the runtime uses stale/zero data and the pipeline produces incorrect results. The cut protocol needs to transmit the union of inputs required by all stages (or otherwise derive downstream values independently).
Useful? React with 👍 / 👎.
| if accesses(m, c).iter().any(|(b, _, d)| *d != Dir::In && needed.contains(*b)) { | ||
| needed.extend(reads(m, c).map(str::to_string)); | ||
| kept.push(c.clone()); | ||
| } |
There was a problem hiding this comment.
Preserve once calls that initialize state
The once-program slicer decides whether to retain a call by examining only buffer accesses and the needed buffer set. A valid once call whose output is an inout/out state therefore cannot make the cut, so it is dropped from every stage even when a stage's serving calls later read that state. Such a pipeline silently runs without the required initialization and can produce incorrect results; state dependencies need to participate in the backward slice.
Useful? React with 👍 / 👎.
| let mut modules: BTreeMap<String, Module> = ops | ||
| .values() | ||
| .flat_map(|o| &o.imp.launches) | ||
| .filter_map(Launch::module) | ||
| .filter(|n| *n != MODULE) | ||
| .map(|n| (n.to_string(), m.modules[n].clone())) | ||
| .collect(); | ||
| modules.insert(MODULE.to_string(), pp.clone()); |
There was a problem hiding this comment.
Avoid overwriting an existing
pp module
Cut stages unconditionally replace the manifest's module named pp with the mailbox cubin. If the original manifest already uses a module with that name for a model operation retained in a stage, its operation declaration still references pp but now resolves to the mailbox binary, causing missing-entry/load failures (or the wrong kernel) at runtime. The cut should reject this namespace collision or use a reserved generated module name.
Useful? React with 👍 / 👎.
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
|
Folded into #44, which now targets master and carries these commits. |
…over eight of them
* manifest: a manifest cuts into pipeline stages at call indices
`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>
* kern-serve: a cut manifest serves prefill as one process per stage
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>
* qwen3: the target kv is one paged state per layer
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>
* manifest: a cut stage carries only the states it touches and says how 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>
* qwen3: the example hash baselines follow the per-layer kv
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>
* dsv41: ar_trace skips the clock fill like the other outputs
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
* runtime: a derive program turns source tensors into weights while they 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>
* manifest: a cut slices derive programs per stage and lets the last stage'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>
* runtime: a runtime zeroes a sequence slot another pool handed out
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>
* kern-serve: a pipeline serves per-sequence state and a chunk program 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>
* manifest: a downstream cut stage without a seq_len fill takes its groups 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>
* kern-serve: a pipeline head without a one-row step feeds the whole prompt 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>
* manifest: a prefill call carries several sequences, split by its cu_seqlens 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>
* kern-serve: the pp head packs the oldest prompts' next rows into each 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>
* runtime: a mark lets another thread wait on the stream while the issuer 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>
* kern-serve: a pipeline stage logs each item's issue and done from the 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>
* manifest: a cut adds transport only, the device clock is gone
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>
* manifest: a var says which axis of a call it counts; the roles no longer 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>
* tools: the qwen example baselines follow the axis field
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>
* docs: PP3 is a boundary kind without edges, and comes before PP2
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>
* kern-serve: a pipeline serves paged state only
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>
* pp: a cut stage is a function; the shell moves what crosses
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>
---------
Signed-off-by: JinYan Su <751080330@qq.com>
Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Prefill-only pipeline parallelism. A forward is a pure function, so a cut 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
kern cut: splits a manifest at call indices. Liveness finds the workspaces live across a cut; the boundary becomes asend/recvpair over a mailbox on the downstream GPU (tools/kernels-src/pp_mailbox.cu: wait empty, put, post / wait full, take, free; spin with timeout and error code, never hang). Each stage keeps only the states its calls touch and reports what share of each it uses.--pp-listen/--pp-head. HandshakeHello/Tableover TCP (fabric handles of the mailboxes, min pages across stages), then the same connection carriesItem { vars, rows }down and tokens back (LengthDelimitedCodec+serde_json). Page ids are shared by every stage. Any error exits the process; the head exits on EOF.Contract changes
cut { id, stage, stages }; aclockfill (per-item recv / send wait from%globaltimer, read back by the stage); schema v5 golden regenerated.tokenfill is required only on stage 0.pp.<edge>group joins two different manifests that agree on the peer'sofbuffer. No runtime change (import_peersalready matches by name and size).kv.N, and argmax / embedding / weight_prep now resolve to the registry builds.Gates (2026-10-07, GB300)
kern test qwen3-4b: PASS, bitwise at every span; per-layer vs single-state manifest bitwise;kern runtext identical; dspark 7-row round and--rows 1text identical.Limits (roadmap PP2)
Prefill only (max_tokens 1); fully paged manifests only (per-sequence state refused); a chunk program that emits a token is refused; a stage frozen after the head has drained leaves the head waiting on TCP; the mailbox cubin is not in the registry yet.
🤖 Generated with Claude Code