Diff Coverage

def _restore_sparse_flash_mla_saved_tensors(ctx):
    """Restore optional tensors from the compact autograd saved-tensor sequence."""
    saved_tensors = iter(ctx.saved_tensors)
    presence = (
        True,
        ctx.has_ori_kv,
        ctx.has_cmp_kv,
        ctx.has_sinks,

        True,
        True,
        ctx.has_cmp_residual,
    )
    values = (next(saved_tensors) if is_present else None for is_present in presence)
    return _SparseFlashMlaSavedTensors(*values)


class NpuDenseLightningIndexerSoftmaxLseDFunction(DFunction):  # pylint: disable=W0221
    """DFunction wrapper for npu_dense_lightning_indexer_softmax_lse on MindSpore.

        ctx.has_cmp_residual = cmp_residual_kv is not None
        # metadata is NOT saved for backward: the grad kernel asserts metadata
        # must be nullptr and re-derives its own tiling internally.  cmp_residual_kv
        # IS saved — the grad kernel requires it for CFA/SCFA with cmp_mask_mode=3.
        saved_tensors = [
            query, ori_kv, cmp_kv, sinks, ori_sparse_indices, cmp_sparse_indices,
            cu_seq_lens_q, cu_seq_lens_ori_kv, cu_seq_lens_cmp_kv,
            attention_out, softmax_lse, cmp_residual_kv,
        ]
        ctx.save_for_backward(*[tensor for tensor in saved_tensors if tensor is not None])
        ctx.softmax_scale = softmax_scale
        ctx.cmp_ratio = cmp_ratio
        ctx.ori_mask_mode = ori_mask_mode
        ctx.cmp_mask_mode = cmp_mask_mode

    @staticmethod
    def backward(ctx, grad_attention_out, grad_softmax_lse):  # pylint: disable=unused-argument
        """Backward pass: calls npu_sparse_flash_mla_grad kernel."""
        state = _restore_sparse_flash_mla_saved_tensors(ctx)
        # metadata MUST be None: the grad kernel asserts it is nullptr and
        # re-derives tiling internally.  cmp_residual_kv is passed through —
        # required for CFA/SCFA (cmp_ratio!=1) with cmp_mask_mode=3.
        grads = _custom_ops.npu_sparse_flash_mla_grad(

            ctx.softmax_scale, ctx.cmp_ratio, ctx.ori_mask_mode, ctx.cmp_mask_mode,
            ctx.ori_win_left, ctx.ori_win_right, ctx.layout_q, ctx.layout_kv,
        )
        d_query = grads[0]
        d_ori_kv = grads[1] if state.ori_kv is not None else None
        d_cmp_kv = grads[2] if state.cmp_kv is not None else None
        d_sinks = grads[3] if state.sinks is not None else None
        # grads[4], grads[5] = ori/cmp_softmax_l1_norm — discarded here.
        # 21 positional forward args (ctx excluded):
        # query, ori_kv, cmp_kv, cu_seq_lens_q, cu_seq_lens_ori_kv, cu_seq_lens_cmp_kv,
        # ori_sparse_indices, cmp_sparse_indices, sinks,