perf: vectorize REINFORCE++ discounted returns

[morluto]

Problem

get_reinforce_plus_plus_returns computes every token return in a Python reverse loop for every sample. That host-side loop becomes expensive for batched long responses, while ppo_utils.py already contains a chunked discounted scan for GAE.

Change

• Extract the existing recurrence into chunked_discounted_returns.
• Right-pad variable-length REINFORCE++ token rewards, scan the batch once, and trim each row before returning or context-parallel slicing.
• Reuse the same helper from chunked_gae.
• Register CPU tests for serial equivalence, chunk boundaries, discount edge cases, variable-length padding, dtype/device preservation, the REINFORCE++ caller, and GAE.

The change does not add a dependency or alter context-parallel gather/slice behavior.

Correctness

The CPU tests compare the scan against the previous reverse recurrence for float32 and float64, discounts 0, 0.5, 0.99, and 1, and lengths around the 128-token chunk boundary. Right-padded positions remain zero and are trimmed before caller-visible results.

PYTHONPATH=. python3 tests/test_discounted_returns.py
45 passed in 2.36s

git diff --check origin/main also passes.

CPU microbenchmark

Environment: 4-core DO-Premium-AMD CPU, Python 3.12.3, PyTorch 2.11.0, one PyTorch thread, random float32 rewards, discount 0.99, chunk size 128, one warmup, median of three repetitions.

The candidate used chunked_gae(rewards, zeros_like(rewards), 0.99, 1.0)[0], which is the pre-existing recurrence extracted by this PR.

Shape	Python reverse loop	Chunked scan	Speedup
B=1, T=1,024	29.688 ms	1.357 ms	21.87x
B=8, T=4,096	1,778.590 ms	4.979 ms	357.21x

This measures CPU return computation only. It is not a GPU or end-to-end training-throughput claim.

Review notes

Suggested order: review chunked_discounted_returns, then the REINFORCE++ batching/trim path, then the CPU regression tests and generated CI entry.

Disclosure

This contribution was prepared with AI-assisted research and implementation, followed by CPU validation and structured code review.

[chatgpt-codex-connector]

You have reached your Codex usage limits for code reviews. You can see your limits in the Codex usage dashboard.