[WS1][kernels] Batch-invariant logprob (Native, Triton)

.gitignore

@@ -206,3 +206,6 @@ cython_debug/
 marimo/_static/
 marimo/_lsp/
 __marimo__/
+
+# Local dev notes (not for upstream)
+_dev_notes/

docs/.nav.yml

@@ -15,6 +15,7 @@ nav:
       - operators/attention.md
       - operators/fused-logp.md
       - operators/linear-logp.md
+      - operators/batch-invariant-logp.md
       - operators/linear-logp-tp-test.md
       - operators/grpo-loss.md
       - operators/lm_head.md

docs/operators/README.md

@@ -22,6 +22,7 @@ Every operator page should include:
 - [Standard Attention](attention.md)
 - [Fused LogP](fused-logp.md)
 - [Fused Linear LogP](linear-logp.md)
+- [Batch-Invariant LogP](batch-invariant-logp.md)
 - [Fused Linear LogP TP Test Runbook](linear-logp-tp-test.md)
 - [GRPO Loss](grpo-loss.md)
 - [RoPE](rope.md)

docs/operators/batch-invariant-logp.md

@@ -0,0 +1,174 @@
+# Batch-Invariant LogP
+
+Batch-Invariant LogP computes selected token log-probabilities from already
+materialized logits:
+
+```text
+out[row] = logits[row, target_ids[row]] - logsumexp(logits[row, :])
+```
+
+It targets RL post-training paths where policy log-probs are compared across
+different packing, padding, and batch layouts. The key contract is
+batch-invariance: for a fixed row of logits and target id, the result must not
+change when that row is evaluated alone, at a different batch position, or with
+different neighboring rows.
+
+Unlike `linear_logp`, this operator does not fuse the LM-head projection. It
+takes `[*, V]` logits as input and returns one selected log-probability per row.
+
+## Entry Point
+
+```python
+from rl_engine.kernels.registry import kernel_registry
+
+batch_invariant_logp = kernel_registry.get_op("batch_invariant_logp")
+
+logp = batch_invariant_logp(
+    logits,       # [B, T, V] or [N, V], differentiable
+    target_ids,   # [B, T] or [N], int
+    ignore_index=-100,
+    validate=False,  # Triton fast path; use True to debug-check target range
+)                # -> [B, T] or [N], float32
+
+logp.sum().backward()  # gradients flow into logits only
+```
+
+## Backends
+
+| Backend | Wrapper | Status |
+| --- | --- | --- |
+| CUDA / ROCm (Triton) | `TritonBatchInvariantLogpOp` | Triton online-softmax forward and tile-wise backward. Requires a GPU tensor. |
+| PyTorch native | `NativeBatchInvariantLogpOp` | FP32 reference path; CPU fallback and Triton-less fallback. |
+
+Current dispatch:
+
+```text
+CUDA / ROCm: Triton -> PyTorch
+CPU:         PyTorch
+```
+
+A compiled CUDA backend and benchmark suite are planned follow-up work.
+Benchmarks are not included in this PR; they will be added alongside the CUDA
+backend in a subsequent PR.
+
+## Tensor Contract
+
+| Argument | Shape | Dtype | Requirements |
+| --- | --- | --- | --- |
+| `logits` | `[N, V]` / `[B, T, V]` / `[*lead, V]` | fp32 / fp16 / bf16 | Differentiable input; last dimension is vocab. |
+| `target_ids` | `[N]` / `[B, T]` / `[*lead]` | int | Same leading shape as `logits`; non-ignored values in `[0, V)`. |
+| `ignore_index` | scalar int | Python int | Default `-100`. Ignored rows output zero and receive zero gradient. |
+| Output | `[N]` / `[B, T]` / `[*lead]` | float32 | Selected log-probability per row. |
+
+`target_ids` is integer and non-differentiable. Gradients flow only into
+`logits`.
+
+## Reference Semantics
+
+For non-ignored rows:
+
+```python
+logits_2d = logits.reshape(-1, logits.size(-1)).float()
+target_1d = target_ids.reshape(-1).long()
+
+log_probs = torch.log_softmax(logits_2d, dim=-1)
+selected = torch.gather(
+    log_probs,
+    dim=-1,
+    index=target_1d.unsqueeze(-1),
+).squeeze(-1)
+
+out = selected.reshape(target_ids.shape)
+```
+
+For ignored rows:
+
+```text
+target_ids[row] == ignore_index
+out[row] = 0.0
+grad_logits[row, :] = 0.0
+```
+
+Non-ignored target ids outside `[0, V)` are invalid. In particular,
+`target=-1` is invalid unless `ignore_index=-1`.
+
+The PyTorch native backend validates target ranges by default. The Triton
+backend defaults to `validate=False` to avoid CUDA stream synchronization in
+training hot paths. Use `validate=True` during debugging or in tests when
+calling the Triton backend with untrusted targets.
+
+## Batch-Invariance
+
+The operator is designed so each row is computed independently:
+
+- The PyTorch path reshapes to `[N, V]` and applies row-wise reductions.
+- The Triton forward uses `grid=(num_tokens,)`, so one program owns exactly one
+  row.
+- Triton vocab traversal uses a fixed `_BLOCK_V=1024` and does not autotune by
+  batch size.
+- Triton forward scans vocab tiles left-to-right using online logsumexp.
+- Triton backward uses `grid=(num_tokens, vocab_tiles)` and writes one row tile
+  per program. It reuses the forward-saved per-row `lse`, so no backward
+  reduction crosses row boundaries.
+- No atomic writes are used.
+
+These constraints ensure the result for a row depends only on that row's logits
+and target id, not on batch size, row position, or neighboring rows.
+
+## Accuracy
+
+Both backends accumulate reductions in float32 and return float32 outputs. Tests
+compare against `torch.log_softmax(...).gather(...)` with dtype-appropriate
+tolerances:
+
+```text
+fp32 forward: atol around 1e-5
+fp16/bf16 forward: atol around 1e-4
+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.
+
+## Minimal Example
+
+```python
+import torch
+
+from rl_engine.kernels.registry import kernel_registry
+
+op = kernel_registry.get_op("batch_invariant_logp")
+
+logits = torch.randn(2, 4, 300, device="cuda", dtype=torch.bfloat16)
+target_ids = torch.randint(0, 300, (2, 4), device="cuda")
+target_ids[0, 0] = -100
+
+out = op(logits, target_ids, ignore_index=-100)
+assert out.shape == target_ids.shape
+assert out.dtype == torch.float32
+assert out[0, 0].item() == 0.0
+
+out.sum().backward()
+```
+
+## Tests
+
+```bash
+python -m pytest tests/test_batch_invariant_logp.py -q -rs
+```
+
+All backends (Native, Triton) 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.
+
+Triton tests skip when Triton or CUDA is unavailable. On Windows, run via
+WSL/Linux with CUDA.
+
+## Implementation Files
+
+- `rl_engine/kernels/ops/pytorch/loss/batch_invariant_logp.py`
+- `rl_engine/kernels/ops/triton/loss/batch_invariant_logp.py`
+- `rl_engine/kernels/registry.py`
+- `tests/test_batch_invariant_logp.py`

rl_engine/kernels/gtest/op_checks.py

@@ -154,9 +154,7 @@ def _run_candidate(
     else:
         case_checks = [_run_case(candidate, case, contract) for case in cases]
     total_outputs = sum(len(case.outputs) for case in case_checks)
-    passed_outputs = sum(
-        1 for case in case_checks for output in case.outputs if output.passed
-    )
+    passed_outputs = sum(1 for case in case_checks for output in case.outputs if output.passed)
     pass_rate = float(passed_outputs / total_outputs) if total_outputs else 0.0
     return CandidateReport(
         candidate_name=candidate.name,
@@ -326,9 +324,7 @@ def _backward_grads(
     grad_outputs: list[torch.Tensor],
 ) -> list[torch.Tensor]:
     if len(outputs) != len(grad_outputs):
-        raise ValueError(
-            f"got {len(grad_outputs)} upstream gradients for {len(outputs)} outputs"
-        )
+        raise ValueError(f"got {len(grad_outputs)} upstream gradients for {len(outputs)} outputs")
     # `ones` makes this equivalent to output.sum().backward(); `random` tests a
     # stricter vector-Jacobian product.
     loss = sum(

rl_engine/kernels/gtest/operator_inputs.py

@@ -8,7 +8,6 @@
 
 import torch
 
-
 DEFAULT_HIDDEN = 4096
 DEFAULT_N_HEADS = 32
 DEFAULT_N_KV_HEADS = 8
@@ -95,9 +94,15 @@ def _make_attention_inputs(
 ) -> dict[str, Any]:
     batch, seq = _batch_seq(args)
     return {
-        "q": _floating_tensor((batch, DEFAULT_N_HEADS, seq, DEFAULT_HEAD_DIM), args, dtype, device, 0),
-        "k": _floating_tensor((batch, DEFAULT_N_KV_HEADS, seq, DEFAULT_HEAD_DIM), args, dtype, device, 1),
-        "v": _floating_tensor((batch, DEFAULT_N_KV_HEADS, seq, DEFAULT_HEAD_DIM), args, dtype, device, 2),
+        "q": _floating_tensor(
+            (batch, DEFAULT_N_HEADS, seq, DEFAULT_HEAD_DIM), args, dtype, device, 0
+        ),
+        "k": _floating_tensor(
+            (batch, DEFAULT_N_KV_HEADS, seq, DEFAULT_HEAD_DIM), args, dtype, device, 1
+        ),
+        "v": _floating_tensor(
+            (batch, DEFAULT_N_KV_HEADS, seq, DEFAULT_HEAD_DIM), args, dtype, device, 2
+        ),
         "causal": True,
     }
 
@@ -132,7 +137,9 @@ def _make_rope_inputs(
 ) -> dict[str, Any]:
     batch, seq = _batch_seq(args)
     return {
-        "x": _floating_tensor((batch, DEFAULT_N_HEADS, seq, DEFAULT_HEAD_DIM), args, dtype, device, 0),
+        "x": _floating_tensor(
+            (batch, DEFAULT_N_HEADS, seq, DEFAULT_HEAD_DIM), args, dtype, device, 0
+        ),
         "positions": torch.arange(seq, device=device, dtype=torch.long),
         "theta": _arg_float(args, "theta", DEFAULT_ROPE_THETA),
     }
@@ -185,11 +192,21 @@ def _make_kv_cache_attention_inputs(
 ) -> dict[str, Any]:
     batch, seq = _batch_seq(args)
     return {
-        "q": _floating_tensor((batch, DEFAULT_N_HEADS, 1, DEFAULT_HEAD_DIM), args, dtype, device, 0),
-        "k_cache": _floating_tensor((batch, DEFAULT_N_KV_HEADS, seq, DEFAULT_HEAD_DIM), args, dtype, device, 1),
-        "v_cache": _floating_tensor((batch, DEFAULT_N_KV_HEADS, seq, DEFAULT_HEAD_DIM), args, dtype, device, 2),
-        "k_new": _floating_tensor((batch, DEFAULT_N_KV_HEADS, 1, DEFAULT_HEAD_DIM), args, dtype, device, 3),
-        "v_new": _floating_tensor((batch, DEFAULT_N_KV_HEADS, 1, DEFAULT_HEAD_DIM), args, dtype, device, 4),
+        "q": _floating_tensor(
+            (batch, DEFAULT_N_HEADS, 1, DEFAULT_HEAD_DIM), args, dtype, device, 0
+        ),
+        "k_cache": _floating_tensor(
+            (batch, DEFAULT_N_KV_HEADS, seq, DEFAULT_HEAD_DIM), args, dtype, device, 1
+        ),
+        "v_cache": _floating_tensor(
+            (batch, DEFAULT_N_KV_HEADS, seq, DEFAULT_HEAD_DIM), args, dtype, device, 2
+        ),
+        "k_new": _floating_tensor(
+            (batch, DEFAULT_N_KV_HEADS, 1, DEFAULT_HEAD_DIM), args, dtype, device, 3
+        ),
+        "v_new": _floating_tensor(
+            (batch, DEFAULT_N_KV_HEADS, 1, DEFAULT_HEAD_DIM), args, dtype, device, 4
+        ),
         "causal": True,
     }
 
@@ -201,7 +218,7 @@ def _floating_tensor(
     device: torch.device,
     offset: int,
 ) -> torch.Tensor:
-    # Example: torch.randn((B, S, V), device="cuda", dtype=torch.bfloat16) 
+    # Example: torch.randn((B, S, V), device="cuda", dtype=torch.bfloat16)
     mode = _arg_str(args, "input_mode", "random")
     if mode == "constant":
         value = _arg_float(args, "constant_value", 0.25) + float(offset) * 0.01

rl_engine/kernels/gtest/operator_specs.py

@@ -10,8 +10,8 @@
 
 import torch
 
-from rl_engine.kernels.gtest.operator_inputs import make_operator_inputs, operator_shape_name
 from rl_engine.kernels.gtest.op_checks import CandidateSpec, OperatorCase
+from rl_engine.kernels.gtest.operator_inputs import make_operator_inputs, operator_shape_name
 
 
 @dataclass(frozen=True)

rl_engine/kernels/gtest/tolerance.py

@@ -7,7 +7,6 @@
 from pathlib import Path
 from typing import Any
 
-
 _CONTRACT_PATH = Path(__file__).with_name("tolerance_contract.json")