(
self,
updates: Sequence["MTPDraftProvider.SampledBatchUpdate"],
/,
)
| 2306 | ) |
| 2307 | |
| 2308 | def process_sampled_batch( |
| 2309 | self, |
| 2310 | updates: Sequence["MTPDraftProvider.SampledBatchUpdate"], |
| 2311 | /, |
| 2312 | ) -> List[np.ndarray]: |
| 2313 | results = [np.array([], dtype=np.intc) for _ in updates] |
| 2314 | if self.num_pred_tokens <= 0 or not updates: |
| 2315 | return results |
| 2316 | h_tgt = llama_cpp_ext.llama_get_embeddings_nextn(self.target_ctx) |
| 2317 | if not h_tgt: |
| 2318 | raise RuntimeError("missing target nextn embeddings for MTP") |
| 2319 | n_target_rows = max( |
| 2320 | ( |
| 2321 | max(update.row_indices) + 1 |
| 2322 | for update in updates |
| 2323 | if update.row_indices |
| 2324 | ), |
| 2325 | default=0, |
| 2326 | ) |
| 2327 | if n_target_rows <= 0: |
| 2328 | return results |
| 2329 | h_tgt_rows = np.ctypeslib.as_array(h_tgt, shape=(n_target_rows, self.n_embd)) |
| 2330 | |
| 2331 | plan = self._build_sampled_batch_plan(updates) |
| 2332 | if not plan.context_rows and not plan.pending_rows: |
| 2333 | return results |
| 2334 | if ( |
| 2335 | len(plan.context_rows) > self.n_batch |
| 2336 | or len(plan.pending_rows) > self.n_batch |
| 2337 | ): |
| 2338 | raise RuntimeError("MTP draft batch capacity exceeded") |
| 2339 | |
| 2340 | self._decode_sampled_context_rows(plan.context_rows, h_tgt_rows) |
| 2341 | sampled_outputs = self._decode_sampled_pending_rows(plan, h_tgt_rows) |
| 2342 | cleanup_keep_len_by_seq: Dict[int, int] = {} |
| 2343 | for output in sampled_outputs: |
| 2344 | update = updates[output.update_index] |
| 2345 | sample_index = update.sample_index |
| 2346 | sample_source_row = update.row_indices[sample_index] |
| 2347 | self.pending_h[output.seq_id] = h_tgt_rows[sample_source_row] |
| 2348 | self.ready[output.seq_id] = True |
| 2349 | self.ready_pos[output.seq_id] = output.ready_pos |
| 2350 | cleanup_keep_len_by_seq[output.seq_id] = output.keep_len |
| 2351 | |
| 2352 | if not self.is_mem_shared: |
| 2353 | for seq_id, keep_len in cleanup_keep_len_by_seq.items(): |
| 2354 | self._truncate_memory(seq_id, keep_len) |
| 2355 | |
| 2356 | if sampled_outputs: |
| 2357 | active = self._start_sampled_draft_states( |
| 2358 | updates, |
| 2359 | sampled_outputs, |
| 2360 | results, |
| 2361 | ) |
| 2362 | self._extend_sampled_draft_states( |
| 2363 | active, |
| 2364 | results, |
| 2365 | cleanup_keep_len_by_seq, |
nothing calls this directly
no test coverage detected