diff --git a/crates/gpu-deepseek41/src/block.rs b/crates/gpu-deepseek41/src/block.rs index d3a8d33f4..eb24fbc34 100644 --- a/crates/gpu-deepseek41/src/block.rs +++ b/crates/gpu-deepseek41/src/block.rs @@ -227,6 +227,12 @@ pub struct KBlockKernels { } impl KBlockKernels { + /// The device bytes the kernels hold beside their modules: the K-gemv's + /// expert ids. + pub fn device_bytes(&self) -> usize { + self.gemv.device_bytes() + } + /// Load the three into `gpu`'s context. Load-time only. /// /// # Errors diff --git a/crates/gpu-gates/src/bin/gate_nemotron_serve.rs b/crates/gpu-gates/src/bin/gate_nemotron_serve.rs index 821699c52..d8d4c6633 100644 --- a/crates/gpu-gates/src/bin/gate_nemotron_serve.rs +++ b/crates/gpu-gates/src/bin/gate_nemotron_serve.rs @@ -163,8 +163,8 @@ mod gate { fn plan_only_names_its_default( dir: &Path, file: &Path, - ok: &mut bool, - ) -> Result, GateError> { + ) -> Result<(Option, bool), GateError> { + let mut ok = true; let limit = trained_ctx(file)?; let (status, log, text) = plan(dir, "plan", file, &[])?; let place = log.first(&record::PLACE_UNSET)?; @@ -233,19 +233,16 @@ mod gate { default_ok = line_ok && held && stands && above_short; } check( - ok, + &mut ok, "plan_only_names_its_default", exits && in_order && word_a && no_load && default_ok, ); - Ok(ctx_line(&text).and_then(|l| field(l, "ctx"))) + Ok((ctx_line(&text).and_then(|l| field(l, "ctx")), ok)) } /// `parallel_above_one_is_refused` (the module header). - fn parallel_above_one_is_refused( - dir: &Path, - file: &Path, - ok: &mut bool, - ) -> Result<(), GateError> { + fn parallel_above_one_is_refused(dir: &Path, file: &Path) -> Result { + let mut ok = true; let (status, log, text) = plan(dir, "refuse-parallel", file, &["--parallel", "2"])?; let named = text.contains(ONE_SEQUENCE); let no_plan = log.first(&record::PLAN)?.is_none(); @@ -258,11 +255,11 @@ mod gate { println!("refuse-parallel stderr:\n{text}"); } check( - ok, + &mut ok, "parallel_above_one_is_refused", !status.success() && named && no_plan, ); - Ok(()) + Ok(ok) } /// The ids `NemotronHModel` gives stepped directly: for each prompt, from an empty cache, the @@ -306,8 +303,8 @@ mod gate { file: &Path, default_ctx: Option, layers: u64, - ok: &mut bool, - ) -> Result, GateError> { + ) -> Result<(Option, bool), GateError> { + let mut ok = true; let d = dir.join("serve"); std::fs::create_dir_all(&d)?; let (cmd, mut all) = command(file)?; @@ -319,8 +316,8 @@ mod gate { Ok(addr) => addr, Err(e) => { println!("the server never listened: {e}"); - check(ok, "load_and_listen", false); - return Ok(None); + check(&mut ok, "load_and_listen", false); + return Ok((None, ok)); } }; let text = std::fs::read_to_string(&err_log).unwrap_or_default(); @@ -345,7 +342,7 @@ mod gate { listen.line() ); check( - ok, + &mut ok, "load_and_listen", same_ctx && loaded && ctx_is_plans && in_order, ); @@ -373,15 +370,18 @@ mod gate { println!("{name}: {e}"); false }); - check(ok, name, pass); + check(&mut ok, name, pass); } println!("server stopped: {}", s.stop()?); - Ok(Some(Answers { - prose, - turn, - prose_answer, - turn_answer, - })) + Ok(( + Some(Answers { + prose, + turn, + prose_answer, + turn_answer, + }), + ok, + )) } /// `edit_resend_is_a_fresh_run` (the module header): the slot holds the prose prompt and its @@ -428,7 +428,8 @@ mod gate { } /// `completion_ids` (the module header), the server stopped and the session on the card. - fn completion_ids(file: &Path, a: &Answers, ok: &mut bool) -> Result<(), GateError> { + fn completion_ids(file: &Path, a: &Answers) -> Result { + let mut ok = true; let ids = session_ids(file, &[a.prose.as_slice(), a.turn.as_slice()])?; let prose = agree(&a.prose_answer.tokens, &a.prose_answer.stop, &ids[0]); let turn = agree(&a.turn_answer.tokens, &a.turn_answer.stop, &ids[1]); @@ -438,8 +439,8 @@ mod gate { a.prose_answer.tokens.len(), a.turn_answer.tokens.len() ); - check(ok, "completion_ids", prose && turn); - Ok(()) + check(&mut ok, "completion_ids", prose && turn); + Ok(ok) } pub fn run() -> Result<(), GateError> { @@ -454,11 +455,13 @@ mod gate { )?; drop(split); let mut ok = true; - let default_ctx = plan_only_names_its_default(&dir, &file, &mut ok)?; - parallel_above_one_is_refused(&dir, &file, &mut ok)?; - let answers = served(&dir, &file, default_ctx, layers, &mut ok)?; + let (default_ctx, pass) = plan_only_names_its_default(&dir, &file)?; + ok &= pass; + ok &= parallel_above_one_is_refused(&dir, &file)?; + let (answers, pass) = served(&dir, &file, default_ctx, layers)?; + ok &= pass; if let Some(a) = answers { - completion_ids(&file, &a, &mut ok)?; + ok &= completion_ids(&file, &a)?; } if !ok { return Err(checks_failed()); diff --git a/crates/gpu-glm5next/src/body.rs b/crates/gpu-glm5next/src/body.rs index 930050d78..d72c8d7d4 100644 --- a/crates/gpu-glm5next/src/body.rs +++ b/crates/gpu-glm5next/src/body.rs @@ -2376,6 +2376,7 @@ impl ChainBody for Body { .nextn .as_deref() .map_or(0, |n| n.resident_bytes() + n.arena_bytes()) + + self.k.block.as_ref().map_or(0, KBlockKernels::device_bytes) } fn layers(&self) -> Range { diff --git a/crates/gpu-mimo2/src/body.rs b/crates/gpu-mimo2/src/body.rs index e930ca76d..37018bcc0 100644 --- a/crates/gpu-mimo2/src/body.rs +++ b/crates/gpu-mimo2/src/body.rs @@ -1215,6 +1215,7 @@ impl ChainBody for Body { .taps .as_ref() .map_or(0, |t| t.iter().map(DeviceBuffer::num_bytes).sum()) + + self.k.block.as_ref().map_or(0, KBlockKernels::device_bytes) } fn layers(&self) -> Range { diff --git a/crates/gpu/src/arch/nemotron_h/body.rs b/crates/gpu/src/arch/nemotron_h/body.rs index a31ffe2f4..d0cd3b0cc 100644 --- a/crates/gpu/src/arch/nemotron_h/body.rs +++ b/crates/gpu/src/arch/nemotron_h/body.rs @@ -74,9 +74,6 @@ impl GpuModel { impl Body { fn load(gpu: &Gpu, file: &Split, w: &mut Weights, ctx: usize) -> Result { let hp = Hparams::read(file).map_err(|e| GpuError::plan(WHAT, e))?; - if ctx == 0 { - return Err(GpuError::shape(WHAT, "a cache of 0 rows")); - } let ssm = hp.ssm; let mamba = mamba2::check( WHAT, @@ -359,7 +356,11 @@ impl ChainBody for Body { Layer::Mlp(_) => 0, }) .sum(); - layers + self.s.bytes() + self.acts.bytes() + self.r.num_bytes() + layers + + self.s.bytes() + + self.acts.bytes() + + self.r.num_bytes() + + self.k.kgemv.device_bytes() } fn layers(&self) -> Range { diff --git a/crates/gpu/src/arch/nemotron_h/step.rs b/crates/gpu/src/arch/nemotron_h/step.rs index b95fe12a0..93f50cbc8 100644 --- a/crates/gpu/src/arch/nemotron_h/step.rs +++ b/crates/gpu/src/arch/nemotron_h/step.rs @@ -191,6 +191,7 @@ impl Scratch { + self.words.num_bytes() + self.pos.num_bytes() + self.n_keys.num_bytes() + + self.mamba.device_bytes() } /// The scratch for `sh`, the widest `ff` an MLP layer has and the diff --git a/crates/gpu/src/arch/qwen3moe/body.rs b/crates/gpu/src/arch/qwen3moe/body.rs index 54356daa8..95da8ec6b 100644 --- a/crates/gpu/src/arch/qwen3moe/body.rs +++ b/crates/gpu/src/arch/qwen3moe/body.rs @@ -1022,6 +1022,7 @@ impl ChainBody for Body { + self.ub.bytes() + self.taps.as_ref().map_or(0, |t| t.buf.num_bytes()) + self.placed.as_ref().map_or(0, Placed::device_bytes) + + self.k.q35.as_ref().map_or(0, |k| k.kgemv.device_bytes()) } fn layers(&self) -> Range { diff --git a/crates/gpu/src/arch/qwen3moe/body35.rs b/crates/gpu/src/arch/qwen3moe/body35.rs index daf5178ea..d52fffedf 100644 --- a/crates/gpu/src/arch/qwen3moe/body35.rs +++ b/crates/gpu/src/arch/qwen3moe/body35.rs @@ -3193,6 +3193,7 @@ impl ChainBody for Body35 { .as_ref() .map_or(0, |d| d.resident_bytes() + d.scratch_bytes()) + self.placed.as_ref().map_or(0, Placed::device_bytes) + + self.k.q35.as_ref().map_or(0, |k| k.kgemv.device_bytes()) } fn layers(&self) -> Range { diff --git a/crates/gpu/src/mamba2.rs b/crates/gpu/src/mamba2.rs index f2b7f9ba7..d4d74146b 100644 --- a/crates/gpu/src/mamba2.rs +++ b/crates/gpu/src/mamba2.rs @@ -473,6 +473,11 @@ pub struct Mamba2Scratch { } impl Mamba2Scratch { + /// The device bytes the scratch holds. + pub(crate) fn device_bytes(&self) -> usize { + self.part.num_bytes() + self.done.num_bytes() + } + /// A zeroed scratch for `shape`. Load-time only. `shape` is checked again /// here, so a conv the host accepted and this kernel refuses never /// allocates. diff --git a/crates/gpu/src/site.rs b/crates/gpu/src/site.rs index c8511b979..482caa25c 100644 --- a/crates/gpu/src/site.rs +++ b/crates/gpu/src/site.rs @@ -336,6 +336,12 @@ pub struct KGemvKernels { } impl KGemvKernels { + /// The device bytes the kernels hold beside their modules: the one + /// expert's ids. + pub fn device_bytes(&self) -> usize { + self.one.num_bytes() + } + /// Load the `_sel` family into `gpu`'s context, raising into its fault /// word, and zero the ids. Load-time only. pub fn load(gpu: &Gpu) -> Result {