Diff Coverage

    if cu_seqlens is None:
        N, NT, chunk_offsets = B, triton.cdiv(T, BT), None
    else:
        N, NT, chunk_offsets = len(cu_seqlens) - 1, len(chunk_indices), prepare_chunk_offsets(cu_seqlens, BT)
    if K > 256:
        raise ValueError("current kernel does not support head dimension larger than 256.")

    h = k.new_empty(B, NT, H, K, V).permute(0, 2, 1, 3, 4).contiguous()
    final_state = k.new_empty(N, H, K, V, dtype=torch.float32) if output_final_state else None

) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    B, T, H, K, V = *q.shape, do.shape[-1]
    # N: the actual number of sequences in the batch with either equal or variable lengths
    BT = 64
    if K > 256:
        raise ValueError("current kernel does not support head dimension being larger than 256.")

    if chunk_indices is None and cu_seqlens is not None:
        chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size)
    if cu_seqlens is None:

    NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)
    if scale is None:
        scale = k.shape[-1] ** -0.5

    output = torch.empty_like(v)
    if cu_seqlens is None:
        N, chunk_offsets = B, None
    else:
        N, chunk_offsets = (

        BT=BT,
        BK=128,
        BV=128,
    )
    return output

bwd_chunk_dqkwg = chunk_bwd_dqkwg
bwd_chunk_dv_local = chunk_bwd_dv_local

# pylint: disable=missing-public-type-hints,invalid-name

"""Fixed-shape Triton-Ascend kernels for GDN state summaries."""

__all__ = ["gdn_packed_state_summary_kernel", "gdn_state_grad_ext_kernel"]

import triton
import triton.language as tl

    logger.info(msg)
    error_rate = get_err_ratio(ref, tri)
    if abs_atol <= err_atol:
        return
    allow_warning = warning or (FLA_CI_ENV and (error_rate < 0.01 or abs_atol <= 0.3))
    if allow_warning:
        if error_rate > ratio:
            warnings.warn(msg)
    else:
        if not error_rate < ratio:
            raise AssertionError(msg)


if hasattr(triton.language, '_experimental_make_tensor_descriptor'):
    # For Triton 3.3.x

    def custom_device_ctx(index: int):
        return device_torch_lib.device(index)
else:
    if device != 'cuda':
        raise AssertionError('Only cuda device is supported for PyTorch version < 2.4.0.')
    autocast_custom_fwd = device_torch_lib.amp.custom_fwd
    autocast_custom_bwd = device_torch_lib.amp.custom_bwd

    def custom_device_ctx(index: int):

# pylint: disable=invalid-name,missing-public-type-hints

"""Fixed-shape Triton-Ascend kernel for a packed KDA state summary."""

__all__ = [
    "kda_split_state_summary_kernel",
    "kda_state_grad_ext_kernel",
]

        chunk_size: int = 64,
):
    """Compute forward intermediates that do not depend on the initial state."""
    g = chunk_local_cumsum(g, chunk_size=chunk_size, cu_seqlens=cu_seqlens, head_first=False)
    matrix_a = chunk_scaled_dot_kkt_fwd(
        k=k,
        g=g,
        beta=beta,
        cu_seqlens=cu_seqlens,

        cu_seqlens=cu_seqlens,
        chunk_size=chunk_size,
        output_dtype=torch.float32,
    )
    matrix_a = solve_tril(A=matrix_a, cu_seqlens=cu_seqlens, output_dtype=k.dtype)
    w, u = recompute_w_u_fwd(
        k=k,
        v=v,
        beta=beta,

        A=matrix_a,
        g=g,
        cu_seqlens=cu_seqlens,
    )
    return g, matrix_a, w, u


def chunk_gated_delta_rule_fwd_apply_state(
        k: torch.Tensor,

        output_final_state: bool,
        cu_seqlens: Optional[torch.LongTensor] = None,
        chunk_size: int = 64,
):
    g, matrix_a, w, u = chunk_gated_delta_rule_fwd_prepare(
        k=k,
        v=v,
        g=g,
        beta=beta,

        output_final_state=output_final_state,
        cu_seqlens=cu_seqlens,
        chunk_size=chunk_size,
    )
    output = chunk_gated_delta_rule_fwd_output(
        q=q,
        k=k,
        v_new=v_new,
        h=h,

        scale=scale,
        cu_seqlens=cu_seqlens,
        chunk_size=chunk_size,
    )
    return g, output, matrix_a, final_state


def chunk_gated_delta_rule_bwd_prepare(
        q: torch.Tensor,

    if use_qk_l2norm_in_kernel:
        q_norm, q_inv_norm = _l2norm(q)
        k_norm, k_inv_norm = _l2norm(k)

    g_cumsum, matrix_a, w, u = chunk_gated_delta_rule_fwd_prepare(
        k=k_norm,
        v=v,
        g=g,
        beta=beta,

            cu_seqlens: Optional[torch.LongTensor] = None,
            use_qk_l2norm_in_kernel: bool = False,
            chunk_size: int = 64,
    ):
        g, output, matrix_a, final_state = chunk_gated_delta_rule_fwd(
            q=q,
            k=k,
            v=v,
            g=g,

        )

        saved_initial_state = initial_state if initial_state is not None else q.new_empty(0)
        saved_cu_seqlens = cu_seqlens if cu_seqlens is not None else q.new_empty(0, dtype=torch.long)
        ctx.save_for_backward(q, k, v, g, beta, matrix_a, saved_initial_state, saved_cu_seqlens)
        ctx.has_initial_state = initial_state is not None
        ctx.has_cu_seqlens = cu_seqlens is not None
        ctx.scale = scale
        ctx.chunk_size = chunk_size
        return output.to(q.dtype), final_state

    @staticmethod
    @input_guard
    @autocast_custom_bwd

            ctx,
            do: torch.Tensor,
            dht: torch.Tensor
    ):
        q, k, v, g, beta, matrix_a, initial_state, cu_seqlens = ctx.saved_tensors
        if not ctx.has_initial_state:
            initial_state = None
        if not ctx.has_cu_seqlens:
            cu_seqlens = None

    if use_qk_l2norm_in_kernel:
        q, _ = _l2norm(q)
        k, _ = _l2norm(k)

    output, final_state = ChunkGatedDeltaRuleFunction.apply(
        q,
        k,
        v,
        g,

        cu_seqlens,
        False,
        chunk_size,
    )
    return output, final_state

"""Affine state-summary operations used by GDN State-P2P."""

# pylint: disable=forbidden-backend-import

__all__ = [
    "apply_gdn_state_gradient_summary",
    "apply_gdn_state_summary",
    "chunk_gated_delta_rule_state_gradient_summary_bwd",
    "chunk_gated_delta_rule_state_summary_fwd",

    chunk_size: int = 64,
    block_size: int = 128,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Return the local affine map ``state_out = M @ state_in + S``."""
    if (key.ndim, w.ndim, u.ndim, g.ndim) != (4, 4, 4, 3):
        raise ValueError("GDN state summary expects key/w/u [B,T,H,D] and g [B,T,H].")
    batch, seq_len, heads, key_dim = key.shape
    value_dim = u.shape[-1]
    _validate_fixed_summary_shape(key_dim, value_dim, chunk_size)

    transition = packed_summary[..., value_dim:].contiguous()
    return state_ext, transition


def _validate_gdn_gradient_summary_inputs(
    query: torch.Tensor,
    key: torch.Tensor,
    w: torch.Tensor,
    g: torch.Tensor,

            f"GDN state-gradient sequence length {seq_len} must be divisible by {chunk_size}."
        )
    qk_shape = (batch, seq_len, heads, key_dim)
    value_shape = (batch, seq_len, heads, value_dim)
    if (key.shape, w.shape, g.shape, grad_output.shape, dv.shape) != (
        qk_shape, qk_shape, qk_shape[:3], value_shape, value_shape
    ):
        raise ValueError(
            "Incompatible GDN state-gradient summary shapes: "

            f"query={tuple(query.shape)}, key={tuple(key.shape)}, "
            f"w={tuple(w.shape)}, g={tuple(g.shape)}, "
            f"grad_output={tuple(grad_output.shape)}, dv={tuple(dv.shape)}."
        )
    return batch, seq_len, heads, key_dim, value_dim


@torch.compiler.disable
def chunk_gated_delta_rule_state_gradient_summary_bwd(
    query: torch.Tensor,
    key: torch.Tensor,
    w: torch.Tensor,
    g: torch.Tensor,

    *,
    chunk_size: int = 64,
) -> torch.Tensor:
    """Return the local-loss contribution to the incoming state gradient."""
    batch, seq_len, heads, key_dim, value_dim = _validate_gdn_gradient_summary_inputs(
        query, key, w, g, grad_output, dv, chunk_size
    )

    query, key, w, g, grad_output, dv = (

from __future__ import annotations

# pylint: disable=forbidden-backend-import

__all__ = ["fused_chunk_kda", "fused_chunk_kda_p2p"]

from typing import Any, Optional

import torch

    kda_state_summary_forward_from_prepared,
)


def _validate_local_shapes(
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    gate: torch.Tensor,

    gate: torch.Tensor,
    beta: torch.Tensor,
) -> tuple[int, int, int, int]:
    """Validate local KDA tensor shapes and return dimensions used by later checks."""
    if (query.ndim, key.ndim, value.ndim, gate.ndim) != (4, 4, 4, 4):
        raise ValueError("Fused KDA expects rank-4 query/key/value/gate tensors.")
    if beta.ndim != 3:
        raise ValueError("Fused KDA expects a rank-3 beta tensor.")
    if query.shape != key.shape:

    if beta.shape != (batch, sequence_length, num_value_heads):
        raise ValueError("Fused KDA beta has an incompatible shape.")
    if num_value_heads % num_query_heads:
        raise ValueError("Fused KDA value heads must be divisible by query heads.")
    return sequence_length, key_dim, value_dim, num_value_heads


def _validate_local_inputs(
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    gate: torch.Tensor,

    chunk_size: int,
    lower_bound: float,
) -> None:
    """Validate the fixed-shape local KDA backend before compiling kernels."""
    sequence_length, key_dim, value_dim, num_value_heads = _validate_local_shapes(
        query, key, value, gate, beta
    )
    if not (
        query.dtype == key.dtype == value.dtype == gate.dtype == beta.dtype

            chunk_size=ctx.chunk_size,
        )
        # These recomputed tensors have no consumers after dhu. Dropping local
        # references lets the allocator reuse their storage on this stream.
        del query_gated, key_gated, w, u, grad_value_local, _
        (
            grad_query,
            grad_key,
            grad_value,

            scale=ctx.scale,
            chunk_size=ctx.chunk_size,
        )
        # Release recomputed states before intra backward allocates its outputs.
        del states, grad_states, value_new
        grad_query, grad_key, grad_beta, grad_gate = ops.chunk_kda_bwd_intra(
            q=query,
            k=key,
            g=gate,

from __future__ import annotations

# pylint: disable=forbidden-backend-import

__all__ = [
    "FLAKDAStagedOps",
    "get_fla_kda_staged_ops",
    "is_fla_triton_kda_available",
    "run_fla_chunk_kda",

"""Affine KDA state summary built from prepared WY intermediates."""

# pylint: disable=forbidden-backend-import

__all__ = [
    "apply_kda_state_gradient_summary",
    "apply_kda_state_summary",
    "kda_state_gradient_summary_from_prepared",
    "kda_state_summary_forward_from_prepared",

    gate: torch.Tensor,
    chunk_size: int,
) -> tuple[int, int, int, int, int]:
    """Validate the fixed Kimi K3 summary contract and return dimensions."""
    if (key.ndim, w.ndim, u.ndim, gate.ndim) != (4, 4, 4, 4):
        raise ValueError("KDA prepared summary expects rank-4 key/w/u/gate tensors.")
    batch, sequence_length, heads, key_dim = key.shape
    value_dim = u.shape[-1]
    if key_dim != 128 or value_dim != 128 or chunk_size != 64:

def _parse_version(version_text: str) -> tuple[int, int, int]:
    """Return a three-component numeric version tuple."""
    parts = []
    for part in version_text.split("+")[0].split(".")[:3]:
        match = re.match(r"\d+", part)
        if match is not None:
            parts.append(int(match.group()))
    return tuple((parts + [0, 0, 0])[:3])


def _is_triton_gdn_input_supported(

    g: Optional[torch.Tensor],
    beta: Optional[torch.Tensor],
) -> bool:
    """Check the fixed Qwen3.5 GDN contract validated by this backend."""
    if any(tensor is None for tensor in (key, value, g, beta)):
        return False
    if not (
        query.device.type == "npu"
        and query.dtype == key.dtype == value.dtype == beta.dtype == torch.bfloat16

    dt_bias: torch.Tensor,
    chunk_size: int,
) -> bool:
    """Check the fixed dense shapes accepted by the Triton KDA path."""
    batch_size, sequence_length, num_query_heads, key_dim = query.shape
    num_value_heads, value_dim = value.shape[2:]
    return all((
        value.shape[:2] == (batch_size, sequence_length),
        gate.shape == (batch_size, sequence_length, num_value_heads, key_dim),
        beta.shape == (batch_size, sequence_length, num_value_heads),
        num_value_heads % num_query_heads == 0,

    )
    if not all(basic_contract):
        return False

    shape_supported = _has_triton_kda_shapes(
        query, value, gate, beta, a_log, dt_bias, chunk_size
    )
    lower_bound_supported = -5.0 <= lower_bound < 0
    if not (shape_supported and lower_bound_supported):
        return False
    return all(tensor.device == query.device for tensor in operands)

    convolution: nn.Conv1d,
    halo: Optional[torch.Tensor] = None,
) -> torch.Tensor:
    """Apply one local ShortConv, optionally with a preceding-rank halo."""
    if halo is None:
        conv_input = tensor.transpose(1, 2)
        padding = convolution.padding
    else:
        conv_input = torch.cat((halo, tensor), dim=1).transpose(1, 2)
        padding = 0
    output = F.conv1d(  # pylint: disable=not-callable
        input=conv_input,
        weight=convolution.weight,
        bias=convolution.bias,
        stride=convolution.stride,

        padding=padding,
        dilation=convolution.dilation,
        groups=convolution.groups,
    )
    if halo is None:
        output = output[:, :, : tensor.shape[1]]
    return F.silu(output).transpose(1, 2)


def _causal_short_convs_with_cp_halo(
    projected: tuple[torch.Tensor, torch.Tensor, torch.Tensor],

    if len(halo_widths) != 1:
        raise ValueError("KDA P2P requires Q/K/V ShortConv halo widths to match.")
    halo_width = halo_widths.pop()
    if halo_width == 0 or cp_size == 1:
        return tuple(
            _run_causal_short_conv(tensor, convolution)
            for tensor, convolution in zip(projected, convolutions)
        )
    if projected[0].shape[1] < halo_width:

        cp_rank,
        cp_size,
    )
    halos = torch.split(packed_halo, channel_sizes, dim=-1)
    return tuple(
        _run_causal_short_conv(tensor, convolution, halo)
        for tensor, halo, convolution in zip(projected, halos, convolutions)
    )

        a_log: torch.Tensor,
        dt_bias: torch.Tensor,
    ) -> None:
        """Validate the projected-tensor boundary and Ulysses divisibility."""
        if (query.dim(), key.dim(), value.dim(), gate.dim()) != (4, 4, 4, 4):
            raise ValueError("query, key, value, and gate must be rank-4 tensors.")
        if beta.dim() != 3:
            raise ValueError("beta must be a rank-3 tensor.")
        if query.shape != key.shape: