gpt-5 / cudadccc70

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

namespace {

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, hidden_size]");
    TORCH_CHECK(hidden_states.size(1) == rmsnorm_h512::HIDDEN_SIZE,
                "hidden_size must be 512, got ", hidden_states.size(1));
    TORCH_CHECK(weight.dim() == 1 && weight.size(0) == rmsnorm_h512::HIDDEN_SIZE,
                "weight must be 1D of size 512");
    TORCH_CHECK(hidden_states.dtype() == at::kBFloat16, "hidden_states must be BF16");
    TORCH_CHECK(weight.dtype() == at::kBFloat16, "weight must be BF16");
}

inline torch::Tensor to_contiguous_if_needed(const torch::Tensor& t) {
    return t.is_contiguous() ? t : t.contiguous();
}

} // anonymous namespace

torch::Tensor run(torch::Tensor hidden_states, torch::Tensor weight) {
    check_inputs(hidden_states, weight);

    // Choose device:
    torch::Device device = hidden_states.is_cuda() ? hidden_states.device()
                          : (weight.is_cuda() ? weight.device()
                                              : torch::Device(torch::kCUDA, 0));

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

    // Move to device if needed
    torch::Tensor hidden_dev = hidden_states.is_cuda() ? hidden_states : hidden_states.to(device, /*non_blocking=*/false);
    torch::Tensor weight_dev = weight.is_cuda() ? weight : weight.to(device, /*non_blocking=*/false);

    // Ensure contiguous
    hidden_dev = to_contiguous_if_needed(hidden_dev);
    weight_dev = to_contiguous_if_needed(weight_dev);

    const int64_t batch_size = hidden_dev.size(0);
    // Allocate output on device
    torch::Tensor output_dev = torch::empty_like(hidden_dev);

    // Get CUDA stream
    cudaStream_t stream = at::cuda::getCurrentCUDAStream(device.index()).stream();

    // Copy weight to constant memory
    const __nv_bfloat16* w_ptr = reinterpret_cast<const __nv_bfloat16*>(weight_dev.data_ptr<at::BFloat16>());
    rmsnorm_h512::set_weight_const(w_ptr, stream);

    // Launch kernel
    const __nv_bfloat16* x_ptr = reinterpret_cast<const __nv_bfloat16*>(hidden_dev.data_ptr<at::BFloat16>());
    __nv_bfloat16* y_ptr = reinterpret_cast<__nv_bfloat16*>(output_dev.data_ptr<at::BFloat16>());
    rmsnorm_h512::launch_forward(x_ptr, y_ptr, static_cast<int>(batch_size), stream);

    // If original input was on CPU, return CPU tensor
    if (!hidden_states.is_cuda()) {
        return output_dev.to(hidden_states.device(), /*non_blocking=*/false);
    }
    return output_dev;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("run", &run,
          "rmsnorm_h512 (BF16) - B200-optimized CUDA kernel",
          pybind11::arg("hidden_states"),
          pybind11::arg("weight"));
}