gemini-2.5-pro / cuda208a66

#include "kernel.h"
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h> // FIX: Added missing header for CUDA stream access
#include <vector>

// Main entry point for the Python extension
torch::Tensor run(
    const torch::Tensor& hidden_states,
    const torch::Tensor& weight) {

    // --- Input Validation ---
    TORCH_CHECK(hidden_states.dim() == 2, "hidden_states must be a 2D tensor");
    TORCH_CHECK(weight.dim() == 1, "weight must be a 1D tensor");

    const int64_t hidden_size = hidden_states.size(1);

    TORCH_CHECK(hidden_size == 1536, "hidden_size must be 1536");
    TORCH_CHECK(weight.size(0) == hidden_size, "weight must have size equal to hidden_size");

    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.dtype() == torch::kBFloat16, "hidden_states must have bfloat16 dtype");
    TORCH_CHECK(weight.dtype() == torch::kBFloat16, "weight must have bfloat16 dtype");

    TORCH_CHECK(hidden_states.is_contiguous(), "hidden_states must be contiguous");
    TORCH_CHECK(weight.is_contiguous(), "weight must be contiguous");

    // --- Output Allocation ---
    auto output = torch::empty_like(hidden_states);

    // --- Kernel Execution ---
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    rmsnorm_h1536_launcher(output, hidden_states, weight, stream);

    return output;
}

// --- Pybind11 Module Definition ---
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("run", &run, "RMSNorm implementation for hidden_size=1536 on B200 (CUDA)",
          py::arg("hidden_states"),
          py::arg("weight"));
}