fix: KV-cache generation at batch size >= 2 (#164)

opt_prime/opt_prime/inference.py

@@ -1057,9 +1057,7 @@ def _generate_with_kv_cache(
 
             # Forward single token via decode_step
             if self.tpl.is_first_stage():
-                token_input = next_token.unsqueeze(0) if next_token.dim() == 1 else next_token
-                if token_input.dim() == 1:
-                    token_input = token_input.unsqueeze(0)  # [1, 1]
+                token_input = next_token.reshape(-1, 1)  # [B] -> [B, 1]
                 logits = self._schedule_gen.decode_step(token_input, cur_pos)
             else:
                 logits = self._schedule_gen.decode_step(None, cur_pos)

opt_prime/opt_prime/schedule_infer.py

@@ -370,7 +370,7 @@ def prefill(
 
         Args:
             input_ids: Input token IDs tensor of shape [batch_size, seq_len]
-            position_ids: Position IDs tensor of shape [1, seq_len] for RoPE.
+            position_ids: Position IDs tensor of shape [batch_size, seq_len] for RoPE.
                           If None, auto-generated as arange(seq_len).
 
         Returns:
@@ -383,9 +383,10 @@ def prefill(
             # Set position_ids for RoPE (critical for LLaMA)
             if position_ids is None and input_ids is not None:
                 seq_len = input_ids.size(1) if input_ids.dim() > 1 else input_ids.size(0)
+                # The export ties position_ids' batch dim to input_ids'
                 position_ids = torch.arange(
                     seq_len, device=input_ids.device
-                ).unsqueeze(0)
+                ).expand(input_ids.size(0) if input_ids.dim() > 1 else 1, -1)
             if position_ids is not None:
                 self.optimus.run_info.env[0]["position_ids"] = \
                     position_ids.to(self.optimus.run_info.device)
@@ -427,8 +428,8 @@ def decode_step(
         if self.optimus.tpl.is_first_stage():
             self.get_input(input_token)
             # Set position_ids for RoPE — must be the absolute position
-            position_ids = torch.tensor(
-                [[position]], device=input_token.device
+            position_ids = torch.full(
+                (input_token.size(0), 1), position, device=input_token.device
             )
             self.optimus.run_info.env[0]["position_ids"] = \
                 position_ids.to(self.optimus.run_info.device)

opt_prime/tests/test_kv_cache_batch.py

@@ -0,0 +1,37 @@
+#
+# Copyright (c) 2026-present, ETRI, All rights reserved.
+#
+# opt_prime.inference: greedy generation with the KV cache (torch.export path) matches HF at batch size 1 and 2.
+import copy
+import os
+import sys
+
+import pytest
+import torch
+
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="Optimus_Inference needs a GPU")
+
+
+@pytest.fixture(scope="module")
+def tiny():
+    from transformers import AutoConfig, AutoModelForCausalLM
+    torch.manual_seed(0)
+    cfg = AutoConfig.for_model("llama", hidden_size=64, intermediate_size=128, num_hidden_layers=2,
+                               num_attention_heads=4, num_key_value_heads=2, vocab_size=128,
+                               attn_implementation="sdpa")
+    return AutoModelForCausalLM.from_config(cfg).eval()
+
+
+@pytest.mark.parametrize("batch", [1, 2])
+def test_kv_cache_generation_matches_hf(tiny, batch):
+    from opt_prime.inference import Optimus_Inference
+    ids = torch.randint(0, 128, (batch, 6), generator=torch.Generator().manual_seed(batch))
+    with torch.no_grad():
+        ref = copy.deepcopy(tiny).cuda().generate(ids.cuda(), max_new_tokens=5, do_sample=False,
+                                                  pad_token_id=0).cpu()
+    eng = Optimus_Inference(copy.deepcopy(tiny), use_kv_cache=True, dynamo_capture=True, dtype=torch.float32)
+    eng.eval()
+    out = eng.generate(ids, max_new_tokens=5, do_sample=False, verbose=False).cpu()
+    assert torch.equal(out, ref)