[WS2] Add TP-aware logprob contract and dispatch metadata

.github/workflows/ci.yml

@@ -76,6 +76,9 @@ 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
+
   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,19 @@ 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).
 
+### WS2 TP-aware dispatch
+
+WS2 distributed callers use a separate contract-aware entry point,
+`kernel_registry.get_logprob_op(contract)`. It validates explicit vocab-shard ownership,
+padded-vs-real vocab metadata, active-token masking, and fixed `(max, sumexp)` merge
+semantics before selecting a backend. Legacy `get_op("batch_invariant_logp")` behavior
+remains unchanged.
+
+The backends above are single-shard (TP=1) references and do not yet export vocab-domain
+LSE or carry vocab-shard metadata, so they are declared incompatible with strict WS2
+requests instead of being selected as a silent fallback. The contract objects are
+documented in `rl_engine.kernels.logprob_contract`.
+
 ## Benchmarks
 
 `benchmarks/benchmark_batch_invariant_logp.py` compares Native, Triton, and the