"""Attention layer replacement system for sageLLM.

This module implements custom attention layers with support for:
- PagedAttention for efficient memory management
- FlashAttention for optimized computation
- Continuous batching with fine-grained KV cache control
- Integration with sagellm-backend AttentionBackend
- Integration with sagellm-kv-cache KVCacheManager

This addresses issue #37 and depends on:
- Issue #35: LayerWise model builder (for layer replacement)
- Issue #36: Kernel injection framework (for operator replacement)
- sagellm-backend: AttentionBackend implementations
- sagellm-kv-cache: KVCacheManager

Example:
    # Replace HF attention with SageAttention
    from sagellm_core.attention import SageAttention, replace_attention_layers

    # Create attention backend
    from sagellm_backend.attention import PagedAttentionBackend
    backend = PagedAttentionBackend(block_size=16)

    # Create KV cache manager
    from sagellm_kv_cache import KVCacheManager
    kv_manager = KVCacheManager(max_tokens=4096, block_size=16)

    # Replace attention in model
    model = replace_attention_layers(
        model,
        attention_backend=backend,
        kv_cache_manager=kv_manager
    )
"""

from __future__ import annotations

from sagellm_core.attention.layer_registry import (
    AttentionLayerFactory,
    AttentionRegistry,
    register_attention_layer,
    replace_attention_layers,
)
from sagellm_core.attention.sage_attention import (
    SageAttentionRuntimeContext,
    SageAttention,
    SageAttentionConfig,
    SageAttentionOutput,
    get_sage_attention_runtime_context,
    use_sage_attention_runtime_context,
)

__all__ = [
    # Core attention layer
    "SageAttention",
    "SageAttentionRuntimeContext",
    "SageAttentionConfig",
    "SageAttentionOutput",
    "get_sage_attention_runtime_context",
    "use_sage_attention_runtime_context",
    # Layer replacement system
    "AttentionLayerFactory",
    "AttentionRegistry",
    "register_attention_layer",
    "replace_attention_layers",
]
