Upload README.md

README.md

@@ -1,3 +1,111 @@
----
-license: mit
----
+---
+library_name: transformers
+tags:
+  - custom_generate
+---
+
+
+## Description
+Implementation of the KV cache quantization method introduced in the [SQuat paper (COLM 2025)](https://arxiv.org/abs/2503.24358). 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](https://github.com/Red-Hat-AI-Innovation-Team/SQuat).
+
+
+## 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
+
+```py
+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))
+```