gpt-5 / cudaa83af5

#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>

#include <iostream>
#include <vector>

#include "kernel.h"

static inline void check_inputs(const torch::Tensor& hidden_states,
                                const torch::Tensor& weight) {
  TORCH_CHECK(hidden_states.dim() == 2, "hidden_states must be 2D [batch_size, 1536]");
  TORCH_CHECK(weight.dim() == 1, "weight must be 1D [1536]");
  TORCH_CHECK(hidden_states.size(1) == RMSNORM_H1536_HIDDEN_SIZE,
              "hidden_states.shape[1] must be 1536");
  TORCH_CHECK(weight.size(0) == RMSNORM_H1536_HIDDEN_SIZE,
              "weight.shape[0] must be 1536");
  TORCH_CHECK(hidden_states.scalar_type() == at::kBFloat16,
              "hidden_states must be bfloat16");
  TORCH_CHECK(weight.scalar_type() == at::kBFloat16,
              "weight must be bfloat16");
}

torch::Tensor run(torch::Tensor hidden_states, torch::Tensor weight) {
  // Validate high-level properties and types
  check_inputs(hidden_states, weight);

  const bool input_on_cuda = hidden_states.is_cuda();
  const bool weight_on_cuda = weight.is_cuda();

  // Choose device for execution
  int device_index = 0;
  if (input_on_cuda) {
    device_index = hidden_states.get_device();
  } else if (weight_on_cuda) {
    device_index = weight.get_device();
  } else {
    // Default to device 0 if both are on CPU
    device_index = 0;
  }

  c10::cuda::CUDAGuard device_guard(device_index);

  // Make contiguous copies on the target device
  torch::Tensor hidden_states_dev = hidden_states;
  torch::Tensor weight_dev = weight;

  if (!hidden_states_dev.is_cuda()) {
    hidden_states_dev = hidden_states_dev.to(torch::kCUDA, /*non_blocking=*/true);
  }
  if (!weight_dev.is_cuda()) {
    weight_dev = weight_dev.to(torch::kCUDA, /*non_blocking=*/true);
  }

  hidden_states_dev = hidden_states_dev.contiguous();
  weight_dev = weight_dev.contiguous();

  // Allocate output on device
  torch::Tensor output_dev = torch::empty_like(hidden_states_dev, hidden_states_dev.options());

  // Launch the CUDA kernel through the launcher
  rmsnorm_h1536_launcher(hidden_states_dev, weight_dev, output_dev);

  // Return result matching the original hidden_states device
  if (!input_on_cuda) {
    return output_dev.to(torch::kCPU, /*non_blocking=*/false);
  } else {
    return output_dev;
  }
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
  m.def("run", &run,
        pybind11::arg("hidden_states"),
        pybind11::arg("weight"),
        "rmsnorm_h1536: B200-optimized RMSNorm (BF16) with hidden_size=1536");
}