diff --git a/crates/kern-manifest/src/protocol.rs b/crates/kern-manifest/src/protocol.rs index ca792aa..c675d54 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 0800af6..8afd00b 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 f00f0fe..0a58d29 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 899915a..0b093ce 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 ac3fa39..b011b7e 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 d2d2cf3..a745811 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 0000000..1833e84 --- /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 0d79822..a7ad4bd 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 899f272..5bc84ca 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 0000000..3265ba3 --- /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 46a4cc6..2cac7c4 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 fd882a9..9be32f7 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 a5d1e61..f664cb1 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 422676d..827a598 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 10abf59..f9b3188 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 b0035fa..35ed9e5 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 ecfd1e6..47499a6 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