e2e generation test: mirror generate() cache shape logic

hf_adapters/hf_common.py

@@ -1324,6 +1324,13 @@ def _resolve_generation_params(model, tokenizer, overrides):
     }
 
 
+def generation_cache_len(prompt_length, max_new_tokens):
+    """Return KV-cache capacity for block-padded prompt and generation tokens."""
+    padded_prompt_len = math.ceil(prompt_length / BLOCK_SIZE) * BLOCK_SIZE
+    padded_generation_len = math.ceil(max_new_tokens / BLOCK_SIZE) * BLOCK_SIZE
+    return padded_prompt_len + padded_generation_len
+
+
 def pad_and_position(input_ids, actual_lengths):
     """Left block-pad ``input_ids`` to a BLOCK_SIZE multiple and build positions.
 
@@ -1490,10 +1497,7 @@ def generate(
 
     # Block-pad to a BLOCK_SIZE multiple; real tokens right-aligned (positions
     # 0..actual_len-1 at padded indices prompt_offsets[b]..padded_len-1).
-    max_cache_len = (
-        math.ceil(prompt_length / BLOCK_SIZE) * BLOCK_SIZE
-        + math.ceil(max_new_tokens / BLOCK_SIZE) * BLOCK_SIZE
-    )
+    max_cache_len = generation_cache_len(prompt_length, max_new_tokens)
     input_ids, padded_len, prompt_offsets, position_ids = pad_and_position(
         input_ids, actual_prompt_lengths
     )

hf_adapters/hf_gemma4_mm.py

@@ -76,8 +76,6 @@
 (``prepare_for_spyre`` / forward raise if audio/video inputs are present).
 """
 
-import math
-
 import torch
 
 from hf_adapters import hf_gemma4
@@ -89,6 +87,7 @@
     build_expansion_mask,
     build_prefill_mask,
     decode_block_walk,
+    generation_cache_len,
     get_backbone,
     get_model_dtype,
     pad_and_position,
@@ -536,10 +535,7 @@ def generate(
     batch_size, prompt_length = input_ids.shape
     actual_prompt_lengths = attention_mask.sum(dim=1)  # [B]
 
-    max_cache_len = (
-        math.ceil(prompt_length / BLOCK_SIZE) * BLOCK_SIZE
-        + math.ceil(max_new_tokens / BLOCK_SIZE) * BLOCK_SIZE
-    )
+    max_cache_len = generation_cache_len(prompt_length, max_new_tokens)
     input_ids, padded_len, prompt_offsets, position_ids = pad_and_position(
         input_ids, actual_prompt_lengths
     )

hf_adapters/hf_granite_vision_mm.py

@@ -54,8 +54,6 @@
 logits cosine ≥ 0.999, argmax match) and stock ``model.generate`` (token-exact).
 """
 
-import math
-
 import torch
 
 from hf_adapters import hf_siglip_vision
@@ -67,6 +65,7 @@
     build_expansion_mask,
     build_prefill_mask,
     decode_block_walk,
+    generation_cache_len,
     get_backbone,
     get_model_dtype,
     make_standard_gqa_block,
@@ -453,10 +452,7 @@ def generate(
     batch_size, prompt_length = input_ids.shape
     actual_prompt_lengths = attention_mask.sum(dim=1)  # [B]
 
-    max_cache_len = (
-        math.ceil(prompt_length / BLOCK_SIZE) * BLOCK_SIZE
-        + math.ceil(max_new_tokens / BLOCK_SIZE) * BLOCK_SIZE
-    )
+    max_cache_len = generation_cache_len(prompt_length, max_new_tokens)
     input_ids, padded_len, prompt_offsets, position_ids = pad_and_position(
         input_ids, actual_prompt_lengths
     )

hf_adapters/hf_mistral3_vision_mm.py

@@ -68,8 +68,6 @@
 for both ``mistral`` and ``ministral3`` text-backbone variants.
 """
 
-import math
-
 import torch
 
 from hf_adapters import hf_pixtral_vision
@@ -81,6 +79,7 @@
     build_expansion_mask,
     build_prefill_mask,
     decode_block_walk,
+    generation_cache_len,
     get_backbone,
     get_model_dtype,
     make_standard_gqa_block,
@@ -456,10 +455,7 @@ def generate(
     batch_size, prompt_length = input_ids.shape
     actual_prompt_lengths = attention_mask.sum(dim=1)  # [B]
 
-    max_cache_len = (
-        math.ceil(prompt_length / BLOCK_SIZE) * BLOCK_SIZE
-        + math.ceil(max_new_tokens / BLOCK_SIZE) * BLOCK_SIZE
-    )
+    max_cache_len = generation_cache_len(prompt_length, max_new_tokens)
     input_ids, padded_len, prompt_offsets, position_ids = pad_and_position(
         input_ids, actual_prompt_lengths
     )

tests/spyre/test_e2e_token_compare_spyre.py

@@ -37,6 +37,7 @@
 from hf_adapters.hf_common import (
     BLOCK_SIZE,
     DEVICE,
+    generation_cache_len,
     get_model_dtype,
     move_model_to_spyre,
 )
@@ -114,10 +115,7 @@ def adapter_greedy_steps(
     position_ids = torch.zeros((batch_size, padded_len), dtype=torch.long)
     position_ids[:, prompt_offset:] = torch.arange(seq_len)
 
-    max_cache_len = (
-        padded_len + math.ceil(num_decode / BLOCK_SIZE) * BLOCK_SIZE + BLOCK_SIZE
-    )
-
+    max_cache_len = generation_cache_len(seq_len, num_decode + 1)
     dtype = get_model_dtype(model)
 
     key_caches, value_caches = allocate_kv_caches(

tests/spyre/test_vlm_e2e_spyre.py

@@ -55,7 +55,6 @@
 """
 
 import gc
-import math
 import types
 from typing import Any
 
@@ -81,6 +80,7 @@
     DEVICE,
     allocate_kv_caches,
     build_expansion_mask,
+    generation_cache_len,
     get_model_dtype,
     pad_and_position,
 )
@@ -154,10 +154,7 @@ def _adapter_teacher_forced_steps(
     actual_prompt_lengths = attention_mask.sum(dim=1)
     n_steps = len(forced_tokens)
 
-    max_cache_len = (
-        math.ceil(prompt_length / BLOCK_SIZE) * BLOCK_SIZE
-        + math.ceil((n_steps + 1) / BLOCK_SIZE) * BLOCK_SIZE
-    )
+    max_cache_len = generation_cache_len(prompt_length, n_steps)
     padded_ids, padded_len, prompt_offsets, position_ids = pad_and_position(
         input_ids, actual_prompt_lengths
     )