gemini-2.5-pro / cuda6a9a99

#include <torch/extension.h>
#include "kernel.h"

// --- Host Function Implementation ---
// This is the C++ function that will be called from Python.
torch::Tensor rmsnorm_h7168(
    torch::Tensor hidden_states,
    torch::Tensor weight) {

  // --- Input Validation ---
  const int HIDDEN_SIZE_CONST = 7168;
  TORCH_CHECK(hidden_states.is_cuda(), "hidden_states must be a CUDA tensor");
  TORCH_CHECK(weight.is_cuda(), "weight must be a CUDA tensor");
  TORCH_CHECK(hidden_states.scalar_type() == torch::kBFloat16, "hidden_states must have bfloat16 type");
  TORCH_CHECK(weight.scalar_type() == torch::kBFloat16, "weight must have bfloat16 type");
  TORCH_CHECK(hidden_states.dim() == 2, "hidden_states must be 2D");
  TORCH_CHECK(hidden_states.size(1) == HIDDEN_SIZE_CONST, "hidden_states must have hidden_size of ", HIDDEN_SIZE_CONST);
  TORCH_CHECK(weight.dim() == 1, "weight must be 1D");
  TORCH_CHECK(weight.size(0) == HIDDEN_SIZE_CONST, "weight must have size of ", HIDDEN_SIZE_CONST);
  TORCH_CHECK(hidden_states.is_contiguous(), "hidden_states must be contiguous");
  TORCH_CHECK(weight.is_contiguous(), "weight must be contiguous");
  
  // --- Output Tensor Allocation ---
  auto output = torch::empty_like(hidden_states);

  // --- Kernel Launch ---
  rmsnorm_h7168_cuda_launcher(output, hidden_states, weight);

  return output;
}

// --- Pybind11 Module Definition ---
// This defines the Python module and exposes the C++ function.
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
  m.def(
      "rmsnorm_h7168",                                  // Python function name
      &rmsnorm_h7168,                                   // C++ function to bind
      "Optimized CUDA RMSNorm for hidden_size=7168 (BFloat16)", // Docstring
      pybind11::arg("hidden_states"),                   // Argument names
      pybind11::arg("weight")
  );
}