From 808ac31d92b32929e695129fb848eb338d2b2d3b Mon Sep 17 00:00:00 2001 From: midagedev Date: Thu, 15 Oct 2026 09:45:09 +0900 Subject: [PATCH 1/4] model, gpu: the Qwen plans count the stage rope table, and the placed loads hold it to the plan The Qwen3 and Qwen3.6/3.8 bodies copy one rope table of rope.dimension_count f32 a position at load (body.rs, body35.rs RopeRows::new); their KvLayout impls now count it through KvBytes::stage_position_bytes, so the plan's CardTotals::stage_rows_bytes carries it on the target card. The MTP draft reads the target's table and keeps the trait's default; the slot plan counts one table for every slot. The placed loads hold the table they made to that count by name (model::refuse_stage_rows_past, the stage-row half of refuse_scratch_past, which the GLM and MiMo loads call): a plan that counts no table, or half of one, refuses the load. The qwen4exp plan tests' boundaries and floors move by the 256 B a position; their PIN lines carry the derivations. Co-Authored-By: Claude Opus 5.5 --- crates/gpu/src/arch/qwen3moe/body.rs | 7 +- crates/gpu/src/arch/qwen3moe/body35.rs | 2 + crates/gpu/src/model.rs | 38 +++++++++-- crates/gpu/src/rope_table.rs | 6 ++ crates/model/src/arch/qwen35moe/place.rs | 81 ++++++++++++++++++++++-- crates/model/src/arch/qwen3moe/place.rs | 51 ++++++++++++++- crates/model/tests/qwen4exp_meta.rs | 36 ++++++++--- crates/model/tests/qwen4exp_mtp_meta.rs | 15 +++-- 8 files changed, 207 insertions(+), 29 deletions(-) diff --git a/crates/gpu/src/arch/qwen3moe/body.rs b/crates/gpu/src/arch/qwen3moe/body.rs index 4a5e37c3f..9a761efb6 100644 --- a/crates/gpu/src/arch/qwen3moe/body.rs +++ b/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::, _>>()?; 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.bytes())?; 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)?, diff --git a/crates/gpu/src/arch/qwen3moe/body35.rs b/crates/gpu/src/arch/qwen3moe/body35.rs index e333dce4e..aaf38f66f 100644 --- a/crates/gpu/src/arch/qwen3moe/body35.rs +++ b/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.bytes())?; 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)?)) diff --git a/crates/gpu/src/model.rs b/crates/gpu/src/model.rs index deb34ca63..da2a38517 100644 --- a/crates/gpu/src/model.rs +++ b/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 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 { + 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. diff --git a/crates/gpu/src/rope_table.rs b/crates/gpu/src/rope_table.rs index 4c2e4087c..6c2b7d529 100644 --- a/crates/gpu/src/rope_table.rs +++ b/crates/gpu/src/rope_table.rs @@ -234,6 +234,12 @@ impl RopeRows { }) } + /// The table's device bytes. + #[must_use] + pub fn bytes(&self) -> u64 { + (self.table.len() * size_of::()) as u64 + } + /// The positions the table holds: the cache's rows. #[must_use] pub fn rows(&self) -> usize { diff --git a/crates/model/src/arch/qwen35moe/place.rs b/crates/model/src/arch/qwen35moe/place.rs index 74a128cd6..3c951090f 100644 --- a/crates/model/src/arch/qwen35moe/place.rs +++ b/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)>, /// 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,57 @@ 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); } + /// The stage rope rows, through [`KvLayout::of`]: `rope.dimension_count` + /// f32 a position, not the head width. The q8 cache keeps the same table. + /// A count of no row, or of half a row, is not the table the body copies. + #[test] + fn the_stage_rope_row_is_rope_dims_not_the_head() { + let narrow = layout(64, 256); + assert_eq!(narrow.stage_position_bytes(8), 64 * 4 * 8); + assert_eq!(narrow.in_q8().stage_position_bytes(8), 64 * 4 * 8); + assert_eq!(layout(128, 256).stage_position_bytes(8), 128 * 4 * 8); + } + + fn layout(rope_dims: usize, head_dim: usize) -> KvLayout { + use super::super::hparams::{FfnKind, Hparams, SectionOrder, Variant}; + KvLayout::of(&Hparams { + variant: Variant::Qwen35Moe, + n_layer: 1, + n_trunk: 1, + n_embd: head_dim, + n_head: 1, + n_head_kv: 1, + head_dim, + rope_dims, + 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(), + }) + } + /// 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 +3126,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 +3243,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; diff --git a/crates/model/src/arch/qwen3moe/place.rs b/crates/model/src/arch/qwen3moe/place.rs index 0596ebede..a15407806 100644 --- a/crates/model/src/arch/qwen3moe/place.rs +++ b/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 @@ -61,6 +61,9 @@ pub struct KvLayout { kv_heads: usize, head_dim: usize, 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. + rope_row: u64, } impl KvLayout { @@ -72,6 +75,7 @@ impl KvLayout { kv_heads: hp.n_head_kv, head_dim: hp.head_dim, q8: false, + rope_row: hp.rope.dims as u64 * F32_BYTES, } } @@ -96,6 +100,10 @@ impl KvBytes for KvLayout { 0 } } + + fn stage_position_bytes(&self, ctx_max: u64) -> u64 { + self.rope_row * ctx_max + } } /// Why a plan of a qwen3moe file was refused. @@ -184,3 +192,44 @@ 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}; + + /// `rope.dimension_count` f32 a position, through [`KvLayout::of`]. The + /// q8 cache keeps the same table. A count of no row, or of half a row, + /// is not the table the body copies. + #[test] + fn the_stage_rope_row_is_the_width_the_body_copies() { + let kv = layout(128); + assert_eq!(kv.stage_position_bytes(10), 128 * 4 * 10); + assert_eq!(kv.in_q8().stage_position_bytes(10), 128 * 4 * 10); + } + + fn layout(head_dim: usize) -> KvLayout { + KvLayout::of(&Hparams { + n_layer: 1, + n_embd: head_dim, + n_head: 1, + n_head_kv: 1, + head_dim, + rope: Rope { + mode: RopeMode::Neox, + base: 1e6, + dims: head_dim, + }, + rms_eps: 1e-6, + n_vocab: 2, + n_ctx_train: head_dim, + experts: Experts { + n_expert: 1, + n_used: 1, + ff: 1, + score: Score::Softmax, + weights_norm: true, + }, + }) + } +} diff --git a/crates/model/tests/qwen4exp_meta.rs b/crates/model/tests/qwen4exp_meta.rs index 461431ca8..66e2f57d8 100644 --- a/crates/model/tests/qwen4exp_meta.rs +++ b/crates/model/tests/qwen4exp_meta.rs @@ -177,16 +177,19 @@ const HOST_TABLES: u64 = 28_800_138_240; // PIN(2026-09-28): KV_AT and CARD_MAX_CTX past the verify's delta lanes [derived: each GDN layer's // state keeps runtime::stores::DELTA_LANES = 4 lanes with a u32 stamp each, 3 · 3,145,728 + 16 = // 9,437,200 B past the one lane, 339,739,200 B over the 36 layers; KV_AT was 246,554,624 and -// 1,061,298,176, CARD_MAX_CTX 1,520,312 and 619,339, each now 339,739,200 / 28,416 = 11,955.9 +// 1,061,298,176, CARD_MAX_CTX 1,520,312 and 619,339, each then 339,739,200 / 28,416 = 11,955.9 // positions lower, rounded to 11,956 (one fewer when the old boundary's slack passed 26,000 B)]. const KV_AT: [(u64, u64); 2] = [(4096, 586_293_824), (32_768, 1_401_037_376)]; // The largest context each card holds, as predicted [derived: (usable − margin − the dense -// granules − context − scratch − 469,901,888 B of recurrent bytes) over 28,416 B a position, the -// last pool counted whole]. Printed beside the boundary the test finds from the plan's own totals; -// the clause holds the placement to that boundary, not to this figure. PIN(2026-09-28): the -// scratch counts the ubatch walk's at U = 4,096 [derived: CARD_SCRATCH's 4,096 row − 64 MiB = -// 2,770,817,024 B over 28,416 B a position, 97,509.05 positions, rounded up to 97,510]. -const CARD_MAX_CTX: [(&str, u64); 2] = [("A6000", 1_410_846), ("3090", 509_873)]; +// granules − context − scratch − 469,901,888 B of recurrent bytes) over 28,672 B a position — +// the cache's 28,416 plus the stage rope table's 256 — the last pool counted whole, then one +// fewer]. Printed beside the boundary the test finds from the plan's own totals; the clause +// holds the placement to that boundary, not to this figure. PIN(2026-09-28): the scratch counts +// the ubatch walk's at U = 4,096 [derived: CARD_SCRATCH's 4,096 row − 64 MiB = 2,770,817,024 B +// over 28,416 B a position, 97,509.05 positions, rounded up to 97,510]. +// PIN(2026-10-15): the stage rope table joins the per-position term. Room 40,090,635,712 (A6000) +// and 14,488,604,096 (3090) over 28,672 is 1,398,250 and 505,322; the pin is one fewer. +const CARD_MAX_CTX: [(&str, u64); 2] = [("A6000", 1_398_249), ("3090", 505_321)]; // PIN(2026-09-28): the card's scratch, the m = 1 scratch and the ubatch walk's at 4,096 // positions [derived: 64 MiB + place::ubatch_scratch_bytes(4096) = 67,108,864 + 4,096 · // 668,277 + 32 MiB]. @@ -469,7 +472,12 @@ fn hw_qwen4exp_plan() { let fixed = c.dense_bytes + c.expert_bytes + c.rounding_bytes + c.scratch_bytes + c.context_bytes; let limit = at.usable_bytes(spec) - spec.margin_bytes; - let kv = |ctx: u64| -> u64 { (0..hp.n_layer).map(|l| inputs.kv.layer_bytes(l, ctx)).sum() }; + let kv = |ctx: u64| -> u64 { + (0..hp.n_layer) + .map(|l| inputs.kv.layer_bytes(l, ctx)) + .sum::() + + inputs.kv.stage_position_bytes(ctx) + }; let (mut lo, mut hi) = (1u64, KERNEL_POSITIONS); while lo < hi { let mid = lo + (hi - lo).div_ceil(2); @@ -848,7 +856,7 @@ fn hw_qwen4exp_card_plan() { use model::arch::models::headpick::HeadSet; use model::arch::qwen35moe::place::{self, Experts, MtpInputs, PlanInputs}; use model::placement::workstation::{A6000, CONTEXT, MARGIN, RTX_3090}; - use model::placement::{CardFormat, Device, Format, Plan, PlanLevers, Role}; + use model::placement::{CardFormat, Device, Format, KvBytes, Plan, PlanLevers, Role}; let mut o = String::new(); let mut b: Vec = Vec::new(); @@ -1005,11 +1013,15 @@ fn hw_qwen4exp_card_plan() { } for card in [A6000, RTX_3090] { let machine = place::machine(card, inputs.hp.n_layer, place::UBATCH_PLANNED); + // PIN(2026-10-15): the floor holds the stage rope table, 64 f32 a position, 256 B at + // 4,096 positions (1,048,576 B), beside the cache. + let stage = inputs.kv.stage_position_bytes(KV_AT[0].0); let floor = CARD_DENSE + CARD_ROUNDING + KV_AT[0].1 + CONTEXT + scratch_at(place::UBATCH_PLANNED) + + stage + MARGIN; let levers = PlanLevers { card_budget_bytes: Some(floor), @@ -1270,7 +1282,7 @@ fn hw_qwen4exp_q3_card_plan() { use model::arch::qwen35moe::mtp::BorrowedHead; use model::arch::qwen35moe::place::{self, Experts, MtpInputs, PlanInputs}; use model::placement::workstation::{A6000, CONTEXT, MARGIN, RTX_3090}; - use model::placement::{CardFormat, Device, Format, Plan, PlanLevers, Role}; + use model::placement::{CardFormat, Device, Format, KvBytes, Plan, PlanLevers, Role}; let mut o = String::new(); let mut b: Vec = Vec::new(); @@ -1418,11 +1430,15 @@ fn hw_qwen4exp_q3_card_plan() { } for card in [A6000, RTX_3090] { let machine = place::machine(card, inputs.hp.n_layer, place::UBATCH_PLANNED); + // PIN(2026-10-15): the floor holds the stage rope table, 64 f32 a position, 256 B at + // 4,096 positions (1,048,576 B), beside the cache. + let stage = inputs.kv.stage_position_bytes(KV_AT[0].0); let floor = CARD_DENSE_Q3 + CARD_ROUNDING_Q3 + KV_AT[0].1 + CONTEXT + scratch_at(place::UBATCH_PLANNED) + + stage + MARGIN; let levers = PlanLevers { card_budget_bytes: Some(floor), diff --git a/crates/model/tests/qwen4exp_mtp_meta.rs b/crates/model/tests/qwen4exp_mtp_meta.rs index fe30d81b2..516befce2 100644 --- a/crates/model/tests/qwen4exp_mtp_meta.rs +++ b/crates/model/tests/qwen4exp_mtp_meta.rs @@ -166,12 +166,14 @@ const DRAFT_KV_ROW: u64 = 2048; const DRAFT_MAP: u64 = 993_280; const LIST_ROWS: u32 = 40_960; // The largest context each card holds with the draft (full head), as predicted [derived: the plan -// test's CARD_MAX_CTX budget less 2,785,017,856 B of draft granules, over 28,416 + 2,048 B a -// position plus the draft program's arena, whose dense flash partials grow 198,144 B a 64-key -// segment (3,096 B a position) — the same budget over 33,560 B a position, its 4.86 MB of fixed -// rows beside the granules; the last pool counted whole]. Printed beside the boundary the test -// finds; the clause holds the plan to that boundary. -const CARD_MAX_CTX_MTP: [(&str, u64); 2] = [("A6000", 1_194_276), ("3090", 431_309)]; +// test's room after its fixed terms, less the draft granules 2,785,017,856 B, the map 993,280 B +// and the arena 53,526,700 B, over 30,720 B a position — the cache's 28,416, the draft store's +// 2,048 and the stage rope table's 256 — the last pool counted whole, then one fewer]. Printed +// beside the boundary the test finds; the clause holds the plan to that boundary. +// PIN(2026-10-15): the arena stays one fixed addend (mtp_arena_bytes at vocab 248,320). Room +// 37,251,097,876 (A6000) and 11,649,066,260 (3090) over 30,720 is 1,212,600 and 379,201; the pin +// is one fewer. +const CARD_MAX_CTX_MTP: [(&str, u64); 2] = [("A6000", 1_212_599), ("3090", 379_200)]; /// The plan with the shared draft (`place::PlanInputs::plan_mtp`): on each /// card at 4,096 and 32,768 positions the target's plan is `plan`'s field @@ -365,6 +367,7 @@ fn hw_qwen4exp_mtp_plan() { .sum::() + ctx * DRAFT_KV_ROW + model::arch::qwen35moe::place::mtp_arena_bytes(u64::from(full.draft.vocab)) + + model::placement::KvBytes::stage_position_bytes(&inputs.kv, ctx) }; let (mut lo, mut hi) = (1u64, KERNEL_POSITIONS); while lo < hi { From 5a9d496c7e6a0e874248c474054bfb247ac84de5 Mon Sep 17 00:00:00 2001 From: midagedev Date: Thu, 15 Oct 2026 10:30:31 +0900 Subject: [PATCH 2/4] model, gpu: re-pin the stage rope plans and share the table byte count CARD_PLANS and CARD_PLANS_Q3 take the derived rows for a 256 B stage rope table. Qwen3 counts that table from head_dim, and the placed loads read the buffer's own byte count. --- crates/gpu/src/arch/qwen3moe/body.rs | 2 +- crates/gpu/src/arch/qwen3moe/body35.rs | 2 +- crates/gpu/src/rope_table.rs | 6 -- crates/model/src/arch/qwen3moe/place.rs | 8 +- crates/model/tests/qwen4exp_meta.rs | 120 ++++++++++++------------ 5 files changed, 65 insertions(+), 73 deletions(-) diff --git a/crates/gpu/src/arch/qwen3moe/body.rs b/crates/gpu/src/arch/qwen3moe/body.rs index 9a761efb6..54356daa8 100644 --- a/crates/gpu/src/arch/qwen3moe/body.rs +++ b/crates/gpu/src/arch/qwen3moe/body.rs @@ -944,7 +944,7 @@ impl Body { 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.bytes())?; + 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)?; diff --git a/crates/gpu/src/arch/qwen3moe/body35.rs b/crates/gpu/src/arch/qwen3moe/body35.rs index aaf38f66f..50a2c2425 100644 --- a/crates/gpu/src/arch/qwen3moe/body35.rs +++ b/crates/gpu/src/arch/qwen3moe/body35.rs @@ -1666,7 +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.bytes())?; + 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)?)) diff --git a/crates/gpu/src/rope_table.rs b/crates/gpu/src/rope_table.rs index 6c2b7d529..4c2e4087c 100644 --- a/crates/gpu/src/rope_table.rs +++ b/crates/gpu/src/rope_table.rs @@ -234,12 +234,6 @@ impl RopeRows { }) } - /// The table's device bytes. - #[must_use] - pub fn bytes(&self) -> u64 { - (self.table.len() * size_of::()) as u64 - } - /// The positions the table holds: the cache's rows. #[must_use] pub fn rows(&self) -> usize { diff --git a/crates/model/src/arch/qwen3moe/place.rs b/crates/model/src/arch/qwen3moe/place.rs index a15407806..555af8b2e 100644 --- a/crates/model/src/arch/qwen3moe/place.rs +++ b/crates/model/src/arch/qwen3moe/place.rs @@ -61,9 +61,6 @@ pub struct KvLayout { kv_heads: usize, head_dim: usize, 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. - rope_row: u64, } impl KvLayout { @@ -75,7 +72,6 @@ impl KvLayout { kv_heads: hp.n_head_kv, head_dim: hp.head_dim, q8: false, - rope_row: hp.rope.dims as u64 * F32_BYTES, } } @@ -102,7 +98,7 @@ impl KvBytes for KvLayout { } fn stage_position_bytes(&self, ctx_max: u64) -> u64 { - self.rope_row * ctx_max + self.head_dim as u64 * F32_BYTES * ctx_max } } @@ -198,7 +194,7 @@ mod tests { use super::super::hparams::{Experts, Hparams, Rope, RopeMode, Score}; use super::{KvBytes, KvLayout}; - /// `rope.dimension_count` f32 a position, through [`KvLayout::of`]. The + /// `head_dim` f32 a position, through [`KvLayout::of`]. The /// q8 cache keeps the same table. A count of no row, or of half a row, /// is not the table the body copies. #[test] diff --git a/crates/model/tests/qwen4exp_meta.rs b/crates/model/tests/qwen4exp_meta.rs index 66e2f57d8..fc404f617 100644 --- a/crates/model/tests/qwen4exp_meta.rs +++ b/crates/model/tests/qwen4exp_meta.rs @@ -654,6 +654,7 @@ type CardPlanRow = (&'static str, u64, bool, u64, u64, usize, u64, u64, u64); // +88,080,384 B, the splits 238/12/237, 255/40/254, 68/2/67, 86/5/85 — the count past the estimate // is where the granules each layer's three stacks round to fall, the bytes within two granules of // the arena's move. +// PIN(2026-10-15): [derived] the stage rope table consumes 256 B per position on each card. const CARD_PLANS: [CardPlanRow; 16] = [ ( "A6000", @@ -705,10 +706,10 @@ const CARD_PLANS: [CardPlanRow; 16] = [ false, 4_096, 258, - 8, + 7, 257, - 38_685_388_800, - 478_888_448, + 38_682_316_800, + 475_668_992, ), ( "A6000", @@ -716,10 +717,10 @@ const CARD_PLANS: [CardPlanRow; 16] = [ false, 512, 275, - 37, + 33, 274, - 41_332_224_000, - 526_893_568, + 41_319_936_000, + 530_792_960, ), ( "A6000", @@ -727,10 +728,10 @@ const CARD_PLANS: [CardPlanRow; 16] = [ true, 4_096, 238, - 12, + 8, 237, - 35_689_164_800, - 570_556_928, + 35_676_876_800, + 574_456_320, ), ( "A6000", @@ -738,10 +739,10 @@ const CARD_PLANS: [CardPlanRow; 16] = [ true, 512, 255, - 40, + 39, 254, - 38_332_928_000, - 615_342_592, + 38_329_856_000, + 612_123_136, ), ( "3090", @@ -760,10 +761,10 @@ const CARD_PLANS: [CardPlanRow; 16] = [ false, 512, 110, - 38, + 37, 109, - 16_515_072_000, - 557_806_080, + 16_512_000_000, + 554_586_624, ), ( "3090", @@ -793,10 +794,10 @@ const CARD_PLANS: [CardPlanRow; 16] = [ false, 4_096, 87, - 29, + 27, 86, - 13_027_123_200, - 535_122_432, + 13_020_979_200, + 532_877_824, ), ( "3090", @@ -804,21 +805,21 @@ const CARD_PLANS: [CardPlanRow; 16] = [ false, 512, 105, - 22, + 21, 104, - 15_713_280_000, - 539_611_648, + 15_710_208_000, + 536_392_192, ), ( "3090", 32_768, true, 4_096, - 68, - 2, 67, - 10_084_659_200, - 573_030_912, + 44, + 66, + 10_065_203_200, + 584_098_304, ), ( "3090", @@ -826,10 +827,10 @@ const CARD_PLANS: [CardPlanRow; 16] = [ true, 512, 86, - 5, + 2, 85, - 12_802_969_600, - 547_463_680, + 12_792_320_000, + 547_627_520, ), ]; // PIN(2026-09-28): the routed stacks no card expert kernel of the program reads, which keep their @@ -1091,6 +1092,7 @@ const CARD_ROUNDING_Q3: u64 = 397_191_680; // (`MtpInputs::arena_bytes`, round q3seat); every drafted row here left 400,148 B or more of slack // before it, so no row moves [derived: the q3seat replica reproduced all eight drafted rows' // expert bytes, the least slack the 3090's at ctx 4,096 and U 512]. +// PIN(2026-10-15): [derived] the Q3 file uses the same 256 B per position stage rope table. const CARD_PLANS_Q3: [CardPlanRow; 16] = [ ( "A6000", @@ -1109,10 +1111,10 @@ const CARD_PLANS_Q3: [CardPlanRow; 16] = [ false, 512, 388, - 21, + 20, 387, - 42_242_585_600, - 586_282_496, + 42_240_409_600, + 584_264_192, ), ( "A6000", @@ -1142,10 +1144,10 @@ const CARD_PLANS_Q3: [CardPlanRow; 16] = [ false, 4_096, 355, - 41, + 40, 354, - 38_688_921_600, - 629_314_048, + 38_686_745_600, + 625_198_592, ), ( "A6000", @@ -1153,10 +1155,10 @@ const CARD_PLANS_Q3: [CardPlanRow; 16] = [ false, 512, 381, - 48, - 381, - 41_540_582_400, - 470_396_416, + 45, + 380, + 41_532_416_000, + 472_271_360, ), ( "A6000", @@ -1164,10 +1166,10 @@ const CARD_PLANS_Q3: [CardPlanRow; 16] = [ true, 4_096, 328, - 46, + 44, 327, - 35_755_980_800, - 657_699_328, + 35_751_628_800, + 649_468_416, ), ( "A6000", @@ -1175,10 +1177,10 @@ const CARD_PLANS_Q3: [CardPlanRow; 16] = [ true, 512, 355, - 7, + 6, 354, - 38_614_118_400, - 490_207_744, + 38_611_942_400, + 486_092_288, ), ( "3090", @@ -1219,10 +1221,10 @@ const CARD_PLANS_Q3: [CardPlanRow; 16] = [ true, 512, 128, - 17, + 16, 127, - 13_885_977_600, - 492_926_464, + 13_883_801_600, + 488_811_008, ), ( "3090", @@ -1230,10 +1232,10 @@ const CARD_PLANS_Q3: [CardPlanRow; 16] = [ false, 4_096, 121, - 29, + 28, 120, - 13_148_876_800, - 565_230_080, + 13_146_700_800, + 561_114_624, ), ( "3090", @@ -1241,10 +1243,10 @@ const CARD_PLANS_Q3: [CardPlanRow; 16] = [ false, 512, 146, - 21, + 17, 145, - 15_857_228_800, - 553_815_552, + 15_848_524_800, + 554_130_944, ), ( "3090", @@ -1252,10 +1254,10 @@ const CARD_PLANS_Q3: [CardPlanRow; 16] = [ true, 4_096, 94, - 34, + 33, 93, - 10_216_755_200, - 592_796_160, + 10_214_579_200, + 588_680_704, ), ( "3090", @@ -1263,10 +1265,10 @@ const CARD_PLANS_Q3: [CardPlanRow; 16] = [ true, 512, 119, - 40, + 36, 118, - 12_955_571_200, - 548_820_480, + 12_946_867_200, + 549_135_872, ), ]; From 8f3894a64c78b663dc900d2d61fa07983b290d33 Mon Sep 17 00:00:00 2001 From: midagedev Date: Thu, 15 Oct 2026 11:13:33 +0900 Subject: [PATCH 3/4] model, gpu: the stage-row refusal stays in the crate, and each rope test builds its layout refuse_stage_rows_past is crate-visible. The two stage-rope tests no longer keep a layout helper that one test calls. --- crates/gpu/src/model.rs | 2 +- crates/model/src/arch/qwen35moe/place.rs | 28 +++++++++++------------- crates/model/src/arch/qwen3moe/place.rs | 23 +++++++------------ 3 files changed, 22 insertions(+), 31 deletions(-) diff --git a/crates/gpu/src/model.rs b/crates/gpu/src/model.rs index da2a38517..52cabaebf 100644 --- a/crates/gpu/src/model.rs +++ b/crates/gpu/src/model.rs @@ -1948,7 +1948,7 @@ pub fn refuse_scratch_past( /// (`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 fn refuse_stage_rows_past( +pub(crate) fn refuse_stage_rows_past( what: &'static str, plan: &Plan<'_>, card: usize, diff --git a/crates/model/src/arch/qwen35moe/place.rs b/crates/model/src/arch/qwen35moe/place.rs index 3c951090f..60578db58 100644 --- a/crates/model/src/arch/qwen35moe/place.rs +++ b/crates/model/src/arch/qwen35moe/place.rs @@ -2290,28 +2290,19 @@ mod tests { assert_eq!(selecting.in_q8().layer_bytes(0, 4096), sel); } - /// The stage rope rows, through [`KvLayout::of`]: `rope.dimension_count` - /// f32 a position, not the head width. The q8 cache keeps the same table. - /// A count of no row, or of half a row, is not the table the body copies. #[test] fn the_stage_rope_row_is_rope_dims_not_the_head() { - let narrow = layout(64, 256); - assert_eq!(narrow.stage_position_bytes(8), 64 * 4 * 8); - assert_eq!(narrow.in_q8().stage_position_bytes(8), 64 * 4 * 8); - assert_eq!(layout(128, 256).stage_position_bytes(8), 128 * 4 * 8); - } - - fn layout(rope_dims: usize, head_dim: usize) -> KvLayout { use super::super::hparams::{FfnKind, Hparams, SectionOrder, Variant}; - KvLayout::of(&Hparams { + + let hp = Hparams { variant: Variant::Qwen35Moe, n_layer: 1, n_trunk: 1, - n_embd: head_dim, + n_embd: 256, n_head: 1, n_head_kv: 1, - head_dim, - rope_dims, + head_dim: 256, + rope_dims: 64, rope_sections: [0, 0, 0, 0], rope_order: SectionOrder::Tyxz, rope_base: 1e7, @@ -2332,7 +2323,14 @@ mod tests { 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 diff --git a/crates/model/src/arch/qwen3moe/place.rs b/crates/model/src/arch/qwen3moe/place.rs index 555af8b2e..becb2f5ff 100644 --- a/crates/model/src/arch/qwen3moe/place.rs +++ b/crates/model/src/arch/qwen3moe/place.rs @@ -194,31 +194,22 @@ mod tests { use super::super::hparams::{Experts, Hparams, Rope, RopeMode, Score}; use super::{KvBytes, KvLayout}; - /// `head_dim` f32 a position, through [`KvLayout::of`]. The - /// q8 cache keeps the same table. A count of no row, or of half a row, - /// is not the table the body copies. #[test] fn the_stage_rope_row_is_the_width_the_body_copies() { - let kv = layout(128); - assert_eq!(kv.stage_position_bytes(10), 128 * 4 * 10); - assert_eq!(kv.in_q8().stage_position_bytes(10), 128 * 4 * 10); - } - - fn layout(head_dim: usize) -> KvLayout { - KvLayout::of(&Hparams { + let kv = KvLayout::of(&Hparams { n_layer: 1, - n_embd: head_dim, + n_embd: 128, n_head: 1, n_head_kv: 1, - head_dim, + head_dim: 128, rope: Rope { mode: RopeMode::Neox, base: 1e6, - dims: head_dim, + dims: 128, }, rms_eps: 1e-6, n_vocab: 2, - n_ctx_train: head_dim, + n_ctx_train: 128, experts: Experts { n_expert: 1, n_used: 1, @@ -226,6 +217,8 @@ mod tests { 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); } } From cb78660141e2f36cf20fbe5cefd1b15ca2d14881 Mon Sep 17 00:00:00 2001 From: midagedev Date: Thu, 15 Oct 2026 12:55:49 +0900 Subject: [PATCH 4/4] model: format the inlined stage-row check contrib-check t0 stopped on the one-line Hparams update in the qwen3.5 fixture. --- crates/model/src/arch/qwen35moe/place.rs | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/crates/model/src/arch/qwen35moe/place.rs b/crates/model/src/arch/qwen35moe/place.rs index 60578db58..b6c707890 100644 --- a/crates/model/src/arch/qwen35moe/place.rs +++ b/crates/model/src/arch/qwen35moe/place.rs @@ -2328,7 +2328,11 @@ mod tests { 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), + KvLayout::of(&Hparams { + rope_dims: 128, + ..hp + }) + .stage_position_bytes(8), 128 * 4 * 8 ); }