diff --git a/astrai/inference/core/cache.py b/astrai/inference/core/cache.py index 2777231..cb6d395 100644 --- a/astrai/inference/core/cache.py +++ b/astrai/inference/core/cache.py @@ -300,7 +300,11 @@ class KVCache(ABC): @abstractmethod def bind_tasks( - self, task_ids: List[str], total_len: int, device: torch.device + self, + task_ids: List[str], + total_len: int, + device: torch.device, + write_positions: Optional[Tensor] = None, ) -> CacheView: ... def task_cached(self, task_id: str) -> int: @@ -399,7 +403,11 @@ class PageCache(KVCache): self._pool.record(page_table[i], prompt_ids, i) def bind_tasks( - self, task_ids: List[str], total_len: int, device: torch.device + self, + task_ids: List[str], + total_len: int, + device: torch.device, + write_positions: Optional[Tensor] = None, ) -> PageCacheView: page_table = self._table.table_tensor(task_ids, device) return PageCacheView(self._storage, page_table, total_len) @@ -409,23 +417,37 @@ class ContiguousCacheView(CacheView): """Contiguous KV-cache view for attention layers.""" def __init__( - self, cache: "ContiguousCache", batch_indices: Tensor, total_len: int = 0 + self, + cache: "ContiguousCache", + batch_indices: Tensor, + total_len: int = 0, + write_positions: Optional[Tensor] = None, ): self._cache = cache self._batch_indices = batch_indices self._total_len = total_len + self._write_positions = write_positions def write(self, layer_id: int, k: Tensor, v: Tensor): seq_len = k.size(1) - start_pos = self._total_len - seq_len indices = self._batch_indices - self._cache.k[layer_id, indices, start_pos : start_pos + seq_len] = k - self._cache.v[layer_id, indices, start_pos : start_pos + seq_len] = v - new_len = start_pos + seq_len - for s in indices.tolist(): - cur = self._cache._slot_len.get(s, 0) - if new_len > cur: - self._cache._slot_len[s] = new_len + if self._write_positions is not None and seq_len == 1: + pos = self._write_positions + self._cache.k[layer_id, indices, pos] = k.squeeze(1) + self._cache.v[layer_id, indices, pos] = v.squeeze(1) + for s, p in zip(indices.tolist(), pos.tolist()): + cur = self._cache._slot_len.get(s, 0) + if p + 1 > cur: + self._cache._slot_len[s] = p + 1 + else: + start_pos = self._total_len - seq_len + self._cache.k[layer_id, indices, start_pos : start_pos + seq_len] = k + self._cache.v[layer_id, indices, start_pos : start_pos + seq_len] = v + new_len = start_pos + seq_len + for s in indices.tolist(): + cur = self._cache._slot_len.get(s, 0) + if new_len > cur: + self._cache._slot_len[s] = new_len def gather(self, layer_id: int) -> Tuple[Tensor, Tensor]: max_len = max( @@ -491,9 +513,21 @@ class ContiguousCache(KVCache): def task_extend(self, task_id: str, pos: int) -> bool: return pos < self.max_seq_len + def task_cached(self, task_id: str) -> int: + slot = self._task_slot.get(task_id) + if slot is None: + return 0 + return self._slot_len.get(slot, 0) + def bind_tasks( - self, task_ids: List[str], total_len: int, device: torch.device + self, + task_ids: List[str], + total_len: int, + device: torch.device, + write_positions: Optional[Tensor] = None, ) -> ContiguousCacheView: slots = [self._task_slot[tid] for tid in task_ids] batch_indices = torch.tensor(slots, dtype=torch.long, device=device) - return ContiguousCacheView(self, batch_indices, total_len) + return ContiguousCacheView( + self, batch_indices, total_len, write_positions=write_positions + ) diff --git a/astrai/inference/core/executor.py b/astrai/inference/core/executor.py index abaa164..6f2601f 100644 --- a/astrai/inference/core/executor.py +++ b/astrai/inference/core/executor.py @@ -106,7 +106,12 @@ class Executor: with torch.inference_mode(): outputs = self.model( input_ids.unsqueeze(1), - paged_cache=self.kv_cache.bind_tasks(task_ids, total_len, self.device), + paged_cache=self.kv_cache.bind_tasks( + task_ids, + total_len, + self.device, + write_positions=position_ids, + ), position_ids=position_ids.unsqueeze(1), ) logits = outputs["logits"][:, -1, :] diff --git a/astrai/inference/core/scheduler.py b/astrai/inference/core/scheduler.py index 66e14fb..3b8c8fa 100644 --- a/astrai/inference/core/scheduler.py +++ b/astrai/inference/core/scheduler.py @@ -138,36 +138,31 @@ class InferenceScheduler: t.task_id, t.prompt_ids, start_logical_page ) - pos_groups: Dict[int, List[Task]] = {} - for t in self._task_mgr.get_active_tasks(): - pos_groups.setdefault(t.next_pos, []).append(t) + decode_tasks = self._task_mgr.get_active_tasks() - for next_pos in sorted(pos_groups.keys()): - group = sorted(pos_groups[next_pos], key=lambda t: t.task_id) + valid: List[Task] = [] + for t in sorted(decode_tasks, key=lambda t: t.task_id): + if cache.task_extend(t.task_id, t.next_pos): + valid.append(t) + else: + t.status = TaskStatus.ABORTED + self._task_mgr.invoke_callback(t.task_id, STOP) - valid: List[Task] = [] - for t in group: - if cache.task_extend(t.task_id, t.next_pos): - valid.append(t) - else: - t.status = TaskStatus.ABORTED + if valid: + next_tokens = self._executor.execute_decode(valid) + + for t, ntok in zip(valid, next_tokens): + t.output_ids.append(ntok) + t.output_tokens += 1 + self._task_mgr.invoke_callback( + t.task_id, + self._task_mgr.tokenizer.decode([ntok]), + ) + + for t in valid: + if t.is_finished(stop_ids): self._task_mgr.invoke_callback(t.task_id, STOP) - if valid: - next_tokens = self._executor.execute_decode(valid) - - for t, ntok in zip(valid, next_tokens): - t.output_ids.append(ntok) - t.output_tokens += 1 - self._task_mgr.invoke_callback( - t.task_id, - self._task_mgr.tokenizer.decode([ntok]), - ) - - for t in valid: - if t.is_finished(stop_ids): - self._task_mgr.invoke_callback(t.task_id, STOP) - except Exception as e: self._stop_event.set() logger.error(f"Scheduler loop crashed: {e}", exc_info=True) diff --git a/scripts/tools/generate.py b/scripts/tools/generate.py index b5acec3..e1f15ba 100644 --- a/scripts/tools/generate.py +++ b/scripts/tools/generate.py @@ -1,8 +1,10 @@ import argparse import json +import time from typing import Optional import torch +from tqdm import tqdm from astrai.inference import InferenceEngine from astrai.model import AutoModel @@ -21,15 +23,26 @@ def processor( max_tokens: Optional[int], batch_size: int, num_samples: int = 1, + cache_len: int = 2048, + frequency_penalty: float = 0.0, + rep_window: int = 64, ): + print(f"Loading model from {param_path} ...") + t0 = time.time() model = AutoModel.from_pretrained(param_path) tokenizer = AutoTokenizer.from_pretrained(param_path) model.to(device="cuda", dtype=torch.bfloat16) + print(f" model loaded in {time.time() - t0:.1f}s") engine = InferenceEngine( - model=model, tokenizer=tokenizer, max_batch_size=batch_size * num_samples + model=model, + tokenizer=tokenizer, + max_batch_size=batch_size * num_samples, + max_seq_len=cache_len, + max_prompt_len=cache_len, ) + print(f"Reading {input_json_file} ...") with open(input_json_file, "r", encoding="utf-8") as f: input_data = [json.loads(line) for line in f] @@ -40,41 +53,69 @@ def processor( ] else: prompts = [item[question_key] for item in input_data] + print(f" {len(prompts)} prompts loaded\n") if max_tokens is None: max_tokens = model.config.max_len - if num_samples > 1: - prompts_expanded = [p for p in prompts for _ in range(num_samples)] - responses = engine.generate( - prompt=prompts_expanded, - stream=False, - max_tokens=max_tokens, - temperature=temperature, - top_p=top_p, - top_k=top_k, - ) - responses = [ - responses[i * num_samples : (i + 1) * num_samples] - for i in range(len(prompts)) - ] - else: - responses = engine.generate( - prompt=prompts, - stream=False, - max_tokens=max_tokens, - temperature=temperature, - top_p=top_p, - top_k=top_k, - ) + chunk_size = max(1, batch_size) with open(output_json_file, "w", encoding="utf-8") as f: - for i, prompt in enumerate(prompts): - if input_data and "messages" in input_data[0]: - output_item = {"response": responses[i]} + pbar = tqdm( + total=len(prompts) * num_samples, + unit="gen", + desc=f" Generating ({num_samples}x/prompt)", + ) + for chunk_start in range(0, len(prompts), chunk_size): + chunk = prompts[chunk_start : chunk_start + chunk_size] + + if num_samples > 1: + chunk_expanded = [p for p in chunk for _ in range(num_samples)] + resp_chunk = engine.generate( + prompt=chunk_expanded, + stream=False, + max_tokens=max_tokens, + temperature=temperature, + top_p=top_p, + top_k=top_k, + frequency_penalty=frequency_penalty, + rep_window=rep_window, + ) + resp_chunk = [ + resp_chunk[i * num_samples : (i + 1) * num_samples] + for i in range(len(chunk)) + ] else: - output_item = {question_key: prompt, response_key: responses[i]} - f.write(json.dumps(output_item, ensure_ascii=False) + "\n") + resp_chunk = engine.generate( + prompt=chunk, + stream=False, + max_tokens=max_tokens, + temperature=temperature, + top_p=top_p, + top_k=top_k, + frequency_penalty=frequency_penalty, + rep_window=rep_window, + ) + + for i, prompt in enumerate(chunk): + if input_data and "messages" in input_data[0]: + output_item = {"response": resp_chunk[i]} + else: + output_item = { + question_key: prompt, + response_key: resp_chunk[i], + } + f.write(json.dumps(output_item, ensure_ascii=False) + "\n") + + pbar.update(len(chunk) * num_samples) + + pbar.close() + + elapsed = time.time() - t0 + print( + f"\nDone! {len(prompts)} prompts x {num_samples} samples -> {output_json_file}" + ) + print(f"Total time: {elapsed:.1f}s ({elapsed / len(prompts):.2f}s/prompt)") engine.shutdown() @@ -145,6 +186,24 @@ if __name__ == "__main__": default=None, help="Maximum tokens to generate (default: model config max_len).", ) + parser.add_argument( + "--cache_len", + type=int, + default=2048, + help="KV cache & prompt truncation length (default: 2048, lower = less memory).", + ) + parser.add_argument( + "--frequency_penalty", + type=float, + default=0.0, + help="Frequency penalty to reduce repetition (default: 0.0, try 0.5-1.0).", + ) + parser.add_argument( + "--rep_window", + type=int, + default=64, + help="Window size for frequency penalty (default: 64).", + ) args = parser.parse_args()