fix bug when num_kv > 1

modeling_RW.py

@@ -290,8 +290,8 @@ class Attention(nn.Module):
 
         if alibi is None:
             query_layer_ = query_layer.reshape(batch_size, self.num_heads, -1, self.head_dim)
-            key_layer_ = key_layer.reshape(batch_size, self.num_kv, -1, self.head_dim)
-            value_layer_ = value_layer.reshape(batch_size, self.num_kv, -1, self.head_dim)
+            key_layer_ = key_layer.reshape(batch_size, self.num_heads, -1, self.head_dim)
+            value_layer_ = value_layer.reshape(batch_size, self.num_heads, -1, self.head_dim)
 
             attn_output = F.scaled_dot_product_attention(
                 query_layer_, key_layer_, value_layer_, None, 0.0, is_causal=True