gpt-5-2025-08-07 / cudad0c81e

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

#include <vector>
#include <random>
#include <chrono>
#include "kernel.h"

using torch::Tensor;

namespace tk129280 {

static inline unsigned long long make_seed(uint64_t user_seed) {
    if (user_seed != 0ULL) return user_seed;
    // Derive a seed from random_device and time to avoid collisions
    std::random_device rd;
    auto t = static_cast<uint64_t>(std::chrono::high_resolution_clock::now().time_since_epoch().count());
    uint64_t s = (static_cast<uint64_t>(rd()) << 32) ^ t ^ 0x9E3779B97F4A7C15ULL;
    return static_cast<unsigned long long>(s);
}

// Core entry point
Tensor run(Tensor probs, Tensor top_k, uint64_t seed /*=0*/) {
    TORCH_CHECK(probs.dim() == 2, "probs must be 2D [batch_size, vocab_size]");
    TORCH_CHECK(probs.dtype() == torch::kFloat32 || probs.dtype() == torch::kHalf || probs.dtype() == torch::kBFloat16,
                "probs must be float32/float16/bfloat16");
    TORCH_CHECK(top_k.dim() == 1, "top_k must be 1D [batch_size]");
    TORCH_CHECK(probs.size(1) == tk129280::VOCAB_SIZE,
                "vocab_size must be exactly 129280, got ", probs.size(1));
    TORCH_CHECK(probs.size(0) == top_k.size(0),
                "batch_size mismatch between probs and top_k");

    // Select CUDA device (if probs on CPU, use current CUDA device)
    torch::Device device = probs.is_cuda() ? probs.device() : torch::Device(torch::kCUDA, at::cuda::current_device());
    c10::cuda::CUDAGuard device_guard(device);
    TORCH_CHECK(device.is_cuda(), "CUDA device required.");

    // Convert to float32 on selected device and make contiguous
    Tensor probs_f32 = probs.to(device, torch::kFloat32, /*non_blocking=*/true).contiguous();
    const int64_t batch_size = probs_f32.size(0);
    const int64_t vocab_size = probs_f32.size(1);

    // Prepare top_k on CPU as int32
    Tensor topk_i32_host = top_k.to(torch::kInt32, /*non_blocking=*/false).to(torch::kCPU);
    auto topk_ptr = topk_i32_host.data_ptr<int32_t>();

    // Output tensor on device
    auto options_out = torch::TensorOptions().dtype(torch::kInt64).device(device);
    Tensor samples_dev = torch::empty({batch_size}, options_out);

    // Workspace buffers (re-used for each row)
    float* vals_buf = nullptr;
    int32_t* idx_buf = nullptr;
    CUDA_CHECK(cudaMalloc(&vals_buf, sizeof(float) * vocab_size));
    CUDA_CHECK(cudaMalloc(&idx_buf, sizeof(int32_t) * vocab_size));

    // Use current CUDA stream
    cudaStream_t cuda_stream = at::cuda::getCurrentCUDAStream().stream();

    // Seed for device RNG
    unsigned long long device_seed = make_seed(seed);

    const float* probs_ptr = probs_f32.data_ptr<float>();
    int64_t* samples_ptr = samples_dev.data_ptr<int64_t>();

    for (int64_t i = 0; i < batch_size; ++i) {
        const float* row_ptr = probs_ptr + i * vocab_size;
        int32_t k = topk_ptr[i];

        if (k > 0 && k < vocab_size) {
            // Copy row to workspace (device-to-device)
            CUDA_CHECK(cudaMemcpyAsync(vals_buf, row_ptr, sizeof(float) * vocab_size,
                                       cudaMemcpyDeviceToDevice, cuda_stream));

            // Prepare index sequence 0..V-1
            tk129280::prepare_index_sequence(idx_buf, static_cast<int>(vocab_size), cuda_stream);

            // Reorder so that first k entries are top-k (we sort descending for robustness)
            tk129280::topk_nth_element_select(vals_buf, idx_buf, static_cast<int>(vocab_size), k, cuda_stream);

            // Sample from the first k entries
            tk129280::sample_from_first_k_kernel_launcher(vals_buf, idx_buf, k, device_seed, samples_ptr + i, static_cast<int>(i), cuda_stream);
        } else {
            // No filtering: sample from the original distribution
            tk129280::sample_row_unfiltered_kernel_launcher(row_ptr, static_cast<int>(vocab_size),
                                                            device_seed, samples_ptr + i, static_cast<int>(i), cuda_stream);
        }
    }

    CUDA_CHECK(cudaGetLastError());
    CUDA_CHECK(cudaStreamSynchronize(cuda_stream));

    // Free workspace
    CUDA_CHECK(cudaFree(vals_buf));
    CUDA_CHECK(cudaFree(idx_buf));

    // Return results to CPU to match reference
    Tensor samples_cpu = samples_dev.to(torch::kCPU);
    return samples_cpu;
}

} // namespace tk129280

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("run",
          &tk129280::run,
          pybind11::arg("probs"),
          pybind11::arg("top_k"),
          pybind11::arg("seed") = 0ULL,
          "Top-k sampling from probabilities (vocab size = 129280). "
          "Arguments:\n"
          "  probs: [B, 129280] float32/fp16/bf16 probabilities (on CPU or CUDA)\n"
          "  top_k: [B] int32 per-row K; filtering occurs if 0 < K < 129280\n"
          "  seed: optional uint64 seed for RNG (0 => random seed)\n"
          "Returns: [B] int64 sampled token indices on CPU");
}