perf(core): reuse a persistent weight-fill pool and tune staging geometry

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).

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 }

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::<bf16>();
-/// 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<usize> = 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<u8>] {
+    // SAFETY: MaybeUninit<u8> 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> {
+        Self::with_workers(fill_threads())
+    }
+
+    fn with_workers(workers: usize) -> Result<Self> {
+        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<u8>]) {
+        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<u8>],
+    ) {
+        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<bf16>,
@@ -39,6 +136,7 @@ pub(crate) struct ColShardPlan {
 pub(crate) struct WeightStager {
     stream: Arc<CudaStream>,
     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<u8>]| 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<u8>]| {
+                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<u8>]),
     ) -> 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::<u8>();
-        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::<MaybeUninit<u8>>(), 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::<u8>(), dst.len()) };
             let mut expect = Vec::with_capacity(rows * take_b);
             for r in 0..rows {
                 for c in 0..take_b {