core: KV-cache generation crashes at batch size >= 2 #164

Problem

With the KV cache on, Optimus_Inference generation fails for batch size 2 or more. Prefill passes position_ids of shape [1,S] while input_ids is [B,S] (the export ties both to one batch dim), and decode feeds the [B] next tokens as [1,B] (unsqueeze(0)) instead of [B,1].

Evidence and reproduction

opt_prime/opt_prime/inference.py:1060, opt_prime/opt_prime/schedule_infer.py:386. Tiny Llama via the torch.export path vs HF greedy: B=1 with KV cache matches, B=2 without KV cache matches, B=2 with KV cache crashes in prefill with "expanded size ... must match".

Environment

optimus-rm @253be20; Python 3.10, torch 2.10.0+cu128, transformers 5.1.0; RNGD+A100 host (CPU / one A100 for the checks).

Scope

inference.py, schedule_infer.py; a parity test.

Done when

B=2 generation with KV cache matches HF greedy; a test covers B=1 and B=2.

Owner

@gspark-etri