diff --git a/csrc/cuda/fused_logp_sm90.cu b/csrc/cuda/fused_logp_sm90.cu index ecf9c799..2341cc97 100644 --- a/csrc/cuda/fused_logp_sm90.cu +++ b/csrc/cuda/fused_logp_sm90.cu @@ -2,17 +2,27 @@ // Copyright (c) 2026 RL-Kernel Contributors #include "../utils/tma_utils.cuh" +#include #include #include -#include +#include -#define TILE_V 4096 +// CUtensorMap box dimensions cannot exceed 256 elements. +#define TILE_V 256 -struct MaxOp { - __device__ __forceinline__ float operator()(float a, float b) const { - return fmaxf(a, b); +__device__ __forceinline__ float warp_reduce_max(float value) { + for (int offset = 16; offset > 0; offset /= 2) { + value = fmaxf(value, __shfl_down_sync(0xffffffff, value, offset)); } -}; + return value; +} + +__device__ __forceinline__ float warp_reduce_sum(float value) { + for (int offset = 16; offset > 0; offset /= 2) { + value += __shfl_down_sync(0xffffffff, value, offset); + } + return value; +} template __global__ void fused_logp_online_tma_kernel( @@ -20,6 +30,8 @@ __global__ void fused_logp_online_tma_kernel( const int* __restrict__ labels, const nv_bfloat16* __restrict__ logits_gmem, float* __restrict__ output_logp, + float* __restrict__ max_out, // Optional [batch_size] + float* __restrict__ logsum_out, // Optional [batch_size] int batch_size, int vocab_size) { @@ -48,13 +60,13 @@ __global__ void fused_logp_online_tma_kernel( if (warp_id == 0) { for (int step = 0; step < num_tiles; ++step) { int col_offset = step * TILE_V; - int current_tile_size = min(TILE_V, vocab_size - col_offset); if (step > 0) mbarrier_wait(mma_mbar_addr, phase ^ 1); if (lane_id == 0) { tma_2d_g2s(smem_addr, &logits_tmap, col_offset, row_idx, tma_mbar_addr); - mbarrier_arrive_expect_tx(tma_mbar_addr, current_tile_size * sizeof(nv_bfloat16)); + // The barrier expects a full tile even when the last tile is partial. + mbarrier_arrive_expect_tx(tma_mbar_addr, TILE_V * sizeof(nv_bfloat16)); } phase ^= 1; } @@ -62,9 +74,10 @@ __global__ void fused_logp_online_tma_kernel( else { const int consumer_tid = (warp_id - 1) * 32 + lane_id; const int num_consumers = (NUM_WARPS - 1) * 32; - - using BlockReduce = cub::BlockReduce; - __shared__ typename BlockReduce::TempStorage temp_storage; + // The producer warp does not enter this branch; reduce across consumers only. + __shared__ float warp_partials[NUM_WARPS - 1]; + __shared__ float s_tile_max; + __shared__ float s_tile_sum; float row_max = -CUDART_INF_F; float row_sum = 0.0f; @@ -79,9 +92,17 @@ __global__ void fused_logp_online_tma_kernel( float val = __bfloat162float(smem_logits[i]); tile_max = max(tile_max, val); } - float block_tile_max = BlockReduce(temp_storage).Reduce(tile_max, MaxOp{}); - __shared__ float s_tile_max; - if (consumer_tid == 0) s_tile_max = block_tile_max; + tile_max = warp_reduce_max(tile_max); + if (lane_id == 0) warp_partials[warp_id - 1] = tile_max; + asm volatile("bar.sync 1, %0;" :: "n"(num_consumers)); + if (consumer_tid == 0) { + float block_tile_max = -CUDART_INF_F; +#pragma unroll + for (int i = 0; i < NUM_WARPS - 1; ++i) { + block_tile_max = max(block_tile_max, warp_partials[i]); + } + s_tile_max = block_tile_max; + } asm volatile("bar.sync 1, %0;" :: "n"(num_consumers)); float tile_sum = 0.0f; @@ -89,9 +110,17 @@ __global__ void fused_logp_online_tma_kernel( float val = __bfloat162float(smem_logits[i]); tile_sum += expf(val - s_tile_max); } - float block_tile_sum = BlockReduce(temp_storage).Sum(tile_sum); - __shared__ float s_tile_sum; - if (consumer_tid == 0) s_tile_sum = block_tile_sum; + tile_sum = warp_reduce_sum(tile_sum); + if (lane_id == 0) warp_partials[warp_id - 1] = tile_sum; + asm volatile("bar.sync 1, %0;" :: "n"(num_consumers)); + if (consumer_tid == 0) { + float block_tile_sum = 0.0f; +#pragma unroll + for (int i = 0; i < NUM_WARPS - 1; ++i) { + block_tile_sum += warp_partials[i]; + } + s_tile_sum = block_tile_sum; + } asm volatile("bar.sync 1, %0;" :: "n"(num_consumers)); if (consumer_tid == 0) { @@ -106,17 +135,33 @@ __global__ void fused_logp_online_tma_kernel( } if (consumer_tid == 0) { + float log_sum = logf(row_sum); + if (max_out != nullptr) { + max_out[row_idx] = row_max; + logsum_out[row_idx] = log_sum; + } int label_idx = labels[row_idx]; - float label_val = __bfloat162float(logits_gmem[row_idx * vocab_size + label_idx]); - output_logp[row_idx] = label_val - row_max - logf(row_sum); + if (label_idx >= 0 && label_idx < vocab_size) { + int64_t label_offset = static_cast(row_idx) * vocab_size + label_idx; + float label_val = __bfloat162float(logits_gmem[label_offset]); + output_logp[row_idx] = label_val - row_max - log_sum; + } else { + output_logp[row_idx] = 0.0f; + } } } } -torch::Tensor fused_logp_sm90_forward(torch::Tensor logits, torch::Tensor labels) { +namespace { + +void launch_fused_logp_sm90( + const torch::Tensor& logits, + const torch::Tensor& labels, + const torch::Tensor& output, + float* max_ptr, + float* logsum_ptr) { int B = logits.size(0); int V = logits.size(1); - auto output = torch::empty({B}, logits.options().dtype(torch::kFloat)); CUtensorMap logits_tmap; init_tensor_map(&logits_tmap, @@ -127,7 +172,25 @@ torch::Tensor fused_logp_sm90_forward(torch::Tensor logits, torch::Tensor labels fused_logp_online_tma_kernel<4><<>>( logits_tmap, labels.data_ptr(), reinterpret_cast(logits.data_ptr()), - output.data_ptr(), B, V + output.data_ptr(), max_ptr, logsum_ptr, B, V ); +} + +} // namespace + +torch::Tensor fused_logp_sm90_forward(torch::Tensor logits, torch::Tensor labels) { + auto output = torch::empty({logits.size(0)}, logits.options().dtype(torch::kFloat)); + launch_fused_logp_sm90(logits, labels, output, nullptr, nullptr); return output; } + +std::tuple fused_logp_sm90_forward_with_lse( + torch::Tensor logits, + torch::Tensor labels) { + auto output = torch::empty({logits.size(0)}, logits.options().dtype(torch::kFloat)); + auto row_max = torch::empty({logits.size(0)}, logits.options().dtype(torch::kFloat)); + auto log_sum = torch::empty({logits.size(0)}, logits.options().dtype(torch::kFloat)); + launch_fused_logp_sm90( + logits, labels, output, row_max.data_ptr(), log_sum.data_ptr()); + return {output, row_max, log_sum}; +} diff --git a/csrc/fused_logp_kernel.cu b/csrc/fused_logp_kernel.cu index b620b047..e72200c8 100644 --- a/csrc/fused_logp_kernel.cu +++ b/csrc/fused_logp_kernel.cu @@ -3,6 +3,7 @@ #include #include #include +#include #include template @@ -113,6 +114,8 @@ __global__ void fused_logp_forward_kernel( const scalar_t* __restrict__ logits, // [TotalTokens, VocabSize] const int64_t* __restrict__ token_ids, // [TotalTokens] output_t* __restrict__ output, // [TotalTokens] + float* __restrict__ max_out, // Optional [TotalTokens] + float* __restrict__ logsum_out, // Optional [TotalTokens] const int64_t* __restrict__ row_indices, // Optional [ValidTokens] int64_t total_tokens, int vocab_size) { @@ -145,10 +148,15 @@ __global__ void fused_logp_forward_kernel( __syncthreads(); if (threadIdx.x == 0) { + float log_sum = logf(res_sum); + if (max_out != nullptr) { + max_out[row] = res_max; + logsum_out[row] = log_sum; + } int64_t target_id = token_ids[row]; if (target_id >= 0 && target_id < vocab_size) { float target_logit = static_cast(row_logits[target_id]); - output[row] = static_cast(target_logit - res_max - logf(res_sum)); + output[row] = static_cast(target_logit - res_max - log_sum); } else { output[row] = static_cast(0.0f); } @@ -160,6 +168,8 @@ __global__ void __launch_bounds__(BlockSize) fused_logp_forward_online_kernel( const scalar_t* __restrict__ logits, // [TotalTokens, VocabSize] const int64_t* __restrict__ token_ids, // [TotalTokens] output_t* __restrict__ output, // [TotalTokens] + float* __restrict__ max_out, // Optional [TotalTokens] + float* __restrict__ logsum_out, // Optional [TotalTokens] const int64_t* __restrict__ row_indices, // Optional [ValidTokens] int64_t total_tokens, int vocab_size) { @@ -195,16 +205,60 @@ __global__ void __launch_bounds__(BlockSize) fused_logp_forward_online_kernel( __syncthreads(); if (threadIdx.x == 0) { + float log_sum = logf(res_sum); + if (max_out != nullptr) { + max_out[row] = res_max; + logsum_out[row] = log_sum; + } int64_t target_id = token_ids[row]; if (target_id >= 0 && target_id < vocab_size) { float target_logit = static_cast(row_logits[target_id]); - output[row] = static_cast(target_logit - res_max - logf(res_sum)); + output[row] = static_cast(target_logit - res_max - log_sum); } else { output[row] = static_cast(0.0f); } } } +template +__global__ void fused_logp_backward_kernel( + const float* __restrict__ grad_out, // [TotalTokens] + const scalar_t* __restrict__ logits, // [TotalTokens, V] + const int64_t* __restrict__ token_ids, // [TotalTokens] + const float* __restrict__ row_max, // [TotalTokens] + const float* __restrict__ log_sum, // [TotalTokens] log(sum(exp(x - max))) + scalar_t* __restrict__ grad_logits, // [TotalTokens, V] + const int64_t* __restrict__ row_indices, // Optional [ValidTokens] + int64_t total_tokens, + int vocab_size) { + + int64_t row = row_indices == nullptr ? blockIdx.x : row_indices[blockIdx.x]; + if (row < 0 || row >= total_tokens) { + return; + } + + const scalar_t* row_logits = logits + row * vocab_size; + scalar_t* row_grad = grad_logits + row * vocab_size; + + float g = grad_out[row]; + float m = row_max[row]; + float ls = log_sum[row]; + int64_t target_id = token_ids[row]; + + if (target_id < 0 || target_id >= vocab_size) { + for (int i = threadIdx.x; i < vocab_size; i += blockDim.x) { + row_grad[i] = static_cast(0.0f); + } + return; + } + + for (int i = threadIdx.x; i < vocab_size; i += blockDim.x) { + float p = expf((static_cast(row_logits[i]) - m) - ls); + float indicator = (i == target_id) ? 1.0f : 0.0f; + row_grad[i] = static_cast(g * (indicator - p)); + } +} + namespace { #ifndef FUSED_LOGP_TWOPASS_BLOCK_SIZE @@ -276,6 +330,8 @@ void launch_fused_logp_online_variant( const torch::Tensor& logits, const torch::Tensor& token_ids, const torch::Tensor& output, + float* max_ptr, + float* logsum_ptr, const int64_t* row_indices_ptr, int64_t launch_rows, int64_t total_tokens, @@ -290,6 +346,8 @@ void launch_fused_logp_online_variant( logits.data_ptr(), token_ids.data_ptr(), output.data_ptr(), + max_ptr, + logsum_ptr, row_indices_ptr, total_tokens, static_cast(vocab_size)); @@ -344,6 +402,8 @@ void launch_fused_logp_kernel( const torch::Tensor& logits, const torch::Tensor& token_ids, const torch::Tensor& output, + float* max_ptr, + float* logsum_ptr, const int64_t* row_indices_ptr, int64_t launch_rows, int64_t total_tokens, @@ -374,6 +434,8 @@ void launch_fused_logp_kernel( logits.data_ptr(), token_ids.data_ptr(), output.data_ptr(), + max_ptr, + logsum_ptr, row_indices_ptr, total_tokens, static_cast(vocab_size)); @@ -387,6 +449,8 @@ void launch_fused_logp_online_kernel( const torch::Tensor& logits, const torch::Tensor& token_ids, const torch::Tensor& output, + float* max_ptr, + float* logsum_ptr, const int64_t* row_indices_ptr, int64_t launch_rows, int64_t total_tokens, @@ -424,6 +488,8 @@ void launch_fused_logp_online_kernel( logits, token_ids, output, + max_ptr, + logsum_ptr, row_indices_ptr, launch_rows, total_tokens, @@ -436,6 +502,8 @@ void launch_fused_logp_online_kernel( logits, token_ids, output, + max_ptr, + logsum_ptr, row_indices_ptr, launch_rows, total_tokens, @@ -447,12 +515,52 @@ void launch_fused_logp_online_kernel( C10_CUDA_KERNEL_LAUNCH_CHECK(); } -} // namespace +void launch_fused_logp_backward_kernel( + const torch::Tensor& grad_out, + const torch::Tensor& logits, + const torch::Tensor& token_ids, + const torch::Tensor& row_max, + const torch::Tensor& log_sum, + const torch::Tensor& grad_logits, + const int64_t* row_indices_ptr, + int64_t launch_rows, + int64_t total_tokens, + int64_t vocab_size) { + if (launch_rows == 0) { + return; + } -torch::Tensor fused_logp_forward_out( - torch::Tensor logits, - torch::Tensor token_ids, - torch::Tensor output) { + AT_DISPATCH_FLOATING_TYPES_AND2( + at::ScalarType::Half, + at::ScalarType::BFloat16, + logits.scalar_type(), + "fused_logp_backward_kernel", + ([&] { + fused_logp_backward_kernel<<< + static_cast(launch_rows), + kFusedLogpTwoPassBlockSize, + 0, + at::cuda::getCurrentCUDAStream()>>>( + grad_out.data_ptr(), + logits.data_ptr(), + token_ids.data_ptr(), + row_max.data_ptr(), + log_sum.data_ptr(), + grad_logits.data_ptr(), + row_indices_ptr, + total_tokens, + static_cast(vocab_size)); + })); + + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +torch::Tensor run_fused_logp_twopass( + const torch::Tensor& logits, + const torch::Tensor& token_ids, + const torch::Tensor& output, + float* max_ptr, + float* logsum_ptr) { check_fused_logp_inputs(logits, token_ids, output); auto logits_contig = logits.contiguous(); @@ -464,6 +572,8 @@ torch::Tensor fused_logp_forward_out( logits_contig, token_ids_contig, output, + max_ptr, + logsum_ptr, nullptr, total_tokens, total_tokens, @@ -472,11 +582,13 @@ torch::Tensor fused_logp_forward_out( return output; } -torch::Tensor fused_logp_forward_indexed_out( - torch::Tensor logits, - torch::Tensor token_ids, - torch::Tensor row_indices, - torch::Tensor output) { +torch::Tensor run_fused_logp_indexed_twopass( + const torch::Tensor& logits, + const torch::Tensor& token_ids, + const torch::Tensor& row_indices, + const torch::Tensor& output, + float* max_ptr, + float* logsum_ptr) { check_fused_logp_inputs(logits, token_ids, output); check_fused_logp_indices(logits, row_indices); @@ -492,6 +604,8 @@ torch::Tensor fused_logp_forward_indexed_out( logits_contig, token_ids_contig, output, + max_ptr, + logsum_ptr, row_indices_contig.data_ptr(), valid_tokens, total_tokens, @@ -500,10 +614,12 @@ torch::Tensor fused_logp_forward_indexed_out( return output; } -torch::Tensor fused_logp_forward_online_out( - torch::Tensor logits, - torch::Tensor token_ids, - torch::Tensor output) { +torch::Tensor run_fused_logp_online( + const torch::Tensor& logits, + const torch::Tensor& token_ids, + const torch::Tensor& output, + float* max_ptr, + float* logsum_ptr) { check_fused_logp_inputs(logits, token_ids, output); auto logits_contig = logits.contiguous(); @@ -515,6 +631,8 @@ torch::Tensor fused_logp_forward_online_out( logits_contig, token_ids_contig, output, + max_ptr, + logsum_ptr, nullptr, total_tokens, total_tokens, @@ -523,11 +641,13 @@ torch::Tensor fused_logp_forward_online_out( return output; } -torch::Tensor fused_logp_forward_online_indexed_out( - torch::Tensor logits, - torch::Tensor token_ids, - torch::Tensor row_indices, - torch::Tensor output) { +torch::Tensor run_fused_logp_online_indexed( + const torch::Tensor& logits, + const torch::Tensor& token_ids, + const torch::Tensor& row_indices, + const torch::Tensor& output, + float* max_ptr, + float* logsum_ptr) { check_fused_logp_inputs(logits, token_ids, output); check_fused_logp_indices(logits, row_indices); @@ -543,6 +663,8 @@ torch::Tensor fused_logp_forward_online_indexed_out( logits_contig, token_ids_contig, output, + max_ptr, + logsum_ptr, row_indices_contig.data_ptr(), valid_tokens, total_tokens, @@ -551,6 +673,70 @@ torch::Tensor fused_logp_forward_online_indexed_out( return output; } +void check_fused_logp_backward_inputs( + const torch::Tensor& grad_out, + const torch::Tensor& logits, + const torch::Tensor& token_ids, + const torch::Tensor& row_max, + const torch::Tensor& log_sum) { + TORCH_CHECK(logits.is_cuda(), "logits must be a CUDA tensor"); + TORCH_CHECK(grad_out.is_cuda(), "grad_out must be a CUDA tensor"); + TORCH_CHECK(token_ids.is_cuda(), "token_ids must be a CUDA tensor"); + TORCH_CHECK(row_max.is_cuda(), "row_max must be a CUDA tensor"); + TORCH_CHECK(log_sum.is_cuda(), "log_sum must be a CUDA tensor"); + TORCH_CHECK( + logits.device() == grad_out.device() && logits.device() == token_ids.device() && + logits.device() == row_max.device() && logits.device() == log_sum.device(), + "grad_out, logits, token_ids, row_max, and log_sum must be on the same CUDA device"); + TORCH_CHECK(logits.dim() == 2, "logits must be a 2D tensor"); + TORCH_CHECK(grad_out.scalar_type() == at::ScalarType::Float, "grad_out must be float32"); + TORCH_CHECK(row_max.scalar_type() == at::ScalarType::Float, "row_max must be float32"); + TORCH_CHECK(log_sum.scalar_type() == at::ScalarType::Float, "log_sum must be float32"); + TORCH_CHECK(token_ids.scalar_type() == at::ScalarType::Long, "token_ids must be int64"); + TORCH_CHECK(grad_out.numel() == logits.size(0), "grad_out length must match logits rows"); + TORCH_CHECK(row_max.numel() == logits.size(0), "row_max length must match logits rows"); + TORCH_CHECK(log_sum.numel() == logits.size(0), "log_sum length must match logits rows"); + TORCH_CHECK(token_ids.numel() == logits.size(0), "token_ids length must match logits rows"); + TORCH_CHECK( + logits.size(0) <= std::numeric_limits::max(), + "logits row count exceeds CUDA grid-x limit"); + TORCH_CHECK( + logits.size(1) <= std::numeric_limits::max(), + "logits vocab dimension exceeds int32 kernel limit"); +} + +} // namespace + +torch::Tensor fused_logp_forward_out( + torch::Tensor logits, + torch::Tensor token_ids, + torch::Tensor output) { + return run_fused_logp_twopass(logits, token_ids, output, nullptr, nullptr); +} + +torch::Tensor fused_logp_forward_indexed_out( + torch::Tensor logits, + torch::Tensor token_ids, + torch::Tensor row_indices, + torch::Tensor output) { + return run_fused_logp_indexed_twopass(logits, token_ids, row_indices, output, nullptr, nullptr); +} + +torch::Tensor fused_logp_forward_online_out( + torch::Tensor logits, + torch::Tensor token_ids, + torch::Tensor output) { + return run_fused_logp_online(logits, token_ids, output, nullptr, nullptr); +} + +torch::Tensor fused_logp_forward_online_indexed_out( + torch::Tensor logits, + torch::Tensor token_ids, + torch::Tensor row_indices, + torch::Tensor output) { + return run_fused_logp_online_indexed(logits, token_ids, row_indices, output, nullptr, nullptr); +} + torch::Tensor fused_logp_forward(torch::Tensor logits, torch::Tensor token_ids) { TORCH_CHECK(logits.dim() == 2, "logits must be a 2D tensor"); auto output = torch::empty({logits.size(0)}, logits.options()); @@ -586,3 +772,127 @@ torch::Tensor fused_logp_forward_online_indexed_fp32( auto output = torch::zeros({logits.size(0)}, logits.options().dtype(at::ScalarType::Float)); return fused_logp_forward_online_indexed_out(logits, token_ids, row_indices, output); } + +std::tuple fused_logp_forward_with_lse( + torch::Tensor logits, + torch::Tensor token_ids) { + TORCH_CHECK(logits.dim() == 2, "logits must be a 2D tensor"); + auto output = torch::empty({logits.size(0)}, logits.options().dtype(at::ScalarType::Float)); + auto row_max = torch::empty({logits.size(0)}, logits.options().dtype(at::ScalarType::Float)); + auto log_sum = torch::empty({logits.size(0)}, logits.options().dtype(at::ScalarType::Float)); + run_fused_logp_twopass( + logits, token_ids, output, row_max.data_ptr(), log_sum.data_ptr()); + return {output, row_max, log_sum}; +} + +torch::Tensor fused_logp_backward( + torch::Tensor grad_out, + torch::Tensor logits, + torch::Tensor token_ids, + torch::Tensor row_max, + torch::Tensor log_sum) { + check_fused_logp_backward_inputs(grad_out, logits, token_ids, row_max, log_sum); + + auto grad_out_contig = grad_out.contiguous(); + auto logits_contig = logits.contiguous(); + auto token_ids_contig = token_ids.contiguous(); + auto row_max_contig = row_max.contiguous(); + auto log_sum_contig = log_sum.contiguous(); + + int64_t total_tokens = logits_contig.size(0); + int64_t vocab_size = logits_contig.size(1); + auto grad_logits = torch::empty_like(logits_contig); + + launch_fused_logp_backward_kernel( + grad_out_contig, + logits_contig, + token_ids_contig, + row_max_contig, + log_sum_contig, + grad_logits, + nullptr, + total_tokens, + total_tokens, + vocab_size); + + return grad_logits; +} + +std::tuple fused_logp_forward_indexed_with_lse( + torch::Tensor logits, + torch::Tensor token_ids, + torch::Tensor row_indices) { + TORCH_CHECK(logits.dim() == 2, "logits must be a 2D tensor"); + auto output = torch::zeros({logits.size(0)}, logits.options().dtype(at::ScalarType::Float)); + auto row_max = torch::zeros({logits.size(0)}, logits.options().dtype(at::ScalarType::Float)); + auto log_sum = torch::zeros({logits.size(0)}, logits.options().dtype(at::ScalarType::Float)); + run_fused_logp_indexed_twopass( + logits, token_ids, row_indices, output, row_max.data_ptr(), + log_sum.data_ptr()); + return {output, row_max, log_sum}; +} + +std::tuple fused_logp_forward_online_with_lse( + torch::Tensor logits, + torch::Tensor token_ids) { + TORCH_CHECK(logits.dim() == 2, "logits must be a 2D tensor"); + auto output = torch::empty({logits.size(0)}, logits.options().dtype(at::ScalarType::Float)); + auto row_max = torch::empty({logits.size(0)}, logits.options().dtype(at::ScalarType::Float)); + auto log_sum = torch::empty({logits.size(0)}, logits.options().dtype(at::ScalarType::Float)); + run_fused_logp_online( + logits, token_ids, output, row_max.data_ptr(), log_sum.data_ptr()); + return {output, row_max, log_sum}; +} + +std::tuple fused_logp_forward_online_indexed_with_lse( + torch::Tensor logits, + torch::Tensor token_ids, + torch::Tensor row_indices) { + TORCH_CHECK(logits.dim() == 2, "logits must be a 2D tensor"); + auto output = torch::zeros({logits.size(0)}, logits.options().dtype(at::ScalarType::Float)); + auto row_max = torch::zeros({logits.size(0)}, logits.options().dtype(at::ScalarType::Float)); + auto log_sum = torch::zeros({logits.size(0)}, logits.options().dtype(at::ScalarType::Float)); + run_fused_logp_online_indexed( + logits, token_ids, row_indices, output, row_max.data_ptr(), + log_sum.data_ptr()); + return {output, row_max, log_sum}; +} + +torch::Tensor fused_logp_backward_indexed( + torch::Tensor grad_out, + torch::Tensor logits, + torch::Tensor token_ids, + torch::Tensor row_max, + torch::Tensor log_sum, + torch::Tensor row_indices) { + check_fused_logp_backward_inputs(grad_out, logits, token_ids, row_max, log_sum); + check_fused_logp_indices(logits, row_indices); + + auto grad_out_contig = grad_out.contiguous(); + auto logits_contig = logits.contiguous(); + auto token_ids_contig = token_ids.contiguous(); + auto row_max_contig = row_max.contiguous(); + auto log_sum_contig = log_sum.contiguous(); + auto row_indices_contig = row_indices.contiguous(); + + int64_t total_tokens = logits_contig.size(0); + int64_t vocab_size = logits_contig.size(1); + int64_t valid_tokens = row_indices_contig.numel(); + // Rows outside row_indices produced a constant-zero forward output, so their + // gradient is exactly zero; the kernel only writes the indexed rows. + auto grad_logits = torch::zeros_like(logits_contig); + + launch_fused_logp_backward_kernel( + grad_out_contig, + logits_contig, + token_ids_contig, + row_max_contig, + log_sum_contig, + grad_logits, + row_indices_contig.data_ptr(), + valid_tokens, + total_tokens, + vocab_size); + + return grad_logits; +} diff --git a/csrc/ops.cpp b/csrc/ops.cpp index 61ba4a3b..5bf23bf7 100644 --- a/csrc/ops.cpp +++ b/csrc/ops.cpp @@ -3,12 +3,16 @@ #include #include +#include // Fused LogP Declarations torch::Tensor fused_logp_forward(torch::Tensor logits, torch::Tensor token_ids); #if defined(__CUDACC__) || defined(KERNEL_ALIGN_WITH_SM90) torch::Tensor fused_logp_sm90_forward(torch::Tensor logits, torch::Tensor labels); +std::tuple fused_logp_sm90_forward_with_lse( + torch::Tensor logits, + torch::Tensor labels); std::vector fused_linear_logp_sm90_forward(torch::Tensor hidden, torch::Tensor weight, torch::Tensor target, @@ -69,6 +73,34 @@ torch::Tensor fused_logp_forward_online_out(torch::Tensor logits, torch::Tensor torch::Tensor fused_logp_forward_online_fp32(torch::Tensor logits, torch::Tensor token_ids); torch::Tensor fused_logp_forward_online_indexed_out(torch::Tensor logits, torch::Tensor token_ids, torch::Tensor row_indices, torch::Tensor output); torch::Tensor fused_logp_forward_online_indexed_fp32(torch::Tensor logits, torch::Tensor token_ids, torch::Tensor row_indices); +std::tuple fused_logp_forward_with_lse( + torch::Tensor logits, + torch::Tensor token_ids); +std::tuple fused_logp_forward_indexed_with_lse( + torch::Tensor logits, + torch::Tensor token_ids, + torch::Tensor row_indices); +std::tuple fused_logp_forward_online_with_lse( + torch::Tensor logits, + torch::Tensor token_ids); +std::tuple +fused_logp_forward_online_indexed_with_lse( + torch::Tensor logits, + torch::Tensor token_ids, + torch::Tensor row_indices); +torch::Tensor fused_logp_backward( + torch::Tensor grad_out, + torch::Tensor logits, + torch::Tensor token_ids, + torch::Tensor row_max, + torch::Tensor log_sum); +torch::Tensor fused_logp_backward_indexed( + torch::Tensor grad_out, + torch::Tensor logits, + torch::Tensor token_ids, + torch::Tensor row_max, + torch::Tensor log_sum, + torch::Tensor row_indices); torch::Tensor deterministic_logp_forward(torch::Tensor logits, torch::Tensor token_ids); torch::Tensor deterministic_logp_forward_out(torch::Tensor logits, torch::Tensor token_ids, torch::Tensor output); torch::Tensor deterministic_logp_forward_fp32(torch::Tensor logits, torch::Tensor token_ids); @@ -129,6 +161,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { #if defined(__CUDACC__) || defined(KERNEL_ALIGN_WITH_SM90) m.def("fused_logp_sm90", &fused_logp_sm90_forward, "TMA-accelerated Online Softmax Fused LogP"); + m.def("fused_logp_sm90_with_lse", &fused_logp_sm90_forward_with_lse, + "TMA fused logp forward returning (logp, row_max, log_sum)"); m.def("fused_linear_logp_sm90", &fused_linear_logp_sm90_forward, "TMA+WGMMA fused linear log-prob (hidden @ W^T -> selected-token logp), SM90"); m.def("fused_linear_logp_sm90_global_target", &fused_linear_logp_sm90_global_target_forward, @@ -161,6 +195,12 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("fused_logp_forward_online_fp32", &fused_logp_forward_online_fp32, "Fused logp online fp32"); m.def("fused_logp_forward_online_indexed_out", &fused_logp_forward_online_indexed_out, "Fused logp online indexed out"); m.def("fused_logp_forward_online_indexed_fp32", &fused_logp_forward_online_indexed_fp32, "Fused logp online indexed fp32"); + m.def("fused_logp_forward_with_lse", &fused_logp_forward_with_lse, "Fused logp fp32 forward with softmax statistics"); + m.def("fused_logp_forward_indexed_with_lse", &fused_logp_forward_indexed_with_lse, "Fused logp indexed fp32 forward with softmax statistics"); + m.def("fused_logp_forward_online_with_lse", &fused_logp_forward_online_with_lse, "Fused logp online fp32 forward with softmax statistics"); + m.def("fused_logp_forward_online_indexed_with_lse", &fused_logp_forward_online_indexed_with_lse, "Fused logp online indexed fp32 forward with softmax statistics"); + m.def("fused_logp_backward", &fused_logp_backward, "Fused logp backward from saved softmax statistics"); + m.def("fused_logp_backward_indexed", &fused_logp_backward_indexed, "Fused logp backward over indexed rows from saved softmax statistics"); m.def("deterministic_logp", &deterministic_logp_forward, "Batch-invariant deterministic logp"); m.def("deterministic_logp_forward_out", &deterministic_logp_forward_out, "Batch-invariant deterministic logp out"); m.def("deterministic_logp_forward_fp32", &deterministic_logp_forward_fp32, "Batch-invariant deterministic logp fp32"); diff --git a/docs/operators/fused-logp.md b/docs/operators/fused-logp.md index 7fc8008a..988f5238 100644 --- a/docs/operators/fused-logp.md +++ b/docs/operators/fused-logp.md @@ -27,19 +27,39 @@ reference = logp_ref.forward_fp32(logits, token_ids) ## Backends -| Backend | Wrapper | Native symbol | Notes | +| Backend | Wrapper | Extension entry point | Notes | | --- | --- | --- | --- | -| CUDA SM90 | `FusedLogpSM90Op` | `_C.fused_logp_sm90` | Experimental TMA-oriented path for 2D contiguous bf16 logits on Hopper-class GPUs. It is disabled by default and requires `RL_KERNEL_ENABLE_EXPERIMENTAL_SM90_LOGP=1`; otherwise the wrapper delegates to the CUDA generic fallback. | +| CUDA SM90 | `FusedLogpSM90Op` | `_C.fused_logp_sm90`, `_C.fused_logp_sm90_with_lse` | Experimental TMA path for eligible bf16 logits on Hopper-class GPUs. | | CUDA generic | `FusedLogpGenericOp` | `_C.fused_logp` | Generic compiled extension fallback. | -| PyTorch native | `NativeLogpOp` | None | PyTorch baseline/reference path. | +| PyTorch native | `NativeLogpOp` | — | PyTorch baseline/reference path. | ## Tensor Contract +Let `N` be the product of the leading dimensions of `logits`. + | Argument | Shape | Dtype | Requirements | | --- | --- | --- | --- | -| `logits` | `[N, V]` | `bfloat16` for the experimental SM90 fast path; fp16/fp32 use generic fallback | Contiguous, on the target device for the experimental SM90 fast path. | -| `token_ids` / `labels` | `[N]` | Converted to `int32` | Same logical device as `logits`. | -| Output | `[N]` | Backend-defined tensor dtype | One selected log probability per row. | +| `logits` | `[..., V]` | Floating point | The generic and native wrappers flatten leading dimensions. The SM90 fast path requires contiguous 2D bf16 `[N, V]`. | +| `token_ids` | `[...]` | Integer | Shape must match the leading dimensions of `logits`, with every value in `[0, V)`. Wrappers move IDs to the logits device and use int64 for native/generic kernels or int32 for the SM90 forward. | +| `row_indices` | `[K]` | Integer | Optional flattened row indices used by indexed variants. Values must be in `[0, N)` and are converted to int64. | +| Output | `[...]` | See below | One selected log probability per input row. | + +`forward` / `apply` return the logits dtype on the native and generic paths, while +`forward_fp32` / `apply_fp32` and the other allocating `*_fp32` variants return +float32. Caller-provided `*_out` variants use the output buffer dtype. An eligible +SM90 TMA call always returns float32. + +The experimental SM90 path is selected only when all of the following hold: + +- the extension was built with `KERNEL_ALIGN_FORCE_SM90=1` for a supported GPU; +- `RL_KERNEL_ENABLE_EXPERIMENTAL_SM90_LOGP=1` is set at runtime; +- `logits` is a contiguous 2D bf16 tensor; and +- `V` is divisible by 8, so each bf16 row has the 16-byte-aligned stride required + by TMA. + +If the build, environment, or GPU requirements are not met, the registry does not +select the SM90 backend. If an input tensor is ineligible, the SM90 wrapper delegates +to the CUDA generic backend. ## Reference Semantics @@ -48,23 +68,69 @@ ref = torch.log_softmax(logits.float(), dim=-1) ref = torch.gather(ref, dim=-1, index=token_ids.unsqueeze(-1).long()).squeeze(-1) ``` +## Backward / Autograd + +The CUDA generic backend is differentiable with respect to `logits`. When +`logits.requires_grad` is set under grad mode, the allocating variants — +`apply` / `apply_fp32` / `indexed_fp32` / `online_fp32` / `online_indexed_fp32` — +route through a `torch.autograd.Function` using the same forward reduction path as +the corresponding no-grad call. It additionally saves separate float32 `row_max` +and `log_sum` statistics, avoiding the precision loss that can occur when a large +constant logit offset is folded into one float32 log-sum-exp value. + +Backward rebuilds probabilities as `exp((logit - row_max) - log_sum)` and computes + +``` +grad_logits[v] = grad_out * (1[v == token_id] - softmax(logits)[v]) +``` + +in a dedicated kernel without materializing another logits-sized intermediate. +Indexed variants only touch selected rows; all other rows receive exactly-zero +gradient. + +The generic CUDA `*_out` variants remain non-differentiable, matching PyTorch's +`out=` convention, and raise `RuntimeError` for grad-requiring `logits`. +`DeterministicLogpCUDAOp` is also forward-only. + +Eligible SM90 calls are differentiable too. Grad mode uses +`fused_logp_sm90_with_lse`, which runs the same TMA reduction as the no-grad entry +point while also returning `row_max` and `log_sum`. It then reuses the generic +elementwise CUDA backward kernel; no SM90-specific backward kernel is needed. +Grad mode changes neither the float32 output contract nor the forward values. + ## Tests ```bash -python -m pytest tests/test_logp.py -q -python -m pytest tests/test_op_accuracy.py -q +python -m pytest tests/test_logp.py -q -rs +python -m pytest tests/test_op_accuracy.py -q -rs +python -m pytest tests/test_fused_logp_backward.py -q -rs ``` `tests/test_logp.py` covers the PyTorch reference contract, dtype behavior, -backward-compatible aliases, batch invariance, and registry dispatch. The existing -operator accuracy tests continue to validate native/CUDA fused API compatibility. +backward-compatible aliases, batch invariance, and registry dispatch. The operator +accuracy tests validate native/CUDA fused API compatibility. +`tests/test_fused_logp_backward.py` covers gradients across dtypes and variants, +train/inference bitwise consistency, forward-only guards, and SM90 multi-tile, +partial-tile, and fallback behavior. + +CUDA cases require the compiled extension. The SM90 cases additionally require a +Hopper GPU and an SM90-enabled build; otherwise pytest reports them as skipped. A +matching H100 build can be installed with: + +```bash +FORCE_CUDA=1 KERNEL_ALIGN_FORCE_SM90=1 TORCH_CUDA_ARCH_LIST="9.0+PTX" \ + python -m pip install --no-build-isolation --no-deps -e . +``` ## Implementation Files - `rl_engine/kernels/registry.py` +- `rl_engine/_C.pyi` - `rl_engine/kernels/ops/pytorch/loss/logp.py` - `rl_engine/kernels/ops/cuda/loss/logp.py` - `csrc/ops.cpp` - `csrc/fused_logp_kernel.cu` - `csrc/cuda/fused_logp_sm90.cu` - `tests/test_logp.py` +- `tests/test_op_accuracy.py` +- `tests/test_fused_logp_backward.py` diff --git a/rl_engine/_C.pyi b/rl_engine/_C.pyi index 6ab592c6..d3f5b606 100644 --- a/rl_engine/_C.pyi +++ b/rl_engine/_C.pyi @@ -4,6 +4,10 @@ import torch def fused_logp(logits: torch.Tensor, token_ids: torch.Tensor) -> torch.Tensor: ... def fused_logp_sm90(logits: torch.Tensor, labels: torch.Tensor) -> torch.Tensor: ... +def fused_logp_sm90_with_lse( + logits: torch.Tensor, + labels: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: ... def fused_linear_logp_sm90( hidden: torch.Tensor, weight: torch.Tensor, @@ -109,6 +113,39 @@ def fused_logp_forward_online_indexed_fp32( token_ids: torch.Tensor, row_indices: torch.Tensor, ) -> torch.Tensor: ... +def fused_logp_forward_with_lse( + logits: torch.Tensor, + token_ids: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: ... +def fused_logp_forward_indexed_with_lse( + logits: torch.Tensor, + token_ids: torch.Tensor, + row_indices: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: ... +def fused_logp_forward_online_with_lse( + logits: torch.Tensor, + token_ids: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: ... +def fused_logp_forward_online_indexed_with_lse( + logits: torch.Tensor, + token_ids: torch.Tensor, + row_indices: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: ... +def fused_logp_backward( + grad_out: torch.Tensor, + logits: torch.Tensor, + token_ids: torch.Tensor, + row_max: torch.Tensor, + log_sum: torch.Tensor, +) -> torch.Tensor: ... +def fused_logp_backward_indexed( + grad_out: torch.Tensor, + logits: torch.Tensor, + token_ids: torch.Tensor, + row_max: torch.Tensor, + log_sum: torch.Tensor, + row_indices: torch.Tensor, +) -> torch.Tensor: ... def deterministic_logp(logits: torch.Tensor, token_ids: torch.Tensor) -> torch.Tensor: ... def deterministic_logp_forward_out( logits: torch.Tensor, diff --git a/rl_engine/kernels/ops/cuda/loss/logp.py b/rl_engine/kernels/ops/cuda/loss/logp.py index 79442531..b852fab7 100644 --- a/rl_engine/kernels/ops/cuda/loss/logp.py +++ b/rl_engine/kernels/ops/cuda/loss/logp.py @@ -9,16 +9,57 @@ from rl_engine.utils.logger import logger +def _grad_requested(logits: torch.Tensor) -> bool: + return torch.is_grad_enabled() and logits.requires_grad + + +_BACKWARD_EXT_SYMBOLS = ( + "fused_logp_forward_with_lse", + "fused_logp_forward_indexed_with_lse", + "fused_logp_forward_online_with_lse", + "fused_logp_forward_online_indexed_with_lse", + "fused_logp_backward", + "fused_logp_backward_indexed", +) + +_SM90_BACKWARD_EXT_SYMBOLS = ( + "fused_logp_sm90_with_lse", + "fused_logp_backward", +) + + +class _FusedLogpSM90Function(torch.autograd.Function): + """SM90 TMA fused forward with the generic CUDA backward.""" + + @staticmethod + def forward( + ctx, + logits: torch.Tensor, + labels_i32: torch.Tensor, + token_ids: torch.Tensor, + ) -> torch.Tensor: + logp, row_max, log_sum = _C.fused_logp_sm90_with_lse(logits, labels_i32) + ctx.save_for_backward(logits, token_ids, row_max, log_sum) + return logp + + @staticmethod + def backward(ctx, grad_out: torch.Tensor): + logits, token_ids, row_max, log_sum = ctx.saved_tensors + grad_logits = _C.fused_logp_backward(grad_out, logits, token_ids, row_max, log_sum) + return grad_logits, None, None + + class FusedLogpSM90Op: - """TMA-accelerated Fused LogP for SM90+ cards.""" + """TMA-accelerated fused logp for supported SM90+ GPUs.""" is_fused_logp = True def __init__(self): if not _EXT_AVAILABLE or not hasattr(_C, "fused_logp_sm90"): raise RuntimeError( - "TMA Fused LogP kernel is not compiled or unsupported on this card architecture. " - "Please rebuild extension using 'pip install -e .'" + "The TMA fused logp kernel is unavailable in this build or on this GPU. " + "Rebuild on a supported GPU with " + "'KERNEL_ALIGN_FORCE_SM90=1 pip install -e .'" ) self.op = _C.fused_logp_sm90 self._fallback = None @@ -38,12 +79,26 @@ def _can_use_sm90(self, logits: torch.Tensor) -> bool: and logits.dim() == 2 and logits.dtype == torch.bfloat16 and logits.is_contiguous() + and logits.size(1) % 8 == 0 ) + def _check_backward_symbols(self) -> None: + missing = [name for name in _SM90_BACKWARD_EXT_SYMBOLS if not hasattr(_C, name)] + if missing: + raise RuntimeError( + "Differentiable SM90 fused logp requires kernels missing from the compiled " + f"extension: {', '.join(missing)}. " + "Rebuild with 'KERNEL_ALIGN_FORCE_SM90=1 pip install -e .'" + ) + def apply(self, logits: torch.Tensor, labels: torch.Tensor) -> torch.Tensor: if not self._can_use_sm90(logits): return self._fallback_op().apply(logits, labels) labels_fused = labels.to(device=logits.device, dtype=torch.int32).contiguous() + if _grad_requested(logits): + self._check_backward_symbols() + token_ids = labels.to(device=logits.device, dtype=torch.long).contiguous() + return _FusedLogpSM90Function.apply(logits, labels_fused, token_ids) return self.op(logits, labels_fused) def apply_fp32(self, logits: torch.Tensor, token_ids: torch.Tensor) -> torch.Tensor: @@ -91,6 +146,49 @@ def online_indexed_fp32( return self._fallback_op().online_indexed_fp32(logits, token_ids, row_indices) +class _FusedLogpFunction(torch.autograd.Function): + """Fused logp variants with CUDA backward from saved softmax statistics.""" + + @staticmethod + def forward( + ctx, + logits: torch.Tensor, + token_ids: torch.Tensor, + row_indices: torch.Tensor | None = None, + online: bool = False, + ) -> torch.Tensor: + # Make logits contiguous once so backward reuses the saved buffer instead + # of materializing another [T, V] copy. + logits = logits.contiguous() + if row_indices is None: + forward = ( + _C.fused_logp_forward_online_with_lse if online else _C.fused_logp_forward_with_lse + ) + logp, row_max, log_sum = forward(logits, token_ids) + ctx.save_for_backward(logits, token_ids, row_max, log_sum) + else: + forward = ( + _C.fused_logp_forward_online_indexed_with_lse + if online + else _C.fused_logp_forward_indexed_with_lse + ) + logp, row_max, log_sum = forward(logits, token_ids, row_indices) + ctx.save_for_backward(logits, token_ids, row_max, log_sum, row_indices) + return logp + + @staticmethod + def backward(ctx, grad_out: torch.Tensor): + if len(ctx.saved_tensors) == 5: + logits, token_ids, row_max, log_sum, row_indices = ctx.saved_tensors + grad_logits = _C.fused_logp_backward_indexed( + grad_out, logits, token_ids, row_max, log_sum, row_indices + ) + else: + logits, token_ids, row_max, log_sum = ctx.saved_tensors + grad_logits = _C.fused_logp_backward(grad_out, logits, token_ids, row_max, log_sum) + return grad_logits, None, None, None + + class FusedLogpGenericOp: """Generic custom CUDA fallback Fused LogP with RL variants.""" @@ -127,19 +225,45 @@ def _prepare_output(self, output: torch.Tensor, orig_shape: torch.Size) -> torch def _prepare_indices(self, row_indices: torch.Tensor, logits: torch.Tensor) -> torch.Tensor: return row_indices.reshape(-1).to(device=logits.device, dtype=torch.long).contiguous() + def _needs_grad(self, logits: torch.Tensor) -> bool: + if not _grad_requested(logits): + return False + missing = [name for name in _BACKWARD_EXT_SYMBOLS if not hasattr(_C, name)] + if missing: + raise RuntimeError( + "Differentiable fused logp requires kernels missing from the compiled " + f"extension: {', '.join(missing)}. " + "Rebuild the extension with 'pip install -e .'" + ) + return True + + def _reject_grad(self, logits: torch.Tensor, variant: str) -> None: + if _grad_requested(logits): + raise RuntimeError( + f"{type(self).__name__}.{variant} is forward-only and does not propagate " + "gradients to logits; call it under torch.no_grad() or detach logits first." + ) + def apply(self, logits: torch.Tensor, token_ids: torch.Tensor) -> torch.Tensor: logits_2d, token_ids_1d, orig_shape = self._prepare_inputs(logits, token_ids) - results = self.op(logits_2d, token_ids_1d) + if self._needs_grad(logits_2d): + results = _FusedLogpFunction.apply(logits_2d, token_ids_1d).to(logits_2d.dtype) + else: + results = self.op(logits_2d, token_ids_1d) return results.view(orig_shape) def apply_fp32(self, logits: torch.Tensor, token_ids: torch.Tensor) -> torch.Tensor: logits_2d, token_ids_1d, orig_shape = self._prepare_inputs(logits, token_ids) - results = self._backend.fused_logp_forward_fp32(logits_2d, token_ids_1d) + if self._needs_grad(logits_2d): + results = _FusedLogpFunction.apply(logits_2d, token_ids_1d) + else: + results = self._backend.fused_logp_forward_fp32(logits_2d, token_ids_1d) return results.view(orig_shape) def out( self, logits: torch.Tensor, token_ids: torch.Tensor, output: torch.Tensor ) -> torch.Tensor: + self._reject_grad(logits, "out") logits_2d, token_ids_1d, orig_shape = self._prepare_inputs(logits, token_ids) output_1d = self._prepare_output(output, orig_shape) results = self._backend.fused_logp_forward_out(logits_2d, token_ids_1d, output_1d) @@ -152,6 +276,7 @@ def indexed_out( row_indices: torch.Tensor, output: torch.Tensor, ) -> torch.Tensor: + self._reject_grad(logits, "indexed_out") logits_2d, token_ids_1d, orig_shape = self._prepare_inputs(logits, token_ids) row_indices_1d = self._prepare_indices(row_indices, logits) output_1d = self._prepare_output(output, orig_shape) @@ -165,14 +290,18 @@ def indexed_fp32( ) -> torch.Tensor: logits_2d, token_ids_1d, orig_shape = self._prepare_inputs(logits, token_ids) row_indices_1d = self._prepare_indices(row_indices, logits) - results = self._backend.fused_logp_forward_indexed_fp32( - logits_2d, token_ids_1d, row_indices_1d - ) + if self._needs_grad(logits_2d): + results = _FusedLogpFunction.apply(logits_2d, token_ids_1d, row_indices_1d) + else: + results = self._backend.fused_logp_forward_indexed_fp32( + logits_2d, token_ids_1d, row_indices_1d + ) return results.view(orig_shape) def online_out( self, logits: torch.Tensor, token_ids: torch.Tensor, output: torch.Tensor ) -> torch.Tensor: + self._reject_grad(logits, "online_out") logits_2d, token_ids_1d, orig_shape = self._prepare_inputs(logits, token_ids) output_1d = self._prepare_output(output, orig_shape) results = self._backend.fused_logp_forward_online_out(logits_2d, token_ids_1d, output_1d) @@ -180,7 +309,10 @@ def online_out( def online_fp32(self, logits: torch.Tensor, token_ids: torch.Tensor) -> torch.Tensor: logits_2d, token_ids_1d, orig_shape = self._prepare_inputs(logits, token_ids) - results = self._backend.fused_logp_forward_online_fp32(logits_2d, token_ids_1d) + if self._needs_grad(logits_2d): + results = _FusedLogpFunction.apply(logits_2d, token_ids_1d, None, True) + else: + results = self._backend.fused_logp_forward_online_fp32(logits_2d, token_ids_1d) return results.view(orig_shape) def online_indexed_out( @@ -190,6 +322,7 @@ def online_indexed_out( row_indices: torch.Tensor, output: torch.Tensor, ) -> torch.Tensor: + self._reject_grad(logits, "online_indexed_out") logits_2d, token_ids_1d, orig_shape = self._prepare_inputs(logits, token_ids) row_indices_1d = self._prepare_indices(row_indices, logits) output_1d = self._prepare_output(output, orig_shape) @@ -203,9 +336,12 @@ def online_indexed_fp32( ) -> torch.Tensor: logits_2d, token_ids_1d, orig_shape = self._prepare_inputs(logits, token_ids) row_indices_1d = self._prepare_indices(row_indices, logits) - results = self._backend.fused_logp_forward_online_indexed_fp32( - logits_2d, token_ids_1d, row_indices_1d - ) + if self._needs_grad(logits_2d): + results = _FusedLogpFunction.apply(logits_2d, token_ids_1d, row_indices_1d, True) + else: + results = self._backend.fused_logp_forward_online_indexed_fp32( + logits_2d, token_ids_1d, row_indices_1d + ) return results.view(orig_shape) diff --git a/tests/test_fused_logp_backward.py b/tests/test_fused_logp_backward.py new file mode 100644 index 00000000..93d8e30a --- /dev/null +++ b/tests/test_fused_logp_backward.py @@ -0,0 +1,512 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Tests for fused logp CUDA backward support.""" + +from __future__ import annotations + +import pytest +import torch + +from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE +from rl_engine.testing.reference_ops import selected_logprobs_reference + +pytestmark = pytest.mark.skipif( + not ( + torch.cuda.is_available() and _EXT_AVAILABLE and hasattr(_C, "fused_logp_backward_indexed") + ), + reason="requires CUDA and the compiled fused logp backward extension", +) + + +def _make_inputs( + rows: int, + vocab: int, + *, + dtype: torch.dtype = torch.float32, + seed: int = 123, +) -> tuple[torch.Tensor, torch.Tensor]: + gen = torch.Generator().manual_seed(seed) + logits = torch.randn(rows, vocab, generator=gen, dtype=dtype).cuda() + token_ids = torch.randint(0, vocab, (rows,), generator=gen, dtype=torch.long).cuda() + return logits, token_ids + + +def _reference_grad( + logits: torch.Tensor, token_ids: torch.Tensor, grad_out: torch.Tensor +) -> torch.Tensor: + ref_logits = logits.detach().clone().requires_grad_(True) + (selected_logprobs_reference(ref_logits, token_ids) * grad_out.float()).sum().backward() + assert ref_logits.grad is not None + return ref_logits.grad + + +def _make_row_indices(rows: int, *, seed: int = 321) -> torch.Tensor: + gen = torch.Generator().manual_seed(seed) + return torch.randperm(rows, generator=gen)[: rows // 2].cuda() + + +def _mask_rows(grad: torch.Tensor, row_indices: torch.Tensor) -> torch.Tensor: + mask = torch.zeros(grad.size(0), dtype=torch.bool, device=grad.device) + mask[row_indices] = True + return grad * mask.unsqueeze(1) + + +_GRAD_TOLERANCES = [ + (torch.float32, 1e-5), + (torch.bfloat16, 2e-2), + (torch.float16, 5e-3), +] + + +class TestFusedLogpForwardWithLse: + def test_logp_matches_reference(self): + logits, token_ids = _make_inputs(8, 257) + logp, _, _ = _C.fused_logp_forward_with_lse(logits, token_ids) + ref = selected_logprobs_reference(logits, token_ids) + assert logp.dtype == torch.float32 + assert torch.allclose(logp, ref, atol=1e-5, rtol=1e-5) + + def test_lse_components_match_reference(self): + # Keep LSE decomposed so a large row_max cannot round away log_sum in float32. + logits, token_ids = _make_inputs(8, 257) + _, row_max, log_sum = _C.fused_logp_forward_with_lse(logits, token_ids) + assert row_max.dtype == torch.float32 + assert log_sum.dtype == torch.float32 + assert torch.equal(row_max, logits.float().amax(dim=-1)) + ref_lse = torch.logsumexp(logits.float(), dim=-1) + assert torch.allclose(row_max + log_sum, ref_lse, atol=1e-5, rtol=1e-5) + + def test_matches_existing_fp32_forward_bitwise(self): + logits, token_ids = _make_inputs(8, 257) + logp, _, _ = _C.fused_logp_forward_with_lse(logits, token_ids) + legacy = _C.fused_logp_forward_fp32(logits, token_ids) + assert torch.equal(logp, legacy) + + +class TestFusedLogpVariantForwardWithLse: + def test_indexed_matches_legacy_bitwise(self): + logits, token_ids = _make_inputs(8, 257) + row_indices = _make_row_indices(8) + logp, _, _ = _C.fused_logp_forward_indexed_with_lse(logits, token_ids, row_indices) + legacy = _C.fused_logp_forward_indexed_fp32(logits, token_ids, row_indices) + assert torch.equal(logp, legacy) + + def test_indexed_unselected_rows_are_zero(self): + logits, token_ids = _make_inputs(8, 257) + row_indices = _make_row_indices(8) + logp, row_max, log_sum = _C.fused_logp_forward_indexed_with_lse( + logits, token_ids, row_indices + ) + selected = torch.zeros(8, dtype=torch.bool, device=logp.device) + selected[row_indices] = True + assert torch.equal(logp[~selected], torch.zeros_like(logp[~selected])) + assert torch.equal(row_max[~selected], torch.zeros_like(row_max[~selected])) + assert torch.equal(log_sum[~selected], torch.zeros_like(log_sum[~selected])) + ref_lse = torch.logsumexp(logits.float(), dim=-1) + lse = row_max + log_sum + assert torch.allclose(lse[selected], ref_lse[selected], atol=1e-5, rtol=1e-5) + + def test_online_matches_legacy_bitwise(self): + logits, token_ids = _make_inputs(8, 257) + logp, _, _ = _C.fused_logp_forward_online_with_lse(logits, token_ids) + legacy = _C.fused_logp_forward_online_fp32(logits, token_ids) + assert torch.equal(logp, legacy) + + def test_online_lse_components_match_reference(self): + logits, token_ids = _make_inputs(8, 257) + _, row_max, log_sum = _C.fused_logp_forward_online_with_lse(logits, token_ids) + ref_lse = torch.logsumexp(logits.float(), dim=-1) + assert torch.allclose(row_max + log_sum, ref_lse, atol=1e-5, rtol=1e-5) + + def test_online_indexed_matches_legacy_bitwise(self): + logits, token_ids = _make_inputs(8, 257) + row_indices = _make_row_indices(8) + logp, _, _ = _C.fused_logp_forward_online_indexed_with_lse(logits, token_ids, row_indices) + legacy = _C.fused_logp_forward_online_indexed_fp32(logits, token_ids, row_indices) + assert torch.equal(logp, legacy) + + +class TestFusedLogpBackwardKernel: + @pytest.mark.parametrize("dtype, atol", _GRAD_TOLERANCES) + def test_grad_matches_autograd_reference(self, dtype, atol): + logits, token_ids = _make_inputs(8, 257, dtype=dtype) + gen = torch.Generator().manual_seed(456) + grad_out = torch.randn(8, generator=gen).cuda() + + _, row_max, log_sum = _C.fused_logp_forward_with_lse(logits, token_ids) + grad_logits = _C.fused_logp_backward(grad_out, logits, token_ids, row_max, log_sum) + + ref_grad = _reference_grad(logits, token_ids, grad_out) + assert grad_logits.dtype == dtype + assert torch.allclose(grad_logits.float(), ref_grad.float(), atol=atol, rtol=0.0) + + def test_grad_rows_sum_to_zero(self): + logits, token_ids = _make_inputs(8, 257) + grad_out = torch.randn(8, generator=torch.Generator().manual_seed(7)).cuda() + _, row_max, log_sum = _C.fused_logp_forward_with_lse(logits, token_ids) + grad_logits = _C.fused_logp_backward(grad_out, logits, token_ids, row_max, log_sum) + row_sums = grad_logits.sum(dim=-1) + assert torch.allclose(row_sums, torch.zeros_like(row_sums), atol=1e-4) + + def test_invalid_target_rows_get_zero_grad(self): + logits, token_ids = _make_inputs(4, 33) + token_ids = token_ids.clone() + token_ids[1] = -1 + grad_out = torch.ones(4).cuda() + _, row_max, log_sum = _C.fused_logp_forward_with_lse(logits, token_ids) + grad_logits = _C.fused_logp_backward(grad_out, logits, token_ids, row_max, log_sum) + assert torch.equal(grad_logits[1], torch.zeros_like(grad_logits[1])) + assert not torch.equal(grad_logits[0], torch.zeros_like(grad_logits[0])) + + @pytest.mark.parametrize("shift", [-1e4, 1e4, 1e8]) + def test_grad_stable_under_constant_shift(self, shift): + # Separate (row_max, log_sum) statistics preserve shift-invariant gradients. + logits, token_ids = _make_inputs(8, 257) + logits = logits + shift + grad_out = torch.randn(8, generator=torch.Generator().manual_seed(456)).cuda() + + _, row_max, log_sum = _C.fused_logp_forward_with_lse(logits, token_ids) + grad_logits = _C.fused_logp_backward(grad_out, logits, token_ids, row_max, log_sum) + + ref_grad = _reference_grad(logits, token_ids, grad_out) + assert torch.allclose(grad_logits, ref_grad, atol=1e-5, rtol=0.0) + + +class TestFusedLogpBackwardIndexedKernel: + def test_grad_matches_masked_reference(self): + logits, token_ids = _make_inputs(8, 257) + row_indices = _make_row_indices(8) + grad_out = torch.randn(8, generator=torch.Generator().manual_seed(456)).cuda() + + _, row_max, log_sum = _C.fused_logp_forward_indexed_with_lse(logits, token_ids, row_indices) + grad_logits = _C.fused_logp_backward_indexed( + grad_out, logits, token_ids, row_max, log_sum, row_indices + ) + + ref_grad = _mask_rows(_reference_grad(logits, token_ids, grad_out), row_indices) + assert torch.allclose(grad_logits, ref_grad, atol=1e-5, rtol=0.0) + + def test_unselected_rows_get_exactly_zero_grad(self): + logits, token_ids = _make_inputs(8, 257) + row_indices = _make_row_indices(8) + grad_out = torch.ones(8).cuda() + _, row_max, log_sum = _C.fused_logp_forward_indexed_with_lse(logits, token_ids, row_indices) + grad_logits = _C.fused_logp_backward_indexed( + grad_out, logits, token_ids, row_max, log_sum, row_indices + ) + selected = torch.zeros(8, dtype=torch.bool, device=logits.device) + selected[row_indices] = True + assert torch.equal(grad_logits[~selected], torch.zeros_like(grad_logits[~selected])) + assert not torch.equal(grad_logits[selected], torch.zeros_like(grad_logits[selected])) + + +class TestFusedLogpOpAutogradRouting: + def _op(self): + from rl_engine.kernels.ops.cuda.loss.logp import FusedLogpGenericOp + + return FusedLogpGenericOp() + + @pytest.mark.parametrize("dtype, atol", _GRAD_TOLERANCES) + def test_apply_backward_matches_reference(self, dtype, atol): + logits, token_ids = _make_inputs(8, 257, dtype=dtype) + logits.requires_grad_(True) + gen = torch.Generator().manual_seed(654) + grad_out = torch.randn(8, generator=gen).cuda().to(dtype) + + out = self._op().apply(logits, token_ids) + assert out.dtype == dtype + out.backward(grad_out) + + ref_grad = _reference_grad(logits, token_ids, grad_out) + assert logits.grad is not None + assert torch.allclose(logits.grad.float(), ref_grad.float(), atol=atol, rtol=0.0) + + def test_apply_fp32_backward_matches_reference(self): + logits, token_ids = _make_inputs(8, 257) + logits.requires_grad_(True) + grad_out = torch.randn(8, generator=torch.Generator().manual_seed(11)).cuda() + + out = self._op().apply_fp32(logits, token_ids) + out.backward(grad_out) + + ref_grad = _reference_grad(logits, token_ids, grad_out) + assert logits.grad is not None + assert torch.allclose(logits.grad, ref_grad, atol=1e-5, rtol=0.0) + + def test_apply_fp32_grad_stable_under_constant_shift(self): + logits, token_ids = _make_inputs(8, 257) + logits = (logits + 1e8).requires_grad_(True) + grad_out = torch.randn(8, generator=torch.Generator().manual_seed(17)).cuda() + + out = self._op().apply_fp32(logits, token_ids) + out.backward(grad_out) + + ref_grad = _reference_grad(logits, token_ids, grad_out) + assert logits.grad is not None + assert torch.allclose(logits.grad, ref_grad, atol=1e-5, rtol=0.0) + + def test_3d_inputs_flow_gradients(self): + gen = torch.Generator().manual_seed(99) + logits = torch.randn(2, 4, 65, generator=gen).cuda().requires_grad_(True) + token_ids = torch.randint(0, 65, (2, 4), generator=gen).cuda() + + out = self._op().apply(logits, token_ids) + assert out.shape == (2, 4) + out.sum().backward() + assert logits.grad is not None + assert logits.grad.shape == logits.shape + + def test_no_grad_path_keeps_input_dtype(self): + logits, token_ids = _make_inputs(8, 257, dtype=torch.bfloat16) + with torch.no_grad(): + out = self._op().apply(logits, token_ids) + assert out.dtype == torch.bfloat16 + + def test_no_grad_path_bitwise_unchanged(self): + logits, token_ids = _make_inputs(8, 257) + with torch.no_grad(): + routed = self._op().apply(logits, token_ids) + legacy = _C.fused_logp(logits, token_ids) + assert torch.equal(routed, legacy) + + def test_grad_and_no_grad_forward_bitwise_equal(self): + logits, token_ids = _make_inputs(8, 257) + op = self._op() + with torch.no_grad(): + rollout = op.apply_fp32(logits, token_ids) + train = op.apply_fp32(logits.clone().requires_grad_(True), token_ids) + assert torch.equal(train.detach(), rollout) + + def test_online_fp32_backward_matches_reference(self): + logits, token_ids = _make_inputs(8, 257) + logits.requires_grad_(True) + grad_out = torch.randn(8, generator=torch.Generator().manual_seed(21)).cuda() + + out = self._op().online_fp32(logits, token_ids) + out.backward(grad_out) + + ref_grad = _reference_grad(logits, token_ids, grad_out) + assert logits.grad is not None + assert torch.allclose(logits.grad, ref_grad, atol=1e-5, rtol=0.0) + + @pytest.mark.parametrize("variant", ["indexed_fp32", "online_indexed_fp32"]) + def test_indexed_variants_backward_matches_masked_reference(self, variant): + logits, token_ids = _make_inputs(8, 257) + logits.requires_grad_(True) + row_indices = _make_row_indices(8) + grad_out = torch.randn(8, generator=torch.Generator().manual_seed(31)).cuda() + + out = getattr(self._op(), variant)(logits, token_ids, row_indices) + out.backward(grad_out) + + ref_grad = _mask_rows(_reference_grad(logits, token_ids, grad_out), row_indices) + assert logits.grad is not None + assert torch.allclose(logits.grad, ref_grad, atol=1e-5, rtol=0.0) + + def test_online_fp32_grad_and_no_grad_forward_bitwise_equal(self): + # The online reduction order differs from two-pass, so consistency must + # hold within the online path itself. + logits, token_ids = _make_inputs(8, 257) + op = self._op() + with torch.no_grad(): + rollout = op.online_fp32(logits, token_ids) + train = op.online_fp32(logits.clone().requires_grad_(True), token_ids) + assert torch.equal(train.detach(), rollout) + + def test_out_variants_reject_grad_logits(self): + # out= style variants write a caller-provided buffer and stay + # non-differentiable, matching PyTorch's own out= convention. + logits, token_ids = _make_inputs(8, 257) + op = self._op() + with torch.no_grad(): + op.out(logits, token_ids, torch.empty(8).cuda()) + logits.requires_grad_(True) + with pytest.raises(RuntimeError, match="forward-only"): + op.out(logits, token_ids, torch.empty(8).cuda()) + with pytest.raises(RuntimeError, match="forward-only"): + op.online_out(logits, token_ids, torch.empty(8).cuda()) + + def test_stale_extension_raises_rebuild_hint(self, monkeypatch): + logits, token_ids = _make_inputs(8, 257) + logits.requires_grad_(True) + op = self._op() + monkeypatch.delattr(_C, "fused_logp_backward") + with pytest.raises(RuntimeError, match="pip install -e ."): + op.apply(logits, token_ids) + + def _sm90_op(self, monkeypatch): + from rl_engine.kernels.ops.cuda.loss.logp import FusedLogpSM90Op + + # Bypass __init__ (the SM90 kernel may not be compiled) to test routing only. + op = FusedLogpSM90Op.__new__(FusedLogpSM90Op) + op._fallback = None + op.op = lambda *_: pytest.fail("no-grad SM90 kernel entry must not run under grad") + monkeypatch.setenv("RL_KERNEL_ENABLE_EXPERIMENTAL_SM90_LOGP", "1") + return op + + def test_sm90_grad_path_keeps_fp32_contract_and_matches_reference(self, monkeypatch): + # Use the generic statistics-returning forward to isolate SM90 autograd routing. + op = self._sm90_op(monkeypatch) + monkeypatch.setattr( + _C, + "fused_logp_sm90_with_lse", + lambda logits, labels: _C.fused_logp_forward_with_lse(logits, labels.long()), + raising=False, + ) + + gen = torch.Generator().manual_seed(3) + logits = torch.randn(4, 72, generator=gen).cuda().to(torch.bfloat16).requires_grad_(True) + token_ids = torch.randint(0, 72, (4,), generator=gen).cuda() + + out = op.apply(logits, token_ids) + assert out.dtype == torch.float32 + assert out.grad_fn is not None + + grad_out = torch.randn(4, generator=gen).cuda() + out.backward(grad_out) + ref_grad = _reference_grad(logits, token_ids, grad_out) + assert logits.grad is not None + assert torch.allclose(logits.grad.float(), ref_grad.float(), atol=2e-2, rtol=0.0) + + def test_sm90_no_grad_path_still_runs_tma_kernel(self, monkeypatch): + from rl_engine.kernels.ops.cuda.loss.logp import FusedLogpSM90Op + + op = FusedLogpSM90Op.__new__(FusedLogpSM90Op) + op._fallback = None + seen = {} + + def fake_kernel(logits, labels): + seen["labels_dtype"] = labels.dtype + return torch.zeros(logits.size(0), device=logits.device) + + op.op = fake_kernel + monkeypatch.setenv("RL_KERNEL_ENABLE_EXPERIMENTAL_SM90_LOGP", "1") + + gen = torch.Generator().manual_seed(3) + logits = torch.randn(4, 72, generator=gen).cuda().to(torch.bfloat16) + token_ids = torch.randint(0, 72, (4,), generator=gen).cuda() + with torch.no_grad(): + out = op.apply(logits, token_ids) + assert seen["labels_dtype"] == torch.int32 + assert out.shape == (4,) + + def test_sm90_grad_path_stale_extension_raises_rebuild_hint(self, monkeypatch): + op = self._sm90_op(monkeypatch) + monkeypatch.delattr(_C, "fused_logp_sm90_with_lse", raising=False) + + gen = torch.Generator().manual_seed(3) + logits = torch.randn(4, 72, generator=gen).cuda().to(torch.bfloat16).requires_grad_(True) + token_ids = torch.randint(0, 72, (4,), generator=gen).cuda() + + with pytest.raises(RuntimeError, match="pip install -e ."): + op.apply(logits, token_ids) + + @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16, torch.float16]) + def test_apply_grad_and_no_grad_forward_bitwise_equal(self, dtype): + # The fp32 grad path must round exactly like the typed forward kernel. + logits, token_ids = _make_inputs(8, 257, dtype=dtype) + op = self._op() + with torch.no_grad(): + rollout = op.apply(logits, token_ids) + train = op.apply(logits.clone().requires_grad_(True), token_ids) + assert train.dtype == rollout.dtype == dtype + assert torch.equal(train.detach(), rollout) + + +_SM90_HOPPER_AVAILABLE = ( + torch.cuda.is_available() + and _EXT_AVAILABLE + and hasattr(_C, "fused_logp_sm90_with_lse") + and torch.cuda.get_device_capability()[0] == 9 +) + + +@pytest.mark.skipif( + not _SM90_HOPPER_AVAILABLE, + reason="requires a Hopper GPU and the KERNEL_ALIGN_FORCE_SM90 build", +) +class TestFusedLogpSM90Kernel: + def _inputs(self) -> tuple[torch.Tensor, torch.Tensor]: + gen = torch.Generator().manual_seed(5) + # Exercise repeated TMA phase transitions across multiple tiles. + logits = torch.randn(4, 8192, generator=gen).to(torch.bfloat16).cuda() + token_ids = torch.randint(0, 8192, (4,), generator=gen, dtype=torch.long).cuda() + return logits, token_ids + + def test_with_lse_logp_matches_plain_sm90_bitwise(self): + logits, token_ids = self._inputs() + labels = token_ids.to(torch.int32) + logp, _, _ = _C.fused_logp_sm90_with_lse(logits, labels) + assert torch.equal(logp, _C.fused_logp_sm90(logits, labels)) + + def test_stats_match_reference(self): + logits, token_ids = self._inputs() + _, row_max, log_sum = _C.fused_logp_sm90_with_lse(logits, token_ids.to(torch.int32)) + assert torch.equal(row_max, logits.float().amax(dim=-1)) + ref_lse = torch.logsumexp(logits.float(), dim=-1) + assert torch.allclose(row_max + log_sum, ref_lse, atol=1e-3, rtol=1e-5) + + def test_partial_last_tile_matches_reference(self): + logits, token_ids = _make_inputs(4, 520, dtype=torch.bfloat16, seed=5) + logp, row_max, log_sum = _C.fused_logp_sm90_with_lse(logits, token_ids.to(torch.int32)) + ref = selected_logprobs_reference(logits, token_ids) + assert torch.allclose(logp, ref, atol=1e-3, rtol=1e-5) + assert torch.allclose( + row_max + log_sum, + torch.logsumexp(logits.float(), dim=-1), + atol=1e-3, + rtol=1e-5, + ) + + def test_invalid_labels_return_zero(self): + logits, token_ids = _make_inputs(4, 520, dtype=torch.bfloat16, seed=5) + labels = token_ids.to(torch.int32) + labels[0] = -1 + labels[1] = logits.size(1) + + logp, _, _ = _C.fused_logp_sm90_with_lse(logits, labels) + assert torch.equal(logp[:2], torch.zeros_like(logp[:2])) + + ref = selected_logprobs_reference(logits, token_ids) + assert torch.allclose(logp[2:], ref[2:], atol=1e-3, rtol=1e-5) + + def test_unaligned_row_stride_uses_generic_fallback(self, monkeypatch): + from rl_engine.kernels.ops.cuda.loss.logp import FusedLogpSM90Op + + monkeypatch.setenv("RL_KERNEL_ENABLE_EXPERIMENTAL_SM90_LOGP", "1") + op = FusedLogpSM90Op() + op.op = lambda *_: pytest.fail("unaligned row stride must not use the TMA kernel") + logits, token_ids = _make_inputs(4, 257, dtype=torch.bfloat16, seed=5) + with torch.no_grad(): + out = op.apply(logits, token_ids) + ref = selected_logprobs_reference(logits, token_ids) + assert torch.allclose(out.float(), ref, atol=2e-2, rtol=0.0) + + def test_op_grad_and_no_grad_forward_bitwise_equal(self, monkeypatch): + from rl_engine.kernels.ops.cuda.loss.logp import FusedLogpSM90Op + + monkeypatch.setenv("RL_KERNEL_ENABLE_EXPERIMENTAL_SM90_LOGP", "1") + op = FusedLogpSM90Op() + logits, token_ids = self._inputs() + with torch.no_grad(): + rollout = op.apply(logits, token_ids) + train = op.apply(logits.clone().requires_grad_(True), token_ids) + assert train.dtype == rollout.dtype == torch.float32 + assert torch.equal(train.detach(), rollout) + + def test_op_grad_matches_reference(self, monkeypatch): + from rl_engine.kernels.ops.cuda.loss.logp import FusedLogpSM90Op + + monkeypatch.setenv("RL_KERNEL_ENABLE_EXPERIMENTAL_SM90_LOGP", "1") + op = FusedLogpSM90Op() + logits, token_ids = self._inputs() + logits.requires_grad_(True) + grad_out = torch.randn(4, generator=torch.Generator().manual_seed(9)).cuda() + + out = op.apply(logits, token_ids) + out.backward(grad_out) + + ref_grad = _reference_grad(logits, token_ids, grad_out) + assert logits.grad is not None + assert torch.allclose(logits.grad.float(), ref_grad.float(), atol=2e-2, rtol=0.0)