gpt-5-2025-08-07 / cuda4194a7

#include "kernel.h"

#include <torch/extension.h>
#include <pybind11/pybind11.h>

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

#include <chrono>
#include <vector>
#include <cstdint>

namespace py = pybind11;

static inline uint64_t seed_from_host() {
  // Simple 64-bit seed from time and address entropy
  uint64_t a = (uint64_t)std::chrono::high_resolution_clock::now().time_since_epoch().count();
  uint64_t b = (uint64_t)reinterpret_cast<uintptr_t>(&a);
  uint64_t c = 0x9E3779B97F4A7C15ULL;
  uint64_t seed = a ^ (b + 0x85ebca6b) ^ (c + (a<<6) + (a>>2));
  return seed;
}

static torch::Tensor top_k_top_p_sampling_from_probs_v151936(
    torch::Tensor probs,
    torch::Tensor top_k,
    torch::Tensor top_p)
{
  TORCH_CHECK(probs.dim() == 2, "probs must be 2D [batch, 151936]");
  TORCH_CHECK(probs.size(1) == VOCAB_SIZE_V151936, "vocab_size must be 151936");
  TORCH_CHECK(probs.dtype() == torch::kFloat32, "probs must be float32");
  TORCH_CHECK(top_k.dim() == 1 && top_k.size(0) == probs.size(0), "top_k shape must be [batch]");
  TORCH_CHECK(top_p.dim() == 1 && top_p.size(0) == probs.size(0), "top_p shape must be [batch]");
  TORCH_CHECK(top_k.dtype() == torch::kInt32, "top_k must be int32");
  TORCH_CHECK(top_p.dtype() == torch::kFloat32, "top_p must be float32");
  TORCH_CHECK(probs.is_cuda(), "probs must be a CUDA tensor");
  TORCH_CHECK(top_k.is_cuda() && top_p.is_cuda(), "top_k and top_p must be CUDA tensors");

  c10::cuda::CUDAGuard device_guard(probs.device());

  // Ensure contiguous
  probs = probs.contiguous();
  top_k = top_k.contiguous();
  top_p = top_p.contiguous();

  const int64_t batch = probs.size(0);
  auto options_out = torch::TensorOptions().dtype(torch::kInt64).device(probs.device());
  torch::Tensor samples = torch::empty({batch}, options_out);

  // Prepare CPU copies for routing decisions (small transfers)
  torch::Tensor top_k_h = top_k.to(torch::kCPU);
  torch::Tensor top_p_h = top_p.to(torch::kCPU);
  const int32_t* top_k_acc = top_k_h.data_ptr<int32_t>();
  const float* top_p_acc = top_p_h.data_ptr<float>();

  std::vector<int32_t> fast_rows;
  std::vector<int32_t> fallback_rows;
  fast_rows.reserve((size_t)batch);
  fallback_rows.reserve((size_t)batch);

  for (int64_t i = 0; i < batch; ++i) {
    int32_t k = top_k_acc[i];
    float p = top_p_acc[i];
    bool apply_top_k = (k > 0 && k < VOCAB_SIZE_V151936);
    bool needs_top_p = (p > 0.f && p < 1.f);
    // Fast path if no top-k and top-p is either <= 0 (argmax) or >= 1 (full sample)
    if (!apply_top_k && !needs_top_p) {
      fast_rows.push_back((int32_t)i);
    } else {
      fallback_rows.push_back((int32_t)i);
    }
  }

  cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream();

  // Seed
  uint64_t seed = seed_from_host();

  // Fast path launch
  if (!fast_rows.empty()) {
    // Move row indices to device
    torch::Tensor rows_fast_dev = torch::from_blob(
        fast_rows.data(),
        {(int64_t)fast_rows.size()},
        torch::TensorOptions().dtype(torch::kInt32).device(torch::kCPU)).clone().to(probs.device());

    launch_fast_path_kernel(
        probs.data_ptr<float>(),
        top_k.data_ptr<int32_t>(),
        top_p.data_ptr<float>(),
        samples.data_ptr<int64_t>(),
        rows_fast_dev.data_ptr<int32_t>(),
        static_cast<int32_t>(fast_rows.size()),
        seed,
        stream);
  }

  // Fallback processing using Thrust host helpers compiled by NVCC
  if (!fallback_rows.empty()) {
    // Pre-generate uniforms for fallback rows on device and copy to host
    auto uniforms_opts = torch::TensorOptions().dtype(torch::kFloat32).device(probs.device());
    torch::Tensor uniforms_dev = torch::empty({(int64_t)fallback_rows.size()}, uniforms_opts);
    launch_fill_uniform(uniforms_dev.data_ptr<float>(), static_cast<int32_t>(fallback_rows.size()), seed ^ 0xCAFEBABE12345678ULL, stream);

    std::vector<float> uniforms_host(fallback_rows.size());
    CUDA_CHECK(cudaMemcpyAsync(
        uniforms_host.data(),
        uniforms_dev.data_ptr<float>(),
        sizeof(float) * fallback_rows.size(),
        cudaMemcpyDeviceToHost,
        stream));
    CUDA_CHECK(cudaStreamSynchronize(stream));

    // Workspaces (reused across rows)
    auto idx_opts = torch::TensorOptions().dtype(torch::kInt32).device(probs.device());
    auto val_opts = torch::TensorOptions().dtype(torch::kFloat32).device(probs.device());
    torch::Tensor idx_buf = torch::empty({VOCAB_SIZE_V151936}, idx_opts);
    torch::Tensor values_buf = torch::empty({VOCAB_SIZE_V151936}, val_opts);
    torch::Tensor cdf_buf = torch::empty({VOCAB_SIZE_V151936}, val_opts);

    int32_t* idx_ptr = idx_buf.data_ptr<int32_t>();
    float* values_ptr = values_buf.data_ptr<float>();
    float* cdf_ptr = cdf_buf.data_ptr<float>();
    const float* probs_ptr = probs.data_ptr<float>();
    int64_t* samples_ptr = samples.data_ptr<int64_t>();

    for (size_t rr = 0; rr < fallback_rows.size(); ++rr) {
      int32_t row = fallback_rows[rr];
      float p = top_p_acc[row];
      int32_t k = top_k_acc[row];
      float u = uniforms_host[rr];

      thrust_process_row(
          probs_ptr,
          row,
          VOCAB_SIZE_V151936,
          k,
          p,
          u,
          idx_ptr,
          values_ptr,
          cdf_ptr,
          samples_ptr,
          stream);
    }
  }

  // Ensure all device work complete before returning
  CUDA_CHECK(cudaGetLastError());
  CUDA_CHECK(cudaStreamSynchronize(stream));

  return samples;
}

// Python binding entry point: run(probs, top_k, top_p)
torch::Tensor run(torch::Tensor probs, torch::Tensor top_k, torch::Tensor top_p) {
  c10::cuda::CUDAGuard device_guard(probs.device());
  return top_k_top_p_sampling_from_probs_v151936(probs, top_k, top_p);
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
  m.def("run", &run, "top_k_top_p_sampling_from_probs_v151936",
        py::arg("probs"), py::arg("top_k"), py::arg("top_p"));
}