Repository navigation
Conversation
…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>
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: e4105690d4
ℹ️ 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".
| pub(crate) fn derive(&mut self, d: Derivation, tensors: &dyn Tensors, stage: &Lanes) -> Result<u64> { | ||
| let vars = Dense::check(&self.manifest, &BTreeMap::new(), &vec![false; self.manifest.vars.len()])?; |
There was a problem hiding this comment.
When a derive call or launch uses a manifest variable, this forces every variable to Var::MIN because Dense::check is given an all-false usage mask and an empty value map. The verifier currently permits variables in derive calls, so a source transformation whose element count, scalar argument, or grid depends on a variable will run only at the minimum shape and can leave the derived carry incomplete or compute the wrong values; either reject variable-dependent derive programs or evaluate them at the appropriate declared bounds.
Useful? React with 👍 / 👎.
| if p.derive && (p.once || p.batch.is_some() || p.graph) { | ||
| errs.push(format!( | ||
| "program `{pname}`: `derive` (run inside weight loading) excludes `once`, `batch` and `graph`" | ||
| )); |
There was a problem hiding this comment.
The new derive execution path runs all derive programs automatically in manifest-name order, but verification still treats every carry as initially written and does not enforce a producer order for derive programs. A manifest can therefore have derive program a read a carry produced only by derive program b (or by a later serving program), pass verification, and execute a against uninitialized memory during load_weights; derive dataflow needs to reject such reads or validate that the producer precedes the consumer.
Useful? React with 👍 / 👎.
| def finish(self, name="load", derive=False): | ||
| """The queued calls as program `name`: once after load, or with | ||
| `derive` inside weight loading (its reads may be `source` buffers).""" |
There was a problem hiding this comment.
Calling Once.finish(..., derive=True) does not make the raw checkpoint buffers created by p1() or cast() into source buffers: Once.derive always creates them as kind: "weight". The runtime consequently allocates and retains every raw tensor as a persistent weight and uploads it through the normal weight path, defeating the staging/reclamation behavior (and potentially causing the memory-saving derive load to OOM); the helper must mark those raw inputs as sources when producing a derive program.
Useful? React with 👍 / 👎.
|
Folded into #44, which now targets master and carries this change. |
1 participant
What
A manifest can declare a checkpoint tensor as a
sourceand the transform that turns it into a weight as aderiveprogram. The runtime runs the derive program call by call insideload_weights. Each source is staged in one shared allocation only between its first and last reader, so peak device memory is the derived weights plus the widest set of live sources, not the checkpoint twice.kern-manifest: asourcebuffer kind and aderive: trueprogram flag. Verification checks that a source binds checkpoint tensors, is read only by derive programs and is never written or exported. It also checks that a derive program touches no state, input, output or peer buffer and runs no collective. The schema v5 golden is regenerated.kern-runtime:derive.rsplans the staging windows as a pure function (first fit in arrival order), andload.rsfills sources lane by lane and runs each derive call once its sources have arrived.tools/once.py/tools/kern_manifest.py: generators can emit a derive program.docs/manifest.md,docs/runtime.md.Why
Full Kimi-K3's mxfp4 experts have to be shuffled into the batched GEMM's layout before any kernel reads them. Before this change the raw tensors loaded as weights and a
onceprogram shuffled them, so both copies were resident. A PP8 stage holds about 188 GB of experts, which does not fit twice in one GB300 (288 GB).Gate
cargo fmt --check,cargo clippy --all-targets -D warnings,cargo test(kern-manifest, kern-pool, kern-runtime, kern-test), the schema golden, the hawk restricted-visibility class, andcargo build -p kern-serveall pass locally.crates/kern-runtime/tests/derive.rs(the staging plan) and verifier cases incrates/kern-manifest/tests/verify.rs.Dependencies
None. This branch is based on
masterand cherry-picks cleanly. The pipeline-parallel K3 work (stacked on #42) builds on it and carries the same commit until this merges.🤖 Generated with Claude Code