diff --git a/pegainfer-gemma4/src/engine.rs b/pegainfer-gemma4/src/engine.rs index 79b13a3b..a66a2b01 100644 --- a/pegainfer-gemma4/src/engine.rs +++ b/pegainfer-gemma4/src/engine.rs @@ -1176,8 +1176,10 @@ impl EngineState { sampler_graphs.push(graph); bucket *= 2; } - ctx.sync()?; } + // The KV pools, arena and warm passes were enqueued on this stream. The lane + // stream and everything after it must see them complete, graphs or not. + ctx.sync()?; let lane = lane_mode .map(|mode| AsyncPrefillLane::new(&ctx, mode)) .transpose()?; diff --git a/pegainfer-gemma4/src/weights/load.rs b/pegainfer-gemma4/src/weights/load.rs index 28bdcc57..8287fb6d 100644 --- a/pegainfer-gemma4/src/weights/load.rs +++ b/pegainfer-gemma4/src/weights/load.rs @@ -35,8 +35,8 @@ use crate::nvfp4::QuantSource; /// redemption, prefetch join and unmap fall between them, and the allocations /// submitted under `record_api_wall_ms` execute under /// `execute_and_drain_wall_ms`. `elapsed_ms` is the submission total: it -/// samples when the call returns, after the A4B expert kernels are enqueued -/// but without draining them. +/// samples when the call returns, after the A4B expert kernels have drained +/// on the loader stream; the unmap's host cost overlaps that drain. struct LoadStats { /// Every required tensor at its dtype. manifest_bytes: usize, @@ -519,6 +519,11 @@ impl Gemma4Weights { // once an executor wants it too. drop(mmaps); + // The expert kernels ran on this stream while the unmap paid its host cost. + // Every later weight consumer uses another stream, so this drain is the handoff. + ctx.sync() + .map_err(|e| anyhow::anyhow!("Gemma 4: cannot drain the expert kernels: {e}"))?; + let stats = LoadStats { manifest_bytes, device_bytes: free_before as i64 - device_free_bytes as i64, @@ -533,7 +538,7 @@ impl Gemma4Weights { "Gemma 4 weights resident: {:.2} GiB manifest, {:.2} GiB device, {:.2} GiB free, \ {} modality tensors skipped; \ {:.0} ms submission total, of which {:.0} validate, {:.0} record-api, \ - {:.0} execute-and-drain (expert kernels enqueued, not drained)", + {:.0} execute-and-drain (expert kernels drained)", gib(stats.manifest_bytes as i64), gib(stats.device_bytes), gib(stats.device_free_bytes as i64),