| """Reference functions synchronized from PyTorch Helion. |
| |
| Generated by scripts/sync_from_helion.py. Do not edit directly. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import torch |
| import torch.nn.functional as F |
| from torch.utils.checkpoint import checkpoint |
|
|
|
|
| def rel_error(a: torch.Tensor | None, b: torch.Tensor | None) -> float: |
| """Relative L2 error between two tensors.""" |
| assert a is not None and b is not None |
| return (a.float() - b.float()).norm().item() / b.float().norm().clamp( |
| min=1e-8 |
| ).item() |
|
|
|
|
| def chunked_linear_attn_reference( |
| q: torch.Tensor, |
| k: torch.Tensor, |
| v: torch.Tensor, |
| g: torch.Tensor, |
| beta: torch.Tensor | None = None, |
| a: torch.Tensor | None = None, |
| C: int = 64, |
| initial_state: torch.Tensor | None = None, |
| ) -> torch.Tensor: |
| """ |
| Pure PyTorch reference for chunked linear attention forward. |
| |
| Supports both scalar decay (g: [B,H,T]) and diagonal decay (g: [B,H,T,D]). |
| For correction case (beta provided): uses sequential recurrence within chunks. |
| For no-correction case: uses parallel matmul-based intra-chunk computation. |
| """ |
| B, H, T, D = q.shape |
| DV = v.shape[-1] |
| N = T // C |
| assert T % C == 0 |
|
|
| input_dtype = q.dtype |
| q, k, v, g = q.float(), k.float(), v.float(), g.float() |
| if beta is not None: |
| beta = beta.float() |
| if a is not None: |
| a = a.float() |
|
|
| diagonal_decay = g.dim() == 4 |
|
|
| qc = q.reshape(B, H, N, C, D) |
| kc = k.reshape(B, H, N, C, D) |
| vc = v.reshape(B, H, N, C, DV) |
| if diagonal_decay: |
| gc = g.reshape(B, H, N, C, D) |
| else: |
| gc = g.reshape(B, H, N, C) |
|
|
| bc = beta.reshape(B, H, N, C) if beta is not None else None |
| ac = a.reshape(B, H, N, C, D) if a is not None else None |
|
|
| outputs = [] |
| if initial_state is not None: |
| state = initial_state.float().clone() |
| else: |
| state = q.new_zeros(B, H, D, DV, dtype=torch.float32) |
|
|
| def correction_chunk( |
| qi: torch.Tensor, |
| ki: torch.Tensor, |
| vi: torch.Tensor, |
| gi: torch.Tensor, |
| bi: torch.Tensor, |
| ai: torch.Tensor, |
| state: torch.Tensor, |
| ) -> tuple[torch.Tensor, torch.Tensor]: |
| """One chunk of the sequential beta-correction recurrence, returning |
| (chunk_output, new_state); tensors are explicit args so it can be |
| gradient-checkpointed.""" |
| s = state |
| chunk_out = [] |
| for j in range(C): |
| g_j = gi[:, :, j] |
| alpha = ( |
| torch.exp(g_j).unsqueeze(-1) |
| if diagonal_decay |
| else torch.exp(g_j).unsqueeze(-1).unsqueeze(-1) |
| ) |
| s = alpha * s |
| kj = ai[:, :, j] |
| vj = vi[:, :, j] |
| bj = bi[:, :, j].unsqueeze(-1).unsqueeze(-1) |
| kts = torch.einsum("bhd,bhdv->bhv", kj, s) |
| s = s - bj * torch.einsum("bhd,bhv->bhdv", kj, kts) |
| s = s + bj * torch.einsum("bhd,bhv->bhdv", kj, vj) |
| chunk_out.append(torch.einsum("bhd,bhdv->bhv", qi[:, :, j], s)) |
| return torch.stack(chunk_out, dim=2), s |
|
|
| for i in range(N): |
| qi = qc[:, :, i] |
| ki = kc[:, :, i] |
| vi = vc[:, :, i] |
| gi = gc[:, :, i] |
|
|
| bi = bc[:, :, i] if bc is not None else None |
| ai = ac[:, :, i] if ac is not None else ki |
|
|
| if bi is not None: |
| |
| |
| chunk_out, state = checkpoint( |
| correction_chunk, |
| qi, |
| ki, |
| vi, |
| gi, |
| bi, |
| ai, |
| state, |
| use_reentrant=False, |
| ) |
| outputs.append(chunk_out.to(input_dtype)) |
|
|
| elif diagonal_decay: |
| g_cs = gi.cumsum(-2) |
| g_last = g_cs[:, :, -1] |
|
|
| q_scaled = qi * torch.exp(g_cs) |
| k_scaled_intra = ki * torch.exp(-g_cs) |
|
|
| o_cross = torch.einsum("bhcd,bhdv->bhcv", q_scaled, state) |
|
|
| attn = torch.einsum("bhid,bhjd->bhij", q_scaled, k_scaled_intra) |
| causal_mask = torch.tril(torch.ones(C, C, device=q.device)) |
| attn = attn * causal_mask |
| o_intra = torch.einsum("bhij,bhjv->bhiv", attn, vi) |
|
|
| outputs.append((o_cross + o_intra).to(input_dtype)) |
|
|
| k_for_state = ki * torch.exp(g_last[:, :, None, :] - g_cs) |
| state = state * torch.exp(g_last).unsqueeze(-1) + torch.einsum( |
| "bhcd,bhcv->bhdv", k_for_state, vi |
| ) |
|
|
| else: |
| g_cs = gi.cumsum(-1) |
| g_last = g_cs[:, :, -1] |
|
|
| o_cross = torch.einsum("bhcd,bhdv->bhcv", qi, state) |
| o_cross = o_cross * torch.exp(g_cs).unsqueeze(-1) |
|
|
| attn = torch.einsum("bhid,bhjd->bhij", qi, ki) |
| decay_mask = torch.exp(g_cs.unsqueeze(-1) - g_cs.unsqueeze(-2)) |
| causal_mask = torch.tril(torch.ones(C, C, device=q.device)) |
| attn = attn * decay_mask * causal_mask |
| o_intra = torch.einsum("bhij,bhjv->bhiv", attn, vi) |
|
|
| outputs.append((o_cross + o_intra).to(input_dtype)) |
|
|
| k_scaled = ki * torch.exp(g_last[..., None, None] - g_cs.unsqueeze(-1)) |
| state = state * torch.exp(g_last).unsqueeze(-1).unsqueeze( |
| -1 |
| ) + torch.einsum("bhcd,bhcv->bhdv", k_scaled, vi) |
|
|
| return torch.stack(outputs, dim=2).reshape(B, H, T, DV).to(input_dtype) |
|
|
|
|
| def naive_recurrent_reference( |
| q: torch.Tensor, |
| k: torch.Tensor, |
| v: torch.Tensor, |
| decay: torch.Tensor, |
| beta: torch.Tensor | None = None, |
| q_scale: float = 1.0, |
| A_log: torch.Tensor | None = None, |
| dt_bias: torch.Tensor | None = None, |
| lower_bound: float | None = None, |
| use_qk_l2norm_in_kernel: bool = False, |
| use_gate_in_kernel: bool = False, |
| use_beta_sigmoid_in_kernel: bool = False, |
| cu_seqlens: torch.Tensor | None = None, |
| ) -> torch.Tensor: |
| """Step-by-step recurrent computation. Slow but correct. |
| |
| The optional arguments apply the input preamble first, so the reference can be |
| fed the same pre-activation tensors a model produces: |
| use_qk_l2norm_in_kernel: q, k = q / ||q||, k / ||k|| |
| use_beta_sigmoid_in_kernel: beta = sigmoid(beta) |
| use_gate_in_kernel: decay = lower_bound * sigmoid(exp(A_log) * |
| (decay + dt_bias)), or -exp(A_log) * |
| softplus(...) when lower_bound is None |
| The names match the flags on the kernels, so a caller forwards one dict to both. |
| |
| cu_seqlens [N+1] marks sequence boundaries in a varlen batch (B == 1), where |
| sequence n spans tokens cu_seqlens[n] : cu_seqlens[n+1]. The state resets to |
| zero at each boundary, so no key from one sequence reaches another's output. |
| """ |
| if use_qk_l2norm_in_kernel: |
| q = F.normalize(q.float(), dim=-1).to(q.dtype) |
| k = F.normalize(k.float(), dim=-1).to(k.dtype) |
| if use_gate_in_kernel: |
| assert A_log is not None |
| H_, D_ = decay.shape[1], decay.shape[-1] |
| gb = decay.float() |
| if dt_bias is not None: |
| gb = gb + dt_bias.view(1, H_, 1, D_) |
| a = torch.exp(A_log.float()).view(1, H_, 1, 1) |
| if lower_bound is not None: |
| decay = lower_bound * torch.sigmoid(a * gb) |
| else: |
| decay = -a * F.softplus(gb) |
| if use_beta_sigmoid_in_kernel: |
| assert beta is not None |
| beta = torch.sigmoid(beta.float()).to(beta.dtype) |
|
|
| B, H, T, D = q.shape |
| DV = v.shape[-1] |
| S = q.new_zeros(B, H, D, DV, dtype=torch.float32) |
| outputs = [] |
| diagonal = decay.dim() == 4 |
| starts = set() if cu_seqlens is None else set(cu_seqlens[:-1].tolist()) |
|
|
| for t in range(T): |
| if t in starts: |
| S = torch.zeros_like(S) |
| qt = q[:, :, t].float() * q_scale |
| kt = k[:, :, t].float() |
| vt = v[:, :, t].float() |
| bt = beta[:, :, t] if beta is not None else None |
|
|
| if diagonal: |
| alpha = torch.exp(decay[:, :, t]) |
| S = alpha.unsqueeze(-1) * S |
| else: |
| alpha = torch.exp(decay[:, :, t]) |
| S = alpha.unsqueeze(-1).unsqueeze(-1) * S |
|
|
| if bt is not None: |
| kts = torch.einsum("bhd,bhdv->bhv", kt, S) |
| S = S - bt.unsqueeze(-1).unsqueeze(-1) * torch.einsum( |
| "bhd,bhv->bhdv", kt, kts |
| ) |
| S = S + bt.unsqueeze(-1).unsqueeze(-1) * torch.einsum( |
| "bhd,bhv->bhdv", kt, vt |
| ) |
| else: |
| S = S + torch.einsum("bhd,bhv->bhdv", kt, vt) |
|
|
| out_t = torch.einsum("bhd,bhdv->bhv", qt, S) |
| outputs.append(out_t) |
|
|
| return torch.stack(outputs, dim=2) |
|
|
|
|
| def recurrent_step_reference( |
| q: torch.Tensor, |
| k: torch.Tensor, |
| v: torch.Tensor, |
| state: torch.Tensor, |
| alpha: float | torch.Tensor = 1.0, |
| beta_val: torch.Tensor | None = None, |
| a_val: torch.Tensor | None = None, |
| ) -> tuple[torch.Tensor, torch.Tensor]: |
| """Single-step recurrent update (pure PyTorch reference). |
| |
| Returns (output, new_state). |
| """ |
| k_sq = k.squeeze(2) |
| v_sq = v.squeeze(2) |
| q_sq = q.squeeze(2) |
|
|
| if isinstance(alpha, torch.Tensor): |
| alpha_sq = alpha.squeeze(2) |
| if alpha_sq.dim() == 3: |
| state = alpha_sq.unsqueeze(-1) * state |
| else: |
| state = alpha_sq.unsqueeze(-1).unsqueeze(-1) * state |
| else: |
| state = alpha * state |
|
|
| if beta_val is not None: |
| b = beta_val.squeeze(2) |
| a_sq = a_val.squeeze(2) if a_val is not None else k_sq |
|
|
| if b.dim() == 3: |
| r = b.shape[-1] |
| for j in range(r): |
| bj = b[:, :, j].unsqueeze(-1).unsqueeze(-1) |
| aj = a_sq[:, :, j] |
| kts = torch.einsum("bhd,bhdv->bhv", aj, state) |
| state = state - bj * torch.einsum("bhd,bhv->bhdv", aj, kts) |
| state = state + bj * torch.einsum("bhd,bhv->bhdv", aj, v_sq) |
| else: |
| b_exp = b.unsqueeze(-1).unsqueeze(-1) |
| kts = torch.einsum("bhd,bhdv->bhv", a_sq, state) |
| state = state - b_exp * torch.einsum("bhd,bhv->bhdv", a_sq, kts) |
| state = state + b_exp * torch.einsum("bhd,bhv->bhdv", a_sq, v_sq) |
| else: |
| state = state + torch.einsum("bhd,bhv->bhdv", k_sq, v_sq) |
|
|
| o = torch.einsum("bhd,bhdv->bhv", q_sq, state) |
| return o.unsqueeze(2), state |
|
|
