gpt-5 / cuda5fa8b5

#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAStream.h>
#include <vector>
#include <stdexcept>
#include "kernel.h"

// Validate shapes and dtypes
static void validate_inputs(const torch::Tensor& hidden_states,
                            const torch::Tensor& weight) {
  TORCH_CHECK(hidden_states.dim() == 2, "hidden_states must be 2D [batch_size, hidden_size]");
  TORCH_CHECK(weight.dim() == 1, "weight must be 1D [hidden_size]");
  TORCH_CHECK(hidden_states.size(1) == 2048, "hidden_size must be 2048");
  TORCH_CHECK(weight.size(0) == 2048, "weight size must be 2048");
  TORCH_CHECK(hidden_states.scalar_type() == at::kBFloat16,
              "hidden_states must be torch.bfloat16");
  TORCH_CHECK(weight.scalar_type() == at::kBFloat16,
              "weight must be torch.bfloat16");
}

// Public entry point exposed to Python
torch::Tensor run(torch::Tensor hidden_states, torch::Tensor weight) {
  validate_inputs(hidden_states, weight);

  bool input_on_cpu = !hidden_states.is_cuda();
  bool weight_on_cpu = !weight.is_cuda();

  // Choose device: current CUDA device
  at::OptionalDeviceGuard device_guard;
  int device_index = 0;
  if (!input_on_cpu) {
    device_index = hidden_states.get_device();
  } else if (!weight_on_cpu) {
    device_index = weight.get_device();
  } else {
    device_index = at::cuda::current_device();
  }
  device_guard.reset_device(at::Device(at::kCUDA, device_index));

  // Move tensors to CUDA and ensure contiguous memory
  torch::Tensor d_hidden = hidden_states.is_cuda()
                               ? hidden_states.contiguous()
                               : hidden_states.to(hidden_states.options().device(at::kCUDA, device_index)).contiguous();

  torch::Tensor d_weight = weight.is_cuda()
                               ? weight.contiguous()
                               : weight.to(weight.options().device(at::kCUDA, device_index)).contiguous();

  auto d_output = torch::empty_like(d_hidden);

  // Use current CUDA stream for proper PyTorch stream semantics
  auto stream = at::cuda::getCurrentCUDAStream();

  // Upload weight to constant memory (BF16)
  set_weight_const_from_device(d_weight.data_ptr(), stream.stream());

  // Launch kernel
  const int64_t batch_size = d_hidden.size(0);
  rmsnorm_h2048_launcher(d_hidden.data_ptr(),
                         d_output.data_ptr(),
                         static_cast<int>(batch_size),
                         stream.stream());

  // If input was on CPU, bring result back; otherwise return device tensor
  if (input_on_cpu) {
    return d_output.to(hidden_states.device(), hidden_states.scalar_type());
  } else {
    return d_output;
  }
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
  m.def("run", &run, pybind11::arg("hidden_states"), pybind11::arg("weight"),
        "RMSNorm (H=2048) kernel optimized for NVIDIA B200. "
        "Inputs: hidden_states [B,2048] bfloat16, weight [2048] bfloat16. Output: [B,2048] bfloat16.");
}