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/gguf/src/quant.rs b/crates/gguf/src/quant.rs index d7e9f9bc8..1439a1b47 100644 --- a/crates/gguf/src/quant.rs +++ b/crates/gguf/src/quant.rs @@ -908,10 +908,14 @@ pub fn fp4_pair(byte: u8) -> [i8; 2] { #[inline] #[must_use] 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/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-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/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/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/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/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/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/gpu/src/upload.rs b/crates/gpu/src/upload.rs index 4de981f7a..d4eca7622 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 @@ -60,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 @@ -248,6 +253,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 +342,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}" + ); + } +} 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..1a9dbdb05 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,9 @@ 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()); + // 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) @@ -1325,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() @@ -1346,10 +1352,17 @@ 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); 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..fe818a2fd 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 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 { @@ -212,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> { @@ -220,9 +275,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 +374,71 @@ 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` 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 { 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)?; + 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 { @@ -969,6 +1102,76 @@ 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"); + 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] 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/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 diff --git a/crates/placement/src/placement.rs b/crates/placement/src/placement.rs index db5f8416a..484d9acc2 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,160 @@ 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: NonZeroU64, rows: u64, buffers: Vec) -> Last { + let (large, small): (Vec, Vec) = + buffers.into_iter().partition(|&b| b >= granule.get()); + Last { + rows: Some(rows), + large: allocator_bytes(granule, large), + 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, 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 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(self.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 @@ -2727,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); @@ -2766,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)) @@ -2879,6 +3011,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 +3021,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 +5073,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) @@ -6195,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 { @@ -6218,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