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..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,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::>(); - 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(), + }; + 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") } 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..1f7b2b64 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,6 +177,21 @@ fn load_safetensor_map(root: &Path, original: &Path) -> Result Result> { + // 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()))?; @@ -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], ) -> 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 +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) 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());