(us, dts, As, Bs, Cs, h0, mask)
| 600 | mask = torch.tril(torch.ones((chunksize, chunksize), device=us.device), diagonal=0) |
| 601 | |
| 602 | def ss_chunk(us, dts, As, Bs, Cs, h0, mask): |
| 603 | # BHLD, BHLN, HND, BHDN |
| 604 | cL = us.shape[2] |
| 605 | _mask = (mask[:cL,:cL].contiguous() if cL < mask.shape[0] else mask).view(1, 1, cL, cL, 1) |
| 606 | |
| 607 | w_log = As[None, :, None, :, :] * (torch.cumsum(dts, dim=2)[..., None, :]) # (B, H, L, Dk, Dv) |
| 608 | v = us * dts # (B,H,L,Dv) |
| 609 | k = Bs # (B,H,L,Dk) |
| 610 | q = Cs # (B,H,L,Dk) |
| 611 | w = w_log.exp() |
| 612 | |
| 613 | k_div_w = k[..., None] / w |
| 614 | q_mul_w = q[..., None] * w |
| 615 | |
| 616 | # h0 independent ==================== |
| 617 | next_h_1 = w[:,:,-1] * torch.einsum("bhlkv,bhlv->bhkv", k_div_w, v) |
| 618 | y_1 = torch.einsum("bhlrv,bhrv->bhlv", torch.einsum("bhlkv,bhrkv->bhlrv", q_mul_w, k_div_w) * _mask, v) |
| 619 | |
| 620 | # h0 dependent ====================== |
| 621 | y_0 = torch.einsum("bhlkv,bhkv->bhlv", q_mul_w, h0) |
| 622 | next_h_0 = w[:,:, -1] * h0 |
| 623 | |
| 624 | next_h = next_h_0 + next_h_1 |
| 625 | y = y_0 + y_1 |
| 626 | |
| 627 | return y, next_h |
| 628 | |
| 629 | dtype = torch.float32 |
| 630 | # dtype = torch.float16 |
nothing calls this directly
no outgoing calls
no test coverage detected