CharlesCNorton

Add logits-processor source, tests, and model card

9aa33e7 24 days ago

591 Bytes

	#include <torch/library.h>

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

	TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
	ops.def("apply_token_bitmask_inplace(Tensor(a!) logits, Tensor bitmask) -> ()");
	ops.def("sample(Tensor logits, Tensor temperature, Tensor top_p, Tensor top_k, "
	"Tensor min_p, int seed, int offset) -> Tensor");

	#if defined(CUDA_KERNEL) \|\| defined(ROCM_KERNEL)
	ops.impl("apply_token_bitmask_inplace", torch::kCUDA, &apply_token_bitmask_inplace);
	ops.impl("sample", torch::kCUDA, &sample);
	#endif
	}

	REGISTER_EXTENSION(TORCH_EXTENSION_NAME)