From e4105690d4770d91626349fd327da0de29d5db1e Mon Sep 17 00:00:00 2001 From: JinYan Su <751080330@qq.com> Date: Wed, 7 Oct 2026 15:55:35 +0000 Subject: [PATCH] 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) --- crates/kern-manifest/src/protocol.rs | 5 + crates/kern-manifest/src/types.rs | 8 +- crates/kern-manifest/src/verify.rs | 60 +++++++-- crates/kern-manifest/tests/verify.rs | 53 ++++++++ crates/kern-run/src/run.rs | 1 + crates/kern-runtime/src/compile.rs | 27 +++- crates/kern-runtime/src/derive.rs | 186 +++++++++++++++++++++++++++ crates/kern-runtime/src/lib.rs | 12 +- crates/kern-runtime/src/load.rs | 110 +++++++++++++--- crates/kern-runtime/tests/derive.rs | 86 +++++++++++++ crates/kern-test/src/diff.rs | 2 +- crates/kern-test/src/harness.rs | 8 +- docs/manifest.md | 25 +++- docs/runtime.md | 3 +- schema/manifest-v5.schema.json | 11 +- tools/kern_manifest.py | 7 +- tools/once.py | 8 +- 17 files changed, 555 insertions(+), 57 deletions(-) create mode 100644 crates/kern-runtime/src/derive.rs create mode 100644 crates/kern-runtime/tests/derive.rs diff --git a/crates/kern-manifest/src/protocol.rs b/crates/kern-manifest/src/protocol.rs index ca792aa7..c675d541 100644 --- a/crates/kern-manifest/src/protocol.rs +++ b/crates/kern-manifest/src/protocol.rs @@ -203,6 +203,9 @@ pub struct Protocol { pub forwards: Vec, /// Programs run once after load, in name order. pub once: Vec, + /// Programs the runtime runs inside weight loading, in name order: no + /// caller runs them. + pub derive: Vec, } impl Protocol { @@ -409,6 +412,7 @@ impl Protocol { // Forwards: the shape, and by dataflow what each hands back. let mut forwards: Vec = Vec::new(); let mut once = Vec::new(); + let derive = m.programs.iter().filter(|(_, p)| p.derive).map(|(n, _)| n.clone()).collect(); let mut span_var: Option = None; let mut context_var: Option = None; for (pname, p) in &m.programs { @@ -599,6 +603,7 @@ impl Protocol { line_tables, forwards, once, + derive, }) } else { Err(ProtocolErrors(errs)) diff --git a/crates/kern-manifest/src/types.rs b/crates/kern-manifest/src/types.rs index 0800af6e..8afd00bf 100644 --- a/crates/kern-manifest/src/types.rs +++ b/crates/kern-manifest/src/types.rs @@ -148,6 +148,9 @@ pub struct Program { /// Run once after load (after every peer is imported), never per step: a tray manifest's collective setup. Takes no per-call input. #[serde(default, skip_serializing_if = "is_false")] pub once: bool, + /// Run by the runtime inside weight loading, call by call, to turn `source` buffers into the `carry` buffers the other programs read (a shuffle into a kernel's layout, a dtype cast). Reads no state, peer, input or output; nothing else may run it. + #[serde(default, skip_serializing_if = "is_false")] + pub derive: bool, /// Driven through a CUDA graph: the runtime captures the call list at each var assignment the program is called with and replays it with one launch. For a program of fixed shape per call whose launches are many and short (a decode step, a speculative round); a program whose rows vary per call (a prefill chunk) runs launch by launch. #[serde(default, skip_serializing_if = "is_false")] pub graph: bool, @@ -353,7 +356,7 @@ pub struct Buffer { /// `peer` buffers only: the topology group the addresses are indexed by, e.g. `"ep"`. #[serde(default, skip_serializing_if = "Option::is_none")] pub group: Option, - /// `weight` buffers only: the checkpoint tensors laid end to end into this buffer, in order, e.g. `[{"tensor": "q_proj.weight"}, {"tensor": "k_proj.weight"}]`. + /// `weight` and `source` buffers only: the checkpoint tensors laid end to end into this buffer, in order, e.g. `[{"tensor": "q_proj.weight"}, {"tensor": "k_proj.weight"}]`. #[serde(default, skip_serializing_if = "Vec::is_empty")] pub bind: Vec, } @@ -579,6 +582,8 @@ pub enum BufferKind { Output, /// Assembled at load time from the checkpoint tensors its `bind` names, e.g. `model.embed_tokens.weight`. Weight, + /// Checkpoint tensors bound like a weight but alive only while a `derive` program reads them during load: device memory holds what is derived from them, never both at once, e.g. an expert stack before its shuffle. + Source, /// Runtime-owned scratch, dead between runs, e.g. `hidden`. Workspace, /// Written by one program and read by another, kept between runs, e.g. the `fc_out` hidden states a draft reads. @@ -593,6 +598,7 @@ impl fmt::Display for BufferKind { BufferKind::Input => "input", BufferKind::Output => "output", BufferKind::Weight => "weight", + BufferKind::Source => "source", BufferKind::Workspace => "workspace", BufferKind::Carry => "carry", BufferKind::Peer => "peer", diff --git a/crates/kern-manifest/src/verify.rs b/crates/kern-manifest/src/verify.rs index f00f0fe3..0a58d296 100644 --- a/crates/kern-manifest/src/verify.rs +++ b/crates/kern-manifest/src/verify.rs @@ -25,9 +25,11 @@ //! 7. calls: op refs resolve, arg/param arity and per-position type match //! against the interface, var ranges fit scalar params //! 8. dataflow per program: no read-before-write, no writes to input, -//! weight or peer buffers, every output / carry buffer written by some -//! program; a program is `once` or has a `batch`, not both; a batch -//! has `groups >= 1` and constant `rows >= 1` or a declared var +//! weight, source or peer buffers, every output / carry buffer written +//! by some program; a program is `once` or has a `batch`, not both; a +//! batch has `groups >= 1` and constant `rows >= 1` or a declared var; +//! a source is read only by `derive` programs, which run inside weight +//! loading and so touch no state, input, output, peer or collective //! 9. no unused declarations; a `fill` sits on an input or output of an //! integer dtype, input roles on inputs and output roles on outputs //! (whether the fills add up to a serving contract is @@ -270,12 +272,15 @@ fn diagnostics(m: &Manifest) -> Vec { errs.push(format!("{ctx}: shared host weights cannot select tensors, rows or columns by rank")); } } - // 9c. bind: a weight is its segments, nothing else has any + // 9c. bind: a weight or a source is its segments, nothing else has any + if b.kind == BufferKind::Source && b.export { + errs.push(format!("{ctx}: a source lives only while weights load; it cannot be exported")); + } match (b.kind, b.bind.is_empty()) { - (BufferKind::Weight, true) => { - errs.push(format!("{ctx}: a weight buffer binds at least one checkpoint tensor (`bind`)")) + (BufferKind::Weight | BufferKind::Source, true) => { + errs.push(format!("{ctx}: a {} buffer binds at least one checkpoint tensor (`bind`)", b.kind)) } - (BufferKind::Weight, false) => { + (BufferKind::Weight | BufferKind::Source, false) => { for (i, s) in b.bind.iter().enumerate() { let sctx = format!("{ctx}: bind[{i}]"); let source = |t: &TensorSource, errs: &mut Vec, used_groups: &mut BTreeSet| match t @@ -340,9 +345,8 @@ fn diagnostics(m: &Manifest) -> Vec { } } } - (kind, false) => { - errs.push(format!("{ctx}: `bind` names checkpoint tensors for a weight, not a {kind} buffer")) - } + (kind, false) => errs + .push(format!("{ctx}: `bind` names checkpoint tensors for a weight or a source, not a {kind} buffer")), (_, true) => {} } // 10b. export / peer @@ -763,7 +767,10 @@ fn diagnostics(m: &Manifest) -> Vec { // per-program dataflow treats them as initially written. // Peer buffers are filled by the runtime when the group's // handles are imported, before any program runs. - matches!(b.kind, BufferKind::Input | BufferKind::Weight | BufferKind::Carry | BufferKind::Peer) + matches!( + b.kind, + BufferKind::Input | BufferKind::Weight | BufferKind::Source | BufferKind::Carry | BufferKind::Peer + ) }) .map(|(n, _)| n.as_str()) .collect(); @@ -775,6 +782,11 @@ fn diagnostics(m: &Manifest) -> Vec { "program `{pname}`: `once` (run after load) and `batch` (driven per step) are exclusive" )); } + 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`" + )); + } if p.graph && p.batch.is_none() { errs.push(format!("program `{pname}`: `graph` (captured per call shape) needs a `batch` (a program a serving loop drives)")); } @@ -820,6 +832,12 @@ fn diagnostics(m: &Manifest) -> Vec { }; used_ops.insert(c.op.clone()); let has_extern = op.imp.launches.iter().any(|l| l.is_extern()); + if p.derive && op.imp.launches.iter().any(|l| l.entry().starts_with("extern:nccl")) { + errs.push(format!( + "{ctx}: a `derive` program runs before any group is joined; op `{}` is a collective", + c.op + )); + } if c.args.len() != op.params.len() { errs.push(format!("{ctx}: op `{}` takes {} params, got {} args", c.op, op.params.len(), c.args.len())); @@ -870,6 +888,16 @@ fn diagnostics(m: &Manifest) -> Vec { } } } + match (p.derive, b.kind) { + (false, BufferKind::Source) => errs.push(format!( + "{actx}: source buffer `{buf}` is read only by a `derive` program; it is gone after load" + )), + (true, BufferKind::Input | BufferKind::Output | BufferKind::Peer) => errs.push(format!( + "{actx}: a `derive` program runs inside weight loading, before any {} buffer such as `{buf}` means anything", + b.kind + )), + _ => {} + } if b.kind == BufferKind::Peer && has_extern { errs.push(format!( "{actx}: peer buffer `{buf}` passed to op `{}`, which has an extern launch; runtime built-ins never receive peer memory", @@ -877,7 +905,10 @@ fn diagnostics(m: &Manifest) -> Vec { )); } if matches!(dir, Dir::Out | Dir::InOut) { - if matches!(b.kind, BufferKind::Input | BufferKind::Weight | BufferKind::Peer) { + if matches!( + b.kind, + BufferKind::Input | BufferKind::Weight | BufferKind::Source | BufferKind::Peer + ) { errs.push(format!("{actx}: op writes to read-only {} buffer `{buf}`", b.kind)); } written.insert(buf.as_str()); @@ -889,6 +920,11 @@ fn diagnostics(m: &Manifest) -> Vec { // runtime-scaled pool; there is no static bound to // check them against. used_states.insert(state.clone()); + if p.derive { + errs.push(format!( + "{actx}: a `derive` program runs inside weight loading, before state `{state}` holds anything" + )); + } if !m.states.contains_key(state) { errs.push(format!("{actx}: unknown state `{state}`")); } diff --git a/crates/kern-manifest/tests/verify.rs b/crates/kern-manifest/tests/verify.rs index 899915a7..0b093cee 100644 --- a/crates/kern-manifest/tests/verify.rs +++ b/crates/kern-manifest/tests/verify.rs @@ -1091,3 +1091,56 @@ fn host_span_ends_at_the_last_element() { let h = HostTensor { dtype: DType::U8, shape: vec![0, 10], strides: vec![16, 1] }; assert_eq!((h.span(1), h.span(3)), (Some(10), Some(42))); } + +/// The base with its weight turned into a source that a derive program +/// shuffles into the carry `decode` embeds from. +fn derived() -> serde_json::Value { + let mut v = base(); + v["buffers"].as_object_mut().unwrap().remove("w"); + v["buffers"]["raw"] = + serde_json::json!({ "dtype": "bf16", "shape": [64, 64], "kind": "source", "bind": [{ "tensor": "raw" }] }); + v["buffers"]["ws"] = serde_json::json!({ "dtype": "bf16", "shape": [64, 64], "kind": "carry" }); + v["ops"]["shuffle"] = serde_json::json!({ + "params": ["in buffer", "out buffer", "i32"], + "impl": { "launches": [{ "module": "toy", "entry": "shuffle_k", "block": [128, 1, 1], "grid": [32, 1, 1] }] } + }); + v["programs"]["load"] = serde_json::json!({ "derive": true, "calls": [ + { "op": "shuffle", "args": [{ "buf": "raw" }, { "buf": "ws" }, { "i32": 4096 }] } + ] }); + v["programs"]["decode"]["calls"][0]["args"][1] = serde_json::json!({ "buf": "ws" }); + v +} + +#[test] +fn a_source_is_read_by_derive_programs_only() { + check(derived()).unwrap(); + let mut v = derived(); + v["programs"]["decode"]["calls"][0]["args"][1] = serde_json::json!({ "buf": "raw" }); + assert_err(v, "source buffer `raw` is read only by a `derive` program"); + let mut v = derived(); + v["programs"]["load"]["calls"][0]["args"] = serde_json::json!([{ "buf": "ws" }, { "buf": "raw" }, { "i32": 1 }]); + assert_err(v, "op writes to read-only source buffer `raw`"); + let mut v = derived(); + v["buffers"]["raw"]["bind"] = serde_json::json!([]); + assert_err(v, "a source buffer binds at least one checkpoint tensor"); + let mut v = derived(); + v["buffers"]["raw"]["export"] = serde_json::json!(true); + assert_err(v, "a source lives only while weights load"); +} + +#[test] +fn a_derive_program_touches_nothing_serving_owns() { + let mut v = derived(); + v["programs"]["load"]["once"] = serde_json::json!(true); + assert_err(v, "`derive` (run inside weight loading) excludes `once`, `batch` and `graph`"); + let mut v = derived(); + v["programs"]["load"]["calls"].as_array_mut().unwrap().push( + serde_json::json!({ "op": "embed", "args": [{ "buf": "x" }, { "buf": "ws" }, { "buf": "h" }, { "i32": 1 }] }), + ); + assert_err(v, "before any input buffer such as `x` means anything"); + let mut v = derived(); + v["programs"]["load"]["calls"].as_array_mut().unwrap().push(serde_json::json!( + { "op": "attn", "args": [{ "buf": "ws" }, { "state": "kv" }, { "buf": "h" }, { "i32": 1 }, { "i64": 0 }] } + )); + assert_err(v, "before state `kv` holds anything"); +} diff --git a/crates/kern-run/src/run.rs b/crates/kern-run/src/run.rs index ac3fa398..b011b7e1 100644 --- a/crates/kern-run/src/run.rs +++ b/crates/kern-run/src/run.rs @@ -226,6 +226,7 @@ fn banner(rt: &Runtime, manifest: &Path, kernels: Option<&Path>, load: Duration) let shape = match &p.batch { Some(b) => format!(", {} × {:?} per call", b.groups, b.rows), None if p.once => ", once after load".into(), + None if p.derive => ", inside weight loading".into(), None => String::new(), }; info!(" program `{name}`: {} calls{shape}", p.calls.len()); diff --git a/crates/kern-runtime/src/compile.rs b/crates/kern-runtime/src/compile.rs index d2d2cf3f..a745811a 100644 --- a/crates/kern-runtime/src/compile.rs +++ b/crates/kern-runtime/src/compile.rs @@ -460,8 +460,9 @@ impl PinnedCheck<'_> { } } -/// Lower every program's call list into a flat launch list. Every launch -/// that receives a peer buffer is SASS-scanned for multicast TMA first. +/// Lower the call list of every program whose `derive` is `derive` into +/// a flat launch list. Every launch that receives a peer buffer is +/// SASS-scanned for multicast TMA first. fn compile_programs( manifest: &Manifest, ops: &BTreeMap, @@ -469,13 +470,14 @@ fn compile_programs( states: &BTreeMap, ranks: &Ranks, blt: &CudaBlasLT, + derive: bool, ) -> Result> { let vars: BTreeMap<&str, usize> = manifest.vars.keys().enumerate().map(|(i, s)| (s.as_str(), i)).collect(); let maxima: Vec = manifest.vars.values().map(|v| v.max).collect(); let pinned = PinnedCheck { blt, maxima: &maxima }; let mut scan = MulticastScan::new(); let mut programs = BTreeMap::new(); - for (pname, p) in &manifest.programs { + for (pname, p) in manifest.programs.iter().filter(|(_, p)| p.derive == derive) { let calls = &p.calls; let mut launches = Vec::new(); let mut call_ranges = Vec::with_capacity(calls.len()); @@ -496,8 +498,8 @@ fn compile_programs( Ok(programs) } -/// [`compile_programs`] with the placement read off the runtime's ranks and -/// its peer buffers. +/// [`compile_programs`] over every program but the derive ones, with the +/// placement read off the runtime's ranks and its peer buffers. pub(crate) fn compile( manifest: &Manifest, ops: &BTreeMap, @@ -508,7 +510,20 @@ pub(crate) fn compile( blt: &CudaBlasLT, ) -> Result> { let peer_buffers: BTreeSet = peers.keys().cloned().collect(); - compile_programs(manifest, ops, buffers, states, &Ranks { ranks, peer_buffers: &peer_buffers }, blt) + compile_programs(manifest, ops, buffers, states, &Ranks { ranks, peer_buffers: &peer_buffers }, blt, false) +} + +/// The derive programs, lowered: they read no state and no peer, so they +/// compile at load whatever the host binds later. +pub(crate) fn compile_derive( + manifest: &Manifest, + ops: &BTreeMap, + buffers: &BTreeMap, + ranks: &BTreeMap, + blt: &CudaBlasLT, +) -> Result> { + let none = BTreeSet::new(); + compile_programs(manifest, ops, buffers, &BTreeMap::new(), &Ranks { ranks, peer_buffers: &none }, blt, true) } /// The impl-private scratch of every resolved op, whose addresses the diff --git a/crates/kern-runtime/src/derive.rs b/crates/kern-runtime/src/derive.rs new file mode 100644 index 00000000..1833e841 --- /dev/null +++ b/crates/kern-runtime/src/derive.rs @@ -0,0 +1,186 @@ +//! Derivation: weights that are a function of checkpoint tensors rather +//! than the tensors themselves (an expert stack shuffled into a GEMM's +//! layout), computed while the weights load. +//! +//! A `source` buffer is checkpoint bytes no serving program reads; a +//! `derive` program turns sources into the carries the other programs +//! read. Device memory holds a source only while derivation needs it: the +//! derive programs run call by call, in program-name order, and each +//! source occupies a window of one shared staging allocation from the +//! call that first reads it to the call that last does. Two sources +//! alive at the same call never share bytes; a source whose last reader +//! has run gives its bytes to the next one. The peak is the derived +//! weights plus the widest set of sources one call needs at once, never +//! the whole checkpoint twice. +//! +//! [`layout`] decides the windows and is pure; the runtime allocates the +//! staging, compiles the derive programs against it, and in +//! [`Runtime::load_weights`] uploads each source just before its first +//! reader. After the last call the staging is freed and the derive +//! programs are dropped with it: nothing can run them again. + +use std::collections::BTreeMap; + +use kern_manifest::types::{Arg, BufferKind, Manifest}; + +use crate::compile::{CompiledProgram, Dense}; +use crate::device::DeviceBuf; +use crate::error::Error; +use crate::load::Lanes; +use crate::weights::Tensors; +use crate::{Result, Runtime}; + +/// Alignment of every source in the staging allocation. +const ALIGN: u64 = 256; + +/// Where each source sits in the staging allocation and when it arrives. +pub(crate) struct Plan { + /// Byte offset of every source. + pub(crate) at: BTreeMap, + /// Bytes of the staging allocation. + pub(crate) bytes: u64, + /// Per derive call, in run order: the sources it reads first. + pub(crate) arrivals: Vec>, +} + +/// What is left to derive once the weights are in: the derive programs +/// lowered against the staging, and the staging itself. +pub(crate) struct Derivation { + pub(crate) programs: Vec, + pub(crate) arrivals: Vec>, + pub(crate) stage: DeviceBuf, +} + +impl Runtime { + /// Run the derive programs call by call, each source uploaded just + /// before its first reader, then free the staging and forget the + /// sources. The stream is drained before a source lands, so no call + /// still reads the bytes it takes over. + pub(crate) fn derive(&mut self, d: Derivation, tensors: &dyn Tensors, stage: &Lanes) -> Result { + let vars = Dense::check(&self.manifest, &BTreeMap::new(), &vec![false; self.manifest.vars.len()])?; + let calls = d.programs.iter().flat_map(|p| p.call_ranges.iter().map(move |&(lo, hi)| &p.launches[lo..hi])); + let mut uploaded = 0; + for (launches, arriving) in calls.zip(&d.arrivals) { + if !arriving.is_empty() { + self.stream.synchronize()?; + uploaded += self.fill_sources(arriving, tensors, stage)?; + } + for l in launches { + self.launch(l, &vars).map_err(|e| Error::Call { context: l.ctx.clone(), source: Box::new(e) })?; + } + } + self.stream.synchronize()?; + for s in d.arrivals.iter().flatten() { + self.buffers.remove(s); + } + drop(d.stage); + Ok(uploaded) + } +} + +/// The derive programs' names, in the order they run. +pub(crate) fn programs(m: &Manifest) -> Vec<&str> { + m.programs.iter().filter(|(_, p)| p.derive).map(|(n, _)| n.as_str()).collect() +} + +/// The staging plan for a manifest whose sources are `bytes` long. +pub(crate) fn plan(m: &Manifest, bytes: &BTreeMap) -> Plan { + let calls: Vec<_> = programs(m).into_iter().flat_map(|p| &m.programs[p].calls).collect(); + let mut live: BTreeMap<&str, (usize, usize)> = BTreeMap::new(); + for (k, c) in calls.iter().enumerate() { + for a in &c.args { + let Arg::Buf { buf, .. } = a else { continue }; + if m.buffers.get(buf).is_some_and(|b| b.kind == BufferKind::Source) { + live.entry(buf.as_str()).and_modify(|(_, last)| *last = k).or_insert((k, k)); + } + } + } + let windows: Vec = + live.iter().map(|(n, &(first, last))| Window { first, last, bytes: bytes[*n] }).collect(); + let (at, total) = layout(&windows); + let mut arrivals = vec![Vec::new(); calls.len()]; + live.iter().for_each(|(n, (first, _))| arrivals[*first].push(n.to_string())); + Plan { at: live.keys().map(|n| n.to_string()).zip(at).collect(), bytes: total, arrivals } +} + +/// A source's lifetime in derive calls, `first..=last`, and its size. +#[derive(Clone, Copy, Debug)] +pub(crate) struct Window { + first: usize, + last: usize, + bytes: u64, +} + +/// First fit in arrival order: each window takes the lowest aligned gap +/// among the windows still alive when it arrives. Returns every window's +/// offset, in input order, and the bytes the layout spans. +pub(crate) fn layout(windows: &[Window]) -> (Vec, u64) { + let mut order: Vec = (0..windows.len()).collect(); + order.sort_by_key(|&i| (windows[i].first, i)); + let mut at = vec![0u64; windows.len()]; + let mut placed: Vec = Vec::new(); + for i in order { + let w = windows[i]; + let mut taken: Vec<(u64, u64)> = placed + .iter() + .filter(|&&j| windows[j].last >= w.first) + .map(|&j| (at[j], at[j] + windows[j].bytes)) + .collect(); + taken.sort(); + let mut lo = 0; + for (a, b) in taken { + if lo + w.bytes <= a { + break; + } + lo = lo.max(b.next_multiple_of(ALIGN)); + } + at[i] = lo; + placed.push(i); + } + let total = windows.iter().zip(&at).map(|(w, a)| a + w.bytes).max().unwrap_or(0); + (at, total) +} + +#[cfg(test)] +mod tests { + use super::{layout, Window, ALIGN}; + + /// Windows alive at one call never overlap, every offset is aligned, + /// and the span is the sum of sizes at most, over random lifetimes. + #[test] + fn live_windows_are_disjoint() { + let mut s = 0x2545_f491_4f6c_dd1du64; + let mut next = |n: u64| { + s ^= s << 13; + s ^= s >> 7; + s ^= s << 17; + s % n + }; + for _ in 0..2000 { + let windows: Vec = (0..next(12)) + .map(|_| { + let first = next(10) as usize; + Window { first, last: first + next(4) as usize, bytes: 1 + next(5000) } + }) + .collect(); + let (at, total) = layout(&windows); + for (i, (a, w)) in at.iter().zip(&windows).enumerate() { + assert_eq!((a % ALIGN, a + w.bytes <= total), (0, true)); + for (b, v) in at.iter().zip(&windows).skip(i + 1) { + let alive_together = w.first <= v.last && v.first <= w.last; + let apart = a + w.bytes <= *b || b + v.bytes <= *a; + assert!(!alive_together || apart, "{w:?} at {a} and {v:?} at {b}"); + } + } + let sum: u64 = windows.iter().map(|w| w.bytes.next_multiple_of(ALIGN)).sum(); + assert!(total <= sum); + } + } + + /// A chain of sources each read by one call reuses one window. + #[test] + fn a_chain_reuses_one_window() { + let windows: Vec = (0..5).map(|k| Window { first: k, last: k, bytes: 1000 + k as u64 }).collect(); + assert_eq!(layout(&windows), (vec![0; 5], 1004)); + } +} diff --git a/crates/kern-runtime/src/lib.rs b/crates/kern-runtime/src/lib.rs index 0d79822a..a7ad4bdb 100644 --- a/crates/kern-runtime/src/lib.rs +++ b/crates/kern-runtime/src/lib.rs @@ -28,6 +28,7 @@ mod compare; mod compile; mod cubin; mod cublas; +mod derive; mod device; mod error; mod exec; @@ -179,6 +180,8 @@ pub struct Runtime { /// ([`Runtime::enqueue_after`]): host work before the program, the /// program before host work after it. joins: [sys::CUevent; 2], + /// The derive programs and their staging, until `load_weights` runs them. + derivation: Option, } impl Drop for Runtime { @@ -208,9 +211,14 @@ impl Runtime { self.n_modules } - /// (name, class, allocated bytes) for every buffer. + /// (name, class, allocated bytes) for every buffer device memory + /// holds: the sources only until `load_weights` has derived from them. pub fn buffer_sizes(&self) -> Vec<(&str, BufferKind, u64)> { - self.manifest.buffers.iter().map(|(n, b)| (n.as_str(), b.kind, self.buffers[n].bytes)).collect() + self.manifest + .buffers + .iter() + .filter_map(|(n, b)| Some((n.as_str(), b.kind, self.buffers.get(n)?.bytes))) + .collect() } /// (name, declaration, allocated bytes) for every state. diff --git a/crates/kern-runtime/src/load.rs b/crates/kern-runtime/src/load.rs index 899f272e..5bc84caf 100644 --- a/crates/kern-runtime/src/load.rs +++ b/crates/kern-runtime/src/load.rs @@ -21,15 +21,16 @@ use kern_manifest::types::{Buffer, BufferKind, DType, Dim, Manifest, Placement, use kern_manifest::Verified; use crate::cublas::Blas; +use crate::derive::Derivation; use crate::device::{ - alloc, alloc_host, alloc_vmm, chunk_granularity, copy_1d, copy_2d, record, Arena, DeviceBuf, Mapper, Physical, - Pinned, Space, + alloc, alloc_host, alloc_uninit, alloc_vmm, chunk_granularity, copy_1d, copy_2d, record, Arena, DeviceBuf, Mapper, + Physical, Pinned, Space, }; use crate::error::{bail, cuda_check}; use crate::lease::Remaps; use crate::peers::PeerSlot; use crate::weights::{Blob, Tensors}; -use crate::{compile, cubin, weights, Capacity, Error, Result, Runtime, Topology, HEADROOM}; +use crate::{compile, cubin, derive, weights, Capacity, Error, Result, Runtime, Topology, HEADROOM}; use kern_pool::{chunks_for, page_unit, Kind, Pool}; impl Runtime { @@ -153,8 +154,13 @@ impl Runtime { let mut buffers = BTreeMap::new(); let mut peers = BTreeMap::new(); let (mut filled, mut kept) = (BTreeSet::new(), 0u64); + let mut sources = BTreeMap::new(); for (name, b) in &manifest.buffers { let bytes = compile::shaped_bytes(&format!("buffer `{name}`"), &b.shape, b.dtype.bytes(), &vars_max)?; + if b.kind == BufferKind::Source { + sources.insert(name.clone(), bytes); + continue; + } let buf = if let Some(r) = resident.as_mut().and_then(|r| r.take(name, b, bytes)) { filled.insert(name.clone()); kept += bytes; @@ -179,6 +185,13 @@ impl Runtime { } buffers.insert(name.clone(), buf); } + // Sources are windows of one staging allocation (`derive`); their + // views are names for the compiler, the staging owns the bytes. + let plan = derive::plan(&manifest, &sources); + let stage = alloc_uninit(&stream, plan.bytes)?; + for (name, at) in &plan.at { + buffers.insert(name.clone(), DeviceBuf::borrowed(&stream, stage.ptr + at, sources[name])); + } let mut staging = BTreeMap::new(); for (name, b) in &manifest.buffers { if b.kind == BufferKind::Input { @@ -276,6 +289,12 @@ impl Runtime { } else { compile::compile(&manifest, &resolved, &buffers, &states, &ranks, &peers, &blt)? }; + let mut derived = compile::compile_derive(&manifest, &resolved, &buffers, &ranks, &blt)?; + let derivation = Derivation { + programs: derive::programs(&manifest).iter().filter_map(|p| derived.remove(*p)).collect(), + arrivals: plan.arrivals, + stage, + }; let provision = Provision { tokens: pool.pages_max() as u64 * page, seq_slots: pool.slots_max() as u64 }; let remaps = Remaps::spawn(Arc::clone(&ctx), mapper)?; @@ -311,6 +330,7 @@ impl Runtime { imports: Vec::new(), compare: Default::default(), joins: [join_event()?, join_event()?], + derivation: Some(derivation), }; rt.zero_fresh(&initial)?; rt.stream.synchronize()?; @@ -322,14 +342,23 @@ impl Runtime { /// bytes, wherever [`Tensors`] says they are: this process's memory /// (safetensors blobs) or memory this device reads from another /// process (a weight cache's buckets, mapped with [`Runtime::map`]). + /// + /// Then the derive programs run, call by call, over the `source` + /// tensors (see `derive`): a runtime with derive programs loads its + /// weights once, since the sources are gone after. pub fn load_weights(&mut self, tensors: &dyn Tensors) -> Result<()> { if self.host_weights_ready && self.buffers.values().any(|b| b.host_weight().is_some()) { bail!(Api, "host weights are an immutable snapshot; use a new runtime and scope to reload"); } + let derivation = self.derivation.take(); + if derivation.is_none() && !derive::programs(&self.manifest).is_empty() { + bail!(Api, "weights are derived once per runtime; load a new runtime to reload"); + } self.ctx.bind_to_thread()?; let gib = |b: u64| b as f64 / (1u64 << 30) as f64; // Copies of a few MiB each are launch-latency bound on one stream. let lanes = (0..LOAD_LANES).map(|_| self.ctx.new_stream()).collect::, _>>()?; + let stage = Lanes::new(&self.ctx, self.gpu as i32)?; let mut planned = Vec::new(); for (name, b) in &self.manifest.buffers { if b.kind != BufferKind::Weight || self.filled.contains(name) { @@ -356,7 +385,7 @@ impl Runtime { device.0 += dst.bytes; } let t = std::time::Instant::now(); - let up = (upload(&self.ctx, self.gpu as i32, &staged)?, t.elapsed().as_secs_f64()); + let up = (upload(&self.ctx, &stage, &staged)?, t.elapsed().as_secs_f64()); for (dst, copies) in planned.iter() { let Some(h) = dst.host_weight() else { continue }; let t = std::time::Instant::now(); @@ -369,8 +398,11 @@ impl Runtime { } lanes.iter().try_for_each(|s| s.synchronize())?; self.stream.synchronize()?; + let t = std::time::Instant::now(); + let derived = derivation.map(|d| self.derive(d, tensors, &stage)).transpose()?; + let sources = (derived.unwrap_or(0), t.elapsed().as_secs_f64()); tracing::info!( - "gpu {} weights: device {:.1} GiB in {} copies ({:.1} GiB resident already; {:.1} GiB staged from host memory in {:.1}s, {:.0} GiB/s), host {:.1} GiB in {} copies of which this rank filled {:.1} GiB in {:.1}s ({:.0} GiB/s); {:.1}s in all", + "gpu {} weights: device {:.1} GiB in {} copies ({:.1} GiB resident already; {:.1} GiB staged from host memory in {:.1}s, {:.0} GiB/s), host {:.1} GiB in {} copies of which this rank filled {:.1} GiB in {:.1}s ({:.0} GiB/s), {:.1} GiB of sources derived from in {:.1}s; {:.1}s in all", self.gpu, gib(device.0), device.1, @@ -383,6 +415,8 @@ impl Runtime { gib(filled.0), filled.1, gib(filled.0) / filled.1.max(1e-9), + gib(sources.0), + sources.1, t0.elapsed().as_secs_f64(), ); self.host_weights_ready = true; @@ -468,13 +502,26 @@ const STAGE: u64 = 16 << 20; /// Threads staging host and file bytes, each on its own stream. const STAGE_LANES: usize = 16; -/// Host and file bytes go up through page-locked staging: [`STAGE_LANES`] -/// threads, each with two blocks and a stream, gather a piece of a copy -/// into one block while the DMA out of the other is in flight, and share -/// the pieces through a counter. A pageable `cuMemcpyHtoD` stages the -/// same way inside the driver, on the calling thread alone. Returns the -/// bytes uploaded. -fn upload(ctx: &Arc, dev: i32, copies: &[(u64, &weights::Copy)]) -> Result { +/// The page-locked staging host and file bytes go up through: a stream +/// and two blocks per lane, made once per `load_weights` and reused by +/// every upload in it (the derive programs upload once per call). +pub(crate) struct Lanes(Vec<(Arc, [Pinned; 2])>); + +impl Lanes { + fn new(ctx: &Arc, dev: i32) -> Result { + ctx.bind_to_thread()?; + let lane = |_| Ok((ctx.new_stream()?, [Pinned::alloc(STAGE, dev)?, Pinned::alloc(STAGE, dev)?])); + (0..STAGE_LANES).map(lane).collect::>().map(Lanes) + } +} + +/// Host and file bytes go up through the [`Lanes`]: one thread per lane +/// gathers a piece of a copy into one block while the DMA out of the +/// other is in flight, and the threads share the pieces through a +/// counter. A pageable `cuMemcpyHtoD` stages the same way inside the +/// driver, on the calling thread alone. Returns the bytes uploaded, all +/// landed. +fn upload(ctx: &Arc, lanes: &Lanes, copies: &[(u64, &weights::Copy)]) -> Result { // A copy landing with a stride goes up row by row, so its pieces end // on row boundaries. let pieces: Vec<(usize, u64, u64)> = copies @@ -487,10 +534,8 @@ fn upload(ctx: &Arc, dev: i32, copies: &[(u64, &weights::Copy)]) -> }) .collect(); let next = AtomicUsize::new(0); - let lane = || -> Result<()> { + let lane = |(stream, blocks): &(Arc, [Pinned; 2])| -> Result<()> { ctx.bind_to_thread()?; - let stream = ctx.new_stream()?; - let blocks = [Pinned::alloc(STAGE, dev)?, Pinned::alloc(STAGE, dev)?]; let mut landed: [Option; 2] = [None, None]; let mut i = 0; let mut run = || -> Result<()> { @@ -518,12 +563,13 @@ fn upload(ctx: &Arc, dev: i32, copies: &[(u64, &weights::Copy)]) -> (hi - lo) / copy.width, )?; } - landed[i] = Some(record(&stream)?); + landed[i] = Some(record(stream)?); i ^= 1; } Ok(()) }; - // Nothing may be in flight out of a block when it is freed. + // Nothing may be in flight out of a block when the next upload + // reuses it, and the bytes have landed when this returns. let r = run(); stream.synchronize()?; for ev in landed.into_iter().flatten() { @@ -532,8 +578,8 @@ fn upload(ctx: &Arc, dev: i32, copies: &[(u64, &weights::Copy)]) -> r }; std::thread::scope(|s| { - let lanes: Vec<_> = (0..STAGE_LANES).map(|_| s.spawn(lane)).collect(); - lanes.into_iter().try_for_each(|h| h.join().expect("an upload lane panicked")) + let running: Vec<_> = lanes.0.iter().map(|l| s.spawn(move || lane(l))).collect(); + running.into_iter().try_for_each(|h| h.join().expect("an upload lane panicked")) })?; Ok(copies.iter().map(|(_, c)| c.width * c.rows).sum()) } @@ -555,6 +601,32 @@ fn gather(c: &weights::Copy, lo: u64, hi: u64, out: &mut [u8]) -> Result<()> { Ok(()) } +impl Runtime { + /// Sources' tensors into their staging windows, landed when this + /// returns. Returns their bytes. + pub(crate) fn fill_sources(&self, names: &[String], tensors: &dyn Tensors, stage: &Lanes) -> Result { + let mut planned = Vec::new(); + for n in names { + let dst = &self.buffers[n]; + let plan = |t: &str| tensors.find(t); + let copies = weights::plan(n, &self.manifest.buffers[n], dst.bytes, plan, |g| self.ranks.get(g).copied())?; + planned.push((dst, copies)); + } + let mut staged = Vec::new(); + for (dst, copies) in &planned { + for c in copies { + match c.src { + Blob::Host(_) | Blob::File { .. } => staged.push((dst.ptr, c)), + Blob::Device { ptr, .. } => copy_device(&self.stream, dst, c, ptr)?, + } + } + } + upload(&self.ctx, stage, &staged)?; + self.stream.synchronize()?; + Ok(planned.iter().map(|(d, _)| d.bytes).sum()) + } +} + /// One planned copy of device bytes at `ptr` (another process's /// allocation mapped here): one device-to-device copy, 2D when the rows /// are strided. diff --git a/crates/kern-runtime/tests/derive.rs b/crates/kern-runtime/tests/derive.rs new file mode 100644 index 00000000..3265ba34 --- /dev/null +++ b/crates/kern-runtime/tests/derive.rs @@ -0,0 +1,86 @@ +//! Derived weights end to end: two sources share one staging window, each +//! uploaded only after the call before it has read its predecessor, and +//! the carries they leave behind are what the serving program reads. + +use std::collections::BTreeMap; + +use half::bf16; +use kern_manifest::types::BufferKind; +use kern_manifest::Verified; +use kern_runtime::{Capacity, Runtime, Safetensors}; + +const N: usize = 16; + +fn tensor(f: impl Fn(usize, usize) -> f32) -> Vec { + (0..N * N).flat_map(|i| bf16::from_f32(f(i / N, i % N)).to_le_bytes()).collect() +} + +fn checkpoint(tensors: &[(&str, &[u8])]) -> Vec { + let (mut meta, mut at) = (serde_json::Map::new(), 0); + for (name, bytes) in tensors { + meta.insert( + name.to_string(), + serde_json::json!({"dtype": "BF16", "shape": [N, N], "data_offsets": [at, at + bytes.len()]}), + ); + at += bytes.len(); + } + let mut header = serde_json::Value::Object(meta).to_string().into_bytes(); + header.resize(header.len().div_ceil(8) * 8, b' '); + let mut data = (header.len() as u64).to_le_bytes().to_vec(); + data.extend(header); + tensors.iter().for_each(|(_, b)| data.extend(*b)); + data +} + +#[test] +#[ignore = "requires a CUDA GPU"] +fn sources_derive_through_one_window_and_are_gone_after() { + let gemm = |a: &str, b: &str, c: &str| serde_json::json!({"op": "gemm", "args": [{"buf": a}, {"buf": b}, {"buf": c}, {"i32": N}, {"i32": N}, {"i32": N}]}); + let square = |kind: &str| serde_json::json!({"kind": kind, "dtype": "bf16", "shape": [N, N]}); + let bound = |kind: &str, t: &str| { + let mut b = square(kind); + b["bind"] = serde_json::json!([{"tensor": t}]); + b + }; + let manifest = serde_json::json!({ + "schema_version": 5, "model": "derive-test", "vars": {}, "states": {}, + "buffers": { + "eye": bound("weight", "eye"), "s1": bound("source", "a"), "s2": bound("source", "b"), + "c1": square("carry"), "c2": square("carry"), "o1": square("output"), "o2": square("output") + }, + "modules": {}, + "ops": {"gemm": { + "params": ["in buffer", "in buffer", "out buffer", "i32", "i32", "i32"], + "impl": {"launches": [{"entry": "extern:cublaslt_bf16_tn"}]} + }}, + "programs": { + "load": {"derive": true, "calls": [gemm("eye", "s1", "c1"), gemm("eye", "s2", "c2")]}, + "probe": {"calls": [gemm("eye", "c1", "o1"), gemm("eye", "c2", "o2")]} + } + }); + let verified = Verified::from_json(&manifest.to_string()).unwrap(); + let dir = std::env::temp_dir().join(format!("kern-derive-test-{}", std::process::id())); + std::fs::create_dir_all(&dir).unwrap(); + let mut rt = Runtime::load(&verified, Some(&dir), 0, Some(Capacity { tokens: Some(1), seqs: 1 }), None).unwrap(); + let sources = |rt: &Runtime| -> Vec { + rt.buffer_sizes().into_iter().filter(|(_, k, _)| *k == BufferKind::Source).map(|(n, ..)| n.into()).collect() + }; + assert_eq!(sources(&rt), ["s1", "s2"]); + + let (eye, a, b) = ( + tensor(|r, c| if r == c { 1.0 } else { 0.0 }), + tensor(|r, c| (r * 3 + c) as f32 % 13.0 - 6.0), + tensor(|r, c| (r + c * 5) as f32 % 11.0 - 5.0), + ); + let data = checkpoint(&[("eye", &eye), ("a", &a), ("b", &b)]); + let tensors = Safetensors::parse(&[&data]).unwrap(); + rt.load_weights(&tensors).unwrap(); + assert_eq!(sources(&rt), Vec::::new()); + assert!(rt.load_weights(&tensors).is_err()); + + rt.run("probe", &BTreeMap::new()).unwrap(); + assert_eq!((rt.read_output("o1").unwrap(), rt.read_output("o2").unwrap()), (a, b)); + assert!(rt.run("load", &BTreeMap::new()).is_err()); + drop(rt); + std::fs::remove_dir_all(dir).unwrap(); +} diff --git a/crates/kern-test/src/diff.rs b/crates/kern-test/src/diff.rs index 46a4cc60..2cac7c4b 100644 --- a/crates/kern-test/src/diff.rs +++ b/crates/kern-test/src/diff.rs @@ -175,7 +175,7 @@ pub fn frontier_inputs(m: &Manifest, prog: &str, calls: Range) -> BTreeSe inputs } -/// Buffers the `once` programs write: load-time constants (a packed +/// Buffers the `once` and `derive` programs write: load-time constants (a packed /// weight, a rope table), each side's own like its weights, never handed /// from A to B. pub fn constants(m: &Manifest, once: &[String]) -> BTreeSet { diff --git a/crates/kern-test/src/harness.rs b/crates/kern-test/src/harness.rs index fd882a95..9be32f7e 100644 --- a/crates/kern-test/src/harness.rs +++ b/crates/kern-test/src/harness.rs @@ -233,7 +233,7 @@ pub struct Recording { pub(crate) ranks: usize, pub(crate) chunk: Option, steps: Vec, - /// Programs either side runs once after load: their spans are each + /// Programs either side runs at load (`once`, `derive`): their spans are each /// side's own setup, not something the workload drives. once: Vec, pub(crate) runs: Vec>, @@ -582,9 +582,9 @@ pub fn record( }; let shared = rec.shared_states(); let pb = Protocol::check(mb).context("B does not fit the serving protocol")?; - let mut fixed = constants(&ma, &pa.once); - fixed.extend(constants(mb, &pb.once)); - rec.once = pa.once.iter().chain(&pb.once).cloned().collect(); + let mut fixed = constants(&ma, &[&pa.once[..], &pa.derive].concat()); + fixed.extend(constants(mb, &[&pb.once[..], &pb.derive].concat())); + rec.once = [&pa.once[..], &pa.derive, &pb.once, &pb.derive].concat(); let mb: &Manifest = mb; let chunk_name = chunk_f.as_ref().map_or("", |f| f.name.as_str()); diff --git a/docs/manifest.md b/docs/manifest.md index a5d1e616..f664cb1f 100644 --- a/docs/manifest.md +++ b/docs/manifest.md @@ -14,7 +14,7 @@ ops..impl.launches[] launch 起一次 module 入口 {"module": "argmax modules. module launch 钉住的工件 {"source": "argmax.cubin", "sha256": "…"} vars. var caller 每次调用供应的标量，有上界 states. state 不透明持久内存，runtime 按字节供应 -buffers. buffer 有类型的张量：input / output / weight / workspace / carry / peer +buffers. buffer 有类型的张量：input / output / weight / source / workspace / carry / peer topology.groups. group 多卡 SPMD 的 rank 组：只有名字和大小 ``` @@ -50,7 +50,8 @@ topology.groups. group 多卡 SPMD 的 rank 组：只有名字和 - `buffers`：`dtype + shape + kind`，可选 `domain`（内容的先验，见下）与 `fill`（在 serving 循环里的角色，见「Serving 协议」）。shape 维度是常量或 var 名；kind 说 的是"谁供应、活多久"：`input`（runtime 写入）/ `output`（runtime 读回） - / `weight`（load 时从 checkpoint 的张量拼出来，`bind` 说是哪些，见「权重」）/ `workspace`（runtime 规划，跨次执行 + / `weight`（load 时从 checkpoint 的张量拼出来，`bind` 说是哪些，见「权重」）/ `source`（同样 `bind` + checkpoint 张量，但只活在 load 里、只给 `derive` program 读，见「权重」）/ `workspace`（runtime 规划，跨次执行 不保留）/ `carry`（一个 program 写、另一个 program 读的交接棒，跨次 执行保留；谁先跑是 caller 契约，verifier 只要求它被某个 program 写到 ——投机解码的 aux 隐状态逼出来的）/ `peer`（runtime 填的地址数组，见 @@ -134,8 +135,8 @@ topology.groups. group 多卡 SPMD 的 rank 组：只有名字和 vLLM attention（unified + reduce_segments）这类"一个逻辑算子 = 多次 launch + 私有中间缓冲"整体折叠成一个 impl，不向调用方泄漏。 - `programs`：每个 program（如 `prefill`/`decode`）是 - `{"batch"?, "once"?, "calls": [...]}`——`batch` / `once` 说它怎么被 - serving 循环调（见「Serving 协议」），`calls` 是一段顺序 call 列表： + `{"batch"?, "once"?, "derive"?, "calls": [...]}`——`batch` / `once` / + `derive` 说它怎么被 serving 循环或 load 调（见「Serving 协议」），`calls` 是一段顺序 call 列表： `op` 名 + 接口实参（buffer/state 实参可带字节 `offset`，默认 0：kernel 收到 base+offset——provider 用它寻址融合 buffer 里的视图如 qkv 的 q/k/v 切片、state 里的逐层区域，offset 是 provider 布局算术的字面量， @@ -316,6 +317,17 @@ cos / sin 表、全 1 的 kv scale、chunk 索引表——不是权重，是 `ca `tools/kernels-src/weight_prep.cu` 的几个 elementwise 核）。拼接方式与派生 表都是模型知识，住在生成器（`tools/qwen_weights.py`），不住在 runtime。 +checkpoint 张量要先变换才能用（K3 的 mxfp4 专家按 GEMM 布局 shuffle），而原样 +那份没人再读时，原样的张量声明成 `source`，变换写成 `derive` program： +runtime 在 `load_weights` 里逐 call 跑它，每个 source 只在第一个读它的 call +之前上传、最后一个读它的 call 之后让出——所有 source 共用一块暂存，同一 +call 上活着的 source 互不重叠（`kern-runtime/src/derive.rs` 的 `layout`）。 +显存峰值是派生结果加上同时活着的最宽一组 source（K3 全量一层的 w13 ≈ 10 +GB），不是整份 checkpoint 两遍。跑完暂存释放、derive program 一起丢掉， +不能再跑；同一个 runtime 不能第二次 `load_weights`。derive program 不碰 +state / input / output / peer、不调 collective（它跑在这些都还没意义的时候）， +与 `once` / `batch` / `graph` 互斥。 + ## Serving 协议：`fill`、`batch`、`once` runtime 之上还有一层契约：一个 serving 循环（`kern run`、kern-serve 的 @@ -350,6 +362,8 @@ program 接受几组几行、跑完从哪读 token。v3 把这层写在 caller 测试）。 - **program 上的 `once`**：装载后跑一次、不再驱动（k3 的 `tp_init` 预填 allreduce 的 poison 值）。与 `batch` 互斥。 +- **program 上的 `derive`**：runtime 自己在 `load_weights` 里跑，caller + 不跑（见「权重」）。 - **program 上的 `graph`**：runtime 按调用形状（这次调用的 var 值）首次 capture 成一张 CUDA graph、之后 replay，一步一次 launch。给形状固定、 launch 多而短的 program（decode 步、投机 round）；行数随调用变的 @@ -410,7 +424,8 @@ kern-serve 只拿 `Protocol` 驱动 runtime，不读 JSON（CI 用 grep 保证 能接 `state`、var/表达式的取值范围必须装进标量参数类型、offset 对齐 且在界内）； 9. 逐 program 数据流：禁止读未写（read-before-write）、禁止写 input/ - weight；output / carry 必须被**某个** program 写到（prefill 这类只落 + weight/source；source 只许 `derive` program 读，`derive` program 不碰 + state / input / output / peer、不调 collective；output / carry 必须被**某个** program 写到（prefill 这类只落 state 的 program 合法地不写任何 output；carry 在每个 program 内视为 已写——它的生产者是另一个 program）；program 的 `once` 与 `batch` 互斥，`batch.groups ≥ 1`、`rows` 是 ≥ 1 的常量或已声明的 var（fill 之间 diff --git a/docs/runtime.md b/docs/runtime.md index 422676d2..827a5987 100644 --- a/docs/runtime.md +++ b/docs/runtime.md @@ -24,7 +24,8 @@ ABI 校验兼做实例选择，绕开了 capture 缺 launch→module 映射的 var max 分配全部 buffer / 分配 state（分页与 per-seq 的走下面的块池）→ 按每个 weight buffer 的 `bind` 从 checkpoint 里拷张量段拼出 buffer （张量来自一个 `Tensors`，见下「权重来源」；scratch 按 impl 声明另行私有分配） -→ 跑 `once` program 算派生表 → 顺序重放 +→ 同一次 `load_weights` 里逐 call 跑 `derive` program，`source` 张量随用随传、 +共用一块暂存（manifest.md「权重」）→ 跑 `once` program 算派生表 → 顺序重放 call 表：接口实参解析一次，逐 launch 按 `args` 连线转发/接 scratch/ 填字面量后 raw `cuLaunchKernel`（实参 staging 成小端 u64 slot；>48KB 动态 shmem 自动 `cuFuncSetAttribute`）。 diff --git a/schema/manifest-v5.schema.json b/schema/manifest-v5.schema.json index 10abf59b..f9b31887 100644 --- a/schema/manifest-v5.schema.json +++ b/schema/manifest-v5.schema.json @@ -282,7 +282,7 @@ "description": "A typed tensor, e.g. `{\"dtype\": \"bf16\", \"shape\": [\"tokens\", 2560], \"kind\": \"workspace\"}`.", "properties": { "bind": { - "description": "`weight` buffers only: the checkpoint tensors laid end to end into this buffer, in order, e.g. `[{\"tensor\": \"q_proj.weight\"}, {\"tensor\": \"k_proj.weight\"}]`.", + "description": "`weight` and `source` buffers only: the checkpoint tensors laid end to end into this buffer, in order, e.g. `[{\"tensor\": \"q_proj.weight\"}, {\"tensor\": \"k_proj.weight\"}]`.", "items": { "$ref": "#/$defs/Segment" }, @@ -373,6 +373,11 @@ "description": "Assembled at load time from the checkpoint tensors its `bind` names, e.g. `model.embed_tokens.weight`.", "type": "string" }, + { + "const": "source", + "description": "Checkpoint tensors bound like a weight but alive only while a `derive` program reads them during load: device memory holds what is derived from them, never both at once, e.g. an expert stack before its shuffle.", + "type": "string" + }, { "const": "workspace", "description": "Runtime-owned scratch, dead between runs, e.g. `hidden`.", @@ -1633,6 +1638,10 @@ }, "type": "array" }, + "derive": { + "description": "Run by the runtime inside weight loading, call by call, to turn `source` buffers into the `carry` buffers the other programs read (a shuffle into a kernel's layout, a dtype cast). Reads no state, peer, input or output; nothing else may run it.", + "type": "boolean" + }, "graph": { "description": "Driven through a CUDA graph: the runtime captures the call list at each var assignment the program is called with and replays it with one launch. For a program of fixed shape per call whose launches are many and short (a decode step, a speculative round); a program whose rows vary per call (a prefill chunk) runs launch by launch.", "type": "boolean" diff --git a/tools/kern_manifest.py b/tools/kern_manifest.py index b0035fa0..35ed9e54 100644 --- a/tools/kern_manifest.py +++ b/tools/kern_manifest.py @@ -40,12 +40,13 @@ _CALL = ["label", "op", "args"] -def program(calls, groups=None, rows=None, span=None, context=None, once=False, graph=False): +def program(calls, groups=None, rows=None, span=None, context=None, once=False, graph=False, derive=False): """A program object of the wire form: a forward of `groups` sequences of `rows` rows each (rows a constant or the name of the var fed per call; `span` the var one sequence's run of rows is sized by; `context` the var the call's longest sequence length goes in), a - once-after-load program, or a plain one. `graph`: the runtime drives + once-after-load program, a derive program the runtime runs inside + weight loading over `source` buffers, or a plain one. `graph`: the runtime drives it through a CUDA graph captured per call shape (a fixed-shape step, never a variable-row prefill).""" p = {} @@ -54,6 +55,8 @@ def program(calls, groups=None, rows=None, span=None, context=None, once=False, **({"context": context} if context else {})} if once: p["once"] = True + if derive: + p["derive"] = True if graph: p["graph"] = True p["calls"] = calls diff --git a/tools/once.py b/tools/once.py index ecfd1e6e..47499a6d 100644 --- a/tools/once.py +++ b/tools/once.py @@ -125,7 +125,9 @@ def rope(self, cos, sin, rows, half, stride, base): self.derive(name) self.call("rope_table", [buf(*cos), buf(*sin), {"i32": rows}, {"i32": half}, {"i32": stride}, {"f32": base}], rows) - def finish(self): + 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).""" for op, n in self.widest.items(): cubin, entry, params = KERNELS[op] grid, block = ([n, 1, 1], [64, 1, 1]) if op == "rope_table" else ([(n + BLOCK - 1) // BLOCK, 1, 1], [BLOCK, 1, 1]) @@ -133,6 +135,6 @@ def finish(self): self.m["ops"][op] = {"params": params, "impl": {"launches": [ {"entry": entry, "params": params, "block": block, "grid": grid, "args": [{"param": i} for i in range(len(params))], **variant(cubin).module}]}} - assert "load" not in self.m["programs"] - self.m["programs"]["load"] = program(self.calls, once=True) + assert name not in self.m["programs"] + self.m["programs"][name] = program(self.calls, once=not derive, derive=derive) return self.m