"""
MK3 Core Module

Complete salience scoring system with normalization invariant
and advanced attention mechanisms.
"""

from .salience_formula import GibbsSalienceFormula
from .salience_layers import (
    SalienceAttentionLayer,
    SalienceSelectiveLayer,
    SalienceTransformerBlock
)
from .continuous_embeddings import ContinuousEmbedding

# Advanced attention mechanisms
from .flash_attention import (
    FlashAttention,
    FlashMultiheadAttention,
    create_flash_attention_layer
)
from .moe_attention import (
    MoEAttention,
    SparseMoEAttention,
    TopKRouter,
    create_moe_attention
)
from .ring_attention import (
    RingAttention,
    BlockwiseAttention,
    StripedAttention,
    create_ring_attention
)
from .sparse_attention import (
    LocalAttention,
    GlobalLocalAttention,
    BigBirdAttention,
    StridedAttention,
    create_sparse_attention
)
from .streaming_llm import (
    StreamingAttention,
    StreamingTransformerBlock,
    StreamingLLM,
    KVCache
)

__all__ = [
    # Original salience components
    'GibbsSalienceFormula',
    'SalienceAttentionLayer',
    'SalienceSelectiveLayer',
    'SalienceTransformerBlock',
    'ContinuousEmbedding',

    # FlashAttention
    'FlashAttention',
    'FlashMultiheadAttention',
    'create_flash_attention_layer',

    # MoE Attention
    'MoEAttention',
    'SparseMoEAttention',
    'TopKRouter',
    'create_moe_attention',

    # Ring Attention
    'RingAttention',
    'BlockwiseAttention',
    'StripedAttention',
    'create_ring_attention',

    # Sparse Attention
    'LocalAttention',
    'GlobalLocalAttention',
    'BigBirdAttention',
    'StridedAttention',
    'create_sparse_attention',

    # StreamingLLM
    'StreamingAttention',
    'StreamingTransformerBlock',
    'StreamingLLM',
    'KVCache',
]
