[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:

README.md

@@ -117,8 +117,12 @@ RL-Kernel sits between high-level alignment libraries and low-level GPU kernels,
 git clone https://github.com/RL-Align/RL-Kernel.git
 cd RL-Kernel
 
-# Install core dependencies (CUDA 12.4+ recommended)
-pip install -e .
+# CPU-only / pure-Python fallback
+python -m pip install -e .
+
+# Native CUDA or ROCm extension (install a matching PyTorch build first)
+RL_KERNEL_REQUIRE_EXT=1 python -m pip install --no-build-isolation -e .
+python -c "import rl_engine._C as _C; assert hasattr(_C, 'fused_logp'); print(_C.__file__)"
 ```
 
 ### Contributions

csrc/deterministic_logp_kernel.cu

@@ -15,6 +15,15 @@ constexpr int kDeterministicLogpMediumVocabLimit = 4096;
 constexpr int kDeterministicLogpWarpSize = 32;
 constexpr float kDeterministicLogpNegInf = -3.4028234663852886e38F;
 
+template <typename T>
+__device__ __forceinline__ T deterministic_logp_shfl_down_32(T value, unsigned int delta) {
+#if defined(__HIPCC__) || defined(__HIP_PLATFORM_AMD__)
+    return __shfl_down(value, delta, kDeterministicLogpWarpSize);
+#else
+    return __shfl_down_sync(0xffffffffu, value, delta, kDeterministicLogpWarpSize);
+#endif
+}
+
 template <int BlockSize>
 struct DeterministicLogpBlockTraits {
     static_assert(
@@ -36,7 +45,7 @@ __device__ __forceinline__ float deterministicBlockReduceMax(float val) {
 
 #pragma unroll
     for (int offset = 16; offset > 0; offset >>= 1) {
-        val = fmaxf(val, __shfl_down_sync(0xffffffff, val, offset));
+        val = fmaxf(val, deterministic_logp_shfl_down_32(val, offset));
     }
 
     if (lane == 0) {
@@ -50,7 +59,7 @@ __device__ __forceinline__ float deterministicBlockReduceMax(float val) {
     if (wid == 0) {
 #pragma unroll
         for (int offset = 16; offset > 0; offset >>= 1) {
-            val = fmaxf(val, __shfl_down_sync(0xffffffff, val, offset));
+            val = fmaxf(val, deterministic_logp_shfl_down_32(val, offset));
         }
     }
     return val;
@@ -66,7 +75,7 @@ __device__ __forceinline__ float deterministicBlockReduceSum(float val) {
 
 #pragma unroll
     for (int offset = 16; offset > 0; offset >>= 1) {
-        val += __shfl_down_sync(0xffffffff, val, offset);
+        val += deterministic_logp_shfl_down_32(val, offset);
     }
 
     if (lane == 0) {
@@ -80,7 +89,7 @@ __device__ __forceinline__ float deterministicBlockReduceSum(float val) {
     if (wid == 0) {
 #pragma unroll
         for (int offset = 16; offset > 0; offset >>= 1) {
-            val += __shfl_down_sync(0xffffffff, val, offset);
+            val += deterministic_logp_shfl_down_32(val, offset);
         }
     }
     return val;

csrc/fused_logp_kernel.cu

@@ -5,6 +5,17 @@
 #include <limits>
 #include <torch/extension.h>
 
+constexpr int kFusedLogpLogicalWarpSize = 32;
+
+template <typename T>
+__device__ __forceinline__ T fused_logp_shfl_down_32(T value, unsigned int delta) {
+#if defined(__HIPCC__) || defined(__HIP_PLATFORM_AMD__)
+    return __shfl_down(value, delta, kFusedLogpLogicalWarpSize);
+#else
+    return __shfl_down_sync(0xffffffffu, value, delta, kFusedLogpLogicalWarpSize);
+#endif
+}
+
 template <typename scalar_t>
 __device__ __forceinline__ scalar_t blockReduceMax(scalar_t val) {
     static __shared__ float shared[32];
@@ -14,15 +25,15 @@ __device__ __forceinline__ scalar_t blockReduceMax(scalar_t val) {
     float f_val = static_cast<float>(val);
 
     for (int offset = 16; offset > 0; offset /= 2)
-        f_val = max(f_val, __shfl_down_sync(0xffffffff, f_val, offset));
+        f_val = max(f_val, fused_logp_shfl_down_32(f_val, offset));
 
     if (lane == 0) shared[wid] = f_val;
     __syncthreads();
 
     f_val = (threadIdx.x < blockDim.x / 32) ? shared[lane] : -1e20f;
     if (wid == 0) {
         for (int offset = 16; offset > 0; offset /= 2)
-            f_val = max(f_val, __shfl_down_sync(0xffffffff, f_val, offset));
+            f_val = max(f_val, fused_logp_shfl_down_32(f_val, offset));
     }
     return static_cast<scalar_t>(f_val);
 }
@@ -36,15 +47,15 @@ __device__ __forceinline__ scalar_t blockReduceSum(scalar_t val) {
     float f_val = static_cast<float>(val);
 
     for (int offset = 16; offset > 0; offset /= 2)
-        f_val += __shfl_down_sync(0xffffffff, f_val, offset);
+        f_val += fused_logp_shfl_down_32(f_val, offset);
 
     if (lane == 0) shared[wid] = f_val;
     __syncthreads();
 
     f_val = (threadIdx.x < blockDim.x / 32) ? shared[lane] : 0.0f;
     if (wid == 0) {
         for (int offset = 16; offset > 0; offset /= 2)
-            f_val += __shfl_down_sync(0xffffffff, f_val, offset);
+            f_val += fused_logp_shfl_down_32(f_val, offset);
     }
     return static_cast<scalar_t>(f_val);
 }
@@ -82,8 +93,8 @@ __device__ __forceinline__ LogSumExpState blockReduceLogSumExp(LogSumExpState st
 
     for (int offset = 16; offset > 0; offset /= 2) {
         LogSumExpState other{
-            __shfl_down_sync(0xffffffff, state.max_val, offset),
-            __shfl_down_sync(0xffffffff, state.sum_exp, offset)};
+            fused_logp_shfl_down_32(state.max_val, offset),
+            fused_logp_shfl_down_32(state.sum_exp, offset)};
         state = merge_logsumexp_state(state, other);
     }
 
@@ -100,8 +111,8 @@ __device__ __forceinline__ LogSumExpState blockReduceLogSumExp(LogSumExpState st
     if (wid == 0) {
         for (int offset = 16; offset > 0; offset /= 2) {
             LogSumExpState other{
-                __shfl_down_sync(0xffffffff, state.max_val, offset),
-                __shfl_down_sync(0xffffffff, state.sum_exp, offset)};
+                fused_logp_shfl_down_32(state.max_val, offset),
+                fused_logp_shfl_down_32(state.sum_exp, offset)};
             state = merge_logsumexp_state(state, other);
         }
     }

csrc/ops.cpp

@@ -253,6 +253,7 @@ std::vector<torch::Tensor> deterministic_attention_backward(
 
 // Prefix-Shared Attention Declarations & Wrappers
 
+#if !defined(USE_ROCM)
 void prefix_shared_attention_forward(
   const __nv_bfloat16 *Q,  // [bs, G, len_q, DIM]
   const __nv_bfloat16 *K,  // [bs, len_kv, DIM]
@@ -296,6 +297,7 @@ at::Tensor prefix_shared_attention(
   return O;
 }
 #endif
+#endif
 
 // PyBind11 Module Registration
 PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
@@ -355,8 +357,10 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
     m.def("deterministic_logp_forward_indexed_out", &deterministic_logp_forward_indexed_out, "Batch-invariant deterministic logp indexed out");
     m.def("deterministic_logp_forward_indexed_fp32", &deterministic_logp_forward_indexed_fp32, "Batch-invariant deterministic logp indexed fp32");
 
-    // registry Prefix-Shared Attention
+    // Prefix-shared attention uses NVIDIA PTX and falls back to PyTorch SDPA on ROCm.
+#if !defined(USE_ROCM)
     m.def("prefix_shared_attention", &prefix_shared_attention, "Prefix-Shared Fused Attention for GRPO");
+#endif
 
     // registry Batch-Invariant Deterministic GEMM
     m.def("det_gemm_fwd", &det_gemm_fwd, "Batch-invariant deterministic GEMM forward (C=A@B)");

docker/Dockerfile.rocm

@@ -0,0 +1,28 @@
+# docker/Dockerfile.rocm
+# base: docker build -f docker/Dockerfile.rocm_base -t rl-kernel:rocm-dev .
+# build: docker build -f docker/Dockerfile.rocm -t <registry>/rl-kernel-ci:rocm .
+# push: docker push <registry>/rl-kernel-ci:rocm
+
+FROM rl-kernel:rocm-dev
+
+# Build a portable extension by default. The list follows the multi-architecture
+# ROCm profile used by vLLM: MI200 (gfx90a), MI300/MI325 (gfx942), MI350/MI355
+# (gfx950), plus supported RDNA 3/4 targets. Override it at build time with
+# --build-arg PYTORCH_ROCM_ARCH=<targets>, or at run time with `docker run -e`.
+# Newer GPU targets need no setup.py change: pass the gfx target supported by the
+# installed PyTorch/ROCm pair.
+ARG PYTORCH_ROCM_ARCH=gfx90a;gfx942;gfx950;gfx1100;gfx1101;gfx1150;gfx1151;gfx1200;gfx1201
+ENV PYTORCH_ROCM_ARCH=${PYTORCH_ROCM_ARCH}
+ENV MAX_JOBS=8
+
+USER root
+WORKDIR /opt/rl-kernel
+
+COPY pyproject.toml setup.py* requirements*.txt ./
+
+RUN pip install --no-cache-dir -U pip \
+ && pip install --no-cache-dir -r requirements.txt \
+ && pip install --no-cache-dir pytest
+
+USER rlkernel
+WORKDIR /workspace/RL-Kernel

docker/Dockerfile.rocm_base

@@ -0,0 +1,52 @@
+# docker/Dockerfile.rocm_base
+# Build: docker build -f docker/Dockerfile.rocm_base -t rl-kernel:rocm-dev .
+#
+# To build only the targets deployed in a particular image, override the default:
+# docker build -f docker/Dockerfile.rocm_base \
+#   --build-arg PYTORCH_ROCM_ARCH='gfx942;gfx950' -t rl-kernel:rocm-dev .
+
+ARG BASE_IMAGE=rocm/dev-ubuntu-22.04:7.2.3-complete
+FROM ${BASE_IMAGE}
+
+# This is a build-target list, not a hardware allow-list. It covers the target
+# families supported by the vLLM ROCm 7.2 reference image: MI200 (gfx90a),
+# MI300/MI325 (gfx942), MI350/MI355 (gfx950), and RDNA 3/4. For future GPUs,
+# pass the target accepted by the selected ROCm/PyTorch toolchain at build time.
+ARG PYTORCH_ROCM_ARCH=gfx90a;gfx942;gfx950;gfx1100;gfx1101;gfx1150;gfx1151;gfx1200;gfx1201
+# Keep this wheel index aligned with BASE_IMAGE's ROCm release when overriding it.
+ARG PYTORCH_INDEX_URL=https://download.pytorch.org/whl/rocm7.2
+ARG PYTORCH_VERSION=2.12.1
+ARG RL_KERNEL_USER=rlkernel
+ARG RL_KERNEL_UID=10001
+ARG RL_KERNEL_GID=10001
+
+ENV PATH=/opt/rocm/llvm/bin:/opt/rocm/bin:${PATH}
+ENV ROCM_PATH=/opt/rocm
+ENV LD_LIBRARY_PATH=/opt/rocm/lib:/usr/local/lib:${LD_LIBRARY_PATH}
+ENV PYTORCH_ROCM_ARCH=${PYTORCH_ROCM_ARCH}
+ENV MAX_JOBS=8
+ENV DEBIAN_FRONTEND=noninteractive
+
+RUN apt-get update -y \
+ && apt-get install -y --no-install-recommends \
+      build-essential \
+      ca-certificates \
+      cmake \
+      git \
+      ninja-build \
+      pkg-config \
+      python3 \
+      python3-dev \
+      python3-pip \
+      python3-venv \
+ && python3 -m pip install --no-cache-dir --upgrade pip setuptools wheel \
+ && python3 -m pip install --no-cache-dir --index-url "${PYTORCH_INDEX_URL}" "torch==${PYTORCH_VERSION}" \
+ && python3 -c "import torch; assert torch.version.hip is not None, torch.__version__" \
+ && groupadd --gid "${RL_KERNEL_GID}" "${RL_KERNEL_USER}" \
+ && useradd --uid "${RL_KERNEL_UID}" --gid "${RL_KERNEL_GID}" --create-home --shell /bin/bash "${RL_KERNEL_USER}" \
+ && install -d --owner "${RL_KERNEL_USER}" --group "${RL_KERNEL_USER}" /workspace/RL-Kernel \
+ && rm -rf /var/lib/apt/lists/*
+
+ENV HOME=/home/${RL_KERNEL_USER}
+WORKDIR /workspace/RL-Kernel
+USER ${RL_KERNEL_USER}

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/getting_started/faq.md

@@ -13,8 +13,8 @@ change more often than RL-Kernel's public API.
 | --- | --- | --- | --- |
 | Read docs or edit docs | `pip install -r requirements-docs.txt` | No | Use `mkdocs build --strict -f mkdocs.yaml` before opening a PR. |
 | Run CPU/mock tests | `pip install -e ".[dev]"` | No | Matches the default CI style: fallback and mocked integration coverage. |
-| Run CUDA operators | `pip install -e ".[cuda]"` | Yes, NVIDIA | Requires a CUDA-enabled PyTorch wheel and a working CUDA toolchain for source builds. |
-| Run ROCm operators | `pip install -e ".[rocm]"` | Yes, AMD | Requires a ROCm-enabled PyTorch wheel and ROCm compiler/runtime environment. |
+| Run CUDA operators | `RL_KERNEL_REQUIRE_EXT=1 pip install --no-build-isolation -e ".[cuda]"` | Yes, NVIDIA | Requires a CUDA-enabled PyTorch wheel and a working CUDA toolchain. |
+| Run ROCm operators | `RL_KERNEL_REQUIRE_EXT=1 pip install --no-build-isolation -e ".[rocm]"` | Yes, AMD | Requires a ROCm-enabled PyTorch wheel and ROCm compiler/runtime environment. |
 | Run real vLLM rollout | `pip install -e ".[vllm]"` | Runtime-dependent | Core tests do not need vLLM; install this only where real vLLM is used. |
 
 Do not install every optional extra by default. Install the smallest environment
@@ -84,7 +84,13 @@ The important rule is that PyTorch must match your runtime:
 ```bash
 git clone https://github.com/RL-Align/RL-Kernel.git
 cd RL-Kernel
+
+# CPU-only / pure-Python fallback
 pip install -e .
+
+# Native CUDA or ROCm extension (after installing a matching PyTorch build)
+RL_KERNEL_REQUIRE_EXT=1 pip install --no-build-isolation -e .
+python -c "import rl_engine._C as _C; assert hasattr(_C, 'fused_logp'); print(_C.__file__)"
 ```
 
 The examples on this page use `python3` for system-level commands. Inside an
@@ -93,9 +99,9 @@ activated virtual environment, `python` is also fine.
 ### Which optional extras exist?
 
 ```bash
-pip install -e ".[cuda]"
-pip install -e ".[rocm]"
-pip install -e ".[vllm]"
+RL_KERNEL_REQUIRE_EXT=1 pip install --no-build-isolation -e ".[cuda]"
+RL_KERNEL_REQUIRE_EXT=1 pip install --no-build-isolation -e ".[rocm]"
+pip install --no-build-isolation -e ".[vllm]"
 pip install -e ".[dev]"
 ```
 

docs/getting_started/installation.md

@@ -14,15 +14,15 @@ git clone https://github.com/RL-Align/RL-Kernel.git
 cd RL-Kernel
 # Optional: pin the compile target. If unset, the build targets your GPU's arch.
 # export TORCH_CUDA_ARCH_LIST="9.0+PTX"   # e.g. Hopper; or "8.6+PTX", "12.0+PTX"
-pip install --no-build-isolation -e .
+RL_KERNEL_REQUIRE_EXT=1 python -m pip install --no-build-isolation -e .
 ```
 
-Without `--no-build-isolation`, PyTorch is invisible to the isolated build
-environment, the extension is silently skipped, and the library falls back to the
-slower pure-PyTorch kernels. Confirm the compiled extension is present with:
+`RL_KERNEL_REQUIRE_EXT=1` makes the build fail if `_C` cannot be compiled. Without
+`--no-build-isolation`, PyTorch is invisible to the isolated build environment.
+Confirm the compiled extension is present with:
 
 ```bash
-python -c "from rl_engine import _C; print('compiled extension OK')"
+python -c "import rl_engine._C as _C; assert hasattr(_C, 'fused_logp'); print(_C.__file__)"
 ```
 
 A CPU-only install (plain `pip install -e .` on a machine with no GPU) remains
@@ -34,15 +34,15 @@ The extras add optional dependencies on top of the compiled package, so they use
 the same `--no-build-isolation` flag as the source build above.
 
 ```bash
-pip install --no-build-isolation -e ".[cuda]"
+RL_KERNEL_REQUIRE_EXT=1 python -m pip install --no-build-isolation -e ".[cuda]"
 ```
 
 ```bash
-pip install --no-build-isolation -e ".[rocm]"
+RL_KERNEL_REQUIRE_EXT=1 python -m pip install --no-build-isolation -e ".[rocm]"
 ```
 
 ```bash
-pip install --no-build-isolation -e ".[vllm]"
+python -m pip install --no-build-isolation -e ".[vllm]"
 ```
 
 Install the vLLM extra only on rollout or benchmark environments that need the