claude-opus-4-1-20250805 / cuda8688f2

#include <torch/extension.h>
#include <cuda_runtime.h>
#include <vector>
#include <stdexcept>
#include "kernel.h"

// Error checking macros
#define CHECK_CUDA(x) TORCH_CHECK(x.device().is_cuda(), #x " must be a CUDA tensor")
#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)

torch::Tensor run(
    torch::Tensor probs,
    torch::Tensor top_k
) {
    // Input validation
    CHECK_INPUT(probs);
    CHECK_INPUT(top_k);
    
    // Check dimensions
    TORCH_CHECK(probs.dim() == 2, "probs must be 2D tensor");
    TORCH_CHECK(top_k.dim() == 1, "top_k must be 1D tensor");
    
    int batch_size = probs.size(0);
    int vocab_size = probs.size(1);
    
    // Validate vocab size
    TORCH_CHECK(vocab_size == 128256, 
                "vocab_size must be 128256, got " + std::to_string(vocab_size));
    
    // Validate batch sizes match
    TORCH_CHECK(top_k.size(0) == batch_size, 
                "top_k batch size must match probs batch size");
    
    // Ensure correct dtypes
    torch::Tensor probs_f32 = probs.to(torch::kFloat32).contiguous();
    torch::Tensor top_k_i32 = top_k.to(torch::kInt32).contiguous();
    
    // Create output tensor
    auto options = torch::TensorOptions()
        .dtype(torch::kInt64)
        .device(probs.device());
    torch::Tensor samples = torch::empty({batch_size}, options);
    
    // Get current CUDA stream
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    
    // Launch kernel
    launchTopKSampling(
        probs_f32.data_ptr<float>(),
        top_k_i32.data_ptr<int32_t>(),
        samples.data_ptr<int64_t>(),
        batch_size,
        stream
    );
    
    // Ensure kernel completion
    cudaError_t error = cudaStreamSynchronize(stream);
    if (error != cudaSuccess) {
        throw std::runtime_error(
            std::string("CUDA synchronization error: ") + cudaGetErrorString(error)
        );
    }
    
    // Check for kernel errors
    error = cudaGetLastError();
    if (error != cudaSuccess) {
        throw std::runtime_error(
            std::string("CUDA kernel error: ") + cudaGetErrorString(error)
        );
    }
    
    return samples;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("run", &run, "Top-k sampling from probability distributions",
          pybind11::arg("probs"), 
          pybind11::arg("top_k"));
}