Phase 2: Add Visualizer tab with tiling animation, online softmax, and memory hierarchy

app.py

@@ -31,6 +31,12 @@ from src.constants import (
     SEQ_LENGTH_OPTIONS,
 )
 from src.models import get_available_models, get_model_memory_footprint
+from src.visualizer import (
+    create_tiling_grid,
+    create_online_softmax_state,
+    create_memory_hierarchy_diagram,
+    get_max_steps,
+)
 
 
 def create_app() -> gr.Blocks:
@@ -59,22 +65,137 @@ def create_app() -> gr.Blocks:
                 ## FlashAttention Visualizer
                 
                 Understand how FlashAttention processes attention in tiles,
-                avoiding the O(N²) memory bottleneck.
-                
-                *Coming in Phase 2...*
+                avoiding the O(N²) memory bottleneck. Step through the algorithm
+                to see how tiles are processed and how online softmax maintains
+                running statistics.
                 """)
                 
+                # Controls
+                with gr.Row():
+                    with gr.Column(scale=1):
+                        seq_len_viz = gr.Slider(
+                            minimum=4,
+                            maximum=16,
+                            step=2,
+                            value=8,
+                            label="Sequence Length (tokens)",
+                        )
+                    with gr.Column(scale=1):
+                        block_size_viz = gr.Slider(
+                            minimum=2,
+                            maximum=4,
+                            step=1,
+                            value=2,
+                            label="Block Size",
+                        )
+                    with gr.Column(scale=1):
+                        causal_viz = gr.Checkbox(
+                            value=False,
+                            label="Causal Masking",
+                        )
+                
+                # Step controls
+                with gr.Row():
+                    step_back_btn = gr.Button("◀ Step Back", size="sm")
+                    step_slider = gr.Slider(
+                        minimum=0,
+                        maximum=15,
+                        step=1,
+                        value=0,
+                        label="Current Step",
+                    )
+                    step_forward_btn = gr.Button("Step Forward ▶", size="sm")
+                    reset_btn = gr.Button("Reset", size="sm", variant="secondary")
+                
+                # Tiling and Online Softmax side by side
                 with gr.Row():
-                    with gr.Column():
-                        gr.Markdown("### Tiling Animation")
-                        gr.Markdown("*Tile-by-tile attention computation will appear here*")
+                    with gr.Column(scale=1):
+                        gr.Markdown("### Attention Matrix Tiling")
+                        tiling_plot = gr.Plot(label="Tiling View")
                     
-                    with gr.Column():
+                    with gr.Column(scale=1):
                         gr.Markdown("### Online Softmax State")
-                        gr.Markdown("*Running max (m), sum (l), and output (O) tracking*")
+                        softmax_plot = gr.Plot(label="Running m and l")
+                        softmax_explanation = gr.Markdown("*Step through to see online softmax updates*")
+                
+                # Memory Hierarchy
+                gr.Markdown("### Memory Hierarchy Comparison")
+                with gr.Row():
+                    algo_choice = gr.Radio(
+                        choices=["flash", "standard"],
+                        value="flash",
+                        label="Algorithm",
+                    )
+                memory_plot = gr.Plot(label="Memory Hierarchy")
+                
+                # Event handlers for visualizer
+                def update_visualizations(seq_len, block_size, causal, step):
+                    """Update all visualizations based on current parameters."""
+                    max_steps = get_max_steps(seq_len, block_size, causal)
+                    # Clamp step to valid range
+                    step = min(step, max_steps - 1)
+                    step = max(step, 0)
+                    
+                    tiling_fig = create_tiling_grid(seq_len, block_size, step, causal)
+                    
+                    # Online softmax uses 4 tiles for the example
+                    num_tiles = seq_len // block_size
+                    softmax_step = min(step, num_tiles - 1)
+                    softmax_fig, explanation = create_online_softmax_state(softmax_step, num_tiles)
+                    
+                    return tiling_fig, softmax_fig, explanation, step
+                
+                def update_memory_hierarchy(algo):
+                    """Update memory hierarchy diagram."""
+                    return create_memory_hierarchy_diagram(algo)
+                
+                def step_forward(seq_len, block_size, causal, current_step):
+                    """Move to next step."""
+                    max_steps = get_max_steps(seq_len, block_size, causal)
+                    new_step = min(current_step + 1, max_steps - 1)
+                    return new_step
+                
+                def step_back(current_step):
+                    """Move to previous step."""
+                    return max(current_step - 1, 0)
+                
+                def reset_step():
+                    """Reset to step 0."""
+                    return 0
+                
+                # Wire up events
+                viz_inputs = [seq_len_viz, block_size_viz, causal_viz, step_slider]
+                viz_outputs = [tiling_plot, softmax_plot, softmax_explanation, step_slider]
+                
+                # Update on parameter change
+                seq_len_viz.change(fn=update_visualizations, inputs=viz_inputs, outputs=viz_outputs)
+                block_size_viz.change(fn=update_visualizations, inputs=viz_inputs, outputs=viz_outputs)
+                causal_viz.change(fn=update_visualizations, inputs=viz_inputs, outputs=viz_outputs)
+                step_slider.change(fn=update_visualizations, inputs=viz_inputs, outputs=viz_outputs)
+                
+                # Step controls
+                step_forward_btn.click(
+                    fn=step_forward,
+                    inputs=[seq_len_viz, block_size_viz, causal_viz, step_slider],
+                    outputs=step_slider
+                )
+                step_back_btn.click(fn=step_back, inputs=step_slider, outputs=step_slider)
+                reset_btn.click(fn=reset_step, outputs=step_slider)
+                
+                # Memory hierarchy
+                algo_choice.change(fn=update_memory_hierarchy, inputs=algo_choice, outputs=memory_plot)
                 
-                gr.Markdown("### Memory Hierarchy")
-                gr.Markdown("*HBM ↔ SRAM data movement visualization*")
+                # Initialize on load
+                demo.load(
+                    fn=update_visualizations,
+                    inputs=viz_inputs,
+                    outputs=viz_outputs
+                )
+                demo.load(
+                    fn=update_memory_hierarchy,
+                    inputs=algo_choice,
+                    outputs=memory_plot
+                )
             
             # Tab 2: Benchmark (Zero GPU)
             with gr.Tab("Benchmark", id="tab-benchmark"):

src/visualizer.py

@@ -0,0 +1,490 @@
+"""
+Visualizer for FlashAttention concepts.
+CPU-only animations showing tiling, online softmax, and memory hierarchy.
+"""
+
+import numpy as np
+import plotly.graph_objects as go
+from plotly.subplots import make_subplots
+
+
+def create_tiling_grid(
+    seq_len: int = 8,
+    block_size: int = 2,
+    current_step: int = 0,
+    causal: bool = False
+) -> go.Figure:
+    """
+    Create a grid visualization showing FlashAttention tile processing.
+    
+    Args:
+        seq_len: Sequence length (number of tokens)
+        block_size: Size of each tile block
+        current_step: Current step in the animation (0-indexed)
+        causal: Whether to use causal masking
+        
+    Returns:
+        Plotly figure with the tiling grid
+    """
+    num_blocks = seq_len // block_size
+    total_tiles = num_blocks * num_blocks if not causal else sum(range(1, num_blocks + 1))
+    
+    # Create figure
+    fig = go.Figure()
+    
+    # Calculate which tiles are done, current, future, or masked
+    tile_idx = 0
+    annotations = []
+    
+    for i in range(num_blocks):  # Query blocks (rows)
+        for j in range(num_blocks):  # Key blocks (columns)
+            x0, x1 = j, j + 1
+            y0, y1 = num_blocks - i - 1, num_blocks - i
+            
+            # Determine tile status
+            if causal and j > i:
+                # Masked tile (future keys for causal attention)
+                color = "rgba(200, 200, 200, 0.3)"
+                status = "masked"
+            elif tile_idx < current_step:
+                # Done
+                color = "rgba(34, 197, 94, 0.6)"  # Green
+                status = "done"
+            elif tile_idx == current_step:
+                # Current
+                color = "rgba(249, 115, 22, 0.8)"  # Orange
+                status = "current"
+            else:
+                # Future
+                color = "rgba(229, 231, 235, 0.5)"  # Light gray
+                status = "pending"
+            
+            # Add rectangle
+            fig.add_shape(
+                type="rect",
+                x0=x0, y0=y0, x1=x1, y1=y1,
+                line=dict(color="rgba(0,0,0,0.3)", width=1),
+                fillcolor=color,
+            )
+            
+            # Add label for current tile
+            if status == "current":
+                annotations.append(dict(
+                    x=(x0 + x1) / 2,
+                    y=(y0 + y1) / 2,
+                    text=f"Q[{i}]×K[{j}]",
+                    showarrow=False,
+                    font=dict(size=10, color="white", weight="bold"),
+                ))
+            
+            if not (causal and j > i):
+                tile_idx += 1
+    
+    # Add axis labels
+    for i in range(num_blocks):
+        # K labels (top)
+        annotations.append(dict(
+            x=i + 0.5,
+            y=num_blocks + 0.2,
+            text=f"K[{i}]",
+            showarrow=False,
+            font=dict(size=9, color="gray"),
+        ))
+        # Q labels (left)
+        annotations.append(dict(
+            x=-0.3,
+            y=num_blocks - i - 0.5,
+            text=f"Q[{i}]",
+            showarrow=False,
+            font=dict(size=9, color="gray"),
+        ))
+    
+    fig.update_layout(
+        annotations=annotations,
+        xaxis=dict(
+            range=[-0.5, num_blocks + 0.5],
+            showgrid=False,
+            zeroline=False,
+            showticklabels=False,
+            title="Key Blocks →",
+        ),
+        yaxis=dict(
+            range=[-0.5, num_blocks + 0.5],
+            showgrid=False,
+            zeroline=False,
+            showticklabels=False,
+            scaleanchor="x",
+            title="← Query Blocks",
+        ),
+        height=350,
+        margin=dict(l=50, r=20, t=40, b=50),
+        title=dict(
+            text=f"Attention Matrix Tiling (Step {current_step + 1}/{tile_idx if current_step >= tile_idx else total_tiles})",
+            x=0.5,
+        ),
+        showlegend=False,
+    )
+    
+    # Add legend manually
+    legend_items = [
+        ("Current", "rgba(249, 115, 22, 0.8)"),
+        ("Done", "rgba(34, 197, 94, 0.6)"),
+        ("Pending", "rgba(229, 231, 235, 0.5)"),
+    ]
+    if causal:
+        legend_items.append(("Masked", "rgba(200, 200, 200, 0.3)"))
+    
+    for idx, (name, color) in enumerate(legend_items):
+        fig.add_trace(go.Scatter(
+            x=[None], y=[None],
+            mode="markers",
+            marker=dict(size=15, color=color, symbol="square"),
+            name=name,
+            showlegend=True,
+        ))
+    
+    fig.update_layout(
+        legend=dict(
+            orientation="h",
+            yanchor="bottom",
+            y=-0.25,
+            xanchor="center",
+            x=0.5,
+        )
+    )
+    
+    return fig
+
+
+def create_online_softmax_state(
+    current_step: int = 0,
+    num_tiles: int = 4,
+) -> tuple[go.Figure, str]:
+    """
+    Create visualization of online softmax state (m, l, O) evolution.
+    
+    Uses a concrete 8-token example with block_size=2.
+    Shows how running max (m) and sum (l) update, with rescaling when max changes.
+    
+    Args:
+        current_step: Current tile being processed (0-indexed)
+        num_tiles: Total number of tiles
+        
+    Returns:
+        Tuple of (Plotly figure, explanation text)
+    """
+    # Pre-computed example values for 8 tokens, block_size=2
+    # Simulating attention scores from Q[0] to all K blocks
+    example_data = [
+        {
+            "tile": 0,
+            "block_max": 2.1,
+            "block_sum_exp": 3.42,
+            "m_before": float("-inf"),
+            "m_after": 2.1,
+            "l_before": 0.0,
+            "l_after": 3.42,
+            "rescale_factor": 1.0,
+            "rescaled": False,
+        },
+        {
+            "tile": 1,
+            "block_max": 3.5,
+            "block_sum_exp": 5.21,
+            "m_before": 2.1,
+            "m_after": 3.5,
+            "l_before": 3.42,
+            "l_after": 1.70,  # 3.42 * exp(2.1-3.5) + 5.21 = 0.85 + 5.21
+            "rescale_factor": 0.247,  # exp(2.1 - 3.5)
+            "rescaled": True,
+        },
+        {
+            "tile": 2,
+            "block_max": 2.8,
+            "block_sum_exp": 4.01,
+            "m_before": 3.5,
+            "m_after": 3.5,  # No change - block_max < m
+            "l_before": 6.06,
+            "l_after": 8.03,  # 6.06 * 1.0 + 4.01 * exp(2.8-3.5)
+            "rescale_factor": 1.0,
+            "rescaled": False,
+        },
+        {
+            "tile": 3,
+            "block_max": 4.2,
+            "block_sum_exp": 6.83,
+            "m_before": 3.5,
+            "m_after": 4.2,
+            "l_before": 8.03,
+            "l_after": 10.79,  # 8.03 * exp(3.5-4.2) + 6.83
+            "rescale_factor": 0.497,  # exp(3.5 - 4.2)
+            "rescaled": True,
+        },
+    ]
+    
+    # Build the visualization
+    step = min(current_step, len(example_data) - 1)
+    current_data = example_data[step]
+    
+    # Create figure with bar chart showing m and l evolution
+    fig = make_subplots(
+        rows=1, cols=2,
+        subplot_titles=("Running Max (m)", "Running Sum (l)"),
+        horizontal_spacing=0.15,
+    )
+    
+    # Get historical values up to current step
+    m_values = [example_data[i]["m_after"] if i <= step else None for i in range(num_tiles)]
+    l_values = [example_data[i]["l_after"] if i <= step else None for i in range(num_tiles)]
+    
+    # Colors - highlight rescaling events
+    m_colors = []
+    l_colors = []
+    for i in range(num_tiles):
+        if i > step:
+            m_colors.append("rgba(200, 200, 200, 0.5)")
+            l_colors.append("rgba(200, 200, 200, 0.5)")
+        elif i == step:
+            m_colors.append("rgba(249, 115, 22, 0.9)")  # Orange for current
+            l_colors.append("rgba(249, 115, 22, 0.9)")
+        elif example_data[i]["rescaled"]:
+            m_colors.append("rgba(239, 68, 68, 0.7)")  # Red for rescale events
+            l_colors.append("rgba(239, 68, 68, 0.7)")
+        else:
+            m_colors.append("rgba(34, 197, 94, 0.7)")  # Green for normal
+            l_colors.append("rgba(34, 197, 94, 0.7)")
+    
+    # Add bars for m
+    fig.add_trace(
+        go.Bar(
+            x=[f"Tile {i}" for i in range(num_tiles)],
+            y=[v if v is not None else 0 for v in m_values],
+            marker_color=m_colors,
+            text=[f"{v:.2f}" if v is not None else "" for v in m_values],
+            textposition="outside",
+            name="m (max)",
+        ),
+        row=1, col=1
+    )
+    
+    # Add bars for l
+    fig.add_trace(
+        go.Bar(
+            x=[f"Tile {i}" for i in range(num_tiles)],
+            y=[v if v is not None else 0 for v in l_values],
+            marker_color=l_colors,
+            text=[f"{v:.2f}" if v is not None else "" for v in l_values],
+            textposition="outside",
+            name="l (sum)",
+        ),
+        row=1, col=2
+    )
+    
+    fig.update_layout(
+        height=280,
+        margin=dict(l=40, r=40, t=60, b=40),
+        showlegend=False,
+    )
+    
+    fig.update_yaxes(range=[0, 12], row=1, col=1)
+    fig.update_yaxes(range=[0, 15], row=1, col=2)
+    
+    # Generate explanation text
+    d = current_data
+    if d["rescaled"]:
+        explanation = f"""**Processing Tile {step} (Keys {step*2}-{step*2+1})**
+
+🔴 **MAX CHANGED!** Block max ({d['block_max']:.2f}) > Previous max ({d['m_before']:.2f})
+
+**Rescaling required:**
+- Rescale factor: exp({d['m_before']:.1f} - {d['block_max']:.1f}) = **{d['rescale_factor']:.3f}**
+- Previous l rescaled: {d['l_before']:.2f} × {d['rescale_factor']:.3f} = {d['l_before'] * d['rescale_factor']:.2f}
+- New l = rescaled + block_sum = **{d['l_after']:.2f}**
+- Previous O also rescaled by {d['rescale_factor']:.3f}
+
+*This is the key insight: when max increases, we must rescale all previous accumulations!*
+"""
+    else:
+        explanation = f"""**Processing Tile {step} (Keys {step*2}-{step*2+1})**
+
+✅ No rescaling needed (block max {d['block_max']:.2f} ≤ current max {d['m_after']:.2f})
+
+**Simple accumulation:**
+- m stays at: **{d['m_after']:.2f}**
+- l += block_sum × exp(block_max - m)
+- l = {d['l_before']:.2f} + {d['block_sum_exp']:.2f} × exp({d['block_max']:.1f} - {d['m_after']:.1f}) = **{d['l_after']:.2f}**
+"""
+    
+    return fig, explanation
+
+
+def create_memory_hierarchy_diagram(
+    algorithm: str = "flash",
+    current_step: int = 0,
+) -> go.Figure:
+    """
+    Create a diagram showing HBM vs SRAM memory hierarchy.
+    
+    Args:
+        algorithm: "standard" or "flash"
+        current_step: For animation purposes
+        
+    Returns:
+        Plotly figure showing memory hierarchy
+    """
+    fig = go.Figure()
+    
+    # Define positions
+    hbm_y = 0.7
+    sram_y = 0.3
+    
+    # HBM box
+    fig.add_shape(
+        type="rect",
+        x0=0.05, y0=0.55, x1=0.95, y1=0.95,
+        fillcolor="rgba(59, 130, 246, 0.1)",
+        line=dict(color="rgba(59, 130, 246, 0.8)", width=2),
+    )
+    
+    # SRAM box
+    fig.add_shape(
+        type="rect",
+        x0=0.2, y0=0.15, x1=0.8, y1=0.45,
+        fillcolor="rgba(34, 197, 94, 0.1)",
+        line=dict(color="rgba(34, 197, 94, 0.8)", width=2),
+    )
+    
+    # HBM matrices (Q, K, V, O)
+    matrix_width = 0.15
+    matrices = ["Q", "K", "V", "O"]
+    hbm_x_start = 0.15
+    
+    for i, name in enumerate(matrices):
+        x = hbm_x_start + i * 0.2
+        fig.add_shape(
+            type="rect",
+            x0=x, y0=0.65, x1=x + matrix_width, y1=0.85,
+            fillcolor="rgba(59, 130, 246, 0.3)",
+            line=dict(color="rgba(59, 130, 246, 0.6)", width=1),
+        )
+        fig.add_annotation(
+            x=x + matrix_width/2, y=0.75,
+            text=f"<b>{name}</b><br>[N, d]",
+            showarrow=False,
+            font=dict(size=11),
+        )
+    
+    # SRAM tiles
+    if algorithm == "flash":
+        tiles = ["Q_tile", "K_tile", "V_tile", "S_tile", "O_tile"]
+        tile_width = 0.1
+        sram_x_start = 0.25
+        
+        for i, name in enumerate(tiles):
+            x = sram_x_start + i * 0.11
+            # Highlight current tile being processed
+            is_active = (i == current_step % len(tiles))
+            fill = "rgba(249, 115, 22, 0.5)" if is_active else "rgba(34, 197, 94, 0.3)"
+            
+            fig.add_shape(
+                type="rect",
+                x0=x, y0=0.22, x1=x + tile_width, y1=0.38,
+                fillcolor=fill,
+                line=dict(color="rgba(34, 197, 94, 0.6)", width=1),
+            )
+            fig.add_annotation(
+                x=x + tile_width/2, y=0.30,
+                text=name.replace("_", "<br>"),
+                showarrow=False,
+                font=dict(size=9),
+            )
+        
+        # Transfer arrows (selective)
+        # Show only tile-sized transfers
+        fig.add_annotation(
+            x=0.5, y=0.50,
+            ax=0.5, ay=0.55,
+            xref="x", yref="y",
+            axref="x", ayref="y",
+            text="",
+            showarrow=True,
+            arrowhead=2,
+            arrowsize=1.5,
+            arrowwidth=2,
+            arrowcolor="rgba(34, 197, 94, 0.8)",
+        )
+        fig.add_annotation(
+            x=0.5, y=0.52,
+            text="O(B) per tile",
+            showarrow=False,
+            font=dict(size=10, color="green"),
+        )
+    else:
+        # Standard attention - full matrix in SRAM (doesn't fit!)
+        fig.add_shape(
+            type="rect",
+            x0=0.3, y0=0.22, x1=0.7, y1=0.38,
+            fillcolor="rgba(239, 68, 68, 0.3)",
+            line=dict(color="rgba(239, 68, 68, 0.6)", width=1, dash="dash"),
+        )
+        fig.add_annotation(
+            x=0.5, y=0.30,
+            text="S[N,N]<br>❌ Doesn't fit!",
+            showarrow=False,
+            font=dict(size=10, color="red"),
+        )
+        
+        # Transfer arrows (full matrix)
+        fig.add_annotation(
+            x=0.5, y=0.50,
+            ax=0.5, ay=0.55,
+            xref="x", yref="y",
+            axref="x", ayref="y",
+            text="",
+            showarrow=True,
+            arrowhead=2,
+            arrowsize=1.5,
+            arrowwidth=2,
+            arrowcolor="rgba(239, 68, 68, 0.8)",
+        )
+        fig.add_annotation(
+            x=0.5, y=0.52,
+            text="O(N²) traffic!",
+            showarrow=False,
+            font=dict(size=10, color="red"),
+        )
+    
+    # Labels
+    fig.add_annotation(
+        x=0.5, y=0.97,
+        text="<b>HBM (High Bandwidth Memory)</b><br>80 GB capacity | 2 TB/s bandwidth | ~400 cycles latency",
+        showarrow=False,
+        font=dict(size=11),
+    )
+    fig.add_annotation(
+        x=0.5, y=0.12,
+        text="<b>SRAM (Shared Memory)</b><br>192 KB capacity | 19 TB/s bandwidth | ~20 cycles latency",
+        showarrow=False,
+        font=dict(size=11),
+    )
+    
+    fig.update_layout(
+        xaxis=dict(range=[0, 1], showgrid=False, zeroline=False, showticklabels=False),
+        yaxis=dict(range=[0, 1], showgrid=False, zeroline=False, showticklabels=False),
+        height=400,
+        margin=dict(l=20, r=20, t=40, b=20),
+        title=dict(
+            text=f"Memory Hierarchy: {'FlashAttention' if algorithm == 'flash' else 'Standard Attention'}",
+            x=0.5,
+        ),
+    )
+    
+    return fig
+
+
+def get_max_steps(seq_len: int, block_size: int, causal: bool) -> int:
+    """Calculate total number of steps for the tiling animation."""
+    num_blocks = seq_len // block_size
+    if causal:
+        return sum(range(1, num_blocks + 1))
+    return num_blocks * num_blocks