engine-cuda: drop host Qwen3 weights and gather embeddings on device

crates/synapse-engine-cuda/Cargo.toml

@@ -17,6 +17,7 @@ anyhow = "1"
 serde = { version = "1", features = ["derive"] }
 serde_json = "1"
 safetensors = "0.6.2"
+memmap2 = "0.9"
 half = "2"
 thiserror = "1"
 sha2 = "0.10"

crates/synapse-engine-cuda/src/cuda.rs

@@ -290,6 +290,9 @@ mod enabled {
     pub struct Qwen3Context {
         binding: DeviceBinding,
         raw: NonNull<c_void>,
+        /// Layer count from the first (upload) forward; later forwards pass no
+        /// weight pointers, so the count must be remembered for the shape key.
+        layer_count: usize,
     }
 
     impl Qwen3Context {
@@ -303,12 +306,16 @@ mod enabled {
             Ok(Self {
                 binding,
                 raw: NonNull::new(raw).ok_or_else(last_error)?,
+                layer_count: 0,
             })
         }
 
+        /// First call uploads the layer weights and the embedding table, then
+        /// the caller drops its host copies. Every later call passes null
+        /// pointers for the upload payloads and the flags as 0.
         pub fn forward(
             &mut self,
-            hidden_states: &mut [f32],
+            token_ids: &[u32],
             attention_mask: &[u8],
             batch: usize,
             seq: usize,
@@ -319,30 +326,55 @@ mod enabled {
             intermediate: usize,
             epsilon: f32,
             rope_theta: f32,
-            layers: &[Qwen3Layer],
-            final_norm: &[f32],
+            layers: Option<&[Qwen3Layer]>,
+            final_norm: Option<&[f32]>,
+            embeddings: Option<&[u16]>,
+            vocab_size: usize,
+            output: &mut [f32],
         ) -> Result<()> {
             self.binding.bind()?;
-            ensure!(hidden_states.len() == batch * seq * hidden);
+            ensure!(token_ids.len() == batch * seq);
             ensure!(attention_mask.len() == batch * seq);
-            ensure!(final_norm.len() == hidden);
-            let params = layers
-                .iter()
-                .map(|layer| Qwen3LayerParams {
-                    input_norm: layer.input_norm.as_ptr(),
-                    post_attention_norm: layer.post_attention_norm.as_ptr(),
-                    q_weight: layer.q_weight.as_ptr(),
-                    q_norm: layer.q_norm.as_ptr(),
-                    k_weight: layer.k_weight.as_ptr(),
-                    k_norm: layer.k_norm.as_ptr(),
-                    v_weight: layer.v_weight.as_ptr(),
-                    o_weight: layer.o_weight.as_ptr(),
-                    gate_weight: layer.gate_weight.as_ptr(),
-                    up_weight: layer.up_weight.as_ptr(),
-                    down_weight: layer.down_weight.as_ptr(),
-                })
-                .collect::<Vec<_>>();
-            let input = crate::encode_f16_bits(hidden_states);
+            ensure!(output.len() == batch * seq * hidden);
+            let params = match layers {
+                Some(layers) => {
+                    ensure!(
+                        final_norm.is_some_and(|norm| norm.len() == hidden),
+                        "Qwen3 weight upload requires a final-norm vector of length {hidden}"
+                    );
+                    layers
+                        .iter()
+                        .map(|layer| Qwen3LayerParams {
+                            input_norm: layer.input_norm.as_ptr(),
+                            post_attention_norm: layer.post_attention_norm.as_ptr(),
+                            q_weight: layer.q_weight.as_ptr(),
+                            q_norm: layer.q_norm.as_ptr(),
+                            k_weight: layer.k_weight.as_ptr(),
+                            k_norm: layer.k_norm.as_ptr(),
+                            v_weight: layer.v_weight.as_ptr(),
+                            o_weight: layer.o_weight.as_ptr(),
+                            gate_weight: layer.gate_weight.as_ptr(),
+                            up_weight: layer.up_weight.as_ptr(),
+                            down_weight: layer.down_weight.as_ptr(),
+                        })
+                        .collect::<Vec<_>>()
+                }
+                None => Vec::new(),
+            };
+            if !params.is_empty() {
+                self.layer_count = params.len();
+            }
+            ensure!(
+                self.layer_count > 0,
+                "Qwen3 CUDA forward requires layer weights to be uploaded on the first call"
+            );
+            let layers_ptr = if params.is_empty() {
+                std::ptr::null()
+            } else {
+                params.as_ptr()
+            };
+            let final_norm_ptr = final_norm.map_or(std::ptr::null(), <[f32]>::as_ptr);
+            let embeddings_ptr = embeddings.map_or(std::ptr::null(), <[u16]>::as_ptr);
             let status = unsafe {
                 synapse_cuda_qwen3_forward(
                     self.raw.as_ptr(),
@@ -353,14 +385,18 @@ mod enabled {
                     kv_heads as u64,
                     head_dim as u64,
                     intermediate as u64,
-                    params.len() as u64,
+                    self.layer_count as u64,
                     epsilon,
                     rope_theta,
-                    input.as_ptr(),
+                    token_ids.as_ptr(),
                     attention_mask.as_ptr(),
-                    params.as_ptr(),
-                    final_norm.as_ptr(),
-                    hidden_states.as_mut_ptr(),
+                    layers_ptr,
+                    final_norm_ptr,
+                    embeddings_ptr,
+                    vocab_size as u64,
+                    i32::from(!params.is_empty()),
+                    i32::from(embeddings.is_some()),
+                    output.as_mut_ptr(),
                 )
             };
             check_status(status, "CUDA Qwen3 encoder")
@@ -489,10 +525,14 @@ mod enabled {
             layer_count: u64,
             epsilon: f32,
             rope_theta: f32,
-            input: *const u16,
+            token_ids: *const u32,
             attention_mask: *const u8,
             layers: *const Qwen3LayerParams,
             final_norm: *const f32,
+            embeddings: *const u16,
+            vocab_size: u64,
+            upload_weights: i32,
+            upload_embeddings: i32,
             output: *mut f32,
         ) -> i32;
         fn synapse_cuda_last_error() -> *const c_char;
@@ -567,7 +607,7 @@ mod enabled {
         #[allow(clippy::too_many_arguments)]
         pub fn forward(
             &mut self,
-            _hidden_states: &mut [f32],
+            _token_ids: &[u32],
             _attention_mask: &[u8],
             _batch: usize,
             _seq: usize,
@@ -578,8 +618,11 @@ mod enabled {
             _intermediate: usize,
             _epsilon: f32,
             _rope_theta: f32,
-            _layers: &[Qwen3Layer],
-            _final_norm: &[f32],
+            _layers: Option<&[Qwen3Layer]>,
+            _final_norm: Option<&[f32]>,
+            _embeddings: Option<&[u16]>,
+            _vocab_size: usize,
+            _output: &mut [f32],
         ) -> Result<()> {
             bail!("owned CUDA is unavailable in this build")
         }

crates/synapse-engine-cuda/src/lib.rs

@@ -540,9 +540,14 @@ pub fn detect_family(model_path: impl AsRef<Path>) -> Result<ModelFamily, CudaEn
 
 fn verify_digest(path: &Path, expected: &str) -> Result<(), EngineError> {
     let expected = expected.strip_prefix("sha256:").unwrap_or(expected);
-    let bytes = std::fs::read(path)
+    // Hash the file through a streaming reader: reading it whole would hold a
+    // second full copy of the model in host RAM for the duration of the hash.
+    let mut file = std::fs::File::open(path)
         .map_err(|error| OwnedCudaEmbedEngine::error(EngineErrorStage::Load, error.to_string()))?;
-    let actual = format!("{:x}", Sha256::digest(bytes));
+    let mut hasher = Sha256::new();
+    std::io::copy(&mut file, &mut hasher)
+        .map_err(|error| OwnedCudaEmbedEngine::error(EngineErrorStage::Load, error.to_string()))?;
+    let actual = format!("{:x}", hasher.finalize());
     if actual == expected {
         Ok(())
     } else {

crates/synapse-engine-cuda/src/model.rs

@@ -123,6 +123,9 @@ pub(crate) struct Qwen3Model {
     pub(crate) embeddings: Tensor,
     pub(crate) layers: Vec<Qwen3Layer>,
     pub(crate) final_norm: Vec<f32>,
+    /// Set once the CUDA context has the layer weights; the host copies are
+    /// then dropped so only the VRAM residency survives.
+    weights_uploaded: bool,
 }
 
 pub(crate) fn resolve_model_root(path: &Path) -> Result<PathBuf> {
@@ -174,6 +177,21 @@ fn load_safetensor_map(root: &Path, original: &Path) -> Result<HashMap<String, T
 }
 
 fn load_safetensors_file(path: &Path) -> Result<HashMap<String, Tensor>> {
+    // The CUDA build maps instead of reading: the file is large and every
+    // tensor is copied out into its own buffer below, so the whole-file
+    // `Vec<u8>` would only ever be a transient second copy of the model in
+    // host RAM. The non-CUDA build cannot use `unsafe` here because the crate
+    // forbids it outside the `cuda` feature, so it keeps the plain read.
+    #[cfg(feature = "cuda")]
+    let bytes = {
+        let file =
+            fs::File::open(path).with_context(|| format!("open safetensors {}", path.display()))?;
+        // SAFETY: the file is opened read-only and is not mutated or truncated
+        // while the mapping is alive (it is dropped at the end of this function).
+        unsafe { memmap2::Mmap::map(&file) }
+            .with_context(|| format!("mmap safetensors {}", path.display()))?
+    };
+    #[cfg(not(feature = "cuda"))]
     let bytes = fs::read(path).with_context(|| format!("read safetensors {}", path.display()))?;
     let tensors = SafeTensors::deserialize(&bytes)
         .map_err(|error| anyhow::anyhow!("load safetensors {}: {error}", path.display()))?;
@@ -631,35 +649,44 @@ impl Qwen3Model {
             eos_token_id: config
                 .eos_token_id
                 .context("Qwen3 config is missing eos_token_id")?,
+            weights_uploaded: false,
             embeddings,
             layers,
             final_norm,
         })
     }
 
+    /// Takes `&mut self` because the first successful forward transfers the
+    /// layer weights and the embedding table to the CUDA context and then
+    /// drops the host copies. Every later call reuses the VRAM residency.
     pub(crate) fn embed(
-        &self,
+        &mut self,
         context: &mut Qwen3Context,
         sequences: &[Vec<u32>],
     ) -> Result<Vec<Vec<f32>>> {
         let real_batch = sequences.len();
         ensure!(real_batch > 0 && sequences.iter().all(|ids| !ids.is_empty()));
         let seq = sequences.iter().map(Vec::len).max().unwrap_or(1);
-        let mut hidden = vec![0.0; real_batch * seq * self.hidden];
+        let mut token_ids = vec![0u32; real_batch * seq];
         let mut mask = vec![0u8; real_batch * seq];
         for (row, ids) in sequences.iter().enumerate() {
             for (position, &token) in ids.iter().enumerate() {
-                let token = token as usize;
-                ensure!(token < self.vocab_size);
-                let destination = (row * seq + position) * self.hidden;
-                hidden[destination..destination + self.hidden].copy_from_slice(
-                    &self.embeddings.data[token * self.hidden..(token + 1) * self.hidden],
-                );
+                ensure!((token as usize) < self.vocab_size);
+                token_ids[row * seq + position] = token;
                 mask[row * seq + position] = 1;
             }
         }
+        let mut hidden = vec![0.0f32; real_batch * seq * self.hidden];
+        let upload = !self.weights_uploaded;
+        // The CUDA worker wants the table as f16; encode once here so the
+        // upload path is a plain memcpy and the f32 table can be freed.
+        let embeddings_f16 = if upload {
+            Some(encode_f16_bits(&self.embeddings.data))
+        } else {
+            None
+        };
         context.forward(
-            &mut hidden,
+            &token_ids,
             &mask,
             real_batch,
             seq,
@@ -670,9 +697,24 @@ impl Qwen3Model {
             self.intermediate,
             self.epsilon,
             self.rope_theta,
-            &self.layers,
-            &self.final_norm,
+            upload.then_some(self.layers.as_slice()),
+            upload.then_some(self.final_norm.as_slice()),
+            embeddings_f16.as_deref(),
+            self.vocab_size,
+            &mut hidden,
         )?;
+        if upload {
+            // CUDA now owns every weight; release the ~2.2 GB of host f32.
+            self.layers = Vec::new();
+            self.layers.shrink_to_fit();
+            self.embeddings = Tensor {
+                shape: Vec::new(),
+                data: Vec::new(),
+            };
+            self.final_norm = Vec::new();
+            self.final_norm.shrink_to_fit();
+            self.weights_uploaded = true;
+        }
         let mut vectors = Vec::with_capacity(real_batch);
         for row in 0..real_batch {
             let last = (0..seq)