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)

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()