Add KV cache support for generation

modeling_nova.py

@@ -2,7 +2,7 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 from transformers import GenerationMixin, PreTrainedModel
-from transformers.modeling_outputs import CausalLMOutput
+from transformers.modeling_outputs import CausalLMOutputWithPast
 
 try:
     from .configuration_nova import NovaConfig
@@ -45,7 +45,7 @@ class NovaBlock(nn.Module):
         self.up = nn.Linear(d, ff, bias=False)
         self.down = nn.Linear(ff, d, bias=False)
 
-    def forward(self, x, cos, sin, attention_mask):
+    def forward(self, x, cos, sin, attention_mask, past_key_value=None, use_cache=False):
         batch, length, _ = x.shape
         h = self.n1(x)
         q = self.wq(h).view(batch, length, self.h, self.hd)
@@ -54,17 +54,26 @@ class NovaBlock(nn.Module):
         q = _rope(self.qn(q), cos, sin).to(v.dtype).transpose(1, 2)
         k = _rope(self.kn(k), cos, sin).to(v.dtype).transpose(1, 2)
         v = v.transpose(1, 2)
+        if past_key_value is not None:
+            k = torch.cat((past_key_value[0], k), dim=-2)
+            v = torch.cat((past_key_value[1], v), dim=-2)
+        present = (k, v) if use_cache else None
         repeat = self.h // self.kv
         k = k.repeat_interleave(repeat, dim=1)
         v = v.repeat_interleave(repeat, dim=1)
-        causal = torch.ones(length, length, dtype=torch.bool, device=x.device).tril()
+        key_length = k.size(-2)
+        past_length = key_length - length
+        query_positions = past_length + torch.arange(length, device=x.device)
+        key_positions = torch.arange(key_length, device=x.device)
+        causal = key_positions[None, :] <= query_positions[:, None]
         mask = causal[None, None, :, :]
         if attention_mask is not None:
             mask = mask & attention_mask[:, None, None, :].to(torch.bool)
         out = F.scaled_dot_product_attention(q, k, v, attn_mask=mask)
         x = x + self.wo(out.transpose(1, 2).reshape(batch, length, -1))
         h = self.n2(x)
-        return x + self.down(F.silu(self.gate(h)) * self.up(h))
+        x = x + self.down(F.silu(self.gate(h)) * self.up(h))
+        return x, present
 
 
 class NovaModel(PreTrainedModel):
@@ -99,25 +108,42 @@ class NovaModel(PreTrainedModel):
         attention_mask=None,
         position_ids=None,
         inputs_embeds=None,
+        past_key_values=None,
+        use_cache=False,
         **kwargs,
     ):
         if inputs_embeds is None:
             inputs_embeds = self.emb(input_ids)
         batch, length, _ = inputs_embeds.shape
         if position_ids is None:
-            position_ids = torch.arange(length, device=inputs_embeds.device).unsqueeze(0).expand(batch, -1)
+            past_length = past_key_values[0][0].size(-2) if past_key_values else 0
+            if attention_mask is not None:
+                position_ids = attention_mask.long().cumsum(-1) - 1
+                position_ids.masked_fill_(attention_mask == 0, 0)
+                position_ids = position_ids[:, -length:]
+            else:
+                position_ids = torch.arange(
+                    past_length, past_length + length, device=inputs_embeds.device
+                ).unsqueeze(0).expand(batch, -1)
         cos = self.cos[position_ids].unsqueeze(-2)
         sin = self.sin[position_ids].unsqueeze(-2)
         hidden = inputs_embeds
-        for block in self.blocks:
-            hidden = block(hidden, cos, sin, attention_mask)
-        return self.norm(hidden)
+        presents = () if use_cache else None
+        for i, block in enumerate(self.blocks):
+            layer_past = past_key_values[i] if past_key_values is not None else None
+            hidden, present = block(
+                hidden, cos, sin, attention_mask, layer_past, use_cache
+            )
+            if use_cache:
+                presents += (present,)
+        return self.norm(hidden), presents
 
 
 class NovaForCausalLM(PreTrainedModel, GenerationMixin):
     config_class = NovaConfig
     base_model_prefix = "model"
     _tied_weights_keys = []
+    _supports_cache_class = False
 
     def __init__(self, config):
         super().__init__(config)
@@ -134,8 +160,25 @@ class NovaForCausalLM(PreTrainedModel, GenerationMixin):
         return None
 
     def prepare_inputs_for_generation(self, input_ids, attention_mask=None, **kwargs):
-        # This model has no KV cache; each generation step recomputes the full prefix.
-        return {"input_ids": input_ids, "attention_mask": attention_mask}
+        past_key_values = kwargs.get("past_key_values")
+        if past_key_values is not None:
+            input_ids = input_ids[:, -1:]
+        if attention_mask is not None:
+            position_ids = attention_mask.long().cumsum(-1) - 1
+            position_ids.masked_fill_(attention_mask == 0, 0)
+            position_ids = position_ids[:, -input_ids.size(1):]
+        else:
+            past_length = past_key_values[0][0].size(-2) if past_key_values else 0
+            position_ids = torch.arange(
+                past_length, past_length + input_ids.size(1), device=input_ids.device
+            ).unsqueeze(0).expand(input_ids.size(0), -1)
+        return {
+            "input_ids": input_ids,
+            "attention_mask": attention_mask,
+            "position_ids": position_ids,
+            "past_key_values": past_key_values,
+            "use_cache": kwargs.get("use_cache", True),
+        }
 
     def forward(
         self,
@@ -144,13 +187,19 @@ class NovaForCausalLM(PreTrainedModel, GenerationMixin):
         position_ids=None,
         inputs_embeds=None,
         labels=None,
+        past_key_values=None,
+        use_cache=None,
         **kwargs,
     ):
-        hidden = self.model(
+        if use_cache is None:
+            use_cache = self.config.use_cache
+        hidden, presents = self.model(
             input_ids=input_ids,
             attention_mask=attention_mask,
             position_ids=position_ids,
             inputs_embeds=inputs_embeds,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
         )
         logits = F.linear(hidden, self.model.emb.weight)
         loss = None
@@ -160,5 +209,7 @@ class NovaForCausalLM(PreTrainedModel, GenerationMixin):
                 labels[..., 1:].contiguous().view(-1),
                 ignore_index=-100,
             )
-        return CausalLMOutput(loss=loss, logits=logits)
+        return CausalLMOutputWithPast(
+            loss=loss, logits=logits, past_key_values=presents
+        )