gpt-5 / cuda8fff8a

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

#include <stdexcept>
#include <string>

static void check_inputs(const torch::Tensor& hidden_states,
                         const torch::Tensor& weight) {
    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 be bfloat16");
    TORCH_CHECK(weight.dtype() == torch::kBFloat16, "weight must be bfloat16");
    TORCH_CHECK(hidden_states.dim() == 2, "hidden_states must be 2D [batch, hidden]");
    TORCH_CHECK(weight.dim() == 1, "weight must be 1D [hidden]");
    TORCH_CHECK(hidden_states.size(1) == RMSNORM_H7168,
                "hidden_states hidden dimension must be 7168, got ", hidden_states.size(1));
    TORCH_CHECK(weight.size(0) == RMSNORM_H7168,
                "weight dimension must be 7168, got ", weight.size(0));
    TORCH_CHECK(hidden_states.is_contiguous(), "hidden_states must be contiguous");
    TORCH_CHECK(weight.is_contiguous(), "weight must be contiguous");
    TORCH_CHECK(hidden_states.get_device() == weight.get_device(),
                "hidden_states and weight must be on the same device");
}

// Entry point exposed to Python
torch::Tensor run(torch::Tensor hidden_states, torch::Tensor weight) {
    check_inputs(hidden_states, weight);

    const int64_t batch_size = hidden_states.size(0);
    if (batch_size == 0) {
        return torch::empty_like(hidden_states);
    }

    // Ensure kernels launch on the correct CUDA device/stream
    at::cuda::CUDAGuard device_guard(hidden_states.device());
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    // Allocate output
    auto output = torch::empty_like(hidden_states);

    // Raw pointers as __nv_bfloat16 (only used as opaque pointers in host code)
    const __nv_bfloat16* hidden_ptr =
        reinterpret_cast<const __nv_bfloat16*>(hidden_states.data_ptr<c10::BFloat16>());
    const __nv_bfloat16* weight_ptr =
        reinterpret_cast<const __nv_bfloat16*>(weight.data_ptr<c10::BFloat16>());
    __nv_bfloat16* output_ptr =
        reinterpret_cast<__nv_bfloat16*>(output.data_ptr<c10::BFloat16>());

    // Launch kernel
    constexpr float eps = 1e-6f;
    launch_rmsnorm_h7168(hidden_ptr, weight_ptr, output_ptr,
                         static_cast<int>(batch_size), stream, eps);

    // Check for asynchronous launch errors
    cudaError_t err = cudaGetLastError();
    TORCH_CHECK(err == cudaSuccess, "rmsnorm_h7168 kernel launch failed: ", cudaGetErrorString(err));

    return output;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("run", &run, "rmsnorm_h7168 (BF16, hidden=7168)",
          pybind11::arg("hidden_states"),
          pybind11::arg("weight"));
}