perf(gemma4): the routed experts load through pinned staging

pegainfer-core/src/weight_loader.rs

@@ -31,6 +31,28 @@ mod staging;
 use staging::ColShardPlan;
 use staging::WeightStager;
 
+/// Raw-byte uploads through the same pinned double buffers used by the BF16
+/// checkpoint loader. Source slices are concatenated into staging chunks.
+pub struct ByteWeightStager {
+    stream: Arc<cudarc::driver::CudaStream>,
+    stager: WeightStager,
+}
+
+impl ByteWeightStager {
+    pub fn new(ctx: &DeviceContext) -> Result<Self> {
+        Ok(Self {
+            stream: ctx.stream.clone(),
+            stager: WeightStager::new(ctx)?,
+        })
+    }
+
+    pub fn upload(&mut self, srcs: &[&[u8]], dst: &mut CudaSlice<u8>) -> Result<()> {
+        let dst_at = staging::prepare_bytes(&self.stream, srcs, dst)?;
+        // SAFETY: prepare_bytes validated the destination and concatenated size.
+        unsafe { self.stager.upload_slices_at(srcs, dst_at) }
+    }
+}
+
 /// Load shard metadata. Returns (shard_file_paths, weight_map: tensor_name -> shard_index)
 pub fn load_shard_info(model_path: &str) -> Result<(Vec<String>, HashMap<String, usize>)> {
     let single_path = format!("{}/model.safetensors", model_path);

pegainfer-core/src/weight_loader/staging.rs

@@ -175,6 +175,31 @@ impl WeightStager {
         Ok(())
     }
 
+    /// # Safety
+    /// `dst_at` must address the sum of `srcs` writable bytes still allocated
+    /// on the stager's stream, as validated by [`prepare_bytes`].
+    pub(crate) unsafe fn upload_slices_at(&mut self, srcs: &[&[u8]], dst_at: u64) -> Result<()> {
+        let mut source = 0;
+        let mut source_at = 0;
+        let mut dst_offset = 0;
+        let mut parts = Vec::new();
+        while source < srcs.len() {
+            let bytes =
+                next_contiguous_chunk(srcs, &mut source, &mut source_at, STAGE_BYTES, &mut parts);
+            let fill = |pool: &FillPool, stage: &mut [MaybeUninit<u8>]| {
+                let mut at = 0;
+                for part in &parts {
+                    pool.copy(part, &mut stage[at..at + part.len()]);
+                    at += part.len();
+                }
+            };
+            // SAFETY: the chunks partition the validated concatenated source.
+            unsafe { self.stage_chunk(bytes, dst_at + dst_offset as u64, fill) }?;
+            dst_offset += bytes;
+        }
+        Ok(())
+    }
+
     /// # Safety
     /// `plan` must come from [`prepare_cols`] for this `src`, with its
     /// destination still allocated on the stager's stream.
@@ -256,9 +281,41 @@ impl WeightStager {
     }
 }
 
+/// Fills `parts` (cleared first, reused across calls) with the source slices
+/// of the next chunk of at most `limit` bytes; returns the chunk's size.
+/// `limit` is a parameter so the boundary coverage below runs on dozens of
+/// bytes instead of a STAGE_BYTES-sized fixture.
+fn next_contiguous_chunk<'a>(
+    srcs: &'a [&'a [u8]],
+    source: &mut usize,
+    source_at: &mut usize,
+    limit: usize,
+    parts: &mut Vec<&'a [u8]>,
+) -> usize {
+    parts.clear();
+    let mut bytes = 0;
+    while *source < srcs.len() && bytes < limit {
+        let src = srcs[*source];
+        if src.is_empty() {
+            *source += 1;
+            *source_at = 0;
+            continue;
+        }
+        let take = (limit - bytes).min(src.len() - *source_at);
+        parts.push(&src[*source_at..*source_at + take]);
+        bytes += take;
+        *source_at += take;
+        if *source_at == src.len() {
+            *source += 1;
+            *source_at = 0;
+        }
+    }
+    bytes
+}
+
 // Both the stager's events and `dst`'s stream-ordered allocation are only
 // ordered against work on `stream`, which must be the stager's own stream.
-fn ensure_uploadable(stream: &Arc<CudaStream>, dst: &CudaSlice<bf16>) -> Result<()> {
+fn ensure_uploadable<T>(stream: &Arc<CudaStream>, dst: &CudaSlice<T>) -> Result<()> {
     anyhow::ensure!(
         Arc::ptr_eq(dst.stream(), stream),
         "staged upload into a buffer allocated on a different stream than the stager's"
@@ -270,6 +327,26 @@ fn ensure_uploadable(stream: &Arc<CudaStream>, dst: &CudaSlice<bf16>) -> Result<
     Ok(())
 }
 
+pub(crate) fn prepare_bytes(
+    stream: &Arc<CudaStream>,
+    srcs: &[&[u8]],
+    dst: &mut CudaSlice<u8>,
+) -> Result<u64> {
+    ensure_uploadable(stream, dst)?;
+    let bytes = srcs.iter().try_fold(0usize, |total, src| {
+        total
+            .checked_add(src.len())
+            .ok_or_else(|| anyhow::anyhow!("staged byte upload source length overflow"))
+    })?;
+    anyhow::ensure!(
+        bytes == dst.len(),
+        "staged byte upload size mismatch: {bytes} source bytes vs dst len {}",
+        dst.len()
+    );
+    let (dst_ptr, _dst_order) = dst.device_ptr_mut(stream);
+    Ok(dst_ptr)
+}
+
 /// Validates a contiguous staged upload and returns the destination device
 /// address for a deferred [`WeightStager::upload_at`].
 pub(crate) fn prepare(
@@ -415,4 +492,26 @@ mod tests {
             );
         }
     }
+
+    #[test]
+    fn contiguous_chunks_preserve_source_bytes() {
+        let limit = 46;
+        let first = vec![1u8; limit - 2];
+        let second = [2u8, 3, 4, 5];
+        let sources: [&[u8]; 3] = [&first, &[], &second];
+        let mut source = 0;
+        let mut source_at = 0;
+        let mut actual = Vec::new();
+        let mut sizes = Vec::new();
+        let mut parts = Vec::new();
+        while source < sources.len() {
+            let bytes =
+                next_contiguous_chunk(&sources, &mut source, &mut source_at, limit, &mut parts);
+            sizes.push(bytes);
+            actual.extend(parts.iter().copied().flatten().copied());
+        }
+        let expected: Vec<u8> = sources.into_iter().flatten().copied().collect();
+        assert_eq!(sizes, [limit, 2]);
+        assert_eq!(actual, expected);
+    }
 }

pegainfer-gemma4/src/weights/load.rs

@@ -5,6 +5,7 @@ use std::time::Instant;
 use anyhow::Result;
 use log::info;
 use pegainfer_core::tensor::DeviceContext;
+use pegainfer_core::weight_loader::ByteWeightStager;
 use pegainfer_core::weight_loader::SlotId;
 use pegainfer_core::weight_loader::StagedWeightLoader;
 use pegainfer_core::weight_loader::VecSlotId;
@@ -70,8 +71,7 @@ struct LayerSlots {
     gate: SlotId,
     up: SlotId,
     down: SlotId,
-    /// The bf16 half of a routed layer. The experts do not travel through the
-    /// staged loader at all.
+    /// The bf16 half of a routed layer.
     moe: Option<MoeSlots>,
 }
 
@@ -164,6 +164,12 @@ fn upload_experts(
     shards: &[SafeTensors],
     manifest: &Manifest,
 ) -> Result<Vec<Option<StackedExperts>>> {
+    // A dense checkpoint (12B, 31B) has no routed layer; skip the stager's
+    // pinned buffers and thread pool instead of allocating them for nothing.
+    if manifest.layers.iter().all(|layer| layer.moe.is_none()) {
+        return Ok(manifest.layers.iter().map(|_| None).collect());
+    }
+    let mut stager = ByteWeightStager::new(ctx)?;
     manifest
         .layers
         .iter()
@@ -173,9 +179,9 @@ fn upload_experts(
                 .as_ref()
                 .map(|moe| {
                     Ok(StackedExperts {
-                        gate: upload_stacked(ctx, shards, &moe.experts, |e| &e.gate)?,
-                        up: upload_stacked(ctx, shards, &moe.experts, |e| &e.up)?,
-                        down: upload_stacked(ctx, shards, &moe.experts, |e| &e.down)?,
+                        gate: upload_stacked(ctx, &mut stager, shards, &moe.experts, |e| &e.gate)?,
+                        up: upload_stacked(ctx, &mut stager, shards, &moe.experts, |e| &e.up)?,
+                        down: upload_stacked(ctx, &mut stager, shards, &moe.experts, |e| &e.down)?,
                     })
                 })
                 .transpose()
@@ -186,12 +192,11 @@ fn upload_experts(
 /// Stack one projection of every expert into a pair of device buffers and
 /// upload it as the checkpoint stores it.
 ///
-/// This bypasses the staged loader on purpose: that path is bf16-typed, and
-/// widening here is exactly what this representation exists to avoid. Each
-/// expert lands in its own row range, so the buffer is already the shape a
-/// batched call wants.
+/// Each expert lands in its own row range, so the buffer is already the shape
+/// a batched call wants.
 fn upload_stacked(
     ctx: &DeviceContext,
+    stager: &mut ByteWeightStager,
     shards: &[SafeTensors],
     experts: &[ExpertTensors],
     pick: fn(&ExpertTensors) -> &QuantMatrix,
@@ -211,8 +216,9 @@ fn upload_stacked(
         .stream
         .alloc_zeros::<u8>(scales_per_expert * experts.len())
         .map_err(|e| anyhow::anyhow!("Gemma 4: cannot hold the stacked block scales: {e}"))?;
+    let mut sources = Vec::with_capacity(experts.len());
     let mut tensor_scales = Vec::with_capacity(experts.len());
-    let mut scale_bytes = Vec::with_capacity(scales_per_expert * experts.len());
+    let mut scale_peak = 0.0f32;
 
     for (index, expert) in experts.iter().enumerate() {
         let plan = pick(expert);
@@ -232,29 +238,27 @@ fn upload_stacked(
             source.packed().len(),
             source.scales().len()
         );
-        let at = index * packed_per_expert;
-        ctx.stream
-            .memcpy_htod(
-                source.packed(),
-                &mut packed.slice_mut(at..at + packed_per_expert),
-            )
-            .map_err(|e| anyhow::anyhow!("Gemma 4: expert {index} weights did not upload: {e}"))?;
-        let at = index * scales_per_expert;
-        ctx.stream
-            .memcpy_htod(
-                source.scales(),
-                &mut scales.slice_mut(at..at + scales_per_expert),
-            )
-            .map_err(|e| anyhow::anyhow!("Gemma 4: expert {index} scales did not upload: {e}"))?;
-        scale_bytes.extend_from_slice(source.scales());
+        scale_peak = source.scales().iter().fold(scale_peak, |peak, byte| {
+            peak.max(crate::nvfp4::decode_e4m3(*byte) * 128.0)
+        });
         tensor_scales.push(source.tensor_scale());
+        sources.push(source);
     }
 
+    let packed_sources: Vec<&[u8]> = sources.iter().map(QuantSource::packed).collect();
+    stager
+        .upload(&packed_sources, &mut packed)
+        .map_err(|e| anyhow::anyhow!("Gemma 4: expert weights did not upload: {e}"))?;
+    let scale_sources: Vec<&[u8]> = sources.iter().map(QuantSource::scales).collect();
+    stager
+        .upload(&scale_sources, &mut scales)
+        .map_err(|e| anyhow::anyhow!("Gemma 4: expert scales did not upload: {e}"))?;
+
     // Marlin reads the block scale as S0E5M3, so every scale is normalized by
     // one shared power of two and the per-tensor scale takes it back. The
     // factor has to be the same across a projection's experts, which is why it
     // is found here rather than per expert.
-    let rescale = marlin_rescale(&scale_bytes);
+    let rescale = marlin_rescale(scale_peak);
     let mut qweight = ctx
         .stream
         .alloc_zeros::<u8>(packed_per_expert * experts.len())
@@ -304,12 +308,8 @@ fn upload_stacked(
 /// The shared power of two that lifts every block scale so its leading bit
 /// survives the S0E5M3 re-encoding. Mirrors vLLM's
 /// `_nvfp4_compute_scale_factor`, whose bound is the e4m3 maximum.
-fn marlin_rescale(scale_bytes: &[u8]) -> f32 {
+fn marlin_rescale(peak: f32) -> f32 {
     const CEILING: f32 = 448.0 * 128.0;
-    let peak = scale_bytes
-        .iter()
-        .map(|byte| crate::nvfp4::decode_e4m3(*byte) * 128.0)
-        .fold(0.0f32, f32::max);
     if peak <= 0.0 || peak >= CEILING {
         return 1.0;
     }