#pragma once

#include <torch/torch.h>
#include <c10/util/Optional.h>

void fused_sample(torch::Tensor &tokens,
                  torch::Tensor const &logits,
                  c10::optional<torch::Tensor> const &prev_tokens,
                  double temperature, int64_t top_k, double top_p, double min_p,
                  double repetition_penalty, int64_t seed, int64_t offset);

void fused_filter(torch::Tensor &out_logits,
                  torch::Tensor const &logits,
                  c10::optional<torch::Tensor> const &prev_tokens,
                  double temperature, int64_t top_k, double top_p, double min_p,
                  double repetition_penalty);

void spec_verify(torch::Tensor &accept_len,
                 torch::Tensor &out_tokens,
                 torch::Tensor const &target_logits,
                 torch::Tensor const &draft_logits,
                 torch::Tensor const &draft_tokens,
                 double temperature, int64_t seed, int64_t offset);