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