feat(alignment): add cross-configuration framework

.github/workflows/ci.yml

@@ -69,14 +69,45 @@ jobs:
           PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 python -m pytest tests/test_attention_correctness.py -q -rs
           python -m pytest tests/test_forward_invariance.py tests/test_tolerance_contract.py tests/test_ws1_workload.py tests/test_gradient_invariance.py tests/test_elementwise_inventory.py tests/test_four_judgment_matrix.py tests/test_op_checks.py tests/test_operator_inputs.py tests/test_profiler.py tests/test_kv_consistency.py tests/test_ws1_qwen3_dense.py tests/test_ws1_chain_integration.py -q
 
+      - name: Run Cross-Configuration Contract Tests (CPU-safe)
+        run: |
+          PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 python -m pytest -q \
+            tests/test_cross_config_*.py \
+            tests/test_stateless_executor.py \
+            tests/test_tolerance_contract.py \
+            tests/test_kernel_registry.py
+
       - name: Run Attention Ground-Truth Tests (CPU-safe)
         run: |
           python -m pytest tests/test_attention.py -v -k "not large and not gpu"
 
+      - name: Run WS2 Attention Cross-Configuration Tests (CPU-safe)
+        run: |
+          PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 python -m pytest -q \
+            tests/test_attention_ablation.py \
+            tests/test_attention_contract.py \
+            tests/test_attention_cross_config_binding.py \
+            tests/test_attention_preprocess.py \
+            tests/test_attention_projection.py \
+            tests/test_cp_attention.py \
+            tests/test_cp_attention_transformer_engine.py
+
       - name: Run KV-Cache Attention Ground-Truth Tests (CPU-safe)
         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 Wrapper Interface and FFN Tests (CPU-safe)
+        run: |
+          python -m pytest -q \
+            tests/test_alignment_wrapper_interfaces.py \
+            tests/test_qwen_ffn.py
+
   docs:
     runs-on: ubuntu-latest
     steps:

.gitignore

@@ -208,6 +208,9 @@ marimo/_static/
 marimo/_lsp/
 __marimo__/
 
+# Cross-configuration alignment local run artifacts
+/runs/
+
 # Local dev notes (not for upstream)
 _dev_notes/
 

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/cuda/gemm/det_gemm_kernel.cu

@@ -8,7 +8,10 @@
 //               fixed K order, NO split-K -> batch-invariant.
 //   Fallback  : naive FP32 scalar kernel (also the correctness ground truth).
 //
-// Both: BF16 in / FP32 accum / no TF32 / no split-K.
+// Both: BF16 in / FP32 accum / BF16 store / no TF32 / no split-K.
+// K is reduced with a mid-split tree. A contiguous half-K GEMM is one child,
+// so simulated TP=2 (a+b) matches TP=1. TP=8 left-fold does not.
+// Leaves stay FP32 (naive: 32-wide MAC; SM90: one BK).
 //   fwd: C = A @ B   |   dA = dC @ B^T   |   dB = A^T @ dC
 // Backward reuses the forward kernel on transposed operands.
 
@@ -24,45 +27,63 @@ namespace {
 
 using nv_bf16 = __nv_bfloat16;
 
-template <typename output_t>
-__device__ __forceinline__ output_t cast_output(float value);
+__host__ __device__ constexpr int cdiv(int a, int b) { return (a + b - 1) / b; }
+
+// Must match SM90 BK so an aligned-K naive tree equals the SM90 tile tree.
+constexpr int K_TREE_LEAF = 32;
 
-template <>
-__device__ __forceinline__ nv_bf16 cast_output<nv_bf16>(float value) {
-  return __float2bfloat16(value);
+__device__ __forceinline__ nv_bf16 bf16_add(nv_bf16 a, nv_bf16 b) {
+  return __float2bfloat16(__bfloat162float(a) + __bfloat162float(b));
 }
 
-template <>
-__device__ __forceinline__ float cast_output<float>(float value) {
-  return value;
+// True iff [lo, hi) is a node of the mid-split tree over [0, n).
+__device__ __forceinline__ bool is_mid_split_node(int lo, int hi, int n) {
+  int a = 0, b = n;
+  while (b - a > 1) {
+    if (a == lo && b == hi) return true;
+    const int m = a + (b - a) / 2;
+    if (hi <= m)
+      b = m;
+    else if (lo >= m)
+      a = m;
+    else
+      return false;
+  }
+  return a == lo && b == hi;
 }
 
-__host__ __device__ constexpr int cdiv(int a, int b) { return (a + b - 1) / b; }
+__device__ nv_bf16 k_tree_naive(const nv_bf16* __restrict__ A, const nv_bf16* __restrict__ B,
+                                int row, int col, int N, int K, int lo, int hi) {
+  if (hi - lo <= K_TREE_LEAF) {
+    float acc = 0.0f;
+    for (int k = lo; k < hi; ++k)
+      acc += __bfloat162float(A[row * K + k]) * __bfloat162float(B[k * N + col]);
+    return __float2bfloat16(acc);
+  }
+  const int mid = lo + (hi - lo) / 2;
+  return bf16_add(k_tree_naive(A, B, row, col, N, K, lo, mid),
+                  k_tree_naive(A, B, row, col, N, K, mid, hi));
+}
 
 // Naive FP32 scalar kernel (fallback + ground truth). Batch-invariant by
-// construction: one thread = one output element, fixed ascending K loop.
+// construction: one thread = one output element, mid-split K tree.
 constexpr int NAIVE_TILE = 16;
 
-template <typename output_t>
 __global__ void det_gemm_naive(const nv_bf16* __restrict__ A,
                                const nv_bf16* __restrict__ B,
-                               output_t* __restrict__ C,
+                               nv_bf16* __restrict__ C,
                                int M, int N, int K) {
   const int row = blockIdx.y * NAIVE_TILE + threadIdx.y;
   const int col = blockIdx.x * NAIVE_TILE + threadIdx.x;
   if (row >= M || col >= N) return;
-  float acc = 0.0f;
-  for (int k = 0; k < K; ++k)
-    acc += __bfloat162float(A[row * K + k]) * __bfloat162float(B[k * N + col]);
-  C[row * N + col] = cast_output<output_t>(acc);
+  C[row * N + col] = k_tree_naive(A, B, row, col, N, K, 0, K);
 }
 
-template <typename output_t>
-void launch_naive(const nv_bf16* A, const nv_bf16* B, output_t* C,
+void launch_naive(const nv_bf16* A, const nv_bf16* B, nv_bf16* C,
                   int M, int N, int K, cudaStream_t stream) {
   dim3 block(NAIVE_TILE, NAIVE_TILE);
   dim3 grid(cdiv(N, NAIVE_TILE), cdiv(M, NAIVE_TILE));
-  det_gemm_naive<output_t><<<grid, block, 0, stream>>>(A, B, C, M, N, K);
+  det_gemm_naive<<<grid, block, 0, stream>>>(A, B, C, M, N, K);
 }
 
 #if defined(RL_KERNEL_ENABLE_SM90)
@@ -72,6 +93,7 @@ void launch_naive(const nv_bf16* A, const nv_bf16* B, output_t* C,
 // passing B^T ([N,K] row-major) so the B smem tile is [BN,BK] (row=n,col=k),
 // matching the validated logp ldmatrix addressing.
 constexpr int BM = 128, BN = 64, BK = 32;
+static_assert(BK == K_TREE_LEAF, "SM90 tile width must match the naive K-tree leaf");
 constexpr int WARPS = 4;
 constexpr int WG_THREADS = WARPS * 32;  // 128
 constexpr int STAGES = 2;
@@ -82,6 +104,7 @@ constexpr int M_TILES = WARP_M / MMA_M;   // 1
 constexpr int N_TILES = BN / MMA_N;       // 8
 constexpr int K_TILES = BK / MMA_K;       // 2
 constexpr int KK_GROUPS = BK / 32;        // 1
+constexpr int TREE_DEPTH = 16;
 
 __device__ __forceinline__ void ldmatrix_x4(uint32_t regs[4], uint32_t addr) {
   asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];"
@@ -96,10 +119,9 @@ __device__ __forceinline__ void mma_m16n8k16(const uint32_t A[4], const uint32_t
                  "f"(D[0]), "f"(D[1]), "f"(D[2]), "f"(D[3]));
 }
 
-template <typename output_t>
 __global__ void det_gemm_sm90_kernel(const __grid_constant__ CUtensorMap a_tmap,
                                      const __grid_constant__ CUtensorMap bt_tmap,
-                                     output_t* __restrict__ C,
+                                     nv_bf16* __restrict__ C,
                                      int M, int N, int K) {
   const int tid = threadIdx.x;
   const int warp = tid / 32;
@@ -141,19 +163,18 @@ __global__ void det_gemm_sm90_kernel(const __grid_constant__ CUtensorMap a_tmap,
 #pragma unroll
   for (int s = 0; s < STAGES; ++s) phase[s] = 0;
 
-  float acc[M_TILES][N_TILES][4];
-#pragma unroll
-  for (int mi = 0; mi < M_TILES; ++mi)
-#pragma unroll
-    for (int n = 0; n < N_TILES; ++n)
-      acc[mi][n][0] = acc[mi][n][1] = acc[mi][n][2] = acc[mi][n][3] = 0.0f;
+  float tile_acc[M_TILES][N_TILES][4];
+  nv_bf16 tree_v[M_TILES][N_TILES][4];
+  nv_bf16 tree_stk[TREE_DEPTH][M_TILES][N_TILES][4];
+  int tree_lo[TREE_DEPTH], tree_hi[TREE_DEPTH];
+  int sp = 0;
 
   if (tid == 0)
 #pragma unroll
     for (int s = 0; s < STAGES - 1; ++s)
       if (s < kd) issue_load(s);
 
-  for (int k = 0; k < kd; ++k) {       // fixed ascending K order, NO split-K
+  for (int k = 0; k < kd; ++k) {       // fixed ascending tile order, NO split-K
     const int buf = k % STAGES;
     if (tid == 0 && k + (STAGES - 1) < kd) issue_load(k + (STAGES - 1));
     det_gemm::mbar_wait(mbar[buf], phase[buf]);
@@ -163,6 +184,12 @@ __global__ void det_gemm_sm90_kernel(const __grid_constant__ CUtensorMap a_tmap,
     const uint32_t sA_buf = sA_base + buf * BM * BK * sizeof(nv_bf16);
     const uint32_t sB_buf = sB_base + buf * BN * BK * sizeof(nv_bf16);
 
+#pragma unroll
+    for (int mi = 0; mi < M_TILES; ++mi)
+#pragma unroll
+      for (int n = 0; n < N_TILES; ++n)
+        tile_acc[mi][n][0] = tile_acc[mi][n][1] = tile_acc[mi][n][2] = tile_acc[mi][n][3] = 0.0f;
+
     uint32_t A[M_TILES][K_TILES][4];
 #pragma unroll
     for (int mi = 0; mi < M_TILES; ++mi) {
@@ -187,12 +214,43 @@ __global__ void det_gemm_sm90_kernel(const __grid_constant__ CUtensorMap a_tmap,
         const uint32_t B1[2] = {b4[2], b4[3]};
 #pragma unroll
         for (int mi = 0; mi < M_TILES; ++mi) {
-          mma_m16n8k16(A[mi][2 * kk + 0], B0, acc[mi][n]);
-          mma_m16n8k16(A[mi][2 * kk + 1], B1, acc[mi][n]);
+          mma_m16n8k16(A[mi][2 * kk + 0], B0, tile_acc[mi][n]);
+          mma_m16n8k16(A[mi][2 * kk + 1], B1, tile_acc[mi][n]);
         }
       }
     }
     __syncthreads();
+
+#pragma unroll
+    for (int mi = 0; mi < M_TILES; ++mi)
+#pragma unroll
+      for (int n = 0; n < N_TILES; ++n)
+#pragma unroll
+        for (int i = 0; i < 4; ++i) tree_v[mi][n][i] = __float2bfloat16(tile_acc[mi][n][i]);
+
+    int lo = k, hi = k + 1;
+    while (sp > 0 && tree_hi[sp - 1] == lo && is_mid_split_node(tree_lo[sp - 1], hi, kd)) {
+#pragma unroll
+      for (int mi = 0; mi < M_TILES; ++mi)
+#pragma unroll
+        for (int n = 0; n < N_TILES; ++n)
+#pragma unroll
+          for (int i = 0; i < 4; ++i)
+            tree_v[mi][n][i] = bf16_add(tree_stk[sp - 1][mi][n][i], tree_v[mi][n][i]);
+      lo = tree_lo[sp - 1];
+      --sp;
+    }
+    if (hi < kd) {
+#pragma unroll
+      for (int mi = 0; mi < M_TILES; ++mi)
+#pragma unroll
+        for (int n = 0; n < N_TILES; ++n)
+#pragma unroll
+          for (int i = 0; i < 4; ++i) tree_stk[sp][mi][n][i] = tree_v[mi][n][i];
+      tree_lo[sp] = lo;
+      tree_hi[sp] = hi;
+      ++sp;
+    }
   }
 
 #pragma unroll
@@ -202,19 +260,18 @@ __global__ void det_gemm_sm90_kernel(const __grid_constant__ CUtensorMap a_tmap,
     for (int n = 0; n < N_TILES; ++n) {
       const int col = col_base + n * MMA_N + (lane % 4) * 2;
       if (row < M && col + 1 < N) {
-        C[row * N + col + 0] = cast_output<output_t>(acc[mi][n][0]);
-        C[row * N + col + 1] = cast_output<output_t>(acc[mi][n][1]);
+        C[row * N + col + 0] = tree_v[mi][n][0];
+        C[row * N + col + 1] = tree_v[mi][n][1];
       }
       if (row + 8 < M && col + 1 < N) {
-        C[(row + 8) * N + col + 0] = cast_output<output_t>(acc[mi][n][2]);
-        C[(row + 8) * N + col + 1] = cast_output<output_t>(acc[mi][n][3]);
+        C[(row + 8) * N + col + 0] = tree_v[mi][n][2];
+        C[(row + 8) * N + col + 1] = tree_v[mi][n][3];
       }
     }
   }
 }
 
-template <typename output_t>
-bool launch_sm90(const nv_bf16* A, const nv_bf16* Bt, output_t* C,
+bool launch_sm90(const nv_bf16* A, const nv_bf16* Bt, nv_bf16* C,
                  int M, int N, int K, cudaStream_t stream) {
   if (M % BM != 0 || N % BN != 0 || K % BK != 0) return false;  // fall back
 
@@ -224,11 +281,11 @@ bool launch_sm90(const nv_bf16* A, const nv_bf16* Bt, output_t* C,
 
   const int smem = STAGES * (BM * BK + BN * BK) * sizeof(nv_bf16) + STAGES * 8;
   if (smem > 48 * 1024)
-    cudaFuncSetAttribute(det_gemm_sm90_kernel<output_t>,
+    cudaFuncSetAttribute(det_gemm_sm90_kernel,
                          cudaFuncAttributeMaxDynamicSharedMemorySize, smem);
 
   dim3 grid(cdiv(N, BN), cdiv(M, BM));
-  det_gemm_sm90_kernel<output_t><<<grid, WG_THREADS, smem, stream>>>(a_tmap, bt_tmap, C, M, N, K);
+  det_gemm_sm90_kernel<<<grid, WG_THREADS, smem, stream>>>(a_tmap, bt_tmap, C, M, N, K);
   return true;
 }
 #endif  // RL_KERNEL_ENABLE_SM90
@@ -249,11 +306,9 @@ void check_in(const torch::Tensor& t, const char* n) {
   TORCH_CHECK(t.scalar_type() == torch::kBFloat16, n, " must be bf16");
 }
 
-torch::Tensor gemm_dispatch(const torch::Tensor& a, const torch::Tensor& b,
-                            bool output_fp32 = false) {
+torch::Tensor gemm_dispatch(const torch::Tensor& a, const torch::Tensor& b) {
   const int M = a.size(0), K = a.size(1), N = b.size(1);
-  auto options = a.options().dtype(output_fp32 ? torch::kFloat32 : torch::kBFloat16);
-  auto c = torch::empty({M, N}, options);
+  auto c = torch::empty({M, N}, a.options());
   auto stream = at::cuda::getCurrentCUDAStream();
 
 #if defined(RL_KERNEL_ENABLE_SM90)
@@ -268,21 +323,15 @@ torch::Tensor gemm_dispatch(const torch::Tensor& a, const torch::Tensor& b,
       a_use = torch::zeros({Mp, K}, a.options());
       a_use.narrow(0, 0, M).copy_(a);
     }
-    torch::Tensor c_use = (Mp != M) ? torch::empty({Mp, N}, options) : c;
+    torch::Tensor c_use = (Mp != M) ? torch::empty({Mp, N}, a.options()) : c;
     auto bt = b.t().contiguous();  // [N,K]
-    const bool launched = output_fp32
-        ? launch_sm90<float>(bf16(a_use), bf16(bt), c_use.data_ptr<float>(), Mp, N, K, stream)
-        : launch_sm90<nv_bf16>(bf16(a_use), bf16(bt), bf16o(c_use), Mp, N, K, stream);
-    if (launched) {
+    if (launch_sm90(bf16(a_use), bf16(bt), bf16o(c_use), Mp, N, K, stream)) {
       if (Mp != M) c.copy_(c_use.narrow(0, 0, M));
       return c;
     }
   }
 #endif
-  if (output_fp32)
-    launch_naive<float>(bf16(a), bf16(b), c.data_ptr<float>(), M, N, K, stream);
-  else
-    launch_naive<nv_bf16>(bf16(a), bf16(b), bf16o(c), M, N, K, stream);
+  launch_naive(bf16(a), bf16(b), bf16o(c), M, N, K, stream);
   return c;
 }
 
@@ -301,7 +350,8 @@ torch::Tensor det_gemm_fwd_fp32(torch::Tensor a, torch::Tensor b) {
   a = a.contiguous(); b = b.contiguous();
   TORCH_CHECK(a.dim() == 2 && b.dim() == 2, "det_gemm_fwd_fp32: expect 2D [M,K]@[K,N]");
   TORCH_CHECK(b.size(0) == a.size(1), "det_gemm_fwd_fp32: K mismatch");
-  return gemm_dispatch(a, b, true);
+  // Keep the FP32 running sum; only the final store is BF16.
+  return gemm_dispatch(a, b);
 }
 
 torch::Tensor det_gemm_da(torch::Tensor dc, torch::Tensor b) {