Kernels:

phanerozoic
/

muon

Kernel card Files Files and versions

muon / torch-ext /torch_binding.cpp

CharlesCNorton

Add kernel source and model card

8791e67 25 days ago

584 Bytes

	#include <torch/library.h>

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

	TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
	ops.def("muon_orthogonalize(Tensor Xb, int steps, float eps) -> Tensor");
	ops.def("muon_sr_update(Tensor(a!) param, Tensor O, float lr, float scale, float wd, "
	"int seed, int step, int base_index) -> ()");

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

	REGISTER_EXTENSION(TORCH_EXTENSION_NAME)