freesolo-chalk 0.1.0__tar.gz

freesolo_chalk-0.1.0/.github/workflows/ci.yml

@@ -0,0 +1,43 @@
+name: ci
+
+on:
+  workflow_dispatch:
+  push:
+    branches: [main, dev]
+  pull_request:
+    branches: [main, dev]
+
+jobs:
+  checkstyle:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.11"
+      - name: Install ruff
+        run: pip install -r dev/fmt-requirements.txt
+      - name: Lint
+        run: ruff check --output-format=concise .
+      - name: Format check
+        run: ruff format --check --diff .
+
+  test:
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: ["3.10", "3.11", "3.12"]
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python-version }}
+      # CPU runner: install the package + pytest only (no torch/triton). GPU-only
+      # kernel correctness tests self-skip when torch/triton/CUDA are unavailable.
+      - name: Install
+        run: |
+          python -m pip install --upgrade pip
+          pip install -e .
+          pip install pytest pytest-cov
+      - name: Test (CPU)
+        run: python -m pytest --disable-warnings test/

freesolo_chalk-0.1.0/.github/workflows/main-source-guard.yml

@@ -0,0 +1,20 @@
+name: Main source guard
+on:
+  pull_request:
+    branches: [main]
+permissions:
+  contents: read
+jobs:
+  source-is-dev:
+    name: Source branch is dev
+    runs-on: ubuntu-latest
+    steps:
+      - name: Require PRs into main to originate from dev
+        env:
+          HEAD_REF: ${{ github.head_ref }}
+        run: |
+          if [ "$HEAD_REF" != "dev" ]; then
+            echo "::error::PRs into main must come from 'dev' (got '$HEAD_REF'). Merge into dev, then promote dev -> main."
+            exit 1
+          fi
+          echo "Source branch '$HEAD_REF' is allowed."

freesolo_chalk-0.1.0/.gitignore

@@ -0,0 +1,28 @@
+__pycache__/
+*.egg-info/
+site/
+.cache/
+.venv/
+venv/
+.ipynb_checkpoints/
+.vscode/
+.idea/
+
+# Misc
+.DS_Store
+
+# Build
+build/
+dist/
+
+# Lockfiles
+uv.lock
+
+# Benchmark images
+benchmark/visualizations
+
+# Coverage
+.coverage
+htmlcov/
+.pytest_cache/
+.ruff_cache/

freesolo_chalk-0.1.0/.pre-commit-config.yaml

@@ -0,0 +1,10 @@
+repos:
+- repo: https://github.com/astral-sh/ruff-pre-commit
+  # Ruff version.
+  rev: v0.14.11
+  hooks:
+    # Run the linter.
+    - id: ruff-check
+      args: [ --fix ]
+    # Run the formatter.
+    - id: ruff-format

freesolo_chalk-0.1.0/LICENSE

@@ -0,0 +1,25 @@
+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.

freesolo_chalk-0.1.0/Makefile

@@ -0,0 +1,41 @@
+.PHONY: test checkstyle coverage run-benchmarks all
+
+all: checkstyle test
+
+# Command to run pytest for correctness tests
+test:
+	python -m pytest --disable-warnings \
+		--cov=src/chalk \
+		--cov-report=term-missing \
+		test/
+
+# Command to run coverage report
+coverage:
+	coverage report -m
+
+# Command to run ruff for linting and formatting code
+checkstyle:
+	ruff check --output-format=concise .; ruff_check_status=$$?; \
+	ruff format --check --diff .; ruff_format_status=$$?; \
+	ruff check . --fix; \
+	ruff format .; \
+	if [ $$ruff_check_status -ne 0 ] || [ $$ruff_format_status -ne 0 ]; then \
+		exit 1; \
+	fi
+
+# Command to run all benchmark scripts and update benchmarking data file
+# By default this doesn't overwrite existing data for the same benchmark experiment
+# run with `make run-benchmarks OVERWRITE=1` to overwrite existing benchmark data
+BENCHMARK_DIR = benchmark/scripts
+BENCHMARK_SCRIPTS = $(wildcard $(BENCHMARK_DIR)/benchmark_*.py)
+OVERWRITE ?= 0
+
+run-benchmarks:
+	@for script in $(BENCHMARK_SCRIPTS); do \
+		echo "Running benchmark: $$script"; \
+		if [ $(OVERWRITE) -eq 1 ]; then \
+			python $$script --overwrite; \
+		else \
+			python $$script; \
+		fi; \
+	done

freesolo_chalk-0.1.0/NOTICE

@@ -0,0 +1,26 @@
+Copyright 2026 Freesolo, Inc.
+All Rights Reserved.
+
+Licensed under the BSD 2-Clause License (the "License"). See LICENSE in the
+project root for license information.
+
+Chalk provides custom Triton/CUDA kernels designed to complement Liger Kernel
+(https://github.com/linkedin/Liger-Kernel, BSD 2-Clause). Its repository layout,
+tooling, and conventions intentionally mirror Liger Kernel.
+
+This product may contain code derived from the following open source projects:
+
+1. Liger Kernel
+   Copyright (c) 2024 LinkedIn Corporation
+   Licensed under the BSD 2-Clause License
+   Source: https://github.com/linkedin/Liger-Kernel
+
+   Repository structure (src/ layout, ops/transformers split), benchmarking
+   harness conventions, and test scaffolding were referenced from this project.
+
+2. Triton
+   Copyright (c) 2023 OpenAI
+   Licensed under the MIT License
+   Source: https://github.com/openai/triton
+
+For full license texts, please refer to the respective project repositories.

freesolo_chalk-0.1.0/PKG-INFO

@@ -0,0 +1,104 @@
+Metadata-Version: 2.4
+Name: freesolo-chalk
+Version: 0.1.0
+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)
+benchmark/       # speed + fp32-correctness A/B harness
+```
+
+## 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,
+  patches only frozen `nn.Linear` layers (never trainable / PEFT-wrapped layers), and
+  silently falls back to the eager / Liger path on any import / compile / self-test failure.
+- **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
+make run-benchmarks
+```
+
+## Status
+
+Intentionally minimal to start — kernels are landed one at a time under `chalk/ops` +
+`chalk/transformers`, each with correctness tests and benchmark evidence.
+
+## License
+
+BSD-2-Clause. See [LICENSE](LICENSE) and [NOTICE](NOTICE).

freesolo_chalk-0.1.0/README.md

@@ -0,0 +1,56 @@
+# 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)
+benchmark/       # speed + fp32-correctness A/B harness
+```
+
+## 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,
+  patches only frozen `nn.Linear` layers (never trainable / PEFT-wrapped layers), and
+  silently falls back to the eager / Liger path on any import / compile / self-test failure.
+- **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
+make run-benchmarks
+```
+
+## Status
+
+Intentionally minimal to start — kernels are landed one at a time under `chalk/ops` +
+`chalk/transformers`, each with correctness tests and benchmark evidence.
+
+## License
+
+BSD-2-Clause. See [LICENSE](LICENSE) and [NOTICE](NOTICE).

freesolo_chalk-0.1.0/benchmark/README.md

@@ -0,0 +1,20 @@
+## Benchmarking Chalk Kernels
+
+Chalk's benchmarking harness mirrors [Liger Kernel's](https://github.com/linkedin/Liger-Kernel/tree/main/benchmark):
+each kernel ships a `benchmark/scripts/benchmark_<kernel>.py` that A/Bs the chalk kernel
+against the eager / Liger / `torch` baseline across model configs and sequence lengths,
+reporting both speed and a fp32-correctness check.
+
+### Running
+
+```bash
+make run-benchmarks            # run every benchmark/scripts/benchmark_*.py
+make run-benchmarks OVERWRITE=1  # overwrite existing recorded data
+```
+
+### Conventions
+
+* Every kept kernel must have **end-to-end** evidence (a real LoRA-SFT / GRPO loss-curve
+  A/B), not just a microbenchmark — a microbenchmark win that disappears E2E is not a win.
+* Record the GPU (A40 / A6000 / H100 / H200) the numbers were measured on; cuBLAS strength
+  differs enough across SKUs that a win on one can vanish on another.

freesolo_chalk-0.1.0/dev/fmt-requirements.txt

@@ -0,0 +1 @@
+ruff>=0.1.6

freesolo_chalk-0.1.0/pyproject.toml

@@ -0,0 +1,84 @@
+[build-system]
+requires = ["setuptools>=42", "wheel", "setuptools-scm"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "freesolo-chalk"
+version = "0.1.0"
+description = "Custom Triton/CUDA kernels that complement Liger Kernel for LLM post-training"
+urls = { "Homepage" = "https://github.com/freesolo-co/chalk" }
+readme = { file = "README.md", content-type = "text/markdown" }
+license = { file = "LICENSE" }
+dynamic = ["dependencies", "optional-dependencies"]
+
+[tool.setuptools]
+package-dir = {"" = "src"}
+
+[tool.setuptools.packages.find]
+where = ["src"]
+include = ["chalk*"]
+namespaces = false
+
+[tool.pytest.ini_options]
+pythonpath = ["src", "."]
+addopts = [
+    "--cov=src/chalk",
+    "--cov-report=term-missing",
+    "--cov-report=html",
+    "--cov-config=pyproject.toml",
+    "--durations=0"
+]
+python_files = "test_*.py"
+testpaths = ["test/"]
+markers = [
+    "gpu: requires a CUDA GPU + torch/triton (skipped on CPU CI)",
+]
+
+[tool.coverage.run]
+branch = true
+parallel = true
+source = ["src/chalk"]
+concurrency = ["multiprocessing"]
+
+[tool.coverage.paths]
+chalk = [
+  "src/chalk",
+  "*/site-packages/chalk"
+]
+
+[tool.coverage.report]
+omit = ["test/*"]
+show_missing = true
+skip_covered = false
+
+
+[tool.ruff]
+line-length = 120
+target-version = "py310"
+respect-gitignore = true
+src = ["src"]
+
+[tool.ruff.lint]
+select = [
+    "E",  # pycodestyle errors
+    "F",  # pyflakes
+    "I",  # isort
+]
+ignore = ["E501", "B006", "E731", "A002", "E203"]
+
+exclude = [
+    ".git",
+    "__pycache__",
+    ".venv",
+]
+
+[tool.ruff.format]
+quote-style = "double"
+indent-style = "space"
+skip-magic-trailing-comma = false
+line-ending = "auto"
+
+[tool.ruff.lint.isort]
+known-first-party = ["chalk"]
+force-single-line = true
+lines-between-types = 1

freesolo_chalk-0.1.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

freesolo_chalk-0.1.0/setup.py

@@ -0,0 +1,102 @@
+# setup.py
+
+import subprocess
+
+from typing import Literal
+
+from setuptools import find_packages
+from setuptools import setup
+
+
+def get_default_dependencies():
+    """Determine the appropriate dependencies based on detected hardware."""
+    platform = get_platform()
+
+    if platform in ["cuda", "cpu"]:
+        return [
+            "torch>=2.1.2",
+            "triton>=2.3.1",
+        ]
+    elif platform == "rocm":
+        return [
+            "triton>=3.0.0",
+        ]
+    elif platform == "xpu":
+        return [
+            "torch>=2.6.0",
+        ]
+
+
+def get_optional_dependencies():
+    """Get optional dependency groups."""
+    dev_deps = [
+        "transformers>=4.52.0",
+        "matplotlib>=3.7.2",
+        "ruff>=0.12.0",
+        "pytest>=7.1.2",
+        "pytest-xdist",
+        "pytest-cov",
+        "datasets>=2.19.2",
+    ]
+    return {
+        "dev": dev_deps,
+    }
+
+
+def is_xpu_available():
+    """Check if Intel XPU is available."""
+    try:
+        subprocess.run(["xpu-smi"], check=True)
+        return True
+    except (subprocess.SubprocessError, FileNotFoundError):
+        pass
+
+    try:
+        result = subprocess.run("sycl-ls", check=True, capture_output=True, shell=True)
+        if "level_zero:gpu" in result.stdout.decode():
+            return True
+    except (subprocess.SubprocessError, FileNotFoundError):
+        pass
+
+    return False
+
+
+def get_platform() -> Literal["cuda", "rocm", "cpu", "xpu"]:
+    """Detect whether the system has NVIDIA or AMD GPU without a torch dependency."""
+    # Try nvidia-smi first
+    try:
+        subprocess.run(["nvidia-smi"], check=True)
+        print("NVIDIA GPU detected")
+        return "cuda"
+    except (subprocess.SubprocessError, FileNotFoundError):
+        # If nvidia-smi fails, check for ROCm
+        try:
+            subprocess.run(["rocm-smi"], check=True)
+            print("ROCm GPU detected")
+            return "rocm"
+        except (subprocess.SubprocessError, FileNotFoundError):
+            if is_xpu_available():
+                print("Intel GPU detected")
+                return "xpu"
+            else:
+                print("No GPU detected")
+                return "cpu"
+
+
+setup(
+    name="freesolo-chalk",
+    package_dir={"": "src"},
+    packages=find_packages(where="src"),
+    install_requires=get_default_dependencies(),
+    extras_require=get_optional_dependencies(),
+    classifiers=[
+        "Development Status :: 3 - Alpha",
+        "Intended Audience :: Developers",
+        "Intended Audience :: Science/Research",
+        "Programming Language :: Python :: 3",
+        "Topic :: Scientific/Engineering :: Artificial Intelligence",
+        "Topic :: Software Development :: Libraries :: Python Modules",
+        "License :: OSI Approved :: BSD License",
+        "Operating System :: OS Independent",
+    ],
+)

freesolo_chalk-0.1.0/src/chalk/ops/__init__.py

@@ -0,0 +1,12 @@
+"""
+Chalk operators — raw Triton/CUDA kernels and their ``torch.autograd.Function`` wrappers.
+
+Mirrors Liger Kernel's layout: ``chalk.ops`` holds the low-level kernel implementations
+(``@triton.jit`` functions, autograd Functions, FP8 GEMM helpers), while ``chalk.transformers``
+holds the model-level installers that monkeypatch these kernels into HuggingFace modules.
+
+This namespace starts empty by design — kernels are landed one at a time, each with its own
+benchmark evidence.
+"""
+
+__all__: list[str] = []