gemini-2.5-pro / cudac4bb10

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

// Helper function to check tensor properties
#define CHECK_TENSOR(x, d) TORCH_CHECK(x.is_cuda() && x.is_contiguous() && x.dtype() == d, \
    #x " must be a contiguous CUDA tensor of type " #d)

/**
 * @brief Python-bindable entry point for the top-p sampling operation.
 *
 * This function serves as the C++/PyTorch interface. It performs tensor checks,
 * extracts necessary metadata and data pointers, and calls the main CUDA host function.
 *
 * @param probs A PyTorch tensor of shape [batch_size, 151936] and dtype float32,
 *              representing the probability distributions.
 * @param top_p A PyTorch tensor of shape [batch_size] and dtype float32,
 *              representing the cumulative probability thresholds.
 * @return A PyTorch tensor of shape [batch_size] and dtype int64, containing
 *         the sampled token indices.
 */
torch::Tensor run(torch::Tensor probs, torch::Tensor top_p) {
    // --- Input Validation ---
    const auto float_type = torch::kFloat32;
    CHECK_TENSOR(probs, float_type);
    CHECK_TENSOR(top_p, float_type);

    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);
    
    TORCH_CHECK(vocab_size == 151936, "vocab_size must be 151936");
    TORCH_CHECK(top_p.dim() == 1, "top_p must be a 1D tensor");
    TORCH_CHECK(top_p.size(0) == batch_size, "top_p must have the same batch_size as probs");

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

    // --- Get CUDA Stream ---
    // Note: c10::cuda::getCurrentCUDAStream() is the modern way to get the stream
    // for the current device.
    cudaStream_t stream = c10::cuda::getCurrentCUDAStream();

    // --- Launch CUDA Kernels ---
    top_p_sampling_from_probs_v151936_cuda(
        samples.data_ptr<long long>(),
        probs.data_ptr<float>(),
        top_p.data_ptr<float>(),
        batch_size,
        stream
    );
    
    // --- Error Checking & Synchronization ---
    // A call to cudaGetLastError() is good practice after a kernel launch
    // to catch any configuration errors. PyTorch's stream handling often
    // makes explicit cudaStreamSynchronize unnecessary here, but it's
    // essential for robust error checking.
    auto err = cudaGetLastError();
    TORCH_CHECK(err == cudaSuccess, "CUDA kernel launch failed: ", cudaGetErrorString(err));

    return samples;
}

// --- Pybind11 Module Definition ---
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("run", &run, "Top-P sampling from probabilities (CUDA implementation for v151936)");
}