Diff Coverage

"""Masked hidden-state aggregation function."""

from typing import Any, Tuple

import torch
import omni_training_custom_ops  # noqa: F401  # pylint: disable=unused-import


class _AggregateHidden(torch.autograd.Function):

    key_rope_tnd = rearrange(key_rope, "b s n d -> (b s) n d")
    # Load only the selected optional backend so either can be used independently.
    # pylint: disable=import-outside-toplevel
    if backend == "cann":
        import torch_npu

        output, softmax_max, softmax_sum = _cann_sparse_attention(
            query_tnd,
            key_tnd,

"""Parameter-sink attention function."""

from typing import Optional, Sequence

import torch
import omni_training_custom_ops  # noqa: F401  # pylint: disable=unused-import


def sink_attention(

from collections.abc import Mapping, Sequence
from typing import Any

# This package provides PyTorch-specific high-performance modules.
import torch
from torch import nn
from torch.nn import functional as F
from hyper_parallel.components.checkpoint.weight_conversion import WeightConverter

    query_ends, key_ends = (lengths.tolist() for lengths in sequence_lengths)
    q_flat, k_flat, v_flat = (tensor.flatten(0, 1) for tensor in (query, key, value))
    num_heads = q_flat.shape[1]
    causal_mask = torch.ones((2048, 2048), dtype=torch.bool, device=query.device).triu(diagonal=1)
    import torch_npu  # pylint: disable=import-outside-toplevel

    result = torch_npu.npu_fusion_attention(
        q_flat.contiguous(), k_flat.contiguous(), v_flat.contiguous(), num_heads, "TND",
        pse=None, padding_mask=None, atten_mask=causal_mask, scale=module.scaling,

from collections.abc import Mapping
from functools import partial
from typing import Any, Optional

import torch
from torch import nn
from transformers.activations import ACT2FN
from hyper_parallel.components.checkpoint.weight_conversion import (
    WeightConverter,
    WeightRenaming,