USN 0.1.0__py3-none-any.whl

usn/__init__.py

@@ -0,0 +1,145 @@
+"""USN - Unified State Network Architecture Library.
+
+A production-grade Python package implementing the Unified State Network (USN)
+architecture for autoregressive sequence modeling with O(n) training complexity
+via associative parallel scan and O(1) inference memory via constant-size state.
+
+Quick Start:
+    >>> import usn
+    >>> model = usn.create_model("tiny")
+    >>> print(model.summary())
+
+Author: BUEORM
+License: MIT
+"""
+
+__version__ = "0.1.0"
+__author__ = "BUEORM"
+
+# Core classes
+# Backends
+from usn.backends import AccelerationLevel, AccelerationManager, DeviceDetector
+from usn.config import USNConfig, USNGenerationConfig, USNTrainingConfig
+from usn.inference import USNGenerator
+from usn.models import create_model
+from usn.models.usn_model import USNModel
+from usn.serialization.export import export_model
+from usn.serialization.reader import USNReader
+
+# Serialization
+from usn.serialization.writer import USNWriter
+from usn.training import USNTrainer
+from usn.utils.counting import count_parameters, estimate_flops, estimate_memory
+
+# Utilities
+from usn.utils.seed import set_seed
+
+
+# High-level API functions
+def save(model, path: str, **kwargs) -> None:
+    """Save a model to .usn format."""
+    writer = USNWriter()
+    config = getattr(model, "config", None)
+    writer.save(path, model, config=config, **kwargs)
+
+
+def load(path: str, map_location=None):
+    """Load a model from .usn format."""
+    reader = USNReader()
+    data = reader.load(path, map_location=map_location)
+    config = data.get("config")
+    if config is None:
+        raise ValueError("No config found in .usn file")
+    model = USNModel(config)
+    weights = data.get("weights", {})
+    if weights:
+        # Filter buffer keys for state_dict loading
+        state_dict = {k: v for k, v in weights.items() if not k.startswith("__buffer__.")}
+        model.load_state_dict(state_dict, strict=False)
+    return model
+
+
+def export(model, format: str, path: str, **kwargs) -> None:
+    """Export model to standard format (onnx, safetensors, state_dict, torchscript)."""
+    config = getattr(model, "config", None)
+    export_model(model, format, path, config=config, **kwargs)
+
+
+def generate(model, prompt: str, max_tokens: int = 256, tokenizer=None, **kwargs) -> str:
+    """Generate text from a model (convenience function)."""
+    if tokenizer is None:
+        raise ValueError("tokenizer is required for generate()")
+    gen = USNGenerator(model, tokenizer)
+    output = gen.generate(prompt, max_new_tokens=max_tokens, **kwargs)
+    return tokenizer.decode(output.token_ids[0].tolist())
+
+
+def train(model, dataset, config=None, **kwargs):
+    """Train a model (convenience function)."""
+    if config is None:
+        config = USNTrainingConfig()
+    trainer = USNTrainer(model, dataset, config, **kwargs)
+    return trainer.train()
+
+
+def summary(model) -> str:
+    """Get model summary string."""
+    if hasattr(model, "summary"):
+        return model.summary()
+    return f"Model with {count_parameters(model):,} parameters"
+
+
+def from_pretrained(path_or_id: str):
+    """Load a pretrained model from path."""
+    return load(path_or_id)
+
+
+def device_info() -> dict:
+    """Get available device information."""
+    return DeviceDetector.detect()
+
+
+def set_acceleration_level(level: int) -> None:
+    """Set the acceleration level manually."""
+    AccelerationManager.set_level(level)
+
+
+def benchmark_acceleration() -> dict:
+    """Compare throughput across acceleration levels."""
+    return {"current_level": AccelerationManager.get_level().name}
+
+
+__all__ = [
+    # Version
+    "__version__",
+    "__author__",
+    # Core classes
+    "USNModel",
+    "USNConfig",
+    "USNTrainingConfig",
+    "USNGenerationConfig",
+    "USNTrainer",
+    "USNGenerator",
+    # Factory
+    "create_model",
+    # High-level API
+    "save",
+    "load",
+    "export",
+    "generate",
+    "train",
+    "summary",
+    "from_pretrained",
+    # Utilities
+    "set_seed",
+    "count_parameters",
+    "estimate_memory",
+    "estimate_flops",
+    "device_info",
+    "set_acceleration_level",
+    "benchmark_acceleration",
+    # Backend
+    "AccelerationLevel",
+    "AccelerationManager",
+    "DeviceDetector",
+]

usn/__init__.pyi

@@ -0,0 +1,39 @@
+"""Type stubs for the USN public API."""
+
+from typing import Any
+
+import torch.nn as nn
+from torch import Tensor
+
+from usn.backends.acceleration import AccelerationLevel as AccelerationLevel
+from usn.backends.acceleration import AccelerationManager as AccelerationManager
+from usn.backends.detection import DeviceDetector as DeviceDetector
+from usn.config.generation_config import USNGenerationConfig as USNGenerationConfig
+from usn.config.model_config import USNConfig as USNConfig
+from usn.config.training_config import USNTrainingConfig as USNTrainingConfig
+from usn.inference.generator import USNGenerator as USNGenerator
+from usn.models import create_model as create_model
+from usn.models.usn_model import USNModel as USNModel
+from usn.training.trainer import USNTrainer as USNTrainer
+from usn.utils.counting import count_parameters as count_parameters
+from usn.utils.counting import estimate_flops as estimate_flops
+from usn.utils.counting import estimate_memory as estimate_memory
+from usn.utils.seed import set_seed as set_seed
+
+__version__: str
+__author__: str
+
+def save(model: nn.Module, path: str, **kwargs: Any) -> None: ...
+def load(path: str, map_location: str | None = None) -> USNModel: ...
+def export(model: nn.Module, format: str, path: str, **kwargs: Any) -> None: ...
+def generate(
+    model: nn.Module, prompt: str, max_tokens: int = 256, tokenizer: Any = None, **kwargs: Any
+) -> str: ...
+def train(
+    model: nn.Module, dataset: Any, config: USNTrainingConfig | None = None, **kwargs: Any
+) -> dict[str, Any]: ...
+def summary(model: nn.Module) -> str: ...
+def from_pretrained(path_or_id: str) -> USNModel: ...
+def device_info() -> dict[str, Any]: ...
+def set_acceleration_level(level: int) -> None: ...
+def benchmark_acceleration() -> dict[str, Any]: ...

usn/backends/__init__.py

@@ -0,0 +1,25 @@
+"""Backend acceleration: device detection, Triton kernels, and fallbacks.
+
+Public API:
+    - DeviceDetector: Hardware auto-detection (detect(), best_device())
+    - AccelerationLevel: Enum of 4 acceleration tiers
+    - AccelerationManager: Manages acceleration level with graceful fallback
+
+Importing this package triggers kernel registration at all acceleration levels:
+    - triton_kernels: Registers Level 1 (TRITON), Level 3 (AUTOGRAD), Level 4 (EAGER)
+    - fallbacks: Registers Level 2 (COMPILE) via torch.compile wrappers
+"""
+
+# Import fallbacks to register Level 2 (COMPILE) kernels
+import usn.backends.fallbacks as _fallbacks  # noqa: F401
+
+# Import triton_kernels to register Level 1, 3, 4 kernels
+import usn.backends.triton_kernels as _triton_kernels  # noqa: F401
+from usn.backends.acceleration import AccelerationLevel, AccelerationManager
+from usn.backends.detection import DeviceDetector
+
+__all__ = [
+    "AccelerationLevel",
+    "AccelerationManager",
+    "DeviceDetector",
+]

usn/backends/acceleration.py

@@ -0,0 +1,236 @@
+"""Acceleration level management for the USN library.
+
+Implements a 4-level acceleration hierarchy with automatic detection and
+graceful fallback. The levels are:
+
+1. TRITON  — Custom Triton fused kernels (fastest, requires triton + CUDA)
+2. COMPILE — torch.compile with inductor backend (requires PyTorch 2.0+ CUDA)
+3. AUTOGRAD — Custom autograd functions (requires CUDA)
+4. EAGER   — Standard PyTorch eager mode (always available, baseline)
+
+The system auto-detects the best level at import time and NEVER fails due to
+unavailable acceleration — it always falls back to EAGER as the last resort.
+"""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Callable
+from enum import IntEnum
+
+import torch
+
+logger = logging.getLogger(__name__)
+
+
+class AccelerationLevel(IntEnum):
+    """4-level acceleration hierarchy for USN kernels.
+
+    Lower numeric value = faster / more specialized.
+    Higher numeric value = more compatible / always available.
+    """
+
+    TRITON = 1
+    """Custom Triton fused kernels — maximum performance."""
+
+    COMPILE = 2
+    """torch.compile with inductor — good performance, no Triton needed."""
+
+    AUTOGRAD = 3
+    """Custom autograd functions — CUDA required, no JIT compilation."""
+
+    EAGER = 4
+    """Standard PyTorch eager mode — always available, baseline performance."""
+
+
+def _detect_best_level() -> AccelerationLevel:
+    """Detect the best available acceleration level.
+
+    Decision tree:
+    1. Is Triton installed AND CUDA available? → TRITON
+    2. Is torch.compile available (PyTorch ≥ 2.0) AND CUDA available? → COMPILE
+    3. Is CUDA available? → AUTOGRAD
+    4. Otherwise → EAGER
+
+    This function NEVER raises — it always returns a valid level.
+    """
+    # Check Level 1: Triton
+    if torch.cuda.is_available():
+        try:
+            import triton  # noqa: F401
+            import triton.language  # noqa: F401
+
+            return AccelerationLevel.TRITON
+        except (ImportError, RuntimeError, OSError):
+            pass
+
+    # Check Level 2: torch.compile with inductor
+    if torch.cuda.is_available() and hasattr(torch, "compile"):
+        try:
+
+            @torch.compile(backend="inductor")
+            def _test_fn(x: torch.Tensor) -> torch.Tensor:
+                return x + 1
+
+            _test_fn(torch.zeros(1, device="cuda"))
+            return AccelerationLevel.COMPILE
+        except Exception:
+            pass
+
+    # Check Level 3: Custom autograd (CUDA available but no compile)
+    if torch.cuda.is_available():
+        return AccelerationLevel.AUTOGRAD
+
+    # Level 4: Eager (CPU or unsupported hardware)
+    return AccelerationLevel.EAGER
+
+
+class AccelerationManager:
+    """Manages the 4-level acceleration hierarchy with graceful fallback.
+
+    This is a class-level (singleton-style) manager. All methods are classmethods
+    that operate on shared state. The acceleration level is auto-detected at
+    import time but can be overridden manually.
+
+    Usage:
+        # Auto-detected at import
+        level = AccelerationManager.get_level()
+
+        # Manual override
+        AccelerationManager.set_level(AccelerationLevel.EAGER)
+
+        # Get kernel for current level
+        kernel = AccelerationManager.get_kernel("projections")
+    """
+
+    _current_level: AccelerationLevel = AccelerationLevel.EAGER
+    _detected_level: AccelerationLevel = AccelerationLevel.EAGER
+    _initialized: bool = False
+    _kernel_registry: dict[str, dict[AccelerationLevel, Callable]] = {}
+
+    @classmethod
+    def _ensure_initialized(cls) -> None:
+        """Lazily initialize by detecting the best level."""
+        if not cls._initialized:
+            cls._detected_level = _detect_best_level()
+            cls._current_level = cls._detected_level
+            cls._initialized = True
+            logger.info(
+                "USN acceleration level auto-detected: %s",
+                cls._current_level.name,
+            )
+
+    @classmethod
+    def detect_best_level(cls) -> AccelerationLevel:
+        """Detect and return the best available acceleration level.
+
+        This re-runs detection (useful if hardware state has changed).
+
+        Returns:
+            The best AccelerationLevel available on this system.
+        """
+        cls._detected_level = _detect_best_level()
+        cls._current_level = cls._detected_level
+        cls._initialized = True
+        return cls._detected_level
+
+    @classmethod
+    def set_level(cls, level: int | AccelerationLevel) -> None:
+        """Manually set the acceleration level.
+
+        Use this to override auto-detection, e.g., for debugging or
+        benchmarking different backends.
+
+        Args:
+            level: The desired acceleration level (int 1-4 or AccelerationLevel).
+
+        Raises:
+            ValueError: If level is not a valid AccelerationLevel value.
+        """
+        cls._ensure_initialized()
+        try:
+            cls._current_level = AccelerationLevel(level)
+        except ValueError:
+            valid = [f"{l.name}={l.value}" for l in AccelerationLevel]
+            raise ValueError(
+                f"Invalid acceleration level: {level!r}. Valid levels: {', '.join(valid)}"
+            ) from None
+        logger.info("USN acceleration level set to: %s", cls._current_level.name)
+
+    @classmethod
+    def get_level(cls) -> AccelerationLevel:
+        """Get the current acceleration level.
+
+        Returns:
+            The currently active AccelerationLevel.
+        """
+        cls._ensure_initialized()
+        return cls._current_level
+
+    @classmethod
+    def get_kernel(cls, kernel_name: str) -> Callable:
+        """Get the appropriate kernel implementation for the current level.
+
+        Looks up the kernel in the registry for the current acceleration level.
+        If the kernel is not available at the current level, falls back to the
+        next available level (higher number = more compatible).
+
+        Args:
+            kernel_name: Name of the kernel to retrieve (e.g., "projections",
+                "temporal_gate", "state_core", "channel_mlp").
+
+        Returns:
+            Callable implementing the requested kernel at the best available level.
+
+        Raises:
+            KeyError: If kernel_name is not registered at any level.
+        """
+        cls._ensure_initialized()
+
+        if kernel_name not in cls._kernel_registry:
+            raise KeyError(
+                f"Unknown kernel '{kernel_name}'. "
+                f"Registered kernels: {list(cls._kernel_registry.keys())}"
+            )
+
+        level_map = cls._kernel_registry[kernel_name]
+
+        # Try current level, then fall back to higher (more compatible) levels
+        for level in sorted(AccelerationLevel, key=lambda l: l.value):
+            if level.value >= cls._current_level.value and level in level_map:
+                return level_map[level]
+
+        # Should not reach here if EAGER is always registered
+        raise KeyError(
+            f"No implementation found for kernel '{kernel_name}' "
+            f"at level {cls._current_level.name} or below."
+        )
+
+    @classmethod
+    def register_kernel(
+        cls,
+        kernel_name: str,
+        level: AccelerationLevel,
+        implementation: Callable,
+    ) -> None:
+        """Register a kernel implementation at a specific acceleration level.
+
+        This is used by the backend modules (triton_kernels, fallbacks) to
+        populate the registry.
+
+        Args:
+            kernel_name: Name of the kernel (e.g., "projections").
+            level: The acceleration level this implementation targets.
+            implementation: The callable implementing the kernel.
+        """
+        if kernel_name not in cls._kernel_registry:
+            cls._kernel_registry[kernel_name] = {}
+        cls._kernel_registry[kernel_name][level] = implementation
+
+    @classmethod
+    def reset(cls) -> None:
+        """Reset manager to uninitialized state. Useful for testing."""
+        cls._current_level = AccelerationLevel.EAGER
+        cls._detected_level = AccelerationLevel.EAGER
+        cls._initialized = False
+        cls._kernel_registry = {}

usn/backends/detection.py

@@ -0,0 +1,92 @@
+"""Hardware device detection for the USN library.
+
+Provides automatic detection of available compute hardware (CUDA GPUs, Apple MPS,
+CPU) and selection of the optimal device for model execution.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+
+
+class DeviceDetector:
+    """Auto-detect available hardware and capabilities.
+
+    Provides static methods for querying the runtime environment's compute
+    resources. Used by AccelerationManager to determine the best acceleration
+    level and by model initialization to select the target device.
+    """
+
+    @staticmethod
+    def detect() -> dict[str, Any]:
+        """Detect available hardware and return a capabilities dictionary.
+
+        Returns:
+            Dictionary containing:
+                - device: str — best available device name ("cuda", "mps", "cpu")
+                - cuda_available: bool — whether CUDA is available
+                - mps_available: bool — whether Apple MPS is available
+                - gpu_name: str | None — name of the CUDA GPU (if available)
+                - gpu_memory: int | None — total GPU memory in bytes (if available)
+                - compute_capability: tuple[int, int] | None — CUDA compute cap
+                - cuda_version: str | None — CUDA toolkit version string
+                - device_count: int — number of CUDA devices
+        """
+        cuda_available = torch.cuda.is_available()
+        mps_available = hasattr(torch.backends, "mps") and torch.backends.mps.is_available()
+
+        gpu_name: str | None = None
+        gpu_memory: int | None = None
+        compute_capability: tuple[int, int] | None = None
+        cuda_version: str | None = None
+        device_count = 0
+
+        if cuda_available:
+            device_count = torch.cuda.device_count()
+            try:
+                gpu_name = torch.cuda.get_device_name(0)
+            except Exception:
+                gpu_name = None
+            try:
+                props = torch.cuda.get_device_properties(0)
+                gpu_memory = props.total_mem
+                compute_capability = (props.major, props.minor)
+            except Exception:
+                pass
+            cuda_version = torch.version.cuda
+
+        # Determine best device string
+        if cuda_available:
+            device = "cuda"
+        elif mps_available:
+            device = "mps"
+        else:
+            device = "cpu"
+
+        return {
+            "device": device,
+            "cuda_available": cuda_available,
+            "mps_available": mps_available,
+            "gpu_name": gpu_name,
+            "gpu_memory": gpu_memory,
+            "compute_capability": compute_capability,
+            "cuda_version": cuda_version,
+            "device_count": device_count,
+        }
+
+    @staticmethod
+    def best_device() -> torch.device:
+        """Return the best available torch.device.
+
+        Priority: CUDA > MPS > CPU.
+
+        Returns:
+            torch.device for the best available hardware.
+        """
+        if torch.cuda.is_available():
+            return torch.device("cuda")
+        if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
+            return torch.device("mps")
+        return torch.device("cpu")

usn/backends/fallbacks.py

@@ -0,0 +1,102 @@
+"""PyTorch fallback implementations for all USN kernels.
+
+Provides Level 2 (torch.compile) wrappers and registers them with
+AccelerationManager. When torch.compile is available, these are
+JIT-compiled for better performance. When unavailable, they delegate
+directly to the Level 4 (eager) implementations.
+"""
+
+from __future__ import annotations
+
+import logging
+
+import torch
+
+from usn.backends.acceleration import AccelerationLevel, AccelerationManager
+from usn.backends.triton_kernels import (
+    eager_channel_mlp,
+    eager_fused_projections,
+    eager_state_core,
+    eager_temporal_gate,
+)
+
+logger = logging.getLogger(__name__)
+
+# Check if torch.compile is available
+_HAS_COMPILE = hasattr(torch, "compile")
+
+
+def _try_compile(fn):
+    """Attempt to wrap a function with torch.compile.
+
+    Uses 'reduce-overhead' mode which is optimized for repeated calls
+    with the same shapes (common in autoregressive generation and
+    fixed-length training batches).
+
+    If torch.compile is unavailable or compilation fails for any reason,
+    returns the original function unchanged — ensuring graceful fallback
+    to eager execution.
+
+    Note: torch.compile is lazy — errors may occur at first invocation
+    rather than at wrap time. The wrapper catches these runtime compilation
+    failures and falls back to the original function transparently.
+
+    Args:
+        fn: The function to compile.
+
+    Returns:
+        Compiled function if torch.compile is available and succeeds,
+        otherwise the original function unchanged.
+    """
+    if _HAS_COMPILE:
+        try:
+            compiled = torch.compile(fn, mode="reduce-overhead")
+        except Exception:
+            logger.debug(
+                "torch.compile wrapping failed for %s, using eager",
+                getattr(fn, "__name__", str(fn)),
+            )
+            return fn
+
+        # Wrap to catch lazy compilation failures at first call
+        def _safe_compiled(*args, **kwargs):
+            try:
+                return compiled(*args, **kwargs)
+            except Exception:
+                logger.debug(
+                    "torch.compile execution failed for %s, falling back to eager",
+                    getattr(fn, "__name__", str(fn)),
+                )
+                # Replace ourselves with the original to avoid repeated failures
+                _safe_compiled.__wrapped_fallback__ = True
+                return fn(*args, **kwargs)
+
+        # Preserve metadata for introspection
+        _safe_compiled.__name__ = getattr(fn, "__name__", "compiled_fn")
+        _safe_compiled.__qualname__ = getattr(fn, "__qualname__", "compiled_fn")
+        _safe_compiled.__doc__ = fn.__doc__
+        _safe_compiled.__wrapped__ = fn
+        _safe_compiled.__wrapped_fallback__ = False
+        return _safe_compiled
+
+    return fn
+
+
+# Compiled versions (Level 2)
+compiled_projections = _try_compile(eager_fused_projections)
+compiled_temporal_gate = _try_compile(eager_temporal_gate)
+compiled_state_core = _try_compile(eager_state_core)
+compiled_channel_mlp = _try_compile(eager_channel_mlp)
+
+# Register Level 2 (COMPILE) kernels
+AccelerationManager.register_kernel("projections", AccelerationLevel.COMPILE, compiled_projections)
+AccelerationManager.register_kernel(
+    "temporal_gate", AccelerationLevel.COMPILE, compiled_temporal_gate
+)
+AccelerationManager.register_kernel("state_core", AccelerationLevel.COMPILE, compiled_state_core)
+AccelerationManager.register_kernel("channel_mlp", AccelerationLevel.COMPILE, compiled_channel_mlp)
+
+logger.debug(
+    "Registered Level 2 (COMPILE) fallback kernels (torch.compile %s)",
+    "available" if _HAS_COMPILE else "NOT available — using eager",
+)