diff --git a/crates/metal/build.rs b/crates/metal/build.rs index 0e08d79bb..b6aa25132 100644 --- a/crates/metal/build.rs +++ b/crates/metal/build.rs @@ -26,7 +26,7 @@ mod geometry { /// The gated delta step's geometry, from the file the host module reads. #[allow( dead_code, - reason = "the host module reads every constant; the header emits the ring length" + reason = "the host module reads every constant; the header emits the ring length and the copies" )] mod delta_geom { include!("src/kernels/delta_geom.rs"); @@ -100,6 +100,7 @@ fn geometry_header() -> String { #define BLOOMERY_DELTA_COLS_PER_BLOCK {per}\n\ // The conv ring, from src/kernels/delta_geom.rs.\n\ #define BLOOMERY_DELTA_RING_ROWS {ring}\n\ + #define BLOOMERY_DELTA_SNAPS {snaps}\n\ // The router, from src/kernels/router_geom.rs.\n\ #define BLOOMERY_ROUTER_ROWS {rrows}\n\ #define BLOOMERY_ROUTER_EXPERTS {rexperts}\n\ @@ -137,6 +138,7 @@ fn geometry_header() -> String { blocks = geometry::PARTIALS_PER_HEAD, per = geometry::HEAD_COLS / geometry::PARTIALS_PER_HEAD, ring = delta_geom::RING_ROWS, + snaps = delta_geom::SNAPS, rrows = router_geom::ROWS, rexperts = router_geom::EXPERTS, rk = router_geom::K, diff --git a/crates/metal/kernels/delta.metal b/crates/metal/kernels/delta.metal index a1265aeae..3a930e56a 100644 --- a/crates/metal/kernels/delta.metal +++ b/crates/metal/kernels/delta.metal @@ -43,6 +43,13 @@ // ring is committed by gdn_ring_rows: the last min(m, ring) positions land in // their slots, after the prepared rows were read. // +// A rows launch may also leave copies after rows inside the chunk (a prompt +// pass's marks): gdn_rows_snap is gdn_rows that writes the lane's registers +// after each named row into that copy's lane, and gdn_ring_snap writes the +// ring a chunk ending after that row would leave into that copy's ring, read +// before gdn_ring_rows overwrites the slots. A copy is the bits a chunk ending +// at its row leaves; the chunk's own outputs do not move. +// // The state is read from the lane the word at `lane_at` names and written to // the lane the word at `lane_to` names. The two words may name one lane: the // step then updates it in place. Two lanes apart, the step leaves the lane it @@ -119,6 +126,18 @@ struct DeltaRingArgs { }; static_assert(sizeof(DeltaRingArgs) == 16, "DeltaRingArgs is 16 bytes: the host packs four words"); + +// The copies a rows launch leaves: how many, the floats of one copy's lane +// and of one copy's ring, and the row each copy follows, ascending. +struct DeltaSnapArgs { + uint n; + uint lane_len; + uint ring_len; + uint after[BLOOMERY_DELTA_SNAPS]; +}; + +static_assert(sizeof(DeltaSnapArgs) == 4 * (3 + BLOOMERY_DELTA_SNAPS), + "DeltaSnapArgs is three words and a row word a copy: the host packs them"); static_assert(BLOOMERY_DELTA_HEAD_COLS / BLOOMERY_DELTA_COLUMN_BLOCKS == BLOOMERY_DELTA_COLS_PER_BLOCK, "a column block is the head divided by the block count"); @@ -724,17 +743,15 @@ kernel void gdn_prep_rows(device const float *x [[buffer(0)]], // in the lane's registers between rows, so m rows leave the bits m stepped // launches leave. A row reads only its prepared row, and writes the column // outputs into the numerator window (gdn_post_rows finishes them). The -// simdgroups of a threadgroup share nothing, so there is no barrier. -kernel void gdn_rows(device const float *prep [[buffer(0)]], - device const uint *lane_buf [[buffer(1)]], - device float *state [[buffer(2)]], - device float *g [[buffer(3)]], - constant DeltaRowsArgs &args [[buffer(4)]], - device atomic_uint *faults [[buffer(5)]], - uint3 tgp [[threadgroup_position_in_grid]], - uint3 tpt [[thread_position_in_threadgroup]], - ushort tiisg [[thread_index_in_simdgroup]], - ushort sgitg [[simdgroup_index_in_threadgroup]]) { +// simdgroups of a threadgroup share nothing, so there is no barrier. With +// `snaps`, the registers after each row it names land in that copy's lane of +// `snap`, as the lane write after the last row would land them: a copy is the +// state a chunk ending at its row leaves, and no output of the chunk moves. +inline void gdn_rows_run(device const float *prep, device const uint *lane_buf, + device float *state, device float *g, constant DeltaRowsArgs &args, + device atomic_uint *faults, device float *snap, + constant DeltaSnapArgs *snaps, uint3 tgp, uint3 tpt, ushort tiisg, + ushort sgitg) { GdnSeat seat = gdn_seat(tgp, tpt, tiisg, sgitg, args.step.n_k, args.step.grouped); uint n_v = BLOOMERY_DELTA_VALUE_HEADS; uint head = BLOOMERY_DELTA_HEAD_COLS; @@ -744,6 +761,9 @@ kernel void gdn_rows(device const float *prep [[buffer(0)]], bool lane_ok = word < args.step.lanes && word_to < args.step.lanes; uint base = ((word * n_v + seat.h) * head + seat.col) * head; uint base_to = ((word_to * n_v + seat.h) * head + seat.col) * head; + uint cell = (seat.h * head + seat.col) * head; + uint copies = snaps == nullptr ? 0u : snaps->n; + uint copy = 0u; float st[16]; gdn_load(state, base, seat.j, lane_ok, st); @@ -780,6 +800,15 @@ kernel void gdn_rows(device const float *prep [[buffer(0)]], } gdn_settle(seat, faults, args.step.layer, row_ok, conv_bad, column_bad, lane_state_bad, oval, sp, st); + if (copy < copies && r == snaps->after[copy]) { + if (lane_ok) { + device float *lane = snap + (ulong)copy * snaps->lane_len + cell; + for (uint i = 0u; i < 16u; ++i) { + lane[key_of(seat.j, i)] = st[i]; + } + } + ++copy; + } } if (lane_ok) { for (uint i = 0u; i < 16u; ++i) { @@ -788,6 +817,37 @@ kernel void gdn_rows(device const float *prep [[buffer(0)]], } } +kernel void gdn_rows(device const float *prep [[buffer(0)]], + device const uint *lane_buf [[buffer(1)]], + device float *state [[buffer(2)]], + device float *g [[buffer(3)]], + constant DeltaRowsArgs &args [[buffer(4)]], + device atomic_uint *faults [[buffer(5)]], + uint3 tgp [[threadgroup_position_in_grid]], + uint3 tpt [[thread_position_in_threadgroup]], + ushort tiisg [[thread_index_in_simdgroup]], + ushort sgitg [[simdgroup_index_in_threadgroup]]) { + gdn_rows_run(prep, lane_buf, state, g, args, faults, nullptr, nullptr, tgp, tpt, tiisg, + sgitg); +} + +// gdn_rows with the copies `snaps` names: copy k's lane is `lane_len` floats +// of `snap` from k·lane_len, laid out as one lane of the state. +kernel void gdn_rows_snap(device const float *prep [[buffer(0)]], + device const uint *lane_buf [[buffer(1)]], + device float *state [[buffer(2)]], + device float *g [[buffer(3)]], + constant DeltaRowsArgs &args [[buffer(4)]], + device atomic_uint *faults [[buffer(5)]], + device float *snap [[buffer(6)]], + constant DeltaSnapArgs &snaps [[buffer(7)]], + uint3 tgp [[threadgroup_position_in_grid]], + uint3 tpt [[thread_position_in_threadgroup]], + ushort tiisg [[thread_index_in_simdgroup]], + ushort sgitg [[simdgroup_index_in_threadgroup]]) { + gdn_rows_run(prep, lane_buf, state, g, args, faults, snap, &snaps, tgp, tpt, tiisg, sgitg); +} + // The numerators and the partial sums of a chunk of rows, from the column // outputs gdn_rows left in the numerator window. Grid: (value heads, rows) // threadgroups of 128 threads, thread t the column t of the head. The @@ -833,3 +893,34 @@ kernel void gdn_ring_rows(device const float *x [[buffer(0)]], uint slot = (args.pos + r) % BLOOMERY_DELTA_RING_ROWS; ring[slot * args.ch + c] = x[r * args.x_stride + c]; } + +// The ring a chunk ending after row `after[k]` leaves, for each copy k: slot s +// holds the conv input of the latest position before the copy's own (the +// chunk's first position plus the row plus one) whose slot is s — a row of the +// chunk when that position is one, else the slot as the chunk found it. It +// reads the ring before gdn_ring_rows writes it, so the launch encodes it +// before that commit's barrier. Grid: the channels across, one row of the grid +// for each slot of each copy. Copy k's ring is `ring_len` floats of `snap` +// from k·ring_len, laid out as the ring. +kernel void gdn_ring_snap(device const float *x [[buffer(0)]], + device const float *ring [[buffer(1)]], + device float *snap [[buffer(2)]], + constant DeltaRingArgs &args [[buffer(3)]], + constant DeltaSnapArgs &snaps [[buffer(4)]], + uint2 gid [[thread_position_in_grid]]) { + uint c = gid.x; + uint s = gid.y % (uint)BLOOMERY_DELTA_RING_ROWS; + uint k = gid.y / (uint)BLOOMERY_DELTA_RING_ROWS; + if (c >= args.ch || k >= snaps.n) { + return; + } + uint last = args.pos + snaps.after[k]; + // The slot's distance back from `last`, reduced first so no position wraps. + uint back = (last % (uint)BLOOMERY_DELTA_RING_ROWS + (uint)BLOOMERY_DELTA_RING_ROWS - s) + % (uint)BLOOMERY_DELTA_RING_ROWS; + float v = ring[s * args.ch + c]; + if (back <= last && last - back >= args.pos) { + v = x[(last - back - args.pos) * args.x_stride + c]; + } + snap[(ulong)k * snaps.ring_len + s * args.ch + c] = v; +} diff --git a/crates/metal/src/arch/qwen3moe/ops.rs b/crates/metal/src/arch/qwen3moe/ops.rs index 3d56a3003..2db0fb705 100644 --- a/crates/metal/src/arch/qwen3moe/ops.rs +++ b/crates/metal/src/arch/qwen3moe/ops.rs @@ -74,9 +74,12 @@ pub(super) enum Feed { Device, } -/// What a pass adds to a step's ops: the rows arena. +/// What a pass adds to a step's ops: the rows arena, and the rows after +/// which the delta layers leave copies of their stores (none for a pass that +/// takes none). pub(super) struct RowsCtx<'a> { pub(super) arena: &'a mut RowsArena, + pub(super) after: &'a [u32], } /// The arm a unit of rows takes. diff --git a/crates/metal/src/arch/qwen3moe/ops/rows.rs b/crates/metal/src/arch/qwen3moe/ops/rows.rs index d04fe6b91..46288df5d 100644 --- a/crates/metal/src/arch/qwen3moe/ops/rows.rs +++ b/crates/metal/src/arch/qwen3moe/ops/rows.rs @@ -23,10 +23,10 @@ use super::super::session::SessionError; use super::{MetalOps, WHAT, all, all_mut}; use crate::kernels::{ AttentionPrefillParams, AttnCache, CONV_TAPS, CombineParams, DeltaRowStrides, DeltaRowsParams, - EXPERT_K, EXPERT_SLOTS, EmbedParams, F32GemmParams, GemmMode, GemmParams, GemmSegParams, - HEAD_COLS, HEAD_DIM, HEAD_SIMD_GROUPS, HeadNormRowsParams, HeadParams, KV_HEADS, LOGIT_STRIDE, - NormRowsParams, Operand, QUERY_HEADS, RouteSelectParams, RouteTableParams, RoutedKind, - RoutedParams, VALUE_HEADS, route_table_words, route_tiles, + DeltaSnaps, EXPERT_K, EXPERT_SLOTS, EmbedParams, F32GemmParams, GemmMode, GemmParams, + GemmSegParams, HEAD_COLS, HEAD_DIM, HEAD_SIMD_GROUPS, HeadNormRowsParams, HeadParams, KV_HEADS, + LOGIT_STRIDE, NormRowsParams, Operand, QUERY_HEADS, RouteSelectParams, RouteTableParams, + RoutedKind, RoutedParams, VALUE_HEADS, route_table_words, route_tiles, state_elems, }; /// The operand type the pass's GEMMs compute in. @@ -239,7 +239,9 @@ impl MetalOps<'_, '_> { Ok(()) } - /// Layer `l`'s delta step over the rows, and the ring's commit after it. + /// Layer `l`'s delta step over the rows, and the ring's commit after it; + /// with the pass's copies, the lane and the ring after each row they + /// name, into the layer's copies. pub(super) fn rows_delta_core(&mut self, l: usize, m: usize) -> Result<(), SessionError> { let slot = self.pass_slot(m); let from = self.pos; @@ -247,6 +249,7 @@ impl MetalOps<'_, '_> { let Some(ctx) = self.rows.as_mut() else { return Self::no_rows(); }; + let after = ctx.after; let rows = &mut *ctx.arena; let Arena { layers, @@ -275,6 +278,25 @@ impl MetalOps<'_, '_> { ..dense }; let tail = |start: usize| -> Elems { Elems(m * stride - start) }; + let snaps = if after.is_empty() { + None + } else { + let copies = rows.snaps.as_mut(); + let at = copies.as_ref().and_then(|c| c.at.get(l).copied().flatten()); + let (Some(c), Some(at), Some(lane_len)) = (copies, at, state_elems(1)) else { + return Err(ChainError::shape( + WHAT, + format!("layer {l}: a pass's copies with none open for the layer"), + ) + .into()); + }; + let n = after.len(); + Some(DeltaSnaps { + after, + lane: c.lane.window_mut(Elems(at.lane), Elems(n * lane_len))?, + ring: c.ring.window_mut(Elems(at.ring), Elems(n * at.ring_len))?, + }) + }; self.kernels.delta_rows( self.enc, DeltaRowsParams { @@ -312,6 +334,7 @@ impl MetalOps<'_, '_> { pos: from, strides, eps: *eps, + snaps, }, )?; self.enc.memory_barrier(); diff --git a/crates/metal/src/arch/qwen3moe/rows.rs b/crates/metal/src/arch/qwen3moe/rows.rs index bffc5b5c2..acdc90612 100644 --- a/crates/metal/src/arch/qwen3moe/rows.rs +++ b/crates/metal/src/arch/qwen3moe/rows.rs @@ -11,6 +11,10 @@ //! //! The expert buffers hold a row per slot (`m * 9`), the pair's output and //! the down's output at their widths: they are most of the arena. +//! +//! The copies a pass leaves after rows inside it ([`SnapBufs`]) are opened +//! on the first pass that takes one, apart from the arena: a pass that takes +//! none never pays for them. use bloomery_device::{Buffer, Elems}; @@ -18,8 +22,8 @@ use super::arena::{Arena, MixerBufs}; use super::session::SessionError; use crate::buffer::MetalBuffer; use crate::kernels::{ - EXPERT_K, EXPERT_SLOTS, LOGIT_STRIDE, PARTIALS_PER_HEAD, VALUE_HEADS, prepared_elems, - route_table_words, + EXPERT_K, EXPERT_SLOTS, LOGIT_STRIDE, PARTIALS_PER_HEAD, SNAPS, VALUE_HEADS, prepared_elems, + route_table_words, state_elems, }; use crate::residency::Alloc; @@ -98,6 +102,71 @@ pub(super) struct RowsArena { pub(super) y: MetalBuffer, /// The bytes the buffers take. pub(super) bytes: u64, + /// The copies of the delta layers' stores, once a pass has taken one. + pub(super) snaps: Option, +} + +/// Where one delta layer's copies sit: its [`SNAPS`] lanes from `lane` in +/// [`SnapBufs::lane`], its [`SNAPS`] rings from `ring` in [`SnapBufs::ring`], +/// a ring being `ring_len` values. +#[derive(Clone, Copy, Debug)] +pub(super) struct SnapAt { + pub(super) lane: usize, + pub(super) ring: usize, + pub(super) ring_len: usize, +} + +/// The copies a pass leaves after rows inside it: for each delta layer, +/// [`SNAPS`] lanes of its state and [`SNAPS`] of its conv ring, each laid out +/// as the store it copies, and the position each copy holds. A copy holds a +/// position only once the pass that wrote it has settled clean, and only +/// until the next pass, cut or reset: its buffers are read for nothing else, +/// so they are not zeroed. +pub(super) struct SnapBufs { + pub(super) lane: MetalBuffer, + pub(super) ring: MetalBuffer, + /// Each layer's copies; none for an attention layer. + pub(super) at: Vec>, + /// The position copy `k` holds: the position after the row it follows. + pub(super) holds: [Option; SNAPS], +} + +impl SnapBufs { + /// The copies of the delta layers of `arena`, opened on `device`. + pub(super) fn open(device: Alloc<'_>, arena: &Arena) -> Result { + let lane_len = state_elems(1).ok_or_else(|| refuse("a lane of the state does not fit"))?; + let (mut lanes, mut rings) = (0usize, 0usize); + let mut at = Vec::with_capacity(arena.layers.len()); + for layer in &arena.layers { + let MixerBufs::Delta(d) = &layer.mixer else { + at.push(None); + continue; + }; + let ring_len = d.ring.len().get(); + at.push(Some(SnapAt { + lane: lanes, + ring: rings, + ring_len, + })); + lanes = product(SNAPS, lane_len, "a layer's lane copies")? + .checked_add(lanes) + .ok_or_else(|| refuse("the lane copies do not fit"))?; + rings = product(SNAPS, ring_len, "a layer's ring copies")? + .checked_add(rings) + .ok_or_else(|| refuse("the ring copies do not fit"))?; + } + if lanes == 0 { + return Err(refuse("copies of a model with no delta layer")); + } + let lane = device.buffer(Elems(lanes))?; + let ring = device.buffer(Elems(rings))?; + Ok(SnapBufs { + lane, + ring, + at, + holds: [None; SNAPS], + }) + } } impl RowsArena { @@ -221,6 +290,7 @@ impl RowsArena { h, y, bytes, + snaps: None, }) } } diff --git a/crates/metal/src/arch/qwen3moe/session.rs b/crates/metal/src/arch/qwen3moe/session.rs index a5e090283..2ec2f417c 100644 --- a/crates/metal/src/arch/qwen3moe/session.rs +++ b/crates/metal/src/arch/qwen3moe/session.rs @@ -31,7 +31,9 @@ //! //! A checkpoint copies each delta layer's state lane at the position's parity //! and its whole conv ring. [`Session::checkpoint`] takes one at the position -//! the session stands at, with no token ahead and the stream idle. +//! the session stands at, with no token ahead and the stream idle; +//! [`Session::checkpoint_copy`] takes one at a position inside the last pass, +//! from the copies that pass left ([`StateOut::after_rows`]). //! [`Session::cut`] back to a point copies that lane and ring back and stands //! there; a cut to 0 is a reset, which drops the points too. A cut the ledger //! refuses leaves the position and every store as they were. @@ -44,7 +46,10 @@ //! leaves is the lane a step at the position after it reads. A pass is a //! command buffer an op, the last one waited: there is no host wait inside the //! pass, and none past the last op's. A prompt call cuts its ids into passes -//! at its checkpoint marks ([`Session::prefill`]). +//! at its checkpoint marks ([`Session::prefill`]). A pass may also leave +//! copies of the delta layers' lanes and rings after rows inside it: each is +//! the state the pass held after that row, and the pass's own bits do not +//! move. //! //! Every buffer the kernels bind is registered into the stream's residency //! set when the session allocates it, and the open waits for the set to be @@ -69,7 +74,7 @@ use std::path::Path; use std::sync::{Mutex, PoisonError}; use bloomery_device::fault::Fault; -use bloomery_device::{DeviceError, Elem, Elems, Stream}; +use bloomery_device::{Buffer, DeviceError, Elem, Elems, Stream}; use chain::ChainError; use chain::arch::qwen3moe::ops::{ChainOps, LayerCtx}; use chain::arch::qwen3moe::plan::{DeltaPlan, FfnPlan, GqaPlan, LayerPlan, SiteTy, kinds35}; @@ -82,13 +87,15 @@ use runtime::seqstate::{self, Kept}; use super::arena::{Arena, LANES, MixerBufs, lane_of, slot_of}; use super::count::{GroupBytes, GroupKind, Layer}; use super::ops::{Feed, MetalOps, RowsCtx}; -use super::rows::RowsArena; +use super::rows::{RowsArena, SnapBufs}; use crate::buffer::MetalBuffer; use crate::checkpoint::{self, Checkpoints, Restore}; use crate::device::MetalDevice; use crate::encode::EncodeCount; use crate::handle::Handle; -use crate::kernels::{EXPERT_SLOTS, HEAD_DIM, Kernels, PASS_ROWS, segments_for_keys, state_elems}; +use crate::kernels::{ + EXPERT_SLOTS, HEAD_DIM, Kernels, PASS_ROWS, SNAPS, segments_for_keys, state_elems, +}; use crate::library::Library; use crate::residency::{Alloc, ResidencyStats}; use crate::stream::{ @@ -593,12 +600,15 @@ pub enum PassOut { Last, } -/// The per-row copies of the recurrent state a pass leaves, named by the -/// rows they follow. A prefill pass takes none: its checkpoints fall between -/// passes. The verify arm's rows are this argument's. +/// The copies of the recurrent state a pass leaves, named by the rows they +/// follow: each delta layer's state lane and conv ring after that row, the +/// bits a pass ending there leaves in its stores, which +/// [`Session::checkpoint_copy`] takes as the point at the position after the +/// row. A prompt call names the message marks inside a pass here. #[derive(Clone, Copy, Debug, Default)] pub struct StateOut<'a> { - /// The rows after which a snapshot is taken; none for a prefill pass. + /// The rows after which a copy is taken, ascending, each a row of the + /// pass; at most [`SNAPS`]. None takes none. pub after_rows: &'a [u32], } @@ -921,6 +931,12 @@ impl Session { /// they were; one in a later layer waits out the layers committed, and a /// failure at the settle leaves the stores as the buffers left them: the /// caller resets. + /// + /// `snaps` names the rows after which the delta layers' lanes and rings + /// are copied ([`StateOut`]); rows out of order, past the pass or more + /// than [`SNAPS`] are refused by name before anything runs. The copies + /// hold their positions once the pass settles clean, and until the next + /// pass, cut or reset. pub fn pass( &mut self, ids: &[u32], @@ -956,12 +972,8 @@ impl Session { profile: Option<&mut Vec>, ) -> Result { self.ready()?; - if !snaps.after_rows.is_empty() { - return Err(refuse( - "a pass with per-row state snapshots: only the verify arm takes them", - )); - } let m = ids.len(); + let after = snaps.after_rows; if m <= PASS_ROWS { return Err(refuse(format!( "a pass of {m} ids; a step's window is {PASS_ROWS} and a pass is more" @@ -971,6 +983,15 @@ impl Session { if m > cap { return Err(refuse(format!("a pass of {m} ids; a pass holds {cap}"))); } + if after.len() > SNAPS + || !after.windows(2).all(|w| w[0] < w[1]) + || after.iter().any(|&r| r as usize >= m) + { + return Err(refuse(format!( + "a pass of {m} ids with copies after rows {after:?}: at most {SNAPS}, ascending, \ + each a row of the pass" + ))); + } let from = self.pos; let rows_n = u32::try_from(m).map_err(|_| refuse(format!("{m} ids do not fit u32")))?; match from.checked_add(rows_n) { @@ -995,6 +1016,21 @@ impl Session { staged.commit(); self.rows = Some(rows); } + if let Some(rows) = self.rows.as_mut() { + if !after.is_empty() && rows.snaps.is_none() { + let residency = self + .stream + .residency() + .ok_or_else(|| refuse("the stream has no residency set"))?; + let staged = residency.stage(); + let copies = SnapBufs::open(Alloc::new(&self.device, &staged), &self.arena)?; + staged.commit(); + rows.snaps = Some(copies); + } + if let Some(copies) = rows.snaps.as_mut() { + copies.holds = [None; SNAPS]; + } + } // Built before the first op commits, so a refused build leaves the // state, rings and cache as they were. self.kernels.prepare_prefill()?; @@ -1013,6 +1049,7 @@ impl Session { kernels: &self.kernels, arena: &mut self.arena, rows, + after, from, head, }; @@ -1038,6 +1075,7 @@ impl Session { kernels: &self.kernels, arena: &mut self.arena, rows, + after, from, head, count: EncodeCount::default(), @@ -1067,6 +1105,11 @@ impl Session { } else { ids[m - 1] }; + if let Some(copies) = self.rows.as_mut().and_then(|r| r.snaps.as_mut()) { + for (held, &r) in copies.holds.iter_mut().zip(after) { + *held = Some(from + r + 1); + } + } self.pos = from + rows_n; Ok(Passed { id, count }) } @@ -1336,9 +1379,17 @@ impl Session { self.pos = 0; self.poisoned = None; self.checkpoints.clear(); + self.drop_copies(); Ok(()) } + /// The last pass's copies hold no position any more. + fn drop_copies(&mut self) { + if let Some(copies) = self.rows.as_mut().and_then(|r| r.snaps.as_mut()) { + copies.holds = [None; SNAPS]; + } + } + /// The ledger's counts. #[must_use] pub fn checkpoint_stats(&self) -> seqstate::Stats { @@ -1387,35 +1438,90 @@ impl Session { /// pages straight into the point's slot. Refused while the session is /// poisoned. pub fn checkpoint(&mut self) -> Result<(), SessionError> { + self.take_point(self.pos, None) + } + + /// A checkpoint at `at`, a position inside the last pass that pass left + /// a copy of ([`StateOut::after_rows`]): each delta layer's lane and ring + /// copied once from that copy into the point's slot, the stream idle. A + /// position no copy holds (another pass, a cut or a reset came since) is + /// refused by name, as is one past the position. Points go in ascending + /// order: a copy's point before the point at the pass's end. + pub fn checkpoint_copy(&mut self, at: u32) -> Result<(), SessionError> { + let held = self + .rows + .as_ref() + .and_then(|r| r.snaps.as_ref()) + .and_then(|c| c.holds.iter().position(|&h| h == Some(at))); + match held { + Some(k) if at <= self.pos => self.take_point(at, Some(k)), + _ => Err(refuse(format!( + "a checkpoint at {at} from a pass's copy: no copy the last pass left holds it \ + (the session stands at {})", + self.pos + ))), + } + } + + /// The point at `at` from each delta layer's live stores (`copy` none: + /// the lane at `at`'s parity, `at` the position) or from the last pass's + /// copy `copy`, after the stream is idle. Refused while poisoned. + fn take_point(&mut self, at: u32, copy: Option) -> Result<(), SessionError> { if let Some(fault) = self.poisoned { return Err(SessionError::Poisoned(fault)); } self.quiesce()?; let lane_len = lane_len()?; - let start = lane_of(self.pos) + let live = lane_of(at) .checked_mul(lane_len) .ok_or_else(|| refuse("a delta state lane's offset overflowed"))?; + let copies = match copy { + None => None, + Some(k) => Some(( + k, + self.rows + .as_ref() + .and_then(|r| r.snaps.as_ref()) + .ok_or_else(|| refuse("a pass's copy with no copies open"))?, + )), + }; let mut views = Vec::new(); - for layer in &self.arena.layers { + for (l, layer) in self.arena.layers.iter().enumerate() { let MixerBufs::Delta(d) = &layer.mixer else { continue; }; + let (pages, lane_at, ring_at) = match copies { + None => ((&d.state, &d.ring), live, 0), + Some((k, c)) => { + let at = + c.at.get(l) + .copied() + .flatten() + .ok_or_else(|| refuse(format!("delta layer {l} has no copies")))?; + ( + (&c.lane, &c.ring), + at.lane + k * lane_len, + at.ring + k * at.ring_len, + ) + } + }; + let ring_len = d.ring.len().get(); // SAFETY: `quiesce` above left the stream idle (no token ahead, - // every committed buffer done), and the views borrow the arena, - // whose device work and host writes all go through `&mut self`, - // until the take below has copied them: nothing writes the pages - // while they are read. - let (state, ring) = unsafe { (d.state.pages().as_slice(), d.ring.pages().as_slice()) }; + // every committed buffer done), and the views borrow the arena + // and the rows arena, whose device work and host writes all go + // through `&mut self`, until the take below has copied them: + // nothing writes the pages while they are read. + let (lane, ring) = unsafe { (pages.0.pages().as_slice(), pages.1.pages().as_slice()) }; views.push( - state - .get(start..start + lane_len) + lane.get(lane_at..lane_at + lane_len) .ok_or_else(|| refuse("a delta state lane lies past its buffer"))?, ); - views.push(ring); + views.push( + ring.get(ring_at..ring_at + ring_len) + .ok_or_else(|| refuse("a delta conv ring lies past its buffer"))?, + ); } - self.checkpoints - .take(self.pos, &views) - .map_err(checkpoint_error) + self.checkpoints.take(at, &views).map_err(checkpoint_error) } /// Cut to position `n`. A cut to 0 resets, which drops the checkpoints. A @@ -1471,6 +1577,7 @@ impl Session { Err(e) => return Err(checkpoint_error(e)), } self.pos = n; + self.drop_copies(); Ok(()) } @@ -1629,6 +1736,8 @@ struct PassLayers<'s> { kernels: &'s Kernels, arena: &'s mut Arena, rows: &'s mut RowsArena, + /// The rows the delta layers copy their stores after. + after: &'s [u32], from: u32, head: bool, } @@ -1646,7 +1755,10 @@ impl PassLayers<'_> { f(&mut MetalOps::rows( encoder, self.arena, - RowsCtx { arena: self.rows }, + RowsCtx { + arena: self.rows, + after: self.after, + }, self.kernels, self.from, self.head, @@ -1782,6 +1894,8 @@ struct PassProfiler<'s> { kernels: &'s Kernels, arena: &'s mut Arena, rows: &'s mut RowsArena, + /// The rows the delta layers copy their stores after. + after: &'s [u32], from: u32, head: bool, count: EncodeCount, @@ -1801,7 +1915,10 @@ impl PassProfiler<'_> { f(&mut MetalOps::rows( &mut enc, self.arena, - RowsCtx { arena: self.rows }, + RowsCtx { + arena: self.rows, + after: self.after, + }, self.kernels, self.from, self.head, diff --git a/crates/metal/src/arch/qwen3moe/target.rs b/crates/metal/src/arch/qwen3moe/target.rs index f46206dc6..7382a8fa0 100644 --- a/crates/metal/src/arch/qwen3moe/target.rs +++ b/crates/metal/src/arch/qwen3moe/target.rs @@ -32,11 +32,21 @@ //! read back after the last as asked. Each mark's checkpoint waits out //! whatever token runs ahead of it. A session in [`PrefillMode::Steps`] //! steps every batch, the same-binary arm and the oracle's. +//! - `prompt_marked`, `prompt` with a point at each of the caller's marks +//! (a server's message starts) and no cut a mark: a mark inside a pass is +//! a copy the pass leaves after the row before it +//! ([`StateOut::after_rows`]), taken as a point once the pass settles +//! ([`Session::checkpoint_copy`]), so the point holds the bits the pass +//! held there; a mark inside a run of steps ends that run, its point the +//! step's own state (a step runs one position whatever ends the run); a +//! mark at a batch's end is that batch's point. A pass takes at most +//! [`SNAPS`] copies: a further mark ends the pass there. A mark out of +//! order or outside the call is refused by name before any step runs. //! //! What it refuses: //! -//! - a prompt of no ids, and one that would pass the cache's last row, -//! before any step runs; +//! - a prompt of no ids, one that would pass the cache's last row, and +//! marks out of order or outside the call, before any step runs; //! - a cut where no checkpoint stands between 0 and the position (the //! ledger's refusal, by name); a cut to the position does nothing and a //! cut to 0 resets; @@ -47,7 +57,7 @@ use runtime::seqstate::Kept; use runtime::{Out, Target, Want}; use super::session::{PassOut, Session, SessionError, StateOut, StepMode}; -use crate::kernels::PASS_ROWS; +use crate::kernels::{PASS_ROWS, SNAPS}; const WHAT: &str = "qwen35moe target"; @@ -71,6 +81,15 @@ impl Target for Session { } fn prompt(&mut self, ids: &[u32], want: Want) -> Result, SessionError> { + Target::prompt_marked(self, ids, &[], want) + } + + fn prompt_marked( + &mut self, + ids: &[u32], + marks: &[u32], + want: Want, + ) -> Result, SessionError> { if ids.is_empty() { return Err(refuse(format!( "a prompt of no ids at position {}", @@ -93,32 +112,42 @@ impl Target for Session { ids.len() )) })?; - // The id at an inner mark is stepped by the next span, after the take. - let marks = self.checkpoint_marks(from, to); - let last = ids.len() - 1; + let inside = marks.first().is_none_or(|&u| u > from) + && marks.last().is_none_or(|&u| u <= to) + && marks.windows(2).all(|w| w[0] < w[1]); + if !inside { + return Err(refuse(format!( + "a prompt of positions {from}..{to} marked at {marks:?}: marks go ascending, \ + past the call's start and at most its end" + ))); + } + // The id at an inner point is stepped by the next span, after the take. + let points = self.checkpoint_marks(from, to); + let call = Call { + ids, + from, + last: ids.len() - 1, + steps: Session::prefill(self) == PrefillMode::Steps, + }; let batch_rows = Session::pass_rows(self); - let steps = Session::prefill(self) == PrefillMode::Steps; let mut at = from; let mut last_pick = None; - for m in marks { - // The runs between two marks, each a pass or its ids' steps. - for batch in call_batches(at, m, &[m], batch_rows) { - let start = (batch.start - from) as usize; - let end_i = (batch.end - from) as usize; - if steps || end_i - start <= PASS_ROWS { - last_pick = Some(step_batch(self, ids, start, end_i, last)?); - } else { - let out = if end_i == ids.len() { - PassOut::Last - } else { - PassOut::None - }; - let passed = Session::pass(self, &ids[start..end_i], out, StateOut::default())?; - last_pick = Some(passed.id); + for p in points { + // The runs between two points, each cut into batches: a pass or + // its ids' steps, with the marks inside it. + for batch in call_batches(at, p, &[p], batch_rows) { + let inner: Vec = marks + .iter() + .copied() + .filter(|&u| batch.start < u && u < batch.end) + .collect(); + last_pick = Some(call.run(self, batch.start, batch.end, &inner)?); + if batch.end != p && marks.contains(&batch.end) { + Session::checkpoint(self)?; } } Session::checkpoint(self)?; - at = m; + at = p; } let pick = last_pick.ok_or_else(|| { refuse(format!( @@ -158,6 +187,56 @@ impl Target for Session { } } +/// A prompt call's ids, the position of its first, the index of its last +/// and whether the session steps every id. +struct Call<'a> { + ids: &'a [u32], + from: u32, + last: usize, + steps: bool, +} + +impl Call<'_> { + /// Positions `a..e` of the call with a point at each of `marks` + /// (ascending, inside the span): a pass taking a copy at each, or the + /// steps of each run between two marks; past [`SNAPS`] marks the span is + /// cut at the next one, whose point the first part's end is. Returns the + /// last pick. + fn run(&self, s: &mut Session, a: u32, e: u32, marks: &[u32]) -> Result { + let i = |u: u32| (u - self.from) as usize; + if self.steps || i(e) - i(a) <= PASS_ROWS { + let mut from = a; + for &u in marks { + step_batch(s, self.ids, i(from), i(u), self.last)?; + Session::checkpoint(s)?; + from = u; + } + return step_batch(s, self.ids, i(from), i(e), self.last); + } + if let Some(&cut) = marks.get(SNAPS) { + self.run(s, a, cut, &marks[..SNAPS])?; + Session::checkpoint(s)?; + return self.run(s, cut, e, &marks[SNAPS + 1..]); + } + let after: Vec = marks.iter().map(|&u| u - a - 1).collect(); + let out = if i(e) == self.ids.len() { + PassOut::Last + } else { + PassOut::None + }; + let passed = Session::pass( + s, + &self.ids[i(a)..i(e)], + out, + StateOut { after_rows: &after }, + )?; + for &u in marks { + s.checkpoint_copy(u)?; + } + Ok(passed.id) + } +} + /// The ids `start..end_i` of a prompt stepped one ahead each where the next /// id is known and inner, none of them but the prompt's last carrying a /// head: every id but the batch's last steps the next id ahead, both tokens diff --git a/crates/metal/src/bin/bloomery_serve_metal.rs b/crates/metal/src/bin/bloomery_serve_metal.rs index c2f1c6939..7b47f1fae 100644 --- a/crates/metal/src/bin/bloomery_serve_metal.rs +++ b/crates/metal/src/bin/bloomery_serve_metal.rs @@ -17,12 +17,12 @@ //! //! One slot: the session is one sequence. A request keeps a prefix back to a //! checkpoint of the delta layers' recurrent state: a prompt call takes one -//! at its start, every `CHECKPOINT_EVERY` positions and at its end, and a -//! prompt split at each message start makes that start a call's end. Any -//! other prefix runs from position 0. A prompt span of more than a step's -//! window runs in passes of up to `Session::pass_rows` ids (`Target::prompt`), -//! a shorter one a host-fed step an id. No prompt cache, no slot files, no -//! image input. +//! at its start, every `CHECKPOINT_EVERY` positions, at each message start +//! the server marks inside it (`Target::prompt_marked`: a copy the pass +//! leaves at that row, the call uncut) and at its end. Any other prefix runs +//! from position 0. A prompt span of more than a step's window runs in passes +//! of up to `Session::pass_rows` ids, a shorter one a host-fed step an id. No +//! prompt cache, no slot files, no image input. //! //! A bad flag or value is a named error and exit 64; a refusal of the file, //! the device or the session is a named error and exit 1; an engine failure diff --git a/crates/metal/src/bin/gate_metal_qwen35moe_e2e/keep.rs b/crates/metal/src/bin/gate_metal_qwen35moe_e2e/keep.rs index 3f243d957..328199b71 100644 --- a/crates/metal/src/bin/gate_metal_qwen35moe_e2e/keep.rs +++ b/crates/metal/src/bin/gate_metal_qwen35moe_e2e/keep.rs @@ -13,8 +13,9 @@ //! the next step the replay's; //! - the engine keeps a prefill's end, and a second request that shares it //! matches a fresh run of that request; -//! - `prefill_splits` returns every mark, and a call ending at one makes it -//! a point. +//! - `prefill_splits` cuts at no mark, `prefill_marked` keeps each of its +//! marks as a point inside one call, and a call ending at a message start +//! makes it a point. //! //! The prompt is 16 ids. The conv ring has 11 rows, and position 4 reads the //! slot the token at 12 writes, so a cut that skips the ring is not the fresh @@ -173,9 +174,20 @@ fn engine_keeps( )?; let marks = [3usize, 7]; let splits = Engine::prefill_splits(&e, 0, ids.len(), &marks); - if splits != marks { - return Err(format!("prefill_splits of {marks:?} returned {splits:?}")); + if !splits.is_empty() { + return Err(format!("prefill_splits of {marks:?} cut at {splits:?}")); } + e.prefill_marked(&ids[..9], &[], &marks) + .map_err(engine_says)?; + for u in marks { + if e.keepable(u) != u { + return Err(format!( + "keepable({u}) is {} after a call marked at {marks:?}", + e.keepable(u) + )); + } + } + e.reset().map_err(engine_says)?; let lead = &ids[..4]; e.prefill(lead).map_err(engine_says)?; diff --git a/crates/metal/src/bin/gate_metal_qwen35moe_e2e/mac.rs b/crates/metal/src/bin/gate_metal_qwen35moe_e2e/mac.rs index 6fbc9e3fa..a8e7376a8 100644 --- a/crates/metal/src/bin/gate_metal_qwen35moe_e2e/mac.rs +++ b/crates/metal/src/bin/gate_metal_qwen35moe_e2e/mac.rs @@ -70,14 +70,17 @@ //! site lists — and `generate` over `Plain` emits the session's own //! host-fed picks (`t_target`'s greedy pattern on that session). //! -//! - `q_short`, `q_state`, `q_chunk`, `q_ref`, `q_free`, `q_fault`, `q_layers`: the -//! batched prompt pass (`pass.rs`): a prompt of at most a step's window is the -//! step path; one pass against the stepped ids leaves layer 0's state and ring -//! within a derived band; the pass size moves no bit; the step at the prompt's -//! depth against ik's `l_out` and logits, the pass run's distance over the -//! stepped run's within a ratio; the greedy tokens teacher-forced through both; -//! an id past the vocabulary in a pass raises the token-id site and poisons; a -//! pass is one command buffer a layer and leaves the profiled pass's bits. +//! - `q_short`, `q_state`, `q_chunk`, `q_ref`, `q_free`, `q_fault`, `q_layers`, +//! `q_mark`, `q_mark_cut`: the batched prompt pass (`pass.rs`): a prompt of at +//! most a step's window is the step path; one pass against the stepped ids +//! leaves layer 0's state and ring within a derived band; the pass size moves +//! no bit; the step at the prompt's depth against ik's `l_out` and logits, the +//! pass run's distance over the stepped run's within a ratio; the greedy tokens +//! teacher-forced through both; an id past the vocabulary in a pass raises the +//! token-id site and poisons; a pass is one command buffer a layer and leaves +//! the profiled pass's bits; a pass's copies at rows inside it move none of its +//! bits and are its state there, and a prompt cut back to a mark inside its +//! pass continues the same every time, within the teacher-forced band. //! //! - `r_resident`: the session's kernel buffers are held in one residency set //! (`resident.rs`): the set holds every byte the open and the first pass @@ -457,7 +460,7 @@ pub fn main() -> ExitCode { println!("{GATE}: {name} takes a value"); return ExitCode::from(64); } - const NAMES: [&str; 21] = [ + const NAMES: [&str; 23] = [ "d_count", "e_token_id", "b_chain", @@ -478,6 +481,8 @@ pub fn main() -> ExitCode { "q_free", "q_fault", "q_layers", + "q_mark", + "q_mark_cut", "r_resident", ]; if args.iter().any(|a| a == "--list") { @@ -604,6 +609,8 @@ pub fn main() -> ExitCode { "q_free" => pass::q_free(&device, library, &path), "q_fault" => pass::q_fault(&device, library, &path), "q_layers" => pass::q_layers(&device, library, &path), + "q_mark" => pass::q_mark(&device, library, &path), + "q_mark_cut" => pass::q_mark_cut(&device, library, &path), "r_resident" => resident::r_resident(&device, library, &path), _ => free::a_free(&mut session), }; diff --git a/crates/metal/src/bin/gate_metal_qwen35moe_e2e/mutants.rs b/crates/metal/src/bin/gate_metal_qwen35moe_e2e/mutants.rs index c921905a6..683fbf79e 100644 --- a/crates/metal/src/bin/gate_metal_qwen35moe_e2e/mutants.rs +++ b/crates/metal/src/bin/gate_metal_qwen35moe_e2e/mutants.rs @@ -101,6 +101,26 @@ pub const MUTANTS: &[Mutant] = &[ 1, )], }, + Mutant { + name: "mark_row_late", + clause: "q_mark", + why: "a pass copies a delta layer's state lane one row after the row its copy names", + edits: &[( + "if (copy < copies && r == snaps->after[copy]) {", + "if (copy < copies && r == snaps->after[copy] + 1u) {", + 1, + )], + }, + Mutant { + name: "mark_ring_dropped", + clause: "q_mark", + why: "a pass leaves a copy's conv ring unwritten", + edits: &[( + " snap[(ulong)k * snaps.ring_len + s * args.ch + c] = v;", + " (void)v;", + 1, + )], + }, Mutant { name: "prefill_ring_last_row", clause: "q_state", diff --git a/crates/metal/src/bin/gate_metal_qwen35moe_e2e/pass.rs b/crates/metal/src/bin/gate_metal_qwen35moe_e2e/pass.rs index 04c3a62b6..d9d29588b 100644 --- a/crates/metal/src/bin/gate_metal_qwen35moe_e2e/pass.rs +++ b/crates/metal/src/bin/gate_metal_qwen35moe_e2e/pass.rs @@ -28,6 +28,16 @@ //! and the stepped run: every layer's output of every token within //! [`FREE_BAND`] of the stepped run's, and the routers' picks differing in //! no more than [`FREE_FLIPS_MAX`] sites. +//! - `q_mark`: a pass's copies of its delta stores (`StateOut::after_rows`) +//! move none of its bits (pick, logits row, every store); the copy after +//! its last row, taken as a point and cut back to, is the pass's own end +//! state bit for bit; a copy's ring slots of the positions before the pass +//! are the ring the pass found; and a copy is the same bits whichever other +//! copies the pass takes. +//! - `q_mark_cut`: a prompt marked at two positions inside a pass, cut back +//! to each mark and continued to the end, twice: the same pick, logits row +//! and stores both times, and the tokens after it within `q_free`'s band of +//! the stepped run's. use std::path::Path; @@ -41,6 +51,7 @@ use bloomery_metal::device::MetalDevice; use bloomery_metal::kernels::{CONV_TAPS, HEAD_COLS, HEAD_DIM, KV_HEADS, RING_ROWS, VALUE_HEADS}; use bloomery_metal::library::Library; use runtime::prompt::PrefillMode; +use runtime::{Target, Want}; use super::chain; use super::flip; @@ -100,6 +111,17 @@ const Q_LAYERS_IDS: usize = 100; /// the sixteen router sites [`FREE_FLIPS_MAX`] is derived over. const FREE_STEPS: usize = 4; +/// The ids `q_mark` passes, and the rows its copies follow: an inner row +/// before the ring's length, one past it, and the last. +const MARK_IDS: usize = 40; +const MARK_ROWS: [u32; 3] = [6, 25, 39]; + +/// The marks `q_mark_cut` puts inside its prompt's second pass (a point every +/// [`MARK_EVERY`]): one where the ring is all that pass's rows, one whose ring +/// the pass's last rows overwrite. +const MARK_CUTS: [u32; 2] = [700, 1017]; +const MARK_EVERY: u32 = 512; + /// The ids of a prompt of `n` ids: the d1k set's first. pub(super) fn prompt_ids(n: usize) -> Result, String> { let man = sets::open(Set::D1k)?; @@ -752,3 +774,169 @@ pub fn q_layers(device: &MetalDevice, library: &Library, path: &Path) -> Clause samples.len() ))) } + +/// The pass of `ids` from zero stores taking copies after the rows `after`, +/// and what it left. +fn copied_pass(s: &mut Session, ids: &[u32], after: &[u32]) -> Result { + s.reset().map_err(says)?; + let passed = s + .pass(ids, PassOut::Last, StateOut { after_rows: after }) + .map_err(says)?; + let logits = bits(s.logits_row().map_err(says)?); + read_stores(s, passed.id, logits) +} + +/// The delta layers' stores at `at`, a point the session holds: one step on +/// from the position, then the cut back to it. The attention rows read are +/// those before `at`, which neither the step nor the cut moves. +fn stores_at_point(s: &mut Session, at: u32, end: &Stores) -> Result { + Target::step(s, end.pick, Want::Argmax).map_err(says)?; + Target::cut(s, at).map_err(says)?; + read_stores(s, end.pick, end.logits.clone()) +} + +/// The first delta layer whose ring in `at`, a copy after row `r` of a pass +/// from 0, differs from `found` (the ring that pass found) in a slot of a +/// position before the pass: slots `r + 1..`, which the copy takes from the +/// ring as it was. +fn found_slots(at: &Stores, found: &Stores, r: usize) -> Option { + at.layers + .iter() + .zip(&found.layers) + .position(|(a, f)| match (a, f) { + (Store::Rec { ring: a, .. }, Store::Rec { ring: f, .. }) => { + let from = (r + 1) * (f.len() / RING_ROWS); + bits(&a[from..]) != bits(&f[from..]) + } + _ => false, + }) +} + +/// `q_mark`: a pass's copies are its own state at their rows and move none of +/// its bits. +pub fn q_mark(device: &MetalDevice, library: &Library, path: &Path) -> Clause { + let ids = prompt_ids(MARK_IDS)?; + let mut s = open( + device, + library, + path, + MARK_IDS + 2, + 2048, + PrefillMode::Batch, + )?; + s.reset().map_err(says)?; + let found = read_stores(&mut s, 0, Vec::new())?; + let plain = copied_pass(&mut s, &ids, &[])?; + let [inner, outer, last] = MARK_ROWS; + let mut seen: Vec<(u32, Stores)> = Vec::new(); + for after in [[inner, last], [outer, last], [inner, outer]] { + let what = format!("the pass copying after rows {after:?}"); + let got = copied_pass(&mut s, &ids, &after)?; + if let Some(why) = differs(&what, &got, &plain) { + return Err(format!("{why}, against the pass with no copies")); + } + for &r in &after { + s.checkpoint_copy(r + 1).map_err(says)?; + } + // Descending, so each cut finds the point it names. + for &r in after.iter().rev() { + let at = stores_at_point(&mut s, r + 1, &got)?; + let held = format!("{what}: the point at {} from its copy", r + 1); + if r == last + && let Some(why) = differs(&held, &at, &plain) + { + return Err(format!("{why}, against the pass's own end")); + } + if (r as usize) + 1 < RING_ROWS + && let Some(l) = found_slots(&at, &found, r as usize) + { + return Err(format!( + "{held}: delta layer {l}'s ring slots before the pass are not the ring the \ + pass found" + )); + } + match seen.iter().find(|(row, _)| *row == r) { + Some((_, first)) => { + if let Some(why) = differs(&held, &at, first) { + return Err(format!("{why}, against the same row's copy beside another")); + } + } + None => seen.push((r, at)), + } + } + } + Ok(Verdict::Noted(format!( + "{MARK_IDS} ids as one pass copying after rows {MARK_ROWS:?} two at a time: the pick, logits \ + row and every store of the pass with none; the copy after the last row is the pass's end, \ + the copy after row {inner} keeps the found ring before the pass, and each copy the same \ + beside either other" + ))) +} + +/// `q_mark_cut`: cut back to a mark inside a pass and continued, twice, and +/// the tokens after it against the stepped run's. +pub fn q_mark_cut(device: &MetalDevice, library: &Library, path: &Path) -> Clause { + let ids = prompt_ids(1024)?; + let ctx = ids.len() + FREE_STEPS + 1; + let mut steps = open(device, library, path, ctx, 2048, PrefillMode::Steps)?; + let (_, pick) = keep::prompt_bits(&mut steps, &ids)?; + let mut want = Vec::with_capacity(FREE_STEPS); + let mut tok = pick; + for _ in 0..FREE_STEPS { + let t = traced_step(&mut steps, tok)?; + tok = t.top; + want.push(t); + } + let mut s = open(device, library, path, ctx, MARK_EVERY, PrefillMode::Batch)?; + s.reset().map_err(says)?; + Target::prompt_marked(&mut s, &ids, &MARK_CUTS, Want::Argmax).map_err(says)?; + if s.rows_arena_bytes().is_none() { + return Err("a prompt of 1,024 ids did not take the pass".to_string()); + } + let layers = s.layer_count(); + let mut notes = Vec::new(); + // Descending: a cut drops the points past it. + for &u in MARK_CUTS.iter().rev() { + let rest = &ids[u as usize..]; + let mut first: Option = None; + for _ in 0..2 { + Target::cut(&mut s, u).map_err(says)?; + let (logits, pick) = keep::prompt_bits(&mut s, rest)?; + let got = read_stores(&mut s, pick, logits)?; + match &first { + Some(a) => { + if let Some(why) = differs(&format!("continued from {u} again"), &got, a) { + return Err(why); + } + } + None => first = Some(got), + } + } + let start = first.as_ref().map_or(pick, |a| a.pick); + let (mut worst, mut flips) = (0f64, 0usize); + let mut tok = pick; + for b in &want { + let a = traced_step(&mut s, tok)?; + for l in 0..layers { + worst = flip::worse(worst, flip::rel_to(&a.outs[l], &b.outs[l])); + flips += + usize::from(a.ids[l][..a.ids[l].len() - 1] != b.ids[l][..b.ids[l].len() - 1]); + } + tok = b.top; + } + let line = format!( + "from {u}: {FREE_STEPS} teacher-forced tokens worst layer output {worst:.3e} of the \ + stepped run's (band {FREE_BAND}), {flips} router sites with other picks (cap \ + {FREE_FLIPS_MAX}), pick {start} against the stepped run's {pick}" + ); + if worst > FREE_BAND || flips > FREE_FLIPS_MAX { + return Err(line); + } + notes.push(line); + } + Ok(Verdict::Noted(format!( + "1,024 ids in passes of {MARK_EVERY} marked at {MARK_CUTS:?}, cut back to each and \ + continued twice: the same pick, logits row and stores; {}", + notes.join("; ") + ))) +} diff --git a/crates/metal/src/bin/gate_metal_qwen35moe_kernels/clauses_rows.rs b/crates/metal/src/bin/gate_metal_qwen35moe_kernels/clauses_rows.rs index 4de7a013a..e05e0a3ee 100644 --- a/crates/metal/src/bin/gate_metal_qwen35moe_kernels/clauses_rows.rs +++ b/crates/metal/src/bin/gate_metal_qwen35moe_kernels/clauses_rows.rs @@ -10,9 +10,10 @@ use bloomery_metal::MetalBuffer; use bloomery_metal::kernels::router; use bloomery_metal::kernels::{ AttentionParams, AttentionPrefillParams, AttentionRowsParams, AttnCache, CONV_TAPS, - DeltaParams, DeltaRowStrides, DeltaRowsParams, HEAD_COLS, HEAD_DIM, HeadMap, KV_HEADS, - NormRowsParams, PARTIALS_PER_HEAD, QUERY_HEADS, RING_ROWS, ROT, VALUE_HEADS, channels, out_len, - partials_len, prepared_elems, q8_code_len, q8_scale_len, qg_len, ring_elems, state_elems, + DeltaParams, DeltaRowStrides, DeltaRowsParams, DeltaSnaps, HEAD_COLS, HEAD_DIM, HeadMap, + KV_HEADS, NormRowsParams, PARTIALS_PER_HEAD, QUERY_HEADS, RING_ROWS, ROT, VALUE_HEADS, + channels, out_len, partials_len, prepared_elems, q8_code_len, q8_scale_len, qg_len, ring_elems, + state_elems, }; use super::clauses_attention::{ @@ -172,11 +173,47 @@ impl Bufs { /// The chunk as one rows launch. fn run_rows(h: &Harness<'_>, fx: &Fixed, c: &Chunk) -> Result { + run_rows_copied(h, fx, c, &[]).map(|(out, _)| out) +} + +/// The lanes and rings a rows launch copied, copy after copy. +struct Copies { + lane: Vec, + ring: Vec, +} + +/// The chunk as one rows launch that copies its stores after the rows +/// `after` (none: no copies). The copies' buffers start NaN, so a copy the +/// launch did not write reads as no number. +fn run_rows_copied( + h: &Harness<'_>, + fx: &Fixed, + c: &Chunk, + after: &[u32], +) -> Result<(Out, Copies), String> { let mut bufs = Bufs::new(h, fx, c)?; let strides = DeltaRowStrides::dense(N_K).ok_or("the strides do not fit")?; let prep_len = c.rows * prep_row()?; + let lane_len = after.len() * elems(state_elems(1), "a lane")?; + let ring_len = after.len() * c.ring.len(); + let poison = |n: usize| -> Result, String> { + h.device + .upload(h.stream, &vec![f32::NAN; n.max(1)]) + .map_err(why) + }; + let mut lane_copies = poison(lane_len)?; + let mut ring_copies = poison(ring_len)?; h.clear_fault()?; h.encode(|encoder| { + let snaps = if after.is_empty() { + None + } else { + Some(DeltaSnaps { + after, + lane: lane_copies.window_mut(Elems(0), Elems(lane_len))?, + ring: ring_copies.window_mut(Elems(0), Elems(ring_len))?, + }) + }; h.kernels.delta_rows( encoder, DeltaRowsParams { @@ -208,10 +245,32 @@ fn run_rows(h: &Harness<'_>, fx: &Fixed, c: &Chunk) -> Result { pos: c.pos, strides, eps: fx.eps, + snaps, }, ) })?; - bufs.read(h, c) + let out = bufs.read(h, c)?; + let mut copies = Copies { + lane: vec![0f32; lane_len], + ring: vec![0f32; ring_len], + }; + if !after.is_empty() { + h.device + .read( + h.stream, + lane_copies.window(Elems(0), Elems(lane_len)).map_err(why)?, + &mut copies.lane, + ) + .map_err(why)?; + h.device + .read( + h.stream, + ring_copies.window(Elems(0), Elems(ring_len)).map_err(why)?, + &mut copies.ring, + ) + .map_err(why)?; + } + Ok((out, copies)) } /// The chunk as `rows` stepped launches in one command buffer, row `r` at @@ -275,7 +334,7 @@ fn same_bits(what: &str, got: &[f32], want: &[f32]) -> Result<(), String> { { None => Ok(()), Some(i) => Err(format!( - "{what}[{i}] is {:?} ({:#010x}), the stepped launch's {:?} ({:#010x})", + "{what}[{i}] is {:?} ({:#010x}), expected {:?} ({:#010x})", got[i], got[i].to_bits(), want[i], @@ -345,6 +404,85 @@ pub fn delta_rows(h: &Harness<'_>, fixture: &Fixture) -> Clause { ))) } +/// The first `rows` rows of `c`: the same state, ring and lane words. +fn head_of(c: &Chunk, rows: usize) -> Result { + let ch = elems(channels(N_K), "channels")?; + Ok(Chunk { + rows, + pos: c.pos, + lane_words: c.lane_words.clone(), + lane_at: c.lane_at, + lane_to: c.lane_to, + x: c.x[..rows * ch].to_vec(), + b: c.b[..rows * VALUE_HEADS].to_vec(), + a: c.a[..rows * VALUE_HEADS].to_vec(), + z: c.z[..rows * G_ROW].to_vec(), + state: c.state.clone(), + ring: c.ring.clone(), + }) +} + +/// A copies case: rows, first position, lane words, the copies' rows. +type SnapCase = (usize, u32, (usize, usize), &'static [u32]); + +/// The cases of the copies clause: rows, first position, the read and the +/// write lane word, and the rows the copies follow. A copy after row 0 takes +/// every ring slot but one from the ring the chunk found; one after a row +/// past the ring's length takes none; a first position of 0 has slots no +/// position wrote; a copy after the last row is the chunk's own end; a +/// chunk ending at the last u32 position takes its ring slots without the +/// slot arithmetic wrapping. +const SNAP_CASES: [SnapCase; 5] = [ + (12, 7, (0, 1), &[3, 9]), + (64, 5, (0, 0), &[0, 40]), + (9, 0, (0, 1), &[1, 8]), + (20, 3, (1, 0), &[12]), + (9, u32::MAX - 9, (0, 1), &[8]), +]; + +/// A rows launch's copies are the bits a chunk ending at their rows leaves: +/// copy `k`'s lane is the write lane, and its ring the ring, of the rows +/// launch over rows `0..=after[k]` of the same inputs; and the copies move +/// none of the launch's own outputs (numerators, partials, state, ring, +/// fault word). +pub fn delta_rows_snap(h: &Harness<'_>, fixture: &Fixture) -> Clause { + let fx = fixed(fixture)?; + let lane_len = elems(state_elems(1), "a lane")?; + let mut rng = oracle::Rng::new(0xD0A7); + let mut parts = Vec::new(); + for (rows, pos, lanes, after) in SNAP_CASES { + let c = chunk(&mut rng, rows, pos, lanes)?; + let at = format!("{rows} rows from {pos}, lanes {lanes:?}, copies after {after:?}"); + let (got, copies) = run_rows_copied(h, &fx, &c, after)?; + if let Some(f) = got.fault { + return Err(format!("{at}: raised {f}")); + } + let plain = run_rows(h, &fx, &c)?; + same_out(&got, &plain).map_err(|e| format!("{at}, against no copies: {e}"))?; + let ring_len = c.ring.len(); + for (k, &r) in after.iter().enumerate() { + let end = run_rows(h, &fx, &head_of(&c, r as usize + 1)?)?; + let lane = &end.state[c.lane_to * lane_len..(c.lane_to + 1) * lane_len]; + same_bits( + &format!("{at}: copy {k}'s lane"), + &copies.lane[k * lane_len..(k + 1) * lane_len], + lane, + )?; + same_bits( + &format!("{at}: copy {k}'s ring"), + &copies.ring[k * ring_len..(k + 1) * ring_len], + &end.ring, + )?; + } + parts.push(format!("{rows}@{pos} after {after:?}")); + } + Ok(Verdict::Noted(format!( + "rows@first position {}: each copy is the chunk ending at its row, bit for bit, and \ + the launch's own outputs are the launch's without copies", + parts.join(", ") + ))) +} + /// What one poisoned chunk raises and where the NaN lands. struct Poison { what: &'static str, @@ -1415,6 +1553,7 @@ fn delta_ref_launch( partials: P_ROW, }, eps: 1e-6, + snaps: None, }, ) } @@ -1980,6 +2119,7 @@ pub fn table<'a>(ctx: &'a ClauseCtx<'a>) -> Vec<(&'static str, ClauseRun<'a>)> { vec![ ("delta_rows", Box::new(|| ctx.with(delta_rows))), ("delta_rows_nan", Box::new(|| ctx.with(delta_rows_nan))), + ("delta_rows_snap", Box::new(|| ctx.with(delta_rows_snap))), ("norm_rope_rows", Box::new(|| ctx.with(norm_rope_rows))), ("attn_rows_v0", Box::new(|| ctx.with(attn_rows_v0))), ("attn_prefill", Box::new(|| ctx.with(attn_prefill))), diff --git a/crates/metal/src/bin/gate_metal_qwen35moe_kernels/mutants.rs b/crates/metal/src/bin/gate_metal_qwen35moe_kernels/mutants.rs index da218a6fe..adfa6d380 100644 --- a/crates/metal/src/bin/gate_metal_qwen35moe_kernels/mutants.rs +++ b/crates/metal/src/bin/gate_metal_qwen35moe_kernels/mutants.rs @@ -601,6 +601,26 @@ pub const MUTANTS: &[Mutant] = &[ 1, )], }, + Mutant { + name: "snap_row_late", + clause: "delta_rows_snap", + why: "the rows launch copies the state lane one row after the row its copy names", + edits: &[( + "if (copy < copies && r == snaps->after[copy]) {", + "if (copy < copies && r == snaps->after[copy] + 1u) {", + 1, + )], + }, + Mutant { + name: "snap_ring_dropped", + clause: "delta_rows_snap", + why: "the rows launch leaves a copy's conv ring unwritten", + edits: &[( + " snap[(ulong)k * snaps.ring_len + s * args.ch + c] = v;", + " (void)v;", + 1, + )], + }, Mutant { name: "rows_lane_parity", clause: "delta_rows", diff --git a/crates/metal/src/bin/metal_rows_time/mac.rs b/crates/metal/src/bin/metal_rows_time/mac.rs index ac6524bbb..b629809b3 100644 --- a/crates/metal/src/bin/metal_rows_time/mac.rs +++ b/crates/metal/src/bin/metal_rows_time/mac.rs @@ -615,6 +615,7 @@ fn delta_arms(run: &Run, device: &MetalDevice, kernels: &Kernels) -> Result<(), pos: 100, strides, eps: 1e-6, + snaps: None, }, ) })?; diff --git a/crates/metal/src/kernels/delta.rs b/crates/metal/src/kernels/delta.rs index 16b45ef9e..01dcde7f2 100644 --- a/crates/metal/src/kernels/delta.rs +++ b/crates/metal/src/kernels/delta.rs @@ -12,7 +12,9 @@ //! every row in parallel into a prepared row (the caller's scratch window), the //! recurrence over the rows in order, and then the numerators and partial sums //! in parallel. The ring commit follows the state-free half, which is the last -//! reader of the pre-chunk slots. +//! reader of the pre-chunk slots. A launch may also leave copies after rows +//! inside its chunk ([`DeltaSnaps`]): each the state lane and the ring a chunk +//! ending at that row leaves, the chunk's own outputs unmoved. use bloomery_device::{DeviceError, Window, WindowMut}; @@ -22,7 +24,7 @@ use crate::fault::MetalFaultWord; use crate::library::{Library, Pipeline}; -use super::delta_geom::{CONV_TAPS, RING_ROWS}; +use super::delta_geom::{CONV_TAPS, RING_ROWS, SNAPS}; use super::geometry::{HEAD_COLS, PARTIALS_PER_HEAD, VALUE_HEADS}; use super::rows::check_strided; use super::{Kernels, check_window_len, pack, word}; @@ -394,6 +396,21 @@ impl DeltaRowStrides { } } +/// The copies a rows launch leaves after rows inside its chunk: copy `k` is +/// the state lane and the ring a chunk ending after row `after[k]` leaves, +/// bit for bit, written into the `k`-th lane and ring of these windows. +pub struct DeltaSnaps<'w> { + /// The rows the copies follow, ascending, each below the chunk's rows; + /// one to [`SNAPS`] of them. + pub after: &'w [u32], + /// One lane of the state a copy, laid out as a lane of the state: + /// `after.len()` times [`state_elems`]`(1)` values. + pub lane: WindowMut<'w, f32, MetalDevice>, + /// One ring a copy, laid out as the ring: `after.len()` times + /// [`ring_elems`] values. + pub ring: WindowMut<'w, f32, MetalDevice>, +} + /// One gated delta step over a chunk of rows, rows in order in one launch. /// The state is read from the lane the word at `lane_at` names before the /// first row and written to the lane the word at `lane_to` names after the @@ -461,15 +478,20 @@ pub struct DeltaRowsParams<'w> { pub strides: DeltaRowStrides, /// The L2 norm's epsilon. pub eps: f32, + /// The copies after rows inside the chunk; none takes none. + pub snaps: Option>, } /// The pipelines of the rows launch: the state-free half, the recurrence, the -/// numerators and partials, and the ring commit. +/// numerators and partials, and the ring commit; and the recurrence and the +/// ring's copy of a launch that leaves copies. pub(crate) struct RowsPipelines { prep: Pipeline, rows: Pipeline, post: Pipeline, ring: Pipeline, + rows_snap: Pipeline, + ring_snap: Pipeline, } impl RowsPipelines { @@ -479,6 +501,8 @@ impl RowsPipelines { rows: library.pipeline("gdn_rows")?, post: library.pipeline("gdn_post_rows")?, ring: library.pipeline("gdn_ring_rows")?, + rows_snap: library.pipeline("gdn_rows_snap")?, + ring_snap: library.pipeline("gdn_ring_snap")?, }) } } @@ -550,16 +574,66 @@ fn check_rows(p: &DeltaRowsParams<'_>) -> Result<(), DeviceError> { range.start.get() ))); } + if let Some(snaps) = &p.snaps { + check_snaps(snaps, p.rows, ring)?; + } Ok(()) } +/// The copies of a chunk of `rows` rows whose ring is `ring` values: one to +/// [`SNAPS`] rows, ascending, inside the chunk, and a lane and a ring a copy. +fn check_snaps(s: &DeltaSnaps<'_>, rows: usize, ring: usize) -> Result<(), DeviceError> { + let n = s.after.len(); + if n == 0 || n > SNAPS { + return Err(refuse(format!( + "copies after {n} rows; a launch takes 1 to {SNAPS}" + ))); + } + if !s.after.windows(2).all(|w| w[0] < w[1]) { + return Err(refuse(format!( + "copies after rows {:?}, not ascending", + s.after + ))); + } + if let Some(&r) = s.after.iter().find(|&&r| r as usize >= rows) { + return Err(refuse(format!( + "a copy after row {r} of a chunk of {rows} rows" + ))); + } + let lane = state_elems(1).ok_or_else(|| refuse("a lane of the state does not fit"))?; + let lanes = n + .checked_mul(lane) + .ok_or_else(|| refuse(format!("{n} copies of a lane do not fit")))?; + let rings = n + .checked_mul(ring) + .ok_or_else(|| refuse(format!("{n} copies of the ring do not fit")))?; + word(WHAT, "a copy's lane length", lane)?; + word(WHAT, "a copy's ring length", ring)?; + check_window_len(WHAT, "the copies' lanes", s.lane.len(), lanes)?; + check_window_len(WHAT, "the copies' rings", s.ring.len(), rings) +} + +/// The copies' argument block: their count, the values of a copy's lane and +/// of its ring, and the row each follows (the words past the count are 0). +fn snap_args(s: &DeltaSnaps<'_>, ring: usize) -> Result<[[u8; 4]; 3 + SNAPS], DeviceError> { + let lane = state_elems(1).ok_or_else(|| refuse("a lane of the state does not fit"))?; + let mut words = [0u32; 3 + SNAPS]; + words[0] = word(WHAT, "the copy count", s.after.len())?; + words[1] = word(WHAT, "a copy's lane length", lane)?; + words[2] = word(WHAT, "a copy's ring length", ring)?; + words[3..3 + s.after.len()].copy_from_slice(s.after); + Ok(words.map(u32::to_ne_bytes)) +} + impl Kernels { /// Encode the step over `p.rows` rows: the state-free half, a memory /// barrier, the ring commit with the recurrence, a barrier, and the /// numerators and partial sums. The shape and every window are checked /// before any dispatch, and a refused launch encodes nothing. The ring /// commit and the recurrence read and write disjoint windows, so they may - /// overlap. + /// overlap. With copies, the ring's copies are encoded beside the + /// state-free half (both only read the ring, before the commit's barrier) + /// and the recurrence writes the lanes' copies as it passes their rows. #[must_use = "a refused launch encoded nothing"] pub fn delta_rows( &self, @@ -596,6 +670,18 @@ impl Kernels { .map(u32::to_ne_bytes) .collect::>(); let args: [[u8; 4]; 22] = args.try_into().expect("twenty-two words"); + let ring_len = ring_elems(p.n_k).expect("check_rows fitted the ring"); + let snap = match &p.snaps { + Some(s) => Some((s, snap_args(s, ring_len)?)), + None => None, + }; + let ring_args = [ + word(WHAT, "channels", ch)?, + word(WHAT, "rows", p.rows)?, + word(WHAT, "the x stride", st.x)?, + p.pos, + ] + .map(u32::to_ne_bytes); encoder.set_pipeline(&self.delta_rows.prep); let mut slot = Bindings::new(encoder); @@ -611,14 +697,22 @@ impl Kernels { slot.set(Binding::bytes(pack(&args))?)?; slot.dispatch_groups(p.n_k, p.rows, GROUP_THREADS)?; + if let Some((s, words)) = &snap { + encoder.set_pipeline(&self.delta_rows.ring_snap); + let mut slot = Bindings::new(encoder); + slot.set(Binding::of(&p.x)?)?; + slot.set(Binding::of_mut(&p.ring)?)?; + slot.set(Binding::of_mut(&s.ring)?)?; + slot.set(Binding::bytes(pack(&ring_args))?)?; + slot.set(Binding::bytes(pack(words))?)?; + slot.dispatch_groups( + ch.div_ceil(RING_THREADS), + s.after.len() * RING_ROWS, + RING_THREADS, + )?; + } + encoder.memory_barrier(); - let ring_args = [ - word(WHAT, "channels", ch)?, - word(WHAT, "rows", p.rows)?, - word(WHAT, "the x stride", st.x)?, - p.pos, - ] - .map(u32::to_ne_bytes); encoder.set_pipeline(&self.delta_rows.ring); let mut slot = Bindings::new(encoder); slot.set(Binding::of(&p.x)?)?; @@ -630,7 +724,11 @@ impl Kernels { RING_THREADS, )?; - encoder.set_pipeline(&self.delta_rows.rows); + let rows = match &snap { + Some(_) => &self.delta_rows.rows_snap, + None => &self.delta_rows.rows, + }; + encoder.set_pipeline(rows); let mut slot = Bindings::new(encoder); slot.set(Binding::of_mut(&p.prep)?)?; slot.set(Binding::of(&p.lane)?)?; @@ -638,6 +736,10 @@ impl Kernels { slot.set(Binding::of_mut(&p.g)?)?; slot.set(Binding::bytes(pack(&args))?)?; slot.set(p.faults.binding())?; + if let Some((s, words)) = &snap { + slot.set(Binding::of_mut(&s.lane)?)?; + slot.set(Binding::bytes(pack(words))?)?; + } slot.dispatch_groups(GROUPS, 1, GROUP_THREADS)?; encoder.memory_barrier(); diff --git a/crates/metal/src/kernels/delta_geom.rs b/crates/metal/src/kernels/delta_geom.rs index 0142bde05..de69af8dc 100644 --- a/crates/metal/src/kernels/delta_geom.rs +++ b/crates/metal/src/kernels/delta_geom.rs @@ -9,3 +9,7 @@ pub const PASS_ROWS: usize = 8; /// Slots of the conv ring: the predecessors plus the positions of one pass. pub const RING_ROWS: usize = CONV_ROWS + PASS_ROWS; + +/// Copies of the state lane and the conv ring one rows launch may take after +/// rows inside it: a prompt call's first and last message start. +pub const SNAPS: usize = 2; diff --git a/crates/metal/src/kernels/mod.rs b/crates/metal/src/kernels/mod.rs index 8137c47e3..9aff4d868 100644 --- a/crates/metal/src/kernels/mod.rs +++ b/crates/metal/src/kernels/mod.rs @@ -49,6 +49,7 @@ pub use attention::{ pub use delta::DeltaParams; pub use delta::DeltaRowStrides; pub use delta::DeltaRowsParams; +pub use delta::DeltaSnaps; pub use delta::HeadMap; pub use delta::Q_SCALE; pub use delta::channels; @@ -60,6 +61,7 @@ pub use delta::state_index; pub use delta_geom::CONV_TAPS; pub use delta_geom::PASS_ROWS; pub use delta_geom::RING_ROWS; +pub use delta_geom::SNAPS; pub use embed::EmbedParams; pub use experts::{FoldParams, PairParams, SharedPlan}; pub use experts_geom::{EXPERT_K, EXPERT_SLOTS, ROUTED_IDS, SHARED_SLOT}; diff --git a/crates/metal/src/server.rs b/crates/metal/src/server.rs index 695f5fa25..fed88f8cb 100644 --- a/crates/metal/src/server.rs +++ b/crates/metal/src/server.rs @@ -6,14 +6,16 @@ //! What it maps: //! //! - `prefill` is the target's prompt call, its pick dropped, and one line -//! on stderr naming its positions, its ids and its wall; `next` is one -//! step, [`Want::Argmax`] for a greedy request (the session chains the -//! pick into the token ahead) and [`Want::Logits`] for one that samples, -//! the row copied into the server's buffer; -//! - `prefill_splits` returns every mark. The target's prompt runs one -//! host-fed step a position, so a cut anywhere in the call is bit-neutral, -//! and each call's end becomes a checkpoint. Splitting at every mark makes -//! the start of each chat message a point; +//! on stderr naming its positions, its ids, the marks it kept and its +//! wall; `next` is one step, [`Want::Argmax`] for a greedy request (the +//! session chains the pick into the token ahead) and [`Want::Logits`] for +//! one that samples, the row copied into the server's buffer; +//! - `prefill_marked` is that call with the server's marks handed to the +//! target ([`Target::prompt_marked`]), and `prefill_splits` cuts nowhere: +//! the target keeps each mark as a point inside the call, so the start of +//! each chat message is one without a call of its own. A cut is not +//! bit-neutral here — a part of at most a step's window steps where the +//! one call passes — so the server asks for none; //! - `keepable`, `cut`, `reset` and `ctx_max` are the target's own rules: a //! prefix the target cannot keep is a reset, and why is the target's //! [`runtime::seqstate::Why`] word. @@ -27,6 +29,7 @@ use std::sync::Arc; use std::time::Instant; use runtime::{Out, Target, Want}; +use serve::media::MediaFeed; use serve::{Engine, EngineError, EngineProps, Tokenizer}; /// A target served with its vocabulary. @@ -73,24 +76,40 @@ impl Engine for TargetEngine { Arc::clone(&self.vocab) } - /// Every mark, in order. The target's prompt is one host-fed step a - /// position, so ending a call at a mark is bit-neutral, and that end is - /// a checkpoint. A chat message's start, passed in as a mark, becomes one. - fn prefill_splits(&self, _first: usize, _end: usize, marks: &[usize]) -> Vec { - marks.to_vec() + fn prefill(&mut self, ids: &[u32]) -> Result<(), EngineError> { + self.prefill_marked(ids, &[], &[]) } - fn prefill(&mut self, ids: &[u32]) -> Result<(), EngineError> { + /// The prompt call with a point at each mark ([`Target::prompt_marked`]). + /// A mark that is no `u32` position is the server's error, by name; the + /// target refuses one out of order or outside the call. The default + /// refuses a call that carries images: this engine takes none. + fn prefill_marked( + &mut self, + ids: &[u32], + media: &[MediaFeed], + marks: &[usize], + ) -> Result<(), EngineError> { + if !media.is_empty() { + return self.prefill_media(ids, media); + } if ids.is_empty() { return Ok(()); } + let at = marks + .iter() + .map(|&u| u32::try_from(u)) + .collect::, _>>() + .map_err(|_| { + EngineError(format!("{}: a prompt call marked at {marks:?}", self.what)) + })?; let from = self.target.pos(); let t0 = Instant::now(); self.target - .prompt(ids, Want::Argmax) + .prompt_marked(ids, &at, Want::Argmax) .map_err(engine_error)?; eprintln!( - "{}: prompt call {from}..{}: {} ids in {:.1} ms", + "{}: prompt call {from}..{}: {} ids, points at {at:?}, in {:.1} ms", self.what, self.target.pos(), ids.len(), @@ -171,10 +190,12 @@ mod tests { use super::TargetEngine; /// A target whose every step writes a row of `row` zeros and picks 0; - /// a prompt of no ids is refused, as the session's is. + /// a prompt of no ids is refused, as the session's is. `calls` holds + /// each prompt's id count. struct Rows { row: Vec, pos: u32, + calls: Vec, } impl Target for Rows { @@ -193,6 +214,7 @@ mod tests { return Err(io::Error::other("a prompt of no ids")); } self.pos += ids.len() as u32; + self.calls.push(ids.len()); Ok(Out::Argmax(0)) } @@ -226,6 +248,7 @@ mod tests { let target = Rows { row: vec![0.0; row], pos: 0, + calls: Vec::new(), }; TargetEngine::new( target, @@ -253,11 +276,27 @@ mod tests { } #[test] - fn prefill_splits_returns_every_mark() { + fn marks_reach_the_target_in_one_call() { let n = MockTokenizer.n_vocab(); - let e = engine(n, n).expect("a head of the vocabulary's rows"); + let mut e = engine(n, n).expect("a head of the vocabulary's rows"); let marks = [2usize, 7]; - assert_eq!(Engine::prefill_splits(&e, 0, 10, &marks), marks); - assert_eq!(Engine::prefill_splits(&e, 1, 4, &[]), Vec::::new()); + assert_eq!( + Engine::prefill_splits(&e, 0, 10, &marks), + Vec::::new() + ); + + // The engine hands the marks to the target's prompt; this target + // keeps them by the default, a prompt a part. + e.prefill_marked(&[1; 10], &[], &marks) + .expect("a call marked inside it"); + assert_eq!(e.target.calls, [2, 5, 3]); + e.prefill(&[1; 4]).expect("a call with no marks"); + assert_eq!(e.target.calls, [2, 5, 3, 4]); + + // A mark that is no position of the target is the server's error. + let refused = e + .prefill_marked(&[1; 3], &[], &[usize::MAX]) + .expect_err("a mark past u32"); + assert!(refused.0.contains("marked at"), "{}", refused.0); } } diff --git a/crates/runtime/src/lib.rs b/crates/runtime/src/lib.rs index b8289898b..98e36f394 100644 --- a/crates/runtime/src/lib.rs +++ b/crates/runtime/src/lib.rs @@ -199,6 +199,41 @@ pub trait Target { /// read back after the last; the target stands `ids.len()` positions on. fn prompt(&mut self, ids: &[u32], want: Want) -> Result, Self::Error>; + /// [`Target::prompt`] that leaves each of `marks` keepable: a later + /// [`Target::cut`] to one keeps it whole. The marks are positions, + /// ascending, each past [`Target::pos`] and at most the call's end. The + /// default runs a [`Target::prompt`] a part, cut at each mark before the + /// end: a target whose calls' ends are its points keeps a mark by ending a + /// call there. A target that keeps a position inside a call answers here + /// with one call. A mark out of order or outside the call is the caller's + /// breach: the default panics by name. + fn prompt_marked( + &mut self, + ids: &[u32], + marks: &[u32], + want: Want, + ) -> Result, Self::Error> { + let from = u64::from(self.pos()); + let to = from + ids.len() as u64; + let mut at = 0; + let mut prev = from; + for &u in marks { + let u = u64::from(u); + assert!( + prev < u && u <= to, + "a prompt of positions {from}..{to} marked at {marks:?}: marks go ascending, \ + past the call's start and at most its end" + ); + prev = u; + let i = usize::try_from(u - from).expect("a mark inside the call is an index of it"); + if i < ids.len() { + self.prompt(&ids[at..i], Want::Argmax)?; + at = i; + } + } + self.prompt(&ids[at..], want) + } + /// Run `id` at [`Target::pos`], one position, and read back after it. fn step(&mut self, id: u32, want: Want) -> Result, Self::Error>; diff --git a/crates/serve/src/engine.rs b/crates/serve/src/engine.rs index 3d85273c9..b4834b030 100644 --- a/crates/serve/src/engine.rs +++ b/crates/serve/src/engine.rs @@ -15,8 +15,12 @@ //! cache already holds and the engine can keep ([`Engine::keepable`]): `cut(k)`, //! or `reset` when `k` is 0; then `prefill(ids[k..n-1])`, and `next(ids[n-1])` //! yields the first generated token and every later `next(prev)` the one after -//! it. An engine that names a prompt quantum `q` ([`Engine::prompt_quantum`]) -//! may see that `prefill` arrive instead as calls ending at `k + j·q`, the grid +//! it. The server runs each prompt call as [`Engine::prefill_marked`] (whose +//! default is that `prefill`), handing it the call's message starts, where a +//! later request may be cut back to, and cutting the call where +//! [`Engine::prefill_splits`] asks. An engine that names a prompt quantum `q` +//! ([`Engine::prompt_quantum`]) may see that `prefill` arrive instead as calls +//! ending at `k + j·q`, the grid //! starting again at each cut [`Engine::prefill_splits`] asks for, one a round //! between the other busy slots' decode rounds; the quantum's contract makes //! the request's bits those of the one call. `timings.cache_n` and `usage.prompt_tokens_details.cached_tokens` are @@ -539,11 +543,37 @@ pub trait Engine: Send { /// that the `marks` (ascending, each inside the call) a later request may /// be cut back to stay keepable: a subset of `marks`, ascending. The /// default cuts nowhere: an engine whose cache keeps every position needs - /// no cut. + /// no cut, and one that keeps a mark inside a call takes it from + /// [`Engine::prefill_marked`]. fn prefill_splits(&self, first: usize, end: usize, marks: &[usize]) -> Vec { let _ = (first, end, marks); Vec::new() } + /// The call the server runs for every part of a prompt call: `ids` from + /// the current position, carrying `media` as [`Engine::prefill_media`]'s + /// feeds, with the prompt call's `marks` that fall in this part (the + /// positions past its first and up to its end, ascending), each a + /// position a later request may be cut back to. A mark at the part's end + /// is that end: a cut [`Engine::prefill_splits`] asked for, or a quantum's. + /// The default is [`Engine::prefill`], or [`Engine::prefill_media`] when + /// `media` is not empty, the marks unread: an engine whose cache keeps + /// every position, or that cuts where a mark needs it, keeps them already. + /// An engine that keeps a position inside a call only where it is told + /// answers here instead and cuts nowhere: one call that keeps its marks + /// in place of a call a mark. + fn prefill_marked( + &mut self, + ids: &[u32], + media: &[MediaFeed], + marks: &[usize], + ) -> Result<(), EngineError> { + let _ = marks; + if media.is_empty() { + self.prefill(ids) + } else { + self.prefill_media(ids, media) + } + } /// The call length `q` at which this engine's prompt work is bit-neutral /// to cut: the calls `ids[a..a + k·q]` then `ids[a + k·q..b]` leave the /// same bits as the one call `ids[a..b]` (every cell of the cache, every diff --git a/crates/serve/src/genloop.rs b/crates/serve/src/genloop.rs index c98cc3507..59315fdb3 100644 --- a/crates/serve/src/genloop.rs +++ b/crates/serve/src/genloop.rs @@ -1175,24 +1175,25 @@ impl Slot { } } - /// Where the engine cuts the prompt call `ids[from..to]` into calls - /// ([`Engine::prefill_splits`]) at the messages the prompt opens - /// ([`Tokenizer::user_start`]): its first and its last inside the range. - /// An answer not among those marks in order is the engine's error. - fn marked_cuts(&self, ids: &[u32], from: usize, to: usize) -> Result, EngineError> { + /// The marks of the prompt call `ids[from..to]`, the messages the prompt + /// opens ([`Tokenizer::user_start`]): its first and its last inside the + /// range; and where the engine cuts the call into calls among them + /// ([`Engine::prefill_splits`]). An answer not among those marks in + /// order is the engine's error. + fn marked_cuts(&self, ids: &[u32], from: usize, to: usize) -> Result { let marks = message_starts(&ids[..to], &self.vocab.user_start(), from); - let at = if marks.is_empty() { + let cuts = if marks.is_empty() { Vec::new() } else { self.engine.prefill_splits(from, to, &marks) }; - if !at.iter().all(|u| marks.contains(u)) || !at.windows(2).all(|w| w[0] < w[1]) { + if !cuts.iter().all(|u| marks.contains(u)) || !cuts.windows(2).all(|w| w[0] < w[1]) { return Err(EngineError(format!( - "the engine cut the prompt call {from}..{to} at {at:?}, not among its marks \ + "the engine cut the prompt call {from}..{to} at {cuts:?}, not among its marks \ {marks:?} in order" ))); } - Ok(at) + Ok(Marks { marks, cuts }) } /// The note of a prompt call `first..end` the engine cut at its marks @@ -1203,17 +1204,18 @@ impl Slot { } } - /// `prefill` of the prompt's positions `at..at + ids.len()` that books - /// what it fed; a call that holds an image span (of `media`, the prompt's - /// feeds) is [`Engine::prefill_media`] with the feeds it holds, at its own - /// positions, and books their spans. A call that would cut a span is the - /// server's error: the cache and the prompt's last position never lie - /// inside one. + /// [`Engine::prefill_marked`] of the prompt's positions + /// `at..at + ids.len()` that books what it fed: the feeds of the image + /// spans it holds (of `media`, the prompt's feeds), at its own positions, + /// and the prompt call's `marks` that fall in it, past `at` and up to its + /// end. A call that would cut a span is the server's error: the cache and + /// the prompt's last position never lie inside one. fn prefill_range( &mut self, ids: &[u32], at: usize, media: &[MediaFeed], + marks: &[usize], ) -> Result<(), EngineError> { let end = at + ids.len(); let mut feeds = Vec::new(); @@ -1234,12 +1236,13 @@ impl Slot { ..f.clone() }); } + let inside: Vec = marks + .iter() + .copied() + .filter(|&u| at < u && u <= end) + .collect(); let held = std::mem::take(&mut self.held); - if feeds.is_empty() { - self.engine.prefill(ids)?; - } else { - self.engine.prefill_media(ids, &feeds)?; - } + self.engine.prefill_marked(ids, &feeds, &inside)?; self.held = held; self.held.ids.extend_from_slice(ids); self.held.media.extend(feeds.iter().map(|f| MediaSpan { @@ -1276,19 +1279,17 @@ impl Slot { } /// The prompt call's last call `ids` at the prompt's positions `at..` - /// (`media`, the prompt's feeds) and the step of its last id, booked: - /// [`Slot::prefill_range`] of all but the last id and [`Slot::next`] of - /// the last. A prompt call takes the checkpoints of its marks, so it is - /// never an answer run's [`Engine::next_after`]. + /// (`media`, the prompt's feeds; `marks`, the prompt call's) and the step + /// of its last id, booked: [`Slot::prefill_range`] of all but the last id + /// and [`Slot::next`] of the last. A prompt call takes the checkpoints of + /// its marks, so it is never an answer run's [`Engine::next_after`]. fn close_call( &mut self, - ids: &[u32], - at: usize, - media: &[MediaFeed], + (ids, at, media, marks): Call<'_>, logits: Option<&mut [f32]>, ) -> Result { let (&last, run) = ids.split_last().expect("the prompt's last id"); - self.prefill_range(run, at, media)?; + self.prefill_range(run, at, media, marks)?; self.next(last, logits) } @@ -1486,6 +1487,17 @@ impl Slot { } } +/// The marks of a prompt call ([`Slot::marked_cuts`]): its message starts, +/// ascending, and the engine's cuts among them. +struct Marks { + marks: Vec, + cuts: Vec, +} + +/// A call of a prompt: its ids, the position of the first, the prompt's +/// feeds and the prompt call's marks. +type Call<'a> = (&'a [u32], usize, &'a [MediaFeed], &'a [usize]); + /// The positions in `from + 1 .. ids.len()` where `marker` starts: the first /// and the last of them (one when they are the same). An empty marker marks /// nothing. @@ -1531,9 +1543,12 @@ pub(crate) struct Prompt { /// quantum's cut inside a span moves to the span's end. media: Vec, kept: Kept, + /// The prompt call's message starts, each call handed those that fall in + /// it ([`Engine::prefill_marked`]). + marks: Vec, /// The engine's cuts at the messages, which the prompt call's note names /// as an uninterrupted call's does: a quantum's cut keeps nothing more. - marks: Vec, + cuts: Vec, /// Every call's end, ascending; the last is the prompt call's, `n - 1`. ends: Vec, /// The calls run. @@ -2013,16 +2028,16 @@ impl Gen { Cow::Owned([ids, &self.lead].concat()) }; let to = ids.len() - 1; - let at = slot.marked_cuts(&ids, kept.cache_n, to)?; + let Marks { marks, cuts } = slot.marked_cuts(&ids, kept.cache_n, to)?; let mut first = kept.cache_n; - for &u in &at { - slot.prefill_range(&ids[first..u], first, media)?; + for &u in &cuts { + slot.prefill_range(&ids[first..u], first, media, &marks)?; first = u; } - slot.note_split(kept.cache_n, to, at); + slot.note_split(kept.cache_n, to, cuts); self.close( slot, - (&ids[first..], first, media), + (&ids[first..], first, media, &marks), kept, t0, p.return_progress, @@ -2085,8 +2100,8 @@ impl Gen { let t = Instant::now(); ids.extend_from_slice(&self.lead); let to = ids.len() - 1; - let marks = slot.marked_cuts(&ids, kept.cache_n, to)?; - let ends = quantum_ends(kept.cache_n, to, &marks, &media, q.get()); + let Marks { marks, cuts } = slot.marked_cuts(&ids, kept.cache_n, to)?; + let ends = quantum_ends(kept.cache_n, to, &cuts, &media, q.get()); self.tim = Timings { prompt_ms: ms_since(t), n_ctx: self.ctx_max, @@ -2101,6 +2116,7 @@ impl Gen { media, kept, marks, + cuts, ends, ran: 0, progress, @@ -2127,16 +2143,16 @@ impl Gen { let t = Instant::now(); p.ran += 1; if p.ran < p.ends.len() { - slot.prefill_range(&p.ids[from..to], from, &p.media)?; + slot.prefill_range(&p.ids[from..to], from, &p.media, &p.marks)?; self.tim.prompt_ms += ms_since(t); self.tim.prompt_n = to.min(self.n) - p.kept.cache_n; self.tim.n_past = to; return Ok(true); } - slot.note_split(p.kept.cache_n, to, std::mem::take(&mut p.marks)); + slot.note_split(p.kept.cache_n, to, std::mem::take(&mut p.cuts)); self.close( slot, - (&p.ids[from..], from, &p.media), + (&p.ids[from..], from, &p.media, &p.marks), p.kept, t, p.progress, @@ -2145,24 +2161,24 @@ impl Gen { Ok(false) } - /// The prompt's end: its last call `call` (the ids, their first position - /// and the prompt's feeds) fed and its last id stepped - /// ([`Slot::close_call`]), and the first token's logits read, its timings, its event and the first token - /// in hand: the lead the call carried ([`Gen::take_lead`]) then the pick - /// from that row, in hand as a pass's tokens are. Its engine time is what - /// the timings hold of it (nothing for a prompt run in one go) and the - /// time since `t`. + /// The prompt's end: its last call `call` (the ids, their first position, + /// the prompt's feeds and the prompt call's marks) fed and its last id + /// stepped ([`Slot::close_call`]), and the first token's logits read, its + /// timings, its event and the first token in hand: the lead the call + /// carried ([`Gen::take_lead`]) then the pick from that row, in hand as a + /// pass's tokens are. Its engine time is what the timings hold of it + /// (nothing for a prompt run in one go) and the time since `t`. fn close( &mut self, slot: &mut Slot, - (ids, at, media): (&[u32], usize, &[MediaFeed]), + call: Call<'_>, kept: Kept, t: Instant, progress: bool, sink: &mut dyn FnMut(Event<'_>) -> io::Result<()>, ) -> Result<(), GenError> { let n = self.n; - let greedy = slot.close_call(ids, at, media, self.logits_out())?; + let greedy = slot.close_call(call, self.logits_out())?; slot.prompts[slot.cur] = n; self.tim = Timings { prompt_n: n - kept.cache_n, @@ -2447,7 +2463,7 @@ impl Gen { }; Some(slot.next_after(ids, logits)?) } else { - slot.prefill_range(ids, slot.held.len(), &[])?; + slot.prefill_range(ids, slot.held.len(), &[], &[])?; None }; span.fed = end; @@ -2696,7 +2712,7 @@ mod cache_tests { CacheNote, Decoder, Engine, EngineError, SamplingParams, SaveInFlight, Saved, StateError, Tokenizer, }; - use crate::media::Held; + use crate::media::{Held, MediaFeed}; use crate::mock::{MockEngine, MockTokenizer, SaveStats}; use crate::reasoning::{ThinkEntry, ThinkTags}; use crate::sampling; @@ -2738,11 +2754,13 @@ mod cache_tests { } /// What a [`Probe`] saw: its notes, the length of every prefill call, - /// and the states it took and took back. + /// the marks every prompt call's part was handed, and the states it took + /// and took back. #[derive(Default)] struct Log { notes: Vec, calls: Vec, + marks: Vec>, /// What each request told [`Engine::will_reply`]. replies: Vec>, snapshots: usize, @@ -2752,7 +2770,8 @@ mod cache_tests { /// The mock engine under a cut rule like V4.1's CED hole: a position /// strictly inside a prompt call is not kept, and a cut there falls to the /// call's start. It caches its states (`budget` bytes) and, with `split`, - /// cuts a prompt call at every mark it is offered. + /// cuts a prompt call at every mark it is offered; with `marked`, it + /// keeps the marks a call is handed as positions inside that one call. struct Probe { inner: MockEngine, tok: Arc, @@ -2762,6 +2781,8 @@ mod cache_tests { split: bool, /// Cut one position past each mark instead: a defective engine. off_mark: bool, + /// Keep the marks [`Engine::prefill_marked`] hands it. + marked: bool, /// Refuse every state it is handed back. refuse: bool, /// How long each snapshot sleeps. @@ -2809,6 +2830,7 @@ mod cache_tests { budget, split, off_mark: false, + marked: false, refuse: false, snapshot_takes: Duration::ZERO, log: Arc::clone(log), @@ -2907,6 +2929,34 @@ mod cache_tests { (false, _) => Vec::new(), } } + fn prefill_marked( + &mut self, + ids: &[u32], + media: &[MediaFeed], + marks: &[usize], + ) -> Result<(), EngineError> { + self.log.lock().expect("the log").marks.push(marks.to_vec()); + if !self.marked { + return if media.is_empty() { + self.prefill(ids) + } else { + self.prefill_media(ids, media) + }; + } + self.prefill(ids)?; + let call = self.calls.pop().expect("the call just fed"); + assert!( + marks.iter().all(|&u| call.start < u && u <= call.end), + "probe: marks {marks:?} outside the call {call:?}" + ); + let mut start = call.start; + for &u in marks.iter().filter(|&&u| u < call.end) { + self.calls.push(start..u); + start = u; + } + self.calls.push(start..call.end); + Ok(()) + } fn note(&self, note: &CacheNote) { self.log.lock().expect("the log").notes.push(note.clone()); } @@ -3030,6 +3080,11 @@ mod cache_tests { [u1, u2 - u1, a.len() - 1 - u2], "the prompt call's pieces" ); + assert_eq!( + log.lock().expect("log").marks[..3], + [vec![u1], vec![u2], vec![]], + "each piece is handed the marks in it: its end" + ); let mut b = system.clone(); b.extend(enc("<｜User｜>a dog ran and the ")); let (kept, tb) = run(&mut slot, &b); @@ -3038,6 +3093,39 @@ mod cache_tests { assert_eq!(run(&mut fresh, &b).1, tb); } + /// An engine that cuts nowhere and keeps the marks it is handed gets the + /// prompt call as one call with its first and last message starts, and a + /// request that shares the prompt up to the first keeps it — with the + /// same ids a fresh prefill gives. + #[test] + fn a_marked_engine_takes_one_call_with_its_marks() { + let log = Arc::new(Mutex::new(Log::default())); + let probe = Probe { + marked: true, + ..Probe::new(0, false, &log) + }; + let mut slot = Slot::new(Box::new(probe)); + let system = enc("be brief and kind. "); + let mut a = system.clone(); + a.extend(enc( + "<｜User｜>the cat sat<｜Assistant｜>on the mat<｜User｜>and the ", + )); + run(&mut slot, &a); + let u1 = system.len(); + let u2 = a.len() - enc("<｜User｜>and the ").len(); + { + let log = log.lock().expect("log"); + assert_eq!(log.calls[..1], [a.len() - 1], "one prompt call"); + assert_eq!(log.marks[0], [u1, u2], "the call's marks"); + } + let mut b = system.clone(); + b.extend(enc("<｜User｜>a dog ran and the ")); + let (kept, tb) = run(&mut slot, &b); + assert_eq!(kept, u1, "the system prompt, a mark inside the first call"); + let (mut fresh, _) = Probe::slot(0, false); + assert_eq!(run(&mut fresh, &b).1, tb); + } + /// An engine that cuts a prompt call anywhere but at the marks it was /// offered fails the request by name; nothing is fed. #[test] diff --git a/crates/serve/src/swap.rs b/crates/serve/src/swap.rs index 20e3f9e54..7f3df2a3c 100644 --- a/crates/serve/src/swap.rs +++ b/crates/serve/src/swap.rs @@ -44,6 +44,7 @@ use crate::engine::{ CacheNote, Drafted, Engine, EngineError, EngineProps, ResidencyReset, Saved, SavedState, SlotRow, StateError, Tokenizer, }; +use crate::media::MediaFeed; /// Tokens a request takes on the engine before the next live request's turn: /// long enough that a turn's steps outweigh the two state copies that end it @@ -187,6 +188,20 @@ impl Engine for SwapEngine { fn prefill_splits(&self, first: usize, end: usize, marks: &[usize]) -> Vec { self.inner.prefill_splits(first, end, marks) } + /// The inner engine's call with its marks; a call that carries images is + /// this engine's [`Engine::prefill_media`], which takes none. + fn prefill_marked( + &mut self, + ids: &[u32], + media: &[MediaFeed], + marks: &[usize], + ) -> Result<(), EngineError> { + if media.is_empty() { + self.inner.prefill_marked(ids, media, marks) + } else { + self.prefill_media(ids, media) + } + } fn note(&self, note: &CacheNote) { self.inner.note(note); }