gpt-5 / cuda727b5d

#include "kernel.h"

#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <vector>
#include <stdexcept>
#include <sstream>

static void check_inputs(const torch::Tensor& hidden_states,
                         const torch::Tensor& residual,
                         const torch::Tensor& weight) {
  TORCH_CHECK(hidden_states.dim() == 2, "hidden_states must be rank-2 [batch_size, 4096]");
  TORCH_CHECK(residual.dim() == 2, "residual must be rank-2 [batch_size, 4096]");
  TORCH_CHECK(weight.dim() == 1, "weight must be rank-1 [4096]");
  TORCH_CHECK(hidden_states.size(1) == HIDDEN_SIZE, "hidden_size must be 4096");
  TORCH_CHECK(residual.size(1) == HIDDEN_SIZE, "hidden_size must be 4096");
  TORCH_CHECK(weight.size(0) == HIDDEN_SIZE, "weight length must be 4096");
  TORCH_CHECK(hidden_states.scalar_type() == at::kBFloat16, "hidden_states must be bfloat16");
  TORCH_CHECK(residual.scalar_type() == at::kBFloat16, "residual must be bfloat16");
  TORCH_CHECK(weight.scalar_type() == at::kBFloat16, "weight must be bfloat16");
}

torch::Tensor run(torch::Tensor hidden_states,
                  torch::Tensor residual,
                  torch::Tensor weight) {
  check_inputs(hidden_states, residual, weight);

  const int64_t batch_size = hidden_states.size(0);

  // Ensure contiguous tensors; move to CUDA if needed
  torch::Tensor hs_cuda = hidden_states.contiguous();
  torch::Tensor rs_cuda = residual.contiguous();
  torch::Tensor w_cuda  = weight.contiguous();

  if (!hs_cuda.is_cuda()) hs_cuda = hs_cuda.to(at::kCUDA, at::kBFloat16, /*non_blocking=*/false, /*copy=*/true);
  if (!rs_cuda.is_cuda()) rs_cuda = rs_cuda.to(at::kCUDA, at::kBFloat16, /*non_blocking=*/false, /*copy=*/true);
  if (!w_cuda.is_cuda())  w_cuda  = w_cuda.to(at::kCUDA,  at::kBFloat16, /*non_blocking=*/false, /*copy=*/true);

  auto opts = hs_cuda.options();
  torch::Tensor out_cuda = torch::empty_like(hs_cuda, opts);

  // Raw pointers
  const __nv_bfloat16* hs_ptr = reinterpret_cast<const __nv_bfloat16*>(hs_cuda.data_ptr<at::BFloat16>());
  const __nv_bfloat16* rs_ptr = reinterpret_cast<const __nv_bfloat16*>(rs_cuda.data_ptr<at::BFloat16>());
  const __nv_bfloat16* w_ptr  = reinterpret_cast<const __nv_bfloat16*>(w_cuda.data_ptr<at::BFloat16>());
  __nv_bfloat16* out_ptr      = reinterpret_cast<__nv_bfloat16*>(out_cuda.data_ptr<at::BFloat16>());

  cudaStream_t stream = at::cuda::getCurrentCUDAStream();

  launch_fused_add_rmsnorm_h4096(hs_ptr, rs_ptr, w_ptr, out_ptr,
                                 static_cast<int>(batch_size), stream);

  // Make sure work is finished before moving data back to CPU
  auto err = cudaStreamSynchronize(stream);
  TORCH_CHECK(err == cudaSuccess, "CUDA stream sync failed: ", cudaGetErrorString(err));

  // Return results to CPU BF16 as in the reference
  torch::Tensor out_cpu = out_cuda.to(at::kCPU, at::kBFloat16);

  return out_cpu;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
  m.def("run",
        &run,
        "fused_add_rmsnorm_h4096 (BF16, B200-optimized)",
        py::arg("hidden_states"),
        py::arg("residual"),
        py::arg("weight"));
}