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)
2026-07-18 08:50:46 +08:00
parent f7df02f9a3
commit a24a7b4da5
4 changed files with 162 additions and 69 deletions
@@ -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
+        )
@@ -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, :]
@@ -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)
@@ -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()