theapemachine's picture
Add sparse transformer v19 with Triton-backed KNN scheduler and various backward modes. Includes utilities for synthetic data generation and model training. Implements chunked sparse updates and integrates with existing sparse linear layers.
bc1b8eb
raw
history blame contribute delete
435 Bytes
from setuptools import setup
from torch.utils.cpp_extension import CppExtension, BuildExtension
setup(
name='sparse_linear',
ext_modules=[
CppExtension(
name='sparse_linear',
sources=['sparse_linear_ops.mm'],
extra_compile_args=['-ObjC++'],
extra_link_args=['-framework', 'Metal', '-framework', 'Foundation']
)
],
cmdclass={'build_ext': BuildExtension}
)