engine-cuda: drop host Qwen3 weights and gather embeddings on device - #15
Conversation
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<f32>, 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.
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.
|
CI coverage note
The |
![cubic-dev-ai[bot]](https://avatars.githubusercontent.com/in/1082092?s=60&v=4){kind=small}
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).
|
Follow-up commit ad3c627 — the two-forward crash case. Testing the RAM-slim change end-to-end on the real RTX 4050 worker surfaced a defect: the first forward (the upload path) returned 1024-dim vectors correctly, but the second forward after host weights were dropped crashed the engine with: Root cause: Fix: Verified live with two consecutive
|
CI coverage on this branchThe Two things CI does not cover here, verified locally instead:
The default-feature path ( |
|
Verified the earlier head That run sent a single Nothing needed from you unless the head moves again. |
Summary
The owned-CUDA Qwen3 embedding worker held ~2.2 GB of host RAM for the entire process lifetime even though CUDA owns every weight after the first forward. The safetensors file was read whole into a
Vec<u8>, every f16/bf16 tensor was expanded toVec<f32>, and both the layer weights and the embedding table stayed alive only to be re-uploaded never.Three orthogonal changes, no numerical change to any output.
Measured
Real worker, real staged package (
qwen3-embedding-0.6b, f16 safetensors), RTX 4050 Laptop, both binaries returning 1024-dim vectors:~2 GB reclaimed, ~24× drop in private bytes. Quality is unchanged: the existing
synapse-worker-cudasmoke path (load +embed_batch) returns the same 1024-dimensional normalised vectors on both binaries.What changed
1. Upload once, then drop the host copy
Qwen3Context::forwardnow takeslayers: Option<&[Qwen3Layer]>,final_norm: Option<&[f32]>,embeddings: Option<&[u16]>plusvocab_size,upload_weights,upload_embeddings. On the first callQwen3Model::embedpasses the host tensors and setsupload_weights: 1; the C++ side uploads to the device as before. Thenmodel.rsemptiesself.layers,self.final_norm,self.embeddings, callsshrink_to_fit, and flips aweights_uploadedflag. Later calls pass null pointers and 0 flags.Qwen3Context::load_weights(incuda_qwen3.cu) previously threw on a second call whenlayer_count != count; it now early-returns unconditionally onceweights_loaded, because dims cannot legitimately change for a given context.2. Gather embeddings on the device
The embedding table now uploads as f16 with a plain memcpy (no f16→f32→f16 round trip). A new
embed_gatherkernel copies one row per position into the hidden-state buffer, replacing the host-side f32 gather that pinned the 0.58 GiB table in RAM. Token ids reach the kernel directly asuint32_t*, soencode_f16_bitsis no longer used on the Qwen3 path.The host
Tensortype and the loader are otherwise untouched — only the Qwen3 code path loses its f32 mirror.3. mmap the safetensors, stream the digest
load_safetensors_file(inmodel.rs) usesmemmap2::Mmap::mapinstead offs::read. Tensor bytes are copied out into per-tensor buffers below, so the whole-fileVec<u8>was only ever a transient second copy of the model.verify_digest(inlib.rs) hashes throughstd::io::copyinstead ofSha256::digest(bytes), so the file is streamed rather than slurped.memmap2 = "0.9"added tocrates/synapse-engine-cuda/Cargo.toml.Load order preserved
The upload still happens in the original order: layers, then final-norm, then embeddings, then
cudaDeviceSynchronize, thenweights_loaded = true.Numerics
Unchanged. Host f32 → CUDA f16 already happened through
copy_weight<half>on the upload path; storing f16 on the host for the embedding table just skips a lossless f16→f32→f16 round trip. Norm vectors are still copied as fp32 (copy_float), and the kernel writes only the hidden-state buffer the forward pass already used.Testing
cargo build -p synapse-worker-cuda --no-default-features --features cuda --release— clean.cargo check -p synapse-engine-cuda(default, nocuda) — clean; the non-cuda stub was updated to match the new signature so the Linux CI lane compiles./health, and returned 1024-dim vectors on a real embed.Follow-up
The token id buffer allocated per forward (
ShapePlan::token_ids) is new but small (rows × 4 bytes) and reuses the existingDeviceAllocationpattern; if the plan cache turns out to be hot we can hoist it, but it was not worth changing the plan shape in this PR.Summary by cubic
The owned-CUDA Qwen3 embedding worker now drops host copies of model weights after uploading them to the device, cutting private memory from ~3.9 GB to ~247 MB on a real worker without changing output vectors.
Vecs; later forwards pass null pointers.memmap2and streams the sha256 digest throughstd::io::copy; non-CUDA builds still load withfs::read.u32, so the f16 encode step is gone from the Qwen3 path.Written for commit ad3c627. Summary will update on new commits.