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

docs/operators/fused-logp.md

@@ -13,21 +13,36 @@ logp_op = kernel_registry.get_op("logp")
 output = logp_op(logits, token_ids)
 ```
 
+The PyTorch native reference also exposes the Issue #108 interface:
+
+```python
+from rl_engine.kernels.ops.pytorch.loss.logp import NativeLogpOp
+
+logp_ref = NativeLogpOp()
+output = logp_ref.forward(logits, token_ids)
+reference = logp_ref.forward_fp32(logits, token_ids)
+```
+
+`apply(...)` and `apply_fp32(...)` remain available as backward-compatible aliases.
+
 ## Backends
 
 | Backend | Wrapper | Native symbol | Notes |
 | --- | --- | --- | --- |
 | CUDA SM90 | `FusedLogpSM90Op` | `_C.fused_logp_sm90` | TMA-oriented path for Hopper-class GPUs. |
 | CUDA generic | `FusedLogpGenericOp` | `_C.fused_logp` | Generic compiled extension fallback. |
-| PyTorch native | `NativeOp` | None | Baseline fallback path. |
+| PyTorch native | `NativeLogpOp` | None | PyTorch baseline/reference path. |
 
 ## Tensor Contract
 
 | Argument | Shape | Dtype | Requirements |
 | --- | --- | --- | --- |
-| `logits` | `[N, V]` | `bfloat16` for SM90 path | Contiguous, on the target device. |
-| `token_ids` / `labels` | `[N]` | Converted to `int32` | Same logical device as `logits`. |
-| Output | `[N]` | Backend-defined tensor dtype | One selected log probability per row. |
+| `logits` | `[..., V]` | Floating point | Contiguous for fused CUDA paths; arbitrary leading dimensions. |
+| `token_ids` / `labels` | `[...]` | Integer | Must match `logits.shape[:-1]`. |
+| Output | `[...]` | See below | One selected log probability per row. |
+
+For `NativeLogpOp`, `forward(...)` returns the input dtype and `forward_fp32(...)`
+returns `torch.float32`.
 
 ## Reference Semantics
 
@@ -39,16 +54,20 @@ ref = torch.gather(ref, dim=-1, index=token_ids.unsqueeze(-1).long()).squeeze(-1
 ## Tests
 
 ```bash
-python tests/test_op_accuracy.py
+python -m pytest tests/test_logp.py -q
+python -m pytest tests/test_op_accuracy.py -q
 ```
 
-The current accuracy test compares the dispatched operator with a PyTorch reference and
-uses a dtype-dependent threshold.
+`tests/test_logp.py` covers the PyTorch reference contract, dtype behavior,
+backward-compatible aliases, batch invariance, and registry dispatch. The existing
+operator accuracy tests continue to validate native/CUDA fused API compatibility.
 
 ## Implementation Files
 
 - `rl_engine/kernels/registry.py`
-- `rl_engine/kernels/ops/cuda.py`
+- `rl_engine/kernels/ops/pytorch/loss/logp.py`
+- `rl_engine/kernels/ops/cuda/loss/logp.py`
 - `csrc/ops.cpp`
 - `csrc/fused_logp_kernel.cu`
 - `csrc/cuda/fused_logp_sm90.cu`
+- `tests/test_logp.py`

rl_engine/kernels/ops/pytorch/loss/logp.py

@@ -1,17 +1,21 @@
 # SPDX-License-Identifier: Apache-2.0
 # Copyright (c) 2026 RL-Kernel Contributors
 
+from __future__ import annotations
+
 import torch
 
 
 class NativeLogpOp:
     """Pure PyTorch native fallback for Fused LogP."""
 
-    def __init__(self):
+    op_class = "logprob"
+
+    def __init__(self) -> None:
         pass
 
     def __call__(self, logits: torch.Tensor, token_ids: torch.Tensor) -> torch.Tensor:
-        return self.apply(logits, token_ids)
+        return self.forward(logits, token_ids)
 
     def _selected_logps(
         self,
@@ -45,14 +49,22 @@ def _validate_output_shape(self, output: torch.Tensor, logits: torch.Tensor) ->
                 f"{tuple(logits.shape[:-1])}"
             )
 
-    def apply(self, logits: torch.Tensor, token_ids: torch.Tensor) -> torch.Tensor:
+    def forward(self, logits: torch.Tensor, token_ids: torch.Tensor) -> torch.Tensor:
         """Baseline selected-token log probability extraction using torch.gather."""
         return self._selected_logps(logits, token_ids, output_dtype=logits.dtype)
 
-    def apply_fp32(self, logits: torch.Tensor, token_ids: torch.Tensor) -> torch.Tensor:
+    def forward_fp32(self, logits: torch.Tensor, token_ids: torch.Tensor) -> torch.Tensor:
         """Same as apply but forces float32 output for numerical stability."""
         return self._selected_logps(logits, token_ids, output_dtype=torch.float32)
 
+    def apply(self, logits: torch.Tensor, token_ids: torch.Tensor) -> torch.Tensor:
+        """Backward-compatible alias for forward."""
+        return self.forward(logits, token_ids)
+
+    def apply_fp32(self, logits: torch.Tensor, token_ids: torch.Tensor) -> torch.Tensor:
+        """Backward-compatible alias for forward_fp32."""
+        return self.forward_fp32(logits, token_ids)
+
     def indexed_out(
         self,
         logits: torch.Tensor,

tests/test_logp.py

@@ -0,0 +1,144 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) 2026 RL-Kernel Contributors
+
+"""Tests for NativeLogpOp, the PyTorch selected-logprob reference."""
+
+from __future__ import annotations
+
+import pytest
+import torch
+
+from rl_engine.kernels.ops.pytorch.loss.logp import NativeLogpOp
+
+
+def _make_inputs(
+    batch: int,
+    seq: int,
+    vocab: int,
+    *,
+    dtype: torch.dtype = torch.float32,
+    seed: int = 123,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    gen = torch.Generator().manual_seed(seed)
+    logits = torch.randn(batch, seq, vocab, generator=gen, dtype=dtype)
+    token_ids = torch.randint(0, vocab, (batch, seq), generator=gen, dtype=torch.long)
+    return logits, token_ids
+
+
+def _reference_selected_logp(logits: torch.Tensor, token_ids: torch.Tensor) -> torch.Tensor:
+    log_probs = torch.log_softmax(logits.float(), dim=-1)
+    return torch.gather(log_probs, dim=-1, index=token_ids.long().unsqueeze(-1)).squeeze(-1)
+
+
+class TestNativeLogpOpCorrectness:
+    def test_output_shape_matches_token_ids(self):
+        op = NativeLogpOp()
+        logits, token_ids = _make_inputs(2, 16, 257)
+        out = op.forward_fp32(logits, token_ids)
+        assert out.shape == token_ids.shape
+
+    def test_forward_fp32_returns_fp32(self):
+        op = NativeLogpOp()
+        logits, token_ids = _make_inputs(2, 16, 257, dtype=torch.bfloat16)
+        out = op.forward_fp32(logits, token_ids)
+        assert out.dtype == torch.float32
+
+    @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16, torch.float16])
+    def test_forward_returns_input_dtype(self, dtype):
+        op = NativeLogpOp()
+        logits, token_ids = _make_inputs(2, 16, 257, dtype=dtype)
+        out = op.forward(logits, token_ids)
+        assert out.dtype == dtype
+
+    def test_call_and_apply_alias_forward(self):
+        op = NativeLogpOp()
+        logits, token_ids = _make_inputs(2, 16, 257)
+        forward = op.forward(logits, token_ids)
+        assert torch.equal(op(logits, token_ids), forward)
+        assert torch.equal(op.apply(logits, token_ids), forward)
+
+    def test_apply_fp32_alias_forward_fp32(self):
+        op = NativeLogpOp()
+        logits, token_ids = _make_inputs(2, 16, 257)
+        assert torch.equal(op.apply_fp32(logits, token_ids), op.forward_fp32(logits, token_ids))
+
+    def test_matches_fp32_reference_bitwise(self):
+        op = NativeLogpOp()
+        logits, token_ids = _make_inputs(2, 16, 257)
+        out = op.forward_fp32(logits, token_ids)
+        ref = _reference_selected_logp(logits, token_ids)
+        assert torch.equal(out, ref)
+
+    def test_pure_function_no_inplace(self):
+        op = NativeLogpOp()
+        logits, token_ids = _make_inputs(2, 16, 257)
+        logits_orig = logits.clone()
+        token_ids_orig = token_ids.clone()
+        _ = op.forward_fp32(logits, token_ids)
+        assert torch.equal(logits, logits_orig)
+        assert torch.equal(token_ids, token_ids_orig)
+
+    def test_op_class_is_logprob(self):
+        assert NativeLogpOp.op_class == "logprob"
+
+    def test_rejects_mismatched_shapes(self):
+        op = NativeLogpOp()
+        logits = torch.randn(2, 3, 5)
+        token_ids = torch.randint(0, 5, (2, 4))
+        with pytest.raises(ValueError, match="must match"):
+            op.forward_fp32(logits, token_ids)
+
+
+class TestNativeLogpOpBatchInvariance:
+    def test_batch1_vs_batchN_bitwise(self):
+        op = NativeLogpOp()
+        logits, token_ids = _make_inputs(4, 16, 257, seed=321)
+        full_out = op.forward_fp32(logits, token_ids)
+        for row in range(logits.shape[0]):
+            single_out = op.forward_fp32(logits[row : row + 1], token_ids[row : row + 1])
+            assert torch.equal(
+                full_out[row], single_out[0]
+            ), f"Batch invariance broken at row {row}"
+
+    def test_batch_invariance_with_padding(self):
+        op = NativeLogpOp()
+        logits_valid, token_ids_valid = _make_inputs(2, 16, 257, seed=456)
+        gen = torch.Generator().manual_seed(789)
+        logits_padding = torch.randn(3, 16, 257, generator=gen)
+        token_padding = torch.randint(0, 257, (3, 16), generator=gen)
+        logits_padded = torch.cat([logits_valid, logits_padding], dim=0)
+        token_ids_padded = torch.cat([token_ids_valid, token_padding], dim=0)
+
+        out_valid = op.forward_fp32(logits_valid, token_ids_valid)
+        out_padded = op.forward_fp32(logits_padded, token_ids_padded)
+        assert torch.equal(out_valid[0], out_padded[0])
+        assert torch.equal(out_valid[1], out_padded[1])
+
+
+class TestNativeLogpOpAccuracy:
+    @pytest.mark.parametrize(
+        "dtype, atol",
+        [
+            (torch.float32, 1e-5),
+            (torch.bfloat16, 2e-2),
+            (torch.float16, 5e-3),
+        ],
+    )
+    def test_forward_vs_fp32_within_tolerance(self, dtype, atol):
+        op = NativeLogpOp()
+        logits, token_ids = _make_inputs(2, 16, 17, dtype=dtype)
+        out_typed = op.forward(logits, token_ids).float()
+        out_fp32 = op.forward_fp32(logits, token_ids)
+        diff = (out_typed - out_fp32).abs().max().item()
+        assert torch.allclose(
+            out_typed, out_fp32, atol=atol, rtol=0.0
+        ), f"dtype={dtype}, max_abs_error={diff:.3e} exceeds atol={atol}"
+
+
+class TestNativeLogpOpRegistry:
+    @pytest.mark.skipif(torch.cuda.is_available(), reason="CUDA dispatch may select fused logp")
+    def test_registry_returns_logp_op(self):
+        from rl_engine.kernels.registry import kernel_registry
+
+        op = kernel_registry.get_op("logp")
+        assert isinstance(op, NativeLogpOp)