Diff Coverage

        Returns:
            A view of the restored output with the paired gradient edge.
        """
        del attention
        ctx.state = state
        return restored.view_as(restored)

    @staticmethod
    @once_differentiable
    def backward(ctx: Any, gradient: torch.Tensor) -> tuple[torch.Tensor, None, None]:

        Returns:
            Head-sharded attention gradient and empty metadata gradients.
        """
        del gradient
        state = ctx.state
        if state.backward_exchange is None:
            raise RuntimeError("MLA output projection did not submit its gradient exchange")
        result = state.backward_exchange.wait()
        state.backward_exchange = None
        return result, None, None


class _OutputLinear(torch.autograd.Function):
    """Launch dO before wgrad while preserving the ordinary Linear and gate math."""

        Returns:
            Gated and projected local hidden states.
        """
        projected = activation if gate is None else activation * gate
        ctx.save_for_backward(activation, projected, weight, gate)
        ctx.state, ctx.has_bias = state, bias is not None
        ctx.device_type = activation.device.type
        ctx.autocast_enabled = torch.is_autocast_enabled(ctx.device_type)
        ctx.autocast_dtype = torch.get_autocast_dtype(ctx.device_type) if ctx.autocast_enabled else None
        return F.linear(projected, weight, bias)  # pylint: disable=not-callable

    @staticmethod
    @once_differentiable
    def backward(ctx: Any, gradient: torch.Tensor) -> tuple[torch.Tensor | None, ...]:

        Returns:
            Gradients for activation, weight, bias and gate; no state gradient.
        """
        with torch.autocast(ctx.device_type, enabled=ctx.autocast_enabled, dtype=ctx.autocast_dtype):
            return _output_linear_backward(ctx, gradient)


def _output_linear_backward(ctx, gradient):
    """Match forward autocast and skip communication when attention is frozen."""
    activation, projected, weight, gate = ctx.saved_tensors
    grad_projected = gradient.matmul(weight).to(projected.dtype)
    grad_activation = None
    # A frozen attention source has no RestoreOutput consumer to wait for dO.
    if ctx.needs_input_grad[0]:
        grad_activation = grad_projected if gate is None else grad_projected * gate
        state = ctx.state
        exchange_input = grad_activation.reshape(*activation.shape[:2], -1, state.head_dim)
        state.backward_exchange = async_ulysses_seq_to_head_launch(exchange_input, 1, 2, state.mesh)
    needs_gate_grad = ctx.needs_input_grad[3] or ctx.state.input_gate is not None
    grad_gate = grad_projected * activation if gate is not None and needs_gate_grad else None
    if ctx.state.input_gate is not None:
        ctx.state.input_gate.gradient = grad_gate
        grad_gate = None
    flat_input = projected.reshape(-1, projected.shape[-1])
    flat_gradient = gradient.reshape(-1, gradient.shape[-1])
    grad_weight = None
    if ctx.needs_input_grad[1]:
        grad_weight = flat_input.transpose(0, 1).matmul(flat_gradient).transpose(0, 1)
    grad_bias = flat_gradient.sum(0) if ctx.has_bias and ctx.needs_input_grad[2] else None
    return grad_activation, grad_weight, grad_bias, grad_gate, None


def restore_mla_output_gradient(attention: torch.Tensor, restored: torch.Tensor,
                                state: MLAOutputExchange) -> torch.Tensor:

        The paired output projection must be the only consumer. Its saved dO
        supplies this edge's complete first-order gradient; partial autograd.grad
        calls or extra losses attached to the restored value are unsupported.
    """
    return _RestoreOutput.apply(attention, restored, state)


def mla_output_projection_request(projection: nn.Linear) -> _ForwardRewriteRequest:
    """Keep parameter materialization hooks while allowing dgrad/wgrad scheduling.

    Returns:
        An atomic child forward rewrite; ordinary calls remain unchanged.
    """
    # Quantized and LoRA subclasses need their own derivative contracts.
    if type(projection) is not nn.Linear:  # pylint: disable=unidiomatic-typecheck
        raise ValueError("MLA backward overlap requires a plain output Linear")
    if getattr(projection.forward, "__func__", None) is not nn.Linear.forward:
        raise ValueError("MLA backward overlap requires an unwrapped plain output Linear")
    original = projection.forward

    @wraps(original)
    def output_projection(*args: torch.Tensor, mla_output_exchange: MLAOutputExchange | None = None,
                          mla_output_gate: torch.Tensor | None = None, **kwargs: Any) -> torch.Tensor:
        """Read managed parameters after their normal module pre-hooks.

        Args:

        Returns:
            Projected hidden states using ordinary or scheduled autograd.
        """
        if mla_output_exchange is None:
            return original(*args, **kwargs)
        # The private scheduled call passes one activation; ordinary calls retain input= support.
        if len(args) != 1 or kwargs:
            raise TypeError("Scheduled MLA output projection requires one positional activation")
        return _OutputLinear.apply(args[0], projection.weight, projection.bias, mla_output_gate, mla_output_exchange)

    return _ForwardRewriteRequest(projection, output_projection)


@dataclass
class MLAInputExchange:

        Returns:
            A view of compute_input retaining both gradient dependencies.
        """
        ctx.state = state
        state.needs_send_grad = ctx.needs_input_grad[0]
        state.forward_exchange = async_ulysses_seq_to_head_launch(send_tensor, 1, 2, state.mesh)
        return compute_input.view_as(compute_input)

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

        Returns:
            Exchanged source gradient, compute gradient and no state gradient.
        """
        state = ctx.state
        send_gradient = None
        if state.needs_send_grad:
            if state.backward_exchange is None:
                raise RuntimeError("MLA input consumer did not submit its inverse exchange")
            send_gradient = state.backward_exchange.wait()
            state.backward_exchange = None
        return send_gradient, compute_gradient, None


class _LatentInputWait(torch.autograd.Function):
    """Launch inverse Q before releasing the local KV projection gradients."""

        Returns:
            Attention inputs in query, key, value order.
        """
        ctx.state = state
        query = state.forward_exchange.wait()
        state.forward_exchange = None
        return query, key.view_as(key), value.view_as(value)

    @staticmethod
    @once_differentiable
    def backward(ctx: Any, grad_query: torch.Tensor, grad_key: torch.Tensor,

        Returns:
            Key and value gradients, followed by no state gradient.
        """
        state = ctx.state
        if state.needs_send_grad:
            state.backward_exchange = async_ulysses_seq_to_head_launch(grad_query, 2, 1, state.mesh)
        return grad_key, grad_value, None


@dataclass
class MLAInputGateExchange:

        Returns:
            The model's post-activation gate in local-token layout.
        """
        if self.projection is None:
            raise RuntimeError("MLA input gate must be projected exactly once")
        self.value = self.projection(source)
        self.hidden, self.projection = None, None
        return self.value

    def take_gradient(self) -> torch.Tensor:
        """Consume the output multiplication's deferred gate gradient once."""
        if self.gradient is None:
            raise RuntimeError("MLA output projection did not release its gate gradient")
        gradient, self.gradient = self.gradient, None
        return gradient


@dataclass
class MLAExpandedInputExchange:

    gate: MLAInputGateExchange | None = None

    def wait(self) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """Consume Q/K/V in attention order after all forward work is submitted."""
        results = []
        for exchange in (self.query, self.key, self.value):
            results.append(exchange.forward_exchange.wait())
            exchange.forward_exchange = None
        return results[0], results[1], results[2]


class _ExpandedInputWait(torch.autograd.Function):
    """Complete V/K/Q after optional gate work, reversing the same dependency."""

        Returns:
            Completed attention inputs in query, key, value order.
        """
        ctx.state = state
        if state.gate is None:
            state.query.needs_send_grad = ctx.needs_input_grad[0]
            state.query.forward_exchange = async_ulysses_seq_to_head_launch(
                final_input, 1, 2, state.query.mesh,
            )
        return state.wait()

    @staticmethod
    @once_differentiable
    def backward(ctx: Any, grad_query: torch.Tensor, grad_key: torch.Tensor,

        Returns:
            Deferred gate or completed Q gradient, and no state gradient.
        """
        state = ctx.state
        # Reverse V/K/Q: optional gate backward hides dQ; Q backward hides dK/dV.
        for exchange, gradient in ((state.query, grad_query), (state.key, grad_key), (state.value, grad_value)):
            if exchange.needs_send_grad:
                exchange.backward_exchange = async_ulysses_seq_to_head_launch(gradient, 2, 1, exchange.mesh)
        if state.gate is not None:
            return state.gate.take_gradient(), None
        query_gradient = None
        if state.query.needs_send_grad:
            query_gradient = state.query.backward_exchange.wait()
            state.query.backward_exchange = None
        return query_gradient, None


def _launch_expanded_input(send_tensor, compute_input, exchange, defer_backward):
    """Use the same forward schedule with native or explicitly paired adjoints."""
    if defer_backward:
        return _InputLaunch.apply(send_tensor, compute_input, exchange)
    exchange.forward_exchange = async_ulysses_seq_to_head_launch(send_tensor, 1, 2, exchange.mesh)
    return compute_input


class _LatentKVLinear(torch.autograd.Function):
    """Schedule the gathered latent adjoint before selected-head weight gradients."""

        Returns:
            Expanded KV heads and the shared key, without duplicating its gradient.
        """
        del payload
        latent = gathered[..., :weight.shape[-1]].contiguous()
        shared_key = gathered[..., weight.shape[-1]:]
        ctx.save_for_backward(latent, weight)
        ctx.mesh, ctx.has_bias = mesh, bias is not None
        ctx.device_type = latent.device.type
        ctx.autocast_enabled = torch.is_autocast_enabled(ctx.device_type)
        ctx.autocast_dtype = torch.get_autocast_dtype(ctx.device_type) if ctx.autocast_enabled else None
        return F.linear(latent, weight, bias), shared_key.view_as(shared_key)  # pylint: disable=not-callable

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

        Returns:
            Local payload, weight and optional bias gradients; no gather or mesh gradient.
        """
        with torch.autocast(ctx.device_type, enabled=ctx.autocast_enabled, dtype=ctx.autocast_dtype):
            return _latent_projection_backward(ctx, gradient, shared_gradient)


def _latent_projection_backward(ctx, gradient, shared_gradient):
    """Keep both reduction buffers alive until their local gradient is consumed."""
    activation, weight = ctx.saved_tensors
    work = None
    if ctx.needs_input_grad[0]:
        grad_latent = gradient.matmul(weight).to(activation.dtype)
        grad_payload = torch.cat((grad_latent, shared_gradient), dim=-1)
        send = grad_payload.movedim(1, 0).contiguous()
        shape = list(send.shape)
        shape[0] //= ctx.mesh.size()
        receive = send.new_empty(shape)
        work = dist.reduce_scatter_tensor(
            receive, send, op=dist.ReduceOp.SUM, group=ctx.mesh.get_group(), async_op=True,
        )
    flat_gradient = gradient.reshape(-1, gradient.shape[-1])
    grad_weight = None
    if ctx.needs_input_grad[1]:
        flat_activation = activation.reshape(-1, activation.shape[-1])
        grad_weight = flat_gradient.transpose(0, 1).matmul(flat_activation)
    grad_bias = flat_gradient.sum(0) if ctx.has_bias and ctx.needs_input_grad[2] else None
    local = None
    if work is not None:
        work.wait()
        local = receive.movedim(0, 1).contiguous()
    return local, grad_weight, grad_bias, None, None

) -> torch.Tensor:
    """Expand and rotate the query through the managed TP-local projection module."""
    projected = module.q_b_proj(latent)
    projected = projected.view(*latent.shape[:2], heads, module.qk_head_dim)
    if positions is None:
        return projected.transpose(1, 2)
    content, rotary = projected.split((module.qk_nope_head_dim, module.qk_rope_head_dim), dim=-1)
    rotary = _apply_mla_rope(rotary, positions, interleaved=module.rotary_interleaved)
    return torch.cat((content, rotary), dim=-1).transpose(1, 2)

) -> tuple[torch.Tensor, torch.Tensor]:
    """Expand the requested KV heads through the managed projection module."""
    projected = (module.kv_b_proj(latent) if head_range is None
                 else module.kv_b_proj(latent, mla_head_range=head_range))
    return _assemble_key_value(module, projected, key_rope, heads)


def _assemble_key_value(module, projected, key_rope, heads):
    """Build per-head keys and values from shared projection results."""
    projected = projected.view(*projected.shape[:2], heads, module.qk_nope_head_dim + module.v_head_dim)
    content, value = projected.split((module.qk_nope_head_dim, module.v_head_dim), dim=-1)
    key = torch.cat((content, key_rope.unsqueeze(2).expand(-1, -1, heads, -1)), dim=-1)
    return key.transpose(1, 2), value.transpose(1, 2)

        _validate_state(module, self.plan, state)
        rotated_key = _apply_mla_rope(
            state.key_rope.unsqueeze(2), positions, interleaved=module.rotary_interleaved,
        ).squeeze(2)
        query, key, value = self._prepare_attention_inputs(
            module, state, positions, rotated_key, partition, schedule_backward, input_gate,
        )
        output = npu_fusion_attention_forward(
            module, query, key, value, attention_mask, scaling=module.scaling,

            module, query, key, value, attention_mask, scaling=module.scaling,
            is_causal=causal, actual_seq_qlen=lengths, actual_seq_kvlen=key_lengths,
            pre_tokens=2147483647, next_tokens=0 if causal else 2147483647,
        )[0]
        if restore_output:
            output = self.restore_output(output)
        return output, None

    def _prepare_attention_inputs(self, module, state, positions, key_rope, partition,
                                  schedule_backward, input_gate):
        """Select input scheduling without changing hybrid output-overlap support."""
        use_overlap = self.plan.overlap and self.plan.cp_degree > 1
        pure_ulysses = self.plan.all_gather_degree == 1
        expanded = self.plan.strategy == "expanded_ulysses"
        if input_gate is not None and not (use_overlap and pure_ulysses and expanded):
            raise ValueError("Input gate overlap requires expanded pure Ulysses with overlap enabled")
        if use_overlap:
            if pure_ulysses and expanded:
                return self._expanded_ulysses_overlap(
                    module, state, positions, key_rope, input_gate, schedule_backward,
                )
            if schedule_backward and pure_ulysses:
                return self._latent_backward_overlap(module, state, positions, key_rope)
            prepare = self._expanded_overlap if expanded else self._latent_overlap
            return prepare(module, state, positions, key_rope, partition)
        query = _query(module, state.query_latent, positions, self.plan.local_heads)
        prepare = self._expanded if expanded or self.plan.cp_degree == 1 else self._latent
        return prepare(module, state, query, key_rope, partition)

    def restore_output(self, output: torch.Tensor) -> torch.Tensor:
        """Restore token ownership after attention, before the model output projection.

        Returns:
            Original local-token, all-local-head BSND tensor.
        """
        batch, sequence = output.shape[:2]
        sequence //= self.plan.ulysses_degree
        if self.plan.ulysses_degree > 1:
            output = ulysses_head_to_seq(output, 1, 2, self.ulysses_mesh)
        expected = (batch, sequence, self.plan.local_heads, self.plan.dimensions.value_dim)
        if tuple(output.shape) != expected:
            raise RuntimeError(f"MLA CP output has shape {tuple(output.shape)}, expected {expected}")
        return output

    def _expanded(self, module, state, query, key_rope, partition=None):
        """Pack unequal-width QKV into one sequence-to-head exchange."""
        key, value = _key_value(module, state.kv_latent, key_rope, self.plan.local_heads, None)

    def _expanded_ulysses_overlap(self, module, state, positions, key_rope, input_gate, schedule_backward):
        """Send V/K/Q in order; an optional model gate extends the final wait window."""
        # Frozen gate inputs retain native autograd ownership of all exchanges.
        paired = schedule_backward and (input_gate is None or input_gate.defer_backward)
        query_exchange = MLAInputExchange(self.ulysses_mesh)
        key_exchange = MLAInputExchange(self.ulysses_mesh)
        value_exchange = MLAInputExchange(self.ulysses_mesh)
        exchanges = MLAExpandedInputExchange(query_exchange, key_exchange, value_exchange, input_gate)
        projected = module.kv_b_proj(state.kv_latent)
        projected = projected.view(*state.kv_latent.shape[:2], self.plan.local_heads,
                                   module.qk_nope_head_dim + module.v_head_dim)
        content, value = projected.split((module.qk_nope_head_dim, module.v_head_dim), dim=-1)
        content = _launch_expanded_input(value, content, value_exchange, paired)
        shared_key = key_rope.unsqueeze(2).expand(-1, -1, self.plan.local_heads, -1)
        key = torch.cat((content, shared_key), dim=-1)
        query_latent = _launch_expanded_input(key, state.query_latent, key_exchange, paired)
        query = _query(module, query_latent, positions, self.plan.local_heads).transpose(1, 2)
        final_input = query
        if input_gate is not None:
            gate_input = _launch_expanded_input(query, input_gate.hidden, query_exchange, paired)
            final_input = input_gate.project(gate_input)
        elif not paired:
            query_exchange.forward_exchange = async_ulysses_seq_to_head_launch(query, 1, 2, self.ulysses_mesh)
        query, key, value = (_ExpandedInputWait.apply(final_input, exchanges) if paired else exchanges.wait())
        return query.transpose(1, 2), key.transpose(1, 2), value.transpose(1, 2)

    def _latent_backward_overlap(self, module, state, positions, key_rope):
        """Cover inverse Q with KV backward, then latent reduction with KV wgrad."""
        payload = torch.cat((state.kv_latent, key_rope), dim=-1)
        pending = async_cp_allgather_launch(payload, 1, self.cp_mesh)
        query = _query(module, state.query_latent, positions, self.plan.local_heads).transpose(1, 2)
        # The managed projection below owns the gather's single gradient edge.
        with torch.no_grad():
            gathered = pending.wait()
        exchange = MLAInputExchange(self.ulysses_mesh)
        bridge = _InputLaunch.apply(query, payload, exchange)
        heads = self.plan.compute_heads
        start = self.head_rank * heads
        projected, shared_key = module.kv_b_proj(
            bridge, mla_head_range=(start, start + heads),
            mla_gathered=gathered, mla_cp_mesh=self.cp_mesh,
        )
        key, value = _assemble_key_value(module, projected, shared_key, heads)
        query, key, value = _LatentInputWait.apply(key, value, exchange)
        return query.transpose(1, 2), key, value


def _head_projection_request(projection: nn.Linear, head_dim: int) -> _ForwardRewriteRequest:
    """Select head rows inside the original Linear call to preserve parameter hooks."""

        Returns:
            Projected activations, or their pair with the shared key for scheduled gather backward.
        """
        if mla_head_range is None:
            return original(input)
        begin, end = mla_head_range
        if not 0 <= begin < end <= projection.weight.shape[0] // head_dim:
            raise ValueError("MLA projection head range is outside the original parameter")
        rows = slice(begin * head_dim, end * head_dim)
        bias = None if projection.bias is None else projection.bias[rows]
        # Read parameters inside the managed module's forward, after pre-hooks.
        # Slicing preserves the Parameter identity and scatters gradients back.

        bias = None if projection.bias is None else projection.bias[rows]
        # Read parameters inside the managed module's forward, after pre-hooks.
        # Slicing preserves the Parameter identity and scatters gradients back.
        if mla_gathered is not None:
            if mla_cp_mesh is None:
                raise ValueError("Scheduled MLA latent projection requires its CP mesh")
            return _LatentKVLinear.apply(input, projection.weight[rows], bias, mla_gathered, mla_cp_mesh)
        return F.linear(input, projection.weight[rows], bias)  # pylint: disable=not-callable

    return _ForwardRewriteRequest(projection, head_projection)

    )
    for projection, in_features, head_dim in projections:
        # Quantized/LoRA subclasses do not obey contiguous Linear head slicing.
        if type(projection) is not nn.Linear:  # pylint: disable=unidiomatic-typecheck
            raise ValueError("MLA CP requires plain nn.Linear up projections")
        weight = projection.weight.to_local() if isinstance(projection.weight, DTensor) else projection.weight
        if weight.shape != (plan.local_heads * head_dim, in_features):
            raise ValueError("MLA up-projection shape must match TP-local heads; apply TP parameter sharding first")
        if getattr(projection.forward, "__func__", None) is not nn.Linear.forward:
            raise ValueError("MLA CP must be installed before other up-projection forward adapters")


def prepare_mla_cp(
    target_module: nn.Module,

        Runtime and child rewrites; the model adapter owns its forward contract.
    """
    for name, axis in (("TP", tp_mesh), ("CP", cp_mesh)):
        if axis is not None and (axis.size() < 1 or getattr(axis, "ndim", 1) != 1):
            raise ValueError(f"mla_cp_wrapper requires a one-dimensional {name} mesh")
    cp_degree = 1 if cp_mesh is None else cp_mesh.size()
    tp_degree = 1 if tp_mesh is None else tp_mesh.size()
    if getattr(target_module, "mla_cp_runtime", None) is not None:
        raise ValueError("MLA CP is already configured on this module")