"""
Copyright (c) 2025 by FlashInfer team.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

  http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""

from typing import Optional

import torch
import torch.nn.functional as F


def exclusive_cumsum(a: list[int]):
    r = [0]
    for v in a:
        r.append(r[-1] + v)
    return r


def matmul(a: torch.Tensor, b: torch.Tensor):
    assert a.dtype in [torch.float16, torch.bfloat16, torch.float32, torch.float64]
    assert b.dtype in [torch.float16, torch.bfloat16, torch.float32, torch.float64]
    if (
        a.dtype == torch.float16
        or b.dtype == torch.float16
        or a.dtype == torch.bfloat16
        or b.dtype == torch.bfloat16
    ):
        a_f32 = a.to(torch.float32)
        b_f32 = b.to(torch.float32)
        c_f32 = a_f32 @ b_f32
        if a.dtype == torch.bfloat16:
            return c_f32
        else:
            return c_f32.to(torch.float16)
    else:
        return a @ b


def LambdaQ(decay_factor, valid_nrows, block_size, device, offset=0):
    e = (
        F.pad(
            torch.arange(valid_nrows, device=device) + offset,
            (0, block_size - valid_nrows),
        )
        .unsqueeze(1)
        .unsqueeze(0)
    )
    return torch.pow(decay_factor, e)


def LambdaK(decay_factor, valid_nrows, block_size, device, offset=0):
    # NOTE: IT IS valid_nrows - ..., NOT block_size - ..., this is crucial for tail blocks
    e = (
        (
            (valid_nrows - offset)
            - F.pad(
                torch.arange(valid_nrows, device=device),
                (0, block_size - valid_nrows),
                value=block_size,
            )
        )
        .unsqueeze(1)
        .unsqueeze(0)
    )
    return torch.pow(decay_factor, e)


# sequence/block level linear attention
def _linear_attention(
    q: torch.Tensor,  # [seq_len, num_heads, head_size]
    k: torch.Tensor,  # [seq_len, num_heads, head_size]
    v: torch.Tensor,  # [seq_len, num_heads, head_size]
    *,
    decay_factor: torch.Tensor | None = None,
    qk_weight: torch.Tensor | None = None,
) -> torch.Tensor:
    # Compute Q @ K^T
    num_qo_heads = q.shape[1]
    num_kv_heads = k.shape[1]
    assert num_qo_heads == num_kv_heads
    q = q.transpose(0, 1)
    k = k.transpose(0, 1)
    v = v.transpose(0, 1)

    # print(q.shape, k.shape, v.shape)
    scores = matmul(q, k.transpose(-2, -1))

    # Create causal mask
    seq_len = q.size(-2)
    mask = torch.tril(
        torch.ones(num_qo_heads, seq_len, seq_len, dtype=q.dtype, device=q.device)
    )
    if decay_factor is not None and (decay_factor != 1.0).any():
        _, sq, sk = mask.shape
        with torch.device(q.device):
            e = (
                torch.arange(sq).unsqueeze(1) - torch.arange(sk).unsqueeze(0)
            ).unsqueeze(0)
            M = torch.pow(decay_factor, e)
            M[mask == 0.0] = 0.0
    elif qk_weight is not None:
        M = qk_weight.clone()
        M[mask == 0.0] = 0.0
    else:
        M = mask

    # Apply mask (Q @ K^T \odot M)
    masked_scores = scores * M

    # Apply to values (Q @ K^T \odot M) V
    out = matmul(masked_scores, v)
    out = out.transpose(0, 1)

    return out


@torch.inference_mode
def blockwise_linear_attention(
    q: torch.Tensor,  # [total_seq_len, num_qo_heads, head_size]
    k: torch.Tensor,  # [total_seq_len, num_kv_heads, head_size]
    v: torch.Tensor,  # [total_seq_len, num_kv_heads, head_size]
    seq_lens: list[int],  # sequence length for each sequence
    block_size: int = 32,
    scale_factor=1.0,
    decay_factor: float
    | torch.Tensor = 1.0,  # float or tensor with num_elems == num_qo_heads
    decay_exponent_offset=0,
    state_dtype: torch.dtype = torch.float32,
) -> torch.Tensor:
    num_qo_heads = q.size(1)
    head_size = q.size(2)
    num_kv_heads = k.size(1)

    if scale_factor != 1.0:
        k = k * scale_factor
    if isinstance(decay_factor, float):
        decay_factor = torch.ones(num_qo_heads) * decay_factor
        decay_factor = decay_factor.to(q.device)
    assert decay_factor.numel() == num_qo_heads
    decay_factor = decay_factor.reshape(num_qo_heads, 1, 1)

    k = k.repeat_interleave(num_qo_heads // num_kv_heads, dim=1)
    v = v.repeat_interleave(num_qo_heads // num_kv_heads, dim=1)

    KVs = []  # FIXME: kernel debug only
    kv = torch.zeros(
        (len(seq_lens), num_qo_heads, head_size, head_size),
        dtype=state_dtype,
        device=q.device,
    )
    output = torch.zeros_like(q)

    seq_offset = exclusive_cumsum(seq_lens)
    for seq_idx, seq_start in enumerate(seq_offset[:-1]):
        seq_end = seq_offset[seq_idx + 1]
        blk_offset = seq_start
        carried_kv = torch.zeros(
            (num_qo_heads, head_size, head_size), dtype=state_dtype, device=q.device
        )
        while blk_offset < seq_end:
            is_full_block = seq_end - blk_offset >= block_size
            valid_len = block_size if is_full_block else seq_end - blk_offset
            o_t = output[blk_offset : min(seq_end, blk_offset + block_size)]
            if is_full_block:
                q_t = q[blk_offset : blk_offset + block_size]
                k_t = k[blk_offset : blk_offset + block_size]
                v_t = v[blk_offset : blk_offset + block_size]
            else:
                q_t = torch.zeros(
                    (block_size, num_qo_heads, head_size),
                    dtype=q.dtype,
                    device=q.device,
                )
                k_t = torch.zeros(
                    (block_size, num_qo_heads, head_size),
                    dtype=q.dtype,
                    device=q.device,
                )
                v_t = torch.zeros(
                    (block_size, num_qo_heads, head_size),
                    dtype=q.dtype,
                    device=q.device,
                )
                q_t[: seq_end - blk_offset] = q[blk_offset:seq_end]
                k_t[: seq_end - blk_offset] = k[blk_offset:seq_end]
                v_t[: seq_end - blk_offset] = v[blk_offset:seq_end]

            Lq = LambdaQ(
                decay_factor,
                valid_len,
                block_size,
                device=q.device,
                offset=decay_exponent_offset,
            )

            o_inter = (
                matmul(
                    q_t.transpose(0, 1).to(torch.float32) * Lq,
                    carried_kv.to(torch.float32),
                )
                .transpose(0, 1)
                .to(q.dtype)
            )
            o_intra = _linear_attention(q_t, k_t, v_t, decay_factor=decay_factor)
            if is_full_block:
                # print(seq_idx, blk_offset, seq_end, o_t.shape, o_inter.shape, o_intra.shape)
                o_t[:] = o_inter + o_intra
            else:
                # print(seq_idx, blk_offset, seq_end, o_t.shape, o_inter.shape, o_intra.shape)
                o_t[:] = (o_inter + o_intra)[: o_t.shape[0]]

            if (decay_factor == 1.0).all():
                inc_kv = matmul(
                    k_t.transpose(0, 1).transpose(-2, -1).to(torch.float32),
                    v_t.transpose(0, 1).to(torch.float32),
                )
                carried_kv = (carried_kv.to(torch.float32) + inc_kv).to(state_dtype)
            else:
                Lk = LambdaK(
                    decay_factor,
                    valid_len,
                    block_size,
                    device=q.device,
                    offset=decay_exponent_offset,
                )
                inc_kv = matmul(
                    (k_t.transpose(0, 1) * Lk).transpose(-2, -1).to(torch.float32),
                    v_t.transpose(0, 1).to(torch.float32),
                )
                block_decay = decay_factor**valid_len
                carried_kv = (block_decay * carried_kv.to(torch.float32) + inc_kv).to(
                    state_dtype
                )
            KVs.append(carried_kv.clone())

            blk_offset += block_size

        # print(kv.shape, carried_kv.shape)
        kv[seq_idx, :, :] = carried_kv

    return output, kv, KVs


def delta_rule(
    q: torch.Tensor,  # [total_seq_len, num_qo_heads, head_size]
    k: torch.Tensor,  # [total_seq_len, num_kv_heads, head_size]
    v: torch.Tensor,  # [total_seq_len, num_kv_heads, head_size]
    seq_lens: list[int],  # sequence length for each sequence
    *,
    alpha: torch.Tensor | None = None,  # [total_seq_len, num_qo_heads]
    beta: torch.Tensor | None = None,  # [total_seq_len, num_qo_heads]
    scale_factor=1.0,
    state_dtype: torch.dtype = torch.float32,
):
    o = []
    kv = []
    total_seqlen = q.size(0)
    num_q_heads = q.size(1)
    num_k_heads = k.size(1)
    num_v_heads = v.size(1)
    num_sab_heads = max(num_q_heads, num_v_heads)
    head_size = k.size(2)

    if alpha is None:
        alpha = torch.ones(
            total_seqlen, num_sab_heads, dtype=torch.float32, device=q.device
        )
    if beta is None:
        beta = torch.ones(
            total_seqlen, num_sab_heads, dtype=torch.float32, device=q.device
        )

    if num_q_heads > num_v_heads:  # GQA
        k = k.repeat_interleave(num_q_heads // num_k_heads, dim=1)
        v = v.repeat_interleave(num_q_heads // num_v_heads, dim=1)
    else:  # GVA
        q = q.repeat_interleave(num_v_heads // num_q_heads, dim=1)
        k = k.repeat_interleave(num_v_heads // num_k_heads, dim=1)

    seq_offset = exclusive_cumsum(seq_lens)
    for seq_idx, seq_start in enumerate(seq_offset[:-1]):
        seq_end = seq_offset[seq_idx + 1]
        seq_len = seq_end - seq_start
        s = slice(seq_start, seq_end)

        # slices
        qs = q[s]
        ks = k[s]
        vs = v[s]
        alphas = alpha[s]
        betas = beta[s]

        state_HKV = torch.zeros(
            num_q_heads, head_size, head_size, dtype=state_dtype, device=q.device
        )
        for i in range(seq_len):
            # var_DS where var is variable basename and DS is the dimensional semantics.
            # Q/K/V are Dq/Dk/Dv respectively
            q_H1Q = qs[i].unsqueeze(1)
            k_H1K = ks[i].unsqueeze(1)
            v_H1V = vs[i].unsqueeze(1)
            alpha_H11 = alphas[i].unsqueeze(1).unsqueeze(2)
            beta_H11 = betas[i].unsqueeze(1).unsqueeze(2)

            ### listed at the bottom of page3 of section 2.2 DELTA NETWORKS: LINEAR ATTENTION WITH DELTA RULE

            # state update rule, use the middle version for clearer dimensional semantics
            # Read state in fp32, compute in fp32, store back in state_dtype
            old_state_HKV = alpha_H11 * state_HKV.to(torch.float32)
            old_v_H1V = matmul(k_H1K, old_state_HKV)
            new_v_H1V = beta_H11 * v_H1V + (1 - beta_H11) * old_v_H1V
            state_remove = torch.einsum("htv,htk->hkv", old_v_H1V, k_H1K)
            state_update = torch.einsum("htv,htk->hkv", new_v_H1V, k_H1K)
            state_HKV[:] = (old_state_HKV - state_remove + state_update).to(state_dtype)

            o_H1V = scale_factor * matmul(q_H1Q, state_HKV.to(torch.float32))
            o.append(o_H1V.squeeze(1))

        kv.append(state_HKV.clone())

    return torch.stack(o), torch.stack(kv)


def _expand_delta_rule_heads(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    alpha: torch.Tensor | None,
    beta: torch.Tensor | None,
):
    total_seqlen = q.size(0)
    num_q_heads = q.size(1)
    num_k_heads = k.size(1)
    num_v_heads = v.size(1)
    num_sab_heads = max(num_q_heads, num_v_heads)

    if alpha is None:
        alpha = torch.ones(
            total_seqlen, num_sab_heads, dtype=torch.float32, device=q.device
        )
    if beta is None:
        beta = torch.ones(
            total_seqlen, num_sab_heads, dtype=torch.float32, device=q.device
        )

    if num_q_heads > num_v_heads:  # GQA
        k = k.repeat_interleave(num_q_heads // num_k_heads, dim=1)
        v = v.repeat_interleave(num_q_heads // num_v_heads, dim=1)
    else:  # GVA
        q = q.repeat_interleave(num_v_heads // num_q_heads, dim=1)
        k = k.repeat_interleave(num_v_heads // num_k_heads, dim=1)

    return q, k, v, alpha, beta


@torch.inference_mode
def cp_delta_rule_pre(
    k: torch.Tensor,  # [chunk_len, num_heads, head_size]
    v: torch.Tensor,  # [chunk_len, num_heads, head_size]
    alpha: torch.Tensor,  # [chunk_len, num_heads]
    beta: torch.Tensor,  # [chunk_len, num_heads]
    *,
    kv_dtype: torch.dtype = torch.float32,
):
    """CP chunk core: compute only this chunk's affine transfer.

    The returned `(M, B)` satisfies `S_out = M @ S_in + B` for this chunk.
    This function intentionally receives no data outside the CP chunk.
    """
    assert k.size(0) == v.size(0) == alpha.size(0) == beta.size(0)
    num_heads = k.size(1)
    head_size = k.size(2)
    eye_HKK = torch.eye(head_size, dtype=kv_dtype, device=k.device).expand(
        num_heads, head_size, head_size
    )
    transfer_HKK = eye_HKK.clone()
    state_HKV = torch.zeros(
        (num_heads, head_size, head_size), dtype=kv_dtype, device=k.device
    )

    for token_idx in range(k.size(0)):
        k_HK = k[token_idx].to(kv_dtype)
        v_HV = v[token_idx].to(kv_dtype)
        alpha_H11 = alpha[token_idx].to(kv_dtype).reshape(num_heads, 1, 1)
        beta_H11 = beta[token_idx].to(kv_dtype).reshape(num_heads, 1, 1)

        kk_HKK = torch.einsum("hk,hl->hkl", k_HK, k_HK)
        kv_HKV = torch.einsum("hk,hv->hkv", k_HK, v_HV)
        token_transfer_HKK = alpha_H11 * (eye_HKK - beta_H11 * kk_HKK)
        token_update_HKV = beta_H11 * kv_HKV

        transfer_HKK = torch.bmm(token_transfer_HKK, transfer_HKK)
        state_HKV = torch.bmm(token_transfer_HKK, state_HKV) + token_update_HKV

    return transfer_HKK, state_HKV


@torch.inference_mode
def blockwise_cp_delta_rule_pre(
    k: torch.Tensor,  # [chunk_len, num_heads, head_size]
    v: torch.Tensor,  # [chunk_len, num_heads, head_size]
    alpha: torch.Tensor,  # [chunk_len, num_heads]
    beta: torch.Tensor,  # [chunk_len, num_heads]
    *,
    block_size: int = 64,
    kv_dtype: torch.dtype = torch.float32,
):
    """CP preprocessing using matrix-form delta-rule blocks.

    Tokens inside each `block_size` block are processed with the same matrix
    formulation as `blockwise_delta_rule`; recurrence is only across blocks.
    Returns `(M, B)` such that `S_out = M @ S_in + B` for the whole CP chunk.
    """
    assert block_size > 0
    assert k.size(0) == v.size(0) == alpha.size(0) == beta.size(0)

    chunk_len = k.size(0)
    num_heads = k.size(1)
    head_size = k.size(2)
    eye_HKK = torch.eye(head_size, dtype=kv_dtype, device=k.device).expand(
        num_heads, head_size, head_size
    )
    transfer_HKK = eye_HKK.clone()
    state_HKV = torch.zeros(
        (num_heads, head_size, head_size), dtype=kv_dtype, device=k.device
    )

    for blk_offset in range(0, chunk_len, block_size):
        valid_len = min(block_size, chunk_len - blk_offset)
        if valid_len == block_size:
            k_SHK = k[blk_offset : blk_offset + block_size]
            v_SHV = v[blk_offset : blk_offset + block_size]
            alpha_SH = alpha[blk_offset : blk_offset + block_size]
            beta_SH = beta[blk_offset : blk_offset + block_size]
        else:
            k_SHK = torch.zeros(
                (block_size, num_heads, head_size), dtype=k.dtype, device=k.device
            )
            v_SHV = torch.zeros(
                (block_size, num_heads, head_size), dtype=v.dtype, device=v.device
            )
            alpha_SH = torch.ones(
                (block_size, num_heads), dtype=alpha.dtype, device=alpha.device
            )
            beta_SH = torch.zeros(
                (block_size, num_heads), dtype=beta.dtype, device=beta.device
            )
            k_SHK[:valid_len] = k[blk_offset:]
            v_SHV[:valid_len] = v[blk_offset:]
            alpha_SH[:valid_len] = alpha[blk_offset:]
            beta_SH[:valid_len] = beta[blk_offset:]

        alpha_HS = alpha_SH.transpose(0, 1)
        beta_HS1 = beta_SH.transpose(0, 1).unsqueeze(2)
        Gamma_HSS, gamma_HS1 = to_logspace_Gamma_and_gamma(alpha_HS)
        block_gamma_H11 = gamma_HS1[:, [valid_len - 1], :]

        k_HSK = k_SHK.transpose(0, 1)
        v_HSV = v_SHV.transpose(0, 1)

        IKK = identity_add_strict_lower_diagonal(
            beta_HS1 * torch.exp(Gamma_HSS) * matmul(k_HSK, k_HSK.transpose(-2, -1))
        )
        T = torch.inverse(IKK) * beta_HS1.transpose(1, 2)
        T = T.to(k.dtype)
        u_HSV = matmul(T, v_HSV)
        w_HSK = matmul(T, torch.exp(gamma_HS1) * k_HSK)

        k_decay_HSK = torch.exp(block_gamma_H11 - gamma_HS1) * k_HSK
        block_transfer_HKK = torch.exp(block_gamma_H11) * eye_HKK - matmul(
            k_decay_HSK.transpose(-2, -1).to(kv_dtype),
            w_HSK.to(kv_dtype),
        )
        block_state_HKV = matmul(
            k_decay_HSK.transpose(-2, -1).to(kv_dtype),
            u_HSV.to(kv_dtype),
        )

        transfer_HKK = torch.bmm(block_transfer_HKK, transfer_HKK)
        state_HKV = torch.bmm(block_transfer_HKK, state_HKV) + block_state_HKV

    return transfer_HKK, state_HKV


@torch.inference_mode
def precompute_blockwise_cp_delta_rule_t(
    k: torch.Tensor,  # [chunk_len, num_heads, head_size]
    beta: torch.Tensor,  # [chunk_len, num_heads]
    *,
    block_size: int = 64,
    kv_dtype: torch.dtype = torch.float32,
    t_dtype: torch.dtype | None = None,
):
    """Precompute the signed, beta-folded triangular solve matrix.

    Returns `T := -(T_clean beta)^T` with shape
    `[num_blocks, num_heads, block_size, block_size]`, in the orientation used
    directly by the CP preprocess WGMMA. Tail blocks are padded with `K = 0`,
    and `beta = 0`, so the returned tail tile is projected as `P T P`.
    """
    assert block_size > 0
    assert k.size(0) == beta.size(0)
    if t_dtype is None:
        t_dtype = k.dtype

    chunk_len = k.size(0)
    num_heads = k.size(1)
    head_size = k.size(2)
    num_blocks = (chunk_len + block_size - 1) // block_size
    t_HSS = torch.empty(
        (num_blocks, num_heads, block_size, block_size), dtype=t_dtype, device=k.device
    )

    for blk_idx, blk_offset in enumerate(range(0, chunk_len, block_size)):
        valid_len = min(block_size, chunk_len - blk_offset)
        if valid_len == block_size:
            k_SHK = k[blk_offset : blk_offset + block_size]
            beta_SH = beta[blk_offset : blk_offset + block_size]
        else:
            k_SHK = torch.zeros(
                (block_size, num_heads, head_size), dtype=k.dtype, device=k.device
            )
            beta_SH = torch.zeros(
                (block_size, num_heads), dtype=beta.dtype, device=beta.device
            )
            k_SHK[:valid_len] = k[blk_offset:]
            beta_SH[:valid_len] = beta[blk_offset:]

        beta_HS1 = beta_SH.transpose(0, 1).unsqueeze(2)
        k_HSK = k_SHK.transpose(0, 1)

        IKK = identity_add_strict_lower_diagonal(
            beta_HS1.to(kv_dtype)
            * matmul(k_HSK.to(kv_dtype), k_HSK.to(kv_dtype).transpose(-2, -1))
        )
        t_clean_HSS = torch.inverse(IKK) * beta_HS1.transpose(1, 2).to(kv_dtype)
        t_HSS[blk_idx] = (-t_clean_HSS.transpose(-2, -1)).to(t_dtype)

    return t_HSS


@torch.inference_mode
def blockwise_cp_delta_rule_pre_transposed(
    k: torch.Tensor,  # [chunk_len, num_heads, head_size]
    v: torch.Tensor,  # [chunk_len, num_heads, head_size]
    alpha: torch.Tensor,  # [chunk_len, num_heads]
    t: torch.Tensor,  # [num_blocks, num_heads, block_size, block_size], signed gamma-sandwiched
    *,
    block_size: int = 64,
    kv_dtype: torch.dtype = torch.float32,
):
    """CP preprocessing in the same orientation as `CPDeltaRuleMNPrecomputeSm90`.

    This path consumes raw `K`, raw `V`, `alpha`, and externally precomputed
    signed gamma-sandwiched `T`. It keeps the running transfer/state transposed internally:
    `Mt_next = Mt @ M_block_t` and `Bt_next = Bt @ M_block_t + B_block_t`.
    The returned tensors use the normal `(M, B)` orientation.
    """
    assert block_size > 0
    assert k.size(0) == v.size(0) == alpha.size(0)

    chunk_len = k.size(0)
    num_heads = k.size(1)
    head_size = k.size(2)
    num_blocks = (chunk_len + block_size - 1) // block_size
    assert t.shape == (num_blocks, num_heads, block_size, block_size)

    eye_HKK = torch.eye(head_size, dtype=kv_dtype, device=k.device).expand(
        num_heads, head_size, head_size
    )
    transfer_t_HKK = eye_HKK.clone()
    state_t_HKV = torch.zeros(
        (num_heads, head_size, head_size), dtype=kv_dtype, device=k.device
    )

    for blk_idx, blk_offset in enumerate(range(0, chunk_len, block_size)):
        valid_len = min(block_size, chunk_len - blk_offset)
        if valid_len == block_size:
            k_SHK = k[blk_offset : blk_offset + block_size]
            v_SHV = v[blk_offset : blk_offset + block_size]
            alpha_SH = alpha[blk_offset : blk_offset + block_size]
        else:
            k_SHK = torch.zeros(
                (block_size, num_heads, head_size), dtype=k.dtype, device=k.device
            )
            v_SHV = torch.zeros(
                (block_size, num_heads, head_size), dtype=v.dtype, device=v.device
            )
            alpha_SH = torch.ones(
                (block_size, num_heads), dtype=alpha.dtype, device=alpha.device
            )
            k_SHK[:valid_len] = k[blk_offset:]
            v_SHV[:valid_len] = v[blk_offset:]
            alpha_SH[:valid_len] = alpha[blk_offset:]

        alpha_HS = alpha_SH.transpose(0, 1)
        _, gamma_HS1 = to_logspace_Gamma_and_gamma(alpha_HS)
        gamma_HS1 = torch.exp(gamma_HS1).to(kv_dtype)
        inv_gamma_H1S = torch.reciprocal(gamma_HS1.transpose(1, 2))
        block_gamma_H11 = gamma_HS1[:, [valid_len - 1], :]

        k_HSK = k_SHK.transpose(0, 1)
        v_HSV = v_SHV.transpose(0, 1)
        k_HSD = k_HSK.to(kv_dtype)
        t_HSS = t[blk_idx].to(kv_dtype)
        v_inv_HDS = v_HSV.transpose(-2, -1).to(kv_dtype) * inv_gamma_H1S
        kt_HDS = k_HSD.transpose(-2, -1)

        block_state_t_HKV = -block_gamma_H11 * matmul(matmul(v_inv_HDS, t_HSS), k_HSD)
        block_transfer_t_HKK = block_gamma_H11 * eye_HKK + block_gamma_H11 * matmul(
            matmul(kt_HDS, t_HSS), k_HSD
        )

        state_t_HKV = matmul(state_t_HKV, block_transfer_t_HKK) + block_state_t_HKV
        transfer_t_HKK = matmul(transfer_t_HKK, block_transfer_t_HKK)

    return transfer_t_HKK.transpose(-2, -1), state_t_HKV.transpose(-2, -1)


@torch.inference_mode
def compose_delta_rule_chunk_range(
    local_transfers: torch.Tensor,  # [num_chunks, num_heads, head_size, head_size]
    local_states: torch.Tensor,  # [num_chunks, num_heads, head_size, head_size]
    start_chunk: int = 0,
    end_chunk: int | None = None,
):
    """Compose chunk-local affine maps over `[start_chunk, end_chunk)`.

    Returns `(M, B)` such that `S_end = M @ S_start + B`.
    """
    if end_chunk is None:
        end_chunk = local_transfers.size(0)
    assert 0 <= start_chunk <= end_chunk <= local_transfers.size(0)

    num_heads = local_transfers.size(1)
    head_size = local_transfers.size(2)
    transfer_HKK = (
        torch.eye(head_size, dtype=local_transfers.dtype, device=local_transfers.device)
        .expand(num_heads, head_size, head_size)
        .clone()
    )
    state_HKV = torch.zeros(
        (local_states.size(1), local_states.size(2), local_states.size(3)),
        dtype=local_states.dtype,
        device=local_states.device,
    )

    for chunk_idx in range(start_chunk, end_chunk):
        chunk_transfer_HKK = local_transfers[chunk_idx]
        chunk_state_HKV = local_states[chunk_idx]
        transfer_HKK = torch.bmm(chunk_transfer_HKK, transfer_HKK)
        state_HKV = torch.bmm(chunk_transfer_HKK, state_HKV) + chunk_state_HKV

    return transfer_HKK, state_HKV


@torch.inference_mode
def cp_delta_rule_fixup(
    local_transfers_by_seq: list[torch.Tensor],
    local_states_by_seq: list[torch.Tensor],
):
    """Fix CP chunk artifacts into global chunk-boundary states.

    This stage consumes only outputs from `blockwise_cp_delta_rule_pre`.
    """
    assert len(local_transfers_by_seq) == len(local_states_by_seq)
    fixed_states_by_seq = []
    final_states = []

    for local_transfers, local_states in zip(  # noqa: B905
        local_transfers_by_seq, local_states_by_seq
    ):
        assert local_transfers.size(0) == local_states.size(0)
        seq_fixed_states = []
        if local_transfers.size(0) == 0:
            fixed_states_by_seq.append(local_states)
            final_states.append(local_states)
            continue

        state_HKV = torch.zeros_like(local_states[0])
        for chunk_idx in range(local_transfers.size(0)):
            state_HKV = (
                torch.bmm(local_transfers[chunk_idx], state_HKV)
                + local_states[chunk_idx]
            )
            seq_fixed_states.append(state_HKV)

        fixed_states = torch.stack(seq_fixed_states)
        fixed_states_by_seq.append(fixed_states)
        final_states.append(state_HKV)

    return torch.stack(final_states), fixed_states_by_seq


@torch.inference_mode
def cp_delta_rule_fixup_transposed(
    local_transfers_by_seq: list[torch.Tensor],
    local_states_by_seq: list[torch.Tensor],
    initial_states_by_seq: list[torch.Tensor] | None = None,
):
    """Fix CP chunk artifacts in CPDeltaRuleFixupSm90 workspace orientation.

    `CPDeltaRuleFixupSm90` consumes transfer/state tensors in transposed
    workspace layout and applies `S_i = S_{i-1} @ M_i + N_i`.
    """
    assert len(local_transfers_by_seq) == len(local_states_by_seq)
    if initial_states_by_seq is not None:
        assert len(initial_states_by_seq) == len(local_transfers_by_seq)
    fixed_states_by_seq = []
    final_states = []

    for seq_idx, (local_transfers, local_states) in enumerate(
        zip(local_transfers_by_seq, local_states_by_seq)  # noqa: B905
    ):
        assert local_transfers.size(0) == local_states.size(0)
        seq_fixed_states = []
        initial_state = (
            initial_states_by_seq[seq_idx]
            if initial_states_by_seq is not None
            else None
        )
        if local_transfers.size(0) == 0:
            fixed_states_by_seq.append(local_states)
            final_states.append(
                local_states if initial_state is None else initial_state
            )
            continue

        state_HKV = (
            torch.zeros_like(local_states[0])
            if initial_state is None
            else initial_state.clone()
        )
        for chunk_idx in range(local_transfers.size(0)):
            state_HKV = (
                torch.bmm(state_HKV, local_transfers[chunk_idx])
                + local_states[chunk_idx]
            )
            seq_fixed_states.append(state_HKV)

        fixed_states = torch.stack(seq_fixed_states)
        fixed_states_by_seq.append(fixed_states)
        final_states.append(state_HKV)

    return torch.stack(final_states), fixed_states_by_seq


@torch.inference_mode
def context_parallel_delta_rule_states(
    q: torch.Tensor,  # [total_seq_len, num_qo_heads, head_size]
    k: torch.Tensor,  # [total_seq_len, num_kv_heads, head_size]
    v: torch.Tensor,  # [total_seq_len, num_kv_heads, head_size]
    seq_lens: list[int],  # sequence length for each sequence
    *,
    alpha: torch.Tensor | None = None,  # [total_seq_len, num_qo_heads]
    beta: torch.Tensor | None = None,  # [total_seq_len, num_qo_heads]
    chunk_size: int = 256,
    pre_block_size: int = 64,
    kv_dtype: torch.dtype = torch.float32,
    return_debug: bool = False,
):
    """Torch reference for context-parallel delta-rule chunk-state fixup.

    Each chunk independently computes an affine transfer `S_out = M S_in + B`.
    The fixup pass applies those chunk transfers to recover true global
    chunk-boundary states. The final fixed state should match `delta_rule`.
    """
    assert chunk_size > 0
    assert pre_block_size > 0

    q, k, v, alpha, beta = _expand_delta_rule_heads(q, k, v, alpha, beta)
    del q  # State fixup only needs K/V and gates.

    local_transfers = []
    local_states = []
    local_transfers_by_seq = []
    local_states_by_seq = []

    seq_offset = exclusive_cumsum(seq_lens)
    for seq_idx, seq_start in enumerate(seq_offset[:-1]):
        seq_end = seq_offset[seq_idx + 1]
        seq_local_transfers = []
        seq_local_states = []

        for chunk_start in range(seq_start, seq_end, chunk_size):
            chunk_end = min(seq_end, chunk_start + chunk_size)
            transfer_HKK, local_state_HKV = blockwise_cp_delta_rule_pre(
                k[chunk_start:chunk_end],
                v[chunk_start:chunk_end],
                alpha[chunk_start:chunk_end],
                beta[chunk_start:chunk_end],
                block_size=pre_block_size,
                kv_dtype=kv_dtype,
            )
            seq_local_transfers.append(transfer_HKK)
            seq_local_states.append(local_state_HKV)
            local_transfers.append(transfer_HKK)
            local_states.append(local_state_HKV)

        local_transfers_by_seq.append(torch.stack(seq_local_transfers))
        local_states_by_seq.append(torch.stack(seq_local_states))

    final_states, fixed_states_by_seq = cp_delta_rule_fixup(
        local_transfers_by_seq, local_states_by_seq
    )

    if not return_debug:
        return final_states, fixed_states_by_seq

    debug = {
        "local_transfers": torch.stack(local_transfers)
        if local_transfers
        else final_states.new_empty(0),
        "local_states": torch.stack(local_states)
        if local_states
        else final_states.new_empty(0),
        "local_transfers_by_seq": local_transfers_by_seq,
        "local_states_by_seq": local_states_by_seq,
        "fixed_states_by_seq": fixed_states_by_seq,
    }
    return final_states, debug


def identity_add_strict_lower_diagonal(m: torch.Tensor):
    SIZE = m.size(-1)
    assert m.size(-2) == SIZE
    with torch.device(m.device):
        m = m.clone()
        mask = torch.arange(SIZE).unsqueeze(1) <= torch.arange(SIZE)
        m[:, mask] = 0.0
        # m[mask.unsqueeze(0)] = 0.0
        m = m + torch.eye(SIZE).unsqueeze(0)
    return m


def to_logspace_Gamma_and_gamma(alpha_HS: torch.Tensor, epsilon=1e-10):
    g = torch.log(alpha_HS + epsilon)
    cu_g = torch.cumsum(g, dim=-1)
    cu_g_HSS = cu_g.unsqueeze(2) - cu_g.unsqueeze(1)
    cu_g_HS1 = cu_g.unsqueeze(2)
    return cu_g_HSS, cu_g_HS1


@torch.inference_mode
def blockwise_delta_rule(
    q: torch.Tensor,  # [total_seq_len, num_qo_heads, head_size]
    k: torch.Tensor,  # [total_seq_len, num_kv_heads, head_size]
    v: torch.Tensor,  # [total_seq_len, num_kv_heads, head_size]
    seq_lens: list[int],  # sequence length for each sequence
    alpha: torch.Tensor | None = None,  # [total_seq_len, num_qo_heads]
    beta: torch.Tensor | None = None,  # [total_seq_len, num_qo_heads]
    block_size: int = 32,
    scale_factor=1.0,
    state_dtype: torch.dtype = torch.float32,
    initial_state: torch.Tensor | None = None,
    # intermediate_outputs = None,  # debug output
) -> torch.Tensor:
    total_seqlen = q.size(0)
    num_q_heads = q.size(1)
    num_k_heads = k.size(1)
    num_v_heads = v.size(1)
    num_sab_heads = max(num_q_heads, num_v_heads)
    head_size = q.size(2)

    if alpha is None:
        alpha = torch.ones(
            total_seqlen, num_sab_heads, dtype=torch.float32, device=q.device
        )
    if beta is None:
        beta = torch.ones(
            total_seqlen, num_sab_heads, dtype=torch.float32, device=q.device
        )

    if num_q_heads > num_v_heads:  # GQA
        num_qkv_heads = num_q_heads
        k = k.repeat_interleave(num_q_heads // num_k_heads, dim=1)
        v = v.repeat_interleave(num_q_heads // num_v_heads, dim=1)
    else:  # GVA
        num_qkv_heads = num_v_heads
        q = q.repeat_interleave(num_v_heads // num_q_heads, dim=1)
        k = k.repeat_interleave(num_v_heads // num_k_heads, dim=1)

    kv = torch.zeros(
        (len(seq_lens), num_sab_heads, head_size, head_size),
        dtype=state_dtype,
        device=q.device,
    )
    if initial_state is not None:
        assert tuple(initial_state.shape) == tuple(kv.shape)
        kv.copy_(initial_state.to(state_dtype))
    output = torch.zeros_like(q)

    seq_offset = exclusive_cumsum(seq_lens)
    for seq_idx, seq_start in enumerate(seq_offset[:-1]):
        seq_end = seq_offset[seq_idx + 1]
        blk_offset = seq_start
        state_HKV = kv[seq_idx].clone()
        while blk_offset < seq_end:
            is_full_block = seq_end - blk_offset >= block_size
            valid_len = block_size if is_full_block else seq_end - blk_offset
            o_t = output[blk_offset : min(seq_end, blk_offset + block_size)]
            if is_full_block:
                q_SHQ = q[blk_offset : blk_offset + block_size]
                k_SHK = k[blk_offset : blk_offset + block_size]
                v_SHV = v[blk_offset : blk_offset + block_size]
                alpha_SH = alpha[blk_offset : blk_offset + block_size]
                beta_SH = beta[blk_offset : blk_offset + block_size]
            else:
                q_SHQ = torch.zeros(
                    (block_size, num_qkv_heads, head_size),
                    dtype=q.dtype,
                    device=q.device,
                )
                k_SHK = torch.zeros(
                    (block_size, num_qkv_heads, head_size),
                    dtype=k.dtype,
                    device=k.device,
                )
                v_SHV = torch.zeros(
                    (block_size, num_qkv_heads, head_size),
                    dtype=v.dtype,
                    device=v.device,
                )
                alpha_SH = torch.ones(
                    block_size, num_sab_heads, dtype=alpha.dtype, device=alpha.device
                )
                beta_SH = torch.zeros(
                    block_size, num_sab_heads, dtype=beta.dtype, device=beta.device
                )
                q_SHQ[:valid_len] = q[blk_offset:seq_end]
                k_SHK[:valid_len] = k[blk_offset:seq_end]
                v_SHV[:valid_len] = v[blk_offset:seq_end]
                alpha_SH[:valid_len] = alpha[blk_offset:seq_end]
                beta_SH[:valid_len] = beta[blk_offset:seq_end]

            alpha_HS = alpha_SH.transpose(0, 1)
            beta_HS1 = beta_SH.transpose(0, 1).unsqueeze(2)
            Gamma_HSS, gamma_HS1 = to_logspace_Gamma_and_gamma(alpha_HS)
            block_gamma = gamma_HS1[:, [valid_len - 1], :]

            q_HSQ = q_SHQ.transpose(0, 1)
            k_HSK = k_SHK.transpose(0, 1)
            v_HSV = v_SHV.transpose(0, 1)

            IKK = identity_add_strict_lower_diagonal(
                beta_HS1 * torch.exp(Gamma_HSS) * matmul(k_HSK, k_HSK.transpose(-2, -1))
            )  # NOTE: beta scale row-wise
            T = torch.inverse(IKK) * beta_HS1.transpose(
                1, 2
            )  # NOTE: beta scale col-wise
            T = T.to(q.dtype)
            # new_v_HSV = matmul(T, (v_HSV - matmul(torch.exp(gamma_HS1) * k_HSK, state_HKV)))
            u_HSV = matmul(T, v_HSV)
            w_HSK = matmul(T, torch.exp(gamma_HS1) * k_HSK)
            new_v_HSV = u_HSV - matmul(
                w_HSK.to(torch.float32), state_HKV.to(torch.float32)
            ).to(u_HSV.dtype)
            new_v_SHV = new_v_HSV.transpose(0, 1)

            # if intermediate_outputs is not None:
            #     intermediate_outputs["G"].append(Gamma_HSS.clone())
            #     intermediate_outputs["g"].append(gamma_HS1.clone())
            #     intermediate_outputs["IKK"].append(IKK.clone())
            #     intermediate_outputs["T"].append(T.clone())
            #     intermediate_outputs["u"].append(u_HSV.clone())
            #     intermediate_outputs["w"].append(w_HSK.clone())
            #     intermediate_outputs["new_v"].append(new_v_HSV.clone())

            o_inter = (
                matmul(
                    torch.exp(gamma_HS1) * q_HSQ.to(torch.float32),
                    state_HKV.to(torch.float32),
                )
                .transpose(0, 1)
                .to(q.dtype)
            )
            o_intra = _linear_attention(
                q_SHQ, k_SHK, new_v_SHV, qk_weight=torch.exp(Gamma_HSS)
            )

            if is_full_block:
                o_t[:] = scale_factor * (o_inter + o_intra)
            else:
                o_t[:] = scale_factor * (o_inter + o_intra)[: o_t.shape[0]]

            inc_HKV = matmul(
                (torch.exp(block_gamma - gamma_HS1) * k_HSK)
                .transpose(-2, -1)
                .to(torch.float32),
                new_v_HSV.to(torch.float32),
            )
            state_HKV = (
                torch.exp(block_gamma) * state_HKV.to(torch.float32) + inc_HKV
            ).to(state_dtype)

            blk_offset += block_size

        kv[seq_idx, :, :, :] = state_HKV

    return output, kv


@torch.inference_mode
def decode_delta_rule(
    q: torch.Tensor,  # [B, num_q_heads, K]
    k: torch.Tensor,  # [B, num_k_heads, K]
    v: torch.Tensor,  # [B, num_v_heads, V]
    state: torch.Tensor,  # [B, num_heads, K, V]
    A_log: torch.Tensor,  # [num_heads] - log decay parameter
    a: torch.Tensor,  # [B, num_heads] - input-dependent decay
    dt_bias: torch.Tensor,  # [num_heads] - decay bias
    b: torch.Tensor,  # [B, num_heads] - update gate input
    scale_factor: float = 1.0,
    softplus_beta: float = 1.0,
    softplus_threshold: float = 20.0,
    use_l2_norm: bool = True,
    state_dtype: torch.dtype = torch.float32,
) -> tuple[torch.Tensor, torch.Tensor]:
    """
    Reference implementation for single-step decode with GDN formula.

    Strictly follows the Triton kernel logic from fused_sigmoid_gating_recurrent.py:
        1. Compute g = -exp(A_log) * softplus(a + dt_bias)
        2. Compute beta = sigmoid(b)
        3. Apply L2 norm to q and k (if enabled)
        4. h *= exp(g)                    # Apply decay to state
        5. v_new = v - k^T @ h            # Delta rule (h is [K,V], k is [K])
        6. v_new *= beta                  # Apply update gate
        7. h += k @ v_new^T               # Update state (outer product)
        8. o = q^T @ h                    # Compute output (h is [K,V], q is [K])

    Args:
        q: Query [B, num_q_heads, K]
        k: Key [B, num_k_heads, K]
        v: Value [B, num_v_heads, V]
        state: Input state [B, num_heads, K, V], where num_heads = num_v_heads
        A_log: Log decay parameter [num_heads]
        a: Input-dependent decay [B, num_heads]
        dt_bias: Decay bias [num_heads]
        b: Update gate input [B, num_heads]
        scale_factor: Scale factor for q
        softplus_beta: Beta parameter for softplus activation
        softplus_threshold: Threshold for softplus numerical stability
        use_l2_norm: Whether to apply L2 normalization to q and k
        state_dtype: Storage dtype for the hidden state (read in fp32, stored in this dtype)

    Returns:
        output: [B, num_heads, V]
        new_state: [B, num_heads, K, V]
    """
    B = q.size(0)
    num_q_heads = q.size(1)
    num_k_heads = k.size(1)
    num_v_heads = v.size(1)
    K = q.size(2)

    device = q.device
    dtype = torch.float32

    # Convert to float32 for computation
    A_log = A_log.to(dtype).to(device)
    a = a.to(dtype).to(device)
    dt_bias = dt_bias.to(dtype).to(device)
    b = b.to(dtype).to(device)

    # ============================================
    # Compute gating values (following Triton kernel exactly)
    # ============================================

    # Step 1: Compute g = -exp(A_log) * softplus(a + dt_bias)
    # Triton kernel lines 100-109
    x = a + dt_bias  # [B, num_heads]
    beta_x = softplus_beta * x

    # Apply softplus with numerical stability
    # softplus(x) = (1/beta) * log(1 + exp(beta*x)) if beta*x <= threshold, else x
    softplus_x = torch.where(
        beta_x <= softplus_threshold,
        (1.0 / softplus_beta) * torch.log(1.0 + torch.exp(beta_x)),
        x,
    )

    # Compute g (log-space decay gate)
    # Triton kernel line 109: b_g = -tl.exp(b_A_log) * softplus_x
    g = -torch.exp(A_log) * softplus_x  # [B, num_heads]

    # Step 2: Compute beta = sigmoid(b)
    # Triton kernel line 112: b_beta = 1.0 / (1.0 + tl.exp(-b_b))
    beta = 1.0 / (1.0 + torch.exp(-b))  # [B, num_heads]

    # Expand heads if needed (for GQA/GVA)
    # The reference works at v_heads level
    # For GQA (num_q_heads > num_v_heads): k and q need to be averaged/pooled per v_head
    # For GVA (num_v_heads > num_q_heads): q and k need to be repeated
    if num_k_heads < num_v_heads:
        k = k.repeat_interleave(num_v_heads // num_k_heads, dim=1)
    if num_q_heads < num_v_heads:
        q = q.repeat_interleave(num_v_heads // num_q_heads, dim=1)
    elif num_q_heads > num_v_heads:
        # GQA: multiple q_heads per v_head, reshape and average
        # [B, num_q_heads, K] -> [B, num_v_heads, num_q_heads//num_v_heads, K]
        q = q.reshape(B, num_v_heads, num_q_heads // num_v_heads, K).mean(dim=2)
        if num_k_heads == num_q_heads:
            k = k.reshape(B, num_v_heads, num_k_heads // num_v_heads, K).mean(dim=2)

    q = q.to(dtype)
    k = k.to(dtype)
    v = v.to(dtype)
    state = state.to(dtype)

    # Apply L2 normalization if requested
    if use_l2_norm:
        q = F.normalize(q, p=2.0, dim=-1)
        k = F.normalize(k, p=2.0, dim=-1)

    # Apply scale to q
    q = q * scale_factor

    # ============================================
    # Process each batch and head
    # ============================================
    # All (batch, head) pairs are independent, so the recurrent update below is
    # vectorized over [B, num_heads].

    # Step 1: Apply gating to hidden state: h *= exp(g)
    h_state = state.to(torch.float32) * torch.exp(g)[:, :, None, None]  # [B,H,K,V]

    # Step 2: Delta rule: v -= k^T @ h  ([B,H,K] x [B,H,K,V] -> [B,H,V])
    v_new = v - torch.einsum("bhk,bhkv->bhv", k, h_state)

    # Step 3: Apply beta gating: v *= beta
    v_new = v_new * beta[:, :, None]

    # Step 4: Update hidden state with outer product: h += k ⊗ v_new
    h_state = h_state + k.unsqueeze(-1) * v_new.unsqueeze(-2)

    # Step 5: Compute output: o = q^T @ h
    output = torch.einsum("bhk,bhkv->bhv", q, h_state).to(dtype)

    # Store updated state (cast back to state_dtype)
    new_state = h_state.to(state_dtype)

    return output, new_state


@torch.inference_mode
def verify_delta_rule(
    q: torch.Tensor,  # [B, T, num_q_heads, K]
    k: torch.Tensor,  # [B, T, num_k_heads, K]
    v: torch.Tensor,  # [B, T, num_v_heads, V]
    state: torch.Tensor,  # [B, num_heads, K, V]
    A_log: torch.Tensor,  # [num_heads] - log decay parameter
    a: torch.Tensor,  # [B, T, num_heads] - input-dependent decay
    dt_bias: torch.Tensor,  # [num_heads] - decay bias
    b: torch.Tensor,  # [B, T, num_heads] - update gate input
    scale_factor: float = 1.0,
    softplus_beta: float = 1.0,
    softplus_threshold: float = 20.0,
    use_l2_norm: bool = True,
    cache_intermediate_states: bool = False,
    state_dtype: torch.dtype = torch.float32,
) -> tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:
    """
    Reference implementation for multi-token (verify mode) delta rule.

    Processes T tokens sequentially, updating the state after each token.
    Optionally caches intermediate states for rollback in speculative decoding.

    Args:
        q: Query tensor [B, T, num_q_heads, K]
        k: Key tensor [B, T, num_k_heads, K]
        v: Value tensor [B, T, num_v_heads, V]
        state: Initial state tensor [B, num_heads, K, V]
        A_log: Log decay parameter [num_heads]
        a: Input-dependent decay [B, T, num_heads]
        dt_bias: Decay bias [num_heads]
        b: Update gate input [B, T, num_heads]
        scale_factor: Scaling factor for queries
        softplus_beta: Beta parameter for softplus
        softplus_threshold: Threshold for softplus approximation
        use_l2_norm: Whether to apply L2 normalization
        cache_intermediate_states: Whether to cache state at each time step
        state_dtype: Storage dtype for the hidden state (read in fp32, stored in this dtype)

    Returns:
        output: Output tensor [B, T, num_heads, V]
        new_state: Final state tensor [B, num_heads, K, V]
        intermediate_states: Cached intermediate states [B, T, num_heads, K, V] or None
    """
    B, T, num_q_heads, K = q.shape
    _, _, num_k_heads, _ = k.shape
    _, _, num_v_heads, V = v.shape
    num_heads = state.shape[1]

    # Handle GQA/GVA: expand or average heads
    if num_q_heads != num_heads:
        # Expand q heads to match num_heads (num_v_heads)
        assert num_heads % num_q_heads == 0
        repeat_factor = num_heads // num_q_heads
        q = q.repeat_interleave(repeat_factor, dim=2)  # [B, T, num_heads, K]

    if num_k_heads != num_heads:
        # Expand k heads to match num_heads (num_v_heads)
        assert num_heads % num_k_heads == 0
        repeat_factor = num_heads // num_k_heads
        k = k.repeat_interleave(repeat_factor, dim=2)  # [B, T, num_heads, K]

    # Convert to float32 for computation
    q = q.float()
    k = k.float()
    v = v.float()
    state = state.float()
    A_log = A_log.float()
    a = a.float()
    dt_bias = dt_bias.float()
    b = b.float()

    # Pre-compute gating values for all time steps
    # Shape: [B, T, num_heads]
    x = a + dt_bias.unsqueeze(0).unsqueeze(0)  # [B, T, num_heads]
    beta_x = softplus_beta * x

    # Softplus with threshold
    softplus_x = torch.where(
        beta_x <= softplus_threshold,
        (1.0 / softplus_beta) * torch.log(1.0 + torch.exp(beta_x)),
        x,
    )

    # Compute g (decay factor, already includes exp)
    g = torch.exp(
        -torch.exp(A_log.unsqueeze(0).unsqueeze(0)) * softplus_x
    )  # [B, T, num_heads]

    # Compute beta (update gate)
    beta = 1.0 / (1.0 + torch.exp(-b))  # [B, T, num_heads]

    # Apply L2 normalization if needed
    if use_l2_norm:
        q = torch.nn.functional.normalize(q, p=2, dim=-1)
        k = torch.nn.functional.normalize(k, p=2, dim=-1)

    # Apply scaling to q
    q = q * scale_factor

    # Initialize output and intermediate states
    output = torch.zeros(B, T, num_heads, V, dtype=torch.float32, device=q.device)
    current_state = state.clone().to(
        state_dtype
    )  # [B, num_heads, K, V] stored in state_dtype

    if cache_intermediate_states:
        intermediate_states = torch.zeros(
            B, T, num_heads, K, V, dtype=state_dtype, device=q.device
        )
    else:
        intermediate_states = None

    # Process each time step sequentially (only t carries a dependency; the
    # per-(batch, head) updates are independent and vectorized over [B, H]).
    for t in range(T):
        # Recurrent update
        # 1. Apply decay (state read as fp32)
        h_state = current_state.to(torch.float32) * g[:, t, :, None, None]

        # 2. Compute prediction error: v - k^T @ h
        v_new = v[:, t] - torch.einsum("bhk,bhkv->bhv", k[:, t], h_state)

        # 3. Apply gating
        v_new = v_new * beta[:, t, :, None]

        # 4. Update state with outer product: h = h + k ⊗ v_new
        h_state = h_state + k[:, t].unsqueeze(-1) * v_new.unsqueeze(-2)

        # 5. Compute output: o = q^T @ h
        output[:, t] = torch.einsum("bhk,bhkv->bhv", q[:, t], h_state)

        # Update current state (cast back to state_dtype)
        current_state = h_state.to(state_dtype)

        # Cache intermediate state if requested
        if cache_intermediate_states:
            intermediate_states[:, t] = current_state

    return output, current_state, intermediate_states
