| |
| |
| |
| |
| |
| |
|
|
| import torch |
|
|
| from ...ops.utils.cumsum import chunk_global_cumsum |
|
|
|
|
| class StatefulNormalizeFunction(torch.autograd.Function): |
| """Normalized linear-attention divide step with a hand-written backward. |
| |
| Forward computes the running key-cumsum z_t and divides ``o`` by ``<q_t, z_t>``. |
| Backward mirrors the pattern used by simple_gla / gla / gated_delta_rule |
| """ |
|
|
| @staticmethod |
| def forward(ctx, o, q, k, scale, z_init, reverse, cu_seqlens): |
| k_cum = chunk_global_cumsum(k, reverse=reverse, cu_seqlens=cu_seqlens, output_dtype=k.dtype) |
| if z_init is not None: |
| if cu_seqlens is not None: |
| seq_lens = (cu_seqlens[1:] - cu_seqlens[:-1]).to(torch.long) |
| z_init_b = z_init.squeeze(1).repeat_interleave(seq_lens, dim=0).unsqueeze(0) |
| else: |
| z_init_b = z_init |
| k_cum = k_cum + z_init_b |
| denom = (q * scale * k_cum).sum(-1, keepdim=True) + 1e-10 |
| o_out = o / denom |
| if cu_seqlens is not None: |
| idx = (cu_seqlens[:-1] if reverse else cu_seqlens[1:] - 1).to(torch.long) |
| z_state_out = k_cum[0, idx].unsqueeze(1).contiguous() |
| else: |
| z_state_out = (k_cum[:, :1] if reverse else k_cum[:, -1:]).contiguous() |
|
|
| ctx.save_for_backward(o, q, k_cum, denom) |
| ctx.scale = scale |
| ctx.reverse = reverse |
| ctx.cu_seqlens = cu_seqlens |
| ctx.has_z_init = z_init is not None |
| ctx.k_dtype = k.dtype |
| return o_out, z_state_out |
|
|
| @staticmethod |
| def backward(ctx, do, dz_state): |
| o, q, k_cum, denom = ctx.saved_tensors |
| scale = ctx.scale |
| reverse = ctx.reverse |
| cu_seqlens = ctx.cu_seqlens |
|
|
| d_o_in = do / denom |
| ddenom = -(do * o).sum(-1, keepdim=True) / (denom * denom) |
|
|
| dq = ddenom * scale * k_cum |
| dk_cum = ddenom * scale * q |
|
|
| if dz_state is not None: |
| if cu_seqlens is not None: |
| idx = (cu_seqlens[:-1] if reverse else cu_seqlens[1:] - 1).to(torch.long) |
| dk_cum[0].index_add_(0, idx, dz_state.squeeze(1).to(dk_cum.dtype)) |
| elif reverse: |
| dk_cum[:, :1].add_(dz_state.to(dk_cum.dtype)) |
| else: |
| dk_cum[:, -1:].add_(dz_state.to(dk_cum.dtype)) |
|
|
| dz_init = None |
| if ctx.has_z_init: |
| if cu_seqlens is not None: |
| seq_lens = (cu_seqlens[1:] - cu_seqlens[:-1]).to(torch.long) |
| seg_id = torch.repeat_interleave( |
| torch.arange(seq_lens.numel(), device=dk_cum.device), seq_lens, |
| ) |
| flat = dk_cum[0] |
| dz_init = torch.zeros( |
| seq_lens.numel(), flat.shape[1], flat.shape[2], |
| device=flat.device, dtype=flat.dtype, |
| ) |
| dz_init.index_add_(0, seg_id, flat) |
| dz_init = dz_init.unsqueeze(1) |
| else: |
| dz_init = dk_cum.sum(dim=1, keepdim=True) |
|
|
| dk = chunk_global_cumsum( |
| dk_cum, reverse=not reverse, cu_seqlens=cu_seqlens, output_dtype=ctx.k_dtype, |
| ) |
| return d_o_in, dq, dk, None, dz_init, None, None |
|
|
|
|
| def normalize_with_z_state( |
| o: torch.Tensor, |
| q: torch.Tensor, |
| k: torch.Tensor, |
| scale: float, |
| z_init: torch.Tensor | None, |
| reverse: bool, |
| cu_seqlens: torch.LongTensor | None, |
| ) -> tuple[torch.Tensor, torch.Tensor]: |
| """Apply running normalization Z_t = sum_{s<=t} k_s (or sum_{s>=t} for reverse). |
| |
| z_init carries the cumulative key from prior chunks; the returned z_state is the |
| boundary needed to chain into the next chunk. Varlen path broadcasts z_init and |
| extracts the boundary per-sequence. |
| """ |
| return StatefulNormalizeFunction.apply(o, q, k, scale, z_init, reverse, cu_seqlens) |
|
|
