Fix Llama4 attention flops

[gagika]

Description

Fix Llama4 attention flops to only count for chunk attention window size for chunk attention layers.

Tests

Comparing total_tflops, learnable_weight_tflops, attention_tflops before and after the fix:

https://diff.googleplex.com/#key=wo0bVmy9wbUc

learnable_weight_tflops are matching
attention_tflops are matching for context length <= 8192
For context length 2* 8192, old attention flops was increasing 4 x over 8192, now
it's increasing 10/4 x (3 chunk attention increasing 2x and 1 global attention increasing 4x).
1187* 10/4 = 2968

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed.

[RissyRan]

LGTM

[gobbleturk]

Can we separate out chunked attention flops into its own method?

[gagika]

done

[gobbleturk]

Thanks Gagik!