Metadata-Version: 2.4
Name: freesolo-chalk
Version: 0.5.7
Summary: Custom Triton/CUDA kernels that complement Liger Kernel for LLM post-training
License: BSD 2-CLAUSE LICENSE
        
        Copyright (c) 2026 Freesolo, Inc.
        All rights reserved.
        
        Redistribution and use in source and binary forms, with or without
        modification, are permitted provided that the following conditions are met:
        
        1. Redistributions of source code must retain the above copyright notice, this
           list of conditions and the following disclaimer.
        
        2. Redistributions in binary form must reproduce the above copyright notice,
           this list of conditions and the following disclaimer in the documentation
           and/or other materials provided with the distribution.
        
        THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
        AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
        IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
        DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
        FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
        DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
        SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
        CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
        OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
        OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
        
Project-URL: Homepage, https://github.com/freesolo-co/chalk
Description-Content-Type: text/markdown
License-File: LICENSE
License-File: NOTICE
Requires-Dist: torch>=2.1.2
Requires-Dist: triton>=2.3.1
Provides-Extra: dev
Requires-Dist: transformers>=4.52.0; extra == "dev"
Requires-Dist: matplotlib>=3.7.2; extra == "dev"
Requires-Dist: ruff>=0.12.0; extra == "dev"
Requires-Dist: pytest>=7.1.2; extra == "dev"
Requires-Dist: pytest-xdist; extra == "dev"
Requires-Dist: pytest-cov; extra == "dev"
Requires-Dist: datasets>=2.19.2; extra == "dev"
Dynamic: license-file
Dynamic: provides-extra
Dynamic: requires-dist

# Chalk

**Custom Triton/CUDA kernels that complement [Liger Kernel](https://github.com/linkedin/Liger-Kernel).**

`pip install freesolo-chalk`

Liger fuses the cross-entropy, activation, and RMSNorm paths. Chalk fills the gaps that
matter for Freesolo's [Flash](https://github.com/freesolo-co/flash) post-training stack —
fused GEMMs, the LoRA-delta matmuls, the QKV norm+RoPE epilogue, embedding gather, and
FP8 frozen-base GEMMs — each behind a documented, benchmarked, opt-in entry point.

Chalk's repository layout and conventions intentionally mirror Liger Kernel.

## Layout

```
src/chalk/
  ops/           # raw Triton/CUDA kernels + autograd.Function wrappers
  transformers/  # model-level installers that monkeypatch kernels into HF modules
  utils.py       # device detection helpers
test/            # correctness + gating tests (mirrors test/ops, test/transformers)
```

## Design principles

- **Worker-side kernel library.** Like Liger, chalk depends on `torch` + `triton` — it is
  meant to be installed where kernels actually run (the GPU worker), so consumers should
  depend on it from a `gpu` extra rather than their base install. Importing the top-level
  `chalk` package is still cheap (kernels lazy-load), so probing `chalk.utils.infer_device()`
  never forces a heavy import.
- **Complements, not replaces, Liger.** Liger fuses CE / activation / RMSNorm; chalk fuses
  the GEMMs, LoRA delta, QKV epilogue, embedding, and FP8 base.
- **Safe fallback.** Every installer is arch-gated, runs a numeric self-test on install, and
  silently falls back to the eager / Liger path on any import / compile / self-test failure.
  All installers but one patch only frozen `nn.Linear` layers (never trainable / PEFT-wrapped
  layers). The exception is the LoRA-delta installer (`install_fused_lora_delta`), the one
  kernel that intentionally accelerates the _trainable_ PEFT LoRA path: it monkeypatches
  PEFT's dense `lora.Linear.forward` to route the adapter delta through the fused kernel. Its
  self-test checks the fused forward _and_ the `x` / `A` / `B` gradients against torch autograd
  (bf16 tolerance) before patching, and at runtime it falls back to stock PEFT per call for
  everything the kernel does not replicate exactly (merged / disabled adapters, mixed-adapter
  batches, LoRA variants, dropout, `lora_bias`, quantized / non-dense bases, and any non-CUDA
  tensor) — so a failed self-test or any unsupported layer leaves training numerics on the
  stock PEFT path.
- **Opt-in & evidence-based.** Kernels are off unless explicitly enabled, and every kept
  kernel has end-to-end loss-curve evidence — not just a microbenchmark.

## Development

```bash
pip install -e '.[dev]'
make checkstyle   # ruff check + format
make test         # pytest with coverage
```

## Status

Intentionally minimal to start — kernels are landed one at a time under `chalk/ops` +
`chalk/transformers`, each with correctness tests.

## License

BSD-2-Clause. See [LICENSE](LICENSE) and [NOTICE](NOTICE).
