|
|
|
@ -81,8 +81,6 @@ class TransformerBackend(ModuleBackend):
|
|
|
|
|
assert new_k.shape[0] == past_k.shape[0] and new_v.shape[0] == past_v.shape[0]
|
|
|
|
|
assert new_k.shape[1] == new_length and new_v.shape[1] == new_length
|
|
|
|
|
assert new_k.shape[2:] == past_k.shape[2:] and new_v.shape[2:] == past_v.shape[2:]
|
|
|
|
|
assert torch.allclose(new_v[:, : past_v.shape[1]], past_v)
|
|
|
|
|
assert torch.allclose(new_k[:, : past_k.shape[1]], past_k)
|
|
|
|
|
cache[0, :, prefix_length:new_length, :] = new_k[:, prefix_length:new_length]
|
|
|
|
|
cache[1, :, prefix_length:new_length, :] = new_v[:, prefix_length:new_length]
|
|
|
|
|
return (hidden_states,)
|
|
|
|
|