gemini-2.5-pro / cuda54f90a

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

// Helper macro for checking tensor properties
#define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x " must be a CUDA tensor")
#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)


/**
 * @brief Python-bindable entry point for the Top-K/Top-P sampling operation.
 *
 * This function handles the boilerplate of converting PyTorch tensors to raw
 * pointers and launching the CUDA implementation.
 *
 * @param probs Probability distributions. Shape: [batch_size, 128256], DType: float32.
 * @param top_k Top-K values. Shape: [batch_size], DType: int32.
 * @param top_p Top-P values. Shape: [batch_size], DType: float32.
 * @return A torch.Tensor containing the sampled token indices. Shape: [batch_size], DType: int64.
 */
torch::Tensor run(
    torch::Tensor probs,
    torch::Tensor top_k,
    torch::Tensor top_p) {

    // --- Input Validation ---
    CHECK_INPUT(probs);
    CHECK_INPUT(top_k);
    CHECK_INPUT(top_p);

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

    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");
    
    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");

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

    // --- Setup CUDA Environment ---
    const at::cuda::OptionalCUDAGuard device_guard(device_of(probs));
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    // --- Launch CUDA Kernel ---
    top_k_top_p_sampling_from_probs_v128256_cuda(
        probs.data_ptr<float>(),
        top_k.data_ptr<int>(),
        top_p.data_ptr<float>(),
        samples.data_ptr<long long>(),
        batch_size,
        stream
    );
    
    // Synchronize to ensure completion before returning to host
    AT_CUDA_CHECK(cudaStreamSynchronize(stream));

    return samples;
}

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