From dc9d4c6d7dd8462c28d9ab446d96304c8e6d1a58 Mon Sep 17 00:00:00 2001 From: nodeeeeee Date: Sun, 30 Aug 2026 13:26:25 +0000 Subject: [PATCH 1/3] feat: add batch-invariant MHC H Aggregate kernel Signed-off-by: nodeeeeee --- csrc/cuda/mhc/mhc_pre_h_aggregate.cu | 62 +++++++++++ csrc/cuda/mhc/mhc_pre_h_aggregate_kernel.cuh | 107 +++++++++++++++++++ csrc/ops.cpp | 8 ++ rl_engine/_C.pyi | 4 + setup.py | 1 + tests/test_mhc_pre_h_aggregate.py | 72 +++++++++++++ 6 files changed, 254 insertions(+) create mode 100644 csrc/cuda/mhc/mhc_pre_h_aggregate.cu create mode 100644 csrc/cuda/mhc/mhc_pre_h_aggregate_kernel.cuh create mode 100644 tests/test_mhc_pre_h_aggregate.py diff --git a/csrc/cuda/mhc/mhc_pre_h_aggregate.cu b/csrc/cuda/mhc/mhc_pre_h_aggregate.cu new file mode 100644 index 00000000..d109952e --- /dev/null +++ b/csrc/cuda/mhc/mhc_pre_h_aggregate.cu @@ -0,0 +1,62 @@ +#include +#include +#include +#include + +#include +#include + +#include +#include + +#include "mhc_pre_h_aggregate_kernel.cuh" + +torch::Tensor mhc_pre_h_aggregate_cuda(torch::Tensor residual, + torch::Tensor pre) { + TORCH_CHECK(residual.is_cuda() && pre.is_cuda(), + "residual and pre must be CUDA tensors"); + TORCH_CHECK(residual.device() == pre.device(), + "residual and pre must be on the same CUDA device"); + TORCH_CHECK(residual.is_contiguous() && pre.is_contiguous(), + "residual and pre must be contiguous"); + TORCH_CHECK(residual.scalar_type() == torch::kBFloat16, + "residual must be bfloat16"); + TORCH_CHECK(pre.scalar_type() == torch::kFloat32, + "pre must be float32"); + TORCH_CHECK(residual.dim() == 3 && residual.size(1) == 4, + "residual must have shape [num_tokens, 4, hidden_size]"); + TORCH_CHECK(pre.dim() == 2 && pre.size(1) == 4, + "pre must have shape [num_tokens, 4]"); + TORCH_CHECK(residual.size(0) == pre.size(0), + "residual and pre must have the same num_tokens"); + + int64_t const num_tokens = residual.size(0); + int64_t const hidden_size = residual.size(2); + TORCH_CHECK(num_tokens <= std::numeric_limits::max(), + "num_tokens exceeds the CUDA grid limit"); + + auto output = torch::empty({num_tokens, hidden_size}, residual.options()); + if (num_tokens == 0 || hidden_size == 0) { + return output; + } + + c10::cuda::CUDAGuard const device_guard(residual.device()); + int device = 0; + int major = 0; + C10_CUDA_CHECK(cudaGetDevice(&device)); + C10_CUDA_CHECK(cudaDeviceGetAttribute( + &major, cudaDevAttrComputeCapabilityMajor, device)); + TORCH_CHECK(major >= 8, "mhc_pre_h_aggregate requires SM80 or newer"); + + auto const* residual_ptr = reinterpret_cast<__nv_bfloat16 const*>( + residual.data_ptr()); + auto const* pre_ptr = pre.data_ptr(); + auto* output_ptr = reinterpret_cast<__nv_bfloat16*>( + output.data_ptr()); + + cudaError_t const status = rl_kernel::mhc::launch_mhc_pre_h_aggregate( + residual_ptr, pre_ptr, output_ptr, num_tokens, hidden_size, + at::cuda::getCurrentCUDAStream(), major >= 9); + C10_CUDA_CHECK(status); + return output; +} diff --git a/csrc/cuda/mhc/mhc_pre_h_aggregate_kernel.cuh b/csrc/cuda/mhc/mhc_pre_h_aggregate_kernel.cuh new file mode 100644 index 00000000..b9c13bbe --- /dev/null +++ b/csrc/cuda/mhc/mhc_pre_h_aggregate_kernel.cuh @@ -0,0 +1,107 @@ +#pragma once + +#include +#include + +#include + +namespace rl_kernel::mhc { + +constexpr int kMhcPreHAggregateDecodeThreads = 1024; +constexpr int kMhcPreHAggregateBatchThreads = 512; + +__global__ void mhc_pre_h_aggregate_kernel(__nv_bfloat16 const* residual, + float const* pre, + __nv_bfloat16* output, + int64_t hidden_size) { +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + cudaGridDependencySynchronize(); +#endif + + __shared__ float weights[4]; + if (threadIdx.x < 4) { + weights[threadIdx.x] = + pre[static_cast(blockIdx.x) * 4 + threadIdx.x]; + } + __syncthreads(); + + int64_t const token_offset = + static_cast(blockIdx.x) * 4 * hidden_size; + int64_t const output_offset = static_cast(blockIdx.x) * hidden_size; + if ((hidden_size & 1) == 0) { + auto const* residual_pairs = + reinterpret_cast<__nv_bfloat162 const*>(residual + token_offset); + auto* output_pairs = + reinterpret_cast<__nv_bfloat162*>(output + output_offset); + int64_t const pair_count = hidden_size / 2; + for (int64_t hidden_pair = threadIdx.x; hidden_pair < pair_count; + hidden_pair += blockDim.x) { + // store two bf16 to a 32-bit reg + float2 const value_0 = __bfloat1622float2(residual_pairs[hidden_pair]); + float2 const value_1 = + __bfloat1622float2(residual_pairs[pair_count + hidden_pair]); + float2 const value_2 = + __bfloat1622float2(residual_pairs[2 * pair_count + hidden_pair]); + float2 const value_3 = + __bfloat1622float2(residual_pairs[3 * pair_count + hidden_pair]); + float2 result; + float const left_x = __fadd_rn(__fmul_rn(weights[0], value_0.x), + __fmul_rn(weights[1], value_1.x)); + float const right_x = __fadd_rn(__fmul_rn(weights[2], value_2.x), + __fmul_rn(weights[3], value_3.x)); + result.x = __fadd_rn(left_x, right_x); + float const left_y = __fadd_rn(__fmul_rn(weights[0], value_0.y), + __fmul_rn(weights[1], value_1.y)); + float const right_y = __fadd_rn(__fmul_rn(weights[2], value_2.y), + __fmul_rn(weights[3], value_3.y)); + result.y = __fadd_rn(left_y, right_y); + output_pairs[hidden_pair] = __floats2bfloat162_rn(result.x, result.y); + } + } else { + for (int64_t hidden = threadIdx.x; hidden < hidden_size; + hidden += blockDim.x) { + float const product_0 = __fmul_rn( + weights[0], __bfloat162float(residual[token_offset + hidden])); + float const product_1 = __fmul_rn( + weights[1], + __bfloat162float(residual[token_offset + hidden_size + hidden])); + float const product_2 = __fmul_rn( + weights[2], __bfloat162float( + residual[token_offset + 2 * hidden_size + hidden])); + float const product_3 = __fmul_rn( + weights[3], __bfloat162float( + residual[token_offset + 3 * hidden_size + hidden])); + float const left = __fadd_rn(product_0, product_1); + float const right = __fadd_rn(product_2, product_3); + output[output_offset + hidden] = __float2bfloat16_rn(__fadd_rn(left, right)); + } + } + +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + cudaTriggerProgrammaticLaunchCompletion(); +#endif +} + +inline cudaError_t launch_mhc_pre_h_aggregate( + __nv_bfloat16 const* residual, float const* pre, __nv_bfloat16* output, + int64_t num_tokens, int64_t hidden_size, cudaStream_t stream, + bool enable_pdl) { + cudaLaunchConfig_t config{}; + config.gridDim = dim3(static_cast(num_tokens)); + config.blockDim = dim3(num_tokens <= 128 ? kMhcPreHAggregateDecodeThreads + : kMhcPreHAggregateBatchThreads); + config.stream = stream; + + cudaLaunchAttribute attribute{}; + if (enable_pdl) { + attribute.id = cudaLaunchAttributeProgrammaticStreamSerialization; + attribute.val.programmaticStreamSerializationAllowed = 1; + config.attrs = &attribute; + config.numAttrs = 1; + } + + return cudaLaunchKernelEx(&config, mhc_pre_h_aggregate_kernel, residual, pre, + output, hidden_size); +} + +} diff --git a/csrc/ops.cpp b/csrc/ops.cpp index aecc30ed..2887dea5 100644 --- a/csrc/ops.cpp +++ b/csrc/ops.cpp @@ -164,6 +164,10 @@ int64_t rmsnorm_backward_dw_chunks_cuda(int64_t rows); void reduce_rows_fp32_left_fold_cuda( torch::Tensor rows, torch::Tensor output); + +torch::Tensor mhc_pre_h_aggregate_cuda( + torch::Tensor residual, + torch::Tensor pre); #endif static void rmsnorm_check_input(const torch::Tensor& x, const char* name) { @@ -491,6 +495,10 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { "reduce_rows_fp32_left_fold", &reduce_rows_fp32_left_fold, "Ascending-row FP32 left-fold reduction CUDA"); + m.def( + "mhc_pre_h_aggregate", + &mhc_pre_h_aggregate_cuda, + "Batch-invariant MHC H Aggregate CUDA"); #endif // registry SiLU / SwiGLU (elementwise activation) diff --git a/rl_engine/_C.pyi b/rl_engine/_C.pyi index 8e0e865a..f2f64c8c 100644 --- a/rl_engine/_C.pyi +++ b/rl_engine/_C.pyi @@ -2,6 +2,10 @@ # This file is a type stub for the compiled C++ extension module. import torch +def mhc_pre_h_aggregate( + residual: torch.Tensor, + pre: torch.Tensor, +) -> torch.Tensor: ... def deterministic_collective_ipc_meta( tensor: torch.Tensor, ) -> tuple[list[int], int]: ... diff --git a/setup.py b/setup.py index 79f882d9..b79aa3ea 100644 --- a/setup.py +++ b/setup.py @@ -145,6 +145,7 @@ def get_extensions(): # 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") + cuda_sources.append("csrc/cuda/mhc/mhc_pre_h_aggregate.cu") nvcc_flags = ["-O3", "-Xfatbin", "-compress-all"] if envs.env_flag(envs.KERNEL_ALIGN_USE_FAST_MATH): diff --git a/tests/test_mhc_pre_h_aggregate.py b/tests/test_mhc_pre_h_aggregate.py new file mode 100644 index 00000000..6563e6a5 --- /dev/null +++ b/tests/test_mhc_pre_h_aggregate.py @@ -0,0 +1,72 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +import pytest +import torch + + +def _kernel_available() -> bool: + if not torch.cuda.is_available(): + return False + if torch.cuda.get_device_capability()[0] < 8: + return False + try: + from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE + except Exception: + return False + return _EXT_AVAILABLE and hasattr(_C, "mhc_pre_h_aggregate") + + +requires_mhc_kernel = pytest.mark.skipif( + not _kernel_available(), + reason="mhc_pre_h_aggregate requires the CUDA extension on SM80 or newer", +) + + +def _same_bytes(left: torch.Tensor, right: torch.Tensor) -> bool: + return torch.equal(left.view(torch.uint8), right.view(torch.uint8)) + + +@requires_mhc_kernel +def test_mhc_pre_h_aggregate_is_batch_invariant_and_matches_pytorch(): + from rl_engine.kernels.ops.base import _C + + num_tokens = 129 + hidden_size = 4096 + num_runs = 100 + generator = torch.Generator(device="cuda").manual_seed(0) + residual = torch.randn( + num_tokens, + 4, + hidden_size, + dtype=torch.bfloat16, + device="cuda", + generator=generator, + ) + pre = torch.rand( + num_tokens, + 4, + dtype=torch.float32, + device="cuda", + generator=generator, + ) + + first = _C.mhc_pre_h_aggregate(residual, pre) + for _ in range(num_runs - 1): + repeated = _C.mhc_pre_h_aggregate(residual, pre) + assert _same_bytes(first, repeated) + + token_by_token = torch.cat( + [ + _C.mhc_pre_h_aggregate( + residual[token : token + 1], pre[token : token + 1] + ) + for token in range(num_tokens) + ] + ) + assert _same_bytes(first, token_by_token) + + reference = torch.sum( + pre.unsqueeze(-1) * residual.to(torch.float32), dim=1 + ).to(torch.bfloat16) + torch.testing.assert_close(first, reference, atol=5e-2, rtol=1e-2) From 1afadb62badaa62ced7c176d20f20cd48a3801a7 Mon Sep 17 00:00:00 2001 From: nodeeeeee Date: Sun, 30 Aug 2026 17:32:48 +0000 Subject: [PATCH 2/3] feat: add MHC H Aggregate backward and gtest coverage Signed-off-by: nodeeeeee --- csrc/cuda/mhc/mhc_pre_h_aggregate.cu | 64 ++++++++++ csrc/cuda/mhc/mhc_pre_h_aggregate_kernel.cuh | 119 +++++++++++++++++++ csrc/ops.cpp | 9 ++ rl_engine/_C.pyi | 5 + rl_engine/kernels/gtest/gradient_adapters.py | 28 +++++ rl_engine/kernels/gtest/operator_inputs.py | 14 +++ rl_engine/kernels/gtest/operator_specs.py | 14 ++- rl_engine/kernels/ops/cuda/mhc.py | 60 ++++++++++ rl_engine/kernels/ops/pytorch/mhc.py | 40 +++++++ tests/test_gradient_invariance.py | 4 + tests/test_mhc_pre_h_aggregate.py | 82 ++++++++++++- tests/test_operator_inputs.py | 23 ++++ 12 files changed, 454 insertions(+), 8 deletions(-) create mode 100644 rl_engine/kernels/ops/cuda/mhc.py create mode 100644 rl_engine/kernels/ops/pytorch/mhc.py diff --git a/csrc/cuda/mhc/mhc_pre_h_aggregate.cu b/csrc/cuda/mhc/mhc_pre_h_aggregate.cu index d109952e..1f8658fc 100644 --- a/csrc/cuda/mhc/mhc_pre_h_aggregate.cu +++ b/csrc/cuda/mhc/mhc_pre_h_aggregate.cu @@ -8,6 +8,7 @@ #include #include +#include #include "mhc_pre_h_aggregate_kernel.cuh" @@ -60,3 +61,66 @@ torch::Tensor mhc_pre_h_aggregate_cuda(torch::Tensor residual, C10_CUDA_CHECK(status); return output; } + +std::vector mhc_pre_h_aggregate_backward_cuda( + torch::Tensor grad_output, torch::Tensor residual, torch::Tensor pre) { + TORCH_CHECK(grad_output.is_cuda() && residual.is_cuda() && pre.is_cuda(), + "grad_output, residual, and pre must be CUDA tensors"); + TORCH_CHECK(grad_output.device() == residual.device() && + residual.device() == pre.device(), + "grad_output, residual, and pre must be on the same CUDA device"); + TORCH_CHECK(grad_output.is_contiguous() && residual.is_contiguous() && + pre.is_contiguous(), + "grad_output, residual, and pre must be contiguous"); + TORCH_CHECK(grad_output.scalar_type() == torch::kBFloat16 && + residual.scalar_type() == torch::kBFloat16, + "grad_output and residual must be bfloat16"); + TORCH_CHECK(pre.scalar_type() == torch::kFloat32, "pre must be float32"); + TORCH_CHECK(residual.dim() == 3 && residual.size(1) == 4, + "residual must have shape [num_tokens, 4, hidden_size]"); + TORCH_CHECK(pre.dim() == 2 && pre.size(1) == 4, + "pre must have shape [num_tokens, 4]"); + TORCH_CHECK(grad_output.dim() == 2 && + grad_output.size(0) == residual.size(0) && + grad_output.size(1) == residual.size(2), + "grad_output must have shape [num_tokens, hidden_size]"); + TORCH_CHECK(pre.size(0) == residual.size(0), + "residual and pre must have the same num_tokens"); + + int64_t const num_tokens = residual.size(0); + int64_t const hidden_size = residual.size(2); + TORCH_CHECK(num_tokens <= std::numeric_limits::max(), + "num_tokens exceeds the CUDA grid limit"); + + auto grad_residual = torch::empty_like(residual); + auto grad_pre = torch::zeros_like(pre); + if (num_tokens == 0 || hidden_size == 0) { + return {grad_residual, grad_pre}; + } + + c10::cuda::CUDAGuard const device_guard(residual.device()); + int device = 0; + int major = 0; + C10_CUDA_CHECK(cudaGetDevice(&device)); + C10_CUDA_CHECK(cudaDeviceGetAttribute( + &major, cudaDevAttrComputeCapabilityMajor, device)); + TORCH_CHECK(major >= 8, + "mhc_pre_h_aggregate_backward requires SM80 or newer"); + + auto const* grad_output_ptr = reinterpret_cast<__nv_bfloat16 const*>( + grad_output.data_ptr()); + auto const* residual_ptr = reinterpret_cast<__nv_bfloat16 const*>( + residual.data_ptr()); + auto const* pre_ptr = pre.data_ptr(); + auto* grad_residual_ptr = reinterpret_cast<__nv_bfloat16*>( + grad_residual.data_ptr()); + auto* grad_pre_ptr = grad_pre.data_ptr(); + + cudaError_t const status = + rl_kernel::mhc::launch_mhc_pre_h_aggregate_backward( + grad_output_ptr, residual_ptr, pre_ptr, grad_residual_ptr, + grad_pre_ptr, num_tokens, hidden_size, + at::cuda::getCurrentCUDAStream(), major >= 9); + C10_CUDA_CHECK(status); + return {grad_residual, grad_pre}; +} diff --git a/csrc/cuda/mhc/mhc_pre_h_aggregate_kernel.cuh b/csrc/cuda/mhc/mhc_pre_h_aggregate_kernel.cuh index b9c13bbe..ef864b28 100644 --- a/csrc/cuda/mhc/mhc_pre_h_aggregate_kernel.cuh +++ b/csrc/cuda/mhc/mhc_pre_h_aggregate_kernel.cuh @@ -9,6 +9,7 @@ namespace rl_kernel::mhc { constexpr int kMhcPreHAggregateDecodeThreads = 1024; constexpr int kMhcPreHAggregateBatchThreads = 512; +constexpr int kMhcPreHAggregateBackwardThreads = 256; __global__ void mhc_pre_h_aggregate_kernel(__nv_bfloat16 const* residual, float const* pre, @@ -82,6 +83,101 @@ __global__ void mhc_pre_h_aggregate_kernel(__nv_bfloat16 const* residual, #endif } +__device__ __forceinline__ float mhc_warp_sum(float value) { +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { + value = __fadd_rn(value, __shfl_down_sync(0xffffffff, value, offset)); + } + return value; +} + +__global__ void mhc_pre_h_aggregate_backward_kernel( + __nv_bfloat16 const* grad_output, __nv_bfloat16 const* residual, + float const* pre, __nv_bfloat16* grad_residual, float* grad_pre, + int64_t hidden_size) { +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + cudaGridDependencySynchronize(); +#endif + + constexpr int kWarpSize = 32; + constexpr int kNumWarps = kMhcPreHAggregateBackwardThreads / kWarpSize; + __shared__ float warp_sums[4][kNumWarps]; + + int64_t const token = static_cast(blockIdx.x); + int64_t const output_offset = token * hidden_size; + int64_t const residual_offset = token * 4 * hidden_size; + float const weight_0 = pre[token * 4]; + float const weight_1 = pre[token * 4 + 1]; + float const weight_2 = pre[token * 4 + 2]; + float const weight_3 = pre[token * 4 + 3]; + float sum_0 = 0.0f; + float sum_1 = 0.0f; + float sum_2 = 0.0f; + float sum_3 = 0.0f; + + for (int64_t hidden = threadIdx.x; hidden < hidden_size; + hidden += blockDim.x) { + float const dy = __bfloat162float(grad_output[output_offset + hidden]); + float const residual_0 = + __bfloat162float(residual[residual_offset + hidden]); + float const residual_1 = + __bfloat162float(residual[residual_offset + hidden_size + hidden]); + float const residual_2 = __bfloat162float( + residual[residual_offset + 2 * hidden_size + hidden]); + float const residual_3 = __bfloat162float( + residual[residual_offset + 3 * hidden_size + hidden]); + + sum_0 = __fadd_rn(sum_0, __fmul_rn(dy, residual_0)); + sum_1 = __fadd_rn(sum_1, __fmul_rn(dy, residual_1)); + sum_2 = __fadd_rn(sum_2, __fmul_rn(dy, residual_2)); + sum_3 = __fadd_rn(sum_3, __fmul_rn(dy, residual_3)); + + grad_residual[residual_offset + hidden] = + __float2bfloat16_rn(__fmul_rn(dy, weight_0)); + grad_residual[residual_offset + hidden_size + hidden] = + __float2bfloat16_rn(__fmul_rn(dy, weight_1)); + grad_residual[residual_offset + 2 * hidden_size + hidden] = + __float2bfloat16_rn(__fmul_rn(dy, weight_2)); + grad_residual[residual_offset + 3 * hidden_size + hidden] = + __float2bfloat16_rn(__fmul_rn(dy, weight_3)); + } + + int const lane = threadIdx.x & (kWarpSize - 1); + int const warp = threadIdx.x / kWarpSize; + sum_0 = mhc_warp_sum(sum_0); + sum_1 = mhc_warp_sum(sum_1); + sum_2 = mhc_warp_sum(sum_2); + sum_3 = mhc_warp_sum(sum_3); + if (lane == 0) { + warp_sums[0][warp] = sum_0; + warp_sums[1][warp] = sum_1; + warp_sums[2][warp] = sum_2; + warp_sums[3][warp] = sum_3; + } + __syncthreads(); + + if (warp == 0) { + sum_0 = lane < kNumWarps ? warp_sums[0][lane] : 0.0f; + sum_1 = lane < kNumWarps ? warp_sums[1][lane] : 0.0f; + sum_2 = lane < kNumWarps ? warp_sums[2][lane] : 0.0f; + sum_3 = lane < kNumWarps ? warp_sums[3][lane] : 0.0f; + sum_0 = mhc_warp_sum(sum_0); + sum_1 = mhc_warp_sum(sum_1); + sum_2 = mhc_warp_sum(sum_2); + sum_3 = mhc_warp_sum(sum_3); + if (lane == 0) { + grad_pre[token * 4] = sum_0; + grad_pre[token * 4 + 1] = sum_1; + grad_pre[token * 4 + 2] = sum_2; + grad_pre[token * 4 + 3] = sum_3; + } + } + +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + cudaTriggerProgrammaticLaunchCompletion(); +#endif +} + inline cudaError_t launch_mhc_pre_h_aggregate( __nv_bfloat16 const* residual, float const* pre, __nv_bfloat16* output, int64_t num_tokens, int64_t hidden_size, cudaStream_t stream, @@ -104,4 +200,27 @@ inline cudaError_t launch_mhc_pre_h_aggregate( output, hidden_size); } +inline cudaError_t launch_mhc_pre_h_aggregate_backward( + __nv_bfloat16 const* grad_output, __nv_bfloat16 const* residual, + float const* pre, __nv_bfloat16* grad_residual, float* grad_pre, + int64_t num_tokens, int64_t hidden_size, cudaStream_t stream, + bool enable_pdl) { + cudaLaunchConfig_t config{}; + config.gridDim = dim3(static_cast(num_tokens)); + config.blockDim = dim3(kMhcPreHAggregateBackwardThreads); + config.stream = stream; + + cudaLaunchAttribute attribute{}; + if (enable_pdl) { + attribute.id = cudaLaunchAttributeProgrammaticStreamSerialization; + attribute.val.programmaticStreamSerializationAllowed = 1; + config.attrs = &attribute; + config.numAttrs = 1; + } + + return cudaLaunchKernelEx(&config, mhc_pre_h_aggregate_backward_kernel, + grad_output, residual, pre, grad_residual, grad_pre, + hidden_size); +} + } diff --git a/csrc/ops.cpp b/csrc/ops.cpp index 2887dea5..5dd74eab 100644 --- a/csrc/ops.cpp +++ b/csrc/ops.cpp @@ -168,6 +168,11 @@ void reduce_rows_fp32_left_fold_cuda( torch::Tensor mhc_pre_h_aggregate_cuda( torch::Tensor residual, torch::Tensor pre); + +std::vector mhc_pre_h_aggregate_backward_cuda( + torch::Tensor grad_output, + torch::Tensor residual, + torch::Tensor pre); #endif static void rmsnorm_check_input(const torch::Tensor& x, const char* name) { @@ -499,6 +504,10 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { "mhc_pre_h_aggregate", &mhc_pre_h_aggregate_cuda, "Batch-invariant MHC H Aggregate CUDA"); + m.def( + "mhc_pre_h_aggregate_backward", + &mhc_pre_h_aggregate_backward_cuda, + "Batch-invariant MHC H Aggregate backward CUDA"); #endif // registry SiLU / SwiGLU (elementwise activation) diff --git a/rl_engine/_C.pyi b/rl_engine/_C.pyi index f2f64c8c..8508151e 100644 --- a/rl_engine/_C.pyi +++ b/rl_engine/_C.pyi @@ -6,6 +6,11 @@ def mhc_pre_h_aggregate( residual: torch.Tensor, pre: torch.Tensor, ) -> torch.Tensor: ... +def mhc_pre_h_aggregate_backward( + grad_output: torch.Tensor, + residual: torch.Tensor, + pre: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: ... def deterministic_collective_ipc_meta( tensor: torch.Tensor, ) -> tuple[list[int], int]: ... diff --git a/rl_engine/kernels/gtest/gradient_adapters.py b/rl_engine/kernels/gtest/gradient_adapters.py index da172821..f7992607 100644 --- a/rl_engine/kernels/gtest/gradient_adapters.py +++ b/rl_engine/kernels/gtest/gradient_adapters.py @@ -98,9 +98,25 @@ def to_dict(self) -> dict[str, Any]: _DGATE = GradientTensorSpec("dgate", "token", "gate") _DUP = GradientTensorSpec("dup", "token", "up") _DW_LINEAR = GradientTensorSpec("dW", "parameter", "lm_head_weight") +_DRESIDUAL = GradientTensorSpec("dresidual", "token", "residual") +_DPRE = GradientTensorSpec("dpre", "token", "pre") GRADIENT_ADAPTERS: dict[str, GradientAdapterSpec] = { + "mhc_pre_h_aggregate": GradientAdapterSpec( + op_name="mhc_pre_h_aggregate", + chain_node="mhc_pre_h_aggregate", + op_class="reduction", + spec_name="mhc_pre_h_aggregate", + tensors=(_DRESIDUAL, _DPRE), + requirement="optional_fused", + source_files=( + "rl_engine/kernels/ops/cuda/mhc.py", + "rl_engine/kernels/ops/pytorch/mhc.py", + "csrc/cuda/mhc/mhc_pre_h_aggregate.cu", + "csrc/cuda/mhc/mhc_pre_h_aggregate_kernel.cuh", + ), + ), "rms_norm": GradientAdapterSpec( op_name="rms_norm", chain_node="rms_norm", @@ -681,6 +697,18 @@ def _row_inputs( "gate": _stack_rows(keys, leading, (hidden,), device=device, dtype=dtype, offset=0), "up": _stack_rows(keys, leading, (hidden,), device=device, dtype=dtype, offset=1), } + if op_name == "mhc_pre_h_aggregate": + return { + "residual": _stack_rows(keys, leading, (4, hidden), device=device, dtype=dtype), + "pre": _stack_rows( + keys, + leading, + (4,), + device=device, + dtype=torch.float32, + offset=1, + ), + } if op_name == "det_gemm": return { "a": _stack_rows(keys, leading, (hidden,), device=device, dtype=dtype), diff --git a/rl_engine/kernels/gtest/operator_inputs.py b/rl_engine/kernels/gtest/operator_inputs.py index 34de892a..b78c5cc5 100644 --- a/rl_engine/kernels/gtest/operator_inputs.py +++ b/rl_engine/kernels/gtest/operator_inputs.py @@ -25,6 +25,7 @@ def make_operator_inputs( device: torch.device, ) -> dict[str, Any]: builders = { + "mhc_pre_h_aggregate": _make_mhc_pre_h_aggregate_inputs, "rms_norm": _make_rms_norm_inputs, "qk_norm": _make_qk_norm_inputs, "pack": _make_pack_inputs, @@ -52,6 +53,7 @@ def operator_shape_name(op_name: str, args: argparse.Namespace) -> str: batch, seq = _batch_seq(args) vocab = _arg_int(args, "vocab", DEFAULT_VOCAB) names = { + "mhc_pre_h_aggregate": f"{batch * seq}x4x{_normalized_dim(args)}", "rms_norm": f"{batch}x{seq}x{_normalized_dim(args)}", "qk_norm": f"{batch}x{seq}x{_arg_int(args, 'n_heads', DEFAULT_N_HEADS)}x" f"{_arg_int(args, 'head_dim', DEFAULT_HEAD_DIM)}", @@ -88,6 +90,18 @@ def _make_rms_norm_inputs( } +def _make_mhc_pre_h_aggregate_inputs( + args: argparse.Namespace, dtype: torch.dtype, device: torch.device +) -> dict[str, Any]: + batch, seq = _batch_seq(args) + num_tokens = batch * seq + hidden = _normalized_dim(args) + return { + "residual": _floating_tensor((num_tokens, 4, hidden), args, dtype, device, offset=0), + "pre": _floating_tensor((num_tokens, 4), args, torch.float32, device, offset=1), + } + + def _make_qk_norm_inputs( args: argparse.Namespace, dtype: torch.dtype, device: torch.device ) -> dict[str, Any]: diff --git a/rl_engine/kernels/gtest/operator_specs.py b/rl_engine/kernels/gtest/operator_specs.py index ca4a462e..7c7bf9b3 100644 --- a/rl_engine/kernels/gtest/operator_specs.py +++ b/rl_engine/kernels/gtest/operator_specs.py @@ -32,6 +32,17 @@ def _load_object(path: str) -> Any: OP_SPECS = { + "mhc_pre_h_aggregate": OperatorSpec( + name="mhc_pre_h_aggregate", + op_class="reduction", + gold_path=("rl_engine.kernels.ops.pytorch.mhc.NativeMHCPreHAggregateOp"), + gold_method="forward_fp32", + candidate_paths={ + "pytorch": ("rl_engine.kernels.ops.pytorch.mhc.NativeMHCPreHAggregateOp"), + "cuda": "rl_engine.kernels.ops.cuda.mhc.MHCPreHAggregateCudaOp", + }, + grad_input_names=("residual", "pre"), + ), "rms_norm": OperatorSpec( name="rms_norm", op_class="reduction", @@ -70,8 +81,7 @@ def _load_object(path: str) -> Any: "TritonBatchInvariantAttentionOp" ), "cuda": ( - "rl_engine.kernels.ops.cuda.attention.deterministic_attn." - "DeterministicAttentionOp" + "rl_engine.kernels.ops.cuda.attention.deterministic_attn.DeterministicAttentionOp" ), }, grad_input_names=("q", "k", "v"), diff --git a/rl_engine/kernels/ops/cuda/mhc.py b/rl_engine/kernels/ops/cuda/mhc.py new file mode 100644 index 00000000..88143bdb --- /dev/null +++ b/rl_engine/kernels/ops/cuda/mhc.py @@ -0,0 +1,60 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from __future__ import annotations + +import torch + +from rl_engine.kernels.ops.backward_runtime import record_backward +from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE + + +def _require_mhc_extension() -> None: + required = ("mhc_pre_h_aggregate", "mhc_pre_h_aggregate_backward") + if not _EXT_AVAILABLE or _C is None or not all(hasattr(_C, name) for name in required): + raise RuntimeError("MHC H Aggregate CUDA symbols are unavailable; rebuild rl_engine._C") + + +class _MHCPreHAggregateFunction(torch.autograd.Function): + @staticmethod + def forward(ctx, residual: torch.Tensor, pre: torch.Tensor) -> torch.Tensor: + residual = residual.contiguous() + pre = pre.contiguous() + output = _C.mhc_pre_h_aggregate(residual, pre) + ctx.save_for_backward(residual, pre) + return output + + @staticmethod + def backward(ctx, grad_output: torch.Tensor): + residual, pre = ctx.saved_tensors + grad_residual = grad_pre = None + if ctx.needs_input_grad[0] or ctx.needs_input_grad[1]: + grads = _C.mhc_pre_h_aggregate_backward(grad_output.contiguous(), residual, pre) + if ctx.needs_input_grad[0]: + grad_residual = grads[0] + if ctx.needs_input_grad[1]: + grad_pre = grads[1] + record_backward( + "mhc_pre_h_aggregate", + kernel_id="rl_engine._C.mhc_pre_h_aggregate_backward", + impl="cuda_fixed_tree_mhc_pre_h_aggregate_backward", + family="cuda", + ) + return grad_residual, grad_pre + + +class MHCPreHAggregateCudaOp: + """Autograd-enabled CUDA MHC weighted collapse.""" + + op_class = "reduction" + is_batch_invariant = True + backward_impl = "cuda_fixed_tree_mhc_pre_h_aggregate_backward" + + def __init__(self) -> None: + _require_mhc_extension() + + def __call__(self, residual: torch.Tensor, pre: torch.Tensor) -> torch.Tensor: + return self.forward(residual, pre) + + def forward(self, residual: torch.Tensor, pre: torch.Tensor) -> torch.Tensor: + return _MHCPreHAggregateFunction.apply(residual, pre) diff --git a/rl_engine/kernels/ops/pytorch/mhc.py b/rl_engine/kernels/ops/pytorch/mhc.py new file mode 100644 index 00000000..ac8c6091 --- /dev/null +++ b/rl_engine/kernels/ops/pytorch/mhc.py @@ -0,0 +1,40 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from __future__ import annotations + +import torch + + +class NativeMHCPreHAggregateOp: + """PyTorch reference for the four-stream MHC weighted collapse.""" + + op_class = "reduction" + + def __call__(self, residual: torch.Tensor, pre: torch.Tensor) -> torch.Tensor: + return self.forward(residual, pre) + + def forward(self, residual: torch.Tensor, pre: torch.Tensor) -> torch.Tensor: + return self._compute(residual, pre).to(residual.dtype) + + def forward_fp32(self, residual: torch.Tensor, pre: torch.Tensor) -> torch.Tensor: + return self._compute(residual, pre) + + @staticmethod + def _compute(residual: torch.Tensor, pre: torch.Tensor) -> torch.Tensor: + if residual.dim() != 3 or residual.shape[1] != 4: + raise ValueError("residual must have shape [num_tokens, 4, hidden_size]") + if pre.shape != residual.shape[:2]: + raise ValueError("pre must have shape [num_tokens, 4]") + if residual.device != pre.device: + raise RuntimeError("residual and pre must be on the same device") + + residual_fp32 = residual.float() + pre_fp32 = pre.float() + left = ( + pre_fp32[:, 0, None] * residual_fp32[:, 0] + pre_fp32[:, 1, None] * residual_fp32[:, 1] + ) + right = ( + pre_fp32[:, 2, None] * residual_fp32[:, 2] + pre_fp32[:, 3, None] * residual_fp32[:, 3] + ) + return left + right diff --git a/tests/test_gradient_invariance.py b/tests/test_gradient_invariance.py index e97a343f..a1f638bd 100644 --- a/tests/test_gradient_invariance.py +++ b/tests/test_gradient_invariance.py @@ -588,6 +588,10 @@ def test_required_ops_are_enumerable(self): assert adapter.shape_dependent_bwd_accum == "forbidden" def test_stable_grad_names(self): + assert tuple(tensor.name for tensor in get_adapter("mhc_pre_h_aggregate").tensors) == ( + "dresidual", + "dpre", + ) assert tuple(t.name for t in get_adapter("rms_norm").tensors) == ("dx", "dweight") assert tuple(t.name for t in get_adapter("det_gemm").tensors) == ("dX", "dW") assert tuple(t.name for t in get_adapter("attention").tensors) == ("dQ", "dK", "dV") diff --git a/tests/test_mhc_pre_h_aggregate.py b/tests/test_mhc_pre_h_aggregate.py index 6563e6a5..6c86edd2 100644 --- a/tests/test_mhc_pre_h_aggregate.py +++ b/tests/test_mhc_pre_h_aggregate.py @@ -58,15 +58,85 @@ def test_mhc_pre_h_aggregate_is_batch_invariant_and_matches_pytorch(): token_by_token = torch.cat( [ - _C.mhc_pre_h_aggregate( - residual[token : token + 1], pre[token : token + 1] - ) + _C.mhc_pre_h_aggregate(residual[token : token + 1], pre[token : token + 1]) for token in range(num_tokens) ] ) assert _same_bytes(first, token_by_token) - reference = torch.sum( - pre.unsqueeze(-1) * residual.to(torch.float32), dim=1 - ).to(torch.bfloat16) + reference = torch.sum(pre.unsqueeze(-1) * residual.to(torch.float32), dim=1).to(torch.bfloat16) torch.testing.assert_close(first, reference, atol=5e-2, rtol=1e-2) + + +@requires_mhc_kernel +def test_mhc_pre_h_aggregate_backward_is_batch_invariant_and_matches_pytorch(): + from rl_engine.kernels.ops.base import _C + from rl_engine.kernels.ops.cuda.mhc import MHCPreHAggregateCudaOp + from rl_engine.kernels.ops.pytorch.mhc import NativeMHCPreHAggregateOp + + num_tokens = 129 + hidden_size = 4096 + generator = torch.Generator(device="cuda").manual_seed(1) + residual = torch.randn( + num_tokens, + 4, + hidden_size, + dtype=torch.bfloat16, + device="cuda", + generator=generator, + ) + pre = torch.rand( + num_tokens, + 4, + dtype=torch.float32, + device="cuda", + generator=generator, + ) + grad_output = torch.randn( + num_tokens, + hidden_size, + dtype=torch.bfloat16, + device="cuda", + generator=generator, + ) + + first = _C.mhc_pre_h_aggregate_backward(grad_output, residual, pre) + for _ in range(99): + repeated = _C.mhc_pre_h_aggregate_backward(grad_output, residual, pre) + assert _same_bytes(first[0], repeated[0]) + assert _same_bytes(first[1], repeated[1]) + + split = [ + _C.mhc_pre_h_aggregate_backward( + grad_output[token : token + 1], + residual[token : token + 1], + pre[token : token + 1], + ) + for token in range(num_tokens) + ] + split_grad_residual = torch.cat([grads[0] for grads in split]) + split_grad_pre = torch.cat([grads[1] for grads in split]) + assert _same_bytes(first[0], split_grad_residual) + assert _same_bytes(first[1], split_grad_pre) + + candidate_residual = residual.detach().requires_grad_(True) + candidate_pre = pre.detach().requires_grad_(True) + candidate_output = MHCPreHAggregateCudaOp()(candidate_residual, candidate_pre) + candidate_grads = torch.autograd.grad( + candidate_output, + (candidate_residual, candidate_pre), + grad_outputs=grad_output, + ) + assert _same_bytes(first[0], candidate_grads[0]) + assert _same_bytes(first[1], candidate_grads[1]) + + reference_residual = residual.detach().requires_grad_(True) + reference_pre = pre.detach().requires_grad_(True) + reference_output = NativeMHCPreHAggregateOp()(reference_residual, reference_pre) + reference_grads = torch.autograd.grad( + reference_output, + (reference_residual, reference_pre), + grad_outputs=grad_output, + ) + torch.testing.assert_close(candidate_grads[0], reference_grads[0], atol=5e-2, rtol=2e-2) + torch.testing.assert_close(candidate_grads[1], reference_grads[1], atol=5e-2, rtol=2e-2) diff --git a/tests/test_operator_inputs.py b/tests/test_operator_inputs.py index 3b92af3b..27105032 100644 --- a/tests/test_operator_inputs.py +++ b/tests/test_operator_inputs.py @@ -41,6 +41,7 @@ def _args(**overrides): @pytest.mark.parametrize( "op_name", [ + "mhc_pre_h_aggregate", "rms_norm", "qk_norm", "pack", @@ -74,6 +75,28 @@ def test_constant_logp_inputs_are_deterministic(): assert torch.equal(inputs["token_ids"], torch.full((1, 2), 3, dtype=torch.long)) +def test_mhc_pre_h_aggregate_inputs_match_mixed_precision_contract(): + args = _args(input_mode="constant", batch=2, seq=3, normalized_dim=128) + inputs = make_operator_inputs("mhc_pre_h_aggregate", args, torch.bfloat16, torch.device("cpu")) + + assert inputs["residual"].shape == (6, 4, 128) + assert inputs["residual"].dtype is torch.bfloat16 + assert inputs["pre"].shape == (6, 4) + assert inputs["pre"].dtype is torch.float32 + assert operator_shape_name("mhc_pre_h_aggregate", args) == "6x4x128" + + +def test_mhc_pre_h_aggregate_operator_spec_registers_both_gradients(): + args = _args(op="mhc_pre_h_aggregate", candidate="pytorch") + + case = make_operator_case(args, torch.bfloat16, torch.device("cpu")) + candidate = make_candidate(args) + + assert case.op_class == "reduction" + assert case.grad_input_names == ("residual", "pre") + assert candidate.name == "pytorch-mhc_pre_h_aggregate" + + def test_constant_batch_invariant_logp_inputs_match_operator_contract(): args = _args(input_mode="constant", constant_value=0.5, token_value=3) inputs = make_operator_inputs("batch_invariant_logp", args, torch.float32, torch.device("cpu")) From 38e3f825d4b7a8c0ac5a24349c54d08bc3791559 Mon Sep 17 00:00:00 2001 From: nodeeeeee Date: Tue, 1 Sep 2026 15:02:30 +0000 Subject: [PATCH 3/3] fix: align MHC H Aggregate with FP32 contract Signed-off-by: nodeeeeee --- csrc/cuda/mhc/mhc_pre_h_aggregate.cu | 28 ++- csrc/cuda/mhc/mhc_pre_h_aggregate_kernel.cuh | 48 ++-- csrc/ops.cpp | 2 +- rl_engine/_C.pyi | 2 +- rl_engine/kernels/gtest/gradient_adapters.py | 78 +++++-- rl_engine/kernels/gtest/operator_inputs.py | 9 +- rl_engine/kernels/ops/cuda/mhc.py | 51 ++--- rl_engine/kernels/ops/pytorch/mhc.py | 101 ++++++++- setup.py | 10 +- tests/test_mhc_pre_h_aggregate.py | 218 ++++++++++++------- tests/test_operator_inputs.py | 4 +- 11 files changed, 376 insertions(+), 175 deletions(-) diff --git a/csrc/cuda/mhc/mhc_pre_h_aggregate.cu b/csrc/cuda/mhc/mhc_pre_h_aggregate.cu index 1f8658fc..960952a2 100644 --- a/csrc/cuda/mhc/mhc_pre_h_aggregate.cu +++ b/csrc/cuda/mhc/mhc_pre_h_aggregate.cu @@ -24,9 +24,12 @@ torch::Tensor mhc_pre_h_aggregate_cuda(torch::Tensor residual, "residual must be bfloat16"); TORCH_CHECK(pre.scalar_type() == torch::kFloat32, "pre must be float32"); - TORCH_CHECK(residual.dim() == 3 && residual.size(1) == 4, - "residual must have shape [num_tokens, 4, hidden_size]"); - TORCH_CHECK(pre.dim() == 2 && pre.size(1) == 4, + TORCH_CHECK(residual.dim() == 3 && + residual.size(1) == rl_kernel::mhc::kMhcPreHcMult && + residual.size(2) == rl_kernel::mhc::kMhcPreHiddenSize, + "residual must have shape [num_tokens, 4, 4096]"); + TORCH_CHECK(pre.dim() == 2 && + pre.size(1) == rl_kernel::mhc::kMhcPreHcMult, "pre must have shape [num_tokens, 4]"); TORCH_CHECK(residual.size(0) == pre.size(0), "residual and pre must have the same num_tokens"); @@ -76,14 +79,17 @@ std::vector mhc_pre_h_aggregate_backward_cuda( residual.scalar_type() == torch::kBFloat16, "grad_output and residual must be bfloat16"); TORCH_CHECK(pre.scalar_type() == torch::kFloat32, "pre must be float32"); - TORCH_CHECK(residual.dim() == 3 && residual.size(1) == 4, - "residual must have shape [num_tokens, 4, hidden_size]"); - TORCH_CHECK(pre.dim() == 2 && pre.size(1) == 4, + TORCH_CHECK(residual.dim() == 3 && + residual.size(1) == rl_kernel::mhc::kMhcPreHcMult && + residual.size(2) == rl_kernel::mhc::kMhcPreHiddenSize, + "residual must have shape [num_tokens, 4, 4096]"); + TORCH_CHECK(pre.dim() == 2 && + pre.size(1) == rl_kernel::mhc::kMhcPreHcMult, "pre must have shape [num_tokens, 4]"); TORCH_CHECK(grad_output.dim() == 2 && grad_output.size(0) == residual.size(0) && - grad_output.size(1) == residual.size(2), - "grad_output must have shape [num_tokens, hidden_size]"); + grad_output.size(1) == rl_kernel::mhc::kMhcPreHiddenSize, + "grad_output must have shape [num_tokens, 4096]"); TORCH_CHECK(pre.size(0) == residual.size(0), "residual and pre must have the same num_tokens"); @@ -92,7 +98,8 @@ std::vector mhc_pre_h_aggregate_backward_cuda( TORCH_CHECK(num_tokens <= std::numeric_limits::max(), "num_tokens exceeds the CUDA grid limit"); - auto grad_residual = torch::empty_like(residual); + auto grad_residual = torch::empty( + residual.sizes(), residual.options().dtype(torch::kFloat32)); auto grad_pre = torch::zeros_like(pre); if (num_tokens == 0 || hidden_size == 0) { return {grad_residual, grad_pre}; @@ -112,8 +119,7 @@ std::vector mhc_pre_h_aggregate_backward_cuda( auto const* residual_ptr = reinterpret_cast<__nv_bfloat16 const*>( residual.data_ptr()); auto const* pre_ptr = pre.data_ptr(); - auto* grad_residual_ptr = reinterpret_cast<__nv_bfloat16*>( - grad_residual.data_ptr()); + auto* grad_residual_ptr = grad_residual.data_ptr(); auto* grad_pre_ptr = grad_pre.data_ptr(); cudaError_t const status = diff --git a/csrc/cuda/mhc/mhc_pre_h_aggregate_kernel.cuh b/csrc/cuda/mhc/mhc_pre_h_aggregate_kernel.cuh index ef864b28..9039c856 100644 --- a/csrc/cuda/mhc/mhc_pre_h_aggregate_kernel.cuh +++ b/csrc/cuda/mhc/mhc_pre_h_aggregate_kernel.cuh @@ -7,9 +7,23 @@ namespace rl_kernel::mhc { +constexpr int64_t kMhcPreHcMult = 4; +constexpr int64_t kMhcPreHiddenSize = 4096; +constexpr int kMhcWarpSize = 32; constexpr int kMhcPreHAggregateDecodeThreads = 1024; constexpr int kMhcPreHAggregateBatchThreads = 512; constexpr int kMhcPreHAggregateBackwardThreads = 256; +constexpr int kMhcPreHAggregateBackwardWarps = + kMhcPreHAggregateBackwardThreads / kMhcWarpSize; + +static_assert(kMhcPreHAggregateBackwardThreads % kMhcWarpSize == 0, + "MHC H Aggregate backward requires complete warps"); +static_assert(kMhcPreHiddenSize % kMhcPreHAggregateBackwardThreads == 0, + "each backward thread must reduce a fixed number of elements"); +static_assert(kMhcPreHAggregateBackwardThreads <= 1024, + "MHC H Aggregate backward exceeds the CUDA block limit"); +static_assert(kMhcPreHAggregateBackwardWarps <= kMhcWarpSize, + "warp 0 must be able to reduce all per-warp partial sums"); __global__ void mhc_pre_h_aggregate_kernel(__nv_bfloat16 const* residual, float const* pre, @@ -93,15 +107,14 @@ __device__ __forceinline__ float mhc_warp_sum(float value) { __global__ void mhc_pre_h_aggregate_backward_kernel( __nv_bfloat16 const* grad_output, __nv_bfloat16 const* residual, - float const* pre, __nv_bfloat16* grad_residual, float* grad_pre, + float const* pre, float* grad_residual, float* grad_pre, int64_t hidden_size) { #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) cudaGridDependencySynchronize(); #endif - constexpr int kWarpSize = 32; - constexpr int kNumWarps = kMhcPreHAggregateBackwardThreads / kWarpSize; - __shared__ float warp_sums[4][kNumWarps]; + __shared__ float + warp_sums[kMhcPreHcMult][kMhcPreHAggregateBackwardWarps]; int64_t const token = static_cast(blockIdx.x); int64_t const output_offset = token * hidden_size; @@ -132,18 +145,17 @@ __global__ void mhc_pre_h_aggregate_backward_kernel( sum_2 = __fadd_rn(sum_2, __fmul_rn(dy, residual_2)); sum_3 = __fadd_rn(sum_3, __fmul_rn(dy, residual_3)); - grad_residual[residual_offset + hidden] = - __float2bfloat16_rn(__fmul_rn(dy, weight_0)); + grad_residual[residual_offset + hidden] = __fmul_rn(dy, weight_0); grad_residual[residual_offset + hidden_size + hidden] = - __float2bfloat16_rn(__fmul_rn(dy, weight_1)); + __fmul_rn(dy, weight_1); grad_residual[residual_offset + 2 * hidden_size + hidden] = - __float2bfloat16_rn(__fmul_rn(dy, weight_2)); + __fmul_rn(dy, weight_2); grad_residual[residual_offset + 3 * hidden_size + hidden] = - __float2bfloat16_rn(__fmul_rn(dy, weight_3)); + __fmul_rn(dy, weight_3); } - int const lane = threadIdx.x & (kWarpSize - 1); - int const warp = threadIdx.x / kWarpSize; + int const lane = threadIdx.x & (kMhcWarpSize - 1); + int const warp = threadIdx.x / kMhcWarpSize; sum_0 = mhc_warp_sum(sum_0); sum_1 = mhc_warp_sum(sum_1); sum_2 = mhc_warp_sum(sum_2); @@ -157,10 +169,14 @@ __global__ void mhc_pre_h_aggregate_backward_kernel( __syncthreads(); if (warp == 0) { - sum_0 = lane < kNumWarps ? warp_sums[0][lane] : 0.0f; - sum_1 = lane < kNumWarps ? warp_sums[1][lane] : 0.0f; - sum_2 = lane < kNumWarps ? warp_sums[2][lane] : 0.0f; - sum_3 = lane < kNumWarps ? warp_sums[3][lane] : 0.0f; + sum_0 = lane < kMhcPreHAggregateBackwardWarps ? warp_sums[0][lane] + : 0.0f; + sum_1 = lane < kMhcPreHAggregateBackwardWarps ? warp_sums[1][lane] + : 0.0f; + sum_2 = lane < kMhcPreHAggregateBackwardWarps ? warp_sums[2][lane] + : 0.0f; + sum_3 = lane < kMhcPreHAggregateBackwardWarps ? warp_sums[3][lane] + : 0.0f; sum_0 = mhc_warp_sum(sum_0); sum_1 = mhc_warp_sum(sum_1); sum_2 = mhc_warp_sum(sum_2); @@ -202,7 +218,7 @@ inline cudaError_t launch_mhc_pre_h_aggregate( inline cudaError_t launch_mhc_pre_h_aggregate_backward( __nv_bfloat16 const* grad_output, __nv_bfloat16 const* residual, - float const* pre, __nv_bfloat16* grad_residual, float* grad_pre, + float const* pre, float* grad_residual, float* grad_pre, int64_t num_tokens, int64_t hidden_size, cudaStream_t stream, bool enable_pdl) { cudaLaunchConfig_t config{}; diff --git a/csrc/ops.cpp b/csrc/ops.cpp index 5dd74eab..809715b7 100644 --- a/csrc/ops.cpp +++ b/csrc/ops.cpp @@ -507,7 +507,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def( "mhc_pre_h_aggregate_backward", &mhc_pre_h_aggregate_backward_cuda, - "Batch-invariant MHC H Aggregate backward CUDA"); + "Batch-invariant MHC H Aggregate backward CUDA (FP32 dR and dPRE)"); #endif // registry SiLU / SwiGLU (elementwise activation) diff --git a/rl_engine/_C.pyi b/rl_engine/_C.pyi index 8508151e..46c83ac6 100644 --- a/rl_engine/_C.pyi +++ b/rl_engine/_C.pyi @@ -10,7 +10,7 @@ def mhc_pre_h_aggregate_backward( grad_output: torch.Tensor, residual: torch.Tensor, pre: torch.Tensor, -) -> tuple[torch.Tensor, torch.Tensor]: ... +) -> tuple[torch.Tensor, torch.Tensor]: ... # (dR FP32, dPRE FP32) def deterministic_collective_ipc_meta( tensor: torch.Tensor, ) -> tuple[list[int], int]: ... diff --git a/rl_engine/kernels/gtest/gradient_adapters.py b/rl_engine/kernels/gtest/gradient_adapters.py index f7992607..cfad6fbb 100644 --- a/rl_engine/kernels/gtest/gradient_adapters.py +++ b/rl_engine/kernels/gtest/gradient_adapters.py @@ -24,6 +24,7 @@ ) from rl_engine.kernels.gtest.operator_specs import OP_SPECS, _load_object from rl_engine.kernels.gtest.tolerance import normalize_dtype_name +from rl_engine.kernels.ops.pytorch.mhc import MHC_PRE_HIDDEN_SIZE from rl_engine.testing.ws1_workload import ( PaddedBatch, PhysicalLayout, @@ -535,7 +536,11 @@ def make_gradient_runner( def run(config: ConfigSpec, **kwargs: Any) -> dict[str, torch.Tensor] | GradientObservation: denom = int(kwargs["active_token_denominator"]) - exec_dtype = torch.float32 if reference else dtype + exec_dtype = ( + torch.bfloat16 + if op_name == "mhc_pre_h_aggregate" + else torch.float32 if reference else dtype + ) grads = _run_adapter( adapter, operator, @@ -596,7 +601,11 @@ def run( config: ConfigSpec, **kwargs: Any ) -> dict[tuple[str, int], torch.Tensor] | RuntimeObservation: del kwargs - exec_dtype = torch.float32 if reference else dtype + exec_dtype = ( + torch.bfloat16 + if op_name == "mhc_pre_h_aggregate" + else torch.float32 if reference else dtype + ) outputs = _run_forward( adapter, operator, @@ -699,7 +708,9 @@ def _row_inputs( } if op_name == "mhc_pre_h_aggregate": return { - "residual": _stack_rows(keys, leading, (4, hidden), device=device, dtype=dtype), + "residual": _stack_rows( + keys, leading, (4, MHC_PRE_HIDDEN_SIZE), device=device, dtype=dtype + ), "pre": _stack_rows( keys, leading, @@ -825,25 +836,48 @@ def _run_row_stream( n_heads=n_heads, head_dim=head_dim, ) - prepared = _requires_grad_inputs(inputs, [spec.source_input for spec in specs]) - raw = _require_differentiable( - adapter.op_name, _first_output(_call_operator(operator, prepared)) - ) - out_rows = _to_rows(adapter.op_name, raw, length) - upstream = _scaled_upstream( - keys, - tokens, - tuple(out_rows.shape[1:]), - active_token_denominator=active_token_denominator, - device=device, - dtype=out_rows.dtype, - ) - grads = torch.autograd.grad( - out_rows, - [prepared[spec.source_input] for spec in specs], - grad_outputs=upstream, - allow_unused=True, - ) + if adapter.op_name == "mhc_pre_h_aggregate": + prepared = inputs + raw = _first_output(_call_operator(operator, prepared)) + out_rows = _to_rows(adapter.op_name, raw, length) + upstream = _scaled_upstream( + keys, + tokens, + tuple(out_rows.shape[1:]), + active_token_denominator=active_token_denominator, + device=device, + dtype=out_rows.dtype, + ) + backward_fp32 = getattr(operator, "backward_fp32", None) + if not callable(backward_fp32): + raise MissingBackwardError( + adapter.op_name, "candidate has no explicit FP32 backward" + ) + grads = backward_fp32( + upstream, + prepared["residual"], + prepared["pre"], + ) + else: + prepared = _requires_grad_inputs(inputs, [spec.source_input for spec in specs]) + raw = _require_differentiable( + adapter.op_name, _first_output(_call_operator(operator, prepared)) + ) + out_rows = _to_rows(adapter.op_name, raw, length) + upstream = _scaled_upstream( + keys, + tokens, + tuple(out_rows.shape[1:]), + active_token_denominator=active_token_denominator, + device=device, + dtype=out_rows.dtype, + ) + grads = torch.autograd.grad( + out_rows, + [prepared[spec.source_input] for spec in specs], + grad_outputs=upstream, + allow_unused=True, + ) contribution_fn = getattr(operator, "parameter_vjp_contributions_fp32", None) contributions = ( contribution_fn(**prepared, grad_output=upstream) if callable(contribution_fn) else None diff --git a/rl_engine/kernels/gtest/operator_inputs.py b/rl_engine/kernels/gtest/operator_inputs.py index b78c5cc5..9b30f0b3 100644 --- a/rl_engine/kernels/gtest/operator_inputs.py +++ b/rl_engine/kernels/gtest/operator_inputs.py @@ -8,6 +8,8 @@ import torch +from rl_engine.kernels.ops.pytorch.mhc import MHC_PRE_HIDDEN_SIZE + DEFAULT_HIDDEN = 4096 DEFAULT_N_HEADS = 32 DEFAULT_N_KV_HEADS = 8 @@ -53,7 +55,7 @@ def operator_shape_name(op_name: str, args: argparse.Namespace) -> str: batch, seq = _batch_seq(args) vocab = _arg_int(args, "vocab", DEFAULT_VOCAB) names = { - "mhc_pre_h_aggregate": f"{batch * seq}x4x{_normalized_dim(args)}", + "mhc_pre_h_aggregate": f"{batch * seq}x4x{MHC_PRE_HIDDEN_SIZE}", "rms_norm": f"{batch}x{seq}x{_normalized_dim(args)}", "qk_norm": f"{batch}x{seq}x{_arg_int(args, 'n_heads', DEFAULT_N_HEADS)}x" f"{_arg_int(args, 'head_dim', DEFAULT_HEAD_DIM)}", @@ -95,9 +97,10 @@ def _make_mhc_pre_h_aggregate_inputs( ) -> dict[str, Any]: batch, seq = _batch_seq(args) num_tokens = batch * seq - hidden = _normalized_dim(args) return { - "residual": _floating_tensor((num_tokens, 4, hidden), args, dtype, device, offset=0), + "residual": _floating_tensor( + (num_tokens, 4, MHC_PRE_HIDDEN_SIZE), args, dtype, device, offset=0 + ), "pre": _floating_tensor((num_tokens, 4), args, torch.float32, device, offset=1), } diff --git a/rl_engine/kernels/ops/cuda/mhc.py b/rl_engine/kernels/ops/cuda/mhc.py index 88143bdb..6b4860c8 100644 --- a/rl_engine/kernels/ops/cuda/mhc.py +++ b/rl_engine/kernels/ops/cuda/mhc.py @@ -5,7 +5,6 @@ import torch -from rl_engine.kernels.ops.backward_runtime import record_backward from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE @@ -15,36 +14,13 @@ def _require_mhc_extension() -> None: raise RuntimeError("MHC H Aggregate CUDA symbols are unavailable; rebuild rl_engine._C") -class _MHCPreHAggregateFunction(torch.autograd.Function): - @staticmethod - def forward(ctx, residual: torch.Tensor, pre: torch.Tensor) -> torch.Tensor: - residual = residual.contiguous() - pre = pre.contiguous() - output = _C.mhc_pre_h_aggregate(residual, pre) - ctx.save_for_backward(residual, pre) - return output - - @staticmethod - def backward(ctx, grad_output: torch.Tensor): - residual, pre = ctx.saved_tensors - grad_residual = grad_pre = None - if ctx.needs_input_grad[0] or ctx.needs_input_grad[1]: - grads = _C.mhc_pre_h_aggregate_backward(grad_output.contiguous(), residual, pre) - if ctx.needs_input_grad[0]: - grad_residual = grads[0] - if ctx.needs_input_grad[1]: - grad_pre = grads[1] - record_backward( - "mhc_pre_h_aggregate", - kernel_id="rl_engine._C.mhc_pre_h_aggregate_backward", - impl="cuda_fixed_tree_mhc_pre_h_aggregate_backward", - family="cuda", - ) - return grad_residual, grad_pre - - class MHCPreHAggregateCudaOp: - """Autograd-enabled CUDA MHC weighted collapse.""" + """Explicit forward/backward CUDA MHC weighted collapse. + + ``backward_fp32`` exposes the aggregate-only composite boundary without + downcasting ``dR_from_aggregate`` before the controller-gradient merge. + This operator intentionally has no standalone autograd path. + """ op_class = "reduction" is_batch_invariant = True @@ -57,4 +33,17 @@ def __call__(self, residual: torch.Tensor, pre: torch.Tensor) -> torch.Tensor: return self.forward(residual, pre) def forward(self, residual: torch.Tensor, pre: torch.Tensor) -> torch.Tensor: - return _MHCPreHAggregateFunction.apply(residual, pre) + if torch.is_grad_enabled() and (residual.requires_grad or pre.requires_grad): + raise RuntimeError( + "MHC H Aggregate does not expose standalone autograd; " + "use the explicit FP32 composite backward" + ) + return _C.mhc_pre_h_aggregate(residual, pre) + + def backward_fp32( + self, + grad_output: torch.Tensor, + residual: torch.Tensor, + pre: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + return tuple(_C.mhc_pre_h_aggregate_backward(grad_output, residual, pre)) diff --git a/rl_engine/kernels/ops/pytorch/mhc.py b/rl_engine/kernels/ops/pytorch/mhc.py index ac8c6091..10e0d682 100644 --- a/rl_engine/kernels/ops/pytorch/mhc.py +++ b/rl_engine/kernels/ops/pytorch/mhc.py @@ -5,6 +5,11 @@ import torch +MHC_PRE_HC_MULT = 4 +MHC_PRE_HIDDEN_SIZE = 4096 +MHC_PRE_H_AGGREGATE_BACKWARD_THREADS = 256 +MHC_WARP_SIZE = 32 + class NativeMHCPreHAggregateOp: """PyTorch reference for the four-stream MHC weighted collapse.""" @@ -20,21 +25,93 @@ def forward(self, residual: torch.Tensor, pre: torch.Tensor) -> torch.Tensor: def forward_fp32(self, residual: torch.Tensor, pre: torch.Tensor) -> torch.Tensor: return self._compute(residual, pre) + def backward_fp32( + self, + grad_output: torch.Tensor, + residual: torch.Tensor, + pre: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Explicit FP32 backward oracle matching the CUDA reduction tree.""" + + self._validate_inputs(residual, pre) + if grad_output.shape != (residual.shape[0], MHC_PRE_HIDDEN_SIZE): + raise ValueError("grad_output must have shape [num_tokens, 4096]") + if grad_output.dtype is not torch.bfloat16: + raise TypeError("grad_output must be bfloat16") + if grad_output.device != residual.device: + raise RuntimeError("grad_output, residual, and pre must be on the same device") + if not grad_output.is_contiguous(): + raise ValueError("grad_output, residual, and pre must be contiguous") + + grad_output_fp32 = grad_output.float() + grad_residual = grad_output_fp32[:, None, :] * pre[:, :, None] + + products = grad_output_fp32[:, None, :] * residual.float() + values_per_thread = MHC_PRE_HIDDEN_SIZE // MHC_PRE_H_AGGREGATE_BACKWARD_THREADS + thread_products = products.reshape( + residual.shape[0], + MHC_PRE_HC_MULT, + values_per_thread, + MHC_PRE_H_AGGREGATE_BACKWARD_THREADS, + ) + thread_sums = torch.zeros_like(thread_products[:, :, 0]) + for index in range(values_per_thread): + thread_sums = thread_sums + thread_products[:, :, index] + + num_warps = MHC_PRE_H_AGGREGATE_BACKWARD_THREADS // MHC_WARP_SIZE + warp_lanes = thread_sums.reshape( + residual.shape[0], MHC_PRE_HC_MULT, num_warps, MHC_WARP_SIZE + ) + warp_sums = self._warp_lane_zero(warp_lanes) + zero_lanes = torch.zeros( + (*warp_sums.shape[:-1], MHC_WARP_SIZE - num_warps), + dtype=torch.float32, + device=residual.device, + ) + block_lanes = torch.cat((warp_sums, zero_lanes), dim=-1) + grad_pre = self._warp_lane_zero(block_lanes) + return grad_residual, grad_pre + @staticmethod def _compute(residual: torch.Tensor, pre: torch.Tensor) -> torch.Tensor: - if residual.dim() != 3 or residual.shape[1] != 4: - raise ValueError("residual must have shape [num_tokens, 4, hidden_size]") - if pre.shape != residual.shape[:2]: + NativeMHCPreHAggregateOp._validate_inputs(residual, pre) + + residual_fp32 = residual.float() + product_0 = pre[:, 0, None] * residual_fp32[:, 0] + product_1 = pre[:, 1, None] * residual_fp32[:, 1] + product_2 = pre[:, 2, None] * residual_fp32[:, 2] + product_3 = pre[:, 3, None] * residual_fp32[:, 3] + left = product_0 + product_1 + right = product_2 + product_3 + return left + right + + @staticmethod + def _validate_inputs(residual: torch.Tensor, pre: torch.Tensor) -> None: + if residual.dim() != 3 or residual.shape[1:] != ( + MHC_PRE_HC_MULT, + MHC_PRE_HIDDEN_SIZE, + ): + raise ValueError("residual must have shape [num_tokens, 4, 4096]") + if pre.shape != (residual.shape[0], MHC_PRE_HC_MULT): raise ValueError("pre must have shape [num_tokens, 4]") + if residual.dtype is not torch.bfloat16: + raise TypeError("residual must be bfloat16") + if pre.dtype is not torch.float32: + raise TypeError("pre must be float32") if residual.device != pre.device: raise RuntimeError("residual and pre must be on the same device") + if not residual.is_contiguous() or not pre.is_contiguous(): + raise ValueError("residual and pre must be contiguous") - residual_fp32 = residual.float() - pre_fp32 = pre.float() - left = ( - pre_fp32[:, 0, None] * residual_fp32[:, 0] + pre_fp32[:, 1, None] * residual_fp32[:, 1] - ) - right = ( - pre_fp32[:, 2, None] * residual_fp32[:, 2] + pre_fp32[:, 3, None] * residual_fp32[:, 3] - ) - return left + right + @staticmethod + def _warp_lane_zero(values: torch.Tensor) -> torch.Tensor: + """Return lane 0 after the CUDA offset=16,8,4,2,1 shuffle tree.""" + + if values.shape[-1] != MHC_WARP_SIZE: + raise ValueError("fixed warp reduction requires exactly 32 lanes") + reduced = values + offset = MHC_WARP_SIZE // 2 + while offset: + reduced = reduced[..., :offset] + reduced[..., offset : 2 * offset] + offset //= 2 + return reduced.squeeze(-1) diff --git a/setup.py b/setup.py index b79aa3ea..c2f801f2 100644 --- a/setup.py +++ b/setup.py @@ -148,8 +148,14 @@ def get_extensions(): cuda_sources.append("csrc/cuda/mhc/mhc_pre_h_aggregate.cu") nvcc_flags = ["-O3", "-Xfatbin", "-compress-all"] - if envs.env_flag(envs.KERNEL_ALIGN_USE_FAST_MATH): - nvcc_flags.append("--use_fast_math") + use_fast_math = envs.env_flag(envs.KERNEL_ALIGN_USE_FAST_MATH) + if use_fast_math and not is_rocm: + raise RuntimeError( + "KERNEL_ALIGN_USE_FAST_MATH is incompatible with the deterministic " + "MHC H Aggregate CUDA contract" + ) + if 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" diff --git a/tests/test_mhc_pre_h_aggregate.py b/tests/test_mhc_pre_h_aggregate.py index 6c86edd2..016543e5 100644 --- a/tests/test_mhc_pre_h_aggregate.py +++ b/tests/test_mhc_pre_h_aggregate.py @@ -4,6 +4,10 @@ import pytest import torch +from rl_engine.kernels.ops.pytorch.mhc import NativeMHCPreHAggregateOp + +HIDDEN_SIZE = 4096 + def _kernel_available() -> bool: if not torch.cuda.is_available(): @@ -14,7 +18,8 @@ def _kernel_available() -> bool: from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE except Exception: return False - return _EXT_AVAILABLE and hasattr(_C, "mhc_pre_h_aggregate") + required = ("mhc_pre_h_aggregate", "mhc_pre_h_aggregate_backward") + return _EXT_AVAILABLE and all(hasattr(_C, name) for name in required) requires_mhc_kernel = pytest.mark.skipif( @@ -24,119 +29,184 @@ def _kernel_available() -> bool: def _same_bytes(left: torch.Tensor, right: torch.Tensor) -> bool: - return torch.equal(left.view(torch.uint8), right.view(torch.uint8)) + return left.dtype == right.dtype and torch.equal( + left.contiguous().view(torch.uint8), right.contiguous().view(torch.uint8) + ) -@requires_mhc_kernel -def test_mhc_pre_h_aggregate_is_batch_invariant_and_matches_pytorch(): - from rl_engine.kernels.ops.base import _C - - num_tokens = 129 - hidden_size = 4096 - num_runs = 100 - generator = torch.Generator(device="cuda").manual_seed(0) +def _make_inputs( + num_tokens: int, *, device: str, seed: int +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + generator = torch.Generator(device=device).manual_seed(seed) residual = torch.randn( num_tokens, 4, - hidden_size, + HIDDEN_SIZE, dtype=torch.bfloat16, - device="cuda", + device=device, generator=generator, ) pre = torch.rand( num_tokens, 4, dtype=torch.float32, - device="cuda", + device=device, generator=generator, ) + grad_output = torch.randn( + num_tokens, + HIDDEN_SIZE, + dtype=torch.bfloat16, + device=device, + generator=generator, + ) + return residual, pre, grad_output + + +def test_native_mhc_pre_h_aggregate_backward_is_explicit_fp32(): + residual, pre, grad_output = _make_inputs(3, device="cpu", seed=7) + + grad_residual, grad_pre = NativeMHCPreHAggregateOp().backward_fp32( + grad_output, residual, pre + ) + + assert grad_residual.dtype is torch.float32 + assert grad_pre.dtype is torch.float32 + expected_grad_residual = grad_output.float()[:, None, :] * pre[:, :, None] + expected_grad_pre = torch.sum( + grad_output.float()[:, None, :] * residual.float(), dim=-1 + ) + assert _same_bytes(grad_residual, expected_grad_residual) + torch.testing.assert_close(grad_pre, expected_grad_pre, atol=5e-4, rtol=1e-5) + + +@requires_mhc_kernel +def test_mhc_pre_h_aggregate_forward_matches_fixed_tree_oracle_and_recomputes(): + from rl_engine.kernels.ops.base import _C + + residual, pre, _grad_output = _make_inputs(129, device="cuda", seed=0) + reference = NativeMHCPreHAggregateOp()(residual, pre) first = _C.mhc_pre_h_aggregate(residual, pre) - for _ in range(num_runs - 1): - repeated = _C.mhc_pre_h_aggregate(residual, pre) - assert _same_bytes(first, repeated) + assert first.dtype is torch.bfloat16 + assert _same_bytes(first, reference) + for _ in range(9): + assert _same_bytes(first, _C.mhc_pre_h_aggregate(residual, pre)) + + +@requires_mhc_kernel +def test_mhc_pre_h_aggregate_forward_is_batch_and_padding_invariant(): + from rl_engine.kernels.ops.base import _C + + residual, pre, _grad_output = _make_inputs(127, device="cuda", seed=1) + pad_residual, pad_pre, _pad_grad_output = _make_inputs(2, device="cuda", seed=2) + padded_residual = torch.cat((pad_residual[:1], residual, pad_residual[1:])) + padded_pre = torch.cat((pad_pre[:1], pre, pad_pre[1:])) + + unpadded = _C.mhc_pre_h_aggregate(residual, pre) + padded = _C.mhc_pre_h_aggregate(padded_residual, padded_pre)[1:-1] token_by_token = torch.cat( [ - _C.mhc_pre_h_aggregate(residual[token : token + 1], pre[token : token + 1]) - for token in range(num_tokens) + _C.mhc_pre_h_aggregate( + residual[token : token + 1], pre[token : token + 1] + ) + for token in range(residual.shape[0]) ] ) - assert _same_bytes(first, token_by_token) - reference = torch.sum(pre.unsqueeze(-1) * residual.to(torch.float32), dim=1).to(torch.bfloat16) - torch.testing.assert_close(first, reference, atol=5e-2, rtol=1e-2) + assert _same_bytes(unpadded, padded) + assert _same_bytes(unpadded, token_by_token) @requires_mhc_kernel -def test_mhc_pre_h_aggregate_backward_is_batch_invariant_and_matches_pytorch(): +def test_mhc_pre_h_aggregate_backward_matches_explicit_fixed_tree_oracle(): from rl_engine.kernels.ops.base import _C - from rl_engine.kernels.ops.cuda.mhc import MHCPreHAggregateCudaOp - from rl_engine.kernels.ops.pytorch.mhc import NativeMHCPreHAggregateOp - num_tokens = 129 - hidden_size = 4096 - generator = torch.Generator(device="cuda").manual_seed(1) - residual = torch.randn( - num_tokens, - 4, - hidden_size, - dtype=torch.bfloat16, - device="cuda", - generator=generator, - ) - pre = torch.rand( - num_tokens, - 4, - dtype=torch.float32, - device="cuda", - generator=generator, - ) - grad_output = torch.randn( - num_tokens, - hidden_size, - dtype=torch.bfloat16, - device="cuda", - generator=generator, + residual, pre, grad_output = _make_inputs(129, device="cuda", seed=3) + reference = NativeMHCPreHAggregateOp().backward_fp32( + grad_output, residual, pre ) first = _C.mhc_pre_h_aggregate_backward(grad_output, residual, pre) - for _ in range(99): + + assert first[0].dtype is torch.float32 + assert first[1].dtype is torch.float32 + assert _same_bytes(first[0], reference[0]) + assert _same_bytes(first[1], reference[1]) + for _ in range(9): repeated = _C.mhc_pre_h_aggregate_backward(grad_output, residual, pre) assert _same_bytes(first[0], repeated[0]) assert _same_bytes(first[1], repeated[1]) + +@requires_mhc_kernel +def test_mhc_pre_h_aggregate_cuda_wrapper_has_no_autograd_fallback(): + from rl_engine.kernels.ops.cuda.mhc import MHCPreHAggregateCudaOp + + residual, pre, _grad_output = _make_inputs(2, device="cuda", seed=8) + residual.requires_grad_(True) + pre.requires_grad_(True) + + with pytest.raises(RuntimeError, match="does not expose standalone autograd"): + MHCPreHAggregateCudaOp()(residual, pre) + + +@requires_mhc_kernel +def test_mhc_pre_h_aggregate_backward_is_batch_and_padding_invariant(): + from rl_engine.kernels.ops.base import _C + + residual, pre, grad_output = _make_inputs(127, device="cuda", seed=4) + pad_residual, pad_pre, pad_grad_output = _make_inputs(2, device="cuda", seed=5) + padded_residual = torch.cat((pad_residual[:1], residual, pad_residual[1:])) + padded_pre = torch.cat((pad_pre[:1], pre, pad_pre[1:])) + padded_grad_output = torch.cat( + (pad_grad_output[:1], grad_output, pad_grad_output[1:]) + ) + + unpadded = _C.mhc_pre_h_aggregate_backward(grad_output, residual, pre) + padded = _C.mhc_pre_h_aggregate_backward( + padded_grad_output, padded_residual, padded_pre + ) split = [ _C.mhc_pre_h_aggregate_backward( grad_output[token : token + 1], residual[token : token + 1], pre[token : token + 1], ) - for token in range(num_tokens) + for token in range(residual.shape[0]) ] split_grad_residual = torch.cat([grads[0] for grads in split]) split_grad_pre = torch.cat([grads[1] for grads in split]) - assert _same_bytes(first[0], split_grad_residual) - assert _same_bytes(first[1], split_grad_pre) - - candidate_residual = residual.detach().requires_grad_(True) - candidate_pre = pre.detach().requires_grad_(True) - candidate_output = MHCPreHAggregateCudaOp()(candidate_residual, candidate_pre) - candidate_grads = torch.autograd.grad( - candidate_output, - (candidate_residual, candidate_pre), - grad_outputs=grad_output, - ) - assert _same_bytes(first[0], candidate_grads[0]) - assert _same_bytes(first[1], candidate_grads[1]) - - reference_residual = residual.detach().requires_grad_(True) - reference_pre = pre.detach().requires_grad_(True) - reference_output = NativeMHCPreHAggregateOp()(reference_residual, reference_pre) - reference_grads = torch.autograd.grad( - reference_output, - (reference_residual, reference_pre), - grad_outputs=grad_output, - ) - torch.testing.assert_close(candidate_grads[0], reference_grads[0], atol=5e-2, rtol=2e-2) - torch.testing.assert_close(candidate_grads[1], reference_grads[1], atol=5e-2, rtol=2e-2) + + assert _same_bytes(unpadded[0], padded[0][1:-1]) + assert _same_bytes(unpadded[1], padded[1][1:-1]) + assert _same_bytes(unpadded[0], split_grad_residual) + assert _same_bytes(unpadded[1], split_grad_pre) + + +@requires_mhc_kernel +def test_mhc_pre_h_aggregate_fails_closed_for_unsupported_contracts(): + from rl_engine.kernels.ops.base import _C + + residual, pre, grad_output = _make_inputs(2, device="cuda", seed=6) + wrong_hidden = residual[:, :, :2048].contiguous() + noncontiguous_residual = torch.empty( + 2, 4, HIDDEN_SIZE * 2, dtype=torch.bfloat16, device="cuda" + )[:, :, ::2] + noncontiguous_grad_output = torch.empty( + 2, HIDDEN_SIZE * 2, dtype=torch.bfloat16, device="cuda" + )[:, ::2] + + with pytest.raises(RuntimeError, match=r"\[num_tokens, 4, 4096\]"): + _C.mhc_pre_h_aggregate(wrong_hidden, pre) + with pytest.raises(RuntimeError, match="residual must be bfloat16"): + _C.mhc_pre_h_aggregate(residual.float(), pre) + with pytest.raises(RuntimeError, match="pre must be float32"): + _C.mhc_pre_h_aggregate(residual, pre.to(torch.bfloat16)) + with pytest.raises(RuntimeError, match="must be contiguous"): + _C.mhc_pre_h_aggregate(noncontiguous_residual, pre) + with pytest.raises(RuntimeError, match="grad_output and residual must be bfloat16"): + _C.mhc_pre_h_aggregate_backward(grad_output.float(), residual, pre) + with pytest.raises(RuntimeError, match="must be contiguous"): + _C.mhc_pre_h_aggregate_backward(noncontiguous_grad_output, residual, pre) diff --git a/tests/test_operator_inputs.py b/tests/test_operator_inputs.py index 27105032..11b1216d 100644 --- a/tests/test_operator_inputs.py +++ b/tests/test_operator_inputs.py @@ -79,11 +79,11 @@ def test_mhc_pre_h_aggregate_inputs_match_mixed_precision_contract(): args = _args(input_mode="constant", batch=2, seq=3, normalized_dim=128) inputs = make_operator_inputs("mhc_pre_h_aggregate", args, torch.bfloat16, torch.device("cpu")) - assert inputs["residual"].shape == (6, 4, 128) + assert inputs["residual"].shape == (6, 4, 4096) assert inputs["residual"].dtype is torch.bfloat16 assert inputs["pre"].shape == (6, 4) assert inputs["pre"].dtype is torch.float32 - assert operator_shape_name("mhc_pre_h_aggregate", args) == "6x4x128" + assert operator_shape_name("mhc_pre_h_aggregate", args) == "6x4x4096" def test_mhc_pre_h_aggregate_operator_spec_registers_both_gradients():