initial commit

.gitattributes

@@ -0,0 +1 @@
+*.safetensors filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,6 @@
+.ruff_cache/
+.venv/
+__pycache__/
+*.swp
+*.kate-swp
+*~

LICENSE

@@ -0,0 +1,474 @@
+Lowdown Labs Lovely License 1.0 (LLLL-1.0)
+
+Everything in this repository (the model weights, the configuration, and the code) is
+released by Lowdown Labs under two licenses that apply at the same time. To use this work
+you must comply with BOTH of them. Where a term in one is stricter than the other, the
+stricter term controls. "Lowdown Labs Lovely License 1.0" is a convenience name for this
+exact pair; it is not a new legal instrument.
+
+SPDX-License-Identifier: CC-BY-NC-4.0 AND LicenseRef-Hippocratic-3.0
+
+Commercial licensing. The grant below is non-commercial only. Commercial licenses are
+sold separately by Lowdown Labs on a per-customer basis. To use this work, its weights, or
+its outputs for any commercial purpose, contact Lowdown Labs to purchase a commercial
+license. A commercial license does not remove the Hippocratic ethical-use obligations in
+Part 2; those apply to commercial licensees as well.
+
+==============================================================================
+Part 1 of 2. Creative Commons Attribution-NonCommercial 4.0 International (CC BY-NC 4.0)
+==============================================================================
+
+You may share and adapt this work for non-commercial purposes, with attribution to
+Lowdown Labs. Commercial use is not granted under this license.
+
+  Full legal text:        https://creativecommons.org/licenses/by-nc/4.0/legalcode
+  Plain-language summary:  https://creativecommons.org/licenses/by-nc/4.0/
+  SPDX-License-Identifier: CC-BY-NC-4.0
+
+==============================================================================
+Part 2 of 2. The Hippocratic License 3.0 (ethical use)
+==============================================================================
+
+Module set enabled: bds, cl, eco, extr, ffd, law, media, mil, my, soc, sup, sv, usta.
+Canonical build: https://firstdonoharm.dev/build/?modules=bds,cl,eco,extr,ffd,law,media,mil,my,soc,sup,sv,usta
+
+The verbatim official Hippocratic License 3.0 text for exactly this module set follows,
+between the markers.
+
+--------------------- BEGIN OFFICIAL HIPPOCRATIC LICENSE 3.0 TEXT ---------------------
+
+HIPPOCRATIC LICENSE
+
+Version 3.0, October 2021
+
+https://firstdonoharm.dev/version/3/0/bds-cl-eco-extr-ffd-law-media-mil-my-soc-sup-sv-usta.txt
+
+TERMS AND CONDITIONS
+
+TERMS AND CONDITIONS FOR USE, COPY, MODIFICATION, PREPARATION OF DERIVATIVE
+WORK, REPRODUCTION, AND DISTRIBUTION:
+
+1. DEFINITIONS:
+
+This section defines certain terms used throughout this license agreement.
+
+1.1. “License” means the terms and conditions, as stated herein, for use, copy,
+modification, preparation of derivative work, reproduction, and distribution of
+Software (as defined below).
+
+1.2. “Licensor” means the copyright and/or patent owner or entity authorized by
+the copyright and/or patent owner that is granting the License.
+
+1.3. “Licensee” means the individual or entity exercising permissions granted by
+this License, including the use, copy, modification, preparation of derivative
+work, reproduction, and distribution of Software (as defined below).
+
+1.4. “Software” means any copyrighted work, including but not limited to
+software code, authored by Licensor and made available under this License.
+
+1.5. “Supply Chain” means the sequence of processes involved in the production
+and/or distribution of a commodity, good, or service offered by the Licensee.
+
+1.6. “Supply Chain Impacted Party” or “Supply Chain Impacted Parties” means any
+person(s) directly impacted by any of Licensee’s Supply Chain, including the
+practices of all persons or entities within the Supply Chain prior to a good or
+service reaching the Licensee.
+
+1.7. “Duty of Care” is defined by its use in tort law, delict law, and/or
+similar bodies of law closely related to tort and/or delict law, including
+without limitation, a requirement to act with the watchfulness, attention,
+caution, and prudence that a reasonable person in the same or similar
+circumstances would use towards any Supply Chain Impacted Party.
+
+1.8. “Worker” is defined to include any and all permanent, temporary, and agency
+workers, as well as piece-rate, salaried, hourly paid, legal young (minors),
+part-time, night, and migrant workers.
+
+2. INTELLECTUAL PROPERTY GRANTS:
+
+This section identifies intellectual property rights granted to a Licensee.
+
+2.1. Grant of Copyright License: Subject to the terms and conditions of this
+License, Licensor hereby grants to Licensee a worldwide, non-exclusive,
+no-charge, royalty-free copyright license to use, copy, modify, prepare
+derivative work, reproduce, or distribute the Software, Licensor authored
+modified software, or other work derived from the Software.
+
+2.2. Grant of Patent License: Subject to the terms and conditions of this
+License, Licensor hereby grants Licensee a worldwide, non-exclusive, no-charge,
+royalty-free patent license to make, have made, use, offer to sell, sell,
+import, and otherwise transfer Software.
+
+3. ETHICAL STANDARDS:
+
+This section lists conditions the Licensee must comply with in order to have
+rights under this License.
+
+The rights granted to the Licensee by this License are expressly made subject to
+the Licensee’s ongoing compliance with the following conditions:
+
+ * 3.1. The Licensee SHALL NOT, whether directly or indirectly, through agents
+   or assigns:
+   
+   * 3.1.1. Infringe upon any person’s right to life or security of person,
+     engage in extrajudicial killings, or commit murder, without lawful cause
+     (See Article 3, United Nations Universal Declaration of Human Rights;
+     Article 6, International Covenant on Civil and Political Rights)
+   
+   * 3.1.2. Hold any person in slavery, servitude, or forced labor (See Article
+     4, United Nations Universal Declaration of Human Rights; Article 8,
+     International Covenant on Civil and Political Rights);
+   
+   * 3.1.3. Contribute to the institution of slavery, slave trading, forced
+     labor, or unlawful child labor (See Article 4, United Nations Universal
+     Declaration of Human Rights; Article 8, International Covenant on Civil and
+     Political Rights);
+   
+   * 3.1.4. Torture or subject any person to cruel, inhumane, or degrading
+     treatment or punishment (See Article 5, United Nations Universal
+     Declaration of Human Rights; Article 7, International Covenant on Civil and
+     Political Rights);
+   
+   * 3.1.5. Discriminate on the basis of sex, gender, sexual orientation, race,
+     ethnicity, nationality, religion, caste, age, medical disability or
+     impairment, and/or any other like circumstances (See Article 7, United
+     Nations Universal Declaration of Human Rights; Article 2, International
+     Covenant on Economic, Social and Cultural Rights; Article 26, International
+     Covenant on Civil and Political Rights);
+   
+   * 3.1.6. Prevent any person from exercising his/her/their right to seek an
+     effective remedy by a competent court or national tribunal (including
+     domestic judicial systems, international courts, arbitration bodies, and
+     other adjudicating bodies) for actions violating the fundamental rights
+     granted to him/her/them by applicable constitutions, applicable laws, or by
+     this License (See Article 8, United Nations Universal Declaration of Human
+     Rights; Articles 9 and 14, International Covenant on Civil and Political
+     Rights);
+   
+   * 3.1.7. Subject any person to arbitrary arrest, detention, or exile (See
+     Article 9, United Nations Universal Declaration of Human Rights; Article 9,
+     International Covenant on Civil and Political Rights);
+   
+   * 3.1.8. Subject any person to arbitrary interference with a person’s
+     privacy, family, home, or correspondence without the express written
+     consent of the person (See Article 12, United Nations Universal Declaration
+     of Human Rights; Article 17, International Covenant on Civil and Political
+     Rights);
+   
+   * 3.1.9. Arbitrarily deprive any person of his/her/their property (See
+     Article 17, United Nations Universal Declaration of Human Rights);
+   
+   * 3.1.10. Forcibly remove indigenous peoples from their lands or territories
+     or take any action with the aim or effect of dispossessing indigenous
+     peoples from their lands, territories, or resources, including without
+     limitation the intellectual property or traditional knowledge of indigenous
+     peoples, without the free, prior, and informed consent of indigenous
+     peoples concerned (See Articles 8 and 10, United Nations Declaration on the
+     Rights of Indigenous Peoples);
+    * 3.1.11. Fossil Fuel Divestment: Be an individual or entity, or a
+      representative, agent, affiliate, successor, attorney, or assign of an
+      individual or entity, on the FFI Solutions Carbon Underground 200 list
+      [https://www.ffisolutions.com/research-analytics-index-solutions/research-screening/the-carbon-underground-200/?cn-reloaded=1];
+   
+    * 3.1.12. Ecocide: Commit ecocide:
+      
+      * 3.1.12.1. For the purpose of this section, “ecocide” means unlawful or
+        wanton acts committed with knowledge that there is a substantial
+        likelihood of severe and either widespread or long-term damage to the
+        environment being caused by those acts;
+      
+      * 3.1.12.2. For the purpose of further defining ecocide and the terms
+        contained in the previous paragraph:
+        
+        * 3.1.12.2.1. “Wanton” means with reckless disregard for damage which
+          would be clearly excessive in relation to the social and economic
+          benefits anticipated;
+        
+        * 3.1.12.2.2. “Severe” means damage which involves very serious adverse
+          changes, disruption, or harm to any element of the environment,
+          including grave impacts on human life or natural, cultural, or
+          economic resources;
+        
+        * 3.1.12.2.3. “Widespread” means damage which extends beyond a limited
+          geographic area, crosses state boundaries, or is suffered by an entire
+          ecosystem or species or a large number of human beings;
+        
+        * 3.1.12.2.4. “Long-term” means damage which is irreversible or which
+          cannot be redressed through natural recovery within a reasonable
+          period of time; and
+        
+        * 3.1.12.2.5. “Environment” means the earth, its biosphere, cryosphere,
+          lithosphere, hydrosphere, and atmosphere, as well as outer space
+          
+          (See Section II, Independent Expert Panel for the Legal Definition of
+          Ecocide, Stop Ecocide Foundation and the Promise Institute for Human
+          Rights at UCLA School of Law, June 2021);
+   
+    * 3.1.13. Extractive Industries: Be an individual or entity, or a
+      representative, agent, affiliate, successor, attorney, or assign of an
+      individual or entity, that engages in fossil fuel or mineral exploration,
+      extraction, development, or sale;
+   
+    * 3.1.14. Boycott / Divestment / Sanctions: Be an individual or entity, or a
+      representative, agent, affiliate, successor, attorney, or assign of an
+      individual or entity, identified by the Boycott, Divestment, Sanctions
+      (“BDS”) movement on its website (https://bdsmovement.net/
+      [https://bdsmovement.net/] and
+      https://bdsmovement.net/get-involved/what-to-boycott
+      [https://bdsmovement.net/get-involved/what-to-boycott]) as a target for
+      boycott;
+   
+    * 3.1.15. Myanmar: Be an individual or entity that:
+      
+      * 3.1.15.1. engages in any commercial transactions with the
+        Myanmar/Burmese military junta; or
+      
+      * 3.1.15.2. is a representative, agent, affiliate, successor, attorney, or
+        assign of the Myanmar/Burmese government;
+   
+    * 3.1.16. US Tariff Act: Be an individual or entity:
+      
+      * 3.1.16.1. which U.S. Customs and Border Protection (CBP) has currently
+        issued a Withhold Release Order (WRO) or finding against based on
+        reasonable suspicion of forced labor; or
+      
+      * 3.1.16.2. that is a representative, agent, affiliate, successor,
+        attorney, or assign of an individual or entity that does business with
+        an individual or entity which currently has a WRO or finding from CBP
+        issued against it based on reasonable suspicion of forced labor;
+   
+    * 3.1.17. Mass Surveillance: Be a government agency or multinational
+      corporation, or a representative, agent, affiliate, successor, attorney,
+      or assign of a government or multinational corporation, which participates
+      in mass surveillance programs;
+   
+    * 3.1.18. Military Activities: Be an entity or a representative, agent,
+      affiliate, successor, attorney, or assign of an entity which conducts
+      military activities;
+   
+    * 3.1.19. Law Enforcement: Be an individual or entity, or a representative,
+      agent, affiliate, successor, attorney, or assign of an individual or
+      entity, that provides good or services to, or otherwise enters into any
+      commercial contracts with, any local, state, or federal law enforcement
+      agency;
+   
+    * 3.1.20. Media: Be an individual or entity, or a representative, agent,
+      affiliate, successor, attorney, or assign of an individual or entity, that
+      broadcasts messages promoting killing, torture, or other forms of extreme
+      violence;
+   
+   * 3.1.21. Interfere with Workers’ free exercise of the right to organize and
+     associate (See Article 20, United Nations Universal Declaration of Human
+     Rights; C087 - Freedom of Association and Protection of the Right to
+     Organise Convention, 1948 (No. 87), International Labour Organization;
+     Article 8, International Covenant on Economic, Social and Cultural Rights);
+     and
+   
+   * 3.1.22. Harm the environment in a manner inconsistent with local, state,
+     national, or international law.
+
+ * 3.2. The Licensee SHALL:
+   
+    * 3.2.1. Social Auditing: Only use social auditing mechanisms that adhere to
+      Worker-Driven Social Responsibility Network’s Statement of Principles
+      (https://wsr-network.org/what-is-wsr/statement-of-principles/
+      [https://wsr-network.org/what-is-wsr/statement-of-principles/]) over
+      traditional social auditing mechanisms, to the extent the Licensee uses
+      any social auditing mechanisms at all;
+   
+    * 3.2.2. Supply Chain: Provide clear, accessible supply chain data to the
+      public in accordance with the following conditions:
+      
+      * 3.2.2.1. All data will be on Licensee’s website and/or, to the extent
+        Licensee is a representative, agent, affiliate, successor, attorney,
+        subsidiary, or assign, on Licensee’s principal’s or parent’s website or
+        some other online platform accessible to the public via an internet
+        search on a common internet search engine; and
+      
+      * 3.2.2.2. Data published will include, where applicable, manufacturers,
+        top tier suppliers, subcontractors, cooperatives, component parts
+        producers, and farms;
+   
+   * 3.2.3. Provide equal pay for equal work where the performance of such work
+     requires equal skill, effort, and responsibility, and which are performed
+     under similar working conditions, except where such payment is made
+     pursuant to:
+     
+     * 3.2.3.1. A seniority system;
+     
+     * 3.2.3.2. A merit system;
+     
+     * 3.2.3.3. A system which measures earnings by quantity or quality of
+       production; or
+     
+     * 3.2.3.4. A differential based on any other factor other than sex, gender,
+       sexual orientation, race, ethnicity, nationality, religion, caste, age,
+       medical disability or impairment, and/or any other like circumstances
+       (See 29 U.S.C.A. § 206(d)(1); Article 23, United Nations Universal
+       Declaration of Human Rights; Article 7, International Covenant on
+       Economic, Social and Cultural Rights; Article 26, International Covenant
+       on Civil and Political Rights); and
+   
+   * 3.2.4. Allow for reasonable limitation of working hours and periodic
+     holidays with pay (See Article 24, United Nations Universal Declaration of
+     Human Rights; Article 7, International Covenant on Economic, Social and
+     Cultural Rights).
+
+4. SUPPLY CHAIN IMPACTED PARTIES:
+
+This section identifies additional individuals or entities that a Licensee could
+harm as a result of violating the Ethical Standards section, the condition that
+the Licensee must voluntarily accept a Duty of Care for those individuals or
+entities, and the right to a private right of action that those individuals or
+entities possess as a result of violations of the Ethical Standards section.
+
+4.1. In addition to the above Ethical Standards, Licensee voluntarily accepts a
+Duty of Care for Supply Chain Impacted Parties of this License, including
+individuals and communities impacted by violations of the Ethical Standards. The
+Duty of Care is breached when a provision within the Ethical Standards section
+is violated by a Licensee, one of its successors or assigns, or by an individual
+or entity that exists within the Supply Chain prior to a good or service
+reaching the Licensee.
+
+4.2. Breaches of the Duty of Care, as stated within this section, shall create a
+private right of action, allowing any Supply Chain Impacted Party harmed by the
+Licensee to take legal action against the Licensee in accordance with applicable
+negligence laws, whether they be in tort law, delict law, and/or similar bodies
+of law closely related to tort and/or delict law, regardless if Licensee is
+directly responsible for the harms suffered by a Supply Chain Impacted Party.
+Nothing in this section shall be interpreted to include acts committed by
+individuals outside of the scope of his/her/their employment.
+
+5. NOTICE: This section explains when a Licensee must notify others of the
+License.
+
+5.1. Distribution of Notice: Licensee must ensure that everyone who receives a
+copy of or uses any part of Software from Licensee, with or without changes,
+also receives the License and the copyright notice included with Software (and
+if included by the Licensor, patent, trademark, and attribution notice).
+Licensee must ensure that License is prominently displayed so that any
+individual or entity seeking to download, copy, use, or otherwise receive any
+part of Software from Licensee is notified of this License and its terms and
+conditions. Licensee must cause any modified versions of the Software to carry
+prominent notices stating that Licensee changed the Software.
+
+5.2. Modified Software: Licensee is free to create modifications of the Software
+and distribute only the modified portion created by Licensee, however, any
+derivative work stemming from the Software or its code must be distributed
+pursuant to this License, including this Notice provision.
+
+5.3. Recipients as Licensees: Any individual or entity that uses, copies,
+modifies, reproduces, distributes, or prepares derivative work based upon the
+Software, all or part of the Software’s code, or a derivative work developed by
+using the Software, including a portion of its code, is a Licensee as defined
+above and is subject to the terms and conditions of this License.
+
+6. REPRESENTATIONS AND WARRANTIES:
+
+6.1. Disclaimer of Warranty: TO THE FULL EXTENT ALLOWED BY LAW, THIS SOFTWARE
+COMES “AS IS,” WITHOUT ANY WARRANTY, EXPRESS OR IMPLIED, AND LICENSOR SHALL NOT
+BE LIABLE TO ANY PERSON OR ENTITY FOR ANY DAMAGES OR OTHER LIABILITY ARISING
+FROM, OUT OF, OR IN CONNECTION WITH THE SOFTWARE OR THIS LICENSE, UNDER ANY
+LEGAL CLAIM.
+
+6.2. Limitation of Liability: LICENSEE SHALL HOLD LICENSOR HARMLESS AGAINST ANY
+AND ALL CLAIMS, DEBTS, DUES, LIABILITIES, LIENS, CAUSES OF ACTION, DEMANDS,
+OBLIGATIONS, DISPUTES, DAMAGES, LOSSES, EXPENSES, ATTORNEYS’ FEES, COSTS,
+LIABILITIES, AND ALL OTHER CLAIMS OF EVERY KIND AND NATURE WHATSOEVER, WHETHER
+KNOWN OR UNKNOWN, ANTICIPATED OR UNANTICIPATED, FORESEEN OR UNFORESEEN, ACCRUED
+OR UNACCRUED, DISCLOSED OR UNDISCLOSED, ARISING OUT OF OR RELATING TO LICENSEE’S
+USE OF THE SOFTWARE. NOTHING IN THIS SECTION SHOULD BE INTERPRETED TO REQUIRE
+LICENSEE TO INDEMNIFY LICENSOR, NOR REQUIRE LICENSOR TO INDEMNIFY LICENSEE.
+
+7. TERMINATION
+
+7.1. Violations of Ethical Standards or Breaching Duty of Care: If Licensee
+violates the Ethical Standards section or Licensee, or any other person or
+entity within the Supply Chain prior to a good or service reaching the Licensee,
+breaches its Duty of Care to Supply Chain Impacted Parties, Licensee must remedy
+the violation or harm caused by Licensee within 30 days of being notified of the
+violation or harm. If Licensee fails to remedy the violation or harm within 30
+days, all rights in the Software granted to Licensee by License will be null and
+void as between Licensor and Licensee.
+
+7.2. Failure of Notice: If any person or entity notifies Licensee in writing
+that Licensee has not complied with the Notice section of this License, Licensee
+can keep this License by taking all practical steps to comply within 30 days
+after the notice of noncompliance. If Licensee does not do so, Licensee’s
+License (and all rights licensed hereunder) will end immediately.
+
+7.3. Judicial Findings: In the event Licensee is found by a civil, criminal,
+administrative, or other court of competent jurisdiction, or some other
+adjudicating body with legal authority, to have committed actions which are in
+violation of the Ethical Standards or Supply Chain Impacted Party sections of
+this License, all rights granted to Licensee by this License will terminate
+immediately.
+
+7.4. Patent Litigation: If Licensee institutes patent litigation against any
+entity (including a cross-claim or counterclaim in a suit) alleging that the
+Software, all or part of the Software’s code, or a derivative work developed
+using the Software, including a portion of its code, constitutes direct or
+contributory patent infringement, then any patent license, along with all other
+rights, granted to Licensee under this License will terminate as of the date
+such litigation is filed.
+
+7.5. Additional Remedies: Termination of the License by failing to remedy harms
+in no way prevents Licensor or Supply Chain Impacted Party from seeking
+appropriate remedies at law or in equity.
+
+8. MISCELLANEOUS:
+
+8.1. Conditions: Sections 3, 4.1, 5.1, 5.2, 7.1, 7.2, 7.3, and 7.4 are
+conditions of the rights granted to Licensee in the License.
+
+8.2. Equitable Relief: Licensor and any Supply Chain Impacted Party shall be
+entitled to equitable relief, including injunctive relief or specific
+performance of the terms hereof, in addition to any other remedy to which they
+are entitled at law or in equity.
+
+8.3. Copyleft: Modified software, source code, or other derivative work must be
+licensed, in its entirety, under the exact same conditions as this License.
+
+8.4. Severability: If any term or provision of this License is determined to be
+invalid, illegal, or unenforceable by a court of competent jurisdiction, any
+such determination of invalidity, illegality, or unenforceability shall not
+affect any other term or provision of this License or invalidate or render
+unenforceable such term or provision in any other jurisdiction. If the
+determination of invalidity, illegality, or unenforceability by a court of
+competent jurisdiction pertains to the terms or provisions contained in the
+Ethical Standards section of this License, all rights in the Software granted to
+Licensee shall be deemed null and void as between Licensor and Licensee.
+
+8.5. Section Titles: Section titles are solely written for organizational
+purposes and should not be used to interpret the language within each section.
+
+8.6. Citations: Citations are solely written to provide context for the source
+of the provisions in the Ethical Standards.
+
+8.7. Section Summaries: Some sections have a brief italicized description which
+is provided for the sole purpose of briefly describing the section and should
+not be used to interpret the terms of the License.
+
+8.8. Entire License: This is the entire License between the Licensor and
+Licensee with respect to the claims released herein and that the consideration
+stated herein is the only consideration or compensation to be paid or exchanged
+between them for this License. This License cannot be modified or amended except
+in a writing signed by Licensor and Licensee.
+
+8.9. Successors and Assigns: This License shall be binding upon and inure to the
+benefit of the Licensor’s and Licensee’s respective heirs, successors, and
+assigns.
+---------------------- END OFFICIAL HIPPOCRATIC LICENSE 3.0 TEXT ----------------------
+
+==============================================================================
+Attribution, commercial use, and warranty
+==============================================================================
+
+Attribution: cite this work as described in the model card (README.md), section
+"How to cite". Attribution to Lowdown Labs is required under CC BY-NC 4.0.
+
+Commercial use: CC BY-NC 4.0 does not grant commercial rights. Commercial licenses are
+sold by Lowdown Labs; contact Lowdown Labs to purchase one.
+
+No warranty: this work is provided as is, without warranty of any kind. See the model card
+for the intended use, the evaluated conditions, and the known limitations.

README.md

@@ -0,0 +1,285 @@
+---
+license: other
+license_name: lowdown-labs-lovely-license-1.0
+license_link: LICENSE
+tags:
+  - fela
+  - fourier-neural-operator
+  - fno
+  - gated-deltanet
+  - cpu
+  - on-device
+  - code
+  - autocomplete
+  - fill-in-the-middle
+  - constant-memory
+library_name: transformers
+pipeline_tag: text-generation
+language:
+  - code
+model-index:
+- name: fela-autocomplete
+  results:
+  - task:
+      type: text-generation
+      name: Fill in the middle single line infilling
+    dataset:
+      type: humaneval-singlelineinfilling
+      name: HumanEval SingleLineInfilling
+    metrics:
+    - type: pass@1
+      value: 54.6
+      name: FIM pass@1
+---
+
+# FELA LLM 1.5: On device code autocomplete (fill in the middle)
+
+FELA Autocomplete is a code completion model. You give it the code before your cursor, and
+optionally more codebase context, and it writes the piece that goes in between. This is
+the same fill in the middle idea that powers the gray autocomplete text in a modern editor.
+It is a 1.79 billion parameter model that runs on a plain CPU with no GPU, so it can sit
+inside a local IDE helper or an on premises code tool and complete code without sending your
+source to the cloud.
+
+It was warm started from Qwen2.5-Coder 1.5B and clean room distilled, then annealed and
+supervised fine tuned for fill in the middle. Our key niche insight is to provide for inline single line
+completion, not full text generation. We feel this a great compromise to have a speedy
+development experience while still having a product that doesn't contribute to cognitive
+atrophy and feels better to use.
+
+# Headline result: fill in the middle
+
+We wanted to apply the FELA methodology as an LLM. Given a tight training budget while we
+study how to scale this new architecture, we thought a smaller model would be a great
+autocomplete tool to demonstrate the power of efficient computing.
+
+We measure functional pass@1 with the HumanEval-SingleLineInfilling methodology, so
+a completion counts only when it makes the code run correctly, not when it matches the
+original character for character.
+
+| Stage | FIM Pass@1 | FIM EM | FIM Edit Sim | MBPP | HumanEval |
+|---|---|---|---|---|---|
+| Base (15B token distill) | 49.6% | 42.0 | 73.6 | 13.0 | 11.0 |
+| Plus fill in the middle anneal | 54.4% | 47.4 | 77.6 | 16.0 | 9.8 |
+| Plus fill in the middle SFT (final) | **54.6%** | 46.8 | **78.4** | **18.0** | 11.0 |
+
+Over half of single line completions functionally work. We lead with FIM pass@1 because that
+is the product metric. The Rust server auto formats the output in microseconds, so we optimize
+for semantic correctness (pass@1 and edit similarity), not whitespace exact match.
+
+MBPP and HumanEval are shown for context only. They ask the model to write a whole function
+from a description, which is the wrong lens for an autocomplete model and low by design at this
+size. We provide it to compare with other language models at size ~1.5B. Do not read the HumanEval number as the quality of the product; read the FIM pass@1.
+
+# How it got here
+
+1. Base: clean room distill from Qwen2.5-Coder 1.5B to 15 billion tokens, on a permissive code
+   mix with fill in the middle in the objective.
+2. Fill in the middle anneal (the big lever): a further billion tokens with a line aware fill
+   in the middle objective (mask whole lines, matching how the eval scores lines) across
+   Python, JavaScript, TypeScript, Java, Go, Rust, and C plus plus. This train and eval
+   alignment gave the largest jump.
+3. Fill in the middle SFT: a further 1.2 billion tokens with the native Qwen fill in the middle
+   tokens and the loss on the middle span. This consolidated the gains and lifted general code
+   (MBPP 16 to 18).
+
+# What goes in, what comes out
+
+- Input: your code as text, split into a prefix (before the cursor) and an optional suffix
+  (after the cursor). The text is turned into tokens with the Qwen2.5-Coder tokenizer, which
+  carries the fill in the middle markers `<|fim_prefix|>`, `<|fim_suffix|>`, and
+  `<|fim_middle|>`.
+- Output: the code that belongs between the prefix and the suffix. By default it writes one
+  line and stops, which is what an inline autocomplete wants. With no suffix it simply
+  continues the prefix.
+- In plain terms: you type `import numpy as ` and it finishes the line with ` np`; you give it
+  the body of `def add(a, b):` with the call `add(2, 3)` waiting below it, and it fills in
+  ` return a + b`.
+
+# Why we built it this way
+
+The model has no growing cache of past tokens, which is the usual reason memory creeps up as
+context gets longer. Its sequence mixers are Fourier Neural Operators (a filter the model
+learns and applies in the frequency domain) for global context, Gated-DeltaNet recall layers
+that overwrite an old memory when a new one arrives (so keys do not interfere), and a Landmark
+routing layer that lets a token look back at summaries of earlier chunks. None of these keep a
+per token attention cache, so the working memory stays small and fixed no matter how much code
+is in the window. That is what lets it decode on a low power CPU next to you, with the source
+never leaving the machine.
+
+The full stack is 28 layers at width 1536, in the repeating pattern of three Fourier mixers
+then one Gated-DeltaNet recall layer, with a Landmark routing layer every seventh layer.
+
+# Performance
+
+The model runs on CPU with no GPU. Two paths run these weights:
+
+| Path | Format | Size on disk | Single core decode |
+|---|---|---|---|
+| Reference (this repo, `modeling.py`) | bf16 | 3.58 GB | about 4 tokens per second |
+| Production (FELA server, separate Rust runtime) | int8 | about 2.08 GB | about 22 tokens per second |
+
+The `modeling.py` loader here is the plain, readable reference path in pure PyTorch. It loads
+the weights, runs a real forward, and decodes on CPU, which is enough for the quickstart and for
+verifying the model. It is a correctness path, not the fast path, so do not read its speed as the
+product speed. Loading `model_int8.safetensors` in Python is smaller in memory but not faster,
+because the pure torch int8 path dequantizes the weights on the fly.
+
+The fast path is the CPU native FELA server, a separate Rust runtime that runs these weights as
+int8 with a constant state streaming decode and no KV cache. On one core it decodes at about 22
+tokens per second, roughly 20 times the plain f32 path, measured at this model's shape with a
+steady state O(1) decode step. Single stream decode is bound by memory bandwidth, since it
+streams the full int8 weight set once per token, so adding cores does not speed up one stream
+(measured flat from one to eight cores); extra cores serve more streams at once and speed up the
+prefill. The server's Rust int8 forward is checked against the PyTorch model with a golden test,
+so the fast serving path produces the same model, not a lookalike.
+
+# What to expect
+
+Some real completions from the final model:
+
+- `import numpy as ` gives ` np`
+- prefix `def add(a, b):` with the call `add(2, 3)` in the suffix gives ` return a + b`
+- prefix `def is_even(n):\n    return n % 2 == ` gives `0`
+
+Where it earns its keep is single line fill in the middle on the patterns developers write all
+day - and it can be made more useful with prefilled repository context.
+It is a 1.5B class model, so it is not a whole function generator or a reasoning model:
+multi line blocks, novel algorithms, and anything that needs deep reasoning about your specific
+code are outside what it is built for. Read completions before you use them.
+
+# How to run it
+
+It runs on CPU; no GPU is required. See `quickstart/` for a runnable example. The short
+version, using the self contained `modeling.py`, `model_cpu_gpt2.py`, and `config.json` in
+this repo:
+
+```python
+from modeling import load_model
+m = load_model("/path/to/this/repo")
+
+# Fill in the middle: code before and after the cursor -> the middle
+r = m.complete("def add(a, b):\n    ", suffix="\nresult = add(2, 3)\n")
+print(r["middle"])          # e.g. " return a + b"
+
+# Plain autocomplete: continue a single line
+print(m.complete("import numpy as ")["middle"])   # e.g. " np"
+
+# Load the smaller int8 weights instead (same completions, smaller on disk)
+m8 = load_model("/path/to/this/repo", quant="int8")
+```
+
+It also loads with transformers directly. The code is custom, so pass `trust_remote_code=True`:
+
+```python
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+model = AutoModelForCausalLM.from_pretrained("lowdown-labs/fela-autocomplete", trust_remote_code=True, dtype=torch.bfloat16)
+tok = AutoTokenizer.from_pretrained("lowdown-labs/fela-autocomplete", trust_remote_code=True)
+prompt = "<|fim_prefix|>def add(a, b):\n    <|fim_suffix|>\nresult = add(2, 3)\n<|fim_middle|>"
+enc = tok(prompt, return_tensors="pt")
+out = model.generate(**enc, max_new_tokens=8, use_cache=False)
+print(tok.decode(out[0, enc["input_ids"].shape[1]:], skip_special_tokens=True))
+```
+
+The loader reads the bf16 `model.safetensors` by default. Pass `quant="int8"` to load the
+smaller `model_int8.safetensors` instead; it dequantizes the linear weights with their per
+channel scales and keeps the Fourier filters in bf16, giving the same completions. Pass
+`quant="auto"` to prefer bf16 when present and fall back to int8.
+
+`complete` is greedy by default, so a skeptic gets the same real output twice. For an
+interactive playground, see the Hugging Face Space in `space/`. For production serving, use
+the CPU native FELA server (`https://github.com/Lowdown-Labs/fela_server`).
+
+## Loading, exports, and verification
+
+- Weights ship as `model.safetensors` in bf16 (no pickle code execution risk). `config.json`
+  holds the architecture hyperparameters. `tokenizer.json` is the Qwen2.5-Coder tokenizer with
+  the fill in the middle tokens. The architecture is defined in the self contained
+  `model_cpu_gpt2.py`, and the pure torch CPU paths for the Gated-DeltaNet and Landmark layers
+  are in `cpu_delta.py`, `cpu_landmark.py`, `cpu_swa.py`, and `cpu_patch.py`, which ship beside
+  it. `modeling.py` wires them into a loader with a `complete` fill in the middle entry point.
+- `verify.py` builds the model, loads the weights, runs a fixed input forward, checks the
+  output shape and the vocabulary size, confirms the logits are finite, and compares the top
+  predicted tokens against a captured reference.
+
+## Formats
+
+- bf16: the reference weights and the default load path in this repo (`model.safetensors`,
+  3.58 GB).
+- int8: the smaller on device format (`model_int8.safetensors`, about 2.08 GB), weight only int8
+  per channel on the large linear layers, with the Fourier filters kept in bf16. Load it in
+  Python with `load_model(dir, quant="int8")`, or serve it with the FELA server. The int8
+  forward matches the bf16 model by the golden test.
+
+## Model details
+
+- Parameters: 1,788,379,536 (about 1.79 billion).
+- Layers: 28, width 1536, 12 heads, feed forward hidden 8960.
+- Sequence mixers: Fourier Neural Operator (512 modes), Gated-DeltaNet recall, and Landmark
+  global routing, in the SSSL pattern (three Fourier mixers then one recall layer), with a
+  Landmark layer every seventh layer.
+- Tokenizer: Qwen2.5-Coder (vocabulary 151936), which carries the fill in the middle tokens.
+- Warm start: Qwen2.5-Coder 1.5B, then clean room distilled, annealed, and fill in the middle
+  supervised fine tuned.
+
+# Intended use, limitations, and safety
+
+What it is for: inline code autocomplete and fill in the middle inside a local or on premises
+developer tool, where running on CPU with no GPU and keeping the source on the machine matter.
+It is a research preview for that use.
+
+What it is not for: it is not a whole function generator, not a chat model, and not a reasoning
+model. Do not rely on its completions for correctness or security without reading them, and do
+not ship it into a product without your own evaluation. It is a 1.5B class model and will be
+wrong on anything past common single line patterns.
+
+Evaluated conditions and known failure modes:
+
+- Strong on single line, high frequency fill in the middle; weaker on multi line blocks and
+  novel logic.
+- The reference `modeling.py` path is bf16 and slow; the int8 FELA server is the fast path.
+- No safety or license filtering of generated code is claimed. Review generated code before you
+  use it.
+
+Privacy: the model runs on the device or on your own CPU, so your source does not have to leave
+the machine and there is no per call API. That is the on device privacy angle.
+
+# Model family
+
+This is part of the FELA family from Lowdown Labs: one Fourier Neural Operator architecture
+across many modalities, all CPU native and subquadratic. This repo is pushed as
+`lowdown-labs/FELA-autocomplete`. Sibling repos are trained independently per modality and
+share no weights, so none carries a `base_model` link.
+
+# Acknowledgements and references
+
+- Fourier Neural Operator: Li, Z., et al. (2021). Fourier Neural Operator for Parametric
+  Partial Differential Equations. ICLR. https://arxiv.org/abs/2010.08895
+- Gated Linear Attention: Yang, S., et al. (2023). Gated Linear Attention Transformers with
+  Hardware-Efficient Training. https://arxiv.org/abs/2312.06635
+- Gated DeltaNet: Yang, S., et al. (2024). Gated Delta Networks: Improving Mamba2 with Delta
+  Rule. https://arxiv.org/abs/2412.06464
+- Landmark Attention: Mohtashami, A., Jaggi, M. (2023). Landmark Attention: Random-Access
+  Infinite Context Length for Transformers. https://arxiv.org/abs/2305.16300
+- Qwen2.5-Coder: Hui, B., et al. (2024). Qwen2.5-Coder Technical Report.
+  https://arxiv.org/abs/2409.12186
+- PyTorch: Paszke, A., et al. (2019). NeurIPS. https://arxiv.org/abs/1912.01703
+- flash-linear-attention: the reference kernels for the Gated-DeltaNet layer.
+  https://github.com/fla-org/flash-linear-attention
+
+# How to cite
+
+```bibtex
+@misc{lowdownlabs_fela_llm,
+  title  = {FELA LLM 1.5: An on device Fourier Neural Operator code autocomplete model},
+  author = {Lowdown Labs},
+  year   = {2026},
+  note   = {Model card}
+}
+```
+
+# License
+
+Released under the Lowdown Labs Lovely License 1.0 (CC BY-NC 4.0 plus Hippocratic License 3.0). See LICENSE. For most LL models, a commercial license may be available; contact Lowdown Labs.

config.json

@@ -0,0 +1,40 @@
+{
+  "model_type": "fela",
+  "library_name": "pytorch",
+  "arch": "CPUGPT (FELA LM)",
+  "edition": "FELA LLM 1.5 (code, 1.5B class)",
+  "note": "FELA code language model for fill in the middle autocomplete. 1.79B parameters. 28 layers of FNO spectral token mixers plus Gated-DeltaNet recall plus Landmark global routing (SSSL pattern). Final fill in the middle SFT model.",
+  "quant": "bf16",
+  "vocab_size": 151936,
+  "n_layer": 28,
+  "n_embd": 1536,
+  "n_head": 12,
+  "ffn_hidden": 8960,
+  "fno_modes": 512,
+  "gla_chunk": 256,
+  "layer_pattern": "SSSL",
+  "gla_delta": true,
+  "landmark_layer_every": 7,
+  "landmark_chunk": 32,
+  "landmark_max": 64,
+  "attn_layer_every": 0,
+  "seq_len": 2048,
+  "n_params": 1788379536,
+  "tokenizer": "Qwen2.5-Coder (tokenizer.json, vocab 151936, carries the FIM tokens <|fim_prefix|> / <|fim_suffix|> / <|fim_middle|>)",
+  "fim_tokens": {
+    "prefix": "<|fim_prefix|>",
+    "suffix": "<|fim_suffix|>",
+    "middle": "<|fim_middle|>"
+  },
+  "input_desc": "token ids from the Qwen2.5-Coder tokenizer. For fill in the middle: <|fim_prefix|> P <|fim_suffix|> S <|fim_middle|>, and the model completes the middle. For plain continuation: the prompt token ids.",
+  "serving": "Runs on a plain CPU with no GPU via the FELA server (int8, constant memory streaming decode, no KV cache). The bf16 model.safetensors is the reference; model_int8.safetensors is the int8 serving export.",
+  "int8_variant": "model_int8.safetensors",
+  "int8_scheme": "weight only int8 per channel symmetric on linear layers; scales fp32; kept tensors bf16; keys prefixed int8 / scale / keep",
+  "architectures": [
+    "FelaForCausalLM"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_fela.FelaConfig",
+    "AutoModelForCausalLM": "modeling_fela.FelaForCausalLM"
+  }
+}

configuration_fela.py

@@ -0,0 +1,43 @@
+from transformers import PretrainedConfig
+
+
+class FelaConfig(PretrainedConfig):
+    model_type = "fela"
+
+    def __init__(
+        self,
+        vocab_size=151936,
+        seq_len=2048,
+        n_layer=28,
+        n_embd=1536,
+        n_head=12,
+        ffn_hidden=8960,
+        layer_pattern="SSSL",
+        gla_delta=True,
+        fno_modes=512,
+        gla_chunk=256,
+        landmark_layer_every=7,
+        landmark_chunk=32,
+        landmark_max=64,
+        attn_layer_every=0,
+        tie_word_embeddings=False,
+        use_cache=False,
+        **kwargs,
+    ):
+        self.vocab_size = vocab_size
+        self.seq_len = seq_len
+        self.n_layer = n_layer
+        self.n_embd = n_embd
+        self.n_head = n_head
+        self.ffn_hidden = ffn_hidden
+        self.layer_pattern = layer_pattern
+        self.gla_delta = gla_delta
+        self.fno_modes = fno_modes
+        self.gla_chunk = gla_chunk
+        self.landmark_layer_every = landmark_layer_every
+        self.landmark_chunk = landmark_chunk
+        self.landmark_max = landmark_max
+        self.attn_layer_every = attn_layer_every
+        self.num_hidden_layers = n_layer
+        self.use_cache = use_cache
+        super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)

cpu_delta.py

@@ -0,0 +1,256 @@
+from __future__ import annotations
+import torch
+import torch.nn.functional as F
+
+
+def gdn_recurrent(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    beta: torch.Tensor,
+    g: torch.Tensor,
+    scale: float,
+    initial_state: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    q, k, v, beta, g = (
+        t.transpose(1, 2).contiguous().to(torch.float32) for t in (q, k, v, beta, g)
+    )
+    B, H, T, K = q.shape
+    V = v.shape[-1]
+    o = torch.zeros(B, H, T, V, dtype=torch.float32, device=q.device)
+    if initial_state is None:
+        h = torch.zeros(B, H, K, V, dtype=torch.float32, device=q.device)
+    else:
+        h = initial_state.to(torch.float32).clone()
+    q = q * scale
+    for i in range(T):
+        b_q = q[:, :, i]
+        b_k = k[:, :, i]
+        b_v = v[:, :, i].clone()
+        h = h * g[:, :, i].exp()[..., None, None]
+        b_v = b_v - (h * b_k[..., None]).sum(-2)
+        b_v = b_v * beta[:, :, i][..., None]
+        h = h + b_k.unsqueeze(-1) * b_v.unsqueeze(-2)
+        o[:, :, i] = torch.einsum("bhd,bhdm->bhm", b_q, h)
+    return (o.transpose(1, 2).contiguous(), h)
+
+
+def gdn_chunk_recurrent(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    beta: torch.Tensor,
+    g: torch.Tensor,
+    scale: float,
+    initial_state: torch.Tensor | None = None,
+    chunk_size: int = 64,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    bt = chunk_size
+    q, k, v, beta, g = (
+        t.transpose(1, 2).contiguous().to(torch.float32) for t in (q, k, v, beta, g)
+    )
+    B, H, T, K = q.shape
+    Vd = v.shape[-1]
+    pad = (bt - T % bt) % bt
+    if pad:
+        q = F.pad(q, (0, 0, 0, pad))
+        k = F.pad(k, (0, 0, 0, pad))
+        v = F.pad(v, (0, 0, 0, pad))
+        beta = F.pad(beta, (0, pad))
+        g = F.pad(g, (0, pad))
+    L = q.shape[2]
+    n = L // bt
+    q = q * scale
+    v = v * beta[..., None]
+    k_beta = k * beta[..., None]
+
+    def _chunks(x):
+        return x.reshape(B, H, n, bt, x.shape[-1])
+
+    q, k, v, k_beta = (_chunks(q), _chunks(k), _chunks(v), _chunks(k_beta))
+    decay = g.reshape(B, H, n, bt).cumsum(-1)
+    decay_exp = decay.exp()[..., None]
+    l_mask = (decay.unsqueeze(-1) - decay.unsqueeze(-2)).tril().exp().tril()
+    eye_mask = torch.triu(
+        torch.ones(bt, bt, dtype=torch.bool, device=q.device), diagonal=0
+    )
+    attn = -(k_beta @ k.transpose(-1, -2) * l_mask).masked_fill(eye_mask, 0)
+    for i in range(1, bt):
+        attn[..., i, :i] = attn[..., i, :i].clone() + (
+            attn[..., i, :i, None].clone() * attn[..., :i, :i].clone()
+        ).sum(-2)
+    attn = attn + torch.eye(bt, dtype=torch.float32, device=q.device)
+    v = attn @ v
+    k_cumdecay = attn @ (k_beta * decay_exp)
+    S = q.new_zeros(B, H, K, Vd)
+    if initial_state is not None:
+        S = initial_state.to(torch.float32).clone()
+    o = torch.zeros_like(v)
+    causal = torch.triu(
+        torch.ones(bt, bt, dtype=torch.bool, device=q.device), diagonal=1
+    )
+    for i in range(n):
+        q_i, k_i, v_i = (q[:, :, i], k[:, :, i], v[:, :, i])
+        a_i = (q_i @ k_i.transpose(-1, -2) * l_mask[:, :, i]).masked_fill(causal, 0)
+        v_new = v_i - k_cumdecay[:, :, i] @ S
+        o_inter = q_i * decay[:, :, i, :, None].exp() @ S
+        o[:, :, i] = o_inter + a_i @ v_new
+        S = (
+            S * decay[:, :, i, -1, None, None].exp()
+            + (
+                k_i * (decay[:, :, i, -1, None] - decay[:, :, i]).exp()[..., None]
+            ).transpose(-1, -2)
+            @ v_new
+        )
+    o = o.reshape(B, H, L, Vd)[:, :, :T]
+    return (o.transpose(1, 2).contiguous(), S)
+
+
+def causal_depthwise_conv1d(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor | None = None,
+    activation: str | None = "silu",
+) -> torch.Tensor:
+    B, T, C = x.shape
+    W = weight.shape[-1]
+    xt = x.transpose(1, 2)
+    xp = F.pad(xt, (W - 1, 0))
+    y = F.conv1d(xp, weight, bias=bias, groups=C)[..., :T].transpose(1, 2)
+    if activation == "silu":
+        y = F.silu(y)
+    elif activation is not None:
+        raise ValueError(f"unsupported activation {activation!r}")
+    return y
+
+
+def _conv_chunk(
+    p: torch.Tensor,
+    weight: torch.Tensor,
+    conv_state: torch.Tensor | None,
+    bias: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    B, L, C = p.shape
+    W = weight.shape[-1]
+    pt = p.transpose(1, 2)
+    if conv_state is None:
+        conv_state = pt.new_zeros(B, C, W - 1)
+    full = torch.cat([conv_state, pt], dim=-1)
+    y = F.silu(F.conv1d(full, weight, bias=bias, groups=C)).transpose(1, 2)
+    return (y, full[..., -(W - 1) :])
+
+
+def gdn_gate(
+    g: torch.Tensor, A_log: torch.Tensor, dt_bias: torch.Tensor | None = None
+) -> torch.Tensor:
+    g = g.float()
+    if dt_bias is not None:
+        g = g + dt_bias.float()
+    return -A_log.float().exp() * F.softplus(g)
+
+
+def gated_rmsnorm(
+    o: torch.Tensor, gate: torch.Tensor, weight: torch.Tensor, eps: float = 1e-05
+) -> torch.Tensor:
+    o = o.float()
+    gate = gate.float()
+    rstd = torch.rsqrt(o.pow(2).mean(-1, keepdim=True) + eps)
+    return o * rstd * weight.float() * (gate * torch.sigmoid(gate))
+
+
+class CPUGatedDeltaNet:
+    def __init__(self, gdn):
+        self.gdn = gdn
+        self.H = gdn.num_heads
+        self.Dk = gdn.head_k_dim
+        self.Dv = gdn.head_v_dim
+        self.C = gdn.value_dim
+        self.W = gdn.conv_size
+        self.scale = self.Dk ** (-0.5)
+        self.eps = getattr(gdn.o_norm, "eps", 1e-05)
+        self.chunk_size = 64
+        assert not gdn.allow_neg_eigval, "allow_neg_eigval not supported"
+        assert gdn.use_gate and gdn.use_short_conv
+
+    def _project(self, x, conv_state):
+        g = self.gdn
+        q, sq = _conv_chunk(
+            g.q_proj(x),
+            g.q_conv1d.weight,
+            None if conv_state is None else conv_state[0],
+        )
+        k, sk = _conv_chunk(
+            g.k_proj(x),
+            g.k_conv1d.weight,
+            None if conv_state is None else conv_state[1],
+        )
+        v, sv = _conv_chunk(
+            g.v_proj(x),
+            g.v_conv1d.weight,
+            None if conv_state is None else conv_state[2],
+        )
+        B, L = (x.shape[0], x.shape[1])
+        q = F.normalize(q.view(B, L, self.H, self.Dk), p=2, dim=-1)
+        k = F.normalize(k.view(B, L, self.H, self.Dk), p=2, dim=-1)
+        v = v.view(B, L, self.H, self.Dv)
+        beta = torch.sigmoid(g.b_proj(x))
+        gate_g = gdn_gate(g.a_proj(x), g.A_log, g.dt_bias)
+        return (q, k, v, beta, gate_g, (sq, sk, sv))
+
+    def _output(self, o, x):
+        g = self.gdn
+        B, L = (x.shape[0], x.shape[1])
+        gate = g.g_proj(x).view(B, L, self.H, self.Dv)
+        o = gated_rmsnorm(o, gate, g.o_norm.weight, self.eps)
+        o = o.reshape(B, L, self.C)
+        return g.o_proj(o)
+
+    def init_state(self, batch_size: int = 1, device=None):
+        if device is None:
+            device = self.gdn.q_proj.weight.device
+        return {
+            "recurrent_state": torch.zeros(
+                batch_size, self.H, self.Dk, self.Dv, device=device
+            ),
+            "conv": tuple(
+                (
+                    torch.zeros(batch_size, self.C, self.W - 1, device=device)
+                    for _ in range(3)
+                )
+            ),
+        }
+
+    def _recurrence(self, q, k, v, beta, g, initial_state):
+        if q.shape[1] >= self.chunk_size:
+            return gdn_chunk_recurrent(
+                q,
+                k,
+                v,
+                beta,
+                g,
+                self.scale,
+                initial_state=initial_state,
+                chunk_size=self.chunk_size,
+            )
+        return gdn_recurrent(q, k, v, beta, g, self.scale, initial_state=initial_state)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        q, k, v, beta, gate_g, _ = self._project(x, None)
+        o, _ = self._recurrence(q, k, v, beta, gate_g, initial_state=None)
+        return self._output(o, x).to(x.dtype)
+
+    def forward_chunk(self, x: torch.Tensor, state):
+        conv_state = None if state is None else state["conv"]
+        rec = None if state is None else state["recurrent_state"]
+        q, k, v, beta, gate_g, new_conv = self._project(x, conv_state)
+        o, rec = self._recurrence(q, k, v, beta, gate_g, initial_state=rec)
+        return (
+            self._output(o, x).to(x.dtype),
+            {"recurrent_state": rec, "conv": new_conv},
+        )
+
+    def step(self, x: torch.Tensor, state):
+        if x.dim() == 2:
+            x = x.unsqueeze(1)
+        o, state = self.forward_chunk(x, state)
+        return (o.squeeze(1), state)

cpu_landmark.py

@@ -0,0 +1,99 @@
+from __future__ import annotations
+import torch
+import torch.nn.functional as F
+
+
+class CPULandmark:
+    def __init__(self, mixer):
+        self.mixer = mixer
+        self.H = mixer.n_head
+        self.D = mixer.d_head
+        self.C = self.H * self.D
+        self.chunk = mixer.chunk
+        self.max_land = mixer.max_land
+
+    def init_state(self, batch_size: int = 1, device=None, c: int | None = None):
+        if device is None:
+            device = self.mixer.q_proj.weight.device
+        return {
+            "k": torch.zeros(batch_size, self.max_land, self.H, self.D, device=device),
+            "v": torch.zeros(batch_size, self.max_land, self.H, self.D, device=device),
+            "n_land": torch.zeros(batch_size, dtype=torch.long, device=device),
+            "acc_sum": torch.zeros(batch_size, self.C, device=device),
+            "acc_cnt": torch.zeros(batch_size, dtype=torch.long, device=device),
+            "c": int(c) if c else self.chunk,
+        }
+
+    def _sink_kv(self, batch_size: int, device):
+        sink = self.mixer.sink
+        sk = self.mixer.k_proj(sink).reshape(1, 1, self.H, self.D)
+        sv = self.mixer.v_proj(sink).reshape(1, 1, self.H, self.D)
+        exp = (batch_size, 1, self.H, self.D)
+        return (sk.to(device).expand(*exp), sv.to(device).expand(*exp))
+
+    def step(self, x: torch.Tensor, state):
+        if x.dim() == 3:
+            x = x.squeeze(1)
+        B = x.shape[0]
+        if state is None:
+            state = self.init_state(B, device=x.device)
+        c = state["c"]
+        n = state["n_land"]
+        Lmax = self.max_land
+        sk, sv = self._sink_kv(B, x.device)
+        keys = torch.cat([sk, state["k"]], dim=1)
+        vals = torch.cat([sv, state["v"]], dim=1)
+        idx = torch.arange(Lmax, device=x.device)[None, :]
+        valid = torch.cat(
+            [torch.ones(B, 1, dtype=torch.bool, device=x.device), idx < n[:, None]],
+            dim=1,
+        )
+        q = self.mixer.q_proj(x).reshape(B, 1, self.H, self.D).transpose(1, 2)
+        o = F.scaled_dot_product_attention(
+            q,
+            keys.transpose(1, 2),
+            vals.transpose(1, 2),
+            attn_mask=valid[:, None, None, :],
+        )
+        o = self.mixer.o_proj(o.transpose(1, 2).reshape(B, self.C))
+        acc_sum = state["acc_sum"] + x
+        acc_cnt = state["acc_cnt"] + 1
+        k, v = (state["k"].clone(), state["v"].clone())
+        fin = acc_cnt == c
+        if bool(fin.any()):
+            mean = acc_sum / c
+            lk = self.mixer.k_proj(mean).reshape(B, self.H, self.D)
+            lv = self.mixer.v_proj(mean).reshape(B, self.H, self.D)
+            rows = torch.arange(B, device=x.device)
+            full = fin & (n >= Lmax)
+            if bool(full.any()):
+                k[full] = torch.cat([k[full][:, 1:], k[full][:, :1]], dim=1)
+                v[full] = torch.cat([v[full][:, 1:], v[full][:, :1]], dim=1)
+            write_pos = n.clamp(max=Lmax - 1)
+            fr, wp = (rows[fin], write_pos[fin])
+            k[fr, wp] = lk[fr]
+            v[fr, wp] = lv[fr]
+            n = (n + fin.long()).clamp(max=Lmax)
+            acc_sum = torch.where(fin[:, None], torch.zeros_like(acc_sum), acc_sum)
+            acc_cnt = torch.where(fin, torch.zeros_like(acc_cnt), acc_cnt)
+        return (
+            o,
+            {
+                "k": k,
+                "v": v,
+                "n_land": n,
+                "acc_sum": acc_sum,
+                "acc_cnt": acc_cnt,
+                "c": c,
+            },
+        )
+
+    def forward_chunk(self, x: torch.Tensor, state):
+        B, L, _ = x.shape
+        if state is None:
+            state = self.init_state(B, device=x.device)
+        outs = []
+        for i in range(L):
+            o, state = self.step(x[:, i], state)
+            outs.append(o)
+        return (torch.stack(outs, dim=1), state)

cpu_patch.py

@@ -0,0 +1,157 @@
+from __future__ import annotations
+import types
+import importlib
+import torch
+
+
+def _sib(mod, *names):
+    try:
+        m = importlib.import_module("." + mod, __package__ or None)
+    except (ImportError, TypeError, ValueError):
+        m = importlib.import_module(mod)
+    return [getattr(m, n) for n in names]
+
+
+(CPUGatedDeltaNet,) = _sib("cpu_delta", "CPUGatedDeltaNet")
+(CPULandmark,) = _sib("cpu_landmark", "CPULandmark")
+CPUSlidingWindow, swa_fused_forward = _sib(
+    "cpu_swa", "CPUSlidingWindow", "swa_fused_forward"
+)
+
+
+def _is_delta_block(block) -> bool:
+    return (
+        getattr(block, "is_gla", False)
+        and getattr(block.mixer, "gla_delta", False)
+        and hasattr(block.mixer, "gdn")
+    )
+
+
+def _is_landmark_block(block) -> bool:
+    return getattr(block, "is_landmark", False)
+
+
+def _has_swa_fused(block) -> bool:
+    return int(getattr(block.mixer, "swa_fused_window", 0)) > 0
+
+
+def _make_block_methods(cpu_gdn, cpu_swa):
+
+    def step(self, x, bstate):
+        h = self.ln1(x)
+        if cpu_swa is not None:
+            h, bstate["swa"] = cpu_swa.step(h, bstate.get("swa"))
+            if h.dim() == 2:
+                h = h.unsqueeze(1)
+        o, bstate["gdn"] = cpu_gdn.step(h, bstate.get("gdn"))
+        x = x + o
+        x = x + self.ffn(self.ln2(x))
+        return (x, bstate)
+
+    def forward_chunk(self, x, bstate):
+        h = self.ln1(x)
+        if cpu_swa is not None:
+            h, bstate["swa"] = cpu_swa.forward_chunk(h, bstate.get("swa"))
+        o, bstate["gdn"] = cpu_gdn.forward_chunk(h, bstate.get("gdn"))
+        x = x + o
+        x = x + self.ffn(self.ln2(x))
+        return (x, bstate)
+
+    return (step, forward_chunk)
+
+
+def _make_landmark_block_methods(cpu_land):
+
+    def step(self, x, bstate):
+        h = self.ln1(x)
+        o, bstate["land"] = cpu_land.step(h, bstate.get("land"))
+        x = x + o
+        x = x + self.ffn(self.ln2(x))
+        return (x, bstate)
+
+    def forward_chunk(self, x, bstate):
+        h = self.ln1(x)
+        o, bstate["land"] = cpu_land.forward_chunk(h, bstate.get("land"))
+        x = x + o
+        x = x + self.ffn(self.ln2(x))
+        return (x, bstate)
+
+    return (step, forward_chunk)
+
+
+def _make_mixer_forward(cpu_gdn, cpu_swa, orig_forward):
+
+    def forward(self, x):
+        if x.is_cuda:
+            return orig_forward(x)
+        if cpu_swa is not None:
+            x = swa_fused_forward(self, x)
+        return cpu_gdn.forward(x)
+
+    return forward
+
+
+def _make_init_state(model):
+    cfg = model.cfg
+    H = cfg.n_head
+    D = cfg.n_embd // H
+    M = cfg.fno_modes
+    C = cfg.n_embd
+
+    def init_state(self, batch_size: int = 1, device=None):
+        if device is None:
+            device = next(self.parameters()).device
+        states = []
+        for block in self.blocks:
+            if _is_delta_block(block):
+                states.append(
+                    {"swa": None, "gdn": None}
+                    if _has_swa_fused(block)
+                    else {"gdn": None}
+                )
+            elif _is_landmark_block(block):
+                states.append({"land": None})
+            elif block.is_gla:
+                states.append(
+                    {
+                        "gla_state": torch.zeros(batch_size, H, D, D, device=device),
+                        "z_norm": torch.zeros(batch_size, H, D, device=device),
+                    }
+                )
+            else:
+                states.append(
+                    {"buf": torch.zeros(batch_size, M, C, device=device), "pos": 0}
+                )
+        return states
+
+    return init_state
+
+
+def enable_cpu_delta(model) -> int:
+    n = 0
+    for block in model.blocks:
+        if _is_delta_block(block):
+            cpu_gdn = CPUGatedDeltaNet(block.mixer.gdn)
+            cpu_swa = CPUSlidingWindow(block.mixer) if _has_swa_fused(block) else None
+            step, forward_chunk = _make_block_methods(cpu_gdn, cpu_swa)
+            block.step = types.MethodType(step, block)
+            block.forward_chunk = types.MethodType(forward_chunk, block)
+            block.mixer.forward = types.MethodType(
+                _make_mixer_forward(cpu_gdn, cpu_swa, block.mixer.forward), block.mixer
+            )
+            block.mixer._cpu_gdn = cpu_gdn
+            block.mixer._cpu_swa = cpu_swa
+            n += 1
+        elif _is_landmark_block(block):
+            cpu_land = CPULandmark(block.mixer)
+            step, forward_chunk = _make_landmark_block_methods(cpu_land)
+            block.step = types.MethodType(step, block)
+            block.forward_chunk = types.MethodType(forward_chunk, block)
+            block.mixer._cpu_land = cpu_land
+            if not hasattr(block.mixer, "prepare_inference"):
+                block.mixer.prepare_inference = types.MethodType(
+                    lambda self: None, block.mixer
+                )
+            n += 1
+    model.init_state = types.MethodType(_make_init_state(model), model)
+    return n

cpu_swa.py

@@ -0,0 +1,89 @@
+from __future__ import annotations
+import torch
+import torch.nn.functional as F
+
+
+def banded_softmax_attention(
+    q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, w: int
+) -> torch.Tensor:
+    T = q.shape[-2]
+    i = torch.arange(T, device=q.device)[:, None]
+    j = torch.arange(T, device=q.device)[None, :]
+    mask = (j <= i) & (j > i - w)
+    return F.scaled_dot_product_attention(q, k, v, attn_mask=mask)
+
+
+def swa_fused_forward(mixer, x: torch.Tensor) -> torch.Tensor:
+    B, T, C = x.shape
+    H, w = (mixer.n_head, mixer.swa_fused_window)
+    D = C // H
+    q = mixer.swa_fused_q(x).reshape(B, T, H, D).transpose(1, 2)
+    k = mixer.swa_fused_k(x).reshape(B, T, H, D).transpose(1, 2)
+    v = mixer.swa_fused_v(x).reshape(B, T, H, D).transpose(1, 2)
+    o = banded_softmax_attention(q, k, v, w)
+    o = o.transpose(1, 2).reshape(B, T, C)
+    return x + mixer.swa_fused_o(o)
+
+
+class CPUSlidingWindow:
+    def __init__(self, mixer):
+        self.mixer = mixer
+        self.H = mixer.n_head
+        self.C = mixer.swa_fused_q.out_features
+        self.D = self.C // self.H
+        self.w = mixer.swa_fused_window
+        self.scale = self.D ** (-0.5)
+
+    def init_state(self, batch_size: int = 1, device=None):
+        if device is None:
+            device = self.mixer.swa_fused_q.weight.device
+        W1 = self.w - 1
+        return {
+            "k": torch.zeros(batch_size, W1, self.H, self.D, device=device),
+            "v": torch.zeros(batch_size, W1, self.H, self.D, device=device),
+            "n": 0,
+        }
+
+    def _project(self, x):
+        B, L, _ = x.shape
+        q = self.mixer.swa_fused_q(x).reshape(B, L, self.H, self.D)
+        k = self.mixer.swa_fused_k(x).reshape(B, L, self.H, self.D)
+        v = self.mixer.swa_fused_v(x).reshape(B, L, self.H, self.D)
+        return (q, k, v)
+
+    def forward_chunk(self, x: torch.Tensor, state):
+        B, L, C = x.shape
+        w = self.w
+        if state is None:
+            state = self.init_state(B, device=x.device)
+        n = state["n"]
+        W1 = w - 1
+        ck = state["k"][:, W1 - n :, :, :] if n > 0 else state["k"][:, :0]
+        cv = state["v"][:, W1 - n :, :, :] if n > 0 else state["v"][:, :0]
+        q, k, v = self._project(x)
+        k_all = torch.cat([ck, k], dim=1)
+        v_all = torch.cat([cv, v], dim=1)
+        Tk = n + L
+        iq = torch.arange(L, device=x.device)[:, None] + n
+        jk = torch.arange(Tk, device=x.device)[None, :]
+        mask = (jk <= iq) & (jk > iq - w)
+        qh = q.transpose(1, 2)
+        kh = k_all.transpose(1, 2)
+        vh = v_all.transpose(1, 2)
+        o = F.scaled_dot_product_attention(qh, kh, vh, attn_mask=mask)
+        o = o.transpose(1, 2).reshape(B, L, C)
+        out = x + self.mixer.swa_fused_o(o)
+        keep = min(W1, Tk)
+        new_k = state["k"].clone()
+        new_v = state["v"].clone()
+        if keep > 0:
+            new_k[:, W1 - keep :, :, :] = k_all[:, Tk - keep :, :, :]
+            new_v[:, W1 - keep :, :, :] = v_all[:, Tk - keep :, :, :]
+        new_state = {"k": new_k, "v": new_v, "n": keep}
+        return (out, new_state)
+
+    def step(self, x: torch.Tensor, state):
+        if x.dim() == 2:
+            x = x.unsqueeze(1)
+        o, state = self.forward_chunk(x, state)
+        return (o.squeeze(1), state)

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0edb28e43cd050f22d6960421f0c359ad23778136929285254cbb2ab084d28a9
+size 3576787488

model_cpu_gpt2.py

@@ -0,0 +1,1262 @@
+from __future__ import annotations
+import os
+from dataclasses import dataclass
+from typing import Optional
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def _detect_cpu_bf16() -> bool:
+    try:
+        with open("/proc/cpuinfo") as f:
+            return "avx512_bf16" in f.read()
+    except Exception:
+        return False
+
+
+_CPU_HAS_BF16: bool = _detect_cpu_bf16()
+try:
+    import pyfftw
+
+    pyfftw.interfaces.cache.enable()
+    pyfftw.interfaces.cache.set_keepalive_time(60.0)
+except ImportError:
+    pass
+try:
+    if os.environ.get("NO_CPP_EXT"):
+        raise ImportError("C++ extensions disabled via NO_CPP_EXT")
+    import torch.utils.cpp_extension as _cpp_ext
+
+    _ext_path = os.path.join(os.path.dirname(__file__), "csrc")
+    _gla_ext = (
+        _cpp_ext.load(
+            name="gla_scan_cpu",
+            sources=[os.path.join(_ext_path, "gla_scan.cpp")],
+            extra_cflags=["-O3", "-fopenmp", "-march=native"],
+            extra_ldflags=["-fopenmp"],
+            verbose=False,
+        )
+        if os.path.exists(os.path.join(_ext_path, "gla_scan.cpp"))
+        else None
+    )
+except Exception:
+    _gla_ext = None
+try:
+    if os.environ.get("NO_CPP_EXT"):
+        raise ImportError("C++ extensions disabled via NO_CPP_EXT")
+    import torch.utils.cpp_extension as _fno_cpp_ext
+
+    _fno_ext_src = os.path.join(os.path.dirname(__file__), "csrc", "fno_conv.cpp")
+    _fno_ext = (
+        _fno_cpp_ext.load(
+            name="fno_conv_cpu",
+            sources=[_fno_ext_src],
+            extra_cflags=["-O3", "-fopenmp", "-march=native"],
+            extra_ldflags=["-fopenmp"],
+            verbose=False,
+        )
+        if os.path.exists(_fno_ext_src)
+        else None
+    )
+except Exception:
+    _fno_ext = None
+try:
+    from flashfftconv import FlashFFTConv as _FlashFFTConv
+
+    _flash_fft_available = True
+except ImportError:
+    _FlashFFTConv = None
+    _flash_fft_available = False
+try:
+    from fla.ops.gla import chunk_gla as _fla_chunk_gla
+
+    _fla_available = True
+    try:
+        from fla.ops import chunk_gated_delta_rule as _fla_chunk_gdr
+        from fla.layers import GatedDeltaNet as _FlaGatedDeltaNet
+    except Exception:
+        _fla_chunk_gdr = None
+        _FlaGatedDeltaNet = None
+except ImportError:
+    _fla_chunk_gla = None
+    _fla_chunk_gdr = None
+    _FlaGatedDeltaNet = None
+    _fla_available = False
+
+
+def _fft_dtype_to_torch(fft_dtype: str):
+    return {"fp32": torch.float32, "bf16": torch.bfloat16, "fp16": torch.float16}[
+        fft_dtype
+    ]
+
+
+def _gla_scan_py(q_k, k_k, v_f, chunk_gate, CS):
+    B, H, T, D = q_k.shape
+    n_chunks = T // CS
+    bf16 = q_k.dtype == torch.bfloat16
+    q_c = q_k.reshape(B, H, n_chunks, CS, D)
+    k_c = k_k.reshape(B, H, n_chunks, CS, D)
+    v_c = v_f.reshape(B, H, n_chunks, CS, D)
+    state = torch.zeros(B, H, D, D, device=q_k.device, dtype=torch.float32)
+    z_norm = torch.zeros(B, H, D, device=q_k.device, dtype=torch.float32)
+    chunks = []
+    for c in range(n_chunks):
+        if bf16:
+            num = q_c[:, :, c] @ state.bfloat16()
+            den = (q_c[:, :, c] @ z_norm.bfloat16().unsqueeze(-1)).clamp(min=1.0)
+            chunks.append(num / den)
+            g = chunk_gate[:, :, c, None, None]
+            state = g * state + (k_c[:, :, c].transpose(-2, -1) @ v_c[:, :, c]).float()
+            z_norm = chunk_gate[:, :, c, None] * z_norm + k_c[:, :, c].sum(-2).float()
+        else:
+            num = q_c[:, :, c] @ state
+            den = (q_c[:, :, c] @ z_norm.unsqueeze(-1)).clamp(min=1.0)
+            chunks.append(num / den)
+            g = chunk_gate[:, :, c, None, None]
+            state = g * state + k_c[:, :, c].transpose(-2, -1) @ v_c[:, :, c]
+            z_norm = chunk_gate[:, :, c, None] * z_norm + k_c[:, :, c].sum(dim=-2)
+    return torch.cat(chunks, dim=2)
+
+
+@torch.compiler.disable
+class _GLAScanFn(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, q_k, k_k, v_f, chunk_gate, CS):
+        ctx.CS = CS
+        with torch.no_grad():
+            out, states, z_norms = _gla_ext.gla_scan_fwd(
+                q_k.contiguous(),
+                k_k.contiguous(),
+                v_f.contiguous(),
+                chunk_gate.contiguous(),
+                CS,
+            )
+        ctx.save_for_backward(q_k, k_k, v_f, chunk_gate, states, z_norms)
+        return out
+
+    @staticmethod
+    def backward(ctx, grad_out):
+        q_k, k_k, v_f, chunk_gate, states, z_norms = ctx.saved_tensors
+        CS = ctx.CS
+        d_q, d_k, d_v, d_gate = _gla_ext.gla_scan_backward(
+            grad_out.float().contiguous(),
+            q_k.contiguous(),
+            k_k.contiguous(),
+            v_f.contiguous(),
+            chunk_gate.contiguous(),
+            states,
+            z_norms,
+            CS,
+        )
+        return (d_q, d_k, d_v, d_gate, None)
+
+
+@dataclass
+class CPUGPTConfig:
+    vocab_size: int = 50257
+    seq_len: int = 1024
+    n_layer: int = 12
+    n_embd: int = 768
+    n_head: int = 12
+    fno_modes: int = 256
+    gla_chunk: int = 64
+    ffn_hidden: int = 2048
+    layer_pattern: str = "SSSL"
+    bias: bool = False
+    dropout: float = 0.0
+    fft_dtype: str = "fp32"
+    use_flash_fft: bool = False
+    gla_delta: bool = False
+    sliding_window: int = 0
+    swa_window: int = 0
+    swa_fused_window: int = 0
+    attn_layer_every: int = 0
+    attn_full: bool = False
+    attn_window: int = 512
+    landmark_layer_every: int = 0
+    landmark_chunk: int = 32
+    landmark_max: int = 64
+    gdn_expand_v: float = 1.0
+    gdn_head_dim: int = 0
+
+    def __post_init__(self):
+        assert self.n_embd % 16 == 0, (
+            f"n_embd={self.n_embd} must be divisible by 16 for AMX BF16"
+        )
+        assert self.ffn_hidden % 16 == 0, (
+            f"ffn_hidden={self.ffn_hidden} must be divisible by 16 for AMX BF16"
+        )
+
+
+def _layer_is_gla(layer_idx: int, n_layer: int, pattern: str) -> bool:
+    if pattern == "SSSL":
+        return layer_idx % 4 == 3
+    elif pattern in ("SGSG", "SLSL"):
+        return layer_idx % 2 == 1
+    elif pattern == "SSLL":
+        return layer_idx % 4 >= 2
+    elif pattern == "SLLL":
+        return layer_idx % 4 >= 1
+    elif pattern == "GLA":
+        return True
+    elif pattern == "FNO":
+        return False
+    return False
+
+
+def rms_norm(x: torch.Tensor) -> torch.Tensor:
+    return F.rms_norm(x, (x.size(-1),))
+
+
+class SwiGLU(nn.Module):
+    def __init__(self, cfg: CPUGPTConfig):
+        super().__init__()
+        h = cfg.ffn_hidden
+        d = cfg.n_embd
+        self.gate = nn.Linear(d, h, bias=cfg.bias)
+        self.up = nn.Linear(d, h, bias=cfg.bias)
+        self.down = nn.Linear(h, d, bias=cfg.bias)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.down(F.silu(self.gate(x)) * self.up(x))
+
+
+@torch.compiler.disable
+class _FNOConvFn(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, x, filter_td, T):
+        x_f = x.float().contiguous()
+        f_f = filter_td.float().detach().contiguous()
+        y, Xf, Hf = _fno_ext.fno_conv_fwd(x_f, f_f, T)
+        ctx.save_for_backward(Xf, Hf)
+        ctx.n_use = min(filter_td.shape[1], T)
+        ctx.M = filter_td.shape[1]
+        return y.to(x.dtype)
+
+    @staticmethod
+    def backward(ctx, grad_out):
+        Xf, Hf = ctx.saved_tensors
+        d_x, d_filter = _fno_ext.fno_conv_backward(
+            grad_out.float().contiguous(), Xf, Hf, ctx.n_use, ctx.M
+        )
+        return (d_x, d_filter, None)
+
+
+class FNOSeqMixer(nn.Module):
+    def __init__(self, cfg: CPUGPTConfig):
+        super().__init__()
+        C = cfg.n_embd
+        M = cfg.fno_modes
+        self.filter_td = nn.Parameter(torch.empty(C, M))
+        self.out_scale = nn.Linear(C, C, bias=cfg.bias)
+        self.register_buffer("_hf_cache", None, persistent=False)
+        self._filter_version: int = -1
+        self._fft_dtype = _fft_dtype_to_torch(cfg.fft_dtype)
+        self._use_flash_fft = cfg.use_flash_fft and _flash_fft_available
+        self._flash_fft_conv = None
+        self._flash_fft_T: int = -1
+        nn.init.normal_(self.filter_td, std=0.02)
+
+    def _flash_conv(self, x: torch.Tensor, T: int, C: int, n_use: int) -> torch.Tensor:
+        if self._flash_fft_conv is None or self._flash_fft_T != T:
+            self._flash_fft_conv = _FlashFFTConv(2 * T, dtype=self._fft_dtype).to(
+                x.device
+            )
+            self._flash_fft_T = T
+        h = self.filter_td.new_zeros(C, 2 * T)
+        h[:, :n_use] = self.filter_td[:, :n_use]
+        xp = F.pad(x.transpose(1, 2), (0, T))
+        y = self._flash_fft_conv(xp.to(self._fft_dtype), h.to(self._fft_dtype))
+        return y[:, :, :T].transpose(1, 2).to(x.dtype)
+
+    def _rfft_conv(self, x: torch.Tensor, T: int, C: int, n_use: int) -> torch.Tensor:
+        dt = self._fft_dtype
+        h = self.filter_td.new_zeros(2 * T, C).to(dt)
+        h[:n_use] = self.filter_td[:, :n_use].T.to(dt)
+        xp = F.pad(x, (0, 0, 0, T)).to(dt)
+        Xf = torch.fft.rfft(xp, dim=1)
+        Hf = torch.fft.rfft(h, dim=0).unsqueeze(0)
+        Xr, Xi = (Xf.real, Xf.imag)
+        Hr, Hi = (Hf.real, Hf.imag)
+        YHf = torch.view_as_complex(
+            torch.stack([Xr * Hr - Xi * Hi, Xr * Hi + Xi * Hr], dim=-1).contiguous()
+        )
+        return torch.fft.irfft(YHf, n=2 * T, dim=1)[:, :T]
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        B, T, C = x.shape
+        n_use = min(self.filter_td.shape[1], T)
+        if _fno_ext is not None:
+            y = _FNOConvFn.apply(x, self.filter_td, T)
+        elif self._use_flash_fft and x.is_cuda:
+            y = self._flash_conv(x, T, C, n_use)
+        elif not self.training:
+            y = self._forward_eval_cached(x, T, C, n_use)
+        else:
+            y = self._rfft_conv(x, T, C, n_use)
+        return self.out_scale(y.to(x.dtype))
+
+    @torch.compiler.disable
+    def _forward_eval_cached(
+        self, x: torch.Tensor, T: int, C: int, n_use: int
+    ) -> torch.Tensor:
+        ver = self.filter_td._version
+        if (
+            self._hf_cache is None
+            or self._filter_version != ver
+            or self._hf_cache.shape[0] != T + 1
+        ):
+            h = self.filter_td.new_zeros(2 * T, C)
+            h[:n_use] = self.filter_td[:, :n_use].T.detach()
+            self._hf_cache = torch.fft.rfft(h.float(), dim=0).contiguous()
+            self._filter_version = ver
+        dt = self._fft_dtype
+        xp = F.pad(x, (0, 0, 0, T)).to(dt)
+        Xf = torch.fft.rfft(xp, dim=1)
+        H = self._hf_cache.unsqueeze(0)
+        Xr, Xi, Hr, Hi = (Xf.real, Xf.imag, H.real, H.imag)
+        YHf = torch.view_as_complex(
+            torch.stack([Xr * Hr - Xi * Hi, Xr * Hi + Xi * Hr], dim=-1).contiguous()
+        )
+        return torch.fft.irfft(YHf, n=2 * T, dim=1)[:, :T]
+
+    def prepare_inference(self):
+        if hasattr(self, "_h_td"):
+            return
+        with torch.no_grad():
+            pass
+
+    def _build_h_td(self, T: int) -> None:
+        if hasattr(self, "_h_td"):
+            return
+        C = self.filter_td.shape[0]
+        M = min(self.filter_td.shape[1], T)
+        with torch.no_grad():
+            h_td = self.filter_td[:, :M].T.float().contiguous()
+        self.register_buffer("_h_td", h_td.contiguous())
+
+    def step(
+        self, x: torch.Tensor, buf: torch.Tensor, pos: int
+    ) -> tuple[torch.Tensor, int]:
+        M = buf.shape[1]
+        self._build_h_td(M)
+        slot = pos % M
+        buf[:, slot, :] = x.detach()
+        indices = torch.arange(M, device=x.device)
+        ordered_idx = (slot + 1 + indices) % M
+        ordered = buf[:, ordered_idx, :]
+        h_flip = self._h_td.flip(0)
+        y = (h_flip.unsqueeze(0) * ordered).sum(1)
+        return (self.out_scale(y.to(x.dtype)), pos + 1)
+
+    def forward_chunk(self, h: torch.Tensor, ctx):
+        B, L, C = h.shape
+        M = self.filter_td.shape[1]
+        if ctx is None:
+            ctx = h.new_zeros(B, M - 1, C)
+        full = torch.cat([ctx, h], dim=1)
+        P = full.shape[1]
+        n_use = min(M, P)
+        dt = self._fft_dtype
+        hk = self.filter_td.new_zeros(2 * P, C).to(dt)
+        hk[:n_use] = self.filter_td[:, :n_use].T.to(dt)
+        xp = F.pad(full, (0, 0, 0, P)).to(dt)
+        Xf = torch.fft.rfft(xp, dim=1)
+        Hf = torch.fft.rfft(hk, dim=0).unsqueeze(0)
+        Xr, Xi, Hr, Hi = (Xf.real, Xf.imag, Hf.real, Hf.imag)
+        YHf = torch.view_as_complex(
+            torch.stack([Xr * Hr - Xi * Hi, Xr * Hi + Xi * Hr], dim=-1).contiguous()
+        )
+        y = torch.fft.irfft(YHf, n=2 * P, dim=1)[:, M - 1 : P]
+        return (self.out_scale(y.to(h.dtype)), full[:, -(M - 1) :, :])
+
+
+class GLAMixer(nn.Module):
+    def __init__(self, cfg: CPUGPTConfig):
+        super().__init__()
+        H = cfg.n_head
+        D = cfg.n_embd // H
+        C = cfg.n_embd
+        CS = cfg.gla_chunk
+        assert C % H == 0, "n_embd must be divisible by n_head"
+        assert cfg.seq_len % CS == 0, (
+            f"seq_len {cfg.seq_len} must be divisible by gla_chunk {CS}"
+        )
+        self.n_head = H
+        self.d_head = D
+        self.chunk = CS
+        self.gla_delta = bool(getattr(cfg, "gla_delta", False))
+        self.sliding_window = int(getattr(cfg, "sliding_window", 0))
+        self.swa_window = int(getattr(cfg, "swa_window", 0))
+        if self.swa_window > 0:
+            self.swa_q = nn.Linear(C, C, bias=False)
+            self.swa_k = nn.Linear(C, C, bias=False)
+            self.swa_v = nn.Linear(C, C, bias=False)
+            self.swa_o = nn.Linear(C, C, bias=False)
+        self.swa_fused_window = int(getattr(cfg, "swa_fused_window", 0))
+        if self.swa_fused_window > 0:
+            self.swa_fused_q = nn.Linear(C, C, bias=False)
+            self.swa_fused_k = nn.Linear(C, C, bias=False)
+            self.swa_fused_v = nn.Linear(C, C, bias=False)
+            self.swa_fused_o = nn.Linear(C, C, bias=False)
+            nn.init.zeros_(self.swa_fused_o.weight)
+        self.gdn_expand_v = float(getattr(cfg, "gdn_expand_v", 1.0))
+        self.gdn_head_dim = int(getattr(cfg, "gdn_head_dim", 0)) or D
+        if not (self.gla_delta and _FlaGatedDeltaNet is not None):
+            self.q_proj = nn.Linear(C, C, bias=False)
+            self.k_proj = nn.Linear(C, C, bias=False)
+            self.v_proj = nn.Linear(C, C, bias=False)
+            self.g_proj = nn.Linear(C, H, bias=True)
+            self.out_proj = nn.Linear(C, C, bias=False)
+            self.gla_scale = nn.Parameter(torch.ones(H) * 0.5)
+            nn.init.constant_(self.g_proj.bias, -4.0)
+        self._swa_mask = None
+        if self.gla_delta and _FlaGatedDeltaNet is not None:
+            self.gdn = _FlaGatedDeltaNet(
+                hidden_size=C,
+                num_heads=H,
+                head_dim=self.gdn_head_dim,
+                expand_v=self.gdn_expand_v,
+                mode="chunk",
+            )
+
+    def _python_scan(
+        self,
+        q_k: torch.Tensor,
+        k_k: torch.Tensor,
+        v: torch.Tensor,
+        chunk_gate: torch.Tensor,
+    ) -> torch.Tensor:
+        B, H, T, D = q_k.shape
+        CS = self.chunk
+        n_chunks = T // CS
+        q_c = q_k.view(B, H, n_chunks, CS, D)
+        k_c = k_k.view(B, H, n_chunks, CS, D)
+        v_c = v.view(B, H, n_chunks, CS, D)
+        state = torch.zeros(B, H, D, D, device=q_k.device, dtype=q_k.dtype)
+        z_norm = torch.zeros(B, H, D, device=q_k.device, dtype=q_k.dtype)
+        out = torch.zeros(B, H, T, D, device=q_k.device, dtype=q_k.dtype)
+        for c in range(n_chunks):
+            num = q_c[:, :, c] @ state
+            den = (q_c[:, :, c] @ z_norm.unsqueeze(-1)).clamp(min=1.0)
+            out[:, :, c * CS : (c + 1) * CS] = num / den
+            g = chunk_gate[:, :, c, None, None]
+            state = g * state + k_c[:, :, c].transpose(-2, -1) @ v_c[:, :, c]
+            z_norm = chunk_gate[:, :, c, None] * z_norm + k_c[:, :, c].sum(dim=-2)
+        return out
+
+    def _swa(self, q, k, v):
+        B, H, T, D = q.shape
+        w = self.sliding_window
+        pad = (w - T % w) % w
+        if pad:
+            q, k, v = (F.pad(t, (0, 0, 0, pad)) for t in (q, k, v))
+        Tp = T + pad
+        nb = Tp // w
+        qb = q.reshape(B, H, nb, w, D)
+        kb = k.reshape(B, H, nb, w, D)
+        vb = v.reshape(B, H, nb, w, D)
+        kc = torch.cat([F.pad(kb, (0, 0, 0, 0, 1, 0))[:, :, :-1], kb], dim=3)
+        vc = torch.cat([F.pad(vb, (0, 0, 0, 0, 1, 0))[:, :, :-1], vb], dim=3)
+        if (
+            self._swa_mask is None
+            or self._swa_mask.shape[-1] != 2 * w
+            or self._swa_mask.device != q.device
+        ):
+            i = torch.arange(w, device=q.device)
+            j = torch.arange(2 * w, device=q.device)
+            self._swa_mask = (j[None, :] < w) | (j[None, :] - w <= i[:, None])
+        out = F.scaled_dot_product_attention(qb, kc, vc, attn_mask=self._swa_mask)
+        return out.reshape(B, H, Tp, D)[:, :, :T]
+
+    def _swa_exact(self, x: torch.Tensor) -> torch.Tensor:
+        B, T, C = x.shape
+        H, D, w = (self.n_head, self.d_head, self.swa_window)
+        q = self.swa_q(x).reshape(B, T, H, D).transpose(1, 2)
+        k = self.swa_k(x).reshape(B, T, H, D).transpose(1, 2)
+        v = self.swa_v(x).reshape(B, T, H, D).transpose(1, 2)
+        i = torch.arange(T, device=x.device)[:, None]
+        j = torch.arange(T, device=x.device)[None, :]
+        mask = (j <= i) & (j > i - w)
+        o = F.scaled_dot_product_attention(q, k, v, attn_mask=mask)
+        o = o.transpose(1, 2).reshape(B, T, C)
+        return self.swa_o(o)
+
+    def _swa_fused(self, x: torch.Tensor) -> torch.Tensor:
+        B, T, C = x.shape
+        H, D, w = (self.n_head, self.d_head, self.swa_fused_window)
+        q = self.swa_fused_q(x).reshape(B, T, H, D).transpose(1, 2)
+        k = self.swa_fused_k(x).reshape(B, T, H, D).transpose(1, 2)
+        v = self.swa_fused_v(x).reshape(B, T, H, D).transpose(1, 2)
+        i = torch.arange(T, device=x.device)[:, None]
+        j = torch.arange(T, device=x.device)[None, :]
+        mask = (j <= i) & (j > i - w)
+        o = F.scaled_dot_product_attention(q, k, v, attn_mask=mask)
+        o = o.transpose(1, 2).reshape(B, T, C)
+        return x + self.swa_fused_o(o)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        B, T, C = x.shape
+        H, D, CS = (self.n_head, self.d_head, self.chunk)
+        if self.gla_delta and hasattr(self, "gdn") and x.is_cuda:
+            gdn_in = self._swa_fused(x) if self.swa_fused_window > 0 else x
+            o = self.gdn(gdn_in)
+            o = o[0] if isinstance(o, tuple) else o
+            if self.swa_window > 0:
+                o = o + self._swa_exact(x)
+            return o
+        q = self.q_proj(x).reshape(B, T, H, D).transpose(1, 2)
+        k = self.k_proj(x).reshape(B, T, H, D).transpose(1, 2)
+        v = self.v_proj(x).reshape(B, T, H, D).transpose(1, 2)
+        if x.is_cuda and _fla_available:
+            log_g = -F.softplus(self.g_proj(x))
+            q_t = q.transpose(1, 2).contiguous()
+            k_t = k.transpose(1, 2).contiguous()
+            v_t = v.transpose(1, 2).contiguous()
+            if self.gla_delta and _fla_chunk_gdr is not None:
+                beta = self.beta_proj(x).float()
+                g_delta = self.a_proj(x).float()
+                y_fla, _ = _fla_chunk_gdr(
+                    q_t,
+                    k_t,
+                    v_t,
+                    g=g_delta,
+                    beta=beta,
+                    scale=D ** (-0.5),
+                    output_final_state=False,
+                    use_beta_sigmoid_in_kernel=True,
+                    use_qk_l2norm_in_kernel=True,
+                )
+            else:
+                g_t = log_g.unsqueeze(-1).expand(B, T, H, D).contiguous().float()
+                y_fla, _ = _fla_chunk_gla(
+                    q_t, k_t, v_t, g=g_t, scale=D ** (-0.5), output_final_state=False
+                )
+            y = y_fla.to(x.dtype).transpose(1, 2).reshape(B, T, C)
+            if self.sliding_window > 0:
+                y = y + self._swa(q, k, v).transpose(1, 2).reshape(B, T, C)
+            return self.out_proj(y)
+        n_chunks = T // CS
+        BH = B * H
+        q_l = q.reshape(BH, n_chunks, CS, D)
+        k_l = k.reshape(BH, n_chunks, CS, D)
+        v_l = v.reshape(BH, n_chunks, CS, D)
+        y_local = F.scaled_dot_product_attention(q_l, k_l, v_l, is_causal=True).reshape(
+            B, H, T, D
+        )
+        if _CPU_HAS_BF16 and q.dtype == torch.bfloat16:
+            q_k = F.elu(q) + 1.0
+            k_k = F.elu(k) + 1.0
+            v_f = v
+        else:
+            q_k = F.elu(q.float()) + 1.0
+            k_k = F.elu(k.float()) + 1.0
+            v_f = v.float()
+        log_g = -F.softplus(self.g_proj(x).float())
+        log_g = log_g.transpose(1, 2)
+        chunk_log_g = log_g.reshape(B, H, n_chunks, CS).sum(-1)
+        chunk_gate = chunk_log_g.exp()
+        if _gla_ext is not None and q_k.dtype == torch.float32 and (not q_k.is_cuda):
+            y_gla = _GLAScanFn.apply(q_k, k_k, v_f, chunk_gate, CS)
+        else:
+            y_gla = _gla_scan_py(q_k, k_k, v_f, chunk_gate, CS)
+        scale = self.gla_scale.reshape(1, H, 1, 1).to(x.dtype)
+        y = (y_local + scale * y_gla.to(x.dtype)).transpose(1, 2).reshape(B, T, C)
+        return self.out_proj(y)
+
+    def step(
+        self, x: torch.Tensor, state: torch.Tensor, z_norm: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        B = x.shape[0]
+        H, D = (self.n_head, self.d_head)
+        q = self.q_proj(x).view(B, H, D)
+        k = self.k_proj(x).view(B, H, D)
+        v = self.v_proj(x).view(B, H, D)
+        log_g = -F.softplus(self.g_proj(x).float())
+        scale = D ** (-0.5)
+        if x.is_cuda and _fla_available:
+            st = (
+                state
+                if torch.is_tensor(state) and state.dtype == torch.float32
+                else None
+            )
+            g_t = log_g.view(B, 1, H, 1).expand(B, 1, H, D).contiguous()
+            y, state = _fla_chunk_gla(
+                q.view(B, 1, H, D),
+                k.view(B, 1, H, D),
+                v.view(B, 1, H, D),
+                g=g_t,
+                scale=scale,
+                initial_state=st,
+                output_final_state=True,
+            )
+            y = y.reshape(B, H * D).to(x.dtype)
+            return (self.out_proj(y), state, z_norm)
+        gate = log_g.exp()
+        qf, kf, vf = (q.float(), k.float(), v.float())
+        if not (torch.is_tensor(state) and state.dim() == 4):
+            state = torch.zeros(B, H, D, D, device=x.device, dtype=torch.float32)
+        state = gate.unsqueeze(-1).unsqueeze(-1) * state + torch.einsum(
+            "bhd,bhe->bhde", kf, vf
+        )
+        y_gla = torch.einsum("bhd,bhde->bhe", qf, state) * scale
+        return (self.out_proj(y_gla.reshape(B, H * D).to(x.dtype)), state, z_norm)
+
+    def forward_chunk(self, h: torch.Tensor, state):
+        B, L, C = h.shape
+        H, D = (self.n_head, self.d_head)
+        q = self.q_proj(h).view(B, L, H, D)
+        k = self.k_proj(h).view(B, L, H, D)
+        v = self.v_proj(h).view(B, L, H, D)
+        log_g = -F.softplus(self.g_proj(h).float())
+        g = log_g.unsqueeze(-1).expand(B, L, H, D).contiguous()
+        st = state if torch.is_tensor(state) and state.dtype == torch.float32 else None
+        y, state = _fla_chunk_gla(
+            q, k, v, g=g, scale=D ** (-0.5), initial_state=st, output_final_state=True
+        )
+        return (self.out_proj(y.reshape(B, L, C).to(h.dtype)), state)
+
+
+class AttnMixer(nn.Module):
+    def __init__(self, cfg: CPUGPTConfig):
+        super().__init__()
+        C, H = (cfg.n_embd, cfg.n_head)
+        assert C % H == 0
+        self.n_head = H
+        self.d_head = C // H
+        self.full = bool(getattr(cfg, "attn_full", False))
+        self.window = int(getattr(cfg, "attn_window", 512))
+        self.q_proj = nn.Linear(C, C, bias=False)
+        self.k_proj = nn.Linear(C, C, bias=False)
+        self.v_proj = nn.Linear(C, C, bias=False)
+        self.o_proj = nn.Linear(C, C, bias=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        B, T, C = x.shape
+        H, D = (self.n_head, self.d_head)
+        q = self.q_proj(x).reshape(B, T, H, D).transpose(1, 2)
+        k = self.k_proj(x).reshape(B, T, H, D).transpose(1, 2)
+        v = self.v_proj(x).reshape(B, T, H, D).transpose(1, 2)
+        if self.full:
+            o = F.scaled_dot_product_attention(q, k, v, is_causal=True)
+        else:
+            w = self.window
+            i = torch.arange(T, device=x.device)[:, None]
+            j = torch.arange(T, device=x.device)[None, :]
+            mask = (j <= i) & (j > i - w)
+            o = F.scaled_dot_product_attention(q, k, v, attn_mask=mask)
+        o = o.transpose(1, 2).reshape(B, T, C)
+        return self.o_proj(o)
+
+
+def _layer_is_attn(layer_idx: int, cfg: CPUGPTConfig) -> bool:
+    k = int(getattr(cfg, "attn_layer_every", 0))
+    if k <= 0:
+        return False
+    return (layer_idx + 1) % k == 0
+
+
+class LandmarkMixer(nn.Module):
+    def __init__(self, cfg: CPUGPTConfig):
+        super().__init__()
+        C, H = (cfg.n_embd, cfg.n_head)
+        assert C % H == 0
+        self.n_head = H
+        self.d_head = C // H
+        self.chunk = max(1, int(getattr(cfg, "landmark_chunk", 32)))
+        self.max_land = max(1, int(getattr(cfg, "landmark_max", 64)))
+        self.q_proj = nn.Linear(C, C, bias=False)
+        self.k_proj = nn.Linear(C, C, bias=False)
+        self.v_proj = nn.Linear(C, C, bias=False)
+        self.o_proj = nn.Linear(C, C, bias=False)
+        self.sink = nn.Parameter(torch.zeros(1, 1, C))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        B, T, C = x.shape
+        H, D = (self.n_head, self.d_head)
+        c = max(self.chunk, (T + self.max_land - 1) // self.max_land)
+        nc = (T + c - 1) // c
+        pad = nc * c - T
+        xp = F.pad(x, (0, 0, 0, pad))
+        land = xp.view(B, nc, c, C).mean(2)
+        land = torch.cat([self.sink.expand(B, 1, C), land], dim=1)
+        q = self.q_proj(x).reshape(B, T, H, D).transpose(1, 2)
+        k = self.k_proj(land).reshape(B, nc + 1, H, D).transpose(1, 2)
+        v = self.v_proj(land).reshape(B, nc + 1, H, D).transpose(1, 2)
+        tok_c = (torch.arange(T, device=x.device) // c)[:, None]
+        land_i = torch.arange(nc, device=x.device)[None, :]
+        sink_ok = torch.ones(T, 1, dtype=torch.bool, device=x.device)
+        chunk_ok = land_i < tok_c
+        mask = torch.cat([sink_ok, chunk_ok], dim=1).view(1, 1, T, nc + 1)
+        o = F.scaled_dot_product_attention(q, k, v, attn_mask=mask)
+        o = o.transpose(1, 2).reshape(B, T, C)
+        return self.o_proj(o)
+
+
+def _layer_is_landmark(layer_idx: int, cfg: CPUGPTConfig) -> bool:
+    k = int(getattr(cfg, "landmark_layer_every", 0))
+    return k > 0 and (layer_idx + 1) % k == 0
+
+
+class CPUGPTBlock(nn.Module):
+    def __init__(self, cfg: CPUGPTConfig, layer_idx: int):
+        super().__init__()
+        is_landmark = _layer_is_landmark(layer_idx, cfg)
+        is_attn = not is_landmark and _layer_is_attn(layer_idx, cfg)
+        is_gla = (
+            not is_landmark
+            and (not is_attn)
+            and _layer_is_gla(layer_idx, cfg.n_layer, cfg.layer_pattern)
+        )
+        if is_landmark:
+            self.mixer = LandmarkMixer(cfg)
+        elif is_attn:
+            self.mixer = AttnMixer(cfg)
+        else:
+            self.mixer = GLAMixer(cfg) if is_gla else FNOSeqMixer(cfg)
+        self.ffn = SwiGLU(cfg)
+        self.ln1 = nn.RMSNorm(cfg.n_embd)
+        self.ln2 = nn.RMSNorm(cfg.n_embd)
+        self.is_gla = is_gla
+        self.is_attn = is_attn
+        self.is_landmark = is_landmark
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = x + self.mixer(self.ln1(x))
+        x = x + self.ffn(self.ln2(x))
+        return x
+
+    def step(self, x: torch.Tensor, bstate: dict) -> tuple:
+        h = self.ln1(x)
+        mixer = self.mixer
+        if isinstance(mixer, (AttnMixer, LandmarkMixer)):
+            raise NotImplementedError(
+                "AttnMixer/LandmarkMixer recurrent step() not implemented; use forward()"
+            )
+        if isinstance(mixer, GLAMixer):
+            h_out, bstate["gla_state"], bstate["z_norm"] = mixer.step(
+                h, bstate["gla_state"], bstate["z_norm"]
+            )
+        else:
+            assert isinstance(mixer, FNOSeqMixer)
+            h_out, bstate["pos"] = mixer.step(h, bstate["buf"], bstate["pos"])
+        x = x + h_out
+        x = x + self.ffn(self.ln2(x))
+        return (x, bstate)
+
+    def forward_chunk(self, x: torch.Tensor, bstate: dict) -> tuple:
+        h = self.ln1(x)
+        mixer = self.mixer
+        if isinstance(mixer, (AttnMixer, LandmarkMixer)):
+            raise NotImplementedError(
+                "AttnMixer/LandmarkMixer forward_chunk() not implemented; use forward()"
+            )
+        if isinstance(mixer, GLAMixer):
+            h_out, bstate["gla_state"] = mixer.forward_chunk(h, bstate.get("gla_state"))
+        else:
+            assert isinstance(mixer, FNOSeqMixer)
+            h_out, bstate["fno_ctx"] = mixer.forward_chunk(h, bstate.get("fno_ctx"))
+        x = x + h_out
+        x = x + self.ffn(self.ln2(x))
+        return (x, bstate)
+
+
+class _ChunkedCEFn(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, x_flat, weight, targets, chunk_size):
+        N = x_flat.shape[0]
+        use_bf16 = x_flat.dtype == torch.bfloat16
+        count = (targets != -1).sum().clamp(min=1)
+        w_g = weight.bfloat16() if use_bf16 else weight.float()
+        loss_sum = torch.zeros((), dtype=torch.float32, device=x_flat.device)
+        for start in range(0, N, chunk_size):
+            end = min(start + chunk_size, N)
+            x_c = x_flat[start:end]
+            t_c = targets[start:end]
+            logits_c = (x_c @ w_g.T).float()
+            logits_c = 15.0 * torch.tanh(logits_c * (1.0 / 15.0))
+            loss_sum = loss_sum + F.cross_entropy(
+                logits_c, t_c, ignore_index=-1, reduction="sum"
+            )
+        ctx.save_for_backward(x_flat, weight, targets)
+        ctx.chunk_size = chunk_size
+        ctx.count = int(count.item())
+        ctx.use_bf16 = use_bf16
+        return loss_sum / count
+
+    @staticmethod
+    def backward(ctx, grad_out):
+        x_flat, weight, targets = ctx.saved_tensors
+        N, C = x_flat.shape
+        chunk_size = ctx.chunk_size
+        use_bf16 = ctx.use_bf16
+        scale = grad_out.item() / ctx.count
+        w_g = weight.bfloat16() if use_bf16 else weight.float()
+        dx = torch.zeros(N, C, dtype=torch.float32, device=x_flat.device)
+        dw = torch.zeros_like(weight, dtype=torch.float32)
+        for start in range(0, N, chunk_size):
+            end = min(start + chunk_size, N)
+            x_c = x_flat[start:end]
+            t_c = targets[start:end]
+            logits_c = (x_c @ w_g.T).float()
+            tanh_c = torch.tanh(logits_c * (1.0 / 15.0))
+            logits_cap = 15.0 * tanh_c
+            probs_c = torch.softmax(logits_cap, dim=-1)
+            valid = t_c != -1
+            probs_c[~valid] = 0.0
+            if valid.any():
+                probs_c[valid, t_c[valid]] -= 1.0
+            d_logits_c = probs_c * (1.0 - tanh_c * tanh_c) * scale
+            if use_bf16:
+                dx[start:end] = (d_logits_c.bfloat16() @ w_g).float()
+            else:
+                dx[start:end] = d_logits_c @ w_g
+            if use_bf16:
+                dw.add_((d_logits_c.bfloat16().T @ x_c).float())
+            else:
+                dw.addmm_(d_logits_c.T, x_c)
+        return (dx.to(x_flat.dtype), dw, None, None)
+
+
+def _chunked_cross_entropy(
+    weight: torch.Tensor, x: torch.Tensor, targets: torch.Tensor, chunk: int = 512
+) -> torch.Tensor:
+    B, T, C = x.shape
+    return _ChunkedCEFn.apply(
+        x.reshape(B * T, C), weight, targets.reshape(B * T), chunk
+    )
+
+
+class CPUGPT(nn.Module):
+    def __init__(self, cfg: CPUGPTConfig):
+        super().__init__()
+        self.cfg = cfg
+        self.wte = nn.Embedding(cfg.vocab_size, cfg.n_embd)
+        self.blocks = nn.ModuleList([CPUGPTBlock(cfg, i) for i in range(cfg.n_layer)])
+        self.ln_out = nn.RMSNorm(cfg.n_embd)
+        self.lm_head = nn.Linear(cfg.n_embd, cfg.vocab_size, bias=False)
+        self.lm_head.weight = self.wte.weight
+        self._init_weights()
+
+    def _init_weights(self):
+        nn.init.normal_(self.wte.weight, std=0.02)
+        for block in self.blocks:
+            if isinstance(block.mixer, GLAMixer) and hasattr(block.mixer, "q_proj"):
+                for proj in [
+                    block.mixer.q_proj,
+                    block.mixer.k_proj,
+                    block.mixer.v_proj,
+                    block.mixer.out_proj,
+                ]:
+                    nn.init.normal_(proj.weight, std=0.02)
+            if isinstance(block.mixer, AttnMixer):
+                for proj in [
+                    block.mixer.q_proj,
+                    block.mixer.k_proj,
+                    block.mixer.v_proj,
+                    block.mixer.o_proj,
+                ]:
+                    nn.init.normal_(proj.weight, std=0.02)
+            nn.init.normal_(block.ffn.gate.weight, std=0.02)
+            nn.init.normal_(block.ffn.up.weight, std=0.02)
+            nn.init.zeros_(block.ffn.down.weight)
+
+    def param_count(self) -> int:
+        return sum((p.numel() for p in self.parameters()))
+
+    def prepare_inference(self):
+        for block in self.blocks:
+            if not block.is_gla and (not getattr(block, "is_attn", False)):
+                block.mixer.prepare_inference()
+
+    def init_state(self, batch_size: int = 1, device=None) -> list:
+        if device is None:
+            device = next(self.parameters()).device
+        H = self.cfg.n_head
+        D = self.cfg.n_embd // H
+        M = self.cfg.fno_modes
+        C = self.cfg.n_embd
+        states = []
+        for block in self.blocks:
+            if getattr(block, "is_attn", False):
+                states.append({})
+            elif block.is_gla:
+                states.append(
+                    {
+                        "gla_state": torch.zeros(batch_size, H, D, D, device=device),
+                        "z_norm": torch.zeros(batch_size, H, D, device=device),
+                    }
+                )
+            else:
+                states.append(
+                    {"buf": torch.zeros(batch_size, M, C, device=device), "pos": 0}
+                )
+        return states
+
+    @torch.no_grad()
+    def step(self, idx: torch.Tensor, states: list):
+        x = self.wte(idx.unsqueeze(1) if idx.dim() == 1 else idx)
+        if x.dim() == 3:
+            x = x.squeeze(1)
+        x = F.rms_norm(x, (x.size(-1),))
+        new_states = []
+        for block, bstate in zip(self.blocks, states):
+            x, bstate = block.step(x, bstate)
+            new_states.append(bstate)
+        x = self.ln_out(x)
+        logits = self.lm_head(x.unsqueeze(1)).squeeze(1).float()
+        logits = 15.0 * torch.tanh(logits / 15.0)
+        return (logits, new_states)
+
+    @torch.no_grad()
+    def prefill_chunked(self, idx: torch.Tensor, chunk_size: int = 512):
+        B, T = idx.shape
+        states = [dict() for _ in self.blocks]
+        last_hidden = None
+        for c0 in range(0, T, chunk_size):
+            x = self.wte(idx[:, c0 : c0 + chunk_size])
+            x = F.rms_norm(x, (x.size(-1),))
+            for l, block in enumerate(self.blocks):
+                x, states[l] = block.forward_chunk(x, states[l])
+            last_hidden = x[:, -1:, :]
+        logits = self.lm_head(self.ln_out(last_hidden)).float()
+        logits = 15.0 * torch.tanh(logits / 15.0)
+        return (logits.squeeze(1), states)
+
+    @torch.no_grad()
+    def generate(
+        self,
+        prompt_ids: torch.Tensor,
+        max_new_tokens: int = 200,
+        temperature: float = 0.8,
+        top_k: int = 50,
+    ) -> torch.Tensor:
+        self.prepare_inference()
+        device = prompt_ids.device
+        states = self.init_state(batch_size=1, device=device)
+        logits = torch.zeros(prompt_ids.shape[0], self.cfg.vocab_size, device=device)
+        for i in range(prompt_ids.shape[1]):
+            tok = prompt_ids[:, i]
+            logits, states = self.step(tok, states)
+        generated = []
+        for _ in range(max_new_tokens):
+            if temperature == 0.0:
+                next_tok = logits.argmax(dim=-1)
+            else:
+                scaled = logits / temperature
+                if top_k > 0:
+                    topk_vals, _ = torch.topk(scaled, top_k)
+                    scaled = scaled.masked_fill(
+                        scaled < topk_vals[:, -1:], float("-inf")
+                    )
+                probs = torch.softmax(scaled, dim=-1)
+                next_tok = torch.multinomial(probs, num_samples=1).squeeze(1)
+            generated.append(next_tok)
+            logits, states = self.step(next_tok, states)
+        return torch.stack(generated, dim=1)
+
+    def forward(
+        self,
+        idx: torch.Tensor,
+        targets: Optional[torch.Tensor] = None,
+        reduction: str = "mean",
+    ) -> torch.Tensor:
+        B, T = idx.shape
+        assert T <= self.cfg.seq_len, f"Sequence {T} > max {self.cfg.seq_len}"
+        x = self.wte(idx)
+        x = F.rms_norm(x, (x.size(-1),))
+        if _CPU_HAS_BF16:
+            with torch.autocast(device_type="cpu", dtype=torch.bfloat16):
+                for block in self.blocks:
+                    x = block(x)
+        else:
+            for block in self.blocks:
+                x = block(x)
+        x = self.ln_out(x)
+        if targets is None:
+            logits = self.lm_head(x).float()
+            logits = 15.0 * torch.tanh(logits / 15.0)
+            return logits
+        return _chunked_cross_entropy(self.lm_head.weight, x, targets)
+
+
+def gpt2_small_config(**overrides) -> CPUGPTConfig:
+    cfg = CPUGPTConfig(
+        vocab_size=50257,
+        seq_len=1024,
+        n_layer=12,
+        n_embd=768,
+        n_head=12,
+        fno_modes=256,
+        gla_chunk=64,
+        ffn_hidden=2048,
+        layer_pattern="SSSL",
+    )
+    for k, v in overrides.items():
+        setattr(cfg, k, v)
+    return cfg
+
+
+def smoke_config(**overrides) -> CPUGPTConfig:
+    cfg = CPUGPTConfig(
+        vocab_size=50257,
+        seq_len=256,
+        n_layer=4,
+        n_embd=256,
+        n_head=4,
+        fno_modes=64,
+        gla_chunk=64,
+        ffn_hidden=512,
+        layer_pattern="SSSL",
+    )
+    for k, v in overrides.items():
+        setattr(cfg, k, v)
+    return cfg
+
+
+def gpt2_8b_config() -> CPUGPTConfig:
+    return CPUGPTConfig(
+        n_layer=32,
+        n_embd=4096,
+        n_head=32,
+        ffn_hidden=14336,
+        fno_modes=512,
+        gla_chunk=256,
+        seq_len=2048,
+        layer_pattern="SSSL",
+    )
+
+
+def gpt2_1b_config() -> CPUGPTConfig:
+    return CPUGPTConfig(
+        n_layer=24,
+        n_embd=2048,
+        n_head=16,
+        ffn_hidden=5632,
+        fno_modes=512,
+        gla_chunk=256,
+        seq_len=2048,
+        layer_pattern="SSSL",
+        vocab_size=50257,
+    )
+
+
+def byte125m_config(**overrides) -> CPUGPTConfig:
+    cfg = CPUGPTConfig(
+        vocab_size=261,
+        seq_len=16384,
+        n_layer=12,
+        n_embd=1024,
+        n_head=16,
+        fno_modes=256,
+        gla_chunk=512,
+        ffn_hidden=2816,
+        layer_pattern="SSSL",
+        fft_dtype="fp32",
+        use_flash_fft=True,
+    )
+    for k, v in overrides.items():
+        setattr(cfg, k, v)
+    return cfg
+
+
+def gpt2_1b_optimized_config() -> CPUGPTConfig:
+    return CPUGPTConfig(
+        n_layer=24,
+        n_embd=2048,
+        n_head=16,
+        ffn_hidden=5632,
+        fno_modes=512,
+        gla_chunk=512,
+        seq_len=2048,
+        layer_pattern="SSSL",
+        vocab_size=50257,
+        fft_dtype="bf16",
+        use_flash_fft=True,
+    )
+
+
+QWEN_CODER_VOCAB = 152064
+QWEN3_4B_VOCAB = 151936
+
+
+def code_3b_config(**overrides) -> CPUGPTConfig:
+    cfg = CPUGPTConfig(
+        vocab_size=QWEN_CODER_VOCAB,
+        seq_len=4096,
+        n_layer=28,
+        n_embd=3072,
+        n_head=24,
+        fno_modes=512,
+        gla_chunk=256,
+        ffn_hidden=8192,
+        layer_pattern="SSSL",
+        gla_delta=True,
+        sliding_window=0,
+    )
+    for k, v in overrides.items():
+        setattr(cfg, k, v)
+    return cfg
+
+
+def code_3b_exact_config(**overrides) -> CPUGPTConfig:
+    cfg = CPUGPTConfig(
+        vocab_size=QWEN_CODER_VOCAB,
+        seq_len=4096,
+        n_layer=28,
+        n_embd=3584,
+        n_head=28,
+        fno_modes=512,
+        gla_chunk=256,
+        ffn_hidden=18944,
+        layer_pattern="SSSL",
+        gla_delta=True,
+        sliding_window=0,
+    )
+    for k, v in overrides.items():
+        setattr(cfg, k, v)
+    return cfg
+
+
+def code_4b_config(**overrides) -> CPUGPTConfig:
+    cfg = CPUGPTConfig(
+        vocab_size=QWEN3_4B_VOCAB,
+        seq_len=4096,
+        n_layer=36,
+        n_embd=2560,
+        n_head=20,
+        fno_modes=512,
+        gla_chunk=256,
+        ffn_hidden=9728,
+        layer_pattern="SSSL",
+        gla_delta=True,
+        sliding_window=0,
+    )
+    for k, v in overrides.items():
+        setattr(cfg, k, v)
+    return cfg
+
+
+def code_1b_config(**overrides) -> CPUGPTConfig:
+    cfg = CPUGPTConfig(
+        vocab_size=QWEN_CODER_VOCAB,
+        seq_len=4096,
+        n_layer=24,
+        n_embd=2048,
+        n_head=16,
+        fno_modes=512,
+        gla_chunk=256,
+        ffn_hidden=5632,
+        layer_pattern="SSSL",
+        gla_delta=True,
+        sliding_window=0,
+    )
+    for k, v in overrides.items():
+        setattr(cfg, k, v)
+    return cfg
+
+
+def code_1p5b_exact_config(**overrides) -> CPUGPTConfig:
+    cfg = CPUGPTConfig(
+        vocab_size=151936,
+        seq_len=4096,
+        n_layer=28,
+        n_embd=1536,
+        n_head=12,
+        fno_modes=512,
+        gla_chunk=256,
+        ffn_hidden=8960,
+        layer_pattern="SSSL",
+        gla_delta=True,
+        sliding_window=0,
+    )
+    for k, v in overrides.items():
+        setattr(cfg, k, v)
+    return cfg
+
+
+def arc_a_swa512_config(**overrides) -> CPUGPTConfig:
+    cfg = code_3b_config(swa_window=512)
+    for k, v in overrides.items():
+        setattr(cfg, k, v)
+    return cfg
+
+
+def arc_c_samba_k6_config(**overrides) -> CPUGPTConfig:
+    cfg = code_3b_config(attn_layer_every=6, attn_full=False, attn_window=512)
+    for k, v in overrides.items():
+        setattr(cfg, k, v)
+    return cfg
+
+
+def arc_c_samba_k3_config(**overrides) -> CPUGPTConfig:
+    cfg = code_3b_config(attn_layer_every=3, attn_full=False, attn_window=512)
+    for k, v in overrides.items():
+        setattr(cfg, k, v)
+    return cfg
+
+
+def arc_ub_full4_config(**overrides) -> CPUGPTConfig:
+    cfg = code_3b_config(attn_layer_every=7, attn_full=True)
+    for k, v in overrides.items():
+        setattr(cfg, k, v)
+    return cfg
+
+
+def arc_d_bigstate_config(**overrides) -> CPUGPTConfig:
+    cfg = code_3b_config(gdn_expand_v=2.0)
+    for k, v in overrides.items():
+        setattr(cfg, k, v)
+    return cfg
+
+
+def arc_swadelta_config(**overrides) -> CPUGPTConfig:
+    cfg = code_3b_config(swa_fused_window=512)
+    for k, v in overrides.items():
+        setattr(cfg, k, v)
+    return cfg
+
+
+def code_smoke_config(**overrides) -> CPUGPTConfig:
+    cfg = CPUGPTConfig(
+        vocab_size=QWEN_CODER_VOCAB,
+        seq_len=256,
+        n_layer=4,
+        n_embd=256,
+        n_head=4,
+        fno_modes=64,
+        gla_chunk=64,
+        ffn_hidden=512,
+        layer_pattern="SSSL",
+        gla_delta=False,
+    )
+    for k, v in overrides.items():
+        setattr(cfg, k, v)
+    return cfg
+
+
+def get_config(name: str) -> CPUGPTConfig:
+    configs = {
+        "gpt2-small": gpt2_small_config,
+        "smoke": smoke_config,
+        "gpt2-8b": gpt2_8b_config,
+        "gpt2-1b": gpt2_1b_config,
+        "gpt2-1b-optimized": gpt2_1b_optimized_config,
+        "byte-125m": byte125m_config,
+        "code-3b": code_3b_config,
+        "code-3b-exact": code_3b_exact_config,
+        "code-4b": code_4b_config,
+        "code-1b": code_1b_config,
+        "code-1.5b-exact": code_1p5b_exact_config,
+        "code-smoke": code_smoke_config,
+        "arc-a-swa512": arc_a_swa512_config,
+        "arc-swadelta": arc_swadelta_config,
+        "arc-c-samba-k6": arc_c_samba_k6_config,
+        "arc-c-samba-k3": arc_c_samba_k3_config,
+        "arc-ub-full4": arc_ub_full4_config,
+        "arc-d-bigstate": arc_d_bigstate_config,
+    }
+    if name not in configs:
+        raise ValueError(f"Unknown config '{name}'. Available: {list(configs.keys())}")
+    return configs[name]()

model_int8.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ba877fe6d1bf500906eef6b7cbbd1fdc10127af8b5286b37d228db94314b2a8b
+size 2081695032

modeling.py

@@ -0,0 +1,230 @@
+from __future__ import annotations
+import json
+import os
+import time
+from typing import Optional
+
+os.environ.setdefault("CUDA_VISIBLE_DEVICES", "")
+os.environ.setdefault("HIP_VISIBLE_DEVICES", "")
+import torch
+import torch.nn as nn
+
+
+def _read_config(weights_dir: str) -> dict:
+    path = os.path.join(weights_dir, "config.json")
+    if not os.path.exists(path):
+        path = weights_dir if weights_dir.endswith(".json") else path
+    with open(path) as f:
+        return json.load(f)
+
+
+def _build_config(cfg_json: dict):
+    from model_cpu_gpt2 import CPUGPTConfig
+
+    return CPUGPTConfig(
+        vocab_size=cfg_json["vocab_size"],
+        seq_len=cfg_json.get("seq_len", 1024),
+        n_layer=cfg_json["n_layer"],
+        n_embd=cfg_json["n_embd"],
+        n_head=cfg_json["n_head"],
+        ffn_hidden=cfg_json["ffn_hidden"],
+        layer_pattern=cfg_json.get("layer_pattern", "SSSL"),
+        gla_delta=cfg_json.get("gla_delta", True),
+        fno_modes=cfg_json.get("fno_modes", 512),
+        gla_chunk=cfg_json.get("gla_chunk", 64),
+        landmark_layer_every=cfg_json.get("landmark_layer_every", 0),
+        landmark_chunk=cfg_json.get("landmark_chunk", 32),
+        landmark_max=cfg_json.get("landmark_max", 64),
+        attn_layer_every=cfg_json.get("attn_layer_every", 0),
+        dropout=0.0,
+    )
+
+
+class FelaLM:
+    def __init__(self, model, cfg, tokenizer, cfg_json):
+        self.model = model
+        self.cfg = cfg
+        self.tok = tokenizer
+        self.cfg_json = cfg_json
+
+        def _tid(name):
+            i = tokenizer.token_to_id(name)
+            return i if i is not None and i >= 0 else None
+
+        self.fim_prefix = _tid("<|fim_prefix|>")
+        self.fim_suffix = _tid("<|fim_suffix|>")
+        self.fim_middle = _tid("<|fim_middle|>")
+        self.fim_pad = _tid("<|fim_pad|>")
+        self.eot = _tid("<|endoftext|>")
+        self.fim_ok = None not in (self.fim_prefix, self.fim_suffix, self.fim_middle)
+        self._stops = {
+            t
+            for t in (
+                self.fim_prefix,
+                self.fim_suffix,
+                self.fim_middle,
+                self.fim_pad,
+                self.eot,
+            )
+            if t is not None
+        }
+
+    @torch.no_grad()
+    def complete(
+        self,
+        prefix: str,
+        suffix: str = "",
+        max_tokens: int = 40,
+        temperature: float = 0.0,
+        single_line: bool = True,
+    ) -> dict:
+        prefix = prefix or ""
+        suffix = suffix or ""
+        used_fim = bool(suffix.strip()) and self.fim_ok
+        if used_fim:
+            ids = (
+                [self.fim_prefix]
+                + self.tok.encode(prefix).ids
+                + [self.fim_suffix]
+                + self.tok.encode(suffix).ids
+                + [self.fim_middle]
+            )
+        else:
+            ids = self.tok.encode(prefix).ids
+        if not ids:
+            ids = [self.eot] if self.eot is not None else [0]
+        t0 = time.perf_counter()
+        states = self.model.init_state(batch_size=1)
+        logits = None
+        for tok_id in ids:
+            logits, states = self.model.step(
+                torch.tensor([tok_id], dtype=torch.long), states
+            )
+        prefill_ms = (time.perf_counter() - t0) * 1000.0
+        out_ids = []
+        td = time.perf_counter()
+        for _ in range(max_tokens):
+            if temperature and temperature > 0:
+                probs = torch.softmax(logits.float().reshape(-1) / temperature, -1)
+                nxt = int(torch.multinomial(probs, 1).item())
+            else:
+                nxt = int(logits.float().reshape(-1).argmax().item())
+            if nxt in self._stops:
+                break
+            out_ids.append(nxt)
+            piece = self.tok.decode(out_ids)
+            if single_line and "\n" in piece:
+                break
+            logits, states = self.model.step(
+                torch.tensor([nxt], dtype=torch.long), states
+            )
+        decode_ms = (time.perf_counter() - td) * 1000.0
+        text = self.tok.decode(out_ids) if out_ids else ""
+        if single_line:
+            text = text.split("\n", 1)[0]
+        n = len(out_ids)
+        return {
+            "middle": text,
+            "n_tokens": n,
+            "used_fim": used_fim,
+            "prompt_tokens": len(ids),
+            "prefill_ms": round(prefill_ms, 1),
+            "decode_ms": round(decode_ms, 1),
+            "tok_per_s": round(n / (decode_ms / 1000.0), 2)
+            if decode_ms > 0 and n
+            else 0.0,
+        }
+
+
+def _read_bf16_state(weights_dir: str) -> dict:
+    from safetensors import safe_open
+
+    st = {}
+    path = os.path.join(weights_dir, "model.safetensors")
+    with safe_open(path, framework="pt", device="cpu") as f:
+        for k in f.keys():
+            st[k] = f.get_tensor(k).float()
+    return st
+
+
+def _read_int8_state(weights_dir: str) -> dict:
+    from safetensors import safe_open
+
+    st = {}
+    path = os.path.join(weights_dir, "model_int8.safetensors")
+    with safe_open(path, framework="pt", device="cpu") as f:
+        keys = list(f.keys())
+        for k in keys:
+            if k.startswith("keep."):
+                st[k[len("keep.") :]] = f.get_tensor(k).float()
+        for k in keys:
+            if k.startswith("int8."):
+                base = k[len("int8.") :]
+                w = f.get_tensor(k).float()
+                s = f.get_tensor("scale." + base).float()
+                st[base] = w * s.reshape([-1] + [1] * (w.dim() - 1))
+    return st
+
+
+def _apply_state(model, st: dict) -> None:
+    params = dict(model.named_parameters())
+    params.update(dict(model.named_buffers()))
+    keys = set(st)
+    for k in keys:
+        dst = params.get(k)
+        if dst is None:
+            raise KeyError(f"Checkpoint key {k!r} has no home in the model")
+        with torch.no_grad():
+            dst.copy_(st[k])
+    missing = set(params) - keys
+    if missing:
+        raise KeyError(f"Missing {len(missing)} params, e.g. {sorted(missing)[:5]}")
+
+
+def _resolve_quant(weights_dir: str, quant: str) -> str:
+    if quant == "auto":
+        has_int8 = os.path.exists(os.path.join(weights_dir, "model_int8.safetensors"))
+        has_bf16 = os.path.exists(os.path.join(weights_dir, "model.safetensors"))
+        return "bf16" if has_bf16 else ("int8" if has_int8 else "bf16")
+    return quant
+
+
+def load_model(
+    weights_dir: str = ".", threads: Optional[int] = None, quant: str = "bf16"
+) -> FelaLM:
+    from model_cpu_gpt2 import CPUGPT
+    from cpu_patch import enable_cpu_delta
+    from tokenizers import Tokenizer
+
+    if threads:
+        torch.set_num_threads(threads)
+    if weights_dir.endswith(".safetensors"):
+        weights_dir = os.path.dirname(os.path.abspath(weights_dir)) or "."
+    cfg_json = _read_config(weights_dir)
+    cfg = _build_config(cfg_json)
+    model = CPUGPT(cfg)
+    model.lm_head = nn.Linear(cfg.n_embd, cfg.vocab_size, bias=False)
+    quant = _resolve_quant(weights_dir, quant)
+    if quant == "int8":
+        st = _read_int8_state(weights_dir)
+    else:
+        st = _read_bf16_state(weights_dir)
+    _apply_state(model, st)
+    model.eval()
+    enable_cpu_delta(model)
+    model.prepare_inference()
+    tok_path = os.path.join(weights_dir, "tokenizer.json")
+    tokenizer = Tokenizer.from_file(tok_path)
+    return FelaLM(model, cfg, tokenizer, cfg_json)
+
+
+def from_pretrained(repo_id: str = "lowdown-labs/FELA-autocomplete") -> FelaLM:
+    from huggingface_hub import hf_hub_download
+
+    d = os.path.dirname(hf_hub_download(repo_id, "config.json"))
+    hf_hub_download(repo_id, "model.safetensors")
+    hf_hub_download(repo_id, "tokenizer.json")
+    hf_hub_download(repo_id, "model_cpu_gpt2.py")
+    for f in ("cpu_delta.py", "cpu_landmark.py", "cpu_swa.py", "cpu_patch.py"):
+        hf_hub_download(repo_id, f)
+    return load_model(d)

modeling_fela.py

@@ -0,0 +1,84 @@
+import os
+import sys
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+import torch
+import torch.nn as nn
+from transformers import PreTrainedModel
+from transformers.modeling_outputs import CausalLMOutput
+
+from .configuration_fela import FelaConfig
+from .cpu_delta import CPUGatedDeltaNet as _cd
+from .cpu_landmark import CPULandmark as _cl
+from .cpu_swa import CPUSlidingWindow as _cs
+from .model_cpu_gpt2 import CPUGPT, CPUGPTConfig
+from .cpu_patch import enable_cpu_delta
+
+_KEEP = (_cd, _cl, _cs)
+
+
+class FelaForCausalLM(PreTrainedModel):
+    config_class = FelaConfig
+    base_model_prefix = "model"
+    _tied_weights_keys = []
+    _no_split_modules = []
+
+    def __init__(self, config):
+        super().__init__(config)
+        cfg = CPUGPTConfig(
+            vocab_size=config.vocab_size,
+            seq_len=config.seq_len,
+            n_layer=config.n_layer,
+            n_embd=config.n_embd,
+            n_head=config.n_head,
+            ffn_hidden=config.ffn_hidden,
+            layer_pattern=config.layer_pattern,
+            gla_delta=config.gla_delta,
+            fno_modes=config.fno_modes,
+            gla_chunk=config.gla_chunk,
+            landmark_layer_every=config.landmark_layer_every,
+            landmark_chunk=config.landmark_chunk,
+            landmark_max=config.landmark_max,
+            attn_layer_every=config.attn_layer_every,
+            dropout=0.0,
+        )
+        self.model = CPUGPT(cfg)
+        self.model.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
+        self._prepared = False
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.model.wte
+
+    def set_input_embeddings(self, value):
+        self.model.wte = value
+
+    def get_output_embeddings(self):
+        return self.model.lm_head
+
+    def set_output_embeddings(self, value):
+        self.model.lm_head = value
+
+    def _ensure_prepared(self):
+        if not self._prepared:
+            enable_cpu_delta(self.model)
+            self.model.prepare_inference()
+            self._prepared = True
+
+    def forward(
+        self, input_ids=None, attention_mask=None, labels=None, use_cache=None, **kwargs
+    ):
+        self._ensure_prepared()
+        logits = self.model(input_ids)
+        loss = None
+        if labels is not None:
+            sl = logits[..., :-1, :].contiguous()
+            lb = labels[..., 1:].contiguous()
+            loss = nn.functional.cross_entropy(sl.view(-1, sl.size(-1)), lb.view(-1))
+        return CausalLMOutput(loss=loss, logits=logits)
+
+    def prepare_inputs_for_generation(self, input_ids, **kwargs):
+        return {"input_ids": input_ids}
+
+    def can_generate(self):
+        return True

quickstart/README.md

@@ -0,0 +1,47 @@
+# Quickstart
+
+Load the FELA LLM 1.5 code model and run a real fill in the middle completion on CPU.
+Uses the self contained loader in `../modeling.py`.
+
+## Steps
+
+1. Install the pinned requirements (CPU PyTorch):
+
+   ```bash
+   pip install -r requirements.txt
+   ```
+
+2. Point the loader at the directory that holds the weights (this repo). The model
+   ships as `model.safetensors` next to `config.json` and `tokenizer.json`:
+
+   ```bash
+   export FELA_LLM_WEIGHTS=/path/to/this/repo
+   python run.py
+   ```
+
+## Load from Python
+
+```python
+from modeling import load_model
+m = load_model("/path/to/this/repo")
+
+# Fill in the middle: give the code before and after the cursor, get the middle
+r = m.complete("def add(a, b):\n    ", suffix="\nresult = add(2, 3)\n")
+print(r["middle"])          # e.g. " return a + b"
+
+# Plain autocomplete: continue a single line
+print(m.complete("import numpy as ")["middle"])   # e.g. " np"
+```
+
+`complete` is greedy by default, so the same prompt gives the same real output every
+time. When you pass a non empty `suffix` and the tokenizer has the FIM tokens, it uses
+the fill in the middle layout `<|fim_prefix|> P <|fim_suffix|> S <|fim_middle|>` and
+writes the code that belongs between them; otherwise it continues the prefix.
+
+## What to expect
+
+This is the final fill in the middle model. Single line patterns land well (imports,
+obvious returns, boilerplate); multi line blocks and novel logic are outside what it is
+built for. Every completion is a genuine forward of the model, never a lookup. The
+reference load path here runs on CPU with no GPU. For fast production serving the same
+weights run int8 on the CPU native FELA server.

quickstart/requirements.txt

@@ -0,0 +1,5 @@
+torch==2.8.0
+safetensors==0.5.3
+tokenizers==0.22.2
+huggingface-hub==0.34.4
+flash-linear-attention==0.5.1

quickstart/run.py

@@ -0,0 +1,44 @@
+import argparse
+import os
+import sys
+import torch
+
+sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), ".."))
+from modeling import load_model
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument(
+        "--weights",
+        default=os.environ.get("FELA_LLM_WEIGHTS", ".."),
+        help="Dir with model.safetensors, config.json, tokenizer.json",
+    )
+    ap.add_argument("--threads", type=int, default=4)
+    args = ap.parse_args()
+    torch.set_num_threads(args.threads)
+    print("Loading FELA LLM 1.5 on CPU (no GPU needed)...")
+    m = load_model(args.weights, threads=args.threads)
+    print(
+        f"Loaded {m.cfg_json.get('n_params', 0) / 1000000000.0:.2f}B params, fim={m.fim_ok}\n"
+    )
+    print("Fill in the middle: prefix + suffix -> the model writes the middle")
+    prefix = "def add(a, b):\n    "
+    suffix = "\nresult = add(2, 3)\n"
+    r = m.complete(prefix, suffix, max_tokens=16)
+    print(f"  Prefix: {prefix!r}")
+    print(f"  Suffix: {suffix!r}")
+    print(
+        f"  Middle: {r['middle']!r}  (fim={r['used_fim']}, {r['n_tokens']} tokens, {r['tok_per_s']} tok/s)\n"
+    )
+    print("Plain autocomplete: continue a single line")
+    for prefix in ["import numpy as ", "from fastapi import ", "for i in range("]:
+        r = m.complete(prefix, "", max_tokens=12)
+        print(f"  {prefix!r} -> {r['middle']!r}")
+    print(
+        "\nThis is the final fill in the middle model. Single line completions land well; multi line blocks and novel logic are outside what it is built for. Every completion is a real forward of the model."
+    )
+
+
+if __name__ == "__main__":
+    main()

requirements.txt

@@ -0,0 +1,5 @@
+torch==2.8.0
+safetensors==0.5.3
+tokenizers==0.22.2
+huggingface-hub==0.34.4
+flash-linear-attention==0.5.1

space/README.md

@@ -0,0 +1,28 @@
+---
+title: FELA LLM 1.5 code completion
+colorFrom: indigo
+colorTo: purple
+sdk: gradio
+sdk_version: 4.44.0
+app_file: app.py
+pinned: false
+license: other
+---
+
+# FELA LLM 1.5 code completion
+
+Type the code before your cursor (Prefix) and optionally the code after it (Suffix), and
+the model writes the line that goes in between. It runs on a plain CPU with no GPU.
+
+## How it works
+
+The model reads your prefix and suffix, and when the suffix is present it uses the fill
+in the middle layout so the completion fits between them. It writes one line at a time,
+greedy by default, so the same prompt gives the same real output every time.
+
+## What to expect
+
+This is the final fill in the middle model. Single line patterns land well (imports,
+obvious returns, boilerplate); multi line blocks and novel logic are outside what it is
+built for. Every completion is a genuine forward of the model, never a lookup. For fast
+production serving the same weights run int8 on the CPU native FELA server.

space/app.py

@@ -0,0 +1,44 @@
+import os
+import sys
+import gradio as gr
+import torch
+
+sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), ".."))
+from modeling import load_model
+
+WEIGHTS = os.environ.get("FELA_LLM_WEIGHTS", "..")
+try:
+    torch.set_num_threads(4)
+    MODEL = load_model(WEIGHTS, threads=4)
+except Exception as e:
+    MODEL = None
+    print(f"[Note] Model not loaded ({e}); the Space reports the input instead.")
+
+
+def complete(prefix, suffix, max_tokens):
+    if MODEL is None:
+        return "Weights not found. Set FELA_LLM_WEIGHTS to the repo directory."
+    if not (prefix or "").strip() and (not (suffix or "").strip()):
+        return "Type some code before the cursor (prefix), and optionally after it (suffix)."
+    r = MODEL.complete(prefix or "", suffix or "", max_tokens=int(max_tokens))
+    tag = "fill in the middle" if r["used_fim"] else "continuation"
+    return f"{r['middle']}\n\n[{tag}, {r['n_tokens']} tokens, {r['tok_per_s']} tok/s on CPU]"
+
+
+with gr.Blocks(title="FELA LLM 1.5 code completion") as demo:
+    gr.Markdown(
+        "# FELA LLM 1.5 code completion\nA CPU native code autocomplete model. Type the code before your cursor in Prefix, and optionally the code after your cursor in Suffix, and it writes the line that goes in between. It runs on a plain CPU with no GPU. This is the final fill in the middle model: single line patterns land well, multi line blocks and novel logic are outside what it is built for. Every completion is a real forward of the model."
+    )
+    prefix = gr.Textbox(
+        label="Prefix (code before the cursor)", value="def add(a, b):\n    ", lines=4
+    )
+    suffix = gr.Textbox(
+        label="Suffix (code after the cursor, optional)",
+        value="\nresult = add(2, 3)\n",
+        lines=2,
+    )
+    max_tokens = gr.Slider(4, 64, value=16, step=1, label="Max tokens")
+    out = gr.Textbox(label="Completed middle", lines=4)
+    gr.Button("Complete").click(complete, [prefix, suffix, max_tokens], out)
+if __name__ == "__main__":
+    demo.launch()

space/requirements.txt

@@ -0,0 +1,6 @@
+gradio==4.44.0
+torch==2.8.0
+safetensors==0.5.3
+tokenizers==0.22.2
+huggingface-hub==0.34.4
+flash-linear-attention==0.5.1

verify.py

@@ -0,0 +1,39 @@
+import argparse
+import sys
+import torch
+from modeling import load_model
+
+VERIFICATION_TOP5 = [319, 516, 22441, 497, 4102]
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--weights", default=".")
+    ap.add_argument("--capture", action="store_true")
+    args = ap.parse_args()
+    torch.set_num_threads(1)
+    m = load_model(args.weights, threads=1)
+    ids = torch.arange(16, dtype=torch.long).unsqueeze(0) + 100
+    with torch.no_grad():
+        logits = m.model(ids)
+    assert logits.dim() == 3 and logits.size(-1) == m.cfg.vocab_size, (
+        f"Bad output shape {tuple(logits.shape)} (expected vocab {m.cfg.vocab_size})"
+    )
+    assert torch.isfinite(logits).all(), "Non finite logits"
+    print(f"Output shape ok {tuple(logits.shape)}")
+    top5 = logits[0, -1].topk(5).indices.tolist()
+    if args.capture:
+        print(f"VERIFICATION_TOP5 = {top5}")
+        return
+    if top5 != VERIFICATION_TOP5:
+        print(
+            f"[WARNING] top 5 tokens {top5} differ from the reference {VERIFICATION_TOP5}; small CPU numeric differences can move ties across machines"
+        )
+    else:
+        print("Verification value ok")
+    r = m.complete("import numpy as ", "", max_tokens=4)
+    print(f"Sample completion: 'import numpy as ' -> {r['middle']!r}")
+
+
+if __name__ == "__main__":
+    main()