| 209 | } |
| 210 | |
| 211 | buffer<bf16> AutoModel::_chunked_insert(chat_meta_info_t& meta_info, std::vector<int>& tokens, std::function<bool()> is_cancelled, void* payload, int first_len_run) { |
| 212 | int max_prefill_len = meta_info.max_prefill_len; |
| 213 | // make max_prefill_len a 2^n |
| 214 | max_prefill_len = 1 << static_cast<int>(std::ceil(std::log2(max_prefill_len))); |
| 215 | buffer<bf16> y; |
| 216 | if (max_prefill_len < 512) { |
| 217 | y = this->lm_engine->prefill(tokens, payload); |
| 218 | } |
| 219 | else{ |
| 220 | if (first_len_run > 0) { |
| 221 | int new_max_len = 1 << static_cast<int>(std::ceil(std::log2(first_len_run))); |
| 222 | if (new_max_len > max_prefill_len) { |
| 223 | max_prefill_len = new_max_len; |
| 224 | } |
| 225 | } |
| 226 | int chunks = (tokens.size() + max_prefill_len - 1) / max_prefill_len; |
| 227 | for (int i = 0; i < chunks; i++) { |
| 228 | if (is_cancelled()) { |
| 229 | meta_info.stop_reason = CANCEL_DETECTED; |
| 230 | // reset stream content |
| 231 | buffer_.clear(); |
| 232 | current_mode_ = StreamEventType::CONTENT; |
| 233 | tool_name_.clear(); |
| 234 | is_in_tool_block_ = false; |
| 235 | break; |
| 236 | } |
| 237 | int start = i * max_prefill_len; |
| 238 | int end = std::min(static_cast<int>(tokens.size()), (i + 1) * max_prefill_len); |
| 239 | std::vector<int> chunk_tokens(tokens.begin() + start, tokens.begin() + end); |
| 240 | header_print("FLM", "Prefill chunk " + std::to_string(i+1) + "/" + std::to_string(chunks) + " with " + std::to_string(chunk_tokens.size()) + " tokens"); |
| 241 | buffer<bf16> chunk_y = this->lm_engine->prefill(chunk_tokens, (i == 0)? payload : nullptr); |
| 242 | if (i == chunks - 1) { |
| 243 | y = chunk_y; |
| 244 | } |
| 245 | } |
| 246 | } |
| 247 | return y; |
| 248 | } |
| 249 | |
| 250 | std::string AutoModel::_shared_generate(chat_meta_info_t& meta_info, int length_limit, std::ostream& os, std::function<bool()> is_cancelled) { |
| 251 | std::vector<int> sampled_tokens; |