atharv6f/flash-attention-explorer
0
1"""2Prefill vs Decode phase comparison module.3 4Demonstrates the key difference between:5- Prefill: Process entire prompt in parallel (N² attention complexity)6- Decode: Generate one token at a time (N attention per token, but sequential)7 8Uses REAL HuggingFace model attention layers for accurate benchmarking.9"""10 11import torch12import torch.nn.functional as F13import numpy as np14import plotly.graph_objects as go15from plotly.subplots import make_subplots16 17from .constants import MODEL_CONFIGS, ATTENTION_BACKENDS18from .models import load_model19from .attention_utils import (20 extract_attention_layer,21 create_attention_inputs,22 benchmark_attention_layer,23 get_model_attention_info,24)25 26 27def get_real_model_config(model_name: str) -> dict:28 """29 Load model and extract ACTUAL config values from model.config.30 31 This function ensures we use real model architecture values,32 NOT hardcoded constants from MODEL_CONFIGS.33 34 Args:35 model_name: Key from MODEL_CONFIGS (e.g., "SmolLM2-360M")36 37 Returns:38 Dict with real model configuration values39 """40 model = load_model(model_name)41 config = model.config42 43 # Extract values directly from model.config44 num_heads = config.num_attention_heads45 num_kv_heads = getattr(config, 'num_key_value_heads', num_heads)46 head_dim = config.hidden_size // num_heads47 48 return {49 "num_layers": config.num_hidden_layers,50 "num_heads": num_heads,51 "num_kv_heads": num_kv_heads,52 "head_dim": head_dim,53 "hidden_size": config.hidden_size,54 "model_type": getattr(config, 'model_type', 'unknown'),55 "gqa_ratio": num_heads // num_kv_heads if num_kv_heads > 0 else 1,56 }57 58 59def run_prefill_with_real_model(60 model,61 attention_layer,62 seq_len: int,63 batch_size: int = 1,64 num_iterations: int = 5,65 use_flash: bool = True,66) -> dict:67 """68 Run prefill phase attention using a REAL model's attention layer.69 70 Prefill processes the entire prompt at once:71 - Hidden states have shape [batch, seq_len, hidden_dim]72 - Full N×N attention matrix computed via the real attention layer73 74 Args:75 model: Loaded HuggingFace model76 attention_layer: Extracted attention module77 seq_len: Sequence length78 batch_size: Batch size79 num_iterations: Number of timed iterations80 use_flash: Whether to use FlashAttention backend81 82 Returns:83 Dict with timing and memory stats84 """85 if not torch.cuda.is_available():86 return {"error": "CUDA not available"}87 88 device = torch.device("cuda")89 dtype = torch.float1690 91 # Create proper inputs for the attention layer92 hidden_states, position_ids = create_attention_inputs(93 model, batch_size, seq_len, device, dtype94 )95 96 # Backend configuration97 backend = "flash" if use_flash else "math"98 99 # Run benchmark using the utility function100 result = benchmark_attention_layer(101 attention_layer=attention_layer,102 hidden_states=hidden_states,103 position_ids=position_ids,104 backend=backend,105 num_iterations=num_iterations,106 warmup_iterations=2,107 )108 109 # Clean up110 del hidden_states, position_ids111 torch.cuda.empty_cache()112 113 # Add phase info to result114 result["seq_len"] = seq_len115 result["phase"] = "prefill"116 result["using_real_model"] = True117 118 return result119 120 121def run_prefill_benchmark(122 model_name: str,123 seq_len: int,124 batch_size: int = 1,125 num_iterations: int = 10,126 use_flash: bool = True,127) -> dict:128 """129 Benchmark prefill phase using F.scaled_dot_product_attention with REAL model dimensions.130 131 This function uses the model's actual configuration (from model.config) to create132 properly-sized Q, K, V tensors, then benchmarks the SDPA operation directly.133 This is more reliable than calling attention layer forward() methods.134 135 Args:136 model_name: Key from MODEL_CONFIGS (model will be loaded to get real config)137 seq_len: Sequence length (prompt tokens)138 batch_size: Batch size139 num_iterations: Number of timed iterations140 use_flash: Whether to use FlashAttention backend141 142 Returns:143 Dict with time_ms, memory_mb, and status144 """145 if not torch.cuda.is_available():146 return {"time_ms": 0, "memory_mb": 0, "status": "error: CUDA not available"}147 148 device = torch.device("cuda")149 dtype = torch.float16150 151 try:152 # Get REAL config from loaded model153 real_config = get_real_model_config(model_name)154 num_heads = real_config["num_heads"]155 head_dim = real_config["head_dim"]156 157 # Create Q, K, V tensors with REAL model dimensions158 # Shape: [batch, num_heads, seq_len, head_dim]159 Q = torch.randn(batch_size, num_heads, seq_len, head_dim, dtype=dtype, device=device)160 K = torch.randn(batch_size, num_heads, seq_len, head_dim, dtype=dtype, device=device)161 V = torch.randn(batch_size, num_heads, seq_len, head_dim, dtype=dtype, device=device)162 163 # Set backend flags164 if use_flash:165 enable_math, enable_flash, enable_mem_efficient = False, True, False166 else:167 enable_math, enable_flash, enable_mem_efficient = True, False, False168 169 # Warmup170 for _ in range(3):171 with torch.backends.cuda.sdp_kernel(172 enable_flash=enable_flash,173 enable_math=enable_math,174 enable_mem_efficient=enable_mem_efficient175 ):176 _ = F.scaled_dot_product_attention(Q, K, V, is_causal=True)177 178 torch.cuda.synchronize()179 torch.cuda.reset_peak_memory_stats()180 181 # Timed runs182 start = torch.cuda.Event(enable_timing=True)183 end = torch.cuda.Event(enable_timing=True)184 185 start.record()186 for _ in range(num_iterations):187 with torch.backends.cuda.sdp_kernel(188 enable_flash=enable_flash,189 enable_math=enable_math,190 enable_mem_efficient=enable_mem_efficient191 ):192 output = F.scaled_dot_product_attention(Q, K, V, is_causal=True)193 end.record()194 195 torch.cuda.synchronize()196 197 time_ms = start.elapsed_time(end) / num_iterations198 memory_mb = torch.cuda.max_memory_allocated() / (1024 * 1024)199 200 # Cleanup201 del Q, K, V, output202 torch.cuda.empty_cache()203 204 return {205 "time_ms": round(time_ms, 3),206 "memory_mb": round(memory_mb, 1),207 "seq_len": seq_len,208 "phase": "prefill",209 "backend": "flash" if use_flash else "math",210 "num_heads": num_heads,211 "head_dim": head_dim,212 "status": "success",213 "using_real_config": True,214 }215 216 except Exception as e:217 return {218 "time_ms": 0,219 "memory_mb": 0,220 "status": f"error: {str(e)[:100]}",221 "phase": "prefill",222 }223 224 225def run_decode_with_real_model(226 model,227 attention_layer,228 kv_cache_len: int,229 num_tokens: int = 10,230 batch_size: int = 1,231 num_iterations: int = 3,232 use_flash: bool = True,233) -> dict:234 """235 Run decode phase attention using a REAL model's attention layer.236 237 Decode generates one token at a time:238 - Single query token attending to all past keys/values239 - Simulates the memory-bound decode phase240 241 Args:242 model: Loaded HuggingFace model243 attention_layer: Extracted attention module244 kv_cache_len: Length of the KV cache (context)245 num_tokens: Number of tokens to simulate generating246 batch_size: Batch size247 num_iterations: Iterations for averaging248 use_flash: Whether to use FlashAttention backend249 250 Returns:251 Dict with per-token timing and memory stats252 """253 if not torch.cuda.is_available():254 return {"error": "CUDA not available"}255 256 device = torch.device("cuda")257 dtype = torch.float16258 259 # Create single-token query input (simulating decode)260 hidden_dim = model.config.hidden_size261 query_hidden = torch.randn(batch_size, 1, hidden_dim, dtype=dtype, device=device)262 position_ids = torch.tensor([[kv_cache_len]], device=device).expand(batch_size, 1)263 264 # Backend flags265 if use_flash:266 enable_math, enable_flash, enable_mem_efficient = False, True, False267 else:268 enable_math, enable_flash, enable_mem_efficient = True, False, False269 270 try:271 # Warmup272 with torch.backends.cuda.sdp_kernel(273 enable_flash=enable_flash,274 enable_math=enable_math,275 enable_mem_efficient=enable_mem_efficient276 ):277 with torch.no_grad():278 for _ in range(2):279 _ = attention_layer(query_hidden, position_ids=position_ids)280 281 torch.cuda.synchronize()282 torch.cuda.reset_peak_memory_stats()283 284 # Time multiple tokens285 start = torch.cuda.Event(enable_timing=True)286 end = torch.cuda.Event(enable_timing=True)287 288 with torch.backends.cuda.sdp_kernel(289 enable_flash=enable_flash,290 enable_math=enable_math,291 enable_mem_efficient=enable_mem_efficient292 ):293 with torch.no_grad():294 start.record()295 for _ in range(num_tokens * num_iterations):296 output = attention_layer(query_hidden, position_ids=position_ids)297 end.record()298 299 torch.cuda.synchronize()300 301 total_time_ms = start.elapsed_time(end)302 time_per_token_ms = total_time_ms / (num_tokens * num_iterations)303 memory_mb = torch.cuda.max_memory_allocated() / (1024 * 1024)304 305 # Clean up306 del query_hidden307 torch.cuda.empty_cache()308 309 return {310 "time_ms_per_token": round(time_per_token_ms, 4),311 "total_time_ms": round(total_time_ms / num_iterations, 3),312 "memory_mb": round(memory_mb, 1),313 "kv_cache_len": kv_cache_len,314 "num_tokens": num_tokens,315 "phase": "decode",316 "using_real_model": True,317 "status": "success",318 }319 320 except Exception as e:321 return {322 "time_ms_per_token": 0,323 "total_time_ms": 0,324 "memory_mb": 0,325 "kv_cache_len": kv_cache_len,326 "num_tokens": num_tokens,327 "phase": "decode",328 "status": f"error: {str(e)[:80]}",329 }330 331 332def run_decode_benchmark(333 model_name: str,334 kv_cache_len: int,335 num_tokens: int = 10,336 batch_size: int = 1,337 num_iterations: int = 5,338 use_flash: bool = True,339) -> dict:340 """341 Benchmark decode phase using F.scaled_dot_product_attention with REAL model dimensions.342 343 Properly simulates decode by:344 - Creating single query token (Q with seq_len=1)345 - Creating KV cache tensors with kv_cache_len tokens346 - Handling GQA by expanding KV heads to match Q heads347 348 Args:349 model_name: Key from MODEL_CONFIGS (model will be loaded to get real config)350 kv_cache_len: Length of KV cache (context length)351 num_tokens: Number of decode tokens to simulate352 batch_size: Batch size353 num_iterations: Iterations for timing354 use_flash: Whether to use FlashAttention backend355 356 Returns:357 Dict with time_ms_per_token, memory_mb, and status358 """359 if not torch.cuda.is_available():360 return {"time_ms_per_token": 0, "memory_mb": 0, "status": "error: CUDA not available"}361 362 device = torch.device("cuda")363 dtype = torch.float16364 365 try:366 # Get REAL config from loaded model367 real_config = get_real_model_config(model_name)368 num_heads = real_config["num_heads"]369 num_kv_heads = real_config["num_kv_heads"]370 head_dim = real_config["head_dim"]371 372 # Single query token: [batch, num_heads, 1, head_dim]373 Q = torch.randn(batch_size, num_heads, 1, head_dim, dtype=dtype, device=device)374 375 # KV cache with real model's KV head count: [batch, num_kv_heads, kv_cache_len, head_dim]376 K_cache = torch.randn(batch_size, num_kv_heads, kv_cache_len, head_dim, dtype=dtype, device=device)377 V_cache = torch.randn(batch_size, num_kv_heads, kv_cache_len, head_dim, dtype=dtype, device=device)378 379 # Handle GQA: expand KV heads to match Q heads if needed380 if num_kv_heads < num_heads:381 repeat_factor = num_heads // num_kv_heads382 K_cache = K_cache.repeat_interleave(repeat_factor, dim=1)383 V_cache = V_cache.repeat_interleave(repeat_factor, dim=1)384 385 # Set backend flags386 if use_flash:387 enable_math, enable_flash_flag, enable_mem_efficient = False, True, False388 else:389 enable_math, enable_flash_flag, enable_mem_efficient = True, False, False390 391 # Warmup392 for _ in range(3):393 with torch.backends.cuda.sdp_kernel(394 enable_flash=enable_flash_flag,395 enable_math=enable_math,396 enable_mem_efficient=enable_mem_efficient397 ):398 _ = F.scaled_dot_product_attention(Q, K_cache, V_cache)399 400 torch.cuda.synchronize()401 torch.cuda.reset_peak_memory_stats()402 403 # Timed runs - simulate generating num_tokens404 start = torch.cuda.Event(enable_timing=True)405 end = torch.cuda.Event(enable_timing=True)406 407 start.record()408 for _ in range(num_tokens * num_iterations):409 with torch.backends.cuda.sdp_kernel(410 enable_flash=enable_flash_flag,411 enable_math=enable_math,412 enable_mem_efficient=enable_mem_efficient413 ):414 output = F.scaled_dot_product_attention(Q, K_cache, V_cache)415 end.record()416 417 torch.cuda.synchronize()418 419 total_time_ms = start.elapsed_time(end)420 time_per_token_ms = total_time_ms / (num_tokens * num_iterations)421 memory_mb = torch.cuda.max_memory_allocated() / (1024 * 1024)422 423 # Cleanup424 del Q, K_cache, V_cache, output425 torch.cuda.empty_cache()426 427 return {428 "time_ms_per_token": round(time_per_token_ms, 4),429 "total_time_ms": round(total_time_ms / num_iterations, 3),430 "memory_mb": round(memory_mb, 1),431 "kv_cache_len": kv_cache_len,432 "num_tokens": num_tokens,433 "phase": "decode",434 "backend": "flash" if use_flash else "math",435 "num_heads": num_heads,436 "num_kv_heads": num_kv_heads,437 "head_dim": head_dim,438 "status": "success",439 "using_real_config": True,440 }441 442 except Exception as e:443 return {444 "time_ms_per_token": 0,445 "total_time_ms": 0,446 "memory_mb": 0,447 "kv_cache_len": kv_cache_len,448 "num_tokens": num_tokens,449 "phase": "decode",450 "status": f"error: {str(e)[:100]}",451 }452 453 454# Legacy function kept for backwards compatibility455def simulate_prefill_attention(456 batch_size: int,457 num_heads: int,458 seq_len: int,459 head_dim: int,460 num_iterations: int = 5,461 use_flash: bool = True,462) -> dict:463 """464 Legacy: Simulate prefill phase attention with random tensors.465 Use run_prefill_with_real_model() for real model benchmarks.466 """467 if not torch.cuda.is_available():468 return {"error": "CUDA not available"}469 470 device = torch.device("cuda")471 dtype = torch.float16472 473 Q = torch.randn(batch_size, num_heads, seq_len, head_dim, device=device, dtype=dtype)474 K = torch.randn(batch_size, num_heads, seq_len, head_dim, device=device, dtype=dtype)475 V = torch.randn(batch_size, num_heads, seq_len, head_dim, device=device, dtype=dtype)476 477 if use_flash:478 enable_math, enable_flash_flag, enable_mem_efficient = False, True, False479 else:480 enable_math, enable_flash_flag, enable_mem_efficient = True, False, False481 482 # Warmup483 for _ in range(2):484 with torch.backends.cuda.sdp_kernel(485 enable_flash=enable_flash_flag, enable_math=enable_math, enable_mem_efficient=enable_mem_efficient486 ):487 try:488 _ = F.scaled_dot_product_attention(Q, K, V)489 except Exception:490 pass491 492 torch.cuda.synchronize()493 torch.cuda.reset_peak_memory_stats()494 495 start = torch.cuda.Event(enable_timing=True)496 end = torch.cuda.Event(enable_timing=True)497 498 start.record()499 for _ in range(num_iterations):500 with torch.backends.cuda.sdp_kernel(501 enable_flash=enable_flash_flag, enable_math=enable_math, enable_mem_efficient=enable_mem_efficient502 ):503 try:504 output = F.scaled_dot_product_attention(Q, K, V)505 except Exception:506 output = F.scaled_dot_product_attention(Q, K, V)507 end.record()508 509 torch.cuda.synchronize()510 511 total_time_ms = start.elapsed_time(end)512 avg_time_ms = total_time_ms / num_iterations513 peak_memory_mb = torch.cuda.max_memory_allocated() / (1024 * 1024)514 515 del Q, K, V, output516 torch.cuda.empty_cache()517 518 return {519 "time_ms": avg_time_ms,520 "memory_mb": peak_memory_mb,521 "seq_len": seq_len,522 "phase": "prefill",523 }524 525 526# Legacy function kept for backwards compatibility527def simulate_decode_attention(528 batch_size: int,529 num_heads: int,530 kv_cache_len: int,531 head_dim: int,532 num_tokens: int = 10,533 use_flash: bool = True,534) -> dict:535 """536 Legacy: Simulate decode phase attention with random tensors.537 Use run_decode_with_real_model() for real model benchmarks.538 """539 if not torch.cuda.is_available():540 return {"error": "CUDA not available"}541 542 device = torch.device("cuda")543 dtype = torch.float16544 545 K_cache = torch.randn(batch_size, num_heads, kv_cache_len, head_dim, device=device, dtype=dtype)546 V_cache = torch.randn(batch_size, num_heads, kv_cache_len, head_dim, device=device, dtype=dtype)547 Q = torch.randn(batch_size, num_heads, 1, head_dim, device=device, dtype=dtype)548 549 if use_flash:550 enable_math, enable_flash_flag, enable_mem_efficient = False, True, False551 else:552 enable_math, enable_flash_flag, enable_mem_efficient = True, False, False553 554 # Warmup555 for _ in range(2):556 with torch.backends.cuda.sdp_kernel(557 enable_flash=enable_flash_flag, enable_math=enable_math, enable_mem_efficient=enable_mem_efficient558 ):559 try:560 _ = F.scaled_dot_product_attention(Q, K_cache, V_cache)561 except Exception:562 pass563 564 torch.cuda.synchronize()565 torch.cuda.reset_peak_memory_stats()566 567 start = torch.cuda.Event(enable_timing=True)568 end = torch.cuda.Event(enable_timing=True)569 570 start.record()571 for _ in range(num_tokens):572 with torch.backends.cuda.sdp_kernel(573 enable_flash=enable_flash_flag, enable_math=enable_math, enable_mem_efficient=enable_mem_efficient574 ):575 try:576 output = F.scaled_dot_product_attention(Q, K_cache, V_cache)577 except Exception:578 output = F.scaled_dot_product_attention(Q, K_cache, V_cache)579 end.record()580 581 torch.cuda.synchronize()582 583 total_time_ms = start.elapsed_time(end)584 avg_time_per_token_ms = total_time_ms / num_tokens585 peak_memory_mb = torch.cuda.max_memory_allocated() / (1024 * 1024)586 587 del Q, K_cache, V_cache, output588 torch.cuda.empty_cache()589 590 return {591 "time_ms_per_token": avg_time_per_token_ms,592 "total_time_ms": total_time_ms,593 "memory_mb": peak_memory_mb,594 "kv_cache_len": kv_cache_len,595 "num_tokens": num_tokens,596 "phase": "decode",597 }598 599 600def run_prefill_decode_comparison(601 model_name: str,602 context_length: int,603 decode_tokens: int = 32,604) -> tuple:605 """606 Run full comparison between prefill and decode phases using REAL HuggingFace model.607 608 Uses F.scaled_dot_product_attention with real model dimensions for reliable benchmarking.609 All config values come from model.config, not constants.610 611 Returns results dict, comparison chart, KV cache chart, and insight text.612 """613 if model_name not in MODEL_CONFIGS:614 return {"error": f"Unknown model: {model_name}"}, None, None, "Error: Unknown model"615 616 # Get REAL config from model.config (not constants)617 try:618 real_config = get_real_model_config(model_name)619 except Exception as e:620 return {"error": f"Failed to load model: {str(e)[:50]}"}, None, None, f"Error: {str(e)[:50]}"621 622 results = {623 "model": model_name,624 "context_length": context_length,625 "decode_tokens": decode_tokens,626 "real_config": real_config,627 "using_real_config": True,628 }629 630 # Run prefill benchmarks using SDPA with REAL model dimensions631 prefill_flash = run_prefill_benchmark(632 model_name=model_name,633 seq_len=context_length,634 batch_size=1,635 use_flash=True,636 )637 638 prefill_math = run_prefill_benchmark(639 model_name=model_name,640 seq_len=context_length,641 batch_size=1,642 use_flash=False,643 )644 645 # Run decode benchmarks using SDPA with proper KV cache simulation646 decode_flash = run_decode_benchmark(647 model_name=model_name,648 kv_cache_len=context_length,649 num_tokens=decode_tokens,650 batch_size=1,651 use_flash=True,652 )653 654 decode_math = run_decode_benchmark(655 model_name=model_name,656 kv_cache_len=context_length,657 num_tokens=decode_tokens,658 batch_size=1,659 use_flash=False,660 )661 662 results["prefill"] = {663 "flash": prefill_flash,664 "math": prefill_math,665 }666 results["decode"] = {667 "flash": decode_flash,668 "math": decode_math,669 }670 671 # Add model info for display672 results["model_info"] = {673 "num_heads": real_config["num_heads"],674 "num_kv_heads": real_config["num_kv_heads"],675 "head_dim": real_config["head_dim"],676 "num_layers": real_config["num_layers"],677 "gqa_ratio": real_config["gqa_ratio"],678 }679 680 # Create comparison chart681 comparison_chart = create_comparison_chart(results)682 683 # Create KV cache growth chart using REAL model config684 kv_cache_chart = create_kv_cache_chart(model_name, context_length, decode_tokens)685 686 # Generate insight687 insight = generate_phase_insight(results)688 689 # Add real model indicator to insight690 if results.get("using_real_config"):691 model_indicator = f"\n\n---\n\n*Benchmarked using real **{model_name}** config ({real_config['num_heads']} heads, {real_config['head_dim']}d, GQA {real_config['gqa_ratio']}:1)*"692 insight = insight + model_indicator693 694 return results, comparison_chart, kv_cache_chart, insight695 696 697def create_comparison_chart(results: dict) -> go.Figure:698 """Create bar chart comparing prefill vs decode timing."""699 700 prefill_flash = results["prefill"]["flash"]701 prefill_math = results["prefill"]["math"]702 decode_flash = results["decode"]["flash"]703 decode_math = results["decode"]["math"]704 705 # Helper to safely get numeric value (handles None)706 def safe_get(d, key, default=0):707 val = d.get(key, default)708 return val if val is not None else default709 710 fig = make_subplots(711 rows=1, cols=2,712 subplot_titles=("<b>Prefill Time</b> (Full Prompt)", "<b>Decode Time</b> (Per Token)"),713 horizontal_spacing=0.15,714 vertical_spacing=0.15,715 )716 717 # Get max values for proper y-axis scaling with headroom for labels718 prefill_math_time = safe_get(prefill_math, "time_ms", 0)719 prefill_flash_time = safe_get(prefill_flash, "time_ms", 0)720 decode_math_time = safe_get(decode_math, "time_ms_per_token", 0)721 decode_flash_time = safe_get(decode_flash, "time_ms_per_token", 0)722 723 prefill_max = max(prefill_math_time, prefill_flash_time)724 decode_max = max(decode_math_time, decode_flash_time)725 726 # Prefill comparison727 fig.add_trace(728 go.Bar(729 x=["Math<br>(Standard)", "Flash<br>Attention"],730 y=[prefill_math_time, prefill_flash_time],731 marker_color=["#ef4444", "#22c55e"],732 text=[f"<b>{prefill_math_time:.2f}ms</b>", f"<b>{prefill_flash_time:.2f}ms</b>"],733 textposition="inside",734 textangle=0,735 insidetextanchor="middle",736 textfont=dict(color="white", size=12),737 name="Prefill",738 showlegend=False,739 ),740 row=1, col=1741 )742 743 # Decode comparison (per token)744 fig.add_trace(745 go.Bar(746 x=["Math<br>(Standard)", "Flash<br>Attention"],747 y=[decode_math_time, decode_flash_time],748 marker_color=["#ef4444", "#22c55e"],749 text=[f"<b>{decode_math_time:.3f}ms</b>", f"<b>{decode_flash_time:.3f}ms</b>"],750 textposition="inside",751 textangle=0,752 insidetextanchor="middle",753 textfont=dict(color="white", size=12),754 name="Decode",755 showlegend=False,756 ),757 row=1, col=2758 )759 760 # Calculate speedups761 if prefill_math_time > 0 and prefill_flash_time > 0:762 prefill_speedup = prefill_math_time / prefill_flash_time763 else:764 prefill_speedup = 1.0765 766 if decode_math_time > 0 and decode_flash_time > 0:767 decode_speedup = decode_math_time / decode_flash_time768 else:769 decode_speedup = 1.0770 771 fig.update_layout(772 title=dict(773 text=f"<b>Prefill vs Decode: FlashAttention Speedup</b><br>"774 f"<span style='font-size:13px;color:#16a34a'>"775 f"Prefill: {prefill_speedup:.1f}× faster | Decode: {decode_speedup:.1f}× faster</span>",776 x=0.5,777 font=dict(size=15),778 ),779 height=380,780 margin=dict(l=60, r=40, t=100, b=60),781 yaxis_title="Time (ms)",782 yaxis2_title="Time (ms)",783 )784 785 # Add more y-axis headroom786 fig.update_yaxes(range=[0, prefill_max * 1.15], row=1, col=1)787 fig.update_yaxes(range=[0, decode_max * 1.15], row=1, col=2)788 789 return fig790 791 792def create_kv_cache_chart(model_name: str, context_length: int, decode_tokens: int) -> go.Figure:793 """794 Create chart showing KV cache growth during generation.795 796 Uses REAL model config values from model.config, not constants.797 798 Args:799 model_name: Model name to load config from800 context_length: Number of context tokens (prefill)801 decode_tokens: Number of decode tokens to generate802 803 Returns:804 Plotly figure showing KV cache growth805 """806 # Get REAL config from loaded model (no constants!)807 real_config = get_real_model_config(model_name)808 809 num_kv_heads = real_config["num_kv_heads"]810 head_dim = real_config["head_dim"]811 num_layers = real_config["num_layers"]812 813 # Calculate KV cache size at each step814 # KV cache per layer: 2 (K+V) × kv_heads × head_dim × 2 (FP16 bytes)815 bytes_per_token_per_layer = 2 * num_kv_heads * head_dim * 2816 total_bytes_per_token = bytes_per_token_per_layer * num_layers817 818 # Generate sequence of token counts819 token_counts = list(range(0, context_length + decode_tokens + 1, max(1, (context_length + decode_tokens) // 50)))820 if token_counts[-1] != context_length + decode_tokens:821 token_counts.append(context_length + decode_tokens)822 823 # Calculate cache sizes in MB824 cache_sizes_mb = [(t * total_bytes_per_token) / (1024 * 1024) for t in token_counts]825 826 fig = go.Figure()827 828 # Prefill region (0 to context_length)829 prefill_tokens = [t for t in token_counts if t <= context_length]830 prefill_sizes = [(t * total_bytes_per_token) / (1024 * 1024) for t in prefill_tokens]831 832 fig.add_trace(go.Scatter(833 x=prefill_tokens,834 y=prefill_sizes,835 mode="lines",836 name="Prefill Phase",837 fill="tozeroy",838 line=dict(color="#3b82f6", width=2),839 fillcolor="rgba(59, 130, 246, 0.3)",840 ))841 842 # Decode region (context_length to end)843 decode_tokens_list = [t for t in token_counts if t >= context_length]844 decode_sizes = [(t * total_bytes_per_token) / (1024 * 1024) for t in decode_tokens_list]845 846 fig.add_trace(go.Scatter(847 x=decode_tokens_list,848 y=decode_sizes,849 mode="lines",850 name="Decode Phase",851 fill="tozeroy",852 line=dict(color="#22c55e", width=2),853 fillcolor="rgba(34, 197, 94, 0.3)",854 ))855 856 # Add vertical line at context boundary857 cache_at_context = (context_length * total_bytes_per_token) / (1024 * 1024)858 fig.add_vline(859 x=context_length,860 line_dash="dash",861 line_color="rgba(0, 0, 0, 0.5)",862 annotation_text=f"Prefill→Decode<br>({cache_at_context:.1f} MB)",863 annotation_position="top",864 )865 866 fig.update_layout(867 title=dict(868 text=f"KV Cache Growth ({num_kv_heads} KV heads × {num_layers} layers)",869 x=0.5,870 ),871 xaxis_title="Tokens Processed",872 yaxis_title="KV Cache Size (MB)",873 height=300,874 margin=dict(l=50, r=50, t=60, b=50),875 legend=dict(876 orientation="h",877 yanchor="bottom",878 y=-0.25,879 xanchor="center",880 x=0.5,881 ),882 yaxis=dict(rangemode='tozero'),883 )884 885 return fig886 887 888def generate_phase_insight(results: dict) -> str:889 """Generate insight text from comparison results."""890 891 prefill_flash = results["prefill"]["flash"]892 prefill_math = results["prefill"]["math"]893 decode_flash = results["decode"]["flash"]894 decode_math = results["decode"]["math"]895 896 # Helper to safely get numeric value (handles None)897 def safe_get(d, key, default=0):898 val = d.get(key, default)899 return val if val is not None else default900 901 prefill_math_time = safe_get(prefill_math, "time_ms", 0)902 prefill_flash_time = safe_get(prefill_flash, "time_ms", 0)903 decode_math_time = safe_get(decode_math, "time_ms_per_token", 0)904 decode_flash_time = safe_get(decode_flash, "time_ms_per_token", 0)905 906 # Calculate speedups907 if prefill_math_time > 0 and prefill_flash_time > 0:908 prefill_speedup = prefill_math_time / prefill_flash_time909 else:910 prefill_speedup = 1.0911 912 if decode_math_time > 0 and decode_flash_time > 0:913 decode_speedup = decode_math_time / decode_flash_time914 else:915 decode_speedup = 1.0916 917 context_length = results["context_length"]918 decode_tokens = results["decode_tokens"]919 920 insight = f"""### Key Observations921 922**Prefill Phase** (processing {context_length} tokens):923- Standard attention: **{prefill_math_time:.2f}ms**924- FlashAttention: **{prefill_flash_time:.2f}ms**925- Speedup: **{prefill_speedup:.1f}×**926 927**Decode Phase** (generating {decode_tokens} tokens):928- Standard attention: **{decode_math_time:.3f}ms/token**929- FlashAttention: **{decode_flash_time:.3f}ms/token**930- Speedup: **{decode_speedup:.1f}×**931 932---933 934### Why the Difference?935 9361. **Prefill is compute-bound** with N² attention operations937 - FlashAttention's memory efficiency provides significant speedup938 - Larger contexts benefit more (quadratic scaling)939 9402. **Decode is memory-bound** with 1×N attention per token941 - Each decode step is fast but sequential942 - KV cache read dominates, limiting FlashAttention's advantage943 9443. **Optimal strategy**: FlashAttention helps most during prefill;945 decode phase benefits from KV cache optimizations (GQA/MQA)946"""947 948 return insight949 950 951def get_attention_pattern_chart(context_length: int) -> go.Figure:952 """Create visualization of prefill vs decode attention patterns using scatter."""953 954 # Calculate FLOPs for insight955 prefill_flops = context_length * context_length # N² attention956 decode_flops_per_token = context_length # 1×N per decode token957 958 fig = make_subplots(959 rows=1, cols=2,960 subplot_titles=(961 f"<b>Prefill:</b> {context_length}×{context_length} = {prefill_flops:,} ops",962 f"<b>Decode:</b> 1×{context_length} = {decode_flops_per_token:,} ops/token"963 ),964 horizontal_spacing=0.15,965 )966 967 # Prefill: Lower triangular pattern (causal mask)968 # Dynamic size based on context length for visual feedback969 if context_length <= 16:970 size = context_length971 elif context_length <= 128:972 size = 12973 elif context_length <= 512:974 size = 10975 else:976 size = 8 # Smaller grid for very large contexts977 978 # Adjust marker size based on grid size979 marker_size = max(10, 22 - size)980 981 # Generate coordinates for filled cells (lower triangular)982 prefill_x = []983 prefill_y = []984 for row in range(size):985 for col in range(row + 1): # Only up to diagonal986 prefill_x.append(col)987 prefill_y.append(row)988 989 fig.add_trace(990 go.Scatter(991 x=prefill_x,992 y=prefill_y,993 mode="markers",994 marker=dict(995 size=marker_size,996 color="#3b82f6",997 symbol="square",998 ),999 name="Attends",1000 showlegend=False,1001 hovertemplate="Query %{y} → Key %{x}<extra></extra>",1002 ),1003 row=1, col=11004 )1005 1006 # Decode: Each step attends to growing sequence1007 num_decode_steps = 61008 base_context = max(4, size - num_decode_steps)1009 1010 decode_x = []1011 decode_y = []1012 for step in range(num_decode_steps):1013 attend_len = base_context + step + 11014 for col in range(min(attend_len, size)):1015 decode_x.append(col)1016 decode_y.append(step)1017 1018 fig.add_trace(1019 go.Scatter(1020 x=decode_x,1021 y=decode_y,1022 mode="markers",1023 marker=dict(1024 size=marker_size + 4,1025 color="#22c55e",1026 symbol="square",1027 ),1028 name="Attends",1029 showlegend=False,1030 hovertemplate="Decode step %{y} → Key %{x}<extra></extra>",1031 ),1032 row=1, col=21033 )1034 1035 # Update axes with proper ranges1036 fig.update_xaxes(1037 title_text="Key positions", 1038 range=[-0.5, size - 0.5],1039 dtick=2,1040 row=1, col=11041 )1042 fig.update_xaxes(1043 title_text="Key positions (KV cache)", 1044 range=[-0.5, size - 0.5],1045 dtick=2,1046 row=1, col=21047 )1048 fig.update_yaxes(1049 title_text="Query positions", 1050 range=[-0.5, size - 0.5],1051 dtick=2,1052 row=1, col=11053 )1054 fig.update_yaxes(1055 title_text="Decode steps", 1056 range=[-0.5, num_decode_steps - 0.5],1057 dtick=1,1058 row=1, col=21059 )1060 1061 fig.update_layout(1062 height=380,1063 margin=dict(l=60, r=30, t=70, b=50),1064 plot_bgcolor="rgba(241, 245, 249, 0.5)",1065 )1066 1067 return fig1068