Diff Coverage

@lru_cache(maxsize=1)
def _load_ops():
    """Keep the optional native runtime out of ordinary framework imports."""
    try:
        ops = importlib.import_module("fla_npu.ops.ascendc")
    except ImportError as error:
        raise RuntimeError("GDN backend='ascendc' requires a compatible fla-npu installation.") from error
    required = {
        "chunk_gated_delta_rule_fwd": (
            "layout", "use_exp2", "cu_seqlens", "chunk_indices", "disable_recompute",
            "use_qk_l2norm_in_kernel", "use_beta_sigmoid_in_kernel", "use_gate_in_kernel",
        ),

            "layout", "use_exp2", "cu_seqlens", "chunk_indices", "use_qk_l2norm_in_kernel",
            "use_beta_sigmoid_in_kernel", "use_gate_in_kernel",
        ),
    }
    for name, parameters in required.items():
        try:
            signature = inspect.signature(getattr(ops, name, None))
        except (TypeError, ValueError) as error:
            raise RuntimeError(f"fla-npu lacks a compatible public {name} API.") from error
        missing = set(parameters).difference(signature.parameters)
        if missing:
            raise RuntimeError(f"fla-npu {name} lacks required parameters {sorted(missing)}.")
    return ops


def _packed_offsets(cu_seqlens, tokens):
    """Freeze validated document offsets without a device-to-host transfer."""

                or cu_seqlens.dtype not in (torch.int32, torch.int64)):
            raise ValueError("Ascend C GDN cu_seqlens must be a one-dimensional CPU integer tensor or sequence.")
        cu_seqlens = cu_seqlens.tolist()
    if not isinstance(cu_seqlens, Sequence):
        raise ValueError("Ascend C GDN cu_seqlens must be a CPU integer tensor or sequence.")
    offsets = tuple(cu_seqlens)
    if (len(offsets) < 2 or offsets[0] != 0 or offsets[-1] != tokens
            or any(not isinstance(value, int) or isinstance(value, bool) for value in offsets)
            or any(start >= end for start, end in zip(offsets, offsets[1:]))):

        # The public ABI keeps ten slots even when optional norm/state outputs
        # are disabled: o, final_state, g_cumsum, A, beta_eff, h, q_hat,
        # k_hat, q_rstd, k_rstd. This wrapper only consumes slots 0, 2 and 3.
        if len(result) != 10 or any(not isinstance(result[index], torch.Tensor) for index in (0, 2, 3)):
            raise RuntimeError("fla-npu GDN forward must return ten slots with output, g_cumsum and A tensors.")
        output, _, gate_cumsum, matrix_a, *_ = result
        # Gate and beta transformations remain in the caller's autograd graph.
        # Native backward consumes cumulative gate but returns the raw-gate gradient.
        ctx.save_for_backward(query, key, value, gate_cumsum, beta, matrix_a)

            use_qk_l2norm_in_kernel=False, use_beta_sigmoid_in_kernel=False,
            use_gate_in_kernel=False, state_v_first=False, **ctx.native_metadata,
        )
        if len(gradients) != 8 or any(not isinstance(gradient, torch.Tensor) for gradient in gradients[:5]):
            raise RuntimeError("fla-npu GDN backward must return eight slots with all five input gradients.")
        dq, dk, dv, dbeta, dgate, *_ = gradients
        return dq, dk, dv, dgate, dbeta, None

        ValueError: Inputs violate the explicit native training contract.
        RuntimeError: The optional native API or device is unsupported.
    """
    if query.ndim != 4:
        raise ValueError("GDN inputs must have shape [B,T,H,K].")
    if cu_seqlens is not None and query.shape[0] != 1:
        raise ValueError("Packed Ascend C GDN requires physical batch size one.")
    metadata = _packed_metadata(cu_seqlens, query.shape[1])
    return _AscendCGDN.apply(query, key, value, g, beta, metadata), None

def _load_ops():
    """Keep the optional native runtime out of ordinary framework imports."""
    try:
        ops = importlib.import_module("fla_npu.ops.ascendc")
    except ImportError as error:
        raise RuntimeError("backend='ascendc' requires a compatible fla-npu installation.") from error
    required = {
        "chunk_kda_fwd_prepare": ("backward_mode", "use_qk_l2norm_in_kernel", "use_beta_sigmoid_in_kernel"),
        "chunk_fwd_h": ("output_final_state", "cu_seqlens", "chunk_indices"),
        "chunk_kda_fwd_finalize": ("output_layout", "cu_seqlens", "chunk_indices"),

    for name, parameters in required.items():
        operation = getattr(ops, name, None)
        try:
            signature = inspect.signature(operation)
        except (TypeError, ValueError) as error:
            raise RuntimeError(f"fla-npu lacks a compatible public {name} API.") from error
        missing = set(parameters).difference(signature.parameters)
        if missing:
            raise RuntimeError(
                f"fla-npu {name} lacks required parameters {sorted(missing)}; "

    Native operators remain responsible for validating device support.
    """
    if device.type != "npu":
        raise ValueError("Ascend C KDA requires NPU tensors before querying native capabilities.")
    return str(torch.npu.get_device_name(device)).startswith(("Ascend910B", "Ascend910C", "Ascend910_93"))


class _KDAPrepared(NamedTuple):
    """Public Prepare save-mode slots, in the native ABI's fixed order."""

@lru_cache(maxsize=1)
def _load_beta_backward():
    """Load optional FLA helpers without requiring its staged KDA backend."""
    try:
        importlib.import_module("triton.backends.ascend")
        sigmoid = importlib.import_module("fla.ops.common.gate").fused_beta_sigmoid_bwd
    except (ImportError, AttributeError) as error:
        raise RuntimeError(
            "backend='ascendc' requires FLA with the Triton-Ascend beta-sigmoid backward helper."
        ) from error
    if not callable(sigmoid):
        raise RuntimeError("The installed FLA beta-sigmoid backward API is incompatible.")
    return sigmoid


def _native_backward_view(tensor: torch.Tensor, packed: bool) -> torch.Tensor:
    """Expose packed NTD views without copying the B=1 Prepare/FwdH storage."""

def _packed_offsets(cu_seqlens, tokens):
    """Snapshot host offsets before deriving native chunk metadata."""
    if isinstance(cu_seqlens, torch.Tensor):
        if (cu_seqlens.device.type != "cpu" or cu_seqlens.ndim != 1
                or cu_seqlens.dtype not in (torch.int32, torch.int64)):
            raise ValueError("Ascend C cu_seqlens must be a one-dimensional CPU integer tensor or sequence.")
        cu_seqlens = cu_seqlens.tolist()
    offsets = tuple(cu_seqlens)
    if (len(offsets) < 2 or offsets[0] != 0 or offsets[-1] != tokens
            or any(not isinstance(x, int) or isinstance(x, bool) for x in offsets)
            or any(begin >= end for begin, end in zip(offsets, offsets[1:]))):

        backward are not exposed. Native Q/K norm backward uses save mode;
        native capability checks select the supported device implementation.
    """
    if query.ndim == 3:
        if cu_seqlens is None:
            raise ValueError("Packed TND inputs require cu_seqlens.")
        return ascendc_chunk_kda(
            *(tensor.unsqueeze(0) for tensor in (query, key, value, gate, beta)),
            a_log=a_log, dt_bias=dt_bias, scale=scale, lower_bound=lower_bound,
            chunk_size=chunk_size, cu_seqlens=cu_seqlens).squeeze(0)
    if query.ndim != 4:
        raise ValueError("KDA inputs must have shape [B,T,H,K] or packed [T,H,K].")
    if chunk_size != 64:
        raise ValueError("Ascend C KDA uses chunk_size=64.")
    metadata = _packed_metadata(cu_seqlens, query.shape[1])
    if metadata and query.shape[0] != 1:
        raise ValueError("Packed Ascend C KDA requires B=1.")
    scale = query.shape[-1] ** -0.5 if scale is None else float(scale)
    tokens = tuple(tensor.transpose(1, 2).contiguous() for tensor in (query, key, value))
    # Gate/beta stay token-major for backward. Native operators validate their
    # dtype, shape and device; only adapter-owned layout/metadata checks live here.

@lru_cache(maxsize=1)
def _load_state_ops():
    """Check every native stage before entering the first state collective."""
    ops = _load_ops()
    required = (
        "pre_process_fwd_kernel_merged", "pre_process_bwd_kernel_merged",
        "merge_fwd_bwd_kernel", "chunk_kda_bwd_prepare",
        "chunk_gated_delta_rule_bwd_dhu", "chunk_kda_bwd_finalize",
    )
    missing = [name for name in required if not callable(getattr(ops, name, None))]
    if missing:
        raise RuntimeError(f"Native KDA state CP requires these fla-npu APIs: {', '.join(missing)}.")
    return ops


def _native_ag_intervals(info, tokens, rank, size):
    """Return tail/head windows and document-bounded exclusive scan lengths."""
    if info is None:
        return (0, tokens), (0, tokens), rank, size - rank - 1
    first_rank = (info.start - info.preceding_sample_tokens) // tokens
    last_rank = (info.global_offsets[bisect_left(info.global_offsets, info.end)] - 1) // tokens
    return ((info.tail.start, info.tail.stop), (info.head.start, info.head.stop),
            rank - first_rank if info.continues_previous else 0,
            last_rank - rank if info.continues_next else 0)

    """

    def __init__(self, mesh: Any, executor: Any, info: Any) -> None:
        """Copy communication metadata and retain only this call's sequence plan."""
        self.executor, self.info = executor, info
        self.rank = mesh.get_local_rank()
        self.size = mesh.size()
        self.ranks = tuple(mesh.rank_list)
        self.group = mesh.get_group()
        self.width = getattr(executor, "width", self.size) if executor is not None else 1
        self.block, self.local = divmod(self.rank, self.width)
        self.owner = self.block * self.width + self.width - 1

    def peer(self, forward: bool, receiving: bool) -> int | None:
        """Return a chronological peer, pruning disconnected ordinary P2P edges."""
        direction = 1 if forward else -1
        peer = self.owner + (-direction if receiving else direction) * self.width
        if not 0 <= peer < self.size:
            return None
        if self.width == 1 and self.info is not None:
            needs_previous = forward == receiving
            connected = self.info.continues_previous if needs_previous else self.info.continues_next
            if not connected:
                return None
        return peer

    def needs_summary(self, forward: bool) -> bool:
        """Avoid producing an ordinary P2P summary without a receiver."""
        return self.width > 1 or self.peer(forward, receiving=False) is not None

    def receive(self, reference: torch.Tensor, forward: bool) -> tuple[torch.Tensor | None, Any]:
        """Post reception before native prepare, independently of summary production."""
        if self.width == self.size:
            return None, None
        incoming = reference.new_zeros((1, reference.shape[1], 128, 128), dtype=torch.float32)
        peer = self.peer(forward, receiving=True)
        handle = None
        if self.rank == self.owner and peer is not None:
            handle = dist.irecv(incoming, src=self.ranks[peer], group=self.group)
        return incoming, handle

    @staticmethod
    def apply(packed: torch.Tensor, incoming: torch.Tensor | None) -> torch.Tensor:
        """Apply forward M or native backward N directly, without transposing N."""
        return apply_kda_state_summary(packed[..., :128], packed[..., 128:], incoming)

    def exchange(self, ops: Any, packed: torch.Tensor | None,
                 reception: tuple[torch.Tensor | None, Any], forward: bool,
                 count: int) -> tuple[torch.Tensor, Any, torch.Tensor | None]:

    def exchange(self, ops: Any, packed: torch.Tensor | None,
                 reception: tuple[torch.Tensor | None, Any], forward: bool,
                 count: int) -> tuple[torch.Tensor, Any, torch.Tensor | None]:
        """Produce the local boundary and retain an asynchronous outgoing payload."""
        if self.width == self.size:
            maps = self.executor.gather_packed(packed.reshape(1, *packed.shape[-3:]))
            incoming = maps.new_zeros((maps.shape[1], 128, 128))
            if count:
                ops.merge_fwd_bwd_kernel(incoming, maps, count, self.rank, forward=forward)
            return incoming.unsqueeze(0), None, None
        incoming, recv_handle = reception
        if self.width > 1:
            return self._grouped_exchange(packed, incoming, recv_handle, forward)
        transfer = self.info is None or self.info.num_sequences == 1
        if recv_handle is not None and transfer:
            recv_handle.wait()
            recv_handle = None
        send_handle, outgoing = None, None
        peer = self.peer(forward, receiving=False)
        if peer is not None:
            outgoing = self.apply(packed, incoming if transfer else None).contiguous()
            send_handle = dist.isend(outgoing, dst=self.ranks[peer], group=self.group)
        # A reset permits sending the independent tail/head before this wait,
        # but local h/dhu still consumes the received first/last fragment state.
        if recv_handle is not None:
            recv_handle.wait()
        return incoming, send_handle, outgoing

    def _grouped_exchange(self, packed, incoming, recv_handle, forward):
        """Gather native maps within a group and relay only its exit state."""
        if self.info is not None:
            if not self.info.transfer_active or not self.info.continues_next:
                packed[..., 128:].zero_()
            connected = self.info.continues_next if forward else self.info.continues_previous
            if not connected:
                packed[..., :128].zero_()
        maps = packed.new_empty((self.width, *packed.shape[-3:]))
        dist.all_gather_into_tensor(maps, packed.reshape(1, *packed.shape[-3:]), group=self.executor.intra)
        send_handle, outgoing, result = None, None, None
        indices = range(self.width) if forward else range(self.width - 1, -1, -1)
        prefix_count = self.local if forward else self.width - self.local - 1
        if self.rank == self.owner:
            if recv_handle is not None:
                recv_handle.wait()
            peer = self.peer(forward, receiving=False)
            if peer is not None:
                outgoing, result = self._scan(maps, incoming, indices, prefix_count)
                send_handle = dist.isend(outgoing, dst=self.ranks[peer], group=self.group)
        if self.peer(forward, receiving=True) is not None:
            dist.broadcast(incoming, src=self.ranks[self.owner], group=self.executor.intra)
        if result is None:
            local_indices = range(self.local) if forward else range(self.width - 1, self.local, -1)
            result, _ = self._scan(maps, incoming, local_indices)
        return result, send_handle, outgoing

    @staticmethod
    def _scan(maps, incoming, indices, prefix_count=None):
        """Keep native M/N orientation while applying a nonzero group entry."""
        state = incoming
        prefix = incoming if prefix_count == 0 else None
        for step, index in enumerate(indices, start=1):
            state = _NativeStateExchange.apply(maps[index:index + 1], state)
            if step == prefix_count:
                prefix = state
        return state, prefix


class _AscendCKDAState(torch.autograd.Function):
    """Use native S/M forward and newly generated G/N backward summaries."""

    def forward(ctx: Any, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor,
                gate: torch.Tensor, beta: torch.Tensor, a_log: torch.Tensor, dt_bias: torch.Tensor,
                scale: float, lower_bound: float, metadata: dict, exchange: _NativeStateExchange) -> torch.Tensor:
        """Share prepare with S/M and local FwdH, then release gathered maps."""
        ops = _load_state_ops()
        inputs = (query, key, value, gate, beta, a_log, dt_bias)
        sequence_info = exchange.info
        reception = exchange.receive(query, True)
        prepared = _prepare_kda(inputs, scale, lower_bound, metadata)
        tail, head, prefix, suffix = _native_ag_intervals(
            sequence_info, query.shape[2], exchange.rank, exchange.size)
        maps = None
        if exchange.needs_summary(True):
            maps = ops.pre_process_fwd_kernel_merged(
                prepared.kg, prepared.w, prepared.u, gk=prepared.gk, cu_seqlens=tail)
        initial, send_handle, outgoing = exchange.exchange(ops, maps, reception, True, prefix)
        del maps, reception
        if sequence_info is not None:
            initial = sequence_info.first_state(initial)
        hidden, v_new, _ = ops.chunk_fwd_h(
            prepared.kg, prepared.w, prepared.u, gk=prepared.gk, initial_state=initial,
            output_final_state=False, use_exp2=True, **metadata)
        output = ops.chunk_kda_fwd_finalize(
            prepared.qg_scaled, prepared.aqk, v_new, hidden, output_layout="BSND", **metadata)
        if send_handle is not None:
            send_handle.wait()
        del outgoing
        prepared = prepared._replace(u=None, qg_scaled=None)
        # The saved hidden tensor contains entry states. Summary protocols
        # produce the outgoing gradient themselves, so native dh0 is unused.
        ctx.save_for_backward(value, gate, beta, a_log, dt_bias, *prepared, hidden, v_new)
        ctx.dtypes = tuple(tensor.dtype for tensor in inputs)
        ctx.scale, ctx.lower_bound, ctx.native_metadata = scale, lower_bound, metadata
        ctx.exchange, ctx.head, ctx.suffix = exchange, head, suffix
        ctx.native_beta_backward = _use_native_beta_backward(query.device)
        if not ctx.native_beta_backward:
            _load_beta_backward()
        return output

    @staticmethod
    @once_differentiable
    def backward(ctx: Any, grad_output: torch.Tensor) -> tuple:

    @staticmethod
    @once_differentiable
    def backward(ctx: Any, grad_output: torch.Tensor) -> tuple:
        """Collect G/N, inject nonzero dHT and finish raw-input gradients."""
        ops = _load_state_ops()
        value, gate, beta_raw, a_log, dt_bias, *saved = ctx.saved_tensors
        (gk, aqk, akk, weights, _, qg, kg, _, query, key,
         query_rstd, key_rstd, beta, hidden, v_new) = saved
        packed = bool(ctx.native_metadata)

        reception = ctx.exchange.receive(query, False)
        grad = grad_output.transpose(1, 2).contiguous()
        d_aqk, dv_local, dq_raw = ops.chunk_kda_bwd_prepare(
            *(_native_backward_view(tensor, packed) for tensor in (aqk, v_new, grad, hidden)), scale=ctx.scale,
            layout="NTD" if packed else "BNSD", chunk_size=64, **ctx.native_metadata)
        dv_local_head = dv_local.unsqueeze(0) if packed else dv_local
        maps = None
        if ctx.exchange.needs_summary(False):
            # The stock summary ABI requires model-dtype gk. Keep the original
            # FP32 cumulative gate for local dhu/finalize; only this input is rounded.
            maps = ops.pre_process_bwd_kernel_merged(
                qg, kg, weights, grad, dv_local_head, ctx.scale, 64,
                gk=gk.to(qg.dtype), cu_seqlens=ctx.head)
        terminal, send_handle, outgoing = ctx.exchange.exchange(ops, maps, reception, False, ctx.suffix)
        del maps, reception
        if ctx.exchange.info is not None:
            terminal = ctx.exchange.info.last_state(terminal)
        # Communication and summary apply retain FP32. Convert once at the
        # stock dhu input; its local state gradient output is not used by CP.
        terminal = None if terminal is None else terminal.to(qg.dtype)
        hidden_grad, _, dv_scan = ops.chunk_gated_delta_rule_bwd_dhu(
            qg, kg, weights, grad, dv_local_head, ctx.scale, 64,
            gK=gk, h0=None, dht=terminal, use_exp2=True, **ctx.native_metadata)
        dq, dk, dv, dbeta, dgate, da_log, ddt_bias = ops.chunk_kda_bwd_finalize(
            *(_native_backward_view(tensor, packed)
              for tensor in (query, key, value, gk, gate.transpose(1, 2).contiguous(), beta)),
            a_log, dt_bias,
            *(_native_backward_view(tensor, packed) for tensor in (akk, v_new, hidden, hidden_grad, dv_scan)),

            q_rstd=_native_backward_view(query_rstd, packed),
            k_rstd=_native_backward_view(key_rstd, packed),
            use_beta_sigmoid_in_kernel=ctx.native_beta_backward,
            output_layout="NTD" if packed else "BNSD", **ctx.native_metadata)
        if packed:
            dq, dk, dv, dbeta, dgate = (tensor.unsqueeze(0) for tensor in (dq, dk, dv, dbeta, dgate))
        dbeta, dgate = dbeta.transpose(1, 2), dgate.transpose(1, 2)
        if not ctx.native_beta_backward:
            dbeta = _load_beta_backward()(beta_raw, dbeta.contiguous())
        if send_handle is not None:
            send_handle.wait()
        del outgoing
        gradients = (dq, dk, dv, dgate, dbeta, da_log, ddt_bias.reshape_as(dt_bias))
        return (*(grad.to(dtype) for grad, dtype in zip(gradients, ctx.dtypes)), None, None, None, None)


def ascendc_state_kda(
    query: torch.Tensor, key: torch.Tensor, value: torch.Tensor,

        backward summary gk and local dhu dHT are converted to model dtype.
        Prepared gates, communicated states and summary apply remain FP32.
        These ABI conversions introduce rounding relative to FP32-input kernels.
    """
    if query.device.type != "npu" or query.dtype != torch.bfloat16:
        raise ValueError("Native KDA state CP requires BF16 NPU inputs.")
    if query.ndim != 4 or query.shape[0] != 1 or query.shape[-1] != 128:
        raise ValueError("Native KDA state CP requires [1,T,H,128] inputs.")
    if any(tensor.shape != query.shape for tensor in (key, value, gate)) or beta.shape != query.shape[:-1]:
        raise ValueError("Native KDA state CP requires matching Q/K/V/gate heads and beta [1,T,H].")
    metadata = _packed_metadata(None if sequence_info is None else sequence_info.local_offsets, query.shape[1])
    scale = query.shape[-1] ** -0.5 if scale is None else float(scale)
    tokens = tuple(tensor.transpose(1, 2).contiguous() for tensor in (query, key, value))
    exchange = _NativeStateExchange(cp_mesh, state_cp, sequence_info)
    return _AscendCKDAState.apply(
        *tokens, gate.contiguous(), beta.contiguous(), a_log.reshape(-1).clone(),
        dt_bias.reshape(-1, query.shape[-1]).clone(), scale, lower_bound, metadata, exchange)

    """
    if initial_state is None:
        return state_ext
    # Model autocast must not silently quantize a communicated FP32 state.
    with torch.autocast(device_type=state_ext.device.type, enabled=False):
        return torch.matmul(transition.float(), initial_state.float()) + state_ext


def apply_kda_state_gradient_summary(
    grad_state_ext: torch.Tensor,

) -> torch.Tensor:
    """Apply the adjoint summary for the reverse P2P state wavefront."""
    if grad_final_state is None:
        return grad_state_ext
    with torch.autocast(device_type=grad_state_ext.device.type, enabled=False):
        return (
            torch.matmul(transition.float().transpose(-2, -1), grad_final_state.float())
            + grad_state_ext
        )

def _normalize_kda_backend(backend: str) -> str:
    """Normalize the selected backend and reject unsupported names."""
    backend = backend.lower()
    if backend not in _KDA_BACKENDS:
        raise ValueError(
            f"unsupported KDA backend {backend!r}; "
            f"expected one of {sorted(_KDA_BACKENDS)}."
        )
    return backend

        ).to(query.dtype)

        if send_handle is not None:
            send_handle.wait()
        del send_buffer

        empty = query.new_empty(0)
        ctx.save_for_backward(
            query_norm,

                f"expected one of {sorted(_GDN_BACKENDS)}."
            )
        if chunk_size <= 0:
            raise ValueError(f"chunk_size must be positive, got {chunk_size}.")
        if backend in ("triton", "ascendc") and chunk_size != 64:
            raise NotImplementedError("Fused GDN backends require chunk_size=64.")
        self.gdn_backend = backend
        self.chunk_size = chunk_size
        self.cp_mesh = None if device_mesh is None else _ensure_1d(device_mesh)
        self.cp_size = 1 if self.cp_mesh is None else self.cp_mesh.size()
        self.cp_rank = 0 if self.cp_mesh is None else self.cp_mesh.get_local_rank()
        self.seq_dim = 1
        self.head_dim = 2
        self._validate_module()
        if self.gdn_backend != "eager":

        if convolution.padding != (convolution.kernel_size[0] - 1,) or convolution.padding_mode != "zeros":
            raise ValueError("GDN fused convolution requires causal zero padding.")

    def _seq_to_head(self, tensor: torch.Tensor) -> torch.Tensor:
        if self.cp_size == 1:
            return tensor
        return _differentiable_all_to_all_shard(
            tensor,
            self.cp_mesh,
            split_dim=self.head_dim,

            concat_dim=self.seq_dim,
        )

    def _head_to_seq(self, tensor: torch.Tensor) -> torch.Tensor:
        if self.cp_size == 1:
            return tensor
        return _differentiable_all_to_all_shard(
            tensor,
            self.cp_mesh,
            split_dim=self.seq_dim,

        if kv_groups > 1:
            query = query.repeat_interleave(kv_groups, dim=2)
            key = key.repeat_interleave(kv_groups, dim=2)

        cu_seqlens = None
        if exchanged_info is not None:
            cu_seqlens = (exchanged_info.local_offsets if self.gdn_backend == "ascendc" else
                          exchanged_info.materialize(query.device, self.chunk_size).cu_seqlens)
        core_attn_out, _ = chunk_gated_delta_rule(
            query,
            key,

    ):
        super().__init__()
        self.module = module
        backend = backend.lower()
        if backend == "ascendc":
            raise NotImplementedError("GDN Ascend C state P2P is not supported; select the pure Ulysses wrapper.")
        if backend not in _GDN_BACKENDS:
            raise ValueError(
                f"unsupported GDN backend {backend!r}; "
                f"expected one of {sorted(_GDN_BACKENDS)}."

    backend: str,
    sequence_info: Optional[PackedSequenceInfo] = None,
) -> torch.Tensor:
    """Run the explicitly selected local KDA implementation."""
    if backend in ("triton", "ascendc") and not safe_gate:
        raise RuntimeError("The fused KDA backend requires the lower-bounded gate.")
    cu_seqlens = None
    if sequence_info is not None:
        # Native APIs consume host metadata; do not materialize an NPU plan
        # only to copy its offsets back to the host at every layer.
        cu_seqlens = (sequence_info.local_offsets if backend == "ascendc" else
                      sequence_info.materialize(query.device, chunk_size).cu_seqlens)
    output, _ = chunk_kda(
        query,
        key,

    sequence_info: Optional[PackedSequenceInfo] = None,
    backend: str = "eager",
) -> torch.Tensor:
    """Apply one local ShortConv, optionally with a preceding-rank halo."""
    if backend != "eager":
        return fused_causal_conv1d(
            tensor, convolution.weight, convolution.bias, sequence_info, halo,
            backend=backend, dilation=convolution.dilation[0])
    if sequence_info is not None:
        return packed_causal_conv1d(tensor, convolution, sequence_info, halo)

        self.cp_mesh = _ensure_1d(device_mesh)
        self.cp_size = self.cp_mesh.size()
        self.cp_rank = self.cp_mesh.get_local_rank()
        self.cp_group = self.cp_mesh.get_group()
        if backend == "ascendc" and self.cp_size != 1 and state_cp_method == "recursive_doubling":
            raise NotImplementedError("Ascend C KDA recursive doubling is not yet supported.")
        if backend == "ascendc" and (chunk_size != 64 or not safe_gate):
            raise ValueError("Ascend C KDA requires safe_gate=True and chunk_size=64.")
        if backend == "eager" and state_cp_method != "p2p":
            raise ValueError("Eager KDA state exchange supports P2P only.")
        self.state_cp = build_kda_state_cp(self.cp_mesh, state_cp_method, allgather_group_size)
        self.prev_rank = _global_peer_rank(
            self.cp_mesh,
            max(self.cp_rank - 1, 0),

                cp_size=self.cp_size,
                sequence_info=sequence_info,
            )

        if self.backend == "ascendc":
            # The native package remains optional; loading it belongs to the selected backend.
            from hyper_parallel.components.functional.kimi_delta_attention_ascendc import (  # pylint: disable=import-outside-toplevel
                ascendc_state_kda,
            )
            return ascendc_state_kda(
                query, key, value, gate, beta, a_log=a_log, dt_bias=dt_bias, scale=scale,
                lower_bound=self.lower_bound, state_cp=self.state_cp, cp_mesh=self.cp_mesh,
                sequence_info=sequence_info)

        device_mesh = _ensure_1d(device_mesh)
        if state_cp_method not in ("p2p", "allgather", "grouped_allgather_p2p", "recursive_doubling"):
            raise ValueError("KDA requires a supported state_cp_method.")
        if backend.lower() == "eager" and state_cp_method != "p2p":
            raise ValueError("Eager KDA state exchange supports P2P only.")
        _validate_cp_degrees(device_mesh.size(), ulysses_degree, allgather_group_size)
        if state_cp_method != "grouped_allgather_p2p" and allgather_group_size != 1:
            raise ValueError("allgather_group_size is only valid for grouped_allgather_p2p.")
        if lower_bound >= 0:

        The autograd caller owns the returned storage and its reverse protocol;
        this metadata-only executor retains no per-invocation tensors.
        """
        output = packed.new_empty((self.size, *packed.shape[1:]))
        dist.all_gather_into_tensor(output, packed, group=self.group)
        return output

    def forward(self, state: torch.Tensor, matrix: torch.Tensor) -> tuple[torch.Tensor, tuple[torch.Tensor, ...]]:
        """Return the incoming state and this invocation's original local M.