diff --git a/astrai/inference/api/protocol.py b/astrai/inference/api/protocol.py index b00bf2d..07c98e9 100644 --- a/astrai/inference/api/protocol.py +++ b/astrai/inference/api/protocol.py @@ -181,12 +181,10 @@ class ProtocolHandler: self, agen: AsyncGenerator, ctx: GenContext, stop_sequences: List[str] ) -> Dict[str, Any]: checker = StopChecker(stop_sequences) - chunks: List[str] = [] body = "" matched = None async for token in agen: - chunks.append(token) body += token matched = checker.check(body) @@ -195,6 +193,5 @@ class ProtocolHandler: ctx.completion_tokens += 1 - content = "".join(chunks) stop = StopInfo(matched=matched, body=body) - return self.builder.format_response(ctx, content, stop) + return self.builder.format_response(ctx, body, stop) diff --git a/astrai/inference/core/cache.py b/astrai/inference/core/cache.py index dd7dcb8..258e987 100644 --- a/astrai/inference/core/cache.py +++ b/astrai/inference/core/cache.py @@ -413,6 +413,8 @@ class PagePool: if slots is None: for p in self._task_pages[task_id]: self._alloc.free(p) + self._task_pages.pop(task_id, None) + self._task_slots.pop(task_id, None) self._req_pool.free([req_idx]) del self._task_req[task_id] return False @@ -427,6 +429,8 @@ class PagePool: self._alloc.free(hp) for np_ in new_pages: self._alloc.free(np_) + self._task_pages.pop(task_id, None) + self._task_slots.pop(task_id, None) self._req_pool.free([req_idx]) del self._task_req[task_id] return False @@ -442,6 +446,7 @@ class PagePool: req_idx = self._task_req.pop(task_id, None) if req_idx is None: return + self._bind_state = None self._task_len.pop(req_idx, None) self._task_cached.pop(task_id, None) @@ -455,6 +460,9 @@ class PagePool: else: for p in self._task_pages.get(task_id, []): self._alloc.free(p) + if self.page_size == 1: + for slot in self._task_slots.get(task_id, []): + self._alloc.free(slot) self._task_pages.pop(task_id, None) self._task_slots.pop(task_id, None) @@ -503,7 +511,7 @@ class PagePool: def task_record_hashes( self, task_id: str, prompt_ids: List[int], start_logical_page: int = 0 ): - if self._prefix is None or self.contiguous: + if self._prefix is None: return pages = self._task_pages.get(task_id, []) full_pages = len(prompt_ids) // self.page_size @@ -635,9 +643,6 @@ class PagePool: req_idx = self._task_req[task_id] total = len(prompt_ids) - if self.contiguous: - return - if self.page_size == 1: slots = self._task_slots.get(task_id, []) all_slots = slots[: total - cached] diff --git a/astrai/inference/core/executor.py b/astrai/inference/core/executor.py index fd8c185..c685c68 100644 --- a/astrai/inference/core/executor.py +++ b/astrai/inference/core/executor.py @@ -19,7 +19,6 @@ from astrai.inference.core.task import Task from astrai.inference.core.workspace import InferenceWorkspace from astrai.inference.sample import sample from astrai.model.automodel import AutoModel -from astrai.tokenize.tokenizer import AutoTokenizer logger = logging.getLogger(__name__) @@ -184,13 +183,11 @@ class Executor: def __init__( self, model: AutoModel, - tokenizer: AutoTokenizer, kv_cache: PagePool, device: Optional[str] = None, dtype: Optional[torch.dtype] = None, ): self.model = model - self.tokenizer = tokenizer self.kv_cache = kv_cache self.device = device or next(model.parameters()).device self.dtype = dtype or next(model.parameters()).dtype @@ -207,7 +204,6 @@ class Executor: config = model.config max_q_heads = config.num_attention_heads head_dim = config.hidden_size // config.num_attention_heads - self._head_dim = head_dim self._graph_supported = CudaBackend.supports(head_dim=head_dim) self._workspace = InferenceWorkspace( max_batch_size=kv_cache.max_batch_size, diff --git a/astrai/inference/core/metrics.py b/astrai/inference/core/metrics.py index 9996367..5aaba38 100644 --- a/astrai/inference/core/metrics.py +++ b/astrai/inference/core/metrics.py @@ -12,7 +12,7 @@ class TaskTiming: """Timestamp snapshots and computed metrics for one generation task. Created by :class:`MetricsCollector` at task-registration time; - updated via ``prefill_scope`` / ``mark_finished``. + updated via ``record`` / ``mark_finished``. """ task_id: str @@ -27,8 +27,6 @@ class TaskTiming: # derived metrics - # derived metrics - @property def queue_wait_ms(self) -> Optional[float]: if self.prefill_start_time is not None: @@ -116,7 +114,7 @@ class MetricsCollector: metrics = MetricsCollector() metrics.register(task_id, arrival_time) - with metrics.prefill_scope(task_ids): + with metrics.record(task_ids, "prefill"): run_prefill(...) metrics.mark_finished(task_id, input_tokens, output_tokens) diff --git a/astrai/inference/core/scheduler.py b/astrai/inference/core/scheduler.py index ac2d204..1bd6c7e 100644 --- a/astrai/inference/core/scheduler.py +++ b/astrai/inference/core/scheduler.py @@ -68,7 +68,6 @@ class InferenceScheduler: self._executor = Executor( model=model, - tokenizer=tokenizer, kv_cache=self._cache, device=self.device, dtype=self.dtype, @@ -210,9 +209,6 @@ class InferenceScheduler: if new_text: self._task_mgr.invoke_callback(t.task_id, new_text) if t.is_finished(stop_ids): - remaining = t.flush_remaining(self._task_mgr.tokenizer) - if remaining: - self._task_mgr.invoke_callback(t.task_id, remaining) self._task_mgr.invoke_callback(t.task_id, STOP) except Exception as e: diff --git a/astrai/inference/core/task.py b/astrai/inference/core/task.py index 57a4e07..cd04ade 100644 --- a/astrai/inference/core/task.py +++ b/astrai/inference/core/task.py @@ -1,4 +1,3 @@ -import logging import threading import time import uuid @@ -11,8 +10,6 @@ from tokenizers.decoders import DecodeStream from astrai.inference.core.metrics import MetricsCollector from astrai.tokenize.tokenizer import AutoTokenizer -logger = logging.getLogger(__name__) - STOP = object() @@ -92,16 +89,6 @@ class Task: self._decoder = StreamDecoder(tokenizer) return self._decoder.push(self.output_ids[-1]) - def flush_remaining(self, tokenizer: AutoTokenizer) -> str: - """Emit any text still buffered in the decoder. - - With the Rust-native DecodeStream, the stream is always in a - correct state — any completed text was already emitted by the - last ``push``. A trailing incomplete multi-byte sequence has no - valid text to emit, so this is a no-op. - """ - return "" - @property def next_pos(self) -> int: # The first output is sampled from prefill and enters KV on the next step. @@ -157,11 +144,6 @@ class TaskManager: if len(prompt_ids) > self.max_seq_len: prompt_ids = prompt_ids[-self.max_seq_len :] - if len(prompt_ids) > self.max_seq_len: - if stream_callback: - stream_callback(STOP) - return task_id - if max_tokens is None: max_tokens = self.max_seq_len - len(prompt_ids) else: diff --git a/astrai/inference/sample.py b/astrai/inference/sample.py index df6a240..b921961 100644 --- a/astrai/inference/sample.py +++ b/astrai/inference/sample.py @@ -363,20 +363,6 @@ def sample( ``True`` — a ``(token_ids, chosen_logprobs)`` tuple where ``chosen_logprobs`` has shape ``[batch]``. """ - greedy = ( - bool((temperature == 0).all()) - if isinstance(temperature, Tensor) - else temperature == 0 - ) - - if greedy: - tokens = logits.argmax(dim=-1) - if not return_logprobs: - return tokens - log_probs = torch.log_softmax(logits.float(), dim=-1) - chosen = torch.gather(log_probs, -1, tokens.unsqueeze(-1)).squeeze(-1) - return tokens, chosen - has_freq = ( (isinstance(frequency_penalty, Tensor) and (frequency_penalty != 0).any()) if isinstance(frequency_penalty, Tensor)