| """ | |
| 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", | |
| ] | |