[Layers] update mode assignment for GDN family layers

Author: sustcsonglin

fla/layers/comba.py

@@ -222,10 +222,9 @@ def forward(
 
         batch_size, q_len, _ = hidden_states.shape
         # change to inference mode.
-        mode = 'fused_recurrent' if q_len <= 64 else self.mode
+        mode = 'fused_recurrent' if (q_len <= 64 and not self.training) else self.mode
         if self.training:
             assert mode == 'chunk', "Only chunk mode is supported in training."
-
         last_state = None
         if past_key_values is not None and len(past_key_values) > self.layer_idx:
             last_state = past_key_values[self.layer_idx]

fla/layers/gated_deltanet.py

@@ -217,7 +217,7 @@ def forward(
 
         batch_size, q_len, _ = hidden_states.shape
         # change to inference mode.
-        mode = 'fused_recurrent' if q_len <= 64 else self.mode
+        mode = 'fused_recurrent' if (q_len <= 64 and not self.training) else self.mode
         if self.training:
             assert mode == 'chunk', "Only chunk mode is supported in training."
 

fla/layers/gated_deltaproduct.py

@@ -174,8 +174,7 @@ def forward(
 
         batch_size, q_len, _ = hidden_states.shape
         # change to inference mode.
-        mode = 'fused_recurrent' if q_len <= 64 else self.mode
-
+        mode = 'fused_recurrent' if (q_len <= 64 and not self.training) else self.mode
         if self.training:
             assert mode == 'chunk', "Only chunk mode is supported in training."
 

fla/layers/kda.py

@@ -172,7 +172,7 @@ def forward(
 
         batch_size, q_len, _ = hidden_states.shape
         # change to inference mode.
-        mode = 'fused_recurrent' if q_len <= 64 and not self.training else self.mode
+        mode = 'fused_recurrent' if (q_len <= 64 and not self.training) else self.mode
         if self.training:
             assert mode == 'chunk', "Only chunk mode is supported in training."
 

fla/layers/mom.py

@@ -437,7 +437,7 @@ def forward(
         if origin_cu_seqlens is not None:
             hidden_states, attention_mask = self.cu2pad(hidden_states, origin_cu_seqlens)
 
-        mode = 'fused_recurrent' if hidden_states.shape[1] <= 64 else self.mode
+        mode = 'fused_recurrent' if (hidden_states.shape[1] <= 64 and not self.training) else self.mode
         if self.training:
             assert mode == 'chunk', "Only chunk mode is supported in training."