Skip to content

Commit b5e15c8

Browse files
kirklandsignfacebook-github-bot
authored andcommitted
Fix dangling pointer in TextTokenGenerator non-kv-cache path
Summary: In the non-kv-cache branch of TextTokenGenerator::generate(), push_back() on token_data can trigger vector reallocation, but the tensor created via from_blob still points to the old data address. resize_tensor_ptr only updates shape metadata, not the data pointer, resulting in a dangling pointer. Fix by pre-allocating the vector with reserve() before creating the tensor, ensuring push_back never triggers reallocation during the generate loop. Differential Revision: D99408541
1 parent 3d2c853 commit b5e15c8

1 file changed

Lines changed: 2 additions & 0 deletions

File tree

extension/llm/runner/text_token_generator.h

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -77,6 +77,8 @@ class ET_EXPERIMENTAL TextTokenGenerator {
7777
} else {
7878
token_data = tokens;
7979
token_shape = {1, static_cast<int>(tokens.size())};
80+
// Prevent reallocation that would invalidate from_blob's data pointer.
81+
token_data.reserve(token_data.size() + max_new_tokens);
8082
}
8183

8284
// initialize tensor wrappers

0 commit comments

Comments
 (0)