test: add selected logprob parity harness

[Txy02]

Summary

Adds a small selected-logprob parity harness for validating that the same effective completion rows produce identical selected logprobs across batch layout changes.

This is intended as a first step toward the rollout-vs-training numerical parity work discussed in #148, #152, and #154.

What is covered

• Batch-position invariance for selected logprobs
• Padding-layout invariance for selected logprobs
• Drift summaries through the existing summarize_kernel_drift
• Optional CUDA FP32/BF16 dtype sweep when CUDA is available

What is intentionally out of scope

• No kernel implementation changes
• No vLLM integration yet
• No benchmark changes
• No external framework integration

Tests

Validated on AutoDL RTX 4090:

• python -m pytest tests/test_logprob_parity.py -v
• python -m pytest tests/test_reference_ops.py tests/test_logprob_parity.py -v

Result:

• 4 passed
• 11 passed

Summary by CodeRabbit

• New Features
  • Added testing helpers to generate padded batch layouts and compare selected log-probability results across different batch/padding layouts.
• Tests
  • Added invariance tests for batch position and padding layout.
  • Added negative tests for invalid row mappings and out-of-range padding token IDs.
  • Added CUDA dtype sweep coverage for float32 and bfloat16, validating consistent parity in a chosen output dtype.
• Chores
  • Re-exported the new testing helpers from the main testing utilities entry point.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: 1142d5d1-35fb-4447-a3d2-a856855a73ed

📥 Commits

Reviewing files that changed from the base of the PR and between 92b9799 and b208efc.

📒 Files selected for processing (3)

• rl_engine/testing/__init__.py
• rl_engine/testing/logprob_parity.py
• tests/test_logprob_parity.py

🚧 Files skipped from review as they are similar to previous changes (2)

• rl_engine/testing/init.py
• tests/test_logprob_parity.py

---

📝 Walkthrough

Walkthrough

Adds a new testing module for selected-logprob layout parity, re-exports its helpers from the testing package, and expands the test suite with permutation, padding, invalid-input, and CUDA dtype checks.

Changes

Logprob Parity Utilities and Tests

Layer / File(s)	Summary
make_padded_batch_layout and compare_selected_logprob_layouts rl_engine/testing/logprob_parity.py	make_padded_batch_layout validates tensor shapes, pad tokens, and destination rows, then builds padded logits, token IDs, and masks. compare_selected_logprob_layouts computes selected logprobs for reference and candidate inputs, remaps candidate rows, and returns kernel-drift summary metrics.
Re-export from testing package rl_engine/testing/__init__.py	Imports and adds compare_selected_logprob_layouts and make_padded_batch_layout to __all__.
Parity tests: permutation, padding, CUDA dtype sweep tests/test_logprob_parity.py	Adds a seeded _case() helper plus tests for batch-row permutation invariance, padded-layout parity, invalid input rejection, and a CUDA-only float32/bfloat16 sweep.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~10 minutes

Poem

🐇 I hopped through rows with padded grace,
And logprobs stayed in perfect place.
Shuffle, pad, and dtype tune—
Zero drift still sings this tune!

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 25.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title accurately summarizes the new selected-logprob parity test harness added in this PR.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@rl_engine/testing/logprob_parity.py`:
- Around line 89-90: Validate candidate_rows before using it in logprob_parity’s
indexing path: in the logic that builds restored from candidate[rows], add the
same row-id checks used by make_padded_batch_layout so each reference batch
entry maps to exactly one candidate row, with no negative or duplicate indices.
Also bounds-check the remapped rows against candidate.shape[0] before indexing,
and fail fast if the layout is invalid.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: ff5ebd1f-7893-473e-a0b0-c33ea78094ba

📥 Commits

Reviewing files that changed from the base of the PR and between 9480500 and f713279.

📒 Files selected for processing (3)

• rl_engine/testing/__init__.py
• rl_engine/testing/logprob_parity.py
• tests/test_logprob_parity.py

[coderabbitai]

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@rl_engine/testing/logprob_parity.py`:
- Line 20: The `pad_token_id` used in `selected_logprobs_reference` and the
related test setup must be validated against the logits vocabulary before any
indexing occurs. Add an explicit check in the `logprob_parity` test helpers to
ensure `pad_token_id` is within the valid token range derived from the logits
shape, and fail fast with a clear assertion if it is not. Apply the same
validation wherever `pad_token_id` is passed through the parity test path so
fully masked rows never rely on an out-of-range index.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: 48ce688f-85e0-4cdf-983c-d78932299446

📥 Commits

Reviewing files that changed from the base of the PR and between f713279 and 7b31bf7.

📒 Files selected for processing (3)

• rl_engine/testing/__init__.py
• rl_engine/testing/logprob_parity.py
• tests/test_logprob_parity.py

🚧 Files skipped from review as they are similar to previous changes (1)

• rl_engine/testing/init.py

[Txy02]

Fixed the pre-commit EOF issue and re-ran local checks:

• pre-commit run --files rl_engine/testing/logprob_parity.py tests/test_logprob_parity.py rl_engine/testing/__init__.py
• python -m pytest tests/test_logprob_parity.py -v (9 passed)
• python -m pytest tests/test_reference_ops.py tests/test_logprob_parity.py -v (16 passed)

[Flink-ddd]

cc @maxiaosong1124 @a-kaa PTAL