model, gpu: the Qwen plans count the stage rope table, and the placed loads hold it to the plan

crates/gpu/src/arch/qwen3moe/body.rs

@@ -31,6 +31,7 @@ use crate::host::swap_source::ResidencyGlue;
 use crate::linear::LinearKernels;
 use crate::model::{
     ChainBody, GpuModel, Instrumented, Rollback, RowHeads, SlotRange, SlotRows, Slots, block_count,
+    refuse_stage_rows_past,
 };
 use crate::q6k_sel::Q6kSelKernels;
 use crate::q38::Q38Kernels;
@@ -937,11 +938,13 @@ impl Body {
             .map(|_| KvPlanes::new(stream, &dims, kv))
             .collect::<Result<Vec<_>, _>>()?;
         let k = Kernels::load(gpu, needs_site_kernels(w, &plans, placed.is_some()))?;
-        let rope = RopeTable::new(&RopeSpec::window(hp.rope.base, hp.rope.dims))?;
+        let table = RopeTable::new(&RopeSpec::window(hp.rope.base, hp.rope.dims))?;
+        let rope = RopeRows::new(stream, &table, HEAD, dims.ctx)?;
         let forms = Forms::of(&plans, &dims);
         let (ub, placed) = match placed {
             None => (Ubatch::new(stream, dims, ubatch, Q32::of(&plans))?, None),
             Some((open, counted)) => {
+                refuse_stage_rows_past(what, open.plan, 0, rope.table.num_bytes() as u64)?;
                 let side = placed_bytes(&dims, plans.len(), open.rows)?;
                 let arena = counted_arena(stream, (dims, forms), ubatch, counted, side)?;
                 let ub = Ubatch::placed(stream, arena, ubatch)?;
@@ -951,7 +954,7 @@ impl Body {
         Ok(Body {
             prefill: Prefill::new(stream, dims, forms)?,
             ub,
-            rope: RopeRows::new(stream, &rope, HEAD, dims.ctx)?,
+            rope,
             s: Arena::new(stream, dims, 1)?,
             sp: StepParams::new(stream, false)?,
             slot_in: SlotIn::new(stream)?,

crates/gpu/src/arch/qwen3moe/body35.rs

@@ -112,6 +112,7 @@ use crate::linear::{self, LinearShape};
 use crate::media::{ImageSpan, MediaBody, MediaSpan, SeqRope};
 use crate::model::{
     ChainBody, GpuModel, HostServed, Instrumented, MAX_PASS_ROWS, Rollback, Rows, SlotRange, Slots,
+    refuse_stage_rows_past,
 };
 use crate::rope_table::{RopeSpec, RopeTable};
 use crate::site::{self, Order, file_site};
@@ -1665,6 +1666,7 @@ impl Body35 {
                 (ubatch_arena(stream, (d, forms), ubatch.get(), free)?, None)
             }
             Some((open, counted)) => {
+                refuse_stage_rows_past(WHAT, open.plan, 0, rope.table.num_bytes() as u64)?;
                 let side = placed_bytes(&d, plans.len(), open.rows)?;
                 let u = counted_arena(stream, (d, forms), ubatch.get(), counted, side)?;
                 (u, Some(Placed::new(gpu, w, &plans, &d, open)?))

crates/gpu/src/model.rs

@@ -1929,11 +1929,7 @@ pub fn refuse_scratch_past(
         .get(card)
         .map(|c| c.scratch_bytes)
         .ok_or_else(|| GpuError::shape(what, format!("the plan has no card {card}")))?;
-    let rows = plan
-        .cards
-        .get(card)
-        .map(|c| c.stage_rows_bytes)
-        .ok_or_else(|| GpuError::shape(what, format!("the plan totals no card {card}")))?;
+    let rows = stage_rows(what, plan, card)?;
     let term = term + rows;
     if made as u64 > term {
         return Err(GpuError::shape(
@@ -1947,6 +1943,38 @@ pub fn refuse_scratch_past(
     Ok(())
 }
 
+/// Refused by name, as `what`'s, when the stage's position rows the load
+/// made, `made` device bytes, pass what card `card` of `plan` counts for them
+/// (`CardTotals::stage_rows_bytes`, the rope tables at the context): the check
+/// of a body whose scratch the plan holds by a term of its own rather than
+/// through [`refuse_scratch_past`].
+pub(crate) fn refuse_stage_rows_past(
+    what: &'static str,
+    plan: &Plan<'_>,
+    card: usize,
+    made: u64,
+) -> Result<(), GpuError> {
+    let rows = stage_rows(what, plan, card)?;
+    if made > rows {
+        return Err(GpuError::shape(
+            what,
+            format!(
+                "the stage's position rows hold {made} device bytes; card {card}'s plan counts \
+                 {rows} for them"
+            ),
+        ));
+    }
+    Ok(())
+}
+
+/// The stage's position rows card `card` of `plan` counts.
+fn stage_rows(what: &'static str, plan: &Plan<'_>, card: usize) -> Result<u64, GpuError> {
+    plan.cards
+        .get(card)
+        .map(|c| c.stage_rows_bytes)
+        .ok_or_else(|| GpuError::shape(what, format!("the plan totals no card {card}")))
+}
+
 /// The placed load after its body, which took `built`: the output head the
 /// card carries, over the body's weights, and the model over all of it, with
 /// the load's phases.

crates/model/src/arch/qwen35moe/place.rs

@@ -64,8 +64,8 @@ use crate::placement::{
 };
 
 use runtime::stores::{
-    DELTA_LANES, delta_lane_bytes, dense_kv_bytes, kv_q8_row_bytes, kv_row_bytes, ple_ring_bytes,
-    recurrent_bytes, selecting_bytes,
+    DELTA_LANES, F32_BYTES, delta_lane_bytes, dense_kv_bytes, kv_q8_row_bytes, kv_row_bytes,
+    ple_ring_bytes, recurrent_bytes, selecting_bytes,
 };
 pub use runtime::stores::{PASS_ROWS, conv_ring_rows, ple_ring_rows};
 
@@ -794,12 +794,14 @@ impl PlanInputs {
         )?;
         let t = &plan.cards[0];
         // The draft's term is the reserve the expert rule spread within, so
-        // the bound and the budget count one number.
+        // the bound and the budget count one number. The stage rope rows are
+        // the target card's, beside the cache.
         let total = [
             t.dense_bytes,
             t.expert_bytes,
             t.rounding_bytes,
             t.kv_bytes,
+            t.stage_rows_bytes,
             t.scratch_bytes,
             t.context_bytes,
             reserve,
@@ -909,12 +911,14 @@ impl PlanInputs {
         plan.cards[0].grow_kv(rows);
         let t = &plan.cards[0];
         // As in [`PlanInputs::plan_mtp_with`]: the draft's term is the
-        // reserve the expert rule spread within.
+        // reserve the expert rule spread within, and the stage rope rows are
+        // the target card's.
         let total = [
             t.dense_bytes,
             t.expert_bytes,
             t.rounding_bytes,
             t.kv_bytes,
+            t.stage_rows_bytes,
             t.scratch_bytes,
             t.context_bytes,
             reserve,
@@ -1447,6 +1451,10 @@ pub struct KvLayout {
     select: Option<(usize, Vec<usize>)>,
     /// A plain attention layer's planes in q8_0's layout.
     q8: bool,
+    /// The stage rope table, in bytes a position: `rope.dimension_count`
+    /// f32. The body copies one table at load, not one per sequence. A media
+    /// sequence's own table is outside this count.
+    rope_row: u64,
 }
 
 impl KvLayout {
@@ -1486,6 +1494,7 @@ impl KvLayout {
             head_dim: hp.head_dim,
             select: exp.map(|e| (e.idx_dim, e.ratios.clone())),
             q8: false,
+            rope_row: hp.rope_dims as u64 * F32_BYTES,
         }
     }
 
@@ -1593,6 +1602,10 @@ impl KvBytes for KvLayout {
         };
         own + ple
     }
+
+    fn stage_position_bytes(&self, ctx_max: u64) -> u64 {
+        self.rope_row * ctx_max
+    }
 }
 
 // ------------------------------------------------------------- the MTP draft
@@ -2102,7 +2115,9 @@ fn draft_machine() -> Machine {
 }
 
 /// The MTP layer's store: every position's K and V
-/// (`runtime::stores::dense_kv_bytes`).
+/// (`runtime::stores::dense_kv_bytes`). The draft program reads the target's
+/// rope table and allocates none of its own, so the stage rows stay the
+/// trait's default.
 #[derive(Clone, Copy, Debug)]
 struct MtpKv {
     kv_heads: usize,
@@ -2141,6 +2156,7 @@ mod tests {
             head_dim: 256,
             select: Some((128, vec![0, 0, 4])),
             q8: false,
+            rope_row: 64 * 4,
         };
         assert_eq!(kv.layer_bytes(0, 4096), 3_596_288 + 9_437_200);
         assert_eq!(kv.layer_bytes(1, 4096), 3_596_288 + 9_437_200 + 696_320);
@@ -2177,6 +2193,7 @@ mod tests {
                 (0..48).map(|l| if l % 4 == 3 { 4 } else { 0 }).collect(),
             )),
             q8: false,
+            rope_row: 64 * 4,
         };
         let checkpoint = 36 * 3_596_288 + 696_320;
         assert_eq!(kv.state_bytes(0), 2 * checkpoint);
@@ -2227,6 +2244,7 @@ mod tests {
             head_dim: 256,
             select: None,
             q8: false,
+            rope_row: 64 * 4,
         };
         assert_eq!(layout(1).layer_bytes(0, 4096), one.0);
         assert_eq!(layout(8).layer_bytes(0, 4096), eight.0);
@@ -2250,6 +2268,7 @@ mod tests {
             head_dim: 128,
             select: None,
             q8: false,
+            rope_row: 128 * 4,
         };
         let f16 = plain.layer_bytes(0, 4096);
         let q8 = plain.in_q8().layer_bytes(0, 4096);
@@ -2265,11 +2284,59 @@ mod tests {
             head_dim: 256,
             select: Some((128, vec![4])),
             q8: false,
+            rope_row: 64 * 4,
         };
         let sel = selecting.layer_bytes(0, 4096);
         assert_eq!(selecting.in_q8().layer_bytes(0, 4096), sel);
     }
 
+    #[test]
+    fn the_stage_rope_row_is_rope_dims_not_the_head() {
+        use super::super::hparams::{FfnKind, Hparams, SectionOrder, Variant};
+
+        let hp = Hparams {
+            variant: Variant::Qwen35Moe,
+            n_layer: 1,
+            n_trunk: 1,
+            n_embd: 256,
+            n_head: 1,
+            n_head_kv: 1,
+            head_dim: 256,
+            rope_dims: 64,
+            rope_sections: [0, 0, 0, 0],
+            rope_order: SectionOrder::Tyxz,
+            rope_base: 1e7,
+            rms_eps: 1e-6,
+            n_vocab: 2,
+            n_ctx_train: 8,
+            n_expert: 0,
+            n_used: 0,
+            expert_ff: 0,
+            shared_ff: None,
+            ff: None,
+            conv: 4,
+            state: 1,
+            v_heads: 1,
+            k_heads: 1,
+            interval: 1,
+            kinds: vec![Kind::Attention],
+            ffns: vec![FfnKind::Dense],
+            exp: None,
+            defaults: Vec::new(),
+        };
+        let narrow = KvLayout::of(&hp);
+        assert_eq!(narrow.stage_position_bytes(8), 64 * 4 * 8);
+        assert_eq!(narrow.in_q8().stage_position_bytes(8), 64 * 4 * 8);
+        assert_eq!(
+            KvLayout::of(&Hparams {
+                rope_dims: 128,
+                ..hp
+            })
+            .stage_position_bytes(8),
+            128 * 4 * 8
+        );
+    }
+
     /// The card rule over Qwen3.8's routed stacks ([`super::routed_rule`]) on
     /// six synthetic layers of the file's per-expert shapes (512 experts; a
     /// Q4_K gate and up of 640 rows of 2560 values, a Q5_1 down of 2560 rows
@@ -3061,6 +3128,11 @@ mod tests {
             fn shadow_bytes(&self, layer: usize, ctx_max: u64) -> u64 {
                 self.kv.shadow_bytes(layer, ctx_max) * self.slots
             }
+
+            /// One copy for every slot: the slots read one stage's rows.
+            fn stage_position_bytes(&self, ctx_max: u64) -> u64 {
+                self.kv.stage_position_bytes(ctx_max)
+            }
         }
 
         /// The target's plan of `slots` slots: each layer's bytes times the
@@ -3173,6 +3245,7 @@ mod tests {
                 + t.expert_bytes
                 + t.rounding_bytes
                 + t.kv_bytes
+                + t.stage_rows_bytes
                 + t.scratch_bytes
                 + t.context_bytes
                 + t.reserve_bytes;

crates/model/src/arch/qwen3moe/place.rs

@@ -23,7 +23,7 @@ use crate::placement::{
     self, Card, KvBytes, Machine, ModelTensors, PlacementError, Plan, PlanLevers, Violation,
     checked, joined,
 };
-use runtime::stores::{kv_q8_row_bytes, kv_row_bytes};
+use runtime::stores::{F32_BYTES, kv_q8_row_bytes, kv_row_bytes};
 
 /// The qwen3moe program, whose card experts read the routed stacks
 /// ([`routed_rule`] of it): by part, a gate and an up
@@ -96,6 +96,10 @@ impl KvBytes for KvLayout {
             0
         }
     }
+
+    fn stage_position_bytes(&self, ctx_max: u64) -> u64 {
+        self.head_dim as u64 * F32_BYTES * ctx_max
+    }
 }
 
 /// Why a plan of a qwen3moe file was refused.
@@ -184,3 +188,37 @@ impl PlanInputs {
         checked(plan).map_err(PlaceError::Broken)
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::super::hparams::{Experts, Hparams, Rope, RopeMode, Score};
+    use super::{KvBytes, KvLayout};
+
+    #[test]
+    fn the_stage_rope_row_is_the_width_the_body_copies() {
+        let kv = KvLayout::of(&Hparams {
+            n_layer: 1,
+            n_embd: 128,
+            n_head: 1,
+            n_head_kv: 1,
+            head_dim: 128,
+            rope: Rope {
+                mode: RopeMode::Neox,
+                base: 1e6,
+                dims: 128,
+            },
+            rms_eps: 1e-6,
+            n_vocab: 2,
+            n_ctx_train: 128,
+            experts: Experts {
+                n_expert: 1,
+                n_used: 1,
+                ff: 1,
+                score: Score::Softmax,
+                weights_norm: true,
+            },
+        });
+        assert_eq!(kv.stage_position_bytes(10), 128 * 4 * 10);
+        assert_eq!(kv.in_q8().stage_position_bytes(10), 128 * 4 * 10);
+    }
+}