Garrison

no human docs, do not use in production!

Garrison is a collection of reusable GPU kernels built with TileLang. Kernel families live in focused subpackages under src/garrison/; backend-specific implementations can sit alongside one another within a family.

The first family is quantized GEMM, with NF4 linear kernels for CUDA and Apple Metal, plus row-wise FP8 linear kernels for TorchAO Float8Tensor operands:

• garrison.quantized_gemm.nf4.cuda provides the CUDA forward and input-gradient kernels, TileLang's built-in autotuner, and an NF4 warmup adapter for CUDA graph capture.
• garrison.quantized_gemm.nf4.metal provides the Metal forward and input-gradient kernels with TileLang autotuning.
• garrison.quantized_gemm.float8.metal provides an inference-only Metal GEMM for E4M3FN PerRow() tensors, with first-use tile autotuning and a cached per-shape winner. Call prepare_float8_weight once for a static weight, then pass the prepared weight and row-wise FP8 activation to linear_float8_tilelang.
• garrison.quantized_gemm.float8.cuda provides the corresponding CUDA GEMM with a K pipeline and TileLang tile autotuning. Call prepare_float8_weight_cuda once before repeated calls to linear_float8_tilelang_cuda.

The FP8 paths currently support two-dimensional FP16-source E4M3FN tensors with row-wise scales. The Metal benchmark is benchmarks/benchmark_float8_metal.py; it checks against TorchAO's eager FP8 linear, persistent dequantized FP16 MPS weights, and a per-call FP8 dequantization plus FP16 MPS linear. It compares the original fixed Metal path with its MPS scale/cast operations against the autotuned kernel using a Metal scale-and-cast epilogue. The benchmark also compares the tuned kernel with both epilogue paths and supports an M sweep with --sweep-rows. Use set_float8_metal_autotune(False) or GARRISON_TILELANG_AUTOTUNE=0 to select the original fixed Metal tile without tuning.

The NF4 backends autotune their supported tile configurations on first use and reuse TileLang's in-process autotune results. CUDA also uses TileLang's disk cache; Metal uses the in-process cache with TileLang's Torch backend. Both NF4 backends use garrison.utils.warmup_kernels to run collections of zero-argument preparation callbacks. The CUDA-specific warmup_nf4_tilelang_cuda(model, rows) adapter applies it to the model's NF4 kernel shapes before CUDA graph capture. TileLang's TILELANG_CACHE_DIR and TILELANG_AUTO_TUNING_DISABLE_CACHE settings control its cache; GARRISON_TILELANG_AUTOTUNE=0 disables Garrison's first-use tuning and uses each backend's default tile.

The package root stays lightweight. Import a backend module directly so that using one backend does not load the others.

Installation

The base package depends on PyTorch and TileLang. The quantized-gemm extra adds TorchAO for code that creates or manages quantized weights:

uv sync --extra quantized-gemm

Development

Install the development tools with uv sync --group dev (add --extra quantized-gemm when working with quantized weights), then run:

uv run ruff check src
uv run ruff format --check src
uv run pyrefly check

The project uses a src/ package layout. Add new kernel families as src/garrison/<family>/, keeping separate backend modules where their runtime requirements differ. The py.typed marker enables type checking for package users.