flash-xlstm-exp 0.1.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (33) hide show
  1. flash_xlstm_exp-0.1.0/LICENSE +102 -0
  2. flash_xlstm_exp-0.1.0/MANIFEST.in +3 -0
  3. flash_xlstm_exp-0.1.0/NOTICE +19 -0
  4. flash_xlstm_exp-0.1.0/PKG-INFO +217 -0
  5. flash_xlstm_exp-0.1.0/README.md +86 -0
  6. flash_xlstm_exp-0.1.0/flash_xlstm_exp.egg-info/PKG-INFO +217 -0
  7. flash_xlstm_exp-0.1.0/flash_xlstm_exp.egg-info/SOURCES.txt +31 -0
  8. flash_xlstm_exp-0.1.0/flash_xlstm_exp.egg-info/dependency_links.txt +1 -0
  9. flash_xlstm_exp-0.1.0/flash_xlstm_exp.egg-info/requires.txt +6 -0
  10. flash_xlstm_exp-0.1.0/flash_xlstm_exp.egg-info/top_level.txt +1 -0
  11. flash_xlstm_exp-0.1.0/fsp/__init__.py +0 -0
  12. flash_xlstm_exp-0.1.0/fsp/benchmark/__init__.py +0 -0
  13. flash_xlstm_exp-0.1.0/fsp/benchmark/bench_autotune_vs_baseline.py +91 -0
  14. flash_xlstm_exp-0.1.0/fsp/kernels/__init__.py +0 -0
  15. flash_xlstm_exp-0.1.0/fsp/kernels/autotune.py +222 -0
  16. flash_xlstm_exp-0.1.0/fsp/vendor/__init__.py +0 -0
  17. flash_xlstm_exp-0.1.0/fsp/vendor/torch_fw/__init__.py +0 -0
  18. flash_xlstm_exp-0.1.0/fsp/vendor/torch_fw/chunkwise_gates.py +95 -0
  19. flash_xlstm_exp-0.1.0/fsp/vendor/torch_fw/fw.py +136 -0
  20. flash_xlstm_exp-0.1.0/fsp/vendor/torch_fw/fw_parallel.py +136 -0
  21. flash_xlstm_exp-0.1.0/fsp/vendor/torch_fw/fw_recurrent.py +144 -0
  22. flash_xlstm_exp-0.1.0/fsp/vendor/triton_kernels/__init__.py +0 -0
  23. flash_xlstm_exp-0.1.0/fsp/vendor/triton_kernels/chunkwise_param_heuristics.py +180 -0
  24. flash_xlstm_exp-0.1.0/fsp/vendor/triton_kernels/fw_kernel_parallel.py +289 -0
  25. flash_xlstm_exp-0.1.0/fsp/vendor/triton_kernels/fw_kernel_recurrent.py +230 -0
  26. flash_xlstm_exp-0.1.0/fsp/vendor/triton_kernels/kernel_param_heuristics.py +16 -0
  27. flash_xlstm_exp-0.1.0/fsp/vendor/triton_kernels/xl_chunk_init.py +9 -0
  28. flash_xlstm_exp-0.1.0/fsp/vendor/utils/__init__.py +0 -0
  29. flash_xlstm_exp-0.1.0/fsp/vendor/utils/kernels.py +7 -0
  30. flash_xlstm_exp-0.1.0/fsp/vendor/utils/torch_utils.py +70 -0
  31. flash_xlstm_exp-0.1.0/pyproject.toml +39 -0
  32. flash_xlstm_exp-0.1.0/setup.cfg +4 -0
  33. flash_xlstm_exp-0.1.0/tests/test_autotune_equivalence.py +95 -0
@@ -0,0 +1,102 @@
1
+ NXAI COMMUNITY LICENSE AGREEMENT
2
+
3
+ Preamble 1
4
+ We are proud to present the NXAI xLSTM 7B model and software, demonstrating the strength of next-generation RNN-based large language models, delivering high-quality performance and fast inference speeds. While xLSTM 7B is freely available for open research and development, we believe that organizations significantly benefiting from our technology should contribute back. Our goal is to support research, small and medium-sized enterprises (SMEs), and open innovation, while ensuring that large enterprises who incorporate xLSTM 7B into commercial products or services fairly compensate the creators for their research and development efforts.
5
+ Linz, December 12, 2024.
6
+
7
+ Preamble 2
8
+ The NXAI COMMUNITY LICENSE AGREEMENT is based on the META LLAMA 3 COMMUNITY LICENSE AGREEMENT and contains some modifications, especially Section 2, “Additional Commercial Terms” is different.
9
+
10
+ “Agreement” means the terms and conditions for use, reproduction, distribution and modification of the NXAI Materials set forth herein.
11
+ “Documentation” means the specifications, manuals and documentation accompanying NXAI Materials distributed by NXAI at https://github.com/NX-AI/.
12
+ “Licensee” or “you” means you, or your employer or any other person or entity (if you are entering into this Agreement on such person or entity’s behalf), of the age required under applicable laws, rules or regulations to provide legal consent and that has legal authority to bind your employer or such other person or entity if you are entering in this Agreement on their behalf.
13
+ “NXAI Materials” means, collectively, NXAI’s proprietary large language models, algorithms and any Software, including machine-learning model code, trained model weights, inference-enabling code, training-enabling code, fine-tuning enabling code and all other work of NXAI in the field of neural networks, Documentation (and any portion thereof) made available under this Agreement.
14
+ “NXAI” or “we” means NXAI GmbH, Linz, Austria.
15
+
16
+ By using or distributing any portion or element of the NXAI Materials, you agree to be bound by this Agreement.
17
+
18
+ 1. License Rights and Redistribution.
19
+
20
+ a. Grant of Rights. You are granted a non-exclusive, worldwide, non-transferable and royalty-free limited license under NXAI’s intellectual property embodied in the NXAI Materials to use, reproduce, distribute, copy, create derivative works of, and make modifications to the NXAI Materials.
21
+
22
+ b. Redistribution and Use.
23
+
24
+ i. If you distribute or make available the NXAI Materials (or any derivative works thereof), or a product or service that uses any of them, including another AI model, you shall (A) provide a copy of this Agreement with any such NXAI Materials; and (B) prominently display “Built with technology from NXAI” on a related website, user interface, blogpost, about page, or product documentation.
25
+
26
+ ii. If you receive NXAI Materials, or any derivative works thereof, from a Licensee as part of an integrated end user product, then Section 2 of this Agreement will not apply to you.
27
+
28
+ iii. You must retain in all copies of the NXAI Materials that you distribute the following attribution notice within a “Notice” text file distributed as a part of such copies: “This product includes materials developed at NXAI that are licensed under the NXAI Community License, Copyright © NXAI GmbH, All Rights Reserved.”
29
+
30
+ 2. Additional Commercial Terms. If (a) the Licensee, on a consolidated basis (including parent, subsidiaries, and affiliates), exceeds the annual revenue of one hundred million Euros (€100,000,000), and (b) the Licensee incorporates NXAI Material, in whole or in part, into a Commercial Product or Service, then the Licensee must obtain a commercial license from NXAI, which NXAI may grant to you in its sole discretion, and you are not authorized to exercise any of the rights under this Agreement unless or until NXAI otherwise expressly grants you such rights
31
+
32
+ 3. Disclaimer of Warranty. UNLESS REQUIRED BY APPLICABLE LAW, THE NXAI MATERIALS AND ANY OUTPUT AND RESULTS THEREFROM ARE PROVIDED ON AN “AS IS” BASIS, WITHOUT WARRANTIES OF ANY KIND, AND NXAI DISCLAIMS ALL WARRANTIES OF ANY KIND, BOTH EXPRESS AND IMPLIED, INCLUDING, WITHOUT LIMITATION, ANY WARRANTIES OF TITLE, NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. YOU ARE SOLELY RESPONSIBLE FOR DETERMINING THE APPROPRIATENESS OF USING OR REDISTRIBUTING THE NXAI MATERIALS AND ASSUME ANY RISKS ASSOCIATED WITH YOUR USE OF THE NXAI MATERIALS AND ANY OUTPUT AND RESULTS.
33
+
34
+ 4. Limitation of Liability. IN NO EVENT WILL NXAI OR ITS AFFILIATES BE LIABLE UNDER ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, TORT, NEGLIGENCE, PRODUCTS LIABILITY, OR OTHERWISE, ARISING OUT OF THIS AGREEMENT, FOR ANY LOST PROFITS OR ANY INDIRECT, SPECIAL, CONSEQUENTIAL, INCIDENTAL, EXEMPLARY OR PUNITIVE DAMAGES, EVEN IF NXAI OR ITS AFFILIATES HAVE BEEN ADVISED OF THE POSSIBILITY OF ANY OF THE FOREGOING.
35
+
36
+ 5. Intellectual Property.
37
+
38
+ a. No trademark licenses are granted under this Agreement, and in connection with the NXAI Materials, neither NXAI nor Licensee may use any name or mark owned by or associated with the other or any of its affiliates, except as required for reasonable and customary use in describing and redistributing the NXAI Materials or as set forth in this Section 5(a). NXAI hereby grants you a license to use “NXAI” (the “Mark”) solely as required to comply with the last sentence of Section 1.b.i. All goodwill arising out of your use of the Mark will insure to the benefit of NXAI.
39
+
40
+ b. Subject to NXAI’s ownership of NXAI Materials and derivatives made by or for NXAI, with respect to any derivative works and modifications of the NXAI Materials that are made by you, as between you and NXAI, you are and will be the owner of such derivative works and modifications.
41
+
42
+ c. If you institute litigation or other proceedings against NXAI or any entity (including a cross-claim or counterclaim in a lawsuit) alleging that the NXAI Materials or models released by NXAI outputs or results, or any portion of any of the foregoing, constitutes infringement of intellectual property or other rights owned or licensable by you, then any licenses granted to you under this Agreement shall terminate as of the date such litigation or claim is filed or instituted. You will indemnify and hold harmless NXAI from and against any claim by any third party arising out of or related to your use or distribution of the NXAI Materials.
43
+
44
+ 6. Term and Termination. The term of this Agreement will commence upon your acceptance of this Agreement or access to the NXAI Materials and will continue in full force and effect until terminated in accordance with the terms and conditions herein. NXAI may terminate this Agreement if you are in breach of any term or condition of this Agreement. Upon termination of this Agreement, you shall delete and cease use of the NXAI Materials. Sections 3, 4 and 7 shall survive the termination of this Agreement.
45
+
46
+ 7. Governing Law and Jurisdiction. This Agreement shall be governed by and construed in accordance with the laws of the Republic of Austria, without regard to its conflict of laws principles. The courts located in Linz, Austria shall have exclusive jurisdiction over any disputes arising out of or in connection with this Agreement.
47
+
48
+ ====================================================================================================
49
+
50
+ This product includes software licensed under the MIT License:
51
+
52
+ MIT License
53
+
54
+ Permission is hereby granted, free of charge, to any person obtaining
55
+ a copy of this software and associated documentation files
56
+ (the "Software"), to deal in the Software without restriction,
57
+ including without limitation the rights to use, copy, modify, merge,
58
+ publish, distribute, sublicense, and/or sell copies of the Software,
59
+ and to permit persons to whom the Software is furnished to do so,
60
+ subject to the following conditions:
61
+
62
+ The above copyright notice and this permission notice shall be
63
+ included in all copies or substantial portions of the Software.
64
+
65
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
66
+ EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
67
+ MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
68
+ IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
69
+ CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
70
+ TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
71
+ SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
72
+
73
+ ====================================================================================================
74
+
75
+ This product includes software licensed under the BSD-3-Clause License.
76
+
77
+ BSD 3-Clause License
78
+
79
+ Redistribution and use in source and binary forms, with or without
80
+ modification, are permitted provided that the following conditions are met:
81
+
82
+ * Redistributions of source code must retain the above copyright notice, this
83
+ list of conditions and the following disclaimer.
84
+
85
+ * Redistributions in binary form must reproduce the above copyright notice,
86
+ this list of conditions and the following disclaimer in the documentation
87
+ and/or other materials provided with the distribution.
88
+
89
+ * Neither the name of the copyright holder nor the names of its
90
+ contributors may be used to endorse or promote products derived from
91
+ this software without specific prior written permission.
92
+
93
+ THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
94
+ AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
95
+ IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
96
+ DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
97
+ FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
98
+ DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
99
+ SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
100
+ CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
101
+ OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
102
+ OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
@@ -0,0 +1,3 @@
1
+ include LICENSE
2
+ include NOTICE
3
+ include README.md
@@ -0,0 +1,19 @@
1
+ This product includes materials developed at NXAI that are licensed under
2
+ the NXAI Community License, Copyright © NXAI GmbH, All Rights Reserved.
3
+
4
+ Portions of this repository (fsp/vendor/**) are vendored, with minimal
5
+ import-path adaptations, from:
6
+
7
+ NX-AI/mlstm_kernels
8
+ https://github.com/NX-AI/mlstm_kernels
9
+ Copyright (c) NXAI GmbH
10
+ Authors: Maximilian Beck, Korbinian Pöppel, Phillip Lippe, and contributors
11
+ Licensed under the NXAI Community License Agreement (see LICENSE)
12
+
13
+ Original paper:
14
+
15
+ Tiled Flash Linear Attention: More Efficient Linear RNN and xLSTM Kernels
16
+ Maximilian Beck, Korbinian Pöppel, Phillip Lippe, Sepp Hochreiter
17
+ https://arxiv.org/abs/2503.14376
18
+
19
+ flash_xlstm_exp (fsp) — Built with technology from NXAI.
@@ -0,0 +1,217 @@
1
+ Metadata-Version: 2.4
2
+ Name: flash-xlstm-exp
3
+ Version: 0.1.0
4
+ Summary: Fork-and-upgrade experiments on top of NX-AI TFLA mLSTM Triton kernels.
5
+ Author: Nata
6
+ License: NXAI COMMUNITY LICENSE AGREEMENT
7
+
8
+ Preamble 1
9
+ We are proud to present the NXAI xLSTM 7B model and software, demonstrating the strength of next-generation RNN-based large language models, delivering high-quality performance and fast inference speeds. While xLSTM 7B is freely available for open research and development, we believe that organizations significantly benefiting from our technology should contribute back. Our goal is to support research, small and medium-sized enterprises (SMEs), and open innovation, while ensuring that large enterprises who incorporate xLSTM 7B into commercial products or services fairly compensate the creators for their research and development efforts.
10
+ Linz, December 12, 2024.
11
+
12
+ Preamble 2
13
+ The NXAI COMMUNITY LICENSE AGREEMENT is based on the META LLAMA 3 COMMUNITY LICENSE AGREEMENT and contains some modifications, especially Section 2, “Additional Commercial Terms” is different.
14
+
15
+ “Agreement” means the terms and conditions for use, reproduction, distribution and modification of the NXAI Materials set forth herein.
16
+ “Documentation” means the specifications, manuals and documentation accompanying NXAI Materials distributed by NXAI at https://github.com/NX-AI/.
17
+ “Licensee” or “you” means you, or your employer or any other person or entity (if you are entering into this Agreement on such person or entity’s behalf), of the age required under applicable laws, rules or regulations to provide legal consent and that has legal authority to bind your employer or such other person or entity if you are entering in this Agreement on their behalf.
18
+ “NXAI Materials” means, collectively, NXAI’s proprietary large language models, algorithms and any Software, including machine-learning model code, trained model weights, inference-enabling code, training-enabling code, fine-tuning enabling code and all other work of NXAI in the field of neural networks, Documentation (and any portion thereof) made available under this Agreement.
19
+ “NXAI” or “we” means NXAI GmbH, Linz, Austria.
20
+
21
+ By using or distributing any portion or element of the NXAI Materials, you agree to be bound by this Agreement.
22
+
23
+ 1. License Rights and Redistribution.
24
+
25
+ a. Grant of Rights. You are granted a non-exclusive, worldwide, non-transferable and royalty-free limited license under NXAI’s intellectual property embodied in the NXAI Materials to use, reproduce, distribute, copy, create derivative works of, and make modifications to the NXAI Materials.
26
+
27
+ b. Redistribution and Use.
28
+
29
+ i. If you distribute or make available the NXAI Materials (or any derivative works thereof), or a product or service that uses any of them, including another AI model, you shall (A) provide a copy of this Agreement with any such NXAI Materials; and (B) prominently display “Built with technology from NXAI” on a related website, user interface, blogpost, about page, or product documentation.
30
+
31
+ ii. If you receive NXAI Materials, or any derivative works thereof, from a Licensee as part of an integrated end user product, then Section 2 of this Agreement will not apply to you.
32
+
33
+ iii. You must retain in all copies of the NXAI Materials that you distribute the following attribution notice within a “Notice” text file distributed as a part of such copies: “This product includes materials developed at NXAI that are licensed under the NXAI Community License, Copyright © NXAI GmbH, All Rights Reserved.”
34
+
35
+ 2. Additional Commercial Terms. If (a) the Licensee, on a consolidated basis (including parent, subsidiaries, and affiliates), exceeds the annual revenue of one hundred million Euros (€100,000,000), and (b) the Licensee incorporates NXAI Material, in whole or in part, into a Commercial Product or Service, then the Licensee must obtain a commercial license from NXAI, which NXAI may grant to you in its sole discretion, and you are not authorized to exercise any of the rights under this Agreement unless or until NXAI otherwise expressly grants you such rights
36
+
37
+ 3. Disclaimer of Warranty. UNLESS REQUIRED BY APPLICABLE LAW, THE NXAI MATERIALS AND ANY OUTPUT AND RESULTS THEREFROM ARE PROVIDED ON AN “AS IS” BASIS, WITHOUT WARRANTIES OF ANY KIND, AND NXAI DISCLAIMS ALL WARRANTIES OF ANY KIND, BOTH EXPRESS AND IMPLIED, INCLUDING, WITHOUT LIMITATION, ANY WARRANTIES OF TITLE, NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. YOU ARE SOLELY RESPONSIBLE FOR DETERMINING THE APPROPRIATENESS OF USING OR REDISTRIBUTING THE NXAI MATERIALS AND ASSUME ANY RISKS ASSOCIATED WITH YOUR USE OF THE NXAI MATERIALS AND ANY OUTPUT AND RESULTS.
38
+
39
+ 4. Limitation of Liability. IN NO EVENT WILL NXAI OR ITS AFFILIATES BE LIABLE UNDER ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, TORT, NEGLIGENCE, PRODUCTS LIABILITY, OR OTHERWISE, ARISING OUT OF THIS AGREEMENT, FOR ANY LOST PROFITS OR ANY INDIRECT, SPECIAL, CONSEQUENTIAL, INCIDENTAL, EXEMPLARY OR PUNITIVE DAMAGES, EVEN IF NXAI OR ITS AFFILIATES HAVE BEEN ADVISED OF THE POSSIBILITY OF ANY OF THE FOREGOING.
40
+
41
+ 5. Intellectual Property.
42
+
43
+ a. No trademark licenses are granted under this Agreement, and in connection with the NXAI Materials, neither NXAI nor Licensee may use any name or mark owned by or associated with the other or any of its affiliates, except as required for reasonable and customary use in describing and redistributing the NXAI Materials or as set forth in this Section 5(a). NXAI hereby grants you a license to use “NXAI” (the “Mark”) solely as required to comply with the last sentence of Section 1.b.i. All goodwill arising out of your use of the Mark will insure to the benefit of NXAI.
44
+
45
+ b. Subject to NXAI’s ownership of NXAI Materials and derivatives made by or for NXAI, with respect to any derivative works and modifications of the NXAI Materials that are made by you, as between you and NXAI, you are and will be the owner of such derivative works and modifications.
46
+
47
+ c. If you institute litigation or other proceedings against NXAI or any entity (including a cross-claim or counterclaim in a lawsuit) alleging that the NXAI Materials or models released by NXAI outputs or results, or any portion of any of the foregoing, constitutes infringement of intellectual property or other rights owned or licensable by you, then any licenses granted to you under this Agreement shall terminate as of the date such litigation or claim is filed or instituted. You will indemnify and hold harmless NXAI from and against any claim by any third party arising out of or related to your use or distribution of the NXAI Materials.
48
+
49
+ 6. Term and Termination. The term of this Agreement will commence upon your acceptance of this Agreement or access to the NXAI Materials and will continue in full force and effect until terminated in accordance with the terms and conditions herein. NXAI may terminate this Agreement if you are in breach of any term or condition of this Agreement. Upon termination of this Agreement, you shall delete and cease use of the NXAI Materials. Sections 3, 4 and 7 shall survive the termination of this Agreement.
50
+
51
+ 7. Governing Law and Jurisdiction. This Agreement shall be governed by and construed in accordance with the laws of the Republic of Austria, without regard to its conflict of laws principles. The courts located in Linz, Austria shall have exclusive jurisdiction over any disputes arising out of or in connection with this Agreement.
52
+
53
+ ====================================================================================================
54
+
55
+ This product includes software licensed under the MIT License:
56
+
57
+ MIT License
58
+
59
+ Permission is hereby granted, free of charge, to any person obtaining
60
+ a copy of this software and associated documentation files
61
+ (the "Software"), to deal in the Software without restriction,
62
+ including without limitation the rights to use, copy, modify, merge,
63
+ publish, distribute, sublicense, and/or sell copies of the Software,
64
+ and to permit persons to whom the Software is furnished to do so,
65
+ subject to the following conditions:
66
+
67
+ The above copyright notice and this permission notice shall be
68
+ included in all copies or substantial portions of the Software.
69
+
70
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
71
+ EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
72
+ MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
73
+ IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
74
+ CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
75
+ TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
76
+ SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
77
+
78
+ ====================================================================================================
79
+
80
+ This product includes software licensed under the BSD-3-Clause License.
81
+
82
+ BSD 3-Clause License
83
+
84
+ Redistribution and use in source and binary forms, with or without
85
+ modification, are permitted provided that the following conditions are met:
86
+
87
+ * Redistributions of source code must retain the above copyright notice, this
88
+ list of conditions and the following disclaimer.
89
+
90
+ * Redistributions in binary form must reproduce the above copyright notice,
91
+ this list of conditions and the following disclaimer in the documentation
92
+ and/or other materials provided with the distribution.
93
+
94
+ * Neither the name of the copyright holder nor the names of its
95
+ contributors may be used to endorse or promote products derived from
96
+ this software without specific prior written permission.
97
+
98
+ THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
99
+ AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
100
+ IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
101
+ DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
102
+ FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
103
+ DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
104
+ SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
105
+ CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
106
+ OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
107
+ OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
108
+
109
+ Project-URL: Repository, https://github.com/Natarizki/flash_xlstm_exp
110
+ Project-URL: Upstream, https://github.com/NX-AI/mlstm_kernels
111
+ Project-URL: Author, https://github.com/Natarizki
112
+ Keywords: xlstm,mlstm,triton,flash-attention,linear-attention,tfla,deep-learning
113
+ Classifier: Development Status :: 2 - Pre-Alpha
114
+ Classifier: Intended Audience :: Developers
115
+ Classifier: Intended Audience :: Science/Research
116
+ Classifier: License :: Other/Proprietary License
117
+ Classifier: Programming Language :: Python :: 3
118
+ Classifier: Programming Language :: Python :: 3.11
119
+ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
120
+ Classifier: Environment :: GPU :: NVIDIA CUDA
121
+ Requires-Python: >=3.11
122
+ Description-Content-Type: text/markdown
123
+ License-File: LICENSE
124
+ License-File: NOTICE
125
+ Requires-Dist: torch>=2.4
126
+ Requires-Dist: triton>=3.0
127
+ Requires-Dist: einops>=0.7
128
+ Provides-Extra: dev
129
+ Requires-Dist: pytest>=8.0; extra == "dev"
130
+ Dynamic: license-file
131
+
132
+ # flash_xlstm_exp (fsp)
133
+
134
+ Experimental fork-and-upgrade of NX-AI's **TFLA** (Tiled Flash Linear
135
+ Attention) mLSTM kernels, targeting further GPU/Triton-level optimization
136
+ on top of the official implementation.
137
+
138
+ Built with technology from NXAI. See [NOTICE](./NOTICE) and
139
+ [LICENSE](./LICENSE) — this repo is derived from
140
+ [`NX-AI/mlstm_kernels`](https://github.com/NX-AI/mlstm_kernels), licensed
141
+ under the **NXAI Community License Agreement**, not a permissive OSS
142
+ license. Redistributing fsp (or derivatives of it) requires keeping the
143
+ LICENSE + NOTICE files and the "Built with technology from NXAI" attribution
144
+ intact. See the NXAI Community License for the (currently inapplicable to
145
+ solo/small-scale use) large-enterprise commercial terms.
146
+
147
+ ## Upgrade roadmap
148
+
149
+ | # | Upgrade | Status |
150
+ |---|---------|--------|
151
+ | 1 | Dynamic tile-size autotuner (`triton.autotune` in place of the static `get_head_dim_block_size` / hardcoded `siz_b_LQ=32,siz_b_LKV=32` heuristic) | **Implemented** — `fsp/kernels/autotune.py` |
152
+ | 2 | Fused recurrent+parallel forward kernel (avoid materializing intermediate chunk states to HBM between the two kernel launches) | Planned |
153
+ | 3 | Quantized (fp8/int8) memory state `C` for long-context inference | Planned |
154
+
155
+ ## Why GPU-only
156
+
157
+ TFLA's kernels are written in [Triton](https://github.com/triton-lang/triton)
158
+ and target NVIDIA/AMD GPUs via CUDA/ROCm. There is currently no Triton
159
+ backend for mobile/ARM64-without-GPU environments (e.g. Termux on Android),
160
+ so fsp's kernels can be *written and code-reviewed* in such an environment,
161
+ but must be **benchmarked and numerically validated on an actual CUDA (or
162
+ ROCm) GPU** — a local NVIDIA GPU, or a rented cloud GPU instance.
163
+
164
+ ## Layout
165
+
166
+ ```
167
+ fsp/
168
+ ├── vendor/ # Baseline TFLA, vendored from NX-AI/mlstm_kernels
169
+ │ ├── triton_kernels/ # - the @triton.jit kernels themselves
170
+ │ ├── torch_fw/ # - PyTorch launchers/wrappers (forward pass only, so far)
171
+ │ └── utils/ # - small shared dtype/validation helpers
172
+ ├── kernels/
173
+ │ └── autotune.py # Upgrade #1: autotuned forward-parallel kernel
174
+ └── benchmark/
175
+ └── bench_autotune_vs_baseline.py
176
+ tests/
177
+ └── test_autotune_equivalence.py
178
+ ```
179
+
180
+ > **Scope note:** only the *forward* pass (recurrent + parallel kernels) is
181
+ > vendored so far. Backward-pass kernels
182
+ > (`bw_kernel_parallel_dK/dQ/dV`, `bw_kernel_recurrent`) will be vendored in
183
+ > a later batch once the forward-pass upgrades are validated on real
184
+ > hardware.
185
+
186
+ ## Install
187
+
188
+ ```bash
189
+ pip install torch triton
190
+ pip install -e .
191
+ ```
192
+
193
+ ## Usage
194
+
195
+ ```python
196
+ from fsp.kernels.autotune import mlstm_chunkwise__parallel_fw_Hintra_autotuned
197
+
198
+ matH_out, vecN_out, vecM_out = mlstm_chunkwise__parallel_fw_Hintra_autotuned(
199
+ matQ, matK, matV, vecI, vecF,
200
+ matC_states, vecN_states, scaMinter_states,
201
+ chunk_size=128,
202
+ )
203
+ ```
204
+
205
+ Drop-in replacement for
206
+ `fsp.vendor.torch_fw.fw_parallel.mlstm_chunkwise__parallel_fw_Hintra` — same
207
+ inputs/outputs, same math, but tile sizes are chosen empirically per shape
208
+ instead of by static heuristic.
209
+
210
+ ## Running tests / benchmarks
211
+
212
+ Both require a CUDA-capable GPU:
213
+
214
+ ```bash
215
+ pytest tests/test_autotune_equivalence.py
216
+ python -m fsp.benchmark.bench_autotune_vs_baseline
217
+ ```
@@ -0,0 +1,86 @@
1
+ # flash_xlstm_exp (fsp)
2
+
3
+ Experimental fork-and-upgrade of NX-AI's **TFLA** (Tiled Flash Linear
4
+ Attention) mLSTM kernels, targeting further GPU/Triton-level optimization
5
+ on top of the official implementation.
6
+
7
+ Built with technology from NXAI. See [NOTICE](./NOTICE) and
8
+ [LICENSE](./LICENSE) — this repo is derived from
9
+ [`NX-AI/mlstm_kernels`](https://github.com/NX-AI/mlstm_kernels), licensed
10
+ under the **NXAI Community License Agreement**, not a permissive OSS
11
+ license. Redistributing fsp (or derivatives of it) requires keeping the
12
+ LICENSE + NOTICE files and the "Built with technology from NXAI" attribution
13
+ intact. See the NXAI Community License for the (currently inapplicable to
14
+ solo/small-scale use) large-enterprise commercial terms.
15
+
16
+ ## Upgrade roadmap
17
+
18
+ | # | Upgrade | Status |
19
+ |---|---------|--------|
20
+ | 1 | Dynamic tile-size autotuner (`triton.autotune` in place of the static `get_head_dim_block_size` / hardcoded `siz_b_LQ=32,siz_b_LKV=32` heuristic) | **Implemented** — `fsp/kernels/autotune.py` |
21
+ | 2 | Fused recurrent+parallel forward kernel (avoid materializing intermediate chunk states to HBM between the two kernel launches) | Planned |
22
+ | 3 | Quantized (fp8/int8) memory state `C` for long-context inference | Planned |
23
+
24
+ ## Why GPU-only
25
+
26
+ TFLA's kernels are written in [Triton](https://github.com/triton-lang/triton)
27
+ and target NVIDIA/AMD GPUs via CUDA/ROCm. There is currently no Triton
28
+ backend for mobile/ARM64-without-GPU environments (e.g. Termux on Android),
29
+ so fsp's kernels can be *written and code-reviewed* in such an environment,
30
+ but must be **benchmarked and numerically validated on an actual CUDA (or
31
+ ROCm) GPU** — a local NVIDIA GPU, or a rented cloud GPU instance.
32
+
33
+ ## Layout
34
+
35
+ ```
36
+ fsp/
37
+ ├── vendor/ # Baseline TFLA, vendored from NX-AI/mlstm_kernels
38
+ │ ├── triton_kernels/ # - the @triton.jit kernels themselves
39
+ │ ├── torch_fw/ # - PyTorch launchers/wrappers (forward pass only, so far)
40
+ │ └── utils/ # - small shared dtype/validation helpers
41
+ ├── kernels/
42
+ │ └── autotune.py # Upgrade #1: autotuned forward-parallel kernel
43
+ └── benchmark/
44
+ └── bench_autotune_vs_baseline.py
45
+ tests/
46
+ └── test_autotune_equivalence.py
47
+ ```
48
+
49
+ > **Scope note:** only the *forward* pass (recurrent + parallel kernels) is
50
+ > vendored so far. Backward-pass kernels
51
+ > (`bw_kernel_parallel_dK/dQ/dV`, `bw_kernel_recurrent`) will be vendored in
52
+ > a later batch once the forward-pass upgrades are validated on real
53
+ > hardware.
54
+
55
+ ## Install
56
+
57
+ ```bash
58
+ pip install torch triton
59
+ pip install -e .
60
+ ```
61
+
62
+ ## Usage
63
+
64
+ ```python
65
+ from fsp.kernels.autotune import mlstm_chunkwise__parallel_fw_Hintra_autotuned
66
+
67
+ matH_out, vecN_out, vecM_out = mlstm_chunkwise__parallel_fw_Hintra_autotuned(
68
+ matQ, matK, matV, vecI, vecF,
69
+ matC_states, vecN_states, scaMinter_states,
70
+ chunk_size=128,
71
+ )
72
+ ```
73
+
74
+ Drop-in replacement for
75
+ `fsp.vendor.torch_fw.fw_parallel.mlstm_chunkwise__parallel_fw_Hintra` — same
76
+ inputs/outputs, same math, but tile sizes are chosen empirically per shape
77
+ instead of by static heuristic.
78
+
79
+ ## Running tests / benchmarks
80
+
81
+ Both require a CUDA-capable GPU:
82
+
83
+ ```bash
84
+ pytest tests/test_autotune_equivalence.py
85
+ python -m fsp.benchmark.bench_autotune_vs_baseline
86
+ ```