gpt-o3 / cuda0743e3

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

/*
 * Python entry point
 *     samples = run(probs, top_k)
 *
 * All tensors must already live on the desired CUDA device.
 */
torch::Tensor run(torch::Tensor probs, torch::Tensor top_k)
{
    TORCH_CHECK(probs.is_cuda(), "probs must reside on CUDA");
    TORCH_CHECK(top_k.is_cuda(), "top_k must reside on CUDA");
    TORCH_CHECK(probs.scalar_type() == torch::kFloat32, "probs must be float32");
    TORCH_CHECK(top_k.scalar_type()  == torch::kInt32 , "top_k must be int32");
    TORCH_CHECK(probs.dim() == 2 && probs.size(1) == 151936,
                "probs must have shape [batch_size, 151936]");

    const int64_t batch_size = probs.size(0);

    auto out_opts = torch::TensorOptions()
                        .dtype(torch::kInt64)
                        .device(probs.device());

    torch::Tensor samples = torch::empty({batch_size}, out_opts);

    /* delegate to CUDA implementation */
    top_k_sampling_from_probs_v151936_cuda(
        probs.contiguous(),
        top_k.contiguous(),
        samples);

    return samples;
}

/* PyBind11 --------------------------------------------------------------- */
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
{
    m.def("run", &run,
          "Top-k sampling for Qwen-3 vocab (CUDA)");
}