[Perf] Simple GLA Performance Optimization

[0xaskr]

Summary

The current implementation of Simple GLA operations demonstrates performance bottlenecks under certain workloads. This issue aims to profile, analyze, and optimize performance.

Type

• Performance regression (was faster before)
• Below expected performance target (not meeting 80% roofline)
• Optimization opportunity

Kernel / Operation

Simple GLA operations under tops/ops/simple_gla/.

Observed Performance

Performance bottlenecks under certain workloads (exact numbers TBD via profiling).

Expected Performance

Concrete measurable speedup over current implementation, targeting 80% of hardware theoretical peak per project standards.

Environment

• Python version:
• JAX version:
• Hardware: CPU / GPU (model) / TPU (version)
• OS:

Reproduction

# TBD: benchmark script

Tasks

• Benchmark existing GLA operations and identify slow paths
• Profile code to pinpoint CPU or memory bottlenecks
• Enhance algorithm efficiency or parallelization strategies
• Evaluate impact of hardware features (e.g., SIMD, cache usage)
• Document improvement progress and test results

Acceptance Criteria

• Concrete measurable speedup over current implementation
• No regressions in accuracy or stability
• Code is properly tested and documented

Additional Context

Priority: P0.