From 123c7a749bc7de43bae1cc144f4a688056ffd0a7 Mon Sep 17 00:00:00 2001 From: juliankang4 <200967947+juliankang4@users.noreply.github.com> Date: Sun, 25 Oct 2026 01:15:25 +0900 Subject: [PATCH 1/8] os: the huge-page advice of a new allocation, moved from gpu host_filled's advice (madvise(MADV_HUGEPAGE) over the whole huge pages of a new allocation of 32 MiB or more) moves to bloomery-os as advise_huge_within, so the loader's own new buffers can take it too. The gpu block moves with it; the unsafe pins follow. --- crates/gpu/src/tensor.rs | 41 +++++----------------------------------- crates/os/src/lib.rs | 33 ++++++++++++++++++++++++++++++++ 2 files changed, 38 insertions(+), 36 deletions(-) diff --git a/crates/gpu/src/tensor.rs b/crates/gpu/src/tensor.rs index cdd5ea659..4ec159098 100644 --- a/crates/gpu/src/tensor.rs +++ b/crates/gpu/src/tensor.rs @@ -7,7 +7,7 @@ use crate::GpuError; use crate::graph::cu; use cuda_core::{CudaContext, CudaStream, DeviceBuffer, DeviceCopy, sys}; use std::marker::PhantomData; -use std::mem::{ManuallyDrop, MaybeUninit, align_of, size_of, size_of_val}; +use std::mem::{ManuallyDrop, align_of, size_of}; use std::ops::{Deref, DerefMut}; use std::sync::Arc; @@ -337,47 +337,16 @@ impl HostFill for u16 { /// reaches pageable memory through the host's page tables (GB10's unified /// memory) faults on each page a copy is the first to write, one page at a /// time; a new allocation nothing has written is such pages. A large one is -/// advised onto huge pages first ([`advise_huge`]). `h2d_probe --arm readback` +/// advised onto huge pages first ([`os::advise_huge_within`]). `h2d_probe --arm readback` /// (`bloomery-gpu-gates`) times a readback into each kind of destination. #[must_use] pub fn host_filled(n: usize) -> Vec { let mut host = Vec::with_capacity(n); - advise_huge(host.spare_capacity_mut()); + os::advise_huge_within(host.spare_capacity_mut()); host.resize(n, T::FILL); host } -/// The huge page the advice asks for. -const HUGE_PAGE: usize = 2 << 20; - -/// The size from which a new allocation is memory nothing has written: -/// below it glibc's allocator mostly hands back heap memory it has written -/// before, which the advice does not change. -const ADVISE_FROM: usize = 32 << 20; - -/// Huge pages for the whole huge pages inside `spare`, a new allocation of -/// at least [`ADVISE_FROM`] bytes, so that its fill faults once each huge -/// page instead of once each base page. -fn advise_huge(spare: &mut [MaybeUninit]) { - let bytes = size_of_val(spare); - if bytes < ADVISE_FROM { - return; - } - let at = spare.as_mut_ptr() as usize; - let (first, end) = ( - at.next_multiple_of(HUGE_PAGE), - (at + bytes) / HUGE_PAGE * HUGE_PAGE, - ); - if end > first { - // SAFETY: `first..end` is whole pages inside the allocation `spare` - // borrows, private anonymous memory its vector owns. - let advised = unsafe { os::advise_huge(first as *mut libc::c_void, end - first) }; - // A refusal (a kernel without transparent huge pages) leaves base - // pages, which the fill writes the same. - drop(advised); - } -} - /// `src` read back into a new [`host_filled`] `Vec`: a readback the caller /// keeps. Blocking. pub fn to_host( @@ -1050,8 +1019,8 @@ mod tests { fn host_filled_memory_holds_the_fill_not_zero() { assert!(host_filled::(5).iter().all(|v| v.is_nan())); assert_eq!(host_filled::(5), [u16::MAX; 5]); - let advised = host_filled::(ADVISE_FROM / 2 + 3); - assert_eq!(advised.len(), ADVISE_FROM / 2 + 3); + let advised = host_filled::(os::ADVISE_FROM / 2 + 3); + assert_eq!(advised.len(), os::ADVISE_FROM / 2 + 3); assert!(advised.iter().all(|&v| v == u16::MAX), "the advised fill"); } diff --git a/crates/os/src/lib.rs b/crates/os/src/lib.rs index 3608e5d3f..856a7407a 100644 --- a/crates/os/src/lib.rs +++ b/crates/os/src/lib.rs @@ -809,6 +809,39 @@ pub unsafe fn advise_huge(_addr: *mut libc::c_void, _len: usize) -> Result<(), O }) } +/// The huge page [`advise_huge_within`] asks for. +const HUGE_PAGE: usize = 2 << 20; + +/// The size from which a new allocation is memory nothing has written: +/// below it glibc's allocator mostly hands back heap memory it has written +/// before, which the advice does not change. +pub const ADVISE_FROM: usize = 32 << 20; + +/// Huge pages for the whole huge pages inside `span`, a new heap allocation +/// of at least [`ADVISE_FROM`] bytes nothing has written yet (a vector's +/// spare capacity, or a new zeroed vector), so that its first writes fault +/// once each huge page instead of once each base page, and its free unmaps +/// huge pages. A refusal (a kernel without transparent huge pages, or a +/// platform without the advice) leaves base pages, which the writes fill +/// the same. +pub fn advise_huge_within(span: &mut [T]) { + let bytes = size_of_val(span); + if bytes < ADVISE_FROM { + return; + } + let at = span.as_mut_ptr().addr(); + let (first, end) = ( + at.next_multiple_of(HUGE_PAGE), + (at + bytes) / HUGE_PAGE * HUGE_PAGE, + ); + if end > first { + // SAFETY: `first..end` is whole pages inside the allocation `span` + // borrows, private anonymous memory its owner holds. + let advised = unsafe { advise_huge(first as *mut libc::c_void, end - first) }; + drop(advised); + } +} + /// Fault the pages of `addr .. addr + len` in. /// /// `advise` is the platform's populate advice call. A refusal that names the From ab5f792620c8c41402b00c50d87e535f403879f3 Mon Sep 17 00:00:00 2001 From: juliankang4 <200967947+juliankang4@users.noreply.github.com> Date: Sun, 25 Oct 2026 01:15:26 +0900 Subject: [PATCH 2/8] model-hf: read a whole stored tensor past the page cache A checkpoint folder's read normalized each tensor from the shard's mapping, so every 4 KiB page of the 78 GB Qwen3.8 FP4-Local checkpoint was a fault on the reading thread, and the pages then sat in the page cache until the load dropped them again. Checkpoint::read now reads each stored tensor of a whole tensor with one direct read through engram's DirectFile into a buffer of its own, and normalizes from those bytes; a part of a tensor (a tensor-parallel rank's rows) still reads its rows through the mapping, and so does a shard whose mount takes no direct reads (DirectFile::open refuses it). The new buffers, and the normalized NVFP4 and FP8 blocks, are advised onto huge pages (os::advise_huge_within), so filling them faults once each 2 MiB. --- Cargo.lock | 2 + crates/hf/Cargo.toml | 2 + crates/hf/src/modelopt.rs | 2 + crates/model-hf/Cargo.toml | 2 + crates/model-hf/src/part.rs | 85 +++++++++++++++++++++++++++++++++++-- 5 files changed, 89 insertions(+), 4 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index e95e1197b..1b960d92a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -341,6 +341,7 @@ name = "bloomery-hf" version = "0.1.0" dependencies = [ "bloomery-gguf", + "bloomery-os", "bloomery-runtime", "bloomery-safetensors", "bloomery-tokenizer", @@ -407,6 +408,7 @@ dependencies = [ name = "bloomery-model-hf" version = "0.1.0" dependencies = [ + "bloomery-engram", "bloomery-gguf", "bloomery-hf", "bloomery-levers", diff --git a/crates/hf/Cargo.toml b/crates/hf/Cargo.toml index 5a19f5067..ab0059bab 100644 --- a/crates/hf/Cargo.toml +++ b/crates/hf/Cargo.toml @@ -14,6 +14,8 @@ workspace = true [dependencies] bloomery-gguf = { path = "../gguf" } +# A normalized tensor's new memory is advised onto huge pages. +bloomery-os.workspace = true # `dspark` returns the draft's geometry as the runtime's `DsparkShape`; the runtime depends on the model # description only, so hf does not close a cycle. bloomery-runtime = { path = "../runtime" } diff --git a/crates/hf/src/modelopt.rs b/crates/hf/src/modelopt.rs index 9d0434f70..e4a16b4c9 100644 --- a/crates/hf/src/modelopt.rs +++ b/crates/hf/src/modelopt.rs @@ -1039,6 +1039,7 @@ fn fp8_append<'a>( .ok_or_else(|| refuse("block bytes size overflow".into()))?; out.try_reserve(more) .map_err(|e| refuse(format!("{more} block bytes: {e}")))?; + os::advise_huge_within(out.spare_capacity_mut()); for (row, values) in w.bytes.chunks_exact(k).enumerate() { let at = out.len(); for (j, block) in values.chunks_exact(blck).enumerate() { @@ -1311,6 +1312,7 @@ fn append<'a>( out.blocks .try_reserve(bytes) .map_err(|e| shape(&wname, format!("cannot reserve {bytes} NVFP4 bytes: {e}")))?; + os::advise_huge_within(out.blocks.spare_capacity_mut()); for (row, (w, s)) in w .bytes .chunks_exact(half) diff --git a/crates/model-hf/Cargo.toml b/crates/model-hf/Cargo.toml index e6bad704d..710ce6884 100644 --- a/crates/model-hf/Cargo.toml +++ b/crates/model-hf/Cargo.toml @@ -19,6 +19,8 @@ bloomery-hf = { path = "../hf" } bloomery-levers = { path = "../levers" } bloomery-model = { path = "../model" } bloomery-safetensors = { path = "../safetensors" } +# A whole stored tensor is read past the page cache through the one direct reader. +bloomery-engram = { path = "../engram" } # The `prepare` bin builds serving's tokenizer from the GGUF header it wrote and checks its split. bloomery-tokenizer = { path = "../tokenizer" } # The `fixture` bin and its test read peak resident sets in bytes on every platform. diff --git a/crates/model-hf/src/part.rs b/crates/model-hf/src/part.rs index cfd87b073..2322a6b8c 100644 --- a/crates/model-hf/src/part.rs +++ b/crates/model-hf/src/part.rs @@ -9,15 +9,18 @@ //! experts or by a list, are the experts' own sources normalized in the part's order. A tensor //! the model moves into its layout after normalizing is made whole, moved, and then cut. -use std::collections::BTreeMap; +use std::collections::{BTreeMap, HashMap}; +use std::ops::Range; +use engram::EngramError; +use engram::direct::{DIRECT_ALIGN, DirectFile, aligned_span}; use gguf::{GgmlType, TensorInfo}; use hf::modelopt::{self, Normal, Use}; use hf::quant_config::{Placed, Plan}; use model::part::{self, At, Part, Planned, Source, Take}; use model::placement::ExpertList; -use model::placement::tp::TpError; -use safetensors::Folder; +use model::placement::tp::{TpError, part_of}; +use safetensors::{Dtype, Folder, Raw}; /// Why a part was refused. #[derive(Debug, thiserror::Error)] @@ -34,6 +37,8 @@ pub enum Error { Release(#[from] model::placement::PlacementError), #[error(transparent)] Tensors(#[from] safetensors::Error), + #[error(transparent)] + Direct(#[from] engram::EngramError), } impl From for Error { @@ -220,9 +225,22 @@ impl Source for Checkpoint { .made .get(name) .ok_or_else(|| refuse(name, "is not one the map makes"))?; - let get = |n: &str| self.folder.raw(n); let placed = self.placed(m); let ty = self.card_type(name, m)?; + // A whole tensor's stored bytes are read into memory of the read's own ([`read_whole`]); + // a part reads only its rows through the mapping. + let mut read = HashMap::new(); + if p.at.is_none() || m.step.is_some() { + for n in stored(&placed) { + if let Some(w) = read_whole(&self.folder, n)? { + read.insert(n, w); + } + } + } + let get = |n: &str| match read.get(n) { + Some(w) => Ok(w.raw()), + None => self.folder.raw(n), + }; let n = match (&p.at, &m.step) { (Some(At::Along(dim, r)), None) => { modelopt::normalize_part(get, &placed, m.role, &m.dims, (*dim, r.clone()))? @@ -306,6 +324,65 @@ impl Source for Checkpoint { } } +/// A stored tensor read whole into memory of its own ([`read_whole`]). +struct Whole { + dtype: Dtype, + shape: Vec, + buf: Vec, + /// The tensor's bytes in `buf`. + at: Range, +} + +impl Whole { + fn raw(&self) -> Raw<'_> { + Raw { + dtype: self.dtype, + shape: &self.shape, + bytes: &self.buf[self.at.clone()], + } + } +} + +/// Stored tensor `name` read whole from its shard past the page cache (engram's direct reader) +/// into new memory advised onto huge pages, so the read takes the drive's rate, faults once each +/// 2 MiB and leaves the page cache as it was. `None` where the shard cannot be opened for direct +/// reads or its mount refuses them ([`DirectFile::open`]): the caller reads the tensor through +/// the mapping. +fn read_whole(folder: &Folder, name: &str) -> Result, Error> { + let (path, range) = folder.file_range(name)?; + let file = match DirectFile::open(&path) { + Ok(file) => file, + Err(EngramError::OpenDirect { .. } | EngramError::DirectProbe { .. }) => return Ok(None), + Err(e) => return Err(e.into()), + }; + let raw = folder.raw(name)?; + let (start, span, need) = aligned_span(range.start, range.end - range.start); + let mut buf = vec![0u8; span + DIRECT_ALIGN]; + os::advise_huge_within(&mut buf); + let lead = buf.as_ptr().align_offset(DIRECT_ALIGN); + if file.read_span(&mut buf[lead..lead + span], start)? < need { + return Err(refuse( + name, + format!("{} ends inside its bytes", path.display()), + )); + } + let skip = need - raw.bytes.len(); + Ok(Some(Whole { + dtype: raw.dtype, + shape: raw.shape.to_vec(), + buf, + at: lead + skip..lead + need, + })) +} + +/// The stored tensors a tensor placed as `placed` is made of: each source, then the companions +/// its place names (scales, the packed shape). +fn stored<'s>(placed: &[(&'s str, &'s Placed)]) -> impl Iterator { + placed.iter().flat_map(|&(source, placed)| { + std::iter::once(source).chain(placed.companions.iter().map(String::as_str)) + }) +} + /// The whole tensor `name` of `ty` and `dims` as a header entry; dims that are not whole blocks /// of the type, or bytes past u64, are refused by name. fn whole(name: &str, ty: GgmlType, dims: &[u64]) -> Result { From d3157048fcc0490ff6d26777bd11d334a2181bac Mon Sep 17 00:00:00 2001 From: juliankang4 <200967947+juliankang4@users.noreply.github.com> Date: Sun, 25 Oct 2026 01:15:26 +0900 Subject: [PATCH 3/8] hf: pack NVFP4 groups a word at a time append moved each NVFP4 group's 16 nibbles into GGUF's order one byte at a time through a per-byte iterator, and decoded every scale byte with powi. The group is now two 32-bit words whose nibbles are spread with three shift-and-mask steps, and the scale bytes' half values come from a table built once a tensor. A byte that is no UE4M3 scale is refused as before, naming the same row and group. The modelopt tests' reference values hold the bytes. --- crates/gguf/src/quant.rs | 13 +++++++++---- crates/hf/src/modelopt.rs | 31 ++++++++++++++++++++++--------- 2 files changed, 31 insertions(+), 13 deletions(-) diff --git a/crates/gguf/src/quant.rs b/crates/gguf/src/quant.rs index d7e9f9bc8..24b4f86c8 100644 --- a/crates/gguf/src/quant.rs +++ b/crates/gguf/src/quant.rs @@ -907,11 +907,16 @@ pub fn fp4_pair(byte: u8) -> [i8; 2] { /// plane runs it, and qdot's AVX2 body is held to it bit for bit. #[inline] #[must_use] +#[inline] pub fn nvfp4_group_to_block(g: &[u8; 8]) -> [u8; 8] { - std::array::from_fn(|j| { - let nib = |x: u8| (x >> (4 * (j % 2))) & 0x0f; - nib(g[j / 2]) | (nib(g[4 + j / 2]) << 4) - }) + let spread = |half: [u8; 4]| { + let x = u64::from(u32::from_le_bytes(half)); + let x = (x | x << 16) & 0x0000_ffff_0000_ffff; + let x = (x | x << 8) & 0x00ff_00ff_00ff_00ff; + (x | x << 4) & 0x0f0f_0f0f_0f0f_0f0f + }; + let [a, b, c, d, e, f, h, i] = *g; + (spread([a, b, c, d]) | spread([e, f, h, i]) << 4).to_le_bytes() } /// The block scale of an MXFP4 block, `ggml_e8m0_to_fp32_half` (ggml-impl.h:40-45): diff --git a/crates/hf/src/modelopt.rs b/crates/hf/src/modelopt.rs index e4a16b4c9..718882df2 100644 --- a/crates/hf/src/modelopt.rs +++ b/crates/hf/src/modelopt.rs @@ -1313,6 +1313,8 @@ fn append<'a>( .try_reserve(bytes) .map_err(|e| shape(&wname, format!("cannot reserve {bytes} NVFP4 bytes: {e}")))?; os::advise_huge_within(out.blocks.spare_capacity_mut()); + // Each scale byte's half value, worked out once; a byte that is no scale is refused below. + let halves: [Option; 256] = std::array::from_fn(|b| ue4m3_to_f32_half(b as u8).ok()); for (row, (w, s)) in w .bytes .chunks_exact(half) @@ -1327,12 +1329,14 @@ fn append<'a>( .enumerate() { for (sub, &byte) in d.iter().enumerate() { - let half_scale = ue4m3_to_f32_half(byte).map_err(|source| Error::Scale { - name: sname.clone(), - row, - group: 4 * block + sub, - source, - })?; + let half_scale = halves[usize::from(byte)] + .map_or_else(|| ue4m3_to_f32_half(byte), Ok) + .map_err(|source| Error::Scale { + name: sname.clone(), + row, + group: 4 * block + sub, + source, + })?; if !(half_scale * 12.0 * scale).is_finite() && w[sub * 8..][..8] .iter() @@ -1348,16 +1352,25 @@ fn append<'a>( )); } } - out.blocks.extend_from_slice(d); - for g in w.as_chunks::<8>().0 { - out.blocks.extend(nvfp4_group_to_block(g)); + let mut blocks = [0u8; 36]; + blocks[..4].copy_from_slice(d); + for (g, o) in w + .as_chunks::<8>() + .0 + .iter() + .zip(blocks[4..].as_chunks_mut::<8>().0) + { + *o = nvfp4_group_to_block(g); } + out.blocks.extend_from_slice(&blocks); } } out.scale.push(scale); Ok([rows, k]) } + + /// The stored tensor `name` of dtype `want` whose bytes fill its shape; bytes that do not are /// `unfilled`'s refusal of the tensor, with what is wrong. fn stored<'a>( From 86aa6cd83d3e957b5186cd5710771f183214db55 Mon Sep 17 00:00:00 2001 From: juliankang4 <200967947+juliankang4@users.noreply.github.com> Date: Sun, 25 Oct 2026 01:15:26 +0900 Subject: [PATCH 4/8] placement: the expert rule's footprint keeps what no count changes spread adds one expert at a time and asks the card's footprint after each one, and each ask replayed every upload's buffers through the allocator model. Qwen3.8 has 48 x 512 routed experts and about 3000 tensors, so one plan took about 3 s on GB10, and the seat plans five times before the load (the placement choice, the context rule, the host need and the final plan). Footprint keeps each step's buffers at its last rows, and what its small buffers did to the shared granules. A buffer of a granule or more takes whole granules wherever it falls in the order, so only the small buffers depend on the granules a step finds, and a step that finds them as it last did, with the same small buffers, leaves them as it last did. The footprint is the same bytes as the replay at any counts; the new test asks it at counts that grow, shrink and cross a granule. --- crates/placement/src/placement.rs | 286 +++++++++++++++++++++++++++--- 1 file changed, 262 insertions(+), 24 deletions(-) diff --git a/crates/placement/src/placement.rs b/crates/placement/src/placement.rs index db5f8416a..73fd6793f 100644 --- a/crates/placement/src/placement.rs +++ b/crates/placement/src/placement.rs @@ -28,6 +28,7 @@ //! expert at a time on the layer that holds the fewest so far, so the //! per-layer totals stay even ([`Plan::tier_n_l`]). +use std::cell::RefCell; use std::fmt; use std::num::NonZeroU64; use std::ops::Range; @@ -2282,29 +2283,171 @@ fn card_uploads<'m>( Ok(out) } -/// The bytes of the granules `uploads` take with `n_l` experts on each layer: -/// every upload's buffers, in order, through the card's allocator ([`Heap`]). -fn footprint( +/// The bytes of the granules a card's uploads take with `n_l` experts on each layer: every +/// upload's buffers, in order, through the card's allocator ([`Heap`]). The expert rule asks it +/// once an expert ([`spread`]), and each ask changes one layer's count, so each step of the +/// uploads keeps what it worked out last: its buffers at its last rows, and what its small +/// buffers did to the shared granules from the granules it found. A buffer of a granule or more +/// takes whole granules of its own wherever it falls in the order, so only the small buffers +/// depend on the granules a step finds; a step that finds them as it last found them, with the +/// same small buffers, leaves them as it last left them. +struct Footprint<'m> { granule: NonZeroU64, - uploads: &[(&ModelTensor, CardFormat, Held)], - n_l: &[u64], -) -> Result { - let mut heap = Heap::new(granule); - for &(t, format, held) in uploads { - let rows = match held { - Held::Rows(rows) => rows, - Held::Experts { - layer, - rows_per_expert, - } => n_l[layer] * rows_per_expert, - }; - if rows > 0 { - for b in card_buffers(t, format, rows)? { - heap.alloc(b); + steps: Vec>, +} + +struct Step<'m> { + what: What<'m>, + last: RefCell, +} + +/// One step of [`Footprint`]'s uploads, in order. +enum What<'m> { + /// Fixed uploads' buffers. + Fixed, + /// An upload whose buffers follow its layer's count: a routed stack, or one whose buffers + /// are refused, which each count refuses in its place. + Upload(&'m ModelTensor, CardFormat, Held), +} + +/// What a step worked out last. +#[derive(Default)] +struct Last { + /// The rows its buffers are of; `None` before the first. + rows: Option, + /// The bytes of its buffers of a granule or more, rounded to whole granules. + large: u64, + /// Its smaller buffers, in order. + small: Vec, + /// The shared granules' bytes left as its small buffers found them, and what the buffers + /// did there; `None` until they are allocated. + effect: Option<(Vec, Effect)>, +} + +/// What a run of small buffers did to the shared granules: the bytes left it set by index, the +/// granules it opened (their bytes left), and the bytes of those granules. +struct Effect { + set: Vec<(usize, u64)>, + opened: Vec, + took: u64, +} + +impl Last { + fn of(granule: u64, rows: u64, buffers: Vec) -> Last { + let (large, small): (Vec, Vec) = buffers.into_iter().partition(|&b| b >= granule); + Last { + rows: Some(rows), + large: large.iter().map(|b| b.div_ceil(granule) * granule).sum(), + small, + effect: None, + } + } +} + +impl<'m> Footprint<'m> { + fn new(granule: NonZeroU64, uploads: &[(&'m ModelTensor, CardFormat, Held)]) -> Footprint<'m> { + let mut steps: Vec> = Vec::new(); + for &(t, format, held) in uploads { + let fixed = match held { + Held::Rows(0) => continue, + Held::Rows(rows) => card_buffers(t, format, rows).ok(), + Held::Experts { .. } => None, + }; + let Some(buffers) = fixed else { + steps.push(Step { + what: What::Upload(t, format, held), + last: RefCell::default(), + }); + continue; + }; + let add = Last::of(granule.get(), 0, buffers); + if let Some(Step { + what: What::Fixed, + last, + }) = steps.last_mut() + { + let last = last.get_mut(); + last.large += add.large; + last.small.extend(add.small); + } else { + steps.push(Step { + what: What::Fixed, + last: RefCell::new(add), + }); + } + } + Footprint { granule, steps } + } + + fn at(&self, n_l: &[u64]) -> Result { + let taken = self.walk(n_l); + if taken.is_err() { + // A walk cut short leaves later steps' granules from an older ask. + for step in &self.steps { + *step.last.borrow_mut() = Last::default(); + } + } + taken + } + + fn walk(&self, n_l: &[u64]) -> Result { + let granule = self.granule.get(); + let mut heap = Heap::new(self.granule); + // Whether the shared granules are as the last ask found them at this step. + let mut same = true; + for step in &self.steps { + let mut last = step.last.borrow_mut(); + if let What::Upload(t, format, held) = step.what { + let rows = match held { + Held::Rows(rows) => rows, + Held::Experts { + layer, + rows_per_expert, + } => n_l[layer] * rows_per_expert, + }; + if last.rows != Some(rows) { + let buffers = if rows > 0 { + card_buffers(t, format, rows)? + } else { + Vec::new() + }; + let now = Last::of(granule, rows, buffers); + let effect = last.effect.take().filter(|_| now.small == last.small); + *last = Last { effect, ..now }; + } + } + let last = &mut *last; + heap.taken += last.large; + match &last.effect { + Some((found, e)) if same || *found == heap.left => { + for &(i, left) in &e.set { + heap.left[i] = left; + } + heap.left.extend_from_slice(&e.opened); + heap.taken += e.took; + same = true; + } + _ => { + let (found, open, taken) = (heap.left.clone(), heap.left.len(), heap.taken); + for &b in &last.small { + heap.alloc(b); + } + let set = (0..open) + .filter(|&i| heap.left[i] != found[i]) + .map(|i| (i, heap.left[i])) + .collect(); + let effect = Effect { + set, + opened: heap.left[open..].to_vec(), + took: heap.taken - taken, + }; + last.effect = Some((found, effect)); + same = false; + } } } + Ok(heap.taken) } - Ok(heap.taken) } /// The placement's levers, which a binary parses once @@ -2879,6 +3022,7 @@ impl Fill<'_, '_> { .routed .unwrap_or(RoutedRule::PerType(CardFormat::of_routed)); let uploads = card_uploads(self.model, self.c, self.rows, self.eligible, format)?; + let footprint = Footprint::new(card.granule_bytes, &uploads); let budget = i128::from(capped(card, self.card_budget)) - i128::from(self.kv) - i128::from( @@ -2888,19 +3032,17 @@ impl Fill<'_, '_> { - i128::from(card.margin_bytes) - i128::from(self.reserve); if let Some(b) = self.card_budget { - let dense = footprint(card.granule_bytes, &uploads, n_l)?; + let dense = footprint.at(n_l)?; check_floor(card, b, dense, self.kv, self.reserve)?; } if let Some(free) = card.free_bytes.filter(|&f| f < card.usable_bytes) { - let dense = footprint(card.granule_bytes, &uploads, n_l)?; + let dense = footprint.at(n_l)?; check_free_floor(card, free, dense, self.kv, self.reserve)?; } let caps: Vec = (0..n_l.len()) .map(|l| self.model.layer_experts(l)) .collect(); - spread(n_l, held, self.eligible, &caps, budget, |n| { - footprint(card.granule_bytes, &uploads, n) - }) + spread(n_l, held, self.eligible, &caps, budget, |n| footprint.at(n)) } } @@ -4942,6 +5084,102 @@ mod tests { } } + /// The footprint the expert rule asks once an expert is the allocator's replay of every + /// buffer in order, whatever counts it is asked at and in whatever order: fixed tensors + /// below and above a granule between stacks whose buffers cross it as their counts grow. + #[test] + fn a_footprint_is_the_replay_of_every_buffer_at_any_counts() { + let tensor = |name: &str, layer, role, ty: GgmlType, dims: Vec| ModelTensor { + name: name.into(), + shard: 0, + layer, + role, + ty, + dims, + file_bytes: 0, + gathered_rows: None, + }; + let (layers, granule) = (3, NonZeroU64::new(4096).expect("nonzero")); + let mut tensors = Vec::new(); + for l in 0..layers { + tensors.push(tensor( + "norm", + Some(l), + Role::Attention, + GgmlType::F32, + vec![256], + )); + tensors.push(tensor( + "big", + Some(l), + Role::Attention, + GgmlType::Q8_0, + vec![256, 20], + )); + tensors.push(tensor( + "small", + Some(l), + Role::Attention, + GgmlType::Q8_0, + vec![256, 3], + )); + tensors.push(tensor( + "stack", + Some(l), + Role::RoutedExperts, + GgmlType::Q4_K, + vec![256, 8, 9], + )); + } + let uploads: Vec<_> = tensors + .iter() + .map(|t| { + let format = CardFormat::of(t.ty).expect("a card format"); + let held = match t.role { + Role::RoutedExperts => Held::Experts { + layer: t.layer.expect("a layer"), + rows_per_expert: 8, + }, + _ => Held::Rows(rows_of(t)), + }; + (t, format, held) + }) + .collect(); + let replay = |n_l: &[u64]| { + let buffers = uploads.iter().flat_map(|&(t, format, held)| { + let rows = match held { + Held::Rows(rows) => rows, + Held::Experts { + layer, + rows_per_expert, + } => n_l[layer] * rows_per_expert, + }; + if rows == 0 { + Vec::new() + } else { + card_buffers(t, format, rows).expect("buffers") + } + }); + allocator_bytes(granule, buffers) + }; + let footprint = Footprint::new(granule, &uploads); + let mut n_l = vec![0u64; layers]; + let mut seed = 7u64; + for ask in 0..400 { + if ask < 27 { + n_l[ask % layers] += 1; + } else { + seed = seed.wrapping_mul(6_364_136_223_846_793_005).wrapping_add(1); + n_l[(seed >> 33) as usize % layers] = (seed >> 40) % 10; + } + assert_eq!( + footprint.at(&n_l).expect("a footprint"), + replay(&n_l), + "{n_l:?}" + ); + } + } + /// A model of 4 layers with every role's bf16 tensor (an attention /// matrix, a hyper-connection, a shared expert and its one-row gate, a /// router, an FFN norm, the head and the token embedding, 256-value rows) From 13e99026cd32917f973f1feb7a6b02f665d4d104 Mon Sep 17 00:00:00 2001 From: juliankang4 <200967947+juliankang4@users.noreply.github.com> Date: Sun, 25 Oct 2026 01:15:38 +0900 Subject: [PATCH 5/8] gpu: read a load's parts a run at a time over the pool parts::load read, normalized and staged one part at a time on one thread, and the host packings of the staged parts (the NVFP4 and MXFP4 repack, the E4M3 and BF16 MXFP8 planes) ran a row at a time on the same thread, so a 78 GB checkpoint load kept one core busy while the drive and the 48.9 GB/s copy lane waited. Qwen3.8 FP4-Local took about 270 s to load on GB10. The parts are now read in runs of at most 8 GiB, each run's parts at once over the resident worker pool (pool_map, each thread taking the next part as it finishes one), then staged in plan order, and the run's parts freed and their pages dropped over the pool. The host packings pack their rows in runs over the pool (pack_planes) into planes advised onto huge pages. A refusal is the first in plan or row order, with the same text as before. A source is Sync for its parts to be read at once. --- crates/gpu-glm5next/src/tp/mod.rs | 2 +- crates/gpu/Cargo.toml | 2 +- crates/gpu/src/mxfp4.rs | 66 ++++++------ crates/gpu/src/q8f32.rs | 161 ++++++++++++++++-------------- crates/gpu/src/upload.rs | 147 +++++++++++++++++++++++++++ 5 files changed, 273 insertions(+), 105 deletions(-) diff --git a/crates/gpu-glm5next/src/tp/mod.rs b/crates/gpu-glm5next/src/tp/mod.rs index 0f12759b5..8f84bc35b 100644 --- a/crates/gpu-glm5next/src/tp/mod.rs +++ b/crates/gpu-glm5next/src/tp/mod.rs @@ -620,7 +620,7 @@ impl Drop for Rank { /// ([`RankOf::take`]), in its type's card format. pub fn load( gpu: &Gpu, - source: &impl Source, + source: &(impl Source + Sync), at: RankOf, layers: usize, heads: usize, diff --git a/crates/gpu/Cargo.toml b/crates/gpu/Cargo.toml index d7d2a1386..c77ad696f 100644 --- a/crates/gpu/Cargo.toml +++ b/crates/gpu/Cargo.toml @@ -34,7 +34,7 @@ bloomery-verbs = { path = "../verbs" } # The expert tier's page shared with a tier node's process through a /dev/shm file (`graph::MappedHost`). memmap2.workspace = true libc.workspace = true -# The huge-page advice for large filled host memory (`tensor::host_filled`). +# The huge-page advice for large new host memory (`tensor::host_filled`, `upload::pack_planes`). bloomery-os.workspace = true # The pinned pool the hybrid host tier waits on and computes with (the CPU # engine dispatches its expert matmuls on the same pool). diff --git a/crates/gpu/src/mxfp4.rs b/crates/gpu/src/mxfp4.rs index 53fef5881..af7fd6a58 100644 --- a/crates/gpu/src/mxfp4.rs +++ b/crates/gpu/src/mxfp4.rs @@ -72,6 +72,7 @@ use crate::cores::{table16, table16_word}; use crate::kquant::q8_0::q8_0_code; use crate::tensor::Q8Act; +use crate::upload::pack_planes; use crate::{DeviceTensor, GpuError, launch_u32}; use cuda_core::{CudaContext, CudaStream, DeviceBuffer, LaunchConfig1D}; use cuda_device::dotprod::dp4a_s32; @@ -615,34 +616,43 @@ pub fn repack(bytes: &[u8], rows: usize, k: usize) -> Result 1 { - format!(" sub-block {s}") - } else { - String::new() - }; - GpuError::shape( - what, - format!( - "row {} block {}{sub}: scale byte {why}", - i / n_blk, - i % n_blk - ), - ) - })?; - let at = i * F::SCALE_BYTES + s; - e[at / 4] |= u32::from(byte) << (8 * (at % 4)); - } - for w in codes.as_chunks::<4>().0 { - qs.push(u32::from_le_bytes(*w)); - } - } + let code_words = (F::BLOCK_BYTES - F::SCALE_BYTES) / 4; + let widths = [n_blk * code_words, n_blk * F::SCALE_BYTES / 4]; + let (qs, e) = pack_planes( + what, + bytes, + (rows, n_blk * F::BLOCK_BYTES), + widths, + |first, src, qs, e| { + for (j, blk) in src.chunks_exact(F::BLOCK_BYTES).enumerate() { + let (scales, codes) = blk.split_at(F::SCALE_BYTES); + for (s, &byte) in scales.iter().enumerate() { + F::scale_ok(byte).map_err(|why| { + let sub = if F::SCALE_BYTES > 1 { + format!(" sub-block {s}") + } else { + String::new() + }; + let i = first * n_blk + j; + format!( + "row {} block {}{sub}: scale byte {why}", + i / n_blk, + i % n_blk + ) + })?; + let at = j * F::SCALE_BYTES + s; + e[at / 4] |= u32::from(byte) << (8 * (at % 4)); + } + for (w, c) in qs[j * code_words..][..code_words] + .iter_mut() + .zip(codes.as_chunks::<4>().0) + { + *w = u32::from_le_bytes(*c); + } + } + Ok(()) + }, + )?; Ok(Planes { qs, e }) } diff --git a/crates/gpu/src/q8f32.rs b/crates/gpu/src/q8f32.rs index 4d27d2623..a54aad1c7 100644 --- a/crates/gpu/src/q8f32.rs +++ b/crates/gpu/src/q8f32.rs @@ -95,6 +95,7 @@ use crate::gemm::{ }; use crate::launch_u32; use crate::tensor::DeviceTensor; +use crate::upload::pack_planes; use cuda_core::{CudaContext, CudaStream, DeviceBuffer, LaunchConfig1D}; use cuda_device::vector::{F32x4, as_vectors}; use cuda_device::{ @@ -3610,45 +3611,48 @@ pub fn e4m3_planes( ), )); } - let mut qs = Vec::with_capacity(rows * k / 4); - let mut e = Vec::with_capacity(rows * k / 128); - for (r, row) in bytes.chunks_exact(row_bytes).enumerate() { - let checked = match scales { - E4m3Scales::Mxfp8 => check_mxfp8(row), - E4m3Scales::Fp8B128 => check_fp8_b128(row), - }; - checked.map_err(|why| GpuError::shape(what, format!("row {r}: {why}")))?; - for group in row.as_chunks::<132>().0 { - match scales { - E4m3Scales::Mxfp8 => { - let mut word = 0u32; - for (i, blk) in group.as_chunks::<33>().0.iter().enumerate() { - word |= u32::from(blk[0]) << (8 * i); - qs.extend( - blk[1..] - .as_chunks::<4>() - .0 - .iter() - .map(|c| u32::from_le_bytes(*c)), - ); + pack_planes( + what, + bytes, + (rows, row_bytes), + [k / 4, k / 128], + |first, src, qs, e| { + for (r, row) in src.chunks_exact(row_bytes).enumerate() { + let checked = match scales { + E4m3Scales::Mxfp8 => check_mxfp8(row), + E4m3Scales::Fp8B128 => check_fp8_b128(row), + }; + checked.map_err(|why| format!("row {}: {why}", first + r))?; + for (g, group) in row.as_chunks::<132>().0.iter().enumerate() { + let at = r * (k / 128) + g; + let words = &mut qs[32 * at..][..32]; + match scales { + E4m3Scales::Mxfp8 => { + let mut word = 0u32; + for (i, blk) in group.as_chunks::<33>().0.iter().enumerate() { + word |= u32::from(blk[0]) << (8 * i); + for (w, c) in words[8 * i..][..8] + .iter_mut() + .zip(blk[1..].as_chunks::<4>().0) + { + *w = u32::from_le_bytes(*c); + } + } + e[at] = word; + } + E4m3Scales::Fp8B128 => { + let (scale, values) = group.split_at(4); + e[at] = u32::from_le_bytes([scale[0], scale[1], scale[2], scale[3]]); + for (w, c) in words.iter_mut().zip(values.as_chunks::<4>().0) { + *w = u32::from_le_bytes(*c); + } + } } - e.push(word); - } - E4m3Scales::Fp8B128 => { - let (scale, values) = group.split_at(4); - e.push(u32::from_le_bytes([scale[0], scale[1], scale[2], scale[3]])); - qs.extend( - values - .as_chunks::<4>() - .0 - .iter() - .map(|c| u32::from_le_bytes(*c)), - ); } } - } - } - Ok((qs, e)) + Ok(()) + }, + ) } /// [`Q8F32Kernels::enqueue_q8_0_gemv_heads_mcol`]'s arguments. The strides @@ -3761,46 +3765,53 @@ pub(crate) fn bf16_mxfp8_planes( format!("{rows} rows of {k} do not match {} BF16 bytes", src.len()), )); } - let mut qs = vec![0u32; count / 4]; - let mut scales = vec![0u32; count / 128]; - for (block, bytes) in src.as_chunks::<64>().0.iter().enumerate() { - let values: [f32; 32] = std::array::from_fn(|i| { - f32::from_bits(u32::from(u16::from_le_bytes([bytes[2 * i], bytes[2 * i + 1]])) << 16) - }); - if !values.iter().all(|v| v.is_finite()) { - return Err(GpuError::shape( - WHAT, - format!("block {block} has a non-finite value"), - )); - } - let amax = values.iter().fold(f32::MIN_POSITIVE, |a, v| a.max(v.abs())); - let exponent = ((amax / 448.0).log2().ceil() as i32).clamp(-127, 127); - scales[block / 4] |= ((exponent + 127) as u32) << (8 * (block % 4)); - let scale = 2.0f32.powi(exponent); - let inverse = 2.0f32.powi(-exponent); - for (i, value) in values.into_iter().enumerate() { - let magnitude = (value * inverse).abs(); - let code = if magnitude < 1.0 / 64.0 { - (magnitude * 512.0).round_ties_even() as u32 - } else { - let bits = magnitude.to_bits() - (120 << 23); - let kept = bits >> 20; - let tail = bits & 0x000f_ffff; - (kept + u32::from(tail > 0x0008_0000 || (tail == 0x0008_0000 && kept & 1 != 0))) - .min(126) - }; - if !(Mxfp8Bytes::value(code) * scale).is_finite() { - return Err(GpuError::shape( - WHAT, - format!("block {block} rounds a value past f32's range"), - )); + let row_blocks = k / 32; + pack_planes( + WHAT, + src, + (rows, 2 * k), + [k / 4, k / 128], + |first, src, qs, scales| { + for (local, bytes) in src.as_chunks::<64>().0.iter().enumerate() { + let block = first * row_blocks + local; + let values: [f32; 32] = std::array::from_fn(|i| { + f32::from_bits( + u32::from(u16::from_le_bytes([bytes[2 * i], bytes[2 * i + 1]])) << 16, + ) + }); + if !values.iter().all(|v| v.is_finite()) { + return Err(format!("block {block} has a non-finite value")); + } + let amax = values.iter().fold(f32::MIN_POSITIVE, |a, v| a.max(v.abs())); + let exponent = ((amax / 448.0).log2().ceil() as i32).clamp(-127, 127); + scales[local / 4] |= ((exponent + 127) as u32) << (8 * (local % 4)); + let scale = 2.0f32.powi(exponent); + let inverse = 2.0f32.powi(-exponent); + for (i, value) in values.into_iter().enumerate() { + let magnitude = (value * inverse).abs(); + let code = if magnitude < 1.0 / 64.0 { + (magnitude * 512.0).round_ties_even() as u32 + } else { + let bits = magnitude.to_bits() - (120 << 23); + let kept = bits >> 20; + let tail = bits & 0x000f_ffff; + (kept + + u32::from( + tail > 0x0008_0000 || (tail == 0x0008_0000 && kept & 1 != 0), + )) + .min(126) + }; + if !(Mxfp8Bytes::value(code) * scale).is_finite() { + return Err(format!("block {block} rounds a value past f32's range")); + } + let signed = code | ((value.to_bits() >> 24) & 0x80); + let at = local * 32 + i; + qs[at / 4] |= signed << (8 * (at % 4)); + } } - let signed = code | ((value.to_bits() >> 24) & 0x80); - let at = block * 32 + i; - qs[at / 4] |= signed << (8 * (at % 4)); - } - } - Ok((qs, scales)) + Ok(()) + }, + ) } #[cfg(test)] diff --git a/crates/gpu/src/upload.rs b/crates/gpu/src/upload.rs index 4de981f7a..155025fa0 100644 --- a/crates/gpu/src/upload.rs +++ b/crates/gpu/src/upload.rs @@ -11,6 +11,8 @@ use crate::GpuError; use crate::tensor::window; use cuda_core::{CudaStream, DeviceCopy, PinnedHostBuffer, sys}; use std::mem::ManuallyDrop; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::{Mutex, PoisonError}; // The staged uploads are the file's bytes read as the device's little-endian // words; a target of another endianness would need a repacking path, not a @@ -248,6 +250,79 @@ fn memcpy_pool(dst: &mut [u8], src: &[u8]) { }); } +/// `f` of each of `items` over the resident worker pool, the results in item +/// order: every thread takes the next item as it finishes one, so items of +/// uneven cost keep every core busy. `f` must not use the pool itself +/// (`for_each_chunk` is not reentrant). +pub(crate) fn pool_map(items: Vec, f: impl Fn(T) -> R + Sync) -> Vec { + let todo: Vec>> = items.into_iter().map(|t| Mutex::new(Some(t))).collect(); + let done: Vec>> = todo.iter().map(|_| Mutex::new(None)).collect(); + let next = AtomicUsize::new(0); + let pool = threads::pool(); + pool.for_each_chunk(pool.threads(), |_| { + loop { + let i = next.fetch_add(1, Ordering::Relaxed); + let Some(slot) = todo.get(i) else { break }; + let item = slot.lock().unwrap_or_else(PoisonError::into_inner).take(); + if let Some(item) = item { + let r = f(item); + *done[i].lock().unwrap_or_else(PoisonError::into_inner) = Some(r); + } + } + }); + done.into_iter() + .filter_map(|d| d.into_inner().unwrap_or_else(PoisonError::into_inner)) + .collect() +} + +/// A host packing of `rows` rows of `row_bytes` bytes each (`src`) into two +/// word planes of `widths[0]` and `widths[1]` words a row, over the resident +/// worker pool ([`pool_map`]): `pack(first, src, qs, e)` packs the run of +/// whole rows from row `first` that its slices hold, into planes zeroed +/// first (new memory on huge pages, [`os::advise_huge_within`]), and its first +/// refusal in row order is `what`'s. A short input packs on the calling +/// thread ([`TEAM_MIN`]). +pub(crate) fn pack_planes( + what: &'static str, + src: &[u8], + (rows, row_bytes): (usize, usize), + widths: [usize; 2], + pack: impl Fn(usize, &[u8], &mut [u32], &mut [u32]) -> Result<(), String> + Sync, +) -> Result<(Vec, Vec), GpuError> { + let words = |w: usize| { + rows.checked_mul(w) + .ok_or_else(|| GpuError::shape(what, format!("{rows} rows of {w} words pass usize"))) + }; + let (mut qs, mut e) = (vec![0u32; words(widths[0])?], vec![0u32; words(widths[1])?]); + os::advise_huge_within(&mut qs); + os::advise_huge_within(&mut e); + if rows.checked_mul(row_bytes) != Some(src.len()) { + return Err(GpuError::shape( + what, + format!("{} bytes are not {rows} rows of {row_bytes}", src.len()), + )); + } + let refused = if src.len() < TEAM_MIN { + pack(0, src, &mut qs, &mut e).err() + } else { + let run = rows.div_ceil(4 * threads::pool().threads()).max(1); + let tasks: Vec<_> = src + .chunks(run * row_bytes) + .zip(qs.chunks_mut((run * widths[0]).max(1))) + .zip(e.chunks_mut((run * widths[1]).max(1))) + .enumerate() + .map(|(i, ((src, qs), e))| (i * run, src, qs, e)) + .collect(); + pool_map(tasks, |(first, src, qs, e)| pack(first, src, qs, e)) + .into_iter() + .find_map(Result::err) + }; + match refused { + Some(detail) => Err(GpuError::shape(what, detail)), + None => Ok((qs, e)), + } +} + /// The destination of [`memcpy_pool`], shared with the pool's workers, each /// of which writes only its own chunk. struct SharedDst(*mut u8); @@ -264,3 +339,75 @@ pub(crate) fn bytes_of(data: &[T]) -> &[u8] { // reading bytes through a lower alignment is always sound. unsafe { std::slice::from_raw_parts(data.as_ptr().cast(), std::mem::size_of_val(data)) } } + +#[cfg(test)] +mod tests { + use super::TEAM_MIN; + use crate::GpuError; + use crate::mxfp4::{Mxfp4, Nvfp4, repack}; + use crate::q8f32::{E4m3Scales, bf16_mxfp8_planes, e4m3_planes}; + + type Packed = Result<(Vec, Vec), GpuError>; + + /// The planes `pack` makes of `rows` rows of `src` at once, past [`TEAM_MIN`] so its runs + /// go to the pool, are those of packing the rows one at a time. + fn pooled_as_by_rows(rows: usize, src: &[u8], pack: impl Fn(&[u8], usize) -> Packed) { + assert!(src.len() >= TEAM_MIN); + let (mut qs, mut e) = (Vec::new(), Vec::new()); + for row in src.chunks(src.len() / rows) { + let (q, x) = pack(row, 1).unwrap_or_else(|err| panic!("{err}")); + qs.extend(q); + e.extend(x); + } + let whole = pack(src, rows).unwrap_or_else(|err| panic!("{err}")); + assert_eq!(whole, (qs, e)); + } + + /// Every host packing over the pool (`pack_planes`) gives the planes of packing its rows one + /// at a time, and names the first refused row of the tensor, not of a run. + #[test] + fn pooled_packings_match_row_by_row_and_name_the_first_refused_row() { + let bytes = |n: usize, f: fn(usize) -> u8| (0..n).map(f).collect::>(); + let fp4 = |i: usize| (i * 7 % 0x7F) as u8; + let rows = 16384; + let mut nv = bytes(rows * 72, fp4); + pooled_as_by_rows(rows, &nv, |b, r| { + repack::(b, r, 128).map(|p| (p.qs, p.e)) + }); + let mx4 = bytes(rows * 68, |i| { + if i % 17 == 0 { + 127 + } else { + (i * 7 % 0x7F) as u8 + } + }); + pooled_as_by_rows(rows, &mx4, |b, r| { + repack::(b, r, 128).map(|p| (p.qs, p.e)) + }); + let mx8 = bytes( + 8192 * 132, + |i| if i % 33 == 0 { 127 } else { (i % 0x70) as u8 }, + ); + pooled_as_by_rows(8192, &mx8, |b, r| e4m3_planes(E4m3Scales::Mxfp8, b, r, 128)); + let b128 = bytes(8192 * 132, |i| match i % 132 { + 0..4 => 1.0f32.to_le_bytes()[i % 132], + _ => (i % 0x70) as u8, + }); + pooled_as_by_rows(8192, &b128, |b, r| { + e4m3_planes(E4m3Scales::Fp8B128, b, r, 128) + }); + let bf16 = bytes(4096 * 256, |i| (i % 0x7E) as u8); + pooled_as_by_rows(4096, &bf16, |b, r| bf16_mxfp8_planes(b, r, 128)); + for row in [rows - 1, rows / 2] { + nv[row * 72 + 36 + 1] = 0xFF; + } + let err = repack::(&nv, rows, 128) + .map(|_| ()) + .unwrap_err() + .to_string(); + assert!( + err.contains(&format!("row {} block 1 sub-block 1", rows / 2)), + "{err}" + ); + } +} From 7480c5f2630f4bb92d502bcfe8135c21abc1e263 Mon Sep 17 00:00:00 2001 From: juliankang4 <200967947+juliankang4@users.noreply.github.com> Date: Sun, 25 Oct 2026 02:54:34 +0900 Subject: [PATCH 6/8] model-hf: keep direct I/O failures visible past a cached shard --- crates/model-hf/src/part.rs | 38 +++++++++++++++++++++++++++++++++---- 1 file changed, 34 insertions(+), 4 deletions(-) diff --git a/crates/model-hf/src/part.rs b/crates/model-hf/src/part.rs index 2322a6b8c..3bdeed6d2 100644 --- a/crates/model-hf/src/part.rs +++ b/crates/model-hf/src/part.rs @@ -345,14 +345,20 @@ impl Whole { /// Stored tensor `name` read whole from its shard past the page cache (engram's direct reader) /// into new memory advised onto huge pages, so the read takes the drive's rate, faults once each -/// 2 MiB and leaves the page cache as it was. `None` where the shard cannot be opened for direct -/// reads or its mount refuses them ([`DirectFile::open`]): the caller reads the tensor through -/// the mapping. +/// 2 MiB and leaves the page cache as it was. `None` only when the shard's filesystem refuses +/// direct reads as invalid or unsupported ([`DirectFile::open`]); other I/O errors propagate. fn read_whole(folder: &Folder, name: &str) -> Result, Error> { let (path, range) = folder.file_range(name)?; let file = match DirectFile::open(&path) { Ok(file) => file, - Err(EngramError::OpenDirect { .. } | EngramError::DirectProbe { .. }) => return Ok(None), + Err(EngramError::OpenDirect { source, .. } | EngramError::DirectProbe { source, .. }) + if matches!( + source.kind(), + std::io::ErrorKind::InvalidInput | std::io::ErrorKind::Unsupported + ) => + { + return Ok(None); + } Err(e) => return Err(e.into()), }; let raw = folder.raw(name)?; @@ -1046,6 +1052,30 @@ pub(crate) mod tests { } } + #[test] + fn a_mapped_shard_does_not_hide_a_direct_open_error() { + let dir = temp_dir("direct-error"); + let c = checkpoint(&dir); + let (path, _) = c.folder.file_range("bf.w").unwrap(); + std::fs::rename(&path, path.with_extension("moved")).unwrap(); + let parts = c.plan(rule("bf.weight", 0, 0, 1)).unwrap(); + assert_eq!(parts.len(), 1); + for p in parts { + match c.read(&p) { + Err(Error::Direct(EngramError::OpenDirect { + path: found, + source, + })) => { + assert_eq!(found, path); + assert_eq!(source.kind(), std::io::ErrorKind::NotFound); + } + Err(e) => panic!("unexpected error: {e}"), + Ok(_) => panic!("the mapping hid the missing shard"), + } + } + std::fs::remove_dir_all(dir).unwrap(); + } + /// Before any value is read, a source the plan does not place, a refusal of the rule and a /// part off its card type's blocks are refused by name. #[test] From 140a21bb611fa78bf94c3f8cb2bf3077f5345bbc Mon Sep 17 00:00:00 2001 From: juliankang4 <200967947+juliankang4@users.noreply.github.com> Date: Sun, 25 Oct 2026 02:59:38 +0900 Subject: [PATCH 7/8] placement: reuse allocator accounting and drop unobserved cache recovery --- crates/placement/src/placement.rs | 23 ++++++----------------- 1 file changed, 6 insertions(+), 17 deletions(-) diff --git a/crates/placement/src/placement.rs b/crates/placement/src/placement.rs index 73fd6793f..f69bca884 100644 --- a/crates/placement/src/placement.rs +++ b/crates/placement/src/placement.rs @@ -2333,11 +2333,12 @@ struct Effect { } impl Last { - fn of(granule: u64, rows: u64, buffers: Vec) -> Last { - let (large, small): (Vec, Vec) = buffers.into_iter().partition(|&b| b >= granule); + fn of(granule: NonZeroU64, rows: u64, buffers: Vec) -> Last { + let (large, small): (Vec, Vec) = + buffers.into_iter().partition(|&b| b >= granule.get()); Last { rows: Some(rows), - large: large.iter().map(|b| b.div_ceil(granule) * granule).sum(), + large: allocator_bytes(granule, large), small, effect: None, } @@ -2360,7 +2361,7 @@ impl<'m> Footprint<'m> { }); continue; }; - let add = Last::of(granule.get(), 0, buffers); + let add = Last::of(granule, 0, buffers); if let Some(Step { what: What::Fixed, last, @@ -2380,18 +2381,6 @@ impl<'m> Footprint<'m> { } fn at(&self, n_l: &[u64]) -> Result { - let taken = self.walk(n_l); - if taken.is_err() { - // A walk cut short leaves later steps' granules from an older ask. - for step in &self.steps { - *step.last.borrow_mut() = Last::default(); - } - } - taken - } - - fn walk(&self, n_l: &[u64]) -> Result { - let granule = self.granule.get(); let mut heap = Heap::new(self.granule); // Whether the shared granules are as the last ask found them at this step. let mut same = true; @@ -2411,7 +2400,7 @@ impl<'m> Footprint<'m> { } else { Vec::new() }; - let now = Last::of(granule, rows, buffers); + let now = Last::of(self.granule, rows, buffers); let effect = last.effect.take().filter(|_| now.small == last.small); *last = Last { effect, ..now }; } From fb859cb4f71610c6b7b915f856e1cfb9cf24a9bc Mon Sep 17 00:00:00 2001 From: juliankang4 <200967947+juliankang4@users.noreply.github.com> Date: Sun, 25 Oct 2026 04:01:51 +0900 Subject: [PATCH 8/8] gpu: reserve parallel part reads in the host load budget --- crates/gguf/src/quant.rs | 1 - crates/gpu-gates/src/bin/bloomery_serve.rs | 84 ++--- .../shared/serve_seats/minimaxm3_folder.rs | 3 +- .../bin/shared/serve_seats/qwen38_folder.rs | 3 +- crates/gpu-gates/src/bind.rs | 1 + crates/gpu-gates/src/model_file.rs | 45 ++- crates/gpu/src/arch/qwen3moe/mtp38.rs | 10 +- crates/gpu/src/folder.rs | 31 +- crates/gpu/src/host/node.rs | 7 +- crates/gpu/src/model.rs | 54 ++- crates/gpu/src/parts.rs | 333 ++++++++++++++++-- crates/gpu/src/upload.rs | 3 + crates/hf/src/modelopt.rs | 2 - crates/model-hf/src/part.rs | 98 +++++- crates/model/src/part.rs | 111 ++++++ crates/placement/src/placement.rs | 12 +- crates/placement/src/placement/workstation.rs | 43 ++- tools/unsafe-ratchet.txt | 6 +- 18 files changed, 723 insertions(+), 124 deletions(-) diff --git a/crates/gguf/src/quant.rs b/crates/gguf/src/quant.rs index 24b4f86c8..1439a1b47 100644 --- a/crates/gguf/src/quant.rs +++ b/crates/gguf/src/quant.rs @@ -907,7 +907,6 @@ pub fn fp4_pair(byte: u8) -> [i8; 2] { /// plane runs it, and qdot's AVX2 body is held to it bit for bit. #[inline] #[must_use] -#[inline] pub fn nvfp4_group_to_block(g: &[u8; 8]) -> [u8; 8] { let spread = |half: [u8; 4]| { let x = u64::from(u32::from_le_bytes(half)); diff --git a/crates/gpu-gates/src/bin/bloomery_serve.rs b/crates/gpu-gates/src/bin/bloomery_serve.rs index 1811c67ca..06012fad1 100644 --- a/crates/gpu-gates/src/bin/bloomery_serve.rs +++ b/crates/gpu-gates/src/bin/bloomery_serve.rs @@ -358,42 +358,43 @@ models (the --model word is optional with a model file: the file's architecture /// the seat words. A qwen35moe file (Qwen3.6) is the qwen3 seat's. fn seat_of_file(generative: bool) -> Result<(Model, String), GateError> { let path = ref_model_path()?; - let split = model_file::header(&path)?; - let (seat, arch) = match split.architecture() { - Some(a) if side_seat::row(a).is_some() => { - return Err(format!( - "{} is a {a} file, an embedding model: serve it beside a chat model with \ + model_file::with_header(&path, |split| { + let (seat, arch) = match split.architecture() { + Some(a) if side_seat::row(a).is_some() => { + return Err(format!( + "{} is a {a} file, an embedding model: serve it beside a chat model with \ --embd-model", - path.display() - ) - .into()); - } - Some(a) if !generative && decide::ROWS.iter().any(|r| r.backbones.contains(&a)) => { - return Err(format!( - "{}: {}", - path.display(), - serve::decide::no_head(a, decide::ROWS) - ) - .into()); - } - Some(a) => ( - Model::ALL - .iter() - .copied() - .find(|m| m.serves(a)) - .ok_or_else(|| no_seat(&path, &split, a))?, - a.to_owned(), - ), - None => { - return Err(format!( - "{}: the first shard names no architecture; the seats are --model {}", - path.display(), - seat_words() - ) - .into()); - } - }; - Ok((seat, arch)) + path.display() + ) + .into()); + } + Some(a) if !generative && decide::ROWS.iter().any(|r| r.backbones.contains(&a)) => { + return Err(format!( + "{}: {}", + path.display(), + serve::decide::no_head(a, decide::ROWS) + ) + .into()); + } + Some(a) => ( + Model::ALL + .iter() + .copied() + .find(|m| m.serves(a)) + .ok_or_else(|| no_seat(&path, split, a))?, + a.to_owned(), + ), + None => { + return Err(format!( + "{}: the first shard names no architecture; the seats are --model {}", + path.display(), + seat_words() + ) + .into()); + } + }; + Ok((seat, arch)) + })? } /// The refusal of a file of architecture `a` that no seat serves: what no @@ -425,12 +426,13 @@ models (the --model word is optional with a model file: the file's architecture /// ([`ref_model_path`]) and its `.decision.type`, when it has one. fn file_facts() -> Result<(String, Option), GateError> { let path = ref_model_path()?; - let split = model_file::header(&path)?; - let decision = bloomery_gpu_gates::model_file::decision_type(&split)?.map(str::to_owned); - Ok(( - split.architecture().unwrap_or("").to_owned(), - decision, - )) + model_file::with_header(&path, |split| { + let decision = model_file::decision_type(split)?.map(str::to_owned); + Ok(( + split.architecture().unwrap_or("").to_owned(), + decision, + )) + })? } /// The head the arguments name ([`serve::decide::pick`]) and the rest: diff --git a/crates/gpu-gates/src/bin/shared/serve_seats/minimaxm3_folder.rs b/crates/gpu-gates/src/bin/shared/serve_seats/minimaxm3_folder.rs index 35e25c253..5758775f9 100644 --- a/crates/gpu-gates/src/bin/shared/serve_seats/minimaxm3_folder.rs +++ b/crates/gpu-gates/src/bin/shared/serve_seats/minimaxm3_folder.rs @@ -32,8 +32,7 @@ pub struct Opened { /// body does not read), is refused by name. #[cfg(feature = "checkpoint")] pub fn open(dir: &Path) -> Result { - let f = model_hf::folder::open(dir) - .map_err(|e| format!("checkpoint folder {}: {e}", dir.display()))?; + let f = bloomery_gpu_gates::model_file::open_folder(dir)?; let arch = refset::arch::minimaxm3::ARCH; if f.header.architecture() != Some(arch) { return Err(format!( diff --git a/crates/gpu-gates/src/bin/shared/serve_seats/qwen38_folder.rs b/crates/gpu-gates/src/bin/shared/serve_seats/qwen38_folder.rs index bc50ba0d2..65171c2ea 100644 --- a/crates/gpu-gates/src/bin/shared/serve_seats/qwen38_folder.rs +++ b/crates/gpu-gates/src/bin/shared/serve_seats/qwen38_folder.rs @@ -24,8 +24,7 @@ pub struct Opened { /// architecture is refused by name. #[cfg(feature = "checkpoint")] pub fn open(path: &Path) -> Result { - let mut f = model_hf::folder::open(path) - .map_err(|e| format!("checkpoint folder {}: {e}", path.display()))?; + let mut f = bloomery_gpu_gates::model_file::open_folder(path)?; if f.header.architecture() != Some("qwen4exp") { return Err(format!( "checkpoint folder {}: a {} model, where this seat serves qwen4exp", diff --git a/crates/gpu-gates/src/bind.rs b/crates/gpu-gates/src/bind.rs index 90545a5ff..9e442d79c 100644 --- a/crates/gpu-gates/src/bind.rs +++ b/crates/gpu-gates/src/bind.rs @@ -3597,6 +3597,7 @@ mod tests { shadows: 0, reserves: 0, extra: 0, + read_ahead: 0, arena: 0, cards: 0, os: 0, diff --git a/crates/gpu-gates/src/model_file.rs b/crates/gpu-gates/src/model_file.rs index e8e941ced..5e0398b81 100644 --- a/crates/gpu-gates/src/model_file.rs +++ b/crates/gpu-gates/src/model_file.rs @@ -29,6 +29,47 @@ pub use fetch::{ pub(crate) const NONE: &str = "no model file: pass -m PATH (the first shard, or a checkpoint folder) \ or --hf [:], or set BLOOMERY_REF_MODEL"; +#[cfg(feature = "checkpoint")] +static KEPT: std::sync::Mutex> = + std::sync::Mutex::new(None); + +/// The folder at `path`, reusing the open whose header selected its server seat. +#[cfg(feature = "checkpoint")] +pub fn open_folder(path: &Path) -> Result { + let kept = KEPT + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .take(); + match kept { + Some((p, opened)) if p == path => Ok(opened), + _ => read_folder(path), + } +} + +#[cfg(feature = "checkpoint")] +fn read_folder(path: &Path) -> Result { + model_hf::folder::open(path) + .map_err(|e| format!("checkpoint folder {}: {e}", path.display()).into()) +} + +/// `f` of the model's header, keeping a checkpoint folder for its seat's open. +pub fn with_header(path: &Path, f: impl FnOnce(&Split) -> R) -> Result { + #[cfg(feature = "checkpoint")] + if path.is_dir() { + let mut kept = KEPT + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let opened = match kept.take() { + Some((p, opened)) if p == path => opened, + _ => read_folder(path)?, + }; + let result = f(&opened.header); + *kept = Some((path.to_path_buf(), opened)); + return Ok(result); + } + header(path).map(|header| f(&header)) +} + /// The header of the model at `path`, which a server seats it by and `-ngl N` is compared with: /// its GGUF split's first shard, or a checkpoint folder's when `path` is a directory, which the /// map its config names reads (`model_hf::folder::open`), so the seat is chosen from the header's @@ -43,9 +84,7 @@ pub fn header(path: &Path) -> Result { /// The header of the checkpoint folder `path`. #[cfg(feature = "checkpoint")] fn folder_header(path: &Path) -> Result { - model_hf::folder::open(path) - .map(|f| f.header) - .map_err(|e| format!("checkpoint folder {}: {e}", path.display()).into()) + read_folder(path).map(|f| f.header) } /// A checkpoint folder in a build without the `checkpoint` feature, refused by name. diff --git a/crates/gpu/src/arch/qwen3moe/mtp38.rs b/crates/gpu/src/arch/qwen3moe/mtp38.rs index 002013566..6c2cead37 100644 --- a/crates/gpu/src/arch/qwen3moe/mtp38.rs +++ b/crates/gpu/src/arch/qwen3moe/mtp38.rs @@ -260,7 +260,7 @@ impl Mtp38 { pub fn open( gpu: &Gpu, target_w: &Weights, - (draft_file, parts): (&Split, Option<&crate::model::PartsLoad<'_>>), + (draft_file, parts): (&Split, Option<&dyn crate::model::PartsLoad>), mtp: &model::arch::qwen35moe::place::MtpInputs, plan: &model::arch::qwen35moe::place::MtpPlan<'_>, card_dontneed: bool, @@ -309,7 +309,13 @@ impl Mtp38 { )?; let rows = file_rows(&plan.draft); let mut w = match parts { - Some(load) => load(stream, &rows, 0, card_dontneed)?, + Some(parts) => { + let available = model::placement::workstation::host_available() + .map_err(|e| GpuError::plan(WHAT, e))?; + let (need, room) = parts.host_need(&rows, 0, 0, available)?; + need.check(available).map_err(|e| GpuError::plan(WHAT, e))?; + parts.load(stream, &rows, 0, card_dontneed, room)? + } None => Weights::load_placed(stream, draft_file, &rows, 0, card_dontneed)?, }; for t in &mtp.tensors { diff --git a/crates/gpu/src/folder.rs b/crates/gpu/src/folder.rs index eef303c97..47163a960 100644 --- a/crates/gpu/src/folder.rs +++ b/crates/gpu/src/folder.rs @@ -11,23 +11,18 @@ //! machine's source reads the experts where the folder stores them, through the located header //! (`crate::host::swap_source::FileSwap`). -use cuda_core::CudaStream; use gguf::Split; use model::part::Source; -use model::placement::Plan; use crate::GpuError; use crate::host::swap::Residency; use crate::model::CardBytes; -use crate::parts::{InPlace, load_placed_in_place}; -use crate::weights::Weights; +use crate::parts::{InPlace, PlacedSource}; /// What the errors name. const WHAT: &str = "folder::Model::into_origin"; -/// The card's tensors of a plan from a part source ([`load_placed_in_place`]). -type CardParts = - Box, usize, bool) -> Result + Send>; +type CardParts = Box; /// The model a load reads: a GGUF split, or a checkpoint folder ([`Folder`]). `X` is what the /// folder's architecture carries beside the folder's parts. @@ -91,7 +86,7 @@ impl From> for Model { /// `gguf::Split::header_only` for a folder whose host reads nothing where the folder stores it), /// the header its plan is made from, whose in-place tensors the host reads where the folder /// stores them ([`InPlace`]), and its card tensors through its part source -/// ([`load_placed_in_place`]). `X` is the architecture's own addition to the load, `()` for +/// ([`crate::parts::load_placed_in_place`]). `X` is the architecture's own addition to the load, `()` for /// none. pub struct Folder { header: Split, @@ -103,17 +98,15 @@ impl Folder { /// The folder of header `header`, its card tensors read from `source`, with the /// architecture's `extra`. A load reads only its plan's tensors of the source, so the source /// may also hold another plan's (a draft's, which loads through [`Origin::bytes`] too). - pub fn with(header: Split, source: S, extra: X) -> Folder { + pub fn with( + header: Split, + source: S, + extra: X, + ) -> Folder { let in_place = InPlace::of(&header); Folder { header, - parts: Box::new(move |stream, plan, card, release| { - let selected = model::part::Filtered { - source: &source, - keep: |name: &str| plan.model.tensors.iter().any(|t| t.name == name), - }; - load_placed_in_place(stream, &selected, plan, card, release, &in_place) - }), + parts: Box::new(PlacedSource { source, in_place }), extra, } } @@ -121,7 +114,7 @@ impl Folder { impl Folder { /// The folder of header `header`, its card tensors read from `source`. - pub fn new(header: Split, source: S) -> Folder { + pub fn new(header: Split, source: S) -> Folder { Folder::with(header, source, ()) } } @@ -175,6 +168,10 @@ mod tests { Err(std::fmt::Error) } + fn read_bytes(&self, _: &Planned) -> Result { + Err(std::fmt::Error) + } + fn read(&self, _: &Planned) -> Result, Self::Error> { Err(std::fmt::Error) } diff --git a/crates/gpu/src/host/node.rs b/crates/gpu/src/host/node.rs index 0d2b3df3b..5c7f2a3a3 100644 --- a/crates/gpu/src/host/node.rs +++ b/crates/gpu/src/host/node.rs @@ -417,7 +417,12 @@ impl NodeOpen { None => { Weights::load_placed(gpu.stream(), file, plan, self.card, card_dontneed)? } - Some(load) => load(gpu.stream(), plan, self.card, card_dontneed)?, + Some(load) => { + let available = + workstation::host_available().map_err(|e| GpuError::plan(WHAT, e))?; + let room = load.host_bytes(plan, self.card, available)?; + load.load(gpu.stream(), plan, self.card, card_dontneed, room)? + } }; let experts = (family.experts)(&gpu, file, &tensors, &w, &self)?; (w, experts) diff --git a/crates/gpu/src/model.rs b/crates/gpu/src/model.rs index 6089bb5c7..8b0a08529 100644 --- a/crates/gpu/src/model.rs +++ b/crates/gpu/src/model.rs @@ -2162,13 +2162,47 @@ pub enum CardBytes<'a> { /// A part source's, for a file that is a checkpoint folder's header /// (`gguf::Split::header_only`): card `card` of `plan` loaded on the stream, each part's file /// pages released when the host config's `card_dontneed` asks (`crate::parts::load_placed_in_place`). - Parts(&'a PartsLoad<'a>), + Parts(&'a dyn PartsLoad), } -/// A part source's load of card `card` of `plan` on the stream, `bool` the host config's -/// `card_dontneed` ([`CardBytes::Parts`]). -pub type PartsLoad<'a> = - dyn Fn(&CudaStream, &Plan<'_>, usize, bool) -> Result + 'a; +/// A part source's host reservation and load of card `card` of `plan`. +pub trait PartsLoad: Send { + /// Plan within offered host room, including pinned upload staging. + fn host_bytes(&self, plan: &Plan<'_>, card: usize, room: u64) -> Result; + /// Load within the checked peak returned by host_bytes. + fn load( + &self, + stream: &CudaStream, + plan: &Plan<'_>, + card: usize, + release: bool, + room: u64, + ) -> Result; + + /// The host admission and exact read peak to pass to load. + fn host_need( + &self, + plan: &Plan<'_>, + card: usize, + extra: u64, + available: u64, + ) -> Result<(HostNeed, u64), GpuError> { + let mut need = HostNeed::of(plan, extra); + let room = available.saturating_sub(need.bytes()); + let peak = self.host_bytes(plan, card, room)?; + if peak > 0 { + let ring = crate::upload::UploadRing::HOST_MAX; + need.read_ahead = peak.checked_sub(ring).ok_or_else(|| { + GpuError::shape("PartsLoad", "read peak is under its staging ring") + })?; + need.shadows = need + .shadows + .checked_add(ring) + .ok_or_else(|| GpuError::shape("PartsLoad", "pinned ring bytes pass u64"))?; + } + Ok((need, peak)) + } +} /// The steps of a placed load before its body: the plan's card and context /// budget, the card, the host's memory against the plan's host need, the @@ -2226,7 +2260,11 @@ fn placed_pre( } else { 0 }; - let grant = HostNeed::of(plan, extra_bytes) + let (need, room) = match &bytes { + CardBytes::File => (HostNeed::of(plan, extra_bytes), 0), + CardBytes::Parts(parts) => parts.host_need(plan, card, extra_bytes, available)?, + }; + let grant = need .grant(available, least) .map_err(|e| GpuError::plan(what, format!("{e}{}", plan.host.room_clause())))?; let t = Instant::now(); @@ -2234,7 +2272,9 @@ fn placed_pre( CardBytes::File => { Weights::load_placed(gpu.stream(), file, plan, card, host.card_dontneed)? } - CardBytes::Parts(load) => load(gpu.stream(), plan, card, host.card_dontneed)?, + CardBytes::Parts(parts) => { + parts.load(gpu.stream(), plan, card, host.card_dontneed, room)? + } }; let upload = t.elapsed(); let t = Instant::now(); diff --git a/crates/gpu/src/parts.rs b/crates/gpu/src/parts.rs index 2e0769b8d..2daadc538 100644 --- a/crates/gpu/src/parts.rs +++ b/crates/gpu/src/parts.rs @@ -1,8 +1,8 @@ //! The one load of a part source onto a card (`model::part::Source`: a GGUF file or a //! checkpoint folder). The caller's rule says which tensors the card holds, whole, the part of //! a rank, or a stack's listed experts, as the source plans them; every planned part's card -//! format is checked before any is read; then each part is read and staged in turn through one -//! staging ring, which one synchronize ends ([`load`]). A rank of a tensor-parallel load of N +//! format is checked before any is read; then parts are read a run at a time over the pool and +//! staged in order through one staging ring, which one synchronize ends ([`load`]). A rank of a tensor-parallel load of N //! ranks passes its family's cut; a one-card load passes its placement plan's card //! ([`OnCard`]), and a folder whose host reads tensors where its located header holds them //! passes those ([`InPlace`]): the card takes its share of a stack the plan splits between it @@ -12,65 +12,227 @@ use std::collections::{BTreeMap, HashMap, HashSet}; use cuda_core::CudaStream; use gguf::{GgmlType, Place, Split, TensorInfo}; -use model::part::{Source, Take}; +use model::part::{Planned, Source, Take, read_ahead}; use model::placement::{CardFormat, Device, ExpertList, Format, ModelTensor, Plan, Segment}; use crate::GpuError; -use crate::upload::UploadRing; +use crate::upload::{UploadRing, pool_map}; use crate::weights::{TensorPart, Weights, part_format, stage_parts}; /// What the errors name. const WHAT: &str = "parts::load"; -/// The parts of `source` that `rule` takes, each read and uploaded to the card of `stream` in -/// the card format `card` names for it, else its type's ([`part_format`]). Every part is planned -/// and its format checked, and refused by name, before any is read; then one part at a time is -/// read and staged, so the load holds one part's bytes of its own beside the card (a file's -/// parts are its mapping's and hold none). The weights are the caller's only once every copy -/// has ended: a refusal at any part drops what was staged. Under `release` each part's file -/// pages are dropped once it is staged ([`Source::release`]), as a split's placed load drops its -/// card segments' (`BLOOMERY_CARD_DONTNEED`). -pub fn load( +/// The parts selected by `rule`, read over the pool in runs bounded by owned read and packing +/// bytes, then staged in plan order. Oversized parts run alone. The placed caller reserves +/// `LoadPlan::host_bytes` in its host need; standalone loads check the live host reading. +/// A refusal drains pending copies before freeing device buffers. Under `release`, each +/// run's source pages are dropped after staging. +pub fn load( stream: &CudaStream, source: &S, rule: impl Fn(&str, GgmlType, &[u64]) -> Result, card: impl Fn(&TensorInfo) -> Option, release: bool, ) -> Result { - let planned = source.plan(rule).map_err(|e| refused(e.to_string()))?; - let mut budget = 0usize; - for p in &planned { - budget = budget - .checked_add(part_format(&p.info)?.1) - .ok_or_else(|| refused("the parts' staging bytes pass usize".into()))?; + let available = + model::placement::workstation::host_available().map_err(|e| refused(e.to_string()))?; + let plan = LoadPlan::new(source, rule, available)?; + load_plan(stream, source, &plan, card, release, available) +} + +fn load_plan( + stream: &CudaStream, + source: &S, + plan: &LoadPlan, + card: impl Fn(&TensorInfo) -> Option, + release: bool, + available: u64, +) -> Result { + plan.check(available)?; + for p in &plan.parts { if let Some(format) = card(&p.info) { card_layout(&p.info, format)?; } } stream.context().bind_to_thread()?; let mut staged = BTreeMap::new(); - // Made after `staged` and dropped before it: its drop drains the copies, so no buffer of a - // refused load is freed while a copy still writes it. - let mut ring = UploadRing::new(stream, budget)?; - for p in &planned { - let part = source.read(p).map_err(|e| refused(e.to_string()))?; - let one = TensorPart { - info: part.info.clone(), - spans: part.spans(), - scale: part.scale.clone(), - }; - let held = |name: &str| staged.contains_key(name); - let added = stage_parts(stream, &mut ring, &[one], held, &card)?; - staged.extend(added); - drop(part); - if release { - source.release(p).map_err(|e| refused(e.to_string()))?; + // Dropped before the staged device buffers: draining the ring protects refused loads. + let mut ring = UploadRing::new(stream, plan.upload_bytes)?; + for run in plan.checked_runs(|| { + model::placement::workstation::host_available().map_err(|e| refused(e.to_string())) + }) { + let run = run?; + let parts = pool_map(run.iter().collect(), |p| { + source.read(p).map_err(|e| e.to_string()) + }); + let mut done = Vec::with_capacity(run.len()); + for (p, part) in run.iter().zip(parts) { + let part = part.map_err(refused)?; + let one = TensorPart { + info: part.info.clone(), + spans: part.spans(), + scale: part.scale.clone(), + }; + let held = |name: &str| staged.contains_key(name); + let added = stage_parts(stream, &mut ring, &[one], held, &card)?; + staged.extend(added); + done.push((p, part)); } + let released = pool_map(done, |(p, part)| { + drop(part); + match release { + true => source.release(p).map_err(|e| e.to_string()), + false => Ok(()), + } + }); + released + .into_iter() + .collect::>() + .map_err(refused)?; } ring.finish(stream)?; Ok(Weights::of_staged(staged)) } +/// A placed source shares one admission and upload implementation across model components. +pub(crate) struct PlacedSource { + pub(crate) source: S, + pub(crate) in_place: InPlace, +} + +impl crate::model::PartsLoad for PlacedSource { + fn host_bytes(&self, plan: &Plan<'_>, card: usize, room: u64) -> Result { + Ok( + LoadPlan::placed(&self.source, plan, card, &self.in_place, room)? + .1 + .host_bytes, + ) + } + + fn load( + &self, + stream: &CudaStream, + plan: &Plan<'_>, + card: usize, + release: bool, + room: u64, + ) -> Result { + load_placed_in_place( + stream, + &self.source, + plan, + card, + release, + &self.in_place, + room, + ) + } +} + +/// Parts, their bounded parallel reads and the peak owned host bytes reserved before loading. +struct LoadPlan { + parts: Vec, + runs: Vec<(std::ops::Range, u64)>, + host_bytes: u64, + upload_bytes: usize, +} + +impl LoadPlan { + fn checked_runs<'a>( + &'a self, + mut available: impl FnMut() -> Result + 'a, + ) -> impl Iterator> + 'a { + self.runs.iter().map(move |(range, bytes)| { + // The ring is already paid for in the live reading. + let live = available()?; + if *bytes > live { + return Err(refused(format!( + "part read-ahead needs {bytes} B, the host has {live} B available" + ))); + } + Ok(&self.parts[range.clone()]) + }) + } + + fn check(&self, available: u64) -> Result<(), GpuError> { + if self.host_bytes > available { + return Err(refused(format!( + "part read-ahead needs {} B, the host has {available} B available", + self.host_bytes + ))); + } + Ok(()) + } + + fn new( + source: &S, + rule: impl Fn(&str, GgmlType, &[u64]) -> Result, + room: u64, + ) -> Result { + let parts = source.plan(rule).map_err(|e| refused(e.to_string()))?; + let mut upload_bytes = 0usize; + let mut costs = Vec::with_capacity(parts.len()); + for p in &parts { + let staged = part_format(&p.info)?.1; + upload_bytes = upload_bytes + .checked_add(staged) + .ok_or_else(|| refused("the parts' staging bytes pass usize".into()))?; + // Packing capacity and scale copies coexist with the entire run's read results. + let cost = source + .read_bytes(p) + .map_err(|e| refused(e.to_string()))? + .checked_add( + (staged as u64).checked_mul(2).ok_or_else(|| { + refused(format!("{}: packing bytes pass u64", p.info.name)) + })?, + ) + .ok_or_else(|| refused(format!("{}: read-ahead bytes pass u64", p.info.name)))?; + costs.push(cost); + } + let mut runs = Vec::new(); + let mut host_bytes = 0; + let mut at = 0; + while at < parts.len() { + let (n, bytes) = read_ahead( + &parts[at..], + &costs[at..], + room.saturating_sub(UploadRing::HOST_MAX), + ) + .map_err(|e| refused(e.to_string()))?; + runs.push((at..at + n, bytes)); + host_bytes = host_bytes.max(bytes); + at += n; + } + if host_bytes > 0 { + host_bytes = host_bytes + .checked_add(UploadRing::HOST_MAX) + .ok_or_else(|| refused("the parts' read-ahead and ring bytes pass u64".into()))?; + } + Ok(Self { + parts, + runs, + host_bytes, + upload_bytes, + }) + } + + fn placed( + source: &S, + plan: &Plan<'_>, + card: usize, + in_place: &InPlace, + room: u64, + ) -> Result<(OnCard, Self), GpuError> { + let on = OnCard::of_in_place(plan, card, in_place)?; + let selected = model::part::Filtered { + source, + keep: |name: &str| plan.model.tensors.iter().any(|t| t.name == name), + }; + let read = Self::new(&selected, |name, _, _| on.take(name), room)?; + Ok((on, read)) + } +} + /// The tensors a located header holds in place (`gguf::Place::File` or `gguf::Place::Stack`), /// which the host reads where the folder stores them: a load from a part source skips such a /// tensor the plan holds on the host, and takes its card's experts of such a stack the plan @@ -293,21 +455,25 @@ impl OnCard { /// Card `card` of `plan` from `source` ([`OnCard::of_in_place`], [`load`], each part's pages /// released under `release`), checked against the plan; the host reads the tensors `in_place` /// names where they are stored. -pub fn load_placed_in_place( +pub fn load_placed_in_place( stream: &CudaStream, source: &S, plan: &Plan<'_>, card: usize, release: bool, in_place: &InPlace, + room: u64, ) -> Result { - let on = OnCard::of_in_place(plan, card, in_place)?; - let w = load( + let (on, read) = LoadPlan::placed(source, plan, card, in_place, room)?; + let available = + model::placement::workstation::host_available().map_err(|e| refused(e.to_string()))?; + let w = load_plan( stream, source, - |name, _, _| on.take(name), + &read, |t| on.format(t), release, + available.min(room), )?; on.check(&w)?; Ok(w) @@ -396,6 +562,95 @@ mod tests { use super::*; use model::placement::Role; + #[test] + fn later_reads_use_their_own_peak_under_a_falling_host_room() { + let mib = 1u64 << 20; + let decls = [64, 8] + .into_iter() + .map(|n| gguf::write::TensorDecl { + name: format!("p{n}.weight"), + dims: vec![64, n * mib / 256], + type_id: GgmlType::F32.as_u32(), + nbytes: n * mib, + }) + .collect(); + let layout = gguf::write::Layout::new(&[], decls).unwrap(); + let source = + Split::header_only(std::path::PathBuf::from("/synthetic/falling-room"), &layout) + .unwrap(); + let plan = + LoadPlan::new(&source, |_, _, _| Ok::<_, String>(Take::Whole), 300 * mib).unwrap(); + assert_eq!(plan.runs, [(0..1, 192 * mib), (1..2, 24 * mib)]); + assert_eq!(plan.host_bytes, 288 * mib); + plan.check(300 * mib).unwrap(); + assert!(plan.check(288 * mib - 1).is_err()); + for readings in [[228, 164], [192, 24]] { + let mut rooms = readings.into_iter(); + let runs = plan + .checked_runs(|| Ok(rooms.next().unwrap() * mib)) + .collect::, _>>() + .expect("each remaining run fits its live room"); + assert_eq!( + runs.iter() + .map(|r| r[0].info.name.as_str()) + .collect::>(), + ["p64.weight", "p8.weight"] + ); + } + let mut rooms = [228, 23].into_iter(); + let mut runs = plan.checked_runs(|| Ok(rooms.next().unwrap() * mib)); + assert!(runs.next().unwrap().is_ok()); + let error = runs.next().unwrap().unwrap_err().to_string(); + assert!(error.contains("needs 25165824 B"), "{error}"); + assert!(error.contains("24117248 B available"), "{error}"); + } + + #[test] + #[ignore = "needs a CUDA card"] + fn hw_a_parts_load_refuses_a_short_host_before_reading() { + struct Owned; + impl Source for Owned { + type Error = std::io::Error; + fn plan( + &self, + rule: impl Fn(&str, GgmlType, &[u64]) -> Result, + ) -> Result, Self::Error> { + let info = TensorInfo { + name: "owned.weight".into(), + ty: GgmlType::BF16, + dims: vec![64, 262144], + offset: 0, + nbytes: 32 << 20, + }; + rule(&info.name, info.ty, &info.dims) + .map_err(|e| std::io::Error::other(e.to_string()))?; + Ok(vec![Planned { at: None, info }; 256]) + } + fn read_bytes(&self, p: &Planned) -> Result { + Ok(p.info.nbytes * 2) + } + fn read(&self, _: &Planned) -> Result, Self::Error> { + Err(std::io::Error::other("read before host admission")) + } + } + let plan = + LoadPlan::new(&Owned, |_, _, _| Ok::<_, String>(Take::Whole), 128 << 20).unwrap(); + assert_eq!(plan.runs.len(), 256); + assert_eq!(plan.host_bytes, (192 + 3 * 32) << 20); + let gpu = crate::Gpu::with_device(0).unwrap(); + let error = load_plan(gpu.stream(), &Owned, &plan, |_| None, false, 128 << 20) + .err() + .expect("a short host must refuse the load") + .to_string(); + assert!( + error.contains("part read-ahead needs 301989888 B"), + "{error}" + ); + assert!(error.contains("134217728 B available"), "{error}"); + assert!(plan.check(plan.host_bytes).is_ok()); + assert!(plan.check(plan.host_bytes - 1).is_err()); + } + /// The rule a load from a part source plans by: a tensor whole on the card is taken in the /// plan's card format, an unused tensor, a row-gathered table on the NVMe tier and a tensor /// whole on another card (that card's load reads it) are skipped, and a tensor the plan diff --git a/crates/gpu/src/upload.rs b/crates/gpu/src/upload.rs index 155025fa0..d4eca7622 100644 --- a/crates/gpu/src/upload.rs +++ b/crates/gpu/src/upload.rs @@ -62,6 +62,9 @@ pub(crate) struct Stage<'a> { } impl UploadRing { + /// Upper bound on this ring's pinned host bytes, including small loads. + pub(crate) const HOST_MAX: u64 = (SLOTS * SLOT_MAX) as u64; + /// A ring for a load of `budget` upload bytes: [`SLOTS`] slots of the /// budget's share, each between [`SLOT_MIN`] and [`SLOT_MAX`]; a load of /// nothing takes no slots. `stream`'s context must be current (the diff --git a/crates/hf/src/modelopt.rs b/crates/hf/src/modelopt.rs index 718882df2..1a9dbdb05 100644 --- a/crates/hf/src/modelopt.rs +++ b/crates/hf/src/modelopt.rs @@ -1369,8 +1369,6 @@ fn append<'a>( Ok([rows, k]) } - - /// The stored tensor `name` of dtype `want` whose bytes fill its shape; bytes that do not are /// `unfilled`'s refusal of the tensor, with what is wrong. fn stored<'a>( diff --git a/crates/model-hf/src/part.rs b/crates/model-hf/src/part.rs index 3bdeed6d2..fe818a2fd 100644 --- a/crates/model-hf/src/part.rs +++ b/crates/model-hf/src/part.rs @@ -19,7 +19,7 @@ use hf::modelopt::{self, Normal, Use}; use hf::quant_config::{Placed, Plan}; use model::part::{self, At, Part, Planned, Source, Take}; use model::placement::ExpertList; -use model::placement::tp::{TpError, part_of}; +use model::placement::tp::TpError; use safetensors::{Dtype, Folder, Raw}; /// Why a part was refused. @@ -217,6 +217,56 @@ impl Source for Checkpoint { Ok(out) } + fn read_bytes(&self, p: &Planned) -> Result { + let name = &p.info.name; + let m = self + .made + .get(name) + .ok_or_else(|| refuse(name, "is not one the map makes"))?; + let placed = self.placed(m); + let ty = self.card_type(name, m)?; + let whole = whole(name, ty, &m.dims)?; + let cut = p.at.is_some() && m.step.is_none(); + let format = modelopt::format_of(&placed)?; + let compressed = matches!( + format, + hf::quant_config::Format::Nvfp4Csf | hf::quant_config::Format::Nvfp4Csf32 + ); + let plain = matches!( + format, + hf::quant_config::Format::Bf16 | hf::quant_config::Format::F32 + ); + // Direct inputs, or growing cut buffers; decoded companion words and widened plain + // values coexist with them. A step may shrink the tensor, so its pre-step output is + // bounded from the stored bytes too. Output capacity can grow to twice its length. + let mut bytes = 0u64; + for n in stored(&placed) { + let companion = !m.sources.iter().any(|source| source == n); + let factor = if cut { 2u64 } else { 1 } + + u64::from(compressed) * 2 + + u64::from(m.step.is_some()) * 3 + + u64::from(plain || companion) * 4; + let (_, range) = self.folder.file_range(n)?; + bytes = range + .end + .checked_sub(range.start) + .and_then(|n| n.checked_mul(factor)) + .and_then(|n| n.checked_add((2 * DIRECT_ALIGN) as u64)) + .and_then(|n| bytes.checked_add(n)) + .ok_or_else(|| refuse(name, "stored read bytes pass u64"))?; + } + let normal = if cut { p.info.nbytes } else { whole.nbytes }; + let moved = if m.step.is_some() { whole.nbytes } else { 0 }; + normal + .checked_mul(2) + .and_then(|n| n.checked_add(moved)) + .and_then(|n| n.checked_add(p.info.nbytes)) + .and_then(|n| n.checked_add(if compressed { whole.nbytes } else { 0 })) + .and_then(|n| n.checked_add((m.sources.len() as u64).checked_mul(8)?)) + .and_then(|n| bytes.checked_add(n)) + .ok_or_else(|| refuse(name, "normalized read bytes pass u64")) + } + /// A tensor its format refuses, or that does not normalize into the planned type and bytes, /// is refused by name. fn read(&self, p: &Planned) -> Result, Error> { @@ -1052,6 +1102,52 @@ pub(crate) mod tests { } } + #[test] + fn read_peaks_cover_the_owned_inputs_and_results() { + let dir = temp_dir("read-peak"); + let mut c = checkpoint(&dir); + let shrinking = c.made.get_mut("mx.weight").unwrap(); + shrinking.dims[1] = 1; + shrinking.step = Some(Box::new(|n| { + let row = n.bytes.len() / n.dims[1] as usize; + n.bytes.truncate(row); + n.dims[1] = 1; + Ok(()) + })); + let planned = c.plan(|_, _, _| Ok::<_, String>(Take::Whole)).unwrap(); + for p in &planned { + let read_bytes = c.read_bytes(p).unwrap(); + let part = c.read(p).unwrap(); + let output: u64 = part.spans().iter().map(|s| s.len() as u64).sum(); + let made = &c.made[&p.info.name]; + let placed = c.placed(made); + let before = modelopt::normalize(|n| c.folder.raw(n), &placed, made.role).unwrap(); + let output = output.max(before.bytes.capacity() as u64); + let stored: u64 = stored(&placed) + .map(|n| { + let (_, r) = c.folder.file_range(n).unwrap(); + let (_, span, _) = aligned_span(r.start, r.end - r.start); + (span + DIRECT_ALIGN) as u64 + }) + .sum(); + assert!( + read_bytes >= stored + output, + "{}: {read_bytes} < {}", + p.info.name, + stored + output + ); + } + let p = planned.iter().find(|p| p.info.name == "v.weight").unwrap(); + c.made.get_mut("v.weight").unwrap().dims = vec![u64::MAX, 2]; + assert!( + c.read_bytes(p) + .unwrap_err() + .to_string() + .contains("v.weight") + ); + std::fs::remove_dir_all(&dir).unwrap(); + } + #[test] fn a_mapped_shard_does_not_hide_a_direct_open_error() { let dir = temp_dir("direct-error"); diff --git a/crates/model/src/part.rs b/crates/model/src/part.rs index 972fb95ae..1d192d5d5 100644 --- a/crates/model/src/part.rs +++ b/crates/model/src/part.rs @@ -159,6 +159,37 @@ pub struct Planned { pub info: TensorInfo, } +/// One read's leading parts, bounded by normalized payload and owned host room. An oversized +/// first part runs alone; the caller checks its returned peak before allocating it. +/// `costs` holds each already-checked part's owned read and packing bytes. +pub fn read_ahead(planned: &[Planned], costs: &[u64], room: u64) -> Result<(usize, u64), Error> { + const PAYLOAD: u64 = 8 << 30; + if planned.len() != costs.len() { + return Err(refuse( + "read-ahead", + "the parts and their costs differ in length", + )); + } + let (mut n, mut bytes, mut payload) = (0, 0u64, 0u64); + for (p, &cost) in planned.iter().zip(costs) { + let next = bytes + .checked_add(cost) + .ok_or_else(|| refuse(&p.info.name, "read-ahead bytes pass u64"))?; + let more = payload + .checked_add(p.info.nbytes) + .ok_or_else(|| refuse(&p.info.name, "read-ahead payload passes u64"))?; + if n > 0 && (more > PAYLOAD || next > room) { + break; + } + (bytes, payload) = (next, more); + n += 1; + if payload > PAYLOAD || bytes > room { + break; + } + } + Ok((n, bytes)) +} + /// A part's bytes in its card type ([`Source::read`]), with its header entry and the tensor /// scales its type reads. pub struct Part<'s> { @@ -226,6 +257,11 @@ pub trait Source { rule: impl Fn(&str, GgmlType, &[u64]) -> Result, ) -> Result, Self::Error>; + /// An upper bound on the owned payload bytes a read of `p` holds, including its result, + /// normalization scratch, stored inputs and a contiguous copy for a packed upload. + /// Borrowed file pages are not owned bytes. + fn read_bytes(&self, p: &Planned) -> Result; + /// The bytes of the part `p` planned ([`Source::plan`]) and the tensor scales its type reads; /// a tensor no longer in the source is refused by name. fn read(&self, p: &Planned) -> Result, Self::Error>; @@ -261,6 +297,10 @@ impl bool> Source for Filtered<'_, S, F> { }) } + fn read_bytes(&self, p: &Planned) -> Result { + self.source.read_bytes(p) + } + fn read(&self, p: &Planned) -> Result, Self::Error> { self.source.read(p) } @@ -291,6 +331,29 @@ impl Source for Split { Ok(out) } + fn read_bytes(&self, p: &Planned) -> Result { + let name = &p.info.name; + let (_, t) = self + .find(name) + .ok_or_else(|| refuse(name, "is not in the file"))?; + // A noncontiguous rank cut needs a contiguous copy when its upload packs planes. + // tensor_scales may also hold the whole scale vector beside a rank's cut of it. + let scales = if t.ty.reads_scale() { + t.dims + .last() + .copied() + .unwrap_or(1) + .checked_mul(8) + .ok_or_else(|| refuse(name, "scale read bytes pass u64"))? + } else { + 0 + }; + p.info + .nbytes + .checked_add(scales) + .ok_or_else(|| refuse(name, "row copy and scale read bytes pass u64")) + } + /// A scale the file lacks or does not hold as its type reads it is refused by name. fn read(&self, p: &Planned) -> Result, Error> { let name = p.info.name.as_str(); @@ -322,6 +385,52 @@ mod tests { use super::*; use gguf::write::{Layout, TensorDecl}; + #[test] + fn a_read_takes_the_parts_that_fit_and_at_least_one() { + let part = |nbytes| Planned { + at: None, + info: TensorInfo { + name: "p".into(), + dims: vec![1], + ty: GgmlType::F32, + offset: 0, + nbytes, + }, + }; + let gib = 1u64 << 30; + let parts = [part(3 * gib), part(5 * gib), part(4)]; + assert_eq!( + read_ahead(&parts, &[6 * gib, 10 * gib, 8], 16 * gib).unwrap(), + (2, 16 * gib) + ); + assert_eq!( + read_ahead(&parts, &[6 * gib, 10 * gib, 8], 16 * gib - 1).unwrap(), + (1, 6 * gib) + ); + assert_eq!( + read_ahead(&[part(9 * gib), part(4)], &[9 * gib, 4], 20 * gib).unwrap(), + (1, 9 * gib) + ); + assert_eq!( + read_ahead(&[part(gib)], &[2 * gib], gib).unwrap(), + (1, 2 * gib) + ); + assert_eq!(read_ahead(&[], &[], 0).unwrap(), (0, 0)); + let small = vec![part(32 << 20); 256]; + let costs = vec![192 << 20; 256]; + assert_eq!( + read_ahead(&small, &costs, (1024 - 96) << 20).unwrap(), + (4, 768 << 20) + ); + assert_eq!( + read_ahead(&small, &costs, 48 * gib).unwrap(), + (256, 48 * gib) + ); + assert!(read_ahead(&parts, &[], gib).is_err()); + assert!(read_ahead(&[part(4), part(4)], &[1, u64::MAX], u64::MAX).is_err()); + assert!(read_ahead(&[part(4), part(u64::MAX)], &[1, 1], u64::MAX).is_err()); + } + #[test] fn shared_parts_are_planned_only_for_the_selected_component() { let names = ["target.weight", "draft.weight"]; @@ -347,6 +456,8 @@ mod tests { let parts = selected .plan(|_, _, _| Ok::<_, std::convert::Infallible>(Take::Whole)) .unwrap(); + assert_eq!(source.read_bytes(&parts[0]).unwrap(), 4); + assert_eq!(selected.read_bytes(&parts[0]).unwrap(), 4); assert_eq!( parts .iter() diff --git a/crates/placement/src/placement.rs b/crates/placement/src/placement.rs index f69bca884..484d9acc2 100644 --- a/crates/placement/src/placement.rs +++ b/crates/placement/src/placement.rs @@ -2859,7 +2859,7 @@ fn resolve_shares( }; let own = eligible(layers, stacks, model, routed)?; let uploads = card_uploads(model, c, rows, &own, routed)?; - let dense = footprint(card.granule_bytes, &uploads, &zeros)?; + let dense = Footprint::new(card.granule_bytes, &uploads).at(&zeros)?; let kv = kv_card.saturating_add(stage_rows); let floor = plan_floor(card, dense, kv, reserve); let cap = capped(card, None); @@ -2898,7 +2898,7 @@ fn resolve_shares( let routable = eligible(0..model.layers, stacks, model, routed)?; let caps: Vec = (0..model.layers).map(|l| model.layer_experts(l)).collect(); let experts = card_uploads(model, cards.len(), rows, &routable, routed)?; - let total = footprint(cards[0].granule_bytes, &experts, &caps)?; + let total = Footprint::new(cards[0].granule_bytes, &experts).at(&caps)?; let per = named .iter() .zip(card_budget::split(&parts, total)) @@ -6422,6 +6422,13 @@ mod tests { need.bytes(), "{how:?}: the classes sum to the need" ); + let reading = HostNeed { + read_ahead: 17, + ..need + }; + let with_reads = reading.classes(&plan, how).expect("a load's read reserve"); + assert_eq!(with_reads.anon, got.anon + 17); + assert_eq!(with_reads.bytes(), reading.bytes()); } // The table: each term's class under both holdings. let locked = Hosting { @@ -6445,6 +6452,7 @@ mod tests { (PageClass::Cache, PageClass::Cache), (PageClass::Cache, PageClass::Cache), (PageClass::Cache, PageClass::Locked), + (PageClass::Anon, PageClass::Anon), (PageClass::Anon, PageClass::Pinned), (PageClass::Unified, PageClass::Unified), ] diff --git a/crates/placement/src/placement/workstation.rs b/crates/placement/src/placement/workstation.rs index eaf29c9b1..6b27baac9 100644 --- a/crates/placement/src/placement/workstation.rs +++ b/crates/placement/src/placement/workstation.rs @@ -854,9 +854,12 @@ pub fn spec_of(card: &Card) -> Option { pub struct HostNeed { pub experts: u64, pub tables: u64, + /// Pinned ring buffers, including the part loader's upload staging at load time. pub shadows: u64, pub reserves: u64, pub extra: u64, + /// Owned part results and their concurrent read/normalization scratch during loading. + pub read_ahead: u64, /// The plan's NVMe expert tier arena (`HostTotals::nvme_arena_bytes`): /// the bytes the tier's slots fill beside these terms, anonymous or /// pinned, 0 without one. [`HostNeed::bytes`] counts it; @@ -897,6 +900,7 @@ impl HostNeed { shadows: totals.shadow_bytes, reserves: totals.reserve_bytes, extra, + read_ahead: 0, arena: totals.nvme_arena_bytes, cards: totals.pool_card_bytes, os: os_reserve(host), @@ -925,6 +929,7 @@ impl HostNeed { self.shadows, self.reserves, self.extra, + self.read_ahead, if arena { self.arena } else { 0 }, self.cards, ] @@ -1089,6 +1094,9 @@ impl fmt::Display for HostShort { if n.extra > 0 { write!(f, " + residency churn pool {} B", n.extra)?; } + if n.read_ahead > 0 { + write!(f, " + part read-ahead {} B", n.read_ahead)?; + } if n.arena > 0 { write!(f, " + NVMe tier arena {} B", n.arena)?; if let Some(floor) = self.floor { @@ -1178,7 +1186,7 @@ pub enum NeedTerm { /// The row-gathered tables the host holds ([`HostNeed::tables`]): the /// same mapping as the experts. Tables, - /// The cards' ring shadows ([`HostNeed::shadows`]). + /// The cards' pinned ring buffers and load-time upload staging ([`HostNeed::shadows`]). Shadows, /// The expert tier's prompt batch rows and places /// ([`TIER_BATCH_HOST_RESERVE`]), allocated with `cuMemAllocHost`. @@ -1199,6 +1207,8 @@ pub enum NeedTerm { /// The residency churn pool ([`HostNeed::extra`]): runs of the host set, /// read in and locked with it. Extra, + /// Owned part results and read/normalization scratch ([`HostNeed::read_ahead`]). + PartReads, /// The NVMe tier's RAM arena ([`HostNeed::arena`]). Arena, /// The cards' share of a unified pool ([`HostNeed::cards`]). @@ -1207,7 +1217,7 @@ pub enum NeedTerm { impl NeedTerm { /// Every term, in the order [`HostNeed::classes`] adds them. - pub const ALL: [NeedTerm; 11] = [ + pub const ALL: [NeedTerm; 12] = [ NeedTerm::Experts, NeedTerm::Tables, NeedTerm::Shadows, @@ -1217,6 +1227,7 @@ impl NeedTerm { NeedTerm::RowCache, NeedTerm::RunAhead, NeedTerm::Extra, + NeedTerm::PartReads, NeedTerm::Arena, NeedTerm::Cards, ]; @@ -1237,6 +1248,7 @@ impl NeedTerm { | NeedTerm::Checkpoints | NeedTerm::PointRows => PageClass::Pinned, NeedTerm::RowCache | NeedTerm::RunAhead => PageClass::Cache, + NeedTerm::PartReads => PageClass::Anon, NeedTerm::Arena => { if how.arena_pinned { PageClass::Pinned @@ -1369,6 +1381,7 @@ impl HostNeed { add(NeedTerm::RowCache, rows); add(NeedTerm::RunAhead, window); add(NeedTerm::Extra, self.extra); + add(NeedTerm::PartReads, self.read_ahead); add(NeedTerm::Arena, self.arena); add(NeedTerm::Cards, self.cards); Ok(classes) @@ -2322,6 +2335,7 @@ mod tests { shadows: 10, reserves: 500, extra: 50, + read_ahead: 0, arena: 0, cards: 0, os: 400, @@ -2352,6 +2366,7 @@ mod tests { shadows: 0, reserves: 11_955_785_728, extra: 0, + read_ahead: 0, arena: 20_335_071_232, cards: 0, os: 6_694_629_376, @@ -2359,6 +2374,29 @@ mod tests { } } + #[test] + fn a_parts_load_reserves_its_read_ahead_before_allocating() { + let base = unpaged_need(); + let need = HostNeed { + read_ahead: 8 << 30, + ..base + }; + let available = base.bytes() + (1 << 30); + assert!(base.check(available).is_ok()); + let short = need + .grant(available, 0) + .expect_err("owned parts exceed the host limit"); + assert_eq!(short.need.read_ahead, 8 << 30); + assert!(short.to_string().contains("part read-ahead 8589934592 B")); + assert!(need.check(base.bytes() + (8 << 30)).is_ok()); + assert!(need.check(base.bytes() + (8 << 30) - 1).is_err()); + let overflow = HostNeed { + read_ahead: u64::MAX, + ..base + }; + assert!(overflow.check(u64::MAX).is_err()); + } + const PAGED_AVAILABLE: u64 = 27_499_458_560; /// A reading under the need by less than the arena's room above its @@ -2439,6 +2477,7 @@ mod tests { shadows: 0, reserves: 8_533_127_168, extra: 0, + read_ahead: 0, arena: 0, cards: 0, os: 6_694_629_376, diff --git a/tools/unsafe-ratchet.txt b/tools/unsafe-ratchet.txt index 6071636fb..a4cdeec0f 100644 --- a/tools/unsafe-ratchet.txt +++ b/tools/unsafe-ratchet.txt @@ -552,7 +552,7 @@ gguf 7 2 0 # PIN(2026-10-24): gpu 2414 -> 2418 blocks, new code, not old code relaxed: the two branch pins the rebase onto main # folded away, the four-slot tier handoff (ds41_ffn_handoff_4_tier, +2) and Muse's attention branch (its fork and join, # +2); 2331 + 83 + 2 + 2, recounted by tools/check-unsafe.sh. -gpu 2418 311 10 +gpu 2417 311 10 # PIN(2026-10-01): gpu-deepseek41 459 -> 464 blocks, 51 -> 52 fns, a new kernel entry, not old code relaxed: # ds41_ffn_card_acc_8_tier (the tier form of the card sum at 8 slots, chain/ffn/batch.rs) adds its two blocks and # card_acc_tier_at (the one unsafe fn) with its three, each with its SAFETY line. @@ -757,7 +757,9 @@ models 0 0 0 # and sysctl reads; model and model-hf come down by the five that moved. # PIN(2026-10-24): os 30 -> 32 blocks, 5 -> 7 fns, new code, not old code relaxed: advise_huge # (madvise(MADV_HUGEPAGE), its Linux and macOS bodies) and its test, each with its SAFETY line. -os 32 7 0 +# PIN(2026-10-25): os 32 -> 33 blocks, code moved in, not old code relaxed: advise_huge_within (moved from gpu's +# host_filled advice, gpu comes down by the one that moved), with its SAFETY line. +os 33 7 0 placement 0 0 0 # PIN(2026-10-02): qdot 69 -> 77 blocks, 36 -> 41 fns, new kernels, not old code relaxed: the Q4_K/Q5_K row # lanes (pack_lanes_avx2, dot_lanes_avx2 and their helpers transpose8_epi32, transpose8_ps, lane_codes: five