Repository: huggingface/nanoVLM
Pull request #69: Implementing KV Cache for inference
Author: login: ariG23498
Created: 2025-05-22T08:13:55Z | Merged: 2025-05-27T14:35:20Z

Running `benchmark.py` (for 1000 tokens) with KV Cache:
```
Using device: cuda
<torch.utils.benchmark.utils.common.Measurement object at 0x750056f6b4d0>
generate_tokens(tokens, image)
setup: from __main__ import generate_tokens
  15.23 s
  1 measurement, 10 runs , 16 threads
```

Running `benchmark` without KV Cache:
```
Using device: cuda
<torch.utils.benchmark.utils.common.Measurement object at 0x7c02b496bec0>
generate_tokens(tokens, image)
setup: from __main__ import generate_tokens
  15.37 s
  1 measurement, 10 runs , 16 threads
```

Speedup: 0.9% (which is not that much). I wonder if I am missing something very big.

Update:
The speedup is about 50% now.

Here are the changes:
1. Sending one query at a time (in an autoregressive manner)
2. Carefully deciding the `is_causal` parameter for sdpa


Diff:
diff --git a/benchmark-inference.py b/benchmark-inference.py
index 49d4eddd..94d7b484 100644
--- a/benchmark-inference.py
+++ b/benchmark-inference.py
@@ -14,7 +14,7 @@
 print(f"Using device: {device}")
 
 def generate_tokens(tokens, image):
-    gen = model.generate(tokens, image, max_new_tokens=100)
+    gen = model.generate(tokens, image, max_new_tokens=1000)
 
 if __name__ == "__main__":
     model = VisionLanguageModel.from_pretrained("lusxvr/nanoVLM-222M").to(device)
diff --git a/benchmark_suite.py b/benchmark_suite.py
index 717b276a..ac98c315 100644
--- a/benchmark_suite.py
+++ b/benchmark_suite.py
@@ -76,7 +76,7 @@ def benchmark_vlm(
             mask = torch.cat((torch.ones((1, img_len), device=device), attention_mask), dim=1)
         outputs = combined
         for _ in range(max_new_tokens):
-            out = model.decoder(outputs, mask)
+            out, _ = model.decoder(outputs, mask)
             logits = out[:, -1, :]
             if not model.decoder.lm_use_tokens:
                 logits = model.decoder.head(logits)
@@ -112,7 +112,7 @@ def benchmark_vlm(
             mask = torch.cat((torch.ones((1, img_emb.size(1)), device=device), attention_mask), dim=1)
         if device.type == 'cuda': torch.cuda.synchronize()
         t0 = time.perf_counter()
-        out = model.decoder(combined, mask)
+        out, _ = model.decoder(combined, mask)
         logits = out[:, -1, :]
         if not model.decoder.lm_use_tokens: logits = model.decoder.head(logits)
         probs = torch.softmax(logits, dim=-1)
@@ -129,7 +129,7 @@ def benchmark_vlm(
         m = mask
         if m is not None: m = torch.cat((m, torch.ones((1,1), device=device)), dim=1)
         for _ in range(1, max_new_tokens):
-            out = model.decoder(seq, m)
+            out, _ = model.decoder(seq, m)
             logits = out[:, -1, :]
             if not model.decoder.lm_use_tokens: logits = model.decoder.head(logits)
             p = torch.softmax(logits, dim=-1)
diff --git a/models/language_model.py b/models/language_model.py
index ce247351..f40d3871 100644
--- a/models/language_model.py
+++ b/models/language_model.py
@@ -112,53 +112,80 @@ def __init__(self, cfg):
         if not self.sdpa:
             print("Warning: scaled dot product attention not available, using standard attention in LM.")
 
-    def forward(self, x, cos, sin, attention_mask=None):
-        B, T, C = x.size()
-
-        q = self.q_proj(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)  # (B, n_heads, T, head_dim)
-        k = self.k_proj(x).view(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2)  # (B, n_kv_heads, T, head_dim)
-        v = self.v_proj(x).view(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2)  # (B, n_kv_heads, T, head_dim)
+    def forward(self, x, cos, sin, attention_mask=None, block_kv_cache=None):
+        is_prefill = block_kv_cache is None
+
+        B, T_curr, C = x.size() # T_curr is the sequence length of the current input x
+
+        q_curr = self.q_proj(x).view(B, T_curr, self.n_heads, self.head_dim).transpose(1, 2)  # (B, n_heads, T_curr, head_dim)
+        k_curr = self.k_proj(x).view(B, T_curr, self.n_kv_heads, self.head_dim).transpose(1, 2) # (B, n_kv_heads, T_curr, head_dim)
+        v_curr = self.v_proj(x).view(B, T_curr, self.n_kv_heads, self.head_dim).transpose(1, 2) # (B, n_kv_heads, T_curr, head_dim)
+
+        # Apply rotary embeddings to the current q and k
+        q, k_rotated = apply_rotary_pos_embd(q_curr, k_curr, cos, sin)
+
+        # Check if we can use cached keys and values
+        if not is_prefill and block_kv_cache['key'] is not None:
+            # Concatenate with cached K, V
+            # k_rotated and v_curr are for the new token(s)
+            k = block_kv_cache['key']
+            v = block_kv_cache['value']
+            k = torch.cat([k, k_rotated], dim=2)
+            v = torch.cat([v, v_curr], dim=2)
+            block_kv_cache['key'] = k
+            block_kv_cache['value'] = v
+        else:
+            # No cache, this is the first pass (prefill)
+            k = k_rotated
+            v = v_curr
+            block_kv_cache = {'key': k, 'value': v}
+
+        # Repeat K, V for Grouped Query Attention
+        k_exp = k.repeat_interleave(self.n_kv_groups, dim=1) # (B, n_heads, T_kv, head_dim)
+        v_exp = v.repeat_interleave(self.n_kv_groups, dim=1) # (B, n_heads, T_kv, head_dim)
         
-        # Use precomputed positional embeddings
-        q, k = apply_rotary_pos_embd(q, k, cos, sin)
-
-        k = k.repeat_interleave(self.n_kv_groups, dim=1)
-        v = v.repeat_interleave(self.n_kv_groups, dim=1)
+        T_kv = k_exp.size(2) # Total sequence length of keys/values
 
-        # Process attention mask if provided
+        # Prepare attention mask for SDPA or manual path
+        # attention_mask is (B, T_kv_total_length), 1 for attend, 0 for pad
+        additive_attn_mask = None
         if attention_mask is not None:
-            # Create a 4D attention mask [batch_size, 1, 1, seq_length], In this format, 1 = attend, 0 = mask
-            attention_mask = attention_mask.unsqueeze(1).unsqueeze(2)  # [B, 1, 1, T]
-            padding_mask = (attention_mask == 0).transpose(-1, -2) # Use this for the manual path
-            # Convert to attention mask where 0 keeps values and -inf masks
-            attention_mask = (1.0 - attention_mask) * torch.finfo(q.dtype).min
+            # The current `attention_mask` parameter is assumed to be `[B, total_sequence_length_kv]`
+            # Let's make it `[B, 1, 1, T_kv]` for SDPA.
+            mask_for_keys = attention_mask[:, :T_kv] # Ensure mask matches key length [B, T_kv]
+            additive_attn_mask = (1.0 - mask_for_keys.unsqueeze(1).unsqueeze(2).float()) * torch.finfo(q.dtype).min
+            # This additive_attn_mask shape is [B, 1, 1, T_kv]
 
         if self.sdpa and x.device.type != 'mps':
+            # During decode, no additional masking needed as [1, T_kv] is naturally causal
+            is_causal = (T_curr == T_kv and T_curr > 1)
             y = torch.nn.functional.scaled_dot_product_attention(
-                q, k, v,
-                attn_mask=attention_mask,
+                q, k_exp, v_exp,
+                attn_mask=additive_attn_mask, 
                 dropout_p=self.dropout if self.training else 0.0,
-                is_causal=True # LM attention is causal (masked)
+                is_causal=is_causal
             )
         else:
-            attn = torch.matmul(q, k.transpose(2, 3)) / math.sqrt(self.head_dim)
-            causal_mask = torch.tril(torch.ones(T, T, device=x.device)).view(1, 1, T, T)
-            attn = attn.masked_fill(causal_mask == 0, float('-inf'))
-            if attention_mask is not None:
-                attn = attn + attention_mask 
+            # Manual attention implementation
+            attn = torch.matmul(q, k_exp.transpose(2, 3)) / math.sqrt(self.head_dim) # (B, n_heads, T_curr, T_kv)
+            # During decode: no additional masking needed as [1, T_kv] is naturally causal
+            if T_curr == T_kv and T_curr > 1:
+                causal_mask_val = torch.tril(torch.ones(T_curr, T_curr, device=x.device, dtype=torch.bool)).view(1, 1, T_curr, T_curr)
+                attn = attn.masked_fill(~causal_mask_val, float('-inf'))
+
+            if additive_attn_mask is not None: # Additive padding mask
+                # additive_attn_mask is [B,1,1,T_kv], needs to be broadcast to [B, n_heads, T_curr, T_kv]
+                attn = attn + additive_attn_mask 
 
             attn = F.softmax(attn, dim=-1)
             attn = self.attn_dropout(attn)
-            y = attn @ v
+            y = attn @ v_exp
             
-            if attention_mask is not None:
-                y = y.masked_fill(padding_mask, 0.0) # Zero out the padded positions in the output
-
-        y = y.transpose(1, 2).contiguous().view(B, T, C)  
+        y = y.transpose(1, 2).contiguous().view(B, T_curr, C)
         y = self.out_proj(y)
         y = self.resid_dropout(y)
 
-        return y
+        return y, block_kv_cache
 
 # https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py#L160
 class LanguageModelMLP(nn.Module):
@@ -188,10 +215,10 @@ def __init__(self, cfg):
         self.norm1 = RMSNorm(cfg) # Input Norm
         self.norm2 = RMSNorm(cfg) # Post Attention Norm
     
-    def forward(self, x, cos, sin, attention_mask=None):
+    def forward(self, x, cos, sin, attention_mask=None, block_kv_cache=None):
         res = x
         x = self.norm1(x)
-        x = self.attn(x, cos, sin, attention_mask)
+        x, block_kv_cache = self.attn(x, cos, sin, attention_mask, block_kv_cache)
         x = res + x
 
         res = x
@@ -199,7 +226,7 @@ def forward(self, x, cos, sin, attention_mask=None):
         x = self.mlp(x)
         x = res + x
 
-        return x
+        return x, block_kv_cache
 
 # https://github.com/meta-llama/llama3/blob/main/llama/model.py#L251
 class LanguageModel(nn.Module):
@@ -231,50 +258,74 @@ def _init_weights(self, module):
         elif isinstance(module, RMSNorm):
             module.weight.data.fill_(1.0)
 
-    def forward(self, x, attention_mask=None):
+    def forward(self, x, attention_mask=None, kv_cache=None, start_pos=0):
         if self.lm_use_tokens:
-            x = self.token_embedding(x) # Only embed the inputs when using tokens
-        
-        B , T, _ = x.size()
+            x = self.token_embedding(x)
+
+        # T_curr is the length of the current input sequence
+        B, T_curr, _ = x.size()
         
-        # Note: You could also cache these input embeddings if you want to avoid recomputing them
-        position_ids = torch.arange(T, device=x.device).unsqueeze(0).expand(B, -1) # Create position ids [0, 1, 2, ..., seq_len-1]
-        cos, sin = self.rotary_embd(position_ids) # Get rotary position embeddings
+        # Create position_ids for the current sequence based on start_pos
+        current_position_ids = torch.arange(start_pos, start_pos + T_curr, device=x.device).unsqueeze(0).expand(B, -1)
+        cos, sin = self.rotary_embd(current_position_ids) # Get rotary position embeddings for current tokens
+
+        # Initialize new KV cache if none provided
+        if kv_cache is None:
+            kv_cache = [None] * len(self.blocks)
+
+        for i, block in enumerate(self.blocks):
+            x, kv_cache[i] = block(x, cos, sin, attention_mask, kv_cache[i])
 
-        for block in self.blocks:
-            x = block(x, cos, sin, attention_mask)
         x = self.norm(x)
 
-        if self.lm_use_tokens:
-            x = self.head(x) # Compute logits if we are using tokens, otherwise stay in the embedding space
+        # Compute logits if we are using tokens, otherwise stay in the embedding space
+        if self.lm_use_tokens: 
+            x = self.head(x) 
 
-        return x
+        return x, kv_cache
 
-    @torch.no_grad()
+
+    @torch.inference_mode()
     def generate(self, inputs, max_new_tokens=20):
         # Add batch dimension if needed
         if inputs.dim() == 1:
             inputs = inputs.unsqueeze(0)
-            
-        generated = inputs.clone()
-        
-        for _ in range(max_new_tokens):
-            # Forward pass through the model
-            outputs = self.forward(generated)
-            last_output = outputs[:, -1, :]
-
+        generated_outputs = inputs.clone()
+
+        prompt_output, kv_cache_list = self.forward(
+            generated_outputs, 
+            attention_mask=None,
+            kv_cache=None,
+            start_pos=0
+        )
+        last_output = prompt_output[:, -1, :]
+
+        # Decode Phase with KV cache
+        for i in range(max_new_tokens):
             if self.lm_use_tokens:
                 # Now the model outputs logits
-                next_token = torch.argmax(last_output, dim=-1, keepdim=True)
-                generated = torch.cat((generated, next_token), dim=-1)
+                next_output = torch.argmax(last_output, dim=-1, keepdim=True)
             else:
                 # Now the model outputs embeddings
-                next_token_embedding = last_output.unsqueeze(1)  # Shape: [batch_size, 1, hidden_dim]
-                generated = torch.cat((generated, next_token_embedding), dim=1)
+                next_output = last_output.unsqueeze(1)
+
+            generated_outputs = torch.cat((generated_outputs, next_output), dim=1)
             
-            #Note: You could enable the generation to break earlier than max_new_tokens when it detects a eos token, but this does not work in batched generation (output tensors need to have the same size)
+            # The token being processed is `next_token`. Its position is `generated_outputs.size(1) - 1`.
+            current_token_start_pos = generated_outputs.size(1) - 1
+
+            if i == max_new_tokens - 1: 
+                break
+
+            decode_step_output, kv_cache_list = self.forward(
+                next_output, 
+                attention_mask=None,
+                kv_cache=kv_cache_list,
+                start_pos=current_token_start_pos
+            )
+            last_output = decode_step_output[:, -1, :] 
     
-        return generated
+        return generated_outputs
 
     # Load the model from a pretrained HuggingFace model (we don't want to have to train the Language Backbone from scratch)
     @classmethod
diff --git a/models/vision_language_model.py b/models/vision_language_model.py
index 83581bc7..c876b458 100644
--- a/models/vision_language_model.py
+++ b/models/vision_language_model.py
@@ -48,7 +48,7 @@ def forward(self, input_ids, image, attention_mask=None, targets=None):
             # Combine image and token attention masks
             attention_mask = torch.cat((image_attention_mask, attention_mask), dim=1)
 
-        logits = self.decoder(combined_embd, attention_mask) # Not logits yet, but easier to return like this
+        logits, _ = self.decoder(combined_embd, attention_mask=attention_mask) # Not logits yet, but easier to return like this
 
         loss = None
         if targets is not None:
@@ -59,59 +59,102 @@ def forward(self, input_ids, image, attention_mask=None, targets=None):
 
         return logits, loss
 
-    @torch.no_grad()
-    def generate(self, input_ids, image, attention_mask=None, max_new_tokens=5, top_k=50, top_p=0.9, temperature=0.5, greedy=False):
-        # Process image through vision encoder and projection
-        image_embd = self.vision_encoder(image)
-        image_embd = self.MP(image_embd)
-        
-        # Embed initial tokens
-        token_embd = self.decoder.token_embedding(input_ids)
-        
-        # Concatenate image embeddings with token embeddings
-        combined_embd = torch.cat((image_embd, token_embd), dim=1)
+    @torch.inference_mode()
+    def generate(self, input_ids, image, attention_mask=None, max_new_tokens=5, top_k=50, top_p=0.9, temperature=0.5, greedy=False, use_kv_cache: bool = True):
+
+        # 1. Process image
+        image_embd = self.vision_encoder(image) # [B, T_img, D_model]
+        image_embd = self.MP(image_embd)      # [B, T_img, D_lm]
+
+        # 2. Embed initial text prompt tokens
+        prompt_token_embeds = self.decoder.token_embedding(input_ids) # [B, T_prompt_text, D_lm]
+
+        # 3. Combine image and text prompt embeddings for prefill
+        initial_combined_embeds = torch.cat((image_embd, prompt_token_embeds), dim=1) # [B, T_img + T_prompt_text, D_lm]
+        current_total_seq_len = initial_combined_embeds.size(1)
 
         batch_size = image_embd.size(0)
-        img_seq_len = image_embd.size(1)
-        # Adjust attention mask to account for image tokens
         if attention_mask is not None:
             # Create mask of 1s for image tokens (all image tokens should be attended to)
+            img_seq_len = image_embd.size(1)
             image_attention_mask = torch.ones((batch_size, img_seq_len), device=attention_mask.device, dtype=attention_mask.dtype)
+            
+            # Combine image and token attention masks
             attention_mask = torch.cat((image_attention_mask, attention_mask), dim=1)
         
-        # Generate from combined embeddings using the decoder
-        # We need to use the decoder's forward function and not its generate method
-        # because we want to keep track of the image prefix
-        outputs = combined_embd
-        generated_tokens = torch.zeros((batch_size, max_new_tokens), device=input_ids.device, dtype=input_ids.dtype)
+        # --- Multimodal Prefill Phase ---
+        prefill_output, kv_cache_list = self.decoder(
+            initial_combined_embeds,
+            attention_mask=attention_mask,
+            kv_cache=None,
+            start_pos=0
+        )
         
-        #Note: Here you could implement improvements like e.g. KV caching
-        for i in range(max_new_tokens):
-            model_out = self.decoder(outputs, attention_mask)
-            
-            # Get predictions for the last token only (normally this is the embedding, not the logits)
-            last_token_logits = model_out[:, -1, :]
-            
-            # Apply head to get logits (if model is in embedding mode)
-            if not self.decoder.lm_use_tokens:
-                last_token_logits = self.decoder.head(last_token_logits)
+        last_token_output_from_prefill = prefill_output[:, -1, :] 
+        
+        if not self.decoder.lm_use_tokens:
+            current_logits = self.decoder.head(last_token_output_from_prefill) 
+        else:
+            current_logits = last_token_output_from_prefill 
+
+        # Store newly generated token IDs
+        newly_generated_ids_list = []
+
+        # --- Decode Phase by sampling tokens autoregressively using the kv-cache ---
+        for _ in range(max_new_tokens):
             if greedy:
-                next_token = torch.argmax(last_token_logits, dim=-1, keepdim=True)
+                next_token_id = torch.argmax(current_logits, dim=-1, keepdim=True)
             else:
-                filtered_logits = top_k_top_p_filtering(last_token_logits, top_k=top_k, top_p=top_p)
-                probs = torch.softmax(filtered_logits/temperature, dim=-1)
-                next_token = torch.multinomial(probs, num_samples=1)
-                
-            generated_tokens[:, i] = next_token.squeeze(-1)
+                filtered_logits = top_k_top_p_filtering(current_logits, top_k=top_k, top_p=top_p)
+                probs = torch.softmax(filtered_logits / temperature, dim=-1)
+                next_token_id = torch.multinomial(probs, num_samples=1)
+            
+            newly_generated_ids_list.append(next_token_id)
             
-            # Convert to embedding and append
-            next_embd = self.decoder.token_embedding(next_token)
-            outputs = torch.cat((outputs, next_embd), dim=1)
+            # Embed the newly generated token
+            next_token_embed = self.decoder.token_embedding(next_token_id) # [B, 1, D_lm]
+            
+            # The start_pos for the new token is the current total sequence length *before* adding this new token
+            current_token_start_pos = current_total_seq_len
+            current_total_seq_len += 1
 
+            # update attention mask
             if attention_mask is not None:
-                attention_mask = torch.cat((attention_mask, torch.ones((batch_size, 1), device=attention_mask.device)), dim=1)
+                attention_mask = torch.cat((attention_mask, torch.ones((batch_size, 1), device=attention_mask.device, dtype=attention_mask.dtype)), dim=1)
+
+            if use_kv_cache:
+                # With KV cache: only process the new token
+                decode_step_output, kv_cache_list = self.decoder(
+                    next_token_embed,
+                    attention_mask=attention_mask,
+                    kv_cache=kv_cache_list,
+                    start_pos=current_token_start_pos
+                )
+            else:
+                # Without KV cache: process the entire sequence from scratch
+                # Reconstruct the full sequence: image + prompt + generated tokens so far
+                generated_token_embeds = torch.cat([self.decoder.token_embedding(tid) for tid in newly_generated_ids_list], dim=1)
+                full_sequence_embeds = torch.cat([initial_combined_embeds, generated_token_embeds], dim=1)
+ 
+                decode_step_output, _ = self.decoder(
+                    full_sequence_embeds,
+                    attention_mask=attention_mask,
+                    kv_cache=None,
+                    start_pos=0
+                )
+                
+            last_token_output = decode_step_output[:, -1, :] 
+            
+            # Apply head to get logits (if model is in embedding mode)
+            if not self.decoder.lm_use_tokens:
+                current_logits = self.decoder.head(last_token_output)
+            else:
+                current_logits = last_token_output
         
-        return generated_tokens
+        if not newly_generated_ids_list: # Handle case where max_new_tokens might be 0
+             return torch.empty((B,0), dtype=torch.long, device=input_ids.device)
+
+        return torch.cat(newly_generated_ids_list, dim=1)
 
     @classmethod
     def from_pretrained(
diff --git a/nanoVLM.ipynb b/nanoVLM.ipynb
index 91022681..23841c08 100644
--- a/nanoVLM.ipynb
+++ b/nanoVLM.ipynb
@@ -261,10 +261,9 @@
         "            input_ids = batch['input_ids'].to(device)\n",
         "            labels = batch['labels'].to(device)\n",
         "            attention_mask = batch['attention_mask'].to(device)\n",
-        "\n",
         "            correct_answer = tokenizer.batch_decode(labels, skip_special_tokens=True)\n",
         "\n",
-        "            gen = model.generate(input_ids, image, attention_mask)\n",
+        "            gen = model.generate(input_ids, image, attention_mask, greedy=True)\n",
         "            model_output = tokenizer.batch_decode(gen, skip_special_tokens=True)\n",
         "\n",
         "            is_correct = utils.check_multiple_choice_with_regex(model_output, correct_answer)\n",
diff --git a/tests/test_language_model.py b/tests/test_language_model.py
new file mode 100644
index 00000000..8ecd2267
--- /dev/null
+++ b/tests/test_language_model.py
@@ -0,0 +1,98 @@
+import torch
+import unittest
+from models.language_model import LanguageModel
+from types import SimpleNamespace
+
+class TestLanguageModel(unittest.TestCase):
+    def setUp(self):
+        # Minimal config for testing
+        self.cfg = SimpleNamespace(
+            lm_hidden_dim=64,
+            lm_inter_dim=128,
+            lm_rms_eps=1e-5,
+            lm_re_base=10000.0,
+            lm_max_position_embeddings=1024,
+            lm_attn_scaling=1.0,
+            lm_vocab_size=100, # Small vocab for testing
+            lm_n_heads=4,
+            lm_n_kv_heads=2,
+            lm_dropout=0.0,
+            lm_n_blocks=2,
+            lm_use_tokens=True,
+            lm_tie_weights=True
+        )
+        self.model = LanguageModel(self.cfg)
+        self.model.eval() # Set model to evaluation mode
+
+    def test_kv_caching_consistency(self):
+        # Input for the model
+        batch_size = 16
+        seq_len = 1000
+        input_ids = torch.randint(0, self.cfg.lm_vocab_size, (batch_size, seq_len))
+
+        # Forward pass without KV caching (prefill)
+        output_no_cache, _ = self.model(input_ids, start_pos=0)
+
+        # Forward pass with KV caching
+        # 1. Prefill phase
+        prefill_output, kv_cache_prefill = self.model(input_ids[:, :-1], start_pos=0)
+        
+        # 2. Decode phase (one token at a time)
+        # We expect the output of the last token from prefill + decode to match the no_cache output
+        # for that same token.
+        
+        # Get the last token's input_id for the decode step
+        last_token_input = input_ids[:, -1].unsqueeze(-1) # Shape: [B, 1]
+        
+        # The start_pos for this token is seq_len - 1
+        output_with_cache_last_token, _ = self.model(
+            last_token_input, 
+            kv_cache=kv_cache_prefill, 
+            start_pos=seq_len - 1
+        )
+
+        # Compare the logits for the last token
+        # output_no_cache is [B, seq_len, vocab_size]
+        # output_with_cache_last_token is [B, 1, vocab_size]
+        
+        # We compare the last token's output from the no_cache run
+        # with the single token output from the with_cache run.
+        logits_no_cache_last_token = output_no_cache[:, -1, :]
+        logits_with_cache_last_token = output_with_cache_last_token[:, 0, :]
+
+        self.assertTrue(
+            torch.allclose(logits_no_cache_last_token, logits_with_cache_last_token, atol=1e-5),
+            "Outputs with and without KV caching do not match for the last token."
+        )
+
+        # Let's also test a multi-step decode to be more thorough
+        # We'll compare the full sequence output if we decode token by token
+        
+        # Reset for a full token-by-token generation using KV cache
+        current_input = input_ids[:, :1] # Start with the first token
+        output_tokens_with_cache_list = []
+        kv_cache_step = None
+
+        for i in range(seq_len):
+            if i > 0:
+                current_input = input_ids[:, i:i+1] # Next token
+            
+            # The start_pos for the current token is simply i
+            output_step, kv_cache_step = self.model(
+                current_input,
+                kv_cache=kv_cache_step,
+                start_pos=i 
+            )
+            # output_step is [B, 1, vocab_size]
+            output_tokens_with_cache_list.append(output_step)
+        
+        # Concatenate all single token outputs
+        output_with_cache_full = torch.cat(output_tokens_with_cache_list, dim=1) # [B, seq_len, vocab_size]
+
+        self.assertTrue(
+            torch.allclose(output_no_cache, output_with_cache_full, atol=1e-5),
+            "Full sequence outputs with and without KV caching do not match."
+        )
+
+if __name__ == '__main__':
+    unittest.main() 
\ No newline at end of file
diff --git a/tests/test_vision_language_model.py b/tests/test_vision_language_model.py
new file mode 100644
index 00000000..9c6fec7f
--- /dev/null
+++ b/tests/test_vision_language_model.py
@@ -0,0 +1,77 @@
+import torch
+import unittest
+from models.vision_language_model import VisionLanguageModel
+from models.config import VLMConfig # Assuming VLMConfig is in models.config
+from types import SimpleNamespace
+
+class TestVisionLanguageModel(unittest.TestCase):
+    def setUp(self):
+        # Minimal config for testing VLM
+        # We need to ensure the sub-configs for ViT and LanguageModel are also present
+        self.cfg = VLMConfig(
+            # ViT specific (minimal)
+            vit_model_type='testing',
+            vit_patch_size=16,
+            vit_hidden_dim=48, # Small for testing
+            vit_inter_dim=96,
+            vit_n_heads=3,
+            vit_n_blocks=1,
+            vit_img_size=32, # Small image size
+            vit_dropout=0.0,
+            # LM specific
+            lm_model_type='testing',
+            lm_hidden_dim=64,
+            lm_inter_dim=128,
+            lm_rms_eps=1e-5,
+            lm_re_base=10000.0,
+            lm_max_position_embeddings=512,
+            lm_attn_scaling=1.0,
+            lm_vocab_size=100, # Small vocab
+            lm_n_heads=4,
+            lm_n_kv_heads=2,
+            lm_dropout=0.0,
+            lm_n_blocks=2,
+            lm_use_tokens=False,
+            lm_tie_weights=True,
+            # MP specific
+            mp_pixel_shuffle_factor=2,
+        )
+        
+        self.model = VisionLanguageModel(self.cfg, load_backbone=False) # Don't load pretrained for unit test
+        self.model.eval() # Set model to evaluation mode
+
+    def test_generate_kv_caching_consistency(self):
+        batch_size = 16
+        prompt_seq_len = 32
+        max_new_tokens = 16 # Generate a few tokens
+
+        # Dummy image (Batch, Channels, Height, Width)
+        image_input = torch.randn(batch_size, 3, self.cfg.vit_img_size, self.cfg.vit_img_size)
+        # Dummy prompt input_ids
+        prompt_ids = torch.randint(0, self.cfg.lm_vocab_size, (batch_size, prompt_seq_len))
+
+        # Generation with KV caching (default)
+        generated_ids_with_cache = self.model.generate(
+            prompt_ids,
+            image_input,
+            max_new_tokens=max_new_tokens,
+            use_kv_cache=True,
+            greedy=True # Use greedy for deterministic output
+        )
+
+        # Generation without KV caching
+        generated_ids_without_cache = self.model.generate(
+            prompt_ids,
+            image_input,
+            max_new_tokens=max_new_tokens,
+            use_kv_cache=False,
+            greedy=True # Use greedy for deterministic output
+        )
+        
+        self.assertTrue(
+            torch.equal(generated_ids_with_cache, generated_ids_without_cache),
+            f"Generated token IDs with and without KV caching do not match.\nWith cache: {generated_ids_with_cache}\nWithout cache: {generated_ids_without_cache}"
+        )
+
+if __name__ == '__main__':
+    unittest.main() 
\ No newline at end of file
diff --git a/train.py b/train.py
index 1079337a..a83e0e4c 100644
--- a/train.py
+++ b/train.py
@@ -165,10 +165,10 @@ def test_mmstar(model, tokenizer, test_loader, device):
             input_ids = batch['input_ids'].to(device)
             labels = batch['labels'].to(device)
             attention_mask = batch['attention_mask'].to(device)
-            
+
             correct_answer = tokenizer.batch_decode(labels, skip_special_tokens=True)
             
-            gen = model.generate(input_ids, image, attention_mask)
+            gen = model.generate(input_ids, image, attention_mask, greedy=True)
             model_output = tokenizer.batch_decode(gen, skip_special_tokens=True)
             
             is_correct = utils.check_multiple_choice_with_regex(model_output, correct_answer)


Reviews:
user: login: ariG23498
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-22T08:24:21Z
user: login: kashif
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-23T08:09:31Z
user: login: kashif
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-23T08:09:43Z
user: login: gante
body: In general looks good, but I think we need to benchmark memory as well: I suspect there is some memory usage degradation (see comments below).

In general, using KV cache should have a very similar memory footprint inside the forward pass, let's double-check that! :D
state: COMMENTED
author_association: NONE
submitted_at: 2025-05-23T09:11:45Z
user: login: ariG23498
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-23T10:04:36Z
user: login: ariG23498
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-23T10:05:07Z
user: login: lusxvr
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-23T12:37:28Z
user: login: kashif
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-23T12:48:15Z
user: login: kashif
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-26T08:00:41Z
user: login: andimarafioti
body: I wrote two tests and they don't pass. The first one that I think should trivially pass can be run with:
```bash
python -m unittest tests.test_language_model
```

The second one is a bit more involved, I had to make small changes to the VLM model file to make the test runnable, and I'm not sure passing 0 is the right way to do it when not using the KV cache. Let me know if I messed up the tests, but also feel free to modify them into something similar that passes. Ideally, the results with/without kv-cache should be identical, right?

Second test is:
```bash
python -m unittest tests.test_vision_language_model
```

state: CHANGES_REQUESTED
author_association: MEMBER
submitted_at: 2025-05-26T10:35:21Z
user: login: andimarafioti
body: Thank you so much for going through this! I'm sorry for the thousand reviews, but I'm super happy with the results!

<img src="https://media2.giphy.com/media/v1.Y2lkPWJkM2VhNTdlaWN4czZob2MwbzB6eGVpYzF3NmM5bGM3bzM4aW1idmhqNTJxbHhzNiZlcD12MV9naWZzX3RyZW5kaW5nJmN0PWc/Lv2VhwHrt6ljhvZ6LF/giphy.gif"/>
state: APPROVED
author_association: MEMBER
submitted_at: 2025-05-27T14:28:11Z

Review comments:
path: benchmark.py
user: login: ariG23498
body: I will remove this in favour of `benchmark-inference.py` in the repo.
created_at: 2025-05-22T08:24:20Z
author_association: COLLABORATOR
diff_hunk: @@ -59,59 +59,81 @@ def forward(self, input_ids, image, attention_mask=None, targets=None):
 
         return logits, loss
 
-    @torch.no_grad()
-    def generate(self, input_ids, image, attention_mask=None, max_new_tokens=5, top_k=50, top_p=0.9, temperature=0.5, greedy=False):
-        # Process image through vision encoder and projection
-        image_embd = self.vision_encoder(image)
-        image_embd = self.MP(image_embd)
-        
-        # Embed initial tokens
-        token_embd = self.decoder.token_embedding(input_ids)
-        
-        # Concatenate image embeddings with token embeddings
-        combined_embd = torch.cat((image_embd, token_embd), dim=1)
+    @torch.inference_mode()
+    def generate(self, input_ids, image, max_new_tokens=5, top_k=50, top_p=0.9, temperature=0.5, greedy=False):
+        # input_ids: [B, T_prompt_text]
+        # image: [B, C, H, W]
+        # max_new_tokens: int, maximum number of tokens to generate
 
-        batch_size = image_embd.size(0)
-        img_seq_len = image_embd.size(1)
-        # Adjust attention mask to account for image tokens
-        if attention_mask is not None:
-            # Create mask of 1s for image tokens (all image tokens should be attended to)
-            image_attention_mask = torch.ones((batch_size, img_seq_len), device=attention_mask.device, dtype=attention_mask.dtype)
-            attention_mask = torch.cat((image_attention_mask, attention_mask), dim=1)
+        B = image.size(0)
+
+        # 1. Process image
+        image_embd = self.vision_encoder(image) # [B, T_img, D_model]
+        image_embd = self.MP(image_embd)      # [B, T_img, D_lm]
+
+        # 2. Embed initial text prompt tokens
+        prompt_token_embeds = self.decoder.token_embedding(input_ids) # [B, T_prompt_text, D_lm]
+
+        # 3. Combine image and text prompt embeddings for prefill
+        initial_combined_embeds = torch.cat((image_embd, prompt_token_embeds), dim=1) # [B, T_img + T_prompt_text, D_lm]
+        current_total_seq_len = initial_combined_embeds.size(1)
+        
+        # --- Multimodal Prefill Phase ---
+        prefill_output, kv_cache_list = self.decoder.forward(
path: models/vision_language_model.py
user: login: kashif
body: ```suggestion
        prefill_output, kv_cache_list = self.decoder(
```
created_at: 2025-05-23T08:09:31Z
author_association: COLLABORATOR
diff_hunk: @@ -59,59 +59,81 @@ def forward(self, input_ids, image, attention_mask=None, targets=None):
 
         return logits, loss
 
-    @torch.no_grad()
-    def generate(self, input_ids, image, attention_mask=None, max_new_tokens=5, top_k=50, top_p=0.9, temperature=0.5, greedy=False):
-        # Process image through vision encoder and projection
-        image_embd = self.vision_encoder(image)
-        image_embd = self.MP(image_embd)
-        
-        # Embed initial tokens
-        token_embd = self.decoder.token_embedding(input_ids)
-        
-        # Concatenate image embeddings with token embeddings
-        combined_embd = torch.cat((image_embd, token_embd), dim=1)
+    @torch.inference_mode()
+    def generate(self, input_ids, image, max_new_tokens=5, top_k=50, top_p=0.9, temperature=0.5, greedy=False):
+        # input_ids: [B, T_prompt_text]
+        # image: [B, C, H, W]
+        # max_new_tokens: int, maximum number of tokens to generate
 
-        batch_size = image_embd.size(0)
-        img_seq_len = image_embd.size(1)
-        # Adjust attention mask to account for image tokens
-        if attention_mask is not None:
-            # Create mask of 1s for image tokens (all image tokens should be attended to)
-            image_attention_mask = torch.ones((batch_size, img_seq_len), device=attention_mask.device, dtype=attention_mask.dtype)
-            attention_mask = torch.cat((image_attention_mask, attention_mask), dim=1)
+        B = image.size(0)
+
+        # 1. Process image
+        image_embd = self.vision_encoder(image) # [B, T_img, D_model]
+        image_embd = self.MP(image_embd)      # [B, T_img, D_lm]
+
+        # 2. Embed initial text prompt tokens
+        prompt_token_embeds = self.decoder.token_embedding(input_ids) # [B, T_prompt_text, D_lm]
+
+        # 3. Combine image and text prompt embeddings for prefill
+        initial_combined_embeds = torch.cat((image_embd, prompt_token_embeds), dim=1) # [B, T_img + T_prompt_text, D_lm]
+        current_total_seq_len = initial_combined_embeds.size(1)
+        
+        # --- Multimodal Prefill Phase ---
+        prefill_output, kv_cache_list = self.decoder.forward(
+            initial_combined_embeds,
+            attention_mask=None,
+            kv_cache=None,
+            start_pos=0
+        )
         
-        # Generate from combined embeddings using the decoder
-        # We need to use the decoder's forward function and not its generate method
-        # because we want to keep track of the image prefix
-        outputs = combined_embd
-        generated_tokens = torch.zeros((batch_size, max_new_tokens), device=input_ids.device, dtype=input_ids.dtype)
+        last_token_output_from_prefill = prefill_output[:, -1, :] 
         
-        #Note: Here you could implement improvements like e.g. KV caching
-        for i in range(max_new_tokens):
-            model_out = self.decoder(outputs, attention_mask)
+        if not self.decoder.lm_use_tokens:
+            current_logits = self.decoder.head(last_token_output_from_prefill) 
+        else:
+            current_logits = last_token_output_from_prefill 
+
+        # Store newly generated token IDs
+        newly_generated_ids_list = []
+
+        # --- Decode Phase with Sampling tokens one by one ---
+        for _ in range(max_new_tokens):
+            if greedy:
+                next_token_id = torch.argmax(current_logits, dim=-1, keepdim=True)
+            else:
+                filtered_logits = top_k_top_p_filtering(current_logits, top_k=top_k, top_p=top_p)
+                probs = torch.softmax(filtered_logits / temperature, dim=-1)
+                next_token_id = torch.multinomial(probs, num_samples=1)
+            
+            newly_generated_ids_list.append(next_token_id)
             
-            # Get predictions for the last token only (normally this is the embedding, not the logits)
-            last_token_logits = model_out[:, -1, :]
+            # Embed the newly generated token
+            next_token_embed = self.decoder.token_embedding(next_token_id) # [B, 1, D_lm]
+            
+            # The start_pos for the new token is the current total sequence length *before* adding this new token
+            current_token_start_pos = current_total_seq_len
+            current_total_seq_len += 1
+
+            # Call decoder.forward with the new token's embedding and the updated KV cache
+            decode_step_output, kv_cache_list = self.decoder.forward(
path: models/vision_language_model.py
user: login: kashif
body: ```suggestion
            decode_step_output, kv_cache_list = self.decoder(
```
created_at: 2025-05-23T08:09:43Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,40 @@
+import torch; torch.manual_seed(0);
+from PIL import Image
+
+from models.vision_language_model import VisionLanguageModel
+from data.processors import get_tokenizer, get_image_processor
+
+from torch.utils import benchmark
+
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print(f"Using device: {device}")
+
+def generate_tokens(tokens, image):
+    gen = model.generate(tokens, image, max_new_tokens=100)
path: benchmark.py
user: login: gante
body: Longer generations are better for measuring speedups! The speedup is larger the larger the context window is
created_at: 2025-05-23T08:41:08Z
author_association: NONE
diff_hunk: @@ -112,53 +112,83 @@ def __init__(self, cfg):
         if not self.sdpa:
             print("Warning: scaled dot product attention not available, using standard attention in LM.")
 
-    def forward(self, x, cos, sin, attention_mask=None):
-        B, T, C = x.size()
-
-        q = self.q_proj(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)  # (B, n_heads, T, head_dim)
-        k = self.k_proj(x).view(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2)  # (B, n_kv_heads, T, head_dim)
-        v = self.v_proj(x).view(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2)  # (B, n_kv_heads, T, head_dim)
+    def forward(self, x, cos, sin, attention_mask=None, kv_cache=None):
+        B, T_curr, C = x.size() # T_curr is the sequence length of the current input x
+
+        q_curr = self.q_proj(x).view(B, T_curr, self.n_heads, self.head_dim).transpose(1, 2)  # (B, n_heads, T_curr, head_dim)
+        k_curr = self.k_proj(x).view(B, T_curr, self.n_kv_heads, self.head_dim).transpose(1, 2) # (B, n_kv_heads, T_curr, head_dim)
+        v_curr = self.v_proj(x).view(B, T_curr, self.n_kv_heads, self.head_dim).transpose(1, 2) # (B, n_kv_heads, T_curr, head_dim)
+
+        # Apply rotary embeddings to the current q and k
+        q, k_rotated = apply_rotary_pos_embd(q_curr, k_curr, cos, sin)
+
+        # Check if we can use cached keys and values
+        if kv_cache is not None and kv_cache['key'] is not None:
+            # Concatenate with cached K, V
+            # k_rotated and v_curr are for the new token(s)
+            k_past = kv_cache['key']
+            v_past = kv_cache['value']
+            k = torch.cat([k_past, k_rotated], dim=2)
+            v = torch.cat([v_past, v_curr], dim=2)
+        else:
+            # No cache, this is the first pass (prefill)
+            k = k_rotated
+            v = v_curr
         
-        # Use precomputed positional embeddings
-        q, k = apply_rotary_pos_embd(q, k, cos, sin)
+        new_kv_cache = {'key': k, 'value': v}
path: models/language_model.py
user: login: gante
body: This implementation concatenates the existing cache into a new tensor (e.g. `k_past` + new data -> `k`). This means we will be allocating ~2x the memory needed for the KV values -- once in `k_past`, and again in `k`. My suggestion would be to reuse the same tensor objects again and again (e.g. rename `k_past` into `k` and then do `k = torch.cat([k, k_rotated], dim=2)`) :) 

Storing the tensors in new dicts/lists should be fine, though.


```py
>>> import torch
>>> torch.cuda.memory_allocated()
0
>>> a = torch.randint(100, (1024, 1024), device="cuda")
>>> torch.cuda.memory_allocated()
8388608
>>> # using the tensor objects in lists/dicts relies on references
>>> b = [a]
>>> torch.cuda.memory_allocated()
8388608
>>> c = {"foo": a}
>>> torch.cuda.memory_allocated()
8388608
>>> # torch.cat returns a new tensor, allocating new memory
>>> b = torch.cat((a, a), dim=0)
>>> torch.cuda.memory_allocated()
25165824
>>> # this is 3x the original memory: 2x in `b`, 1x in `a`
```

```py
>>> import torch
>>> torch.cuda.memory_allocated()
0
>>> a = torch.randint(100, (1024, 1024), device="cuda")
>>> torch.cuda.memory_allocated()
8388608
>>> a = torch.cat((a, a), dim=0)
>>> torch.cuda.memory_allocated()
16777216
>>> # concatenating into the new `a` removes the reference to the old `a` tensor 
```
created_at: 2025-05-23T08:56:20Z
author_association: NONE
diff_hunk: @@ -231,50 +261,76 @@ def _init_weights(self, module):
         elif isinstance(module, RMSNorm):
             module.weight.data.fill_(1.0)
 
-    def forward(self, x, attention_mask=None):
+    def forward(self, x, attention_mask=None, kv_cache=None, start_pos=0):
         if self.lm_use_tokens:
-            x = self.token_embedding(x) # Only embed the inputs when using tokens
-        
-        B , T, _ = x.size()
+            x = self.token_embedding(x)
+
+        # T_curr is the length of the current input sequence
+        B, T_curr, _ = x.size()
         
-        # Note: You could also cache these input embeddings if you want to avoid recomputing them
-        position_ids = torch.arange(T, device=x.device).unsqueeze(0).expand(B, -1) # Create position ids [0, 1, 2, ..., seq_len-1]
-        cos, sin = self.rotary_embd(position_ids) # Get rotary position embeddings
+        # Create position_ids for the current sequence based on start_pos
+        current_position_ids = torch.arange(start_pos, start_pos + T_curr, device=x.device).unsqueeze(0).expand(B, -1)
+        cos, sin = self.rotary_embd(current_position_ids) # Get rotary position embeddings for current tokens
+
+        # Initialize new KV cache if none provided
+        new_kv_cache_list = []
+        if kv_cache is None:
+            kv_cache = [None] * len(self.blocks)
+
+        for i, block in enumerate(self.blocks):
+            x, block_kv_cache = block(x, cos, sin, attention_mask, kv_cache[i])
+            new_kv_cache_list.append(block_kv_cache)
path: models/language_model.py
user: login: gante
body: At this point, we will have references to `kv_cache` (input KV cache) and `new_kv_cache_list` (output KV cache), i.e. references to the old and new tensors. I think we will have the same effect as above, i.e. duplicated memory.

I see two options:
- Rely on side-effects: the attention layer updates the cache object.
- Do something like `kv_cache[i] = block_kv_cache` on this line. We would still get a small memory peak, as we would have two cache objects for a given layer here. It would still be much better than having a copy for the whole cache :D
created_at: 2025-05-23T09:09:39Z
author_association: NONE
diff_hunk: @@ -112,53 +112,83 @@ def __init__(self, cfg):
         if not self.sdpa:
             print("Warning: scaled dot product attention not available, using standard attention in LM.")
 
-    def forward(self, x, cos, sin, attention_mask=None):
-        B, T, C = x.size()
-
-        q = self.q_proj(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)  # (B, n_heads, T, head_dim)
-        k = self.k_proj(x).view(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2)  # (B, n_kv_heads, T, head_dim)
-        v = self.v_proj(x).view(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2)  # (B, n_kv_heads, T, head_dim)
+    def forward(self, x, cos, sin, attention_mask=None, kv_cache=None):
+        B, T_curr, C = x.size() # T_curr is the sequence length of the current input x
+
+        q_curr = self.q_proj(x).view(B, T_curr, self.n_heads, self.head_dim).transpose(1, 2)  # (B, n_heads, T_curr, head_dim)
+        k_curr = self.k_proj(x).view(B, T_curr, self.n_kv_heads, self.head_dim).transpose(1, 2) # (B, n_kv_heads, T_curr, head_dim)
+        v_curr = self.v_proj(x).view(B, T_curr, self.n_kv_heads, self.head_dim).transpose(1, 2) # (B, n_kv_heads, T_curr, head_dim)
+
+        # Apply rotary embeddings to the current q and k
+        q, k_rotated = apply_rotary_pos_embd(q_curr, k_curr, cos, sin)
+
+        # Check if we can use cached keys and values
+        if kv_cache is not None and kv_cache['key'] is not None:
+            # Concatenate with cached K, V
+            # k_rotated and v_curr are for the new token(s)
+            k_past = kv_cache['key']
+            v_past = kv_cache['value']
+            k = torch.cat([k_past, k_rotated], dim=2)
+            v = torch.cat([v_past, v_curr], dim=2)
+        else:
+            # No cache, this is the first pass (prefill)
+            k = k_rotated
+            v = v_curr
         
-        # Use precomputed positional embeddings
-        q, k = apply_rotary_pos_embd(q, k, cos, sin)
+        new_kv_cache = {'key': k, 'value': v}
path: models/language_model.py
user: login: ariG23498
body: Very helpful. Implemented it -> https://github.com/huggingface/nanoVLM/pull/69/commits/469022915c4d6eb7aff4e418b5436275f8e55441
created_at: 2025-05-23T10:04:36Z
author_association: COLLABORATOR
diff_hunk: @@ -231,50 +261,76 @@ def _init_weights(self, module):
         elif isinstance(module, RMSNorm):
             module.weight.data.fill_(1.0)
 
-    def forward(self, x, attention_mask=None):
+    def forward(self, x, attention_mask=None, kv_cache=None, start_pos=0):
         if self.lm_use_tokens:
-            x = self.token_embedding(x) # Only embed the inputs when using tokens
-        
-        B , T, _ = x.size()
+            x = self.token_embedding(x)
+
+        # T_curr is the length of the current input sequence
+        B, T_curr, _ = x.size()
         
-        # Note: You could also cache these input embeddings if you want to avoid recomputing them
-        position_ids = torch.arange(T, device=x.device).unsqueeze(0).expand(B, -1) # Create position ids [0, 1, 2, ..., seq_len-1]
-        cos, sin = self.rotary_embd(position_ids) # Get rotary position embeddings
+        # Create position_ids for the current sequence based on start_pos
+        current_position_ids = torch.arange(start_pos, start_pos + T_curr, device=x.device).unsqueeze(0).expand(B, -1)
+        cos, sin = self.rotary_embd(current_position_ids) # Get rotary position embeddings for current tokens
+
+        # Initialize new KV cache if none provided
+        new_kv_cache_list = []
+        if kv_cache is None:
+            kv_cache = [None] * len(self.blocks)
+
+        for i, block in enumerate(self.blocks):
+            x, block_kv_cache = block(x, cos, sin, attention_mask, kv_cache[i])
+            new_kv_cache_list.append(block_kv_cache)
path: models/language_model.py
user: login: ariG23498
body: I think the lowest hanging fruit is `kv_cahce[i]=block_kv_cache`.

Thanks for catching this.
created_at: 2025-05-23T10:05:07Z
author_association: COLLABORATOR
diff_hunk: @@ -231,50 +264,74 @@ def _init_weights(self, module):
         elif isinstance(module, RMSNorm):
             module.weight.data.fill_(1.0)
 
-    def forward(self, x, attention_mask=None):
+    def forward(self, x, attention_mask=None, kv_cache=None, start_pos=0):
         if self.lm_use_tokens:
-            x = self.token_embedding(x) # Only embed the inputs when using tokens
-        
-        B , T, _ = x.size()
+            x = self.token_embedding(x)
+
+        # T_curr is the length of the current input sequence
+        B, T_curr, _ = x.size()
         
-        # Note: You could also cache these input embeddings if you want to avoid recomputing them
-        position_ids = torch.arange(T, device=x.device).unsqueeze(0).expand(B, -1) # Create position ids [0, 1, 2, ..., seq_len-1]
-        cos, sin = self.rotary_embd(position_ids) # Get rotary position embeddings
+        # Create position_ids for the current sequence based on start_pos
+        current_position_ids = torch.arange(start_pos, start_pos + T_curr, device=x.device).unsqueeze(0).expand(B, -1)
+        cos, sin = self.rotary_embd(current_position_ids) # Get rotary position embeddings for current tokens
+
+        # Initialize new KV cache if none provided
+        if kv_cache is None:
+            kv_cache = [None] * len(self.blocks)
+
+        for i, block in enumerate(self.blocks):
+            x, kv_cache[i] = block(x, cos, sin, attention_mask, kv_cache[i])
 
-        for block in self.blocks:
-            x = block(x, cos, sin, attention_mask)
         x = self.norm(x)
 
-        if self.lm_use_tokens:
-            x = self.head(x) # Compute logits if we are using tokens, otherwise stay in the embedding space
+        # Compute logits if we are using tokens, otherwise stay in the embedding space
+        if self.lm_use_tokens: 
+            x = self.head(x) 
 
-        return x
+        return x, kv_cache
 
-    @torch.no_grad()
+
+    @torch.inference_mode()
path: models/language_model.py
user: login: lusxvr
body: What is the difference between `torch.no_grad()` and `torch.inference_mode()`
created_at: 2025-05-23T12:37:28Z
author_association: CONTRIBUTOR
diff_hunk: @@ -231,50 +264,74 @@ def _init_weights(self, module):
         elif isinstance(module, RMSNorm):
             module.weight.data.fill_(1.0)
 
-    def forward(self, x, attention_mask=None):
+    def forward(self, x, attention_mask=None, kv_cache=None, start_pos=0):
         if self.lm_use_tokens:
-            x = self.token_embedding(x) # Only embed the inputs when using tokens
-        
-        B , T, _ = x.size()
+            x = self.token_embedding(x)
+
+        # T_curr is the length of the current input sequence
+        B, T_curr, _ = x.size()
         
-        # Note: You could also cache these input embeddings if you want to avoid recomputing them
-        position_ids = torch.arange(T, device=x.device).unsqueeze(0).expand(B, -1) # Create position ids [0, 1, 2, ..., seq_len-1]
-        cos, sin = self.rotary_embd(position_ids) # Get rotary position embeddings
+        # Create position_ids for the current sequence based on start_pos
+        current_position_ids = torch.arange(start_pos, start_pos + T_curr, device=x.device).unsqueeze(0).expand(B, -1)
+        cos, sin = self.rotary_embd(current_position_ids) # Get rotary position embeddings for current tokens
+
+        # Initialize new KV cache if none provided
+        if kv_cache is None:
+            kv_cache = [None] * len(self.blocks)
+
+        for i, block in enumerate(self.blocks):
+            x, kv_cache[i] = block(x, cos, sin, attention_mask, kv_cache[i])
 
-        for block in self.blocks:
-            x = block(x, cos, sin, attention_mask)
         x = self.norm(x)
 
-        if self.lm_use_tokens:
-            x = self.head(x) # Compute logits if we are using tokens, otherwise stay in the embedding space
+        # Compute logits if we are using tokens, otherwise stay in the embedding space
+        if self.lm_use_tokens: 
+            x = self.head(x) 
 
-        return x
+        return x, kv_cache
 
-    @torch.no_grad()
+
+    @torch.inference_mode()
path: models/language_model.py
user: login: kashif
body: torch.inference_mode() does a few more things: https://docs.pytorch.org/docs/stable/generated/torch.autograd.grad_mode.inference_mode.html
created_at: 2025-05-23T12:48:15Z
author_association: COLLABORATOR
diff_hunk: @@ -59,59 +59,81 @@ def forward(self, input_ids, image, attention_mask=None, targets=None):
 
         return logits, loss
 
-    @torch.no_grad()
-    def generate(self, input_ids, image, attention_mask=None, max_new_tokens=5, top_k=50, top_p=0.9, temperature=0.5, greedy=False):
-        # Process image through vision encoder and projection
-        image_embd = self.vision_encoder(image)
-        image_embd = self.MP(image_embd)
-        
-        # Embed initial tokens
-        token_embd = self.decoder.token_embedding(input_ids)
-        
-        # Concatenate image embeddings with token embeddings
-        combined_embd = torch.cat((image_embd, token_embd), dim=1)
+    @torch.inference_mode()
+    def generate(self, input_ids, image, max_new_tokens=5, top_k=50, top_p=0.9, temperature=0.5, greedy=False):
+        # input_ids: [B, T_prompt_text]
+        # image: [B, C, H, W]
+        # max_new_tokens: int, maximum number of tokens to generate
 
-        batch_size = image_embd.size(0)
-        img_seq_len = image_embd.size(1)
-        # Adjust attention mask to account for image tokens
-        if attention_mask is not None:
-            # Create mask of 1s for image tokens (all image tokens should be attended to)
-            image_attention_mask = torch.ones((batch_size, img_seq_len), device=attention_mask.device, dtype=attention_mask.dtype)
-            attention_mask = torch.cat((image_attention_mask, attention_mask), dim=1)
+        B = image.size(0)
+
+        # 1. Process image
+        image_embd = self.vision_encoder(image) # [B, T_img, D_model]
+        image_embd = self.MP(image_embd)      # [B, T_img, D_lm]
+
+        # 2. Embed initial text prompt tokens
+        prompt_token_embeds = self.decoder.token_embedding(input_ids) # [B, T_prompt_text, D_lm]
+
+        # 3. Combine image and text prompt embeddings for prefill
+        initial_combined_embeds = torch.cat((image_embd, prompt_token_embeds), dim=1) # [B, T_img + T_prompt_text, D_lm]
+        current_total_seq_len = initial_combined_embeds.size(1)
+        
+        # --- Multimodal Prefill Phase ---
+        prefill_output, kv_cache_list = self.decoder(
+            initial_combined_embeds,
+            attention_mask=None,
+            kv_cache=None,
+            start_pos=0
+        )
         
-        # Generate from combined embeddings using the decoder
-        # We need to use the decoder's forward function and not its generate method
-        # because we want to keep track of the image prefix
-        outputs = combined_embd
-        generated_tokens = torch.zeros((batch_size, max_new_tokens), device=input_ids.device, dtype=input_ids.dtype)
+        last_token_output_from_prefill = prefill_output[:, -1, :] 
         
-        #Note: Here you could implement improvements like e.g. KV caching
-        for i in range(max_new_tokens):
-            model_out = self.decoder(outputs, attention_mask)
+        if not self.decoder.lm_use_tokens:
+            current_logits = self.decoder.head(last_token_output_from_prefill) 
+        else:
+            current_logits = last_token_output_from_prefill 
+
+        # Store newly generated token IDs
+        newly_generated_ids_list = []
+
+        # --- Decode Phase with Sampling tokens one by one ---
path: models/vision_language_model.py
user: login: kashif
body: ```suggestion
        # --- Decode Phase by sampling tokens autoregressively using the kv-cache ---
```
created_at: 2025-05-26T08:00:41Z
author_association: COLLABORATOR

PR comments:
user: login: lusxvr
created_at: 2025-05-22T09:07:54Z
body: For how many tokens did you run the generate? If it was just a few we might not see the difference, try with a pretty big number, if the time still stays the same there is some issue
author_association: CONTRIBUTOR
user: login: ariG23498
created_at: 2025-05-22T09:39:16Z
body: @lusxvr I tried it with 1000 tokens.
author_association: COLLABORATOR
user: login: ariG23498
created_at: 2025-05-23T06:17:28Z
body: With the current implementation from @kashif we get 50% speedups!
```
Using device: cuda
<torch.utils.benchmark.utils.common.Measurement object at 0x7138c8315bb0>
generate_tokens(tokens, image)
setup: from __main__ import generate_tokens
  9.44 s
  1 measurement, 10 runs , 16 threads
```
author_association: COLLABORATOR
user: login: ariG23498
created_at: 2025-05-23T08:10:25Z
body: A note from @kashif 

We should aslo handle training logic as the `decoder.forward` now outputs two things instead of one.

After the initial review, I will quickly work on this.
author_association: COLLABORATOR
user: login: ariG23498
created_at: 2025-05-23T10:20:10Z
body: @gante after adding some print statements
```
print(f"Phase ->{torch.cuda.memory_allocated()/1024**2:0.2f} MB\n")
```
 I get the following

```shell
Using device: cuda
Loading weights from: lusxvr/nanoVLM-222M

Input:
  What is this? 

Outputs:
Before generation -> 871.11 MB
Before prefill phase -> 879.48 MB
After prefill phase-> 882.13 MB

Before next token decoding-> 882.69 MB
After next token decoding-> 882.73 MB

Before next token decoding-> 882.73 MB
After next token decoding-> 882.79 MB

Before next token decoding-> 882.79 MB
After next token decoding-> 882.82 MB

Before next token decoding-> 882.82 MB
After next token decoding-> 882.88 MB

Before next token decoding-> 882.87 MB
After next token decoding-> 882.91 MB

Before next token decoding-> 882.90 MB
After next token decoding-> 882.97 MB

Before next token decoding-> 882.96 MB
After next token decoding-> 883.00 MB

Before next token decoding-> 882.99 MB
After next token decoding-> 883.05 MB

Before next token decoding-> 883.05 MB
After next token decoding-> 883.08 MB

Before next token decoding-> 883.08 MB
After next token decoding-> 883.14 MB

Before next token decoding-> 883.14 MB
After next token decoding-> 883.17 MB

Before next token decoding-> 883.17 MB
After next token decoding-> 883.23 MB

Before next token decoding-> 883.23 MB
After next token decoding-> 883.26 MB

Before next token decoding-> 883.26 MB
After next token decoding-> 883.32 MB

Before next token decoding-> 883.32 MB
After next token decoding-> 883.35 MB

Before next token decoding-> 883.35 MB
After next token decoding-> 883.41 MB

Before next token decoding-> 883.41 MB
After next token decoding-> 883.44 MB

Before next token decoding-> 883.44 MB
After next token decoding-> 883.50 MB

Before next token decoding-> 883.50 MB
After next token decoding-> 883.53 MB

Before next token decoding-> 883.53 MB
After next token decoding-> 883.59 MB
```
This is in alignment to what you have said. Would this be enough to conclude that memory requirements are okay?
author_association: COLLABORATOR
user: login: lusxvr
created_at: 2025-05-23T10:56:06Z
body: > We should aslo handle training logic as the decoder.forward now outputs two things instead of one.

Indeed! We need to make sure we don't break the training pipeline when we merge this! I can help with testing the training  pipeline if you need.

But super great progress on this, very impressed @ariG23498 @kashif :)
author_association: CONTRIBUTOR
user: login: kashif
created_at: 2025-05-23T11:05:22Z
body: @lusxvr all good with regards to the training... the kv_cache output is now ignored
author_association: COLLABORATOR
user: login: lusxvr
created_at: 2025-05-23T13:00:02Z
body: > With the current implementation from @kashif we get 50% speedups!

How can I interpret the `benchmark-inference.py` results? When I try for myself on one H100 I get this
```
Using device: cuda
<torch.utils.benchmark.utils.common.Measurement object at 0x7f07c24e6360>
generate_tokens(tokens, image)
setup: from __main__ import generate_tokens
  18.71 s
  1 measurement, 10 runs , 12 threads
```
author_association: CONTRIBUTOR
user: login: kashif
created_at: 2025-05-23T13:07:46Z
body: then on `main`, change the `max_new_tokens=1000` and try for a comparison of the default

author_association: COLLABORATOR
user: login: lusxvr
created_at: 2025-05-23T13:09:19Z
body: And unfortunately when I launch a training run I get an error:
```
--- VLM Config ---
VLMConfig(vit_hidden_dim=768, vit_inter_dim=3072, vit_patch_size=16, vit_img_size=224, vit_n_heads=12, vit_dropout=0.0, vit_n_blocks=12, vit_ln_eps=1e-06, vit_cls_flag=False, vit_model_type='google/siglip-base-patch16-224', lm_hidden_dim=576, lm_inter_dim=1536, lm_rms_eps=1e-05, lm_re_base=100000, lm_max_position_embeddings=8192, lm_vocab_size=49152, lm_n_heads=9, lm_n_kv_heads=3, lm_dropout=0.0, lm_n_blocks=30, lm_attn_scaling=1.0, IMAGE_TOKEN_LENGTH=49, TOTAL_SEQUENCE_LENGTH=128, lm_max_length=79, lm_use_tokens=False, lm_tie_weights=True, lm_model_type='HuggingFaceTB/SmolLM2-135M', lm_tokenizer='HuggingFaceTB/cosmo2-tokenizer', lm_eos_token_id=0, mp_pixel_shuffle_factor=2, vlm_load_backbone_weights=True, vlm_checkpoint_path='checkpoints/nanoVLM-222M')
--- Train Config ---
TrainConfig(lr_mp=0.002, lr_backbones=0.0001, data_cutoff_idx=None, val_ratio=0.025, batch_size=256, gradient_accumulation_steps=1, mmstar_batch_size=32, eval_in_epochs=True, eval_interval=250, epochs=5, compile=False, resume_from_vlm_checkpoint=False, train_dataset_path='HuggingFaceM4/the_cauldron', train_dataset_name=('ai2d', 'aokvqa', 'chart2text', 'chartqa', 'clevr', 'cocoqa', 'datikz', 'diagram_image_to_text', 'docvqa', 'dvqa', 'figureqa', 'finqa', 'geomverse', 'hateful_memes', 'hitab', 'iam', 'iconqa', 'infographic_vqa', 'intergps', 'localized_narratives', 'mapqa', 'multihiertt', 'ocrvqa', 'plotqa', 'raven', 'rendered_text', 'robut_sqa', 'robut_wikisql', 'robut_wtq', 'scienceqa', 'screen2words', 'st_vqa', 'tabmwp', 'tallyqa', 'tat_qa', 'textcaps', 'textvqa', 'tqa', 'vistext', 'visual7w', 'visualmrc', 'vqarad', 'vqav2', 'vsr', 'websight'), test_dataset_path='Lin-Chen/MMStar', wandb_entity='HuggingFace', log_wandb=True)
Resolving data files: 100%|█████████████████████████████████████████████████████████████████████████████████████████| 24/24 [00:00<00:00, 470.52it/s]
Loading dataset shards: 100%|████████████████████████████████████████████████████████████████████████████████████████| 26/26 [00:01<00:00, 21.04it/s]
Resolving data files: 100%|█████████████████████████████████████████████████████████████████████████████████████████| 43/43 [00:00<00:00, 259.60it/s]
Loading dataset shards: 100%|████████████████████████████████████████████████████████████████████████████████████████| 44/44 [00:01<00:00, 22.46it/s]
Loading dataset shards: 100%|████████████████████████████████████████████████████████████████████████████████████████| 17/17 [00:00<00:00, 18.49it/s]
Resolving data files: 100%|██████████████████████████████████████████████████████████████████████████████████████| 23/23 [00:00<00:00, 236443.61it/s]
Loading dataset shards: 100%|████████████████████████████████████████████████████████████████████████████████████████| 20/20 [00:00<00:00, 39.05it/s]
Resolving data files: 100%|█████████████████████████████████████████████████████████████████████████████████████████| 22/22 [00:00<00:00, 577.23it/s]
Loading dataset shards: 100%|████████████████████████████████████████████████████████████████████████████████████████| 27/27 [00:00<00:00, 28.12it/s]
wandb: Currently logged in as: luis-wiedmann (huggingface) to https://api.wandb.ai. Use `wandb login --relogin` to force relogin
wandb: Tracking run with wandb version 0.19.11
wandb: Run data is saved locally in /fsx/luis_wiedmann/nanoVLM/wandb/run-20250523_130634-233m391e
wandb: Run `wandb offline` to turn off syncing.
wandb: Syncing run nanoVLM_1xGPU_1630251samples_bs256_ep5_lr0.0001-0.002_0523
wandb: ⭐️ View project at https://wandb.ai/huggingface/nanoVLM
wandb: 🚀 View run at https://wandb.ai/huggingface/nanoVLM/runs/233m391e
Loading from backbone weights
Successfully loaded google/siglip-base-patch16-224 weights from safetensors. Model has 85,797,120 parameters.
Successfully loaded HuggingFaceTB/SmolLM2-135M weights from safetensors. Model has 134,515,008 parameters.
nanoVLM initialized with 222,081,600 parameters
Training summary: 1630251 samples, 6368 batches/epoch, batch size 256
Validation summary: 41801 samples, 163 batches/epoch, batch size 256
Traceback (most recent call last):
  File "/fsx/luis_wiedmann/nanoVLM/train.py", line 432, in <module>
    main()
  File "/fsx/luis_wiedmann/nanoVLM/train.py", line 426, in main
    train(train_cfg, vlm_cfg)
  File "/fsx/luis_wiedmann/nanoVLM/train.py", line 332, in train
    epoch_accuracy = test_mmstar(eval_model, tokenizer, test_loader, device)
                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/fsx/luis_wiedmann/nanoVLM/train.py", line 171, in test_mmstar
    gen = model.generate(input_ids, image, attention_mask)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/fsx/luis_wiedmann/nanoVLM/.venv/lib/python3.12/site-packages/torch/utils/_contextlib.py", line 116, in decorate_context
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/fsx/luis_wiedmann/nanoVLM/models/vision_language_model.py", line 100, in generate
    for _ in range(max_new_tokens):
             ^^^^^^^^^^^^^^^^^^^^^
TypeError: only integer tensors of a single element can be converted to an index
```
author_association: CONTRIBUTOR
user: login: kashif
created_at: 2025-05-23T13:10:01Z
body: right let me fix the training
author_association: COLLABORATOR
user: login: lusxvr
created_at: 2025-05-23T13:10:42Z
body: > then on main, change the max_new_tokens=1000 and try for a comparison of the default

What `main` do you mean? The default is already at `max_new_tokens=1000`, I pulled the PR as it is in the current state
author_association: CONTRIBUTOR
user: login: kashif
created_at: 2025-05-23T13:14:50Z
body: i mean in the `main` branch if you run the benchmark, but change the max_new_tokens=1000, you should see slower times 
author_association: COLLABORATOR
user: login: kashif
created_at: 2025-05-23T13:18:33Z
body: can you kindly pull and test training
author_association: COLLABORATOR
user: login: lusxvr
created_at: 2025-05-23T14:18:53Z
body: > i mean in the main branch if you run the benchmark, but change the max_new_tokens=1000, you should see slower times

Ah, yes makes sense. I see a ~20% improvement from my side
author_association: CONTRIBUTOR
user: login: gante
created_at: 2025-05-23T14:23:55Z
body: @ariG23498 

> This is in alignment to what you have said. Would this be enough to conclude that memory requirements are okay?

yes!
author_association: NONE
user: login: lusxvr
created_at: 2025-05-23T15:03:17Z
body: <img width="1166" alt="image" src="https://github.com/user-attachments/assets/03b5b5bb-30c1-4503-be43-65b1ad57b3e5" />

Training and val loss looks great. Accuracy is noticeably lower though (evaluation is generation based), so this might either be a bug in the code or due to @andimarafioti recently merged greedy decoding. I launched a new run that takes this into account and will keep you updated
author_association: CONTRIBUTOR
user: login: lusxvr
created_at: 2025-05-23T15:52:06Z
body: <img width="1075" alt="image" src="https://github.com/user-attachments/assets/4c4db162-25a1-4583-ab38-d062bbb3eb95" />

Okay it looks like there is still a bit of a discrepancy between the generation with the new KV-Cache. I also re-ran the experiments, the difference in accuracy stayed, so it was not due to some randomness. But interesstingly even though the performance was worse, the KV-cache implementation was more consistent in the accuracy.

Red: no KV-cache, no greedy decoding (Baseline)
Blue: KV-cache, greedy decoding
Turquoise: no KV-cache, greedy decoding

(Loss for blue and turquoise are basically identical, this is why we don't really see two distinct lines there)

The evaluation is generation based, so it's the only metric we have in this case, but it is also not a very good metric to start with. If you believe the KV implementation is correct then I am almost inclined to say the difference in accuracy is due to it being a non-optimal metric / implementation in our case and would be fine with merging it. We already started giving it a bit less importance while thinking of a way to do better evaluation.

We talked a lot about this already, whats your opinion on this @andimarafioti. Should we take this accuracy difference into account or not worry about it too much since it is quite an unstable metric for us anyway?
author_association: CONTRIBUTOR
user: login: kashif
created_at: 2025-05-23T16:10:19Z
body: @lusxvr let me also check I was not expecting the ACC to go down with kv-cache
author_association: COLLABORATOR
user: login: arpitsinghgautam
created_at: 2025-05-23T18:40:09Z
body: Any experiments I can do to help here!
author_association: NONE
user: login: kashif
created_at: 2025-05-26T07:58:13Z
body: @lusxvr I think we can merge this and i can look at the metrics for training after that? I get the same generation output when I use the same seed with and without kv-caching so there is no issue with this PR as far as I can tell
author_association: COLLABORATOR
user: login: andimarafioti
created_at: 2025-05-26T09:41:58Z
body: @lusxvr are you sure you were doing greedy decoding? The default in train.py now is to not use greedy decoding.
author_association: MEMBER
user: login: lusxvr
created_at: 2025-05-26T09:46:29Z
body: > are you sure you were doing greedy decoding? The default in train.py now is to not use greedy decoding.

great catch, I did not change the default (somehow I just assumed it would be enabled)
author_association: CONTRIBUTOR
user: login: lusxvr
created_at: 2025-05-26T09:46:56Z
body: @kashif sounds good to me, do you have anything against me squashing the commits when I merge?
author_association: CONTRIBUTOR
user: login: kashif
created_at: 2025-05-26T10:08:47Z
body: squash away! I love squash
author_association: COLLABORATOR
user: login: ariG23498
created_at: 2025-05-26T12:48:19Z
body: Hi @andimarafioti 

Thanks for the tests! I think they might have uncovered something severe. The norms used in the language models are causing the problem and not the KV Cache in place.

If you run the following, it will fail with the current implementation. My idea here is to test the model's outputs if you provide the entire sequence and the input sequence with the last input token left out.
```py
import torch
import unittest
from models.language_model import LanguageModel
from types import SimpleNamespace

class TestLanguageModel(unittest.TestCase):
    def setUp(self):
        # Minimal config for testing
        self.cfg = SimpleNamespace(
            lm_hidden_dim=64,
            lm_inter_dim=128,
            lm_rms_eps=1e-5,
            lm_re_base=10000.0,
            lm_max_position_embeddings=512,
            lm_attn_scaling=1.0,
            lm_vocab_size=100, # Small vocab for testing
            lm_n_heads=4,
            lm_n_kv_heads=2,
            lm_dropout=0.0,
            lm_n_blocks=2,
            lm_use_tokens=True,
            lm_tie_weights=True
        )
        self.model = LanguageModel(self.cfg)
        self.model.eval() # Set model to evaluation mode

    def test_kv_caching_consistency(self):
        # Input for the model
        batch_size = 16
        seq_len = 1000
        input_ids = torch.randint(0, self.cfg.lm_vocab_size, (batch_size, seq_len))

        # Forward pass without KV caching (prefill)
        output_no_cache1, _ = self.model(input_ids, start_pos=0)
        output_no_cache2, _ = self.model(input_ids, start_pos=0)

        self.assertTrue(
            torch.allclose(output_no_cache1, output_no_cache2, atol=1e-5),
            "Outputs do not match"
        )

        output_no_cache1, _ = self.model(input_ids, start_pos=0)
        output_no_cache2, _ = self.model(input_ids[:, :-1], start_pos=0)

        self.assertTrue(
            torch.allclose(output_no_cache1[:, :-1, :], output_no_cache2, atol=1e-5),
            "Outputs do not match"
        )

if __name__ == '__main__':
    unittest.main() 
```

Doing a similar test on a language model from transformers pass
```py
from transformers import AutoTokenizer, AutoModel

tok = AutoTokenizer.from_pretrained("meta-llama/Llama-3.2-1B")
model = AutoModel.from_pretrained("meta-llama/Llama-3.2-1B")

inputs = tok("Hello, my dog is ", return_tensors="pt")
out1 = model(inputs["input_ids"]).last_hidden_state
out2 = model(inputs["input_ids"][:, :-1]).last_hidden_state
torch.testing.assert_close(
    out1[:, :-1, :],
    out2,
)
```

If you comment out all the norm calls in the language model and run the tests that you have written, the tests pass. We might need to create another issue for this and figure out the problem (my bet is on the norms).
All the norm calls are:
1. https://github.com/huggingface/nanoVLM/blob/1ce10cf53ac23af8693a462e7ae83efe9a8211ef/models/language_model.py#L246
2. https://github.com/huggingface/nanoVLM/blob/1ce10cf53ac23af8693a462e7ae83efe9a8211ef/models/language_model.py#L193
3. https://github.com/huggingface/nanoVLM/blob/1ce10cf53ac23af8693a462e7ae83efe9a8211ef/models/language_model.py#L198

CC: @kashif 

author_association: COLLABORATOR
user: login: andimarafioti
created_at: 2025-05-26T13:18:44Z
body: Hi, those norms follow quite closely the implementation in transformers' `modeling_llama.py` file. I just checked. I'll see if I can do a similar experiment with a transformers model. 
author_association: MEMBER
user: login: andimarafioti
created_at: 2025-05-26T14:42:31Z
body: FYI, I tried the same with transformers and the versions with and without cache are identical:

```python
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
import warnings

# Suppress a specific UserWarning from transformers
warnings.filterwarnings("ignore", message="The QuietFreezeWrapper class will be deprecated")
warnings.filterwarnings("ignore", message="Passing the following arguments to `inputs` is deprecated")


def test_kv_caching_consistency_hf(model_name="meta-llama/Llama-2-7b-chat-hf", device="cuda:0"):
    """
    Tests KV caching consistency for a Hugging Face transformer model.
    1. Compares last token logits: (full pass) vs (prefill + 1 decode step).
    2. Compares full sequence logits: (full pass) vs (token-by-token decode).
    """
    print(f"Loading tokenizer and model: {model_name}...")
    # It's good practice to use trust_remote_code=True for some models,
    # though for standard Llama it might not be strictly necessary.
    # Adding it preemptively.
    tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
    model = AutoModelForCausalLM.from_pretrained(
        model_name,
        torch_dtype=torch.bfloat16,
        trust_remote_code=True
    ).to(device)
    model.eval()

    # Test parameters
    batch_size = 16 # Keep small to avoid OOM, can be increased
    seq_len = 320   # Keep small for faster testing, can be increased
    vocab_size = tokenizer.vocab_size

    print(f"Generating random input_ids (batch_size={batch_size}, seq_len={seq_len}, vocab_size={vocab_size})...")
    input_ids = torch.randint(0, vocab_size, (batch_size, seq_len), device=model.device)
    
    print("\n--- Test 1: Last Token Consistency ---")
    # Forward pass without explicit KV caching reuse (but cache is computed)
    print("Running full forward pass (reference)...")
    with torch.no_grad():
        outputs_no_cache = model(input_ids, use_cache=True) # use_cache=True to get past_key_values if needed later
        logits_no_cache = outputs_no_cache.logits

    # Forward pass with KV caching
    # 1. Prefill phase
    print("Running prefill phase (all but last token)...")
    prefill_input_ids = input_ids[:, :-1]
    with torch.no_grad():
        outputs_prefill = model(prefill_input_ids, use_cache=True)
        kv_cache_prefill = outputs_prefill.past_key_values

    # 2. Decode phase (last token)
    print("Running decode phase (last token using KV cache)...")
    last_token_input_ids = input_ids[:, -1].unsqueeze(-1) # Shape: [B, 1]
    with torch.no_grad():
        outputs_with_cache_last_token = model(
            last_token_input_ids,
            past_key_values=kv_cache_prefill,
            use_cache=True
        )
        logits_with_cache_last_token = outputs_with_cache_last_token.logits

    # Compare the logits for the last token
    logits_no_cache_last_token = logits_no_cache[:, -1, :]
    # logits_with_cache_last_token is [B, 1, vocab_size], so squeeze or select
    logits_with_cache_last_token_squeezed = logits_with_cache_last_token[:, 0, :]

    atol = 1e-6
    last_token_match = torch.allclose(logits_no_cache_last_token, logits_with_cache_last_token_squeezed, atol=atol)
    print(f"Logits for the last token match: {last_token_match}")
    if not last_token_match:
        print("Max difference (last token):", (logits_no_cache_last_token - logits_with_cache_last_token_squeezed).abs().max())


    print("\n--- Test 2: Full Sequence Consistency (Token-by-Token) ---")
    # Reset for a full token-by-token generation using KV cache
    print("Running token-by-token generation with KV cache...")
    output_logits_list_step_by_step = []
    current_kv_cache = None

    for i in range(seq_len):
        current_input_id_step = input_ids[:, i:i+1] # Current token
        with torch.no_grad():
            outputs_step = model(
                current_input_id_step,
                past_key_values=current_kv_cache,
                use_cache=True
            )
            # outputs_step.logits is [B, 1, vocab_size]
            output_logits_list_step_by_step.append(outputs_step.logits)
            current_kv_cache = outputs_step.past_key_values
            if i % 5 == 0 or i == seq_len -1 :
                print(f"  Processed token {i+1}/{seq_len}")

    # Concatenate all single token logits
    logits_with_cache_full_sequence = torch.cat(output_logits_list_step_by_step, dim=1) # [B, seq_len, vocab_size]

    full_sequence_match = torch.allclose(logits_no_cache, logits_with_cache_full_sequence, atol=atol)
    print(f"Full sequence logits match (token-by-token vs full pass): {full_sequence_match}")
    if not full_sequence_match:
        print("Max difference (full sequence):", (logits_no_cache - logits_with_cache_full_sequence).abs().max())

    print("\nTest completed.")

if __name__ == '__main__':
    model_to_test = "HuggingFaceTB/SmolLM2-360M-Instruct"
    selected_device = "cuda:0" if torch.cuda.is_available() else "cpu"
    print(f"Using device: {selected_device}")
    
    test_kv_caching_consistency_hf(model_name=model_to_test, device=selected_device)

```
author_association: MEMBER
user: login: andimarafioti
created_at: 2025-05-26T15:00:49Z
body: Looking further, the norms in transformers do convert the values to float32 to process. I applied that change and it didn't change the results. I agree that your test @ariG23498 should also work. The reason why it does without norms is that the values are in a different range and then the absolute tolerance you accept is sufficient. 
author_association: MEMBER
user: login: andimarafioti
created_at: 2025-05-26T15:12:20Z
body: Ok, I checked Aritra's test on main and it also fails ! 
author_association: MEMBER
user: login: andimarafioti
created_at: 2025-05-26T15:25:23Z
body: Figured out why Aritra's test was failing. If `seq_len > lm_max_position_embeddings`, then the RoPE scale is computed on the seq_len. Because they changed between the two examples, the values are different. Setting lm_max_position_embeddings to 1024 fixes it. 
author_association: MEMBER
user: login: andimarafioti
created_at: 2025-05-26T15:34:47Z
body: Some more info: 
- Loading backbones the test still fails, even more than with the small test
- Setting both times to use or not use the cache gives us deterministic results, so the difference is between having the cache or not. 
- passing`start_pos=current_token_start_pos` instead of zero when not using kv_cache still fails 
author_association: MEMBER
user: login: andimarafioti
created_at: 2025-05-26T21:51:07Z
body: I can't figure out why, but I'm also getting the same issue that @lusxvr had. Basically, training on main or on this branch gives very different accuracies
![Screenshot 2025-05-26 at 23 50 30](https://github.com/user-attachments/assets/b0ac69b8-3bf3-41cc-aedf-90a49682344d)

author_association: MEMBER
user: login: andimarafioti
created_at: 2025-05-26T21:52:59Z
body: I created a script to test MMStar directly on a pretrained model. And evaluating it on main or on this branch gives different results. I changed the test_mmstar to use greedy decoding for my tests for reproducibility. On this branch, with or without kv-cache it gives me the same (worse) results. Here is my script:

```python
import argparse
import torch
import os
import numpy
import random
from datasets import load_dataset
from torch.utils.data import DataLoader

# Imports from the project
from models.vision_language_model import VisionLanguageModel
from data.processors import get_tokenizer, get_image_processor
from data.datasets import MMStarDataset
from data.collators import MMStarCollator
from train import test_mmstar # Using the function from train.py
# models.utils is implicitly used by test_mmstar

# Set TOKENIZERS_PARALLELISM to avoid warnings
os.environ["TOKENIZERS_PARALLELISM"] = "false"

# Seed for reproducibility
torch.manual_seed(0)
if torch.cuda.is_available():
    torch.cuda.manual_seed_all(0)
numpy.random.seed(0)
random.seed(0)

def seed_worker(worker_id):
    worker_seed = torch.initial_seed() % 2**32
    numpy.random.seed(worker_seed)
    random.seed(worker_seed)

def parse_args():
    parser = argparse.ArgumentParser(description="Evaluate a nanoVLM model on the MMStar benchmark, similar to train.py.")
    parser.add_argument(
        "--checkpoint", type=str, default=None,
        help="Path to a local checkpoint (directory or safetensors/pth). If omitted, --hf_model is used."
    )
    parser.add_argument(
        "--hf_model", type=str, default="lusxvr/nanoVLM-222M",
        help="HuggingFace repo ID to download from if --checkpoint isn't set."
    )
    parser.add_argument(
        "--dataset_path", type=str, default="Lin-Chen/MMStar", # Matches typical train_cfg.test_dataset_path
        help="Path or HuggingFace dataset ID for the MMStar test dataset. Script will use the 'val' split."
    )
    parser.add_argument(
        "--mmstar_batch_size", type=int, default=256,
        help="Batch size for MMStar evaluation."
    )
    parser.add_argument(
        "--num_workers", type=int, default=4,
        help="Number of workers for the DataLoader."
    )
    return parser.parse_args()

def main():
    args = parse_args()

    # Setup device
    if torch.cuda.is_available():
        device = torch.device("cuda")
    elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
        device = torch.device("mps")
        if device.type == "mps": # Ensure mps fallback if needed
             torch.backends.mps.enable_fallback_to_cpu = True
    else:
        device = torch.device("cpu")
    print(f"Using device: {device}")

    # Load model
    source = args.checkpoint if args.checkpoint else args.hf_model
    print(f"Loading model weights from: {source}")
    model = VisionLanguageModel.from_pretrained(source).to(device)
    model.eval() # Set model to evaluation mode

    # Get tokenizer and image processor from model's config
    tokenizer = get_tokenizer(model.cfg.lm_tokenizer)
    image_processor = get_image_processor(model.cfg.vit_img_size)

    # Load test dataset (using 'val' split, similar to train.py for MMStar)
    print(f"Loading dataset: {args.dataset_path}, attempting to use 'val' split.")
    try:
        dataset_dict = load_dataset(args.dataset_path)
        if 'val' not in dataset_dict:
            raise KeyError(f"The 'val' split was not found in dataset {args.dataset_path}. Available splits: {list(dataset_dict.keys())}")
        test_ds_split = dataset_dict['val']
    except Exception as e:
        print(f"Error loading dataset or finding 'val' split: {e}")
        return
        
    test_dataset = MMStarDataset(test_ds_split, tokenizer, image_processor)

    # Create MMStar collator
    mmstar_collator = MMStarCollator(tokenizer)

    # Create DataLoader
    g = torch.Generator()
    g.manual_seed(0) # Seed for DataLoader

    test_loader = DataLoader(
        test_dataset,
        batch_size=args.mmstar_batch_size,
        shuffle=False, # Evaluation is typically not shuffled
        collate_fn=mmstar_collator,
        num_workers=args.num_workers,
        pin_memory=True, # Matches train.py test_loader configuration
        worker_init_fn=seed_worker if args.num_workers > 0 else None,
        generator=g
    )

    # Run MMStar evaluation
    print("Starting MMStar evaluation...")
    accuracy = test_mmstar(model, tokenizer, test_loader, device)

    print(f"MMStar Test Accuracy on 'val' split of {args.dataset_path}: {accuracy:.4f}")

if __name__ == "__main__":
    main() 
 ```

I run it with:
```bash
python evaluate.py --hf_model andito/nanoVLM
```
author_association: MEMBER
user: login: ariG23498
created_at: 2025-05-27T04:08:08Z
body: Thanks for the script @andimarafioti 
Some more info:
1. I ran a one to one check of the inputs, and they were the same
2. Iterating on each batch, I found out that some batches have lower hits, and I am currently looking into them
3. I also tried using `use_kv_cache=False` with the `kv-cache` branch, and it resulted in an OOM (which signifies that there is more memory being used). This intuitively makes sense because with `use_kv_cache=False` we are still creating the cache, but also not using it. This needs to be looked into
4. Using `sdpa=False` in `main` branch works, while `sdpa=False` with `kv-cache` branch leads to OOM

author_association: COLLABORATOR
user: login: ddiddi
created_at: 2025-05-27T05:34:29Z
body: LGTM! Appreciate the efforts on this one. 
author_association: NONE
user: login: lusxvr
created_at: 2025-05-27T14:41:54Z
body: Crazy effort by you guys, kudos for pushing through! 🫶
author_association: CONTRIBUTOR