File size: 1,183 Bytes
388fd6e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 | """
Gamma Space Model: A lightweight PyTorch SSM block based on HiPPO-Gamma state space models.
Features TileLang-accelerated parallel scan kernels for GPU acceleration with optional Triton optimization.
Reference: https://github.com/state-spaces/mamba
"""
__version__ = "0.1.0"
# Core modules
from gamma_space_model.modules import (
SSMGamma,
GammaSingleBlock,
SSMGammaS4,
GammaS4Block,
GammaS4MinimalBlock,
S4TernaryDPLRSSM,
S4TernaryDPLRBlock,
LayerNorm,
RMSNorm,
)
# Optimized operations
try:
from gamma_space_model.ops import (
ssm_gamma_forward,
selective_scan_fwd,
HAS_TILELANG_OPS,
TILELANG_BACKEND,
)
_OPS_AVAILABLE = True
except ImportError:
_OPS_AVAILABLE = False
HAS_TILELANG_OPS = False
TILELANG_BACKEND = "unavailable"
__all__ = [
"SSMGamma",
"GammaSingleBlock",
"SSMGammaS4",
"GammaS4Block",
"GammaS4MinimalBlock",
"S4TernaryDPLRSSM",
"S4TernaryDPLRBlock",
"LayerNorm",
"RMSNorm",
"ssm_gamma_forward",
"selective_scan_fwd",
"HAS_TILELANG_OPS",
"TILELANG_BACKEND",
]
|