gemini-2.5-pro / cuda4b5e32

#include "kernel.h"
#include <torch/extension.h>
#include <pybind11/pybind11.h>
#include <c10/cuda/CUDAStream.h>
#include <string>
#include <stdexcept>

namespace py = pybind11;

// Helper to check tensor properties for robust error handling.
void check_tensor(const torch::Tensor& t, const std::string& name, torch::ScalarType dtype, torch::DeviceType device, bool is_contiguous = true) {
    if (!t.defined()) {
        throw std::runtime_error(name + " is not defined");
    }
    if (t.scalar_type() != dtype) {
        throw std::runtime_error(name + " has incorrect dtype. Expected " + std::string(c10::toString(dtype)) + ", but got " + std::string(c10::toString(t.scalar_type())));
    }
    if (t.device().type() != device) {
        throw std::runtime_error(name + " is not on the correct device. Expected " + std::string(c10::toString(device)));
    }
    if (is_contiguous && !t.is_contiguous()) {
        throw std::runtime_error(name + " is not contiguous");
    }
}

// Main `run` function exposed to Python.
torch::Tensor run(
    torch::Tensor probs,
    torch::Tensor top_p) {

    // --- Input Validation ---
    const auto device = torch::kCUDA;
    check_tensor(probs, "probs", torch::kFloat32, device);
    check_tensor(top_p, "top_p", torch::kFloat32, device);

    const int batch_size = probs.size(0);
    const int vocab_size = probs.size(1);

    if (probs.dim() != 2) {
        throw std::runtime_error("probs must be a 2D tensor, but got " + std::to_string(probs.dim()) + " dimensions.");
    }
    if (vocab_size != 128256) {
        throw std::runtime_error("vocab_size must be 128256, but got " + std::to_string(vocab_size));
    }
    if (top_p.dim() != 1 || top_p.size(0) != batch_size) {
        throw std::runtime_error("top_p must be a 1D tensor of size batch_size (" + std::to_string(batch_size) + "), but got shape " + c10::IntArrayRef(top_p.sizes()).str());
    }

    // --- Output Allocation ---
    auto samples = torch::empty({batch_size}, torch::TensorOptions()
                                    .dtype(torch::kInt64)
                                    .device(device)
                                    .memory_format(torch::MemoryFormat::Contiguous));

    // --- Launch CUDA Kernels via the host function ---
    top_p_sampling_from_probs_v128256_cuda(probs, top_p, samples);

    return samples;
}

// --- Pybind11 Module Definition ---
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("run", &run, "Top-P Sampling from Probabilities (CUDA v128256)",
          py::arg("probs"), py::arg("top_p"));
}