cache: (batch, seq_len, 2, kv_heads, head_dim)
(cache: Optional[Tensor])
| 17 | CACHE_FILL_VALUE = -1 |
| 18 | |
| 19 | def get_cache_len(cache: Optional[Tensor]) -> int: |
| 20 | """ |
| 21 | cache: (batch, seq_len, 2, kv_heads, head_dim) |
| 22 | """ |
| 23 | if cache is None: |
| 24 | return 0 |
| 25 | nonzeros = T.any(cache.flatten(2) != CACHE_FILL_VALUE, dim=-1) |
| 26 | length = nonzeros.sum(dim=-1).int() |
| 27 | assert T.all(length == length[0]) |
| 28 | return length[0] |
| 29 | |
| 30 | |
| 31 | def rotate_half(x): |