[Token Mixing] Remove the head_first arg from token mixing layers (#347)

Author: yzhangcs

fla/layers/abc.py

@@ -194,7 +194,6 @@ def forward(
             s=s,
             initial_state=recurrent_state,
             output_final_state=use_cache,
-            head_first=False
         )
         if past_key_values is not None:
             past_key_values.update(

fla/layers/based.py

@@ -54,21 +54,19 @@ def forward(self, hidden_states: torch.Tensor, **kwargs):
         q, k, v = map(lambda x: rearrange(x, "... (h d) -> ... h d", d=self.head_dim), [q, k, v])
         if mode == "fused_chunk":
             q, k = self.feature_map(q), self.feature_map(k)
-            o, _ = fused_chunk_linear_attn(q, k, v, normalize=True, scale=1, head_first=False)
+            o, _ = fused_chunk_linear_attn(q, k, v, normalize=True, scale=1)
         elif mode == 'chunk':
             q, k = self.feature_map(q), self.feature_map(k)
-            o, _ = chunk_linear_attn(q, k, v, normalize=True, scale=1, head_first=False)
+            o, _ = chunk_linear_attn(q, k, v, normalize=True, scale=1)
         elif mode == 'parallel':
             assert q.shape[-1] <= 128
-            o = parallel_based(q, k, v, scale=1, use_norm=True, head_first=False)
+            o = parallel_based(q, k, v, scale=1, use_norm=True)
         o = rearrange(o, 'b t h d -> b t (h d)')
         o = self.o_proj(o)
         o = self.dropout(o)
         return o
 
-    # https://github.com/HazyResearch/zoology/blob/main/zoology/mixers/based.py#L119
-
-    def forward_reference(self, hidden_states: torch.Tensor, filters: torch.Tensor = None, *args, **kwargs):
+    def forward_reference(self, hidden_states: torch.Tensor, **kwargs):
         """
         x (torch.Tensor): tensor of shape (b, d, t)
         y (torch.Tensor): tensor of shape (b, d, t)

fla/layers/delta_net.py

@@ -249,7 +249,6 @@ def forward(
                 initial_state=recurrent_state,
                 output_final_state=use_cache,
                 cu_seqlens=cu_seqlens,
-                head_first=False,
                 use_qk_l2norm_in_kernel=True if self.qk_norm == 'l2' else False
             )
         elif mode == 'chunk':
@@ -261,7 +260,6 @@ def forward(
                 initial_state=recurrent_state,
                 output_final_state=use_cache,
                 cu_seqlens=cu_seqlens,
-                head_first=False,
                 use_qk_l2norm_in_kernel=True if self.qk_norm == 'l2' else False
             )
         else:

fla/layers/gated_deltanet.py

@@ -256,7 +256,6 @@ def forward(
                 initial_state=recurrent_state,
                 output_final_state=use_cache,
                 cu_seqlens=cu_seqlens,
-                head_first=False,
                 use_qk_l2norm_in_kernel=True
             )
         elif mode == 'fused_recurrent':
@@ -269,7 +268,6 @@ def forward(
                 initial_state=recurrent_state,
                 output_final_state=use_cache,
                 cu_seqlens=cu_seqlens,
-                head_first=False,
                 use_qk_l2norm_in_kernel=True
             )
         else:

fla/layers/gated_deltaproduct.py

@@ -307,7 +307,6 @@ def forward(
                     initial_state=recurrent_state,
                     output_final_state=use_cache,
                     cu_seqlens=offsets,
-                    head_first=False,
                     use_qk_l2norm_in_kernel=True
                 )
             else:
@@ -319,7 +318,6 @@ def forward(
                     initial_state=recurrent_state,
                     output_final_state=use_cache,
                     cu_seqlens=offsets,
-                    head_first=False,
                     use_qk_l2norm_in_kernel=True
                 )
         else:

fla/layers/gla.py

@@ -235,7 +235,6 @@ def forward(
                 initial_state=recurrent_state,
                 output_final_state=use_cache,
                 cu_seqlens=cu_seqlens,
-                head_first=False
             )
         elif mode == 'fused_chunk':
             o, recurrent_state = fused_chunk_gla(
@@ -245,7 +244,6 @@ def forward(
                 g=gk,
                 initial_state=recurrent_state,
                 output_final_state=use_cache,
-                head_first=False
             )
         elif mode == 'chunk':
             o, recurrent_state = chunk_gla(
@@ -256,7 +254,6 @@ def forward(
                 initial_state=recurrent_state,
                 output_final_state=use_cache,
                 cu_seqlens=cu_seqlens,
-                head_first=False
             )
         else:
             raise NotImplementedError(f"Not supported mode `{mode}`.")

fla/layers/gsa.py

@@ -190,7 +190,6 @@ def forward(
                 output_final_state=use_cache,
                 scale=self.scale,
                 cu_seqlens=cu_seqlens,
-                head_first=False
             )
         elif mode == 'chunk':
             o, recurrent_state = chunk_gsa(
@@ -203,7 +202,6 @@ def forward(
                 output_final_state=use_cache,
                 scale=self.scale,
                 cu_seqlens=cu_seqlens,
-                head_first=False
             )
         else:
             raise NotImplementedError(f"Not supported mode `{mode}`.")

fla/layers/hgrn2.py

@@ -162,7 +162,6 @@ def forward(
                 initial_state=recurrent_state,
                 output_final_state=use_cache,
                 cu_seqlens=cu_seqlens,
-                head_first=False
             )
         elif mode == 'fused_chunk':
             o, recurrent_state = fused_chunk_gla(
@@ -172,7 +171,6 @@ def forward(
                 g=g,
                 initial_state=recurrent_state,
                 output_final_state=use_cache,
-                head_first=False
             )
         elif mode == 'chunk':
             o, recurrent_state = chunk_gla(
@@ -183,7 +181,6 @@ def forward(
                 initial_state=recurrent_state,
                 output_final_state=use_cache,
                 cu_seqlens=cu_seqlens,
-                head_first=False
             )
         else:
             raise NotImplementedError(f"Not supported mode `{mode}`.")

fla/layers/lightnet.py

@@ -168,7 +168,6 @@ def forward(
                 initial_state=recurrent_state,
                 output_final_state=use_cache,
                 cu_seqlens=cu_seqlens,
-                head_first=False
             )
         elif mode == 'chunk':
             o, recurrent_state = chunk_gla(
@@ -179,7 +178,6 @@ def forward(
                 initial_state=recurrent_state,
                 output_final_state=use_cache,
                 cu_seqlens=cu_seqlens,
-                head_first=False
             )
         else:
             raise NotImplementedError(f"Not supported mode `{mode}`.")

fla/layers/linear_attn.py

@@ -142,7 +142,6 @@ def forward(
                 k=k,
                 v=v,
                 normalize=self.do_feature_map_norm,
-                head_first=False
             )
         elif mode == 'fused_chunk':
             o, final_state = fused_chunk_linear_attn(