From a126372a20165f4840bcba141c9459ad21f6104b Mon Sep 17 00:00:00 2001 From: JinYan Su <751080330@qq.com> Date: Wed, 7 Oct 2026 15:33:39 +0000 Subject: [PATCH 01/22] manifest: a manifest cuts into pipeline stages at call indices `kern cut --at --module pp_mailbox.cubin --out ` splits a verified manifest into S stage manifests. Liveness across each cut point finds the boundary workspaces; they go through a send/recv op pair over a mailbox on the downstream GPU (tools/kernels-src/pp_mailbox.cu: full/empty counters, system-scope release/acquire, a device-side spin that traps after 60 s). Each stage carries `cut {id, stage, stages}` and a `clock` output fill the mailbox kernels stamp per item with globaltimer. A stage's group `pp.` joins two different manifests that agree only on the peer's `of` buffer; import_peers already matches by name and size, so the runtime is unchanged. Protocol requires a token fill on stage 0 only. Gate: a property test (400 random manifests, symbolic interpreter) shows every cut that is made computes the whole; qwen3-4b cut at l18 into two stages, run in sequence (examples/pp_chain), gives bitwise identical logits and tokens to the whole manifest on 7 prompts, chunk 512 and 64. Co-Authored-By: Claude Opus 5.5 (1M context) --- crates/kern-manifest/src/constants.rs | 2 + crates/kern-manifest/src/cut.rs | 485 ++++++++++++++++++++++++++ crates/kern-manifest/src/lib.rs | 1 + crates/kern-manifest/src/protocol.rs | 18 +- crates/kern-manifest/src/types.rs | 20 +- crates/kern-manifest/src/verify.rs | 11 +- crates/kern-manifest/tests/cut.rs | 294 ++++++++++++++++ crates/kern-run/examples/pp_chain.rs | 233 +++++++++++++ crates/kern-run/src/bin/kern.rs | 48 ++- crates/kern-run/src/lib.rs | 82 +++-- crates/kern-serve/src/tray.rs | 2 +- docs/manifest.md | 41 ++- schema/manifest-v5.schema.json | 66 +++- tools/kernels-src/pp_mailbox.cu | 150 ++++++++ 14 files changed, 1420 insertions(+), 33 deletions(-) create mode 100644 crates/kern-manifest/src/cut.rs create mode 100644 crates/kern-manifest/tests/cut.rs create mode 100644 crates/kern-run/examples/pp_chain.rs create mode 100644 tools/kernels-src/pp_mailbox.cu diff --git a/crates/kern-manifest/src/constants.rs b/crates/kern-manifest/src/constants.rs index 0838f6d1..a77b9dec 100644 --- a/crates/kern-manifest/src/constants.rs +++ b/crates/kern-manifest/src/constants.rs @@ -18,6 +18,8 @@ struct LiteralManifest { #[serde(default)] topology: Option, #[serde(default)] + cut: Option, + #[serde(default)] vars: BTreeMap, #[serde(default)] states: BTreeMap, diff --git a/crates/kern-manifest/src/cut.rs b/crates/kern-manifest/src/cut.rs new file mode 100644 index 00000000..1023c6fe --- /dev/null +++ b/crates/kern-manifest/src/cut.rs @@ -0,0 +1,485 @@ +//! Cutting a manifest into a pipeline: stage `s` runs the calls between +//! two cut points of every program a serving loop drives, and the +//! activations live across a cut travel through a mailbox. +//! +//! A forward is a pure function of its inputs and its states, so a cut +//! changes no arithmetic: the stages run the same calls in the same order +//! on the same values, and the pipeline's outputs are the unsplit +//! manifest's byte for byte. What a cut adds is transport. At a cut point +//! the buffers live across it are the ones the calls after it read before +//! they write them whole (an `out` at offset 0 is taken to write the whole +//! buffer, as every generator writes a view at its base) and that the +//! calls before it wrote; each must be a workspace, since an output is read +//! from the last stage and a carry would cross programs. The upstream stage +//! ends every program by putting them into a box on the downstream GPU +//! (`tools/kernels-src/pp_mailbox.cu`): wait until the box is empty, put, +//! post. The downstream stage starts by waiting for the post, taking them +//! out and freeing the box. One box per edge, laid out per program at the +//! bounds of its vars; the box sits on the downstream GPU and the upstream +//! reaches it through a `peer` buffer of a two-member topology group +//! `pp.`, its own copy of the box unused. +//! +//! A state never travels. Every region of a state (a state and a call's +//! offset into it) belongs to the one stage whose calls touch it, which is +//! what makes the stages' page tables interchangeable: a sequence's pages +//! hold the first stage's layers on one GPU and the next stage's on +//! another, under the same page ids. A region two stages touch is refused. +//! A program run once after load is sliced per stage to the calls the +//! stage's own calls need, so every stage derives its own tables. +//! +//! The prompt reaches the step program's first token the way it does +//! unsplit; a chunk program that hands a token back itself is refused, so +//! a pipeline's first token always comes from a step. Everything a stage +//! does not use is dropped, and every stage verifies. + +use std::collections::{BTreeMap, BTreeSet}; + +use serde_json::json; + +use crate::types::*; +use crate::verify::{verify, Verified}; + +/// How long a mailbox wait spins before the stage gives up and traps. +pub const TIMEOUT_NS: i64 = 60_000_000_000; +/// Bytes of a box before its data: the two counters, a cache line each. +const HEADER: u64 = 256; +/// Each carried buffer starts at a multiple of this in the box. +const ALIGN: u64 = 256; +/// Entries of a stage's clock. +const RING: u64 = 256; +/// The module name the pipeline kernels go under in every stage. +const MODULE: &str = "pp"; +/// The stage's clock (fill `clock`). +const CLOCK: &str = "pp.clock"; + +/// Every reason a manifest cannot be cut where asked, reported together. +#[derive(Debug)] +pub struct CutErrors(pub Vec); + +impl std::fmt::Display for CutErrors { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str("manifest cannot be cut:")?; + for e in &self.0 { + write!(f, "\n - {e}")?; + } + Ok(()) + } +} + +impl std::error::Error for CutErrors {} + +/// A buffer an edge carries in one program: where in the box, and how +/// many bytes per call. +#[derive(Debug, Clone)] +struct Carried { + buffer: String, + dtype: DType, + offset: u64, + bytes: Expr, +} + +/// Cut `m` before each call index in `at` of every program with a +/// `batch` (all of them must be longer than the last point), wiring in +/// the mailbox kernels of `pp`. Returns the stages in order, each +/// verified. +pub fn cut(m: &Verified, at: &[usize], pp: &Module) -> Result, CutErrors> { + let mut errs = Vec::new(); + if m.topology.is_some() { + errs.push("the manifest has a topology; a cut takes a single-rank manifest".to_string()); + } + if m.cut.is_some() { + errs.push("the manifest is already a stage of a cut".to_string()); + } + if at.is_empty() || at[0] == 0 || at.windows(2).any(|w| w[0] >= w[1]) { + errs.push(format!("cut points {at:?}: at least one, increasing, none at 0")); + } + let driven: Vec<(&String, &Program)> = m.programs.iter().filter(|(_, p)| p.batch.is_some()).collect(); + if driven.is_empty() { + errs.push("no program declares a `batch`: nothing to cut".to_string()); + } + for (name, p) in &driven { + if let Some(&last) = at.last() { + if last >= p.calls.len() { + errs.push(format!("program `{name}` has {} calls; it cannot be cut before call {last}", p.calls.len())); + } + } + let chunk = p.batch.as_ref().is_some_and(|b| matches!(b.rows, Dim::Var(_))); + if chunk && writes(m, &p.calls).iter().any(|b| m.buffers[*b].fill == Some(Fill::Tokens)) { + errs.push(format!( + "program `{name}` takes a chunk and hands a token back; a pipeline takes its first token from a step" + )); + } + } + if !errs.is_empty() { + return Err(CutErrors(errs)); + } + + let stages = at.len() + 1; + let bounds: Vec = std::iter::once(0).chain(at.iter().copied()).collect(); + let span = |p: &Program, s: usize| bounds[s]..bounds.get(s + 1).copied().unwrap_or(p.calls.len()); + let max: BTreeMap = m.vars.iter().map(|(k, v)| (k.clone(), v.max)).collect(); + + // What each edge carries in each program, and each box's size. + let mut carried: BTreeMap<(&str, usize), Vec> = BTreeMap::new(); + for (name, p) in &driven { + for (e, &a) in at.iter().enumerate() { + let before = writes(m, &p.calls[..a]); + let mut offset = 0; + let mut list = Vec::new(); + for b in live_in(m, &p.calls[a..]).into_iter().filter(|b| before.contains(b)) { + let buf = &m.buffers[b]; + if buf.kind != BufferKind::Workspace { + errs.push(format!( + "program `{name}`: {} buffer `{b}` is live across the cut before call {a}; only a workspace travels", + buf.kind + )); + continue; + } + match bytes_of(buf) { + Some(bytes) => { + let at_max = bytes.eval(&max).unwrap_or(u64::MAX); + list.push(Carried { buffer: b.to_string(), dtype: buf.dtype, offset, bytes }); + offset += at_max.div_ceil(ALIGN) * ALIGN; + } + None => errs.push(format!( + "program `{name}`: buffer `{b}` is shaped {:?}; a carried buffer has at most one var dim", + buf.shape + )), + } + } + carried.insert((name.as_str(), e), list); + } + for (i, c) in p.calls.iter().enumerate().take(bounds[stages - 1]) { + for (b, _, dir) in accesses(m, c) { + if dir != Dir::In && matches!(m.buffers[b].kind, BufferKind::Output | BufferKind::Carry) { + errs.push(format!( + "program `{name}` call {i} writes {} buffer `{b}` before the last stage; a pipeline's outputs come from its last stage", + m.buffers[b].kind + )); + } + } + } + } + let mut owner: BTreeMap<(String, u64), usize> = BTreeMap::new(); + for (name, p) in &driven { + for s in 0..stages { + for c in &p.calls[span(p, s)] { + for (st, off) in state_args(c) { + match owner.insert((st.to_string(), off), s) { + Some(o) if o != s => errs.push(format!( + "program `{name}`: state `{st}`+{off} is touched by stages {o} and {s}; a state region belongs to one stage" + )), + _ => {} + } + } + } + } + } + if !errs.is_empty() { + return Err(CutErrors(errs)); + } + let boxes: Vec = (0..at.len()) + .map(|e| { + let data = carried + .iter() + .filter(|((_, ee), _)| *ee == e) + .map(|(_, list)| list.last().map_or(0, |c| c.offset + c.bytes.eval(&max).unwrap_or(0))) + .max() + .unwrap_or(0); + HEADER + data.max(1).div_ceil(ALIGN) * ALIGN + }) + .collect(); + + let id = id_of(m, at); + let built: Vec = (0..stages) + .map(|s| { + let programs: BTreeMap = driven + .iter() + .map(|(name, p)| { + let into = (s > 0).then(|| &carried[&(name.as_str(), s - 1)][..]); + let out = (s + 1 < stages).then(|| &carried[&(name.as_str(), s)][..]); + let calls = + receive(s, into).into_iter().chain(p.calls[span(p, s)].iter().cloned()).chain(send(s, out)); + ((*name).clone(), Program { calls: calls.collect(), ..(*p).clone() }) + }) + .collect(); + stage(m, programs, s, stages, &boxes, &id, pp) + }) + .collect(); + built + .into_iter() + .enumerate() + .map(|(s, st)| verify(st).map_err(|e| CutErrors(e.0.into_iter().map(|d| format!("stage {s}: {d}")).collect()))) + .collect() +} + +fn box_name(e: usize) -> String { + format!("pp.{e}.box") +} + +fn peer_name(e: usize) -> String { + format!("pp.{e}.peer") +} + +fn call(op: &str, args: serde_json::Value) -> Call { + serde_json::from_value(json!({"label": op, "op": op, "args": args})).expect("a well-formed call") +} + +/// The calls that start stage `s`: take what the edge into it carries, or +/// stamp the start of a first stage. +fn receive(s: usize, into: Option<&[Carried]>) -> Vec { + let Some(list) = into else { + return vec![call("pp.stamp_first", json!([{"buf": peer_name(s)}, {"buf": CLOCK}]))]; + }; + let b = box_name(s - 1); + std::iter::once(call("pp.wait_full", json!([{"buf": b}, {"buf": CLOCK}]))) + .chain(list.iter().map(|c| { + call( + &format!("pp.take.{}", c.dtype), + json!([{"buf": b}, {"buf": c.buffer}, {"expr": c.bytes}, {"i64": c.offset}]), + ) + })) + .chain(std::iter::once(call("pp.post_empty", json!([{"buf": b}])))) + .collect() +} + +/// The calls that end stage `s`: put what the edge out of it carries, or +/// stamp the end of a last stage. +fn send(s: usize, out: Option<&[Carried]>) -> Vec { + let Some(list) = out else { + return vec![call("pp.stamp_last", json!([{"buf": box_name(s - 1)}, {"buf": CLOCK}]))]; + }; + let p = peer_name(s); + std::iter::once(call("pp.wait_empty", json!([{"buf": p}, {"buf": CLOCK}]))) + .chain(list.iter().map(|c| { + call( + &format!("pp.put.{}", c.dtype), + json!([{"buf": p}, {"buf": c.buffer}, {"expr": c.bytes}, {"i64": c.offset}]), + ) + })) + .chain(std::iter::once(call("pp.post_full", json!([{"buf": p}])))) + .collect() +} + +/// The pipeline ops a stage's calls name: `pp.put.` and +/// `pp.take.` per carried dtype, the waits, posts and stamps. +fn pp_op(name: &str) -> Op { + let single = |entry: &str, params: serde_json::Value| { + json!({"params": params, "impl": {"launches": [ + {"module": MODULE, "entry": entry, "block": [1, 1, 1], "grid": [1, 1, 1]}]}}) + }; + let wait = |entry: &str, first: &str| { + json!({"params": [first, "out buffer"], "impl": {"launches": [ + {"module": MODULE, "entry": entry, "block": [1, 1, 1], "grid": [1, 1, 1], + "params": [first, "out buffer", "i64"], + "args": [{"param": 0}, {"param": 1}, {"i64": TIMEOUT_NS}]}]}}) + }; + let copy = |entry: &str, params: serde_json::Value| { + json!({"params": params, "impl": {"launches": [ + {"module": MODULE, "entry": entry, "block": [256, 1, 1], "grid": [264, 1, 1]}]}}) + }; + let v = match name.split('.').collect::>().as_slice() { + ["pp", "wait_empty"] => wait("kern_pp_wait_empty", "in buffer"), + ["pp", "wait_full"] => wait("kern_pp_wait_full", "in buffer"), + ["pp", "post_full"] => single("kern_pp_post_full", json!(["in buffer"])), + ["pp", "post_empty"] => single("kern_pp_post_empty", json!(["inout buffer"])), + ["pp", "stamp_first"] => single("kern_pp_stamp_first", json!(["in buffer", "out buffer"])), + ["pp", "stamp_last"] => single("kern_pp_stamp_last", json!(["in buffer", "out buffer"])), + ["pp", "put", dt] => copy("kern_pp_put", json!(["in buffer", format!("in buffer<{dt}>"), "i64", "i64"])), + ["pp", "take", dt] => copy("kern_pp_take", json!(["in buffer", format!("out buffer<{dt}>"), "i64", "i64"])), + _ => unreachable!("the cut names only the pipeline ops above"), + }; + serde_json::from_value(v).expect("a well-formed pipeline op") +} + +/// Stage `s` of `stages`: `programs` (already cut and wired), the once +/// programs sliced to what they need, and every declaration they use. +fn stage( + m: &Manifest, + mut programs: BTreeMap, + s: usize, + stages: usize, + boxes: &[u64], + id: &str, + pp: &Module, +) -> Manifest { + let mut needed: BTreeSet = + programs.values().flat_map(|p| &p.calls).flat_map(|c| reads(m, c)).map(str::to_string).collect(); + for (name, p) in m.programs.iter().filter(|(_, p)| p.once) { + let mut kept = Vec::new(); + for c in p.calls.iter().rev() { + if accesses(m, c).iter().any(|(b, _, d)| *d != Dir::In && needed.contains(*b)) { + needed.extend(reads(m, c).map(str::to_string)); + kept.push(c.clone()); + } + } + if !kept.is_empty() { + kept.reverse(); + programs.insert(name.clone(), Program { calls: kept, ..p.clone() }); + } + } + + let mut buffers: BTreeMap = BTreeMap::new(); + let u8_box = |bytes: u64, kind: &str| -> Buffer { + serde_json::from_value(json!({"dtype": "u8", "shape": [bytes], "kind": kind, "export": true})).expect("a box") + }; + buffers.insert( + CLOCK.to_string(), + serde_json::from_value(json!({"dtype": "i64", "shape": [RING * 5], "kind": "output", "fill": "clock"})) + .expect("a clock"), + ); + if s > 0 { + buffers.insert(box_name(s - 1), u8_box(boxes[s - 1], "carry")); + } + let mut topology = None; + if s + 1 < stages { + buffers.insert(box_name(s), u8_box(boxes[s], "workspace")); + buffers.insert( + peer_name(s), + serde_json::from_value( + json!({"dtype": "u64", "shape": [2], "kind": "peer", "of": box_name(s), "group": format!("pp.{s}")}), + ) + .expect("a peer"), + ); + topology = Some(Topology { groups: BTreeMap::from([(format!("pp.{s}"), 2)]) }); + } + let calls: Vec<&Call> = programs.values().flat_map(|p| &p.calls).collect(); + for c in &calls { + for a in &c.args { + if let Arg::Buf { buf, .. } = a { + if let Some(b) = m.buffers.get(buf) { + buffers.insert(buf.clone(), b.clone()); + } + } + } + } + let states: BTreeMap = + calls.iter().flat_map(|c| state_args(c)).map(|(st, _)| (st.to_string(), m.states[st].clone())).collect(); + let names: BTreeSet = buffers.keys().chain(states.keys()).cloned().collect(); + for b in buffers.values_mut() { + if b.domain.as_ref().and_then(|d| d.index_into.as_ref()).is_some_and(|t| !names.contains(t)) { + b.domain = None; + } + } + let ops: BTreeMap = + calls.iter().map(|c| (c.op.clone(), m.ops.get(&c.op).cloned().unwrap_or_else(|| pp_op(&c.op)))).collect(); + let mut modules: BTreeMap = ops + .values() + .flat_map(|o| &o.imp.launches) + .filter_map(Launch::module) + .filter(|n| *n != MODULE) + .map(|n| (n.to_string(), m.modules[n].clone())) + .collect(); + modules.insert(MODULE.to_string(), pp.clone()); + // A var stays when some kept declaration names it: shapes, grids, + // call args, batches and domains all spell a var as a bare string. + let mut strings = BTreeSet::new(); + for v in [serde_json::to_value(&buffers), serde_json::to_value(&ops), serde_json::to_value(&programs)] { + collect_strings(&v.expect("serializable"), &mut strings); + } + let vars = m.vars.iter().filter(|(k, _)| strings.contains(*k)).map(|(k, v)| (k.clone(), v.clone())).collect(); + Manifest { + schema_version: m.schema_version, + model: m.model.clone(), + topology, + cut: Some(Cut { id: id.to_string(), stage: s as u64, stages: stages as u64 }), + vars, + states, + buffers, + modules, + ops, + programs, + } +} + +fn collect_strings(v: &serde_json::Value, out: &mut BTreeSet) { + match v { + serde_json::Value::String(s) => { + out.insert(s.clone()); + } + serde_json::Value::Array(a) => a.iter().for_each(|x| collect_strings(x, out)), + serde_json::Value::Object(o) => o.values().for_each(|x| collect_strings(x, out)), + _ => {} + } +} + +/// Bytes one call moves of `b`: its size at the call's vars, when at most +/// one dim is a var. +fn bytes_of(b: &Buffer) -> Option { + let constant: u64 = b.shape.iter().filter_map(|d| if let Dim::Const(c) = d { Some(*c) } else { None }).product(); + let vars: Vec<&String> = b.shape.iter().filter_map(|d| if let Dim::Var(v) = d { Some(v) } else { None }).collect(); + let c = constant * b.dtype.bytes(); + match vars.as_slice() { + [] => Some(Expr::Const(c)), + [v] => Some(Expr::Mul { mul: (Box::new(Expr::Var((*v).clone())), c) }), + _ => None, + } +} + +/// Every buffer argument of `c`: name, offset and the direction its op +/// declares. +fn accesses<'a>(m: &'a Manifest, c: &'a Call) -> Vec<(&'a str, u64, Dir)> { + let params = m.ops.get(&c.op).map(|o| o.params.as_slice()).unwrap_or(&[]); + c.args + .iter() + .zip(params) + .filter_map(|(a, p)| match a { + Arg::Buf { buf, offset } => Some((buf.as_str(), *offset, p.dir().unwrap_or(Dir::In))), + _ => None, + }) + .collect() +} + +/// The buffers `c` reads, by name; a pipeline call reads its buffers too. +fn reads<'a>(m: &'a Manifest, c: &'a Call) -> impl Iterator { + let known = m.ops.contains_key(&c.op); + let acc = accesses(m, c); + let pipeline: Vec<&str> = if known { + Vec::new() + } else { + c.args.iter().filter_map(|a| if let Arg::Buf { buf, .. } = a { Some(buf.as_str()) } else { None }).collect() + }; + acc.into_iter().filter(|(_, _, d)| *d != Dir::Out).map(|(b, _, _)| b).chain(pipeline) +} + +/// The buffers `calls` write. +fn writes<'a>(m: &'a Manifest, calls: &'a [Call]) -> BTreeSet<&'a str> { + calls.iter().flat_map(|c| accesses(m, c)).filter(|(_, _, d)| *d != Dir::In).map(|(b, _, _)| b).collect() +} + +/// The buffers `calls` read before writing them whole: what must already +/// hold its value when they start. +fn live_in<'a>(m: &'a Manifest, calls: &'a [Call]) -> BTreeSet<&'a str> { + let mut killed = BTreeSet::new(); + let mut live = BTreeSet::new(); + for c in calls { + let acc = accesses(m, c); + for &(b, _, d) in &acc { + if d != Dir::Out && !killed.contains(b) { + live.insert(b); + } + } + for &(b, off, d) in &acc { + if d == Dir::Out && off == 0 { + killed.insert(b); + } + } + } + live +} + +/// Every state region `c` touches. +fn state_args(c: &Call) -> impl Iterator { + c.args.iter().filter_map(|a| match a { + Arg::State { state, offset } => Some((state.as_str(), *offset)), + _ => None, + }) +} + +/// The cut's name: FNV-1a over the manifest and the cut points, the same +/// bytes for the same cut. +fn id_of(m: &Manifest, at: &[usize]) -> String { + let bytes = m.to_json().into_bytes().into_iter().chain(at.iter().flat_map(|a| (*a as u64).to_le_bytes())); + let h = bytes.fold(0xcbf2_9ce4_8422_2325u64, |h, b| (h ^ b as u64).wrapping_mul(0x0000_0100_0000_01b3)); + format!("{h:016x}") +} diff --git a/crates/kern-manifest/src/lib.rs b/crates/kern-manifest/src/lib.rs index d5a684b5..1093c58a 100644 --- a/crates/kern-manifest/src/lib.rs +++ b/crates/kern-manifest/src/lib.rs @@ -17,6 +17,7 @@ mod constants; mod host; pub use constants::json_schema; +pub mod cut; pub mod protocol; pub mod types; pub mod values; diff --git a/crates/kern-manifest/src/protocol.rs b/crates/kern-manifest/src/protocol.rs index ca792aa7..5f28a060 100644 --- a/crates/kern-manifest/src/protocol.rs +++ b/crates/kern-manifest/src/protocol.rs @@ -322,6 +322,7 @@ impl Protocol { Fill::Tokens => axis == Axis::Groups || (axis == Axis::Tray && width == 1), Fill::SpanAt | Fill::Error => axis == Axis::Fixed(1) && b.dtype == DType::I32, Fill::Blocks => matches!(axis, Axis::Fixed(n) if n >= 2) && b.dtype == DType::I32, + Fill::Clock => matches!(axis, Axis::Fixed(n) if n > 0 && n % 5 == 0) && b.dtype == DType::I64, }; if !ok { errs.push(format!( @@ -335,6 +336,7 @@ impl Protocol { Fill::Tokens => "expected [groups], [groups, w] or [tray]", Fill::SpanAt | Fill::Error => "expected i32 [1]", Fill::Blocks => "expected i32 [members + 1]", + Fill::Clock => "expected i64 [ring × 5]", } )); continue; @@ -347,14 +349,24 @@ impl Protocol { errs.push(format!("fill `{fill}` over {axis:?} is on {n} buffers, expected one")); } } - if !fills.iter().any(|f| f.fill == Fill::Token && f.axis != Axis::Groups) { + // A stage past the first of a cut takes activations, not tokens. + let fed = m.cut.as_ref().is_none_or(|c| c.stage == 0); + if fed && !fills.iter().any(|f| f.fill == Fill::Token && f.axis != Axis::Groups) { errs.push("no input has fill `token` over the rows: nothing carries the tokens a call feeds".into()); } if fills.iter().filter(|f| f.fill == Fill::Token && f.axis == Axis::Groups).count() > 1 { errs.push("fill `token` over the sequences is on more than one buffer".into()); } - for fill in [Fill::Position, Fill::Valid, Fill::CuSeqlens, Fill::SpanAt, Fill::Blocks, Fill::Count, Fill::Error] - { + for fill in [ + Fill::Position, + Fill::Valid, + Fill::CuSeqlens, + Fill::SpanAt, + Fill::Blocks, + Fill::Count, + Fill::Error, + Fill::Clock, + ] { if one(fill).is_none() && m.buffers.values().any(|b| b.fill == Some(fill)) { errs.push(format!("fill `{fill}` is on more than one buffer")); } diff --git a/crates/kern-manifest/src/types.rs b/crates/kern-manifest/src/types.rs index 0800af6e..af89c7ce 100644 --- a/crates/kern-manifest/src/types.rs +++ b/crates/kern-manifest/src/types.rs @@ -71,9 +71,12 @@ pub struct Manifest { pub schema_version: u32, /// Free-form model label, e.g. `"qwen3-4b"`. pub model: String, - /// Rank groups a multi-GPU manifest is SPMD over, e.g. `{"groups": {"ep": 4}}`; every rank loads the same manifest. + /// Rank groups of a multi-GPU manifest, e.g. `{"groups": {"ep": 4}}`: every rank of an SPMD group loads the same manifest; the two stages a cut's group joins load their own, agreeing only on the peer's `of` buffer. #[serde(default, skip_serializing_if = "Option::is_none")] pub topology: Option, + /// The pipeline stage this manifest is, when it is one piece of a cut manifest, e.g. `{"id": "9f2c…", "stage": 1, "stages": 2}`. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub cut: Option, /// Per-call scalars the caller supplies, e.g. `{"tokens": {"max": 2048}}`. #[serde(default, deserialize_with = "unique_map", skip_serializing_if = "BTreeMap::is_empty")] pub vars: BTreeMap, @@ -197,6 +200,8 @@ pub enum Fill { Count, /// Output: a one-word error flag of the collectives, read after every call; nonzero is a failed step. Error, + /// Output: a pipeline stage's clock, `i64 [ring × 5]`: entry `item % ring` is {item + 1, receive wait begin, end, send wait begin, end} in globaltimer ns, 0 where the stage has no such edge. A cut declares it; the stage reads it when it pleases. + Clock, } impl fmt::Display for Fill { @@ -213,6 +218,7 @@ impl fmt::Display for Fill { Fill::Tokens => "tokens", Fill::Count => "count", Fill::Error => "error", + Fill::Clock => "clock", }) } } @@ -226,6 +232,18 @@ pub struct Topology { pub groups: BTreeMap, } +/// One stage of a manifest cut into a pipeline (`kern cut`): every stage of one cut carries the same `id`; stage `stage` of `stages` receives its predecessor's activations and sends its own to the next. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)] +#[serde(deny_unknown_fields)] +pub struct Cut { + /// Names the cut: a hash of the manifest it was cut from and where. + pub id: String, + /// This stage's place in the pipeline, from 0. + pub stage: u64, + /// Stages in the pipeline, at least 2. + pub stages: u64, +} + /// A per-call scalar the caller supplies, bounded `1..=max`; the only kind of number that may size a shape or a grid. #[derive(Debug, Clone, Serialize, Deserialize, schemars::JsonSchema)] #[serde(deny_unknown_fields)] diff --git a/crates/kern-manifest/src/verify.rs b/crates/kern-manifest/src/verify.rs index f00f0fe3..3aa8bd64 100644 --- a/crates/kern-manifest/src/verify.rs +++ b/crates/kern-manifest/src/verify.rs @@ -37,6 +37,7 @@ //! `of` and `group` only where they mean something; `{"rank": g}` //! binds only to i32/i64 params; an op with an extern launch never //! receives a peer buffer (runtime built-ins may not touch peer memory) +//! 11. cut: at least two stages, the stage below their count, an id //! //! What this deliberately cannot check: kernel *behavior*, and the //! *semantics* of interface params (that a replacement implementation @@ -195,6 +196,14 @@ fn diagnostics(m: &Manifest) -> Vec { } } } + if let Some(c) = &m.cut { + if c.id.is_empty() || c.stages < 2 || c.stage >= c.stages { + errs.push(format!( + "cut: stage {} of {} with id `{}`; a cut has at least 2 stages, a stage below their count and an id", + c.stage, c.stages, c.id + )); + } + } let group_ctx = |g: &str, errs: &mut Vec, used_groups: &mut BTreeSet, ctx: &str| -> Option { match m.group_size(g) { Some(sz) => { @@ -241,7 +250,7 @@ fn diagnostics(m: &Manifest) -> Vec { let ctx = format!("buffer `{name}`"); // 9b. fill if let Some(fill) = b.fill { - let output = matches!(fill, Fill::Tokens | Fill::Count | Fill::Error); + let output = matches!(fill, Fill::Tokens | Fill::Count | Fill::Error | Fill::Clock); match (b.kind, output) { (BufferKind::Input, false) | (BufferKind::Output, true) => {} (BufferKind::Input | BufferKind::Output, _) => errs.push(format!( diff --git a/crates/kern-manifest/tests/cut.rs b/crates/kern-manifest/tests/cut.rs new file mode 100644 index 00000000..748f1135 --- /dev/null +++ b/crates/kern-manifest/tests/cut.rs @@ -0,0 +1,294 @@ +//! A cut changes transport, never arithmetic: run symbolically, the stages +//! of any cut that verifies compute what the whole manifest computes. + +use std::collections::BTreeMap; + +use kern_manifest::cut::cut; +use kern_manifest::types::{Arg, Dir, Manifest, Module, ParamType}; +use kern_manifest::{verify, Protocol, Verified}; +use serde_json::{json, Value}; + +const QWEN3: &str = include_str!("../../../examples/qwen3-4b.json"); + +fn pp() -> Module { + Module { source: "pp_mailbox.cubin".into(), sha256: "ab".repeat(32) } +} + +struct Rand(u64); + +impl Rand { + fn next(&mut self, n: usize) -> usize { + self.0 ^= self.0 << 13; + self.0 ^= self.0 >> 7; + self.0 ^= self.0 << 17; + (self.0 >> 33) as usize % n + } +} + +fn hash(parts: &[u64]) -> u64 { + parts + .iter() + .flat_map(|p| p.to_le_bytes()) + .fold(0xcbf2_9ce4_8422_2325, |h, b| (h ^ b as u64).wrapping_mul(0x100_0000_01b3)) +} + +fn name_hash(s: &str) -> u64 { + hash(&s.bytes().map(u64::from).collect::>()) +} + +/// A random straight-line program over four workspaces, an input, a carry +/// the once program derives and one paged state in per-layer regions, its +/// calls reading only what is written. Ops are opaque functions of what +/// they read. +fn random_manifest(r: &mut Rand) -> Manifest { + let ops = json!({ + "f1": ["in buffer", "out buffer"], + "f2": ["in buffer", "in buffer", "out buffer"], + "g": ["inout buffer"], + "sw": ["in buffer", "inout state"], + "sr": ["in state", "in buffer", "out buffer"], + "z": ["out buffer"] + }); + let ws = ["w0", "w1", "w2", "w3"]; + let mut readable = vec!["x", "c"]; + let mut calls: Vec = vec![json!({"op": "f1", "args": [{"buf": "x"}, {"buf": ws[r.next(4)]}]})]; + let pick = |r: &mut Rand, from: &[&'static str]| from[r.next(from.len())]; + readable.push(match &calls[0]["args"][1]["buf"] { + Value::String(s) => ws.into_iter().find(|w| w == s).unwrap(), + _ => unreachable!(), + }); + for i in 1..4 + r.next(9) { + // A layer is three calls and owns one region of the state. + let region = json!({"state": "s", "offset": 4 * (i / 3)}); + let out = ws[r.next(4)]; + let c = match r.next(5) { + 0 => json!({"op": "f1", "args": [{"buf": pick(r, &readable)}, {"buf": out}]}), + 1 => json!({"op": "f2", "args": [{"buf": pick(r, &readable)}, {"buf": pick(r, &readable)}, {"buf": out}]}), + 2 => { + let b = pick(r, &readable[2..]); + json!({"op": "g", "args": [{"buf": b}]}) + } + 3 => json!({"op": "sw", "args": [{"buf": pick(r, &readable)}, region]}), + _ => json!({"op": "sr", "args": [region, {"buf": pick(r, &readable)}, {"buf": out}]}), + }; + if c["op"] != "g" && c["op"] != "sw" && !readable.contains(&out) { + readable.push(out); + } + calls.push(c); + } + calls.push(json!({"op": "f1", "args": [{"buf": pick(r, &readable)}, {"buf": "y"}]})); + for (i, c) in calls.iter_mut().enumerate() { + c["label"] = json!(format!("c{i}")); + } + let used: Vec = calls.iter().map(|c| c["op"].as_str().unwrap().to_string()).chain(["z".into()]).collect(); + let buf = |kind: &str, shape: Value| json!({"dtype": "bf16", "shape": shape, "kind": kind}); + let mut buffers = json!({ + "x": buf("input", json!(["tokens", 8])), "y": buf("output", json!(["tokens", 8])), + "c": buf("carry", json!([8])), "d": buf("carry", json!([8])) + }); + let named = serde_json::to_string(&calls).unwrap(); + for w in ws.iter().filter(|w| named.contains(&format!("\"{w}\""))) { + buffers[*w] = buf("workspace", json!(["tokens", 8])); + } + let ops: serde_json::Map = ops + .as_object() + .unwrap() + .iter() + .filter(|(k, _)| used.contains(k)) + .map(|(k, p)| (k.clone(), json!({"params": p, "impl": {"launches": [{"entry": "extern:x"}]}}))) + .collect(); + let uses_state = named.contains("\"state\""); + Manifest::from_json( + &json!({ + "schema_version": 5, "model": "t", "vars": {"tokens": {"max": 4}}, + "states": if uses_state { json!({"s": {"bytes_per_token": 16}}) } else { json!({}) }, + "buffers": buffers, "modules": {}, "ops": ops, + "programs": { + "init": {"once": true, "calls": [ + {"label": "i0", "op": "z", "args": [{"buf": "c"}]}, + {"label": "i1", "op": "z", "args": [{"buf": "d"}]}]}, + "p": {"batch": {"groups": 1, "rows": 1}, "calls": calls} + } + }) + .to_string(), + ) + .unwrap() +} + +/// What one stage (or the whole manifest) leaves behind: its buffers and +/// the state regions it wrote. +#[derive(Default)] +struct Machine { + bufs: BTreeMap, + states: BTreeMap<(String, u64), u64>, +} + +/// The edge a pipeline buffer belongs to: `pp..box` / `pp..peer`. +fn edge(name: &str) -> u64 { + name.split('.').nth(1).unwrap().parse().unwrap() +} + +impl Machine { + /// Run `program` of `m`: every op a hash of what it reads into what it + /// writes; the pipeline ops move values through `mail`. + fn run(&mut self, m: &Manifest, program: &str, mail: &mut BTreeMap<(u64, i64), u64>) { + for c in &m.programs[program].calls { + let label = c.label.clone().unwrap(); + let buf = |a: &Arg| match a { + Arg::Buf { buf, .. } => buf.clone(), + _ => unreachable!(), + }; + let off = |a: &Arg| match a { + Arg::I64 { i64: v } => *v, + _ => unreachable!(), + }; + if c.op.starts_with("pp.put.") { + mail.insert((edge(&buf(&c.args[0])), off(&c.args[3])), self.bufs[&buf(&c.args[1])]); + continue; + } + if c.op.starts_with("pp.take.") { + self.bufs.insert(buf(&c.args[1]), mail[&(edge(&buf(&c.args[0])), off(&c.args[3]))]); + continue; + } + if c.op.starts_with("pp.") { + continue; + } + let params = &m.ops[&c.op].params; + let mut read = vec![name_hash(&c.op), name_hash(&label)]; + for (a, p) in c.args.iter().zip(params) { + match (a, p) { + (Arg::Buf { buf, .. }, ParamType::Buf { dir: Dir::In | Dir::InOut, .. }) => { + read.push(*self.bufs.get(buf).unwrap_or_else(|| panic!("{label} reads unset `{buf}`"))) + } + (Arg::State { state, offset }, ParamType::State { dir: Dir::In | Dir::InOut }) => { + read.push(*self.states.get(&(state.clone(), *offset)).unwrap_or(&0)) + } + _ => {} + } + } + let h = hash(&read); + for (j, (a, p)) in c.args.iter().zip(params).enumerate() { + let v = hash(&[h, j as u64]); + match (a, p) { + (Arg::Buf { buf, .. }, ParamType::Buf { dir: Dir::Out | Dir::InOut, .. }) => { + self.bufs.insert(buf.clone(), v); + } + (Arg::State { state, offset }, ParamType::State { dir: Dir::Out | Dir::InOut }) => { + self.states.insert((state.clone(), *offset), v); + } + _ => {} + } + } + } + } + + fn start(m: &Manifest) -> Machine { + let mut me = Machine::default(); + me.bufs.insert("x".into(), name_hash("x")); + if m.programs.contains_key("init") { + me.run(m, "init", &mut BTreeMap::new()); + } + me + } +} + +#[test] +fn every_cut_that_verifies_computes_the_whole() { + let (mut cuts, mut refused) = (0, 0); + for seed in 1..=400u64 { + let mut r = Rand(seed.wrapping_mul(0x9e37_79b9_7f4a_7c15)); + let m = verify(random_manifest(&mut r)).unwrap_or_else(|e| panic!("seed {seed}: {e}")); + let n = m.programs["p"].calls.len(); + let a = 1 + r.next(n - 1); + let at = if a + 1 < n && r.next(2) == 0 { vec![a, a + 1 + r.next(n - a - 1)] } else { vec![a] }; + + let mut whole = Machine::start(&m); + whole.run(&m, "p", &mut BTreeMap::new()); + + let stages = match cut(&m, &at, &pp()) { + Ok(s) => s, + Err(e) => { + assert!(e.0.iter().all(|d| d.contains("touched by stages")), "seed {seed} at {at:?}: {e}"); + refused += 1; + continue; + } + }; + cuts += 1; + assert_eq!(stages.len(), at.len() + 1); + let mut mail = BTreeMap::new(); + let mut states = BTreeMap::new(); + let mut y = None; + for (s, st) in stages.iter().enumerate() { + assert_eq!(st.cut.as_ref().map(|c| (c.stage, c.stages)), Some((s as u64, stages.len() as u64))); + let mut me = Machine::start(st); + me.run(st, "p", &mut mail); + for (k, v) in me.states { + assert!(states.insert(k, v).is_none(), "seed {seed}: a state region written by two stages"); + } + y = me.bufs.get("y").copied(); + } + assert_eq!((y, states), (whole.bufs.get("y").copied(), whole.states), "seed {seed} at {at:?}"); + } + assert!(cuts > 200 && refused > 0, "{cuts} cuts, {refused} refused"); +} + +fn qwen3() -> Verified { + verify(Manifest::from_json(QWEN3).unwrap()).unwrap() +} + +/// The buffers stage `s`'s program puts into its outgoing box. +fn carried(st: &Manifest, program: &str) -> Vec { + st.programs[program] + .calls + .iter() + .filter(|c| c.op.starts_with("pp.put.")) + .map(|c| match &c.args[1] { + Arg::Buf { buf, .. } => buf.clone(), + _ => unreachable!(), + }) + .collect() +} + +#[test] +fn qwen3_cuts_at_a_layer_into_stages_that_each_serve() { + // l18.qkv_proj: embed, l0 norm, 18 layers of 12 calls. + let at = 2 + 18 * 12; + assert_eq!(qwen3().programs["prefill"].calls[at].label.as_deref(), Some("l18.qkv_proj")); + let stages = cut(&qwen3(), &[at], &pp()).unwrap_or_else(|e| panic!("{e}")); + let (head, tail) = (&stages[0], &stages[1]); + // The residual stream and the normed input are all that crosses, in + // every program. + for p in ["prefill", "decode", "decode_batch"] { + assert_eq!(carried(head, p), ["residual", "x"]); + } + // Each stage holds its own weights and its own calls of the KV. + assert!( + head.buffers.contains_key("model.embed_tokens.weight") + && !tail.buffers.contains_key("model.embed_tokens.weight") + ); + assert!(tail.buffers.contains_key("lm_head.weight") && !head.buffers.contains_key("lm_head.weight")); + assert!(!head.buffers.contains_key("model.layers.18.self_attn.qkv_proj.weight")); + // The head stages tokens and hands nothing back; the tail is fed + // activations and samples. + let ph = Protocol::check_unsampled(head).unwrap_or_else(|e| panic!("{e}")); + assert!(ph.forwards.iter().all(|f| f.emits.is_none())); + let pt = Protocol::check(tail).unwrap_or_else(|e| panic!("{e}")); + assert_eq!(pt.forward(1, kern_manifest::protocol::Rows::Const(1)).map(|f| f.emits.is_some()), Some(true)); + assert_eq!(head.cut.as_ref().map(|c| &c.id), tail.cut.as_ref().map(|c| &c.id)); + // The same cut is the same bytes. + let again = cut(&qwen3(), &[at], &pp()).unwrap(); + assert_eq!(again[1].to_json(), tail.to_json()); + + let three = cut(&qwen3(), &[2 + 12 * 12, 2 + 24 * 12], &pp()).unwrap_or_else(|e| panic!("{e}")); + assert_eq!(three.len(), 3); + Protocol::check_unsampled(&three[1]).unwrap_or_else(|e| panic!("{e}")); + assert_ne!(three[0].cut.as_ref().map(|c| &c.id), head.cut.as_ref().map(|c| &c.id)); +} + +#[test] +fn a_cut_says_why_it_cannot_be_made() { + let n = qwen3().programs["prefill"].calls.len(); + for at in [vec![], vec![0], vec![5, 5], vec![n]] { + assert!(cut(&qwen3(), &at, &pp()).is_err(), "{at:?}"); + } +} diff --git a/crates/kern-run/examples/pp_chain.rs b/crates/kern-run/examples/pp_chain.rs new file mode 100644 index 00000000..d9f87e43 --- /dev/null +++ b/crates/kern-run/examples/pp_chain.rs @@ -0,0 +1,233 @@ +//! The cut gate: the stage manifests of `kern cut`, run in sequence in +//! one process, compute what the unsplit manifest computes, bit for bit. +//! +//! pp_chain --manifest whole.json --stages s0.json,s1.json --kernels dir +//! --weights dir --gpus 0,1,2 --prompts ids.txt [--chunk 512] +//! [--steps 8] [--capacity 65536] +//! +//! The first GPU runs the whole manifest, the rest one stage each, the +//! stages wired through their mailboxes as a pipeline is (here with +//! in-process handles). Each line of `--prompts` is one prompt's token +//! ids. A prompt but its last token goes through the chunk program in +//! `--chunk` pieces, then `--steps` one-row steps feed back what they +//! pick. After every step the input of the call writing the tokens (the +//! logits) and the tokens themselves must be equal on the whole and on the +//! last stage. All stages are issued before any is waited on: a stage's +//! device waits for its predecessor's mailbox on its own. + +use std::collections::BTreeMap; +use std::path::PathBuf; + +use anyhow::{bail, ensure, Context, Result}; +use kern_manifest::protocol::{Forward, Rows}; +use kern_manifest::types::Arg; +use kern_manifest::{Protocol, Verified}; +use kern_pool::Lease; +use kern_run::{fills, run_once, write_named, Vars}; +use kern_runtime::{Capacity, GroupRank, PeerHandle, Runtime, Topology}; + +struct Opts { + manifest: PathBuf, + stages: Vec, + kernels: PathBuf, + weights: Vec, + gpus: Vec, + prompts: PathBuf, + chunk: usize, + steps: usize, + capacity: u64, +} + +fn opts() -> Result { + let mut o = Opts { + manifest: PathBuf::new(), + stages: Vec::new(), + kernels: PathBuf::new(), + weights: Vec::new(), + gpus: Vec::new(), + prompts: PathBuf::new(), + chunk: 512, + steps: 8, + capacity: 65536, + }; + let mut args = std::env::args().skip(1); + while let Some(a) = args.next() { + let v = args.next().with_context(|| format!("{a}: a value"))?; + match a.as_str() { + "--manifest" => o.manifest = v.into(), + "--stages" => o.stages = v.split(',').map(PathBuf::from).collect(), + "--kernels" => o.kernels = v.into(), + "--weights" => o.weights.push(v), + "--gpus" => o.gpus = v.split(',').map(str::parse).collect::>()?, + "--prompts" => o.prompts = v.into(), + "--chunk" => o.chunk = v.parse()?, + "--steps" => o.steps = v.parse()?, + "--capacity" => o.capacity = v.parse()?, + _ => bail!("unknown argument {a}"), + } + } + ensure!(o.gpus.len() == o.stages.len() + 1, "one GPU for the whole manifest and one per stage"); + Ok(o) +} + +fn load(path: &PathBuf, o: &Opts, gpu: usize, weights: &kern_run::Weights) -> Result<(Runtime, Protocol)> { + let m = Verified::from_json(&std::fs::read_to_string(path)?).with_context(|| format!("{}", path.display()))?; + let topology = m.topology.as_ref().map(|t| Topology { + groups: t.groups.iter().map(|(g, &size)| (g.clone(), GroupRank { index: 0, size })).collect(), + }); + let capacity = Capacity { tokens: Some(o.capacity), seqs: 2 }; + let mut rt = Runtime::load(&m, Some(&o.kernels), gpu, Some(capacity), topology.as_ref())?; + weights.bind(&mut rt)?; + let p = Protocol::check_unsampled(&rt.manifest)?; + Ok((rt, p)) +} + +/// The buffer the call writing `tokens` reads: the logits it picks from. +fn logits(m: &Verified, program: &str, tokens: &str) -> Result { + let call = m.programs[program] + .calls + .iter() + .rev() + .find(|c| c.args.iter().any(|a| matches!(a, Arg::Buf { buf, .. } if buf == tokens))) + .with_context(|| format!("`{program}` writes no `{tokens}`"))?; + call.args + .iter() + .find_map(|a| match a { + Arg::Buf { buf, .. } if buf != tokens => Some(buf.clone()), + _ => None, + }) + .with_context(|| format!("the call writing `{tokens}` reads no buffer")) +} + +/// Stage one call of `ids` at `pos` into `rt` from the values `p` and +/// `lease` give. +fn item(p: &Protocol, lease: &Lease, pos: usize, ids: &[i64]) -> Result<(Vars, BTreeMap>)> { + let (vars, values) = fills(p, std::slice::from_ref(lease), &[pos], ids.len(), ids); + let mut rows: BTreeMap> = values.into_iter().map(|(f, v)| (f.name.clone(), v)).collect(); + for t in &p.page_tables { + let mut v = Vec::new(); + lease.extend_row(&t.name, &mut v)?; + rows.insert(t.name.clone(), v.into_iter().map(i64::from).collect()); + } + Ok((vars, rows)) +} + +fn issue(rt: &mut Runtime, p: &Protocol, program: &str, vars: &Vars, rows: &BTreeMap>) -> Result<()> { + let vars: Vars = + vars.iter().filter(|(k, _)| rt.manifest.vars.contains_key(*k)).map(|(k, v)| (k.clone(), *v)).collect(); + for (name, v) in rows { + write_named(rt, p, name, v, &vars)?; + } + Ok(rt.issue(program, &vars)?) +} + +fn forwards(p: &Protocol) -> Result<(Forward, Forward)> { + Ok(( + p.chunk().cloned().context("no chunk program")?, + p.forward(1, Rows::Const(1)).cloned().context("no one-row program")?, + )) +} + +fn main() -> Result<()> { + tracing_subscriber::fmt().with_env_filter(tracing_subscriber::EnvFilter::from_default_env()).init(); + let o = opts()?; + let weights = kern_run::Weights::parse(&o.weights)?; + let (mut whole, pw) = load(&o.manifest, &o, o.gpus[0], &weights)?; + let mut stages: Vec<(Runtime, Protocol)> = + o.stages.iter().zip(&o.gpus[1..]).map(|(s, &g)| load(s, &o, g, &weights)).collect::>()?; + let exports: Vec> = + stages.iter().map(|(rt, _)| rt.export_handles()).collect::>()?; + for (s, (rt, _)) in stages.iter_mut().enumerate() { + let groups: Vec = rt.manifest.topology.iter().flat_map(|t| t.groups.keys().cloned()).collect(); + for g in groups { + rt.import_peers(&g, &[exports[s].clone(), exports[s + 1].clone()])?; + } + } + run_once(&whole, &pw)?; + for (rt, p) in &stages { + run_once(rt, p)?; + } + + let (wchunk, wstep) = forwards(&pw)?; + let (_, sstep) = forwards(&stages[0].1)?; + let tokens = &pw.fills[wstep.emits.context("the whole manifest's step hands back no tokens")?]; + let (last, lp) = stages.last().unwrap(); + let ltokens = lp.fills[lp + .forwards + .iter() + .find(|f| f.name == sstep.name) + .and_then(|f| f.emits) + .context("the last stage's step hands back no tokens")?] + .clone(); + let wlogits = logits(&whole.manifest, &wstep.name, &tokens.name)?; + let llogits = logits(&last.manifest, &sstep.name, <okens.name)?; + ensure!(wlogits == llogits, "logits are `{wlogits}` whole, `{llogits}` on the last stage"); + + let wlease = whole.lease(o.capacity as usize)?; + let slease = stages[0].0.lease(o.capacity as usize)?; + let prompts: Vec> = std::fs::read_to_string(&o.prompts)? + .lines() + .filter(|l| !l.trim().is_empty()) + .map(|l| l.split_whitespace().map(str::parse).collect::>()) + .collect::>()?; + let mut failed = 0; + for (n, prompt) in prompts.iter().enumerate() { + let mut pos = 0; + let mut calls: Vec<(String, usize, Vec)> = Vec::new(); + while pos < prompt.len() - 1 { + let c = (prompt.len() - 1 - pos).min(o.chunk); + calls.push((wchunk.name.clone(), pos, prompt[pos..pos + c].to_vec())); + pos += c; + } + for (program, pos, ids) in &calls { + let (vars, rows) = item(&pw, &wlease, *pos, ids)?; + issue(&mut whole, &pw, program, &vars, &rows)?; + let (vars, rows) = item(&stages[0].1, &slease, *pos, ids)?; + for (rt, p) in &mut stages { + issue(rt, p, program, &vars, &rows)?; + } + } + let mut tok = *prompt.last().unwrap(); + let mut picked = Vec::new(); + let mut same = true; + for step in 0..o.steps { + let at = prompt.len() - 1 + step; + let (vars, rows) = item(&pw, &wlease, at, &[tok])?; + issue(&mut whole, &pw, &wstep.name, &vars, &rows)?; + let (vars, rows) = item(&stages[0].1, &slease, at, &[tok])?; + for (rt, p) in &mut stages { + issue(rt, p, &sstep.name, &vars, &rows)?; + } + for (rt, _) in &stages { + rt.synchronize()?; + } + let (last, _) = stages.last().unwrap(); + let wt = tokens.decode(&whole.read_output(&tokens.name)?)[0]; + let st = ltokens.decode(&last.read_output(<okens.name)?)[0]; + let wl = whole.read_buffer(&wlogits)?; + let sl = last.read_buffer(&llogits)?; + if wt != st || wl != sl { + let differ = wl.iter().zip(&sl).filter(|(a, b)| a != b).count(); + println!( + "prompt {n} step {step}: token {wt} whole, {st} staged; logits bytes differing {differ}/{}", + wl.len() + ); + same = false; + } + picked.push(wt); + tok = wt; + } + println!( + "prompt {n}: {} tokens, {} chunk calls, {} steps: {} {:?}", + prompt.len(), + calls.len(), + o.steps, + if same { "IDENTICAL" } else { "DIFFERENT" }, + picked + ); + failed += usize::from(!same); + } + ensure!(failed == 0, "{failed} of {} prompts differ", prompts.len()); + println!("PASS: {} prompts bitwise identical", prompts.len()); + Ok(()) +} diff --git a/crates/kern-run/src/bin/kern.rs b/crates/kern-run/src/bin/kern.rs index 771159e5..d98e1ee0 100644 --- a/crates/kern-run/src/bin/kern.rs +++ b/crates/kern-run/src/bin/kern.rs @@ -71,6 +71,23 @@ enum Cmd { /// Manifest JSON to verify (does not read kern.toml) manifest: PathBuf, }, + /// Cut a manifest into pipeline stages before the given call indices + /// of every program a serving loop drives; writes `stage.json` + /// into `--out`, each verified. No GPU + Cut { + /// Manifest JSON to cut (does not read kern.toml) + manifest: PathBuf, + /// Call indices to cut before, increasing + #[arg(long, value_delimiter = ',', required = true)] + at: Vec, + /// The mailbox kernels' cubin (`tools/kernels-src/pp_mailbox.cu`), + /// pinned by its sha256; it goes into each stage's kernels dir + #[arg(long)] + module: PathBuf, + /// Directory the stages are written to + #[arg(long)] + out: PathBuf, + }, } fn main() -> Result<()> { @@ -85,7 +102,7 @@ fn main() -> Result<()> { .init(); let cli = Cli::parse(); let cfg = match &cli.cmd { - Cmd::Verify { .. } => None, + Cmd::Verify { .. } | Cmd::Cut { .. } => None, _ => Config::find(cli.config.as_deref())?, }; match cli.cmd { @@ -118,9 +135,38 @@ fn main() -> Result<()> { } Ok(()) } + Cmd::Cut { manifest, at, module, out } => cut(&manifest, &at, &module, &out), } } +/// `kern cut`: the stages of a manifest cut before `at`, written as +/// `stage.json`. +fn cut(manifest: &Path, at: &[usize], module: &Path, out: &Path) -> Result<()> { + use sha2::Digest; + let json = std::fs::read_to_string(manifest).with_context(|| format!("reading {}", manifest.display()))?; + let m = kern_manifest::Verified::from_json(&json)?; + let bytes = std::fs::read(module).with_context(|| format!("reading {}", module.display()))?; + let pp = kern_manifest::types::Module { + source: module.file_name().context("--module names a file")?.to_string_lossy().into_owned(), + sha256: hex::encode(sha2::Sha256::digest(&bytes)), + }; + let stages = kern_manifest::cut::cut(&m, at, &pp)?; + std::fs::create_dir_all(out).with_context(|| format!("creating {}", out.display()))?; + for (s, st) in stages.iter().enumerate() { + let path = out.join(format!("stage{s}.json")); + std::fs::write(&path, st.to_json()).with_context(|| format!("writing {}", path.display()))?; + let calls: Vec = st.programs.iter().map(|(n, p)| format!("{n} {}", p.calls.len())).collect(); + tracing::info!( + "stage {s}: {} ({} weights; calls: {})", + path.display(), + st.buffers.values().filter(|b| b.kind == kern_manifest::types::BufferKind::Weight).count(), + calls.join(", ") + ); + } + tracing::info!("cut `{}` into {} stages", stages[0].cut.as_ref().map_or("", |c| c.id.as_str()), stages.len()); + Ok(()) +} + /// The target a command runs on: the named one, else the only one, else /// none at all — every command takes its inputs from flags when there is /// no `kern.toml` to look in. diff --git a/crates/kern-run/src/lib.rs b/crates/kern-run/src/lib.rs index 4f63c713..a451b617 100644 --- a/crates/kern-run/src/lib.rs +++ b/crates/kern-run/src/lib.rs @@ -26,7 +26,7 @@ use std::sync::LazyLock; use std::time::Duration; use anyhow::{bail, ensure, Context, Result}; -use kern_manifest::protocol::{Axis, Forward, LineTable, PageTable, Rows}; +use kern_manifest::protocol::{Axis, Filled, Forward, LineTable, PageTable, Rows}; use kern_manifest::types::Fill; use kern_manifest::{Protocol, Verified}; use kern_pool::Lease; @@ -94,8 +94,8 @@ pub fn le_bytes_i32(v: &[i32]) -> Vec { v.iter().flat_map(|x| x.to_le_bytes()).collect() } -/// The var vars of one call. -pub(crate) type Vars = BTreeMap; +/// The var values of one call. +pub type Vars = BTreeMap; /// The safetensors a `--weights` entry stands for: the file itself, or /// every `*.safetensors` under a directory (a checkpoint's shards) in @@ -134,11 +134,60 @@ fn map_file(f: &std::path::Path) -> Result { unsafe { memmap2::Mmap::map(&file) }.with_context(|| format!("mapping weights {}", f.display())) } -/// Stage one call: `rows` rows for each lease, at that sequence's -/// position, `ids` in row order. Returns the call's vars. +/// What one call writes into each fill: `rows` rows for each lease, at +/// that sequence's position, `ids` in row order; and the call's vars. /// -/// Every fill a caller of one rank produces is written; `blocks` is the +/// Every fill a caller of one rank produces is there; `blocks` is the /// tray batch's and only a caller that spans a rank group can fill it. +pub fn fills<'p>( + p: &'p Protocol, + leases: &[Lease], + positions: &[usize], + rows: usize, + ids: &[i64], +) -> (Vars, Vec<(&'p Filled, Vec)>) { + let b = leases.len(); + let longest = positions.iter().map(|&pos| pos + rows).max().unwrap_or(rows); + let vars = p.vars(b as u64, rows as u64, (b * rows) as u64, longest as u64); + let values = p + .fills + .iter() + .filter_map(|f| { + let v: Vec = match (f.fill, f.axis) { + // Each sequence's first token: the anchor a drafting program + // splices its own rows from. + (Fill::Token, Axis::Groups) => ids.chunks(rows).map(|x| x[0]).collect(), + (Fill::Token, _) => ids.to_vec(), + (Fill::Valid, _) => vec![1; b * rows], + (Fill::Position, _) => positions.iter().flat_map(|&pos| (pos..pos + rows).map(|v| v as i64)).collect(), + (Fill::Slot, _) => { + leases.iter().zip(positions).flat_map(|(l, &pos)| l.slots(pos..pos + rows)).collect() + } + (Fill::SeqLen, _) => positions.iter().map(|&pos| (pos + rows) as i64).collect(), + (Fill::CuSeqlens, _) => (0..=b).map(|i| (i * rows) as i64).collect(), + (Fill::SpanAt, _) => vec![0], + (Fill::Blocks | Fill::Tokens | Fill::Count | Fill::Error | Fill::Clock, _) => return None, + }; + Some((f, v)) + }) + .collect(); + (vars, values) +} + +/// Write `values` into input `name` of `rt` the way `p` reads it: a fill +/// in its dtype, a page table as i32. `false` when the manifest has no +/// such input (a pipeline stage that never reads it). +pub fn write_named(rt: &mut Runtime, p: &Protocol, name: &str, values: &[i64], vars: &Vars) -> Result { + let bytes = match (p.fills.iter().find(|f| f.name == name), p.page_tables.iter().any(|t| t.name == name)) { + (Some(f), _) => f.encode(values), + (None, true) => values.iter().flat_map(|&v| (v as i32).to_le_bytes()).collect(), + (None, false) => return Ok(false), + }; + rt.write_input_at(name, &bytes, vars)?; + Ok(true) +} + +/// Stage one call (see [`fills`]). Returns the call's vars. fn stage( rt: &mut Runtime, p: &Protocol, @@ -147,23 +196,8 @@ fn stage( rows: usize, ids: &[i64], ) -> Result { - let b = leases.len(); - let longest = positions.iter().map(|&pos| pos + rows).max().unwrap_or(rows); - let vars = p.vars(b as u64, rows as u64, (b * rows) as u64, longest as u64); - for f in &p.fills { - let v: Vec = match (f.fill, f.axis) { - // Each sequence's first token: the anchor a drafting program - // splices its own rows from. - (Fill::Token, Axis::Groups) => ids.chunks(rows).map(|x| x[0]).collect(), - (Fill::Token, _) => ids.to_vec(), - (Fill::Valid, _) => vec![1; b * rows], - (Fill::Position, _) => positions.iter().flat_map(|&pos| (pos..pos + rows).map(|v| v as i64)).collect(), - (Fill::Slot, _) => leases.iter().zip(positions).flat_map(|(l, &pos)| l.slots(pos..pos + rows)).collect(), - (Fill::SeqLen, _) => positions.iter().map(|&pos| (pos + rows) as i64).collect(), - (Fill::CuSeqlens, _) => (0..=b).map(|i| (i * rows) as i64).collect(), - (Fill::SpanAt, _) => vec![0], - (Fill::Blocks | Fill::Tokens | Fill::Count | Fill::Error, _) => continue, - }; + let (vars, values) = fills(p, leases, positions, rows, ids); + for (f, v) in values { rt.write_input_at(&f.name, &f.encode(&v), &vars)?; } Ok(vars) @@ -197,7 +231,7 @@ fn line_rows(t: &LineTable, leases: &[Lease], cols: usize) -> Result> { /// What the manifest runs once after load (the derived tables a weight /// prep program computes), with every var at 1: a once program takes no /// call shape. -fn run_once(rt: &Runtime, p: &Protocol) -> Result<()> { +pub fn run_once(rt: &Runtime, p: &Protocol) -> Result<()> { let vars: Vars = rt.manifest.vars.keys().map(|v| (v.clone(), 1)).collect(); p.once.iter().try_for_each(|name| rt.run(name, &vars).with_context(|| format!("`{name}`"))) } diff --git a/crates/kern-serve/src/tray.rs b/crates/kern-serve/src/tray.rs index f20e6a74..d80fab00 100644 --- a/crates/kern-serve/src/tray.rs +++ b/crates/kern-serve/src/tray.rs @@ -822,7 +822,7 @@ impl Tray { (Fill::CuSeqlens, _) => (0..=b as i64).map(|i| i * per as i64).collect(), (Fill::SpanAt, _) => vec![span_at], (Fill::Blocks, _) => l.offsets(&groups, q).into_iter().map(|o| (o * per) as i64).collect(), - (Fill::Tokens | Fill::Count | Fill::Error, _) => continue, + (Fill::Tokens | Fill::Count | Fill::Error | Fill::Clock, _) => continue, (Fill::Token, Axis::Fixed(_)) => unreachable!("the protocol checks fill shapes"), }; writes.push((f, v)); diff --git a/docs/manifest.md b/docs/manifest.md index a5d1e616..fb438dab 100644 --- a/docs/manifest.md +++ b/docs/manifest.md @@ -16,6 +16,7 @@ vars. var caller 每次调用供应的标量，有 states. state 不透明持久内存，runtime 按字节供应 buffers. buffer 有类型的张量：input / output / weight / workspace / carry / peer topology.groups. group 多卡 SPMD 的 rank 组：只有名字和大小 +cut cut 流水线的一段：id、第几段、共几段 ``` `call` 调 `op`，`op` 由若干 `launch` 实现，`launch` 起 `module` 里的 @@ -68,6 +69,10 @@ topology.groups. group 多卡 SPMD 的 rank 组：只有名字和 / TRT-LLM / DeepEP intranode 都是这个形状）。 - `{"rank": "ep"}`：call 或 launch 的标量实参来源，本 rank 在组里的 下标，load 时烧成常量，只能接 `i32`/`i64` 参。 +- `cut`（可选，`kern cut` 写的）：`{"id": "6a23…", "stage": 1, "stages": + 2}`——这份 manifest 是一份整 manifest 在某些 call 下标处切开后的第 + `stage` 段，同一刀切出的各段 `id` 相同（整 manifest 的 JSON 与切点的 + hash）。见「流水线切分」。 - `modules`：manifest 的依赖清单（必填）——每个 kernel launch 钉住的代码工件： `source`（本地文件名 `argmax.cubin`，或 registry ref `hf://[@revision]`）+ `sha256`。**身份是 sha256， @@ -385,6 +390,40 @@ kern-serve 只拿 `Protocol` 驱动 runtime，不读 JSON（CI 用 grep 保证 段代码，`--rows` 选每组行数（缺省取 manifest 声明的最宽）。设计与取舍见 [v4-design.md](v4-design.md)。 +## 流水线切分：`kern cut` + +`kern cut --at [,…] --module pp_mailbox.cubin --out +` 把一份单卡 manifest 在 call 下标处切成 S 段（`tools/kernels-src/pp_mailbox.cu` +是邮箱核，与模型无关；`kern_manifest::cut::cut` 是纯函数，不碰 GPU）： + +- 每个 batch program 在同一下标处切开（下标须小于每个 batch program 的 + 长度）；once program 在每段按"本段读到的东西"往回切片。活跃性分析找出 + 切点两侧都活着的 buffer：只许是 workspace、至多一个 var 维（字节数 + `var × 常数`）；output / carry 只许在末段写；一个 (state, offset) 区域 + 只许一段碰（按层切 KV 时各段各持自己那些层）。做不到的切点报错并说明 + 哪条不满足，不会静默切出错的东西。 +- 边界处插一对 send / recv：上游段末尾 `pp.wait_empty → pp.put.… + → pp.post_full`，下游段开头 `pp.wait_full → pp.take.… → + pp.post_empty`；邮箱 `pp..box`（`u8[256 + 最大载荷]`）在下游 GPU 上， + `full` / `empty` 两个计数在头里，上游经 `pp..peer` 远程写。等待在设备 + 上自旋，超时（60 s）`__trap`——进程死于 launch failure，不会挂住。 +- 每段声明 `pp.clock`（`fill: clock`，`i64 [256 × 5]` output）：第 k 项 = + {k+1, 收等开始, 收等结束, 发等开始, 发等结束}（globaltimer ns），首段的 + "收等结束"是本段开始、末段的"发等开始"是本段结束，于是每段 busy = 发等 + 开始 − 收等结束。item 下标由邮箱计数在设备上导出，不需要输入。 +- 非末段声明拓扑组 `pp.`（大小 2，本段是 0 号）。**这里放宽了 SPMD**： + 组两端装的是不同的 manifest，只在 `of` buffer（同名同字节数的 + `pp..box`，上游那份是不用的 workspace + export）上对上—— + `import_peers` 本来就只按名字和字节数匹配，runtime 不用改。 +- 首段之外不需要 `token` fill（`Protocol` 只对 `cut.stage == 0` 要求它）。 + 一个 chunk program 若自己吐 token（hybrid GDN）则拒切：prefill-only 的 + 流水线约定首个 token 从一行的 step program 出。 + +门禁：`crates/kern-manifest/tests/cut.rs` 用随机 manifest（400 个种子） ++ 符号解释器验证"能切的切点，各段依次跑 == 整份跑"，且确实有切点被拒； +GPU 上 `crates/kern-run/examples/pp_chain.rs` 把各段在同一进程里串起来， +logits 与 token 与整份逐位同（serve.md「流水线」）。 + ## Verifier（`kern-manifest`） `verify()` 收集全部错误一次报告（`VerifyErrors`）： @@ -423,7 +462,7 @@ kern-serve 只拿 `Protocol` 驱动 runtime，不读 JSON（CI 用 grep 保证 `export`；`of`/`group` 只许出现在 peer 上；peer 对 op 只读且视为初始 已写；`{"rank": g}` 只接 `i32`/`i64` 参且 g 已声明；**带 extern launch 的 op 不得收到 peer buffer**——runtime 内置（cublasLt）永远不碰 peer - 内存； + 内存；`cut` 有 id、`stages ≥ 2`、`stage < stages`； 12. pack / tensormap / cluster：`bytes` 只许作 launch 参（不进 extern）、 只接 `{"pack"}` 实参且 `size == n`，字段都在 image 内、互不重叠，引用的 接口参 / scratch / var / group 都存在；引用 `out` 接口参的指针字段算 diff --git a/schema/manifest-v5.schema.json b/schema/manifest-v5.schema.json index 10abf59b..28e7d4a7 100644 --- a/schema/manifest-v5.schema.json +++ b/schema/manifest-v5.schema.json @@ -419,6 +419,54 @@ ], "type": "object" }, + "Cut": { + "additionalProperties": false, + "description": "One stage of a manifest cut into a pipeline (`kern cut`): every stage of one cut carries the same `id`; stage `stage` of `stages` receives its predecessor's activations and sends its own to the next.", + "properties": { + "id": { + "description": "Names the cut: a hash of the manifest it was cut from and where.", + "type": "string" + }, + "stage": { + "anyOf": [ + { + "description": "This stage's place in the pipeline, from 0.", + "format": "uint64", + "minimum": 0, + "type": "integer" + }, + { + "description": "Name of a numeric constant declared in constants.", + "minLength": 1, + "type": "string" + } + ], + "description": "This stage's place in the pipeline, from 0." + }, + "stages": { + "anyOf": [ + { + "description": "Stages in the pipeline, at least 2.", + "format": "uint64", + "minimum": 0, + "type": "integer" + }, + { + "description": "Name of a numeric constant declared in constants.", + "minLength": 1, + "type": "string" + } + ], + "description": "Stages in the pipeline, at least 2." + } + }, + "required": [ + "id", + "stage", + "stages" + ], + "type": "object" + }, "DType": { "description": "Element type of a buffer or scratch, e.g. `\"bf16\"`.", "enum": [ @@ -959,6 +1007,11 @@ "const": "error", "description": "Output: a one-word error flag of the collectives, read after every call; nonzero is a failed step.", "type": "string" + }, + { + "const": "clock", + "description": "Output: a pipeline stage's clock, `i64 [ring × 5]`: entry `item % ring` is {item + 1, receive wait begin, end, send wait begin, end} in globaltimer ns, 0 where the stage has no such edge. A cut declares it; the stage reads it when it pleases.", + "type": "string" } ] }, @@ -2154,6 +2207,17 @@ }, "type": "object" }, + "cut": { + "anyOf": [ + { + "$ref": "#/$defs/Cut" + }, + { + "type": "null" + } + ], + "description": "The pipeline stage this manifest is, when it is one piece of a cut manifest, e.g. `{\"id\": \"9f2c…\", \"stage\": 1, \"stages\": 2}`." + }, "model": { "description": "Free-form model label, e.g. `\"qwen3-4b\"`.", "type": "string" @@ -2201,7 +2265,7 @@ "type": "null" } ], - "description": "Rank groups a multi-GPU manifest is SPMD over, e.g. `{\"groups\": {\"ep\": 4}}`; every rank loads the same manifest." + "description": "Rank groups of a multi-GPU manifest, e.g. `{\"groups\": {\"ep\": 4}}`: every rank of an SPMD group loads the same manifest; the two stages a cut's group joins load their own, agreeing only on the peer's `of` buffer." }, "vars": { "additionalProperties": { diff --git a/tools/kernels-src/pp_mailbox.cu b/tools/kernels-src/pp_mailbox.cu new file mode 100644 index 00000000..c3bcb897 --- /dev/null +++ b/tools/kernels-src/pp_mailbox.cu @@ -0,0 +1,150 @@ +// The pipeline edge between two stages of a cut manifest (`kern cut`): one +// mailbox on the downstream GPU, written by the upstream stage through a +// peer pointer. +// +// The box is `u8[256 + data]`. Its header holds two item counters, `full` +// at byte 0 (items the upstream has posted) and `empty` at byte 64 (items +// the downstream has taken out). Both live on the downstream GPU, so the +// downstream spins locally and the upstream reads and writes remote +// memory; neither side keeps a counter of its own. An upstream item: wait +// until `empty == full` (the box is free), put each boundary buffer at its +// offset, post `full + 1`. A downstream item: wait until `full > empty`, +// take each buffer out, post `empty + 1`. Each step is its own launch, so +// stream order puts every copy of an item between its two flags; a post +// fences system-wide before its release store. +// +// A wait gives up after `timeout_ns` of globaltimer and traps: the process +// dies with a launch failure instead of hanging on a dead stage. +// +// Every stage also keeps a clock: `clock` is `i64[RING * 5]`, entry `item % +// RING` = {item + 1, recv wait begin, recv wait end, send wait begin, send +// wait end} in globaltimer ns. The item is read off the counters, never +// passed: a downstream stage's is its box's `empty` before the take, an +// upstream stage's the downstream box's `full` before the post. The waits +// stamp their own spin; a first stage stamps its start as "recv end" and a +// last stage its end as "send begin", so busy = send begin - recv end on +// every stage. One thread writes each field. +// +// kern_pp_wait_empty (in u64 peers[2], out i64 clock[], i64 timeout_ns) +// kern_pp_put (in u64 peers[2], in u8 src[], i64 bytes, i64 offset) +// kern_pp_post_full (in u64 peers[2]) +// kern_pp_wait_full (in u8 box[], out i64 clock[], i64 timeout_ns) +// kern_pp_take (in u8 box[], out u8 dst[], i64 bytes, i64 offset) +// kern_pp_post_empty (inout u8 box[]) +// kern_pp_stamp_first(in u64 peers[2], out i64 clock[]) +// kern_pp_stamp_last (in u8 box[], out i64 clock[]) +// waits, posts, stamps: block [1,1,1], grid [1,1,1]; put / take: any grid. +// peers[1] is the downstream box; the upstream's own box (peers[0]) is unused. + +#define PP_HEADER 256 +#define PP_RING 256 + +typedef unsigned long long u64; + +__device__ __forceinline__ u64 gtimer() { + u64 t; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t)); + return t; +} + +__device__ __forceinline__ u64 ld_acquire_sys(const u64* p) { + u64 v; + asm volatile("ld.acquire.sys.global.u64 %0, [%1];" : "=l"(v) : "l"(p) : "memory"); + return v; +} + +__device__ __forceinline__ void st_release_sys(u64* p, u64 v) { + asm volatile("st.release.sys.global.u64 [%0], %1;" ::"l"(p), "l"(v) : "memory"); +} + +__device__ __forceinline__ u64* full_of(const unsigned char* box) { + return reinterpret_cast(const_cast(box)); +} + +__device__ __forceinline__ u64* empty_of(const unsigned char* box) { + return reinterpret_cast(const_cast(box) + 64); +} + +__device__ __forceinline__ unsigned char* downstream(const u64* peers) { + return reinterpret_cast(peers[1]); +} + +__device__ __forceinline__ long long* entry(long long* clock, u64 item) { + long long* e = clock + (item % PP_RING) * 5; + e[0] = (long long)item + 1; + return e; +} + +__device__ void copy(unsigned char* dst, const unsigned char* src, long long bytes) { + const long long stride = (long long)gridDim.x * blockDim.x; + const long long t = (long long)blockIdx.x * blockDim.x + threadIdx.x; + const long long vecs = bytes / 16; + const uint4* s = reinterpret_cast(src); + uint4* d = reinterpret_cast(dst); + for (long long i = t; i < vecs; i += stride) { + d[i] = s[i]; + } + for (long long i = vecs * 16 + t; i < bytes; i += stride) { + dst[i] = src[i]; + } +} + +extern "C" __global__ void kern_pp_wait_empty(const u64* peers, long long* clock, long long timeout_ns) { + const unsigned char* box = downstream(peers); + const u64 item = ld_acquire_sys(full_of(box)); + long long* e = entry(clock, item); + const u64 t0 = gtimer(); + e[3] = (long long)t0; + while (ld_acquire_sys(empty_of(box)) != item) { + if ((long long)(gtimer() - t0) > timeout_ns) { + printf("kern pp: item %llu waited %lld ns for the downstream stage to take the last one\n", item, + timeout_ns); + __trap(); + } + } + e[4] = (long long)gtimer(); +} + +extern "C" __global__ void kern_pp_put(const u64* peers, const unsigned char* src, long long bytes, + long long offset) { + copy(downstream(peers) + PP_HEADER + offset, src, bytes); +} + +extern "C" __global__ void kern_pp_post_full(const u64* peers) { + const unsigned char* box = downstream(peers); + __threadfence_system(); + st_release_sys(full_of(box), ld_acquire_sys(full_of(box)) + 1); +} + +extern "C" __global__ void kern_pp_wait_full(const unsigned char* box, long long* clock, long long timeout_ns) { + const u64 item = ld_acquire_sys(empty_of(box)); + long long* e = entry(clock, item); + const u64 t0 = gtimer(); + e[1] = (long long)t0; + while (ld_acquire_sys(full_of(box)) <= item) { + if ((long long)(gtimer() - t0) > timeout_ns) { + printf("kern pp: item %llu waited %lld ns for the upstream stage\n", item, timeout_ns); + __trap(); + } + } + __threadfence_system(); + e[2] = (long long)gtimer(); +} + +extern "C" __global__ void kern_pp_take(const unsigned char* box, unsigned char* dst, long long bytes, + long long offset) { + copy(dst, box + PP_HEADER + offset, bytes); +} + +extern "C" __global__ void kern_pp_post_empty(unsigned char* box) { + __threadfence_system(); + st_release_sys(empty_of(box), *empty_of(box) + 1); +} + +extern "C" __global__ void kern_pp_stamp_first(const u64* peers, long long* clock) { + entry(clock, ld_acquire_sys(full_of(downstream(peers))))[2] = (long long)gtimer(); +} + +extern "C" __global__ void kern_pp_stamp_last(const unsigned char* box, long long* clock) { + entry(clock, ld_acquire_sys(empty_of(box)) - 1)[3] = (long long)gtimer(); +} From b2e23399fa4db76164ddcefe2608ed8c3f78f4d6 Mon Sep 17 00:00:00 2001 From: JinYan Su <751080330@qq.com> Date: Wed, 7 Oct 2026 15:33:39 +0000 Subject: [PATCH 02/22] kern-serve: a cut manifest serves prefill as one process per stage Stage 0 (`--pp-listen`) is the head: HTTP, scheduler and the only pool, sized to the fewest pages any stage has. Other stages (`--pp-head`) send Hello {stage, cut, pages, mailbox}, get Table {downstream_mailbox, pages}, then replay the head's Item {program, vars, rows} in order; the last stage answers tokens. Activations go GPU to GPU through the mailbox. Any error exits the process; a closed connection ends the other side. Each stage logs one `pp item` line per item from its clock: rows, busy, recv_wait, send_wait. New dependencies: tokio-util (codec) for length-delimited framing, futures-util (sink/stream) to drive it, bytes for its frames; all three were already in the lock through pegainfer. Gate (tray07, qwen3-4b, 2 stages): first tokens identical to single-GPU kern-serve on 7 prompts; 8193-token prompt TTFT 313 -> 174 ms at conc1. Co-Authored-By: Claude Opus 5.5 (1M context) --- Cargo.lock | 5 + Cargo.toml | 3 + crates/kern-serve/Cargo.toml | 7 +- crates/kern-serve/src/lib.rs | 117 ++++++++--- crates/kern-serve/src/pp/head.rs | 322 ++++++++++++++++++++++++++++++ crates/kern-serve/src/pp/mod.rs | 210 +++++++++++++++++++ crates/kern-serve/src/pp/stage.rs | 143 +++++++++++++ docs/roadmap.md | 7 + docs/serve.md | 42 ++++ 9 files changed, 823 insertions(+), 33 deletions(-) create mode 100644 crates/kern-serve/src/pp/head.rs create mode 100644 crates/kern-serve/src/pp/mod.rs create mode 100644 crates/kern-serve/src/pp/stage.rs diff --git a/Cargo.lock b/Cargo.lock index f7448052..859ca407 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2292,14 +2292,19 @@ name = "kern-serve" version = "0.2.3" dependencies = [ "anyhow", + "bytes", "clap", + "futures-util", + "hex", "kern-manifest", "kern-pool", "kern-run", "kern-runtime", "pegainfer-frontend", + "serde", "serde_json", "tokio", + "tokio-util", "tracing", "tracing-subscriber", "vllm-chat", diff --git a/Cargo.toml b/Cargo.toml index 098d40f7..ab21a18e 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -37,6 +37,7 @@ pegainfer-frontend = { git = "https://github.com/pegainfer-project/pegainfer", r vllm-chat = { git = "https://github.com/vllm-project/vllm.git", rev = "89dbb2644552d6e473a7e97da0ce8f0aa8e32c9d" } anyhow = "1" +bytes = "1" # Recorded references are ragged tables (a row per token position, list # columns for the top-k), which parquet holds and pandas / duckdb read as is. arrow-array = "=60.0.0" @@ -46,6 +47,7 @@ clap = { version = "4", features = ["derive"] } # The CUDA API kern binds is a fact about the binary, not about the build # machine: `cuda-13000` pins it, and no toolkit is needed to build. cudarc = { version = "=0.19.7", features = ["cuda-13000", "f16", "nccl"] } +futures-util = { version = "0.3", default-features = false, features = ["sink"] } half = "=2.7.1" hex = "0.4" libc = "0.2" @@ -58,6 +60,7 @@ sha2 = "0.11" thiserror = "2" tokenizers = "=0.23.2" tokio = { version = "1", features = ["rt-multi-thread", "macros"] } +tokio-util = { version = "0.7", features = ["codec"] } toml = "1.1" tracing = "0.1" tracing-subscriber = { version = "0.3", features = ["env-filter"] } diff --git a/crates/kern-serve/Cargo.toml b/crates/kern-serve/Cargo.toml index 13eeae10..a5f5e8c1 100644 --- a/crates/kern-serve/Cargo.toml +++ b/crates/kern-serve/Cargo.toml @@ -12,9 +12,14 @@ kern-run.workspace = true pegainfer-frontend.workspace = true anyhow.workspace = true +bytes.workspace = true clap.workspace = true +futures-util.workspace = true +hex.workspace = true +serde.workspace = true serde_json.workspace = true -tokio.workspace = true +tokio = { workspace = true, features = ["net", "time"] } +tokio-util.workspace = true tracing.workspace = true tracing-subscriber.workspace = true diff --git a/crates/kern-serve/src/lib.rs b/crates/kern-serve/src/lib.rs index 211ee83f..3b3218ec 100644 --- a/crates/kern-serve/src/lib.rs +++ b/crates/kern-serve/src/lib.rs @@ -11,6 +11,7 @@ #![deny(unsafe_code)] pub mod logline; +mod pp; mod scheduler; mod tray; @@ -22,7 +23,7 @@ use clap::Args; use kern_manifest::Verified; use kern_runtime::{Capacity, Runtime, Topology}; use pegainfer_frontend::engine::{ - drive, scheduler_pair, Engine, EngineInfo, KvCapacity, LaunchedEngine, LiveScheduler, + drive, scheduler_pair, Engine, EngineInfo, KvCapacity, LaunchedEngine, LiveScheduler, Scheduler, }; use pegainfer_frontend::vllm; use tracing::info; @@ -90,6 +91,15 @@ pub struct ServeOpts { /// (0: off) #[arg(long, default_value_t = 0.0)] pub host_gib: f64, + + /// Serve stage 0 of a cut manifest (`kern cut`): HTTP, scheduler and + /// pool, waiting at this host:port for every other stage to join + #[arg(long, conflicts_with = "pp_head")] + pub pp_listen: Option, + + /// Serve a later stage of a cut manifest, joining the head at host:port + #[arg(long)] + pub pp_head: Option, } pub fn serve(o: ServeOpts, art: Artifacts) -> Result<()> { @@ -102,6 +112,26 @@ pub fn serve(o: ServeOpts, art: Artifacts) -> Result<()> { "loading" ); let gpus = if o.gpus.is_empty() { vec![0] } else { o.gpus.clone() }; + let manifest_json = std::fs::read_to_string(&art.manifest) + .with_context(|| format!("reading manifest {}", art.manifest.display()))?; + let manifest = + Verified::from_json(&manifest_json).with_context(|| format!("manifest {}", art.manifest.display()))?; + anyhow::ensure!( + manifest.cut.is_none() || (gpus.len() == 1 && (o.pp_listen.is_some() || o.pp_head.is_some())), + "a pipeline stage runs on one GPU, with `--pp-listen` (stage 0) or `--pp-head` (the others)" + ); + if let Some(head) = &o.pp_head { + let stage = pp::stage::Stage { + head: head.clone(), + gpu: gpus[0], + kernels: art.kernels.clone(), + capacity: o.capacity, + max_seqs: o.max_seqs, + eager: o.eager, + }; + return pp::stage::run(manifest, &weights, stage); + } + let model_path = weights.dirs().remove(0); let mut stop_tokens: Vec = kern_run::eos_ids(&model_path).into_iter().map(|x| x as u32).collect(); stop_tokens.extend(&o.stop_tokens); @@ -113,52 +143,75 @@ pub fn serve(o: ServeOpts, art: Artifacts) -> Result<()> { model_path.display() ); info!(ids = ?stop_tokens, "stop tokens"); - - let manifest_json = std::fs::read_to_string(&art.manifest) - .with_context(|| format!("reading manifest {}", art.manifest.display()))?; - let manifest = - Verified::from_json(&manifest_json).with_context(|| format!("manifest {}", art.manifest.display()))?; let served_name = o.served_model_name.clone().unwrap_or_else(|| manifest.model.clone()); + + if let Some(listen) = o.pp_listen.clone() { + let head = pp::head::Head { + listen, + gpu: gpus[0], + kernels: art.kernels.clone(), + capacity: o.capacity, + chunk: o.chunk, + max_seqs: o.max_seqs, + eager: o.eager, + stop_tokens, + }; + let load = move || -> Result<(pp::head::PpHead, scheduler::Facts)> { + let h = pp::head::load(manifest, &weights, head)?; + let facts = h.facts(); + Ok((h, facts)) + }; + return launch(load, served_name, model_path, o.port); + } + // Every sequence of a tray batch group holds a token slot on each of // its `t` ranks, and each rank its pad. let t = manifest.group_size("tp").unwrap_or(1) as usize; let capacity = Capacity { tokens: o.capacity, seqs: ((o.max_seqs + 1) * t) as u64 }; - - // The scheduler thread owns the tray for its whole life: load there, - // report readiness, then drive. - let (handle, backend) = scheduler_pair(); - let (ready_tx, ready_rx) = std::sync::mpsc::channel::>(); let host_bytes = (o.host_gib * (1u64 << 30) as f64) as u64; let policy = Policy { chunk: o.chunk.map(|c| c as usize), max_seqs: o.max_seqs, stop_tokens, rows: o.rows, host_bytes }; + let (kernels, eager) = (art.kernels.clone(), o.eager); + let load = move || -> Result<(KernScheduler, scheduler::Facts)> { + let t0 = Instant::now(); + let bind = |rt: &mut Runtime, _: &Topology| weights.bind(rt); + let tray = Tray::load(&manifest, kernels.as_deref(), &gpus, capacity, &bind, host_bytes, eager)?; + info!(model = %tray.manifest().model, gpus = ?gpus, load_s = logline::secs(t0.elapsed()), "tray loaded"); + let sched = KernScheduler::new(tray, policy)?; + let facts = sched.facts(); + Ok((sched, facts)) + }; + launch(load, served_name, model_path, o.port) +} + +/// Run the scheduler `load` builds on a thread of its own, which owns it +/// (and its GPUs) for its whole life, and serve HTTP once it is ready. +fn launch(load: L, served_name: String, model_path: PathBuf, port: u16) -> Result<()> +where + S: Scheduler + 'static, + L: FnOnce() -> Result<(S, scheduler::Facts)> + Send + 'static, +{ + let (handle, backend) = scheduler_pair(); + let (ready_tx, ready_rx) = std::sync::mpsc::channel::>(); let join = std::thread::Builder::new() .name("kern-scheduler".into()) - .spawn(move || { - let load = || -> Result { - let t0 = Instant::now(); - let bind = |rt: &mut Runtime, _: &Topology| weights.bind(rt); - let tray = Tray::load(&manifest, art.kernels.as_deref(), &gpus, capacity, &bind, host_bytes, o.eager)?; - info!(model = %tray.manifest().model, gpus = ?gpus, load_s = logline::secs(t0.elapsed()), "tray loaded"); - KernScheduler::new(tray, policy) - }; - match load() { - Ok(sched) => { - let _ = ready_tx.send(Ok(sched.facts())); - // A scheduler that panicked would leave the port open - // and every request hanging: the process goes with it. - if std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| drive(sched, backend))).is_err() { - std::process::exit(101); - } - } - Err(e) => { - let _ = ready_tx.send(Err(e)); + .spawn(move || match load() { + Ok((sched, facts)) => { + let _ = ready_tx.send(Ok(facts)); + // A scheduler that panicked would leave the port open + // and every request hanging: the process goes with it. + if std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| drive(sched, backend))).is_err() { + std::process::exit(101); } } + Err(e) => { + let _ = ready_tx.send(Err(e)); + } }) .context("spawning the scheduler thread")?; // The port opens once the engine is up: `serving` follows readiness. - let (model, model_dir, port) = (served_name.clone(), model_path.clone(), o.port); + let (model, model_dir) = (served_name.clone(), model_path.clone()); let engine = async move { let facts = tokio::task::spawn_blocking(move || ready_rx.recv()) .await @@ -179,6 +232,6 @@ pub fn serve(o: ServeOpts, art: Artifacts) -> Result<()> { rt.block_on(async move { // Needs the runtime: it spawns the signal listener. let shutdown = vllm::shutdown_token_from_ctrl_c(); - vllm::serve_with_engine_count(engine, &model_path, vec![served_name], o.port, None, 1, shutdown).await + vllm::serve_with_engine_count(engine, &model_path, vec![served_name], port, None, 1, shutdown).await }) } diff --git a/crates/kern-serve/src/pp/head.rs b/crates/kern-serve/src/pp/head.rs new file mode 100644 index 00000000..0af0399c --- /dev/null +++ b/crates/kern-serve/src/pp/head.rs @@ -0,0 +1,322 @@ +//! The head: stage 0 of a pipeline, the scheduler every stage follows. +//! +//! A request leases its prompt's pages from the head's pool, the one pool +//! of the pipeline, and goes out as items: the prompt but its last token +//! in chunks through the chunk program, then the last token through the +//! one-row step, whose tokens the last stage hands back. Each item is sent +//! to every stage before the head stages and issues it itself, so a stage +//! waiting on its device for an item always has it on the way. The first +//! token finishes the request. Nothing waits on the head's device but the +//! clock read, done when nothing is in flight or the clock ring is half +//! full. + +use std::collections::{BTreeMap, VecDeque}; +use std::sync::mpsc; + +use anyhow::{bail, ensure, Context, Result}; +use kern_manifest::protocol::{Forward, Rows}; +use kern_manifest::{Protocol, Verified}; +use kern_pool::{Denied, Lease}; +use kern_run::Weights; +use kern_runtime::{Capacity, Error, Runtime}; +use pegainfer_frontend::engine::{ + FinishReason, QueuedRequest, RejectReason, RequestId, RequestLedger, Scheduler, SchedulerMetrics, +}; +use tokio::net::TcpListener; +use tracing::{error, info}; + +use super::{connect_downstream, recv, send, topology, wire, Clock, Handles, Hello, Item, Reader, Table, Writer}; +use crate::scheduler::Facts; + +/// How the head loads and what it may run at once. +pub struct Head { + pub listen: String, + pub gpu: usize, + pub kernels: Option, + pub capacity: Option, + pub chunk: Option, + pub max_seqs: usize, + pub eager: bool, + pub stop_tokens: Vec, +} + +/// A request whose items are out and whose token is not back yet. +struct Inflight { + id: RequestId, + _lease: Lease, +} + +pub struct PpHead { + rt: Runtime, + p: Protocol, + prefill: Forward, + step: Forward, + chunk: usize, + max_seqs: usize, + stop_tokens: Vec, + io: tokio::runtime::Runtime, + stages: Vec, + replies: mpsc::Receiver>, + waiting: VecDeque, + inflight: VecDeque, + clock: Clock, +} + +/// Wait for every other stage, load stage 0 at the pages they all have, +/// map stage 1's mailbox and answer each stage with its table. +pub fn load(m: Verified, weights: &Weights, h: Head) -> Result { + let cut = m.cut.clone().context("the head serves a cut manifest")?; + ensure!(cut.stage == 0, "`--pp-listen` serves stage 0; this manifest is stage {} of {}", cut.stage, cut.stages); + let io = super::io()?; + let listener = io.block_on(TcpListener::bind(&h.listen)).with_context(|| format!("binding {}", h.listen))?; + info!(listen = %h.listen, stages = cut.stages, cut = %cut.id, "waiting for the stages"); + let mut hellos: BTreeMap = BTreeMap::new(); + while (hellos.len() as u64) < cut.stages - 1 { + let (s, addr) = io.block_on(listener.accept())?; + let (mut r, w) = wire(s); + let hello: Hello = io.block_on(recv(&mut r))?.with_context(|| format!("{addr} closed before its hello"))?; + ensure!( + hello.cut == cut.id, + "{addr}: stage {} is of cut `{}`, the head of `{}`", + hello.stage, + hello.cut, + cut.id + ); + ensure!( + (1..cut.stages).contains(&hello.stage) && !hellos.contains_key(&hello.stage), + "{addr}: stage {} is not a stage still missing of 1..{}", + hello.stage, + cut.stages + ); + info!(stage = hello.stage, %addr, pages = hello.pages, "stage joined"); + hellos.insert(hello.stage, (hello, r, w)); + } + let unit = kern_pool::page_unit(&m); + let fewest = hellos.values().map(|(x, ..)| x.pages).min().expect("at least one other stage"); + let tokens = h.capacity.map_or(fewest * unit, |c| c.min(fewest * unit)); + let seqs = (h.max_seqs + 1) as u64; + let mut rt = Runtime::load( + &m, + h.kernels.as_deref(), + h.gpu, + Some(Capacity { tokens: Some(tokens), seqs }), + topology(&m).as_ref(), + )?; + rt.set_eager(h.eager); + weights.bind(&mut rt)?; + connect_downstream(&mut rt, &hellos[&1].0.mailbox)?; + let p = Protocol::check_unsampled(&rt.manifest)?; + kern_run::run_once(&rt, &p)?; + ensure!(p.line_tables.is_empty(), "a pipeline serves paged state only; this manifest has per-sequence states"); + let pages = rt.pages_total() as u64; + + let downstream: Vec> = + (1..cut.stages).map(|s| hellos.get(&(s + 1)).map(|(x, ..)| x.mailbox.clone())).collect(); + let (tx, replies) = mpsc::channel(); + let mut stages = Vec::new(); + for ((s, (_, r, mut w)), down) in hellos.into_iter().zip(downstream) { + io.block_on(send(&mut w, &Table { downstream_mailbox: down, pages }))?; + let tx = (s == cut.stages - 1).then(|| tx.clone()); + io.spawn(listen(s, r, tx)); + stages.push(w); + } + + let prefill = + p.chunk().cloned().context("a pipeline prefills through a program taking one sequence's rows as fed")?; + let step = p.forward(1, Rows::Const(1)).cloned().context("no program takes one sequence of one row")?; + let max = p.rows.max as usize; + let chunk = match h.chunk { + None => max, + Some(c) if (1..=max as u64).contains(&c) => c as usize, + Some(c) => bail!("--chunk {c}: the manifest's `tokens` bound is {max}; a chunk can only be smaller"), + }; + let clock = Clock::new(&p)?; + info!(pages, page = rt.page(), chunk, prefill = %prefill.name, step = %step.name, "pipeline head ready"); + Ok(PpHead { + rt, + p, + prefill, + step, + chunk, + max_seqs: h.max_seqs, + stop_tokens: h.stop_tokens, + io, + stages, + replies, + waiting: VecDeque::new(), + inflight: VecDeque::new(), + clock, + }) +} + +/// Read one stage's connection to its end: the last stage's tokens go to +/// the scheduler; anything closing ends the head. +async fn listen(stage: u64, mut r: Reader, tokens: Option>>) { + let why = loop { + match (recv::>(&mut r).await, &tokens) { + (Ok(Some(t)), Some(tx)) => { + if tx.send(t).is_err() { + return; + } + } + (Ok(Some(_)), None) => break "a stage other than the last handed tokens back".to_string(), + (Ok(None), _) => break "stage closed its connection; the pipeline is gone".to_string(), + (Err(e), _) => break format!("stage connection: {e:#}"), + } + }; + error!(stage, "{why}"); + std::process::exit(1); +} + +impl PpHead { + pub fn facts(&self) -> Facts { + Facts { + total_blocks: self.rt.pages_total(), + block_size: self.rt.page() as usize, + max_request_tokens: self.rt.max_seq_tokens(), + } + } + + /// Lease each waiting prompt in order and send its items out; one that + /// cannot be seated stops the scan. + fn admit(&mut self, ledger: &mut RequestLedger) -> Result<()> { + while let Some(q) = self.waiting.front() { + let id = q.id; + if ledger.is_aborted(id) { + ledger.retire(id); + self.waiting.pop_front(); + continue; + } + if self.inflight.len() >= self.max_seqs { + break; + } + let (prompt, max_tokens) = (q.request.prompt_tokens.len(), q.request.max_tokens); + let limit = self.rt.max_seq_tokens(); + let leased = + if prompt == 0 { Err(Error::Denied(Denied::ExceedsRow { limit })) } else { self.rt.lease(prompt) }; + let reject = match leased { + Ok(lease) => { + let q = self.waiting.pop_front().unwrap(); + ledger.admit(id); + ledger.set_cached_tokens(id, 0); + self.send_prompt(&q, &lease)?; + self.inflight.push_back(Inflight { id, _lease: lease }); + continue; + } + Err(Error::Denied(Denied::Busy | Denied::Remapping)) => break, + Err(Error::Denied(Denied::ExceedsRow { limit })) => { + Some(RejectReason::ContextLength { prompt_tokens: prompt, max_tokens, limit }) + } + Err(Error::Denied(Denied::ExceedsPool)) => { + Some(RejectReason::KvBudget { prompt_tokens: prompt, worst_case_tokens: prompt }) + } + Err(e) => return Err(e.into()), + }; + if let Some(r) = reject { + ledger.reject(id, r); + self.waiting.pop_front(); + } + } + Ok(()) + } + + /// The prompt as items: all but its last token in chunks, the last + /// through the step, whose token comes back. + fn send_prompt(&mut self, q: &QueuedRequest, lease: &Lease) -> Result<()> { + let ids: Vec = q.request.prompt_tokens.iter().map(|&t| t as i64).collect(); + let n = ids.len(); + let (prefill, step) = (self.prefill.clone(), self.step.clone()); + let mut pos = 0; + while pos < n - 1 { + let c = (n - 1 - pos).min(self.chunk); + self.send_item(&prefill, lease, pos, &ids[pos..pos + c])?; + pos += c; + } + self.send_item(&step, lease, n - 1, &ids[n - 1..]) + } + + fn send_item(&mut self, f: &Forward, lease: &Lease, pos: usize, ids: &[i64]) -> Result<()> { + let (vars, values) = kern_run::fills(&self.p, std::slice::from_ref(lease), &[pos], ids.len(), ids); + let mut rows: BTreeMap> = values.into_iter().map(|(f, v)| (f.name.clone(), v)).collect(); + for t in &self.p.page_tables { + let mut v = Vec::with_capacity(t.width); + lease.extend_row(&t.name, &mut v)?; + rows.insert(t.name.clone(), v.into_iter().map(i64::from).collect()); + } + let item = Item { program: f.name.clone(), vars, rows }; + for w in &mut self.stages { + self.io.block_on(send(w, &item))?; + } + for (name, v) in &item.rows { + kern_run::write_named(&mut self.rt, &self.p, name, v, &item.vars)?; + } + self.rt.issue(&item.program, &item.vars).with_context(|| format!("`{}`", item.program))?; + self.clock.issued(ids.len() as u64); + Ok(()) + } + + /// Admit, collect, and read the clock whenever this step issued + /// nothing: the read waits for the head's device, which is also how a + /// mailbox wait that gave up on a dead stage surfaces here. + fn advance(&mut self, ledger: &mut RequestLedger) -> Result<()> { + let issued = self.clock.pending(); + self.admit(ledger)?; + self.collect(ledger)?; + if self.clock.pending() == issued || self.clock.half_full() { + self.clock.flush(&self.rt, 0)?; + } + Ok(()) + } + + /// Every token back since the last step finishes its request, oldest + /// first: the last stage answers in issue order. + fn collect(&mut self, ledger: &mut RequestLedger) -> Result<()> { + loop { + let t = match self.replies.try_recv() { + Ok(t) => t, + Err(mpsc::TryRecvError::Empty) => return Ok(()), + Err(mpsc::TryRecvError::Disconnected) => bail!("the last stage's connection is gone"), + }; + let Some(r) = self.inflight.pop_front() else { bail!("a token came back for no request") }; + let Some(&tok) = t.first() else { bail!("the last stage handed back no token") }; + if ledger.is_aborted(r.id) { + ledger.retire(r.id); + continue; + } + let tok = tok as u32; + ledger.push_tokens(r.id, &[tok], &[]); + let why = if self.stop_tokens.contains(&tok) { FinishReason::Stop } else { FinishReason::Length }; + ledger.finish(r.id, why); + } + } +} + +// Built on the scheduler thread and used there only, like the `Tray`. +#[allow(unsafe_code)] +unsafe impl Send for PpHead {} + +impl Scheduler for PpHead { + fn submit(&mut self, request: QueuedRequest) { + self.waiting.push_back(request); + } + + /// A failed step ends the process: the stages see the connection close + /// and follow, and no half-fed pipeline lingers behind an open port. + fn step(&mut self, ledger: &mut RequestLedger) -> Result<()> { + if let Err(e) = self.advance(ledger) { + error!("pipeline head: {e:#}"); + std::process::exit(1); + } + Ok(()) + } + + fn metrics(&self) -> SchedulerMetrics { + SchedulerMetrics { + kv_used_blocks: self.rt.pages_used() as u64, + kv_total_blocks: self.rt.pages_total() as u64, + num_running_reqs: self.inflight.len() as u64, + num_waiting_reqs: self.waiting.len() as u64, + spec_decode: None, + } + } +} diff --git a/crates/kern-serve/src/pp/mod.rs b/crates/kern-serve/src/pp/mod.rs new file mode 100644 index 00000000..0c0eaf4c --- /dev/null +++ b/crates/kern-serve/src/pp/mod.rs @@ -0,0 +1,210 @@ +//! Pipeline serving: one process per stage of a cut manifest (`kern cut`), +//! prefill only — a request gets its first token and finishes. +//! +//! Stage 0 is the head: the HTTP front end, the scheduler and the one +//! [`kern_pool::Pool`] that decides every lease. The other stages are +//! executors that hold no accounting of their own. Each stage's KV pages +//! hold that stage's layers under the head's page ids, which is why the +//! head's pool is sized to the smallest stage: the stages say how many +//! pages they have in their [`Hello`], the head loads with the minimum. +//! +//! One TCP connection per downstream stage, framed by `tokio-util`'s +//! length-delimited codec, each frame one JSON message. A stage connects +//! and sends its [`Hello`]: its place, the cut it belongs to, its pages and +//! the export handles of its allocations (its mailbox among them). Once +//! every stage has said hello the head loads, maps the first stage's +//! mailbox and answers each stage with a [`Table`]: the next stage's +//! handles, to map its own outgoing peer from, and the pages. From then on +//! the connection carries the log: the head sends every stage every +//! [`Item`] in issue order — the program, its vars and the values of each +//! input it staged — and the last stage answers each item whose program +//! hands tokens back with the tokens. TCP keeps the order; there is no +//! sequence number, no stop message and no fault message. A stage that +//! fails exits; a closed connection ends the process on the other side. +//! +//! The activations never touch the connection: they go GPU to GPU through +//! the mailbox the cut wired in, and each stage's device waits for its +//! predecessor on its own. A stage only enqueues; it synchronizes when it +//! has nothing to enqueue (to read its clock) or when it must read tokens. + +pub mod head; +pub mod stage; + +use std::collections::{BTreeMap, VecDeque}; + +use anyhow::{bail, Context, Result}; +use bytes::Bytes; +use futures_util::{SinkExt, StreamExt}; +use kern_manifest::protocol::Axis; +use kern_manifest::types::Fill; +use kern_manifest::Protocol; +use kern_runtime::{GroupRank, PeerHandle, Runtime, Topology}; +use serde::{Deserialize, Serialize}; +use tokio::net::tcp::{OwnedReadHalf, OwnedWriteHalf}; +use tokio::net::TcpStream; +use tokio_util::codec::{FramedRead, FramedWrite, LengthDelimitedCodec}; + +/// A stage introducing itself to the head. +#[derive(Debug, Serialize, Deserialize)] +pub struct Hello { + pub stage: u64, + pub cut: String, + pub pages: u64, + /// Every allocation the stage exports, by name, as hex fabric handles. + pub mailbox: Handles, +} + +/// The head's answer: the next stage's exports (none for the last stage) +/// and the pages every stage addresses. +#[derive(Debug, Serialize, Deserialize)] +pub struct Table { + pub downstream_mailbox: Option, + pub pages: u64, +} + +/// One call, as the head staged it. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct Item { + pub program: String, + pub vars: BTreeMap, + /// Each input's values by name: the fills and the page tables. + pub rows: BTreeMap>, +} + +pub type Handles = BTreeMap; + +/// This runtime's exports as wire handles; a stage's box must be mappable +/// from another process, so a local handle is an error. +fn handles(rt: &Runtime) -> Result { + rt.export_handles()? + .into_iter() + .map(|(name, h)| { + let b = h.to_bytes().with_context(|| { + format!("`{name}` has no fabric handle; a pipeline stage maps its neighbour across processes") + })?; + Ok((name, hex::encode(b))) + }) + .collect() +} + +fn peer_handles(h: &Handles) -> Result> { + h.iter() + .map(|(name, x)| { + let b = hex::decode(x).with_context(|| format!("handle of `{name}`"))?; + Ok((name.clone(), PeerHandle::from_bytes(&b).with_context(|| format!("handle of `{name}`"))?)) + }) + .collect() +} + +/// A stage's place in the group of its outgoing edge: always the +/// upstream member, rank 0 of 2. +fn topology(rt_manifest: &kern_manifest::Manifest) -> Option { + rt_manifest.topology.as_ref().map(|t| Topology { + groups: t.groups.iter().map(|(g, &size)| (g.clone(), GroupRank { index: 0, size })).collect(), + }) +} + +/// Map the next stage's box into this stage's outgoing peer buffer. +fn connect_downstream(rt: &mut Runtime, downstream: &Handles) -> Result<()> { + let groups: Vec = rt.manifest.topology.iter().flat_map(|t| t.groups.keys().cloned()).collect(); + let own = rt.export_handles()?; + let theirs = peer_handles(downstream)?; + for g in groups { + rt.import_peers(&g, &[own.clone(), theirs.clone()]) + .with_context(|| format!("mapping the next stage over `{g}`"))?; + } + let pending = rt.pending_peers(); + if !pending.is_empty() { + bail!("peer buffers {pending:?} still unfilled after the next stage was mapped"); + } + Ok(()) +} + +type Reader = FramedRead; +type Writer = FramedWrite; + +fn wire(s: TcpStream) -> (Reader, Writer) { + let _ = s.set_nodelay(true); + let (r, w) = s.into_split(); + (FramedRead::new(r, LengthDelimitedCodec::new()), FramedWrite::new(w, LengthDelimitedCodec::new())) +} + +async fn send(w: &mut Writer, msg: &T) -> Result<()> { + w.send(Bytes::from(serde_json::to_vec(msg)?)).await.context("sending") +} + +/// The next message; `None` when the other side closed. +async fn recv Deserialize<'de>>(r: &mut Reader) -> Result> { + match r.next().await { + None => Ok(None), + Some(frame) => Ok(Some(serde_json::from_slice(&frame.context("receiving")?)?)), + } +} + +/// The IO runtime a process's connections live on, one thread beside the +/// one driving the GPU. +fn io() -> Result { + Ok(tokio::runtime::Builder::new_multi_thread().worker_threads(1).enable_all().build()?) +} + +/// What a stage's device spent per item, read off its clock output: one +/// log line per item, busy and waiting on either side. +struct Clock { + name: String, + ring: u64, + logged: u64, + /// Rows of each item issued and not yet logged, oldest first. + rows: VecDeque, +} + +impl Clock { + fn new(p: &Protocol) -> Result { + let f = p.any(Fill::Clock).context("a cut stage declares a `clock` fill")?; + let Axis::Fixed(n) = f.axis else { bail!("`{}`: a clock is a fixed ring", f.name) }; + Ok(Clock { name: f.name.clone(), ring: n / 5, logged: 0, rows: VecDeque::new() }) + } + + fn issued(&mut self, rows: u64) { + self.rows.push_back(rows); + } + + /// Items issued and not yet logged. + fn pending(&self) -> usize { + self.rows.len() + } + + /// Whether the oldest unread entry is half a ring from being overwritten. + fn half_full(&self) -> bool { + self.rows.len() as u64 >= self.ring / 2 + } + + /// Read the clock (waiting for the stage's stream) and log every item + /// issued since the last read. + fn flush(&mut self, rt: &Runtime, stage: u64) -> Result<()> { + if self.rows.is_empty() { + return Ok(()); + } + let raw = rt.read_output(&self.name)?; + let e: Vec = raw.as_chunks::<8>().0.iter().map(|c| i64::from_le_bytes(*c)).collect(); + let span = |a: i64, b: i64| (a != 0 && b != 0).then_some(b - a); + for rows in self.rows.drain(..) { + let k = self.logged; + self.logged += 1; + let x = &e[(k % self.ring) as usize * 5..][..5]; + if x[0] != k as i64 + 1 { + tracing::warn!(stage, item = k, "pp clock entry overwritten before it was read"); + continue; + } + tracing::info!( + stage, + item = k, + rows, + busy_ns = span(x[2], x[3]), + recv_wait_ns = span(x[1], x[2]), + send_wait_ns = span(x[3], x[4]), + "pp item" + ); + } + Ok(()) + } +} diff --git a/crates/kern-serve/src/pp/stage.rs b/crates/kern-serve/src/pp/stage.rs new file mode 100644 index 00000000..64ddb35f --- /dev/null +++ b/crates/kern-serve/src/pp/stage.rs @@ -0,0 +1,143 @@ +//! A downstream stage: loads its slice of the model, says hello to the +//! head, maps the next stage's mailbox from the table it gets back, then +//! replays the head's items in order. It keeps no accounting: the page +//! ids, positions and slots of every item are the head's. The last stage +//! hands back the tokens of every item whose program emits them. + +use std::sync::mpsc; +use std::time::{Duration, Instant}; + +use anyhow::{bail, ensure, Context, Result}; +use kern_manifest::{Protocol, Verified}; +use kern_run::{Vars, Weights}; +use kern_runtime::{Capacity, Runtime}; +use tokio::net::TcpStream; +use tracing::{error, info}; + +use super::{connect_downstream, handles, recv, send, topology, wire, Clock, Hello, Item, Table}; + +/// How a stage loads and where its head listens. +pub struct Stage { + pub head: String, + pub gpu: usize, + pub kernels: Option, + pub capacity: Option, + pub max_seqs: usize, + pub eager: bool, +} + +/// How long a stage keeps trying to reach a head that is not up yet. +const CONNECT: Duration = Duration::from_secs(600); + +pub fn run(m: Verified, weights: &Weights, o: Stage) -> Result<()> { + let cut = m.cut.clone().context("`--pp-head` serves a stage of a cut manifest")?; + ensure!(cut.stage > 0, "stage 0 is the head; it takes `--pp-listen`"); + let last = cut.stage + 1 == cut.stages; + let t0 = Instant::now(); + let capacity = Capacity { tokens: o.capacity, seqs: (o.max_seqs + 1) as u64 }; + let mut rt = Runtime::load(&m, o.kernels.as_deref(), o.gpu, Some(capacity), topology(&m).as_ref())?; + rt.set_eager(o.eager); + weights.bind(&mut rt)?; + let p = Protocol::check_unsampled(&rt.manifest)?; + info!( + stage = cut.stage, + gpu = o.gpu, + pages = rt.pages_total(), + load_s = t0.elapsed().as_secs_f64(), + "stage loaded" + ); + + let io = super::io()?; + let stream = io.block_on(connect(&o.head))?; + let (mut r, mut w) = wire(stream); + let hello = Hello { stage: cut.stage, cut: cut.id.clone(), pages: rt.pages_total() as u64, mailbox: handles(&rt)? }; + io.block_on(send(&mut w, &hello))?; + let table: Table = io.block_on(recv(&mut r))?.context("the head closed before its table")?; + ensure!( + table.pages <= rt.pages_total() as u64, + "the head addresses {} pages, this stage has {}", + table.pages, + rt.pages_total() + ); + match (&table.downstream_mailbox, last) { + (Some(d), false) => connect_downstream(&mut rt, d)?, + (None, true) => {} + (d, _) => bail!( + "stage {} of {}: the head sent {} next stage's mailbox", + cut.stage, + cut.stages, + if d.is_some() { "a" } else { "no" } + ), + } + kern_run::run_once(&rt, &p)?; + info!(stage = cut.stage, head = %o.head, last, "stage serving"); + + let (tx, items) = mpsc::channel::(); + io.spawn(async move { + loop { + match recv::(&mut r).await { + Ok(Some(item)) => { + if tx.send(item).is_err() { + return; + } + } + Ok(None) => { + info!("the head closed its connection"); + return; + } + Err(e) => { + error!("head connection: {e:#}"); + return; + } + } + } + }); + + let mut clock = Clock::new(&p)?; + loop { + let item = match items.try_recv() { + Ok(item) => item, + Err(mpsc::TryRecvError::Empty) => { + clock.flush(&rt, cut.stage)?; + match items.recv() { + Ok(item) => item, + Err(_) => break, + } + } + Err(mpsc::TryRecvError::Disconnected) => break, + }; + let vars: Vars = item.vars.into_iter().filter(|(k, _)| rt.manifest.vars.contains_key(k)).collect(); + for (name, v) in &item.rows { + kern_run::write_named(&mut rt, &p, name, v, &vars)?; + } + rt.issue(&item.program, &vars).with_context(|| format!("`{}`", item.program))?; + clock.issued(vars.get(&p.rows.var).copied().unwrap_or(0)); + let emits = p.forwards.iter().find(|f| f.name == item.program).and_then(|f| f.emits); + if let Some(i) = emits.filter(|_| last) { + let t = &p.fills[i]; + let mut tokens = t.decode(&rt.read_output(&t.name)?); + tokens.truncate(t.width as usize); + io.block_on(send(&mut w, &tokens))?; + } + if clock.half_full() { + clock.flush(&rt, cut.stage)?; + } + } + clock.flush(&rt, cut.stage)?; + Ok(()) +} + +/// Reach the head, retrying while it is still loading. +async fn connect(head: &str) -> Result { + let t0 = Instant::now(); + loop { + match TcpStream::connect(head).await { + Ok(s) => return Ok(s), + Err(e) if t0.elapsed() < CONNECT => { + tracing::debug!("head {head}: {e}"); + tokio::time::sleep(Duration::from_millis(500)).await; + } + Err(e) => return Err(e).with_context(|| format!("connecting to the head at {head}")), + } + } +} diff --git a/docs/roadmap.md b/docs/roadmap.md index 2f41e89e..872c5e50 100644 --- a/docs/roadmap.md +++ b/docs/roadmap.md @@ -58,6 +58,13 @@ kern 在内部 RSI 里是 kernel agent 的反馈环境（`kern test` 裁判、`k | R2 ✅ | blob store `Pegainfer/kern-kernels`（private，`blobs/` 平铺）+ `$KERN_INDEX_DIR/.toml` + `tools/kernels/abi/.py` + `check.py`（CI）；三个预编译核（trtllm fmha ctx、DSL mla decode、flash_kda）与手写核迁入，`handwritten.hw()` / `prebuilt()` 收编为 index 查询，`tools/kernels-bin/` 与生成器里的 nvcc 路径删除 | 重生成的 `examples/k3-*.json` 除 module `source` 外逐字节相同；4 层 TP4 prefill c8/c40/c7 与 EP4 c8/c40+free16 全 exact（同 P 行口径）；`check.py --online` 全部命中。**2026-09-19 tray14 过**：10 个 example 重生成 0 行非 source 差异；43 族 49 变体 53 blob 上传，`check.py` 离线/在线 0 问题；k3_golden 用空 `--cubins` 目录跑 TP4 c8/c40/c7、EP4 c8/c40+free16、decode span8/plain 七项，exact 数与 09-18 逐项相同（span8 的 1 wrong 是同一个近平局）；`~/bench_results/2026-09-19-registry/` | | R3 ✅ | 第一个新族 `trtllm_bmm_mxe4m3_mxe2m1_mxe4m3`（flashinfer-cubin 0.6.18 的 trtllm-gen batched GEMM，`siTuGlu` 变体即 K3 的激活）：整族导入、ABI 从 bundle 的 `KernelParams.h` + `flashinferMetaInfo.h` 写、vLLM 容器里 `trtllm_fp4_block_scale_moe` 探针 capture 验证；K3 prefill 的 MoE 换成 allgather(fp8 latent) → CTA 表 → FC1 → FC2 → finalize → reduce-scatter，MegaMoE 留给 decode 形态；`[[pick]]` 由 `kern bench` 对 16k/4 行的 K3 形状扫 tile 变体写入 | 4 层 TP4 prefill 门禁不变；93 层 12.9k 判定点对；16k over 0 < 780 ms（MegaMoE 231 ms 里 GEMM 之外的 86 ms 是靶，新增两次集合通信预算 ~40 ms）。**2026-09-19 tray03 过**：两族 215 变体导入、ABI 移植对 capture 逐字段 0 差异、probe 门禁 `out` 逐位同；`kern bench` 扫 18 对 tile 变体（`k3-prefill-16k-ep4`）写入 pick（tile 128：FC1 0.85 / FC2 0.40 ms 每层，捕获的 tile 64 是 1.23 / 0.58）；4 层 TP4 c40 1/1、c7 6/6，c8 4/5（step 15 是 4 ulp 近平局，逐层对比 MegaMoE 形态：同专家集合的行差 ~5% 即两条路各自的 fp8 中间量化）；93 层 12.9k 单块 2/2、c8 16/17（之前 13/17）；**16k over 0 = 749.5 ms**（之前 780；MoE 本体 231 → ~156，reduce-scatter 多出的 54 ms 是等最慢 rank 的专家负载）；`~/bench_results/2026-09-19-registry/` | +## 流水线线（PP） + +| 级 | 内容 | 门禁 | +|---|---|---| +| PP1 ✅ | `kern cut`（按 call 下标切、活跃性找边界、send/recv 邮箱核、`clock` fill）+ kern-serve 一段一进程的 prefill-only 流水（head 持唯一的池，TCP 发 item，GPU 邮箱传激活） | **2026-10-07 tray07 过**：qwen3-4b 2 段，各段依次跑 logits/token 与整份逐位同（chunk 512/64）；两进程首 token 与单卡同；8193 token conc1 TTFT 313 → 174 ms（serve.md「流水线」） | +| PP2 | 每段只分自己那些层的 KV（现在按整份 `bytes_per_token`）；decode 走流水；per-seq state（GDN）与吐 token 的 chunk program | | + ## 协议线 | 级 | 内容 | 门禁 | diff --git a/docs/serve.md b/docs/serve.md index d9edbccb..02c8e101 100644 --- a/docs/serve.md +++ b/docs/serve.md @@ -567,6 +567,48 @@ toy 抓到的几条： toy 不测 kernel 数值和性能，也没有 tray（EP4 的 toy 要一个走 peer 指针的 collective，还没写）。 +## 流水线（prefill-only PP v1，2026-10-07，tray07 GB300） + +一段一个进程：`kern cut` 切出的 stage 0 是 head（`--pp-listen host:port`：HTTP、scheduler、唯一的 +`Pool`），其余段 `--pp-head host:port`。握手：每段连上 head 发 `Hello { stage, cut, pages, mailbox }` +（export handle 的 hex），head 收齐 S−1 个、按最少的页数装自己的池，回 `Table { downstream_mailbox, +pages }`；之后同一条 TCP（tokio-util `LengthDelimitedCodec` + JSON）上 head 按发射顺序给每段发 +`Item { program, vars, rows }`（fill 与页表的值，页号全流水共用），末段对吐 token 的 item 回 token。 +激活不走 TCP，走 GPU 上的邮箱（manifest.md「流水线切分」），各段设备自己等上游。没有序号、 +stop、fault 消息：任何错误进程退出，连接断了对端退出；设备上的邮箱等待 60 s 超时 trap。 +只做 prefill：prompt 除末 token 按 chunk 走 chunk program，末 token 走一行的 step，首 token 即结束 +（`max_tokens` 视为 1）。每段每个 item 一行 `pp item` 日志：rows、busy_ns、recv_wait_ns、send_wait_ns +（读 `clock` fill，空闲时读，不加协议消息）。 + +门禁（qwen3-4b，切在 `l18.qkv_proj` 前，2 段）： +- 只切不分进程：`pp_chain` 各段依次跑 vs 整份，7 条 prompt（1–2049 token），chunk 512 × 8 步、 + chunk 64 × 16 步（2049 token 32 个 chunk），每步 logits 与 token **逐位同**。 +- 两进程 serve：7 条 gate prompt 首 token 与单卡 kern-serve 全同；杀末段 → head 立即退出；冻住末段 + （SIGSTOP）→ head 61 s 后 trap 退出，恢复的末段 61 s 后同样退出。 + +扫描（chunk 512，`--max-seqs 16`，每格 conc1 4 条 / conc4 16 条不同 prompt；TTFT = max_tokens=1 的 +请求延迟；busy/wait 是每 item 均值）： + +| conc | prompt | chunks | TTFT unsplit ms | TTFT pp ms | speedup | stage | chunk busy ms | step busy ms | recv_wait ms/item | send_wait ms/item | +|---|---:|---:|---:|---:|---:|---|---:|---:|---:|---:| +| c1 | 513 | 1 | 8.1 | 8.5 | 0.95x | 0 | 2.99 | 1.32 | - | 0.00 | +| | | | | | | 1 | 2.90 | 1.46 | 1.48 | - | +| c1 | 2049 | 4 | 38.2 | 25.3 | 1.51x | 0 | 4.27 | 1.46 | - | 0.00 | +| | | | | | | 1 | 4.21 | 1.60 | 1.16 | - | +| c1 | 8193 | 16 | 313.4 | 173.8 | 1.80x | 0 | 9.57 | 1.99 | - | 0.00 | +| | | | | | | 1 | 9.52 | 2.13 | 1.01 | - | +| c4 | 513 | 1 | 29.7 | 17.6 | 1.68x | 0 | 2.96 | 1.32 | - | 0.00 | +| | | | | | | 1 | 2.91 | 1.46 | 0.09 | - | +| c4 | 2049 | 4 | 139.6 | 78.5 | 1.78x | 0 | 4.26 | 1.46 | - | 0.19 | +| | | | | | | 1 | 4.20 | 1.60 | 0.30 | - | +| c4 | 8193 | 16 | 1194.0 | 664.5 | 1.80x | 0 | 9.57 | 1.99 | - | 0.61 | +| | | | | | | 1 | 9.52 | 2.13 | 0.73 | - | + +单请求内 chunk 之间就流水起来（stage 0 算第 k+1 块时 stage 1 算第 k 块），所以长 prompt conc1 也 +接近 2×；1 块的 prompt 没有可重叠的东西，略慢（邮箱 + 两段 launch）。v1 的限制：每段按整份的 +`bytes_per_token` 分 KV（只用自己那些层，浪费一半）；只认纯 paged 的 manifest；chunk program +自己吐 token 的模型拒切；上游那份 box 是不用的占位。 + ## 没做（按需要加） span 长的预算 policy（K5 D2：按稠密 / attention 预算定 c，现在是 `--chunk` 上限）、抢占 / 动态页分配、 From bbc94b8ec94412f05f464991adbd77448a2dbabc Mon Sep 17 00:00:00 2001 From: JinYan Su <751080330@qq.com> Date: Wed, 7 Oct 2026 15:44:10 +0000 Subject: [PATCH 03/22] qwen3: the target kv is one paged state per layer gen_qwen3_decode.py emits `kv.0`..`kv.35` (vLLM's per-layer pool layout, 4096 bytes per token each) instead of one layer-interleaved `kv`; the pool hands every paged state the same page ids, so nothing else changes. The draft's `draft_kv` stays interleaved. The generator had not run since the registry move (dead `handwritten` import, `variant` unimported, launch `label` dropped); it runs again, so the regenerated examples also take the registry builds of argmax, embedding and weight_prep. Gate (tray07): `kern test qwen3-4b` bit-identical at every span; new vs old qwen3-4b.json bit-identical at every span; `kern run` 32 steps text identical; dspark round and --rows 1, 128 steps, text identical to before. Co-Authored-By: Claude Opus 5.5 (1M context) --- crates/kern-manifest/tests/qwen3_example.rs | 6 +- examples/qwen3-4b-dspark.json | 2468 +++++++++---------- examples/qwen3-4b-silu-mined.json | 1668 ++++++------- examples/qwen3-4b.json | 1668 ++++++------- tools/gen_qwen3_decode.py | 43 +- 5 files changed, 2789 insertions(+), 3064 deletions(-) diff --git a/crates/kern-manifest/tests/qwen3_example.rs b/crates/kern-manifest/tests/qwen3_example.rs index d51b1589..2510d3d0 100644 --- a/crates/kern-manifest/tests/qwen3_example.rs +++ b/crates/kern-manifest/tests/qwen3_example.rs @@ -24,8 +24,10 @@ fn qwen3_decode_mined_verifies() { // prefill: same forward minus the final_norm/lm_head/sample tail — the // last prompt token goes through `decode` instead. assert_eq!(m.programs["prefill"].calls.len(), 2 + 36 * 12 - 1); - // The runtime's entire knowledge of the KV cache: a byte count. - assert_eq!(m.states["kv"].bytes_per_token, 36 * 2 * 1024 * 2); + // The runtime's entire knowledge of the KV cache: a byte count per + // layer's state, every one paged under the same ids. + assert_eq!(m.states.len(), 36); + assert!(m.states.values().all(|s| s.bytes_per_token == 2 * 1024 * 2)); let again = Manifest::from_json(&m.to_json()).expect("reparse"); assert_eq!(m.to_json(), again.to_json()); diff --git a/examples/qwen3-4b-dspark.json b/examples/qwen3-4b-dspark.json index ed29590d..78cb3ff9 100644 --- a/examples/qwen3-4b-dspark.json +++ b/examples/qwen3-4b-dspark.json @@ -22,7 +22,42 @@ "DRAFT_KV_SCALE_COUNT": 10, "MAX_TOKENS": 2048, "MAX_SEQS": 256, - "KV_BYTES_PER_TOKEN": 147456, + "KV.0_BYTES_PER_TOKEN": 4096, + "KV.1_BYTES_PER_TOKEN": 4096, + "KV.2_BYTES_PER_TOKEN": 4096, + "KV.3_BYTES_PER_TOKEN": 4096, + "KV.4_BYTES_PER_TOKEN": 4096, + "KV.5_BYTES_PER_TOKEN": 4096, + "KV.6_BYTES_PER_TOKEN": 4096, + "KV.7_BYTES_PER_TOKEN": 4096, + "KV.8_BYTES_PER_TOKEN": 4096, + "KV.9_BYTES_PER_TOKEN": 4096, + "KV.10_BYTES_PER_TOKEN": 4096, + "KV.11_BYTES_PER_TOKEN": 4096, + "KV.12_BYTES_PER_TOKEN": 4096, + "KV.13_BYTES_PER_TOKEN": 4096, + "KV.14_BYTES_PER_TOKEN": 4096, + "KV.15_BYTES_PER_TOKEN": 4096, + "KV.16_BYTES_PER_TOKEN": 4096, + "KV.17_BYTES_PER_TOKEN": 4096, + "KV.18_BYTES_PER_TOKEN": 4096, + "KV.19_BYTES_PER_TOKEN": 4096, + "KV.20_BYTES_PER_TOKEN": 4096, + "KV.21_BYTES_PER_TOKEN": 4096, + "KV.22_BYTES_PER_TOKEN": 4096, + "KV.23_BYTES_PER_TOKEN": 4096, + "KV.24_BYTES_PER_TOKEN": 4096, + "KV.25_BYTES_PER_TOKEN": 4096, + "KV.26_BYTES_PER_TOKEN": 4096, + "KV.27_BYTES_PER_TOKEN": 4096, + "KV.28_BYTES_PER_TOKEN": 4096, + "KV.29_BYTES_PER_TOKEN": 4096, + "KV.30_BYTES_PER_TOKEN": 4096, + "KV.31_BYTES_PER_TOKEN": 4096, + "KV.32_BYTES_PER_TOKEN": 4096, + "KV.33_BYTES_PER_TOKEN": 4096, + "KV.34_BYTES_PER_TOKEN": 4096, + "KV.35_BYTES_PER_TOKEN": 4096, "DRAFT_KV_BYTES_PER_TOKEN": 20480, "RMS_NORM_EPS": 9.999999974752427e-07, "NUM_HEADS": 32, @@ -41,8 +76,113 @@ } }, "states": { - "kv": { - "bytes_per_token": "KV_BYTES_PER_TOKEN" + "kv.0": { + "bytes_per_token": "KV.0_BYTES_PER_TOKEN" + }, + "kv.1": { + "bytes_per_token": "KV.1_BYTES_PER_TOKEN" + }, + "kv.2": { + "bytes_per_token": "KV.2_BYTES_PER_TOKEN" + }, + "kv.3": { + "bytes_per_token": "KV.3_BYTES_PER_TOKEN" + }, + "kv.4": { + "bytes_per_token": "KV.4_BYTES_PER_TOKEN" + }, + "kv.5": { + "bytes_per_token": "KV.5_BYTES_PER_TOKEN" + }, + "kv.6": { + "bytes_per_token": "KV.6_BYTES_PER_TOKEN" + }, + "kv.7": { + "bytes_per_token": "KV.7_BYTES_PER_TOKEN" + }, + "kv.8": { + "bytes_per_token": "KV.8_BYTES_PER_TOKEN" + }, + "kv.9": { + "bytes_per_token": "KV.9_BYTES_PER_TOKEN" + }, + "kv.10": { + "bytes_per_token": "KV.10_BYTES_PER_TOKEN" + }, + "kv.11": { + "bytes_per_token": "KV.11_BYTES_PER_TOKEN" + }, + "kv.12": { + "bytes_per_token": "KV.12_BYTES_PER_TOKEN" + }, + "kv.13": { + "bytes_per_token": "KV.13_BYTES_PER_TOKEN" + }, + "kv.14": { + "bytes_per_token": "KV.14_BYTES_PER_TOKEN" + }, + "kv.15": { + "bytes_per_token": "KV.15_BYTES_PER_TOKEN" + }, + "kv.16": { + "bytes_per_token": "KV.16_BYTES_PER_TOKEN" + }, + "kv.17": { + "bytes_per_token": "KV.17_BYTES_PER_TOKEN" + }, + "kv.18": { + "bytes_per_token": "KV.18_BYTES_PER_TOKEN" + }, + "kv.19": { + "bytes_per_token": "KV.19_BYTES_PER_TOKEN" + }, + "kv.20": { + "bytes_per_token": "KV.20_BYTES_PER_TOKEN" + }, + "kv.21": { + "bytes_per_token": "KV.21_BYTES_PER_TOKEN" + }, + "kv.22": { + "bytes_per_token": "KV.22_BYTES_PER_TOKEN" + }, + "kv.23": { + "bytes_per_token": "KV.23_BYTES_PER_TOKEN" + }, + "kv.24": { + "bytes_per_token": "KV.24_BYTES_PER_TOKEN" + }, + "kv.25": { + "bytes_per_token": "KV.25_BYTES_PER_TOKEN" + }, + "kv.26": { + "bytes_per_token": "KV.26_BYTES_PER_TOKEN" + }, + "kv.27": { + "bytes_per_token": "KV.27_BYTES_PER_TOKEN" + }, + "kv.28": { + "bytes_per_token": "KV.28_BYTES_PER_TOKEN" + }, + "kv.29": { + "bytes_per_token": "KV.29_BYTES_PER_TOKEN" + }, + "kv.30": { + "bytes_per_token": "KV.30_BYTES_PER_TOKEN" + }, + "kv.31": { + "bytes_per_token": "KV.31_BYTES_PER_TOKEN" + }, + "kv.32": { + "bytes_per_token": "KV.32_BYTES_PER_TOKEN" + }, + "kv.33": { + "bytes_per_token": "KV.33_BYTES_PER_TOKEN" + }, + "kv.34": { + "bytes_per_token": "KV.34_BYTES_PER_TOKEN" + }, + "kv.35": { + "bytes_per_token": "KV.35_BYTES_PER_TOKEN" }, "draft_kv": { "bytes_per_token": "DRAFT_KV_BYTES_PER_TOKEN" @@ -79,7 +219,7 @@ "kind": "input", "fill": "slot", "domain": { - "index_into": "kv" + "index_into": "kv.0" } }, "block_table": { @@ -90,7 +230,7 @@ ], "kind": "input", "domain": { - "index_into": "kv", + "index_into": "kv.0", "stride": "KV_BLOCK_SIZE" } }, @@ -5051,20 +5191,20 @@ }, "modules": { "argmax": { - "source": "argmax.cubin", - "sha256": "25374575bd35dbd9fe07f90982d0da0c2d4a72cc6ac88df5ec92cdb62a3beb16" + "source": "hf:Pegainfer/kern-kernels/blobs/ebe4ba79544ddfac43ee2dcf6a7f5299aba1124e176c201133d4ef9400e3dc25", + "sha256": "ebe4ba79544ddfac43ee2dcf6a7f5299aba1124e176c201133d4ef9400e3dc25" }, "copy_rows": { - "source": "copy_rows.cubin", - "sha256": "a49dcdf873038dbf4de19e8a324ac420afc669dca2cf122d601e034bcab10edd" + "source": "hf:Pegainfer/kern-kernels/blobs/f301274a5957faac2f0891ea9d6e80e0dabdb56ee4c14c3795e836f6bf31fb54", + "sha256": "f301274a5957faac2f0891ea9d6e80e0dabdb56ee4c14c3795e836f6bf31fb54" }, "embedding": { - "source": "embedding.cubin", - "sha256": "1c3dd3e591d872aed259d7ef59d44649061f77a4247a5df864a7fbf0ccbde92f" + "source": "hf:Pegainfer/kern-kernels/blobs/868e93ae6e46e36a11c56af1fae49aeff9fcdfdc79f64ca99475a309e721e9e7", + "sha256": "868e93ae6e46e36a11c56af1fae49aeff9fcdfdc79f64ca99475a309e721e9e7" }, "markov_rows": { - "source": "markov_rows.cubin", - "sha256": "73e7038bafa6c50023ad2db1a1534b83d59890c80573a23c79e16a77390fe72b" + "source": "hf:Pegainfer/kern-kernels/blobs/307c90a3eb9e8219ffa9045134edfc881d7c6d74e8a32b5ec1e36ef5b8421f87", + "sha256": "307c90a3eb9e8219ffa9045134edfc881d7c6d74e8a32b5ec1e36ef5b8421f87" }, "reduce_segments": { "source": "reduce_segments.cubin", @@ -5075,8 +5215,8 @@ "sha256": "1e5912bed482f8555e91a145a3cff4a1ce9b0c4240e23aa21b184c32d4716f13" }, "spec_round": { - "source": "spec_round.cubin", - "sha256": "83a1cd23bb15864787adb483d4a8b0737fb2c1c4d3a875750306d9a0f1ac0b8a" + "source": "hf:Pegainfer/kern-kernels/blobs/f8b61d2c217c24d55b020825ae4af59e80c6a08165c901e1385d12b8ffefec8d", + "sha256": "f8b61d2c217c24d55b020825ae4af59e80c6a08165c901e1385d12b8ffefec8d" }, "unified_causal": { "source": "unified_causal.cubin", @@ -5103,7 +5243,7 @@ "sha256": "116de812ac5a3a186f76fa239d4c16abbed6b8b0af06d16e1baec0aefe8fe8fa" }, "weight_prep": { - "source": "weight_prep.cubin", + "source": "hf:Pegainfer/kern-kernels/blobs/040f51b01c421b9c69a2090034ff1bc4741e40c02f5ea52bd04a3e62e5c1af4d", "sha256": "040f51b01c421b9c69a2090034ff1bc4741e40c02f5ea52bd04a3e62e5c1af4d" } }, @@ -5683,7 +5823,7 @@ "param": 8 }, { - "i64": 1179648 + "i64": 32768 }, { "i64": 2048 @@ -5692,7 +5832,7 @@ "i64": 256 }, { - "i64": 1179648 + "i64": 32768 }, { "i64": 2048 @@ -5914,7 +6054,7 @@ "param": 8 }, { - "i64": 1179648 + "i64": 32768 }, { "i64": 2048 @@ -5923,7 +6063,7 @@ "i64": 256 }, { - "i64": 1179648 + "i64": 32768 }, { "i64": 2048 @@ -6074,7 +6214,7 @@ "param": 8 }, { - "i64": 1179648 + "i64": 32768 }, { "i64": 2048 @@ -6083,7 +6223,7 @@ "i64": 256 }, { - "i64": 1179648 + "i64": 32768 }, { "i64": 2048 @@ -6664,7 +6804,7 @@ "param": 8 }, { - "i64": 1179648 + "i64": 32768 }, { "i64": 2048 @@ -6673,7 +6813,7 @@ "i64": 256 }, { - "i64": 1179648 + "i64": 32768 }, { "i64": 2048 @@ -7477,10 +7617,10 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv" + "state": "kv.0" }, { - "state": "kv", + "state": "kv.0", "offset": 256 }, { @@ -7497,7 +7637,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -7512,10 +7652,10 @@ "buf": "q_n" }, { - "state": "kv" + "state": "kv.0" }, { - "state": "kv", + "state": "kv.0", "offset": 256 }, { @@ -7802,12 +7942,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 65536 + "state": "kv.1" }, { - "state": "kv", - "offset": 65792 + "state": "kv.1", + "offset": 256 }, { "buf": "slot_mapping" @@ -7824,7 +7963,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -7839,12 +7978,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 65536 + "state": "kv.1" }, { - "state": "kv", - "offset": 65792 + "state": "kv.1", + "offset": 256 }, { "buf": "block_table" @@ -8107,12 +8245,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 131072 + "state": "kv.2" }, { - "state": "kv", - "offset": 131328 + "state": "kv.2", + "offset": 256 }, { "buf": "slot_mapping" @@ -8129,7 +8266,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -8144,12 +8281,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 131072 + "state": "kv.2" }, { - "state": "kv", - "offset": 131328 + "state": "kv.2", + "offset": 256 }, { "buf": "block_table" @@ -8412,12 +8548,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 196608 + "state": "kv.3" }, { - "state": "kv", - "offset": 196864 + "state": "kv.3", + "offset": 256 }, { "buf": "slot_mapping" @@ -8434,7 +8569,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -8449,12 +8584,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 196608 + "state": "kv.3" }, { - "state": "kv", - "offset": 196864 + "state": "kv.3", + "offset": 256 }, { "buf": "block_table" @@ -8717,12 +8851,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 262144 + "state": "kv.4" }, { - "state": "kv", - "offset": 262400 + "state": "kv.4", + "offset": 256 }, { "buf": "slot_mapping" @@ -8739,7 +8872,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -8754,12 +8887,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 262144 + "state": "kv.4" }, { - "state": "kv", - "offset": 262400 + "state": "kv.4", + "offset": 256 }, { "buf": "block_table" @@ -9022,12 +9154,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 327680 + "state": "kv.5" }, { - "state": "kv", - "offset": 327936 + "state": "kv.5", + "offset": 256 }, { "buf": "slot_mapping" @@ -9044,7 +9175,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -9059,12 +9190,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 327680 + "state": "kv.5" }, { - "state": "kv", - "offset": 327936 + "state": "kv.5", + "offset": 256 }, { "buf": "block_table" @@ -9327,12 +9457,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 393216 + "state": "kv.6" }, { - "state": "kv", - "offset": 393472 + "state": "kv.6", + "offset": 256 }, { "buf": "slot_mapping" @@ -9349,7 +9478,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -9364,12 +9493,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 393216 + "state": "kv.6" }, { - "state": "kv", - "offset": 393472 + "state": "kv.6", + "offset": 256 }, { "buf": "block_table" @@ -9632,12 +9760,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 458752 + "state": "kv.7" }, { - "state": "kv", - "offset": 459008 + "state": "kv.7", + "offset": 256 }, { "buf": "slot_mapping" @@ -9654,7 +9781,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -9669,12 +9796,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 458752 + "state": "kv.7" }, { - "state": "kv", - "offset": 459008 + "state": "kv.7", + "offset": 256 }, { "buf": "block_table" @@ -9937,12 +10063,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 524288 + "state": "kv.8" }, { - "state": "kv", - "offset": 524544 + "state": "kv.8", + "offset": 256 }, { "buf": "slot_mapping" @@ -9959,7 +10084,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -9974,12 +10099,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 524288 + "state": "kv.8" }, { - "state": "kv", - "offset": 524544 + "state": "kv.8", + "offset": 256 }, { "buf": "block_table" @@ -10260,12 +10384,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 589824 + "state": "kv.9" }, { - "state": "kv", - "offset": 590080 + "state": "kv.9", + "offset": 256 }, { "buf": "slot_mapping" @@ -10282,7 +10405,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -10297,12 +10420,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 589824 + "state": "kv.9" }, { - "state": "kv", - "offset": 590080 + "state": "kv.9", + "offset": 256 }, { "buf": "block_table" @@ -10565,12 +10687,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 655360 + "state": "kv.10" }, { - "state": "kv", - "offset": 655616 + "state": "kv.10", + "offset": 256 }, { "buf": "slot_mapping" @@ -10587,7 +10708,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -10602,12 +10723,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 655360 + "state": "kv.10" }, { - "state": "kv", - "offset": 655616 + "state": "kv.10", + "offset": 256 }, { "buf": "block_table" @@ -10870,12 +10990,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 720896 + "state": "kv.11" }, { - "state": "kv", - "offset": 721152 + "state": "kv.11", + "offset": 256 }, { "buf": "slot_mapping" @@ -10892,7 +11011,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -10907,12 +11026,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 720896 + "state": "kv.11" }, { - "state": "kv", - "offset": 721152 + "state": "kv.11", + "offset": 256 }, { "buf": "block_table" @@ -11175,12 +11293,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 786432 + "state": "kv.12" }, { - "state": "kv", - "offset": 786688 + "state": "kv.12", + "offset": 256 }, { "buf": "slot_mapping" @@ -11197,7 +11314,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -11212,12 +11329,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 786432 + "state": "kv.12" }, { - "state": "kv", - "offset": 786688 + "state": "kv.12", + "offset": 256 }, { "buf": "block_table" @@ -11480,12 +11596,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 851968 + "state": "kv.13" }, { - "state": "kv", - "offset": 852224 + "state": "kv.13", + "offset": 256 }, { "buf": "slot_mapping" @@ -11502,7 +11617,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -11517,12 +11632,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 851968 + "state": "kv.13" }, { - "state": "kv", - "offset": 852224 + "state": "kv.13", + "offset": 256 }, { "buf": "block_table" @@ -11785,12 +11899,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 917504 + "state": "kv.14" }, { - "state": "kv", - "offset": 917760 + "state": "kv.14", + "offset": 256 }, { "buf": "slot_mapping" @@ -11807,7 +11920,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -11822,12 +11935,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 917504 + "state": "kv.14" }, { - "state": "kv", - "offset": 917760 + "state": "kv.14", + "offset": 256 }, { "buf": "block_table" @@ -12090,12 +12202,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 983040 + "state": "kv.15" }, { - "state": "kv", - "offset": 983296 + "state": "kv.15", + "offset": 256 }, { "buf": "slot_mapping" @@ -12112,7 +12223,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -12127,12 +12238,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 983040 + "state": "kv.15" }, { - "state": "kv", - "offset": 983296 + "state": "kv.15", + "offset": 256 }, { "buf": "block_table" @@ -12395,12 +12505,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1048576 + "state": "kv.16" }, { - "state": "kv", - "offset": 1048832 + "state": "kv.16", + "offset": 256 }, { "buf": "slot_mapping" @@ -12417,7 +12526,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -12432,12 +12541,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1048576 + "state": "kv.16" }, { - "state": "kv", - "offset": 1048832 + "state": "kv.16", + "offset": 256 }, { "buf": "block_table" @@ -12718,12 +12826,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1114112 + "state": "kv.17" }, { - "state": "kv", - "offset": 1114368 + "state": "kv.17", + "offset": 256 }, { "buf": "slot_mapping" @@ -12740,7 +12847,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -12755,12 +12862,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1114112 + "state": "kv.17" }, { - "state": "kv", - "offset": 1114368 + "state": "kv.17", + "offset": 256 }, { "buf": "block_table" @@ -13023,12 +13129,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1179648 + "state": "kv.18" }, { - "state": "kv", - "offset": 1179904 + "state": "kv.18", + "offset": 256 }, { "buf": "slot_mapping" @@ -13045,7 +13150,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -13060,12 +13165,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1179648 + "state": "kv.18" }, { - "state": "kv", - "offset": 1179904 + "state": "kv.18", + "offset": 256 }, { "buf": "block_table" @@ -13328,12 +13432,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1245184 + "state": "kv.19" }, { - "state": "kv", - "offset": 1245440 + "state": "kv.19", + "offset": 256 }, { "buf": "slot_mapping" @@ -13350,7 +13453,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -13365,12 +13468,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1245184 + "state": "kv.19" }, { - "state": "kv", - "offset": 1245440 + "state": "kv.19", + "offset": 256 }, { "buf": "block_table" @@ -13633,12 +13735,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1310720 + "state": "kv.20" }, { - "state": "kv", - "offset": 1310976 + "state": "kv.20", + "offset": 256 }, { "buf": "slot_mapping" @@ -13655,7 +13756,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -13670,12 +13771,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1310720 + "state": "kv.20" }, { - "state": "kv", - "offset": 1310976 + "state": "kv.20", + "offset": 256 }, { "buf": "block_table" @@ -13938,12 +14038,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1376256 + "state": "kv.21" }, { - "state": "kv", - "offset": 1376512 + "state": "kv.21", + "offset": 256 }, { "buf": "slot_mapping" @@ -13960,7 +14059,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -13975,12 +14074,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1376256 + "state": "kv.21" }, { - "state": "kv", - "offset": 1376512 + "state": "kv.21", + "offset": 256 }, { "buf": "block_table" @@ -14243,12 +14341,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1441792 + "state": "kv.22" }, { - "state": "kv", - "offset": 1442048 + "state": "kv.22", + "offset": 256 }, { "buf": "slot_mapping" @@ -14265,7 +14362,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -14280,12 +14377,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1441792 + "state": "kv.22" }, { - "state": "kv", - "offset": 1442048 + "state": "kv.22", + "offset": 256 }, { "buf": "block_table" @@ -14548,12 +14644,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1507328 + "state": "kv.23" }, { - "state": "kv", - "offset": 1507584 + "state": "kv.23", + "offset": 256 }, { "buf": "slot_mapping" @@ -14570,7 +14665,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -14585,12 +14680,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1507328 + "state": "kv.23" }, { - "state": "kv", - "offset": 1507584 + "state": "kv.23", + "offset": 256 }, { "buf": "block_table" @@ -14853,12 +14947,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1572864 + "state": "kv.24" }, { - "state": "kv", - "offset": 1573120 + "state": "kv.24", + "offset": 256 }, { "buf": "slot_mapping" @@ -14875,7 +14968,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -14890,12 +14983,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1572864 + "state": "kv.24" }, { - "state": "kv", - "offset": 1573120 + "state": "kv.24", + "offset": 256 }, { "buf": "block_table" @@ -15176,12 +15268,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1638400 + "state": "kv.25" }, { - "state": "kv", - "offset": 1638656 + "state": "kv.25", + "offset": 256 }, { "buf": "slot_mapping" @@ -15198,7 +15289,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -15213,12 +15304,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1638400 + "state": "kv.25" }, { - "state": "kv", - "offset": 1638656 + "state": "kv.25", + "offset": 256 }, { "buf": "block_table" @@ -15481,12 +15571,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1703936 + "state": "kv.26" }, { - "state": "kv", - "offset": 1704192 + "state": "kv.26", + "offset": 256 }, { "buf": "slot_mapping" @@ -15503,7 +15592,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -15518,12 +15607,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1703936 + "state": "kv.26" }, { - "state": "kv", - "offset": 1704192 + "state": "kv.26", + "offset": 256 }, { "buf": "block_table" @@ -15786,12 +15874,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1769472 + "state": "kv.27" }, { - "state": "kv", - "offset": 1769728 + "state": "kv.27", + "offset": 256 }, { "buf": "slot_mapping" @@ -15808,7 +15895,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -15823,12 +15910,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1769472 + "state": "kv.27" }, { - "state": "kv", - "offset": 1769728 + "state": "kv.27", + "offset": 256 }, { "buf": "block_table" @@ -16091,12 +16177,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1835008 + "state": "kv.28" }, { - "state": "kv", - "offset": 1835264 + "state": "kv.28", + "offset": 256 }, { "buf": "slot_mapping" @@ -16113,7 +16198,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -16128,12 +16213,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1835008 + "state": "kv.28" }, { - "state": "kv", - "offset": 1835264 + "state": "kv.28", + "offset": 256 }, { "buf": "block_table" @@ -16396,12 +16480,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1900544 + "state": "kv.29" }, { - "state": "kv", - "offset": 1900800 + "state": "kv.29", + "offset": 256 }, { "buf": "slot_mapping" @@ -16418,7 +16501,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -16433,12 +16516,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1900544 + "state": "kv.29" }, { - "state": "kv", - "offset": 1900800 + "state": "kv.29", + "offset": 256 }, { "buf": "block_table" @@ -16701,12 +16783,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1966080 + "state": "kv.30" }, { - "state": "kv", - "offset": 1966336 + "state": "kv.30", + "offset": 256 }, { "buf": "slot_mapping" @@ -16723,7 +16804,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -16738,12 +16819,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1966080 + "state": "kv.30" }, { - "state": "kv", - "offset": 1966336 + "state": "kv.30", + "offset": 256 }, { "buf": "block_table" @@ -17006,12 +17086,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 2031616 + "state": "kv.31" }, { - "state": "kv", - "offset": 2031872 + "state": "kv.31", + "offset": 256 }, { "buf": "slot_mapping" @@ -17028,7 +17107,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -17043,12 +17122,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 2031616 + "state": "kv.31" }, { - "state": "kv", - "offset": 2031872 + "state": "kv.31", + "offset": 256 }, { "buf": "block_table" @@ -17311,12 +17389,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 2097152 + "state": "kv.32" }, { - "state": "kv", - "offset": 2097408 + "state": "kv.32", + "offset": 256 }, { "buf": "slot_mapping" @@ -17333,7 +17410,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -17348,12 +17425,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 2097152 + "state": "kv.32" }, { - "state": "kv", - "offset": 2097408 + "state": "kv.32", + "offset": 256 }, { "buf": "block_table" @@ -17634,12 +17710,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 2162688 + "state": "kv.33" }, { - "state": "kv", - "offset": 2162944 + "state": "kv.33", + "offset": 256 }, { "buf": "slot_mapping" @@ -17656,7 +17731,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -17671,12 +17746,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 2162688 + "state": "kv.33" }, { - "state": "kv", - "offset": 2162944 + "state": "kv.33", + "offset": 256 }, { "buf": "block_table" @@ -17939,12 +18013,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 2228224 + "state": "kv.34" }, { - "state": "kv", - "offset": 2228480 + "state": "kv.34", + "offset": 256 }, { "buf": "slot_mapping" @@ -17961,7 +18034,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -17976,12 +18049,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 2228224 + "state": "kv.34" }, { - "state": "kv", - "offset": 2228480 + "state": "kv.34", + "offset": 256 }, { "buf": "block_table" @@ -18244,12 +18316,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 2293760 + "state": "kv.35" }, { - "state": "kv", - "offset": 2294016 + "state": "kv.35", + "offset": 256 }, { "buf": "slot_mapping" @@ -18266,7 +18337,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -18281,12 +18352,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 2293760 + "state": "kv.35" }, { - "state": "kv", - "offset": 2294016 + "state": "kv.35", + "offset": 256 }, { "buf": "block_table" @@ -19159,10 +19229,10 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv" + "state": "kv.0" }, { - "state": "kv", + "state": "kv.0", "offset": 256 }, { @@ -19179,7 +19249,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -19194,10 +19264,10 @@ "buf": "q_n" }, { - "state": "kv" + "state": "kv.0" }, { - "state": "kv", + "state": "kv.0", "offset": 256 }, { @@ -19457,12 +19527,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 65536 + "state": "kv.1" }, { - "state": "kv", - "offset": 65792 + "state": "kv.1", + "offset": 256 }, { "buf": "slot_mapping" @@ -19479,7 +19548,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -19494,12 +19563,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 65536 + "state": "kv.1" }, { - "state": "kv", - "offset": 65792 + "state": "kv.1", + "offset": 256 }, { "buf": "block_table" @@ -19759,12 +19827,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 131072 + "state": "kv.2" }, { - "state": "kv", - "offset": 131328 + "state": "kv.2", + "offset": 256 }, { "buf": "slot_mapping" @@ -19781,7 +19848,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -19796,12 +19863,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 131072 + "state": "kv.2" }, { - "state": "kv", - "offset": 131328 + "state": "kv.2", + "offset": 256 }, { "buf": "block_table" @@ -20061,12 +20127,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 196608 + "state": "kv.3" }, { - "state": "kv", - "offset": 196864 + "state": "kv.3", + "offset": 256 }, { "buf": "slot_mapping" @@ -20083,7 +20148,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -20098,12 +20163,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 196608 + "state": "kv.3" }, { - "state": "kv", - "offset": 196864 + "state": "kv.3", + "offset": 256 }, { "buf": "block_table" @@ -20363,12 +20427,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 262144 + "state": "kv.4" }, { - "state": "kv", - "offset": 262400 + "state": "kv.4", + "offset": 256 }, { "buf": "slot_mapping" @@ -20385,7 +20448,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -20400,12 +20463,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 262144 + "state": "kv.4" }, { - "state": "kv", - "offset": 262400 + "state": "kv.4", + "offset": 256 }, { "buf": "block_table" @@ -20665,12 +20727,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 327680 + "state": "kv.5" }, { - "state": "kv", - "offset": 327936 + "state": "kv.5", + "offset": 256 }, { "buf": "slot_mapping" @@ -20687,7 +20748,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -20702,12 +20763,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 327680 + "state": "kv.5" }, { - "state": "kv", - "offset": 327936 + "state": "kv.5", + "offset": 256 }, { "buf": "block_table" @@ -20967,12 +21027,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 393216 + "state": "kv.6" }, { - "state": "kv", - "offset": 393472 + "state": "kv.6", + "offset": 256 }, { "buf": "slot_mapping" @@ -20989,7 +21048,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -21004,12 +21063,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 393216 + "state": "kv.6" }, { - "state": "kv", - "offset": 393472 + "state": "kv.6", + "offset": 256 }, { "buf": "block_table" @@ -21269,12 +21327,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 458752 + "state": "kv.7" }, { - "state": "kv", - "offset": 459008 + "state": "kv.7", + "offset": 256 }, { "buf": "slot_mapping" @@ -21291,7 +21348,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -21306,12 +21363,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 458752 + "state": "kv.7" }, { - "state": "kv", - "offset": 459008 + "state": "kv.7", + "offset": 256 }, { "buf": "block_table" @@ -21571,12 +21627,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 524288 + "state": "kv.8" }, { - "state": "kv", - "offset": 524544 + "state": "kv.8", + "offset": 256 }, { "buf": "slot_mapping" @@ -21593,7 +21648,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -21608,12 +21663,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 524288 + "state": "kv.8" }, { - "state": "kv", - "offset": 524544 + "state": "kv.8", + "offset": 256 }, { "buf": "block_table" @@ -21873,12 +21927,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 589824 + "state": "kv.9" }, { - "state": "kv", - "offset": 590080 + "state": "kv.9", + "offset": 256 }, { "buf": "slot_mapping" @@ -21895,7 +21948,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -21910,12 +21963,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 589824 + "state": "kv.9" }, { - "state": "kv", - "offset": 590080 + "state": "kv.9", + "offset": 256 }, { "buf": "block_table" @@ -22175,12 +22227,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 655360 + "state": "kv.10" }, { - "state": "kv", - "offset": 655616 + "state": "kv.10", + "offset": 256 }, { "buf": "slot_mapping" @@ -22197,7 +22248,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -22212,12 +22263,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 655360 + "state": "kv.10" }, { - "state": "kv", - "offset": 655616 + "state": "kv.10", + "offset": 256 }, { "buf": "block_table" @@ -22477,12 +22527,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 720896 + "state": "kv.11" }, { - "state": "kv", - "offset": 721152 + "state": "kv.11", + "offset": 256 }, { "buf": "slot_mapping" @@ -22499,7 +22548,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -22514,12 +22563,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 720896 + "state": "kv.11" }, { - "state": "kv", - "offset": 721152 + "state": "kv.11", + "offset": 256 }, { "buf": "block_table" @@ -22779,12 +22827,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 786432 + "state": "kv.12" }, { - "state": "kv", - "offset": 786688 + "state": "kv.12", + "offset": 256 }, { "buf": "slot_mapping" @@ -22801,7 +22848,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -22816,12 +22863,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 786432 + "state": "kv.12" }, { - "state": "kv", - "offset": 786688 + "state": "kv.12", + "offset": 256 }, { "buf": "block_table" @@ -23081,12 +23127,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 851968 + "state": "kv.13" }, { - "state": "kv", - "offset": 852224 + "state": "kv.13", + "offset": 256 }, { "buf": "slot_mapping" @@ -23103,7 +23148,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -23118,12 +23163,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 851968 + "state": "kv.13" }, { - "state": "kv", - "offset": 852224 + "state": "kv.13", + "offset": 256 }, { "buf": "block_table" @@ -23383,12 +23427,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 917504 + "state": "kv.14" }, { - "state": "kv", - "offset": 917760 + "state": "kv.14", + "offset": 256 }, { "buf": "slot_mapping" @@ -23405,7 +23448,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -23420,12 +23463,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 917504 + "state": "kv.14" }, { - "state": "kv", - "offset": 917760 + "state": "kv.14", + "offset": 256 }, { "buf": "block_table" @@ -23685,12 +23727,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 983040 + "state": "kv.15" }, { - "state": "kv", - "offset": 983296 + "state": "kv.15", + "offset": 256 }, { "buf": "slot_mapping" @@ -23707,7 +23748,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -23722,12 +23763,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 983040 + "state": "kv.15" }, { - "state": "kv", - "offset": 983296 + "state": "kv.15", + "offset": 256 }, { "buf": "block_table" @@ -23987,12 +24027,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1048576 + "state": "kv.16" }, { - "state": "kv", - "offset": 1048832 + "state": "kv.16", + "offset": 256 }, { "buf": "slot_mapping" @@ -24009,7 +24048,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -24024,12 +24063,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1048576 + "state": "kv.16" }, { - "state": "kv", - "offset": 1048832 + "state": "kv.16", + "offset": 256 }, { "buf": "block_table" @@ -24289,12 +24327,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1114112 + "state": "kv.17" }, { - "state": "kv", - "offset": 1114368 + "state": "kv.17", + "offset": 256 }, { "buf": "slot_mapping" @@ -24311,7 +24348,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -24326,12 +24363,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1114112 + "state": "kv.17" }, { - "state": "kv", - "offset": 1114368 + "state": "kv.17", + "offset": 256 }, { "buf": "block_table" @@ -24591,12 +24627,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1179648 + "state": "kv.18" }, { - "state": "kv", - "offset": 1179904 + "state": "kv.18", + "offset": 256 }, { "buf": "slot_mapping" @@ -24613,7 +24648,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -24628,12 +24663,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1179648 + "state": "kv.18" }, { - "state": "kv", - "offset": 1179904 + "state": "kv.18", + "offset": 256 }, { "buf": "block_table" @@ -24893,12 +24927,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1245184 + "state": "kv.19" }, { - "state": "kv", - "offset": 1245440 + "state": "kv.19", + "offset": 256 }, { "buf": "slot_mapping" @@ -24915,7 +24948,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -24930,12 +24963,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1245184 + "state": "kv.19" }, { - "state": "kv", - "offset": 1245440 + "state": "kv.19", + "offset": 256 }, { "buf": "block_table" @@ -25195,12 +25227,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1310720 + "state": "kv.20" }, { - "state": "kv", - "offset": 1310976 + "state": "kv.20", + "offset": 256 }, { "buf": "slot_mapping" @@ -25217,7 +25248,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -25232,12 +25263,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1310720 + "state": "kv.20" }, { - "state": "kv", - "offset": 1310976 + "state": "kv.20", + "offset": 256 }, { "buf": "block_table" @@ -25497,12 +25527,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1376256 + "state": "kv.21" }, { - "state": "kv", - "offset": 1376512 + "state": "kv.21", + "offset": 256 }, { "buf": "slot_mapping" @@ -25519,7 +25548,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -25534,12 +25563,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1376256 + "state": "kv.21" }, { - "state": "kv", - "offset": 1376512 + "state": "kv.21", + "offset": 256 }, { "buf": "block_table" @@ -25799,12 +25827,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1441792 + "state": "kv.22" }, { - "state": "kv", - "offset": 1442048 + "state": "kv.22", + "offset": 256 }, { "buf": "slot_mapping" @@ -25821,7 +25848,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -25836,12 +25863,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1441792 + "state": "kv.22" }, { - "state": "kv", - "offset": 1442048 + "state": "kv.22", + "offset": 256 }, { "buf": "block_table" @@ -26101,12 +26127,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1507328 + "state": "kv.23" }, { - "state": "kv", - "offset": 1507584 + "state": "kv.23", + "offset": 256 }, { "buf": "slot_mapping" @@ -26123,7 +26148,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -26138,12 +26163,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1507328 + "state": "kv.23" }, { - "state": "kv", - "offset": 1507584 + "state": "kv.23", + "offset": 256 }, { "buf": "block_table" @@ -26403,12 +26427,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1572864 + "state": "kv.24" }, { - "state": "kv", - "offset": 1573120 + "state": "kv.24", + "offset": 256 }, { "buf": "slot_mapping" @@ -26425,7 +26448,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -26440,12 +26463,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1572864 + "state": "kv.24" }, { - "state": "kv", - "offset": 1573120 + "state": "kv.24", + "offset": 256 }, { "buf": "block_table" @@ -26705,12 +26727,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1638400 + "state": "kv.25" }, { - "state": "kv", - "offset": 1638656 + "state": "kv.25", + "offset": 256 }, { "buf": "slot_mapping" @@ -26727,7 +26748,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -26742,12 +26763,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1638400 + "state": "kv.25" }, { - "state": "kv", - "offset": 1638656 + "state": "kv.25", + "offset": 256 }, { "buf": "block_table" @@ -27007,12 +27027,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1703936 + "state": "kv.26" }, { - "state": "kv", - "offset": 1704192 + "state": "kv.26", + "offset": 256 }, { "buf": "slot_mapping" @@ -27029,7 +27048,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -27044,12 +27063,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1703936 + "state": "kv.26" }, { - "state": "kv", - "offset": 1704192 + "state": "kv.26", + "offset": 256 }, { "buf": "block_table" @@ -27309,12 +27327,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1769472 + "state": "kv.27" }, { - "state": "kv", - "offset": 1769728 + "state": "kv.27", + "offset": 256 }, { "buf": "slot_mapping" @@ -27331,7 +27348,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -27346,12 +27363,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1769472 + "state": "kv.27" }, { - "state": "kv", - "offset": 1769728 + "state": "kv.27", + "offset": 256 }, { "buf": "block_table" @@ -27611,12 +27627,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1835008 + "state": "kv.28" }, { - "state": "kv", - "offset": 1835264 + "state": "kv.28", + "offset": 256 }, { "buf": "slot_mapping" @@ -27633,7 +27648,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -27648,12 +27663,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1835008 + "state": "kv.28" }, { - "state": "kv", - "offset": 1835264 + "state": "kv.28", + "offset": 256 }, { "buf": "block_table" @@ -27913,12 +27927,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1900544 + "state": "kv.29" }, { - "state": "kv", - "offset": 1900800 + "state": "kv.29", + "offset": 256 }, { "buf": "slot_mapping" @@ -27935,7 +27948,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -27950,12 +27963,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1900544 + "state": "kv.29" }, { - "state": "kv", - "offset": 1900800 + "state": "kv.29", + "offset": 256 }, { "buf": "block_table" @@ -28215,12 +28227,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1966080 + "state": "kv.30" }, { - "state": "kv", - "offset": 1966336 + "state": "kv.30", + "offset": 256 }, { "buf": "slot_mapping" @@ -28237,7 +28248,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -28252,12 +28263,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1966080 + "state": "kv.30" }, { - "state": "kv", - "offset": 1966336 + "state": "kv.30", + "offset": 256 }, { "buf": "block_table" @@ -28517,12 +28527,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 2031616 + "state": "kv.31" }, { - "state": "kv", - "offset": 2031872 + "state": "kv.31", + "offset": 256 }, { "buf": "slot_mapping" @@ -28539,7 +28548,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -28554,12 +28563,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 2031616 + "state": "kv.31" }, { - "state": "kv", - "offset": 2031872 + "state": "kv.31", + "offset": 256 }, { "buf": "block_table" @@ -28819,12 +28827,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 2097152 + "state": "kv.32" }, { - "state": "kv", - "offset": 2097408 + "state": "kv.32", + "offset": 256 }, { "buf": "slot_mapping" @@ -28841,7 +28848,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -28856,12 +28863,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 2097152 + "state": "kv.32" }, { - "state": "kv", - "offset": 2097408 + "state": "kv.32", + "offset": 256 }, { "buf": "block_table" @@ -29121,12 +29127,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 2162688 + "state": "kv.33" }, { - "state": "kv", - "offset": 2162944 + "state": "kv.33", + "offset": 256 }, { "buf": "slot_mapping" @@ -29143,7 +29148,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -29158,12 +29163,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 2162688 + "state": "kv.33" }, { - "state": "kv", - "offset": 2162944 + "state": "kv.33", + "offset": 256 }, { "buf": "block_table" @@ -29423,12 +29427,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 2228224 + "state": "kv.34" }, { - "state": "kv", - "offset": 2228480 + "state": "kv.34", + "offset": 256 }, { "buf": "slot_mapping" @@ -29445,7 +29448,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -29460,12 +29463,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 2228224 + "state": "kv.34" }, { - "state": "kv", - "offset": 2228480 + "state": "kv.34", + "offset": 256 }, { "buf": "block_table" @@ -29725,12 +29727,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 2293760 + "state": "kv.35" }, { - "state": "kv", - "offset": 2294016 + "state": "kv.35", + "offset": 256 }, { "buf": "slot_mapping" @@ -29747,7 +29748,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -29762,12 +29763,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 2293760 + "state": "kv.35" }, { - "state": "kv", - "offset": 2294016 + "state": "kv.35", + "offset": 256 }, { "buf": "block_table" @@ -30108,10 +30108,10 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv" + "state": "kv.0" }, { - "state": "kv", + "state": "kv.0", "offset": 256 }, { @@ -30128,7 +30128,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -30143,10 +30143,10 @@ "buf": "q_n" }, { - "state": "kv" + "state": "kv.0" }, { - "state": "kv", + "state": "kv.0", "offset": 256 }, { @@ -30409,12 +30409,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 65536 + "state": "kv.1" }, { - "state": "kv", - "offset": 65792 + "state": "kv.1", + "offset": 256 }, { "buf": "slot_mapping" @@ -30431,7 +30430,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -30446,12 +30445,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 65536 + "state": "kv.1" }, { - "state": "kv", - "offset": 65792 + "state": "kv.1", + "offset": 256 }, { "buf": "block_table" @@ -30714,12 +30712,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 131072 + "state": "kv.2" }, { - "state": "kv", - "offset": 131328 + "state": "kv.2", + "offset": 256 }, { "buf": "slot_mapping" @@ -30736,7 +30733,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -30751,12 +30748,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 131072 + "state": "kv.2" }, { - "state": "kv", - "offset": 131328 + "state": "kv.2", + "offset": 256 }, { "buf": "block_table" @@ -31019,12 +31015,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 196608 + "state": "kv.3" }, { - "state": "kv", - "offset": 196864 + "state": "kv.3", + "offset": 256 }, { "buf": "slot_mapping" @@ -31041,7 +31036,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -31056,12 +31051,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 196608 + "state": "kv.3" }, { - "state": "kv", - "offset": 196864 + "state": "kv.3", + "offset": 256 }, { "buf": "block_table" @@ -31324,12 +31318,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 262144 + "state": "kv.4" }, { - "state": "kv", - "offset": 262400 + "state": "kv.4", + "offset": 256 }, { "buf": "slot_mapping" @@ -31346,7 +31339,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -31361,12 +31354,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 262144 + "state": "kv.4" }, { - "state": "kv", - "offset": 262400 + "state": "kv.4", + "offset": 256 }, { "buf": "block_table" @@ -31629,12 +31621,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 327680 + "state": "kv.5" }, { - "state": "kv", - "offset": 327936 + "state": "kv.5", + "offset": 256 }, { "buf": "slot_mapping" @@ -31651,7 +31642,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -31666,12 +31657,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 327680 + "state": "kv.5" }, { - "state": "kv", - "offset": 327936 + "state": "kv.5", + "offset": 256 }, { "buf": "block_table" @@ -31934,12 +31924,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 393216 + "state": "kv.6" }, { - "state": "kv", - "offset": 393472 + "state": "kv.6", + "offset": 256 }, { "buf": "slot_mapping" @@ -31956,7 +31945,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -31971,12 +31960,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 393216 + "state": "kv.6" }, { - "state": "kv", - "offset": 393472 + "state": "kv.6", + "offset": 256 }, { "buf": "block_table" @@ -32239,12 +32227,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 458752 + "state": "kv.7" }, { - "state": "kv", - "offset": 459008 + "state": "kv.7", + "offset": 256 }, { "buf": "slot_mapping" @@ -32261,7 +32248,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -32276,12 +32263,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 458752 + "state": "kv.7" }, { - "state": "kv", - "offset": 459008 + "state": "kv.7", + "offset": 256 }, { "buf": "block_table" @@ -32544,12 +32530,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 524288 + "state": "kv.8" }, { - "state": "kv", - "offset": 524544 + "state": "kv.8", + "offset": 256 }, { "buf": "slot_mapping" @@ -32566,7 +32551,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -32581,12 +32566,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 524288 + "state": "kv.8" }, { - "state": "kv", - "offset": 524544 + "state": "kv.8", + "offset": 256 }, { "buf": "block_table" @@ -32849,12 +32833,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 589824 + "state": "kv.9" }, { - "state": "kv", - "offset": 590080 + "state": "kv.9", + "offset": 256 }, { "buf": "slot_mapping" @@ -32871,7 +32854,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -32886,12 +32869,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 589824 + "state": "kv.9" }, { - "state": "kv", - "offset": 590080 + "state": "kv.9", + "offset": 256 }, { "buf": "block_table" @@ -33154,12 +33136,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 655360 + "state": "kv.10" }, { - "state": "kv", - "offset": 655616 + "state": "kv.10", + "offset": 256 }, { "buf": "slot_mapping" @@ -33176,7 +33157,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -33191,12 +33172,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 655360 + "state": "kv.10" }, { - "state": "kv", - "offset": 655616 + "state": "kv.10", + "offset": 256 }, { "buf": "block_table" @@ -33459,12 +33439,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 720896 + "state": "kv.11" }, { - "state": "kv", - "offset": 721152 + "state": "kv.11", + "offset": 256 }, { "buf": "slot_mapping" @@ -33481,7 +33460,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -33496,12 +33475,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 720896 + "state": "kv.11" }, { - "state": "kv", - "offset": 721152 + "state": "kv.11", + "offset": 256 }, { "buf": "block_table" @@ -33764,12 +33742,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 786432 + "state": "kv.12" }, { - "state": "kv", - "offset": 786688 + "state": "kv.12", + "offset": 256 }, { "buf": "slot_mapping" @@ -33786,7 +33763,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -33801,12 +33778,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 786432 + "state": "kv.12" }, { - "state": "kv", - "offset": 786688 + "state": "kv.12", + "offset": 256 }, { "buf": "block_table" @@ -34069,12 +34045,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 851968 + "state": "kv.13" }, { - "state": "kv", - "offset": 852224 + "state": "kv.13", + "offset": 256 }, { "buf": "slot_mapping" @@ -34091,7 +34066,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -34106,12 +34081,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 851968 + "state": "kv.13" }, { - "state": "kv", - "offset": 852224 + "state": "kv.13", + "offset": 256 }, { "buf": "block_table" @@ -34374,12 +34348,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 917504 + "state": "kv.14" }, { - "state": "kv", - "offset": 917760 + "state": "kv.14", + "offset": 256 }, { "buf": "slot_mapping" @@ -34396,7 +34369,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -34411,12 +34384,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 917504 + "state": "kv.14" }, { - "state": "kv", - "offset": 917760 + "state": "kv.14", + "offset": 256 }, { "buf": "block_table" @@ -34679,12 +34651,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 983040 + "state": "kv.15" }, { - "state": "kv", - "offset": 983296 + "state": "kv.15", + "offset": 256 }, { "buf": "slot_mapping" @@ -34701,7 +34672,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -34716,12 +34687,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 983040 + "state": "kv.15" }, { - "state": "kv", - "offset": 983296 + "state": "kv.15", + "offset": 256 }, { "buf": "block_table" @@ -34984,12 +34954,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1048576 + "state": "kv.16" }, { - "state": "kv", - "offset": 1048832 + "state": "kv.16", + "offset": 256 }, { "buf": "slot_mapping" @@ -35006,7 +34975,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -35021,12 +34990,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1048576 + "state": "kv.16" }, { - "state": "kv", - "offset": 1048832 + "state": "kv.16", + "offset": 256 }, { "buf": "block_table" @@ -35289,12 +35257,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1114112 + "state": "kv.17" }, { - "state": "kv", - "offset": 1114368 + "state": "kv.17", + "offset": 256 }, { "buf": "slot_mapping" @@ -35311,7 +35278,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -35326,12 +35293,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1114112 + "state": "kv.17" }, { - "state": "kv", - "offset": 1114368 + "state": "kv.17", + "offset": 256 }, { "buf": "block_table" @@ -35594,12 +35560,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1179648 + "state": "kv.18" }, { - "state": "kv", - "offset": 1179904 + "state": "kv.18", + "offset": 256 }, { "buf": "slot_mapping" @@ -35616,7 +35581,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -35631,12 +35596,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1179648 + "state": "kv.18" }, { - "state": "kv", - "offset": 1179904 + "state": "kv.18", + "offset": 256 }, { "buf": "block_table" @@ -35899,12 +35863,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1245184 + "state": "kv.19" }, { - "state": "kv", - "offset": 1245440 + "state": "kv.19", + "offset": 256 }, { "buf": "slot_mapping" @@ -35921,7 +35884,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -35936,12 +35899,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1245184 + "state": "kv.19" }, { - "state": "kv", - "offset": 1245440 + "state": "kv.19", + "offset": 256 }, { "buf": "block_table" @@ -36204,12 +36166,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1310720 + "state": "kv.20" }, { - "state": "kv", - "offset": 1310976 + "state": "kv.20", + "offset": 256 }, { "buf": "slot_mapping" @@ -36226,7 +36187,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -36241,12 +36202,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1310720 + "state": "kv.20" }, { - "state": "kv", - "offset": 1310976 + "state": "kv.20", + "offset": 256 }, { "buf": "block_table" @@ -36509,12 +36469,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1376256 + "state": "kv.21" }, { - "state": "kv", - "offset": 1376512 + "state": "kv.21", + "offset": 256 }, { "buf": "slot_mapping" @@ -36531,7 +36490,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -36546,12 +36505,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1376256 + "state": "kv.21" }, { - "state": "kv", - "offset": 1376512 + "state": "kv.21", + "offset": 256 }, { "buf": "block_table" @@ -36814,12 +36772,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1441792 + "state": "kv.22" }, { - "state": "kv", - "offset": 1442048 + "state": "kv.22", + "offset": 256 }, { "buf": "slot_mapping" @@ -36836,7 +36793,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -36851,12 +36808,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1441792 + "state": "kv.22" }, { - "state": "kv", - "offset": 1442048 + "state": "kv.22", + "offset": 256 }, { "buf": "block_table" @@ -37119,12 +37075,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1507328 + "state": "kv.23" }, { - "state": "kv", - "offset": 1507584 + "state": "kv.23", + "offset": 256 }, { "buf": "slot_mapping" @@ -37141,7 +37096,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -37156,12 +37111,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1507328 + "state": "kv.23" }, { - "state": "kv", - "offset": 1507584 + "state": "kv.23", + "offset": 256 }, { "buf": "block_table" @@ -37424,12 +37378,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1572864 + "state": "kv.24" }, { - "state": "kv", - "offset": 1573120 + "state": "kv.24", + "offset": 256 }, { "buf": "slot_mapping" @@ -37446,7 +37399,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -37461,12 +37414,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1572864 + "state": "kv.24" }, { - "state": "kv", - "offset": 1573120 + "state": "kv.24", + "offset": 256 }, { "buf": "block_table" @@ -37729,12 +37681,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1638400 + "state": "kv.25" }, { - "state": "kv", - "offset": 1638656 + "state": "kv.25", + "offset": 256 }, { "buf": "slot_mapping" @@ -37751,7 +37702,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -37766,12 +37717,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1638400 + "state": "kv.25" }, { - "state": "kv", - "offset": 1638656 + "state": "kv.25", + "offset": 256 }, { "buf": "block_table" @@ -38034,12 +37984,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1703936 + "state": "kv.26" }, { - "state": "kv", - "offset": 1704192 + "state": "kv.26", + "offset": 256 }, { "buf": "slot_mapping" @@ -38056,7 +38005,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -38071,12 +38020,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1703936 + "state": "kv.26" }, { - "state": "kv", - "offset": 1704192 + "state": "kv.26", + "offset": 256 }, { "buf": "block_table" @@ -38339,12 +38287,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1769472 + "state": "kv.27" }, { - "state": "kv", - "offset": 1769728 + "state": "kv.27", + "offset": 256 }, { "buf": "slot_mapping" @@ -38361,7 +38308,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -38376,12 +38323,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1769472 + "state": "kv.27" }, { - "state": "kv", - "offset": 1769728 + "state": "kv.27", + "offset": 256 }, { "buf": "block_table" @@ -38644,12 +38590,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1835008 + "state": "kv.28" }, { - "state": "kv", - "offset": 1835264 + "state": "kv.28", + "offset": 256 }, { "buf": "slot_mapping" @@ -38666,7 +38611,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -38681,12 +38626,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1835008 + "state": "kv.28" }, { - "state": "kv", - "offset": 1835264 + "state": "kv.28", + "offset": 256 }, { "buf": "block_table" @@ -38949,12 +38893,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1900544 + "state": "kv.29" }, { - "state": "kv", - "offset": 1900800 + "state": "kv.29", + "offset": 256 }, { "buf": "slot_mapping" @@ -38971,7 +38914,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -38986,12 +38929,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1900544 + "state": "kv.29" }, { - "state": "kv", - "offset": 1900800 + "state": "kv.29", + "offset": 256 }, { "buf": "block_table" @@ -39254,12 +39196,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1966080 + "state": "kv.30" }, { - "state": "kv", - "offset": 1966336 + "state": "kv.30", + "offset": 256 }, { "buf": "slot_mapping" @@ -39276,7 +39217,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -39291,12 +39232,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1966080 + "state": "kv.30" }, { - "state": "kv", - "offset": 1966336 + "state": "kv.30", + "offset": 256 }, { "buf": "block_table" @@ -39559,12 +39499,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 2031616 + "state": "kv.31" }, { - "state": "kv", - "offset": 2031872 + "state": "kv.31", + "offset": 256 }, { "buf": "slot_mapping" @@ -39581,7 +39520,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -39596,12 +39535,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 2031616 + "state": "kv.31" }, { - "state": "kv", - "offset": 2031872 + "state": "kv.31", + "offset": 256 }, { "buf": "block_table" @@ -39864,12 +39802,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 2097152 + "state": "kv.32" }, { - "state": "kv", - "offset": 2097408 + "state": "kv.32", + "offset": 256 }, { "buf": "slot_mapping" @@ -39886,7 +39823,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -39901,12 +39838,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 2097152 + "state": "kv.32" }, { - "state": "kv", - "offset": 2097408 + "state": "kv.32", + "offset": 256 }, { "buf": "block_table" @@ -40169,12 +40105,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 2162688 + "state": "kv.33" }, { - "state": "kv", - "offset": 2162944 + "state": "kv.33", + "offset": 256 }, { "buf": "slot_mapping" @@ -40191,7 +40126,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -40206,12 +40141,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 2162688 + "state": "kv.33" }, { - "state": "kv", - "offset": 2162944 + "state": "kv.33", + "offset": 256 }, { "buf": "block_table" @@ -40474,12 +40408,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 2228224 + "state": "kv.34" }, { - "state": "kv", - "offset": 2228480 + "state": "kv.34", + "offset": 256 }, { "buf": "slot_mapping" @@ -40496,7 +40429,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -40511,12 +40444,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 2228224 + "state": "kv.34" }, { - "state": "kv", - "offset": 2228480 + "state": "kv.34", + "offset": 256 }, { "buf": "block_table" @@ -40779,12 +40711,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 2293760 + "state": "kv.35" }, { - "state": "kv", - "offset": 2294016 + "state": "kv.35", + "offset": 256 }, { "buf": "slot_mapping" @@ -40801,7 +40732,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -40816,12 +40747,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 2293760 + "state": "kv.35" }, { - "state": "kv", - "offset": 2294016 + "state": "kv.35", + "offset": 256 }, { "buf": "block_table" @@ -43132,10 +43062,10 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv" + "state": "kv.0" }, { - "state": "kv", + "state": "kv.0", "offset": 256 }, { @@ -43152,7 +43082,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -43167,10 +43097,10 @@ "buf": "q_n" }, { - "state": "kv" + "state": "kv.0" }, { - "state": "kv", + "state": "kv.0", "offset": 256 }, { @@ -43457,12 +43387,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 65536 + "state": "kv.1" }, { - "state": "kv", - "offset": 65792 + "state": "kv.1", + "offset": 256 }, { "buf": "slot_mapping" @@ -43479,7 +43408,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -43494,12 +43423,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 65536 + "state": "kv.1" }, { - "state": "kv", - "offset": 65792 + "state": "kv.1", + "offset": 256 }, { "buf": "block_table" @@ -43762,12 +43690,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 131072 + "state": "kv.2" }, { - "state": "kv", - "offset": 131328 + "state": "kv.2", + "offset": 256 }, { "buf": "slot_mapping" @@ -43784,7 +43711,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -43799,12 +43726,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 131072 + "state": "kv.2" }, { - "state": "kv", - "offset": 131328 + "state": "kv.2", + "offset": 256 }, { "buf": "block_table" @@ -44067,12 +43993,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 196608 + "state": "kv.3" }, { - "state": "kv", - "offset": 196864 + "state": "kv.3", + "offset": 256 }, { "buf": "slot_mapping" @@ -44089,7 +44014,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -44104,12 +44029,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 196608 + "state": "kv.3" }, { - "state": "kv", - "offset": 196864 + "state": "kv.3", + "offset": 256 }, { "buf": "block_table" @@ -44372,12 +44296,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 262144 + "state": "kv.4" }, { - "state": "kv", - "offset": 262400 + "state": "kv.4", + "offset": 256 }, { "buf": "slot_mapping" @@ -44394,7 +44317,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -44409,12 +44332,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 262144 + "state": "kv.4" }, { - "state": "kv", - "offset": 262400 + "state": "kv.4", + "offset": 256 }, { "buf": "block_table" @@ -44677,12 +44599,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 327680 + "state": "kv.5" }, { - "state": "kv", - "offset": 327936 + "state": "kv.5", + "offset": 256 }, { "buf": "slot_mapping" @@ -44699,7 +44620,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -44714,12 +44635,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 327680 + "state": "kv.5" }, { - "state": "kv", - "offset": 327936 + "state": "kv.5", + "offset": 256 }, { "buf": "block_table" @@ -44982,12 +44902,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 393216 + "state": "kv.6" }, { - "state": "kv", - "offset": 393472 + "state": "kv.6", + "offset": 256 }, { "buf": "slot_mapping" @@ -45004,7 +44923,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -45019,12 +44938,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 393216 + "state": "kv.6" }, { - "state": "kv", - "offset": 393472 + "state": "kv.6", + "offset": 256 }, { "buf": "block_table" @@ -45287,12 +45205,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 458752 + "state": "kv.7" }, { - "state": "kv", - "offset": 459008 + "state": "kv.7", + "offset": 256 }, { "buf": "slot_mapping" @@ -45309,7 +45226,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -45324,12 +45241,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 458752 + "state": "kv.7" }, { - "state": "kv", - "offset": 459008 + "state": "kv.7", + "offset": 256 }, { "buf": "block_table" @@ -45592,12 +45508,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 524288 + "state": "kv.8" }, { - "state": "kv", - "offset": 524544 + "state": "kv.8", + "offset": 256 }, { "buf": "slot_mapping" @@ -45614,7 +45529,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -45629,12 +45544,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 524288 + "state": "kv.8" }, { - "state": "kv", - "offset": 524544 + "state": "kv.8", + "offset": 256 }, { "buf": "block_table" @@ -45915,12 +45829,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 589824 + "state": "kv.9" }, { - "state": "kv", - "offset": 590080 + "state": "kv.9", + "offset": 256 }, { "buf": "slot_mapping" @@ -45937,7 +45850,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -45952,12 +45865,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 589824 + "state": "kv.9" }, { - "state": "kv", - "offset": 590080 + "state": "kv.9", + "offset": 256 }, { "buf": "block_table" @@ -46220,12 +46132,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 655360 + "state": "kv.10" }, { - "state": "kv", - "offset": 655616 + "state": "kv.10", + "offset": 256 }, { "buf": "slot_mapping" @@ -46242,7 +46153,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -46257,12 +46168,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 655360 + "state": "kv.10" }, { - "state": "kv", - "offset": 655616 + "state": "kv.10", + "offset": 256 }, { "buf": "block_table" @@ -46525,12 +46435,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 720896 + "state": "kv.11" }, { - "state": "kv", - "offset": 721152 + "state": "kv.11", + "offset": 256 }, { "buf": "slot_mapping" @@ -46547,7 +46456,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -46562,12 +46471,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 720896 + "state": "kv.11" }, { - "state": "kv", - "offset": 721152 + "state": "kv.11", + "offset": 256 }, { "buf": "block_table" @@ -46830,12 +46738,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 786432 + "state": "kv.12" }, { - "state": "kv", - "offset": 786688 + "state": "kv.12", + "offset": 256 }, { "buf": "slot_mapping" @@ -46852,7 +46759,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -46867,12 +46774,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 786432 + "state": "kv.12" }, { - "state": "kv", - "offset": 786688 + "state": "kv.12", + "offset": 256 }, { "buf": "block_table" @@ -47135,12 +47041,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 851968 + "state": "kv.13" }, { - "state": "kv", - "offset": 852224 + "state": "kv.13", + "offset": 256 }, { "buf": "slot_mapping" @@ -47157,7 +47062,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -47172,12 +47077,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 851968 + "state": "kv.13" }, { - "state": "kv", - "offset": 852224 + "state": "kv.13", + "offset": 256 }, { "buf": "block_table" @@ -47440,12 +47344,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 917504 + "state": "kv.14" }, { - "state": "kv", - "offset": 917760 + "state": "kv.14", + "offset": 256 }, { "buf": "slot_mapping" @@ -47462,7 +47365,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -47477,12 +47380,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 917504 + "state": "kv.14" }, { - "state": "kv", - "offset": 917760 + "state": "kv.14", + "offset": 256 }, { "buf": "block_table" @@ -47745,12 +47647,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 983040 + "state": "kv.15" }, { - "state": "kv", - "offset": 983296 + "state": "kv.15", + "offset": 256 }, { "buf": "slot_mapping" @@ -47767,7 +47668,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -47782,12 +47683,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 983040 + "state": "kv.15" }, { - "state": "kv", - "offset": 983296 + "state": "kv.15", + "offset": 256 }, { "buf": "block_table" @@ -48050,12 +47950,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1048576 + "state": "kv.16" }, { - "state": "kv", - "offset": 1048832 + "state": "kv.16", + "offset": 256 }, { "buf": "slot_mapping" @@ -48072,7 +47971,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -48087,12 +47986,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1048576 + "state": "kv.16" }, { - "state": "kv", - "offset": 1048832 + "state": "kv.16", + "offset": 256 }, { "buf": "block_table" @@ -48373,12 +48271,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1114112 + "state": "kv.17" }, { - "state": "kv", - "offset": 1114368 + "state": "kv.17", + "offset": 256 }, { "buf": "slot_mapping" @@ -48395,7 +48292,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -48410,12 +48307,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1114112 + "state": "kv.17" }, { - "state": "kv", - "offset": 1114368 + "state": "kv.17", + "offset": 256 }, { "buf": "block_table" @@ -48678,12 +48574,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1179648 + "state": "kv.18" }, { - "state": "kv", - "offset": 1179904 + "state": "kv.18", + "offset": 256 }, { "buf": "slot_mapping" @@ -48700,7 +48595,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -48715,12 +48610,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1179648 + "state": "kv.18" }, { - "state": "kv", - "offset": 1179904 + "state": "kv.18", + "offset": 256 }, { "buf": "block_table" @@ -48983,12 +48877,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1245184 + "state": "kv.19" }, { - "state": "kv", - "offset": 1245440 + "state": "kv.19", + "offset": 256 }, { "buf": "slot_mapping" @@ -49005,7 +48898,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -49020,12 +48913,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1245184 + "state": "kv.19" }, { - "state": "kv", - "offset": 1245440 + "state": "kv.19", + "offset": 256 }, { "buf": "block_table" @@ -49288,12 +49180,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1310720 + "state": "kv.20" }, { - "state": "kv", - "offset": 1310976 + "state": "kv.20", + "offset": 256 }, { "buf": "slot_mapping" @@ -49310,7 +49201,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -49325,12 +49216,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1310720 + "state": "kv.20" }, { - "state": "kv", - "offset": 1310976 + "state": "kv.20", + "offset": 256 }, { "buf": "block_table" @@ -49593,12 +49483,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1376256 + "state": "kv.21" }, { - "state": "kv", - "offset": 1376512 + "state": "kv.21", + "offset": 256 }, { "buf": "slot_mapping" @@ -49615,7 +49504,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -49630,12 +49519,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1376256 + "state": "kv.21" }, { - "state": "kv", - "offset": 1376512 + "state": "kv.21", + "offset": 256 }, { "buf": "block_table" @@ -49898,12 +49786,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1441792 + "state": "kv.22" }, { - "state": "kv", - "offset": 1442048 + "state": "kv.22", + "offset": 256 }, { "buf": "slot_mapping" @@ -49920,7 +49807,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -49935,12 +49822,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1441792 + "state": "kv.22" }, { - "state": "kv", - "offset": 1442048 + "state": "kv.22", + "offset": 256 }, { "buf": "block_table" @@ -50203,12 +50089,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1507328 + "state": "kv.23" }, { - "state": "kv", - "offset": 1507584 + "state": "kv.23", + "offset": 256 }, { "buf": "slot_mapping" @@ -50225,7 +50110,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -50240,12 +50125,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1507328 + "state": "kv.23" }, { - "state": "kv", - "offset": 1507584 + "state": "kv.23", + "offset": 256 }, { "buf": "block_table" @@ -50508,12 +50392,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1572864 + "state": "kv.24" }, { - "state": "kv", - "offset": 1573120 + "state": "kv.24", + "offset": 256 }, { "buf": "slot_mapping" @@ -50530,7 +50413,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -50545,12 +50428,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1572864 + "state": "kv.24" }, { - "state": "kv", - "offset": 1573120 + "state": "kv.24", + "offset": 256 }, { "buf": "block_table" @@ -50831,12 +50713,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1638400 + "state": "kv.25" }, { - "state": "kv", - "offset": 1638656 + "state": "kv.25", + "offset": 256 }, { "buf": "slot_mapping" @@ -50853,7 +50734,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -50868,12 +50749,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1638400 + "state": "kv.25" }, { - "state": "kv", - "offset": 1638656 + "state": "kv.25", + "offset": 256 }, { "buf": "block_table" @@ -51136,12 +51016,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1703936 + "state": "kv.26" }, { - "state": "kv", - "offset": 1704192 + "state": "kv.26", + "offset": 256 }, { "buf": "slot_mapping" @@ -51158,7 +51037,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -51173,12 +51052,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1703936 + "state": "kv.26" }, { - "state": "kv", - "offset": 1704192 + "state": "kv.26", + "offset": 256 }, { "buf": "block_table" @@ -51441,12 +51319,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1769472 + "state": "kv.27" }, { - "state": "kv", - "offset": 1769728 + "state": "kv.27", + "offset": 256 }, { "buf": "slot_mapping" @@ -51463,7 +51340,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -51478,12 +51355,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1769472 + "state": "kv.27" }, { - "state": "kv", - "offset": 1769728 + "state": "kv.27", + "offset": 256 }, { "buf": "block_table" @@ -51746,12 +51622,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1835008 + "state": "kv.28" }, { - "state": "kv", - "offset": 1835264 + "state": "kv.28", + "offset": 256 }, { "buf": "slot_mapping" @@ -51768,7 +51643,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -51783,12 +51658,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1835008 + "state": "kv.28" }, { - "state": "kv", - "offset": 1835264 + "state": "kv.28", + "offset": 256 }, { "buf": "block_table" @@ -52051,12 +51925,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1900544 + "state": "kv.29" }, { - "state": "kv", - "offset": 1900800 + "state": "kv.29", + "offset": 256 }, { "buf": "slot_mapping" @@ -52073,7 +51946,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -52088,12 +51961,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1900544 + "state": "kv.29" }, { - "state": "kv", - "offset": 1900800 + "state": "kv.29", + "offset": 256 }, { "buf": "block_table" @@ -52356,12 +52228,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1966080 + "state": "kv.30" }, { - "state": "kv", - "offset": 1966336 + "state": "kv.30", + "offset": 256 }, { "buf": "slot_mapping" @@ -52378,7 +52249,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -52393,12 +52264,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1966080 + "state": "kv.30" }, { - "state": "kv", - "offset": 1966336 + "state": "kv.30", + "offset": 256 }, { "buf": "block_table" @@ -52661,12 +52531,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 2031616 + "state": "kv.31" }, { - "state": "kv", - "offset": 2031872 + "state": "kv.31", + "offset": 256 }, { "buf": "slot_mapping" @@ -52683,7 +52552,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -52698,12 +52567,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 2031616 + "state": "kv.31" }, { - "state": "kv", - "offset": 2031872 + "state": "kv.31", + "offset": 256 }, { "buf": "block_table" @@ -52966,12 +52834,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 2097152 + "state": "kv.32" }, { - "state": "kv", - "offset": 2097408 + "state": "kv.32", + "offset": 256 }, { "buf": "slot_mapping" @@ -52988,7 +52855,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -53003,12 +52870,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 2097152 + "state": "kv.32" }, { - "state": "kv", - "offset": 2097408 + "state": "kv.32", + "offset": 256 }, { "buf": "block_table" @@ -53289,12 +53155,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 2162688 + "state": "kv.33" }, { - "state": "kv", - "offset": 2162944 + "state": "kv.33", + "offset": 256 }, { "buf": "slot_mapping" @@ -53311,7 +53176,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -53326,12 +53191,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 2162688 + "state": "kv.33" }, { - "state": "kv", - "offset": 2162944 + "state": "kv.33", + "offset": 256 }, { "buf": "block_table" @@ -53594,12 +53458,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 2228224 + "state": "kv.34" }, { - "state": "kv", - "offset": 2228480 + "state": "kv.34", + "offset": 256 }, { "buf": "slot_mapping" @@ -53616,7 +53479,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -53631,12 +53494,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 2228224 + "state": "kv.34" }, { - "state": "kv", - "offset": 2228480 + "state": "kv.34", + "offset": 256 }, { "buf": "block_table" @@ -53899,12 +53761,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 2293760 + "state": "kv.35" }, { - "state": "kv", - "offset": 2294016 + "state": "kv.35", + "offset": 256 }, { "buf": "slot_mapping" @@ -53921,7 +53782,7 @@ "i64": 6144 }, { - "i64": 1179648 + "i64": 32768 } ] }, @@ -53936,12 +53797,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 2293760 + "state": "kv.35" }, { - "state": "kv", - "offset": 2294016 + "state": "kv.35", + "offset": 256 }, { "buf": "block_table" diff --git a/examples/qwen3-4b-silu-mined.json b/examples/qwen3-4b-silu-mined.json index a607f8e9..48086007 100644 --- a/examples/qwen3-4b-silu-mined.json +++ b/examples/qwen3-4b-silu-mined.json @@ -17,7 +17,42 @@ "KV_SCALE_COUNT": 72, "MAX_TOKENS": 2048, "MAX_SEQS": 256, - "KV_BYTES_PER_TOKEN": 147456, + "KV.0_BYTES_PER_TOKEN": 4096, + "KV.1_BYTES_PER_TOKEN": 4096, + "KV.2_BYTES_PER_TOKEN": 4096, + "KV.3_BYTES_PER_TOKEN": 4096, + "KV.4_BYTES_PER_TOKEN": 4096, + "KV.5_BYTES_PER_TOKEN": 4096, + "KV.6_BYTES_PER_TOKEN": 4096, + "KV.7_BYTES_PER_TOKEN": 4096, + "KV.8_BYTES_PER_TOKEN": 4096, + "KV.9_BYTES_PER_TOKEN": 4096, + "KV.10_BYTES_PER_TOKEN": 4096, + "KV.11_BYTES_PER_TOKEN": 4096, + "KV.12_BYTES_PER_TOKEN": 4096, + "KV.13_BYTES_PER_TOKEN": 4096, + "KV.14_BYTES_PER_TOKEN": 4096, + "KV.15_BYTES_PER_TOKEN": 4096, + "KV.16_BYTES_PER_TOKEN": 4096, + "KV.17_BYTES_PER_TOKEN": 4096, + "KV.18_BYTES_PER_TOKEN": 4096, + "KV.19_BYTES_PER_TOKEN": 4096, + "KV.20_BYTES_PER_TOKEN": 4096, + "KV.21_BYTES_PER_TOKEN": 4096, + "KV.22_BYTES_PER_TOKEN": 4096, + "KV.23_BYTES_PER_TOKEN": 4096, + "KV.24_BYTES_PER_TOKEN": 4096, + "KV.25_BYTES_PER_TOKEN": 4096, + "KV.26_BYTES_PER_TOKEN": 4096, + "KV.27_BYTES_PER_TOKEN": 4096, + "KV.28_BYTES_PER_TOKEN": 4096, + "KV.29_BYTES_PER_TOKEN": 4096, + "KV.30_BYTES_PER_TOKEN": 4096, + "KV.31_BYTES_PER_TOKEN": 4096, + "KV.32_BYTES_PER_TOKEN": 4096, + "KV.33_BYTES_PER_TOKEN": 4096, + "KV.34_BYTES_PER_TOKEN": 4096, + "KV.35_BYTES_PER_TOKEN": 4096, "RMS_NORM_EPS": 9.999999974752427e-07, "NUM_HEADS": 32, "NUM_KV_HEADS": 8, @@ -34,8 +69,113 @@ } }, "states": { - "kv": { - "bytes_per_token": "KV_BYTES_PER_TOKEN" + "kv.0": { + "bytes_per_token": "KV.0_BYTES_PER_TOKEN" + }, + "kv.1": { + "bytes_per_token": "KV.1_BYTES_PER_TOKEN" + }, + "kv.2": { + "bytes_per_token": "KV.2_BYTES_PER_TOKEN" + }, + "kv.3": { + "bytes_per_token": "KV.3_BYTES_PER_TOKEN" + }, + "kv.4": { + "bytes_per_token": "KV.4_BYTES_PER_TOKEN" + }, + "kv.5": { + "bytes_per_token": "KV.5_BYTES_PER_TOKEN" + }, + "kv.6": { + "bytes_per_token": "KV.6_BYTES_PER_TOKEN" + }, + "kv.7": { + "bytes_per_token": "KV.7_BYTES_PER_TOKEN" + }, + "kv.8": { + "bytes_per_token": "KV.8_BYTES_PER_TOKEN" + }, + "kv.9": { + "bytes_per_token": "KV.9_BYTES_PER_TOKEN" + }, + "kv.10": { + "bytes_per_token": "KV.10_BYTES_PER_TOKEN" + }, + "kv.11": { + "bytes_per_token": "KV.11_BYTES_PER_TOKEN" + }, + "kv.12": { + "bytes_per_token": "KV.12_BYTES_PER_TOKEN" + }, + "kv.13": { + "bytes_per_token": "KV.13_BYTES_PER_TOKEN" + }, + "kv.14": { + "bytes_per_token": "KV.14_BYTES_PER_TOKEN" + }, + "kv.15": { + "bytes_per_token": "KV.15_BYTES_PER_TOKEN" + }, + "kv.16": { + "bytes_per_token": "KV.16_BYTES_PER_TOKEN" + }, + "kv.17": { + "bytes_per_token": "KV.17_BYTES_PER_TOKEN" + }, + "kv.18": { + "bytes_per_token": "KV.18_BYTES_PER_TOKEN" + }, + "kv.19": { + "bytes_per_token": "KV.19_BYTES_PER_TOKEN" + }, + "kv.20": { + "bytes_per_token": "KV.20_BYTES_PER_TOKEN" + }, + "kv.21": { + "bytes_per_token": "KV.21_BYTES_PER_TOKEN" + }, + "kv.22": { + "bytes_per_token": "KV.22_BYTES_PER_TOKEN" + }, + "kv.23": { + "bytes_per_token": "KV.23_BYTES_PER_TOKEN" + }, + "kv.24": { + "bytes_per_token": "KV.24_BYTES_PER_TOKEN" + }, + "kv.25": { + "bytes_per_token": "KV.25_BYTES_PER_TOKEN" + }, + "kv.26": { + "bytes_per_token": "KV.26_BYTES_PER_TOKEN" + }, + "kv.27": { + "bytes_per_token": "KV.27_BYTES_PER_TOKEN" + }, + "kv.28": { + "bytes_per_token": "KV.28_BYTES_PER_TOKEN" + }, + "kv.29": { + "bytes_per_token": "KV.29_BYTES_PER_TOKEN" + }, + "kv.30": { + "bytes_per_token": "KV.30_BYTES_PER_TOKEN" + }, + "kv.31": { + "bytes_per_token": "KV.31_BYTES_PER_TOKEN" + }, + "kv.32": { + "bytes_per_token": "KV.32_BYTES_PER_TOKEN" + }, + "kv.33": { + "bytes_per_token": "KV.33_BYTES_PER_TOKEN" + }, + "kv.34": { + "bytes_per_token": "KV.34_BYTES_PER_TOKEN" + }, + "kv.35": { + "bytes_per_token": "KV.35_BYTES_PER_TOKEN" } }, "buffers": { @@ -69,7 +209,7 @@ "kind": "input", "fill": "slot", "domain": { - "index_into": "kv" + "index_into": "kv.0" } }, "block_table": { @@ -80,7 +220,7 @@ ], "kind": "input", "domain": { - "index_into": "kv", + "index_into": "kv.0", "stride": "KV_BLOCK_SIZE" } }, @@ -4182,12 +4322,12 @@ }, "modules": { "argmax": { - "source": "argmax.cubin", - "sha256": "25374575bd35dbd9fe07f90982d0da0c2d4a72cc6ac88df5ec92cdb62a3beb16" + "source": "hf:Pegainfer/kern-kernels/blobs/ebe4ba79544ddfac43ee2dcf6a7f5299aba1124e176c201133d4ef9400e3dc25", + "sha256": "ebe4ba79544ddfac43ee2dcf6a7f5299aba1124e176c201133d4ef9400e3dc25" }, "embedding": { - "source": "embedding.cubin", - "sha256": "1c3dd3e591d872aed259d7ef59d44649061f77a4247a5df864a7fbf0ccbde92f" + "source": "hf:Pegainfer/kern-kernels/blobs/868e93ae6e46e36a11c56af1fae49aeff9fcdfdc79f64ca99475a309e721e9e7", + "sha256": "868e93ae6e46e36a11c56af1fae49aeff9fcdfdc79f64ca99475a309e721e9e7" }, "reduce_segments": { "source": "reduce_segments.cubin", @@ -4218,7 +4358,7 @@ "sha256": "116de812ac5a3a186f76fa239d4c16abbed6b8b0af06d16e1baec0aefe8fe8fa" }, "weight_prep": { - "source": "weight_prep.cubin", + "source": "hf:Pegainfer/kern-kernels/blobs/040f51b01c421b9c69a2090034ff1bc4741e40c02f5ea52bd04a3e62e5c1af4d", "sha256": "040f51b01c421b9c69a2090034ff1bc4741e40c02f5ea52bd04a3e62e5c1af4d" } }, @@ -4621,7 +4761,7 @@ "i64": "QKV_DIM" }, { - "i64": 1179648 + "i64": 32768 }, { "i64": 256 @@ -4791,7 +4931,7 @@ "param": 8 }, { - "i64": 1179648 + "i64": 32768 }, { "i64": 2048 @@ -4800,7 +4940,7 @@ "i64": 256 }, { - "i64": 1179648 + "i64": 32768 }, { "i64": 2048 @@ -5022,7 +5162,7 @@ "param": 8 }, { - "i64": 1179648 + "i64": 32768 }, { "i64": 2048 @@ -5031,7 +5171,7 @@ "i64": 256 }, { - "i64": 1179648 + "i64": 32768 }, { "i64": 2048 @@ -5182,7 +5322,7 @@ "param": 8 }, { - "i64": 1179648 + "i64": 32768 }, { "i64": 2048 @@ -5191,7 +5331,7 @@ "i64": 256 }, { - "i64": 1179648 + "i64": 32768 }, { "i64": 2048 @@ -5727,10 +5867,10 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv" + "state": "kv.0" }, { - "state": "kv", + "state": "kv.0", "offset": 256 }, { @@ -5756,10 +5896,10 @@ "buf": "q_n" }, { - "state": "kv" + "state": "kv.0" }, { - "state": "kv", + "state": "kv.0", "offset": 256 }, { @@ -6007,12 +6147,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 65536 + "state": "kv.1" }, { - "state": "kv", - "offset": 65792 + "state": "kv.1", + "offset": 256 }, { "buf": "slot_mapping" @@ -6038,12 +6177,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 65536 + "state": "kv.1" }, { - "state": "kv", - "offset": 65792 + "state": "kv.1", + "offset": 256 }, { "buf": "block_table" @@ -6291,12 +6429,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 131072 + "state": "kv.2" }, { - "state": "kv", - "offset": 131328 + "state": "kv.2", + "offset": 256 }, { "buf": "slot_mapping" @@ -6322,12 +6459,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 131072 + "state": "kv.2" }, { - "state": "kv", - "offset": 131328 + "state": "kv.2", + "offset": 256 }, { "buf": "block_table" @@ -6575,12 +6711,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 196608 + "state": "kv.3" }, { - "state": "kv", - "offset": 196864 + "state": "kv.3", + "offset": 256 }, { "buf": "slot_mapping" @@ -6606,12 +6741,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 196608 + "state": "kv.3" }, { - "state": "kv", - "offset": 196864 + "state": "kv.3", + "offset": 256 }, { "buf": "block_table" @@ -6859,12 +6993,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 262144 + "state": "kv.4" }, { - "state": "kv", - "offset": 262400 + "state": "kv.4", + "offset": 256 }, { "buf": "slot_mapping" @@ -6890,12 +7023,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 262144 + "state": "kv.4" }, { - "state": "kv", - "offset": 262400 + "state": "kv.4", + "offset": 256 }, { "buf": "block_table" @@ -7143,12 +7275,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 327680 + "state": "kv.5" }, { - "state": "kv", - "offset": 327936 + "state": "kv.5", + "offset": 256 }, { "buf": "slot_mapping" @@ -7174,12 +7305,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 327680 + "state": "kv.5" }, { - "state": "kv", - "offset": 327936 + "state": "kv.5", + "offset": 256 }, { "buf": "block_table" @@ -7427,12 +7557,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 393216 + "state": "kv.6" }, { - "state": "kv", - "offset": 393472 + "state": "kv.6", + "offset": 256 }, { "buf": "slot_mapping" @@ -7458,12 +7587,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 393216 + "state": "kv.6" }, { - "state": "kv", - "offset": 393472 + "state": "kv.6", + "offset": 256 }, { "buf": "block_table" @@ -7711,12 +7839,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 458752 + "state": "kv.7" }, { - "state": "kv", - "offset": 459008 + "state": "kv.7", + "offset": 256 }, { "buf": "slot_mapping" @@ -7742,12 +7869,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 458752 + "state": "kv.7" }, { - "state": "kv", - "offset": 459008 + "state": "kv.7", + "offset": 256 }, { "buf": "block_table" @@ -7995,12 +8121,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 524288 + "state": "kv.8" }, { - "state": "kv", - "offset": 524544 + "state": "kv.8", + "offset": 256 }, { "buf": "slot_mapping" @@ -8026,12 +8151,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 524288 + "state": "kv.8" }, { - "state": "kv", - "offset": 524544 + "state": "kv.8", + "offset": 256 }, { "buf": "block_table" @@ -8279,12 +8403,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 589824 + "state": "kv.9" }, { - "state": "kv", - "offset": 590080 + "state": "kv.9", + "offset": 256 }, { "buf": "slot_mapping" @@ -8310,12 +8433,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 589824 + "state": "kv.9" }, { - "state": "kv", - "offset": 590080 + "state": "kv.9", + "offset": 256 }, { "buf": "block_table" @@ -8563,12 +8685,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 655360 + "state": "kv.10" }, { - "state": "kv", - "offset": 655616 + "state": "kv.10", + "offset": 256 }, { "buf": "slot_mapping" @@ -8594,12 +8715,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 655360 + "state": "kv.10" }, { - "state": "kv", - "offset": 655616 + "state": "kv.10", + "offset": 256 }, { "buf": "block_table" @@ -8847,12 +8967,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 720896 + "state": "kv.11" }, { - "state": "kv", - "offset": 721152 + "state": "kv.11", + "offset": 256 }, { "buf": "slot_mapping" @@ -8878,12 +8997,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 720896 + "state": "kv.11" }, { - "state": "kv", - "offset": 721152 + "state": "kv.11", + "offset": 256 }, { "buf": "block_table" @@ -9131,12 +9249,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 786432 + "state": "kv.12" }, { - "state": "kv", - "offset": 786688 + "state": "kv.12", + "offset": 256 }, { "buf": "slot_mapping" @@ -9162,12 +9279,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 786432 + "state": "kv.12" }, { - "state": "kv", - "offset": 786688 + "state": "kv.12", + "offset": 256 }, { "buf": "block_table" @@ -9415,12 +9531,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 851968 + "state": "kv.13" }, { - "state": "kv", - "offset": 852224 + "state": "kv.13", + "offset": 256 }, { "buf": "slot_mapping" @@ -9446,12 +9561,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 851968 + "state": "kv.13" }, { - "state": "kv", - "offset": 852224 + "state": "kv.13", + "offset": 256 }, { "buf": "block_table" @@ -9699,12 +9813,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 917504 + "state": "kv.14" }, { - "state": "kv", - "offset": 917760 + "state": "kv.14", + "offset": 256 }, { "buf": "slot_mapping" @@ -9730,12 +9843,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 917504 + "state": "kv.14" }, { - "state": "kv", - "offset": 917760 + "state": "kv.14", + "offset": 256 }, { "buf": "block_table" @@ -9983,12 +10095,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 983040 + "state": "kv.15" }, { - "state": "kv", - "offset": 983296 + "state": "kv.15", + "offset": 256 }, { "buf": "slot_mapping" @@ -10014,12 +10125,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 983040 + "state": "kv.15" }, { - "state": "kv", - "offset": 983296 + "state": "kv.15", + "offset": 256 }, { "buf": "block_table" @@ -10267,12 +10377,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1048576 + "state": "kv.16" }, { - "state": "kv", - "offset": 1048832 + "state": "kv.16", + "offset": 256 }, { "buf": "slot_mapping" @@ -10298,12 +10407,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1048576 + "state": "kv.16" }, { - "state": "kv", - "offset": 1048832 + "state": "kv.16", + "offset": 256 }, { "buf": "block_table" @@ -10551,12 +10659,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1114112 + "state": "kv.17" }, { - "state": "kv", - "offset": 1114368 + "state": "kv.17", + "offset": 256 }, { "buf": "slot_mapping" @@ -10582,12 +10689,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1114112 + "state": "kv.17" }, { - "state": "kv", - "offset": 1114368 + "state": "kv.17", + "offset": 256 }, { "buf": "block_table" @@ -10835,12 +10941,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1179648 + "state": "kv.18" }, { - "state": "kv", - "offset": 1179904 + "state": "kv.18", + "offset": 256 }, { "buf": "slot_mapping" @@ -10866,12 +10971,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1179648 + "state": "kv.18" }, { - "state": "kv", - "offset": 1179904 + "state": "kv.18", + "offset": 256 }, { "buf": "block_table" @@ -11119,12 +11223,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1245184 + "state": "kv.19" }, { - "state": "kv", - "offset": 1245440 + "state": "kv.19", + "offset": 256 }, { "buf": "slot_mapping" @@ -11150,12 +11253,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1245184 + "state": "kv.19" }, { - "state": "kv", - "offset": 1245440 + "state": "kv.19", + "offset": 256 }, { "buf": "block_table" @@ -11403,12 +11505,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1310720 + "state": "kv.20" }, { - "state": "kv", - "offset": 1310976 + "state": "kv.20", + "offset": 256 }, { "buf": "slot_mapping" @@ -11434,12 +11535,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1310720 + "state": "kv.20" }, { - "state": "kv", - "offset": 1310976 + "state": "kv.20", + "offset": 256 }, { "buf": "block_table" @@ -11687,12 +11787,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1376256 + "state": "kv.21" }, { - "state": "kv", - "offset": 1376512 + "state": "kv.21", + "offset": 256 }, { "buf": "slot_mapping" @@ -11718,12 +11817,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1376256 + "state": "kv.21" }, { - "state": "kv", - "offset": 1376512 + "state": "kv.21", + "offset": 256 }, { "buf": "block_table" @@ -11971,12 +12069,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1441792 + "state": "kv.22" }, { - "state": "kv", - "offset": 1442048 + "state": "kv.22", + "offset": 256 }, { "buf": "slot_mapping" @@ -12002,12 +12099,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1441792 + "state": "kv.22" }, { - "state": "kv", - "offset": 1442048 + "state": "kv.22", + "offset": 256 }, { "buf": "block_table" @@ -12255,12 +12351,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1507328 + "state": "kv.23" }, { - "state": "kv", - "offset": 1507584 + "state": "kv.23", + "offset": 256 }, { "buf": "slot_mapping" @@ -12286,12 +12381,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1507328 + "state": "kv.23" }, { - "state": "kv", - "offset": 1507584 + "state": "kv.23", + "offset": 256 }, { "buf": "block_table" @@ -12539,12 +12633,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1572864 + "state": "kv.24" }, { - "state": "kv", - "offset": 1573120 + "state": "kv.24", + "offset": 256 }, { "buf": "slot_mapping" @@ -12570,12 +12663,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1572864 + "state": "kv.24" }, { - "state": "kv", - "offset": 1573120 + "state": "kv.24", + "offset": 256 }, { "buf": "block_table" @@ -12823,12 +12915,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1638400 + "state": "kv.25" }, { - "state": "kv", - "offset": 1638656 + "state": "kv.25", + "offset": 256 }, { "buf": "slot_mapping" @@ -12854,12 +12945,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1638400 + "state": "kv.25" }, { - "state": "kv", - "offset": 1638656 + "state": "kv.25", + "offset": 256 }, { "buf": "block_table" @@ -13107,12 +13197,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1703936 + "state": "kv.26" }, { - "state": "kv", - "offset": 1704192 + "state": "kv.26", + "offset": 256 }, { "buf": "slot_mapping" @@ -13138,12 +13227,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1703936 + "state": "kv.26" }, { - "state": "kv", - "offset": 1704192 + "state": "kv.26", + "offset": 256 }, { "buf": "block_table" @@ -13391,12 +13479,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1769472 + "state": "kv.27" }, { - "state": "kv", - "offset": 1769728 + "state": "kv.27", + "offset": 256 }, { "buf": "slot_mapping" @@ -13422,12 +13509,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1769472 + "state": "kv.27" }, { - "state": "kv", - "offset": 1769728 + "state": "kv.27", + "offset": 256 }, { "buf": "block_table" @@ -13675,12 +13761,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1835008 + "state": "kv.28" }, { - "state": "kv", - "offset": 1835264 + "state": "kv.28", + "offset": 256 }, { "buf": "slot_mapping" @@ -13706,12 +13791,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1835008 + "state": "kv.28" }, { - "state": "kv", - "offset": 1835264 + "state": "kv.28", + "offset": 256 }, { "buf": "block_table" @@ -13959,12 +14043,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1900544 + "state": "kv.29" }, { - "state": "kv", - "offset": 1900800 + "state": "kv.29", + "offset": 256 }, { "buf": "slot_mapping" @@ -13990,12 +14073,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1900544 + "state": "kv.29" }, { - "state": "kv", - "offset": 1900800 + "state": "kv.29", + "offset": 256 }, { "buf": "block_table" @@ -14243,12 +14325,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1966080 + "state": "kv.30" }, { - "state": "kv", - "offset": 1966336 + "state": "kv.30", + "offset": 256 }, { "buf": "slot_mapping" @@ -14274,12 +14355,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1966080 + "state": "kv.30" }, { - "state": "kv", - "offset": 1966336 + "state": "kv.30", + "offset": 256 }, { "buf": "block_table" @@ -14527,12 +14607,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 2031616 + "state": "kv.31" }, { - "state": "kv", - "offset": 2031872 + "state": "kv.31", + "offset": 256 }, { "buf": "slot_mapping" @@ -14558,12 +14637,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 2031616 + "state": "kv.31" }, { - "state": "kv", - "offset": 2031872 + "state": "kv.31", + "offset": 256 }, { "buf": "block_table" @@ -14811,12 +14889,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 2097152 + "state": "kv.32" }, { - "state": "kv", - "offset": 2097408 + "state": "kv.32", + "offset": 256 }, { "buf": "slot_mapping" @@ -14842,12 +14919,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 2097152 + "state": "kv.32" }, { - "state": "kv", - "offset": 2097408 + "state": "kv.32", + "offset": 256 }, { "buf": "block_table" @@ -15095,12 +15171,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 2162688 + "state": "kv.33" }, { - "state": "kv", - "offset": 2162944 + "state": "kv.33", + "offset": 256 }, { "buf": "slot_mapping" @@ -15126,12 +15201,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 2162688 + "state": "kv.33" }, { - "state": "kv", - "offset": 2162944 + "state": "kv.33", + "offset": 256 }, { "buf": "block_table" @@ -15379,12 +15453,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 2228224 + "state": "kv.34" }, { - "state": "kv", - "offset": 2228480 + "state": "kv.34", + "offset": 256 }, { "buf": "slot_mapping" @@ -15410,12 +15483,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 2228224 + "state": "kv.34" }, { - "state": "kv", - "offset": 2228480 + "state": "kv.34", + "offset": 256 }, { "buf": "block_table" @@ -15663,12 +15735,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 2293760 + "state": "kv.35" }, { - "state": "kv", - "offset": 2294016 + "state": "kv.35", + "offset": 256 }, { "buf": "slot_mapping" @@ -15694,12 +15765,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 2293760 + "state": "kv.35" }, { - "state": "kv", - "offset": 2294016 + "state": "kv.35", + "offset": 256 }, { "buf": "block_table" @@ -15974,10 +16044,10 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv" + "state": "kv.0" }, { - "state": "kv", + "state": "kv.0", "offset": 256 }, { @@ -16003,10 +16073,10 @@ "buf": "q_n" }, { - "state": "kv" + "state": "kv.0" }, { - "state": "kv", + "state": "kv.0", "offset": 256 }, { @@ -16251,12 +16321,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 65536 + "state": "kv.1" }, { - "state": "kv", - "offset": 65792 + "state": "kv.1", + "offset": 256 }, { "buf": "slot_mapping" @@ -16282,12 +16351,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 65536 + "state": "kv.1" }, { - "state": "kv", - "offset": 65792 + "state": "kv.1", + "offset": 256 }, { "buf": "block_table" @@ -16532,12 +16600,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 131072 + "state": "kv.2" }, { - "state": "kv", - "offset": 131328 + "state": "kv.2", + "offset": 256 }, { "buf": "slot_mapping" @@ -16563,12 +16630,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 131072 + "state": "kv.2" }, { - "state": "kv", - "offset": 131328 + "state": "kv.2", + "offset": 256 }, { "buf": "block_table" @@ -16813,12 +16879,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 196608 + "state": "kv.3" }, { - "state": "kv", - "offset": 196864 + "state": "kv.3", + "offset": 256 }, { "buf": "slot_mapping" @@ -16844,12 +16909,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 196608 + "state": "kv.3" }, { - "state": "kv", - "offset": 196864 + "state": "kv.3", + "offset": 256 }, { "buf": "block_table" @@ -17094,12 +17158,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 262144 + "state": "kv.4" }, { - "state": "kv", - "offset": 262400 + "state": "kv.4", + "offset": 256 }, { "buf": "slot_mapping" @@ -17125,12 +17188,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 262144 + "state": "kv.4" }, { - "state": "kv", - "offset": 262400 + "state": "kv.4", + "offset": 256 }, { "buf": "block_table" @@ -17375,12 +17437,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 327680 + "state": "kv.5" }, { - "state": "kv", - "offset": 327936 + "state": "kv.5", + "offset": 256 }, { "buf": "slot_mapping" @@ -17406,12 +17467,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 327680 + "state": "kv.5" }, { - "state": "kv", - "offset": 327936 + "state": "kv.5", + "offset": 256 }, { "buf": "block_table" @@ -17656,12 +17716,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 393216 + "state": "kv.6" }, { - "state": "kv", - "offset": 393472 + "state": "kv.6", + "offset": 256 }, { "buf": "slot_mapping" @@ -17687,12 +17746,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 393216 + "state": "kv.6" }, { - "state": "kv", - "offset": 393472 + "state": "kv.6", + "offset": 256 }, { "buf": "block_table" @@ -17937,12 +17995,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 458752 + "state": "kv.7" }, { - "state": "kv", - "offset": 459008 + "state": "kv.7", + "offset": 256 }, { "buf": "slot_mapping" @@ -17968,12 +18025,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 458752 + "state": "kv.7" }, { - "state": "kv", - "offset": 459008 + "state": "kv.7", + "offset": 256 }, { "buf": "block_table" @@ -18218,12 +18274,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 524288 + "state": "kv.8" }, { - "state": "kv", - "offset": 524544 + "state": "kv.8", + "offset": 256 }, { "buf": "slot_mapping" @@ -18249,12 +18304,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 524288 + "state": "kv.8" }, { - "state": "kv", - "offset": 524544 + "state": "kv.8", + "offset": 256 }, { "buf": "block_table" @@ -18499,12 +18553,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 589824 + "state": "kv.9" }, { - "state": "kv", - "offset": 590080 + "state": "kv.9", + "offset": 256 }, { "buf": "slot_mapping" @@ -18530,12 +18583,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 589824 + "state": "kv.9" }, { - "state": "kv", - "offset": 590080 + "state": "kv.9", + "offset": 256 }, { "buf": "block_table" @@ -18780,12 +18832,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 655360 + "state": "kv.10" }, { - "state": "kv", - "offset": 655616 + "state": "kv.10", + "offset": 256 }, { "buf": "slot_mapping" @@ -18811,12 +18862,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 655360 + "state": "kv.10" }, { - "state": "kv", - "offset": 655616 + "state": "kv.10", + "offset": 256 }, { "buf": "block_table" @@ -19061,12 +19111,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 720896 + "state": "kv.11" }, { - "state": "kv", - "offset": 721152 + "state": "kv.11", + "offset": 256 }, { "buf": "slot_mapping" @@ -19092,12 +19141,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 720896 + "state": "kv.11" }, { - "state": "kv", - "offset": 721152 + "state": "kv.11", + "offset": 256 }, { "buf": "block_table" @@ -19342,12 +19390,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 786432 + "state": "kv.12" }, { - "state": "kv", - "offset": 786688 + "state": "kv.12", + "offset": 256 }, { "buf": "slot_mapping" @@ -19373,12 +19420,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 786432 + "state": "kv.12" }, { - "state": "kv", - "offset": 786688 + "state": "kv.12", + "offset": 256 }, { "buf": "block_table" @@ -19623,12 +19669,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 851968 + "state": "kv.13" }, { - "state": "kv", - "offset": 852224 + "state": "kv.13", + "offset": 256 }, { "buf": "slot_mapping" @@ -19654,12 +19699,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 851968 + "state": "kv.13" }, { - "state": "kv", - "offset": 852224 + "state": "kv.13", + "offset": 256 }, { "buf": "block_table" @@ -19904,12 +19948,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 917504 + "state": "kv.14" }, { - "state": "kv", - "offset": 917760 + "state": "kv.14", + "offset": 256 }, { "buf": "slot_mapping" @@ -19935,12 +19978,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 917504 + "state": "kv.14" }, { - "state": "kv", - "offset": 917760 + "state": "kv.14", + "offset": 256 }, { "buf": "block_table" @@ -20185,12 +20227,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 983040 + "state": "kv.15" }, { - "state": "kv", - "offset": 983296 + "state": "kv.15", + "offset": 256 }, { "buf": "slot_mapping" @@ -20216,12 +20257,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 983040 + "state": "kv.15" }, { - "state": "kv", - "offset": 983296 + "state": "kv.15", + "offset": 256 }, { "buf": "block_table" @@ -20466,12 +20506,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1048576 + "state": "kv.16" }, { - "state": "kv", - "offset": 1048832 + "state": "kv.16", + "offset": 256 }, { "buf": "slot_mapping" @@ -20497,12 +20536,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1048576 + "state": "kv.16" }, { - "state": "kv", - "offset": 1048832 + "state": "kv.16", + "offset": 256 }, { "buf": "block_table" @@ -20747,12 +20785,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1114112 + "state": "kv.17" }, { - "state": "kv", - "offset": 1114368 + "state": "kv.17", + "offset": 256 }, { "buf": "slot_mapping" @@ -20778,12 +20815,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1114112 + "state": "kv.17" }, { - "state": "kv", - "offset": 1114368 + "state": "kv.17", + "offset": 256 }, { "buf": "block_table" @@ -21028,12 +21064,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1179648 + "state": "kv.18" }, { - "state": "kv", - "offset": 1179904 + "state": "kv.18", + "offset": 256 }, { "buf": "slot_mapping" @@ -21059,12 +21094,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1179648 + "state": "kv.18" }, { - "state": "kv", - "offset": 1179904 + "state": "kv.18", + "offset": 256 }, { "buf": "block_table" @@ -21309,12 +21343,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1245184 + "state": "kv.19" }, { - "state": "kv", - "offset": 1245440 + "state": "kv.19", + "offset": 256 }, { "buf": "slot_mapping" @@ -21340,12 +21373,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1245184 + "state": "kv.19" }, { - "state": "kv", - "offset": 1245440 + "state": "kv.19", + "offset": 256 }, { "buf": "block_table" @@ -21590,12 +21622,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1310720 + "state": "kv.20" }, { - "state": "kv", - "offset": 1310976 + "state": "kv.20", + "offset": 256 }, { "buf": "slot_mapping" @@ -21621,12 +21652,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1310720 + "state": "kv.20" }, { - "state": "kv", - "offset": 1310976 + "state": "kv.20", + "offset": 256 }, { "buf": "block_table" @@ -21871,12 +21901,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1376256 + "state": "kv.21" }, { - "state": "kv", - "offset": 1376512 + "state": "kv.21", + "offset": 256 }, { "buf": "slot_mapping" @@ -21902,12 +21931,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1376256 + "state": "kv.21" }, { - "state": "kv", - "offset": 1376512 + "state": "kv.21", + "offset": 256 }, { "buf": "block_table" @@ -22152,12 +22180,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1441792 + "state": "kv.22" }, { - "state": "kv", - "offset": 1442048 + "state": "kv.22", + "offset": 256 }, { "buf": "slot_mapping" @@ -22183,12 +22210,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1441792 + "state": "kv.22" }, { - "state": "kv", - "offset": 1442048 + "state": "kv.22", + "offset": 256 }, { "buf": "block_table" @@ -22433,12 +22459,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1507328 + "state": "kv.23" }, { - "state": "kv", - "offset": 1507584 + "state": "kv.23", + "offset": 256 }, { "buf": "slot_mapping" @@ -22464,12 +22489,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1507328 + "state": "kv.23" }, { - "state": "kv", - "offset": 1507584 + "state": "kv.23", + "offset": 256 }, { "buf": "block_table" @@ -22714,12 +22738,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1572864 + "state": "kv.24" }, { - "state": "kv", - "offset": 1573120 + "state": "kv.24", + "offset": 256 }, { "buf": "slot_mapping" @@ -22745,12 +22768,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1572864 + "state": "kv.24" }, { - "state": "kv", - "offset": 1573120 + "state": "kv.24", + "offset": 256 }, { "buf": "block_table" @@ -22995,12 +23017,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1638400 + "state": "kv.25" }, { - "state": "kv", - "offset": 1638656 + "state": "kv.25", + "offset": 256 }, { "buf": "slot_mapping" @@ -23026,12 +23047,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1638400 + "state": "kv.25" }, { - "state": "kv", - "offset": 1638656 + "state": "kv.25", + "offset": 256 }, { "buf": "block_table" @@ -23276,12 +23296,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1703936 + "state": "kv.26" }, { - "state": "kv", - "offset": 1704192 + "state": "kv.26", + "offset": 256 }, { "buf": "slot_mapping" @@ -23307,12 +23326,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1703936 + "state": "kv.26" }, { - "state": "kv", - "offset": 1704192 + "state": "kv.26", + "offset": 256 }, { "buf": "block_table" @@ -23557,12 +23575,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1769472 + "state": "kv.27" }, { - "state": "kv", - "offset": 1769728 + "state": "kv.27", + "offset": 256 }, { "buf": "slot_mapping" @@ -23588,12 +23605,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1769472 + "state": "kv.27" }, { - "state": "kv", - "offset": 1769728 + "state": "kv.27", + "offset": 256 }, { "buf": "block_table" @@ -23838,12 +23854,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1835008 + "state": "kv.28" }, { - "state": "kv", - "offset": 1835264 + "state": "kv.28", + "offset": 256 }, { "buf": "slot_mapping" @@ -23869,12 +23884,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1835008 + "state": "kv.28" }, { - "state": "kv", - "offset": 1835264 + "state": "kv.28", + "offset": 256 }, { "buf": "block_table" @@ -24119,12 +24133,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1900544 + "state": "kv.29" }, { - "state": "kv", - "offset": 1900800 + "state": "kv.29", + "offset": 256 }, { "buf": "slot_mapping" @@ -24150,12 +24163,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1900544 + "state": "kv.29" }, { - "state": "kv", - "offset": 1900800 + "state": "kv.29", + "offset": 256 }, { "buf": "block_table" @@ -24400,12 +24412,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1966080 + "state": "kv.30" }, { - "state": "kv", - "offset": 1966336 + "state": "kv.30", + "offset": 256 }, { "buf": "slot_mapping" @@ -24431,12 +24442,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1966080 + "state": "kv.30" }, { - "state": "kv", - "offset": 1966336 + "state": "kv.30", + "offset": 256 }, { "buf": "block_table" @@ -24681,12 +24691,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 2031616 + "state": "kv.31" }, { - "state": "kv", - "offset": 2031872 + "state": "kv.31", + "offset": 256 }, { "buf": "slot_mapping" @@ -24712,12 +24721,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 2031616 + "state": "kv.31" }, { - "state": "kv", - "offset": 2031872 + "state": "kv.31", + "offset": 256 }, { "buf": "block_table" @@ -24962,12 +24970,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 2097152 + "state": "kv.32" }, { - "state": "kv", - "offset": 2097408 + "state": "kv.32", + "offset": 256 }, { "buf": "slot_mapping" @@ -24993,12 +25000,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 2097152 + "state": "kv.32" }, { - "state": "kv", - "offset": 2097408 + "state": "kv.32", + "offset": 256 }, { "buf": "block_table" @@ -25243,12 +25249,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 2162688 + "state": "kv.33" }, { - "state": "kv", - "offset": 2162944 + "state": "kv.33", + "offset": 256 }, { "buf": "slot_mapping" @@ -25274,12 +25279,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 2162688 + "state": "kv.33" }, { - "state": "kv", - "offset": 2162944 + "state": "kv.33", + "offset": 256 }, { "buf": "block_table" @@ -25524,12 +25528,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 2228224 + "state": "kv.34" }, { - "state": "kv", - "offset": 2228480 + "state": "kv.34", + "offset": 256 }, { "buf": "slot_mapping" @@ -25555,12 +25558,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 2228224 + "state": "kv.34" }, { - "state": "kv", - "offset": 2228480 + "state": "kv.34", + "offset": 256 }, { "buf": "block_table" @@ -25805,12 +25807,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 2293760 + "state": "kv.35" }, { - "state": "kv", - "offset": 2294016 + "state": "kv.35", + "offset": 256 }, { "buf": "slot_mapping" @@ -25836,12 +25837,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 2293760 + "state": "kv.35" }, { - "state": "kv", - "offset": 2294016 + "state": "kv.35", + "offset": 256 }, { "buf": "block_table" @@ -26167,10 +26167,10 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv" + "state": "kv.0" }, { - "state": "kv", + "state": "kv.0", "offset": 256 }, { @@ -26196,10 +26196,10 @@ "buf": "q_n" }, { - "state": "kv" + "state": "kv.0" }, { - "state": "kv", + "state": "kv.0", "offset": 256 }, { @@ -26447,12 +26447,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 65536 + "state": "kv.1" }, { - "state": "kv", - "offset": 65792 + "state": "kv.1", + "offset": 256 }, { "buf": "slot_mapping" @@ -26478,12 +26477,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 65536 + "state": "kv.1" }, { - "state": "kv", - "offset": 65792 + "state": "kv.1", + "offset": 256 }, { "buf": "block_table" @@ -26731,12 +26729,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 131072 + "state": "kv.2" }, { - "state": "kv", - "offset": 131328 + "state": "kv.2", + "offset": 256 }, { "buf": "slot_mapping" @@ -26762,12 +26759,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 131072 + "state": "kv.2" }, { - "state": "kv", - "offset": 131328 + "state": "kv.2", + "offset": 256 }, { "buf": "block_table" @@ -27015,12 +27011,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 196608 + "state": "kv.3" }, { - "state": "kv", - "offset": 196864 + "state": "kv.3", + "offset": 256 }, { "buf": "slot_mapping" @@ -27046,12 +27041,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 196608 + "state": "kv.3" }, { - "state": "kv", - "offset": 196864 + "state": "kv.3", + "offset": 256 }, { "buf": "block_table" @@ -27299,12 +27293,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 262144 + "state": "kv.4" }, { - "state": "kv", - "offset": 262400 + "state": "kv.4", + "offset": 256 }, { "buf": "slot_mapping" @@ -27330,12 +27323,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 262144 + "state": "kv.4" }, { - "state": "kv", - "offset": 262400 + "state": "kv.4", + "offset": 256 }, { "buf": "block_table" @@ -27583,12 +27575,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 327680 + "state": "kv.5" }, { - "state": "kv", - "offset": 327936 + "state": "kv.5", + "offset": 256 }, { "buf": "slot_mapping" @@ -27614,12 +27605,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 327680 + "state": "kv.5" }, { - "state": "kv", - "offset": 327936 + "state": "kv.5", + "offset": 256 }, { "buf": "block_table" @@ -27867,12 +27857,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 393216 + "state": "kv.6" }, { - "state": "kv", - "offset": 393472 + "state": "kv.6", + "offset": 256 }, { "buf": "slot_mapping" @@ -27898,12 +27887,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 393216 + "state": "kv.6" }, { - "state": "kv", - "offset": 393472 + "state": "kv.6", + "offset": 256 }, { "buf": "block_table" @@ -28151,12 +28139,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 458752 + "state": "kv.7" }, { - "state": "kv", - "offset": 459008 + "state": "kv.7", + "offset": 256 }, { "buf": "slot_mapping" @@ -28182,12 +28169,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 458752 + "state": "kv.7" }, { - "state": "kv", - "offset": 459008 + "state": "kv.7", + "offset": 256 }, { "buf": "block_table" @@ -28435,12 +28421,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 524288 + "state": "kv.8" }, { - "state": "kv", - "offset": 524544 + "state": "kv.8", + "offset": 256 }, { "buf": "slot_mapping" @@ -28466,12 +28451,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 524288 + "state": "kv.8" }, { - "state": "kv", - "offset": 524544 + "state": "kv.8", + "offset": 256 }, { "buf": "block_table" @@ -28719,12 +28703,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 589824 + "state": "kv.9" }, { - "state": "kv", - "offset": 590080 + "state": "kv.9", + "offset": 256 }, { "buf": "slot_mapping" @@ -28750,12 +28733,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 589824 + "state": "kv.9" }, { - "state": "kv", - "offset": 590080 + "state": "kv.9", + "offset": 256 }, { "buf": "block_table" @@ -29003,12 +28985,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 655360 + "state": "kv.10" }, { - "state": "kv", - "offset": 655616 + "state": "kv.10", + "offset": 256 }, { "buf": "slot_mapping" @@ -29034,12 +29015,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 655360 + "state": "kv.10" }, { - "state": "kv", - "offset": 655616 + "state": "kv.10", + "offset": 256 }, { "buf": "block_table" @@ -29287,12 +29267,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 720896 + "state": "kv.11" }, { - "state": "kv", - "offset": 721152 + "state": "kv.11", + "offset": 256 }, { "buf": "slot_mapping" @@ -29318,12 +29297,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 720896 + "state": "kv.11" }, { - "state": "kv", - "offset": 721152 + "state": "kv.11", + "offset": 256 }, { "buf": "block_table" @@ -29571,12 +29549,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 786432 + "state": "kv.12" }, { - "state": "kv", - "offset": 786688 + "state": "kv.12", + "offset": 256 }, { "buf": "slot_mapping" @@ -29602,12 +29579,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 786432 + "state": "kv.12" }, { - "state": "kv", - "offset": 786688 + "state": "kv.12", + "offset": 256 }, { "buf": "block_table" @@ -29855,12 +29831,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 851968 + "state": "kv.13" }, { - "state": "kv", - "offset": 852224 + "state": "kv.13", + "offset": 256 }, { "buf": "slot_mapping" @@ -29886,12 +29861,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 851968 + "state": "kv.13" }, { - "state": "kv", - "offset": 852224 + "state": "kv.13", + "offset": 256 }, { "buf": "block_table" @@ -30139,12 +30113,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 917504 + "state": "kv.14" }, { - "state": "kv", - "offset": 917760 + "state": "kv.14", + "offset": 256 }, { "buf": "slot_mapping" @@ -30170,12 +30143,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 917504 + "state": "kv.14" }, { - "state": "kv", - "offset": 917760 + "state": "kv.14", + "offset": 256 }, { "buf": "block_table" @@ -30423,12 +30395,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 983040 + "state": "kv.15" }, { - "state": "kv", - "offset": 983296 + "state": "kv.15", + "offset": 256 }, { "buf": "slot_mapping" @@ -30454,12 +30425,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 983040 + "state": "kv.15" }, { - "state": "kv", - "offset": 983296 + "state": "kv.15", + "offset": 256 }, { "buf": "block_table" @@ -30707,12 +30677,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1048576 + "state": "kv.16" }, { - "state": "kv", - "offset": 1048832 + "state": "kv.16", + "offset": 256 }, { "buf": "slot_mapping" @@ -30738,12 +30707,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1048576 + "state": "kv.16" }, { - "state": "kv", - "offset": 1048832 + "state": "kv.16", + "offset": 256 }, { "buf": "block_table" @@ -30991,12 +30959,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1114112 + "state": "kv.17" }, { - "state": "kv", - "offset": 1114368 + "state": "kv.17", + "offset": 256 }, { "buf": "slot_mapping" @@ -31022,12 +30989,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1114112 + "state": "kv.17" }, { - "state": "kv", - "offset": 1114368 + "state": "kv.17", + "offset": 256 }, { "buf": "block_table" @@ -31275,12 +31241,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1179648 + "state": "kv.18" }, { - "state": "kv", - "offset": 1179904 + "state": "kv.18", + "offset": 256 }, { "buf": "slot_mapping" @@ -31306,12 +31271,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1179648 + "state": "kv.18" }, { - "state": "kv", - "offset": 1179904 + "state": "kv.18", + "offset": 256 }, { "buf": "block_table" @@ -31559,12 +31523,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1245184 + "state": "kv.19" }, { - "state": "kv", - "offset": 1245440 + "state": "kv.19", + "offset": 256 }, { "buf": "slot_mapping" @@ -31590,12 +31553,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1245184 + "state": "kv.19" }, { - "state": "kv", - "offset": 1245440 + "state": "kv.19", + "offset": 256 }, { "buf": "block_table" @@ -31843,12 +31805,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1310720 + "state": "kv.20" }, { - "state": "kv", - "offset": 1310976 + "state": "kv.20", + "offset": 256 }, { "buf": "slot_mapping" @@ -31874,12 +31835,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1310720 + "state": "kv.20" }, { - "state": "kv", - "offset": 1310976 + "state": "kv.20", + "offset": 256 }, { "buf": "block_table" @@ -32127,12 +32087,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1376256 + "state": "kv.21" }, { - "state": "kv", - "offset": 1376512 + "state": "kv.21", + "offset": 256 }, { "buf": "slot_mapping" @@ -32158,12 +32117,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1376256 + "state": "kv.21" }, { - "state": "kv", - "offset": 1376512 + "state": "kv.21", + "offset": 256 }, { "buf": "block_table" @@ -32411,12 +32369,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1441792 + "state": "kv.22" }, { - "state": "kv", - "offset": 1442048 + "state": "kv.22", + "offset": 256 }, { "buf": "slot_mapping" @@ -32442,12 +32399,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1441792 + "state": "kv.22" }, { - "state": "kv", - "offset": 1442048 + "state": "kv.22", + "offset": 256 }, { "buf": "block_table" @@ -32695,12 +32651,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1507328 + "state": "kv.23" }, { - "state": "kv", - "offset": 1507584 + "state": "kv.23", + "offset": 256 }, { "buf": "slot_mapping" @@ -32726,12 +32681,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1507328 + "state": "kv.23" }, { - "state": "kv", - "offset": 1507584 + "state": "kv.23", + "offset": 256 }, { "buf": "block_table" @@ -32979,12 +32933,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1572864 + "state": "kv.24" }, { - "state": "kv", - "offset": 1573120 + "state": "kv.24", + "offset": 256 }, { "buf": "slot_mapping" @@ -33010,12 +32963,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1572864 + "state": "kv.24" }, { - "state": "kv", - "offset": 1573120 + "state": "kv.24", + "offset": 256 }, { "buf": "block_table" @@ -33263,12 +33215,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1638400 + "state": "kv.25" }, { - "state": "kv", - "offset": 1638656 + "state": "kv.25", + "offset": 256 }, { "buf": "slot_mapping" @@ -33294,12 +33245,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1638400 + "state": "kv.25" }, { - "state": "kv", - "offset": 1638656 + "state": "kv.25", + "offset": 256 }, { "buf": "block_table" @@ -33547,12 +33497,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1703936 + "state": "kv.26" }, { - "state": "kv", - "offset": 1704192 + "state": "kv.26", + "offset": 256 }, { "buf": "slot_mapping" @@ -33578,12 +33527,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1703936 + "state": "kv.26" }, { - "state": "kv", - "offset": 1704192 + "state": "kv.26", + "offset": 256 }, { "buf": "block_table" @@ -33831,12 +33779,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1769472 + "state": "kv.27" }, { - "state": "kv", - "offset": 1769728 + "state": "kv.27", + "offset": 256 }, { "buf": "slot_mapping" @@ -33862,12 +33809,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1769472 + "state": "kv.27" }, { - "state": "kv", - "offset": 1769728 + "state": "kv.27", + "offset": 256 }, { "buf": "block_table" @@ -34115,12 +34061,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1835008 + "state": "kv.28" }, { - "state": "kv", - "offset": 1835264 + "state": "kv.28", + "offset": 256 }, { "buf": "slot_mapping" @@ -34146,12 +34091,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1835008 + "state": "kv.28" }, { - "state": "kv", - "offset": 1835264 + "state": "kv.28", + "offset": 256 }, { "buf": "block_table" @@ -34399,12 +34343,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1900544 + "state": "kv.29" }, { - "state": "kv", - "offset": 1900800 + "state": "kv.29", + "offset": 256 }, { "buf": "slot_mapping" @@ -34430,12 +34373,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1900544 + "state": "kv.29" }, { - "state": "kv", - "offset": 1900800 + "state": "kv.29", + "offset": 256 }, { "buf": "block_table" @@ -34683,12 +34625,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1966080 + "state": "kv.30" }, { - "state": "kv", - "offset": 1966336 + "state": "kv.30", + "offset": 256 }, { "buf": "slot_mapping" @@ -34714,12 +34655,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1966080 + "state": "kv.30" }, { - "state": "kv", - "offset": 1966336 + "state": "kv.30", + "offset": 256 }, { "buf": "block_table" @@ -34967,12 +34907,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 2031616 + "state": "kv.31" }, { - "state": "kv", - "offset": 2031872 + "state": "kv.31", + "offset": 256 }, { "buf": "slot_mapping" @@ -34998,12 +34937,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 2031616 + "state": "kv.31" }, { - "state": "kv", - "offset": 2031872 + "state": "kv.31", + "offset": 256 }, { "buf": "block_table" @@ -35251,12 +35189,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 2097152 + "state": "kv.32" }, { - "state": "kv", - "offset": 2097408 + "state": "kv.32", + "offset": 256 }, { "buf": "slot_mapping" @@ -35282,12 +35219,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 2097152 + "state": "kv.32" }, { - "state": "kv", - "offset": 2097408 + "state": "kv.32", + "offset": 256 }, { "buf": "block_table" @@ -35535,12 +35471,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 2162688 + "state": "kv.33" }, { - "state": "kv", - "offset": 2162944 + "state": "kv.33", + "offset": 256 }, { "buf": "slot_mapping" @@ -35566,12 +35501,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 2162688 + "state": "kv.33" }, { - "state": "kv", - "offset": 2162944 + "state": "kv.33", + "offset": 256 }, { "buf": "block_table" @@ -35819,12 +35753,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 2228224 + "state": "kv.34" }, { - "state": "kv", - "offset": 2228480 + "state": "kv.34", + "offset": 256 }, { "buf": "slot_mapping" @@ -35850,12 +35783,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 2228224 + "state": "kv.34" }, { - "state": "kv", - "offset": 2228480 + "state": "kv.34", + "offset": 256 }, { "buf": "block_table" @@ -36103,12 +36035,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 2293760 + "state": "kv.35" }, { - "state": "kv", - "offset": 2294016 + "state": "kv.35", + "offset": 256 }, { "buf": "slot_mapping" @@ -36134,12 +36065,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 2293760 + "state": "kv.35" }, { - "state": "kv", - "offset": 2294016 + "state": "kv.35", + "offset": 256 }, { "buf": "block_table" diff --git a/examples/qwen3-4b.json b/examples/qwen3-4b.json index 21931136..a58ab33c 100644 --- a/examples/qwen3-4b.json +++ b/examples/qwen3-4b.json @@ -17,7 +17,42 @@ "KV_SCALE_COUNT": 72, "MAX_TOKENS": 2048, "MAX_SEQS": 256, - "KV_BYTES_PER_TOKEN": 147456, + "KV.0_BYTES_PER_TOKEN": 4096, + "KV.1_BYTES_PER_TOKEN": 4096, + "KV.2_BYTES_PER_TOKEN": 4096, + "KV.3_BYTES_PER_TOKEN": 4096, + "KV.4_BYTES_PER_TOKEN": 4096, + "KV.5_BYTES_PER_TOKEN": 4096, + "KV.6_BYTES_PER_TOKEN": 4096, + "KV.7_BYTES_PER_TOKEN": 4096, + "KV.8_BYTES_PER_TOKEN": 4096, + "KV.9_BYTES_PER_TOKEN": 4096, + "KV.10_BYTES_PER_TOKEN": 4096, + "KV.11_BYTES_PER_TOKEN": 4096, + "KV.12_BYTES_PER_TOKEN": 4096, + "KV.13_BYTES_PER_TOKEN": 4096, + "KV.14_BYTES_PER_TOKEN": 4096, + "KV.15_BYTES_PER_TOKEN": 4096, + "KV.16_BYTES_PER_TOKEN": 4096, + "KV.17_BYTES_PER_TOKEN": 4096, + "KV.18_BYTES_PER_TOKEN": 4096, + "KV.19_BYTES_PER_TOKEN": 4096, + "KV.20_BYTES_PER_TOKEN": 4096, + "KV.21_BYTES_PER_TOKEN": 4096, + "KV.22_BYTES_PER_TOKEN": 4096, + "KV.23_BYTES_PER_TOKEN": 4096, + "KV.24_BYTES_PER_TOKEN": 4096, + "KV.25_BYTES_PER_TOKEN": 4096, + "KV.26_BYTES_PER_TOKEN": 4096, + "KV.27_BYTES_PER_TOKEN": 4096, + "KV.28_BYTES_PER_TOKEN": 4096, + "KV.29_BYTES_PER_TOKEN": 4096, + "KV.30_BYTES_PER_TOKEN": 4096, + "KV.31_BYTES_PER_TOKEN": 4096, + "KV.32_BYTES_PER_TOKEN": 4096, + "KV.33_BYTES_PER_TOKEN": 4096, + "KV.34_BYTES_PER_TOKEN": 4096, + "KV.35_BYTES_PER_TOKEN": 4096, "RMS_NORM_EPS": 9.999999974752427e-07, "NUM_HEADS": 32, "NUM_KV_HEADS": 8, @@ -34,8 +69,113 @@ } }, "states": { - "kv": { - "bytes_per_token": "KV_BYTES_PER_TOKEN" + "kv.0": { + "bytes_per_token": "KV.0_BYTES_PER_TOKEN" + }, + "kv.1": { + "bytes_per_token": "KV.1_BYTES_PER_TOKEN" + }, + "kv.2": { + "bytes_per_token": "KV.2_BYTES_PER_TOKEN" + }, + "kv.3": { + "bytes_per_token": "KV.3_BYTES_PER_TOKEN" + }, + "kv.4": { + "bytes_per_token": "KV.4_BYTES_PER_TOKEN" + }, + "kv.5": { + "bytes_per_token": "KV.5_BYTES_PER_TOKEN" + }, + "kv.6": { + "bytes_per_token": "KV.6_BYTES_PER_TOKEN" + }, + "kv.7": { + "bytes_per_token": "KV.7_BYTES_PER_TOKEN" + }, + "kv.8": { + "bytes_per_token": "KV.8_BYTES_PER_TOKEN" + }, + "kv.9": { + "bytes_per_token": "KV.9_BYTES_PER_TOKEN" + }, + "kv.10": { + "bytes_per_token": "KV.10_BYTES_PER_TOKEN" + }, + "kv.11": { + "bytes_per_token": "KV.11_BYTES_PER_TOKEN" + }, + "kv.12": { + "bytes_per_token": "KV.12_BYTES_PER_TOKEN" + }, + "kv.13": { + "bytes_per_token": "KV.13_BYTES_PER_TOKEN" + }, + "kv.14": { + "bytes_per_token": "KV.14_BYTES_PER_TOKEN" + }, + "kv.15": { + "bytes_per_token": "KV.15_BYTES_PER_TOKEN" + }, + "kv.16": { + "bytes_per_token": "KV.16_BYTES_PER_TOKEN" + }, + "kv.17": { + "bytes_per_token": "KV.17_BYTES_PER_TOKEN" + }, + "kv.18": { + "bytes_per_token": "KV.18_BYTES_PER_TOKEN" + }, + "kv.19": { + "bytes_per_token": "KV.19_BYTES_PER_TOKEN" + }, + "kv.20": { + "bytes_per_token": "KV.20_BYTES_PER_TOKEN" + }, + "kv.21": { + "bytes_per_token": "KV.21_BYTES_PER_TOKEN" + }, + "kv.22": { + "bytes_per_token": "KV.22_BYTES_PER_TOKEN" + }, + "kv.23": { + "bytes_per_token": "KV.23_BYTES_PER_TOKEN" + }, + "kv.24": { + "bytes_per_token": "KV.24_BYTES_PER_TOKEN" + }, + "kv.25": { + "bytes_per_token": "KV.25_BYTES_PER_TOKEN" + }, + "kv.26": { + "bytes_per_token": "KV.26_BYTES_PER_TOKEN" + }, + "kv.27": { + "bytes_per_token": "KV.27_BYTES_PER_TOKEN" + }, + "kv.28": { + "bytes_per_token": "KV.28_BYTES_PER_TOKEN" + }, + "kv.29": { + "bytes_per_token": "KV.29_BYTES_PER_TOKEN" + }, + "kv.30": { + "bytes_per_token": "KV.30_BYTES_PER_TOKEN" + }, + "kv.31": { + "bytes_per_token": "KV.31_BYTES_PER_TOKEN" + }, + "kv.32": { + "bytes_per_token": "KV.32_BYTES_PER_TOKEN" + }, + "kv.33": { + "bytes_per_token": "KV.33_BYTES_PER_TOKEN" + }, + "kv.34": { + "bytes_per_token": "KV.34_BYTES_PER_TOKEN" + }, + "kv.35": { + "bytes_per_token": "KV.35_BYTES_PER_TOKEN" } }, "buffers": { @@ -69,7 +209,7 @@ "kind": "input", "fill": "slot", "domain": { - "index_into": "kv" + "index_into": "kv.0" } }, "block_table": { @@ -80,7 +220,7 @@ ], "kind": "input", "domain": { - "index_into": "kv", + "index_into": "kv.0", "stride": "KV_BLOCK_SIZE" } }, @@ -4186,12 +4326,12 @@ "sha256": "73748b54059552f5983322f7dedc36ed349b38ad6fb9318301bb4965b1fe49aa" }, "argmax": { - "source": "argmax.cubin", - "sha256": "25374575bd35dbd9fe07f90982d0da0c2d4a72cc6ac88df5ec92cdb62a3beb16" + "source": "hf:Pegainfer/kern-kernels/blobs/ebe4ba79544ddfac43ee2dcf6a7f5299aba1124e176c201133d4ef9400e3dc25", + "sha256": "ebe4ba79544ddfac43ee2dcf6a7f5299aba1124e176c201133d4ef9400e3dc25" }, "embedding": { - "source": "embedding.cubin", - "sha256": "1c3dd3e591d872aed259d7ef59d44649061f77a4247a5df864a7fbf0ccbde92f" + "source": "hf:Pegainfer/kern-kernels/blobs/868e93ae6e46e36a11c56af1fae49aeff9fcdfdc79f64ca99475a309e721e9e7", + "sha256": "868e93ae6e46e36a11c56af1fae49aeff9fcdfdc79f64ca99475a309e721e9e7" }, "reduce_segments": { "source": "reduce_segments.cubin", @@ -4218,7 +4358,7 @@ "sha256": "116de812ac5a3a186f76fa239d4c16abbed6b8b0af06d16e1baec0aefe8fe8fa" }, "weight_prep": { - "source": "weight_prep.cubin", + "source": "hf:Pegainfer/kern-kernels/blobs/040f51b01c421b9c69a2090034ff1bc4741e40c02f5ea52bd04a3e62e5c1af4d", "sha256": "040f51b01c421b9c69a2090034ff1bc4741e40c02f5ea52bd04a3e62e5c1af4d" } }, @@ -4621,7 +4761,7 @@ "i64": "QKV_DIM" }, { - "i64": 1179648 + "i64": 32768 }, { "i64": 256 @@ -4791,7 +4931,7 @@ "param": 8 }, { - "i64": 1179648 + "i64": 32768 }, { "i64": 2048 @@ -4800,7 +4940,7 @@ "i64": 256 }, { - "i64": 1179648 + "i64": 32768 }, { "i64": 2048 @@ -5022,7 +5162,7 @@ "param": 8 }, { - "i64": 1179648 + "i64": 32768 }, { "i64": 2048 @@ -5031,7 +5171,7 @@ "i64": 256 }, { - "i64": 1179648 + "i64": 32768 }, { "i64": 2048 @@ -5182,7 +5322,7 @@ "param": 8 }, { - "i64": 1179648 + "i64": 32768 }, { "i64": 2048 @@ -5191,7 +5331,7 @@ "i64": 256 }, { - "i64": 1179648 + "i64": 32768 }, { "i64": 2048 @@ -5715,10 +5855,10 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv" + "state": "kv.0" }, { - "state": "kv", + "state": "kv.0", "offset": 256 }, { @@ -5744,10 +5884,10 @@ "buf": "q_n" }, { - "state": "kv" + "state": "kv.0" }, { - "state": "kv", + "state": "kv.0", "offset": 256 }, { @@ -5995,12 +6135,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 65536 + "state": "kv.1" }, { - "state": "kv", - "offset": 65792 + "state": "kv.1", + "offset": 256 }, { "buf": "slot_mapping" @@ -6026,12 +6165,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 65536 + "state": "kv.1" }, { - "state": "kv", - "offset": 65792 + "state": "kv.1", + "offset": 256 }, { "buf": "block_table" @@ -6279,12 +6417,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 131072 + "state": "kv.2" }, { - "state": "kv", - "offset": 131328 + "state": "kv.2", + "offset": 256 }, { "buf": "slot_mapping" @@ -6310,12 +6447,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 131072 + "state": "kv.2" }, { - "state": "kv", - "offset": 131328 + "state": "kv.2", + "offset": 256 }, { "buf": "block_table" @@ -6563,12 +6699,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 196608 + "state": "kv.3" }, { - "state": "kv", - "offset": 196864 + "state": "kv.3", + "offset": 256 }, { "buf": "slot_mapping" @@ -6594,12 +6729,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 196608 + "state": "kv.3" }, { - "state": "kv", - "offset": 196864 + "state": "kv.3", + "offset": 256 }, { "buf": "block_table" @@ -6847,12 +6981,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 262144 + "state": "kv.4" }, { - "state": "kv", - "offset": 262400 + "state": "kv.4", + "offset": 256 }, { "buf": "slot_mapping" @@ -6878,12 +7011,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 262144 + "state": "kv.4" }, { - "state": "kv", - "offset": 262400 + "state": "kv.4", + "offset": 256 }, { "buf": "block_table" @@ -7131,12 +7263,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 327680 + "state": "kv.5" }, { - "state": "kv", - "offset": 327936 + "state": "kv.5", + "offset": 256 }, { "buf": "slot_mapping" @@ -7162,12 +7293,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 327680 + "state": "kv.5" }, { - "state": "kv", - "offset": 327936 + "state": "kv.5", + "offset": 256 }, { "buf": "block_table" @@ -7415,12 +7545,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 393216 + "state": "kv.6" }, { - "state": "kv", - "offset": 393472 + "state": "kv.6", + "offset": 256 }, { "buf": "slot_mapping" @@ -7446,12 +7575,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 393216 + "state": "kv.6" }, { - "state": "kv", - "offset": 393472 + "state": "kv.6", + "offset": 256 }, { "buf": "block_table" @@ -7699,12 +7827,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 458752 + "state": "kv.7" }, { - "state": "kv", - "offset": 459008 + "state": "kv.7", + "offset": 256 }, { "buf": "slot_mapping" @@ -7730,12 +7857,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 458752 + "state": "kv.7" }, { - "state": "kv", - "offset": 459008 + "state": "kv.7", + "offset": 256 }, { "buf": "block_table" @@ -7983,12 +8109,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 524288 + "state": "kv.8" }, { - "state": "kv", - "offset": 524544 + "state": "kv.8", + "offset": 256 }, { "buf": "slot_mapping" @@ -8014,12 +8139,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 524288 + "state": "kv.8" }, { - "state": "kv", - "offset": 524544 + "state": "kv.8", + "offset": 256 }, { "buf": "block_table" @@ -8267,12 +8391,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 589824 + "state": "kv.9" }, { - "state": "kv", - "offset": 590080 + "state": "kv.9", + "offset": 256 }, { "buf": "slot_mapping" @@ -8298,12 +8421,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 589824 + "state": "kv.9" }, { - "state": "kv", - "offset": 590080 + "state": "kv.9", + "offset": 256 }, { "buf": "block_table" @@ -8551,12 +8673,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 655360 + "state": "kv.10" }, { - "state": "kv", - "offset": 655616 + "state": "kv.10", + "offset": 256 }, { "buf": "slot_mapping" @@ -8582,12 +8703,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 655360 + "state": "kv.10" }, { - "state": "kv", - "offset": 655616 + "state": "kv.10", + "offset": 256 }, { "buf": "block_table" @@ -8835,12 +8955,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 720896 + "state": "kv.11" }, { - "state": "kv", - "offset": 721152 + "state": "kv.11", + "offset": 256 }, { "buf": "slot_mapping" @@ -8866,12 +8985,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 720896 + "state": "kv.11" }, { - "state": "kv", - "offset": 721152 + "state": "kv.11", + "offset": 256 }, { "buf": "block_table" @@ -9119,12 +9237,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 786432 + "state": "kv.12" }, { - "state": "kv", - "offset": 786688 + "state": "kv.12", + "offset": 256 }, { "buf": "slot_mapping" @@ -9150,12 +9267,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 786432 + "state": "kv.12" }, { - "state": "kv", - "offset": 786688 + "state": "kv.12", + "offset": 256 }, { "buf": "block_table" @@ -9403,12 +9519,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 851968 + "state": "kv.13" }, { - "state": "kv", - "offset": 852224 + "state": "kv.13", + "offset": 256 }, { "buf": "slot_mapping" @@ -9434,12 +9549,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 851968 + "state": "kv.13" }, { - "state": "kv", - "offset": 852224 + "state": "kv.13", + "offset": 256 }, { "buf": "block_table" @@ -9687,12 +9801,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 917504 + "state": "kv.14" }, { - "state": "kv", - "offset": 917760 + "state": "kv.14", + "offset": 256 }, { "buf": "slot_mapping" @@ -9718,12 +9831,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 917504 + "state": "kv.14" }, { - "state": "kv", - "offset": 917760 + "state": "kv.14", + "offset": 256 }, { "buf": "block_table" @@ -9971,12 +10083,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 983040 + "state": "kv.15" }, { - "state": "kv", - "offset": 983296 + "state": "kv.15", + "offset": 256 }, { "buf": "slot_mapping" @@ -10002,12 +10113,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 983040 + "state": "kv.15" }, { - "state": "kv", - "offset": 983296 + "state": "kv.15", + "offset": 256 }, { "buf": "block_table" @@ -10255,12 +10365,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1048576 + "state": "kv.16" }, { - "state": "kv", - "offset": 1048832 + "state": "kv.16", + "offset": 256 }, { "buf": "slot_mapping" @@ -10286,12 +10395,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1048576 + "state": "kv.16" }, { - "state": "kv", - "offset": 1048832 + "state": "kv.16", + "offset": 256 }, { "buf": "block_table" @@ -10539,12 +10647,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1114112 + "state": "kv.17" }, { - "state": "kv", - "offset": 1114368 + "state": "kv.17", + "offset": 256 }, { "buf": "slot_mapping" @@ -10570,12 +10677,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1114112 + "state": "kv.17" }, { - "state": "kv", - "offset": 1114368 + "state": "kv.17", + "offset": 256 }, { "buf": "block_table" @@ -10823,12 +10929,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1179648 + "state": "kv.18" }, { - "state": "kv", - "offset": 1179904 + "state": "kv.18", + "offset": 256 }, { "buf": "slot_mapping" @@ -10854,12 +10959,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1179648 + "state": "kv.18" }, { - "state": "kv", - "offset": 1179904 + "state": "kv.18", + "offset": 256 }, { "buf": "block_table" @@ -11107,12 +11211,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1245184 + "state": "kv.19" }, { - "state": "kv", - "offset": 1245440 + "state": "kv.19", + "offset": 256 }, { "buf": "slot_mapping" @@ -11138,12 +11241,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1245184 + "state": "kv.19" }, { - "state": "kv", - "offset": 1245440 + "state": "kv.19", + "offset": 256 }, { "buf": "block_table" @@ -11391,12 +11493,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1310720 + "state": "kv.20" }, { - "state": "kv", - "offset": 1310976 + "state": "kv.20", + "offset": 256 }, { "buf": "slot_mapping" @@ -11422,12 +11523,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1310720 + "state": "kv.20" }, { - "state": "kv", - "offset": 1310976 + "state": "kv.20", + "offset": 256 }, { "buf": "block_table" @@ -11675,12 +11775,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1376256 + "state": "kv.21" }, { - "state": "kv", - "offset": 1376512 + "state": "kv.21", + "offset": 256 }, { "buf": "slot_mapping" @@ -11706,12 +11805,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1376256 + "state": "kv.21" }, { - "state": "kv", - "offset": 1376512 + "state": "kv.21", + "offset": 256 }, { "buf": "block_table" @@ -11959,12 +12057,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1441792 + "state": "kv.22" }, { - "state": "kv", - "offset": 1442048 + "state": "kv.22", + "offset": 256 }, { "buf": "slot_mapping" @@ -11990,12 +12087,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1441792 + "state": "kv.22" }, { - "state": "kv", - "offset": 1442048 + "state": "kv.22", + "offset": 256 }, { "buf": "block_table" @@ -12243,12 +12339,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1507328 + "state": "kv.23" }, { - "state": "kv", - "offset": 1507584 + "state": "kv.23", + "offset": 256 }, { "buf": "slot_mapping" @@ -12274,12 +12369,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1507328 + "state": "kv.23" }, { - "state": "kv", - "offset": 1507584 + "state": "kv.23", + "offset": 256 }, { "buf": "block_table" @@ -12527,12 +12621,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1572864 + "state": "kv.24" }, { - "state": "kv", - "offset": 1573120 + "state": "kv.24", + "offset": 256 }, { "buf": "slot_mapping" @@ -12558,12 +12651,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1572864 + "state": "kv.24" }, { - "state": "kv", - "offset": 1573120 + "state": "kv.24", + "offset": 256 }, { "buf": "block_table" @@ -12811,12 +12903,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1638400 + "state": "kv.25" }, { - "state": "kv", - "offset": 1638656 + "state": "kv.25", + "offset": 256 }, { "buf": "slot_mapping" @@ -12842,12 +12933,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1638400 + "state": "kv.25" }, { - "state": "kv", - "offset": 1638656 + "state": "kv.25", + "offset": 256 }, { "buf": "block_table" @@ -13095,12 +13185,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1703936 + "state": "kv.26" }, { - "state": "kv", - "offset": 1704192 + "state": "kv.26", + "offset": 256 }, { "buf": "slot_mapping" @@ -13126,12 +13215,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1703936 + "state": "kv.26" }, { - "state": "kv", - "offset": 1704192 + "state": "kv.26", + "offset": 256 }, { "buf": "block_table" @@ -13379,12 +13467,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1769472 + "state": "kv.27" }, { - "state": "kv", - "offset": 1769728 + "state": "kv.27", + "offset": 256 }, { "buf": "slot_mapping" @@ -13410,12 +13497,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1769472 + "state": "kv.27" }, { - "state": "kv", - "offset": 1769728 + "state": "kv.27", + "offset": 256 }, { "buf": "block_table" @@ -13663,12 +13749,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1835008 + "state": "kv.28" }, { - "state": "kv", - "offset": 1835264 + "state": "kv.28", + "offset": 256 }, { "buf": "slot_mapping" @@ -13694,12 +13779,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1835008 + "state": "kv.28" }, { - "state": "kv", - "offset": 1835264 + "state": "kv.28", + "offset": 256 }, { "buf": "block_table" @@ -13947,12 +14031,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1900544 + "state": "kv.29" }, { - "state": "kv", - "offset": 1900800 + "state": "kv.29", + "offset": 256 }, { "buf": "slot_mapping" @@ -13978,12 +14061,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1900544 + "state": "kv.29" }, { - "state": "kv", - "offset": 1900800 + "state": "kv.29", + "offset": 256 }, { "buf": "block_table" @@ -14231,12 +14313,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1966080 + "state": "kv.30" }, { - "state": "kv", - "offset": 1966336 + "state": "kv.30", + "offset": 256 }, { "buf": "slot_mapping" @@ -14262,12 +14343,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1966080 + "state": "kv.30" }, { - "state": "kv", - "offset": 1966336 + "state": "kv.30", + "offset": 256 }, { "buf": "block_table" @@ -14515,12 +14595,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 2031616 + "state": "kv.31" }, { - "state": "kv", - "offset": 2031872 + "state": "kv.31", + "offset": 256 }, { "buf": "slot_mapping" @@ -14546,12 +14625,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 2031616 + "state": "kv.31" }, { - "state": "kv", - "offset": 2031872 + "state": "kv.31", + "offset": 256 }, { "buf": "block_table" @@ -14799,12 +14877,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 2097152 + "state": "kv.32" }, { - "state": "kv", - "offset": 2097408 + "state": "kv.32", + "offset": 256 }, { "buf": "slot_mapping" @@ -14830,12 +14907,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 2097152 + "state": "kv.32" }, { - "state": "kv", - "offset": 2097408 + "state": "kv.32", + "offset": 256 }, { "buf": "block_table" @@ -15083,12 +15159,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 2162688 + "state": "kv.33" }, { - "state": "kv", - "offset": 2162944 + "state": "kv.33", + "offset": 256 }, { "buf": "slot_mapping" @@ -15114,12 +15189,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 2162688 + "state": "kv.33" }, { - "state": "kv", - "offset": 2162944 + "state": "kv.33", + "offset": 256 }, { "buf": "block_table" @@ -15367,12 +15441,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 2228224 + "state": "kv.34" }, { - "state": "kv", - "offset": 2228480 + "state": "kv.34", + "offset": 256 }, { "buf": "slot_mapping" @@ -15398,12 +15471,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 2228224 + "state": "kv.34" }, { - "state": "kv", - "offset": 2228480 + "state": "kv.34", + "offset": 256 }, { "buf": "block_table" @@ -15651,12 +15723,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 2293760 + "state": "kv.35" }, { - "state": "kv", - "offset": 2294016 + "state": "kv.35", + "offset": 256 }, { "buf": "slot_mapping" @@ -15682,12 +15753,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 2293760 + "state": "kv.35" }, { - "state": "kv", - "offset": 2294016 + "state": "kv.35", + "offset": 256 }, { "buf": "block_table" @@ -15962,10 +16032,10 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv" + "state": "kv.0" }, { - "state": "kv", + "state": "kv.0", "offset": 256 }, { @@ -15991,10 +16061,10 @@ "buf": "q_n" }, { - "state": "kv" + "state": "kv.0" }, { - "state": "kv", + "state": "kv.0", "offset": 256 }, { @@ -16239,12 +16309,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 65536 + "state": "kv.1" }, { - "state": "kv", - "offset": 65792 + "state": "kv.1", + "offset": 256 }, { "buf": "slot_mapping" @@ -16270,12 +16339,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 65536 + "state": "kv.1" }, { - "state": "kv", - "offset": 65792 + "state": "kv.1", + "offset": 256 }, { "buf": "block_table" @@ -16520,12 +16588,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 131072 + "state": "kv.2" }, { - "state": "kv", - "offset": 131328 + "state": "kv.2", + "offset": 256 }, { "buf": "slot_mapping" @@ -16551,12 +16618,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 131072 + "state": "kv.2" }, { - "state": "kv", - "offset": 131328 + "state": "kv.2", + "offset": 256 }, { "buf": "block_table" @@ -16801,12 +16867,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 196608 + "state": "kv.3" }, { - "state": "kv", - "offset": 196864 + "state": "kv.3", + "offset": 256 }, { "buf": "slot_mapping" @@ -16832,12 +16897,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 196608 + "state": "kv.3" }, { - "state": "kv", - "offset": 196864 + "state": "kv.3", + "offset": 256 }, { "buf": "block_table" @@ -17082,12 +17146,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 262144 + "state": "kv.4" }, { - "state": "kv", - "offset": 262400 + "state": "kv.4", + "offset": 256 }, { "buf": "slot_mapping" @@ -17113,12 +17176,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 262144 + "state": "kv.4" }, { - "state": "kv", - "offset": 262400 + "state": "kv.4", + "offset": 256 }, { "buf": "block_table" @@ -17363,12 +17425,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 327680 + "state": "kv.5" }, { - "state": "kv", - "offset": 327936 + "state": "kv.5", + "offset": 256 }, { "buf": "slot_mapping" @@ -17394,12 +17455,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 327680 + "state": "kv.5" }, { - "state": "kv", - "offset": 327936 + "state": "kv.5", + "offset": 256 }, { "buf": "block_table" @@ -17644,12 +17704,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 393216 + "state": "kv.6" }, { - "state": "kv", - "offset": 393472 + "state": "kv.6", + "offset": 256 }, { "buf": "slot_mapping" @@ -17675,12 +17734,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 393216 + "state": "kv.6" }, { - "state": "kv", - "offset": 393472 + "state": "kv.6", + "offset": 256 }, { "buf": "block_table" @@ -17925,12 +17983,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 458752 + "state": "kv.7" }, { - "state": "kv", - "offset": 459008 + "state": "kv.7", + "offset": 256 }, { "buf": "slot_mapping" @@ -17956,12 +18013,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 458752 + "state": "kv.7" }, { - "state": "kv", - "offset": 459008 + "state": "kv.7", + "offset": 256 }, { "buf": "block_table" @@ -18206,12 +18262,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 524288 + "state": "kv.8" }, { - "state": "kv", - "offset": 524544 + "state": "kv.8", + "offset": 256 }, { "buf": "slot_mapping" @@ -18237,12 +18292,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 524288 + "state": "kv.8" }, { - "state": "kv", - "offset": 524544 + "state": "kv.8", + "offset": 256 }, { "buf": "block_table" @@ -18487,12 +18541,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 589824 + "state": "kv.9" }, { - "state": "kv", - "offset": 590080 + "state": "kv.9", + "offset": 256 }, { "buf": "slot_mapping" @@ -18518,12 +18571,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 589824 + "state": "kv.9" }, { - "state": "kv", - "offset": 590080 + "state": "kv.9", + "offset": 256 }, { "buf": "block_table" @@ -18768,12 +18820,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 655360 + "state": "kv.10" }, { - "state": "kv", - "offset": 655616 + "state": "kv.10", + "offset": 256 }, { "buf": "slot_mapping" @@ -18799,12 +18850,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 655360 + "state": "kv.10" }, { - "state": "kv", - "offset": 655616 + "state": "kv.10", + "offset": 256 }, { "buf": "block_table" @@ -19049,12 +19099,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 720896 + "state": "kv.11" }, { - "state": "kv", - "offset": 721152 + "state": "kv.11", + "offset": 256 }, { "buf": "slot_mapping" @@ -19080,12 +19129,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 720896 + "state": "kv.11" }, { - "state": "kv", - "offset": 721152 + "state": "kv.11", + "offset": 256 }, { "buf": "block_table" @@ -19330,12 +19378,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 786432 + "state": "kv.12" }, { - "state": "kv", - "offset": 786688 + "state": "kv.12", + "offset": 256 }, { "buf": "slot_mapping" @@ -19361,12 +19408,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 786432 + "state": "kv.12" }, { - "state": "kv", - "offset": 786688 + "state": "kv.12", + "offset": 256 }, { "buf": "block_table" @@ -19611,12 +19657,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 851968 + "state": "kv.13" }, { - "state": "kv", - "offset": 852224 + "state": "kv.13", + "offset": 256 }, { "buf": "slot_mapping" @@ -19642,12 +19687,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 851968 + "state": "kv.13" }, { - "state": "kv", - "offset": 852224 + "state": "kv.13", + "offset": 256 }, { "buf": "block_table" @@ -19892,12 +19936,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 917504 + "state": "kv.14" }, { - "state": "kv", - "offset": 917760 + "state": "kv.14", + "offset": 256 }, { "buf": "slot_mapping" @@ -19923,12 +19966,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 917504 + "state": "kv.14" }, { - "state": "kv", - "offset": 917760 + "state": "kv.14", + "offset": 256 }, { "buf": "block_table" @@ -20173,12 +20215,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 983040 + "state": "kv.15" }, { - "state": "kv", - "offset": 983296 + "state": "kv.15", + "offset": 256 }, { "buf": "slot_mapping" @@ -20204,12 +20245,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 983040 + "state": "kv.15" }, { - "state": "kv", - "offset": 983296 + "state": "kv.15", + "offset": 256 }, { "buf": "block_table" @@ -20454,12 +20494,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1048576 + "state": "kv.16" }, { - "state": "kv", - "offset": 1048832 + "state": "kv.16", + "offset": 256 }, { "buf": "slot_mapping" @@ -20485,12 +20524,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1048576 + "state": "kv.16" }, { - "state": "kv", - "offset": 1048832 + "state": "kv.16", + "offset": 256 }, { "buf": "block_table" @@ -20735,12 +20773,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1114112 + "state": "kv.17" }, { - "state": "kv", - "offset": 1114368 + "state": "kv.17", + "offset": 256 }, { "buf": "slot_mapping" @@ -20766,12 +20803,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1114112 + "state": "kv.17" }, { - "state": "kv", - "offset": 1114368 + "state": "kv.17", + "offset": 256 }, { "buf": "block_table" @@ -21016,12 +21052,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1179648 + "state": "kv.18" }, { - "state": "kv", - "offset": 1179904 + "state": "kv.18", + "offset": 256 }, { "buf": "slot_mapping" @@ -21047,12 +21082,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1179648 + "state": "kv.18" }, { - "state": "kv", - "offset": 1179904 + "state": "kv.18", + "offset": 256 }, { "buf": "block_table" @@ -21297,12 +21331,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1245184 + "state": "kv.19" }, { - "state": "kv", - "offset": 1245440 + "state": "kv.19", + "offset": 256 }, { "buf": "slot_mapping" @@ -21328,12 +21361,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1245184 + "state": "kv.19" }, { - "state": "kv", - "offset": 1245440 + "state": "kv.19", + "offset": 256 }, { "buf": "block_table" @@ -21578,12 +21610,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1310720 + "state": "kv.20" }, { - "state": "kv", - "offset": 1310976 + "state": "kv.20", + "offset": 256 }, { "buf": "slot_mapping" @@ -21609,12 +21640,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1310720 + "state": "kv.20" }, { - "state": "kv", - "offset": 1310976 + "state": "kv.20", + "offset": 256 }, { "buf": "block_table" @@ -21859,12 +21889,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1376256 + "state": "kv.21" }, { - "state": "kv", - "offset": 1376512 + "state": "kv.21", + "offset": 256 }, { "buf": "slot_mapping" @@ -21890,12 +21919,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1376256 + "state": "kv.21" }, { - "state": "kv", - "offset": 1376512 + "state": "kv.21", + "offset": 256 }, { "buf": "block_table" @@ -22140,12 +22168,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1441792 + "state": "kv.22" }, { - "state": "kv", - "offset": 1442048 + "state": "kv.22", + "offset": 256 }, { "buf": "slot_mapping" @@ -22171,12 +22198,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1441792 + "state": "kv.22" }, { - "state": "kv", - "offset": 1442048 + "state": "kv.22", + "offset": 256 }, { "buf": "block_table" @@ -22421,12 +22447,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1507328 + "state": "kv.23" }, { - "state": "kv", - "offset": 1507584 + "state": "kv.23", + "offset": 256 }, { "buf": "slot_mapping" @@ -22452,12 +22477,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1507328 + "state": "kv.23" }, { - "state": "kv", - "offset": 1507584 + "state": "kv.23", + "offset": 256 }, { "buf": "block_table" @@ -22702,12 +22726,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1572864 + "state": "kv.24" }, { - "state": "kv", - "offset": 1573120 + "state": "kv.24", + "offset": 256 }, { "buf": "slot_mapping" @@ -22733,12 +22756,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1572864 + "state": "kv.24" }, { - "state": "kv", - "offset": 1573120 + "state": "kv.24", + "offset": 256 }, { "buf": "block_table" @@ -22983,12 +23005,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1638400 + "state": "kv.25" }, { - "state": "kv", - "offset": 1638656 + "state": "kv.25", + "offset": 256 }, { "buf": "slot_mapping" @@ -23014,12 +23035,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1638400 + "state": "kv.25" }, { - "state": "kv", - "offset": 1638656 + "state": "kv.25", + "offset": 256 }, { "buf": "block_table" @@ -23264,12 +23284,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1703936 + "state": "kv.26" }, { - "state": "kv", - "offset": 1704192 + "state": "kv.26", + "offset": 256 }, { "buf": "slot_mapping" @@ -23295,12 +23314,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1703936 + "state": "kv.26" }, { - "state": "kv", - "offset": 1704192 + "state": "kv.26", + "offset": 256 }, { "buf": "block_table" @@ -23545,12 +23563,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1769472 + "state": "kv.27" }, { - "state": "kv", - "offset": 1769728 + "state": "kv.27", + "offset": 256 }, { "buf": "slot_mapping" @@ -23576,12 +23593,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1769472 + "state": "kv.27" }, { - "state": "kv", - "offset": 1769728 + "state": "kv.27", + "offset": 256 }, { "buf": "block_table" @@ -23826,12 +23842,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1835008 + "state": "kv.28" }, { - "state": "kv", - "offset": 1835264 + "state": "kv.28", + "offset": 256 }, { "buf": "slot_mapping" @@ -23857,12 +23872,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1835008 + "state": "kv.28" }, { - "state": "kv", - "offset": 1835264 + "state": "kv.28", + "offset": 256 }, { "buf": "block_table" @@ -24107,12 +24121,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1900544 + "state": "kv.29" }, { - "state": "kv", - "offset": 1900800 + "state": "kv.29", + "offset": 256 }, { "buf": "slot_mapping" @@ -24138,12 +24151,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1900544 + "state": "kv.29" }, { - "state": "kv", - "offset": 1900800 + "state": "kv.29", + "offset": 256 }, { "buf": "block_table" @@ -24388,12 +24400,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1966080 + "state": "kv.30" }, { - "state": "kv", - "offset": 1966336 + "state": "kv.30", + "offset": 256 }, { "buf": "slot_mapping" @@ -24419,12 +24430,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1966080 + "state": "kv.30" }, { - "state": "kv", - "offset": 1966336 + "state": "kv.30", + "offset": 256 }, { "buf": "block_table" @@ -24669,12 +24679,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 2031616 + "state": "kv.31" }, { - "state": "kv", - "offset": 2031872 + "state": "kv.31", + "offset": 256 }, { "buf": "slot_mapping" @@ -24700,12 +24709,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 2031616 + "state": "kv.31" }, { - "state": "kv", - "offset": 2031872 + "state": "kv.31", + "offset": 256 }, { "buf": "block_table" @@ -24950,12 +24958,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 2097152 + "state": "kv.32" }, { - "state": "kv", - "offset": 2097408 + "state": "kv.32", + "offset": 256 }, { "buf": "slot_mapping" @@ -24981,12 +24988,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 2097152 + "state": "kv.32" }, { - "state": "kv", - "offset": 2097408 + "state": "kv.32", + "offset": 256 }, { "buf": "block_table" @@ -25231,12 +25237,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 2162688 + "state": "kv.33" }, { - "state": "kv", - "offset": 2162944 + "state": "kv.33", + "offset": 256 }, { "buf": "slot_mapping" @@ -25262,12 +25267,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 2162688 + "state": "kv.33" }, { - "state": "kv", - "offset": 2162944 + "state": "kv.33", + "offset": 256 }, { "buf": "block_table" @@ -25512,12 +25516,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 2228224 + "state": "kv.34" }, { - "state": "kv", - "offset": 2228480 + "state": "kv.34", + "offset": 256 }, { "buf": "slot_mapping" @@ -25543,12 +25546,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 2228224 + "state": "kv.34" }, { - "state": "kv", - "offset": 2228480 + "state": "kv.34", + "offset": 256 }, { "buf": "block_table" @@ -25793,12 +25795,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 2293760 + "state": "kv.35" }, { - "state": "kv", - "offset": 2294016 + "state": "kv.35", + "offset": 256 }, { "buf": "slot_mapping" @@ -25824,12 +25825,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 2293760 + "state": "kv.35" }, { - "state": "kv", - "offset": 2294016 + "state": "kv.35", + "offset": 256 }, { "buf": "block_table" @@ -26155,10 +26155,10 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv" + "state": "kv.0" }, { - "state": "kv", + "state": "kv.0", "offset": 256 }, { @@ -26184,10 +26184,10 @@ "buf": "q_n" }, { - "state": "kv" + "state": "kv.0" }, { - "state": "kv", + "state": "kv.0", "offset": 256 }, { @@ -26435,12 +26435,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 65536 + "state": "kv.1" }, { - "state": "kv", - "offset": 65792 + "state": "kv.1", + "offset": 256 }, { "buf": "slot_mapping" @@ -26466,12 +26465,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 65536 + "state": "kv.1" }, { - "state": "kv", - "offset": 65792 + "state": "kv.1", + "offset": 256 }, { "buf": "block_table" @@ -26719,12 +26717,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 131072 + "state": "kv.2" }, { - "state": "kv", - "offset": 131328 + "state": "kv.2", + "offset": 256 }, { "buf": "slot_mapping" @@ -26750,12 +26747,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 131072 + "state": "kv.2" }, { - "state": "kv", - "offset": 131328 + "state": "kv.2", + "offset": 256 }, { "buf": "block_table" @@ -27003,12 +26999,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 196608 + "state": "kv.3" }, { - "state": "kv", - "offset": 196864 + "state": "kv.3", + "offset": 256 }, { "buf": "slot_mapping" @@ -27034,12 +27029,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 196608 + "state": "kv.3" }, { - "state": "kv", - "offset": 196864 + "state": "kv.3", + "offset": 256 }, { "buf": "block_table" @@ -27287,12 +27281,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 262144 + "state": "kv.4" }, { - "state": "kv", - "offset": 262400 + "state": "kv.4", + "offset": 256 }, { "buf": "slot_mapping" @@ -27318,12 +27311,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 262144 + "state": "kv.4" }, { - "state": "kv", - "offset": 262400 + "state": "kv.4", + "offset": 256 }, { "buf": "block_table" @@ -27571,12 +27563,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 327680 + "state": "kv.5" }, { - "state": "kv", - "offset": 327936 + "state": "kv.5", + "offset": 256 }, { "buf": "slot_mapping" @@ -27602,12 +27593,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 327680 + "state": "kv.5" }, { - "state": "kv", - "offset": 327936 + "state": "kv.5", + "offset": 256 }, { "buf": "block_table" @@ -27855,12 +27845,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 393216 + "state": "kv.6" }, { - "state": "kv", - "offset": 393472 + "state": "kv.6", + "offset": 256 }, { "buf": "slot_mapping" @@ -27886,12 +27875,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 393216 + "state": "kv.6" }, { - "state": "kv", - "offset": 393472 + "state": "kv.6", + "offset": 256 }, { "buf": "block_table" @@ -28139,12 +28127,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 458752 + "state": "kv.7" }, { - "state": "kv", - "offset": 459008 + "state": "kv.7", + "offset": 256 }, { "buf": "slot_mapping" @@ -28170,12 +28157,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 458752 + "state": "kv.7" }, { - "state": "kv", - "offset": 459008 + "state": "kv.7", + "offset": 256 }, { "buf": "block_table" @@ -28423,12 +28409,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 524288 + "state": "kv.8" }, { - "state": "kv", - "offset": 524544 + "state": "kv.8", + "offset": 256 }, { "buf": "slot_mapping" @@ -28454,12 +28439,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 524288 + "state": "kv.8" }, { - "state": "kv", - "offset": 524544 + "state": "kv.8", + "offset": 256 }, { "buf": "block_table" @@ -28707,12 +28691,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 589824 + "state": "kv.9" }, { - "state": "kv", - "offset": 590080 + "state": "kv.9", + "offset": 256 }, { "buf": "slot_mapping" @@ -28738,12 +28721,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 589824 + "state": "kv.9" }, { - "state": "kv", - "offset": 590080 + "state": "kv.9", + "offset": 256 }, { "buf": "block_table" @@ -28991,12 +28973,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 655360 + "state": "kv.10" }, { - "state": "kv", - "offset": 655616 + "state": "kv.10", + "offset": 256 }, { "buf": "slot_mapping" @@ -29022,12 +29003,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 655360 + "state": "kv.10" }, { - "state": "kv", - "offset": 655616 + "state": "kv.10", + "offset": 256 }, { "buf": "block_table" @@ -29275,12 +29255,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 720896 + "state": "kv.11" }, { - "state": "kv", - "offset": 721152 + "state": "kv.11", + "offset": 256 }, { "buf": "slot_mapping" @@ -29306,12 +29285,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 720896 + "state": "kv.11" }, { - "state": "kv", - "offset": 721152 + "state": "kv.11", + "offset": 256 }, { "buf": "block_table" @@ -29559,12 +29537,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 786432 + "state": "kv.12" }, { - "state": "kv", - "offset": 786688 + "state": "kv.12", + "offset": 256 }, { "buf": "slot_mapping" @@ -29590,12 +29567,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 786432 + "state": "kv.12" }, { - "state": "kv", - "offset": 786688 + "state": "kv.12", + "offset": 256 }, { "buf": "block_table" @@ -29843,12 +29819,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 851968 + "state": "kv.13" }, { - "state": "kv", - "offset": 852224 + "state": "kv.13", + "offset": 256 }, { "buf": "slot_mapping" @@ -29874,12 +29849,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 851968 + "state": "kv.13" }, { - "state": "kv", - "offset": 852224 + "state": "kv.13", + "offset": 256 }, { "buf": "block_table" @@ -30127,12 +30101,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 917504 + "state": "kv.14" }, { - "state": "kv", - "offset": 917760 + "state": "kv.14", + "offset": 256 }, { "buf": "slot_mapping" @@ -30158,12 +30131,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 917504 + "state": "kv.14" }, { - "state": "kv", - "offset": 917760 + "state": "kv.14", + "offset": 256 }, { "buf": "block_table" @@ -30411,12 +30383,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 983040 + "state": "kv.15" }, { - "state": "kv", - "offset": 983296 + "state": "kv.15", + "offset": 256 }, { "buf": "slot_mapping" @@ -30442,12 +30413,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 983040 + "state": "kv.15" }, { - "state": "kv", - "offset": 983296 + "state": "kv.15", + "offset": 256 }, { "buf": "block_table" @@ -30695,12 +30665,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1048576 + "state": "kv.16" }, { - "state": "kv", - "offset": 1048832 + "state": "kv.16", + "offset": 256 }, { "buf": "slot_mapping" @@ -30726,12 +30695,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1048576 + "state": "kv.16" }, { - "state": "kv", - "offset": 1048832 + "state": "kv.16", + "offset": 256 }, { "buf": "block_table" @@ -30979,12 +30947,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1114112 + "state": "kv.17" }, { - "state": "kv", - "offset": 1114368 + "state": "kv.17", + "offset": 256 }, { "buf": "slot_mapping" @@ -31010,12 +30977,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1114112 + "state": "kv.17" }, { - "state": "kv", - "offset": 1114368 + "state": "kv.17", + "offset": 256 }, { "buf": "block_table" @@ -31263,12 +31229,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1179648 + "state": "kv.18" }, { - "state": "kv", - "offset": 1179904 + "state": "kv.18", + "offset": 256 }, { "buf": "slot_mapping" @@ -31294,12 +31259,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1179648 + "state": "kv.18" }, { - "state": "kv", - "offset": 1179904 + "state": "kv.18", + "offset": 256 }, { "buf": "block_table" @@ -31547,12 +31511,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1245184 + "state": "kv.19" }, { - "state": "kv", - "offset": 1245440 + "state": "kv.19", + "offset": 256 }, { "buf": "slot_mapping" @@ -31578,12 +31541,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1245184 + "state": "kv.19" }, { - "state": "kv", - "offset": 1245440 + "state": "kv.19", + "offset": 256 }, { "buf": "block_table" @@ -31831,12 +31793,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1310720 + "state": "kv.20" }, { - "state": "kv", - "offset": 1310976 + "state": "kv.20", + "offset": 256 }, { "buf": "slot_mapping" @@ -31862,12 +31823,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1310720 + "state": "kv.20" }, { - "state": "kv", - "offset": 1310976 + "state": "kv.20", + "offset": 256 }, { "buf": "block_table" @@ -32115,12 +32075,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1376256 + "state": "kv.21" }, { - "state": "kv", - "offset": 1376512 + "state": "kv.21", + "offset": 256 }, { "buf": "slot_mapping" @@ -32146,12 +32105,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1376256 + "state": "kv.21" }, { - "state": "kv", - "offset": 1376512 + "state": "kv.21", + "offset": 256 }, { "buf": "block_table" @@ -32399,12 +32357,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1441792 + "state": "kv.22" }, { - "state": "kv", - "offset": 1442048 + "state": "kv.22", + "offset": 256 }, { "buf": "slot_mapping" @@ -32430,12 +32387,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1441792 + "state": "kv.22" }, { - "state": "kv", - "offset": 1442048 + "state": "kv.22", + "offset": 256 }, { "buf": "block_table" @@ -32683,12 +32639,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1507328 + "state": "kv.23" }, { - "state": "kv", - "offset": 1507584 + "state": "kv.23", + "offset": 256 }, { "buf": "slot_mapping" @@ -32714,12 +32669,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1507328 + "state": "kv.23" }, { - "state": "kv", - "offset": 1507584 + "state": "kv.23", + "offset": 256 }, { "buf": "block_table" @@ -32967,12 +32921,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1572864 + "state": "kv.24" }, { - "state": "kv", - "offset": 1573120 + "state": "kv.24", + "offset": 256 }, { "buf": "slot_mapping" @@ -32998,12 +32951,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1572864 + "state": "kv.24" }, { - "state": "kv", - "offset": 1573120 + "state": "kv.24", + "offset": 256 }, { "buf": "block_table" @@ -33251,12 +33203,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1638400 + "state": "kv.25" }, { - "state": "kv", - "offset": 1638656 + "state": "kv.25", + "offset": 256 }, { "buf": "slot_mapping" @@ -33282,12 +33233,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1638400 + "state": "kv.25" }, { - "state": "kv", - "offset": 1638656 + "state": "kv.25", + "offset": 256 }, { "buf": "block_table" @@ -33535,12 +33485,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1703936 + "state": "kv.26" }, { - "state": "kv", - "offset": 1704192 + "state": "kv.26", + "offset": 256 }, { "buf": "slot_mapping" @@ -33566,12 +33515,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1703936 + "state": "kv.26" }, { - "state": "kv", - "offset": 1704192 + "state": "kv.26", + "offset": 256 }, { "buf": "block_table" @@ -33819,12 +33767,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1769472 + "state": "kv.27" }, { - "state": "kv", - "offset": 1769728 + "state": "kv.27", + "offset": 256 }, { "buf": "slot_mapping" @@ -33850,12 +33797,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1769472 + "state": "kv.27" }, { - "state": "kv", - "offset": 1769728 + "state": "kv.27", + "offset": 256 }, { "buf": "block_table" @@ -34103,12 +34049,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1835008 + "state": "kv.28" }, { - "state": "kv", - "offset": 1835264 + "state": "kv.28", + "offset": 256 }, { "buf": "slot_mapping" @@ -34134,12 +34079,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1835008 + "state": "kv.28" }, { - "state": "kv", - "offset": 1835264 + "state": "kv.28", + "offset": 256 }, { "buf": "block_table" @@ -34387,12 +34331,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1900544 + "state": "kv.29" }, { - "state": "kv", - "offset": 1900800 + "state": "kv.29", + "offset": 256 }, { "buf": "slot_mapping" @@ -34418,12 +34361,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1900544 + "state": "kv.29" }, { - "state": "kv", - "offset": 1900800 + "state": "kv.29", + "offset": 256 }, { "buf": "block_table" @@ -34671,12 +34613,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 1966080 + "state": "kv.30" }, { - "state": "kv", - "offset": 1966336 + "state": "kv.30", + "offset": 256 }, { "buf": "slot_mapping" @@ -34702,12 +34643,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 1966080 + "state": "kv.30" }, { - "state": "kv", - "offset": 1966336 + "state": "kv.30", + "offset": 256 }, { "buf": "block_table" @@ -34955,12 +34895,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 2031616 + "state": "kv.31" }, { - "state": "kv", - "offset": 2031872 + "state": "kv.31", + "offset": 256 }, { "buf": "slot_mapping" @@ -34986,12 +34925,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 2031616 + "state": "kv.31" }, { - "state": "kv", - "offset": 2031872 + "state": "kv.31", + "offset": 256 }, { "buf": "block_table" @@ -35239,12 +35177,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 2097152 + "state": "kv.32" }, { - "state": "kv", - "offset": 2097408 + "state": "kv.32", + "offset": 256 }, { "buf": "slot_mapping" @@ -35270,12 +35207,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 2097152 + "state": "kv.32" }, { - "state": "kv", - "offset": 2097408 + "state": "kv.32", + "offset": 256 }, { "buf": "block_table" @@ -35523,12 +35459,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 2162688 + "state": "kv.33" }, { - "state": "kv", - "offset": 2162944 + "state": "kv.33", + "offset": 256 }, { "buf": "slot_mapping" @@ -35554,12 +35489,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 2162688 + "state": "kv.33" }, { - "state": "kv", - "offset": 2162944 + "state": "kv.33", + "offset": 256 }, { "buf": "block_table" @@ -35807,12 +35741,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 2228224 + "state": "kv.34" }, { - "state": "kv", - "offset": 2228480 + "state": "kv.34", + "offset": 256 }, { "buf": "slot_mapping" @@ -35838,12 +35771,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 2228224 + "state": "kv.34" }, { - "state": "kv", - "offset": 2228480 + "state": "kv.34", + "offset": 256 }, { "buf": "block_table" @@ -36091,12 +36023,11 @@ "offset": "QKV_V_OFFSET_BYTES" }, { - "state": "kv", - "offset": 2293760 + "state": "kv.35" }, { - "state": "kv", - "offset": 2294016 + "state": "kv.35", + "offset": 256 }, { "buf": "slot_mapping" @@ -36122,12 +36053,11 @@ "buf": "q_n" }, { - "state": "kv", - "offset": 2293760 + "state": "kv.35" }, { - "state": "kv", - "offset": 2294016 + "state": "kv.35", + "offset": 256 }, { "buf": "block_table" diff --git a/tools/gen_qwen3_decode.py b/tools/gen_qwen3_decode.py index 03d673ab..f9b12f3f 100644 --- a/tools/gen_qwen3_decode.py +++ b/tools/gen_qwen3_decode.py @@ -16,9 +16,9 @@ 被 Triton 特化进 binary 不再传参）——这里 pin 的是 decode 实例。 - GEMM 是 runtime 特判（symbol 前缀 `extern:`，cublasLt）；embedding 是 待写的 Triton 占位。 -- KV state 布局从 vLLM 的逐层池改为层交织 `[page][layer][16][8][2][128]` - （同一批 kernel，靠 stride 参数 ×LAYERS 和 state offset 字面量适配）， - bytes_per_token = 36*2*8*128*2 = 147456。 +- KV 是每层一个 paged state `kv.`（vLLM 的逐层池布局，bytes_per_token + = 2*8*128*2 = 4096），页号由 pool 统一发；一份切开的 manifest 每段只带 + 自己那些层的 state。 - 发射前对挖矿数据做结构断言：q/k/v 在 qkv 中的视图偏移、residual 全程 同址、逐层权重互异、KV 池 k/v 相距 256B、cache/attention 共享同一 KV 池与 scale 指针、unified 与 reduce 共享 segm 缓冲——连线是手写的， @@ -65,7 +65,7 @@ sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent)) import mine_capture as mc -from handwritten import hw # tools/handwritten.py: build + pin handwritten cubins +from kernels.index import variant # noqa: E402 HIDDEN = 2560 LAYERS = 36 @@ -96,12 +96,13 @@ TAPS = {0: 0, 8: 1, 16: 2, 24: 3, 32: 4} # target_layer_ids [1,9,17,25,33]-1 DRAFT_KV_DIM = DRAFT_LAYERS * 2 * KV_DIM # 融合 KV GEMM 输出行宽 10240 -# 层交织 KV 布局（vLLM 逐层池 -> 我们的单 state）： -# 一个 block(16 token) 在一层里占 16*8*2*128 = 32768 elems；层间连续。 +# target KV：每层一个 paged state `kv.`，即 vLLM 的逐层池布局 +# `[page][16][8][2][128]`；pool 给所有 paged state 同一套页号。一个 +# block(16 token) 在一层里占 16*8*2*128 = 32768 elems。draft 仍是一个 +# 层交织的 state（`draft_kv`，stride × 层数）。 BLOCK_ELEMS_PER_LAYER = 16 * KV_HEADS * 2 * HEAD_DIM -BLOCK_STRIDE = LAYERS * BLOCK_ELEMS_PER_LAYER # 传给 kernel 的 elems -LAYER_KV_BYTES = BLOCK_ELEMS_PER_LAYER * BF16 # state offset 步长 -KV_BYTES_PER_TOKEN = LAYERS * 2 * KV_DIM * BF16 # 147456 +LAYER_KV_BYTES = BLOCK_ELEMS_PER_LAYER * BF16 # draft_kv 的层 offset 步长 +LAYER_KV_BYTES_PER_TOKEN = 2 * KV_DIM * BF16 # 4096 DRAFT_BLOCK_STRIDE = DRAFT_LAYERS * BLOCK_ELEMS_PER_LAYER DRAFT_KV_BYTES_PER_TOKEN = DRAFT_LAYERS * 2 * KV_DIM * BF16 # 20480 V_BYTE_OFF = 2 * HEAD_DIM # 挖矿实测 k/v 相距 256B @@ -388,7 +389,7 @@ def scr(name): def step(symbol, params, block, grid, args, shared_mem=None, cubin=None, - sha256=None): + sha256=None, label=None): s = {"entry": symbol, "params": params, "block": block, "grid": [_e(g) for g in grid], "args": args} if shared_mem is not None: @@ -397,16 +398,18 @@ def step(symbol, params, block, grid, args, shared_mem=None, cubin=None, s["cubin"] = cubin if sha256 is not None: s["sha256"] = sha256 + if label is not None: + s["label"] = label return s def single(symbol, params, block, grid, shared_mem=None, cubin=None, - sha256=None): + sha256=None, label=None): """单步实现，恒等布线：接口即该核的 launch ABI。""" return {"params": params, "impl": {"launches": [step(symbol, params, block, grid, [a(i) for i in range(len(params))], - shared_mem, cubin, sha256)]}} + shared_mem, cubin, sha256, label)]}} # 结构输入的先验（domain）：接模型的人才知道 buffer 是页表不是激活。 @@ -415,8 +418,8 @@ def single(symbol, params, block, grid, shared_mem=None, cubin=None, DOMAINS = { "token_ids": TOKEN_DOMAIN, "positions": {"index_into": "rope.cos_sin_cache"}, - "slot_mapping": {"index_into": "kv"}, - "block_table": {"index_into": "kv", "stride": 16}, # vLLM block_size + "slot_mapping": {"index_into": "kv.0"}, + "block_table": {"index_into": "kv.0", "stride": 16}, # vLLM block_size "seq_lens": {"min": 1}, "cu_seqlens_q": {"min": 0, "max": "tokens", "monotone": True}, "next_token": TOKEN_DOMAIN, @@ -822,8 +825,8 @@ def fused(label, x, w): [buf(x), i64(HIDDEN), buf("residual"), buf(w), f32(eps), T, i32(HIDDEN), i64(HIDDEN)]) - def forward(attn_kernel, tail, mp="model.", layers=LAYERS, kv_state="kv", - block_stride=BLOCK_STRIDE, scales="kv_scales", taps=False, + def forward(attn_kernel, tail, mp="model.", layers=LAYERS, kv_state=None, + block_stride=BLOCK_ELEMS_PER_LAYER, scales="kv_scales", taps=False, lm_head_w="lm_head.weight", final_norm_w="model.norm.weight", ids="token_ids"): """embed + 逐层的直线 dispatch 表，target/draft 共用（几何同构， @@ -848,7 +851,7 @@ def forward(attn_kernel, tail, mp="model.", layers=LAYERS, kv_state="kv", for i in range(layers): p = f"{mp}layers.{i}." l = f"l{i}." - koff = i * LAYER_KV_BYTES + st, koff = (f"kv.{i}", 0) if kv_state is None else (kv_state, i * LAYER_KV_BYTES) ks, vs = buf(scales, i * 8), buf(scales, i * 8 + 4) last = i + 1 == layers ds += [ @@ -867,14 +870,14 @@ def forward(attn_kernel, tail, mp="model.", layers=LAYERS, kv_state="kv", # [8]=value 行距=QKV_DIM（v 视图在融合 qkv 里；decode # capture 的 1024 同样是 tokens=1 下的假常量） [buf("k_n"), buf("qkv", (Q_DIM + KV_DIM) * BF16), - state(kv_state, koff), state(kv_state, koff + V_BYTE_OFF), + state(st, koff), state(st, koff + V_BYTE_OFF), buf("slot_mapping"), ks, vs, i64(KV_DIM), i64(QKV_DIM), i64(block_stride), i64(2 * HEAD_DIM), i64(0), i64(0), i64(KV_HEADS * 2 * HEAD_DIM), i64(0), i64(0)]), d(l + "attn", attn_kernel, [buf("attn_out"), buf("q_n"), - state(kv_state, koff), state(kv_state, koff + V_BYTE_OFF), + state(st, koff), state(st, koff + V_BYTE_OFF), buf("block_table"), buf("seq_lens"), f32(scale), ks, vs, f32(1.0), f32(0.0), i64(MAX_BLOCKS), i64(Q_DIM), i64(HEAD_DIM), i64(Q_DIM), @@ -1000,7 +1003,7 @@ def draft_precompute(): ] return ds - states = {"kv": {"bytes_per_token": KV_BYTES_PER_TOKEN}} + states = {f"kv.{i}": {"bytes_per_token": LAYER_KV_BYTES_PER_TOKEN} for i in range(LAYERS)} programs = { # prefill: one sequence, as many rows as the chunk holds (tokens); # decode: the bs=1 contract, 3D split-KV kernels; decode_batch: up From 59590c7821c7509846445cae0ce2705f50876a98 Mon Sep 17 00:00:00 2001 From: JinYan Su <751080330@qq.com> Date: Wed, 7 Oct 2026 15:44:10 +0000 Subject: [PATCH 04/22] manifest: a cut stage carries only the states it touches and says how much of each it uses A stage keeps a state whole or drops it, so a whole-model state is carried by every stage that touches a piece of it. `kern cut` now prints, per stage, the regions (distinct call offsets) of each kept state its calls touch against the whole manifest's; a share below one is memory held for other stages, reported, not refused. A page or slot table whose state a stage dropped indexes a kept state of the same kind instead of losing its domain (and with it its place in the protocol). Gate (tray07, qwen3-4b per-layer kv cut before l18): each stage keeps 18 states, all used whole (the single-state manifest: `kv` 36 of 72 regions on both stages); pages per stage without --capacity 122993 -> 245760; the chain gate (7 prompts, chunk 512 and 64) and the 2-process serve gate are bitwise identical as before; sweep unchanged within noise. Co-Authored-By: Claude Opus 5.5 (1M context) --- crates/kern-manifest/src/cut.rs | 38 ++++++++++++++++++++++++++++--- crates/kern-manifest/tests/cut.rs | 10 +++++++- crates/kern-run/src/bin/kern.rs | 9 ++++++++ docs/manifest.md | 7 ++++++ docs/roadmap.md | 3 ++- docs/serve.md | 8 ++++--- 6 files changed, 67 insertions(+), 8 deletions(-) diff --git a/crates/kern-manifest/src/cut.rs b/crates/kern-manifest/src/cut.rs index 1023c6fe..c5b3e4ea 100644 --- a/crates/kern-manifest/src/cut.rs +++ b/crates/kern-manifest/src/cut.rs @@ -30,7 +30,13 @@ //! The prompt reaches the step program's first token the way it does //! unsplit; a chunk program that hands a token back itself is refused, so //! a pipeline's first token always comes from a step. Everything a stage -//! does not use is dropped, and every stage verifies. +//! does not use is dropped, and every stage verifies. A state is kept whole +//! or not at all, so a state laid out over the whole model (one KV state +//! for every layer) is carried whole by every stage that touches a piece of +//! it; [`state_use`] says how much of each one a stage uses, and a +//! generator that gives each piece its own state (one per layer) lets a +//! stage carry only its own. A table whose state a stage dropped indexes a +//! kept state of the same kind: a pool hands every paged state the same ids. use std::collections::{BTreeMap, BTreeSet}; @@ -356,9 +362,17 @@ fn stage( let states: BTreeMap = calls.iter().flat_map(|c| state_args(c)).map(|(st, _)| (st.to_string(), m.states[st].clone())).collect(); let names: BTreeSet = buffers.keys().chain(states.keys()).cloned().collect(); + // A table hands the same ids to every state of its kind: one naming a + // state this stage dropped names a kept state of that kind instead. + let kind = |s: &State| (s.bytes_per_token > 0, s.bytes_per_seq > 0, s.is_owned()); for b in buffers.values_mut() { - if b.domain.as_ref().and_then(|d| d.index_into.as_ref()).is_some_and(|t| !names.contains(t)) { - b.domain = None; + let Some(t) = b.domain.as_ref().and_then(|d| d.index_into.clone()).filter(|t| !names.contains(t)) else { + continue; + }; + let to = m.states.get(&t).and_then(|old| states.iter().find(|(_, s)| kind(s) == kind(old))).map(|(n, _)| n); + match (to, b.domain.as_mut()) { + (Some(n), Some(d)) => d.index_into = Some(n.clone()), + _ => b.domain = None, } } let ops: BTreeMap = @@ -468,6 +482,24 @@ fn live_in<'a>(m: &'a Manifest, calls: &'a [Call]) -> BTreeSet<&'a str> { live } +/// How much of each state a stage keeps it uses: per state, the regions +/// (distinct call offsets) the stage's calls touch and those the whole +/// manifest's calls touch. A stage allocates a state whole, so a share +/// below one is memory it holds for another stage's regions: a state laid +/// out over the whole model rather than one per piece. +pub fn state_use(whole: &Manifest, stage: &Manifest) -> Vec<(String, usize, usize)> { + let regions = |m: &Manifest, st: &str| -> BTreeSet { + m.programs + .values() + .flat_map(|p| &p.calls) + .flat_map(state_args) + .filter(|(s, _)| *s == st) + .map(|(_, o)| o) + .collect() + }; + stage.states.keys().map(|st| (st.clone(), regions(stage, st).len(), regions(whole, st).len())).collect() +} + /// Every state region `c` touches. fn state_args(c: &Call) -> impl Iterator { c.args.iter().filter_map(|a| match a { diff --git a/crates/kern-manifest/tests/cut.rs b/crates/kern-manifest/tests/cut.rs index 748f1135..51acd8e4 100644 --- a/crates/kern-manifest/tests/cut.rs +++ b/crates/kern-manifest/tests/cut.rs @@ -3,7 +3,7 @@ use std::collections::BTreeMap; -use kern_manifest::cut::cut; +use kern_manifest::cut::{cut, state_use}; use kern_manifest::types::{Arg, Dir, Manifest, Module, ParamType}; use kern_manifest::{verify, Protocol, Verified}; use serde_json::{json, Value}; @@ -268,6 +268,14 @@ fn qwen3_cuts_at_a_layer_into_stages_that_each_serve() { ); assert!(tail.buffers.contains_key("lm_head.weight") && !head.buffers.contains_key("lm_head.weight")); assert!(!head.buffers.contains_key("model.layers.18.self_attn.qkv_proj.weight")); + // A KV state per layer: each stage keeps its 18 and uses them whole, + // and its tables index one it kept. + for (st, first) in [(head, "kv.0"), (tail, "kv.18")] { + let used = state_use(&qwen3(), st); + assert_eq!((used.len(), used.iter().all(|(_, a, b)| a == b)), (18, true)); + let table = |b: &str| st.buffers[b].domain.as_ref().and_then(|d| d.index_into.clone()); + assert_eq!((table("block_table"), table("slot_mapping")), (Some(first.into()), Some(first.into()))); + } // The head stages tokens and hands nothing back; the tail is fed // activations and samples. let ph = Protocol::check_unsampled(head).unwrap_or_else(|e| panic!("{e}")); diff --git a/crates/kern-run/src/bin/kern.rs b/crates/kern-run/src/bin/kern.rs index d98e1ee0..f9214c1d 100644 --- a/crates/kern-run/src/bin/kern.rs +++ b/crates/kern-run/src/bin/kern.rs @@ -162,6 +162,15 @@ fn cut(manifest: &Path, at: &[usize], module: &Path, out: &Path) -> Result<()> { st.buffers.values().filter(|b| b.kind == kern_manifest::types::BufferKind::Weight).count(), calls.join(", ") ); + let (whole, part): (Vec<_>, Vec<_>) = + kern_manifest::cut::state_use(&m, st).into_iter().partition(|(_, a, b)| a == b); + tracing::info!("stage {s}: {} of {} states used whole", whole.len(), whole.len() + part.len()); + for (name, a, b) in part { + tracing::warn!( + "stage {s}: state `{name}` is allocated whole and its calls touch {a} of its {b} regions (~{:.0}% of its bytes)", + 100.0 * a as f64 / b as f64 + ); + } } tracing::info!("cut `{}` into {} stages", stages[0].cut.as_ref().map_or("", |c| c.id.as_str()), stages.len()); Ok(()) diff --git a/docs/manifest.md b/docs/manifest.md index fb438dab..220eae69 100644 --- a/docs/manifest.md +++ b/docs/manifest.md @@ -402,6 +402,13 @@ kern-serve 只拿 `Protocol` 驱动 runtime，不读 JSON（CI 用 grep 保证 `var × 常数`）；output / carry 只许在末段写；一个 (state, offset) 区域 只许一段碰（按层切 KV 时各段各持自己那些层）。做不到的切点报错并说明 哪条不满足，不会静默切出错的东西。 +- state 整个留或整个丢：一段只要碰到一个 state 的一块就分配整个。`kern cut` + 逐段打印每个 state 被本段碰到的区域数 / 整份碰到的区域数（区域 = 不同的 + call offset，按等大估字节），不满的就是给别段占着的内存（报告，不拒）。 + 所以生成器把 KV 写成每层一个 state（qwen3-4b：`kv.0`…`kv.35`），各段只带 + 自己那些层；生成器不知道流水线。页表 / slot 的 `index_into` 指向一个本段 + 丢掉的 state 时改指本段留下的同类 state（pool 给所有 paged state 同一套 + 页号）。 - 边界处插一对 send / recv：上游段末尾 `pp.wait_empty → pp.put.… → pp.post_full`，下游段开头 `pp.wait_full → pp.take.… → pp.post_empty`；邮箱 `pp..box`（`u8[256 + 最大载荷]`）在下游 GPU 上， diff --git a/docs/roadmap.md b/docs/roadmap.md index 872c5e50..c44e78ff 100644 --- a/docs/roadmap.md +++ b/docs/roadmap.md @@ -63,7 +63,8 @@ kern 在内部 RSI 里是 kernel agent 的反馈环境（`kern test` 裁判、`k | 级 | 内容 | 门禁 | |---|---|---| | PP1 ✅ | `kern cut`（按 call 下标切、活跃性找边界、send/recv 邮箱核、`clock` fill）+ kern-serve 一段一进程的 prefill-only 流水（head 持唯一的池，TCP 发 item，GPU 邮箱传激活） | **2026-10-07 tray07 过**：qwen3-4b 2 段，各段依次跑 logits/token 与整份逐位同（chunk 512/64）；两进程首 token 与单卡同；8193 token conc1 TTFT 313 → 174 ms（serve.md「流水线」） | -| PP2 | 每段只分自己那些层的 KV（现在按整份 `bytes_per_token`）；decode 走流水；per-seq state（GDN）与吐 token 的 chunk program | | +| PP1.1 ✅ | qwen3-4b 的 KV 每层一个 state，`kern cut` 每段只带自己碰的 state、逐段报告 state 用了几成 | **2026-10-07 tray07 过**：`kern test qwen3-4b` 位一致、`kern run` 文本同；chain / serve gate 位一致；每段页数 122,993 → 245,760 | +| PP2 | decode 走流水；K3 的 state 拆分；per-seq state（GDN）与吐 token 的 chunk program | | ## 协议线 diff --git a/docs/serve.md b/docs/serve.md index 02c8e101..a3cbc822 100644 --- a/docs/serve.md +++ b/docs/serve.md @@ -605,9 +605,11 @@ stop、fault 消息：任何错误进程退出，连接断了对端退出；设 | | | | | | | 1 | 9.52 | 2.13 | 0.73 | - | 单请求内 chunk 之间就流水起来（stage 0 算第 k+1 块时 stage 1 算第 k 块），所以长 prompt conc1 也 -接近 2×；1 块的 prompt 没有可重叠的东西，略慢（邮箱 + 两段 launch）。v1 的限制：每段按整份的 -`bytes_per_token` 分 KV（只用自己那些层，浪费一半）；只认纯 paged 的 manifest；chunk program -自己吐 token 的模型拒切；上游那份 box 是不用的占位。 +接近 2×；1 块的 prompt 没有可重叠的东西，略慢（邮箱 + 两段 launch）。KV 每层一个 state 之后（同日， +`examples/qwen3-4b.json` 重生成，`kern test qwen3-4b` 位一致、`kern run` 文本同），每段只带自己的 +18 层：不给 `--capacity` 时每段 122,993 → 245,760 页（2.0×）；切点、gate、扫描数字与上表一致 +（`~/bench_results/2026-10-07-kern-pp-v1/perlayer/`）。v1 的限制：只认纯 paged 的 manifest； +chunk program 自己吐 token 的模型拒切；上游那份 box 是不用的占位。 ## 没做（按需要加） From 0bbf81a88f5ca85e069e9ab2dfce5972a494ac4c Mon Sep 17 00:00:00 2001 From: JinYan Su <751080330@qq.com> Date: Wed, 7 Oct 2026 15:50:11 +0000 Subject: [PATCH 05/22] qwen3: the example hash baselines follow the per-layer kv The regenerated qwen3-4b examples carry one kv state per layer and the registry builds of argmax, embedding and weight_prep; kern test and the pipeline gates on PR #42 cover the change. Co-Authored-By: Claude Opus 5.5 (1M context) --- tools/test_qwen_constants.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/tools/test_qwen_constants.py b/tools/test_qwen_constants.py index 2b67087d..9ffe542e 100644 --- a/tools/test_qwen_constants.py +++ b/tools/test_qwen_constants.py @@ -9,12 +9,13 @@ ROOT = pathlib.Path(__file__).resolve().parent.parent # Canonical JSON hashes of the checked-in examples (last: unified attention -# reads the kv state, `inout` -> `in`, 2026-09-16). Changes to kernel wiring, +# reads the kv state, `inout` -> `in`, 2026-09-16; qwen3-4b: one kv state per +# layer and registry builds of argmax / embedding / weight_prep, 2026-10-07). Changes to kernel wiring, # capacities, float literals or even artifact pins require review. BASELINES = { - "qwen3-4b-dspark.json": "444001bf6fa6991128128a878c0cc139b8ecf71800c6ccfea0d344c1ad169133", - "qwen3-4b-silu-mined.json": "e14198af226849f46826bc624c4a6e8a208a4e8d7b160ffc611b817d2915d7b8", - "qwen3-4b.json": "792666b476938520a21598bec6bdf5461cb62dd8fdf52cd5998b33a7388c2e69", + "qwen3-4b-dspark.json": "c72a03b5c0f817e742befcffa1e8241bc13f92c4f3458743a8ee27ccc6e4d30c", + "qwen3-4b-silu-mined.json": "3cd79ecd0d0d54c830f842644138b5dafb92c88efcf283709535b096665fe23e", + "qwen3-4b.json": "b54ac6fa2766667782626ec56b8d02ceb46b250eee7f3cf1a6d79561e5d58e84", "qwen3.8-27b-dflash2.json": "134839dd1b25fc8fb5a5cb97391a02576d2624f2f0ccc5431af6e10c9dd679dd", "qwen3.8-27b.json": "a9cc996a7548dabc756b30f9ef848459c2a5c2043075b79e11091054d1a30443" } From f83b8ace6eaaa07a8066e7f947ad09b6c1b28a93 Mon Sep 17 00:00:00 2001 From: JinYan Su <751080330@qq.com> Date: Wed, 7 Oct 2026 15:54:05 +0000 Subject: [PATCH 06/22] dsv41: ar_trace skips the clock fill like the other outputs Co-Authored-By: Claude Opus 5.5 (1M context) --- tools/dsv41/attention/ar_trace/src/main.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tools/dsv41/attention/ar_trace/src/main.rs b/tools/dsv41/attention/ar_trace/src/main.rs index 3c56175a..7a9015b0 100644 --- a/tools/dsv41/attention/ar_trace/src/main.rs +++ b/tools/dsv41/attention/ar_trace/src/main.rs @@ -104,7 +104,7 @@ fn main() -> anyhow::Result<()> { Fill::CuSeqlens => vec![0, rows as i64], Fill::SpanAt => vec![0], Fill::Blocks => anyhow::bail!("unexpected collective row blocks in DP manifest"), - Fill::Tokens | Fill::Count | Fill::Error => continue, + Fill::Tokens | Fill::Count | Fill::Error | Fill::Clock => continue, }; rt.write_input_at(&f.name, &f.encode(&values), &vars)?; } @@ -159,7 +159,7 @@ fn main() -> anyhow::Result<()> { Fill::CuSeqlens => vec![0, 6], Fill::SpanAt => vec![0], Fill::Blocks => anyhow::bail!("unexpected blocks"), - Fill::Tokens | Fill::Count | Fill::Error => continue, + Fill::Tokens | Fill::Count | Fill::Error | Fill::Clock => continue, }; rt.write_input_at(&f.name, &f.encode(&values), &dv)?; } From 2969958b0f39ebdb8ea7f2dce404a61d0ccf6713 Mon Sep 17 00:00:00 2001 From: JinYan Su <751080330@qq.com> Date: Wed, 7 Oct 2026 15:55:35 +0000 Subject: [PATCH 07/22] 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 5f28a060..1bd412bd 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 { @@ -421,6 +424,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 { @@ -611,6 +615,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 af89c7ce..a8a27e2f 100644 --- a/crates/kern-manifest/src/types.rs +++ b/crates/kern-manifest/src/types.rs @@ -151,6 +151,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, @@ -371,7 +374,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, } @@ -597,6 +600,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. @@ -611,6 +616,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 3aa8bd64..0afc4860 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 @@ -279,12 +281,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 @@ -349,9 +354,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 @@ -772,7 +776,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(); @@ -784,6 +791,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)")); } @@ -829,6 +841,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())); @@ -879,6 +897,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", @@ -886,7 +914,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()); @@ -898,6 +929,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 220eae69..6edab8b1 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 组：只有名字和大小 cut cut 流水线的一段：id、第几段、共几段 ``` @@ -51,7 +51,8 @@ cut cut 流水线的一段：id、第几段、共 - `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 填的地址数组，见 @@ -139,8 +140,8 @@ cut cut 流水线的一段：id、第几段、共 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 布局算术的字面量， @@ -321,6 +322,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 的 @@ -355,6 +367,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）；行数随调用变的 @@ -456,7 +470,8 @@ logits 与 token 与整份逐位同（serve.md「流水线」）。 能接 `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 28e7d4a7..e44071a7 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`.", @@ -1686,6 +1691,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 From 57b8bfeb035460acb0a369c16348b9e15e3c9fd7 Mon Sep 17 00:00:00 2001 From: JinYan Su <751080330@qq.com> Date: Wed, 7 Oct 2026 16:01:49 +0000 Subject: [PATCH 08/22] manifest: a cut slices derive programs per stage and lets the last stage's chunk hand back the token A stage derives only the weights its own calls read, like the tables of a once program. A chunk program that writes the first token (a model whose chunked kernels must see every prompt token, K3) is no longer refused: only the last stage writes an output, so the token comes from it either way. Signed-off-by: JinYan Su <751080330@qq.com> Co-Authored-By: Claude Opus 5.5 (1M context) --- crates/kern-manifest/src/cut.rs | 22 +++++++++------------- 1 file changed, 9 insertions(+), 13 deletions(-) diff --git a/crates/kern-manifest/src/cut.rs b/crates/kern-manifest/src/cut.rs index c5b3e4ea..f0c892ab 100644 --- a/crates/kern-manifest/src/cut.rs +++ b/crates/kern-manifest/src/cut.rs @@ -24,12 +24,13 @@ //! what makes the stages' page tables interchangeable: a sequence's pages //! hold the first stage's layers on one GPU and the next stage's on //! another, under the same page ids. A region two stages touch is refused. -//! A program run once after load is sliced per stage to the calls the -//! stage's own calls need, so every stage derives its own tables. +//! A program run once after load, and one deriving weights while they +//! load, is sliced per stage to the calls the stage's own calls need, so +//! every stage derives its own tables and weights. //! -//! The prompt reaches the step program's first token the way it does -//! unsplit; a chunk program that hands a token back itself is refused, so -//! a pipeline's first token always comes from a step. Everything a stage +//! The prompt reaches its first token the way it does unsplit, from the +//! step or from a chunk program that hands one back; either way the last +//! stage writes it, since only the last stage writes an output. Everything a stage //! does not use is dropped, and every stage verifies. A state is kept whole //! or not at all, so a state laid out over the whole model (one KV state //! for every layer) is carried whole by every stage that touches a piece of @@ -109,12 +110,6 @@ pub fn cut(m: &Verified, at: &[usize], pp: &Module) -> Result, Cut errs.push(format!("program `{name}` has {} calls; it cannot be cut before call {last}", p.calls.len())); } } - let chunk = p.batch.as_ref().is_some_and(|b| matches!(b.rows, Dim::Var(_))); - if chunk && writes(m, &p.calls).iter().any(|b| m.buffers[*b].fill == Some(Fill::Tokens)) { - errs.push(format!( - "program `{name}` takes a chunk and hands a token back; a pipeline takes its first token from a step" - )); - } } if !errs.is_empty() { return Err(CutErrors(errs)); @@ -299,7 +294,8 @@ fn pp_op(name: &str) -> Op { } /// Stage `s` of `stages`: `programs` (already cut and wired), the once -/// programs sliced to what they need, and every declaration they use. +/// and derive programs sliced to what they need, and every declaration +/// they use. fn stage( m: &Manifest, mut programs: BTreeMap, @@ -311,7 +307,7 @@ fn stage( ) -> Manifest { let mut needed: BTreeSet = programs.values().flat_map(|p| &p.calls).flat_map(|c| reads(m, c)).map(str::to_string).collect(); - for (name, p) in m.programs.iter().filter(|(_, p)| p.once) { + for (name, p) in m.programs.iter().filter(|(_, p)| p.once || p.derive) { let mut kept = Vec::new(); for c in p.calls.iter().rev() { if accesses(m, c).iter().any(|(b, _, d)| *d != Dir::In && needed.contains(*b)) { From dfd0b9c225dfca83a79da975348a51f5c849657e Mon Sep 17 00:00:00 2001 From: JinYan Su <751080330@qq.com> Date: Wed, 7 Oct 2026 16:01:49 +0000 Subject: [PATCH 09/22] runtime: a runtime zeroes a sequence slot another pool handed out A pipeline stage holds the head's slot ids for its own layers' recurrent state, so it zeroes a fresh sequence's slot on the head's word rather than through a lease of its own. kern-run writes a line table by name, whole. Signed-off-by: JinYan Su <751080330@qq.com> Co-Authored-By: Claude Opus 5.5 (1M context) --- crates/kern-run/src/lib.rs | 14 ++++++++++---- crates/kern-runtime/src/lease.rs | 22 ++++++++++++++++++++++ 2 files changed, 32 insertions(+), 4 deletions(-) diff --git a/crates/kern-run/src/lib.rs b/crates/kern-run/src/lib.rs index a451b617..19cea7c3 100644 --- a/crates/kern-run/src/lib.rs +++ b/crates/kern-run/src/lib.rs @@ -175,12 +175,18 @@ pub fn fills<'p>( } /// Write `values` into input `name` of `rt` the way `p` reads it: a fill -/// in its dtype, a page table as i32. `false` when the manifest has no -/// such input (a pipeline stage that never reads it). +/// in its dtype, a page table as i32, a line table as i32 and whole (its +/// rows sit at the bound of their var, see [`line_rows`]). `false` when the +/// manifest has no such input (a pipeline stage that never reads it). pub fn write_named(rt: &mut Runtime, p: &Protocol, name: &str, values: &[i64], vars: &Vars) -> Result { + let i32s = |v: &[i64]| -> Vec { v.iter().flat_map(|&x| (x as i32).to_le_bytes()).collect() }; + if p.line_tables.iter().any(|t| t.name == name) { + rt.write_input(name, &i32s(values))?; + return Ok(true); + } let bytes = match (p.fills.iter().find(|f| f.name == name), p.page_tables.iter().any(|t| t.name == name)) { (Some(f), _) => f.encode(values), - (None, true) => values.iter().flat_map(|&v| (v as i32).to_le_bytes()).collect(), + (None, true) => i32s(values), (None, false) => return Ok(false), }; rt.write_input_at(name, &bytes, vars)?; @@ -217,7 +223,7 @@ fn page_rows(t: &PageTable, leases: &[Lease], rows: usize) -> Result> { /// line in entry 0 and zeros through the rest of a wide cell (a program /// that moves along one does so on the device). Columns past the leases /// given repeat the last. -fn line_rows(t: &LineTable, leases: &[Lease], cols: usize) -> Result> { +pub fn line_rows(t: &LineTable, leases: &[Lease], cols: usize) -> Result> { let mut v = Vec::with_capacity(t.lines * cols * t.width); for line in 0..t.lines { for c in 0..cols { diff --git a/crates/kern-runtime/src/lease.rs b/crates/kern-runtime/src/lease.rs index 038fb14b..c2e1981e 100644 --- a/crates/kern-runtime/src/lease.rs +++ b/crates/kern-runtime/src/lease.rs @@ -128,6 +128,28 @@ impl Runtime { Ok(lease) } + /// Zero sequence slot `slot` of every per-sequence state on the + /// stream, the way a lease zeroes its own: for a runtime whose slots + /// another runtime's pool hands out (a pipeline stage holds the head's + /// slot ids for its own layers' recurrent state). + pub fn zero_slot(&mut self, slot: i32) -> Result<()> { + self.ctx.bind_to_thread()?; + let live = self.pool.slots(); + if slot < 1 || slot as usize >= live { + bail!(Api, "slot {slot}: this runtime has slots 1..{live}"); + } + for (name, st) in &self.manifest.states { + if !st.is_per_seq() { + continue; + } + let b = st.bytes_per_seq as usize; + let s = slot as usize; + let mut view = self.states.get_mut(name).unwrap().view(s * b..(s + 1) * b)?; + self.stream.memset_zeros(&mut view)?; + } + Ok(()) + } + /// The first `len` tokens of `lease` as a [`Checkpoint`] its /// sequence keeps running past: pages shared, the per-sequence state /// copied into a fresh slot ([`Error::Denied`] when none is free). From 41fe30e62d3da53681114c87c8790c3179556129 Mon Sep 17 00:00:00 2001 From: JinYan Su <751080330@qq.com> Date: Wed, 7 Oct 2026 16:01:50 +0000 Subject: [PATCH 10/22] kern-serve: a pipeline serves per-sequence state and a chunk program that emits the token An item carries the sequence's line tables and the slot to zero on its first call, and says whether it finishes the prompt; the last stage answers only that one. Without a one-row step the whole prompt goes through the chunk program and the last chunk's token is the answer (full K3 PP8). Signed-off-by: JinYan Su <751080330@qq.com> Co-Authored-By: Claude Opus 5.5 (1M context) --- crates/kern-serve/src/pp/head.rs | 61 ++++++++++++++++++++++--------- crates/kern-serve/src/pp/mod.rs | 13 +++++-- crates/kern-serve/src/pp/stage.rs | 10 +++-- 3 files changed, 61 insertions(+), 23 deletions(-) diff --git a/crates/kern-serve/src/pp/head.rs b/crates/kern-serve/src/pp/head.rs index 0af0399c..413b37f3 100644 --- a/crates/kern-serve/src/pp/head.rs +++ b/crates/kern-serve/src/pp/head.rs @@ -1,9 +1,12 @@ //! The head: stage 0 of a pipeline, the scheduler every stage follows. //! -//! A request leases its prompt's pages from the head's pool, the one pool -//! of the pipeline, and goes out as items: the prompt but its last token -//! in chunks through the chunk program, then the last token through the -//! one-row step, whose tokens the last stage hands back. Each item is sent +//! A request leases its prompt's pages (and slot) from the head's pool, +//! the one pool of the pipeline, and goes out as items: the prompt but its +//! last token in chunks through the chunk program, then the last token +//! through the one-row step, whose tokens the last stage hands back; or, +//! when the chunk program hands a token back itself (a model whose chunked +//! kernels must see every prompt token), the whole prompt in chunks, the +//! last chunk's token the answer. Each item is sent //! to every stage before the head stages and issues it itself, so a stage //! waiting on its device for an item always has it on the way. The first //! token finishes the request. Nothing waits on the head's device but the @@ -14,7 +17,7 @@ use std::collections::{BTreeMap, VecDeque}; use std::sync::mpsc; use anyhow::{bail, ensure, Context, Result}; -use kern_manifest::protocol::{Forward, Rows}; +use kern_manifest::protocol::{Axis, Forward, Rows}; use kern_manifest::{Protocol, Verified}; use kern_pool::{Denied, Lease}; use kern_run::Weights; @@ -50,7 +53,9 @@ pub struct PpHead { rt: Runtime, p: Protocol, prefill: Forward, - step: Forward, + /// The one-row step the last prompt token goes through; none when the + /// chunk program hands the first token back itself. + step: Option, chunk: usize, max_seqs: usize, stop_tokens: Vec, @@ -107,7 +112,6 @@ pub fn load(m: Verified, weights: &Weights, h: Head) -> Result { connect_downstream(&mut rt, &hellos[&1].0.mailbox)?; let p = Protocol::check_unsampled(&rt.manifest)?; kern_run::run_once(&rt, &p)?; - ensure!(p.line_tables.is_empty(), "a pipeline serves paged state only; this manifest has per-sequence states"); let pages = rt.pages_total() as u64; let downstream: Vec> = @@ -123,7 +127,10 @@ pub fn load(m: Verified, weights: &Weights, h: Head) -> Result { let prefill = p.chunk().cloned().context("a pipeline prefills through a program taking one sequence's rows as fed")?; - let step = p.forward(1, Rows::Const(1)).cloned().context("no program takes one sequence of one row")?; + let step = match prefill.emits { + Some(_) => None, + None => Some(p.forward(1, Rows::Const(1)).cloned().context("no program takes one sequence of one row")?), + }; let max = p.rows.max as usize; let chunk = match h.chunk { None => max, @@ -131,7 +138,8 @@ pub fn load(m: Verified, weights: &Weights, h: Head) -> Result { Some(c) => bail!("--chunk {c}: the manifest's `tokens` bound is {max}; a chunk can only be smaller"), }; let clock = Clock::new(&p)?; - info!(pages, page = rt.page(), chunk, prefill = %prefill.name, step = %step.name, "pipeline head ready"); + let step_name = step.as_ref().map_or("-", |f| f.name.as_str()).to_string(); + info!(pages, page = rt.page(), chunk, prefill = %prefill.name, step = %step_name, "pipeline head ready"); Ok(PpHead { rt, p, @@ -221,21 +229,27 @@ impl PpHead { } /// The prompt as items: all but its last token in chunks, the last - /// through the step, whose token comes back. + /// through the step, whose token comes back; or all of it in chunks + /// when the chunk program hands the token back. fn send_prompt(&mut self, q: &QueuedRequest, lease: &Lease) -> Result<()> { let ids: Vec = q.request.prompt_tokens.iter().map(|&t| t as i64).collect(); let n = ids.len(); let (prefill, step) = (self.prefill.clone(), self.step.clone()); + let chunked = if step.is_some() { n - 1 } else { n }; let mut pos = 0; - while pos < n - 1 { - let c = (n - 1 - pos).min(self.chunk); - self.send_item(&prefill, lease, pos, &ids[pos..pos + c])?; + while pos < chunked { + let c = (chunked - pos).min(self.chunk); + let reply = step.is_none() && pos + c == n; + self.send_item(&prefill, lease, pos, &ids[pos..pos + c], reply)?; pos += c; } - self.send_item(&step, lease, n - 1, &ids[n - 1..]) + match &step { + Some(f) => self.send_item(f, lease, n - 1, &ids[n - 1..], true), + None => Ok(()), + } } - fn send_item(&mut self, f: &Forward, lease: &Lease, pos: usize, ids: &[i64]) -> Result<()> { + fn send_item(&mut self, f: &Forward, lease: &Lease, pos: usize, ids: &[i64], reply: bool) -> Result<()> { let (vars, values) = kern_run::fills(&self.p, std::slice::from_ref(lease), &[pos], ids.len(), ids); let mut rows: BTreeMap> = values.into_iter().map(|(f, v)| (f.name.clone(), v)).collect(); for t in &self.p.page_tables { @@ -243,9 +257,22 @@ impl PpHead { lease.extend_row(&t.name, &mut v)?; rows.insert(t.name.clone(), v.into_iter().map(i64::from).collect()); } - let item = Item { program: f.name.clone(), vars, rows }; + // A sequence's line tables and slot are the same on every item: its + // first item carries them and zeroes the slot on every stage (the + // head's own lease zeroed its slot already). + let first = pos == 0; + for t in self.p.line_tables.iter().filter(|_| first) { + let cols = match t.axis { + Axis::Tray => self.p.tray.as_ref().map_or(1, |b| b.max), + _ => self.p.groups.max, + }; + let v = kern_run::line_rows(t, std::slice::from_ref(lease), cols as usize)?; + rows.insert(t.name.clone(), v.into_iter().map(i64::from).collect()); + } + let item = Item { program: f.name.clone(), vars, rows, zero_slot: None, reply }; + let out = Item { zero_slot: lease.seq_slot().filter(|_| first), ..item.clone() }; for w in &mut self.stages { - self.io.block_on(send(w, &item))?; + self.io.block_on(send(w, &out))?; } for (name, v) in &item.rows { kern_run::write_named(&mut self.rt, &self.p, name, v, &item.vars)?; diff --git a/crates/kern-serve/src/pp/mod.rs b/crates/kern-serve/src/pp/mod.rs index 0c0eaf4c..c43ca2b1 100644 --- a/crates/kern-serve/src/pp/mod.rs +++ b/crates/kern-serve/src/pp/mod.rs @@ -17,8 +17,8 @@ //! handles, to map its own outgoing peer from, and the pages. From then on //! the connection carries the log: the head sends every stage every //! [`Item`] in issue order — the program, its vars and the values of each -//! input it staged — and the last stage answers each item whose program -//! hands tokens back with the tokens. TCP keeps the order; there is no +//! input it staged — and the last stage answers the item that finishes a +//! prompt with its tokens. TCP keeps the order; there is no //! sequence number, no stop message and no fault message. A stage that //! fails exits; a closed connection ends the process on the other side. //! @@ -67,8 +67,15 @@ pub struct Table { pub struct Item { pub program: String, pub vars: BTreeMap, - /// Each input's values by name: the fills and the page tables. + /// Each input's values by name: the fills, the page tables and, on a + /// sequence's first item, its line tables. pub rows: BTreeMap>, + /// The sequence slot to zero before the call: on its first item, the + /// recurrent state a fresh lease starts from. + pub zero_slot: Option, + /// Whether the last stage hands this call's tokens back: the item + /// that finishes a prompt. + pub reply: bool, } pub type Handles = BTreeMap; diff --git a/crates/kern-serve/src/pp/stage.rs b/crates/kern-serve/src/pp/stage.rs index 64ddb35f..990ef529 100644 --- a/crates/kern-serve/src/pp/stage.rs +++ b/crates/kern-serve/src/pp/stage.rs @@ -1,8 +1,9 @@ //! A downstream stage: loads its slice of the model, says hello to the //! head, maps the next stage's mailbox from the table it gets back, then //! replays the head's items in order. It keeps no accounting: the page -//! ids, positions and slots of every item are the head's. The last stage -//! hands back the tokens of every item whose program emits them. +//! ids, positions and slots of every item are the head's, and so is the +//! zeroing of a fresh sequence's slot. The last stage hands back the +//! tokens of the item that finishes a prompt. use std::sync::mpsc; use std::time::{Duration, Instant}; @@ -107,13 +108,16 @@ pub fn run(m: Verified, weights: &Weights, o: Stage) -> Result<()> { Err(mpsc::TryRecvError::Disconnected) => break, }; let vars: Vars = item.vars.into_iter().filter(|(k, _)| rt.manifest.vars.contains_key(k)).collect(); + if let Some(slot) = item.zero_slot { + rt.zero_slot(slot)?; + } for (name, v) in &item.rows { kern_run::write_named(&mut rt, &p, name, v, &vars)?; } rt.issue(&item.program, &vars).with_context(|| format!("`{}`", item.program))?; clock.issued(vars.get(&p.rows.var).copied().unwrap_or(0)); let emits = p.forwards.iter().find(|f| f.name == item.program).and_then(|f| f.emits); - if let Some(i) = emits.filter(|_| last) { + if let Some(i) = emits.filter(|_| last && item.reply) { let t = &p.fills[i]; let mut tokens = t.decode(&rt.read_output(&t.name)?); tokens.truncate(t.width as usize); From 18c14c21c286fe2edd29e74b79e6153e9ebdcfba Mon Sep 17 00:00:00 2001 From: JinYan Su <751080330@qq.com> Date: Wed, 7 Oct 2026 16:07:08 +0000 Subject: [PATCH 11/22] manifest: a downstream cut stage without a seq_len fill takes its groups from the paged table Stage k>0 of a cut chunk program gets the sequence count from the mailbox, not a seq_len input, so the protocol takes the groups axis from the first dim of the unfilled input that indexes a paged state. Signed-off-by: JinYan Su <751080330@qq.com> Co-Authored-By: Claude Opus 5.5 (1M context) --- crates/kern-manifest/src/protocol.rs | 21 ++++++++++++++++++++- 1 file changed, 20 insertions(+), 1 deletion(-) diff --git a/crates/kern-manifest/src/protocol.rs b/crates/kern-manifest/src/protocol.rs index 1bd412bd..80de8c9a 100644 --- a/crates/kern-manifest/src/protocol.rs +++ b/crates/kern-manifest/src/protocol.rs @@ -265,7 +265,26 @@ impl Protocol { } }; let rows = axis_of(Fill::Slot, &mut errs); - let groups = axis_of(Fill::SeqLen, &mut errs); + // A stage past the first of a cut may read no sequence length (its + // attention lengths came through the mailbox): its sequences are + // then the rows of its page table, or none of its own. + let downstream = m.cut.as_ref().is_some_and(|c| c.stage > 0); + let groups = if downstream && !m.buffers.values().any(|b| b.fill == Some(Fill::SeqLen)) { + let paged = |b: &Buffer| { + b.domain.as_ref().and_then(|d| d.index_into.as_deref()).and_then(|s| m.states.get(s)).is_some_and( + |s| s.bytes_per_token > 0, + ) + }; + let var = m + .buffers + .values() + .filter(|b| b.kind == BufferKind::Input && b.fill.is_none() && paged(b)) + .find_map(axis_var) + .unwrap_or_default(); + Some(Bound { max: var_max(&var).unwrap_or(1), var }) + } else { + axis_of(Fill::SeqLen, &mut errs) + }; if let (Some(r), Some(g)) = (&rows, &groups) { if r.var == g.var { errs.push(format!( From 7ec497c4ba2ae7acf0832d085ab6cbe4852f8a7d Mon Sep 17 00:00:00 2001 From: JinYan Su <751080330@qq.com> Date: Wed, 7 Oct 2026 16:10:59 +0000 Subject: [PATCH 12/22] kern-serve: a pipeline head without a one-row step feeds the whole prompt as chunks The head runs stage 0's manifest, which cannot tell whether the last stage's chunk emits the token, so the step is optional rather than inferred from the head's own chunk. Fallible setup also runs before the stage listeners start, so its error is not masked by their exit. Signed-off-by: JinYan Su <751080330@qq.com> Co-Authored-By: Claude Opus 5.5 (1M context) --- crates/kern-serve/src/pp/head.rs | 25 ++++++++++++------------- 1 file changed, 12 insertions(+), 13 deletions(-) diff --git a/crates/kern-serve/src/pp/head.rs b/crates/kern-serve/src/pp/head.rs index 413b37f3..0bcf2707 100644 --- a/crates/kern-serve/src/pp/head.rs +++ b/crates/kern-serve/src/pp/head.rs @@ -114,6 +114,18 @@ pub fn load(m: Verified, weights: &Weights, h: Head) -> Result { kern_run::run_once(&rt, &p)?; let pages = rt.pages_total() as u64; + let prefill = + p.chunk().cloned().context("a pipeline prefills through a program taking one sequence's rows as fed")?; + // The head is stage 0: whether the last stage's chunk emits is not in its manifest, so + // a pipeline with no one-row step feeds every prompt row as chunks and the last stage answers. + let step = p.forward(1, Rows::Const(1)).cloned(); + let max = p.rows.max as usize; + let chunk = match h.chunk { + None => max, + Some(c) if (1..=max as u64).contains(&c) => c as usize, + Some(c) => bail!("--chunk {c}: the manifest's `tokens` bound is {max}; a chunk can only be smaller"), + }; + let clock = Clock::new(&p)?; let downstream: Vec> = (1..cut.stages).map(|s| hellos.get(&(s + 1)).map(|(x, ..)| x.mailbox.clone())).collect(); let (tx, replies) = mpsc::channel(); @@ -125,19 +137,6 @@ pub fn load(m: Verified, weights: &Weights, h: Head) -> Result { stages.push(w); } - let prefill = - p.chunk().cloned().context("a pipeline prefills through a program taking one sequence's rows as fed")?; - let step = match prefill.emits { - Some(_) => None, - None => Some(p.forward(1, Rows::Const(1)).cloned().context("no program takes one sequence of one row")?), - }; - let max = p.rows.max as usize; - let chunk = match h.chunk { - None => max, - Some(c) if (1..=max as u64).contains(&c) => c as usize, - Some(c) => bail!("--chunk {c}: the manifest's `tokens` bound is {max}; a chunk can only be smaller"), - }; - let clock = Clock::new(&p)?; let step_name = step.as_ref().map_or("-", |f| f.name.as_str()).to_string(); info!(pages, page = rt.page(), chunk, prefill = %prefill.name, step = %step_name, "pipeline head ready"); Ok(PpHead { From 8bf4d9ea6592154671441bb43ea78b5eae9aea7e Mon Sep 17 00:00:00 2001 From: JinYan Su <751080330@qq.com> Date: Wed, 7 Oct 2026 17:27:27 +0000 Subject: [PATCH 13/22] manifest: a prefill call carries several sequences, split by its cu_seqlens fill A call whose rows are set per call may now span several groups when a `cu_seqlens` fill says where each sequence's rows start; without one it is still an error. The context var becomes the call's sequences' lengths summed (a packed prefill expands every one of them), equal to the old longest-sequence value wherever groups is 1, the only place one is named today. `kern_run::fills` takes per-sequence lengths and lays the sequences back to back; `vars` takes the call's total rows. Gate: K3 PP8 packed prefill (16 sequences per 8192-row item), 192 frozen prompts conc16 on two trays: 47.0k tok/s vs 42.6k unpacked; single-sequence first tokens identical, 187/192 under load with the 5 others near-ties. Signed-off-by: JinYan Su <751080330@qq.com> Co-Authored-By: Claude Opus 5.5 (1M context) --- crates/kern-manifest/src/protocol.rs | 39 +++++++++++++++----------- crates/kern-manifest/tests/protocol.rs | 10 +++++-- crates/kern-run/examples/pp_chain.rs | 2 +- crates/kern-run/src/lib.rs | 39 ++++++++++++++------------ crates/kern-serve/src/pp/head.rs | 2 +- crates/kern-serve/src/tray.rs | 4 +-- tools/kern_manifest.py | 2 +- 7 files changed, 56 insertions(+), 42 deletions(-) diff --git a/crates/kern-manifest/src/protocol.rs b/crates/kern-manifest/src/protocol.rs index 80de8c9a..f4e279e5 100644 --- a/crates/kern-manifest/src/protocol.rs +++ b/crates/kern-manifest/src/protocol.rs @@ -45,9 +45,10 @@ //! //! A kernel sized by the context rather than by the call (a prefill that //! expands every cached row of its sequence) runs over a var the program's -//! `batch` names as `context`: the length of the call's longest sequence -//! once its rows are in, which the driver already knows from the lengths -//! it stages. One var bounds every context in a manifest. +//! `batch` names as `context`: the call's sequences' lengths once their +//! rows are in, summed (a call of several sequences expands them all), which +//! the driver already knows from the lengths it stages. One var bounds every +//! context in a manifest. use crate::types::*; use crate::Verified; @@ -127,7 +128,8 @@ pub struct LineTable { pub enum Rows { /// Exactly this many, the layout the program's kernels expect. Const(u64), - /// As many as the call feeds: one sequence, the rows var set per call. + /// As many as the call feeds, the rows var set per call: one + /// sequence's, or several back to back split by the `cu_seqlens` fill. Var, } @@ -271,9 +273,11 @@ impl Protocol { let downstream = m.cut.as_ref().is_some_and(|c| c.stage > 0); let groups = if downstream && !m.buffers.values().any(|b| b.fill == Some(Fill::SeqLen)) { let paged = |b: &Buffer| { - b.domain.as_ref().and_then(|d| d.index_into.as_deref()).and_then(|s| m.states.get(s)).is_some_and( - |s| s.bytes_per_token > 0, - ) + b.domain + .as_ref() + .and_then(|d| d.index_into.as_deref()) + .and_then(|s| m.states.get(s)) + .is_some_and(|s| s.bytes_per_token > 0) }; let var = m .buffers @@ -469,8 +473,11 @@ impl Protocol { rows.var )); } - if batch.groups != 1 { - errs.push(format!("{ctx}: rows set per call means one sequence, not {} groups", batch.groups)); + if batch.groups != 1 && one(Fill::CuSeqlens).is_none() { + errs.push(format!( + "{ctx}: rows set per call over {} groups needs a `cu_seqlens` fill to split them", + batch.groups + )); } Rows::Var } @@ -677,18 +684,18 @@ impl Protocol { self.forwards.iter().filter(|f| !f.span && f.rows == rows).map(|f| f.groups).max().unwrap_or(0) } - /// The var vars of a call: `b` sequences of `per` rows on this rank, - /// `tray` rows in the whole tray batch (the sum of its members' blocks; - /// this rank's `b * per` when it is alone), `longest` the length of - /// the call's longest sequence once its rows are in (only read into a + /// The var vars of a call: `b` sequences of `rows` rows in all on this + /// rank, `tray` rows in the whole tray batch (the sum of its members' + /// blocks; this rank's `rows` when it is alone), `context` the call's + /// sequences' lengths once their rows are in, summed (only read into a /// context var). - pub fn vars(&self, b: u64, per: u64, tray: u64, longest: u64) -> BTreeMap { - let mut vars = BTreeMap::from([(self.rows.var.clone(), b * per), (self.groups.var.clone(), b)]); + pub fn vars(&self, b: u64, rows: u64, tray: u64, context: u64) -> BTreeMap { + let mut vars = BTreeMap::from([(self.rows.var.clone(), rows), (self.groups.var.clone(), b)]); if let Some(t) = &self.tray { vars.insert(t.var.clone(), tray); } if let Some(c) = &self.context { - vars.insert(c.var.clone(), longest); + vars.insert(c.var.clone(), context); } vars } diff --git a/crates/kern-manifest/tests/protocol.rs b/crates/kern-manifest/tests/protocol.rs index d6d68e9c..51c77a89 100644 --- a/crates/kern-manifest/tests/protocol.rs +++ b/crates/kern-manifest/tests/protocol.rs @@ -112,7 +112,7 @@ fn plain_contract() { assert_eq!(p.forward(5, Rows::Const(1)), None); assert_eq!(p.chunk().map(|f| f.name.as_str()), Some("prefill")); assert_eq!((p.row_shapes(), p.max_groups(Rows::Const(1))), (vec![1], 4)); - assert_eq!(p.vars(3, 2, 6, 6), BTreeMap::from([("tokens".into(), 6), ("seqs".into(), 3)])); + assert_eq!(p.vars(3, 6, 6, 6), BTreeMap::from([("tokens".into(), 6), ("seqs".into(), 3)])); } #[test] @@ -200,7 +200,7 @@ fn tray_contract() { assert_eq!(p.line_tables[0].axis, Axis::Tray); assert_eq!(p.any(Fill::Error).map(|f| f.name.as_str()), Some("tp_err")); assert_eq!(p.once, vec!["tp_init".to_string()]); - assert_eq!(p.vars(2, 1, 8, 9), BTreeMap::from([("tokens".into(), 2), ("seqs".into(), 2), ("rows".into(), 8)])); + assert_eq!(p.vars(2, 2, 8, 9), BTreeMap::from([("tokens".into(), 2), ("seqs".into(), 2), ("rows".into(), 8)])); } #[test] @@ -254,7 +254,11 @@ fn batch_rules() { let mut m = plain(); m.programs.get_mut("prefill").unwrap().batch = Some(Batch { groups: 2, rows: Dim::Var("tokens".into()), span: None, context: None }); - rejects(&m, "one sequence, not 2 groups"); + let p = protocol(&m).unwrap(); + let prefill = p.forwards.iter().find(|f| f.name == "prefill").unwrap(); + assert_eq!((prefill.groups, prefill.rows), (2, Rows::Var)); + m.buffers.get_mut("cu_seqlens_q").unwrap().fill = None; + rejects(&m, "needs a `cu_seqlens` fill to split them"); let mut m = plain(); m.programs.get_mut("prefill").unwrap().batch = Some(Batch { groups: 1, rows: Dim::Var("seqs".into()), span: None, context: None }); diff --git a/crates/kern-run/examples/pp_chain.rs b/crates/kern-run/examples/pp_chain.rs index d9f87e43..1d9dd429 100644 --- a/crates/kern-run/examples/pp_chain.rs +++ b/crates/kern-run/examples/pp_chain.rs @@ -102,7 +102,7 @@ fn logits(m: &Verified, program: &str, tokens: &str) -> Result { /// Stage one call of `ids` at `pos` into `rt` from the values `p` and /// `lease` give. fn item(p: &Protocol, lease: &Lease, pos: usize, ids: &[i64]) -> Result<(Vars, BTreeMap>)> { - let (vars, values) = fills(p, std::slice::from_ref(lease), &[pos], ids.len(), ids); + let (vars, values) = fills(p, std::slice::from_ref(lease), &[pos], &[ids.len()], ids); let mut rows: BTreeMap> = values.into_iter().map(|(f, v)| (f.name.clone(), v)).collect(); for t in &p.page_tables { let mut v = Vec::new(); diff --git a/crates/kern-run/src/lib.rs b/crates/kern-run/src/lib.rs index 19cea7c3..de224fc9 100644 --- a/crates/kern-run/src/lib.rs +++ b/crates/kern-run/src/lib.rs @@ -21,6 +21,7 @@ pub mod weights; pub use inputs::{Given, Inputs}; pub use weights::Weights; +use std::borrow::Borrow; use std::collections::BTreeMap; use std::sync::LazyLock; use std::time::Duration; @@ -134,21 +135,25 @@ fn map_file(f: &std::path::Path) -> Result { unsafe { memmap2::Mmap::map(&file) }.with_context(|| format!("mapping weights {}", f.display())) } -/// What one call writes into each fill: `rows` rows for each lease, at -/// that sequence's position, `ids` in row order; and the call's vars. +/// What one call writes into each fill: `lens[i]` rows for lease `i`, at +/// that sequence's position, `ids` in row order (the sequences back to +/// back); and the call's vars. /// /// Every fill a caller of one rank produces is there; `blocks` is the /// tray batch's and only a caller that spans a rank group can fill it. -pub fn fills<'p>( +pub fn fills<'p, L: Borrow>( p: &'p Protocol, - leases: &[Lease], + leases: &[L], positions: &[usize], - rows: usize, + lens: &[usize], ids: &[i64], ) -> (Vars, Vec<(&'p Filled, Vec)>) { let b = leases.len(); - let longest = positions.iter().map(|&pos| pos + rows).max().unwrap_or(rows); - let vars = p.vars(b as u64, rows as u64, (b * rows) as u64, longest as u64); + let rows: usize = lens.iter().sum(); + let context: usize = positions.iter().zip(lens).map(|(pos, n)| pos + n).sum(); + let vars = p.vars(b as u64, rows as u64, rows as u64, context as u64); + let starts: Vec = lens.iter().scan(0, |at, n| Some(std::mem::replace(at, *at + n))).collect(); + let spans = || positions.iter().zip(lens).map(|(&pos, &n)| pos..pos + n); let values = p .fills .iter() @@ -156,15 +161,13 @@ pub fn fills<'p>( let v: Vec = match (f.fill, f.axis) { // Each sequence's first token: the anchor a drafting program // splices its own rows from. - (Fill::Token, Axis::Groups) => ids.chunks(rows).map(|x| x[0]).collect(), + (Fill::Token, Axis::Groups) => starts.iter().map(|&s| ids[s]).collect(), (Fill::Token, _) => ids.to_vec(), - (Fill::Valid, _) => vec![1; b * rows], - (Fill::Position, _) => positions.iter().flat_map(|&pos| (pos..pos + rows).map(|v| v as i64)).collect(), - (Fill::Slot, _) => { - leases.iter().zip(positions).flat_map(|(l, &pos)| l.slots(pos..pos + rows)).collect() - } - (Fill::SeqLen, _) => positions.iter().map(|&pos| (pos + rows) as i64).collect(), - (Fill::CuSeqlens, _) => (0..=b).map(|i| (i * rows) as i64).collect(), + (Fill::Valid, _) => vec![1; rows], + (Fill::Position, _) => spans().flat_map(|r| r.map(|v| v as i64)).collect(), + (Fill::Slot, _) => leases.iter().zip(spans()).flat_map(|(l, r)| l.borrow().slots(r)).collect(), + (Fill::SeqLen, _) => spans().map(|r| r.end as i64).collect(), + (Fill::CuSeqlens, _) => starts.iter().chain([&rows]).map(|&s| s as i64).collect(), (Fill::SpanAt, _) => vec![0], (Fill::Blocks | Fill::Tokens | Fill::Count | Fill::Error | Fill::Clock, _) => return None, }; @@ -202,7 +205,7 @@ fn stage( rows: usize, ids: &[i64], ) -> Result { - let (vars, values) = fills(p, leases, positions, rows, ids); + let (vars, values) = fills(p, leases, positions, &vec![rows; leases.len()], ids); for (f, v) in values { rt.write_input_at(&f.name, &f.encode(&v), &vars)?; } @@ -223,11 +226,11 @@ fn page_rows(t: &PageTable, leases: &[Lease], rows: usize) -> Result> { /// line in entry 0 and zeros through the rest of a wide cell (a program /// that moves along one does so on the device). Columns past the leases /// given repeat the last. -pub fn line_rows(t: &LineTable, leases: &[Lease], cols: usize) -> Result> { +pub fn line_rows>(t: &LineTable, leases: &[L], cols: usize) -> Result> { let mut v = Vec::with_capacity(t.lines * cols * t.width); for line in 0..t.lines { for c in 0..cols { - v.push(leases[c.min(leases.len() - 1)].seq_line(&t.name, line)?); + v.push(leases[c.min(leases.len() - 1)].borrow().seq_line(&t.name, line)?); v.extend(std::iter::repeat_n(0, t.width - 1)); } } diff --git a/crates/kern-serve/src/pp/head.rs b/crates/kern-serve/src/pp/head.rs index 0bcf2707..edefce71 100644 --- a/crates/kern-serve/src/pp/head.rs +++ b/crates/kern-serve/src/pp/head.rs @@ -249,7 +249,7 @@ impl PpHead { } fn send_item(&mut self, f: &Forward, lease: &Lease, pos: usize, ids: &[i64], reply: bool) -> Result<()> { - let (vars, values) = kern_run::fills(&self.p, std::slice::from_ref(lease), &[pos], ids.len(), ids); + let (vars, values) = kern_run::fills(&self.p, std::slice::from_ref(lease), &[pos], &[ids.len()], ids); let mut rows: BTreeMap> = values.into_iter().map(|(f, v)| (f.name.clone(), v)).collect(); for t in &self.p.page_tables { let mut v = Vec::with_capacity(t.width); diff --git a/crates/kern-serve/src/tray.rs b/crates/kern-serve/src/tray.rs index d80fab00..5405bf10 100644 --- a/crates/kern-serve/src/tray.rs +++ b/crates/kern-serve/src/tray.rs @@ -763,8 +763,8 @@ impl Tray { bail!("{} rows in the tray batch, the manifest allows {}", layout.tray * per, t.max); } } - let longest = cells.iter().map(|c| c.pos + c.ids.len()).max().unwrap_or(per); - let mut vars = self.protocol.vars(b as u64, per as u64, (layout.tray * per) as u64, longest as u64); + let context: usize = cells.iter().map(|c| c.pos + c.ids.len()).sum(); + let mut vars = self.protocol.vars(b as u64, (b * per) as u64, (layout.tray * per) as u64, context as u64); if let (Some(s), Some(c)) = (&self.protocol.span, run) { vars.insert(s.var.clone(), c as u64); } diff --git a/tools/kern_manifest.py b/tools/kern_manifest.py index 35ed9e54..52dc6fbe 100644 --- a/tools/kern_manifest.py +++ b/tools/kern_manifest.py @@ -44,7 +44,7 @@ def program(calls, groups=None, rows=None, span=None, context=None, once=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 + the call's sequences' lengths go in, summed), a 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, From 1589fb9559f2b04e53b63c938e4d0d1b6f6a8e9d Mon Sep 17 00:00:00 2001 From: JinYan Su <751080330@qq.com> Date: Wed, 7 Oct 2026 17:27:28 +0000 Subject: [PATCH 14/22] kern-serve: the pp head packs the oldest prompts' next rows into each prefill item A pipeline fed one prompt per item pays a full weight pass on every stage for each prompt's last partial chunk. The head now keeps the admitted prompts active and fills each item with `pack()`: oldest first, partial chunks allowed, up to the call's rows, sequences and context; a pure function with a property test that every small input drains within its bounds, in order, every row once. A full item goes out at once, a partial one only while fewer items than stages are in flight. The last stage answers every item, so the head counts what is in the pipeline and finishes each prompt from its entry in the reply; an item zeroes the slot of each sequence it starts, and a line table is sent only when it changes. Gate: K3 PP8, 192 prompts of 4.5-9.5k tokens, conc16: 47.0k tok/s (5.88k/GPU) vs 42.6k unpacked, 161 items instead of 219, slowest stage 50.1k tok/s while busy. Signed-off-by: JinYan Su <751080330@qq.com> Co-Authored-By: Claude Opus 5.5 (1M context) --- crates/kern-serve/src/pp/head.rs | 222 +++++++++++++++++++++--------- crates/kern-serve/src/pp/mod.rs | 18 ++- crates/kern-serve/src/pp/pack.rs | 103 ++++++++++++++ crates/kern-serve/src/pp/stage.rs | 21 +-- 4 files changed, 283 insertions(+), 81 deletions(-) create mode 100644 crates/kern-serve/src/pp/pack.rs diff --git a/crates/kern-serve/src/pp/head.rs b/crates/kern-serve/src/pp/head.rs index edefce71..02b354e6 100644 --- a/crates/kern-serve/src/pp/head.rs +++ b/crates/kern-serve/src/pp/head.rs @@ -6,12 +6,18 @@ //! through the one-row step, whose tokens the last stage hands back; or, //! when the chunk program hands a token back itself (a model whose chunked //! kernels must see every prompt token), the whole prompt in chunks, the -//! last chunk's token the answer. Each item is sent -//! to every stage before the head stages and issues it itself, so a stage -//! waiting on its device for an item always has it on the way. The first -//! token finishes the request. Nothing waits on the head's device but the -//! clock read, done when nothing is in flight or the clock ring is half -//! full. +//! last chunk's token the answer. A chunk program over several sequences +//! takes the admitted prompts' next rows packed back to back ([`pack`]). +//! +//! Each item is sent to every stage before the head stages and issues it +//! itself, so a stage waiting on its device for an item always has it on +//! the way. The last stage answers every item, so the head knows how many +//! are in the pipeline: a full item (all its rows or all its sequences) +//! goes out at once, a partial one only +//! while fewer items than stages are in flight, so packing waits for rows +//! only when the stages have work. The first token finishes the request. +//! Nothing waits on the head's device but the clock read, done when +//! nothing is in flight or the clock ring is half full. use std::collections::{BTreeMap, VecDeque}; use std::sync::mpsc; @@ -28,6 +34,7 @@ use pegainfer_frontend::engine::{ use tokio::net::TcpListener; use tracing::{error, info}; +use super::pack::{pack, Bounds, Waiting}; use super::{connect_downstream, recv, send, topology, wire, Clock, Handles, Hello, Item, Reader, Table, Writer}; use crate::scheduler::Facts; @@ -43,8 +50,18 @@ pub struct Head { pub stop_tokens: Vec, } -/// A request whose items are out and whose token is not back yet. -struct Inflight { +/// An admitted prompt with rows still to send. +struct Active { + id: RequestId, + lease: Lease, + ids: Vec, + pos: usize, +} + +/// A prompt whose last rows are out: its token is entry `at` of the reply +/// to the item that carried them, and its lease lives until then. +struct Finish { + at: usize, id: RequestId, _lease: Lease, } @@ -56,14 +73,20 @@ pub struct PpHead { /// The one-row step the last prompt token goes through; none when the /// chunk program hands the first token back itself. step: Option, - chunk: usize, + bounds: Bounds, + /// Items in flight below which a partial item goes out: the stages. + depth: usize, max_seqs: usize, stop_tokens: Vec, io: tokio::runtime::Runtime, stages: Vec, replies: mpsc::Receiver>, waiting: VecDeque, - inflight: VecDeque, + active: VecDeque, + /// Per item in flight, oldest first, the prompts its reply finishes. + due: VecDeque>, + /// The line tables as the stages last had them written. + lines: BTreeMap>, clock: Clock, } @@ -125,6 +148,8 @@ pub fn load(m: Verified, weights: &Weights, h: Head) -> Result { Some(c) if (1..=max as u64).contains(&c) => c as usize, Some(c) => bail!("--chunk {c}: the manifest's `tokens` bound is {max}; a chunk can only be smaller"), }; + let context = p.context.as_ref().map_or(usize::MAX, |c| c.max as usize); + let bounds = Bounds { rows: chunk, seqs: prefill.groups as usize, context }; let clock = Clock::new(&p)?; let downstream: Vec> = (1..cut.stages).map(|s| hellos.get(&(s + 1)).map(|(x, ..)| x.mailbox.clone())).collect(); @@ -138,25 +163,29 @@ pub fn load(m: Verified, weights: &Weights, h: Head) -> Result { } let step_name = step.as_ref().map_or("-", |f| f.name.as_str()).to_string(); - info!(pages, page = rt.page(), chunk, prefill = %prefill.name, step = %step_name, "pipeline head ready"); + info!(pages, page = rt.page(), chunk, seqs = bounds.seqs, prefill = %prefill.name, step = %step_name, + "pipeline head ready"); Ok(PpHead { rt, p, prefill, step, - chunk, + bounds, + depth: cut.stages as usize, max_seqs: h.max_seqs, stop_tokens: h.stop_tokens, io, stages, replies, waiting: VecDeque::new(), - inflight: VecDeque::new(), + active: VecDeque::new(), + due: VecDeque::new(), + lines: BTreeMap::new(), clock, }) } -/// Read one stage's connection to its end: the last stage's tokens go to +/// Read one stage's connection to its end: the last stage's answers go to /// the scheduler; anything closing ends the head. async fn listen(stage: u64, mut r: Reader, tokens: Option>>) { let why = loop { @@ -166,7 +195,7 @@ async fn listen(stage: u64, mut r: Reader, tokens: Option> return; } } - (Ok(Some(_)), None) => break "a stage other than the last handed tokens back".to_string(), + (Ok(Some(_)), None) => break "a stage other than the last answered an item".to_string(), (Ok(None), _) => break "stage closed its connection; the pipeline is gone".to_string(), (Err(e), _) => break format!("stage connection: {e:#}"), } @@ -184,8 +213,13 @@ impl PpHead { } } - /// Lease each waiting prompt in order and send its items out; one that - /// cannot be seated stops the scan. + /// Prompts leased and not finished. + fn live(&self) -> usize { + self.active.len() + self.due.iter().map(Vec::len).sum::() + } + + /// Lease each waiting prompt in order; one that cannot be seated stops + /// the scan. fn admit(&mut self, ledger: &mut RequestLedger) -> Result<()> { while let Some(q) = self.waiting.front() { let id = q.id; @@ -194,7 +228,7 @@ impl PpHead { self.waiting.pop_front(); continue; } - if self.inflight.len() >= self.max_seqs { + if self.live() >= self.max_seqs { break; } let (prompt, max_tokens) = (q.request.prompt_tokens.len(), q.request.max_tokens); @@ -206,8 +240,8 @@ impl PpHead { let q = self.waiting.pop_front().unwrap(); ledger.admit(id); ledger.set_cached_tokens(id, 0); - self.send_prompt(&q, &lease)?; - self.inflight.push_back(Inflight { id, _lease: lease }); + let ids = q.request.prompt_tokens.iter().map(|&t| t as i64).collect(); + self.active.push_back(Active { id, lease, ids, pos: 0 }); continue; } Err(Error::Denied(Denied::Busy | Denied::Remapping)) => break, @@ -227,49 +261,106 @@ impl PpHead { Ok(()) } - /// The prompt as items: all but its last token in chunks, the last - /// through the step, whose token comes back; or all of it in chunks - /// when the chunk program hands the token back. - fn send_prompt(&mut self, q: &QueuedRequest, lease: &Lease) -> Result<()> { - let ids: Vec = q.request.prompt_tokens.iter().map(|&t| t as i64).collect(); - let n = ids.len(); - let (prefill, step) = (self.prefill.clone(), self.step.clone()); - let chunked = if step.is_some() { n - 1 } else { n }; - let mut pos = 0; - while pos < chunked { - let c = (chunked - pos).min(self.chunk); - let reply = step.is_none() && pos + c == n; - self.send_item(&prefill, lease, pos, &ids[pos..pos + c], reply)?; - pos += c; + /// The rows of a prompt that go through the chunk program: all of them, + /// or all but the last when the step takes that one. + fn chunked(&self, a: &Active) -> usize { + a.ids.len() - usize::from(self.step.is_some()) + } + + /// Send every item that is ready: packed chunks while they are full or + /// the pipeline has room, then the step of each prompt whose chunks are out. + fn issue(&mut self) -> Result<()> { + loop { + let waiting: Vec = + self.active.iter().map(|a| Waiting { pos: a.pos, left: self.chunked(a) - a.pos }).collect(); + let plan = pack(&waiting, self.bounds); + let rows: usize = plan.iter().map(|&(_, n)| n).sum(); + let full = rows == self.bounds.rows || plan.len() == self.bounds.seqs; + if plan.is_empty() || (!full && self.due.len() >= self.depth) { + break; + } + self.send_chunk(&plan)?; } - match &step { - Some(f) => self.send_item(f, lease, n - 1, &ids[n - 1..], true), - None => Ok(()), + while let Some(i) = self.step.as_ref().and(self.active.iter().position(|a| a.pos == self.chunked(a))) { + let a = self.active.remove(i).expect("found"); + let f = self.step.clone().expect("checked"); + let pos = a.pos; + let finish = Finish { at: 0, id: a.id, _lease: a.lease }; + self.send_item(&f, &[(&finish._lease, pos, &a.ids[pos..])])?; + self.due.push_back(vec![finish]); } + Ok(()) } - fn send_item(&mut self, f: &Forward, lease: &Lease, pos: usize, ids: &[i64], reply: bool) -> Result<()> { - let (vars, values) = kern_run::fills(&self.p, std::slice::from_ref(lease), &[pos], &[ids.len()], ids); + /// One packed chunk item: each planned prompt's next rows; the prompts + /// it finishes leave the active set for the item's reply. + fn send_chunk(&mut self, plan: &[(usize, usize)]) -> Result<()> { + let f = self.prefill.clone(); + let active = std::mem::take(&mut self.active); + let parts: Vec<(&Lease, usize, &[i64])> = plan + .iter() + .map(|&(i, n)| { + let a = &active[i]; + (&a.lease, a.pos, &a.ids[a.pos..a.pos + n]) + }) + .collect(); + let sent = self.send_item(&f, &parts); + drop(parts); + self.active = active; + sent?; + let emits = self.step.is_none(); + let mut finished = Vec::new(); + for (at, &(i, n)) in plan.iter().enumerate().rev() { + self.active[i].pos += n; + if emits && self.active[i].pos == self.active[i].ids.len() { + finished.push((at, i)); + } + } + // Highest index first, so each removal leaves the next one's index in place. + let mut due: Vec = finished + .into_iter() + .map(|(at, i)| { + let a = self.active.remove(i).expect("planned"); + Finish { at, id: a.id, _lease: a.lease } + }) + .collect(); + due.reverse(); + self.due.push_back(due); + Ok(()) + } + + /// Stage one item of `parts` (a lease, its position, its rows) on every + /// stage and on the head. A sequence's first item zeroes its slot on the + /// stages (the head's own lease zeroed it already); a line table goes + /// out only when the item's sequences change it. + fn send_item(&mut self, f: &Forward, parts: &[(&Lease, usize, &[i64])]) -> Result<()> { + let leases: Vec<&Lease> = parts.iter().map(|&(l, ..)| l).collect(); + let positions: Vec = parts.iter().map(|&(_, pos, _)| pos).collect(); + let lens: Vec = parts.iter().map(|(.., ids)| ids.len()).collect(); + let ids: Vec = parts.iter().flat_map(|(.., ids)| ids.iter().copied()).collect(); + let (vars, values) = kern_run::fills(&self.p, &leases, &positions, &lens, &ids); let mut rows: BTreeMap> = values.into_iter().map(|(f, v)| (f.name.clone(), v)).collect(); for t in &self.p.page_tables { - let mut v = Vec::with_capacity(t.width); - lease.extend_row(&t.name, &mut v)?; + let mut v = Vec::with_capacity(t.width * leases.len()); + for l in &leases { + l.extend_row(&t.name, &mut v)?; + } rows.insert(t.name.clone(), v.into_iter().map(i64::from).collect()); } - // A sequence's line tables and slot are the same on every item: its - // first item carries them and zeroes the slot on every stage (the - // head's own lease zeroed its slot already). - let first = pos == 0; - for t in self.p.line_tables.iter().filter(|_| first) { + for t in &self.p.line_tables { let cols = match t.axis { Axis::Tray => self.p.tray.as_ref().map_or(1, |b| b.max), _ => self.p.groups.max, }; - let v = kern_run::line_rows(t, std::slice::from_ref(lease), cols as usize)?; - rows.insert(t.name.clone(), v.into_iter().map(i64::from).collect()); + let v: Vec = kern_run::line_rows(t, &leases, cols as usize)?.into_iter().map(i64::from).collect(); + if self.lines.get(&t.name) != Some(&v) { + self.lines.insert(t.name.clone(), v.clone()); + rows.insert(t.name.clone(), v); + } } - let item = Item { program: f.name.clone(), vars, rows, zero_slot: None, reply }; - let out = Item { zero_slot: lease.seq_slot().filter(|_| first), ..item.clone() }; + let zero_slots = leases.iter().zip(&positions).filter(|(_, &pos)| pos == 0).filter_map(|(l, _)| l.seq_slot()); + let item = Item { program: f.name.clone(), vars, rows, zero_slots: Vec::new() }; + let out = Item { zero_slots: zero_slots.collect(), ..item.clone() }; for w in &mut self.stages { self.io.block_on(send(w, &out))?; } @@ -281,12 +372,13 @@ impl PpHead { Ok(()) } - /// Admit, collect, and read the clock whenever this step issued + /// Admit, send, collect, and read the clock whenever this step issued /// nothing: the read waits for the head's device, which is also how a /// mailbox wait that gave up on a dead stage surfaces here. fn advance(&mut self, ledger: &mut RequestLedger) -> Result<()> { let issued = self.clock.pending(); self.admit(ledger)?; + self.issue()?; self.collect(ledger)?; if self.clock.pending() == issued || self.clock.half_full() { self.clock.flush(&self.rt, 0)?; @@ -294,8 +386,8 @@ impl PpHead { Ok(()) } - /// Every token back since the last step finishes its request, oldest - /// first: the last stage answers in issue order. + /// Every answer back since the last step retires its item and finishes + /// the prompts it carried last: the last stage answers in issue order. fn collect(&mut self, ledger: &mut RequestLedger) -> Result<()> { loop { let t = match self.replies.try_recv() { @@ -303,16 +395,20 @@ impl PpHead { Err(mpsc::TryRecvError::Empty) => return Ok(()), Err(mpsc::TryRecvError::Disconnected) => bail!("the last stage's connection is gone"), }; - let Some(r) = self.inflight.pop_front() else { bail!("a token came back for no request") }; - let Some(&tok) = t.first() else { bail!("the last stage handed back no token") }; - if ledger.is_aborted(r.id) { - ledger.retire(r.id); - continue; + let Some(due) = self.due.pop_front() else { bail!("an answer came back for no item") }; + for r in due { + let Some(&tok) = t.get(r.at) else { + bail!("the last stage handed back {} tokens, no token {}", t.len(), r.at) + }; + if ledger.is_aborted(r.id) { + ledger.retire(r.id); + continue; + } + let tok = tok as u32; + ledger.push_tokens(r.id, &[tok], &[]); + let why = if self.stop_tokens.contains(&tok) { FinishReason::Stop } else { FinishReason::Length }; + ledger.finish(r.id, why); } - let tok = tok as u32; - ledger.push_tokens(r.id, &[tok], &[]); - let why = if self.stop_tokens.contains(&tok) { FinishReason::Stop } else { FinishReason::Length }; - ledger.finish(r.id, why); } } } @@ -340,7 +436,7 @@ impl Scheduler for PpHead { SchedulerMetrics { kv_used_blocks: self.rt.pages_used() as u64, kv_total_blocks: self.rt.pages_total() as u64, - num_running_reqs: self.inflight.len() as u64, + num_running_reqs: (self.active.len() + self.due.iter().map(Vec::len).sum::()) as u64, num_waiting_reqs: self.waiting.len() as u64, spec_decode: None, } diff --git a/crates/kern-serve/src/pp/mod.rs b/crates/kern-serve/src/pp/mod.rs index c43ca2b1..6b157686 100644 --- a/crates/kern-serve/src/pp/mod.rs +++ b/crates/kern-serve/src/pp/mod.rs @@ -17,8 +17,8 @@ //! handles, to map its own outgoing peer from, and the pages. From then on //! the connection carries the log: the head sends every stage every //! [`Item`] in issue order — the program, its vars and the values of each -//! input it staged — and the last stage answers the item that finishes a -//! prompt with its tokens. TCP keeps the order; there is no +//! input it staged — and the last stage answers every item with its tokens +//! (none for a program that hands none back). TCP keeps the order; there is no //! sequence number, no stop message and no fault message. A stage that //! fails exits; a closed connection ends the process on the other side. //! @@ -28,6 +28,7 @@ //! has nothing to enqueue (to read its clock) or when it must read tokens. pub mod head; +mod pack; pub mod stage; use std::collections::{BTreeMap, VecDeque}; @@ -67,15 +68,12 @@ pub struct Table { pub struct Item { pub program: String, pub vars: BTreeMap, - /// Each input's values by name: the fills, the page tables and, on a - /// sequence's first item, its line tables. + /// Each input's values by name: the fills, the page tables and the + /// line tables when the item's sequences change them. pub rows: BTreeMap>, - /// The sequence slot to zero before the call: on its first item, the - /// recurrent state a fresh lease starts from. - pub zero_slot: Option, - /// Whether the last stage hands this call's tokens back: the item - /// that finishes a prompt. - pub reply: bool, + /// The sequence slots to zero before the call, the ones of sequences + /// on their first item: the recurrent state a fresh lease starts from. + pub zero_slots: Vec, } pub type Handles = BTreeMap; diff --git a/crates/kern-serve/src/pp/pack.rs b/crates/kern-serve/src/pp/pack.rs new file mode 100644 index 00000000..fae6634a --- /dev/null +++ b/crates/kern-serve/src/pp/pack.rs @@ -0,0 +1,103 @@ +//! What the next item of a packed prefill carries: the oldest sequences' +//! next rows, back to back, until the call is full. +//! +//! An item reads a stage's weights once whatever its rows, so a pipeline +//! fed one prompt per item pays a whole weight pass for a prompt's last +//! few hundred rows. Packing fills each item instead: sequences in +//! admission order each give their next rows, the last one a partial chunk +//! if that is what fits, until the item has its rows, its sequences or its +//! context (the sum of its sequences' lengths after the call, what a +//! program sized by its context expands). A sequence never skips ahead of +//! an older one, so no prompt starves behind shorter ones arriving later. + +/// One sequence waiting to be fed: `pos` rows already sent, `left` still to send. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct Waiting { + pub pos: usize, + pub left: usize, +} + +/// What one call may hold. +#[derive(Debug, Clone, Copy)] +pub struct Bounds { + pub rows: usize, + pub seqs: usize, + pub context: usize, +} + +/// The next item as `(index into seqs, rows)`, oldest first; empty when +/// nothing is waiting or the oldest cannot advance at all. +pub fn pack(seqs: &[Waiting], b: Bounds) -> Vec<(usize, usize)> { + let mut out = Vec::new(); + let (mut rows, mut context) = (0, 0); + for (i, s) in seqs.iter().enumerate().filter(|(_, s)| s.left > 0) { + let room = b.context.saturating_sub(context + s.pos); + let n = s.left.min(b.rows - rows).min(room); + if n == 0 || out.len() == b.seqs { + break; + } + out.push((i, n)); + rows += n; + context += s.pos + n; + } + out +} + +#[cfg(test)] +mod tests { + use super::*; + + const B: Bounds = Bounds { rows: 8, seqs: 3, context: 20 }; + + fn w(pos: usize, left: usize) -> Waiting { + Waiting { pos, left } + } + + #[test] + fn fills_the_call_oldest_first() { + assert_eq!(pack(&[w(0, 5), w(0, 5), w(0, 5)], B), vec![(0, 5), (1, 3)]); + assert_eq!(pack(&[w(0, 20)], B), vec![(0, 8)]); + assert_eq!(pack(&[w(0, 1), w(0, 1), w(0, 1), w(0, 1)], B), vec![(0, 1), (1, 1), (2, 1)]); + assert_eq!(pack(&[w(4, 0), w(0, 2)], B), vec![(1, 2)]); + assert_eq!(pack(&[], B), vec![]); + } + + #[test] + fn a_long_context_takes_the_room() { + assert_eq!(pack(&[w(16, 8), w(0, 8)], B), vec![(0, 4)]); + assert_eq!(pack(&[w(10, 2), w(0, 8)], B), vec![(0, 2), (1, 6)]); + assert_eq!(pack(&[w(20, 2)], B), vec![]); + } + + /// Every small input against the rules restated plainly: within bounds, + /// in order, no sequence skipped while it could still take a row, and + /// feeding the plan repeatedly sends every row exactly once. + #[test] + fn every_small_input_drains_within_bounds() { + let shapes = (0..4).flat_map(|a| (0..4).flat_map(move |b| (0..10).map(move |c| [a, b, c]))); + for lefts in shapes.flat_map(|l| (0..3).map(move |p| (l, p))) { + let (lefts, p0) = lefts; + let mut seqs: Vec = + lefts.iter().enumerate().map(|(i, &l)| w(if i == 0 { p0 } else { 0 }, l)).collect(); + let total: usize = seqs.iter().map(|s| s.left).sum(); + let mut sent = 0; + for _ in 0..=total { + let plan = pack(&seqs, B); + if plan.is_empty() { + break; + } + let rows: usize = plan.iter().map(|&(_, n)| n).sum(); + let context: usize = plan.iter().map(|&(i, n)| seqs[i].pos + n).sum(); + assert!(rows <= B.rows && plan.len() <= B.seqs && context <= B.context, "{seqs:?} {plan:?}"); + assert!(plan.windows(2).all(|x| x[0].0 < x[1].0), "{plan:?}"); + let first = seqs.iter().position(|s| s.left > 0).unwrap(); + assert_eq!(plan[0].0, first, "{seqs:?} {plan:?}"); + for &(i, n) in &plan { + seqs[i] = w(seqs[i].pos + n, seqs[i].left - n); + sent += n; + } + } + assert_eq!((sent, seqs.iter().all(|s| s.left == 0)), (total, true), "{lefts:?} from {p0}"); + } + } +} diff --git a/crates/kern-serve/src/pp/stage.rs b/crates/kern-serve/src/pp/stage.rs index 990ef529..ef262939 100644 --- a/crates/kern-serve/src/pp/stage.rs +++ b/crates/kern-serve/src/pp/stage.rs @@ -2,8 +2,8 @@ //! head, maps the next stage's mailbox from the table it gets back, then //! replays the head's items in order. It keeps no accounting: the page //! ids, positions and slots of every item are the head's, and so is the -//! zeroing of a fresh sequence's slot. The last stage hands back the -//! tokens of the item that finishes a prompt. +//! zeroing of a fresh sequence's slot. The last stage answers every item +//! with its tokens. use std::sync::mpsc; use std::time::{Duration, Instant}; @@ -108,7 +108,7 @@ pub fn run(m: Verified, weights: &Weights, o: Stage) -> Result<()> { Err(mpsc::TryRecvError::Disconnected) => break, }; let vars: Vars = item.vars.into_iter().filter(|(k, _)| rt.manifest.vars.contains_key(k)).collect(); - if let Some(slot) = item.zero_slot { + for &slot in &item.zero_slots { rt.zero_slot(slot)?; } for (name, v) in &item.rows { @@ -116,11 +116,16 @@ pub fn run(m: Verified, weights: &Weights, o: Stage) -> Result<()> { } rt.issue(&item.program, &vars).with_context(|| format!("`{}`", item.program))?; clock.issued(vars.get(&p.rows.var).copied().unwrap_or(0)); - let emits = p.forwards.iter().find(|f| f.name == item.program).and_then(|f| f.emits); - if let Some(i) = emits.filter(|_| last && item.reply) { - let t = &p.fills[i]; - let mut tokens = t.decode(&rt.read_output(&t.name)?); - tokens.truncate(t.width as usize); + if last { + let emits = p.forwards.iter().find(|f| f.name == item.program).and_then(|f| f.emits); + let tokens = match emits.map(|i| &p.fills[i]) { + Some(t) => { + let mut v = t.decode(&rt.read_output(&t.name)?); + v.truncate(vars.get(&p.groups.var).copied().unwrap_or(1) as usize * t.width as usize); + v + } + None => Vec::new(), + }; io.block_on(send(&mut w, &tokens))?; } if clock.half_full() { From e1e47649dd0878f0844ed6d35e0eb5e2d8e86943 Mon Sep 17 00:00:00 2001 From: JinYan Su <751080330@qq.com> Date: Thu, 8 Oct 2026 14:13:27 +0000 Subject: [PATCH 15/22] runtime: a mark lets another thread wait on the stream while the issuer runs on Runtime::mark records a blocking-sync event on the compute stream; Mark::wait binds the context on whatever thread holds it and returns the fault of any launch before it. A caller learns when the device got somewhere without the issuing thread ever waiting. Co-Authored-By: Claude Opus 5.5 (1M context) Signed-off-by: JinYan Su <751080330@qq.com> --- crates/kern-runtime/src/exec.rs | 45 ++++++++++++++++++++++++++++++++- crates/kern-runtime/src/lib.rs | 1 + 2 files changed, 45 insertions(+), 1 deletion(-) diff --git a/crates/kern-runtime/src/exec.rs b/crates/kern-runtime/src/exec.rs index b0b3184e..37e710fc 100644 --- a/crates/kern-runtime/src/exec.rs +++ b/crates/kern-runtime/src/exec.rs @@ -19,18 +19,61 @@ //! Ranks whose kernels wait on each other (an EP dispatch, a tray //! collective) must all be issued before any is waited for: `enqueue` or //! `issue` each, then [`Runtime::synchronize`] each. +//! +//! A caller that wants to know when the device got somewhere without +//! stopping the thread that issues takes a [`Mark`] and waits on it from +//! another thread. use std::collections::BTreeMap; use std::os::raw::c_void; +use std::sync::Arc; -use cudarc::driver::sys; +use cudarc::driver::{sys, CudaContext}; use crate::compile::{CompiledProgram, Dense, Launch, LaunchKind, RVal, Slot}; use crate::cublas::{gemm_bf16_tn, gemm_bf16_tn_f32, gemm_bf16_tn_pinned, gemm_fp8_tn}; use crate::error::{bail, cuda_check}; use crate::{Error, Result, Runtime}; +/// A point on the compute stream: everything issued before +/// [`Runtime::mark`]. [`Mark::wait`] sleeps until the device has passed it, +/// on any thread, and returns the fault of any launch before it. +pub struct Mark { + ctx: Arc, + event: sys::CUevent, +} + +// The event is a handle into the context the mark carries; waiting binds +// that context, so the mark may be waited on from any thread. +unsafe impl Send for Mark {} + +impl Mark { + pub fn wait(self) -> Result<()> { + self.ctx.bind_to_thread()?; + cuda_check(unsafe { sys::cuEventSynchronize(self.event) }, "cuEventSynchronize") + } +} + +impl Drop for Mark { + fn drop(&mut self) { + unsafe { sys::cuEventDestroy_v2(self.event) }; + } +} + impl Runtime { + /// Mark the compute stream where it is now. The event blocks rather + /// than spins, so a thread waiting on it costs no core. + pub fn mark(&self) -> Result { + self.ctx.bind_to_thread()?; + let mut event: sys::CUevent = std::ptr::null_mut(); + let flags = + sys::CUevent_flags::CU_EVENT_DISABLE_TIMING as u32 | sys::CUevent_flags::CU_EVENT_BLOCKING_SYNC as u32; + cuda_check(unsafe { sys::cuEventCreate(&mut event, flags) }, "cuEventCreate")?; + let mark = Mark { ctx: self.ctx.clone(), event }; + cuda_check(unsafe { sys::cuEventRecord(mark.event, self.stream.cu_stream()) }, "cuEventRecord")?; + Ok(mark) + } + fn require_ready(&self) -> Result<()> { if self.awaits_host() { bail!(Api, "the manifest's host states are not bound; bind_host before executing programs"); diff --git a/crates/kern-runtime/src/lib.rs b/crates/kern-runtime/src/lib.rs index a7ad4bdb..0ac87b64 100644 --- a/crates/kern-runtime/src/lib.rs +++ b/crates/kern-runtime/src/lib.rs @@ -58,6 +58,7 @@ use device::{alloc, DeviceBuf, Pinned}; pub use device::{device_uuid, Mapped, PeerHandle}; use error::{bail, cuda_check}; pub use error::{Error, Result}; +pub use exec::Mark; pub use harness::Scratch; pub use host::HostRegion; pub use host_weights::HostWeights; From eaecc5299a1bc40dcadc6dc02348a11c1b1be948 Mon Sep 17 00:00:00 2001 From: JinYan Su <751080330@qq.com> Date: Thu, 8 Oct 2026 14:13:27 +0000 Subject: [PATCH 16/22] kern-serve: a pipeline stage logs each item's issue and done from the host Every stage marks each item it issues and a thread beside it waits on the mark: one `pp item` line per item with its number (its place in the TCP order, the same on every stage), rows, issued_us and done_us. Nothing on the issuing thread waits any more; the device clock's half-ring read was a stream sync every 128 items. A device fault (a mailbox wait that gave up) surfaces at the item's mark and ends the process, which is how a dead stage is noticed. tools/pp_timeline.py joins the stages' lines by item into per-stage period, held and wait, and a gantt. Gate (qwen3-4b, 2 stages, tray07): gate prompts' first tokens equal to unsplit; head and stage logged 498 items each; last stage killed -> head exits; last stage frozen -> head exits at 60 s on its item's mark, the resumed stage 62 s later the same way. Co-Authored-By: Claude Opus 5.5 (1M context) Signed-off-by: JinYan Su <751080330@qq.com> --- crates/kern-serve/src/pp/head.rs | 22 ++---- crates/kern-serve/src/pp/mod.rs | 116 ++++++++++++++---------------- crates/kern-serve/src/pp/stage.rs | 23 ++---- tools/pp_timeline.py | 99 +++++++++++++++++++++++++ 4 files changed, 164 insertions(+), 96 deletions(-) create mode 100755 tools/pp_timeline.py diff --git a/crates/kern-serve/src/pp/head.rs b/crates/kern-serve/src/pp/head.rs index 02b354e6..93aa9361 100644 --- a/crates/kern-serve/src/pp/head.rs +++ b/crates/kern-serve/src/pp/head.rs @@ -16,8 +16,7 @@ //! goes out at once, a partial one only //! while fewer items than stages are in flight, so packing waits for rows //! only when the stages have work. The first token finishes the request. -//! Nothing waits on the head's device but the clock read, done when -//! nothing is in flight or the clock ring is half full. +//! Nothing waits on the head's device; its [`Watch`] logs each item. use std::collections::{BTreeMap, VecDeque}; use std::sync::mpsc; @@ -35,7 +34,7 @@ use tokio::net::TcpListener; use tracing::{error, info}; use super::pack::{pack, Bounds, Waiting}; -use super::{connect_downstream, recv, send, topology, wire, Clock, Handles, Hello, Item, Reader, Table, Writer}; +use super::{connect_downstream, recv, send, topology, wire, Handles, Hello, Item, Reader, Table, Watch, Writer}; use crate::scheduler::Facts; /// How the head loads and what it may run at once. @@ -87,7 +86,7 @@ pub struct PpHead { due: VecDeque>, /// The line tables as the stages last had them written. lines: BTreeMap>, - clock: Clock, + watch: Watch, } /// Wait for every other stage, load stage 0 at the pages they all have, @@ -150,7 +149,6 @@ pub fn load(m: Verified, weights: &Weights, h: Head) -> Result { }; let context = p.context.as_ref().map_or(usize::MAX, |c| c.max as usize); let bounds = Bounds { rows: chunk, seqs: prefill.groups as usize, context }; - let clock = Clock::new(&p)?; let downstream: Vec> = (1..cut.stages).map(|s| hellos.get(&(s + 1)).map(|(x, ..)| x.mailbox.clone())).collect(); let (tx, replies) = mpsc::channel(); @@ -181,7 +179,7 @@ pub fn load(m: Verified, weights: &Weights, h: Head) -> Result { active: VecDeque::new(), due: VecDeque::new(), lines: BTreeMap::new(), - clock, + watch: Watch::new(0), }) } @@ -368,22 +366,14 @@ impl PpHead { kern_run::write_named(&mut self.rt, &self.p, name, v, &item.vars)?; } self.rt.issue(&item.program, &item.vars).with_context(|| format!("`{}`", item.program))?; - self.clock.issued(ids.len() as u64); + self.watch.issued(&self.rt, ids.len() as u64)?; Ok(()) } - /// Admit, send, collect, and read the clock whenever this step issued - /// nothing: the read waits for the head's device, which is also how a - /// mailbox wait that gave up on a dead stage surfaces here. fn advance(&mut self, ledger: &mut RequestLedger) -> Result<()> { - let issued = self.clock.pending(); self.admit(ledger)?; self.issue()?; - self.collect(ledger)?; - if self.clock.pending() == issued || self.clock.half_full() { - self.clock.flush(&self.rt, 0)?; - } - Ok(()) + self.collect(ledger) } /// Every answer back since the last step retires its item and finishes diff --git a/crates/kern-serve/src/pp/mod.rs b/crates/kern-serve/src/pp/mod.rs index 6b157686..ac58ca94 100644 --- a/crates/kern-serve/src/pp/mod.rs +++ b/crates/kern-serve/src/pp/mod.rs @@ -18,28 +18,31 @@ //! the connection carries the log: the head sends every stage every //! [`Item`] in issue order — the program, its vars and the values of each //! input it staged — and the last stage answers every item with its tokens -//! (none for a program that hands none back). TCP keeps the order; there is no -//! sequence number, no stop message and no fault message. A stage that -//! fails exits; a closed connection ends the process on the other side. +//! (none for a program that hands none back). TCP keeps the order, so an +//! item's number is its place in it and every stage counts the same items +//! alike; there is no sequence number, no stop message and no fault +//! message. A stage that fails exits; a closed connection ends the process +//! on the other side. //! //! The activations never touch the connection: they go GPU to GPU through //! the mailbox the cut wired in, and each stage's device waits for its -//! predecessor on its own. A stage only enqueues; it synchronizes when it -//! has nothing to enqueue (to read its clock) or when it must read tokens. +//! predecessor on its own. A stage only enqueues; the last one waits for +//! its tokens, nothing else waits on a device. A [`Watch`] beside each +//! stage logs when every item left the device, and is where a device +//! fault (a mailbox wait that gave up on a dead stage) ends the process. pub mod head; mod pack; pub mod stage; -use std::collections::{BTreeMap, VecDeque}; +use std::collections::BTreeMap; +use std::sync::mpsc; +use std::time::{SystemTime, UNIX_EPOCH}; use anyhow::{bail, Context, Result}; use bytes::Bytes; use futures_util::{SinkExt, StreamExt}; -use kern_manifest::protocol::Axis; -use kern_manifest::types::Fill; -use kern_manifest::Protocol; -use kern_runtime::{GroupRank, PeerHandle, Runtime, Topology}; +use kern_runtime::{GroupRank, Mark, PeerHandle, Runtime, Topology}; use serde::{Deserialize, Serialize}; use tokio::net::tcp::{OwnedReadHalf, OwnedWriteHalf}; use tokio::net::TcpStream; @@ -152,64 +155,55 @@ fn io() -> Result { Ok(tokio::runtime::Builder::new_multi_thread().worker_threads(1).enable_all().build()?) } -/// What a stage's device spent per item, read off its clock output: one -/// log line per item, busy and waiting on either side. -struct Clock { - name: String, - ring: u64, - logged: u64, - /// Rows of each item issued and not yet logged, oldest first. - rows: VecDeque, +/// When each item a stage issued left its device: one `pp item` line per +/// item (its number, its rows, `issued_us` and `done_us` of the wall +/// clock), logged by a +/// thread that waits on the item's [`Mark`] so the issuing thread never +/// does. A stage's period is done to done; against what the stage alone +/// takes for the item's shape, the rest is its pipeline bubble. A fault on +/// the device surfaces at its item's mark and ends the process. +struct Watch { + marks: Option>, + thread: Option>, + next: u64, } -impl Clock { - fn new(p: &Protocol) -> Result { - let f = p.any(Fill::Clock).context("a cut stage declares a `clock` fill")?; - let Axis::Fixed(n) = f.axis else { bail!("`{}`: a clock is a fixed ring", f.name) }; - Ok(Clock { name: f.name.clone(), ring: n / 5, logged: 0, rows: VecDeque::new() }) - } - - fn issued(&mut self, rows: u64) { - self.rows.push_back(rows); +impl Watch { + fn new(stage: u64) -> Watch { + let (marks, rx) = mpsc::channel::<(u64, u64, u64, Mark)>(); + let thread = std::thread::spawn(move || { + for (item, rows, issued_us, mark) in rx { + if let Err(e) = mark.wait() { + tracing::error!(stage, item, "device: {e:#}"); + std::process::exit(1); + } + tracing::info!(stage, item, rows, issued_us, done_us = now_us(), "pp item"); + } + }); + Watch { marks: Some(marks), thread: Some(thread), next: 0 } } - /// Items issued and not yet logged. - fn pending(&self) -> usize { - self.rows.len() + /// Mark the item just issued, of `rows` rows. + fn issued(&mut self, rt: &Runtime, rows: u64) -> Result<()> { + let marks = self.marks.as_ref().expect("taken only on drop"); + if marks.send((self.next, rows, now_us(), rt.mark()?)).is_err() { + bail!("the watch thread is gone"); + } + self.next += 1; + Ok(()) } +} - /// Whether the oldest unread entry is half a ring from being overwritten. - fn half_full(&self) -> bool { - self.rows.len() as u64 >= self.ring / 2 - } +fn now_us() -> u64 { + SystemTime::now().duration_since(UNIX_EPOCH).map_or(0, |d| d.as_micros() as u64) +} - /// Read the clock (waiting for the stage's stream) and log every item - /// issued since the last read. - fn flush(&mut self, rt: &Runtime, stage: u64) -> Result<()> { - if self.rows.is_empty() { - return Ok(()); +/// The items still in flight are logged before the stage goes. +impl Drop for Watch { + fn drop(&mut self) { + drop(self.marks.take()); + if let Some(t) = self.thread.take() { + let _ = t.join(); } - let raw = rt.read_output(&self.name)?; - let e: Vec = raw.as_chunks::<8>().0.iter().map(|c| i64::from_le_bytes(*c)).collect(); - let span = |a: i64, b: i64| (a != 0 && b != 0).then_some(b - a); - for rows in self.rows.drain(..) { - let k = self.logged; - self.logged += 1; - let x = &e[(k % self.ring) as usize * 5..][..5]; - if x[0] != k as i64 + 1 { - tracing::warn!(stage, item = k, "pp clock entry overwritten before it was read"); - continue; - } - tracing::info!( - stage, - item = k, - rows, - busy_ns = span(x[2], x[3]), - recv_wait_ns = span(x[1], x[2]), - send_wait_ns = span(x[3], x[4]), - "pp item" - ); - } - Ok(()) } } diff --git a/crates/kern-serve/src/pp/stage.rs b/crates/kern-serve/src/pp/stage.rs index ef262939..43f6c52c 100644 --- a/crates/kern-serve/src/pp/stage.rs +++ b/crates/kern-serve/src/pp/stage.rs @@ -15,7 +15,7 @@ use kern_runtime::{Capacity, Runtime}; use tokio::net::TcpStream; use tracing::{error, info}; -use super::{connect_downstream, handles, recv, send, topology, wire, Clock, Hello, Item, Table}; +use super::{connect_downstream, handles, recv, send, topology, wire, Hello, Item, Table, Watch}; /// How a stage loads and where its head listens. pub struct Stage { @@ -94,19 +94,8 @@ pub fn run(m: Verified, weights: &Weights, o: Stage) -> Result<()> { } }); - let mut clock = Clock::new(&p)?; - loop { - let item = match items.try_recv() { - Ok(item) => item, - Err(mpsc::TryRecvError::Empty) => { - clock.flush(&rt, cut.stage)?; - match items.recv() { - Ok(item) => item, - Err(_) => break, - } - } - Err(mpsc::TryRecvError::Disconnected) => break, - }; + let mut watch = Watch::new(cut.stage); + for item in items { let vars: Vars = item.vars.into_iter().filter(|(k, _)| rt.manifest.vars.contains_key(k)).collect(); for &slot in &item.zero_slots { rt.zero_slot(slot)?; @@ -115,7 +104,7 @@ pub fn run(m: Verified, weights: &Weights, o: Stage) -> Result<()> { kern_run::write_named(&mut rt, &p, name, v, &vars)?; } rt.issue(&item.program, &vars).with_context(|| format!("`{}`", item.program))?; - clock.issued(vars.get(&p.rows.var).copied().unwrap_or(0)); + watch.issued(&rt, vars.get(&p.rows.var).copied().unwrap_or(0))?; if last { let emits = p.forwards.iter().find(|f| f.name == item.program).and_then(|f| f.emits); let tokens = match emits.map(|i| &p.fills[i]) { @@ -128,11 +117,7 @@ pub fn run(m: Verified, weights: &Weights, o: Stage) -> Result<()> { }; io.block_on(send(&mut w, &tokens))?; } - if clock.half_full() { - clock.flush(&rt, cut.stage)?; - } } - clock.flush(&rt, cut.stage)?; Ok(()) } diff --git a/tools/pp_timeline.py b/tools/pp_timeline.py new file mode 100755 index 00000000..a5e5be23 --- /dev/null +++ b/tools/pp_timeline.py @@ -0,0 +1,99 @@ +#!/usr/bin/env python3 +"""A pipeline's timeline from its stages' `pp item` log lines (kern-serve --pp-*). + +Every stage logs when it issued each item and when the item left its +device (`issued_us`, `done_us`); the item number is the same on every +stage. A stage starts item k once it has issued it, is done with k - 1 and +its upstream is done with k, so per stage and item: + + wait = start - done[k-1] (idle: nothing issued yet, or the upstream had not delivered) + held = done - start (its own work, plus any wait for the downstream to take) + +and done to done is the stage's period. Across trays `done_us` is two wall +clocks; `--skew S=us` shifts stage S and every later one. + + pp_timeline.py stage*.log [--skew 4=-1200] [--from K] [--to K] [--gantt] +""" +import argparse +import re +import statistics +from collections import defaultdict + +LINE = re.compile(r"pp item stage=(\d+) item=(\d+) rows=(\d+) issued_us=(\d+) done_us=(\d+)") + + +def parse(paths): + issued, done, rows = defaultdict(dict), defaultdict(dict), {} + for p in paths: + for m in map(LINE.search, open(p, errors="replace")): + if m: + s, k, r, i, t = map(int, m.groups()) + issued[s][k], done[s][k], rows[k] = i, t, r + return issued, done, rows + + +def timeline(issued, done, skew): + stages = sorted(done) + shift = {s: sum(v for at, v in skew.items() if at <= s) for s in stages} + d = {s: {k: t + shift[s] for k, t in done[s].items()} for s in stages} + i = {s: {k: t + shift[s] for k, t in issued[s].items()} for s in stages} + items = sorted(set.intersection(*(set(v) for v in d.values()))) + start = {s: {} for s in stages} + for s in stages: + for k in items: + ready = [i[s][k], d[s].get(k - 1, 0)] + ([d[s - 1][k]] if s - 1 in d else []) + start[s][k] = max(ready) + return stages, items, d, start + + +def pct(xs, p): + xs = sorted(xs) + return xs[min(len(xs) - 1, int(p * len(xs)))] if xs else float("nan") + + +def report(stages, items, d, start, rows, lo, hi): + ks = [k for k in items if lo <= k <= hi and k - 1 in d[stages[0]]] + print(f"items {ks[0]}..{ks[-1]} ({len(ks)}), rows/item mean {statistics.mean(rows[k] for k in ks):.0f}") + print(f"{'stage':>5} {'period p50':>11} {'p90':>8} {'held p50':>9} {'wait p50':>9} {'wait share':>11} (ms)") + for s in stages: + period = [(d[s][k] - d[s][k - 1]) / 1e3 for k in ks] + held = [(d[s][k] - start[s][k]) / 1e3 for k in ks] + wait = [(start[s][k] - d[s][k - 1]) / 1e3 for k in ks] + share = sum(wait) / sum(period) if sum(period) else 0 + print(f"{s:>5} {pct(period, .5):>11.2f} {pct(period, .9):>8.2f} {pct(held, .5):>9.2f} " + f"{pct(wait, .5):>9.2f} {share:>10.1%}") + + +def gantt(stages, items, d, start, lo, hi, width=120): + ks = [k for k in items if lo <= k <= hi] + t0 = min(start[s][ks[0]] for s in stages) + t1 = max(d[s][ks[-1]] for s in stages) + scale = (t1 - t0) / width or 1 + for s in stages: + row = ["."] * width + for k in ks: + a, b = int((start[s][k] - t0) / scale), int((d[s][k] - t0) / scale) + for x in range(max(a, 0), min(max(b, a + 1), width)): + row[x] = str(k % 10) + print(f"S{s} |{''.join(row)}|") + print(f" {(t1 - t0) / 1e3:.1f} ms across, '.' idle, digit = item number mod 10") + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("logs", nargs="+") + ap.add_argument("--skew", action="append", default=[], help="S=us: shift stage S and later by us") + ap.add_argument("--from", dest="lo", type=int, default=1) + ap.add_argument("--to", dest="hi", type=int, default=1 << 62) + ap.add_argument("--gantt", action="store_true") + a = ap.parse_args() + issued, done, rows = parse(a.logs) + skew = {int(s): int(v) for s, v in (x.split("=") for x in a.skew)} + stages, items, d, start = timeline(issued, done, skew) + report(stages, items, d, start, rows, a.lo, a.hi) + if a.gantt: + gantt(stages, items, d, start, a.lo, min(a.hi, a.lo + 19)) + + +if __name__ == "__main__": + main() From 945ddbeae142bf0a0fc3f4a5bb4b061cb0449afa Mon Sep 17 00:00:00 2001 From: JinYan Su <751080330@qq.com> Date: Thu, 8 Oct 2026 14:13:27 +0000 Subject: [PATCH 17/22] manifest: a cut adds transport only, the device clock is gone The `clock` fill made a stage's output carry globaltimer stamps, so a forward was no longer a pure function of its inputs, and its [ring x 5] layout was spread over the mailbox kernel, the cut, the protocol and kern-serve. The host now knows when each item left a device (previous commit), so the fill, `pp.clock`, the stamp kernels and the waits' clock argument are deleted; a first stage has no receive calls, a last stage no send calls. Gate: kern-manifest tests (cut property test included) pass; pp_chain on qwen3-4b cut at 218, 7 prompts, chunk 512 x 8 steps and chunk 64 x 16 steps: logits and tokens bitwise identical to the unsplit manifest. Co-Authored-By: Claude Opus 5.5 (1M context) Signed-off-by: JinYan Su <751080330@qq.com> --- crates/kern-manifest/src/cut.rs | 38 ++++++-------------- crates/kern-manifest/src/protocol.rs | 14 ++------ crates/kern-manifest/src/types.rs | 3 -- crates/kern-manifest/src/verify.rs | 2 +- crates/kern-run/src/lib.rs | 2 +- crates/kern-serve/src/tray.rs | 2 +- docs/manifest.md | 7 ++-- docs/roadmap.md | 2 +- docs/serve.md | 10 ++++-- schema/manifest-v5.schema.json | 5 --- tools/dsv41/attention/ar_trace/src/main.rs | 4 +-- tools/kernels-src/pp_mailbox.cu | 42 +++------------------- 12 files changed, 35 insertions(+), 96 deletions(-) diff --git a/crates/kern-manifest/src/cut.rs b/crates/kern-manifest/src/cut.rs index f0c892ab..8a754457 100644 --- a/crates/kern-manifest/src/cut.rs +++ b/crates/kern-manifest/src/cut.rs @@ -52,12 +52,8 @@ pub const TIMEOUT_NS: i64 = 60_000_000_000; const HEADER: u64 = 256; /// Each carried buffer starts at a multiple of this in the box. const ALIGN: u64 = 256; -/// Entries of a stage's clock. -const RING: u64 = 256; /// The module name the pipeline kernels go under in every stage. const MODULE: &str = "pp"; -/// The stage's clock (fill `clock`). -const CLOCK: &str = "pp.clock"; /// Every reason a manifest cannot be cut where asked, reported together. #[derive(Debug)] @@ -226,14 +222,12 @@ fn call(op: &str, args: serde_json::Value) -> Call { serde_json::from_value(json!({"label": op, "op": op, "args": args})).expect("a well-formed call") } -/// The calls that start stage `s`: take what the edge into it carries, or -/// stamp the start of a first stage. +/// The calls that start stage `s`: take what the edge into it carries +/// (none on a first stage). fn receive(s: usize, into: Option<&[Carried]>) -> Vec { - let Some(list) = into else { - return vec![call("pp.stamp_first", json!([{"buf": peer_name(s)}, {"buf": CLOCK}]))]; - }; + let Some(list) = into else { return Vec::new() }; let b = box_name(s - 1); - std::iter::once(call("pp.wait_full", json!([{"buf": b}, {"buf": CLOCK}]))) + std::iter::once(call("pp.wait_full", json!([{"buf": b}]))) .chain(list.iter().map(|c| { call( &format!("pp.take.{}", c.dtype), @@ -244,14 +238,12 @@ fn receive(s: usize, into: Option<&[Carried]>) -> Vec { .collect() } -/// The calls that end stage `s`: put what the edge out of it carries, or -/// stamp the end of a last stage. +/// The calls that end stage `s`: put what the edge out of it carries +/// (none on a last stage). fn send(s: usize, out: Option<&[Carried]>) -> Vec { - let Some(list) = out else { - return vec![call("pp.stamp_last", json!([{"buf": box_name(s - 1)}, {"buf": CLOCK}]))]; - }; + let Some(list) = out else { return Vec::new() }; let p = peer_name(s); - std::iter::once(call("pp.wait_empty", json!([{"buf": p}, {"buf": CLOCK}]))) + std::iter::once(call("pp.wait_empty", json!([{"buf": p}]))) .chain(list.iter().map(|c| { call( &format!("pp.put.{}", c.dtype), @@ -263,17 +255,16 @@ fn send(s: usize, out: Option<&[Carried]>) -> Vec { } /// The pipeline ops a stage's calls name: `pp.put.` and -/// `pp.take.` per carried dtype, the waits, posts and stamps. +/// `pp.take.` per carried dtype, the waits and the posts. fn pp_op(name: &str) -> Op { let single = |entry: &str, params: serde_json::Value| { json!({"params": params, "impl": {"launches": [ {"module": MODULE, "entry": entry, "block": [1, 1, 1], "grid": [1, 1, 1]}]}}) }; let wait = |entry: &str, first: &str| { - json!({"params": [first, "out buffer"], "impl": {"launches": [ + json!({"params": [first], "impl": {"launches": [ {"module": MODULE, "entry": entry, "block": [1, 1, 1], "grid": [1, 1, 1], - "params": [first, "out buffer", "i64"], - "args": [{"param": 0}, {"param": 1}, {"i64": TIMEOUT_NS}]}]}}) + "params": [first, "i64"], "args": [{"param": 0}, {"i64": TIMEOUT_NS}]}]}}) }; let copy = |entry: &str, params: serde_json::Value| { json!({"params": params, "impl": {"launches": [ @@ -284,8 +275,6 @@ fn pp_op(name: &str) -> Op { ["pp", "wait_full"] => wait("kern_pp_wait_full", "in buffer"), ["pp", "post_full"] => single("kern_pp_post_full", json!(["in buffer"])), ["pp", "post_empty"] => single("kern_pp_post_empty", json!(["inout buffer"])), - ["pp", "stamp_first"] => single("kern_pp_stamp_first", json!(["in buffer", "out buffer"])), - ["pp", "stamp_last"] => single("kern_pp_stamp_last", json!(["in buffer", "out buffer"])), ["pp", "put", dt] => copy("kern_pp_put", json!(["in buffer", format!("in buffer<{dt}>"), "i64", "i64"])), ["pp", "take", dt] => copy("kern_pp_take", json!(["in buffer", format!("out buffer<{dt}>"), "i64", "i64"])), _ => unreachable!("the cut names only the pipeline ops above"), @@ -325,11 +314,6 @@ fn stage( let u8_box = |bytes: u64, kind: &str| -> Buffer { serde_json::from_value(json!({"dtype": "u8", "shape": [bytes], "kind": kind, "export": true})).expect("a box") }; - buffers.insert( - CLOCK.to_string(), - serde_json::from_value(json!({"dtype": "i64", "shape": [RING * 5], "kind": "output", "fill": "clock"})) - .expect("a clock"), - ); if s > 0 { buffers.insert(box_name(s - 1), u8_box(boxes[s - 1], "carry")); } diff --git a/crates/kern-manifest/src/protocol.rs b/crates/kern-manifest/src/protocol.rs index f4e279e5..1caf3f53 100644 --- a/crates/kern-manifest/src/protocol.rs +++ b/crates/kern-manifest/src/protocol.rs @@ -348,7 +348,6 @@ impl Protocol { Fill::Tokens => axis == Axis::Groups || (axis == Axis::Tray && width == 1), Fill::SpanAt | Fill::Error => axis == Axis::Fixed(1) && b.dtype == DType::I32, Fill::Blocks => matches!(axis, Axis::Fixed(n) if n >= 2) && b.dtype == DType::I32, - Fill::Clock => matches!(axis, Axis::Fixed(n) if n > 0 && n % 5 == 0) && b.dtype == DType::I64, }; if !ok { errs.push(format!( @@ -362,7 +361,6 @@ impl Protocol { Fill::Tokens => "expected [groups], [groups, w] or [tray]", Fill::SpanAt | Fill::Error => "expected i32 [1]", Fill::Blocks => "expected i32 [members + 1]", - Fill::Clock => "expected i64 [ring × 5]", } )); continue; @@ -383,16 +381,8 @@ impl Protocol { if fills.iter().filter(|f| f.fill == Fill::Token && f.axis == Axis::Groups).count() > 1 { errs.push("fill `token` over the sequences is on more than one buffer".into()); } - for fill in [ - Fill::Position, - Fill::Valid, - Fill::CuSeqlens, - Fill::SpanAt, - Fill::Blocks, - Fill::Count, - Fill::Error, - Fill::Clock, - ] { + for fill in [Fill::Position, Fill::Valid, Fill::CuSeqlens, Fill::SpanAt, Fill::Blocks, Fill::Count, Fill::Error] + { if one(fill).is_none() && m.buffers.values().any(|b| b.fill == Some(fill)) { errs.push(format!("fill `{fill}` is on more than one buffer")); } diff --git a/crates/kern-manifest/src/types.rs b/crates/kern-manifest/src/types.rs index a8a27e2f..957904c8 100644 --- a/crates/kern-manifest/src/types.rs +++ b/crates/kern-manifest/src/types.rs @@ -203,8 +203,6 @@ pub enum Fill { Count, /// Output: a one-word error flag of the collectives, read after every call; nonzero is a failed step. Error, - /// Output: a pipeline stage's clock, `i64 [ring × 5]`: entry `item % ring` is {item + 1, receive wait begin, end, send wait begin, end} in globaltimer ns, 0 where the stage has no such edge. A cut declares it; the stage reads it when it pleases. - Clock, } impl fmt::Display for Fill { @@ -221,7 +219,6 @@ impl fmt::Display for Fill { Fill::Tokens => "tokens", Fill::Count => "count", Fill::Error => "error", - Fill::Clock => "clock", }) } } diff --git a/crates/kern-manifest/src/verify.rs b/crates/kern-manifest/src/verify.rs index 0afc4860..d00ac0df 100644 --- a/crates/kern-manifest/src/verify.rs +++ b/crates/kern-manifest/src/verify.rs @@ -252,7 +252,7 @@ fn diagnostics(m: &Manifest) -> Vec { let ctx = format!("buffer `{name}`"); // 9b. fill if let Some(fill) = b.fill { - let output = matches!(fill, Fill::Tokens | Fill::Count | Fill::Error | Fill::Clock); + let output = matches!(fill, Fill::Tokens | Fill::Count | Fill::Error); match (b.kind, output) { (BufferKind::Input, false) | (BufferKind::Output, true) => {} (BufferKind::Input | BufferKind::Output, _) => errs.push(format!( diff --git a/crates/kern-run/src/lib.rs b/crates/kern-run/src/lib.rs index de224fc9..c0f33415 100644 --- a/crates/kern-run/src/lib.rs +++ b/crates/kern-run/src/lib.rs @@ -169,7 +169,7 @@ pub fn fills<'p, L: Borrow>( (Fill::SeqLen, _) => spans().map(|r| r.end as i64).collect(), (Fill::CuSeqlens, _) => starts.iter().chain([&rows]).map(|&s| s as i64).collect(), (Fill::SpanAt, _) => vec![0], - (Fill::Blocks | Fill::Tokens | Fill::Count | Fill::Error | Fill::Clock, _) => return None, + (Fill::Blocks | Fill::Tokens | Fill::Count | Fill::Error, _) => return None, }; Some((f, v)) }) diff --git a/crates/kern-serve/src/tray.rs b/crates/kern-serve/src/tray.rs index 5405bf10..3d144408 100644 --- a/crates/kern-serve/src/tray.rs +++ b/crates/kern-serve/src/tray.rs @@ -822,7 +822,7 @@ impl Tray { (Fill::CuSeqlens, _) => (0..=b as i64).map(|i| i * per as i64).collect(), (Fill::SpanAt, _) => vec![span_at], (Fill::Blocks, _) => l.offsets(&groups, q).into_iter().map(|o| (o * per) as i64).collect(), - (Fill::Tokens | Fill::Count | Fill::Error | Fill::Clock, _) => continue, + (Fill::Tokens | Fill::Count | Fill::Error, _) => continue, (Fill::Token, Axis::Fixed(_)) => unreachable!("the protocol checks fill shapes"), }; writes.push((f, v)); diff --git a/docs/manifest.md b/docs/manifest.md index 6edab8b1..44657b98 100644 --- a/docs/manifest.md +++ b/docs/manifest.md @@ -428,10 +428,9 @@ kern-serve 只拿 `Protocol` 驱动 runtime，不读 JSON（CI 用 grep 保证 pp.post_empty`；邮箱 `pp..box`（`u8[256 + 最大载荷]`）在下游 GPU 上， `full` / `empty` 两个计数在头里，上游经 `pp..peer` 远程写。等待在设备 上自旋，超时（60 s）`__trap`——进程死于 launch failure，不会挂住。 -- 每段声明 `pp.clock`（`fill: clock`，`i64 [256 × 5]` output）：第 k 项 = - {k+1, 收等开始, 收等结束, 发等开始, 发等结束}（globaltimer ns），首段的 - "收等结束"是本段开始、末段的"发等开始"是本段结束，于是每段 busy = 发等 - 开始 − 收等结束。item 下标由邮箱计数在设备上导出，不需要输入。 +- 切分只加传输，不加观测：首段没有 recv、末段没有 send，各段的输出都是 + forward 的纯函数。一个 item 什么时候离开设备是 serving 外壳的事（host 上 + 等 event，见 serve.md「流水线」）。 - 非末段声明拓扑组 `pp.`（大小 2，本段是 0 号）。**这里放宽了 SPMD**： 组两端装的是不同的 manifest，只在 `of` buffer（同名同字节数的 `pp..box`，上游那份是不用的 workspace + export）上对上—— diff --git a/docs/roadmap.md b/docs/roadmap.md index c44e78ff..f5b008bd 100644 --- a/docs/roadmap.md +++ b/docs/roadmap.md @@ -62,7 +62,7 @@ kern 在内部 RSI 里是 kernel agent 的反馈环境（`kern test` 裁判、`k | 级 | 内容 | 门禁 | |---|---|---| -| PP1 ✅ | `kern cut`（按 call 下标切、活跃性找边界、send/recv 邮箱核、`clock` fill）+ kern-serve 一段一进程的 prefill-only 流水（head 持唯一的池，TCP 发 item，GPU 邮箱传激活） | **2026-10-07 tray07 过**：qwen3-4b 2 段，各段依次跑 logits/token 与整份逐位同（chunk 512/64）；两进程首 token 与单卡同；8193 token conc1 TTFT 313 → 174 ms（serve.md「流水线」） | +| PP1 ✅ | `kern cut`（按 call 下标切、活跃性找边界、send/recv 邮箱核）+ kern-serve 一段一进程的 prefill-only 流水（head 持唯一的池，TCP 发 item，GPU 邮箱传激活） | **2026-10-07 tray07 过**：qwen3-4b 2 段，各段依次跑 logits/token 与整份逐位同（chunk 512/64）；两进程首 token 与单卡同；8193 token conc1 TTFT 313 → 174 ms（serve.md「流水线」） | | PP1.1 ✅ | qwen3-4b 的 KV 每层一个 state，`kern cut` 每段只带自己碰的 state、逐段报告 state 用了几成 | **2026-10-07 tray07 过**：`kern test qwen3-4b` 位一致、`kern run` 文本同；chain / serve gate 位一致；每段页数 122,993 → 245,760 | | PP2 | decode 走流水；K3 的 state 拆分；per-seq state（GDN）与吐 token 的 chunk program | | diff --git a/docs/serve.md b/docs/serve.md index a3cbc822..068dab16 100644 --- a/docs/serve.md +++ b/docs/serve.md @@ -577,8 +577,14 @@ pages }`；之后同一条 TCP（tokio-util `LengthDelimitedCodec` + JSON）上 激活不走 TCP，走 GPU 上的邮箱（manifest.md「流水线切分」），各段设备自己等上游。没有序号、 stop、fault 消息：任何错误进程退出，连接断了对端退出；设备上的邮箱等待 60 s 超时 trap。 只做 prefill：prompt 除末 token 按 chunk 走 chunk program，末 token 走一行的 step，首 token 即结束 -（`max_tokens` 视为 1）。每段每个 item 一行 `pp item` 日志：rows、busy_ns、recv_wait_ns、send_wait_ns -（读 `clock` fill，空闲时读，不加协议消息）。 +（`max_tokens` 视为 1）。item 的编号是它在 TCP 上的次序，各段数得一样。每段 issue 完一个 item 就打一个 +`Runtime::mark`（event），旁边一个线程等它，落地时记一行 `pp item`：stage、item、rows、`issued_us`、 +`done_us`（墙钟）。host 上谁也不为它停：段的周期 = 相邻 done 之差；按 item 号把各段的行拼起来，段 s 开始 +item k = max(本段 issue k、本段 done k−1、上游 done k)，于是每段每 item 的空等与占用都有了 +（`tools/pp_timeline.py`，带 `--gantt`；跨 tray 的段用 `--skew` 扣两台机器的时钟差）。占用里含等下游取走的 +时间，要拆出纯计算就拿该段单独跑这个形状的时间（`kern bench` 段 manifest）去比。 +设备上的 fault（邮箱等待超时 trap）在该 item 的 mark 上冒出来，进程退出。下表的 busy / wait 列是当时 +设备上 globaltimer 打点量的（`clock` fill，已删：它让输出不再是 forward 的纯函数）。 门禁（qwen3-4b，切在 `l18.qkv_proj` 前，2 段）： - 只切不分进程：`pp_chain` 各段依次跑 vs 整份，7 条 prompt（1–2049 token），chunk 512 × 8 步、 diff --git a/schema/manifest-v5.schema.json b/schema/manifest-v5.schema.json index e44071a7..d57a5348 100644 --- a/schema/manifest-v5.schema.json +++ b/schema/manifest-v5.schema.json @@ -1012,11 +1012,6 @@ "const": "error", "description": "Output: a one-word error flag of the collectives, read after every call; nonzero is a failed step.", "type": "string" - }, - { - "const": "clock", - "description": "Output: a pipeline stage's clock, `i64 [ring × 5]`: entry `item % ring` is {item + 1, receive wait begin, end, send wait begin, end} in globaltimer ns, 0 where the stage has no such edge. A cut declares it; the stage reads it when it pleases.", - "type": "string" } ] }, diff --git a/tools/dsv41/attention/ar_trace/src/main.rs b/tools/dsv41/attention/ar_trace/src/main.rs index 7a9015b0..3c56175a 100644 --- a/tools/dsv41/attention/ar_trace/src/main.rs +++ b/tools/dsv41/attention/ar_trace/src/main.rs @@ -104,7 +104,7 @@ fn main() -> anyhow::Result<()> { Fill::CuSeqlens => vec![0, rows as i64], Fill::SpanAt => vec![0], Fill::Blocks => anyhow::bail!("unexpected collective row blocks in DP manifest"), - Fill::Tokens | Fill::Count | Fill::Error | Fill::Clock => continue, + Fill::Tokens | Fill::Count | Fill::Error => continue, }; rt.write_input_at(&f.name, &f.encode(&values), &vars)?; } @@ -159,7 +159,7 @@ fn main() -> anyhow::Result<()> { Fill::CuSeqlens => vec![0, 6], Fill::SpanAt => vec![0], Fill::Blocks => anyhow::bail!("unexpected blocks"), - Fill::Tokens | Fill::Count | Fill::Error | Fill::Clock => continue, + Fill::Tokens | Fill::Count | Fill::Error => continue, }; rt.write_input_at(&f.name, &f.encode(&values), &dv)?; } diff --git a/tools/kernels-src/pp_mailbox.cu b/tools/kernels-src/pp_mailbox.cu index c3bcb897..e5ba84eb 100644 --- a/tools/kernels-src/pp_mailbox.cu +++ b/tools/kernels-src/pp_mailbox.cu @@ -16,28 +16,16 @@ // A wait gives up after `timeout_ns` of globaltimer and traps: the process // dies with a launch failure instead of hanging on a dead stage. // -// Every stage also keeps a clock: `clock` is `i64[RING * 5]`, entry `item % -// RING` = {item + 1, recv wait begin, recv wait end, send wait begin, send -// wait end} in globaltimer ns. The item is read off the counters, never -// passed: a downstream stage's is its box's `empty` before the take, an -// upstream stage's the downstream box's `full` before the post. The waits -// stamp their own spin; a first stage stamps its start as "recv end" and a -// last stage its end as "send begin", so busy = send begin - recv end on -// every stage. One thread writes each field. -// -// kern_pp_wait_empty (in u64 peers[2], out i64 clock[], i64 timeout_ns) +// kern_pp_wait_empty (in u64 peers[2], i64 timeout_ns) // kern_pp_put (in u64 peers[2], in u8 src[], i64 bytes, i64 offset) // kern_pp_post_full (in u64 peers[2]) -// kern_pp_wait_full (in u8 box[], out i64 clock[], i64 timeout_ns) +// kern_pp_wait_full (in u8 box[], i64 timeout_ns) // kern_pp_take (in u8 box[], out u8 dst[], i64 bytes, i64 offset) // kern_pp_post_empty (inout u8 box[]) -// kern_pp_stamp_first(in u64 peers[2], out i64 clock[]) -// kern_pp_stamp_last (in u8 box[], out i64 clock[]) -// waits, posts, stamps: block [1,1,1], grid [1,1,1]; put / take: any grid. +// waits and posts: block [1,1,1], grid [1,1,1]; put / take: any grid. // peers[1] is the downstream box; the upstream's own box (peers[0]) is unused. #define PP_HEADER 256 -#define PP_RING 256 typedef unsigned long long u64; @@ -69,12 +57,6 @@ __device__ __forceinline__ unsigned char* downstream(const u64* peers) { return reinterpret_cast(peers[1]); } -__device__ __forceinline__ long long* entry(long long* clock, u64 item) { - long long* e = clock + (item % PP_RING) * 5; - e[0] = (long long)item + 1; - return e; -} - __device__ void copy(unsigned char* dst, const unsigned char* src, long long bytes) { const long long stride = (long long)gridDim.x * blockDim.x; const long long t = (long long)blockIdx.x * blockDim.x + threadIdx.x; @@ -89,12 +71,10 @@ __device__ void copy(unsigned char* dst, const unsigned char* src, long long byt } } -extern "C" __global__ void kern_pp_wait_empty(const u64* peers, long long* clock, long long timeout_ns) { +extern "C" __global__ void kern_pp_wait_empty(const u64* peers, long long timeout_ns) { const unsigned char* box = downstream(peers); const u64 item = ld_acquire_sys(full_of(box)); - long long* e = entry(clock, item); const u64 t0 = gtimer(); - e[3] = (long long)t0; while (ld_acquire_sys(empty_of(box)) != item) { if ((long long)(gtimer() - t0) > timeout_ns) { printf("kern pp: item %llu waited %lld ns for the downstream stage to take the last one\n", item, @@ -102,7 +82,6 @@ extern "C" __global__ void kern_pp_wait_empty(const u64* peers, long long* clock __trap(); } } - e[4] = (long long)gtimer(); } extern "C" __global__ void kern_pp_put(const u64* peers, const unsigned char* src, long long bytes, @@ -116,11 +95,9 @@ extern "C" __global__ void kern_pp_post_full(const u64* peers) { st_release_sys(full_of(box), ld_acquire_sys(full_of(box)) + 1); } -extern "C" __global__ void kern_pp_wait_full(const unsigned char* box, long long* clock, long long timeout_ns) { +extern "C" __global__ void kern_pp_wait_full(const unsigned char* box, long long timeout_ns) { const u64 item = ld_acquire_sys(empty_of(box)); - long long* e = entry(clock, item); const u64 t0 = gtimer(); - e[1] = (long long)t0; while (ld_acquire_sys(full_of(box)) <= item) { if ((long long)(gtimer() - t0) > timeout_ns) { printf("kern pp: item %llu waited %lld ns for the upstream stage\n", item, timeout_ns); @@ -128,7 +105,6 @@ extern "C" __global__ void kern_pp_wait_full(const unsigned char* box, long long } } __threadfence_system(); - e[2] = (long long)gtimer(); } extern "C" __global__ void kern_pp_take(const unsigned char* box, unsigned char* dst, long long bytes, @@ -140,11 +116,3 @@ extern "C" __global__ void kern_pp_post_empty(unsigned char* box) { __threadfence_system(); st_release_sys(empty_of(box), *empty_of(box) + 1); } - -extern "C" __global__ void kern_pp_stamp_first(const u64* peers, long long* clock) { - entry(clock, ld_acquire_sys(full_of(downstream(peers))))[2] = (long long)gtimer(); -} - -extern "C" __global__ void kern_pp_stamp_last(const unsigned char* box, long long* clock) { - entry(clock, ld_acquire_sys(empty_of(box)) - 1)[3] = (long long)gtimer(); -} From a7f59049bb38af8a32bda75a8560a23c29b74c2a Mon Sep 17 00:00:00 2001 From: JinYan Su <751080330@qq.com> Date: Thu, 8 Oct 2026 15:08:38 +0000 Subject: [PATCH 18/22] manifest: a var says which axis of a call it counts; the roles no longer imply the axes `vars..axis` is `rows`, `groups` or `tray`, at most one var each. Protocol reads a call's axes off the vars. Before, it reverse-engineered them from the roles: the rows var was whichever var the `slot` fill's buffer was shaped over, the groups var the `seq_len` fill's, the tray var whatever third var a fill or line table used. So every forward had to carry `slot`, `seq_len` and `token`, which is a serving loop's habit written into the function's type: a pipeline stage past the first takes activations, not tokens, and reads no length, yet has both axes. The cut-stage special cases in Protocol (`cut.stage > 0`: no token required, groups guessed from the page table) are gone with the inference; Protocol does not read `cut` at all. `token_rows()` is an Option now and the drivers that feed tokens (kern run, kern test) say so when a manifest takes none; `slots()` / `seq_lens()` go, `any(Fill::...)` covers them. Schema golden regenerated; every generator and example manifest with a batch names its axes (k3: `rows` is the tray axis, as the old inference found it). Gate: kern-manifest / kern-pool / kern-run / kern-test / kern-serve tests pass; qwen3-4b cut at 218 with the axis manifest: pp_chain chunk 512 x 8 and 64 x 16 bitwise identical to unsplit, serve gate first tokens equal, killed stage -> head exits, frozen stage -> head exits at 61 s. The K3 PP8 downstream stages (no token, no seq_len) pass Protocol without a special case. Co-Authored-By: Claude Fable 5.1 Signed-off-by: JinYan Su <751080330@qq.com> --- crates/kern-manifest/src/protocol.rs | 154 +++++------------- crates/kern-manifest/src/types.rs | 25 +++ crates/kern-manifest/src/verify.rs | 8 +- .../kern-manifest/tests/fixtures/plain.json | 2 +- crates/kern-manifest/tests/fixtures/tray.json | 2 +- crates/kern-manifest/tests/protocol.rs | 39 +++-- crates/kern-manifest/tests/verify.rs | 8 + crates/kern-run/examples/agentx_replay.rs | 2 +- crates/kern-run/src/lib.rs | 6 +- crates/kern-test/src/workload.rs | 4 +- crates/kern-test/tests/common/mod.rs | 8 +- docs/manifest.md | 20 ++- docs/roadmap.md | 3 + examples/k3-4l-ep1.json | 6 +- examples/k3-4l-ep4-tp4-prefill.json | 6 +- examples/k3-4l-ep4-tp4.json | 6 +- examples/k3-4l-ep4.json | 6 +- examples/k3-4l-ep8.json | 6 +- examples/k3-ep4-tp4.json | 6 +- examples/k3-ep4.json | 6 +- examples/k3-ep8.json | 6 +- examples/qwen3-4b-dspark.json | 4 +- examples/qwen3-4b-silu-mined.json | 4 +- examples/qwen3-4b.json | 4 +- examples/qwen3.8-27b-dflash2.json | 4 +- examples/qwen3.8-27b.json | 4 +- schema/manifest-v5.schema.json | 31 ++++ tools/dsv41/auxiliary/serving.py | 2 +- tools/gen_k3.py | 3 +- tools/gen_qwen35.py | 2 +- tools/gen_qwen3_decode.py | 2 +- tools/toy/gen.py | 2 +- 32 files changed, 207 insertions(+), 184 deletions(-) diff --git a/crates/kern-manifest/src/protocol.rs b/crates/kern-manifest/src/protocol.rs index 1caf3f53..dfafd965 100644 --- a/crates/kern-manifest/src/protocol.rs +++ b/crates/kern-manifest/src/protocol.rs @@ -16,14 +16,17 @@ //! it hands out is typed. //! //! The axes. A call has `groups` sequences of `rows` rows each on this -//! rank. The var the rows go in is the one the `slot` fill's buffer is -//! over ([`Axis::Rows`]); the var the sequences go in is the `seq_len` -//! buffer's ([`Axis::Groups`]). A manifest whose tray batch spans several -//! ranks (a `tp` group) lays some buffers out over a third var, the whole -//! tray's rows with this rank's first ([`Axis::Tray`]); any fill or line -//! table over a var that is neither of the first two names it. Fixed -//! lengths ([`Axis::Fixed`]) are for what is declared at its bound: the -//! `cu_seqlens` buffer, the collectives' one-word `error` flag. +//! rank, and the vars that count them are declared as such: the var with +//! `axis: rows` ([`Axis::Rows`]), the var with `axis: groups` +//! ([`Axis::Groups`]). A manifest whose tray batch spans several ranks (a +//! `tp` group) lays some buffers out over a third var, the whole tray's +//! rows with this rank's first, declared `axis: tray` ([`Axis::Tray`]). +//! A fill or a line table is over one of these three; fixed lengths +//! ([`Axis::Fixed`]) are for what is declared at its bound: the +//! `cu_seqlens` buffer, the collectives' one-word `error` flag. The roles +//! say nothing about the axes: a manifest that reads no sequence length +//! still has a sequence axis, and a pipeline stage past the first takes +//! no tokens at all. //! //! A program with a `batch` is a [`Forward`]. Which tokens it hands back is //! not a role but dataflow: the `tokens` output it writes (one per @@ -228,10 +231,6 @@ impl Protocol { fn check_with(m: &Verified, sampled: bool) -> Result { let mut errs = Vec::new(); let var_max = |v: &str| m.vars.get(v).map(|v| v.max); - let axis_var = |b: &Buffer| match b.shape.first() { - Some(Dim::Var(v)) => Some(v.clone()), - _ => None, - }; let one = |fill: Fill| -> Option<(&String, &Buffer)> { let mut it = m.buffers.iter().filter(|(_, b)| b.fill == Some(fill)); let first = it.next(); @@ -242,97 +241,47 @@ impl Protocol { } }; - // The axes come from the two fills every call needs. - let axis_of = |fill: Fill, errs: &mut Vec| -> Option { - let all: Vec<&String> = m.buffers.iter().filter(|(_, b)| b.fill == Some(fill)).map(|(n, _)| n).collect(); - match all.as_slice() { - [] => { - errs.push(format!("no input has fill `{fill}`")); - None - } - [name] => { - let b = &m.buffers[*name]; - match (b.shape.as_slice(), axis_var(b).and_then(|v| var_max(&v).map(|max| Bound { var: v, max }))) { - ([Dim::Var(_)], Some(bound)) => Some(bound), - _ => { - errs.push(format!("`{name}` (fill `{fill}`) is shaped {:?}, expected []", b.shape)); - None - } - } - } - _ => { - errs.push(format!("fill `{fill}` is on {} buffers ({all:?}), expected one", all.len())); - None - } - } - }; - let rows = axis_of(Fill::Slot, &mut errs); - // A stage past the first of a cut may read no sequence length (its - // attention lengths came through the mailbox): its sequences are - // then the rows of its page table, or none of its own. - let downstream = m.cut.as_ref().is_some_and(|c| c.stage > 0); - let groups = if downstream && !m.buffers.values().any(|b| b.fill == Some(Fill::SeqLen)) { - let paged = |b: &Buffer| { - b.domain - .as_ref() - .and_then(|d| d.index_into.as_deref()) - .and_then(|s| m.states.get(s)) - .is_some_and(|s| s.bytes_per_token > 0) - }; - let var = m - .buffers - .values() - .filter(|b| b.kind == BufferKind::Input && b.fill.is_none() && paged(b)) - .find_map(axis_var) - .unwrap_or_default(); - Some(Bound { max: var_max(&var).unwrap_or(1), var }) - } else { - axis_of(Fill::SeqLen, &mut errs) + // The axes are the vars declared to count them (verified: one each). + let axis_var = |axis: VarAxis| -> Option { + m.vars.iter().find(|(_, v)| v.axis == Some(axis)).map(|(n, v)| Bound { var: n.clone(), max: v.max }) }; - if let (Some(r), Some(g)) = (&rows, &groups) { - if r.var == g.var { - errs.push(format!( - "`slot` and `seq_len` are both over var `{}`; rows and sequences need their own", - r.var - )); + let (rows, groups, tray) = (axis_var(VarAxis::Rows), axis_var(VarAxis::Groups), axis_var(VarAxis::Tray)); + for (axis, found) in [(VarAxis::Rows, &rows), (VarAxis::Groups, &groups)] { + if found.is_none() { + errs.push(format!("no var has `axis: {axis}`: nothing counts a call's {axis}")); } } let (Some(rows), Some(groups)) = (rows, groups) else { return Err(ProtocolErrors(errs)); }; - - // A third var, over which something spans the tray batch. - let mut tray_vars: BTreeSet = BTreeSet::new(); - for b in m.buffers.values().filter(|b| b.fill.is_some() || is_line_table(m, b)) { - let col = if is_line_table(m, b) { b.shape.get(1) } else { b.shape.first() }; - if let Some(Dim::Var(v)) = col { - if *v != rows.var && *v != groups.var { - tray_vars.insert(v.clone()); - } - } - } - if tray_vars.len() > 1 { - errs.push(format!("fills and line tables span {} vars besides rows and sequences ({tray_vars:?}), at most one names the tray batch", tray_vars.len())); - } - let tray = tray_vars.first().and_then(|v| var_max(v).map(|max| Bound { var: v.clone(), max })); - let classify = |v: &str| -> Axis { + let classify = |v: &str| -> Option { if v == rows.var { - Axis::Rows + Some(Axis::Rows) } else if v == groups.var { - Axis::Groups + Some(Axis::Groups) + } else if tray.as_ref().is_some_and(|t| t.var == v) { + Some(Axis::Tray) } else { - Axis::Tray + None + } + }; + let axis_of = |name: &str, what: &str, v: &str, errs: &mut Vec| -> Option { + let axis = classify(v); + if axis.is_none() { + errs.push(format!("{what} `{name}` is over var `{v}`, which counts no axis of a call")); } + axis }; // Fills. let mut fills: Vec = Vec::new(); for (name, b) in &m.buffers { let Some(fill) = b.fill else { continue }; + let what = format!("fill `{fill}`"); let (axis, width) = match b.shape.as_slice() { - [Dim::Var(v)] => (classify(v), 1), - [Dim::Const(c)] => (Axis::Fixed(*c), 1), - [Dim::Var(v), Dim::Const(w)] => (classify(v), *w), + [Dim::Var(v)] => (axis_of(name, &what, v, &mut errs), 1), + [Dim::Const(c)] => (Some(Axis::Fixed(*c)), 1), + [Dim::Var(v), Dim::Const(w)] => (axis_of(name, &what, v, &mut errs), *w), s => { errs.push(format!( "`{name}` (fill `{fill}`) is shaped {s:?}, expected [], [] or [, w]" @@ -340,6 +289,7 @@ impl Protocol { continue; } }; + let Some(axis) = axis else { continue }; let ok = match fill { Fill::Token => width == 1 && matches!(axis, Axis::Rows | Axis::Tray | Axis::Groups), Fill::Position | Fill::Valid | Fill::Slot => width == 1 && axis == Axis::Rows, @@ -373,11 +323,6 @@ impl Protocol { errs.push(format!("fill `{fill}` over {axis:?} is on {n} buffers, expected one")); } } - // A stage past the first of a cut takes activations, not tokens. - let fed = m.cut.as_ref().is_none_or(|c| c.stage == 0); - if fed && !fills.iter().any(|f| f.fill == Fill::Token && f.axis != Axis::Groups) { - errs.push("no input has fill `token` over the rows: nothing carries the tokens a call feeds".into()); - } if fills.iter().filter(|f| f.fill == Fill::Token && f.axis == Axis::Groups).count() > 1 { errs.push("fill `token` over the sequences is on more than one buffer".into()); } @@ -418,7 +363,7 @@ impl Protocol { continue; } }; - let axis = classify(cols); + let Some(axis) = axis_of(name, "line table", cols, &mut errs) else { continue }; if axis == Axis::Rows { errs.push(format!("line table `{name}` has a column per row (`{cols}`); lines are per sequence")); continue; @@ -701,26 +646,9 @@ impl Protocol { self.fills.iter().find(|f| f.fill == fill) } - /// The tokens fed, one per row: over this rank's rows or the tray's. - pub fn token_rows(&self) -> &Filled { - self.filled(Fill::Token, Axis::Rows).or_else(|| self.filled(Fill::Token, Axis::Tray)).expect("checked") - } - - pub fn slots(&self) -> &Filled { - self.filled(Fill::Slot, Axis::Rows).expect("checked") + /// The tokens fed, one per row (over this rank's rows or the tray's), + /// when the manifest takes any: a pipeline stage past the first does not. + pub fn token_rows(&self) -> Option<&Filled> { + self.filled(Fill::Token, Axis::Rows).or_else(|| self.filled(Fill::Token, Axis::Tray)) } - - pub fn seq_lens(&self) -> &Filled { - self.filled(Fill::SeqLen, Axis::Groups).expect("checked") - } -} - -/// An input indexing a per-sequence state. -fn is_line_table(m: &Manifest, b: &Buffer) -> bool { - b.kind == BufferKind::Input - && b.domain - .as_ref() - .and_then(|d| d.index_into.as_deref()) - .and_then(|s| m.states.get(s)) - .is_some_and(State::is_per_seq) } diff --git a/crates/kern-manifest/src/types.rs b/crates/kern-manifest/src/types.rs index 957904c8..9c415c47 100644 --- a/crates/kern-manifest/src/types.rs +++ b/crates/kern-manifest/src/types.rs @@ -250,6 +250,31 @@ pub struct Cut { pub struct Var { /// Upper bound, e.g. `2048` for the token count of a prefill chunk. pub max: u64, + /// The axis of a call this var counts, for the serving loop: every buffer shaped over it has one entry per element of that axis, e.g. `"rows"`. At most one var per axis; a manifest with a `batch` names a `rows` var and a `groups` var. Absent for a var that sizes something else (a run's length, a context). + #[serde(default, skip_serializing_if = "Option::is_none")] + pub axis: Option, +} + +/// What a var counts in one call of a program: this rank's rows, its sequences, or the rows of the whole tray batch. The runtime never reads it. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)] +#[serde(rename_all = "snake_case")] +pub enum VarAxis { + /// One per row of this rank's call, e.g. `tokens`. + Rows, + /// One per sequence of this rank's call, e.g. `seqs`. + Groups, + /// One per row of the whole tray batch, this rank's rows first, e.g. a `tp` manifest's `rows`. + Tray, +} + +impl fmt::Display for VarAxis { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str(match self { + VarAxis::Rows => "rows", + VarAxis::Groups => "groups", + VarAxis::Tray => "tray", + }) + } } impl Var { diff --git a/crates/kern-manifest/src/verify.rs b/crates/kern-manifest/src/verify.rs index d00ac0df..2c1b2e4e 100644 --- a/crates/kern-manifest/src/verify.rs +++ b/crates/kern-manifest/src/verify.rs @@ -4,7 +4,7 @@ //! //! Checks: //! 1. schema_version -//! 2. vars: max > 0 +//! 2. vars: max > 0, at most one var per axis //! 3. states: exactly one of bytes_per_token / bytes / bytes_per_seq is non-zero, //! or none of them and a well-formed `host` layout; a host state is never //! exported (`of`); a domain may index it, declaring the host's ids @@ -187,6 +187,12 @@ fn diagnostics(m: &Manifest) -> Vec { errs.push(format!("var `{name}`: max must be >= {}", Var::MIN)); } } + for axis in [VarAxis::Rows, VarAxis::Groups, VarAxis::Tray] { + let named: Vec<&String> = m.vars.iter().filter(|(_, v)| v.axis == Some(axis)).map(|(n, _)| n).collect(); + if named.len() > 1 { + errs.push(format!("vars {named:?} all count the `{axis}` axis; one var does")); + } + } let vars_max: BTreeMap = m.vars.iter().map(|(k, v)| (k.clone(), v.max)).collect(); let vars_min: BTreeMap = m.vars.keys().map(|k| (k.clone(), Var::MIN)).collect(); diff --git a/crates/kern-manifest/tests/fixtures/plain.json b/crates/kern-manifest/tests/fixtures/plain.json index 5bdf3985..f5c4022c 100644 --- a/crates/kern-manifest/tests/fixtures/plain.json +++ b/crates/kern-manifest/tests/fixtures/plain.json @@ -1,6 +1,6 @@ { "schema_version": 5, "model": "t", - "vars": {"tokens": {"max": 8}, "seqs": {"max": 4}}, + "vars": {"tokens": {"max": 8, "axis": "rows"}, "seqs": {"max": 4, "axis": "groups"}}, "states": {"kv": {"bytes_per_token": 1}, "gdn": {"bytes_per_seq": 24}}, "buffers": { "token_ids": {"kind": "input", "dtype": "i64", "shape": ["tokens"], "fill": "token"}, diff --git a/crates/kern-manifest/tests/fixtures/tray.json b/crates/kern-manifest/tests/fixtures/tray.json index 9879f39b..f23d705f 100644 --- a/crates/kern-manifest/tests/fixtures/tray.json +++ b/crates/kern-manifest/tests/fixtures/tray.json @@ -1,6 +1,6 @@ { "schema_version": 5, "model": "t", - "vars": {"tokens": {"max": 8}, "seqs": {"max": 8}, "rows": {"max": 32}}, + "vars": {"tokens": {"max": 8, "axis": "rows"}, "seqs": {"max": 8, "axis": "groups"}, "rows": {"max": 32, "axis": "tray"}}, "topology": {"groups": {"ep": 4, "tp": 4}}, "states": {"kv": {"bytes_per_token": 1}, "kda": {"bytes_per_seq": 24}}, "buffers": { diff --git a/crates/kern-manifest/tests/protocol.rs b/crates/kern-manifest/tests/protocol.rs index 51c77a89..03a34a42 100644 --- a/crates/kern-manifest/tests/protocol.rs +++ b/crates/kern-manifest/tests/protocol.rs @@ -89,7 +89,11 @@ fn plain_contract() { assert_eq!((p.rows.var.as_str(), p.rows.max, p.groups.var.as_str(), p.groups.max), ("tokens", 8, "seqs", 4)); assert_eq!(p.tray, None); assert_eq!( - (p.token_rows().name.as_str(), p.slots().name.as_str(), p.seq_lens().name.as_str()), + ( + p.token_rows().unwrap().name.as_str(), + p.any(Fill::Slot).unwrap().name.as_str(), + p.any(Fill::SeqLen).unwrap().name.as_str() + ), ("token_ids", "slot_mapping", "seq_lens") ); assert_eq!(p.any(Fill::CuSeqlens).map(|f| f.axis), Some(Axis::Fixed(5))); @@ -196,7 +200,7 @@ fn span_rules() { fn tray_contract() { let p = protocol(&tray()).unwrap(); assert_eq!(p.tray, Some(Bound { var: "rows".into(), max: 32 })); - assert_eq!(p.token_rows().axis, Axis::Tray); + assert_eq!(p.token_rows().unwrap().axis, Axis::Tray); assert_eq!(p.line_tables[0].axis, Axis::Tray); assert_eq!(p.any(Fill::Error).map(|f| f.name.as_str()), Some("tp_err")); assert_eq!(p.once, vec!["tp_init".to_string()]); @@ -206,27 +210,36 @@ fn tray_contract() { #[test] fn encodes_by_dtype() { let p = protocol(&plain()).unwrap(); - assert_eq!(p.seq_lens().encode(&[3, -1]), vec![3, 0, 0, 0, 255, 255, 255, 255]); - assert_eq!(p.slots().encode(&[2]), 2i64.to_le_bytes()); - assert_eq!(p.seq_lens().decode(&[7, 0, 0, 0]), vec![7]); + let (slots, seq_lens) = (p.any(Fill::Slot).unwrap(), p.any(Fill::SeqLen).unwrap()); + assert_eq!(seq_lens.encode(&[3, -1]), vec![3, 0, 0, 0, 255, 255, 255, 255]); + assert_eq!(slots.encode(&[2]), 2i64.to_le_bytes()); + assert_eq!(seq_lens.decode(&[7, 0, 0, 0]), vec![7]); } +/// The axes are the vars', not the roles': a manifest may read no slot and +/// no length (a pipeline stage past the first) and still have both axes. #[test] -fn missing_pieces_are_all_named() { +fn axes_are_declared_on_the_vars() { let mut m = plain(); - m.buffers.get_mut("slot_mapping").unwrap().fill = None; m.buffers.get_mut("seq_lens").unwrap().fill = None; + m.buffers.get_mut("token_ids").unwrap().fill = None; + let p = protocol(&m).unwrap(); + assert_eq!((p.rows.var.as_str(), p.groups.var.as_str(), p.token_rows()), ("tokens", "seqs", None)); + let mut m = plain(); + m.vars.get_mut("tokens").unwrap().axis = None; + m.vars.get_mut("seqs").unwrap().axis = None; let Err(e) = protocol(&m) else { panic!() }; assert_eq!(e.len(), 2); - rejects(&m, "no input has fill `slot`"); - rejects(&m, "no input has fill `seq_len`"); + rejects(&m, "no var has `axis: rows`"); + rejects(&m, "no var has `axis: groups`"); + let mut m = plain(); + m.vars.insert("other".into(), serde_json::from_str(r#"{"max": 8}"#).unwrap()); + m.buffers.get_mut("positions").unwrap().shape = vec![Dim::Var("other".into())]; + rejects(&m, "fill `position` `positions` is over var `other`, which counts no axis"); } #[test] fn shape_rules() { - let mut m = plain(); - m.buffers.get_mut("token_ids").unwrap().shape = vec![Dim::Var("seqs".into())]; - rejects(&m, "no input has fill `token` over the rows"); let mut m = plain(); m.buffers.get_mut("cu_seqlens_q").unwrap().shape = vec![Dim::Const(4)]; rejects(&m, "expected [n] with n >= groups + 1"); @@ -238,7 +251,7 @@ fn shape_rules() { rejects(&m, "column per row"); let mut m = plain(); m.buffers.get_mut("seq_lens").unwrap().shape = vec![Dim::Var("tokens".into())]; - rejects(&m, "both over var `tokens`"); + rejects(&m, "expected [groups]"); } #[test] diff --git a/crates/kern-manifest/tests/verify.rs b/crates/kern-manifest/tests/verify.rs index 0b093cee..f2a7714f 100644 --- a/crates/kern-manifest/tests/verify.rs +++ b/crates/kern-manifest/tests/verify.rs @@ -80,6 +80,14 @@ fn assert_err(v: serde_json::Value, needle: &str) { assert!(errs.iter().any(|e| e.contains(needle)), "no error containing `{needle}` in {errs:#?}"); } +#[test] +fn one_var_per_axis() { + let mut v = base(); + v["vars"]["tokens"]["axis"] = serde_json::json!("rows"); + v["vars"]["more"] = serde_json::json!({"max": 4, "axis": "rows"}); + assert_err(v, "all count the `rows` axis"); +} + #[test] fn base_manifest_verifies() { check(base()).unwrap(); diff --git a/crates/kern-run/examples/agentx_replay.rs b/crates/kern-run/examples/agentx_replay.rs index 2f8f8c79..1e3d17f2 100644 --- a/crates/kern-run/examples/agentx_replay.rs +++ b/crates/kern-run/examples/agentx_replay.rs @@ -144,7 +144,7 @@ fn manifest(unit: usize, kv_bytes: u64, state: u64, slots: usize, row: usize) -> }; Manifest::from_json(&format!( r#"{{ - "schema_version": 5, "model": "replay", "vars": {{"tokens": {{"max": 1}}, "seqs": {{"max": {seqs}}}}}, + "schema_version": 5, "model": "replay", "vars": {{"tokens": {{"max": 1, "axis": "rows"}}, "seqs": {{"max": {seqs}, "axis": "groups"}}}}, "states": {{"kv": {{"bytes_per_token": {kv_bytes}}}{state_json}}}, "buffers": {{ "block_table": {{"kind": "input", "dtype": "i32", "shape": ["seqs", {row}], "domain": {{"index_into": "kv", "stride": {unit}}}}}{line_json} diff --git a/crates/kern-run/src/lib.rs b/crates/kern-run/src/lib.rs index c0f33415..ff83bfbc 100644 --- a/crates/kern-run/src/lib.rs +++ b/crates/kern-run/src/lib.rs @@ -463,9 +463,9 @@ impl Caller { /// Vocabulary size as declared by the token fill's domain (1000 if none). fn vocab(&self) -> u64 { let m = &self.rt.manifest; - m.buffers[&self.protocol.token_rows().name] - .domain - .as_ref() + self.protocol + .token_rows() + .and_then(|t| m.buffers[&t.name].domain.as_ref()) .and_then(|d| d.resolve(m, &self.protocol.vars(1, 1, 1, 1), &self.rt.provision()).ok()) .and_then(|r| r.hi) .map_or(1000, |hi| hi as u64 + 1) diff --git a/crates/kern-test/src/workload.rs b/crates/kern-test/src/workload.rs index 12e668df..3a964a9d 100644 --- a/crates/kern-test/src/workload.rs +++ b/crates/kern-test/src/workload.rs @@ -50,7 +50,9 @@ pub fn sample(o: &Options, m: &Manifest, pr: &Protocol, p: Provision, page: u64) let capacity = p.tokens; let mut rng = Rng(o.seed ^ 0x776f_726b_6c6f_6164); let tmax = pr.rows.max.max(1); - let tokens = &pr.token_rows().name; + let Some(tokens) = pr.token_rows().map(|t| &t.name) else { + anyhow::bail!("the manifest takes no tokens; a workload feeds some") + }; let vocab = m.buffers[tokens] .domain .as_ref() diff --git a/crates/kern-test/tests/common/mod.rs b/crates/kern-test/tests/common/mod.rs index a53162a1..ab844713 100644 --- a/crates/kern-test/tests/common/mod.rs +++ b/crates/kern-test/tests/common/mod.rs @@ -228,7 +228,7 @@ impl Fixture { } let mut m = serde_json::json!({ "schema_version": 5, "model": "fake", - "vars": {"tokens": {"max": 8}, "seqs": {"max": 1}}, + "vars": {"tokens": {"max": 8, "axis": "rows"}, "seqs": {"max": 1, "axis": "groups"}}, "states": {"kv": {"bytes_per_token": LAYERS * D * 4}}, "buffers": { "token_ids": {"kind": "input", "dtype": "i64", "shape": ["tokens"], "fill": "token", @@ -495,9 +495,9 @@ impl Side for Fake { r.bufs.get_mut(&f.name).unwrap()[..b.len()].copy_from_slice(&b); } }; - put(p.token_rows(), ids); - put(p.slots(), &(pos..pos + c as i64).collect::>()); - put(p.seq_lens(), &[pos + c as i64]); + put(p.token_rows().expect("the fixture takes tokens"), ids); + put(p.any(Fill::Slot).expect("the fixture takes slots"), &(pos..pos + c as i64).collect::>()); + put(p.any(Fill::SeqLen).expect("the fixture takes lengths"), &[pos + c as i64]); if let Some(f) = p.filled(Fill::Position, Axis::Rows) { put(f, &(pos..pos + c as i64).collect::>()); } diff --git a/docs/manifest.md b/docs/manifest.md index 44657b98..0130999a 100644 --- a/docs/manifest.md +++ b/docs/manifest.md @@ -38,7 +38,9 @@ cut cut 流水线的一段：id、第几段、共 - `vars`：caller 每次调用时提供的标量（如 `tokens`），声明 `max`，下界 恒为 1；所有静态校验在界上进行，运行时拒绝越界值。**会改变内存尺寸或 launch 几何的标量才是 var**；别的标量（temperature 之类）是数据，走 - `[1]` 形状的 input buffer。 + `[1]` 形状的 input buffer。可选的 `axis`（`rows` / `groups` / `tray`） + 说这个 var 数的是一次调用的哪根轴——本卡的行、本卡的序列、整个 tray + batch 的行——每根轴至多一个 var；runtime 不读它（见「Serving 协议」）。 - `states`：不透明持久内存。**runtime 只知道字节数**——`bytes_per_token` （按 token 容量伸缩：paged KV）、`bytes_per_seq`（每个活跃序列一个 slot：GDN 的 conv + SSM 递归状态；runtime 供应 `seqs.max + 2` 个 slot， @@ -379,11 +381,14 @@ program 接受几组几行、跑完从哪读 token。v3 把这层写在 caller 值一张图是 manifest 作者量过才知道、runtime 猜不到的事，所以写在这里； `--eager` 只是调试开关，全部按 launch 跑。要求有 `batch`。 -轴全部从 fill 派生：行轴的 var 是 `slot` fill 的维（`tokens`），组轴的 -var 是 `seq_len` fill 的维（`seqs`），fill 或 line 表跨过的第三个 var 是 -tray 轴（k3 的 `rows`）；manifest 里没有一个名字是 caller 认得的。派生 -这层的是 `kern_manifest::Protocol::check(&Manifest)`（纯函数，不碰 GPU， -verify 之后的第二遍）：每个 fill 至多一个 buffer、形状与角色相符、 +轴是 var 自己声明的（`axis: rows` 的 `tokens`、`axis: groups` 的 `seqs`、 +`axis: tray` 的 k3 `rows`），fill 只说角色，不定轴：一份不读 `seq_len` +的 manifest 照样有序列轴，流水线首段之外的段根本不吃 token（v5 之前轴从 +`slot` / `seq_len` fill 所在的维反推，于是每个 forward 都"必须"有这两个 +角色——把 scheduler 的习惯写成了函数的类型）。manifest 里没有一个名字是 +caller 认得的。读这层的是 `kern_manifest::Protocol::check(&Manifest)` +（纯函数，不碰 GPU，verify 之后的第二遍）：有 `batch` 就得有 rows 与 +groups 两根轴的 var、每个 fill 至多一个 buffer、形状与角色相符、 `[seqs, r]` 的 `tokens` 的 r 是某个 `batch` 的 `rows`、`count` 只在有它时 合法、两个 program 不得同形状（span 算形状的一部分）、`groups × rows ≤ tokens.max`、`span` 只落在 `rows: 1` 的 program 上、全 manifest 一个 span @@ -435,7 +440,8 @@ kern-serve 只拿 `Protocol` 驱动 runtime，不读 JSON（CI 用 grep 保证 组两端装的是不同的 manifest，只在 `of` buffer（同名同字节数的 `pp..box`，上游那份是不用的 workspace + export）上对上—— `import_peers` 本来就只按名字和字节数匹配，runtime 不用改。 -- 首段之外不需要 `token` fill（`Protocol` 只对 `cut.stage == 0` 要求它）。 +- 首段之外的段没有 `token` fill，也常常没有 `seq_len`（长度随激活从邮箱来）， + `Protocol` 对它们没有特判：轴写在 var 上，角色有就用。 一个 chunk program 若自己吐 token（hybrid GDN）则拒切：prefill-only 的 流水线约定首个 token 从一行的 step program 出。 diff --git a/docs/roadmap.md b/docs/roadmap.md index f5b008bd..1351dea1 100644 --- a/docs/roadmap.md +++ b/docs/roadmap.md @@ -65,12 +65,15 @@ kern 在内部 RSI 里是 kernel agent 的反馈环境（`kern test` 裁判、`k | PP1 ✅ | `kern cut`（按 call 下标切、活跃性找边界、send/recv 邮箱核）+ kern-serve 一段一进程的 prefill-only 流水（head 持唯一的池，TCP 发 item，GPU 邮箱传激活） | **2026-10-07 tray07 过**：qwen3-4b 2 段，各段依次跑 logits/token 与整份逐位同（chunk 512/64）；两进程首 token 与单卡同；8193 token conc1 TTFT 313 → 174 ms（serve.md「流水线」） | | PP1.1 ✅ | qwen3-4b 的 KV 每层一个 state，`kern cut` 每段只带自己碰的 state、逐段报告 state 用了几成 | **2026-10-07 tray07 过**：`kern test qwen3-4b` 位一致、`kern run` 文本同；chain / serve gate 位一致；每段页数 122,993 → 245,760 | | PP2 | decode 走流水；K3 的 state 拆分；per-seq state（GDN）与吐 token 的 chunk program | | +| PP3 | manifest 说得出"值从另一份 manifest 来"：跨切点的 buffer 是一种边界 kind（运行前从上游到达 / 运行后交给下游），边是 manifest 之间的有向关系，runtime 负责等、拷、通知；`kern cut` 只切 call 列表、标边界，不再生成 `pp.*` op，`pp_mailbox.cu` 退成 runtime 的一种传输（NVLink 邮箱 / RDMA / 同进程拷贝可换） | 切出的段里没有 `pp.*` op；chain / serve gate 位一致 | +| EP-PP | EP 也一卡一进程：head 持 N 个池、全连握手分发 handle 与 NcclId、每步给每个 rank 发它自己的 item、token 回 head；`Tray` 的单进程多 runtime 路径删除；manifest 不变 | K3 EP8 decode 跨两 tray 进 kern-serve，conc1 文本与单 tray 同；每 rank 的 `pp item` 时间线 | ## 协议线 | 级 | 内容 | 门禁 | |---|---|---| | V4 ✅ | serving 协议进 manifest（schema v4，设计 v4-design.md）：buffer 的 `fill`、program 的 `batch` / `once`，`spec` 块删除；`Verified` newtype + `Protocol::check` 投影，kern-run / kern-serve 不读 JSON、不认名字（CI grep）；投机轮统一成 `round` program（dspark 也补了：splice / 计数 / prefill 收编 head 与 precompute 全在设备上），`--spec` → `--rows`；`kern verify` 打印协议事实 | **2026-09-03 tray03 通过**：`kern test` 位一致；qwen3-4b / qwen3.8-27b / dspark / dflash2 四份 conc1 与 `kern run` 逐字同；dflash2 与 v3 `--spec` 逐字节同；dspark 7 行 round 在一个 0.125 的 bf16 平局上与 v3 的 8 行 verify 分叉（核噪声）；conc32 接受率 34% / 24% 不塌，5930 / 1709 tok/s。记录 v4-design.md §9、serve.md | +| V5.1 ✅ | 轴写在 var 上（`axis: rows / groups / tray`），`Protocol` 不再从 `slot` / `seq_len` fill 反推行轴与组轴、不再从"第三个 var"猜 tray 轴、不读 `cut`；`token` / `slot` / `seq_len` 从每个 forward 必有的角色变成有就用（流水线首段之外的段没有它们） | **2026-10-08**：kern-manifest / kern-run / kern-test / kern-serve 测试过；qwen3-4b 2 段 chain / serve gate 位一致；K3 PP8 的下游段（无 token、无 seq_len）不特判通过 `Protocol` | | W1 ✅ | 权重就是 HF checkpoint：weight buffer 的 `bind` 列出拼成它的张量段（qkv / gate_up 拼接、draft fc 列块、tied lm_head），runtime 只读 safetensors header、按段 memcpy / 2D copy；派生表（Gemma +1、A_log f32、rope、kv_scales、chunk 索引）是 `carry`，由 `load` once program 用 `weight_prep.cu` 算；`tools/export_*.py` 与自定义 .safetensors 全删，`--weights` 给 snapshot 目录（K3 的 bind 未写，examples 仍是逐 buffer 同名的平凡 bind） | **2026-09-08 tray03 通过**：qwen3.8-27b / dflash2 / qwen3-4b-dspark 的每种 bind 形态（拼接、列块、按层 cat、tied）与每种派生表的 dump 对旧导出逐字节相同（qwen3-4b 的 rope 表 2078/5.2M 元素差 1 bf16 ulp：旧表是 CPU torch 算的）；`kern test qwen3-4b` 位一致，`kern run` 文本同；qwen3.8-27b plain 与 dflash2 8 行 round（`qwen38_compare.py`，5 条 prompt × 400 token）与旧导出在同一 binary 上逐 token 相同，对 vLLM 参考的分叉点两者一样、全在 margin ≤ 0.125 的近平局上 | | H1 ✅ | 宿主模式：manifest 的 `host` state（宿主分配、kern 验 layout）、`bind_host`（绑定时编译，可重绑）/ `enqueue_after`（event 挂宿主 stream，可进宿主的图）/ `region`；`crates/kern-py`（`import kern`）；`python/kern_vllm` 把 manifest 装成 vLLM MRv2 的模型，KV / 调度 / graph / 采样归 vLLM（vllm.md） | **2026-09-25 单卡 GB300 过**：Qwen3.8-27B 在 vLLM 里 eager 与 FULL decode graph 下 32 个 greedy token 与 `kern run` 逐个同；32 条并发对串行分叉点 margin 全 ≤ 0.25；在线 A/B 与 TTFT 扫描在 vllm.md | | H2 ✅ | SGLang 宿主：`python/kern_sglang`（`sglang.srt.plugins` 入口，沿用 checkpoint 的 architecture 名，SGLang 按 config 建池、kern 适配池）；`tools/qwen38_sglang.py` 把 K/V/conv/ssm 拆成每层四族 host state，碰 conv state 的核另编 `-DCONV_DIM_MAJOR` 版（源码与 cubin 在 HF artifact repo，不带宏编译字节不变）；examples 与 vLLM RSI 优化版两种 base 都能转；prefix cache 用 `extra_buffer`：extend 在快照点切成两次 `prefill`，decode 在图内用 SGLang 的 masked copy 存快照（sglang.md） | **2026-10-05 单卡 GB300（tray04）过**：SGLang nightly `f70e8c68` 里 greedy 与 `kern run` 一致到参考自身 top-2 精确平局的一步；prefill 快照与 decode 快照命中后的输出与 flush 后冷算逐 token 同，8 条并发分叉点 margin 全 0.125；RSI 版 `kern test` 对其 vLLM 版每个 span 位一致、greedy 48/48 同 `kern run`；AgentX 一小时对同参数（page 64、trtllm_mha）的 SGLang 多服务 10.5% 请求，表在 sglang.md | diff --git a/examples/k3-4l-ep1.json b/examples/k3-4l-ep1.json index 8d13225f..626eca3a 100644 --- a/examples/k3-4l-ep1.json +++ b/examples/k3-4l-ep1.json @@ -3,13 +3,13 @@ "model": "kimi-k3-pruned-75pct/4l/ep1", "vars": { "tokens": { - "max": 8 + "max": 8, "axis": "rows" }, "seqs": { - "max": 8 + "max": 8, "axis": "groups" }, "rows": { - "max": 8 + "max": 8, "axis": "tray" }, "span": { "max": 8 diff --git a/examples/k3-4l-ep4-tp4-prefill.json b/examples/k3-4l-ep4-tp4-prefill.json index 39932d55..65c8b87c 100644 --- a/examples/k3-4l-ep4-tp4-prefill.json +++ b/examples/k3-4l-ep4-tp4-prefill.json @@ -3,13 +3,13 @@ "model": "kimi-k3-pruned-75pct/4l/ep4-tp4-prefill", "vars": { "tokens": { - "max": 256 + "max": 256, "axis": "rows" }, "seqs": { - "max": 1 + "max": 1, "axis": "groups" }, "rows": { - "max": 256 + "max": 256, "axis": "tray" }, "ctx": { "max": 4096 diff --git a/examples/k3-4l-ep4-tp4.json b/examples/k3-4l-ep4-tp4.json index 935bb7fd..c9052fbc 100644 --- a/examples/k3-4l-ep4-tp4.json +++ b/examples/k3-4l-ep4-tp4.json @@ -3,13 +3,13 @@ "model": "kimi-k3-pruned-75pct/4l/ep4-tp4", "vars": { "tokens": { - "max": 8 + "max": 8, "axis": "rows" }, "seqs": { - "max": 8 + "max": 8, "axis": "groups" }, "rows": { - "max": 32 + "max": 32, "axis": "tray" }, "span": { "max": 8 diff --git a/examples/k3-4l-ep4.json b/examples/k3-4l-ep4.json index b944a1b0..c6dec452 100644 --- a/examples/k3-4l-ep4.json +++ b/examples/k3-4l-ep4.json @@ -3,13 +3,13 @@ "model": "kimi-k3-pruned-75pct/4l/ep4", "vars": { "tokens": { - "max": 8 + "max": 8, "axis": "rows" }, "seqs": { - "max": 8 + "max": 8, "axis": "groups" }, "rows": { - "max": 8 + "max": 8, "axis": "tray" }, "span": { "max": 8 diff --git a/examples/k3-4l-ep8.json b/examples/k3-4l-ep8.json index 4668516c..1ea243b5 100644 --- a/examples/k3-4l-ep8.json +++ b/examples/k3-4l-ep8.json @@ -3,13 +3,13 @@ "model": "kimi-k3-pruned-75pct/4l/ep8", "vars": { "tokens": { - "max": 8 + "max": 8, "axis": "rows" }, "seqs": { - "max": 8 + "max": 8, "axis": "groups" }, "rows": { - "max": 8 + "max": 8, "axis": "tray" }, "span": { "max": 8 diff --git a/examples/k3-ep4-tp4.json b/examples/k3-ep4-tp4.json index 4adda940..445a1ef2 100644 --- a/examples/k3-ep4-tp4.json +++ b/examples/k3-ep4-tp4.json @@ -3,13 +3,13 @@ "model": "kimi-k3-pruned-75pct/93l/ep4-tp4", "vars": { "tokens": { - "max": 64 + "max": 64, "axis": "rows" }, "seqs": { - "max": 64 + "max": 64, "axis": "groups" }, "rows": { - "max": 256 + "max": 256, "axis": "tray" }, "span": { "max": 64 diff --git a/examples/k3-ep4.json b/examples/k3-ep4.json index ee5f8c86..0fae16c0 100644 --- a/examples/k3-ep4.json +++ b/examples/k3-ep4.json @@ -3,13 +3,13 @@ "model": "kimi-k3-pruned-75pct/93l/ep4", "vars": { "tokens": { - "max": 64 + "max": 64, "axis": "rows" }, "seqs": { - "max": 64 + "max": 64, "axis": "groups" }, "rows": { - "max": 64 + "max": 64, "axis": "tray" }, "span": { "max": 64 diff --git a/examples/k3-ep8.json b/examples/k3-ep8.json index 7823e75a..e84aa7d3 100644 --- a/examples/k3-ep8.json +++ b/examples/k3-ep8.json @@ -3,13 +3,13 @@ "model": "kimi-k3-pruned-75pct/93l/ep8", "vars": { "tokens": { - "max": 64 + "max": 64, "axis": "rows" }, "seqs": { - "max": 64 + "max": 64, "axis": "groups" }, "rows": { - "max": 64 + "max": 64, "axis": "tray" }, "span": { "max": 64 diff --git a/examples/qwen3-4b-dspark.json b/examples/qwen3-4b-dspark.json index 78cb3ff9..6aa51c5f 100644 --- a/examples/qwen3-4b-dspark.json +++ b/examples/qwen3-4b-dspark.json @@ -69,10 +69,10 @@ }, "vars": { "tokens": { - "max": "MAX_TOKENS" + "max": "MAX_TOKENS", "axis": "rows" }, "seqs": { - "max": "MAX_SEQS" + "max": "MAX_SEQS", "axis": "groups" } }, "states": { diff --git a/examples/qwen3-4b-silu-mined.json b/examples/qwen3-4b-silu-mined.json index 48086007..9eec2107 100644 --- a/examples/qwen3-4b-silu-mined.json +++ b/examples/qwen3-4b-silu-mined.json @@ -62,10 +62,10 @@ }, "vars": { "tokens": { - "max": "MAX_TOKENS" + "max": "MAX_TOKENS", "axis": "rows" }, "seqs": { - "max": "MAX_SEQS" + "max": "MAX_SEQS", "axis": "groups" } }, "states": { diff --git a/examples/qwen3-4b.json b/examples/qwen3-4b.json index a58ab33c..39d017e6 100644 --- a/examples/qwen3-4b.json +++ b/examples/qwen3-4b.json @@ -62,10 +62,10 @@ }, "vars": { "tokens": { - "max": "MAX_TOKENS" + "max": "MAX_TOKENS", "axis": "rows" }, "seqs": { - "max": "MAX_SEQS" + "max": "MAX_SEQS", "axis": "groups" } }, "states": { diff --git a/examples/qwen3.8-27b-dflash2.json b/examples/qwen3.8-27b-dflash2.json index 3b26a6eb..a4dd0ccb 100644 --- a/examples/qwen3.8-27b-dflash2.json +++ b/examples/qwen3.8-27b-dflash2.json @@ -56,10 +56,10 @@ }, "vars": { "tokens": { - "max": "MAX_TOKENS" + "max": "MAX_TOKENS", "axis": "rows" }, "seqs": { - "max": "MAX_SEQS" + "max": "MAX_SEQS", "axis": "groups" } }, "states": { diff --git a/examples/qwen3.8-27b.json b/examples/qwen3.8-27b.json index bb56007b..31c56678 100644 --- a/examples/qwen3.8-27b.json +++ b/examples/qwen3.8-27b.json @@ -35,10 +35,10 @@ }, "vars": { "tokens": { - "max": "MAX_TOKENS" + "max": "MAX_TOKENS", "axis": "rows" }, "seqs": { - "max": "MAX_SEQS" + "max": "MAX_SEQS", "axis": "groups" } }, "states": { diff --git a/schema/manifest-v5.schema.json b/schema/manifest-v5.schema.json index d57a5348..77b3f039 100644 --- a/schema/manifest-v5.schema.json +++ b/schema/manifest-v5.schema.json @@ -2118,6 +2118,17 @@ "additionalProperties": false, "description": "A per-call scalar the caller supplies, bounded `1..=max`; the only kind of number that may size a shape or a grid.", "properties": { + "axis": { + "anyOf": [ + { + "$ref": "#/$defs/VarAxis" + }, + { + "type": "null" + } + ], + "description": "The axis of a call this var counts, for the serving loop: every buffer shaped over it has one entry per element of that axis, e.g. `\"rows\"`. At most one var per axis; a manifest with a `batch` names a `rows` var and a `groups` var. Absent for a var that sizes something else (a run's length, a context)." + }, "max": { "anyOf": [ { @@ -2140,6 +2151,26 @@ ], "type": "object" }, + "VarAxis": { + "description": "What a var counts in one call of a program: this rank's rows, its sequences, or the rows of the whole tray batch. The runtime never reads it.", + "oneOf": [ + { + "const": "rows", + "description": "One per row of this rank's call, e.g. `tokens`.", + "type": "string" + }, + { + "const": "groups", + "description": "One per sequence of this rank's call, e.g. `seqs`.", + "type": "string" + }, + { + "const": "tray", + "description": "One per row of the whole tray batch, this rank's rows first, e.g. a `tp` manifest's `rows`.", + "type": "string" + } + ] + }, "When": { "additionalProperties": false, "description": "The inclusive range of a var a launch runs in, e.g. `{\"var\": \"tokens\", \"min\": 17}`; an end left out is open.", diff --git a/tools/dsv41/auxiliary/serving.py b/tools/dsv41/auxiliary/serving.py index 6afc3e74..1eb70d91 100644 --- a/tools/dsv41/auxiliary/serving.py +++ b/tools/dsv41/auxiliary/serving.py @@ -248,7 +248,7 @@ def rope_name(compress_ratio, *, split=False): def build(cubin: Path, layout=Layout(), peers=None): """Return concrete definitions and lowering helpers (weights supplied outside).""" - variables = {"tokens": {"max": layout.max_tokens}, "seqs": {"max": layout.max_seqs}} + variables = {"tokens": {"max": layout.max_tokens, "axis": "rows"}, "seqs": {"max": layout.max_seqs, "axis": "groups"}} buffers = { "input_ids": {"kind": "input", "dtype": "i64", "shape": ["tokens"], "fill": "token"}, "anchor_token": {"kind": "input", "dtype": "i64", "shape": ["seqs"], "fill": "token"}, diff --git a/tools/gen_k3.py b/tools/gen_k3.py index 7913bb54..a699ea42 100755 --- a/tools/gen_k3.py +++ b/tools/gen_k3.py @@ -1077,7 +1077,8 @@ def emit(span, chunk=False): "schema_version": kern_manifest.SCHEMA_VERSION, "model": f"kimi-k3-pruned-75pct/{layers}l/ep{ranks}" + (f"-tp{tp}" if tp > 1 else "") + ("-prefill" if chunk else ""), - "vars": {T: {"max": t_max}, "seqs": {"max": seqs_max}, R: {"max": rows_max}, + "vars": {T: {"max": t_max, "axis": "rows"}, "seqs": {"max": seqs_max, "axis": "groups"}, + R: {"max": rows_max, "axis": "tray"}, **({SP: {"max": span_max}} if span_max else {}), **({CTX: {"max": max_ctx}} if chunk else {})}, "topology": {"groups": groups}, "states": states, diff --git a/tools/gen_qwen35.py b/tools/gen_qwen35.py index 74dfec3c..19c5568c 100644 --- a/tools/gen_qwen35.py +++ b/tools/gen_qwen35.py @@ -1517,7 +1517,7 @@ def pre(prefix, calls): return normalize(head | { # bs=1; prefill per chunk (tokens <= CHUNK_MAX) over *all* prompt # tokens (it emits next_token), decode at tokens=1 - "vars": {"tokens": {"max": CHUNK_MAX}, "seqs": {"max": MAX_SEQS}}, + "vars": {"tokens": {"max": CHUNK_MAX, "axis": "rows"}, "seqs": {"max": MAX_SEQS, "axis": "groups"}}, "states": states, "buffers": buffers, "ops": kernels, diff --git a/tools/gen_qwen3_decode.py b/tools/gen_qwen3_decode.py index f9b12f3f..006c3e4f 100644 --- a/tools/gen_qwen3_decode.py +++ b/tools/gen_qwen3_decode.py @@ -1058,7 +1058,7 @@ def pre(prefix, calls): return normalize({ "schema_version": SCHEMA_VERSION, "model": "qwen3-4b-dspark" if spec else "qwen3-4b", - "vars": {"tokens": {"max": CHUNK_MAX}, "seqs": {"max": MAX_SEQS}}, + "vars": {"tokens": {"max": CHUNK_MAX, "axis": "rows"}, "seqs": {"max": MAX_SEQS, "axis": "groups"}}, "states": states, "buffers": buffers, "ops": kernels, diff --git a/tools/toy/gen.py b/tools/toy/gen.py index 755b44e2..6ca443c9 100644 --- a/tools/toy/gen.py +++ b/tools/toy/gen.py @@ -119,7 +119,7 @@ def op(entry: str, grid: str, iface: list[str], extra: list[dict]) -> dict: return { "schema_version": 5, "model": name, - "vars": {"tokens": {"max": tokens_max}, "seqs": {"max": seqs_max}}, + "vars": {"tokens": {"max": tokens_max, "axis": "rows"}, "seqs": {"max": seqs_max, "axis": "groups"}}, "states": states, "buffers": buffers, "modules": {"toy": {"source": src, "sha256": sha}}, From c02b14d6cfae500260b42afa0d962779f4c15e8a Mon Sep 17 00:00:00 2001 From: JinYan Su <751080330@qq.com> Date: Thu, 8 Oct 2026 15:16:48 +0000 Subject: [PATCH 19/22] tools: the qwen example baselines follow the axis field The five hashed examples gained `axis` on their vars with the manifest change; the gate re-pins them. Co-Authored-By: Claude Fable 5.1 Signed-off-by: JinYan Su <751080330@qq.com> --- tools/test_qwen_constants.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/tools/test_qwen_constants.py b/tools/test_qwen_constants.py index 9ffe542e..d3904c70 100644 --- a/tools/test_qwen_constants.py +++ b/tools/test_qwen_constants.py @@ -10,14 +10,15 @@ ROOT = pathlib.Path(__file__).resolve().parent.parent # Canonical JSON hashes of the checked-in examples (last: unified attention # reads the kv state, `inout` -> `in`, 2026-09-16; qwen3-4b: one kv state per -# layer and registry builds of argmax / embedding / weight_prep, 2026-10-07). Changes to kernel wiring, +# layer and registry builds of argmax / embedding / weight_prep, 2026-10-07; all: vars +# declare their `axis`, 2026-10-08). Changes to kernel wiring, # capacities, float literals or even artifact pins require review. BASELINES = { - "qwen3-4b-dspark.json": "c72a03b5c0f817e742befcffa1e8241bc13f92c4f3458743a8ee27ccc6e4d30c", - "qwen3-4b-silu-mined.json": "3cd79ecd0d0d54c830f842644138b5dafb92c88efcf283709535b096665fe23e", - "qwen3-4b.json": "b54ac6fa2766667782626ec56b8d02ceb46b250eee7f3cf1a6d79561e5d58e84", - "qwen3.8-27b-dflash2.json": "134839dd1b25fc8fb5a5cb97391a02576d2624f2f0ccc5431af6e10c9dd679dd", - "qwen3.8-27b.json": "a9cc996a7548dabc756b30f9ef848459c2a5c2043075b79e11091054d1a30443" + "qwen3-4b-dspark.json": "40349cc8712c2efde92a04072ebe6444ec0718e90989e7d5702e21b07a5d0543", + "qwen3-4b-silu-mined.json": "014b6caf3b7588ad28c19a8b439357e1b83cf5e42011ba7c363416192d815f5a", + "qwen3-4b.json": "dac04a3a6f9793e10ee2d89e27d9252998c4b376d2d91cfb42eec1d10de73262", + "qwen3.8-27b-dflash2.json": "7500f3047dd5bd68a21efccbe04ba83a3f100b727722947cbe84cdf4a7ee632f", + "qwen3.8-27b.json": "0f404c48c06a310d894dfb6ab8e6600b4badaa65b0c55fdd14e4ef97d5603175" } From 9b88f6288e075fd12456352ac2489da0cf5d0213 Mon Sep 17 00:00:00 2001 From: JinYan Su <751080330@qq.com> Date: Thu, 8 Oct 2026 15:22:15 +0000 Subject: [PATCH 20/22] docs: PP3 is a boundary kind without edges, and comes before PP2 Wiring is deployment, not manifest; the kind carries only direction, dtype and shape so EP all-to-all can be its second user. Decode over the pipeline waits for it rather than building on `pp.*` ops. Co-Authored-By: Claude Fable 5.1 Signed-off-by: JinYan Su <751080330@qq.com> --- docs/roadmap.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/roadmap.md b/docs/roadmap.md index 1351dea1..40bc6376 100644 --- a/docs/roadmap.md +++ b/docs/roadmap.md @@ -64,8 +64,8 @@ kern 在内部 RSI 里是 kernel agent 的反馈环境（`kern test` 裁判、`k |---|---|---| | PP1 ✅ | `kern cut`（按 call 下标切、活跃性找边界、send/recv 邮箱核）+ kern-serve 一段一进程的 prefill-only 流水（head 持唯一的池，TCP 发 item，GPU 邮箱传激活） | **2026-10-07 tray07 过**：qwen3-4b 2 段，各段依次跑 logits/token 与整份逐位同（chunk 512/64）；两进程首 token 与单卡同；8193 token conc1 TTFT 313 → 174 ms（serve.md「流水线」） | | PP1.1 ✅ | qwen3-4b 的 KV 每层一个 state，`kern cut` 每段只带自己碰的 state、逐段报告 state 用了几成 | **2026-10-07 tray07 过**：`kern test qwen3-4b` 位一致、`kern run` 文本同；chain / serve gate 位一致；每段页数 122,993 → 245,760 | -| PP2 | decode 走流水；K3 的 state 拆分；per-seq state（GDN）与吐 token 的 chunk program | | -| PP3 | manifest 说得出"值从另一份 manifest 来"：跨切点的 buffer 是一种边界 kind（运行前从上游到达 / 运行后交给下游），边是 manifest 之间的有向关系，runtime 负责等、拷、通知；`kern cut` 只切 call 列表、标边界，不再生成 `pp.*` op，`pp_mailbox.cu` 退成 runtime 的一种传输（NVLink 邮箱 / RDMA / 同进程拷贝可换） | 切出的段里没有 `pp.*` op；chain / serve gate 位一致 | +| PP3 | manifest 说得出"值从另一份 manifest 来"：跨切点的 buffer 是一种边界 kind（`edge: in` 运行前已到达，`edge: out` 运行后送出），字段只有方向、dtype 与 shape，不认"第几段"；连到谁是部署的事，不进 manifest（`cut {id, stage, stages}` 只做握手核对）；runtime 负责等、拷、通知，`pp_mailbox.cu` 退成它的一种传输（NVLink 邮箱 / RDMA / 同进程拷贝可换）；`kern cut` 只切 call 列表、标边界，不再生成 `pp.*` op，不再认识任何 kernel。EP 的 all-to-all 是这个 kind 的第二个用户，字段名为它留门 | 切出的段里没有 `pp.*` op，kern-manifest 里没有 kernel 名；verify 拒绝本段写 `in` 边界、本段不写 `out` 边界；chain / serve gate 位一致 | +| PP2 | decode 走流水；K3 的 state 拆分；per-seq state（GDN）与吐 token 的 chunk program | 在 PP3 之后：decode 的流水用边界 kind，不用 `pp.*` op | | EP-PP | EP 也一卡一进程：head 持 N 个池、全连握手分发 handle 与 NcclId、每步给每个 rank 发它自己的 item、token 回 head；`Tray` 的单进程多 runtime 路径删除；manifest 不变 | K3 EP8 decode 跨两 tray 进 kern-serve，conc1 文本与单 tray 同；每 rank 的 `pp item` 时间线 | ## 协议线 From 3005ed9bbdac816d051a08161ce060732307a1c0 Mon Sep 17 00:00:00 2001 From: JinYan Su <751080330@qq.com> Date: Thu, 8 Oct 2026 15:51:46 +0000 Subject: [PATCH 21/22] kern-serve: a pipeline serves paged state only The slot zeroing and line tables an item carried for a per-sequence state had no model through the pipeline (qwen3 and K3 have none); they come back with the model that needs them and the gate that proves it. Co-Authored-By: Claude Fable 5.1 Signed-off-by: JinYan Su <751080330@qq.com> --- crates/kern-run/src/lib.rs | 4 ++-- crates/kern-runtime/src/lease.rs | 22 ---------------------- crates/kern-serve/src/pp/head.rs | 26 ++++---------------------- crates/kern-serve/src/pp/mod.rs | 6 +----- crates/kern-serve/src/pp/stage.rs | 10 ++++------ 5 files changed, 11 insertions(+), 57 deletions(-) diff --git a/crates/kern-run/src/lib.rs b/crates/kern-run/src/lib.rs index ff83bfbc..a53334de 100644 --- a/crates/kern-run/src/lib.rs +++ b/crates/kern-run/src/lib.rs @@ -226,11 +226,11 @@ fn page_rows(t: &PageTable, leases: &[Lease], rows: usize) -> Result> { /// line in entry 0 and zeros through the rest of a wide cell (a program /// that moves along one does so on the device). Columns past the leases /// given repeat the last. -pub fn line_rows>(t: &LineTable, leases: &[L], cols: usize) -> Result> { +fn line_rows(t: &LineTable, leases: &[Lease], cols: usize) -> Result> { let mut v = Vec::with_capacity(t.lines * cols * t.width); for line in 0..t.lines { for c in 0..cols { - v.push(leases[c.min(leases.len() - 1)].borrow().seq_line(&t.name, line)?); + v.push(leases[c.min(leases.len() - 1)].seq_line(&t.name, line)?); v.extend(std::iter::repeat_n(0, t.width - 1)); } } diff --git a/crates/kern-runtime/src/lease.rs b/crates/kern-runtime/src/lease.rs index c2e1981e..038fb14b 100644 --- a/crates/kern-runtime/src/lease.rs +++ b/crates/kern-runtime/src/lease.rs @@ -128,28 +128,6 @@ impl Runtime { Ok(lease) } - /// Zero sequence slot `slot` of every per-sequence state on the - /// stream, the way a lease zeroes its own: for a runtime whose slots - /// another runtime's pool hands out (a pipeline stage holds the head's - /// slot ids for its own layers' recurrent state). - pub fn zero_slot(&mut self, slot: i32) -> Result<()> { - self.ctx.bind_to_thread()?; - let live = self.pool.slots(); - if slot < 1 || slot as usize >= live { - bail!(Api, "slot {slot}: this runtime has slots 1..{live}"); - } - for (name, st) in &self.manifest.states { - if !st.is_per_seq() { - continue; - } - let b = st.bytes_per_seq as usize; - let s = slot as usize; - let mut view = self.states.get_mut(name).unwrap().view(s * b..(s + 1) * b)?; - self.stream.memset_zeros(&mut view)?; - } - Ok(()) - } - /// The first `len` tokens of `lease` as a [`Checkpoint`] its /// sequence keeps running past: pages shared, the per-sequence state /// copied into a fresh slot ([`Error::Denied`] when none is free). diff --git a/crates/kern-serve/src/pp/head.rs b/crates/kern-serve/src/pp/head.rs index 93aa9361..7b6c911a 100644 --- a/crates/kern-serve/src/pp/head.rs +++ b/crates/kern-serve/src/pp/head.rs @@ -22,7 +22,7 @@ use std::collections::{BTreeMap, VecDeque}; use std::sync::mpsc; use anyhow::{bail, ensure, Context, Result}; -use kern_manifest::protocol::{Axis, Forward, Rows}; +use kern_manifest::protocol::{Forward, Rows}; use kern_manifest::{Protocol, Verified}; use kern_pool::{Denied, Lease}; use kern_run::Weights; @@ -84,8 +84,6 @@ pub struct PpHead { active: VecDeque, /// Per item in flight, oldest first, the prompts its reply finishes. due: VecDeque>, - /// The line tables as the stages last had them written. - lines: BTreeMap>, watch: Watch, } @@ -178,7 +176,6 @@ pub fn load(m: Verified, weights: &Weights, h: Head) -> Result { waiting: VecDeque::new(), active: VecDeque::new(), due: VecDeque::new(), - lines: BTreeMap::new(), watch: Watch::new(0), }) } @@ -328,9 +325,7 @@ impl PpHead { } /// Stage one item of `parts` (a lease, its position, its rows) on every - /// stage and on the head. A sequence's first item zeroes its slot on the - /// stages (the head's own lease zeroed it already); a line table goes - /// out only when the item's sequences change it. + /// stage and on the head. fn send_item(&mut self, f: &Forward, parts: &[(&Lease, usize, &[i64])]) -> Result<()> { let leases: Vec<&Lease> = parts.iter().map(|&(l, ..)| l).collect(); let positions: Vec = parts.iter().map(|&(_, pos, _)| pos).collect(); @@ -345,22 +340,9 @@ impl PpHead { } rows.insert(t.name.clone(), v.into_iter().map(i64::from).collect()); } - for t in &self.p.line_tables { - let cols = match t.axis { - Axis::Tray => self.p.tray.as_ref().map_or(1, |b| b.max), - _ => self.p.groups.max, - }; - let v: Vec = kern_run::line_rows(t, &leases, cols as usize)?.into_iter().map(i64::from).collect(); - if self.lines.get(&t.name) != Some(&v) { - self.lines.insert(t.name.clone(), v.clone()); - rows.insert(t.name.clone(), v); - } - } - let zero_slots = leases.iter().zip(&positions).filter(|(_, &pos)| pos == 0).filter_map(|(l, _)| l.seq_slot()); - let item = Item { program: f.name.clone(), vars, rows, zero_slots: Vec::new() }; - let out = Item { zero_slots: zero_slots.collect(), ..item.clone() }; + let item = Item { program: f.name.clone(), vars, rows }; for w in &mut self.stages { - self.io.block_on(send(w, &out))?; + self.io.block_on(send(w, &item))?; } for (name, v) in &item.rows { kern_run::write_named(&mut self.rt, &self.p, name, v, &item.vars)?; diff --git a/crates/kern-serve/src/pp/mod.rs b/crates/kern-serve/src/pp/mod.rs index ac58ca94..10ada697 100644 --- a/crates/kern-serve/src/pp/mod.rs +++ b/crates/kern-serve/src/pp/mod.rs @@ -71,12 +71,8 @@ pub struct Table { pub struct Item { pub program: String, pub vars: BTreeMap, - /// Each input's values by name: the fills, the page tables and the - /// line tables when the item's sequences change them. + /// Each input's values by name: the fills and the page tables. pub rows: BTreeMap>, - /// The sequence slots to zero before the call, the ones of sequences - /// on their first item: the recurrent state a fresh lease starts from. - pub zero_slots: Vec, } pub type Handles = BTreeMap; diff --git a/crates/kern-serve/src/pp/stage.rs b/crates/kern-serve/src/pp/stage.rs index 43f6c52c..3e17c8d9 100644 --- a/crates/kern-serve/src/pp/stage.rs +++ b/crates/kern-serve/src/pp/stage.rs @@ -1,9 +1,10 @@ //! A downstream stage: loads its slice of the model, says hello to the //! head, maps the next stage's mailbox from the table it gets back, then //! replays the head's items in order. It keeps no accounting: the page -//! ids, positions and slots of every item are the head's, and so is the -//! zeroing of a fresh sequence's slot. The last stage answers every item -//! with its tokens. +//! ids and positions of every item are the head's. A stage serves paged +//! state only; a per-sequence state (its slot zeroed on a sequence's +//! first item) waits for a model that needs it. The last stage answers +//! every item with its tokens. use std::sync::mpsc; use std::time::{Duration, Instant}; @@ -97,9 +98,6 @@ pub fn run(m: Verified, weights: &Weights, o: Stage) -> Result<()> { let mut watch = Watch::new(cut.stage); for item in items { let vars: Vars = item.vars.into_iter().filter(|(k, _)| rt.manifest.vars.contains_key(k)).collect(); - for &slot in &item.zero_slots { - rt.zero_slot(slot)?; - } for (name, v) in &item.rows { kern_run::write_named(&mut rt, &p, name, v, &vars)?; } From 31d55c0e136239cb09f379e38789d9dda0b963f0 Mon Sep 17 00:00:00 2001 From: JinYan Su <751080330@qq.com> Date: Fri, 9 Oct 2026 02:29:56 +0000 Subject: [PATCH 22/22] pp: a cut stage is a function; the shell moves what crosses A buffer crossing a cut is an output on the stage before and an input or an inout (a new general buffer kind: an input the program rewrites in place) on the stage after, exported. `kern cut` only changes signatures; the manifest carries nothing pipeline-specific: the mailbox kernel, the six `pp.*` ops and the per-edge peer group are gone. Transport belongs to the serving shell: the next stage's arrivals are mapped into this one, `Runtime::push` copies on the compute stream, `Runtime::mark` says when it landed, and item numbers go over TCP. The next stage's host does not wait for the data: it issues item k behind `Runtime::hold`, a kernel spinning on a pinned count that `Gate::open` raises when "k arrived" comes in. Every signal the waiting thread handles stays off the driver: marks are a kernel writing a sequence number into pinned memory, outputs come back through `Runtime::fetch` into pooled pinned buffers, faults are the issuing thread's `Runtime::fault`. The four deadlocks that forced each of these are in docs/lessons.md. Gate (2026-10-08): qwen3-4b pp2 chain and serve first tokens identical to the unsplit manifest, faults as before; K3 PP8 c16 44.3k tok/s (mailbox 44.9-46.8k); qwen3 conc1 TTFT unchanged, conc4 5-13 ms behind the mailbox (docs/serve.md). Co-Authored-By: Claude Fable 5.1 Signed-off-by: JinYan Su <751080330@qq.com> --- Cargo.lock | 3 - Cargo.toml | 3 - crates/kern-manifest/src/cut.rs | 303 +++++++---------------- crates/kern-manifest/src/types.rs | 5 +- crates/kern-manifest/src/verify.rs | 12 +- crates/kern-manifest/tests/cut.rs | 98 +++----- crates/kern-run/examples/pp_chain.rs | 81 +++--- crates/kern-run/src/bin/kern.rs | 27 +- crates/kern-run/src/lib.rs | 18 +- crates/kern-runtime/src/device.rs | 30 +++ crates/kern-runtime/src/exec.rs | 263 ++++++++++++++++++-- crates/kern-runtime/src/lib.rs | 7 +- crates/kern-runtime/src/load.rs | 3 +- crates/kern-runtime/src/peers.rs | 18 ++ crates/kern-serve/Cargo.toml | 5 +- crates/kern-serve/src/pp/head.rs | 74 +++--- crates/kern-serve/src/pp/mod.rs | 358 +++++++++++++++++++-------- crates/kern-serve/src/pp/stage.rs | 101 ++++---- docs/lessons.md | 36 +++ docs/manifest.md | 51 ++-- docs/roadmap.md | 2 +- docs/serve.md | 35 ++- schema/manifest-v5.schema.json | 7 +- tools/kernels-src/pp_mailbox.cu | 118 --------- 24 files changed, 940 insertions(+), 718 deletions(-) delete mode 100644 tools/kernels-src/pp_mailbox.cu diff --git a/Cargo.lock b/Cargo.lock index 859ca407..88a8caac 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2292,9 +2292,7 @@ name = "kern-serve" version = "0.2.3" dependencies = [ "anyhow", - "bytes", "clap", - "futures-util", "hex", "kern-manifest", "kern-pool", @@ -2304,7 +2302,6 @@ dependencies = [ "serde", "serde_json", "tokio", - "tokio-util", "tracing", "tracing-subscriber", "vllm-chat", diff --git a/Cargo.toml b/Cargo.toml index ab21a18e..098d40f7 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -37,7 +37,6 @@ pegainfer-frontend = { git = "https://github.com/pegainfer-project/pegainfer", r vllm-chat = { git = "https://github.com/vllm-project/vllm.git", rev = "89dbb2644552d6e473a7e97da0ce8f0aa8e32c9d" } anyhow = "1" -bytes = "1" # Recorded references are ragged tables (a row per token position, list # columns for the top-k), which parquet holds and pandas / duckdb read as is. arrow-array = "=60.0.0" @@ -47,7 +46,6 @@ clap = { version = "4", features = ["derive"] } # The CUDA API kern binds is a fact about the binary, not about the build # machine: `cuda-13000` pins it, and no toolkit is needed to build. cudarc = { version = "=0.19.7", features = ["cuda-13000", "f16", "nccl"] } -futures-util = { version = "0.3", default-features = false, features = ["sink"] } half = "=2.7.1" hex = "0.4" libc = "0.2" @@ -60,7 +58,6 @@ sha2 = "0.11" thiserror = "2" tokenizers = "=0.23.2" tokio = { version = "1", features = ["rt-multi-thread", "macros"] } -tokio-util = { version = "0.7", features = ["codec"] } toml = "1.1" tracing = "0.1" tracing-subscriber = { version = "0.3", features = ["env-filter"] } diff --git a/crates/kern-manifest/src/cut.rs b/crates/kern-manifest/src/cut.rs index 8a754457..7b2404c2 100644 --- a/crates/kern-manifest/src/cut.rs +++ b/crates/kern-manifest/src/cut.rs @@ -1,23 +1,28 @@ //! Cutting a manifest into a pipeline: stage `s` runs the calls between -//! two cut points of every program a serving loop drives, and the -//! activations live across a cut travel through a mailbox. +//! two cut points of every program a serving loop drives, and what it +//! needs from the stage before it, or leaves for the stage after it, is +//! an input or an output of its own like any other. //! //! A forward is a pure function of its inputs and its states, so a cut //! changes no arithmetic: the stages run the same calls in the same order //! on the same values, and the pipeline's outputs are the unsplit -//! manifest's byte for byte. What a cut adds is transport. At a cut point -//! the buffers live across it are the ones the calls after it read before -//! they write them whole (an `out` at offset 0 is taken to write the whole -//! buffer, as every generator writes a view at its base) and that the -//! calls before it wrote; each must be a workspace, since an output is read -//! from the last stage and a carry would cross programs. The upstream stage -//! ends every program by putting them into a box on the downstream GPU -//! (`tools/kernels-src/pp_mailbox.cu`): wait until the box is empty, put, -//! post. The downstream stage starts by waiting for the post, taking them -//! out and freeing the box. One box per edge, laid out per program at the -//! bounds of its vars; the box sits on the downstream GPU and the upstream -//! reaches it through a `peer` buffer of a two-member topology group -//! `pp.`, its own copy of the box unused. +//! manifest's byte for byte. What a cut adds is a signature per stage. At +//! a cut point the buffers live across it are the ones the calls after it +//! read before they write them whole (an `out` at offset 0 is taken to +//! write the whole buffer, as every generator writes a view at its base) +//! and that the calls before it wrote; each must be a workspace, since an +//! output is read from the last stage and a carry would cross programs. +//! On the stage before the cut such a buffer is an `output`; on the stage +//! after it an `input` when that stage only reads it (on its way to a +//! later stage or not), an `inout` when the stage rewrites it in place. +//! A buffer that arrives is exported, so the stage before can map it and +//! copy into it. How the value travels (a memcpy over NVLink, RDMA, a +//! copy within a process) and when (the stage before has finished the +//! item, the stage after is done with the one before it) is the serving +//! shell's, which wires the stages by name: every arriving buffer of a +//! stage is one its predecessor holds under the same name. A value a +//! stage neither reads nor writes on its way to a later stage is +//! refused: a stage carries only what it uses. //! //! A state never travels. Every region of a state (a state and a call's //! offset into it) belongs to the one stage whose calls touch it, which is @@ -30,31 +35,21 @@ //! //! The prompt reaches its first token the way it does unsplit, from the //! step or from a chunk program that hands one back; either way the last -//! stage writes it, since only the last stage writes an output. Everything a stage -//! does not use is dropped, and every stage verifies. A state is kept whole -//! or not at all, so a state laid out over the whole model (one KV state -//! for every layer) is carried whole by every stage that touches a piece of -//! it; [`state_use`] says how much of each one a stage uses, and a -//! generator that gives each piece its own state (one per layer) lets a -//! stage carry only its own. A table whose state a stage dropped indexes a -//! kept state of the same kind: a pool hands every paged state the same ids. +//! stage writes it, since only the last stage writes an output. Everything +//! a stage does not use is dropped, and every stage verifies. A state is +//! kept whole or not at all, so a state laid out over the whole model (one +//! KV state for every layer) is carried whole by every stage that touches +//! a piece of it; [`state_use`] says how much of each one a stage uses, +//! and a generator that gives each piece its own state (one per layer) +//! lets a stage carry only its own. A table whose state a stage dropped +//! indexes a kept state of the same kind: a pool hands every paged state +//! the same ids. use std::collections::{BTreeMap, BTreeSet}; -use serde_json::json; - use crate::types::*; use crate::verify::{verify, Verified}; -/// How long a mailbox wait spins before the stage gives up and traps. -pub const TIMEOUT_NS: i64 = 60_000_000_000; -/// Bytes of a box before its data: the two counters, a cache line each. -const HEADER: u64 = 256; -/// Each carried buffer starts at a multiple of this in the box. -const ALIGN: u64 = 256; -/// The module name the pipeline kernels go under in every stage. -const MODULE: &str = "pp"; - /// Every reason a manifest cannot be cut where asked, reported together. #[derive(Debug)] pub struct CutErrors(pub Vec); @@ -71,21 +66,10 @@ impl std::fmt::Display for CutErrors { impl std::error::Error for CutErrors {} -/// A buffer an edge carries in one program: where in the box, and how -/// many bytes per call. -#[derive(Debug, Clone)] -struct Carried { - buffer: String, - dtype: DType, - offset: u64, - bytes: Expr, -} - /// Cut `m` before each call index in `at` of every program with a -/// `batch` (all of them must be longer than the last point), wiring in -/// the mailbox kernels of `pp`. Returns the stages in order, each -/// verified. -pub fn cut(m: &Verified, at: &[usize], pp: &Module) -> Result, CutErrors> { +/// `batch` (all of them must be longer than the last point). Returns the +/// stages in order, each verified. +pub fn cut(m: &Verified, at: &[usize]) -> Result, CutErrors> { let mut errs = Vec::new(); if m.topology.is_some() { errs.push("the manifest has a topology; a cut takes a single-rank manifest".to_string()); @@ -114,37 +98,21 @@ pub fn cut(m: &Verified, at: &[usize], pp: &Module) -> Result, Cut let stages = at.len() + 1; let bounds: Vec = std::iter::once(0).chain(at.iter().copied()).collect(); let span = |p: &Program, s: usize| bounds[s]..bounds.get(s + 1).copied().unwrap_or(p.calls.len()); - let max: BTreeMap = m.vars.iter().map(|(k, v)| (k.clone(), v.max)).collect(); - // What each edge carries in each program, and each box's size. - let mut carried: BTreeMap<(&str, usize), Vec> = BTreeMap::new(); + // What crosses each edge: live after it, written before it. + let mut crossing: Vec> = vec![BTreeSet::new(); at.len()]; for (name, p) in &driven { for (e, &a) in at.iter().enumerate() { let before = writes(m, &p.calls[..a]); - let mut offset = 0; - let mut list = Vec::new(); for b in live_in(m, &p.calls[a..]).into_iter().filter(|b| before.contains(b)) { - let buf = &m.buffers[b]; - if buf.kind != BufferKind::Workspace { + if m.buffers[b].kind != BufferKind::Workspace { errs.push(format!( "program `{name}`: {} buffer `{b}` is live across the cut before call {a}; only a workspace travels", - buf.kind + m.buffers[b].kind )); - continue; - } - match bytes_of(buf) { - Some(bytes) => { - let at_max = bytes.eval(&max).unwrap_or(u64::MAX); - list.push(Carried { buffer: b.to_string(), dtype: buf.dtype, offset, bytes }); - offset += at_max.div_ceil(ALIGN) * ALIGN; - } - None => errs.push(format!( - "program `{name}`: buffer `{b}` is shaped {:?}; a carried buffer has at most one var dim", - buf.shape - )), } + crossing[e].insert(b); } - carried.insert((name.as_str(), e), list); } for (i, c) in p.calls.iter().enumerate().take(bounds[stages - 1]) { for (b, _, dir) in accesses(m, c) { @@ -157,6 +125,33 @@ pub fn cut(m: &Verified, at: &[usize], pp: &Module) -> Result, Cut } } } + // Each stage's signature: what arrives, what it writes, what leaves. + let mut kinds: Vec> = Vec::with_capacity(stages); + for s in 0..stages { + let arriving = if s > 0 { crossing[s - 1].clone() } else { BTreeSet::new() }; + let leaving = crossing.get(s).cloned().unwrap_or_default(); + let (mut touched, mut written) = (BTreeSet::new(), BTreeSet::new()); + for (_, p) in &driven { + let calls = &p.calls[span(p, s)]; + touched.extend(calls.iter().flat_map(|c| accesses(m, c)).map(|(b, ..)| b)); + written.extend(writes(m, calls)); + } + let mut k = BTreeMap::new(); + for &b in arriving.union(&leaving) { + k.insert( + b, + match (arriving.contains(b), written.contains(b)) { + (true, true) => BufferKind::Inout, + (true, false) => BufferKind::Input, + (false, _) => BufferKind::Output, + }, + ); + if !touched.contains(b) { + errs.push(format!("stage {s}: `{b}` passes through untouched; a stage carries only what it uses")); + } + } + kinds.push(k); + } let mut owner: BTreeMap<(String, u64), usize> = BTreeMap::new(); for (name, p) in &driven { for s in 0..stages { @@ -175,124 +170,30 @@ pub fn cut(m: &Verified, at: &[usize], pp: &Module) -> Result, Cut if !errs.is_empty() { return Err(CutErrors(errs)); } - let boxes: Vec = (0..at.len()) - .map(|e| { - let data = carried - .iter() - .filter(|((_, ee), _)| *ee == e) - .map(|(_, list)| list.last().map_or(0, |c| c.offset + c.bytes.eval(&max).unwrap_or(0))) - .max() - .unwrap_or(0); - HEADER + data.max(1).div_ceil(ALIGN) * ALIGN - }) - .collect(); let id = id_of(m, at); - let built: Vec = (0..stages) + (0..stages) .map(|s| { - let programs: BTreeMap = driven + let programs = driven .iter() - .map(|(name, p)| { - let into = (s > 0).then(|| &carried[&(name.as_str(), s - 1)][..]); - let out = (s + 1 < stages).then(|| &carried[&(name.as_str(), s)][..]); - let calls = - receive(s, into).into_iter().chain(p.calls[span(p, s)].iter().cloned()).chain(send(s, out)); - ((*name).clone(), Program { calls: calls.collect(), ..(*p).clone() }) - }) + .map(|(name, p)| ((*name).clone(), Program { calls: p.calls[span(p, s)].to_vec(), ..(*p).clone() })) .collect(); - stage(m, programs, s, stages, &boxes, &id, pp) + let st = stage(m, programs, &kinds[s], s, stages, &id); + verify(st).map_err(|e| CutErrors(e.0.into_iter().map(|d| format!("stage {s}: {d}")).collect())) }) - .collect(); - built - .into_iter() - .enumerate() - .map(|(s, st)| verify(st).map_err(|e| CutErrors(e.0.into_iter().map(|d| format!("stage {s}: {d}")).collect()))) - .collect() -} - -fn box_name(e: usize) -> String { - format!("pp.{e}.box") -} - -fn peer_name(e: usize) -> String { - format!("pp.{e}.peer") -} - -fn call(op: &str, args: serde_json::Value) -> Call { - serde_json::from_value(json!({"label": op, "op": op, "args": args})).expect("a well-formed call") -} - -/// The calls that start stage `s`: take what the edge into it carries -/// (none on a first stage). -fn receive(s: usize, into: Option<&[Carried]>) -> Vec { - let Some(list) = into else { return Vec::new() }; - let b = box_name(s - 1); - std::iter::once(call("pp.wait_full", json!([{"buf": b}]))) - .chain(list.iter().map(|c| { - call( - &format!("pp.take.{}", c.dtype), - json!([{"buf": b}, {"buf": c.buffer}, {"expr": c.bytes}, {"i64": c.offset}]), - ) - })) - .chain(std::iter::once(call("pp.post_empty", json!([{"buf": b}])))) - .collect() -} - -/// The calls that end stage `s`: put what the edge out of it carries -/// (none on a last stage). -fn send(s: usize, out: Option<&[Carried]>) -> Vec { - let Some(list) = out else { return Vec::new() }; - let p = peer_name(s); - std::iter::once(call("pp.wait_empty", json!([{"buf": p}]))) - .chain(list.iter().map(|c| { - call( - &format!("pp.put.{}", c.dtype), - json!([{"buf": p}, {"buf": c.buffer}, {"expr": c.bytes}, {"i64": c.offset}]), - ) - })) - .chain(std::iter::once(call("pp.post_full", json!([{"buf": p}])))) .collect() } -/// The pipeline ops a stage's calls name: `pp.put.` and -/// `pp.take.` per carried dtype, the waits and the posts. -fn pp_op(name: &str) -> Op { - let single = |entry: &str, params: serde_json::Value| { - json!({"params": params, "impl": {"launches": [ - {"module": MODULE, "entry": entry, "block": [1, 1, 1], "grid": [1, 1, 1]}]}}) - }; - let wait = |entry: &str, first: &str| { - json!({"params": [first], "impl": {"launches": [ - {"module": MODULE, "entry": entry, "block": [1, 1, 1], "grid": [1, 1, 1], - "params": [first, "i64"], "args": [{"param": 0}, {"i64": TIMEOUT_NS}]}]}}) - }; - let copy = |entry: &str, params: serde_json::Value| { - json!({"params": params, "impl": {"launches": [ - {"module": MODULE, "entry": entry, "block": [256, 1, 1], "grid": [264, 1, 1]}]}}) - }; - let v = match name.split('.').collect::>().as_slice() { - ["pp", "wait_empty"] => wait("kern_pp_wait_empty", "in buffer"), - ["pp", "wait_full"] => wait("kern_pp_wait_full", "in buffer"), - ["pp", "post_full"] => single("kern_pp_post_full", json!(["in buffer"])), - ["pp", "post_empty"] => single("kern_pp_post_empty", json!(["inout buffer"])), - ["pp", "put", dt] => copy("kern_pp_put", json!(["in buffer", format!("in buffer<{dt}>"), "i64", "i64"])), - ["pp", "take", dt] => copy("kern_pp_take", json!(["in buffer", format!("out buffer<{dt}>"), "i64", "i64"])), - _ => unreachable!("the cut names only the pipeline ops above"), - }; - serde_json::from_value(v).expect("a well-formed pipeline op") -} - -/// Stage `s` of `stages`: `programs` (already cut and wired), the once -/// and derive programs sliced to what they need, and every declaration -/// they use. +/// Stage `s` of `stages`: `programs` (already cut), the once and derive +/// programs sliced to what they need, and every declaration they use, +/// the buffers crossing its edges re-declared as `kinds` says. fn stage( m: &Manifest, mut programs: BTreeMap, + kinds: &BTreeMap<&str, BufferKind>, s: usize, stages: usize, - boxes: &[u64], id: &str, - pp: &Module, ) -> Manifest { let mut needed: BTreeSet = programs.values().flat_map(|p| &p.calls).flat_map(|c| reads(m, c)).map(str::to_string).collect(); @@ -310,26 +211,8 @@ fn stage( } } - let mut buffers: BTreeMap = BTreeMap::new(); - let u8_box = |bytes: u64, kind: &str| -> Buffer { - serde_json::from_value(json!({"dtype": "u8", "shape": [bytes], "kind": kind, "export": true})).expect("a box") - }; - if s > 0 { - buffers.insert(box_name(s - 1), u8_box(boxes[s - 1], "carry")); - } - let mut topology = None; - if s + 1 < stages { - buffers.insert(box_name(s), u8_box(boxes[s], "workspace")); - buffers.insert( - peer_name(s), - serde_json::from_value( - json!({"dtype": "u64", "shape": [2], "kind": "peer", "of": box_name(s), "group": format!("pp.{s}")}), - ) - .expect("a peer"), - ); - topology = Some(Topology { groups: BTreeMap::from([(format!("pp.{s}"), 2)]) }); - } let calls: Vec<&Call> = programs.values().flat_map(|p| &p.calls).collect(); + let mut buffers: BTreeMap = BTreeMap::new(); for c in &calls { for a in &c.args { if let Arg::Buf { buf, .. } = a { @@ -339,6 +222,11 @@ fn stage( } } } + for (name, &kind) in kinds { + let b = buffers.get_mut(*name).expect("a crossing buffer is one the stage's calls touch"); + b.kind = kind; + b.export = kind != BufferKind::Output; + } let states: BTreeMap = calls.iter().flat_map(|c| state_args(c)).map(|(st, _)| (st.to_string(), m.states[st].clone())).collect(); let names: BTreeSet = buffers.keys().chain(states.keys()).cloned().collect(); @@ -355,16 +243,13 @@ fn stage( _ => b.domain = None, } } - let ops: BTreeMap = - calls.iter().map(|c| (c.op.clone(), m.ops.get(&c.op).cloned().unwrap_or_else(|| pp_op(&c.op)))).collect(); - let mut modules: BTreeMap = ops + let ops: BTreeMap = calls.iter().map(|c| (c.op.clone(), m.ops[&c.op].clone())).collect(); + let modules: BTreeMap = ops .values() .flat_map(|o| &o.imp.launches) .filter_map(Launch::module) - .filter(|n| *n != MODULE) .map(|n| (n.to_string(), m.modules[n].clone())) .collect(); - modules.insert(MODULE.to_string(), pp.clone()); // A var stays when some kept declaration names it: shapes, grids, // call args, batches and domains all spell a var as a bare string. let mut strings = BTreeSet::new(); @@ -375,7 +260,7 @@ fn stage( Manifest { schema_version: m.schema_version, model: m.model.clone(), - topology, + topology: None, cut: Some(Cut { id: id.to_string(), stage: s as u64, stages: stages as u64 }), vars, states, @@ -397,19 +282,6 @@ fn collect_strings(v: &serde_json::Value, out: &mut BTreeSet) { } } -/// Bytes one call moves of `b`: its size at the call's vars, when at most -/// one dim is a var. -fn bytes_of(b: &Buffer) -> Option { - let constant: u64 = b.shape.iter().filter_map(|d| if let Dim::Const(c) = d { Some(*c) } else { None }).product(); - let vars: Vec<&String> = b.shape.iter().filter_map(|d| if let Dim::Var(v) = d { Some(v) } else { None }).collect(); - let c = constant * b.dtype.bytes(); - match vars.as_slice() { - [] => Some(Expr::Const(c)), - [v] => Some(Expr::Mul { mul: (Box::new(Expr::Var((*v).clone())), c) }), - _ => None, - } -} - /// Every buffer argument of `c`: name, offset and the direction its op /// declares. fn accesses<'a>(m: &'a Manifest, c: &'a Call) -> Vec<(&'a str, u64, Dir)> { @@ -424,16 +296,9 @@ fn accesses<'a>(m: &'a Manifest, c: &'a Call) -> Vec<(&'a str, u64, Dir)> { .collect() } -/// The buffers `c` reads, by name; a pipeline call reads its buffers too. +/// The buffers `c` reads, by name. fn reads<'a>(m: &'a Manifest, c: &'a Call) -> impl Iterator { - let known = m.ops.contains_key(&c.op); - let acc = accesses(m, c); - let pipeline: Vec<&str> = if known { - Vec::new() - } else { - c.args.iter().filter_map(|a| if let Arg::Buf { buf, .. } = a { Some(buf.as_str()) } else { None }).collect() - }; - acc.into_iter().filter(|(_, _, d)| *d != Dir::Out).map(|(b, _, _)| b).chain(pipeline) + accesses(m, c).into_iter().filter(|(_, _, d)| *d != Dir::Out).map(|(b, _, _)| b) } /// The buffers `calls` write. diff --git a/crates/kern-manifest/src/types.rs b/crates/kern-manifest/src/types.rs index 9c415c47..4aece24c 100644 --- a/crates/kern-manifest/src/types.rs +++ b/crates/kern-manifest/src/types.rs @@ -71,7 +71,7 @@ pub struct Manifest { pub schema_version: u32, /// Free-form model label, e.g. `"qwen3-4b"`. pub model: String, - /// Rank groups of a multi-GPU manifest, e.g. `{"groups": {"ep": 4}}`: every rank of an SPMD group loads the same manifest; the two stages a cut's group joins load their own, agreeing only on the peer's `of` buffer. + /// Rank groups a multi-GPU manifest is SPMD over, e.g. `{"groups": {"ep": 4}}`; every rank loads the same manifest. #[serde(default, skip_serializing_if = "Option::is_none")] pub topology: Option, /// The pipeline stage this manifest is, when it is one piece of a cut manifest, e.g. `{"id": "9f2c…", "stage": 1, "stages": 2}`. @@ -620,6 +620,8 @@ pub enum BufferKind { Input, /// Read back by the runtime after each run, e.g. `next_token`. Output, + /// Both: holds a value before the run, which the programs may rewrite in place, and is read back after, e.g. the `hidden` a pipeline stage takes from the stage before it and hands to the next. + Inout, /// 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. @@ -637,6 +639,7 @@ impl fmt::Display for BufferKind { f.write_str(match self { BufferKind::Input => "input", BufferKind::Output => "output", + BufferKind::Inout => "inout", BufferKind::Weight => "weight", BufferKind::Source => "source", BufferKind::Workspace => "workspace", diff --git a/crates/kern-manifest/src/verify.rs b/crates/kern-manifest/src/verify.rs index 2c1b2e4e..0a4928b2 100644 --- a/crates/kern-manifest/src/verify.rs +++ b/crates/kern-manifest/src/verify.rs @@ -26,7 +26,8 @@ //! against the interface, var ranges fit scalar params //! 8. dataflow per program: no read-before-write, no writes to input, //! weight, source or peer buffers, every output / carry buffer written -//! by some program; a program is `once` or has a `batch`, not both; a +//! by some program (an inout holds a value before the run and need +//! not be); 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 @@ -784,7 +785,12 @@ fn diagnostics(m: &Manifest) -> Vec { // handles are imported, before any program runs. matches!( b.kind, - BufferKind::Input | BufferKind::Weight | BufferKind::Source | BufferKind::Carry | BufferKind::Peer + BufferKind::Input + | BufferKind::Inout + | BufferKind::Weight + | BufferKind::Source + | BufferKind::Carry + | BufferKind::Peer ) }) .map(|(n, _)| n.as_str()) @@ -907,7 +913,7 @@ fn diagnostics(m: &Manifest) -> Vec { (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!( + (true, BufferKind::Input | BufferKind::Output | BufferKind::Inout | BufferKind::Peer) => errs.push(format!( "{actx}: a `derive` program runs inside weight loading, before any {} buffer such as `{buf}` means anything", b.kind )), diff --git a/crates/kern-manifest/tests/cut.rs b/crates/kern-manifest/tests/cut.rs index 51acd8e4..b75a7aa8 100644 --- a/crates/kern-manifest/tests/cut.rs +++ b/crates/kern-manifest/tests/cut.rs @@ -1,19 +1,17 @@ -//! A cut changes transport, never arithmetic: run symbolically, the stages -//! of any cut that verifies compute what the whole manifest computes. +//! A cut changes a stage's signature, never arithmetic: run symbolically, +//! the stages of any cut that verifies, each fed what the one before it +//! holds under the names it takes in, compute what the whole manifest +//! computes. use std::collections::BTreeMap; use kern_manifest::cut::{cut, state_use}; -use kern_manifest::types::{Arg, Dir, Manifest, Module, ParamType}; +use kern_manifest::types::{Arg, BufferKind, Dir, Manifest, ParamType}; use kern_manifest::{verify, Protocol, Verified}; use serde_json::{json, Value}; const QWEN3: &str = include_str!("../../../examples/qwen3-4b.json"); -fn pp() -> Module { - Module { source: "pp_mailbox.cubin".into(), sha256: "ab".repeat(32) } -} - struct Rand(u64); impl Rand { @@ -123,36 +121,12 @@ struct Machine { states: BTreeMap<(String, u64), u64>, } -/// The edge a pipeline buffer belongs to: `pp..box` / `pp..peer`. -fn edge(name: &str) -> u64 { - name.split('.').nth(1).unwrap().parse().unwrap() -} - impl Machine { /// Run `program` of `m`: every op a hash of what it reads into what it - /// writes; the pipeline ops move values through `mail`. - fn run(&mut self, m: &Manifest, program: &str, mail: &mut BTreeMap<(u64, i64), u64>) { + /// writes. + fn run(&mut self, m: &Manifest, program: &str) { for c in &m.programs[program].calls { let label = c.label.clone().unwrap(); - let buf = |a: &Arg| match a { - Arg::Buf { buf, .. } => buf.clone(), - _ => unreachable!(), - }; - let off = |a: &Arg| match a { - Arg::I64 { i64: v } => *v, - _ => unreachable!(), - }; - if c.op.starts_with("pp.put.") { - mail.insert((edge(&buf(&c.args[0])), off(&c.args[3])), self.bufs[&buf(&c.args[1])]); - continue; - } - if c.op.starts_with("pp.take.") { - self.bufs.insert(buf(&c.args[1]), mail[&(edge(&buf(&c.args[0])), off(&c.args[3]))]); - continue; - } - if c.op.starts_with("pp.") { - continue; - } let params = &m.ops[&c.op].params; let mut read = vec![name_hash(&c.op), name_hash(&label)]; for (a, p) in c.args.iter().zip(params) { @@ -186,7 +160,7 @@ impl Machine { let mut me = Machine::default(); me.bufs.insert("x".into(), name_hash("x")); if m.programs.contains_key("init") { - me.run(m, "init", &mut BTreeMap::new()); + me.run(m, "init"); } me } @@ -203,29 +177,41 @@ fn every_cut_that_verifies_computes_the_whole() { let at = if a + 1 < n && r.next(2) == 0 { vec![a, a + 1 + r.next(n - a - 1)] } else { vec![a] }; let mut whole = Machine::start(&m); - whole.run(&m, "p", &mut BTreeMap::new()); + whole.run(&m, "p"); - let stages = match cut(&m, &at, &pp()) { + let stages = match cut(&m, &at) { Ok(s) => s, Err(e) => { - assert!(e.0.iter().all(|d| d.contains("touched by stages")), "seed {seed} at {at:?}: {e}"); + let known = |d: &String| d.contains("touched by stages") || d.contains("passes through untouched"); + assert!(e.0.iter().all(known), "seed {seed} at {at:?}: {e}"); refused += 1; continue; } }; cuts += 1; assert_eq!(stages.len(), at.len() + 1); - let mut mail = BTreeMap::new(); + let mut held: BTreeMap = BTreeMap::new(); let mut states = BTreeMap::new(); let mut y = None; for (s, st) in stages.iter().enumerate() { assert_eq!(st.cut.as_ref().map(|c| (c.stage, c.stages)), Some((s as u64, stages.len() as u64))); + assert!(st.topology.is_none() && st.ops.keys().all(|o| m.ops.contains_key(o)), "seed {seed}: a cut adds"); let mut me = Machine::start(st); - me.run(st, "p", &mut mail); + // The shell's wiring: what a stage takes in, the stage before holds by name. + for (n, b) in &st.buffers { + if matches!(b.kind, BufferKind::Input | BufferKind::Inout) { + assert!(n == "x" || (b.export && held.contains_key(n)), "seed {seed}: `{n}` arrives from nowhere"); + if let Some(&v) = held.get(n) { + me.bufs.insert(n.clone(), v); + } + } + } + me.run(st, "p"); for (k, v) in me.states { assert!(states.insert(k, v).is_none(), "seed {seed}: a state region written by two stages"); } y = me.bufs.get("y").copied(); + held = me.bufs; } assert_eq!((y, states), (whole.bufs.get("y").copied(), whole.states), "seed {seed} at {at:?}"); } @@ -236,17 +222,9 @@ fn qwen3() -> Verified { verify(Manifest::from_json(QWEN3).unwrap()).unwrap() } -/// The buffers stage `s`'s program puts into its outgoing box. -fn carried(st: &Manifest, program: &str) -> Vec { - st.programs[program] - .calls - .iter() - .filter(|c| c.op.starts_with("pp.put.")) - .map(|c| match &c.args[1] { - Arg::Buf { buf, .. } => buf.clone(), - _ => unreachable!(), - }) - .collect() +/// The buffers of `st` of one kind. +fn of_kind(st: &Manifest, kind: BufferKind) -> Vec<&str> { + st.buffers.iter().filter(|(_, b)| b.kind == kind).map(|(n, _)| n.as_str()).collect() } #[test] @@ -254,13 +232,15 @@ fn qwen3_cuts_at_a_layer_into_stages_that_each_serve() { // l18.qkv_proj: embed, l0 norm, 18 layers of 12 calls. let at = 2 + 18 * 12; assert_eq!(qwen3().programs["prefill"].calls[at].label.as_deref(), Some("l18.qkv_proj")); - let stages = cut(&qwen3(), &[at], &pp()).unwrap_or_else(|e| panic!("{e}")); + let stages = cut(&qwen3(), &[at]).unwrap_or_else(|e| panic!("{e}")); let (head, tail) = (&stages[0], &stages[1]); - // The residual stream and the normed input are all that crosses, in - // every program. - for p in ["prefill", "decode", "decode_batch"] { - assert_eq!(carried(head, p), ["residual", "x"]); - } + // The residual stream and the normed input are all that crosses: the + // head's outputs, rewritten in place on the tail, which exports them + // for the head to copy into. + assert_eq!((of_kind(head, BufferKind::Output), of_kind(head, BufferKind::Inout)), (vec!["residual", "x"], vec![])); + let exported: Vec<&str> = tail.buffers.iter().filter(|(_, b)| b.export).map(|(n, _)| n.as_str()).collect(); + assert_eq!((of_kind(tail, BufferKind::Inout), exported), (vec!["residual", "x"], vec!["residual", "x"])); + assert!(!head.buffers["x"].export); // Each stage holds its own weights and its own calls of the KV. assert!( head.buffers.contains_key("model.embed_tokens.weight") @@ -284,10 +264,10 @@ fn qwen3_cuts_at_a_layer_into_stages_that_each_serve() { assert_eq!(pt.forward(1, kern_manifest::protocol::Rows::Const(1)).map(|f| f.emits.is_some()), Some(true)); assert_eq!(head.cut.as_ref().map(|c| &c.id), tail.cut.as_ref().map(|c| &c.id)); // The same cut is the same bytes. - let again = cut(&qwen3(), &[at], &pp()).unwrap(); + let again = cut(&qwen3(), &[at]).unwrap(); assert_eq!(again[1].to_json(), tail.to_json()); - let three = cut(&qwen3(), &[2 + 12 * 12, 2 + 24 * 12], &pp()).unwrap_or_else(|e| panic!("{e}")); + let three = cut(&qwen3(), &[2 + 12 * 12, 2 + 24 * 12]).unwrap_or_else(|e| panic!("{e}")); assert_eq!(three.len(), 3); Protocol::check_unsampled(&three[1]).unwrap_or_else(|e| panic!("{e}")); assert_ne!(three[0].cut.as_ref().map(|c| &c.id), head.cut.as_ref().map(|c| &c.id)); @@ -297,6 +277,6 @@ fn qwen3_cuts_at_a_layer_into_stages_that_each_serve() { fn a_cut_says_why_it_cannot_be_made() { let n = qwen3().programs["prefill"].calls.len(); for at in [vec![], vec![0], vec![5, 5], vec![n]] { - assert!(cut(&qwen3(), &at, &pp()).is_err(), "{at:?}"); + assert!(cut(&qwen3(), &at).is_err(), "{at:?}"); } } diff --git a/crates/kern-run/examples/pp_chain.rs b/crates/kern-run/examples/pp_chain.rs index 1d9dd429..8838f21b 100644 --- a/crates/kern-run/examples/pp_chain.rs +++ b/crates/kern-run/examples/pp_chain.rs @@ -5,15 +5,14 @@ //! --weights dir --gpus 0,1,2 --prompts ids.txt [--chunk 512] //! [--steps 8] [--capacity 65536] //! -//! The first GPU runs the whole manifest, the rest one stage each, the -//! stages wired through their mailboxes as a pipeline is (here with -//! in-process handles). Each line of `--prompts` is one prompt's token -//! ids. A prompt but its last token goes through the chunk program in -//! `--chunk` pieces, then `--steps` one-row steps feed back what they -//! pick. After every step the input of the call writing the tokens (the -//! logits) and the tokens themselves must be equal on the whole and on the -//! last stage. All stages are issued before any is waited on: a stage's -//! device waits for its predecessor's mailbox on its own. +//! The first GPU runs the whole manifest, the rest one stage each, each +//! stage's arriving buffers mapped into the stage before it as a pipeline +//! does (here with in-process handles) and pushed after every call. Each +//! line of `--prompts` is one prompt's token ids. A prompt but its last +//! token goes through the chunk program in `--chunk` pieces, then +//! `--steps` one-row steps feed back what they pick. After every step the +//! input of the call writing the tokens (the logits) and the tokens +//! themselves must be equal on the whole and on the last stage. use std::collections::BTreeMap; use std::path::PathBuf; @@ -23,8 +22,8 @@ use kern_manifest::protocol::{Forward, Rows}; use kern_manifest::types::Arg; use kern_manifest::{Protocol, Verified}; use kern_pool::Lease; -use kern_run::{fills, run_once, write_named, Vars}; -use kern_runtime::{Capacity, GroupRank, PeerHandle, Runtime, Topology}; +use kern_run::{arrivals, fills, run_once, write_named, Vars}; +use kern_runtime::{Capacity, Mapped, Runtime}; struct Opts { manifest: PathBuf, @@ -72,11 +71,8 @@ fn opts() -> Result { fn load(path: &PathBuf, o: &Opts, gpu: usize, weights: &kern_run::Weights) -> Result<(Runtime, Protocol)> { let m = Verified::from_json(&std::fs::read_to_string(path)?).with_context(|| format!("{}", path.display()))?; - let topology = m.topology.as_ref().map(|t| Topology { - groups: t.groups.iter().map(|(g, &size)| (g.clone(), GroupRank { index: 0, size })).collect(), - }); let capacity = Capacity { tokens: Some(o.capacity), seqs: 2 }; - let mut rt = Runtime::load(&m, Some(&o.kernels), gpu, Some(capacity), topology.as_ref())?; + let mut rt = Runtime::load(&m, Some(&o.kernels), gpu, Some(capacity), None)?; weights.bind(&mut rt)?; let p = Protocol::check_unsampled(&rt.manifest)?; Ok((rt, p)) @@ -121,6 +117,33 @@ fn issue(rt: &mut Runtime, p: &Protocol, program: &str, vars: &Vars, rows: &BTre Ok(rt.issue(program, &vars)?) } +/// One call on every stage in order, each pushing what the next takes +/// in once the next is done with the call before, and landing before +/// the next issues: the pipeline's rule, with the host between every +/// pair. +fn chain( + stages: &mut [(Runtime, Protocol)], + boxes: &[Vec<(String, Mapped)>], + program: &str, + vars: &Vars, + rows: &BTreeMap>, +) -> Result<()> { + for s in 0..stages.len() { + let (rt, p) = &mut stages[s]; + issue(rt, p, program, vars, rows)?; + let Some(b) = boxes.get(s) else { continue }; + let vars: Vars = + vars.iter().filter(|(k, _)| rt.manifest.vars.contains_key(*k)).map(|(k, v)| (k.clone(), *v)).collect(); + stages[s + 1].0.synchronize()?; + let (rt, _) = &stages[s]; + for (name, to) in b { + rt.push(name, &vars, to)?; + } + rt.synchronize()?; + } + Ok(()) +} + fn forwards(p: &Protocol) -> Result<(Forward, Forward)> { Ok(( p.chunk().cloned().context("no chunk program")?, @@ -135,13 +158,15 @@ fn main() -> Result<()> { let (mut whole, pw) = load(&o.manifest, &o, o.gpus[0], &weights)?; let mut stages: Vec<(Runtime, Protocol)> = o.stages.iter().zip(&o.gpus[1..]).map(|(s, &g)| load(s, &o, g, &weights)).collect::>()?; - let exports: Vec> = - stages.iter().map(|(rt, _)| rt.export_handles()).collect::>()?; - for (s, (rt, _)) in stages.iter_mut().enumerate() { - let groups: Vec = rt.manifest.topology.iter().flat_map(|t| t.groups.keys().cloned()).collect(); - for g in groups { - rt.import_peers(&g, &[exports[s].clone(), exports[s + 1].clone()])?; - } + let mut boxes: Vec> = Vec::new(); + for s in 0..stages.len() - 1 { + let (next, np) = &stages[s + 1]; + let handles = next.export_handles()?; + let b = arrivals(&next.manifest, np) + .into_iter() + .map(|n| Ok((n.clone(), stages[s].0.map(&handles[&n], &format!("stage {} `{n}`", s + 1))?))) + .collect::>>()?; + boxes.push(b); } run_once(&whole, &pw)?; for (rt, p) in &stages { @@ -183,9 +208,7 @@ fn main() -> Result<()> { let (vars, rows) = item(&pw, &wlease, *pos, ids)?; issue(&mut whole, &pw, program, &vars, &rows)?; let (vars, rows) = item(&stages[0].1, &slease, *pos, ids)?; - for (rt, p) in &mut stages { - issue(rt, p, program, &vars, &rows)?; - } + chain(&mut stages, &boxes, program, &vars, &rows)?; } let mut tok = *prompt.last().unwrap(); let mut picked = Vec::new(); @@ -195,12 +218,8 @@ fn main() -> Result<()> { let (vars, rows) = item(&pw, &wlease, at, &[tok])?; issue(&mut whole, &pw, &wstep.name, &vars, &rows)?; let (vars, rows) = item(&stages[0].1, &slease, at, &[tok])?; - for (rt, p) in &mut stages { - issue(rt, p, &sstep.name, &vars, &rows)?; - } - for (rt, _) in &stages { - rt.synchronize()?; - } + chain(&mut stages, &boxes, &sstep.name, &vars, &rows)?; + stages.last().unwrap().0.synchronize()?; let (last, _) = stages.last().unwrap(); let wt = tokens.decode(&whole.read_output(&tokens.name)?)[0]; let st = ltokens.decode(&last.read_output(<okens.name)?)[0]; diff --git a/crates/kern-run/src/bin/kern.rs b/crates/kern-run/src/bin/kern.rs index f9214c1d..1ff4e182 100644 --- a/crates/kern-run/src/bin/kern.rs +++ b/crates/kern-run/src/bin/kern.rs @@ -80,10 +80,6 @@ enum Cmd { /// Call indices to cut before, increasing #[arg(long, value_delimiter = ',', required = true)] at: Vec, - /// The mailbox kernels' cubin (`tools/kernels-src/pp_mailbox.cu`), - /// pinned by its sha256; it goes into each stage's kernels dir - #[arg(long)] - module: PathBuf, /// Directory the stages are written to #[arg(long)] out: PathBuf, @@ -135,32 +131,33 @@ fn main() -> Result<()> { } Ok(()) } - Cmd::Cut { manifest, at, module, out } => cut(&manifest, &at, &module, &out), + Cmd::Cut { manifest, at, out } => cut(&manifest, &at, &out), } } /// `kern cut`: the stages of a manifest cut before `at`, written as /// `stage.json`. -fn cut(manifest: &Path, at: &[usize], module: &Path, out: &Path) -> Result<()> { - use sha2::Digest; +fn cut(manifest: &Path, at: &[usize], out: &Path) -> Result<()> { let json = std::fs::read_to_string(manifest).with_context(|| format!("reading {}", manifest.display()))?; let m = kern_manifest::Verified::from_json(&json)?; - let bytes = std::fs::read(module).with_context(|| format!("reading {}", module.display()))?; - let pp = kern_manifest::types::Module { - source: module.file_name().context("--module names a file")?.to_string_lossy().into_owned(), - sha256: hex::encode(sha2::Sha256::digest(&bytes)), - }; - let stages = kern_manifest::cut::cut(&m, at, &pp)?; + let stages = kern_manifest::cut::cut(&m, at)?; std::fs::create_dir_all(out).with_context(|| format!("creating {}", out.display()))?; for (s, st) in stages.iter().enumerate() { let path = out.join(format!("stage{s}.json")); std::fs::write(&path, st.to_json()).with_context(|| format!("writing {}", path.display()))?; let calls: Vec = st.programs.iter().map(|(n, p)| format!("{n} {}", p.calls.len())).collect(); + let crossing: Vec = st + .buffers + .iter() + .filter(|(_, b)| b.export || b.kind == kern_manifest::types::BufferKind::Output) + .map(|(n, b)| format!("{n} {}", b.kind)) + .collect(); tracing::info!( - "stage {s}: {} ({} weights; calls: {})", + "stage {s}: {} ({} weights; calls: {}; crossing: {})", path.display(), st.buffers.values().filter(|b| b.kind == kern_manifest::types::BufferKind::Weight).count(), - calls.join(", ") + calls.join(", "), + crossing.join(", ") ); let (whole, part): (Vec<_>, Vec<_>) = kern_manifest::cut::state_use(&m, st).into_iter().partition(|(_, a, b)| a == b); diff --git a/crates/kern-run/src/lib.rs b/crates/kern-run/src/lib.rs index a53334de..667ef98c 100644 --- a/crates/kern-run/src/lib.rs +++ b/crates/kern-run/src/lib.rs @@ -28,7 +28,7 @@ use std::time::Duration; use anyhow::{bail, ensure, Context, Result}; use kern_manifest::protocol::{Axis, Filled, Forward, LineTable, PageTable, Rows}; -use kern_manifest::types::Fill; +use kern_manifest::types::{BufferKind, Fill}; use kern_manifest::{Protocol, Verified}; use kern_pool::Lease; use kern_runtime::{connect_nccl, GroupRank, PeerHandle, Runtime, Topology}; @@ -196,6 +196,22 @@ pub fn write_named(rt: &mut Runtime, p: &Protocol, name: &str, values: &[i64], v Ok(true) } +/// The inputs of `m` nobody on the host stages: no fill, no table. On a +/// pipeline stage they are what the stage before it holds under the +/// same names and copies in ([`kern_runtime::Runtime::push`]). +pub fn arrivals(m: &kern_manifest::types::Manifest, p: &Protocol) -> Vec { + let staged = |n: &str| { + p.fills.iter().any(|f| f.name == n) + || p.page_tables.iter().any(|t| t.name == n) + || p.line_tables.iter().any(|t| t.name == n) + }; + m.buffers + .iter() + .filter(|(n, b)| matches!(b.kind, BufferKind::Input | BufferKind::Inout) && !staged(n)) + .map(|(n, _)| n.clone()) + .collect() +} + /// Stage one call (see [`fills`]). Returns the call's vars. fn stage( rt: &mut Runtime, diff --git a/crates/kern-runtime/src/device.rs b/crates/kern-runtime/src/device.rs index bf278246..b9792105 100644 --- a/crates/kern-runtime/src/device.rs +++ b/crates/kern-runtime/src/device.rs @@ -408,10 +408,40 @@ impl Pinned { Ok(Pinned { ptr, bytes }) } + /// `bytes` of pinned memory the device can address too, for words the + /// host and a kernel share. + pub(crate) fn alloc_mapped(bytes: u64, dev: i32) -> Result { + let node = numa_node(dev); + if let Some(n) = node { + mempolicy(libc::MPOL_BIND, Some(n)); + } + let mut ptr: *mut c_void = std::ptr::null_mut(); + let flags = sys::CU_MEMHOSTALLOC_DEVICEMAP; + let r = cuda_check( + unsafe { sys::cuMemHostAlloc(&mut ptr, bytes.max(1) as usize, flags) }, + "cuMemHostAlloc(mapped)", + ); + if node.is_some() { + mempolicy(libc::MPOL_DEFAULT, None); + } + r?; + Ok(Pinned { ptr, bytes }) + } + pub(crate) fn ptr(&self) -> u64 { self.ptr as u64 } + /// The address a kernel reads this memory at (an `alloc_mapped` one). + pub(crate) fn device(&self) -> Result { + let mut device = 0; + cuda_check( + unsafe { sys::cuMemHostGetDevicePointer_v2(&mut device, self.ptr, 0) }, + "cuMemHostGetDevicePointer", + )?; + Ok(device) + } + pub(crate) fn bytes(&self) -> u64 { self.bytes } diff --git a/crates/kern-runtime/src/exec.rs b/crates/kern-runtime/src/exec.rs index 37e710fc..8f138553 100644 --- a/crates/kern-runtime/src/exec.rs +++ b/crates/kern-runtime/src/exec.rs @@ -25,53 +25,274 @@ //! another thread. use std::collections::BTreeMap; +use std::ffi::CStr; use std::os::raw::c_void; -use std::sync::Arc; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::{Arc, Mutex}; use cudarc::driver::{sys, CudaContext}; use crate::compile::{CompiledProgram, Dense, Launch, LaunchKind, RVal, Slot}; use crate::cublas::{gemm_bf16_tn, gemm_bf16_tn_f32, gemm_bf16_tn_pinned, gemm_fp8_tn}; +use crate::device::Pinned; use crate::error::{bail, cuda_check}; use crate::{Error, Result, Runtime}; +use kern_manifest::types::BufferKind; /// A point on the compute stream: everything issued before -/// [`Runtime::mark`]. [`Mark::wait`] sleeps until the device has passed it, -/// on any thread, and returns the fault of any launch before it. +/// [`Runtime::mark`]. [`Mark::wait`] blocks until the device has passed +/// it, on any thread. +/// +/// Passing is read off a word in page-locked memory that a kernel at the +/// mark writes, never off the driver: the thread that waits on a mark may +/// be one the issuing thread is waiting for (it reports items done to the +/// stage before, which pushes the next item, which opens the hold the +/// issuing thread's launches sit behind), and any driver call it made +/// could block on the context lock that thread holds inside a blocking +/// launch (lessons.md, 2026-10-08). So a mark carries no event and says +/// nothing about faults; the issuing thread asks [`Runtime::fault`]. pub struct Mark { - ctx: Arc, - event: sys::CUevent, + seq: u64, + passed: Arc, } -// The event is a handle into the context the mark carries; waiting binds -// that context, so the mark may be waited on from any thread. -unsafe impl Send for Mark {} +/// How long a mark is waited for before the device is given up on. +const TIMEOUT: std::time::Duration = std::time::Duration::from_secs(60); impl Mark { + /// Whether the device has passed the mark. + pub fn passed(&self) -> bool { + unsafe { &*(self.passed.ptr() as *const AtomicU64) }.load(Ordering::Acquire) >= self.seq + } + + /// Spins for the first moments, then sleeps in short steps. pub fn wait(self) -> Result<()> { - self.ctx.bind_to_thread()?; - cuda_check(unsafe { sys::cuEventSynchronize(self.event) }, "cuEventSynchronize") + let t = std::time::Instant::now(); + while !self.passed() { + if t.elapsed() > TIMEOUT { + bail!(Cuda, "the device has not passed mark {} in {TIMEOUT:?}", self.seq); + } + if t.elapsed() > std::time::Duration::from_millis(2) { + std::thread::sleep(std::time::Duration::from_micros(20)); + } else { + std::hint::spin_loop(); + } + } + Ok(()) + } +} + +/// A count in page-locked host memory the compute stream can wait on +/// ([`Runtime::hold`]) and any thread can raise ([`Gate::open`]): the +/// mirror of a [`Mark`]. With it a caller issues work before the data it +/// reads is there and the device, not the host, holds the launches back. +pub struct Gate { + host: Pinned, + device: u64, +} + +impl Gate { + /// Raise the count to `n`; every hold for at most `n` lets go. + pub fn open(&self, n: u64) { + unsafe { &*(self.host.ptr() as *const AtomicU64) }.store(n, Ordering::Release); + } +} + +/// The two kernels the runtime signals with, a few lines of PTX loaded +/// once: `hold` spins on a gate's count, `flag` writes a mark's number. +/// The hold is a kernel and not a stream memory operation because +/// `cuStreamWaitValue` stalls the hardware channel its stream is +/// multiplexed onto, and cuBLAS's own streams share those channels: a +/// GEMM issued behind one deadlocked against the next hold (tray08, +/// 2026-10-08). A spinning kernel holds one SM and nothing else. +pub(crate) struct Signals { + module: sys::CUmodule, + hold: sys::CUfunction, + flag: sys::CUfunction, + passed: Arc, + passed_device: u64, + seq: AtomicU64, + fetches: Arc>>, +} + +unsafe impl Send for Signals {} +unsafe impl Sync for Signals {} + +const SIGNALS_PTX: &CStr = c" +.version 8.0 +.target sm_90 +.address_size 64 +.visible .entry hold(.param .u64 gate, .param .u64 count) +{ + .reg .b64 %g, %c, %v; + .reg .pred %p; + ld.param.u64 %g, [gate]; + ld.param.u64 %c, [count]; + cvta.to.global.u64 %g, %g; +L: + ld.acquire.sys.global.u64 %v, [%g]; + setp.ge.u64 %p, %v, %c; + @%p bra D; + nanosleep.u32 256; + bra L; +D: + ret; +} +.visible .entry flag(.param .u64 word, .param .u64 value) +{ + .reg .b64 %w, %v; + ld.param.u64 %w, [word]; + ld.param.u64 %v, [value]; + cvta.to.global.u64 %w, %w; + st.release.sys.global.u64 [%w], %v; + ret; +} +"; + +impl Signals { + pub(crate) fn load(ctx: &Arc, gpu: usize) -> Result { + ctx.bind_to_thread()?; + let mut module: sys::CUmodule = std::ptr::null_mut(); + cuda_check( + unsafe { sys::cuModuleLoadData(&mut module, SIGNALS_PTX.as_ptr() as *const c_void) }, + "cuModuleLoadData(signals)", + )?; + let function = |name: &CStr| { + let mut f: sys::CUfunction = std::ptr::null_mut(); + cuda_check( + unsafe { sys::cuModuleGetFunction(&mut f, module, name.as_ptr()) }, + "cuModuleGetFunction(signals)", + ) + .map(|()| f) + }; + let loaded = function(c"hold").and_then(|hold| Ok((hold, function(c"flag")?))); + let (hold, flag) = match loaded { + Ok(fs) => fs, + Err(e) => { + unsafe { sys::cuModuleUnload(module) }; + return Err(e); + } + }; + let passed = Arc::new(Pinned::alloc_mapped(8, gpu as i32)?); + unsafe { &*(passed.ptr() as *const AtomicU64) }.store(0, Ordering::Release); + let passed_device = passed.device()?; + Ok(Signals { module, hold, flag, passed, passed_device, seq: AtomicU64::new(0), fetches: Arc::default() }) } } -impl Drop for Mark { +impl Drop for Signals { fn drop(&mut self) { - unsafe { sys::cuEventDestroy_v2(self.event) }; + unsafe { sys::cuModuleUnload(self.module) }; + } +} + +/// An output's bytes on their way to the host: a copy the compute stream +/// makes into page-locked memory, complete once a [`Mark`] taken after it +/// has passed. The thread that issues never waits for it. Dropped, the +/// memory goes back to the runtime for the next fetch rather than to the +/// driver: freeing page-locked memory synchronizes the device, and the +/// thread reading a fetch is one that must not (lessons.md, 2026-10-08). +pub struct Fetch { + pinned: Option, + len: usize, + pool: Arc>>, +} + +impl Fetch { + /// The bytes; meaningful only after a mark taken after the fetch passed. + pub fn bytes(&self) -> &[u8] { + let p = self.pinned.as_ref().expect("taken only on drop"); + unsafe { std::slice::from_raw_parts(p.ptr() as *const u8, self.len) } + } +} + +impl Drop for Fetch { + fn drop(&mut self) { + if let Some(p) = self.pinned.take() { + self.pool.lock().unwrap_or_else(|e| e.into_inner()).push(p); + } } } impl Runtime { - /// Mark the compute stream where it is now. The event blocks rather - /// than spins, so a thread waiting on it costs no core. + fn signal(&self, f: sys::CUfunction, word: u64, value: u64, what: &str) -> Result<()> { + self.ctx.bind_to_thread()?; + let mut params = [&word as *const u64 as *mut c_void, &value as *const u64 as *mut c_void]; + let r = unsafe { + sys::cuLaunchKernel( + f, + 1, + 1, + 1, + 1, + 1, + 1, + 0, + self.stream.cu_stream(), + params.as_mut_ptr(), + std::ptr::null_mut(), + ) + }; + cuda_check(r, what) + } + + /// Start fetching `name`, an output buffer, as it is after everything + /// issued so far. + pub fn fetch(&self, name: &str) -> Result { + self.ctx.bind_to_thread()?; + let Some(b) = self.manifest.buffers.get(name) else { + bail!(Api, "no buffer `{name}`"); + }; + if !matches!(b.kind, BufferKind::Output | BufferKind::Inout) { + bail!(Api, "buffer `{name}` is {}, not output", b.kind); + } + let src = &self.buffers[name]; + let pool = &self.signals.fetches; + let spare = { + let mut v = pool.lock().unwrap_or_else(|e| e.into_inner()); + v.iter().position(|p| p.bytes() >= src.bytes).map(|i| v.swap_remove(i)) + }; + let pinned = match spare { + Some(p) => p, + None => Pinned::alloc(src.bytes, self.gpu as i32)?, + }; + let r = unsafe { + sys::cuMemcpyDtoHAsync_v2(pinned.ptr() as *mut c_void, src.ptr, src.bytes as usize, self.stream.cu_stream()) + }; + cuda_check(r, "cuMemcpyDtoHAsync")?; + Ok(Fetch { pinned: Some(pinned), len: src.bytes as usize, pool: pool.clone() }) + } + + /// Mark the compute stream where it is now. pub fn mark(&self) -> Result { + let seq = self.signals.seq.fetch_add(1, Ordering::Relaxed) + 1; + self.signal(self.signals.flag, self.signals.passed_device, seq, "cuLaunchKernel(flag)")?; + Ok(Mark { seq, passed: self.signals.passed.clone() }) + } + + /// The fault of any launch so far, without waiting for the rest: the + /// issuing thread's way of hearing what a mark it never waits on would + /// have said. + pub fn fault(&self) -> Result<()> { self.ctx.bind_to_thread()?; - let mut event: sys::CUevent = std::ptr::null_mut(); - let flags = - sys::CUevent_flags::CU_EVENT_DISABLE_TIMING as u32 | sys::CUevent_flags::CU_EVENT_BLOCKING_SYNC as u32; - cuda_check(unsafe { sys::cuEventCreate(&mut event, flags) }, "cuEventCreate")?; - let mark = Mark { ctx: self.ctx.clone(), event }; - cuda_check(unsafe { sys::cuEventRecord(mark.event, self.stream.cu_stream()) }, "cuEventRecord")?; - Ok(mark) + match unsafe { sys::cuStreamQuery(self.stream.cu_stream()) } { + sys::CUresult::CUDA_SUCCESS | sys::CUresult::CUDA_ERROR_NOT_READY => Ok(()), + r => cuda_check(r, "cuStreamQuery"), + } + } + + /// A gate at count zero. + pub fn gate(&self) -> Result { + let host = Pinned::alloc_mapped(8, self.gpu as i32)?; + let device = host.device()?; + let gate = Gate { host, device }; + gate.open(0); + Ok(gate) + } + + /// Hold the compute stream where it is now until `gate` reaches `n`. + pub fn hold(&self, gate: &Gate, n: u64) -> Result<()> { + self.signal(self.signals.hold, gate.device, n, "cuLaunchKernel(hold)") } fn require_ready(&self) -> Result<()> { diff --git a/crates/kern-runtime/src/lib.rs b/crates/kern-runtime/src/lib.rs index 0ac87b64..38bc2d65 100644 --- a/crates/kern-runtime/src/lib.rs +++ b/crates/kern-runtime/src/lib.rs @@ -58,7 +58,7 @@ use device::{alloc, DeviceBuf, Pinned}; pub use device::{device_uuid, Mapped, PeerHandle}; use error::{bail, cuda_check}; pub use error::{Error, Result}; -pub use exec::Mark; +pub use exec::{Fetch, Gate, Mark}; pub use harness::Scratch; pub use host::HostRegion; pub use host_weights::HostWeights; @@ -135,6 +135,7 @@ pub struct Runtime { buffers: BTreeMap, /// No kernels may read mapped host bytes before checkpoint binding completes. host_weights_ready: bool, + signals: exec::Signals, /// Weight buffers holding their checkpoint bytes: taken from a /// [`Resident`] at load, or copied by `load_weights`. filled: BTreeSet, @@ -269,7 +270,7 @@ impl Runtime { let Some(b) = self.manifest.buffers.get(name) else { bail!(Api, "no buffer `{name}`"); }; - if b.kind != BufferKind::Input { + if !matches!(b.kind, BufferKind::Input | BufferKind::Inout) { bail!(Api, "buffer `{name}` is {}, not input", b.kind); } if data.len() as u64 > self.buffers[name].bytes { @@ -292,7 +293,7 @@ impl Runtime { let Some(b) = self.manifest.buffers.get(name) else { bail!(Api, "no buffer `{name}`"); }; - if b.kind != BufferKind::Output { + if !matches!(b.kind, BufferKind::Output | BufferKind::Inout) { bail!(Api, "buffer `{name}` is {}, not output", b.kind); } Ok(self.stream.clone_dtoh(&self.buffers[name])?) diff --git a/crates/kern-runtime/src/load.rs b/crates/kern-runtime/src/load.rs index 5bc84caf..e527c200 100644 --- a/crates/kern-runtime/src/load.rs +++ b/crates/kern-runtime/src/load.rs @@ -194,7 +194,7 @@ impl Runtime { } let mut staging = BTreeMap::new(); for (name, b) in &manifest.buffers { - if b.kind == BufferKind::Input { + if matches!(b.kind, BufferKind::Input | BufferKind::Inout) { let mut pinned = unsafe { ctx.alloc_pinned::(buffers[name].bytes.max(1) as usize)? }; pinned.as_mut_slice()?.fill(0); staging.insert(name.clone(), pinned); @@ -300,6 +300,7 @@ impl Runtime { let remaps = Remaps::spawn(Arc::clone(&ctx), mapper)?; let mut rt = Runtime { host_weights_ready: !manifest.buffers.values().any(|b| b.placement == Placement::Host), + signals: crate::exec::Signals::load(&ctx, gpu)?, manifest, filled, kept, diff --git a/crates/kern-runtime/src/peers.rs b/crates/kern-runtime/src/peers.rs index 9c88fdc8..35b568b3 100644 --- a/crates/kern-runtime/src/peers.rs +++ b/crates/kern-runtime/src/peers.rs @@ -18,6 +18,7 @@ use std::collections::BTreeMap; +use crate::compile; use crate::device::{self, Mapped, PeerHandle}; use crate::error::bail; use crate::{Error, Result, Runtime}; @@ -43,6 +44,23 @@ impl Runtime { Ok(Mapped::new(device::import(&self.stream, self.gpu as i32, handle, what)?)) } + /// Copy buffer `name`, as long as its shape is at `vars`, into `to`: + /// another runtime's buffer of the same shape, mapped. On the compute + /// stream, so it follows every launch issued before it and precedes + /// every one issued after; a [`Mark`](crate::Mark) taken next says + /// when it has landed. + pub fn push(&self, name: &str, vars: &BTreeMap, to: &Mapped) -> Result<()> { + self.ctx.bind_to_thread()?; + let Some(b) = self.manifest.buffers.get(name) else { + bail!(Api, "no buffer `{name}`"); + }; + let bytes = compile::shaped_bytes(&format!("buffer `{name}`"), &b.shape, b.dtype.bytes(), vars)?; + if bytes > to.bytes() { + bail!(Api, "buffer `{name}` is {bytes} bytes at these vars; the mapping it goes to holds {}", to.bytes()); + } + device::copy_1d(self.stream.cu_stream(), to.ptr(), self.buffers[name].ptr, bytes) + } + /// A handle for every `export` buffer and every state, by name: what /// the other ranks pass to [`Runtime::import_peers`]. pub fn export_handles(&self) -> Result> { diff --git a/crates/kern-serve/Cargo.toml b/crates/kern-serve/Cargo.toml index a5f5e8c1..8a948f1d 100644 --- a/crates/kern-serve/Cargo.toml +++ b/crates/kern-serve/Cargo.toml @@ -12,14 +12,11 @@ kern-run.workspace = true pegainfer-frontend.workspace = true anyhow.workspace = true -bytes.workspace = true clap.workspace = true -futures-util.workspace = true hex.workspace = true serde.workspace = true serde_json.workspace = true -tokio = { workspace = true, features = ["net", "time"] } -tokio-util.workspace = true +tokio.workspace = true tracing.workspace = true tracing-subscriber.workspace = true diff --git a/crates/kern-serve/src/pp/head.rs b/crates/kern-serve/src/pp/head.rs index 7b6c911a..11c393f0 100644 --- a/crates/kern-serve/src/pp/head.rs +++ b/crates/kern-serve/src/pp/head.rs @@ -10,15 +10,16 @@ //! takes the admitted prompts' next rows packed back to back ([`pack`]). //! //! Each item is sent to every stage before the head stages and issues it -//! itself, so a stage waiting on its device for an item always has it on -//! the way. The last stage answers every item, so the head knows how many -//! are in the pipeline: a full item (all its rows or all its sequences) -//! goes out at once, a partial one only -//! while fewer items than stages are in flight, so packing waits for rows -//! only when the stages have work. The first token finishes the request. -//! Nothing waits on the head's device; its [`Watch`] logs each item. +//! itself, so a stage waiting for an item always has it on the way. The +//! last stage answers every item, so the head knows how many are in the +//! pipeline: a full item (all its rows or all its sequences) goes out at +//! once, a partial one only while fewer items than stages are in flight, +//! so packing waits for rows only when the stages have work. The first +//! token finishes the request. The head's own device is waited on only +//! by its [`Flow`], to push each item to stage 1. use std::collections::{BTreeMap, VecDeque}; +use std::net::{TcpListener, TcpStream}; use std::sync::mpsc; use anyhow::{bail, ensure, Context, Result}; @@ -30,11 +31,10 @@ use kern_runtime::{Capacity, Error, Runtime}; use pegainfer_frontend::engine::{ FinishReason, QueuedRequest, RejectReason, RequestId, RequestLedger, Scheduler, SchedulerMetrics, }; -use tokio::net::TcpListener; use tracing::{error, info}; use super::pack::{pack, Bounds, Waiting}; -use super::{connect_downstream, recv, send, topology, wire, Handles, Hello, Item, Reader, Table, Watch, Writer}; +use super::{recv, send, Flow, Hello, Item, Next, Table}; use crate::scheduler::Facts; /// How the head loads and what it may run at once. @@ -77,29 +77,27 @@ pub struct PpHead { depth: usize, max_seqs: usize, stop_tokens: Vec, - io: tokio::runtime::Runtime, - stages: Vec, + stages: Vec, replies: mpsc::Receiver>, waiting: VecDeque, active: VecDeque, /// Per item in flight, oldest first, the prompts its reply finishes. due: VecDeque>, - watch: Watch, + flow: Flow, } /// Wait for every other stage, load stage 0 at the pages they all have, -/// map stage 1's mailbox and answer each stage with its table. +/// answer each stage with its table and join stage 1. pub fn load(m: Verified, weights: &Weights, h: Head) -> Result { let cut = m.cut.clone().context("the head serves a cut manifest")?; ensure!(cut.stage == 0, "`--pp-listen` serves stage 0; this manifest is stage {} of {}", cut.stage, cut.stages); - let io = super::io()?; - let listener = io.block_on(TcpListener::bind(&h.listen)).with_context(|| format!("binding {}", h.listen))?; + let listener = TcpListener::bind(&h.listen).with_context(|| format!("binding {}", h.listen))?; info!(listen = %h.listen, stages = cut.stages, cut = %cut.id, "waiting for the stages"); - let mut hellos: BTreeMap = BTreeMap::new(); + let mut hellos: BTreeMap = BTreeMap::new(); while (hellos.len() as u64) < cut.stages - 1 { - let (s, addr) = io.block_on(listener.accept())?; - let (mut r, w) = wire(s); - let hello: Hello = io.block_on(recv(&mut r))?.with_context(|| format!("{addr} closed before its hello"))?; + let (mut s, addr) = listener.accept()?; + let _ = s.set_nodelay(true); + let hello: Hello = recv(&mut s)?.with_context(|| format!("{addr} closed before its hello"))?; ensure!( hello.cut == cut.id, "{addr}: stage {} is of cut `{}`, the head of `{}`", @@ -114,22 +112,16 @@ pub fn load(m: Verified, weights: &Weights, h: Head) -> Result { cut.stages ); info!(stage = hello.stage, %addr, pages = hello.pages, "stage joined"); - hellos.insert(hello.stage, (hello, r, w)); + let reach = format!("{}:{}", addr.ip(), hello.port); + hellos.insert(hello.stage, (hello, s, reach)); } let unit = kern_pool::page_unit(&m); let fewest = hellos.values().map(|(x, ..)| x.pages).min().expect("at least one other stage"); let tokens = h.capacity.map_or(fewest * unit, |c| c.min(fewest * unit)); let seqs = (h.max_seqs + 1) as u64; - let mut rt = Runtime::load( - &m, - h.kernels.as_deref(), - h.gpu, - Some(Capacity { tokens: Some(tokens), seqs }), - topology(&m).as_ref(), - )?; + let mut rt = Runtime::load(&m, h.kernels.as_deref(), h.gpu, Some(Capacity { tokens: Some(tokens), seqs }), None)?; rt.set_eager(h.eager); weights.bind(&mut rt)?; - connect_downstream(&mut rt, &hellos[&1].0.mailbox)?; let p = Protocol::check_unsampled(&rt.manifest)?; kern_run::run_once(&rt, &p)?; let pages = rt.pages_total() as u64; @@ -147,16 +139,21 @@ pub fn load(m: Verified, weights: &Weights, h: Head) -> Result { }; let context = p.context.as_ref().map_or(usize::MAX, |c| c.max as usize); let bounds = Bounds { rows: chunk, seqs: prefill.groups as usize, context }; - let downstream: Vec> = - (1..cut.stages).map(|s| hellos.get(&(s + 1)).map(|(x, ..)| x.mailbox.clone())).collect(); + let next_of = |s: u64| hellos.get(&(s + 1)).map(|(x, _, reach)| (reach.clone(), x.arrivals.clone())); + let tables: Vec = (1..cut.stages).map(|s| Table { pages, next: next_of(s) }).collect(); + let first = next_of(0).expect("stage 1 said hello"); let (tx, replies) = mpsc::channel(); let mut stages = Vec::new(); - for ((s, (_, r, mut w)), down) in hellos.into_iter().zip(downstream) { - io.block_on(send(&mut w, &Table { downstream_mailbox: down, pages }))?; + for ((s, (_, mut w, _)), table) in hellos.into_iter().zip(tables) { + send(&mut w, &table)?; let tx = (s == cut.stages - 1).then(|| tx.clone()); - io.spawn(listen(s, r, tx)); + std::thread::spawn({ + let r = w.try_clone()?; + move || listen(s, r, tx) + }); stages.push(w); } + let next = Next::connect(&rt, &first.0, &first.1)?; let step_name = step.as_ref().map_or("-", |f| f.name.as_str()).to_string(); info!(pages, page = rt.page(), chunk, seqs = bounds.seqs, prefill = %prefill.name, step = %step_name, @@ -170,21 +167,20 @@ pub fn load(m: Verified, weights: &Weights, h: Head) -> Result { depth: cut.stages as usize, max_seqs: h.max_seqs, stop_tokens: h.stop_tokens, - io, stages, replies, waiting: VecDeque::new(), active: VecDeque::new(), due: VecDeque::new(), - watch: Watch::new(0), + flow: Flow::new(0, None, Some(next), None), }) } /// Read one stage's connection to its end: the last stage's answers go to /// the scheduler; anything closing ends the head. -async fn listen(stage: u64, mut r: Reader, tokens: Option>>) { +fn listen(stage: u64, mut r: TcpStream, tokens: Option>>) { let why = loop { - match (recv::>(&mut r).await, &tokens) { + match (recv::>(&mut r), &tokens) { (Ok(Some(t)), Some(tx)) => { if tx.send(t).is_err() { return; @@ -342,13 +338,13 @@ impl PpHead { } let item = Item { program: f.name.clone(), vars, rows }; for w in &mut self.stages { - self.io.block_on(send(w, &item))?; + send(w, &item)?; } for (name, v) in &item.rows { kern_run::write_named(&mut self.rt, &self.p, name, v, &item.vars)?; } self.rt.issue(&item.program, &item.vars).with_context(|| format!("`{}`", item.program))?; - self.watch.issued(&self.rt, ids.len() as u64)?; + self.flow.issued(&self.rt, &item.vars, ids.len() as u64, None)?; Ok(()) } diff --git a/crates/kern-serve/src/pp/mod.rs b/crates/kern-serve/src/pp/mod.rs index 10ada697..4cc2d875 100644 --- a/crates/kern-serve/src/pp/mod.rs +++ b/crates/kern-serve/src/pp/mod.rs @@ -8,45 +8,55 @@ //! head's pool is sized to the smallest stage: the stages say how many //! pages they have in their [`Hello`], the head loads with the minimum. //! -//! One TCP connection per downstream stage, framed by `tokio-util`'s -//! length-delimited codec, each frame one JSON message. A stage connects -//! and sends its [`Hello`]: its place, the cut it belongs to, its pages and -//! the export handles of its allocations (its mailbox among them). Once -//! every stage has said hello the head loads, maps the first stage's -//! mailbox and answers each stage with a [`Table`]: the next stage's -//! handles, to map its own outgoing peer from, and the pages. From then on -//! the connection carries the log: the head sends every stage every -//! [`Item`] in issue order — the program, its vars and the values of each -//! input it staged — and the last stage answers every item with its tokens -//! (none for a program that hands none back). TCP keeps the order, so an -//! item's number is its place in it and every stage counts the same items -//! alike; there is no sequence number, no stop message and no fault -//! message. A stage that fails exits; a closed connection ends the process -//! on the other side. +//! Two kinds of connection, each a TCP stream of frames (a `u32` length, +//! then JSON). Every stage connects to the head and sends its [`Hello`]: +//! its place, the cut it belongs to, its pages, the port its predecessor +//! may reach it on and the handles of its arriving buffers. Once every +//! stage has said hello the head loads and answers each with a [`Table`]: +//! the pages, and where the next stage listens with what arrives there. +//! From then on the head sends every stage every [`Item`] in issue order +//! — the program, its vars and the values of each input it staged — and +//! the last stage answers every item with its tokens (none for a program +//! that hands none back). TCP keeps the order, so an item's number is its +//! place in it and every stage counts the same items alike; there is no +//! sequence number, no stop message and no fault message. A stage that +//! fails exits; a closed connection ends the process on the other side. //! -//! The activations never touch the connection: they go GPU to GPU through -//! the mailbox the cut wired in, and each stage's device waits for its -//! predecessor on its own. A stage only enqueues; the last one waits for -//! its tokens, nothing else waits on a device. A [`Watch`] beside each -//! stage logs when every item left the device, and is where a device -//! fault (a mailbox wait that gave up on a dead stage) ends the process. +//! The second kind joins each stage to the next, and carries item +//! numbers both ways. A stage's arriving buffers (`kern_run::arrivals`: +//! the inputs nobody on the host stages) are exported, and the stage +//! before maps them at start ([`Next`]). Having issued item `k`, that +//! stage waits until the next one reports item `k - 1` done, copies each +//! arriving buffer over on its compute stream ([`Runtime::push`]) and +//! marks; when the mark lands its [`Watch`] thread reports `k` arrived to +//! the next stage and `k` done to the previous one ([`Flow`]). The next +//! stage's host does not wait for that: it issues `k` behind a hold on +//! its count of arrived items ([`Runtime::hold`]), the launches are in +//! before the data, and the report opens the gate ([`Gate::open`]). So +//! each crossing buffer has one copy per stage and the host is in the +//! loop once per edge and item, off the device's path. The last stage +//! reads its tokens, and every stage logs when each item left it. A next +//! stage that reports nothing done for [`TIMEOUT`] is taken for dead; a +//! dead stage before shows as the head going. pub mod head; mod pack; pub mod stage; use std::collections::BTreeMap; -use std::sync::mpsc; -use std::time::{SystemTime, UNIX_EPOCH}; +use std::io::{Read, Write}; +use std::net::{TcpListener, TcpStream}; +use std::sync::{mpsc, Arc}; +use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; use anyhow::{bail, Context, Result}; -use bytes::Bytes; -use futures_util::{SinkExt, StreamExt}; -use kern_runtime::{GroupRank, Mark, PeerHandle, Runtime, Topology}; +use kern_manifest::protocol::Filled; +use kern_manifest::Protocol; +use kern_runtime::{Fetch, Gate, Mapped, Mark, PeerHandle, Runtime}; use serde::{Deserialize, Serialize}; -use tokio::net::tcp::{OwnedReadHalf, OwnedWriteHalf}; -use tokio::net::TcpStream; -use tokio_util::codec::{FramedRead, FramedWrite, LengthDelimitedCodec}; + +/// How long a stage waits for its neighbour to deliver or take an item. +const TIMEOUT: Duration = Duration::from_secs(60); /// A stage introducing itself to the head. #[derive(Debug, Serialize, Deserialize)] @@ -54,16 +64,18 @@ pub struct Hello { pub stage: u64, pub cut: String, pub pages: u64, - /// Every allocation the stage exports, by name, as hex fabric handles. - pub mailbox: Handles, + /// The port the stage before it connects to. + pub port: u16, + /// Its arriving buffers, by name, as hex fabric handles. + pub arrivals: Handles, } -/// The head's answer: the next stage's exports (none for the last stage) -/// and the pages every stage addresses. +/// The head's answer: the pages every stage addresses, and where the next +/// stage listens with what arrives there (none for the last stage). #[derive(Debug, Serialize, Deserialize)] pub struct Table { - pub downstream_mailbox: Option, pub pages: u64, + pub next: Option<(String, Handles)>, } /// One call, as the head staged it. @@ -77,115 +89,249 @@ pub struct Item { pub type Handles = BTreeMap; -/// This runtime's exports as wire handles; a stage's box must be mappable -/// from another process, so a local handle is an error. -fn handles(rt: &Runtime) -> Result { - rt.export_handles()? +/// This runtime's arriving buffers as wire handles; a stage maps its +/// neighbour across processes, so a local handle is an error. +fn arrivals(rt: &Runtime, p: &Protocol) -> Result { + let all = rt.export_handles()?; + kern_run::arrivals(&rt.manifest, p) .into_iter() - .map(|(name, h)| { - let b = h.to_bytes().with_context(|| { - format!("`{name}` has no fabric handle; a pipeline stage maps its neighbour across processes") - })?; + .map(|name| { + let h = all + .get(&name) + .with_context(|| format!("`{name}` arrives from the stage before but is not exported"))?; + let b = h.to_bytes().with_context(|| format!("`{name}` has no fabric handle"))?; Ok((name, hex::encode(b))) }) .collect() } -fn peer_handles(h: &Handles) -> Result> { - h.iter() - .map(|(name, x)| { - let b = hex::decode(x).with_context(|| format!("handle of `{name}`"))?; - Ok((name.clone(), PeerHandle::from_bytes(&b).with_context(|| format!("handle of `{name}`"))?)) - }) - .collect() +fn send(w: &mut TcpStream, msg: &T) -> Result<()> { + let body = serde_json::to_vec(msg)?; + w.write_all(&(body.len() as u32).to_le_bytes()).and_then(|()| w.write_all(&body)).context("sending") } -/// A stage's place in the group of its outgoing edge: always the -/// upstream member, rank 0 of 2. -fn topology(rt_manifest: &kern_manifest::Manifest) -> Option { - rt_manifest.topology.as_ref().map(|t| Topology { - groups: t.groups.iter().map(|(g, &size)| (g.clone(), GroupRank { index: 0, size })).collect(), - }) +/// The next message; `None` when the other side closed. +fn recv Deserialize<'de>>(r: &mut TcpStream) -> Result> { + let mut len = [0u8; 4]; + match r.read_exact(&mut len) { + Err(e) if e.kind() == std::io::ErrorKind::UnexpectedEof => return Ok(None), + done => done.context("receiving")?, + } + let mut body = vec![0u8; u32::from_le_bytes(len) as usize]; + r.read_exact(&mut body).context("receiving")?; + Ok(Some(serde_json::from_slice(&body)?)) } -/// Map the next stage's box into this stage's outgoing peer buffer. -fn connect_downstream(rt: &mut Runtime, downstream: &Handles) -> Result<()> { - let groups: Vec = rt.manifest.topology.iter().flat_map(|t| t.groups.keys().cloned()).collect(); - let own = rt.export_handles()?; - let theirs = peer_handles(downstream)?; - for g in groups { - rt.import_peers(&g, &[own.clone(), theirs.clone()]) - .with_context(|| format!("mapping the next stage over `{g}`"))?; +/// Reach `addr`, retrying while the other side is still coming up. +fn connect(addr: &str, patience: Duration) -> Result { + let t0 = Instant::now(); + loop { + match TcpStream::connect(addr) { + Ok(s) => { + let _ = s.set_nodelay(true); + return Ok(s); + } + Err(e) if t0.elapsed() < patience => { + tracing::debug!("{addr}: {e}"); + std::thread::sleep(Duration::from_millis(500)); + } + Err(e) => return Err(e).with_context(|| format!("connecting to {addr}")), + } } - let pending = rt.pending_peers(); - if !pending.is_empty() { - bail!("peer buffers {pending:?} still unfilled after the next stage was mapped"); +} + +/// Item numbers between neighbours: one `u64` a frame. +fn tell(w: &mut TcpStream, item: u64) -> Result<()> { + w.write_all(&item.to_le_bytes()).context("telling the neighbour") +} + +/// Read item numbers off `r` until the neighbour goes away. +fn hear(mut r: TcpStream, mut heard: impl FnMut(u64) -> bool) { + let mut item = [0u8; 8]; + while r.read_exact(&mut item).is_ok() && heard(u64::from_le_bytes(item)) {} +} + +/// Wait for `rx` to say `k`, the next item in order. +fn wait(rx: &mpsc::Receiver, k: u64, what: &str) -> Result<()> { + match rx.recv_timeout(TIMEOUT) { + Ok(n) if n == k => Ok(()), + Ok(n) => bail!("{what} item {n}, expected {k}"), + Err(mpsc::RecvTimeoutError::Timeout) => bail!("{what} nothing in {TIMEOUT:?} on item {k}; the stage is gone"), + Err(mpsc::RecvTimeoutError::Disconnected) => bail!("{what}: the connection is gone"), } - Ok(()) } -type Reader = FramedRead; -type Writer = FramedWrite; +/// The edge out of a stage: the next stage's arriving buffers mapped +/// (each one this stage holds under the same name), and the items it +/// reports done. +struct Next { + boxes: Vec<(String, Mapped)>, + done: mpsc::Receiver, +} -fn wire(s: TcpStream) -> (Reader, Writer) { - let _ = s.set_nodelay(true); - let (r, w) = s.into_split(); - (FramedRead::new(r, LengthDelimitedCodec::new()), FramedWrite::new(w, LengthDelimitedCodec::new())) +impl Next { + /// Map the next stage's arrivals and connect to it. Returns the edge + /// and the write half the watch reports arrivals on. + fn connect(rt: &Runtime, addr: &str, arrivals: &Handles) -> Result<(Next, TcpStream)> { + let mut boxes = Vec::new(); + for (name, x) in arrivals { + if !rt.manifest.buffers.contains_key(name) { + bail!("the next stage takes `{name}` from this one, which has no such buffer"); + } + let b = hex::decode(x).with_context(|| format!("handle of `{name}`"))?; + let h = PeerHandle::from_bytes(&b).with_context(|| format!("handle of `{name}`"))?; + boxes.push((name.clone(), rt.map(&h, &format!("the next stage's `{name}`"))?)); + } + let s = connect(addr, TIMEOUT)?; + let (tx, done) = mpsc::channel(); + std::thread::spawn({ + let r = s.try_clone()?; + move || hear(r, |k| tx.send(k).is_ok()) + }); + Ok((Next { boxes, done }, s)) + } } -async fn send(w: &mut Writer, msg: &T) -> Result<()> { - w.send(Bytes::from(serde_json::to_vec(msg)?)).await.context("sending") +/// The edge into a stage: how many items the stage before reports +/// arrived, as a gate the device waits on. +struct Prev { + arrived: Arc, } -/// The next message; `None` when the other side closed. -async fn recv Deserialize<'de>>(r: &mut Reader) -> Result> { - match r.next().await { - None => Ok(None), - Some(frame) => Ok(Some(serde_json::from_slice(&frame.context("receiving")?)?)), +impl Prev { + /// Accept the stage before. Returns the edge and the write half the + /// watch reports taken items on. + fn accept(rt: &Runtime, listener: &TcpListener) -> Result<(Prev, TcpStream)> { + let (s, _) = listener.accept().context("accepting the stage before")?; + let _ = s.set_nodelay(true); + let arrived = Arc::new(rt.gate()?); + std::thread::spawn({ + let r = s.try_clone()?; + let gate = arrived.clone(); + move || { + hear(r, |k| { + gate.open(k + 1); + true + }) + } + }); + Ok((Prev { arrived }, s)) } } -/// The IO runtime a process's connections live on, one thread beside the -/// one driving the GPU. -fn io() -> Result { - Ok(tokio::runtime::Builder::new_multi_thread().worker_threads(1).enable_all().build()?) +/// What the last stage answers an item with: its token output on the way +/// to the host, decoded and cut to the item's groups once the item is done. +pub(super) struct Answer { + pub fetch: Fetch, + pub fill: Filled, + pub keep: usize, +} + +/// A stage's place in the flow of items: what it waits for before +/// issuing one and what it does after. +pub(super) struct Flow { + prev: Option, + next: Option, + watch: Watch, + item: u64, +} + +impl Flow { + /// `head` is the connection the last stage answers items on. + fn new( + stage: u64, + prev: Option<(Prev, TcpStream)>, + next: Option<(Next, TcpStream)>, + head: Option, + ) -> Flow { + let (prev, prev_w) = prev.map_or((None, None), |(p, w)| (Some(p), Some(w))); + let (next, next_w) = next.map_or((None, None), |(n, w)| (Some(n), Some(w))); + Flow { prev, next, watch: Watch::new(stage, prev_w, next_w, head), item: 0 } + } + + /// Hold the compute stream until the next item has arrived from the + /// stage before; the host goes on issuing. + fn arrived(&self, rt: &Runtime) -> Result<()> { + match &self.prev { + Some(p) => Ok(rt.hold(&p.arrived, self.item + 1)?), + None => Ok(()), + } + } + + /// The next item is issued, of `rows` rows at `vars`: once the next + /// stage has taken the one before it, push it on, and mark it. + fn issued(&mut self, rt: &Runtime, vars: &BTreeMap, rows: u64, answer: Option) -> Result<()> { + let k = self.item; + if let Some(n) = &self.next { + if k > 0 { + wait(&n.done, k - 1, "the next stage reported")?; + } + for (name, to) in &n.boxes { + rt.push(name, vars, to).with_context(|| format!("pushing `{name}` to the next stage"))?; + } + } + rt.fault()?; + self.watch.issued(Issued { item: k, rows, issued_us: now_us(), mark: rt.mark()?, answer })?; + self.item += 1; + Ok(()) + } } /// When each item a stage issued left its device: one `pp item` line per /// item (its number, its rows, `issued_us` and `done_us` of the wall -/// clock), logged by a -/// thread that waits on the item's [`Mark`] so the issuing thread never -/// does. A stage's period is done to done; against what the stage alone -/// takes for the item's shape, the rest is its pipeline bubble. A fault on -/// the device surfaces at its item's mark and ends the process. +/// clock), logged by a thread that waits on the item's [`Mark`] so the +/// issuing thread never does, and tells the neighbours: the next stage +/// that the item arrived, the previous that it was taken. On the last +/// stage it also answers the head with the item's tokens. A stage's +/// period is done to done; against what the stage alone takes for the +/// item's shape, the rest is its pipeline bubble. A fault on the device +/// surfaces at its item's mark and ends the process. struct Watch { - marks: Option>, - thread: Option>, - next: u64, + issued: Option>, +} + +struct Issued { + item: u64, + rows: u64, + issued_us: u64, + mark: Mark, + answer: Option, } impl Watch { - fn new(stage: u64) -> Watch { - let (marks, rx) = mpsc::channel::<(u64, u64, u64, Mark)>(); - let thread = std::thread::spawn(move || { - for (item, rows, issued_us, mark) in rx { + fn new(stage: u64, mut prev: Option, mut next: Option, mut head: Option) -> Watch { + let (issued, rx) = mpsc::channel::(); + std::thread::spawn(move || { + for Issued { item, rows, issued_us, mark, answer } in rx { if let Err(e) = mark.wait() { tracing::error!(stage, item, "device: {e:#}"); std::process::exit(1); } tracing::info!(stage, item, rows, issued_us, done_us = now_us(), "pp item"); + let tokens = answer.map_or(Vec::new(), |a| { + let mut v = a.fill.decode(a.fetch.bytes()); + v.truncate(a.keep); + v + }); + let told = next + .as_mut() + .map_or(Ok(()), |w| tell(w, item)) + .and(prev.as_mut().map_or(Ok(()), |w| tell(w, item))) + .and(head.as_mut().map_or(Ok(()), |h| send(h, &tokens))); + if let Err(e) = told { + tracing::error!(stage, item, "{e:#}"); + std::process::exit(1); + } } }); - Watch { marks: Some(marks), thread: Some(thread), next: 0 } + Watch { issued: Some(issued) } } - /// Mark the item just issued, of `rows` rows. - fn issued(&mut self, rt: &Runtime, rows: u64) -> Result<()> { - let marks = self.marks.as_ref().expect("taken only on drop"); - if marks.send((self.next, rows, now_us(), rt.mark()?)).is_err() { + fn issued(&self, i: Issued) -> Result<()> { + let issued = self.issued.as_ref().expect("taken only on drop"); + if issued.send(i).is_err() { bail!("the watch thread is gone"); } - self.next += 1; Ok(()) } } @@ -194,12 +340,10 @@ fn now_us() -> u64 { SystemTime::now().duration_since(UNIX_EPOCH).map_or(0, |d| d.as_micros() as u64) } -/// The items still in flight are logged before the stage goes. +/// The watch goes with the process: an item still in flight when the head +/// is gone may be holding for data that will never come. impl Drop for Watch { fn drop(&mut self) { - drop(self.marks.take()); - if let Some(t) = self.thread.take() { - let _ = t.join(); - } + drop(self.issued.take()); } } diff --git a/crates/kern-serve/src/pp/stage.rs b/crates/kern-serve/src/pp/stage.rs index 3e17c8d9..04863f61 100644 --- a/crates/kern-serve/src/pp/stage.rs +++ b/crates/kern-serve/src/pp/stage.rs @@ -1,11 +1,12 @@ //! A downstream stage: loads its slice of the model, says hello to the -//! head, maps the next stage's mailbox from the table it gets back, then -//! replays the head's items in order. It keeps no accounting: the page -//! ids and positions of every item are the head's. A stage serves paged -//! state only; a per-sequence state (its slot zeroed on a sequence's -//! first item) waits for a model that needs it. The last stage answers -//! every item with its tokens. +//! head, joins the stage after it and the one before from the table it +//! gets back, then replays the head's items in order as they arrive. It +//! keeps no accounting: the page ids and positions of every item are the +//! head's. A stage serves paged state only; a per-sequence state (its +//! slot zeroed on a sequence's first item) waits for a model that needs +//! it. The last stage answers every item with its tokens. +use std::net::TcpListener; use std::sync::mpsc; use std::time::{Duration, Instant}; @@ -13,10 +14,9 @@ use anyhow::{bail, ensure, Context, Result}; use kern_manifest::{Protocol, Verified}; use kern_run::{Vars, Weights}; use kern_runtime::{Capacity, Runtime}; -use tokio::net::TcpStream; use tracing::{error, info}; -use super::{connect_downstream, handles, recv, send, topology, wire, Hello, Item, Table, Watch}; +use super::{arrivals, connect, recv, send, Answer, Flow, Hello, Item, Next, Prev, Table}; /// How a stage loads and where its head listens. pub struct Stage { @@ -37,7 +37,7 @@ pub fn run(m: Verified, weights: &Weights, o: Stage) -> Result<()> { let last = cut.stage + 1 == cut.stages; let t0 = Instant::now(); let capacity = Capacity { tokens: o.capacity, seqs: (o.max_seqs + 1) as u64 }; - let mut rt = Runtime::load(&m, o.kernels.as_deref(), o.gpu, Some(capacity), topology(&m).as_ref())?; + let mut rt = Runtime::load(&m, o.kernels.as_deref(), o.gpu, Some(capacity), None)?; rt.set_eager(o.eager); weights.bind(&mut rt)?; let p = Protocol::check_unsampled(&rt.manifest)?; @@ -49,35 +49,42 @@ pub fn run(m: Verified, weights: &Weights, o: Stage) -> Result<()> { "stage loaded" ); - let io = super::io()?; - let stream = io.block_on(connect(&o.head))?; - let (mut r, mut w) = wire(stream); - let hello = Hello { stage: cut.stage, cut: cut.id.clone(), pages: rt.pages_total() as u64, mailbox: handles(&rt)? }; - io.block_on(send(&mut w, &hello))?; - let table: Table = io.block_on(recv(&mut r))?.context("the head closed before its table")?; + let listener = TcpListener::bind("0.0.0.0:0")?; + let mut head = connect(&o.head, CONNECT)?; + let hello = Hello { + stage: cut.stage, + cut: cut.id.clone(), + pages: rt.pages_total() as u64, + port: listener.local_addr()?.port(), + arrivals: arrivals(&rt, &p)?, + }; + send(&mut head, &hello)?; + let table: Table = recv(&mut head)?.context("the head closed before its table")?; ensure!( table.pages <= rt.pages_total() as u64, "the head addresses {} pages, this stage has {}", table.pages, rt.pages_total() ); - match (&table.downstream_mailbox, last) { - (Some(d), false) => connect_downstream(&mut rt, d)?, - (None, true) => {} - (d, _) => bail!( - "stage {} of {}: the head sent {} next stage's mailbox", + let next = match (table.next, last) { + (Some((addr, h)), false) => Some(Next::connect(&rt, &addr, &h)?), + (None, true) => None, + (n, _) => bail!( + "stage {} of {}: the head sent {} next stage", cut.stage, cut.stages, - if d.is_some() { "a" } else { "no" } + if n.is_some() { "a" } else { "no" } ), - } + }; + let prev = Prev::accept(&rt, &listener)?; kern_run::run_once(&rt, &p)?; info!(stage = cut.stage, head = %o.head, last, "stage serving"); let (tx, items) = mpsc::channel::(); - io.spawn(async move { - loop { - match recv::(&mut r).await { + std::thread::spawn({ + let mut r = head.try_clone()?; + move || loop { + match recv::(&mut r) { Ok(Some(item)) => { if tx.send(item).is_err() { return; @@ -95,41 +102,27 @@ pub fn run(m: Verified, weights: &Weights, o: Stage) -> Result<()> { } }); - let mut watch = Watch::new(cut.stage); + let answers = last.then(|| head.try_clone()).transpose()?; + let mut flow = Flow::new(cut.stage, Some(prev), next, answers); for item in items { let vars: Vars = item.vars.into_iter().filter(|(k, _)| rt.manifest.vars.contains_key(k)).collect(); + // The fills and tables are stream-ordered behind the item before, + // so they go over while the activations are still on their way. for (name, v) in &item.rows { kern_run::write_named(&mut rt, &p, name, v, &vars)?; } + flow.arrived(&rt)?; rt.issue(&item.program, &vars).with_context(|| format!("`{}`", item.program))?; - watch.issued(&rt, vars.get(&p.rows.var).copied().unwrap_or(0))?; - if last { - let emits = p.forwards.iter().find(|f| f.name == item.program).and_then(|f| f.emits); - let tokens = match emits.map(|i| &p.fills[i]) { - Some(t) => { - let mut v = t.decode(&rt.read_output(&t.name)?); - v.truncate(vars.get(&p.groups.var).copied().unwrap_or(1) as usize * t.width as usize); - v - } - None => Vec::new(), - }; - io.block_on(send(&mut w, &tokens))?; - } + let emits = p.forwards.iter().find(|f| f.name == item.program).and_then(|f| f.emits); + let answer = emits + .filter(|_| last) + .map(|i| { + let fill = p.fills[i].clone(); + let keep = vars.get(&p.groups.var).copied().unwrap_or(1) as usize * fill.width as usize; + Ok::<_, anyhow::Error>(Answer { fetch: rt.fetch(&fill.name)?, fill, keep }) + }) + .transpose()?; + flow.issued(&rt, &vars, vars.get(&p.rows.var).copied().unwrap_or(0), answer)?; } Ok(()) } - -/// Reach the head, retrying while it is still loading. -async fn connect(head: &str) -> Result { - let t0 = Instant::now(); - loop { - match TcpStream::connect(head).await { - Ok(s) => return Ok(s), - Err(e) if t0.elapsed() < CONNECT => { - tracing::debug!("head {head}: {e}"); - tokio::time::sleep(Duration::from_millis(500)).await; - } - Err(e) => return Err(e).with_context(|| format!("connecting to the head at {head}")), - } - } -} diff --git a/docs/lessons.md b/docs/lessons.md index 2af6be24..b67dfc49 100644 --- a/docs/lessons.md +++ b/docs/lessons.md @@ -3,6 +3,42 @@ 设计写在 runtime.md / serve.md / multi-gpu.md 里；这里只记那些"不写下来下次还会 再踩一遍"的事，以及它们落到了哪条规则上。 +## 2026-10-08，流水线段的 host 提前 launch：四个死锁 + +段的 host 想在数据到达前就把 item 的 launch 发出去，让设备自己等。独立的 C 程序里每一种写法都 +能醒，kern 进程里每一种都在前几个 item 就死掉，一共踩了四次，规则在最后。 + +**`cuStreamWaitValue` 卡的是硬件 channel，不是 stream。** 第一版用流内存操作等 pinned 内存里的 +到达计数。它在 channel 上是一个 semaphore acquire，多路复用到同一 channel 的其他 stream 全排在 +它后面（`CUDA_DEVICE_MAX_CONNECTIONS=1` 的 C 程序就复现：旁路 stream 的 8 字节拷贝永远不完成）， +kern 进程里 cuBLAS 自己的 stream 和下一个 hold 就这样绕成了环。设备侧的等待改成一个自旋 kernel +（占一个 SM，channel 照常往下走）。 + +**往 pageable 内存的拷贝在 API 调用里阻塞，并且攥着 context 锁。** 末段主线程 `read_output` 走 +`cuMemcpyDtoHAsync` 到 `Vec`，driver 在调用内部等 DMA 完成，这段等待不像 `cuStreamSynchronize` +那样放锁，watch 线程的 `cuEventSynchronize` 于是在 futex 里睡死，"k 取走了"发不出去，上游不推 +k+1，主线程等的 DMA 永远排在 hold 后面。线程表一看就明白：主线程 90% CPU、watch 线程 0% 在 +futex_wait。输出改走 `Runtime::fetch`（stream 上拷进 pinned，watch 线程在 mark 落地后读）。 + +**cuBLAS 的 eager launch 自己也会这样阻塞。** 连发两个 chunk，第二个 prefill 的第一条 +cublasLt GEMM 在调用里自旋，直到前一个 item 的 hold 打开（driver 给 pageable 参数拷贝的 staging +池用完了，猜测，但自测可复现：hold 后面的第一条 GEMM 等 300 ms 直到门开），同样攥着锁； +watch 线程此时只要碰 driver（哪怕是 mark 落地之后查 fault 的 `cuEventSynchronize`、或 drop +event）就进同一个环。于是 mark 不再带 event：设备在 mark 处用一个 kernel 把序号写进 pinned +内存，watch 线程只读内存；fault 由发射线程自己每个 item 用 `cuStreamQuery`（不阻塞）问一次。 + +**释放 pinned 内存也是同步。** 改完之后 qwen3 的 2 段流水能跑 18 个 item，K3 的 8 段在末段的第一个 +item 上又死了：watch 线程记完 "pp item" 就没了下文，socket 上一个字节没发出去，线程 100% CPU。 +日志行和 tell 之间只剩 `Answer` 的 drop——`Fetch` 里那块 pinned 内存的 `cuMemFreeHost`，它要等设备上 +所有在飞的工作结束（含下一个 item 没开的 hold），而且是自旋等。qwen3 没踩中只是时序：2 段的 hop 短， +free 的时候下一个 hold 多半还没发出去。`Fetch` 的内存现在回收进 runtime 的池子，下次 fetch 复用， +runtime 没了才还给 driver。 + +规则：**提前 launch 之后，等待设备的线程一个 driver 调用都不许做**——到达用 `Runtime::hold` +（自旋 kernel）、完成用 `Runtime::mark`（设备写 pinned 字）、输出用 `Runtime::fetch`、fault 用 +发射线程的 `Runtime::fault`；`read_output` 与 `cuStreamWaitValue` 只给没人提前跑的调用方。 +主线程被 cuBLAS 卡住本身无害（门一开就走），但它决定了 host 最多只能领先一个 item。"一个 driver 调用都不许做"包括 drop：任何带 CUDA 资源的类型都不能在这种线程上析构。 + ## 2026-09-22，cuBLASLt fp8 在图捕获里报 "green context" **capture 里的错误往往出在报错的前一条。** `extern:cublaslt_fp8_tn` eager 全对，`capture` 一进 diff --git a/docs/manifest.md b/docs/manifest.md index 0130999a..c3618bb2 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 / source / workspace / carry / peer +buffers. buffer 有类型的张量：input / output / inout / weight / source / workspace / carry / peer topology.groups. group 多卡 SPMD 的 rank 组：只有名字和大小 cut cut 流水线的一段：id、第几段、共几段 ``` @@ -411,16 +411,23 @@ kern-serve 只拿 `Protocol` 驱动 runtime，不读 JSON（CI 用 grep 保证 ## 流水线切分：`kern cut` -`kern cut --at [,…] --module pp_mailbox.cubin --out -` 把一份单卡 manifest 在 call 下标处切成 S 段（`tools/kernels-src/pp_mailbox.cu` -是邮箱核，与模型无关；`kern_manifest::cut::cut` 是纯函数，不碰 GPU）： +`kern cut --at [,…] --out ` 把一份单卡 manifest +在 call 下标处切成 S 段（`kern_manifest::cut::cut` 是纯函数，不碰 GPU）。 +**一段就是一个函数**：从上一段来的值是它的 input，给下一段的值是它的 +output，manifest 里没有任何流水线专有的东西；怎么传、何时传是 serving +外壳的事（serve.md「流水线」）。 - 每个 batch program 在同一下标处切开（下标须小于每个 batch program 的 - 长度）；once program 在每段按"本段读到的东西"往回切片。活跃性分析找出 - 切点两侧都活着的 buffer：只许是 workspace、至多一个 var 维（字节数 - `var × 常数`）；output / carry 只许在末段写；一个 (state, offset) 区域 - 只许一段碰（按层切 KV 时各段各持自己那些层）。做不到的切点报错并说明 - 哪条不满足，不会静默切出错的东西。 + 长度）；once / derive program 在每段按"本段读到的东西"往回切片。活跃性 + 分析找出切点两侧都活着的 buffer：只许是 workspace；output / carry 只许在 + 末段写；一个 (state, offset) 区域只许一段碰（按层切 KV 时各段各持自己 + 那些层）。做不到的切点报错并说明哪条不满足，不会静默切出错的东西。 +- 跨切点的 buffer 在上游段是 `output`，在下游段是 `input`（本段只读它， + 无论是不是还要传给再下一段）或 `inout`（本段原地改写它：中间段的 + `hidden`）。`inout` 是一般概念，不是流水线概念——运行前有值、可写、不 + 要求必写、运行后可读回。到达的 buffer（input / inout）带 `export: true`， + 上游段映射它、往里拷。一段既不读也不写、只是路过的值被拒：一段只携带 + 它用到的东西。 - state 整个留或整个丢：一段只要碰到一个 state 的一块就分配整个。`kern cut` 逐段打印每个 state 被本段碰到的区域数 / 整份碰到的区域数（区域 = 不同的 call offset，按等大估字节），不满的就是给别段占着的内存（报告，不拒）。 @@ -428,26 +435,18 @@ kern-serve 只拿 `Protocol` 驱动 runtime，不读 JSON（CI 用 grep 保证 自己那些层；生成器不知道流水线。页表 / slot 的 `index_into` 指向一个本段 丢掉的 state 时改指本段留下的同类 state（pool 给所有 paged state 同一套 页号）。 -- 边界处插一对 send / recv：上游段末尾 `pp.wait_empty → pp.put.… - → pp.post_full`，下游段开头 `pp.wait_full → pp.take.… → - pp.post_empty`；邮箱 `pp..box`（`u8[256 + 最大载荷]`）在下游 GPU 上， - `full` / `empty` 两个计数在头里，上游经 `pp..peer` 远程写。等待在设备 - 上自旋，超时（60 s）`__trap`——进程死于 launch failure，不会挂住。 -- 切分只加传输，不加观测：首段没有 recv、末段没有 send，各段的输出都是 - forward 的纯函数。一个 item 什么时候离开设备是 serving 外壳的事（host 上 - 等 event，见 serve.md「流水线」）。 -- 非末段声明拓扑组 `pp.`（大小 2，本段是 0 号）。**这里放宽了 SPMD**： - 组两端装的是不同的 manifest，只在 `of` buffer（同名同字节数的 - `pp..box`，上游那份是不用的 workspace + export）上对上—— - `import_peers` 本来就只按名字和字节数匹配，runtime 不用改。 -- 首段之外的段没有 `token` fill，也常常没有 `seq_len`（长度随激活从邮箱来）， - `Protocol` 对它们没有特判：轴写在 var 上，角色有就用。 - 一个 chunk program 若自己吐 token（hybrid GDN）则拒切：prefill-only 的 - 流水线约定首个 token 从一行的 step program 出。 +- 切分只改签名，不加 op、不加观测：各段的输出都是 forward 的纯函数。一个 + item 什么时候离开设备是 serving 外壳的事（host 上等 event）。每段带 + `cut {id, stage, stages}`，只用于握手核对同一次切分。 +- 首段之外的段没有 `token` fill，也常常没有 `seq_len`（长度随激活来）， + `Protocol` 对它们没有特判：轴写在 var 上，角色有就用。首 token 从一行 + 的 step program 出，或从自己吐 token 的 chunk program 出（K3）；两者都 + 只在末段写 output。 门禁：`crates/kern-manifest/tests/cut.rs` 用随机 manifest（400 个种子） + 符号解释器验证"能切的切点，各段依次跑 == 整份跑"，且确实有切点被拒； -GPU 上 `crates/kern-run/examples/pp_chain.rs` 把各段在同一进程里串起来， +GPU 上 `crates/kern-run/examples/pp_chain.rs` 把各段在同一进程里串起来（下一段 +到达的 buffer 映射进上一段、每个 call 后 `Runtime::push`）， logits 与 token 与整份逐位同（serve.md「流水线」）。 ## Verifier（`kern-manifest`） diff --git a/docs/roadmap.md b/docs/roadmap.md index 40bc6376..fe88e1cd 100644 --- a/docs/roadmap.md +++ b/docs/roadmap.md @@ -64,7 +64,7 @@ kern 在内部 RSI 里是 kernel agent 的反馈环境（`kern test` 裁判、`k |---|---|---| | PP1 ✅ | `kern cut`（按 call 下标切、活跃性找边界、send/recv 邮箱核）+ kern-serve 一段一进程的 prefill-only 流水（head 持唯一的池，TCP 发 item，GPU 邮箱传激活） | **2026-10-07 tray07 过**：qwen3-4b 2 段，各段依次跑 logits/token 与整份逐位同（chunk 512/64）；两进程首 token 与单卡同；8193 token conc1 TTFT 313 → 174 ms（serve.md「流水线」） | | PP1.1 ✅ | qwen3-4b 的 KV 每层一个 state，`kern cut` 每段只带自己碰的 state、逐段报告 state 用了几成 | **2026-10-07 tray07 过**：`kern test qwen3-4b` 位一致、`kern run` 文本同；chain / serve gate 位一致；每段页数 122,993 → 245,760 | -| PP3 | manifest 说得出"值从另一份 manifest 来"：跨切点的 buffer 是一种边界 kind（`edge: in` 运行前已到达，`edge: out` 运行后送出），字段只有方向、dtype 与 shape，不认"第几段"；连到谁是部署的事，不进 manifest（`cut {id, stage, stages}` 只做握手核对）；runtime 负责等、拷、通知，`pp_mailbox.cu` 退成它的一种传输（NVLink 邮箱 / RDMA / 同进程拷贝可换）；`kern cut` 只切 call 列表、标边界，不再生成 `pp.*` op，不再认识任何 kernel。EP 的 all-to-all 是这个 kind 的第二个用户，字段名为它留门 | 切出的段里没有 `pp.*` op，kern-manifest 里没有 kernel 名；verify 拒绝本段写 `in` 边界、本段不写 `out` 边界；chain / serve gate 位一致 | +| PP3 ✅ | 一段就是一个函数：跨段的 buffer 在上游是 `output`、在下游是 `input` 或 `inout`（原地改写的 input，一般概念）；manifest 里没有流水线专有的东西，`kern cut` 只切 call、标 kind、给到达的 buffer `export`；传输是 serving 外壳的：下一段的 arrivals 映射进上一段，`Runtime::push` + `Runtime::mark`，信号走 TCP，下游在 stream 上 `Runtime::hold` 一个到达计数的门、host 提前 launch，mark / fetch / fault 都不让等待线程碰 driver；`pp_mailbox.cu`、`pp.*` op、每条边的 peer group 删除 | **2026-10-08**：切出的段里没有 `pp.*` op 也没有 topology；kern-manifest 对 kernel 一无所知；qwen3 pp2 chain / serve gate 位一致、faults 同前；K3 PP8 容量见 serve.md | | PP2 | decode 走流水；K3 的 state 拆分；per-seq state（GDN）与吐 token 的 chunk program | 在 PP3 之后：decode 的流水用边界 kind，不用 `pp.*` op | | EP-PP | EP 也一卡一进程：head 持 N 个池、全连握手分发 handle 与 NcclId、每步给每个 rank 发它自己的 item、token 回 head；`Tray` 的单进程多 runtime 路径删除；manifest 不变 | K3 EP8 decode 跨两 tray 进 kern-serve，conc1 文本与单 tray 同；每 rank 的 `pp item` 时间线 | diff --git a/docs/serve.md b/docs/serve.md index 068dab16..d97ac16f 100644 --- a/docs/serve.md +++ b/docs/serve.md @@ -570,12 +570,23 @@ toy 不测 kernel 数值和性能，也没有 tray（EP4 的 toy 要一个走 pe ## 流水线（prefill-only PP v1，2026-10-07，tray07 GB300） 一段一个进程：`kern cut` 切出的 stage 0 是 head（`--pp-listen host:port`：HTTP、scheduler、唯一的 -`Pool`），其余段 `--pp-head host:port`。握手：每段连上 head 发 `Hello { stage, cut, pages, mailbox }` -（export handle 的 hex），head 收齐 S−1 个、按最少的页数装自己的池，回 `Table { downstream_mailbox, -pages }`；之后同一条 TCP（tokio-util `LengthDelimitedCodec` + JSON）上 head 按发射顺序给每段发 -`Item { program, vars, rows }`（fill 与页表的值，页号全流水共用），末段对吐 token 的 item 回 token。 -激活不走 TCP，走 GPU 上的邮箱（manifest.md「流水线切分」），各段设备自己等上游。没有序号、 -stop、fault 消息：任何错误进程退出，连接断了对端退出；设备上的邮箱等待 60 s 超时 trap。 +`Pool`），其余段 `--pp-head host:port`。握手：每段连上 head 发 `Hello { stage, cut, pages, port, arrivals }` +（`arrivals` = 本段 host 不填的 input / inout，即上一段要拷进来的，export handle 的 hex；`port` 是上一段 +连本段的端口），head 收齐 S−1 个、按最少的页数装自己的池，回 `Table { pages, next }`（下一段的地址 +与 arrivals）；之后同一条 TCP（u32 长度 + JSON）上 head 按发射顺序给每段发 `Item { program, vars, rows }` +（fill 与页表的值，页号全流水共用），末段对吐 token 的 item 回 token。 +激活走 NVLink，信号走 host：每段把下一段的 arrivals 映射进来（`Runtime::map`），issue 完 item k 后等 +下一段报告 k−1 已取走，再在 compute stream 上 `Runtime::push` 同名 buffer、打 `Runtime::mark`；mark 落地 +时 watch 线程告诉下一段"k 到了"、上一段"k 取走了"。下一段的 host 不等这句话：它先在 compute stream 上 +放一个"到达数 ≥ k+1"的门（`Runtime::hold`：一个在 pinned 计数上自旋的 kernel，runtime 自带的几行 +PTX），接着就把 k 的 launch 全发出去，收到"到了"时写计数（`Gate::open`）放行——等数据的是设备，host 的 +launch 不在关键路径上（qwen3 的 prefill 不走图，218 个 launch 要 0.6 ms，等到了再发会让 conc4 慢 +3–30%；cuBLAS 的 eager launch 排在没开的门后面时会在调用里自旋到门开，所以 host 最多领先一个 +item）。mark 也不走 driver：设备在 mark 处用一个 kernel 把序号写进 pinned 内存，watch 线程只读内存； +末段的 token 同样 issue 完就 `Runtime::fetch`（stream 上拷进 pinned），watch 线程在 mark 落地后读出来 +答给 head；fault 由发射线程每个 item `Runtime::fault`（`cuStreamQuery`，不阻塞）问一次。为什么等待 +设备的线程一个 driver 调用都不能做，见 lessons.md 2026-10-08。这换掉了设备上的邮箱核与 `pp.*` op：manifest 里只剩 stage 的签名。没有序号、stop、fault +消息：任何错误进程退出，连接断了对端退出；下一段 60 s 不报"取走"视为死，上一段死了表现为 head 退出。 只做 prefill：prompt 除末 token 按 chunk 走 chunk program，末 token 走一行的 step，首 token 即结束 （`max_tokens` 视为 1）。item 的编号是它在 TCP 上的次序，各段数得一样。每段 issue 完一个 item 就打一个 `Runtime::mark`（event），旁边一个线程等它，落地时记一行 `pp item`：stage、item、rows、`issued_us`、 @@ -583,14 +594,14 @@ stop、fault 消息：任何错误进程退出，连接断了对端退出；设 item k = max(本段 issue k、本段 done k−1、上游 done k)，于是每段每 item 的空等与占用都有了 （`tools/pp_timeline.py`，带 `--gantt`；跨 tray 的段用 `--skew` 扣两台机器的时钟差）。占用里含等下游取走的 时间，要拆出纯计算就拿该段单独跑这个形状的时间（`kern bench` 段 manifest）去比。 -设备上的 fault（邮箱等待超时 trap）在该 item 的 mark 上冒出来，进程退出。下表的 busy / wait 列是当时 +设备上的 fault 在该 item 的 mark 上冒出来，进程退出。下表的 busy / wait 列是当时 设备上 globaltimer 打点量的（`clock` fill，已删：它让输出不再是 forward 的纯函数）。 门禁（qwen3-4b，切在 `l18.qkv_proj` 前，2 段）： - 只切不分进程：`pp_chain` 各段依次跑 vs 整份，7 条 prompt（1–2049 token），chunk 512 × 8 步、 chunk 64 × 16 步（2049 token 32 个 chunk），每步 logits 与 token **逐位同**。 - 两进程 serve：7 条 gate prompt 首 token 与单卡 kern-serve 全同；杀末段 → head 立即退出；冻住末段 - （SIGSTOP）→ head 61 s 后 trap 退出，恢复的末段 61 s 后同样退出。 + （SIGSTOP）→ head 60 s 后退出（等"取走"超时），恢复的末段随后退出。 扫描（chunk 512，`--max-seqs 16`，每格 conc1 4 条 / conc4 16 条不同 prompt；TTFT = max_tokens=1 的 请求延迟；busy/wait 是每 item 均值）： @@ -617,6 +628,14 @@ item k = max(本段 issue k、本段 done k−1、上游 done k)，于是每段 （`~/bench_results/2026-10-07-kern-pp-v1/perlayer/`）。v1 的限制：只认纯 paged 的 manifest； chunk program 自己吐 token 的模型拒切；上游那份 box 是不用的占位。 +传输搬到外壳之后（2026-10-08，`~/bench_results/2026-10-08-pp-inout/`）：门禁同上（tray08：8 条 gate +prompt 首 token 与单卡全同、498 item 两侧一致、杀 / 冻末段同前，head 退出后末段也退出）；K3 PP8 +容量（tray12+17，c16、rate inf、200 条 2k–32k）44.3k tok/s、2.47 req/s（邮箱 44.9–46.8k），每段 +period 175–179 ms 由算力定，host 不在关键路径上。qwen3 扫描中位 TTFT ms（邮箱 → 现在）：conc1 +8.3 / 25.2 / 172.6 → 8.7 / 25.3 / 172.4，conc4 17.6 / 79.6 / 664.7 → 22.7 / 90.1 / 677.4——conc4 还差 +5–13 ms：host 等下一段"取走"再 push 是一次 TCP 往返，而 cuBLAS 的 launch 排在没开的门后会阻塞，门一开 +host 还得把剩下的 launch 发完，设备追上 host 就空转；要关上这段得弄清 cuBLAS 在 launch 里等什么。 + ## 没做（按需要加） span 长的预算 policy（K5 D2：按稠密 / attention 预算定 c，现在是 `--chunk` 上限）、抢占 / 动态页分配、 diff --git a/schema/manifest-v5.schema.json b/schema/manifest-v5.schema.json index 77b3f039..07094369 100644 --- a/schema/manifest-v5.schema.json +++ b/schema/manifest-v5.schema.json @@ -368,6 +368,11 @@ "description": "Read back by the runtime after each run, e.g. `next_token`.", "type": "string" }, + { + "const": "inout", + "description": "Both: holds a value before the run, which the programs may rewrite in place, and is read back after, e.g. the `hidden` a pipeline stage takes from the stage before it and hands to the next.", + "type": "string" + }, { "const": "weight", "description": "Assembled at load time from the checkpoint tensors its `bind` names, e.g. `model.embed_tokens.weight`.", @@ -2300,7 +2305,7 @@ "type": "null" } ], - "description": "Rank groups of a multi-GPU manifest, e.g. `{\"groups\": {\"ep\": 4}}`: every rank of an SPMD group loads the same manifest; the two stages a cut's group joins load their own, agreeing only on the peer's `of` buffer." + "description": "Rank groups a multi-GPU manifest is SPMD over, e.g. `{\"groups\": {\"ep\": 4}}`; every rank loads the same manifest." }, "vars": { "additionalProperties": { diff --git a/tools/kernels-src/pp_mailbox.cu b/tools/kernels-src/pp_mailbox.cu deleted file mode 100644 index e5ba84eb..00000000 --- a/tools/kernels-src/pp_mailbox.cu +++ /dev/null @@ -1,118 +0,0 @@ -// The pipeline edge between two stages of a cut manifest (`kern cut`): one -// mailbox on the downstream GPU, written by the upstream stage through a -// peer pointer. -// -// The box is `u8[256 + data]`. Its header holds two item counters, `full` -// at byte 0 (items the upstream has posted) and `empty` at byte 64 (items -// the downstream has taken out). Both live on the downstream GPU, so the -// downstream spins locally and the upstream reads and writes remote -// memory; neither side keeps a counter of its own. An upstream item: wait -// until `empty == full` (the box is free), put each boundary buffer at its -// offset, post `full + 1`. A downstream item: wait until `full > empty`, -// take each buffer out, post `empty + 1`. Each step is its own launch, so -// stream order puts every copy of an item between its two flags; a post -// fences system-wide before its release store. -// -// A wait gives up after `timeout_ns` of globaltimer and traps: the process -// dies with a launch failure instead of hanging on a dead stage. -// -// kern_pp_wait_empty (in u64 peers[2], i64 timeout_ns) -// kern_pp_put (in u64 peers[2], in u8 src[], i64 bytes, i64 offset) -// kern_pp_post_full (in u64 peers[2]) -// kern_pp_wait_full (in u8 box[], i64 timeout_ns) -// kern_pp_take (in u8 box[], out u8 dst[], i64 bytes, i64 offset) -// kern_pp_post_empty (inout u8 box[]) -// waits and posts: block [1,1,1], grid [1,1,1]; put / take: any grid. -// peers[1] is the downstream box; the upstream's own box (peers[0]) is unused. - -#define PP_HEADER 256 - -typedef unsigned long long u64; - -__device__ __forceinline__ u64 gtimer() { - u64 t; - asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t)); - return t; -} - -__device__ __forceinline__ u64 ld_acquire_sys(const u64* p) { - u64 v; - asm volatile("ld.acquire.sys.global.u64 %0, [%1];" : "=l"(v) : "l"(p) : "memory"); - return v; -} - -__device__ __forceinline__ void st_release_sys(u64* p, u64 v) { - asm volatile("st.release.sys.global.u64 [%0], %1;" ::"l"(p), "l"(v) : "memory"); -} - -__device__ __forceinline__ u64* full_of(const unsigned char* box) { - return reinterpret_cast(const_cast(box)); -} - -__device__ __forceinline__ u64* empty_of(const unsigned char* box) { - return reinterpret_cast(const_cast(box) + 64); -} - -__device__ __forceinline__ unsigned char* downstream(const u64* peers) { - return reinterpret_cast(peers[1]); -} - -__device__ void copy(unsigned char* dst, const unsigned char* src, long long bytes) { - const long long stride = (long long)gridDim.x * blockDim.x; - const long long t = (long long)blockIdx.x * blockDim.x + threadIdx.x; - const long long vecs = bytes / 16; - const uint4* s = reinterpret_cast(src); - uint4* d = reinterpret_cast(dst); - for (long long i = t; i < vecs; i += stride) { - d[i] = s[i]; - } - for (long long i = vecs * 16 + t; i < bytes; i += stride) { - dst[i] = src[i]; - } -} - -extern "C" __global__ void kern_pp_wait_empty(const u64* peers, long long timeout_ns) { - const unsigned char* box = downstream(peers); - const u64 item = ld_acquire_sys(full_of(box)); - const u64 t0 = gtimer(); - while (ld_acquire_sys(empty_of(box)) != item) { - if ((long long)(gtimer() - t0) > timeout_ns) { - printf("kern pp: item %llu waited %lld ns for the downstream stage to take the last one\n", item, - timeout_ns); - __trap(); - } - } -} - -extern "C" __global__ void kern_pp_put(const u64* peers, const unsigned char* src, long long bytes, - long long offset) { - copy(downstream(peers) + PP_HEADER + offset, src, bytes); -} - -extern "C" __global__ void kern_pp_post_full(const u64* peers) { - const unsigned char* box = downstream(peers); - __threadfence_system(); - st_release_sys(full_of(box), ld_acquire_sys(full_of(box)) + 1); -} - -extern "C" __global__ void kern_pp_wait_full(const unsigned char* box, long long timeout_ns) { - const u64 item = ld_acquire_sys(empty_of(box)); - const u64 t0 = gtimer(); - while (ld_acquire_sys(full_of(box)) <= item) { - if ((long long)(gtimer() - t0) > timeout_ns) { - printf("kern pp: item %llu waited %lld ns for the upstream stage\n", item, timeout_ns); - __trap(); - } - } - __threadfence_system(); -} - -extern "C" __global__ void kern_pp_take(const unsigned char* box, unsigned char* dst, long long bytes, - long long offset) { - copy(dst, box + PP_HEADER + offset, bytes); -} - -extern "C" __global__ void kern_pp_post_empty(unsigned char* box) { - __threadfence_system(); - st_release_sys(empty_of(box), *empty_of(box) + 1); -}