(self)
| 13361 | return shifted - math.log(float(np.sum(np.exp(shifted, dtype=np.float64)))) |
| 13362 | |
| 13363 | def step(self) -> bool: |
| 13364 | step_started_at = time.perf_counter() |
| 13365 | if self.closed: |
| 13366 | return False |
| 13367 | try: |
| 13368 | self.admit_waiting() |
| 13369 | batch_items = self.build_batch() |
| 13370 | if not batch_items: |
| 13371 | if self.maybe_fill_batched_draft_tokens(): |
| 13372 | self.finalize_cancelled() |
| 13373 | return True |
| 13374 | return self.finalize_cancelled() |
| 13375 | draft_processing_enabled = self.batch_needs_draft_processing(batch_items) |
| 13376 | self.defer_sampled_draft_processing = ( |
| 13377 | draft_processing_enabled |
| 13378 | and self.supports_sampled_draft_processing |
| 13379 | and all(item.kind == "decode" for item in batch_items) |
| 13380 | ) |
| 13381 | self.model.set_draft_processing_enabled(draft_processing_enabled) |
| 13382 | self.model.clear_batch() |
| 13383 | for item in batch_items: |
| 13384 | prefill = item.prefill_state |
| 13385 | if prefill is not None and prefill.embeddings is not None: |
| 13386 | if prefill.positions is None: |
| 13387 | raise RuntimeError("embedding batch item is missing positions") |
| 13388 | self.model.add_batch_embeddings( |
| 13389 | seq_id=item.seq_id, |
| 13390 | embeddings=prefill.embeddings, |
| 13391 | positions=prefill.positions, |
| 13392 | output_indices=item.output_indices, |
| 13393 | ) |
| 13394 | else: |
| 13395 | self.model.add_batch_tokens( |
| 13396 | seq_id=item.seq_id, |
| 13397 | start_pos=item.llama_start_pos, |
| 13398 | tokens=item.tokens, |
| 13399 | output_indices=item.output_indices, |
| 13400 | ) |
| 13401 | decode_started_at = time.perf_counter() |
| 13402 | try: |
| 13403 | non_causal_decode = any( |
| 13404 | item.prefill_state is not None and item.prefill_state.non_causal |
| 13405 | for item in batch_items |
| 13406 | ) |
| 13407 | try: |
| 13408 | if non_causal_decode: |
| 13409 | llama_cpp.llama_set_causal_attn(self.model.ctx, False) |
| 13410 | self.model.decode() |
| 13411 | finally: |
| 13412 | if non_causal_decode: |
| 13413 | llama_cpp.llama_set_causal_attn(self.model.ctx, True) |
| 13414 | decode_elapsed = time.perf_counter() - decode_started_at |
| 13415 | if draft_processing_enabled and not self.defer_sampled_draft_processing: |
| 13416 | draft_process_started_at = time.perf_counter() |
| 13417 | self.model.process_draft_batch() |
| 13418 | draft_process_elapsed = time.perf_counter() - draft_process_started_at |
| 13419 | self.metrics.draft_process_seconds_total += draft_process_elapsed |
| 13420 | self.metrics.draft_process_calls_total += 1 |
no test coverage detected