[WS1] Add batch-invariant h_aggregate kernel

benchmarks/benchmark_mhc_pre_h_aggregate.py

@@ -0,0 +1,94 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) 2026 RL-Kernel Contributors
+
+"""Compare explicit MHC forward/backward latency and peak allocated memory."""
+
+import argparse
+import pathlib
+import statistics
+import sys
+from functools import partial
+
+import torch
+
+sys.path.insert(0, str(pathlib.Path(__file__).resolve().parents[1]))
+
+from rl_engine.kernels.ops.cuda.mhc import MHCPreHAggregateCudaOp  # noqa: E402
+from rl_engine.kernels.ops.pytorch.mhc import NativeMHCPreHAggregateOp  # noqa: E402
+from rl_engine.kernels.ops.triton.mhc import MHCPreHAggregateTritonOp  # noqa: E402
+from rl_engine.mhc import oracle  # noqa: E402
+
+
+def assert_bytes(actual, expected):
+    assert actual.dtype == expected.dtype and actual.shape == expected.shape
+    assert torch.equal(actual.contiguous().view(torch.uint8),
+                       expected.contiguous().view(torch.uint8)), "P1 oracle bytes differ"
+
+
+def measure(fn, warmup, repeat):
+    for _ in range(warmup):
+        fn()
+    torch.cuda.synchronize()
+    times = []
+    for _ in range(repeat):
+        start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)
+        start.record()
+        result = fn()
+        end.record()
+        end.synchronize()
+        times.append(start.elapsed_time(end) * 1000)
+        del result
+    torch.cuda.reset_peak_memory_stats()
+    baseline = torch.cuda.memory_allocated()
+    result = fn()
+    torch.cuda.synchronize()
+    extra_mib = (torch.cuda.max_memory_allocated() - baseline) / 2**20
+    del result
+    return statistics.median(times), extra_mib
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--tokens", type=int, nargs="+", default=[1, 16, 128, 129, 1024])
+    parser.add_argument("--hidden", type=int, choices=(128, 4096), default=4096)
+    parser.add_argument("--grad-dtype", choices=("fp32", "bf16"), default="fp32")
+    parser.add_argument("--warmup", type=int, default=3)
+    parser.add_argument("--repeat", type=int, default=10)
+    args = parser.parse_args()
+    if min(args.tokens) <= 0 or args.warmup < 0 or args.repeat <= 0:
+        parser.error("tokens and repeat must be positive; warmup must be nonnegative")
+    ops = {"PyTorch": NativeMHCPreHAggregateOp(), "CUDA": MHCPreHAggregateCudaOp(),
+           "Triton": MHCPreHAggregateTritonOp()}
+    grad_dtype = {"fp32": torch.float32, "bf16": torch.bfloat16}[args.grad_dtype]
+    print(f"GPU: {torch.cuda.get_device_name()}; PyTorch: {torch.__version__}; "
+          f"CUDA: {torch.version.cuda}; D={args.hidden}; dH={args.grad_dtype}")
+    print(f"Median CUDA-event latency; warmup={args.warmup}, repeat={args.repeat}. "
+          "Extra peak allocated memory includes outputs and temporaries.")
+    print("| T | Backend | Forward (us) | Backward (us) | Forward extra (MiB) | "
+          "Backward extra (MiB) | Oracle bytes |")
+    print("|---:|:---|---:|---:|---:|---:|:---|")
+    torch.manual_seed(366)
+    with torch.no_grad():
+        for tokens in args.tokens:
+            residual = torch.randn(tokens, 4, args.hidden, device="cuda", dtype=torch.bfloat16)
+            pre = torch.randn(tokens, 4, device="cuda")
+            dh = torch.randn(tokens, args.hidden, device="cuda", dtype=grad_dtype)
+            expected = oracle.h_aggregate_fwd(pre, residual)
+            expected_dr, expected_dp = oracle.h_aggregate_bwd(dh, pre, residual)
+            for name, op in ops.items():
+                assert_bytes(op(residual, pre), expected)
+                dr, dp = op.backward_fp32(dh, residual, pre)
+                assert_bytes(dr, expected_dr)
+                assert_bytes(dp, expected_dp)
+                del dr, dp
+                forward_us, forward_mem = measure(
+                    partial(op, residual, pre), args.warmup, args.repeat)
+                backward_us, backward_mem = measure(
+                    partial(op.backward_fp32, dh, residual, pre), args.warmup, args.repeat)
+                print(f"| {tokens} | {name} | {forward_us:.2f} | {backward_us:.2f} | "
+                      f"{forward_mem:.3f} | {backward_mem:.3f} | PASS |", flush=True)
+
+
+if __name__ == "__main__":
+    main()

csrc/cuda/mhc/mhc_pre_h_aggregate.cu

@@ -0,0 +1,140 @@
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <c10/cuda/CUDAException.h>
+#include <torch/extension.h>
+
+#include <cuda_bf16.h>
+#include <cuda_runtime.h>
+
+#include <cstdint>
+#include <limits>
+#include <vector>
+
+#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) == rl_kernel::mhc::kMhcPreHcMult &&
+                  (residual.size(2) == 128 ||
+                   residual.size(2) == rl_kernel::mhc::kMhcPreHiddenSize),
+              "residual must have shape [num_tokens, 4, hidden_size], hidden_size must be 128 or 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");
+
+  int64_t const num_tokens = residual.size(0);
+  int64_t const hidden_size = residual.size(2);
+  TORCH_CHECK(num_tokens <= std::numeric_limits<unsigned int>::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<at::BFloat16>());
+  auto const* pre_ptr = pre.data_ptr<float>();
+  auto* output_ptr = reinterpret_cast<__nv_bfloat16*>(
+      output.data_ptr<at::BFloat16>());
+
+  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;
+}
+
+std::vector<torch::Tensor> 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(residual.scalar_type() == torch::kBFloat16, "residual must be bfloat16");
+  TORCH_CHECK(grad_output.scalar_type() == torch::kBFloat16 ||
+                  grad_output.scalar_type() == torch::kFloat32,
+              "grad_output must be bfloat16 or float32");
+  TORCH_CHECK(pre.scalar_type() == torch::kFloat32, "pre must be float32");
+  TORCH_CHECK(residual.dim() == 3 &&
+                  residual.size(1) == rl_kernel::mhc::kMhcPreHcMult &&
+                  (residual.size(2) == 128 ||
+                   residual.size(2) == rl_kernel::mhc::kMhcPreHiddenSize),
+              "residual must have shape [num_tokens, 4, hidden_size], hidden_size must be 128 or 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]");
+  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<unsigned int>::max(),
+              "num_tokens exceeds the CUDA grid limit");
+
+  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};
+  }
+
+  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* residual_ptr = reinterpret_cast<__nv_bfloat16 const*>(
+      residual.data_ptr<at::BFloat16>());
+  auto const* pre_ptr = pre.data_ptr<float>();
+  auto* grad_residual_ptr = grad_residual.data_ptr<float>();
+  auto* grad_pre_ptr = grad_pre.data_ptr<float>();
+
+  cudaError_t status;
+  if (grad_output.scalar_type() == torch::kFloat32) {
+    status = rl_kernel::mhc::launch_mhc_pre_h_aggregate_backward(
+        grad_output.data_ptr<float>(), residual_ptr, pre_ptr, grad_residual_ptr,
+        grad_pre_ptr, num_tokens, hidden_size, at::cuda::getCurrentCUDAStream(), major >= 9);
+  } else {
+    auto const* grad_output_ptr = reinterpret_cast<__nv_bfloat16 const*>(
+        grad_output.data_ptr<at::BFloat16>());
+    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};
+}