gemini-2.5-pro / cudad79ad6

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

// For CUDA stream management
#include <ATen/cuda/CUDAContext.h>

#include <stdexcept>
#include <string>

// C++ entry point function that interfaces with Python.
torch::Tensor run(torch::Tensor probs, torch::Tensor top_k) {
    // --- Input Validation ---
    TORCH_CHECK(probs.is_cuda(), "Input tensor 'probs' must be on a CUDA device");
    TORCH_CHECK(top_k.is_cuda(), "Input tensor 'top_k' must be on a CUDA device");

    TORCH_CHECK(probs.dim() == 2, "Input tensor 'probs' must be 2-dimensional");
    TORCH_CHECK(top_k.dim() == 1, "Input tensor 'top_k' must be 1-dimensional");

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

    TORCH_CHECK(vocab_size == 151936, "Vocabulary size (dim 1 of probs) must be 151936");
    TORCH_CHECK(top_k.size(0) == batch_size, "Dimension 0 of 'top_k' must match batch_size");

    TORCH_CHECK(probs.scalar_type() == torch::kFloat32, "Input tensor 'probs' must have dtype float32");
    TORCH_CHECK(top_k.scalar_type() == torch::kInt32, "Input tensor 'top_k' must have dtype int32");

    // Ensure tensors are contiguous in memory for efficient CUDA access.
    auto probs_c = probs.contiguous();
    auto top_k_c = top_k.contiguous();

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

    // --- Kernel Execution ---
    // Get the current CUDA stream from PyTorch to enqueue work.
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    // Call the host launcher function defined in kernel.cu.
    top_k_sampling_from_probs_v151936_launch(
        probs_c.data_ptr<float>(),
        top_k_c.data_ptr<int>(),
        samples.data_ptr<int64_t>(), // Corrected to int64_t
        batch_size,
        stream
    );
    
    // --- Error Handling ---
    // 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 failed: ") + cudaGetErrorString(err));
    }
    
    return samples;
}

// --- Pybind11 Module Definition ---
// This creates the Python module 'TORCH_EXTENSION_NAME' that can be imported.
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("run", &run, "Top-K sampling from probabilities (CUDA implementation for v151936)");
}