Divide FLOPs by two due to causal masks

[NuojCheng]

Description

As pointed out by this issue, the current FLOPs calculation does not consider the impact from causal masks. This PR fixes this issue.

FIXES: b/431892390

Tests

Model	Max Target Length	TFLOPs Type	Before	After
llama2-7b	8192	Total	1721.22	1510.11
		Learnable Weight	1299.00	1299.00
		Attention	422.21	211.11
deepseek2-16b	8192	Total	704.60	593.27
		Learnable Weight	481.95	481.95
		Attention	222.65	111.33
llama3-8b	8192	Total	1897.69	1686.58
		Learnable Weight	1475.48	1475.48
		Attention	422.21	211.11
llama4-17b-16e	8192	Total	3964.49	3568.67
		Learnable Weight	3172.84	3172.84
		Attention	791.65	395.82

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]

Could you run some quick tests similar like this PR, so we could have some idea on flops impact?

[RissyRan]

LGTM! @hengtaoguo @aireenmei @gagika you will need to add a config to opt out this calculation for bi-directional mask in Llama4

[RissyRan]

Shall we keep this as one line for easy copy/paste?

[NuojCheng]

I tried but it won't pass linter test. Maybe using short-URL? Does google have this function?

[gobbleturk]

ya I'd keep as one line, just add

pylint: disable=line-too-long

to the previous line

[gobbleturk]

Thanks noujin!