feat(ws1):Add LogProb reference operator interface(reuse)

[a-kaa]

Summary

Add the reference-operator interface to NativeLogpOp
This keeps the existing selected-token log probability implementation and adds the unified baseline contract:

• forward(logits, token_ids)
• forward_fp32(logits, token_ids)
• op_class = "logprob"

The existing apply(...) and apply_fp32(...) APIs remain available as backward-
compatible aliases.

Implementation

• Updated rl_engine/kernels/ops/pytorch/loss/logp.py
• Added tests/test_logp.py
• Updated docs/operators/fused-logp.md

Summary by CodeRabbit

Release Notes

• New Features

  • Added dedicated forward(...) and forward_fp32(...) methods for the LogP operator with explicit dtype behavior.
  • Kept apply(...) / apply_fp32(...) as backward-compatible aliases of the new methods.

• Documentation

  • Updated Fused LogP operator docs with refreshed interface examples and generalized tensor shape/dtype contract details.
  • Revised test and implementation references in the documentation.

• Tests

  • Added a new NativeLogpOp test suite validating output shape/dtype, aliasing behavior, purity, registry dispatch, input validation, batch invariance, and numeric accuracy.

[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: 1bbd2933-97e6-4eee-8c0a-96014283bc5e

📥 Commits

Reviewing files that changed from the base of the PR and between 2915fb1 and 646c887.

📒 Files selected for processing (1)

• tests/test_logp.py

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

• tests/test_logp.py

---

📝 Walkthrough

Walkthrough

NativeLogpOp gains forward and forward_fp32 as its primary API methods delegating to _selected_logps, with __call__ rerouted through forward and apply/apply_fp32 reduced to backward-compatible aliases. A new test file covers correctness, aliasing, batch invariance, dtype accuracy, and registry dispatch. The operator documentation is updated accordingly.

Changes

NativeLogpOp API refactor, tests, and docs

Layer / File(s)	Summary
NativeLogpOp forward/apply refactor rl_engine/kernels/ops/pytorch/loss/logp.py	Adds op_class = "logprob", introduces forward and forward_fp32 as primary implementations (both delegate to _selected_logps with different output_dtype), rewires __call__ to call forward, and converts apply/apply_fp32 into thin aliases.
NativeLogpOp test suite tests/test_logp.py	Adds _make_inputs and _reference_selected_logp helpers; TestNativeLogpOpCorrectness verifies shape/dtype contracts, aliasing, bitwise reference equality, input non-mutation, and error handling; TestNativeLogpOpBatchInvariance checks per-row independence; TestNativeLogpOpAccuracy checks typed forward vs forward_fp32 tolerances; TestNativeLogpOpRegistry checks kernel registry dispatch.
Fused-logp documentation update docs/operators/fused-logp.md	Updates the usage snippet to show NativeLogpOp, generalizes the tensor contract from [N, V]/SM90-specific to [..., V], revises test commands to point at tests/test_logp.py, and updates the implementation files list.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~10 minutes

Suggested reviewers

• EthanZero2Hero
• inaniloquentee
• Flink-ddd

Poem

🐇 Hop, hop, forward leads the way,
apply just whispers "forward, hey!"
fp32 or native dtype true,
batch invariant through and through.
The registry knows where logprobs dwell —
tests confirm it all works well! 🌟

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 16.67% 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 'feat(ws1):Add LogProb reference operator interface(reuse)' directly describes the main objective of introducing a reference operator interface to NativeLogpOp, which aligns with the PR's core purpose of adding forward/forward_fp32 methods and setting op_class.
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.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ 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 and usage tips.}

[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 `@tests/test_logp.py`:
- Around line 96-99: The assert statements checking torch.equal results in the
test_logp.py file need to be reformatted to comply with Black's code formatting
standards. Locate the assert statements in the forward_fp32 test block (around
line 96) and at line 130-133 that check torch.equal(full_out[row],
single_out[0]) and reformat them according to Black's multi-line formatting
rules to ensure the parentheses, line breaks, and indentation match what Black
expects, which will unblock the linting CI.

🪄 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: 55f63297-1a32-43d5-bd32-272905316723

📥 Commits

Reviewing files that changed from the base of the PR and between 51b8b21 and 2915fb1.

📒 Files selected for processing (3)

• docs/operators/fused-logp.md
• rl_engine/kernels/ops/pytorch/loss/logp.py
• tests/test_logp.py