Hi I don’t really understand when pytorch internally allocates memory / does copies.
I’m using the fairseq multihead attention code https://github.com/pytorch/fairseq/blob/master/fairseq/modules/multihead_attention.py to build a transformer model that can do incremental encoding. However, I notice encoding with incremental state is 5x slower than encoding the entire sequence at once.
My current hypothesis is its because of this line

  
      github.com
  
  
    pytorch/fairseq/blob/master/fairseq/modules/multihead_attention.py#L266

if saved_state is not None:    # saved states are stored with shape (bsz, num_heads, seq_len, head_dim)    if "prev_key" in saved_state:        _prev_key = saved_state["prev_key"]        assert _prev_key is not None        prev_key = _prev_key.view(bsz * self.num_heads, -1, self.head_dim)        if static_kv:            k = prev_key        else:            assert k is not None            k = torch.cat([prev_key, k], dim=1)    if "prev_value" in saved_state:        _prev_value = saved_state["prev_value"]        assert _prev_value is not None        prev_value = _prev_value.view(bsz * self.num_heads, -1, self.head_dim)        if static_kv:            v = prev_value        else:            assert v is not None            v = torch.cat([prev_value, v], dim=1)    prev_key_padding_mask: Optional[Tensor] = None



  
  
    
    
  
  


This line cats the new key row to the end of the previously saved key matrix. When this operation is done, does pytorch internally do a malloc, copy the old matrix over, and then copy the new row over? If so, is there a way to pre allocate space for rows like this to avoid a copy when cat is called?