Metadata-Version: 2.4
Name: mlx-iqk
Version: 0.1.3
Summary: IQ_K codecs, relayouts, and Metal kernels for MLX on Apple silicon.
License-Expression: MIT
License-File: LICENSE
License-File: NOTICE
Requires-Python: >=3.11
Requires-Dist: mlx>=0.31.2
Requires-Dist: numpy>=1.26
Description-Content-Type: text/markdown

# mlx-iqk

`mlx-iqk` provides IQ_K codecs, byte-preserving relayouts, and Metal kernels
for MLX on Apple silicon. It supports stacked routed-expert projections and
ordinary dense projections. Codec selection, tensor allocation, and model
policy remain the responsibility of the caller.

## Supported formats

| use | members |
|---|---|
| routed experts | `IQ1_S_R4`, `IQ2_KS`, `IQ2_K`, `IQ3_K` |
| dense tensors | `IQ4_KS`, `IQ4_K`, `IQ5_KS`, `IQ5_K`, `IQ6_K` |

For every member, the library provides:

- a CPU encoder and dequantizer based on pinned ik_llama.cpp code;
- conversion between ik wire bytes and a reduction-contiguous layout;
- a NumPy reference decoder;
- Metal decode and dequantization kernels;
- geometry and representation checks before kernel compilation.

The relayout changes where fields are stored, not their encoded values or
total byte count. `mlx_iqk.format.bits_per_weight()` reports the exact storage
cost for a member and row width, including row metadata.

## Installation

```bash
uv add mlx-iqk
```

The runtime requires Apple silicon and MLX. Python 3.11 or newer is supported.
The CPU codec builds its vendored C++ sources on first use if the corresponding
shared library is absent.

## Conversion

`mlx_iqk.codec` converts float32 rows to ik wire bytes. `mlx_iqk.format` packs
those bytes into the layout consumed by the Metal kernels.

```python
from mlx_iqk import codec
from mlx_iqk import format as fmt

wire = codec.quantize("iq2_ks", weights, imatrix)
streams = fmt.pack("iq2_ks", wire, weights.shape[-1])

# Both checks are byte-exact.
assert fmt.relayout_row_bytes("iq2_ks", weights.shape[-1]) == wire.shape[-1]
assert (fmt.unpack("iq2_ks", streams, weights.shape[-1]) == wire).all()
```

`imatrix` is one float32 importance vector with the same length as a weight
row. Passing `None` selects the codec's unweighted objective; it is not
equivalent to calibrated conversion.

## Routed experts

`IqkSwitchLinear` stores one stacked expert projection. Small token-expert
batches use a packed GEMV. Larger or explicitly sorted batches dequantize the
expert stack to fp16 and use `mx.gather_mm`.

```python
import mlx.core as mx

from mlx_iqk import IqkSwitchLinear

projection = IqkSwitchLinear(
    "iq2_ks",
    num_experts=num_experts,
    out_features=out_features,
    in_features=in_features,
)
projection.load_streams({name: mx.array(value) for name, value in streams.items()})
output = projection(hidden, expert_indices)
```

The routed kernels currently support these reduction widths:

| members | `in_features` |
|---|---|
| `IQ2_KS`, `IQ2_K` | 768, 2048, 2560, 4096 |
| `IQ3_K` | 768, 2560 |
| `IQ1_S_R4` | 640, 2048, 2560, 4096 |

The generic packed GEMV requires `out_features` to be a positive multiple of
16. Dequantization accepts any positive output size.

`mlx_iqk.routed` also contains a specialized two-dispatch MoE executor for its
declared geometry. It fuses gate, up, and SwiGLU in the first dispatch, then
fuses down projection, route weighting, and ordered reduction in the second.
Unsupported geometries and codec tuples are rejected before compilation.

## Dense tensors

Dense entry points accept either named relayout streams or one row-packed
`uint8` matrix:

```python
from mlx_iqk import dense_gemv, dense_linear

decoded = dense_gemv(member, hidden, streams, out_features, in_features)
automatic = dense_linear(member, hidden, streams, out_features, in_features)
```

The `dense_*_packed` variants operate directly on a row-packed MLX allocation.
Full and row-range dequantization are available for prefill paths. The dense
kernels support reduction widths 512, 768, 1024, 2048, 2560, 4096, 6144,
8192, and 10240, with arbitrary positive output sizes.

`dense_linear` uses a token-count crossover between packed GEMV and
dequantize-plus-matmul. Callers with measured scheduling policy can invoke
`dense_gemv`, `dense_dequantized`, or their packed and range variants
directly.

## Architecture adapters

Model-specific geometry is isolated from the format and kernel core.
`mlx_iqk.qwen` provides the logical-to-stored width mapping and zero-padding
helpers required by the supported Qwen hybrid architecture. The adapter does
not select codecs or tensor allocations.

## Correctness

The CPU dequantizers define the arithmetic reference. Relayout tests compare
their float32 output bit for bit, and Metal dequantization tests compare fp16
output bit for bit. GEMV tests use independent decoded-weight products and
format-specific tolerances where accumulation order differs.

```bash
uv sync --locked --group dev
uv run --locked --group dev pytest
```

The optional upstream cross-check requires an ik_llama.cpp checkout at the
revision recorded in `vendor/ik_llama/UPSTREAM_REVISION`:

```bash
IK_LLAMA_DIR=/path/to/ik_llama.cpp \
  uv run --locked --group dev pytest tests/test_vendor_upstream.py
```

It rebuilds its comparator in a temporary directory and verifies quantization,
dequantization, row sizes, and relayouts against upstream `libggml`.

## Benchmarks

`bench/RESULTS.md` records kernel measurements and the corresponding address
floors. Benchmark drivers validate their outputs before timing and emit JSON
for reproducibility.

## Licensing

The project is MIT licensed. Vendored codec code retains the upstream MIT
terms. `NOTICE` and `ATTRIBUTION.md` identify copied and adapted work, source
revisions, symbols, and local changes.
