From 3a4fb723f008f1a1a95761ac80d888ac9c357bf9 Mon Sep 17 00:00:00 2001 From: gspark Date: Fri, 23 Oct 2026 09:54:40 +0900 Subject: [PATCH] fix: KV-cache generation at batch size >= 2 (#164) The exported graph ties position_ids' batch dim to input_ids'. Prefill passed position_ids [1, S] and decode [1, 1], and decode fed the [B] sampled tokens as [1, B]. Prefill now builds [B, S] positions, decode [B, 1] positions and a [B, 1] token input. Adds tests/test_kv_cache_batch.py: greedy KV-cache generation matches HF at B=1 and B=2 (GPU only). Co-Authored-By: Claude Opus 5.5 --- opt_prime/opt_prime/inference.py | 4 +-- opt_prime/opt_prime/schedule_infer.py | 9 ++++--- opt_prime/tests/test_kv_cache_batch.py | 37 ++++++++++++++++++++++++++ 3 files changed, 43 insertions(+), 7 deletions(-) create mode 100644 opt_prime/tests/test_kv_cache_batch.py diff --git a/opt_prime/opt_prime/inference.py b/opt_prime/opt_prime/inference.py index 164a3acd..cce789ca 100644 --- a/opt_prime/opt_prime/inference.py +++ b/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) diff --git a/opt_prime/opt_prime/schedule_infer.py b/opt_prime/opt_prime/schedule_infer.py index 10bf1576..e68f79ed 100644 --- a/opt_prime/opt_prime/schedule_infer.py +++ b/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) diff --git a/opt_prime/tests/test_kv_cache_batch.py b/opt_prime/tests/test_kv_cache_batch.py new file mode 100644 index 00000000..d1a094c3 --- /dev/null +++ b/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)