Upload folder using huggingface_hub

jet_block.py

@@ -46,7 +46,7 @@ class JetBlockConfig():
     head_dim: int = 256
     norm_eps: float = 1e-5
     conv_size: int = 4
-    dconv_generator_reduction: int = None
+    dconv_generator_reduction: int = 8
     dconv_implementation: str = 'triton'
 
 
@@ -180,24 +180,21 @@ class JetBlock(nn.Module):
         if attention_mask is not None and q_len > 1:
             indices, cu_seqlens, _ = get_unpad_data(attention_mask[:, -q_len:])
 
-        conv_state = None
-
         conv_mask = attention_mask[:, -hidden_states.shape[1]:] if attention_mask is not None else None
 
         q = F.silu(self.q_proj(hidden_states))
         k = F.silu(self.k_proj(hidden_states))
 
-        conv_state_v = None
+        conv_state = None
         if last_state is not None:
-            conv_state_v = last_state['conv_state'][-1]
-        v, conv_state_v = self.dynamic_conv1d(
+            conv_state = last_state['conv_state']
+        v, conv_state = self.dynamic_conv1d(
             x=self.v_proj(hidden_states),
             generator_input=hidden_states,
             mask=conv_mask,
-            cache=conv_state_v,
+            cache=conv_state,
             output_final_state=use_cache,
         )
-        conv_state = conv_state + (conv_state_v,) if conv_state is not None else (conv_state_v,)
 
         if attention_mask is not None and q_len > 1:
             q = index_first_axis(rearrange(q, "b s ... -> (b s) ..."), indices).unsqueeze(0)