feat(gla): fused forward+backward kernels for g_gamma mode#122

Open

sii-xinglong wants to merge 16 commits intomainfrom

feat/chunk-gla-fused-kernels

Commits on Mar 31, 2026

feat(gla): add fused forward kernel for g_gamma mode

sii-xinglong
and
claude
committed
feat(gla): add fused backward kernel for g_gamma mode

sii-xinglong
and
claude
committed
style(chunk_fused_kernels): use exp() utility consistently in backward kernel

sii-xinglong
and
claude
committed
feat(gla): dispatch to fused kernels for g_gamma mode on TPU

sii-xinglong
and
claude
committed
bench(gla): add chunk_fused and chunk_fused_bwd benchmark providers

sii-xinglong
and
claude
committed
fix(chunk_fused_kernels): cast float32 scratch to h_ref dtype on save

sii-xinglong
and
claude
committed
fix(chunk_fused_kernels): make scale a static JIT argument

sii-xinglong
and
claude
committed
fix(chunk_fused_kernels): use parallel semantics for K/V grid dims

sii-xinglong
and
claude
committed
fix(chunk_fused_kernels): use f32 inputs for all matmuls in Pallas kernels

sii-xinglong
and
claude
committed
fix(gla): correct cu_seqlens parameter names in chunk kernel calls

sii-xinglong
and
claude
committed
fix(tests): use bfloat16 inputs for both fused and reference paths

sii-xinglong
and
claude
committed
fix(tests): use naive recurrent reference instead of non-fused Pallas

sii-xinglong
and
claude
committed
fix: scale g_gamma in E2E test to prevent numerical overflow

sii-xinglong
and
claude
committed
fix: address PR #122 review comments

sii-xinglong
and
claude
committed
perf(chunk_gla): eliminate jnp.flip copies in fused backward via reverse index_map

sii-xinglong
and
claude
committed
bench(gla): add memory profiling script for fused vs non-fused

sii-xinglong
and
claude
committed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

feat(gla): fused forward+backward kernels for g_gamma mode#122

feat(gla): fused forward+backward kernels for g_gamma mode#122
sii-xinglong wants to merge 16 commits intomainfrom
feat/chunk-gla-fused-kernels

Commits on Mar 31, 2026

feat(gla): add fused forward kernel for g_gamma mode

feat(gla): add fused backward kernel for g_gamma mode

style(chunk_fused_kernels): use exp() utility consistently in backward kernel

feat(gla): dispatch to fused kernels for g_gamma mode on TPU

bench(gla): add chunk_fused and chunk_fused_bwd benchmark providers

fix(chunk_fused_kernels): cast float32 scratch to h_ref dtype on save

fix(chunk_fused_kernels): make scale a static JIT argument

fix(chunk_fused_kernels): use parallel semantics for K/V grid dims

fix(chunk_fused_kernels): use f32 inputs for all matmuls in Pallas kernels

fix(gla): correct cu_seqlens parameter names in chunk kernel calls

fix(tests): use bfloat16 inputs for both fused and reference paths

fix(tests): use naive recurrent reference instead of non-fused Pallas

fix: scale g_gamma in E2E test to prevent numerical overflow

fix: address PR #122 review comments

perf(chunk_gla): eliminate jnp.flip copies in fused backward via reverse index_map

bench(gla): add memory profiling script for fused vs non-fused