perf: merge decode batch for 10x throughput

- merge all active decode tasks into single forward pass (was grouped by next_pos)
- add per-task write_positions to ContiguousCacheView for correct KV writes
- override ContiguousCache.task_cached (base returned 0, caused prefill loops)
- add --cache_len/--frequency_penalty/--rep_window to generate.py
- chunked batch processing with tqdm progress

bench (1.2B model, 128 prompts, 64 tok, batch=128):
  before: 77.2s, ~111 tok/s
  after:   7.1s, ~1210 tok/s (10.9x)

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
+        if self._write_positions is not None and seq_len == 1:
-        self._cache.v[layer_id, indices, start_pos : start_pos + seq_len] = v
+            pos = self._write_positions
-        new_len = start_pos + seq_len
+            self._cache.k[layer_id, indices, pos] = k.squeeze(1)
-        for s in indices.tolist():
+            self._cache.v[layer_id, indices, pos] = v.squeeze(1)
-            cur = self._cache._slot_len.get(s, 0)
+            for s, p in zip(indices.tolist(), pos.tolist()):
-            if new_len > cur:
+                cur = self._cache._slot_len.get(s, 0)
-                self._cache._slot_len[s] = new_len
+                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
+        )

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, :]

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]] = {}
+                decode_tasks = self._task_mgr.get_active_tasks()
-                for t in self._task_mgr.get_active_tasks():
-                    pos_groups.setdefault(t.next_pos, []).append(t)
 
-                for next_pos in sorted(pos_groups.keys()):
+                valid: List[Task] = []
-                    group = sorted(pos_groups[next_pos], key=lambda t: t.task_id)
+                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] = []
+                if valid:
-                    for t in group:
+                    next_tokens = self._executor.execute_decode(valid)
-                        if cache.task_extend(t.task_id, t.next_pos):
+
-                            valid.append(t)
+                    for t, ntok in zip(valid, next_tokens):
-                        else:
+                        t.output_ids.append(ntok)
-                            t.status = TaskStatus.ABORTED
+                        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)

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:
+    chunk_size = max(1, batch_size)
-        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,
-        )
 
     with open(output_json_file, "w", encoding="utf-8") as f:
-        for i, prompt in enumerate(prompts):
+        pbar = tqdm(
-            if input_data and "messages" in input_data[0]:
+            total=len(prompts) * num_samples,
-                output_item = {"response": responses[i]}
+            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]}
+                resp_chunk = engine.generate(
-            f.write(json.dumps(output_item, ensure_ascii=False) + "\n")
+                    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()