gpt-o3 / cuda0002fb

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

#include "kernel.h"

/* -------------------------------------------------------------------------- */
/*  Python visible entry point                                                */
/* -------------------------------------------------------------------------- */
torch::Tensor run(torch::Tensor probs,
                  torch::Tensor top_k)
{
    /* ------------------- basic validation ------------------------------ */
    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 must be 2-D [batch , vocab]");
    TORCH_CHECK(probs.size(1) == TOPK_V128256_VOCAB_SIZE,
                "vocab dimension must be 128256");

    TORCH_CHECK(probs.size(0) == top_k.size(0),
                "batch dimension mismatch between probs and top_k");

    const int64_t batch_size = probs.size(0);

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

    /* ------------------- launch CUDA code ------------------------------ */
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    top_k_sampling_from_probs_v128256(probs.data_ptr<float>(),
                                      top_k.data_ptr<int32_t>(),
                                      samples.data_ptr<int64_t>(),
                                      static_cast<int>(batch_size),
                                      stream);

    /* ------------------- check for runtime failures -------------------- */
    TORCH_CHECK(cudaStreamSynchronize(stream) == cudaSuccess,
                "CUDA execution failed");

    return samples;
}

/* -------------------------------------------------------------------------- */
/*  PyBind11 module                                                           */
/* -------------------------------------------------------------------------- */
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.doc() = "Top-k sampling (fixed vocab 128256, B200 optimised baseline)";
    m.def("run",
          &run,
          pybind11::arg("probs"),
          pybind11::arg("top_k"));
}