[gspark-etri]

Background

With the KV cache on, Optimus_Inference.generate crashed for batch size >= 2 (#164). The torch.export graph ties the batch dim of position_ids to that of input_ids, but prefill passed position_ids [1, S], decode passed [1, 1], and decode fed the [B] sampled tokens as [1, B] (unsqueeze(0)).

Changes

• opt_prime/opt_prime/schedule_infer.py ScheduleGeneration.prefill: default position_ids are arange(S) expanded to [B, S] (docstring updated).
• opt_prime/opt_prime/schedule_infer.py ScheduleGeneration.decode_step: position_ids are [B, 1] (torch.full) instead of [[position]].
• opt_prime/opt_prime/inference.py decode loop (was :1060): token input is next_token.reshape(-1, 1) ([B, 1]); replaces three lines of unsqueeze(0) handling.
• New opt_prime/tests/test_kv_cache_batch.py: tiny random Llama (sdpa), greedy, 5 new tokens, KV cache on, export path (dynamo_capture=True); output must equal HF generate exactly for B=1 and B=2. Skipped without CUDA.

Verification

Environment: Python 3.10, torch 2.10.0+cu128, transformers 5.1.0, one A100 80GB (CUDA_VISIBLE_DEVICES=1). Base: optimus-rm @253be20.

cd opt_prime && TORCH_DEVICE_BACKEND_AUTOLOAD=0 CUDA_VISIBLE_DEVICES=1 PYTHONPATH=$PWD python3 -m pytest -q tests/test_kv_cache_batch.py

• Before the fix: 1 failed, 1 passed. B=1 passes; B=2 fails in prefill with RuntimeError: The expanded size of the tensor (2) must match the existing size (4) at non-singleton dimension 1. Target sizes: [1, 2, 2, 6, 16].
• Fixing only prefill and the token shape is not enough: decode then fails the same way (Target sizes: [1, 2, 2, 1, 16]), so the decode position_ids change is needed.
• After the fix: 2 passed.

Full suite (cd opt_prime && TORCH_DEVICE_BACKEND_AUTOLOAD=0 CUDA_VISIBLE_DEVICES=1 PYTHONPATH=$PWD python3 -m pytest -q tests):

• Before (253be20): 400 passed, 23 skipped, 135 subtests passed
• After: 402 passed, 23 skipped, 135 subtests passed

Compatibility and limits

• B=1 behavior is unchanged (positions and token input have the same shapes as before).
• A caller that passes its own position_ids to prefill must now pass [B, S] for B >= 2 (the docstring said [1, S]); in-tree callers pass none.
• Checked on a single process (one stage). Multi-stage PP and TP generation with B >= 2 were not run here; the TP placeholder token on middle stages (torch.zeros(1)) is untouched.
• Left-padded prompts of different lengths are out of scope: positions are arange(S) for every row, as before.

Related Issue

Closes #164

🤖 Generated with Claude Code