metal: message marks kept inside one prompt call (prefill_marked)

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,

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;
+}

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.

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,14 +239,17 @@ 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;
         let to = from + u32::try_from(m).expect("a pass's rows fit u32");
         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();

crates/metal/src/arch/qwen3moe/rows.rs

@@ -11,15 +11,19 @@
 //!
 //! 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};
 
 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<f32>,
     /// 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<SnapBufs>,
+}
+
+/// 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<f32>,
+    pub(super) ring: MetalBuffer<f32>,
+    /// Each layer's copies; none for an attention layer.
+    pub(super) at: Vec<Option<SnapAt>>,
+    /// The position copy `k` holds: the position after the row it follows.
+    pub(super) holds: [Option<u32>; SNAPS],
+}
+
+impl SnapBufs {
+    /// The copies of the delta layers of `arena`, opened on `device`.
+    pub(super) fn open(device: Alloc<'_>, arena: &Arena) -> Result<SnapBufs, SessionError> {
+        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,
         })
     }
 }