| #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); |
|
|
