Skip to content

Commit 6340427

Browse files
protobird-gitcopybara-github
authored andcommitted
Fix a bug of wrong lock type in embeddings
- Write is the correct type as contents are supposed to be written to the tensor. - Lock once to reduce IO overhead which might be slight worse if IO happens <= 1. LiteRT-LM-PiperOrigin-RevId: 863415711
1 parent 40d3a47 commit 6340427

File tree

1 file changed

+4
-7
lines changed

1 file changed

+4
-7
lines changed

runtime/components/embedding_lookup/embedding_lookup_manager.cc

Lines changed: 4 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -195,18 +195,15 @@ absl::Status EmbeddingLookupManager::LookupPrefill(
195195
// If fully_supports_multi_modal_ is false, then we need to fill in the
196196
// missing embeddings with the default embedding vector.
197197
const size_t bytes_per_token = floats_per_token * sizeof(float);
198+
auto output_tensor_lock_and_addr = ::litert::TensorBufferScopedLock::Create(
199+
*output_tensor, ::litert::TensorBuffer::LockMode::kWrite);
200+
auto output_tensor_ptr =
201+
reinterpret_cast<uint8_t*>(output_tensor_lock_and_addr->second);
198202
for (int i = 0; i < tokens.size(); ++i) {
199203
if (tokens[i] >= 0) {
200204
continue;
201205
}
202206
size_t byte_offset_for_token = byte_offset + i * bytes_per_token;
203-
204-
auto output_tensor_lock_and_addr =
205-
::litert::TensorBufferScopedLock::Create(
206-
*output_tensor, ::litert::TensorBuffer::LockMode::kRead);
207-
auto output_tensor_ptr =
208-
reinterpret_cast<uint8_t*>(output_tensor_lock_and_addr->second);
209-
210207
memcpy(output_tensor_ptr + byte_offset_for_token,
211208
default_embedding_vector_.data(),
212209
default_embedding_vector_.size() * sizeof(float));

0 commit comments

Comments
 (0)