diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index c7e9040a..7ada3eb6 100644 --- a/.github/workflows/ci.yml +++ b/.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: diff --git a/.gitignore b/.gitignore index 5c38edec..bab026fa 100644 --- a/.gitignore +++ b/.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/ diff --git a/README.md b/README.md index f1ab58d8..0be7704e 100644 --- a/README.md +++ b/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 diff --git a/csrc/cuda/gemm/det_gemm_kernel.cu b/csrc/cuda/gemm/det_gemm_kernel.cu index cd92fc9d..04281f6a 100644 --- a/csrc/cuda/gemm/det_gemm_kernel.cu +++ b/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 -__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(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 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 __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(acc); + C[row * N + col] = k_tree_naive(A, B, row, col, N, K, 0, K); } -template -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<<>>(A, B, C, M, N, K); + det_gemm_naive<<>>(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 __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(acc[mi][n][0]); - C[row * N + col + 1] = cast_output(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(acc[mi][n][2]); - C[(row + 8) * N + col + 1] = cast_output(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 -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, + cudaFuncSetAttribute(det_gemm_sm90_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem); dim3 grid(cdiv(N, BN), cdiv(M, BM)); - det_gemm_sm90_kernel<<>>(a_tmap, bt_tmap, C, M, N, K); + det_gemm_sm90_kernel<<>>(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(bf16(a_use), bf16(bt), c_use.data_ptr(), Mp, N, K, stream) - : launch_sm90(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(bf16(a), bf16(b), c.data_ptr(), M, N, K, stream); - else - launch_naive(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) { diff --git a/csrc/deterministic_logp_kernel.cu b/csrc/deterministic_logp_kernel.cu index 1ca23287..c238e11d 100644 --- a/csrc/deterministic_logp_kernel.cu +++ b/csrc/deterministic_logp_kernel.cu @@ -15,6 +15,15 @@ constexpr int kDeterministicLogpMediumVocabLimit = 4096; constexpr int kDeterministicLogpWarpSize = 32; constexpr float kDeterministicLogpNegInf = -3.4028234663852886e38F; +template +__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 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; diff --git a/csrc/fused_logp_kernel.cu b/csrc/fused_logp_kernel.cu index b620b047..679a6a30 100644 --- a/csrc/fused_logp_kernel.cu +++ b/csrc/fused_logp_kernel.cu @@ -5,6 +5,17 @@ #include #include +constexpr int kFusedLogpLogicalWarpSize = 32; + +template +__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 __device__ __forceinline__ scalar_t blockReduceMax(scalar_t val) { static __shared__ float shared[32]; @@ -14,7 +25,7 @@ __device__ __forceinline__ scalar_t blockReduceMax(scalar_t val) { float f_val = static_cast(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(); @@ -22,7 +33,7 @@ __device__ __forceinline__ scalar_t blockReduceMax(scalar_t val) { 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(f_val); } @@ -36,7 +47,7 @@ __device__ __forceinline__ scalar_t blockReduceSum(scalar_t val) { float f_val = static_cast(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(); @@ -44,7 +55,7 @@ __device__ __forceinline__ scalar_t blockReduceSum(scalar_t val) { 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(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); } } diff --git a/csrc/ops.cpp b/csrc/ops.cpp index e0e86c84..fd612b1e 100644 --- a/csrc/ops.cpp +++ b/csrc/ops.cpp @@ -71,7 +71,6 @@ torch::Tensor lm_head_sm90_forward(torch::Tensor hidden, torch::Tensor lm_head_sm90_forward_fp32(torch::Tensor hidden, torch::Tensor weight, torch::optional bias); -torch::Tensor det_gemm_rowwise_fwd_fp32(torch::Tensor a, torch::Tensor b); #endif #if defined(__CUDACC__) || defined(KERNEL_ALIGN_WITH_CUDA) @@ -105,7 +104,6 @@ void deterministic_collective_all_gather(int64_t handle, torch::Tensor& output); // Batch-Invariant Deterministic GEMM Declarations torch::Tensor det_gemm_fwd(torch::Tensor a, torch::Tensor b); -torch::Tensor det_gemm_fwd_fp32(torch::Tensor a, torch::Tensor b); torch::Tensor det_gemm_da(torch::Tensor dc, torch::Tensor b); torch::Tensor det_gemm_db(torch::Tensor a, torch::Tensor dc); // SiLU / SwiGLU Declarations (elementwise activation, general CUDA) @@ -257,14 +255,6 @@ std::vector deterministic_attention_forward( double scale, torch::optional key_padding_mask); -std::vector deterministic_attention_forward_fp32( - torch::Tensor q, - torch::Tensor k, - torch::Tensor v, - bool causal, - double scale, - torch::optional key_padding_mask); - std::vector deterministic_attention_backward( torch::Tensor grad_output, torch::Tensor q, @@ -277,6 +267,7 @@ std::vector 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] @@ -320,6 +311,7 @@ at::Tensor prefix_shared_attention( return O; } #endif +#endif // PyBind11 Module Registration PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { @@ -362,8 +354,6 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { "Single-card SM90 batch-invariant LM-head forward"); m.def("lm_head_sm90_forward_fp32", &lm_head_sm90_forward_fp32, "Single-card SM90 batch-invariant LM-head forward with fp32 output"); - m.def("det_gemm_rowwise_fwd_fp32", &det_gemm_rowwise_fwd_fp32, - "SM90 deterministic rowwise GEMM with FP32 inputs/accumulation/output"); #endif #if defined(__CUDACC__) || defined(KERNEL_ALIGN_WITH_CUDA) @@ -411,13 +401,13 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { &deterministic_collective_all_gather, "Run the TP=8 deterministic rank-ordered all-gather kernel"); - // 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)"); - m.def("det_gemm_fwd_fp32", &det_gemm_fwd_fp32, - "Batch-invariant deterministic GEMM forward with FP32 output"); m.def("det_gemm_da", &det_gemm_da, "Batch-invariant deterministic GEMM backward dA (dC@B^T)"); m.def("det_gemm_db", &det_gemm_db, "Batch-invariant deterministic GEMM backward dB (A^T@dC)"); // registry RMSNorm @@ -436,10 +426,6 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { "deterministic_attention_forward", &deterministic_attention_forward, "Deterministic standard softmax attention forward (out, lse)"); - m.def( - "deterministic_attention_forward_fp32", - &deterministic_attention_forward_fp32, - "Deterministic standard softmax attention forward with FP32 output"); m.def( "deterministic_attention_backward", &deterministic_attention_backward, diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm new file mode 100644 index 00000000..93ec25f8 --- /dev/null +++ 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 /rl-kernel-ci:rocm . +# push: docker push /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=, 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 diff --git a/docker/Dockerfile.rocm_base b/docker/Dockerfile.rocm_base new file mode 100644 index 00000000..ab577175 --- /dev/null +++ b/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} diff --git a/docs/assets/ws2-cross-config-before-after.png b/docs/assets/ws2-cross-config-before-after.png deleted file mode 100644 index c83db1ce..00000000 Binary files a/docs/assets/ws2-cross-config-before-after.png and /dev/null differ diff --git a/docs/design/cross_config_implementation_report.md b/docs/design/cross_config_implementation_report.md new file mode 100644 index 00000000..7239f9d3 --- /dev/null +++ b/docs/design/cross_config_implementation_report.md @@ -0,0 +1,185 @@ +# Cross-Configuration Alignment Implementation Report + +Status: V1 framework snapshot, 2026-07-19 + +Related work: + +- [Roadmap #83](https://github.com/RL-Align/RL-Kernel/issues/83) +- [Cross-configuration alignment #111](https://github.com/RL-Align/RL-Kernel/issues/111) +- [Numerical contract #108](https://github.com/RL-Align/RL-Kernel/issues/108) +- [V1 contract](cross_config_logprob_drift_contract.md) + +## Result and claim boundary + +This change provides a small framework for planning and executing paired +rollout/training logprob comparisons across controlled configuration changes. It +includes strict configuration loading, bounded case planning, exact semantic +operator selection, lifecycle-aware runtime materialization, paired read-only +scoring, fixed-contract comparison, append-only artifacts, and validated resume. + +The included executable path is deliberately CPU-only. It validates framework +plumbing with a synthetic model and temporary selected-logprob backends; it does +not claim production vLLM, FSDP, TP, CP, accelerator, or distributed numerical +alignment. The S1, S2, and S3 examples are plans, not execution evidence. + +## Architecture + +The implementation keeps configuration, operator resolution, runtime ownership, +and execution separate: + +```text +JSON -> ExperimentConfig -> Planner -> ExperimentPlan + | + build_execution_plan + | + operator-bound ExecutionPlan + / \ + operator session RuntimeMaterializer + | + RuntimeBinding + | + ArtifactStore <- PairedRunner -> fixed comparator +``` + +| Boundary | Responsibility | +| --- | --- | +| `config.py` and `planner.py` | Load strict, versioned JSON; normalize the ten supported knobs; emit an `ExperimentPlan` containing a baseline plus declared OAT or explicit pairwise cases under a fixed 256-case cap; compute stable semantic case IDs without importing a runtime. | +| `build_execution_plan` | Resolve rollout/training selections, bind them into immutable case identity, and emit canonical operator-bound `ExecutionPlan` rows shared by planning and execution. | +| `SemanticOperatorCatalog` | Store immutable backend descriptors and their target, device, dtype, per-target required topology, lifecycle, factory, and observability constraints. | +| `OperatorSession` | Resolve and instantiate exact rollout/training implementations for one case, cache only within that case, and produce concrete provenance. | +| `RuntimeMaterializer` | Apply each normalized knob through an owning adapter and report requested, materialized, and actual values with status and lifecycle evidence. | +| `RuntimeBinding` | Carry only backend-neutral batch, side-configuration, topology, scorer, operator-backend, and runtime-kind mappings. Runtime-specific engine objects stay behind the adapter boundary. | +| `PairedRunner` | Supervise isolated rollout/training scoring children, enforce timeout and read-only model state, validate ranks and exact operator instances, compare selected logprobs, and coordinate resume/publication. | +| `ArtifactStore` | Publish immutable attempt directories and write `COMPLETE` last with SHA-256 seals for every required payload. Resume accepts only an attempt whose identity, execution, provenance, tensors, and comparison still validate. | + +Runtime bindings deep-freeze their execution handoff. Rollout topology contains +only rollout-owned world/TP/CP state, while training topology contains only +training-owned world/sharding state; neither side is padded with fields owned by +the other. + +The fixed-threshold comparator applies the repository numerical contract only to +active selected tokens. Identity violations, invalid artifacts, non-finite +scores, and zero active tokens cannot become passes. Diagnostics remain separate +from the pass/fail rule. + +## Configuration and operator selection + +Execution controls do not live in experiment JSON. `scenario` is metadata; the +CLI chooses planning versus execution, the runtime adapter, temporary-operator +authorization, timeout, and resume policy. + +`logp.backend` is the concise choice when rollout and training use the same +selected-logprob implementation. The optional top-level `operators` mapping is +the extension point for independent sides and per-backend options: + +```json +{ + "baseline": { + "logp": {"backend": "rlkernel.reference_logp"} + }, + "operators": { + "selected_logprob": { + "rollout": "rlkernel.reference_logp", + "training": { + "backend": "vendor.training_logp", + "options": {"mode": "exact"} + } + } + } +} +``` + +The explicit rollout backend must agree with the baseline `logp.backend`. +Explicit operators cannot be combined with `logp.backend` interventions because +that would make the planned knob differ from the implementation actually used. +Unknown fields, threshold overrides, hidden execution controls, duplicate JSON +keys, and non-finite values are rejected. + +## CLI + +Planning validates and records a plan without constructing a runtime: + +```bash +python -m rl_engine.alignment.cross_config plan \ + examples/cross_config_s3_qwen3_8b_tp4_cp4_bf16.json +``` + +The only shipped execution adapter is the explicit CPU smoke runtime: + +```bash +python -m rl_engine.alignment.cross_config run \ + examples/cross_config_s0_cpu_smoke.json \ + --runtime cpu-smoke \ + --allow-smoke-operators +``` + +Both commands accept `--output-root`. `run` also exposes a per-attempt timeout +and `--no-resume`; these policies are intentionally absent from the JSON schema. + +## CPU testing adapter + +`rl_engine.alignment.testing.cpu_cross_config` owns the synthetic causal model, +stateless CPU scorer, `CpuSmokeMaterializer`, canonical batch construction, and +CPU experiment helpers. Keeping these objects outside the core package prevents +test hardware and model assumptions from becoming runtime abstractions. + +Temporary selected-logprob implementations live under +`rl_engine/alignment/testing/smoke_ops`. They advertise CPU as their only device, +are not registered by default, and require explicit policy authorization. The +reference backend performs `log_softmax` plus gather; the offset backend is used +only by focused mismatch tests. The shipped S0 example selects reference on both +sides and contains one baseline case. + +## Scenario evidence + +| Scenario | Planned cases | Current evidence | +| --- | ---: | --- | +| S0 CPU framework smoke | 1 | One reference/reference case passes on CPU; a second invocation validates and resumes the same complete attempt. | +| S1 distributed smoke | 5 | Configuration loading and planning only. | +| S2 vLLM TP versus FSDP | 10 | Configuration loading and planning only. | +| S3 Qwen3-8B TP=4, CP=4, BF16 | 11 | Configuration loading and planning only. | + +Planning success does not imply that the requested production topology can be +materialized. Unsupported or unobservable runtime settings fail strict execution +instead of silently falling back. + +## Validation snapshot + +```text +Focused contract/runtime/runner/CLI and existing regression tests: +60 passed in 2.20s + +CPU-collectable repository suite: +418 passed, 242 skipped in 15.83s + +Named scenario checks: +S0 run: 1 pass, then 1 validated resume +S1/S2/S3 plan: 5 / 10 / 11 cases, no runtime constructed +``` + +The final review also checks JSON syntax, formatting, static typing, strict +documentation build, and `git diff --check`. + +The full CPU command excludes `test_grpo_loss.py` and `test_ratio_kl.py` +because those modules require Triton during collection. It ran outside the +restricted sandbox so Gloo and POSIX shared-memory tests could use host +resources. + +## Extension path and known gaps + +A production backend extends the semantic catalog with a descriptor and factory, +then supplies injection/read-back hooks through its runtime adapter. The planner +and runner do not need backend-specific branches. Operator correctness remains +owned by the operator implementation; the framework verifies exact selection, +materialization evidence, paired identity, comparison, and provenance. + +Production execution still requires: + +- verified selected-logprob injection and read-back for the rollout engine; +- a read-only pre-update training scorer for FSDP and distributed rank evidence; +- context-parallel application/read-back and process-group orchestration; +- accelerator-backed lifecycle and cleanup tests; and +- the production kernels tracked by their owning workstreams. + +Until those adapters exist, S1-S3 remain reproducible planning inputs and the CPU +smoke remains a framework claim only. diff --git a/docs/design/cross_config_logprob_drift_contract.md b/docs/design/cross_config_logprob_drift_contract.md new file mode 100644 index 00000000..3f06bc18 --- /dev/null +++ b/docs/design/cross_config_logprob_drift_contract.md @@ -0,0 +1,308 @@ +# Cross-Configuration Logprob Drift Contract + +Status: V1 implementation contract + +Related work: + +- [Roadmap #83](https://github.com/RL-Align/RL-Kernel/issues/83) +- [Cross-configuration alignment #111](https://github.com/RL-Align/RL-Kernel/issues/111) +- [Numerical contract #108](https://github.com/RL-Align/RL-Kernel/issues/108) + +## Goal and boundary + +This framework isolates configuration changes that can make rollout-selected +log probabilities differ from training-side recomputation. It provides typed +plans, lifecycle-aware runtime materialization, exact semantic-operator +selection, paired read-only scoring, append-only artifacts, and safe resume. + +It does not implement production AG, RS, GEMM, attention, logprob, TP-invariant, +CP-aware, or deterministic collective kernels. Those implementations remain +owned by their operator workstreams and integrate through the semantic operator +catalog described below. + +## The only pass/fail rule + +For every active selected response/action token: + +```text +abs(training_logprob - rollout_logprob) > fixed_threshold +``` + +The fixed threshold is loaded from the repository numerical contract. It is not +a config field, CLI flag, experiment axis, workload policy, or operator option. +Equality with the threshold is not a mismatch. + +```python +mismatch_mask = active_mask & ( + torch.abs(training_logprobs - rollout_logprobs) > fixed_threshold +) +``` + +Mean, percentiles, maximum absolute difference, mismatch ratio, worst-token +location, and approximate KL are diagnostics only. They never change pass/fail. +The token-level artifact persists both logprob tensors, the active mask, and the +resolved threshold so the mask can be recomputed offline. + +Zero active tokens produce `ZERO_ACTIVE_TOKENS`, never a pass. Non-finite or +non-floating active scores produce `INVALID_ARTIFACT`. + +## Identity before numerics + +A comparison is valid only when both scorers use the same logical input: + +- immutable checkpoint and model version; +- tokenizer ID and tokenization policy; +- generated token IDs and selected-token IDs; +- active and attention masks; +- pre-update model state; +- required position, cache, and packing metadata. + +The training scorer teacher-forces the already generated sequence. It cannot +generate replacement tokens, use a KV cache when the frozen identity forbids +one, own an optimizer, update parameters or buffers, or leave the model in a +different mode. An identity violation is `INVALID_IDENTITY`, not numerical +drift. + +## V1 configuration + +The user supplies one explicit baseline plus declared interventions. Lists do +not imply a Cartesian product. + +```json +{ + "schema_version": "cross_config.experiment_config.v1", + "experiment_id": "qwen3-8b-alignment", + "scenario_id": "qwen3-8b-tp4-cp4-bf16", + "contract_source": "ws1", + "contract_version": "current", + "strategy": "one_at_a_time", + "strict_fallback": true, + "identity": {"...": "frozen scoring identity"}, + "baseline": { + "batch": {"size": 8}, + "rollout": { + "tensor_parallel_size": 4, + "context_parallel_size": 4, + "dtype": "bfloat16", + "enable_prefix_caching": true, + "enforce_eager": false + }, + "training": { + "attention_backend": "flash_attention_2", + "compute_dtype": "bfloat16", + "sharding": "fsdp" + }, + "logp": {"backend": "native"} + }, + "interventions": [ + {"path": "batch.size", "values": [1]}, + {"path": "logp.backend", "values": ["rlkernel.reference_logp"]} + ], + "scenario": {"level": "S3", "device": "cuda"} +} +``` + +`scenario` is metadata. Execution mode and authorization policy belong to the +CLI, so a config cannot hide `plan_only`, expected test outcomes, or permission +to activate temporary operators. + +### Exact knob allowlist + +| Knob | Minimum lifecycle | Meaning | +|---|---|---| +| `batch.size` | request | Canonical sample chunking only. | +| `rollout.tensor_parallel_size` | process | Rollout TP world. | +| `rollout.context_parallel_size` | process | Rollout CP world. | +| `rollout.dtype` | engine construction | Rollout numerical dtype. | +| `rollout.enable_prefix_caching` | engine construction | Engine cache policy. | +| `rollout.enforce_eager` | engine construction | Eager versus optimized/graph path. | +| `training.attention_backend` | engine construction | Training scorer attention implementation. | +| `training.compute_dtype` | engine construction | Training scorer compute dtype. | +| `logp.backend` | engine construction | Both-sides selected-logprob shortcut. | +| `training.sharding` | process | Training topology, such as unsharded or FSDP. | + +TP/vocabulary layout is derived and recorded, not user-settable. Tokenization, +masks, positions, checkpoint identity, and pre-update state are invariants, not +ordinary knobs. Quantization, FP8, MoE, speculative decoding, pipeline +parallelism, and arbitrary runtime fields are deferred. + +### Planning + +`one_at_a_time` emits one baseline and cases that change exactly one declared +path. `pairwise` is opt-in and expands only explicitly listed path pairs. The +planner normalizes aliases, validates the allowlist and capability constraints, +and reports structured issues without creating engines. A fixed 256-case +framework cap stops OAT or pairwise expansion before unbounded accumulation. + +Stable case IDs hash normalized requested values, identity, contract version, +and scenario definition. Runtime readback never rewrites a case ID. A retry gets +a new attempt ID under the same case. + +## Architecture and extension points + +The core has one-way responsibilities: + +```text +strict config -> Planner -> ExperimentPlan -> build_execution_plan + | + operator-bound ExecutionPlan + | + runtime adapter -> RuntimeBinding + | + paired runner + / \ + artifact store comparator +``` + +- `config.py` owns the external schema and strict JSON loading. +- `schema.py` owns immutable, versioned domain records. +- `planner.py` owns normalization, the knob catalog, OAT, and pairwise cases. +- `execution_plan.py` binds the selected rollout/training operators into each + immutable case and produces canonical rows shared by planning and execution. +- `runtime.py` owns the adapter protocol, three-stage materialization, lifecycle + fingerprints, and a backend-neutral execution binding. +- `comparison.py` owns identity validation and the fixed-threshold result. +- `runner.py` coordinates paired execution and atomic publication; private + execution, provenance, and resume modules isolate process supervision and + validation details. +- `artifacts.py` owns append-only attempts and resume discovery. +- `semantic_registry.py` owns generic operator descriptors and case-local + resolution sessions; it is shared by future alignment features. + +The package root exposes only the common planning and execution facade. Runtime, +artifact, schema, and operator internals remain in their owning modules. + +### Runtime adapters + +A runtime adapter receives the normalized case and returns one application +record per knob: + +```text +requested -> materialized -> actual +``` + +Each record includes status (`applied`, `fallback`, `unsupported`, +`unobservable`, or `error`), evidence, and lifecycle. The facade derives +construction, distributed-context, and process fingerprints. Reuse is allowed +only when all relevant fingerprints and operator bindings match. + +Adapters may construct repository-native vLLM, training, or stateless config +objects internally. The core runner receives only backend-neutral batch, side +configuration, topology, scorer, operator-backend, and runtime-kind mappings, +so adding a runtime does not add branches to the planner or runner. + +Strict execution rejects fallback, ignored settings, unobservable critical +values, stale registry state, and incompatible reuse. Fallback is measurable +only when it is itself the declared intervention. + +## Semantic operator selection + +The first semantic operator is `selected_logprob`. Rollout and training can +select implementations independently: + +```json +{ + "operators": { + "selected_logprob": { + "rollout": "rlkernel.reference_logp", + "training": { + "backend": "rlkernel.reference_logp", + "options": {} + } + } + } +} +``` + +When `operators` is absent, `logp.backend` selects the same backend on both +sides. An explicit mapping is bound into execution identity before a runtime is +created. A `logp.backend` intervention cannot be combined with a fixed explicit +mapping because that would create a knob that no longer changes execution. + +Each backend descriptor declares: + +- semantic operation and backend ID; +- supported target tags, devices, dtypes, and per-target required topology values; +- alignment properties and lifecycle; +- implementation factory and version/build fingerprint; +- explicit fallback policy and temporary-test marker. + +`SemanticOperatorCatalog` stores immutable descriptors. Each case creates an +`OperatorSession` for resolution, instantiation, caching, and provenance. Failed +or cached state cannot leak into the next case. Strict resolution never invokes +legacy priority fallback. + +Adding a production implementation requires its existing semantic interface, +one descriptor, runtime injection hooks where needed, operator-owned correctness +tests, and one framework case. It does not require a planner change. + +## Artifacts and resume + +Attempts are append-only: + +```text +runs// + experiment.json + plan.jsonl + cases/// + requested.json + materialized.json + actual.json + identity.json + score_rollout.pt + score_training.pt + comparison.json + token_diffs.pt + COMPLETE +``` + +`COMPLETE` is published last and seals every required payload with a SHA-256 +digest. Resume accepts only a complete attempt whose case, identity, +materialization, scorer, operator, environment, comparison, and tensor artifacts +match the current execution key. Partial, malformed, or tampered attempts are +ignored; an older valid attempt may still be reused. Existing files are never +overwritten. + +## CPU smoke boundary + +The only executable adapter delivered here is under +`rl_engine.alignment.testing.cpu_cross_config`. It is explicitly CPU-only and +uses a deterministic synthetic model plus read-only stateless scoring. Named +distributed and accelerator scenarios are configuration/plan coverage only. + +Temporary selected-logprob backends live together under +`rl_engine/alignment/testing/smoke_ops`: + +- `smoke_only.logp_reference`: PyTorch `log_softmax` plus gather; +- `smoke_only.logp_offset`: the same result with an authorized deterministic + offset used to prove mismatch detection. + +They are CPU-only, marked `is_smoke_only`, unregistered by default, and require +both explicit registration and execution policy authorization. Their exact +removal procedure is in `SMOKE_OPERATORS.md`. Remove them when equivalent +production operators pass the same framework cases, then remove the opt-in flag +and temporary test marker. + +## Scenario levels and claims + +| Level | Purpose | Current claim | +|---|---|---| +| S0 | Local CPU framework smoke | Executable: config, planner, operator selection, paired scoring, comparison, artifacts, resume. | +| S1 | Small distributed lifecycle smoke | Plan only until suitable hardware/runtime adapters exist. | +| S2 | Named vLLM TP versus training FSDP comparison | Plan only. | +| S3 | Qwen3-8B TP=4, CP=4, BF16 milestone | Plan only; no production alignment claim. | + +Run the shipped examples with: + +```bash +python -m rl_engine.alignment.cross_config plan \ + examples/cross_config_s3_qwen3_8b_tp4_cp4_bf16.json + +python -m rl_engine.alignment.cross_config run \ + examples/cross_config_s0_cpu_smoke.json \ + --runtime cpu-smoke \ + --allow-smoke-operators +``` + +Passing S0 proves framework plumbing only. It does not prove accelerator, +distributed, production-operator, or roadmap numerical alignment. diff --git a/docs/design/runtime-dispatch.md b/docs/design/runtime-dispatch.md index bedf3475..23c1586a 100644 --- a/docs/design/runtime-dispatch.md +++ b/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 | diff --git a/docs/design/ws2-attention-cross-config-integration.md b/docs/design/ws2-attention-cross-config-integration.md new file mode 100644 index 00000000..6ddd45a0 --- /dev/null +++ b/docs/design/ws2-attention-cross-config-integration.md @@ -0,0 +1,185 @@ +# WS2 Attention Cross-Configuration Integration + +Implements PR4 of [#235](https://github.com/RL-Align/RL-Kernel/issues/235): wiring the +CP attention path into the cross-configuration planner/runtime for the Qwen3-8B +TP=2 CP=2 BF16 target. + +Builds on [#236](https://github.com/RL-Align/RL-Kernel/pull/236) (attention contract +and dispatch metadata), [#238](https://github.com/RL-Align/RL-Kernel/pull/238) +(deterministic CP reference) and [#230](https://github.com/RL-Align/RL-Kernel/pull/230) +(cross-configuration framework). + +## What "bind to the same contract" means here + +The PR4 acceptance criteria say rollout and training descriptors must "bind to the +same semantic attention contract". Under the frozen deployment the two sides can +never produce identical `AttentionContract` instances: + +| | training (Megatron) | rollout (vLLM) | +| --- | --- | --- | +| mode | full-sequence prefill | chunked prefill, later decode | +| CP | `context_parallel_size`, whole forward | `prefill_context_parallel_size`, prefill only | +| KV | no paging | paged KV with a block table | +| backend vocabulary | `AttnBackend{flash,fused,unfused,local,auto}` | `AttentionBackendEnum` | + +Read literally, the criterion is unsatisfiable. It is therefore implemented as three +tiers, in `rl_engine/alignment/cross_config/attention_binding.py`: + +| tier | fields | rule | failure | +| --- | --- | --- | --- | +| `IDENTICAL` | checkpoint/model/token identity plus complete TP/CP GQA head and sequence ownership | equal bit for bit | `comparable=False`; no drift number from the pair means anything | +| `SEMANTIC` | reduction contract, dtype, exported LSE, first-class Split-KV request, complete actual Split-KV plan sets, CUDA QK-Norm/RoPE identity, cross-side determinism | both sides equal **and** equal to the WS2 mandate | fail closed | +| `RECORDED` | `mode`, `backend_id`, `reduction.engine`, RoPE materialization state and fusion boundary, KV-cache paging | free to differ | recorded into provenance and measured | + +Two placements are load-bearing: + +* **`reduction.engine` is `RECORDED`, not `SEMANTIC`.** Training may run the in-op + deterministic reference while rollout runs a Transformer Engine merge oracle. + Forcing them equal would defeat the oracle comparison that #235 PR2/PR3/PR5/PR6 + depend on. +* **TP/CP topology is `IDENTICAL`, not `RECORDED`.** TP selects local Qwen3 GQA head + ownership and CP selects local sequence ownership. Different topology is a + different local attention problem, not a backend detail. +* **`reduction.order`, `reduction.acc_dtype`, and actual Split-KV schedules are + `SEMANTIC`.** Both runtimes must export the complete batch x TP x CP x KV-owner + plan set, including logical boundaries, merge order, FP32 accumulation, final + downcast, and fallback state. Configured policy alone never passes strict binding. + +`comparable` and `passed` are separate flags. A pair with mismatched identity is not +comparable. A pair that is comparable but violates the reduction mandate is still +rejected -- the drift would be real but attributable to the wrong thing. + +## H100 Attention input and projection boundary + +Megatron/TE and vLLM/FlashInfer are the first-choice implementations. +`H100AttentionPreprocessor` runs a same-input H100 bitwise probe against the +deterministic RL-Kernel path. An unavailable native callable, a native exception, +or a failed probe switches both sides to `RMSNormCudaOp` + `RoPESM90Op` and +records the fallback reason and probe ID. The launcher passes the returned +backend evidence into `AttentionRuntimeReadback`: + +```python +from rl_engine.kernels.attention_preprocess import H100AttentionPreprocessor + +prepared = H100AttentionPreprocessor(device)( + q, k, q_norm_weight, k_norm_weight, position_ids +) +readback = AttentionRuntimeReadback( + # contract, knobs, Split-KV plan set, source, and scope fields omitted here + **prepared.readback_fields(), +) +``` + +Strict binding rejects a missing or unknown backend and rejects mixed native / +fallback execution. If both sides fall back, they must report the same deterministic +backend IDs and policy ID. Printing a configured backend without executing the +probe is not evidence. + +The Attention boundary includes QKV projection, Q/K RMSNorm, RoPE, core +attention, KV-cache access, CP `(Out, LSE)` communication/merge, and o_proj. +`AttentionProjectionOp` freezes QKV/o_proj to BF16 input and output, FP32 +accumulation, ascending-K reduction, and Split-K disabled. Native projection +callables are accepted only after a bitwise probe against `DetGemmOp`; otherwise +both sides use the deterministic fallback. Its collective contract records QKV +column-parallel plus backward TP all-reduce, o_proj row-parallel partial output, +and the SP all-gather/reduce-scatter directions. The model input RMSNorm and +residual add remain outside this Attention experiment. + +## Determinism is not one thing + +`rl_engine/alignment/cross_config/determinism.py` probes both sides and compares +them, because the two frameworks mean different things by "deterministic": + +| | Megatron `deterministic_mode` | vLLM `VLLM_BATCH_INVARIANT` | +| --- | --- | --- | +| `NCCL_ALGO` | asserts membership in a five-value set | hard-sets `allreduce:tree` | +| `NCCL_PROTO`, channels, threads | not managed | hard-set (`Simple`, `1`, `1`) | +| TF32 | **not managed at all** | disabled (`fp32_precision="ieee"`) | +| BF16 reduced-precision reduction | not managed | disabled | +| cuBLAS workspace / BLAS library | not managed | `:4096:8`, cuBLASLt | +| GEMM | cuBLAS / TE | Triton `matmul_persistent` | +| FlashAttention | forbidden | permitted | + +`NCCL_ALGO`, `NCCL_PROTO` and `CUBLAS_WORKSPACE_CONFIG` change arithmetic, so a +mismatch there is blocking. The remaining differences -- including the TF32 and +BF16-reduction asymmetry, which under a pure BF16 GEMM path does not fire -- are +recorded so the asymmetry appears in every artifact rather than being invisible. + +## Runtime adapters + +Before this PR the only `RuntimeMaterializer` in the repository was +`CpuSmokeMaterializer` over a synthetic CPU model, and every named scenario +(`S1`/`S2`/`S3`) was planning-only. This PR adds the first two framework-shaped +adapters: + +* `adapters/megatron.py` -- `MegatronProvenanceAdapter` (construction and + distributed-context fingerprints, determinism probe, frozen-scope assertions) and + `MegatronAttentionMaterializer`. +* `adapters/vllm.py` -- `VllmProvenanceAdapter` (including diagnostic vLLM split + limits) and `VllmRolloutMaterializer`. +* `AttentionRuntimeReadback` -- the explicit handoff from an executed engine. It + carries the reconstructed actual contract, actual knob values, frozen-scope + verification, executed CUDA QK-Norm/RoPE identities and fallback state, and the + complete Split-KV runtime plan set. + +Constructing a contract is not runtime verification. Without a readback, adapter +applications are `UNOBSERVABLE`; only matching values reconstructed from a real +Megatron or vLLM execution are `APPLIED`. `bind_attention_runtime_readbacks` is the +strict public entry point used after both framework launchers collect that evidence. + +Neither module imports `megatron` or `vllm`; configs are duck-typed, so the binding +rules are exercised on CPU in CI rather than only on a 2-node cluster. + +## Fail closed, never substitute + +`unsupported_reduction_reason` rejects requests that #236 cannot express, instead of +collapsing them onto the supported value: + +| request | status | why | +| --- | --- | --- | +| `attention.reduction_order=arrival` | `UNSUPPORTED` | the control group must stay distinguishable from the treatment | +| `attention.reduction_downcast_at=per_block` | `UNSUPPORTED` | `DowncastPoint` declares only `final_write` | +| `attention.reduction_engine=te_oracle` | `UNSUPPORTED` | the TE merge oracle lands in #235 PR2/PR3; PR4's TE plan is provenance only | +| `attention.reduction_acc_dtype=bf16` | `UNSUPPORTED` | the CP `(out, lse)` merge accumulates in FP32 | +| configured contract without runtime readback | `UNOBSERVABLE` | requested values do not prove what executed | +| `rollout.context_parallel_size>1` with effective decode CP=1 | `ERROR`/`FALLBACK` | strict TP=2/CP=2 acceptance rejects the topology change | +| missing/mismatched/fallback Split-KV plan set | binding failure | Split-KV provenance must cover every batch/TP/CP/owner coordinate | +| missing/unknown QK-Norm or RoPE backend, or mixed native/fallback sides | binding failure | both sides must execute the same verified native policy or the common RL-Kernel CUDA fallback | + +## Knobs + +`adapters/knobs.py` extends `V1_KNOBS` additively. Added: training-side +`tensor_parallel_size` / `context_parallel_size` / `deterministic_mode` / +`cp_comm_type`, `rollout.batch_invariant` / `rollout.kv_block_size`, and the +reduction axis (`acc_dtype`, `order`, `downcast_at`, `engine`) plus +`attention.fusion_boundary` and `attention.split_kv_policy`. + +`training.attention_backend` keeps its path but its value domain is replaced with +Megatron's `AttnBackend`; the HuggingFace names have no Megatron counterpart, so this +is a replacement rather than a mapping. + +Not done here, because both change `V1_KNOBS` itself and would break existing +cross-config tests: removing `training.sharding` (Megatron has no such concept, and +DP=1 makes it moot) and renaming `rollout.context_parallel_size` to reflect that it +binds to `prefill_context_parallel_size`. + +## Scenario + +`examples/cross_config_qwen3_8b_megatron_tp2_cp2_vllm.json` supersedes +`cross_config_s1_distributed_smoke.json` and +`cross_config_s3_qwen3_8b_tp4_cp4_bf16.json`, whose training sides used `sdpa` / +`flash_attention_2` and `sharding: fsdp` -- none of which exist under Megatron -- and +whose TP=4/CP=4 topology does not match the target. +`cross_config_s2_vllm_tp_vs_fsdp.json` has no Megatron-only counterpart and should be +retired rather than rewritten. + +## Out of scope + +Deliberately not in this PR: + +* launching `torchrun`, initializing process groups, or executing core attention; +* pre-attention model RMSNorm and residual add; +* decode-mode materialization, which needs the validated `KVCacheSpec` from #235 PR6 + and is refused with that reference rather than stubbed; +* distributed drift benchmarks and report artifacts (#235 PR5); +* fused production backend alignment (#235 PR7) and backward (#235 PR8). diff --git a/docs/design/ws2-attention-debug-matrix.md b/docs/design/ws2-attention-debug-matrix.md new file mode 100644 index 00000000..32dacb79 --- /dev/null +++ b/docs/design/ws2-attention-debug-matrix.md @@ -0,0 +1,82 @@ +# Attention Debug Matrix + +The post-training Attention tool replays one frozen rollout, keeps the sample +identity fixed, changes one factor at a time, and reports both mismatch metrics +and a small set of invariance controls. +This is intentionally separate from the runtime knob catalog in +`examples/cross_config_qwen3_8b_megatron_tp2_cp2_vllm.json`. + +## Matrix shape + +`rl_engine.kernels.ops.pytorch.attention.debug_matrix` is the source of truth for +the compact matrix manifest: + +| Rows | Role | Meaning | +| --- | --- | --- | +| `A0` | baseline | Strict deterministic replay with the same train/rollout identity. | +| `A1`-`A3`, `A5`-`A7` | root-cause probes | One representative probe for each numerical or stateful Attention category. | +| `A4` | comparability gate | Head/sequence ownership must match before a numerical comparison is meaningful. | +| `C0`-`C2` | invariant controls | Cases that must remain exactly zero; a nonzero result invalidates the run. | + +The seven `A` rows are: + +| Row | Category | Representative probe | +| --- | --- | --- | +| `A1` | Position / RoPE | `position_ids` | +| `A2` | Q/K preprocessing | `qk_norm_disabled` | +| `A3` | Mask / sequence boundary | `causal_mask` | +| `A4` | Topology / head ownership | `tp_head_ownership`; reject as `comparable=false`, do not report a drift scalar. | +| `A5` | KV-cache identity / layout | `kv_page_order` | +| `A6` | Numerical policy | `accum_dtype` | +| `A7` | Distributed schedule | `merge_order` | + +The representative row is a fast first-line diagnosis. The existing taxonomy +keeps the secondary probes for a second pass, so users do not need to run all +21 probes for every post-training incident. + +## Replay contract + +Every row reuses the same: + +- checkpoint and model version; +- token IDs, selected-token IDs, masks, and positions; +- KV-cache and packing metadata; +- pre-update model state; +- train/rollout sample ordering. + +The matrix is one-at-a-time. It does not create a Cartesian product of Attention +knobs. Each diagnostic row records its own phase-local baseline, which allows +rollout-only and train-only debugging to be compared independently. A changed topology is a +gate failure (`comparable=false`), not a meaningful numerical drift sample. A changed +Split-KV plan or merge order is useful as a diagnostic injection (`A7`), but is not an +accepted production comparison until the actual plan again matches on both sides. + +## Metrics + +The replay report should include selected-token mismatch metrics: +`train_rollout_logprob_abs_diff`, forward `mismatch_kl`, and `mismatch_k3_kl`. +The Attention artifact additionally records `out`, `lse`, `dQ`, `dK`, and `dV` +maximum absolute drift. These are diagnostics; pass/fail still uses the fixed +repository numerical contract and never a user-supplied tolerance. + +`C0` (`tp_partition_control`), `C1` (`batch_composition_control`), and `C2` +(`prefill_decode_tail_control`) must be bitwise zero. They catch accidental +batch dependence, invalid preserved head ownership, and cache-position mistakes +before a root-cause row is trusted. + +## Runtime knobs versus debug probes + +The cross-configuration planner continues to own executable lifecycle knobs such +as TP/CP topology, CP communication, fusion boundary, and reduction policy. +Those knobs are materialized and read back by the runtime adapters. The compact +`A0`-`A7` matrix is a triage layer over that execution evidence; it does not +pretend that unsupported values such as `reduction_order=arrival` or BF16 CP +accumulation are production implementations. + +The portable manifest is available from: + +```python +from rl_engine.kernels.ops.pytorch.attention.debug_matrix import attention_debug_matrix + +manifest = attention_debug_matrix() +``` diff --git a/docs/design/ws2-cp-attention-contract.md b/docs/design/ws2-cp-attention-contract.md new file mode 100644 index 00000000..95320101 --- /dev/null +++ b/docs/design/ws2-cp-attention-contract.md @@ -0,0 +1,227 @@ +# WS2 CP-Aware Attention Contract + +Status: PR1 contract and dispatch metadata + +Tracking and shared contracts: + +- [#235: CP-aware deterministic Attention](https://github.com/RL-Align/RL-Kernel/issues/235) +- [#83: WS2 roadmap](https://github.com/RL-Align/RL-Kernel/issues/83) +- [#108: WS1 numerical contract](https://github.com/RL-Align/RL-Kernel/issues/108) +- [#111: WS2 cross-config alignment](https://github.com/RL-Align/RL-Kernel/issues/111) +- [#207: cross-config logprob drift contract](https://github.com/RL-Align/RL-Kernel/issues/207) + +## Scope + +This contract describes the logical inputs and deterministic reduction semantics for standard +softmax Attention under tensor parallelism (TP) and context parallelism (CP). It lets runtime +dispatch reject a backend whose numerical semantics do not match the requested layout. + +This PR1 layer does not shard tensors, launch a collective, merge CP partial states, or implement +a fused kernel. The deterministic CP reference implementation and its distributed numerical tests +belong to later work in #235. + +## Contract Objects + +`rl_engine.kernels.attention_contract` defines: + +- `AttentionContract`: role, mode, dtype, causal metadata, sharding, reduction, and optional cache + identity; +- `ShardingSpec`: TP-local head ownership and CP block-to-token ownership; +- `ReductionSpec`: fixed `(out, lse)` merge semantics; +- `KVCacheSpec`: decode replay cache identity; +- `RoPESpec`: Qwen3 RoPE state, position identity, and fused/unfused boundary metadata; +- `AttentionBackendCapability`: the layouts and semantics a backend explicitly supports. + +Construction performs validation immediately. A structurally valid contract means that the +request is complete and internally consistent; it does not mean that an installed backend can +materialize it. + +`AttentionContract.batch_size` is the logical sequence count. For packed varlen input it must +equal `len(packed_sequence_offsets) - 1`; it is not the physical leading dimension of a flattened +token tensor. + +For full `prefill`, `query_sequence_length` equals the local sequence length described by +`ShardingSpec`. Chunked prefill and decode may use shorter query lengths than their available KV +context. + +## Qwen3-8B TP=2 CP=2 Example + +```python +from rl_engine.kernels.attention_contract import ( + AttentionContract, + ReductionSpec, + RoPESpec, + ShardingSpec, +) + +sharding = ShardingSpec( + tp_rank=0, + tp_world_size=2, + cp_rank=0, + cp_world_size=2, + global_q_heads=32, + global_kv_heads=8, + local_q_head_start=0, + local_q_heads=16, + local_kv_head_start=0, + local_kv_heads=4, + global_sequence_length=4096, + local_sequence_length=2048, + global_block_indices=(0,), + global_block_token_starts=(0,), + local_block_offsets=(0, 2048), +) + +contract = AttentionContract( + role="infer", + mode="prefill", + dtype="bf16", + batch_size=1, + query_sequence_length=2048, + head_dim=128, + causal=True, + causal_offsets=(0,), + sharding=sharding, + reduction=ReductionSpec(), + rope=RoPESpec( + q_state="post_rope", + k_state="post_rope", + k_cache_state="post_rope", + theta=1.0e6, + rotary_dim=128, + query_position_offsets=(0,), + key_position_offsets=(0,), + cast_at="after_rope", + output_dtype="bf16", + fusion_boundary="unfused_rope_attention", + ), +) +``` + +The TP fields preserve the global Qwen3 GQA mapping: each rank owns 16 of 32 query heads and 4 of +8 KV heads. The CP fields map local tensor slices to stable logical global block ids. A rank that +owns non-contiguous blocks uses one global token start per block and one extra local boundary: + +```python +global_block_indices=(0, 3) +global_block_token_starts=(0, 3072) +local_block_offsets=(0, 1024, 2048) +``` + +This metadata is sufficient for a later implementation to restore logical global order without +using ring arrival order. + +## RoPE / Position Semantics + +RoPE is part of the attention contract because rollout can materialize +`RoPE+Attention` as a fused or cache-aware path while training may materialize +`RoPE -> Attention` as separate operators. PR1 does not execute the RoPE kernel, +but it records the metadata required to prove both materializations use the same +model semantics. + +`RoPESpec` records: + +- whether Q, K, and cached K are `pre_rope` or `post_rope`; +- `theta`, optional `rope_scaling`, and `rotary_dim`; +- dense `position_ids` or per-sequence `query_position_offsets` / + `key_position_offsets`; +- the RoPE cast point and output dtype; +- `fusion_boundary`, either `unfused_rope_attention` or `fused_rope_attention`. + +When RoPE metadata is present, construction validates that rotary dimensions fit +the attention head dimension and that offset metadata matches the logical batch +shape. Backends must declare RoPE support through `AttentionBackendCapability`; +a backend that cannot consume RoPE/position metadata or cannot support a fused +RoPE+Attention boundary is rejected before dispatch. + +## Reduction Semantics + +The only PR1 reduction contract is: + +```text +partial state: (out, attention-domain lse) +merge: online_softmax_lse +acc_dtype: fp32 +order: global_block_index +downcast_at: final_write +engine: in_op_reference +``` + +CP output is not a plain sum. A backend that cannot export attention-domain LSE or cannot merge +partial states in fixed logical order is incompatible with this contract. + +The acceptable output and selected-logprob drift thresholds remain owned by #108. This contract +does not introduce another tolerance table. When connected to the rollout/training chain, the +selected-token metric remains the #207 convention: + +```text +dlogp = training-side recomputed logp - rollout-side old logp +``` + +## Mode-Specific Metadata + +All causal calls provide `causal_offsets`. Packed varlen calls provide one causal offset per +packed sequence and validated `packed_sequence_offsets`. + +Decode additionally requires `KVCacheSpec` with: + +- one cache position and KV sequence length per logical sequence; +- a block/page table; +- the physical page size; +- global token positions for every logical cached token; +- a prefix-cache key and explicit shared-prefix page count when prefix caching is enabled. + +Within each logical sequence, global token positions must be strictly increasing. Block-table +padding must be trailing, the active page count must match `ceil(kv_seq_len / page_size)`, and a +sequence cannot repeat one physical page id. Different sequences may share physical pages for an +equivalent prefix only when those pages are declared by `shared_prefix_page_count`, use the same +leading page ids and logical positions, and are fully populated. Declared shared prefix pages are +read-only; all suffix pages are exclusive to one sequence, providing the contract boundary needed +for copy-on-write before divergent decode. When prefix caching is disabled, no active page may be +shared across sequences. Missing or inconsistent decode cache identity is an error at contract +construction time. + +Each `cache_positions` entry is the terminal logical position already present in that sequence's +KV cache, so it must equal the final corresponding `global_token_positions` entry. It is not the +next position to be written. + +## Contract-Aware Dispatch + +Legacy callers continue to use `KernelRegistry.get_op()`. WS2 callers use: + +```python +result = kernel_registry.get_attention_op(contract) +op = result.op +provenance = result.provenance +``` + +Dispatch considers only backends with an `AttentionBackendCapability`. It checks role, attention +mode, dtype, TP/CP degree, LSE export, deterministic CP merge, packed varlen, and KV-cache support. +When RoPE metadata is present, dispatch also checks whether the backend explicitly supports +RoPE/position metadata and fused RoPE+Attention boundaries. +An undeclared or incompatible backend is skipped with an explicit rejection reason. + +The current WS1 PyTorch Attention implementations support local reference math but do not export +attention-domain LSE or materialize deterministic CP merge. Strict WS2 requests therefore fail +clearly today. A later deterministic backend becomes selectable by registering a capability that +truthfully declares those features; no grid-planner branch or silent fallback is required. + +Successful dispatch provenance records: + +- requested and actual backend ids; +- platform and fallback status; +- prior candidate rejection reasons; +- the complete requested contract; +- the selected backend capability descriptor. + +## Validation + +Contract and dispatch behavior are covered by: + +```bash +python -m pytest tests/test_attention_contract.py -q +``` + +The tests include Qwen3 TP=2/CP=2 construction, GQA ownership errors, non-contiguous CP blocks, +packed varlen metadata, decode cache identity, undeclared backend rejection, no incompatible +fallback, RoPE metadata validation, and JSON-compatible provenance. diff --git a/docs/design/ws2-module-debug-matrix.md b/docs/design/ws2-module-debug-matrix.md new file mode 100644 index 00000000..89302c9c --- /dev/null +++ b/docs/design/ws2-module-debug-matrix.md @@ -0,0 +1,53 @@ +# Module Debug Matrix + +The cross-configuration debug surface uses a fixed replay and changes exactly +one factor at a time. It compares two edges of the same frozen rollout: + +1. training score versus rollout prefill; +2. rollout prefill versus rollout decode. + +The source of truth is +`rl_engine.alignment.cross_config.debug_matrix.module_debug_matrix`. The +manifest is reporting metadata, not a second collection of runtime flags. + +## Comparison gates + +Before reporting numerical drift, the run must prove that the selected tokens, +active mask, model state, and logical ownership map are unchanged. A failed +gate records `comparable=false`; it does not produce a drift number. This keeps +an input or sharding error from being misdiagnosed as kernel rounding. + +## First-line axes + +| Module | Baseline | Gate rows | Diagnostic rows | Invariant control | +| --- | --- | --- | --- | --- | +| Attention | `A0` | `A4` head/sequence ownership | `A1` position/RoPE, `A2` Q/K preprocessing, `A3` mask boundary, `A5` KV state, `A6` precision/rounding, `A7` block plan/merge | `C0`-`C2` | +| FFN / GEMM | `F0` | `F1` TP weight ownership | `F2` SwiGLU rounding, `F3` K-reduction/Split-K, `F4` token collective | `FC0` batch-invariant row replay | +| Selected-token logp | `L0` | `L1` vocabulary ownership/padding domain, `L2` selected token/active mask | `L3` vocabulary LSE tile/merge | `LC0` batch-invariant row replay | + +The logp wrapper receives local logits and therefore owns the vocabulary +log-sum-exp, selected-token gather, mask, padding, and TP merge. It does not +own the language-model-head GEMM. The FFN/GEMM wrapper owns the K reduction and +SwiGLU rounding; both must be diagnosed before attributing a logp difference to +the vocabulary merge. + +## Numerical program + +For each diagnostic row, record the same inputs and constants, intermediate +precision and final write, reduction grouping and logical order, and every +state handoff. Scheduler choices that do not change the numerical program or +logical ownership are intentionally not exposed as mismatch axes. + +Use the compact row first. The operator-specific secondary probes remain +available for a focused follow-up only after the representative row reproduces +the mismatch. + +## Report artifacts + +The matrix reporter is owned by RL-Kernel because it consumes the sealed +cross-configuration attempt contract. It reads `actual.json`, +`comparison.json`, and `token_diffs.pt` only after `COMPLETE` validation, so +the report displays actual selected operators and the materialized topology, +not merely requested flags. See +[Cross-Configuration Drift Report](../usage/cross-config-drift-report.md) for +the `.rlk-drift` bundle, desktop viewer, static image, and trace entry points. diff --git a/docs/design/ws2_cross_config_logprob_drift_contract.md b/docs/design/ws2_cross_config_logprob_drift_contract.md deleted file mode 100644 index 3d7fac06..00000000 --- a/docs/design/ws2_cross_config_logprob_drift_contract.md +++ /dev/null @@ -1,874 +0,0 @@ -# WS2 Cross-Config Logprob Drift Contract - -Status: RFC - -Tracking issues: - -- [#111: WS2 cross-config alignment](https://github.com/RL-Align/RL-Kernel/issues/111) -- [#108: WS1 numerical contract](https://github.com/RL-Align/RL-Kernel/issues/108) - -## Motivation - -WS2 covers rollout and training paths that use different parallelism strategies, such as -rollout tensor parallelism and training FSDP. The alignment problem is not a single-op -accuracy check. It is end-to-end floating-point drift across tokenizer, masks, serving, -rollout, and training recomputation before any optimizer update. - -For PPO, GRPO, and related RL post-training algorithms, the most direct pre-update signal -is selected-token log probability drift. If rollout-side `old_logprobs` and train-side -recomputed log probabilities disagree for the same checkpoint, same token ids, same masks, -and same model version, classify the failure as infrastructure, precision, mask, -tokenizer, or serving-path drift. Do not classify that failure as an algorithm or reward -problem until pre-update logprob alignment is clean. - -Aggregate KL-style diagnostics are useful but not sufficient as the primary WS2 contract. -In training-inference mismatch cases, KL estimates can stay flat or fail to expose the -early failure phase, because the first-order issue is token-level rollout-vs-training -probability disagreement before the optimizer update, not necessarily a large aggregate -policy-space shift. - -## Framework Upgrade Overview - -The upgrade moves cross-config validation from two separately configured execution paths -that require a manual comparison into a shared runtime flow. `RuntimeTools` coordinates -the rollout and training configurations, while `PairedRunner` collects their selected -logprobs and performs the comparison automatically. This keeps both sides aligned on the -same inputs and makes drift visible as part of the run rather than as a follow-up manual -check. - -![Before and after framework upgrade](../assets/ws2-cross-config-before-after.png) - -The left side shows the pre-upgrade flow, where `VLLMSamplerConfig` and -`TorchRLTrainingConfig` feed independent executors and the results are manually compared. -The right side shows the upgraded flow, where `RuntimeTools` and `PairedRunner` connect the -two paths and produce an automatic comparison while preserving the shared -`KernelRegistry` contract. - -## Scope - -This RFC defines what WS2 cross-config alignment measures, how failures are classified, -and the modular implementation roadmap for making that contract executable. It does not -itself add a test harness, distributed tests, runtime gates, layer-wise probes, or -distributed fixes. - -Out of scope for this document: - -- Implementing multi-GPU test infrastructure. -- Adding runtime pass/fail gates. -- Adding automatic layer-wise drift probes. -- Fixing TP, FSDP, SP, cache, mask, tokenizer, or serving-path bugs. -- Defining a second numerical tolerance table. -- Reimplementing work owned by the adjacent TP, SP, collective, training-integration, or - layer-probe issues referenced by the roadmap below. - -## Measurement Contract - -The primary metric is selected-token logprob drift: - -```text -dlogp = train_recomputed_logp - rollout_old_logp -``` - -Compute `dlogp` only on active response/action tokens. Prompt tokens, padding tokens, and -masked-out response positions are excluded from every aggregate metric. - -The comparison must use teacher-forcing scoring on the training side. The scored sequence -is the already-sampled rollout sequence; the training path must not resample or regenerate -tokens for this contract. - -The rollout and training values are comparable only when they share the same logical -inputs: - -- Same checkpoint and same model version. -- Same input token ids. -- Same selected response/action token ids. -- Same attention mask and action mask. -- Same tokenizer version and tokenization policy. -- Same padding layout semantics, including left-padding or right-padding behavior. -- Same pre-update state, before any optimizer step, weight sync, or policy mutation that - belongs to the next training step. - -If the implementation has explicit position ids, cache-position metadata, sequence ids, or -packed-sequence metadata, those inputs are part of the comparison contract as well. - -## Primary Failure Signal - -The pass/fail decision starts from `dlogp` over active tokens. Reward, gradnorm, -weightnorm, and update norm are downstream symptoms. They are useful for debugging and -triage, but they are not the primary contract for cross-config alignment. - -The zero-update expectation is: - -```text -train_recomputed_logp ~= rollout_old_logp -ratio0 ~= 1 -approx_kl0 ~= 0 -``` - -The acceptable meaning of `~=` is defined by the WS1 per-dtype numerical threshold table -from [#108](https://github.com/RL-Align/RL-Kernel/issues/108). This RFC defines the -measurement surface and classification rules only. - -## Diagnostics - -All diagnostics are computed on active response/action tokens only. - -| Metric | Definition | Purpose | -| --- | --- | --- | -| `ratio0` | `exp(dlogp)` | Zero-update policy ratio implied by train-vs-rollout logprob drift. | -| `clipfrac0` | Mean indicator that `ratio0` falls outside the configured PPO/GRPO clip range. | Detects whether drift alone would trigger clipping before any update. | -| `approx_kl0` | Masked mean of `exp(dlogp) - 1 - dlogp`. | Zero-update approximate KL implied by logprob drift. | -| `mean_abs_dlogp` | Mean of `abs(dlogp)`. | Average selected-token drift. | -| `p95_abs_dlogp` | 95th percentile of `abs(dlogp)`. | Tail drift below outliers. | -| `p99_abs_dlogp` | 99th percentile of `abs(dlogp)`. | High-tail drift. | -| `max_abs_dlogp` | Maximum of `abs(dlogp)`. | Worst selected-token mismatch. | - -When the run is distributed, report optional per-rank versions of the same metrics. The -per-rank view should preserve enough metadata to identify the rollout rank, training rank, -parallelism mode, dtype, padding side, cache mode, and local active-token count for that -rank. - -## Tolerance Source - -This RFC does not define a separate numerical tolerance table. The single source of truth -for acceptable numerical drift is the per-dtype threshold table owned by -[#108](https://github.com/RL-Align/RL-Kernel/issues/108). - -For WS2, acceptable numerical drift means that `max_abs_dlogp` over active -response/action tokens satisfies the WS1 per-dtype threshold from #108. If the #108 table -changes, WS2 inherits that policy without editing this document or maintaining a second -table. - -## Tolerance Interpretation and Effect-Based Validation - -Numerical tolerances in this RFC are infrastructure contract thresholds, not a universal -statement of algorithmic harmlessness. There is no model-independent scale that proves a -given train-vs-rollout logprob difference is harmless for every algorithm, reward model, -prompt distribution, sequence length, or optimization schedule. Any hand-written threshold -encodes a prior about acceptable numerical error. WS2 therefore does not introduce an -additional algorithmic noise budget, nor does it define a new estimator for tolerable -logprob noise. - -The #108 threshold defines whether rollout and training paths are numerically aligned -enough to continue debugging the failure as an algorithmic or reward problem. It does not -prove that all smaller drift is behaviorally irrelevant, and it does not imply that all -larger drift is the only cause of downstream failure. - -When downstream model-effect validation is available, such as reward trajectory, train KL, -eval win rate, collapse rate, policy regression tests, or task-specific success metrics, -use it as a severity and root-cause prioritization signal. It must not replace the -pre-update selected-token logprob contract. A run can be numerically out of contract even -if a short downstream run appears healthy, and a run can be numerically in contract while -still failing because of algorithmic tuning, reward hacking, insufficient KL control, or -data issues. - -The intended interpretation is: - -```text -#108 per-dtype threshold: - numerical infrastructure contract - -selected-token dlogp: - primary WS2 train-vs-rollout drift surface - -downstream model effect: - practical severity and algorithmic relevance signal - -KL / ratio / percentile diagnostics: - debugging and triage signals, not replacement pass/fail criteria -``` - -## Drift Source Taxonomy - -Before treating train-vs-rollout drift as generic algorithmic noise, WS2 should classify -likely sources of mismatch. At minimum, the following source classes should be considered -separately. - -### Arithmetic Schedule Drift - -Arithmetic schedule drift comes from different floating-point operation order between -rollout and training. This includes different kernels, fused vs unfused implementations, -compiler-generated graph rewrites, attention implementation differences, matmul epilogue -differences, accumulation dtype differences, and changes introduced by advanced compilers -or graph optimizers. - -This class answers the question: - -```text -Do rollout and training compute mathematically equivalent expressions using different -floating-point schedules? -``` - -Examples include: - -- Fused attention vs unfused attention. -- Different FlashAttention or SDPA backends. -- Fused RMSNorm or LayerNorm vs decomposed normalization. -- Compiler-reordered graph segments. -- Different matmul epilogues or activation fusion. -- Different accumulation precision in otherwise equivalent kernels. - -### Reduction and Collective Drift - -Reduction drift comes from operations whose floating-point result depends on reduction -order, parallel topology, or concurrent execution. This includes local reductions, -cross-rank reductions, all-reduce, reduce-scatter, gather/scatter patterns, sharded logits -or loss computation, tensor-parallel collectives, FSDP reductions, and nondeterministic -reduction scheduling. - -This class answers the question: - -```text -Does the mismatch appear because rollout and training aggregate partial results in -different orders or across different rank topologies? -``` - -Examples include: - -- TP logits produced through a different collective path from the training path. -- FSDP reduce-scatter or all-gather changing accumulation order. -- Per-rank partial reductions with different shard boundaries. -- Loss or logprob reductions performed before vs after cross-rank communication. -- Nondeterministic collective algorithms or concurrent reductions. - -### Quantization and Dequantization Drift - -Quantization drift comes from representing weights, activations, KV cache, logits, or -intermediate tensors with different quantization policies between rollout and training. -Quantization is not merely a floating-point ordering issue; it introduces representation -noise through scales, zero points, clipping, grouping, calibration, and dequantization -paths. - -This class answers the question: - -```text -Does the mismatch appear because rollout and training use different numerical -representations or quantization policies? -``` - -Examples include: - -- Rollout uses weight-only quantization while training recomputation uses bf16/fp16 - weights. -- Different quantization group sizes. -- Different activation quantization or KV-cache quantization policy. -- Different scale computation or calibration data. -- Different dequantization placement relative to fused kernels. -- Serving-path quantization that is absent from the training path. - -### Logical Input and Metadata Drift - -Logical input mismatch must be ruled out before interpreting any result as numerical -drift. This class includes tokenizer version, tokenization policy, attention mask, action -mask, padding side, explicit position ids, cache positions, sequence ids, packed-sequence -metadata, and serving-path request formatting. - -This class answers the question: - -```text -Are rollout and training actually scoring the same logical sequence under the same masking -and positional semantics? -``` - -If this class is not clean, the comparison is invalid rather than merely noisy. - -## Decision Rule - -Use this order when classifying a cross-config failure: - -1. If pre-update selected-token logprobs do not match under the same checkpoint, same token - ids, same masks, and same model version, treat the failure as infrastructure, - precision, mask, tokenizer, or serving-path drift. -2. If `max_abs_dlogp` violates the #108 threshold but downstream metrics look healthy in a - short run, keep the issue classified as infrastructure drift. Short-horizon model - health does not prove the drift is safe. -3. If KL or ratio diagnostics move before gradnorm or update norm moves, treat the failure - as likely infrastructure or logprob plumbing. -4. If gradnorm or update norm moves first and KL moves later, treat the failure as more - likely algorithmic tuning, such as learning rate, KL beta, reward scale, or advantage - outliers. -5. If only some ranks drift, treat the failure as distributed infrastructure until rank - placement, shard boundaries, collective algorithms, local active-token counts, masks, - and cache-position issues are ruled out. -6. If reward rises and then collapses while pre-update logprob alignment is clean, treat - the failure as more likely algorithmic, reward hacking, data-related, or insufficient - KL constraint. - -This classification does not prove root cause by itself. It defines the first branch in -the debugging tree so WS2 bugs do not get misfiled as reward or algorithm regressions -before the zero-update logprob contract is satisfied. - -## Layered Ablation Strategy - -WS2 should not treat train-vs-rollout mismatch as a single undifferentiated error source. -Later tests should use a layered ablation strategy that changes one source class at a time -whenever the implementation allows it. - -The minimum useful ablation structure is: - -```text -A0. Fully aligned reference - Same checkpoint, same dtype policy, same kernels where possible, same reduction - topology where possible, same quantization policy, same tokenizer, same masks, same - padding, same cache/position metadata. - -A1. Arithmetic-schedule-only mismatch - Keep logical inputs, reduction topology, and quantization policy aligned. Allow only - kernel, fusion, compiler, or graph execution differences. - -A2. Reduction-topology-only mismatch - Keep logical inputs, kernel policy, and quantization policy aligned. Allow only - reduction order, collective topology, sharding, or rank placement differences. - -A3. Quantization-only mismatch - Keep logical inputs, kernel policy, and reduction topology aligned. Allow only - quantization, dequantization, scale, group size, or representation differences. - -A4. Pairwise mismatches - Enable two mismatch classes at a time: - arithmetic + reduction - arithmetic + quantization - reduction + quantization - -A5. Full production mismatch - Use the real rollout and training configurations, including all production - differences. -``` - -Each ablation should collect the same primary and diagnostic metrics: - -```text -primary: - dlogp over active response/action tokens - max_abs_dlogp - -diagnostics: - mean_abs_dlogp - p95_abs_dlogp - p99_abs_dlogp - ratio0 - clipfrac0 - approx_kl0 - per-rank versions when distributed - -metadata: - dtype - kernel/backend choices - fusion/compiler mode - reduction/collective topology - quantization policy - padding side - cache mode - position/cache-position metadata - active-token count -``` - -When downstream model-effect validation is available, the same ablations should also -record practical training outcomes, for example reward trajectory, training KL, entropy, -clip fraction, update norm, collapse rate, and task-specific evaluation metrics. These -downstream metrics are not the WS2 pass/fail contract, but they help rank which numerical -mismatch class matters most for the workload. - -## Ablation Interpretation Rules - -Use these rules when reading the ablation matrix: - -1. If the fully aligned reference fails, the issue is not a cross-config mismatch yet. - First debug the base scoring path, masks, tokenizer, position metadata, checkpoint - identity, or implementation correctness. -2. If a single-source ablation fails the `max_abs_dlogp` contract, that source class is - sufficient to create unacceptable train-vs-rollout drift under the tested workload. For - example, if only quantization is misaligned and the run fails, quantization is a - dominant source candidate for that task and configuration. -3. If all single-source ablations pass, but pairwise or full-production mismatches fail, - the failure is likely an interaction effect. Identify the minimal failing pair before - attributing the issue to any single subsystem. -4. If one single-source ablation passes the numerical contract but shows materially worse - downstream model effect, record it as behaviorally sensitive even if it remains - numerically in contract. This is a signal that the #108 infrastructure tolerance may be - sufficient for numerical alignment but not necessarily predictive of algorithmic - robustness for that workload. -5. If pre-update logprob alignment is clean but downstream training still collapses, - classify the failure as more likely algorithmic, reward-related, data-related, or - KL-control-related rather than cross-config numerical drift. - -## Minimal and Layered Alignment Principle - -The governing principle of WS2 is **minimal alignment**: - -> Keep rollout and training semantically identical, then align only the smallest numerical -> layer needed to satisfy the selected-token logprob contract. - -WS2 does not require every internal tensor, kernel, reduction, or execution schedule to be -identical. If the production rollout and training paths already satisfy the #108 -`logprob` tolerance, no numerical alignment change is required. Different engines are -allowed to keep different high-performance implementations. - -Minimal alignment does not relax logical correctness. Checkpoint/version, token ids, -masks, tokenizer semantics, and required position metadata must match exactly. A logical -input mismatch invalidates the experiment; it is not acceptable numerical drift. - -### Alignment Ladder - -Use the following ladder in order and stop at the first level that satisfies the contract: - -| Level | Action | Production implication | -| --- | --- | --- | -| L0: semantic identity | Make logical inputs and model version exactly comparable. | Mandatory for every case. | -| L1: observable contract | Keep both production paths unchanged and compare selected-token logprobs. | Stop here if #108 passes. | -| L2: source isolation | Change one declared knob at a time to locate the smallest sufficient drift source. | Diagnostic only; do not change production yet. | -| L3: local alignment | Align or fix one operator, collective, metadata field, or representation policy. | Preferred production fix when L1 fails. | -| L4: layered alignment | Align the smallest interacting pair or contiguous layer boundary that is required. | Use only when no single local change is sufficient. | -| L5: full/bitwise alignment | Force broad identical paths or reference implementations. | Diagnostic fallback, not the default WS2 exit criterion. | - -The chosen fix should minimize, in order: - -1. semantic scope changed; -2. number of aligned knobs; -3. performance and memory overhead; -4. engine-specific intrusion; -5. maintenance burden. - -A fix is incomplete if it proves only that the fully aligned reference passes. It must -also show that unrelated rollout/training differences can remain enabled. Conversely, WS2 -must not reject a configuration merely because internal tensors are not bitwise equal when -the selected-token contract passes. - -## Controller-Centered Design - -The central feature is an ablation controller, not a hard-coded list of distributed -tests. It separates experiment planning from engine-specific knob application. - -```mermaid -flowchart LR - Definition["ExperimentDefinition
identity + baseline + axes + constraints"] - Planner["GridPlanner
product / one-at-a-time / pairwise"] - Isolation["IsolationValidator
declared deltas only"] - Definition --> Planner --> Isolation - - Isolation --> Cases["ExperimentCase[]
stable ids + provenance"] - - subgraph Materializers["Knob materializers"] - Rollout["vLLM adapter"] - Training["stateless / FSDP adapter"] - Kernel["kernel policy adapter"] - Environment["process/build environment adapter"] - end - - Cases --> Materializers - Materializers --> Runner["isolated paired runner"] - Runner --> Samples["canonical alignment samples"] - Samples --> Comparator["identity validator + dlogp comparator"] - Comparator --> Cube["result cube
axes + per-rank reports + cost"] - Cube --> Analyzer["minimal sufficient alignment analyzer"] -``` - -### Core Objects - -The implementation should expose a small typed model rather than passing more loose -dictionaries through the current executors: - -- `SemanticIdentitySpec`: checkpoint/weight version, tokenizer, fixed token sequences, - masks, and position metadata that must match. -- `ScorerSpec`: rollout or training engine, world size, device/dtype, and immutable engine - construction settings. -- `KnobDefinition`: one controllable source of variation. -- `ExperimentDefinition`: baseline scorers plus axes, constraints, and measurement policy. -- `ExperimentCase`: one fully materialized grid point with a stable content-derived id. -- `AlignmentSample`: logical tensors, selected logprobs, and actual runtime provenance. -- `AlignmentResult`: global/per-rank drift, pass/fail, actual applied knobs, and optional - cost metrics. -- `ResultCube`: results indexed by normalized knob values, independent of execution order. - -Every `KnobDefinition` must declare: - -```text -name: - stable dotted name, for example rollout.tensor_parallel_size - -source_class: - logical-layout | arithmetic | reduction | representation | execution - -lifecycle: - request | engine-construction | process-start | build - -targets: - rollout | training | both | kernel - -domain: - allowed typed values - -capability: - how an adapter proves that a value is supported - -constraints: - incompatible or conditional combinations - -apply: - engine-specific materialization hook - -provenance: - how the actual applied value is read back and reported -``` - -The controller must compare requested and actual provenance. A silent runtime fallback is -an invalid ablation unless the fallback itself is the declared knob under test. - -### Grid Composition - -`GridPlanner` should support the following modes over the same typed axes: - -- `product`: full Cartesian grid; -- `one_at_a_time`: baseline plus one changed factor per case; -- `pairwise`: covering pairs without requiring the full Cartesian product; -- `zip`: paired values such as compatible model/dtype artifacts; -- fixed overrides and named slices; -- capability and compatibility constraints; -- deterministic case ids, filtering, resume, and retry. - -A normal workflow starts with `one_at_a_time`, expands to `pairwise` only when single -factors do not explain the failure, and uses `product` for an explicit grid search. CI -runs a named slice of the same definition rather than maintaining a separate handwritten -test matrix. - -For every generated case, `IsolationValidator` compares its normalized spec with the -baseline and rejects undeclared changes. This is what makes an arithmetic-only, -reduction-only, or quantization-only claim trustworthy. - -### Minimal Sufficient Alignment Analysis - -The analyzer treats "align this knob between rollout and training" as an intervention. It -reports the smallest passing intervention set found by the executed grid: - -```text -production mismatch: - fail - -align attention backend only: - fail - -align logp reduction only: - pass - -minimal sufficient alignment candidate: - {logp.reduction_policy} - -unrelated differences left enabled: - attention backend, cache policy, TP/FSDP topology -``` - -This result is evidence for the smallest effective intervention, not automatic proof of -root cause. A later fix PR still needs the smallest reproducer and a local regression. - -## Mapping to Current Code - -The controller should initially map to existing configuration surfaces instead of -introducing a second execution stack. - -| High-level knob | Current code path | Required adapter behavior | -| --- | --- | --- | -| `rollout.tensor_parallel_size` | `VLLMSamplerConfig.engine_kwargs` | Materialize `tensor_parallel_size` before vLLM engine construction and read it back from runtime metadata. | -| `rollout.dtype` | `VLLMSamplerConfig.engine_kwargs["dtype"]` | Normalize string/torch dtype and record the actual engine dtype. | -| `sampling.temperature` | `VLLMSamplerConfig.sampling_params` | Apply per request; require the same scoring semantics on both sides. | -| `execution.prefix_cache` | `VLLMSamplerConfig.enable_prefix_caching` | Treat as engine-construction-time, not a request toggle. | -| `training.attention_backend` | `StatelessForwardConfig.attention_backend` | Apply before forward and report requested backend plus any actual fallback. | -| `training.output_dtype` | `StatelessForwardConfig.output_dtype` | Keep observation dtype separate from model compute dtype. | -| `training.compute_dtype` | `TorchRLTrainingConfig.dtype` and FSDP model construction | Materialize before wrapping/sharding the model. | -| `logp.backend` | `RolloutExecutor` / `TorchRLTrainingConfig.logp_backend` | Reuse `resolve_logp_op_type()` aliases and report the resolved op type and concrete backend class. | -| `logp.deterministic` | `require_batch_invariant_logp` | Express policy intent; do not hard-code a CUDA implementation in the controller. | -| `training.sharding` | new score-only FSDP adapter | Materialize world size and sharding strategy before process-group/model construction. | -| `logp.tp_layout` | `linear_logp` `tp_group`, `vocab_start_index`, `global_vocab_size` | Record shard boundaries and reject incomplete ownership metadata. | -| `kernel.fast_math` | `KERNEL_ALIGN_USE_FAST_MATH` | Treat as build-time and bind the case to a distinct built artifact. | -| `kernel.sm90_path` | `KERNEL_ALIGN_FORCE_SM90` and compiled extension | Capability-gate by architecture and build artifact; never switch it after import. | - -The existing `KernelRegistry` caches instances and resolves priority maps during -initialization. vLLM TP, dtype, and prefix caching also belong to engine construction. -Therefore the runner must not mutate these values in a long-lived process and assume the -next case is isolated. - -Cases may share a worker only when their engine-construction and process-start -fingerprints are identical. Request-time knobs may reuse that worker. Build-time knobs -always select a prebuilt artifact and a separate process. The artifact id and extension -build metadata are part of result provenance. - -## Kernel Integration Contract - -Kernel work may require a new or rewritten implementation, but the ablation controller -must not know CUDA/Triton class names or kernel launch details. - -The kernel boundary should expose a backend descriptor with: - -- stable backend id and semantic operator name; -- supported device architectures, dtypes, shapes, and parallel layouts; -- determinism/alignment properties; -- required TP/SP metadata and collectives; -- configuration lifecycle, including build-time flags; -- concrete implementation selected at runtime; -- fallback behavior; -- version/build fingerprint. - -The controller requests a policy such as `production`, `reference`, `deterministic`, or a -stable backend id. The kernel adapter resolves that policy through `KernelRegistry` and -records the concrete implementation. Strict WS2 cases reject an undeclared fallback. - -A rewritten kernel integrates cleanly by: - -1. implementing the existing operator semantic interface; -2. registering a new stable backend descriptor; -3. passing #108 operator accuracy and batch-invariance checks; -4. declaring TP/SP metadata and supported lifecycle knobs; -5. adding one isolated end-to-end controller case; -6. reporting performance/memory overhead against the production backend. - -It should not require a new branch in `GridPlanner`. If a framework cannot inject the -kernel through a supported hook, its engine adapter reports the knob as unsupported; it -must not claim that the ablation ran. - -## Repository Fit - -The current repository already provides useful pieces: - -- #108 owns `tolerance_contract.json`. -- `VLLMSamplerConfig` exposes loose `engine_kwargs`, `sampling_params`, and prefix-cache - configuration. -- `StatelessForwardConfig` exposes attention backend, temperature, and output dtype. -- `TorchRLTrainingConfig` exposes compute dtype, `logp_backend`, and the deterministic - requirement. -- `resolve_logp_op_type()` already separates user-facing logp policy from registry op type. -- TP `linear_logp` already accepts explicit process group and vocab-shard metadata. -- `RolloutStageResult` and the weight bridge carry iteration/weight version. -- `StatelessForwardExecutor` is a reusable no-update teacher-forcing scorer. - -The missing pieces are the typed experiment model, actual-value provenance, strict scoring -payload, FSDP score-only adapter, lifecycle-aware knob materializers, grid planner, and -result cube. - -`DeepSpeedTrainingWorker.train()` still performs backward/step and constructs its current -objective's `old_logps` from recomputed values. It is not a WS2 comparator. A later -DeepSpeed scorer must be a separate read-only adapter. - -## Ownership Boundaries - -| Issue | Boundary | -| --- | --- | -| [#108](https://github.com/RL-Align/RL-Kernel/issues/108) | Owns numerical thresholds. | -| [#109](https://github.com/RL-Align/RL-Kernel/issues/109) | Owns deterministic TP reduction implementations. | -| [#110](https://github.com/RL-Align/RL-Kernel/issues/110) | Owns SP-aware operators and reductions. | -| [#112](https://github.com/RL-Align/RL-Kernel/issues/112) | Owns deterministic collective implementations. | -| [#113](https://github.com/RL-Align/RL-Kernel/issues/113) | Owns the later distributed forward/backward chain gate. | -| [#116](https://github.com/RL-Align/RL-Kernel/issues/116) | Shares the tolerance/report foundation implemented by B1. | -| [#127](https://github.com/RL-Align/RL-Kernel/issues/127) | Owns the pinned multi-GPU dual-engine environment. | -| [#130](https://github.com/RL-Align/RL-Kernel/issues/130) | Owns full FSDP/Megatron training integration and backward. | -| [#131](https://github.com/RL-Align/RL-Kernel/issues/131) | Owns the later production cross-benchmark command. | -| [#136](https://github.com/RL-Align/RL-Kernel/issues/136) | Owns automatic layer-wise probes. | - -## Revised Modular PR Roadmap - -The identifiers below are roadmap labels, not existing GitHub PR numbers. The former -Phases A, B, and C are consolidated because they jointly form the baseline infrastructure. - -### Phase 1: Baseline Infrastructure - -#### B1 — Alignment contract, comparator, and report - -**Scope:** Expose #108 tolerance lookup; add canonical identity/provenance/sample types, -logical comparability validation, active-token drift metrics, and one JSON/human report. - -**Acceptance:** CPU tests cover identity mismatch, masks, percentiles, zero active tokens, -worst-token metadata, and dtype-specific pass/fail without copying threshold values. - -**Why one PR:** These types form one public contract and cannot provide useful independent -behavior when landed separately. - -#### B2 — Exact rollout and teacher-forcing scoring adapters - -**Scope:** Normalize vLLM sampled-token logprobs and rollout provenance; add strict -rollout-to-teacher-forcing collation; define the read-only scorer protocol and adapt -`StatelessForwardExecutor`. - -**Acceptance:** A fixture round trip preserves prompt/generated ids, masks, selected -logprobs, weight version, and available position metadata. Missing identity data or -undeclared backend fallback fails explicitly. Repeated scoring does not change model state. - -**Non-goal:** No FSDP, subprocess runner, or grid planner. - -#### B3 — Score-only FSDP adapter and baseline controls - -**Scope:** Add a PyTorch FSDP scorer with no optimizer/backward, then add A0 identical -stateless scoring and unsharded-vs-FSDP controls. - -**Acceptance:** A0 passes on CPU; a labeled two-GPU/NCCL control proves FSDP recomputation -is clean and model state is unchanged. - -**Non-goal:** Full training integration remains in #130. - -#### B4 — Paired runner, artifacts, and rank aggregation - -**Scope:** Launch rollout/training scorers with independent world sizes; write versioned -canonical artifacts; enforce timeout/cleanup; aggregate deterministic per-rank/global -reports. - -**Acceptance:** CPU fixtures cover child failure, timeout, malformed artifact, duplicate -or missing ranks, weight-version mismatch, and global worst-token selection. - -**Design requirement:** The runner accepts separate construction/process/build -fingerprints so the later controller can isolate cases correctly. - -### Phase 2: Composable Ablation Controller - -#### C1 — Typed experiment model, knob registry, and grid planner - -**Scope:** Implement `ExperimentDefinition`, typed knob descriptors, constraints, -capability declarations, stable case ids, and `product`, `one_at_a_time`, `pairwise`, and -`zip` planners. - -**Acceptance:** Pure CPU tests generate deterministic grids, reject invalid combinations, -resume by case id, and prove each one-at-a-time case changes exactly one declared knob. - -**Non-goal:** Do not launch engines in this PR. - -#### C2 — Runtime knob materializers and capability checks - -**Scope:** Map controller knobs to current vLLM, stateless, and FSDP configuration -surfaces. Separate request-time, engine-construction, and process-start application. Read -back actual values and construction fingerprints. - -**Acceptance:** Fake engine adapters prove every requested value is either applied and -reported or rejected as unsupported. No silent fallback is accepted in strict cases. - -#### C3 — Kernel policy bridge - -**Scope:** Add the backend descriptor and kernel materializer boundary described above. -Adapt existing logp policy aliases and TP metadata without changing kernel math. - -**Acceptance:** The same experiment definition can select production/reference/ -deterministic logp policies and report the concrete registry backend. A fake rewritten -kernel registers without a controller code change. - -**Non-goal:** Kernel rewrites discovered later remain one-root-cause fix PRs. - -#### C4 — Grid executor and result cube - -**Scope:** Execute C1 cases through B4, pool only workers with identical lifecycle -fingerprints, select build artifacts, persist results, and expose filtering/resume plus a -machine-readable result cube. - -**Acceptance:** An interrupted fake grid resumes without rerunning completed cases; -requested and actual provenance are queryable for every axis. - -### Phase 3: Core Scenario and Minimal Alignment - -#### M1 — TP=2 rollout versus FSDP diagnostic grid - -**Scope:** Define the first real experiment using the controller: fixed model/tokenizer/ -tokens, vLLM TP=2 rollout, FSDP recomputation, bf16, and production defaults. Generate the -production point plus one-at-a-time alignment interventions. - -**Acceptance:** Execution and reports succeed on the pinned #127 environment. Numerical -failure is recorded without weakening #108. - -#### M-FIX-N — One minimal root cause per PR - -Each fix PR consumes the smallest controller case that exposes one problem. A kernel -rewrite, collective change, metadata fix, or adapter fix remains separate. - -A fix must show: - -- the failing production or isolated case; -- the smallest intervention that makes it pass; -- one local implementation change; -- A0 and unrelated-knob regressions; -- actual backend provenance; -- performance/memory cost when applicable. - -#### M2 — Promote the minimally aligned core case to a gate - -**Scope:** After required M-FIX PRs, gate TP=2/FSDP using the smallest passing alignment -set, not a fully reference configuration. - -**Acceptance:** The report names which knobs were aligned, which differences remained -enabled, and why a broader alignment level was unnecessary. - -### Phase 4: Grid Coverage and Ablation Closure - -#### G1 — Required composable grid - -**Scope:** Add the required batch-size, padding/layout, dtype, and cache/position axes as -declarative knob values and constraints. Allow full product, named slices, and -one-at-a-time views from the same definition. - -**Acceptance:** A user can request, for example: - -```text -batch_size = [1, 8] -padding_side = [left, right] -dtype = [fp32, bf16, fp16] -prefix_cache = [off, on] -logp.backend = [production, deterministic] -``` - -without writing a new test function. Unsupported combinations are capability-filtered -with explicit reasons, and every result is indexed in the same cube. - -#### G2 — A0-A5 profile and minimal-alignment wrapper - -**Scope:** Express the RFC's A0-A5 ablations as presets over C1 rather than separate test -implementations: - -- A0: fully aligned diagnostic reference; -- A1: arithmetic one-at-a-time; -- A2: reduction/topology one-at-a-time; -- A3: representation/quantization one-at-a-time; -- A4: pairwise expansion only when needed; -- A5: production mismatch. - -Add a CLI/config wrapper that selects profiles, axes, filters, and output location. - -**Acceptance:** Phase F behavior is only a planner/profile layer over G1. It adds no -engine-specific branching. - -#### G3 — Targeted GPU CI and downstream handoff - -**Scope:** Run a curated named slice of the same grid in labeled GPU CI, upload the result -cube, and expose fixtures/reports to #113/#131. - -**Acceptance:** CI distinguishes launch/environment/numerical failure, always cleans up, -and does not maintain a second handwritten matrix. - -## PR Sizing Rules - -The consolidated roadmap uses fewer baseline PRs, but later numerical fixes remain small: - -1. B1-B4 may each land one cohesive baseline subsystem. -2. C1-C4 each own one controller layer: planning, runtime materialization, kernel policy, - or execution/results. -3. Adding a new ordinary knob changes one descriptor and one engine adapter, not the - planner. -4. Adding a rewritten kernel changes the kernel implementation and its backend descriptor, - not the controller. -5. M-FIX PRs contain one root cause only. -6. G1/G2 add declarative grids/profiles and must not include numerical fixes. -7. No PR adds or copies a tolerance value. - -## Completion Criteria for #111 - -#111 is complete when: - -1. semantic identity validation is strict and independent of numerical alignment; -2. the controller can compose, filter, resume, and report a multidimensional configuration - grid; -3. requested knobs are verified against actual runtime/kernel provenance; -4. the TP=2/FSDP production case is brought into #108 contract using the documented - smallest sufficient alignment set; -5. batch, padding/layout, dtype, and cache/position axes are available through G1 without - new test functions; -6. at least one production, one arithmetic, and one reduction/topology slice run through - the same controller/report path; -7. a rewritten kernel can register through C3 without changing grid-planner code; -8. the stable core and selected grid slice run in targeted GPU CI; -9. forward fixtures and result cubes are reusable by #113 and #131. - -SP, additional TP sizes, quantization variants, exhaustive product grids, and downstream -training effects remain extensions unless maintainers promote named grid slices into the -required gate. Full/bitwise internal alignment is not a completion criterion unless a -separate contract explicitly requires it. diff --git a/docs/getting_started/faq.md b/docs/getting_started/faq.md index 878c8c14..27cdd0c1 100644 --- a/docs/getting_started/faq.md +++ b/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]" ``` diff --git a/docs/getting_started/installation.md b/docs/getting_started/installation.md index 71fe227d..0e8bad46 100644 --- a/docs/getting_started/installation.md +++ b/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 diff --git a/docs/operators/batch-invariant-logp.md b/docs/operators/batch-invariant-logp.md index 7e93a60d..bb05c038 100644 --- a/docs/operators/batch-invariant-logp.md +++ b/docs/operators/batch-invariant-logp.md @@ -27,7 +27,7 @@ logp = batch_invariant_logp( logits, # [B, T, V] or [N, V], differentiable target_ids, # [B, T] or [N], int ignore_index=-100, - validate=False, # Accelerated fast path; use True to check target range + 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 @@ -38,7 +38,6 @@ logp.sum().backward() # gradients flow into logits only | Backend | Wrapper | Status | | --- | --- | --- | | CUDA (SM90 TMA) | `BatchInvariantLogpSM90Op` | Hopper TMA online-softmax forward. | -| Ascend (CANN) | `BatchInvariantLogpAscendOp` | Ascend C two-pass streaming forward; PyTorch-formula backward. | | 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. | @@ -47,7 +46,6 @@ Current dispatch: ```text CUDA (Hopper, SM90 kernel compiled): CUDA (SM90 TMA) -> Triton -> PyTorch CUDA / ROCm (otherwise): Triton -> PyTorch -Ascend NPU: Ascend -> PyTorch CPU: PyTorch ``` @@ -56,28 +54,74 @@ 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). -The Ascend backend lives on the `npu` platform key and is available when the -extension exposes `_C_npu.batch_invariant_logp_ascend` (built with -`KERNEL_ALIGN_FORCE_ASCEND=1` on a CANN + torch_npu host; `npu-arch` defaults -to `dav-2201`, override with `KERNEL_ALIGN_ASCEND_ARCH`). When the extension is -not compiled, instantiation fails and dispatch falls through to PyTorch native. -bf16/fp32 NPU tensors run the Ascend C kernel; anything else (e.g. fp16) -silently falls back to the native op. +## Tensor Parallel + +`VocabParallelLogprobOp` +(`rl_engine/kernels/ops/pytorch/loss/vocab_parallel_logp.py`) +**TP=1, TP=2, and TP=4 produce bit-identical results.** + +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`. + +The TP-aware implementation uses the following fixed-order construction: + +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) +``` + +### Vime CP=2 runtime provider + +The optional Vime adapter is owned by RL-Kernel and can be selected without +patching Megatron or vLLM: + +```text +--selected-logprob-provider rl_engine.integrations.vime.logp.provider +--selected-logprob-provider-mode strict +``` + +Vime passes the local `[T, V_local]` logits, shifted targets, TP subgroup, +and CP row-ownership metadata. The provider builds the same `LogprobContract` +used by the distributed report, dispatches the explicit +`pytorch-vocab-parallel-logp-ws2` backend, and returns selected logp as `[T, 1]`. +When entropy is requested, it uses the same fixed TP-rank order and returns +full-vocabulary entropy for the existing loss surface. CP rank/layout are +recorded in provenance and never participate in the vocabulary LSE merge. + +The provider fails closed for undeclared real/padded vocabulary sizes, TP/CP +metadata mismatches, unsupported top-p replay masks, and backend fallback. +`auto` mode may then use Vime's native path; `strict` mode reports the +configuration error. This adapter does not import Vime. ## Benchmarks -`benchmarks/benchmark_batch_invariant_logp.py` compares Native, Triton when -available, and the active device's CUDA SM90 or Ascend backend (forward latency -and peak device memory across a vocab sweep, bf16): +`benchmarks/benchmark_batch_invariant_logp.py` compares Native, Triton, and the +CUDA SM90 backend (forward latency and peak VRAM across a vocab sweep, bf16): ```bash python benchmarks/benchmark_batch_invariant_logp.py python benchmarks/benchmark_batch_invariant_logp.py --configs "4096,128256;8192,151936" ``` -The hardware-specific column is shown only when the matching kernel for the -active device is compiled in. An NPU run never selects a CUDA kernel, even on a -host where both device types are visible. +The CUDA column is only shown when the SM90 kernel is compiled in; otherwise the +benchmark reports Native vs Triton only. ### Measured results @@ -153,11 +197,10 @@ 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. Accelerated -backends default to `validate=False` to avoid device synchronization in training -hot paths. With validation disabled, every non-ignored target must already be in -`[0, V)`; violating this precondition has undefined results and may fail during -backward. Use `validate=True` during debugging or with untrusted targets. +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 @@ -172,10 +215,6 @@ The operator is designed so each row is computed independently: - 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. -- The Ascend forward strides rows across blocks, so one AI core block owns - exactly one row; the vocab is scanned left-to-right in fixed - `TILE_LENGTH=4096` tiles with a two-pass (max, then sum-exp) fixed-order - reduction. - No atomic writes are used. These constraints ensure the result for a row depends only on that row's logits @@ -223,24 +262,24 @@ out.sum().backward() python -m pytest tests/test_batch_invariant_logp.py -q -rs ``` -All backends (Native, Triton, SM90, Ascend) are tested in a single file. -Coverage includes: correctness, empty batches, leading-shape preservation, -batch-invariance (bitwise), validation, ignore-index behavior, backward -correctness, registry dispatch, and dtype- and backend-specific smoke cases. +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. SM90 tests skip without a -Hopper build; Ascend tests skip without an NPU + `_C_npu` build. On Windows, run -via WSL/Linux with CUDA. +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/ops/cuda/loss/batch_invariant_logp.py` -- `rl_engine/kernels/ops/ascend/loss/batch_invariant_logp.py` - `csrc/cuda/batch_invariant_logp_kernel_sm90.cu` -- `csrc/ascend/batch_invariant_logp_ascend.asc` - `rl_engine/kernels/registry.py` -- `rl_engine/platforms/device.py` - `tests/test_batch_invariant_logp.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` diff --git a/docs/usage/cross-config-drift-report.md b/docs/usage/cross-config-drift-report.md new file mode 100644 index 00000000..3902bc5c --- /dev/null +++ b/docs/usage/cross-config-drift-report.md @@ -0,0 +1,40 @@ +# Cross-Configuration Drift Report + +The cross-configuration runner writes a sealed, append-only attempt directory. +The report tool reads that directory after `COMPLETE` has been published and +renders the recorded comparison rather than re-running a model. It shows actual +operator provenance, requested TP/CP and dtype axes, the fixed comparison +threshold, and the worst active selected-token delta. + +Generate the offline desktop bundle from one completed attempt: + +```bash +python -m rl_engine.alignment.cross_config report \ + runs//cases//attempt-0001 \ + --output /tmp/qwen3-tp2-cp2.rlk-drift +``` + +The `.rlk-drift` file contains sanitized report JSON, a Chrome Trace Event JSON +trace, and a PNG preview. It excludes checkpoints, prompts, and raw score +tensors. Install the optional local viewer and open it without a browser: + +```bash +python -m pip install "rl-engine[drift-viewer]" +rlk-drift-view /tmp/qwen3-tp2-cp2.rlk-drift +``` + +The viewer has an expandable track tree, horizontal zoom/scroll, selectable +events, and an event-details panel. The horizontal scale is explicitly marked +as a sample ordinal unless the source artifacts contain real timestamps. + +Other output suffixes expose the same sealed evidence in a form suitable for a +specific workflow: + +| Suffix | Artifact | Intended use | +| --- | --- | --- | +| `.rlk-drift` | Offline desktop bundle | Interactive post-training triage | +| `.png`, `.jpg` | Static summary | Pull request or issue attachment | +| `.json` | Chrome Trace Event | Perfetto or trace-viewer inspection | + +An unsealed or malformed attempt is rejected. A failed identity gate stays +`not comparable`; the tool does not turn it into a numerical drift value. diff --git a/envs.py b/envs.py index 833f00ac..77dcb8e3 100644 --- a/envs.py +++ b/envs.py @@ -26,5 +26,7 @@ def env_flag(name: str, default: bool = False) -> bool: KERNEL_ALIGN_NCU_LINEINFO = "KERNEL_ALIGN_NCU_LINEINFO" KERNEL_ALIGN_ALLOW_UNSUPPORTED_MSVC = "KERNEL_ALIGN_ALLOW_UNSUPPORTED_MSVC" KERNEL_ALIGN_FORCE_SM90 = "KERNEL_ALIGN_FORCE_SM90" +RL_KERNEL_REQUIRE_EXT = "RL_KERNEL_REQUIRE_EXT" + KERNEL_ALIGN_FORCE_ASCEND = "KERNEL_ALIGN_FORCE_ASCEND" KERNEL_ALIGN_ASCEND_ARCH = "KERNEL_ALIGN_ASCEND_ARCH" diff --git a/examples/cross_config_qwen3_8b_megatron_tp2_cp2_vllm.json b/examples/cross_config_qwen3_8b_megatron_tp2_cp2_vllm.json new file mode 100644 index 00000000..ed668820 --- /dev/null +++ b/examples/cross_config_qwen3_8b_megatron_tp2_cp2_vllm.json @@ -0,0 +1,98 @@ +{ + "experiment_id": "ws2-qwen3-8b-attention-tp2-cp2", + "scenario_id": "qwen3_8b_megatron_tp2_cp2_vllm", + "contract_source": "ws1", + "contract_version": "current", + "strategy": "one_at_a_time", + "strict_fallback": true, + "scenario": { + "issue": "https://github.com/RL-Align/RL-Kernel/issues/235", + "pull_request": "PR4 -- cross-config integration", + "model": "Qwen3-8B dense", + "training_framework": "megatron", + "rollout_framework": "vllm", + "topology": "2 nodes x 2 GPUs, TP=2 CP=2 PP=1 DP=1, BF16, SM90", + "notes": [ + "Supersedes cross_config_s1_distributed_smoke.json and", + "cross_config_s3_qwen3_8b_tp4_cp4_bf16.json, whose training side used", + "HuggingFace attention backends and FSDP sharding. Neither exists in", + "Megatron, and DP=1 makes the sharding knob meaningless.", + "cross_config_s2_vllm_tp_vs_fsdp.json has no Megatron-only counterpart at", + "all and should be retired rather than rewritten.", + "rollout.context_parallel_size binds to vLLM", + "ParallelConfig.prefill_context_parallel_size and therefore applies to", + "prefill only; strict PR4 acceptance covers CP=2 prefill/chunked prefill.", + "A decode request that becomes CP=1 is a blocking fallback and is tested", + "separately by the PR6 logical KV replay harness." + ], + "debug_matrix": { + "schema_version": "rlkernel.debug_matrix.v1", + "method": "fixed_replay_one_at_a_time", + "module_manifest": "rl_engine.alignment.cross_config.debug_matrix", + "modules": ["attention", "ffn", "logp"], + "comparison_edges": ["train_vs_rollout_prefill", "rollout_prefill_vs_decode"], + "replay_identity": "fixed rollout tokens, masks, positions, cache metadata, and pre-update state", + "cartesian_product": false + } + }, + "baseline": { + "batch": { + "size": 2 + }, + "rollout": { + "tensor_parallel_size": 2, + "context_parallel_size": 2, + "dtype": "bfloat16", + "enable_prefix_caching": false, + "enforce_eager": true, + "batch_invariant": true, + "kv_block_size": 16 + }, + "training": { + "tensor_parallel_size": 2, + "context_parallel_size": 2, + "attention_backend": "unfused", + "compute_dtype": "bfloat16", + "deterministic_mode": true, + "cp_comm_type": "p2p", + "sharding": "unsharded" + }, + "attention": { + "reduction_acc_dtype": "fp32", + "reduction_order": "global_block_index", + "reduction_downcast_at": "final_write", + "reduction_engine": "in_op_reference", + "fusion_boundary": "unfused_rope_attention", + "split_kv_policy": 32 + }, + "logp": { + "backend": "native" + } + }, + "interventions": [ + { + "path": "training.context_parallel_size", + "values": [1, 2] + }, + { + "path": "training.tensor_parallel_size", + "values": [1, 2] + }, + { + "path": "attention.fusion_boundary", + "values": ["unfused_rope_attention", "fused_rope_attention"] + }, + { + "path": "training.cp_comm_type", + "values": ["p2p", "all_gather"] + }, + { + "path": "attention.reduction_order", + "values": ["global_block_index", "arrival"] + }, + { + "path": "attention.reduction_acc_dtype", + "values": ["fp32", "bf16"] + } + ] +} diff --git a/examples/cross_config_s0_cpu_smoke.json b/examples/cross_config_s0_cpu_smoke.json new file mode 100644 index 00000000..0720c67d --- /dev/null +++ b/examples/cross_config_s0_cpu_smoke.json @@ -0,0 +1,75 @@ +{ + "schema_version": "cross_config.experiment_config.v1", + "experiment_id": "cross-config-s0-cpu-smoke-v1", + "scenario_id": "cross_config.s0.cpu_smoke.v1", + "contract_source": "ws1", + "contract_version": "current", + "strategy": "one_at_a_time", + "strict_fallback": true, + "identity": { + "checkpoint_id": "cross_config.synthetic.cpu_logits.v1", + "model_version": "immutable:cross-config-synthetic-cpu-v1", + "tokenizer_id": "cross_config.synthetic_tokenizer.v1", + "tokenizer_policy": "pretokenized_fixture.v1;add_special_tokens=false;padding_side=right", + "token_ids": [ + [11, 12, 13, 21, 22, 23], + [31, 32, 33, 41, 42, 43] + ], + "selected_token_ids": [ + [11, 12, 13, 21, 22, 23], + [31, 32, 33, 41, 42, 43] + ], + "active_mask": [ + [false, false, false, true, true, true], + [false, false, false, true, true, true] + ], + "attention_mask": [ + [true, true, true, true, true, true], + [true, true, true, true, true, true] + ], + "position_ids": [ + [0, 1, 2, 3, 4, 5], + [0, 1, 2, 3, 4, 5] + ], + "pre_update_state": "synthetic_read_only:no_parameters:no_optimizer", + "cache_metadata": { + "use_cache": false + }, + "packing_metadata": { + "packed": false + } + }, + "baseline": { + "batch": { + "size": 2 + }, + "rollout": { + "tensor_parallel_size": 1, + "context_parallel_size": 1, + "dtype": "float32", + "enable_prefix_caching": false, + "enforce_eager": true + }, + "training": { + "attention_backend": "eager", + "compute_dtype": "float32", + "sharding": "unsharded" + }, + "logp": { + "backend": "smoke_only.logp_reference" + } + }, + "interventions": [], + "operators": { + "selected_logprob": { + "rollout": "smoke_only.logp_reference", + "training": "smoke_only.logp_reference" + } + }, + "scenario": { + "level": "S0", + "name": "CPU framework smoke", + "device": "cpu", + "hardware_required": false + } +} diff --git a/examples/cross_config_s1_distributed_smoke.json b/examples/cross_config_s1_distributed_smoke.json new file mode 100644 index 00000000..e7112c0d --- /dev/null +++ b/examples/cross_config_s1_distributed_smoke.json @@ -0,0 +1,98 @@ +{ + "schema_version": "cross_config.experiment_config.v1", + "experiment_id": "cross-config-s1-distributed-smoke-v1", + "scenario_id": "cross_config.s1.distributed_smoke.v1", + "contract_source": "ws1", + "contract_version": "current", + "strategy": "one_at_a_time", + "strict_fallback": true, + "identity": { + "checkpoint_id": "Qwen/Qwen3-0.6B", + "model_version": "c1899de289a04d12100db370d81485cdf75e47ca", + "tokenizer_id": "Qwen/Qwen3-0.6B@c1899de289a04d12100db370d81485cdf75e47ca", + "tokenizer_policy": "pretokenized_fixture.v1;add_special_tokens=false;padding_side=left", + "token_ids": [ + [151643, 8948, 198, 11, 12, 13, 14, 15], + [151643, 8948, 198, 21, 22, 23, 24, 25] + ], + "selected_token_ids": [ + [151643, 8948, 198, 11, 12, 13, 14, 15], + [151643, 8948, 198, 21, 22, 23, 24, 25] + ], + "active_mask": [ + [false, false, false, false, true, true, true, true], + [false, false, false, false, true, true, true, true] + ], + "attention_mask": [ + [true, true, true, true, true, true, true, true], + [true, true, true, true, true, true, true, true] + ], + "position_ids": [ + [0, 1, 2, 3, 4, 5, 6, 7], + [0, 1, 2, 3, 4, 5, 6, 7] + ], + "pre_update_state": "checkpoint_revision:c1899de289a04d12100db370d81485cdf75e47ca;optimizer_steps=0", + "cache_metadata": { + "position_policy": "absolute" + }, + "packing_metadata": { + "packed": false + } + }, + "baseline": { + "batch": { + "size": 2 + }, + "rollout": { + "tensor_parallel_size": 1, + "context_parallel_size": 1, + "dtype": "bfloat16", + "enable_prefix_caching": true, + "enforce_eager": false + }, + "training": { + "attention_backend": "sdpa", + "compute_dtype": "bfloat16", + "sharding": "unsharded" + }, + "logp": { + "backend": "native" + } + }, + "interventions": [ + { + "path": "batch.size", + "values": [ + 1 + ] + }, + { + "path": "rollout.tensor_parallel_size", + "values": [ + 2 + ] + }, + { + "path": "rollout.context_parallel_size", + "values": [ + 2 + ] + }, + { + "path": "training.sharding", + "values": [ + "fsdp" + ] + } + ], + "scenario": { + "level": "S1", + "name": "Smallest distributed lifecycle smoke", + "device": "cuda", + "hardware_required": true, + "model_id": "Qwen/Qwen3-0.6B", + "model_revision": "c1899de289a04d12100db370d81485cdf75e47ca", + "rollout_engine": "vllm", + "training_engine": "fsdp_score_only" + } +} diff --git a/examples/cross_config_s2_vllm_tp_vs_fsdp.json b/examples/cross_config_s2_vllm_tp_vs_fsdp.json new file mode 100644 index 00000000..10a44cf3 --- /dev/null +++ b/examples/cross_config_s2_vllm_tp_vs_fsdp.json @@ -0,0 +1,128 @@ +{ + "schema_version": "cross_config.experiment_config.v1", + "experiment_id": "cross-config-s2-vllm-tp-vs-fsdp-v1", + "scenario_id": "cross_config.s2.vllm_tp_vs_fsdp.v1", + "contract_source": "ws1", + "contract_version": "current", + "strategy": "one_at_a_time", + "strict_fallback": true, + "identity": { + "checkpoint_id": "Qwen/Qwen3-8B", + "model_version": "b968826d9c46dd6066d109eabc6255188de91218", + "tokenizer_id": "Qwen/Qwen3-8B@b968826d9c46dd6066d109eabc6255188de91218", + "tokenizer_policy": "pretokenized_fixture.v1;add_special_tokens=false;padding_side=left", + "token_ids": [ + [151643, 8948, 198, 11, 12, 13, 14, 15], + [151643, 8948, 198, 21, 22, 23, 24, 25] + ], + "selected_token_ids": [ + [151643, 8948, 198, 11, 12, 13, 14, 15], + [151643, 8948, 198, 21, 22, 23, 24, 25] + ], + "active_mask": [ + [false, false, false, false, true, true, true, true], + [false, false, false, false, true, true, true, true] + ], + "attention_mask": [ + [true, true, true, true, true, true, true, true], + [true, true, true, true, true, true, true, true] + ], + "position_ids": [ + [0, 1, 2, 3, 4, 5, 6, 7], + [0, 1, 2, 3, 4, 5, 6, 7] + ], + "pre_update_state": "checkpoint_revision:b968826d9c46dd6066d109eabc6255188de91218;optimizer_steps=0", + "cache_metadata": { + "position_policy": "absolute" + }, + "packing_metadata": { + "packed": false + } + }, + "baseline": { + "batch": { + "size": 2 + }, + "rollout": { + "tensor_parallel_size": 2, + "context_parallel_size": 1, + "dtype": "bfloat16", + "enable_prefix_caching": true, + "enforce_eager": false + }, + "training": { + "attention_backend": "flash_attention_2", + "compute_dtype": "bfloat16", + "sharding": "fsdp" + }, + "logp": { + "backend": "native" + } + }, + "interventions": [ + { + "path": "batch.size", + "values": [ + 1 + ] + }, + { + "path": "rollout.tensor_parallel_size", + "values": [ + 1 + ] + }, + { + "path": "rollout.dtype", + "values": [ + "float32" + ] + }, + { + "path": "rollout.enable_prefix_caching", + "values": [ + false + ] + }, + { + "path": "rollout.enforce_eager", + "values": [ + true + ] + }, + { + "path": "training.attention_backend", + "values": [ + "eager" + ] + }, + { + "path": "training.compute_dtype", + "values": [ + "float32" + ] + }, + { + "path": "logp.backend", + "values": [ + "rlkernel.reference_logp" + ] + }, + { + "path": "training.sharding", + "values": [ + "unsharded" + ] + } + ], + "scenario": { + "level": "S2", + "name": "Issue 111 vLLM TP=2 versus training FSDP", + "device": "cuda", + "hardware_required": true, + "model_id": "Qwen/Qwen3-8B", + "model_revision": "b968826d9c46dd6066d109eabc6255188de91218", + "rollout_engine": "vllm", + "training_engine": "fsdp_score_only" + } +} diff --git a/examples/cross_config_s3_qwen3_8b_tp4_cp4_bf16.json b/examples/cross_config_s3_qwen3_8b_tp4_cp4_bf16.json new file mode 100644 index 00000000..2dcf79fb --- /dev/null +++ b/examples/cross_config_s3_qwen3_8b_tp4_cp4_bf16.json @@ -0,0 +1,134 @@ +{ + "schema_version": "cross_config.experiment_config.v1", + "experiment_id": "cross-config-s3-qwen3-8b-tp4-cp4-bf16-v1", + "scenario_id": "cross_config.s3.qwen3_8b_tp4_cp4_bf16.v1", + "contract_source": "ws1", + "contract_version": "current", + "strategy": "one_at_a_time", + "strict_fallback": true, + "identity": { + "checkpoint_id": "Qwen/Qwen3-8B", + "model_version": "b968826d9c46dd6066d109eabc6255188de91218", + "tokenizer_id": "Qwen/Qwen3-8B@b968826d9c46dd6066d109eabc6255188de91218", + "tokenizer_policy": "pretokenized_fixture.v1;add_special_tokens=false;padding_side=left", + "token_ids": [ + [151643, 8948, 198, 11, 12, 13, 14, 15], + [151643, 8948, 198, 21, 22, 23, 24, 25] + ], + "selected_token_ids": [ + [151643, 8948, 198, 11, 12, 13, 14, 15], + [151643, 8948, 198, 21, 22, 23, 24, 25] + ], + "active_mask": [ + [false, false, false, false, true, true, true, true], + [false, false, false, false, true, true, true, true] + ], + "attention_mask": [ + [true, true, true, true, true, true, true, true], + [true, true, true, true, true, true, true, true] + ], + "position_ids": [ + [0, 1, 2, 3, 4, 5, 6, 7], + [0, 1, 2, 3, 4, 5, 6, 7] + ], + "pre_update_state": "checkpoint_revision:b968826d9c46dd6066d109eabc6255188de91218;optimizer_steps=0", + "cache_metadata": { + "position_policy": "absolute" + }, + "packing_metadata": { + "packed": false + } + }, + "baseline": { + "batch": { + "size": 2 + }, + "rollout": { + "tensor_parallel_size": 4, + "context_parallel_size": 4, + "dtype": "bfloat16", + "enable_prefix_caching": true, + "enforce_eager": false + }, + "training": { + "attention_backend": "flash_attention_2", + "compute_dtype": "bfloat16", + "sharding": "fsdp" + }, + "logp": { + "backend": "native" + } + }, + "interventions": [ + { + "path": "batch.size", + "values": [ + 1 + ] + }, + { + "path": "rollout.tensor_parallel_size", + "values": [ + 1 + ] + }, + { + "path": "rollout.context_parallel_size", + "values": [ + 1 + ] + }, + { + "path": "rollout.dtype", + "values": [ + "float32" + ] + }, + { + "path": "rollout.enable_prefix_caching", + "values": [ + false + ] + }, + { + "path": "rollout.enforce_eager", + "values": [ + true + ] + }, + { + "path": "training.attention_backend", + "values": [ + "eager" + ] + }, + { + "path": "training.compute_dtype", + "values": [ + "float32" + ] + }, + { + "path": "logp.backend", + "values": [ + "rlkernel.reference_logp" + ] + }, + { + "path": "training.sharding", + "values": [ + "unsharded" + ] + } + ], + "scenario": { + "level": "S3", + "name": "Roadmap Qwen3-8B TP=4 CP=4 BF16 milestone", + "device": "cuda", + "hardware_required": true, + "model_id": "Qwen/Qwen3-8B", + "model_revision": "b968826d9c46dd6066d109eabc6255188de91218", + "rollout_engine": "vllm", + "training_engine": "fsdp_score_only" + } +} diff --git a/examples/vime_qwen3_8b_tp2_cp2/README.md b/examples/vime_qwen3_8b_tp2_cp2/README.md new file mode 100644 index 00000000..2039f118 --- /dev/null +++ b/examples/vime_qwen3_8b_tp2_cp2/README.md @@ -0,0 +1,83 @@ +# Vime Qwen3-8B TP=2 CP=2 validation + +This example is the recommended reproducible entry point for the Vime-side +selected-logprob integration. It keeps framework glue in Vime and keeps the +numerical provider, contract, provenance, and report in RL-Kernel. + +The example is deliberately strict: + +- Megatron training uses `TP=2`, `CP=2`, `PP=1`, and four actor ranks. +- vLLM rollout uses processed logprobs and `top_p=1.0`. +- Vime must load `rl_engine.integrations.vime.logp.provider` in `strict` mode. +- A native fallback or a missing provider marker is not reported as a pass. +- Attention and FFN are not declared consistent from configuration alone. They + require executed Megatron and vLLM readbacks, so the report marks them + `unclaimed` until those artifacts are supplied. The readback must use + `rlkernel.operator_runtime_evidence.v1` and report exact-zero comparison + metrics for both sides. + +The Vime companion must be installed or checked out separately. This example +does not modify `vllm-project/vime`. + +The executable entry point is the Vime script +`scripts/run-qwen3-8B-rlkernel-tp2-cp2.sh`. Its default topology is an +8-GPU H100 node with four Megatron actor GPUs and four vLLM rollout GPUs. The +script refuses to start on a different GPU count or GPU class. Set +`COLOCATE=1` only when intentionally testing the colocated path; that mode is +not the default 8-GPU train/infer split. + +## Dry run + +```bash +python examples/vime_qwen3_8b_tp2_cp2/run.py \ + --vime-root /path/to/RL-Align/vime \ + --rl-kernel-root /path/to/RL-Kernel \ + --output reports/qwen3_8b_tp2_cp2.validation.json +``` + +## Execute + +The Vime script expects model/checkpoint/data paths through environment +variables. Override them before adding `--run`: + +```bash +export MODEL_ROOT=/models/Qwen3-8B +export TORCH_DIST_ROOT=/models/Qwen3-8B_torch_dist +export PROMPT_DATA=/data/dapo-math-17k.jsonl +export RL_KERNEL_ROOT=/path/to/RL-Kernel +export MEGATRON_ROOT=/path/to/Megatron-LM + +python examples/vime_qwen3_8b_tp2_cp2/run.py \ + --vime-root /path/to/RL-Align/vime \ + --rl-kernel-root "$RL_KERNEL_ROOT" \ + --output reports/qwen3_8b_tp2_cp2.validation.json \ + --run +``` + +For a real 8xH100 run, the model, Megatron torch-dist checkpoint, prompt data, +and Megatron checkout must already exist on the host. The first run can omit +`VIME_CKPT`; the script will initialize from `TORCH_DIST_ROOT` and save the +Vime checkpoint there. Use `NUM_ROLLOUT=1` for the integration smoke test and +increase it only after the provider marker is observed. + +When the Megatron/vLLM launch also emits the operator readback artifact, pass +it explicitly: + +```bash +python examples/vime_qwen3_8b_tp2_cp2/run.py \ + --vime-root /path/to/RL-Align/vime \ + --rl-kernel-root "$RL_KERNEL_ROOT" \ + --runtime-evidence reports/qwen3_8b_tp2_cp2.runtime-evidence.json \ + --output reports/qwen3_8b_tp2_cp2.validation.json \ + --run +``` + +The evidence file is intentionally post-execution. It must include training +and rollout identities for `attention` and `ffn`, plus `passed: true` and +exact-zero `out`, backward, and (for attention) `LSE` comparison metrics. A +configured backend without this readback remains `unclaimed`. + +The runner writes a JSON report and a sibling combined log. The report records +the exact command, both repository revisions, provider backend identity, strict +fallback status, and the claim boundary. It does not fabricate numerical drift +when the GPU run was not executed. diff --git a/examples/vime_qwen3_8b_tp2_cp2/qwen3_8b_tp2_cp2.json b/examples/vime_qwen3_8b_tp2_cp2/qwen3_8b_tp2_cp2.json new file mode 100644 index 00000000..93332cbb --- /dev/null +++ b/examples/vime_qwen3_8b_tp2_cp2/qwen3_8b_tp2_cp2.json @@ -0,0 +1,43 @@ +{ + "schema_version": "rlkernel.vime_qwen3_8b_tp2_cp2.v1", + "model": "Qwen/Qwen3-8B", + "training": { + "framework": "megatron", + "tensor_model_parallel_size": 2, + "context_parallel_size": 2, + "pipeline_model_parallel_size": 1, + "world_size": 4, + "dtype": "bf16" + }, + "rollout": { + "framework": "vllm", + "top_p": 1.0, + "logprobs_mode": "processed_logprobs" + }, + "selected_logprob_provider": { + "path": "rl_engine.integrations.vime.logp.provider", + "mode": "strict", + "backend_id": "pytorch-vocab-parallel-logp-ws2", + "real_vocab_size": 151936, + "padded_vocab_size": 152064, + "num_vocab_tiles": 64 + }, + "operator_evidence": { + "logp": { + "training": "rl-kernel-provider", + "rollout": "vllm-native-processed-logprobs", + "required_runtime_marker": "Selected-logprob provider active" + }, + "attention": { + "training": "runtime-readback-required", + "rollout": "runtime-readback-required", + "status": "not_claimed_without_Megatron_and_vLLM_readback" + }, + "ffn": { + "training": "runtime-readback-required", + "rollout": "runtime-readback-required", + "status": "not_claimed_without_Megatron_and_vLLM_readback" + } + }, + "vime_script": "scripts/run-qwen3-8B-rlkernel-tp2-cp2.sh" +} diff --git a/examples/vime_qwen3_8b_tp2_cp2/run.py b/examples/vime_qwen3_8b_tp2_cp2/run.py new file mode 100644 index 00000000..72e3744b --- /dev/null +++ b/examples/vime_qwen3_8b_tp2_cp2/run.py @@ -0,0 +1,271 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Run and archive the Vime Qwen3-8B TP=2/CP=2 validation entry point. + +This is an integration example, not a synthetic pass generator. A dry run +only records the exact launch contract. ``--run`` executes Vime and records +whether the strict RL-Kernel provider was actually observed in the log. The +report deliberately leaves attention/FFN unclaimed until both framework +readbacks are supplied. +""" + +from __future__ import annotations + +import argparse +import json +import os +import subprocess +from datetime import datetime, timezone +from pathlib import Path +from typing import Any, Mapping + +DEFAULT_CONFIG = Path(__file__).with_name("qwen3_8b_tp2_cp2.json") +PROVIDER_MARKER = "Selected-logprob provider active" +FALLBACK_MARKERS = ("using native path", "fallback=True", "fallback=true") +RUNTIME_EVIDENCE_SCHEMA = "rlkernel.operator_runtime_evidence.v1" +_OPERATOR_METRICS = { + "attention": ("out_max_abs", "lse_max_abs", "dq_max_abs", "dk_max_abs", "dv_max_abs"), + "ffn": ("out_max_abs", "dx_max_abs", "dw_max_abs"), +} + + +def load_config(path: Path) -> dict[str, Any]: + with path.open(encoding="utf-8") as handle: + value = json.load(handle) + if not isinstance(value, dict): + raise ValueError("example config must contain a JSON object") + return value + + +def validate_config(config: Mapping[str, Any]) -> None: + training = config.get("training") + rollout = config.get("rollout") + provider = config.get("selected_logprob_provider") + if ( + not isinstance(training, Mapping) + or not isinstance(rollout, Mapping) + or not isinstance(provider, Mapping) + ): + raise ValueError("training, rollout, and selected_logprob_provider sections are required") + expected = { + "tensor_model_parallel_size": 2, + "context_parallel_size": 2, + "pipeline_model_parallel_size": 1, + "world_size": 4, + } + for name, value in expected.items(): + if training.get(name) != value: + raise ValueError(f"training.{name} must be {value!r}") + if rollout.get("top_p") != 1.0: + raise ValueError("rollout.top_p must remain 1.0 for the strict provider contract") + if provider.get("mode") != "strict": + raise ValueError("selected_logprob_provider.mode must be strict") + if provider.get("path") != "rl_engine.integrations.vime.logp.provider": + raise ValueError("example must use the RL-Kernel Vime provider") + if provider.get("backend_id") != "pytorch-vocab-parallel-logp-ws2": + raise ValueError("example must pin the WS2 vocab-parallel backend") + + +def load_runtime_evidence(path: Path | None) -> dict[str, Any] | None: + """Load post-execution readback without treating configuration as evidence.""" + + if path is None: + return None + with path.open(encoding="utf-8") as handle: + value = json.load(handle) + if not isinstance(value, dict) or value.get("schema_version") != RUNTIME_EVIDENCE_SCHEMA: + raise ValueError(f"runtime evidence must use schema {RUNTIME_EVIDENCE_SCHEMA!r}") + return value + + +def _operator_evidence_status(evidence: Mapping[str, Any] | None, operator: str) -> str: + if evidence is None: + return "unclaimed" + operators = evidence.get("operators") + item = operators.get(operator) if isinstance(operators, Mapping) else None + if not isinstance(item, Mapping): + return "unclaimed" + training = item.get("training") + rollout = item.get("rollout") + comparison = item.get("comparison") + if not isinstance(training, Mapping) or not isinstance(rollout, Mapping): + return "unclaimed" + if not isinstance(comparison, Mapping) or comparison.get("passed") is not True: + return "failed" + required_identity = ("implementation_id", "backend_id", "contract_id") + if any(not training.get(name) or not rollout.get(name) for name in required_identity): + return "failed" + if training["implementation_id"] != rollout["implementation_id"]: + return "failed" + for metric in _OPERATOR_METRICS[operator]: + value = comparison.get(metric) + if not isinstance(value, (int, float)) or isinstance(value, bool) or value != 0.0: + return "failed" + return "passed" + + +def validate_runtime_evidence(evidence: Mapping[str, Any] | None) -> None: + """Reject malformed evidence before it can affect a report.""" + + if evidence is None: + return + for operator in _OPERATOR_METRICS: + status = _operator_evidence_status(evidence, operator) + if status == "failed": + raise ValueError(f"runtime evidence for {operator} is incomplete or non-zero") + + +def _revision(path: Path) -> str | None: + try: + result = subprocess.run( + ["git", "-C", str(path), "rev-parse", "HEAD"], + check=True, + capture_output=True, + text=True, + ) + except (OSError, subprocess.CalledProcessError): + return None + return result.stdout.strip() or None + + +def build_environment(vime_root: Path, rl_kernel_root: Path) -> dict[str, str]: + env = dict(os.environ) + existing = [str(vime_root), str(rl_kernel_root), "/root/Megatron-LM"] + if env.get("PYTHONPATH"): + existing.append(env["PYTHONPATH"]) + env["PYTHONPATH"] = os.pathsep.join(existing) + env["RL_KERNEL_ROOT"] = str(rl_kernel_root) + env["TP_SIZE"] = "2" + env["CP_SIZE"] = "2" + env["ROLLOUT_TOP_P"] = "1.0" + return env + + +def build_command(config: Mapping[str, Any], vime_root: Path) -> list[str]: + script = vime_root / str(config.get("vime_script", "")) + if not script.is_file(): + raise FileNotFoundError(f"Vime entry script does not exist: {script}") + return ["bash", str(script)] + + +def build_report( + config: Mapping[str, Any], + *, + vime_root: Path, + rl_kernel_root: Path, + command: list[str], + status: str, + returncode: int | None, + log_text: str, + log_path: Path | None, + runtime_evidence: Mapping[str, Any] | None = None, + runtime_evidence_path: Path | None = None, +) -> dict[str, Any]: + provider_active = PROVIDER_MARKER in log_text + fallback_observed = any(marker in log_text for marker in FALLBACK_MARKERS) + strict_provider_passed = status == "passed" and provider_active and not fallback_observed + effective_status = ( + "passed" if strict_provider_passed else ("failed" if status == "passed" else status) + ) + attention_status = _operator_evidence_status(runtime_evidence, "attention") + ffn_status = _operator_evidence_status(runtime_evidence, "ffn") + return { + "schema_version": "rlkernel.vime_validation_report.v1", + "created_at": datetime.now(timezone.utc).isoformat(), + "status": effective_status, + "claim_boundary": { + "qwen3_8b_tp2_cp2_vime_training": strict_provider_passed, + "attention_train_infer_consistency": attention_status, + "ffn_train_infer_consistency": ffn_status, + "reason": ( + "attention and FFN require executed Megatron/vLLM runtime readbacks; " + "the evidence contract accepts only exact-zero comparison metrics" + ), + }, + "config": dict(config), + "topology": config["training"], + "provider": { + "configured_path": config["selected_logprob_provider"]["path"], + "configured_mode": config["selected_logprob_provider"]["mode"], + "backend_id": config["selected_logprob_provider"]["backend_id"], + "active_observed": provider_active, + "fallback_observed": fallback_observed, + }, + "command": command, + "returncode": returncode, + "artifacts": { + "log": None if log_path is None else str(log_path), + "runtime_evidence": ( + None if runtime_evidence_path is None else str(runtime_evidence_path) + ), + }, + "runtime_evidence": None if runtime_evidence is None else dict(runtime_evidence), + "revisions": { + "vime": _revision(vime_root), + "rl_kernel": _revision(rl_kernel_root), + }, + } + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--config", type=Path, default=DEFAULT_CONFIG) + parser.add_argument("--vime-root", type=Path, default=Path(os.environ.get("VIME_ROOT", "."))) + parser.add_argument( + "--rl-kernel-root", type=Path, default=Path(os.environ.get("RL_KERNEL_ROOT", ".")) + ) + parser.add_argument("--output", type=Path, default=Path("qwen3_8b_tp2_cp2.validation.json")) + parser.add_argument( + "--runtime-evidence", + type=Path, + default=None, + help="post-execution Megatron/vLLM operator readback JSON (strict exact-zero contract)", + ) + parser.add_argument("--run", action="store_true", help="execute the Vime script") + args = parser.parse_args(argv) + + config = load_config(args.config) + validate_config(config) + runtime_evidence = load_runtime_evidence(args.runtime_evidence) + validate_runtime_evidence(runtime_evidence) + vime_root = args.vime_root.resolve() + rl_kernel_root = args.rl_kernel_root.resolve() + command = build_command(config, vime_root) + + status = "not_run" + returncode: int | None = None + log_text = "" + log_path: Path | None = None + if args.run: + args.output.parent.mkdir(parents=True, exist_ok=True) + log_path = args.output.with_suffix(".log") + env = build_environment(vime_root, rl_kernel_root) + with log_path.open("w", encoding="utf-8") as log_handle: + process = subprocess.run( + command, cwd=vime_root, env=env, stdout=log_handle, stderr=subprocess.STDOUT + ) + returncode = process.returncode + log_text = log_path.read_text(encoding="utf-8", errors="replace") + status = "passed" if returncode == 0 else "failed" + + report = build_report( + config, + vime_root=vime_root, + rl_kernel_root=rl_kernel_root, + command=command, + status=status, + returncode=returncode, + log_text=log_text, + log_path=log_path, + runtime_evidence=runtime_evidence, + runtime_evidence_path=args.runtime_evidence, + ) + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(report, indent=2, sort_keys=True) + "\n", encoding="utf-8") + print(json.dumps(report, indent=2, sort_keys=True)) + return 0 if report["status"] in {"passed", "not_run"} else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/pyproject.toml b/pyproject.toml index b0b80a1a..3aa5ded1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,45 +1,52 @@ -[build-system] -requires = ["setuptools>=64", "wheel"] -build-backend = "setuptools.build_meta" - -[project] -name = "RL-Kernel" -version = "0.1.0" -description = "High-performance RL training engine focused on kernel fusion and memory efficiency." -readme = "README.md" -requires-python = ">=3.10" -license = {text = "Apache-2.0"} -authors = [ - {name = "RL-Kernel Contributors"} -] -dependencies = [ - "torch>=2.4.1", - "tabulate", - "numpy", - "accelerate", - "transformers==5.13.1", -] - -[project.optional-dependencies] -cuda = ["flashinfer-python>=0.1.6", "nvidia-ml-py"] +[build-system] +requires = ["setuptools>=64", "wheel"] +build-backend = "setuptools.build_meta" + +[project] +name = "RL-Kernel" +version = "0.1.0" +description = "High-performance RL training engine focused on kernel fusion and memory efficiency." +readme = "README.md" +requires-python = ">=3.10" +license = {text = "Apache-2.0"} +authors = [ + {name = "RL-Kernel Contributors"} +] +dependencies = [ + "torch>=2.4.1", + "tabulate", + "numpy", + "accelerate", + "transformers==5.13.1", +] + +[project.optional-dependencies] +cuda = ["flashinfer-python>=0.1.6", "nvidia-ml-py"] rocm = ["aiter"] vllm = ["vllm>=0.6.0"] +drift-viewer = ["Pillow>=10", "PySide6>=6.6"] dev = ["pytest", "black", "isort", "ruff", "mypy", "pre-commit"] - -[tool.setuptools.packages.find] -where = ["."] -include = ["rl_engine*"] - -[tool.ruff] -line-length = 100 - -[tool.ruff.lint] -select = ["E", "F", "B"] -ignore = [] - -[tool.ruff.lint.per-file-ignores] -"__init__.py" = ["F401"] - -[tool.mypy] -ignore_missing_imports = true -follow_imports = "silent" + +[tool.setuptools.packages.find] +where = ["."] +include = ["rl_engine*"] + +[tool.ruff] +line-length = 100 + +[tool.ruff.lint] +select = ["E", "F", "B"] +ignore = [] + +[tool.ruff.lint.per-file-ignores] +"__init__.py" = ["F401"] + +[tool.mypy] +ignore_missing_imports = true +follow_imports = "silent" + +[tool.pytest.ini_options] +markers = [ + "smoke_operator: temporary smoke-only operator plumbing tests", + "unit: CPU-safe unit tests", +] diff --git a/rl_engine/__init__.py b/rl_engine/__init__.py index e69de29b..306d6813 100644 --- a/rl_engine/__init__.py +++ b/rl_engine/__init__.py @@ -0,0 +1,4 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +import torch # noqa: F401 # Load torch shared libraries before importing rl_engine._C. diff --git a/rl_engine/alignment/cross_config/__init__.py b/rl_engine/alignment/cross_config/__init__.py new file mode 100644 index 00000000..26a28e8d --- /dev/null +++ b/rl_engine/alignment/cross_config/__init__.py @@ -0,0 +1,40 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Plan and run cross-configuration alignment experiments. + +The package root is intentionally small and lazily loads execution code. Extension +authors import adapter, artifact, operator, or schema details from their owning +submodule. +""" + +from importlib import import_module +from typing import Any + +from rl_engine.alignment.cross_config.comparison import compare_score_artifacts +from rl_engine.alignment.cross_config.config import ExperimentConfig, load_config +from rl_engine.alignment.cross_config.execution_plan import ExecutionPlan, build_execution_plan +from rl_engine.alignment.cross_config.planner import ExperimentPlan, Planner + + +def __getattr__(name: str) -> Any: + if name in {"PairedRunResult", "PairedRunner"}: + return getattr(import_module("rl_engine.alignment.cross_config.runner"), name) + if name in {"RuntimeMaterializer", "RuntimeTools"}: + return getattr(import_module("rl_engine.alignment.cross_config.runtime"), name) + raise AttributeError(name) + + +__all__ = [ + "ExperimentConfig", + "ExperimentPlan", + "ExecutionPlan", + "PairedRunResult", + "PairedRunner", + "Planner", + "RuntimeMaterializer", + "RuntimeTools", + "build_execution_plan", + "compare_score_artifacts", + "load_config", +] diff --git a/rl_engine/alignment/cross_config/__main__.py b/rl_engine/alignment/cross_config/__main__.py new file mode 100644 index 00000000..611dddd2 --- /dev/null +++ b/rl_engine/alignment/cross_config/__main__.py @@ -0,0 +1,195 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Command-line interface for cross-configuration experiments.""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path +from typing import Any, Optional, Sequence + +from rl_engine.alignment.cross_config.artifacts import ArtifactStore +from rl_engine.alignment.cross_config.config import ExperimentConfig, load_config +from rl_engine.alignment.cross_config.execution_plan import build_execution_plan + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + commands = parser.add_subparsers(dest="command", required=True) + + plan = commands.add_parser("plan", help="validate and persist a plan without execution") + _add_common_arguments(plan) + + run = commands.add_parser("run", help="execute a plan with an explicit runtime adapter") + _add_common_arguments(run) + run.add_argument( + "--runtime", + required=True, + choices=("cpu-smoke",), + help="Runtime adapter; only the temporary CPU smoke adapter ships in V1", + ) + run.add_argument( + "--allow-smoke-operators", + action="store_true", + help="Explicitly authorize temporary smoke-only operator backends", + ) + run.add_argument( + "--timeout-seconds", + type=float, + default=30.0, + help="Per paired-scoring attempt deadline", + ) + run.add_argument( + "--no-resume", + action="store_true", + help="Create new attempts even when matching COMPLETE artifacts exist", + ) + + report = commands.add_parser( + "report", + help="render a sealed attempt as an offline drift report", + ) + report.add_argument("attempt_dir", type=Path, help="Completed cross-config attempt directory") + report.add_argument( + "--output", + "-o", + type=Path, + required=True, + help="Output .rlk-drift bundle, .png/.jpg image, or .json trace", + ) + report.add_argument("--title", default=None, help="Optional report title") + report.add_argument( + "--no-preview", action="store_true", help="Omit PNG from a .rlk-drift bundle" + ) + return parser + + +def _add_common_arguments(parser: argparse.ArgumentParser) -> None: + parser.add_argument("config", type=Path, help="Versioned experiment JSON") + parser.add_argument( + "--output-root", + type=Path, + default=Path("runs"), + help="Append-only artifact root (default: runs)", + ) + + +def main(argv: Optional[Sequence[str]] = None) -> int: + args = build_parser().parse_args(argv) + try: + if args.command == "report": + summary = _report(args) + else: + config = load_config(args.config) + if args.command == "plan": + summary = record_plan(config, args.output_root) + else: + summary = _run(config, args) + except Exception as exc: + summary = { + "schema_version": "cross_config.cli_summary.v1", + "status": "error", + "error_type": f"{type(exc).__module__}.{type(exc).__qualname__}", + "error": str(exc), + } + print(json.dumps(summary, sort_keys=True)) + print(f"cross-configuration error: {type(exc).__name__}: {exc}", file=sys.stderr) + return 2 + + print(json.dumps(summary, sort_keys=True)) + if args.command == "plan": + print( + f"planned {summary['planned_case_count']} cases; no runtime was created", + file=sys.stderr, + ) + return 0 + if args.command == "report": + print( + f"drift report: {summary['status']} ({summary['output']})", + file=sys.stderr, + ) + return 0 + print( + f"CPU smoke: {summary['status']} ({len(summary['cases'])} cases)", + file=sys.stderr, + ) + for case in summary["cases"]: + print( + f" {case['case_id']}: {case['status']}; actual backends " + f"rollout={case['rollout_backend']}, training={case['training_backend']}; " + f"worst sample/token={case['worst_token_index']}; " + f"mismatches={case['mismatch_count']}; resumed={case['resumed']}", + file=sys.stderr, + ) + return 0 if summary["status"] == "pass" else 1 + + +def record_plan(config: ExperimentConfig, output_root: Path) -> dict[str, Any]: + plan = build_execution_plan(config) + store = ArtifactStore(output_root) + experiment_dir = store.initialize_experiment( + config.definition.experiment_id, + experiment=plan.experiment, + plan=plan.rows(), + ) + return { + "schema_version": "cross_config.cli_summary.v1", + "status": "planned", + "experiment_id": config.definition.experiment_id, + "scenario_id": config.definition.scenario_id, + "planned_case_count": len(plan.entries), + "planning_issues": [issue.to_dict() for issue in plan.issues], + "artifact_dir": str(experiment_dir), + } + + +def _run(config: ExperimentConfig, args: argparse.Namespace) -> dict[str, Any]: + if args.runtime != "cpu-smoke": # pragma: no cover - argparse owns choices + raise ValueError(f"unsupported runtime {args.runtime!r}") + from rl_engine.alignment.testing.cpu_cross_config import run_cpu_experiment + + return run_cpu_experiment( + config, + output_root=args.output_root, + allow_smoke_operators=args.allow_smoke_operators, + timeout_seconds=args.timeout_seconds, + resume=not args.no_resume, + ) + + +def _report(args: argparse.Namespace) -> dict[str, Any]: + from rl_engine.alignment.cross_config.drift_report import ( + build_cross_config_attempt_report, + write_drift_bundle, + write_drift_report_image, + write_drift_trace, + ) + + report = build_cross_config_attempt_report(args.attempt_dir, title=args.title) + suffix = args.output.suffix.lower() + if suffix == ".rlk-drift": + output = write_drift_bundle(report, args.output, include_preview=not args.no_preview) + artifact_type = "bundle" + elif suffix in {".png", ".jpg", ".jpeg"}: + output = write_drift_report_image(report, args.output) + artifact_type = "image" + elif suffix == ".json": + output = write_drift_trace(report, args.output) + artifact_type = "trace" + else: + raise ValueError("report output must use .rlk-drift, .png/.jpg, or .json") + return { + "schema_version": "cross_config.drift_report_summary.v1", + "status": report["status"], + "artifact_type": artifact_type, + "output": str(output), + "case_id": report["axes"].get("case_id"), + "attempt_id": report["axes"].get("attempt_id"), + } + + +if __name__ == "__main__": # pragma: no cover + raise SystemExit(main()) diff --git a/rl_engine/alignment/cross_config/_execution.py b/rl_engine/alignment/cross_config/_execution.py new file mode 100644 index 00000000..a860b951 --- /dev/null +++ b/rl_engine/alignment/cross_config/_execution.py @@ -0,0 +1,617 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Private scoring contracts and child-process supervision.""" + +from __future__ import annotations + +import hashlib +import inspect +import json +import multiprocessing as mp +import os +import tempfile +import time +import traceback +from contextlib import contextmanager +from dataclasses import dataclass, field +from enum import Enum +from pathlib import Path +from typing import Any, Iterator, Mapping, Optional, Protocol, Sequence + +import torch + +from rl_engine.alignment.cross_config._json import strict_json_loads +from rl_engine.alignment.cross_config.schema import CanonicalScoringBatch, ScorerSpec, ScoreSide + + +class PairedRunnerError(RuntimeError): + """Base error for a paired scoring attempt.""" + + +class OperatorExecutionError(PairedRunnerError): + """Raised when exact operator evidence cannot authorize execution.""" + + +class ChildScoringError(PairedRunnerError): + """Raised when a scoring child exits without a valid result.""" + + +class ScoringTimeoutError(PairedRunnerError): + """Raised after all scoring children are stopped at the deadline.""" + + +class RankCompletenessError(PairedRunnerError): + """Raised when rank results are missing, duplicated, or inconsistent.""" + + +class ScorerIdentityError(PairedRunnerError): + """Raised when paired scorer model state is not logically identical.""" + + +@dataclass(frozen=True) +class RankScore: + """One rank's full canonical selected-logprob observation.""" + + rank: int + world_size: int + selected_logprobs: torch.Tensor + metadata: Mapping[str, Any] = field(default_factory=dict) + + def __post_init__(self) -> None: + if self.rank < 0: + raise ValueError("rank must be >= 0") + if self.world_size < 1: + raise ValueError("world_size must be >= 1") + if self.rank >= self.world_size: + raise ValueError("rank must be less than world_size") + if not isinstance(self.selected_logprobs, torch.Tensor): + raise TypeError("selected_logprobs must be a torch.Tensor") + object.__setattr__( + self, + "selected_logprobs", + self.selected_logprobs.detach().to(device="cpu").clone(), + ) + object.__setattr__(self, "metadata", dict(self.metadata)) + + +class PairedScorer(Protocol): + """Small injection boundary used by the paired-runner control plane.""" + + spec: ScorerSpec + + def score( + self, + batch: CanonicalScoringBatch, + *, + batch_size: int, + operator: Any, + ) -> torch.Tensor | RankScore | Sequence[RankScore]: ... + + +class ChildSupervisor: + """Own the lifecycle of the two isolated scoring children.""" + + def __init__(self, start_method: Optional[str] = None): + available = mp.get_all_start_methods() + resolved = start_method or ("fork" if "fork" in available else "spawn") + if resolved not in available: + raise ValueError(f"multiprocessing start method is unavailable: {resolved}") + self.start_method = resolved + self._active_processes: list[mp.Process] = [] + + @property + def active_child_pids(self) -> tuple[int, ...]: + return tuple( + process.pid + for process in self._active_processes + if process.pid is not None and process.is_alive() + ) + + def run( + self, + attempt_dir: Path, + batch: CanonicalScoringBatch, + *, + batch_size: int, + scorers: Mapping[str, PairedScorer], + specs: Mapping[str, ScorerSpec], + instances: Mapping[str, Any], + timeout_seconds: float, + ) -> dict[str, Mapping[str, Any]]: + context: Any = mp.get_context(self.start_method) + processes: dict[str, mp.Process] = {} + with tempfile.TemporaryDirectory(prefix=".paired-runner-", dir=attempt_dir) as tmp: + temporary_dir = Path(tmp) + try: + for target in ("rollout", "training"): + process = context.Process( + target=_score_child, + name=f"cross-config-{target}", + args=( + temporary_dir / f"{target}.pt", + temporary_dir / f"{target}.error.json", + scorers[target], + specs[target], + batch, + batch_size, + instances[target], + ), + ) + process.start() + processes[target] = process + self._active_processes = list(processes.values()) + self._wait( + processes, + temporary_dir, + timeout_seconds=timeout_seconds, + ) + return { + target: _load_child_result(temporary_dir / f"{target}.pt") + for target in ("rollout", "training") + } + finally: + _stop_processes(tuple(processes.values())) + self._active_processes = [] + + @staticmethod + def _wait( + processes: Mapping[str, mp.Process], + temporary_dir: Path, + *, + timeout_seconds: float, + ) -> None: + deadline = time.monotonic() + timeout_seconds + unfinished = set(processes) + while unfinished: + for target in tuple(unfinished): + process = processes[target] + process.join(timeout=0.01) + if process.is_alive(): + continue + unfinished.remove(target) + if process.exitcode != 0: + detail = _child_error_detail(temporary_dir / f"{target}.error.json") + raise ChildScoringError( + f"{target} scoring child failed with exit code " + f"{process.exitcode}: {detail}" + ) + if unfinished and time.monotonic() >= deadline: + labels = ", ".join(sorted(unfinished)) + raise ScoringTimeoutError( + f"paired scoring exceeded {timeout_seconds:.3f}s; " + f"stopped children: {labels}" + ) + + +def _score_child( + result_path: Path, + error_path: Path, + scorer: PairedScorer, + spec: ScorerSpec, + batch: CanonicalScoringBatch, + batch_size: int, + operator: Any, +) -> None: + try: + with _read_only_scoring_guard(scorer, verify_state=True) as evidence: + with torch.no_grad(): + output = scorer.score(batch, batch_size=batch_size, operator=operator) + ranks = _coerce_rank_scores(output, spec.world_size) + payload = { + "schema_version": 1, + "guard_evidence": evidence, + "ranks": [ + { + "rank": rank.rank, + "world_size": rank.world_size, + "selected_logprobs": rank.selected_logprobs, + "metadata": json_safe(rank.metadata), + } + for rank in ranks + ], + } + temporary = result_path.with_suffix(".tmp") + torch.save(payload, temporary) + os.replace(temporary, result_path) + except BaseException as exc: + error = { + "type": f"{type(exc).__module__}.{type(exc).__qualname__}", + "message": str(exc), + "traceback": traceback.format_exc(), + } + error_path.write_text(json.dumps(error, sort_keys=True), encoding="utf-8") + raise SystemExit(1) from None + + +@contextmanager +def _read_only_scoring_guard( + scorer: PairedScorer, + *, + verify_state: bool, +) -> Iterator[dict[str, Any]]: + model = scorer_model(scorer) + if verify_state and getattr(scorer, "optimizer", None) is not None: + raise ValueError("scorer must not own an active optimizer") + if model is None: + yield { + "model_state_verified": False, + "model_eval": False, + "no_grad": True, + "optimizer_step": False, + } + return + + modes = tuple((module, module.training) for module in model.modules()) + snapshot = _module_tensor_snapshot(model) if verify_state else None + model.eval() + evidence = { + "model_state_verified": verify_state, + "model_eval": True, + "no_grad": True, + "optimizer_step": False, + "model_modes_restored": False, + "model_state_unchanged": False if verify_state else None, + } + try: + yield evidence + finally: + for module, was_training in modes: + module.training = was_training + evidence["model_modes_restored"] = True + if snapshot is not None: + mutations = _module_state_mutations(model, snapshot) + if mutations: + raise RuntimeError( + "read-only scorer mutated model parameters/buffers: " + ", ".join(mutations) + ) + evidence["model_state_unchanged"] = True + + +def _module_tensor_snapshot(model: torch.nn.Module) -> dict[str, torch.Tensor]: + values = { + f"parameter:{name}": tensor.detach().to(device="cpu").clone() + for name, tensor in model.named_parameters() + } + values.update( + { + f"buffer:{name}": tensor.detach().to(device="cpu").clone() + for name, tensor in model.named_buffers() + } + ) + return values + + +def _module_state_mutations( + model: torch.nn.Module, + before: Mapping[str, torch.Tensor], +) -> list[str]: + after = _module_tensor_snapshot(model) + mutations: list[str] = [] + for name in sorted(set(before) | set(after)): + left = before.get(name) + right = after.get(name) + if left is None or right is None: + mutations.append(name) + continue + if left.dtype != right.dtype or left.shape != right.shape or not torch.equal(left, right): + mutations.append(name) + return mutations + + +def scorer_model(scorer: PairedScorer) -> Optional[torch.nn.Module]: + if isinstance(scorer, torch.nn.Module): + return scorer + candidate = getattr(scorer, "model", None) + return candidate if isinstance(candidate, torch.nn.Module) else None + + +def paired_model_state_fingerprints( + rollout_scorer: PairedScorer, + training_scorer: PairedScorer, +) -> dict[str, Optional[str]]: + fingerprints = { + "rollout": _scorer_model_state_fingerprint(rollout_scorer), + "training": _scorer_model_state_fingerprint(training_scorer), + } + if fingerprints["rollout"] is None or fingerprints["training"] is None: + raise ScorerIdentityError( + "rollout and training model state fingerprints must both be observable" + ) + if fingerprints["rollout"] != fingerprints["training"]: + raise ScorerIdentityError( + "rollout and training model state fingerprints differ before scoring" + ) + return fingerprints + + +def _scorer_model_state_fingerprint(scorer: PairedScorer) -> Optional[str]: + declared = getattr(scorer, "model_state_fingerprint", None) + if declared is not None and (not isinstance(declared, str) or not declared): + raise ScorerIdentityError("scorer model_state_fingerprint must be a non-empty string") + model = scorer_model(scorer) + if model is None: + return declared + observed_model = _module_state_fingerprint(model) + if declared is not None and declared != observed_model: + raise ScorerIdentityError( + "declared scorer model_state_fingerprint does not match observed model state" + ) + return observed_model + + +def scorer_implementation_fingerprint(scorer: PairedScorer) -> str: + declared = getattr(scorer, "implementation_fingerprint", None) + if declared is not None and (not isinstance(declared, str) or not declared): + raise ScorerIdentityError("scorer implementation_fingerprint must be a non-empty string") + scorer_type = f"{type(scorer).__module__}.{type(scorer).__qualname__}" + score_source = _source_text(getattr(type(scorer), "score", None)) + return canonical_fingerprint( + { + "declared_implementation": declared, + "scorer_type": scorer_type, + "score_source_fingerprint": hashlib.sha256(score_source.encode("utf-8")).hexdigest(), + } + ) + + +def _module_state_fingerprint(model: torch.nn.Module) -> str: + digest = hashlib.sha256() + digest.update(f"{type(model).__module__}.{type(model).__qualname__}".encode("utf-8")) + digest.update(_source_text(type(model)).encode("utf-8")) + tensors = tuple( + (f"parameter:{name}", tensor) for name, tensor in model.named_parameters() + ) + tuple((f"buffer:{name}", tensor) for name, tensor in model.named_buffers()) + for name, tensor in tensors: + snapshot = tensor.detach().to(device="cpu") + if snapshot.is_sparse: + snapshot = snapshot.to_dense() + snapshot = snapshot.contiguous() + digest.update(name.encode("utf-8")) + digest.update(str(snapshot.dtype).encode("utf-8")) + digest.update(str(tuple(snapshot.shape)).encode("utf-8")) + digest.update(snapshot.reshape(-1).view(torch.uint8).numpy().tobytes()) + return digest.hexdigest() + + +def _source_text(value: Any) -> str: + try: + return inspect.getsource(value) + except (OSError, TypeError): + return repr(value) + + +def _coerce_rank_scores( + output: torch.Tensor | RankScore | Sequence[RankScore], + expected_world_size: int, +) -> tuple[RankScore, ...]: + if isinstance(output, torch.Tensor): + if expected_world_size != 1: + raise RankCompletenessError( + "a bare tensor result is valid only for a world_size=1 scorer" + ) + return (RankScore(rank=0, world_size=1, selected_logprobs=output),) + if isinstance(output, RankScore): + return (output,) + if not isinstance(output, Sequence) or isinstance(output, (str, bytes)): + raise TypeError("scorer must return a tensor, RankScore, or sequence of RankScore") + values = tuple(output) + if not all(isinstance(value, RankScore) for value in values): + raise TypeError("every scorer sequence item must be a RankScore") + return values + + +def _load_child_result(path: Path) -> Mapping[str, Any]: + if not path.is_file(): + raise ChildScoringError(f"scoring child produced no result artifact: {path.name}") + try: + payload = torch.load(path, map_location="cpu", weights_only=True) + except Exception as exc: + raise ChildScoringError(f"failed to load scoring child result {path.name}: {exc}") from exc + if not isinstance(payload, Mapping) or payload.get("schema_version") != 1: + raise ChildScoringError(f"malformed scoring child result: {path.name}") + return payload + + +def validate_rank_outputs( + payload: Mapping[str, Any], + spec: ScorerSpec, + *, + expected_shape: torch.Size, + target: str, +) -> dict[int, RankScore]: + raw_ranks = payload.get("ranks") + if not isinstance(raw_ranks, Sequence): + raise RankCompletenessError(f"{target} child result has no rank sequence") + ranks: dict[int, RankScore] = {} + duplicates: list[int] = [] + for raw in raw_ranks: + if not isinstance(raw, Mapping): + raise RankCompletenessError(f"{target} rank result must be a mapping") + rank_score = RankScore( + rank=int(raw["rank"]), + world_size=int(raw["world_size"]), + selected_logprobs=raw["selected_logprobs"], + metadata=raw.get("metadata", {}), + ) + if rank_score.rank in ranks: + duplicates.append(rank_score.rank) + ranks[rank_score.rank] = rank_score + if duplicates: + raise RankCompletenessError(f"{target} returned duplicate ranks: {sorted(set(duplicates))}") + expected = set(range(spec.world_size)) + actual = set(ranks) + if actual != expected: + missing = sorted(expected - actual) + unexpected = sorted(actual - expected) + raise RankCompletenessError( + f"{target} rank set is incomplete; missing={missing}, unexpected={unexpected}" + ) + for rank_index, value in ranks.items(): + if value.world_size != spec.world_size: + raise RankCompletenessError( + f"{target} rank {rank_index} reported world_size={value.world_size}, " + f"expected {spec.world_size}" + ) + if value.selected_logprobs.shape != expected_shape: + raise RankCompletenessError( + f"{target} rank {rank_index} selected_logprobs shape " + f"{tuple(value.selected_logprobs.shape)} does not match " + f"canonical shape {tuple(expected_shape)}" + ) + expected_dtype = torch_dtype(spec.dtype) + if ( + not value.selected_logprobs.is_floating_point() + or value.selected_logprobs.dtype != expected_dtype + ): + raise RankCompletenessError( + f"{target} rank {rank_index} selected_logprobs dtype " + f"{value.selected_logprobs.dtype} does not match scorer dtype {expected_dtype}" + ) + rank_zero = ranks[0].selected_logprobs + for rank_index, value in ranks.items(): + if rank_index == 0: + continue + if value.selected_logprobs.dtype != rank_zero.dtype or not torch.equal( + value.selected_logprobs, + rank_zero, + ): + raise RankCompletenessError( + f"{target} rank {rank_index} selected_logprobs diverge from rank 0" + ) + return ranks + + +def scorer_spec(scorer: PairedScorer, expected_side: ScoreSide) -> ScorerSpec: + spec = getattr(scorer, "spec", None) + if not isinstance(spec, ScorerSpec): + raise TypeError("paired scorer must expose a ScorerSpec as .spec") + if spec.side is not expected_side: + raise ValueError(f"scorer side {spec.side.value!r} does not match {expected_side.value!r}") + model = scorer_model(scorer) + if model is not None: + _require_module_on_device(model, device_type(spec.device)) + _require_module_float_dtype(model, torch_dtype(spec.dtype)) + return spec + + +def validate_scorer_identity( + specs: Mapping[str, ScorerSpec], + batch: CanonicalScoringBatch, +) -> None: + identity = batch.identity + expected = { + "checkpoint_id": identity.checkpoint_id, + "model_version": identity.model_version, + "pre_update_state": identity.pre_update_state, + } + for target in ("rollout", "training"): + observed = specs[target].construction_options + mismatches = [key for key, value in expected.items() if observed.get(key) != value] + if mismatches: + raise ScorerIdentityError( + f"{target} scorer construction identity differs from canonical identity: " + + ", ".join(mismatches) + ) + + +def _child_error_detail(path: Path) -> str: + try: + value = _read_json_object(path) + except (OSError, ValueError, json.JSONDecodeError): + return "child did not publish structured error evidence" + return f"{value.get('type', 'error')}: {value.get('message', '')}" + + +def _read_json_object(path: Path) -> dict[str, Any]: + value = strict_json_loads(path.read_text(encoding="utf-8")) + if not isinstance(value, dict): + raise ValueError(f"expected a JSON object: {path}") + return value + + +def _stop_processes(processes: Iterator[mp.Process] | Sequence[mp.Process]) -> None: + values = tuple(processes) + for process in values: + if process.is_alive(): + process.terminate() + for process in values: + if process.pid is not None: + process.join(timeout=1.0) + for process in values: + if process.is_alive() and hasattr(process, "kill"): + process.kill() + process.join(timeout=1.0) + + +def _require_module_on_device(model: torch.nn.Module, expected: str) -> None: + for name, tensor in tuple(model.named_parameters()) + tuple(model.named_buffers()): + if tensor.device.type != expected: + raise ValueError( + f"scorer model tensor {name!r} is on {tensor.device}; expected {expected}" + ) + + +def _require_module_float_dtype(model: torch.nn.Module, expected: torch.dtype) -> None: + mismatches = [ + f"{name}={tensor.dtype}" + for name, tensor in tuple(model.named_parameters()) + tuple(model.named_buffers()) + if tensor.is_floating_point() and tensor.dtype != expected + ] + if mismatches: + raise ValueError( + f"scorer floating model state must use {expected}: " + ", ".join(mismatches) + ) + + +def device_type(value: str) -> str: + try: + return torch.device(value).type + except (TypeError, RuntimeError) as exc: + raise ValueError(f"invalid scorer device: {value!r}") from exc + + +def normalized_dtype(value: str) -> str: + dtype = torch_dtype(value) + return str(dtype).removeprefix("torch.") + + +def torch_dtype(value: str) -> torch.dtype: + normalized = str(value).strip().lower().replace("torch.", "") + dtypes = { + "float32": torch.float32, + "fp32": torch.float32, + "float16": torch.float16, + "fp16": torch.float16, + "bfloat16": torch.bfloat16, + "bf16": torch.bfloat16, + } + try: + return dtypes[normalized] + except KeyError as exc: + raise ValueError(f"unsupported stateless scorer dtype: {value!r}") from exc + + +def json_safe(value: Any) -> Any: + if isinstance(value, Enum): + return value.value + if isinstance(value, Mapping): + return {str(key): json_safe(item) for key, item in value.items()} + if isinstance(value, (set, frozenset, tuple, list)): + items = [json_safe(item) for item in value] + if isinstance(value, (set, frozenset)): + return sorted(items, key=lambda item: json.dumps(item, sort_keys=True)) + return items + if isinstance(value, (str, int, float, bool)) or value is None: + return value + return repr(value) + + +def canonical_fingerprint(value: Any) -> str: + serialized = json.dumps( + json_safe(value), + sort_keys=True, + separators=(",", ":"), + allow_nan=False, + ) + return hashlib.sha256(serialized.encode("utf-8")).hexdigest() diff --git a/rl_engine/alignment/cross_config/_json.py b/rl_engine/alignment/cross_config/_json.py new file mode 100644 index 00000000..fe8eb876 --- /dev/null +++ b/rl_engine/alignment/cross_config/_json.py @@ -0,0 +1,44 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Fail-closed JSON decoding shared by configs and artifacts.""" + +from __future__ import annotations + +import json +import math +from typing import Any + + +def strict_json_loads(value: str) -> Any: + """Decode RFC JSON while rejecting duplicate keys and non-finite numbers.""" + + return json.loads( + value, + object_pairs_hook=_unique_object, + parse_constant=_reject_json_constant, + parse_float=_parse_finite_float, + ) + + +def _unique_object(pairs: list[tuple[str, Any]]) -> dict[str, Any]: + result: dict[str, Any] = {} + for key, value in pairs: + if key in result: + raise ValueError(f"duplicate JSON key {key!r}") + result[key] = value + return result + + +def _reject_json_constant(value: str) -> Any: + raise ValueError(f"non-finite JSON constant {value!r} is forbidden") + + +def _parse_finite_float(value: str) -> float: + result = float(value) + if not math.isfinite(result): + raise ValueError(f"non-finite JSON number {value!r} is forbidden") + return result + + +__all__ = ["strict_json_loads"] diff --git a/rl_engine/alignment/cross_config/_provenance.py b/rl_engine/alignment/cross_config/_provenance.py new file mode 100644 index 00000000..05817b17 --- /dev/null +++ b/rl_engine/alignment/cross_config/_provenance.py @@ -0,0 +1,327 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Private execution identity and provenance construction.""" + +from __future__ import annotations + +import hashlib +import importlib.metadata +import json +import platform +from dataclasses import replace +from pathlib import Path +from typing import Any, Mapping, Optional + +import torch + +from rl_engine.alignment.cross_config._execution import ( + OperatorExecutionError, + PairedRunnerError, + canonical_fingerprint, + device_type, + json_safe, +) +from rl_engine.alignment.cross_config.runtime import RuntimeMaterialization +from rl_engine.alignment.cross_config.schema import ( + MaterializationStatus, + RuntimeProvenance, + ScoreArtifact, + ScorerSpec, + ScoreSide, +) +from rl_engine.kernels.semantic_registry import OperatorInstanceProvenance, OperatorResolution + +PAIRED_RUNNER_IMPLEMENTATION_FINGERPRINT = "cross_config.paired_runner.v2" + + +def effective_runtime_status( + materialization: RuntimeMaterialization, +) -> MaterializationStatus: + """Aggregate runtime status after exact resolution supersedes logp readback.""" + + statuses = [ + application.status + for application in materialization.applications + if not ( + application.path == "logp.backend" + and application.status is MaterializationStatus.UNOBSERVABLE + ) + ] + precedence = ( + MaterializationStatus.ERROR, + MaterializationStatus.UNSUPPORTED, + MaterializationStatus.UNOBSERVABLE, + MaterializationStatus.FALLBACK, + MaterializationStatus.APPLIED, + ) + return next( + (status for status in precedence if status in statuses), + MaterializationStatus.APPLIED, + ) + + +def side_provenance( + base: RuntimeProvenance, + resolution: OperatorResolution, + instance: OperatorInstanceProvenance, + child_payload: Mapping[str, Any], + spec: ScorerSpec, + *, + status: MaterializationStatus, + factory_options: Mapping[str, Any], + model_state_fingerprint: Optional[str], + scorer_implementation_fingerprint: str, +) -> RuntimeProvenance: + payload = base.to_dict() + actual = dict(payload["actual"]) + actual["operators"] = { + "selected_logprob": { + "backend_id": instance.backend_id, + "descriptor_fingerprint": instance.descriptor_fingerprint, + "implementation_fingerprint": instance.implementation_fingerprint, + "instance_fingerprint": instance.instance_fingerprint, + "concrete_implementation": instance.concrete_implementation, + "factory_options": json_safe(factory_options), + "factory_options_fingerprint": factory_options_fingerprint(factory_options), + } + } + actual["model_state_fingerprint"] = model_state_fingerprint + actual["scorer_implementation_fingerprint"] = scorer_implementation_fingerprint + evidence = dict(payload["evidence"]) + evidence.update( + { + "operator_resolution": resolution.to_dict(), + "operator_instance": instance.to_dict(), + "operator_factory_options": json_safe(factory_options), + "scoring_guard": json_safe(child_payload.get("guard_evidence", {})), + "rank_metadata": [ + json_safe(rank.get("metadata", {})) + for rank in child_payload.get("ranks", ()) + if isinstance(rank, Mapping) + ], + "model_state_fingerprint": model_state_fingerprint, + "scorer_implementation_fingerprint": scorer_implementation_fingerprint, + } + ) + implementation_fingerprint = hashlib.sha256( + f"{base.implementation_fingerprint}:{instance.instance_fingerprint}".encode("utf-8") + ).hexdigest() + return RuntimeProvenance( + requested=payload["requested"], + normalized=payload["normalized"], + materialized=payload["materialized"], + actual=actual, + status=status, + construction_fingerprint=base.construction_fingerprint, + distributed_context_fingerprint=base.distributed_context_fingerprint, + process_fingerprint=base.process_fingerprint, + implementation_fingerprint=implementation_fingerprint, + evidence=evidence, + rank=0, + world_size=spec.world_size, + ) + + +def concrete_scorer_spec( + spec: ScorerSpec, + instance: OperatorInstanceProvenance, +) -> ScorerSpec: + overrides = dict(spec.operator_overrides) + overrides["selected_logprob"] = instance.backend_id + return replace(spec, operator_overrides=overrides) + + +def score_metadata(artifact: ScoreArtifact) -> dict[str, Any]: + value = artifact.to_dict() + value.pop("selected_logprobs", None) + value.pop("active_mask", None) + return { + "case_id": artifact.case_id, + "attempt_id": artifact.attempt_id, + "side": artifact.side.value, + "score_artifact": value, + } + + +def execution_fingerprint( + materialization: RuntimeMaterialization, + *, + specs: Mapping[str, ScorerSpec], + instance_provenance: Mapping[str, OperatorInstanceProvenance], + operator_factory_options: Optional[Mapping[str | ScoreSide, Mapping[str, Any]]], + model_state_fingerprints: Mapping[str, Optional[str]], + scorer_implementation_fingerprints: Mapping[str, str], + environment: Mapping[str, Any], +) -> str: + payload = { + "schema_version": "cross_config.execution_identity.v1", + "runner_implementation_fingerprint": PAIRED_RUNNER_IMPLEMENTATION_FINGERPRINT, + "environment": environment, + "materialized_case": materialization.materialized_case.to_dict(), + "runtime_provenance": materialization.provenance.to_dict(), + "runtime_binding": materialization.binding.to_dict(), + "applications": [application.to_dict() for application in materialization.applications], + "targets": { + target: { + "scorer": concrete_scorer_spec( + specs[target], + instance_provenance[target], + ).to_dict(), + "operator_instance": instance_provenance[target].to_dict(), + "operator_factory_options": json_safe( + target_factory_options(operator_factory_options, target) + ), + "model_state_fingerprint": model_state_fingerprints[target], + "scorer_implementation_fingerprint": (scorer_implementation_fingerprints[target]), + } + for target in ("rollout", "training") + }, + } + return canonical_fingerprint(payload) + + +def mapping_target(mapping: Mapping[str | ScoreSide, Any], target: str) -> Any: + if target in mapping: + return mapping[target] + side = ScoreSide(target) + return mapping.get(side) + + +def target_factory_options( + options: Optional[Mapping[str | ScoreSide, Mapping[str, Any]]], + target: str, +) -> Mapping[str, Any]: + if options is None: + return {} + value = mapping_target(options, target) + if value is None: + return {} + if not isinstance(value, Mapping): + raise OperatorExecutionError(f"{target} operator factory options must be a mapping") + return dict(value) + + +def factory_options_fingerprint(options: Mapping[str, Any]) -> str: + return canonical_fingerprint(options) + + +def runtime_adapter_fingerprint(materialization: RuntimeMaterialization) -> str: + observed = materialization.provenance.evidence.get("adapter_implementation_fingerprint") + if isinstance(observed, str) and observed: + return observed + return materialization.provenance.implementation_fingerprint + + +def execution_environment_provenance( + specs: Mapping[str, ScorerSpec], + *, + runtime_adapter_fingerprint: str, + operator_implementation_fingerprints: Mapping[str, str], +) -> dict[str, Any]: + source_root = Path(__file__).resolve().parents[3] + try: + package_version = importlib.metadata.version("rl-kernel") + except importlib.metadata.PackageNotFoundError: + package_version = None + torch_config = torch.__config__.show() + execution_devices = {target: device_type(spec.device) for target, spec in sorted(specs.items())} + return { + "schema_version": "cross_config.environment.v1", + "execution_devices": execution_devices, + "python": { + "implementation": platform.python_implementation(), + "version": platform.python_version(), + }, + "torch": { + "version": str(torch.__version__), + "git_version": getattr(torch.version, "git_version", None), + "cuda_build": getattr(torch.version, "cuda", None), + "hip_build": getattr(torch.version, "hip", None), + "debug_build": bool(getattr(torch.version, "debug", False)), + "config_fingerprint": hashlib.sha256(torch_config.encode("utf-8")).hexdigest(), + }, + "host_runtime": { + "platform": platform.platform(), + "machine": platform.machine(), + "processor": platform.processor(), + "mkldnn_available": bool(torch.backends.mkldnn.is_available()), + "mkl_available": bool(torch.backends.mkl.is_available()), + }, + "rl_kernel": { + "package_version": package_version, + "git_revision": _git_revision(source_root), + "source_tree_fingerprint": _cross_config_source_tree_fingerprint( + source_root, + implementation_fingerprints={ + "runtime_adapter": runtime_adapter_fingerprint, + "operators": dict(operator_implementation_fingerprints), + }, + ), + }, + } + + +def _git_revision(source_root: Path) -> Optional[str]: + git_dir = source_root / ".git" + try: + if git_dir.is_file(): + marker = git_dir.read_text(encoding="utf-8").strip() + if not marker.startswith("gitdir: "): + return None + resolved = Path(marker.removeprefix("gitdir: ")) + git_dir = resolved if resolved.is_absolute() else source_root / resolved + head = (git_dir / "HEAD").read_text(encoding="utf-8").strip() + if not head.startswith("ref: "): + return head or None + reference = head.removeprefix("ref: ") + loose_ref = git_dir / reference + if loose_ref.is_file(): + return loose_ref.read_text(encoding="utf-8").strip() or None + packed_refs = git_dir / "packed-refs" + if packed_refs.is_file(): + suffix = f" {reference}" + for line in packed_refs.read_text(encoding="utf-8").splitlines(): + if line.endswith(suffix): + return line.split(" ", 1)[0] + except OSError: + return None + return None + + +def _cross_config_source_tree_fingerprint( + source_root: Path, + *, + implementation_fingerprints: Mapping[str, Any], +) -> str: + paths = list((source_root / "rl_engine/alignment/cross_config").glob("*.py")) + paths.extend( + source_root / relative + for relative in ( + "rl_engine/executors/stateless_executor.py", + "rl_engine/kernels/gtest/tolerance.py", + "rl_engine/kernels/registry.py", + "rl_engine/kernels/semantic_registry.py", + ) + ) + digest = hashlib.sha256() + for path in sorted(set(paths)): + try: + content = path.read_bytes() + except OSError as exc: + raise PairedRunnerError( + f"cannot fingerprint cross-configuration source file {path}: {exc}" + ) from exc + digest.update(str(path.relative_to(source_root)).encode("utf-8")) + digest.update(b"\0") + digest.update(content) + digest.update(b"\0") + digest.update( + json.dumps( + json_safe(implementation_fingerprints), + sort_keys=True, + separators=(",", ":"), + allow_nan=False, + ).encode("utf-8") + ) + return digest.hexdigest() diff --git a/rl_engine/alignment/cross_config/_resume.py b/rl_engine/alignment/cross_config/_resume.py new file mode 100644 index 00000000..f3d31e70 --- /dev/null +++ b/rl_engine/alignment/cross_config/_resume.py @@ -0,0 +1,397 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Private validation for append-only attempt resume.""" + +from __future__ import annotations + +import hashlib +import json +import math +from pathlib import Path +from typing import Any, Mapping, Optional + +import torch + +from rl_engine.alignment.cross_config._execution import ( + canonical_fingerprint, + json_safe, + torch_dtype, +) +from rl_engine.alignment.cross_config._json import strict_json_loads +from rl_engine.alignment.cross_config._provenance import ( + PAIRED_RUNNER_IMPLEMENTATION_FINGERPRINT, + concrete_scorer_spec, + effective_runtime_status, + factory_options_fingerprint, + target_factory_options, +) +from rl_engine.alignment.cross_config.artifacts import REQUIRED_CASE_ARTIFACTS +from rl_engine.alignment.cross_config.comparison import recompute_mismatch_mask +from rl_engine.alignment.cross_config.runtime import RuntimeMaterialization +from rl_engine.alignment.cross_config.schema import ( + CanonicalScoringBatch, + ExperimentCase, + ScorerSpec, + ScoreSide, +) +from rl_engine.kernels.gtest.tolerance import ( + resolve_logprob_threshold, + tolerance_contract_fingerprint, +) +from rl_engine.kernels.semantic_registry import OperatorInstanceProvenance + +_COMPLETE_KEYS = frozenset( + { + "schema_version", + "case_id", + "attempt_id", + "status", + "comparable", + "passed", + "active_token_count", + "mismatch_count", + "worst_token_index", + "max_abs_diff", + "rollout_backend", + "training_backend", + "execution_fingerprint", + "environment_fingerprint", + "runner_implementation_fingerprint", + "artifact_sha256", + } +) + + +def completed_attempt_matches( + attempt_dir: Path, + case: ExperimentCase, + batch: CanonicalScoringBatch, + *, + materialization: RuntimeMaterialization, + specs: Mapping[str, ScorerSpec], + instance_provenance: Mapping[str, OperatorInstanceProvenance], + operator_factory_options: Optional[Mapping[str | ScoreSide, Mapping[str, Any]]], + model_state_fingerprints: Mapping[str, Optional[str]], + scorer_implementation_fingerprints: Mapping[str, str], + environment: Mapping[str, Any], + execution_fingerprint: str, +) -> bool: + try: + identity = read_json_object(attempt_dir / "identity.json") + requested = read_json_object(attempt_dir / "requested.json") + actual = read_json_object(attempt_dir / "actual.json") + marker = read_json_object(attempt_dir / "COMPLETE") + except (OSError, ValueError, json.JSONDecodeError): + return False + if not ( + set(marker) == _COMPLETE_KEYS + and identity.get("schema_version") == "cross_config.identity_envelope.v1" + and requested.get("schema_version") == "cross_config.requested.v1" + and actual.get("schema_version") == "cross_config.actual.v1" + and marker.get("schema_version") == "cross_config.complete.v1" + and actual.get("runner_implementation_fingerprint") + == PAIRED_RUNNER_IMPLEMENTATION_FINGERPRINT + and actual.get("environment") == environment + and actual.get("environment_fingerprint") == canonical_fingerprint(environment) + and marker.get("runner_implementation_fingerprint") + == PAIRED_RUNNER_IMPLEMENTATION_FINGERPRINT + and marker.get("environment_fingerprint") == canonical_fingerprint(environment) + and isinstance(marker.get("artifact_sha256"), Mapping) + and set(marker["artifact_sha256"]) == set(REQUIRED_CASE_ARTIFACTS) + ): + return False + if not ( + identity.get("case_id") == case.case_id + and identity.get("identity") == batch.identity.to_dict() + and requested.get("case") == case.to_dict() + and marker.get("execution_fingerprint") == execution_fingerprint + and actual.get("execution_fingerprint") == execution_fingerprint + and marker.get("rollout_backend") == instance_provenance["rollout"].backend_id + and marker.get("training_backend") == instance_provenance["training"].backend_id + ): + return False + + runtime = materialization.provenance.to_dict() + effective_status = effective_runtime_status(materialization).value + score_tensors: dict[str, Mapping[str, torch.Tensor]] = {} + for target in ("rollout", "training"): + prior = actual.get(target) + if not isinstance(prior, Mapping): + return False + instance = instance_provenance[target] + options = target_factory_options(operator_factory_options, target) + expected_operator = { + "backend_id": instance.backend_id, + "descriptor_fingerprint": instance.descriptor_fingerprint, + "implementation_fingerprint": instance.implementation_fingerprint, + "instance_fingerprint": instance.instance_fingerprint, + "concrete_implementation": instance.concrete_implementation, + "factory_options": json_safe(options), + "factory_options_fingerprint": factory_options_fingerprint(options), + } + prior_actual = prior.get("actual") + if not isinstance(prior_actual, Mapping): + return False + if any(prior_actual.get(key) != value for key, value in runtime["actual"].items()): + return False + if prior_actual.get("operators", {}).get("selected_logprob") != expected_operator: + return False + if prior_actual.get("model_state_fingerprint") != model_state_fingerprints[target]: + return False + if ( + prior_actual.get("scorer_implementation_fingerprint") + != scorer_implementation_fingerprints[target] + ): + return False + expected_implementation = hashlib.sha256( + ( + f"{materialization.provenance.implementation_fingerprint}:" + f"{instance.instance_fingerprint}" + ).encode("utf-8") + ).hexdigest() + for key, expected in ( + ("requested", runtime["requested"]), + ("normalized", runtime["normalized"]), + ("materialized", runtime["materialized"]), + ("status", effective_status), + ( + "construction_fingerprint", + materialization.provenance.construction_fingerprint, + ), + ( + "distributed_context_fingerprint", + materialization.provenance.distributed_context_fingerprint, + ), + ("process_fingerprint", materialization.provenance.process_fingerprint), + ("implementation_fingerprint", expected_implementation), + ("world_size", specs[target].world_size), + ): + if prior.get(key) != expected: + return False + try: + score_payload = _load_resume_tensor_bundle(attempt_dir / f"score_{target}.pt") + score_artifact = score_payload["metadata"]["score_artifact"] + prior_scorer = score_artifact["scorer"] + except (OSError, KeyError, TypeError, RuntimeError, ValueError): + return False + expected_scorer = concrete_scorer_spec(specs[target], instance).to_dict() + if prior_scorer != expected_scorer: + return False + if ( + score_artifact.get("schema_version") != "cross_config.score_artifact.v1" + or score_artifact.get("case_id") != case.case_id + or score_artifact.get("attempt_id") != attempt_dir.name + or score_artifact.get("side") != target + or score_artifact.get("identity") != batch.identity.to_dict() + or score_artifact.get("provenance") != prior + ): + return False + tensors = score_payload["tensors"] + selected = tensors.get("selected_logprobs") + active_mask = tensors.get("active_mask") + expected_dtype = torch_dtype(specs[target].dtype) + if ( + not isinstance(selected, torch.Tensor) + or not isinstance(active_mask, torch.Tensor) + or selected.shape != batch.input_ids.shape + or active_mask.shape != batch.input_ids.shape + or selected.dtype != expected_dtype + or active_mask.dtype != torch.bool + or not torch.equal(active_mask, batch.active_mask.to(device="cpu")) + ): + return False + metadata = score_payload["metadata"] + if ( + metadata.get("case_id") != case.case_id + or metadata.get("attempt_id") != attempt_dir.name + or metadata.get("side") != target + ): + return False + score_tensors[target] = tensors + return _resume_comparison_matches( + attempt_dir, + case, + batch, + marker, + score_tensors, + specs, + ) + + +def _load_resume_tensor_bundle(path: Path) -> Mapping[str, Any]: + payload = torch.load(path, map_location="cpu", weights_only=True) + if not isinstance(payload, Mapping) or payload.get("schema_version") != 1: + raise ValueError(f"invalid tensor bundle schema: {path}") + tensors = payload.get("tensors") + metadata = payload.get("metadata") + if not isinstance(tensors, Mapping) or not all( + isinstance(tensor, torch.Tensor) for tensor in tensors.values() + ): + raise ValueError(f"invalid tensor bundle payload: {path}") + if not isinstance(metadata, Mapping): + raise ValueError(f"invalid tensor bundle metadata: {path}") + return payload + + +def _resume_comparison_matches( + attempt_dir: Path, + case: ExperimentCase, + batch: CanonicalScoringBatch, + marker: Mapping[str, Any], + scores: Mapping[str, Mapping[str, torch.Tensor]], + specs: Mapping[str, ScorerSpec], +) -> bool: + try: + comparison = read_json_object(attempt_dir / "comparison.json") + token_bundle = _load_resume_tensor_bundle(attempt_dir / "token_diffs.pt") + except (OSError, ValueError, RuntimeError, json.JSONDecodeError): + return False + required_token_keys = { + "rollout_logprobs", + "training_logprobs", + "active_mask", + "absolute_diff", + "mismatch_mask", + } + token_tensors = token_bundle["tensors"] + if not required_token_keys.issubset(token_tensors): + return False + rollout = scores["rollout"]["selected_logprobs"] + training = scores["training"]["selected_logprobs"] + active_mask = batch.active_mask.to(device="cpu", dtype=torch.bool) + if not bool(torch.isfinite(rollout[active_mask]).all().item()) or not bool( + torch.isfinite(training[active_mask]).all().item() + ): + return False + rollout_threshold = resolve_logprob_threshold(specs["rollout"].dtype) + training_threshold = resolve_logprob_threshold(specs["training"].dtype) + if rollout_threshold != training_threshold: + return False + fixed_threshold = rollout_threshold + rollout = rollout.masked_fill(~active_mask, 0.0) + training = training.masked_fill(~active_mask, 0.0) + absolute_diff = torch.abs(training - rollout) + mismatch_mask = recompute_mismatch_mask( + rollout, + training, + active_mask, + fixed_threshold, + ) + expected_tensors = { + "rollout_logprobs": rollout, + "training_logprobs": training, + "active_mask": active_mask, + "absolute_diff": absolute_diff, + "mismatch_mask": mismatch_mask, + } + if any( + token_tensors[name].dtype != expected.dtype + or token_tensors[name].shape != expected.shape + or not torch.equal(token_tensors[name], expected) + for name, expected in expected_tensors.items() + ): + return False + token_metadata = token_bundle["metadata"] + active_count = int(active_mask.sum().item()) + if active_count == 0: + return False + mismatch_count = int(mismatch_mask.sum().item()) + passed = mismatch_count == 0 + status = "pass" if passed else "fail" + if ( + token_metadata.get("case_id") != case.case_id + or token_metadata.get("attempt_id") != attempt_dir.name + or token_metadata.get("status") != status + or token_metadata.get("fixed_threshold") != fixed_threshold + ): + return False + diagnostics = _comparison_diagnostics( + rollout, + training, + active_mask, + absolute_diff, + mismatch_count, + ) + expected_comparison = { + "schema_version": "cross_config.alignment_result.v1", + "case_id": case.case_id, + "attempt_id": attempt_dir.name, + "status": status, + "comparable": True, + "passed": passed, + "active_token_count": active_count, + "mismatch_count": mismatch_count, + "contract_fingerprint": tolerance_contract_fingerprint(), + "fixed_threshold": fixed_threshold, + "identity_errors": [], + "artifact_errors": [], + "diagnostics": diagnostics, + "token_artifact": { + **{name: _serialized_tensor(tensor) for name, tensor in expected_tensors.items()}, + "fixed_threshold": fixed_threshold, + "schema_version": "cross_config.token_comparison.v1", + }, + } + if comparison != expected_comparison: + return False + if ( + marker.get("case_id") != case.case_id + or marker.get("attempt_id") != attempt_dir.name + or marker.get("status") != status + or marker.get("comparable") is not True + or marker.get("passed") is not passed + or marker.get("active_token_count") != active_count + or marker.get("mismatch_count") != mismatch_count + or marker.get("max_abs_diff") != diagnostics["max_abs_diff"] + or marker.get("worst_token_index") != diagnostics["worst_token_index"] + ): + return False + return True + + +def _comparison_diagnostics( + rollout: torch.Tensor, + training: torch.Tensor, + active_mask: torch.Tensor, + absolute_diff: torch.Tensor, + mismatch_count: int, +) -> dict[str, Any]: + active_diff = absolute_diff[active_mask].float() + delta = (training[active_mask] - rollout[active_mask]).float() + worst_index = int(torch.argmax(active_diff).item()) + coordinates = torch.nonzero(active_mask, as_tuple=False) + worst_token = [int(item) for item in coordinates[worst_index].tolist()] + approximate_kl = torch.exp(delta.double()) - delta.double() - 1.0 + approximate_kl_mean = _finite_float(approximate_kl.mean()) + active_count = int(active_diff.numel()) + return { + "mean_abs_diff": _finite_float(active_diff.mean()), + "p95_abs_diff": _finite_float(torch.quantile(active_diff, 0.95)), + "p99_abs_diff": _finite_float(torch.quantile(active_diff, 0.99)), + "max_abs_diff": _finite_float(active_diff.max()), + "mismatch_ratio": mismatch_count / active_count, + "approximate_kl_mean": approximate_kl_mean, + "approximate_kl_finite": approximate_kl_mean is not None, + "worst_token_index": worst_token, + } + + +def _finite_float(value: torch.Tensor) -> Optional[float]: + result = float(value.item()) + return result if math.isfinite(result) else None + + +def _serialized_tensor(tensor: torch.Tensor) -> dict[str, Any]: + return { + "dtype": str(tensor.dtype).removeprefix("torch."), + "shape": list(tensor.shape), + "values": tensor.tolist(), + } + + +def read_json_object(path: Path) -> dict[str, Any]: + value = strict_json_loads(path.read_text(encoding="utf-8")) + if not isinstance(value, dict): + raise ValueError(f"expected a JSON object: {path}") + return value diff --git a/rl_engine/alignment/cross_config/adapters/__init__.py b/rl_engine/alignment/cross_config/adapters/__init__.py new file mode 100644 index 00000000..f02b9b38 --- /dev/null +++ b/rl_engine/alignment/cross_config/adapters/__init__.py @@ -0,0 +1,38 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Runtime adapters for the WS2 Qwen3-8B Megatron + vLLM cross-config target.""" + +from rl_engine.alignment.cross_config.adapters._common import ( + QWEN3_8B, + AttentionRuntimeReadback, + Qwen3ModelSpec, +) +from rl_engine.alignment.cross_config.adapters.knobs import ( + MEGATRON_ATTENTION_BACKENDS, + WS2_ATTENTION_KNOB_DESCRIPTORS, + WS2_ATTENTION_KNOBS, + WS2_ATTENTION_NORMALIZERS, +) +from rl_engine.alignment.cross_config.adapters.megatron import ( + MegatronAttentionMaterializer, + MegatronProvenanceAdapter, +) +from rl_engine.alignment.cross_config.adapters.vllm import ( + VllmProvenanceAdapter, + VllmRolloutMaterializer, +) + +__all__ = [ + "MEGATRON_ATTENTION_BACKENDS", + "AttentionRuntimeReadback", + "MegatronAttentionMaterializer", + "MegatronProvenanceAdapter", + "QWEN3_8B", + "Qwen3ModelSpec", + "VllmProvenanceAdapter", + "VllmRolloutMaterializer", + "WS2_ATTENTION_KNOBS", + "WS2_ATTENTION_KNOB_DESCRIPTORS", + "WS2_ATTENTION_NORMALIZERS", +] diff --git a/rl_engine/alignment/cross_config/adapters/_common.py b/rl_engine/alignment/cross_config/adapters/_common.py new file mode 100644 index 00000000..9a3c5384 --- /dev/null +++ b/rl_engine/alignment/cross_config/adapters/_common.py @@ -0,0 +1,281 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Shared pieces for the Megatron and vLLM WS2 attention adapters.""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Any + +from rl_engine.alignment.cross_config.attention_binding import AttentionRuntimeReadback +from rl_engine.alignment.cross_config.runtime import KnobApplication +from rl_engine.alignment.cross_config.schema import ( + IsolationScope, + KnobDescriptor, + MaterializationStatus, +) +from rl_engine.kernels.attention_contract import ( + AttentionDType, + AttentionMerge, + DowncastPoint, + ReductionEngine, + ReductionOrder, + ReductionSpec, + ShardingSpec, + SplitKVSpec, +) + +__all__ = [ + "QWEN3_8B", + "Qwen3ModelSpec", + "AttentionRuntimeReadback", + "application", + "attention_dtype", + "build_reduction_spec", + "build_sharding_spec", + "causal_offsets_for", + "flatten", + "split_kv_spec", + "unsupported_reduction_reason", +] + + +@dataclass(frozen=True) +class Qwen3ModelSpec: + """Architecture constants for the frozen dense target. + + These are *not* knobs. #235/#239/#241 all fix Qwen3-8B dense, so they belong to + the scenario, and both sides must agree on them or the comparison is void. + """ + + name: str = "qwen3-8b" + hidden_size: int = 4096 + ffn_hidden_size: int = 12288 + num_layers: int = 36 + q_heads: int = 32 + kv_heads: int = 8 + head_dim: int = 128 + real_vocab_size: int = 151936 + rope_theta: float = 1.0e6 + rotary_dim: int = 128 + rope_scaling: str | None = None + qk_layernorm: bool = True + + def identity_fields(self) -> dict[str, Any]: + """The subset of :data:`IDENTITY_FIELDS` this spec is responsible for.""" + + return { + "q_heads": self.q_heads, + "kv_heads": self.kv_heads, + "head_dim": self.head_dim, + "rope_theta": self.rope_theta, + "rope_scaling": self.rope_scaling, + "rotary_dim": self.rotary_dim, + "qk_layernorm": self.qk_layernorm, + } + + +QWEN3_8B = Qwen3ModelSpec() + + +#: The planner normalizes dtype knobs to torch spellings (``bfloat16``), while +#: :class:`AttentionDType` uses short spellings (``bf16``). Passing a normalized knob +#: straight into the enum raises, so every adapter must translate here rather than +#: each inventing its own mapping. +_DTYPE_ALIASES: Mapping[str, AttentionDType] = { + "bf16": AttentionDType.BF16, + "bfloat16": AttentionDType.BF16, + "fp16": AttentionDType.FP16, + "float16": AttentionDType.FP16, + "half": AttentionDType.FP16, + "fp32": AttentionDType.FP32, + "float32": AttentionDType.FP32, + "float": AttentionDType.FP32, +} + + +def attention_dtype(value: Any, *, field: str) -> AttentionDType: + """Translate a normalized knob dtype into an :class:`AttentionDType`.""" + + if isinstance(value, AttentionDType): + return value + key = str(value).strip().lower().replace("torch.", "") + try: + return _DTYPE_ALIASES[key] + except KeyError as exc: + raise ValueError( + f"{field}={value!r} is not a supported attention dtype; " + f"expected one of {sorted(set(_DTYPE_ALIASES))}" + ) from exc + + +def split_kv_spec(flat: Mapping[str, Any]) -> SplitKVSpec: + """Build the first-class logical Split-KV request. + + The integer is a fixed logical KV chunk size in tokens. It is intentionally + not vLLM's ``flash_attn_max_num_splits_for_cuda_graph``: that setting is only + an upper bound and cannot prove which runtime boundaries executed. + """ + + split_size = flat.get("attention.split_kv_policy") + if split_size is None: + return SplitKVSpec.disabled() + return SplitKVSpec.fixed(int(split_size)) + + +def flatten(value: Mapping[str, Any], prefix: str = "") -> dict[str, Any]: + """Flatten nested knob mappings into dotted paths.""" + + flat: dict[str, Any] = {} + for key, child in value.items(): + path = f"{prefix}{key}" + if isinstance(child, Mapping): + flat.update(flatten(child, f"{path}.")) + else: + flat[path] = child + return flat + + +def application( + descriptor: KnobDescriptor, + requested: Any, + materialized: Any, + actual: Any, + status: MaterializationStatus, + reason: str, + **evidence: Any, +) -> KnobApplication: + return KnobApplication( + path=descriptor.path, + requested=requested, + materialized=materialized, + actual=actual, + lifecycle=descriptor.lifecycle, + status=status, + evidence={"reason": reason, **evidence}, + critical=descriptor.critical, + ) + + +def unsupported_reduction_reason(flat: Mapping[str, Any]) -> str | None: + """Return why the requested reduction cannot be materialized, if it cannot. + + #236 declares single-member enums for merge order, downcast point and reduction + engine, so the alternative knob values exist only as control groups. Requesting + one must fail loudly rather than quietly collapse onto the supported value -- + silently substituting ``global_block_index`` for a requested ``arrival`` would + make the control group indistinguishable from the treatment. + """ + + order = flat.get("attention.reduction_order") + if order is not None and order != ReductionOrder.GLOBAL_BLOCK_INDEX.value: + return ( + f"attention.reduction_order={order!r} has no backend; #236 ReductionOrder " + f"declares only {ReductionOrder.GLOBAL_BLOCK_INDEX.value!r}" + ) + downcast = flat.get("attention.reduction_downcast_at") + if downcast is not None and downcast != DowncastPoint.FINAL_WRITE.value: + return ( + f"attention.reduction_downcast_at={downcast!r} has no backend; #236 " + f"DowncastPoint declares only {DowncastPoint.FINAL_WRITE.value!r}" + ) + engine = flat.get("attention.reduction_engine") + if engine is not None and engine != ReductionEngine.IN_OP_REFERENCE.value: + return ( + f"attention.reduction_engine={engine!r} has no backend; the Transformer " + "Engine merge oracle lands in #235 PR2/PR3, not here" + ) + acc_dtype = flat.get("attention.reduction_acc_dtype") + if ( + acc_dtype is not None + and attention_dtype(acc_dtype, field="attention.reduction_acc_dtype") + is not AttentionDType.FP32 + ): + return ( + f"attention.reduction_acc_dtype={acc_dtype!r} violates the WS2 mandate; " + "the CP (out, lse) merge accumulates in fp32" + ) + return None + + +def build_reduction_spec(flat: Mapping[str, Any]) -> ReductionSpec: + """Build the reduction spec, having already rejected unsupported requests.""" + + return ReductionSpec( + merge=AttentionMerge.ONLINE_SOFTMAX_LSE, + acc_dtype=AttentionDType.FP32, + order=ReductionOrder.GLOBAL_BLOCK_INDEX, + downcast_at=DowncastPoint.FINAL_WRITE, + engine=ReductionEngine.IN_OP_REFERENCE, + ) + + +def build_sharding_spec( + *, + model: Qwen3ModelSpec, + tp_rank: int, + tp_world_size: int, + cp_rank: int, + cp_world_size: int, + global_sequence_length: int, +) -> ShardingSpec: + """Build a CP/TP sharding spec for one rank of the frozen layout. + + TP splits heads, CP splits the sequence. The #239 rank layout fixes + ``rank = cp_rank * tp_world_size + tp_rank`` for a 2-node x 2-GPU deployment, + but nothing here depends on that mapping: ownership is derived from the ranks + themselves so the same builder serves CP=1 baselines. + """ + + if model.q_heads % tp_world_size or model.kv_heads % tp_world_size: + raise ValueError( + f"Qwen3 GQA heads ({model.q_heads}/{model.kv_heads}) must divide evenly " + f"across tp_world_size={tp_world_size}" + ) + if global_sequence_length % cp_world_size: + raise ValueError( + f"global_sequence_length={global_sequence_length} must divide evenly " + f"across cp_world_size={cp_world_size}" + ) + + local_q_heads = model.q_heads // tp_world_size + local_kv_heads = model.kv_heads // tp_world_size + local_sequence_length = global_sequence_length // cp_world_size + + return ShardingSpec( + tp_rank=tp_rank, + tp_world_size=tp_world_size, + cp_rank=cp_rank, + cp_world_size=cp_world_size, + global_q_heads=model.q_heads, + global_kv_heads=model.kv_heads, + local_q_head_start=tp_rank * local_q_heads, + local_q_heads=local_q_heads, + local_kv_head_start=tp_rank * local_kv_heads, + local_kv_heads=local_kv_heads, + global_sequence_length=global_sequence_length, + local_sequence_length=local_sequence_length, + # One contiguous CP block per rank. The merge order key is the global block + # index, never the arrival order of the CP exchange. + global_block_indices=(cp_rank,), + global_block_token_starts=(cp_rank * local_sequence_length,), + local_block_offsets=(0, local_sequence_length), + ) + + +def causal_offsets_for(sharding: ShardingSpec, batch_size: int) -> tuple[int, ...]: + """Causal offsets for one CP shard, one entry per batch entry. + + Under CP the local query block does not start at global position zero, so the + causal mask has to be shifted by the number of preceding global tokens. Taking + that from ``global_block_token_starts`` rather than recomputing + ``cp_rank * local_sequence_length`` keeps uneven CP splits correct. + """ + + offset = sharding.global_block_token_starts[0] + return (offset,) * batch_size + + +_PROCESS_SCOPES = (IsolationScope.PROCESS, IsolationScope.DISTRIBUTED_CONTEXT) diff --git a/rl_engine/alignment/cross_config/adapters/knobs.py b/rl_engine/alignment/cross_config/adapters/knobs.py new file mode 100644 index 00000000..33c29820 --- /dev/null +++ b/rl_engine/alignment/cross_config/adapters/knobs.py @@ -0,0 +1,161 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""WS2 attention knobs for the Qwen3-8B TP=2 CP=2 Megatron + vLLM target. + +``V1_KNOBS`` was written against a HuggingFace/FSDP rollout-vs-training pair. Three +of its entries do not survive contact with the frozen Megatron + vLLM target: + +* ``training.sharding`` takes ``unsharded``/``fsdp``, neither of which exists in + Megatron, and is meaningless at DP=1 anyway; +* ``training.attention_backend`` takes HuggingFace names + (``flash_attention_2``/``sdpa``/``eager``/``model_default``) while Megatron's + ``AttnBackend`` is ``flash``/``fused``/``unfused``/``local``/``auto``; +* there is no training-side ``tensor_parallel_size`` or ``context_parallel_size`` + at all, so the target configuration cannot even be expressed. + +This module is deliberately **additive**: it extends ``V1_KNOBS`` rather than +editing it, and overrides only the normalizer for ``training.attention_backend``. +Deleting the two dead knobs changes ``V1_KNOBS`` itself and would break existing +cross-config tests, so it is left to a follow-up on the framework PR. +""" + +from __future__ import annotations + +from collections.abc import Mapping + +from rl_engine.alignment.cross_config.planner import ( + _NORMALIZERS, + V1_KNOBS, + Normalizer, + _normalize_choice, + _positive_int, + _strict_bool, +) +from rl_engine.alignment.cross_config.schema import IsolationScope, KnobDescriptor + +__all__ = [ + "MEGATRON_ATTENTION_BACKENDS", + "WS2_ATTENTION_KNOBS", + "WS2_ATTENTION_KNOB_DESCRIPTORS", + "WS2_ATTENTION_NORMALIZERS", +] + + +#: ``megatron.core.transformer.enums.AttnBackend``. +MEGATRON_ATTENTION_BACKENDS: tuple[str, ...] = ( + "flash", + "fused", + "unfused", + "local", + "auto", +) + + +WS2_ATTENTION_KNOB_DESCRIPTORS: tuple[KnobDescriptor, ...] = ( + # -- training-side parallelism: the target configuration itself ------------ + KnobDescriptor( + "training.tensor_parallel_size", + IsolationScope.PROCESS, + ("training",), + ), + KnobDescriptor( + "training.context_parallel_size", + IsolationScope.PROCESS, + ("training",), + ), + # -- determinism switches, one per framework ------------------------------ + KnobDescriptor( + "training.deterministic_mode", + IsolationScope.PROCESS, + ("training",), + ), + KnobDescriptor( + "rollout.batch_invariant", + IsolationScope.PROCESS, + ("rollout",), + ), + # -- reduction knobs: the "turn the noise sources on and off" axis --------- + # These are what make drift attributable. ``reduction.order=arrival`` in + # particular is a control group, not a supported production value. + KnobDescriptor( + "attention.reduction_acc_dtype", + IsolationScope.ENGINE_CONSTRUCTION, + ("rollout", "training"), + allowed_values=("fp32", "bf16"), + ), + KnobDescriptor( + "attention.reduction_order", + IsolationScope.ENGINE_CONSTRUCTION, + ("rollout", "training"), + allowed_values=("global_block_index", "arrival"), + ), + KnobDescriptor( + "attention.reduction_downcast_at", + IsolationScope.ENGINE_CONSTRUCTION, + ("rollout", "training"), + allowed_values=("final_write", "per_block"), + ), + KnobDescriptor( + "attention.reduction_engine", + IsolationScope.ENGINE_CONSTRUCTION, + ("rollout", "training"), + allowed_values=("in_op_reference", "te_oracle"), + ), + # -- materialization knobs: differences the experiment measures ------------ + KnobDescriptor( + "attention.fusion_boundary", + IsolationScope.ENGINE_CONSTRUCTION, + ("rollout", "training"), + allowed_values=("unfused_rope_attention", "fused_rope_attention"), + ), + KnobDescriptor( + # Shared logical KV chunk size. Runtime adapters must separately report + # the actual per-owner boundaries; a configured value is not evidence. + "attention.split_kv_policy", + IsolationScope.ENGINE_CONSTRUCTION, + ("rollout", "training"), + ), + KnobDescriptor( + # vLLM: CacheConfig.block_size -> AttentionContract.kv_cache.page_size + "rollout.kv_block_size", + IsolationScope.ENGINE_CONSTRUCTION, + ("rollout",), + ), + # The CP communication group cannot be reconfigured once built; it is bound to + # the distributed context, not merely to engine construction. + KnobDescriptor( + "training.cp_comm_type", + IsolationScope.DISTRIBUTED_CONTEXT, + ("training",), + allowed_values=("p2p", "all_gather", "a2a", "a2a+p2p"), + ), +) + + +WS2_ATTENTION_KNOBS: Mapping[str, KnobDescriptor] = { + **V1_KNOBS, + **{descriptor.path: descriptor for descriptor in WS2_ATTENTION_KNOB_DESCRIPTORS}, +} + + +WS2_ATTENTION_NORMALIZERS: Mapping[str, Normalizer] = { + **_NORMALIZERS, + # Replace, not map: the HuggingFace names have no Megatron counterpart. + "training.attention_backend": _normalize_choice(*MEGATRON_ATTENTION_BACKENDS), + "training.tensor_parallel_size": _positive_int, + "training.context_parallel_size": _positive_int, + "training.deterministic_mode": _strict_bool, + "rollout.batch_invariant": _strict_bool, + # AttentionDType values, not torch dtype names -- these feed ReductionSpec directly. + "attention.reduction_acc_dtype": _normalize_choice("fp32", "bf16"), + "attention.reduction_order": _normalize_choice("global_block_index", "arrival"), + "attention.reduction_downcast_at": _normalize_choice("final_write", "per_block"), + "attention.reduction_engine": _normalize_choice("in_op_reference", "te_oracle"), + "attention.fusion_boundary": _normalize_choice( + "unfused_rope_attention", "fused_rope_attention" + ), + "attention.split_kv_policy": _positive_int, + "rollout.kv_block_size": _positive_int, + "training.cp_comm_type": _normalize_choice("p2p", "all_gather", "a2a", "a2a+p2p"), +} diff --git a/rl_engine/alignment/cross_config/adapters/megatron.py b/rl_engine/alignment/cross_config/adapters/megatron.py new file mode 100644 index 00000000..55903cc8 --- /dev/null +++ b/rl_engine/alignment/cross_config/adapters/megatron.py @@ -0,0 +1,481 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Training-side (Megatron) runtime adapter for WS2 attention cross-config. + +Two things live here: + +``MegatronProvenanceAdapter`` + Read-only. Turns a Megatron config object into the construction and + distributed-context fingerprints the cross-config framework already expects, + plus the determinism probe. It never imports ``megatron`` -- every accessor is + duck-typed -- so this module is importable and testable on a laptop. + +``MegatronAttentionMaterializer`` + Implements the ``RuntimeMaterializer`` protocol. Before this PR the only + implementation in the repository was ``CpuSmokeMaterializer`` over a synthetic + CPU model, so nothing had ever materialized a real distributed runtime. + +Scope boundary: materialization builds and validates the training-side +:class:`AttentionContract` and reports what would be constructed. Without an +``AttentionRuntimeReadback`` it reports ``UNOBSERVABLE``, never ``APPLIED``. It +does not launch ``torchrun``, initialize process groups, or execute attention; +the 2-node x 2-GPU launcher must inject readback collected after execution. +""" + +from __future__ import annotations + +import hashlib +import json +from collections.abc import Mapping +from typing import Any, Optional + +from rl_engine.alignment.cross_config.adapters._common import ( + QWEN3_8B, + AttentionRuntimeReadback, + Qwen3ModelSpec, + application, + attention_dtype, + build_reduction_spec, + build_sharding_spec, + causal_offsets_for, + flatten, + split_kv_spec, + unsupported_reduction_reason, +) +from rl_engine.alignment.cross_config.determinism import ( + DeterminismProbe, + megatron_probe_from_config, +) +from rl_engine.alignment.cross_config.runtime import ( + AdapterMaterialization, + KnobApplication, + RuntimeBinding, +) +from rl_engine.alignment.cross_config.schema import KnobDescriptor, MaterializationStatus +from rl_engine.kernels.attention_contract import ( + AttentionContract, + AttentionContractError, + AttentionMode, + AttentionRole, + RoPEFusionBoundary, + RoPESpec, + RoPEState, +) +from rl_engine.kernels.semantic_registry import implementation_fingerprint + +__all__ = [ + "MEGATRON_CONSTRUCTION_KEYS", + "MEGATRON_DISTRIBUTED_KEYS", + "MegatronAttentionMaterializer", + "MegatronProvenanceAdapter", +] + + +#: ``TransformerConfig`` fields that change attention arithmetic. Hashed into the +#: construction fingerprint. Deliberately excludes MoE, Mamba, MLA and sparse +#: attention fields: the frozen target is Qwen3-8B dense, and those are asserted +#: off rather than recorded. +MEGATRON_CONSTRUCTION_KEYS: tuple[str, ...] = ( + "attention_backend", + "attention_softmax_in_fp32", + "apply_query_key_layer_scaling", + "apply_rope_fusion", + "masked_softmax_fusion", + "bias_activation_fusion", + "bias_dropout_fusion", + "gradient_accumulation_fusion", + "cross_entropy_loss_fusion", + "cross_entropy_fusion_impl", + "recompute_granularity", + "recompute_method", + "recompute_num_layers", + "recompute_modules", + "rotary_base", + "rotary_percent", + "rotary_interleaved", + "rotary_scaling_factor", + "qk_layernorm", + "hidden_dropout", + "attention_dropout", + "params_dtype", + "bf16", + "fp16", + "fp8", + "deterministic_mode", +) + + +#: ``ModelParallelConfig`` fields that define the distributed context. +MEGATRON_DISTRIBUTED_KEYS: tuple[str, ...] = ( + "tensor_model_parallel_size", + "pipeline_model_parallel_size", + "virtual_pipeline_model_parallel_size", + "context_parallel_size", + "hierarchical_context_parallel_sizes", + "expert_model_parallel_size", + "expert_tensor_parallel_size", + "sequence_parallel", + "cp_comm_type", + "tp_comm_overlap", + "use_te_rng_tracker", +) + + +#: Fields that must hold these values for the frozen dense target. A mismatch is a +#: hard stop, not a recorded difference -- see the exclusion list in the WS2 scope. +MEGATRON_FROZEN_ASSERTIONS: Mapping[str, Any] = { + "pipeline_model_parallel_size": 1, + "expert_model_parallel_size": 1, + "sequence_parallel": False, + "fp8": None, + "hidden_dropout": 0.0, + "attention_dropout": 0.0, +} + + +def _value(config: Any, name: str) -> Any: + raw = getattr(config, name, None) + return getattr(raw, "value", raw) if raw is not None else None + + +def _fingerprint(payload: Mapping[str, Any]) -> str: + encoded = json.dumps(payload, sort_keys=True, separators=(",", ":"), default=str) + return hashlib.sha256(encoded.encode("utf-8")).hexdigest() + + +class MegatronProvenanceAdapter: + """Extract fingerprints and determinism evidence from a Megatron config. + + ``config`` may be a real ``TransformerConfig``/``ModelParallelConfig``, a merged + namespace, or a test double. Missing attributes read as ``None`` and are + recorded as such rather than raising: an absent field is itself provenance. + """ + + framework = "megatron" + + def __init__(self, config: Any, *, env: Optional[Mapping[str, str]] = None): + self.config = config + self.env = dict(env or {}) + + def construction_view(self) -> dict[str, Any]: + return {name: _value(self.config, name) for name in MEGATRON_CONSTRUCTION_KEYS} + + def distributed_view(self) -> dict[str, Any]: + return {name: _value(self.config, name) for name in MEGATRON_DISTRIBUTED_KEYS} + + @property + def construction_fingerprint(self) -> str: + return _fingerprint(self.construction_view()) + + @property + def distributed_context_fingerprint(self) -> str: + return _fingerprint(self.distributed_view()) + + def determinism_probe(self) -> DeterminismProbe: + return megatron_probe_from_config(self.config, env=self.env) + + def frozen_scope_violations(self) -> tuple[str, ...]: + """Return the frozen-scope assertions this config violates.""" + + violations: list[str] = [] + for name, expected in MEGATRON_FROZEN_ASSERTIONS.items(): + actual = _value(self.config, name) + if actual is None: + # Not declared. Treated as unknown rather than as satisfied, because + # a silently-absent MoE or FP8 setting is exactly the case that would + # otherwise slip past a dense-only claim. + violations.append(f"{name} is not declared (expected {expected!r})") + elif actual != expected: + violations.append(f"{name}={actual!r} (expected {expected!r})") + return tuple(violations) + + def to_dict(self) -> dict[str, Any]: + return { + "framework": self.framework, + "construction": self.construction_view(), + "distributed_context": self.distributed_view(), + "construction_fingerprint": self.construction_fingerprint, + "distributed_context_fingerprint": self.distributed_context_fingerprint, + "frozen_scope_violations": list(self.frozen_scope_violations()), + "determinism": self.determinism_probe().to_dict(), + } + + +class MegatronAttentionMaterializer: + """Materialize the training-side attention runtime for the WS2 target.""" + + runtime_kind = "megatron_attention" + + def __init__( + self, + *, + model: Qwen3ModelSpec = QWEN3_8B, + global_sequence_length: int = 4096, + tp_rank: int = 0, + cp_rank: int = 0, + backend_id: str = "rlkernel.cp_attention_reference", + provenance: Optional[MegatronProvenanceAdapter] = None, + runtime_readback: Optional[AttentionRuntimeReadback] = None, + ): + self.model = model + self.global_sequence_length = global_sequence_length + self.tp_rank = tp_rank + self.cp_rank = cp_rank + self.backend_id = backend_id + self.provenance = provenance + self.runtime_readback = runtime_readback + + @property + def implementation_fingerprint(self) -> str: + return implementation_fingerprint( + type(self), + instance=self, + entrypoints=("materialize", "build_contract"), + ) + + def build_contract(self, flat: Mapping[str, Any]) -> AttentionContract: + """Build the training-side contract. Raises on an unusable request.""" + + tp_world_size = int(flat.get("training.tensor_parallel_size", 1)) + cp_world_size = int(flat.get("training.context_parallel_size", 1)) + sharding = build_sharding_spec( + model=self.model, + tp_rank=self.tp_rank, + tp_world_size=tp_world_size, + cp_rank=self.cp_rank, + cp_world_size=cp_world_size, + global_sequence_length=self.global_sequence_length, + ) + fusion = RoPEFusionBoundary( + flat.get( + "attention.fusion_boundary", + RoPEFusionBoundary.UNFUSED_ROPE_ATTENTION.value, + ) + ) + rope = RoPESpec( + q_state=RoPEState.POST_ROPE, + k_state=RoPEState.POST_ROPE, + k_cache_state=RoPEState.POST_ROPE, + theta=self.model.rope_theta, + rotary_dim=self.model.rotary_dim, + rope_scaling=self.model.rope_scaling, + fusion_boundary=fusion, + ) + batch_size = int(flat.get("batch.size", 1)) + return AttentionContract( + role=AttentionRole.TRAIN, + mode=AttentionMode.PREFILL, + dtype=attention_dtype( + flat.get("training.compute_dtype", "bf16"), field="training.compute_dtype" + ), + batch_size=batch_size, + query_sequence_length=sharding.local_sequence_length, + head_dim=self.model.head_dim, + causal=True, + causal_offsets=causal_offsets_for(sharding, batch_size), + sharding=sharding, + reduction=build_reduction_spec(flat), + split_kv=split_kv_spec(flat), + rope=rope, + export_lse=True, + ) + + def materialize( + self, + normalized: Mapping[str, Any], + descriptors: Mapping[str, KnobDescriptor], + ) -> AdapterMaterialization: + flat = flatten(normalized) + applications: list[KnobApplication] = [] + + blocked = unsupported_reduction_reason(flat) + scope_violations = ( + self.provenance.frozen_scope_violations() if self.provenance is not None else () + ) + + contract: AttentionContract | None = None + contract_error: str | None = None + if blocked is None: + try: + contract = self.build_contract(flat) + except (AttentionContractError, ValueError) as exc: + contract_error = str(exc) + + for path, requested in flat.items(): + descriptor = descriptors.get(path) + if descriptor is None or "training" not in descriptor.targets: + continue + if blocked is not None and path.startswith("attention.reduction"): + applications.append( + application( + descriptor, + requested, + None, + None, + MaterializationStatus.UNSUPPORTED, + blocked, + ) + ) + continue + if contract_error is not None: + applications.append( + application( + descriptor, + requested, + None, + None, + MaterializationStatus.ERROR, + contract_error, + ) + ) + continue + if contract is None: + applications.append( + application( + descriptor, + requested, + None, + None, + MaterializationStatus.ERROR, + f"attention contract is unavailable: {blocked}", + ) + ) + continue + applications.append( + self._runtime_application( + descriptor, + requested, + contract=contract, + scope_violations=scope_violations, + ) + ) + + tp_world_size = int(flat.get("training.tensor_parallel_size", 1)) + cp_world_size = int(flat.get("training.context_parallel_size", 1)) + side_config: dict[str, Any] = { + "framework": "megatron", + "attention_backend": flat.get("training.attention_backend"), + "compute_dtype": flat.get("training.compute_dtype"), + "deterministic_mode": flat.get("training.deterministic_mode"), + "cp_comm_type": flat.get("training.cp_comm_type"), + "contract": contract.to_dict() if contract is not None else None, + "contract_error": contract_error or blocked, + "runtime_readback": ( + None if self.runtime_readback is None else self.runtime_readback.to_dict() + ), + "frozen_scope_violations": list(scope_violations), + } + if self.provenance is not None: + side_config["provenance"] = self.provenance.to_dict() + + return AdapterMaterialization( + applications=tuple(applications), + binding=RuntimeBinding( + batch_size=int(flat.get("batch.size", 1)), + side_configs={"training": side_config, "rollout": {}}, + topology={ + "training": { + "world_size": tp_world_size * cp_world_size, + "tensor_parallel_size": tp_world_size, + "context_parallel_size": cp_world_size, + "pipeline_parallel_size": 1, + "data_parallel_size": 1, + }, + "rollout": {"world_size": 1}, + }, + scorer={ + "mode": "teacher_forcing", + "framework": "megatron", + "export_lse": True, + }, + operator_backends={ + "training": self.backend_id, + "rollout": self.backend_id, + }, + runtime_kind=self.runtime_kind, + ), + ) + + def _runtime_application( + self, + descriptor: KnobDescriptor, + requested: Any, + *, + contract: AttentionContract, + scope_violations: tuple[str, ...], + ) -> KnobApplication: + readback = self.runtime_readback + if readback is None: + return application( + descriptor, + requested, + requested, + None, + MaterializationStatus.UNOBSERVABLE, + ( + "configured in the training contract, but no Megatron runtime " + "readback was supplied" + ), + frozen_scope_violations=list(scope_violations), + ) + if scope_violations or not readback.frozen_scope_verified: + return application( + descriptor, + requested, + requested, + readback.actual_knobs.get(descriptor.path), + MaterializationStatus.UNOBSERVABLE, + "Megatron frozen-scope assertions were not all verified by runtime readback", + runtime_readback_source=readback.source, + frozen_scope_violations=list(scope_violations), + ) + if readback.split_kv_fallback: + return application( + descriptor, + requested, + requested, + readback.actual_knobs.get(descriptor.path), + MaterializationStatus.FALLBACK, + "Megatron runtime reported a Split-KV fallback", + runtime_readback_source=readback.source, + ) + if readback.contract != contract: + return application( + descriptor, + requested, + requested, + readback.actual_knobs.get(descriptor.path), + MaterializationStatus.FALLBACK, + "Megatron runtime contract differs from the requested contract", + runtime_readback_source=readback.source, + ) + if descriptor.path not in readback.actual_knobs: + return application( + descriptor, + requested, + requested, + None, + MaterializationStatus.UNOBSERVABLE, + "Megatron runtime readback does not expose this knob", + runtime_readback_source=readback.source, + ) + actual = readback.actual_knobs[descriptor.path] + status = ( + MaterializationStatus.APPLIED if actual == requested else MaterializationStatus.FALLBACK + ) + reason = ( + "verified from the executed Megatron runtime" + if status is MaterializationStatus.APPLIED + else "Megatron runtime value differs from the requested value" + ) + return application( + descriptor, + requested, + requested, + actual, + status, + reason, + runtime_readback_source=readback.source, + frozen_scope_violations=list(scope_violations), + ) diff --git a/rl_engine/alignment/cross_config/adapters/vllm.py b/rl_engine/alignment/cross_config/adapters/vllm.py new file mode 100644 index 00000000..f89ab232 --- /dev/null +++ b/rl_engine/alignment/cross_config/adapters/vllm.py @@ -0,0 +1,540 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Rollout-side (vLLM) runtime adapter for WS2 attention cross-config. + +Mirrors :mod:`.megatron`, with three differences that come straight from what vLLM +actually is: + +* vLLM's context parallelism is ``prefill_context_parallel_size`` -- it applies to + prefill only, so a decode-mode contract must declare ``cp_world_size == 1`` + regardless of what the prefill knob says. +* ``CacheConfig.block_size`` is the paged-KV page size, and it feeds + ``KVCacheSpec.page_size`` directly rather than being invented here. +* Determinism comes from the ``VLLM_BATCH_INVARIANT`` environment variable rather + than from a config field, because vLLM applies it inside + ``init_batch_invariance()`` at worker startup. + +Like the Megatron adapter, nothing here imports ``vllm``; configs are duck-typed so +the module is importable anywhere. Configured-only values remain ``UNOBSERVABLE``; +``APPLIED`` requires an explicit post-execution ``AttentionRuntimeReadback``. +""" + +from __future__ import annotations + +import hashlib +import json +from collections.abc import Mapping +from typing import Any, Optional + +from rl_engine.alignment.cross_config.adapters._common import ( + QWEN3_8B, + AttentionRuntimeReadback, + Qwen3ModelSpec, + application, + attention_dtype, + build_reduction_spec, + build_sharding_spec, + causal_offsets_for, + flatten, + split_kv_spec, + unsupported_reduction_reason, +) +from rl_engine.alignment.cross_config.determinism import DeterminismProbe, vllm_probe_from_env +from rl_engine.alignment.cross_config.runtime import ( + AdapterMaterialization, + KnobApplication, + RuntimeBinding, +) +from rl_engine.alignment.cross_config.schema import KnobDescriptor, MaterializationStatus +from rl_engine.kernels.attention_contract import ( + AttentionContract, + AttentionContractError, + AttentionMode, + AttentionRole, + RoPEFusionBoundary, + RoPESpec, + RoPEState, +) +from rl_engine.kernels.semantic_registry import implementation_fingerprint + +__all__ = [ + "VLLM_ATTENTION_KEYS", + "VLLM_CACHE_KEYS", + "VLLM_FROZEN_ASSERTIONS", + "VLLM_MODEL_KEYS", + "VLLM_PARALLEL_KEYS", + "VllmProvenanceAdapter", + "VllmRolloutMaterializer", +] + + +VLLM_MODEL_KEYS: tuple[str, ...] = ( + "dtype", + "seed", + "quantization", + "enforce_eager", + "max_logprobs", + "disable_cascade_attn", + "max_model_len", +) + +VLLM_CACHE_KEYS: tuple[str, ...] = ( + "block_size", + "cache_dtype", + "enable_prefix_caching", + "prefix_caching_hash_algo", + "calculate_kv_scales", + "sliding_window", +) + +VLLM_ATTENTION_KEYS: tuple[str, ...] = ( + "backend", + "flash_attn_version", + "use_prefill_decode_attention", + "flash_attn_max_num_splits_for_cuda_graph", + "use_cudnn_prefill", + "disable_flashinfer_prefill", + "use_non_causal", +) + +VLLM_PARALLEL_KEYS: tuple[str, ...] = ( + "tensor_parallel_size", + "pipeline_parallel_size", + "prefill_context_parallel_size", + "data_parallel_size", +) + + +#: Frozen dense-target assertions on the rollout side. ``cache_dtype`` must stay +#: ``auto`` because an FP8 KV cache is a representation-drift problem tracked +#: separately, and ``disable_cascade_attn`` must stay ``True`` because cascade +#: attention changes the block-merge structure the contract pins down. +VLLM_FROZEN_ASSERTIONS: Mapping[str, Any] = { + "quantization": None, + "cache_dtype": "auto", + "calculate_kv_scales": False, + "disable_cascade_attn": True, + "pipeline_parallel_size": 1, + "data_parallel_size": 1, + "sliding_window": None, +} + + +def _value(config: Any, name: str) -> Any: + raw = getattr(config, name, None) + return getattr(raw, "value", raw) if raw is not None else None + + +def _fingerprint(payload: Mapping[str, Any]) -> str: + encoded = json.dumps(payload, sort_keys=True, separators=(",", ":"), default=str) + return hashlib.sha256(encoded.encode("utf-8")).hexdigest() + + +class VllmProvenanceAdapter: + """Extract fingerprints and determinism evidence from vLLM configs.""" + + framework = "vllm" + + def __init__( + self, + *, + model_config: Any = None, + cache_config: Any = None, + attention_config: Any = None, + parallel_config: Any = None, + env: Optional[Mapping[str, str]] = None, + ): + self.model_config = model_config + self.cache_config = cache_config + self.attention_config = attention_config + self.parallel_config = parallel_config + self.env = dict(env or {}) + + def construction_view(self) -> dict[str, Any]: + view: dict[str, Any] = {} + for prefix, config, keys in ( + ("model", self.model_config, VLLM_MODEL_KEYS), + ("cache", self.cache_config, VLLM_CACHE_KEYS), + ("attention", self.attention_config, VLLM_ATTENTION_KEYS), + ): + for name in keys: + view[f"{prefix}.{name}"] = _value(config, name) + return view + + def distributed_view(self) -> dict[str, Any]: + return { + f"parallel.{name}": _value(self.parallel_config, name) for name in VLLM_PARALLEL_KEYS + } + + @property + def construction_fingerprint(self) -> str: + return _fingerprint(self.construction_view()) + + @property + def distributed_context_fingerprint(self) -> str: + return _fingerprint(self.distributed_view()) + + def determinism_probe(self) -> DeterminismProbe: + return vllm_probe_from_env(self.env, model_config=self.model_config) + + def frozen_scope_violations(self) -> tuple[str, ...]: + sources = { + "quantization": self.model_config, + "disable_cascade_attn": self.model_config, + "cache_dtype": self.cache_config, + "calculate_kv_scales": self.cache_config, + "sliding_window": self.cache_config, + "pipeline_parallel_size": self.parallel_config, + "data_parallel_size": self.parallel_config, + } + violations: list[str] = [] + for name, expected in VLLM_FROZEN_ASSERTIONS.items(): + config = sources.get(name) + if config is None: + continue + actual = _value(config, name) + if actual != expected: + violations.append(f"{name}={actual!r} (expected {expected!r})") + return tuple(violations) + + @property + def kv_page_size(self) -> Optional[int]: + """vLLM's paged-KV block size, which is the contract's ``page_size``.""" + + block_size = _value(self.cache_config, "block_size") + return int(block_size) if block_size is not None else None + + @property + def split_kv_policy(self) -> Optional[int]: + """Diagnostic vLLM maximum split count, not the logical chunk-size contract.""" + + splits = _value(self.attention_config, "flash_attn_max_num_splits_for_cuda_graph") + return int(splits) if splits is not None else None + + def to_dict(self) -> dict[str, Any]: + return { + "framework": self.framework, + "construction": self.construction_view(), + "distributed_context": self.distributed_view(), + "construction_fingerprint": self.construction_fingerprint, + "distributed_context_fingerprint": self.distributed_context_fingerprint, + "frozen_scope_violations": list(self.frozen_scope_violations()), + "kv_page_size": self.kv_page_size, + "flash_attn_max_num_splits_for_cuda_graph": self.split_kv_policy, + "determinism": self.determinism_probe().to_dict(), + } + + +class VllmRolloutMaterializer: + """Materialize the rollout-side attention runtime for the WS2 target.""" + + runtime_kind = "vllm_attention" + + def __init__( + self, + *, + model: Qwen3ModelSpec = QWEN3_8B, + global_sequence_length: int = 4096, + tp_rank: int = 0, + cp_rank: int = 0, + mode: AttentionMode = AttentionMode.CHUNKED_PREFILL, + backend_id: str = "vllm.flash_attn", + provenance: Optional[VllmProvenanceAdapter] = None, + runtime_readback: Optional[AttentionRuntimeReadback] = None, + ): + self.model = model + self.global_sequence_length = global_sequence_length + self.tp_rank = tp_rank + self.cp_rank = cp_rank + self.mode = mode + self.backend_id = backend_id + self.provenance = provenance + self.runtime_readback = runtime_readback + + @property + def implementation_fingerprint(self) -> str: + return implementation_fingerprint( + type(self), + instance=self, + entrypoints=("materialize", "build_contract"), + ) + + def effective_cp_world_size(self, flat: Mapping[str, Any]) -> int: + """CP applies to prefill only; decode always runs at CP=1.""" + + requested = int(flat.get("rollout.context_parallel_size", 1)) + if self.mode is AttentionMode.DECODE: + return 1 + return requested + + def build_contract(self, flat: Mapping[str, Any]) -> AttentionContract: + if self.mode is AttentionMode.DECODE: + # Decode replay needs a validated KVCacheSpec (cache positions, page + # ownership, prefix-cache identity). That is #235 PR6's contract surface, + # and inventing a placeholder here would let an unvalidated decode case + # look bound. Fail instead. + raise AttentionContractError( + "decode-mode materialization requires KV-cache identity from #235 PR6; " + "this adapter covers prefill and chunked prefill" + ) + tp_world_size = int(flat.get("rollout.tensor_parallel_size", 1)) + cp_world_size = self.effective_cp_world_size(flat) + sharding = build_sharding_spec( + model=self.model, + tp_rank=self.tp_rank, + tp_world_size=tp_world_size, + cp_rank=self.cp_rank if cp_world_size > 1 else 0, + cp_world_size=cp_world_size, + global_sequence_length=self.global_sequence_length, + ) + fusion = RoPEFusionBoundary( + flat.get( + "attention.fusion_boundary", + RoPEFusionBoundary.FUSED_ROPE_ATTENTION.value, + ) + ) + rope = RoPESpec( + q_state=RoPEState.POST_ROPE, + k_state=RoPEState.POST_ROPE, + # vLLM stores post-RoPE K in the cache; recorded, not asserted equal to + # the training side, because it is a materialization fact. + k_cache_state=RoPEState.POST_ROPE, + theta=self.model.rope_theta, + rotary_dim=self.model.rotary_dim, + rope_scaling=self.model.rope_scaling, + fusion_boundary=fusion, + ) + batch_size = int(flat.get("batch.size", 1)) + return AttentionContract( + role=AttentionRole.INFER, + mode=self.mode, + dtype=attention_dtype(flat.get("rollout.dtype", "bf16"), field="rollout.dtype"), + batch_size=batch_size, + query_sequence_length=sharding.local_sequence_length, + head_dim=self.model.head_dim, + causal=True, + causal_offsets=causal_offsets_for(sharding, batch_size), + sharding=sharding, + reduction=build_reduction_spec(flat), + split_kv=split_kv_spec(flat), + rope=rope, + export_lse=True, + ) + + def materialize( + self, + normalized: Mapping[str, Any], + descriptors: Mapping[str, KnobDescriptor], + ) -> AdapterMaterialization: + flat = flatten(normalized) + applications: list[KnobApplication] = [] + + blocked = unsupported_reduction_reason(flat) + scope_violations = ( + self.provenance.frozen_scope_violations() if self.provenance is not None else () + ) + + contract: AttentionContract | None = None + contract_error: str | None = None + if blocked is None: + try: + contract = self.build_contract(flat) + except (AttentionContractError, ValueError) as exc: + contract_error = str(exc) + + requested_cp = int(flat.get("rollout.context_parallel_size", 1)) + effective_cp = self.effective_cp_world_size(flat) + + for path, requested in flat.items(): + descriptor = descriptors.get(path) + if descriptor is None or "rollout" not in descriptor.targets: + continue + if blocked is not None and path.startswith("attention.reduction"): + applications.append( + application( + descriptor, + requested, + None, + None, + MaterializationStatus.UNSUPPORTED, + blocked, + ) + ) + continue + if contract_error is not None: + applications.append( + application( + descriptor, + requested, + None, + None, + MaterializationStatus.ERROR, + contract_error, + ) + ) + continue + if contract is None: + applications.append( + application( + descriptor, + requested, + None, + None, + MaterializationStatus.ERROR, + f"attention contract is unavailable: {blocked}", + ) + ) + continue + if path == "rollout.context_parallel_size" and effective_cp != requested_cp: + applications.append( + application( + descriptor, + requested, + effective_cp, + effective_cp, + MaterializationStatus.FALLBACK, + ( + "vLLM context parallelism covers prefill only; a decode-mode " + f"contract runs at cp_world_size=1, not {requested_cp}" + ), + vllm_field="ParallelConfig.prefill_context_parallel_size", + ) + ) + continue + applications.append( + self._runtime_application( + descriptor, + requested, + contract=contract, + scope_violations=scope_violations, + ) + ) + + tp_world_size = int(flat.get("rollout.tensor_parallel_size", 1)) + side_config: dict[str, Any] = { + "framework": "vllm", + "dtype": flat.get("rollout.dtype"), + "enforce_eager": flat.get("rollout.enforce_eager"), + "enable_prefix_caching": flat.get("rollout.enable_prefix_caching"), + "batch_invariant": flat.get("rollout.batch_invariant"), + "kv_block_size": flat.get("rollout.kv_block_size"), + "split_kv_policy": flat.get("attention.split_kv_policy"), + "attention_mode": self.mode.value, + "contract": contract.to_dict() if contract is not None else None, + "contract_error": contract_error or blocked, + "runtime_readback": ( + None if self.runtime_readback is None else self.runtime_readback.to_dict() + ), + "frozen_scope_violations": list(scope_violations), + } + if self.provenance is not None: + side_config["provenance"] = self.provenance.to_dict() + + return AdapterMaterialization( + applications=tuple(applications), + binding=RuntimeBinding( + batch_size=int(flat.get("batch.size", 1)), + side_configs={"rollout": side_config, "training": {}}, + topology={ + "rollout": { + "world_size": tp_world_size * effective_cp, + "tensor_parallel_size": tp_world_size, + "context_parallel_size": effective_cp, + "pipeline_parallel_size": 1, + "data_parallel_size": 1, + }, + "training": {"world_size": 1}, + }, + scorer={ + "mode": "rollout_logprob", + "framework": "vllm", + "export_lse": True, + }, + operator_backends={ + "rollout": self.backend_id, + "training": self.backend_id, + }, + runtime_kind=self.runtime_kind, + ), + ) + + def _runtime_application( + self, + descriptor: KnobDescriptor, + requested: Any, + *, + contract: AttentionContract, + scope_violations: tuple[str, ...], + ) -> KnobApplication: + readback = self.runtime_readback + if readback is None: + return application( + descriptor, + requested, + requested, + None, + MaterializationStatus.UNOBSERVABLE, + "configured in the rollout contract, but no vLLM runtime readback was supplied", + frozen_scope_violations=list(scope_violations), + ) + if scope_violations or not readback.frozen_scope_verified: + return application( + descriptor, + requested, + requested, + readback.actual_knobs.get(descriptor.path), + MaterializationStatus.UNOBSERVABLE, + "vLLM frozen-scope assertions were not all verified by runtime readback", + runtime_readback_source=readback.source, + frozen_scope_violations=list(scope_violations), + ) + if readback.split_kv_fallback: + return application( + descriptor, + requested, + requested, + readback.actual_knobs.get(descriptor.path), + MaterializationStatus.FALLBACK, + "vLLM runtime reported a Split-KV fallback", + runtime_readback_source=readback.source, + ) + if readback.contract != contract: + return application( + descriptor, + requested, + requested, + readback.actual_knobs.get(descriptor.path), + MaterializationStatus.FALLBACK, + "vLLM runtime contract differs from the requested contract", + runtime_readback_source=readback.source, + ) + if descriptor.path not in readback.actual_knobs: + return application( + descriptor, + requested, + requested, + None, + MaterializationStatus.UNOBSERVABLE, + "vLLM runtime readback does not expose this knob", + runtime_readback_source=readback.source, + ) + actual = readback.actual_knobs[descriptor.path] + status = ( + MaterializationStatus.APPLIED if actual == requested else MaterializationStatus.FALLBACK + ) + reason = ( + "verified from the executed vLLM runtime" + if status is MaterializationStatus.APPLIED + else "vLLM runtime value differs from the requested value" + ) + return application( + descriptor, + requested, + requested, + actual, + status, + reason, + runtime_readback_source=readback.source, + frozen_scope_violations=list(scope_violations), + ) diff --git a/rl_engine/alignment/cross_config/artifacts.py b/rl_engine/alignment/cross_config/artifacts.py new file mode 100644 index 00000000..296c2809 --- /dev/null +++ b/rl_engine/alignment/cross_config/artifacts.py @@ -0,0 +1,508 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Append-only, crash-safe artifacts for cross-configuration runs.""" + +from __future__ import annotations + +import hashlib +import json +import os +import tempfile +from pathlib import Path +from typing import Any, Iterable, Mapping, Optional + +import torch + +from rl_engine.alignment.cross_config._json import strict_json_loads + +REQUIRED_CASE_ARTIFACTS = frozenset( + { + "requested.json", + "materialized.json", + "actual.json", + "identity.json", + "score_rollout.pt", + "score_training.pt", + "comparison.json", + "token_diffs.pt", + } +) +_JSON_SCHEMAS = { + "requested.json": "cross_config.requested.v1", + "materialized.json": "cross_config.materialized_envelope.v1", + "actual.json": "cross_config.actual.v1", + "identity.json": "cross_config.identity_envelope.v1", + "comparison.json": "cross_config.alignment_result.v1", +} +_JSON_REQUIRED_KEYS = { + "requested.json": frozenset({"case"}), + "materialized.json": frozenset({"materialized_case"}), + "actual.json": frozenset({"rollout", "training"}), + "identity.json": frozenset({"identity"}), + "comparison.json": frozenset({"status", "comparable", "passed"}), +} +_TENSOR_REQUIRED_KEYS = { + "score_rollout.pt": frozenset({"selected_logprobs", "active_mask"}), + "score_training.pt": frozenset({"selected_logprobs", "active_mask"}), + "token_diffs.pt": frozenset( + { + "rollout_logprobs", + "training_logprobs", + "active_mask", + "absolute_diff", + "mismatch_mask", + } + ), +} + + +class ArtifactError(RuntimeError): + """Raised when an artifact is incomplete, malformed, or would be overwritten.""" + + +class ArtifactStore: + """Persist immutable attempt directories and atomically mark completed attempts.""" + + def __init__(self, root: str | Path): + self.root = Path(root) + + def experiment_dir(self, experiment_id: str) -> Path: + return self.root / _safe_component(experiment_id, "experiment_id") + + def initialize_experiment( + self, + experiment_id: str, + *, + experiment: Mapping[str, Any], + plan: Iterable[Mapping[str, Any]], + ) -> Path: + """Create immutable experiment metadata, or verify an identical resume target.""" + + directory = self.experiment_dir(experiment_id) + directory.mkdir(parents=True, exist_ok=True) + self._write_or_verify_json(directory / "experiment.json", experiment) + plan_text = "".join(_canonical_json(item) + "\n" for item in plan) + self._write_or_verify_text(directory / "plan.jsonl", plan_text) + return directory + + def create_attempt( + self, + experiment_id: str, + case_id: str, + *, + attempt_id: Optional[str] = None, + ) -> Path: + """Allocate an append-only attempt directory for a case.""" + + case_dir = ( + self.experiment_dir(experiment_id) / "cases" / _safe_component(case_id, "case_id") + ) + case_dir.mkdir(parents=True, exist_ok=True) + if attempt_id is not None: + attempt_dir = case_dir / _safe_component(attempt_id, "attempt_id") + try: + attempt_dir.mkdir() + except FileExistsError as exc: + raise ArtifactError(f"attempt already exists: {attempt_dir}") from exc + _fsync_directory(case_dir) + return attempt_dir + + # Another controller can win after _next_attempt_id() observes the + # directory. mkdir is the atomic allocator; retry rather than aliasing or + # overwriting the winning attempt. + while True: + resolved_attempt_id = self._next_attempt_id(case_dir) + attempt_dir = case_dir / resolved_attempt_id + try: + attempt_dir.mkdir() + except FileExistsError: + continue + _fsync_directory(case_dir) + return attempt_dir + + def write_json(self, attempt_dir: str | Path, name: str, value: Mapping[str, Any]) -> Path: + path = self._attempt_path(attempt_dir, name, suffix=".json") + self._write_new_text(path, _canonical_json(value) + "\n") + return path + + def write_tensor_bundle( + self, + attempt_dir: str | Path, + name: str, + tensors: Mapping[str, torch.Tensor], + *, + metadata: Optional[Mapping[str, Any]] = None, + ) -> Path: + """Write CPU tensor payloads that can be loaded with ``weights_only=True``.""" + + path = self._attempt_path(attempt_dir, name, suffix=".pt") + payload: dict[str, Any] = { + "schema_version": 1, + "tensors": { + key: tensor.detach().to(device="cpu").contiguous() + for key, tensor in tensors.items() + }, + "metadata": dict(metadata or {}), + } + self._atomic_torch_save(path, payload) + return path + + def load_tensor_bundle(self, path: str | Path) -> dict[str, Any]: + try: + payload = torch.load(Path(path), map_location="cpu", weights_only=True) + except Exception as exc: + raise ArtifactError(f"failed to load tensor artifact {path}: {exc}") from exc + if not isinstance(payload, dict) or payload.get("schema_version") != 1: + raise ArtifactError(f"unsupported tensor artifact schema: {path}") + tensors = payload.get("tensors") + if not isinstance(tensors, dict) or not all( + isinstance(value, torch.Tensor) for value in tensors.values() + ): + raise ArtifactError(f"malformed tensor payload: {path}") + return payload + + def complete_attempt( + self, + attempt_dir: str | Path, + *, + summary: Mapping[str, Any], + required: Iterable[str] = REQUIRED_CASE_ARTIFACTS, + ) -> Path: + """Validate all payloads before publishing an atomic ``COMPLETE`` marker.""" + + directory = Path(attempt_dir) + required_names = frozenset(required) + missing = sorted(name for name in required_names if not (directory / name).is_file()) + if missing: + raise ArtifactError(f"cannot complete {directory}; missing artifacts: {missing}") + self._validate_machine_artifacts(directory) + marker_value = dict(summary) + marker_value["artifact_sha256"] = { + name: _sha256_file(directory / name) for name in sorted(required_names) + } + self._validate_complete_summary(directory, marker_value) + marker = directory / "COMPLETE" + self._write_new_text(marker, _canonical_json(marker_value) + "\n") + return marker + + def completed_attempt( + self, + experiment_id: str, + case_id: str, + *, + required: Iterable[str] = REQUIRED_CASE_ARTIFACTS, + ) -> Optional[Path]: + """Return the newest valid completed attempt, ignoring partial attempts.""" + + case_dir = ( + self.experiment_dir(experiment_id) / "cases" / _safe_component(case_id, "case_id") + ) + if not case_dir.is_dir(): + return None + for attempt_dir in sorted( + case_dir.iterdir(), + key=_attempt_sort_key, + reverse=True, + ): + if not attempt_dir.is_dir() or not (attempt_dir / "COMPLETE").is_file(): + continue + try: + self.validate_completed_attempt( + attempt_dir, + required=required, + expected_case_id=case_id, + ) + except ArtifactError: + continue + return attempt_dir + return None + + def validate_completed_attempt( + self, + attempt_dir: str | Path, + *, + required: Iterable[str] = REQUIRED_CASE_ARTIFACTS, + expected_case_id: Optional[str] = None, + ) -> None: + directory = Path(attempt_dir) + required_names = frozenset(required) + marker = directory / "COMPLETE" + if not marker.is_file(): + raise ArtifactError(f"missing COMPLETE marker: {directory}") + try: + marker_value = strict_json_loads(marker.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError, ValueError) as exc: + raise ArtifactError(f"malformed COMPLETE marker: {marker}") from exc + if not isinstance(marker_value, dict): + raise ArtifactError(f"COMPLETE marker must contain a JSON object: {marker}") + self._validate_complete_summary( + directory, + marker_value, + expected_case_id=expected_case_id, + ) + missing = sorted(name for name in required_names if not (directory / name).is_file()) + if missing: + raise ArtifactError(f"completed attempt is missing artifacts: {missing}") + self._validate_artifact_hashes(directory, marker_value, required_names) + self._validate_machine_artifacts(directory, expected_case_id=expected_case_id) + + @staticmethod + def _validate_artifact_hashes( + directory: Path, + marker: Mapping[str, Any], + required: frozenset[str], + ) -> None: + recorded = marker.get("artifact_sha256") + if not isinstance(recorded, Mapping) or set(recorded) != set(required): + raise ArtifactError(f"COMPLETE marker has invalid artifact hashes: {directory}") + for name in sorted(required): + expected = recorded.get(name) + if not isinstance(expected, str) or expected != _sha256_file(directory / name): + raise ArtifactError(f"artifact hash does not match COMPLETE: {directory / name}") + + def _validate_machine_artifacts( + self, + directory: Path, + *, + expected_case_id: Optional[str] = None, + ) -> None: + tensor_payloads: dict[str, dict[str, Any]] = {} + for name in ("score_rollout.pt", "score_training.pt", "token_diffs.pt"): + path = directory / name + if path.exists(): + payload = self.load_tensor_bundle(path) + tensor_payloads[name] = payload + tensors = payload["tensors"] + missing_tensor_keys = sorted(_TENSOR_REQUIRED_KEYS[name].difference(tensors)) + if missing_tensor_keys: + raise ArtifactError( + f"tensor artifact {name} is missing keys: {missing_tensor_keys}" + ) + metadata = payload.get("metadata", {}) + if not isinstance(metadata, Mapping): + raise ArtifactError(f"tensor artifact metadata must be an object: {path}") + if expected_case_id is not None and metadata.get("case_id") != expected_case_id: + raise ArtifactError( + f"tensor artifact case_id does not match {expected_case_id!r}: {path}" + ) + if metadata.get("attempt_id") != directory.name: + raise ArtifactError( + f"tensor artifact attempt_id does not match {directory.name!r}: {path}" + ) + json_payloads: dict[str, dict[str, Any]] = {} + for name in ( + "requested.json", + "materialized.json", + "actual.json", + "identity.json", + "comparison.json", + ): + path = directory / name + if not path.exists(): + continue + try: + value = strict_json_loads(path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError, ValueError) as exc: + raise ArtifactError(f"malformed JSON artifact: {path}") from exc + if not isinstance(value, dict): + raise ArtifactError(f"JSON artifact must contain an object: {path}") + json_payloads[name] = value + if expected_case_id is not None and value.get("case_id") != expected_case_id: + raise ArtifactError( + f"JSON artifact case_id does not match {expected_case_id!r}: {path}" + ) + if value.get("schema_version") != _JSON_SCHEMAS[name]: + raise ArtifactError(f"JSON artifact has an unsupported schema: {path}") + missing_json_keys = sorted(_JSON_REQUIRED_KEYS[name].difference(value)) + if missing_json_keys: + raise ArtifactError(f"JSON artifact {name} is missing keys: {missing_json_keys}") + if value.get("attempt_id") != directory.name: + raise ArtifactError( + f"JSON artifact attempt_id does not match {directory.name!r}: {path}" + ) + + if expected_case_id is None and json_payloads: + case_ids = {payload.get("case_id") for payload in json_payloads.values()} + if len(case_ids) != 1 or None in case_ids: + raise ArtifactError("JSON artifacts must declare one consistent case_id") + inferred_case_id = next(iter(case_ids)) + for name, payload in tensor_payloads.items(): + if payload["metadata"].get("case_id") != inferred_case_id: + raise ArtifactError( + f"tensor artifact case_id does not match {inferred_case_id!r}: " + f"{directory / name}" + ) + + @staticmethod + def _validate_complete_summary( + directory: Path, + summary: Mapping[str, Any], + *, + expected_case_id: Optional[str] = None, + ) -> None: + if summary.get("schema_version") != "cross_config.complete.v1": + raise ArtifactError(f"COMPLETE marker has an unsupported schema: {directory}") + case_id = summary.get("case_id") + if not isinstance(case_id, str) or not case_id: + raise ArtifactError(f"COMPLETE marker is missing case_id: {directory}") + if expected_case_id is not None and case_id != expected_case_id: + raise ArtifactError( + f"COMPLETE marker case_id does not match {expected_case_id!r}: {directory}" + ) + if summary.get("attempt_id") != directory.name: + raise ArtifactError( + f"COMPLETE marker attempt_id does not match {directory.name!r}: {directory}" + ) + if not isinstance(summary.get("status"), str): + raise ArtifactError(f"COMPLETE marker is missing status: {directory}") + if not isinstance(summary.get("artifact_sha256"), Mapping): + raise ArtifactError(f"COMPLETE marker is missing artifact hashes: {directory}") + + def _write_or_verify_json(self, path: Path, value: Mapping[str, Any]) -> None: + self._write_or_verify_text(path, _canonical_json(value) + "\n") + + def _write_or_verify_text(self, path: Path, text: str) -> None: + if path.exists(): + self._verify_existing_text(path, text) + return + try: + self._atomic_write_text(path, text) + except ArtifactError: + # A concurrent writer may have atomically published the same immutable + # experiment metadata. Accept only byte-identical content. + if not path.exists(): + raise + self._verify_existing_text(path, text) + + @staticmethod + def _verify_existing_text(path: Path, text: str) -> None: + try: + existing = path.read_text(encoding="utf-8") + except OSError as exc: + raise ArtifactError(f"failed to read existing artifact {path}: {exc}") from exc + if existing != text: + raise ArtifactError(f"resume metadata differs from existing artifact: {path}") + + def _write_new_text(self, path: Path, text: str) -> None: + if path.exists(): + raise ArtifactError(f"refusing to overwrite artifact: {path}") + self._atomic_write_text(path, text) + + def _atomic_write_text(self, path: Path, text: str) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + fd, temporary_name = tempfile.mkstemp(prefix=f".{path.name}.", dir=path.parent) + temporary = Path(temporary_name) + try: + with os.fdopen(fd, "w", encoding="utf-8") as handle: + handle.write(text) + handle.flush() + os.fsync(handle.fileno()) + _publish_new_file(temporary, path) + except Exception: + temporary.unlink(missing_ok=True) + raise + + def _atomic_torch_save(self, path: Path, payload: Mapping[str, Any]) -> None: + if path.exists(): + raise ArtifactError(f"refusing to overwrite artifact: {path}") + path.parent.mkdir(parents=True, exist_ok=True) + fd, temporary_name = tempfile.mkstemp(prefix=f".{path.name}.", dir=path.parent) + os.close(fd) + temporary = Path(temporary_name) + try: + torch.save(dict(payload), temporary) + with temporary.open("rb") as handle: + os.fsync(handle.fileno()) + _publish_new_file(temporary, path) + except Exception: + temporary.unlink(missing_ok=True) + raise + + @staticmethod + def _next_attempt_id(case_dir: Path) -> str: + indices: list[int] = [] + for child in case_dir.iterdir(): + if not child.is_dir() or not child.name.startswith("attempt-"): + continue + suffix = child.name.removeprefix("attempt-") + if suffix.isdigit(): + indices.append(int(suffix)) + return f"attempt-{max(indices, default=0) + 1:04d}" + + @staticmethod + def _attempt_path(attempt_dir: str | Path, name: str, *, suffix: str) -> Path: + directory = Path(attempt_dir) + safe_name = _safe_component(name, "artifact name") + if not safe_name.endswith(suffix): + safe_name += suffix + return directory / safe_name + + +def _canonical_json(value: Mapping[str, Any]) -> str: + try: + return json.dumps( + value, + sort_keys=True, + separators=(",", ":"), + ensure_ascii=False, + allow_nan=False, + ) + except (TypeError, ValueError) as exc: + raise ArtifactError(f"artifact is not strict JSON: {exc}") from exc + + +def _sha256_file(path: Path) -> str: + digest = hashlib.sha256() + try: + with path.open("rb") as handle: + for chunk in iter(lambda: handle.read(1024 * 1024), b""): + digest.update(chunk) + except OSError as exc: + raise ArtifactError(f"failed to hash artifact {path}: {exc}") from exc + return digest.hexdigest() + + +def _safe_component(value: str, label: str) -> str: + if not value or value in {".", ".."} or Path(value).name != value: + raise ValueError(f"{label} must be a single non-empty path component") + return value + + +def _attempt_sort_key(path: Path) -> tuple[int, int, str]: + """Order standard attempt IDs numerically and retain nonstandard fallbacks.""" + + prefix = "attempt-" + suffix = path.name.removeprefix(prefix) + if path.name.startswith(prefix) and suffix.isdigit(): + return (1, int(suffix), path.name) + return (0, -1, path.name) + + +def _publish_new_file(temporary: Path, destination: Path) -> None: + """Atomically publish without ever replacing an existing artifact.""" + + try: + os.link(temporary, destination) + except FileExistsError as exc: + raise ArtifactError(f"refusing to overwrite artifact: {destination}") from exc + temporary.unlink() + _fsync_directory(destination.parent) + + +def _fsync_directory(directory: Path) -> None: + """Persist directory entry changes where the host filesystem supports it.""" + + try: + descriptor = os.open(directory, os.O_RDONLY) + except OSError: + return + try: + os.fsync(descriptor) + except OSError: + pass + finally: + os.close(descriptor) + + +__all__ = ["ArtifactError", "ArtifactStore", "REQUIRED_CASE_ARTIFACTS"] diff --git a/rl_engine/alignment/cross_config/attention_binding.py b/rl_engine/alignment/cross_config/attention_binding.py new file mode 100644 index 00000000..f6a40ff5 --- /dev/null +++ b/rl_engine/alignment/cross_config/attention_binding.py @@ -0,0 +1,1430 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Three-tier binding between rollout-side and training-side attention contracts. + +Issue #235 PR4 requires that "rollout and training descriptors bind to the same +semantic attention contract". Under the frozen Megatron + vLLM deployment the two +sides can never produce *identical* :class:`AttentionContract` instances: training +runs full-sequence prefill over a CP-sharded sequence, while rollout runs vLLM +paged-KV chunked prefill and decode. Taking "same contract" literally would make +the target configuration permanently unbindable. + +This module therefore splits binding into three tiers: + +``IDENTICAL`` + Logical identity. Both sides must agree bit for bit, otherwise the pair is not + comparable at all and no drift number from it means anything. + +``SEMANTIC`` + The WS2 numerical claim: merge semantics, accumulation dtype, reduction order + and downcast point are decided by the contract, not by the implementation. + Both sides must carry the same values *and* those values must match the WS2 + mandate, otherwise the comparison fails closed. + +``RECORDED`` + Materialization facts that the two sides are expected to differ on -- attention + mode, RoPE fusion boundary, KV-cache paging, backend id, reduction engine. These + differences are exactly what the experiment measures, so they are recorded into + provenance rather than rejected. + +Deliberately *not* in ``SEMANTIC``: ``engine``. Training may run the in-op +deterministic reference while rollout runs a Transformer Engine merge oracle; forcing +those equal would defeat the purpose of the oracle comparison in #235 PR2/3/5/6. +""" + +from __future__ import annotations + +import hashlib +import json +from collections.abc import Mapping, Sequence +from dataclasses import dataclass, field +from enum import Enum +from types import MappingProxyType +from typing import Any, Optional + +from rl_engine.kernels.attention_contract import ( + AttentionContract, + AttentionContractError, + AttentionDType, + AttentionMerge, + AttentionRole, + DowncastPoint, + ReductionOrder, + SplitKVRuntimePlanSet, + validate_split_kv_plan_set_alignment, +) +from rl_engine.kernels.attention_preprocess import ( + ALLOWED_ATTENTION_PREPROCESS_BACKENDS, + MANDATED_ATTENTION_PREPROCESS_BACKENDS, + PREPROCESS_POLICY_ID, +) +from rl_engine.kernels.attention_projection import ( + CUDA_DETERMINISTIC_PROJECTION_BACKEND_ID, + O_PROJ_COLLECTIVE_CONTRACT, + PROJECTION_POLICY_ID, + QKV_COLLECTIVE_CONTRACT, + ROCM_DETERMINISTIC_PROJECTION_BACKEND_ID, +) + +__all__ = [ + "ATTENTION_LSE_DOMAIN", + "AttentionBindingError", + "AttentionBindingResult", + "AttentionRuntimeReadback", + "BindingErrorCode", + "BindingIssue", + "BindingTier", + "IDENTITY_FIELDS", + "NULLABLE_IDENTITY_FIELDS", + "RECORDED_FIELDS", + "SEMANTIC_CONTRACT_FIELDS", + "SEMANTIC_REDUCTION_FIELDS", + "TOPOLOGY_FIELDS", + "WS2_ATTENTION_REDUCTION_MANDATE", + "bind_attention_contracts", + "bind_attention_runtime_readbacks", + "first_blocking_issue", + "identity_fingerprint", + "summarize_binding", +] + + +class AttentionBindingError(ValueError): + """Raised when a caller supplies structurally unusable binding inputs.""" + + +class BindingTier(str, Enum): + """Which rule a field is governed by.""" + + IDENTICAL = "identical" + SEMANTIC = "semantic" + RECORDED = "recorded" + + +class BindingErrorCode(str, Enum): + """Stable, machine-readable reasons a binding is rejected. + + Callers branch on these; they are part of the artifact schema and must not be + renamed without a schema version bump. + """ + + IDENTITY_MISSING = "IDENTITY_MISSING" + IDENTITY_MISMATCH = "IDENTITY_MISMATCH" + REDUCTION_SEMANTIC_MISMATCH = "REDUCTION_SEMANTIC_MISMATCH" + REDUCTION_MANDATE_VIOLATION = "REDUCTION_MANDATE_VIOLATION" + LSE_NOT_EXPORTED = "LSE_NOT_EXPORTED" + ROLE_COLLISION = "ROLE_COLLISION" + DETERMINISM_INCOMPATIBLE = "DETERMINISM_INCOMPATIBLE" + TOPOLOGY_MISMATCH = "TOPOLOGY_MISMATCH" + SPLIT_KV_RUNTIME_MISSING = "SPLIT_KV_RUNTIME_MISSING" + SPLIT_KV_MISMATCH = "SPLIT_KV_MISMATCH" + SPLIT_KV_FALLBACK = "SPLIT_KV_FALLBACK" + ATTENTION_PREPROCESS_MISSING = "ATTENTION_PREPROCESS_MISSING" + ATTENTION_PREPROCESS_MISMATCH = "ATTENTION_PREPROCESS_MISMATCH" + ATTENTION_PREPROCESS_FALLBACK = "ATTENTION_PREPROCESS_FALLBACK" + ATTENTION_PROJECTION_MISSING = "ATTENTION_PROJECTION_MISSING" + ATTENTION_PROJECTION_MISMATCH = "ATTENTION_PROJECTION_MISMATCH" + ATTENTION_CORE_MISSING = "ATTENTION_CORE_MISSING" + ATTENTION_CORE_MISMATCH = "ATTENTION_CORE_MISMATCH" + ATTENTION_CORE_SCHEDULE = "ATTENTION_CORE_SCHEDULE" + ATTENTION_NATIVE_ARITHMETIC = "ATTENTION_NATIVE_ARITHMETIC" + ATTENTION_CORE_SPLIT_K = "ATTENTION_CORE_SPLIT_K" + ATTENTION_BACKEND_MISSING = "ATTENTION_BACKEND_MISSING" + ATTENTION_NOT_PRODUCTION_READY = "ATTENTION_NOT_PRODUCTION_READY" + + +@dataclass(frozen=True) +class AttentionRuntimeReadback: + """Actual attention contract and all-rank Split-KV evidence from one engine.""" + + contract: AttentionContract + actual_knobs: Mapping[str, Any] + split_kv_plan_set: SplitKVRuntimePlanSet + source: str + frozen_scope_verified: bool + preprocess_backends: Mapping[str, str] = field(default_factory=dict) + preprocess_fallback: bool = False + preprocess_fallback_reason: str | None = None + preprocess_probe_id: str = "" + preprocess_policy_id: str = PREPROCESS_POLICY_ID + projection_plans: Mapping[str, Mapping[str, Any]] = field(default_factory=dict) + strict_mode: bool = False + strict_core_id: str | None = None + strict_schedule: str | None = None + native_attention_arithmetic: bool = True + strict_split_kv_policy: str | None = None + actual_backend: str | None = None + communication_backend: str | None = None + production_ready: bool = False + attention_fallback: bool = False + reference_only: bool = False + + def __post_init__(self) -> None: + if not isinstance(self.contract, AttentionContract): + raise TypeError("runtime readback contract must be an AttentionContract") + if not isinstance(self.actual_knobs, Mapping): + raise TypeError("runtime readback actual_knobs must be a mapping") + if not isinstance(self.split_kv_plan_set, SplitKVRuntimePlanSet): + raise TypeError("runtime readback requires a complete SplitKVRuntimePlanSet") + if not isinstance(self.source, str) or not self.source.strip(): + raise ValueError("runtime readback source must be a non-empty string") + if not isinstance(self.frozen_scope_verified, bool): + raise TypeError("frozen_scope_verified must be a bool") + if not isinstance(self.preprocess_backends, Mapping): + raise TypeError("runtime readback preprocess_backends must be a mapping") + for name, backend in self.preprocess_backends.items(): + if not isinstance(name, str) or not name.strip(): + raise ValueError("preprocess backend names must be non-empty strings") + if not isinstance(backend, str) or not backend.strip(): + raise ValueError("preprocess backend IDs must be non-empty strings") + if not isinstance(self.preprocess_fallback, bool): + raise TypeError("preprocess_fallback must be a bool") + if self.preprocess_fallback and not self.preprocess_fallback_reason: + raise ValueError( + "preprocess_fallback_reason is required when preprocess_fallback is true" + ) + if not isinstance(self.preprocess_probe_id, str): + raise TypeError("preprocess_probe_id must be a string") + if not isinstance(self.preprocess_policy_id, str) or not self.preprocess_policy_id.strip(): + raise ValueError("preprocess_policy_id must be a non-empty string") + if not isinstance(self.projection_plans, Mapping): + raise TypeError("projection_plans must be a mapping") + if not isinstance(self.strict_mode, bool): + raise TypeError("strict_mode must be a bool") + if self.strict_core_id is not None and ( + not isinstance(self.strict_core_id, str) or not self.strict_core_id.strip() + ): + raise ValueError("strict_core_id must be a non-empty string when provided") + if self.strict_schedule is not None and ( + not isinstance(self.strict_schedule, str) or not self.strict_schedule.strip() + ): + raise ValueError("strict_schedule must be a non-empty string when provided") + if not isinstance(self.native_attention_arithmetic, bool): + raise TypeError("native_attention_arithmetic must be a bool") + if self.strict_split_kv_policy is not None and self.strict_split_kv_policy not in { + "disabled", + "fixed", + "auto", + }: + raise ValueError("strict_split_kv_policy must be disabled, fixed, or auto") + for name, value in ( + ("actual_backend", self.actual_backend), + ("communication_backend", self.communication_backend), + ): + if value is not None and (not isinstance(value, str) or not value.strip()): + raise ValueError(f"{name} must be a non-empty string when provided") + if not isinstance(self.production_ready, bool): + raise TypeError("production_ready must be a bool") + if not isinstance(self.attention_fallback, bool): + raise TypeError("attention_fallback must be a bool") + if not isinstance(self.reference_only, bool): + raise TypeError("reference_only must be a bool") + normalized_projection_plans: dict[str, Mapping[str, Any]] = {} + for name, plan in self.projection_plans.items(): + if not isinstance(name, str) or not isinstance(plan, Mapping): + raise TypeError("projection_plans must map projection names to mappings") + normalized_projection_plans[name] = MappingProxyType(dict(plan)) + + plan_error = _split_kv_plan_contract_error(self.contract, self.split_kv_plan_set) + if plan_error is not None: + raise ValueError(plan_error) + object.__setattr__(self, "actual_knobs", MappingProxyType(dict(self.actual_knobs))) + object.__setattr__( + self, + "preprocess_backends", + MappingProxyType(dict(self.preprocess_backends)), + ) + object.__setattr__( + self, + "projection_plans", + MappingProxyType(normalized_projection_plans), + ) + + @property + def split_kv_fallback(self) -> bool: + return bool(_split_kv_fallbacks(self.split_kv_plan_set)) + + def to_dict(self) -> dict[str, Any]: + return { + "source": self.source, + "frozen_scope_verified": self.frozen_scope_verified, + "contract": self.contract.to_dict(), + "actual_knobs": dict(self.actual_knobs), + "attention_preprocess": { + "backends": dict(self.preprocess_backends), + "fallback": self.preprocess_fallback, + "fallback_reason": self.preprocess_fallback_reason, + "probe_id": self.preprocess_probe_id, + "policy_id": self.preprocess_policy_id, + }, + "attention_projections": { + name: dict(plan) for name, plan in self.projection_plans.items() + }, + "strict_attention": { + "enabled": self.strict_mode, + "core_id": self.strict_core_id, + "schedule": self.strict_schedule, + "native_attention_arithmetic": self.native_attention_arithmetic, + "split_kv_policy": self.strict_split_kv_policy, + }, + "runtime_backend": { + "actual_backend": self.actual_backend, + "communication_backend": self.communication_backend, + "production_ready": self.production_ready, + "fallback": self.attention_fallback, + "reference_only": self.reference_only, + }, + "split_kv_runtime_plan_set": self.split_kv_plan_set.to_dict(), + } + + +#: Attention exports attention-domain LSE, never vocab-logprob LSE (#235). +#: Recorded explicitly so a future ``LogprobContract`` binding cannot be confused +#: with this one purely because both set ``export_lse=True``. +ATTENTION_LSE_DOMAIN = "attention" + + +#: Fields both sides must agree on bit for bit before any comparison is meaningful. +#: Sourced from #235 "Numerical Contract" preconditions plus the vime-owned rollout +#: provenance (weight version, sampling, padding) that the issue assumes but does +#: not enumerate. +IDENTITY_FIELDS: tuple[str, ...] = ( + "checkpoint_id", + "model_version", + "weight_version", + "tokenizer_fingerprint", + "token_ids_fingerprint", + "active_mask_fingerprint", + "position_ids_fingerprint", + "padding_side", + "pre_update_state", + # model semantics that decide what attention *means* + "q_heads", + "kv_heads", + "head_dim", + "rope_theta", + "rope_scaling", + "rotary_dim", + "qk_layernorm", + # batch composition: batch-invariance is a claim about results not changing with + # batch makeup, so two sides scoring different batches are not comparable at all + "batch_size", + # decode replay identity (#235 PR6) + "global_token_positions_fingerprint", + "kv_seq_lens_fingerprint", +) + + +#: Reduction fields that decide the numerical result. Both sides must carry the +#: same value, and that value must satisfy :data:`WS2_ATTENTION_REDUCTION_MANDATE`. +SEMANTIC_REDUCTION_FIELDS: tuple[str, ...] = ( + "merge", + "acc_dtype", + "order", + "downcast_at", +) + + +#: Contract fields outside ``ReductionSpec`` that still decide the numerical result. +#: ``dtype`` is here rather than in :data:`RECORDED_FIELDS` because comparing a BF16 +#: rollout against an FP16 training pass produces a real drift number attributable to +#: nothing. #235 PR5 does sweep BF16 against an FP32 reference; that sweep opts in via +#: ``allow_dtype_difference`` instead of loosening the default. +SEMANTIC_CONTRACT_FIELDS: tuple[str, ...] = ("dtype",) + + +#: Sharding fields that determine local GQA head and sequence ownership. These are +#: comparison preconditions, not harmless backend provenance: a TP/CP mismatch +#: means the two ranks did not evaluate the same local attention problem. +TOPOLOGY_FIELDS: tuple[str, ...] = ( + "tp_rank", + "tp_world_size", + "cp_rank", + "cp_world_size", + "global_q_heads", + "global_kv_heads", + "local_q_head_start", + "local_q_heads", + "local_kv_head_start", + "local_kv_heads", + "global_sequence_length", + "local_sequence_length", + "global_block_indices", + "global_block_token_starts", + "local_block_offsets", + "packed_sequence_offsets", +) + + +#: The WS2 mandate itself. ``#236`` currently declares single-member enums for +#: ``merge`` / ``order`` / ``downcast_at``, so those checks are tautological today; +#: they are written out anyway so that widening any of those enums later fails here +#: instead of silently admitting a non-conforming backend. +WS2_ATTENTION_REDUCTION_MANDATE: Mapping[str, str] = { + "merge": AttentionMerge.ONLINE_SOFTMAX_LSE.value, + "acc_dtype": AttentionDType.FP32.value, + "order": ReductionOrder.GLOBAL_BLOCK_INDEX.value, + "downcast_at": DowncastPoint.FINAL_WRITE.value, +} + + +#: Materialization facts the two sides are expected to differ on. Recorded into +#: provenance; never a rejection reason. +RECORDED_FIELDS: tuple[str, ...] = ( + "mode", + "backend_id", + "reduction.engine", + "rope.fusion_boundary", + "rope.q_state", + "rope.k_state", + "rope.k_cache_state", + "rope.cast_at", + "rope.output_dtype", + "preprocess.qk_rmsnorm", + "preprocess.rope", + "preprocess.fallback", + "kv_cache.page_size", + "kv_cache.prefix_cache_enabled", + "kv_cache.block_table_shape", +) + + +@dataclass(frozen=True) +class BindingIssue: + """One reason a binding is not comparable or not admissible.""" + + code: BindingErrorCode + tier: BindingTier + field: str + rollout: Any = None + training: Any = None + message: str = "" + + def to_dict(self) -> dict[str, Any]: + return { + "code": self.code.value, + "tier": self.tier.value, + "field": self.field, + "rollout": self.rollout, + "training": self.training, + "message": self.message, + } + + +@dataclass(frozen=True) +class AttentionBindingResult: + """Outcome of binding one rollout contract to one training contract. + + ``comparable`` and ``passed`` are deliberately separate. A pair whose identity + does not match is *not comparable* -- reporting a drift number for it would be + meaningless. A pair that is comparable but violates the reduction mandate *is* + comparable yet must still fail closed, because the whole WS2 claim is that + reduction order and accumulation precision come from the contract. + """ + + comparable: bool + passed: bool + issues: tuple[BindingIssue, ...] = () + identity_fingerprint: str = "" + reduction_fingerprint: str = "" + binding_fingerprint: str = "" + recorded_differences: Mapping[str, Mapping[str, Any]] = field(default_factory=dict) + provenance: Mapping[str, Any] = field(default_factory=dict) + schema_version: str = "cross_config.attention_binding.v3" + + def issues_by_code(self, code: BindingErrorCode) -> tuple[BindingIssue, ...]: + return tuple(issue for issue in self.issues if issue.code is code) + + def to_dict(self) -> dict[str, Any]: + return { + "schema_version": self.schema_version, + "comparable": self.comparable, + "passed": self.passed, + "issues": [issue.to_dict() for issue in self.issues], + "identity_fingerprint": self.identity_fingerprint, + "reduction_fingerprint": self.reduction_fingerprint, + "binding_fingerprint": self.binding_fingerprint, + "recorded_differences": { + key: dict(value) for key, value in self.recorded_differences.items() + }, + "provenance": dict(self.provenance), + } + + +def _canonical_fingerprint(payload: Any) -> str: + encoded = json.dumps(payload, sort_keys=True, separators=(",", ":"), default=str) + return hashlib.sha256(encoded.encode("utf-8")).hexdigest() + + +def identity_fingerprint(identity: Mapping[str, Any]) -> str: + """Fingerprint only the declared :data:`IDENTITY_FIELDS`, in a fixed order. + + Extra keys in ``identity`` are ignored on purpose: callers pass whole + provenance bundles, and the fingerprint must not drift when an unrelated + diagnostic field is added. + """ + + return _canonical_fingerprint({name: identity.get(name) for name in IDENTITY_FIELDS}) + + +def _reduction_view(contract: AttentionContract) -> dict[str, Any]: + reduction = contract.reduction + return { + "merge": reduction.merge.value, + "acc_dtype": reduction.acc_dtype.value, + "order": reduction.order.value, + "downcast_at": reduction.downcast_at.value, + "engine": reduction.engine.value, + } + + +def _recorded_view( + contract: AttentionContract, + extra: Optional[Mapping[str, Any]] = None, +) -> dict[str, Any]: + rope = contract.rope + kv_cache = contract.kv_cache + view: dict[str, Any] = { + "mode": contract.mode.value, + "backend_id": None, + "reduction.engine": contract.reduction.engine.value, + } + if rope is not None: + view.update( + { + "rope.fusion_boundary": rope.fusion_boundary.value, + "rope.q_state": rope.q_state.value, + "rope.k_state": rope.k_state.value, + "rope.k_cache_state": rope.k_cache_state.value, + "rope.cast_at": rope.cast_at.value, + "rope.output_dtype": rope.output_dtype.value, + } + ) + if kv_cache is not None: + view.update( + { + "kv_cache.page_size": kv_cache.page_size, + "kv_cache.prefix_cache_enabled": kv_cache.prefix_cache_enabled, + "kv_cache.block_table_shape": [ + len(kv_cache.block_table), + max((len(row) for row in kv_cache.block_table), default=0), + ], + } + ) + if extra: + view.update(extra) + return view + + +def _topology_view(contract: AttentionContract) -> dict[str, Any]: + sharding = contract.sharding + return {name: getattr(sharding, name) for name in TOPOLOGY_FIELDS} + + +def _split_kv_fallbacks(plan_set: SplitKVRuntimePlanSet) -> list[dict[str, Any]]: + return [ + entry.to_dict() + for entry in plan_set.entries + if entry.execution.fallback + or entry.execution.actual_mode is None + or entry.execution.actual_mode is not entry.execution.requested_mode + or entry.execution.actual_split_size != entry.execution.requested_split_size + ] + + +def _split_kv_plan_contract_error( + contract: AttentionContract, + plan_set: SplitKVRuntimePlanSet, +) -> str | None: + sharding = contract.sharding + expected_topology = ( + contract.batch_size, + sharding.tp_world_size, + sharding.cp_world_size, + ) + actual_topology = ( + plan_set.batch_size, + plan_set.tp_world_size, + plan_set.cp_world_size, + ) + if actual_topology != expected_topology: + return ( + "Split-KV plan-set batch/TP/CP topology does not match the attention " + f"contract: actual={actual_topology}, expected={expected_topology}" + ) + if contract.mode.value in {"prefill", "chunked_prefill"}: + expected_totals = (sharding.global_sequence_length,) * contract.batch_size + if plan_set.total_kv_tokens != expected_totals: + return ( + "Split-KV plan-set KV lengths do not match the prefill attention " + f"contract: actual={plan_set.total_kv_tokens}, expected={expected_totals}" + ) + elif contract.kv_cache is not None: + expected_totals = contract.kv_cache.kv_seq_lens + if plan_set.total_kv_tokens != expected_totals: + return ( + "Split-KV plan-set KV lengths do not match decode KV-cache lengths: " + f"actual={plan_set.total_kv_tokens}, expected={expected_totals}" + ) + for entry in plan_set.entries: + execution = entry.execution + if ( + execution.requested_mode is not contract.split_kv.mode + or execution.requested_split_size != contract.split_kv.fixed_split_size + ): + return ( + "Split-KV runtime request does not match the first-class attention " + f"contract at {entry.coordinate}" + ) + return None + + +#: Identity fields where ``None`` is a real value rather than an omission. Qwen3-8B +#: applies no RoPE scaling, so ``rope_scaling=None`` must not read as "undeclared" -- +#: both sides still have to agree on it, which the equality pass below handles. +NULLABLE_IDENTITY_FIELDS: frozenset[str] = frozenset({"rope_scaling"}) + + +def _missing_identity_fields(identity: Mapping[str, Any]) -> tuple[str, ...]: + return tuple( + name + for name in IDENTITY_FIELDS + if name not in NULLABLE_IDENTITY_FIELDS and identity.get(name) is None + ) + + +def bind_attention_contracts( + *, + rollout_contract: AttentionContract, + training_contract: AttentionContract, + rollout_identity: Mapping[str, Any], + training_identity: Mapping[str, Any], + rollout_backend_id: str, + training_backend_id: str, + determinism_issues: Sequence[BindingIssue] = (), + require_full_identity: bool = True, + allow_dtype_difference: bool = False, + rollout_split_kv_plan_set: Optional[SplitKVRuntimePlanSet] = None, + training_split_kv_plan_set: Optional[SplitKVRuntimePlanSet] = None, + rollout_recorded_extra: Optional[Mapping[str, Any]] = None, + training_recorded_extra: Optional[Mapping[str, Any]] = None, +) -> AttentionBindingResult: + """Bind a rollout attention contract to a training attention contract. + + ``determinism_issues`` is threaded in from + :mod:`rl_engine.alignment.cross_config.determinism` rather than computed here, + so that this module stays free of framework probing and remains testable + without Megatron or vLLM present. + + ``require_full_identity`` exists for the single-GPU harness in #235 PR2, which + legitimately has no KV-cache or decode identity to declare. Distributed callers + must leave it at ``True``. + + ``allow_dtype_difference`` exists for the #235 PR5 sweep that deliberately scores + a BF16 path against an FP32 reference. It must stay ``False`` everywhere else. + + Strict binding requires complete actual Split-KV plan sets from both runtimes. + A configured policy is insufficient because auto-selection, graph capture, and + backend fallbacks can change the executed boundaries. The plan sets cover the + complete batch x TP x CP x KV-owner Cartesian product. + + ``rollout_recorded_extra`` / ``training_recorded_extra`` are diagnostic-only + backend facts. They can never make a semantic mismatch admissible. + """ + + if rollout_contract.role is not AttentionRole.INFER: + raise AttentionBindingError( + f"rollout_contract.role must be {AttentionRole.INFER.value!r}, " + f"got {rollout_contract.role.value!r}" + ) + if training_contract.role is not AttentionRole.TRAIN: + raise AttentionBindingError( + f"training_contract.role must be {AttentionRole.TRAIN.value!r}, " + f"got {training_contract.role.value!r}" + ) + + issues: list[BindingIssue] = [] + + # ---- tier 1: identity, bit for bit ------------------------------------- + if require_full_identity: + for side, identity in (("rollout", rollout_identity), ("training", training_identity)): + for name in _missing_identity_fields(identity): + issues.append( + BindingIssue( + code=BindingErrorCode.IDENTITY_MISSING, + tier=BindingTier.IDENTICAL, + field=f"{side}.{name}", + message=f"{side} identity does not declare {name!r}", + ) + ) + + for name in IDENTITY_FIELDS: + rollout_value = rollout_identity.get(name) + training_value = training_identity.get(name) + if rollout_value != training_value: + issues.append( + BindingIssue( + code=BindingErrorCode.IDENTITY_MISMATCH, + tier=BindingTier.IDENTICAL, + field=name, + rollout=rollout_value, + training=training_value, + message=( + f"{name!r} differs between sides; the pair is not comparable " + "and any drift computed from it is meaningless" + ), + ) + ) + + comparable = not any(issue.tier is BindingTier.IDENTICAL for issue in issues) + + rollout_topology = _topology_view(rollout_contract) + training_topology = _topology_view(training_contract) + for name in TOPOLOGY_FIELDS: + if rollout_topology[name] != training_topology[name]: + issues.append( + BindingIssue( + code=BindingErrorCode.TOPOLOGY_MISMATCH, + tier=BindingTier.IDENTICAL, + field=f"sharding.{name}", + rollout=rollout_topology[name], + training=training_topology[name], + message=( + f"sharding.{name} changes TP/CP ownership; the pair is not " + "the same local attention problem" + ), + ) + ) + + comparable = not any(issue.tier is BindingTier.IDENTICAL for issue in issues) + + # ---- tier 2: reduction semantics, and the WS2 mandate ------------------- + rollout_reduction = _reduction_view(rollout_contract) + training_reduction = _reduction_view(training_contract) + + for name in SEMANTIC_REDUCTION_FIELDS: + if rollout_reduction[name] != training_reduction[name]: + issues.append( + BindingIssue( + code=BindingErrorCode.REDUCTION_SEMANTIC_MISMATCH, + tier=BindingTier.SEMANTIC, + field=f"reduction.{name}", + rollout=rollout_reduction[name], + training=training_reduction[name], + message=( + f"reduction.{name!r} must be decided by the contract, not by the " + "backend; the two sides disagree" + ), + ) + ) + mandated = WS2_ATTENTION_REDUCTION_MANDATE[name] + for side, view in (("rollout", rollout_reduction), ("training", training_reduction)): + if view[name] != mandated: + issues.append( + BindingIssue( + code=BindingErrorCode.REDUCTION_MANDATE_VIOLATION, + tier=BindingTier.SEMANTIC, + field=f"{side}.reduction.{name}", + rollout=rollout_reduction[name], + training=training_reduction[name], + message=( + f"WS2 requires reduction.{name} == {mandated!r}; " + f"{side} declares {view[name]!r}" + ), + ) + ) + + if not allow_dtype_difference and rollout_contract.dtype is not training_contract.dtype: + issues.append( + BindingIssue( + code=BindingErrorCode.REDUCTION_SEMANTIC_MISMATCH, + tier=BindingTier.SEMANTIC, + field="dtype", + rollout=rollout_contract.dtype.value, + training=training_contract.dtype.value, + message=( + "the two sides compute in different dtypes; the resulting drift is " + "not attributable. Pass allow_dtype_difference=True only for a " + "deliberate precision sweep" + ), + ) + ) + + if rollout_contract.split_kv != training_contract.split_kv: + issues.append( + BindingIssue( + code=BindingErrorCode.SPLIT_KV_MISMATCH, + tier=BindingTier.SEMANTIC, + field="split_kv", + rollout=rollout_contract.split_kv.to_dict(), + training=training_contract.split_kv.to_dict(), + message="training and rollout must request the same first-class Split-KV policy", + ) + ) + + for side, plan_set in ( + ("rollout", rollout_split_kv_plan_set), + ("training", training_split_kv_plan_set), + ): + if plan_set is None: + issues.append( + BindingIssue( + code=BindingErrorCode.SPLIT_KV_RUNTIME_MISSING, + tier=BindingTier.SEMANTIC, + field=f"{side}.split_kv_runtime_plan_set", + message=( + f"{side} did not report a complete actual Split-KV plan set; " + "configured policy alone is not runtime evidence" + ), + ) + ) + continue + contract = rollout_contract if side == "rollout" else training_contract + contract_error = _split_kv_plan_contract_error(contract, plan_set) + if contract_error is not None: + issues.append( + BindingIssue( + code=BindingErrorCode.SPLIT_KV_MISMATCH, + tier=BindingTier.SEMANTIC, + field=f"{side}.split_kv_runtime_plan_set", + rollout=plan_set.to_dict() if side == "rollout" else None, + training=plan_set.to_dict() if side == "training" else None, + message=contract_error, + ) + ) + fallbacks = _split_kv_fallbacks(plan_set) + if fallbacks: + issues.append( + BindingIssue( + code=BindingErrorCode.SPLIT_KV_FALLBACK, + tier=BindingTier.SEMANTIC, + field=f"{side}.split_kv_runtime_plan_set", + rollout=fallbacks if side == "rollout" else None, + training=fallbacks if side == "training" else None, + message=f"{side} Split-KV runtime used an unknown or fallback plan", + ) + ) + + if rollout_split_kv_plan_set is not None and training_split_kv_plan_set is not None: + try: + validate_split_kv_plan_set_alignment( + training_split_kv_plan_set, + rollout_split_kv_plan_set, + ) + except AttentionContractError as exc: + issues.append( + BindingIssue( + code=BindingErrorCode.SPLIT_KV_MISMATCH, + tier=BindingTier.SEMANTIC, + field="split_kv_runtime_plan_set", + rollout=rollout_split_kv_plan_set.to_dict(), + training=training_split_kv_plan_set.to_dict(), + message=str(exc), + ) + ) + + for side, contract in (("rollout", rollout_contract), ("training", training_contract)): + if not contract.export_lse: + issues.append( + BindingIssue( + code=BindingErrorCode.LSE_NOT_EXPORTED, + tier=BindingTier.SEMANTIC, + field=f"{side}.export_lse", + message=( + "attention-domain LSE must be exported; without it the deterministic " + "CP merge cannot be validated" + ), + ) + ) + + if rollout_backend_id == training_backend_id and rollout_backend_id: + # Not an error, but worth surfacing: an identical backend on both sides means + # the experiment is not actually measuring a cross-implementation difference. + pass + + issues.extend(determinism_issues) + + # ---- tier 3: recorded differences -------------------------------------- + rollout_recorded = _recorded_view(rollout_contract, rollout_recorded_extra) + rollout_recorded["backend_id"] = rollout_backend_id + training_recorded = _recorded_view(training_contract, training_recorded_extra) + training_recorded["backend_id"] = training_backend_id + + recorded_differences: dict[str, dict[str, Any]] = {} + for name in RECORDED_FIELDS: + rollout_value = rollout_recorded.get(name) + training_value = training_recorded.get(name) + if rollout_value != training_value: + recorded_differences[name] = { + "rollout": rollout_value, + "training": training_value, + } + + identity_fp = identity_fingerprint(training_identity if comparable else rollout_identity) + reduction_fp = _canonical_fingerprint( + {name: training_reduction[name] for name in SEMANTIC_REDUCTION_FIELDS} + ) + passed = comparable and not any(issue.tier is BindingTier.SEMANTIC for issue in issues) + + provenance = { + "lse_domain": ATTENTION_LSE_DOMAIN, + "dtype": training_contract.dtype.value, + "split_kv_runtime": { + "rollout": ( + None if rollout_split_kv_plan_set is None else rollout_split_kv_plan_set.to_dict() + ), + "training": ( + None if training_split_kv_plan_set is None else training_split_kv_plan_set.to_dict() + ), + }, + "rollout": { + "contract": rollout_contract.to_dict(), + "backend_id": rollout_backend_id, + "recorded": rollout_recorded, + }, + "training": { + "contract": training_contract.to_dict(), + "backend_id": training_backend_id, + "recorded": training_recorded, + }, + } + + return AttentionBindingResult( + comparable=comparable, + passed=passed, + issues=tuple(issues), + identity_fingerprint=identity_fp, + reduction_fingerprint=reduction_fp, + binding_fingerprint=_canonical_fingerprint( + { + "identity": identity_fp, + "reduction": reduction_fp, + "topology": training_topology, + "split_kv": provenance["split_kv_runtime"], + "lse_domain": ATTENTION_LSE_DOMAIN, + "rollout_backend": rollout_backend_id, + "training_backend": training_backend_id, + "attention_preprocess": { + "rollout": { + name: rollout_recorded.get(f"preprocess.{name}") + for name in ( + *MANDATED_ATTENTION_PREPROCESS_BACKENDS, + "fallback", + ) + }, + "training": { + name: training_recorded.get(f"preprocess.{name}") + for name in ( + *MANDATED_ATTENTION_PREPROCESS_BACKENDS, + "fallback", + ) + }, + }, + } + ), + recorded_differences=recorded_differences, + provenance=provenance, + ) + + +def bind_attention_runtime_readbacks( + *, + rollout: AttentionRuntimeReadback, + training: AttentionRuntimeReadback, + rollout_identity: Mapping[str, Any], + training_identity: Mapping[str, Any], + rollout_backend_id: str, + training_backend_id: str, + determinism_issues: Sequence[BindingIssue] = (), +) -> AttentionBindingResult: + """Strict public handoff from executed framework runtimes to PR4 binding. + + The Megatron/vLLM launchers remain environment-owned. Once both launchers have + reconstructed their actual contracts and all-rank Split-KV reports, this entry + point performs the complete comparison without accepting configured-only data. + """ + + missing_scope_evidence = [] + for side, readback in (("rollout", rollout), ("training", training)): + if not readback.frozen_scope_verified: + missing_scope_evidence.append( + BindingIssue( + code=BindingErrorCode.DETERMINISM_INCOMPATIBLE, + tier=BindingTier.SEMANTIC, + field=f"{side}.frozen_scope_verified", + message=f"{side} runtime did not verify the frozen attention scope", + ) + ) + missing_scope_evidence.extend(_attention_preprocess_issues(side, readback)) + missing_scope_evidence.extend(_attention_projection_issues(side, readback)) + missing_scope_evidence.extend(_strict_attention_core_issues(side, readback)) + missing_scope_evidence.extend(_strict_attention_core_pair_issues(rollout, training)) + if rollout.preprocess_fallback != training.preprocess_fallback: + missing_scope_evidence.append( + BindingIssue( + code=BindingErrorCode.ATTENTION_PREPROCESS_MISMATCH, + tier=BindingTier.SEMANTIC, + field="preprocess.fallback", + rollout=rollout.preprocess_fallback, + training=training.preprocess_fallback, + message=( + "both runtimes must either pass the platform vendor bitwise probe or " + "use the same deterministic preprocess fallback" + ), + ) + ) + if rollout.preprocess_policy_id != training.preprocess_policy_id: + missing_scope_evidence.append( + BindingIssue( + code=BindingErrorCode.ATTENTION_PREPROCESS_MISMATCH, + tier=BindingTier.SEMANTIC, + field="preprocess.policy_id", + rollout=rollout.preprocess_policy_id, + training=training.preprocess_policy_id, + message="QK-Norm/RoPE policy IDs differ between runtimes", + ) + ) + for name in MANDATED_ATTENTION_PREPROCESS_BACKENDS: + rollout_backend = rollout.preprocess_backends.get(name) + training_backend = training.preprocess_backends.get(name) + if rollout_backend != training_backend: + missing_scope_evidence.append( + BindingIssue( + code=BindingErrorCode.ATTENTION_PREPROCESS_MISMATCH, + tier=BindingTier.SEMANTIC, + field=f"preprocess.{name}", + rollout=rollout_backend, + training=training_backend, + message="training and rollout must execute the same preprocess backend", + ) + ) + for projection in ("qkv", "o_proj"): + rollout_plan = rollout.projection_plans.get(projection, {}) + training_plan = training.projection_plans.get(projection, {}) + rollout_fallback = bool(rollout_plan.get("fallback", False)) + training_fallback = bool(training_plan.get("fallback", False)) + if rollout_fallback != training_fallback: + missing_scope_evidence.append( + BindingIssue( + code=BindingErrorCode.ATTENTION_PROJECTION_MISMATCH, + tier=BindingTier.SEMANTIC, + field=f"projection.{projection}.fallback", + rollout=rollout_fallback, + training=training_fallback, + message=( + "QKV/o_proj must use the same native-or-deterministic " "path on both sides" + ), + ) + ) + for field_name in ("backend_id", "policy_id", "split_k", "reduction_order"): + if rollout_plan.get(field_name) != training_plan.get(field_name): + missing_scope_evidence.append( + BindingIssue( + code=BindingErrorCode.ATTENTION_PROJECTION_MISMATCH, + tier=BindingTier.SEMANTIC, + field=f"projection.{projection}.{field_name}", + rollout=rollout_plan.get(field_name), + training=training_plan.get(field_name), + message="projection execution evidence differs between sides", + ) + ) + return bind_attention_contracts( + rollout_contract=rollout.contract, + training_contract=training.contract, + rollout_identity=rollout_identity, + training_identity=training_identity, + rollout_backend_id=rollout_backend_id, + training_backend_id=training_backend_id, + determinism_issues=tuple(determinism_issues) + tuple(missing_scope_evidence), + rollout_split_kv_plan_set=rollout.split_kv_plan_set, + training_split_kv_plan_set=training.split_kv_plan_set, + rollout_recorded_extra={ + **{ + f"preprocess.{name}": backend + for name, backend in rollout.preprocess_backends.items() + }, + "preprocess.fallback": rollout.preprocess_fallback, + "preprocess.fallback_reason": rollout.preprocess_fallback_reason, + "preprocess.probe_id": rollout.preprocess_probe_id, + "preprocess.policy_id": rollout.preprocess_policy_id, + "strict.enabled": rollout.strict_mode, + "strict.core_id": rollout.strict_core_id, + "strict.schedule": rollout.strict_schedule, + "strict.native_attention_arithmetic": rollout.native_attention_arithmetic, + "strict.split_kv_policy": rollout.strict_split_kv_policy, + "runtime.actual_backend": rollout.actual_backend, + "runtime.communication_backend": rollout.communication_backend, + "runtime.production_ready": rollout.production_ready, + "runtime.fallback": rollout.attention_fallback, + "runtime.reference_only": rollout.reference_only, + **{ + f"projection.{projection}": dict(plan) + for projection, plan in rollout.projection_plans.items() + }, + }, + training_recorded_extra={ + **{ + f"preprocess.{name}": backend + for name, backend in training.preprocess_backends.items() + }, + "preprocess.fallback": training.preprocess_fallback, + "preprocess.fallback_reason": training.preprocess_fallback_reason, + "preprocess.probe_id": training.preprocess_probe_id, + "preprocess.policy_id": training.preprocess_policy_id, + "strict.enabled": training.strict_mode, + "strict.core_id": training.strict_core_id, + "strict.schedule": training.strict_schedule, + "strict.native_attention_arithmetic": training.native_attention_arithmetic, + "strict.split_kv_policy": training.strict_split_kv_policy, + "runtime.actual_backend": training.actual_backend, + "runtime.communication_backend": training.communication_backend, + "runtime.production_ready": training.production_ready, + "runtime.fallback": training.attention_fallback, + "runtime.reference_only": training.reference_only, + **{ + f"projection.{projection}": dict(plan) + for projection, plan in training.projection_plans.items() + }, + }, + ) + + +def _strict_attention_core_issues( + side: str, + readback: AttentionRuntimeReadback, +) -> list[BindingIssue]: + if not readback.strict_mode: + return [] + issues = [] + if not readback.actual_backend: + issues.append( + BindingIssue( + code=BindingErrorCode.ATTENTION_BACKEND_MISSING, + tier=BindingTier.SEMANTIC, + field=f"{side}.runtime.actual_backend", + rollout=readback.actual_backend if side == "rollout" else None, + training=readback.actual_backend if side == "training" else None, + message=(f"{side} strict Attention did not report its executed production core"), + ) + ) + supported_communication = {"self_owned_cuda_ag_rs", "cuda_ag_rs", "rccl_ag_rs"} + if ( + readback.contract.sharding.cp_world_size > 1 + and readback.communication_backend not in supported_communication + ): + issues.append( + BindingIssue( + code=BindingErrorCode.ATTENTION_BACKEND_MISSING, + tier=BindingTier.SEMANTIC, + field=f"{side}.runtime.communication_backend", + rollout=readback.communication_backend if side == "rollout" else None, + training=readback.communication_backend if side == "training" else None, + message=( + f"{side} strict CP Attention did not execute a supported self-owned AG/RS path" + ), + ) + ) + if not readback.production_ready: + issues.append( + BindingIssue( + code=BindingErrorCode.ATTENTION_NOT_PRODUCTION_READY, + tier=BindingTier.SEMANTIC, + field=f"{side}.runtime.production_ready", + rollout=False if side == "rollout" else None, + training=False if side == "training" else None, + message=(f"{side} evidence is reference-only and cannot close the production gate"), + ) + ) + if not readback.strict_core_id: + issues.append( + BindingIssue( + code=BindingErrorCode.ATTENTION_CORE_MISSING, + tier=BindingTier.SEMANTIC, + field=f"{side}.strict.core_id", + message=(f"{side} strict Attention did not report its exact shared core identity"), + ) + ) + if not readback.strict_schedule: + issues.append( + BindingIssue( + code=BindingErrorCode.ATTENTION_CORE_SCHEDULE, + tier=BindingTier.SEMANTIC, + field=f"{side}.strict.schedule", + message=(f"{side} strict Attention did not report its exact reduction schedule"), + ) + ) + if readback.attention_fallback or readback.reference_only: + issues.append( + BindingIssue( + code=BindingErrorCode.ATTENTION_BACKEND_MISSING, + tier=BindingTier.SEMANTIC, + field=f"{side}.runtime.production_path", + message=f"{side} strict Attention executed a fallback or reference-only path", + ) + ) + if ( + readback.strict_split_kv_policy != "disabled" + or readback.contract.split_kv.mode.value != "disabled" + ): + issues.append( + BindingIssue( + code=BindingErrorCode.ATTENTION_CORE_SPLIT_K, + tier=BindingTier.SEMANTIC, + field=f"{side}.strict.split_kv_policy", + message=( + f"{side} strict Attention did not prove Split-KV disabled " + "in both runtime evidence and AttentionContract" + ), + ) + ) + return issues + + +def _strict_attention_core_pair_issues( + rollout: AttentionRuntimeReadback, + training: AttentionRuntimeReadback, +) -> list[BindingIssue]: + if rollout.strict_mode != training.strict_mode: + return [ + BindingIssue( + code=BindingErrorCode.ATTENTION_CORE_MISMATCH, + tier=BindingTier.SEMANTIC, + field="strict.enabled", + rollout=rollout.strict_mode, + training=training.strict_mode, + message="training and rollout must use the same strict Attention mode", + ) + ] + if rollout.strict_mode and rollout.strict_core_id != training.strict_core_id: + return [ + BindingIssue( + code=BindingErrorCode.ATTENTION_CORE_MISMATCH, + tier=BindingTier.SEMANTIC, + field="strict.core_id", + rollout=rollout.strict_core_id, + training=training.strict_core_id, + message="training and rollout executed different Attention cores", + ) + ] + if rollout.strict_mode and rollout.strict_schedule != training.strict_schedule: + return [ + BindingIssue( + code=BindingErrorCode.ATTENTION_CORE_SCHEDULE, + tier=BindingTier.SEMANTIC, + field="strict.schedule", + rollout=rollout.strict_schedule, + training=training.strict_schedule, + message="training and rollout executed different strict Attention schedules", + ) + ] + if ( + rollout.strict_mode + and rollout.native_attention_arithmetic != training.native_attention_arithmetic + ): + return [ + BindingIssue( + code=BindingErrorCode.ATTENTION_NATIVE_ARITHMETIC, + tier=BindingTier.SEMANTIC, + field="strict.native_attention_arithmetic", + rollout=rollout.native_attention_arithmetic, + training=training.native_attention_arithmetic, + message="training and rollout disagree on vendor Attention arithmetic", + ) + ] + if rollout.strict_mode and rollout.actual_backend != training.actual_backend: + return [ + BindingIssue( + code=BindingErrorCode.ATTENTION_CORE_MISMATCH, + tier=BindingTier.SEMANTIC, + field="runtime.actual_backend", + rollout=rollout.actual_backend, + training=training.actual_backend, + message="training and rollout executed different Attention backends", + ) + ] + if rollout.strict_mode and rollout.communication_backend != training.communication_backend: + return [ + BindingIssue( + code=BindingErrorCode.ATTENTION_CORE_MISMATCH, + tier=BindingTier.SEMANTIC, + field="runtime.communication_backend", + rollout=rollout.communication_backend, + training=training.communication_backend, + message="training and rollout executed different Attention communication backends", + ) + ] + return [] + + +def _attention_preprocess_issues( + side: str, + readback: AttentionRuntimeReadback, +) -> list[BindingIssue]: + issues: list[BindingIssue] = [] + for name, mandated in MANDATED_ATTENTION_PREPROCESS_BACKENDS.items(): + actual = readback.preprocess_backends.get(name) + if actual is None: + issues.append( + BindingIssue( + code=BindingErrorCode.ATTENTION_PREPROCESS_MISSING, + tier=BindingTier.SEMANTIC, + field=f"{side}.preprocess.{name}", + message=( + f"{side} did not report the executed {name} backend; " + "runtime-native execution cannot validate the Attention input boundary" + ), + ) + ) + elif actual not in ALLOWED_ATTENTION_PREPROCESS_BACKENDS[name]: + issues.append( + BindingIssue( + code=BindingErrorCode.ATTENTION_PREPROCESS_MISMATCH, + tier=BindingTier.SEMANTIC, + field=f"{side}.preprocess.{name}", + rollout=actual if side == "rollout" else None, + training=actual if side == "training" else None, + message=( + f"{side} executed {actual!r}; " + f"the experiment requires {mandated!r} or a verified platform backend" + ), + ) + ) + if readback.preprocess_fallback and any( + str(readback.preprocess_backends.get(name, "")).startswith("transformer_engine.") + for name in MANDATED_ATTENTION_PREPROCESS_BACKENDS + ): + issues.append( + BindingIssue( + code=BindingErrorCode.ATTENTION_PREPROCESS_FALLBACK, + tier=BindingTier.SEMANTIC, + field=f"{side}.preprocess.fallback", + message=( + f"{side} reported fallback while still claiming a vendor preprocess backend" + ), + ) + ) + return issues + + +def _attention_projection_issues( + side: str, + readback: AttentionRuntimeReadback, +) -> list[BindingIssue]: + issues: list[BindingIssue] = [] + expected_collectives = { + "qkv": QKV_COLLECTIVE_CONTRACT.to_dict(), + "o_proj": O_PROJ_COLLECTIVE_CONTRACT.to_dict(), + } + fixed_fields = { + "input_dtype": "torch.bfloat16", + "weight_dtype": "torch.bfloat16", + "output_dtype": "torch.bfloat16", + "accumulation_dtype": "torch.float32", + "reduction_order": "k_ascending", + "split_k": False, + "policy_id": PROJECTION_POLICY_ID, + } + for projection, expected_collective in expected_collectives.items(): + plan = readback.projection_plans.get(projection) + if plan is None: + issues.append( + BindingIssue( + code=BindingErrorCode.ATTENTION_PROJECTION_MISSING, + tier=BindingTier.SEMANTIC, + field=f"{side}.projection.{projection}", + message=f"{side} did not report the executed {projection} projection plan", + ) + ) + continue + for field_name, expected in fixed_fields.items(): + actual = plan.get(field_name) + if actual != expected: + issues.append( + BindingIssue( + code=BindingErrorCode.ATTENTION_PROJECTION_MISMATCH, + tier=BindingTier.SEMANTIC, + field=f"{side}.projection.{projection}.{field_name}", + rollout=actual if side == "rollout" else None, + training=actual if side == "training" else None, + message=f"{side} {projection} {field_name} must be {expected!r}", + ) + ) + collective = plan.get("collective") + if not isinstance(collective, Mapping) or dict(collective) != expected_collective: + issues.append( + BindingIssue( + code=BindingErrorCode.ATTENTION_PROJECTION_MISMATCH, + tier=BindingTier.SEMANTIC, + field=f"{side}.projection.{projection}.collective", + message=f"{side} {projection} TP/SP collective directions are invalid", + ) + ) + backend_id = plan.get("backend_id") + if not isinstance(backend_id, str) or not backend_id.strip(): + issues.append( + BindingIssue( + code=BindingErrorCode.ATTENTION_PROJECTION_MISSING, + tier=BindingTier.SEMANTIC, + field=f"{side}.projection.{projection}.backend_id", + message=f"{side} {projection} backend identity is missing", + ) + ) + if not isinstance(plan.get("probe_id"), str) or not plan.get("probe_id"): + issues.append( + BindingIssue( + code=BindingErrorCode.ATTENTION_PROJECTION_MISSING, + tier=BindingTier.SEMANTIC, + field=f"{side}.projection.{projection}.probe_id", + message=f"{side} {projection} bitwise probe identity is missing", + ) + ) + if plan.get("fallback"): + deterministic_backends = { + CUDA_DETERMINISTIC_PROJECTION_BACKEND_ID, + ROCM_DETERMINISTIC_PROJECTION_BACKEND_ID, + } + if backend_id not in deterministic_backends or not plan.get("fallback_reason"): + issues.append( + BindingIssue( + code=BindingErrorCode.ATTENTION_PROJECTION_MISMATCH, + tier=BindingTier.SEMANTIC, + field=f"{side}.projection.{projection}.fallback", + message=( + f"{side} {projection} fallback must execute DetGemmOp and record why" + ), + ) + ) + return issues + + +def summarize_binding(result: AttentionBindingResult) -> str: + """One-line human summary for CLI output and failure messages.""" + + if result.passed: + return ( + f"attention binding OK " + f"(identity={result.identity_fingerprint[:12]}, " + f"{len(result.recorded_differences)} recorded difference(s))" + ) + if not result.comparable: + fields = ", ".join( + sorted({issue.field for issue in result.issues if issue.tier is BindingTier.IDENTICAL}) + ) + return f"attention binding NOT COMPARABLE; identity problems: {fields}" + fields = ", ".join( + sorted({issue.field for issue in result.issues if issue.tier is BindingTier.SEMANTIC}) + ) + return f"attention binding FAILED CLOSED; semantic problems: {fields}" + + +def first_blocking_issue( + result: AttentionBindingResult, +) -> Optional[BindingIssue]: + """Return the issue a caller should report, preferring identity over semantics.""" + + for tier in (BindingTier.IDENTICAL, BindingTier.SEMANTIC): + for issue in result.issues: + if issue.tier is tier: + return issue + return None diff --git a/rl_engine/alignment/cross_config/comparison.py b/rl_engine/alignment/cross_config/comparison.py new file mode 100644 index 00000000..135dbd75 --- /dev/null +++ b/rl_engine/alignment/cross_config/comparison.py @@ -0,0 +1,321 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Fixed-contract selected-token comparison for cross-configuration cases.""" + +from __future__ import annotations + +import math +from dataclasses import fields +from typing import Any + +import torch + +from rl_engine.alignment.cross_config.schema import ( + AlignmentResult, + AlignmentStatus, + ScoreArtifact, + ScoreSide, + SemanticIdentitySpec, + TokenComparisonArtifact, +) +from rl_engine.kernels.gtest.tolerance import ( + resolve_logprob_threshold, + tolerance_contract_fingerprint, +) + + +def semantic_identity_errors( + rollout: SemanticIdentitySpec, + training: SemanticIdentitySpec, +) -> tuple[str, ...]: + """Return every logical identity field that differs between the two sides.""" + + return tuple( + item.name + for item in fields(SemanticIdentitySpec) + if getattr(rollout, item.name) != getattr(training, item.name) + ) + + +def recompute_mismatch_mask( + rollout_logprobs: torch.Tensor, + training_logprobs: torch.Tensor, + active_mask: torch.Tensor, + fixed_threshold: float, +) -> torch.Tensor: + """Recompute the sole token mismatch signal from persisted tensors.""" + + if rollout_logprobs.shape != training_logprobs.shape: + raise ValueError("rollout and training logprobs must have identical shapes") + if active_mask.shape != rollout_logprobs.shape: + raise ValueError("active_mask shape must match selected logprobs") + if fixed_threshold < 0.0: + raise ValueError("fixed_threshold must be non-negative") + active = active_mask.to(device=rollout_logprobs.device, dtype=torch.bool) + training = training_logprobs.to(device=rollout_logprobs.device) + return active & (torch.abs(training - rollout_logprobs) > fixed_threshold) + + +class FixedThresholdComparator: + """Compare paired selected logprobs using only the current WS1 contract.""" + + def compare(self, rollout: ScoreArtifact, training: ScoreArtifact) -> AlignmentResult: + contract_fingerprint = tolerance_contract_fingerprint() + artifact_errors = _artifact_errors(rollout, training) + if artifact_errors: + return AlignmentResult( + case_id=rollout.case_id, + attempt_id=rollout.attempt_id, + status=AlignmentStatus.INVALID_ARTIFACT, + comparable=False, + passed=False, + active_token_count=0, + mismatch_count=0, + contract_fingerprint=contract_fingerprint, + artifact_errors=artifact_errors, + ) + + identity_errors = list(semantic_identity_errors(rollout.identity, training.identity)) + identity_errors.extend(_artifact_identity_errors(rollout, training)) + if identity_errors: + return AlignmentResult( + case_id=rollout.case_id, + attempt_id=rollout.attempt_id, + status=AlignmentStatus.INVALID_IDENTITY, + comparable=False, + passed=False, + active_token_count=0, + mismatch_count=0, + contract_fingerprint=contract_fingerprint, + identity_errors=tuple(dict.fromkeys(identity_errors)), + ) + + threshold, threshold_error = _resolve_fixed_threshold(rollout, training) + if threshold_error is not None: + return AlignmentResult( + case_id=rollout.case_id, + attempt_id=rollout.attempt_id, + status=AlignmentStatus.INVALID_ARTIFACT, + comparable=False, + passed=False, + active_token_count=0, + mismatch_count=0, + contract_fingerprint=contract_fingerprint, + artifact_errors=(threshold_error,), + ) + assert threshold is not None + fixed_threshold = threshold + + rollout_logprobs = rollout.selected_logprobs.detach().cpu() + training_logprobs = training.selected_logprobs.detach().cpu() + active_mask = rollout.active_mask.detach().cpu().to(dtype=torch.bool) + active_token_count = int(active_mask.sum().item()) + if active_token_count: + active_rollout = rollout_logprobs[active_mask] + active_training = training_logprobs[active_mask] + if not bool(torch.isfinite(active_rollout).all().item()) or not bool( + torch.isfinite(active_training).all().item() + ): + return AlignmentResult( + case_id=rollout.case_id, + attempt_id=rollout.attempt_id, + status=AlignmentStatus.INVALID_ARTIFACT, + comparable=False, + passed=False, + active_token_count=active_token_count, + mismatch_count=0, + contract_fingerprint=contract_fingerprint, + fixed_threshold=fixed_threshold, + artifact_errors=("active selected logprobs must be finite",), + ) + + # Inactive positions are outside the numerical contract. Canonicalize + # them before persistence so an ignored NaN/Inf cannot break strict JSON + # serialization or make resume artifacts non-reproducible. + rollout_logprobs = rollout_logprobs.masked_fill(~active_mask, 0.0) + training_logprobs = training_logprobs.masked_fill(~active_mask, 0.0) + absolute_diff = torch.abs(training_logprobs - rollout_logprobs) + mismatch_mask = recompute_mismatch_mask( + rollout_logprobs, + training_logprobs, + active_mask, + fixed_threshold, + ) + token_artifact = TokenComparisonArtifact( + rollout_logprobs=rollout_logprobs, + training_logprobs=training_logprobs, + active_mask=active_mask, + absolute_diff=absolute_diff, + mismatch_mask=mismatch_mask, + fixed_threshold=fixed_threshold, + ) + if active_token_count == 0: + return AlignmentResult( + case_id=rollout.case_id, + attempt_id=rollout.attempt_id, + status=AlignmentStatus.ZERO_ACTIVE_TOKENS, + comparable=False, + passed=False, + active_token_count=0, + mismatch_count=0, + contract_fingerprint=contract_fingerprint, + fixed_threshold=fixed_threshold, + token_artifact=token_artifact, + ) + + mismatch_count = int(mismatch_mask.sum().item()) + passed = mismatch_count == 0 + return AlignmentResult( + case_id=rollout.case_id, + attempt_id=rollout.attempt_id, + status=AlignmentStatus.PASS if passed else AlignmentStatus.FAIL, + comparable=True, + passed=passed, + active_token_count=active_token_count, + mismatch_count=mismatch_count, + contract_fingerprint=contract_fingerprint, + fixed_threshold=fixed_threshold, + diagnostics=_diagnostics( + rollout_logprobs, + training_logprobs, + active_mask, + absolute_diff, + mismatch_count, + ), + token_artifact=token_artifact, + ) + + +def compare_score_artifacts( + rollout: ScoreArtifact, + training: ScoreArtifact, +) -> AlignmentResult: + """Convenience wrapper whose API deliberately exposes no threshold override.""" + + return FixedThresholdComparator().compare(rollout, training) + + +def _resolve_fixed_threshold( + rollout: ScoreArtifact, + training: ScoreArtifact, +) -> tuple[float | None, str | None]: + """Resolve one WS1 threshold, rejecting any mixed-dtype ambiguity.""" + + try: + rollout_threshold = resolve_logprob_threshold(rollout.scorer.dtype) + training_threshold = resolve_logprob_threshold(training.scorer.dtype) + except ValueError as exc: + return None, f"fixed WS1 threshold is unavailable: {exc}" + if rollout_threshold != training_threshold: + return ( + None, + "fixed WS1 threshold is ambiguous for scorer dtypes " + f"rollout={rollout.scorer.dtype!r}, training={training.scorer.dtype!r}", + ) + return rollout_threshold, None + + +def _artifact_errors(rollout: ScoreArtifact, training: ScoreArtifact) -> tuple[str, ...]: + errors: list[str] = [] + if rollout.side is not ScoreSide.ROLLOUT: + errors.append("first artifact side must be rollout") + if training.side is not ScoreSide.TRAINING: + errors.append("second artifact side must be training") + if rollout.case_id != training.case_id: + errors.append("case_id") + if rollout.attempt_id != training.attempt_id: + errors.append("attempt_id") + if rollout.selected_logprobs.shape != training.selected_logprobs.shape: + errors.append("selected_logprobs shape") + for label, artifact in (("rollout", rollout), ("training", training)): + expected_dtype = _score_dtype(artifact.scorer.dtype) + if not artifact.selected_logprobs.is_floating_point(): + errors.append(f"{label}.selected_logprobs must be floating point") + elif expected_dtype is None: + errors.append(f"{label}.scorer dtype is unsupported") + elif artifact.selected_logprobs.dtype != expected_dtype: + errors.append( + f"{label}.selected_logprobs dtype does not match scorer dtype " + f"({artifact.selected_logprobs.dtype} != {expected_dtype})" + ) + return tuple(errors) + + +def _score_dtype(value: str) -> torch.dtype | None: + normalized = str(value).strip().lower().removeprefix("torch.") + return { + "float32": torch.float32, + "fp32": torch.float32, + "float16": torch.float16, + "fp16": torch.float16, + "bfloat16": torch.bfloat16, + "bf16": torch.bfloat16, + "float64": torch.float64, + }.get(normalized) + + +def _artifact_identity_errors( + rollout: ScoreArtifact, + training: ScoreArtifact, +) -> tuple[str, ...]: + errors: list[str] = [] + rollout_identity_mask = _identity_mask(rollout.identity) + training_identity_mask = _identity_mask(training.identity) + rollout_mask = rollout.active_mask.detach().cpu().to(dtype=torch.bool) + training_mask = training.active_mask.detach().cpu().to(dtype=torch.bool) + if rollout_mask.shape != rollout_identity_mask.shape or not torch.equal( + rollout_mask, rollout_identity_mask + ): + errors.append("rollout.active_mask") + if training_mask.shape != training_identity_mask.shape or not torch.equal( + training_mask, training_identity_mask + ): + errors.append("training.active_mask") + if rollout_mask.shape != training_mask.shape or not torch.equal(rollout_mask, training_mask): + errors.append("active_mask") + return tuple(errors) + + +def _identity_mask(identity: SemanticIdentitySpec) -> torch.Tensor: + return torch.tensor(identity.active_mask, dtype=torch.bool) + + +def _diagnostics( + rollout_logprobs: torch.Tensor, + training_logprobs: torch.Tensor, + active_mask: torch.Tensor, + absolute_diff: torch.Tensor, + mismatch_count: int, +) -> dict[str, Any]: + active_diff = absolute_diff[active_mask].float() + delta = (training_logprobs[active_mask] - rollout_logprobs[active_mask]).float() + worst_active_index = int(torch.argmax(active_diff).item()) + active_coordinates = torch.nonzero(active_mask, as_tuple=False) + worst_coordinate = tuple(int(item) for item in active_coordinates[worst_active_index].tolist()) + approximate_kl = torch.exp(delta.double()) - delta.double() - 1.0 + approximate_kl_mean = _finite_float_or_none(approximate_kl.mean()) + active_count = int(active_diff.numel()) + return { + "mean_abs_diff": _finite_float_or_none(active_diff.mean()), + "p95_abs_diff": _finite_float_or_none(torch.quantile(active_diff, 0.95)), + "p99_abs_diff": _finite_float_or_none(torch.quantile(active_diff, 0.99)), + "max_abs_diff": _finite_float_or_none(active_diff.max()), + "mismatch_ratio": mismatch_count / active_count, + "approximate_kl_mean": approximate_kl_mean, + "approximate_kl_finite": approximate_kl_mean is not None, + "worst_token_index": worst_coordinate, + } + + +def _finite_float_or_none(value: torch.Tensor) -> float | None: + result = float(value.item()) + return result if math.isfinite(result) else None + + +__all__ = [ + "FixedThresholdComparator", + "compare_score_artifacts", + "recompute_mismatch_mask", + "semantic_identity_errors", +] diff --git a/rl_engine/alignment/cross_config/config.py b/rl_engine/alignment/cross_config/config.py new file mode 100644 index 00000000..e02c2198 --- /dev/null +++ b/rl_engine/alignment/cross_config/config.py @@ -0,0 +1,424 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Strict, dependency-free experiment configuration.""" + +from __future__ import annotations + +import hashlib +import json +from dataclasses import dataclass, field +from dataclasses import fields as dataclass_fields +from pathlib import Path +from types import MappingProxyType +from typing import TYPE_CHECKING, Any, Mapping + +from rl_engine.alignment.cross_config._json import strict_json_loads +from rl_engine.alignment.cross_config.planner import normalize_backend_id +from rl_engine.alignment.cross_config.schema import ( + ExperimentCase, + ExperimentDefinition, + InterventionSpec, + PlanningStrategy, + SemanticIdentitySpec, +) + +if TYPE_CHECKING: + from rl_engine.alignment.cross_config.planner import ExperimentPlan + + +CONFIG_SCHEMA_VERSION = "cross_config.experiment_config.v1" +_FORBIDDEN_THRESHOLD_KEYS = frozenset({"threshold", "fixed_threshold", "tolerance", "atol", "rtol"}) +_TOP_LEVEL_KEYS = frozenset( + { + "schema_version", + "experiment_id", + "scenario_id", + "contract_source", + "contract_version", + "strategy", + "strict_fallback", + "identity", + "baseline", + "interventions", + "pairwise_paths", + "operators", + "scenario", + } +) +_IDENTITY_KEYS = frozenset( + item.name for item in dataclass_fields(SemanticIdentitySpec) if item.name != "schema_version" +) +_INTERVENTION_KEYS = frozenset({"path", "values"}) +_OPERATOR_NAMES = frozenset({"selected_logprob"}) +_OPERATOR_TARGETS = frozenset({"rollout", "training"}) +_OPERATOR_BINDING_KEYS = frozenset({"backend", "options"}) + + +@dataclass(frozen=True) +class OperatorSelection: + """Concrete selected-logprob implementation requested for each scorer side. + + ``logp.backend`` remains the concise both-sides shortcut. This explicit form + is needed only when rollout and training intentionally use different + implementations. + """ + + rollout_backend: str + training_backend: str + rollout_options: Mapping[str, Any] = field(default_factory=dict) + training_options: Mapping[str, Any] = field(default_factory=dict) + + def __post_init__(self) -> None: + object.__setattr__(self, "rollout_backend", normalize_backend_id(self.rollout_backend)) + object.__setattr__( + self, + "training_backend", + normalize_backend_id(self.training_backend), + ) + object.__setattr__(self, "rollout_options", _freeze_mapping(self.rollout_options)) + object.__setattr__(self, "training_options", _freeze_mapping(self.training_options)) + + def backend_for(self, target: str) -> str: + if target == "rollout": + return self.rollout_backend + if target == "training": + return self.training_backend + raise ValueError("operator target must be 'rollout' or 'training'") + + def options_for(self, target: str) -> Mapping[str, Any]: + if target == "rollout": + return self.rollout_options + if target == "training": + return self.training_options + raise ValueError("operator target must be 'rollout' or 'training'") + + def to_dict(self) -> dict[str, Any]: + return { + "selected_logprob": { + "rollout": { + "backend": self.rollout_backend, + "options": _plain_value(self.rollout_options), + }, + "training": { + "backend": self.training_backend, + "options": _plain_value(self.training_options), + }, + } + } + + +@dataclass(frozen=True) +class ExperimentConfig: + """Loaded experiment plus optional target-specific operator selection.""" + + definition: ExperimentDefinition + source_path: Path + operators: OperatorSelection | None = None + schema_version: str = CONFIG_SCHEMA_VERSION + + def to_dict(self) -> dict[str, Any]: + """Return the normalized, portable experiment-config representation.""" + + payload = self.definition.to_dict() + payload["schema_version"] = self.schema_version + if self.operators is not None: + payload["operators"] = self.operators.to_dict() + return payload + + def plan(self) -> ExperimentPlan: + """Build the deterministic plan without importing a runtime backend.""" + + from rl_engine.alignment.cross_config.planner import Planner + + return Planner().plan(self.definition) + + def operators_for(self, case: ExperimentCase) -> OperatorSelection: + """Resolve the concise ``logp.backend`` shortcut for one planned case.""" + + backend = _case_logp_backend(case) + if self.operators is None: + return OperatorSelection(backend, backend) + if self.operators.rollout_backend != backend: + raise ValueError( + "operators.selected_logprob.rollout must match the planned " + f"logp.backend: {self.operators.rollout_backend!r} != {backend!r}" + ) + return self.operators + + +def load_config(path: str | Path) -> ExperimentConfig: + """Load one versioned JSON experiment with no threshold override surface.""" + + source = Path(path) + try: + raw = strict_json_loads(source.read_text(encoding="utf-8")) + except (OSError, ValueError) as exc: + raise ValueError(f"failed to load cross-configuration config {source}: {exc}") from exc + if not isinstance(raw, dict): + raise ValueError("cross-configuration config must contain a JSON object") + _reject_unknown_keys(raw, _TOP_LEVEL_KEYS, "config") + _reject_threshold_keys(raw) + if raw.get("schema_version") != CONFIG_SCHEMA_VERSION: + raise ValueError( + f"unsupported cross-configuration config schema {raw.get('schema_version')!r}; " + f"expected {CONFIG_SCHEMA_VERSION!r}" + ) + + identity_raw = _required_mapping(raw, "identity") + _reject_unknown_keys(identity_raw, _IDENTITY_KEYS, "identity") + scenario = _optional_mapping(raw, "scenario") + _reject_scenario_controls(scenario) + + interventions_raw = raw.get("interventions", []) + if not isinstance(interventions_raw, list): + raise ValueError("interventions must be a list") + interventions = tuple(_load_intervention(item) for item in interventions_raw) + + pairwise_raw = raw.get("pairwise_paths", []) + if not isinstance(pairwise_raw, list): + raise ValueError("pairwise_paths must be a list") + pairwise_paths = tuple(_load_pair(item) for item in pairwise_raw) + + strict_fallback = raw.get("strict_fallback", True) + if not isinstance(strict_fallback, bool): + raise ValueError("strict_fallback must be a JSON boolean") + + definition = ExperimentDefinition( + experiment_id=_required_string(raw, "experiment_id"), + scenario_id=_required_string(raw, "scenario_id"), + identity=SemanticIdentitySpec(**identity_raw), + baseline=_required_mapping(raw, "baseline"), + interventions=interventions, + scenario=scenario, + strategy=PlanningStrategy(raw.get("strategy", "one_at_a_time")), + strict_fallback=strict_fallback, + pairwise_paths=pairwise_paths, + contract_source=raw.get("contract_source", "ws1"), + contract_version=raw.get("contract_version", "current"), + ) + operators = _load_operators(raw.get("operators")) + if operators is not None: + if any(item.path == "logp.backend" for item in interventions): + raise ValueError( + "explicit operators cannot be combined with logp.backend interventions; " + "use the shortcut or one fixed target mapping" + ) + baseline_backend = _definition_logp_backend(definition) + if operators.rollout_backend != baseline_backend: + raise ValueError( + "operators.selected_logprob.rollout must match baseline logp.backend: " + f"{operators.rollout_backend!r} != {baseline_backend!r}" + ) + + return ExperimentConfig( + definition=definition, + operators=operators, + source_path=source, + ) + + +def bind_operator_selection( + case: ExperimentCase, + selection: OperatorSelection, +) -> ExperimentCase: + """Bind target-specific operators into the execution identity. + + Planning remains semantic-operator agnostic; the immutable binding extends + the case and resume key before any runtime is created. + """ + + requested_backend = _case_logp_backend(case) + if selection.rollout_backend != requested_backend: + raise ValueError( + "rollout operator must match the planned logp.backend: " + f"{selection.rollout_backend!r} != {requested_backend!r}" + ) + binding = selection.to_dict() + payload = { + "base_case_id": case.case_id, + "base_scenario_fingerprint": case.scenario_fingerprint, + "operators": binding, + } + serialized = json.dumps(payload, sort_keys=True, separators=(",", ":")).encode("utf-8") + operator_fingerprint = hashlib.sha256(serialized).hexdigest() + case_hash = hashlib.sha256( + json.dumps( + {"base_case_id": case.case_id, "operator_fingerprint": operator_fingerprint}, + sort_keys=True, + separators=(",", ":"), + ).encode("utf-8") + ).hexdigest()[:24] + scenario_fingerprint = hashlib.sha256( + f"{case.scenario_fingerprint}:{operator_fingerprint}".encode("utf-8") + ).hexdigest() + return ExperimentCase( + case_id=f"cross-config-{case_hash}", + experiment_id=case.experiment_id, + scenario_id=case.scenario_id, + identity=case.identity, + requested=case.requested, + execution_binding={"operators": binding}, + changed_paths=case.changed_paths, + contract_fingerprint=case.contract_fingerprint, + scenario_fingerprint=scenario_fingerprint, + ) + + +def _load_intervention(value: Any) -> InterventionSpec: + if not isinstance(value, Mapping): + raise ValueError("each intervention must be an object") + _reject_unknown_keys(value, _INTERVENTION_KEYS, "intervention") + values = value.get("values") + if not isinstance(values, list): + raise ValueError("intervention values must be a list") + return InterventionSpec(path=_required_string(value, "path"), values=tuple(values)) + + +def _load_pair(value: Any) -> tuple[str, str]: + if ( + not isinstance(value, list) + or len(value) != 2 + or not all(isinstance(item, str) for item in value) + ): + raise ValueError("each pairwise_paths entry must contain exactly two string paths") + return value[0], value[1] + + +def _load_operators(value: Any) -> OperatorSelection | None: + if value is None: + return None + if not isinstance(value, Mapping): + raise ValueError("operators must be an object") + _reject_unknown_keys(value, _OPERATOR_NAMES, "operators") + selected = value.get("selected_logprob") + if not isinstance(selected, Mapping): + raise ValueError("operators.selected_logprob must be an object") + _reject_unknown_keys(selected, _OPERATOR_TARGETS, "operators.selected_logprob") + rollout_backend, rollout_options = _load_operator_binding(selected, "rollout") + training_backend, training_options = _load_operator_binding(selected, "training") + return OperatorSelection( + rollout_backend=rollout_backend, + training_backend=training_backend, + rollout_options=rollout_options, + training_options=training_options, + ) + + +def _load_operator_binding( + value: Mapping[str, Any], + target: str, +) -> tuple[str, Mapping[str, Any]]: + binding = value.get(target) + if isinstance(binding, str): + if not binding.strip(): + raise ValueError(f"operators.selected_logprob.{target} must not be empty") + return binding, {} + if not isinstance(binding, Mapping): + raise ValueError(f"operators.selected_logprob.{target} must be a backend string or object") + _reject_unknown_keys(binding, _OPERATOR_BINDING_KEYS, f"{target} operator binding") + return _required_string(binding, "backend"), _optional_mapping(binding, "options") + + +def _case_logp_backend(case: ExperimentCase) -> str: + logp = case.requested.get("logp") + backend = logp.get("backend") if isinstance(logp, Mapping) else None + if not isinstance(backend, str) or not backend: + raise ValueError("planned cases must contain a non-empty string logp.backend") + return normalize_backend_id(backend) + + +def _definition_logp_backend(definition: ExperimentDefinition) -> str: + logp = definition.baseline.get("logp") + backend = logp.get("backend") if isinstance(logp, Mapping) else None + if not isinstance(backend, str) or not backend: + raise ValueError("baseline must contain a non-empty string logp.backend") + return normalize_backend_id(backend) + + +def _reject_scenario_controls(scenario: Mapping[str, Any]) -> None: + behavior_keys = sorted( + set(scenario).intersection( + {"execution", "plan_only", "operator_cases", "expected_status", "allow_smoke_operators"} + ) + ) + if behavior_keys: + raise ValueError( + "scenario is metadata only; move execution and operator policy to the CLI/config: " + f"{behavior_keys}" + ) + + +def _reject_threshold_keys(value: Any, prefix: str = "") -> None: + if isinstance(value, Mapping): + for key, child in value.items(): + normalized = str(key).strip().lower() + path = f"{prefix}.{key}" if prefix else str(key) + if normalized in _FORBIDDEN_THRESHOLD_KEYS: + raise ValueError( + f"{path} is forbidden: the fixed numerical-contract threshold is imported" + ) + _reject_threshold_keys(child, path) + elif isinstance(value, list): + for index, child in enumerate(value): + _reject_threshold_keys(child, f"{prefix}[{index}]") + + +def _reject_unknown_keys( + value: Mapping[str, Any], + allowed: frozenset[str], + label: str, +) -> None: + unknown = sorted(set(value).difference(allowed)) + if unknown: + raise ValueError(f"unknown {label} keys: {unknown}") + + +def _required_mapping(value: Mapping[str, Any], key: str) -> dict[str, Any]: + child = value.get(key) + if not isinstance(child, Mapping): + raise ValueError(f"{key} must be an object") + return dict(child) + + +def _optional_mapping(value: Mapping[str, Any], key: str) -> dict[str, Any]: + child = value.get(key, {}) + if not isinstance(child, Mapping): + raise ValueError(f"{key} must be an object") + return dict(child) + + +def _required_string(value: Mapping[str, Any], key: str) -> str: + child = value.get(key) + if not isinstance(child, str) or not child.strip(): + raise ValueError(f"{key} must be a non-empty string") + return child.strip() + + +def _freeze_mapping(value: Mapping[str, Any]) -> Mapping[str, Any]: + return MappingProxyType({str(key): _freeze_value(child) for key, child in value.items()}) + + +def _freeze_value(value: Any) -> Any: + if isinstance(value, Mapping): + return _freeze_mapping(value) + if isinstance(value, list): + return tuple(_freeze_value(child) for child in value) + return value + + +def _plain_value(value: Any) -> Any: + if isinstance(value, Mapping): + return {str(key): _plain_value(child) for key, child in value.items()} + if isinstance(value, (tuple, list)): + return [_plain_value(child) for child in value] + return value + + +__all__ = [ + "CONFIG_SCHEMA_VERSION", + "ExperimentConfig", + "OperatorSelection", + "bind_operator_selection", + "load_config", +] diff --git a/rl_engine/alignment/cross_config/debug_matrix.py b/rl_engine/alignment/cross_config/debug_matrix.py new file mode 100644 index 00000000..32a49470 --- /dev/null +++ b/rl_engine/alignment/cross_config/debug_matrix.py @@ -0,0 +1,183 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Compact, module-level mismatch axes for post-training drift triage. + +The manifest is deliberately a reporting contract. It names the first +diagnostic probe for each semantic operator without turning the probes into +runtime configuration or expanding a Cartesian product of settings. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from types import MappingProxyType +from typing import Any + +from rl_engine.kernels.ops.pytorch.attention.debug_matrix import ( + ATTENTION_DEBUG_MATRIX, + ATTENTION_DEBUG_MATRIX_SCHEMA_VERSION, +) +from rl_engine.kernels.ops.pytorch.attention.debug_taxonomy import ATTENTION_DEBUG_AXES + +DEBUG_MATRIX_SCHEMA_VERSION = "rlkernel.debug_matrix.v1" + + +@dataclass(frozen=True) +class ModuleDebugAxis: + """One first-line axis exposed by a semantic operator.""" + + module: str + axis_id: str + label: str + representative_probe: str + kind: str = "diagnostic" + + def __post_init__(self) -> None: + for field in ("module", "axis_id", "label", "representative_probe"): + value = getattr(self, field) + if not isinstance(value, str) or not value.strip(): + raise ValueError(f"module debug {field} must be a non-empty string") + if self.kind not in {"diagnostic", "gate"}: + raise ValueError(f"unknown module debug axis kind {self.kind!r}") + + def to_dict(self, *, row_id: str) -> dict[str, str]: + return { + "row": row_id, + "id": self.axis_id, + "label": self.label, + "representative_probe": self.representative_probe, + "kind": self.kind, + } + + +def _attention_axes() -> tuple[ModuleDebugAxis, ...]: + return tuple( + ModuleDebugAxis( + module="attention", + axis_id=axis.axis_id, + label=axis.label, + representative_probe=axis.representative_subprobe, + kind="gate" if axis.axis_id == "topology_head_ownership" else "diagnostic", + ) + for axis in ATTENTION_DEBUG_AXES + ) + + +MODULE_DEBUG_AXES: MappingProxyType[str, tuple[ModuleDebugAxis, ...]] = MappingProxyType( + { + "attention": _attention_axes(), + "ffn": ( + ModuleDebugAxis( + "ffn", + "weight_shard_ownership", + "Weight shard / TP ownership", + "tp_weight_ownership", + "gate", + ), + ModuleDebugAxis( + "ffn", + "swiglu_rounding", + "SwiGLU intermediate rounding", + "swiglu_one_round", + ), + ModuleDebugAxis( + "ffn", + "gemm_reduction", + "GEMM K-reduction / Split-K policy", + "k_reduction_split_k", + ), + ModuleDebugAxis( + "ffn", + "token_collective", + "Token gather / reduce-scatter", + "sequence_parallel", + ), + ), + "logp": ( + ModuleDebugAxis( + "logp", + "vocab_shard_ownership", + "Vocabulary shard / TP ownership", + "vocab_shard_bounds", + "gate", + ), + ModuleDebugAxis( + "logp", + "selected_token_identity", + "Selected-token and active-mask identity", + "selected_token_active_mask", + "gate", + ), + ModuleDebugAxis( + "logp", + "vocab_lse_reduction", + "Vocabulary LSE tile / merge policy", + "vocab_tile_merge", + ), + ), + } +) + + +_MODULE_DEBUG_AXIS_BY_ID = MappingProxyType( + {(axis.module, axis.axis_id): axis for axes in MODULE_DEBUG_AXES.values() for axis in axes} +) + + +def _module_rows(module: str) -> dict[str, Any]: + axes = MODULE_DEBUG_AXES[module] + baseline = {"attention": "A0", "ffn": "F0", "logp": "L0"}[module] + controls = { + "attention": ["C0", "C1", "C2"], + "ffn": ["FC0"], + "logp": ["LC0"], + }[module] + rows = [baseline] + [f"{module[0].upper()}{index}" for index in range(1, len(axes) + 1)] + return { + "baseline_row": baseline, + "rows": rows, + "invariant_controls": controls, + "axes": [axis.to_dict(row_id=rows[index]) for index, axis in enumerate(axes, start=1)], + } + + +def module_debug_matrix() -> dict[str, Any]: + """Return the portable matrix manifest shared by all three operators.""" + + return { + "schema_version": DEBUG_MATRIX_SCHEMA_VERSION, + "method": "fixed_replay_one_at_a_time", + "cartesian_product": False, + "comparison_edges": [ + "train_vs_rollout_prefill", + "rollout_prefill_vs_decode", + ], + "replay_identity": ( + "same checkpoint, token IDs, selected-token IDs, masks, positions, " + "cache metadata, and pre-update model state" + ), + "modules": {module: _module_rows(module) for module in MODULE_DEBUG_AXES}, + "attention_compatibility": { + "schema_version": ATTENTION_DEBUG_MATRIX_SCHEMA_VERSION, + "rows": [row.to_dict() for row in ATTENTION_DEBUG_MATRIX], + }, + } + + +def module_debug_axis(module: str, axis_id: str) -> ModuleDebugAxis: + """Look up one stable axis by module and identifier.""" + + try: + return _MODULE_DEBUG_AXIS_BY_ID[(module.strip(), axis_id.strip())] + except (AttributeError, KeyError) as exc: + raise ValueError(f"unknown module debug axis {module!r}/{axis_id!r}") from exc + + +__all__ = [ + "DEBUG_MATRIX_SCHEMA_VERSION", + "MODULE_DEBUG_AXES", + "ModuleDebugAxis", + "module_debug_axis", + "module_debug_matrix", +] diff --git a/rl_engine/alignment/cross_config/determinism.py b/rl_engine/alignment/cross_config/determinism.py new file mode 100644 index 00000000..be81654e --- /dev/null +++ b/rl_engine/alignment/cross_config/determinism.py @@ -0,0 +1,305 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Cross-side determinism probing for the Megatron + vLLM cross-config target. + +Both frameworks ship a "make this deterministic" switch, but they mean different +things by it, and neither knows the other exists: + +``Megatron`` ``ModelParallelConfig.deterministic_mode`` + Asserts ``NCCL_ALGO`` is one of five values, forbids FlashAttention and fused + cross-entropy, calls ``torch.use_deterministic_algorithms(True)``, and requires + ``NVTE_ALLOW_NONDETERMINISTIC_ALGO == 0``. It does **not** touch TF32, BF16 + reduced-precision reduction, cuBLAS workspace, NCCL protocol, or NCCL channel + counts. + +``vLLM`` ``VLLM_BATCH_INVARIANT`` + Replaces ``aten::mm/addmm/matmul/linear/bmm``, ``log_softmax``/``softmax``, + ``mean.dim`` and ``rms_norm`` with Triton kernels, disables TF32 and BF16/FP16 + reduced-precision reduction, pins cuBLAS workspace and the BLAS library, and + hard-sets ten NCCL environment variables. + +So a run can have both switches on and still be comparing two different notions of +determinism. This module makes that difference explicit and, where it changes the +numerics, blocking. It never imports Megatron or vLLM: probes are built from plain +mappings so the logic is testable on any machine. +""" + +from __future__ import annotations + +import hashlib +import json +from collections.abc import Mapping +from dataclasses import dataclass, field +from typing import Any, Optional + +from rl_engine.alignment.cross_config.attention_binding import ( + BindingErrorCode, + BindingIssue, + BindingTier, +) + +__all__ = [ + "COMPARED_NCCL_KEYS", + "DeterminismProbe", + "DeterminismReport", + "compare_determinism", + "megatron_probe_from_config", + "vllm_probe_from_env", +] + + +#: Environment keys whose value can change a reduction result. Compared across +#: sides; a difference is reported, and a difference in the *arithmetic* subset is +#: blocking. Ordering is fixed so the fingerprint is stable. +COMPARED_NCCL_KEYS: tuple[str, ...] = ( + "NCCL_ALGO", + "NCCL_PROTO", + "NCCL_MIN_NCHANNELS", + "NCCL_MAX_NCHANNELS", + "NCCL_NTHREADS", + "NCCL_SOCKET_NTHREADS", + "NCCL_COLLNET_ENABLE", + "NCCL_NVLS_ENABLE", + "NCCL_P2P_NET_DISABLE", + "NCCL_LAUNCH_MODE", + "CUBLAS_WORKSPACE_CONFIG", +) + + +#: The subset above that changes arithmetic rather than only scheduling. A mismatch +#: here fails the binding closed; a mismatch in the remainder is recorded only. +_ARITHMETIC_NCCL_KEYS: frozenset[str] = frozenset( + {"NCCL_ALGO", "NCCL_PROTO", "CUBLAS_WORKSPACE_CONFIG"} +) + + +@dataclass(frozen=True) +class DeterminismProbe: + """What one side actually has switched on. + + ``tf32_disabled`` and ``bf16_reduced_precision_reduction`` are tri-state on + purpose: ``None`` means "the framework does not manage this", which is exactly + Megatron's situation and is itself the finding. + """ + + side: str + framework: str + mode_flag: str + enabled: bool + env: Mapping[str, Any] = field(default_factory=dict) + tf32_disabled: Optional[bool] = None + bf16_reduced_precision_reduction: Optional[bool] = None + forbids_flash_attention: Optional[bool] = None + evidence: Mapping[str, Any] = field(default_factory=dict) + schema_version: str = "cross_config.determinism_probe.v1" + + def __post_init__(self) -> None: + if self.side not in ("rollout", "training"): + raise ValueError("side must be 'rollout' or 'training'") + if not self.framework: + raise ValueError("framework must not be empty") + object.__setattr__(self, "env", dict(self.env)) + object.__setattr__(self, "evidence", dict(self.evidence)) + + @property + def env_fingerprint(self) -> str: + payload = {key: self.env.get(key) for key in COMPARED_NCCL_KEYS} + encoded = json.dumps(payload, sort_keys=True, separators=(",", ":"), default=str) + return hashlib.sha256(encoded.encode("utf-8")).hexdigest() + + def to_dict(self) -> dict[str, Any]: + return { + "schema_version": self.schema_version, + "side": self.side, + "framework": self.framework, + "mode_flag": self.mode_flag, + "enabled": self.enabled, + "env": {key: self.env.get(key) for key in COMPARED_NCCL_KEYS}, + "env_fingerprint": self.env_fingerprint, + "tf32_disabled": self.tf32_disabled, + "bf16_reduced_precision_reduction": self.bf16_reduced_precision_reduction, + "forbids_flash_attention": self.forbids_flash_attention, + "evidence": dict(self.evidence), + } + + +@dataclass(frozen=True) +class DeterminismReport: + """Cross-side comparison result.""" + + rollout: DeterminismProbe + training: DeterminismProbe + issues: tuple[BindingIssue, ...] = () + differences: Mapping[str, Mapping[str, Any]] = field(default_factory=dict) + schema_version: str = "cross_config.determinism_report.v1" + + @property + def compatible(self) -> bool: + return not self.issues + + def to_dict(self) -> dict[str, Any]: + return { + "schema_version": self.schema_version, + "compatible": self.compatible, + "rollout": self.rollout.to_dict(), + "training": self.training.to_dict(), + "issues": [issue.to_dict() for issue in self.issues], + "differences": {key: dict(value) for key, value in self.differences.items()}, + } + + +def megatron_probe_from_config( + config: Any, + env: Optional[Mapping[str, str]] = None, +) -> DeterminismProbe: + """Build a training-side probe from a Megatron config object. + + ``config`` is duck-typed (anything exposing ``deterministic_mode`` and + optionally ``attention_backend`` / ``cross_entropy_loss_fusion``) so this works + against a real ``ModelParallelConfig``, a test double, or a plain namespace, + and so importing this module never requires Megatron. + + ``tf32_disabled`` and ``bf16_reduced_precision_reduction`` are reported as + ``None`` because Megatron does not manage them -- a ``grep`` for ``allow_tf32`` + and ``fp32_precision`` across ``megatron/`` returns nothing. That asymmetry + against vLLM is the point of :func:`compare_determinism`. + """ + + environ = dict(env or {}) + enabled = bool(getattr(config, "deterministic_mode", False)) + return DeterminismProbe( + side="training", + framework="megatron", + mode_flag="deterministic_mode", + enabled=enabled, + env={key: environ.get(key) for key in COMPARED_NCCL_KEYS}, + tf32_disabled=None, + bf16_reduced_precision_reduction=None, + forbids_flash_attention=enabled, + evidence={ + "nvte_allow_nondeterministic_algo": environ.get("NVTE_ALLOW_NONDETERMINISTIC_ALGO"), + "cross_entropy_loss_fusion": getattr(config, "cross_entropy_loss_fusion", None), + "attention_backend": _enum_value(getattr(config, "attention_backend", None)), + "tensor_model_parallel_size": getattr(config, "tensor_model_parallel_size", None), + "context_parallel_size": getattr(config, "context_parallel_size", None), + "sequence_parallel": getattr(config, "sequence_parallel", None), + "manages_tf32": False, + "manages_bf16_reduced_precision_reduction": False, + }, + ) + + +def vllm_probe_from_env( + env: Mapping[str, str], + *, + model_config: Any = None, +) -> DeterminismProbe: + """Build a rollout-side probe from the vLLM process environment. + + ``VLLM_BATCH_INVARIANT`` is read from ``env`` rather than ``vllm.envs`` so the + probe can be constructed from a remote worker's reported environment, which is + how vime's Ray actors expose it. + """ + + enabled = str(env.get("VLLM_BATCH_INVARIANT", "0")).strip() in ("1", "true", "True") + return DeterminismProbe( + side="rollout", + framework="vllm", + mode_flag="VLLM_BATCH_INVARIANT", + enabled=enabled, + env={key: env.get(key) for key in COMPARED_NCCL_KEYS}, + # vLLM sets both to "ieee"/disabled inside init_batch_invariance(). + tf32_disabled=enabled or None, + bf16_reduced_precision_reduction=(False if enabled else None), + forbids_flash_attention=False, + evidence={ + "vllm_allreduce_use_symm_mem": env.get("VLLM_ALLREDUCE_USE_SYMM_MEM"), + "vllm_use_aot_compile": env.get("VLLM_USE_AOT_COMPILE"), + "enforce_eager": getattr(model_config, "enforce_eager", None), + "disable_cascade_attn": getattr(model_config, "disable_cascade_attn", None), + "quantization": getattr(model_config, "quantization", None), + "manages_tf32": True, + "manages_bf16_reduced_precision_reduction": True, + }, + ) + + +def _enum_value(value: Any) -> Any: + return getattr(value, "value", value) + + +def compare_determinism( + *, + rollout: DeterminismProbe, + training: DeterminismProbe, +) -> DeterminismReport: + """Compare two probes and produce blocking issues plus recorded differences.""" + + if rollout.side != "rollout" or training.side != "training": + raise ValueError("compare_determinism expects one rollout probe and one training probe") + + issues: list[BindingIssue] = [] + differences: dict[str, dict[str, Any]] = {} + + for probe in (rollout, training): + if not probe.enabled: + issues.append( + BindingIssue( + code=BindingErrorCode.DETERMINISM_INCOMPATIBLE, + tier=BindingTier.SEMANTIC, + field=f"{probe.side}.{probe.mode_flag}", + rollout=rollout.enabled, + training=training.enabled, + message=( + f"{probe.framework} {probe.mode_flag} is not enabled; the " + f"{probe.side} side is not batch-invariant and cannot anchor a " + "cross-config comparison" + ), + ) + ) + + for key in COMPARED_NCCL_KEYS: + rollout_value = rollout.env.get(key) + training_value = training.env.get(key) + if rollout_value == training_value: + continue + differences[key] = {"rollout": rollout_value, "training": training_value} + if key in _ARITHMETIC_NCCL_KEYS: + issues.append( + BindingIssue( + code=BindingErrorCode.DETERMINISM_INCOMPATIBLE, + tier=BindingTier.SEMANTIC, + field=f"env.{key}", + rollout=rollout_value, + training=training_value, + message=( + f"{key} differs between sides; the two sides would reduce with " + "different arithmetic and the resulting drift is not attributable" + ), + ) + ) + + # Megatron reports None for these because it does not manage them at all. That is + # recorded rather than blocking: under a pure BF16 GEMM path TF32 does not fire, and + # forcing Megatron to manage it is out of scope for this PR. It is surfaced so the + # asymmetry appears in every artifact instead of being invisible. + for name in ("tf32_disabled", "bf16_reduced_precision_reduction"): + rollout_value = getattr(rollout, name) + training_value = getattr(training, name) + if rollout_value != training_value: + differences[name] = { + "rollout": rollout_value, + "training": training_value, + "note": ( + "megatron does not manage this setting; vllm sets it inside " + "init_batch_invariance()" + ), + } + + return DeterminismReport( + rollout=rollout, + training=training, + issues=tuple(issues), + differences=differences, + ) diff --git a/rl_engine/alignment/cross_config/drift_report.py b/rl_engine/alignment/cross_config/drift_report.py new file mode 100644 index 00000000..7f11227a --- /dev/null +++ b/rl_engine/alignment/cross_config/drift_report.py @@ -0,0 +1,1309 @@ +# flake8: noqa: E501 +"""Build self-contained profiler-style reports from cross-config artifacts. + +The reporter is deliberately downstream of execution. It validates a sealed +attempt directory, then renders the recorded comparison, token deltas, actual +operator provenance, and materialized topology without changing either the +training or rollout path. +""" + +from __future__ import annotations + +import html +import io +import json +import zipfile +from collections.abc import Mapping +from pathlib import Path +from typing import Any + +import torch + +from rl_engine.alignment.cross_config.artifacts import ArtifactStore + +REPORT_SCHEMA_VERSION = 2 +_STATUS_LABELS = {"pass": "PASS", "warning": "WARN", "failure": "FAIL", "info": "INFO"} +_STATUS_COLORS = { + "pass": "#39d98a", + "warning": "#f4b942", + "failure": "#ff5c77", + "info": "#70a7ff", +} +_REPORT_IMAGE_COLORS = { + "background": "#eef0f2", + "panel": "#ffffff", + "panel_alt": "#f7f8f9", + "track": "#c8cdd2", + "line": "#cbd1d7", + "grid": "#e5e8eb", + "text": "#20252b", + "muted": "#66707a", + "green": "#43a64b", + "yellow": "#d79700", + "red": "#d64242", + "blue": "#3677c8", + "purple": "#7659bb", + "pass": "#43a64b", + "warning": "#d79700", + "failure": "#d64242", + "info": "#3677c8", +} + + +def build_cross_config_attempt_report( + attempt_dir: str | Path, + *, + title: str | None = None, + validate: bool = True, +) -> dict[str, Any]: + """Render one immutable cross-configuration attempt into a drift report. + + ``attempt_dir`` must be a directory emitted by :class:`PairedRunner`. The + report never derives a numerical claim from an incomplete attempt. A + caller can set ``validate=False`` only when it has already validated the + same sealed directory in the current process. + """ + + directory = Path(attempt_dir) + store = ArtifactStore(directory.parent) + if validate: + store.validate_completed_attempt(directory) + + comparison = _read_attempt_json(directory / "comparison.json") + actual = _read_attempt_json(directory / "actual.json") + materialized = _read_attempt_json(directory / "materialized.json") + identity = _read_attempt_json(directory / "identity.json") + token_payload = store.load_tensor_bundle(directory / "token_diffs.pt") + tensors = token_payload["tensors"] + active_mask = tensors["active_mask"].to(dtype=torch.bool) + absolute_diff = tensors["absolute_diff"].to(dtype=torch.float32) + active_count = int(active_mask.sum().item()) + worst_token = _worst_token(absolute_diff, active_mask) + + materialized_case = _plain_mapping(materialized.get("materialized_case")) + case = _plain_mapping(materialized_case.get("case")) + requested = _plain_mapping(case.get("requested")) + rollout_requested = _plain_mapping(requested.get("rollout")) + training_requested = _plain_mapping(requested.get("training")) + rollout_actual = _plain_mapping(actual.get("rollout")) + training_actual = _plain_mapping(actual.get("training")) + diagnostics = _plain_mapping(comparison.get("diagnostics")) + + failures: list[dict[str, Any]] = [] + if not bool(comparison.get("comparable")): + failures.append( + { + "code": "not_comparable", + "message": "The sealed attempt failed an identity or artifact comparison gate.", + } + ) + for field in ("identity_errors", "artifact_errors"): + for value in comparison.get(field) or (): + failures.append({"code": field, "message": str(value)}) + if not bool(comparison.get("passed")) and not failures: + failures.append( + { + "code": str(comparison.get("status", "comparison_failed")), + "message": "The selected-token comparison exceeded its fixed contract.", + } + ) + + metrics = { + "active_token_count": active_count, + "mismatch_count": int(comparison.get("mismatch_count", 0)), + "max_abs_dlogp": 0.0 if worst_token is None else worst_token["abs_dlogp"], + "fixed_threshold": comparison.get("fixed_threshold"), + "comparison_status": comparison.get("status"), + "comparable": bool(comparison.get("comparable")), + "passed": bool(comparison.get("passed")), + } + axes = { + "case_id": comparison.get("case_id"), + "attempt_id": comparison.get("attempt_id"), + "rollout_tp": rollout_requested.get("tensor_parallel_size"), + "rollout_cp": rollout_requested.get("context_parallel_size"), + "rollout_dtype": rollout_requested.get("dtype"), + "training_dtype": training_requested.get("compute_dtype"), + "training_sharding": training_requested.get("sharding"), + "contract_fingerprint": comparison.get("contract_fingerprint"), + } + rollout_backend = _first_present( + rollout_actual.get("backend_id"), rollout_actual.get("actual_backend") + ) + training_backend = _first_present( + training_actual.get("backend_id"), training_actual.get("actual_backend") + ) + provenance = { + "operator_source": actual.get("operator_source"), + "execution_fingerprint": actual.get("execution_fingerprint"), + "environment_fingerprint": actual.get("environment_fingerprint"), + "actual_backend": _backend_label(rollout_backend, training_backend), + "rollout_backend": rollout_backend, + "training_backend": training_backend, + "rollout": rollout_actual, + "training": training_actual, + } + samples = _samples_from_identity(identity, active_mask) + replay_manifest = { + "mode": "cross_config_fixed_replay", + "samples": samples, + "batch_invariance_cases": (), + "validation": {"warnings": (), "failures": failures}, + "runtime_provenance": provenance, + } + result_cube = { + "mode": "cross_config_fixed_replay", + "axes": axes, + "metrics": metrics, + "worst_token": worst_token or {}, + "metadata_validation": replay_manifest["validation"], + "runtime_provenance": provenance, + "diagnostics": diagnostics, + } + return build_drift_report( + replay_manifest=replay_manifest, + result_cube=result_cube, + runtime_provenance=provenance, + title=title + or f"Cross-config drift: {comparison.get('case_id', directory.parent.name)} / {directory.name}", + ) + + +def build_drift_report( + *, + replay_manifest: Mapping[str, Any] | None = None, + result_cube: Mapping[str, Any] | None = None, + runtime_provenance: Mapping[str, Any] | None = None, + title: str | None = None, +) -> dict[str, Any]: + """Normalize existing audit artifacts into a visual diagnostic report.""" + + manifest = _plain_mapping(replay_manifest) + cube = _plain_mapping(result_cube) + metrics = _plain_mapping(cube.get("metrics")) + axes = _plain_mapping(cube.get("axes")) + validation = _plain_mapping(cube.get("metadata_validation") or manifest.get("validation")) + provenance = _plain_mapping( + runtime_provenance or cube.get("runtime_provenance") or manifest.get("runtime_provenance") + ) + samples = [ + _plain_mapping(item) for item in manifest.get("samples", []) if isinstance(item, Mapping) + ] + warnings = [ + _plain_mapping(item) for item in validation.get("warnings", []) if isinstance(item, Mapping) + ] + failures = [ + _plain_mapping(item) for item in validation.get("failures", []) if isinstance(item, Mapping) + ] + + max_abs_dlogp = _number(metrics.get("max_abs_dlogp")) + warning_count = _number(metrics.get("warning_count"), default=0.0) or 0.0 + metadata_warning_count = ( + _number(metrics.get("metadata_warning_count"), default=float(len(warnings))) or 0.0 + ) + metadata_failure_count = ( + _number(metrics.get("metadata_failure_count"), default=float(len(failures))) or 0.0 + ) + runtime_fallback = bool(metrics.get("runtime_fallback") or provenance.get("fallback")) + strict_failure = bool(metrics.get("runtime_strict_failure") or provenance.get("strict_failure")) + comparison_passed = metrics.get("passed") is True + + if strict_failure or metadata_failure_count > 0: + status = "failure" + elif warning_count > 0 or metadata_warning_count > 0 or runtime_fallback: + status = "warning" + elif comparison_passed: + status = "pass" + elif (max_abs_dlogp or 0.0) > 0.0: + status = "warning" + else: + status = "pass" + + span = max(1.0, float(len(samples))) + has_timestamps = any(_number(sample.get("start_ts")) is not None for sample in samples) + timeline_mode = "timestamp" if has_timestamps else "ordinal_diagnostic" + events: list[dict[str, Any]] = [] + + train_status = ( + "failure" if status == "failure" else "warning" if status == "warning" else "pass" + ) + events.append( + { + "id": "train-audit", + "kind": "bar", + "lane": "Training audit", + "start": 0.0, + "end": span, + "label": "training-side audit", + "status": train_status, + "details": { + "mode": manifest.get("mode", cube.get("mode", "unknown")), + "rank": cube.get("rank"), + }, + } + ) + for position, sample in enumerate(samples): + sample_label = sample.get("sample_index") + if sample_label is None: + sample_label = sample.get("rollout_id") + if sample_label is None: + sample_label = position + start = _number(sample.get("start_ts"), default=float(position)) + end = _number(sample.get("end_ts"), default=float(start or position) + 0.82) + if end is None or start is None or end <= start: + start, end = float(position), float(position) + 0.82 + events.append( + { + "id": f"rollout-{position}", + "kind": "bar", + "lane": "Rollout samples", + "start": start, + "end": end, + "label": f"sample {sample_label}", + "status": "info", + "details": sample, + } + ) + + actual_backend = _first_present( + provenance.get("actual_backend"), + provenance.get("backend_id"), + ( + cube.get("axes", {}).get("logp_backend") + if isinstance(cube.get("axes"), Mapping) + else None + ), + provenance.get("requested_backend"), + "unknown backend", + ) + operator_status = "failure" if strict_failure else "warning" if runtime_fallback else "pass" + events.append( + { + "id": "operator-backend", + "kind": "bar", + "lane": "Operator / backend", + "start": 0.12, + "end": max(0.94, span - 0.12), + "label": str(actual_backend), + "status": operator_status, + "details": provenance, + } + ) + + worst_token = _plain_mapping(cube.get("worst_token")) + if worst_token: + marker_position = ( + _number(worst_token.get("sample_position"), default=max(0.0, span - 0.5)) or 0.0 + ) + events.append( + { + "id": "worst-drift", + "kind": "marker", + "lane": "Drift markers", + "start": marker_position + 0.41, + "end": marker_position + 0.41, + "label": f"|dlogp| {_format_number(worst_token.get('abs_dlogp'))}", + "status": ( + "failure" + if status == "failure" + else "warning" if status == "warning" else "pass" + ), + "details": worst_token, + } + ) + if warnings or failures or warning_count: + marker_status = "failure" if failures else "warning" + events.append( + { + "id": "validation-marker", + "kind": "marker", + "lane": "Drift markers", + "start": max(0.2, span - 0.22), + "end": max(0.2, span - 0.22), + "label": f"{len(failures)} failures / {len(warnings)} warnings", + "status": marker_status, + "details": { + "warnings": warnings, + "failures": failures, + "dlogp_warning_count": warning_count, + }, + } + ) + + return { + "schema_version": REPORT_SCHEMA_VERSION, + "title": title or "RL-Kernel cross-config drift report", + "status": status, + "status_label": _STATUS_LABELS[status], + "timeline_mode": timeline_mode, + "timeline_note": ( + "Real artifact timestamps are shown." + if timeline_mode == "timestamp" + else "No artifact timestamps were available; positions are stable sample ordinals, not elapsed time." + ), + "lanes": [ + "Training audit", + "Rollout samples", + "Operator / backend", + "Drift markers", + ], + "events": events, + "axes": axes, + "metrics": metrics, + "worst_token": worst_token, + "validation": validation, + "runtime_provenance": provenance, + "sample_count": len(samples), + "replay_case_count": len(manifest.get("batch_invariance_cases", []) or []), + "manifest_fingerprint": manifest.get("fingerprint"), + "cube_fingerprint": cube.get("fingerprint"), + } + + +def build_drift_trace(report: Mapping[str, Any]) -> dict[str, Any]: + """Convert a consistency report into Chrome Trace Event JSON. + + The resulting file can be opened directly in Perfetto and expanded by + process/thread track. It deliberately uses the report's timestamp mode; + ordinal reports remain diagnostic sample positions, never fabricated time. + """ + + normalized = _plain_mapping(report) + events = [_plain_mapping(event) for event in normalized.get("events", [])] + status = str(normalized.get("status", "info")) + timeline_mode = str(normalized.get("timeline_mode", "ordinal_diagnostic")) + lanes = [ + ("Audit", None), + ("Training audit", "Training audit"), + ("Rollout samples", "Rollout samples"), + ("Execution", None), + ("Operator / backend", "Operator / backend"), + ("Token comparison", "Token comparison"), + ("Validation", None), + ("Drift markers", "Drift markers"), + ] + lane_ids = {name: 100 + index for index, (name, _) in enumerate(lanes)} + process_id = 1 + trace_events: list[dict[str, Any]] = [ + { + "name": "process_name", + "ph": "M", + "pid": process_id, + "args": {"name": "RL-Kernel cross-config drift"}, + }, + ] + for lane, event_lane in lanes: + if event_lane is None: + continue + trace_events.append( + { + "name": "thread_name", + "ph": "M", + "pid": process_id, + "tid": lane_ids[lane], + "args": {"name": lane}, + } + ) + + def trace_time(value: Any) -> float: + number = _number(value, default=0.0) or 0.0 + # Chrome Trace timestamps are microseconds. In ordinal mode the same + # scale is retained only to make adjacent sample positions visible. + return number * 1_000_000.0 + + event_colors = { + "pass": "good", + "warning": "terrible", + "failure": "bad", + "info": "thread_state_running", + } + for event in events: + event_lane = str(event.get("lane", "Drift markers")) + tid = lane_ids.get(event_lane, lane_ids["Drift markers"]) + start = _number(event.get("start"), default=0.0) or 0.0 + end = _number(event.get("end"), default=start) or start + details = event.get("details") if isinstance(event.get("details"), Mapping) else {} + args = { + "event_id": str(event.get("id", "")), + "status": str(event.get("status", status)), + "timeline_mode": timeline_mode, + "timeline_note": normalized.get("timeline_note", ""), + "details": details, + } + color = event_colors.get(str(event.get("status", "info")), "thread_state_running") + if event.get("kind") == "marker": + trace_events.append( + { + "name": str(event.get("label", "marker")), + "cat": "consistency.drift", + "ph": "I", + "s": "t", + "pid": process_id, + "tid": tid, + "ts": trace_time(start), + "cname": color, + "args": args, + } + ) + continue + trace_events.append( + { + "name": str(event.get("label", event.get("id", "event"))), + "cat": "consistency.audit", + "ph": "X", + "pid": process_id, + "tid": tid, + "ts": trace_time(start), + "dur": max(1.0, trace_time(end) - trace_time(start)), + "cname": color, + "args": args, + } + ) + + metrics = _plain_mapping(normalized.get("metrics")) + max_abs_dlogp = _number(metrics.get("max_abs_dlogp")) + if max_abs_dlogp is not None: + trace_events.append( + { + "name": "max |dlogp|", + "cat": "consistency.metric", + "ph": "C", + "pid": process_id, + "tid": lane_ids["Token comparison"], + "ts": 0.0, + "args": {"max_abs_dlogp": max_abs_dlogp}, + } + ) + + return { + "traceEvents": trace_events, + "displayTimeUnit": "ms", + "metadata": { + "report_title": normalized.get("title", "RL-Kernel cross-config drift report"), + "status": status, + "timeline_mode": timeline_mode, + "timeline_note": normalized.get("timeline_note", ""), + "schema_version": normalized.get("schema_version", REPORT_SCHEMA_VERSION), + }, + } + + +def write_drift_trace(report: Mapping[str, Any], path: str | Path) -> Path: + """Write a Chrome Trace Event JSON file for Perfetto or trace viewers.""" + + output = Path(path) + output.parent.mkdir(parents=True, exist_ok=True) + output.write_text( + json.dumps(build_drift_trace(report), ensure_ascii=True, indent=2), + encoding="utf-8", + ) + return output + + +def write_drift_bundle( + report: Mapping[str, Any], + path: str | Path, + *, + include_preview: bool = True, +) -> Path: + """Write a self-contained ``.rlk-drift`` desktop-viewer bundle. + + The bundle contains sanitized report JSON and a portable trace. The PNG + preview is included for PR/issue sharing but is not required by the viewer. + Raw train dumps are intentionally not copied into the bundle. + """ + + output = Path(path) + output.parent.mkdir(parents=True, exist_ok=True) + normalized = _plain_mapping(report) + trace = build_drift_trace(normalized) + manifest = { + "format": "rl_kernel.cross_config_drift", + "bundle_version": 1, + "report_schema_version": normalized.get("schema_version", REPORT_SCHEMA_VERSION), + "files": [ + "manifest.json", + "report.json", + "trace.json", + *(["preview.png"] if include_preview else []), + ], + "preview_included": bool(include_preview), + "title": normalized.get("title", "RL-Kernel cross-config drift report"), + "status": normalized.get("status", "info"), + } + with zipfile.ZipFile(output, "w", compression=zipfile.ZIP_DEFLATED) as archive: + archive.writestr("manifest.json", json.dumps(manifest, ensure_ascii=True, indent=2)) + archive.writestr("report.json", json.dumps(normalized, ensure_ascii=True, indent=2)) + archive.writestr("trace.json", json.dumps(trace, ensure_ascii=True, indent=2)) + if include_preview: + preview = render_drift_report_image(normalized) + buffer = io.BytesIO() + preview.save(buffer, format="PNG", optimize=True) + archive.writestr("preview.png", buffer.getvalue()) + return output + + +def load_drift_bundle(path: str | Path) -> dict[str, Any]: + """Load a ``.rlk-drift`` bundle without importing the optional GUI.""" + + with zipfile.ZipFile(path, "r") as archive: + names = set(archive.namelist()) + if "report.json" not in names or "trace.json" not in names: + raise ValueError( + "invalid RL-Kernel drift bundle: report.json and trace.json are required" + ) + report = json.loads(archive.read("report.json")) + trace = json.loads(archive.read("trace.json")) + manifest = json.loads(archive.read("manifest.json")) if "manifest.json" in names else {} + return {"manifest": manifest, "report": report, "trace": trace} + + +def render_drift_report(report: Mapping[str, Any]) -> str: + """Render a report as a self-contained HTML document with an SVG timeline.""" + + normalized = _plain_mapping(report) + events = [_plain_mapping(event) for event in normalized.get("events", [])] + lanes = [str(lane) for lane in normalized.get("lanes", [])] + width = 1180 + left = 190 + right = 28 + top = 46 + row_height = 54 + timeline_width = width - left - right + timeline_span = max( + 1.0, + max( + (_number(event.get("end"), default=1.0) or 1.0 for event in events), + default=1.0, + ), + ) + svg_height = top + row_height * len(lanes) + 42 + event_map = {str(event.get("id")): event for event in events} + + def x(value: float) -> float: + return left + max(0.0, min(timeline_span, value)) / timeline_span * timeline_width + + svg_parts = [ + f'', + f'', + ] + grid_steps = min(12, max(2, int(timeline_span) + 1)) + for index in range(grid_steps + 1): + value = timeline_span * index / grid_steps + xpos = x(value) + svg_parts.append( + f'' + ) + svg_parts.append( + f'{html.escape(_format_number(value))}' + ) + for lane_index, lane in enumerate(lanes): + ypos = top + lane_index * row_height + svg_parts.append( + f'{html.escape(lane)}' + ) + svg_parts.append( + f'' + ) + + for event in events: + lane_index = lanes.index(str(event.get("lane"))) if str(event.get("lane")) in lanes else 0 + ypos = top + lane_index * row_height + status = str(event.get("status", "info")) + color = _STATUS_COLORS.get(status, _STATUS_COLORS["info"]) + event_id = html.escape(str(event.get("id")), quote=True) + label = html.escape(_truncate(str(event.get("label", "event")), 28)) + title = html.escape( + f"{event.get('label', 'event')} [{_STATUS_LABELS.get(status, status.upper())}]", + quote=True, + ) + if event.get("kind") == "marker": + xpos = x(_number(event.get("start"), default=0.0) or 0.0) + points = f"{xpos:.2f},{ypos + 7} {xpos + 9:.2f},{ypos + 16} {xpos:.2f},{ypos + 25} {xpos - 9:.2f},{ypos + 16}" + svg_parts.append( + f'{title}' + ) + svg_parts.append( + f'{label}' + ) + else: + start = _number(event.get("start"), default=0.0) or 0.0 + end = _number(event.get("end"), default=start + 0.5) or start + 0.5 + xpos = x(start) + event_width = max(8.0, x(end) - xpos) + svg_parts.append( + f'{title}' + ) + if event_width > 60: + svg_parts.append( + f'{label}' + ) + else: + svg_parts.append( + f'{label}' + ) + svg_parts.append("") + svg = "".join(svg_parts) + + metrics = _plain_mapping(normalized.get("metrics")) + metric_cards = [ + ("Max |dlogp|", _format_number(metrics.get("max_abs_dlogp"))), + ("Active tokens", _format_number(metrics.get("active_token_count"))), + ("Warnings", _format_number(metrics.get("warning_count"), default="0")), + ("Replay cases", str(normalized.get("replay_case_count", 0))), + ] + cards_html = "".join( + f'
{html.escape(name)}
{html.escape(value)}
' + for name, value in metric_cards + ) + axes_html = _render_key_value_table( + normalized.get("axes"), empty="No normalized axes were recorded." + ) + provenance_html = _render_key_value_table( + normalized.get("runtime_provenance"), + empty="No runtime provenance was recorded.", + ) + validation_html = _render_validation(normalized.get("validation")) + event_json = json.dumps(event_map, ensure_ascii=True, separators=(",", ":")).replace( + "<", "\\u003c" + ) + status = str(normalized.get("status", "info")) + status_color = _STATUS_COLORS.get(status, _STATUS_COLORS["info"]) + title = html.escape(str(normalized.get("title", "RL-Kernel cross-config drift report"))) + timeline_note = html.escape(str(normalized.get("timeline_note", ""))) + + return f""" + +{title} +
+

{title}

{timeline_note}
{html.escape(_STATUS_LABELS.get(status, status.upper()))}
+
{cards_html}
+

Operator drift timeline {html.escape(str(normalized.get("timeline_mode", "diagnostic")))}

{svg}
+

Selected event

Select a bar or marker in the timeline.
+

Normalized axes

{axes_html}

Runtime provenance

{provenance_html}
+

Validation

{validation_html}
+
""" + + +def write_drift_report(report: Mapping[str, Any], path: str | Path) -> Path: + """Write a self-contained HTML report and return its resolved path.""" + + output = Path(path) + output.parent.mkdir(parents=True, exist_ok=True) + output.write_text(render_drift_report(report), encoding="utf-8") + return output + + +def render_drift_report_image(report: Mapping[str, Any], *, width: int = 2400) -> Any: + """Render a static profiler-style report image. + + The image is intentionally self-contained and suitable for attaching to a + PR, issue, or debug artifact. It uses a light, high-density profiler + layout: a compact tool chrome, a hierarchical track tree, a fine-grained + ruler/grid, thin event bars, and tabular details below the timeline. + """ + + from PIL import Image, ImageDraw + + normalized = _plain_mapping(report) + metrics = _plain_mapping(normalized.get("metrics")) + events = [_plain_mapping(event) for event in normalized.get("events", [])] + status = str(normalized.get("status", "info")) + status_color = _REPORT_IMAGE_COLORS.get(status, _REPORT_IMAGE_COLORS["info"]) + bg = _REPORT_IMAGE_COLORS["background"] + panel_bg = _REPORT_IMAGE_COLORS["panel"] + panel_alt = _REPORT_IMAGE_COLORS["panel_alt"] + line = _REPORT_IMAGE_COLORS["line"] + text_color = _REPORT_IMAGE_COLORS["text"] + muted = _REPORT_IMAGE_COLORS["muted"] + + height = 1680 + image = Image.new("RGB", (width, height), bg) + draw = ImageDraw.Draw(image) + regular = _load_report_font(22) + small = _load_report_font(18) + tiny = _load_report_font(15) + label_font = _load_report_font(18, bold=True) + section_font = _load_report_font(20, bold=True) + title_font = _load_report_font(30, bold=True) + metric_font = _load_report_font(25, bold=True) + mono = _load_report_font(16, mono=True) + + def rect( + box: tuple[int, int, int, int], + fill: str, + radius: int = 2, + outline: str | None = None, + ) -> None: + draw.rectangle(box, fill=fill, outline=outline, width=1 if outline else 1) + + def write( + x: int, + y: int, + value: Any, + font: Any = regular, + fill: str = text_color, + anchor: str | None = None, + ) -> None: + draw.text((x, y), str(value), font=font, fill=fill, anchor=anchor) + + def fit(value: Any, limit: int) -> str: + return _truncate(str(value), limit) + + def key_value_rows(values: Any, limit: int = 33) -> list[tuple[str, str]]: + if not isinstance(values, Mapping): + return [] + return [ + (fit(key, limit), fit(_format_value(values[key]), 47)) + for key in sorted(values, key=str) + ] + + # Nsight-like tool chrome: compact title, run metadata, and controls. + draw.rectangle((0, 0, width, 42), fill="#252a30") + write(30, 12, "RL-KERNEL CROSS-CONFIG ANALYSIS", tiny, "#e9edf1") + write( + width - 30, + 12, + "static report | schema v" + str(normalized.get("schema_version", REPORT_SCHEMA_VERSION)), + tiny, + "#b7c0c8", + anchor="ra", + ) + write( + 34, + 60, + normalized.get("title", "RL-Kernel cross-config drift report"), + title_font, + text_color, + ) + write(34, 101, "run / audit replay", tiny, muted) + write(180, 101, normalized.get("timeline_note", ""), tiny, muted) + badge_text = _STATUS_LABELS.get(status, status.upper()) + draw.rectangle((width - 188, 61, width - 34, 101), fill=panel_bg, outline=status_color, width=2) + write(width - 111, 81, badge_text, label_font, status_color, anchor="mm") + + # Compact metric strip. These are values, not dashboard cards. + strip_y = 127 + draw.rectangle((34, strip_y, width - 34, strip_y + 64), fill=panel_bg, outline=line, width=1) + strip_items = [ + ("MAX |DLOGP|", _format_number(metrics.get("max_abs_dlogp")), status_color), + ( + "ACTIVE TOKENS", + _format_number(metrics.get("active_token_count")), + _REPORT_IMAGE_COLORS["blue"], + ), + ( + "WARNINGS", + _format_number(metrics.get("warning_count"), default="0"), + _REPORT_IMAGE_COLORS["yellow"], + ), + ( + "REPLAY CASES", + normalized.get("replay_case_count", 0), + _REPORT_IMAGE_COLORS["purple"], + ), + ] + strip_width = (width - 68) // len(strip_items) + for index, (name, value, accent) in enumerate(strip_items): + x0 = 34 + index * strip_width + if index: + draw.line((x0, strip_y + 10, x0, strip_y + 54), fill=line, width=1) + draw.rectangle((x0, strip_y, x0 + 4, strip_y + 64), fill=accent) + write(x0 + 18, strip_y + 12, name, tiny, muted) + write(x0 + 18, strip_y + 34, value, metric_font, text_color) + + # Timeline panel: strict left track tree + dense aligned lanes. + timeline_y = 216 + timeline_h = 620 + rect((34, timeline_y, width - 34, timeline_y + timeline_h), panel_bg, 2, line) + write(52, timeline_y + 14, "CAPTURED EXECUTION", section_font, text_color) + mode = str(normalized.get("timeline_mode", "diagnostic")) + mode_color = ( + _REPORT_IMAGE_COLORS["yellow"] if mode != "timestamp" else _REPORT_IMAGE_COLORS["blue"] + ) + write(width - 52, timeline_y + 18, mode, tiny, mode_color, anchor="ra") + + lanes = [ + ("Audit", None), + ("Training audit", "Training audit"), + ("Rollout samples", "Rollout samples"), + ("Execution", None), + ("Operator / backend", "Operator / backend"), + ("Token comparison", "Token comparison"), + ("Validation", None), + ("Drift markers", "Drift markers"), + ] + left = 370 + right = width - 54 + ruler_y = timeline_y + 59 + chart_top = timeline_y + 91 + row_height = 62 + chart_bottom = chart_top + row_height * len(lanes) + max_end = max((_number(event.get("end"), default=1.0) or 1.0 for event in events), default=1.0) + span = max(1.0, max_end) + + def x_for(value: float) -> int: + return int(left + max(0.0, min(span, value)) / span * (right - left)) + + # Group header, ruler, minor grid, and hierarchical track labels. + draw.rectangle((34, timeline_y + 42, width - 34, timeline_y + 80), fill="#f1f3f5") + write(52, timeline_y + 54, "TRACKS", tiny, muted) + write(left + 8, timeline_y + 54, "TIME / SAMPLE ORDINAL", tiny, muted) + draw.line((left, ruler_y + 16, right, ruler_y + 16), fill=line, width=1) + for tick in range(0, 33): + value = span * tick / 32 + xpos = x_for(value) + major = tick % 4 == 0 + draw.line( + (xpos, ruler_y + (7 if major else 13), xpos, chart_bottom), + fill=line if major else _REPORT_IMAGE_COLORS["grid"], + width=1, + ) + if major: + write(xpos, ruler_y - 3, _format_number(value), tiny, muted, anchor="ma") + + for lane_index, (lane, event_lane) in enumerate(lanes): + y0 = chart_top + lane_index * row_height + draw.rectangle( + (34, y0, right, y0 + row_height), + fill=panel_alt if lane_index % 2 else panel_bg, + ) + draw.line((34, y0 + row_height, right, y0 + row_height), fill=line, width=1) + if event_lane is None: + draw.rectangle((34, y0, left, y0 + row_height), fill="#e8ebee") + write(52, y0 + 20, lane.upper(), label_font, text_color) + else: + write(52, y0 + 22, "|--", tiny, muted) + write(92, y0 + 21, lane, tiny, text_color) + + event_colors = { + "pass": _REPORT_IMAGE_COLORS["green"], + "warning": _REPORT_IMAGE_COLORS["yellow"], + "failure": _REPORT_IMAGE_COLORS["red"], + "info": _REPORT_IMAGE_COLORS["blue"], + } + for event in events: + event_lane = str(event.get("lane", "Drift markers")) + lane_index = next( + (idx for idx, (_, name) in enumerate(lanes) if name == event_lane), + len(lanes) - 1, + ) + y0 = chart_top + lane_index * row_height + color = event_colors.get(str(event.get("status", "info")), event_colors["info"]) + start = _number(event.get("start"), default=0.0) or 0.0 + end = _number(event.get("end"), default=start + 0.5) or start + 0.5 + x0 = x_for(start) + x1 = max(x0 + 10, x_for(end)) + if event.get("kind") == "marker": + mid = x_for(start) + draw.line((mid, y0 + 6, mid, y0 + row_height - 6), fill=color, width=3) + draw.polygon([(mid - 7, y0 + 7), (mid + 7, y0 + 7), (mid, y0 + 17)], fill=color) + marker_label = fit(event.get("label", "marker"), 34) + label_x = mid + 14 if mid <= right - 250 else max(left + 10, mid - 240) + write(label_x, y0 + 22, marker_label, tiny, color) + else: + bar_y = y0 + 22 + draw.rectangle((x0, bar_y, x1, bar_y + 18), fill=color) + if x1 - x0 >= 130: + write( + x0 + 10, + bar_y + 2, + fit(event.get("label", "event"), 42), + tiny, + "#ffffff", + ) + else: + write( + min(x1 + 10, right - 260), + bar_y + 2, + fit(event.get("label", "event"), 34), + tiny, + text_color, + ) + + # Keep the comparison track meaningful even when the dump only contains a + # scalar worst-token summary instead of per-token samples. + worst = _plain_mapping(normalized.get("worst_token")) + comparison_y = chart_top + 5 * row_height + 31 + draw.line( + (left + 18, comparison_y, right - 18, comparison_y), + fill=_REPORT_IMAGE_COLORS["track"], + width=3, + ) + write(left + 18, comparison_y - 22, "train vs rollout", tiny, muted) + if worst: + sample_position = _number(worst.get("sample_position"), default=0.0) or 0.0 + comparison_x = x_for(sample_position + 0.5) + draw.line( + (comparison_x, comparison_y - 15, comparison_x, comparison_y + 15), + fill=status_color, + width=3, + ) + write( + comparison_x + 10, + comparison_y - 10, + f"delta={_format_number(worst.get('abs_dlogp'))}", + tiny, + status_color, + ) + + # A compact legend makes the static image readable without hover state. + legend_y = timeline_y + timeline_h - 30 + for x0, label, color in ( + (52, "PASS", _REPORT_IMAGE_COLORS["green"]), + (130, "WARN", _REPORT_IMAGE_COLORS["yellow"]), + (214, "FAIL", _REPORT_IMAGE_COLORS["red"]), + (298, "INFO", _REPORT_IMAGE_COLORS["blue"]), + ): + draw.rectangle((x0, legend_y + 3, x0 + 14, legend_y + 15), fill=color) + write(x0 + 22, legend_y, label, tiny, muted) + write( + right, + legend_y, + "positions are sample ordinals when timestamps are absent", + tiny, + muted, + anchor="ra", + ) + + # Diagnostic summary panels. + summary_y = 850 + summary_h = 220 + half_gap = 18 + half_width = (width - 68 - half_gap) // 2 + rect((34, summary_y, 34 + half_width, summary_y + summary_h), panel_bg, 2, line) + rect( + (34 + half_width + half_gap, summary_y, width - 34, summary_y + summary_h), + panel_bg, + 2, + line, + ) + write(54, summary_y + 18, "DRIFT SUMMARY", section_font) + write(54, summary_y + 57, "observed maximum", tiny, muted) + max_value = _number(metrics.get("max_abs_dlogp"), default=0.0) or 0.0 + bar_x = 54 + bar_y = summary_y + 91 + bar_width = half_width - 110 + rect((bar_x, bar_y, bar_x + bar_width, bar_y + 24), _REPORT_IMAGE_COLORS["track"], 4) + fill_width = int(min(1.0, max_value / max(max_value, 1.0)) * bar_width) if max_value else 0 + if fill_width: + rect((bar_x, bar_y, bar_x + max(8, fill_width), bar_y + 24), status_color, 4) + write(bar_x, bar_y + 30, "0", tiny, muted) + write( + bar_x + bar_width, + bar_y + 30, + _format_number(max(max_value, 1.0)), + tiny, + muted, + anchor="ra", + ) + worst = _plain_mapping(normalized.get("worst_token")) + worst_text = ( + f"worst token: sample={worst.get('sample_position', '-')} token={worst.get('token_position', '-')} " + f"|dlogp|={_format_number(worst.get('abs_dlogp'))}" + ) + write(54, summary_y + 157, fit(worst_text, 92), mono, text_color) + right_x = 34 + half_width + half_gap + 20 + write(right_x, summary_y + 18, "SELECTED EVENT", section_font) + warning_items = _plain_mapping(normalized.get("validation")) + warnings = warning_items.get("warnings") or [] + failures = warning_items.get("failures") or [] + issue_text = "No anomaly recorded." + if failures or warnings: + first = (failures or warnings)[0] + first = first if isinstance(first, Mapping) else {"message": first} + issue_text = f"{str(first.get('code', 'validation'))}: {str(first.get('message', ''))}" + elif worst: + issue_text = ( + f"worst token: sample={worst.get('sample_position', '-')} " + f"token={worst.get('token_position', '-')} " + f"|dlogp|={_format_number(worst.get('abs_dlogp'))}" + ) + write(right_x, summary_y + 57, fit(issue_text, 95), small, status_color) + write(right_x, summary_y + 101, "worst token / validation marker", tiny, muted) + write( + right_x, + summary_y + 140, + f"status={badge_text} warnings={len(warnings)} failures={len(failures)}", + mono, + text_color, + ) + + # Bottom tables: axis capture and runtime provenance are the actionable part + # of the image for a post-training user. + table_y = 1095 + table_h = 400 + rect((34, table_y, 34 + half_width, table_y + table_h), panel_bg, 2, line) + rect( + (34 + half_width + half_gap, table_y, width - 34, table_y + table_h), + panel_bg, + 2, + line, + ) + write(54, table_y + 18, "CAPTURED EXECUTION AXES", section_font) + write( + 34 + half_width + half_gap + 20, + table_y + 18, + "RUNTIME PROVENANCE", + section_font, + ) + axes_rows = key_value_rows(normalized.get("axes")) + provenance_rows = key_value_rows(normalized.get("runtime_provenance")) + + def table(rows: list[tuple[str, str]], x0: int, y0: int, w: int, max_rows: int = 8) -> None: + row_y = y0 + for index, (key, value) in enumerate(rows[:max_rows]): + if index % 2 == 0: + draw.rectangle((x0, row_y - 3, x0 + w, row_y + 35), fill=panel_alt) + write(x0 + 14, row_y + 8, key, tiny, muted) + write(x0 + 270, row_y + 8, value, mono, text_color) + draw.line((x0, row_y + 38, x0 + w, row_y + 38), fill=line, width=1) + row_y += 39 + if not rows: + write(x0 + 14, row_y + 8, "not recorded", small, muted) + + table(axes_rows, 54, table_y + 66, half_width - 40) + table(provenance_rows, 34 + half_width + half_gap + 20, table_y + 66, half_width - 40) + validation_label = "VALIDATION: PASS" if status == "pass" else f"VALIDATION: {badge_text}" + if failures or warnings: + validation_label = f"VALIDATION: {badge_text} | {fit(issue_text, 86)}" + write(70, height - 58, validation_label, tiny, status_color) + footer = "static diagnostic image | schema v" + str( + normalized.get("schema_version", REPORT_SCHEMA_VERSION) + ) + write(width - 54, height - 28, footer, tiny, muted, anchor="ra") + return image + + +def write_drift_report_image(report: Mapping[str, Any], path: str | Path) -> Path: + """Write a static PNG/JPEG consistency drift report image.""" + + output = Path(path) + output.parent.mkdir(parents=True, exist_ok=True) + image = render_drift_report_image(report) + suffix = output.suffix.lower() + image_format = "JPEG" if suffix in {".jpg", ".jpeg"} else "PNG" + if image_format == "JPEG": + image.save(output, format=image_format, quality=95, optimize=True) + else: + image.save(output, format=image_format, optimize=True) + return output + + +def _read_attempt_json(path: Path) -> dict[str, Any]: + try: + value = json.loads(path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as exc: + raise ValueError(f"invalid cross-config report input: {path}") from exc + if not isinstance(value, Mapping): + raise ValueError(f"cross-config report input must be a JSON object: {path}") + return _plain_mapping(value) + + +def _worst_token( + absolute_diff: torch.Tensor, + active_mask: torch.Tensor, +) -> dict[str, Any] | None: + if absolute_diff.shape != active_mask.shape: + raise ValueError("token_diffs absolute_diff and active_mask must have matching shapes") + if not bool(active_mask.any()): + return None + values = absolute_diff.masked_fill(~active_mask, -1.0) + flat_index = int(values.reshape(-1).argmax().item()) + coordinates = list(torch.unravel_index(torch.tensor(flat_index), values.shape)) + sample_position = int(coordinates[0].item()) if coordinates else 0 + token_position = int(coordinates[-1].item()) if coordinates else flat_index + return { + "sample_position": sample_position, + "token_position": token_position, + "abs_dlogp": float(values.reshape(-1)[flat_index].item()), + } + + +def _samples_from_identity( + identity: Mapping[str, Any], + active_mask: torch.Tensor, +) -> list[dict[str, Any]]: + logical_identity = _plain_mapping(identity.get("identity")) + token_ids = logical_identity.get("token_ids") + samples: list[dict[str, Any]] = [] + for index in range(active_mask.shape[0] if active_mask.ndim else 0): + row = token_ids[index] if isinstance(token_ids, list) and index < len(token_ids) else None + samples.append( + { + "sample_position": index, + "sample_index": index, + "token_count": (len(row) if isinstance(row, list) else int(active_mask.shape[-1])), + "active_token_count": int(active_mask[index].sum().item()), + } + ) + return samples + + +def _render_key_value_table(values: Any, *, empty: str) -> str: + if not isinstance(values, Mapping) or not values: + return f'
{html.escape(empty)}
' + rows = [] + for key in sorted(values, key=str): + rows.append( + f"{html.escape(str(key))}{html.escape(_format_value(values[key]))}" + ) + return "" + "".join(rows) + "
" + + +def _render_validation(validation: Any) -> str: + if not isinstance(validation, Mapping): + return '
No validation record was captured.
' + warnings = validation.get("warnings") or [] + failures = validation.get("failures") or [] + if not warnings and not failures: + return '
PASS: no metadata validation warnings or failures.
' + rows = [] + for kind, items in (("failure", failures), ("warning", warnings)): + for item in items: + item = item if isinstance(item, Mapping) else {"message": item} + rows.append( + f'{kind.upper()}{html.escape(str(item.get("code", "")))}{html.escape(str(item.get("message", "")))}' + ) + return ( + "" + + "".join(rows) + + "
StatusCodeMessage
" + ) + + +def _plain_mapping(value: Any) -> dict[str, Any]: + return ( + {str(key): _json_safe(item) for key, item in value.items()} + if isinstance(value, Mapping) + else {} + ) + + +def _json_safe(value: Any) -> Any: + if value is None or isinstance(value, (str, int, float, bool)): + return value + if isinstance(value, Mapping): + return {str(key): _json_safe(item) for key, item in value.items()} + if isinstance(value, (list, tuple)): + return [_json_safe(item) for item in value] + return str(value) + + +def _first_present(*values: Any) -> Any: + for value in values: + if value not in (None, ""): + return value + return None + + +def _backend_label(rollout: Any, training: Any) -> str | None: + """Show one actual backend only when both execution sides agree.""" + + if rollout in (None, "") and training in (None, ""): + return None + if rollout in (None, ""): + return f"training={training}" + if training in (None, ""): + return f"rollout={rollout}" + if rollout == training: + return str(rollout) + return f"rollout={rollout}; training={training}" + + +def _number(value: Any, default: float | None = None) -> float | None: + if value is None or isinstance(value, bool): + return default + try: + return float(value) + except (TypeError, ValueError): + return default + + +def _format_number(value: Any, default: str = "-") -> str: + number = _number(value) + if number is None: + return default + if abs(number) >= 1000 or (abs(number) < 0.001 and number != 0): + return f"{number:.3e}" + return f"{number:.6f}".rstrip("0").rstrip(".") + + +def _format_value(value: Any) -> str: + if isinstance(value, (dict, list, tuple)): + return json.dumps(value, ensure_ascii=True, sort_keys=True) + return str(value) + + +def _truncate(value: str, limit: int) -> str: + return value if len(value) <= limit else value[: max(1, limit - 3)] + "..." + + +def _load_report_font(size: int, *, bold: bool = False, mono: bool = False) -> Any: + """Load a platform font with deterministic fallbacks for report images.""" + + from PIL import ImageFont + + candidates = [] + if mono: + candidates.extend( + [ + "C:/Windows/Fonts/consola.ttf", + "/usr/share/fonts/truetype/dejavu/DejaVuSansMono.ttf", + ] + ) + elif bold: + candidates.extend( + [ + "C:/Windows/Fonts/segoeuib.ttf", + "C:/Windows/Fonts/arialbd.ttf", + "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", + ] + ) + else: + candidates.extend( + [ + "C:/Windows/Fonts/segoeui.ttf", + "C:/Windows/Fonts/arial.ttf", + "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", + ] + ) + for candidate in candidates: + try: + return ImageFont.truetype(candidate, size=size) + except OSError: + continue + return ImageFont.load_default() diff --git a/rl_engine/alignment/cross_config/drift_viewer.py b/rl_engine/alignment/cross_config/drift_viewer.py new file mode 100644 index 00000000..80d8bd65 --- /dev/null +++ b/rl_engine/alignment/cross_config/drift_viewer.py @@ -0,0 +1,288 @@ +"""Optional offline Qt viewer for ``.rlk-drift`` cross-config bundles.""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path +from typing import Any + +from rl_engine.alignment.cross_config.drift_report import load_drift_bundle + +_LANES = [ + ("Audit", None), + ("Training audit", "Training audit"), + ("Rollout samples", "Rollout samples"), + ("Execution", None), + ("Operator / backend", "Operator / backend"), + ("Token comparison", "Token comparison"), + ("Validation", None), + ("Drift markers", "Drift markers"), +] +_COLORS = { + "pass": "#54a64f", + "warning": "#e0a000", + "failure": "#d53f3f", + "info": "#72a6c8", +} + + +def _load_qt() -> tuple[Any, Any, Any]: + try: + from PySide6 import QtCore, QtGui, QtWidgets + except ImportError as exc: # pragma: no cover - depends on workstation extras + raise RuntimeError( + "The offline viewer requires the optional GUI dependency. " + "Install it with: pip install 'rl-engine[drift-viewer]'" + ) from exc + return QtCore, QtGui, QtWidgets + + +def _format_event(event: dict[str, Any]) -> str: + details = event.get("details") if isinstance(event.get("details"), dict) else {} + payload = { + "id": event.get("id"), + "lane": event.get("lane"), + "label": event.get("label"), + "status": event.get("status"), + "start": event.get("start"), + "end": event.get("end"), + "details": details, + } + return json.dumps(payload, ensure_ascii=True, indent=2, sort_keys=True) + + +def _run_qt(bundle_path: Path) -> int: # pragma: no cover - GUI path + QtCore, QtGui, QtWidgets = _load_qt() + bundle = load_drift_bundle(bundle_path) + report = bundle["report"] + events = [event for event in report.get("events", []) if isinstance(event, dict)] + title = str(report.get("title", "RL-Kernel cross-config drift report")) + status = str(report.get("status", "info")) + status_color = _COLORS.get(status, _COLORS["info"]) + + class EventItem(QtWidgets.QGraphicsRectItem): # type: ignore[name-defined] + def __init__(self, rect: Any, event: dict[str, Any], color: str) -> None: + super().__init__(rect) + self.event = event + self.setBrush(QtGui.QColor(color)) + self.setPen(QtGui.QPen(QtGui.QColor(color))) + self.setFlag(QtWidgets.QGraphicsItem.GraphicsItemFlag.ItemIsSelectable, True) + self.setToolTip(str(event.get("label", event.get("id", "event")))) + + class TimelineView(QtWidgets.QGraphicsView): # type: ignore[name-defined] + def __init__(self, scene: Any) -> None: + super().__init__(scene) + self.setRenderHint(QtGui.QPainter.RenderHint.Antialiasing, False) + self.setDragMode(QtWidgets.QGraphicsView.DragMode.ScrollHandDrag) + self.setTransformationAnchor(QtWidgets.QGraphicsView.ViewportAnchor.AnchorUnderMouse) + self.setResizeAnchor(QtWidgets.QGraphicsView.ViewportAnchor.AnchorViewCenter) + self.setHorizontalScrollBarPolicy(QtCore.Qt.ScrollBarPolicy.ScrollBarAlwaysOn) + self.setVerticalScrollBarPolicy(QtCore.Qt.ScrollBarPolicy.ScrollBarAlwaysOff) + self.setStyleSheet("QGraphicsView { border: 0; background: #ffffff; }") + + def wheelEvent(self, event: Any) -> None: + factor = 1.18 if event.angleDelta().y() > 0 else 1.0 / 1.18 + self.scale(factor, 1.0) + event.accept() + + class Window(QtWidgets.QMainWindow): # type: ignore[name-defined] + def __init__(self) -> None: + super().__init__() + self.setWindowTitle(f"RL-Kernel Cross-Config Drift - {title}") + self.resize(1500, 900) + self._scene_width = 2200.0 + self._row_height = 52.0 + self._left_label_width = 8.0 + self._span = max( + 1.0, + max( + (float(event.get("end", 1.0) or 1.0) for event in events), + default=1.0, + ), + ) + self._scene = QtWidgets.QGraphicsScene(self) + self._scene.setBackgroundBrush(QtGui.QColor("#ffffff")) + self._view = TimelineView(self._scene) + self._tree = QtWidgets.QTreeWidget() + self._tree.setHeaderLabel("Tracks") + self._tree.setMinimumWidth(245) + self._tree.setStyleSheet( + "QTreeWidget { background: #f4f5f6; border: 0; " + "font-family: 'Segoe UI'; font-size: 10pt; }" + ) + self._details = QtWidgets.QPlainTextEdit() + self._details.setReadOnly(True) + self._details.setPlaceholderText("Select an event to inspect its audit details.") + self._details.setStyleSheet( + "QPlainTextEdit { background: #ffffff; border-top: 1px solid #c9ced3; " + "font-family: Consolas; font-size: 10pt; padding: 8px; }" + ) + self._build_tree() + self._build_scene() + self._scene.selectionChanged.connect(self._show_selected) + self._tree.itemSelectionChanged.connect(self._select_lane) + self._build_layout() + + def _build_tree(self) -> None: + groups: dict[str, Any] = {} + for label, event_lane in _LANES: + if event_lane is None: + item = QtWidgets.QTreeWidgetItem([label]) + item.setFlags(item.flags() & ~QtCore.Qt.ItemFlag.ItemIsSelectable) + self._tree.addTopLevelItem(item) + groups[label] = item + continue + parent = groups.get( + "Audit" if event_lane in {"Training audit", "Rollout samples"} else "Execution" + ) + if event_lane == "Drift markers": + parent = groups.get("Validation") + item = QtWidgets.QTreeWidgetItem([label]) + item.setData(0, QtCore.Qt.ItemDataRole.UserRole, event_lane) + if parent is not None: + parent.addChild(item) + else: + self._tree.addTopLevelItem(item) + self._tree.expandAll() + + def _x_for(self, value: float) -> float: + return 120.0 + max(0.0, min(self._span, value)) / self._span * self._scene_width + + def _build_scene(self) -> None: + font = QtGui.QFont("Segoe UI", 9) + muted = QtGui.QColor("#69737c") + line = QtGui.QColor("#d6dade") + grid = QtGui.QColor("#e8eaec") + for tick in range(0, 17): + x = self._x_for(self._span * tick / 16.0) + self._scene.addLine( + x, + 0, + x, + len(_LANES) * self._row_height, + QtGui.QPen(line if tick % 4 == 0 else grid), + ) + label = self._scene.addText(f"{self._span * tick / 16.0:g}", font) + label.setDefaultTextColor(muted) + label.setPos(x + 3, -24) + + lane_index = {name: index for index, (_, name) in enumerate(_LANES) if name is not None} + for index, (label, event_lane) in enumerate(_LANES): + y = index * self._row_height + self._scene.addLine( + 0, + y + self._row_height, + self._scene_width + 140, + y + self._row_height, + QtGui.QPen(line), + ) + if event_lane is None: + block = self._scene.addRect( + 0, + y, + self._scene_width + 140, + self._row_height, + QtGui.QPen(), + QtGui.QBrush(QtGui.QColor("#eef0f2")), + ) + block.setZValue(-2) + else: + text = self._scene.addText(label, font) + text.setDefaultTextColor(QtGui.QColor("#252a2f")) + text.setPos(12, y + 17) + + for event in events: + event_lane = str(event.get("lane", "Drift markers")) + if event_lane not in lane_index: + continue + y = lane_index[event_lane] * self._row_height + 16 + start = float(event.get("start", 0.0) or 0.0) + end = float(event.get("end", start) or start) + x0 = self._x_for(start) + x1 = max(x0 + 10.0, self._x_for(end)) + color = _COLORS.get(str(event.get("status", "info")), _COLORS["info"]) + if event.get("kind") == "marker": + item = EventItem(QtCore.QRectF(x0 - 5, y + 5, 10, 10), event, color) + else: + item = EventItem(QtCore.QRectF(x0, y, x1 - x0, 18), event, color) + self._scene.addItem(item) + if x1 - x0 > 130 and event.get("kind") != "marker": + text = self._scene.addText(str(event.get("label", "event")), font) + text.setDefaultTextColor(QtGui.QColor("#ffffff")) + text.setPos(x0 + 8, y - 1) + text.setZValue(1) + self._scene.setSceneRect( + 0, -28, self._scene_width + 140, len(_LANES) * self._row_height + 40 + ) + + def _build_layout(self) -> None: + toolbar = QtWidgets.QToolBar() + toolbar.setMovable(False) + toolbar.setStyleSheet( + "QToolBar { background: #ffffff; border-bottom: 1px solid #c9ced3; }" + ) + status_label = QtWidgets.QLabel(f" {title} STATUS: {status.upper()} ") + status_label.setStyleSheet(f"color: {status_color}; font-weight: 600; padding: 5px;") + toolbar.addWidget(status_label) + fit_button = QtWidgets.QAction("Fit", self) + fit_button.triggered.connect( + lambda: self._view.fitInView( + self._scene.sceneRect(), QtCore.Qt.AspectRatioMode.KeepAspectRatio + ) + ) + toolbar.addAction(fit_button) + toolbar.addWidget(QtWidgets.QLabel(f" {report.get('timeline_note', '')}")) + self.addToolBar(toolbar) + + right = QtWidgets.QSplitter(QtCore.Qt.Orientation.Vertical) + right.addWidget(self._view) + right.addWidget(self._details) + right.setStretchFactor(0, 4) + right.setStretchFactor(1, 1) + root = QtWidgets.QSplitter(QtCore.Qt.Orientation.Horizontal) + root.addWidget(self._tree) + root.addWidget(right) + root.setStretchFactor(1, 1) + self.setCentralWidget(root) + + def _select_lane(self) -> None: + items = self._tree.selectedItems() + if not items: + return + lane = items[0].data(0, QtCore.Qt.ItemDataRole.UserRole) + if not lane: + return + for item in self._scene.items(): + if isinstance(item, EventItem): + item.setOpacity(1.0 if item.event.get("lane") == lane else 0.28) + + def _show_selected(self) -> None: + selected = self._scene.selectedItems() + self._details.setPlainText(_format_event(selected[0].event) if selected else "") + + app = QtWidgets.QApplication.instance() or QtWidgets.QApplication(sys.argv) + window = Window() + window.show() + window._view.fitInView(window._scene.sceneRect(), QtCore.Qt.AspectRatioMode.KeepAspectRatio) + return app.exec() + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser( + description="Open an RL-Kernel drift bundle in the offline desktop viewer." + ) + parser.add_argument("bundle", type=Path, help=".rlk-drift bundle") + args = parser.parse_args(argv) + if args.bundle.suffix.lower() != ".rlk-drift": + parser.error("the offline viewer expects a .rlk-drift bundle") + try: + return _run_qt(args.bundle) + except RuntimeError as exc: + parser.error(str(exc)) + return 2 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/rl_engine/alignment/cross_config/execution_plan.py b/rl_engine/alignment/cross_config/execution_plan.py new file mode 100644 index 00000000..d3be05b2 --- /dev/null +++ b/rl_engine/alignment/cross_config/execution_plan.py @@ -0,0 +1,77 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Canonical, runtime-independent execution-plan construction.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Mapping + +from rl_engine.alignment.cross_config.config import ( + ExperimentConfig, + OperatorSelection, + bind_operator_selection, +) +from rl_engine.alignment.cross_config.planner import PlanningIssue +from rl_engine.alignment.cross_config.schema import ExperimentCase + + +@dataclass(frozen=True) +class ExecutionPlanEntry: + """One operator-bound case and its resolved operator selection.""" + + case: ExperimentCase + operators: OperatorSelection + schema_version: str = "cross_config.execution_plan_entry.v1" + + def to_dict(self) -> dict[str, Any]: + """Return the canonical append-only plan row.""" + + return { + "schema_version": self.schema_version, + "case": self.case.to_dict(), + "operators": self.operators.to_dict(), + } + + +@dataclass(frozen=True) +class ExecutionPlan: + """Canonical metadata shared by planning and every runtime adapter.""" + + experiment: Mapping[str, Any] + entries: tuple[ExecutionPlanEntry, ...] + issues: tuple[PlanningIssue, ...] = () + schema_version: str = "cross_config.execution_plan.v1" + + def rows(self) -> tuple[dict[str, Any], ...]: + """Serialize all plan entries in deterministic execution order.""" + + return tuple(entry.to_dict() for entry in self.entries) + + +def build_execution_plan(config: ExperimentConfig) -> ExecutionPlan: + """Plan, resolve operators, and bind them into immutable case identities.""" + + planned = config.plan() + entries: list[ExecutionPlanEntry] = [] + for case in planned.cases: + operators = config.operators_for(case) + entries.append( + ExecutionPlanEntry( + case=bind_operator_selection(case, operators), + operators=operators, + ) + ) + return ExecutionPlan( + experiment=config.to_dict(), + entries=tuple(entries), + issues=planned.issues, + ) + + +__all__ = [ + "ExecutionPlan", + "ExecutionPlanEntry", + "build_execution_plan", +] diff --git a/rl_engine/alignment/cross_config/operators.py b/rl_engine/alignment/cross_config/operators.py new file mode 100644 index 00000000..2b359219 --- /dev/null +++ b/rl_engine/alignment/cross_config/operators.py @@ -0,0 +1,307 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Target-specific semantic operator selection for alignment cases.""" + +from __future__ import annotations + +import math +from dataclasses import dataclass, replace +from typing import Any, Literal, Mapping, Optional, cast + +import torch + +from rl_engine.kernels.logprob_contract import LogprobContract +from rl_engine.kernels.semantic_registry import ( + OperatorInstanceProvenance, + OperatorRequirements, + OperatorResolution, + OperatorResolutionPolicy, + OperatorSession, + SemanticOperatorCatalog, +) + +OperatorTarget = Literal["rollout", "training", "both"] +ConcreteOperatorTarget = Literal["rollout", "training"] + + +@dataclass(frozen=True) +class OperatorOverride: + """Backend overrides for one semantic operator on either or both sides.""" + + semantic_op: str + rollout_backend: Optional[str] = None + training_backend: Optional[str] = None + + def __post_init__(self) -> None: + semantic_op = self.semantic_op.strip() + if not semantic_op: + raise ValueError("semantic_op must not be empty") + rollout_backend = _normalized_optional_backend(self.rollout_backend) + training_backend = _normalized_optional_backend(self.training_backend) + if rollout_backend is None and training_backend is None: + raise ValueError("operator override must select rollout, training, or both") + object.__setattr__(self, "semantic_op", semantic_op) + object.__setattr__(self, "rollout_backend", rollout_backend) + object.__setattr__(self, "training_backend", training_backend) + + @classmethod + def for_target( + cls, + *, + semantic_op: str, + backend_id: str, + target: OperatorTarget, + ) -> OperatorOverride: + """Create a rollout-only, training-only, or dual-side override.""" + + normalized_target = target.strip().lower() + if normalized_target == "rollout": + return cls(semantic_op=semantic_op, rollout_backend=backend_id) + if normalized_target == "training": + return cls(semantic_op=semantic_op, training_backend=backend_id) + if normalized_target == "both": + return cls( + semantic_op=semantic_op, + rollout_backend=backend_id, + training_backend=backend_id, + ) + raise ValueError("target must be 'rollout', 'training', or 'both'") + + def backend_for(self, target: ConcreteOperatorTarget) -> Optional[str]: + normalized_target = _concrete_target(target) + if normalized_target == "rollout": + return self.rollout_backend + return self.training_backend + + def to_dict(self) -> dict[str, Any]: + return { + "semantic_op": self.semantic_op, + "rollout_backend": self.rollout_backend, + "training_backend": self.training_backend, + } + + +@dataclass(frozen=True) +class ResolvedOperatorOverride: + """Target-specific exact resolutions produced from an operator override.""" + + semantic_op: str + rollout: Optional[OperatorResolution] = None + training: Optional[OperatorResolution] = None + + def for_target(self, target: ConcreteOperatorTarget) -> Optional[OperatorResolution]: + normalized_target = _concrete_target(target) + return self.rollout if normalized_target == "rollout" else self.training + + def to_dict(self) -> dict[str, Any]: + return { + "semantic_op": self.semantic_op, + "rollout": None if self.rollout is None else self.rollout.to_dict(), + "training": None if self.training is None else self.training.to_dict(), + } + + +class OperatorBridge: + """Resolve and instantiate semantic operator overrides without planner branches.""" + + def __init__( + self, + catalog: Optional[SemanticOperatorCatalog | OperatorSession] = None, + *, + policy: Optional[OperatorResolutionPolicy] = None, + ): + """Create a bridge backed by one case-local operator session.""" + + if isinstance(catalog, OperatorSession): + self.catalog = catalog.catalog + self.session = catalog + self.policy = policy or catalog.policy + else: + if catalog is None: + # Built-in descriptors are repository integration details; the + # generic semantic catalog itself remains backend-neutral. + from rl_engine.kernels.registry import kernel_registry + + catalog = kernel_registry.semantic + if not isinstance(catalog, SemanticOperatorCatalog): + raise TypeError("catalog must be a SemanticOperatorCatalog or OperatorSession") + self.catalog = catalog + self.policy = policy or OperatorResolutionPolicy() + self.session = self.catalog.session(self.policy) + + def resolve_override( + self, + override: OperatorOverride, + *, + requirements: Mapping[str, OperatorRequirements], + strict: bool = True, + ) -> ResolvedOperatorOverride: + """Resolve only the sides explicitly selected by ``override``.""" + + resolved: dict[ConcreteOperatorTarget, OperatorResolution] = {} + targets: tuple[ConcreteOperatorTarget, ...] = ("rollout", "training") + for target in targets: + backend_id = override.backend_for(target) + if backend_id is None: + continue + target_requirements = requirements.get(target) + if target_requirements is None: + raise ValueError(f"missing operator requirements for target {target!r}") + target_policy = replace(self.policy, strict=strict) + resolved[target] = self.session.resolve( + semantic_op=override.semantic_op, + requested_backend=backend_id, + target=target, + requirements=target_requirements, + policy=target_policy, + ) + return ResolvedOperatorOverride( + semantic_op=override.semantic_op, + rollout=resolved.get("rollout"), + training=resolved.get("training"), + ) + + def instantiate( + self, + resolved: ResolvedOperatorOverride, + *, + target: ConcreteOperatorTarget, + factory_kwargs: Optional[Mapping[str, Any]] = None, + cache: bool = False, + ) -> Any: + """Instantiate one resolved side; rollout and training remain independent.""" + + resolution = resolved.for_target(target) + if resolution is None: + raise ValueError(f"operator override does not select target {target!r}") + return self.session.instantiate( + resolution, + factory_kwargs=factory_kwargs, + cache=cache, + ) + + def instance_provenance( + self, + resolved: ResolvedOperatorOverride, + *, + target: ConcreteOperatorTarget, + instance: Any, + ) -> OperatorInstanceProvenance: + resolution = resolved.for_target(target) + if resolution is None: + raise ValueError(f"operator override does not select target {target!r}") + return self.session.instance_provenance(resolution, instance) + + +def selected_logprobs_with_operator( + operator: Any, + logits: torch.Tensor, + token_ids: torch.Tensor, + *, + active_mask: Optional[torch.Tensor] = None, + temperature: float = 1.0, + output_dtype: torch.dtype = torch.float32, + contract: Optional[LogprobContract] = None, + tp_group: Any = None, + num_vocab_tiles: Optional[int] = None, + deterministic: bool = True, + validate: bool = True, +) -> torch.Tensor: + """Apply legacy or contract-aware selected-logprob operators uniformly.""" + + if not math.isfinite(temperature) or temperature <= 0.0: + raise ValueError("temperature must be finite and greater than zero") + if logits.shape[:-1] != token_ids.shape: + raise ValueError( + f"logits leading shape {tuple(logits.shape[:-1])} must match " + f"token_ids shape {tuple(token_ids.shape)}" + ) + mask: Optional[torch.Tensor] = None + safe_token_ids = token_ids.to(device=logits.device, dtype=torch.long) + if active_mask is not None: + if active_mask.shape != token_ids.shape: + raise ValueError("active_mask shape must match token_ids shape") + mask = active_mask.to(device=logits.device, dtype=torch.bool) + safe_token_ids = safe_token_ids.masked_fill(~mask, 0) + + if contract is not None: + if not isinstance(contract, LogprobContract): + raise TypeError("contract must be a LogprobContract") + contract_mask = torch.tensor(contract.mask.active_mask, dtype=torch.bool) + if contract_mask.numel() != token_ids.numel(): + raise ValueError("contract active_mask size must match token_ids") + if mask is not None and not torch.equal(contract_mask, mask.detach().cpu().reshape(-1)): + raise ValueError("active_mask must match contract.mask.active_mask") + kwargs: dict[str, Any] = { + "contract": contract, + "tp_group": tp_group, + "deterministic": deterministic, + "validate": validate, + } + if num_vocab_tiles is not None: + kwargs["num_vocab_tiles"] = num_vocab_tiles + if not callable(operator): + raise TypeError("contract-aware selected-logprob operator must be callable") + result = operator( + (logits / float(temperature)).reshape(-1, logits.shape[-1]), + safe_token_ids.reshape(-1), + **kwargs, + ) + if isinstance(result, tuple): + if len(result) != 2: + raise TypeError("contract-aware selected-logprob operator must return (logp, lse)") + selected, lse = result + if not isinstance(selected, torch.Tensor) or not isinstance(lse, torch.Tensor): + raise TypeError("contract-aware selected-logprob operator must return (logp, lse)") + if lse.shape != safe_token_ids.reshape(-1).shape: + raise ValueError("vocab LSE output shape must match token_ids") + else: + selected = result + if isinstance(selected, torch.Tensor): + selected = selected.reshape(token_ids.shape) + else: + scaled_logits = logits.float() / float(temperature) + if hasattr(operator, "apply_fp32") and callable(operator.apply_fp32): + selected = operator.apply_fp32(scaled_logits, safe_token_ids) + elif callable(operator): + selected = operator(scaled_logits, safe_token_ids) + else: + raise TypeError("selected-logprob operator must be callable or expose apply_fp32") + if not isinstance(selected, torch.Tensor): + raise TypeError("selected-logprob operator must return a torch.Tensor") + if selected.shape != token_ids.shape: + raise ValueError( + f"selected-logprob output shape {tuple(selected.shape)} must match " + f"token_ids shape {tuple(token_ids.shape)}" + ) + selected = selected.to(device=logits.device, dtype=output_dtype) + if mask is not None: + selected = selected.masked_fill(~mask, 0.0) + return selected + + +def _normalized_optional_backend(value: Optional[str]) -> Optional[str]: + if value is None: + return None + normalized = value.strip() + if not normalized: + raise ValueError("backend_id must not be empty") + return normalized + + +def _concrete_target(target: ConcreteOperatorTarget) -> ConcreteOperatorTarget: + normalized = target.strip().lower() + if normalized not in {"rollout", "training"}: + raise ValueError("target must be 'rollout' or 'training'") + return cast(ConcreteOperatorTarget, normalized) + + +__all__ = [ + "ConcreteOperatorTarget", + "OperatorBridge", + "OperatorOverride", + "OperatorTarget", + "ResolvedOperatorOverride", + "selected_logprobs_with_operator", +] diff --git a/rl_engine/alignment/cross_config/planner.py b/rl_engine/alignment/cross_config/planner.py new file mode 100644 index 00000000..476da5de --- /dev/null +++ b/rl_engine/alignment/cross_config/planner.py @@ -0,0 +1,573 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Typed baseline, one-at-a-time, and explicitly bounded pairwise planning.""" + +from __future__ import annotations + +import hashlib +import itertools +import json +from dataclasses import dataclass +from typing import Any, Callable, Mapping, Optional, Sequence + +from rl_engine.alignment.cross_config.schema import ( + ExperimentCase, + ExperimentDefinition, + IsolationScope, + KnobDescriptor, + PlanningStrategy, +) +from rl_engine.kernels.gtest.tolerance import tolerance_contract_fingerprint + +Normalizer = Callable[[Any], Any] +Constraint = Callable[[str, Any, Mapping[str, Any]], Optional["PlanningIssue"]] +MAX_PLAN_CASES = 256 + + +@dataclass(frozen=True) +class PlanningIssue: + """Structured planning rejection that callers can persist or display.""" + + code: str + reason: str + path: Optional[str] = None + value: Any = None + + def to_dict(self) -> dict[str, Any]: + return { + "code": self.code, + "reason": self.reason, + "path": self.path, + "value": self.value, + } + + +class PlanningError(ValueError): + """Raised for an invalid experiment definition with structured issues.""" + + def __init__(self, issues: Sequence[PlanningIssue]): + self.issues = tuple(issues) + message = "; ".join( + f"{issue.code}{f'[{issue.path}]' if issue.path else ''}: {issue.reason}" + for issue in self.issues + ) + super().__init__(message) + + +@dataclass(frozen=True) +class ExperimentPlan: + """A deterministic plan plus non-fatal capability findings.""" + + definition: ExperimentDefinition + cases: tuple[ExperimentCase, ...] + issues: tuple[PlanningIssue, ...] = () + + def to_dict(self) -> dict[str, Any]: + return { + "schema_version": "cross_config.experiment_plan.v1", + "experiment_id": self.definition.experiment_id, + "cases": [case.to_dict() for case in self.cases], + "issues": [issue.to_dict() for issue in self.issues], + } + + +def _positive_int(value: Any) -> int: + if isinstance(value, bool) or not isinstance(value, int) or value <= 0: + raise ValueError("must be a positive integer") + return value + + +def _strict_bool(value: Any) -> bool: + if not isinstance(value, bool): + raise ValueError("must be a JSON boolean") + return value + + +def _normalize_dtype(value: Any) -> str: + if not isinstance(value, str): + raise ValueError("must be a dtype string") + normalized = value.strip().lower().replace("torch.", "") + aliases = { + "bf16": "bfloat16", + "bfloat16": "bfloat16", + "fp16": "float16", + "half": "float16", + "float16": "float16", + "fp32": "float32", + "float": "float32", + "float32": "float32", + } + try: + return aliases[normalized] + except KeyError as exc: + raise ValueError(f"unsupported dtype {value!r}") from exc + + +def _normalize_choice(*choices: str) -> Normalizer: + allowed = frozenset(choices) + + def normalize(value: Any) -> str: + if not isinstance(value, str): + raise ValueError("must be a string") + normalized = value.strip().lower().replace("-", "_") + if normalized not in allowed: + raise ValueError(f"must be one of {sorted(allowed)}") + return normalized + + return normalize + + +def normalize_backend_id(value: Any) -> str: + """Normalize the public selected-logprob backend shortcut.""" + + if not isinstance(value, str) or not value.strip(): + raise ValueError("must be a non-empty backend ID") + normalized = value.strip().lower().replace("-", "_") + aliases = { + "auto": "native", + "default": "native", + "pytorch": "rlkernel.reference_logp", + "reference": "rlkernel.reference_logp", + } + return aliases.get(normalized, normalized) + + +V1_KNOB_DESCRIPTORS: tuple[KnobDescriptor, ...] = ( + KnobDescriptor("batch.size", IsolationScope.REQUEST, ("rollout", "training")), + KnobDescriptor("rollout.tensor_parallel_size", IsolationScope.PROCESS, ("rollout",)), + KnobDescriptor("rollout.context_parallel_size", IsolationScope.PROCESS, ("rollout",)), + KnobDescriptor("rollout.dtype", IsolationScope.ENGINE_CONSTRUCTION, ("rollout",)), + KnobDescriptor( + "rollout.enable_prefix_caching", + IsolationScope.ENGINE_CONSTRUCTION, + ("rollout",), + ), + KnobDescriptor("rollout.enforce_eager", IsolationScope.ENGINE_CONSTRUCTION, ("rollout",)), + KnobDescriptor( + "training.attention_backend", + IsolationScope.ENGINE_CONSTRUCTION, + ("training",), + allowed_values=("flash_attention_2", "sdpa", "eager", "model_default"), + ), + KnobDescriptor("training.compute_dtype", IsolationScope.ENGINE_CONSTRUCTION, ("training",)), + KnobDescriptor( + "logp.backend", + IsolationScope.ENGINE_CONSTRUCTION, + ("rollout", "training"), + derived=True, + ), + KnobDescriptor( + "training.sharding", + IsolationScope.PROCESS, + ("training",), + allowed_values=("unsharded", "fsdp"), + ), +) + +V1_KNOBS: Mapping[str, KnobDescriptor] = { + descriptor.path: descriptor for descriptor in V1_KNOB_DESCRIPTORS +} + +_NORMALIZERS: Mapping[str, Normalizer] = { + "batch.size": _positive_int, + "rollout.tensor_parallel_size": _positive_int, + "rollout.context_parallel_size": _positive_int, + "rollout.dtype": _normalize_dtype, + "rollout.enable_prefix_caching": _strict_bool, + "rollout.enforce_eager": _strict_bool, + "training.attention_backend": _normalize_choice( + "flash_attention_2", "sdpa", "eager", "model_default" + ), + "training.compute_dtype": _normalize_dtype, + "logp.backend": normalize_backend_id, + "training.sharding": _normalize_choice("unsharded", "fsdp"), +} + + +class Planner: + """Generate a bounded plan without importing runtime- or operator-specific branches.""" + + def __init__( + self, + *, + knobs: Mapping[str, KnobDescriptor] = V1_KNOBS, + normalizers: Mapping[str, Normalizer] = _NORMALIZERS, + constraints: Sequence[Constraint] = (), + ): + self.knobs = dict(knobs) + self.normalizers = dict(normalizers) + self.constraints = tuple(constraints) + + def plan(self, definition: ExperimentDefinition) -> ExperimentPlan: + issues = self._validate_definition(definition) + if issues: + raise PlanningError(issues) + + baseline = self.normalize_requested(definition.baseline) + requested_cases: list[tuple[dict[str, Any], tuple[str, ...]]] = [(baseline, ())] + intervention_values: dict[str, tuple[Any, ...]] = {} + + def append_requested(requested: dict[str, Any], changed_paths: tuple[str, ...]) -> None: + if len(requested_cases) >= MAX_PLAN_CASES: + raise PlanningError( + ( + PlanningIssue( + code="PLAN_TOO_LARGE", + reason=f"a plan may contain at most {MAX_PLAN_CASES} cases", + value=MAX_PLAN_CASES, + ), + ) + ) + requested_cases.append((requested, changed_paths)) + + for intervention in definition.interventions: + if len(intervention.values) > MAX_PLAN_CASES: + raise PlanningError( + ( + PlanningIssue( + code="PLAN_TOO_LARGE", + reason=(f"an intervention may contain at most {MAX_PLAN_CASES} values"), + path=intervention.path, + value=len(intervention.values), + ), + ) + ) + normalized_values = tuple( + self._normalize_value(intervention.path, value) for value in intervention.values + ) + intervention_values[intervention.path] = normalized_values + baseline_value = _get_path(baseline, intervention.path) + for value in normalized_values: + if value == baseline_value: + continue + requested = _deep_copy_mapping(baseline) + _set_path(requested, intervention.path, value) + append_requested(requested, (intervention.path,)) + + if definition.strategy == PlanningStrategy.PAIRWISE: + for first_path, second_path in definition.pairwise_paths: + first_baseline = _get_path(baseline, first_path) + second_baseline = _get_path(baseline, second_path) + for first_value, second_value in itertools.product( + intervention_values[first_path], intervention_values[second_path] + ): + if first_value == first_baseline or second_value == second_baseline: + continue + requested = _deep_copy_mapping(baseline) + _set_path(requested, first_path, first_value) + _set_path(requested, second_path, second_value) + append_requested(requested, tuple(sorted((first_path, second_path)))) + + contract_fingerprint = tolerance_contract_fingerprint() + scenario_fingerprint = _fingerprint( + {"scenario_id": definition.scenario_id, "scenario": definition.scenario} + ) + cases: list[ExperimentCase] = [] + seen_ids: set[str] = set() + capability_issues: list[PlanningIssue] = [] + for requested, changed_paths in requested_cases: + case_issues = self._apply_constraints(requested, changed_paths) + capability_issues.extend(case_issues) + case_id = self._case_id( + definition, + requested, + contract_fingerprint=contract_fingerprint, + scenario_fingerprint=scenario_fingerprint, + ) + if case_id in seen_ids: + continue + seen_ids.add(case_id) + cases.append( + ExperimentCase( + case_id=case_id, + experiment_id=definition.experiment_id, + scenario_id=definition.scenario_id, + identity=definition.identity, + requested=requested, + changed_paths=changed_paths, + contract_fingerprint=contract_fingerprint, + scenario_fingerprint=scenario_fingerprint, + ) + ) + + return ExperimentPlan( + definition=definition, + cases=tuple(cases), + issues=tuple(capability_issues), + ) + + def normalize_requested(self, requested: Mapping[str, Any]) -> dict[str, Any]: + flattened = _flatten(requested) + issues: list[PlanningIssue] = [] + normalized: dict[str, Any] = {} + for path, value in flattened.items(): + if path not in self.knobs: + code = "DERIVED_KNOB" if path == "logp.tp_layout" else "UNSUPPORTED_PATH" + issues.append( + PlanningIssue( + code=code, + path=path, + value=value, + reason="path is not a user-settable V1 knob", + ) + ) + continue + try: + normalized[path] = self._normalize_value(path, value) + except (TypeError, ValueError) as exc: + issues.append( + PlanningIssue( + code="UNSUPPORTED_VALUE", + path=path, + value=value, + reason=str(exc), + ) + ) + if issues: + raise PlanningError(issues) + result: dict[str, Any] = {} + for path, value in normalized.items(): + _set_path(result, path, value) + return result + + def isolation_for(self, changed_paths: Sequence[str]) -> IsolationScope: + if not changed_paths: + return IsolationScope.REQUEST + order = { + IsolationScope.REQUEST: 0, + IsolationScope.ENGINE_CONSTRUCTION: 1, + IsolationScope.DISTRIBUTED_CONTEXT: 2, + IsolationScope.PROCESS: 3, + } + return max((self.knobs[path].lifecycle for path in changed_paths), key=order.__getitem__) + + def _validate_definition(self, definition: ExperimentDefinition) -> list[PlanningIssue]: + issues: list[PlanningIssue] = [] + try: + baseline = self.normalize_requested(definition.baseline) + except PlanningError as exc: + return list(exc.issues) + baseline_paths = set(_flatten(baseline)) + for path in sorted(set(self.knobs).difference(baseline_paths)): + issues.append( + PlanningIssue( + code="MISSING_BASELINE_VALUE", + path=path, + reason="strict baselines must declare every allowlisted knob", + ) + ) + declared_paths: set[str] = set() + for intervention in definition.interventions: + path = intervention.path + if path not in self.knobs: + issues.append( + PlanningIssue( + code="UNSUPPORTED_PATH", + path=path, + reason="intervention path is not in the V1 allowlist", + ) + ) + continue + if path in declared_paths: + issues.append( + PlanningIssue( + code="DUPLICATE_INTERVENTION", + path=path, + reason="each intervention path must be declared once", + ) + ) + declared_paths.add(path) + if not intervention.values: + issues.append( + PlanningIssue( + code="EMPTY_INTERVENTION", + path=path, + reason="intervention values cannot be empty", + ) + ) + try: + _get_path(baseline, path) + except KeyError: + issues.append( + PlanningIssue( + code="MISSING_BASELINE_VALUE", + path=path, + reason="every intervention path must exist in baseline", + ) + ) + for value in intervention.values: + try: + self._normalize_value(path, value) + except (TypeError, ValueError) as exc: + issues.append( + PlanningIssue( + code="UNSUPPORTED_VALUE", + path=path, + value=value, + reason=str(exc), + ) + ) + + if definition.strategy == PlanningStrategy.ONE_AT_A_TIME and definition.pairwise_paths: + issues.append( + PlanningIssue( + code="PAIRWISE_NOT_ENABLED", + reason="pairwise_paths require strategy='pairwise'", + ) + ) + if definition.strategy == PlanningStrategy.PAIRWISE and not definition.pairwise_paths: + issues.append( + PlanningIssue( + code="PAIRWISE_PATHS_REQUIRED", + reason="pairwise strategy requires at least one explicit path pair", + ) + ) + seen_pairs: set[tuple[str, str]] = set() + for pair in definition.pairwise_paths: + if len(pair) != 2: + issues.append( + PlanningIssue( + code="INVALID_PAIR", + reason="each pairwise entry must contain exactly two paths", + value=pair, + ) + ) + continue + first, second = pair + canonical_pair = (first, second) if first < second else (second, first) + if first == second: + issues.append( + PlanningIssue( + code="INVALID_PAIR", + reason="pairwise paths must be distinct", + value=pair, + ) + ) + elif first not in declared_paths or second not in declared_paths: + issues.append( + PlanningIssue( + code="UNDECLARED_PAIR_PATH", + reason="pairwise paths must both have declared interventions", + value=pair, + ) + ) + elif canonical_pair in seen_pairs: + issues.append( + PlanningIssue( + code="DUPLICATE_PAIR", + reason="pairwise path pair is duplicated", + value=pair, + ) + ) + seen_pairs.add(canonical_pair) + return issues + + def _normalize_value(self, path: str, value: Any) -> Any: + try: + normalizer = self.normalizers[path] + except KeyError as exc: + raise ValueError(f"no normalizer registered for {path}") from exc + return normalizer(value) + + def _apply_constraints( + self, requested: Mapping[str, Any], changed_paths: Sequence[str] + ) -> list[PlanningIssue]: + issues: list[PlanningIssue] = [] + paths = changed_paths or tuple(_flatten(requested)) + for path in paths: + value = _get_path(requested, path) + for constraint in self.constraints: + issue = constraint(path, value, requested) + if issue is not None: + issues.append(issue) + return issues + + @staticmethod + def _case_id( + definition: ExperimentDefinition, + requested: Mapping[str, Any], + *, + contract_fingerprint: str, + scenario_fingerprint: str, + ) -> str: + payload = { + "requested": requested, + "identity": definition.identity.to_dict(), + "contract": { + "source": definition.contract_source, + "version": definition.contract_version, + "fingerprint": contract_fingerprint, + }, + "scenario_id": definition.scenario_id, + "scenario_fingerprint": scenario_fingerprint, + } + return f"cross-config-{_fingerprint(payload)[:20]}" + + +def _flatten(value: Mapping[str, Any], prefix: str = "") -> dict[str, Any]: + flattened: dict[str, Any] = {} + for key, child in value.items(): + if not isinstance(key, str) or not key: + raise PlanningError( + [PlanningIssue(code="INVALID_PATH", reason="configuration keys must be strings")] + ) + path = f"{prefix}.{key}" if prefix else key + if isinstance(child, Mapping): + flattened.update(_flatten(child, path)) + else: + flattened[path] = child + return flattened + + +def _get_path(value: Mapping[str, Any], path: str) -> Any: + current: Any = value + for part in path.split("."): + if not isinstance(current, Mapping) or part not in current: + raise KeyError(path) + current = current[part] + return current + + +def _set_path(value: dict[str, Any], path: str, child: Any) -> None: + current = value + parts = path.split(".") + for part in parts[:-1]: + existing = current.setdefault(part, {}) + if not isinstance(existing, dict): + raise ValueError(f"configuration path collision at {path}") + current = existing + current[parts[-1]] = child + + +def _deep_copy_mapping(value: Mapping[str, Any]) -> dict[str, Any]: + return json.loads(json.dumps(value)) + + +def _fingerprint(value: Mapping[str, Any]) -> str: + payload = json.dumps( + _json_plain(value), + sort_keys=True, + separators=(",", ":"), + ensure_ascii=False, + ).encode("utf-8") + return hashlib.sha256(payload).hexdigest() + + +def _json_plain(value: Any) -> Any: + if isinstance(value, Mapping): + return {str(key): _json_plain(child) for key, child in value.items()} + if isinstance(value, (tuple, list)): + return [_json_plain(child) for child in value] + return value + + +__all__ = [ + "ExperimentPlan", + "Planner", + "PlanningError", + "PlanningIssue", + "V1_KNOBS", + "V1_KNOB_DESCRIPTORS", + "normalize_backend_id", +] diff --git a/rl_engine/alignment/cross_config/runner.py b/rl_engine/alignment/cross_config/runner.py new file mode 100644 index 00000000..1d9418b2 --- /dev/null +++ b/rl_engine/alignment/cross_config/runner.py @@ -0,0 +1,714 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Control plane for paired read-only scoring runs.""" + +from __future__ import annotations + +import importlib +import math +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any, Mapping, Optional + +import torch + +from rl_engine.alignment.cross_config._execution import ( + ChildScoringError, + ChildSupervisor, + OperatorExecutionError, + PairedRunnerError, + PairedScorer, + RankCompletenessError, + RankScore, + ScorerIdentityError, + ScoringTimeoutError, + canonical_fingerprint, + device_type, + json_safe, + normalized_dtype, + paired_model_state_fingerprints, + scorer_implementation_fingerprint, + scorer_spec, + validate_rank_outputs, + validate_scorer_identity, +) +from rl_engine.alignment.cross_config._provenance import ( + PAIRED_RUNNER_IMPLEMENTATION_FINGERPRINT, + concrete_scorer_spec, + effective_runtime_status, + execution_environment_provenance, + execution_fingerprint, + factory_options_fingerprint, + mapping_target, + runtime_adapter_fingerprint, + score_metadata, + side_provenance, + target_factory_options, +) +from rl_engine.alignment.cross_config._resume import completed_attempt_matches, read_json_object +from rl_engine.alignment.cross_config.artifacts import ArtifactStore +from rl_engine.alignment.cross_config.comparison import compare_score_artifacts +from rl_engine.alignment.cross_config.operators import ResolvedOperatorOverride +from rl_engine.alignment.cross_config.runtime import ( + RuntimeMaterialization, + RuntimeMaterializationError, +) +from rl_engine.alignment.cross_config.schema import ( + AlignmentResult, + CanonicalScoringBatch, + ExperimentCase, + MaterializationStatus, + ScoreArtifact, + ScorerSpec, + ScoreSide, +) +from rl_engine.kernels.semantic_registry import ( + OperatorInstanceProvenance, + OperatorResolution, + operator_implementation_fingerprint, + operator_instance_fingerprint, +) + + +@dataclass(frozen=True) +class PairedRunResult: + """Completed attempt or a validated resume hit.""" + + case_id: str + attempt_id: str + attempt_dir: Path + resumed: bool + rollout_score: Optional[ScoreArtifact] = None + training_score: Optional[ScoreArtifact] = None + alignment: Optional[AlignmentResult] = None + summary: Mapping[str, Any] = field(default_factory=dict) + + +class PairedRunner: + """Supervise paired scorers and publish one append-only attempt.""" + + def __init__( + self, + artifact_store: ArtifactStore, + *, + timeout_seconds: float = 30.0, + start_method: Optional[str] = None, + ): + if not math.isfinite(timeout_seconds) or timeout_seconds <= 0.0: + raise ValueError("timeout_seconds must be finite and greater than zero") + self.artifact_store = artifact_store + self.timeout_seconds = float(timeout_seconds) + self._child_supervisor = ChildSupervisor(start_method) + self.start_method = self._child_supervisor.start_method + + @property + def active_child_pids(self) -> tuple[int, ...]: + return self._child_supervisor.active_child_pids + + def run( + self, + case: ExperimentCase, + materialization: RuntimeMaterialization, + batch: CanonicalScoringBatch, + rollout_scorer: PairedScorer, + training_scorer: PairedScorer, + resolved_override: ResolvedOperatorOverride, + operator_instances: Mapping[str | ScoreSide, Any], + operator_instance_provenance: Mapping[ + str | ScoreSide, + OperatorInstanceProvenance, + ], + *, + operator_factory_options: Optional[Mapping[str | ScoreSide, Mapping[str, Any]]] = None, + strict: bool = True, + timeout_seconds: Optional[float] = None, + resume: bool = True, + ) -> PairedRunResult: + """Run both sides against one canonical batch and persist the comparison.""" + + deadline_seconds = self.timeout_seconds if timeout_seconds is None else timeout_seconds + if not math.isfinite(deadline_seconds) or deadline_seconds <= 0.0: + raise ValueError("timeout_seconds must be finite and greater than zero") + self._validate_case_inputs(case, materialization, batch) + rollout_spec = scorer_spec(rollout_scorer, ScoreSide.ROLLOUT) + training_spec = scorer_spec(training_scorer, ScoreSide.TRAINING) + specs = {"rollout": rollout_spec, "training": training_spec} + validate_scorer_identity(specs, batch) + model_state_fingerprints = paired_model_state_fingerprints( + rollout_scorer, + training_scorer, + ) + scorer_implementation_fingerprints = { + "rollout": scorer_implementation_fingerprint(rollout_scorer), + "training": scorer_implementation_fingerprint(training_scorer), + } + resolutions, instances, instance_provenance = _validate_exact_operators( + materialization, + resolved_override, + operator_instances, + operator_instance_provenance, + operator_factory_options=operator_factory_options, + specs=specs, + strict=strict, + ) + environment = execution_environment_provenance( + specs, + runtime_adapter_fingerprint=runtime_adapter_fingerprint(materialization), + operator_implementation_fingerprints={ + target: instance_provenance[target].implementation_fingerprint + for target in ("rollout", "training") + }, + ) + environment_fingerprint = canonical_fingerprint(environment) + _require_materialization_executable(materialization, strict=strict) + current_execution_fingerprint = execution_fingerprint( + materialization, + specs=specs, + instance_provenance=instance_provenance, + operator_factory_options=operator_factory_options, + model_state_fingerprints=model_state_fingerprints, + scorer_implementation_fingerprints=scorer_implementation_fingerprints, + environment=environment, + ) + + if resume: + completed = self.artifact_store.completed_attempt(case.experiment_id, case.case_id) + if completed is not None and completed_attempt_matches( + completed, + case, + batch, + materialization=materialization, + specs=specs, + instance_provenance=instance_provenance, + operator_factory_options=operator_factory_options, + model_state_fingerprints=model_state_fingerprints, + scorer_implementation_fingerprints=scorer_implementation_fingerprints, + environment=environment, + execution_fingerprint=current_execution_fingerprint, + ): + summary = read_json_object(completed / "COMPLETE") + return PairedRunResult( + case_id=case.case_id, + attempt_id=completed.name, + attempt_dir=completed, + resumed=True, + summary=summary, + ) + + attempt_dir = self.artifact_store.create_attempt(case.experiment_id, case.case_id) + attempt_id = attempt_dir.name + self._write_attempt_inputs(attempt_dir, attempt_id, case, materialization, batch) + + child_results = self._child_supervisor.run( + attempt_dir, + batch, + batch_size=materialization.binding.batch_size, + scorers={ + "rollout": rollout_scorer, + "training": training_scorer, + }, + specs=specs, + instances=instances, + timeout_seconds=float(deadline_seconds), + ) + rollout_ranks = validate_rank_outputs( + child_results["rollout"], + rollout_spec, + expected_shape=batch.input_ids.shape, + target="rollout", + ) + training_ranks = validate_rank_outputs( + child_results["training"], + training_spec, + expected_shape=batch.input_ids.shape, + target="training", + ) + + rollout_provenance = side_provenance( + materialization.provenance, + resolutions["rollout"], + instance_provenance["rollout"], + child_results["rollout"], + rollout_spec, + status=effective_runtime_status(materialization), + factory_options=target_factory_options(operator_factory_options, "rollout"), + model_state_fingerprint=model_state_fingerprints["rollout"], + scorer_implementation_fingerprint=scorer_implementation_fingerprints["rollout"], + ) + training_provenance = side_provenance( + materialization.provenance, + resolutions["training"], + instance_provenance["training"], + child_results["training"], + training_spec, + status=effective_runtime_status(materialization), + factory_options=target_factory_options(operator_factory_options, "training"), + model_state_fingerprint=model_state_fingerprints["training"], + scorer_implementation_fingerprint=scorer_implementation_fingerprints["training"], + ) + rollout_artifact = ScoreArtifact( + case_id=case.case_id, + attempt_id=attempt_id, + side=ScoreSide.ROLLOUT, + identity=batch.identity, + scorer=concrete_scorer_spec( + rollout_spec, + instance_provenance["rollout"], + ), + selected_logprobs=rollout_ranks[0].selected_logprobs, + active_mask=batch.active_mask, + provenance=rollout_provenance, + ) + training_artifact = ScoreArtifact( + case_id=case.case_id, + attempt_id=attempt_id, + side=ScoreSide.TRAINING, + identity=batch.identity, + scorer=concrete_scorer_spec( + training_spec, + instance_provenance["training"], + ), + selected_logprobs=training_ranks[0].selected_logprobs, + active_mask=batch.active_mask, + provenance=training_provenance, + ) + alignment = compare_score_artifacts(rollout_artifact, training_artifact) + self._write_attempt_results( + attempt_dir, + rollout_artifact, + training_artifact, + alignment, + execution_fingerprint=current_execution_fingerprint, + environment=environment, + environment_fingerprint=environment_fingerprint, + ) + summary = { + "schema_version": "cross_config.complete.v1", + "case_id": case.case_id, + "attempt_id": attempt_id, + "status": alignment.status.value, + "comparable": alignment.comparable, + "passed": alignment.passed, + "active_token_count": alignment.active_token_count, + "mismatch_count": alignment.mismatch_count, + "worst_token_index": alignment.diagnostics.get("worst_token_index"), + "max_abs_diff": alignment.diagnostics.get("max_abs_diff"), + "rollout_backend": instance_provenance["rollout"].backend_id, + "training_backend": instance_provenance["training"].backend_id, + "execution_fingerprint": current_execution_fingerprint, + "environment_fingerprint": environment_fingerprint, + "runner_implementation_fingerprint": PAIRED_RUNNER_IMPLEMENTATION_FINGERPRINT, + } + marker = self.artifact_store.complete_attempt(attempt_dir, summary=summary) + summary = read_json_object(marker) + return PairedRunResult( + case_id=case.case_id, + attempt_id=attempt_id, + attempt_dir=attempt_dir, + resumed=False, + rollout_score=rollout_artifact, + training_score=training_artifact, + alignment=alignment, + summary=summary, + ) + + @staticmethod + def _validate_case_inputs( + case: ExperimentCase, + materialization: RuntimeMaterialization, + batch: CanonicalScoringBatch, + ) -> None: + materialized_case = materialization.materialized_case.case + if materialized_case != case: + raise ValueError("materialization case does not exactly match the requested case") + if batch.identity != case.identity: + raise ValueError("canonical scoring batch identity does not match the case identity") + if batch.input_ids.shape[0] < 1: + raise ValueError("canonical scoring batch must contain at least one sequence") + + def _write_attempt_inputs( + self, + attempt_dir: Path, + attempt_id: str, + case: ExperimentCase, + materialization: RuntimeMaterialization, + batch: CanonicalScoringBatch, + ) -> None: + envelope = {"case_id": case.case_id, "attempt_id": attempt_id} + self.artifact_store.write_json( + attempt_dir, + "requested", + {**envelope, "schema_version": "cross_config.requested.v1", "case": case.to_dict()}, + ) + self.artifact_store.write_json( + attempt_dir, + "materialized", + { + **envelope, + "schema_version": "cross_config.materialized_envelope.v1", + "materialized_case": materialization.materialized_case.to_dict(), + }, + ) + self.artifact_store.write_json( + attempt_dir, + "identity", + { + **envelope, + "schema_version": "cross_config.identity_envelope.v1", + "identity": batch.identity.to_dict(), + }, + ) + + def _write_attempt_results( + self, + attempt_dir: Path, + rollout: ScoreArtifact, + training: ScoreArtifact, + alignment: AlignmentResult, + *, + execution_fingerprint: str, + environment: Mapping[str, Any], + environment_fingerprint: str, + ) -> None: + self.artifact_store.write_json( + attempt_dir, + "actual", + { + "case_id": rollout.case_id, + "attempt_id": rollout.attempt_id, + "schema_version": "cross_config.actual.v1", + "execution_fingerprint": execution_fingerprint, + "environment": environment, + "environment_fingerprint": environment_fingerprint, + "runner_implementation_fingerprint": PAIRED_RUNNER_IMPLEMENTATION_FINGERPRINT, + "operator_source": "exact_resolution_and_instance", + "rollout": rollout.provenance.to_dict(), + "training": training.provenance.to_dict(), + }, + ) + self.artifact_store.write_tensor_bundle( + attempt_dir, + "score_rollout", + { + "selected_logprobs": rollout.selected_logprobs, + "active_mask": rollout.active_mask, + }, + metadata=score_metadata(rollout), + ) + self.artifact_store.write_tensor_bundle( + attempt_dir, + "score_training", + { + "selected_logprobs": training.selected_logprobs, + "active_mask": training.active_mask, + }, + metadata=score_metadata(training), + ) + self.artifact_store.write_json( + attempt_dir, + "comparison", + alignment.to_dict(), + ) + token_artifact = alignment.token_artifact + if token_artifact is None: + empty = torch.empty((0,), dtype=torch.float32) + token_tensors = { + "rollout_logprobs": empty, + "training_logprobs": empty, + "active_mask": torch.empty((0,), dtype=torch.bool), + "absolute_diff": empty, + "mismatch_mask": torch.empty((0,), dtype=torch.bool), + } + else: + token_tensors = { + "rollout_logprobs": token_artifact.rollout_logprobs, + "training_logprobs": token_artifact.training_logprobs, + "active_mask": token_artifact.active_mask, + "absolute_diff": token_artifact.absolute_diff, + "mismatch_mask": token_artifact.mismatch_mask, + } + self.artifact_store.write_tensor_bundle( + attempt_dir, + "token_diffs", + token_tensors, + metadata={ + "case_id": alignment.case_id, + "attempt_id": alignment.attempt_id, + "status": alignment.status.value, + "fixed_threshold": alignment.fixed_threshold, + }, + ) + + +def _validate_exact_operators( + materialization: RuntimeMaterialization, + resolved: ResolvedOperatorOverride, + instances: Mapping[str | ScoreSide, Any], + instance_provenance: Mapping[str | ScoreSide, OperatorInstanceProvenance], + *, + operator_factory_options: Optional[Mapping[str | ScoreSide, Mapping[str, Any]]], + specs: Mapping[str, ScorerSpec], + strict: bool, +) -> tuple[ + dict[str, OperatorResolution], + dict[str, Any], + dict[str, OperatorInstanceProvenance], +]: + if resolved.semantic_op != "selected_logprob": + raise OperatorExecutionError("PairedRunner V1 requires semantic_op='selected_logprob'") + resolutions: dict[str, OperatorResolution] = {} + concrete_instances: dict[str, Any] = {} + provenance: dict[str, OperatorInstanceProvenance] = {} + for target in ("rollout", "training"): + resolution = resolved.for_target(target) # type: ignore[arg-type] + if resolution is None: + raise OperatorExecutionError(f"missing exact {target} operator resolution") + if resolution.target != target: + raise OperatorExecutionError( + f"{target} operator resolution reports target={resolution.target!r}" + ) + if ( + resolution.descriptor.semantic_op != "selected_logprob" + or resolution.trace.semantic_op != "selected_logprob" + ): + raise OperatorExecutionError(f"{target} resolution does not describe selected_logprob") + if resolution.trace.status != "resolved" or resolution.trace.concrete_backend is None: + raise OperatorExecutionError( + f"{target} operator is not exactly observable: {resolution.trace.status}" + ) + if resolution.trace.fallback_attempts: + raise OperatorExecutionError(f"{target} operator resolution attempted fallback") + if strict and not resolution.strict: + raise OperatorExecutionError(f"{target} operator was not resolved in strict mode") + if resolution.trace.concrete_backend != resolution.descriptor.backend_id: + raise OperatorExecutionError(f"{target} resolution backend evidence is inconsistent") + if resolution.trace.descriptor_fingerprint != resolution.descriptor.descriptor_fingerprint: + raise OperatorExecutionError( + f"{target} resolution descriptor fingerprint is inconsistent" + ) + if device_type(resolution.requirements.device) != device_type(specs[target].device): + raise OperatorExecutionError( + f"{target} operator resolution device does not match scorer device" + ) + if normalized_dtype(resolution.requirements.dtype) != normalized_dtype(specs[target].dtype): + raise OperatorExecutionError(f"{target} resolution dtype does not match scorer dtype") + _validate_exact_topology( + materialization, + resolution, + specs[target], + target=target, + ) + instance = mapping_target(instances, target) + if instance is None: + raise OperatorExecutionError(f"missing instantiated {target} operator") + _require_instance_matches_resolution(resolution, instance, target=target) + instance_evidence = mapping_target(instance_provenance, target) + if not isinstance(instance_evidence, OperatorInstanceProvenance): + raise OperatorExecutionError(f"missing sealed {target} operator instance provenance") + _validate_instance_provenance( + resolution, + instance, + instance_evidence, + factory_options=target_factory_options(operator_factory_options, target), + target=target, + ) + if instance_evidence.backend_id != resolution.trace.concrete_backend: + raise OperatorExecutionError(f"{target} instance backend does not match resolution") + declared = materialization.binding.operator_backends.get(target) + if declared is not None and declared != instance_evidence.backend_id: + raise OperatorExecutionError( + f"{target} exact backend {instance_evidence.backend_id!r} does not match " + f"declared override {declared!r}" + ) + resolutions[target] = resolution + concrete_instances[target] = instance + provenance[target] = instance_evidence + requested_logp = materialization.materialized_case.case.requested.get("logp") + requested_backend = ( + requested_logp.get("backend") if isinstance(requested_logp, Mapping) else None + ) + if requested_backend != provenance["rollout"].backend_id: + raise OperatorExecutionError( + "exact rollout operator does not match the public logp.backend request: " + f"{provenance['rollout'].backend_id!r} != {requested_backend!r}" + ) + return resolutions, concrete_instances, provenance + + +def _require_materialization_executable( + materialization: RuntimeMaterialization, + *, + strict: bool, +) -> None: + rejected = { + MaterializationStatus.ERROR, + MaterializationStatus.UNSUPPORTED, + MaterializationStatus.UNOBSERVABLE, + } + if strict: + rejected.add(MaterializationStatus.FALLBACK) + if not materialization.applications and materialization.materialized_case.status in rejected: + raise RuntimeMaterializationError( + f"case {materialization.materialized_case.case.case_id} is not executable: " + f"materialization={materialization.materialized_case.status.value}" + ) + problems = [] + for application in materialization.applications: + if ( + application.path == "logp.backend" + and application.status is MaterializationStatus.UNOBSERVABLE + ): + continue + if application.status in rejected: + problems.append( + f"{application.path}={application.status.value}: " + f"{application.evidence.get('reason', 'no evidence')}" + ) + if problems: + raise RuntimeMaterializationError( + f"case {materialization.materialized_case.case.case_id} is not executable: " + + "; ".join(problems) + ) + + +def _validate_exact_topology( + materialization: RuntimeMaterialization, + resolution: OperatorResolution, + spec: ScorerSpec, + *, + target: str, +) -> None: + bound_topology = mapping_target(materialization.binding.topology, target) + if not isinstance(bound_topology, Mapping): + raise OperatorExecutionError(f"{target} materialized topology is missing") + expected = dict(bound_topology) + topology_paths = { + "rollout": ( + ("rollout.tensor_parallel_size", "tensor_parallel_size"), + ("rollout.context_parallel_size", "context_parallel_size"), + ), + "training": (("training.sharding", "sharding"),), + } + required_keys = {"world_size", *(key for _, key in topology_paths[target])} + missing_keys = sorted(required_keys.difference(expected)) + if missing_keys: + raise OperatorExecutionError( + f"{target} materialized topology is missing required keys: {missing_keys!r}" + ) + if expected.get("world_size") != spec.world_size: + raise OperatorExecutionError( + f"{target} scorer world_size does not match materialized topology" + ) + if dict(resolution.requirements.topology) != expected: + raise OperatorExecutionError( + f"{target} resolution topology does not match materialized topology" + ) + if dict(spec.topology) != expected: + raise OperatorExecutionError( + f"{target} scorer topology does not match materialized topology" + ) + + actual_by_path = { + application.path: application.actual for application in materialization.applications + } + for path, key in topology_paths[target]: + if actual_by_path.get(path) != expected[key]: + raise OperatorExecutionError( + f"{target} actual {path} does not match exact operator topology" + ) + + +def _require_instance_matches_resolution( + resolution: OperatorResolution, + instance: Any, + *, + target: str, +) -> None: + implementation = resolution.descriptor.implementation_class_or_factory + factory = implementation + if isinstance(implementation, str): + try: + module_name, object_name = implementation.rsplit(".", 1) + factory = getattr(importlib.import_module(module_name), object_name) + except (ValueError, ImportError, AttributeError, ModuleNotFoundError) as exc: + raise OperatorExecutionError( + f"{target} exact operator factory cannot be verified: {exc}" + ) from exc + if isinstance(factory, type) and not isinstance(instance, factory): + raise OperatorExecutionError( + f"{target} operator instance type {type(instance).__qualname__!r} " + f"does not match resolved factory {factory.__qualname__!r}" + ) + if not callable(instance) and not callable(getattr(instance, "apply_fp32", None)): + raise OperatorExecutionError( + f"{target} selected-logprob operator instance is not executable" + ) + + +def _validate_instance_provenance( + resolution: OperatorResolution, + instance: Any, + provenance: OperatorInstanceProvenance, + *, + factory_options: Mapping[str, Any], + target: str, +) -> None: + expected_concrete = f"{type(instance).__module__}.{type(instance).__qualname__}" + expected_factory = resolution.descriptor.implementation_reference + if expected_factory is None: + raise OperatorExecutionError(f"{target} resolved operator has no factory reference") + implementation = resolution.descriptor.implementation_class_or_factory + if implementation is None: + raise OperatorExecutionError(f"{target} resolved operator has no implementation") + mismatches = [] + if provenance.semantic_op != resolution.descriptor.semantic_op: + mismatches.append("semantic_op") + if provenance.backend_id != resolution.descriptor.backend_id: + mismatches.append("backend_id") + if provenance.target != target: + mismatches.append("target") + if provenance.factory_reference != expected_factory: + mismatches.append("factory_reference") + if provenance.concrete_implementation != expected_concrete: + mismatches.append("concrete_implementation") + if provenance.descriptor_fingerprint != resolution.descriptor.descriptor_fingerprint: + mismatches.append("descriptor_fingerprint") + observed_implementation_fingerprint = operator_implementation_fingerprint( + implementation, + instance, + ) + if provenance.implementation_fingerprint != observed_implementation_fingerprint: + mismatches.append("implementation_fingerprint") + + if json_safe(provenance.factory_options) != json_safe(factory_options): + mismatches.append("factory_options") + if provenance.factory_options_fingerprint != factory_options_fingerprint(factory_options): + mismatches.append("factory_options_fingerprint") + expected_instance_fingerprint = operator_instance_fingerprint( + descriptor_fingerprint=resolution.descriptor.descriptor_fingerprint, + factory_reference=expected_factory, + concrete_implementation=expected_concrete, + implementation_fingerprint=observed_implementation_fingerprint, + factory_options_fingerprint=factory_options_fingerprint(factory_options), + ) + if provenance.instance_fingerprint != expected_instance_fingerprint: + mismatches.append("instance_fingerprint") + if mismatches: + raise OperatorExecutionError( + f"{target} operator instance provenance is inconsistent: " + ", ".join(mismatches) + ) + + +__all__ = [ + "ChildScoringError", + "OperatorExecutionError", + "PairedRunResult", + "PairedRunner", + "PairedRunnerError", + "PairedScorer", + "RankCompletenessError", + "RankScore", + "ScorerIdentityError", + "ScoringTimeoutError", +] diff --git a/rl_engine/alignment/cross_config/runtime.py b/rl_engine/alignment/cross_config/runtime.py new file mode 100644 index 00000000..3eab0352 --- /dev/null +++ b/rl_engine/alignment/cross_config/runtime.py @@ -0,0 +1,465 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Thin runtime materialization facade for the V1 allowlist.""" + +from __future__ import annotations + +import hashlib +import json +from dataclasses import dataclass, field +from types import MappingProxyType +from typing import Any, Iterable, Mapping, Protocol, Sequence + +from rl_engine.alignment.cross_config.planner import V1_KNOBS +from rl_engine.alignment.cross_config.schema import ( + ExperimentCase, + IsolationScope, + KnobDescriptor, + MaterializationStatus, + MaterializedCase, + RuntimeProvenance, + SerializableModel, +) + + +@dataclass(frozen=True) +class KnobApplication(SerializableModel): + """One adapter's requested, materialized, and observed value.""" + + path: str + requested: Any + materialized: Any + actual: Any + lifecycle: IsolationScope + status: MaterializationStatus + evidence: Mapping[str, Any] = field(default_factory=dict) + critical: bool = True + schema_version: str = "cross_config.knob_application.v1" + + def __post_init__(self) -> None: + object.__setattr__(self, "lifecycle", IsolationScope(self.lifecycle)) + object.__setattr__(self, "status", MaterializationStatus(self.status)) + for name in ("requested", "materialized", "actual"): + object.__setattr__(self, name, _freeze_value(getattr(self, name))) + object.__setattr__(self, "evidence", _freeze_mapping(self.evidence)) + + +@dataclass(frozen=True) +class RuntimeBinding: + """Small, backend-neutral handoff from materialization to execution. + + Runtime adapters may construct repository-specific objects internally, but + the core runner sees only the values required to create scorers and validate + lifecycle identity. New vLLM, FSDP, or other adapters therefore do not + change the runner's type surface. + """ + + batch_size: int + side_configs: Mapping[str, Mapping[str, Any]] + topology: Mapping[str, Mapping[str, Any]] + scorer: Mapping[str, Any] + operator_backends: Mapping[str, str] + runtime_kind: str + + def __post_init__(self) -> None: + if isinstance(self.batch_size, bool) or not isinstance(self.batch_size, int): + raise TypeError("batch_size must be an integer") + if self.batch_size < 1: + raise ValueError("batch_size must be greater than zero") + if not isinstance(self.runtime_kind, str) or not self.runtime_kind.strip(): + raise ValueError("runtime_kind must be a non-empty string") + for name, value in ( + ("side_configs", self.side_configs), + ("topology", self.topology), + ): + for target in ("rollout", "training"): + if not isinstance(value.get(target), Mapping): + raise ValueError(f"{name} must define a {target} mapping") + for target in ("rollout", "training"): + world_size = self.topology[target].get("world_size") + if isinstance(world_size, bool) or not isinstance(world_size, int) or world_size < 1: + raise ValueError(f"{target} topology must define a positive integer world_size") + backend = self.operator_backends.get(target) + if not isinstance(backend, str) or not backend.strip(): + raise ValueError(f"operator_backends must define a non-empty {target} backend") + object.__setattr__(self, "side_configs", _freeze_mapping(self.side_configs)) + object.__setattr__(self, "topology", _freeze_mapping(self.topology)) + object.__setattr__(self, "scorer", _freeze_mapping(self.scorer)) + object.__setattr__(self, "operator_backends", _freeze_mapping(self.operator_backends)) + + def to_dict(self) -> dict[str, Any]: + return { + "schema_version": "cross_config.runtime_binding.v1", + "runtime_kind": self.runtime_kind, + "batch_size": self.batch_size, + "side_configs": _plain_mapping(self.side_configs), + "topology": _plain_mapping(self.topology), + "scorer": _plain_mapping(self.scorer), + "operators": dict(self.operator_backends), + } + + +@dataclass(frozen=True) +class AdapterMaterialization: + """Output of a typed runtime adapter before the facade adds fingerprints.""" + + applications: tuple[KnobApplication, ...] + binding: RuntimeBinding + + def __post_init__(self) -> None: + object.__setattr__(self, "applications", tuple(self.applications)) + + +class RuntimeMaterializer(Protocol): + """Adapter boundary used by the small ``RuntimeTools`` facade. + + The declared implementation fingerprint must deterministically identify the + executable materialization path and change when that implementation changes. + """ + + runtime_kind: str + + @property + def implementation_fingerprint(self) -> str: ... + + def materialize( + self, + normalized: Mapping[str, Any], + descriptors: Mapping[str, KnobDescriptor], + ) -> AdapterMaterialization: ... + + +@dataclass(frozen=True) +class RuntimeMaterialization: + materialized_case: MaterializedCase + provenance: RuntimeProvenance + applications: tuple[KnobApplication, ...] + binding: RuntimeBinding + + @property + def executable_in_strict_mode(self) -> bool: + return self.materialized_case.status is MaterializationStatus.APPLIED + + +class RuntimeMaterializationError(RuntimeError): + pass + + +class RuntimeTools: + """Materialize cases and compute reuse fingerprints without owning execution.""" + + def __init__(self, descriptors: Mapping[str, KnobDescriptor] = V1_KNOBS): + self.descriptors = dict(descriptors) + + def materialize( + self, + case: ExperimentCase, + adapter: RuntimeMaterializer, + ) -> RuntimeMaterialization: + runtime_kind = _adapter_identity(adapter, "runtime_kind") + adapter_implementation_fingerprint = _adapter_identity( + adapter, + "implementation_fingerprint", + ) + normalized = _plain_mapping(case.requested) + adapter_result = adapter.materialize(normalized, self.descriptors) + if not isinstance(adapter_result, AdapterMaterialization): + raise RuntimeMaterializationError("runtime adapter must return AdapterMaterialization") + if not isinstance(adapter_result.binding, RuntimeBinding): + raise RuntimeMaterializationError("runtime adapter must return a RuntimeBinding") + if adapter_result.binding.runtime_kind != runtime_kind: + raise RuntimeMaterializationError( + "runtime binding kind must match the materializer runtime_kind" + ) + applications = tuple(adapter_result.applications) + _validate_application_contract(normalized, applications, self.descriptors) + materialized = _mapping_from_applications(applications, "materialized") + actual = _mapping_from_applications(applications, "actual") + status = _aggregate_status(application.status for application in applications) + construction_fingerprint = _scope_fingerprint( + runtime_kind, + adapter_implementation_fingerprint, + applications, + scopes=( + IsolationScope.ENGINE_CONSTRUCTION, + IsolationScope.DISTRIBUTED_CONTEXT, + IsolationScope.PROCESS, + ), + ) + distributed_fingerprint = _scope_fingerprint( + runtime_kind, + adapter_implementation_fingerprint, + applications, + scopes=(IsolationScope.DISTRIBUTED_CONTEXT, IsolationScope.PROCESS), + ) + process_fingerprint = _scope_fingerprint( + runtime_kind, + adapter_implementation_fingerprint, + applications, + scopes=(IsolationScope.PROCESS,), + ) + isolation_scope = _strongest_scope( + [ + self.descriptors[path].lifecycle + for path in (case.changed_paths or tuple(_flatten(normalized))) + ] + ) + evidence = { + "runtime_kind": runtime_kind, + "execution_binding": case.execution_binding, + "adapter_implementation_fingerprint": adapter_implementation_fingerprint, + "binding_fingerprint": _fingerprint(adapter_result.binding.to_dict()), + "applications": { + application.path: application.to_dict() for application in applications + }, + } + materialized_case = MaterializedCase( + case=case, + normalized=normalized, + materialized=materialized, + isolation_scope=isolation_scope, + construction_fingerprint=construction_fingerprint, + distributed_context_fingerprint=distributed_fingerprint, + process_fingerprint=process_fingerprint, + status=status, + evidence=evidence, + ) + provenance = RuntimeProvenance( + requested=_plain_mapping(case.requested), + normalized=normalized, + materialized=materialized, + actual=actual, + status=status, + construction_fingerprint=construction_fingerprint, + distributed_context_fingerprint=distributed_fingerprint, + process_fingerprint=process_fingerprint, + implementation_fingerprint=adapter_implementation_fingerprint, + evidence=evidence, + ) + return RuntimeMaterialization( + materialized_case=materialized_case, + provenance=provenance, + applications=applications, + binding=adapter_result.binding, + ) + + @staticmethod + def require_executable( + materialization: RuntimeMaterialization, + *, + strict: bool, + ) -> None: + status = materialization.materialized_case.status + rejected = { + MaterializationStatus.ERROR, + MaterializationStatus.UNSUPPORTED, + MaterializationStatus.UNOBSERVABLE, + } + if strict: + rejected.add(MaterializationStatus.FALLBACK) + if status in rejected: + problems = [ + f"{application.path}={application.status.value}: " + f"{application.evidence.get('reason', 'no evidence')}" + for application in materialization.applications + if application.status is not MaterializationStatus.APPLIED + ] + raise RuntimeMaterializationError( + f"case {materialization.materialized_case.case.case_id} is not executable: " + + "; ".join(problems) + ) + + @staticmethod + def can_reuse(previous: RuntimeMaterialization, current: RuntimeMaterialization) -> bool: + """Reuse only exact semantic and implementation identities with matching state.""" + + previous_case = previous.materialized_case + current_case = current.materialized_case + return ( + previous_case.status is MaterializationStatus.APPLIED + and current_case.status is MaterializationStatus.APPLIED + and previous_case.case.identity == current_case.case.identity + and previous_case.case.execution_binding == current_case.case.execution_binding + and previous.provenance.implementation_fingerprint + == current.provenance.implementation_fingerprint + and previous_case.process_fingerprint == current_case.process_fingerprint + and previous_case.distributed_context_fingerprint + == current_case.distributed_context_fingerprint + and previous_case.construction_fingerprint == current_case.construction_fingerprint + ) + + +def _aggregate_status(statuses: Iterable[MaterializationStatus]) -> MaterializationStatus: + priority = ( + MaterializationStatus.ERROR, + MaterializationStatus.UNSUPPORTED, + MaterializationStatus.UNOBSERVABLE, + MaterializationStatus.FALLBACK, + MaterializationStatus.APPLIED, + ) + status_set = set(statuses) + if not status_set: + return MaterializationStatus.ERROR + return next(status for status in priority if status in status_set) + + +def _validate_application_contract( + normalized: Mapping[str, Any], + applications: tuple[KnobApplication, ...], + descriptors: Mapping[str, KnobDescriptor], +) -> None: + expected = _flatten(normalized) + observed_paths = [application.path for application in applications] + duplicate_paths = sorted(path for path in set(observed_paths) if observed_paths.count(path) > 1) + missing_paths = sorted(set(expected).difference(observed_paths)) + unknown_paths = sorted(set(observed_paths).difference(expected)) + missing_descriptors = sorted(set(expected).difference(descriptors)) + problems: list[str] = [] + if missing_paths: + problems.append(f"missing paths={missing_paths!r}") + if duplicate_paths: + problems.append(f"duplicate paths={duplicate_paths!r}") + if unknown_paths: + problems.append(f"unknown paths={unknown_paths!r}") + if missing_descriptors: + problems.append(f"missing descriptors={missing_descriptors!r}") + for application in applications: + descriptor = descriptors.get(application.path) + if descriptor is None or application.path not in expected: + continue + if _plain_value(application.requested) != _plain_value(expected[application.path]): + problems.append(f"{application.path} requested value differs from normalized case") + if application.lifecycle is not descriptor.lifecycle: + problems.append(f"{application.path} lifecycle differs from descriptor") + if application.critical is not descriptor.critical: + problems.append(f"{application.path} critical flag differs from descriptor") + if application.status is MaterializationStatus.APPLIED: + if _plain_value(application.actual) != _plain_value(application.materialized): + problems.append(f"{application.path} applied actual differs from materialized") + if not descriptor.derived and _plain_value(application.materialized) != _plain_value( + expected[application.path] + ): + problems.append( + f"{application.path} applied materialized value differs from normalized case" + ) + if problems: + raise RuntimeMaterializationError( + "runtime adapter returned invalid V1 knob applications: " + "; ".join(problems) + ) + + +def _mapping_from_applications( + applications: Sequence[KnobApplication], attribute: str +) -> dict[str, Any]: + result: dict[str, Any] = {} + for application in applications: + _set_path(result, application.path, getattr(application, attribute)) + return result + + +def _scope_fingerprint( + runtime_kind: str, + implementation_fingerprint: str, + applications: Sequence[KnobApplication], + *, + scopes: Sequence[IsolationScope], +) -> str: + scope_set = set(scopes) + values = { + application.path: application.materialized + for application in applications + if application.lifecycle in scope_set + } + return _fingerprint( + { + "runtime_kind": runtime_kind, + "implementation_fingerprint": implementation_fingerprint, + "values": values, + } + ) + + +def _strongest_scope(scopes: Sequence[IsolationScope]) -> IsolationScope: + order = { + IsolationScope.REQUEST: 0, + IsolationScope.ENGINE_CONSTRUCTION: 1, + IsolationScope.DISTRIBUTED_CONTEXT: 2, + IsolationScope.PROCESS: 3, + } + return max(scopes, key=order.__getitem__, default=IsolationScope.REQUEST) + + +def _flatten(value: Mapping[str, Any], prefix: str = "") -> dict[str, Any]: + result: dict[str, Any] = {} + for key, child in value.items(): + path = f"{prefix}.{key}" if prefix else key + if isinstance(child, Mapping): + result.update(_flatten(child, path)) + else: + result[path] = child + return result + + +def _set_path(value: dict[str, Any], path: str, child: Any) -> None: + current = value + parts = path.split(".") + for part in parts[:-1]: + current = current.setdefault(part, {}) + current[parts[-1]] = child + + +def _plain_mapping(value: Mapping[str, Any]) -> dict[str, Any]: + return {str(key): _plain_value(item) for key, item in value.items()} + + +def _plain_value(value: Any) -> Any: + if isinstance(value, Mapping): + return {str(key): _plain_value(item) for key, item in value.items()} + if isinstance(value, (tuple, list)): + return [_plain_value(item) for item in value] + if isinstance(value, (set, frozenset)): + return [_plain_value(item) for item in sorted(value, key=repr)] + return value + + +def _freeze_value(value: Any) -> Any: + if isinstance(value, Mapping): + return MappingProxyType({str(key): _freeze_value(item) for key, item in value.items()}) + if isinstance(value, (tuple, list)): + return tuple(_freeze_value(item) for item in value) + if isinstance(value, (set, frozenset)): + return frozenset(_freeze_value(item) for item in value) + return value + + +def _freeze_mapping(value: Mapping[str, Any]) -> Mapping[str, Any]: + return _freeze_value(value) + + +def _fingerprint(value: Mapping[str, Any]) -> str: + payload = json.dumps( + _plain_mapping(value), + sort_keys=True, + separators=(",", ":"), + ensure_ascii=False, + ).encode("utf-8") + return hashlib.sha256(payload).hexdigest() + + +def _adapter_identity(adapter: RuntimeMaterializer, attribute: str) -> str: + value = getattr(adapter, attribute, None) + if not isinstance(value, str) or not value.strip(): + raise RuntimeMaterializationError(f"runtime adapter {attribute} must be a non-empty string") + return value.strip() + + +__all__ = [ + "AdapterMaterialization", + "KnobApplication", + "RuntimeBinding", + "RuntimeMaterialization", + "RuntimeMaterializationError", + "RuntimeMaterializer", + "RuntimeTools", +] diff --git a/rl_engine/alignment/cross_config/schema.py b/rl_engine/alignment/cross_config/schema.py new file mode 100644 index 00000000..27f668db --- /dev/null +++ b/rl_engine/alignment/cross_config/schema.py @@ -0,0 +1,583 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Stable, versioned domain schema for cross-configuration alignment.""" + +from __future__ import annotations + +import json +import math +from dataclasses import dataclass, field, fields +from enum import Enum +from pathlib import Path +from types import MappingProxyType +from typing import Any, Mapping, Optional, Sequence + +import torch + + +class ScoreSide(str, Enum): + ROLLOUT = "rollout" + TRAINING = "training" + + +class IsolationScope(str, Enum): + REQUEST = "request" + ENGINE_CONSTRUCTION = "engine_construction" + DISTRIBUTED_CONTEXT = "distributed_context" + PROCESS = "process" + + +class PlanningStrategy(str, Enum): + ONE_AT_A_TIME = "one_at_a_time" + PAIRWISE = "pairwise" + + +class MaterializationStatus(str, Enum): + UNSUPPORTED = "unsupported" + APPLIED = "applied" + FALLBACK = "fallback" + UNOBSERVABLE = "unobservable" + ERROR = "error" + + +class AlignmentStatus(str, Enum): + PASS = "pass" + FAIL = "fail" + INVALID_IDENTITY = "invalid_identity" + INVALID_ARTIFACT = "invalid_artifact" + ZERO_ACTIVE_TOKENS = "zero_active_tokens" + + +class SerializableModel: + """Mixin providing a stable JSON-compatible representation.""" + + def to_dict(self) -> dict[str, Any]: + return { + item.name: _serialize_value(getattr(self, item.name)) + for item in fields(self) # type: ignore[arg-type] + } + + def to_json(self, *, indent: Optional[int] = None) -> str: + return json.dumps(self.to_dict(), indent=indent, sort_keys=True) + + +def _serialize_value(value: Any) -> Any: + if isinstance(value, Enum): + return value.value + if isinstance(value, SerializableModel): + return value.to_dict() + if isinstance(value, torch.Tensor): + snapshot = value.detach().cpu() + return { + "dtype": str(snapshot.dtype).replace("torch.", ""), + "shape": list(snapshot.shape), + "values": snapshot.tolist(), + } + if isinstance(value, Mapping): + return {str(key): _serialize_value(item) for key, item in value.items()} + if isinstance(value, (tuple, list)): + return [_serialize_value(item) for item in value] + if isinstance(value, (set, frozenset)): + return [_serialize_value(item) for item in sorted(value, key=repr)] + if isinstance(value, Path): + return str(value) + if isinstance(value, torch.dtype): + return str(value).replace("torch.", "") + return value + + +def _freeze_value(value: Any) -> Any: + if isinstance(value, Mapping): + return MappingProxyType({str(key): _freeze_value(item) for key, item in value.items()}) + if isinstance(value, (tuple, list)): + return tuple(_freeze_value(item) for item in value) + if isinstance(value, (set, frozenset)): + return tuple(sorted((_freeze_value(item) for item in value), key=repr)) + return value + + +def _freeze_mapping(value: Mapping[str, Any]) -> Mapping[str, Any]: + return _freeze_value(value) + + +def _coerce_enum(value: Any, enum_type: type[Enum]) -> Enum: + if isinstance(value, enum_type): + return value + return enum_type(value) + + +def _int_matrix(value: Sequence[Sequence[int]]) -> tuple[tuple[int, ...], ...]: + rows: list[tuple[int, ...]] = [] + for row in value: + normalized: list[int] = [] + for item in row: + if isinstance(item, bool) or not isinstance(item, int): + raise ValueError("integer identity matrices accept JSON integers only") + normalized.append(item) + rows.append(tuple(normalized)) + return tuple(rows) + + +def _bool_matrix(value: Sequence[Sequence[bool]]) -> tuple[tuple[bool, ...], ...]: + rows: list[tuple[bool, ...]] = [] + for row in value: + normalized: list[bool] = [] + for item in row: + if not isinstance(item, bool): + raise ValueError("boolean identity matrices accept JSON booleans only") + normalized.append(item) + rows.append(tuple(normalized)) + return tuple(rows) + + +def _validate_rectangular(name: str, value: tuple[tuple[Any, ...], ...]) -> None: + if not value: + return + width = len(value[0]) + if any(len(row) != width for row in value): + raise ValueError(f"{name} must be rectangular") + + +def _validate_same_matrix_shape( + left_name: str, + left: tuple[tuple[Any, ...], ...], + right_name: str, + right: tuple[tuple[Any, ...], ...], +) -> None: + if left and right and (len(left), len(left[0])) != (len(right), len(right[0])): + raise ValueError(f"{left_name} shape must match {right_name} shape") + + +def _snapshot_tensor(value: torch.Tensor, *, dtype: Optional[torch.dtype] = None) -> torch.Tensor: + if not isinstance(value, torch.Tensor): + raise TypeError(f"expected torch.Tensor, got {type(value)!r}") + snapshot = value.detach().clone() + return snapshot.to(dtype=dtype) if dtype is not None else snapshot + + +@dataclass(frozen=True) +class SemanticIdentitySpec(SerializableModel): + """Logical inputs that must match before numerical comparison is meaningful.""" + + checkpoint_id: str + model_version: str + tokenizer_policy: str + token_ids: tuple[tuple[int, ...], ...] + selected_token_ids: tuple[tuple[int, ...], ...] + active_mask: tuple[tuple[bool, ...], ...] + pre_update_state: str + tokenizer_id: str = "" + attention_mask: tuple[tuple[bool, ...], ...] = () + position_ids: tuple[tuple[int, ...], ...] = () + cache_metadata: Mapping[str, Any] = field(default_factory=dict) + packing_metadata: Mapping[str, Any] = field(default_factory=dict) + schema_version: str = "cross_config.semantic_identity.v1" + + def __post_init__(self) -> None: + if self.schema_version != "cross_config.semantic_identity.v1": + raise ValueError("unsupported SemanticIdentitySpec schema_version") + if not self.checkpoint_id: + raise ValueError("checkpoint_id must not be empty") + if not self.model_version: + raise ValueError("model_version must not be empty") + if not self.tokenizer_policy: + raise ValueError("tokenizer_policy must not be empty") + if not self.pre_update_state: + raise ValueError("pre_update_state must not be empty") + + object.__setattr__(self, "token_ids", _int_matrix(self.token_ids)) + object.__setattr__(self, "selected_token_ids", _int_matrix(self.selected_token_ids)) + object.__setattr__(self, "active_mask", _bool_matrix(self.active_mask)) + object.__setattr__(self, "attention_mask", _bool_matrix(self.attention_mask)) + object.__setattr__(self, "position_ids", _int_matrix(self.position_ids)) + object.__setattr__(self, "cache_metadata", _freeze_mapping(self.cache_metadata)) + object.__setattr__(self, "packing_metadata", _freeze_mapping(self.packing_metadata)) + + if not self.token_ids or not self.token_ids[0]: + raise ValueError("token_ids must contain at least one token") + if not self.selected_token_ids: + raise ValueError("selected_token_ids must not be empty") + if not self.active_mask: + raise ValueError("active_mask must not be empty") + if not self.attention_mask: + raise ValueError("attention_mask must not be empty") + for name in ( + "token_ids", + "selected_token_ids", + "active_mask", + "attention_mask", + "position_ids", + ): + _validate_rectangular(name, getattr(self, name)) + _validate_same_matrix_shape( + "token_ids", + self.token_ids, + "selected_token_ids", + self.selected_token_ids, + ) + _validate_same_matrix_shape( + "selected_token_ids", self.selected_token_ids, "active_mask", self.active_mask + ) + _validate_same_matrix_shape( + "token_ids", self.token_ids, "attention_mask", self.attention_mask + ) + _validate_same_matrix_shape("token_ids", self.token_ids, "position_ids", self.position_ids) + + +@dataclass(frozen=True) +class ScorerSpec(SerializableModel): + side: ScoreSide + backend_id: str + dtype: str + device: str = "cpu" + world_size: int = 1 + topology: Mapping[str, Any] = field(default_factory=dict) + construction_options: Mapping[str, Any] = field(default_factory=dict) + operator_overrides: Mapping[str, str] = field(default_factory=dict) + schema_version: str = "cross_config.scorer.v1" + + def __post_init__(self) -> None: + object.__setattr__(self, "side", _coerce_enum(self.side, ScoreSide)) + if not self.backend_id: + raise ValueError("backend_id must not be empty") + if not self.dtype: + raise ValueError("dtype must not be empty") + if self.world_size < 1: + raise ValueError("world_size must be >= 1") + object.__setattr__(self, "topology", _freeze_mapping(self.topology)) + object.__setattr__(self, "construction_options", _freeze_mapping(self.construction_options)) + object.__setattr__(self, "operator_overrides", _freeze_mapping(self.operator_overrides)) + + +@dataclass(frozen=True) +class KnobDescriptor(SerializableModel): + path: str + lifecycle: IsolationScope + targets: tuple[str, ...] + allowed_values: tuple[Any, ...] = () + derived: bool = False + critical: bool = True + schema_version: str = "cross_config.knob_descriptor.v1" + + def __post_init__(self) -> None: + object.__setattr__(self, "lifecycle", _coerce_enum(self.lifecycle, IsolationScope)) + object.__setattr__(self, "targets", tuple(str(target) for target in self.targets)) + object.__setattr__(self, "allowed_values", tuple(_freeze_value(self.allowed_values))) + if not self.path: + raise ValueError("path must not be empty") + if not self.targets: + raise ValueError("targets must not be empty") + + +@dataclass(frozen=True) +class InterventionSpec(SerializableModel): + path: str + values: tuple[Any, ...] + schema_version: str = "cross_config.intervention.v1" + + def __post_init__(self) -> None: + if not self.path: + raise ValueError("path must not be empty") + object.__setattr__(self, "values", tuple(_freeze_value(self.values))) + if not self.values: + raise ValueError("values must not be empty") + + +@dataclass(frozen=True) +class ExperimentDefinition(SerializableModel): + experiment_id: str + scenario_id: str + identity: SemanticIdentitySpec + baseline: Mapping[str, Any] + interventions: tuple[InterventionSpec, ...] = () + scenario: Mapping[str, Any] = field(default_factory=dict) + strategy: PlanningStrategy = PlanningStrategy.ONE_AT_A_TIME + strict_fallback: bool = True + pairwise_paths: tuple[tuple[str, str], ...] = () + contract_source: str = "ws1" + contract_version: str = "current" + schema_version: str = "cross_config.experiment_definition.v1" + + def __post_init__(self) -> None: + if not self.experiment_id: + raise ValueError("experiment_id must not be empty") + if not self.scenario_id: + raise ValueError("scenario_id must not be empty") + if self.contract_source != "ws1": + raise ValueError("Cross-configuration alignment V1 requires contract_source='ws1'") + if self.contract_version != "current": + raise ValueError("Cross-configuration alignment V1 requires contract_version='current'") + object.__setattr__(self, "baseline", _freeze_mapping(self.baseline)) + object.__setattr__(self, "scenario", _freeze_mapping(self.scenario)) + object.__setattr__(self, "interventions", tuple(self.interventions)) + object.__setattr__(self, "strategy", _coerce_enum(self.strategy, PlanningStrategy)) + normalized_pairs: list[tuple[str, str]] = [] + for pair in self.pairwise_paths: + if len(pair) != 2: + raise ValueError("each pairwise_paths entry must contain exactly two paths") + normalized_pairs.append((str(pair[0]), str(pair[1]))) + object.__setattr__(self, "pairwise_paths", tuple(normalized_pairs)) + + +@dataclass(frozen=True) +class ExperimentCase(SerializableModel): + case_id: str + experiment_id: str + scenario_id: str + identity: SemanticIdentitySpec + requested: Mapping[str, Any] + execution_binding: Mapping[str, Any] = field(default_factory=dict) + changed_paths: tuple[str, ...] = () + contract_fingerprint: str = "" + scenario_fingerprint: str = "" + schema_version: str = "cross_config.experiment_case.v1" + + def __post_init__(self) -> None: + if not self.case_id: + raise ValueError("case_id must not be empty") + if not self.experiment_id: + raise ValueError("experiment_id must not be empty") + if not self.scenario_id: + raise ValueError("scenario_id must not be empty") + object.__setattr__(self, "requested", _freeze_mapping(self.requested)) + object.__setattr__(self, "execution_binding", _freeze_mapping(self.execution_binding)) + object.__setattr__(self, "changed_paths", tuple(str(path) for path in self.changed_paths)) + + +@dataclass(frozen=True) +class MaterializedCase(SerializableModel): + case: ExperimentCase + normalized: Mapping[str, Any] + materialized: Mapping[str, Any] + isolation_scope: IsolationScope + construction_fingerprint: str = "" + distributed_context_fingerprint: str = "" + process_fingerprint: str = "" + status: MaterializationStatus = MaterializationStatus.APPLIED + evidence: Mapping[str, Any] = field(default_factory=dict) + schema_version: str = "cross_config.materialized_case.v1" + + def __post_init__(self) -> None: + object.__setattr__(self, "normalized", _freeze_mapping(self.normalized)) + object.__setattr__(self, "materialized", _freeze_mapping(self.materialized)) + object.__setattr__( + self, + "isolation_scope", + _coerce_enum(self.isolation_scope, IsolationScope), + ) + object.__setattr__(self, "status", _coerce_enum(self.status, MaterializationStatus)) + object.__setattr__(self, "evidence", _freeze_mapping(self.evidence)) + + +@dataclass(frozen=True) +class CanonicalScoringBatch(SerializableModel): + identity: SemanticIdentitySpec + input_ids: torch.Tensor + selected_token_ids: torch.Tensor + active_mask: torch.Tensor + attention_mask: torch.Tensor + position_ids: Optional[torch.Tensor] = None + metadata: Mapping[str, Any] = field(default_factory=dict) + schema_version: str = "cross_config.canonical_scoring_batch.v1" + + def __post_init__(self) -> None: + object.__setattr__(self, "input_ids", _snapshot_tensor(self.input_ids, dtype=torch.long)) + object.__setattr__( + self, "selected_token_ids", _snapshot_tensor(self.selected_token_ids, dtype=torch.long) + ) + object.__setattr__( + self, "active_mask", _snapshot_tensor(self.active_mask, dtype=torch.bool) + ) + object.__setattr__( + self, "attention_mask", _snapshot_tensor(self.attention_mask, dtype=torch.bool) + ) + if self.position_ids is not None: + object.__setattr__( + self, "position_ids", _snapshot_tensor(self.position_ids, dtype=torch.long) + ) + object.__setattr__(self, "metadata", _freeze_mapping(self.metadata)) + + shape = self.input_ids.shape + if self.input_ids.ndim != 2: + raise ValueError("input_ids must have shape [batch, sequence]") + for name in ("selected_token_ids", "active_mask", "attention_mask"): + if getattr(self, name).shape != shape: + raise ValueError(f"{name} shape must match input_ids shape") + if self.position_ids is not None and self.position_ids.shape != shape: + raise ValueError("position_ids shape must match input_ids shape") + + _require_tensor_matches_matrix("input_ids", self.input_ids, self.identity.token_ids) + _require_tensor_matches_matrix( + "selected_token_ids", self.selected_token_ids, self.identity.selected_token_ids + ) + _require_tensor_matches_matrix("active_mask", self.active_mask, self.identity.active_mask) + _require_tensor_matches_matrix( + "attention_mask", self.attention_mask, self.identity.attention_mask + ) + if self.identity.position_ids: + if self.position_ids is None: + raise ValueError("position_ids are required by the semantic identity") + _require_tensor_matches_matrix( + "position_ids", self.position_ids, self.identity.position_ids + ) + elif self.position_ids is not None: + raise ValueError("position_ids were supplied but are absent from semantic identity") + + +def _require_tensor_matches_matrix( + name: str, + tensor: torch.Tensor, + matrix: tuple[tuple[Any, ...], ...], +) -> None: + expected = torch.tensor(matrix, dtype=tensor.dtype, device=tensor.device) + if expected.shape != tensor.shape or not torch.equal(tensor, expected): + raise ValueError(f"{name} does not match the semantic identity") + + +@dataclass(frozen=True) +class RuntimeProvenance(SerializableModel): + requested: Mapping[str, Any] + normalized: Mapping[str, Any] + materialized: Mapping[str, Any] + actual: Mapping[str, Any] + status: MaterializationStatus = MaterializationStatus.APPLIED + construction_fingerprint: str = "" + distributed_context_fingerprint: str = "" + process_fingerprint: str = "" + implementation_fingerprint: str = "" + evidence: Mapping[str, Any] = field(default_factory=dict) + rank: int = 0 + world_size: int = 1 + schema_version: str = "cross_config.runtime_provenance.v1" + + def __post_init__(self) -> None: + object.__setattr__(self, "requested", _freeze_mapping(self.requested)) + object.__setattr__(self, "normalized", _freeze_mapping(self.normalized)) + object.__setattr__(self, "materialized", _freeze_mapping(self.materialized)) + object.__setattr__(self, "actual", _freeze_mapping(self.actual)) + object.__setattr__(self, "status", _coerce_enum(self.status, MaterializationStatus)) + object.__setattr__(self, "evidence", _freeze_mapping(self.evidence)) + if self.rank < 0: + raise ValueError("rank must be >= 0") + if self.world_size < 1: + raise ValueError("world_size must be >= 1") + if self.rank >= self.world_size: + raise ValueError("rank must be less than world_size") + + +@dataclass(frozen=True) +class ScoreArtifact(SerializableModel): + case_id: str + attempt_id: str + side: ScoreSide + identity: SemanticIdentitySpec + scorer: ScorerSpec + selected_logprobs: torch.Tensor + active_mask: torch.Tensor + provenance: RuntimeProvenance + schema_version: str = "cross_config.score_artifact.v1" + + def __post_init__(self) -> None: + object.__setattr__(self, "side", _coerce_enum(self.side, ScoreSide)) + if not self.case_id: + raise ValueError("case_id must not be empty") + if not self.attempt_id: + raise ValueError("attempt_id must not be empty") + if self.scorer.side is not self.side: + raise ValueError("scorer side must match score artifact side") + object.__setattr__(self, "selected_logprobs", _snapshot_tensor(self.selected_logprobs)) + object.__setattr__( + self, "active_mask", _snapshot_tensor(self.active_mask, dtype=torch.bool) + ) + if self.selected_logprobs.shape != self.active_mask.shape: + raise ValueError("selected_logprobs shape must match active_mask shape") + + +@dataclass(frozen=True) +class TokenComparisonArtifact(SerializableModel): + rollout_logprobs: torch.Tensor + training_logprobs: torch.Tensor + active_mask: torch.Tensor + absolute_diff: torch.Tensor + mismatch_mask: torch.Tensor + fixed_threshold: float + schema_version: str = "cross_config.token_comparison.v1" + + def __post_init__(self) -> None: + object.__setattr__(self, "rollout_logprobs", _snapshot_tensor(self.rollout_logprobs)) + object.__setattr__(self, "training_logprobs", _snapshot_tensor(self.training_logprobs)) + object.__setattr__( + self, "active_mask", _snapshot_tensor(self.active_mask, dtype=torch.bool) + ) + object.__setattr__(self, "absolute_diff", _snapshot_tensor(self.absolute_diff)) + object.__setattr__( + self, "mismatch_mask", _snapshot_tensor(self.mismatch_mask, dtype=torch.bool) + ) + shape = self.rollout_logprobs.shape + for name in ( + "training_logprobs", + "active_mask", + "absolute_diff", + "mismatch_mask", + ): + if getattr(self, name).shape != shape: + raise ValueError(f"{name} shape must match rollout_logprobs shape") + if not math.isfinite(self.fixed_threshold) or self.fixed_threshold < 0.0: + raise ValueError("fixed_threshold must be finite and non-negative") + + +@dataclass(frozen=True) +class AlignmentResult(SerializableModel): + case_id: str + attempt_id: str + status: AlignmentStatus + comparable: bool + passed: bool + active_token_count: int + mismatch_count: int + contract_fingerprint: str + fixed_threshold: Optional[float] = None + identity_errors: tuple[str, ...] = () + artifact_errors: tuple[str, ...] = () + diagnostics: Mapping[str, Any] = field(default_factory=dict) + token_artifact: Optional[TokenComparisonArtifact] = None + schema_version: str = "cross_config.alignment_result.v1" + + def __post_init__(self) -> None: + object.__setattr__(self, "status", _coerce_enum(self.status, AlignmentStatus)) + object.__setattr__(self, "identity_errors", tuple(self.identity_errors)) + object.__setattr__(self, "artifact_errors", tuple(self.artifact_errors)) + object.__setattr__(self, "diagnostics", _freeze_mapping(self.diagnostics)) + if self.active_token_count < 0: + raise ValueError("active_token_count must be non-negative") + if self.mismatch_count < 0: + raise ValueError("mismatch_count must be non-negative") + if self.mismatch_count > self.active_token_count: + raise ValueError("mismatch_count cannot exceed active_token_count") + if self.status is AlignmentStatus.PASS and not self.passed: + raise ValueError("PASS result must set passed=True") + if self.status is not AlignmentStatus.PASS and self.passed: + raise ValueError("only PASS results may set passed=True") + + +__all__ = [ + "AlignmentResult", + "AlignmentStatus", + "CanonicalScoringBatch", + "ExperimentCase", + "ExperimentDefinition", + "InterventionSpec", + "IsolationScope", + "KnobDescriptor", + "MaterializationStatus", + "MaterializedCase", + "PlanningStrategy", + "RuntimeProvenance", + "ScoreArtifact", + "ScoreSide", + "ScorerSpec", + "SemanticIdentitySpec", + "SerializableModel", + "TokenComparisonArtifact", +] diff --git a/rl_engine/alignment/testing/__init__.py b/rl_engine/alignment/testing/__init__.py new file mode 100644 index 00000000..2254e26a --- /dev/null +++ b/rl_engine/alignment/testing/__init__.py @@ -0,0 +1,18 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Test-only integration helpers for the alignment framework.""" + +from .smoke_ops import ( + SMOKE_ONLY_LOGP_OFFSET_BACKEND_ID, + SMOKE_ONLY_LOGP_REFERENCE_BACKEND_ID, + register_smoke_operators, + smoke_operator_descriptors, +) + +__all__ = [ + "SMOKE_ONLY_LOGP_OFFSET_BACKEND_ID", + "SMOKE_ONLY_LOGP_REFERENCE_BACKEND_ID", + "register_smoke_operators", + "smoke_operator_descriptors", +] diff --git a/rl_engine/alignment/testing/cpu_cross_config.py b/rl_engine/alignment/testing/cpu_cross_config.py new file mode 100644 index 00000000..9c9e2ce2 --- /dev/null +++ b/rl_engine/alignment/testing/cpu_cross_config.py @@ -0,0 +1,696 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""CPU-only adapters for cross-configuration smoke execution. + +This module is deliberately outside the production framework package. It gives +the CLI and tests a deterministic execution target without implying CUDA, +distributed, vLLM, or training-runtime support. +""" + +from __future__ import annotations + +from pathlib import Path +from types import SimpleNamespace +from typing import Any, Mapping, Optional + +import torch + +from rl_engine.alignment.cross_config.artifacts import ArtifactStore +from rl_engine.alignment.cross_config.config import ExperimentConfig, OperatorSelection +from rl_engine.alignment.cross_config.execution_plan import build_execution_plan +from rl_engine.alignment.cross_config.operators import ( + OperatorBridge, + OperatorOverride, + selected_logprobs_with_operator, +) +from rl_engine.alignment.cross_config.runner import PairedRunner, PairedRunResult, RankScore +from rl_engine.alignment.cross_config.runtime import ( + AdapterMaterialization, + KnobApplication, + RuntimeBinding, + RuntimeTools, +) +from rl_engine.alignment.cross_config.schema import ( + CanonicalScoringBatch, + ExperimentCase, + KnobDescriptor, + MaterializationStatus, + ScorerSpec, + ScoreSide, +) +from rl_engine.executors.stateless_executor import ( + StatelessForwardConfig, + StatelessForwardExecutor, + StatelessForwardInputs, +) +from rl_engine.kernels.registry import kernel_registry +from rl_engine.kernels.semantic_registry import ( + OperatorRequirements, + OperatorResolutionPolicy, + SemanticOperatorCatalog, +) +from rl_engine.kernels.semantic_registry import ( + implementation_fingerprint as fingerprint_implementation, +) + +CPU_SCORER_IMPLEMENTATION_FINGERPRINT = "cross_config.cpu_stateless_scorer.v1" + + +class SyntheticCpuCausalLM(torch.nn.Module): + """Deterministic parameter-free model for the named CPU smoke scenario.""" + + def __init__(self, vocab_size: int): + super().__init__() + self.vocab_axis: torch.Tensor + self.register_buffer( + "vocab_axis", + torch.arange(vocab_size, dtype=torch.float32), + persistent=False, + ) + self.config = SimpleNamespace(use_cache=False, _attn_implementation="eager") + self.generation_config = SimpleNamespace(use_cache=False) + + def forward( + self, + input_ids: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.Tensor] = None, + use_cache: Optional[bool] = None, + ) -> Any: + del attention_mask + if use_cache not in {None, False}: + raise ValueError("CPU smoke scoring forbids KV-cache generation") + if input_ids.device.type != "cpu": + raise ValueError("the synthetic smoke model accepts CPU tensors only") + if position_ids is None: + position_ids = torch.arange(input_ids.shape[1], device="cpu").expand_as(input_ids) + centers = torch.remainder(input_ids + position_ids + 1, self.vocab_axis.numel()).float() + logits = -torch.abs(self.vocab_axis.view(1, 1, -1) - centers.unsqueeze(-1)) * 0.125 + return SimpleNamespace(logits=logits, past_key_values=None) + + +class CpuStatelessScorer: + """Read-only teacher-forcing adapter over ``StatelessForwardExecutor``.""" + + optimizer = None + implementation_fingerprint = CPU_SCORER_IMPLEMENTATION_FINGERPRINT + + def __init__( + self, + model: torch.nn.Module, + spec: ScorerSpec, + config: Optional[StatelessForwardConfig] = None, + ): + if spec.world_size != 1: + raise ValueError("CpuStatelessScorer supports only world_size=1") + if _device_type(spec.device) != "cpu": + raise ValueError("CpuStatelessScorer is explicitly CPU-only") + resolved_config = config or StatelessForwardConfig( + mode="reference", + attention_backend="eager", + output_dtype=_torch_dtype(spec.dtype), + ) + if resolved_config.mode not in {"reference", "both"}: + raise ValueError("CpuStatelessScorer requires reference scoring mode") + expected_dtype = _torch_dtype(spec.dtype) + if resolved_config.output_dtype is not expected_dtype: + raise ValueError("stateless output_dtype must match the scorer dtype") + _require_module_on_cpu(model) + _require_module_float_dtype(model, expected_dtype) + self.model = model + self.spec = spec + self.config = resolved_config + + def score( + self, + batch: CanonicalScoringBatch, + *, + batch_size: int, + operator: Any, + ) -> tuple[RankScore, ...]: + if batch_size < 1: + raise ValueError("batch_size must be greater than zero") + + def selected_logprob_fn( + logits: torch.Tensor, + token_ids: torch.Tensor, + *, + mask: Optional[torch.Tensor] = None, + temperature: float = 1.0, + output_dtype: torch.dtype = torch.float32, + ) -> torch.Tensor: + return selected_logprobs_with_operator( + operator, + logits, + token_ids, + active_mask=mask, + temperature=temperature, + output_dtype=output_dtype, + ) + + executor = StatelessForwardExecutor( + self.model, + self.config, + selected_logprob_fn=selected_logprob_fn, + ) + chunks: list[torch.Tensor] = [] + observed_ranges: list[tuple[int, int]] = [] + for start in range(0, batch.input_ids.shape[0], batch_size): + stop = min(start + batch_size, batch.input_ids.shape[0]) + inputs = StatelessForwardInputs( + input_ids=batch.input_ids[start:stop], + attention_mask=batch.attention_mask[start:stop], + completion_mask=batch.active_mask[start:stop], + labels=batch.selected_token_ids[start:stop], + position_ids=( + None if batch.position_ids is None else batch.position_ids[start:stop] + ), + ) + result = executor.score(inputs) + if result.reference_logps is None: # pragma: no cover - guarded by config mode + raise RuntimeError("stateless scorer returned no selected logprobs") + chunks.append(result.reference_logps.detach().to(device="cpu")) + observed_ranges.append((start, stop)) + selected = torch.cat(chunks, dim=0) + return ( + RankScore( + rank=0, + world_size=1, + selected_logprobs=selected, + metadata={ + "device": "cpu", + "teacher_forcing": True, + "use_cache": False, + "optimizer_step": False, + "batch_ranges": observed_ranges, + }, + ), + ) + + +class CpuSmokeMaterializer: + """Materialize the exact single-process CPU surface used by smoke tests.""" + + runtime_kind = "cpu_smoke" + + @property + def implementation_fingerprint(self) -> str: + """Seal the adapter's concrete class and materialization entry point.""" + + return fingerprint_implementation( + type(self), + instance=self, + entrypoints=("materialize",), + ) + + def __init__( + self, + *, + requested_operator_backends: Optional[Mapping[str, str]] = None, + actual_operator_backends: Optional[Mapping[str, str]] = None, + ): + self.requested_operator_backends = dict(requested_operator_backends or {}) + self.actual_operator_backends = dict(actual_operator_backends or {}) + + def materialize( + self, + normalized: Mapping[str, Any], + descriptors: Mapping[str, KnobDescriptor], + ) -> AdapterMaterialization: + flat = _flatten(normalized) + applications = tuple( + self._application(path, value, descriptors[path]) for path, value in flat.items() + ) + batch_size = int(flat["batch.size"]) + requested_logp = str(flat["logp.backend"]) + operator_backends = self.requested_operator_backends or { + "rollout": requested_logp, + "training": requested_logp, + } + return AdapterMaterialization( + applications=applications, + binding=RuntimeBinding( + batch_size=batch_size, + side_configs={ + "rollout": { + "device": "cpu", + "dtype": "float32", + "enable_prefix_caching": False, + "enforce_eager": True, + }, + "training": { + "device": "cpu", + "dtype": "float32", + "attention_backend": "eager", + "sharding": "unsharded", + }, + }, + topology={ + "rollout": { + "world_size": 1, + "tensor_parallel_size": 1, + "context_parallel_size": 1, + }, + "training": {"world_size": 1, "sharding": "unsharded"}, + }, + scorer={ + "mode": "reference", + "use_cache": False, + "attention_backend": "eager", + "output_dtype": "float32", + }, + operator_backends=operator_backends, + runtime_kind=self.runtime_kind, + ), + ) + + def _application( + self, + path: str, + requested: Any, + descriptor: KnobDescriptor, + ) -> KnobApplication: + fixed_values = { + "rollout.tensor_parallel_size": 1, + "rollout.context_parallel_size": 1, + "rollout.dtype": "float32", + "rollout.enable_prefix_caching": False, + "rollout.enforce_eager": True, + "training.attention_backend": "eager", + "training.compute_dtype": "float32", + "training.sharding": "unsharded", + } + if path == "batch.size": + return _application( + descriptor, + requested, + requested, + requested, + MaterializationStatus.APPLIED, + "canonical batch is partitioned at scorer invocation", + ) + if path == "logp.backend": + requested_backends = self.requested_operator_backends or { + "rollout": requested, + "training": requested, + } + if requested_backends.get("rollout") != requested: + return _application( + descriptor, + requested, + requested_backends, + None, + MaterializationStatus.ERROR, + "rollout operator conflicts with public logp.backend", + ) + actual_backends = { + "rollout": self.actual_operator_backends.get("rollout"), + "training": self.actual_operator_backends.get("training"), + } + if None in actual_backends.values(): + return _application( + descriptor, + requested, + requested_backends, + None, + MaterializationStatus.UNOBSERVABLE, + "operator resolution trace has not been supplied", + ) + status = ( + MaterializationStatus.APPLIED + if actual_backends == requested_backends + else MaterializationStatus.FALLBACK + ) + return _application( + descriptor, + requested, + requested_backends, + actual_backends, + status, + "concrete CPU backends were read from exact resolution traces", + ) + + actual = fixed_values[path] + status = ( + MaterializationStatus.APPLIED + if requested == actual + else MaterializationStatus.UNSUPPORTED + ) + reason = ( + "read back from the single-process CPU scorer" + if status is MaterializationStatus.APPLIED + else f"CPU smoke supports only {path}={actual!r}" + ) + return _application(descriptor, requested, requested, actual, status, reason) + + +def run_cpu_experiment( + config: ExperimentConfig, + *, + output_root: str | Path, + allow_smoke_operators: bool = False, + timeout_seconds: float = 30.0, + resume: bool = True, +) -> dict[str, Any]: + """Run every planned case through the explicit CPU smoke adapter.""" + + scenario_device = str(config.definition.scenario.get("device", "")).strip().lower() + if scenario_device != "cpu": + raise ValueError("the CPU runtime requires scenario.device='cpu'") + plan = build_execution_plan(config) + + store = ArtifactStore(output_root) + experiment_dir = store.initialize_experiment( + config.definition.experiment_id, + experiment=plan.experiment, + plan=plan.rows(), + ) + batch = canonical_cpu_batch(config) + runs = [ + run_cpu_case( + store, + entry.case, + batch, + entry.operators, + allow_smoke_operators=allow_smoke_operators, + strict=config.definition.strict_fallback, + timeout_seconds=timeout_seconds, + resume=resume, + ) + for entry in plan.entries + ] + cases = [ + { + "case_id": run.case_id, + "attempt_id": run.attempt_id, + "status": str(run.summary["status"]), + "rollout_backend": run.summary["rollout_backend"], + "training_backend": run.summary["training_backend"], + "mismatch_count": run.summary.get("mismatch_count"), + "worst_token_index": run.summary.get("worst_token_index"), + "resumed": run.resumed, + "attempt_dir": str(run.attempt_dir), + } + for run in runs + ] + return { + "schema_version": "cross_config.cli_summary.v1", + "status": "pass" if all(item["status"] == "pass" for item in cases) else "fail", + "experiment_id": config.definition.experiment_id, + "scenario_id": config.definition.scenario_id, + "runtime": "cpu-smoke", + "artifact_dir": str(experiment_dir), + "cases": cases, + } + + +def run_cpu_case( + store: ArtifactStore, + case: ExperimentCase, + batch: CanonicalScoringBatch, + selection: OperatorSelection, + *, + allow_smoke_operators: bool, + strict: bool, + timeout_seconds: float, + resume: bool, +) -> PairedRunResult: + """Execute one already-bound CPU case with case-local operator state.""" + + catalog = SemanticOperatorCatalog(kernel_registry.semantic.backend_descriptors()) + if allow_smoke_operators: + from rl_engine.alignment.testing.smoke_ops import register_smoke_operators + + register_smoke_operators(catalog, allow_smoke_operators=True) + bridge = OperatorBridge( + catalog, + policy=OperatorResolutionPolicy( + strict=strict, + allow_test_backends=allow_smoke_operators, + ), + ) + override = OperatorOverride( + semantic_op="selected_logprob", + rollout_backend=selection.rollout_backend, + training_backend=selection.training_backend, + ) + topologies: dict[str, Mapping[str, Any]] = { + "rollout": { + "world_size": 1, + "tensor_parallel_size": 1, + "context_parallel_size": 1, + }, + "training": {"world_size": 1, "sharding": "unsharded"}, + } + rollout_dtype = str(case.requested["rollout"]["dtype"]) + training_dtype = str(case.requested["training"]["compute_dtype"]) + requirements = { + "rollout": OperatorRequirements( + device="cpu", + dtype=rollout_dtype, + topology=topologies["rollout"], + alignment_properties={"deterministic": True}, + ), + "training": OperatorRequirements( + device="cpu", + dtype=training_dtype, + topology=topologies["training"], + alignment_properties={"deterministic": True}, + ), + } + resolved = bridge.resolve_override(override, requirements=requirements, strict=strict) + options = { + target: _factory_options( + selection.backend_for(target), + selection.options_for(target), + allow_smoke_operators=allow_smoke_operators, + ) + for target in ("rollout", "training") + } + instances = { + target: bridge.instantiate( + resolved, + target=target, # type: ignore[arg-type] + factory_kwargs=options[target], + ) + for target in ("rollout", "training") + } + provenance = { + target: bridge.instance_provenance( + resolved, + target=target, # type: ignore[arg-type] + instance=instances[target], + ) + for target in ("rollout", "training") + } + actual_backends = {target: provenance[target].backend_id for target in provenance} + materialization = RuntimeTools().materialize( + case, + CpuSmokeMaterializer( + requested_operator_backends={ + "rollout": selection.rollout_backend, + "training": selection.training_backend, + }, + actual_operator_backends=actual_backends, + ), + ) + RuntimeTools.require_executable(materialization, strict=strict) + + minimum_token_id = min( + int(batch.input_ids.min().item()), + int(batch.selected_token_ids.min().item()), + ) + if minimum_token_id < 0: + raise ValueError("CPU smoke token IDs must be non-negative") + vocab_size = ( + max( + int(batch.input_ids.max().item()), + int(batch.selected_token_ids.max().item()), + ) + + 17 + ) + scorers = { + "rollout": CpuStatelessScorer( + SyntheticCpuCausalLM(vocab_size), + _scorer_spec( + ScoreSide.ROLLOUT, + rollout_dtype, + selection.rollout_backend, + topologies["rollout"], + case, + ), + ), + "training": CpuStatelessScorer( + SyntheticCpuCausalLM(vocab_size), + _scorer_spec( + ScoreSide.TRAINING, + training_dtype, + selection.training_backend, + topologies["training"], + case, + ), + ), + } + return PairedRunner(store, timeout_seconds=timeout_seconds).run( + case, + materialization, + batch, + scorers["rollout"], + scorers["training"], + resolved, + instances, + provenance, + operator_factory_options=options, + strict=strict, + timeout_seconds=timeout_seconds, + resume=resume, + ) + + +def canonical_cpu_batch(config: ExperimentConfig) -> CanonicalScoringBatch: + """Build the immutable CPU tensors frozen by an experiment identity.""" + + identity = config.definition.identity + position_ids = ( + torch.tensor(identity.position_ids, dtype=torch.long, device="cpu") + if identity.position_ids + else None + ) + return CanonicalScoringBatch( + identity=identity, + input_ids=torch.tensor(identity.token_ids, dtype=torch.long, device="cpu"), + selected_token_ids=torch.tensor( + identity.selected_token_ids, + dtype=torch.long, + device="cpu", + ), + active_mask=torch.tensor(identity.active_mask, dtype=torch.bool, device="cpu"), + attention_mask=torch.tensor( + identity.attention_mask, + dtype=torch.bool, + device="cpu", + ), + position_ids=position_ids, + metadata={"source": "named_json", "device": "cpu"}, + ) + + +def _factory_options( + backend_id: str, + configured: Mapping[str, Any], + *, + allow_smoke_operators: bool, +) -> dict[str, Any]: + options = dict(configured) + if backend_id == "smoke_only.logp_offset": + if not allow_smoke_operators: + raise PermissionError("smoke offset requires explicit test authorization") + options["allow_smoke_operators"] = True + return options + + +def _scorer_spec( + side: ScoreSide, + dtype: str, + backend_id: str, + topology: Mapping[str, Any], + case: ExperimentCase, +) -> ScorerSpec: + identity = case.identity + return ScorerSpec( + side=side, + backend_id="cpu_stateless_teacher_forcing", + dtype=dtype, + device="cpu", + world_size=1, + topology=topology, + construction_options={ + "checkpoint_id": identity.checkpoint_id, + "model_version": identity.model_version, + "pre_update_state": identity.pre_update_state, + "teacher_forcing": True, + "use_cache": False, + }, + operator_overrides={"selected_logprob": backend_id}, + ) + + +def _application( + descriptor: KnobDescriptor, + requested: Any, + materialized: Any, + actual: Any, + status: MaterializationStatus, + reason: str, +) -> KnobApplication: + return KnobApplication( + path=descriptor.path, + requested=requested, + materialized=materialized, + actual=actual, + lifecycle=descriptor.lifecycle, + status=status, + critical=descriptor.critical, + evidence={"reason": reason}, + ) + + +def _flatten(value: Mapping[str, Any], prefix: str = "") -> dict[str, Any]: + result: dict[str, Any] = {} + for key, child in value.items(): + path = f"{prefix}.{key}" if prefix else key + if isinstance(child, Mapping): + result.update(_flatten(child, path)) + else: + result[path] = child + return result + + +def _require_module_on_cpu(model: torch.nn.Module) -> None: + tensors = tuple(model.parameters()) + tuple(model.buffers()) + if any(tensor.device.type != "cpu" for tensor in tensors): + raise ValueError("CPU smoke models must remain on CPU") + + +def _require_module_float_dtype(model: torch.nn.Module, expected: torch.dtype) -> None: + tensors = tuple(model.parameters()) + tuple(model.buffers()) + mismatched = sorted( + { + str(tensor.dtype).replace("torch.", "") + for tensor in tensors + if tensor.is_floating_point() and tensor.dtype is not expected + } + ) + if mismatched: + raise ValueError( + f"CPU smoke model floating dtype must be {expected}; observed {mismatched}" + ) + + +def _device_type(value: str) -> str: + return value.split(":", 1)[0].strip().lower() + + +def _torch_dtype(value: str) -> torch.dtype: + normalized = value.strip().lower().replace("torch.", "") + aliases = {"fp32": "float32", "bf16": "bfloat16", "fp16": "float16"} + normalized = aliases.get(normalized, normalized) + try: + return { + "float32": torch.float32, + "bfloat16": torch.bfloat16, + "float16": torch.float16, + }[normalized] + except KeyError as exc: + raise ValueError(f"unsupported scorer dtype {value!r}") from exc + + +__all__ = [ + "CpuSmokeMaterializer", + "CpuStatelessScorer", + "SyntheticCpuCausalLM", + "canonical_cpu_batch", + "run_cpu_case", + "run_cpu_experiment", +] diff --git a/rl_engine/alignment/testing/smoke_ops/SMOKE_OPERATORS.md b/rl_engine/alignment/testing/smoke_ops/SMOKE_OPERATORS.md new file mode 100644 index 00000000..0ff9e6ea --- /dev/null +++ b/rl_engine/alignment/testing/smoke_ops/SMOKE_OPERATORS.md @@ -0,0 +1,30 @@ +# Cross-configuration smoke-only operators + +These files are temporary test scaffolding. They validate operator selection, +strict resolution, active-token scoring, and provenance; they do not establish +production numerical alignment. + +| File | Backend | Purpose | Replacement owner / issue | +| --- | --- | --- | --- | +| `smoke_only_logp_reference.py` | `smoke_only.logp_reference` | CPU PyTorch `log_softmax` plus gather reference for rollout/training injection tests. | Production selected-logprob operator workstream; roadmap issue #83 / WS1 contract issue #108. | +| `smoke_only_logp_offset.py` | `smoke_only.logp_offset` | Adds an explicit deterministic active-token offset so comparator mismatch detection can be tested. | Test-only fault injection; no production replacement should preserve the offset. | +| `__init__.py` | registration boundary | Keeps registration disabled by default and requires `allow_smoke_operators=True`. | Remove with both smoke implementations. | + +Removal trigger: delete this package once equivalent production RL-Kernel +selected-logprob operators are integrated and the same framework tests pass using +those production backends on both rollout and training sides. + +Exact deletion steps: + +1. Change `tests/test_cross_config_runtime.py` to exercise the production backend + IDs while preserving disabled/unavailable, capability, paired-output, and + provenance coverage. +2. Remove the `smoke_operator` test marker if no other temporary smoke operator + tests use it. +3. Delete `rl_engine/alignment/testing/smoke_ops/` and remove its exports from + `rl_engine/alignment/testing/__init__.py`. +4. Search the repository for `smoke_only.`, `allow_smoke_operators`, and + `RL_KERNEL_ALLOW_SMOKE_OPS`; remove configuration and documentation references + that no longer describe an active test boundary. +5. Run the cross-configuration contract, runtime, runner, and production-backend + tests before merging the deletion. diff --git a/rl_engine/alignment/testing/smoke_ops/__init__.py b/rl_engine/alignment/testing/smoke_ops/__init__.py new file mode 100644 index 00000000..5b6296ae --- /dev/null +++ b/rl_engine/alignment/testing/smoke_ops/__init__.py @@ -0,0 +1,75 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Opt-in registration for temporary Cross-configuration alignment smoke-only operators.""" + +from __future__ import annotations + +from typing import Any + +from rl_engine.kernels.semantic_registry import OperatorBackendDescriptor, SemanticOperatorCatalog + +SMOKE_ONLY_LOGP_REFERENCE_BACKEND_ID = "smoke_only.logp_reference" +SMOKE_ONLY_LOGP_OFFSET_BACKEND_ID = "smoke_only.logp_offset" + + +def smoke_operator_descriptors() -> tuple[OperatorBackendDescriptor, ...]: + """Build smoke descriptors lazily without registering them globally.""" + + from .smoke_only_logp_offset import SMOKE_ONLY_LOGP_OFFSET_DESCRIPTOR + from .smoke_only_logp_reference import SMOKE_ONLY_LOGP_REFERENCE_DESCRIPTOR + + return ( + SMOKE_ONLY_LOGP_REFERENCE_DESCRIPTOR, + SMOKE_ONLY_LOGP_OFFSET_DESCRIPTOR, + ) + + +def register_smoke_operators( + catalog: SemanticOperatorCatalog, + *, + allow_smoke_operators: bool = False, + replace: bool = False, +) -> tuple[OperatorBackendDescriptor, ...]: + """Register every smoke backend after an explicit per-call opt-in. + + Importing this package never mutates a catalog. Resolution independently + requires an ``OperatorResolutionPolicy`` that allows test backends; this + registration guard is the first fail-closed boundary. + """ + + if allow_smoke_operators is not True: + raise PermissionError( + "smoke operator registration requires explicit " "allow_smoke_operators=True" + ) + if not isinstance(catalog, SemanticOperatorCatalog): + raise TypeError("catalog must be a SemanticOperatorCatalog") + + descriptors = smoke_operator_descriptors() + for descriptor in descriptors: + catalog.register_backend(descriptor, replace=replace) + return descriptors + + +def __getattr__(name: str) -> Any: + """Lazily expose implementation classes without default torch imports.""" + + if name == "SmokeOnlyLogpReference": + from .smoke_only_logp_reference import SmokeOnlyLogpReference + + return SmokeOnlyLogpReference + if name == "SmokeOnlyLogpOffset": + from .smoke_only_logp_offset import SmokeOnlyLogpOffset + + return SmokeOnlyLogpOffset + raise AttributeError(name) + + +__all__ = [ + "SMOKE_ONLY_LOGP_OFFSET_BACKEND_ID", + "SMOKE_ONLY_LOGP_REFERENCE_BACKEND_ID", + "SmokeOnlyLogpOffset", + "SmokeOnlyLogpReference", + "register_smoke_operators", + "smoke_operator_descriptors", +] diff --git a/rl_engine/alignment/testing/smoke_ops/smoke_only_logp_offset.py b/rl_engine/alignment/testing/smoke_ops/smoke_only_logp_offset.py new file mode 100644 index 00000000..9bbe5e0e --- /dev/null +++ b/rl_engine/alignment/testing/smoke_ops/smoke_only_logp_offset.py @@ -0,0 +1,98 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""TEMPORARY TEST SCAFFOLD - NOT A PRODUCTION RL-KERNEL OPERATOR""" + +from __future__ import annotations + +import math +from typing import Optional + +import torch + +from rl_engine.kernels.semantic_registry import ( + OperatorBackendDescriptor, + OperatorFallbackPolicy, + OperatorLifecycle, +) + +from . import SMOKE_ONLY_LOGP_OFFSET_BACKEND_ID +from .smoke_only_logp_reference import SmokeOnlyLogpReference + + +class SmokeOnlyLogpOffset(SmokeOnlyLogpReference): + """CPU reference plus a deterministic test-only active-token offset. + + Inputs and output follow :class:`SmokeOnlyLogpReference`. ``offset`` is zero + by default. A non-zero value requires the constructor's explicit + ``allow_smoke_operators=True`` guard. The cross-configuration bridge masks + inactive output positions after invocation, so drift applies only to active + selected tokens. + """ + + def __init__( + self, + offset: float = 0.0, + *, + allow_smoke_operators: bool = False, + ) -> None: + normalized_offset = float(offset) + if not math.isfinite(normalized_offset): + raise ValueError("offset must be finite") + if normalized_offset != 0.0 and allow_smoke_operators is not True: + raise PermissionError( + "a non-zero smoke offset requires explicit " "allow_smoke_operators=True" + ) + self.offset = normalized_offset + + def apply_fp32( + self, + logits: torch.Tensor, + token_ids: torch.Tensor, + active_mask: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + selected = super().apply_fp32(logits, token_ids, active_mask=active_mask) + if self.offset == 0.0: + return selected + if active_mask is None: + return selected + self.offset + mask = active_mask.to(device=selected.device, dtype=torch.bool) + return selected + mask.to(dtype=selected.dtype) * self.offset + + +SMOKE_ONLY_LOGP_OFFSET_DESCRIPTOR = OperatorBackendDescriptor( + semantic_op="selected_logprob", + backend_id=SMOKE_ONLY_LOGP_OFFSET_BACKEND_ID, + supported_targets=frozenset({"rollout", "training"}), + supported_devices=frozenset({"cpu"}), + supported_dtypes=frozenset({"bfloat16", "float16", "float32"}), + supported_topologies={ + "rollout": { + "world_size": (1,), + "tensor_parallel_size": (1,), + "context_parallel_size": (1,), + }, + "training": { + "world_size": (1,), + "sharding": ("unsharded",), + }, + }, + determinism_or_alignment_properties={ + "algorithm": "pytorch.log_softmax_gather_plus_test_offset", + "batch_invariant": True, + "deterministic": True, + "strict_observable": True, + "test_offset_configurable": True, + }, + lifecycle=OperatorLifecycle.ENGINE_CONSTRUCTION, + implementation_class_or_factory=SmokeOnlyLogpOffset, + fallback_policy=OperatorFallbackPolicy.ERROR, + version_or_build_fingerprint="cross-config-smoke-only-logp-offset-v1", + is_smoke_only=True, +) + + +__all__ = [ + "SMOKE_ONLY_LOGP_OFFSET_DESCRIPTOR", + "SmokeOnlyLogpOffset", +] diff --git a/rl_engine/alignment/testing/smoke_ops/smoke_only_logp_reference.py b/rl_engine/alignment/testing/smoke_ops/smoke_only_logp_reference.py new file mode 100644 index 00000000..e64d37fa --- /dev/null +++ b/rl_engine/alignment/testing/smoke_ops/smoke_only_logp_reference.py @@ -0,0 +1,113 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""TEMPORARY TEST SCAFFOLD - NOT A PRODUCTION RL-KERNEL OPERATOR""" + +from __future__ import annotations + +from typing import Optional + +import torch + +from rl_engine.kernels.semantic_registry import ( + OperatorBackendDescriptor, + OperatorFallbackPolicy, + OperatorLifecycle, +) + +from . import SMOKE_ONLY_LOGP_REFERENCE_BACKEND_ID + + +class SmokeOnlyLogpReference: + """CPU selected-logprob reference used only to test operator plumbing. + + The semantic inputs are logits shaped ``[..., vocabulary]`` and selected + token IDs shaped ``[...]``. The result is one float32 log probability per + selected token. When supplied, ``active_mask`` has the token-ID shape and + inactive output positions are exactly zero. The cross-configuration bridge + applies the same masking rule when invoking the two-argument interface. + """ + + op_class = "logprob" + + def __call__( + self, + logits: torch.Tensor, + token_ids: torch.Tensor, + active_mask: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + return self.apply_fp32(logits, token_ids, active_mask=active_mask) + + def apply_fp32( + self, + logits: torch.Tensor, + token_ids: torch.Tensor, + active_mask: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """Compute CPU log-softmax/gather output with optional active masking.""" + + _validate_inputs(logits, token_ids, active_mask) + selected_ids = token_ids.to(device=logits.device, dtype=torch.long) + mask = None + if active_mask is not None: + mask = active_mask.to(device=logits.device, dtype=torch.bool) + selected_ids = selected_ids.masked_fill(~mask, 0) + + log_probs = torch.log_softmax(logits.float(), dim=-1) + selected = torch.gather(log_probs, dim=-1, index=selected_ids.unsqueeze(-1)).squeeze(-1) + if mask is not None: + selected = selected.masked_fill(~mask, 0.0) + return selected + + +def _validate_inputs( + logits: torch.Tensor, + token_ids: torch.Tensor, + active_mask: Optional[torch.Tensor], +) -> None: + if logits.device.type != "cpu": + raise ValueError("smoke-only logprob operators support CPU tensors only") + if logits.shape[:-1] != token_ids.shape: + raise ValueError( + f"logits leading shape {tuple(logits.shape[:-1])} must match " + f"token_ids shape {tuple(token_ids.shape)}" + ) + if active_mask is not None and active_mask.shape != token_ids.shape: + raise ValueError("active_mask shape must match token_ids shape") + + +SMOKE_ONLY_LOGP_REFERENCE_DESCRIPTOR = OperatorBackendDescriptor( + semantic_op="selected_logprob", + backend_id=SMOKE_ONLY_LOGP_REFERENCE_BACKEND_ID, + supported_targets=frozenset({"rollout", "training"}), + supported_devices=frozenset({"cpu"}), + supported_dtypes=frozenset({"bfloat16", "float16", "float32"}), + supported_topologies={ + "rollout": { + "world_size": (1,), + "tensor_parallel_size": (1,), + "context_parallel_size": (1,), + }, + "training": { + "world_size": (1,), + "sharding": ("unsharded",), + }, + }, + determinism_or_alignment_properties={ + "algorithm": "pytorch.log_softmax_gather", + "batch_invariant": True, + "deterministic": True, + "strict_observable": True, + }, + lifecycle=OperatorLifecycle.ENGINE_CONSTRUCTION, + implementation_class_or_factory=SmokeOnlyLogpReference, + fallback_policy=OperatorFallbackPolicy.ERROR, + version_or_build_fingerprint="cross-config-smoke-only-logp-reference-v1", + is_smoke_only=True, +) + + +__all__ = [ + "SMOKE_ONLY_LOGP_REFERENCE_DESCRIPTOR", + "SmokeOnlyLogpReference", +] diff --git a/rl_engine/executors/stateless_executor.py b/rl_engine/executors/stateless_executor.py index 2047218f..70a25f99 100644 --- a/rl_engine/executors/stateless_executor.py +++ b/rl_engine/executors/stateless_executor.py @@ -16,6 +16,7 @@ StatelessForwardMode = Literal["reference", "reward", "both"] StatelessAttentionBackend = Literal["flash_attention_2", "sdpa", "eager", "model_default"] RewardAdapter = Callable[["StatelessForwardOutputs", "StatelessForwardInputs"], torch.Tensor] +SelectedLogprobCallable = Callable[..., torch.Tensor] _MISSING = object() @@ -57,6 +58,7 @@ class StatelessForwardInputs: attention_mask: torch.Tensor completion_mask: torch.Tensor labels: Optional[torch.Tensor] = None + position_ids: Optional[torch.Tensor] = None @dataclass(frozen=True) @@ -105,10 +107,12 @@ def __init__( config: Optional[StatelessForwardConfig] = None, *, reward_adapter: Optional[RewardAdapter] = None, + selected_logprob_fn: Optional[SelectedLogprobCallable] = None, ): self.model = model self.config = config or StatelessForwardConfig() self.reward_adapter = reward_adapter or default_reward_adapter + self.selected_logprob_fn = selected_logprob_fn def score(self, inputs: StatelessForwardInputs) -> StatelessForwardResult: _validate_inputs(inputs, self.config) @@ -170,6 +174,7 @@ def score(self, inputs: StatelessForwardInputs) -> StatelessForwardResult: inputs, temperature=self.config.temperature, output_dtype=self.config.output_dtype, + selected_logprob_fn=self.selected_logprob_fn, ) if self.config.return_token_scores: token_scores = reference_logps @@ -215,8 +220,14 @@ def score_reference_logprobs( *, temperature: float = 1.0, output_dtype: torch.dtype = torch.float32, + selected_logprob_fn: Optional[SelectedLogprobCallable] = None, ) -> torch.Tensor: - """Compute causal next-token selected logprobs aligned to ``[B, S]`` masks.""" + """Compute causal next-token selected logprobs aligned to ``[B, S]`` masks. + + ``selected_logprob_fn`` is an exact injection seam with the same callable + contract as :func:`selected_logprobs_reference`. Leaving it unset preserves + the historical PyTorch-reference behavior. + """ if logits.ndim != 3: raise ValueError(f"reference logits must have shape [B, S, V], got {tuple(logits.shape)}") @@ -237,13 +248,21 @@ def score_reference_logprobs( shifted_logits = logits[:, :-1, :] shifted_labels = labels[:, 1:] shifted_mask = _bool_mask(inputs.completion_mask[:, 1:], device=logits.device) - shifted_logps = selected_logprobs_reference( + scorer = selected_logprob_fn or selected_logprobs_reference + shifted_logps = scorer( shifted_logits, shifted_labels.to(device=logits.device), mask=shifted_mask, temperature=temperature, output_dtype=output_dtype, ) + if not isinstance(shifted_logps, torch.Tensor): + raise TypeError("selected_logprob_fn must return a torch.Tensor") + if shifted_logps.shape != shifted_labels.shape: + raise ValueError( + "selected_logprob_fn output shape must match selected token IDs, got " + f"{tuple(shifted_logps.shape)} and {tuple(shifted_labels.shape)}" + ) result = torch.zeros( inputs.input_ids.shape, device=logits.device, @@ -372,11 +391,20 @@ def _temporarily_configure_stateless_model( config: StatelessForwardConfig, ) -> Iterator[dict[str, float | int | str | bool]]: saved = _model_config_snapshot(model, config) - policy = configure_stateless_model(model, config) + saved_training_modes = tuple((module, module.training) for module in model.modules()) + model.eval() try: + policy = configure_stateless_model(model, config) + policy["model_eval_during_forward"] = True yield policy finally: - _restore_model_config_snapshot(saved) + try: + _restore_model_config_snapshot(saved) + finally: + # Restore each module directly. Calling ``model.train(...)`` would + # flatten intentionally mixed child-module modes. + for module, was_training in saved_training_modes: + module.training = was_training def extract_kv_cache_outputs(raw_outputs: Any) -> Optional[Any]: @@ -450,12 +478,16 @@ def _validate_inputs(inputs: StatelessForwardInputs, config: StatelessForwardCon raise ValueError("completion_mask shape must match input_ids shape") if inputs.labels is not None and inputs.labels.shape != input_ids.shape: raise ValueError("labels shape must match input_ids shape") + if inputs.position_ids is not None and inputs.position_ids.shape != input_ids.shape: + raise ValueError("position_ids shape must match input_ids shape") if attention_mask.device != input_ids.device: raise ValueError("attention_mask device must match input_ids device") if completion_mask.device != input_ids.device: raise ValueError("completion_mask device must match input_ids device") if inputs.labels is not None and inputs.labels.device != input_ids.device: raise ValueError("labels device must match input_ids device") + if inputs.position_ids is not None and inputs.position_ids.device != input_ids.device: + raise ValueError("position_ids device must match input_ids device") if config.max_batch_size is not None and input_ids.shape[0] > config.max_batch_size: raise ValueError( f"batch size {input_ids.shape[0]} exceeds max_batch_size {config.max_batch_size}" @@ -479,6 +511,10 @@ def _run_no_cache_forward( "input_ids": inputs.input_ids, "attention_mask": inputs.attention_mask, } + if inputs.position_ids is not None: + if not _call_accepts_keyword(model, "position_ids"): + raise ValueError("model does not accept the canonical batch position_ids") + kwargs["position_ids"] = inputs.position_ids if _call_accepts_keyword(model, "use_cache"): kwargs["use_cache"] = False return model(**kwargs), True diff --git a/rl_engine/integrations/__init__.py b/rl_engine/integrations/__init__.py new file mode 100644 index 00000000..24eac10c --- /dev/null +++ b/rl_engine/integrations/__init__.py @@ -0,0 +1,24 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Framework integration entry points owned by RL-Kernel.""" + +from rl_engine.integrations.ablation import ( + Implementation, + IntegrationPlan, + OperatorAblationCase, + operator_ablation_case, + operator_ablation_cases, +) +from rl_engine.integrations.megatron import MegatronIntegration +from rl_engine.integrations.vllm import VllmIntegration + +__all__ = [ + "Implementation", + "IntegrationPlan", + "MegatronIntegration", + "OperatorAblationCase", + "VllmIntegration", + "operator_ablation_case", + "operator_ablation_cases", +] diff --git a/rl_engine/integrations/ablation.py b/rl_engine/integrations/ablation.py new file mode 100644 index 00000000..a2ccdb82 --- /dev/null +++ b/rl_engine/integrations/ablation.py @@ -0,0 +1,130 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Executable production/RL-Kernel routes for the module debug matrix.""" + +from __future__ import annotations + +from dataclasses import dataclass +from enum import Enum +from types import MappingProxyType +from typing import Any, Mapping + +from rl_engine.alignment.cross_config.debug_matrix import module_debug_matrix + + +class Implementation(str, Enum): + PRODUCTION = "production" + RL_KERNEL = "rl_kernel" + + +@dataclass(frozen=True) +class OperatorAblationCase: + module: str + case_id: str + training: Implementation + rollout: Implementation + purpose: str + diagnostic_axes: tuple[str, ...] + + def implementation_for(self, target: str) -> Implementation: + normalized = target.strip().lower() + if normalized == "training": + return self.training + if normalized == "rollout": + return self.rollout + raise ValueError("target must be 'training' or 'rollout'") + + def to_dict(self) -> dict[str, Any]: + return { + "module": self.module, + "case_id": self.case_id, + "training_implementation": self.training.value, + "rollout_implementation": self.rollout.value, + "purpose": self.purpose, + "diagnostic_axes": list(self.diagnostic_axes), + } + + +_CASE_DEFINITIONS = ( + ("P/P", Implementation.PRODUCTION, Implementation.PRODUCTION, "native baseline"), + ("R/R", Implementation.RL_KERNEL, Implementation.RL_KERNEL, "RL-Kernel control"), + ("P/R", Implementation.PRODUCTION, Implementation.RL_KERNEL, "rollout-only mismatch"), + ("R/P", Implementation.RL_KERNEL, Implementation.PRODUCTION, "training-only mismatch"), +) + + +def operator_ablation_cases(module: str) -> tuple[OperatorAblationCase, ...]: + normalized = module.strip().lower() + matrix = module_debug_matrix() + try: + axes = tuple(str(axis["id"]) for axis in matrix["modules"][normalized]["axes"]) + except KeyError as exc: + raise ValueError(f"unknown ablation module {module!r}") from exc + return tuple( + OperatorAblationCase(normalized, case_id, training, rollout, purpose, axes) + for case_id, training, rollout, purpose in _CASE_DEFINITIONS + ) + + +def operator_ablation_case(module: str, case_id: str) -> OperatorAblationCase: + normalized = case_id.strip().upper() + for case in operator_ablation_cases(module): + if case.case_id == normalized: + return case + raise ValueError(f"unknown ablation case {case_id!r}") + + +@dataclass(frozen=True) +class IntegrationPlan: + """One independently selectable P/R case for Attention, FFN, and Logp.""" + + cases: Mapping[str, OperatorAblationCase] + + def __post_init__(self) -> None: + normalized = dict(self.cases) + if set(normalized) != {"attention", "ffn", "logp"}: + raise ValueError("integration plan must define attention, ffn, and logp") + for module, case in normalized.items(): + if not isinstance(case, OperatorAblationCase) or case.module != module: + raise ValueError(f"invalid integration case for {module!r}") + object.__setattr__(self, "cases", MappingProxyType(normalized)) + + @classmethod + def from_case_ids( + cls, + *, + attention: str = "P/P", + ffn: str = "P/P", + logp: str = "P/P", + ) -> "IntegrationPlan": + return cls( + { + "attention": operator_ablation_case("attention", attention), + "ffn": operator_ablation_case("ffn", ffn), + "logp": operator_ablation_case("logp", logp), + } + ) + + def implementation_for(self, module: str, target: str) -> Implementation: + try: + case = self.cases[module.strip().lower()] + except KeyError as exc: + raise ValueError(f"unknown integration module {module!r}") from exc + return case.implementation_for(target) + + def to_dict(self) -> dict[str, Any]: + matrix = module_debug_matrix() + return { + "schema_version": matrix["schema_version"], + "cases": {module: case.to_dict() for module, case in self.cases.items()}, + } + + +__all__ = [ + "Implementation", + "IntegrationPlan", + "OperatorAblationCase", + "operator_ablation_case", + "operator_ablation_cases", +] diff --git a/rl_engine/integrations/megatron.py b/rl_engine/integrations/megatron.py new file mode 100644 index 00000000..104bf623 --- /dev/null +++ b/rl_engine/integrations/megatron.py @@ -0,0 +1,39 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Megatron-side operator boundary owned by RL-Kernel.""" + +from __future__ import annotations + +from collections.abc import Callable, Mapping +from typing import Any + +from rl_engine.integrations.ablation import IntegrationPlan +from rl_engine.integrations.runtime import FrameworkOperatorIntegration + + +class MegatronIntegration(FrameworkOperatorIntegration): + def __init__( + self, + plan: IntegrationPlan, + *, + rl_kernel_operators: Mapping[str, Callable[..., Any]], + ) -> None: + super().__init__( + framework="megatron", + target="training", + plan=plan, + rl_kernel_operators=rl_kernel_operators, + ) + + def attention(self, native: Callable[..., Any], *args: Any, **kwargs: Any) -> Any: + return self.execute("attention", native, *args, **kwargs) + + def ffn(self, native: Callable[..., Any], *args: Any, **kwargs: Any) -> Any: + return self.execute("ffn", native, *args, **kwargs) + + def logp(self, native: Callable[..., Any], *args: Any, **kwargs: Any) -> Any: + return self.execute("logp", native, *args, **kwargs) + + +__all__ = ["MegatronIntegration"] diff --git a/rl_engine/integrations/runtime.py b/rl_engine/integrations/runtime.py new file mode 100644 index 00000000..7ede1224 --- /dev/null +++ b/rl_engine/integrations/runtime.py @@ -0,0 +1,106 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Fail-closed operator routing shared by framework integrations.""" + +from __future__ import annotations + +from collections import Counter +from collections.abc import Callable, Mapping +from dataclasses import dataclass +from threading import Lock +from typing import Any + +from rl_engine.integrations.ablation import Implementation, IntegrationPlan + + +@dataclass(frozen=True) +class OperatorReadback: + framework: str + target: str + module: str + case_id: str + implementation: str + backend_id: str + call_count: int + + def to_dict(self) -> dict[str, Any]: + return self.__dict__.copy() + + +class FrameworkOperatorIntegration: + """Route framework calls without importing or modifying framework packages.""" + + def __init__( + self, + *, + framework: str, + target: str, + plan: IntegrationPlan, + rl_kernel_operators: Mapping[str, Callable[..., Any]], + ) -> None: + self.framework = framework + self.target = target + self.plan = plan + self._rl_kernel_operators = dict(rl_kernel_operators) + self._counts: Counter[str] = Counter() + self._readbacks: dict[str, OperatorReadback] = {} + self._lock = Lock() + + def execute( + self, + module: str, + native: Callable[..., Any], + *args: Any, + **kwargs: Any, + ) -> Any: + normalized = module.strip().lower() + if not callable(native): + raise TypeError("native operator must be callable") + implementation = self.plan.implementation_for(normalized, self.target) + selected: Callable[..., Any] + if implementation is Implementation.PRODUCTION: + selected = native + else: + rl_kernel_operator = self._rl_kernel_operators.get(normalized) + if rl_kernel_operator is None: + raise RuntimeError( + f"{self.framework} {normalized} selected RL-Kernel " + "but no operator was installed" + ) + selected = rl_kernel_operator + result = selected(*args, **kwargs) + backend_id = getattr(selected, "backend_id", None) + if not isinstance(backend_id, str) or not backend_id.strip(): + backend_id = ( + f"{self.framework}.production.{normalized}" + if implementation is Implementation.PRODUCTION + else f"rlkernel.{normalized}.unidentified" + ) + with self._lock: + self._counts[normalized] += 1 + case = self.plan.cases[normalized] + self._readbacks[normalized] = OperatorReadback( + framework=self.framework, + target=self.target, + module=normalized, + case_id=case.case_id, + implementation=implementation.value, + backend_id=backend_id, + call_count=self._counts[normalized], + ) + return result + + def readback(self) -> dict[str, Any]: + with self._lock: + return { + "framework": self.framework, + "target": self.target, + "plan": self.plan.to_dict(), + "operators": { + module: readback.to_dict() for module, readback in self._readbacks.items() + }, + } + + +__all__ = ["FrameworkOperatorIntegration", "OperatorReadback"] diff --git a/rl_engine/integrations/vime/__init__.py b/rl_engine/integrations/vime/__init__.py new file mode 100644 index 00000000..3c35c88f --- /dev/null +++ b/rl_engine/integrations/vime/__init__.py @@ -0,0 +1,8 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Vime adapter entry points without a Vime runtime dependency.""" + +from .logp import ProviderResult, SelectedLogprobProviderUnavailable, provider + +__all__ = ["ProviderResult", "SelectedLogprobProviderUnavailable", "provider"] diff --git a/rl_engine/integrations/vime/logp.py b/rl_engine/integrations/vime/logp.py new file mode 100644 index 00000000..ef1b6d19 --- /dev/null +++ b/rl_engine/integrations/vime/logp.py @@ -0,0 +1,263 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Runtime selected-logprob provider for Vime's Megatron backend. + +The adapter intentionally accepts and returns structural objects: RL-Kernel +never imports Vime. Vime remains responsible for constructing locally owned +CP token rows and response masks; this provider owns only the TP-vocabulary +reduction. CP rank and layout are recorded and validated as row ownership +metadata, never passed to the numerical merge. +""" + +from __future__ import annotations + +import os +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Any + +import torch + +from rl_engine.kernels.logprob_contract import ( + LogprobContract, + LogprobDType, + LogprobRole, + MaskSpec, + ReductionSpec, + ShardingSpec, +) +from rl_engine.kernels.ops.pytorch.loss.vocab_parallel_logp import ( + BACKEND_ID, + DEFAULT_NUM_VOCAB_TILES, +) +from rl_engine.kernels.registry import kernel_registry + + +class SelectedLogprobProviderUnavailable(RuntimeError): + """Request Vime's native provider fallback in ``auto`` mode. + + Vime recognizes the marker instead of importing this class, which keeps + the dependency direction from Vime to RL-Kernel at runtime only. + """ + + selected_logprob_provider_unavailable = True + + +@dataclass(frozen=True) +class ProviderResult: + """Structural result understood by the Vime provider boundary.""" + + selected_logprobs: torch.Tensor + entropy: torch.Tensor | None + backend_id: str + contract_id: str + provenance: Mapping[str, Any] + + +def _as_positive_int(value: Any, name: str) -> int: + if isinstance(value, bool) or not isinstance(value, int) or value <= 0: + raise SelectedLogprobProviderUnavailable( + f"{name} must be a positive integer; got {value!r}" + ) + return value + + +def _metadata(request: Any) -> Mapping[str, Any]: + value = getattr(request, "metadata", None) + if not isinstance(value, Mapping): + raise SelectedLogprobProviderUnavailable( + "request.metadata must provide vocab-parallel metadata" + ) + return value + + +def _request_tensor(request: Any, name: str) -> torch.Tensor: + value = getattr(request, name, None) + if not isinstance(value, torch.Tensor): + raise SelectedLogprobProviderUnavailable(f"request.{name} must be a torch.Tensor") + return value + + +def _tp_coordinates(tp_group: Any) -> tuple[int, int]: + if tp_group is not None and hasattr(tp_group, "rank") and hasattr(tp_group, "size"): + return int(tp_group.rank()), int(tp_group.size()) + + import torch.distributed as dist + + if dist.is_available() and dist.is_initialized(): + return dist.get_rank(group=tp_group), dist.get_world_size(group=tp_group) + return 0, 1 + + +def _tile_count(metadata: Mapping[str, Any], padded_vocab_size: int) -> int: + configured = metadata.get("num_vocab_tiles", os.getenv("RL_KERNEL_LOGPROB_NUM_VOCAB_TILES")) + if configured is None or configured == "": + configured = DEFAULT_NUM_VOCAB_TILES + try: + tiles = int(configured) + except (TypeError, ValueError) as exc: + raise SelectedLogprobProviderUnavailable( + f"num_vocab_tiles must be an integer; got {configured!r}" + ) from exc + if tiles <= 0 or padded_vocab_size % tiles: + raise SelectedLogprobProviderUnavailable( + f"num_vocab_tiles={tiles} must divide padded_vocab_size={padded_vocab_size}" + ) + return tiles + + +def _contract_for_request(request: Any) -> tuple[LogprobContract, int]: + logits = _request_tensor(request, "logits") + targets = _request_tensor(request, "target_ids") + metadata = _metadata(request) + if logits.ndim != 2 or targets.shape != (logits.shape[0],): + raise SelectedLogprobProviderUnavailable( + "request must contain local [T, V] logits and aligned [T] targets" + ) + if logits.dtype not in (torch.bfloat16, torch.float16, torch.float32): + raise SelectedLogprobProviderUnavailable(f"unsupported logit dtype {logits.dtype}") + if targets.device != logits.device: + raise SelectedLogprobProviderUnavailable("target_ids must share the local logits device") + + cp = getattr(request, "context_parallel", None) + cp_world_size = _as_positive_int(getattr(cp, "world_size", None), "context_parallel.world_size") + cp_rank = getattr(cp, "rank", None) + if ( + isinstance(cp_rank, bool) + or not isinstance(cp_rank, int) + or not 0 <= cp_rank < cp_world_size + ): + raise SelectedLogprobProviderUnavailable( + f"context_parallel.rank={cp_rank!r} is invalid for CP={cp_world_size}" + ) + if getattr(cp, "layout", None) not in ( + {"single"} if cp_world_size == 1 else {"zigzag", "allgather"} + ): + raise SelectedLogprobProviderUnavailable( + "context_parallel layout does not describe local CP token ownership" + ) + + tp_rank, tp_world_size = _tp_coordinates(getattr(request, "tensor_parallel_group", None)) + declared_tp_rank = metadata.get("tp_rank") + declared_tp_world_size = metadata.get("tp_world_size") + if declared_tp_rank is not None and declared_tp_rank != tp_rank: + raise SelectedLogprobProviderUnavailable( + f"metadata tp_rank={declared_tp_rank} disagrees with TP group rank={tp_rank}" + ) + if declared_tp_world_size is not None and declared_tp_world_size != tp_world_size: + raise SelectedLogprobProviderUnavailable( + f"metadata tp_world_size={declared_tp_world_size} disagrees with " + f"TP group size={tp_world_size}" + ) + + real_vocab_size = _as_positive_int(metadata.get("real_vocab_size"), "real_vocab_size") + padded_vocab_size = _as_positive_int(metadata.get("padded_vocab_size"), "padded_vocab_size") + if logits.shape[1] * tp_world_size != padded_vocab_size: + raise SelectedLogprobProviderUnavailable( + "local vocab width and TP group do not cover padded_vocab_size exactly: " + f"{logits.shape[1]} * {tp_world_size} != {padded_vocab_size}" + ) + if real_vocab_size > padded_vocab_size: + raise SelectedLogprobProviderUnavailable( + "real_vocab_size must not exceed padded_vocab_size" + ) + + bounds = tuple( + (rank * logits.shape[1], (rank + 1) * logits.shape[1]) for rank in range(tp_world_size) + ) + active_mask = (True,) * logits.shape[0] + contract = LogprobContract( + role=LogprobRole.TRAIN, + dtype={ + torch.bfloat16: LogprobDType.BF16, + torch.float16: LogprobDType.FP16, + torch.float32: LogprobDType.FP32, + }[logits.dtype], + mask=MaskSpec(num_tokens=logits.shape[0], active_mask=active_mask), + sharding=ShardingSpec( + tp_rank=tp_rank, + tp_world_size=tp_world_size, + vocab_shard_bounds=bounds, + real_vocab_size=real_vocab_size, + padded_vocab_size=padded_vocab_size, + cp_rank=cp_rank, + cp_world_size=cp_world_size, + ), + reduction=ReductionSpec(), + ) + return contract, _tile_count(metadata, padded_vocab_size) + + +def provider(request: Any) -> ProviderResult: + """Compute Vime selected logprobs on the explicit WS2 TP/CP contract. + + Top-p replay is deliberately unavailable until it has a separately + validated fixed-order mask contract. In Vime ``auto`` mode this signals + native execution; in ``strict`` mode it fails instead of changing sampled + distribution semantics. + """ + + if getattr(request, "log_prob_keep_mask", None) is not None: + raise SelectedLogprobProviderUnavailable( + "RL-Kernel WS2 logprob does not yet materialize Vime top-p replay masks" + ) + + contract, num_vocab_tiles = _contract_for_request(request) + dispatch = kernel_registry.get_logprob_op(contract, requested_backend=BACKEND_ID) + if dispatch.provenance["actual_backend"] != BACKEND_ID or dispatch.provenance["fallback"]: + raise RuntimeError("explicit WS2 backend dispatch changed during materialization") + if getattr(request, "with_entropy", False): + selected_logp, _lse, entropy = dispatch.op.apply_with_entropy( + request.logits, + request.target_ids, + contract=contract, + tp_group=getattr(request, "tensor_parallel_group", None), + num_vocab_tiles=num_vocab_tiles, + with_entropy_grad=bool(getattr(request, "with_entropy_grad", False)), + ) + else: + selected_logp, _lse = dispatch.op( + request.logits, + request.target_ids, + contract=contract, + tp_group=getattr(request, "tensor_parallel_group", None), + num_vocab_tiles=num_vocab_tiles, + ) + entropy = None + provenance = dict(dispatch.provenance) + provenance["request"] = { + "logits_shape": list(request.logits.shape), + "logits_dtype": str(request.logits.dtype).replace("torch.", ""), + "target_shape": list(request.target_ids.shape), + "target_dtype": str(request.target_ids.dtype).replace("torch.", ""), + "real_vocab_size": contract.sharding.real_vocab_size, + "padded_vocab_size": contract.sharding.padded_vocab_size, + "tp_rank": contract.sharding.tp_rank, + "tp_world_size": contract.sharding.tp_world_size, + "cp_rank": contract.sharding.cp_rank, + "cp_world_size": contract.sharding.cp_world_size, + } + provenance["execution"] = { + "role": "vime_training_selected_logprob", + "strict_backend": True, + "top_p_replay": False, + } + provenance["cp_row_ownership"] = { + "cp_rank": contract.sharding.cp_rank, + "cp_world_size": contract.sharding.cp_world_size, + "layout": getattr(request.context_parallel, "layout"), + "local_token_rows": int(request.logits.shape[0]), + "cp_is_merge_axis": False, + } + provenance["num_vocab_tiles"] = num_vocab_tiles + return ProviderResult( + selected_logprobs=selected_logp.unsqueeze(-1), + entropy=entropy, + backend_id=dispatch.capability.backend_id, + contract_id=contract.cross_rank_fingerprint(), + provenance=provenance, + ) + + +__all__ = ["ProviderResult", "SelectedLogprobProviderUnavailable", "provider"] diff --git a/rl_engine/integrations/vllm.py b/rl_engine/integrations/vllm.py new file mode 100644 index 00000000..1d760850 --- /dev/null +++ b/rl_engine/integrations/vllm.py @@ -0,0 +1,39 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""vLLM-side operator boundary owned by RL-Kernel.""" + +from __future__ import annotations + +from collections.abc import Callable, Mapping +from typing import Any + +from rl_engine.integrations.ablation import IntegrationPlan +from rl_engine.integrations.runtime import FrameworkOperatorIntegration + + +class VllmIntegration(FrameworkOperatorIntegration): + def __init__( + self, + plan: IntegrationPlan, + *, + rl_kernel_operators: Mapping[str, Callable[..., Any]], + ) -> None: + super().__init__( + framework="vllm", + target="rollout", + plan=plan, + rl_kernel_operators=rl_kernel_operators, + ) + + def attention(self, native: Callable[..., Any], *args: Any, **kwargs: Any) -> Any: + return self.execute("attention", native, *args, **kwargs) + + def ffn(self, native: Callable[..., Any], *args: Any, **kwargs: Any) -> Any: + return self.execute("ffn", native, *args, **kwargs) + + def logp(self, native: Callable[..., Any], *args: Any, **kwargs: Any) -> Any: + return self.execute("logp", native, *args, **kwargs) + + +__all__ = ["VllmIntegration"] diff --git a/rl_engine/kernels/attention_contract.py b/rl_engine/kernels/attention_contract.py new file mode 100644 index 00000000..8b5f588b --- /dev/null +++ b/rl_engine/kernels/attention_contract.py @@ -0,0 +1,1532 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Typed WS2 contract for context-parallel standard softmax attention. + +The objects in this module describe a distributed attention invocation. They +do not shard tensors, launch collectives, or implement the ``(out, lse)`` +merge. Keeping description and materialization separate lets dispatch reject +an incompatible backend before any numerically different path is launched. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from enum import Enum +from typing import Any, Iterable, TypeVar + +_EnumT = TypeVar("_EnumT", bound=Enum) + + +STRICT_ATTENTION_CORE_ID = "rlkernel.attention.deterministic_core.v1" +STRICT_ATTENTION_SCHEDULE_ID = "single_batch_single_query_global_kv_blocks" + + +class AttentionContractError(ValueError): + """Raised when attention metadata does not describe a valid invocation.""" + + +class AttentionRole(str, Enum): + TRAIN = "train" + INFER = "infer" + + +class AttentionMode(str, Enum): + PREFILL = "prefill" + CHUNKED_PREFILL = "chunked_prefill" + DECODE = "decode" + + +class AttentionDType(str, Enum): + BF16 = "bf16" + FP16 = "fp16" + FP32 = "fp32" + + +class AttentionMerge(str, Enum): + ONLINE_SOFTMAX_LSE = "online_softmax_lse" + + +class ReductionOrder(str, Enum): + GLOBAL_BLOCK_INDEX = "global_block_index" + + +class DowncastPoint(str, Enum): + FINAL_WRITE = "final_write" + + +class ReductionEngine(str, Enum): + IN_OP_REFERENCE = "in_op_reference" + + +class SplitKVMode(str, Enum): + DISABLED = "disabled" + FIXED = "fixed" + AUTO = "auto" + + +class RoPEState(str, Enum): + PRE_ROPE = "pre_rope" + POST_ROPE = "post_rope" + + +class RoPECastPoint(str, Enum): + NONE = "none" + BEFORE_ROPE = "before_rope" + AFTER_ROPE = "after_rope" + + +class RoPEFusionBoundary(str, Enum): + UNFUSED_ROPE_ATTENTION = "unfused_rope_attention" + FUSED_ROPE_ATTENTION = "fused_rope_attention" + + +def _enum_value(enum_type: type[_EnumT], value: Any, field: str) -> _EnumT: + try: + return enum_type(value) + except (TypeError, ValueError) as exc: + allowed = ", ".join(item.value for item in enum_type) + raise AttentionContractError(f"{field} must be one of: {allowed}; got {value!r}") from exc + + +def _positive_int(value: Any, field: str) -> int: + if isinstance(value, bool) or not isinstance(value, int) or value <= 0: + raise AttentionContractError(f"{field} must be a positive integer; got {value!r}") + return value + + +def _non_negative_int(value: Any, field: str) -> int: + if isinstance(value, bool) or not isinstance(value, int) or value < 0: + raise AttentionContractError(f"{field} must be a non-negative integer; got {value!r}") + return value + + +def _integer_tuple(values: Iterable[int], field: str) -> tuple[int, ...]: + try: + result = tuple(values) + except TypeError as exc: + raise AttentionContractError(f"{field} must be an iterable of integers") from exc + for index, value in enumerate(result): + if isinstance(value, bool) or not isinstance(value, int): + raise AttentionContractError(f"{field}[{index}] must be an integer; got {value!r}") + return result + + +@dataclass(frozen=True) +class ShardingSpec: + """Logical TP/CP ownership for one attention invocation. + + TP head shards are currently required to be equal and contiguous. CP + sequence ownership may be uneven, but every local block must carry a stable + logical global index so a later implementation can merge by logical order + instead of collective arrival order. + """ + + tp_rank: int + tp_world_size: int + cp_rank: int + cp_world_size: int + global_q_heads: int + global_kv_heads: int + local_q_head_start: int + local_q_heads: int + local_kv_head_start: int + local_kv_heads: int + global_sequence_length: int + local_sequence_length: int + global_block_indices: tuple[int, ...] + global_block_token_starts: tuple[int, ...] + local_block_offsets: tuple[int, ...] + packed_sequence_offsets: tuple[int, ...] | None = None + + def __post_init__(self) -> None: + tp_world_size = _positive_int(self.tp_world_size, "tp_world_size") + cp_world_size = _positive_int(self.cp_world_size, "cp_world_size") + tp_rank = _non_negative_int(self.tp_rank, "tp_rank") + cp_rank = _non_negative_int(self.cp_rank, "cp_rank") + if tp_rank >= tp_world_size: + raise AttentionContractError( + f"tp_rank={tp_rank} must be smaller than tp_world_size={tp_world_size}" + ) + if cp_rank >= cp_world_size: + raise AttentionContractError( + f"cp_rank={cp_rank} must be smaller than cp_world_size={cp_world_size}" + ) + + global_q_heads = _positive_int(self.global_q_heads, "global_q_heads") + global_kv_heads = _positive_int(self.global_kv_heads, "global_kv_heads") + if global_q_heads % global_kv_heads != 0: + raise AttentionContractError( + f"global_q_heads={global_q_heads} must be divisible by " + f"global_kv_heads={global_kv_heads} for GQA" + ) + if global_q_heads % tp_world_size != 0 or global_kv_heads % tp_world_size != 0: + raise AttentionContractError( + "global Q/KV heads must be evenly divisible by tp_world_size; " + f"got Hq={global_q_heads}, Hkv={global_kv_heads}, TP={tp_world_size}" + ) + + local_q_heads = _positive_int(self.local_q_heads, "local_q_heads") + local_kv_heads = _positive_int(self.local_kv_heads, "local_kv_heads") + expected_q_heads = global_q_heads // tp_world_size + expected_kv_heads = global_kv_heads // tp_world_size + if local_q_heads != expected_q_heads or local_kv_heads != expected_kv_heads: + raise AttentionContractError( + "local TP head counts do not preserve the global Q/KV mapping; " + f"expected ({expected_q_heads}, {expected_kv_heads}), got " + f"({local_q_heads}, {local_kv_heads})" + ) + + local_q_head_start = _non_negative_int(self.local_q_head_start, "local_q_head_start") + local_kv_head_start = _non_negative_int(self.local_kv_head_start, "local_kv_head_start") + expected_q_start = tp_rank * expected_q_heads + expected_kv_start = tp_rank * expected_kv_heads + if local_q_head_start != expected_q_start or local_kv_head_start != expected_kv_start: + raise AttentionContractError( + "local TP head starts do not match contiguous rank ownership; " + f"expected ({expected_q_start}, {expected_kv_start}), got " + f"({local_q_head_start}, {local_kv_head_start})" + ) + + global_sequence_length = _positive_int( + self.global_sequence_length, "global_sequence_length" + ) + local_sequence_length = _positive_int(self.local_sequence_length, "local_sequence_length") + + block_indices = _integer_tuple(self.global_block_indices, "global_block_indices") + if not block_indices: + raise AttentionContractError("global_block_indices must not be empty") + if any(index < 0 for index in block_indices): + raise AttentionContractError("global_block_indices must be non-negative") + if any( + left >= right for left, right in zip(block_indices, block_indices[1:], strict=False) + ): + raise AttentionContractError( + "global_block_indices must be unique and strictly increasing" + ) + object.__setattr__(self, "global_block_indices", block_indices) + + block_token_starts = _integer_tuple( + self.global_block_token_starts, "global_block_token_starts" + ) + local_block_offsets = _integer_tuple(self.local_block_offsets, "local_block_offsets") + if len(block_token_starts) != len(block_indices): + raise AttentionContractError( + "global_block_token_starts must contain one entry per global_block_indices entry" + ) + if any(start < 0 for start in block_token_starts): + raise AttentionContractError("global_block_token_starts must be non-negative") + if len(local_block_offsets) != len(block_indices) + 1: + raise AttentionContractError( + "local_block_offsets must contain one boundary more than global_block_indices" + ) + if local_block_offsets[0] != 0 or local_block_offsets[-1] != local_sequence_length: + raise AttentionContractError( + "local_block_offsets must start at 0 and end at local_sequence_length" + ) + if any( + left >= right + for left, right in zip(local_block_offsets, local_block_offsets[1:], strict=False) + ): + raise AttentionContractError("local_block_offsets must be strictly increasing") + + previous_global_end = 0 + for index, (global_start, local_start, local_end) in enumerate( + zip( + block_token_starts, + local_block_offsets[:-1], + local_block_offsets[1:], + strict=True, + ) + ): + global_end = global_start + (local_end - local_start) + if global_end > global_sequence_length: + raise AttentionContractError( + f"global block {block_indices[index]} exceeds global_sequence_length" + ) + if index > 0 and global_start < previous_global_end: + raise AttentionContractError( + "global block token ranges must be non-overlapping and ordered" + ) + previous_global_end = global_end + object.__setattr__(self, "global_block_token_starts", block_token_starts) + object.__setattr__(self, "local_block_offsets", local_block_offsets) + + if self.packed_sequence_offsets is not None: + offsets = _integer_tuple(self.packed_sequence_offsets, "packed_sequence_offsets") + if len(offsets) < 2 or offsets[0] != 0: + raise AttentionContractError( + "packed_sequence_offsets must start at 0 and contain an end offset" + ) + if any(left >= right for left, right in zip(offsets, offsets[1:], strict=False)): + raise AttentionContractError("packed_sequence_offsets must be strictly increasing") + if offsets[-1] != local_sequence_length: + raise AttentionContractError( + "the final packed_sequence_offsets value must equal local_sequence_length; " + f"got {offsets[-1]} and {local_sequence_length}" + ) + object.__setattr__(self, "packed_sequence_offsets", offsets) + + +@dataclass(frozen=True) +class ReductionSpec: + """Deterministic CP ``(out, lse)`` merge semantics.""" + + merge: AttentionMerge = AttentionMerge.ONLINE_SOFTMAX_LSE + acc_dtype: AttentionDType = AttentionDType.FP32 + order: ReductionOrder = ReductionOrder.GLOBAL_BLOCK_INDEX + downcast_at: DowncastPoint = DowncastPoint.FINAL_WRITE + engine: ReductionEngine = ReductionEngine.IN_OP_REFERENCE + + def __post_init__(self) -> None: + object.__setattr__(self, "merge", _enum_value(AttentionMerge, self.merge, "merge")) + object.__setattr__( + self, "acc_dtype", _enum_value(AttentionDType, self.acc_dtype, "acc_dtype") + ) + object.__setattr__(self, "order", _enum_value(ReductionOrder, self.order, "order")) + object.__setattr__( + self, "downcast_at", _enum_value(DowncastPoint, self.downcast_at, "downcast_at") + ) + object.__setattr__(self, "engine", _enum_value(ReductionEngine, self.engine, "engine")) + if self.acc_dtype is not AttentionDType.FP32: + raise AttentionContractError( + f"CP attention accumulation must be fp32; got {self.acc_dtype.value}" + ) + + +@dataclass(frozen=True) +class SplitKVExecutionPlan: + """Actual backend-local Split-KV schedule emitted by a runtime. + + CP ownership is intentionally separate from this plan. ``boundaries`` + describe the canonical logical KV ranges reduced by one backend invocation; + CP may transport those partial states between ranks but must not change the + FP32 merge contract recorded here. + """ + + requested_mode: SplitKVMode + requested_split_size: int | None + actual_mode: SplitKVMode | None + actual_split_size: int | None + boundaries: tuple[tuple[int, int], ...] + merge_order: ReductionOrder = ReductionOrder.GLOBAL_BLOCK_INDEX + acc_dtype: AttentionDType = AttentionDType.FP32 + downcast_at: DowncastPoint = DowncastPoint.FINAL_WRITE + backend: str = "reference" + source: str = "contract_exact" + fallback: bool = False + fallback_reason: str | None = None + + def __post_init__(self) -> None: + object.__setattr__( + self, + "requested_mode", + _enum_value(SplitKVMode, self.requested_mode, "requested_mode"), + ) + if self.actual_mode is not None: + object.__setattr__( + self, + "actual_mode", + _enum_value(SplitKVMode, self.actual_mode, "actual_mode"), + ) + for field_name in ("requested_split_size", "actual_split_size"): + value = getattr(self, field_name) + if value is not None: + _positive_int(value, field_name) + if self.requested_mode is SplitKVMode.FIXED: + if self.requested_split_size is None: + raise AttentionContractError( + "requested fixed Split-KV mode requires requested_split_size" + ) + elif self.requested_split_size is not None: + raise AttentionContractError( + "requested_split_size is only valid for requested fixed Split-KV mode" + ) + try: + boundaries = tuple(tuple(boundary) for boundary in self.boundaries) + except TypeError as exc: + raise AttentionContractError( + "Split-KV boundaries must be an iterable of (start, end) pairs" + ) from exc + previous_end = 0 + for index, boundary in enumerate(boundaries): + if len(boundary) != 2: + raise AttentionContractError( + f"Split-KV boundary {index} must contain exactly start and end" + ) + start, end = boundary + if ( + isinstance(start, bool) + or isinstance(end, bool) + or not isinstance(start, int) + or not isinstance(end, int) + or start < 0 + or end <= start + ): + raise AttentionContractError("Split-KV boundaries must satisfy 0 <= start < end") + if index > 0 and start != previous_end: + raise AttentionContractError( + "Split-KV boundaries must be contiguous and in logical KV order" + ) + previous_end = end + if self.actual_mode is None and boundaries: + raise AttentionContractError("unknown actual Split-KV plan cannot declare boundaries") + if self.actual_mode is not None and not boundaries: + raise AttentionContractError("actual Split-KV plan must declare logical boundaries") + if self.actual_mode is SplitKVMode.FIXED: + if self.actual_split_size is None: + raise AttentionContractError( + "actual fixed Split-KV mode requires actual_split_size" + ) + widths = tuple(end - start for start, end in boundaries) + if any(width != self.actual_split_size for width in widths[:-1]) or ( + widths and widths[-1] > self.actual_split_size + ): + raise AttentionContractError( + "fixed Split-KV boundaries must use actual_split_size except " + "for a shorter final split" + ) + elif self.actual_split_size is not None: + raise AttentionContractError( + "actual_split_size is only valid for actual fixed Split-KV mode" + ) + if self.actual_mode is SplitKVMode.DISABLED and len(boundaries) != 1: + raise AttentionContractError( + "disabled Split-KV execution must contain exactly one boundary" + ) + object.__setattr__(self, "boundaries", boundaries) + object.__setattr__( + self, + "merge_order", + _enum_value(ReductionOrder, self.merge_order, "merge_order"), + ) + object.__setattr__( + self, + "acc_dtype", + _enum_value(AttentionDType, self.acc_dtype, "acc_dtype"), + ) + object.__setattr__( + self, + "downcast_at", + _enum_value(DowncastPoint, self.downcast_at, "downcast_at"), + ) + if self.acc_dtype is not AttentionDType.FP32: + raise AttentionContractError("Split-KV partial states must be merged in fp32") + if not isinstance(self.backend, str) or not self.backend.strip(): + raise AttentionContractError("Split-KV backend must be a non-empty string") + if not isinstance(self.source, str) or not self.source.strip(): + raise AttentionContractError("Split-KV plan source must be a non-empty string") + if not isinstance(self.fallback, bool): + raise AttentionContractError("Split-KV fallback must be a bool") + if self.fallback and not self.fallback_reason: + raise AttentionContractError("Split-KV fallback_reason is required for a fallback") + if not self.fallback and self.fallback_reason is not None: + raise AttentionContractError( + "Split-KV fallback_reason must be None when fallback=False" + ) + if not self.fallback and self.actual_mode is not None: + if self.actual_mode is not self.requested_mode: + raise AttentionContractError( + "actual Split-KV mode may differ from requested mode only for a fallback" + ) + if self.actual_split_size != self.requested_split_size: + raise AttentionContractError( + "actual Split-KV size may differ from requested size only for a fallback" + ) + + @property + def actual_split_count(self) -> int | None: + return None if self.actual_mode is None else len(self.boundaries) + + def to_dict(self) -> dict[str, Any]: + return { + "requested_split_kv_policy": self.requested_mode.value, + "requested_split_kv_size": self.requested_split_size, + "actual_split_kv_policy": ( + None if self.actual_mode is None else self.actual_mode.value + ), + "actual_split_kv_size": self.actual_split_size, + "actual_split_kv_count": self.actual_split_count, + "actual_split_boundaries": [list(boundary) for boundary in self.boundaries], + "split_kv_merge_order": self.merge_order.value, + "split_kv_accum_dtype": self.acc_dtype.value, + "split_kv_downcast_at": self.downcast_at.value, + "split_kv_backend": self.backend, + "split_kv_plan_source": self.source, + "split_kv_fallback": self.fallback, + "split_kv_fallback_reason": self.fallback_reason, + } + + +@dataclass(frozen=True) +class SplitKVSpec: + """Requested Split-KV policy shared by training and rollout paths.""" + + mode: SplitKVMode = SplitKVMode.DISABLED + fixed_split_size: int | None = None + strict_consistency: bool = True + + @classmethod + def disabled(cls, *, strict_consistency: bool = True) -> "SplitKVSpec": + return cls(mode=SplitKVMode.DISABLED, strict_consistency=strict_consistency) + + @classmethod + def fixed(cls, split_size: int, *, strict_consistency: bool = True) -> "SplitKVSpec": + return cls( + mode=SplitKVMode.FIXED, + fixed_split_size=split_size, + strict_consistency=strict_consistency, + ) + + @classmethod + def auto(cls, *, strict_consistency: bool = False) -> "SplitKVSpec": + return cls(mode=SplitKVMode.AUTO, strict_consistency=strict_consistency) + + def __post_init__(self) -> None: + object.__setattr__(self, "mode", _enum_value(SplitKVMode, self.mode, "split_kv.mode")) + if not isinstance(self.strict_consistency, bool): + raise AttentionContractError("split_kv.strict_consistency must be a bool") + if self.mode is SplitKVMode.FIXED: + if self.fixed_split_size is None: + raise AttentionContractError("fixed Split-KV policy requires fixed_split_size") + _positive_int(self.fixed_split_size, "split_kv.fixed_split_size") + elif self.fixed_split_size is not None: + raise AttentionContractError("fixed_split_size is only valid for fixed Split-KV policy") + if self.strict_consistency and self.mode is SplitKVMode.AUTO: + raise AttentionContractError( + "auto Split-KV is runtime-shape dependent and is not allowed in strict consistency" + ) + + def resolve(self, total_kv_tokens: int, *, backend: str) -> SplitKVExecutionPlan: + """Resolve policies whose logical schedule is fully known by contract.""" + + total_kv_tokens = _positive_int(total_kv_tokens, "total_kv_tokens") + if self.mode is SplitKVMode.AUTO: + return SplitKVExecutionPlan( + requested_mode=self.mode, + requested_split_size=None, + actual_mode=None, + actual_split_size=None, + boundaries=(), + backend=backend, + source="runtime_required", + ) + split_size = total_kv_tokens if self.mode is SplitKVMode.DISABLED else self.fixed_split_size + assert split_size is not None + boundaries = tuple( + (start, min(start + split_size, total_kv_tokens)) + for start in range(0, total_kv_tokens, split_size) + ) + return SplitKVExecutionPlan( + requested_mode=self.mode, + requested_split_size=self.fixed_split_size, + actual_mode=self.mode, + actual_split_size=self.fixed_split_size, + boundaries=boundaries, + backend=backend, + source="contract_exact", + ) + + def to_dict(self) -> dict[str, Any]: + return { + "mode": self.mode.value, + "fixed_split_size": self.fixed_split_size, + "strict_consistency": self.strict_consistency, + } + + +def validate_split_kv_alignment( + training: SplitKVExecutionPlan, + rollout: SplitKVExecutionPlan, +) -> None: + """Fail closed unless train and rollout executed the same logical plan.""" + + if training.actual_mode is None or rollout.actual_mode is None: + raise AttentionContractError( + "strict Split-KV alignment requires actual runtime plans from both sides" + ) + fields = ( + "requested_mode", + "requested_split_size", + "actual_mode", + "actual_split_size", + "boundaries", + "merge_order", + "acc_dtype", + "downcast_at", + "fallback", + ) + mismatches = [ + field_name + for field_name in fields + if getattr(training, field_name) != getattr(rollout, field_name) + ] + if mismatches: + raise AttentionContractError( + "training/rollout Split-KV execution plans differ: " + ", ".join(mismatches) + ) + + +@dataclass(frozen=True, order=True) +class SplitKVRuntimeCoordinate: + """Identity of one batch/rank/owner Split-KV runtime plan.""" + + batch_index: int + tp_rank: int + cp_rank: int + owner_cp_rank: int + + def validate( + self, + *, + batch_size: int, + tp_world_size: int, + cp_world_size: int, + ) -> None: + batch_index = _non_negative_int(self.batch_index, "batch_index") + tp_rank = _non_negative_int(self.tp_rank, "tp_rank") + cp_rank = _non_negative_int(self.cp_rank, "cp_rank") + owner_cp_rank = _non_negative_int(self.owner_cp_rank, "owner_cp_rank") + if batch_index >= batch_size: + raise AttentionContractError("Split-KV batch_index is outside batch_size") + if tp_rank >= tp_world_size: + raise AttentionContractError("Split-KV tp_rank is outside tp_world_size") + if cp_rank >= cp_world_size: + raise AttentionContractError("Split-KV cp_rank is outside cp_world_size") + if owner_cp_rank >= cp_world_size: + raise AttentionContractError("Split-KV owner_cp_rank is outside cp_world_size") + + def to_dict(self) -> dict[str, int]: + return { + "batch_index": self.batch_index, + "tp_rank": self.tp_rank, + "cp_rank": self.cp_rank, + "owner_cp_rank": self.owner_cp_rank, + } + + +@dataclass(frozen=True) +class SplitKVRuntimePlanEntry: + """Actual plan for one batch/TP/CP consumer and logical KV owner.""" + + coordinate: SplitKVRuntimeCoordinate + expected_kv_range: tuple[int, int] + execution: SplitKVExecutionPlan + + def validate( + self, + *, + batch_size: int, + tp_world_size: int, + cp_world_size: int, + total_kv_tokens: int, + ) -> None: + self.coordinate.validate( + batch_size=batch_size, + tp_world_size=tp_world_size, + cp_world_size=cp_world_size, + ) + try: + start, end = self.expected_kv_range + except (TypeError, ValueError) as exc: + raise AttentionContractError( + "expected_kv_range must contain exactly (start, end)" + ) from exc + if ( + isinstance(start, bool) + or isinstance(end, bool) + or not isinstance(start, int) + or not isinstance(end, int) + or start < 0 + or end <= start + or end > total_kv_tokens + ): + raise AttentionContractError( + "expected_kv_range must satisfy 0 <= start < end <= total_kv_tokens" + ) + if self.execution.actual_mode is None: + raise AttentionContractError("complete Split-KV plan sets require actual runtime plans") + if self.execution.boundaries[0][0] != start or self.execution.boundaries[-1][1] != end: + raise AttentionContractError( + "Split-KV execution boundaries must exactly cover expected_kv_range" + ) + if any( + boundary_start < start or boundary_end > end + for boundary_start, boundary_end in self.execution.boundaries + ): + raise AttentionContractError("Split-KV execution boundary escapes expected_kv_range") + + def to_dict(self) -> dict[str, Any]: + return { + **self.coordinate.to_dict(), + "expected_kv_range": list(self.expected_kv_range), + **self.execution.to_dict(), + } + + +@dataclass(frozen=True) +class SplitKVRuntimePlanSet: + """Complete actual Split-KV plans across batch, TP, CP, and KV owners. + + Every CP consumer must report the plan used for every CP-owned KV range. + This duplicates owner plans across consumers intentionally: it detects one + rank silently choosing a different Split-K schedule or merge policy. + """ + + batch_size: int + tp_world_size: int + cp_world_size: int + total_kv_tokens: tuple[int, ...] + entries: tuple[SplitKVRuntimePlanEntry, ...] + + def __post_init__(self) -> None: + batch_size = _positive_int(self.batch_size, "batch_size") + tp_world_size = _positive_int(self.tp_world_size, "tp_world_size") + cp_world_size = _positive_int(self.cp_world_size, "cp_world_size") + totals = _integer_tuple(self.total_kv_tokens, "total_kv_tokens") + if len(totals) != batch_size or any(total <= 0 for total in totals): + raise AttentionContractError( + "total_kv_tokens must contain one positive length per batch item" + ) + object.__setattr__(self, "total_kv_tokens", totals) + entries = tuple(self.entries) + object.__setattr__(self, "entries", entries) + expected_coordinates = { + SplitKVRuntimeCoordinate(batch_index, tp_rank, cp_rank, owner_cp_rank) + for batch_index in range(batch_size) + for tp_rank in range(tp_world_size) + for cp_rank in range(cp_world_size) + for owner_cp_rank in range(cp_world_size) + } + actual_coordinates = [entry.coordinate for entry in entries] + if len(set(actual_coordinates)) != len(actual_coordinates): + raise AttentionContractError("Split-KV runtime plan set contains duplicate coordinates") + missing = expected_coordinates.difference(actual_coordinates) + extra = set(actual_coordinates).difference(expected_coordinates) + if missing or extra: + raise AttentionContractError( + "Split-KV runtime plan set coordinate coverage is incomplete; " + f"missing={_format_split_kv_coordinates(missing)}, " + f"extra={_format_split_kv_coordinates(extra)}" + ) + for entry in entries: + entry.validate( + batch_size=batch_size, + tp_world_size=tp_world_size, + cp_world_size=cp_world_size, + total_kv_tokens=totals[entry.coordinate.batch_index], + ) + self._validate_owner_coverage() + self._validate_rank_invariance() + + def _validate_owner_coverage(self) -> None: + for batch_index, total in enumerate(self.total_kv_tokens): + for tp_rank in range(self.tp_world_size): + ranges = [] + for owner_cp_rank in range(self.cp_world_size): + matches = [ + entry + for entry in self.entries + if entry.coordinate.batch_index == batch_index + and entry.coordinate.tp_rank == tp_rank + and entry.coordinate.cp_rank == 0 + and entry.coordinate.owner_cp_rank == owner_cp_rank + ] + ranges.append(matches[0].expected_kv_range) + previous_end = 0 + for start, end in ranges: + if start != previous_end: + raise AttentionContractError( + "Split-KV owner ranges must be gap-free in CP owner order" + ) + previous_end = end + if previous_end != total: + raise AttentionContractError( + "Split-KV owner ranges do not cover total_kv_tokens" + ) + + def _validate_rank_invariance(self) -> None: + for batch_index in range(self.batch_size): + for owner_cp_rank in range(self.cp_world_size): + entries = sorted( + ( + entry + for entry in self.entries + if entry.coordinate.batch_index == batch_index + and entry.coordinate.owner_cp_rank == owner_cp_rank + ), + key=lambda entry: ( + entry.coordinate.tp_rank, + entry.coordinate.cp_rank, + ), + ) + reference = entries[0] + for entry in entries[1:]: + if entry.expected_kv_range != reference.expected_kv_range: + raise AttentionContractError( + "Split-KV owner range differs across TP/CP consumers" + ) + try: + validate_split_kv_alignment( + reference.execution, + entry.execution, + ) + except AttentionContractError as exc: + raise AttentionContractError( + "Split-KV plan differs across TP/CP consumers for " + f"batch={batch_index}, owner_cp={owner_cp_rank}: {exc}" + ) from exc + + def to_dict(self) -> dict[str, Any]: + return { + "batch_size": self.batch_size, + "tp_world_size": self.tp_world_size, + "cp_world_size": self.cp_world_size, + "total_kv_tokens": list(self.total_kv_tokens), + "entries": [ + entry.to_dict() + for entry in sorted(self.entries, key=lambda entry: entry.coordinate) + ], + "coverage": "complete_batch_tp_cp_owner_cartesian_product", + } + + +def validate_split_kv_plan_set_alignment( + training: SplitKVRuntimePlanSet, + rollout: SplitKVRuntimePlanSet, +) -> None: + """Fail closed unless complete train/rollout runtime plan sets align.""" + + topology_fields = ( + "batch_size", + "tp_world_size", + "cp_world_size", + "total_kv_tokens", + ) + topology_mismatches = [ + field_name + for field_name in topology_fields + if getattr(training, field_name) != getattr(rollout, field_name) + ] + if topology_mismatches: + raise AttentionContractError( + "training/rollout Split-KV plan-set topology differs: " + ", ".join(topology_mismatches) + ) + training_by_coordinate = {entry.coordinate: entry for entry in training.entries} + rollout_by_coordinate = {entry.coordinate: entry for entry in rollout.entries} + if training_by_coordinate.keys() != rollout_by_coordinate.keys(): + raise AttentionContractError("training/rollout Split-KV plan-set coordinates differ") + for coordinate in sorted(training_by_coordinate): + train_entry = training_by_coordinate[coordinate] + rollout_entry = rollout_by_coordinate[coordinate] + if train_entry.expected_kv_range != rollout_entry.expected_kv_range: + raise AttentionContractError( + f"training/rollout expected KV range differs at {coordinate}" + ) + try: + validate_split_kv_alignment( + train_entry.execution, + rollout_entry.execution, + ) + except AttentionContractError as exc: + raise AttentionContractError( + f"training/rollout Split-KV plan differs at {coordinate}: {exc}" + ) from exc + + +def build_split_kv_runtime_plan_set( + total_kv_tokens: Iterable[int], + *, + tp_world_size: int, + cp_world_size: int, + split_kv: SplitKVSpec, + backend: str = "contract_reference", +) -> SplitKVRuntimePlanSet: + """Build a complete owner-local plan set for contract tests and adapters.""" + + totals = _integer_tuple(total_kv_tokens, "total_kv_tokens") + if not totals or any(total < cp_world_size for total in totals): + raise AttentionContractError( + "contract plan sets require at least one KV token per CP owner" + ) + tp_world_size = _positive_int(tp_world_size, "tp_world_size") + cp_world_size = _positive_int(cp_world_size, "cp_world_size") + if not isinstance(split_kv, SplitKVSpec): + raise AttentionContractError("split_kv must be a SplitKVSpec") + + entries: list[SplitKVRuntimePlanEntry] = [] + for batch_index, total in enumerate(totals): + base = total // cp_world_size + remainder = total % cp_world_size + owner_ranges: list[tuple[int, int]] = [] + start = 0 + for owner_cp_rank in range(cp_world_size): + end = start + base + (1 if owner_cp_rank < remainder else 0) + owner_ranges.append((start, end)) + start = end + for tp_rank in range(tp_world_size): + for cp_rank in range(cp_world_size): + for owner_cp_rank, (owner_start, owner_end) in enumerate(owner_ranges): + local_total = owner_end - owner_start + local = split_kv.resolve(local_total, backend=backend) + execution = SplitKVExecutionPlan( + requested_mode=local.requested_mode, + requested_split_size=local.requested_split_size, + actual_mode=local.actual_mode, + actual_split_size=local.actual_split_size, + boundaries=tuple( + (owner_start + start, owner_start + end) + for start, end in local.boundaries + ), + merge_order=local.merge_order, + acc_dtype=local.acc_dtype, + downcast_at=local.downcast_at, + backend=local.backend, + source=local.source, + fallback=local.fallback, + fallback_reason=local.fallback_reason, + ) + entries.append( + SplitKVRuntimePlanEntry( + coordinate=SplitKVRuntimeCoordinate( + batch_index=batch_index, + tp_rank=tp_rank, + cp_rank=cp_rank, + owner_cp_rank=owner_cp_rank, + ), + expected_kv_range=(owner_start, owner_end), + execution=execution, + ) + ) + return SplitKVRuntimePlanSet( + batch_size=len(totals), + tp_world_size=tp_world_size, + cp_world_size=cp_world_size, + total_kv_tokens=totals, + entries=tuple(entries), + ) + + +def _format_split_kv_coordinates( + coordinates: Iterable[SplitKVRuntimeCoordinate], +) -> list[dict[str, int]]: + return [coordinate.to_dict() for coordinate in sorted(coordinates)] + + +@dataclass(frozen=True) +class KVCacheSpec: + """Logical identity of the paged/block KV cache used for replay.""" + + cache_positions: tuple[int, ...] + kv_seq_lens: tuple[int, ...] + block_table: tuple[tuple[int, ...], ...] + global_token_positions: tuple[int, ...] + page_size: int + prefix_cache_enabled: bool = False + prefix_cache_key: str | None = None + shared_prefix_page_count: int = 0 + + def __post_init__(self) -> None: + cache_positions = _integer_tuple(self.cache_positions, "cache_positions") + kv_seq_lens = _integer_tuple(self.kv_seq_lens, "kv_seq_lens") + page_size = _positive_int(self.page_size, "page_size") + global_token_positions = _integer_tuple( + self.global_token_positions, "global_token_positions" + ) + if not cache_positions or any(position < 0 for position in cache_positions): + raise AttentionContractError("cache_positions must contain non-negative positions") + if not kv_seq_lens or any(length <= 0 for length in kv_seq_lens): + raise AttentionContractError("kv_seq_lens must contain positive sequence lengths") + if len(cache_positions) != len(kv_seq_lens): + raise AttentionContractError( + "cache_positions must contain one entry per kv_seq_lens entry" + ) + if not global_token_positions or any(position < 0 for position in global_token_positions): + raise AttentionContractError( + "global_token_positions must contain non-negative positions" + ) + if len(global_token_positions) != sum(kv_seq_lens): + raise AttentionContractError( + "global_token_positions must describe every logical cached token; " + f"expected {sum(kv_seq_lens)}, got {len(global_token_positions)}" + ) + token_offset = 0 + sequence_position_rows: list[tuple[int, ...]] = [] + for sequence_index, sequence_length in enumerate(kv_seq_lens): + sequence_positions = global_token_positions[ + token_offset : token_offset + sequence_length + ] + if any( + left >= right + for left, right in zip(sequence_positions, sequence_positions[1:], strict=False) + ): + raise AttentionContractError( + "global_token_positions must be strictly increasing within each sequence; " + f"sequence {sequence_index} is invalid" + ) + sequence_position_rows.append(sequence_positions) + token_offset += sequence_length + + for sequence_index, (cache_position, sequence_positions) in enumerate( + zip(cache_positions, sequence_position_rows, strict=True) + ): + terminal_position = sequence_positions[-1] + if cache_position != terminal_position: + raise AttentionContractError( + "cache_positions must equal the terminal global token position for each " + f"sequence; sequence {sequence_index} expected {terminal_position}, " + f"got {cache_position}" + ) + + try: + block_table = tuple(tuple(row) for row in self.block_table) + except TypeError as exc: + raise AttentionContractError( + "block_table must be a two-dimensional integer table" + ) from exc + if len(block_table) != len(kv_seq_lens) or any(not row for row in block_table): + raise AttentionContractError( + "block_table must contain one non-empty row per kv_seq_lens entry" + ) + active_block_rows: list[tuple[int, ...]] = [] + for row_index, (row, sequence_length) in enumerate( + zip(block_table, kv_seq_lens, strict=True) + ): + row_active_blocks: list[int] = [] + saw_padding = False + for column_index, block in enumerate(row): + if isinstance(block, bool) or not isinstance(block, int) or block < -1: + raise AttentionContractError( + "block_table entries must be integer block ids or -1 padding; " + f"got block_table[{row_index}][{column_index}]={block!r}" + ) + if block == -1: + saw_padding = True + continue + if saw_padding: + raise AttentionContractError( + "block_table -1 padding must be trailing; " + f"row {row_index} contains an active block after padding" + ) + row_active_blocks.append(block) + + expected_blocks = (sequence_length + page_size - 1) // page_size + if len(row_active_blocks) != expected_blocks: + raise AttentionContractError( + "block_table active page count must match kv_seq_lens and page_size; " + f"row {row_index} expected {expected_blocks}, got {len(row_active_blocks)}" + ) + if len(set(row_active_blocks)) != len(row_active_blocks): + raise AttentionContractError( + f"block_table row {row_index} contains duplicate active page ids" + ) + active_block_rows.append(tuple(row_active_blocks)) + + if not isinstance(self.prefix_cache_enabled, bool): + raise AttentionContractError("prefix_cache_enabled must be a bool") + shared_prefix_page_count = _non_negative_int( + self.shared_prefix_page_count, "shared_prefix_page_count" + ) + if self.prefix_cache_enabled and not self.prefix_cache_key: + raise AttentionContractError( + "prefix_cache_key is required when prefix_cache_enabled=True" + ) + if not self.prefix_cache_enabled and self.prefix_cache_key is not None: + raise AttentionContractError( + "prefix_cache_key must be None when prefix_cache_enabled=False" + ) + if not self.prefix_cache_enabled and shared_prefix_page_count != 0: + raise AttentionContractError( + "shared_prefix_page_count must be 0 when prefix_cache_enabled=False" + ) + + shared_prefix_pages: tuple[int, ...] = () + if shared_prefix_page_count > 0: + if any(len(row_blocks) < shared_prefix_page_count for row_blocks in active_block_rows): + raise AttentionContractError( + "shared_prefix_page_count exceeds an active block-table row" + ) + shared_prefix_token_count = shared_prefix_page_count * page_size + if any(length < shared_prefix_token_count for length in kv_seq_lens): + raise AttentionContractError( + "shared prefix pages must be fully populated and read-only" + ) + + shared_prefix_pages = active_block_rows[0][:shared_prefix_page_count] + shared_prefix_positions = sequence_position_rows[0][:shared_prefix_token_count] + for sequence_index, (row_blocks, positions) in enumerate( + zip(active_block_rows[1:], sequence_position_rows[1:], strict=True), start=1 + ): + if row_blocks[:shared_prefix_page_count] != shared_prefix_pages: + raise AttentionContractError( + "shared prefix page ids must match across every sequence; " + f"sequence {sequence_index} is inconsistent" + ) + if positions[:shared_prefix_token_count] != shared_prefix_positions: + raise AttentionContractError( + "shared prefix token positions must match across every sequence; " + f"sequence {sequence_index} is inconsistent" + ) + + exclusive_page_owners: dict[int, int] = {} + shared_prefix_page_ids = set(shared_prefix_pages) + for sequence_index, row_blocks in enumerate(active_block_rows): + for page_id in row_blocks[shared_prefix_page_count:]: + if page_id in shared_prefix_page_ids: + raise AttentionContractError( + "a writable suffix page cannot alias a read-only shared prefix page" + ) + previous_owner = exclusive_page_owners.get(page_id) + if previous_owner is not None: + raise AttentionContractError( + "active pages may be shared across sequences only when declared as " + "read-only prefix pages; " + f"page {page_id} is used by sequences {previous_owner} and " + f"{sequence_index}" + ) + exclusive_page_owners[page_id] = sequence_index + + object.__setattr__(self, "cache_positions", cache_positions) + object.__setattr__(self, "kv_seq_lens", kv_seq_lens) + object.__setattr__(self, "block_table", block_table) + object.__setattr__(self, "global_token_positions", global_token_positions) + + +@dataclass(frozen=True) +class RoPESpec: + """Qwen3 RoPE identity for attention inputs and KV-cache replay. + + The CP attention reference consumes attention inputs. This object records + whether those inputs are pre- or post-RoPE and pins the position/cache + metadata needed to compare fused ``RoPE+Attention`` and unfused + ``RoPE -> Attention`` materializations. + """ + + q_state: RoPEState = RoPEState.POST_ROPE + k_state: RoPEState = RoPEState.POST_ROPE + k_cache_state: RoPEState = RoPEState.POST_ROPE + theta: float = 1.0e6 + rotary_dim: int = 128 + rope_scaling: str | None = None + position_ids: tuple[int, ...] | None = None + query_position_offsets: tuple[int, ...] | None = None + key_position_offsets: tuple[int, ...] | None = None + cast_at: RoPECastPoint = RoPECastPoint.AFTER_ROPE + output_dtype: AttentionDType = AttentionDType.BF16 + fusion_boundary: RoPEFusionBoundary = RoPEFusionBoundary.UNFUSED_ROPE_ATTENTION + + def __post_init__(self) -> None: + object.__setattr__(self, "q_state", _enum_value(RoPEState, self.q_state, "q_state")) + object.__setattr__(self, "k_state", _enum_value(RoPEState, self.k_state, "k_state")) + object.__setattr__( + self, "k_cache_state", _enum_value(RoPEState, self.k_cache_state, "k_cache_state") + ) + if isinstance(self.theta, bool) or not isinstance(self.theta, (float, int)): + raise AttentionContractError(f"theta must be a positive number; got {self.theta!r}") + theta = float(self.theta) + if theta <= 0.0: + raise AttentionContractError(f"theta must be a positive number; got {self.theta!r}") + object.__setattr__(self, "theta", theta) + _positive_int(self.rotary_dim, "rotary_dim") + if self.rope_scaling is not None and ( + not isinstance(self.rope_scaling, str) or not self.rope_scaling.strip() + ): + raise AttentionContractError("rope_scaling must be a non-empty string when provided") + for position_field in ("position_ids", "query_position_offsets", "key_position_offsets"): + values = getattr(self, position_field) + if values is None: + continue + normalized = _integer_tuple(values, position_field) + if not normalized or any(value < 0 for value in normalized): + raise AttentionContractError( + f"{position_field} must contain non-negative positions" + ) + object.__setattr__(self, position_field, normalized) + object.__setattr__(self, "cast_at", _enum_value(RoPECastPoint, self.cast_at, "cast_at")) + object.__setattr__( + self, "output_dtype", _enum_value(AttentionDType, self.output_dtype, "output_dtype") + ) + object.__setattr__( + self, + "fusion_boundary", + _enum_value(RoPEFusionBoundary, self.fusion_boundary, "fusion_boundary"), + ) + + +@dataclass(frozen=True) +class AttentionContract: + """Complete semantic request consumed by contract-aware dispatch.""" + + role: AttentionRole + mode: AttentionMode + dtype: AttentionDType + batch_size: int + query_sequence_length: int + head_dim: int + causal: bool + causal_offsets: tuple[int, ...] | None + sharding: ShardingSpec + reduction: ReductionSpec + split_kv: SplitKVSpec = field(default_factory=SplitKVSpec.disabled) + kv_cache: KVCacheSpec | None = None + rope: RoPESpec | None = None + export_lse: bool = True + + def __post_init__(self) -> None: + object.__setattr__(self, "role", _enum_value(AttentionRole, self.role, "role")) + object.__setattr__(self, "mode", _enum_value(AttentionMode, self.mode, "mode")) + object.__setattr__(self, "dtype", _enum_value(AttentionDType, self.dtype, "dtype")) + batch_size = _positive_int(self.batch_size, "batch_size") + query_sequence_length = _positive_int(self.query_sequence_length, "query_sequence_length") + _positive_int(self.head_dim, "head_dim") + if not isinstance(self.sharding, ShardingSpec): + raise AttentionContractError("sharding must be a ShardingSpec") + if not isinstance(self.reduction, ReductionSpec): + raise AttentionContractError("reduction must be a ReductionSpec") + if not isinstance(self.split_kv, SplitKVSpec): + raise AttentionContractError("split_kv must be a SplitKVSpec") + if ( + self.mode is AttentionMode.PREFILL + and query_sequence_length != self.sharding.local_sequence_length + ): + raise AttentionContractError( + "prefill query_sequence_length must equal sharding.local_sequence_length; " + f"got {query_sequence_length} and {self.sharding.local_sequence_length}" + ) + if self.sharding.packed_sequence_offsets is not None: + packed_sequence_count = len(self.sharding.packed_sequence_offsets) - 1 + if packed_sequence_count != batch_size: + raise AttentionContractError( + "packed sequence count must equal logical batch_size; " + f"got {packed_sequence_count} packed sequences and batch_size={batch_size}" + ) + if not isinstance(self.causal, bool): + raise AttentionContractError("causal must be a bool") + if self.causal: + if self.causal_offsets is None: + raise AttentionContractError("causal_offsets are required for causal attention") + if self.causal_offsets is not None: + causal_offsets = _integer_tuple(self.causal_offsets, "causal_offsets") + if not causal_offsets or any(offset < 0 for offset in causal_offsets): + raise AttentionContractError("causal_offsets must contain non-negative offsets") + if self.sharding.packed_sequence_offsets is not None: + expected_causal_offsets = batch_size + offset_owner = "packed sequence" + else: + expected_causal_offsets = batch_size + offset_owner = "batch entry" + if len(causal_offsets) != expected_causal_offsets: + raise AttentionContractError( + f"causal_offsets must contain one entry per {offset_owner}" + ) + object.__setattr__(self, "causal_offsets", causal_offsets) + + if not isinstance(self.export_lse, bool) or not self.export_lse: + raise AttentionContractError( + "export_lse must be True for the WS2 attention-domain LSE contract" + ) + + if self.mode is AttentionMode.DECODE and self.kv_cache is None: + raise AttentionContractError("kv_cache metadata is required for decode attention") + if self.kv_cache is not None and not isinstance(self.kv_cache, KVCacheSpec): + raise AttentionContractError("kv_cache must be a KVCacheSpec when provided") + if self.rope is not None: + if not isinstance(self.rope, RoPESpec): + raise AttentionContractError("rope must be a RoPESpec when provided") + if self.rope.rotary_dim > self.head_dim: + raise AttentionContractError( + f"rotary_dim={self.rope.rotary_dim} must not exceed head_dim={self.head_dim}" + ) + if self.rope.position_ids is not None and len(self.rope.position_ids) not in { + query_sequence_length, + self.sharding.local_sequence_length, + }: + raise AttentionContractError( + "position_ids must describe the local query sequence or full local " + "sequence length" + ) + for position_field in ("query_position_offsets", "key_position_offsets"): + offsets = getattr(self.rope, position_field) + if offsets is not None and len(offsets) != batch_size: + raise AttentionContractError( + f"{position_field} must contain one entry per logical batch entry" + ) + if self.mode is AttentionMode.DECODE and self.kv_cache is not None: + if len(self.kv_cache.kv_seq_lens) != batch_size: + raise AttentionContractError( + "decode kv_seq_lens must contain one entry per batch entry" + ) + if len(self.kv_cache.cache_positions) != batch_size: + raise AttentionContractError( + "decode cache_positions must contain one entry per batch entry" + ) + + def to_dict(self) -> dict[str, Any]: + """Return stable, JSON-compatible requested-contract provenance.""" + + sharding = { + "tp_rank": self.sharding.tp_rank, + "tp_world_size": self.sharding.tp_world_size, + "cp_rank": self.sharding.cp_rank, + "cp_world_size": self.sharding.cp_world_size, + "global_q_heads": self.sharding.global_q_heads, + "global_kv_heads": self.sharding.global_kv_heads, + "local_q_head_start": self.sharding.local_q_head_start, + "local_q_heads": self.sharding.local_q_heads, + "local_kv_head_start": self.sharding.local_kv_head_start, + "local_kv_heads": self.sharding.local_kv_heads, + "global_sequence_length": self.sharding.global_sequence_length, + "local_sequence_length": self.sharding.local_sequence_length, + "global_block_indices": list(self.sharding.global_block_indices), + "global_block_token_starts": list(self.sharding.global_block_token_starts), + "local_block_offsets": list(self.sharding.local_block_offsets), + "packed_sequence_offsets": ( + list(self.sharding.packed_sequence_offsets) + if self.sharding.packed_sequence_offsets is not None + else None + ), + } + reduction = { + "merge": self.reduction.merge.value, + "acc_dtype": self.reduction.acc_dtype.value, + "order": self.reduction.order.value, + "downcast_at": self.reduction.downcast_at.value, + "engine": self.reduction.engine.value, + } + kv_cache = None + if self.kv_cache is not None: + kv_cache = { + "cache_positions": list(self.kv_cache.cache_positions), + "kv_seq_lens": list(self.kv_cache.kv_seq_lens), + "block_table": [list(row) for row in self.kv_cache.block_table], + "global_token_positions": list(self.kv_cache.global_token_positions), + "page_size": self.kv_cache.page_size, + "prefix_cache_enabled": self.kv_cache.prefix_cache_enabled, + "prefix_cache_key": self.kv_cache.prefix_cache_key, + "shared_prefix_page_count": self.kv_cache.shared_prefix_page_count, + } + rope = None + if self.rope is not None: + rope = { + "q_state": self.rope.q_state.value, + "k_state": self.rope.k_state.value, + "k_cache_state": self.rope.k_cache_state.value, + "theta": self.rope.theta, + "rotary_dim": self.rope.rotary_dim, + "rope_scaling": self.rope.rope_scaling, + "position_ids": ( + list(self.rope.position_ids) if self.rope.position_ids is not None else None + ), + "query_position_offsets": ( + list(self.rope.query_position_offsets) + if self.rope.query_position_offsets is not None + else None + ), + "key_position_offsets": ( + list(self.rope.key_position_offsets) + if self.rope.key_position_offsets is not None + else None + ), + "cast_at": self.rope.cast_at.value, + "output_dtype": self.rope.output_dtype.value, + "fusion_boundary": self.rope.fusion_boundary.value, + } + return { + "semantic_operator": "standard_softmax_attention", + "role": self.role.value, + "mode": self.mode.value, + "dtype": self.dtype.value, + "batch_size": self.batch_size, + "query_sequence_length": self.query_sequence_length, + "head_dim": self.head_dim, + "causal": self.causal, + "causal_offsets": ( + list(self.causal_offsets) if self.causal_offsets is not None else None + ), + "export_lse": self.export_lse, + "lse_domain": "attention", + "sharding": sharding, + "reduction": reduction, + "split_kv": self.split_kv.to_dict(), + "kv_cache": kv_cache, + "rope": rope, + } + + +@dataclass(frozen=True) +class AttentionBackendCapability: + """Capabilities a concrete backend declares to contract-aware dispatch.""" + + backend_id: str + roles: frozenset[AttentionRole] + modes: frozenset[AttentionMode] + dtypes: frozenset[AttentionDType] + cp_world_sizes: tuple[int, ...] + tp_world_sizes: tuple[int, ...] | None = None + exports_attention_lse: bool = False + deterministic_cp_merge: bool = False + supports_packed_varlen: bool = False + supports_kv_cache: bool = False + supports_rope_metadata: bool = False + supports_fused_rope_attention: bool = False + supports_split_kv_disabled: bool = True + supports_split_kv_fixed: bool = False + supports_split_kv_auto: bool = False + reports_actual_split_kv_plan: bool = False + implementation_kind: str = "production" + + def __post_init__(self) -> None: + if not isinstance(self.backend_id, str) or not self.backend_id.strip(): + raise AttentionContractError("backend_id must be a non-empty string") + roles = frozenset(_enum_value(AttentionRole, value, "roles") for value in self.roles) + modes = frozenset(_enum_value(AttentionMode, value, "modes") for value in self.modes) + dtypes = frozenset(_enum_value(AttentionDType, value, "dtypes") for value in self.dtypes) + if not roles or not modes or not dtypes: + raise AttentionContractError("backend roles, modes, and dtypes must not be empty") + cp_world_sizes = _integer_tuple(self.cp_world_sizes, "cp_world_sizes") + if not cp_world_sizes or any(size <= 0 for size in cp_world_sizes): + raise AttentionContractError("cp_world_sizes must contain positive values") + if len(set(cp_world_sizes)) != len(cp_world_sizes): + raise AttentionContractError("cp_world_sizes must not contain duplicates") + tp_world_sizes = None + if self.tp_world_sizes is not None: + tp_world_sizes = _integer_tuple(self.tp_world_sizes, "tp_world_sizes") + if not tp_world_sizes or any(size <= 0 for size in tp_world_sizes): + raise AttentionContractError("tp_world_sizes must contain positive values") + if len(set(tp_world_sizes)) != len(tp_world_sizes): + raise AttentionContractError("tp_world_sizes must not contain duplicates") + for capability_field in ( + "exports_attention_lse", + "deterministic_cp_merge", + "supports_packed_varlen", + "supports_kv_cache", + "supports_rope_metadata", + "supports_fused_rope_attention", + "supports_split_kv_disabled", + "supports_split_kv_fixed", + "supports_split_kv_auto", + "reports_actual_split_kv_plan", + ): + if not isinstance(getattr(self, capability_field), bool): + raise AttentionContractError(f"{capability_field} must be a bool") + if self.implementation_kind not in {"production", "reference", "deterministic"}: + raise AttentionContractError( + "implementation_kind must be production, reference, or deterministic" + ) + object.__setattr__(self, "roles", roles) + object.__setattr__(self, "modes", modes) + object.__setattr__(self, "dtypes", dtypes) + object.__setattr__(self, "cp_world_sizes", cp_world_sizes) + object.__setattr__(self, "tp_world_sizes", tp_world_sizes) + + def incompatibilities(self, contract: AttentionContract) -> tuple[str, ...]: + """Explain every reason this backend cannot materialize ``contract``.""" + + reasons: list[str] = [] + if contract.role not in self.roles: + reasons.append(f"role={contract.role.value} is unsupported") + if contract.mode not in self.modes: + reasons.append(f"mode={contract.mode.value} is unsupported") + if contract.dtype not in self.dtypes: + reasons.append(f"dtype={contract.dtype.value} is unsupported") + tp_size = contract.sharding.tp_world_size + cp_size = contract.sharding.cp_world_size + if self.tp_world_sizes is not None and tp_size not in self.tp_world_sizes: + reasons.append(f"TP={tp_size} is unsupported") + if cp_size not in self.cp_world_sizes: + reasons.append(f"CP={cp_size} is unsupported") + if contract.export_lse and not self.exports_attention_lse: + reasons.append("attention-domain LSE export is unsupported") + if cp_size > 1 and not self.deterministic_cp_merge: + reasons.append("deterministic CP (out, lse) merge is unsupported") + if ( + contract.sharding.packed_sequence_offsets is not None + and not self.supports_packed_varlen + ): + reasons.append("packed varlen layout is unsupported") + if contract.kv_cache is not None and not self.supports_kv_cache: + reasons.append("KV-cache identity materialization is unsupported") + if contract.rope is not None and not self.supports_rope_metadata: + reasons.append("RoPE/position metadata is unsupported") + if ( + contract.rope is not None + and contract.rope.fusion_boundary is RoPEFusionBoundary.FUSED_ROPE_ATTENTION + and not self.supports_fused_rope_attention + ): + reasons.append("fused RoPE+Attention boundary is unsupported") + split_support = { + SplitKVMode.DISABLED: self.supports_split_kv_disabled, + SplitKVMode.FIXED: self.supports_split_kv_fixed, + SplitKVMode.AUTO: self.supports_split_kv_auto, + } + if not split_support[contract.split_kv.mode]: + reasons.append(f"Split-KV policy={contract.split_kv.mode.value} is unsupported") + if contract.split_kv.strict_consistency and not self.reports_actual_split_kv_plan: + reasons.append("actual Split-KV execution-plan provenance is unsupported") + return tuple(reasons) + + def supports(self, contract: AttentionContract) -> bool: + return not self.incompatibilities(contract) + + def to_dict(self) -> dict[str, Any]: + return { + "backend_id": self.backend_id, + "roles": sorted(role.value for role in self.roles), + "modes": sorted(mode.value for mode in self.modes), + "dtypes": sorted(dtype.value for dtype in self.dtypes), + "tp_world_sizes": list(self.tp_world_sizes) if self.tp_world_sizes else None, + "cp_world_sizes": list(self.cp_world_sizes), + "exports_attention_lse": self.exports_attention_lse, + "deterministic_cp_merge": self.deterministic_cp_merge, + "supports_packed_varlen": self.supports_packed_varlen, + "supports_kv_cache": self.supports_kv_cache, + "supports_rope_metadata": self.supports_rope_metadata, + "supports_fused_rope_attention": self.supports_fused_rope_attention, + "supports_split_kv_disabled": self.supports_split_kv_disabled, + "supports_split_kv_fixed": self.supports_split_kv_fixed, + "supports_split_kv_auto": self.supports_split_kv_auto, + "reports_actual_split_kv_plan": self.reports_actual_split_kv_plan, + "implementation_kind": self.implementation_kind, + } + + +@dataclass(frozen=True) +class AttentionDispatchResult: + """A concrete backend plus the actual provenance bound to the request.""" + + op: Any + capability: AttentionBackendCapability + provenance: dict[str, Any] + + +__all__ = [ + "AttentionContract", + "AttentionContractError", + "AttentionBackendCapability", + "AttentionDispatchResult", + "AttentionDType", + "AttentionMerge", + "AttentionMode", + "AttentionRole", + "DowncastPoint", + "KVCacheSpec", + "ReductionEngine", + "ReductionOrder", + "ReductionSpec", + "RoPECastPoint", + "RoPEFusionBoundary", + "RoPESpec", + "RoPEState", + "ShardingSpec", + "build_split_kv_runtime_plan_set", + "SplitKVExecutionPlan", + "SplitKVMode", + "SplitKVRuntimeCoordinate", + "SplitKVRuntimePlanEntry", + "SplitKVRuntimePlanSet", + "SplitKVSpec", + "STRICT_ATTENTION_CORE_ID", + "STRICT_ATTENTION_SCHEDULE_ID", + "validate_split_kv_alignment", + "validate_split_kv_plan_set_alignment", +] diff --git a/rl_engine/kernels/attention_preprocess.py b/rl_engine/kernels/attention_preprocess.py new file mode 100644 index 00000000..746990de --- /dev/null +++ b/rl_engine/kernels/attention_preprocess.py @@ -0,0 +1,493 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Bitwise-bound QK-Norm and RoPE handoff for WS2 Attention. + +Transformer Engine RMSNorm is the first-choice QK-Norm implementation on CUDA +and ROCm. It is admitted only after a same-input bitwise probe against the +platform RL-Kernel path. RoPE remains on the shared RL-Kernel implementation so +training and paged-cache rollout expose the same post-RoPE boundary. +""" + +from __future__ import annotations + +import hashlib +import json +from dataclasses import dataclass +from functools import lru_cache +from types import MappingProxyType +from typing import Any, Callable, Mapping + +import torch +from torch import Tensor +from torch.autograd import Function + +QK_RMSNORM_BACKEND_ID = "rlkernel.cuda.rmsnorm" +ROPE_BACKEND_ID = "rlkernel.cuda.rope_sm90" +ROCM_QK_RMSNORM_BACKEND_ID = "rlkernel.rocm.triton_rmsnorm" +ROCM_ROPE_BACKEND_ID = "rlkernel.rocm.deterministic_rope" +TE_CUDA_QK_RMSNORM_BACKEND_ID = "transformer_engine.cuda.rmsnorm" +TE_ROCM_QK_RMSNORM_BACKEND_ID = "transformer_engine.rocm.rmsnorm" +NATIVE_QK_RMSNORM_BACKEND_ID = TE_CUDA_QK_RMSNORM_BACKEND_ID +NATIVE_ROPE_BACKEND_ID = "native.rope" +PREPROCESS_POLICY_ID = "ws2.attention.preprocess.v3" +MANDATED_ATTENTION_PREPROCESS_BACKENDS: Mapping[str, str] = MappingProxyType( + { + "qk_rmsnorm": QK_RMSNORM_BACKEND_ID, + "rope": ROPE_BACKEND_ID, + } +) +ROCM_ATTENTION_PREPROCESS_BACKENDS: Mapping[str, str] = MappingProxyType( + { + "qk_rmsnorm": ROCM_QK_RMSNORM_BACKEND_ID, + "rope": ROCM_ROPE_BACKEND_ID, + } +) +ALLOWED_ATTENTION_PREPROCESS_BACKENDS: Mapping[str, frozenset[str]] = MappingProxyType( + { + "qk_rmsnorm": frozenset( + { + QK_RMSNORM_BACKEND_ID, + ROCM_QK_RMSNORM_BACKEND_ID, + TE_CUDA_QK_RMSNORM_BACKEND_ID, + TE_ROCM_QK_RMSNORM_BACKEND_ID, + } + ), + "rope": frozenset({ROPE_BACKEND_ID, ROCM_ROPE_BACKEND_ID, NATIVE_ROPE_BACKEND_ID}), + } +) + + +class TransformerEngineRMSNormUnavailable(RuntimeError): + """Raised when the exact TE RMSNorm functional contract is unavailable.""" + + +@lru_cache(maxsize=1) +def _load_transformer_engine_rmsnorm(): + try: + import transformer_engine + import transformer_engine.pytorch # noqa: F401 - loads the platform extension + from transformer_engine.pytorch.constants import TE_DType + from transformer_engine.pytorch.cpp_extensions import rmsnorm_bwd, rmsnorm_fwd + except (ImportError, OSError, RuntimeError) as exc: + raise TransformerEngineRMSNormUnavailable( + "Transformer Engine RMSNorm forward/backward is unavailable" + ) from exc + return rmsnorm_fwd, rmsnorm_bwd, TE_DType, getattr(transformer_engine, "__version__", "unknown") + + +class _TransformerEngineRMSNormFunction(Function): + @staticmethod + def forward(ctx, x: Tensor, weight: Tensor, eps: float) -> Tensor: + rmsnorm_fwd, rmsnorm_bwd, te_dtype, _version = _load_transformer_engine_rmsnorm() + original_shape = x.shape + hidden = original_shape[-1] + x_2d = x.contiguous().view(-1, hidden) + weight_1d = weight.contiguous().view(hidden) + try: + out, _unused, rsigma = rmsnorm_fwd( + x_2d, + weight_1d, + float(eps), + None, + None, + te_dtype[x.dtype], + 0, + False, + ) + except (KeyError, TypeError, ValueError, RuntimeError) as exc: + raise TransformerEngineRMSNormUnavailable( + "installed Transformer Engine has an incompatible RMSNorm API" + ) from exc + ctx.save_for_backward(x_2d, rsigma, weight_1d) + ctx.rmsnorm_bwd = rmsnorm_bwd + ctx.original_shape = original_shape + return out.view(original_shape) + + @staticmethod + def backward(ctx, grad_output: Tensor): + x_2d, rsigma, weight = ctx.saved_tensors + dy = grad_output.contiguous().view_as(x_2d) + try: + dx, dw = ctx.rmsnorm_bwd(dy, x_2d, rsigma, weight, 0, False) + except (TypeError, ValueError, RuntimeError) as exc: + raise TransformerEngineRMSNormUnavailable( + "installed Transformer Engine has an incompatible RMSNorm backward API" + ) from exc + return dx.view(ctx.original_shape), dw.view_as(weight), None + + +class TransformerEngineRMSNormOp: + """External-weight TE RMSNorm used by both training and rollout adapters.""" + + def __init__(self) -> None: + _fwd, _bwd, _dtype, version = _load_transformer_engine_rmsnorm() + platform = "rocm" if torch.version.hip is not None else "cuda" + self.backend_id = f"transformer_engine.{platform}.rmsnorm" + self.package_version = version + + def __call__(self, x: Tensor, weight: Tensor, *, eps: float = 1.0e-6) -> Tensor: + if x.dtype not in (torch.float16, torch.bfloat16) or weight.dtype != x.dtype: + raise TypeError( + "Transformer Engine RMSNorm requires matching FP16/BF16 input and weight" + ) + if not x.is_cuda or not weight.is_cuda or x.device != weight.device: + raise ValueError("Transformer Engine RMSNorm requires input and weight on one GPU") + if weight.shape != (x.shape[-1],): + raise ValueError("RMSNorm weight must match the input hidden dimension") + return _TransformerEngineRMSNormFunction.apply(x, weight, float(eps)) + + +@dataclass(frozen=True) +class AttentionPreprocessResult: + """Post-QK-Norm, post-RoPE tensors plus executed backend evidence. + + ``probe_id`` identifies the probe configuration, not tensor contents, and + must never be used as an admission-result cache key. + """ + + q: Tensor + k: Tensor + backend_ids: Mapping[str, str] + fallback: bool + device_capability: tuple[int, int] + fallback_reason: str | None = None + probe_id: str = "" + policy_id: str = PREPROCESS_POLICY_ID + + def __post_init__(self) -> None: + object.__setattr__(self, "backend_ids", MappingProxyType(dict(self.backend_ids))) + + def evidence(self) -> dict[str, Any]: + return { + "backends": dict(self.backend_ids), + "fallback": self.fallback, + "device_capability": list(self.device_capability), + } + + def readback_fields(self) -> dict[str, Any]: + """Keyword fields consumed by ``AttentionRuntimeReadback``.""" + + return { + "preprocess_backends": dict(self.backend_ids), + "preprocess_fallback": self.fallback, + "preprocess_fallback_reason": self.fallback_reason, + "preprocess_probe_id": self.probe_id, + "preprocess_policy_id": self.policy_id, + } + + +class H100AttentionPreprocessor: + """Reuse TE QK-Norm with the RL-Kernel SM90 RoPE boundary. + + ``native_qk_norm`` and ``native_rope`` are framework-owned callables. They + are intentionally injected instead of importing TE/vLLM here, so the same + policy can be used by both runtimes. The deterministic callables default to + RL-Kernel's CUDA operators and are always run to establish the probe oracle. + """ + + def __init__( + self, + device: torch.device | str | int | None = None, + *, + native_qk_norm: Callable[..., Tensor] | None = None, + native_rope: Callable[..., Tensor] | None = None, + native_qk_norm_backend_id: str = NATIVE_QK_RMSNORM_BACKEND_ID, + native_rope_backend_id: str = NATIVE_ROPE_BACKEND_ID, + reuse_transformer_engine_qk_norm: bool = True, + require_transformer_engine_qk_norm: bool = True, + policy_id: str = PREPROCESS_POLICY_ID, + ) -> None: + if not torch.cuda.is_available(): + raise RuntimeError("H100AttentionPreprocessor requires an available CUDA runtime") + + current_device = torch.cuda.current_device() + self.device = torch.device("cuda", current_device) + if device is not None: + self.device = ( + torch.device("cuda", device) if isinstance(device, int) else torch.device(device) + ) + if self.device.type != "cuda": + raise RuntimeError(f"H100AttentionPreprocessor requires CUDA, got {self.device}") + if self.device.index is None: + self.device = torch.device("cuda", current_device) + + capability = torch.cuda.get_device_capability(self.device) + self.device_capability: tuple[int, int] = (int(capability[0]), int(capability[1])) + if self.device_capability[0] != 9: + raise RuntimeError( + "H100AttentionPreprocessor requires Hopper SM90; " + f"got sm_{self.device_capability[0]}{self.device_capability[1]}" + ) + + # Import only after the hardware gate so CPU tools can inspect the module. + from rl_engine.kernels.ops.cuda.norm.rmsnorm import RMSNormCudaOp + from rl_engine.kernels.ops.cuda.rotary_embedding.rope import RoPESM90Op + + self.rmsnorm: Callable[..., Tensor] = RMSNormCudaOp() + self.rope: Callable[..., Tensor] = RoPESM90Op() + self.deterministic_backend_ids = MANDATED_ATTENTION_PREPROCESS_BACKENDS + if not isinstance(policy_id, str) or not policy_id.strip(): + raise ValueError("policy_id must be a non-empty string") + if native_qk_norm is None and reuse_transformer_engine_qk_norm: + try: + native_qk_norm = TransformerEngineRMSNormOp() + native_qk_norm_backend_id = native_qk_norm.backend_id + except TransformerEngineRMSNormUnavailable: + if require_transformer_engine_qk_norm: + raise + self.native_qk_norm = native_qk_norm + self.native_rope = native_rope + self.require_native_qk_norm = bool( + require_transformer_engine_qk_norm and reuse_transformer_engine_qk_norm + ) + self.native_qk_norm_backend_id = native_qk_norm_backend_id + self.native_rope_backend_id = native_rope_backend_id + self.policy_id = policy_id + + def __call__( + self, + q: Tensor, + k: Tensor, + q_weight: Tensor, + k_weight: Tensor, + positions: Tensor, + *, + eps: float = 1.0e-6, + theta: float = 1_000_000.0, + ) -> AttentionPreprocessResult: + return self.forward( + q, + k, + q_weight, + k_weight, + positions, + eps=eps, + theta=theta, + ) + + def forward( + self, + q: Tensor, + k: Tensor, + q_weight: Tensor, + k_weight: Tensor, + positions: Tensor, + *, + eps: float = 1.0e-6, + theta: float = 1_000_000.0, + ) -> AttentionPreprocessResult: + _validate_inputs(q, k, q_weight, k_weight, positions, self.device) + q_norm_det = self.rmsnorm(q, q_weight, eps=eps) + k_norm_det = self.rmsnorm(k, k_weight, eps=eps) + q_det = _apply_deterministic_rope(self.rope, q_norm_det, positions, theta) + k_det = _apply_deterministic_rope(self.rope, k_norm_det, positions, theta) + probe_id = _probe_configuration_id(q, k, q_weight, k_weight, positions, eps, theta) + + native_qk_norm = self.native_qk_norm + native_rope = self.native_rope + if native_qk_norm is not None: + try: + q_norm_native = native_qk_norm(q, q_weight, eps=eps) + k_norm_native = native_qk_norm(k, k_weight, eps=eps) + norm_matches = torch.equal(q_norm_native, q_norm_det) and torch.equal( + k_norm_native, k_norm_det + ) + if norm_matches and native_rope is None: + return AttentionPreprocessResult( + q=_apply_deterministic_rope(self.rope, q_norm_native, positions, theta), + k=_apply_deterministic_rope(self.rope, k_norm_native, positions, theta), + backend_ids=MappingProxyType( + { + "qk_rmsnorm": self.native_qk_norm_backend_id, + "rope": self.deterministic_backend_ids["rope"], + } + ), + fallback=False, + device_capability=self.device_capability, + probe_id=probe_id, + policy_id=self.policy_id, + ) + if norm_matches and native_rope is not None: + q_native = native_rope(q_norm_native, positions, theta=theta) + k_native = native_rope(k_norm_native, positions, theta=theta) + if torch.equal(q_native, q_det) and torch.equal(k_native, k_det): + return AttentionPreprocessResult( + q=q_native, + k=k_native, + backend_ids=MappingProxyType( + { + "qk_rmsnorm": self.native_qk_norm_backend_id, + "rope": self.native_rope_backend_id, + } + ), + fallback=False, + device_capability=self.device_capability, + probe_id=probe_id, + policy_id=self.policy_id, + ) + fallback_reason = "native_preprocess_bitwise_probe_failed" + except Exception as exc: # framework backend failures use the common path + fallback_reason = f"native_preprocess_unavailable:{type(exc).__name__}" + else: + fallback_reason = "native_preprocess_not_supplied" + if self.require_native_qk_norm: + raise TransformerEngineRMSNormUnavailable(fallback_reason) + return AttentionPreprocessResult( + q=q_det, + k=k_det, + backend_ids=self.deterministic_backend_ids, + fallback=True, + device_capability=self.device_capability, + fallback_reason=fallback_reason, + probe_id=probe_id, + policy_id=self.policy_id, + ) + + +class RocmAttentionPreprocessor(H100AttentionPreprocessor): + """Reuse TE QK-Norm with the RL-Kernel deterministic ROCm RoPE boundary.""" + + def __init__( + self, + device: torch.device | str | int | None = None, + *, + native_qk_norm: Callable[..., Tensor] | None = None, + reuse_transformer_engine_qk_norm: bool = True, + require_transformer_engine_qk_norm: bool = True, + policy_id: str = PREPROCESS_POLICY_ID, + ) -> None: + if torch.version.hip is None or not torch.cuda.is_available(): + raise RuntimeError("RocmAttentionPreprocessor requires an available ROCm runtime") + current_device = torch.cuda.current_device() + self.device = torch.device("cuda", current_device) + if device is not None: + self.device = ( + torch.device("cuda", device) if isinstance(device, int) else torch.device(device) + ) + if self.device.type != "cuda": + raise RuntimeError(f"RocmAttentionPreprocessor requires ROCm, got {self.device}") + if self.device.index is None: + self.device = torch.device("cuda", current_device) + + from rl_engine.kernels.ops.cuda.rotary_embedding import rope as rope_module + from rl_engine.kernels.ops.triton.rmsnorm_triton import RMSNormTritonOp + + rocm_rope_type = getattr(rope_module, "RocmDeterministicRoPEOp", None) + if rocm_rope_type is None: + raise RuntimeError( + "ROCm Attention preprocessing requires RocmDeterministicRoPEOp " + "from the ROCm Attention integration" + ) + self.rmsnorm = RMSNormTritonOp() + self.rope = rocm_rope_type() + self.deterministic_backend_ids = ROCM_ATTENTION_PREPROCESS_BACKENDS + self.device_capability = (0, 0) + if native_qk_norm is None and reuse_transformer_engine_qk_norm: + try: + native_qk_norm = TransformerEngineRMSNormOp() + except TransformerEngineRMSNormUnavailable: + if require_transformer_engine_qk_norm: + raise + self.native_qk_norm = native_qk_norm + self.native_rope = None + self.require_native_qk_norm = bool( + require_transformer_engine_qk_norm and reuse_transformer_engine_qk_norm + ) + self.native_qk_norm_backend_id = ( + getattr(native_qk_norm, "backend_id", TE_ROCM_QK_RMSNORM_BACKEND_ID) + if native_qk_norm is not None + else TE_ROCM_QK_RMSNORM_BACKEND_ID + ) + self.native_rope_backend_id = ROCM_ROPE_BACKEND_ID + self.policy_id = policy_id + + +def _apply_deterministic_rope( + rope: Callable[..., Tensor], + x: Tensor, + positions: Tensor, + theta: float, +) -> Tensor: + """Adapt per-sample positions to the CUDA RoPE operator's 1-D contract.""" + + if positions.dim() == 1: + return rope(x, positions, theta=theta) + return torch.cat( + [rope(x[index : index + 1], positions[index], theta=theta) for index in range(x.shape[0])], + dim=0, + ) + + +def _probe_configuration_id( + q: Tensor, + k: Tensor, + q_weight: Tensor, + k_weight: Tensor, + positions: Tensor, + eps: float, + theta: float, +) -> str: + payload = { + "q_shape": list(q.shape), + "k_shape": list(k.shape), + "q_dtype": str(q.dtype), + "k_dtype": str(k.dtype), + "weight_dtype": str(q_weight.dtype), + "positions_shape": list(positions.shape), + "positions_sha256": hashlib.sha256(positions.detach().cpu().numpy().tobytes()).hexdigest(), + "eps": float(eps), + "theta": float(theta), + } + return hashlib.sha256(json.dumps(payload, sort_keys=True).encode("utf-8")).hexdigest()[:16] + + +def _validate_inputs( + q: Tensor, + k: Tensor, + q_weight: Tensor, + k_weight: Tensor, + positions: Tensor, + device: torch.device, +) -> None: + if q.dim() != 4 or k.dim() != 4: + raise ValueError("q and k must use [B, H, S, D] layout") + if q.shape[0] != k.shape[0] or q.shape[-2:] != k.shape[-2:]: + raise ValueError("q and k must have the same batch, sequence, and head dimensions") + if q.dtype is not torch.bfloat16 or k.dtype is not torch.bfloat16: + raise TypeError("the frozen Attention experiment requires BF16 q and k") + if q.device != device or k.device != device: + raise ValueError(f"q and k must both be on the configured device {device}") + for name, weight in (("q_weight", q_weight), ("k_weight", k_weight)): + if weight.shape != (q.shape[-1],): + raise ValueError(f"{name} must have shape ({q.shape[-1]},)") + if weight.device != device or weight.dtype is not torch.bfloat16: + raise ValueError(f"{name} must be BF16 on {device}") + if positions.device != device: + raise ValueError(f"positions must be on {device}") + if positions.dtype not in (torch.int32, torch.int64): + raise TypeError("positions must use int32 or int64 global token indices") + expected = (q.shape[-2],) if positions.dim() == 1 else (q.shape[0], q.shape[-2]) + if positions.dim() not in (1, 2) or tuple(positions.shape) != expected: + raise ValueError(f"positions must have shape [S] or [B, S], expected {expected}") + + +__all__ = [ + "ALLOWED_ATTENTION_PREPROCESS_BACKENDS", + "AttentionPreprocessResult", + "H100AttentionPreprocessor", + "MANDATED_ATTENTION_PREPROCESS_BACKENDS", + "ROCM_ATTENTION_PREPROCESS_BACKENDS", + "RocmAttentionPreprocessor", + "QK_RMSNORM_BACKEND_ID", + "ROCM_QK_RMSNORM_BACKEND_ID", + "ROCM_ROPE_BACKEND_ID", + "ROPE_BACKEND_ID", + "NATIVE_QK_RMSNORM_BACKEND_ID", + "NATIVE_ROPE_BACKEND_ID", + "PREPROCESS_POLICY_ID", + "TE_CUDA_QK_RMSNORM_BACKEND_ID", + "TE_ROCM_QK_RMSNORM_BACKEND_ID", + "TransformerEngineRMSNormOp", + "TransformerEngineRMSNormUnavailable", +] diff --git a/rl_engine/kernels/attention_projection.py b/rl_engine/kernels/attention_projection.py new file mode 100644 index 00000000..1e23ceda --- /dev/null +++ b/rl_engine/kernels/attention_projection.py @@ -0,0 +1,264 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""QKV and output-projection boundaries for the WS2 Attention experiment. + +The framework still owns the native TE/vLLM implementation. This wrapper only +freezes the semantics that must be shared by training and inference: BF16 I/O, +FP32 accumulation, ascending K reduction, and no Split-K. A native callable is +accepted only when its result is bitwise equal to the deterministic fallback. +""" + +from __future__ import annotations + +import hashlib +import json +from dataclasses import dataclass, field +from types import MappingProxyType +from typing import Any, Callable, Mapping + +import torch +from torch import Tensor + +from rl_engine.kernels.ops.cuda.matmul.det_gemm import DetGemmOp + +ProjectionCallable = Callable[[Tensor, Tensor], Tensor] + +QKV_PROJECTION = "qkv" +O_PROJ_PROJECTION = "o_proj" +PROJECTION_POLICY_ID = "ws2.attention.projection.v1" +CUDA_DETERMINISTIC_PROJECTION_BACKEND_ID = "rlkernel.cuda.det_gemm" +ROCM_DETERMINISTIC_PROJECTION_BACKEND_ID = "rlkernel.rocm.triton_det_gemm" + + +@dataclass(frozen=True) +class ProjectionCollectiveContract: + """TP/SP directions fixed by the Attention table.""" + + projection: str + tp_forward: str + tp_backward: str + sp_forward: str + sp_backward: str + reduction_forward: str + reduction_backward: str + + def __post_init__(self) -> None: + if self.projection not in {QKV_PROJECTION, O_PROJ_PROJECTION}: + raise ValueError(f"unsupported projection {self.projection!r}") + + def to_dict(self) -> dict[str, str]: + return { + "projection": self.projection, + "tp_forward": self.tp_forward, + "tp_backward": self.tp_backward, + "sp_forward": self.sp_forward, + "sp_backward": self.sp_backward, + "reduction_forward": self.reduction_forward, + "reduction_backward": self.reduction_backward, + } + + +QKV_COLLECTIVE_CONTRACT = ProjectionCollectiveContract( + projection=QKV_PROJECTION, + tp_forward="column_parallel", + tp_backward="all_reduce", + sp_forward="all_gather", + sp_backward="reduce_scatter", + reduction_forward="none", + reduction_backward="none", +) +O_PROJ_COLLECTIVE_CONTRACT = ProjectionCollectiveContract( + projection=O_PROJ_PROJECTION, + tp_forward="row_parallel", + tp_backward="none", + sp_forward="reduce_scatter", + sp_backward="all_gather", + reduction_forward="all_reduce", + reduction_backward="none", +) + + +@dataclass(frozen=True) +class ProjectionPlan: + projection: str + backend_id: str + fallback: bool + fallback_reason: str | None + probe_id: str + input_dtype: str = "torch.bfloat16" + weight_dtype: str = "torch.bfloat16" + output_dtype: str = "torch.bfloat16" + accumulation_dtype: str = "torch.float32" + reduction_order: str = "k_ascending" + split_k: bool = False + policy_id: str = PROJECTION_POLICY_ID + collective: Mapping[str, str] = field(default_factory=dict) + + def __post_init__(self) -> None: + if self.projection not in {QKV_PROJECTION, O_PROJ_PROJECTION}: + raise ValueError(f"unsupported projection {self.projection!r}") + if self.input_dtype != "torch.bfloat16" or self.weight_dtype != "torch.bfloat16": + raise ValueError("Attention projections require BF16 input and weight") + if self.output_dtype != "torch.bfloat16" or self.accumulation_dtype != "torch.float32": + raise ValueError("Attention projections require FP32 accumulation and BF16 output") + if self.reduction_order != "k_ascending" or self.split_k: + raise ValueError( + "Attention projections require ascending K reduction with Split-K disabled" + ) + object.__setattr__(self, "collective", MappingProxyType(dict(self.collective))) + + def to_dict(self) -> dict[str, Any]: + return { + "projection": self.projection, + "backend_id": self.backend_id, + "fallback": self.fallback, + "fallback_reason": self.fallback_reason, + "probe_id": self.probe_id, + "input_dtype": self.input_dtype, + "weight_dtype": self.weight_dtype, + "output_dtype": self.output_dtype, + "accumulation_dtype": self.accumulation_dtype, + "reduction_order": self.reduction_order, + "split_k": self.split_k, + "policy_id": self.policy_id, + "collective": dict(self.collective), + } + + +@dataclass(frozen=True) +class ProjectionResult: + output: Tensor + plan: ProjectionPlan + + def to_readback(self) -> dict[str, Any]: + return self.plan.to_dict() + + +class AttentionProjectionOp: + """Native-first projection wrapper with a deterministic common fallback.""" + + def __init__( + self, + projection: str, + *, + native: ProjectionCallable | None = None, + native_backend_id: str | None = None, + deterministic: ProjectionCallable | None = None, + deterministic_backend_id: str | None = None, + policy_id: str = PROJECTION_POLICY_ID, + ) -> None: + if projection not in {QKV_PROJECTION, O_PROJ_PROJECTION}: + raise ValueError(f"unsupported projection {projection!r}") + self.projection = projection + self.native = native + self.native_backend_id = native_backend_id or f"native.{projection}" + self.deterministic = deterministic or DetGemmOp() + self.deterministic_backend_id = deterministic_backend_id or ( + ROCM_DETERMINISTIC_PROJECTION_BACKEND_ID + if torch.version.hip is not None + else CUDA_DETERMINISTIC_PROJECTION_BACKEND_ID + ) + self.policy_id = policy_id + self.collective = ( + QKV_COLLECTIVE_CONTRACT if projection == QKV_PROJECTION else O_PROJ_COLLECTIVE_CONTRACT + ) + + def __call__(self, x: Tensor, weight: Tensor) -> ProjectionResult: + _validate_projection_inputs(x, weight) + deterministic_out = self.deterministic(x, weight) + if deterministic_out.dtype is not torch.bfloat16: + deterministic_out = deterministic_out.to(torch.bfloat16) + probe_id = _probe_id(x, weight) + + if self.native is not None: + try: + native_out = self.native(x, weight) + if native_out.dtype is not torch.bfloat16: + native_out = native_out.to(torch.bfloat16) + if torch.equal(native_out, deterministic_out): + return ProjectionResult( + native_out, + ProjectionPlan( + projection=self.projection, + backend_id=self.native_backend_id, + fallback=False, + fallback_reason=None, + probe_id=probe_id, + policy_id=self.policy_id, + collective=self.collective.to_dict(), + ), + ) + reason = "native_projection_bitwise_probe_failed" + except Exception as exc: # framework backend failure: use common fallback + reason = f"native_projection_unavailable:{type(exc).__name__}" + else: + reason = "native_projection_not_supplied" + + return ProjectionResult( + deterministic_out, + ProjectionPlan( + projection=self.projection, + backend_id=self.deterministic_backend_id, + fallback=True, + fallback_reason=reason, + probe_id=probe_id, + policy_id=self.policy_id, + collective=self.collective.to_dict(), + ), + ) + + +def split_qkv( + projected_qkv: Tensor, + q_heads: int, + kv_heads: int, + head_dim: int, +) -> tuple[Tensor, Tensor, Tensor]: + """Split a [Q, K, V] projection in the fixed contiguous Q/K/V order.""" + + if projected_qkv.dim() != 2: + raise ValueError("projected QKV must be [tokens, features]") + q_width = q_heads * head_dim + kv_width = kv_heads * head_dim + expected = q_width + kv_width + kv_width + if projected_qkv.shape[-1] != expected: + raise ValueError(f"projected QKV width must be {expected}, got {projected_qkv.shape[-1]}") + q, k, v = projected_qkv.split((q_width, kv_width, kv_width), dim=-1) + return q, k, v + + +def _validate_projection_inputs(x: Tensor, weight: Tensor) -> None: + if x.dim() != 2 or weight.dim() != 2: + raise ValueError("projection inputs must be [tokens, K] and [K, N]") + if x.shape[-1] != weight.shape[0]: + raise ValueError("projection K dimensions must match") + if x.dtype is not torch.bfloat16 or weight.dtype is not torch.bfloat16: + raise TypeError("Attention projections require BF16 inputs and weights") + if x.device != weight.device: + raise ValueError("projection inputs must be on the same device") + + +def _probe_id(x: Tensor, weight: Tensor) -> str: + payload = { + "x_shape": list(x.shape), + "weight_shape": list(weight.shape), + "device": str(x.device), + } + return hashlib.sha256(json.dumps(payload, sort_keys=True).encode("utf-8")).hexdigest()[:16] + + +__all__ = [ + "AttentionProjectionOp", + "CUDA_DETERMINISTIC_PROJECTION_BACKEND_ID", + "O_PROJ_COLLECTIVE_CONTRACT", + "O_PROJ_PROJECTION", + "PROJECTION_POLICY_ID", + "ProjectionCollectiveContract", + "ProjectionPlan", + "ProjectionResult", + "QKV_COLLECTIVE_CONTRACT", + "QKV_PROJECTION", + "ROCM_DETERMINISTIC_PROJECTION_BACKEND_ID", + "split_qkv", +] diff --git a/rl_engine/kernels/gtest/tolerance.py b/rl_engine/kernels/gtest/tolerance.py index f4cf7a45..4fb5bcbf 100644 --- a/rl_engine/kernels/gtest/tolerance.py +++ b/rl_engine/kernels/gtest/tolerance.py @@ -14,6 +14,7 @@ from __future__ import annotations +import hashlib import json import math from dataclasses import asdict, dataclass @@ -1000,6 +1001,27 @@ def normalize_dtype_name(dtype: str | Any) -> str: raise ContractResolveError(f"unsupported dtype: {dtype!r}") +def resolve_logprob_threshold(dtype: Any) -> float: + dtype_name = normalize_dtype_name(dtype) + try: + raw_threshold = load_contract()["accuracy"]["default"]["logprob"][dtype_name]["atol"] + except (KeyError, TypeError) as exc: + raise ValueError(f"WS1 has no logprob threshold for dtype {dtype_name!r}") from exc + if isinstance(raw_threshold, bool) or not isinstance(raw_threshold, (int, float)): + raise ValueError(f"invalid WS1 logprob threshold for dtype {dtype_name!r}") + threshold = float(raw_threshold) + if not math.isfinite(threshold) or threshold < 0.0: + raise ValueError(f"invalid WS1 logprob threshold for dtype {dtype_name!r}") + return threshold + + +def tolerance_contract_fingerprint() -> str: + canonical = json.dumps( + load_contract(), ensure_ascii=True, separators=(",", ":"), sort_keys=True + ).encode("utf-8") + return hashlib.sha256(canonical).hexdigest() + + # Private compatibility alias for callers outside this package that have not # migrated to the public normalizer yet. _dtype_name = normalize_dtype_name @@ -1030,6 +1052,8 @@ def normalize_dtype_name(dtype: str | Any) -> str: "resolve_chain_aggregate_thresholds", "resolve_comparison_roles", "resolve_dtype_policy", + "resolve_logprob_threshold", + "tolerance_contract_fingerprint", "resolve_tolerance", "resolve_tolerance_support", "validate_backend_provenance", diff --git a/rl_engine/kernels/logprob_contract.py b/rl_engine/kernels/logprob_contract.py new file mode 100644 index 00000000..b0bcd7ab --- /dev/null +++ b/rl_engine/kernels/logprob_contract.py @@ -0,0 +1,652 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Typed WS2 contract for TP-aware selected-token log-probability. + +The objects in this module describe a vocab-parallel logprob invocation: + +``selected_logp[t] = logits[t, target[t]] - logsumexp_vocab(logits[t, :])`` + +Under vocab-parallel tensor parallelism the vocabulary-wide ``logsumexp`` +requires cross-rank reduction. This module only *describes* that invocation +(shard ownership, merge semantics, mask/ignore-index metadata); it does not +shard tensors, launch collectives, or implement the ``(max, sumexp)`` merge. +Keeping description and materialization separate lets dispatch reject an +incompatible backend before any numerically different path is launched. + +Context parallelism is a declared non-merge axis: CP partitions tokens, never +the vocabulary, so the logprob reduction spans TP vocab shards only. CP rank +metadata is carried for provenance and must never widen the merge. +""" + +from __future__ import annotations + +import hashlib +import json +from dataclasses import dataclass, field +from enum import Enum +from typing import Any, TypeVar + +_EnumT = TypeVar("_EnumT", bound=Enum) + +# Policy keywords accepted by KernelRegistry.get_logprob_op; a backend id must +# never shadow one of these, or it becomes unselectable by id. +RESERVED_DISPATCH_POLICIES = frozenset({"auto", "production", "reference", "deterministic"}) +# Backend tiers; determinism is a separate axis (DeterminismScope). +IMPLEMENTATION_KINDS = frozenset({"production", "reference"}) + + +class LogprobContractError(ValueError): + """Raised when logprob metadata does not describe a valid invocation.""" + + +class LogprobRole(str, Enum): + TRAIN = "train" + INFER = "infer" + + +class LogprobDType(str, Enum): + BF16 = "bf16" + FP16 = "fp16" + FP32 = "fp32" + + +class LogprobMerge(str, Enum): + """Merge primitive for per-shard ``(local_max, local_sumexp)`` partials.""" + + MAX_SUMEXP = "max_sumexp" + + +class MergeAxis(str, Enum): + """The only reduction axis of this contract; CP is a non-merge axis.""" + + TP_VOCAB = "tp_vocab" + + +class ReductionOrder(str, Enum): + GLOBAL_VOCAB_SHARD_INDEX = "global_vocab_shard_index" + + +class ReductionTransport(str, Enum): + """Collectives move partial states only; they never reduce numerically.""" + + ALL_GATHER = "all_gather" + + +class DowncastPoint(str, Enum): + FINAL_WRITE = "final_write" + + +class ReductionEngine(str, Enum): + IN_OP_REFERENCE = "in_op_reference" + + +class DeterminismScope(str, Enum): + """Strength of the reduction's determinism guarantee. + + ``fixed_topology``: bitwise-reproducible for one fixed TP degree; results + at different TP degrees are compared against the #108 tolerance table. + + ``cross_tp_bitwise``: additionally bitwise-equal across TP degrees. This + requires the entire reduction to follow a global tile-level structure that + is independent of TP partitioning: a fixed tile decomposition of the + vocabulary plus a fixed merge order and rescaling tree over those tiles, + identical at every TP degree, so the TP degree only selects which rank + computes which tiles and never changes the floating-point grouping. + Fixed shard-order merging alone is not sufficient, because shard + boundaries would still group the combines differently across degrees. + """ + + FIXED_TOPOLOGY = "fixed_topology" + CROSS_TP_BITWISE = "cross_tp_bitwise" + + +class MaskMode(str, Enum): + """How a backend consumes inactive-token information. + + The contract permits inactive targets that do not hold ``ignore_index``, + so an ``ignore_index``-only backend cannot serve a contract with inactive + tokens. + """ + + EXPLICIT_ACTIVE_MASK = "explicit_active_mask" + IGNORE_INDEX = "ignore_index" + + +class TPPlacement(str, Enum): + REPLICATED = "replicated" + + +def _enum_value(enum_type: type[_EnumT], value: Any, field: str) -> _EnumT: + try: + return enum_type(value) + except (TypeError, ValueError) as exc: + allowed = ", ".join(item.value for item in enum_type) + raise LogprobContractError(f"{field} must be one of: {allowed}; got {value!r}") from exc + + +def _positive_int(value: Any, field: str) -> int: + if isinstance(value, bool) or not isinstance(value, int) or value <= 0: + raise LogprobContractError(f"{field} must be a positive integer; got {value!r}") + return value + + +def _non_negative_int(value: Any, field: str) -> int: + if isinstance(value, bool) or not isinstance(value, int) or value < 0: + raise LogprobContractError(f"{field} must be a non-negative integer; got {value!r}") + return value + + +def _plain_int(value: Any, field: str) -> int: + if isinstance(value, bool) or not isinstance(value, int): + raise LogprobContractError(f"{field} must be an integer; got {value!r}") + return value + + +@dataclass(frozen=True) +class ShardingSpec: + """Logical vocab-parallel TP ownership for one logprob invocation. + + ``vocab_shard_bounds`` lists every TP rank's half-open ``[start, end)`` + vocab range, indexed by rank; the full table is required on every rank and + must form a contiguous ``[0, padded_vocab_size)`` partition. + ``padded_vocab_size`` is the shard-covered (weight) vocabulary, + ``real_vocab_size`` the tokenizer vocabulary; padding columns occupy + ``[real_vocab_size, padded_vocab_size)``. + """ + + tp_rank: int + tp_world_size: int + vocab_shard_bounds: tuple[tuple[int, int], ...] + real_vocab_size: int + padded_vocab_size: int + cp_rank: int = 0 + cp_world_size: int = 1 + + def __post_init__(self) -> None: + tp_world_size = _positive_int(self.tp_world_size, "tp_world_size") + tp_rank = _non_negative_int(self.tp_rank, "tp_rank") + if tp_rank >= tp_world_size: + raise LogprobContractError( + f"tp_rank={tp_rank} must be smaller than tp_world_size={tp_world_size}" + ) + cp_world_size = _positive_int(self.cp_world_size, "cp_world_size") + cp_rank = _non_negative_int(self.cp_rank, "cp_rank") + if cp_rank >= cp_world_size: + raise LogprobContractError( + f"cp_rank={cp_rank} must be smaller than cp_world_size={cp_world_size}" + ) + + real_vocab_size = _positive_int(self.real_vocab_size, "real_vocab_size") + padded_vocab_size = _positive_int(self.padded_vocab_size, "padded_vocab_size") + if padded_vocab_size < real_vocab_size: + raise LogprobContractError( + f"padded_vocab_size={padded_vocab_size} must not be smaller than " + f"real_vocab_size={real_vocab_size}" + ) + + try: + bounds = tuple((pair[0], pair[1]) for pair in self.vocab_shard_bounds) + except (TypeError, IndexError) as exc: + raise LogprobContractError( + "vocab_shard_bounds must be an iterable of (start, end) integer pairs" + ) from exc + if len(bounds) != tp_world_size: + raise LogprobContractError( + "vocab_shard_bounds must declare exactly one (start, end) pair per TP rank; " + f"got {len(bounds)} pairs for tp_world_size={tp_world_size}" + ) + expected_start = 0 + for rank, (start, end) in enumerate(bounds): + start = _plain_int(start, f"vocab_shard_bounds[{rank}][0]") + end = _plain_int(end, f"vocab_shard_bounds[{rank}][1]") + if end <= start: + raise LogprobContractError( + f"vocab_shard_bounds[{rank}] must satisfy end > start; got [{start}, {end})" + ) + if start != expected_start: + raise LogprobContractError( + "vocab_shard_bounds must form a contiguous [0, padded_vocab_size) " + f"partition in TP-rank order; rank {rank} starts at {start}, " + f"expected {expected_start}" + ) + expected_start = end + if expected_start != padded_vocab_size: + raise LogprobContractError( + "vocab_shard_bounds must cover padded_vocab_size exactly; " + f"covered {expected_start}, declared {padded_vocab_size}" + ) + object.__setattr__(self, "vocab_shard_bounds", bounds) + + @property + def local_vocab_start(self) -> int: + return self.vocab_shard_bounds[self.tp_rank][0] + + @property + def local_vocab_end(self) -> int: + return self.vocab_shard_bounds[self.tp_rank][1] + + @property + def local_vocab_size(self) -> int: + start, end = self.vocab_shard_bounds[self.tp_rank] + return end - start + + def owner_rank(self, token_id: int) -> int: + """Return the unique TP rank owning ``token_id``; error outside real vocab.""" + + token_id = _plain_int(token_id, "token_id") + if token_id < 0 or token_id >= self.real_vocab_size: + raise LogprobContractError( + f"token_id={token_id} is outside the real vocabulary " + f"[0, {self.real_vocab_size}); mask it as inactive instead" + ) + for rank, (start, end) in enumerate(self.vocab_shard_bounds): + if start <= token_id < end: + return rank + raise LogprobContractError( + f"token_id={token_id} is not covered by any declared vocab shard" + ) + + +@dataclass(frozen=True) +class MaskSpec: + """Active-token mask and ignore index for one logprob invocation. + + Inactive tokens are excluded from drift aggregates and from the + single-owner target gather; their targets may legally hold ``ignore_index``. + """ + + num_tokens: int + active_mask: tuple[bool, ...] + ignore_index: int = -100 + _active_token_count: int = field(init=False, repr=False, compare=False) + _active_mask_sha256: str = field(init=False, repr=False, compare=False) + + def __post_init__(self) -> None: + num_tokens = _positive_int(self.num_tokens, "num_tokens") + _plain_int(self.ignore_index, "ignore_index") + try: + active_mask = tuple(self.active_mask) + except TypeError as exc: + raise LogprobContractError("active_mask must be an iterable of booleans") from exc + for index, value in enumerate(active_mask): + if not isinstance(value, bool): + raise LogprobContractError(f"active_mask[{index}] must be a bool; got {value!r}") + if len(active_mask) != num_tokens: + raise LogprobContractError( + "active_mask must contain exactly one entry per token; " + f"got {len(active_mask)} entries for num_tokens={num_tokens}" + ) + object.__setattr__(self, "active_mask", active_mask) + object.__setattr__(self, "_active_token_count", sum(active_mask)) + object.__setattr__( + self, "_active_mask_sha256", hashlib.sha256(bytes(active_mask)).hexdigest() + ) + + @property + def active_token_count(self) -> int: + return self._active_token_count + + @property + def active_mask_sha256(self) -> str: + """Compact mask identity for provenance and cross-rank agreement.""" + return self._active_mask_sha256 + + +@dataclass(frozen=True) +class ReductionSpec: + """Deterministic TP-vocab ``(max, sumexp)`` merge semantics. + + Every rank first masks local columns whose global id lies in + ``[real_vocab_size, padded_vocab_size)`` to ``-inf`` (padding never + contributes to the logsumexp), then computes ``m_l = max(local_logits)`` + and ``s_l = sum(exp(local_logits - m_l))`` in fp32. Partials travel by + all-gather -- collectives are transport only, never a numerical + reduction -- and every rank merges in fixed global vocab-shard index + order:: + + M = max_l(m_l) + S = sum_l(s_l * exp(m_l - M)) + LSE = M + log(S) + selected_logp = target_logit - LSE + + The selected target logit comes from a masked single-owner gather; + downcast happens only at the final write. The identity partial for a + padding-only shard, or a row whose local columns are all ``-inf`` after + masking, is ``(m_l, s_l) = (-inf, 0)``: a partial with ``s_l = 0`` + contributes nothing to the merge regardless of its ``m_l``, and + implementations must use this identity directly rather than evaluate + ``exp(-inf - (-inf))``, which would poison the merge with NaN. Averaging + per-rank logsumexp values, or letting a collective reduce numerically, is + never conformant at either determinism scope. + """ + + merge: LogprobMerge = LogprobMerge.MAX_SUMEXP + merge_axis: MergeAxis = MergeAxis.TP_VOCAB + acc_dtype: LogprobDType = LogprobDType.FP32 + order: ReductionOrder = ReductionOrder.GLOBAL_VOCAB_SHARD_INDEX + transport: ReductionTransport = ReductionTransport.ALL_GATHER + downcast_at: DowncastPoint = DowncastPoint.FINAL_WRITE + engine: ReductionEngine = ReductionEngine.IN_OP_REFERENCE + determinism_scope: DeterminismScope = DeterminismScope.CROSS_TP_BITWISE + + def __post_init__(self) -> None: + object.__setattr__( + self, + "determinism_scope", + _enum_value(DeterminismScope, self.determinism_scope, "determinism_scope"), + ) + object.__setattr__(self, "merge", _enum_value(LogprobMerge, self.merge, "merge")) + object.__setattr__( + self, "merge_axis", _enum_value(MergeAxis, self.merge_axis, "merge_axis") + ) + object.__setattr__( + self, "acc_dtype", _enum_value(LogprobDType, self.acc_dtype, "acc_dtype") + ) + object.__setattr__(self, "order", _enum_value(ReductionOrder, self.order, "order")) + object.__setattr__( + self, "transport", _enum_value(ReductionTransport, self.transport, "transport") + ) + object.__setattr__( + self, "downcast_at", _enum_value(DowncastPoint, self.downcast_at, "downcast_at") + ) + object.__setattr__(self, "engine", _enum_value(ReductionEngine, self.engine, "engine")) + if self.acc_dtype is not LogprobDType.FP32: + raise LogprobContractError( + f"TP logprob accumulation must be fp32; got {self.acc_dtype.value}" + ) + + +@dataclass(frozen=True) +class LogprobOutputSpec: + """Output surface every conforming backend must produce: fp32 selected + logprob and fp32 vocab-domain LSE, replicated across the TP group.""" + + selected_logp_dtype: LogprobDType = LogprobDType.FP32 + lse_dtype: LogprobDType = LogprobDType.FP32 + tp_placement: TPPlacement = TPPlacement.REPLICATED + + def __post_init__(self) -> None: + object.__setattr__( + self, + "selected_logp_dtype", + _enum_value(LogprobDType, self.selected_logp_dtype, "selected_logp_dtype"), + ) + object.__setattr__( + self, "lse_dtype", _enum_value(LogprobDType, self.lse_dtype, "lse_dtype") + ) + object.__setattr__( + self, "tp_placement", _enum_value(TPPlacement, self.tp_placement, "tp_placement") + ) + if self.selected_logp_dtype is not LogprobDType.FP32: + raise LogprobContractError( + f"selected logprob output must be fp32; got {self.selected_logp_dtype.value}" + ) + if self.lse_dtype is not LogprobDType.FP32: + raise LogprobContractError(f"vocab LSE output must be fp32; got {self.lse_dtype.value}") + + +@dataclass(frozen=True) +class LogprobContract: + """Complete semantic request consumed by contract-aware dispatch.""" + + role: LogprobRole + dtype: LogprobDType + mask: MaskSpec + sharding: ShardingSpec + reduction: ReductionSpec + output: LogprobOutputSpec = field(default_factory=LogprobOutputSpec) + export_lse: bool = True + + def __post_init__(self) -> None: + object.__setattr__(self, "role", _enum_value(LogprobRole, self.role, "role")) + object.__setattr__(self, "dtype", _enum_value(LogprobDType, self.dtype, "dtype")) + if not isinstance(self.mask, MaskSpec): + raise LogprobContractError("mask must be a MaskSpec") + if not isinstance(self.sharding, ShardingSpec): + raise LogprobContractError("sharding must be a ShardingSpec") + if not isinstance(self.reduction, ReductionSpec): + raise LogprobContractError("reduction must be a ReductionSpec") + if not isinstance(self.output, LogprobOutputSpec): + raise LogprobContractError("output must be a LogprobOutputSpec") + if not isinstance(self.export_lse, bool) or not self.export_lse: + raise LogprobContractError( + "export_lse must be True for the WS2 vocab-domain LSE drift contract" + ) + if 0 <= self.mask.ignore_index < self.sharding.real_vocab_size: + raise LogprobContractError( + f"ignore_index={self.mask.ignore_index} must not collide with the real " + f"vocabulary [0, {self.sharding.real_vocab_size})" + ) + + def to_dict(self) -> dict[str, Any]: + """Return stable, JSON-compatible requested-contract provenance.""" + + sharding = { + "tp_rank": self.sharding.tp_rank, + "tp_world_size": self.sharding.tp_world_size, + "cp_rank": self.sharding.cp_rank, + "cp_world_size": self.sharding.cp_world_size, + "vocab_shard_bounds": [list(pair) for pair in self.sharding.vocab_shard_bounds], + "real_vocab_size": self.sharding.real_vocab_size, + "padded_vocab_size": self.sharding.padded_vocab_size, + "local_vocab_start": self.sharding.local_vocab_start, + "local_vocab_end": self.sharding.local_vocab_end, + } + reduction = { + "merge": self.reduction.merge.value, + "merge_axis": self.reduction.merge_axis.value, + "acc_dtype": self.reduction.acc_dtype.value, + "order": self.reduction.order.value, + "transport": self.reduction.transport.value, + "downcast_at": self.reduction.downcast_at.value, + "engine": self.reduction.engine.value, + "determinism_scope": self.reduction.determinism_scope.value, + "cp_is_merge_axis": False, + } + # The digest stands in for the raw per-token mask, which would + # dominate the provenance size. + mask = { + "num_tokens": self.mask.num_tokens, + "active_token_count": self.mask.active_token_count, + "active_mask_sha256": self.mask.active_mask_sha256, + "ignore_index": self.mask.ignore_index, + } + output = { + "selected_logp_dtype": self.output.selected_logp_dtype.value, + "lse_dtype": self.output.lse_dtype.value, + "tp_placement": self.output.tp_placement.value, + } + return { + "semantic_operator": "selected_token_logprob", + "role": self.role.value, + "dtype": self.dtype.value, + "export_lse": self.export_lse, + "lse_domain": "vocab", + "mask": mask, + "sharding": sharding, + "reduction": reduction, + "output": output, + } + + def cross_rank_fingerprint(self) -> str: + """Rank-independent identity for preflight agreement across ranks. + + Excludes ``tp_rank``/``cp_rank`` (and their derived local bounds) so + every rank of one logical invocation computes the same value. + All-gathering this fingerprint together with the resolved backend id + and aborting on mismatch is the documented preflight for distributed + dispatch; ``requested_backend="auto"`` is not distributed-safe + without it. + """ + + payload = self.to_dict() + payload["sharding"] = { + key: value + for key, value in payload["sharding"].items() + if key not in {"tp_rank", "cp_rank", "local_vocab_start", "local_vocab_end"} + } + encoded = json.dumps(payload, sort_keys=True, separators=(",", ":")) + return hashlib.sha256(encoded.encode("utf-8")).hexdigest() + + +@dataclass(frozen=True) +class LogprobBackendCapability: + """Capabilities a concrete backend declares to contract-aware dispatch.""" + + backend_id: str + roles: frozenset[LogprobRole] + dtypes: frozenset[LogprobDType] + tp_world_sizes: tuple[int, ...] | None = None + cp_world_sizes: tuple[int, ...] | None = None + supports_vocab_padding: bool = False + mask_modes: frozenset[MaskMode] = frozenset() + exports_vocab_lse: bool = False + determinism_scopes: frozenset[DeterminismScope] = frozenset() + implementation_kind: str = "production" + + def __post_init__(self) -> None: + if not isinstance(self.backend_id, str) or not self.backend_id.strip(): + raise LogprobContractError("backend_id must be a non-empty string") + if self.backend_id.strip().lower() in RESERVED_DISPATCH_POLICIES: + raise LogprobContractError( + f"backend_id={self.backend_id!r} shadows a reserved dispatch policy keyword" + ) + object.__setattr__(self, "backend_id", self.backend_id.strip()) + try: + roles = frozenset(_enum_value(LogprobRole, value, "roles") for value in self.roles) + dtypes = frozenset(_enum_value(LogprobDType, value, "dtypes") for value in self.dtypes) + except TypeError as exc: + raise LogprobContractError("roles and dtypes must be iterables of enum values") from exc + if not roles or not dtypes: + raise LogprobContractError("backend roles and dtypes must not be empty") + tp_world_sizes = self._validated_world_sizes(self.tp_world_sizes, "tp_world_sizes") + cp_world_sizes = self._validated_world_sizes(self.cp_world_sizes, "cp_world_sizes") + try: + mask_modes = frozenset( + _enum_value(MaskMode, value, "mask_modes") for value in self.mask_modes + ) + determinism_scopes = frozenset( + _enum_value(DeterminismScope, value, "determinism_scopes") + for value in self.determinism_scopes + ) + except TypeError as exc: + raise LogprobContractError( + "mask_modes and determinism_scopes must be iterables of enum values" + ) from exc + for flag_name in ("supports_vocab_padding", "exports_vocab_lse"): + if not isinstance(getattr(self, flag_name), bool): + raise LogprobContractError(f"{flag_name} must be a bool") + if self.implementation_kind not in IMPLEMENTATION_KINDS: + raise LogprobContractError( + f"implementation_kind must be one of: {', '.join(sorted(IMPLEMENTATION_KINDS))}" + ) + object.__setattr__(self, "roles", roles) + object.__setattr__(self, "dtypes", dtypes) + object.__setattr__(self, "tp_world_sizes", tp_world_sizes) + object.__setattr__(self, "cp_world_sizes", cp_world_sizes) + object.__setattr__(self, "mask_modes", mask_modes) + object.__setattr__(self, "determinism_scopes", determinism_scopes) + + @staticmethod + def _validated_world_sizes( + values: tuple[int, ...] | None, field: str + ) -> tuple[int, ...] | None: + if values is None: + return None + try: + sizes = tuple(values) + except TypeError as exc: + raise LogprobContractError(f"{field} must be an iterable of integers") from exc + if not sizes: + raise LogprobContractError(f"{field} must not be empty; use None for unrestricted") + for value in sizes: + if isinstance(value, bool) or not isinstance(value, int) or value <= 0: + raise LogprobContractError(f"{field} must contain positive values; got {value!r}") + if len(set(sizes)) != len(sizes): + raise LogprobContractError(f"{field} must not contain duplicates") + return sizes + + def incompatibilities(self, contract: LogprobContract) -> tuple[str, ...]: + """Explain every reason this backend cannot materialize ``contract``.""" + + reasons: list[str] = [] + if contract.role not in self.roles: + reasons.append(f"role={contract.role.value} is unsupported") + if contract.dtype not in self.dtypes: + reasons.append(f"dtype={contract.dtype.value} is unsupported") + tp_size = contract.sharding.tp_world_size + cp_size = contract.sharding.cp_world_size + if self.tp_world_sizes is not None and tp_size not in self.tp_world_sizes: + reasons.append(f"TP={tp_size} is unsupported") + if self.cp_world_sizes is not None and cp_size not in self.cp_world_sizes: + reasons.append(f"CP={cp_size} is unsupported") + if ( + contract.sharding.padded_vocab_size != contract.sharding.real_vocab_size + and not self.supports_vocab_padding + ): + reasons.append("padded-vs-real vocab masking is unsupported") + if ( + contract.mask.active_token_count != contract.mask.num_tokens + and MaskMode.EXPLICIT_ACTIVE_MASK not in self.mask_modes + ): + # Inactive targets need not hold ignore_index (see MaskMode). + reasons.append("explicit active-token masking is unsupported") + if contract.export_lse and not self.exports_vocab_lse: + reasons.append("vocab-domain LSE export is unsupported") + if contract.reduction.determinism_scope not in self.determinism_scopes: + reasons.append( + f"determinism_scope={contract.reduction.determinism_scope.value} is unsupported" + ) + return tuple(reasons) + + def supports(self, contract: LogprobContract) -> bool: + return not self.incompatibilities(contract) + + def to_dict(self) -> dict[str, Any]: + return { + "backend_id": self.backend_id, + "roles": sorted(role.value for role in self.roles), + "dtypes": sorted(dtype.value for dtype in self.dtypes), + "tp_world_sizes": list(self.tp_world_sizes) if self.tp_world_sizes else None, + "cp_world_sizes": list(self.cp_world_sizes) if self.cp_world_sizes else None, + "supports_vocab_padding": self.supports_vocab_padding, + "mask_modes": sorted(mode.value for mode in self.mask_modes), + "exports_vocab_lse": self.exports_vocab_lse, + "determinism_scopes": sorted(scope.value for scope in self.determinism_scopes), + "implementation_kind": self.implementation_kind, + } + + +@dataclass(frozen=True) +class LogprobDispatchResult: + """A concrete backend plus the actual provenance bound to the request.""" + + op: Any + capability: LogprobBackendCapability + provenance: dict[str, Any] + + +__all__ = [ + "IMPLEMENTATION_KINDS", + "RESERVED_DISPATCH_POLICIES", + "DeterminismScope", + "DowncastPoint", + "LogprobBackendCapability", + "LogprobContract", + "LogprobContractError", + "LogprobDType", + "LogprobDispatchResult", + "LogprobMerge", + "LogprobOutputSpec", + "LogprobRole", + "MaskMode", + "MaskSpec", + "MergeAxis", + "ReductionEngine", + "ReductionOrder", + "ReductionSpec", + "ReductionTransport", + "ShardingSpec", + "TPPlacement", +] diff --git a/rl_engine/kernels/ops/pytorch/attention/__init__.py b/rl_engine/kernels/ops/pytorch/attention/__init__.py index d2454e67..977ab14c 100644 --- a/rl_engine/kernels/ops/pytorch/attention/__init__.py +++ b/rl_engine/kernels/ops/pytorch/attention/__init__.py @@ -4,6 +4,12 @@ import torch import torch.nn.functional as F +from rl_engine.kernels.ops.pytorch.attention.ablation import ( + AttentionAblationConfig, + AttentionAblationOp, + AttentionAblationResult, +) + class NativeAttentionOp: """PyTorch SDPA fallback for FlashAttention-layout tensors.""" @@ -46,4 +52,9 @@ def __call__( return out.transpose(1, 2) -__all__ = ["NativeAttentionOp"] +__all__ = [ + "AttentionAblationConfig", + "AttentionAblationOp", + "AttentionAblationResult", + "NativeAttentionOp", +] diff --git a/rl_engine/kernels/ops/pytorch/attention/ablation.py b/rl_engine/kernels/ops/pytorch/attention/ablation.py new file mode 100644 index 00000000..bea6fb52 --- /dev/null +++ b/rl_engine/kernels/ops/pytorch/attention/ablation.py @@ -0,0 +1,667 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Unified Attention entry point for the PR230 cross-configuration matrix. + +The matrix needs one stable callable shape even though training and rollout may +materialize different Attention backends. This adapter owns the common +contract checks and provenance only; numerical work remains in the qualified +CUDA/ROCm production cores or the explicit RL-Kernel reference core. +""" + +from __future__ import annotations + +import hashlib +import inspect +import json +import math +from dataclasses import dataclass, field +from types import MappingProxyType +from typing import Any, Callable, Mapping, cast + +import torch +from torch import Tensor + +from rl_engine.kernels.attention_contract import ( + STRICT_ATTENTION_CORE_ID, + STRICT_ATTENTION_SCHEDULE_ID, + AttentionContract, + AttentionContractError, + AttentionDType, + SplitKVMode, +) + +BACKEND_ID = "rlkernel.attention.deterministic.v1" +REFERENCE_BACKEND_ID = "rlkernel.attention.reference.v1" +_STRICT_AG_RS_BACKENDS = frozenset({"self_owned_cuda_ag_rs", "cuda_ag_rs", "rccl_ag_rs"}) + +_TORCH_DTYPES = { + AttentionDType.BF16: torch.bfloat16, + AttentionDType.FP16: torch.float16, + AttentionDType.FP32: torch.float32, +} + + +@dataclass(frozen=True) +class AttentionAblationConfig: + """Per-invocation settings materialized by the ablation runner.""" + + backend: str = "auto" + deterministic: bool = True + communication_backend: str = "none" + return_lse: bool = True + return_gradients: bool = False + strict_core_id: str = STRICT_ATTENTION_CORE_ID + strict_schedule: str = STRICT_ATTENTION_SCHEDULE_ID + validate: bool = True + + def __post_init__(self) -> None: + if not isinstance(self.backend, str) or not self.backend.strip(): + raise AttentionContractError("Attention backend must be a non-empty string") + if not isinstance(self.deterministic, bool): + raise AttentionContractError("deterministic must be a bool") + if ( + not isinstance(self.communication_backend, str) + or not self.communication_backend.strip() + ): + raise AttentionContractError("communication_backend must be a non-empty string") + for name in ("return_lse", "return_gradients", "validate"): + if not isinstance(getattr(self, name), bool): + raise AttentionContractError(f"{name} must be a bool") + if not isinstance(self.strict_core_id, str) or not self.strict_core_id.strip(): + raise AttentionContractError("strict_core_id must be a non-empty string") + if not isinstance(self.strict_schedule, str) or not self.strict_schedule.strip(): + raise AttentionContractError("strict_schedule must be a non-empty string") + object.__setattr__(self, "backend", self.backend.strip().lower()) + object.__setattr__(self, "communication_backend", self.communication_backend.strip()) + + +@dataclass(frozen=True) +class AttentionAblationResult: + """Standardized Attention result consumed by cross-config artifacts.""" + + out: Tensor + lse: Tensor | None + dq: Tensor | None = None + dk: Tensor | None = None + dv: Tensor | None = None + backend_id: str = BACKEND_ID + deterministic: bool = True + provenance: Mapping[str, Any] = field(default_factory=dict) + + def __post_init__(self) -> None: + if not isinstance(self.out, Tensor): + raise TypeError("Attention result out must be a torch.Tensor") + if self.lse is not None and not isinstance(self.lse, Tensor): + raise TypeError("Attention result lse must be a torch.Tensor or None") + for name in ("dq", "dk", "dv"): + value = getattr(self, name) + if value is not None and not isinstance(value, Tensor): + raise TypeError(f"Attention result {name} must be a torch.Tensor or None") + if not isinstance(self.backend_id, str) or not self.backend_id.strip(): + raise ValueError("Attention result backend_id must be non-empty") + if not isinstance(self.deterministic, bool): + raise TypeError("Attention result deterministic must be a bool") + if not isinstance(self.provenance, Mapping): + raise TypeError("Attention result provenance must be a mapping") + object.__setattr__(self, "provenance", MappingProxyType(dict(self.provenance))) + + @property + def out_lse(self) -> tuple[Tensor, Tensor | None]: + """Compatibility tuple for callers that consume ``(out, lse)``.""" + + return self.out, self.lse + + def readback(self) -> dict[str, Any]: + """Return JSON-compatible execution evidence for PR230 artifacts.""" + + return { + "backend_id": self.backend_id, + "deterministic": self.deterministic, + "out_shape": list(self.out.shape), + "out_dtype": str(self.out.dtype).replace("torch.", ""), + "lse_shape": None if self.lse is None else list(self.lse.shape), + "lse_dtype": None if self.lse is None else str(self.lse.dtype).replace("torch.", ""), + "gradients": { + "dq": self.dq is not None, + "dk": self.dk is not None, + "dv": self.dv is not None, + }, + "provenance": dict(self.provenance), + } + + +class AttentionAblationOp: + """PR230/PR314-style unified Attention wrapper. + + ``core`` and ``reference`` are injectable so the wrapper is usable by the + semantic operator session without importing CUDA at construction time. + ``core`` should expose ``forward_with_lse``; ``reference`` is the existing + pure-PyTorch CP reference with the same method. + """ + + op_class = "attention" + is_batch_invariant = True + backend_id = BACKEND_ID + + def __init__( + self, + *, + core: Any | None = None, + reference: Any | None = None, + native: Any | None = None, + cp_backend: Any | None = None, + communication_backend: str = "none", + ) -> None: + if not isinstance(communication_backend, str) or not communication_backend.strip(): + raise AttentionContractError("communication_backend must be a non-empty string") + self.core = core + self.reference = reference + self.native = native + # CP production execution is injected by the runtime adapter. Keeping + # it separate from the single-device core prevents an accidental + # fallback to the PyTorch reference when AG/RS is required. + self.cp_backend = cp_backend + self.communication_backend = communication_backend.strip() + + def __call__( + self, + q: Tensor, + k: Tensor, + v: Tensor, + *, + contract: AttentionContract, + config: AttentionAblationConfig | Mapping[str, Any] | None = None, + backend: str | Callable[..., Any] | None = None, + deterministic: bool | None = None, + return_lse: bool | None = None, + return_gradients: bool | None = None, + dout: Tensor | None = None, + communication_backend: str | None = None, + validate: bool | None = None, + **kwargs: Any, + ) -> AttentionAblationResult: + return self.apply( + q, + k, + v, + contract=contract, + config=config, + backend=backend, + deterministic=deterministic, + return_lse=return_lse, + return_gradients=return_gradients, + dout=dout, + communication_backend=communication_backend, + validate=validate, + **kwargs, + ) + + def apply( + self, + q: Tensor, + k: Tensor, + v: Tensor, + *, + contract: AttentionContract, + config: AttentionAblationConfig | Mapping[str, Any] | None = None, + backend: str | Callable[..., Any] | None = None, + deterministic: bool | None = None, + return_lse: bool | None = None, + return_gradients: bool | None = None, + dout: Tensor | None = None, + communication_backend: str | None = None, + validate: bool | None = None, + **kwargs: Any, + ) -> AttentionAblationResult: + if not isinstance(contract, AttentionContract): + raise AttentionContractError("contract must be an AttentionContract") + backend_request = ( + backend + if callable(backend) + or hasattr(backend, "forward_with_lse") + or hasattr(backend, "apply") + else None + ) + cfg = _resolve_config( + config, + backend=backend if isinstance(backend, str) else None, + deterministic=deterministic, + return_lse=return_lse, + return_gradients=return_gradients, + communication_backend=( + communication_backend + if communication_backend is not None + else self.communication_backend + ), + validate=validate, + ) + if cfg.validate: + self._validate_inputs(q, k, v, contract) + if cfg.deterministic and contract.split_kv.mode is not SplitKVMode.DISABLED: + raise AttentionContractError( + "strict deterministic Attention requires Split-KV to be disabled" + ) + if ( + cfg.deterministic + and contract.sharding.cp_world_size > 1 + and cfg.communication_backend not in _STRICT_AG_RS_BACKENDS + ): + raise AttentionContractError( + "strict CP Attention requires an explicit CUDA AG/RS or ROCm RCCL AG/RS backend" + ) + if cfg.deterministic and not cfg.return_lse: + raise AttentionContractError("strict deterministic Attention must return LSE") + if cfg.return_gradients and dout is None: + raise AttentionContractError("dout is required when return_gradients=True") + if dout is not None and dout.shape != q.shape: + raise AttentionContractError("dout must have the same shape as q") + + requested = backend_request if backend_request is not None else cfg.backend + selected, selected_id = self._select_backend( + requested, + q, + contract, + communication_backend=cfg.communication_backend, + ) + if cfg.deterministic and selected_id == "native": + raise AttentionContractError( + "deterministic=True cannot execute an unverified native Attention backend" + ) + selected_core_id = getattr(selected, "core_id", None) + selected_schedule = getattr(selected, "strict_schedule", None) + if cfg.deterministic and selected_id not in {BACKEND_ID, REFERENCE_BACKEND_ID}: + if selected_core_id != cfg.strict_core_id or selected_schedule != cfg.strict_schedule: + raise AttentionContractError( + "deterministic Attention requires the shared strict core and schedule" + ) + + call_kwargs = dict(kwargs) + call_kwargs.setdefault("causal", contract.causal) + call_kwargs.setdefault("scale", 1.0 / math.sqrt(contract.head_dim)) + call_kwargs.setdefault("cp_world_size", contract.sharding.cp_world_size) + if contract.split_kv.mode is SplitKVMode.FIXED: + call_kwargs.setdefault("kv_chunk_size", contract.split_kv.fixed_split_size) + out, lse, backend_provenance = self._invoke(selected, q, k, v, call_kwargs, contract) + if cfg.validate: + self._validate_outputs(out, lse, q, contract) + _validate_runtime_provenance( + selected, + selected_id, + backend_provenance, + cfg, + ) + + dq = dk = dv = None + if cfg.return_gradients: + dq, dk, dv = self._backward(selected, q, k, v, out, dout, call_kwargs) + + provenance = { + "schema_version": "rlkernel.attention.ablation_result.v1", + "semantic_operator": "attention", + "backend_id": selected_id, + "deterministic": cfg.deterministic, + "strict_core_id": (cfg.strict_core_id if cfg.deterministic else None), + "strict_schedule": cfg.strict_schedule if cfg.deterministic else None, + "core_id": cfg.strict_core_id if cfg.deterministic else selected_id, + "backend_deterministic": cfg.deterministic, + "native_attention_arithmetic": False if cfg.deterministic else selected_id == "native", + "communication_backend": cfg.communication_backend, + "communication_executed": bool(getattr(selected, "communication_executed", False)), + "split_kv": contract.split_kv.to_dict(), + "actual_split_kv": _actual_split_provenance( + contract, + total_kv_tokens=k.size(2), + backend=selected_id, + ), + "reduction": _reduction_provenance(contract), + "contract_fingerprint": _contract_fingerprint(contract), + "return_lse": cfg.return_lse, + "return_gradients": cfg.return_gradients, + } + provenance.update(backend_provenance) + provenance.setdefault("actual_backend", selected_id) + provenance.setdefault("communication_backend", cfg.communication_backend) + provenance.setdefault("production_ready", False) + return AttentionAblationResult( + out=out, + lse=lse if cfg.return_lse else None, + dq=dq, + dk=dk, + dv=dv, + backend_id=selected_id, + deterministic=cfg.deterministic, + provenance=provenance, + ) + + def apply_fp32(self, *args: Any, **kwargs: Any) -> AttentionAblationResult: + """Stable fingerprint entry point used by ``OperatorSession``.""" + + return self.apply(*args, **kwargs) + + def _select_backend( + self, + requested: str | Callable[..., Any], + q: Tensor, + contract: AttentionContract, + *, + communication_backend: str, + ) -> tuple[Any, str]: + if ( + callable(requested) + or hasattr(requested, "forward_with_lse") + or hasattr(requested, "apply") + ): + return requested, _callable_backend_id(requested) + normalized = str(requested).strip().lower() + if normalized in {"native", "te", "flashinfer"}: + if self.native is None: + raise AttentionContractError( + "native Attention backend was requested but no native callable was injected" + ) + return self.native, "native" + if normalized in {"reference", "pytorch_reference"}: + return self._reference_backend(), REFERENCE_BACKEND_ID + if normalized not in {"auto", "deterministic", "rlkernel"}: + raise AttentionContractError(f"unsupported Attention backend {requested!r}") + if contract.sharding.cp_world_size > 1: + if communication_backend in _STRICT_AG_RS_BACKENDS: + if self.cp_backend is None: + raise AttentionContractError( + "CP production Attention requires an injected AG/RS backend" + ) + return self.cp_backend, _callable_backend_id(self.cp_backend) + return self._reference_backend(), REFERENCE_BACKEND_ID + if q.device.type != "cuda" or torch.version.hip is not None: + return self._reference_backend(), REFERENCE_BACKEND_ID + if self.core is None: + from rl_engine.kernels.ops.cuda.attention.deterministic_attn import ( + DeterministicAttentionOp, + ) + + self.core = DeterministicAttentionOp() + return self.core, BACKEND_ID + + def _reference_backend(self) -> Any: + if self.reference is None: + from rl_engine.kernels.ops.pytorch.attention.cp_attention import ( + DeterministicCPAttentionReferenceOp, + ) + + self.reference = DeterministicCPAttentionReferenceOp(strict_bitwise=True) + return self.reference + + @staticmethod + def _validate_inputs(q: Tensor, k: Tensor, v: Tensor, contract: AttentionContract) -> None: + if q.dim() != 4 or k.dim() != 4 or v.dim() != 4: + raise AttentionContractError("q, k, and v must use [B, H, S, D] layout") + expected_dtype = _TORCH_DTYPES[contract.dtype] + if ( + q.dtype is not expected_dtype + or k.dtype is not expected_dtype + or v.dtype is not expected_dtype + ): + raise AttentionContractError( + f"q, k, and v must match contract dtype {contract.dtype.value}" + ) + if q.device != k.device or q.device != v.device: + raise AttentionContractError("q, k, and v must be on the same device") + batch, q_heads, q_seq, dim = q.shape + if batch != contract.batch_size: + raise AttentionContractError( + f"q batch={batch} does not match contract batch_size={contract.batch_size}" + ) + if q_seq != contract.query_sequence_length: + raise AttentionContractError( + "q sequence length does not match AttentionContract query_sequence_length" + ) + sharding = contract.sharding + if q_heads != sharding.local_q_heads or k.shape[1] != sharding.local_kv_heads: + raise AttentionContractError("q/k head counts do not match TP sharding in contract") + if ( + k.shape[0] != batch + or v.shape[:3] != k.shape[:3] + or k.shape[-1] != dim + or v.shape[-1] != dim + ): + raise AttentionContractError("q, k, and v shapes are inconsistent") + if dim != contract.head_dim: + raise AttentionContractError("tensor head_dim does not match AttentionContract") + + @staticmethod + def _validate_outputs(out: Tensor, lse: Tensor, q: Tensor, contract: AttentionContract) -> None: + if out.shape != q.shape: + raise AttentionContractError( + f"Attention output shape {tuple(out.shape)} does not match q {tuple(q.shape)}" + ) + expected_lse = (q.shape[0], q.shape[1], q.shape[2]) + if lse.shape != expected_lse: + raise AttentionContractError( + f"attention-domain LSE shape {tuple(lse.shape)} does not match {expected_lse}" + ) + if lse.dtype is not torch.float32: + raise AttentionContractError("attention-domain LSE must remain fp32") + expected_dtype = _TORCH_DTYPES[contract.dtype] + if out.dtype is not expected_dtype: + raise AttentionContractError( + f"Attention output must be written in {contract.dtype.value}, got {out.dtype}" + ) + + @staticmethod + def _invoke( + backend: Any, + q: Tensor, + k: Tensor, + v: Tensor, + kwargs: Mapping[str, Any], + contract: AttentionContract, + ) -> tuple[Tensor, Tensor, dict[str, Any]]: + method = getattr(backend, "forward_with_lse", None) + if not callable(method): + method = getattr(backend, "apply", None) + if not callable(method): + method = backend if callable(backend) else None + if method is None: + raise AttentionContractError( + "Attention backend must expose forward_with_lse, apply, or __call__" + ) + accepted = _accepted_kwargs(method, kwargs) + if contract.sharding.cp_world_size > 1 and not _declares_keyword(method, "cp_world_size"): + raise AttentionContractError( + "CP>1 requires an Attention backend that explicitly accepts cp_world_size" + ) + result = method(q, k, v, **accepted) + backend_provenance: dict[str, Any] = {} + out: Any + lse: Any + if isinstance(result, AttentionAblationResult): + out, lse = result.out, result.lse + elif isinstance(result, tuple) and len(result) == 2: + out, lse = result + else: + out = getattr(result, "out", None) + lse = getattr(result, "lse", None) + raw_provenance = getattr(result, "provenance", {}) + if isinstance(raw_provenance, Mapping): + backend_provenance = dict(raw_provenance) + if out is None or lse is None: + raise AttentionContractError( + "Attention backend must return (out, lse), AttentionAblationResult, " + "or an object with out/lse/provenance" + ) + if isinstance(result, AttentionAblationResult): + backend_provenance = dict(result.provenance) + if not isinstance(out, Tensor) or not isinstance(lse, Tensor): + raise AttentionContractError("Attention backend returned non-tensor output or LSE") + return out, lse, backend_provenance + + @staticmethod + def _backward( + backend: Any, + q: Tensor, + k: Tensor, + v: Tensor, + out: Tensor, + dout: Tensor | None, + kwargs: Mapping[str, Any], + ) -> tuple[Tensor, Tensor, Tensor]: + backward = getattr(backend, "backward_reference", None) + if callable(backward): + result = backward(q, k, v, dout, **_accepted_kwargs(backward, kwargs)) + gradients = getattr(result, "gradients", None) + if gradients is not None: + return gradients.dq, gradients.dk, gradients.dv + if dout is None: + raise AttentionContractError("dout is required to compute Attention gradients") + if not out.requires_grad: + raise AttentionContractError( + "Attention backend did not retain an autograd graph for gradients" + ) + return cast( + tuple[Tensor, Tensor, Tensor], + torch.autograd.grad( + out, + (q, k, v), + grad_outputs=dout.to(dtype=out.dtype), + allow_unused=False, + retain_graph=True, + ), + ) + + +def _resolve_config( + config: AttentionAblationConfig | Mapping[str, Any] | None, + **overrides: Any, +) -> AttentionAblationConfig: + if config is None: + values: dict[str, Any] = {} + elif isinstance(config, AttentionAblationConfig): + values = { + name: getattr(config, name) + for name in ( + "backend", + "deterministic", + "communication_backend", + "return_lse", + "return_gradients", + "strict_core_id", + "strict_schedule", + "validate", + ) + } + elif isinstance(config, Mapping): + values = dict(config) + else: + raise AttentionContractError("config must be AttentionAblationConfig, mapping, or None") + values.update({name: value for name, value in overrides.items() if value is not None}) + return AttentionAblationConfig(**values) + + +def _accepted_kwargs(method: Callable[..., Any], kwargs: Mapping[str, Any]) -> dict[str, Any]: + try: + signature = inspect.signature(method) + except (TypeError, ValueError): + return dict(kwargs) + parameters = signature.parameters.values() + if any(parameter.kind is inspect.Parameter.VAR_KEYWORD for parameter in parameters): + return dict(kwargs) + return {name: value for name, value in kwargs.items() if name in signature.parameters} + + +def _declares_keyword(method: Callable[..., Any], name: str) -> bool: + """Return whether ``method`` explicitly declares a keyword-capable parameter.""" + + try: + parameter = inspect.signature(method).parameters.get(name) + except (TypeError, ValueError): + return False + return parameter is not None and parameter.kind in { + inspect.Parameter.POSITIONAL_OR_KEYWORD, + inspect.Parameter.KEYWORD_ONLY, + } + + +def _callable_backend_id(value: Any) -> str: + explicit = getattr(value, "backend_id", None) or getattr( + value, "__attention_backend_id__", None + ) + if isinstance(explicit, str) and explicit.strip(): + return explicit.strip() + return f"injected.{type(value).__module__}.{type(value).__qualname__}" + + +def _validate_runtime_provenance( + selected: Any, + selected_id: str, + runtime: Mapping[str, Any], + config: AttentionAblationConfig, +) -> None: + """Fail closed when a production strict backend did not prove its identity.""" + + if not config.deterministic: + return + if selected_id == "native": + raise AttentionContractError( + "deterministic Attention cannot execute native Attention arithmetic" + ) + + external_production_backend = selected_id not in {BACKEND_ID, REFERENCE_BACKEND_ID} + if not external_production_backend: + return + + expected = { + "strict_core_id": config.strict_core_id, + "strict_schedule": config.strict_schedule, + "actual_backend": selected_id, + "production_ready": True, + "fallback": False, + "reference_only": False, + } + if config.communication_backend in _STRICT_AG_RS_BACKENDS: + expected["communication_backend"] = config.communication_backend + mismatches = [name for name, value in expected.items() if runtime.get(name) != value] + if not isinstance(runtime.get("native_attention_arithmetic"), bool): + mismatches.append("native_attention_arithmetic") + if mismatches: + raise AttentionContractError( + "production Attention runtime provenance is incomplete or mismatched: " + + ", ".join(mismatches) + ) + + +def _actual_split_provenance( + contract: AttentionContract, + *, + total_kv_tokens: int, + backend: str = BACKEND_ID, +) -> dict[str, Any]: + plan = contract.split_kv.resolve( + total_kv_tokens, + backend=backend, + ) + return plan.to_dict() + + +def _reduction_provenance(contract: AttentionContract) -> dict[str, Any]: + reduction = contract.reduction + return { + "merge": reduction.merge.value, + "acc_dtype": reduction.acc_dtype.value, + "order": reduction.order.value, + "downcast_at": reduction.downcast_at.value, + "engine": reduction.engine.value, + } + + +def _contract_fingerprint(contract: AttentionContract) -> str: + payload = json.dumps(contract.to_dict(), sort_keys=True, separators=(",", ":")) + return hashlib.sha256(payload.encode("utf-8")).hexdigest() + + +__all__ = [ + "AttentionAblationConfig", + "AttentionAblationOp", + "AttentionAblationResult", + "BACKEND_ID", + "REFERENCE_BACKEND_ID", +] diff --git a/rl_engine/kernels/ops/pytorch/attention/cp_attention.py b/rl_engine/kernels/ops/pytorch/attention/cp_attention.py new file mode 100644 index 00000000..194c80d5 --- /dev/null +++ b/rl_engine/kernels/ops/pytorch/attention/cp_attention.py @@ -0,0 +1,1182 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Deterministic context-parallel attention reference. + +This module is the correctness-first WS2 reference for CP-aware standard +softmax attention. It intentionally stays in PyTorch and uses fp32 partial +states so fused CUDA/Triton backends can validate their CP/LSE merge semantics +against a small, inspectable implementation. +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import Optional, Sequence + +import torch + +from rl_engine.kernels.attention_contract import ( + SplitKVExecutionPlan, + SplitKVMode, + SplitKVRuntimeCoordinate, + SplitKVRuntimePlanEntry, + SplitKVRuntimePlanSet, +) +from rl_engine.kernels.ops.pytorch.attention.standard_attn import NativeAttentionOp + + +@dataclass(frozen=True) +class AttentionPartialState: + """One KV block's attention state before deterministic LSE merge. + + ``out`` is already normalized within the local KV block and has shape + ``[B, Hq, Sq, D]``. ``lse`` is the local attention-domain log-sum-exp with + shape ``[B, Hq, Sq]``. ``block_start`` / ``block_end`` are logical global KV + positions and define the canonical merge order. + """ + + out: torch.Tensor + lse: torch.Tensor + block_start: int + block_end: int + + def __post_init__(self) -> None: + if self.out.ndim != 4: + raise ValueError("partial attention out must have shape [B, Hq, Sq, D]") + if self.lse.shape != self.out.shape[:3]: + raise ValueError("partial attention lse must have shape [B, Hq, Sq]") + if self.out.device != self.lse.device: + raise ValueError("partial attention out/lse must be on the same device") + if self.out.dtype is not torch.float32 or self.lse.dtype is not torch.float32: + raise ValueError("partial attention out/lse must remain FP32 before merge") + if self.block_start < 0: + raise ValueError("block_start must be non-negative") + if self.block_end < self.block_start: + raise ValueError("block_end must be >= block_start") + + +@dataclass(frozen=True) +class AttentionBackwardGradients: + """Training-side gradients emitted by the CP attention backward reference.""" + + dq: torch.Tensor + dk: torch.Tensor + dv: torch.Tensor + + +@dataclass(frozen=True) +class AttentionBackwardPathResult: + """One materialized CP attention backward path.""" + + name: str + out: torch.Tensor + lse: torch.Tensor + gradients: AttentionBackwardGradients + provenance: dict[str, object] + + +@dataclass(frozen=True) +class GradientDriftStats: + """Shape-aware absolute drift summary for backward validation reports.""" + + max_abs: float + mean_abs: float + p95_abs: float + p99_abs: float + active_count: int + + def to_dict(self) -> dict[str, object]: + return { + "max_abs": self.max_abs, + "mean_abs": self.mean_abs, + "p95_abs": self.p95_abs, + "p99_abs": self.p99_abs, + "active_count": self.active_count, + } + + +@dataclass(frozen=True) +class AttentionBackwardRankDrift: + """Backward drift for one logical CP rank's sequence ownership.""" + + rank: int + dq: GradientDriftStats + dk: GradientDriftStats + dv: GradientDriftStats + + def to_dict(self) -> dict[str, object]: + return { + "rank": self.rank, + "dq": self.dq.to_dict(), + "dk": self.dk.to_dict(), + "dv": self.dv.to_dict(), + } + + +@dataclass(frozen=True) +class AttentionBackwardPathDrift: + """Candidate-vs-reference backward drift for one CP path.""" + + candidate_name: str + dq: GradientDriftStats + dk: GradientDriftStats + dv: GradientDriftStats + out: GradientDriftStats + lse: GradientDriftStats + per_rank: tuple[AttentionBackwardRankDrift, ...] + provenance: dict[str, object] + + def to_dict(self) -> dict[str, object]: + return { + "candidate_name": self.candidate_name, + "dq": self.dq.to_dict(), + "dk": self.dk.to_dict(), + "dv": self.dv.to_dict(), + "out": self.out.to_dict(), + "lse": self.lse.to_dict(), + "per_rank": [item.to_dict() for item in self.per_rank], + "provenance": self.provenance, + } + + +@dataclass(frozen=True) +class AttentionBackwardComparisonReport: + """Structured PR8 report for CP attention gradient drift validation.""" + + reference_name: str + drifts: tuple[AttentionBackwardPathDrift, ...] + + def to_dict(self) -> dict[str, object]: + return { + "reference_name": self.reference_name, + "drifts": [drift.to_dict() for drift in self.drifts], + } + + +def merge_attention_partial_states( + states: Sequence[AttentionPartialState], +) -> AttentionPartialState: + """Merge CP/chunk partial states in logical block order. + + The merge is the online-softmax/LSE merge used by attention, not a plain + sum. The input order is deliberately ignored: states are sorted by logical + ``block_start`` so the result depends on global block indices rather than + arrival order. + """ + + if not states: + raise ValueError("at least one attention partial state is required") + + ordered = sorted(states, key=lambda item: (item.block_start, item.block_end)) + _validate_merge_shapes_and_ranges(ordered) + + merged = ordered[0] + merged_out = merged.out.float() + merged_lse = merged.lse.float() + for state in ordered[1:]: + merged_out, merged_lse = _merge_two_states( + merged_out, + merged_lse, + state.out.float(), + state.lse.float(), + ) + + return AttentionPartialState( + out=merged_out, + lse=merged_lse, + block_start=ordered[0].block_start, + block_end=ordered[-1].block_end, + ) + + +class DeterministicCPAttentionReferenceOp: + """Correctness-first CP attention reference for prefill and chunked prefill. + + The reference consumes attention-ready Q/K. For Qwen3 WS2 this means Q/K + have already passed QK-Norm and RoPE unless an outer contract explicitly + marks them as pre-RoPE. RoPE is intentionally kept outside this CP merge + implementation so fused and unfused ``RoPE+Attention`` paths can compare the + same post-RoPE Q/K boundary before validating CP communication. + + The op emulates CP by splitting query and KV sequence dimensions into + logical CP shards. Each query shard computes one partial attention state per + KV block, then merges those states in fixed global-block order using fp32 + LSE arithmetic. ``forward`` returns the input dtype after the final write; + ``forward_fp32`` keeps the fp32 merged output. + """ + + op_class = "attention" + + def __init__(self, *, strict_bitwise: bool = False) -> None: + if not isinstance(strict_bitwise, bool): + raise TypeError("strict_bitwise must be a bool") + self.strict_bitwise = strict_bitwise + + @staticmethod + def split_kv_execution_plans( + total_kv_tokens: int, + *, + cp_world_size: int = 1, + kv_chunk_size: Optional[int] = None, + ) -> list[dict[str, object]]: + """Export the actual logical Split-KV plan before execution.""" + + return split_kv_execution_plan_provenance( + total_kv_tokens, + cp_world_size=cp_world_size, + kv_chunk_size=kv_chunk_size, + backend="deterministic_cp_reference", + ) + + def __call__( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + causal: bool = True, + scale: Optional[float] = None, + key_padding_mask: Optional[torch.Tensor] = None, + query_position_offsets: Optional[torch.Tensor] = None, + key_position_offsets: Optional[torch.Tensor] = None, + cp_world_size: int = 1, + kv_chunk_size: Optional[int] = None, + ) -> torch.Tensor: + return self.forward( + q, + k, + v, + causal=causal, + scale=scale, + key_padding_mask=key_padding_mask, + query_position_offsets=query_position_offsets, + key_position_offsets=key_position_offsets, + cp_world_size=cp_world_size, + kv_chunk_size=kv_chunk_size, + ) + + def forward( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + causal: bool = True, + scale: Optional[float] = None, + key_padding_mask: Optional[torch.Tensor] = None, + query_position_offsets: Optional[torch.Tensor] = None, + key_position_offsets: Optional[torch.Tensor] = None, + cp_world_size: int = 1, + kv_chunk_size: Optional[int] = None, + ) -> torch.Tensor: + """Compute CP attention with fp32 accumulation and final input-dtype write.""" + + out, _ = self.forward_with_lse( + q, + k, + v, + causal=causal, + scale=scale, + key_padding_mask=key_padding_mask, + query_position_offsets=query_position_offsets, + key_position_offsets=key_position_offsets, + cp_world_size=cp_world_size, + kv_chunk_size=kv_chunk_size, + output_dtype=q.dtype, + ) + return out + + def forward_fp32( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + causal: bool = True, + scale: Optional[float] = None, + key_padding_mask: Optional[torch.Tensor] = None, + query_position_offsets: Optional[torch.Tensor] = None, + key_position_offsets: Optional[torch.Tensor] = None, + cp_world_size: int = 1, + kv_chunk_size: Optional[int] = None, + ) -> torch.Tensor: + """Compute CP attention with fp32 accumulation and fp32 output.""" + + out, _ = self.forward_fp32_with_lse( + q, + k, + v, + causal=causal, + scale=scale, + key_padding_mask=key_padding_mask, + query_position_offsets=query_position_offsets, + key_position_offsets=key_position_offsets, + cp_world_size=cp_world_size, + kv_chunk_size=kv_chunk_size, + ) + return out + + def forward_with_lse( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + causal: bool = True, + scale: Optional[float] = None, + key_padding_mask: Optional[torch.Tensor] = None, + query_position_offsets: Optional[torch.Tensor] = None, + key_position_offsets: Optional[torch.Tensor] = None, + cp_world_size: int = 1, + kv_chunk_size: Optional[int] = None, + output_dtype: Optional[torch.dtype] = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Return ``(out, lse)`` for the CP reference path. + + ``lse`` is always fp32 and in the attention domain. ``out`` is fp32 + until the final write, then downcast to ``output_dtype``. When omitted, + ``output_dtype`` defaults to the input dtype. + """ + + resolved_output_dtype = q.dtype if output_dtype is None else output_dtype + _validate_output_dtype(resolved_output_dtype) + if self.strict_bitwise: + out, lse = self._forward_strict_bitwise( + q, + k, + v, + causal=causal, + scale=scale, + key_padding_mask=key_padding_mask, + query_position_offsets=query_position_offsets, + key_position_offsets=key_position_offsets, + kv_chunk_size=kv_chunk_size, + ) + else: + out, lse = self._forward_impl( + q, + k, + v, + causal=causal, + scale=scale, + key_padding_mask=key_padding_mask, + query_position_offsets=query_position_offsets, + key_position_offsets=key_position_offsets, + cp_world_size=cp_world_size, + kv_chunk_size=kv_chunk_size, + ) + out = out.to(resolved_output_dtype) + return out, lse + + def forward_fp32_with_lse( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + causal: bool = True, + scale: Optional[float] = None, + key_padding_mask: Optional[torch.Tensor] = None, + query_position_offsets: Optional[torch.Tensor] = None, + key_position_offsets: Optional[torch.Tensor] = None, + cp_world_size: int = 1, + kv_chunk_size: Optional[int] = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Return fp32 ``(out, lse)`` for the CP reference path.""" + + return self.forward_with_lse( + q, + k, + v, + causal=causal, + scale=scale, + key_padding_mask=key_padding_mask, + query_position_offsets=query_position_offsets, + key_position_offsets=key_position_offsets, + cp_world_size=cp_world_size, + kv_chunk_size=kv_chunk_size, + output_dtype=torch.float32, + ) + + def _forward_strict_bitwise( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + causal: bool, + scale: Optional[float], + key_padding_mask: Optional[torch.Tensor], + query_position_offsets: Optional[torch.Tensor], + key_position_offsets: Optional[torch.Tensor], + kv_chunk_size: Optional[int], + ) -> tuple[torch.Tensor, torch.Tensor]: + """Execute one batch/CP-independent arithmetic schedule.""" + + _validate_qkv(q, k, v) + _validate_scale(scale) + batch, hq, sq, dim = q.shape + skv = k.size(2) + if key_padding_mask is not None: + if key_padding_mask.shape != (batch, skv): + raise ValueError("key_padding_mask must have shape [B, Skv]") + if key_padding_mask.dtype != torch.bool: + raise ValueError("key_padding_mask must be bool") + query_offsets = _normalize_position_offsets( + query_position_offsets, + batch, + q.device, + default=skv - sq, + name="query_position_offsets", + ) + key_offsets = _normalize_position_offsets( + key_position_offsets, + batch, + q.device, + default=0, + name="key_position_offsets", + ) + kv_bounds = _kv_block_bounds(skv, 1, kv_chunk_size) + out_rows: list[torch.Tensor] = [] + lse_rows: list[torch.Tensor] = [] + for batch_index in range(batch): + q_batch = q[batch_index : batch_index + 1].contiguous() + k_batch = k[batch_index : batch_index + 1].contiguous() + v_batch = v[batch_index : batch_index + 1].contiguous() + pad_batch = ( + None + if key_padding_mask is None + else key_padding_mask[batch_index : batch_index + 1].contiguous() + ) + query_offset = query_offsets[batch_index : batch_index + 1] + key_offset = key_offsets[batch_index : batch_index + 1] + query_rows: list[torch.Tensor] = [] + lse_query_rows: list[torch.Tensor] = [] + for query_index in range(sq): + q_row = q_batch[:, :, query_index : query_index + 1, :].contiguous() + states = [ + self.local_partial_state( + q_row, + k_batch[:, :, key_start:key_end, :].contiguous(), + v_batch[:, :, key_start:key_end, :].contiguous(), + q_start=query_index, + k_start=key_start, + total_kv_len=skv, + total_query_len=sq, + causal=causal, + scale=scale, + key_padding_mask=( + None + if pad_batch is None + else pad_batch[:, key_start:key_end].contiguous() + ), + query_position_offsets=query_offset, + key_position_offsets=key_offset, + ) + for key_start, key_end in kv_bounds + if key_start != key_end + ] + merged = merge_attention_partial_states(states) + query_rows.append(merged.out) + lse_query_rows.append(merged.lse) + out_rows.append(torch.cat(query_rows, dim=2)) + lse_rows.append(torch.cat(lse_query_rows, dim=2)) + return torch.cat(out_rows, dim=0), torch.cat(lse_rows, dim=0) + + def backward_reference( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + dout: torch.Tensor, + *, + causal: bool = True, + scale: Optional[float] = None, + key_padding_mask: Optional[torch.Tensor] = None, + query_position_offsets: Optional[torch.Tensor] = None, + key_position_offsets: Optional[torch.Tensor] = None, + cp_world_size: int = 1, + kv_chunk_size: Optional[int] = None, + output_dtype: Optional[torch.dtype] = torch.float32, + name: Optional[str] = None, + ) -> AttentionBackwardPathResult: + """Run the deterministic training-side backward validation path. + + The semantic backward input is ``dout`` plus the forward attention state + produced from the same Q/K/V, masks, position offsets, CP world, and KV + block order. The reference keeps the softmax/merge math in fp32 and + records the final-write dtype in provenance; decode backward is + intentionally out of scope for PR8. + """ + + _validate_qkv(q, k, v) + if dout.shape != q.shape: + raise ValueError("dout must have shape [B, Hq, Sq, D], matching q") + if not torch.is_floating_point(dout) or torch.is_complex(dout): + raise ValueError("dout must be a real floating-point tensor") + if dout.device != q.device: + raise ValueError("dout must be on the same device as q, k, and v") + if dout.dtype != q.dtype: + raise ValueError("dout must have the same dtype as q") + q_leaf = q.detach().clone().requires_grad_(True) + k_leaf = k.detach().clone().requires_grad_(True) + v_leaf = v.detach().clone().requires_grad_(True) + + resolved_output_dtype = q.dtype if output_dtype is None else output_dtype + _validate_output_dtype(resolved_output_dtype) + out, lse = self.forward_with_lse( + q_leaf, + k_leaf, + v_leaf, + causal=causal, + scale=scale, + key_padding_mask=key_padding_mask, + query_position_offsets=query_position_offsets, + key_position_offsets=key_position_offsets, + cp_world_size=cp_world_size, + kv_chunk_size=kv_chunk_size, + output_dtype=resolved_output_dtype, + ) + torch.autograd.backward(out, dout.to(dtype=out.dtype)) + if q_leaf.grad is None or k_leaf.grad is None or v_leaf.grad is None: + raise RuntimeError("CP attention backward did not produce dq/dk/dv") + + return AttentionBackwardPathResult( + name=name + or _backward_path_name(cp_world_size=cp_world_size, kv_chunk_size=kv_chunk_size), + out=out.detach(), + lse=lse.detach(), + gradients=AttentionBackwardGradients( + dq=q_leaf.grad.detach(), + dk=k_leaf.grad.detach(), + dv=v_leaf.grad.detach(), + ), + provenance={ + "attention_mode": "prefill" if kv_chunk_size is None else "chunked_prefill", + "gradient_mode": "training_backward", + "gradient_inputs": ["q", "k", "v"], + "gradient_outputs": ["out"], + "saved_forward_state": [ + "out", + "attention_lse", + "causal_mask", + "key_padding_mask", + "query_position_offsets", + "key_position_offsets", + "global_block_index", + ], + "cp_world_size": cp_world_size, + "kv_chunk_size": kv_chunk_size, + "requested_split_kv_policy": ("disabled" if kv_chunk_size is None else "fixed"), + "requested_split_kv_size": kv_chunk_size, + "actual_split_kv_plans": split_kv_execution_plan_provenance( + k.size(2), + cp_world_size=cp_world_size, + kv_chunk_size=kv_chunk_size, + backend="deterministic_cp_backward_reference", + ), + "merge_order": "global_block_index", + "accum_dtype": "fp32", + "downcast_at": "final_write", + "output_dtype": str(resolved_output_dtype).replace("torch.", ""), + "q_dtype": str(q.dtype).replace("torch.", ""), + "k_dtype": str(k.dtype).replace("torch.", ""), + "v_dtype": str(v.dtype).replace("torch.", ""), + "dout_dtype": str(dout.dtype).replace("torch.", ""), + "te_backward_oracle": "not_used", + "decode_backward": "not_supported", + }, + ) + + def local_partial_state( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + q_start: int, + k_start: int, + total_kv_len: int, + total_query_len: Optional[int] = None, + causal: bool = True, + scale: Optional[float] = None, + key_padding_mask: Optional[torch.Tensor] = None, + query_position_offsets: Optional[torch.Tensor] = None, + key_position_offsets: Optional[torch.Tensor] = None, + ) -> AttentionPartialState: + """Compute one query shard against one logical KV block. + + ``query_position_offsets`` and ``key_position_offsets`` are optional + per-batch-row base positions. They let the reference express varlen or + packed metadata while retaining the dense [B, H, S, D] tensor layout. + For post-RoPE Q/K, these offsets must describe the same absolute token + positions used when RoPE was applied. + """ + + _validate_qkv(q, k, v) + _validate_scale(scale) + if q_start < 0 or k_start < 0: + raise ValueError("q_start and k_start must be non-negative") + if total_kv_len < k_start + k.size(2): + raise ValueError("total_kv_len must cover the local KV block") + if total_query_len is None: + total_query_len = q.size(2) + if total_query_len < q_start + q.size(2): + raise ValueError("total_query_len must cover the local query block") + if key_padding_mask is not None: + if key_padding_mask.shape != (q.size(0), k.size(2)): + raise ValueError("local key_padding_mask must have shape [B, local_skv]") + if key_padding_mask.dtype != torch.bool: + raise ValueError("local key_padding_mask must be bool") + query_offsets = _normalize_position_offsets( + query_position_offsets, + q.size(0), + q.device, + default=total_kv_len - total_query_len, + name="query_position_offsets", + ) + key_offsets = _normalize_position_offsets( + key_position_offsets, + q.size(0), + q.device, + default=0, + name="key_position_offsets", + ) + + ctx = NativeAttentionOp._strict_fp32_math(q.device.type) + with ctx: + qf = q.float() + kf = k.float() + vf = v.float() + hq, sq, dim = qf.shape[1], qf.shape[2], qf.shape[3] + hkv, skv = kf.shape[1], kf.shape[2] + if hkv != hq: + repeat = hq // hkv + kf = kf.repeat_interleave(repeat, dim=1) + vf = vf.repeat_interleave(repeat, dim=1) + + if skv == 0: + zero_dep = _zero_dependency(qf, kf, vf) + return AttentionPartialState( + out=torch.zeros(q.size(0), hq, sq, dim, device=q.device, dtype=torch.float32) + + zero_dep, + lse=torch.full( + (q.size(0), hq, sq), + float("-inf"), + device=q.device, + dtype=torch.float32, + ) + + zero_dep, + block_start=k_start, + block_end=k_start, + ) + + scale_value = scale if scale is not None else (1.0 / math.sqrt(dim)) + scores = torch.matmul(qf, kf.transpose(-1, -2)) * scale_value + if causal: + query_base = query_offsets[:, None] + q_start + key_base = key_offsets[:, None] + k_start + q_pos = torch.arange(sq, device=q.device, dtype=torch.long) + query_base + k_pos = torch.arange(skv, device=q.device, dtype=torch.long) + key_base + causal_mask = k_pos[:, None, :] > q_pos[:, :, None] + scores = scores.masked_fill(causal_mask[:, None, :, :], float("-inf")) + if key_padding_mask is not None: + scores = scores.masked_fill(~key_padding_mask[:, None, None, :], float("-inf")) + + lse = torch.logsumexp(scores, dim=-1) + finite_lse = torch.isfinite(lse) + weights = torch.exp(scores - lse.unsqueeze(-1)) + weights = torch.where(finite_lse.unsqueeze(-1), weights, torch.zeros_like(weights)) + out = torch.matmul(weights, vf) + return AttentionPartialState( + out=out, + lse=lse, + block_start=k_start, + block_end=k_start + skv, + ) + + def _forward_impl( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + causal: bool, + scale: Optional[float], + key_padding_mask: Optional[torch.Tensor], + query_position_offsets: Optional[torch.Tensor], + key_position_offsets: Optional[torch.Tensor], + cp_world_size: int, + kv_chunk_size: Optional[int], + ) -> tuple[torch.Tensor, torch.Tensor]: + _validate_qkv(q, k, v) + _validate_scale(scale) + if ( + isinstance(cp_world_size, bool) + or not isinstance(cp_world_size, int) + or cp_world_size < 1 + ): + raise ValueError("cp_world_size must be >= 1") + if kv_chunk_size is not None and ( + isinstance(kv_chunk_size, bool) + or not isinstance(kv_chunk_size, int) + or kv_chunk_size < 1 + ): + raise ValueError("kv_chunk_size must be >= 1 when provided") + + batch, hq, sq, dim = q.shape + skv = k.size(2) + if key_padding_mask is not None: + if key_padding_mask.shape != (batch, skv): + raise ValueError("key_padding_mask must have shape [B, Skv]") + if key_padding_mask.dtype != torch.bool: + raise ValueError("key_padding_mask must be bool") + query_offsets = _normalize_position_offsets( + query_position_offsets, + batch, + q.device, + default=skv - sq, + name="query_position_offsets", + ) + key_offsets = _normalize_position_offsets( + key_position_offsets, + batch, + q.device, + default=0, + name="key_position_offsets", + ) + + q_bounds = _split_bounds(sq, cp_world_size) + kv_bounds = _kv_block_bounds(skv, cp_world_size, kv_chunk_size) + out_chunks: list[torch.Tensor] = [] + lse_chunks: list[torch.Tensor] = [] + for q_start, q_end in q_bounds: + if q_start == q_end: + continue + q_block = q[:, :, q_start:q_end, :] + states = [ + self.local_partial_state( + q_block, + k[:, :, k_start:k_end, :], + v[:, :, k_start:k_end, :], + q_start=q_start, + k_start=k_start, + total_kv_len=skv, + total_query_len=sq, + causal=causal, + scale=scale, + key_padding_mask=( + None if key_padding_mask is None else key_padding_mask[:, k_start:k_end] + ), + query_position_offsets=query_offsets, + key_position_offsets=key_offsets, + ) + for k_start, k_end in kv_bounds + if k_start != k_end + ] + if states: + merged = merge_attention_partial_states(states) + out_chunks.append(merged.out) + lse_chunks.append(merged.lse) + else: + zero_dep = _zero_dependency(q_block.float(), k.float(), v.float()) + out_chunks.append( + torch.zeros(batch, hq, q_end - q_start, dim, device=q.device) + zero_dep + ) + lse_chunks.append( + torch.full( + (batch, hq, q_end - q_start), + float("-inf"), + device=q.device, + dtype=torch.float32, + ) + + zero_dep + ) + + if not out_chunks: + zero_dep = _zero_dependency(q.float(), k.float(), v.float()) + return ( + torch.empty(batch, hq, 0, dim, device=q.device, dtype=torch.float32) + zero_dep, + torch.empty(batch, hq, 0, device=q.device, dtype=torch.float32) + zero_dep, + ) + return torch.cat(out_chunks, dim=2), torch.cat(lse_chunks, dim=2) + + +def compare_cp_attention_backward( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + dout: torch.Tensor, + *, + causal: bool = True, + scale: Optional[float] = None, + key_padding_mask: Optional[torch.Tensor] = None, + query_position_offsets: Optional[torch.Tensor] = None, + key_position_offsets: Optional[torch.Tensor] = None, + candidate_cp_world_size: int = 2, + candidate_kv_chunk_size: Optional[int] = None, + output_dtype: Optional[torch.dtype] = torch.float32, +) -> AttentionBackwardComparisonReport: + """Compare CP=1 backward with a CP/chunked-prefill candidate. + + The report includes whole-tensor ``dq/dk/dv`` drift and per-logical-CP-rank + slices. It is a validation/reporting helper, not a separate production + backward kernel. + """ + + op = DeterministicCPAttentionReferenceOp() + reference = op.backward_reference( + q, + k, + v, + dout, + causal=causal, + scale=scale, + key_padding_mask=key_padding_mask, + query_position_offsets=query_position_offsets, + key_position_offsets=key_position_offsets, + cp_world_size=1, + kv_chunk_size=None, + output_dtype=output_dtype, + name="cp1_backward_reference", + ) + candidate = op.backward_reference( + q, + k, + v, + dout, + causal=causal, + scale=scale, + key_padding_mask=key_padding_mask, + query_position_offsets=query_position_offsets, + key_position_offsets=key_position_offsets, + cp_world_size=candidate_cp_world_size, + kv_chunk_size=candidate_kv_chunk_size, + output_dtype=output_dtype, + ) + return AttentionBackwardComparisonReport( + reference_name=reference.name, + drifts=(_compare_backward_path(candidate, reference),), + ) + + +def _compare_backward_path( + candidate: AttentionBackwardPathResult, + reference: AttentionBackwardPathResult, +) -> AttentionBackwardPathDrift: + cp_world_size = _provenance_int(candidate.provenance, "cp_world_size") + return AttentionBackwardPathDrift( + candidate_name=candidate.name, + dq=_drift_stats(candidate.gradients.dq, reference.gradients.dq), + dk=_drift_stats(candidate.gradients.dk, reference.gradients.dk), + dv=_drift_stats(candidate.gradients.dv, reference.gradients.dv), + out=_drift_stats(candidate.out, reference.out), + lse=_drift_stats(candidate.lse, reference.lse), + per_rank=_per_rank_backward_drifts(candidate, reference, cp_world_size), + provenance=candidate.provenance, + ) + + +def _per_rank_backward_drifts( + candidate: AttentionBackwardPathResult, + reference: AttentionBackwardPathResult, + cp_world_size: int, +) -> tuple[AttentionBackwardRankDrift, ...]: + q_bounds = _split_bounds(candidate.gradients.dq.size(2), cp_world_size) + kv_bounds = _split_bounds(candidate.gradients.dk.size(2), cp_world_size) + per_rank = [] + for rank, ((q_start, q_end), (kv_start, kv_end)) in enumerate(zip(q_bounds, kv_bounds)): + per_rank.append( + AttentionBackwardRankDrift( + rank=rank, + dq=_drift_stats( + candidate.gradients.dq[:, :, q_start:q_end, :], + reference.gradients.dq[:, :, q_start:q_end, :], + ), + dk=_drift_stats( + candidate.gradients.dk[:, :, kv_start:kv_end, :], + reference.gradients.dk[:, :, kv_start:kv_end, :], + ), + dv=_drift_stats( + candidate.gradients.dv[:, :, kv_start:kv_end, :], + reference.gradients.dv[:, :, kv_start:kv_end, :], + ), + ) + ) + return tuple(per_rank) + + +def _drift_stats(candidate: torch.Tensor, reference: torch.Tensor) -> GradientDriftStats: + if candidate.shape != reference.shape: + raise ValueError( + f"candidate shape {tuple(candidate.shape)} must match " + f"reference shape {tuple(reference.shape)}" + ) + diff = (candidate.float() - reference.float()).abs().reshape(-1) + active_count = int(diff.numel()) + if active_count == 0: + return GradientDriftStats(0.0, 0.0, 0.0, 0.0, 0) + return GradientDriftStats( + max_abs=float(diff.max().item()), + mean_abs=float(diff.mean().item()), + p95_abs=float(torch.quantile(diff, 0.95).item()), + p99_abs=float(torch.quantile(diff, 0.99).item()), + active_count=active_count, + ) + + +def _backward_path_name(*, cp_world_size: int, kv_chunk_size: Optional[int]) -> str: + prefix = f"cp{cp_world_size}" + if kv_chunk_size is None: + return f"{prefix}_backward" + return f"{prefix}_chunked_backward" + + +def _provenance_int(provenance: dict[str, object], key: str) -> int: + value = provenance[key] + if isinstance(value, bool) or not isinstance(value, int): + raise ValueError(f"provenance field {key!r} must be an int") + return value + + +def _merge_two_states( + out_a: torch.Tensor, + lse_a: torch.Tensor, + out_b: torch.Tensor, + lse_b: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + merged_lse = torch.logaddexp(lse_a, lse_b) + finite = torch.isfinite(merged_lse) + weight_a = torch.where(finite, torch.exp(lse_a - merged_lse), torch.zeros_like(merged_lse)) + weight_b = torch.where(finite, torch.exp(lse_b - merged_lse), torch.zeros_like(merged_lse)) + merged_out = weight_a.unsqueeze(-1) * out_a + weight_b.unsqueeze(-1) * out_b + return merged_out, merged_lse + + +def _validate_merge_shapes_and_ranges(states: Sequence[AttentionPartialState]) -> None: + first = states[0] + previous_end = first.block_end + for state in states[1:]: + if state.out.shape != first.out.shape or state.lse.shape != first.lse.shape: + raise ValueError("all partial states must have matching out/lse shapes") + if state.block_start != previous_end: + raise ValueError("partial state block ranges must be gap-free and non-overlapping") + previous_end = state.block_end + + +def _validate_qkv(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> None: + if q.ndim != 4 or k.ndim != 4 or v.ndim != 4: + raise ValueError("q, k, and v must have shape [B, H, S, D]") + if k.shape != v.shape: + raise ValueError("k and v must have the same shape") + if q.size(0) != k.size(0) or q.size(3) != k.size(3): + raise ValueError("q, k, and v must share batch size and head dim") + if q.size(1) < 1 or k.size(1) < 1 or q.size(3) < 1: + raise ValueError("q, k, and v must have positive head counts and head dim") + if not all(torch.is_floating_point(tensor) for tensor in (q, k, v)) or any( + torch.is_complex(tensor) for tensor in (q, k, v) + ): + raise ValueError("q, k, and v must be real floating-point tensors") + if q.dtype != k.dtype or q.dtype != v.dtype: + raise ValueError("q, k, and v must have the same dtype") + if q.device != k.device or q.device != v.device: + raise ValueError("q, k, and v must be on the same device") + if q.size(1) % k.size(1) != 0: + raise ValueError(f"Hq={q.size(1)} not divisible by Hkv={k.size(1)} (GQA group)") + + +def _validate_scale(scale: Optional[float]) -> None: + if scale is None: + return + if isinstance(scale, bool) or not isinstance(scale, (int, float)): + raise ValueError("scale must be a positive finite number") + if not math.isfinite(float(scale)) or float(scale) <= 0: + raise ValueError("scale must be a positive finite number") + + +def _validate_output_dtype(output_dtype: torch.dtype) -> None: + if not isinstance(output_dtype, torch.dtype): + raise ValueError("output_dtype must be a real floating-point torch dtype") + probe = torch.empty((), dtype=output_dtype) + if not torch.is_floating_point(probe) or torch.is_complex(probe): + raise ValueError("output_dtype must be a real floating-point torch dtype") + + +def _zero_dependency(*tensors: torch.Tensor) -> torch.Tensor: + total = torch.tensor(0.0, device=tensors[0].device) + for tensor in tensors: + total = total + tensor.sum() + return total * 0.0 + + +def _normalize_position_offsets( + offsets: Optional[torch.Tensor], + batch: int, + device: torch.device, + *, + default: int, + name: str, +) -> torch.Tensor: + if offsets is None: + return torch.full((batch,), default, dtype=torch.long, device=device) + if offsets.ndim != 1 or offsets.numel() != batch: + raise ValueError(f"{name} must have shape [B]") + if torch.is_floating_point(offsets) or torch.is_complex(offsets) or offsets.dtype == torch.bool: + raise ValueError(f"{name} must contain integer positions") + return offsets.to(device=device, dtype=torch.long) + + +def _split_bounds(length: int, parts: int) -> list[tuple[int, int]]: + base, extra = divmod(length, parts) + bounds: list[tuple[int, int]] = [] + start = 0 + for index in range(parts): + width = base + (1 if index < extra else 0) + end = start + width + bounds.append((start, end)) + start = end + return bounds + + +def _kv_block_bounds( + length: int, + cp_world_size: int, + kv_chunk_size: Optional[int], +) -> list[tuple[int, int]]: + bounds: list[tuple[int, int]] = [] + for start, end in _split_bounds(length, cp_world_size): + if kv_chunk_size is None: + bounds.append((start, end)) + continue + cursor = start + while cursor < end: + chunk_end = min(cursor + kv_chunk_size, end) + bounds.append((cursor, chunk_end)) + cursor = chunk_end + return bounds + + +def split_kv_execution_plan_provenance( + length: int, + *, + cp_world_size: int, + kv_chunk_size: Optional[int], + backend: str, +) -> list[dict[str, object]]: + """Return the actual backend-local Split-KV plan for every CP owner.""" + + if length < 1: + raise ValueError("Split-KV sequence length must be >= 1") + if cp_world_size < 1: + raise ValueError("cp_world_size must be >= 1") + if kv_chunk_size is not None and kv_chunk_size < 1: + raise ValueError("kv_chunk_size must be >= 1 when provided") + result: list[dict[str, object]] = [] + for owner_cp_rank, (rank_start, rank_end) in enumerate(_split_bounds(length, cp_world_size)): + if rank_start == rank_end: + continue + boundaries: tuple[tuple[int, int], ...] + if kv_chunk_size is None: + boundaries = ((rank_start, rank_end),) + mode = SplitKVMode.DISABLED + else: + boundaries = tuple( + (start, min(start + kv_chunk_size, rank_end)) + for start in range(rank_start, rank_end, kv_chunk_size) + ) + mode = SplitKVMode.FIXED + plan = SplitKVExecutionPlan( + requested_mode=mode, + requested_split_size=kv_chunk_size, + actual_mode=mode, + actual_split_size=kv_chunk_size, + boundaries=boundaries, + backend=backend, + source="reference_execution", + ) + result.append({"owner_cp_rank": owner_cp_rank, **plan.to_dict()}) + return result + + +def build_reference_split_kv_runtime_plan_set( + total_kv_tokens: Sequence[int], + *, + tp_world_size: int, + cp_world_size: int, + kv_chunk_size: Optional[int], + backend: str = "deterministic_cp_reference", +) -> SplitKVRuntimePlanSet: + """Build complete per-batch/TP/CP/owner plans for the reference path.""" + + totals = tuple(total_kv_tokens) + if not totals or any(total < cp_world_size for total in totals): + raise ValueError("reference runtime plan sets require at least one KV token per CP owner") + if tp_world_size < 1 or cp_world_size < 1: + raise ValueError("TP and CP world sizes must be >= 1") + if kv_chunk_size is not None and kv_chunk_size < 1: + raise ValueError("kv_chunk_size must be >= 1 when provided") + + entries: list[SplitKVRuntimePlanEntry] = [] + for batch_index, total in enumerate(totals): + owner_ranges = _split_bounds(total, cp_world_size) + for tp_rank in range(tp_world_size): + for cp_rank in range(cp_world_size): + for owner_cp_rank, (owner_start, owner_end) in enumerate(owner_ranges): + boundaries: tuple[tuple[int, int], ...] + if kv_chunk_size is None: + mode = SplitKVMode.DISABLED + boundaries = ((owner_start, owner_end),) + else: + mode = SplitKVMode.FIXED + boundaries = tuple( + (start, min(start + kv_chunk_size, owner_end)) + for start in range(owner_start, owner_end, kv_chunk_size) + ) + execution = SplitKVExecutionPlan( + requested_mode=mode, + requested_split_size=kv_chunk_size, + actual_mode=mode, + actual_split_size=kv_chunk_size, + boundaries=boundaries, + backend=backend, + source="reference_execution", + ) + entries.append( + SplitKVRuntimePlanEntry( + coordinate=SplitKVRuntimeCoordinate( + batch_index=batch_index, + tp_rank=tp_rank, + cp_rank=cp_rank, + owner_cp_rank=owner_cp_rank, + ), + expected_kv_range=(owner_start, owner_end), + execution=execution, + ) + ) + return SplitKVRuntimePlanSet( + batch_size=len(totals), + tp_world_size=tp_world_size, + cp_world_size=cp_world_size, + total_kv_tokens=totals, + entries=tuple(entries), + ) + + +CPAttentionReferenceOp = DeterministicCPAttentionReferenceOp + +__all__ = [ + "AttentionBackwardComparisonReport", + "AttentionBackwardGradients", + "AttentionBackwardPathDrift", + "AttentionBackwardPathResult", + "AttentionBackwardRankDrift", + "AttentionPartialState", + "build_reference_split_kv_runtime_plan_set", + "CPAttentionReferenceOp", + "DeterministicCPAttentionReferenceOp", + "GradientDriftStats", + "compare_cp_attention_backward", + "merge_attention_partial_states", + "split_kv_execution_plan_provenance", +] diff --git a/rl_engine/kernels/ops/pytorch/attention/debug_matrix.py b/rl_engine/kernels/ops/pytorch/attention/debug_matrix.py new file mode 100644 index 00000000..ee76e5f8 --- /dev/null +++ b/rl_engine/kernels/ops/pytorch/attention/debug_matrix.py @@ -0,0 +1,210 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Replay matrix for post-training Attention drift triage. + +The matrix deliberately separates a fixed replay baseline, one-at-a-time root- +cause probes, and invariant controls. It is a reporting/debug contract, not a +second runtime knob catalog and it never implies a Cartesian product. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from types import MappingProxyType +from typing import Any + +from rl_engine.kernels.attention_contract import AttentionContractError +from rl_engine.kernels.ops.pytorch.attention.debug_taxonomy import ( + ATTENTION_DEBUG_AXES, + ATTENTION_INVARIANT_CONTROLS, +) + +ATTENTION_DEBUG_MATRIX_SCHEMA_VERSION = "rlkernel.attention.debug_matrix.v1" + + +@dataclass(frozen=True) +class AttentionDebugMatrixRow: + """One independent case in the replay matrix.""" + + row_id: str + label: str + category: str + probe: str | None = None + root_cause_axis: str | None = None + expected: str = "diagnostic" + + def __post_init__(self) -> None: + if not self.row_id.strip() or not self.label.strip(): + raise ValueError("Attention debug matrix rows need an id and label") + if self.category not in { + "baseline", + "root_cause", + "comparability_gate", + "invariant_control", + }: + raise ValueError(f"unknown Attention debug matrix category {self.category!r}") + if self.category == "baseline": + if self.probe is not None or self.root_cause_axis is not None: + raise ValueError("baseline rows cannot name a probe or root-cause axis") + if self.expected != "baseline": + raise ValueError("baseline rows must use expected='baseline'") + elif not self.probe or not self.probe.strip(): + raise ValueError("non-baseline rows need a probe") + if self.category in {"root_cause", "comparability_gate"}: + if not self.root_cause_axis: + raise ValueError("root-cause and gate rows need root_cause_axis") + if self.category == "root_cause": + if self.expected != "diagnostic": + raise ValueError("root-cause rows must use expected='diagnostic'") + if self.category == "comparability_gate" and self.expected != "rejected": + raise ValueError("comparability gates must use expected='rejected'") + if self.category == "invariant_control": + if self.expected != "exact_zero": + raise ValueError("invariant controls must use expected='exact_zero'") + + def to_dict(self) -> dict[str, Any]: + return { + "id": self.row_id, + "label": self.label, + "category": self.category, + "probe": self.probe, + "root_cause_axis": self.root_cause_axis, + "expected": self.expected, + } + + +def _build_rows() -> tuple[AttentionDebugMatrixRow, ...]: + rows: list[AttentionDebugMatrixRow] = [ + AttentionDebugMatrixRow( + row_id="A0", + label="Strict replay baseline", + category="baseline", + expected="baseline", + ) + ] + for index, axis in enumerate(ATTENTION_DEBUG_AXES, start=1): + is_gate = axis.axis_id == "topology_head_ownership" + rows.append( + AttentionDebugMatrixRow( + row_id=f"A{index}", + label=axis.label, + category="comparability_gate" if is_gate else "root_cause", + probe=axis.representative_subprobe, + root_cause_axis=axis.axis_id, + expected="rejected" if is_gate else "diagnostic", + ) + ) + for index, probe in enumerate(ATTENTION_INVARIANT_CONTROLS): + rows.append( + AttentionDebugMatrixRow( + row_id=f"C{index}", + label="Invariant control", + category="invariant_control", + probe=probe, + expected="exact_zero", + ) + ) + return tuple(rows) + + +ATTENTION_DEBUG_MATRIX = _build_rows() +_ROWS_BY_ID = MappingProxyType({row.row_id: row for row in ATTENTION_DEBUG_MATRIX}) + + +def validate_attention_debug_matrix() -> None: + """Validate matrix coverage and the no-Cartesian-product invariant.""" + + if len(_ROWS_BY_ID) != len(ATTENTION_DEBUG_MATRIX): + raise RuntimeError("Attention debug matrix row IDs must be unique") + baselines = [row for row in ATTENTION_DEBUG_MATRIX if row.category == "baseline"] + if len(baselines) != 1 or baselines[0].row_id != "A0": + raise RuntimeError("Attention debug matrix must contain exactly one A0 baseline") + + axis_ids = {axis.axis_id for axis in ATTENTION_DEBUG_AXES} + diagnostic_rows = [ + row + for row in ATTENTION_DEBUG_MATRIX + if row.category in {"root_cause", "comparability_gate"} + ] + if {row.root_cause_axis for row in diagnostic_rows} != axis_ids: + raise RuntimeError("every root-cause axis must have one representative row") + if len(diagnostic_rows) != len(axis_ids): + raise RuntimeError("root-cause representatives must be one-at-a-time") + gates = [row for row in diagnostic_rows if row.category == "comparability_gate"] + if [row.root_cause_axis for row in gates] != ["topology_head_ownership"]: + raise RuntimeError("only topology/head ownership is a comparability gate") + + controls = [row for row in ATTENTION_DEBUG_MATRIX if row.category == "invariant_control"] + if {row.probe for row in controls} != set(ATTENTION_INVARIANT_CONTROLS): + raise RuntimeError("invariant-control coverage is incomplete") + if any(row.expected != "exact_zero" for row in controls): + raise RuntimeError("invariant controls must be exact-zero checks") + + axis_by_id = {axis.axis_id: axis for axis in ATTENTION_DEBUG_AXES} + for row in diagnostic_rows: + axis_id = row.root_cause_axis + if axis_id is None: + raise RuntimeError(f"diagnostic row {row.row_id} must name a root-cause axis") + axis = axis_by_id[axis_id] + if row.probe != axis.representative_subprobe: + raise RuntimeError( + f"matrix row {row.row_id} is not the representative probe for {axis.axis_id}" + ) + + +validate_attention_debug_matrix() + + +def attention_debug_matrix() -> dict[str, Any]: + """Return the portable matrix manifest used by reports and tooling.""" + + return { + "schema_version": ATTENTION_DEBUG_MATRIX_SCHEMA_VERSION, + "method": "fixed_replay_one_at_a_time", + "baseline_row": "A0", + "cartesian_product": False, + "replay_identity": ( + "same checkpoint, token IDs, selected-token IDs, masks, positions, " + "cache metadata, and pre-update model state" + ), + "row_baseline": "each diagnostic row is compared with its own phase-local A0 baseline", + "metrics": [ + "train_rollout_logprob_abs_diff", + "mismatch_kl", + "mismatch_k3_kl", + "out_max_abs", + "lse_max_abs", + "dq_max_abs", + "dk_max_abs", + "dv_max_abs", + ], + "rows": [row.to_dict() for row in ATTENTION_DEBUG_MATRIX], + "topology_gate": [ + "checkpoint/model/token identity", + "TP head ownership", + "CP sequence ownership", + "actual Split-KV plan", + ], + } + + +def attention_debug_matrix_row(row_id: str) -> AttentionDebugMatrixRow: + """Look up a stable matrix row by its report ID.""" + + if not isinstance(row_id, str) or not row_id.strip(): + raise AttentionContractError("Attention debug matrix row ID must be non-empty") + try: + return _ROWS_BY_ID[row_id.strip()] + except KeyError as exc: + raise AttentionContractError(f"unknown Attention debug matrix row {row_id!r}") from exc + + +__all__ = [ + "ATTENTION_DEBUG_MATRIX", + "ATTENTION_DEBUG_MATRIX_SCHEMA_VERSION", + "AttentionDebugMatrixRow", + "attention_debug_matrix", + "attention_debug_matrix_row", + "validate_attention_debug_matrix", +] diff --git a/rl_engine/kernels/ops/pytorch/attention/debug_taxonomy.py b/rl_engine/kernels/ops/pytorch/attention/debug_taxonomy.py new file mode 100644 index 00000000..02bc021c --- /dev/null +++ b/rl_engine/kernels/ops/pytorch/attention/debug_taxonomy.py @@ -0,0 +1,176 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Compact root-cause taxonomy for post-training Attention drift reports.""" + +from __future__ import annotations + +from dataclasses import dataclass +from types import MappingProxyType +from typing import Any + +from rl_engine.kernels.attention_contract import AttentionContractError + +ATTENTION_DEBUG_SCHEMA_VERSION = "rlkernel.attention.debug_taxonomy.v1" + + +@dataclass(frozen=True) +class AttentionDebugAxisSpec: + """One first-line post-training Attention drift category.""" + + axis_id: str + label: str + representative_subprobe: str + subprobes: tuple[str, ...] + + def __post_init__(self) -> None: + for name in ("axis_id", "label", "representative_subprobe"): + value = getattr(self, name) + if not isinstance(value, str) or not value.strip(): + raise ValueError(f"Attention debug {name} must be a non-empty string") + if not self.subprobes or any( + not isinstance(probe, str) or not probe.strip() for probe in self.subprobes + ): + raise ValueError("Attention debug subprobes must be non-empty strings") + if len(set(self.subprobes)) != len(self.subprobes): + raise ValueError(f"Attention debug axis {self.axis_id!r} has duplicate subprobes") + if self.representative_subprobe not in self.subprobes: + raise ValueError( + f"Attention debug representative {self.representative_subprobe!r} " + f"is not assigned to {self.axis_id!r}" + ) + + +ATTENTION_DEBUG_AXES = ( + AttentionDebugAxisSpec( + axis_id="position_rope", + label="Position / RoPE", + representative_subprobe="position_ids", + subprobes=( + "position_ids", + "rope_theta", + "position_offsets", + "decode_cache_position", + ), + ), + AttentionDebugAxisSpec( + axis_id="qk_preprocessing", + label="Q/K preprocessing", + representative_subprobe="qk_norm_disabled", + subprobes=( + "qk_norm_eps", + "qk_norm_disabled", + "qk_norm_weight", + "attention_scale", + "scale_placement", + ), + ), + AttentionDebugAxisSpec( + axis_id="mask_sequence_boundary", + label="Mask / sequence boundary", + representative_subprobe="causal_mask", + subprobes=("causal_mask", "key_padding_mask"), + ), + AttentionDebugAxisSpec( + axis_id="topology_head_ownership", + label="Topology / head ownership", + representative_subprobe="tp_head_ownership", + subprobes=("tp_head_ownership",), + ), + AttentionDebugAxisSpec( + axis_id="kv_cache_identity_layout", + label="KV-cache identity / layout", + representative_subprobe="kv_page_order", + subprobes=("kv_page_order", "kv_cache_content"), + ), + AttentionDebugAxisSpec( + axis_id="numerical_policy", + label="Numerical policy", + representative_subprobe="accum_dtype", + subprobes=( + "accum_dtype", + "execution_dtype", + "final_write_dtype", + "early_downcast", + ), + ), + AttentionDebugAxisSpec( + axis_id="distributed_schedule", + label="Distributed schedule", + representative_subprobe="merge_order", + subprobes=("nonstrict_cp_degree", "split_kv", "merge_order"), + ), +) + +ATTENTION_INVARIANT_CONTROLS = ( + "tp_partition_control", + "batch_composition_control", + "prefill_decode_tail_control", +) + +_ATTENTION_DEBUG_AXIS_BY_ID = MappingProxyType( + {axis.axis_id: axis for axis in ATTENTION_DEBUG_AXES} +) +_ATTENTION_DEBUG_AXIS_BY_PROBE = MappingProxyType( + {probe: axis for axis in ATTENTION_DEBUG_AXES for probe in axis.subprobes} +) +if len(_ATTENTION_DEBUG_AXIS_BY_ID) != len(ATTENTION_DEBUG_AXES): + raise RuntimeError("Attention debug axis IDs must be unique") +if len(_ATTENTION_DEBUG_AXIS_BY_PROBE) != sum(len(axis.subprobes) for axis in ATTENTION_DEBUG_AXES): + raise RuntimeError("Attention debug subprobes must belong to one root-cause axis") +if set(_ATTENTION_DEBUG_AXIS_BY_PROBE) & set(ATTENTION_INVARIANT_CONTROLS): + raise RuntimeError("Attention debug subprobes cannot also be invariant controls") + + +def attention_debug_probe_metadata(probe: str) -> dict[str, Any]: + """Classify one stable debug probe for report aggregation.""" + + if not isinstance(probe, str) or not probe.strip(): + raise AttentionContractError("Attention debug probe must be a non-empty string") + normalized = probe.strip() + if normalized in ATTENTION_INVARIANT_CONTROLS: + return { + "category": "invariant_control", + "root_cause_axis": None, + "root_cause_label": None, + "representative": False, + } + axis = _ATTENTION_DEBUG_AXIS_BY_PROBE.get(normalized) + if axis is None: + raise AttentionContractError(f"unknown Attention debug probe {probe!r}") + return { + "category": "root_cause_subprobe", + "root_cause_axis": axis.axis_id, + "root_cause_label": axis.label, + "representative": normalized == axis.representative_subprobe, + } + + +def attention_debug_taxonomy() -> dict[str, Any]: + """Return the compact JSON schema used by post-training drift reports.""" + + return { + "schema_version": ATTENTION_DEBUG_SCHEMA_VERSION, + "root_cause_axis_count": len(ATTENTION_DEBUG_AXES), + "subprobe_count": len(_ATTENTION_DEBUG_AXIS_BY_PROBE), + "invariant_control_count": len(ATTENTION_INVARIANT_CONTROLS), + "root_cause_axes": { + axis.axis_id: { + "label": axis.label, + "representative_subprobe": axis.representative_subprobe, + "subprobes": list(axis.subprobes), + } + for axis in ATTENTION_DEBUG_AXES + }, + "invariant_controls": list(ATTENTION_INVARIANT_CONTROLS), + } + + +__all__ = [ + "ATTENTION_DEBUG_AXES", + "ATTENTION_DEBUG_SCHEMA_VERSION", + "ATTENTION_INVARIANT_CONTROLS", + "AttentionDebugAxisSpec", + "attention_debug_probe_metadata", + "attention_debug_taxonomy", +] diff --git a/rl_engine/kernels/ops/pytorch/ffn/__init__.py b/rl_engine/kernels/ops/pytorch/ffn/__init__.py new file mode 100644 index 00000000..19a54d23 --- /dev/null +++ b/rl_engine/kernels/ops/pytorch/ffn/__init__.py @@ -0,0 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from .ffn import BACKEND_ID, Qwen3FFNOp, qwen3_ffn + +__all__ = ["BACKEND_ID", "Qwen3FFNOp", "qwen3_ffn"] diff --git a/rl_engine/kernels/ops/pytorch/ffn/ffn.py b/rl_engine/kernels/ops/pytorch/ffn/ffn.py new file mode 100644 index 00000000..c2507a44 --- /dev/null +++ b/rl_engine/kernels/ops/pytorch/ffn/ffn.py @@ -0,0 +1,477 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Bias-free gated FFN assembled from deterministic CUDA kernels.""" + +from __future__ import annotations + +from typing import Any + +import torch +from torch import Tensor + +from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE + +QWEN3_8B_HIDDEN_SIZE = 4096 +QWEN3_8B_INTERMEDIATE_SIZE = 12288 +BACKEND_ID = "rlkernel.ffn.qwen3.deterministic.v1" + +_DET_GEMM_SYMBOLS = ( + "det_gemm_fwd", + "det_gemm_db", +) +_SWIGLU_SYMBOLS = ( + "swiglu_forward", + "swiglu_backward", +) +_REQUIRED_SYMBOLS = _DET_GEMM_SYMBOLS + _SWIGLU_SYMBOLS +_COLLECTIVE_MIN_CAPACITY_BYTES = 64 * 1024 * 1024 +_COLLECTIVES: dict[tuple[int, int, int, int], Any] = {} + + +def _require_ffn_kernels(*, disable_split_k: bool) -> None: + required = _REQUIRED_SYMBOLS if disable_split_k else _SWIGLU_SYMBOLS + missing = [name for name in required if not hasattr(_C, name)] + if not _EXT_AVAILABLE or _C is None or missing: + suffix = f" Missing symbols: {', '.join(missing)}." if missing else "" + needed = ( + "compiled deterministic GEMM and SwiGLU CUDA kernels" + if disable_split_k + else "compiled SwiGLU CUDA kernels" + ) + raise RuntimeError(f"qwen3_ffn requires the {needed}.{suffix}") + + +def _gemm_fwd(a: Tensor, b: Tensor, *, disable_split_k: bool) -> Tensor: + if disable_split_k: + return _C.det_gemm_fwd(a, b) + # cuBLASLt / CUTLASS: may use split-K. Detach so Autograd.Function owns backward. + with torch.no_grad(): + return torch.matmul(a, b) + + +def _gemm_db(a: Tensor, grad_output: Tensor, *, disable_split_k: bool) -> Tensor: + if disable_split_k: + return _C.det_gemm_db(a, grad_output) + with torch.no_grad(): + return torch.matmul(a.t().contiguous(), grad_output) + + +def _require_parallel_group(group: Any, name: str): + if group is None: + return None + + import torch.distributed as dist + + if not dist.is_available(): + raise RuntimeError(f"{name}-parallel FFN requires torch.distributed.") + if not dist.is_initialized(): + raise RuntimeError(f"{name}-parallel FFN requires an initialized process group.") + if dist.get_world_size(group=group) <= 1: + raise ValueError(f"{name}_group must contain at least two ranks.") + return dist + + +def _validate_ffn_inputs( + rmsnorm_output: Tensor, + gate_weight: Tensor, + up_weight: Tensor, + down_weight: Tensor, +) -> None: + tensors = { + "rmsnorm_output": rmsnorm_output, + "gate_weight": gate_weight, + "up_weight": up_weight, + "down_weight": down_weight, + } + for name, tensor in tensors.items(): + if not isinstance(tensor, Tensor): + raise TypeError(f"{name} must be a torch.Tensor, got {type(tensor)!r}.") + + if rmsnorm_output.dim() < 1: + raise ValueError("rmsnorm_output must have at least one dimension.") + if rmsnorm_output.numel() == 0: + raise ValueError("rmsnorm_output must contain at least one token.") + for name, weight in ( + ("gate_weight", gate_weight), + ("up_weight", up_weight), + ("down_weight", down_weight), + ): + if weight.dim() != 2: + raise ValueError(f"{name} must be 2-D, got shape {tuple(weight.shape)}.") + + hidden_size = rmsnorm_output.size(-1) + intermediate_size = gate_weight.size(0) + expected_shapes = { + "gate_weight": (intermediate_size, hidden_size), + "up_weight": (intermediate_size, hidden_size), + "down_weight": (hidden_size, intermediate_size), + } + for name, expected in expected_shapes.items(): + actual = tuple(tensors[name].shape) + if actual != expected: + raise ValueError(f"{name} must have shape {expected}, got {actual}.") + + for name, tensor in tensors.items(): + if tensor.dtype != torch.bfloat16: + raise TypeError(f"{name} must have dtype bfloat16, got {tensor.dtype}.") + if not tensor.is_cuda: + raise RuntimeError(f"{name} must be on a CUDA device, got '{tensor.device}'.") + if tensor.device != rmsnorm_output.device: + raise RuntimeError( + f"all FFN inputs must be on {rmsnorm_output.device}, " + f"got {name} on {tensor.device}." + ) + + +def _collective_for_group(group: Any, *, min_size_bytes: int): + if group is None: + return None + + import torch.distributed as dist + + from rl_engine.distributed import DeterministicCollective + + rank = dist.get_rank(group=group) + world_size = dist.get_world_size(group=group) + device_index = torch.cuda.current_device() + key = (id(group), rank, world_size, device_index) + cached = _COLLECTIVES.get(key) + if cached is not None and cached.max_size_bytes >= min_size_bytes: + return cached + if cached is not None: + cached.close() + + collective = DeterministicCollective( + group=group, + max_size_bytes=max(_COLLECTIVE_MIN_CAPACITY_BYTES, min_size_bytes), + ) + _COLLECTIVES[key] = collective + return collective + + +def _all_gather_tokens(tensor: Tensor, collective: Any) -> Tensor: + return collective.all_gather(tensor.contiguous()) + + +def _reduce_scatter_tokens(tensor: Tensor, collective: Any) -> Tensor: + world_size = collective.world_size + if tensor.size(0) % world_size != 0: + raise ValueError( + "the gathered token count must be divisible by the tensor-parallel " + f"world size, got {tensor.size(0)} and {world_size}." + ) + return collective.reduce_scatter(tensor.contiguous()) + + +def _all_reduce_inplace(tensor: Tensor, collective: Any) -> Tensor: + return collective.all_reduce(tensor, out=tensor) + + +class _DeterministicFFNFunction(torch.autograd.Function): + @staticmethod + def forward( + ctx, + rmsnorm_output: Tensor, + gate_weight: Tensor, + up_weight: Tensor, + down_weight: Tensor, + tp_group: Any, + cp_group: Any, + sequence_parallel: bool, + disable_split_k: bool, + ) -> Tensor: + tp_dist = _require_parallel_group(tp_group, "tensor") + _require_parallel_group(cp_group, "context") + if sequence_parallel and tp_dist is None: + raise ValueError("sequence_parallel requires a tensor-parallel group.") + + input_shape = rmsnorm_output.shape + rmsnorm_output_2d = rmsnorm_output.reshape(-1, input_shape[-1]).contiguous() + tp_world = tp_dist.get_world_size(group=tp_group) if tp_dist is not None else 1 + gemm_tokens = rmsnorm_output_2d.size(0) * (tp_world if sequence_parallel else 1) + element_size = rmsnorm_output_2d.element_size() + min_size_bytes = max( + gemm_tokens * rmsnorm_output_2d.size(1) * element_size, + gemm_tokens * gate_weight.size(0) * element_size, + gate_weight.numel() * element_size, + up_weight.numel() * element_size, + down_weight.numel() * element_size, + ) + # Create TP before CP so every rank follows the same group order. + tp_collective = _collective_for_group(tp_group, min_size_bytes=min_size_bytes) + cp_collective = _collective_for_group(cp_group, min_size_bytes=min_size_bytes) + + if sequence_parallel: + rmsnorm_output_2d = _all_gather_tokens(rmsnorm_output_2d, tp_collective) + + # The model stores projection weights as [out, in]; GEMM consumes [K, N]. + gate = _gemm_fwd( + rmsnorm_output_2d, + gate_weight.t().contiguous(), + disable_split_k=disable_split_k, + ) + up = _gemm_fwd( + rmsnorm_output_2d, + up_weight.t().contiguous(), + disable_split_k=disable_split_k, + ) + activated = _C.swiglu_forward(gate, up) + output = _gemm_fwd(activated, down_weight.t().contiguous(), disable_split_k=disable_split_k) + + if sequence_parallel: + output = _reduce_scatter_tokens(output, tp_collective) + elif tp_collective is not None: + output = _all_reduce_inplace(output, tp_collective) + + ctx.save_for_backward( + rmsnorm_output_2d, + gate, + up, + activated, + gate_weight, + up_weight, + down_weight, + ) + ctx.input_shape = input_shape + ctx.tp_collective = tp_collective + ctx.cp_collective = cp_collective + ctx.sequence_parallel = sequence_parallel + ctx.disable_split_k = disable_split_k + return output.reshape(*input_shape[:-1], output.size(-1)) + + @staticmethod + def backward(ctx, grad_output: Tensor): + ( + rmsnorm_output, + gate, + up, + activated, + gate_weight, + up_weight, + down_weight, + ) = ctx.saved_tensors + tp_collective = ctx.tp_collective + cp_collective = ctx.cp_collective + disable_split_k = ctx.disable_split_k + grad_output = grad_output.reshape(-1, grad_output.size(-1)).contiguous() + if ctx.sequence_parallel: + grad_output = _all_gather_tokens(grad_output, tp_collective) + + # Down weight gradients must see every CP token so gemm_db's K-tree + # matches CP=1. Local dW + AllReduce is a different parenthesization + # whenever T is not a complete mid-split tree of 32-wide leaves. + if cp_collective is not None: + activated_full = _all_gather_tokens(activated, cp_collective) + grad_output_full = _all_gather_tokens(grad_output, cp_collective) + grad_down_weight = ( + _gemm_db(activated_full, grad_output_full, disable_split_k=disable_split_k) + .t() + .contiguous() + ) + else: + grad_down_weight = ( + _gemm_db(activated, grad_output, disable_split_k=disable_split_k).t().contiguous() + ) + + # Down input-gradient shards concatenate across TP; no TP reduction. + grad_activated = _gemm_fwd(grad_output, down_weight, disable_split_k=disable_split_k) + grad_gate, grad_up = _C.swiglu_backward(grad_activated, gate, up) + + if cp_collective is not None: + rmsnorm_full = _all_gather_tokens(rmsnorm_output, cp_collective) + grad_gate_full = _all_gather_tokens(grad_gate, cp_collective) + grad_up_full = _all_gather_tokens(grad_up, cp_collective) + grad_gate_weight = ( + _gemm_db(rmsnorm_full, grad_gate_full, disable_split_k=disable_split_k) + .t() + .contiguous() + ) + grad_up_weight = ( + _gemm_db(rmsnorm_full, grad_up_full, disable_split_k=disable_split_k) + .t() + .contiguous() + ) + else: + grad_gate_weight = ( + _gemm_db(rmsnorm_output, grad_gate, disable_split_k=disable_split_k) + .t() + .contiguous() + ) + grad_up_weight = ( + _gemm_db(rmsnorm_output, grad_up, disable_split_k=disable_split_k).t().contiguous() + ) + + # Gate/Up input gradients reduce across TP, then add locally. + grad_rmsnorm_from_gate = _gemm_fwd(grad_gate, gate_weight, disable_split_k=disable_split_k) + if ctx.sequence_parallel: + grad_rmsnorm_from_gate = _reduce_scatter_tokens( + grad_rmsnorm_from_gate, + tp_collective, + ) + elif tp_collective is not None: + grad_rmsnorm_from_gate = _all_reduce_inplace( + grad_rmsnorm_from_gate, + tp_collective, + ) + + grad_rmsnorm_from_up = _gemm_fwd(grad_up, up_weight, disable_split_k=disable_split_k) + if ctx.sequence_parallel: + grad_rmsnorm_from_up = _reduce_scatter_tokens( + grad_rmsnorm_from_up, + tp_collective, + ) + elif tp_collective is not None: + grad_rmsnorm_from_up = _all_reduce_inplace( + grad_rmsnorm_from_up, + tp_collective, + ) + + grad_rmsnorm_output = grad_rmsnorm_from_gate.add_(grad_rmsnorm_from_up) + return ( + grad_rmsnorm_output.reshape(ctx.input_shape), + grad_gate_weight, + grad_up_weight, + grad_down_weight, + None, + None, + None, + None, + ) + + +def qwen3_ffn( + rmsnorm_output: Tensor, + gate_weight: Tensor, + up_weight: Tensor, + down_weight: Tensor, + *, + tp_group: Any = None, + cp_group: Any = None, + sequence_parallel: bool = False, + deterministic: bool | None = None, + disable_split_k: bool | None = None, +) -> Tensor: + """Apply a bias-free SiLU-gated FFN with deterministic backward kernels. + + Args: + rmsnorm_output: RMSNorm output, shape ``[..., H]``. + gate_weight: Gate projection weight in ``[out, in]`` layout, shape + ``[I_local, H]``. + up_weight: Up projection weight in ``[out, in]`` layout, shape + ``[I_local, H]``. + down_weight: Down projection weight in ``[out, in]`` layout, shape + ``[H, I_local]``. + tp_group: Optional tensor-parallel process group. Gate and Up are + column-parallel; Down is row-parallel. Reductions use the + deterministic fixed-tree collectives rather than NCCL. + cp_group: Optional context-parallel process group. Each rank owns + different token rows and the same local weight shards. Weight + gradients AllGather tokens along CP and run the full-token + ``det_gemm_db`` so they match CP=1 bitwise. + sequence_parallel: Whether ``rmsnorm_output`` and the returned output + are sharded on the flattened token dimension across ``tp_group``. + Token gather/scatter use the deterministic AllGather and + ReduceScatter. + deterministic: Select the RL-Kernel fixed-reduction GEMM when True + (default), or the production ``torch.matmul`` GEMM when False. + disable_split_k: Compatibility alias for ``deterministic``. New code + should use ``deterministic`` because Split-K is only one possible + implementation detail of the production GEMM. + + Returns: + FFN output with shape ``[..., H]``. + """ + if not isinstance(sequence_parallel, bool): + raise TypeError("sequence_parallel must be a bool.") + deterministic = _resolve_deterministic_mode(deterministic, disable_split_k) + _validate_ffn_inputs(rmsnorm_output, gate_weight, up_weight, down_weight) + _require_ffn_kernels(disable_split_k=deterministic) + return _DeterministicFFNFunction.apply( + rmsnorm_output, + gate_weight, + up_weight, + down_weight, + tp_group, + cp_group, + sequence_parallel, + deterministic, + ) + + +def _resolve_deterministic_mode( + deterministic: bool | None, + disable_split_k: bool | None, +) -> bool: + if deterministic is not None and not isinstance(deterministic, bool): + raise TypeError("deterministic must be a bool or None.") + if disable_split_k is not None and not isinstance(disable_split_k, bool): + raise TypeError("disable_split_k must be a bool or None.") + if ( + deterministic is not None + and disable_split_k is not None + and deterministic != disable_split_k + ): + raise ValueError("deterministic and disable_split_k select conflicting FFN backends.") + if deterministic is not None: + return deterministic + if disable_split_k is not None: + return disable_split_k + return True + + +class Qwen3FFNOp: + """Instantiable Qwen3 FFN wrapper for semantic operator dispatch.""" + + op_class = "ffn" + is_batch_invariant = True + backend_id = BACKEND_ID + + def __call__( + self, + rmsnorm_output: Tensor, + gate_weight: Tensor, + up_weight: Tensor, + down_weight: Tensor, + *, + tp_group: Any = None, + cp_group: Any = None, + sequence_parallel: bool = False, + deterministic: bool | None = None, + disable_split_k: bool | None = None, + ) -> Tensor: + return self.apply( + rmsnorm_output, + gate_weight, + up_weight, + down_weight, + tp_group=tp_group, + cp_group=cp_group, + sequence_parallel=sequence_parallel, + deterministic=deterministic, + disable_split_k=disable_split_k, + ) + + def apply( + self, + rmsnorm_output: Tensor, + gate_weight: Tensor, + up_weight: Tensor, + down_weight: Tensor, + *, + tp_group: Any = None, + cp_group: Any = None, + sequence_parallel: bool = False, + deterministic: bool | None = None, + disable_split_k: bool | None = None, + ) -> Tensor: + return qwen3_ffn( + rmsnorm_output, + gate_weight, + up_weight, + down_weight, + tp_group=tp_group, + cp_group=cp_group, + sequence_parallel=sequence_parallel, + deterministic=deterministic, + disable_split_k=disable_split_k, + ) diff --git a/rl_engine/kernels/ops/pytorch/loss/vocab_parallel_logp.py b/rl_engine/kernels/ops/pytorch/loss/vocab_parallel_logp.py new file mode 100644 index 00000000..3745382a --- /dev/null +++ b/rl_engine/kernels/ops/pytorch/loss/vocab_parallel_logp.py @@ -0,0 +1,590 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Deterministic vocab-parallel TP selected-token logprob reference (issue #241 PR3). + +Implements the WS2 contract in ``rl_engine.kernels.logprob_contract`` with a +TP-independent vocab tile decomposition: the padded vocabulary is split into +``num_vocab_tiles`` fixed tiles, every tile's fp32 ``(max, sumexp)`` partial is +computed from a contiguous ``[n, tile]`` tensor, all tile partials travel by +all-gather (transport only), and every rank merges them in global tile-index +order over a fixed ``[n, num_vocab_tiles]`` shape. The TP degree only decides +which rank computes which tiles and never changes any floating-point grouping, +so outputs and gradients are bitwise-identical across TP degrees +(``DeterminismScope.CROSS_TP_BITWISE``) as long as ``num_vocab_tiles`` is held +fixed. A fixed per-shard merge order alone cannot provide this property: +shard boundaries would regroup the combines differently at each degree. + +Consequences of the tile structure: + +- ``num_vocab_tiles`` is part of the numerical identity. It must be pinned + across ranks (enforced by the preflight) and across the TP degrees being + compared; it is never derived from the shard layout. +- Every shard boundary must be tile-aligned; misalignment fails loudly. +- At TP=1 the result matches the WS1 ``NativeBatchInvariantLogpOp`` only + within the #108 logprob tolerance, not bitwise — the WS1 op reduces the + whole ``[n, V]`` row at once, which groups the sums differently. + +Preconditions: logits over the real vocabulary must be finite. A row whose +real-vocab logits are all ``-inf`` has no finite logsumexp; with +``validate=True`` such a row fails loudly if it is active. + +The selected logprob is zero-filled at inactive rows (``MaskSpec.active_mask`` +is the sole authority; with validation enabled an active row can never legally +hold ``ignore_index``). The vocab-domain LSE is returned for every row and is +differentiable everywhere, including inactive rows. + +``deterministic=False`` trades the guarantee for speed: each rank reduces its +whole shard in one pass and the per-shard partials are merged in shard order, +so the floating-point grouping changes with the TP degree and nothing is +promised about reproducibility. ``num_vocab_tiles`` is ignored (no tile +alignment is required), and a contract declaring +``determinism_scope=cross_tp_bitwise`` is rejected loudly — the fast path +cannot honor it. Batch invariance is unaffected: rows never mix either way. +""" + +from __future__ import annotations + +from typing import Any + +import torch + +from rl_engine.kernels.logprob_contract import ( + DeterminismScope, + LogprobContract, + LogprobContractError, + LogprobDType, +) + +BACKEND_ID = "pytorch-vocab-parallel-logp-ws2" +DEFAULT_NUM_VOCAB_TILES = 64 + +_TORCH_TO_CONTRACT_DTYPE = { + torch.bfloat16: LogprobDType.BF16, + torch.float16: LogprobDType.FP16, + torch.float32: LogprobDType.FP32, +} + + +def _require_distributed_initialized(): + import torch.distributed as dist + + if not dist.is_available(): + raise LogprobContractError("vocab-parallel logprob requires torch.distributed.") + if not dist.is_initialized(): + raise LogprobContractError( + "vocab-parallel logprob requires an initialized process group when " + "the contract declares tp_world_size > 1." + ) + return dist + + +def _tile_size(contract: LogprobContract, num_vocab_tiles: int) -> int: + if isinstance(num_vocab_tiles, bool) or not isinstance(num_vocab_tiles, int): + raise LogprobContractError( + f"num_vocab_tiles must be a positive integer; got {num_vocab_tiles!r}" + ) + if num_vocab_tiles <= 0: + raise LogprobContractError( + f"num_vocab_tiles must be a positive integer; got {num_vocab_tiles}" + ) + padded = contract.sharding.padded_vocab_size + if padded % num_vocab_tiles != 0: + raise LogprobContractError( + f"num_vocab_tiles={num_vocab_tiles} must divide " f"padded_vocab_size={padded} exactly" + ) + tile = padded // num_vocab_tiles + for rank, (start, end) in enumerate(contract.sharding.vocab_shard_bounds): + if start % tile != 0 or end % tile != 0: + raise LogprobContractError( + f"vocab_shard_bounds[{rank}]=[{start}, {end}) is not aligned to the " + f"vocab tile size {tile} (num_vocab_tiles={num_vocab_tiles}); " + "cross-TP bitwise determinism requires tile-aligned shard bounds" + ) + return tile + + +def _validate_invocation( + local_logits: torch.Tensor, + target_ids: torch.Tensor, + contract: LogprobContract, + tp_group: Any, +) -> None: + if not isinstance(contract, LogprobContract): + raise LogprobContractError("contract must be a LogprobContract") + if local_logits.dim() != 2: + raise LogprobContractError( + f"local_logits must be 2-D [num_tokens, local_vocab]; got {local_logits.dim()}-D" + ) + if target_ids.dim() != 1 or target_ids.shape[0] != local_logits.shape[0]: + raise LogprobContractError( + f"target_ids must be 1-D with one entry per token; got shape " + f"{tuple(target_ids.shape)} for {local_logits.shape[0]} tokens" + ) + sharding = contract.sharding + if local_logits.shape[1] != sharding.local_vocab_size: + raise LogprobContractError( + f"local_logits has {local_logits.shape[1]} vocab columns but the contract " + f"declares local shard [{sharding.local_vocab_start}, " + f"{sharding.local_vocab_end}) of size {sharding.local_vocab_size}" + ) + if local_logits.shape[0] != contract.mask.num_tokens: + raise LogprobContractError( + f"local_logits has {local_logits.shape[0]} tokens but MaskSpec declares " + f"num_tokens={contract.mask.num_tokens}" + ) + declared = _TORCH_TO_CONTRACT_DTYPE.get(local_logits.dtype) + if declared is not contract.dtype: + raise LogprobContractError( + f"local_logits dtype {local_logits.dtype} does not match the contract " + f"dtype {contract.dtype.value}" + ) + if sharding.tp_world_size > 1: + dist = _require_distributed_initialized() + group_rank = dist.get_rank(group=tp_group) + group_world = dist.get_world_size(group=tp_group) + if group_world != sharding.tp_world_size: + raise LogprobContractError( + f"tp_group world size {group_world} does not match the contract " + f"tp_world_size={sharding.tp_world_size}; pass the TP subgroup, " + "not the global group" + ) + if group_rank != sharding.tp_rank: + raise LogprobContractError( + f"tp_group rank {group_rank} does not match the contract " + f"tp_rank={sharding.tp_rank}" + ) + + +def _validate_active_targets( + target_1d: torch.Tensor, active_mask: torch.Tensor, real_vocab_size: int +) -> None: + bad = active_mask & ((target_1d < 0) | (target_1d >= real_vocab_size)) + if bool(bad.any().item()): + bad_values = target_1d[bad] + raise LogprobContractError( + "active target_ids must lie in the real vocabulary " + f"[0, {real_vocab_size}); got values in " + f"[{int(bad_values.min().item())}, {int(bad_values.max().item())}] " + "on active rows" + ) + + +def _preflight_cross_rank_agreement( + contract: LogprobContract, tp_group: Any, num_vocab_tiles: int, deterministic: bool +) -> None: + """All-gather (fingerprint, backend id, tile count, mode) and abort on mismatch. + + The tile count travels as ``None`` when ``deterministic=False``: the fast + path never uses it, so ranks must not fail preflight over an irrelevant + value — but they must never disagree on the mode itself, or they would + issue different collectives. + """ + + dist = _require_distributed_initialized() + payload = ( + contract.cross_rank_fingerprint(), + BACKEND_ID, + int(num_vocab_tiles) if deterministic else None, + bool(deterministic), + ) + world = dist.get_world_size(group=tp_group) + gathered: list[Any] = [None] * world + dist.all_gather_object(gathered, payload, group=tp_group) + mismatched = [(rank, other) for rank, other in enumerate(gathered) if other != payload] + if mismatched: + rank, other = mismatched[0] + raise LogprobContractError( + "cross-rank preflight failed: rank " + f"{contract.sharding.tp_rank} has {payload} but rank {rank} has {other}; " + "all TP ranks must agree on the contract fingerprint, backend id, " + "num_vocab_tiles, and deterministic mode before any collective" + ) + + +def _local_tile_stats(z_masked: torch.Tensor, tile: int) -> tuple[torch.Tensor, torch.Tensor]: + """fp32 per-tile ``(max, sumexp)`` partials for this rank's shard. + + Each tile is reduced as a contiguous ``[n, tile]`` tensor so the reduction + shape and layout are identical no matter which rank computes the tile or + what the local shard size is. An all-``-inf`` (padding-only) tile yields + the identity partial ``(-inf, 0)`` without evaluating ``exp(-inf - (-inf))``. + """ + + n, local_vocab = z_masked.shape + m_parts: list[torch.Tensor] = [] + s_parts: list[torch.Tensor] = [] + for tile_index in range(local_vocab // tile): + block = z_masked[:, tile_index * tile : (tile_index + 1) * tile].contiguous() + m_t = block.max(dim=-1).values + finite = m_t > float("-inf") + m_safe = torch.where(finite, m_t, torch.zeros_like(m_t)) + s_t = (block - m_safe.unsqueeze(-1)).exp().sum(dim=-1) + s_t = torch.where(finite, s_t, torch.zeros_like(s_t)) + m_parts.append(m_t) + s_parts.append(s_t) + return torch.stack(m_parts, dim=1), torch.stack(s_parts, dim=1) + + +def _gather_tile_stats( + local_m: torch.Tensor, + local_s: torch.Tensor, + contract: LogprobContract, + tp_group: Any, + tile_counts: list[int], +) -> tuple[torch.Tensor, torch.Tensor]: + """Assemble every rank's partials in global shard order. + + ``tile_counts`` holds each rank's partial count: one per tile on the + deterministic path, exactly one per shard on the fast path. + """ + + sharding = contract.sharding + if sharding.tp_world_size == 1: + return local_m.contiguous(), local_s.contiguous() + + dist = _require_distributed_initialized() + n = local_m.shape[0] + max_tiles = max(tile_counts) + packed = local_m.new_zeros((n, max_tiles, 2)) + packed[:, : local_m.shape[1], 0] = local_m + packed[:, : local_s.shape[1], 1] = local_s + packed = packed.contiguous() + gathered = [torch.empty_like(packed) for _ in range(sharding.tp_world_size)] + dist.all_gather(gathered, packed, group=tp_group) + + m_parts = [gathered[rank][:, : tile_counts[rank], 0] for rank in range(len(tile_counts))] + s_parts = [gathered[rank][:, : tile_counts[rank], 1] for rank in range(len(tile_counts))] + return torch.cat(m_parts, dim=1).contiguous(), torch.cat(s_parts, dim=1).contiguous() + + +def _gather_target_logit( + z_masked: torch.Tensor, + safe_target: torch.Tensor, + contract: LogprobContract, + tp_group: Any, +) -> torch.Tensor: + """Exact selected-target logit via a select-by-owner copy.""" + + sharding = contract.sharding + n = z_masked.shape[0] + start = sharding.local_vocab_start + local_vocab = sharding.local_vocab_size + local_idx = (safe_target - start).clamp(0, max(local_vocab - 1, 0)) + owns = (safe_target >= start) & (safe_target < sharding.local_vocab_end) + rows = torch.arange(n, device=z_masked.device) + local_contrib = torch.where( + owns, z_masked[rows, local_idx], torch.zeros_like(safe_target, dtype=z_masked.dtype) + ).contiguous() + + if sharding.tp_world_size == 1: + stacked = local_contrib.unsqueeze(0) + else: + dist = _require_distributed_initialized() + gathered = [torch.empty_like(local_contrib) for _ in range(sharding.tp_world_size)] + dist.all_gather(gathered, local_contrib, group=tp_group) + stacked = torch.stack(gathered, dim=0) + + starts = torch.tensor( + [bound_start for bound_start, _ in sharding.vocab_shard_bounds], + device=safe_target.device, + dtype=torch.long, + ) + owner = torch.bucketize(safe_target, starts, right=True) - 1 + return stacked[owner, rows] + + +def _gather_entropy_partials( + local_entropy: torch.Tensor, + contract: LogprobContract, + tp_group: Any, +) -> torch.Tensor: + """Merge per-shard entropy contributions in TP-rank order. + + Entropy is not part of the WS2 selected-logprob contract, but the Vime + adapter needs it for the existing training loss surface. The collective + transports independent per-shard contributions; every TP rank performs + the same explicit rank-ordered sum afterwards. + """ + + if contract.sharding.tp_world_size == 1: + return local_entropy + + dist = _require_distributed_initialized() + gathered = [torch.empty_like(local_entropy) for _ in range(contract.sharding.tp_world_size)] + dist.all_gather(gathered, local_entropy.contiguous(), group=tp_group) + merged = gathered[0].clone() + for partial in gathered[1:]: + merged = merged + partial + return merged + + +def _merge_tile_partials(m_all: torch.Tensor, s_all: torch.Tensor) -> torch.Tensor: + """Fixed-order (max, sumexp) merge over [n, num_vocab_tiles].""" + + M = m_all.max(dim=1).values + finite = M > float("-inf") + M_safe = torch.where(finite, M, torch.zeros_like(M)) + terms = s_all * (m_all - M_safe.unsqueeze(1)).exp() + S = terms.sum(dim=1) + return M + S.log() + + +class _VocabParallelLogprobFunction(torch.autograd.Function): + @staticmethod + def forward( + ctx, + local_logits, + target_1d, + active_mask, + contract, + tp_group, + tile, + with_entropy, + with_entropy_grad, + ): + z_masked = local_logits.float() + sharding = contract.sharding + global_ids = torch.arange( + sharding.local_vocab_start, sharding.local_vocab_end, device=z_masked.device + ) + padding_cols = global_ids >= sharding.real_vocab_size + if bool(padding_cols.any()): + z_masked = z_masked.masked_fill(padding_cols.unsqueeze(0), float("-inf")) + + safe_target = torch.where(active_mask, target_1d, torch.zeros_like(target_1d)) + + if tile is not None: + local_m, local_s = _local_tile_stats(z_masked, tile) + tile_counts = [(end - start) // tile for start, end in sharding.vocab_shard_bounds] + else: + # Fast path: one (max, sumexp) partial over the whole shard. The + # grouping now depends on the shard layout, so no cross-TP claim. + local_m, local_s = _local_tile_stats(z_masked, z_masked.shape[1]) + tile_counts = [1] * sharding.tp_world_size + m_all, s_all = _gather_tile_stats(local_m, local_s, contract, tp_group, tile_counts) + target_logit = _gather_target_logit(z_masked, safe_target, contract, tp_group) + lse = _merge_tile_partials(m_all, s_all) + + selected_logp = torch.where(active_mask, target_logit - lse, torch.zeros_like(lse)) + + if with_entropy: + finite_row = torch.isfinite(lse) + lse_safe = torch.where(finite_row, lse, torch.zeros_like(lse)) + probabilities = (z_masked - lse_safe.unsqueeze(1)).exp() + probabilities = torch.where( + finite_row.unsqueeze(1), probabilities, torch.zeros_like(probabilities) + ) + finite_logits = torch.isfinite(z_masked) + log_gap = torch.where( + finite_logits, + lse_safe.unsqueeze(1) - z_masked, + torch.zeros_like(z_masked), + ) + local_entropy = (probabilities * log_gap).sum(dim=1) + entropy = _gather_entropy_partials(local_entropy, contract, tp_group) + else: + entropy = local_logits.new_empty((0,), dtype=torch.float32) + + ctx.save_for_backward(z_masked, lse, safe_target, active_mask, padding_cols, entropy) + ctx.local_vocab_start = sharding.local_vocab_start + ctx.local_vocab_size = sharding.local_vocab_size + ctx.input_dtype = local_logits.dtype + ctx.with_entropy_grad = bool(with_entropy and with_entropy_grad) + ctx.set_materialize_grads(False) + if with_entropy and not ctx.with_entropy_grad: + ctx.mark_non_differentiable(entropy) + return selected_logp, lse, entropy + + @staticmethod + def backward(ctx, grad_logp, grad_lse, grad_entropy): + if not ctx.needs_input_grad[0] or ( + grad_logp is None and grad_lse is None and grad_entropy is None + ): + return None, None, None, None, None, None, None, None + + z_masked, lse, safe_target, active_mask, padding_cols, entropy = ctx.saved_tensors + n, local_vocab = z_masked.shape + finite_row = torch.isfinite(lse) + lse_safe = torch.where(finite_row, lse, torch.zeros_like(lse)) + p = (z_masked - lse_safe.unsqueeze(1)).exp() + p = torch.where(finite_row.unsqueeze(1), p, torch.zeros_like(p)) + + grad = torch.zeros_like(z_masked) + if grad_logp is not None: + local_idx = (safe_target - ctx.local_vocab_start).clamp(0, max(local_vocab - 1, 0)) + owns = (safe_target >= ctx.local_vocab_start) & ( + safe_target < ctx.local_vocab_start + local_vocab + ) + onehot = torch.zeros_like(z_masked) + hit = owns & active_mask + rows = torch.arange(n, device=z_masked.device)[hit] + onehot[rows, local_idx[hit]] = 1.0 + g_logp = torch.where(active_mask, grad_logp, torch.zeros_like(grad_logp)) + grad = grad + g_logp.unsqueeze(1) * (onehot - p) + if grad_lse is not None: + grad = grad + grad_lse.unsqueeze(1) * p + if ctx.with_entropy_grad and grad_entropy is not None: + entropy_input = lse_safe.unsqueeze(1) - z_masked - entropy.unsqueeze(1) + entropy_input = torch.where( + torch.isfinite(z_masked), entropy_input, torch.zeros_like(entropy_input) + ) + grad = grad + grad_entropy.unsqueeze(1) * p * entropy_input + if bool(padding_cols.any()): + grad = grad.masked_fill(padding_cols.unsqueeze(0), 0.0) + return grad.to(ctx.input_dtype), None, None, None, None, None, None, None + + +class VocabParallelLogprobOp: + """Vocab-parallel selected-token logprob (WS2 reference). + + Deterministic by default (cross-TP bitwise, tile-ordered merge); + ``deterministic=False`` selects the faster whole-shard reduction with no + reproducibility guarantee. + """ + + op_class = "logprob" + is_batch_invariant = True + backend_id = BACKEND_ID + + def __init__(self) -> None: + pass + + def __call__( + self, + local_logits: torch.Tensor, + target_ids: torch.Tensor, + *, + contract: LogprobContract, + tp_group: Any = None, + num_vocab_tiles: int = DEFAULT_NUM_VOCAB_TILES, + validate: bool = True, + deterministic: bool = True, + ) -> tuple[torch.Tensor, torch.Tensor]: + return self.apply( + local_logits, + target_ids, + contract=contract, + tp_group=tp_group, + num_vocab_tiles=num_vocab_tiles, + validate=validate, + deterministic=deterministic, + ) + + def apply( + self, + local_logits: torch.Tensor, + target_ids: torch.Tensor, + *, + contract: LogprobContract, + tp_group: Any = None, + num_vocab_tiles: int = DEFAULT_NUM_VOCAB_TILES, + validate: bool = True, + deterministic: bool = True, + ) -> tuple[torch.Tensor, torch.Tensor]: + selected_logp, lse, _ = self._apply( + local_logits, + target_ids, + contract=contract, + tp_group=tp_group, + num_vocab_tiles=num_vocab_tiles, + validate=validate, + deterministic=deterministic, + with_entropy=False, + with_entropy_grad=False, + ) + return selected_logp, lse + + def apply_with_entropy( + self, + local_logits: torch.Tensor, + target_ids: torch.Tensor, + *, + contract: LogprobContract, + tp_group: Any = None, + num_vocab_tiles: int = DEFAULT_NUM_VOCAB_TILES, + validate: bool = True, + with_entropy_grad: bool = True, + deterministic: bool = True, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Return selected logprob, vocabulary LSE, and full-vocabulary entropy. + + The method is intentionally separate from :meth:`apply` so the WS2 + selected-logprob surface remains unchanged for existing callers. + """ + + return self._apply( + local_logits, + target_ids, + contract=contract, + tp_group=tp_group, + num_vocab_tiles=num_vocab_tiles, + validate=validate, + deterministic=deterministic, + with_entropy=True, + with_entropy_grad=with_entropy_grad, + ) + + def _apply( + self, + local_logits: torch.Tensor, + target_ids: torch.Tensor, + *, + contract: LogprobContract, + tp_group: Any, + num_vocab_tiles: int, + validate: bool, + deterministic: bool, + with_entropy: bool, + with_entropy_grad: bool, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + if not isinstance(contract, LogprobContract): + raise LogprobContractError("contract must be a LogprobContract") + if not isinstance(deterministic, bool): + raise LogprobContractError(f"deterministic must be a bool; got {deterministic!r}") + if deterministic: + tile = _tile_size(contract, num_vocab_tiles) + elif contract.reduction.determinism_scope is DeterminismScope.CROSS_TP_BITWISE: + raise LogprobContractError( + "deterministic=False cannot honor determinism_scope=cross_tp_bitwise: " + "the fast path reduces each shard in one piece, so its floating-point " + "grouping changes with the TP degree; keep deterministic=True or relax " + "the contract to determinism_scope=fixed_topology" + ) + else: + tile = None + _validate_invocation(local_logits, target_ids, contract, tp_group) + + target_1d = target_ids.reshape(-1).to(device=local_logits.device, dtype=torch.long) + active_mask = torch.tensor( + contract.mask.active_mask, dtype=torch.bool, device=local_logits.device + ) + if validate: + _validate_active_targets(target_1d, active_mask, contract.sharding.real_vocab_size) + if contract.sharding.tp_world_size > 1: + _preflight_cross_rank_agreement(contract, tp_group, num_vocab_tiles, deterministic) + + selected_logp, lse, entropy = _VocabParallelLogprobFunction.apply( + local_logits, + target_1d, + active_mask, + contract, + tp_group, + tile, + with_entropy, + with_entropy_grad, + ) + + if validate and bool((~torch.isfinite(lse) & active_mask).any().item()): + raise LogprobContractError( + "non-finite logsumexp on an active row: logits over the real " + "vocabulary must be finite for every active token" + ) + return selected_logp, lse, entropy + + +__all__ = [ + "BACKEND_ID", + "DEFAULT_NUM_VOCAB_TILES", + "VocabParallelLogprobOp", +] diff --git a/rl_engine/kernels/registry.py b/rl_engine/kernels/registry.py index e7114827..8a6bb74c 100644 --- a/rl_engine/kernels/registry.py +++ b/rl_engine/kernels/registry.py @@ -1,6 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 RL-Kernel Contributors +from __future__ import annotations + import importlib import os from enum import Enum, EnumMeta @@ -8,6 +10,32 @@ import torch +from rl_engine.kernels.attention_contract import ( + AttentionBackendCapability, + AttentionContract, + AttentionContractError, + AttentionDispatchResult, + AttentionDType, + AttentionMode, + AttentionRole, +) +from rl_engine.kernels.logprob_contract import ( + IMPLEMENTATION_KINDS, + DeterminismScope, + LogprobBackendCapability, + LogprobContract, + LogprobContractError, + LogprobDispatchResult, + LogprobDType, + LogprobRole, + MaskMode, +) +from rl_engine.kernels.semantic_registry import ( + OperatorBackendDescriptor, + OperatorFallbackPolicy, + OperatorLifecycle, + SemanticOperatorCatalog, +) from rl_engine.platforms.device import device_ctx from rl_engine.utils.logger import logger @@ -18,9 +46,11 @@ class _KernelEnumMeta(EnumMeta): def __getitem__(cls, name: str): try: return super().__getitem__(name) - except KeyError as e: + except KeyError as exc: valid_ops = ", ".join(cls.__members__.keys()) - raise ValueError(f"Operator '{name}' not found. Supported backends: {valid_ops}") from e + raise ValueError( + f"Operator '{name}' not found. Supported backends: {valid_ops}" + ) from exc class OpBackend(Enum, metaclass=_KernelEnumMeta): @@ -74,8 +104,9 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta): CUDA_BATCH_INVARIANT_LOGP_SM90 = ( "rl_engine.kernels.ops.cuda.loss.batch_invariant_logp.BatchInvariantLogpSM90Op" ) - ASCEND_BATCH_INVARIANT_LOGP = ( - "rl_engine.kernels.ops.ascend.loss.batch_invariant_logp.BatchInvariantLogpAscendOp" + # Deterministic vocab-parallel TP logprob reference (WS2 #241 PR3) + PYTORCH_VOCAB_PARALLEL_LOGP = ( + "rl_engine.kernels.ops.pytorch.loss.vocab_parallel_logp.VocabParallelLogprobOp" ) # RMSNorm(pre-norm / QK-Norm) - pure Pytorch reference(ws1 ground-truth) @@ -106,6 +137,10 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta): PYTORCH_NATIVE_KV_CACHE_ATTN = ( "rl_engine.kernels.ops.pytorch.attention.kv_cache.NativeKVCacheAttnOp" ) + PYTORCH_CP_ATTENTION = ( + "rl_engine.kernels.ops.pytorch.attention.cp_attention." + "DeterministicCPAttentionReferenceOp" + ) # WS1 pure-PyTorch ground-truth linear ops PYTORCH_NATIVE_LM_HEAD = "rl_engine.kernels.ops.pytorch.linear.lm_head.NativeLMHeadOp" # WS1 pure-PyTorch ground-truth embedding ops @@ -114,6 +149,151 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta): CUDA_SM90_EMBEDDING = "rl_engine.kernels.ops.cuda.linear.embedding.SM90EmbeddingOp" +def _default_semantic_descriptors() -> tuple[OperatorBackendDescriptor, ...]: + return ( + OperatorBackendDescriptor( + semantic_op="selected_logprob", + backend_id="rlkernel.reference_logp", + supported_targets=frozenset({"rollout", "training"}), + supported_devices=frozenset({"cpu", "cuda", "rocm"}), + supported_dtypes=frozenset({"float32", "bfloat16", "float16"}), + supported_topologies={ + "rollout": { + "world_size": (1,), + "tensor_parallel_size": (1,), + "context_parallel_size": (1,), + }, + "training": { + "world_size": (1,), + "sharding": ("unsharded",), + }, + }, + determinism_or_alignment_properties={ + "algorithm": "pytorch.log_softmax_gather", + "batch_invariant": True, + "deterministic": True, + "strict_observable": True, + }, + lifecycle=OperatorLifecycle.ENGINE_CONSTRUCTION, + implementation_class_or_factory=( + "rl_engine.kernels.ops.pytorch.loss.logp.NativeLogpOp" + ), + fallback_policy=OperatorFallbackPolicy.ERROR, + version_or_build_fingerprint="NativeLogpOp-selected-logprob-v1", + ), + OperatorBackendDescriptor( + semantic_op="selected_logprob", + backend_id="native", + supported_targets=frozenset({"rollout", "training"}), + supported_devices=frozenset({"*"}), + supported_dtypes=frozenset({"*"}), + supported_topologies={"*": "*"}, + determinism_or_alignment_properties={ + "selection": "runtime_native", + "strict_observable": False, + }, + lifecycle=OperatorLifecycle.ENGINE_CONSTRUCTION, + implementation_class_or_factory=None, + fallback_policy=OperatorFallbackPolicy.RUNTIME_MANAGED, + version_or_build_fingerprint="runtime-native-unresolved-v1", + ), + OperatorBackendDescriptor( + semantic_op="selected_logprob", + backend_id="pytorch-vocab-parallel-logp-ws2", + supported_targets=frozenset({"rollout", "training"}), + supported_devices=frozenset({"cpu", "cuda", "rocm"}), + supported_dtypes=frozenset({"float32", "bfloat16", "float16"}), + supported_topologies={"*": "*"}, + determinism_or_alignment_properties={ + "algorithm": "fixed_global_vocab_tiles", + "deterministic": True, + "cross_tp_bitwise": True, + "exports_vocab_lse": True, + "strict_observable": True, + }, + lifecycle=OperatorLifecycle.DISTRIBUTED_CONTEXT, + implementation_class_or_factory=( + "rl_engine.kernels.ops.pytorch.loss.vocab_parallel_logp." "VocabParallelLogprobOp" + ), + fallback_policy=OperatorFallbackPolicy.ERROR, + version_or_build_fingerprint="VocabParallelLogprobOp-fixed-tiles-v1", + ), + OperatorBackendDescriptor( + semantic_op="attention", + backend_id="rlkernel.attention.deterministic.v1", + supported_targets=frozenset({"rollout", "training"}), + supported_devices=frozenset({"cpu", "cuda", "rocm"}), + supported_dtypes=frozenset({"float32", "bfloat16", "float16"}), + supported_topologies={"*": "*"}, + determinism_or_alignment_properties={ + "algorithm": "standard_softmax_attention", + "batch_invariant": True, + "deterministic": True, + "split_kv": "contract_bound", + "reduction_order": "global_block_index", + "strict_schedule": "single_batch_single_query_global_kv_blocks", + "strict_observable": True, + }, + lifecycle=OperatorLifecycle.ENGINE_CONSTRUCTION, + implementation_class_or_factory=( + "rl_engine.kernels.ops.pytorch.attention.ablation.AttentionAblationOp" + ), + fallback_policy=OperatorFallbackPolicy.ERROR, + version_or_build_fingerprint="AttentionAblationOp-bitwise-v2", + ), + OperatorBackendDescriptor( + semantic_op="attention", + backend_id="native", + supported_targets=frozenset({"rollout", "training"}), + supported_devices=frozenset({"*"}), + supported_dtypes=frozenset({"*"}), + supported_topologies={"*": "*"}, + determinism_or_alignment_properties={ + "selection": "runtime_native", + "strict_observable": False, + }, + lifecycle=OperatorLifecycle.ENGINE_CONSTRUCTION, + implementation_class_or_factory=None, + fallback_policy=OperatorFallbackPolicy.RUNTIME_MANAGED, + version_or_build_fingerprint="runtime-native-attention-unresolved-v1", + ), + OperatorBackendDescriptor( + semantic_op="ffn", + backend_id="rlkernel.ffn.qwen3.deterministic.v1", + supported_targets=frozenset({"rollout", "training"}), + supported_devices=frozenset({"cuda"}), + supported_dtypes=frozenset({"bfloat16"}), + supported_topologies={"*": "*"}, + determinism_or_alignment_properties={ + "algorithm": "qwen3_swiglu_fixed_reduction_gemm", + "batch_invariant": True, + "deterministic": True, + "strict_observable": True, + }, + lifecycle=OperatorLifecycle.DISTRIBUTED_CONTEXT, + implementation_class_or_factory=("rl_engine.kernels.ops.pytorch.ffn.ffn.Qwen3FFNOp"), + fallback_policy=OperatorFallbackPolicy.ERROR, + version_or_build_fingerprint="Qwen3FFNOp-fixed-reduction-v1", + ), + OperatorBackendDescriptor( + semantic_op="ffn", + backend_id="native", + supported_targets=frozenset({"rollout", "training"}), + supported_devices=frozenset({"*"}), + supported_dtypes=frozenset({"*"}), + supported_topologies={"*": "*"}, + determinism_or_alignment_properties={ + "selection": "runtime_native", + "strict_observable": False, + }, + lifecycle=OperatorLifecycle.DISTRIBUTED_CONTEXT, + implementation_class_or_factory=None, + fallback_policy=OperatorFallbackPolicy.RUNTIME_MANAGED, + version_or_build_fingerprint="runtime-native-ffn-unresolved-v1", + ), + ) + + def resolve_logp_op_type( logp_backend: Optional[str] = None, *, @@ -162,14 +342,101 @@ def resolve_logp_op_type( class KernelRegistry: - """ - Central dispatcher for high-performance kernels. - Handles dynamic routing between ROCm and CUDA backends at runtime. - """ + """Legacy hardware dispatcher plus a composed semantic operator catalog.""" def __init__(self): self._instance_cache: Dict[str, Any] = {} self._failed_backends: Set[str] = set() + self.semantic = SemanticOperatorCatalog(_default_semantic_descriptors()) + + common_roles = frozenset({AttentionRole.TRAIN, AttentionRole.INFER}) + common_dtypes = frozenset({AttentionDType.BF16, AttentionDType.FP16, AttentionDType.FP32}) + self._attention_capabilities = { + OpBackend.PYTORCH_NATIVE_ATTENTION: AttentionBackendCapability( + backend_id="pytorch-native-attention-ws1", + roles=common_roles, + modes=frozenset({AttentionMode.PREFILL, AttentionMode.CHUNKED_PREFILL}), + dtypes=common_dtypes, + cp_world_sizes=(1,), + exports_attention_lse=False, + deterministic_cp_merge=False, + supports_packed_varlen=False, + supports_kv_cache=False, + implementation_kind="reference", + ), + OpBackend.PYTORCH_NATIVE_KV_CACHE_ATTN: AttentionBackendCapability( + backend_id="pytorch-native-kv-cache-attention-ws1", + roles=common_roles, + modes=frozenset({AttentionMode.DECODE}), + dtypes=common_dtypes, + cp_world_sizes=(1,), + exports_attention_lse=False, + deterministic_cp_merge=False, + supports_packed_varlen=False, + supports_kv_cache=False, + implementation_kind="reference", + ), + OpBackend.PYTORCH_CP_ATTENTION: AttentionBackendCapability( + backend_id="pytorch-deterministic-cp-attention-reference", + roles=common_roles, + modes=frozenset({AttentionMode.PREFILL, AttentionMode.CHUNKED_PREFILL}), + dtypes=common_dtypes, + tp_world_sizes=(1, 2), + cp_world_sizes=(1, 2), + exports_attention_lse=True, + deterministic_cp_merge=True, + supports_packed_varlen=False, + supports_kv_cache=False, + supports_rope_metadata=False, + supports_fused_rope_attention=False, + supports_split_kv_disabled=True, + supports_split_kv_fixed=True, + supports_split_kv_auto=False, + reports_actual_split_kv_plan=True, + implementation_kind="deterministic", + ), + } + + # Truthful descriptors for the existing WS1 batch-invariant logp + # implementations: single-shard (TP=1), ignore-index masking only, no + # vocab-shard metadata, no vocab-domain LSE export. + common_logprob_roles = frozenset({LogprobRole.TRAIN, LogprobRole.INFER}) + common_logprob_dtypes = frozenset({LogprobDType.BF16, LogprobDType.FP16, LogprobDType.FP32}) + base_logprob_capabilities = { + OpBackend.PYTORCH_BATCH_INVARIANT_LOGP: LogprobBackendCapability( + backend_id="pytorch-batch-invariant-logp-ws1", + roles=common_logprob_roles, + dtypes=common_logprob_dtypes, + tp_world_sizes=(1,), + supports_vocab_padding=False, + mask_modes=frozenset({MaskMode.IGNORE_INDEX}), + exports_vocab_lse=False, + determinism_scopes=frozenset({DeterminismScope.FIXED_TOPOLOGY}), + implementation_kind="reference", + ), + OpBackend.TRITON_BATCH_INVARIANT_LOGP: LogprobBackendCapability( + backend_id="triton-batch-invariant-logp-ws1", + roles=common_logprob_roles, + dtypes=common_logprob_dtypes, + tp_world_sizes=(1,), + supports_vocab_padding=False, + mask_modes=frozenset({MaskMode.IGNORE_INDEX}), + exports_vocab_lse=False, + determinism_scopes=frozenset({DeterminismScope.FIXED_TOPOLOGY}), + implementation_kind="production", + ), + OpBackend.CUDA_BATCH_INVARIANT_LOGP_SM90: LogprobBackendCapability( + backend_id="cuda-batch-invariant-logp-sm90-ws1", + roles=common_logprob_roles, + dtypes=frozenset({LogprobDType.BF16, LogprobDType.FP32}), + tp_world_sizes=(1,), + supports_vocab_padding=False, + mask_modes=frozenset({MaskMode.IGNORE_INDEX}), + exports_vocab_lse=False, + determinism_scopes=frozenset({DeterminismScope.FIXED_TOPOLOGY}), + implementation_kind="production", + ), + } self._priority_map = { "cuda": { @@ -199,11 +466,21 @@ def __init__(self): OpBackend.CUDA_DETERMINISTIC_LOGP, OpBackend.PYTORCH_NATIVE, ], - "attn": [OpBackend.FLASH_ATTN, OpBackend.TRITON_GENERIC, OpBackend.PYTORCH_ATTN], + "attn": [ + OpBackend.FLASH_ATTN, + OpBackend.TRITON_GENERIC, + OpBackend.PYTORCH_ATTN, + ], "attention": [ OpBackend.CUDA_DETERMINISTIC_ATTENTION, OpBackend.PYTORCH_NATIVE_ATTENTION, ], + "cp_attention": [OpBackend.PYTORCH_CP_ATTENTION], + "ws2_attention": [ + OpBackend.PYTORCH_CP_ATTENTION, + OpBackend.CUDA_DETERMINISTIC_ATTENTION, + OpBackend.PYTORCH_NATIVE_ATTENTION, + ], "kv_cache_attention": [OpBackend.PYTORCH_NATIVE_KV_CACHE_ATTN], "grpo_loss": [OpBackend.TRITON_GRPO_LOSS, OpBackend.PYTORCH_GRPO_LOSS], "linear_logp": [ @@ -252,10 +529,18 @@ def __init__(self): OpBackend.TRITON_GENERIC, ], "attention": [OpBackend.PYTORCH_NATIVE_ATTENTION], + "cp_attention": [OpBackend.PYTORCH_CP_ATTENTION], + "ws2_attention": [ + OpBackend.PYTORCH_CP_ATTENTION, + OpBackend.PYTORCH_NATIVE_ATTENTION, + ], "kv_cache_attention": [OpBackend.PYTORCH_NATIVE_KV_CACHE_ATTN], "grpo_loss": [OpBackend.TRITON_GRPO_LOSS, OpBackend.PYTORCH_GRPO_LOSS], "rope": [OpBackend.TRITON_ROPE, OpBackend.PYTORCH_NATIVE_ROPE], - "linear_logp": [OpBackend.TRITON_LINEAR_LOGP, OpBackend.PYTORCH_LINEAR_LOGP], + "linear_logp": [ + OpBackend.TRITON_LINEAR_LOGP, + OpBackend.PYTORCH_LINEAR_LOGP, + ], "ratio_kl": [OpBackend.TRITON_RATIO_KL, OpBackend.PYTORCH_RATIO_KL], "pack": [OpBackend.PYTORCH_PACK], "det_gemm": [OpBackend.TRITON_DET_GEMM], @@ -276,6 +561,11 @@ def __init__(self): "logp_deterministic_indexed": [OpBackend.PYTORCH_NATIVE], "attn": [OpBackend.PYTORCH_ATTN], "attention": [OpBackend.PYTORCH_NATIVE_ATTENTION], + "cp_attention": [OpBackend.PYTORCH_CP_ATTENTION], + "ws2_attention": [ + OpBackend.PYTORCH_CP_ATTENTION, + OpBackend.PYTORCH_NATIVE_ATTENTION, + ], "kv_cache_attention": [OpBackend.PYTORCH_NATIVE_KV_CACHE_ATTN], "grpo_loss": [OpBackend.PYTORCH_GRPO_LOSS], "rope": [OpBackend.PYTORCH_NATIVE_ROPE], @@ -291,19 +581,51 @@ def __init__(self): "swiglu": [OpBackend.PYTORCH_NATIVE_SWIGLU], }, } - # Preserve the former CPU fallback behavior for every operator on NPU, - # then override only the operator with an Ascend-specific backend. - self._priority_map["npu"] = { - op_type: candidates.copy() for op_type, candidates in self._priority_map["cpu"].items() - } - self._priority_map["npu"]["batch_invariant_logp"] = [ - OpBackend.ASCEND_BATCH_INVARIANT_LOGP, - OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, - ] logger.info(f"KernelRegistry initialized for {device_ctx.device_type}") self._adjust_priority_for_hardware() self._adjust_priority_from_env() + # WS2 dispatch owns its candidate list, seeded from the legacy + # batch_invariant_logp priority but decoupled afterwards: neither + # path's registrations may affect the other. + self._logprob_candidates: Dict[str, list] = { + platform: list(ops.get("batch_invariant_logp", [])) + for platform, ops in self._priority_map.items() + } + # Capabilities are scoped per platform: the same backend enum may + # truthfully declare different support on cuda vs rocm vs cpu. + self._logprob_capabilities: Dict[str, Dict[OpBackend, LogprobBackendCapability]] = { + platform: { + backend: base_logprob_capabilities[backend] + for backend in candidates + if backend in base_logprob_capabilities + } + for platform, candidates in self._logprob_candidates.items() + } + + # deterministic vocab-parallel TP logprob reference. + ws2_tp_logprob_capability = LogprobBackendCapability( + backend_id="pytorch-vocab-parallel-logp-ws2", + roles=common_logprob_roles, + dtypes=common_logprob_dtypes, + tp_world_sizes=None, + cp_world_sizes=None, + supports_vocab_padding=True, + mask_modes=frozenset({MaskMode.EXPLICIT_ACTIVE_MASK, MaskMode.IGNORE_INDEX}), + exports_vocab_lse=True, + determinism_scopes=frozenset( + {DeterminismScope.CROSS_TP_BITWISE, DeterminismScope.FIXED_TOPOLOGY} + ), + implementation_kind="reference", + ) + for ws2_platform in self._priority_map: + self.register_logprob_backend( + OpBackend.PYTORCH_VOCAB_PARALLEL_LOGP, + ws2_tp_logprob_capability, + platform=ws2_platform, + prepend=True, + ) + def _adjust_priority_from_env(self): rocm_attn_backend = os.getenv("RL_KERNEL_ROCM_ATTN_BACKEND", "").strip().lower() if rocm_attn_backend in {"flash_attn", "flash-attn", "flash_attention"}: @@ -329,7 +651,8 @@ def _adjust_priority_from_env(self): ) def _adjust_priority_for_hardware(self): - """Adjust CUDA priorities for hardware-gated experimental and production kernels.""" + """Adjust CUDA priorities for hardware-gated kernels.""" + if device_ctx.device_type != "cuda": return try: @@ -351,8 +674,6 @@ def _adjust_priority_for_hardware(self): if OpBackend.CUDA_FUSED_LOGP_SM90 not in logp_list: logp_list.insert(0, OpBackend.CUDA_FUSED_LOGP_SM90) - # The fused linear-logp SM90 kernel uses TMA bulk-tensor copies built - # for sm_90a -- gate strictly on cc_major == 9 (Hopper), not >= 9. linear_logp_compiled = _EXT_AVAILABLE and hasattr(_C, "fused_linear_logp_sm90") if linear_logp_compiled and cc_major == 9: ll_list = self._priority_map["cuda"]["linear_logp"] @@ -370,7 +691,6 @@ def _adjust_priority_for_hardware(self): f"SM{cc}: fused linear-logp SM90 kernel not compiled into _C; " "using generic linear-logp backend." ) - sm90_embedding_compiled = _EXT_AVAILABLE and hasattr(_C, "embedding_sm90_forward") if sm90_embedding_compiled and cc_major == 9: embedding_list = self._priority_map["cuda"]["embedding"] @@ -382,47 +702,25 @@ def _adjust_priority_for_hardware(self): lm_head_list = self._priority_map["cuda"]["lm_head"] if OpBackend.CUDA_SM90_LM_HEAD not in lm_head_list: lm_head_list.insert(0, OpBackend.CUDA_SM90_LM_HEAD) - except Exception as e: - logger.warning(f"Failed to probe device capability: {e}") + except Exception as exc: + logger.warning(f"Failed to probe device capability: {exc}") def get_op(self, op_type: str, device: torch.device | str | None = None) -> Any: - """Core distribution logic: Automatically select the best operator - based on hardware and priority. - """ + """Select the best legacy operator for the requested device.""" + platform = self._platform_for_device(device) candidates = self._priority_map.get(platform, {}).get(op_type, [OpBackend.PYTORCH_NATIVE]) for backend in candidates: - if backend.name in self._instance_cache: - return self._instance_cache[backend.name] - - if backend.name in self._failed_backends: - continue - - op_class = self._load_backend(backend) - if op_class: - try: - op_instance = op_class() - self._instance_cache[backend.name] = op_instance - return op_instance - except Exception as e: - logger.error(f"Failed to instantiate {backend.name}: {e}") - self._failed_backends.add(backend.name) - else: - self._failed_backends.add(backend.name) + op_instance = self._get_or_create_backend(backend) + if op_instance is not None: + return op_instance raise RuntimeError(f"No functional backend found for {op_type} on {platform}") def _platform_for_device(self, device: torch.device | str | None) -> str: if device is None: - if device_ctx.is_rocm: - return "rocm" - if device_ctx.device_type == "cuda": - return "cuda" - if device_ctx.device_type == "npu": - return "npu" - return "cpu" - + return self._platform() resolved = torch.device(device) if resolved.type == "cuda": return "rocm" if torch.version.hip is not None else "cuda" @@ -430,24 +728,286 @@ def _platform_for_device(self, device: torch.device | str | None) -> str: return resolved.type return "cpu" - def _load_backend(self, backend: OpBackend) -> Optional[Type]: - """Dynamic loading technique: Import modules only when needed - and check environment dependencies. + def register_logprob_backend( + self, + backend: OpBackend, + capability: LogprobBackendCapability, + *, + platform: Optional[str] = None, + prepend: bool = False, + ) -> None: + """Register (or replace) a backend for WS2 contract-aware logprob dispatch. + + This is the supported seam for making a new backend selectable by + ``get_logprob_op`` (e.g. the deterministic vocab-parallel TP reference + from issue #241 PR 3) without touching the legacy ``get_op`` priority + lists. Registering the same backend again replaces its capability + without duplicating the candidate entry. + """ + + if not isinstance(backend, OpBackend): + raise LogprobContractError("backend must be an OpBackend") + if not isinstance(capability, LogprobBackendCapability): + raise LogprobContractError("capability must be a LogprobBackendCapability") + resolved_platform = platform if platform is not None else self._platform() + if resolved_platform not in self._priority_map: + raise LogprobContractError( + f"unsupported platform {resolved_platform!r}; expected one of " + f"{sorted(self._priority_map)}" + ) + candidates = self._logprob_candidates.setdefault(resolved_platform, []) + self._logprob_capabilities.setdefault(resolved_platform, {})[backend] = capability + if backend not in candidates: + if prepend: + candidates.insert(0, backend) + else: + candidates.append(backend) + + def get_logprob_op( + self, + contract: LogprobContract, + *, + requested_backend: str = "auto", + ) -> LogprobDispatchResult: + """Resolve only a backend that explicitly supports the WS2 logprob contract. + + This entry point is intentionally separate from legacy ``get_op`` so + existing callers retain their current behavior while WS2 callers cannot + silently fall back to a backend with different distributed semantics. + + ``requested_backend`` is either a case-insensitive policy keyword + (``auto`` | ``production`` | ``reference`` | ``deterministic``) or an + exact, case-sensitive stable backend id. Strictness comes from the + contract's capability checks, not from this policy string, so the + default is ``auto``. """ + + if not isinstance(contract, LogprobContract): + raise LogprobContractError("contract must be a LogprobContract") + if not isinstance(requested_backend, str) or not requested_backend.strip(): + raise LogprobContractError("requested_backend must be a non-empty string") + requested_backend = requested_backend.strip() + if requested_backend.lower() == "deterministic": + raise LogprobContractError( + 'requested_backend="deterministic" is not a dispatch policy; request ' + "determinism through ReductionSpec.determinism_scope and match it against " + "backend determinism_scopes instead" + ) + + platform = self._platform() + candidates = self._logprob_candidates.get(platform, []) + rejected: list[str] = [] + # provenance["fallback"] reports only capability/load rejections of + # otherwise-eligible candidates; skips caused purely by the caller's + # own requested_backend policy filter are not fallbacks. + capability_rejections = 0 + + platform_capabilities = self._logprob_capabilities.get(platform, {}) + for backend in candidates: + capability = platform_capabilities.get(backend) + if capability is None: + rejected.append(f"{backend.name}: no LogprobBackendCapability declared") + capability_rejections += 1 + continue + policy_mismatch = self._logprob_policy_mismatch(requested_backend, capability) + if policy_mismatch is not None: + # Excluded by the caller's own policy: never a fallback, even + # if the candidate would also have failed capability checks. + rejected.append(f"{backend.name}: {policy_mismatch}") + continue + capability_incompat = list(capability.incompatibilities(contract)) + if capability_incompat: + rejected.append(f"{backend.name}: " + "; ".join(capability_incompat)) + capability_rejections += 1 + continue + + op = self._get_or_create_backend(backend) + if op is None: + rejected.append(f"{backend.name}: backend could not be loaded or instantiated") + capability_rejections += 1 + continue + + provenance = { + "requested_backend": requested_backend, + "actual_backend": capability.backend_id, + "backend_enum": backend.name, + "platform": platform, + "fallback": capability_rejections > 0, + "prior_rejections": list(rejected), + "contract": contract.to_dict(), + "capability": capability.to_dict(), + } + return LogprobDispatchResult( + op=op, + capability=capability, + provenance=provenance, + ) + + details = " | ".join(rejected) if rejected else "no candidates registered" + requested = contract.to_dict() + raise RuntimeError( + "No logprob backend supports the requested WS2 contract on " + f"{platform}: role={requested['role']}, dtype={requested['dtype']}, " + f"TP={contract.sharding.tp_world_size}, CP={contract.sharding.cp_world_size}, " + f"padded_vocab={contract.sharding.padded_vocab_size}, " + f"real_vocab={contract.sharding.real_vocab_size}. Rejections: {details}" + ) + + @staticmethod + def _logprob_policy_mismatch( + requested_backend: str, + capability: LogprobBackendCapability, + ) -> str | None: + policy = requested_backend.lower() + if policy == "auto": + return None + if policy in IMPLEMENTATION_KINDS: + if capability.implementation_kind == policy: + return None + return ( + f"implementation_kind={capability.implementation_kind} does not satisfy " + f"requested_backend={policy}" + ) + if capability.backend_id == requested_backend: + return None + return ( + f"backend_id={capability.backend_id} does not match " + f"requested_backend={requested_backend}" + ) + + def get_attention_op( + self, + contract: AttentionContract, + *, + requested_backend: str = "deterministic", + ) -> AttentionDispatchResult: + """Resolve only a backend that explicitly supports the WS2 contract.""" + + if not isinstance(contract, AttentionContract): + raise AttentionContractError("contract must be an AttentionContract") + if not isinstance(requested_backend, str) or not requested_backend.strip(): + raise AttentionContractError("requested_backend must be a non-empty string") + requested_backend = requested_backend.strip().lower() + + platform = self._platform() + candidates = self._priority_map.get(platform, {}).get("ws2_attention", []) + rejected: list[str] = [] + + for backend in candidates: + capability = self._attention_capabilities.get(backend) + if capability is None: + rejected.append(f"{backend.name}: no AttentionBackendCapability declared") + continue + incompatibilities = list(capability.incompatibilities(contract)) + policy_mismatch = self._attention_policy_mismatch(requested_backend, capability) + if policy_mismatch is not None: + incompatibilities.append(policy_mismatch) + if incompatibilities: + rejected.append(f"{backend.name}: " + "; ".join(incompatibilities)) + continue + + op = self._get_or_create_backend(backend) + if op is None: + rejected.append(f"{backend.name}: backend could not be loaded or instantiated") + continue + + return AttentionDispatchResult( + op=op, + capability=capability, + provenance={ + "requested_backend": requested_backend, + "actual_backend": capability.backend_id, + "backend_enum": backend.name, + "platform": platform, + "fallback": bool(rejected), + "prior_rejections": list(rejected), + "contract": contract.to_dict(), + "capability": capability.to_dict(), + }, + ) + + details = " | ".join(rejected) if rejected else "no candidates registered" + requested = contract.to_dict() + raise RuntimeError( + "No attention backend supports the requested WS2 contract on " + f"{platform}: role={requested['role']}, mode={requested['mode']}, " + f"dtype={requested['dtype']}, TP={contract.sharding.tp_world_size}, " + f"CP={contract.sharding.cp_world_size}. Rejections: {details}" + ) + + @staticmethod + def _attention_policy_mismatch( + requested_backend: str, + capability: AttentionBackendCapability, + ) -> Optional[str]: + if requested_backend == "auto": + return None + if requested_backend in {"production", "reference", "deterministic"}: + if capability.implementation_kind == requested_backend: + return None + return ( + f"implementation_kind={capability.implementation_kind} does not satisfy " + f"requested_backend={requested_backend}" + ) + if capability.backend_id == requested_backend: + return None + return ( + f"backend_id={capability.backend_id} does not match " + f"requested_backend={requested_backend}" + ) + + @staticmethod + def _platform() -> str: + if device_ctx.is_rocm: + return "rocm" + if device_ctx.device_type == "cuda": + return "cuda" + return "cpu" + + def _get_or_create_backend(self, backend: OpBackend) -> Optional[Any]: + if backend.name in self._instance_cache: + return self._instance_cache[backend.name] + if backend.name in self._failed_backends: + return None + + op_class = self._load_backend(backend) + if op_class is None: + self._failed_backends.add(backend.name) + return None + try: + op = op_class() + except Exception as exc: + logger.error(f"Failed to instantiate {backend.name}: {exc}") + self._failed_backends.add(backend.name) + return None + self._instance_cache[backend.name] = op + return op + + def _load_backend(self, backend: OpBackend) -> Optional[Type]: + """Import a legacy backend and distinguish wrapper bugs from absence.""" + module_path, class_name = backend.value.rsplit(".", 1) try: module = importlib.import_module(module_path) return getattr(module, class_name) - except (ImportError, AttributeError, ModuleNotFoundError) as e: - missing_module = str(e.name) if hasattr(e, "name") else "" + except (ImportError, AttributeError, ModuleNotFoundError) as exc: + missing_module = str(exc.name) if hasattr(exc, "name") else "" is_missing_backend = missing_module and ( missing_module == module_path or module_path.startswith(missing_module) ) if missing_module and "rl_engine" in missing_module and not is_missing_backend: - logger.critical(f"Internal wrapper implementation bug in '{module_path}': {e}") - raise e - logger.warning(f"Backend {backend.name} unavailable: {e}. Falling back...") + logger.critical(f"Internal wrapper implementation bug in '{module_path}': {exc}") + raise + logger.warning(f"Backend {backend.name} unavailable: {exc}. Falling back...") return None kernel_registry = KernelRegistry() + + +__all__ = [ + "KernelRegistry", + "OpBackend", + "kernel_registry", + "resolve_logp_op_type", +] diff --git a/rl_engine/kernels/semantic_registry.py b/rl_engine/kernels/semantic_registry.py new file mode 100644 index 00000000..99e152d9 --- /dev/null +++ b/rl_engine/kernels/semantic_registry.py @@ -0,0 +1,790 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Exact semantic-operator catalog with case-local instantiation state.""" + +from __future__ import annotations + +import hashlib +import importlib +import inspect +import json +from dataclasses import dataclass, field, fields, replace +from enum import Enum +from pathlib import Path +from types import CodeType, MappingProxyType +from typing import Any, Callable, Iterable, Mapping, Optional, Sequence, cast + + +class OperatorLifecycle(str, Enum): + REQUEST = "request" + ENGINE_CONSTRUCTION = "engine_construction" + DISTRIBUTED_CONTEXT = "distributed_context" + PROCESS = "process" + + +class OperatorFallbackPolicy(str, Enum): + ERROR = "error" + DECLARED = "declared" + RUNTIME_MANAGED = "runtime_managed" + + +@dataclass(frozen=True) +class OperatorResolutionPolicy: + strict: bool = True + allow_test_backends: bool = False + + +_Policy = Optional[OperatorResolutionPolicy] + + +class _JsonRecord: + def to_dict(self) -> dict[str, Any]: + return { + item.name: _json_value(getattr(self, item.name)) for item in fields(cast(Any, self)) + } + + +@dataclass(frozen=True) +class OperatorRequirements(_JsonRecord): + device: str + dtype: str + topology: Mapping[str, Any] = field(default_factory=dict) + alignment_properties: Mapping[str, Any] = field(default_factory=dict) + schema_version: str = "rlkernel.semantic_operator.requirements.v1" + + def __post_init__(self) -> None: + normalized = ( + ("device", _normalize_device(self.device)), + ("dtype", _normalize_dtype(self.dtype)), + ("topology", _freeze(self.topology)), + ("alignment_properties", _freeze(self.alignment_properties)), + ) + for name, value in normalized: + object.__setattr__(self, name, value) + + +@dataclass(frozen=True) +class OperatorBackendDescriptor(_JsonRecord): + semantic_op: str + backend_id: str + supported_targets: frozenset[str] + supported_devices: frozenset[str] + supported_dtypes: frozenset[str] + supported_topologies: Mapping[str, Any] + determinism_or_alignment_properties: Mapping[str, Any] + lifecycle: OperatorLifecycle + implementation_class_or_factory: Optional[str | Callable[..., Any]] + fallback_policy: OperatorFallbackPolicy + version_or_build_fingerprint: str + is_smoke_only: bool = False + schema_version: str = "rlkernel.semantic_operator.backend_descriptor.v1" + + def __post_init__(self) -> None: + values = { + "semantic_op": self.semantic_op.strip(), + "backend_id": self.backend_id.strip(), + "supported_targets": _normalized_values(self.supported_targets, str), + "supported_devices": _normalized_values(self.supported_devices, _normalize_device), + "supported_dtypes": _normalized_values(self.supported_dtypes, _normalize_dtype), + } + for name, value in values.items(): + if not value: + raise ValueError(f"{name} must not be empty") + if not self.version_or_build_fingerprint.strip(): + raise ValueError("version_or_build_fingerprint must not be empty") + values.update( + supported_topologies=_freeze(self.supported_topologies), + determinism_or_alignment_properties=_freeze(self.determinism_or_alignment_properties), + lifecycle=OperatorLifecycle(self.lifecycle), + fallback_policy=OperatorFallbackPolicy(self.fallback_policy), + ) + for name, value in values.items(): + object.__setattr__(self, name, value) + + @property + def implementation_reference(self) -> Optional[str]: + return _reference(self.implementation_class_or_factory) + + @property + def is_strictly_observable(self) -> bool: + return bool( + self.determinism_or_alignment_properties.get( + "strict_observable", self.implementation_class_or_factory is not None + ) + ) + + @property + def descriptor_fingerprint(self) -> str: + return _fingerprint(self.to_dict(include_descriptor_fingerprint=False)) + + def to_dict(self, *, include_descriptor_fingerprint: bool = True) -> dict[str, Any]: + result = super().to_dict() + result["implementation_class_or_factory"] = self.implementation_reference + if include_descriptor_fingerprint: + result["descriptor_fingerprint"] = self.descriptor_fingerprint + return result + + +@dataclass(frozen=True) +class OperatorCapabilityDecision(_JsonRecord): + capability: str + requested: Any + supported: Any + passed: bool + reason: str + + +@dataclass(frozen=True) +class OperatorResolutionTrace(_JsonRecord): + semantic_op: str + requested_backend: str + target: str + strict: bool + status: str + concrete_backend: Optional[str] + implementation_reference: Optional[str] + descriptor_fingerprint: Optional[str] + capability_decisions: tuple[OperatorCapabilityDecision, ...] + fallback_attempts: tuple[str, ...] = () + schema_version: str = "rlkernel.semantic_operator.resolution_trace.v1" + + +@dataclass(frozen=True) +class OperatorResolution(_JsonRecord): + descriptor: OperatorBackendDescriptor + requirements: OperatorRequirements + target: str + strict: bool + trace: OperatorResolutionTrace + schema_version: str = "rlkernel.semantic_operator.resolution.v1" + + +@dataclass(frozen=True) +class OperatorInstanceProvenance(_JsonRecord): + semantic_op: str + backend_id: str + target: str + factory_reference: str + concrete_implementation: str + descriptor_fingerprint: str + implementation_fingerprint: str + instance_fingerprint: str + factory_options: Mapping[str, Any] = field(default_factory=dict) + factory_options_fingerprint: str = "" + schema_version: str = "rlkernel.semantic_operator.instance_provenance.v1" + + def __post_init__(self) -> None: + object.__setattr__(self, "factory_options", _freeze(self.factory_options)) + + +@dataclass(frozen=True) +class _InstanceRecord: + instance: Any + descriptor_fingerprint: str + target: str + factory: Callable[..., Any] + factory_options: Mapping[str, Any] + + +class OperatorRegistrationError(ValueError): + pass + + +class OperatorResolutionError(RuntimeError): + def __init__(self, message: str, trace: OperatorResolutionTrace): + super().__init__(message) + self.trace = trace + + +class OperatorInstantiationError(RuntimeError): + pass + + +class SemanticOperatorCatalog: + def __init__(self, descriptors: Iterable[OperatorBackendDescriptor] = ()): + self._descriptors: dict[tuple[str, str], OperatorBackendDescriptor] = {} + for descriptor in descriptors: + self.register_backend(descriptor) + + def register_backend( + self, + descriptor: OperatorBackendDescriptor, + *, + replace: bool = False, + ) -> None: + if not isinstance(descriptor, OperatorBackendDescriptor): + raise TypeError("descriptor must be an OperatorBackendDescriptor") + key = (descriptor.semantic_op, descriptor.backend_id) + if key in self._descriptors and not replace: + raise OperatorRegistrationError(f"operator backend is already registered: {key!r}") + self._descriptors[key] = descriptor + + def backend_descriptor( + self, + semantic_op: str, + backend_id: str, + ) -> Optional[OperatorBackendDescriptor]: + return self._descriptors.get((semantic_op.strip(), backend_id.strip())) + + def backend_descriptors( + self, + semantic_op: Optional[str] = None, + ) -> tuple[OperatorBackendDescriptor, ...]: + values: Iterable[OperatorBackendDescriptor] = self._descriptors.values() + if semantic_op is not None: + normalized = semantic_op.strip() + values = (value for value in values if value.semantic_op == normalized) + return tuple(sorted(values, key=lambda value: (value.semantic_op, value.backend_id))) + + def session(self, policy: _Policy = None) -> OperatorSession: + return OperatorSession(self, policy=policy) + + def _resolve( + self, + *, + semantic_op: str, + requested_backend: str, + target: str, + requirements: OperatorRequirements, + policy: OperatorResolutionPolicy, + ) -> OperatorResolution: + semantic_op = semantic_op.strip() + requested_backend = requested_backend.strip() + target = target.strip().lower() + if not semantic_op or not requested_backend or not target: + raise ValueError("semantic_op, requested_backend, and target must not be empty") + if not isinstance(requirements, OperatorRequirements): + raise TypeError("requirements must be an OperatorRequirements") + + descriptor = self.backend_descriptor(semantic_op, requested_backend) + if descriptor is None: + decision = _decision( + "registration", + requested_backend, + [item.backend_id for item in self.backend_descriptors(semantic_op)], + passed=False, + ) + trace = _trace( + semantic_op, + requested_backend, + target, + policy, + "unsupported", + (decision,), + ) + raise OperatorResolutionError( + f"exact operator backend {requested_backend!r} is not registered; " + "no fallback was attempted", + trace, + ) + + topology_capabilities = _target_topology_capabilities( + descriptor.supported_topologies, + target, + ) + topology_ok = topology_capabilities is not None and _supports_complete_mapping( + topology_capabilities, + requirements.topology, + ) + decisions = ( + _decision("target", target, descriptor.supported_targets), + _decision( + "smoke_opt_in", + descriptor.is_smoke_only, + policy.allow_test_backends, + not descriptor.is_smoke_only or policy.allow_test_backends, + ), + _decision("device", requirements.device, descriptor.supported_devices), + _decision("dtype", requirements.dtype, descriptor.supported_dtypes), + _decision( + "topology", + requirements.topology, + topology_capabilities, + topology_ok, + ), + _decision( + "alignment_properties", + requirements.alignment_properties, + descriptor.determinism_or_alignment_properties, + ), + _decision( + "strict_observability", + policy.strict, + descriptor.is_strictly_observable, + not policy.strict or descriptor.is_strictly_observable, + ), + _decision( + "fallback_policy", + "error" if policy.strict else "declared", + descriptor.fallback_policy.value, + not policy.strict or descriptor.fallback_policy is OperatorFallbackPolicy.ERROR, + ), + ) + failed = tuple(item for item in decisions if not item.passed) + observable = descriptor.is_strictly_observable + status = "unsupported" if failed else ("resolved" if observable else "unobservable") + trace = _trace( + semantic_op, + requested_backend, + target, + policy, + status, + decisions, + descriptor, + ) + if failed: + raise OperatorResolutionError( + f"exact operator backend {requested_backend!r} is unsupported: " + + "; ".join(item.reason for item in failed), + trace, + ) + return OperatorResolution(descriptor, requirements, target, policy.strict, trace) + + +class OperatorSession: + def __init__(self, catalog: SemanticOperatorCatalog, policy: _Policy = None): + if not isinstance(catalog, SemanticOperatorCatalog): + raise TypeError("catalog must be a SemanticOperatorCatalog") + self.catalog = catalog + self.policy = policy or OperatorResolutionPolicy() + self._cache: dict[str, Any] = {} + self._records: dict[int, _InstanceRecord] = {} + + def resolve( + self, + *, + semantic_op: str, + requested_backend: str, + target: str, + requirements: OperatorRequirements, + policy: _Policy = None, + strict: Optional[bool] = None, + ) -> OperatorResolution: + return self.catalog._resolve( + semantic_op=semantic_op, + requested_backend=requested_backend, + target=target, + requirements=requirements, + policy=_resolve_policy(policy or self.policy, strict), + ) + + def instantiate( + self, + resolution: OperatorResolution, + *, + factory_kwargs: Optional[Mapping[str, Any]] = None, + cache: bool = False, + ) -> Any: + if not isinstance(resolution, OperatorResolution): + raise TypeError("resolution must be an OperatorResolution") + descriptor = resolution.descriptor + implementation = descriptor.implementation_class_or_factory + if resolution.trace.status != "resolved" or implementation is None: + raise OperatorInstantiationError( + f"backend {descriptor.backend_id!r} has no exact implementation" + ) + options = dict(factory_kwargs or {}) + cache_key = _fingerprint( + { + "descriptor": descriptor.descriptor_fingerprint, + "target": resolution.target, + "requirements": resolution.requirements.to_dict(), + "options": options, + } + ) + if cache and cache_key in self._cache: + return self._cache[cache_key] + factory = _load_factory(implementation) + try: + instance = factory(**options) + except Exception as exc: + raise OperatorInstantiationError( + f"failed to instantiate backend {descriptor.backend_id!r}: {exc}" + ) from exc + if instance is None: + raise OperatorInstantiationError("operator factory returned None") + self._records[id(instance)] = _InstanceRecord( + instance, + descriptor.descriptor_fingerprint, + resolution.target, + factory, + _freeze(options), + ) + if cache: + self._cache[cache_key] = instance + return instance + + def instance_provenance( + self, + resolution: OperatorResolution, + instance: Any, + ) -> OperatorInstanceProvenance: + descriptor = resolution.descriptor + record = self._records.get(id(instance)) + if ( + record is None + or record.instance is not instance + or record.descriptor_fingerprint != descriptor.descriptor_fingerprint + or record.target != resolution.target + ): + raise OperatorInstantiationError( + "operator instance does not match this session resolution" + ) + factory_reference = descriptor.implementation_reference + concrete = _reference(type(instance)) + if factory_reference is None or concrete is None: + raise OperatorInstantiationError("operator implementation is not observable") + options_fingerprint = _fingerprint(record.factory_options) + implementation_fingerprint = operator_implementation_fingerprint( + record.factory, + instance, + ) + instance_fingerprint = operator_instance_fingerprint( + descriptor_fingerprint=descriptor.descriptor_fingerprint, + factory_reference=factory_reference, + concrete_implementation=concrete, + implementation_fingerprint=implementation_fingerprint, + factory_options_fingerprint=options_fingerprint, + ) + return OperatorInstanceProvenance( + descriptor.semantic_op, + descriptor.backend_id, + resolution.target, + factory_reference, + concrete, + descriptor.descriptor_fingerprint, + implementation_fingerprint, + instance_fingerprint, + record.factory_options, + options_fingerprint, + ) + + def clear_instance_cache(self) -> None: + self._cache.clear() + + +def operator_implementation_fingerprint( + implementation: str | Callable[..., Any], + instance: Any, +) -> str: + return implementation_fingerprint( + implementation, + instance=instance, + entrypoints=("apply_fp32", "__call__"), + ) + + +def implementation_fingerprint( + implementation: str | Callable[..., Any], + *, + instance: Any = None, + entrypoints: Sequence[str] = (), +) -> str: + """Fingerprint executable code, not only its import reference. + + The identity includes source or bytecode for the resolved factory, its + concrete class, the defining modules, and explicitly named runtime entry + points. Module content covers helper functions called by an entry point; + callable identities additionally make in-process replacements observable. + """ + + factory = _load_factory(implementation) + concrete_type = type(instance) if instance is not None else None + runtime_entrypoints = {} + if instance is not None: + for name in sorted(set(entrypoints)): + value = getattr(instance, name, None) + if callable(value): + runtime_entrypoints[name] = _callable_identity(value) + return _fingerprint( + { + "factory": _implementation_identity(factory), + "concrete_type": ( + _implementation_identity(concrete_type) if concrete_type is not None else None + ), + "runtime_entrypoints": runtime_entrypoints, + } + ) + + +def operator_instance_fingerprint(**identity: str) -> str: + return _fingerprint(identity) + + +def _trace( + semantic_op: str, + backend: str, + target: str, + policy: OperatorResolutionPolicy, + status: str, + decisions: tuple[OperatorCapabilityDecision, ...], + descriptor: Optional[OperatorBackendDescriptor] = None, +) -> OperatorResolutionTrace: + observable = descriptor is not None and descriptor.is_strictly_observable + return OperatorResolutionTrace( + semantic_op, + backend, + target, + policy.strict, + status, + ( + descriptor.backend_id + if descriptor is not None and observable and status != "unsupported" + else None + ), + descriptor.implementation_reference if descriptor else None, + descriptor.descriptor_fingerprint if descriptor else None, + decisions, + ) + + +def _decision( + capability: str, + requested: Any, + supported: Any, + passed: Optional[bool] = None, +) -> OperatorCapabilityDecision: + passed = _supports(supported, requested) if passed is None else passed + actionable = { + "smoke_opt_in": "smoke backend use requires explicit opt-in", + "strict_observability": "runtime-native implementation is not exactly observable", + "fallback_policy": "strict resolution forbids declared or runtime fallback", + } + return OperatorCapabilityDecision( + capability, + requested, + supported, + passed, + ( + f"{capability} is supported" + if passed + else actionable.get(capability, f"{capability} is unsupported") + ), + ) + + +def _supports(supported: Any, requested: Any) -> bool: + if isinstance(supported, str) and supported in {"*", "any"}: + return True + if isinstance(supported, Mapping): + if not isinstance(requested, Mapping): + return False + wildcard = supported.get("*") + return all( + _supports(supported.get(key, wildcard), value) + for key, value in requested.items() + if key in supported or wildcard is not None + ) and all(key in supported or wildcard is not None for key in requested) + if isinstance(supported, (set, frozenset, tuple, list)): + if isinstance(requested, (set, frozenset, tuple, list)): + return all(any(_supports(item, value) for item in supported) for value in requested) + return any(_supports(item, requested) for item in supported) + return supported == requested + + +def _target_topology_capabilities(supported: Any, target: str) -> Any: + if not isinstance(supported, Mapping): + return supported + targeted = any(key in supported for key in ("rollout", "training")) + if not targeted: + return supported + return supported.get(target, supported.get("*")) + + +def _supports_complete_mapping(supported: Any, requested: Any) -> bool: + if isinstance(supported, Mapping) and "*" not in supported: + if not isinstance(requested, Mapping) or any(key not in requested for key in supported): + return False + return _supports(supported, requested) + + +def _resolve_policy(policy: _Policy, strict: Optional[bool]) -> OperatorResolutionPolicy: + policy = policy or OperatorResolutionPolicy() + return policy if strict is None else replace(policy, strict=strict) + + +def _load_factory(value: str | Callable[..., Any]) -> Callable[..., Any]: + if callable(value): + return value + try: + module_name, attribute = value.rsplit(".", 1) + factory = getattr(importlib.import_module(module_name), attribute) + except (ValueError, ImportError, AttributeError, ModuleNotFoundError) as exc: + raise OperatorInstantiationError(f"operator factory {value!r} is unavailable") from exc + if not callable(factory): + raise OperatorInstantiationError(f"operator factory {value!r} is not callable") + return factory + + +def _reference(value: Any) -> Optional[str]: + if value is None or isinstance(value, str): + return value + module = getattr(value, "__module__", type(value).__module__) + qualname = getattr(value, "__qualname__", type(value).__qualname__) + return f"{module}.{qualname}" + + +def _normalize_device(value: Any) -> str: + value = str(value).strip().lower() + if value.startswith("torch.device("): + value = value.removeprefix("torch.device(").removesuffix(")").strip("'\"") + if value.startswith("cuda:"): + return "cuda" + return {"gpu": "cuda", "hip": "rocm"}.get(value, value) + + +def _normalize_dtype(value: Any) -> str: + value = str(value).strip().lower().replace("torch.", "") + return { + "fp32": "float32", + "float": "float32", + "bf16": "bfloat16", + "fp16": "float16", + "half": "float16", + }.get(value, value) + + +def _normalized_values(values: Iterable[Any], normalize: Callable[[Any], str]) -> frozenset[str]: + return frozenset(value for item in values if (value := normalize(item).strip().lower())) + + +def _freeze(value: Any) -> Any: + if isinstance(value, Mapping): + return MappingProxyType({str(key): _freeze(item) for key, item in value.items()}) + if isinstance(value, (tuple, list)): + return tuple(_freeze(item) for item in value) + if isinstance(value, (set, frozenset)): + return frozenset(_freeze(item) for item in value) + return value + + +def _json_value(value: Any) -> Any: + if isinstance(value, _JsonRecord): + return value.to_dict() + if isinstance(value, Enum): + return value.value + if isinstance(value, Mapping): + return {str(key): _json_value(item) for key, item in sorted(value.items())} + if isinstance(value, (set, frozenset)): + return sorted((_json_value(item) for item in value), key=repr) + if isinstance(value, (tuple, list)): + return [_json_value(item) for item in value] + if callable(value): + return _reference(value) + return value + + +def _implementation_identity(value: Any) -> Mapping[str, Any]: + reference = _reference(value) + identity: dict[str, Any] = { + "reference": reference, + "kind": "class" if inspect.isclass(value) else "callable", + "callable": _callable_identity(value), + "module": _module_identity(getattr(value, "__module__", None)), + } + if inspect.isclass(value): + identity["members"] = { + name: _callable_identity(member) + for name, raw_member in sorted(vars(value).items()) + if (member := _descriptor_callable(raw_member)) is not None + } + return identity + + +def _descriptor_callable(value: Any) -> Optional[Callable[..., Any]]: + if isinstance(value, (classmethod, staticmethod)): + value = value.__func__ + elif isinstance(value, property): + return None + return value if callable(value) else None + + +def _callable_identity(value: Any) -> Mapping[str, Any]: + if inspect.ismethod(value): + value = value.__func__ + try: + unwrapped = inspect.unwrap(value) + except (TypeError, ValueError): + unwrapped = value + code = getattr(unwrapped, "__code__", None) + try: + source = inspect.getsource(unwrapped) + except (OSError, TypeError): + source = None + identity: dict[str, Any] = { + "reference": _reference(unwrapped), + "source_sha256": ( + hashlib.sha256(source.encode("utf-8")).hexdigest() if source is not None else None + ), + "code_sha256": _code_fingerprint(code) if isinstance(code, CodeType) else None, + } + if isinstance(code, CodeType): + identity["defaults"] = _code_value(getattr(unwrapped, "__defaults__", None)) + identity["keyword_defaults"] = _code_value(getattr(unwrapped, "__kwdefaults__", None)) + return identity + + +def _code_fingerprint(code: CodeType) -> str: + return _fingerprint( + { + "bytecode": code.co_code.hex(), + "constants": tuple(_code_value(value) for value in code.co_consts), + "names": code.co_names, + "variables": code.co_varnames, + "free_variables": code.co_freevars, + "cell_variables": code.co_cellvars, + "positional_arguments": code.co_argcount, + "positional_only_arguments": code.co_posonlyargcount, + "keyword_only_arguments": code.co_kwonlyargcount, + "flags": code.co_flags, + } + ) + + +def _code_value(value: Any) -> Any: + if isinstance(value, CodeType): + return {"nested_code_sha256": _code_fingerprint(value)} + if isinstance(value, bytes): + return {"bytes_sha256": hashlib.sha256(value).hexdigest()} + if isinstance(value, Mapping): + return { + str(key): _code_value(item) + for key, item in sorted(value.items(), key=lambda pair: repr(pair[0])) + } + if isinstance(value, (tuple, list)): + return [_code_value(item) for item in value] + if isinstance(value, (set, frozenset)): + return sorted((_code_value(item) for item in value), key=repr) + if value is None or isinstance(value, (bool, int, float, str)): + return value + return {"type": _reference(type(value)), "repr": repr(value)} + + +def _module_identity(module_name: Optional[str]) -> Optional[Mapping[str, Any]]: + if not module_name: + return None + try: + module = importlib.import_module(module_name) + except (ImportError, ModuleNotFoundError): + return {"name": module_name, "content_sha256": None} + module_file = getattr(module, "__file__", None) + if not module_file: + return {"name": module_name, "content_sha256": None} + path = Path(module_file) + try: + digest = hashlib.sha256() + with path.open("rb") as source: + for chunk in iter(lambda: source.read(1024 * 1024), b""): + digest.update(chunk) + except OSError: + return {"name": module_name, "content_sha256": None} + return { + "name": module_name, + "content_sha256": digest.hexdigest(), + } + + +def _fingerprint(value: Any) -> str: + encoded = json.dumps(_json_value(value), sort_keys=True, separators=(",", ":")) + return hashlib.sha256(encoded.encode("utf-8")).hexdigest() diff --git a/setup.py b/setup.py index 57c98070..79f882d9 100644 --- a/setup.py +++ b/setup.py @@ -1,330 +1,313 @@ -# SPDX-License-Identifier: Apache-2.0 -# Copyright (c) 2026 RL-Kernel Contributors - -import importlib.util -import os -import sysconfig -from distutils.errors import CompileError -from distutils.spawn import find_executable -from pathlib import Path - -from setuptools import Extension, find_packages, setup - - -def _load_envs_module(): - envs_path = Path(__file__).with_name("envs.py") - spec = importlib.util.spec_from_file_location("_rl_kernel_envs", envs_path) - if spec is None or spec.loader is None: - raise RuntimeError(f"failed to load environment helpers from {envs_path}") - module = importlib.util.module_from_spec(spec) - spec.loader.exec_module(module) - return module - - -envs = _load_envs_module() - - -def _load_torch_extension_tools(): - try: - import torch - from torch.utils.cpp_extension import BuildExtension, CUDAExtension - except ImportError: - return None, None, None, None - - try: - from torch.utils.cpp_extension import ROCMExtension - except ImportError: - ROCMExtension = None - - return torch, BuildExtension, CUDAExtension, ROCMExtension - - -def _cuda_define_from_env(name: str, macro: str) -> list[str]: - value = os.environ.get(name) - if value is None: - return [] - parsed = int(value) - if parsed <= 0: - raise ValueError(f"{name} must be positive, got {value!r}") - return [f"-D{macro}={parsed}"] - - -def get_extensions(): - torch, _, CUDAExtension, ROCMExtension = _load_torch_extension_tools() - if torch is None: - return [] - - extensions = [] - torch_lib_dir = os.path.join(os.path.dirname(torch.__file__), "lib") - torch_rpath = ["-Wl,-rpath,$ORIGIN/../torch/lib"] - if os.environ.get("KERNEL_ALIGN_DEV_RPATH") == "1": - torch_rpath.append(f"-Wl,-rpath,{torch_lib_dir}") - is_rocm = torch.version.hip is not None - - if is_rocm and ROCMExtension is not None: - extensions.append( - ROCMExtension( - name="rl_engine._C", - sources=[ - "csrc/ops.cpp", - "csrc/fused_logp_kernel.cpp", - ], - extra_compile_args={ - "cxx": ["-O3", "-std=c++17"], - "hipcc": ["-O3", "--use_fast_math", "-Xhipcc", "-compress-all"], - }, - extra_link_args=list(torch_rpath), - ) - ) - elif torch.cuda.is_available(): - cuda_sources = [ - "csrc/ops.cpp", - "csrc/fused_logp_kernel.cu", - "csrc/deterministic_logp_kernel.cu", - "csrc/cuda/attention/prefix_shared_attention.cu", - "csrc/cuda/gemm/det_gemm_kernel.cu", - "csrc/cuda/rmsnorm.cu", - "csrc/cuda/activation.cu", - "csrc/cuda/attention/deterministic_attention.cu", - "csrc/cuda/distributed/deterministic_collective.cu", - ] - - cc_major, cc_minor = torch.cuda.get_device_capability() - enable_sm90 = os.environ.get("KERNEL_ALIGN_FORCE_SM90") == "1" - nvcc_flags = ["-O3", "-Xfatbin", "-compress-all"] - if envs.env_flag(envs.KERNEL_ALIGN_USE_FAST_MATH): - nvcc_flags.append("--use_fast_math") - if not enable_sm90: - # SM90 build emits 90a below; mixing plain compute_90 breaks TMA ptxas. - nvcc_flags.append( - f"-gencode=arch=compute_{cc_major}{cc_minor},code=sm_{cc_major}{cc_minor}" - ) - nvcc_flags.append("--expt-relaxed-constexpr") - nvcc_flags.append("--expt-extended-lambda") - nvcc_flags.extend( - _cuda_define_from_env( - "FUSED_LOGP_TWOPASS_BLOCK_SIZE", - "FUSED_LOGP_TWOPASS_BLOCK_SIZE", - ) - ) - nvcc_flags.extend( - _cuda_define_from_env( - "FUSED_LOGP_ONLINE_BLOCK_SIZE", - "FUSED_LOGP_ONLINE_BLOCK_SIZE", - ) - ) - nvcc_flags.extend( - _cuda_define_from_env( - "FUSED_LOGP_ONLINE_SPARSE_LARGE_VOCAB_BLOCK_SIZE", - "FUSED_LOGP_ONLINE_SPARSE_LARGE_VOCAB_BLOCK_SIZE", - ) - ) - nvcc_flags.extend( - _cuda_define_from_env( - "FUSED_LOGP_ONLINE_LARGE_ROW_BYTES_THRESHOLD", - "FUSED_LOGP_ONLINE_LARGE_ROW_BYTES_THRESHOLD", - ) - ) - nvcc_flags.extend( - _cuda_define_from_env( - "FUSED_LOGP_ONLINE_SPARSE_DENSITY_NUMERATOR", - "FUSED_LOGP_ONLINE_SPARSE_DENSITY_NUMERATOR", - ) - ) - nvcc_flags.extend( - _cuda_define_from_env( - "FUSED_LOGP_ONLINE_SPARSE_DENSITY_DENOMINATOR", - "FUSED_LOGP_ONLINE_SPARSE_DENSITY_DENOMINATOR", - ) - ) - nvcc_flags.extend( - _cuda_define_from_env( - "FUSED_LOGP_ONLINE_MIN_BLOCKS_PER_SM", - "FUSED_LOGP_ONLINE_MIN_BLOCKS_PER_SM", - ) - ) - if envs.env_flag(envs.KERNEL_ALIGN_NCU_LINEINFO): - nvcc_flags.append("-lineinfo") - if os.name == "nt" and envs.env_flag(envs.KERNEL_ALIGN_ALLOW_UNSUPPORTED_MSVC): - nvcc_flags.append("-allow-unsupported-compiler") - nvcc_flags.append("-D_ALLOW_COMPILER_AND_STL_VERSION_MISMATCH") - - cxx_flags = ["-O3", "-std=c++17", "-DKERNEL_ALIGN_WITH_CUDA"] - extra_link_args = list(torch_rpath) - if os.name != "nt": - # CUDA IPC metadata queries use the driver API (cuPointerGetAttribute). - extra_link_args.append("-lcuda") - - sm90_srcs = [ - "csrc/cuda/fused_logp_sm90.cu", - "csrc/cuda/fused_linear_logp_sm90.cu", # TMA + WGMMA fused linear log-prob - "csrc/cuda/batch_invariant_logp_kernel_sm90.cu", # TMA batch-invariant logp - "csrc/cuda/rope_sm90.cu", # RoPE rotate-half apply, gated to SM90 build - "csrc/cuda/embedding_lm_head_sm90.cu", # single-card batch-invariant embedding/lm-head - ] - enable_sm90 = envs.env_flag(envs.KERNEL_ALIGN_FORCE_SM90) - present_sm90 = [s for s in sm90_srcs if os.path.exists(s)] - if enable_sm90 and present_sm90: - tma_arch = f"{cc_major}{cc_minor}a" # WGMMA/TMA require the arch-native 'a' variant - cuda_sources.extend(present_sm90) - nvcc_flags.append(f"-gencode=arch=compute_{tma_arch},code=sm_{tma_arch}") - cxx_flags.append("-DKERNEL_ALIGN_WITH_SM90") - if "-lcuda" not in extra_link_args: - extra_link_args.append("-lcuda") - - # det_gemm SM90 (mma.sync + TMA) path: independent of the fused_logp - # SM90 sources, which currently fail ptxas on CUDA 12.4 (shared::cta in - # the shared tma_utils.cuh). det_gemm uses its own gemm/det_gemm_tma.cuh. - enable_det_gemm_sm90 = os.environ.get("KERNEL_ALIGN_DET_GEMM_SM90") == "1" - if enable_det_gemm_sm90: - tma_arch = f"{cc_major}{cc_minor}a" - arch_flag = f"-gencode=arch=compute_{tma_arch},code=sm_{tma_arch}" - if arch_flag not in nvcc_flags: - nvcc_flags.append(arch_flag) - if "-lcuda" not in extra_link_args: - extra_link_args.append("-lcuda") - nvcc_flags.append("-DRL_KERNEL_ENABLE_SM90") - cxx_flags.append("-DRL_KERNEL_ENABLE_SM90") - - extensions.append( - CUDAExtension( - name="rl_engine._C", - sources=cuda_sources, - include_dirs=[], - extra_compile_args={ - "cxx": cxx_flags, - "nvcc": nvcc_flags, - }, - extra_link_args=extra_link_args, - ) - ) - extensions.extend(_ascend_extensions()) - return extensions - - -def _ascend_extensions(): - """Ascend C (CANN) kernels, built with bisheng. Gated on KERNEL_ALIGN_FORCE_ASCEND=1. - - Follows the official torch_npu cpp_extension_asc pattern: .asc sources - (kernel + host + pybind) are compiled by the CANN bisheng compiler into a - single rl_engine._C_npu extension module. Requires CANN toolkit (bisheng on - PATH or ASCEND_HOME_PATH set) and torch_npu. - """ - if not envs.env_flag(envs.KERNEL_ALIGN_FORCE_ASCEND): - return [] - try: - import torch # noqa: F401 - import torch_npu # noqa: F401 - except ImportError as e: - raise RuntimeError( - "KERNEL_ALIGN_FORCE_ASCEND=1 requires torch and torch_npu to be installed" - ) from e - - asc_srcs = sorted(str(p) for p in Path("csrc/ascend").glob("*.asc")) - if not asc_srcs: - raise RuntimeError("KERNEL_ALIGN_FORCE_ASCEND=1 but no .asc sources under csrc/ascend/") - return [Extension(name="rl_engine._C_npu", sources=asc_srcs, language="asc")] - - -def _bisheng_compile_cmd(ext, ext_fullpath): - """Single-command bisheng build for an Ascend C extension (see op-plugin example).""" - import torch - import torch.utils.cpp_extension as cpp_extension - import torch_npu - - if find_executable("bisheng") is None: - raise RuntimeError( - "bisheng compiler not found on PATH; source the CANN toolkit environment first" - ) - - soc = os.environ.get(envs.KERNEL_ALIGN_ASCEND_ARCH, "dav-2201") # A2/A3; A5: dav-3510 - abi_value = "1" if torch._C._GLIBCXX_USE_CXX11_ABI else "0" - module_name = ext.name.rsplit(".", 1)[-1] - - torch_npu_dir = os.path.dirname(os.path.realpath(torch_npu.__file__)) - ascend_home = os.environ.get("ASCEND_HOME_PATH", "/usr/local/Ascend/ascend-toolkit/latest") - - include_dirs = [ - *cpp_extension.include_paths(), - sysconfig.get_config_var("INCLUDEPY"), - os.path.join(torch_npu_dir, "include"), - os.path.join(torch_npu_dir, "include", "third_party", "acl", "inc"), - os.path.join(ascend_home, "include"), - ] - lib_dirs = [ - sysconfig.get_config_var("LIBDIR"), - os.path.join(os.path.dirname(torch.__file__), "lib"), - os.path.join(torch_npu_dir, "lib"), - os.path.join(ascend_home, "lib64"), - ] - - cmd = [ - "bisheng", - "-x", - "asc", - f"--npu-arch={soc}", - "-shared", - "-fPIC", - "-std=c++17", - "-O2", - f"-D_GLIBCXX_USE_CXX11_ABI={abi_value}", - f"-DTORCH_EXTENSION_NAME={module_name}", - "-lascendcl", - "-ltorch_npu", - "-ltorch", - "-ltorch_cpu", - "-ltorch_python", - "-lc10", - *ext.sources, - "-o", - ext_fullpath, - ] - cmd += [f"-I{d}" for d in include_dirs if d] - cmd += [f"-L{d}" for d in lib_dirs if d] - return cmd - - -def get_cmdclass(): - _, BuildExtension, _, _ = _load_torch_extension_tools() - if BuildExtension is None: - return {} - - class AscendBuildExtension(BuildExtension): - """torch BuildExtension + bisheng path for language="asc" extensions.""" - - def build_extension(self, ext): - if getattr(ext, "language", None) != "asc": - super().build_extension(ext) - return - ext_fullpath = self.get_ext_fullpath(ext.name) - os.makedirs(os.path.dirname(ext_fullpath), exist_ok=True) - try: - self.spawn(_bisheng_compile_cmd(ext, ext_fullpath)) - except Exception as e: - raise CompileError(str(e)) from e - - return {"build_ext": AscendBuildExtension} - - -setup( - name="rl-engine", - version="0.1.0", - packages=find_packages(include=["rl_engine", "rl_engine.*"]), - install_requires=[ - "torch>=2.4.1", - "tabulate", - "numpy", - "accelerate", - "transformers==5.13.1", - ], - ext_modules=get_extensions(), - cmdclass=get_cmdclass(), - extras_require={ - "cuda": ["flashinfer"], - "rocm": ["aiter"], - "vllm": ["vllm>=0.6.0"], - }, - python_requires=">=3.10", - include_package_data=True, - zip_safe=False, -) +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +import importlib.util +import os +import warnings +from pathlib import Path + +from setuptools import find_packages, setup + + +def _load_envs_module(): + envs_path = Path(__file__).with_name("envs.py") + spec = importlib.util.spec_from_file_location("_rl_kernel_envs", envs_path) + if spec is None or spec.loader is None: + raise RuntimeError(f"failed to load environment helpers from {envs_path}") + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +envs = _load_envs_module() + + +def _load_torch_extension_tools(): + try: + import torch + except ModuleNotFoundError as exc: + if exc.name != "torch": + raise + return None, None, None + + from torch.utils.cpp_extension import BuildExtension, CUDAExtension + + # CUDAExtension is also the supported extension entry point for ROCm + # PyTorch builds. BuildExtension dispatches .cu/.hip sources to hipcc when + # torch.version.hip is set. + return torch, BuildExtension, CUDAExtension + + +def _native_extension_required() -> bool: + """Whether the caller explicitly requested a native extension build.""" + return ( + envs.env_flag(envs.RL_KERNEL_REQUIRE_EXT) + or bool(os.environ.get("PYTORCH_ROCM_ARCH", "").strip()) + or bool(os.environ.get("TORCH_CUDA_ARCH_LIST", "").strip()) + or envs.env_flag("FORCE_CUDA") + ) + + +def _cuda_define_from_env(name: str, macro: str) -> list[str]: + value = os.environ.get(name) + if value is None: + return [] + parsed = int(value) + if parsed <= 0: + raise ValueError(f"{name} must be positive, got {value!r}") + return [f"-D{macro}={parsed}"] + + +_ROCM_UNSUPPORTED_NVCC_FLAG_PREFIXES = ( + "-Xfatbin", + "-compress-all", + "-gencode", + "--generate-code", + "--expt-", + "-lineinfo", + "-allow-unsupported-compiler", + "-D_ALLOW_COMPILER_AND_STL_VERSION_MISMATCH", +) +_ROCM_NVCC_FLAGS_WITH_SEPARATE_VALUE = { + "-Xfatbin", + "-gencode", + "--generate-code", +} + + +def _filter_rocm_incompatible_nvcc_flags(flags: list[str]) -> list[str]: + """Remove CUDA-only device compiler flags before BuildExtension calls hipcc.""" + filtered_flags = [] + skip_next = False + for flag in flags: + if skip_next: + skip_next = False + continue + if flag in _ROCM_NVCC_FLAGS_WITH_SEPARATE_VALUE: + skip_next = True + continue + if flag.startswith(_ROCM_UNSUPPORTED_NVCC_FLAG_PREFIXES): + continue + filtered_flags.append(flag) + return filtered_flags + + +def get_extensions(): + torch, _, CUDAExtension = _load_torch_extension_tools() + if torch is None: + message = ( + "PyTorch is unavailable, so rl_engine._C cannot be built. Install a matching " + "CUDA/ROCm PyTorch build first, then run " + "`RL_KERNEL_REQUIRE_EXT=1 python -m pip install --no-build-isolation -e .`." + ) + if _native_extension_required(): + raise RuntimeError(message) + warnings.warn( + f"{message} Continuing with the pure-Python fallback because no native extension " + "was explicitly requested.", + RuntimeWarning, + stacklevel=2, + ) + return [] + + extensions = [] + torch_lib_dir = os.path.join(os.path.dirname(torch.__file__), "lib") + torch_rpath = ["-Wl,-rpath,$ORIGIN/../torch/lib"] + if os.environ.get("KERNEL_ALIGN_DEV_RPATH") == "1": + torch_rpath.append(f"-Wl,-rpath,{torch_lib_dir}") + is_rocm = getattr(torch.version, "hip", None) is not None + + # CUDAExtension is intentionally used for both CUDA and ROCm. On ROCm, + # PyTorch's BuildExtension hipifies CUDA sources and invokes hipcc; it also + # consumes PYTORCH_ROCM_ARCH (one or more ';'-separated gfx targets) to add + # --offload-arch. Do not require a visible GPU when a ROCm target was + # explicitly selected. + no_rocm_arch = not os.environ.get("PYTORCH_ROCM_ARCH", "").strip() + if is_rocm and no_rocm_arch and torch.cuda.device_count() == 0: + raise RuntimeError( + "ROCm builds without a visible GPU require PYTORCH_ROCM_ARCH. " + "Set one or more ';'-separated targets, for example " + "PYTORCH_ROCM_ARCH='gfx942;gfx950'." + ) + + if is_rocm or torch.cuda.is_available(): + cuda_sources = [ + "csrc/ops.cpp", + "csrc/fused_logp_kernel.cu", + "csrc/deterministic_logp_kernel.cu", + "csrc/cuda/gemm/det_gemm_kernel.cu", + "csrc/cuda/rmsnorm.cu", + "csrc/cuda/activation.cu", + "csrc/cuda/attention/deterministic_attention.cu", + "csrc/cuda/distributed/deterministic_collective.cu", + ] + if not is_rocm: + # This source contains NVIDIA PTX (cp.async, ldmatrix, and mma.sync). + # The ROCm dispatcher falls back to PyTorch SDPA for this operator. + cuda_sources.append("csrc/cuda/attention/prefix_shared_attention.cu") + + nvcc_flags = ["-O3", "-Xfatbin", "-compress-all"] + if envs.env_flag(envs.KERNEL_ALIGN_USE_FAST_MATH): + nvcc_flags.append("--use_fast_math") + if not is_rocm: + cc_major, cc_minor = torch.cuda.get_device_capability() + enable_sm90 = os.environ.get("KERNEL_ALIGN_FORCE_SM90") == "1" + if not enable_sm90: + # SM90 build emits 90a below; mixing plain compute_90 breaks TMA ptxas. + nvcc_flags.append( + f"-gencode=arch=compute_{cc_major}{cc_minor},code=sm_{cc_major}{cc_minor}" + ) + nvcc_flags.append("--expt-relaxed-constexpr") + nvcc_flags.append("--expt-extended-lambda") + nvcc_flags.extend( + _cuda_define_from_env( + "FUSED_LOGP_TWOPASS_BLOCK_SIZE", + "FUSED_LOGP_TWOPASS_BLOCK_SIZE", + ) + ) + nvcc_flags.extend( + _cuda_define_from_env( + "FUSED_LOGP_ONLINE_BLOCK_SIZE", + "FUSED_LOGP_ONLINE_BLOCK_SIZE", + ) + ) + nvcc_flags.extend( + _cuda_define_from_env( + "FUSED_LOGP_ONLINE_SPARSE_LARGE_VOCAB_BLOCK_SIZE", + "FUSED_LOGP_ONLINE_SPARSE_LARGE_VOCAB_BLOCK_SIZE", + ) + ) + nvcc_flags.extend( + _cuda_define_from_env( + "FUSED_LOGP_ONLINE_LARGE_ROW_BYTES_THRESHOLD", + "FUSED_LOGP_ONLINE_LARGE_ROW_BYTES_THRESHOLD", + ) + ) + nvcc_flags.extend( + _cuda_define_from_env( + "FUSED_LOGP_ONLINE_SPARSE_DENSITY_NUMERATOR", + "FUSED_LOGP_ONLINE_SPARSE_DENSITY_NUMERATOR", + ) + ) + nvcc_flags.extend( + _cuda_define_from_env( + "FUSED_LOGP_ONLINE_SPARSE_DENSITY_DENOMINATOR", + "FUSED_LOGP_ONLINE_SPARSE_DENSITY_DENOMINATOR", + ) + ) + nvcc_flags.extend( + _cuda_define_from_env( + "FUSED_LOGP_ONLINE_MIN_BLOCKS_PER_SM", + "FUSED_LOGP_ONLINE_MIN_BLOCKS_PER_SM", + ) + ) + if not is_rocm and envs.env_flag(envs.KERNEL_ALIGN_NCU_LINEINFO): + nvcc_flags.append("-lineinfo") + if ( + not is_rocm + and os.name == "nt" + and envs.env_flag(envs.KERNEL_ALIGN_ALLOW_UNSUPPORTED_MSVC) + ): + nvcc_flags.append("-allow-unsupported-compiler") + nvcc_flags.append("-D_ALLOW_COMPILER_AND_STL_VERSION_MISMATCH") + + cxx_flags = ["-O3", "-std=c++17", "-DKERNEL_ALIGN_WITH_CUDA"] + extra_link_args = list(torch_rpath) + if os.name != "nt": + # CUDA IPC metadata queries use the driver API (cuPointerGetAttribute). + extra_link_args.append("-lcuda") + + if not is_rocm: + sm90_srcs = [ + "csrc/cuda/fused_logp_sm90.cu", + "csrc/cuda/fused_linear_logp_sm90.cu", # TMA + WGMMA fused linear log-prob + "csrc/cuda/batch_invariant_logp_kernel_sm90.cu", # TMA batch-invariant logp + "csrc/cuda/rope_sm90.cu", # RoPE rotate-half apply, gated to SM90 build + # Single-card batch-invariant embedding/lm-head. + "csrc/cuda/embedding_lm_head_sm90.cu", + ] + enable_sm90 = envs.env_flag(envs.KERNEL_ALIGN_FORCE_SM90) + present_sm90 = [s for s in sm90_srcs if os.path.exists(s)] + if enable_sm90 and present_sm90: + tma_arch = f"{cc_major}{cc_minor}a" # WGMMA/TMA require the arch-native 'a' variant + cuda_sources.extend(present_sm90) + nvcc_flags.append(f"-gencode=arch=compute_{tma_arch},code=sm_{tma_arch}") + cxx_flags.append("-DKERNEL_ALIGN_WITH_SM90") + if "-lcuda" not in extra_link_args: + extra_link_args.append("-lcuda") + + # det_gemm SM90 (mma.sync + TMA) path: independent of the fused_logp + # SM90 sources, which currently fail ptxas on CUDA 12.4 (shared::cta in + # the shared tma_utils.cuh). det_gemm uses its own gemm/det_gemm_tma.cuh. + enable_det_gemm_sm90 = os.environ.get("KERNEL_ALIGN_DET_GEMM_SM90") == "1" + if enable_det_gemm_sm90: + tma_arch = f"{cc_major}{cc_minor}a" + arch_flag = f"-gencode=arch=compute_{tma_arch},code=sm_{tma_arch}" + if arch_flag not in nvcc_flags: + nvcc_flags.append(arch_flag) + if "-lcuda" not in extra_link_args: + extra_link_args.append("-lcuda") + nvcc_flags.append("-DRL_KERNEL_ENABLE_SM90") + cxx_flags.append("-DRL_KERNEL_ENABLE_SM90") + + if is_rocm: + nvcc_flags = _filter_rocm_incompatible_nvcc_flags(nvcc_flags) + + extensions.append( + CUDAExtension( + name="rl_engine._C", + sources=cuda_sources, + include_dirs=[], + extra_compile_args={ + "cxx": cxx_flags, + "nvcc": nvcc_flags, + }, + extra_link_args=extra_link_args, + ) + ) + + if _native_extension_required() and not extensions: + raise RuntimeError( + "rl_engine._C was requested but no CUDA/ROCm build environment is available. " + "Use a matching GPU-enabled PyTorch build; for a GPU-less ROCm build, set " + "PYTORCH_ROCM_ARCH to the target architecture." + ) + + return extensions + + +def get_cmdclass(): + _, BuildExtension, _ = _load_torch_extension_tools() + if BuildExtension is None: + return {} + return {"build_ext": BuildExtension} + + +setup( + name="rl-engine", + version="0.1.0", + packages=find_packages(include=["rl_engine", "rl_engine.*"]), + install_requires=[ + "torch>=2.4.1", + "tabulate", + "numpy", + "accelerate", + "transformers==5.13.1", + ], + ext_modules=get_extensions(), + cmdclass=get_cmdclass(), + extras_require={ + "cuda": ["flashinfer"], + "rocm": ["aiter"], + "vllm": ["vllm>=0.6.0"], + "drift-viewer": ["Pillow>=10", "PySide6>=6.6"], + }, + entry_points={ + "console_scripts": [ + "rlk-drift-view=rl_engine.alignment.cross_config.drift_viewer:main", + ], + }, + python_requires=">=3.10", + include_package_data=True, + zip_safe=False, +) diff --git a/tests/distributed/test_det_gemm_simulated_tp.py b/tests/distributed/test_det_gemm_simulated_tp.py new file mode 100644 index 00000000..c4824e85 --- /dev/null +++ b/tests/distributed/test_det_gemm_simulated_tp.py @@ -0,0 +1,85 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Simulate TP by splitting GEMM K, without a real process group. + +TP=2 is one BF16 add of two shards (`a+b == b+a`), so it can match TP=1. +TP=8 left-fold is a different parenthesization from the kernel K-tree, so it +must not match. Real NCCL / custom AllReduce is out of scope here. +""" + +from __future__ import annotations + +import pytest +import torch + +from rl_engine.kernels.ops.cuda.matmul import deterministic_gemm + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] < 8, + reason="det_gemm requires CUDA SM80+", +) + +DEV = "cuda" + + +def _rand(*shape): + return torch.randn(*shape, device=DEV, dtype=torch.bfloat16) + + +def _k_shards(a: torch.Tensor, b: torch.Tensor, tp: int) -> list[torch.Tensor]: + k = a.shape[1] + assert k % tp == 0 + width = k // tp + return [ + deterministic_gemm( + a[:, i * width : (i + 1) * width].contiguous(), + b[i * width : (i + 1) * width].contiguous(), + ) + for i in range(tp) + ] + + +def _left_fold(parts: list[torch.Tensor]) -> torch.Tensor: + acc = parts[0] + for part in parts[1:]: + acc = acc + part + return acc + + +def test_simulated_tp2_matches_full(): + # Two shards: AllReduce is a+b, and BF16 add is commutative. + torch.manual_seed(8) + m, k, n = 16, 256, 64 + a, b = _rand(m, k), _rand(k, n) + left, right = _k_shards(a, b, 2) + full = deterministic_gemm(a, b) + assert torch.equal(full, left + right) + assert torch.equal(left + right, right + left) + + +def test_simulated_tp8_left_fold_does_not_match_full(): + # Eight shards left-folded: ((((s0+s1)+s2)+...)+s7) is not the kernel tree + # ((s0+s1)+(s2+s3))+((s4+s5)+(s6+s7)), so this must diverge. + torch.manual_seed(8) + m, k, n = 16, 256, 64 + a, b = _rand(m, k), _rand(k, n) + full = deterministic_gemm(a, b) + folded = _left_fold(_k_shards(a, b, 8)) + n_mismatch = int((full != folded).sum().item()) + assert n_mismatch > 0, "TP=8 left-fold unexpectedly matched TP=1" + + +def test_simulated_tp2_is_batch_invariant(): + torch.manual_seed(9) + k, n = 256, 64 + b = _rand(k, n) + row = _rand(1, k) + big = _rand(32, k) + big[0] = row[0] + + def tp2(x): + left, right = _k_shards(x, b, 2) + return left + right + + assert torch.equal(tp2(row)[0], tp2(big)[0]) + assert torch.equal(tp2(row)[0], deterministic_gemm(big, b)[0]) diff --git a/tests/test_alignment_wrapper_interfaces.py b/tests/test_alignment_wrapper_interfaces.py new file mode 100644 index 00000000..c9fbad70 --- /dev/null +++ b/tests/test_alignment_wrapper_interfaces.py @@ -0,0 +1,159 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Common invocation-surface tests for Attention, FFN, and logprob wrappers.""" + +from __future__ import annotations + +import inspect + +import pytest +import torch + +from rl_engine.alignment.cross_config.operators import selected_logprobs_with_operator +from rl_engine.kernels.logprob_contract import ( + LogprobContract, + MaskSpec, + ReductionSpec, + ShardingSpec, +) +from rl_engine.kernels.ops.pytorch.attention import AttentionAblationOp +from rl_engine.kernels.ops.pytorch.ffn import Qwen3FFNOp +from rl_engine.kernels.ops.pytorch.loss.vocab_parallel_logp import VocabParallelLogprobOp +from rl_engine.kernels.registry import KernelRegistry +from rl_engine.kernels.semantic_registry import OperatorRequirements + + +def _logprob_contract( + *, active: tuple[bool, ...], determinism_scope: str = "cross_tp_bitwise" +) -> LogprobContract: + return LogprobContract( + role="train", + dtype="fp32", + mask=MaskSpec(num_tokens=len(active), active_mask=active), + sharding=ShardingSpec( + tp_rank=0, + tp_world_size=1, + vocab_shard_bounds=((0, 8),), + real_vocab_size=7, + padded_vocab_size=8, + ), + reduction=ReductionSpec(determinism_scope=determinism_scope), + ) + + +@pytest.mark.parametrize( + "wrapper", + [AttentionAblationOp, Qwen3FFNOp, VocabParallelLogprobOp], +) +def test_alignment_wrappers_share_the_deterministic_switch(wrapper): + parameters = inspect.signature(wrapper.__call__).parameters + + assert "deterministic" in parameters + assert parameters["deterministic"].kind is inspect.Parameter.KEYWORD_ONLY + assert parameters["deterministic"].default in {None, True} + + +@pytest.mark.parametrize( + ("wrapper", "backend_id"), + [ + (AttentionAblationOp, "rlkernel.attention.deterministic.v1"), + (Qwen3FFNOp, "rlkernel.ffn.qwen3.deterministic.v1"), + (VocabParallelLogprobOp, "pytorch-vocab-parallel-logp-ws2"), + ], +) +def test_alignment_wrapper_backend_ids_are_stable(wrapper, backend_id): + assert wrapper.backend_id == backend_id + + +@pytest.mark.parametrize("deterministic", [True, False]) +def test_contract_aware_logprob_bridge_accepts_logp_lse_result(deterministic): + active = (True, False, True, True) + contract = _logprob_contract( + active=active, + determinism_scope="cross_tp_bitwise" if deterministic else "fixed_topology", + ) + logits = torch.tensor( + [ + [[1.0, 2.0, 0.0, -1.0, 3.0, 0.5, -0.5, 100.0], [0.0] * 8], + [[2.0, 0.0, 1.0, -2.0, 0.5, 3.0, -1.0, 100.0], [0.0] * 8], + ] + ) + targets = torch.tensor([[4, -100], [5, 0]]) + mask = torch.tensor(active).reshape_as(targets) + + actual = selected_logprobs_with_operator( + VocabParallelLogprobOp(), + logits, + targets, + active_mask=mask, + contract=contract, + num_vocab_tiles=4, + deterministic=deterministic, + ) + + real_logits = logits[..., :7] + safe_targets = targets.masked_fill(~mask, 0) + expected = torch.gather( + torch.log_softmax(real_logits, dim=-1), + -1, + safe_targets.unsqueeze(-1), + ).squeeze(-1) + expected = expected.masked_fill(~mask, 0.0) + torch.testing.assert_close(actual, expected, atol=1e-6, rtol=1e-6) + + +def test_contract_aware_logprob_bridge_rejects_a_different_active_mask(): + contract = _logprob_contract(active=(True, False)) + + with pytest.raises(ValueError, match="active_mask must match"): + selected_logprobs_with_operator( + VocabParallelLogprobOp(), + torch.zeros((2, 8)), + torch.tensor([0, 0]), + active_mask=torch.tensor([True, True]), + contract=contract, + num_vocab_tiles=4, + ) + + +def test_semantic_catalog_exposes_exact_rlkernel_and_native_axes(): + catalog = KernelRegistry().semantic + + for semantic_op, backend_id in ( + ("attention", "rlkernel.attention.deterministic.v1"), + ("ffn", "rlkernel.ffn.qwen3.deterministic.v1"), + ("selected_logprob", "pytorch-vocab-parallel-logp-ws2"), + ): + rlkernel = catalog.backend_descriptor(semantic_op, backend_id) + native = catalog.backend_descriptor(semantic_op, "native") + assert rlkernel is not None + assert rlkernel.determinism_or_alignment_properties["deterministic"] is True + assert rlkernel.is_strictly_observable + assert native is not None + assert native.fallback_policy.value == "runtime_managed" + + +def test_semantic_ffn_and_logprob_wrappers_are_instantiable(): + catalog = KernelRegistry().semantic + session = catalog.session() + + ffn = session.instantiate( + session.resolve( + semantic_op="ffn", + requested_backend="rlkernel.ffn.qwen3.deterministic.v1", + target="training", + requirements=OperatorRequirements(device="cuda", dtype="bfloat16"), + ) + ) + logprob = session.instantiate( + session.resolve( + semantic_op="selected_logprob", + requested_backend="pytorch-vocab-parallel-logp-ws2", + target="training", + requirements=OperatorRequirements(device="cpu", dtype="float32"), + ) + ) + + assert isinstance(ffn, Qwen3FFNOp) + assert isinstance(logprob, VocabParallelLogprobOp) diff --git a/tests/test_attention_ablation.py b/tests/test_attention_ablation.py new file mode 100644 index 00000000..e07201da --- /dev/null +++ b/tests/test_attention_ablation.py @@ -0,0 +1,477 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch + +from rl_engine.kernels.attention_contract import ( + STRICT_ATTENTION_CORE_ID, + STRICT_ATTENTION_SCHEDULE_ID, + AttentionContract, + AttentionContractError, + AttentionDType, + AttentionMode, + AttentionRole, + ReductionSpec, + ShardingSpec, + SplitKVSpec, +) +from rl_engine.kernels.ops.pytorch.attention.ablation import ( + BACKEND_ID, + REFERENCE_BACKEND_ID, + AttentionAblationConfig, + AttentionAblationOp, +) + + +def _contract(*, split_kv: SplitKVSpec | None = None) -> AttentionContract: + sharding = ShardingSpec( + tp_rank=0, + tp_world_size=1, + cp_rank=0, + cp_world_size=1, + global_q_heads=2, + global_kv_heads=1, + local_q_head_start=0, + local_q_heads=2, + local_kv_head_start=0, + local_kv_heads=1, + global_sequence_length=4, + local_sequence_length=4, + global_block_indices=(0,), + global_block_token_starts=(0,), + local_block_offsets=(0, 4), + ) + return AttentionContract( + role=AttentionRole.TRAIN, + mode=AttentionMode.PREFILL, + dtype=AttentionDType.BF16, + batch_size=1, + query_sequence_length=4, + head_dim=4, + causal=True, + causal_offsets=(0,), + sharding=sharding, + reduction=ReductionSpec(), + split_kv=split_kv or SplitKVSpec.disabled(), + ) + + +def _qkv(): + torch.manual_seed(0) + return ( + torch.randn(1, 2, 4, 4, dtype=torch.bfloat16), + torch.randn(1, 1, 4, 4, dtype=torch.bfloat16), + torch.randn(1, 1, 4, 4, dtype=torch.bfloat16), + ) + + +def _cp2_contract() -> AttentionContract: + return AttentionContract( + role=AttentionRole.TRAIN, + mode=AttentionMode.PREFILL, + dtype=AttentionDType.BF16, + batch_size=1, + query_sequence_length=2, + head_dim=4, + causal=True, + causal_offsets=(0,), + sharding=ShardingSpec( + tp_rank=0, + tp_world_size=1, + cp_rank=0, + cp_world_size=2, + global_q_heads=2, + global_kv_heads=1, + local_q_head_start=0, + local_q_heads=2, + local_kv_head_start=0, + local_kv_heads=1, + global_sequence_length=4, + local_sequence_length=2, + global_block_indices=(0,), + global_block_token_starts=(0,), + local_block_offsets=(0, 2), + ), + reduction=ReductionSpec(), + split_kv=SplitKVSpec.disabled(), + ) + + +def test_attention_wrapper_has_unified_result_and_provenance(): + q, k, v = _qkv() + result = AttentionAblationOp()(q, k, v, contract=_contract()) + + assert result.backend_id == REFERENCE_BACKEND_ID + assert result.deterministic + assert result.out.shape == q.shape + assert result.lse is not None + assert result.lse.dtype is torch.float32 + assert result.provenance["semantic_operator"] == "attention" + assert result.provenance["split_kv"]["mode"] == "disabled" + assert result.provenance["strict_core_id"] == STRICT_ATTENTION_CORE_ID + assert result.provenance["strict_schedule"] == STRICT_ATTENTION_SCHEDULE_ID + assert result.readback()["out_shape"] == list(q.shape) + + +def test_attention_wrapper_supports_explicit_injected_backend(): + q, k, v = _qkv() + + class FakeBackend: + backend_id = "test.attention.backend" + + def forward_with_lse(self, q, k, v, *, causal, scale): + del k, v, causal, scale + return q.clone(), torch.zeros(q.shape[:3], dtype=torch.float32) + + result = AttentionAblationOp()( + q, + k, + v, + contract=_contract(), + backend=FakeBackend(), + deterministic=False, + ) + assert result.backend_id == "test.attention.backend" + assert torch.equal(result.out, q) + + +def test_wrapper_owned_deterministic_core_does_not_require_external_provenance(): + q, k, v = _qkv() + + class WrapperOwnedCore: + backend_id = BACKEND_ID + + def __call__(self, q, k, v, *, causal, scale): + del k, v, causal, scale + return q.clone(), torch.zeros(q.shape[:3], dtype=torch.float32) + + result = AttentionAblationOp()( + q, + k, + v, + contract=_contract(), + backend=WrapperOwnedCore(), + ) + + assert result.backend_id == BACKEND_ID + assert result.deterministic is True + + +def test_cp_production_configuration_fails_closed_without_ag_rs_backend(): + q, k, v = _qkv() + with pytest.raises(AttentionContractError, match="injected AG/RS backend"): + AttentionAblationOp(communication_backend="self_owned_cuda_ag_rs")( + q[:, :, :2], k[:, :, :2], v[:, :, :2], contract=_cp2_contract() + ) + + +def test_cp_backend_must_explicitly_declare_cp_world_size(): + q, k, v = _qkv() + + class KwargsOnlyBackend: + backend_id = "test.kwargs_only" + core_id = STRICT_ATTENTION_CORE_ID + strict_schedule = STRICT_ATTENTION_SCHEDULE_ID + + def __call__(self, q, k, v, **kwargs): + del k, v, kwargs + return q, torch.zeros(q.shape[:3], dtype=torch.float32) + + with pytest.raises(AttentionContractError, match="explicitly accepts cp_world_size"): + AttentionAblationOp( + cp_backend=KwargsOnlyBackend(), + communication_backend="cuda_ag_rs", + )( + q[:, :, :2], + k[:, :, :2], + v[:, :, :2], + contract=_cp2_contract(), + ) + + +@pytest.mark.parametrize("communication_backend", ["cuda_ag_rs", "rccl_ag_rs"]) +def test_cp_wrapper_accepts_exact_platform_vendor_core(communication_backend): + q, k, v = _qkv() + + class VendorCPBackend: + backend_id = "vendor.strict_attention" + core_id = "vendor.strict_core.v1" + strict_schedule = "vendor.strict_schedule.v1" + + def __call__(self, q, k, v, *, causal, scale, cp_world_size): + del k, v, causal, scale + assert cp_world_size == 2 + return SimpleNamespace( + out=q.clone(), + lse=torch.zeros(q.shape[:3], dtype=torch.float32), + provenance={ + "strict_core_id": self.core_id, + "strict_schedule": self.strict_schedule, + "actual_backend": self.backend_id, + "communication_backend": communication_backend, + "production_ready": True, + "native_attention_arithmetic": True, + "fallback": False, + "reference_only": False, + }, + ) + + result = AttentionAblationOp( + cp_backend=VendorCPBackend(), + communication_backend=communication_backend, + )( + q[:, :, :2], + k[:, :, :2], + v[:, :, :2], + contract=_cp2_contract(), + config=AttentionAblationConfig( + strict_core_id=VendorCPBackend.core_id, + strict_schedule=VendorCPBackend.strict_schedule, + ), + ) + + assert result.provenance["actual_backend"] == VendorCPBackend.backend_id + assert result.provenance["communication_backend"] == communication_backend + assert result.provenance["native_attention_arithmetic"] is True + + +def test_cp_production_wrapper_preserves_runtime_backend_provenance(): + q, k, v = _qkv() + + class StrictCPBackend: + backend_id = "injected.strict_cp_backend" + core_id = STRICT_ATTENTION_CORE_ID + strict_schedule = STRICT_ATTENTION_SCHEDULE_ID + + def __call__(self, q, k, v, *, causal, scale): + del k, v, causal, scale + return SimpleNamespace( + out=q.clone(), + lse=torch.zeros(q.shape[:3], dtype=torch.float32), + provenance={ + "strict_core_id": STRICT_ATTENTION_CORE_ID, + "strict_schedule": STRICT_ATTENTION_SCHEDULE_ID, + "actual_backend": self.backend_id, + "communication_backend": "self_owned_cuda_ag_rs", + "production_ready": True, + "native_attention_arithmetic": False, + "fallback": False, + "reference_only": False, + }, + ) + + result = AttentionAblationOp( + cp_backend=StrictCPBackend(), + communication_backend="self_owned_cuda_ag_rs", + )( + q, + k, + v, + contract=_contract(), + backend=StrictCPBackend(), + ) + assert result.provenance["actual_backend"] == StrictCPBackend.backend_id + assert result.provenance["communication_backend"] == "self_owned_cuda_ag_rs" + assert result.provenance["production_ready"] is True + + +@pytest.mark.parametrize( + ("missing_field", "replacement"), + [ + ("production_ready", None), + ("fallback", True), + ("reference_only", True), + ], +) +def test_vendor_production_core_fails_closed_without_exact_provenance(missing_field, replacement): + q, k, v = _qkv() + + class VendorCore: + backend_id = "vendor.strict_attention" + core_id = "vendor.strict_core.v1" + strict_schedule = "vendor.strict_schedule.v1" + + def __call__(self, q, k, v, *, causal, scale): + del k, v, causal, scale + provenance = { + "strict_core_id": self.core_id, + "strict_schedule": self.strict_schedule, + "actual_backend": self.backend_id, + "production_ready": True, + "native_attention_arithmetic": True, + "fallback": False, + "reference_only": False, + } + if replacement is None: + del provenance[missing_field] + else: + provenance[missing_field] = replacement + return SimpleNamespace( + out=q.clone(), + lse=torch.zeros(q.shape[:3], dtype=torch.float32), + provenance=provenance, + ) + + with pytest.raises(AttentionContractError, match="runtime provenance"): + AttentionAblationOp()( + q, + k, + v, + contract=_contract(), + backend=VendorCore(), + config=AttentionAblationConfig( + strict_core_id=VendorCore.core_id, + strict_schedule=VendorCore.strict_schedule, + ), + ) + + +@pytest.mark.parametrize( + "split_kv", + [SplitKVSpec.auto(strict_consistency=False), SplitKVSpec.fixed(2)], +) +def test_deterministic_attention_requires_split_kv_disabled(split_kv): + q, k, v = _qkv() + contract = _contract(split_kv=split_kv) + with pytest.raises(AttentionContractError, match="Split-KV to be disabled"): + AttentionAblationOp()(q, k, v, contract=contract) + + +def test_explicit_cp_callable_cannot_bypass_ag_rs_requirement(): + q, k, v = _qkv() + + class VendorCPBackend: + backend_id = "vendor.strict_attention" + core_id = "vendor.strict_core.v1" + strict_schedule = "vendor.strict_schedule.v1" + + def __call__(self, q, k, v, *, causal, scale, cp_world_size): + del k, v, causal, scale, cp_world_size + return q, torch.zeros(q.shape[:3], dtype=torch.float32) + + with pytest.raises(AttentionContractError, match="requires an explicit CUDA AG/RS"): + AttentionAblationOp()( + q[:, :, :2], + k[:, :, :2], + v[:, :, :2], + contract=_cp2_contract(), + backend=VendorCPBackend(), + config=AttentionAblationConfig( + strict_core_id=VendorCPBackend.core_id, + strict_schedule=VendorCPBackend.strict_schedule, + ), + ) + + +def test_vendor_core_actual_backend_must_match_selected_backend(): + q, k, v = _qkv() + + class VendorCore: + backend_id = "vendor.strict_attention" + core_id = "vendor.strict_core.v1" + strict_schedule = "vendor.strict_schedule.v1" + + def __call__(self, q, k, v, *, causal, scale): + del k, v, causal, scale + return SimpleNamespace( + out=q.clone(), + lse=torch.zeros(q.shape[:3], dtype=torch.float32), + provenance={ + "strict_core_id": self.core_id, + "strict_schedule": self.strict_schedule, + "actual_backend": "vendor.other_attention", + "production_ready": True, + "native_attention_arithmetic": True, + "fallback": False, + "reference_only": False, + }, + ) + + with pytest.raises(AttentionContractError, match="actual_backend"): + AttentionAblationOp()( + q, + k, + v, + contract=_contract(), + backend=VendorCore(), + config=AttentionAblationConfig( + strict_core_id=VendorCore.core_id, + strict_schedule=VendorCore.strict_schedule, + ), + ) + + +def test_deterministic_native_backend_requires_explicit_native_callable(): + q, k, v = _qkv() + with pytest.raises(AttentionContractError, match="native Attention backend"): + AttentionAblationOp()(q, k, v, contract=_contract(), backend="native") + + +def test_attention_wrapper_can_return_dq_dk_dv_from_reference_backend(): + q, k, v = _qkv() + result = AttentionAblationOp()( + q, + k, + v, + contract=_contract(), + return_gradients=True, + dout=torch.ones_like(q), + ) + + assert result.dq is not None and result.dq.shape == q.shape + assert result.dk is not None and result.dk.shape == k.shape + assert result.dv is not None and result.dv.shape == v.shape + + +def test_attention_backend_is_registered_for_pr230_semantic_resolution(): + from rl_engine.kernels.registry import kernel_registry + from rl_engine.kernels.semantic_registry import OperatorRequirements + + session = kernel_registry.semantic.session() + resolution = session.resolve( + semantic_op="attention", + requested_backend=BACKEND_ID, + target="training", + requirements=OperatorRequirements( + device="cpu", + dtype="bfloat16", + topology={"world_size": 1, "tensor_parallel_size": 1, "context_parallel_size": 1}, + alignment_properties={"deterministic": True}, + ), + ) + instance = session.instantiate(resolution) + assert isinstance(instance, AttentionAblationOp) + provenance = session.instance_provenance(resolution, instance) + assert provenance.backend_id == BACKEND_ID + + +def test_strict_wrapper_is_bitwise_invariant_to_batch_shape(): + q, k, v = _qkv() + noise_q, noise_k, noise_v = _qkv() + contract = _contract() + batch_contract = AttentionContract( + role=contract.role, + mode=contract.mode, + dtype=contract.dtype, + batch_size=2, + query_sequence_length=contract.query_sequence_length, + head_dim=contract.head_dim, + causal=contract.causal, + causal_offsets=(0, 0), + sharding=contract.sharding, + reduction=contract.reduction, + split_kv=contract.split_kv, + ) + single = AttentionAblationOp()(q, k, v, contract=contract) + batched = AttentionAblationOp()( + torch.cat((q, noise_q), dim=0), + torch.cat((k, noise_k), dim=0), + torch.cat((v, noise_v), dim=0), + contract=batch_contract, + ) + assert torch.equal(single.out[0], batched.out[0]) + assert torch.equal(single.lse[0], batched.lse[0]) diff --git a/tests/test_attention_contract.py b/tests/test_attention_contract.py new file mode 100644 index 00000000..350545a1 --- /dev/null +++ b/tests/test_attention_contract.py @@ -0,0 +1,683 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""WS2 Attention CP contract and contract-aware dispatch tests (issue #235).""" + +from __future__ import annotations + +import json +from dataclasses import replace + +import pytest + +from rl_engine.kernels.attention_contract import ( + AttentionBackendCapability, + AttentionContract, + AttentionContractError, + AttentionDType, + AttentionMode, + AttentionRole, + KVCacheSpec, + ReductionSpec, + RoPEFusionBoundary, + RoPESpec, + ShardingSpec, + SplitKVExecutionPlan, + SplitKVRuntimePlanSet, + SplitKVSpec, + build_split_kv_runtime_plan_set, + validate_split_kv_alignment, + validate_split_kv_plan_set_alignment, +) +from rl_engine.kernels.registry import KernelRegistry, OpBackend + + +def _sharding( + *, + tp_rank: int = 0, + tp_world_size: int = 2, + cp_rank: int = 0, + cp_world_size: int = 2, + global_sequence_length: int = 4096, + local_sequence_length: int = 2048, + global_block_indices: tuple[int, ...] = (0,), + global_block_token_starts: tuple[int, ...] = (0,), + local_block_offsets: tuple[int, ...] = (0, 2048), + packed_sequence_offsets: tuple[int, ...] | None = None, +) -> ShardingSpec: + local_q_heads = 32 // tp_world_size + local_kv_heads = 8 // tp_world_size + return ShardingSpec( + tp_rank=tp_rank, + tp_world_size=tp_world_size, + cp_rank=cp_rank, + cp_world_size=cp_world_size, + global_q_heads=32, + global_kv_heads=8, + local_q_head_start=tp_rank * local_q_heads, + local_q_heads=local_q_heads, + local_kv_head_start=tp_rank * local_kv_heads, + local_kv_heads=local_kv_heads, + global_sequence_length=global_sequence_length, + local_sequence_length=local_sequence_length, + global_block_indices=global_block_indices, + global_block_token_starts=global_block_token_starts, + local_block_offsets=local_block_offsets, + packed_sequence_offsets=packed_sequence_offsets, + ) + + +def _contract( + *, + role: str = "infer", + mode: str = "prefill", + sharding: ShardingSpec | None = None, + kv_cache: KVCacheSpec | None = None, + causal_offsets: tuple[int, ...] = (0,), + batch_size: int = 1, + query_sequence_length: int | None = None, + rope: RoPESpec | None = None, +) -> AttentionContract: + resolved_sharding = sharding or _sharding() + return AttentionContract( + role=role, + mode=mode, + dtype="bf16", + batch_size=batch_size, + query_sequence_length=( + query_sequence_length + if query_sequence_length is not None + else (1 if mode == "decode" else resolved_sharding.local_sequence_length) + ), + head_dim=128, + causal=True, + causal_offsets=causal_offsets, + sharding=resolved_sharding, + reduction=ReductionSpec(), + kv_cache=kv_cache, + rope=rope, + ) + + +def _declared_cp_backend() -> AttentionBackendCapability: + return AttentionBackendCapability( + backend_id="test-deterministic-cp-attention", + roles=frozenset({AttentionRole.TRAIN, AttentionRole.INFER}), + modes=frozenset( + {AttentionMode.PREFILL, AttentionMode.CHUNKED_PREFILL, AttentionMode.DECODE} + ), + dtypes=frozenset({AttentionDType.BF16}), + tp_world_sizes=(2,), + cp_world_sizes=(1, 2), + exports_attention_lse=True, + deterministic_cp_merge=True, + supports_packed_varlen=True, + supports_kv_cache=True, + supports_split_kv_fixed=True, + reports_actual_split_kv_plan=True, + implementation_kind="deterministic", + ) + + +def test_qwen3_tp2_cp2_contract_is_representable_and_serializable(): + contract = _contract() + + assert contract.sharding.local_q_heads == 16 + assert contract.sharding.local_kv_heads == 4 + assert contract.sharding.tp_world_size == 2 + assert contract.sharding.cp_world_size == 2 + assert contract.reduction.acc_dtype is AttentionDType.FP32 + assert contract.to_dict()["reduction"] == { + "merge": "online_softmax_lse", + "acc_dtype": "fp32", + "order": "global_block_index", + "downcast_at": "final_write", + "engine": "in_op_reference", + } + json.dumps(contract.to_dict()) + + +def test_rope_metadata_is_part_of_attention_contract_provenance(): + rope = RoPESpec( + q_state="post_rope", + k_state="post_rope", + k_cache_state="post_rope", + theta=1.0e6, + rotary_dim=128, + query_position_offsets=(0,), + key_position_offsets=(0,), + cast_at="after_rope", + output_dtype="bf16", + fusion_boundary="unfused_rope_attention", + ) + + contract = _contract(rope=rope) + payload = contract.to_dict() + + assert payload["rope"] == { + "q_state": "post_rope", + "k_state": "post_rope", + "k_cache_state": "post_rope", + "theta": 1.0e6, + "rotary_dim": 128, + "rope_scaling": None, + "position_ids": None, + "query_position_offsets": [0], + "key_position_offsets": [0], + "cast_at": "after_rope", + "output_dtype": "bf16", + "fusion_boundary": "unfused_rope_attention", + } + json.dumps(payload) + + +def test_rope_position_metadata_is_validated_against_contract_shape(): + with pytest.raises(AttentionContractError, match="rotary_dim=256"): + _contract(rope=RoPESpec(rotary_dim=256)) + + with pytest.raises(AttentionContractError, match="query_position_offsets"): + _contract(batch_size=2, causal_offsets=(0, 0), rope=RoPESpec(query_position_offsets=(0,))) + + with pytest.raises(AttentionContractError, match="position_ids"): + _contract(rope=RoPESpec(position_ids=(0, 1, 2))) + + valid = _contract(rope=RoPESpec(position_ids=tuple(range(2048)))) + assert valid.rope is not None + assert valid.rope.position_ids == tuple(range(2048)) + + +@pytest.mark.parametrize( + ("field", "value", "message"), + [ + ("tp_rank", 2, "tp_rank=2"), + ("cp_rank", 2, "cp_rank=2"), + ("global_block_indices", (), "must not be empty"), + ("global_block_indices", (1, 0), "strictly increasing"), + ], +) +def test_invalid_rank_and_cp_order_metadata_fail_loudly(field, value, message): + values = { + "tp_rank": 0, + "tp_world_size": 2, + "cp_rank": 0, + "cp_world_size": 2, + "global_q_heads": 32, + "global_kv_heads": 8, + "local_q_head_start": 0, + "local_q_heads": 16, + "local_kv_head_start": 0, + "local_kv_heads": 4, + "global_sequence_length": 4096, + "local_sequence_length": 2048, + "global_block_indices": (0,), + "global_block_token_starts": (0,), + "local_block_offsets": (0, 2048), + } + values[field] = value + + with pytest.raises(AttentionContractError, match=message): + ShardingSpec(**values) + + +def test_tp_local_heads_must_preserve_global_gqa_mapping(): + with pytest.raises(AttentionContractError, match="local TP head counts"): + replace(_sharding(), local_q_heads=7) + + with pytest.raises(AttentionContractError, match="head starts"): + replace(_sharding(tp_rank=1), local_q_head_start=0) + + +def test_sequence_range_and_packed_offsets_are_validated(): + with pytest.raises(AttentionContractError, match="exceeds global_sequence_length"): + _sharding(global_block_token_starts=(4000,)) + + with pytest.raises(AttentionContractError, match="final packed_sequence_offsets"): + _sharding(packed_sequence_offsets=(0, 512)) + + sharding = _sharding(packed_sequence_offsets=(0, 512, 2048)) + assert sharding.packed_sequence_offsets == (0, 512, 2048) + + +def test_non_contiguous_cp_blocks_have_explicit_global_and_local_offsets(): + sharding = _sharding( + global_block_indices=(0, 3), + global_block_token_starts=(0, 3072), + local_block_offsets=(0, 1024, 2048), + ) + + assert sharding.global_block_indices == (0, 3) + assert sharding.global_block_token_starts == (0, 3072) + assert sharding.local_block_offsets == (0, 1024, 2048) + + with pytest.raises(AttentionContractError, match="non-overlapping and ordered"): + _sharding( + global_block_indices=(0, 1), + global_block_token_starts=(0, 512), + local_block_offsets=(0, 1024, 2048), + ) + + +def test_reduction_requires_fp32_accumulation(): + with pytest.raises(AttentionContractError, match="must be fp32"): + ReductionSpec(acc_dtype="bf16") + + +def test_split_kv_policy_is_a_first_class_strict_contract(): + contract = _contract() + assert contract.to_dict()["split_kv"] == { + "mode": "disabled", + "fixed_split_size": None, + "strict_consistency": True, + } + + fixed = replace(contract, split_kv=SplitKVSpec.fixed(128)) + assert fixed.to_dict()["split_kv"]["mode"] == "fixed" + assert fixed.to_dict()["split_kv"]["fixed_split_size"] == 128 + + with pytest.raises(AttentionContractError, match="auto Split-KV"): + SplitKVSpec.auto(strict_consistency=True) + + +def test_split_kv_execution_plan_records_actual_logical_schedule(): + plan = SplitKVSpec.fixed(4).resolve(10, backend="training-reference") + + assert plan.actual_split_count == 3 + assert plan.to_dict()["actual_split_boundaries"] == [[0, 4], [4, 8], [8, 10]] + assert plan.to_dict()["split_kv_merge_order"] == "global_block_index" + assert plan.to_dict()["split_kv_accum_dtype"] == "fp32" + assert plan.to_dict()["split_kv_downcast_at"] == "final_write" + + +def test_strict_split_kv_alignment_rejects_unknown_or_mismatched_actual_plan(): + training = SplitKVSpec.fixed(4).resolve(10, backend="training") + rollout = SplitKVSpec.fixed(4).resolve(10, backend="rollout") + validate_split_kv_alignment(training, rollout) + + unknown = SplitKVSpec.auto().resolve(10, backend="rollout") + with pytest.raises(AttentionContractError, match="actual runtime plans"): + validate_split_kv_alignment(training, unknown) + + mismatched = SplitKVSpec.fixed(5).resolve(10, backend="rollout") + with pytest.raises(AttentionContractError, match="differ"): + validate_split_kv_alignment(training, mismatched) + + with pytest.raises(AttentionContractError, match="contiguous"): + SplitKVExecutionPlan( + requested_mode="fixed", + requested_split_size=4, + actual_mode="fixed", + actual_split_size=4, + boundaries=((0, 4), (5, 10)), + ) + + +def test_complete_split_kv_plan_set_covers_batch_tp_cp_and_owner_coordinates(): + plan_set = build_split_kv_runtime_plan_set( + (8, 10), + tp_world_size=2, + cp_world_size=2, + split_kv=SplitKVSpec.fixed(2), + backend="training-reference", + ) + + assert len(plan_set.entries) == 16 + assert plan_set.to_dict()["coverage"] == ("complete_batch_tp_cp_owner_cartesian_product") + assert { + tuple(entry["expected_kv_range"]) + for entry in plan_set.to_dict()["entries"] + if entry["batch_index"] == 0 + } == {(0, 4), (4, 8)} + + +def test_split_kv_plan_set_alignment_rejects_missing_and_mismatched_rank_plans(): + training = build_split_kv_runtime_plan_set( + (8,), + tp_world_size=2, + cp_world_size=2, + split_kv=SplitKVSpec.fixed(2), + backend="training", + ) + rollout = build_split_kv_runtime_plan_set( + (8,), + tp_world_size=2, + cp_world_size=2, + split_kv=SplitKVSpec.fixed(2), + backend="rollout", + ) + validate_split_kv_plan_set_alignment(training, rollout) + + with pytest.raises(AttentionContractError, match="coordinate coverage is incomplete"): + SplitKVRuntimePlanSet( + batch_size=training.batch_size, + tp_world_size=training.tp_world_size, + cp_world_size=training.cp_world_size, + total_kv_tokens=training.total_kv_tokens, + entries=training.entries[:-1], + ) + + mismatched = build_split_kv_runtime_plan_set( + (8,), + tp_world_size=2, + cp_world_size=2, + split_kv=SplitKVSpec.fixed(1), + backend="rollout", + ) + with pytest.raises(AttentionContractError, match="plan differs"): + validate_split_kv_plan_set_alignment(training, mismatched) + + +def test_backend_must_support_policy_and_actual_plan_provenance(): + fixed = replace(_contract(), split_kv=SplitKVSpec.fixed(128)) + capability = replace( + _declared_cp_backend(), + supports_split_kv_fixed=False, + reports_actual_split_kv_plan=False, + ) + + assert capability.incompatibilities(fixed)[-2:] == ( + "Split-KV policy=fixed is unsupported", + "actual Split-KV execution-plan provenance is unsupported", + ) + + +def test_causal_attention_requires_explicit_offset(): + contract = _contract() + with pytest.raises(AttentionContractError, match="causal_offsets are required"): + replace(contract, causal_offsets=None) + + +def test_full_prefill_query_length_must_match_local_sequence_length(): + with pytest.raises(AttentionContractError, match="prefill query_sequence_length must equal"): + _contract(mode="prefill", query_sequence_length=1024) + + chunked = _contract(mode="chunked_prefill", query_sequence_length=512) + decode = _contract( + mode="decode", + query_sequence_length=1, + kv_cache=KVCacheSpec( + cache_positions=(16,), + kv_seq_lens=(17,), + block_table=((0, 1),), + global_token_positions=tuple(range(17)), + page_size=16, + ), + ) + assert chunked.query_sequence_length == 512 + assert decode.query_sequence_length == 1 + + +def test_decode_requires_complete_kv_cache_identity(): + with pytest.raises(AttentionContractError, match="kv_cache metadata is required"): + _contract(mode="decode") + + cache = KVCacheSpec( + cache_positions=(16,), + kv_seq_lens=(17,), + block_table=((0, 1, -1),), + global_token_positions=tuple(range(17)), + page_size=16, + prefix_cache_enabled=True, + prefix_cache_key="prefix:sample-0", + ) + contract = _contract(mode="decode", kv_cache=cache) + assert contract.to_dict()["kv_cache"]["block_table"] == [[0, 1, -1]] + + +def test_prefix_cache_key_is_required_only_when_prefix_cache_is_enabled(): + with pytest.raises(AttentionContractError, match="prefix_cache_key is required"): + KVCacheSpec( + cache_positions=(16,), + kv_seq_lens=(17,), + block_table=((0, 1),), + global_token_positions=tuple(range(17)), + page_size=16, + prefix_cache_enabled=True, + ) + + +def test_cache_positions_must_match_kv_sequence_count(): + with pytest.raises(AttentionContractError, match="one entry per kv_seq_lens"): + KVCacheSpec( + cache_positions=(1,), + kv_seq_lens=(2, 2), + block_table=((0,), (1,)), + global_token_positions=(0, 1, 0, 1), + page_size=2, + ) + + +def test_cache_position_must_match_terminal_global_token_position(): + with pytest.raises(AttentionContractError, match="terminal global token position"): + KVCacheSpec( + cache_positions=(999,), + kv_seq_lens=(17,), + block_table=((0, 1),), + global_token_positions=tuple(range(17)), + page_size=16, + ) + + +@pytest.mark.parametrize("positions", [(7, 6), (7, 7)]) +def test_kv_cache_positions_must_be_strictly_increasing_per_sequence(positions): + with pytest.raises(AttentionContractError, match="strictly increasing"): + KVCacheSpec( + cache_positions=(7,), + kv_seq_lens=(2,), + block_table=((0,),), + global_token_positions=positions, + page_size=2, + ) + + +@pytest.mark.parametrize( + ("block_table", "message"), + [ + ((0, -1, 1), "padding must be trailing"), + ((0, 0, -1), "duplicate active page ids"), + ((0, -1, -1), "active page count"), + ], +) +def test_kv_cache_block_table_page_mapping_is_validated(block_table, message): + with pytest.raises(AttentionContractError, match=message): + KVCacheSpec( + cache_positions=(16,), + kv_seq_lens=(17,), + block_table=(block_table,), + global_token_positions=tuple(range(17)), + page_size=16, + ) + + +def test_prefix_pages_may_be_shared_across_sequences(): + cache = KVCacheSpec( + cache_positions=(1, 1), + kv_seq_lens=(2, 2), + block_table=((3,), (3,)), + global_token_positions=(0, 1, 0, 1), + page_size=2, + prefix_cache_enabled=True, + prefix_cache_key="shared-prefix", + shared_prefix_page_count=1, + ) + + assert cache.block_table == ((3,), (3,)) + assert cache.shared_prefix_page_count == 1 + + +def test_non_prefix_cache_rejects_cross_sequence_page_sharing(): + with pytest.raises(AttentionContractError, match="only when declared"): + KVCacheSpec( + cache_positions=(1, 1), + kv_seq_lens=(2, 2), + block_table=((3,), (3,)), + global_token_positions=(0, 1, 0, 1), + page_size=2, + prefix_cache_enabled=False, + ) + + +def test_prefix_cache_requires_explicit_shared_page_count(): + with pytest.raises(AttentionContractError, match="only when declared"): + KVCacheSpec( + cache_positions=(1, 1), + kv_seq_lens=(2, 2), + block_table=((3,), (3,)), + global_token_positions=(0, 1, 0, 1), + page_size=2, + prefix_cache_enabled=True, + prefix_cache_key="shared-prefix", + shared_prefix_page_count=0, + ) + + +def test_prefix_cache_rejects_shared_writable_suffix_pages(): + with pytest.raises(AttentionContractError, match="only when declared"): + KVCacheSpec( + cache_positions=(3, 3), + kv_seq_lens=(4, 4), + block_table=((3, 4), (3, 4)), + global_token_positions=(0, 1, 2, 3, 0, 1, 2, 3), + page_size=2, + prefix_cache_enabled=True, + prefix_cache_key="shared-prefix", + shared_prefix_page_count=1, + ) + + +def test_shared_prefix_identity_must_match_pages_and_positions(): + with pytest.raises(AttentionContractError, match="page ids must match"): + KVCacheSpec( + cache_positions=(3, 3), + kv_seq_lens=(4, 4), + block_table=((3, 4), (5, 6)), + global_token_positions=(0, 1, 2, 3, 0, 1, 2, 3), + page_size=2, + prefix_cache_enabled=True, + prefix_cache_key="shared-prefix", + shared_prefix_page_count=1, + ) + + with pytest.raises(AttentionContractError, match="token positions must match"): + KVCacheSpec( + cache_positions=(3, 13), + kv_seq_lens=(4, 4), + block_table=((3, 4), (3, 5)), + global_token_positions=(0, 1, 2, 3, 10, 11, 12, 13), + page_size=2, + prefix_cache_enabled=True, + prefix_cache_key="shared-prefix", + shared_prefix_page_count=1, + ) + + +def test_shared_prefix_pages_must_be_fully_populated(): + with pytest.raises(AttentionContractError, match="fully populated and read-only"): + KVCacheSpec( + cache_positions=(0, 0), + kv_seq_lens=(1, 1), + block_table=((3,), (3,)), + global_token_positions=(0, 0), + page_size=2, + prefix_cache_enabled=True, + prefix_cache_key="partial-prefix-page", + shared_prefix_page_count=1, + ) + + +def test_current_ws1_backend_rejects_strict_cp_contract_without_fallback(): + registry = KernelRegistry() + platform = registry._platform() + registry._priority_map[platform]["ws2_attention"] = [OpBackend.PYTORCH_NATIVE_ATTENTION] + + with pytest.raises(RuntimeError) as exc_info: + registry.get_attention_op(_contract()) + + message = str(exc_info.value) + assert "CP=2 is unsupported" in message + assert "attention-domain LSE export is unsupported" in message + assert "deterministic CP (out, lse) merge is unsupported" in message + + +def test_undeclared_backend_capability_is_never_selected(): + registry = KernelRegistry() + platform = registry._platform() + registry._priority_map[platform]["ws2_attention"] = [OpBackend.PYTORCH_ATTN] + + with pytest.raises(RuntimeError, match="no AttentionBackendCapability declared"): + registry.get_attention_op(_contract()) + + +def test_declared_compatible_backend_resolves_and_records_provenance(): + registry = KernelRegistry() + platform = registry._platform() + registry._priority_map[platform]["ws2_attention"] = [OpBackend.PYTORCH_NATIVE_ATTENTION] + registry._attention_capabilities[OpBackend.PYTORCH_NATIVE_ATTENTION] = _declared_cp_backend() + + result = registry.get_attention_op(_contract(), requested_backend="deterministic") + + assert result.op is not None + assert result.capability.backend_id == "test-deterministic-cp-attention" + assert result.provenance["requested_backend"] == "deterministic" + assert result.provenance["actual_backend"] == "test-deterministic-cp-attention" + assert result.provenance["fallback"] is False + assert result.provenance["contract"]["sharding"]["tp_world_size"] == 2 + assert result.provenance["contract"]["sharding"]["cp_world_size"] == 2 + json.dumps(result.provenance) + + +def test_requested_stable_backend_id_is_enforced(): + registry = KernelRegistry() + platform = registry._platform() + registry._priority_map[platform]["ws2_attention"] = [OpBackend.PYTORCH_NATIVE_ATTENTION] + registry._attention_capabilities[OpBackend.PYTORCH_NATIVE_ATTENTION] = _declared_cp_backend() + + with pytest.raises(RuntimeError, match="does not match requested_backend=another-backend"): + registry.get_attention_op(_contract(), requested_backend="another-backend") + + result = registry.get_attention_op( + _contract(), requested_backend="test-deterministic-cp-attention" + ) + assert result.provenance["actual_backend"] == "test-deterministic-cp-attention" + + +def test_packed_layout_requires_declared_backend_support(): + capability = replace(_declared_cp_backend(), supports_packed_varlen=False) + contract = _contract( + sharding=_sharding(packed_sequence_offsets=(0, 512, 2048)), + causal_offsets=(0, 0), + batch_size=2, + ) + + assert capability.incompatibilities(contract) == ("packed varlen layout is unsupported",) + + +def test_rope_contract_requires_declared_backend_support(): + contract = _contract(rope=RoPESpec()) + capability = _declared_cp_backend() + + assert capability.incompatibilities(contract) == ("RoPE/position metadata is unsupported",) + + supported = replace(capability, supports_rope_metadata=True) + assert supported.incompatibilities(contract) == () + + +def test_fused_rope_attention_boundary_requires_declared_backend_support(): + contract = _contract(rope=RoPESpec(fusion_boundary=RoPEFusionBoundary.FUSED_ROPE_ATTENTION)) + capability = replace(_declared_cp_backend(), supports_rope_metadata=True) + + assert capability.incompatibilities(contract) == ( + "fused RoPE+Attention boundary is unsupported", + ) + + supported = replace(capability, supports_fused_rope_attention=True) + assert supported.incompatibilities(contract) == () + + +def test_packed_sequence_count_must_match_logical_batch_size(): + sharding = _sharding(packed_sequence_offsets=(0, 512, 2048)) + + with pytest.raises(AttentionContractError, match="must equal logical batch_size"): + _contract(sharding=sharding, causal_offsets=(0, 0), batch_size=1) + + contract = _contract(sharding=sharding, causal_offsets=(0, 0), batch_size=2) + assert contract.batch_size == 2 diff --git a/tests/test_attention_cross_config_binding.py b/tests/test_attention_cross_config_binding.py new file mode 100644 index 00000000..3e12ad34 --- /dev/null +++ b/tests/test_attention_cross_config_binding.py @@ -0,0 +1,1376 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Tests for #235 PR4: rollout/training attention contract binding. + +Every test here runs on CPU without Megatron or vLLM installed. That is the point: +the binding rules are contract logic, and contract logic that can only be exercised +on a 2-node x 2-GPU cluster would never be exercised. +""" + +from __future__ import annotations + +import json +from dataclasses import replace +from enum import Enum +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from rl_engine.alignment.cross_config.adapters import ( + QWEN3_8B, + WS2_ATTENTION_KNOBS, + AttentionRuntimeReadback, + MegatronAttentionMaterializer, + MegatronProvenanceAdapter, + VllmProvenanceAdapter, + VllmRolloutMaterializer, +) +from rl_engine.alignment.cross_config.attention_binding import ( + ATTENTION_LSE_DOMAIN, + AttentionBindingError, + BindingErrorCode, + BindingTier, + bind_attention_contracts, + bind_attention_runtime_readbacks, + first_blocking_issue, + identity_fingerprint, + summarize_binding, +) +from rl_engine.alignment.cross_config.determinism import ( + compare_determinism, + megatron_probe_from_config, + vllm_probe_from_env, +) +from rl_engine.alignment.cross_config.schema import MaterializationStatus +from rl_engine.kernels.attention_contract import ( + STRICT_ATTENTION_CORE_ID, + STRICT_ATTENTION_SCHEDULE_ID, + AttentionContractError, + AttentionMode, + AttentionRole, + KVCacheSpec, + SplitKVExecutionPlan, + SplitKVRuntimePlanEntry, + SplitKVRuntimePlanSet, + SplitKVSpec, + build_split_kv_runtime_plan_set, +) +from rl_engine.kernels.attention_preprocess import MANDATED_ATTENTION_PREPROCESS_BACKENDS +from rl_engine.kernels.attention_projection import ( + O_PROJ_COLLECTIVE_CONTRACT, + QKV_COLLECTIVE_CONTRACT, + ProjectionPlan, +) + +pytestmark = pytest.mark.unit + + +TRAINING_KNOBS = { + "batch.size": 2, + "training.tensor_parallel_size": 2, + "training.context_parallel_size": 2, + "training.compute_dtype": "bf16", + "attention.split_kv_policy": 32, +} + +ROLLOUT_KNOBS = { + "batch.size": 2, + "rollout.tensor_parallel_size": 2, + "rollout.context_parallel_size": 2, + "rollout.dtype": "bf16", + "attention.split_kv_policy": 32, +} + + +def _identity(**overrides): + identity = { + "checkpoint_id": "qwen3-8b", + "model_version": "v1", + "weight_version": 7, + "tokenizer_fingerprint": "tokenizer-abc", + "token_ids_fingerprint": "tokens-abc", + "active_mask_fingerprint": "mask-abc", + "position_ids_fingerprint": "pos-abc", + "padding_side": "right", + "pre_update_state": "pre_update", + "batch_size": 2, + "global_token_positions_fingerprint": "gtp-abc", + "kv_seq_lens_fingerprint": "kvlen-abc", + } + identity.update(QWEN3_8B.identity_fields()) + identity.update(overrides) + return identity + + +def _contracts(): + training = MegatronAttentionMaterializer().build_contract(TRAINING_KNOBS) + rollout = VllmRolloutMaterializer().build_contract(ROLLOUT_KNOBS) + return rollout, training + + +def _plan_set(contract, *, backend): + return build_split_kv_runtime_plan_set( + (contract.sharding.global_sequence_length,) * contract.batch_size, + tp_world_size=contract.sharding.tp_world_size, + cp_world_size=contract.sharding.cp_world_size, + split_kv=contract.split_kv, + backend=backend, + ) + + +def _bind(rollout_identity=None, training_identity=None, **kwargs): + rollout, training = _contracts() + rollout = kwargs.pop("rollout_contract", rollout) + training = kwargs.pop("training_contract", training) + return bind_attention_contracts( + rollout_contract=rollout, + training_contract=training, + rollout_identity=rollout_identity if rollout_identity is not None else _identity(), + training_identity=training_identity if training_identity is not None else _identity(), + rollout_backend_id="vllm.flash_attn", + training_backend_id="rlkernel.cp_attention_reference", + rollout_split_kv_plan_set=kwargs.pop( + "rollout_split_kv_plan_set", _plan_set(rollout, backend="vllm.readback") + ), + training_split_kv_plan_set=kwargs.pop( + "training_split_kv_plan_set", _plan_set(training, backend="megatron.readback") + ), + **kwargs, + ) + + +# -------------------------------------------------------------------------- +# tier 1: identity +# -------------------------------------------------------------------------- + + +def test_matching_identity_and_topology_bind_despite_different_materialization(): + """The core claim of PR4: same identity + same reduction, different runtimes.""" + + result = _bind() + + assert result.comparable + assert result.passed + assert result.issues == () + # Attention mode is a framework materialization difference, while both sides + # execute the same TP=2/CP=2 local ownership and Split-K schedule. + assert "mode" in result.recorded_differences + assert result.recorded_differences["mode"] == { + "rollout": "chunked_prefill", + "training": "prefill", + } + + +def test_weight_version_mismatch_is_not_comparable(): + result = _bind(rollout_identity=_identity(weight_version=6)) + + assert not result.comparable + assert not result.passed + codes = {issue.code for issue in result.issues} + assert BindingErrorCode.IDENTITY_MISMATCH in codes + blocking = first_blocking_issue(result) + assert blocking is not None and blocking.tier is BindingTier.IDENTICAL + assert "NOT COMPARABLE" in summarize_binding(result) + + +def test_rope_theta_mismatch_is_not_comparable(): + """RoPE math constants are identity, not materialization.""" + + result = _bind(training_identity=_identity(rope_theta=10000.0)) + + assert not result.comparable + assert any(issue.field == "rope_theta" for issue in result.issues) + + +def test_null_rope_scaling_is_a_value_not_an_omission(): + """Qwen3-8B applies no RoPE scaling; ``None`` must not read as undeclared.""" + + result = _bind() + + assert not result.issues_by_code(BindingErrorCode.IDENTITY_MISSING) + + +def test_missing_identity_field_is_reported_per_side(): + identity = _identity() + del identity["padding_side"] + result = _bind(rollout_identity=identity, training_identity=identity) + + missing = result.issues_by_code(BindingErrorCode.IDENTITY_MISSING) + assert {issue.field for issue in missing} == { + "rollout.padding_side", + "training.padding_side", + } + assert not result.comparable + + +def test_single_gpu_harness_may_waive_full_identity(): + """#235 PR2 has no KV-cache identity to declare; it opts out explicitly.""" + + identity = _identity() + del identity["global_token_positions_fingerprint"] + del identity["kv_seq_lens_fingerprint"] + + strict = _bind(rollout_identity=identity, training_identity=identity) + waived = _bind( + rollout_identity=identity, + training_identity=identity, + require_full_identity=False, + ) + + assert not strict.comparable + assert waived.comparable and waived.passed + + +def test_identity_fingerprint_ignores_undeclared_extra_keys(): + base = _identity() + decorated = dict(base, diagnostic_note="added later") + + assert identity_fingerprint(base) == identity_fingerprint(decorated) + + +# -------------------------------------------------------------------------- +# tier 2: reduction semantics +# -------------------------------------------------------------------------- + + +def test_reduction_semantics_are_bound_and_fingerprinted(): + result = _bind() + + reduction = result.provenance["training"]["contract"]["reduction"] + assert reduction["merge"] == "online_softmax_lse" + assert reduction["acc_dtype"] == "fp32" + assert reduction["order"] == "global_block_index" + assert reduction["downcast_at"] == "final_write" + assert result.reduction_fingerprint + + +def test_reduction_engine_difference_is_recorded_not_rejected(): + """A TE merge oracle on one side must not fail the binding.""" + + from rl_engine.alignment.cross_config.attention_binding import ( + RECORDED_FIELDS, + SEMANTIC_REDUCTION_FIELDS, + ) + + assert "reduction.engine" in RECORDED_FIELDS + assert "engine" not in SEMANTIC_REDUCTION_FIELDS + + +def test_lse_domain_is_recorded_as_attention_domain(): + """#235: attention exports attention-domain LSE, not vocab-logprob LSE.""" + + result = _bind() + + assert result.provenance["lse_domain"] == ATTENTION_LSE_DOMAIN == "attention" + + +def test_mixed_dtypes_fail_closed(): + """BF16 rollout against FP16 training produces an unattributable number.""" + + rollout = VllmRolloutMaterializer().build_contract( + {**ROLLOUT_KNOBS, "rollout.dtype": "float16"} + ) + result = _bind(rollout_contract=rollout) + + assert result.comparable # identity is fine + assert not result.passed + assert any(issue.field == "dtype" for issue in result.issues) + + +def test_precision_sweep_may_opt_into_mixed_dtypes(): + """#235 PR5 sweeps BF16 against an FP32 reference; it says so explicitly.""" + + training = MegatronAttentionMaterializer().build_contract( + {**TRAINING_KNOBS, "training.compute_dtype": "float32"} + ) + result = _bind(training_contract=training, allow_dtype_difference=True) + + assert result.passed + assert result.provenance["dtype"] == "fp32" + + +def test_batch_size_mismatch_is_not_comparable(): + """Batch invariance is a claim about batch makeup, so it belongs to identity.""" + + result = _bind(rollout_identity=_identity(batch_size=4)) + + assert not result.comparable + assert any(issue.field == "batch_size" for issue in result.issues) + + +def test_missing_split_kv_runtime_evidence_fails_closed(): + result = _bind(rollout_split_kv_plan_set=None) + + assert result.comparable + assert not result.passed + assert result.issues_by_code(BindingErrorCode.SPLIT_KV_RUNTIME_MISSING) + + +def test_split_kv_requested_policy_mismatch_fails_closed(): + rollout, training = _contracts() + rollout = replace(rollout, split_kv=SplitKVSpec.fixed(16)) + result = _bind(rollout_contract=rollout) + + assert result.comparable + assert not result.passed + assert any(issue.field == "split_kv" for issue in result.issues) + + +def test_split_kv_runtime_boundary_mismatch_fails_closed(): + rollout, training = _contracts() + mismatched_rollout = build_split_kv_runtime_plan_set( + (rollout.sharding.global_sequence_length,) * rollout.batch_size, + tp_world_size=2, + cp_world_size=2, + split_kv=SplitKVSpec.fixed(16), + backend="vllm.readback", + ) + result = _bind(rollout_split_kv_plan_set=mismatched_rollout) + + assert not result.passed + issues = result.issues_by_code(BindingErrorCode.SPLIT_KV_MISMATCH) + assert any(issue.field == "split_kv_runtime_plan_set" for issue in issues) + + +def test_split_kv_plan_set_must_match_its_own_contract_topology(): + rollout, _ = _contracts() + wrong_topology = build_split_kv_runtime_plan_set( + (rollout.sharding.global_sequence_length,) * rollout.batch_size, + tp_world_size=1, + cp_world_size=2, + split_kv=rollout.split_kv, + backend="vllm.readback", + ) + + result = _bind(rollout_split_kv_plan_set=wrong_topology) + + assert not result.passed + assert any( + issue.field == "rollout.split_kv_runtime_plan_set" + for issue in result.issues_by_code(BindingErrorCode.SPLIT_KV_MISMATCH) + ) + + +@pytest.mark.parametrize( + ("field_name", "corrupt_value"), + [ + ("merge_order", Enum("BadOrder", {"ARRIVAL": "arrival"}).ARRIVAL), + ("acc_dtype", Enum("BadDType", {"BF16": "bf16"}).BF16), + ("downcast_at", Enum("BadDowncast", {"PER_BLOCK": "per_block"}).PER_BLOCK), + ], +) +def test_split_kv_runtime_merge_semantic_corruption_fails_closed(field_name, corrupt_value): + rollout, _ = _contracts() + corrupted = _plan_set(rollout, backend="vllm.readback") + # Runtime reports are deserialized at this boundary. Simulate a corrupted + # report after construction to prove binding compares the actual fields. + object.__setattr__(corrupted.entries[0].execution, field_name, corrupt_value) + + result = _bind(rollout_split_kv_plan_set=corrupted) + + assert not result.passed + assert result.issues_by_code(BindingErrorCode.SPLIT_KV_MISMATCH) + + +def test_split_kv_runtime_fallback_fails_closed(): + rollout, _ = _contracts() + plan_set = _plan_set(rollout, backend="vllm.readback") + fallback_entries = [] + for entry in plan_set.entries: + execution = entry.execution + fallback_entries.append( + SplitKVRuntimePlanEntry( + coordinate=entry.coordinate, + expected_kv_range=entry.expected_kv_range, + execution=SplitKVExecutionPlan( + requested_mode=execution.requested_mode, + requested_split_size=execution.requested_split_size, + actual_mode=execution.actual_mode, + actual_split_size=execution.actual_split_size, + boundaries=execution.boundaries, + merge_order=execution.merge_order, + acc_dtype=execution.acc_dtype, + downcast_at=execution.downcast_at, + backend=execution.backend, + source="runtime_fallback", + fallback=True, + fallback_reason="backend substituted a runtime plan", + ), + ) + ) + fallback = SplitKVRuntimePlanSet( + batch_size=plan_set.batch_size, + tp_world_size=plan_set.tp_world_size, + cp_world_size=plan_set.cp_world_size, + total_kv_tokens=plan_set.total_kv_tokens, + entries=tuple(fallback_entries), + ) + + result = _bind(rollout_split_kv_plan_set=fallback) + + assert not result.passed + assert result.issues_by_code(BindingErrorCode.SPLIT_KV_FALLBACK) + + +@pytest.mark.parametrize( + "rollout_overrides", + [ + {"rollout.tensor_parallel_size": 1}, + {"rollout.context_parallel_size": 1}, + ], +) +def test_tp_or_cp_topology_mismatch_is_not_comparable(rollout_overrides): + rollout = VllmRolloutMaterializer().build_contract({**ROLLOUT_KNOBS, **rollout_overrides}) + result = _bind(rollout_contract=rollout) + + assert not result.comparable + assert result.issues_by_code(BindingErrorCode.TOPOLOGY_MISMATCH) + + +# -------------------------------------------------------------------------- +# role and input validation +# -------------------------------------------------------------------------- + + +def test_swapped_roles_are_rejected_outright(): + rollout, training = _contracts() + + with pytest.raises(AttentionBindingError): + bind_attention_contracts( + rollout_contract=training, + training_contract=rollout, + rollout_identity=_identity(), + training_identity=_identity(), + rollout_backend_id="a", + training_backend_id="b", + ) + + +# -------------------------------------------------------------------------- +# determinism cross-check +# -------------------------------------------------------------------------- + + +def _megatron_env(): + return {"NCCL_ALGO": "Tree", "NVTE_ALLOW_NONDETERMINISTIC_ALGO": "0"} + + +def _vllm_env(**overrides): + env = { + "VLLM_BATCH_INVARIANT": "1", + "NCCL_ALGO": "allreduce:tree", + "NCCL_PROTO": "Simple", + "NCCL_MIN_NCHANNELS": "1", + "NCCL_MAX_NCHANNELS": "1", + "CUBLAS_WORKSPACE_CONFIG": ":4096:8", + } + env.update(overrides) + return env + + +def test_nccl_algo_mismatch_blocks_the_binding(): + """Megatron asserts NCCL_ALGO; vLLM hard-sets a different value.""" + + training = megatron_probe_from_config( + SimpleNamespace(deterministic_mode=True), env=_megatron_env() + ) + rollout = vllm_probe_from_env(_vllm_env()) + + report = compare_determinism(rollout=rollout, training=training) + + assert not report.compatible + fields = {issue.field for issue in report.issues} + assert "env.NCCL_ALGO" in fields + assert "env.NCCL_PROTO" in fields + + +def test_matching_nccl_settings_are_compatible(): + shared = {"NCCL_ALGO": "allreduce:tree", "NCCL_PROTO": "Simple"} + training = megatron_probe_from_config( + SimpleNamespace(deterministic_mode=True), env=dict(shared) + ) + rollout = vllm_probe_from_env({**_vllm_env(**shared), "CUBLAS_WORKSPACE_CONFIG": None}) + + report = compare_determinism(rollout=rollout, training=training) + + assert report.compatible, [issue.to_dict() for issue in report.issues] + + +def test_determinism_switch_off_on_either_side_blocks(): + training = megatron_probe_from_config( + SimpleNamespace(deterministic_mode=False), env=_megatron_env() + ) + rollout = vllm_probe_from_env(_vllm_env(VLLM_BATCH_INVARIANT="0")) + + report = compare_determinism(rollout=rollout, training=training) + + fields = {issue.field for issue in report.issues} + assert "training.deterministic_mode" in fields + assert "rollout.VLLM_BATCH_INVARIANT" in fields + + +def test_tf32_asymmetry_is_recorded_not_blocking(): + """Megatron does not manage TF32 at all; vLLM disables it. Record the gap.""" + + training = megatron_probe_from_config( + SimpleNamespace(deterministic_mode=True), env=_megatron_env() + ) + rollout = vllm_probe_from_env(_vllm_env()) + + report = compare_determinism(rollout=rollout, training=training) + + assert training.tf32_disabled is None + assert rollout.tf32_disabled is True + assert "tf32_disabled" in report.differences + assert not any(issue.field == "tf32_disabled" for issue in report.issues) + + +def test_determinism_issues_flow_into_the_binding(): + training = megatron_probe_from_config( + SimpleNamespace(deterministic_mode=True), env=_megatron_env() + ) + rollout = vllm_probe_from_env(_vllm_env()) + report = compare_determinism(rollout=rollout, training=training) + + result = _bind(determinism_issues=report.issues) + + assert result.comparable # identity is fine + assert not result.passed # but the reduction environment is not + assert result.issues_by_code(BindingErrorCode.DETERMINISM_INCOMPATIBLE) + assert "FAILED CLOSED" in summarize_binding(result) + + +# -------------------------------------------------------------------------- +# sharding derived from the frozen #239 layout +# -------------------------------------------------------------------------- + + +@pytest.mark.parametrize("cp_rank", [0, 1]) +def test_cp_shards_cover_the_global_sequence_without_overlap(cp_rank): + contract = MegatronAttentionMaterializer( + cp_rank=cp_rank, global_sequence_length=4096 + ).build_contract(TRAINING_KNOBS) + + sharding = contract.sharding + assert sharding.local_sequence_length == 2048 + assert sharding.global_block_indices == (cp_rank,) + assert sharding.global_block_token_starts == (cp_rank * 2048,) + # The causal offset must be the number of preceding *global* tokens, otherwise + # rank 1 would mask as if its shard started at position zero. + assert contract.causal_offsets == (cp_rank * 2048, cp_rank * 2048) + + +def test_tp_head_shards_split_qwen3_gqa_evenly(): + contract = MegatronAttentionMaterializer(tp_rank=1).build_contract(TRAINING_KNOBS) + + sharding = contract.sharding + assert (sharding.global_q_heads, sharding.global_kv_heads) == (32, 8) + assert (sharding.local_q_heads, sharding.local_kv_heads) == (16, 4) + assert (sharding.local_q_head_start, sharding.local_kv_head_start) == (16, 4) + + +@pytest.mark.parametrize("tp_world_size", [2, 4, 8]) +def test_supported_tp_degrees_shard_qwen3_gqa(tp_world_size): + """Qwen3-8B has 32 Q heads and 8 KV heads, so TP in {2, 4, 8} all divide.""" + + contract = MegatronAttentionMaterializer().build_contract( + {**TRAINING_KNOBS, "training.tensor_parallel_size": tp_world_size} + ) + + sharding = contract.sharding + assert sharding.local_q_heads == 32 // tp_world_size + assert sharding.local_kv_heads == 8 // tp_world_size + + +@pytest.mark.parametrize( + ("knob_value", "expected"), + [("bfloat16", "bf16"), ("float16", "fp16"), ("float32", "fp32"), ("fp16", "fp16")], +) +def test_planner_normalized_dtypes_reach_the_contract(knob_value, expected): + """The planner emits torch spellings; AttentionDType uses short ones.""" + + contract = MegatronAttentionMaterializer().build_contract( + {**TRAINING_KNOBS, "training.compute_dtype": knob_value} + ) + + assert contract.dtype.value == expected + + +def test_unknown_dtype_is_rejected_with_the_offending_field(): + with pytest.raises(ValueError, match="training.compute_dtype"): + MegatronAttentionMaterializer().build_contract( + {**TRAINING_KNOBS, "training.compute_dtype": "int8"} + ) + + +def test_indivisible_tp_is_rejected(): + with pytest.raises(ValueError, match="divide evenly"): + MegatronAttentionMaterializer().build_contract( + {**TRAINING_KNOBS, "training.tensor_parallel_size": 3} + ) + + +def test_indivisible_cp_sequence_is_rejected(): + with pytest.raises(ValueError, match="divide evenly"): + MegatronAttentionMaterializer(global_sequence_length=4097).build_contract(TRAINING_KNOBS) + + +# -------------------------------------------------------------------------- +# materialization: fail closed rather than silently substitute +# -------------------------------------------------------------------------- + + +def _statuses(materialization, path): + return [app.status for app in materialization.applications if app.path == path] + + +def _readback(materializer, flat, *, source): + contract = materializer.build_contract(flat) + projection_plans = { + name: ProjectionPlan( + projection=name, + backend_id="rlkernel.cuda.det_gemm", + fallback=True, + fallback_reason="native projection probe failed", + probe_id=f"{source}-{name}", + collective=collective.to_dict(), + ).to_dict() + for name, collective in ( + ("qkv", QKV_COLLECTIVE_CONTRACT), + ("o_proj", O_PROJ_COLLECTIVE_CONTRACT), + ) + } + return AttentionRuntimeReadback( + contract=contract, + actual_knobs=dict(flat), + split_kv_plan_set=_plan_set(contract, backend=source), + source=source, + frozen_scope_verified=True, + preprocess_backends=MANDATED_ATTENTION_PREPROCESS_BACKENDS, + preprocess_fallback=False, + projection_plans=projection_plans, + actual_backend=f"reference.{source}", + communication_backend="none", + production_ready=False, + ) + + +def _strict_readback(materializer, flat, *, source): + readback = _readback(materializer, flat, source=source) + contract = replace(readback.contract, split_kv=SplitKVSpec.disabled()) + actual_knobs = dict(readback.actual_knobs) + actual_knobs["attention.split_kv_policy"] = "disabled" + return replace( + readback, + contract=contract, + actual_knobs=actual_knobs, + split_kv_plan_set=_plan_set(contract, backend=source), + strict_mode=True, + strict_core_id=STRICT_ATTENTION_CORE_ID, + strict_schedule=STRICT_ATTENTION_SCHEDULE_ID, + native_attention_arithmetic=False, + strict_split_kv_policy="disabled", + actual_backend="rlkernel.cuda.deterministic_attention", + communication_backend="self_owned_cuda_ag_rs", + production_ready=True, + ) + + +def test_configured_contract_without_runtime_readback_is_unobservable(): + normalized = { + "batch": {"size": 2}, + "training": { + "tensor_parallel_size": 2, + "context_parallel_size": 2, + "compute_dtype": "bf16", + }, + "attention": {"split_kv_policy": 32}, + } + materialization = MegatronAttentionMaterializer().materialize(normalized, WS2_ATTENTION_KNOBS) + + assert materialization.applications + assert {app.status for app in materialization.applications} == { + MaterializationStatus.UNOBSERVABLE + } + assert all(app.actual is None for app in materialization.applications) + + +@pytest.mark.parametrize( + ("materializer_type", "flat", "source"), + [ + (MegatronAttentionMaterializer, TRAINING_KNOBS, "megatron.runtime_readback"), + (VllmRolloutMaterializer, ROLLOUT_KNOBS, "vllm.runtime_readback"), + ], +) +def test_runtime_readback_can_verify_materialized_knobs(materializer_type, flat, source): + configured = materializer_type() + readback = _readback(configured, flat, source=source) + materializer = materializer_type(runtime_readback=readback) + normalized = {} + for path, value in flat.items(): + section, key = path.split(".", 1) + normalized.setdefault(section, {})[key] = value + + materialization = materializer.materialize(normalized, WS2_ATTENTION_KNOBS) + + assert materialization.applications + assert {app.status for app in materialization.applications} == {MaterializationStatus.APPLIED} + side = "training" if materializer_type is MegatronAttentionMaterializer else "rollout" + assert materialization.binding.side_configs[side]["runtime_readback"]["source"] == source + + +def test_runtime_readback_mismatch_is_a_fallback(): + configured = MegatronAttentionMaterializer() + readback = _readback(configured, TRAINING_KNOBS, source="megatron.runtime_readback") + actual = dict(readback.actual_knobs) + actual["training.context_parallel_size"] = 1 + mismatched = AttentionRuntimeReadback( + contract=readback.contract, + actual_knobs=actual, + split_kv_plan_set=readback.split_kv_plan_set, + source=readback.source, + frozen_scope_verified=True, + preprocess_backends=readback.preprocess_backends, + preprocess_fallback=readback.preprocess_fallback, + ) + normalized = { + "batch": {"size": 2}, + "training": { + "tensor_parallel_size": 2, + "context_parallel_size": 2, + "compute_dtype": "bf16", + }, + "attention": {"split_kv_policy": 32}, + } + materialization = MegatronAttentionMaterializer(runtime_readback=mismatched).materialize( + normalized, WS2_ATTENTION_KNOBS + ) + + assert _statuses(materialization, "training.context_parallel_size") == [ + MaterializationStatus.FALLBACK + ] + + +def test_decode_split_kv_plan_set_must_match_kv_cache_lengths(): + rollout, _ = _contracts() + kv_cache = KVCacheSpec( + cache_positions=(3, 5), + kv_seq_lens=(4, 6), + block_table=((0, 1, -1), (2, 3, 4)), + global_token_positions=tuple(range(4)) + tuple(range(6)), + page_size=2, + ) + decode = replace( + rollout, + role=AttentionRole.INFER, + mode=AttentionMode.DECODE, + query_sequence_length=1, + causal_offsets=(3, 5), + kv_cache=kv_cache, + ) + wrong_lengths = build_split_kv_runtime_plan_set( + (4, 8), + tp_world_size=2, + cp_world_size=2, + split_kv=decode.split_kv, + backend="vllm.decode.readback", + ) + + with pytest.raises(ValueError, match="KV-cache lengths"): + AttentionRuntimeReadback( + contract=decode, + actual_knobs={}, + split_kv_plan_set=wrong_lengths, + source="vllm.decode.readback", + frozen_scope_verified=True, + ) + + +def test_strict_runtime_readback_entrypoint_binds_executed_evidence(): + rollout_materializer = VllmRolloutMaterializer() + training_materializer = MegatronAttentionMaterializer() + rollout = _readback(rollout_materializer, ROLLOUT_KNOBS, source="vllm.runtime_readback") + training = _readback( + training_materializer, + TRAINING_KNOBS, + source="megatron.runtime_readback", + ) + + result = bind_attention_runtime_readbacks( + rollout=rollout, + training=training, + rollout_identity=_identity(), + training_identity=_identity(), + rollout_backend_id="vllm.flash_attn", + training_backend_id="rlkernel.cp_attention_reference", + ) + + assert result.passed + assert result.provenance["split_kv_runtime"]["rollout"]["coverage"] == ( + "complete_batch_tp_cp_owner_cartesian_product" + ) + assert result.provenance["rollout"]["recorded"]["preprocess.qk_rmsnorm"] == ( + MANDATED_ATTENTION_PREPROCESS_BACKENDS["qk_rmsnorm"] + ) + + +@pytest.mark.parametrize( + ("changes", "expected_code"), + [ + ( + {"strict_core_id": "different.core"}, + BindingErrorCode.ATTENTION_CORE_MISMATCH, + ), + ( + {"strict_schedule": "different_schedule"}, + BindingErrorCode.ATTENTION_CORE_SCHEDULE, + ), + ( + {"strict_split_kv_policy": "fixed"}, + BindingErrorCode.ATTENTION_CORE_SPLIT_K, + ), + ], +) +def test_strict_runtime_readback_rejects_non_shared_attention_arithmetic(changes, expected_code): + rollout = replace( + _strict_readback(VllmRolloutMaterializer(), ROLLOUT_KNOBS, source="vllm.runtime_readback"), + **changes, + ) + training = _strict_readback( + MegatronAttentionMaterializer(), + TRAINING_KNOBS, + source="megatron.runtime_readback", + ) + + result = bind_attention_runtime_readbacks( + rollout=rollout, + training=training, + rollout_identity=_identity(), + training_identity=_identity(), + rollout_backend_id="rlkernel.cuda.deterministic_attention", + training_backend_id="rlkernel.cuda.deterministic_attention", + ) + + assert not result.passed + assert result.issues_by_code(expected_code) + + +@pytest.mark.parametrize( + ("actual_backend", "core_id", "schedule", "communication_backend"), + [ + ( + "flash_attention_4.cute", + "rlkernel.attention.flash_attention4.num_splits1.v1", + "single_batch_flash_attention4_num_splits1", + "cuda_ag_rs", + ), + ( + "aiter.rocm.ck_dense_mha", + "rlkernel.attention.aiter_ck.num_splits1.v1", + "single_batch_aiter_ck_num_splits1", + "rccl_ag_rs", + ), + ], +) +def test_strict_runtime_readback_accepts_same_qualified_vendor_core( + actual_backend, core_id, schedule, communication_backend +): + rollout = replace( + _strict_readback(VllmRolloutMaterializer(), ROLLOUT_KNOBS, source="vllm.runtime_readback"), + strict_core_id=core_id, + strict_schedule=schedule, + native_attention_arithmetic=True, + actual_backend=actual_backend, + communication_backend=communication_backend, + ) + training = replace( + _strict_readback( + MegatronAttentionMaterializer(), TRAINING_KNOBS, source="megatron.runtime_readback" + ), + strict_core_id=core_id, + strict_schedule=schedule, + native_attention_arithmetic=True, + actual_backend=actual_backend, + communication_backend=communication_backend, + ) + + result = bind_attention_runtime_readbacks( + rollout=rollout, + training=training, + rollout_identity=_identity(), + training_identity=_identity(), + rollout_backend_id=actual_backend, + training_backend_id=actual_backend, + ) + + assert result.passed + + +def test_strict_runtime_readback_rejects_vendor_arithmetic_mismatch(): + rollout = _strict_readback( + VllmRolloutMaterializer(), ROLLOUT_KNOBS, source="vllm.runtime_readback" + ) + training = replace( + _strict_readback( + MegatronAttentionMaterializer(), TRAINING_KNOBS, source="megatron.runtime_readback" + ), + native_attention_arithmetic=True, + ) + + result = bind_attention_runtime_readbacks( + rollout=rollout, + training=training, + rollout_identity=_identity(), + training_identity=_identity(), + rollout_backend_id=rollout.actual_backend, + training_backend_id=training.actual_backend, + ) + + assert not result.passed + assert result.issues_by_code(BindingErrorCode.ATTENTION_NATIVE_ARITHMETIC) + + +def test_strict_runtime_readback_accepts_shared_no_split_k_core(): + rollout = _strict_readback( + VllmRolloutMaterializer(), ROLLOUT_KNOBS, source="vllm.runtime_readback" + ) + training = _strict_readback( + MegatronAttentionMaterializer(), + TRAINING_KNOBS, + source="megatron.runtime_readback", + ) + + result = bind_attention_runtime_readbacks( + rollout=rollout, + training=training, + rollout_identity=_identity(), + training_identity=_identity(), + rollout_backend_id="rlkernel.cuda.deterministic_attention", + training_backend_id="rlkernel.cuda.deterministic_attention", + ) + + assert result.passed + assert result.provenance["rollout"]["recorded"]["strict.core_id"] == (STRICT_ATTENTION_CORE_ID) + assert result.provenance["rollout"]["recorded"]["runtime.actual_backend"] == ( + "rlkernel.cuda.deterministic_attention" + ) + assert result.provenance["rollout"]["recorded"]["runtime.communication_backend"] == ( + "self_owned_cuda_ag_rs" + ) + assert result.provenance["rollout"]["recorded"]["runtime.production_ready"] is True + + +@pytest.mark.parametrize( + ("field", "value", "error_code"), + [ + ( + "reference_only", + True, + BindingErrorCode.ATTENTION_BACKEND_MISSING, + ), + ("attention_fallback", True, BindingErrorCode.ATTENTION_BACKEND_MISSING), + ( + "communication_backend", + "p2p_nccl_reference", + BindingErrorCode.ATTENTION_BACKEND_MISSING, + ), + ("production_ready", False, BindingErrorCode.ATTENTION_NOT_PRODUCTION_READY), + ], +) +def test_strict_runtime_readback_rejects_reference_only_evidence(field, value, error_code): + rollout = replace( + _strict_readback(VllmRolloutMaterializer(), ROLLOUT_KNOBS, source="vllm.runtime_readback"), + **{field: value}, + ) + training = _strict_readback( + MegatronAttentionMaterializer(), + TRAINING_KNOBS, + source="megatron.runtime_readback", + ) + result = bind_attention_runtime_readbacks( + rollout=rollout, + training=training, + rollout_identity=_identity(), + training_identity=_identity(), + rollout_backend_id="rlkernel.cuda.deterministic_attention", + training_backend_id="rlkernel.cuda.deterministic_attention", + ) + assert not result.passed + assert result.issues_by_code(error_code) + + +def test_strict_runtime_readback_accepts_common_deterministic_preprocess_fallback(): + rollout = replace( + _readback(VllmRolloutMaterializer(), ROLLOUT_KNOBS, source="vllm.runtime_readback"), + preprocess_fallback=True, + preprocess_fallback_reason="native probe failed", + preprocess_probe_id="rollout-probe", + ) + training = replace( + _readback( + MegatronAttentionMaterializer(), TRAINING_KNOBS, source="megatron.runtime_readback" + ), + preprocess_fallback=True, + preprocess_fallback_reason="native probe failed", + preprocess_probe_id="training-probe", + ) + + result = bind_attention_runtime_readbacks( + rollout=rollout, + training=training, + rollout_identity=_identity(), + training_identity=_identity(), + rollout_backend_id="vllm.flash_attn", + training_backend_id="rlkernel.cp_attention_reference", + ) + + assert result.passed + + +def test_strict_runtime_readback_rejects_distinct_native_preprocess_backends(): + rollout = replace( + _readback(VllmRolloutMaterializer(), ROLLOUT_KNOBS, source="vllm.runtime_readback"), + preprocess_backends={ + "qk_rmsnorm": "transformer_engine.rocm.rmsnorm", + "rope": "rlkernel.rocm.deterministic_rope", + }, + preprocess_probe_id="rollout-probe", + ) + training = replace( + _readback( + MegatronAttentionMaterializer(), TRAINING_KNOBS, source="megatron.runtime_readback" + ), + preprocess_backends={ + "qk_rmsnorm": "transformer_engine.cuda.rmsnorm", + "rope": "rlkernel.cuda.rope_sm90", + }, + preprocess_probe_id="training-probe", + ) + + result = bind_attention_runtime_readbacks( + rollout=rollout, + training=training, + rollout_identity=_identity(), + training_identity=_identity(), + rollout_backend_id="vllm.flash_attn", + training_backend_id="megatron.te", + ) + + assert not result.passed + assert result.issues_by_code(BindingErrorCode.ATTENTION_PREPROCESS_MISMATCH) + + +@pytest.mark.parametrize( + ("backends", "fallback", "expected_code"), + [ + ( + {"rope": MANDATED_ATTENTION_PREPROCESS_BACKENDS["rope"]}, + False, + BindingErrorCode.ATTENTION_PREPROCESS_MISSING, + ), + ( + {**MANDATED_ATTENTION_PREPROCESS_BACKENDS, "qk_rmsnorm": "unknown.backend"}, + False, + BindingErrorCode.ATTENTION_PREPROCESS_MISMATCH, + ), + ( + dict(MANDATED_ATTENTION_PREPROCESS_BACKENDS), + True, + BindingErrorCode.ATTENTION_PREPROCESS_MISMATCH, + ), + ], +) +def test_strict_runtime_readback_rejects_unverified_preprocess_backend( + backends, fallback, expected_code +): + rollout = _readback(VllmRolloutMaterializer(), ROLLOUT_KNOBS, source="vllm.runtime_readback") + rollout = replace( + rollout, + preprocess_backends=backends, + preprocess_fallback=fallback, + preprocess_fallback_reason=("test fallback" if fallback else None), + ) + training = _readback( + MegatronAttentionMaterializer(), + TRAINING_KNOBS, + source="megatron.runtime_readback", + ) + + result = bind_attention_runtime_readbacks( + rollout=rollout, + training=training, + rollout_identity=_identity(), + training_identity=_identity(), + rollout_backend_id="vllm.flash_attn", + training_backend_id="rlkernel.cp_attention_reference", + ) + + assert result.comparable + assert not result.passed + assert result.issues_by_code(expected_code) + + +def test_strict_runtime_readback_rejects_projection_split_k_or_missing_plan(): + rollout = _readback(VllmRolloutMaterializer(), ROLLOUT_KNOBS, source="vllm.runtime_readback") + plans = {name: dict(plan) for name, plan in rollout.projection_plans.items()} + plans["qkv"]["split_k"] = True + plans.pop("o_proj") + rollout = replace(rollout, projection_plans=plans) + training = _readback( + MegatronAttentionMaterializer(), TRAINING_KNOBS, source="megatron.runtime_readback" + ) + + result = bind_attention_runtime_readbacks( + rollout=rollout, + training=training, + rollout_identity=_identity(), + training_identity=_identity(), + rollout_backend_id="vllm.flash_attn", + training_backend_id="megatron.te", + ) + + assert not result.passed + assert result.issues_by_code(BindingErrorCode.ATTENTION_PROJECTION_MISMATCH) + assert result.issues_by_code(BindingErrorCode.ATTENTION_PROJECTION_MISSING) + + +def test_strict_runtime_readback_entrypoint_rejects_unverified_frozen_scope(): + rollout_materializer = VllmRolloutMaterializer() + training_materializer = MegatronAttentionMaterializer() + rollout = _readback(rollout_materializer, ROLLOUT_KNOBS, source="vllm.runtime_readback") + rollout = replace(rollout, frozen_scope_verified=False) + training = _readback( + training_materializer, + TRAINING_KNOBS, + source="megatron.runtime_readback", + ) + + result = bind_attention_runtime_readbacks( + rollout=rollout, + training=training, + rollout_identity=_identity(), + training_identity=_identity(), + rollout_backend_id="vllm.flash_attn", + training_backend_id="rlkernel.cp_attention_reference", + ) + + assert result.comparable + assert not result.passed + assert any(issue.field == "rollout.frozen_scope_verified" for issue in result.issues) + + +def test_arrival_merge_order_is_unsupported_not_silently_corrected(): + """The control group must stay distinguishable from the treatment.""" + + normalized = { + "batch": {"size": 2}, + "training": {"tensor_parallel_size": 2, "context_parallel_size": 2}, + "attention": {"reduction_order": "arrival"}, + } + materialization = MegatronAttentionMaterializer().materialize(normalized, WS2_ATTENTION_KNOBS) + + assert _statuses(materialization, "attention.reduction_order") == [ + MaterializationStatus.UNSUPPORTED + ] + assert _statuses(materialization, "training.tensor_parallel_size") == [ + MaterializationStatus.ERROR + ] + assert materialization.binding.side_configs["training"]["contract"] is None + assert "arrival" in materialization.binding.side_configs["training"]["contract_error"] + + +def test_unsupported_reduction_invalidates_vllm_contract_applications(): + normalized = { + "batch": {"size": 2}, + "rollout": {"tensor_parallel_size": 2, "context_parallel_size": 2}, + "attention": {"reduction_order": "arrival"}, + } + materialization = VllmRolloutMaterializer().materialize(normalized, WS2_ATTENTION_KNOBS) + + assert _statuses(materialization, "attention.reduction_order") == [ + MaterializationStatus.UNSUPPORTED + ] + assert _statuses(materialization, "rollout.tensor_parallel_size") == [ + MaterializationStatus.ERROR + ] + assert materialization.binding.side_configs["rollout"]["contract"] is None + + +def test_bf16_reduction_accumulation_is_unsupported(): + normalized = { + "batch": {"size": 2}, + "training": {"tensor_parallel_size": 2, "context_parallel_size": 2}, + "attention": {"reduction_acc_dtype": "bf16"}, + } + materialization = MegatronAttentionMaterializer().materialize(normalized, WS2_ATTENTION_KNOBS) + + assert _statuses(materialization, "attention.reduction_acc_dtype") == [ + MaterializationStatus.UNSUPPORTED + ] + + +def test_te_oracle_engine_is_unsupported_until_pr2_pr3(): + normalized = { + "batch": {"size": 2}, + "training": {"tensor_parallel_size": 2, "context_parallel_size": 2}, + "attention": {"reduction_engine": "te_oracle"}, + } + materialization = MegatronAttentionMaterializer().materialize(normalized, WS2_ATTENTION_KNOBS) + + assert _statuses(materialization, "attention.reduction_engine") == [ + MaterializationStatus.UNSUPPORTED + ] + + +def test_vllm_cp_falls_back_to_one_in_decode_and_says_why(): + materializer = VllmRolloutMaterializer(mode=AttentionMode.DECODE) + normalized = { + "batch": {"size": 2}, + "rollout": {"tensor_parallel_size": 2, "context_parallel_size": 2}, + } + + assert materializer.effective_cp_world_size({"rollout.context_parallel_size": 2}) == 1 + + materialization = materializer.materialize(normalized, WS2_ATTENTION_KNOBS) + contract_error = materialization.binding.side_configs["rollout"]["contract_error"] + assert "#235 PR6" in contract_error + assert MaterializationStatus.APPLIED not in {app.status for app in materialization.applications} + + +def test_decode_contract_is_refused_without_kv_cache_identity(): + with pytest.raises(AttentionContractError, match="PR6"): + VllmRolloutMaterializer(mode=AttentionMode.DECODE).build_contract(ROLLOUT_KNOBS) + + +def test_materializers_expose_distinct_implementation_fingerprints(): + megatron = MegatronAttentionMaterializer().implementation_fingerprint + vllm = VllmRolloutMaterializer().implementation_fingerprint + + assert megatron and vllm and megatron != vllm + + +def test_runtime_binding_reports_the_frozen_topology(): + normalized = { + "batch": {"size": 2}, + "training": {"tensor_parallel_size": 2, "context_parallel_size": 2}, + } + binding = MegatronAttentionMaterializer().materialize(normalized, WS2_ATTENTION_KNOBS).binding + + topology = binding.topology["training"] + assert topology["tensor_parallel_size"] == 2 + assert topology["context_parallel_size"] == 2 + assert topology["world_size"] == 4 + assert topology["pipeline_parallel_size"] == 1 + assert topology["data_parallel_size"] == 1 + + +# -------------------------------------------------------------------------- +# provenance adapters +# -------------------------------------------------------------------------- + + +def test_megatron_provenance_flags_undeclared_frozen_scope_fields(): + adapter = MegatronProvenanceAdapter(SimpleNamespace(deterministic_mode=True)) + + violations = adapter.frozen_scope_violations() + + # Nothing is declared, so every assertion reads as unknown rather than as met. + assert any("expert_model_parallel_size" in text for text in violations) + assert any("fp8" in text for text in violations) + + +def test_megatron_provenance_accepts_a_conforming_dense_config(): + adapter = MegatronProvenanceAdapter( + SimpleNamespace( + deterministic_mode=True, + pipeline_model_parallel_size=1, + expert_model_parallel_size=1, + sequence_parallel=False, + fp8=None, + hidden_dropout=0.0, + attention_dropout=0.0, + ) + ) + + assert adapter.frozen_scope_violations() == ("fp8 is not declared (expected None)",) + + +def test_megatron_construction_fingerprint_tracks_fusion_changes(): + base = SimpleNamespace(deterministic_mode=True, apply_rope_fusion=False) + fused = SimpleNamespace(deterministic_mode=True, apply_rope_fusion=True) + + assert ( + MegatronProvenanceAdapter(base).construction_fingerprint + != MegatronProvenanceAdapter(fused).construction_fingerprint + ) + + +def test_vllm_provenance_reads_page_size_and_split_kv_policy(): + adapter = VllmProvenanceAdapter( + cache_config=SimpleNamespace(block_size=16, cache_dtype="auto"), + attention_config=SimpleNamespace(flash_attn_max_num_splits_for_cuda_graph=32), + ) + + assert adapter.kv_page_size == 16 + assert adapter.split_kv_policy == 32 + assert adapter.to_dict()["flash_attn_max_num_splits_for_cuda_graph"] == 32 + + +def test_vllm_provenance_flags_fp8_kv_cache_and_cascade_attention(): + adapter = VllmProvenanceAdapter( + model_config=SimpleNamespace(quantization=None, disable_cascade_attn=False), + cache_config=SimpleNamespace( + cache_dtype="fp8", calculate_kv_scales=False, sliding_window=None + ), + parallel_config=SimpleNamespace(pipeline_parallel_size=1, data_parallel_size=1), + ) + + violations = adapter.frozen_scope_violations() + + assert any("cache_dtype" in text for text in violations) + assert any("disable_cascade_attn" in text for text in violations) + + +# -------------------------------------------------------------------------- +# scenario config +# -------------------------------------------------------------------------- + + +SCENARIO = ( + Path(__file__).resolve().parents[1] + / "examples" + / "cross_config_qwen3_8b_megatron_tp2_cp2_vllm.json" +) + + +def test_scenario_uses_megatron_vocabulary_only(): + config = json.loads(SCENARIO.read_text(encoding="utf-8")) + training = config["baseline"]["training"] + + assert training["attention_backend"] in {"flash", "fused", "unfused", "local", "auto"} + assert training["tensor_parallel_size"] == 2 + assert training["context_parallel_size"] == 2 + assert config["baseline"]["rollout"]["batch_invariant"] is True + assert config["scenario"]["debug_matrix"]["modules"] == ["attention", "ffn", "logp"] + assert config["scenario"]["debug_matrix"]["cartesian_product"] is False + + +def test_scenario_knob_paths_all_exist(): + config = json.loads(SCENARIO.read_text(encoding="utf-8")) + + def paths(mapping, prefix=""): + for key, value in mapping.items(): + path = f"{prefix}{key}" + if isinstance(value, dict): + yield from paths(value, f"{path}.") + else: + yield path + + declared = set(paths(config["baseline"])) + unknown = declared - set(WS2_ATTENTION_KNOBS) + assert not unknown, f"scenario declares unknown knobs: {sorted(unknown)}" + + for intervention in config["interventions"]: + assert intervention["path"] in WS2_ATTENTION_KNOBS diff --git a/tests/test_attention_debug_matrix.py b/tests/test_attention_debug_matrix.py new file mode 100644 index 00000000..eb23e513 --- /dev/null +++ b/tests/test_attention_debug_matrix.py @@ -0,0 +1,78 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import pytest + +from rl_engine.kernels.attention_contract import AttentionContractError +from rl_engine.kernels.ops.pytorch.attention.debug_matrix import ( + ATTENTION_DEBUG_MATRIX, + ATTENTION_DEBUG_MATRIX_SCHEMA_VERSION, + attention_debug_matrix, + attention_debug_matrix_row, +) + + +def test_attention_debug_matrix_is_fixed_replay_oat_with_controls(): + manifest = attention_debug_matrix() + + assert manifest["schema_version"] == ATTENTION_DEBUG_MATRIX_SCHEMA_VERSION + assert manifest["method"] == "fixed_replay_one_at_a_time" + assert manifest["baseline_row"] == "A0" + assert manifest["cartesian_product"] is False + assert [row["id"] for row in manifest["rows"]] == [ + "A0", + "A1", + "A2", + "A3", + "A4", + "A5", + "A6", + "A7", + "C0", + "C1", + "C2", + ] + assert sum(row["category"] == "root_cause" for row in manifest["rows"]) == 6 + assert sum(row["category"] == "comparability_gate" for row in manifest["rows"]) == 1 + assert sum(row["category"] == "invariant_control" for row in manifest["rows"]) == 3 + + +def test_attention_debug_matrix_rows_have_stable_semantics(): + baseline = attention_debug_matrix_row("A0") + assert baseline.category == "baseline" + assert baseline.expected == "baseline" + + root = attention_debug_matrix_row("A1") + assert root.category == "root_cause" + assert root.probe == "position_ids" + assert root.root_cause_axis == "position_rope" + assert root.expected == "diagnostic" + + topology = attention_debug_matrix_row("A4") + assert topology.category == "comparability_gate" + assert topology.expected == "rejected" + + control = attention_debug_matrix_row("C0") + assert control.category == "invariant_control" + assert control.probe == "tp_partition_control" + assert control.expected == "exact_zero" + + with pytest.raises(AttentionContractError, match="unknown Attention debug matrix row"): + attention_debug_matrix_row("A8") + + +def test_matrix_manifest_matches_taxonomy_representatives(): + from rl_engine.kernels.ops.pytorch.attention.debug_taxonomy import ATTENTION_DEBUG_AXES + + root_rows = [ + row + for row in ATTENTION_DEBUG_MATRIX + if row.category in {"root_cause", "comparability_gate"} + ] + assert [row.root_cause_axis for row in root_rows] == [ + axis.axis_id for axis in ATTENTION_DEBUG_AXES + ] + assert [row.probe for row in root_rows] == [ + axis.representative_subprobe for axis in ATTENTION_DEBUG_AXES + ] diff --git a/tests/test_attention_debug_taxonomy.py b/tests/test_attention_debug_taxonomy.py new file mode 100644 index 00000000..1d783132 --- /dev/null +++ b/tests/test_attention_debug_taxonomy.py @@ -0,0 +1,48 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import pytest + +from rl_engine.kernels.attention_contract import AttentionContractError +from rl_engine.kernels.ops.pytorch.attention.debug_taxonomy import ( + ATTENTION_DEBUG_SCHEMA_VERSION, + attention_debug_probe_metadata, + attention_debug_taxonomy, +) + + +def test_attention_debug_taxonomy_is_compact_complete_and_unambiguous(): + taxonomy = attention_debug_taxonomy() + + assert taxonomy["schema_version"] == ATTENTION_DEBUG_SCHEMA_VERSION + assert taxonomy["root_cause_axis_count"] == 7 + assert taxonomy["subprobe_count"] == 21 + assert taxonomy["invariant_control_count"] == 3 + assert set(taxonomy["root_cause_axes"]) == { + "position_rope", + "qk_preprocessing", + "mask_sequence_boundary", + "topology_head_ownership", + "kv_cache_identity_layout", + "numerical_policy", + "distributed_schedule", + } + assert "strict_cp_degree_control" not in { + probe for axis in taxonomy["root_cause_axes"].values() for probe in axis["subprobes"] + } + + +def test_attention_debug_probe_metadata_separates_subprobes_and_controls(): + assert attention_debug_probe_metadata("position_ids") == { + "category": "root_cause_subprobe", + "root_cause_axis": "position_rope", + "root_cause_label": "Position / RoPE", + "representative": True, + } + control = attention_debug_probe_metadata("tp_partition_control") + assert control["category"] == "invariant_control" + assert control["root_cause_axis"] is None + + with pytest.raises(AttentionContractError, match="unknown Attention debug probe"): + attention_debug_probe_metadata("strict_cp_degree_control") diff --git a/tests/test_attention_preprocess.py b/tests/test_attention_preprocess.py new file mode 100644 index 00000000..9b800dca --- /dev/null +++ b/tests/test_attention_preprocess.py @@ -0,0 +1,194 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from __future__ import annotations + +import pytest +import torch + +from rl_engine.kernels.attention_preprocess import ( + MANDATED_ATTENTION_PREPROCESS_BACKENDS, + TE_ROCM_QK_RMSNORM_BACKEND_ID, + H100AttentionPreprocessor, +) +from rl_engine.kernels.ops.pytorch.norm.rms_norm import NativeRMSNormOp +from rl_engine.kernels.ops.pytorch.rotary_embedding.rope import NativeRoPEOp + + +def _has_h100_preprocess() -> bool: + if not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 9: + return False + try: + from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE + + required = ( + "rmsnorm_forward", + "rmsnorm_backward_dx", + "rmsnorm_backward_dw", + "rope_apply_sm90", + ) + return bool(_EXT_AVAILABLE and all(hasattr(_C, name) for name in required)) + except ImportError: + return False + + +requires_h100_preprocess = pytest.mark.skipif( + not _has_h100_preprocess(), + reason="Hopper with compiled RMSNorm and RoPE CUDA kernels is required", +) + + +def test_h100_preprocessor_uses_common_backend_ids_for_fallback(): + assert dict(MANDATED_ATTENTION_PREPROCESS_BACKENDS) == { + "qk_rmsnorm": "rlkernel.cuda.rmsnorm", + "rope": "rlkernel.cuda.rope_sm90", + } + + +def test_preprocessor_can_reuse_qk_norm_without_fusing_rope(): + def rmsnorm(x, weight, *, eps): + del eps + return (x.float() * weight.float()).to(torch.bfloat16) + + def rope(x, positions, *, theta): + del positions, theta + return x + + op = object.__new__(H100AttentionPreprocessor) + op.device = torch.device("cpu") + op.device_capability = (0, 0) + op.rmsnorm = rmsnorm + op.rope = rope + op.deterministic_backend_ids = { + "qk_rmsnorm": "rlkernel.rocm.triton_rmsnorm", + "rope": "rlkernel.rocm.deterministic_rope", + } + op.native_qk_norm = rmsnorm + op.native_rope = None + op.require_native_qk_norm = True + op.native_qk_norm_backend_id = TE_ROCM_QK_RMSNORM_BACKEND_ID + op.native_rope_backend_id = "unused" + op.policy_id = "test" + q = torch.randn(2, 4, 3, 8, dtype=torch.bfloat16) + k = torch.randn(2, 2, 3, 8, dtype=torch.bfloat16) + weight = torch.randn(8, dtype=torch.bfloat16) + positions = torch.arange(3, dtype=torch.int64) + + result = op(q, k, weight, weight, positions) + + assert result.fallback is False + assert result.backend_ids == { + "qk_rmsnorm": TE_ROCM_QK_RMSNORM_BACKEND_ID, + "rope": "rlkernel.rocm.deterministic_rope", + } + + +def test_preprocessor_falls_back_atomically_when_vendor_qk_norm_drifts(): + def deterministic(x, weight, *, eps): + del weight, eps + return x + + def drifting(x, weight, *, eps): + del weight, eps + return (x.float() + 1).to(torch.bfloat16) + + def rope(x, positions, *, theta): + del positions, theta + return x + + op = object.__new__(H100AttentionPreprocessor) + op.device = torch.device("cpu") + op.device_capability = (0, 0) + op.rmsnorm = deterministic + op.rope = rope + op.deterministic_backend_ids = dict(MANDATED_ATTENTION_PREPROCESS_BACKENDS) + op.native_qk_norm = drifting + op.native_rope = None + op.require_native_qk_norm = False + op.native_qk_norm_backend_id = "transformer_engine.cuda.rmsnorm" + op.native_rope_backend_id = "unused" + op.policy_id = "test" + q = torch.randn(1, 2, 3, 8, dtype=torch.bfloat16) + k = torch.randn(1, 1, 3, 8, dtype=torch.bfloat16) + weight = torch.randn(8, dtype=torch.bfloat16) + positions = torch.arange(3, dtype=torch.int64) + + result = op(q, k, weight, weight, positions) + + assert result.fallback is True + assert result.backend_ids == MANDATED_ATTENTION_PREPROCESS_BACKENDS + assert result.fallback_reason == "native_preprocess_bitwise_probe_failed" + + +def test_h100_preprocessor_fails_before_dispatch_without_cuda(monkeypatch): + monkeypatch.setattr(torch.cuda, "is_available", lambda: False) + with pytest.raises(RuntimeError, match="requires an available CUDA runtime"): + H100AttentionPreprocessor() + + +def test_h100_preprocessor_rejects_non_hopper_cuda(monkeypatch): + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "current_device", lambda: 0) + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda _device: (8, 0)) + with pytest.raises(RuntimeError, match="requires Hopper SM90"): + H100AttentionPreprocessor() + + +def _inputs(): + torch.manual_seed(7) + device = torch.device("cuda") + q = torch.randn(2, 4, 8, 128, device=device, dtype=torch.bfloat16) + k = torch.randn(2, 2, 8, 128, device=device, dtype=torch.bfloat16) + q_weight = torch.randn(128, device=device, dtype=torch.bfloat16) + k_weight = torch.randn(128, device=device, dtype=torch.bfloat16) + positions = torch.tensor( + [[0, 7, 2, 9, 4, 11, 6, 13], [100, 107, 102, 109, 104, 111, 106, 113]], + device=device, + dtype=torch.int64, + ) + return q, k, q_weight, k_weight, positions + + +@requires_h100_preprocess +def test_h100_preprocessor_executes_cuda_qk_norm_and_zigzag_rope(): + q, k, q_weight, k_weight, positions = _inputs() + result = H100AttentionPreprocessor(reuse_transformer_engine_qk_norm=False)( + q, k, q_weight, k_weight, positions + ) + + norm = NativeRMSNormOp() + rope = NativeRoPEOp() + q_ref = rope(norm(q, q_weight), positions) + k_ref = rope(norm(k, k_weight), positions) + + assert result.fallback is True + assert dict(result.backend_ids) == dict(MANDATED_ATTENTION_PREPROCESS_BACKENDS) + assert result.fallback_reason == "native_preprocess_not_supplied" + assert result.probe_id + assert result.readback_fields() == { + "preprocess_backends": dict(MANDATED_ATTENTION_PREPROCESS_BACKENDS), + "preprocess_fallback": True, + "preprocess_fallback_reason": "native_preprocess_not_supplied", + "preprocess_probe_id": result.probe_id, + "preprocess_policy_id": result.policy_id, + } + torch.testing.assert_close(result.q.float(), q_ref.float(), atol=2e-2, rtol=2e-2) + torch.testing.assert_close(result.k.float(), k_ref.float(), atol=2e-2, rtol=2e-2) + + +@requires_h100_preprocess +def test_h100_preprocessor_is_bitwise_batch_invariant_for_2d_positions(): + q, k, q_weight, k_weight, positions = _inputs() + op = H100AttentionPreprocessor(reuse_transformer_engine_qk_norm=False) + full = op(q, k, q_weight, k_weight, positions) + + for batch_index in range(q.shape[0]): + single = op( + q[batch_index : batch_index + 1], + k[batch_index : batch_index + 1], + q_weight, + k_weight, + positions[batch_index : batch_index + 1], + ) + assert torch.equal(full.q[batch_index], single.q[0]) + assert torch.equal(full.k[batch_index], single.k[0]) diff --git a/tests/test_attention_projection.py b/tests/test_attention_projection.py new file mode 100644 index 00000000..05d02598 --- /dev/null +++ b/tests/test_attention_projection.py @@ -0,0 +1,110 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import pytest +import torch + +from rl_engine.kernels.attention_projection import ( + O_PROJ_COLLECTIVE_CONTRACT, + QKV_COLLECTIVE_CONTRACT, + ROCM_DETERMINISTIC_PROJECTION_BACKEND_ID, + AttentionProjectionOp, + split_qkv, +) + + +def _deterministic(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + return torch.mm(x.float(), weight.float()).to(torch.bfloat16) + + +def _inputs(): + torch.manual_seed(19) + return ( + torch.randn(7, 8, dtype=torch.bfloat16), + torch.randn(8, 12, dtype=torch.bfloat16), + ) + + +@pytest.mark.parametrize( + ("projection", "collective"), + [("qkv", QKV_COLLECTIVE_CONTRACT), ("o_proj", O_PROJ_COLLECTIVE_CONTRACT)], +) +def test_projection_falls_back_to_common_deterministic_path(projection, collective): + x, weight = _inputs() + result = AttentionProjectionOp(projection, deterministic=_deterministic)(x, weight) + + assert torch.equal(result.output, _deterministic(x, weight)) + assert result.plan.backend_id == "rlkernel.cuda.det_gemm" + assert result.plan.fallback is True + assert result.plan.fallback_reason == "native_projection_not_supplied" + assert result.plan.split_k is False + assert result.plan.accumulation_dtype == "torch.float32" + assert dict(result.plan.collective) == collective.to_dict() + + +def test_projection_accepts_native_only_after_bitwise_probe(): + x, weight = _inputs() + result = AttentionProjectionOp( + "qkv", + native=_deterministic, + native_backend_id="megatron.te.qkv", + deterministic=_deterministic, + )(x, weight) + + assert torch.equal(result.output, _deterministic(x, weight)) + assert result.plan.backend_id == "megatron.te.qkv" + assert result.plan.fallback is False + assert result.plan.fallback_reason is None + + +def test_projection_records_rocm_deterministic_backend(): + x, weight = _inputs() + result = AttentionProjectionOp( + "qkv", + deterministic=_deterministic, + deterministic_backend_id=ROCM_DETERMINISTIC_PROJECTION_BACKEND_ID, + )(x, weight) + + assert result.plan.backend_id == ROCM_DETERMINISTIC_PROJECTION_BACKEND_ID + assert result.plan.fallback is True + + +def test_projection_rejects_native_drift_and_records_reason(): + x, weight = _inputs() + + def drifting_native(a, b): + return (_deterministic(a, b).float() + 1.0).to(torch.bfloat16) + + result = AttentionProjectionOp("o_proj", native=drifting_native, deterministic=_deterministic)( + x, weight + ) + + assert result.plan.fallback is True + assert result.plan.fallback_reason == "native_projection_bitwise_probe_failed" + assert torch.equal(result.output, _deterministic(x, weight)) + + +def test_split_qkv_is_fixed_contiguous_q_k_v_order(): + projected = torch.arange(2 * 16, dtype=torch.bfloat16).reshape(2, 16) + q, k, v = split_qkv(projected, q_heads=2, kv_heads=1, head_dim=4) + + assert q.shape == (2, 8) + assert k.shape == (2, 4) + assert v.shape == (2, 4) + assert torch.equal(torch.cat((q, k, v), dim=-1), projected) + + +def test_projection_requires_bf16_and_compatible_k(): + x, weight = _inputs() + with pytest.raises(TypeError, match="BF16"): + AttentionProjectionOp("qkv", deterministic=_deterministic)(x.float(), weight) + with pytest.raises(ValueError, match="K dimensions"): + AttentionProjectionOp("qkv", deterministic=_deterministic)(x, weight[:-1]) + + +def test_o_proj_collective_contract_includes_sp_scatter_gather_and_tp_reduction(): + assert O_PROJ_COLLECTIVE_CONTRACT.sp_forward == "reduce_scatter" + assert O_PROJ_COLLECTIVE_CONTRACT.sp_backward == "all_gather" + assert O_PROJ_COLLECTIVE_CONTRACT.reduction_forward == "all_reduce" + assert O_PROJ_COLLECTIVE_CONTRACT.reduction_backward == "none" diff --git a/tests/test_cp_attention.py b/tests/test_cp_attention.py new file mode 100644 index 00000000..7dff11c0 --- /dev/null +++ b/tests/test_cp_attention.py @@ -0,0 +1,712 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Tests for the WS2 deterministic CP attention reference. + +The implementation is a correctness-first prefill/chunked-prefill reference: +local KV blocks produce ``(out, lse)`` partial states and CP merges those states +with fp32 online-softmax arithmetic in logical global-block order. +""" + +import contextlib +import json +import math + +import pytest +import torch + +from rl_engine.kernels.ops.pytorch.attention.cp_attention import ( + AttentionPartialState, + DeterministicCPAttentionReferenceOp, + compare_cp_attention_backward, + merge_attention_partial_states, + split_kv_execution_plan_provenance, +) +from rl_engine.kernels.ops.pytorch.attention.standard_attn import NativeAttentionOp +from rl_engine.kernels.ops.pytorch.rotary_embedding.rope import NativeRoPEOp +from rl_engine.kernels.registry import kernel_registry + +_N_HEADS = 32 +_N_KV = 8 +_HEAD_DIM = 128 +_ATOL = 3.0e-6 +_GRAD_ATOL = 1.0e-5 + + +@contextlib.contextmanager +def _single_thread(): + prev = torch.get_num_threads() + torch.set_num_threads(1) + try: + yield + finally: + torch.set_num_threads(prev) + + +def _qkv( + batch, + sq, + skv, + *, + seed, + dtype=torch.float32, + heads=_N_HEADS, + kv_heads=_N_KV, + dim=_HEAD_DIM, +): + gen = torch.Generator().manual_seed(seed) + q = torch.randn(batch, heads, sq, dim, generator=gen, dtype=dtype) + k = torch.randn(batch, kv_heads, skv, dim, generator=gen, dtype=dtype) + v = torch.randn(batch, kv_heads, skv, dim, generator=gen, dtype=dtype) + return q, k, v + + +def _full_lse(q, k, *, causal, scale=None, key_padding_mask=None): + qf, kf = q.float(), k.float() + hq, sq, dim = qf.shape[1], qf.shape[2], qf.shape[3] + hkv, skv = kf.shape[1], kf.shape[2] + if hq % hkv != 0: + raise ValueError("invalid GQA shape") + if hq != hkv: + kf = kf.repeat_interleave(hq // hkv, dim=1) + scores = torch.matmul(qf, kf.transpose(-1, -2)) * ( + scale if scale is not None else 1.0 / math.sqrt(dim) + ) + if causal: + query_pos = torch.arange(skv - sq, skv) + key_pos = torch.arange(skv) + scores = scores.masked_fill( + (key_pos[None, :] > query_pos[:, None])[None, None, :, :], + float("-inf"), + ) + if key_padding_mask is not None: + scores = scores.masked_fill(~key_padding_mask[:, None, None, :], float("-inf")) + return torch.logsumexp(scores, dim=-1) + + +def test_cp1_matches_native_attention_and_exports_lse(): + op = DeterministicCPAttentionReferenceOp() + native = NativeAttentionOp() + q, k, v = _qkv(2, 8, 8, seed=1) + + with _single_thread(): + out, lse = op.forward_fp32_with_lse(q, k, v, causal=True, cp_world_size=1) + want = native.forward_fp32(q, k, v, causal=True) + want_lse = _full_lse(q, k, causal=True) + + torch.testing.assert_close(out, want, atol=_ATOL, rtol=0.0) + torch.testing.assert_close(lse, want_lse, atol=_ATOL, rtol=0.0) + assert lse.dtype == torch.float32 + assert lse.shape == q.shape[:3] + + +def test_cp2_prefill_matches_cp1_reference(): + op = DeterministicCPAttentionReferenceOp() + q, k, v = _qkv(2, 9, 9, seed=2) + + with _single_thread(): + out1, lse1 = op.forward_fp32_with_lse(q, k, v, causal=True, cp_world_size=1) + out2, lse2 = op.forward_fp32_with_lse(q, k, v, causal=True, cp_world_size=2) + + torch.testing.assert_close(out2, out1, atol=_ATOL, rtol=0.0) + torch.testing.assert_close(lse2, lse1, atol=_ATOL, rtol=0.0) + + +def test_cp2_consumes_post_rope_qk_with_shared_global_position_metadata(): + op = DeterministicCPAttentionReferenceOp() + rope = NativeRoPEOp() + pre_rope_q, pre_rope_k, v = _qkv(2, 7, 7, seed=14, heads=4, kv_heads=2, dim=8) + position_offsets = torch.tensor([17, 103], dtype=torch.long) + positions = position_offsets[:, None] + torch.arange(pre_rope_q.size(2), dtype=torch.long) + q = rope.forward_fp32(pre_rope_q, positions, theta=1_000_000.0) + k = rope.forward_fp32(pre_rope_k, positions, theta=1_000_000.0) + + assert not torch.equal(q, pre_rope_q.float()) + assert not torch.equal(k, pre_rope_k.float()) + + with _single_thread(): + out1, lse1 = op.forward_fp32_with_lse( + q, + k, + v, + causal=True, + query_position_offsets=position_offsets, + key_position_offsets=position_offsets, + cp_world_size=1, + ) + out2, lse2 = op.forward_fp32_with_lse( + q, + k, + v, + causal=True, + query_position_offsets=position_offsets, + key_position_offsets=position_offsets, + cp_world_size=2, + kv_chunk_size=2, + ) + + torch.testing.assert_close(out2, out1, atol=_ATOL, rtol=0.0) + torch.testing.assert_close(lse2, lse1, atol=_ATOL, rtol=0.0) + + +def test_chunked_prefill_replay_matches_unchunked_cp2(): + op = DeterministicCPAttentionReferenceOp() + q, k, v = _qkv(2, 10, 10, seed=3) + + with _single_thread(): + unchunked_out, unchunked_lse = op.forward_fp32_with_lse( + q, + k, + v, + causal=True, + cp_world_size=2, + ) + chunked_out, chunked_lse = op.forward_fp32_with_lse( + q, + k, + v, + causal=True, + cp_world_size=2, + kv_chunk_size=3, + ) + + torch.testing.assert_close(chunked_out, unchunked_out, atol=_ATOL, rtol=0.0) + torch.testing.assert_close(chunked_lse, unchunked_lse, atol=_ATOL, rtol=0.0) + + +def test_causal_mask_uses_global_positions_across_cp_boundary(): + op = DeterministicCPAttentionReferenceOp() + batch, heads, kv_heads, seq, dim = 1, 2, 1, 5, 3 + q = torch.zeros(batch, heads, seq, dim) + k = torch.zeros(batch, kv_heads, seq, dim) + v = torch.arange(seq * dim, dtype=torch.float32).reshape(1, 1, seq, dim) + out = op.forward_fp32(q, k, v, causal=True, cp_world_size=2) + + expected = torch.stack([v[0, 0, : index + 1].mean(dim=0) for index in range(seq)]) + expected = expected.reshape(1, 1, seq, dim).repeat(1, heads, 1, 1) + torch.testing.assert_close(out, expected, atol=1.0e-6, rtol=0.0) + + +def test_position_offsets_apply_varlen_causal_metadata_per_batch_row(): + op = DeterministicCPAttentionReferenceOp() + q = torch.zeros(2, 2, 2, 1) + k = torch.zeros(2, 1, 4, 1) + v = torch.arange(8, dtype=torch.float32).reshape(2, 1, 4, 1) + + out, lse = op.forward_fp32_with_lse( + q, + k, + v, + causal=True, + query_position_offsets=torch.tensor([0, 11]), + key_position_offsets=torch.tensor([0, 10]), + cp_world_size=2, + kv_chunk_size=1, + ) + + expected = torch.tensor([0.0, 0.5, 4.5, 5.0]).reshape(2, 1, 2, 1).repeat(1, 2, 1, 1) + expected_lse = torch.log(torch.tensor([1.0, 2.0, 2.0, 3.0])).reshape(2, 1, 2) + expected_lse = expected_lse.repeat(1, 2, 1) + torch.testing.assert_close(out, expected, atol=1.0e-6, rtol=0.0) + torch.testing.assert_close(lse, expected_lse, atol=1.0e-6, rtol=0.0) + + +def test_merge_order_uses_global_block_index_not_arrival_order(): + op = DeterministicCPAttentionReferenceOp() + q, k, v = _qkv(1, 6, 6, seed=4) + first = op.local_partial_state( + q, + k[:, :, :3], + v[:, :, :3], + q_start=0, + k_start=0, + total_kv_len=6, + causal=True, + ) + second = op.local_partial_state( + q, + k[:, :, 3:], + v[:, :, 3:], + q_start=0, + k_start=3, + total_kv_len=6, + causal=True, + ) + + forward = merge_attention_partial_states([first, second]) + reversed_arrival = merge_attention_partial_states([second, first]) + assert torch.equal(forward.out, reversed_arrival.out) + assert torch.equal(forward.lse, reversed_arrival.lse) + + +def test_key_padding_mask_and_all_masked_rows_are_stable(): + op = DeterministicCPAttentionReferenceOp() + native = NativeAttentionOp() + q, k, v = _qkv(2, 6, 6, seed=5) + mask = torch.tensor( + [ + [True, True, True, False, False, False], + [False, False, False, False, False, False], + ], + dtype=torch.bool, + ) + + with _single_thread(): + out, lse = op.forward_fp32_with_lse( + q, + k, + v, + causal=False, + key_padding_mask=mask, + cp_world_size=2, + kv_chunk_size=2, + ) + want = native.forward_fp32(q, k, v, causal=False, key_padding_mask=mask) + + torch.testing.assert_close(out[:1], want[:1], atol=_ATOL, rtol=0.0) + assert torch.equal(out[1], torch.zeros_like(out[1])) + assert torch.isneginf(lse[1]).all() + assert torch.isfinite(out).all() + + +def test_empty_query_and_empty_kv_edges_are_stable(): + op = DeterministicCPAttentionReferenceOp() + q_empty = torch.randn(1, 2, 0, 4, requires_grad=True) + k_empty = torch.randn(1, 1, 0, 4, requires_grad=True) + v_empty = torch.randn(1, 1, 0, 4, requires_grad=True) + out, lse = op.forward_fp32_with_lse(q_empty, k_empty, v_empty, cp_world_size=2) + assert out.shape == (1, 2, 0, 4) + assert lse.shape == (1, 2, 0) + assert out.requires_grad + out.sum().backward() + assert torch.equal(q_empty.grad, torch.zeros_like(q_empty)) + assert torch.equal(k_empty.grad, torch.zeros_like(k_empty)) + assert torch.equal(v_empty.grad, torch.zeros_like(v_empty)) + + q = torch.randn(1, 2, 3, 4) + out, lse = op.forward_fp32_with_lse(q, k_empty, v_empty, causal=False, cp_world_size=4) + assert torch.equal(out, torch.zeros_like(out)) + assert torch.isneginf(lse).all() + + +def test_empty_kv_backward_returns_zero_grads(): + op = DeterministicCPAttentionReferenceOp() + q = torch.randn(1, 2, 3, 4, requires_grad=True) + k = torch.randn(1, 1, 0, 4, requires_grad=True) + v = torch.randn(1, 1, 0, 4, requires_grad=True) + + out = op.forward_fp32(q, k, v, causal=False, cp_world_size=4) + assert out.requires_grad + out.sum().backward() + + assert torch.equal(q.grad, torch.zeros_like(q)) + assert torch.equal(k.grad, torch.zeros_like(k)) + assert torch.equal(v.grad, torch.zeros_like(v)) + + +def test_bf16_forward_uses_fp32_merge_then_final_write(): + op = DeterministicCPAttentionReferenceOp() + q, k, v = _qkv(2, 8, 8, seed=6, dtype=torch.bfloat16) + + out, lse = op.forward_with_lse(q, k, v, causal=True, cp_world_size=2, kv_chunk_size=2) + fp32_out, fp32_lse = op.forward_fp32_with_lse( + q, + k, + v, + causal=True, + cp_world_size=2, + kv_chunk_size=2, + ) + + assert out.dtype == torch.bfloat16 + assert lse.dtype == torch.float32 + assert torch.equal(out, fp32_out.to(torch.bfloat16)) + assert torch.equal(lse, fp32_lse) + + +def test_cp2_chunked_gradients_match_cp1_reference(): + op = DeterministicCPAttentionReferenceOp() + q, k, v = _qkv(1, 5, 5, seed=12, heads=4, kv_heads=2, dim=8) + q_ref = q.detach().clone().requires_grad_(True) + k_ref = k.detach().clone().requires_grad_(True) + v_ref = v.detach().clone().requires_grad_(True) + q_cp = q.detach().clone().requires_grad_(True) + k_cp = k.detach().clone().requires_grad_(True) + v_cp = v.detach().clone().requires_grad_(True) + gen = torch.Generator().manual_seed(13) + dy = torch.randn(1, 4, 5, 8, generator=gen) + + with _single_thread(): + out_ref = op.forward_fp32(q_ref, k_ref, v_ref, causal=True, cp_world_size=1) + out_cp = op.forward_fp32( + q_cp, + k_cp, + v_cp, + causal=True, + cp_world_size=2, + kv_chunk_size=2, + ) + out_ref.backward(dy) + out_cp.backward(dy) + + torch.testing.assert_close(out_cp, out_ref, atol=_ATOL, rtol=0.0) + torch.testing.assert_close(q_cp.grad, q_ref.grad, atol=1.0e-5, rtol=0.0) + torch.testing.assert_close(k_cp.grad, k_ref.grad, atol=1.0e-5, rtol=0.0) + torch.testing.assert_close(v_cp.grad, v_ref.grad, atol=1.0e-5, rtol=0.0) + + +def test_backward_report_cp2_prefill_matches_cp1_reference(): + q, k, v = _qkv(1, 5, 5, seed=15, heads=4, kv_heads=2, dim=8) + dout = torch.randn(1, 4, 5, 8, generator=torch.Generator().manual_seed(16)) + + with _single_thread(): + report = compare_cp_attention_backward( + q, + k, + v, + dout, + causal=True, + candidate_cp_world_size=2, + output_dtype=torch.float32, + ) + + assert report.reference_name == "cp1_backward_reference" + drift = report.drifts[0] + assert drift.candidate_name == "cp2_backward" + assert drift.dq.max_abs <= _GRAD_ATOL + assert drift.dk.max_abs <= _GRAD_ATOL + assert drift.dv.max_abs <= _GRAD_ATOL + assert drift.out.max_abs <= _ATOL + assert drift.lse.max_abs <= _ATOL + assert len(drift.per_rank) == 2 + assert drift.per_rank[0].dq.active_count > 0 + assert drift.per_rank[1].dk.active_count > 0 + assert drift.provenance["saved_forward_state"][0] == "out" + assert drift.provenance["merge_order"] == "global_block_index" + assert drift.provenance["te_backward_oracle"] == "not_used" + assert drift.provenance["decode_backward"] == "not_supported" + json.dumps(report.to_dict()) + + +def test_backward_report_cp2_chunked_prefill_matches_cp1_reference(): + q, k, v = _qkv(1, 6, 6, seed=17, heads=4, kv_heads=2, dim=8) + dout = torch.randn(1, 4, 6, 8, generator=torch.Generator().manual_seed(18)) + + with _single_thread(): + report = compare_cp_attention_backward( + q, + k, + v, + dout, + causal=True, + candidate_cp_world_size=2, + candidate_kv_chunk_size=2, + output_dtype=torch.float32, + ) + + drift = report.drifts[0] + assert drift.candidate_name == "cp2_chunked_backward" + assert drift.dq.max_abs <= _GRAD_ATOL + assert drift.dk.max_abs <= _GRAD_ATOL + assert drift.dv.max_abs <= _GRAD_ATOL + assert drift.provenance["attention_mode"] == "chunked_prefill" + assert drift.provenance["kv_chunk_size"] == 2 + assert drift.provenance["requested_split_kv_policy"] == "fixed" + assert drift.provenance["actual_split_kv_plans"] == [ + { + "owner_cp_rank": 0, + "requested_split_kv_policy": "fixed", + "requested_split_kv_size": 2, + "actual_split_kv_policy": "fixed", + "actual_split_kv_size": 2, + "actual_split_kv_count": 2, + "actual_split_boundaries": [[0, 2], [2, 3]], + "split_kv_merge_order": "global_block_index", + "split_kv_accum_dtype": "fp32", + "split_kv_downcast_at": "final_write", + "split_kv_backend": "deterministic_cp_backward_reference", + "split_kv_plan_source": "reference_execution", + "split_kv_fallback": False, + "split_kv_fallback_reason": None, + }, + { + "owner_cp_rank": 1, + "requested_split_kv_policy": "fixed", + "requested_split_kv_size": 2, + "actual_split_kv_policy": "fixed", + "actual_split_kv_size": 2, + "actual_split_kv_count": 2, + "actual_split_boundaries": [[3, 5], [5, 6]], + "split_kv_merge_order": "global_block_index", + "split_kv_accum_dtype": "fp32", + "split_kv_downcast_at": "final_write", + "split_kv_backend": "deterministic_cp_backward_reference", + "split_kv_plan_source": "reference_execution", + "split_kv_fallback": False, + "split_kv_fallback_reason": None, + }, + ] + + +def test_split_kv_plan_never_crosses_cp_owner_boundaries(): + plans = split_kv_execution_plan_provenance( + 10, + cp_world_size=3, + kv_chunk_size=3, + backend="test-reference", + ) + + assert [plan["actual_split_boundaries"] for plan in plans] == [ + [[0, 3], [3, 4]], + [[4, 7]], + [[7, 10]], + ] + assert [plan["owner_cp_rank"] for plan in plans] == [0, 1, 2] + + +def test_backward_report_preserves_post_rope_position_metadata(): + rope = NativeRoPEOp() + pre_rope_q, pre_rope_k, v = _qkv(2, 5, 5, seed=19, heads=4, kv_heads=2, dim=8) + position_offsets = torch.tensor([23, 101], dtype=torch.long) + positions = position_offsets[:, None] + torch.arange(pre_rope_q.size(2), dtype=torch.long) + q = rope.forward_fp32(pre_rope_q, positions, theta=1_000_000.0) + k = rope.forward_fp32(pre_rope_k, positions, theta=1_000_000.0) + dout = torch.randn(2, 4, 5, 8, generator=torch.Generator().manual_seed(20)) + + with _single_thread(): + report = compare_cp_attention_backward( + q, + k, + v, + dout, + causal=True, + query_position_offsets=position_offsets, + key_position_offsets=position_offsets, + candidate_cp_world_size=2, + candidate_kv_chunk_size=2, + output_dtype=torch.float32, + ) + + drift = report.drifts[0] + assert drift.dq.max_abs <= _GRAD_ATOL + assert drift.dk.max_abs <= _GRAD_ATOL + assert drift.dv.max_abs <= _GRAD_ATOL + + +def test_qwen3_8b_local_tp2_cp2_bf16_backward_report_smoke(): + # Qwen3-8B global Hq/Hkv is 32/8. A TP=2 local shard owns 16/4 heads. + q, k, v = _qkv( + 1, + 4, + 4, + seed=21, + dtype=torch.bfloat16, + heads=16, + kv_heads=4, + dim=_HEAD_DIM, + ) + dout = torch.randn( + 1, + 16, + 4, + _HEAD_DIM, + generator=torch.Generator().manual_seed(22), + dtype=torch.bfloat16, + ) + + with _single_thread(): + report = compare_cp_attention_backward( + q, + k, + v, + dout, + causal=True, + candidate_cp_world_size=2, + candidate_kv_chunk_size=2, + output_dtype=torch.bfloat16, + ) + + drift = report.drifts[0] + assert drift.provenance["q_dtype"] == "bfloat16" + assert drift.provenance["output_dtype"] == "bfloat16" + assert drift.provenance["downcast_at"] == "final_write" + assert drift.dq.max_abs <= 5.0e-2 + assert drift.dk.max_abs <= 5.0e-2 + assert drift.dv.max_abs <= 5.0e-2 + + +def test_backward_report_validates_dout_shape_and_dtype(): + op = DeterministicCPAttentionReferenceOp() + q, k, v = _qkv(1, 4, 4, seed=23, heads=4, kv_heads=2, dim=8) + + with pytest.raises(ValueError, match="dout must have shape"): + op.backward_reference(q, k, v, torch.randn(1, 4, 3, 8), cp_world_size=2) + + with pytest.raises(ValueError, match="dout must be a real floating-point tensor"): + op.backward_reference( + q, + k, + v, + torch.ones(1, 4, 4, 8, dtype=torch.long), + cp_world_size=2, + ) + + with pytest.raises(ValueError, match="dout must have the same dtype"): + op.backward_reference( + q.to(torch.bfloat16), + k.to(torch.bfloat16), + v.to(torch.bfloat16), + torch.ones_like(q), + cp_world_size=2, + ) + + +def test_inputs_are_not_mutated(): + op = DeterministicCPAttentionReferenceOp() + q, k, v = _qkv(2, 6, 6, seed=7) + mask = torch.ones(2, 6, dtype=torch.bool) + qc, kc, vc, mc = q.clone(), k.clone(), v.clone(), mask.clone() + + op.forward_fp32_with_lse(q, k, v, causal=True, key_padding_mask=mask, cp_world_size=2) + + assert torch.equal(q, qc) + assert torch.equal(k, kc) + assert torch.equal(v, vc) + assert torch.equal(mask, mc) + + +@pytest.mark.parametrize( + ("kwargs", "message"), + [ + ({"cp_world_size": 0}, "cp_world_size"), + ({"cp_world_size": 2, "kv_chunk_size": 0}, "kv_chunk_size"), + ], +) +def test_invalid_parallelism_arguments_raise(kwargs, message): + op = DeterministicCPAttentionReferenceOp() + q, k, v = _qkv(1, 4, 4, seed=8) + with pytest.raises(ValueError, match=message): + op.forward_fp32_with_lse(q, k, v, causal=True, **kwargs) + + +def test_invalid_gqa_and_mask_shapes_raise(): + op = DeterministicCPAttentionReferenceOp() + q = torch.randn(1, 6, 4, _HEAD_DIM) + k = torch.randn(1, 4, 4, _HEAD_DIM) + v = torch.randn(1, 4, 4, _HEAD_DIM) + with pytest.raises(ValueError, match="not divisible"): + op.forward_fp32_with_lse(q, k, v, causal=True) + + q, k, v = _qkv(1, 4, 4, seed=9) + with pytest.raises(ValueError, match="key_padding_mask"): + op.forward_fp32_with_lse(q, k, v, key_padding_mask=torch.ones(1, 3, dtype=torch.bool)) + with pytest.raises(ValueError, match="key_padding_mask"): + op.forward_fp32_with_lse(q, k, v, key_padding_mask=torch.ones(1, 4)) + with pytest.raises(ValueError, match="query_position_offsets"): + op.forward_fp32_with_lse( + q, + k, + v, + query_position_offsets=torch.ones(2, dtype=torch.long), + ) + with pytest.raises(ValueError, match="key_position_offsets"): + op.forward_fp32_with_lse( + q, + k, + v, + key_position_offsets=torch.ones(1, dtype=torch.float32), + ) + + +@pytest.mark.parametrize("scale", [0.0, -1.0, float("nan"), float("inf"), True, "bad"]) +def test_invalid_scale_fails_before_attention_math(scale): + op = DeterministicCPAttentionReferenceOp() + q, k, v = _qkv(1, 4, 4, seed=24, heads=4, kv_heads=2, dim=8) + + with pytest.raises(ValueError, match="scale must be a positive finite number"): + op.forward_fp32_with_lse(q, k, v, scale=scale) + + +def test_qkv_dtype_and_floating_contract_fails_closed(): + op = DeterministicCPAttentionReferenceOp() + q, k, v = _qkv(1, 4, 4, seed=25, heads=4, kv_heads=2, dim=8) + + with pytest.raises(ValueError, match="same dtype"): + op.forward_fp32_with_lse(q, k.to(torch.bfloat16), v) + with pytest.raises(ValueError, match="real floating-point"): + op.forward_fp32_with_lse(q.to(torch.long), k.to(torch.long), v.to(torch.long)) + + +@pytest.mark.parametrize("kwargs", [{"cp_world_size": True}, {"kv_chunk_size": True}]) +def test_boolean_parallelism_arguments_fail_closed(kwargs): + op = DeterministicCPAttentionReferenceOp() + q, k, v = _qkv(1, 4, 4, seed=26, heads=4, kv_heads=2, dim=8) + + with pytest.raises(ValueError): + op.forward_fp32_with_lse(q, k, v, **kwargs) + + +@pytest.mark.parametrize("output_dtype", [torch.long, torch.complex64, "fp32"]) +def test_nonfloating_output_dtype_fails_closed(output_dtype): + op = DeterministicCPAttentionReferenceOp() + q, k, v = _qkv(1, 4, 4, seed=27, heads=4, kv_heads=2, dim=8) + + with pytest.raises(ValueError, match="output_dtype must be a real floating-point"): + op.forward_with_lse(q, k, v, output_dtype=output_dtype) + + +def test_partial_states_must_remain_fp32_and_colocated(): + out = torch.zeros(1, 1, 1, 1, dtype=torch.bfloat16) + lse = torch.zeros(1, 1, 1, dtype=torch.float32) + + with pytest.raises(ValueError, match="must remain FP32"): + AttentionPartialState(out=out, lse=lse, block_start=0, block_end=1) + + +def test_overlapping_partial_ranges_raise(): + out = torch.zeros(1, 1, 1, 1) + lse = torch.zeros(1, 1, 1) + with pytest.raises(ValueError, match="overlap"): + merge_attention_partial_states( + [ + AttentionPartialState(out=out, lse=lse, block_start=0, block_end=3), + AttentionPartialState(out=out, lse=lse, block_start=2, block_end=4), + ] + ) + + +def test_gapped_partial_ranges_raise(): + out = torch.zeros(1, 1, 1, 1) + lse = torch.zeros(1, 1, 1) + with pytest.raises(ValueError, match="gap-free"): + merge_attention_partial_states( + [ + AttentionPartialState(out=out, lse=lse, block_start=0, block_end=2), + AttentionPartialState(out=out, lse=lse, block_start=3, block_end=4), + ] + ) + + +def test_registry_dispatches_cp_attention_reference(): + assert isinstance(kernel_registry.get_op("cp_attention"), DeterministicCPAttentionReferenceOp) + + +def test_strict_reference_is_bitwise_across_batch_cp_and_backward(): + q, k, v = _qkv(2, 4, 8, seed=41, heads=4, kv_heads=2, dim=8) + dout = torch.randn_like(q) + op = DeterministicCPAttentionReferenceOp(strict_bitwise=True) + + cp1_out, cp1_lse = op.forward_with_lse(q, k, v, cp_world_size=1, kv_chunk_size=3) + cp2_out, cp2_lse = op.forward_with_lse(q, k, v, cp_world_size=2, kv_chunk_size=3) + single_out, single_lse = op.forward_with_lse( + q[:1], k[:1], v[:1], cp_world_size=1, kv_chunk_size=3 + ) + assert torch.equal(cp1_out, cp2_out) + assert torch.equal(cp1_lse, cp2_lse) + assert torch.equal(cp1_out[:1], single_out) + assert torch.equal(cp1_lse[:1], single_lse) + + cp1 = op.backward_reference(q, k, v, dout, cp_world_size=1, kv_chunk_size=3) + cp2 = op.backward_reference(q, k, v, dout, cp_world_size=2, kv_chunk_size=3) + assert torch.equal(cp1.gradients.dq, cp2.gradients.dq) + assert torch.equal(cp1.gradients.dk, cp2.gradients.dk) + assert torch.equal(cp1.gradients.dv, cp2.gradients.dv) diff --git a/tests/test_cp_attention_transformer_engine.py b/tests/test_cp_attention_transformer_engine.py new file mode 100644 index 00000000..d98e31a1 --- /dev/null +++ b/tests/test_cp_attention_transformer_engine.py @@ -0,0 +1,48 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Optional Transformer Engine oracle tests for CP attention merging.""" + +from __future__ import annotations + +import importlib + +import pytest +import torch + +from rl_engine.kernels.ops.pytorch.attention.cp_attention import ( + AttentionPartialState, + merge_attention_partial_states, +) + + +def _te_context_parallel_module(): + try: + return importlib.import_module( + "transformer_engine.pytorch.attention.dot_product_attention.context_parallel" + ) + except (ImportError, OSError, RuntimeError) as exc: + pytest.skip(f"Transformer Engine context-parallel attention is unavailable: {exc}") + + +def test_cp_attention_merge_matches_transformer_engine_corrections(): + te_cp = _te_context_parallel_module() + gen = torch.Generator().manual_seed(238) + out_a = torch.randn(2, 3, 5, 4, generator=gen) + out_b = torch.randn(2, 3, 5, 4, generator=gen) + lse_a = torch.randn(2, 3, 5, generator=gen) + lse_b = torch.randn(2, 3, 5, generator=gen) + + ours = merge_attention_partial_states( + [ + AttentionPartialState(out=out_b, lse=lse_b, block_start=5, block_end=9), + AttentionPartialState(out=out_a, lse=lse_a, block_start=0, block_end=5), + ] + ) + + te_lse = lse_a.clone() + te_cp.flash_attn_fwd_softmax_lse_correction(te_lse, lse_b) + te_out = te_cp.flash_attn_fwd_out_correction_init(out_a.clone(), te_lse, lse_a, seq_dim=2) + te_cp.flash_attn_fwd_out_correction(te_out, out_b, te_lse, lse_b, seq_dim=2) + + torch.testing.assert_close(ours.lse, te_lse, atol=1.0e-6, rtol=0.0) + torch.testing.assert_close(ours.out, te_out, atol=1.0e-6, rtol=0.0) diff --git a/tests/test_cross_config_cli.py b/tests/test_cross_config_cli.py new file mode 100644 index 00000000..d5a3ce4a --- /dev/null +++ b/tests/test_cross_config_cli.py @@ -0,0 +1,145 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from __future__ import annotations + +import json +import sys +from pathlib import Path + +import pytest +import torch + +import rl_engine.alignment.cross_config.__main__ as cli_main +from rl_engine.alignment.cross_config.artifacts import ArtifactStore + +_REPOSITORY_ROOT = Path(__file__).resolve().parents[1] +_EXAMPLES = _REPOSITORY_ROOT / "examples" +_CPU_RUNTIME_MODULE = "rl_engine.alignment.testing.cpu_cross_config" + + +def _summary(captured: str) -> dict: + summaries = [] + for line in captured.splitlines(): + try: + value = json.loads(line) + except json.JSONDecodeError: + continue + if value.get("schema_version") == "cross_config.cli_summary.v1": + summaries.append(value) + assert len(summaries) == 1 + return summaries[0] + + +def test_run_uses_only_cpu_and_resumes_when_cuda_is_available(tmp_path, capsys, monkeypatch): + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + config_path = _EXAMPLES / "cross_config_s0_cpu_smoke.json" + plan_argv = [ + "plan", + str(config_path), + "--output-root", + str(tmp_path), + ] + argv = [ + "run", + str(config_path), + "--runtime", + "cpu-smoke", + "--allow-smoke-operators", + "--output-root", + str(tmp_path), + "--timeout-seconds", + "10", + ] + + assert cli_main.main(plan_argv) == 0 + planned = _summary(capsys.readouterr().out) + experiment_path = Path(planned["artifact_dir"]) / "experiment.json" + plan_path = Path(planned["artifact_dir"]) / "plan.jsonl" + planned_experiment = experiment_path.read_bytes() + planned_cases = plan_path.read_bytes() + stored_config = json.loads(planned_experiment) + stored_row = json.loads(planned_cases) + assert stored_config["schema_version"] == "cross_config.experiment_config.v1" + assert stored_row["schema_version"] == "cross_config.execution_plan_entry.v1" + assert stored_row["case"]["execution_binding"]["operators"] == stored_row["operators"] + + assert cli_main.main(argv) == 0 + captured = capsys.readouterr() + first = _summary(captured.out) + assert first["status"] == "pass" + assert first["runtime"] == "cpu-smoke" + assert "actual backends rollout=smoke_only.logp_reference" in captured.err + assert "training=smoke_only.logp_reference" in captured.err + assert "worst sample/token=[0, 3]" in captured.err + assert first["cases"] + assert all(case["status"] == "pass" for case in first["cases"]) + assert all(case["resumed"] is False for case in first["cases"]) + assert experiment_path.read_bytes() == planned_experiment + assert plan_path.read_bytes() == planned_cases + + store = ArtifactStore(tmp_path) + for case in first["cases"]: + attempt_dir = Path(case["attempt_dir"]) + assert (attempt_dir / "COMPLETE").is_file() + actual = json.loads((attempt_dir / "actual.json").read_text(encoding="utf-8")) + assert actual["environment"]["execution_devices"] == { + "rollout": "cpu", + "training": "cpu", + } + for name in ("score_rollout.pt", "score_training.pt"): + bundle = store.load_tensor_bundle(attempt_dir / name) + assert bundle["tensors"]["selected_logprobs"].device.type == "cpu" + + assert cli_main.main(argv) == 0 + resumed = _summary(capsys.readouterr().out) + assert resumed["status"] == "pass" + assert [case["attempt_id"] for case in resumed["cases"]] == [ + case["attempt_id"] for case in first["cases"] + ] + assert all(case["resumed"] is True for case in resumed["cases"]) + + +@pytest.mark.parametrize( + ("filename", "expected_cases"), + [ + ("cross_config_s1_distributed_smoke.json", 5), + ("cross_config_s2_vllm_tp_vs_fsdp.json", 10), + ("cross_config_s3_qwen3_8b_tp4_cp4_bf16.json", 11), + ], +) +def test_plan_records_named_scenarios_without_loading_a_runtime( + tmp_path, + capsys, + monkeypatch, + filename, + expected_cases, +): + monkeypatch.delitem(sys.modules, _CPU_RUNTIME_MODULE, raising=False) + + def runtime_must_not_run(*args, **kwargs): + raise AssertionError(f"plan unexpectedly invoked the CPU runtime: {args!r}, {kwargs!r}") + + monkeypatch.setattr(cli_main, "_run", runtime_must_not_run) + assert ( + cli_main.main( + [ + "plan", + str(_EXAMPLES / filename), + "--output-root", + str(tmp_path), + ] + ) + == 0 + ) + + summary = _summary(capsys.readouterr().out) + artifact_dir = Path(summary["artifact_dir"]) + assert summary["status"] == "planned" + assert summary["planned_case_count"] == expected_cases + assert (artifact_dir / "experiment.json").is_file() + assert len((artifact_dir / "plan.jsonl").read_text(encoding="utf-8").splitlines()) == ( + expected_cases + ) + assert not list(artifact_dir.glob("cases/*/*")) + assert _CPU_RUNTIME_MODULE not in sys.modules diff --git a/tests/test_cross_config_contract.py b/tests/test_cross_config_contract.py new file mode 100644 index 00000000..817de626 --- /dev/null +++ b/tests/test_cross_config_contract.py @@ -0,0 +1,440 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from __future__ import annotations + +import json +from collections.abc import Mapping +from dataclasses import replace +from pathlib import Path +from typing import Any + +import pytest +import torch + +from rl_engine.alignment.cross_config.comparison import ( + compare_score_artifacts, + recompute_mismatch_mask, +) +from rl_engine.alignment.cross_config.config import ( + CONFIG_SCHEMA_VERSION, + bind_operator_selection, + load_config, +) +from rl_engine.alignment.cross_config.planner import MAX_PLAN_CASES, Planner, PlanningError +from rl_engine.alignment.cross_config.schema import ( + AlignmentStatus, + ExperimentDefinition, + InterventionSpec, + PlanningStrategy, + RuntimeProvenance, + ScoreArtifact, + ScorerSpec, + ScoreSide, + SemanticIdentitySpec, + TokenComparisonArtifact, +) +from rl_engine.kernels.gtest.tolerance import resolve_logprob_threshold + + +def _identity( + *, + checkpoint_id: str = "tiny-checkpoint", + tokenizer_policy: str = "tokenizer-v1:right-padding", + active_mask: tuple[tuple[bool, ...], ...] = ((True, False, True),), +) -> SemanticIdentitySpec: + return SemanticIdentitySpec( + checkpoint_id=checkpoint_id, + model_version="weights-v7", + tokenizer_id="tiny-tokenizer", + tokenizer_policy=tokenizer_policy, + token_ids=((11, 12, 13),), + selected_token_ids=((12, 13, 14),), + active_mask=active_mask, + attention_mask=((True, True, True),), + position_ids=((0, 1, 2),), + pre_update_state="state-before-step-9", + cache_metadata={"use_cache": False}, + packing_metadata={"packed": False}, + ) + + +def _score( + side: ScoreSide, + values: torch.Tensor, + *, + identity: SemanticIdentitySpec | None = None, + active_mask: torch.Tensor | None = None, +) -> ScoreArtifact: + identity = identity or _identity() + backend = f"test.{side.value}.selected_logprob" + return ScoreArtifact( + case_id="case-001", + attempt_id="attempt-001", + side=side, + identity=identity, + scorer=ScorerSpec( + side=side, + backend_id=f"{side.value}-scorer", + dtype="float32", + operator_overrides={"selected_logprob": backend}, + ), + selected_logprobs=values, + active_mask=( + active_mask + if active_mask is not None + else torch.tensor(identity.active_mask, dtype=torch.bool) + ), + provenance=RuntimeProvenance( + requested={"logp": {"backend": backend}}, + normalized={"logp": {"backend": backend}}, + materialized={"logp": {"backend": backend}}, + actual={"logp": {"backend": backend}}, + implementation_fingerprint=f"{side.value}-implementation-v1", + ), + ) + + +def _tensor_from_payload(payload: Mapping[str, Any]) -> torch.Tensor: + return torch.tensor(payload["values"], dtype=getattr(torch, str(payload["dtype"]))).reshape( + payload["shape"] + ) + + +def _baseline() -> dict[str, Any]: + return { + "batch": {"size": 8}, + "rollout": { + "tensor_parallel_size": 1, + "context_parallel_size": 1, + "dtype": "float32", + "enable_prefix_caching": False, + "enforce_eager": True, + }, + "training": { + "sharding": "unsharded", + "attention_backend": "eager", + "compute_dtype": "float32", + }, + "logp": {"backend": "rlkernel.reference_logp"}, + } + + +def _definition( + *, + strategy: PlanningStrategy = PlanningStrategy.ONE_AT_A_TIME, + pairwise_paths: tuple[tuple[str, str], ...] = (), +) -> ExperimentDefinition: + return ExperimentDefinition( + experiment_id="planner-test", + scenario_id="cpu-contract", + scenario={"model": "synthetic", "device": "cpu"}, + identity=_identity(), + baseline=_baseline(), + interventions=( + InterventionSpec("batch.size", (1, 4)), + InterventionSpec("rollout.dtype", ("bfloat16",)), + InterventionSpec("training.attention_backend", ("sdpa",)), + ), + strategy=strategy, + pairwise_paths=pairwise_paths, + ) + + +def _flatten(value: Mapping[str, Any], prefix: str = "") -> dict[str, Any]: + result: dict[str, Any] = {} + for key, child in value.items(): + path = f"{prefix}.{key}" if prefix else key + if isinstance(child, Mapping): + result.update(_flatten(child, path)) + else: + result[path] = child + return result + + +def _config() -> dict[str, Any]: + return { + "schema_version": CONFIG_SCHEMA_VERSION, + "experiment_id": "cpu-config-test", + "scenario_id": "cpu-smoke", + "contract_source": "ws1", + "contract_version": "current", + "strategy": "one_at_a_time", + "strict_fallback": True, + "identity": { + "checkpoint_id": "tiny", + "model_version": "weights-v1", + "tokenizer_policy": "synthetic-v1", + "token_ids": [[1, 2, 3]], + "selected_token_ids": [[0, 2, 3]], + "active_mask": [[False, True, True]], + "attention_mask": [[True, True, True]], + "pre_update_state": "iteration-0", + }, + "baseline": { + **_baseline(), + "batch": {"size": 1}, + }, + "interventions": [{"path": "batch.size", "values": [2]}], + "operators": { + "selected_logprob": { + "rollout": "rlkernel.reference_logp", + "training": { + "backend": "smoke_only.logp_offset", + "options": {"offset": 0.1}, + }, + } + }, + "scenario": {"device": "cpu", "workload": "tiny"}, + } + + +def _write_config(tmp_path: Path, value: Mapping[str, Any], name: str = "config.json") -> Path: + path = tmp_path / name + path.write_text(json.dumps(value), encoding="utf-8") + return path + + +def test_fixed_threshold_uses_only_active_tokens_and_is_reproducible_offline(): + threshold = resolve_logprob_threshold("float32") + rollout_values = torch.zeros((1, 3), dtype=torch.float32) + training_values = torch.tensor( + [[threshold * 2.0, 1_000.0, threshold * 0.5]], + dtype=torch.float32, + ) + + result = compare_score_artifacts( + _score(ScoreSide.ROLLOUT, rollout_values), + _score(ScoreSide.TRAINING, training_values), + ) + + assert result.status is AlignmentStatus.FAIL + assert result.active_token_count == 2 + assert result.mismatch_count == 1 + assert result.fixed_threshold == threshold + assert result.token_artifact is not None + assert result.token_artifact.mismatch_mask.tolist() == [[True, False, False]] + + payload = json.loads(json.dumps(result.token_artifact.to_dict())) + offline = recompute_mismatch_mask( + _tensor_from_payload(payload["rollout_logprobs"]), + _tensor_from_payload(payload["training_logprobs"]), + _tensor_from_payload(payload["active_mask"]), + float(payload["fixed_threshold"]), + ) + assert torch.equal(offline, result.token_artifact.mismatch_mask) + assert not recompute_mismatch_mask( + torch.zeros(1, dtype=torch.float64), + torch.tensor([threshold], dtype=torch.float64), + torch.tensor([True]), + threshold, + ).item() + + inactive_nonfinite = training_values.clone() + inactive_nonfinite[0, 1] = float("nan") + sanitized = compare_score_artifacts( + _score(ScoreSide.ROLLOUT, rollout_values), + _score(ScoreSide.TRAINING, inactive_nonfinite), + ) + assert sanitized.status is result.status + assert sanitized.mismatch_count == result.mismatch_count + assert sanitized.token_artifact is not None + assert sanitized.token_artifact.training_logprobs[0, 1].item() == 0.0 + json.dumps(sanitized.to_dict(), allow_nan=False) + + +def test_zero_tokens_identity_mismatch_and_invalid_scores_are_not_numerical_failures(): + empty_identity = _identity(active_mask=((False, False, False),)) + empty = compare_score_artifacts( + _score(ScoreSide.ROLLOUT, torch.zeros((1, 3)), identity=empty_identity), + _score(ScoreSide.TRAINING, torch.zeros((1, 3)), identity=empty_identity), + ) + assert empty.status is AlignmentStatus.ZERO_ACTIVE_TOKENS + assert empty.comparable is False + assert empty.passed is False + + identity = _identity() + changed_identity = replace(identity, tokenizer_policy="different-policy") + mismatched = compare_score_artifacts( + _score(ScoreSide.ROLLOUT, torch.zeros((1, 3)), identity=identity), + _score(ScoreSide.TRAINING, torch.ones((1, 3)), identity=changed_identity), + ) + assert mismatched.status is AlignmentStatus.INVALID_IDENTITY + assert "tokenizer_policy" in mismatched.identity_errors + + invalid = compare_score_artifacts( + _score(ScoreSide.ROLLOUT, torch.zeros((1, 3))), + _score(ScoreSide.TRAINING, torch.tensor([[float("nan"), 0.0, 0.0]])), + ) + assert invalid.status is AlignmentStatus.INVALID_ARTIFACT + assert invalid.comparable is False + + with pytest.raises(ValueError, match="finite and non-negative"): + TokenComparisonArtifact( + rollout_logprobs=torch.zeros(1), + training_logprobs=torch.zeros(1), + active_mask=torch.ones(1, dtype=torch.bool), + absolute_diff=torch.zeros(1), + mismatch_mask=torch.zeros(1, dtype=torch.bool), + fixed_threshold=float("nan"), + ) + + +def test_planner_emits_one_stable_baseline_and_one_change_per_oat_case(): + definition = _definition() + plan = Planner().plan(definition) + baseline = _flatten(plan.cases[0].requested) + + assert len(plan.cases) == 5 + assert sum(not case.changed_paths for case in plan.cases) == 1 + for case in plan.cases[1:]: + requested = _flatten(case.requested) + changed = {path for path, value in requested.items() if value != baseline[path]} + assert changed == set(case.changed_paths) + assert len(changed) == 1 + + reordered = replace( + definition, + experiment_id="same-plan-from-another-run", + baseline={ + "logp": {"backend": "reference"}, + "training": { + "compute_dtype": "fp32", + "attention_backend": "eager", + "sharding": "unsharded", + }, + "rollout": { + "enforce_eager": True, + "enable_prefix_caching": False, + "dtype": "fp32", + "context_parallel_size": 1, + "tensor_parallel_size": 1, + }, + "batch": {"size": 8}, + }, + ) + assert [case.case_id for case in plan.cases] == [ + case.case_id for case in Planner().plan(reordered).cases + ] + + +def test_pairwise_is_explicit_and_planning_errors_remain_structured(): + pairwise = _definition( + strategy=PlanningStrategy.PAIRWISE, + pairwise_paths=(("batch.size", "rollout.dtype"),), + ) + pairwise_cases = [ + case for case in Planner().plan(pairwise).cases if len(case.changed_paths) == 2 + ] + assert len(pairwise_cases) == 2 + assert all(case.changed_paths == ("batch.size", "rollout.dtype") for case in pairwise_cases) + + with pytest.raises(PlanningError) as not_enabled: + Planner().plan( + replace( + _definition(), + pairwise_paths=(("batch.size", "rollout.dtype"),), + ) + ) + assert {issue.code for issue in not_enabled.value.issues} == {"PAIRWISE_NOT_ENABLED"} + + invalid_requests = ( + ({"logp": {"tp_layout": "arbitrary"}}, "DERIVED_KNOB"), + ({"batch": {"size": True}}, "UNSUPPORTED_VALUE"), + ({"rollout": {"unknown": 1}}, "UNSUPPORTED_PATH"), + ) + for requested, expected_code in invalid_requests: + with pytest.raises(PlanningError) as invalid: + Planner().normalize_requested(requested) + assert invalid.value.issues[0].code == expected_code + + incomplete = replace( + _definition(), + baseline={key: value for key, value in _baseline().items() if key != "training"}, + ) + with pytest.raises(PlanningError) as missing: + Planner().plan(incomplete) + assert {issue.path for issue in missing.value.issues} == { + "training.attention_backend", + "training.compute_dtype", + "training.sharding", + } + assert all(issue.code == "MISSING_BASELINE_VALUE" for issue in missing.value.issues) + + oversized = replace( + _definition(), + interventions=(InterventionSpec("batch.size", tuple(range(1, MAX_PLAN_CASES + 2))),), + ) + with pytest.raises(PlanningError) as too_large: + Planner().plan(oversized) + assert too_large.value.issues[0].code == "PLAN_TOO_LARGE" + + +def test_versioned_config_loads_and_binds_target_specific_operators(tmp_path: Path): + loaded = load_config(_write_config(tmp_path, _config())) + base_case = loaded.plan().cases[0] + selection = loaded.operators_for(base_case) + bound = bind_operator_selection(base_case, selection) + + assert loaded.schema_version == CONFIG_SCHEMA_VERSION + assert loaded.definition.strategy is PlanningStrategy.ONE_AT_A_TIME + assert selection.rollout_backend == "rlkernel.reference_logp" + assert selection.training_backend == "smoke_only.logp_offset" + assert selection.training_options == {"offset": 0.1} + assert bound == bind_operator_selection(base_case, selection) + assert bound.case_id != base_case.case_id + assert bound.requested == base_case.requested + assert bound.execution_binding["operators"] == selection.to_dict() + + +def test_config_rejects_schema_escape_hatches_and_incomplete_operator_coverage( + tmp_path: Path, +): + wrong_schema = _config() + wrong_schema["schema_version"] = "cross_config.experiment_config.v999" + + unknown_key = _config() + unknown_key["strict_falback"] = True + + threshold_override = _config() + threshold_override["scenario"]["nested"] = {"threshold": 999.0} + + scenario_policy = _config() + scenario_policy["scenario"]["execution"] = "run" + + conflicting_axis = _config() + conflicting_axis["interventions"].append( + {"path": "logp.backend", "values": ["smoke_only.logp_offset"]} + ) + + incomplete_targets = _config() + del incomplete_targets["operators"]["selected_logprob"]["training"] + + invalid_configs = ( + (wrong_schema, "unsupported cross-configuration config schema"), + (unknown_key, "unknown config keys"), + (threshold_override, "fixed numerical-contract threshold"), + (scenario_policy, "scenario is metadata only"), + (conflicting_axis, "cannot be combined with logp.backend interventions"), + (incomplete_targets, "selected_logprob.training"), + ) + for index, (value, message) in enumerate(invalid_configs): + with pytest.raises(ValueError, match=message): + load_config(_write_config(tmp_path, value, f"invalid-{index}.json")) + + duplicate_key = tmp_path / "duplicate.json" + duplicate_key.write_text( + '{"schema_version":"cross_config.experiment_config.v1",' + '"schema_version":"cross_config.experiment_config.v1"}', + encoding="utf-8", + ) + with pytest.raises(ValueError, match="duplicate JSON key"): + load_config(duplicate_key) + + overflow = tmp_path / "overflow.json" + overflow.write_text( + json.dumps(_config()).replace('"size": 1', '"size": 1e400'), + encoding="utf-8", + ) + with pytest.raises(ValueError, match="non-finite JSON number"): + load_config(overflow) diff --git a/tests/test_cross_config_drift_report.py b/tests/test_cross_config_drift_report.py new file mode 100644 index 00000000..9996591b --- /dev/null +++ b/tests/test_cross_config_drift_report.py @@ -0,0 +1,334 @@ +from __future__ import annotations + +import json +from pathlib import Path + +import pytest +import torch + +import rl_engine.alignment.cross_config.__main__ as cli_main +from rl_engine.alignment.cross_config.artifacts import ArtifactStore +from rl_engine.alignment.cross_config.drift_report import ( + build_cross_config_attempt_report, + build_drift_report, + build_drift_trace, + load_drift_bundle, + render_drift_report, + render_drift_report_image, + write_drift_bundle, + write_drift_report, + write_drift_report_image, + write_drift_trace, +) + + +def _artifacts(*, with_timestamp: bool = False): + sample = { + "sample_position": 0, + "sample_index": 17, + "rollout_id": 23, + "batch_layout_fingerprint": "layout-1", + } + if with_timestamp: + sample.update(start_ts=10.0, end_ts=10.5) + manifest = { + "mode": "audit", + "samples": [sample], + "batch_invariance_cases": [{"case": "same_sample_alone"}], + "validation": {"warnings": [], "failures": []}, + "runtime_provenance": {"operator": "linear_logp"}, + } + cube = { + "mode": "audit", + "rank": 0, + "axes": {"dtype": "bf16", "cp": 1, "logp_backend": "rlk.linear_logp.fast"}, + "metrics": { + "active_token_count": 2, + "max_abs_dlogp": 0.125, + "warning_count": 1, + "metadata_warning_count": 0, + "metadata_failure_count": 0, + }, + "worst_token": {"abs_dlogp": 0.125, "sample_position": 0, "token_position": 4}, + "metadata_validation": {"warnings": [], "failures": []}, + "runtime_provenance": { + "actual_backend": "rl_engine.linear_logp", + "fallback": False, + }, + } + return manifest, cube + + +def _completed_attempt(root: Path) -> Path: + store = ArtifactStore(root) + attempt = store.create_attempt("drift-report", "case-1") + envelope = {"case_id": "case-1", "attempt_id": attempt.name} + requested = { + "rollout": { + "tensor_parallel_size": 2, + "context_parallel_size": 2, + "dtype": "bfloat16", + }, + "training": {"compute_dtype": "bfloat16", "sharding": "tp2-cp2"}, + } + store.write_json( + attempt, + "requested", + { + **envelope, + "schema_version": "cross_config.requested.v1", + "case": {"requested": requested}, + }, + ) + store.write_json( + attempt, + "materialized", + { + **envelope, + "schema_version": "cross_config.materialized_envelope.v1", + "materialized_case": {"case": {"requested": requested}}, + }, + ) + store.write_json( + attempt, + "identity", + { + **envelope, + "schema_version": "cross_config.identity_envelope.v1", + "identity": {"token_ids": [[11, 12], [21, 22]]}, + }, + ) + store.write_json( + attempt, + "actual", + { + **envelope, + "schema_version": "cross_config.actual.v1", + "operator_source": "exact_resolution_and_instance", + "execution_fingerprint": "execution-sha", + "environment_fingerprint": "environment-sha", + "rollout": {"backend_id": "rlkernel.attention.deterministic.v1"}, + "training": {"backend_id": "rlkernel.ffn.qwen3.deterministic.v1"}, + }, + ) + store.write_json( + attempt, + "comparison", + { + **envelope, + "schema_version": "cross_config.alignment_result.v1", + "status": "pass", + "comparable": True, + "passed": True, + "mismatch_count": 0, + "fixed_threshold": 0.0, + "contract_fingerprint": "contract-sha", + "diagnostics": {}, + }, + ) + scores = { + "selected_logprobs": torch.zeros((2, 2)), + "active_mask": torch.tensor([[False, True], [False, True]]), + } + + def write_tensor_bundle(name: str, tensors: dict[str, torch.Tensor]) -> None: + torch.save( + {"schema_version": 1, "tensors": tensors, "metadata": envelope}, + attempt / f"{name}.pt", + ) + + for name in ("score_rollout", "score_training"): + write_tensor_bundle(name, scores) + write_tensor_bundle( + "token_diffs", + { + "rollout_logprobs": torch.zeros((2, 2)), + "training_logprobs": torch.zeros((2, 2)), + "active_mask": scores["active_mask"], + "absolute_diff": torch.zeros((2, 2)), + "mismatch_mask": torch.zeros((2, 2), dtype=torch.bool), + }, + ) + store.complete_attempt( + attempt, + summary={ + **envelope, + "schema_version": "cross_config.complete.v1", + "status": "pass", + }, + ) + return attempt + + +@pytest.mark.unit +def test_report_uses_ordinal_timeline_without_fabricating_timestamps(): + manifest, cube = _artifacts() + + report = build_drift_report(replay_manifest=manifest, result_cube=cube) + + assert report["timeline_mode"] == "ordinal_diagnostic" + assert "not elapsed time" in report["timeline_note"] + assert {event["lane"] for event in report["events"]} == { + "Training audit", + "Rollout samples", + "Operator / backend", + "Drift markers", + } + assert report["status"] == "warning" + + +@pytest.mark.unit +def test_report_prefers_actual_backend_and_timestamp_mode(): + manifest, cube = _artifacts(with_timestamp=True) + cube["runtime_provenance"] = { + "requested_backend": "registry", + "actual_backend": "native.linear_logp", + "fallback": True, + } + + report = build_drift_report(replay_manifest=manifest, result_cube=cube) + operator = next(event for event in report["events"] if event["id"] == "operator-backend") + + assert report["timeline_mode"] == "timestamp" + assert operator["label"] == "native.linear_logp" + assert operator["status"] == "warning" + + +@pytest.mark.unit +def test_rendered_report_is_self_contained_and_escapes_details(tmp_path: Path): + manifest, cube = _artifacts() + manifest["validation"]["warnings"] = [{"code": "bad<&", "message": "value "}] + cube["metadata_validation"] = manifest["validation"] + + report = build_drift_report( + replay_manifest=manifest, + result_cube=cube, + title="", + ) + html = render_drift_report(report) + output = write_drift_report(report, tmp_path / "drift.html") + + assert output.exists() + assert "" not in html + assert "value " not in html + assert "ordinal_diagnostic" in html + assert "Operator / backend" in html + assert "http://" not in html + assert "https://" not in html + assert "detail.innerHTML" not in html + assert "detail.replaceChildren" in html + + +@pytest.mark.unit +def test_static_report_image_is_shareable_png(tmp_path: Path): + pytest.importorskip("PIL") + manifest, cube = _artifacts() + report = build_drift_report(replay_manifest=manifest, result_cube=cube) + + image = render_drift_report_image(report) + assert image.size == (2400, 1680) + + output = write_drift_report_image(report, tmp_path / "drift.png") + assert output.read_bytes()[:8] == b"\x89PNG\r\n\x1a\n" + + +@pytest.mark.unit +def test_consistency_trace_is_expandable_chrome_trace_json(tmp_path: Path): + manifest, cube = _artifacts() + report = build_drift_report(replay_manifest=manifest, result_cube=cube) + + trace = build_drift_trace(report) + assert trace["metadata"]["timeline_mode"] == "ordinal_diagnostic" + assert any( + event.get("ph") == "M" and event.get("name") == "thread_name" + for event in trace["traceEvents"] + ) + assert any( + event.get("ph") == "X" and event.get("cat") == "consistency.audit" + for event in trace["traceEvents"] + ) + assert any( + event.get("ph") == "I" and event.get("cat") == "consistency.drift" + for event in trace["traceEvents"] + ) + + output = write_drift_trace(report, tmp_path / "drift.json") + assert output.read_text(encoding="utf-8").startswith('{\n "traceEvents"') + + +@pytest.mark.unit +def test_consistency_bundle_contains_report_trace_and_preview(tmp_path: Path): + manifest, cube = _artifacts() + report = build_drift_report(replay_manifest=manifest, result_cube=cube) + + output = write_drift_bundle(report, tmp_path / "drift.rlk-drift") + assert output.exists() + + import zipfile + + with zipfile.ZipFile(output) as archive: + assert set(archive.namelist()) == { + "manifest.json", + "report.json", + "trace.json", + "preview.png", + } + + loaded = load_drift_bundle(output) + assert loaded["manifest"]["format"] == "rl_kernel.cross_config_drift" + assert loaded["manifest"]["preview_included"] is True + assert loaded["report"]["status"] == "warning" + assert loaded["trace"]["metadata"]["timeline_mode"] == "ordinal_diagnostic" + + +@pytest.mark.unit +def test_consistency_bundle_can_omit_preview(tmp_path: Path): + manifest, cube = _artifacts() + report = build_drift_report(replay_manifest=manifest, result_cube=cube) + + output = write_drift_bundle( + report, + tmp_path / "drift-no-preview.rlk-drift", + include_preview=False, + ) + + import zipfile + + with zipfile.ZipFile(output) as archive: + assert "preview.png" not in archive.namelist() + bundle_manifest = json.loads(archive.read("manifest.json")) + assert bundle_manifest["preview_included"] is False + assert bundle_manifest["files"] == ["manifest.json", "report.json", "trace.json"] + + +@pytest.mark.unit +def test_attempt_report_reads_only_sealed_cross_config_artifacts(tmp_path: Path): + attempt_dir = _completed_attempt(tmp_path) + + report = build_cross_config_attempt_report(attempt_dir) + + assert report["status"] == "pass" + assert report["axes"]["rollout_tp"] == 2 + assert report["axes"]["rollout_cp"] == 2 + assert report["metrics"]["active_token_count"] == 2 + assert report["runtime_provenance"]["operator_source"] == "exact_resolution_and_instance" + + +@pytest.mark.unit +def test_attempt_report_rejects_unsealed_artifact_directory(tmp_path: Path): + store = ArtifactStore(tmp_path) + attempt = store.create_attempt("unfinished", "case-1") + + with pytest.raises(Exception, match="COMPLETE"): + build_cross_config_attempt_report(attempt) + + +@pytest.mark.unit +def test_cli_writes_offline_bundle_from_sealed_attempt(tmp_path: Path, capsys): + attempt = _completed_attempt(tmp_path) + output = tmp_path / "report.rlk-drift" + + assert cli_main.main(["report", str(attempt), "--output", str(output), "--no-preview"]) == 0 + + assert output.is_file() + assert "drift report: pass" in capsys.readouterr().err diff --git a/tests/test_cross_config_runner.py b/tests/test_cross_config_runner.py new file mode 100644 index 00000000..21ad8505 --- /dev/null +++ b/tests/test_cross_config_runner.py @@ -0,0 +1,659 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from __future__ import annotations + +import json +import time +from pathlib import Path + +import pytest +import torch + +from rl_engine.alignment.cross_config.artifacts import ArtifactError, ArtifactStore +from rl_engine.alignment.cross_config.comparison import recompute_mismatch_mask +from rl_engine.alignment.cross_config.config import OperatorSelection, bind_operator_selection +from rl_engine.alignment.cross_config.operators import OperatorBridge, OperatorOverride +from rl_engine.alignment.cross_config.runner import ( + ChildScoringError, + PairedRunner, + RankCompletenessError, + RankScore, + ScoringTimeoutError, +) +from rl_engine.alignment.cross_config.runtime import RuntimeTools +from rl_engine.alignment.cross_config.schema import ( + CanonicalScoringBatch, + ExperimentCase, + ScorerSpec, + ScoreSide, + SemanticIdentitySpec, +) +from rl_engine.alignment.testing.cpu_cross_config import CpuSmokeMaterializer, run_cpu_case +from rl_engine.alignment.testing.smoke_ops import ( + SMOKE_ONLY_LOGP_OFFSET_BACKEND_ID, + SMOKE_ONLY_LOGP_REFERENCE_BACKEND_ID, +) +from rl_engine.alignment.testing.smoke_ops.smoke_only_logp_reference import SmokeOnlyLogpReference +from rl_engine.kernels.gtest.tolerance import resolve_logprob_threshold +from rl_engine.kernels.semantic_registry import OperatorRequirements + +_JSON_ARTIFACTS = ( + "requested.json", + "materialized.json", + "actual.json", + "identity.json", + "comparison.json", +) +_TENSOR_ARTIFACTS = ( + "score_rollout.pt", + "score_training.pt", + "token_diffs.pt", +) + + +class FixedRankScorer: + optimizer = None + model_state_fingerprint = "fixed-rank-scorer-state-v1" + + def __init__(self, spec: ScorerSpec, ranks): + self.spec = spec + self.ranks = tuple(ranks) + + def score(self, batch, *, batch_size, operator): + del batch_size, operator + return tuple( + RankScore( + rank=rank, + world_size=self.spec.world_size, + selected_logprobs=torch.zeros_like(batch.input_ids, dtype=torch.float32), + ) + for rank in self.ranks + ) + + +class FailingScorer(FixedRankScorer): + def score(self, batch, *, batch_size, operator): + del batch, batch_size, operator + raise RuntimeError("intentional scorer failure") + + +class SlowScorer(FixedRankScorer): + def score(self, batch, *, batch_size, operator): + time.sleep(2.0) + return super().score(batch, batch_size=batch_size, operator=operator) + + +def _identity() -> SemanticIdentitySpec: + token_ids = ( + (1, 2, 3, 4), + (2, 3, 4, 5), + (3, 4, 5, 6), + ) + selected = ( + (0, 2, 3, 4), + (0, 3, 4, 5), + (0, 4, 5, 6), + ) + active = tuple((False, True, True, True) for _ in token_ids) + attention = tuple((True, True, True, True) for _ in token_ids) + return SemanticIdentitySpec( + checkpoint_id="tiny-cpu-checkpoint", + model_version="weights-v1", + tokenizer_policy="synthetic-tokenizer-v1", + token_ids=token_ids, + selected_token_ids=selected, + active_mask=active, + attention_mask=attention, + pre_update_state="iteration-0", + ) + + +def _batch() -> CanonicalScoringBatch: + identity = _identity() + return CanonicalScoringBatch( + identity=identity, + input_ids=torch.tensor(identity.token_ids, device="cpu"), + selected_token_ids=torch.tensor(identity.selected_token_ids, device="cpu"), + active_mask=torch.tensor(identity.active_mask, device="cpu"), + attention_mask=torch.tensor(identity.attention_mask, device="cpu"), + metadata={"source": "runner-test", "device": "cpu"}, + ) + + +def _requested(*, backend: str = "rlkernel.reference_logp") -> dict[str, object]: + return { + "batch": {"size": 2}, + "rollout": { + "tensor_parallel_size": 1, + "context_parallel_size": 1, + "dtype": "float32", + "enable_prefix_caching": False, + "enforce_eager": True, + }, + "training": { + "attention_backend": "eager", + "compute_dtype": "float32", + "sharding": "unsharded", + }, + "logp": {"backend": backend}, + } + + +def _case( + *, + case_id: str = "case-runner", + backend: str = "rlkernel.reference_logp", +) -> ExperimentCase: + return ExperimentCase( + case_id=case_id, + experiment_id="runner-test", + scenario_id="S0", + identity=_identity(), + requested=_requested(backend=backend), + contract_fingerprint="contract-sha", + scenario_fingerprint="scenario-sha", + ) + + +def _topology(side: ScoreSide) -> dict[str, object]: + if side is ScoreSide.ROLLOUT: + return { + "world_size": 1, + "tensor_parallel_size": 1, + "context_parallel_size": 1, + } + return {"world_size": 1, "sharding": "unsharded"} + + +def _requirements(side: ScoreSide) -> OperatorRequirements: + return OperatorRequirements( + device="cpu", + dtype="float32", + topology=_topology(side), + alignment_properties={"deterministic": True}, + ) + + +def _operators(): + bridge = OperatorBridge() + resolved = bridge.resolve_override( + OperatorOverride.for_target( + semantic_op="selected_logprob", + backend_id="rlkernel.reference_logp", + target="both", + ), + requirements={ + "rollout": _requirements(ScoreSide.ROLLOUT), + "training": _requirements(ScoreSide.TRAINING), + }, + strict=True, + ) + instances = { + target: bridge.instantiate(resolved, target=target) for target in ("rollout", "training") + } + provenance = { + target: bridge.instance_provenance( + resolved, + target=target, + instance=instances[target], + ) + for target in ("rollout", "training") + } + return resolved, instances, provenance + + +def _materialization(case: ExperimentCase): + backend = str(case.requested["logp"]["backend"]) + backends = {"rollout": backend, "training": backend} + return RuntimeTools().materialize( + case, + CpuSmokeMaterializer( + requested_operator_backends=backends, + actual_operator_backends=backends, + ), + ) + + +def _spec(side: ScoreSide) -> ScorerSpec: + identity = _identity() + return ScorerSpec( + side=side, + backend_id="fixed_cpu_teacher_forcing", + dtype="float32", + device="cpu", + world_size=1, + topology=_topology(side), + construction_options={ + "checkpoint_id": identity.checkpoint_id, + "model_version": identity.model_version, + "pre_update_state": identity.pre_update_state, + "teacher_forcing": True, + "use_cache": False, + }, + operator_overrides={"selected_logprob": "rlkernel.reference_logp"}, + ) + + +def _bound_smoke_case(scenario: str) -> tuple[ExperimentCase, OperatorSelection]: + threshold_offset = resolve_logprob_threshold("float32") * 4.0 + if scenario == "reference-reference": + rollout_backend = SMOKE_ONLY_LOGP_REFERENCE_BACKEND_ID + training_backend = SMOKE_ONLY_LOGP_REFERENCE_BACKEND_ID + rollout_options = {} + training_options = {} + elif scenario == "reference-offset": + rollout_backend = SMOKE_ONLY_LOGP_REFERENCE_BACKEND_ID + training_backend = SMOKE_ONLY_LOGP_OFFSET_BACKEND_ID + rollout_options = {} + training_options = {"offset": threshold_offset} + elif scenario == "offset-offset": + rollout_backend = SMOKE_ONLY_LOGP_OFFSET_BACKEND_ID + training_backend = SMOKE_ONLY_LOGP_OFFSET_BACKEND_ID + rollout_options = {"offset": threshold_offset} + training_options = {"offset": threshold_offset} + else: # pragma: no cover - test helper contract + raise ValueError(f"unknown scenario: {scenario}") + selection = OperatorSelection( + rollout_backend=rollout_backend, + training_backend=training_backend, + rollout_options=rollout_options, + training_options=training_options, + ) + case = bind_operator_selection( + _case(case_id=f"case-{scenario}", backend=rollout_backend), + selection, + ) + return case, selection + + +def _write_required_artifacts( + store: ArtifactStore, + attempt_dir: Path, + *, + case_id: str = "case-1", + omit: frozenset[str] = frozenset(), + rollout_logprobs: torch.Tensor | None = None, + training_logprobs: torch.Tensor | None = None, + active_mask: torch.Tensor | None = None, + threshold: float = 0.05, +) -> None: + attempt_id = attempt_dir.name + json_values = { + "requested.json": { + "schema_version": "cross_config.requested.v1", + "case_id": case_id, + "attempt_id": attempt_id, + "case": {"case_id": case_id}, + }, + "materialized.json": { + "schema_version": "cross_config.materialized_envelope.v1", + "case_id": case_id, + "attempt_id": attempt_id, + "materialized_case": {"case": {"case_id": case_id}}, + }, + "actual.json": { + "schema_version": "cross_config.actual.v1", + "case_id": case_id, + "attempt_id": attempt_id, + "rollout": {}, + "training": {}, + }, + "identity.json": { + "schema_version": "cross_config.identity_envelope.v1", + "case_id": case_id, + "attempt_id": attempt_id, + "identity": {"checkpoint_id": "tiny"}, + }, + "comparison.json": { + "schema_version": "cross_config.alignment_result.v1", + "case_id": case_id, + "attempt_id": attempt_id, + "status": "pass", + "comparable": True, + "passed": True, + }, + } + for name, value in json_values.items(): + if name not in omit: + store.write_json(attempt_dir, name, value) + + rollout = rollout_logprobs if rollout_logprobs is not None else torch.tensor([-1.0, -2.0, -3.0]) + training = training_logprobs if training_logprobs is not None else rollout.clone() + active = active_mask if active_mask is not None else torch.tensor([True, True, True]) + mismatch = recompute_mismatch_mask(rollout, training, active, threshold) + tensor_values = { + "score_rollout.pt": { + "selected_logprobs": rollout, + "active_mask": active, + }, + "score_training.pt": { + "selected_logprobs": training, + "active_mask": active, + }, + "token_diffs.pt": { + "rollout_logprobs": rollout, + "training_logprobs": training, + "active_mask": active, + "absolute_diff": torch.abs(training - rollout), + "mismatch_mask": mismatch, + }, + } + for name, tensors in tensor_values.items(): + if name not in omit: + store.write_tensor_bundle( + attempt_dir, + name, + tensors, + metadata={ + "case_id": case_id, + "attempt_id": attempt_id, + "artifact": name, + "fixed_threshold": threshold, + }, + ) + + +def _complete_attempt( + store: ArtifactStore, + *, + experiment_id: str = "experiment-1", + case_id: str = "case-1", + **artifact_options, +) -> Path: + attempt_dir = store.create_attempt(experiment_id, case_id) + _write_required_artifacts(store, attempt_dir, case_id=case_id, **artifact_options) + store.complete_attempt( + attempt_dir, + summary={ + "schema_version": "cross_config.complete.v1", + "case_id": case_id, + "attempt_id": attempt_dir.name, + "status": "pass", + }, + ) + return attempt_dir + + +@pytest.mark.smoke_operator +def test_cpu_smoke_cases_preserve_read_only_scoring_and_exact_provenance(tmp_path: Path): + store = ArtifactStore(tmp_path) + batch = _batch() + inputs_before = batch.input_ids.clone() + expected = { + "reference-reference": (True, 0), + "reference-offset": (False, int(batch.active_mask.sum().item())), + "offset-offset": (True, 0), + } + + for scenario, (expected_pass, expected_mismatches) in expected.items(): + case, selection = _bound_smoke_case(scenario) + result = run_cpu_case( + store, + case, + batch, + selection, + allow_smoke_operators=True, + strict=True, + timeout_seconds=5.0, + resume=False, + ) + + assert result.resumed is False + assert result.alignment is not None + assert result.alignment.passed is expected_pass + assert result.alignment.mismatch_count == expected_mismatches + assert result.rollout_score is not None + assert result.training_score is not None + assert result.rollout_score.selected_logprobs.device.type == "cpu" + assert result.training_score.selected_logprobs.device.type == "cpu" + assert result.rollout_score.scorer.device == "cpu" + assert result.training_score.scorer.device == "cpu" + + guard = result.training_score.provenance.evidence["scoring_guard"] + assert guard == { + "model_state_verified": True, + "model_eval": True, + "no_grad": True, + "optimizer_step": False, + "model_modes_restored": True, + "model_state_unchanged": True, + } + assert result.rollout_score.provenance.evidence["scoring_guard"] == guard + assert result.training_score.provenance.evidence["rank_metadata"][0]["batch_ranges"] == ( + (0, 2), + (2, 3), + ) + rollout_state = result.rollout_score.provenance.evidence["model_state_fingerprint"] + training_state = result.training_score.provenance.evidence["model_state_fingerprint"] + assert rollout_state == training_state + + actual = json.loads((result.attempt_dir / "actual.json").read_text(encoding="utf-8")) + assert actual["operator_source"] == "exact_resolution_and_instance" + for target, backend in ( + ("rollout", selection.rollout_backend), + ("training", selection.training_backend), + ): + operator = actual[target]["actual"]["operators"]["selected_logprob"] + assert operator["backend_id"] == backend + assert operator["descriptor_fingerprint"] + assert operator["implementation_fingerprint"] + assert operator["instance_fingerprint"] + complete = result.attempt_dir / "COMPLETE" + assert complete.is_file() + assert result.summary == json.loads(complete.read_text(encoding="utf-8")) + + assert torch.equal(batch.input_ids, inputs_before) + + +@pytest.mark.smoke_operator +def test_runner_resumes_valid_attempt_and_retries_after_identity_or_tensor_change( + tmp_path: Path, + monkeypatch, +): + store = ArtifactStore(tmp_path) + case, selection = _bound_smoke_case("reference-reference") + batch = _batch() + + def run(): + return run_cpu_case( + store, + case, + batch, + selection, + allow_smoke_operators=True, + strict=True, + timeout_seconds=5.0, + resume=True, + ) + + first = run() + resumed = run() + assert first.attempt_id == "attempt-0001" + assert resumed.resumed is True + assert resumed.attempt_id == first.attempt_id + assert resumed.rollout_score is None + assert resumed.summary == first.summary + + token_path = first.attempt_dir / "token_diffs.pt" + payload = torch.load(token_path, map_location="cpu", weights_only=True) + payload["tensors"]["mismatch_mask"] = torch.ones_like(payload["tensors"]["mismatch_mask"]) + torch.save(payload, token_path) + + retried = run() + assert retried.resumed is False + assert retried.attempt_id == "attempt-0002" + assert (retried.attempt_dir / "COMPLETE").is_file() + + original_apply_fp32 = SmokeOnlyLogpReference.apply_fp32 + + def equivalent_apply_fp32(self, logits, token_ids, active_mask=None): + return original_apply_fp32(self, logits, token_ids, active_mask=active_mask) + + monkeypatch.setattr(SmokeOnlyLogpReference, "apply_fp32", equivalent_apply_fp32) + implementation_changed = run() + assert implementation_changed.resumed is False + assert implementation_changed.attempt_id == "attempt-0003" + before = json.loads((retried.attempt_dir / "actual.json").read_text(encoding="utf-8")) + after = json.loads( + (implementation_changed.attempt_dir / "actual.json").read_text(encoding="utf-8") + ) + assert ( + before["rollout"]["actual"]["operators"]["selected_logprob"]["implementation_fingerprint"] + != after["rollout"]["actual"]["operators"]["selected_logprob"]["implementation_fingerprint"] + ) + attempts = sorted(path.name for path in retried.attempt_dir.parent.iterdir()) + assert attempts == ["attempt-0001", "attempt-0002", "attempt-0003"] + + +@pytest.mark.parametrize( + ("mode", "error_type", "message"), + [ + ("failure", ChildScoringError, "intentional scorer failure"), + ("timeout", ScoringTimeoutError, "stopped children"), + ("missing-rank", RankCompletenessError, r"missing=\[0\]"), + ("duplicate-rank", RankCompletenessError, "duplicate ranks"), + ], +) +def test_runner_supervision_fails_closed_and_cleans_children( + tmp_path: Path, + mode: str, + error_type: type[Exception], + message: str, +): + case = _case(case_id=f"case-{mode}") + resolved, instances, provenance = _operators() + if mode == "failure": + rollout = FailingScorer(_spec(ScoreSide.ROLLOUT), ()) + training = FixedRankScorer(_spec(ScoreSide.TRAINING), (0,)) + timeout = 5.0 + elif mode == "timeout": + rollout = SlowScorer(_spec(ScoreSide.ROLLOUT), (0,)) + training = SlowScorer(_spec(ScoreSide.TRAINING), (0,)) + timeout = 0.1 + elif mode == "missing-rank": + rollout = FixedRankScorer(_spec(ScoreSide.ROLLOUT), ()) + training = FixedRankScorer(_spec(ScoreSide.TRAINING), (0,)) + timeout = 5.0 + else: + rollout = FixedRankScorer(_spec(ScoreSide.ROLLOUT), (0, 0)) + training = FixedRankScorer(_spec(ScoreSide.TRAINING), (0,)) + timeout = 5.0 + + runner = PairedRunner(ArtifactStore(tmp_path), timeout_seconds=timeout) + with pytest.raises(error_type, match=message): + runner.run( + case, + _materialization(case), + _batch(), + rollout, + training, + resolved, + instances, + provenance, + timeout_seconds=timeout, + ) + + assert runner.active_child_pids == () + attempt_dir = tmp_path / case.experiment_id / "cases" / case.case_id / "attempt-0001" + assert attempt_dir.is_dir() + assert not (attempt_dir / "COMPLETE").exists() + assert not list(attempt_dir.glob(".paired-runner-*")) + + +def test_artifacts_are_append_only_and_complete_marker_is_published_last(tmp_path: Path): + store = ArtifactStore(tmp_path) + attempt_dir = store.create_attempt("experiment-1", "case-1") + _write_required_artifacts( + store, + attempt_dir, + omit=frozenset({"token_diffs.pt"}), + ) + requested_path = attempt_dir / "requested.json" + rollout_path = attempt_dir / "score_rollout.pt" + requested_before = requested_path.read_bytes() + rollout_before = rollout_path.read_bytes() + + with pytest.raises(ArtifactError, match="refusing to overwrite"): + store.write_json(attempt_dir, "requested", {"case_id": "changed"}) + with pytest.raises(ArtifactError, match="refusing to overwrite"): + store.write_tensor_bundle( + attempt_dir, + "score_rollout", + {"selected_logprobs": torch.tensor([0.0])}, + ) + assert requested_path.read_bytes() == requested_before + assert rollout_path.read_bytes() == rollout_before + + summary = { + "schema_version": "cross_config.complete.v1", + "case_id": "case-1", + "attempt_id": attempt_dir.name, + "status": "pass", + } + with pytest.raises(ArtifactError, match=r"missing artifacts.*token_diffs\.pt"): + store.complete_attempt(attempt_dir, summary=summary) + assert not (attempt_dir / "COMPLETE").exists() + + _write_required_artifacts( + store, + attempt_dir, + omit=frozenset(_JSON_ARTIFACTS + _TENSOR_ARTIFACTS[:-1]), + ) + marker = store.complete_attempt(attempt_dir, summary=summary) + store.validate_completed_attempt(attempt_dir, expected_case_id="case-1") + payload_times = [ + (attempt_dir / name).stat().st_mtime_ns for name in _JSON_ARTIFACTS + _TENSOR_ARTIFACTS + ] + assert marker.stat().st_mtime_ns >= max(payload_times) + marker_value = json.loads(marker.read_text(encoding="utf-8")) + artifact_hashes = marker_value.pop("artifact_sha256") + assert marker_value == summary + assert set(artifact_hashes) == set(_JSON_ARTIFACTS + _TENSOR_ARTIFACTS) + assert all(len(value) == 64 for value in artifact_hashes.values()) + assert not list(attempt_dir.glob(".COMPLETE.*")) + with pytest.raises(ArtifactError, match="refusing to overwrite"): + store.complete_attempt(attempt_dir, summary=summary) + + next_attempt = store.create_attempt("experiment-1", "case-1") + assert next_attempt.name == "attempt-0002" + + +def test_resume_uses_newest_valid_attempt_and_tensors_support_offline_recompute(tmp_path: Path): + store = ArtifactStore(tmp_path) + rollout = torch.tensor([-1.0, -2.0, -3.0]) + training = torch.tensor([-1.01, -2.20, -2.50]) + active = torch.tensor([True, True, False]) + older = _complete_attempt( + store, + rollout_logprobs=rollout, + training_logprobs=training, + active_mask=active, + threshold=0.05, + ) + newer = _complete_attempt(store) + partial = store.create_attempt("experiment-1", "case-1") + store.write_json(partial, "requested", {"case_id": "case-1"}) + + assert store.completed_attempt("experiment-1", "case-1") == newer + (newer / "COMPLETE").write_text("{not-json", encoding="utf-8") + assert store.completed_attempt("experiment-1", "case-1") == older + + token_payload = store.load_tensor_bundle(older / "token_diffs.pt") + tensors = token_payload["tensors"] + recomputed = recompute_mismatch_mask( + tensors["rollout_logprobs"], + tensors["training_logprobs"], + tensors["active_mask"], + token_payload["metadata"]["fixed_threshold"], + ) + assert torch.equal(recomputed, tensors["mismatch_mask"]) + assert torch.equal(recomputed, torch.tensor([False, True, False])) + assert all(tensor.device.type == "cpu" for tensor in tensors.values()) + assert partial.name == "attempt-0003" + + materialized_path = older / "materialized.json" + materialized = json.loads(materialized_path.read_text(encoding="utf-8")) + materialized["materialized_case"]["case"]["case_id"] = "tampered" + materialized_path.write_text(json.dumps(materialized), encoding="utf-8") + assert store.completed_attempt("experiment-1", "case-1") is None diff --git a/tests/test_cross_config_runtime.py b/tests/test_cross_config_runtime.py new file mode 100644 index 00000000..a95ddd52 --- /dev/null +++ b/tests/test_cross_config_runtime.py @@ -0,0 +1,780 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from __future__ import annotations + +import inspect +from dataclasses import replace +from pathlib import Path + +import pytest +import torch + +from rl_engine.alignment.cross_config.operators import ( + OperatorBridge, + OperatorOverride, + selected_logprobs_with_operator, +) +from rl_engine.alignment.cross_config.planner import V1_KNOBS +from rl_engine.alignment.cross_config.runtime import ( + AdapterMaterialization, + KnobApplication, + RuntimeBinding, + RuntimeMaterializationError, + RuntimeTools, +) +from rl_engine.alignment.cross_config.schema import ( + ExperimentCase, + IsolationScope, + MaterializationStatus, + SemanticIdentitySpec, +) +from rl_engine.alignment.testing.cpu_cross_config import CpuSmokeMaterializer +from rl_engine.alignment.testing.smoke_ops import ( + SMOKE_ONLY_LOGP_OFFSET_BACKEND_ID, + SMOKE_ONLY_LOGP_REFERENCE_BACKEND_ID, + SmokeOnlyLogpOffset, + register_smoke_operators, +) +from rl_engine.kernels.ops.pytorch.loss.logp import NativeLogpOp +from rl_engine.kernels.semantic_registry import ( + OperatorRequirements, + OperatorResolutionError, + OperatorResolutionPolicy, + SemanticOperatorCatalog, +) +from rl_engine.kernels.semantic_registry import ( + implementation_fingerprint as fingerprint_implementation, +) +from rl_engine.testing import selected_logprobs_reference + +_REPOSITORY_ROOT = Path(__file__).resolve().parents[1] +_TEMPORARY_DOCSTRING = "TEMPORARY TEST SCAFFOLD - NOT A PRODUCTION RL-KERNEL OPERATOR" +_TOPOLOGIES = { + "rollout": { + "world_size": 1, + "tensor_parallel_size": 1, + "context_parallel_size": 1, + }, + "training": {"world_size": 1, "sharding": "unsharded"}, +} + + +def _identity() -> SemanticIdentitySpec: + return SemanticIdentitySpec( + checkpoint_id="tiny-cpu-checkpoint", + model_version="weights-v1", + tokenizer_policy="synthetic-tokenizer-v1", + token_ids=((1, 2, 3),), + selected_token_ids=((0, 2, 3),), + active_mask=((False, True, True),), + attention_mask=((True, True, True),), + pre_update_state="iteration-0", + ) + + +def _requested(**overrides): + requested = { + "batch": {"size": 2}, + "rollout": { + "tensor_parallel_size": 1, + "context_parallel_size": 1, + "dtype": "float32", + "enable_prefix_caching": False, + "enforce_eager": True, + }, + "training": { + "attention_backend": "eager", + "compute_dtype": "float32", + "sharding": "unsharded", + }, + "logp": {"backend": "rlkernel.reference_logp"}, + } + for path, value in overrides.items(): + current = requested + parts = path.split(".") + for part in parts[:-1]: + current = current[part] + current[parts[-1]] = value + return requested + + +def _case( + *, + case_id: str = "case-1", + changed_paths=(), + requested=None, + execution_binding=None, +) -> ExperimentCase: + return ExperimentCase( + case_id=case_id, + experiment_id="runtime-test", + scenario_id="S0", + identity=_identity(), + requested=requested or _requested(), + changed_paths=changed_paths, + execution_binding=execution_binding or {}, + contract_fingerprint="contract-sha", + scenario_fingerprint="scenario-sha", + ) + + +def _value_at(requested, path: str): + current = requested + for part in path.split("."): + current = current[part] + return current + + +def _readback(requested): + return {path: _value_at(requested, path) for path in V1_KNOBS if path != "batch.size"} + + +class _RuntimeTestAdapter: + """Small observable fake kept beside the lifecycle tests that need it.""" + + runtime_kind = "test_runtime" + + def __init__(self, *, actual_readback=None): + self.actual_readback = dict(actual_readback or {}) + + @property + def implementation_fingerprint(self): + return fingerprint_implementation( + type(self), + instance=self, + entrypoints=("materialize",), + ) + + def materialize(self, normalized, descriptors): + applications = [] + for path, descriptor in descriptors.items(): + requested = _value_at(normalized, path) + materialized = requested + actual = self.actual_readback.get(path) + status = MaterializationStatus.UNOBSERVABLE + reason = "no runtime readback is available" + + unsupported = (path == "rollout.context_parallel_size" and requested != 1) or ( + path == "training.sharding" and requested != "unsharded" + ) + if unsupported: + materialized = actual = None + status = MaterializationStatus.UNSUPPORTED + reason = "the test adapter does not support this topology" + elif path == "batch.size": + actual = requested + status = MaterializationStatus.APPLIED + reason = "batch size is observed at scorer invocation" + elif path in self.actual_readback: + status = ( + MaterializationStatus.APPLIED + if actual == requested + else MaterializationStatus.FALLBACK + ) + reason = "runtime readback was captured" + + applications.append( + KnobApplication( + path=path, + requested=requested, + materialized=materialized, + actual=actual, + lifecycle=descriptor.lifecycle, + status=status, + evidence={"reason": reason}, + critical=descriptor.critical, + ) + ) + + backend = _value_at(normalized, "logp.backend") + return AdapterMaterialization( + applications=tuple(applications), + binding=RuntimeBinding( + batch_size=_value_at(normalized, "batch.size"), + side_configs={"rollout": {}, "training": {}}, + topology={ + "rollout": { + "world_size": 1, + "tensor_parallel_size": _value_at( + normalized, "rollout.tensor_parallel_size" + ), + "context_parallel_size": _value_at( + normalized, "rollout.context_parallel_size" + ), + }, + "training": { + "world_size": 1, + "sharding": _value_at(normalized, "training.sharding"), + }, + }, + scorer={}, + operator_backends={"rollout": backend, "training": backend}, + runtime_kind=self.runtime_kind, + ), + ) + + +def _cpu_materializer() -> CpuSmokeMaterializer: + backends = { + "rollout": "rlkernel.reference_logp", + "training": "rlkernel.reference_logp", + } + return CpuSmokeMaterializer( + requested_operator_backends=backends, + actual_operator_backends=backends, + ) + + +def _requirements( + *, + target: str = "rollout", + device: str = "cpu", +) -> OperatorRequirements: + return OperatorRequirements( + device=device, + dtype="float32", + topology=_TOPOLOGIES[target], + alignment_properties={"deterministic": True}, + ) + + +def _catalog() -> SemanticOperatorCatalog: + """Clone repository descriptors so each test owns registration state.""" + + return SemanticOperatorCatalog(OperatorBridge().catalog.backend_descriptors()) + + +def test_cpu_materialization_records_all_ten_knobs_across_three_stages(): + case = _case() + materialization = RuntimeTools().materialize(case, _cpu_materializer()) + applications = {application.path: application for application in materialization.applications} + + assert len(V1_KNOBS) == 10 + assert set(applications) == set(V1_KNOBS) + assert materialization.materialized_case.status is MaterializationStatus.APPLIED + assert materialization.executable_in_strict_mode + RuntimeTools.require_executable(materialization, strict=True) + + for path, descriptor in V1_KNOBS.items(): + application = applications[path] + assert application.requested == _value_at(case.requested, path) + assert application.lifecycle is descriptor.lifecycle + assert application.status is MaterializationStatus.APPLIED + assert application.evidence["reason"] + + provenance = materialization.provenance + assert provenance.requested == case.requested + assert provenance.normalized == case.requested + assert provenance.materialized["batch"]["size"] == 2 + assert provenance.materialized["rollout"] == { + "tensor_parallel_size": 1, + "context_parallel_size": 1, + "dtype": "float32", + "enable_prefix_caching": False, + "enforce_eager": True, + } + assert provenance.materialized["training"] == { + "attention_backend": "eager", + "compute_dtype": "float32", + "sharding": "unsharded", + } + assert provenance.materialized["logp"]["backend"] == { + "rollout": "rlkernel.reference_logp", + "training": "rlkernel.reference_logp", + } + assert provenance.actual == provenance.materialized + assert provenance.implementation_fingerprint == _cpu_materializer().implementation_fingerprint + assert provenance.evidence["adapter_implementation_fingerprint"] == ( + provenance.implementation_fingerprint + ) + + binding = materialization.binding + assert binding.runtime_kind == "cpu_smoke" + assert binding.side_configs["rollout"]["device"] == "cpu" + assert binding.side_configs["training"]["device"] == "cpu" + assert binding.side_configs["training"]["dtype"] == "float32" + assert binding.side_configs["rollout"] == { + "device": "cpu", + "dtype": "float32", + "enable_prefix_caching": False, + "enforce_eager": True, + } + assert binding.topology["rollout"] == { + "world_size": 1, + "tensor_parallel_size": 1, + "context_parallel_size": 1, + } + assert binding.topology["training"] == {"world_size": 1, "sharding": "unsharded"} + assert binding.scorer == { + "mode": "reference", + "use_cache": False, + "attention_backend": "eager", + "output_dtype": "float32", + } + with pytest.raises(TypeError): + binding.side_configs["rollout"]["device"] = "cuda" + with pytest.raises(TypeError): + applications["batch.size"].evidence["reason"] = "changed after fingerprinting" + + +def test_lifecycle_fingerprints_allow_request_reuse_and_isolate_engine_and_process_changes( + monkeypatch, +): + tools = RuntimeTools() + baseline_case = _case(case_id="baseline") + baseline_adapter = _RuntimeTestAdapter(actual_readback=_readback(baseline_case.requested)) + baseline = tools.materialize( + baseline_case, + baseline_adapter, + ) + + batch_requested = _requested(**{"batch.size": 1}) + batch = tools.materialize( + _case( + case_id="batch", + requested=batch_requested, + changed_paths=("batch.size",), + ), + _RuntimeTestAdapter(actual_readback=_readback(batch_requested)), + ) + assert batch.materialized_case.isolation_scope is IsolationScope.REQUEST + assert batch.materialized_case.construction_fingerprint == ( + baseline.materialized_case.construction_fingerprint + ) + assert batch.materialized_case.distributed_context_fingerprint == ( + baseline.materialized_case.distributed_context_fingerprint + ) + assert batch.materialized_case.process_fingerprint == ( + baseline.materialized_case.process_fingerprint + ) + assert tools.can_reuse(baseline, batch) + + dtype_requested = _requested(**{"rollout.dtype": "bfloat16"}) + dtype = tools.materialize( + _case( + case_id="dtype", + requested=dtype_requested, + changed_paths=("rollout.dtype",), + ), + _RuntimeTestAdapter(actual_readback=_readback(dtype_requested)), + ) + assert dtype.materialized_case.isolation_scope is IsolationScope.ENGINE_CONSTRUCTION + assert dtype.materialized_case.construction_fingerprint != ( + baseline.materialized_case.construction_fingerprint + ) + assert dtype.materialized_case.process_fingerprint == ( + baseline.materialized_case.process_fingerprint + ) + assert not tools.can_reuse(baseline, dtype) + + topology_requested = _requested(**{"rollout.tensor_parallel_size": 2}) + topology = tools.materialize( + _case( + case_id="topology", + requested=topology_requested, + changed_paths=("rollout.tensor_parallel_size",), + ), + _RuntimeTestAdapter(actual_readback=_readback(topology_requested)), + ) + assert topology.materialized_case.isolation_scope is IsolationScope.PROCESS + assert topology.materialized_case.process_fingerprint != ( + baseline.materialized_case.process_fingerprint + ) + assert not tools.can_reuse(baseline, topology) + + rebound_case = replace( + baseline_case, + case_id="rebound", + execution_binding={"operator_case": {"rollout_options": {"offset": 0.1}}}, + ) + rebound = tools.materialize( + rebound_case, + _RuntimeTestAdapter(actual_readback=_readback(rebound_case.requested)), + ) + assert rebound.materialized_case.construction_fingerprint == ( + baseline.materialized_case.construction_fingerprint + ) + assert not tools.can_reuse(baseline, rebound) + + changed_identity_case = replace( + baseline_case, + case_id="changed-identity", + identity=replace(baseline_case.identity, pre_update_state="iteration-1"), + ) + changed_identity = tools.materialize(changed_identity_case, baseline_adapter) + assert changed_identity.materialized_case.construction_fingerprint == ( + baseline.materialized_case.construction_fingerprint + ) + assert not tools.can_reuse(baseline, changed_identity) + + original_materialize = _RuntimeTestAdapter.materialize + + def materialize_with_same_result(self, normalized, descriptors): + return original_materialize(self, normalized, descriptors) + + monkeypatch.setattr( + _RuntimeTestAdapter, + "materialize", + materialize_with_same_result, + ) + changed_adapter = tools.materialize( + baseline_case, + _RuntimeTestAdapter(actual_readback=_readback(baseline_case.requested)), + ) + assert changed_adapter.provenance.actual == baseline.provenance.actual + assert ( + changed_adapter.provenance.implementation_fingerprint + != baseline.provenance.implementation_fingerprint + ) + assert not tools.can_reuse(baseline, changed_adapter) + + +def test_materialization_fails_closed_for_fallback_unobservable_and_unsupported_paths(): + tools = RuntimeTools() + + incomplete_descriptors = { + path: descriptor for path, descriptor in V1_KNOBS.items() if path != "batch.size" + } + with pytest.raises(RuntimeMaterializationError, match="missing descriptors"): + RuntimeTools(incomplete_descriptors).materialize( + _case(case_id="missing-descriptor"), + _RuntimeTestAdapter(actual_readback=_readback(_requested())), + ) + + fallback_requested = _requested(**{"training.attention_backend": "flash_attention_2"}) + fallback_readback = _readback(fallback_requested) + fallback_readback["training.attention_backend"] = "eager" + fallback = tools.materialize( + _case( + case_id="fallback", + requested=fallback_requested, + changed_paths=("training.attention_backend",), + ), + _RuntimeTestAdapter(actual_readback=fallback_readback), + ) + assert fallback.materialized_case.status is MaterializationStatus.FALLBACK + with pytest.raises(RuntimeMaterializationError, match=r"training\.attention_backend"): + tools.require_executable(fallback, strict=True) + tools.require_executable(fallback, strict=False) + + unobservable = tools.materialize( + _case(case_id="unobservable"), + _RuntimeTestAdapter(), + ) + assert unobservable.materialized_case.status is MaterializationStatus.UNOBSERVABLE + with pytest.raises(RuntimeMaterializationError, match="no runtime readback"): + tools.require_executable(unobservable, strict=False) + + unsupported_requested = _requested( + **{ + "rollout.context_parallel_size": 4, + "training.sharding": "fsdp", + } + ) + unsupported = tools.materialize( + _case( + case_id="unsupported", + requested=unsupported_requested, + changed_paths=("rollout.context_parallel_size", "training.sharding"), + ), + _RuntimeTestAdapter(), + ) + unsupported_paths = { + application.path + for application in unsupported.applications + if application.status is MaterializationStatus.UNSUPPORTED + } + assert unsupported_paths == {"rollout.context_parallel_size", "training.sharding"} + assert unsupported.materialized_case.status is MaterializationStatus.UNSUPPORTED + with pytest.raises(RuntimeMaterializationError, match=r"rollout\.context_parallel_size"): + tools.require_executable(unsupported, strict=False) + + cpu_only_requested = _requested(**{"rollout.tensor_parallel_size": 2}) + cpu_only = tools.materialize( + _case( + case_id="cpu-only", + requested=cpu_only_requested, + changed_paths=("rollout.tensor_parallel_size",), + ), + _cpu_materializer(), + ) + assert cpu_only.materialized_case.status is MaterializationStatus.UNSUPPORTED + assert cpu_only.binding.side_configs["rollout"]["device"] == "cpu" + assert cpu_only.binding.side_configs["training"]["device"] == "cpu" + + +def test_operator_binding_selects_rollout_training_and_both_without_side_leakage(): + bridge = OperatorBridge() + requirements = { + "rollout": _requirements(target="rollout"), + "training": _requirements(target="training"), + } + assert requirements["rollout"].to_dict()["schema_version"] == ( + "rlkernel.semantic_operator.requirements.v1" + ) + + for target, expected_targets in ( + ("rollout", {"rollout"}), + ("training", {"training"}), + ("both", {"rollout", "training"}), + ): + resolved = bridge.resolve_override( + OperatorOverride.for_target( + semantic_op="selected_logprob", + backend_id="rlkernel.reference_logp", + target=target, + ), + requirements=requirements, + strict=True, + ) + selected_targets = { + side for side in ("rollout", "training") if resolved.for_target(side) is not None + } + assert selected_targets == expected_targets + + instances = {} + for side in expected_targets: + resolution = resolved.for_target(side) + assert resolution is not None + assert resolution.to_dict()["schema_version"] == ( + "rlkernel.semantic_operator.resolution.v1" + ) + assert resolution.descriptor.to_dict()["schema_version"] == ( + "rlkernel.semantic_operator.backend_descriptor.v1" + ) + assert resolution.trace.to_dict()["schema_version"] == ( + "rlkernel.semantic_operator.resolution_trace.v1" + ) + instance = bridge.instantiate(resolved, target=side) + instances[side] = instance + assert isinstance(instance, NativeLogpOp) + provenance = bridge.instance_provenance( + resolved, + target=side, + instance=instance, + ) + assert provenance.backend_id == "rlkernel.reference_logp" + assert provenance.target == side + assert provenance.instance_fingerprint + assert provenance.to_dict()["schema_version"] == ( + "rlkernel.semantic_operator.instance_provenance.v1" + ) + with pytest.raises(TypeError): + provenance.factory_options["unexpected"] = True + if target == "both": + assert instances["rollout"] is not instances["training"] + + catalog = _catalog() + descriptor = catalog.backend_descriptor("selected_logprob", "rlkernel.reference_logp") + assert descriptor is not None + catalog.register_backend( + replace(descriptor, supported_topologies={"*": "*"}), + replace=True, + ) + asymmetric = OperatorBridge(catalog).resolve_override( + OperatorOverride.for_target( + semantic_op="selected_logprob", + backend_id="rlkernel.reference_logp", + target="both", + ), + requirements={ + "rollout": OperatorRequirements( + device="cpu", + dtype="float32", + topology={"world_size": 2, "tensor_parallel_size": 2}, + ), + "training": OperatorRequirements( + device="cpu", + dtype="float32", + topology={"world_size": 1, "sharding": "fsdp"}, + ), + }, + ) + assert asymmetric.rollout is not None and asymmetric.training is not None + assert asymmetric.rollout.requirements.topology != asymmetric.training.requirements.topology + + strict_session = _catalog().session() + with pytest.raises(OperatorResolutionError, match="topology"): + strict_session.resolve( + semantic_op="selected_logprob", + requested_backend="rlkernel.reference_logp", + target="rollout", + requirements=OperatorRequirements(device="cpu", dtype="float32", topology={}), + strict=True, + ) + session = catalog.session() + with pytest.raises(OperatorResolutionError, match="not registered") as unsupported: + session.resolve( + semantic_op="selected_logprob", + requested_backend="missing.backend", + target="rollout", + requirements=_requirements(), + strict=True, + ) + assert unsupported.value.trace.status == "unsupported" + assert unsupported.value.trace.fallback_attempts == () + + native_requirements = OperatorRequirements( + device="cpu", + dtype="float32", + topology=_TOPOLOGIES["training"], + ) + with pytest.raises(OperatorResolutionError, match="not exactly observable"): + session.resolve( + semantic_op="selected_logprob", + requested_backend="native", + target="training", + requirements=native_requirements, + strict=True, + ) + native = session.resolve( + semantic_op="selected_logprob", + requested_backend="native", + target="training", + requirements=native_requirements, + strict=False, + ) + assert native.trace.status == "unobservable" + assert native.trace.concrete_backend is None + + +@pytest.mark.smoke_operator +def test_smoke_package_is_temporary_cpu_only_and_disabled_without_two_explicit_opt_ins(): + from rl_engine.alignment.testing.smoke_ops import ( + smoke_only_logp_offset, + smoke_only_logp_reference, + ) + + assert inspect.getdoc(smoke_only_logp_reference) == _TEMPORARY_DOCSTRING + assert inspect.getdoc(smoke_only_logp_offset) == _TEMPORARY_DOCSTRING + + manifest = ( + _REPOSITORY_ROOT / "rl_engine/alignment/testing/smoke_ops/SMOKE_OPERATORS.md" + ).read_text(encoding="utf-8") + for required_text in ( + "temporary test scaffolding", + "smoke_only_logp_reference.py", + "smoke_only_logp_offset.py", + "allow_smoke_operators=True", + "delete this package", + ): + assert required_text in manifest + + catalog = _catalog() + for backend_id in ( + SMOKE_ONLY_LOGP_REFERENCE_BACKEND_ID, + SMOKE_ONLY_LOGP_OFFSET_BACKEND_ID, + ): + assert catalog.backend_descriptor("selected_logprob", backend_id) is None + + with pytest.raises(PermissionError, match="allow_smoke_operators=True"): + register_smoke_operators(catalog) + + descriptors = register_smoke_operators(catalog, allow_smoke_operators=True) + assert {descriptor.backend_id for descriptor in descriptors} == { + SMOKE_ONLY_LOGP_REFERENCE_BACKEND_ID, + SMOKE_ONLY_LOGP_OFFSET_BACKEND_ID, + } + for descriptor in descriptors: + assert descriptor.supported_devices == frozenset({"cpu"}) + assert descriptor.is_smoke_only is True + disabled_session = catalog.session( + OperatorResolutionPolicy(strict=True, allow_test_backends=False) + ) + with pytest.raises(OperatorResolutionError, match="explicit opt-in"): + disabled_session.resolve( + semantic_op="selected_logprob", + requested_backend=descriptor.backend_id, + target="training", + requirements=_requirements(target="training"), + ) + enabled_session = catalog.session( + OperatorResolutionPolicy(strict=True, allow_test_backends=True) + ) + with pytest.raises(OperatorResolutionError) as error: + enabled_session.resolve( + semantic_op="selected_logprob", + requested_backend=descriptor.backend_id, + target="training", + requirements=_requirements(target="training", device="cuda"), + ) + failed = { + decision.capability + for decision in error.value.trace.capability_decisions + if not decision.passed + } + assert failed == {"device"} + + assert SmokeOnlyLogpOffset().offset == 0.0 + with pytest.raises(PermissionError, match="allow_smoke_operators=True"): + SmokeOnlyLogpOffset(offset=0.01) + + +@pytest.mark.smoke_operator +def test_explicit_smoke_opt_in_runs_both_sides_on_cpu_with_sealed_provenance(): + catalog = _catalog() + register_smoke_operators(catalog, allow_smoke_operators=True) + bridge = OperatorBridge( + catalog, + policy=OperatorResolutionPolicy(strict=True, allow_test_backends=True), + ) + resolved = bridge.resolve_override( + OperatorOverride.for_target( + semantic_op="selected_logprob", + backend_id=SMOKE_ONLY_LOGP_REFERENCE_BACKEND_ID, + target="both", + ), + requirements={ + "rollout": _requirements(target="rollout"), + "training": _requirements(target="training"), + }, + strict=True, + ) + instances = { + target: bridge.instantiate(resolved, target=target) for target in ("rollout", "training") + } + logits = torch.tensor( + [[[1.0, 2.0, -1.0], [0.0, 3.0, 1.0], [2.0, 0.0, 4.0]]], + device="cpu", + ) + token_ids = torch.tensor([[1, -100, 2]], device="cpu") + active_mask = torch.tensor([[True, False, True]], device="cpu") + expected = selected_logprobs_reference(logits, token_ids, mask=active_mask) + + outputs = {} + for target, instance in instances.items(): + output = selected_logprobs_with_operator( + instance, + logits, + token_ids, + active_mask=active_mask, + ) + outputs[target] = output + assert output.device.type == "cpu" + torch.testing.assert_close(output, expected, atol=0.0, rtol=0.0) + assert torch.count_nonzero(output[~active_mask]).item() == 0 + + provenance = bridge.instance_provenance( + resolved, + target=target, + instance=instance, + ) + assert provenance.backend_id == SMOKE_ONLY_LOGP_REFERENCE_BACKEND_ID + assert provenance.target == target + assert provenance.concrete_implementation.endswith("SmokeOnlyLogpReference") + assert provenance.descriptor_fingerprint + assert provenance.instance_fingerprint + + for invalid_temperature in (float("nan"), float("inf"), float("-inf")): + with pytest.raises(ValueError, match="finite and greater than zero"): + selected_logprobs_with_operator( + instances["rollout"], + logits, + token_ids, + active_mask=active_mask, + temperature=invalid_temperature, + ) + + assert instances["rollout"] is not instances["training"] + torch.testing.assert_close(outputs["rollout"], outputs["training"], atol=0.0, rtol=0.0) diff --git a/tests/test_det_gemm.py b/tests/test_det_gemm.py index 42b73e1d..715f7170 100644 --- a/tests/test_det_gemm.py +++ b/tests/test_det_gemm.py @@ -38,6 +38,24 @@ def _rand(*shape): return torch.randn(*shape, device=DEV, dtype=torch.bfloat16) +# Matches det_gemm_kernel.cu: mid-split K-tree, FP32 leaf width 32, BF16 internal adds. +_K_TREE_LEAF = 32 + + +def _k_tree_gemm(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: + a = a.detach().contiguous() + b = b.detach().contiguous() + k = a.shape[1] + + def rec(lo: int, hi: int) -> torch.Tensor: + if hi - lo <= _K_TREE_LEAF: + return (a[:, lo:hi].float() @ b[lo:hi, :].float()).to(dtype=torch.bfloat16) + mid = lo + (hi - lo) // 2 + return rec(lo, mid) + rec(mid, hi) + + return rec(0, k) + + @pytest.mark.parametrize("name,gemm", _BACKENDS) def test_forward_batch_invariance(name, gemm): # A row's output must not change when other rows join the batch. @@ -81,7 +99,10 @@ def test_forward_correctness(name, gemm): M, K, N = 128, 2048, 2048 a, b = _rand(M, K), _rand(K, N) out = gemm(a, b).float() - ref = a.float() @ b.float() + if name == "cuda": + ref = _k_tree_gemm(a, b).float() + else: + ref = a.float() @ b.float() contract = load_contract() thresholds = contract["accuracy"]["default"]["reduction"]["bfloat16"] torch.testing.assert_close(out, ref, atol=thresholds["atol"], rtol=thresholds["rtol"]) @@ -111,6 +132,18 @@ def test_backward_correctness(name, gemm): b = _rand(K, N).requires_grad_(True) g = _rand(M, N) gemm(a, b).backward(g) + if name == "cuda": + da = _k_tree_gemm(g, b.t().contiguous()) + db = _k_tree_gemm(a.detach().t().contiguous(), g) + contract = load_contract() + thresholds = contract["accuracy"]["default"]["reduction"]["bfloat16"] + torch.testing.assert_close( + a.grad.float(), da.float(), atol=thresholds["atol"], rtol=thresholds["rtol"] + ) + torch.testing.assert_close( + b.grad.float(), db.float(), atol=thresholds["atol"], rtol=thresholds["rtol"] + ) + return af = a.detach().float().requires_grad_(True) bf = b.detach().float().requires_grad_(True) (af @ bf).backward(g.float()) diff --git a/tests/test_framework_operator_integrations.py b/tests/test_framework_operator_integrations.py new file mode 100644 index 00000000..55f63aa8 --- /dev/null +++ b/tests/test_framework_operator_integrations.py @@ -0,0 +1,69 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from __future__ import annotations + +import pytest + +from rl_engine.integrations import IntegrationPlan, MegatronIntegration, VllmIntegration + + +class FakeOperator: + def __init__(self, value: str, backend_id: str): + self.value = value + self.backend_id = backend_id + + def __call__(self, payload: str) -> str: + return f"{self.value}:{payload}" + + +def _operators(): + return { + "attention": FakeOperator("rlk-attn", "rlkernel.attention.deterministic.v1"), + "ffn": FakeOperator("rlk-ffn", "rlkernel.ffn.qwen3.deterministic.v1"), + "logp": FakeOperator("rlk-logp", "pytorch-vocab-parallel-logp-ws2"), + } + + +def test_plan_uses_module_matrix_cases_without_cartesian_expansion(): + plan = IntegrationPlan.from_case_ids(attention="P/R", ffn="R/P", logp="R/R") + + assert plan.implementation_for("attention", "training").value == "production" + assert plan.implementation_for("attention", "rollout").value == "rl_kernel" + assert plan.implementation_for("ffn", "training").value == "rl_kernel" + assert plan.implementation_for("ffn", "rollout").value == "production" + assert plan.to_dict()["schema_version"] == "rlkernel.debug_matrix.v1" + + +def test_megatron_and_vllm_route_the_same_plan_on_opposite_sides(): + plan = IntegrationPlan.from_case_ids(attention="P/R", ffn="R/P", logp="R/R") + megatron = MegatronIntegration(plan, rl_kernel_operators=_operators()) + vllm = VllmIntegration(plan, rl_kernel_operators=_operators()) + native = FakeOperator("native", "production.backend") + + assert megatron.attention(native, "x") == "native:x" + assert vllm.attention(native, "x") == "rlk-attn:x" + assert megatron.ffn(native, "x") == "rlk-ffn:x" + assert vllm.ffn(native, "x") == "native:x" + assert megatron.logp(native, "x") == "rlk-logp:x" + assert vllm.logp(native, "x") == "rlk-logp:x" + + assert megatron.readback()["operators"]["attention"]["backend_id"] == "production.backend" + assert vllm.readback()["operators"]["attention"]["backend_id"] == ( + "rlkernel.attention.deterministic.v1" + ) + + +def test_rl_kernel_selection_fails_closed_when_operator_is_missing(): + plan = IntegrationPlan.from_case_ids(attention="R/R") + integration = MegatronIntegration(plan, rl_kernel_operators={}) + + with pytest.raises(RuntimeError, match="no operator was installed"): + integration.attention(lambda value: value, "x") + + +def test_integration_modules_do_not_import_optional_frameworks(): + import sys + + assert "megatron" not in sys.modules + assert "vllm" not in sys.modules diff --git a/tests/test_logprob_contract.py b/tests/test_logprob_contract.py new file mode 100644 index 00000000..a961f312 --- /dev/null +++ b/tests/test_logprob_contract.py @@ -0,0 +1,616 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""WS2 TP-aware logprob contract and contract-aware dispatch tests (issue #241).""" + +from __future__ import annotations + +import json +from dataclasses import replace + +import pytest + +from rl_engine.kernels.logprob_contract import ( + DeterminismScope, + LogprobBackendCapability, + LogprobContract, + LogprobContractError, + LogprobDType, + LogprobOutputSpec, + LogprobRole, + MaskMode, + MaskSpec, + ReductionSpec, + ShardingSpec, + TPPlacement, +) +from rl_engine.kernels.registry import KernelRegistry, OpBackend + +QWEN3_REAL_VOCAB = 151936 +QWEN3_PADDED_VOCAB = 152064 + + +def _even_bounds(padded_vocab: int, tp_world_size: int) -> tuple[tuple[int, int], ...]: + shard = padded_vocab // tp_world_size + return tuple( + (rank * shard, padded_vocab if rank == tp_world_size - 1 else (rank + 1) * shard) + for rank in range(tp_world_size) + ) + + +def _sharding( + *, + tp_rank: int = 0, + tp_world_size: int = 2, + cp_rank: int = 0, + cp_world_size: int = 2, + real_vocab_size: int = QWEN3_REAL_VOCAB, + padded_vocab_size: int = QWEN3_PADDED_VOCAB, + vocab_shard_bounds: tuple[tuple[int, int], ...] | None = None, +) -> ShardingSpec: + return ShardingSpec( + tp_rank=tp_rank, + tp_world_size=tp_world_size, + vocab_shard_bounds=( + vocab_shard_bounds + if vocab_shard_bounds is not None + else _even_bounds(padded_vocab_size, tp_world_size) + ), + real_vocab_size=real_vocab_size, + padded_vocab_size=padded_vocab_size, + cp_rank=cp_rank, + cp_world_size=cp_world_size, + ) + + +def _mask( + *, + num_tokens: int = 8, + active_mask: tuple[bool, ...] | None = None, + ignore_index: int = -100, +) -> MaskSpec: + return MaskSpec( + num_tokens=num_tokens, + active_mask=( + active_mask + if active_mask is not None + else (False, False, True, True, True, True, True, False) + ), + ignore_index=ignore_index, + ) + + +def _contract( + *, + role: str = "train", + dtype: str = "bf16", + mask: MaskSpec | None = None, + sharding: ShardingSpec | None = None, + reduction: ReductionSpec | None = None, +) -> LogprobContract: + return LogprobContract( + role=role, + dtype=dtype, + mask=mask if mask is not None else _mask(), + sharding=sharding if sharding is not None else _sharding(), + reduction=reduction if reduction is not None else ReductionSpec(), + ) + + +def _declared_tp_backend() -> LogprobBackendCapability: + return LogprobBackendCapability( + backend_id="test-deterministic-tp-logprob", + roles=frozenset({LogprobRole.TRAIN, LogprobRole.INFER}), + dtypes=frozenset({LogprobDType.BF16}), + tp_world_sizes=(1, 2, 4), + cp_world_sizes=None, + supports_vocab_padding=True, + mask_modes=frozenset({MaskMode.EXPLICIT_ACTIVE_MASK, MaskMode.IGNORE_INDEX}), + exports_vocab_lse=True, + determinism_scopes=frozenset( + {DeterminismScope.CROSS_TP_BITWISE, DeterminismScope.FIXED_TOPOLOGY} + ), + implementation_kind="reference", + ) + + +def test_qwen3_tp2_bf16_contract_is_representable_and_serializable(): + contract = _contract() + + assert contract.sharding.tp_world_size == 2 + assert contract.sharding.cp_world_size == 2 + assert contract.sharding.local_vocab_start == 0 + assert contract.sharding.local_vocab_end == QWEN3_PADDED_VOCAB // 2 + assert contract.sharding.local_vocab_size == QWEN3_PADDED_VOCAB // 2 + assert contract.mask.active_token_count == 5 + assert contract.reduction.acc_dtype is LogprobDType.FP32 + assert contract.to_dict()["reduction"] == { + "merge": "max_sumexp", + "merge_axis": "tp_vocab", + "acc_dtype": "fp32", + "order": "global_vocab_shard_index", + "transport": "all_gather", + "downcast_at": "final_write", + "engine": "in_op_reference", + "determinism_scope": "cross_tp_bitwise", + "cp_is_merge_axis": False, + } + json.dumps(contract.to_dict()) + + +@pytest.mark.parametrize("tp_world_size", [1, 2, 4]) +def test_pr4_sweep_tp_degrees_are_representable(tp_world_size): + sharding = _sharding(tp_world_size=tp_world_size, cp_world_size=1) + + assert len(sharding.vocab_shard_bounds) == tp_world_size + assert sharding.vocab_shard_bounds[-1][1] == QWEN3_PADDED_VOCAB + assert sharding.owner_rank(QWEN3_REAL_VOCAB - 1) == tp_world_size - 1 + + +@pytest.mark.parametrize( + ("field", "value", "message"), + [ + ("tp_rank", 2, "tp_rank=2"), + ("cp_rank", 2, "cp_rank=2"), + ("real_vocab_size", 0, "positive integer"), + ("padded_vocab_size", QWEN3_REAL_VOCAB - 1, "must not be smaller"), + ], +) +def test_invalid_rank_and_vocab_metadata_fail_loudly(field, value, message): + values = { + "tp_rank": 0, + "tp_world_size": 2, + "cp_rank": 0, + "cp_world_size": 2, + "real_vocab_size": QWEN3_REAL_VOCAB, + "padded_vocab_size": QWEN3_PADDED_VOCAB, + "vocab_shard_bounds": _even_bounds(QWEN3_PADDED_VOCAB, 2), + } + values[field] = value + + with pytest.raises(LogprobContractError, match=message): + ShardingSpec(**values) + + +@pytest.mark.parametrize( + ("bounds", "message"), + [ + ((), "one \\(start, end\\) pair per TP rank"), + (((0, 76032),), "one \\(start, end\\) pair per TP rank"), + (((0, 76032), (76032, 76032)), "end > start"), + (((0, 76000), (76032, 152064)), "contiguous"), + (((0, 76064), (76032, 152064)), "contiguous"), + (((0, 76032), (76032, 152000)), "cover padded_vocab_size exactly"), + ], +) +def test_incomplete_or_overlapping_vocab_shard_bounds_fail_loudly(bounds, message): + with pytest.raises(LogprobContractError, match=message): + _sharding(vocab_shard_bounds=bounds) + + +def test_owner_rank_is_unique_and_rejects_out_of_real_vocab_targets(): + sharding = _sharding() + + assert sharding.owner_rank(0) == 0 + assert sharding.owner_rank(QWEN3_PADDED_VOCAB // 2 - 1) == 0 + assert sharding.owner_rank(QWEN3_PADDED_VOCAB // 2) == 1 + assert sharding.owner_rank(QWEN3_REAL_VOCAB - 1) == 1 + + with pytest.raises(LogprobContractError, match="outside the real vocabulary"): + sharding.owner_rank(-1) + with pytest.raises(LogprobContractError, match="outside the real vocabulary"): + sharding.owner_rank(QWEN3_REAL_VOCAB) + + +def test_active_token_mask_metadata_is_validated(): + with pytest.raises(LogprobContractError, match="one entry per token"): + _mask(num_tokens=4) + + with pytest.raises(LogprobContractError, match="must be a bool"): + MaskSpec(num_tokens=2, active_mask=(True, 1)) + + all_inactive = _mask(num_tokens=3, active_mask=(False, False, False)) + assert all_inactive.active_token_count == 0 + + +def test_reduction_requires_fp32_accumulation_and_known_semantics(): + with pytest.raises(LogprobContractError, match="must be fp32"): + ReductionSpec(acc_dtype="bf16") + + with pytest.raises(LogprobContractError, match="merge must be one of"): + ReductionSpec(merge="lse_average") + + with pytest.raises(LogprobContractError, match="transport must be one of"): + ReductionSpec(transport="all_reduce") + + +def test_contract_component_types_and_lse_export_are_enforced(): + with pytest.raises(LogprobContractError, match="mask must be a MaskSpec"): + LogprobContract( + role="train", + dtype="bf16", + mask=None, + sharding=_sharding(), + reduction=ReductionSpec(), + ) + + with pytest.raises(LogprobContractError, match="export_lse must be True"): + replace(_contract(), export_lse=False) + + +def test_ignore_index_must_not_collide_with_the_real_vocabulary(): + with pytest.raises(LogprobContractError, match="must not collide"): + _contract(mask=_mask(ignore_index=5)) + + padding_column = QWEN3_REAL_VOCAB + 1 + contract = _contract(mask=_mask(ignore_index=padding_column)) + assert contract.mask.ignore_index == padding_column + + +def _restrict_to_ws1_candidates(registry: KernelRegistry) -> None: + """Drop the #241 PR3 vocab-parallel reference so only WS1 backends remain.""" + + platform = registry._platform() + registry._logprob_candidates[platform] = [ + backend + for backend in registry._logprob_candidates[platform] + if backend is not OpBackend.PYTORCH_VOCAB_PARALLEL_LOGP + ] + + +def test_current_ws1_backend_rejects_strict_tp_contract_without_fallback(): + registry = KernelRegistry() + _restrict_to_ws1_candidates(registry) + + with pytest.raises(RuntimeError) as exc_info: + registry.get_logprob_op(_contract()) + + message = str(exc_info.value) + assert "TP=2 is unsupported" in message + assert "vocab-domain LSE export is unsupported" in message + assert "determinism_scope=cross_tp_bitwise is unsupported" in message + assert "padded-vs-real vocab masking is unsupported" in message + + +def test_current_ws1_backend_rejects_padded_vocab_even_at_tp1(): + registry = KernelRegistry() + _restrict_to_ws1_candidates(registry) + contract = _contract(sharding=_sharding(tp_world_size=1, cp_world_size=1)) + + with pytest.raises(RuntimeError) as exc_info: + registry.get_logprob_op(contract) + + message = str(exc_info.value) + assert "TP=1 is unsupported" not in message + assert "padded-vs-real vocab masking is unsupported" in message + + +def test_ws1_rejections_recorded_when_vocab_parallel_reference_resolves(): + """The WS1 backends still reject strict contracts; they are skipped with + recorded reasons while dispatch resolves the #241 PR3 reference.""" + + registry = KernelRegistry() + platform = registry._platform() + # Order the WS1 backends ahead of the reference so their rejections are + # exercised on the way to a successful resolution. + candidates = registry._logprob_candidates[platform] + candidates.remove(OpBackend.PYTORCH_VOCAB_PARALLEL_LOGP) + candidates.append(OpBackend.PYTORCH_VOCAB_PARALLEL_LOGP) + + result = registry.get_logprob_op(_contract()) + assert result.capability.backend_id == "pytorch-vocab-parallel-logp-ws2" + assert result.provenance["fallback"] is True + rejections = " | ".join(result.provenance["prior_rejections"]) + assert "vocab-domain LSE export is unsupported" in rejections + + +def test_undeclared_backend_capability_is_never_selected(): + registry = KernelRegistry() + platform = registry._platform() + registry._logprob_candidates[platform] = [OpBackend.PYTORCH_NATIVE] + + with pytest.raises(RuntimeError, match="no LogprobBackendCapability declared"): + registry.get_logprob_op(_contract()) + + +def test_declared_compatible_backend_resolves_and_records_provenance(): + registry = KernelRegistry() + platform = registry._platform() + registry._logprob_candidates[platform] = [] + registry.register_logprob_backend( + OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, _declared_tp_backend(), platform=platform + ) + + result = registry.get_logprob_op(_contract(), requested_backend="reference") + + assert result.op is not None + assert result.capability.backend_id == "test-deterministic-tp-logprob" + assert result.provenance["requested_backend"] == "reference" + assert result.provenance["actual_backend"] == "test-deterministic-tp-logprob" + assert result.provenance["fallback"] is False + assert result.provenance["contract"]["sharding"]["tp_world_size"] == 2 + assert result.provenance["contract"]["sharding"]["real_vocab_size"] == QWEN3_REAL_VOCAB + assert result.provenance["contract"]["reduction"]["cp_is_merge_axis"] is False + json.dumps(result.provenance) + + +def test_requested_stable_backend_id_is_enforced(): + registry = KernelRegistry() + platform = registry._platform() + registry._logprob_candidates[platform] = [] + registry.register_logprob_backend( + OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, _declared_tp_backend(), platform=platform + ) + + with pytest.raises(RuntimeError, match="does not match requested_backend=another-backend"): + registry.get_logprob_op(_contract(), requested_backend="another-backend") + + result = registry.get_logprob_op(_contract(), requested_backend="test-deterministic-tp-logprob") + assert result.provenance["actual_backend"] == "test-deterministic-tp-logprob" + + +def test_cp_is_a_non_merge_axis_and_cp_agnostic_backends_accept_any_cp_degree(): + capability = _declared_tp_backend() + cp2_contract = _contract(sharding=_sharding(cp_world_size=2, cp_rank=1)) + + assert capability.incompatibilities(cp2_contract) == () + + cp_restricted = replace(capability, cp_world_sizes=(1,)) + assert cp_restricted.incompatibilities(cp2_contract) == ("CP=2 is unsupported",) + + +def test_inactive_tokens_require_explicit_active_mask_support(): + capability = replace(_declared_tp_backend(), mask_modes=frozenset({MaskMode.IGNORE_INDEX})) + contract = _contract() + + assert "explicit active-token masking is unsupported" in ( + capability.incompatibilities(contract) + ) + + fully_active = _contract(mask=_mask(num_tokens=3, active_mask=(True, True, True))) + assert capability.incompatibilities(fully_active) == () + + +def test_backend_id_must_not_shadow_a_reserved_policy_keyword(): + with pytest.raises(LogprobContractError, match="reserved dispatch policy keyword"): + replace(_declared_tp_backend(), backend_id="Deterministic") + + +def test_default_auto_policy_resolves_any_compatible_implementation_kind(): + registry = KernelRegistry() + platform = registry._platform() + registry._logprob_candidates[platform] = [] + registry.register_logprob_backend( + OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, + replace(_declared_tp_backend(), implementation_kind="reference"), + platform=platform, + ) + + result = registry.get_logprob_op(_contract()) + + assert result.provenance["requested_backend"] == "auto" + assert result.capability.implementation_kind == "reference" + + +def test_policy_keywords_are_case_insensitive_but_backend_ids_are_exact(): + registry = KernelRegistry() + platform = registry._platform() + registry._logprob_candidates[platform] = [] + registry.register_logprob_backend( + OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, _declared_tp_backend(), platform=platform + ) + + result = registry.get_logprob_op(_contract(), requested_backend="REFERENCE") + assert result.capability.backend_id == "test-deterministic-tp-logprob" + + with pytest.raises(RuntimeError, match="does not match requested_backend"): + registry.get_logprob_op(_contract(), requested_backend="Test-Deterministic-TP-Logprob") + + +def test_policy_only_skips_are_not_reported_as_fallback(): + registry = KernelRegistry() + platform = registry._platform() + registry._logprob_candidates[platform] = [] + registry.register_logprob_backend( + OpBackend.TRITON_BATCH_INVARIANT_LOGP, + replace(_declared_tp_backend(), backend_id="other-compatible-backend"), + platform=platform, + ) + registry.register_logprob_backend( + OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, _declared_tp_backend(), platform=platform + ) + + result = registry.get_logprob_op(_contract(), requested_backend="test-deterministic-tp-logprob") + + assert result.provenance["fallback"] is False + assert len(result.provenance["prior_rejections"]) == 1 + + +def test_capability_rejections_are_reported_as_fallback(): + registry = KernelRegistry() + platform = registry._platform() + registry._logprob_candidates[platform] = [] + registry.register_logprob_backend( + OpBackend.TRITON_BATCH_INVARIANT_LOGP, + replace(_declared_tp_backend(), backend_id="tp1-only-backend", tp_world_sizes=(1,)), + platform=platform, + ) + registry.register_logprob_backend( + OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, _declared_tp_backend(), platform=platform + ) + + result = registry.get_logprob_op(_contract()) + + assert result.provenance["fallback"] is True + assert "TP=2 is unsupported" in result.provenance["prior_rejections"][0] + + +def test_ws2_candidate_list_is_decoupled_from_the_legacy_priority_map(): + registry = KernelRegistry() + platform = registry._platform() + registry._logprob_candidates[platform].insert(0, OpBackend.PYTORCH_NATIVE) + + legacy = registry._priority_map[platform]["batch_invariant_logp"] + assert OpBackend.PYTORCH_NATIVE not in legacy + + legacy.insert(0, OpBackend.PYTORCH_GEMM) + assert OpBackend.PYTORCH_GEMM not in registry._logprob_candidates[platform] + + +def test_register_logprob_backend_is_the_public_registration_seam(): + registry = KernelRegistry() + platform = registry._platform() + registry._logprob_candidates[platform] = [] + capability = _declared_tp_backend() + + registry.register_logprob_backend( + OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, capability, platform=platform + ) + registry.register_logprob_backend( + OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, + replace(capability, backend_id="replacement-backend"), + platform=platform, + ) + + assert registry._logprob_candidates[platform] == [OpBackend.PYTORCH_BATCH_INVARIANT_LOGP] + result = registry.get_logprob_op(_contract()) + assert result.capability.backend_id == "replacement-backend" + + with pytest.raises(LogprobContractError, match="capability must be"): + registry.register_logprob_backend(OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, None) + + +def test_backend_id_whitespace_is_normalized_for_dispatch(): + capability = replace(_declared_tp_backend(), backend_id=" padded-id ") + assert capability.backend_id == "padded-id" + + +def test_capabilities_are_scoped_per_platform(): + registry = KernelRegistry() + platform = registry._platform() + other = "rocm" if platform != "rocm" else "cpu" + registry._logprob_candidates[platform] = [] + registry.register_logprob_backend( + OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, _declared_tp_backend(), platform=platform + ) + registry.register_logprob_backend( + OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, + replace(_declared_tp_backend(), backend_id="other-platform-backend"), + platform=other, + ) + + result = registry.get_logprob_op(_contract()) + + assert result.capability.backend_id == "test-deterministic-tp-logprob" + assert ( + registry._logprob_capabilities[other][OpBackend.PYTORCH_BATCH_INVARIANT_LOGP].backend_id + == "other-platform-backend" + ) + + +def test_register_logprob_backend_rejects_unknown_platform(): + registry = KernelRegistry() + + with pytest.raises(LogprobContractError, match="unsupported platform"): + registry.register_logprob_backend( + OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, + _declared_tp_backend(), + platform="cuda-typo", + ) + + +def test_non_iterable_roles_and_dtypes_raise_contract_errors(): + with pytest.raises(LogprobContractError, match="roles and dtypes must be iterables"): + replace(_declared_tp_backend(), roles=None) + + with pytest.raises(LogprobContractError, match="roles and dtypes must be iterables"): + replace(_declared_tp_backend(), dtypes=42) + + +def test_requested_deterministic_policy_is_a_loud_error(): + registry = KernelRegistry() + + with pytest.raises(LogprobContractError, match="determinism_scope"): + registry.get_logprob_op(_contract(), requested_backend="deterministic") + + +def test_determinism_scope_is_part_of_the_typed_contract(): + fixed_only = replace( + _declared_tp_backend(), + determinism_scopes=frozenset({DeterminismScope.FIXED_TOPOLOGY}), + ) + + assert "determinism_scope=cross_tp_bitwise is unsupported" in ( + fixed_only.incompatibilities(_contract()) + ) + + relaxed = _contract(reduction=ReductionSpec(determinism_scope="fixed_topology")) + assert fixed_only.incompatibilities(relaxed) == () + + +def test_policy_filtered_candidates_never_count_toward_fallback(): + registry = KernelRegistry() + platform = registry._platform() + registry._logprob_candidates[platform] = [] + registry.register_logprob_backend( + OpBackend.TRITON_BATCH_INVARIANT_LOGP, + replace(_declared_tp_backend(), backend_id="tp1-only-backend", tp_world_sizes=(1,)), + platform=platform, + ) + registry.register_logprob_backend( + OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, _declared_tp_backend(), platform=platform + ) + + result = registry.get_logprob_op(_contract(), requested_backend="test-deterministic-tp-logprob") + + assert result.provenance["fallback"] is False + assert len(result.provenance["prior_rejections"]) == 1 + + +def test_output_spec_is_pinned_to_fp32_replicated(): + with pytest.raises(LogprobContractError, match="must be fp32"): + LogprobOutputSpec(selected_logp_dtype="bf16") + with pytest.raises(LogprobContractError, match="must be fp32"): + LogprobOutputSpec(lse_dtype="bf16") + + assert LogprobOutputSpec().tp_placement is TPPlacement.REPLICATED + assert _contract().to_dict()["output"] == { + "selected_logp_dtype": "fp32", + "lse_dtype": "fp32", + "tp_placement": "replicated", + } + + +def test_cross_rank_fingerprint_is_rank_independent_and_content_sensitive(): + rank0 = _contract(sharding=_sharding(tp_rank=0)) + rank1 = _contract(sharding=_sharding(tp_rank=1, cp_rank=1)) + + assert rank0.cross_rank_fingerprint() == rank1.cross_rank_fingerprint() + + different_mask = _contract( + mask=_mask(active_mask=(True, True, True, True, True, True, True, False)) + ) + assert rank0.cross_rank_fingerprint() != different_mask.cross_rank_fingerprint() + + +def test_provenance_records_the_active_mask_digest(): + provenance_mask = _contract().to_dict()["mask"] + + assert provenance_mask["active_mask_sha256"] == _mask().active_mask_sha256 + assert len(provenance_mask["active_mask_sha256"]) == 64 + + same_count_different_mask = _mask( + active_mask=(True, True, True, True, True, False, False, False) + ) + assert same_count_different_mask.active_token_count == _mask().active_token_count + assert same_count_different_mask.active_mask_sha256 != _mask().active_mask_sha256 + + +def test_padding_only_shard_is_constructible_for_the_identity_partial(): + sharding = _sharding( + vocab_shard_bounds=((0, QWEN3_REAL_VOCAB), (QWEN3_REAL_VOCAB, QWEN3_PADDED_VOCAB)), + ) + + assert sharding.local_vocab_start == 0 + assert sharding.vocab_shard_bounds[1] == (QWEN3_REAL_VOCAB, QWEN3_PADDED_VOCAB) + assert sharding.owner_rank(QWEN3_REAL_VOCAB - 1) == 0 diff --git a/tests/test_module_debug_matrix.py b/tests/test_module_debug_matrix.py new file mode 100644 index 00000000..3830b783 --- /dev/null +++ b/tests/test_module_debug_matrix.py @@ -0,0 +1,61 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import pytest + +from rl_engine.alignment.cross_config.debug_matrix import ( + DEBUG_MATRIX_SCHEMA_VERSION, + MODULE_DEBUG_AXES, + module_debug_axis, + module_debug_matrix, +) + + +def test_module_matrix_is_fixed_replay_and_compact(): + manifest = module_debug_matrix() + + assert manifest["schema_version"] == DEBUG_MATRIX_SCHEMA_VERSION + assert manifest["method"] == "fixed_replay_one_at_a_time" + assert manifest["cartesian_product"] is False + assert manifest["comparison_edges"] == [ + "train_vs_rollout_prefill", + "rollout_prefill_vs_decode", + ] + assert set(manifest["modules"]) == {"attention", "ffn", "logp"} + assert manifest["modules"]["attention"]["rows"] == [ + "A0", + "A1", + "A2", + "A3", + "A4", + "A5", + "A6", + "A7", + ] + assert manifest["modules"]["ffn"]["rows"] == ["F0", "F1", "F2", "F3", "F4"] + assert manifest["modules"]["logp"]["rows"] == ["L0", "L1", "L2", "L3"] + assert [row["row"] for row in manifest["modules"]["ffn"]["axes"]] == [ + "F1", + "F2", + "F3", + "F4", + ] + + +def test_module_axes_mark_identity_and_topology_as_gates(): + assert module_debug_axis("attention", "topology_head_ownership").kind == "gate" + assert module_debug_axis("ffn", "weight_shard_ownership").kind == "gate" + assert module_debug_axis("logp", "vocab_shard_ownership").kind == "gate" + assert module_debug_axis("logp", "selected_token_identity").kind == "gate" + assert module_debug_axis("ffn", "gemm_reduction").kind == "diagnostic" + + with pytest.raises(ValueError, match="unknown module debug axis"): + module_debug_axis("ffn", "not_an_axis") + + +def test_module_axes_have_unique_ids_and_probes(): + for module, axes in MODULE_DEBUG_AXES.items(): + assert all(axis.module == module for axis in axes) + assert len({axis.axis_id for axis in axes}) == len(axes) + assert len({axis.representative_probe for axis in axes}) == len(axes) diff --git a/tests/test_qwen_ffn.py b/tests/test_qwen_ffn.py new file mode 100644 index 00000000..9b2b1d21 --- /dev/null +++ b/tests/test_qwen_ffn.py @@ -0,0 +1,1112 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Tests for the deterministic Qwen3 dense FFN. + +Covers single-GPU correctness, token boundaries, Qwen3-8B shapes, TP/CP/SP +bitwise alignment, and DeterministicCollective cache lifetime. +""" + +from __future__ import annotations + +import queue +import tempfile +import traceback +from datetime import timedelta +from pathlib import Path + +import pytest +import torch +import torch.multiprocessing as mp +import torch.nn.functional as F + +import rl_engine.kernels.ops.pytorch.ffn.ffn as ffn_module +from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE +from rl_engine.kernels.ops.pytorch.ffn.ffn import ( + QWEN3_8B_HIDDEN_SIZE, + QWEN3_8B_INTERMEDIATE_SIZE, + qwen3_ffn, +) + +_REQUIRED_SYMBOLS = ( + "det_gemm_fwd", + "det_gemm_db", + "swiglu_forward", + "swiglu_backward", +) +_HIDDEN = 64 +_INTERMEDIATE = 512 +_TOPOLOGY_TOKENS = 256 +_CP_TOKEN_COUNTS = (8, 32, 64, 96, 128, 256) +_TOKEN_BOUNDARY_COUNTS = (8, 31, 32, 33, 64, 96, 128) +_WORLD2_CONFIGS = ( + ("tp2_sp", 2, 1, True, _TOPOLOGY_TOKENS), + ("cp2", 1, 2, False, _TOPOLOGY_TOKENS), + *((f"cp2_T{token_count}", 1, 2, False, token_count) for token_count in _CP_TOKEN_COUNTS), +) +_WORLD4_CONFIGS = ( + ("tp4", 4, 1, False, _TOPOLOGY_TOKENS), + ("cp4", 1, 4, False, _TOPOLOGY_TOKENS), + ("tp2_cp2", 2, 2, False, _TOPOLOGY_TOKENS), + ("tp2_cp2_sp", 2, 2, True, _TOPOLOGY_TOKENS), + *((f"cp4_T{token_count}", 1, 4, False, token_count) for token_count in _CP_TOKEN_COUNTS), +) +_WORLD8_WORLD_GROUP_CONFIGS = ( + ("tp8", 8, 1, False, _TOPOLOGY_TOKENS), + ("tp8_sp", 8, 1, True, _TOPOLOGY_TOKENS), + ("cp8", 1, 8, False, _TOPOLOGY_TOKENS), + *((f"cp8_T{token_count}", 1, 8, False, token_count) for token_count in _CP_TOKEN_COUNTS), +) +_WORLD8_TP2_CP4_CONFIGS = (("tp2_cp4", 2, 4, False, _TOPOLOGY_TOKENS),) +_WORLD8_TP4_CP2_CONFIGS = ( + ("tp4_cp2", 4, 2, False, _TOPOLOGY_TOKENS), + ("tp4_cp2_sp", 4, 2, True, _TOPOLOGY_TOKENS), +) + + +def _has_sm90_ffn_devices(count: int) -> bool: + return ( + _EXT_AVAILABLE + and torch.distributed.is_available() + and torch.distributed.is_nccl_available() + and torch.cuda.device_count() >= count + and all(torch.cuda.get_device_capability(index)[0] == 9 for index in range(count)) + and all(hasattr(_C, name) for name in _REQUIRED_SYMBOLS) + ) + + +def _has_sm90_ffn() -> bool: + return ( + torch.cuda.is_available() + and torch.cuda.get_device_capability()[0] == 9 + and _EXT_AVAILABLE + and all(hasattr(_C, name) for name in _REQUIRED_SYMBOLS) + ) + + +requires_cuda_ffn = pytest.mark.skipif( + not _has_sm90_ffn(), + reason="FFN optimized-path validation requires SM90 and the GEMM/SwiGLU extension", +) + + +class _TorchKernelStub: + def __init__(self) -> None: + self.calls: list[str] = [] + + def det_gemm_fwd(self, a, b): + self.calls.append("det_gemm_fwd") + return a @ b + + def det_gemm_db(self, a, grad_output): + self.calls.append("det_gemm_db") + return a.t().contiguous() @ grad_output + + def swiglu_forward(self, gate, up): + self.calls.append("swiglu_forward") + return gate * torch.sigmoid(gate) * up + + def swiglu_backward(self, grad_output, gate, up): + self.calls.append("swiglu_backward") + sigmoid = torch.sigmoid(gate) + grad_gate = grad_output * up * sigmoid * (1.0 + gate * (1.0 - sigmoid)) + grad_up = grad_output * gate * sigmoid + return grad_gate, grad_up + + +def _reference(hidden_states, gate_weight, up_weight, down_weight): + gate = hidden_states @ gate_weight.t() + up = hidden_states @ up_weight.t() + activated = F.silu(gate) * up + return (activated @ down_weight.t()), gate, up, activated + + +def _randn(shape, *, seed, device="cpu", dtype=torch.float32): + generator = torch.Generator(device="cpu").manual_seed(seed) + value = torch.randn(*shape, generator=generator, dtype=torch.float32) * 0.02 + return value.to(device=device, dtype=dtype) + + +def _close_ffn_collectives() -> None: + for collective in list(ffn_module._COLLECTIVES.values()): + collective.close() + ffn_module._COLLECTIVES.clear() + + +def _shard_ranges( + rank: int, + *, + tp_size: int, + cp_size: int, + sequence_parallel: bool, + token_count: int, + intermediate_size: int, +) -> tuple[int, int, int, int]: + tp_rank = rank % tp_size + cp_rank = rank // tp_size + cp_tokens = token_count // cp_size + local_tokens = cp_tokens // tp_size if sequence_parallel else cp_tokens + token_start = cp_rank * cp_tokens + if sequence_parallel: + token_start += tp_rank * local_tokens + token_end = token_start + local_tokens + local_i = intermediate_size // tp_size + feat_start = tp_rank * local_i + feat_end = feat_start + local_i + return token_start, token_end, feat_start, feat_end + + +def _spawn_nccl_workers(worker, world_size: int, worker_args=(), *, timeout: int = 180) -> None: + if not _has_sm90_ffn_devices(world_size): + pytest.skip(f"requires {world_size} SM90 GPUs, NCCL, and the GEMM/SwiGLU extension") + + ctx = mp.get_context("spawn") + with tempfile.TemporaryDirectory() as tmpdir: + init_method = (Path(tmpdir) / "nccl_init").as_uri() + result_queue = ctx.Queue() + processes = [ + ctx.Process( + target=worker, + args=(rank, world_size, init_method, result_queue, *worker_args), + ) + for rank in range(world_size) + ] + for process in processes: + process.start() + results = [] + try: + for _ in processes: + results.append(result_queue.get(timeout=timeout)) + except queue.Empty: + for process in processes: + if process.is_alive(): + process.terminate() + pytest.fail(f"timed out waiting for {world_size} FFN workers") + finally: + for process in processes: + process.join(timeout=10) + if process.is_alive(): + process.terminate() + + for result in sorted(results, key=lambda item: item["rank"]): + assert result["ok"], result.get("traceback") or result.get("failures") + for process in processes: + assert process.exitcode == 0 + + +def _distributed_ffn_backward_nccl_worker( + rank, + world_size, + init_method, + result_queue, + cp_size, + sequence_parallel, +): + try: + import torch.distributed as dist + + torch.cuda.set_device(rank) + dist.init_process_group( + backend="nccl", + init_method=init_method, + rank=rank, + world_size=world_size, + ) + + tp_size = world_size // cp_size + tp_groups = [ + dist.new_group(list(range(cp_rank * tp_size, (cp_rank + 1) * tp_size))) + for cp_rank in range(cp_size) + ] + cp_groups = [ + dist.new_group([cp_rank * tp_size + tp_rank for cp_rank in range(cp_size)]) + for tp_rank in range(tp_size) + ] + tp_rank = rank % tp_size + cp_rank = rank // tp_size + tp_group = tp_groups[cp_rank] + cp_group = cp_groups[tp_rank] if cp_size > 1 else None + + token_count, hidden_size, intermediate_size = 8, 64, 128 + token_start, token_end, feature_start, feature_end = _shard_ranges( + rank, + tp_size=tp_size, + cp_size=cp_size, + sequence_parallel=sequence_parallel, + token_count=token_count, + intermediate_size=intermediate_size, + ) + local_tokens = token_end - token_start + + device = torch.device("cuda", rank) + rmsnorm_output = _randn( + (token_count, hidden_size), seed=40, device=device, dtype=torch.bfloat16 + ) + gate_weight = _randn( + (intermediate_size, hidden_size), + seed=41, + device=device, + dtype=torch.bfloat16, + ) + up_weight = _randn( + (intermediate_size, hidden_size), + seed=42, + device=device, + dtype=torch.bfloat16, + ) + down_weight = _randn( + (hidden_size, intermediate_size), + seed=43, + device=device, + dtype=torch.bfloat16, + ) + grad_output = _randn( + (token_count, hidden_size), seed=44, device=device, dtype=torch.bfloat16 + ) + + reference_inputs = [ + value.detach().float().requires_grad_(True) + for value in (rmsnorm_output, gate_weight, up_weight, down_weight) + ] + reference_output, _, _, _ = _reference(*reference_inputs) + reference_output.backward(grad_output.float()) + + local_grad_output = grad_output[token_start:token_end].contiguous() + actual_inputs = [ + value.detach().clone().requires_grad_(True) + for value in ( + rmsnorm_output[token_start:token_end].contiguous(), + gate_weight[feature_start:feature_end].contiguous(), + up_weight[feature_start:feature_end].contiguous(), + down_weight[:, feature_start:feature_end].contiguous(), + ) + ] + actual_output = qwen3_ffn( + *actual_inputs, + tp_group=tp_group, + cp_group=cp_group, + sequence_parallel=sequence_parallel, + ) + actual_output.backward(local_grad_output) + + expected_grads = ( + reference_inputs[0].grad[token_start:token_end], + reference_inputs[1].grad[feature_start:feature_end], + reference_inputs[2].grad[feature_start:feature_end], + reference_inputs[3].grad[:, feature_start:feature_end], + ) + torch.testing.assert_close( + actual_output.float(), + reference_output[token_start:token_end].detach(), + atol=5e-2, + rtol=2e-2, + ) + for actual, expected in zip(actual_inputs, expected_grads, strict=True): + torch.testing.assert_close( + actual.grad.float(), + expected, + atol=5e-2, + rtol=2e-2, + ) + + slice_size = max(1, local_tokens // 2) + slice_start = (local_tokens - slice_size) // 2 + slice_end = slice_start + slice_size + slice_inputs = [ + value.detach().clone().requires_grad_(True) + for value in ( + actual_inputs[0][slice_start:slice_end].contiguous(), + actual_inputs[1], + actual_inputs[2], + actual_inputs[3], + ) + ] + slice_output = qwen3_ffn( + *slice_inputs, + tp_group=tp_group, + cp_group=cp_group, + sequence_parallel=sequence_parallel, + ) + slice_output.backward(local_grad_output[slice_start:slice_end]) + + assert torch.equal( + slice_output, + actual_output[slice_start:slice_end], + ), "FFN output changed with the local token batch size" + assert torch.equal( + slice_inputs[0].grad, + actual_inputs[0].grad[slice_start:slice_end], + ), "FFN input gradient changed with the local token batch size" + result_queue.put({"ok": True, "rank": rank}) + except Exception: # pragma: no cover - forwarded to the parent process. + result_queue.put({"ok": False, "rank": rank, "traceback": traceback.format_exc()}) + raise + finally: + _close_ffn_collectives() + if torch.distributed.is_available() and torch.distributed.is_initialized(): + torch.distributed.destroy_process_group() + + +def _tp1_vs_tpn_train_infer_worker(rank, world_size, init_method, result_queue, expect_match): + try: + import torch.distributed as dist + + torch.cuda.set_device(rank) + dist.init_process_group( + backend="nccl", + init_method=init_method, + rank=rank, + world_size=world_size, + ) + device = torch.device("cuda", rank) + token_count, hidden_size, intermediate_size = 16, 64, 256 + hidden = _randn((token_count, hidden_size), seed=50, device=device, dtype=torch.bfloat16) + gate_weight = _randn( + (intermediate_size, hidden_size), + seed=51, + device=device, + dtype=torch.bfloat16, + ) + up_weight = _randn( + (intermediate_size, hidden_size), + seed=52, + device=device, + dtype=torch.bfloat16, + ) + down_weight = _randn( + (hidden_size, intermediate_size), + seed=53, + device=device, + dtype=torch.bfloat16, + ) + grad_output = _randn( + (token_count, hidden_size), seed=54, device=device, dtype=torch.bfloat16 + ) + + with torch.no_grad(): + infer_tp1 = qwen3_ffn(hidden, gate_weight, up_weight, down_weight) + + tp1_inputs = [ + value.detach().clone().requires_grad_(True) + for value in (hidden, gate_weight, up_weight, down_weight) + ] + train_tp1 = qwen3_ffn(*tp1_inputs) + train_tp1.backward(grad_output) + assert torch.equal(infer_tp1, train_tp1.detach()), "TP=1 train/infer forward mismatch" + + local_i = intermediate_size // world_size + feat_start = rank * local_i + feat_end = feat_start + local_i + shard = ( + hidden, + gate_weight[feat_start:feat_end].contiguous(), + up_weight[feat_start:feat_end].contiguous(), + down_weight[:, feat_start:feat_end].contiguous(), + ) + with torch.no_grad(): + infer_tpn = qwen3_ffn(*shard, tp_group=dist.group.WORLD) + + tpn_inputs = [value.detach().clone().requires_grad_(True) for value in shard] + train_tpn = qwen3_ffn(*tpn_inputs, tp_group=dist.group.WORLD) + train_tpn.backward(grad_output) + assert torch.equal( + infer_tpn, train_tpn.detach() + ), f"TP={world_size} train/infer forward mismatch" + + infer_match = torch.equal(infer_tp1, infer_tpn) + train_match = torch.equal(train_tp1.detach(), train_tpn.detach()) + hidden_match = torch.equal(tp1_inputs[0].grad, tpn_inputs[0].grad) + if expect_match: + assert infer_match, f"TP=1 vs TP={world_size} infer forward mismatch" + assert train_match, f"TP=1 vs TP={world_size} train forward mismatch" + assert hidden_match, f"TP=1 vs TP={world_size} hidden grad mismatch" + else: + assert not infer_match, f"TP=1 vs TP={world_size} infer forward unexpectedly matched" + assert not train_match, f"TP=1 vs TP={world_size} train forward unexpectedly matched" + assert not hidden_match, f"TP=1 vs TP={world_size} hidden grad unexpectedly matched" + + assert torch.equal( + tp1_inputs[1].grad[feat_start:feat_end], tpn_inputs[1].grad + ), f"TP=1 vs TP={world_size} gate weight grad mismatch" + assert torch.equal( + tp1_inputs[2].grad[feat_start:feat_end], tpn_inputs[2].grad + ), f"TP=1 vs TP={world_size} up weight grad mismatch" + assert torch.equal( + tp1_inputs[3].grad[:, feat_start:feat_end], tpn_inputs[3].grad + ), f"TP=1 vs TP={world_size} down weight grad mismatch" + + result_queue.put({"ok": True, "rank": rank}) + except Exception: # pragma: no cover - forwarded to the parent process. + result_queue.put({"ok": False, "rank": rank, "traceback": traceback.format_exc()}) + raise + finally: + _close_ffn_collectives() + if torch.distributed.is_available() and torch.distributed.is_initialized(): + torch.distributed.destroy_process_group() + + +def _make_topology_inputs(token_count, device): + hidden = _randn((token_count, _HIDDEN), seed=60, device=device, dtype=torch.bfloat16) + gate = _randn((_INTERMEDIATE, _HIDDEN), seed=61, device=device, dtype=torch.bfloat16) + up = _randn((_INTERMEDIATE, _HIDDEN), seed=62, device=device, dtype=torch.bfloat16) + down = _randn((_HIDDEN, _INTERMEDIATE), seed=63, device=device, dtype=torch.bfloat16) + grad = _randn((token_count, _HIDDEN), seed=64, device=device, dtype=torch.bfloat16) + return hidden, gate, up, down, grad + + +def _canonical(hidden, gate, up, down, grad): + with torch.no_grad(): + infer = qwen3_ffn(hidden, gate, up, down) + inputs = [value.detach().clone().requires_grad_(True) for value in (hidden, gate, up, down)] + train = qwen3_ffn(*inputs) + train.backward(grad) + return infer, train, inputs + + +def _mesh_groups(dist, tp_size, cp_size): + world_size = dist.get_world_size() + if tp_size == world_size and cp_size == 1: + return [dist.group.WORLD], [] + if cp_size == world_size and tp_size == 1: + return [], [dist.group.WORLD] + tp_groups = [] + if tp_size > 1: + for cp_rank in range(cp_size): + ranks = list(range(cp_rank * tp_size, (cp_rank + 1) * tp_size)) + tp_groups.append(dist.new_group(ranks)) + cp_groups = [] + if cp_size > 1: + for tp_rank in range(tp_size): + ranks = [cp_rank * tp_size + tp_rank for cp_rank in range(cp_size)] + cp_groups.append(dist.new_group(ranks)) + return tp_groups, cp_groups + + +def _run_topology_config( + rank, + dist, + meshes, + *, + name, + tp_size, + cp_size, + sequence_parallel, + hidden, + gate, + up, + down, + grad, + infer_ref, + train_ref, + ref_inputs, +): + key = (tp_size, cp_size) + if key not in meshes: + meshes[key] = _mesh_groups(dist, tp_size, cp_size) + tp_groups, cp_groups = meshes[key] + tp_rank = rank % tp_size + cp_rank = rank // tp_size + tp_group = tp_groups[cp_rank] if tp_size > 1 else None + cp_group = cp_groups[tp_rank] if cp_size > 1 else None + token_start, token_end, feat_start, feat_end = _shard_ranges( + rank, + tp_size=tp_size, + cp_size=cp_size, + sequence_parallel=sequence_parallel, + token_count=hidden.size(0), + intermediate_size=_INTERMEDIATE, + ) + shard = ( + hidden[token_start:token_end].contiguous(), + gate[feat_start:feat_end].contiguous(), + up[feat_start:feat_end].contiguous(), + down[:, feat_start:feat_end].contiguous(), + ) + with torch.no_grad(): + infer = qwen3_ffn( + *shard, + tp_group=tp_group, + cp_group=cp_group, + sequence_parallel=sequence_parallel, + ) + inputs = [value.detach().clone().requires_grad_(True) for value in shard] + train = qwen3_ffn( + *inputs, + tp_group=tp_group, + cp_group=cp_group, + sequence_parallel=sequence_parallel, + ) + train.backward(grad[token_start:token_end].contiguous()) + + assert torch.equal(infer, train.detach()), f"{name}: train/infer forward mismatch" + assert torch.equal( + infer, infer_ref[token_start:token_end] + ), f"{name}: infer forward mismatch vs TP=1/CP=1" + assert torch.equal( + train.detach(), train_ref.detach()[token_start:token_end] + ), f"{name}: train forward mismatch vs TP=1/CP=1" + assert torch.equal( + inputs[0].grad, ref_inputs[0].grad[token_start:token_end] + ), f"{name}: hidden grad mismatch vs TP=1/CP=1" + + weight_checks = ( + (1, ref_inputs[1].grad[feat_start:feat_end], "gate"), + (2, ref_inputs[2].grad[feat_start:feat_end], "up"), + (3, ref_inputs[3].grad[:, feat_start:feat_end], "down"), + ) + for index, expected, label in weight_checks: + assert torch.equal( + inputs[index].grad, expected + ), f"{name}: {label} weight grad mismatch vs TP=1/CP=1" + + +def _topology_worker(rank, world_size, init_method, result_queue, configs): + try: + import torch.distributed as dist + + torch.cuda.set_device(rank) + dist.init_process_group( + backend="nccl", + init_method=init_method, + rank=rank, + world_size=world_size, + timeout=timedelta(minutes=5), + device_id=torch.device("cuda", rank), + ) + device = torch.device("cuda", rank) + meshes = {} + canonical = {} + for name, tp_size, cp_size, sequence_parallel, token_count in configs: + if token_count not in canonical: + tensors = _make_topology_inputs(token_count, device) + canonical[token_count] = (*tensors, *_canonical(*tensors)) + hidden, gate, up, down, grad, infer_ref, train_ref, ref_inputs = canonical[token_count] + _run_topology_config( + rank, + dist, + meshes, + name=name, + tp_size=tp_size, + cp_size=cp_size, + sequence_parallel=sequence_parallel, + hidden=hidden, + gate=gate, + up=up, + down=down, + grad=grad, + infer_ref=infer_ref, + train_ref=train_ref, + ref_inputs=ref_inputs, + ) + result_queue.put({"ok": True, "rank": rank}) + except Exception: # pragma: no cover - forwarded to the parent process. + result_queue.put({"ok": False, "rank": rank, "traceback": traceback.format_exc()}) + raise + finally: + _close_ffn_collectives() + if torch.distributed.is_available() and torch.distributed.is_initialized(): + torch.distributed.destroy_process_group() + + +def _ffn_tensors(token_count, device, seed, *, hidden=_HIDDEN, intermediate=_INTERMEDIATE): + rmsnorm = _randn((token_count, hidden), seed=seed, device=device, dtype=torch.bfloat16) + gate = _randn((intermediate, hidden), seed=seed + 1, device=device, dtype=torch.bfloat16) + up = _randn((intermediate, hidden), seed=seed + 2, device=device, dtype=torch.bfloat16) + down = _randn((hidden, intermediate), seed=seed + 3, device=device, dtype=torch.bfloat16) + return rmsnorm, gate, up, down + + +def _cache_worker(rank, world_size, init_method, result_queue): + try: + import torch.distributed as dist + + torch.cuda.set_device(rank) + dist.init_process_group( + backend="nccl", + init_method=init_method, + rank=rank, + world_size=world_size, + ) + device = torch.device("cuda", rank) + ffn_module._COLLECTIVE_MIN_CAPACITY_BYTES = 64 + _close_ffn_collectives() + + small = _ffn_tensors(8, device, seed=100, intermediate=128) + first = qwen3_ffn(*small, tp_group=dist.group.WORLD) + assert len(ffn_module._COLLECTIVES) == 1 + ((cache_key, first_collective),) = ffn_module._COLLECTIVES.items() + first_handle = first_collective._handle + first_capacity = first_collective.max_size_bytes + assert first_handle != 0 + + repeated = qwen3_ffn(*small, tp_group=dist.group.WORLD) + assert torch.equal(first, repeated) + assert len(ffn_module._COLLECTIVES) == 1 + assert ffn_module._COLLECTIVES[cache_key] is first_collective + assert first_collective._handle == first_handle + + large = _ffn_tensors(256, device, seed=110, intermediate=128) + grown = qwen3_ffn(*large, tp_group=dist.group.WORLD) + assert grown.shape[0] == 256 + assert len(ffn_module._COLLECTIVES) == 1 + grown_collective = next(iter(ffn_module._COLLECTIVES.values())) + assert grown_collective is not first_collective + assert first_collective._handle == 0 + assert grown_collective.max_size_bytes > first_capacity + assert grown_collective._handle != 0 + + _close_ffn_collectives() + assert ffn_module._COLLECTIVES == {} + recreated = qwen3_ffn(*small, tp_group=dist.group.WORLD) + assert torch.equal(recreated, first) + assert len(ffn_module._COLLECTIVES) == 1 + recreated_collective = next(iter(ffn_module._COLLECTIVES.values())) + assert recreated_collective is not grown_collective + assert recreated_collective._handle != 0 + + rebuilt_group = dist.new_group(ranks=[0, 1]) + rebuilt = qwen3_ffn(*small, tp_group=rebuilt_group) + assert torch.equal(rebuilt, first) + assert len(ffn_module._COLLECTIVES) == 2 + + _close_ffn_collectives() + result_queue.put({"ok": True, "rank": rank}) + except Exception: # pragma: no cover - forwarded to the parent process. + result_queue.put({"ok": False, "rank": rank, "traceback": traceback.format_exc()}) + raise + finally: + _close_ffn_collectives() + if torch.distributed.is_available() and torch.distributed.is_initialized(): + torch.distributed.destroy_process_group() + + +def _uneven_sp_worker(rank, world_size, init_method, result_queue): + try: + import torch.distributed as dist + + torch.cuda.set_device(rank) + dist.init_process_group( + backend="nccl", + init_method=init_method, + rank=rank, + world_size=world_size, + ) + device = torch.device("cuda", rank) + hidden, gate, up, down = _ffn_tensors(2, device, seed=130, intermediate=128) + local_i = 128 // world_size + feat_start = rank * local_i + feat_end = feat_start + local_i + local_hidden = hidden[:2] if rank == 0 else hidden[:1] + try: + qwen3_ffn( + local_hidden, + gate[feat_start:feat_end].contiguous(), + up[feat_start:feat_end].contiguous(), + down[:, feat_start:feat_end].contiguous(), + tp_group=dist.group.WORLD, + sequence_parallel=True, + ) + result_queue.put( + { + "ok": False, + "rank": rank, + "failures": "uneven SP tokens should have raised", + } + ) + except ValueError as exc: + message = str(exc) + if "matching shapes" not in message and "world_size" not in message: + raise + result_queue.put({"ok": True, "rank": rank}) + except Exception: # pragma: no cover - forwarded to the parent process. + result_queue.put({"ok": False, "rank": rank, "traceback": traceback.format_exc()}) + raise + finally: + _close_ffn_collectives() + if torch.distributed.is_available() and torch.distributed.is_initialized(): + torch.distributed.destroy_process_group() + + +def _qwen3_8b_weights(device): + hidden = _randn((8, QWEN3_8B_HIDDEN_SIZE), seed=90, device=device, dtype=torch.bfloat16) + gate = _randn( + (QWEN3_8B_INTERMEDIATE_SIZE, QWEN3_8B_HIDDEN_SIZE), + seed=91, + device=device, + dtype=torch.bfloat16, + ) + up = _randn( + (QWEN3_8B_INTERMEDIATE_SIZE, QWEN3_8B_HIDDEN_SIZE), + seed=92, + device=device, + dtype=torch.bfloat16, + ) + down = _randn( + (QWEN3_8B_HIDDEN_SIZE, QWEN3_8B_INTERMEDIATE_SIZE), + seed=93, + device=device, + dtype=torch.bfloat16, + ) + grad = _randn((8, QWEN3_8B_HIDDEN_SIZE), seed=94, device=device, dtype=torch.bfloat16) + return hidden, gate, up, down, grad + + +def _qwen3_8b_tp2_worker(rank, world_size, init_method, result_queue): + try: + import torch.distributed as dist + + torch.cuda.set_device(rank) + dist.init_process_group( + backend="nccl", + init_method=init_method, + rank=rank, + world_size=world_size, + ) + device = torch.device("cuda", rank) + hidden, gate, up, down, grad = _qwen3_8b_weights(device) + with torch.no_grad(): + infer_tp1 = qwen3_ffn(hidden, gate, up, down) + tp1_inputs = [ + value.detach().clone().requires_grad_(True) for value in (hidden, gate, up, down) + ] + train_tp1 = qwen3_ffn(*tp1_inputs) + train_tp1.backward(grad) + + local_i = QWEN3_8B_INTERMEDIATE_SIZE // world_size + feat_start = rank * local_i + feat_end = feat_start + local_i + shard = ( + hidden, + gate[feat_start:feat_end].contiguous(), + up[feat_start:feat_end].contiguous(), + down[:, feat_start:feat_end].contiguous(), + ) + with torch.no_grad(): + infer_tp2 = qwen3_ffn(*shard, tp_group=dist.group.WORLD) + tp2_inputs = [value.detach().clone().requires_grad_(True) for value in shard] + train_tp2 = qwen3_ffn(*tp2_inputs, tp_group=dist.group.WORLD) + train_tp2.backward(grad) + + assert torch.equal(infer_tp1, infer_tp2) + assert torch.equal(train_tp1.detach(), train_tp2.detach()) + assert torch.equal(tp1_inputs[0].grad, tp2_inputs[0].grad) + assert torch.equal(tp1_inputs[1].grad[feat_start:feat_end], tp2_inputs[1].grad) + assert torch.equal(tp1_inputs[2].grad[feat_start:feat_end], tp2_inputs[2].grad) + assert torch.equal(tp1_inputs[3].grad[:, feat_start:feat_end], tp2_inputs[3].grad) + result_queue.put({"ok": True, "rank": rank}) + except Exception: # pragma: no cover - forwarded to the parent process. + result_queue.put({"ok": False, "rank": rank, "traceback": traceback.format_exc()}) + raise + finally: + _close_ffn_collectives() + if torch.distributed.is_available() and torch.distributed.is_initialized(): + torch.distributed.destroy_process_group() + + +def test_qwen_ffn_qwen3_8b_dimensions_are_pinned(): + assert QWEN3_8B_HIDDEN_SIZE == 4096 + assert QWEN3_8B_INTERMEDIATE_SIZE == 12288 + + +def test_qwen_ffn_backward_matches_autograd_reference(monkeypatch): + stub = _TorchKernelStub() + monkeypatch.setattr(ffn_module, "_C", stub) + monkeypatch.setattr(ffn_module, "_EXT_AVAILABLE", True) + monkeypatch.setattr(ffn_module, "_validate_ffn_inputs", lambda *args: None) + + hidden = _randn((2, 3, 8), seed=0) + gate_weight = _randn((12, 8), seed=1) + up_weight = _randn((12, 8), seed=2) + down_weight = _randn((8, 12), seed=3) + grad_output = _randn(hidden.shape, seed=4) + + ref_inputs = [ + value.detach().clone().requires_grad_(True) + for value in (hidden, gate_weight, up_weight, down_weight) + ] + expected, _, _, _ = _reference(*ref_inputs) + expected.backward(grad_output) + + actual_inputs = [ + value.detach().clone().requires_grad_(True) + for value in (hidden, gate_weight, up_weight, down_weight) + ] + actual = qwen3_ffn(*actual_inputs) + actual.backward(grad_output) + + torch.testing.assert_close(actual, expected.detach()) + for actual_input, reference in zip(actual_inputs, ref_inputs, strict=True): + torch.testing.assert_close(actual_input.grad, reference.grad) + + assert stub.calls.count("det_gemm_fwd") == 6 + assert stub.calls.count("det_gemm_db") == 3 + assert stub.calls.count("swiglu_forward") == 1 + assert stub.calls.count("swiglu_backward") == 1 + + +def test_qwen_ffn_disable_split_k_false_uses_torch_matmul(monkeypatch): + stub = _TorchKernelStub() + monkeypatch.setattr(ffn_module, "_C", stub) + monkeypatch.setattr(ffn_module, "_EXT_AVAILABLE", True) + monkeypatch.setattr(ffn_module, "_validate_ffn_inputs", lambda *args: None) + + hidden = _randn((2, 3, 8), seed=0) + gate_weight = _randn((12, 8), seed=1) + up_weight = _randn((12, 8), seed=2) + down_weight = _randn((8, 12), seed=3) + grad_output = _randn(hidden.shape, seed=4) + + ref_inputs = [ + value.detach().clone().requires_grad_(True) + for value in (hidden, gate_weight, up_weight, down_weight) + ] + expected, _, _, _ = _reference(*ref_inputs) + expected.backward(grad_output) + + actual_inputs = [ + value.detach().clone().requires_grad_(True) + for value in (hidden, gate_weight, up_weight, down_weight) + ] + actual = qwen3_ffn(*actual_inputs, disable_split_k=False) + actual.backward(grad_output) + + torch.testing.assert_close(actual, expected.detach()) + for actual_input, reference in zip(actual_inputs, ref_inputs, strict=True): + torch.testing.assert_close(actual_input.grad, reference.grad) + + assert stub.calls.count("det_gemm_fwd") == 0 + assert stub.calls.count("det_gemm_db") == 0 + assert stub.calls.count("swiglu_forward") == 1 + assert stub.calls.count("swiglu_backward") == 1 + + +def test_qwen_ffn_deterministic_false_uses_production_gemm(monkeypatch): + modes = [] + monkeypatch.setattr(ffn_module, "_validate_ffn_inputs", lambda *args: None) + monkeypatch.setattr(ffn_module, "_require_ffn_kernels", lambda **kwargs: modes.append(kwargs)) + monkeypatch.setattr( + ffn_module._DeterministicFFNFunction, + "apply", + lambda *args: args[-1], + ) + + tensors = [torch.empty(1)] * 4 + assert qwen3_ffn(*tensors, deterministic=False) is False + assert modes == [{"disable_split_k": False}] + + +def test_qwen_ffn_rejects_conflicting_backend_switches(): + tensors = [torch.empty(1)] * 4 + + with pytest.raises(ValueError, match="conflicting FFN backends"): + qwen3_ffn(*tensors, deterministic=True, disable_split_k=False) + + +def test_qwen_ffn_rejects_non_bool_deterministic(): + tensors = [torch.empty(1)] * 4 + + with pytest.raises(TypeError, match="deterministic must be a bool or None"): + qwen3_ffn(*tensors, deterministic=1) # type: ignore[arg-type] + + +def test_qwen_ffn_rejects_non_bool_disable_split_k(): + hidden = torch.empty((2, 8), dtype=torch.bfloat16) + gate_weight = torch.empty((12, 8), dtype=torch.bfloat16) + up_weight = torch.empty((12, 8), dtype=torch.bfloat16) + down_weight = torch.empty((8, 12), dtype=torch.bfloat16) + with pytest.raises(TypeError, match="disable_split_k must be a bool"): + qwen3_ffn( + hidden, + gate_weight, + up_weight, + down_weight, + disable_split_k=1, # type: ignore[arg-type] + ) + + +def test_qwen_ffn_rejects_non_huggingface_weight_layout(): + hidden = torch.empty((2, 8), dtype=torch.bfloat16) + gate_weight = torch.empty((8, 12), dtype=torch.bfloat16) + up_weight = torch.empty((12, 8), dtype=torch.bfloat16) + down_weight = torch.empty((8, 12), dtype=torch.bfloat16) + + with pytest.raises(ValueError, match="gate_weight must have shape"): + qwen3_ffn(hidden, gate_weight, up_weight, down_weight) + + +@requires_cuda_ffn +@pytest.mark.parametrize("disable_split_k", [True, False]) +def test_qwen_ffn_cuda_forward_backward_matches_fp32_reference(disable_split_k): + hidden = _randn((2, 3, 64), seed=10, device="cuda", dtype=torch.bfloat16) + gate_weight = _randn((128, 64), seed=11, device="cuda", dtype=torch.bfloat16) + up_weight = _randn((128, 64), seed=12, device="cuda", dtype=torch.bfloat16) + down_weight = _randn((64, 128), seed=13, device="cuda", dtype=torch.bfloat16) + grad_output = _randn(hidden.shape, seed=14, device="cuda", dtype=torch.bfloat16) + + ref_inputs = [ + value.detach().cpu().float().requires_grad_(True) + for value in (hidden, gate_weight, up_weight, down_weight) + ] + expected, _, _, _ = _reference(*ref_inputs) + expected.backward(grad_output.cpu().float()) + + actual_inputs = [ + value.detach().clone().requires_grad_(True) + for value in (hidden, gate_weight, up_weight, down_weight) + ] + actual = qwen3_ffn(*actual_inputs, disable_split_k=disable_split_k) + actual.backward(grad_output) + + torch.testing.assert_close( + actual.cpu().float(), + expected.detach(), + atol=5e-2, + rtol=2e-2, + ) + for actual_input, reference in zip(actual_inputs, ref_inputs, strict=True): + torch.testing.assert_close( + actual_input.grad.cpu().float(), + reference.grad, + atol=5e-2, + rtol=2e-2, + ) + + +@requires_cuda_ffn +def test_qwen_ffn_cuda_forward_and_input_gradient_are_batch_invariant(): + gate_weight = _randn((128, 64), seed=20, device="cuda", dtype=torch.bfloat16) + up_weight = _randn((128, 64), seed=21, device="cuda", dtype=torch.bfloat16) + down_weight = _randn((64, 128), seed=22, device="cuda", dtype=torch.bfloat16) + hidden = _randn((6, 64), seed=23, device="cuda", dtype=torch.bfloat16) + grad_output = _randn(hidden.shape, seed=24, device="cuda", dtype=torch.bfloat16) + + full_hidden = hidden.detach().clone().requires_grad_(True) + full_output = qwen3_ffn(full_hidden, gate_weight, up_weight, down_weight) + full_output.backward(grad_output) + + slice_hidden = hidden[2:4].detach().clone().requires_grad_(True) + slice_output = qwen3_ffn(slice_hidden, gate_weight, up_weight, down_weight) + slice_output.backward(grad_output[2:4]) + + assert torch.equal(slice_output, full_output[2:4]) + assert torch.equal(slice_hidden.grad, full_hidden.grad[2:4]) + + +@requires_cuda_ffn +def test_qwen_ffn_cuda_train_and_infer_forward_are_bitwise_identical(): + hidden = _randn((16, 64), seed=30, device="cuda", dtype=torch.bfloat16) + gate_weight = _randn((128, 64), seed=31, device="cuda", dtype=torch.bfloat16) + up_weight = _randn((128, 64), seed=32, device="cuda", dtype=torch.bfloat16) + down_weight = _randn((64, 128), seed=33, device="cuda", dtype=torch.bfloat16) + with torch.no_grad(): + infer = qwen3_ffn(hidden, gate_weight, up_weight, down_weight) + train_hidden = hidden.detach().clone().requires_grad_(True) + train = qwen3_ffn(train_hidden, gate_weight, up_weight, down_weight) + assert torch.equal(infer, train.detach()) + + +@requires_cuda_ffn +@pytest.mark.parametrize("token_count", _TOKEN_BOUNDARY_COUNTS) +def test_qwen_ffn_output_and_hidden_grad_are_batch_invariant(token_count): + device = torch.device("cuda", 0) + gate = _randn((256, _HIDDEN), seed=70, device=device, dtype=torch.bfloat16) + up = _randn((256, _HIDDEN), seed=71, device=device, dtype=torch.bfloat16) + down = _randn((_HIDDEN, 256), seed=72, device=device, dtype=torch.bfloat16) + hidden = _randn((token_count, _HIDDEN), seed=73, device=device, dtype=torch.bfloat16) + grad = _randn((token_count, _HIDDEN), seed=74, device=device, dtype=torch.bfloat16) + + full_hidden = hidden.detach().clone().requires_grad_(True) + full_output = qwen3_ffn(full_hidden, gate, up, down) + full_output.backward(grad) + slice_end = min(8, token_count) + slice_hidden = hidden[:slice_end].detach().clone().requires_grad_(True) + slice_output = qwen3_ffn(slice_hidden, gate, up, down) + slice_output.backward(grad[:slice_end]) + + assert torch.equal(slice_output, full_output[:slice_end]) + assert torch.equal(slice_hidden.grad, full_hidden.grad[:slice_end]) + + +@requires_cuda_ffn +def test_qwen_ffn_qwen3_8b_shapes_run_and_are_batch_invariant(): + device = torch.device("cuda", 0) + hidden, gate, up, down, grad = _qwen3_8b_weights(device) + + with torch.no_grad(): + infer = qwen3_ffn(hidden, gate, up, down) + full_hidden = hidden.detach().clone().requires_grad_(True) + full_gate = gate.detach().clone().requires_grad_(True) + full_up = up.detach().clone().requires_grad_(True) + full_down = down.detach().clone().requires_grad_(True) + train = qwen3_ffn(full_hidden, full_gate, full_up, full_down) + train.backward(grad) + assert torch.equal(infer, train.detach()) + + slice_hidden = hidden[:4].detach().clone().requires_grad_(True) + slice_out = qwen3_ffn(slice_hidden, full_gate, full_up, full_down) + slice_out.backward(grad[:4]) + assert torch.equal(slice_out, train.detach()[:4]) + assert torch.equal(slice_hidden.grad, full_hidden.grad[:4]) + + +@requires_cuda_ffn +def test_qwen_ffn_sequence_parallel_requires_tensor_parallel_group(): + device = torch.device("cuda", 0) + hidden, gate, up, down = _ffn_tensors(8, device, seed=120, intermediate=128) + with pytest.raises(ValueError, match="sequence_parallel requires a tensor-parallel group"): + qwen3_ffn(hidden, gate, up, down, sequence_parallel=True) + + +def test_qwen_ffn_tp_correctness_and_batch_invariance(): + _spawn_nccl_workers(_distributed_ffn_backward_nccl_worker, 2, (1, False), timeout=90) + + +def test_qwen_ffn_tp_sp_correctness_and_batch_invariance(): + _spawn_nccl_workers(_distributed_ffn_backward_nccl_worker, 2, (1, True), timeout=90) + + +def test_qwen_ffn_tp_cp_correctness_and_batch_invariance(): + _spawn_nccl_workers(_distributed_ffn_backward_nccl_worker, 4, (2, False), timeout=90) + + +def test_qwen_ffn_tp_cp_sp_correctness_and_batch_invariance(): + _spawn_nccl_workers(_distributed_ffn_backward_nccl_worker, 4, (2, True), timeout=90) + + +def test_qwen_ffn_tp1_vs_tp2_train_infer_bitwise_identical(): + _spawn_nccl_workers(_tp1_vs_tpn_train_infer_worker, 2, (True,), timeout=120) + + +def test_qwen_ffn_tp1_vs_tp8_train_infer_bitwise_identical(): + _spawn_nccl_workers(_tp1_vs_tpn_train_infer_worker, 8, (True,), timeout=120) + + +def test_qwen_ffn_world2_tp_sp_and_cp_match_tp1_cp1_bitwise(): + _spawn_nccl_workers(_topology_worker, 2, (_WORLD2_CONFIGS,), timeout=120) + + +def test_qwen_ffn_world4_tp_cp_sp_match_tp1_cp1_bitwise(): + _spawn_nccl_workers(_topology_worker, 4, (_WORLD4_CONFIGS,), timeout=120) + + +def test_qwen_ffn_world8_tp8_and_cp8_match_tp1_cp1_bitwise(): + _spawn_nccl_workers(_topology_worker, 8, (_WORLD8_WORLD_GROUP_CONFIGS,), timeout=120) + + +def test_qwen_ffn_world8_tp2_cp4_match_tp1_cp1_bitwise(): + _spawn_nccl_workers(_topology_worker, 8, (_WORLD8_TP2_CP4_CONFIGS,), timeout=120) + + +def test_qwen_ffn_world8_tp4_cp2_match_tp1_cp1_bitwise(): + _spawn_nccl_workers(_topology_worker, 8, (_WORLD8_TP4_CP2_CONFIGS,), timeout=120) + + +def test_qwen_ffn_collective_cache_reuses_grows_closes_and_rebuilds_group(): + _spawn_nccl_workers(_cache_worker, 2, timeout=120) + + +def test_qwen_ffn_sequence_parallel_rejects_uneven_tokens(): + _spawn_nccl_workers(_uneven_sp_worker, 2, timeout=90) + + +def test_qwen_ffn_qwen3_8b_shapes_tp2_matches_tp1_bitwise(): + _spawn_nccl_workers(_qwen3_8b_tp2_worker, 2, timeout=180) diff --git a/tests/test_stateless_executor.py b/tests/test_stateless_executor.py index fc6d9b3a..851402f7 100644 --- a/tests/test_stateless_executor.py +++ b/tests/test_stateless_executor.py @@ -172,6 +172,90 @@ def test_executor_runs_full_sequence_forward_with_use_cache_false_and_detaches_o assert not hasattr(model.generation_config, "attn_implementation") +def test_executor_exact_selected_logprob_callable_is_optional_and_injected(): + inputs = _inputs() + logits = _logits_for(inputs) + calls = [] + + def selected_logprob_fn( + shifted_logits, + shifted_labels, + *, + mask, + temperature, + output_dtype, + ): + calls.append((shifted_logits, shifted_labels, mask, temperature, output_dtype)) + reference = selected_logprobs_reference( + shifted_logits, + shifted_labels, + mask=mask, + temperature=temperature, + output_dtype=output_dtype, + ) + return reference + mask.to(dtype=output_dtype) * 0.25 + + default = StatelessForwardExecutor( + FakeReferenceModel(logits), + StatelessForwardConfig(mode="reference"), + ).score(inputs) + injected = StatelessForwardExecutor( + FakeReferenceModel(logits), + StatelessForwardConfig(mode="reference"), + selected_logprob_fn=selected_logprob_fn, + ).score(inputs) + + assert default.reference_logps is not None + assert injected.reference_logps is not None + assert len(calls) == 1 + assert torch.equal(calls[0][2], inputs.completion_mask[:, 1:]) + torch.testing.assert_close( + injected.reference_logps[inputs.completion_mask], + default.reference_logps[inputs.completion_mask] + 0.25, + ) + assert torch.equal( + injected.reference_logps[~inputs.completion_mask], + torch.zeros_like(injected.reference_logps[~inputs.completion_mask]), + ) + + +def test_executor_uses_eval_no_grad_and_restores_mixed_module_modes_read_only(): + inputs = _inputs() + + class ReadOnlyModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.weight = torch.nn.Parameter(torch.tensor(2.0)) + self.register_buffer("counter", torch.tensor(3.0)) + self.child = torch.nn.Linear(1, 1) + + def forward(self, input_ids, attention_mask=None, use_cache=None): + del attention_mask + assert self.training is False + assert self.child.training is False + assert torch.is_grad_enabled() is False + assert use_cache is False + batch, sequence = input_ids.shape + logits = torch.zeros(batch, sequence, 8) + return {"logits": logits + self.weight * 0.0 + self.counter * 0.0} + + model = ReadOnlyModel() + model.train() + model.child.eval() + state_before = {name: value.detach().clone() for name, value in model.state_dict().items()} + + result = StatelessForwardExecutor( + model, + StatelessForwardConfig(mode="reference", attention_backend="eager"), + ).score(inputs) + + assert result.reference_logps is not None + assert result.metrics["model_eval_during_forward"] is True + assert model.training is True + assert model.child.training is False + assert all(torch.equal(state_before[name], value) for name, value in model.state_dict().items()) + + def test_executor_falls_back_for_models_without_use_cache_argument(): inputs = _inputs() executor = StatelessForwardExecutor( diff --git a/tests/test_vime_logprob_provider.py b/tests/test_vime_logprob_provider.py new file mode 100644 index 00000000..021f4e9f --- /dev/null +++ b/tests/test_vime_logprob_provider.py @@ -0,0 +1,93 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""CPU coverage for the optional Vime WS2 selected-logprob adapter.""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch + +from rl_engine.integrations.vime.logp import SelectedLogprobProviderUnavailable, provider + + +def _request(*, cp_rank: int = 0, with_entropy: bool = False, keep_mask=None): + logits = torch.tensor( + [[0.25, -0.5, 1.0, 0.1, -0.3, 0.6, -0.7, 0.4] for _ in range(3)], + dtype=torch.float32, + requires_grad=True, + ) + return SimpleNamespace( + logits=logits, + target_ids=torch.tensor([2, 5, 0]), + tensor_parallel_group=None, + context_parallel=SimpleNamespace( + world_size=2, + rank=cp_rank, + layout="zigzag", + ), + with_entropy=with_entropy, + with_entropy_grad=with_entropy, + log_prob_keep_mask=keep_mask, + metadata={ + "real_vocab_size": 7, + "padded_vocab_size": 8, + "tp_rank": 0, + "tp_world_size": 1, + "num_vocab_tiles": 4, + }, + ) + + +def test_provider_runs_locally_with_cp2_row_metadata(): + request = _request(cp_rank=1) + + result = provider(request) + reference = torch.log_softmax(request.logits[:, :7], dim=-1)[ + torch.arange(request.logits.size(0)), request.target_ids + ] + + assert result.selected_logprobs.shape == (3, 1) + torch.testing.assert_close(result.selected_logprobs.squeeze(-1), reference) + assert result.backend_id == "pytorch-vocab-parallel-logp-ws2" + assert result.provenance["cp_row_ownership"] == { + "cp_rank": 1, + "cp_world_size": 2, + "layout": "zigzag", + "local_token_rows": 3, + "cp_is_merge_axis": False, + } + + +def test_provider_entropy_preserves_vime_semantics_and_autograd(): + request = _request(with_entropy=True) + + result = provider(request) + reference_logits = request.logits.detach().clone().requires_grad_(True) + log_probs = torch.log_softmax(reference_logits[:, :7], dim=-1) + reference_logp = log_probs[torch.arange(reference_logits.size(0)), request.target_ids] + reference_entropy = -(log_probs.exp() * log_probs).sum(dim=-1) + + torch.testing.assert_close(result.selected_logprobs.squeeze(-1), reference_logp) + torch.testing.assert_close(result.entropy, reference_entropy) + (result.selected_logprobs.sum() + result.entropy.sum()).backward() + (reference_logp.sum() + reference_entropy.sum()).backward() + torch.testing.assert_close(request.logits.grad[:, :7], reference_logits.grad[:, :7]) + assert bool((request.logits.grad[:, 7] == 0).all()) + + +def test_provider_rejects_top_p_replay_without_changing_its_semantics(): + request = _request(keep_mask=torch.ones((3, 8), dtype=torch.bool)) + + with pytest.raises(SelectedLogprobProviderUnavailable, match="top-p replay"): + provider(request) + + +def test_provider_rejects_local_vocab_metadata_that_cannot_describe_tp_ownership(): + request = _request() + request.metadata["padded_vocab_size"] = 16 + + with pytest.raises(SelectedLogprobProviderUnavailable, match="cover padded_vocab_size"): + provider(request) diff --git a/tests/test_vime_qwen3_example.py b/tests/test_vime_qwen3_example.py new file mode 100644 index 00000000..91d9e2a2 --- /dev/null +++ b/tests/test_vime_qwen3_example.py @@ -0,0 +1,140 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from pathlib import Path + +import pytest + +from examples.vime_qwen3_8b_tp2_cp2.run import ( + build_report, + load_config, + validate_config, + validate_runtime_evidence, +) + +ROOT = Path(__file__).parents[1] +CONFIG = ROOT / "examples" / "vime_qwen3_8b_tp2_cp2" / "qwen3_8b_tp2_cp2.json" + + +def test_qwen3_example_config_is_strict_and_explicit(): + config = load_config(CONFIG) + validate_config(config) + assert config["training"]["tensor_model_parallel_size"] == 2 + assert config["training"]["context_parallel_size"] == 2 + assert config["selected_logprob_provider"]["mode"] == "strict" + + +def test_qwen3_example_report_does_not_claim_unread_back_attention_or_ffn(tmp_path): + config = load_config(CONFIG) + report = build_report( + config, + vime_root=tmp_path / "vime", + rl_kernel_root=tmp_path / "rl-kernel", + command=["bash", "run.sh"], + status="passed", + returncode=0, + log_text="Selected-logprob provider active: backend_id=pytorch-vocab-parallel-logp-ws2", + log_path=tmp_path / "run.log", + ) + assert report["status"] == "passed" + assert report["claim_boundary"]["qwen3_8b_tp2_cp2_vime_training"] is True + assert report["claim_boundary"]["attention_train_infer_consistency"] == "unclaimed" + assert report["claim_boundary"]["ffn_train_infer_consistency"] == "unclaimed" + assert report["provider"]["fallback_observed"] is False + + +def test_qwen3_example_fails_closed_when_provider_marker_is_missing(tmp_path): + config = load_config(CONFIG) + report = build_report( + config, + vime_root=tmp_path / "vime", + rl_kernel_root=tmp_path / "rl-kernel", + command=["bash", "run.sh"], + status="passed", + returncode=0, + log_text="training completed without provider provenance", + log_path=None, + ) + assert report["status"] == "failed" + assert report["claim_boundary"]["qwen3_8b_tp2_cp2_vime_training"] is False + + +def _runtime_evidence(): + return { + "schema_version": "rlkernel.operator_runtime_evidence.v1", + "operators": { + "attention": { + "training": { + "implementation_id": "rlk.attn", + "backend_id": "rlk", + "contract_id": "a", + }, + "rollout": { + "implementation_id": "rlk.attn", + "backend_id": "rlk", + "contract_id": "a", + }, + "comparison": { + "passed": True, + "out_max_abs": 0.0, + "lse_max_abs": 0.0, + "dq_max_abs": 0.0, + "dk_max_abs": 0.0, + "dv_max_abs": 0.0, + }, + }, + "ffn": { + "training": { + "implementation_id": "rlk.ffn", + "backend_id": "rlk", + "contract_id": "f", + }, + "rollout": { + "implementation_id": "rlk.ffn", + "backend_id": "rlk", + "contract_id": "f", + }, + "comparison": { + "passed": True, + "out_max_abs": 0.0, + "dx_max_abs": 0.0, + "dw_max_abs": 0.0, + }, + }, + }, + } + + +def test_qwen3_example_accepts_only_exact_zero_runtime_evidence(tmp_path): + evidence = _runtime_evidence() + validate_runtime_evidence(evidence) + config = load_config(CONFIG) + report = build_report( + config, + vime_root=tmp_path / "vime", + rl_kernel_root=tmp_path / "rl-kernel", + command=["bash", "run.sh"], + status="passed", + returncode=0, + log_text="Selected-logprob provider active: backend_id=pytorch-vocab-parallel-logp-ws2", + log_path=None, + runtime_evidence=evidence, + ) + assert report["claim_boundary"]["attention_train_infer_consistency"] == "passed" + assert report["claim_boundary"]["ffn_train_infer_consistency"] == "passed" + + +def test_qwen3_example_rejects_nonzero_runtime_evidence(): + evidence = _runtime_evidence() + evidence["operators"]["attention"]["comparison"]["out_max_abs"] = 1e-6 + with pytest.raises(ValueError, match="attention"): + validate_runtime_evidence(evidence) + + +@pytest.mark.parametrize("bad_path", ["", "other.provider"]) +def test_qwen3_example_rejects_non_rlkernel_provider(bad_path): + config = load_config(CONFIG) + config["selected_logprob_provider"]["path"] = bad_path + with pytest.raises(ValueError, match="RL-Kernel Vime provider"): + validate_config(config) diff --git a/tests/test_vocab_parallel_logp.py b/tests/test_vocab_parallel_logp.py new file mode 100644 index 00000000..eefc1d6a --- /dev/null +++ b/tests/test_vocab_parallel_logp.py @@ -0,0 +1,581 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Deterministic vocab-parallel TP logprob reference tests (issue #241 PR3). + +Bit-level determinism assertions compare raw bit patterns via +``tensor.view(torch.int32)`` rather than ``torch.equal``: value equality +treats ``-0.0 == 0.0`` as equal and ``NaN != NaN`` as different, neither of +which is what a bitwise claim means. +""" + +from __future__ import annotations + +import queue +import tempfile +import traceback +from pathlib import Path + +import pytest +import torch +import torch.multiprocessing as mp + +from rl_engine.kernels.gtest.tolerance import load_contract +from rl_engine.kernels.logprob_contract import ( + DeterminismScope, + LogprobContract, + LogprobContractError, + MaskSpec, + ReductionSpec, + ShardingSpec, +) +from rl_engine.kernels.ops.pytorch.loss.batch_invariant_logp import NativeBatchInvariantLogpOp +from rl_engine.kernels.ops.pytorch.loss.vocab_parallel_logp import ( + BACKEND_ID, + VocabParallelLogprobOp, +) +from rl_engine.kernels.registry import KernelRegistry, OpBackend + +REAL_VOCAB = 27 +PADDED_VOCAB = 32 +NUM_TILES = 8 +NUM_TOKENS = 6 +ACTIVE = (True, True, True, True, True, False) + + +def _even_bounds(padded: int, world: int) -> tuple[tuple[int, int], ...]: + shard = padded // world + return tuple( + (rank * shard, padded if rank == world - 1 else (rank + 1) * shard) for rank in range(world) + ) + + +def _contract( + *, + tp_rank: int = 0, + tp_world_size: int = 1, + bounds: tuple[tuple[int, int], ...] | None = None, + real_vocab: int = REAL_VOCAB, + padded_vocab: int = PADDED_VOCAB, + num_tokens: int = NUM_TOKENS, + active: tuple[bool, ...] = ACTIVE, + dtype: str = "fp32", + determinism_scope: str = "cross_tp_bitwise", +) -> LogprobContract: + return LogprobContract( + role="train", + dtype=dtype, + mask=MaskSpec(num_tokens=num_tokens, active_mask=active), + sharding=ShardingSpec( + tp_rank=tp_rank, + tp_world_size=tp_world_size, + vocab_shard_bounds=( + bounds if bounds is not None else _even_bounds(padded_vocab, tp_world_size) + ), + real_vocab_size=real_vocab, + padded_vocab_size=padded_vocab, + ), + reduction=ReductionSpec(determinism_scope=determinism_scope), + ) + + +def _inputs(dtype=torch.float32, seed: int = 2026): + torch.manual_seed(seed) + logits = torch.randn(NUM_TOKENS, PADDED_VOCAB, dtype=torch.float32).to(dtype) + targets = torch.tensor([1, 5, REAL_VOCAB - 1, 0, 13, -100]) + return logits, targets + + +def _bits(tensor: torch.Tensor) -> torch.Tensor: + view_dtype = {torch.float32: torch.int32, torch.bfloat16: torch.int16}[tensor.dtype] + return tensor.contiguous().view(view_dtype) + + +def _bitwise_equal(a: torch.Tensor, b: torch.Tensor) -> bool: + return a.shape == b.shape and bool((_bits(a) == _bits(b)).all()) + + +def _case_shard_size_mismatch(): + logits, targets = _inputs() + return logits, targets, _contract(tp_rank=0, tp_world_size=2), NUM_TILES, "vocab columns" + + +def _case_mask_length_mismatch(): + logits, targets = _inputs() + contract = _contract(num_tokens=NUM_TOKENS + 1, active=ACTIVE + (True,)) + return logits, targets, contract, NUM_TILES, "num_tokens" + + +def _case_dtype_mismatch(): + logits, targets = _inputs() + return logits, targets, _contract(dtype="bf16"), NUM_TILES, "dtype" + + +def _case_tile_misaligned_bounds(): + # Tile size is 32/8 = 4; a boundary at 6 is misaligned. + logits, targets = _inputs() + contract = _contract(tp_world_size=2, bounds=((0, 6), (6, 32))) + return logits[:, :6], targets, contract, NUM_TILES, "tile" + + +def _case_bad_num_vocab_tiles(): + logits, targets = _inputs() + return logits, targets, _contract(), 7, "num_vocab_tiles" + + +def _case_active_target_out_of_real_vocab(): + logits, targets = _inputs() + bad_targets = targets.clone() + bad_targets[0] = REAL_VOCAB # padding column, active row + return logits, bad_targets, _contract(), NUM_TILES, "real vocabulary" + + +def _case_all_inf_active_row(): + logits, targets = _inputs() + poisoned = logits.clone() + poisoned[0, :] = float("-inf") + return poisoned, targets, _contract(), NUM_TILES, "non-finite" + + +@pytest.mark.parametrize( + "case", + [ + _case_shard_size_mismatch, + _case_mask_length_mismatch, + _case_dtype_mismatch, + _case_tile_misaligned_bounds, + _case_bad_num_vocab_tiles, + _case_active_target_out_of_real_vocab, + _case_all_inf_active_row, + ], + ids=lambda fn: fn.__name__.removeprefix("_case_"), +) +def test_invalid_invocations_fail_loudly(case): + logits, targets, contract, num_tiles, match = case() + with pytest.raises(LogprobContractError, match=match): + VocabParallelLogprobOp()(logits, targets, contract=contract, num_vocab_tiles=num_tiles) + + +class TestSingleRank: + def test_repeated_runs_are_bitwise_identical(self): + contract = _contract() + logits, targets = _inputs() + op = VocabParallelLogprobOp() + logp_a, lse_a = op(logits, targets, contract=contract, num_vocab_tiles=NUM_TILES) + logp_b, lse_b = op(logits, targets, contract=contract, num_vocab_tiles=NUM_TILES) + assert _bitwise_equal(logp_a, logp_b) + assert _bitwise_equal(lse_a, lse_b) + + def test_batch_invariance_same_row_any_context(self): + contract_full = _contract() + logits, targets = _inputs() + op = VocabParallelLogprobOp() + logp_full, lse_full = op(logits, targets, contract=contract_full, num_vocab_tiles=NUM_TILES) + + contract_single = _contract(num_tokens=1, active=(True,)) + logp_one, lse_one = op( + logits[2:3], targets[2:3], contract=contract_single, num_vocab_tiles=NUM_TILES + ) + assert _bitwise_equal(logp_full[2:3], logp_one) + assert _bitwise_equal(lse_full[2:3], lse_one) + + def test_matches_ws1_batch_invariant_logp_within_contract_tolerance(self): + tolerance = load_contract()["accuracy"]["default"]["logprob"]["float32"] + contract = _contract(padded_vocab=REAL_VOCAB + 5) + # Use a real==padded contract so the WS1 op sees identical logits. + contract = _contract(real_vocab=PADDED_VOCAB, padded_vocab=PADDED_VOCAB) + logits, targets = _inputs() + logp, _ = VocabParallelLogprobOp()( + logits, targets, contract=contract, num_vocab_tiles=NUM_TILES + ) + ws1 = NativeBatchInvariantLogpOp().apply(logits, targets) + active = torch.tensor(ACTIVE) + assert torch.allclose( + logp[active], ws1[active], atol=tolerance["atol"], rtol=tolerance["rtol"] + ) + + def test_padding_columns_are_excluded_and_finite(self): + contract = _contract() + logits, targets = _inputs() + boosted = logits.clone() + boosted[:, REAL_VOCAB:] = 1e4 # huge padding logits must not leak into LSE + logp, lse = VocabParallelLogprobOp()( + boosted, targets, contract=contract, num_vocab_tiles=NUM_TILES + ) + ref_lse = torch.logsumexp(boosted[:, :REAL_VOCAB].float(), dim=-1) + assert torch.isfinite(logp).all() and torch.isfinite(lse).all() + assert torch.allclose(lse, ref_lse, atol=1e-5) + + def test_inactive_rows_zero_filled_lse_still_exported(self): + contract = _contract() + logits, targets = _inputs() + logp, lse = VocabParallelLogprobOp()( + logits, targets, contract=contract, num_vocab_tiles=NUM_TILES + ) + assert logp[-1].item() == 0.0 + assert torch.isfinite(lse[-1]) + + +class TestNonDeterministicPath: + """deterministic=False: the fast whole-shard reduction with no guarantee.""" + + def test_rejects_a_cross_tp_bitwise_contract(self): + logits, targets = _inputs() + with pytest.raises(LogprobContractError, match="cross_tp_bitwise"): + VocabParallelLogprobOp()(logits, targets, contract=_contract(), deterministic=False) + + def test_rejects_a_non_bool_flag(self): + logits, targets = _inputs() + with pytest.raises(LogprobContractError, match="deterministic must be a bool"): + VocabParallelLogprobOp()(logits, targets, contract=_contract(), deterministic=1) + + def test_matches_the_deterministic_path_within_tolerance(self): + tolerance = load_contract()["accuracy"]["default"]["logprob"]["float32"] + logits, targets = _inputs() + relaxed = _contract(determinism_scope="fixed_topology") + logp_fast, lse_fast = VocabParallelLogprobOp()( + logits, targets, contract=relaxed, deterministic=False + ) + logp_det, lse_det = VocabParallelLogprobOp()( + logits, targets, contract=_contract(), num_vocab_tiles=NUM_TILES + ) + assert torch.allclose(logp_fast, logp_det, atol=tolerance["atol"], rtol=tolerance["rtol"]) + assert torch.allclose(lse_fast, lse_det, atol=tolerance["atol"], rtol=tolerance["rtol"]) + + def test_ignores_tile_constraints(self): + # 7 does not divide the padded vocab; the deterministic path rejects it, + # the fast path never looks at it. + logits, targets = _inputs() + relaxed = _contract(determinism_scope="fixed_topology") + logp, lse = VocabParallelLogprobOp()( + logits, targets, contract=relaxed, num_vocab_tiles=7, deterministic=False + ) + ref_lse = torch.logsumexp(logits[:, :REAL_VOCAB].float(), dim=-1) + assert torch.allclose(lse, ref_lse, atol=1e-5) + assert torch.isfinite(logp).all() + + def test_grads_match_autograd_oracle(self): + tolerance = load_contract()["accuracy"]["default"]["logprob"]["float32"] + relaxed = _contract(determinism_scope="fixed_topology") + logits, targets = _inputs() + x = logits.clone().requires_grad_(True) + logp, lse = VocabParallelLogprobOp()(x, targets, contract=relaxed, deterministic=False) + (logp.sum() + 0.5 * lse.sum()).backward() + + y = logits.clone().requires_grad_(True) + ref_lse = torch.logsumexp(y[:, :REAL_VOCAB].float(), dim=-1) + safe = targets.clamp(0, REAL_VOCAB - 1) + ref_logp = y[torch.arange(NUM_TOKENS), safe].float() - ref_lse + ref_logp = torch.where(torch.tensor(ACTIVE), ref_logp, torch.zeros_like(ref_logp)) + (ref_logp.sum() + 0.5 * ref_lse.sum()).backward() + + assert torch.allclose(x.grad, y.grad, atol=tolerance["atol"], rtol=tolerance["rtol"]) + assert bool((x.grad[:, REAL_VOCAB:] == 0).all()) + + +class TestBackward: + def test_grads_match_autograd_oracle(self): + tolerance = load_contract()["accuracy"]["default"]["logprob"]["float32"] + contract = _contract() + logits, targets = _inputs() + x = logits.clone().requires_grad_(True) + logp, lse = VocabParallelLogprobOp()( + x, targets, contract=contract, num_vocab_tiles=NUM_TILES + ) + (logp.sum() + 0.5 * lse.sum()).backward() + + y = logits.clone().requires_grad_(True) + ref_lse = torch.logsumexp(y[:, :REAL_VOCAB].float(), dim=-1) + safe = targets.clamp(0, REAL_VOCAB - 1) + ref_logp = y[torch.arange(NUM_TOKENS), safe].float() - ref_lse + ref_logp = torch.where(torch.tensor(ACTIVE), ref_logp, torch.zeros_like(ref_logp)) + (ref_logp.sum() + 0.5 * ref_lse.sum()).backward() + + assert torch.allclose(x.grad, y.grad, atol=tolerance["atol"], rtol=tolerance["rtol"]) + assert bool((x.grad[:, REAL_VOCAB:] == 0).all()) + + # No grad requested -> outputs detached from autograd entirely. + logp_ng, lse_ng = VocabParallelLogprobOp()( + logits, targets, contract=contract, num_vocab_tiles=NUM_TILES + ) + assert not logp_ng.requires_grad and not lse_ng.requires_grad + + def test_inactive_rows_grad_asymmetry(self): + """The logp term is zeroed on inactive rows; the lse term still flows — + lse is a row property exported (and differentiable) for every row.""" + + contract = _contract() + logits, targets = _inputs() + + x = logits.clone().requires_grad_(True) + _, lse = VocabParallelLogprobOp()(x, targets, contract=contract, num_vocab_tiles=NUM_TILES) + lse.sum().backward() + assert bool((x.grad[-1, :REAL_VOCAB].abs() > 0).any()) + + z = logits.clone().requires_grad_(True) + logp, _ = VocabParallelLogprobOp()(z, targets, contract=contract, num_vocab_tiles=NUM_TILES) + logp.sum().backward() + assert bool((z.grad[-1] == 0).all()) + + def test_entropy_matches_full_vocab_oracle_and_backpropagates(self): + contract = _contract() + logits, targets = _inputs() + x = logits.clone().requires_grad_(True) + logp, _lse, entropy = VocabParallelLogprobOp().apply_with_entropy( + x, + targets, + contract=contract, + num_vocab_tiles=NUM_TILES, + ) + + y = logits.clone().requires_grad_(True) + ref_log_probs = torch.log_softmax(y[:, :REAL_VOCAB].float(), dim=-1) + safe = targets.clamp(0, REAL_VOCAB - 1) + ref_logp = ref_log_probs[torch.arange(NUM_TOKENS), safe] + ref_logp = torch.where(torch.tensor(ACTIVE), ref_logp, torch.zeros_like(ref_logp)) + ref_entropy = -(ref_log_probs.exp() * ref_log_probs).sum(dim=-1) + + torch.testing.assert_close(entropy, ref_entropy) + (logp.sum() + entropy.sum()).backward() + (ref_logp.sum() + ref_entropy.sum()).backward() + torch.testing.assert_close(x.grad, y.grad, atol=2e-5, rtol=2e-5) + + +def test_dispatch_resolves_reference_and_leaves_legacy_untouched(): + registry = KernelRegistry() + contract = _contract() + + result = registry.get_logprob_op(contract) + assert result.capability.backend_id == BACKEND_ID + assert result.provenance["fallback"] is False + assert isinstance(result.op, VocabParallelLogprobOp) + assert ( + result.provenance["contract"]["reduction"]["determinism_scope"] + == DeterminismScope.CROSS_TP_BITWISE.value + ) + + by_id = registry.get_logprob_op(contract, requested_backend=BACKEND_ID) + assert by_id.capability.backend_id == BACKEND_ID + by_kind = registry.get_logprob_op(contract, requested_backend="reference") + assert by_kind.capability.backend_id == BACKEND_ID + + for ops in registry._priority_map.values(): + for candidates in ops.values(): + assert OpBackend.PYTORCH_VOCAB_PARALLEL_LOGP not in candidates + + +# --------------------------------------------------------------------------- +# Multi-rank gloo tests (spawn pattern from tests/test_linear_logp.py) +# --------------------------------------------------------------------------- + + +def _gloo_available() -> bool: + return torch.distributed.is_available() and torch.distributed.is_gloo_available() + + +requires_gloo = pytest.mark.skipif( + not _gloo_available(), reason="requires torch.distributed with the gloo backend" +) + +_WORLD_SIZE = 4 +_UNEVEN_BOUNDS = ((0, 4), (4, 16), (16, 24), (24, 32)) # tile-aligned (tile=4) + + +def _tp_worker(rank, world_size, init_method, result_queue, scenario): + import torch.distributed as dist + + torch.set_num_threads(1) + try: + dist.init_process_group( + backend="gloo", init_method=init_method, rank=rank, world_size=world_size + ) + dtype = torch.bfloat16 if scenario == "bf16" else torch.float32 + dtype_name = "bf16" if scenario == "bf16" else "fp32" + bounds = _UNEVEN_BOUNDS if scenario == "uneven" else _even_bounds(PADDED_VOCAB, world_size) + logits, targets = _inputs(dtype=dtype) + start, end = bounds[rank] + + op = VocabParallelLogprobOp() + tiles = 16 if scenario == "preflight" and rank == 0 else NUM_TILES + contract_tp = _contract( + tp_rank=rank, tp_world_size=world_size, bounds=bounds, dtype=dtype_name + ) + + if scenario == "preflight": + try: + op( + logits[:, start:end].contiguous().clone(), + targets, + contract=contract_tp, + tp_group=dist.group.WORLD, + num_vocab_tiles=tiles, + ) + result_queue.put({"ok": False, "rank": rank, "traceback": "no error raised"}) + except LogprobContractError: + result_queue.put({"ok": True, "rank": rank}) + return + + if scenario == "mode_mismatch": + # Same relaxed contract everywhere; only rank 0 runs deterministic. + relaxed = _contract( + tp_rank=rank, + tp_world_size=world_size, + bounds=bounds, + determinism_scope="fixed_topology", + ) + try: + op( + logits[:, start:end].contiguous().clone(), + targets, + contract=relaxed, + tp_group=dist.group.WORLD, + num_vocab_tiles=NUM_TILES, + deterministic=(rank == 0), + ) + result_queue.put({"ok": False, "rank": rank, "traceback": "no error raised"}) + except LogprobContractError: + result_queue.put({"ok": True, "rank": rank}) + return + + if scenario == "nondeterministic": + relaxed = _contract( + tp_rank=rank, + tp_world_size=world_size, + bounds=bounds, + determinism_scope="fixed_topology", + ) + shard = logits[:, start:end].contiguous().clone().requires_grad_(True) + logp_tp, lse_tp = op( + shard, targets, contract=relaxed, tp_group=dist.group.WORLD, deterministic=False + ) + (logp_tp.sum() + 0.5 * lse_tp.sum()).backward() + + full = logits.clone().requires_grad_(True) + relaxed_tp1 = _contract(determinism_scope="fixed_topology") + logp_one, lse_one = op(full, targets, contract=relaxed_tp1, deterministic=False) + (logp_one.sum() + 0.5 * lse_one.sum()).backward() + + result_queue.put( + { + "ok": True, + "rank": rank, + "logp_close": torch.allclose(logp_tp, logp_one, atol=1e-6), + "lse_close": torch.allclose(lse_tp, lse_one, atol=1e-6), + "grad_close": torch.allclose(shard.grad, full.grad[:, start:end], atol=1e-6), + "logp": logp_tp.detach().float(), + "lse": lse_tp.detach().float(), + } + ) + return + + shard = logits[:, start:end].contiguous().clone().requires_grad_(True) + logp_tp, lse_tp = op( + shard, + targets, + contract=contract_tp, + tp_group=dist.group.WORLD, + num_vocab_tiles=NUM_TILES, + ) + (logp_tp.sum() + 0.5 * lse_tp.sum()).backward() + + # In-process TP=1 run of the same op on the full logits: the cross-TP + # bitwise claim is TP=n output == TP=1 output, bit for bit. + full = logits.clone().requires_grad_(True) + contract_tp1 = _contract(dtype=dtype_name) + logp_one, lse_one = op(full, targets, contract=contract_tp1, num_vocab_tiles=NUM_TILES) + (logp_one.sum() + 0.5 * lse_one.sum()).backward() + + result_queue.put( + { + "ok": True, + "rank": rank, + "logp_bits_match": _bitwise_equal(logp_tp, logp_one), + "lse_bits_match": _bitwise_equal(lse_tp, lse_one), + "grad_bits_match": _bitwise_equal(shard.grad, full.grad[:, start:end]), + "logp": logp_tp.detach().float(), + "lse": lse_tp.detach().float(), + } + ) + except Exception: + result_queue.put({"ok": False, "rank": rank, "traceback": traceback.format_exc()}) + raise + finally: + if torch.distributed.is_initialized(): + torch.distributed.destroy_process_group() + + +def _run_gloo_scenario(scenario): + ctx = mp.get_context("spawn") + with tempfile.TemporaryDirectory() as tmpdir: + init_method = (Path(tmpdir) / "gloo_init").as_uri() + result_queue = ctx.Queue() + processes = [ + ctx.Process( + target=_tp_worker, + args=(rank, _WORLD_SIZE, init_method, result_queue, scenario), + ) + for rank in range(_WORLD_SIZE) + ] + results = [] + try: + for process in processes: + process.start() + for _ in range(_WORLD_SIZE): + try: + results.append(result_queue.get(timeout=60)) + except queue.Empty: + for process in processes: + if process.is_alive(): + process.terminate() + pytest.fail("timed out waiting for vocab-parallel gloo workers") + finally: + for process in processes: + process.join(timeout=10) + if process.is_alive(): + process.terminate() + results.sort(key=lambda item: item["rank"]) + for result in results: + assert result["ok"], result.get("traceback") + for process in processes: + assert process.exitcode == 0 + return results + + +@requires_gloo +@pytest.mark.parametrize("scenario", ["even", "uneven", "bf16"]) +def test_tp4_bitwise_identical_to_tp1(scenario): + results = _run_gloo_scenario(scenario) + for result in results: + assert result["logp_bits_match"], f"rank {result['rank']} logp bits differ from TP=1" + assert result["lse_bits_match"], f"rank {result['rank']} lse bits differ from TP=1" + assert result["grad_bits_match"], f"rank {result['rank']} grad bits differ from TP=1" + # Outputs are replicated: every rank must hold identical bits. + for other in results[1:]: + assert _bitwise_equal(results[0]["logp"], other["logp"]) + assert _bitwise_equal(results[0]["lse"], other["lse"]) + + +@requires_gloo +def test_preflight_rejects_mismatched_num_vocab_tiles(): + _run_gloo_scenario("preflight") + + +@requires_gloo +def test_preflight_rejects_mixed_deterministic_modes(): + _run_gloo_scenario("mode_mismatch") + + +@requires_gloo +def test_tp4_nondeterministic_matches_tp1_within_tolerance(): + """No bitwise claim: the fast path's grouping changes with the TP degree. + The values must still agree within fp32 tolerance, and the outputs stay + replicated bitwise across ranks — every rank merges the same partials.""" + + results = _run_gloo_scenario("nondeterministic") + for result in results: + assert result["logp_close"], f"rank {result['rank']} logp differs from TP=1 beyond atol" + assert result["lse_close"], f"rank {result['rank']} lse differs from TP=1 beyond atol" + assert result["grad_close"], f"rank {result['rank']} grads differ from TP=1 beyond atol" + for other in results[1:]: + assert _bitwise_equal(results[0]["logp"], other["logp"]) + assert _bitwise_equal(results[0]["lse"], other["lse"])