LFM2.5 2.6B: Liquid Foundation Models

A minimal PyTorch implementation of LiquidAI/LFM2.5-2.6B. LFM2.5 uses the LFM2 hybrid architecture and updates the vocabulary, context length, RoPE base, tokenizer, and post-training behavior.

📖 LMF2 Docs: lfm2 technical docs

Setup

Python 3.11 and all runtime dependencies are managed with uv.

uv sync
uv run python main.py

The first run downloads the two Hugging Face weight shards. main.py uses the CPU by default; change device to torch.device("cuda") when CUDA is available.

Web Chat UI

The project includes a local FastAPI chat service and a responsive single-page interface in static/index.html. Start it with:

uv run --with-requirements requirements-web.txt \
  uvicorn web_app:app --host 127.0.0.1 --port 8000

Then open http://127.0.0.1:8000. The server loads the same local LFM2.5-2.6B checkpoint used by main.py, automatically selecting CUDA, Apple MPS, or CPU. Initial startup can take some time while the checkpoint is loaded into memory.

The Web UI provides:

• Character-by-character streamed responses with lightweight Markdown formatting.
• Generated-token throughput and total response time shown after each reply.
• Multi-turn conversations that include saved messages in subsequent prompts.
• A history sidebar for restoring previous chats by UUID.
• Collapsed Thinking panels for text generated inside <think>...</think>; the main reply displays only the answer produced after </think>.
• Local JSON persistence in logs/<uuid>.json, with separate thinking and content fields for assistant messages.

Set LFM_MAX_NEW_TOKENS to change the maximum response length. For example:

LFM_MAX_NEW_TOKENS=2048 uv run --with-requirements requirements-web.txt \
  uvicorn web_app:app --host 127.0.0.1 --port 8000

LFM2.5-2.6B Config

The values below come from the checkpoint's config.json.

import torch

LFM25_CONFIG = {
    "vocab_size": 128_000,
    "context_length": 131_072,
    "emb_dim": 2_048,
    "n_heads": 32,
    "n_kv_heads": 8,
    "n_layers": 30,
    "hidden_dim": 10_752,
    "head_dim": 64,
    "conv_kernel_size": 3,
    "norm_eps": 1e-5,
    "rope_base": 10_000_000.0,
    "bos_token_id": 124_894,
    "eos_token_id": 124_900,
    "pad_token_id": 124_893,
    "dtype": torch.bfloat16,
}

The model contains 30 hybrid decoder blocks:

• 22 causal, double-gated, depthwise convolution blocks with kernel size 3.
• 8 grouped-query attention blocks with 32 query heads and 8 key/value heads.
• Every block contains pre-normalization, a residual mixer, and a SwiGLU feed-forward network.
• Attention uses query/key RMSNorm, 64-dimensional heads, causal masking, and RoPE with a base of 10,000,000.
• Token embeddings and the language-model output head share weights.
• Autoregressive generation caches attention keys/values and convolution state.

The full layer schedule is defined in config.py. Architecture references are available in the LFM2 technical report.

Code Pointers

File	Responsibility
config.py	LFM2.5-2.6B dimensions, layer schedule, token IDs, RoPE, and dtype
lfm2.py	Decoder blocks, model assembly, causal masks, and hybrid cache handling
modules.py	RMSNorm, RoPE, SwiGLU, and causal double-gated convolution
attention.py	Grouped-query causal attention and attention KV cache
triton_kernels.py	Optional CUDA inference kernels with PyTorch fallbacks
utils.py	Hugging Face checkpoint-to-model weight mapping
tokenizer.py	Special tokens and the LFM2.5 reasoning chat template
sampling.py	Temperature, top-k, top-p, repetition penalty, and streaming decode
main.py	End-to-end checkpoint download, loading, tokenization, and inference

Load The Model

Instantiate the local architecture using LFM2Config. All model parameters use the torch.bfloat16 dtype declared by config.dtype.

import torch

from config import LFM2Config
from lfm2 import LFM2ForCausalLM

config = LFM2Config()
model = LFM2ForCausalLM(config)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

Load Hugging Face Weights

The checkpoint is sharded, so download the snapshot and merge its Safetensors files before mapping them into the local model.

from pathlib import Path

from huggingface_hub import snapshot_download
from safetensors.torch import load_file

from utils import load_weights_into_lfm2

repo_id = "LiquidAI/LFM2.5-2.6B"
local_dir = Path(repo_id).name

model_path = Path(
    snapshot_download(
        repo_id=repo_id,
        local_dir=local_dir,
        allow_patterns=["*.safetensors", "*.json"],
    )
)

weights_dict = {}
for weights_file in sorted(model_path.glob("*.safetensors")):
    weights_dict.update(load_file(str(weights_file)))

load_weights_into_lfm2(model, weights_dict)
model.to(device=device)
model.eval()
del weights_dict

Loading all shards into one dictionary temporarily requires memory for both the checkpoint tensors and the instantiated model.

Load The Tokenizer

snapshot_download also downloads tokenizer.json and tokenizer_config.json. The wrapper validates that its BOS, EOS, and padding IDs match the model config.

from tokenizer import LFM2Tokenizer

tokenizer = LFM2Tokenizer(
    tokenizer_file_path=str(model_path / "tokenizer.json"),
    tokenizer_config_path=str(model_path / "tokenizer_config.json"),
    model_config=config,
)

LFM2.5-2.6B is a reasoning model. Its chat template starts every assistant turn with <think>:

<|startoftext|><|im_start|>user
What is climate change?<|im_end|>
<|im_start|>assistant
<think>

See the official chat_template.jinja for multi-turn, tool-calling, and structured-content formatting.

Inference And Sampling

The recommended checkpoint defaults are temperature 0.1, top-k 50, and repetition penalty 1.1.

from sampling import advance_decoding

prompt = "Please explain climate change and how it impacts our future."
formatted_prompt = tokenizer.apply_chat_template(prompt)
input_token_ids = tokenizer.encode(formatted_prompt)
input_ids = torch.tensor(input_token_ids, device=device).unsqueeze(0)

print(f"Prompt: {prompt}")
print("Assistant: ", end="", flush=True)

for token in advance_decoding(
    model=model,
    token_ids=input_ids,
    max_new_tokens=512,
    eos_token_id=tokenizer.eos_id,
    temperature=0.1,
    top_k=50,
    top_p=None,
    repetition_penalty=1.1,
    window_size=64,
):
    token_id = token.squeeze().tolist()
    token_ids = [token_id] if isinstance(token_id, int) else token_id
    print(tokenizer.decode(token_ids), end="", flush=True)

The complete runnable version is in main.py.

Optimization Changelog

2026-09-01: Non-Triton Generation Optimizations

This pass optimized the existing PyTorch generation path without adding Triton kernels or changing the model architecture, checkpoint weights, sampling settings, or public generation API.

Changes:

1. Keep RoPE tables on the model device. The precomputed FP32 cosine and sine tables are now non-persistent model buffers. They move to the selected device once instead of copying approximately 64 MiB from CPU on every model forward.
2. Calculate only the required prompt logit. During generation, final normalization and the 128,000-token output projection now run only for the last prompt position. For the 128-token benchmark prompt, this reduces the temporary BF16 logits output from 31.25 MiB to about 0.24 MiB.
3. Keep repetition penalty processing on the GPU. The previous recent_tokens.tolist() and per-token Python loop were replaced by batched gather, penalty, and scatter tensor operations. This removes repeated CPU-GPU synchronizations from every generated token.
4. Use a fixed recent-token ring buffer. Generation no longer grows the complete token tensor with torch.cat after every step. Repetition history now uses fixed memory controlled by window_size (64 tokens by default).

Measured Result

The benchmark used an NVIDIA GeForce RTX 4060 8 GB, PyTorch 2.13.0 with CUDA 13.0, BF16 model weights, a 128-token prompt, 256 generated tokens, 8 warm-up tokens, temperature 0.1, top-k 50, repetition penalty 1.1, and a 64-token repetition window. Timings include an explicit CUDA synchronization.

Measurement	Before	After	Improvement
Generation time	9.593 s	6.815 s average	29.0% lower
Generation throughput	26.69 tok/s	37.56 tok/s average	40.8% higher
Web-like synchronized streaming	26-29 tok/s historical	36.67 tok/s	26-41% higher
Peak allocated CUDA memory	5.139 GiB	5.117 GiB	22.6 MiB lower

The optimized throughput was stable across two measured runs: 37.56 and 37.57 tokens per second. Peak memory is reported using torch.cuda.max_memory_allocated() and does not include CUDA context or reserved allocator memory. Longer prompts should benefit more from computing only the last prompt logit.

2026-09-01: SDPA-GQA Prefill And Fixed KV Cache

This pass optimized attention without adding a custom Triton kernel. Prefill is the first model pass over the complete prompt; it produces the first output logit and builds the key/value history used by subsequent one-token decoding.

Changes:

1. Use fused scaled dot-product attention for CUDA prefill. The explicit query-key matrix multiplication, causal masking, softmax, and value matrix multiplication were replaced with scaled_dot_product_attention(..., enable_gqa=True). On the benchmark GPU, PyTorch selects its Flash Attention GQA backend.
2. Do not expand grouped keys and values on CUDA. The 32 query heads now address the 8 KV heads through native GQA instead of physically creating four copies with repeat_interleave.
3. Use fixed-capacity attention caches during generation. Each attention layer allocates storage once for prompt length + max new tokens. New K/V entries are copied into the next cache position instead of concatenating and copying the entire history after every generated token.
4. Avoid unnecessary masks and RoPE copies. Square prefill uses SDPA's causal mode, single-token decode requires no causal mask, and full-head RoPE no longer concatenates an empty partial-RoPE suffix.
5. Preserve non-CUDA behavior. CPU and MPS retain an expanded-KV SDPA fallback, and legacy tuple caches remain accepted for compatibility.

Prefill Result

The benchmark used the same RTX 4060, BF16 model, and synchronized timing as the previous entry. Each result is the average of two runs.

Prompt length	Before	After	Time reduction
128 tokens	37.5 ms	37.9 ms	Within measurement noise
512 tokens	107.9 ms	104.4 ms	3.2%
1,024 tokens	230.7 ms	208.1 ms	9.8%
2,048 tokens	532.9 ms	440.2 ms	17.4%

At 2,048 prompt tokens, peak allocated CUDA memory decreased from 5.848 GiB to 5.452 GiB, a reduction of approximately 406 MiB. The improvement grows with prompt length because the removed attention-score tensors grow quadratically.

Decode And KV-Cache Result

Starting context	Before	After	Throughput gain
128 tokens	37.90 tok/s	37.77 tok/s	Within measurement noise
2,048 tokens	36.30 tok/s	37.63 tok/s	3.7%
8,192 tokens	32.37 tok/s	40.18 tok/s	24.1%

The 8,192-token comparison isolates 256 decode steps using preinitialized caches so both implementations process the same cache length. Peak CUDA allocation decreased from 5.548 GiB to 5.225 GiB, approximately 331 MiB lower. The full 8,192-token prefill did not fit on the shared 8 GB benchmark GPU due to the remaining feed-forward activation memory; that requires a separate FFN chunking optimization and is not addressed by this attention change.

2026-09-01: Optional Triton Inference Kernels

This pass added src/triton_kernels.py while keeping the existing modules, parameter names, checkpoint loader, and generation API. The requested convolution decode kernel is deliberately not included.

Changes:

1. Triton RMSNorm. Contiguous CUDA BF16 inference tensors now use one kernel with FP32 variance and normalization accumulation. Training, unsupported layouts, CPU, and MPS use the original PyTorch calculation.
2. Triton GQA decode with online softmax. One-token attention directly maps each query head to its KV group without score-tensor materialization. On the RTX 4060 it is used through 256 cached tokens, where it is faster. Longer contexts automatically retain Flash SDPA, which benchmarked faster than the custom kernel.
3. Triton sampling transforms. Temperature scaling uses one linear logits pass. Repetition penalty uses a small sparse kernel that deduplicates the recent-token window and updates each selected vocabulary position once.
4. Portable fallbacks. Importing Triton is optional. Non-CUDA execution, unsupported dtypes and shapes, and gradient-enabled calls retain PyTorch implementations.
5. Explicit dependency. Linux installations now declare Triton 3.3 or newer directly; non-Linux installations keep the CPU/MPS fallback without pulling in the Linux-only package.

Kernel Microbenchmarks

Operation	PyTorch/SDPA	Triton	Speedup
RMSNorm, 1 x 2,048	0.0225 ms	0.0051 ms	4.39x
RMSNorm, 32 x 64	0.0215 ms	0.0051 ms	4.20x
RMSNorm, 1,024 x 2,048	0.2230 ms	0.0379 ms	5.89x
Temperature + repetition penalty	0.0217 ms	0.0152 ms	1.43x
GQA decode, 128-token cache	0.0137 ms	0.0094 ms	1.45x
GQA decode, 256-token cache	0.0222 ms	0.0164 ms	1.36x

At 512 cached tokens and above, Flash SDPA is faster, so the GQA wrapper uses SDPA instead of forcing the Triton path.

Full-Model Result

The controlled 128-token prompt and 256-token generation benchmark improved from 37.69 to 40.22 tokens per second, a 6.7% throughput increase. Average generation time decreased from 6.792 to 6.366 seconds, or 6.3%. Peak allocated CUDA memory was 5.121 GiB and remained dominated by model weights and caches.

Relative to the original 26.69-token-per-second baseline before any optimizations, the current result is approximately 50.7% faster.

Triton RMSNorm also changed synchronized prefill latency as follows:

Prompt length	Before Triton	After Triton	Time reduction
128 tokens	37.9 ms	34.6 ms	8.7%
1,024 tokens	208.1 ms	206.0 ms	1.0%
2,048 tokens	440.2 ms	415.9 ms	5.5%