Skip to content

Latest commit

 

History

History
117 lines (82 loc) · 3.26 KB

File metadata and controls

117 lines (82 loc) · 3.26 KB

Contributing to mlx-kvarn

Thank you for your interest in KVarN for Apple Silicon! This project exists to make efficient KV-cache quantization accessible to everyone running models on Macs.

Getting Started

Prerequisites

  • Apple Silicon Mac (M1/M2/M3/M4)
  • Python ≥ 3.10
  • uv (recommended) or pip

Setup

# Clone and set up the environment
git clone https://github.com/YOUR_USERNAME/mlx-kvarn.git
cd mlx-kvarn

# Create virtual environment
uv venv --python 3.12
source .venv/bin/activate

# Install in development mode
uv pip install -e .

# Install test dependencies
uv pip install pytest

Running Tests

# Quick unit tests (no model download needed)
python -m pytest tests/ -v

# Long-context correctness test (requires cached model)
HF_HUB_OFFLINE=1 python tools/test_long_context.py

# Benchmark comparison (requires cached models)
python benchmarks/compare.py

How to Contribute

Reporting Issues

  • Bugs: Include the model name, macOS version, and a minimal reproducer
  • Performance: Include the model, context length, and python benchmarks/compare.py output
  • Feature requests: Explain your use case — we care about real-world Apple Silicon workflows

Pull Requests

  1. Fork and create a branch
  2. Make your change
  3. Run the test suite: python -m pytest tests/ -v
  4. If changing performance, run: python benchmarks/compare.py
  5. Submit the PR

What We're Looking For

  • More model support: Testing KVarN with different architectures
  • Performance improvements: Any optimization that doesn't break correctness
  • Documentation: Better examples, clearer explanations
  • Bug fixes: Especially for edge cases we haven't tested
  • Benchmarks: Real-world usage patterns on different Mac hardware

What We Won't Merge

  • Changes that break correctness (greedy output must match FP16 for short contexts)
  • Added dependencies without strong justification
  • Changes that slow down decode without clear benefit

Development Guidelines

Code Style

  • Match the existing code's style
  • Docstrings on public functions and classes
  • Comments explaining why, not what

Testing

  • New features should include tests
  • Performance changes should include benchmark results
  • Accuracy claims must be backed by data (greedy comparison to FP16)

Commit Messages

Use conventional commit format:

feat: add support for 8-bit values
fix: correct scale indexing in batch dequant
docs: update README with M4 benchmark results

Architecture Overview

The codebase is intentionally small and focused:

mlx_kvarn/
├── cache.py              # KVarNCache — the main class
├── kvarn_batch_dequant.py # Metal kernel (the performance-critical path)
├── models.py             # SDPA patching (the integration layer)
├── config.py             # Configuration and presets
├── quant.py              # Quantization primitives
├── sinkhorn.py           # Variance normalization
├── hadamard.py           # Hadamard rotation
└── _reference.py         # Per-tile reference implementations (tests only)

The design philosophy: minimal surface area, maximum correctness.

Questions?

Open an issue. We're happy to help you get KVarN working on your Mac.