From faefd43f180631aaec188cda16162efbf18584ba Mon Sep 17 00:00:00 2001 From: Ston Jia Date: Thu, 6 Nov 2025 16:51:55 +0800 Subject: [PATCH] fix Qwen3Embed panic when concurrency Original code get the rwlock only when embedding is done When we have multiple concurrent embedding at the same time, kvcache from one request may pollute another embedding process Like: ```txt Failed to process embedding for "/zion0/smbtest/1111 (1).txt": embed_file panicked: called `Result::unwrap()` on an `Err` value: shape mismatch in broadcast_add, lhs: [10, 16, 335, 654], rhs: [10, 1, 335, 335] ``` Now we hold the write lock during the whole procedure, allow preprocessing to be performed concurrently; but vector generation will not be concurrent. --- rust/src/embeddings/local/qwen3.rs | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/rust/src/embeddings/local/qwen3.rs b/rust/src/embeddings/local/qwen3.rs index 5bb17e35..c903ee01 100644 --- a/rust/src/embeddings/local/qwen3.rs +++ b/rust/src/embeddings/local/qwen3.rs @@ -126,12 +126,20 @@ impl Qwen3Embed for Qwen3Embedder { for mini_text_batch in text_batch.chunks(batch_size) { let (token_ids, attention_mask) = tokenize_batch(&self.tokenizer, mini_text_batch, &self.device)?; + let embeddings: Tensor = { - let mut model = self.model.write().map_err(|e| anyhow::anyhow!("Lock poisoned: {}", e))?; - model.forward(&token_ids, &attention_mask, 0)?.to_dtype(DType::F32)? + let mut model = self + .model + .write() + .map_err(|e| anyhow::anyhow!("Lock poisoned: {}", e))?; + let result = model + .forward(&token_ids, &attention_mask, 0)? + .to_dtype(DType::F32)?; + model.clear_kv_cache(); + + result }; - self.model.write().map_err(|e| anyhow::anyhow!("Lock poisoned: {}", e))?.clear_kv_cache(); let attention_mask = PooledOutputType::from(attention_mask); let attention_mask = Some(&attention_mask); let model_output = ModelOutput::Tensor(embeddings.clone());