Intermediate Materialization Avoidance

[thanay-sisir]

⚡ Optimization Summary: compute_synchronisation

1. Technical Mechanism

• Vectorization: Refactors the operation from a dense $O(N^2)$ outer product to a sparse, index-driven Hadamard product using torch.triu_indices.
• Memory Efficiency: Eliminates the intermediate $(B, N, N)$ tensor allocation, reducing auxiliary memory complexity from $O(B \cdot N^2)$ to $O(1)$.
• Device Awareness: Enforces explicit device placement for indices to prevent implicit host-to-device transfer overhead.

2. Stability & Scalability

• Prevents "Quadratic Trap": Removes the bottleneck where memory usage scales quadratically with $d_{model}$, preventing OOM errors on larger model configurations.
• Eliminates Compute Waste: Bypasses the calculation of the symmetric lower triangle, saving $\approx 50%$ of FLOPs and memory writes previously wasted on discarded data.

[lukedarlow]

This PR does follow the rules.

However, I would like you to add commentary above these altered lines that shows the old version, explaining their equivalence.

The reason for this is simply to aid readers in understanding what the code is actually doing in relation to the paper.

[thanay-sisir]

@lukedarlow when you are going to change that luke.......
Try to merge it as you become free ....😊
Thanks in advance !!!

[lukedarlow]

I already requested changes from you.

[thanay-sisir]

@lukedarlow ok luke how about that ....?