Diff Coverage

    if IS_VARLEN:
        v = v + (i_h * T_all + bos) * V
        k = k + (i_h * T_all + bos) * K
        w = w + (i_h * T_all + bos) * K
        if USE_G:
            g = g + i_h * T_all + bos
        h = h + (i_h * NT_all + boh) * K * V
        if SAVE_NEW_VALUE:
            v_new_base = v_new + (i_h * T_all + bos) * V
    else:

    else:
        v = v + (i_n * H + i_h) * T * V
        k = k + (i_n * H + i_h) * T * K
        w = w + (i_n * H + i_h) * T * K
        if USE_G:
            g = g + (i_n * H + i_h) * T
        h = h + (i_n * H + i_h) * NT * K * V
        if SAVE_NEW_VALUE:
            v_new_base = v_new + (i_n * H + i_h) * T * V

                b_h4 *= b_g_last_exp

        # Apply key gate gk
        if USE_GK:
            o_k1 = tl.arange(0, 64)
            if IS_VARLEN:
                gk_base_ptr = gk + (i_h * T_all + bos) * K
            else:
                gk_base_ptr = gk + (i_n * H + i_h) * T * K
            b_gk_last1 = tl.load(gk_base_ptr + last_idx * K + o_k1, mask=(o_k1 < K), other=0.)
            b_h1 *= tl.exp(b_gk_last1)[:, None]
            if K > 64:
                o_k2 = 64 + o_k1

    v_new = torch.empty_like(u).permute(0, 2, 1, 3).contiguous() if save_new_value else None
    k = k.permute(0, 2, 1, 3).contiguous()
    w = w.permute(0, 2, 1, 3).contiguous()
    u = u.permute(0, 2, 1, 3).contiguous()
    g = g.permute(0, 2, 1).contiguous() if g is not None else None
    # Match the head-major layout used by the forward kernel's token strides.
    gk = gk.permute(0, 2, 1, 3).contiguous() if gk is not None else None
    chunk_gated_delta_rule_fwd_kernel_h_blockdim64[(triton.cdiv(V, BV), N * H)](
        k=k,
        v=u,
        w=w,

        BV=BV,
        NT=NT,
    )
    h = h.permute(0, 2, 1, 3, 4).contiguous()
    v_new = v_new.permute(0, 2, 1, 3).contiguous() if v_new is not None else None
    return h, v_new, final_state


@triton.heuristics({

        )
        p_w2 = tl.make_block_ptr(
            w, (T, K), (stride_k, 1), (i_t * BT, 64), (BT, 64), (1, 0)
        )
        b_w1 = tl.load(p_w1, boundary_check=(0, 1), padding_option='zero')
        b_w2 = tl.load(p_w2, boundary_check=(0, 1), padding_option='zero')
        b_v = tl.dot(b_w1, b_h1.to(b_w1.dtype))
        b_v += tl.dot(b_w2, b_h2.to(b_w2.dtype))

        p_u = tl.make_block_ptr(

            (i_t * BT, i_v * BV),
            (BT, BV),
            (1, 0),
        )
        b_v = tl.load(p_u, boundary_check=(0, 1), padding_option='zero') - b_v

        last_idx = min((i_t + 1) * BT, T) - 1
        token = i_t * BT + tl.arange(0, BT)
        mask = token < T

        )
        p_k2 = tl.make_block_ptr(
            k, (K, T), (1, stride_k), (64, i_t * BT), (64, BT), (0, 1)
        )
        b_h1 += tl.dot(tl.load(p_k1, boundary_check=(0, 1), padding_option='zero'), b_v)
        b_h2 += tl.dot(tl.load(p_k2, boundary_check=(0, 1), padding_option='zero'), b_v)

    p_out1 = tl.make_block_ptr(
        packed_summary,
        (K, V + K),

        )
        p_k2 = tl.make_block_ptr(
            k, (T, K), (stride_k, 1), (i_t * BT, 64), (BT, 64), (1, 0)
        )
        b_k1 = tl.load(p_k1, boundary_check=(0, 1), padding_option='zero')
        b_k2 = tl.load(p_k2, boundary_check=(0, 1), padding_option='zero')
        b_dv = tl.dot(b_k1, b_dh1.to(b_k1.dtype))
        b_dv += tl.dot(b_k2, b_dh2.to(b_k2.dtype))
        b_dv *= tl.where(mask, tl.exp(b_g_last - b_g), 0.0)[:, None]

            (i_t * BT, i_v * BV),
            (BT, BV),
            (1, 0),
        )
        b_dv += tl.load(p_dv, boundary_check=(0, 1), padding_option='zero')

        p_do = tl.make_block_ptr(
            do,
            (T, V),

            (i_t * BT, i_v * BV),
            (BT, BV),
            (1, 0),
        )
        b_do = tl.load(p_do, boundary_check=(0, 1), padding_option='zero')
        decay = tl.exp(b_g_last)
        b_dh1 *= decay
        b_dh2 *= decay

        )
        p_w2 = tl.make_block_ptr(
            w, (K, T), (1, stride_k), (64, i_t * BT), (64, BT), (0, 1)
        )
        b_q1 = tl.load(p_q1, boundary_check=(0, 1), padding_option='zero')
        b_q2 = tl.load(p_q2, boundary_check=(0, 1), padding_option='zero')
        b_w1 = tl.load(p_w1, boundary_check=(0, 1), padding_option='zero')
        b_w2 = tl.load(p_w2, boundary_check=(0, 1), padding_option='zero')
        gate = tl.exp(b_g)[None, :]
        b_q1 *= gate
        b_q2 *= gate
        # Match native dhu's rounding: form the signed update before adding it

        b_q1 *= gate
        b_q2 *= gate
        # Match native dhu's rounding: form the signed update before adding it
        # to the recurrent accumulator, which is rounded again next chunk.
        b_dh1 += tl.dot(b_q1, b_do.to(b_q1.dtype)) * scale - tl.dot(b_w1, b_dv.to(b_w1.dtype))
        b_dh2 += tl.dot(b_q2, b_do.to(b_q2.dtype)) * scale - tl.dot(b_w2, b_dv.to(b_w2.dtype))

    p_out1 = tl.make_block_ptr(
        grad_state_ext, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0)
    )

    num_chunks = tl.cdiv(T, BT)
    for chunk_idx in range(num_chunks):
        w1_ptr = tl.make_block_ptr(w, (T, K), (stride_w, 1), (chunk_idx * BT, 0), (BT, 64), (1, 0))
        w2_ptr = tl.make_block_ptr(w, (T, K), (stride_w, 1), (chunk_idx * BT, 64), (BT, 64), (1, 0))
        w1 = tl.load(w1_ptr, boundary_check=(0, 1), padding_option='zero')
        w2 = tl.load(w2_ptr, boundary_check=(0, 1), padding_option='zero')
        if OUTPUT_MODE == 1:
            # Match Ascend FLA h: state and residual feedback stay in FP32.
            w1 = w1.to(tl.float32)
            w2 = w2.to(tl.float32)

        state2 *= _exp2(gate_last2)[:, None]

        key1_ptr = tl.make_block_ptr(key, (K, T), (1, stride_k), (0, chunk_idx * BT), (64, BT), (0, 1))
        key2_ptr = tl.make_block_ptr(key, (K, T), (1, stride_k), (64, chunk_idx * BT), (64, BT), (0, 1))
        key1 = tl.load(key1_ptr, boundary_check=(0, 1), padding_option='zero')
        key2 = tl.load(key2_ptr, boundary_check=(0, 1), padding_option='zero')
        if OUTPUT_MODE == 1:
            key1 = key1.to(tl.float32)
            key2 = key2.to(tl.float32)
        state1 += tl.dot(key1, value_new)

    # Ascend cannot carry the Cube result directly into the next Cube input.
    # Each program owns disjoint output columns, so this FP32 feedback needs
    # neither an inter-program barrier nor a per-chunk matrix workspace.
    tl.store(transition + output_offsets, matrix)
    for chunk in range(tl.cdiv(T, BT)):
        offsets = ((batch * T + chunk * BT + tokens[:, None]) * H + local_head) * K + rows[None, :]
        key_tile = tl.load(key + offsets, mask=chunk * BT + tokens[:, None] < T, other=0.0)
        weight_tile = tl.load(w + offsets, mask=chunk * BT + tokens[:, None] < T, other=0.0)
        last = ((batch * T + tl.minimum((chunk + 1) * BT, T) - 1) * H + local_head) * K + rows
        decay = tl.exp2(tl.load(gate + last).to(tl.float32))
        diagonal = tl.where(rows[:, None] == rows[None, :], decay[:, None], 0.0)
        interaction = tl.dot(tl.trans(key_tile), weight_tile)
        chunk_transition = diagonal - interaction

        chunk_idx = num_chunks - 1 - reverse_idx

        key1_ptr = tl.make_block_ptr(key, (T, K), (stride_k, 1), (chunk_idx * BT, 0), (BT, 64), (1, 0))
        key2_ptr = tl.make_block_ptr(key, (T, K), (stride_k, 1), (chunk_idx * BT, 64), (BT, 64), (1, 0))
        key1 = tl.load(key1_ptr, boundary_check=(0, 1), padding_option='zero')
        key2 = tl.load(key2_ptr, boundary_check=(0, 1), padding_option='zero')
        value_grad = tl.dot(key1, state1.to(key1.dtype))
        value_grad += tl.dot(key2, state2.to(key2.dtype))

        grad_value_ptr = tl.make_block_ptr(

            (chunk_idx * BT, i_col * BV),
            (BT, BV),
            (1, 0),
        )
        value_grad += tl.load(grad_value_ptr, boundary_check=(0, 1), padding_option='zero')

        last_idx = min((chunk_idx + 1) * BT, T) - 1
        gate_last_ptr = gate + last_idx * H_TOTAL * K
        gate_last1 = tl.load(gate_last_ptr + row1, mask=row1 < K, other=0.0)

            (chunk_idx * BT, i_col * BV),
            (BT, BV),
            (1, 0),
        )
        output_grad = tl.load(output_grad_ptr, boundary_check=(0, 1), padding_option='zero')
        # Match FLA dhu's scale-before-cast contract for the Q contribution.
        output_grad = output_grad * scale
        output_grad_second = output_grad + 0.0

        query1_ptr = tl.make_block_ptr(query, (K, T), (1, stride_k), (0, chunk_idx * BT), (64, BT), (0, 1))
        query2_ptr = tl.make_block_ptr(query, (K, T), (1, stride_k), (64, chunk_idx * BT), (64, BT), (0, 1))
        w1_ptr = tl.make_block_ptr(w, (K, T), (1, stride_k), (0, chunk_idx * BT), (64, BT), (0, 1))
        w2_ptr = tl.make_block_ptr(w, (K, T), (1, stride_k), (64, chunk_idx * BT), (64, BT), (0, 1))
        query1 = tl.load(query1_ptr, boundary_check=(0, 1), padding_option='zero')
        query2 = tl.load(query2_ptr, boundary_check=(0, 1), padding_option='zero')
        w1 = tl.load(w1_ptr, boundary_check=(0, 1), padding_option='zero')
        w2 = tl.load(w2_ptr, boundary_check=(0, 1), padding_option='zero')
        state1 += tl.dot(query1, output_grad.to(query1.dtype)) - tl.dot(w1, value_grad.to(w1.dtype))
        state2 += tl.dot(query2, output_grad_second.to(query2.dtype)) - tl.dot(w2, value_grad.to(w2.dtype))

    output1_ptr = tl.make_block_ptr(grad_state_ext, (K, V), (V, 1), (0, i_col * BV), (64, BV), (1, 0))

    value_dim = u.shape[-1]
    _validate_fixed_summary_shape(key_dim, value_dim, chunk_size)
    if block_size not in (64, 128):
        raise ValueError(f"GDN state-summary block_size must be 64 or 128, got {block_size}.")
    if seq_len <= 0:
        raise ValueError("State summary requires a nonempty sample fragment.")
    if w.shape != key.shape or u.shape[:3] != key.shape[:3] or g.shape != key.shape[:3]:
        raise ValueError(
            "Incompatible GDN state-summary shapes: "
            f"key={tuple(key.shape)}, w={tuple(w.shape)}, "

        BV=block_size,
        NT=triton.cdiv(seq_len, chunk_size),
    )
    state_ext = packed_summary[..., :value_dim].contiguous()
    if not output_transition:
        return state_ext, None
    transition = packed_summary[..., value_dim:].contiguous()
    return state_ext, transition

    """Validate the reverse summary inputs before launching its kernel."""
    batch, seq_len, heads, key_dim = query.shape
    value_dim = grad_output.shape[-1]
    _validate_fixed_summary_shape(key_dim, value_dim, chunk_size)
    if seq_len <= 0:
        raise ValueError("State summary requires a nonempty sample fragment.")
    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

import torch
import torch.distributed as dist

from hyper_parallel.components.functional.linear_attention_packed import PackedSequenceInfo, resolve_packed_sequence

from .kimi_delta_attention_fla_adapter import get_fla_kda_staged_ops, run_fla_chunk_kda
from .kimi_delta_attention_state_summary import (
    apply_kda_state_gradient_summary,

    chunk_size: int,
    lower_bound: float,
) -> None:
    """Validate the fixed-shape local KDA backend before compiling kernels."""
    _, 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

            cp_size: Number of state partitions.
            state_cp: Optional gather executor; None selects P2P.
        """
        ops = get_fla_kda_staged_ops()
        plan = None if sequence_info is None else sequence_info.materialize(query.device, chunk_size)
        varlen = {} if plan is None else plan.kwargs()
        recurrent_varlen = dict(varlen)
        if plan is not None:
            recurrent_varlen["chunk_offsets"] = plan.chunk_offsets

        # FLA consumes base-2 log gates: exp(x) = exp2(x / ln(2)).
        # Use the same conversion when recomputing gates in backward.
        rcp_ln2 = 1.4426950408889634

            **varlen,
        )

        state_shape = (query.shape[0], value.shape[2], query.shape[-1], value.shape[-1])
        receive_previous = cp_rank > 0 and (sequence_info is None or sequence_info.continues_previous)
        send_next = cp_rank < cp_size - 1 and (sequence_info is None or sequence_info.continues_next)
        initial_state = None
        recv_handle = None
        if state_cp is None and receive_previous:
            initial_state = torch.empty(state_shape, device=query.device, dtype=torch.float32)
            recv_handle = dist.irecv(initial_state, src=prev_rank, group=cp_group)

        w, u, _, kg, attention_qk, attention_kk = ops.chunk_kda_fwd_intra(

            **varlen,
        )
        state_ext = None
        transition = None
        if sequence_info is not None:
            if state_cp is not None or send_next:
                tail = sequence_info.tail
                state_ext, transition = kda_state_summary_forward_from_prepared(
                    kg[:, tail], w[:, tail], u[:, tail], gate[:, tail], chunk_size=chunk_size,
                    output_transition=sequence_info.transfer_active,
                )
                if state_cp is not None and not sequence_info.transfer_active:
                    # Collective protocols still require an explicit zero transfer.
                    transition = state_ext.new_zeros((*state_shape[:-1], query.shape[-1]))
        else:
            if state_cp is None and cp_rank == 0 and cp_size > 1:
                # Zero entry uses S directly, and rank 0 has no reverse-state recipient.
                state_ext, transition = kda_state_summary_forward_from_prepared(
                    kg, w, u, gate, chunk_size=chunk_size, output_transition=False,
                )
            elif state_cp is not None or cp_rank < cp_size - 1:
                state_ext, transition = kda_state_summary_forward_from_prepared(kg, w, u, gate, chunk_size=chunk_size)
        # A reset makes the outgoing summary independent of the incoming state.
        # Wait before apply only when it actually reads that state; local h still waits below.
        if recv_handle is not None and (sequence_info is None or sequence_info.transfer_active):
            recv_handle.wait()
            recv_handle = None

        send_handle = None
        send_state = None
        if state_cp is not None:

        send_handle = None
        send_state = None
        if state_cp is not None:
            initial_state = state_cp.forward(state_ext, transition)
        elif send_next:
            summary_initial = initial_state if transition is not None else None
            final_state = apply_kda_state_summary(state_ext, transition, summary_initial)
            send_state = final_state.contiguous()
            send_handle = dist.isend(send_state, dst=next_rank, group=cp_group)

        if recv_handle is not None:
            recv_handle.wait()
        kernel_initial_state = initial_state if sequence_info is None else sequence_info.first_state(initial_state)
        states, value_new, _ = ops.chunk_gated_delta_rule_fwd_h(
            k=kg, w=w, u=u, gk=gate, initial_state=kernel_initial_state,
            output_final_state=False, chunk_size=chunk_size,
            **recurrent_varlen,

        saved_initial_state = initial_state
        if saved_initial_state is None:
            saved_initial_state = query.new_empty(0, dtype=torch.float32)
        packed_tensors = () if plan is None else (
            plan.cu_seqlens, plan.chunk_indices, plan.chunk_offsets)
        ctx.save_for_backward(
            query, query_rstd, key, key_rstd, value, gate_raw, beta_raw, a_log, dt_bias, attention_qk,
            attention_kk, transition if transition is not None else query.new_empty(0), saved_initial_state,

        ctx.next_rank = next_rank
        ctx.cp_rank = cp_rank
        ctx.cp_size = cp_size
        ctx.state_cp = state_cp
        ctx.receive_previous = receive_previous
        ctx.send_next = send_next
        ctx.sequence_info = sequence_info
        return output.to(value.dtype)

    @staticmethod
    def backward(  # pylint: disable=arguments-differ,too-many-locals

        (
            query, query_rstd, key, key_rstd, value, gate_raw, beta_raw, a_log, dt_bias,
            attention_qk, attention_kk, transition, saved_initial_state, *packed_tensors,
        ) = ctx.saved_tensors
        sequence_info = ctx.sequence_info
        cu_seqlens = chunk_indices = chunk_offsets = None
        if sequence_info is not None:
            cu_seqlens, chunk_indices, chunk_offsets = packed_tensors
        varlen = {} if sequence_info is None else {"cu_seqlens": cu_seqlens, "chunk_indices": chunk_indices}
        recurrent_varlen = dict(varlen)
        if sequence_info is not None:
            recurrent_varlen["chunk_offsets"] = chunk_offsets
        initial_state = saved_initial_state if saved_initial_state.numel() else None
        kernel_initial_state = initial_state if sequence_info is None else sequence_info.first_state(initial_state)
        grad_output = grad_output.contiguous()

        grad_final_state = None
        recv_handle = None
        if ctx.state_cp is None and ctx.send_next:
            grad_final_state = torch.empty(
                (query.shape[0], value.shape[2], query.shape[-1], value.shape[-1]), device=query.device,
                dtype=torch.float32,
            )

            chunk_size=ctx.chunk_size,
            **varlen,
        )
        grad_state_ext = None
        if ctx.state_cp is not None or ctx.receive_previous:
            if sequence_info is not None and not sequence_info.continues_previous:
                grad_state_ext = query.new_zeros(
                    (1, value.shape[2], query.shape[-1], value.shape[-1]), dtype=torch.float32)
            else:
                head = slice(None) if sequence_info is None else sequence_info.head
                grad_state_ext = kda_state_gradient_summary_from_prepared(
                    query_gated[:, head], key_gated[:, head], w[:, head], gate[:, head],
                    grad_output[:, head], grad_value_local[:, head], ctx.scale, chunk_size=ctx.chunk_size,
                )
        if recv_handle is not None and (sequence_info is None or sequence_info.transfer_active):
            recv_handle.wait()
            recv_handle = None
        send_handle = None
        send_state_gradient = None
        if ctx.state_cp is not None:
            grad_final_state = ctx.state_cp.backward(grad_state_ext, transition)
        elif ctx.receive_previous:
            summary_final_gradient = grad_final_state if transition.numel() else None
            grad_initial_state = apply_kda_state_gradient_summary(grad_state_ext, transition, summary_final_gradient)
            send_state_gradient = grad_initial_state.contiguous()
            send_handle = dist.isend(send_state_gradient, dst=ctx.prev_rank, group=ctx.cp_group)

        if recv_handle is not None:
            recv_handle.wait()
        kernel_grad_final = grad_final_state if sequence_info is None else sequence_info.last_state(grad_final_state)
        grad_states, _, grad_value = ops.chunk_gated_delta_rule_bwd_dhu(
            q=query_gated, k=key_gated, w=w, gk=gate, h0=kernel_initial_state, dht=kernel_grad_final, do=grad_output,
            dv=grad_value_local, scale=ctx.scale, chunk_size=ctx.chunk_size,
            **recurrent_varlen,

                dim=3,
            )
            grad_key = grad_key.view(*grad_key.shape[:2], num_query_heads, groups, grad_key.shape[-1]).sum(dim=3)

        grad_gate = ops.chunk_local_cumsum(grad_gate, chunk_size=ctx.chunk_size, reverse=True, **varlen)
        grad_gate, grad_a_log, grad_dt_bias = ops.kda_gate_bwd(
            g=gate_raw, A_log=a_log, dt_bias=dt_bias, dyg=grad_gate, lower_bound=ctx.lower_bound,
        )
        grad_beta = ops.fused_beta_sigmoid_bwd(beta_raw, grad_beta)

        lower_bound=lower_bound,
        cp_rank=cp_rank,
        cp_size=cp_size,
    )
    if sequence_info is not None:
        if query.shape[0] != 1:
            raise ValueError("Packed state CP requires batch dimension one.")
        resolve_packed_sequence({"packed_sequence_info": sequence_info}, query.shape[1],
                                cp_rank=cp_rank, cp_size=cp_size)
    effective_scale = query.shape[-1] ** -0.5 if scale is None else float(scale)
    return _KDAStateP2PFunction.apply(
        query,

        raise NotImplementedError(
            "Triton KDA state summary requires key_dim=value_dim=128 and "
            "chunk_size=64."
        )
    if sequence_length <= 0:
        raise ValueError("State summary requires a nonempty sample fragment.")
    if w.shape != key.shape or gate.shape != key.shape:
        raise ValueError(
            "KDA prepared key, w, and gate must have identical shapes, got "
            f"key={tuple(key.shape)}, w={tuple(w.shape)}, gate={tuple(gate.shape)}."

        raise NotImplementedError(
            "Triton KDA state-gradient summary requires key_dim=value_dim=128 "
            "and chunk_size=64."
        )
    if sequence_length <= 0:
        raise ValueError("State summary requires a nonempty sample fragment.")
    prepared_shape = (batch, sequence_length, heads, key_dim)
    value_shape = (batch, sequence_length, heads, value_dim)
    if key.shape != prepared_shape or w.shape != prepared_shape:
        raise ValueError(

    Raises:
        ValueError: If the backend is unsupported.
    """
    if backend not in ("eager", "triton", "ascendc"):
        raise ValueError(f"Unsupported convolution backend {backend!r}; expected eager, triton, or ascendc.")
    return backend


class _AscendCConvolution(torch.autograd.Function):

        offsets: tuple[int, ...] | None,
    ) -> torch.Tensor:
        """Save preactivation and per-call metadata for joint backward."""
        # Optional native operators must not be loaded by ordinary framework imports.
        from fla_npu.ops.ascendc import causal_conv1d_fn  # pylint: disable=import-outside-toplevel

        op_inputs = inputs if offsets is None else inputs.reshape(-1, inputs.shape[-1])
        op_inputs = op_inputs.contiguous()
        op_weight = weight[:, 0].transpose(0, 1).contiguous()
        cpu_offsets = None if offsets is None else torch.tensor(offsets, dtype=torch.int64, device="cpu")
        preactivation = causal_conv1d_fn(
            op_inputs, op_weight, bias, conv_states=None, activation=None,
            query_start_loc_cpu=cpu_offsets, head_num=0,
        )
        ctx.save_for_backward(op_inputs, op_weight, preactivation, cpu_offsets)
        ctx.input_shape = tuple(inputs.shape)
        ctx.has_bias = bias is not None
        ctx.input_layout = "BSND" if offsets is None else "TND"
        return F.silu(preactivation).reshape_as(inputs)

    @staticmethod
    def backward(
        ctx: Any, grad_output: torch.Tensor,

        ctx: Any, grad_output: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None, None]:
        """Differentiate SiLU once inside the native convolution backward."""
        # Keep the optional native dependency lazy in backward as well as forward.
        from fla_npu.ops.ascendc import causal_conv1d_bwd  # pylint: disable=import-outside-toplevel

        inputs, weight, preactivation, cpu_offsets = ctx.saved_tensors
        grad_input, grad_weight, grad_bias, _ = causal_conv1d_bwd(
            x=inputs, y=preactivation, weight=weight,
            dy=grad_output.reshape_as(inputs).contiguous(),
            query_start_loc=cpu_offsets, activation=1, input_layout=ctx.input_layout,
        )
        return (
            grad_input.reshape(ctx.input_shape),
            grad_weight.transpose(0, 1).unsqueeze(1).contiguous(),
            grad_bias if ctx.has_bias else None,
            None,

    info: PackedSequenceInfo | None,
    width: int,
) -> tuple[torch.Tensor, tuple[int, ...] | None, int]:
    """Prepend only the first fragment's true history without inventing a reset."""
    if info is None:
        return inputs, None, 0
    if halo is None:
        return inputs, info.local_offsets, 0
    prefix = min(width - 1, info.preceding_sample_tokens)
    # The empty slice preserves a zero halo gradient when a real document resets.
    extended = torch.cat((halo[:, width - 1 - prefix:], inputs), dim=1)
    offsets = (0, *(offset + prefix for offset in info.local_offsets[1:]))
    return extended, offsets, prefix


def _triton_convolution(
    inputs: torch.Tensor,

    use_plan: bool,
) -> torch.Tensor:
    """Use one FLA entry for dense batches and immutable packed document plans."""
    # FLA is optional and must not be imported for the default eager backend.
    from fla.modules.conv import causal_conv1d  # pylint: disable=import-outside-toplevel

    if info is None:
        cu_seqlens, cu_seqlens_cpu, chunk_indices = None, None, None
    elif use_plan:
        plan = info.materialize(inputs.device, 64)
        cu_seqlens, cu_seqlens_cpu = plan.cu_seqlens, plan.cu_seqlens_cpu
        chunk_indices = plan.chunk_indices
    else:
        cu_seqlens_cpu = torch.tensor(offsets, dtype=torch.int64, device="cpu")
        cu_seqlens = cu_seqlens_cpu.to(device=inputs.device, dtype=torch.int32)
        chunk_indices = None
    # Contiguous BTD enables the NPU kernel without its strided fallback workspace.
    output, _ = causal_conv1d(
        inputs.contiguous(), weight=weight[:, 0], bias=bias, activation=None, backend="triton",
        cu_seqlens=cu_seqlens, cu_seqlens_cpu=cu_seqlens_cpu, chunk_indices=chunk_indices,
    )
    return F.silu(output)


def _validate_fused_inputs(
    inputs: torch.Tensor,

    info: PackedSequenceInfo | None,
    halo: torch.Tensor | None,
) -> None:
    """Check layouts before a selected backend is allowed to launch kernels."""
    if inputs.ndim != 3 or any(size < 1 for size in inputs.shape):
        raise ValueError("Fused convolution expects nonempty [B,S,C] or packed [1,T,C] inputs.")
    if weight.ndim != 3 or weight.shape[:2] != (inputs.shape[-1], 1):
        raise ValueError("Fused convolution weight must have shape [C,1,W] matching input channels.")
    if bias is not None and bias.shape != (inputs.shape[-1],):
        raise ValueError("Fused convolution bias must have shape [C] matching input channels.")
    if weight.device != inputs.device or (bias is not None and bias.device != inputs.device):
        raise ValueError("Fused convolution input, weight and bias must be on the same device.")
    if info is not None and inputs.shape[:2] != (1, info.length):
        raise ValueError("Packed convolution expects [1,T,C] matching its token interval.")
    if halo is not None:
        _validate_fused_history(inputs, halo, info, weight.shape[-1])


def _validate_fused_history(
    inputs: torch.Tensor,

    info: PackedSequenceInfo | None,
    width: int,
) -> None:
    """Reject history that cannot represent the same document and activation layout."""
    if info is None:
        raise ValueError("Fused convolution halo requires explicit packed sequence metadata.")
    if halo.shape != (1, width - 1, inputs.shape[-1]):
        raise ValueError("Fused convolution halo has an incompatible shape.")
    if halo.dtype != inputs.dtype or halo.device != inputs.device:
        raise ValueError("Fused convolution halo must match the input dtype and device.")


def fused_causal_conv1d(
    inputs: torch.Tensor,

        preactivation. Unlike eager or Triton, its backward does not round the SiLU
        gradient to the input dtype before convolution backward. Dense batch samples
        remain separate; only packed inputs with explicit metadata are flattened.
    """
    if backend not in ("triton", "ascendc"):
        raise ValueError(f"Fused convolution requires triton or ascendc, got {backend!r}.")
    if dilation != 1:
        raise ValueError("Fused convolution requires dilation=1.")
    if inputs.device.type != "npu":
        raise ValueError("Fused convolution currently supports NPU tensors only.")
    _validate_fused_inputs(inputs, weight, bias, info, halo)
    width = weight.shape[-1]
    if backend == "ascendc":
        if width not in (2, 3, 4) or inputs.shape[-1] % 16:
            raise ValueError("AscendC convolution requires width 2/3/4 and channels divisible by 16.")
        if inputs.dtype not in (torch.bfloat16, torch.float16):
            raise ValueError("AscendC convolution requires BF16 or FP16 activations.")
        if weight.dtype != inputs.dtype or (bias is not None and bias.dtype != inputs.dtype):
            raise ValueError("AscendC convolution requires matching input/weight/bias dtypes; no implicit cast.")
    extended, offsets, prefix = _with_history(inputs, halo, info, width)
    if backend == "triton":
        output = _triton_convolution(extended, weight, bias, info, offsets, halo is None)
    else:
        output = _AscendCConvolution.apply(extended, weight, bias, offsets)
    return output[:, prefix:]

    chunk_offsets: torch.Tensor

    def kwargs(self) -> dict[str, torch.Tensor]:
        """Return the common FLA varlen arguments without copying device tensors."""
        return {"cu_seqlens": self.cu_seqlens, "chunk_indices": self.chunk_indices}


@dataclass(frozen=True)
class PackedSequenceInfo:

            not isinstance(self.start, int) or isinstance(self.start, bool)
            or not isinstance(self.end, int) or isinstance(self.end, bool)
            or not 0 <= self.start < self.end <= offsets[-1]
        ):
            raise ValueError("Packed token interval must be nonempty and covered by global sample offsets.")
        local = (0, *(x - self.start for x in offsets if self.start < x < self.end), self.end - self.start)
        object.__setattr__(self, "global_offsets", offsets)
        object.__setattr__(self, "local_offsets", local)

    @property
    def continues_next(self) -> bool:
        """Whether the last fragment feeds the following interval's state."""
        return self.end not in self.global_offsets

    @property
    def transfer_active(self) -> bool:
        """Whether the preceding interval can affect this interval's exit."""

    @property
    def head(self) -> slice:
        """Return the first fragment's local token slice."""
        return slice(0, self.local_offsets[1])

    @property
    def tail(self) -> slice:
        """Return the last fragment's local token slice."""
        return slice(self.local_offsets[-2], self.length)

    @property
    def preceding_sample_tokens(self) -> int:
        """Return how many earlier tokens belong to the first local sample."""
        index = bisect_right(self.global_offsets, self.start) - 1
        return self.start - self.global_offsets[index]

    def interval(self, start: int, end: int) -> PackedSequenceInfo:
        """Describe another contiguous interval without retaining artificial cuts."""
        key = (start, end)

    def interval(self, start: int, end: int) -> PackedSequenceInfo:
        """Describe another contiguous interval without retaining artificial cuts."""
        key = (start, end)
        if key == (self.start, self.end):
            return self
        if key not in self._intervals:
            self._intervals[key] = PackedSequenceInfo(self.global_offsets, start, end)
        return self._intervals[key]

    def ulysses(self, degree: int, local_rank: int) -> PackedSequenceInfo:
        """Describe the token interval after consecutive-rank AllToAll."""
        if degree < 1 or not 0 <= local_rank < degree:
            raise ValueError("Invalid Ulysses degree or local rank.")
        start = self.start - local_rank * self.length
        return self.interval(start, start + degree * self.length)

    def materialize(self, device: torch.device, chunk_size: int) -> PackedChunkPlan:

        cache_key = (device, chunk_size)
        if cache_key in self._plans:
            return self._plans[cache_key]
        if chunk_size < 1:
            raise ValueError("chunk_size must be positive.")
        counts = [
            (end - start + chunk_size - 1) // chunk_size
            for start, end in zip(self.local_offsets, self.local_offsets[1:])
        ]