#include <torch/library.h>

#include "registration.h"
#include "torch_binding.h"

TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
  ops.def("fused_sample(Tensor! tokens, Tensor logits, Tensor? prev_tokens, float temperature, int top_k, float top_p, float min_p, float repetition_penalty, int seed, int offset) -> ()");
  ops.def("fused_filter(Tensor! out_logits, Tensor logits, Tensor? prev_tokens, float temperature, int top_k, float top_p, float min_p, float repetition_penalty) -> ()");
  ops.def("spec_verify(Tensor! accept_len, Tensor! out_tokens, Tensor target_logits, Tensor draft_logits, Tensor draft_tokens, float temperature, int seed, int offset) -> ()");

#if defined(CUDA_KERNEL) || defined(ROCM_KERNEL)
  ops.impl("fused_sample", torch::kCUDA, &fused_sample);
  ops.impl("fused_filter", torch::kCUDA, &fused_filter);
  ops.impl("spec_verify",  torch::kCUDA, &spec_verify);
#endif
}

REGISTER_EXTENSION(TORCH_EXTENSION_NAME)