You can deploy this model and start earning money today!

jak-pan/
kalm-reranker-v1-nano-mlx

Name: kalm-reranker-v1-nano-mlx
Availability: InStock

KaLM-Embedding/KaLM-Reranker-V1-Nano

finetune

apache-2.0

Views

No views yet

KaLM-Reranker-V1-Nano — MLX (Apple Silicon)

A hand-written, torch-free MLX inference implementation of KaLM-Embedding/KaLM-Reranker-V1-Nano, KaLM's "fast but not late interaction" (FBNL) document reranker. It runs natively on the Apple Silicon GPU via Metal.

Weights are NOT redistributed here. This repo ships only the MLX forward pass (modeling_kalm_mlx.py) and this card. The actual model weights are loaded at runtime from the base_model above, out of your local HuggingFace cache. Download them separately (see Usage). The base weights and their license belong to the KaLM team and Google (see License).

What this is

KaLM-Reranker-V1 is not a standard cross-encoder reranker. It is a T5-Gemma2 encoder-decoder with a non-standard relevance pipeline:

The encoder pre-encodes each passage once and chunk/Matryoshka-pools it (chunk size 4 tokens, masked-mean per chunk) into a small set of cached document vectors.
The decoder models the system instruction + user instruction + query, and reads relevance via a merged self+cross attention op — a single fused scaled_dot_product_attention that concatenates the decoder's self-attention KV with cross-attention KV computed over the cached passage encodings.
The relevance score is softmax([logit_yes, logit_no])[yes], read from the decoder's final position.

This decoupling (encode passages once, cross-attend at query time) is what makes the architecture "fast but not late interaction." Standard MLX reranker recipes don't cover it — hence this hand-written port.

The forward pass is config-generalized: every dim, head count, layer count, and RoPE parameter is read from the model's own config.json (encoder.text_config + decoder + rope_parameters). The same modeling_kalm_mlx.py file therefore serves both Nano (0.27B activated, H=640, 18 layers) and Small (1B activated, H=1152, 26 layers) — just pass the model id.

Derivation

Weights: KaLM-Reranker-V1-Nano's original bf16 safetensors, used as-is, loaded directly via mx.load() (no torch, no conversion). The tokenizer is the original, loaded via HuggingFace AutoTokenizer (the tokenizers library only — no torch model).
Forward pass: the MLX implementation in modeling_kalm_mlx.py is hand-written from the architecture — encoder, chunk/Matryoshka pool, and the merged self+cross attention decoder. RoPE inverse-frequency tables are computed from config and validated identical to the on-disk buffers.

Nothing about the model's behavior is changed; this is the same model, re-expressed in MLX. The only differences from PyTorch are float accumulation order (see Validation).

Files

File	Purpose
`modeling_kalm_mlx.py`	Torch-free MLX inference module. Exposes `load(model_id)` and `rerank(query, docs, top_n)`. Config-generalized — serves both Nano and Small.
`requirements.txt`	`mlx`, `safetensors`, `transformers`, `numpy`.
`.gitignore`	Ignores `__pycache__/` and `*.safetensors` (weights are never committed here).
`README.md`	This card.

Speed

Measured on the machine below, on a real 100-document reranking payload:

~0.75 s end-to-end / 100 docs (length-sorted sub-batch encode, B=25, then one big decode); ~0.36 s decode-only with a warm encode cache.
~2.6× faster than the local Nemotron-1B MLX reranker (~1.43 s / 100 docs).

Because the encoder pools each passage independently and once, those passage encodings are precompute-cacheable. For persistent-memory / repeated-query workloads (the same document set reranked against many queries), you pay the encode cost once and only the ~0.36 s decode per query — this is the regime where this model is the strongest fit.

Notes from sweeping the MLX stack: length-sorting before batching was the only clean speed lever. Quantization (int8/int4) is a loss on Metal here (the forward is compute-bound, not memory-bound); mx.compile, concurrency, and precompute all landed within measurement noise.

Validation

The MLX port is validated by single-query numerical match against the original PyTorch implementation: max relevance-score difference < 0.0008, with identical top-3 / top-5 ranking. This confirms the port is faithful to the original weights — it is the same model.

What this does NOT prove: that the port is faithful does not establish a task-accuracy number on any end-to-end benchmark. See the honesty note in Quality.

Quality — please read honestly

This repo makes no task-accuracy claim for KaLM-Reranker-V1-Nano. It ships a numerically-faithful MLX port of the forward pass (see Validation) — that proves the port matches the original weights, not that it hits any particular accuracy on your task. If you need a quality number for your use case, measure it yourself.

The strong public reranking results you may have seen — BEIR, MIRACL, LMEB — are the KaLM team's own published numbers for the parent model, not our measurement and not specific to this MLX port. From the parent model card: KaLM-Reranker-V1 reports state-of-the-art / competitive performance on BEIR, strong multilingual results on MIRACL, and on LMEB the 0.27B Nano "remains competitive with 7–12B embedding models." See the parent card and the KaLM-Reranker-V1 paper for the authoritative figures. Cite those as the KaLM team's claim about the parent — not as a benchmark of this port.

Machine

All speed and validation figures were produced on:


Model	MacBook Pro (MacBookPro18,4)
Chip	Apple M1 Max (10 cores: 8 performance + 2 efficiency)
GPU	Apple Silicon unified-memory GPU, driven via Metal (MLX runs natively on it)
Memory	64 GB unified
macOS	26.5.1 (build 25F80)
Python MLX (`mlx.core`)	0.31.2

Usage

Run it fast (Rust): this model also runs in the mlx-rerank server — torch-free, native Metal, /rerank API: RERANK_MODEL=KaLM-Embedding/KaLM-Reranker-V1-Nano rerank --serve. The Python below is the readable reference implementation.

bash

pip install -r requirements.txt
# Download the base weights once (NOT included in this repo):
huggingface-cli download KaLM-Embedding/KaLM-Reranker-V1-Nano

python

import modeling_kalm_mlx as kalm

# Loads weights + tokenizer from your local HF cache (defaults to the Nano model id;
# pass "KaLM-Embedding/KaLM-Reranker-V1-Small" — or set KALM_MODEL — to run Small instead).
kalm.load("KaLM-Embedding/KaLM-Reranker-V1-Nano")

results = kalm.rerank(
    "How long is my commute?",
    [
        "My commute is 45 minutes each way.",
        "I like coffee.",
        "The weather was nice today.",
    ],
    top_n=3,
)
# -> [{"index": 0, "relevance_score": 0.98...},
#     {"index": ..., "relevance_score": ...}, ...]  highest first, print(results)

Or run the built-in smoke test directly:

python modeling_kalm_mlx.py KaLM-Embedding/KaLM-Reranker-V1-Nano

Environment knobs: KALM_MODEL (model id), KALM_SUBBATCH (encode sub-batch size, default 25), KALM_MEM_GB / KALM_CACHE_GB (MLX memory / cache caps).

License

This MLX inference code (modeling_kalm_mlx.py): released under Apache-2.0, matching the parent model.
The model weights (loaded at runtime, not distributed here) are KaLM-Embedding/KaLM-Reranker-V1-Nano, released by the KaLM team under Apache-2.0.
The base weights derive from Google's Gemma (google/t5gemma-2-270m-270m), so Google's Gemma Terms of Use also apply to the underlying base model, in addition to KaLM's Apache-2.0 release. You are responsible for complying with both when you download and use the weights.

Acknowledgements

KaLM-Embedding for KaLM-Reranker-V1 and the released weights.
Google for the Gemma / T5-Gemma2 base model.
The MLX team at Apple.

jak-pan/kalm-reranker-v1-nano-mlx