feat: add setup.py for local CUDA development builds

Local build alternative to kernel-builder (nix) for development and
training yaml workflows. Builds _activation extension with nvcc for
sm_80/89/90/100 (B200 support with CUDA 12.8+).

Usage: pip install --no-build-isolation -e .

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

registration.h

@@ -0,0 +1,14 @@
+#pragma once
+
+// Local build compatibility shim for kernel-builder's registration.h
+
+#include <torch/library.h>
+#include <torch/extension.h>
+
+// TORCH_LIBRARY_EXPAND may not be defined in all PyTorch versions
+#ifndef TORCH_LIBRARY_EXPAND
+#define TORCH_LIBRARY_EXPAND(ns, m) TORCH_LIBRARY(ns, m)
+#endif
+
+// Generate the PyInit_<name> entry point for the shared library
+#define REGISTER_EXTENSION(name) PYBIND11_MODULE(name, m) {}

setup.py

@@ -0,0 +1,81 @@
+"""Local CUDA build for activation kernels.
+
+Usage:
+    pip install -e .          # editable install
+    python setup.py build_ext --inplace  # build only
+
+The built extension is named '_activation' and can be loaded via:
+    import _activation
+    torch.ops._activation.rms_norm(...)
+"""
+
+import os
+from pathlib import Path
+
+import torch
+from setuptools import setup
+from torch.utils.cpp_extension import BuildExtension, CUDAExtension
+
+ROOT = Path(__file__).parent
+
+CUDA_SOURCES = [
+    "activation/poly_norm.cu",
+    "activation/fused_mul_poly_norm.cu",
+    "activation/rms_norm.cu",
+    "activation/fused_add_rms_norm.cu",
+    "activation/grouped_poly_norm.cu",
+]
+
+CPP_SOURCES = [
+    "torch-ext/torch_binding.cpp",
+]
+
+# Include dirs: project root (for registration.h, activation/*.h)
+# and torch-ext/ (for torch_binding.h)
+INCLUDE_DIRS = [
+    str(ROOT),
+    str(ROOT / "activation"),
+    str(ROOT / "torch-ext"),
+]
+
+# CUDA flags matching the existing kernel style
+NVCC_FLAGS = [
+    "-O3",
+    "--use_fast_math",
+    "-std=c++17",
+    # Generate code for common architectures
+    "-gencode=arch=compute_80,code=sm_80",   # A100
+    "-gencode=arch=compute_89,code=sm_89",   # L40/4090
+    "-gencode=arch=compute_90,code=sm_90",   # H100
+]
+
+# Check for B200 support (sm_100, requires CUDA 12.8+)
+cuda_version = tuple(int(x) for x in torch.version.cuda.split(".")[:2])
+if cuda_version >= (12, 8):
+    NVCC_FLAGS.append("-gencode=arch=compute_100,code=sm_100")
+
+CXX_FLAGS = ["-O3", "-std=c++17"]
+
+ext_modules = [
+    CUDAExtension(
+        name="_activation",
+        sources=[str(ROOT / s) for s in CPP_SOURCES + CUDA_SOURCES],
+        include_dirs=INCLUDE_DIRS,
+        extra_compile_args={
+            "cxx": CXX_FLAGS,
+            "nvcc": NVCC_FLAGS,
+        },
+    ),
+]
+
+setup(
+    name="activation",
+    version="0.1.0",
+    description="Custom CUDA normalization kernels for LLM training",
+    ext_modules=ext_modules,
+    cmdclass={"build_ext": BuildExtension},
+    packages=["activation"],
+    package_dir={"activation": "torch-ext/activation"},
+    python_requires=">=3.10",
+    install_requires=["torch>=2.7"],
+)