Error while updating prompt

[Fqlox]

When I use stream.update_prompt(prompt)

I get :

ValueError: Incompatible shapes for attention inputs:
  query.shape: torch.Size([64, 4096, 40])
  key.shape  : torch.Size([32, 77, 40])
  value.shape: torch.Size([32, 77, 40])
HINT: We don't support broadcasting, please use `expand` yourself before calling `memory_efficient_attention` if you need to

tried without xFormers and got:

RuntimeError: The size of tensor a (8192) must match the size of tensor b (4096) at non-singleton dimension 1

How can I prevent this ?