replicate
/

mrope_get_position_ids

Model card Files Files and versions

mrope_get_position_ids / torch-ext /torch_binding.cpp

danieldk's picture

danieldk HF Staff

Sync with latest kernel-builder changes

b865e0f over 1 year ago

Raw History Blame Contribute Delete

343 Bytes

	#include "torch_binding.h"
	#include "registration.h"
	#include <torch/library.h>

	TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
	ops.def("get_position_ids(Tensor out, Tensor input_ids, Tensor "
	"image_grid_thw) -> ()");
	ops.impl("get_position_ids", torch::kCUDA, &get_position_ids);
	}

	REGISTER_EXTENSION(TORCH_EXTENSION_NAME)