gpt-5-2025-08-07 / cudad68ec9

#include <torch/extension.h>
#include <cuda_runtime.h>
#include <vector>
#include <chrono>
#include <c10/cuda/CUDAStream.h>

#include "kernel.h"

namespace py = pybind11;

// Helper to get CUDA stream from PyTorch
static inline cudaStream_t get_cuda_stream() {
  return c10::cuda::getCurrentCUDAStream().stream();
}

static inline void check_input(const torch::Tensor& probs,
                               const torch::Tensor& top_k) {
  TORCH_CHECK(probs.is_cuda(), "probs must be a CUDA tensor");
  TORCH_CHECK(top_k.is_cuda(), "top_k must be a CUDA tensor");
  TORCH_CHECK(probs.dtype() == torch::kFloat32, "probs must be float32");
  TORCH_CHECK(top_k.dtype() == torch::kInt32, "top_k must be int32");

  TORCH_CHECK(probs.dim() == 2, "probs must be 2D [batch_size, vocab_size]");
  TORCH_CHECK(top_k.dim() == 1, "top_k must be 1D [batch_size]");

  int64_t batch_size = probs.size(0);
  int64_t vocab_size = probs.size(1);
  TORCH_CHECK(vocab_size == VOCAB_SIZE_CONST,
              "vocab_size must be exactly ", VOCAB_SIZE_CONST, ", got ", vocab_size);
  TORCH_CHECK(top_k.size(0) == batch_size, "top_k length must equal batch_size");
  TORCH_CHECK(probs.is_contiguous(), "probs must be contiguous");
  TORCH_CHECK(top_k.is_contiguous(), "top_k must be contiguous");
}

torch::Tensor run(torch::Tensor probs,
                  torch::Tensor top_k,
                  c10::optional<int64_t> seed_opt = c10::nullopt) {
  check_input(probs, top_k);

  const int64_t batch_size = probs.size(0);
  const int64_t vocab_size = probs.size(1);
  if (batch_size == 0) {
    return torch::empty({0}, probs.options().dtype(torch::kInt64));
  }

  auto options_int32 = torch::TensorOptions().dtype(torch::kInt32).device(probs.device());
  auto options_float = torch::TensorOptions().dtype(torch::kFloat32).device(probs.device());
  auto options_long  = torch::TensorOptions().dtype(torch::kInt64).device(probs.device());

  // Output tensor: sampled token indices
  torch::Tensor samples = torch::empty({batch_size}, options_long);

  // Determine if any row requires top-k filtering (0 < k < V)
  bool need_sort = false;
  {
    std::vector<int32_t> h_top_k(batch_size);
    CUDA_CHECK(cudaMemcpyAsync(h_top_k.data(),
                               top_k.data_ptr<int32_t>(),
                               sizeof(int32_t) * batch_size,
                               cudaMemcpyDeviceToHost,
                               get_cuda_stream()));
    CUDA_CHECK(cudaStreamSynchronize(get_cuda_stream()));
    for (int64_t i = 0; i < batch_size; ++i) {
      int k = h_top_k[i];
      if (k > 0 && k < vocab_size) { need_sort = true; break; }
    }
  }

  // Temporary buffers for sorting
  torch::Tensor indices_in;
  torch::Tensor keys_sorted;
  torch::Tensor vals_sorted;

  if (need_sort) {
    const int64_t total_elems = batch_size * vocab_size;
    indices_in  = torch::empty({total_elems}, options_int32);
    keys_sorted = torch::empty({total_elems}, options_float);
    vals_sorted = torch::empty({total_elems}, options_int32);

    // Build segment offsets [0, V, 2V, ..., B*V]
    std::vector<int32_t> h_offsets(batch_size + 1);
    for (int64_t i = 0; i <= batch_size; ++i) {
      int64_t off = i * vocab_size;
      h_offsets[i] = static_cast<int32_t>(off);
    }
    torch::Tensor d_offsets = torch::empty({batch_size + 1}, options_int32);
    CUDA_CHECK(cudaMemcpyAsync(d_offsets.data_ptr<int32_t>(),
                               h_offsets.data(),
                               sizeof(int32_t) * (batch_size + 1),
                               cudaMemcpyHostToDevice,
                               get_cuda_stream()));

    // Initialize indices [0..V-1] repeated per row
    launch_init_indices_kernel(
        indices_in.data_ptr<int32_t>(),
        static_cast<int>(batch_size),
        static_cast<int>(vocab_size),
        get_cuda_stream());

    // Perform segmented descending sort (pairs: probs with indices)
    segmented_sort_desc_pairs(
        probs.data_ptr<float>(),
        indices_in.data_ptr<int32_t>(),
        keys_sorted.data_ptr<float>(),
        vals_sorted.data_ptr<int32_t>(),
        d_offsets.data_ptr<int32_t>(),
        static_cast<int>(batch_size),
        static_cast<int>(vocab_size),
        get_cuda_stream());
  } else {
    // Allocate minimal dummies to satisfy kernel signature; they won't be used
    keys_sorted = torch::empty({1}, options_float);
    vals_sorted = torch::empty({1}, options_int32);
  }

  // Seed handling: default is time-based if not provided
  unsigned long long seed;
  if (seed_opt.has_value()) {
    seed = static_cast<unsigned long long>(seed_opt.value());
  } else {
    seed = static_cast<unsigned long long>(
        std::chrono::high_resolution_clock::now().time_since_epoch().count());
  }

  // Launch sampling kernel: per-row block
  launch_sample_sorted_or_full_kernel(
      probs.data_ptr<float>(),
      keys_sorted.data_ptr<float>(),
      vals_sorted.data_ptr<int32_t>(),
      top_k.data_ptr<int32_t>(),
      samples.data_ptr<int64_t>(),
      static_cast<int>(batch_size),
      static_cast<int>(vocab_size),
      seed,
      get_cuda_stream());

  // Synchronize stream to ensure completion before returning to Python
  CUDA_CHECK(cudaStreamSynchronize(get_cuda_stream()));
  return samples;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
  m.def("run", &run,
        "top_k_sampling_from_probs_v128256 (CUDA, B200-optimized)",
        py::arg("probs"),
        py::arg("top_k"),
        py::arg("seed") = py::none());
}