happyz
/
gemma.cpp
mirror of https://github.com/google/gemma.cpp.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #394 from ufownl:bugfix/prefix_lm 

PiperOrigin-RevId: 678710685
This commit is contained in:
Copybara-Service 2024-09-25 08:25:31 -07:00
parent e70e686805 d1010337c3
commit e3507190ae
1 changed files with 2 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
gemma/gemma-inl.h
View File

@ -916,7 +916,8 @@ HWY_NOINLINE void Prefill(
} }
// In prefix-LM mode, we need to look at all the tokens for the prefix in // In prefix-LM mode, we need to look at all the tokens for the prefix in
// one iteration through the layers, so we need a large enough batch size. // one iteration through the layers, so we need a large enough batch size.
HWY_ASSERT(max_tbatch_size >= prefill_this_query); HWY_ASSERT(prefix_end_this_query == 0 ||
max_tbatch_size >= prefill_this_query);
// For each batch of tokens in the query: // For each batch of tokens in the query:
for (size_t tbatch_start = 0; tbatch_start < prefill_this_query; for (size_t tbatch_start = 0; tbatch_start < prefill_this_query;