gemini-2.5-pro / cuda0ade42

#include <torch/extension.h>
#include <c10/cuda/CUDAStream.h>
#include <ATen/cuda/CUDAContext.h>

#include "kernel.h"

#include <string>
#include <stdexcept>
#include <chrono>

// Helper to check common tensor properties to reduce boilerplate code.
void check_tensor(const torch::Tensor& tensor, torch::ScalarType dtype, torch::DeviceType device, const std::string& name) {
    TORCH_CHECK(tensor.is_contiguous(), name, " must be contiguous");
    TORCH_CHECK(tensor.scalar_type() == dtype, name, " must have ", dtype, " dtype, but got ", tensor.scalar_type());
    TORCH_CHECK(tensor.device().type() == device, name, " must be on ", device, " device");
}

/**
 * @brief Python-bindable 'run' function for Top-K sampling.
 *
 * This function serves as the C++/Python bridge. It receives PyTorch tensors,
 * performs validation, obtains raw data pointers, and calls the CUDA kernel launcher.
 *
 * @param probs A [batch_size, 128256] float32 tensor of probabilities.
 * @param top_k A [batch_size] int32 tensor of k values.
 * @return A [batch_size] int64 tensor of sampled token indices.
 */
torch::Tensor run(torch::Tensor probs, torch::Tensor top_k) {
    // --- Input Validation ---
    const auto device = torch::kCUDA;
    check_tensor(probs, torch::kFloat32, device, "probs");
    check_tensor(top_k, torch::kInt32, device, "top_k");

    TORCH_CHECK(probs.dim() == 2, "probs must be a 2D tensor");
    const int64_t batch_size = probs.size(0);
    const int64_t vocab_size = probs.size(1);

    TORCH_CHECK(vocab_size == 128256, "vocab_size must be 128256, but got ", vocab_size);
    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");

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

    // --- Kernel Execution ---
    if (batch_size > 0) {
        // Get the current CUDA stream from PyTorch's context to ensure proper ordering.
        cudaStream_t stream = c10::cuda::getCurrentCUDAStream();

        // Generate a seed for the random number generator. Using a time-based seed
        // provides different random sequences for different runs.
        uint64_t seed = std::chrono::high_resolution_clock::now().time_since_epoch().count();

        top_k_sampling_from_probs_v128256_launch(
            probs.data_ptr<float>(),
            top_k.data_ptr<int>(),
            samples.data_ptr<int64_t>(),
            static_cast<int>(batch_size),
            seed,
            stream
        );
    }

    return samples;
}

// --- Pybind11 Module Definition ---
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("run", &run, "Top-K sampling from probability distributions (CUDA v128256)");
}