diff --git a/Cargo.lock b/Cargo.lock index 95c51517..d661dfcf 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2293,11 +2293,13 @@ version = "0.2.3" dependencies = [ "anyhow", "clap", + "hex", "kern-manifest", "kern-pool", "kern-run", "kern-runtime", "pegainfer-frontend", + "serde", "serde_json", "tokio", "tracing", 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..7b2404c2 --- /dev/null +++ b/crates/kern-manifest/src/cut.rs @@ -0,0 +1,362 @@ +//! Cutting a manifest into a pipeline: stage `s` runs the calls between +//! 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 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 +//! 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, 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 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. + +use std::collections::{BTreeMap, BTreeSet}; + +use crate::types::*; +use crate::verify::{verify, Verified}; + +/// 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 {} + +/// Cut `m` before each call index in `at` of every program with a +/// `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()); + } + 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())); + } + } + } + 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()); + + // 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]); + for b in live_in(m, &p.calls[a..]).into_iter().filter(|b| before.contains(b)) { + 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", + m.buffers[b].kind + )); + } + crossing[e].insert(b); + } + } + 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 + )); + } + } + } + } + // 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 { + 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 id = id_of(m, at); + (0..stages) + .map(|s| { + let programs = driven + .iter() + .map(|(name, p)| ((*name).clone(), Program { calls: p.calls[span(p, s)].to_vec(), ..(*p).clone() })) + .collect(); + 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() +} + +/// 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, + id: &str, +) -> 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 || 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)) { + 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 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 { + if let Some(b) = m.buffers.get(buf) { + buffers.insert(buf.clone(), b.clone()); + } + } + } + } + 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(); + // 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() { + 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 = 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) + .map(|n| (n.to_string(), m.modules[n].clone())) + .collect(); + // 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: None, + 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)), + _ => {} + } +} + +/// 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. +fn reads<'a>(m: &'a Manifest, c: &'a Call) -> impl Iterator { + accesses(m, c).into_iter().filter(|(_, _, d)| *d != Dir::Out).map(|(b, _, _)| b) +} + +/// 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 +} + +/// 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 { + 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..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 @@ -45,9 +48,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 +131,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, } @@ -203,6 +208,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 { @@ -223,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(); @@ -237,76 +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 - } - } + // 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 }) }; - let rows = axis_of(Fill::Slot, &mut errs); - let groups = axis_of(Fill::SeqLen, &mut errs); - 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]" @@ -314,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, @@ -347,9 +323,6 @@ 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) { - 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()); } @@ -390,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; @@ -409,6 +382,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 { @@ -434,8 +408,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 } @@ -599,6 +576,7 @@ impl Protocol { line_tables, forwards, once, + derive, }) } else { Err(ProtocolErrors(errs)) @@ -641,18 +619,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 } @@ -668,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 0800af6e..4aece24c 100644 --- a/crates/kern-manifest/src/types.rs +++ b/crates/kern-manifest/src/types.rs @@ -74,6 +74,9 @@ pub struct Manifest { /// 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}`. + #[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, @@ -148,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, @@ -226,12 +232,49 @@ 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)] 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 { @@ -353,7 +396,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, } @@ -577,8 +620,12 @@ 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. + 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. @@ -592,7 +639,9 @@ 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", BufferKind::Carry => "carry", BufferKind::Peer => "peer", diff --git a/crates/kern-manifest/src/verify.rs b/crates/kern-manifest/src/verify.rs index f00f0fe3..0a4928b2 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 @@ -25,9 +25,12 @@ //! 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 (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 //! 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 @@ -37,6 +40,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 @@ -184,6 +188,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(); @@ -195,6 +205,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) => { @@ -270,12 +288,15 @@ fn diagnostics(m: &Manifest) -> Vec { errs.push(format!("{ctx}: shared host weights cannot select tensors, rows or columns by rank")); } } - // 9c. bind: a weight is its segments, nothing else has any + // 9c. bind: a weight or a source is its segments, nothing else has any + if b.kind == BufferKind::Source && b.export { + errs.push(format!("{ctx}: a source lives only while weights load; it cannot be exported")); + } match (b.kind, b.bind.is_empty()) { - (BufferKind::Weight, true) => { - errs.push(format!("{ctx}: a weight buffer binds at least one checkpoint tensor (`bind`)")) + (BufferKind::Weight | BufferKind::Source, true) => { + errs.push(format!("{ctx}: a {} buffer binds at least one checkpoint tensor (`bind`)", b.kind)) } - (BufferKind::Weight, false) => { + (BufferKind::Weight | BufferKind::Source, false) => { for (i, s) in b.bind.iter().enumerate() { let sctx = format!("{ctx}: bind[{i}]"); let source = |t: &TensorSource, errs: &mut Vec, used_groups: &mut BTreeSet| match t @@ -340,9 +361,8 @@ fn diagnostics(m: &Manifest) -> Vec { } } } - (kind, false) => { - errs.push(format!("{ctx}: `bind` names checkpoint tensors for a weight, not a {kind} buffer")) - } + (kind, false) => errs + .push(format!("{ctx}: `bind` names checkpoint tensors for a weight or a source, not a {kind} buffer")), (_, true) => {} } // 10b. export / peer @@ -763,7 +783,15 @@ 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::Inout + | BufferKind::Weight + | BufferKind::Source + | BufferKind::Carry + | BufferKind::Peer + ) }) .map(|(n, _)| n.as_str()) .collect(); @@ -775,6 +803,11 @@ fn diagnostics(m: &Manifest) -> Vec { "program `{pname}`: `once` (run after load) and `batch` (driven per step) are exclusive" )); } + if p.derive && (p.once || p.batch.is_some() || p.graph) { + errs.push(format!( + "program `{pname}`: `derive` (run inside weight loading) excludes `once`, `batch` and `graph`" + )); + } if p.graph && p.batch.is_none() { errs.push(format!("program `{pname}`: `graph` (captured per call shape) needs a `batch` (a program a serving loop drives)")); } @@ -820,6 +853,12 @@ fn diagnostics(m: &Manifest) -> Vec { }; used_ops.insert(c.op.clone()); let has_extern = op.imp.launches.iter().any(|l| l.is_extern()); + if p.derive && op.imp.launches.iter().any(|l| l.entry().starts_with("extern:nccl")) { + errs.push(format!( + "{ctx}: a `derive` program runs before any group is joined; op `{}` is a collective", + c.op + )); + } if c.args.len() != op.params.len() { errs.push(format!("{ctx}: op `{}` takes {} params, got {} args", c.op, op.params.len(), c.args.len())); @@ -870,6 +909,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::Inout | BufferKind::Peer) => errs.push(format!( + "{actx}: a `derive` program runs inside weight loading, before any {} buffer such as `{buf}` means anything", + b.kind + )), + _ => {} + } if b.kind == BufferKind::Peer && has_extern { errs.push(format!( "{actx}: peer buffer `{buf}` passed to op `{}`, which has an extern launch; runtime built-ins never receive peer memory", @@ -877,7 +926,10 @@ fn diagnostics(m: &Manifest) -> Vec { )); } if matches!(dir, Dir::Out | Dir::InOut) { - if matches!(b.kind, BufferKind::Input | BufferKind::Weight | BufferKind::Peer) { + if matches!( + b.kind, + BufferKind::Input | BufferKind::Weight | BufferKind::Source | BufferKind::Peer + ) { errs.push(format!("{actx}: op writes to read-only {} buffer `{buf}`", b.kind)); } written.insert(buf.as_str()); @@ -889,6 +941,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/cut.rs b/crates/kern-manifest/tests/cut.rs new file mode 100644 index 00000000..b75a7aa8 --- /dev/null +++ b/crates/kern-manifest/tests/cut.rs @@ -0,0 +1,282 @@ +//! 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, BufferKind, Dir, Manifest, ParamType}; +use kern_manifest::{verify, Protocol, Verified}; +use serde_json::{json, Value}; + +const QWEN3: &str = include_str!("../../../examples/qwen3-4b.json"); + +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>, +} + +impl Machine { + /// Run `program` of `m`: every op a hash of what it reads into what it + /// writes. + fn run(&mut self, m: &Manifest, program: &str) { + for c in &m.programs[program].calls { + let label = c.label.clone().unwrap(); + 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"); + } + 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"); + + let stages = match cut(&m, &at) { + Ok(s) => s, + Err(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 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); + // 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:?}"); + } + assert!(cuts > 200 && refused > 0, "{cuts} cuts, {refused} refused"); +} + +fn qwen3() -> Verified { + verify(Manifest::from_json(QWEN3).unwrap()).unwrap() +} + +/// 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] +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]).unwrap_or_else(|e| panic!("{e}")); + let (head, tail) = (&stages[0], &stages[1]); + // 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") + && !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")); + // 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}")); + 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]).unwrap(); + assert_eq!(again[1].to_json(), tail.to_json()); + + 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)); +} + +#[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).is_err(), "{at:?}"); + } +} 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 d6d68e9c..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))); @@ -112,7 +116,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] @@ -196,37 +200,46 @@ 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()]); - 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] 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] @@ -254,7 +267,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-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/crates/kern-manifest/tests/verify.rs b/crates/kern-manifest/tests/verify.rs index 899915a7..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(); @@ -1091,3 +1099,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/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/examples/pp_chain.rs b/crates/kern-run/examples/pp_chain.rs new file mode 100644 index 00000000..8838f21b --- /dev/null +++ b/crates/kern-run/examples/pp_chain.rs @@ -0,0 +1,252 @@ +//! 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, 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; + +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::{arrivals, fills, run_once, write_named, Vars}; +use kern_runtime::{Capacity, Mapped, Runtime}; + +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 capacity = Capacity { tokens: Some(o.capacity), seqs: 2 }; + 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)) +} + +/// 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)?) +} + +/// 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")?, + 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 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 { + 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)?; + chain(&mut stages, &boxes, 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])?; + 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]; + 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..1ff4e182 100644 --- a/crates/kern-run/src/bin/kern.rs +++ b/crates/kern-run/src/bin/kern.rs @@ -71,6 +71,19 @@ 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, + /// Directory the stages are written to + #[arg(long)] + out: PathBuf, + }, } fn main() -> Result<()> { @@ -85,7 +98,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,7 +131,46 @@ fn main() -> Result<()> { } Ok(()) } + 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], 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 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: {}; crossing: {})", + path.display(), + st.buffers.values().filter(|b| b.kind == kern_manifest::types::BufferKind::Weight).count(), + calls.join(", "), + crossing.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(()) } /// The target a command runs on: the named one, else the only one, else diff --git a/crates/kern-run/src/lib.rs b/crates/kern-run/src/lib.rs index 4f63c713..667ef98c 100644 --- a/crates/kern-run/src/lib.rs +++ b/crates/kern-run/src/lib.rs @@ -21,13 +21,14 @@ 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; use anyhow::{bail, ensure, Context, Result}; -use kern_manifest::protocol::{Axis, Forward, LineTable, PageTable, Rows}; -use kern_manifest::types::Fill; +use kern_manifest::protocol::{Axis, Filled, Forward, LineTable, PageTable, Rows}; +use kern_manifest::types::{BufferKind, Fill}; use kern_manifest::{Protocol, Verified}; use kern_pool::Lease; use kern_runtime::{connect_nccl, GroupRank, PeerHandle, Runtime, Topology}; @@ -94,8 +95,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 +135,84 @@ 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: `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 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, L: Borrow>( + p: &'p Protocol, + leases: &[L], + positions: &[usize], + lens: &[usize], + ids: &[i64], +) -> (Vars, Vec<(&'p Filled, Vec)>) { + let b = leases.len(); + 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() + .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) => starts.iter().map(|&s| ids[s]).collect(), + (Fill::Token, _) => ids.to_vec(), + (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, _) => 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, 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) => i32s(values), + (None, false) => return Ok(false), + }; + rt.write_input_at(name, &bytes, vars)?; + 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, p: &Protocol, @@ -147,23 +221,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, &vec![rows; leases.len()], ids); + for (f, v) in values { rt.write_input_at(&f.name, &f.encode(&v), &vars)?; } Ok(vars) @@ -197,7 +256,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}`"))) } @@ -420,9 +479,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-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/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 b0b3184e..8f138553 100644 --- a/crates/kern-runtime/src/exec.rs +++ b/crates/kern-runtime/src/exec.rs @@ -19,18 +19,282 @@ //! 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::ffi::CStr; use std::os::raw::c_void; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::{Arc, Mutex}; -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::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`] 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 { + seq: u64, + passed: Arc, +} + +/// 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<()> { + 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 Signals { + fn drop(&mut self) { + 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 { + 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()?; + 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<()> { 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 0d79822a..38bc2d65 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; @@ -57,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::{Fetch, Gate, Mark}; pub use harness::Scratch; pub use host::HostRegion; pub use host_weights::HostWeights; @@ -133,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, @@ -179,6 +182,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 +213,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. @@ -260,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 { @@ -283,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 899f272e..e527c200 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,9 +185,16 @@ 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 { + 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); @@ -276,11 +289,18 @@ 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)?; 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, @@ -311,6 +331,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 +343,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 +386,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 +399,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 +416,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 +503,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 +535,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 +564,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 +579,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 +602,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/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-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-serve/Cargo.toml b/crates/kern-serve/Cargo.toml index 13eeae10..8a948f1d 100644 --- a/crates/kern-serve/Cargo.toml +++ b/crates/kern-serve/Cargo.toml @@ -13,6 +13,8 @@ pegainfer-frontend.workspace = true anyhow.workspace = true clap.workspace = true +hex.workspace = true +serde.workspace = true serde_json.workspace = true tokio.workspace = true tracing.workspace = true diff --git a/crates/kern-serve/src/lib.rs b/crates/kern-serve/src/lib.rs index a97a4aab..756910d7 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::{self, ParserSelection}; use tracing::info; @@ -86,6 +87,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<()> { @@ -98,52 +108,94 @@ pub fn serve(o: ServeOpts, art: Artifacts) -> Result<()> { "loading" ); let gpus = if o.gpus.is_empty() { vec![0] } else { o.gpus.clone() }; - let model_path = weights.dirs().remove(0); - 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 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, + }; + 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 }; + 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, 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) +} - // The scheduler thread owns the tray for its whole life: load there, - // report readiness, then drive. +/// 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 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, rows: o.rows, host_bytes }; 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 @@ -169,7 +221,7 @@ pub fn serve(o: ServeOpts, art: Artifacts) -> Result<()> { &model_path, vec![served_name], ParserSelection::default(), - o.port, + port, None, 1, shutdown, diff --git a/crates/kern-serve/src/pp/head.rs b/crates/kern-serve/src/pp/head.rs new file mode 100644 index 00000000..9cb6b6ec --- /dev/null +++ b/crates/kern-serve/src/pp/head.rs @@ -0,0 +1,413 @@ +//! The head: stage 0 of a pipeline, the scheduler every stage follows. +//! +//! 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. 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 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}; +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, StopPolicy, +}; +use tracing::{error, info}; + +use super::pack::{pack, Bounds, Waiting}; +use super::{recv, send, Flow, Hello, Item, Next, Table}; +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, +} + +/// An admitted prompt with rows still to send. +struct Active { + id: RequestId, + lease: Lease, + ids: Vec, + pos: usize, + stop: StopPolicy, +} + +/// 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, + stop: StopPolicy, +} + +pub struct PpHead { + rt: Runtime, + p: Protocol, + prefill: Forward, + /// The one-row step the last prompt token goes through; none when the + /// chunk program hands the first token back itself. + step: Option, + bounds: Bounds, + /// Items in flight below which a partial item goes out: the stages. + depth: usize, + max_seqs: usize, + stages: Vec, + replies: mpsc::Receiver>, + waiting: VecDeque, + active: VecDeque, + /// Per item in flight, oldest first, the prompts its reply finishes. + due: VecDeque>, + flow: Flow, +} + +/// Wait for every other stage, load stage 0 at the pages they all have, +/// 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 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(); + while (hellos.len() as u64) < cut.stages - 1 { + 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 `{}`", + 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"); + 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 }), None)?; + rt.set_eager(h.eager); + weights.bind(&mut rt)?; + let p = Protocol::check_unsampled(&rt.manifest)?; + 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 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 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, (_, mut w, _)), table) in hellos.into_iter().zip(tables) { + send(&mut w, &table)?; + let tx = (s == cut.stages - 1).then(|| tx.clone()); + 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, + "pipeline head ready"); + Ok(PpHead { + rt, + p, + prefill, + step, + bounds, + depth: cut.stages as usize, + max_seqs: h.max_seqs, + stages, + replies, + waiting: VecDeque::new(), + active: VecDeque::new(), + due: VecDeque::new(), + 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. +fn listen(stage: u64, mut r: TcpStream, tokens: Option>>) { + let why = loop { + match (recv::>(&mut r), &tokens) { + (Ok(Some(t)), Some(tx)) => { + if tx.send(t).is_err() { + return; + } + } + (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:#}"), + } + }; + 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(), + } + } + + /// 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; + if ledger.is_aborted(id) { + ledger.retire(id); + self.waiting.pop_front(); + continue; + } + if self.live() >= 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); + let ids = q.request.prompt_tokens.iter().map(|&t| t as i64).collect(); + let stop = q.request.stop_policy.clone(); + self.active.push_back(Active { id, lease, ids, pos: 0, stop }); + 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 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)?; + } + 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, stop: a.stop }; + self.send_item(&f, &[(&finish._lease, pos, &a.ids[pos..])])?; + self.due.push_back(vec![finish]); + } + Ok(()) + } + + /// 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, stop: a.stop } + }) + .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. + 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 * 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()); + } + let item = Item { program: f.name.clone(), vars, rows }; + for w in &mut self.stages { + 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.flow.issued(&self.rt, &item.vars, ids.len() as u64, None)?; + Ok(()) + } + + fn advance(&mut self, ledger: &mut RequestLedger) -> Result<()> { + self.admit(ledger)?; + self.issue()?; + self.collect(ledger) + } + + /// 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() { + Ok(t) => t, + Err(mpsc::TryRecvError::Empty) => return Ok(()), + Err(mpsc::TryRecvError::Disconnected) => bail!("the last stage's connection is gone"), + }; + 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 r.stop.classify(tok, |_| false).is_some() { 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.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 new file mode 100644 index 00000000..4cc2d875 --- /dev/null +++ b/crates/kern-serve/src/pp/mod.rs @@ -0,0 +1,349 @@ +//! 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. +//! +//! 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 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::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 kern_manifest::protocol::Filled; +use kern_manifest::Protocol; +use kern_runtime::{Fetch, Gate, Mapped, Mark, PeerHandle, Runtime}; +use serde::{Deserialize, Serialize}; + +/// 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)] +pub struct Hello { + pub stage: u64, + pub cut: String, + pub pages: u64, + /// 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 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 pages: u64, + pub next: Option<(String, Handles)>, +} + +/// 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 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| { + 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 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") +} + +/// 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)?)) +} + +/// 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}")), + } + } +} + +/// 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"), + } +} + +/// 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, +} + +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)) + } +} + +/// The edge into a stage: how many items the stage before reports +/// arrived, as a gate the device waits on. +struct Prev { + arrived: Arc, +} + +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)) + } +} + +/// 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, 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 { + issued: Option>, +} + +struct Issued { + item: u64, + rows: u64, + issued_us: u64, + mark: Mark, + answer: Option, +} + +impl Watch { + 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 { issued: Some(issued) } + } + + 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"); + } + Ok(()) + } +} + +fn now_us() -> u64 { + SystemTime::now().duration_since(UNIX_EPOCH).map_or(0, |d| d.as_micros() as u64) +} + +/// 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.issued.take()); + } +} 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 new file mode 100644 index 00000000..04863f61 --- /dev/null +++ b/crates/kern-serve/src/pp/stage.rs @@ -0,0 +1,128 @@ +//! A downstream stage: loads its slice of the model, says hello to the +//! 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}; + +use anyhow::{bail, ensure, Context, Result}; +use kern_manifest::{Protocol, Verified}; +use kern_run::{Vars, Weights}; +use kern_runtime::{Capacity, Runtime}; +use tracing::{error, info}; + +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 { + 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), None)?; + 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 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() + ); + 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 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::(); + 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; + } + } + Ok(None) => { + info!("the head closed its connection"); + return; + } + Err(e) => { + error!("head connection: {e:#}"); + return; + } + } + } + }); + + 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))?; + 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(()) +} diff --git a/crates/kern-serve/src/tray.rs b/crates/kern-serve/src/tray.rs index 97015d0d..3bedb8cc 100644 --- a/crates/kern-serve/src/tray.rs +++ b/crates/kern-serve/src/tray.rs @@ -766,8 +766,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/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/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/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 a5d1e616..c3618bb2 100644 --- a/docs/manifest.md +++ b/docs/manifest.md @@ -14,8 +14,9 @@ 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 / inout / weight / source / workspace / carry / peer topology.groups. group 多卡 SPMD 的 rank 组：只有名字和大小 +cut cut 流水线的一段：id、第几段、共几段 ``` `call` 调 `op`，`op` 由若干 `launch` 实现，`launch` 起 `module` 里的 @@ -37,7 +38,9 @@ topology.groups. group 多卡 SPMD 的 rank 组：只有名字和 - `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， @@ -50,7 +53,8 @@ topology.groups. group 多卡 SPMD 的 rank 组：只有名字和 - `buffers`：`dtype + shape + kind`，可选 `domain`（内容的先验，见下）与 `fill`（在 serving 循环里的角色，见「Serving 协议」）。shape 维度是常量或 var 名；kind 说 的是"谁供应、活多久"：`input`（runtime 写入）/ `output`（runtime 读回） - / `weight`（load 时从 checkpoint 的张量拼出来，`bind` 说是哪些，见「权重」）/ `workspace`（runtime 规划，跨次执行 + / `weight`（load 时从 checkpoint 的张量拼出来，`bind` 说是哪些，见「权重」）/ `source`（同样 `bind` + checkpoint 张量，但只活在 load 里、只给 `derive` program 读，见「权重」）/ `workspace`（runtime 规划，跨次执行 不保留）/ `carry`（一个 program 写、另一个 program 读的交接棒，跨次 执行保留；谁先跑是 caller 契约，verifier 只要求它被某个 program 写到 ——投机解码的 aux 隐状态逼出来的）/ `peer`（runtime 填的地址数组，见 @@ -68,6 +72,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， @@ -134,8 +142,8 @@ topology.groups. group 多卡 SPMD 的 rank 组：只有名字和 vLLM attention（unified + reduce_segments）这类"一个逻辑算子 = 多次 launch + 私有中间缓冲"整体折叠成一个 impl，不向调用方泄漏。 - `programs`：每个 program（如 `prefill`/`decode`）是 - `{"batch"?, "once"?, "calls": [...]}`——`batch` / `once` 说它怎么被 - serving 循环调（见「Serving 协议」），`calls` 是一段顺序 call 列表： + `{"batch"?, "once"?, "derive"?, "calls": [...]}`——`batch` / `once` / + `derive` 说它怎么被 serving 循环或 load 调（见「Serving 协议」），`calls` 是一段顺序 call 列表： `op` 名 + 接口实参（buffer/state 实参可带字节 `offset`，默认 0：kernel 收到 base+offset——provider 用它寻址融合 buffer 里的视图如 qkv 的 q/k/v 切片、state 里的逐层区域，offset 是 provider 布局算术的字面量， @@ -316,6 +324,17 @@ cos / sin 表、全 1 的 kv scale、chunk 索引表——不是权重，是 `ca `tools/kernels-src/weight_prep.cu` 的几个 elementwise 核）。拼接方式与派生 表都是模型知识，住在生成器（`tools/qwen_weights.py`），不住在 runtime。 +checkpoint 张量要先变换才能用（K3 的 mxfp4 专家按 GEMM 布局 shuffle），而原样 +那份没人再读时，原样的张量声明成 `source`，变换写成 `derive` program： +runtime 在 `load_weights` 里逐 call 跑它，每个 source 只在第一个读它的 call +之前上传、最后一个读它的 call 之后让出——所有 source 共用一块暂存，同一 +call 上活着的 source 互不重叠（`kern-runtime/src/derive.rs` 的 `layout`）。 +显存峰值是派生结果加上同时活着的最宽一组 source（K3 全量一层的 w13 ≈ 10 +GB），不是整份 checkpoint 两遍。跑完暂存释放、derive program 一起丢掉， +不能再跑；同一个 runtime 不能第二次 `load_weights`。derive program 不碰 +state / input / output / peer、不调 collective（它跑在这些都还没意义的时候）， +与 `once` / `batch` / `graph` 互斥。 + ## Serving 协议：`fill`、`batch`、`once` runtime 之上还有一层契约：一个 serving 循环（`kern run`、kern-serve 的 @@ -350,6 +369,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）；行数随调用变的 @@ -360,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 @@ -385,6 +409,46 @@ kern-serve 只拿 `Protocol` 驱动 runtime，不读 JSON（CI 用 grep 保证 段代码，`--rows` 选每组行数（缺省取 manifest 声明的最宽）。设计与取舍见 [v4-design.md](v4-design.md)。 +## 流水线切分：`kern cut` + +`kern cut --at [,…] --out ` 把一份单卡 manifest +在 call 下标处切成 S 段（`kern_manifest::cut::cut` 是纯函数，不碰 GPU）。 +**一段就是一个函数**：从上一段来的值是它的 input，给下一段的值是它的 +output，manifest 里没有任何流水线专有的东西；怎么传、何时传是 serving +外壳的事（serve.md「流水线」）。 + +- 每个 batch program 在同一下标处切开（下标须小于每个 batch program 的 + 长度）；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，按等大估字节），不满的就是给别段占着的内存（报告，不拒）。 + 所以生成器把 KV 写成每层一个 state（qwen3-4b：`kv.0`…`kv.35`），各段只带 + 自己那些层；生成器不知道流水线。页表 / slot 的 `index_into` 指向一个本段 + 丢掉的 state 时改指本段留下的同类 state（pool 给所有 paged state 同一套 + 页号）。 +- 切分只改签名，不加 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` 把各段在同一进程里串起来（下一段 +到达的 buffer 映射进上一段、每个 call 后 `Runtime::push`）， +logits 与 token 与整份逐位同（serve.md「流水线」）。 + ## Verifier（`kern-manifest`） `verify()` 收集全部错误一次报告（`VerifyErrors`）： @@ -410,7 +474,8 @@ kern-serve 只拿 `Protocol` 驱动 runtime，不读 JSON（CI 用 grep 保证 能接 `state`、var/表达式的取值范围必须装进标量参数类型、offset 对齐 且在界内）； 9. 逐 program 数据流：禁止读未写（read-before-write）、禁止写 input/ - weight；output / carry 必须被**某个** program 写到（prefill 这类只落 + weight/source；source 只许 `derive` program 读，`derive` program 不碰 + state / input / output / peer、不调 collective；output / carry 必须被**某个** program 写到（prefill 这类只落 state 的 program 合法地不写任何 output；carry 在每个 program 内视为 已写——它的生产者是另一个 program）；program 的 `once` 与 `batch` 互斥，`batch.groups ≥ 1`、`rows` 是 ≥ 1 的常量或已声明的 var（fill 之间 @@ -423,7 +488,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/docs/roadmap.md b/docs/roadmap.md index 6b1a5784..c61b49f0 100644 --- a/docs/roadmap.md +++ b/docs/roadmap.md @@ -59,11 +59,22 @@ 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 邮箱核）+ 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 ✅ | 一段就是一个函数：跨段的 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` 时间线 | + ## 协议线 | 级 | 内容 | 门禁 | |---|---|---| | 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/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/docs/serve.md b/docs/serve.md index cb290833..e0f3e79c 100644 --- a/docs/serve.md +++ b/docs/serve.md @@ -584,6 +584,75 @@ 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, 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`、 +`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 在该 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 60 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）。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 是不用的占位。 + +传输搬到外壳之后（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/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 ed29590d..6aa51c5f 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, @@ -34,15 +69,120 @@ }, "vars": { "tokens": { - "max": "MAX_TOKENS" + "max": "MAX_TOKENS", "axis": "rows" }, "seqs": { - "max": "MAX_SEQS" + "max": "MAX_SEQS", "axis": "groups" } }, "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..9eec2107 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, @@ -27,15 +62,120 @@ }, "vars": { "tokens": { - "max": "MAX_TOKENS" + "max": "MAX_TOKENS", "axis": "rows" }, "seqs": { - "max": "MAX_SEQS" + "max": "MAX_SEQS", "axis": "groups" } }, "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..39d017e6 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, @@ -27,15 +62,120 @@ }, "vars": { "tokens": { - "max": "MAX_TOKENS" + "max": "MAX_TOKENS", "axis": "rows" }, "seqs": { - "max": "MAX_SEQS" + "max": "MAX_SEQS", "axis": "groups" } }, "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/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 10abf59b..07094369 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" }, @@ -368,11 +368,21 @@ "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`.", "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`.", @@ -419,6 +429,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": [ @@ -1633,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" @@ -2061,6 +2123,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": [ { @@ -2083,6 +2156,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.", @@ -2154,6 +2247,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" 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 03d673ab..006c3e4f 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 @@ -1055,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/kern_manifest.py b/tools/kern_manifest.py index b0035fa0..52dc6fbe 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 + 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, 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 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() diff --git a/tools/test_qwen_constants.py b/tools/test_qwen_constants.py index 2b67087d..d3904c70 100644 --- a/tools/test_qwen_constants.py +++ b/tools/test_qwen_constants.py @@ -9,14 +9,16 @@ 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; 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": "444001bf6fa6991128128a878c0cc139b8ecf71800c6ccfea0d344c1ad169133", - "qwen3-4b-silu-mined.json": "e14198af226849f46826bc624c4a6e8a208a4e8d7b160ffc611b817d2915d7b8", - "qwen3-4b.json": "792666b476938520a21598bec6bdf5461cb62dd8fdf52cd5998b33a7388c2e69", - "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" } 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}},