fix gla

[czy1101]

PR Category

Type of Change

Description

Issue

Progress

• Change is properly reviewed (1 reviewer required, 2 recommended).
• Change is responded to an issue.
• Change is fully covered by a UT.

Performance

Benchmark

Chunk GLA Benchmark — Forward

Provider: FLA Chunk vs FlagGems Chunk (Forward)
dtype: torch.bfloat16

B	T	H	D	dtype	fla(ms)	gems(ms)	gems/fla
1	4096	32	512	bfloat16	2.441	2.068	1.18x
2	2048	16	512	bfloat16	1.243	1.054	1.18x
1	8192	96	128	bfloat16	2.491	1.957	1.27x
2	16384	16	128	bfloat16	1.893	1.541	1.23x
4	2048	16	128	bfloat16	0.454	0.365	1.24x
4	4096	64	128	bfloat16	3.312	2.600	1.27x
8	2048	32	256	bfloat16	3.711	3.046	1.22x
2	2048	16	512	bfloat16	1.244	1.053	1.18x
4	1024	8	512	bfloat16	0.653	0.550	1.19x
8	1024	8	64	bfloat16	0.288	0.241	1.20x

Chunk GLA Benchmark — Forward + Backward

Provider: FLA Chunk vs FlagGems Chunk (Forward + Backward)
dtype: torch.bfloat16

B	T	H	D	dtype	fla(ms)	gems(ms)	gems/fla
1	4096	32	512	bfloat16	11.833	11.466	1.03x
2	2048	16	512	bfloat16	5.975	5.785	1.03x
1	8192	96	128	bfloat16	12.184	11.638	1.05x
2	16384	16	128	bfloat16	8.726	8.371	1.04x
4	2048	16	128	bfloat16	2.150	2.064	1.04x
4	4096	64	128	bfloat16	16.327	15.615	1.05x
8	2048	32	256	bfloat16	18.400	17.739	1.04x
2	2048	16	512	bfloat16	5.978	5.792	1.03x
4	1024	8	512	bfloat16	3.082	2.979	1.03x
8	1024	8	64	bfloat16	1.144	1.022	1.12x

Performance Summary

Mode	Average Speedup (gems/fla)	Maximum Speedup	Minimum Speedup
Forward	~1.21x	1.27x	1.18x
Forward + Backward	~1.04x	1.12x	1.03x

Observations

• Forward kernel achieves consistent performance improvement:

  • Speedup range: 1.18x ~ 1.27x
  • Average improvement: ~21%

• Forward + Backward overall improvement is smaller:

  • Speedup range: 1.03x ~ 1.12x
  • Average improvement: ~4%

pytest benchmark/test_FLA/test_chunk_gla_perf.py::test_perf_chunk_gla -s --level core --iter 200 --warmup 100