Coverage for transformer_lens/benchmarks/generation.py: 45%
64 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-21 19:27 +0000
1"""Generation and KV cache benchmarks for TransformerBridge."""
3from typing import Any
5from transformer_lens.benchmarks.utils import (
6 BenchmarkResult,
7 BenchmarkSeverity,
8 deterministic_rng,
9 is_tiny_test_model,
10)
11from transformer_lens.model_bridge import TransformerBridge
14def resolve_text_generator(bridge: Any):
15 """Return the callable that produces text for this architecture, or None
16 (diffusion LMs delegate to their native sampler)."""
17 if getattr(bridge.adapter, "supports_generation", True):
18 return bridge.generate
19 if getattr(bridge.adapter, "native_sampler", None):
20 return bridge.diffusion_generate
21 return None
24def benchmark_generation(
25 bridge: TransformerBridge,
26 test_text: str,
27 max_new_tokens: int = 10,
28) -> BenchmarkResult:
29 """Benchmark basic text generation."""
30 try:
31 if is_tiny_test_model(getattr(bridge.cfg, "model_name", "") or ""): 31 ↛ 32line 31 didn't jump to line 32 because the condition on line 31 was never true
32 return BenchmarkResult(
33 name="generation",
34 severity=BenchmarkSeverity.INFO,
35 message="Skipped for tiny/test model (random weights produce degenerate generation)",
36 )
37 generator = resolve_text_generator(bridge)
38 if generator is None: 38 ↛ 39line 38 didn't jump to line 39 because the condition on line 38 was never true
39 return BenchmarkResult(
40 name="generation",
41 severity=BenchmarkSeverity.INFO,
42 message="Skipped: architecture supports no text generation",
43 )
44 # Greedy (deterministic, no seeding): tests the loop's mechanics, not
45 # sampling quality. stop_at_eos=False because some models argmax EOS on a
46 # bare prompt (EXAONE-4, raw HF) — a choice to stop, not a stall. BOS is
47 # left alone: forcing it would override default_prepend_bos=False adapters
48 # and derail checkpoints whose BOS token is their EOS.
49 gen_kwargs: dict[str, Any] = {"max_new_tokens": max_new_tokens, "temperature": 0.0}
50 if getattr(bridge.adapter, "supports_generation", True): 50 ↛ 53line 50 didn't jump to line 53 because the condition on line 50 was always true
51 # Native diffusion samplers take neither kwarg through **kwargs.
52 gen_kwargs["stop_at_eos"] = False
53 output = generator(test_text, **gen_kwargs)
55 if not isinstance(output, str): 55 ↛ 56line 55 didn't jump to line 56 because the condition on line 55 was never true
56 return BenchmarkResult(
57 name="generation",
58 severity=BenchmarkSeverity.DANGER,
59 message="Generated output is not a string",
60 passed=False,
61 )
63 # Check token count instead of character count to handle whitespace-only generation
64 input_tokens = bridge.to_tokens(test_text)
65 output_tokens = bridge.to_tokens(output)
67 # Strip leading BOS token if present for fair comparison
68 input_len = input_tokens.shape[-1]
69 output_len = output_tokens.shape[-1]
71 # Compare like with like. decode->re-tokenize is lossy when the prompt
72 # holds out-of-vocabulary characters (a DNA model given English drops
73 # them as [UNK]), which reads as "no new tokens" though generation ran.
74 prompt_roundtrip = test_text
75 if bridge.tokenizer is not None: 75 ↛ 77line 75 didn't jump to line 77 because the condition on line 75 was always true
76 prompt_roundtrip = bridge.tokenizer.decode(input_tokens[0], skip_special_tokens=True)
77 if output_len <= input_len and len(output) <= len(prompt_roundtrip):
78 return BenchmarkResult(
79 name="generation",
80 severity=BenchmarkSeverity.DANGER,
81 message="Generated text has no new tokens",
82 details={
83 "input_tokens": input_len,
84 "output_tokens": output_len,
85 "input_chars": len(test_text),
86 "output_chars": len(output),
87 },
88 passed=False,
89 )
91 return BenchmarkResult(
92 name="generation",
93 severity=BenchmarkSeverity.INFO,
94 message=f"Generation successful: {input_len} -> {output_len} tokens ({len(test_text)} -> {len(output)} chars)",
95 details={
96 "input_tokens": input_len,
97 "output_tokens": output_len,
98 "input_chars": len(test_text),
99 "output_chars": len(output),
100 "max_new_tokens": max_new_tokens,
101 },
102 )
104 except Exception as e:
105 return BenchmarkResult(
106 name="generation",
107 severity=BenchmarkSeverity.ERROR,
108 message=f"Generation failed: {str(e)}",
109 passed=False,
110 )
113def benchmark_generation_with_kv_cache(
114 bridge: TransformerBridge,
115 test_text: str,
116 max_new_tokens: int = 10,
117) -> BenchmarkResult:
118 """Benchmark text generation with KV caching enabled.
120 This ensures that the KV cache is properly passed through attention layers
121 during generation, and that the cache update logic works correctly.
123 Args:
124 bridge: TransformerBridge model to test
125 test_text: Input text for generation
126 max_new_tokens: Number of tokens to generate
128 Returns:
129 BenchmarkResult with generation details
130 """
131 try:
132 if is_tiny_test_model(getattr(bridge.cfg, "model_name", "") or ""):
133 return BenchmarkResult(
134 name="generation_with_kv_cache",
135 severity=BenchmarkSeverity.INFO,
136 message="Skipped for tiny/test model (random weights produce degenerate generation)",
137 )
139 # Cache-free architectures (RWKV, HyenaDNA) would "pass" this vacuously —
140 # the recompute path produces output while exercising no cache at all.
141 # Diffusion samplers have no KV cache concept either.
142 if not getattr(bridge.adapter, "supports_kv_cache", True) or not getattr(
143 bridge.adapter, "supports_generation", True
144 ):
145 return BenchmarkResult(
146 name="generation_with_kv_cache",
147 severity=BenchmarkSeverity.INFO,
148 message="Skipped: architecture has no KV cache (generation recomputes each step)",
149 )
151 # Generate with KV cache (should be enabled by default for max_new_tokens > 1)
152 with deterministic_rng():
153 output = bridge.generate(
154 test_text,
155 max_new_tokens=max_new_tokens,
156 temperature=0.7,
157 prepend_bos=True,
158 )
160 if output is None or len(output) == 0:
161 return BenchmarkResult(
162 name="generation_with_kv_cache",
163 severity=BenchmarkSeverity.DANGER,
164 message="Generation with KV cache produced no output",
165 passed=False,
166 )
168 return BenchmarkResult(
169 name="generation_with_kv_cache",
170 severity=BenchmarkSeverity.INFO,
171 message=f"KV cache generation successful ({len(output)} chars)",
172 details={"output_len": len(output), "max_new_tokens": max_new_tokens},
173 )
175 except Exception as e:
176 return BenchmarkResult(
177 name="generation_with_kv_cache",
178 severity=BenchmarkSeverity.ERROR,
179 message=f"KV cache generation failed: {str(e)}",
180 passed=False,
181 )
184def benchmark_multiple_generation_calls(
185 bridge: TransformerBridge,
186 test_prompts: list,
187 max_new_tokens: int = 5,
188) -> BenchmarkResult:
189 """Benchmark multiple generation calls to ensure KV cache handling is robust.
191 Args:
192 bridge: TransformerBridge model to test
193 test_prompts: List of input prompts for generation
194 max_new_tokens: Number of tokens to generate per prompt
196 Returns:
197 BenchmarkResult with multiple generation details
198 """
199 try:
200 if is_tiny_test_model(getattr(bridge.cfg, "model_name", "") or ""):
201 return BenchmarkResult(
202 name="multiple_generation_calls",
203 severity=BenchmarkSeverity.INFO,
204 message="Skipped for tiny/test model (random weights produce degenerate generation)",
205 )
207 generator = resolve_text_generator(bridge)
208 if generator is None:
209 return BenchmarkResult(
210 name="multiple_generation_calls",
211 severity=BenchmarkSeverity.INFO,
212 message="Skipped: architecture supports no text generation",
213 )
215 outputs = []
216 with deterministic_rng():
217 for prompt in test_prompts:
218 output = generator(
219 prompt,
220 max_new_tokens=max_new_tokens,
221 temperature=0.7,
222 prepend_bos=True,
223 )
224 if output is None or len(output) == 0:
225 return BenchmarkResult(
226 name="multiple_generation_calls",
227 severity=BenchmarkSeverity.DANGER,
228 message=f"Generation failed for prompt: {prompt[:50]}...",
229 passed=False,
230 )
231 outputs.append(output)
233 return BenchmarkResult(
234 name="multiple_generation_calls",
235 severity=BenchmarkSeverity.INFO,
236 message=f"All {len(test_prompts)} generation calls successful",
237 details={
238 "prompt_count": len(test_prompts),
239 "max_new_tokens": max_new_tokens,
240 "output_lens": [len(out) for out in outputs],
241 },
242 )
244 except Exception as e:
245 return BenchmarkResult(
246 name="multiple_generation_calls",
247 severity=BenchmarkSeverity.ERROR,
248 message=f"Multiple generation calls failed: {str(e)}",
249 passed=False,
250 )