gemini-2.5-pro / cudafc4885

#include "kernel.h"
#include <torch/extension.h>
#include <vector>

// The "run" function is the entry point exposed to Python.
// It handles tensor validation and calls the CUDA implementation.
torch::Tensor run(torch::Tensor probs, torch::Tensor top_k, torch::Tensor top_p) {
    // Input validation: Device
    TORCH_CHECK(probs.is_cuda(), "probs tensor must be on a CUDA device");
    TORCH_CHECK(top_k.device() == probs.device(), "top_k must be on the same device as probs");
    TORCH_CHECK(top_p.device() == probs.device(), "top_p must be on the same device as probs");

    // Input validation: Dtypes
    TORCH_CHECK(probs.scalar_type() == torch::kFloat32, "probs must be of type float32");
    TORCH_CHECK(top_k.scalar_type() == torch::kInt32, "top_k must be of type int32");
    TORCH_CHECK(top_p.scalar_type() == torch::kFloat32, "top_p must be of type float32");

    // Input validation: Shapes
    TORCH_CHECK(probs.dim() == 2, "probs must be a 2D tensor");
    const int batch_size = probs.size(0);
    const int vocab_size = probs.size(1);
    
    const int required_vocab_size = 151936;
    TORCH_CHECK(vocab_size == required_vocab_size, "vocab_size must be ", required_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");

    // Ensure tensors are contiguous for safe data_ptr access
    probs = probs.contiguous();
    top_k = top_k.contiguous();
    top_p = top_p.contiguous();
    
    // Create the output tensor
    auto samples = torch::empty({(long)batch_size}, torch::TensorOptions().dtype(torch::kInt64).device(probs.device()));

    // Call the CUDA kernel launcher
    top_k_top_p_sampling_from_probs_v151936_cuda(probs, top_k, top_p, samples);

    return samples;
}

// PYBIND11 module definition to expose the 'run' function to Python
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("run", &run, "Top-K/Top-P sampling from probability distributions (CUDA implementation for v151936)");
}