Diff Coverage

def _l2norm(x: torch.Tensor, dim: int = -1, eps: float = 1e-6) -> torch.Tensor:
    """L2-normalize along ``dim``."""
    # ``(x * x).sum`` (MulBackward) instead of ``x.pow(2).sum`` (PowBackward):
    # equal math, NPU yields ULP-different gradients across the two ops.
    return x * torch.rsqrt((x * x).sum(dim=dim, keepdim=True) + eps)


def torch_chunk_gated_delta_rule(
    query: torch.Tensor,

    beta:         ``(B, S, num_v_heads)``
    """
    initial_dtype = query.dtype
    if use_qk_l2norm_in_kernel:
        query = _l2norm(query, dim=-1, eps=1e-6)
        key = _l2norm(key, dim=-1, eps=1e-6)
    # (B, S, H, D) → (B, H, S, D), fp32
    query, key, value, beta, g = [
        x.transpose(1, 2).contiguous().to(torch.float32)
        for x in (query, key, value, beta, g)

            + (k_i * (g[:, :, i, -1, None] - g[:, :, i]).exp()[..., None]).transpose(-1, -2) @ v_new
        )

    if not output_final_state:
        last_recurrent_state = None
    core_attn_out = core_attn_out.reshape(
        core_attn_out.shape[0], core_attn_out.shape[1], -1, core_attn_out.shape[-1],
    )
    core_attn_out = core_attn_out[:, :, :sequence_length]

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


def _is_triton_gdn_input_supported(
    query: torch.Tensor,

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

        and query.shape[:3] == value.shape[:3] == g.shape == beta.shape
        and query.shape[-1] == value.shape[-1] == _TRITON_GDN_HEAD_DIM
        and query.shape[1] % _TRITON_GDN_CHUNK_SIZE == 0
    ):
        return False
    return all(
        tensor.device == query.device
        for tensor in (key, value, g, beta)
    )

    beta: Optional[torch.Tensor] = None,
    chunk_size: int = _TRITON_GDN_CHUNK_SIZE,
) -> bool:
    """Return whether the validated Triton-Ascend GDN backend is available."""
    if chunk_size != _TRITON_GDN_CHUNK_SIZE:
        return False
    if query is not None and not _is_triton_gdn_input_supported(
        query, key, value, g, beta
    ):
        return False
    try:
        version_text = importlib.metadata.version("triton-ascend")
    except importlib.metadata.PackageNotFoundError:
        return False
    version = _parse_version(version_text)
    if version < _MIN_TRITON_ASCEND_VERSION:
        return False
    try:
        importlib.import_module("triton")
        return importlib.util.find_spec("triton.backends.ascend") is not None
    except (AttributeError, ImportError, ModuleNotFoundError):
        return False


def chunk_gated_delta_rule(
    query: torch.Tensor,

    use_qk_l2norm_in_kernel: bool = False,
    backend: str = "eager",
):
    """Dispatch GDN to an explicitly selected eager or Triton backend."""
    backend = backend.lower()
    if backend not in _GDN_BACKENDS:
        raise ValueError(
            f"unsupported GDN backend {backend!r}; "
            f"expected one of {sorted(_GDN_BACKENDS)}."
        )
    if backend == "eager":
        return torch_chunk_gated_delta_rule(
            query,
            key,
            value,
            g=g,

            initial_state=initial_state,
            output_final_state=output_final_state,
            use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
        )
    if not is_triton_gdn_available(
        query, key, value, g, beta, chunk_size=chunk_size
    ):
        raise RuntimeError(
            "GDN backend='triton' requires triton-ascend >= 3.2.1 "
            "installation with the Ascend backend and NPU inputs "
            "q/k/v/beta=bf16, g=fp32, head_k_dim=head_v_dim=128, "
            "chunk_size=64, and sequence length divisible by 64."

            "q/k/v/beta=bf16, g=fp32, head_k_dim=head_v_dim=128, "
            "chunk_size=64, and sequence length divisible by 64."
        )

    from hyper_parallel.platform.torch.custom_ops.gdn.chunk_gated_delta_rule import (  # pylint: disable=import-outside-toplevel
        chunk_gated_delta_rule as triton_chunk_gated_delta_rule,
    )

    return triton_chunk_gated_delta_rule(
        query,
        key,
        value,
        g,

def _global_peer_rank(cp_mesh: DeviceMesh, local_rank: int) -> int:
    """Map a CP-local rank index to its global distributed rank."""
    return int(cp_mesh.rank_list[local_rank])


def _slice_local_cp(
    tensor: torch.Tensor,

    cp_rank: int,
    cp_size: int,
) -> torch.Tensor:
    """Return this CP rank's contiguous slice along ``dim``."""
    dim_size = tensor.shape[dim]
    if dim_size % cp_size != 0:
        raise ValueError(
            f"linear attention CP expects dim size {dim_size} "
            f"to be divisible by cp_size {cp_size}."
        )
    chunk = dim_size // cp_size
    return tensor.narrow(dim, cp_rank * chunk, chunk)


def _slice_qkv_local_cp(
    tensor: torch.Tensor,

    cp_rank: int,
    cp_size: int,
) -> torch.Tensor:
    """Slice a fused ``[Q, K, V]`` tensor on the Q/K/V channel dimension."""
    q, k, v = torch.split(tensor, [key_dim, key_dim, value_dim], dim=dim)
    return torch.cat(
        (
            _slice_local_cp(q, dim, cp_rank, cp_size),
            _slice_local_cp(k, dim, cp_rank, cp_size),
            _slice_local_cp(v, dim, cp_rank, cp_size),

    activations as local sequence shards. If an upstream wrapper passes that
    shard as a DTensor, use its local tensor and continue with the same
    ``[B, S_local, H]`` boundary contract.
    """
    if isinstance(tensor, DTensor):
        return tensor.to_local()
    return tensor


def _kv_groups(module: nn.Module) -> int:
    """Return the Q/K replication factor used by the official HF GDN layer."""
    return int(module.num_v_heads // module.num_k_heads)


def _apply_local_padding_mask(
    hidden_states: torch.Tensor,

    hidden_states: torch.Tensor,
    attention_mask: Optional[torch.Tensor],
) -> torch.Tensor:
    """Apply the supported local 2-D padding mask without silent fallback."""
    if attention_mask is None:
        return hidden_states
    if attention_mask.ndim != 2:
        raise NotImplementedError(
            "Qwen3.5 GDN Context Parallel supports only a local 2-D padding mask."
        )
    if tuple(attention_mask.shape) != tuple(hidden_states.shape[:2]):
        raise ValueError(
            "Qwen3.5 GDN Context Parallel expects attention_mask shape "
            f"{tuple(hidden_states.shape[:2])}, got {tuple(attention_mask.shape)}."
        )
    return hidden_states * attention_mask[:, :, None].to(hidden_states.dtype)


def _validate_gdn_activation(module: nn.Module) -> None:
    """Reject model variants whose Conv activation is not implemented here."""
    activation = getattr(module, "activation", "silu")
    if activation not in ("silu", "swish"):
        raise NotImplementedError(
            "Qwen3.5 GDN Context Parallel currently supports only SiLU/Swish "
            f"Conv activation, got {activation!r}."
        )

    cp_rank: int,
    cp_size: int,
) -> torch.Tensor:
    """Send a convolution halo only to the next rank using differentiable A2AV."""
    if cp_size == 1:
        return torch.zeros_like(tail)

    cp_group = cp_mesh.get_group()
    group_ranks = tuple(int(rank) for rank in dist.get_process_group_ranks(cp_group))
    rank_list = tuple(int(rank) for rank in cp_mesh.rank_list)
    rank_to_group_index = {rank: index for index, rank in enumerate(group_ranks)}
    halo_width = tail.shape[1]

    input_splits = [0] * cp_size
    exchange_input = tail.permute(1, 0, 2).contiguous()
    if cp_rank < cp_size - 1:
        input_splits[rank_to_group_index[rank_list[cp_rank + 1]]] = halo_width
    else:
        exchange_input = exchange_input[:0]

    output_splits = [0] * cp_size
    if cp_rank > 0:
        output_splits[rank_to_group_index[rank_list[cp_rank - 1]]] = halo_width

    exchange_output = platform.differentiable_all_to_all_single(
        exchange_input,
        input_splits,
        output_splits,
        group=cp_group,

        input_splits,
        output_splits,
        group=cp_group,
    )
    if cp_rank == 0:
        return torch.zeros_like(tail) + exchange_output.sum().to(tail.dtype) * 0
    return exchange_output.permute(1, 0, 2).contiguous()


def _causal_conv1d_with_cp_halo(
    mixed_qkv: torch.Tensor,

    cp_rank: int,
    cp_size: int,
) -> torch.Tensor:
    """Run causal depthwise Conv1d with only the previous rank's boundary."""
    kernel_size = conv1d.kernel_size[0]
    dilation = conv1d.dilation[0]
    halo_width = (kernel_size - 1) * dilation
    if halo_width == 0 or cp_size == 1:
        conv_out = conv1d(mixed_qkv.transpose(1, 2))
        return F.silu(conv_out[:, :, : mixed_qkv.shape[1]]).transpose(1, 2)

    if mixed_qkv.shape[1] < halo_width:
        raise ValueError(
            "linear attention CP conv halo requires local_seq_len >= "
            f"{halo_width}, got {mixed_qkv.shape[1]}."
        )

    halo = _all_to_all_previous_rank_halo(
        mixed_qkv[:, -halo_width:, :].contiguous(),
        cp_mesh,
        cp_rank,
        cp_size,

        cp_mesh,
        cp_rank,
        cp_size,
    )
    conv_input = torch.cat((halo, mixed_qkv), dim=1).transpose(1, 2)
    conv_out = F.conv1d(
        input=conv_input,
        weight=conv1d.weight,
        bias=conv1d.bias,
        stride=conv1d.stride,

        padding=0,
        dilation=conv1d.dilation,
        groups=conv1d.groups,
    )
    return F.silu(conv_out).transpose(1, 2)


def _l2norm_torch(x: torch.Tensor, dim: int = -1, eps: float = 1e-6) -> torch.Tensor:
    """Match the pure torch GDN reference l2norm helper."""
    return x * torch.rsqrt((x * x).sum(dim=dim, keepdim=True) + eps)


class _GDNPreparedChunks(NamedTuple):
    """Reusable chunk intermediates shared by state-summary CP modes."""

    chunk_size: int = 64,
    use_qk_l2norm_in_kernel: bool = False,
) -> _GDNPreparedChunks:
    """Prepare GDN chunk intermediates shared by summary and local output."""
    initial_dtype = query.dtype
    if use_qk_l2norm_in_kernel:
        query = _l2norm_torch(query, dim=-1, eps=1e-6)
        key = _l2norm_torch(key, dim=-1, eps=1e-6)

    query, key, value, beta, g = [
        x.transpose(1, 2).contiguous().to(torch.float32)
        for x in (query, key, value, beta, g)
    ]

    sequence_length = key.shape[2]
    pad_size = (chunk_size - sequence_length % chunk_size) % chunk_size
    query = F.pad(query, (0, 0, 0, pad_size))
    key = F.pad(key, (0, 0, 0, pad_size))
    value = F.pad(value, (0, 0, 0, pad_size))
    beta = F.pad(beta, (0, pad_size))
    g = F.pad(g, (0, pad_size))
    total_sequence_length = sequence_length + pad_size

    query = query * (1 / (query.shape[-1] ** 0.5))
    v_beta = value * beta.unsqueeze(-1)
    k_beta = key * beta.unsqueeze(-1)
    query, key, k_beta, v_beta = [
        x.reshape(x.shape[0], x.shape[1], -1, chunk_size, x.shape[-1])
        for x in (query, key, k_beta, v_beta)
    ]
    g = g.reshape(g.shape[0], g.shape[1], -1, chunk_size)

    mask = torch.triu(
        torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=query.device),
        diagonal=0,
    )
    g = g.cumsum(dim=-1)
    decay_mask = ((g.unsqueeze(-1) - g.unsqueeze(-2)).tril().exp().float()).tril()
    attn = -((k_beta @ key.transpose(-1, -2)) * decay_mask).masked_fill(mask, 0)
    for row_idx in range(1, chunk_size):
        row = attn[..., row_idx, :row_idx].clone()
        sub = attn[..., :row_idx, :row_idx].clone()
        attn[..., row_idx, :row_idx] = row + (row.unsqueeze(-1) * sub).sum(-2)
    attn = attn + torch.eye(chunk_size, dtype=attn.dtype, device=attn.device)

    chunk_value = attn @ v_beta
    k_cumdecay = attn @ (k_beta * g.exp().unsqueeze(-1))
    return _GDNPreparedChunks(
        initial_dtype=initial_dtype,
        query=query,
        key=key,
        chunk_value=chunk_value,

def _compute_gdn_state_summary_from_prepared(
    prepared: _GDNPreparedChunks,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Compute ``state_out = M @ state_in + S`` from prepared GDN chunks."""
    key = prepared.key
    batch_size, num_heads, _, _, k_head_dim = key.shape
    v_head_dim = prepared.chunk_value.shape[-1]
    eye = torch.eye(k_head_dim, device=key.device, dtype=torch.float32).reshape(
        1, 1, k_head_dim, k_head_dim
    )
    state_ext = torch.zeros(
        batch_size,
        num_heads,
        k_head_dim,
        v_head_dim,

        v_head_dim,
        device=key.device,
        dtype=torch.float32,
    )
    transition = eye.expand(batch_size, num_heads, -1, -1).clone()

    for chunk_idx in range(key.shape[2]):
        key_i = key[:, :, chunk_idx]
        value_i = prepared.chunk_value[:, :, chunk_idx]
        w_i = prepared.k_cumdecay[:, :, chunk_idx]
        g_i = prepared.g[:, :, chunk_idx]
        decay = g_i[:, :, -1].exp()
        key_decay = key_i * (g_i[:, :, -1, None] - g_i).exp()[..., None]

        transition_i = (
            decay[:, :, None, None] * eye
            - key_decay.transpose(-1, -2) @ w_i
        )
        state_ext_i = key_decay.transpose(-1, -2) @ value_i
        state_ext = transition_i @ state_ext + state_ext_i
        transition = transition_i @ transition

    return state_ext, transition


def _checkpoint_gdn_state_summary(
    prepared: _GDNPreparedChunks,

def _checkpoint_gdn_state_summary(
    prepared: _GDNPreparedChunks,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Compute a state summary without retaining its per-chunk autograd graph."""
    if not torch.is_grad_enabled():
        return _compute_gdn_state_summary_from_prepared(prepared)

    def recompute(
        key: torch.Tensor,
        chunk_value: torch.Tensor,
        g: torch.Tensor,
        k_cumdecay: torch.Tensor,

        g: torch.Tensor,
        k_cumdecay: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """Rebuild a prepared view from explicit checkpoint inputs."""
        checkpoint_prepared = prepared._replace(
            key=key,
            chunk_value=chunk_value,
            g=g,
            k_cumdecay=k_cumdecay,

            chunk_value=chunk_value,
            g=g,
            k_cumdecay=k_cumdecay,
        )
        return _compute_gdn_state_summary_from_prepared(checkpoint_prepared)

    return checkpoint(
        recompute,
        prepared.key,
        prepared.chunk_value,
        prepared.g,

    prepared: _GDNPreparedChunks,
    initial_state: Optional[torch.Tensor],
) -> torch.Tensor:
    """Run local GDN output using already prepared chunk intermediates."""
    query = prepared.query
    key = prepared.key
    chunk_value = prepared.chunk_value
    batch_size, num_heads, _, _, k_head_dim = key.shape
    v_head_dim = chunk_value.shape[-1]
    recurrent_state = (
        torch.zeros(
            batch_size,
            num_heads,
            k_head_dim,

        )
        if initial_state is None
        else initial_state.to(chunk_value)
    )
    core_attn_out = torch.zeros_like(chunk_value)

    for chunk_idx in range(0, prepared.total_sequence_length // prepared.chunk_size):
        q_i = query[:, :, chunk_idx]
        k_i = key[:, :, chunk_idx]
        v_i = chunk_value[:, :, chunk_idx]
        attn = q_i @ k_i.transpose(-1, -2) * prepared.decay_mask[:, :, chunk_idx]
        v_prime = prepared.k_cumdecay[:, :, chunk_idx] @ recurrent_state
        v_new = v_i - v_prime
        attn_inter = (
            q_i * prepared.g[:, :, chunk_idx, :, None].exp()
        ) @ recurrent_state
        core_attn_out[:, :, chunk_idx] = attn_inter + attn @ v_new
        recurrent_state = (
            recurrent_state * prepared.g[:, :, chunk_idx, -1, None, None].exp()
            + (
                k_i
                * (

                ).exp()[..., None]
            ).transpose(-1, -2) @ v_new
        )

    core_attn_out = core_attn_out.reshape(
        core_attn_out.shape[0],
        core_attn_out.shape[1],
        -1,
        core_attn_out.shape[-1],

        core_attn_out.shape[1],
        -1,
        core_attn_out.shape[-1],
    )
    core_attn_out = core_attn_out[:, :, :prepared.sequence_length]
    return core_attn_out.transpose(1, 2).contiguous().to(prepared.initial_dtype)


class _RecvInitialStateP2PFunction(torch.autograd.Function):
    """Receive the recurrent initial state; send its gradient in backward."""

        prev_rank: int,
        state_shape: tuple[int, ...],
    ) -> torch.Tensor:
        """Receive the initial state from the preceding CP rank."""
        state = torch.empty(state_shape, device=anchor.device, dtype=torch.float32)
        dist.recv(state, src=prev_rank, group=cp_group)
        ctx.cp_group = cp_group
        ctx.prev_rank = prev_rank
        return state

    @staticmethod
    def backward(ctx, grad_state: Optional[torch.Tensor]):
        if grad_state is None:
            raise RuntimeError("linear attention P2P backward missing initial-state grad.")
        dist.send(grad_state.contiguous(), dst=ctx.prev_rank, group=ctx.cp_group)
        return None, None, None, None


class _SendFinalStateP2PFunction(torch.autograd.Function):
    """Send the recurrent final state; receive its gradient in backward."""

        cp_group,
        next_rank: int,
    ) -> torch.Tensor:
        """Send the final state to the succeeding CP rank."""
        dist.send(final_state.contiguous(), dst=next_rank, group=cp_group)
        ctx.cp_group = cp_group
        ctx.next_rank = next_rank
        ctx.state_shape = tuple(final_state.shape)
        ctx.state_dtype = final_state.dtype
        return final_state.new_zeros(())

    @staticmethod
    def backward(ctx, grad_token: torch.Tensor):
        grad_state = torch.empty(
            ctx.state_shape,
            device=grad_token.device,
            dtype=ctx.state_dtype,
        )
        dist.recv(grad_state, src=ctx.next_rank, group=ctx.cp_group)
        return grad_state, None, None


def _apply_gdn_state_summary(
    state_ext: torch.Tensor,

    transition: torch.Tensor,
    initial_state: Optional[torch.Tensor],
) -> torch.Tensor:
    """Apply ``state_out = M @ state_in + S`` to an incoming GDN state."""
    if initial_state is None:
        return state_ext
    return transition @ initial_state.to(transition) + state_ext


def _gdn_state_p2p_summary(
    query: torch.Tensor,

    The rank-ordered critical path then contains only ``M @ state + S`` and
    the small state transfer. Token outputs retain the ordinary PyTorch graph,
    while the two custom autograd boundaries reverse the state communication.
    """
    if cp_size == 1:
        core_attn_out, _ = torch_chunk_gated_delta_rule(
            query,
            key,
            value,
            g=g,

            output_final_state=False,
            chunk_size=chunk_size,
            use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
        )
        return core_attn_out

    cp_group = cp_mesh.get_group()
    prev_rank = _global_peer_rank(cp_mesh, cp_rank - 1) if cp_rank > 0 else -1
    next_rank = _global_peer_rank(cp_mesh, cp_rank + 1) if cp_rank < cp_size - 1 else -1
    state_shape = (query.shape[0], value.shape[2], query.shape[3], value.shape[3])

    prepared = _prepare_gdn_chunks_for_summary(
        query,
        key,
        value,
        g,

        beta,
        chunk_size=chunk_size,
        use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
    )
    state_ext = None
    transition = None
    if cp_rank < cp_size - 1:
        state_ext, transition = _checkpoint_gdn_state_summary(prepared)

    initial_state = None
    if cp_rank > 0:
        initial_state = _RecvInitialStateP2PFunction.apply(
            query,
            cp_group,
            prev_rank,
            state_shape,

            prev_rank,
            state_shape,
        )

    send_token = None
    if cp_rank < cp_size - 1:
        final_state = _apply_gdn_state_summary(
            state_ext,
            transition,
            initial_state,
        )
        send_token = _SendFinalStateP2PFunction.apply(final_state, cp_group, next_rank)

    core_attn_out = _run_prepared_gdn_chunks(prepared, initial_state)
    if send_token is not None:
        core_attn_out = core_attn_out + send_token.to(core_attn_out.dtype) * 0

    return core_attn_out


class _GDNStateP2PTritonFunction(torch.autograd.Function):
    """Pipeline fused affine GDN states over sequence-sharded CP ranks."""

        prev_rank: int,
        next_rank: int,
    ) -> torch.Tensor:
        """Run fused local GDN and forward its affine state across CP ranks."""
        from hyper_parallel.platform.torch.custom_ops.gdn.chunk_gated_delta_rule import (  # pylint: disable=import-outside-toplevel
            chunk_gated_delta_rule_fwd_apply_state_saved,
            chunk_gated_delta_rule_fwd_output_saved,
            chunk_gated_delta_rule_fwd_prepare_saved,
        )
        from hyper_parallel.platform.torch.custom_ops.gdn.state_summary import (  # pylint: disable=import-outside-toplevel
            apply_gdn_state_summary,
            chunk_gated_delta_rule_state_summary_fwd,
        )

        (
            query_norm,
            key_norm,
            _,
            _,

            g,
            beta,
            use_qk_l2norm_in_kernel=False,
        )
        initial_state = None
        recv_buffer = None
        recv_work = None
        if cp_rank > 0:
            recv_buffer = torch.empty(
                (query.shape[0], query.shape[2], query.shape[3], value.shape[3]),
                device=query.device,
                dtype=torch.float32,
            )
            recv_work = dist.irecv(recv_buffer, src=prev_rank, group=cp_group)

        state_ext = None
        transition = None
        if cp_rank < cp_size - 1:
            state_ext, transition = chunk_gated_delta_rule_state_summary_fwd(
                key_norm,
                w,
                u,
                g_cumsum,

                u,
                g_cumsum,
            )

        if recv_work is not None:
            recv_work.wait()
            initial_state = recv_buffer

        send_buffer = None
        send_work = None
        if cp_rank < cp_size - 1:
            send_buffer = apply_gdn_state_summary(
                state_ext,
                transition,
                initial_state,
            ).contiguous()
            send_work = dist.isend(send_buffer, dst=next_rank, group=cp_group)

        h, v_new, _ = chunk_gated_delta_rule_fwd_apply_state_saved(
            key_norm,
            g_cumsum,
            w,
            u,

            u,
            initial_state=initial_state,
            output_final_state=False,
        )
        output = chunk_gated_delta_rule_fwd_output_saved(
            query_norm,
            key_norm,
            g_cumsum,
            h,

            v_new,
            scale,
        ).to(query.dtype)

        if send_work is not None:
            send_work.wait()

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

            matrix_a,
            initial_state if initial_state is not None else empty,
            transition if transition is not None else empty,
        )
        ctx.has_initial_state = initial_state is not None
        ctx.cp_rank = cp_rank
        ctx.cp_size = cp_size
        ctx.cp_group = cp_group
        ctx.prev_rank = prev_rank
        ctx.next_rank = next_rank
        ctx.scale = scale
        return output

    @staticmethod
    def backward(ctx, grad_output: torch.Tensor):  # pylint: disable=too-many-locals
        """Backpropagate local GDN tensors and the state gradient wavefront."""
        from hyper_parallel.platform.torch.custom_ops.gdn.chunk_gated_delta_rule import (  # pylint: disable=import-outside-toplevel
            chunk_gated_delta_rule_bwd_finish_saved,
            chunk_gated_delta_rule_bwd_prepare_saved,
            chunk_gated_delta_rule_bwd_state_saved,
        )
        from hyper_parallel.platform.torch.custom_ops.gdn.state_summary import (  # pylint: disable=import-outside-toplevel
            apply_gdn_state_gradient_summary,
            chunk_gated_delta_rule_state_gradient_summary_bwd,
        )

        (
            query,
            key,
            value,
            g_cumsum,

            matrix_a,
            initial_state,
            transition,
        ) = ctx.saved_tensors
        if not ctx.has_initial_state:
            initial_state = None

        w, h, v_new, dv = chunk_gated_delta_rule_bwd_prepare_saved(
            query,
            key,
            value,
            g_cumsum,

            grad_output,
            ctx.scale,
        )

        grad_state_ext = None
        if ctx.cp_rank > 0:
            grad_state_ext = chunk_gated_delta_rule_state_gradient_summary_bwd(
                query,
                key,
                w,
                g_cumsum,

                dv,
                ctx.scale,
            )

        grad_final_state = None
        recv_work = None
        if ctx.cp_rank < ctx.cp_size - 1:
            recv_buffer = torch.empty(
                (query.shape[0], query.shape[2], query.shape[3], value.shape[3]),
                device=grad_output.device,
                dtype=torch.float32,
            )
            recv_work = dist.irecv(
                recv_buffer,
                src=ctx.next_rank,
                group=ctx.cp_group,
            )
        if recv_work is not None:
            recv_work.wait()
            grad_final_state = recv_buffer

        send_buffer = None
        send_work = None
        if ctx.cp_rank > 0:
            send_buffer = apply_gdn_state_gradient_summary(
                grad_state_ext,
                transition,
                grad_final_state,
            ).contiguous()
            send_work = dist.isend(
                send_buffer,
                dst=ctx.prev_rank,
                group=ctx.cp_group,
            )

                dst=ctx.prev_rank,
                group=ctx.cp_group,
            )

        dh, _, dv = chunk_gated_delta_rule_bwd_state_saved(
            query,
            key,
            g_cumsum,
            w,

            grad_output,
            dv,
            ctx.scale,
        )
        empty = query.new_empty(0)
        dq, dk, dv, dg, dbeta = chunk_gated_delta_rule_bwd_finish_saved(
            query,
            key,
            query,
            key,

            ctx.scale,
            use_qk_l2norm_in_kernel=False,
        )

        if send_work is not None:
            send_work.wait()
        return dq, dk, dv, dg, dbeta, None, None, None, None, None


def _gdn_state_p2p_triton(
    query: torch.Tensor,

    cp_rank: int,
    cp_size: int,
) -> torch.Tensor:
    """Run fused local GDN with an affine state wavefront."""
    if cp_size == 1:
        output, _ = chunk_gated_delta_rule(
            query,
            key,
            value,
            g=g,

            output_final_state=False,
            use_qk_l2norm_in_kernel=True,
            backend="triton",
        )
        return output

    query = _l2norm_torch(query)
    key = _l2norm_torch(key)
    prev_rank = _global_peer_rank(cp_mesh, cp_rank - 1) if cp_rank > 0 else -1
    next_rank = (
        _global_peer_rank(cp_mesh, cp_rank + 1) if cp_rank < cp_size - 1 else -1
    )
    return _GDNStateP2PTritonFunction.apply(
        query,
        key,
        value,
        g,

    Shard(split_dim)`` redistribution for a 1-D mesh. It uses platform-level
    differentiable all-to-all directly to avoid wrapping each activation in a
    temporary DTensor.
    """
    split_count = device_mesh.size()
    if split_count == 1:
        return tensor

    original_shape = tuple(tensor.shape)
    dim_size = original_shape[split_dim]
    if dim_size % split_count != 0:
        raise ValueError(
            f"linear attention all-to-all split dim {split_dim} with size "
            f"{dim_size} must be divisible by cp_size {split_count}."
        )

    split_size = dim_size // split_count
    final_shape = list(original_shape)
    if split_dim != concat_dim:
        final_shape[split_dim] = split_size
        final_shape[concat_dim] = final_shape[concat_dim] * split_count
    final_shape = tuple(final_shape)

    reshape_dims = list(original_shape)
    reshape_dims[split_dim] = split_count
    reshape_dims.insert(split_dim + 1, split_size)

    trans_dims = list(range(len(reshape_dims)))
    trans_dims.remove(split_dim)
    trans_dims.insert(0, split_dim)

    a2a_input = tensor.reshape(reshape_dims).permute(trans_dims).contiguous()
    reshape_shape = list(a2a_input.shape)
    reshape_shape[0] = reshape_shape[0] * reshape_shape[1]
    reshape_shape.pop(1)
    a2a_input = a2a_input.reshape(reshape_shape)

    a2a_input = a2a_input.contiguous()
    split_len = a2a_input.shape[0] // split_count
    input_splits = [split_len] * split_count
    output_splits = [split_len] * split_count
    output = platform.differentiable_all_to_all_single(
        a2a_input,
        input_splits,
        output_splits,
        group=device_mesh.get_group(),

        output_splits,
        group=device_mesh.get_group(),
    )

    output_reshape = list(output.shape)
    output_reshape[0] = split_count
    output_reshape.insert(1, output.shape[0] // split_count)

    out_trans_dims = list(range(len(output_reshape)))
    first_dim = out_trans_dims.pop(0)
    if concat_dim >= len(out_trans_dims):
        out_trans_dims.append(first_dim)
    else:
        out_trans_dims.insert(concat_dim, first_dim)

    final_output = output.reshape(output_reshape).permute(out_trans_dims).contiguous()
    final_reshape = list(final_output.shape)
    if concat_dim < len(final_reshape) - 1:
        final_reshape[concat_dim] = (
            final_reshape[concat_dim] * final_reshape[concat_dim + 1]
        )
        final_reshape.pop(concat_dim + 1)

    return final_output.reshape(final_reshape).view(final_shape)


class GatedDeltaNetUlyssesCP(nn.Module):
    """Pure-Ulysses CP execution wrapper for a Qwen3.5 Gated DeltaNet module.

        *,
        backend: str = "eager",
        chunk_size: int = 64,
    ):
        super().__init__()
        self.module = module
        backend = backend.lower()
        if backend not in _GDN_BACKENDS:
            raise ValueError(
                f"unsupported GDN backend {backend!r}; "
                f"expected one of {sorted(_GDN_BACKENDS)}."
            )
        if chunk_size <= 0:
            raise ValueError(f"chunk_size must be positive, got {chunk_size}.")
        if backend == "triton" and chunk_size != 64:
            raise NotImplementedError("GDN Triton backend requires chunk_size=64.")
        self.gdn_backend = backend
        self.chunk_size = chunk_size
        self.cp_mesh = _ensure_1d(device_mesh)
        self.cp_size = self.cp_mesh.size()
        self.cp_rank = self.cp_mesh.get_local_rank()
        self.seq_dim = 1
        self.head_dim = 2
        self._validate_module()

    def _validate_module(self) -> None:
        """Validate the wrapped GDN module and Ulysses layout constraints."""
        required = (
            "in_proj_qkv", "in_proj_z", "in_proj_b", "in_proj_a",
            "conv1d", "A_log", "dt_bias", "norm", "out_proj",
            "num_k_heads", "num_v_heads", "head_k_dim", "head_v_dim",
            "key_dim", "value_dim", "conv_dim",

            "conv1d", "A_log", "dt_bias", "norm", "out_proj",
            "num_k_heads", "num_v_heads", "head_k_dim", "head_v_dim",
            "key_dim", "value_dim", "conv_dim",
        )
        missing = [name for name in required if not hasattr(self.module, name)]
        if missing:
            raise TypeError(
                "Qwen3.5 GDN layer is missing required attributes: "
                + ", ".join(missing)
            )
        if self.cp_size <= 1:
            return
        if self.module.num_k_heads % self.cp_size != 0:
            raise ValueError(
                f"linear attention num_k_heads ({self.module.num_k_heads}) must be "
                f"divisible by cp_size ({self.cp_size}) for Ulysses CP."
            )
        if self.module.num_v_heads % self.cp_size != 0:
            raise ValueError(
                f"linear attention num_v_heads ({self.module.num_v_heads}) must be "
                f"divisible by cp_size ({self.cp_size}) for Ulysses CP."
            )
        if self.module.num_v_heads % self.module.num_k_heads:
            raise ValueError("GDN num_v_heads must be divisible by num_k_heads.")
        _validate_gdn_activation(self.module)

    def _seq_to_head(self, tensor: torch.Tensor) -> torch.Tensor:
        return _differentiable_all_to_all_shard(
            tensor,
            self.cp_mesh,
            split_dim=self.head_dim,
            concat_dim=self.seq_dim,