Skip to content

README.md

Latest commit

 

History

History
101 lines (74 loc) · 3.39 KB

README.md

File metadata and controls

101 lines (74 loc) · 3.39 KB
Raw

tf-kernel

English | 中文

License Python CUDA PyTorch

tf-kernel provides optimized CUDA operations for TeleFuser, including fused elementwise operations, quantized GEMM, SageAttention, and block-sparse attention kernels. It supports SM80, SM90, and SM100 GPU families.

Important

No prebuilt tf-kernel package is currently published. Build and install the extension from source with this directory's Makefile. Direct pip install . and pip install -e . source builds are intentionally rejected.

Requirements

  • Python 3.10 or newer
  • PyTorch 2.11.0
  • CUDA Toolkit 12.8 or newer
  • CMake 3.26 or newer
  • An NVIDIA GPU in the SM80, SM90, or SM100 family

Build and Install

git clone https://github.com/Tele-AI/TeleFuser.git
cd TeleFuser/tf-kernel
make build-auto PYTHON=/path/to/venv/bin/python

build-auto detects the local GPU architecture. The Makefile builds a correctly tagged wheel under dist/ and installs it into the interpreter selected by PYTHON. This does not install or depend on the TeleFuser Python package.

Use an explicit architecture target for reproducible builds:

Target GPU family
make build-sm80 Ampere and Ada
make build-sm90 Hopper, including H100
make build-sm100 Blackwell
make build All supported architectures

For example:

make build-sm90 PYTHON=/path/to/venv/bin/python

Parallel Compilation

MAX_JOBS controls concurrent build jobs. TF_KERNEL_COMPILE_THREADS controls NVCC threads within each job:

make build-auto \
  PYTHON=/path/to/venv/bin/python \
  MAX_JOBS=16 \
  TF_KERNEL_COMPILE_THREADS=4

Higher values can reduce build time on a sufficiently provisioned host, but also increase CPU and memory pressure. For a resource-constrained build, start with MAX_JOBS=2 TF_KERNEL_COMPILE_THREADS=1.

Verify

Run the smoke test with the same interpreter passed to Make:

/path/to/venv/bin/python - <<'PY'
from pathlib import Path

import torch
import tf_kernel

print("tf-kernel:", tf_kernel.__version__)
print("PyTorch:", torch.__version__)
print("GPU:", torch.cuda.get_device_name())
print("extension:", Path(tf_kernel.common_ops.__file__).resolve())

x = torch.randn(8, 1024, device="cuda", dtype=torch.float16)
weight = torch.ones(1024, device="cuda", dtype=torch.float16)
assert torch.isfinite(tf_kernel.rmsnorm(x, weight)).all()
print("RMSNorm smoke test: OK")
PY

FP4 operators require SM100 or newer, so an FP4-unavailable import message is expected on Ampere and Hopper. The currently validated H100 build has a known failure in the architecture-selected SM90 SageAttention path; use another TeleFuser attention backend until its focused GPU test passes.

Development

make test PYTHON=/path/to/venv/bin/python
make format-check PYTHON=/path/to/venv/bin/python
make docs PYTHON=/path/to/venv/bin/python

See CONTRIBUTING.md for development and future release procedures. See the full installation and usage guide for compatibility, API examples, and troubleshooting.