gemini-2.5-pro / cuda8cde37

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

// Helper function to check tensor properties for robust error handling
void check_tensor(const torch::Tensor& tensor, const std::string& name, torch::ScalarType dtype, bool is_cuda) {
    TORCH_CHECK(tensor.is_contiguous(), name, " must be contiguous");
    if (is_cuda) {
        TORCH_CHECK(tensor.is_cuda(), name, " must be a CUDA tensor");
    }
    TORCH_CHECK(tensor.scalar_type() == dtype, name, " must have ", dtype, " dtype, but got ", tensor.scalar_type());
}

/**
 * @brief Python-bindable entry point for the top-k sampling operation.
 *
 * This function serves as the C++/PyTorch interface. It performs tensor validation,
 * extracts data pointers, and calls the CUDA kernel launcher on the current PyTorch stream.
 *
 * @param probs A [batch_size, vocab_size] float32 CUDA tensor of probabilities.
 * @param top_k A [batch_size] int32 CUDA tensor of top-k values.
 * @return A [batch_size] int64 CUDA tensor of sampled token indices.
 */
torch::Tensor run(
    const torch::Tensor& probs,
    const torch::Tensor& top_k) {

    // --- Input Validation ---
    const int64_t batch_size = probs.size(0);
    const int64_t vocab_size = probs.size(1);

    TORCH_CHECK(probs.dim() == 2, "probs must be a 2D tensor");
    TORCH_CHECK(vocab_size == 129280, "probs vocab_size must be 129280, but got ", vocab_size);
    check_tensor(probs, "probs", torch::kFloat32, true);

    TORCH_CHECK(top_k.dim() == 1, "top_k must be a 1D tensor");
    TORCH_CHECK(top_k.size(0) == batch_size, "top_k batch size must match probs batch size");
    check_tensor(top_k, "top_k", torch::kInt32, true);

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

    // Early exit if batch_size is zero
    if (batch_size == 0) {
        return samples;
    }

    // --- Kernel Execution ---
    // Ensure that the device for the current context matches the tensor's device
    c10::cuda::CUDAGuard device_guard(probs.device());
    
    // Get the current CUDA stream from PyTorch to ensure proper synchronization.
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    top_k_sampling_from_probs_v129280_launcher(
        probs.data_ptr<float>(),
        top_k.data_ptr<int>(),
        samples.data_ptr<int64_t>(),
        static_cast<int>(batch_size),
        stream
    );
    
    // Check for any asynchronous CUDA errors that might have occurred during kernel execution.
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    return samples;
}

// --- Pybind11 Module Definition ---
// Exposes the C++ `run` function to Python.
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("run", &run, "Top-K sampling from probability distributions (CUDA implementation for v129280)");
}