Fix the kv-cache dimensions

Hello!
I have noticed that the dimension of the kv-cache here is weird, and does not match the hugginface transformers modeling_bloom.py file.
Is the departure from the bloom dimension intended?
Judging from the copy-pasted comments, it looks like a bug - also, `_convert_to_rw_cache` & its `_convert_to_standard_cache` counterpart matches bloom dimensions.

modelling_RW.py

@@ -271,7 +271,7 @@ class Attention(nn.Module):
             # concatenate along seq_length dimension:
             #  - key: [batch_size * self.num_heads, head_dim, kv_length]
             #  - value: [batch_size * self.num_heads, kv_length, head_dim]
-            key_layer = torch.cat((past_key, key_layer), dim=1)
+            key_layer = torch.cat((past_key, key_layer), dim=2)
             value_layer = torch.cat((past_value, value_layer), dim=1)
 
         _, kv_length, _ = key_layer.shape