From 2915fb19657ccb9bf4cffc48f2a1c2bbc26e5f7e Mon Sep 17 00:00:00 2001 From: a-kaa Date: Mon, 22 Jun 2026 23:48:22 +0800 Subject: [PATCH 1/2] Add LogProb reference operator interface(reuse) --- docs/operators/fused-logp.md | 35 +++-- rl_engine/kernels/ops/pytorch/loss/logp.py | 20 ++- tests/test_logp.py | 144 +++++++++++++++++++++ 3 files changed, 187 insertions(+), 12 deletions(-) create mode 100644 tests/test_logp.py diff --git a/docs/operators/fused-logp.md b/docs/operators/fused-logp.md index d5008e5c5..ccc0f61ad 100644 --- a/docs/operators/fused-logp.md +++ b/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` diff --git a/rl_engine/kernels/ops/pytorch/loss/logp.py b/rl_engine/kernels/ops/pytorch/loss/logp.py index c79192766..4ef1f2430 100644 --- a/rl_engine/kernels/ops/pytorch/loss/logp.py +++ b/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, diff --git a/tests/test_logp.py b/tests/test_logp.py new file mode 100644 index 000000000..48bb950f3 --- /dev/null +++ b/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) From 646c887c24597a52275e634be0eb3e722bc73df9 Mon Sep 17 00:00:00 2001 From: a-kaa Date: Tue, 23 Jun 2026 00:00:57 +0800 Subject: [PATCH 2/2] format with pre-commit --- tests/test_logp.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/tests/test_logp.py b/tests/test_logp.py index 48bb950f3..42492fbbb 100644 --- a/tests/test_logp.py +++ b/tests/test_logp.py @@ -96,9 +96,9 @@ def test_batch1_vs_batchN_bitwise(self): 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}" - ) + 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() @@ -130,9 +130,9 @@ def test_forward_vs_fp32_within_tolerance(self, dtype, atol): 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}" - ) + 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: