diff --git a/Cargo.lock b/Cargo.lock index 1eff1890a..fdf519f0d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3157,6 +3157,7 @@ dependencies = [ "opentelemetry-otlp", "opentelemetry_sdk", "parking_lot", + "rayon", "safetensors", "serde_json", "tokio", diff --git a/docs/subsystems/frontend/startup-time.md b/docs/subsystems/frontend/startup-time.md index cfb31255d..4068ee888 100644 --- a/docs/subsystems/frontend/startup-time.md +++ b/docs/subsystems/frontend/startup-time.md @@ -30,10 +30,22 @@ Consequences to keep in mind: - The graceful ctrl-c handler is installed only **after** the engine load resolves (`cancel_token_on_ctrl_c`). The blocking load can't be cancelled, so during load SIGINT keeps its default kill behavior — same as before the change. (When testing this: background jobs of non-interactive shells inherit SIGINT=SIG_IGN; use a spawner that restores default dispositions.) - LoRA mode stays sequential: the LoRA routes need the handle when the router is built. -## Pinned-staging H2D upload (rejected 2026-06, adopted 2026-07) +## Pinned-staging H2D upload (rejected 2026-06, adopted 2026-07, retuned 2026-07) The 2026-06 rejection reasoned from the numbers above: upload ran at ~7GB/s pageable, and with the engine path (~1.35s) barely above the concurrent frontend path (~1.25s), a faster upload stood to gain ≲0.1s of HTTP-ready. By 2026-07 the engine path had grown far past that floor (warm HTTP-ready ~5.2s on sm_89), so the frontend overlap no longer hides the upload: the pinned double-buffer pipeline (`WeightStager` in openinfer-core) cuts the warm Qwen3-4B load phase ~1.26s → ~0.69s and HTTP-ready 5.22s → 4.66s — the phase saving passes through in full. (The timed load phase ends at the `GPU model loaded` log, before the mmap teardown; HTTP-ready includes everything.) Cold ready stays storage-bound (5.39s → 5.33s on local NVMe). On dual-GH200 Qwen3-14B TP2 the picture is platform-dependent: pageable copies over NVLink-C2C are already fast, so the pipeline wins there only once the strided column-shard gather lands with it (warm rank-0 load 1.19s → 0.69s vs main; the intermediate gather-only state regresses to 1.33s), validated by the TP2 golden gate with the page cache warm and cold. +### Staging geometry and the fill team + +The first rollout shipped 64 MiB slots and four workers spawned per chunk. The 64 MiB event-sync rationale did not survive the later sensitivity sweep, which found the constant flat from 16 to 64 MiB; the four workers did hold up — on sm_89 that is the smallest team whose fill saturates the H2D path. A ten-arm sweep (two-GPU GH200, Qwen3-32B, TP2, 144 cores, five interleaved repetitions) moved the shipping point to 32 MiB slots filled by a resident rayon team of eight, halving pinned memory to 64 MiB across both buffers. + +The geometry alone is dominated by pooling at the old geometry: spawn 8/32 MiB reaches 8272 ms to ready at 28.84 s of CPU while pool 4/64 MiB reaches 8262 ms at 24.25 s, better on both axes with no constant touched. Pooling alone is not a dead arm — against main's 8903 ms at 21.20 s it buys 641 ms for 3.05 CPU-s. The resident team is what turns the smaller chunk from a loss into a gain: spawning at eight workers, 64→32 MiB costs 7993→8272 ms to ready; pooling at eight, it buys 7651→7356 ms. CPU is bought by the worker count, not the chunk size: at fixed chunk, four workers to eight costs 12.6–14.8% more CPU on the pooled arms; at fixed workers, 64→32 MiB costs 0.3–2.2%. + +**`FILL_THREADS` is per rank**, but ranks do not load concurrently: `Qwen3Executor::from_runtime_with_lora_options` finishes one `Qwen3Model` before starting the next, and each rank's loader — which owns the fill pool — is dropped at the end of its own load. One team of eight is therefore the most that ever fills, at any world size. Concurrency knees at that eight (2/4/8/12 workers → 5482/3583/2720/2701 ms of the two-rank load marker, the last for 20.7% more CPU to ready), so dividing a host budget by `world_size` only shrinks the one active team: at a fixed budget of 16 — four workers per rank at TP4, two at TP8 — eight per rank wins HTTP-ready 5/5 at both widths, by paired medians of 268 ms and 1119 ms, at the cost of 0.31 s and 0.78 s more process CPU (eight-GPU sm_120 host, Qwen3-4B, five order-rotated pairs per width). A host-level bound becomes relevant only where rank-local fill pools run concurrently. GLM5.2 does that under `--glm52-weight-staging`, submitting every rank's load before receiving any result: four fill workers per local rank, so 16 on a single-host EP4 run, and none at all on its default pageable path. On multi-host EP the host total follows the ranks local to that host, not the global EP width. + +On the post-#752 base, where uploads run back to back from `finish()`, five interleaved repetitions pair to medians of -259 ms of GPU-model load and -541 ms to ready for +4.94 s of process CPU — about nine CPU-seconds per second of startup saved. That load figure is the model's own `GPU model loaded in` line, not the two-rank marker above; the two are not comparable. On a single sm_89 GPU the change is flat to ready, measured on the pre-#752 base and not repeated since: that host is transfer-bound at a 23.3 GiB/s H2D roofline against 24.6 GiB/s of fill. + +Halving the buffer also halves the column-shard ceiling `prepare_cols` enforces, from 33.5 M to 16.7 M columns per shard; no supported Qwen3 checkpoint approaches it. + ## Next The warm HTTP-ready floor is now the engine's own post-load startup work (profile, warmup, graph capture — ~3.7s of the ~4.7s), no longer the frontend path; that is where the next startup-time win lives. Cold-start stays bandwidth-bound (per roadmap out of scope). diff --git a/openinfer-core/Cargo.toml b/openinfer-core/Cargo.toml index e1f643301..6a2ba7a6f 100644 --- a/openinfer-core/Cargo.toml +++ b/openinfer-core/Cargo.toml @@ -22,6 +22,7 @@ opentelemetry_sdk = { workspace = true } openinfer-engine = { workspace = true } openinfer-kernels = { workspace = true } parking_lot = { workspace = true } +rayon = { workspace = true } safetensors = { workspace = true } serde_json = { workspace = true } tokio = { workspace = true } diff --git a/openinfer-core/src/weight_loader/staging.rs b/openinfer-core/src/weight_loader/staging.rs index 5154f89e7..a4250c455 100644 --- a/openinfer-core/src/weight_loader/staging.rs +++ b/openinfer-core/src/weight_loader/staging.rs @@ -1,3 +1,4 @@ +use std::mem::MaybeUninit; use std::sync::Arc; use anyhow::Result; @@ -10,15 +11,111 @@ use cudarc::driver::result::memcpy_htod_async; use cudarc::driver::sys::CUevent_flags; use half::bf16; use log::error; +use rayon::ThreadPool; +use rayon::ThreadPoolBuilder; use crate::tensor::DeviceContext; const BF16_SIZE: usize = std::mem::size_of::(); -/// Bytes per pinned staging buffer: 64 MiB amortizes the per-chunk event sync -/// while capping pinned memory at 128 MiB across both buffers. -const STAGE_BYTES: usize = 64 << 20; -/// A single memcpy thread cannot keep up with the pinned H2D copy rate. -const FILL_THREADS: usize = 4; +/// Per-buffer staging chunk. The measured 32 MiB geometry improves overlap and +/// limits the two pinned buffers to 64 MiB. +const STAGE_BYTES: usize = 32 << 20; +/// Widest fill team per rank, before the core-count cap. +const FILL_THREADS: usize = 8; +/// Small tails do not amortize dispatching work across the fill team. +const PARALLEL_FILL_MIN_BYTES: usize = 1 << 20; + +fn fill_threads() -> usize { + static WIDTH: std::sync::OnceLock = std::sync::OnceLock::new(); + *WIDTH.get_or_init(|| { + let width = std::thread::available_parallelism() + .expect("query available CPU parallelism") + .get() + .min(FILL_THREADS); + if width < FILL_THREADS { + log::info!("weight fill: {width} threads, capped by the core count"); + } + width + }) +} + +fn as_uninit(src: &[u8]) -> &[MaybeUninit] { + // SAFETY: MaybeUninit shares u8's layout and adds no validity + // requirement, so widening an initialized slice is sound. + unsafe { std::slice::from_raw_parts(src.as_ptr().cast(), src.len()) } +} + +struct FillPool { + pool: ThreadPool, + workers: usize, +} + +impl FillPool { + fn new() -> Result { + Self::with_workers(fill_threads()) + } + + fn with_workers(workers: usize) -> Result { + let pool = ThreadPoolBuilder::new() + .num_threads(workers) + .thread_name(|worker| format!("weight-fill-{worker}")) + .build() + .map_err(|e| anyhow::anyhow!("build weight-fill pool failed: {e}"))?; + Ok(Self { pool, workers }) + } + + fn copy(&self, src: &[u8], dst: &mut [MaybeUninit]) { + debug_assert_eq!(src.len(), dst.len()); + if src.len() < PARALLEL_FILL_MIN_BYTES || self.workers == 1 { + dst.copy_from_slice(as_uninit(src)); + return; + } + let per = src.len().div_ceil(self.workers); + self.pool.scope(|scope| { + for (src_part, dst_part) in src.chunks(per).zip(dst.chunks_mut(per)) { + scope.spawn(move |_| dst_part.copy_from_slice(as_uninit(src_part))); + } + }); + } + + fn gather_cols( + &self, + src: &[u8], + stride_b: usize, + off_b: usize, + take_b: usize, + rows: usize, + dst: &mut [MaybeUninit], + ) { + debug_assert!(off_b + take_b <= stride_b); + debug_assert!(rows * stride_b <= src.len()); + debug_assert_eq!(rows * take_b, dst.len()); + if dst.len() < PARALLEL_FILL_MIN_BYTES || self.workers == 1 { + for (src_row, dst_row) in src[..rows * stride_b] + .chunks(stride_b) + .zip(dst.chunks_mut(take_b)) + { + dst_row.copy_from_slice(as_uninit(&src_row[off_b..off_b + take_b])); + } + return; + } + let rows_per = rows.div_ceil(self.workers); + self.pool.scope(|scope| { + for (src_rows, dst_rows) in src[..rows * stride_b] + .chunks(rows_per * stride_b) + .zip(dst.chunks_mut(rows_per * take_b)) + { + scope.spawn(move |_| { + for (src_row, dst_row) in + src_rows.chunks(stride_b).zip(dst_rows.chunks_mut(take_b)) + { + dst_row.copy_from_slice(as_uninit(&src_row[off_b..off_b + take_b])); + } + }); + } + }); + } +} struct StagingBuf { pinned: PinnedHostSlice, @@ -39,6 +136,7 @@ pub(crate) struct ColShardPlan { pub(crate) struct WeightStager { stream: Arc, bufs: [StagingBuf; 2], + fill: FillPool, next: usize, } @@ -58,6 +156,7 @@ impl WeightStager { Ok(Self { stream: ctx.stream.clone(), bufs: [make()?, make()?], + fill: FillPool::new()?, next: 0, }) } @@ -68,11 +167,7 @@ impl WeightStager { pub(crate) unsafe fn upload_at(&mut self, src: &[u8], dst_at: u64) -> Result<()> { for (i, chunk) in src.chunks(STAGE_BYTES).enumerate() { let chunk_at = dst_at + (i * STAGE_BYTES) as u64; - let fill = |stage: *mut u8| { - // SAFETY: `chunk.len() <= STAGE_BYTES`, and the privately - // owned buffer cannot overlap `chunk`. - unsafe { fill_pinned(chunk, stage) }; - }; + let fill = |pool: &FillPool, stage: &mut [MaybeUninit]| pool.copy(chunk, stage); // SAFETY: the chunks partition `src`, so `chunk_at` stays inside // the validated destination range, with `chunk.len() <= STAGE_BYTES`. unsafe { self.stage_chunk(chunk.len(), chunk_at, fill) }?; @@ -89,20 +184,15 @@ impl WeightStager { while row < plan.rows { let chunk_rows = rows_per_chunk.min(plan.rows - row); let dst_at = plan.dst_at + (row * plan.take_b) as u64; - let fill = |stage: *mut u8| { - // SAFETY: the privately owned buffer cannot overlap `src`, - // and the subslice covers `chunk_rows` full rows since - // `off_b + take_b <= stride_b`. - unsafe { - fill_pinned_strided( - &src[row * plan.stride_b..], - plan.stride_b, - plan.off_b, - plan.take_b, - chunk_rows, - stage, - ); - } + let fill = |pool: &FillPool, stage: &mut [MaybeUninit]| { + pool.gather_cols( + &src[row * plan.stride_b..], + plan.stride_b, + plan.off_b, + plan.take_b, + chunk_rows, + stage, + ); }; // SAFETY: the destination rows lie inside the validated range per // the rows x take bound, with `chunk_rows * take_b <= STAGE_BYTES`. @@ -119,7 +209,7 @@ impl WeightStager { &mut self, bytes: usize, dst_at: u64, - fill: impl FnOnce(*mut u8), + fill: impl FnOnce(&FillPool, &mut [MaybeUninit]), ) -> Result<()> { let idx = self.next; self.next = (self.next + 1) % self.bufs.len(); @@ -132,7 +222,11 @@ impl WeightStager { .as_mut_ptr() .map_err(|e| anyhow::anyhow!("staging pointer failed: {e}"))? .cast::(); - fill(stage); + // SAFETY: the pinned allocation contains STAGE_BYTES bytes and this + // function requires bytes <= STAGE_BYTES. + let staged = + unsafe { std::slice::from_raw_parts_mut(stage.cast::>(), bytes) }; + fill(&self.fill, staged); // SAFETY: `fill` initialized `bytes` at `stage` and `dst_at` is valid // per the contract; the buffer outlives the copy (`dma_done` or the // drain-or-abort branches), and the event synchronize above bound the @@ -278,82 +372,22 @@ pub(super) fn drain_or_abort(stream: &CudaStream, context: &str) { } } -/// # Safety -/// `dst` must hold `src.len()` writable bytes without overlapping `src`. -unsafe fn fill_pinned(src: &[u8], dst: *mut u8) { - if src.is_empty() { - return; - } - let per = src.len().div_ceil(FILL_THREADS); - let dst_addr = dst as usize; - std::thread::scope(|scope| { - for (i, part) in src.chunks(per).enumerate() { - scope.spawn(move || { - // SAFETY: disjoint per-thread ranges within `dst`. - unsafe { - std::ptr::copy_nonoverlapping( - part.as_ptr(), - (dst_addr as *mut u8).add(i * per), - part.len(), - ); - } - }); - } - }); -} - -/// # Safety -/// `dst` must hold `rows * take_b` writable bytes without overlapping `src`; -/// every requested source row slice must exist. -unsafe fn fill_pinned_strided( - src: &[u8], - stride_b: usize, - off_b: usize, - take_b: usize, - rows: usize, - dst: *mut u8, -) { - if rows == 0 { - return; - } - let rows_per = rows.div_ceil(FILL_THREADS); - let dst_addr = dst as usize; - std::thread::scope(|scope| { - for t in 0..FILL_THREADS.min(rows) { - let start = t * rows_per; - let end = rows.min(start + rows_per); - if start >= end { - break; - } - scope.spawn(move || { - for row in start..end { - // SAFETY: disjoint per-thread row ranges; source rows - // exist per the contract. - unsafe { - std::ptr::copy_nonoverlapping( - src.as_ptr().add(row * stride_b + off_b), - (dst_addr as *mut u8).add(row * take_b), - take_b, - ); - } - } - }); - } - }); -} - #[cfg(test)] mod tests { use super::*; #[test] - fn fill_pinned_strided_matches_scalar_gather() { + fn fill_pool_strided_matches_scalar_gather() { + // Fixed width: `FillPool::new()` would collapse to the serial path on a + // single-CPU runner, so the parallel case below would not be exercised. + let pool = FillPool::with_workers(2).expect("build fill pool"); for &(rows, total_cols, col_offset, take) in &[ (1usize, 7usize, 0usize, 7usize), (3, 8, 2, 5), (4, 5, 1, 4), (9, 6, 3, 3), - (17, 4, 0, 1), + // Past the parallel threshold, with rows that do not divide evenly. + (601, 2560, 640, 1280), ] { let (stride_b, off_b, take_b) = ( total_cols * BF16_SIZE, @@ -365,14 +399,10 @@ mod tests { *b = (i % 251) as u8; } let src = &buf[..]; - let mut dst = vec![0u8; rows * take_b]; - // SAFETY: `dst` holds `rows * take_b` bytes and does not overlap - // `src`; `src` holds `rows * stride_b` bytes, covering - // `(rows - 1) * stride_b + off_b + take_b` since - // `off_b + take_b <= stride_b`. - unsafe { - fill_pinned_strided(src, stride_b, off_b, take_b, rows, dst.as_mut_ptr()); - } + let mut dst = vec![MaybeUninit::uninit(); rows * take_b]; + pool.gather_cols(src, stride_b, off_b, take_b, rows, &mut dst); + // SAFETY: gather_cols wrote every byte of `dst`. + let dst = unsafe { std::slice::from_raw_parts(dst.as_ptr().cast::(), dst.len()) }; let mut expect = Vec::with_capacity(rows * take_b); for r in 0..rows { for c in 0..take_b {