Remove Visualizer tab - focus on GPU benchmarks and real model analysis

app.py

@@ -31,12 +31,6 @@ 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,
-)
 from src.prefill_decode import (
     get_attention_pattern_chart,
 )
@@ -63,167 +57,14 @@ def create_app() -> gr.Blocks:
         # FlashAttention Explorer
         
         Interactive demonstrations for the **Attention Optimizations** article series.
-        Explore FlashAttention tiling, online softmax, GQA/MQA, and memory budgets
-        using real HuggingFace models.
+        Benchmark attention backends, explore GQA/MQA, prefill vs decode phases, 
+        and memory budgets using real HuggingFace models on GPU.
         
         **Models used:** SmolLM2-360M, Qwen2.5-0.5B, Llama-3.2-1B
         """)
         
         with gr.Tabs():
-            # Tab 1: Visualizer (CPU)
-            with gr.Tab("Visualizer", id="tab-visualizer"):
-                gr.Markdown("""
-                ## FlashAttention Visualizer
-                
-                Understand how FlashAttention processes attention in tiles,
-                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(scale=1):
-                        gr.Markdown("### Attention Matrix Tiling")
-                        tiling_plot = gr.Plot(label="Tiling View")
-                    
-                    with gr.Column(scale=1):
-                        gr.Markdown("### Online Softmax State")
-                        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)
-                    
-                    # Only return plots, not step - step is controlled by buttons
-                    return tiling_fig, softmax_fig, explanation
-                
-                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(seq_len, block_size, causal):
-                    """Reset to step 0 and update visualizations."""
-                    step = 0
-                    tiling_fig = create_tiling_grid(seq_len, block_size, step, causal)
-                    num_tiles = seq_len // block_size
-                    softmax_fig, explanation = create_online_softmax_state(step, num_tiles)
-                    return step, tiling_fig, softmax_fig, explanation
-                
-                # Wire up events
-                viz_inputs = [seq_len_viz, block_size_viz, causal_viz, step_slider]
-                viz_outputs = [tiling_plot, softmax_plot, softmax_explanation]
-                
-                # Update on parameter change - these don't modify step
-                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 toggle resets step to 0 since the grid structure changes
-                causal_viz.change(
-                    fn=reset_step,
-                    inputs=[seq_len_viz, block_size_viz, causal_viz],
-                    outputs=[step_slider, tiling_plot, softmax_plot, softmax_explanation]
-                )
-                # Step slider changes only update plots, not the slider itself
-                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,
-                    inputs=[seq_len_viz, block_size_viz, causal_viz],
-                    outputs=[step_slider, tiling_plot, softmax_plot, softmax_explanation]
-                )
-                
-                # Memory hierarchy
-                algo_choice.change(fn=update_memory_hierarchy, inputs=algo_choice, outputs=memory_plot)
-                
-                # 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)
+            # Tab 1: Benchmark (Zero GPU)
             with gr.Tab("Benchmark", id="tab-benchmark"):
                 gr.Markdown("""
                 ## Attention Backend Benchmark