diff --git a/pegainfer-core/src/weight_loader.rs b/pegainfer-core/src/weight_loader.rs index 6c6a263c3..a6849153d 100644 --- a/pegainfer-core/src/weight_loader.rs +++ b/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, + stager: WeightStager, +} + +impl ByteWeightStager { + pub fn new(ctx: &DeviceContext) -> Result { + Ok(Self { + stream: ctx.stream.clone(), + stager: WeightStager::new(ctx)?, + }) + } + + pub fn upload(&mut self, srcs: &[&[u8]], dst: &mut CudaSlice) -> 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, HashMap)> { let single_path = format!("{}/model.safetensors", model_path); diff --git a/pegainfer-core/src/weight_loader/staging.rs b/pegainfer-core/src/weight_loader/staging.rs index a4250c455..341080e1a 100644 --- a/pegainfer-core/src/weight_loader/staging.rs +++ b/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]| { + 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, dst: &CudaSlice) -> Result<()> { +fn ensure_uploadable(stream: &Arc, dst: &CudaSlice) -> 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, dst: &CudaSlice) -> Result< Ok(()) } +pub(crate) fn prepare_bytes( + stream: &Arc, + srcs: &[&[u8]], + dst: &mut CudaSlice, +) -> Result { + 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 = sources.into_iter().flatten().copied().collect(); + assert_eq!(sizes, [limit, 2]); + assert_eq!(actual, expected); + } } diff --git a/pegainfer-gemma4/src/weights/load.rs b/pegainfer-gemma4/src/weights/load.rs index 626c56189..c01df2725 100644 --- a/pegainfer-gemma4/src/weights/load.rs +++ b/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, } @@ -164,6 +164,12 @@ fn upload_experts( shards: &[SafeTensors], manifest: &Manifest, ) -> Result>> { + // 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::(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::(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; }