gpt-5-2025-08-07_cuda_371839

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

// pybind11 alias
#include <pybind11/pybind11.h>
namespace py = pybind11;

static void check_inputs(const torch::Tensor& probs, const torch::Tensor& top_p) {
    TORCH_CHECK(probs.dim() == 2, "probs must be 2D [batch_size, 129280]");
    TORCH_CHECK(probs.size(1) == VOCAB_SIZE_V129280, "vocab_size must be 129280");
    TORCH_CHECK(probs.scalar_type() == at::kFloat, "probs must be float32");
    TORCH_CHECK(top_p.dim() == 1, "top_p must be 1D [batch_size]");
    TORCH_CHECK(top_p.size(0) == probs.size(0), "top_p.shape[0] must equal probs.shape[0]");
    TORCH_CHECK(top_p.scalar_type() == at::kFloat, "top_p must be float32");
}

// Entry point exposed to Python
torch::Tensor run(torch::Tensor probs, torch::Tensor top_p, py::kwargs kwargs) {
    (void)kwargs; // currently unused, kept for API compatibility
    check_inputs(probs, top_p);
    const auto batch_size = probs.size(0);
    const auto vocab_size = probs.size(1);

    // Select device
    int device_index = 0;
    if (probs.is_cuda()) {
        device_index = probs.get_device();
    } else if (top_p.is_cuda()) {
        device_index = top_p.get_device();
    } else {
        device_index = at::cuda::current_device();
    }

    c10::cuda::CUDAGuard device_guard(device_index);
    auto stream = at::cuda::getCurrentCUDAStream();

    // Move inputs to GPU (contiguous)
    torch::Tensor probs_gpu = probs.is_cuda() ? probs.contiguous() : probs.to(torch::TensorOptions().device(at::kCUDA, device_index)).contiguous();
    torch::Tensor top_p_gpu = top_p.is_cuda() ? top_p.contiguous() : top_p.to(torch::TensorOptions().device(at::kCUDA, device_index)).contiguous();

    // Allocate output on GPU
    auto out_gpu = torch::empty({batch_size}, probs_gpu.options().dtype(at::kLong));

    // Launch CUDA implementation
    top_p_sampling_from_probs_v129280_cuda(
        probs_gpu.data_ptr<float>(),
        top_p_gpu.data_ptr<float>(),
        batch_size,
        vocab_size,
        out_gpu.data_ptr<int64_t>(),
        stream.stream());

    // Check for kernel errors and synchronize (ensures completion before returning to CPU)
    C10_CUDA_CHECK(cudaGetLastError());
    C10_CUDA_CHECK(cudaStreamSynchronize(stream.stream()));

    // Return results to CPU
    auto out_cpu = out_gpu.to(torch::kCPU);
    return out_cpu;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("run", &run, "top_p_sampling_from_probs_v129280 (CUDA)",
          py::arg("probs"),
          py::arg("top_p"));
}