runtime: a derive program turns source tensors into weights while they load

crates/kern-manifest/src/protocol.rs

@@ -203,6 +203,9 @@ pub struct Protocol {
     pub forwards: Vec<Forward>,
     /// Programs run once after load, in name order.
     pub once: Vec<String>,
+    /// Programs the runtime runs inside weight loading, in name order: no
+    /// caller runs them.
+    pub derive: Vec<String>,
 }
 
 impl Protocol {
@@ -409,6 +412,7 @@ impl Protocol {
         // Forwards: the shape, and by dataflow what each hands back.
         let mut forwards: Vec<Forward> = 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<Bound> = None;
         let mut context_var: Option<Bound> = None;
         for (pname, p) in &m.programs {
@@ -599,6 +603,7 @@ impl Protocol {
                 line_tables,
                 forwards,
                 once,
+                derive,
             })
         } else {
             Err(ProtocolErrors(errs))

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<String>,
-    /// `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<Segment>,
 }
@@ -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",

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<String> {
                 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<String>, used_groups: &mut BTreeSet<String>| match t
@@ -340,9 +345,8 @@ fn diagnostics(m: &Manifest) -> Vec<String> {
                     }
                 }
             }
-            (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<String> {
             // 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<String> {
                 "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<String> {
             };
             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,14 +888,27 @@ fn diagnostics(m: &Manifest) -> Vec<String> {
                                 }
                             }
                         }
+                        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",
                                 c.op
                             ));
                         }
                         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<String> {
                         // 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}`"));
                         }

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<bf16>", "out buffer<bf16>", "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");
+}

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());

crates/kern-runtime/src/compile.rs

@@ -460,22 +460,24 @@ 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<String, ResolvedOp>,
     buffers: &BTreeMap<String, DeviceBuf>,
     states: &BTreeMap<String, DeviceBuf>,
     ranks: &Ranks,
     blt: &CudaBlasLT,
+    derive: bool,
 ) -> Result<BTreeMap<String, CompiledProgram>> {
     let vars: BTreeMap<&str, usize> = manifest.vars.keys().enumerate().map(|(i, s)| (s.as_str(), i)).collect();
     let maxima: Vec<u64> = 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<String, ResolvedOp>,
@@ -508,7 +510,20 @@ pub(crate) fn compile(
     blt: &CudaBlasLT,
 ) -> Result<BTreeMap<String, CompiledProgram>> {
     let peer_buffers: BTreeSet<String> = 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<String, ResolvedOp>,
+    buffers: &BTreeMap<String, DeviceBuf>,
+    ranks: &BTreeMap<String, u64>,
+    blt: &CudaBlasLT,
+) -> Result<BTreeMap<String, CompiledProgram>> {
+    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