[Simple GLA] Do not pass head_first in the layer

yzhangcs · yzhangcs · commit 3a7ecbfcbd81 · 2025-04-13T15:43:28.000+08:00
diff --git a/fla/layers/simple_gla.py b/fla/layers/simple_gla.py
@@ -217,7 +217,6 @@ def forward(
                 initial_state=recurrent_state,
                 output_final_state=use_cache,
                 cu_seqlens=cu_seqlens,
-                head_first=False
             )
         elif mode == 'fused_recurrent':
             o, recurrent_state = fused_recurrent_simple_gla(

Original file line number	Diff line number	Diff line change
`@@ -217,7 +217,6 @@ def forward(`
`217`	`217`	`initial_state=recurrent_state,`
`218`	`218`	`output_final_state=use_cache,`
`219`	`219`	`cu_seqlens=cu_seqlens,`
`220`		`- head_first=False`
`221`	`220`	`)`
`222`	`221`	`elif mode == 'fused_recurrent':`
`223`	`222`	`o, recurrent_state = fused_recurrent_simple_gla(`