ligongh/squat

Description

Implementation of the KV cache quantization method introduced in the SQuat paper (COLM 2025). SQuat (Subspace-orthogonal KV cache quantization) reduces the memory and compute cost of storing the KV cache by carefully quantizing the key tensors. It constructs a task-relevant subspace and ensures that quantization errors remain orthogonal to it, thereby minimizing their impact on attention outputs. SQuat is training-free, calibration-free, and operates on-the-fly, with strong theoretical grounding and state-of-the-art empirical results.

This repo provides a partial implementation of SQuat via a custom SQuatCache class. It requires passing an additional query_states input to .update(). To support this, you can monkey patch the LlamaAttention.forward method—see the example usage below.

For the full implementation, please refer to the original repository.

Base model:

meta-llama/Llama-3.1-8B-Instruct

Model compatibility

Most models. More specifically, any transformer LLM/VLM trained for causal language modeling.

Additional Arguments

• backend (str, optional): quantization backend, default is quanto
• nbits (int, optional): number of bits for quantization, default is 2
• quant_group_size (int, optional): quantization group size, default is 64
• residual_length (int, optional): residual length, default is 32
• squat_lambda (float, optional): squat lambda, default is 0.001
• subspace_dim (int, optional): subspace dimension, default is 10
• shared_svd (bool, optional): if use shared svd, default is True

Output Type changes

(none)

Example usage

import torch
from typing import Callable, Optional, Tuple
from transformers.cache_utils import Cache
from transformers.models.llama.modeling_llama import apply_rotary_pos_emb, eager_attention_forward
from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
from transformers.modeling_flash_attention_utils import FlashAttentionKwargs
from transformers.processing_utils import Unpack
import transformers

from transformers import AutoTokenizer, AutoModelForCausalLM

def llama_attn_forward(
    self,
    hidden_states: torch.Tensor,
    position_embeddings: Tuple[torch.Tensor, torch.Tensor],
    attention_mask: Optional[torch.Tensor],
    past_key_value: Optional[Cache] = None,
    cache_position: Optional[torch.LongTensor] = None,
    **kwargs: Unpack[FlashAttentionKwargs],
) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:

    input_shape = hidden_states.shape[:-1]
    hidden_shape = (*input_shape, -1, self.head_dim)

    query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
    key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
    value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)

    cos, sin = position_embeddings
    query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)

    if past_key_value is not None:
        # sin and cos are specific to RoPE models; cache_position needed for the static cache
        cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position, "query_states": query_states, "attention_mask": attention_mask}
        key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)

    attention_interface: Callable = eager_attention_forward

    if self.config._attn_implementation != "eager":
        if self.config._attn_implementation == "sdpa" and kwargs.get("output_attentions", False):
            logger.warning_once(
                "`torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True`. Falling back to "
                'eager attention. This warning can be removed using the argument `attn_implementation="eager"` when loading the model.'
            )
        else:
            attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]

    attn_output, attn_weights = attention_interface(
        self,
        query_states,
        key_states,
        value_states,
        attention_mask,
        dropout=0.0 if not self.training else self.attention_dropout,
        scaling=self.scaling,
        **kwargs,
    )

    attn_output = attn_output.reshape(*input_shape, -1).contiguous()
    attn_output = self.o_proj(attn_output)
    return attn_output, attn_weights

def replace_llama():
    transformers.models.llama.modeling_llama.LlamaAttention.forward = llama_attn_forward

replace_llama()

tokenizer = AutoTokenizer.from_pretrained('meta-llama/Llama-3.1-8B-Instruct')
model = AutoModelForCausalLM.from_pretrained('meta-llama/Llama-3.1-8B-Instruct', device_map="auto")

inputs = tokenizer(["I like rock music because"], return_tensors="pt").to(model.device)

gen_out = model.generate(**inputs, custom_generate="ligongh/squat", trust_remote_code=True)
print(tokenizer.batch_decode(gen_out))