From afb106370a6b4b1be313fd44d7d65dce16f91b66 Mon Sep 17 00:00:00 2001 From: Qiiks Date: Wed, 16 Sep 2026 18:32:58 +0530 Subject: [PATCH 1/3] engine-cuda: drop host Qwen3 weights and gather embeddings on device The owned-CUDA Qwen3 embedder held ~2.2 GB of host RSS for the whole process lifetime even though CUDA owns the weights after the first forward: the safetensors were slurped whole, every f16/bf16 tensor was upcast to Vec, and both the layer weights and the embedding table stayed alive only to be re-uploaded never. Three orthogonal changes, none of which alter numerics: 1. Upload once, then drop. Qwen3Context::forward now takes optional layer/final-norm/embedding payloads plus upload_weights / upload_embeddings flags; Qwen3Model::embed passes them on the first call, then empties self.layers / self.embeddings / self.final_norm and flips weights_uploaded. Later calls pass null and 0. 2. Gather embeddings on the device. The embedding table uploads as f16 (memcpy straight through) and a new embed_gather kernel copies one row per position into the hidden-state buffer, replacing the host-side f32 table lookup that pinned the 0.58 GiB table in RAM. 3. mmap the safetensors and stream the digest. load_safetensors_file maps the file instead of reading it whole (tensors are copied out anyway), and verify_digest hashes through std::io::copy instead of holding a second full copy of the model. The token ids now reach the kernel directly, so encode_f16_bits is no longer needed on the Qwen3 path. Measured on the real worker against the same staged package (qwen3-embedding-0.6b, f16 safetensors, RTX 4050 Laptop): before WS 2,045.9 MB private 3,894.8 MB after WS 158.9 MB private 246.9 MB Both binaries returned 1024-dimensional vectors; the ladders, pooling and normalisation are unchanged. --- Cargo.lock | 1 + crates/synapse-engine-cuda/Cargo.toml | 1 + crates/synapse-engine-cuda/src/cuda.rs | 92 +++++++++++++------ crates/synapse-engine-cuda/src/lib.rs | 9 +- crates/synapse-engine-cuda/src/model.rs | 58 +++++++++--- .../src/port/cuda_qwen3.cu | 70 +++++++++++--- 6 files changed, 173 insertions(+), 58 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 43b99e95..6433d67d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6784,6 +6784,7 @@ dependencies = [ "anyhow", "cc", "half", + "memmap2", "safetensors 0.6.2", "serde", "serde_json", diff --git a/crates/synapse-engine-cuda/Cargo.toml b/crates/synapse-engine-cuda/Cargo.toml index 1e08280b..aded0d65 100644 --- a/crates/synapse-engine-cuda/Cargo.toml +++ b/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" diff --git a/crates/synapse-engine-cuda/src/cuda.rs b/crates/synapse-engine-cuda/src/cuda.rs index 9f4b6736..f7ac9105 100644 --- a/crates/synapse-engine-cuda/src/cuda.rs +++ b/crates/synapse-engine-cuda/src/cuda.rs @@ -306,9 +306,12 @@ mod enabled { }) } + /// 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 +322,48 @@ 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::>(); - 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::>() + } + None => Vec::new(), + }; + 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(), @@ -356,11 +377,15 @@ mod enabled { params.len() 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 +514,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 +596,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 +607,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") } diff --git a/crates/synapse-engine-cuda/src/lib.rs b/crates/synapse-engine-cuda/src/lib.rs index 08259415..6de832d3 100644 --- a/crates/synapse-engine-cuda/src/lib.rs +++ b/crates/synapse-engine-cuda/src/lib.rs @@ -540,9 +540,14 @@ pub fn detect_family(model_path: impl AsRef) -> Result 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 { diff --git a/crates/synapse-engine-cuda/src/model.rs b/crates/synapse-engine-cuda/src/model.rs index 02a81a90..a3db729a 100644 --- a/crates/synapse-engine-cuda/src/model.rs +++ b/crates/synapse-engine-cuda/src/model.rs @@ -123,6 +123,9 @@ pub(crate) struct Qwen3Model { pub(crate) embeddings: Tensor, pub(crate) layers: Vec, pub(crate) final_norm: Vec, + /// 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 { @@ -174,7 +177,14 @@ fn load_safetensor_map(root: &Path, original: &Path) -> Result Result> { - let bytes = fs::read(path).with_context(|| format!("read safetensors {}", path.display()))?; + // Map instead of reading: the file is large and every tensor is copied out + // into its own buffer below, so the whole-file `Vec` would only ever be + // a transient second copy of the model in host RAM. + 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). + let bytes = unsafe { memmap2::Mmap::map(&file) } + .with_context(|| format!("mmap safetensors {}", path.display()))?; let tensors = SafeTensors::deserialize(&bytes) .map_err(|error| anyhow::anyhow!("load safetensors {}: {error}", path.display()))?; let mut result = HashMap::new(); @@ -631,35 +641,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], ) -> Result>> { 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 +689,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) diff --git a/crates/synapse-engine-cuda/src/port/cuda_qwen3.cu b/crates/synapse-engine-cuda/src/port/cuda_qwen3.cu index 76412d8b..02cd7c9d 100644 --- a/crates/synapse-engine-cuda/src/port/cuda_qwen3.cu +++ b/crates/synapse-engine-cuda/src/port/cuda_qwen3.cu @@ -194,6 +194,20 @@ __global__ void to_float(const half *input, float *output, int count) { if (index < count) output[index] = __half2float(input[index]); } +// Copies one embedding row per padded sequence position out of the device +// table. Masked positions still read their token id (which the caller +// zero-pads) but never contribute to attention because causal_softmax treats +// them as -10000. +__global__ void embed_gather(const uint32_t *token_ids, const half *table, half *output, int rows, int width) { + int row = blockIdx.x; + if (row >= rows) return; + const half *source = table + static_cast(token_ids[row]) * width; + half *target = output + static_cast(row) * width; + for (int column = threadIdx.x; column < width; column += blockDim.x) { + target[column] = source[column]; + } +} + struct QwenContext; struct ShapePlan { @@ -203,6 +217,7 @@ struct ShapePlan { size_t arena_bytes = 0; DeviceAllocation arena, workspace; DeviceAllocation mask; + DeviceAllocation token_ids; DeviceAllocation cosine, sine, output; half *x0 = nullptr, *x1 = nullptr, *normed = nullptr; half *q_raw = nullptr, *k_raw = nullptr, *v_raw = nullptr; @@ -216,8 +231,8 @@ struct ShapePlan { ShapePlan(QwenContext *owner, int b, int s, int h, int qh_count, int kvh, int hd, int inter, int layers, float eps, float theta); ~ShapePlan(); void compute(StageProfile *profile = nullptr); - void initialize_and_verify(const uint16_t *input, const uint8_t *host_mask); - void run(const uint16_t *input, const uint8_t *host_mask, float *host_output); + void initialize_and_verify(const uint32_t *host_ids, const uint8_t *host_mask); + void run(const uint32_t *host_ids, const uint8_t *host_mask, float *host_output); }; struct QwenContext { @@ -228,6 +243,8 @@ struct QwenContext { int hidden = 0, query_heads = 0, kv_heads = 0, head_dim = 0, intermediate = 0, layer_count = 0; std::vector layers; DeviceAllocation final_norm; + DeviceAllocation embeddings; + bool embeddings_loaded = false; std::unordered_map> plans; explicit QwenContext(bool graphs) : graphs_enabled(graphs) { @@ -246,6 +263,7 @@ struct QwenContext { if (hidden != h || query_heads != qh_count || kv_heads != kvh || head_dim != hd || intermediate != inter || layer_count != count) throw std::runtime_error("Qwen3 CUDA model dimensions changed"); return; } + if (!params || !host_final_norm) throw std::runtime_error("Qwen3 CUDA load_weights received null layer pointers"); hidden = h; query_heads = qh_count; kv_heads = kvh; head_dim = hd; intermediate = inter; layer_count = count; int q_width = qh_count * hd; int kv_width = kvh * hd; @@ -270,6 +288,18 @@ struct QwenContext { weights_loaded = true; std::fprintf(stderr, "CUDA Qwen3 persistent weights: layers=%d dtype=f16 accum=fp32 norm_params=fp32\n", count); } + + void load_embeddings(const uint16_t *host_embeddings, int vocab, int width) { + if (embeddings_loaded) { + if (vocab != static_cast(embeddings.count / static_cast(width)) || width != hidden) throw std::runtime_error("Qwen3 CUDA embedding table dimensions changed"); + return; + } + size_t total = static_cast(vocab) * width; + embeddings.allocate(total); + FAMILY_CUDA_CHECK(cudaMemcpy(embeddings.pointer, host_embeddings, total * sizeof(half), cudaMemcpyHostToDevice)); + embeddings_loaded = true; + std::fprintf(stderr, "CUDA Qwen3 persistent embeddings: vocab=%d hidden=%d bytes=%zu\n", vocab, width, total * sizeof(half)); + } }; ShapePlan::ShapePlan(QwenContext *owner, int b, int s, int h, int qh_count, int kvh, int hd, int inter, int layers_count, float eps, float theta) @@ -284,6 +314,7 @@ ShapePlan::ShapePlan(QwenContext *owner, int b, int s, int h, int qh_count, int arena_bytes = total * sizeof(half) + 20 * 256; arena.allocate(arena_bytes); mask.allocate(rows); + token_ids.allocate(rows); output.allocate(hidden_values); unsigned char *cursor = arena.pointer; auto take = [&](size_t count) { @@ -343,6 +374,9 @@ void ShapePlan::compute(StageProfile *profile) { size_t score_group_values = static_cast(batch) * kv_heads * seq * seq; auto begin = [&](const char *name) { if (profile) profile->begin(name, context->stream); }; auto end = [&] { if (profile) profile->end(context->stream); }; + begin("pointwise_layout"); + embed_gather<<stream>>>(token_ids.pointer, context->embeddings.pointer, x0, rows, hidden); + end(); for (int index = 0; index < layer_count; ++index) { DeviceLayer &layer = context->layers[index]; begin("pointwise_layout"); @@ -402,10 +436,10 @@ void ShapePlan::compute(StageProfile *profile) { FAMILY_CUDA_CHECK(cudaGetLastError()); } -void ShapePlan::initialize_and_verify(const uint16_t *input, const uint8_t *host_mask) { - size_t input_bytes = static_cast(batch) * seq * hidden * sizeof(half); +void ShapePlan::initialize_and_verify(const uint32_t *host_ids, const uint8_t *host_mask) { + size_t ids_bytes = static_cast(batch) * seq * sizeof(uint32_t); size_t mask_bytes = static_cast(batch) * seq; - FAMILY_CUDA_CHECK(cudaMemcpyAsync(x0, input, input_bytes, cudaMemcpyHostToDevice, context->stream)); + FAMILY_CUDA_CHECK(cudaMemcpyAsync(token_ids.pointer, host_ids, ids_bytes, cudaMemcpyHostToDevice, context->stream)); FAMILY_CUDA_CHECK(cudaMemcpyAsync(mask.pointer, host_mask, mask_bytes, cudaMemcpyHostToDevice, context->stream)); StageProfile profile; compute(&profile); @@ -417,7 +451,7 @@ void ShapePlan::initialize_and_verify(const uint16_t *input, const uint8_t *host compute(); FAMILY_CUDA_CHECK(cudaStreamEndCapture(context->stream, &graph)); FAMILY_CUDA_CHECK(cudaGraphInstantiate(&graph_exec, graph, nullptr, nullptr, 0)); - FAMILY_CUDA_CHECK(cudaMemcpyAsync(x0, input, input_bytes, cudaMemcpyHostToDevice, context->stream)); + FAMILY_CUDA_CHECK(cudaMemcpyAsync(token_ids.pointer, host_ids, ids_bytes, cudaMemcpyHostToDevice, context->stream)); FAMILY_CUDA_CHECK(cudaMemcpyAsync(mask.pointer, host_mask, mask_bytes, cudaMemcpyHostToDevice, context->stream)); FAMILY_CUDA_CHECK(cudaGraphLaunch(graph_exec, context->stream)); FAMILY_CUDA_CHECK(cudaStreamSynchronize(context->stream)); @@ -427,10 +461,10 @@ void ShapePlan::initialize_and_verify(const uint16_t *input, const uint8_t *host std::fprintf(stderr, "CUDA Qwen3 shape %dx%d: arena=%zu workspace=%zu captured_exact=true launches=%d gqa=two-group-strided kv_repeat_bytes=0 stage_projection_mlp_gemm=%.3fms stage_attention_gemm=%.3fms stage_score_softmax=%.3fms stage_pointwise_layout=%.3fms stage_final_norm_output=%.3fms\n", batch, seq, arena_bytes, workspace.count, layer_count * (15 + 2 * (query_heads / kv_heads)) + 2, stage_ms["projection_mlp_gemm"], stage_ms["attention_gemm"], stage_ms["score_softmax"], stage_ms["pointwise_layout"], stage_ms["final_norm_output"]); } -void ShapePlan::run(const uint16_t *input, const uint8_t *host_mask, float *host_output) { - size_t input_bytes = static_cast(batch) * seq * hidden * sizeof(half); +void ShapePlan::run(const uint32_t *host_ids, const uint8_t *host_mask, float *host_output) { + size_t ids_bytes = static_cast(batch) * seq * sizeof(uint32_t); size_t mask_bytes = static_cast(batch) * seq; - FAMILY_CUDA_CHECK(cudaMemcpyAsync(x0, input, input_bytes, cudaMemcpyHostToDevice, context->stream)); + FAMILY_CUDA_CHECK(cudaMemcpyAsync(token_ids.pointer, host_ids, ids_bytes, cudaMemcpyHostToDevice, context->stream)); FAMILY_CUDA_CHECK(cudaMemcpyAsync(mask.pointer, host_mask, mask_bytes, cudaMemcpyHostToDevice, context->stream)); if (context->graphs_enabled) FAMILY_CUDA_CHECK(cudaGraphLaunch(graph_exec, context->stream)); else compute(); @@ -467,25 +501,33 @@ int32_t synapse_cuda_qwen3_forward( uint64_t layer_count, float epsilon, float rope_theta, - const uint16_t *input, + const uint32_t *token_ids, const uint8_t *attention_mask, const Qwen3LayerParams *layers, const float *final_norm, + const uint16_t *embeddings, + uint64_t vocab_size, + int32_t upload_weights, + int32_t upload_embeddings, float *output ) { try { - if (!raw_context || !input || !attention_mask || !layers || !final_norm || !output) throw std::runtime_error("Qwen3 CUDA received a null pointer"); + if (!raw_context || !token_ids || !attention_mask || !output) throw std::runtime_error("Qwen3 CUDA received a null pointer"); + if (upload_weights && (!layers || !final_norm)) throw std::runtime_error("Qwen3 CUDA weight upload requires layer and final-norm pointers"); + if (upload_embeddings && !embeddings) throw std::runtime_error("Qwen3 CUDA embedding upload requires a host table pointer"); if (!batch || !seq || !hidden || !query_heads || !kv_heads || query_heads % kv_heads || !head_dim || !layer_count) throw std::runtime_error("Qwen3 CUDA received invalid dimensions"); QwenContext *context = static_cast(raw_context); - context->load_weights(hidden, query_heads, kv_heads, head_dim, intermediate, layer_count, layers, final_norm); + if (upload_weights) context->load_weights(hidden, query_heads, kv_heads, head_dim, intermediate, layer_count, layers, final_norm); + if (upload_embeddings) context->load_embeddings(embeddings, static_cast(vocab_size), static_cast(hidden)); + if (!context->weights_loaded || !context->embeddings_loaded) throw std::runtime_error("Qwen3 CUDA forward called before weights and embeddings were uploaded"); std::string key = shape_key(batch, seq); auto found = context->plans.find(key); if (found == context->plans.end()) { auto plan = std::make_unique(context, batch, seq, hidden, query_heads, kv_heads, head_dim, intermediate, layer_count, epsilon, rope_theta); - plan->initialize_and_verify(input, attention_mask); + plan->initialize_and_verify(token_ids, attention_mask); found = context->plans.emplace(key, std::move(plan)).first; } - found->second->run(input, attention_mask, output); + found->second->run(token_ids, attention_mask, output); return 0; } catch (const std::exception &error) { synapse_cuda_set_last_error(error.what()); From f28da35bf654c20abf8de0520e7b6ad036b23da8 Mon Sep 17 00:00:00 2001 From: Qiiks Date: Wed, 16 Sep 2026 18:37:36 +0530 Subject: [PATCH 2/3] engine-cuda: gate the safetensors mmap behind the cuda feature The crate root forbids unsafe_code outside `all(feature = "cuda", not(macos))`, so an unconditional `unsafe { memmap2::Mmap::map }` broke the non-cuda build (linux CI clippy: `forbid(unsafe_code)` at lib.rs:3). Branch on the feature: the CUDA build maps the safetensors and keeps the streaming digest, the non-CUDA build falls back to `fs::read` and never reaches for `unsafe`. Also applies rustfmt's line wrap in the same function. Verified: `cargo clippy -p synapse-engine-cuda --all-targets` (default, no `cuda`) clean; `cargo build -p synapse-worker-cuda --no-default-features --features cuda --release` EXITCODE=0; `cargo fmt --all --check` clean. --- crates/synapse-engine-cuda/src/model.rs | 24 ++++++++++++++++-------- 1 file changed, 16 insertions(+), 8 deletions(-) diff --git a/crates/synapse-engine-cuda/src/model.rs b/crates/synapse-engine-cuda/src/model.rs index a3db729a..1f7b2b64 100644 --- a/crates/synapse-engine-cuda/src/model.rs +++ b/crates/synapse-engine-cuda/src/model.rs @@ -177,14 +177,22 @@ fn load_safetensor_map(root: &Path, original: &Path) -> Result Result> { - // Map instead of reading: the file is large and every tensor is copied out - // into its own buffer below, so the whole-file `Vec` would only ever be - // a transient second copy of the model in host RAM. - 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). - let bytes = unsafe { memmap2::Mmap::map(&file) } - .with_context(|| format!("mmap safetensors {}", path.display()))?; + // 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` 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()))?; let mut result = HashMap::new(); From ad3c62794fce851b181faccb4b777bbac82872fd Mon Sep 17 00:00:00 2001 From: Qiiks Date: Wed, 16 Sep 2026 18:58:17 +0530 Subject: [PATCH 3/3] engine-cuda: remember the uploaded layer count across forwards MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The crash case: the first forward uploads the layer weights, drops the host copies, and flips weights_uploaded. Every later forward passes layers=None, so `params` is empty and `params.len()` (the value handed to the FFI as layer_count) becomes 0 — and cuda_qwen3.cu:518 rejects layer_count == 0 with "Qwen3 CUDA received invalid dimensions". The first forward worked; the second forward after the host weights were dropped died with engine_crashed. Fix: Qwen3Context now stores the layer count it uploaded the first time and passes that stored value on every subsequent forward. The guard makes the "forward before upload" case an explicit error instead of a size-0 crash. Verified live: two consecutive /v1/embeddings calls through the gateway both returned 1024-dim vectors, and the worker's working set stayed at 384 MB (vs 2655 MB before the RAM-slim change). --- crates/synapse-engine-cuda/src/cuda.rs | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/crates/synapse-engine-cuda/src/cuda.rs b/crates/synapse-engine-cuda/src/cuda.rs index f7ac9105..7ba10341 100644 --- a/crates/synapse-engine-cuda/src/cuda.rs +++ b/crates/synapse-engine-cuda/src/cuda.rs @@ -290,6 +290,9 @@ mod enabled { pub struct Qwen3Context { binding: DeviceBinding, raw: NonNull, + /// 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,6 +306,7 @@ mod enabled { Ok(Self { binding, raw: NonNull::new(raw).ok_or_else(last_error)?, + layer_count: 0, }) } @@ -357,6 +361,13 @@ mod enabled { } 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 { @@ -374,7 +385,7 @@ 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, token_ids.as_ptr(),