[WS2][PR4][Logp] feat: add distributed TP/CP logprob drift report

.github/workflows/ci.yml

@@ -76,6 +76,17 @@ jobs:
         run: |
           python -m pytest tests/test_kv_cache_attention.py -v -k "not large and not gpu"
 
+      - name: Run WS2 Logprob Contract Tests (CPU-safe)
+        run: python -m pytest tests/test_logprob_contract.py -v
+
+      - name: Run WS2 Vocab-Parallel Logprob Tests (CPU-safe)
+        run: python -m pytest tests/test_vocab_parallel_logp.py -v
+
+      - name: Run WS2 Logprob Comparison Tests (CPU-safe)
+        run: |
+          python -m pytest tests/test_logprob_comparison.py -v
+          python -m pytest tests/test_distributed_logprob_comparison.py -v
+
   docs:
     runs-on: ubuntu-latest
     steps:

docs/design/runtime-dispatch.md

@@ -11,6 +11,14 @@ logical type, and the registry selects the first available backend for the curre
 4. Cache successfully constructed operator instances.
 5. Skip backends that already failed in the current process.
 
+WS2 TP-aware logprob uses the stricter `KernelRegistry.get_logprob_op(contract)` path. In
+addition to platform priority, this path requires a backend capability descriptor and checks
+the requested role, dtype, TP/CP layout, padded-vs-real vocab masking, inactive-token
+support, vocab-domain LSE export, and deterministic TP merge semantics. Incompatible
+candidates produce explicit rejection reasons and are never used as an undeclared fallback.
+The contract objects and their normative reduction semantics are documented in
+`rl_engine.kernels.logprob_contract`.
+
 ## LogP Priority
 
 | Platform | Priority |

docs/operators/batch-invariant-logp.md

@@ -54,6 +54,35 @@ CUDA priority list when the extension exposes `_C.batch_invariant_logp_sm90`
 (built with `KERNEL_ALIGN_FORCE_SM90=1`) on an SM90 device. On any other build
 or device, dispatch is unchanged (Triton -> PyTorch).
 
+## Tensor Parallel
+
+`VocabParallelLogprobOp`
+(`rl_engine/kernels/ops/pytorch/loss/vocab_parallel_logp.py`)
+defines a cross-TP bitwise contract for TP=1, TP=2, and TP=4 when
+`num_vocab_tiles` is fixed and every vocabulary-shard boundary is tile-aligned.
+The complete BF16 CUDA/NCCL validation matrix for this contract is tracked by
+issue #241 PR4.
+
+1. Split the padded vocabulary into `num_vocab_tiles` fixed tiles.
+2. Each rank computes fp32 `(max, sumexp)` for the tiles it owns. Every tile
+   is reduced as the same contiguous `[n, tile]` shape, on any rank.
+3. All tile partials are shared with `all_gather`. The collective only moves
+   bytes; it never does math, so it cannot round anything.
+4. Every rank merges all tiles in the same fixed order, over the same
+   `[n, num_vocab_tiles]` shape. `LSE = M + log(sum(s_t * exp(m_t - M)))`.
+5. The target logit is copied from the rank that owns it (never summed).
+6. `logp = target_logit - LSE`. Inactive rows become `0.0`.
+
+Usage goes through the contract-aware entry point:
+
+```python
+from rl_engine.kernels.registry import kernel_registry
+
+result = kernel_registry.get_logprob_op(contract)   # LogprobContract from
+op = result.op                                      # rl_engine.kernels.logprob_contract
+logp, lse = op(local_logits, target_ids, contract=contract, tp_group=tp_group)
+```
+
 ## Benchmarks
 
 `benchmarks/benchmark_batch_invariant_logp.py` compares Native, Triton, and the
@@ -179,6 +208,160 @@ 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.
+
+## Distributed WS2 Drift Report
+
+The issue #241 PR4 runner materializes one TP/CP topology per `torchrun`
+invocation. TP partitions the vocabulary and is the only numerical merge axis;
+CP partitions token rows and is recorded in provenance without participating in
+the vocab-domain LSE merge. For global rank `r`:
+
+```text
+tp_rank = r % tp_world_size
+cp_rank = r // tp_world_size
+```
+
+Every case generates the same seeded FP32 logical logits, targets, and active
+mask. The candidate receives a BF16 token/vocab shard through the explicit
+`pytorch-vocab-parallel-logp-ws2` backend, while the independent oracle computes
+`torch.logsumexp` over the complete real-vocab FP32 token slice. Distributed
+dispatch rejects `auto`, capability fallback, topology mismatches, non-tileable
+vocabularies, and incomplete materialization.
+
+Reports follow the issue #116 fields and contain per-rank and aggregate LSE and
+active-token dlogp summaries: max/mean/p95/p99 absolute drift, max relative
+drift, worst global token position, target id, target owner rank, #108 tolerance,
+and pass/fail. Provenance includes TP/CP topology, dtype, shard bounds, backend
+capability, contract fingerprint, reduction spec, merge order, transport, and
+the exact launch command. Replicated TP outputs are checked bitwise before one
+representative per CP shard is included in aggregate statistics.
+
+Print the scoped TP=1/2/4 x CP=1/2 launch matrix without starting workers:
+
+```bash
+python rl_engine/testing/distributed_logprob_comparison.py \
+  --plan \
+  --device cuda \
+  --dtype bf16 \
+  --output artifacts/ws2-logprob/report.json
+```
+
+Run one TP=2, CP=2 Qwen3-vocab case on four local GPUs:
+
+```bash
+torchrun --standalone --nproc-per-node=4 \
+  rl_engine/testing/distributed_logprob_comparison.py \
+  --tp 2 \
+  --cp 2 \
+  --dtype bf16 \
+  --backend pytorch-vocab-parallel-logp-ws2 \
+  --real-vocab 151936 \
+  --padded-vocab 151936 \
+  --num-vocab-tiles 64 \
+  --batch 2 \
+  --seq 16 \
+  --prompt-tokens 8 \
+  --output artifacts/ws2-logprob/tp2-cp2.json
+```
+
+The full matrix requires up to eight ranks for TP=4, CP=2. CPU/Gloo cases are
+available for topology and artifact validation; the scoped numerical gate is
+BF16 on CUDA/NCCL.
+
 ## Minimal Example
 
 ```python
@@ -206,11 +389,17 @@ 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.
+`tests/test_distributed_logprob_comparison.py` covers topology planning, TP/CP
+rank mapping, token/vocab sharding, explicit backend materialization, #116 JSON
+artifacts, and a real four-process TP=2, CP=2 Gloo smoke case.
 
 Triton tests skip when Triton or CUDA is unavailable. On Windows, run via
 WSL/Linux with CUDA.
@@ -223,4 +412,11 @@ 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`
+- `rl_engine/testing/logprob_drift.py`
+- `rl_engine/testing/distributed_logprob_comparison.py`
+- `tests/test_distributed_logprob_comparison.py`
 - `benchmarks/benchmark_batch_invariant_logp.py`
+- `rl_engine/kernels/ops/pytorch/loss/vocab_parallel_logp.py`
+- `rl_engine/kernels/logprob_contract.py`
+- `tests/test_vocab_parallel_logp.py`