Upload modeling_rwkv6qwen2.py with huggingface_hub

modeling_rwkv6qwen2.py

@@ -204,6 +204,7 @@ class RWKV6State(Cache):
     #         self.value_cache[layer_idx] = self.value_cache[layer_idx][indices, ...]
 
 from fla.ops.gla.chunk import chunk_gla
+from fla.ops.gla.fused_recurrent import fused_recurrent_gla
 
 class RWKV6Attention(nn.Module):
     def __init__(self, config, layer_idx: Optional[int] = None):
@@ -360,8 +361,8 @@ class RWKV6Attention(nn.Module):
         scale = query_states.shape[-1] ** -0.5
         output_final_state = not self.training and use_cache and past_key_value is not None
         #attn_output, output_kv_state = ChunkGLAFunction.apply(query_states, key_states, value_states, decay_states_log.float(), scale, input_kv_state, output_final_state)
-        attn_output, output_kv_state = chunk_gla(query_states, key_states, value_states, decay_states_log, scale, input_kv_state, output_final_state)
-        #attn_output, output_kv_state = fused_recurrent_gla(query_states, key_states, value_states, decay_states_log, input_kv_state, output_final_state)
+        #attn_output, output_kv_state = chunk_gla(query_states, key_states, value_states, decay_states_log, scale, input_kv_state, output_final_state)
+        attn_output, output_kv_state = fused_recurrent_gla(query_states, key_states, value_states, decay_states_log, None, scale, input_kv_state, output_final_state)
 
         if output_final_state:
             past_key_value.update(output_kv_state, output_shift_state, q_len, self.layer_idx)
@@ -1207,4 +1208,4 @@ class RWKV6Qwen2ForQuestionAnswering(RWKV6Qwen2PreTrainedModel):
             end_logits=end_logits,
             hidden_states=outputs.hidden_states,
             attentions=outputs.attentions,
-        )
+        )