gemini-2.5-pro / cuda39bc7a

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

// Helper to check tensor properties for robust error handling
void check_tensor(const torch::Tensor& tensor, const std::string& name, torch::ScalarType dtype) {
    if (!tensor.defined()) {
        throw std::runtime_error(name + " is not defined");
    }
    if (!tensor.is_cuda()) {
        throw std::runtime_error(name + " must be a CUDA tensor");
    }
    if (tensor.scalar_type() != dtype) {
        throw std::runtime_error(name + " must have dtype " + std::string(c10::toString(dtype)));
    }
    if (!tensor.is_contiguous()) {
        throw std::runtime_error(name + " must be contiguous");
    }
}

// Main C++ entry point, called from Python
torch::Tensor top_k_top_p_sampling_from_probs_v129280(
    torch::Tensor probs,
    torch::Tensor top_k,
    torch::Tensor top_p
) {
    // --- Input Validation ---
    check_tensor(probs, "probs", torch::kFloat32);
    check_tensor(top_k, "top_k", torch::kInt32);
    check_tensor(top_p, "top_p", torch::kFloat32);

    TORCH_CHECK(probs.dim() == 2, "probs must be a 2D tensor of shape [batch_size, vocab_size]");
    const int batch_size = probs.size(0);
    const int vocab_size = probs.size(1);
    TORCH_CHECK(vocab_size == 129280, "vocab_size must be 129280, but got ", vocab_size);

    TORCH_CHECK(top_k.dim() == 1 && top_k.size(0) == batch_size, "top_k must be a 1D tensor of size batch_size");
    TORCH_CHECK(top_p.dim() == 1 && top_p.size(0) == batch_size, "top_p must be a 1D tensor of size batch_size");

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

    // --- Kernel Execution ---
    // Get the current CUDA stream from PyTorch to ensure proper synchronization
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    launch_top_k_top_p_sampling_kernel(
        probs.data_ptr<float>(),
        top_k.data_ptr<int>(),
        top_p.data_ptr<float>(),
        samples.data_ptr<long long>(),
        batch_size,
        vocab_size,
        stream
    );

    // Check for any asynchronous errors from the kernel launch
    cudaError_t err = cudaGetLastError();
    if (err != cudaSuccess) {
        throw std::runtime_error(std::string("CUDA kernel launch error: ") + cudaGetErrorString(err));
    }

    return samples;
}

// --- Pybind11 Module Definition ---
// Exposes the C++ function to Python
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("run", &top_k_top_p_sampling_from_probs_v129280, "Top-K Top-P Sampling from Probabilities (CUDA)");
}