[WS2] feat: add TP=1 logprob comparison harness

docs/operators/batch-invariant-logp.md

@@ -179,6 +179,101 @@ fp16/bf16 backward: checked against fp32 reference with relaxed tolerance
 CPU-vs-CUDA comparisons use tolerance-based checks; batch-invariance checks
 within the same backend use exact equality where appropriate.
 
+## TP=1 Comparison Harness
+
+The single-GPU comparison harness is the TP=1 registration and regression guard
+for issue #241. It uses the batch-invariant PyTorch implementation as the
+reference and compares exact `pytorch`, `triton`, or `cuda-sm90` backends before
+distributed communication is introduced.
+
+Each backend exposes a diagnostic-only entry point while the production contract
+remains unchanged:
+
+```text
+op(logits, target_ids)                  -> logp
+op.forward_with_lse(logits, target_ids) -> (logp, lse)
+```
+
+The harness reports LSE drift over every logical token row and selected-logprob
+drift over active response/action tokens only. Drift summaries contain max,
+mean, p95, p99, and the number of compared values. Reports also record requested
+and actual backends, implementation, direct-LSE provenance, input shape and
+dtype, `tp_world=1`, and `communication=none`.
+
+Backend selection is exact and does not use registry fallback. In particular,
+an explicit `cuda-sm90` comparison fails unless the compiled SM90 extension,
+Hopper hardware, input dtype, and vocab row stride satisfy the kernel contract.
+
+Run the PyTorch TP=1 guard directly from the kernel-specific testing module:
+
+```bash
+python rl_engine/testing/logprob_comparison.py \
+  --candidate pytorch \
+  --device cpu \
+  --dtype fp32 \
+  --batch 2 \
+  --seq 16 \
+  --vocab 257
+```
+
+On a GPU, repeat `--candidate` to compare multiple exact backends:
+
+```bash
+python rl_engine/testing/logprob_comparison.py \
+  --candidate triton \
+  --candidate cuda-sm90 \
+  --device cuda \
+  --dtype bf16 \
+  --batch 2 \
+  --seq 16 \
+  --vocab 151936
+```
+
+The command writes structured JSON to stdout and routes backend diagnostics to
+stderr. The harness does not implement vocab sharding, collective communication,
+cross-rank LSE merging, or CP reconstruction.
+
+### SM90 validation
+
+SM90 validation requires a Hopper GPU, CUDA-enabled PyTorch, and an `nvcc`
+toolkit matching `torch.version.cuda`. Build the extension with:
+
+```bash
+export FORCE_CUDA=1
+export KERNEL_ALIGN_FORCE_SM90=1
+export TORCH_CUDA_ARCH_LIST="9.0+PTX"
+
+python -m pip install --no-build-isolation --no-deps -e .
+```
+
+Run the focused harness tests, the complete operator suite, and an explicit
+SM90 comparison:
+
+```bash
+python -m pytest \
+  tests/test_logprob_comparison.py \
+  tests/test_operator_inputs.py \
+  tests/test_op_checks.py -q
+
+python -m pytest tests/test_batch_invariant_logp.py -q
+
+python rl_engine/testing/logprob_comparison.py \
+  --candidate cuda-sm90 \
+  --device cuda \
+  --dtype bf16 \
+  --batch 2 \
+  --seq 16 \
+  --vocab 151936 \
+  --prompt-tokens 8 \
+  --seed 241
+```
+
+The PR2 path was validated on an NVIDIA H800 PCIe with PyTorch 2.11.0+cu128,
+CUDA 12.8, and Triton 3.6.0. The focused tests passed 41 cases and the complete
+batch-invariant suite passed 67 cases. For BF16 shape `[2, 16, 151936]`, both
+LSE and active-token dlogp had maximum absolute drift
+`9.5367431640625e-07` against the PyTorch reference, with no backend fallback.
+
 ## Minimal Example
 
 ```python
@@ -206,11 +301,14 @@ out.sum().backward()
 python -m pytest tests/test_batch_invariant_logp.py -q -rs
 ```
 
-All backends (Native, Triton) are tested in a single file. Coverage includes:
+All production backends are tested in a single file. Coverage includes
 correctness, leading-shape preservation, batch-invariance (bitwise), validation,
 ignore-index behavior, backward correctness, CUDA smoke cases, registry
 dispatch, and Triton-specific fp32/fp16/bf16 correctness, large vocab, backward
-gradient batch-invariance, and ignored-row zero gradients.
+gradient batch-invariance, and ignored-row zero gradients. The focused
+`tests/test_logprob_comparison.py` suite covers TP=1 bitwise regression, direct
+LSE identity, active-token drift statistics, structured serialization, exact
+backend diagnostics, and fail-closed provenance.
 
 Triton tests skip when Triton or CUDA is unavailable. On Windows, run via
 WSL/Linux with CUDA.
@@ -223,4 +321,5 @@ WSL/Linux with CUDA.
 - `csrc/cuda/batch_invariant_logp_kernel_sm90.cu`
 - `rl_engine/kernels/registry.py`
 - `tests/test_batch_invariant_logp.py`
+- `tests/test_logprob_comparison.py`
 - `benchmarks/benchmark_batch_invariant_logp.py`

rl_engine/kernels/ops/cuda/loss/batch_invariant_logp.py

@@ -161,3 +161,48 @@ def apply(
                 )
 
         return _BatchInvariantLogpSM90Function.apply(logits, target_ids, ignore_index)
+
+    def forward_with_lse(
+        self,
+        logits: torch.Tensor,
+        target_ids: torch.Tensor,
+        ignore_index: int = -100,
+        *,
+        validate: bool = True,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Run the exact SM90 path and return its direct FP32 logprob/LSE outputs.
+
+        Unlike the production ``apply`` method, this diagnostic entry point never
+        falls back to Triton or PyTorch, so comparison provenance stays truthful.
+        """
+        if logits.dim() < 2:
+            raise ValueError(
+                f"logits must be at least 2-D ([*lead, V]), got shape {tuple(logits.shape)}"
+            )
+        if logits.shape[:-1] != target_ids.shape:
+            raise ValueError(
+                f"logits leading shape {tuple(logits.shape[:-1])} must match "
+                f"target_ids shape {tuple(target_ids.shape)}"
+            )
+        if not _sm90_supported(logits):
+            raise RuntimeError(
+                "exact cuda-sm90 logprob diagnostics require Hopper, CUDA BF16/FP32 logits, "
+                "and a 16-byte-aligned vocab row stride; fallback is disabled"
+            )
+        if validate:
+            vocab_size = logits.size(-1)
+            valid_targets = target_ids.reshape(-1)
+            valid_targets = valid_targets[valid_targets != ignore_index]
+            if valid_targets.numel() and (
+                (valid_targets < 0).any() or (valid_targets >= vocab_size).any()
+            ):
+                bad = valid_targets[(valid_targets < 0) | (valid_targets >= vocab_size)]
+                raise ValueError(
+                    f"target_ids contains values outside [0, {vocab_size}): {bad.tolist()}"
+                )
+
+        lead_shape = logits.shape[:-1]
+        logits_2d = logits.reshape(-1, logits.size(-1)).contiguous()
+        target_1d = target_ids.reshape(-1).to(device=logits.device, dtype=torch.int64).contiguous()
+        logp, lse = _C.batch_invariant_logp_sm90(logits_2d, target_1d, int(ignore_index))
+        return logp.reshape(lead_shape), lse.reshape(lead_shape)

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

@@ -45,24 +45,42 @@ def apply(
         logits_2d = logits.reshape(-1, vocab_size).float()
         target_1d = target_ids.reshape(-1).to(logits.device, dtype=torch.long)
 
-        selected_logp = self._row_wise_selected_logprob(
+        selected_logp, _ = self._row_wise_selected_logprob_with_lse(
             logits_2d, target_1d, ignore_index=ignore_index, validate=validate
         )
 
         return selected_logp.reshape(lead_shape)
 
+    def forward_with_lse(
+        self,
+        logits: torch.Tensor,
+        target_ids: torch.Tensor,
+        ignore_index: int = -100,
+        *,
+        validate: bool = True,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Return selected logprob and the FP32 vocab-domain LSE for diagnostics."""
+        self._validate_shapes(logits, target_ids)
+        lead_shape = logits.shape[:-1]
+        logits_2d = logits.reshape(-1, logits.size(-1)).float()
+        target_1d = target_ids.reshape(-1).to(logits.device, dtype=torch.long)
+        logp, lse = self._row_wise_selected_logprob_with_lse(
+            logits_2d, target_1d, ignore_index=ignore_index, validate=validate
+        )
+        return logp.reshape(lead_shape), lse.reshape(lead_shape)
+
     # ---------------------------------------------------------------------- #
     # Core Computation
     # ---------------------------------------------------------------------- #
     @staticmethod
-    def _row_wise_selected_logprob(
+    def _row_wise_selected_logprob_with_lse(
         logits_2d: torch.Tensor,
         target_1d: torch.Tensor,
         *,
         ignore_index: int,
         validate: bool = True,
-    ) -> torch.Tensor:
-        """Per-row selected logprob with locked reduction order.
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Per-row selected logprob and LSE with locked reduction order.
 
         The three reduction steps (max, sum-exp, gather) operate on each row
         independently.  PyTorch's ``max(dim=-1)`` and ``sum(dim=-1)`` iterate
@@ -104,7 +122,7 @@ def _row_wise_selected_logprob(
 
         selected_logp = selected_logp.where(valid_mask, torch.zeros_like(selected_logp))
 
-        return selected_logp
+        return selected_logp, log_sum_exp
 
     # ---------------------------------------------------------------------- #
     # Helper

rl_engine/kernels/ops/triton/loss/batch_invariant_logp.py

@@ -10,6 +10,27 @@
 _BLOCK_V: int = 1024
 
 
+def _launch_batch_invariant_logp(
+    logits_2d: torch.Tensor, target_1d: torch.Tensor, ignore_index: int
+) -> tuple[torch.Tensor, torch.Tensor]:
+    num_tokens = logits_2d.shape[0]
+    vocab_size = logits_2d.shape[1]
+    output = torch.empty(num_tokens, device=logits_2d.device, dtype=torch.float32)
+    lse = torch.empty(num_tokens, device=logits_2d.device, dtype=torch.float32)
+    _batch_invariant_logp_kernel[(num_tokens,)](
+        logits_2d,
+        target_1d,
+        output,
+        lse,
+        num_tokens,
+        vocab_size,
+        logits_2d.stride(0),
+        ignore_index=ignore_index,
+        BLOCK_V=_BLOCK_V,
+    )
+    return output, lse
+
+
 @triton.jit
 def _batch_invariant_logp_kernel(
     logits_ptr,  # logits [N, V]
@@ -126,22 +147,7 @@ def forward(ctx, logits, target_ids, ignore_index):
         logits_2d = logits.reshape(-1, vocab_size).contiguous()
         target_1d = target_ids.reshape(-1).to(device=logits.device, dtype=torch.int64).contiguous()
 
-        num_tokens = logits_2d.shape[0]
-        output = torch.empty(num_tokens, device=logits.device, dtype=torch.float32)
-        lse = torch.empty(num_tokens, device=logits.device, dtype=torch.float32)
-
-        grid = (num_tokens,)
-        _batch_invariant_logp_kernel[grid](
-            logits_2d,
-            target_1d,
-            output,
-            lse,
-            num_tokens,
-            vocab_size,
-            logits_2d.stride(0),
-            ignore_index=ignore_index,
-            BLOCK_V=_BLOCK_V,
-        )
+        output, lse = _launch_batch_invariant_logp(logits_2d, target_1d, ignore_index)
 
         ctx.save_for_backward(logits_2d, target_1d, lse)
         ctx.ignore_index = ignore_index
@@ -237,3 +243,53 @@ def apply(
                 )
 
         return _BatchInvariantLogpFunction.apply(logits, target_ids, ignore_index)
+
+    def forward_with_lse(
+        self,
+        logits: torch.Tensor,
+        target_ids: torch.Tensor,
+        ignore_index: int = -100,
+        *,
+        validate: bool = True,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Return direct FP32 logprob/LSE outputs without an autograd wrapper."""
+        self._validate_inputs(logits, target_ids, ignore_index=ignore_index, validate=validate)
+        lead_shape = logits.shape[:-1]
+        logits_2d = logits.reshape(-1, logits.size(-1)).contiguous()
+        target_1d = target_ids.reshape(-1).to(device=logits.device, dtype=torch.int64).contiguous()
+        logp, lse = _launch_batch_invariant_logp(logits_2d, target_1d, ignore_index)
+        return logp.reshape(lead_shape), lse.reshape(lead_shape)
+
+    @staticmethod
+    def _validate_inputs(
+        logits: torch.Tensor,
+        target_ids: torch.Tensor,
+        *,
+        ignore_index: int,
+        validate: bool,
+    ) -> None:
+        if logits.device.type not in ("cuda", "xpu", "hip"):
+            raise RuntimeError(
+                "TritonBatchInvariantLogpOp requires a GPU tensor "
+                f"(CUDA / ROCm / XPU), got device '{logits.device}'."
+            )
+        if logits.dim() < 2:
+            raise ValueError(
+                f"logits must be at least 2-D ([*lead, V]), got shape {tuple(logits.shape)}"
+            )
+        if logits.shape[:-1] != target_ids.shape:
+            raise ValueError(
+                f"logits leading shape {tuple(logits.shape[:-1])} must match "
+                f"target_ids shape {tuple(target_ids.shape)}"
+            )
+        if validate:
+            vocab_size = logits.size(-1)
+            valid_targets = target_ids.reshape(-1)
+            valid_targets = valid_targets[valid_targets != ignore_index]
+            if valid_targets.numel() and (
+                (valid_targets < 0).any() or (valid_targets >= vocab_size).any()
+            ):
+                bad = valid_targets[(valid_targets < 0) | (valid_targets >= vocab_size)]
+                raise ValueError(
+                    f"target_ids contains values outside [0, {vocab_size}): {bad.tolist()}"
+                )

rl_engine/testing/__init__.py

@@ -3,6 +3,14 @@
 
 """Testing helpers for RL-shaped kernel validation."""
 
+from .logprob_comparison import (
+    LogprobBackendUnavailable,
+    LogprobCandidate,
+    LogprobComparisonInputs,
+    LogprobComparisonReport,
+    compare_single_gpu_logprob,
+    make_logprob_candidate,
+)
 from .reference_ops import (
     active_token_count,
     compute_policy_ratio,
@@ -15,10 +23,16 @@
 from .rl_batch import SyntheticRLKernelBatch, make_synthetic_rl_kernel_batch
 
 __all__ = [
+    "LogprobBackendUnavailable",
+    "LogprobCandidate",
+    "LogprobComparisonInputs",
+    "LogprobComparisonReport",
     "SyntheticRLKernelBatch",
     "active_token_count",
+    "compare_single_gpu_logprob",
     "compute_policy_ratio",
     "compute_reference_kl",
+    "make_logprob_candidate",
     "make_synthetic_rl_kernel_batch",
     "masked_mean",
     "masked_sum",