Meta-Llama-3.1-8B-Instruct-Butler

Llama-3.1-8B-Instruct with TokenButler

This repository releases Llama-3.1-8B-Instruct augmented with a TokenButler sparse-attention predictor, packaged as a Hugging Face trust_remote_code model and loadable directly with the standard Transformers API.

• Paper: TokenButler: Token Importance is Predictable
• Code: https://github.com/abdelfattah-lab/tokenbutler

TokenButler is a lightweight, query-aware token-importance predictor that works alongside a frozen large language model. At inference time it identifies the small subset of tokens in the KV cache that contribute meaningfully to each decoding step, enabling sparse attention with a fixed budget while preserving the full KV cache. The base model is left untouched; only the predictor parameters (a small fraction of the total) are added.

Method

Let G denote the producer frequency. Producer layers are the layers at indices 0, G, 2G, …. Each producer reads the layer's hidden states H ∈ ℝ^B×L×E and predicts low-dimensional importance queries for itself and the next G − 1 consumer layers:

Q_imp = f_θ(LN(H)) ∈ ℝ^{(B·H)×G×L×d′}

For each layer ℓ the real (post-RoPE) cached keys are projected into the same d′-dimensional space via a learned matrix W_K^(ℓ) ∈ ℝ^D×d′:

K_imp^(ℓ) = K^(ℓ) · W_K^(ℓ)

A consumer at layer ℓ uses slot (ℓ − 1) mod G of its producer's importance queries together with K_imp^(ℓ) to score every cached token and select a top-x subset under a fixed token budget, augmented by a sink prefix and a recent local window. Only the predictor parameters are trained, by distilling the masked causal attention distribution of the frozen base model.

For the full method, training procedure, and evaluation details see the TokenButler paper.

Configuration

The predictor adds approximately 0.15 GB on top of the 16 GB base model weights (about 1% parameter overhead).

Usage

Standard AutoTokenizer / AutoModelForCausalLM calls with trust_remote_code=True. Loading in bfloat16 is recommended; the default float32 precision needs roughly twice the GPU memory.

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline

model_name = "Alwahsh/Meta-Llama-3.1-8B-Instruct-Butler"
tokenizer  = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
model      = AutoModelForCausalLM.from_pretrained(
    model_name,
    trust_remote_code=True,
    torch_dtype=torch.bfloat16,
).to("cuda")

generator = pipeline("text-generation", model=model, tokenizer=tokenizer)
out = generator(
    "Long-context inference is challenging because ",
    max_new_tokens=200, do_sample=True, top_p=0.9, temperature=0.6,
)
print(out[0]["generated_text"])

Memory requirements

On Google Colab the free T4 instance has 16 GB of VRAM, which is too tight for an 8B model plus a KV cache. Use a Colab Pro / Pro+ runtime (A100, L4, or V100) and pass torch_dtype=torch.bfloat16 as shown above.

The default per-layer sparsity is 50% (every layer except layer 0 retains the top 50% of tokens, plus the sink prefix and the sliding window). To change the sparsity at runtime:

def set_sparsity(model, sparsity: str):
    for module in model.modules():
        if module.__class__.__name__.endswith("AttentionExperimental"):
            module.token_sparse_method = sparsity
            module.set_token_sparsity()
    return model

model = set_sparsity(model, "fixed_70pc")  # retain 30% of tokens per layer

The fixed_<N>pc policy is the supported strategy. Sink length and sliding window can be adjusted in the same way by setting module.min_sparse_index and module.sliding_window on every *AttentionExperimental module.

Files

Requirements

• transformers >= 4.45
• torch with bfloat16 support recommended

Citation

If you use TokenButler in your work, please cite the paper:

@misc{akhauri2025tokenbutlertokenimportancepredictable,
      title={TokenButler: Token Importance is Predictable}, 
      author={Yash Akhauri and Ahmed F AbouElhamayed and Yifei Gao and Chi-Chih Chang and Nilesh Jain and Mohamed S. Abdelfattah},
      year={2025},
      eprint={2503.07518},
      archivePrefix={arXiv},
      primaryClass={cs.CL},
      url={https://arxiv.org/abs/2503.07518}, 
}

License

This release is governed by the Llama 3.1 Community License Agreement, inherited from the base model.