| from .selection import SelectionAttention, selection_attention_wrapper | |
| from .compression import CompressionAttention, compression_attention_wrapper | |
| from .sliding_window_attention import ( | |
| SlidingWindowAttention, | |
| sliding_window_attention_wrapper, | |
| ) | |
| from .top_k import TopKReduction, topk_reduction_wrapper | |
| class NSANamespace: | |
| SelectionAttention = staticmethod(SelectionAttention) | |
| selection_attention_wrapper = staticmethod(selection_attention_wrapper) | |
| SlidingWindowAttention = staticmethod(SlidingWindowAttention) | |
| sliding_window_attention_wrapper = staticmethod(sliding_window_attention_wrapper) | |
| CompressionAttention = staticmethod(CompressionAttention) | |
| compression_attention_wrapper = staticmethod(compression_attention_wrapper) | |
| TopKReduction = staticmethod(TopKReduction) | |
| topk_reduction_wrapper = staticmethod(topk_reduction_wrapper) | |
| NSA = NSANamespace() | |
| __all__ = [ | |
| "NSA", | |
| ] | |