flash-head 0.1.0__tar.gz

flash_head-0.1.0/LICENSE

@@ -0,0 +1,117 @@
+EMBEDL MODELS COMMUNITY LICENSE AGREEMENT
+v.1.0
+
+Embedl is pleased to make available Embedl-optimized AI models and accompanying tools, SDKs, and source code files under this Community License Agreement (the "License" or the "Terms"). Our goal is to foster a vibrant community of developers, researchers, and organizations who can freely use, modify, and build upon these optimized AI models and tools. To that end, this license permits unlimited use, distribution and commercialization for most individuals and organizations, subject only to the limited restrictions described below.
+
+To sustain our ability to continue developing and releasing optimized AI models, we have established three limitations on use. First, competitors to Embedl may not use the licensed materials, as permitting such use would undermine our ability to invest in future development. Second, larger organizations—those with more than 250 employees or annual revenue exceeding €10 million—can use the licensed materials for internal evaluation purposes, but must obtain a separate Commercial Use license before using the licensed materials for commercial purposes; this ensures that entities with greater resources contribute appropriately to the ecosystem. Third, no licensee may host the licensed materials as a service, which protects both Embedl's interests and the integrity of the community by preventing commoditization of these models (although specific use cases are possible with prior approval from us).
+
+We recognize that individual circumstances vary, and we welcome inquiries from those who wish to discuss alternative licensing arrangements that fall outside these terms. Please contact Embedl at  legal@embedl.com to discuss your specific use case. We would be happy to work with you.
+
+    1.0 DEFINED TERMS
+
+"Commercial Use": (a) any provision of the functionality of the Embedl Models, in whole or in part, in any product, offering, consultancy, or as a service; (b) any Distribution; (c) any use of the Embedl Models in whole or in part for direct or indirect strategic or commercial gain or advantage; or (d) any use the Embedl Models or any of its outputs or results, in whole or in part, for purposes of training or improving any other large language model.
+
+"Competitor": any person or entity in the business of selling access to AI development and optimization tools.
+
+"Contribution": any contribution to the Embedl Models that is intentionally submitted for inclusion in the Embedl Models by the contribution owner (copyright and otherwise), or by a person or entity authorized to submit the contribution on behalf of such owner. For the purposes of this definition, "submitted" means any form of electronic, verbal, or written communication sent to us, including but not limited to communication on electronic mailing lists, source code control systems, and issue tracking systems, but excluding communication that is conspicuously marked or otherwise designated in writing by the contribution owner as "Not a Contribution."
+
+"Contributor": any person or entity submitting a Contribution that has been accepted by us and incorporated within the Embedl Models.
+
+"Distribution" or "Distribute":  means any transmission, distribution, deployment, publication, or other sharing of the Embedl Models in whole or in part, whether standalone or as incorporated or bundled with other works, services, or products, to any third party, including by providing or making the Embedl Models or its functionality available as a hosted service via API, web access, or any other electronic or remote means.
+
+"Embedl Models":  the AI models, software, algorithms, code, and weights, licensed by us under these Terms, as any of the foregoing may be updated or modified by us on one or more occasions.
+
+"Third Party Terms": terms and policies applicable to Llama, Gemma, and Qwen, including terms currently available at https://github.com/meta-llama/llama-models/blob/main/models/llama3_2/LICENSE, https://ai.google.dev/gemma/terms, and https://www.apache.org/licenses/LICENSE-2.0 (for Qwen materials), as well as the policies set forth in section 3.9.
+
+"We," "us," "our," or "Embedl": Embedl AB, organized under the laws of Sweden.
+
+"You" or "your": the person, and any entity such person works for or on behalf of, exercising the rights of this license, excluding any Competitor. 
+
+2.0	LICENSE RIGHTS
+
+2.1	License Grant. Subject to your compliance with section 3.0, including the Commercial Use and Competitor limitations below, and the Third Party Terms, we hereby grant to you a worldwide, royalty-free, sublicensable, non-transferable, and non-exclusive copyright license to use, reproduce, distribute, display, perform, modify, and create derivative works of the Embedl Models in whole or in part. 
+
+2.2	Contributor License Grant. Subject to your compliance with these Terms, including the Commercial Use and Competitor limitations below, and the Third Party Terms, each Contributor hereby grants to you a worldwide, non-exclusive, royalty-free, sublicensable, and non-transferable (a) copyright license to use, reproduce, distribute, display, perform, modify, and create derivative works of the Embedl Models in whole or in part; and (b) patent license to make, have made, use, offer to sell, sell, import and otherwise transfer the Embedl Models in whole or in part, but only to the extent of patent claims licensable by such Contributor that are necessarily infringed by the exercise of such rights, including for standalone Contributions and Contributions in combination with the Embedl Models.
+
+2.3	Reservation of Rights.  All rights not expressly granted in this section are reserved to us. You will have no right or license to the Embedl Models other than the rights set forth in this section 2.0, and no other right or license will be implied, by estoppel, conduct, or otherwise. If you are an entity, you may exercise the rights granted in this section solely via your own employees and contractors. We and our licensors retain all right, title and interest in the Embedl Models and associated intellectual property rights, and all copies of the Embedl Models in whole or in part, and we grant no rights to our patents.
+
+3.0	LICENSE CONDITIONS
+
+3.1	No Use by Competitors. If you are, or your employer or client for whom you work is, a Competitor, you are not licensed to download, license, install or use the Embedl Models. You represent and warrant that you are not a Competitor and are not working for a Competitor. If you are in breach of this warranty, now or in the future, you are not licensed to use the Embedl Models; doing so anyway constitutes breach of these Terms and infringement of our intellectual property rights.
+
+3.2	Commercial Use Limitation.
+
+a)	The licensed rights set forth in section 2.0 exclude Commercial Use if you are an entity, or an affiliate or member of a company group, that has more than 250 employees or annual revenue of €10,000,000 or more ("Large Entity"). If you are or become a Large Entity or part of a group of one, Commercial Use is not or is no longer licensed to you under these Terms; you must contact us for a Commercial Use license. You represent and warrant that you are not a Large Entity and are not working for a Large Entity. If you are in breach of this warranty, now or in the future, you are not licensed to use the Embedl Models for Commercial Use; doing so anyway constitutes breach of these Terms and infringement of our intellectual property rights.
+
+b)	In addition to any other remedies, if you violate section 3.2(a), you hereby grant to Embedl on behalf of yourself and each of your affiliates, employees and contractors, a non-exclusive, worldwide, perpetual, irrevocable, non-terminable, sublicensable (through multiple tiers), fully paid-up, royalty-free, assignable and transferable intellectual property rights license as to Distributions, to make, use, reproduce, modify, adapt, create derivative works of, improve, extend, enhance, translate, distribute (directly and indirectly, in any medium, under terms of choice), combine, compile, transmit, display and perform publicly, license, rent, lease, and manufacture, sell, offer to sell, and import, your Distributions, for any commercial or non-commercial purpose.
+
+3.3	No Use as a Hosted Service. You may not, without our explicit permission, use the Embedl Models, any derivative works thereof, or any models created or trained via the use or incorporation of the Embedl Models, to provide a hosted service, software-as-a-service offering, API, or any other mechanism that allows third parties to access or use the functionality of such materials, or to submit inputs to, receive outputs from, the Embedl Models, whether or not for a fee. Notwithstanding the foregoing, you may deploy the Embedl Models on an internal network accessible solely to your employees and contractors performing work exclusively for you, solely for your internal business purposes.
+
+In addition, we are open to hosted use in support of a specific function or feature in a separate application, depending on the context, on a case-by-base basis. Please contact us, and we would be happy to explore that.
+
+3.4	Required Notices and Pass-Throughs. 
+
+a)	You will preserve and not suppress proprietary notices, markings, and branding visible in or associated with or displayed via the use of the Embedl Models, and in any Distributions.
+
+b)	You will give each recipient of your Distribution a copy of these Terms and the following notice: "Optimized Embedl models are licensed under the Embedl Models Community Source License v.1 available at [___], and the terms of Embedl's licensors. Copyright © Embedl AB."
+
+c)	You must cause any modified files in your Distributions to carry prominent notices explaining the changes that were made. 
+
+d)	You must cause each Distribution to prominently display the following notices, as applicable (depending on which models are selected and included in any Distribution):
+
+Llama 3.2 is licensed under the Llama 3.2 Community License, Copyright © Meta Platforms, Inc. All Rights Reserved. Use is subject to applicable terms and policies, including terms currently available at https://github.com/meta-llama/llama-models/blob/main/models/llama3_2/LICENSE
+
+Gemma is provided under and subject to the Gemma Terms of Use found at https://ai.google.dev/gemma/terms.
+
+Qwen is licensed under the Apache License 2.0, Copyright © Alibaba Cloud. All Rights Reserved.
+
+e) 	For Llama-based Distributions: (a) you must prominently display "Built with Llama" on a related website, user interface, blogpost, about page, or product documentation; and (b) if you use the Embedl Models or any outputs or results of the Embedl Models to create, train, fine tune, or otherwise improve any AI model that is distributed or made available, you shall also include "Llama" at the beginning of any such AI model name.
+
+f)	You must require each recipient to a Distribution to agree in writing (or through another mechanism establishing enforceable assent) to be bound by the terms and restrictions of this section 3.0 and section 5.2, and to impose the same requirement on any subsequent recipients, so that each recipient in the chain of distribution and/or use is bound by such terms and restrictions.
+
+3.5	Contributions. Each Contributor hereby grants to us a worldwide, non-exclusive, royalty-free, sublicensable, and non-transferable (a) copyright license to use, reproduce, distribute, display, perform, modify, and create derivative works of each Contribution; and (b) patent license to make, have made, use, offer to sell, sell, import and otherwise transfer the Embedl Models, but only to the extent of patent claims licensable by such Contributor that are necessarily infringed by the exercise of such rights, for standalone Contributions or in Contributions in combination with the Embedl Models. By submitting a Contribution, you represent and warrant that you have sufficient rights to grant the licenses under these Terms, and that doing so does not violate any obligation to a third party.
+
+3.6	Our Marks. You agree that any use of our marks, branding and logos, including Embedl Models ("Marks"), whether permitted or otherwise, will inure to our sole benefit. You will not directly or indirectly: (a) file or prepare any application for registration of any Marks; (b) assert any right, title, license to, or interest in the Marks; or (c) adopt, use, file for registration, or register, in whole or in part, any trademark, service mark, trade name, logo, or domain name which may be confusingly similar to or an infringement of the Marks or any of our domain names.
+
+3.7	Feedback. You agree that we may freely use and exploit in perpetuity any feedback, requirements, recommendations, ideas, bug fixes, reviews, ratings, comments, suggestions, or improvements, that you, or any employee or agent thereof, may at any time disclose or submit to us relating to the Embedl Models for our business purposes, including for product, services or solution sales, licensing, support and development, without any obligation or payment to you. 
+
+3.8	Unauthorized Use. You will not use the Embedl Models or participate in any activities via the Embedl Models in a manner that is likely to be prohibited by law or these Terms or violative of third party rights in any applicable jurisdiction, including intellectual property and data privacy rights. Your use of the Embedl Models must be in full compliance with applicable law and the Third Party Terms. You acknowledge that use of the Embedl Models may be subject to the EU AI Act and other AI-specific laws and regulations. You are solely responsible for determining whether your use or Distribution constitutes a high-risk AI system or is otherwise subject to regulatory obligations, and for compliance with all applicable requirements.
+
+3.9	Acceptable Use. You hereby accept and shall fully comply with the following Acceptable Use Policies:
+
+a) for the Llama-derived materials: https://www.llama.com/llama3_2/use-policy; and
+
+b) for the Gemma-derived materials: https://ai.google.dev/gemma/prohibited_use_policy.
+
+3.10	Third Party Restrictions. You acknowledge that certain components derived from Gemma may be subject to remote restriction by Google if Google believes usage violates the Gemma Terms of Use. To the maximum extent permitted by law, we are not liable for any such restriction.
+
+3.11	Notifications to Us. If you believe that you are entitled or obligated to act contrary to these Terms under any mandatory or applicable law, you agree to provide us with detailed and substantiated explanation of your reasons in writing at least thirty days before you act, to allow us to assess whether we may, at our sole discretion, provide an alternative remedy for the situation, though we are under no obligation to do so. To be effective, notices to us must be sent to legal@embedl.com.
+
+4.0	DISCLAIMER OF WARRANTIES; LIMITATION OF LIABILITY
+
+4.1	Disclaimer. TO THE MAXIMUM EXTENT POSSIBLE UNDER APPLICABLE LAW, we provide the Embedl Models on an as-is, as-available basis with all faults, and WE DISCLAIM ALL WARRANTIES OF ANY KIND WITH RESPECT TO THE EMBEDL MODELS, WHETHER EXPRESS OR IMPLIED, INCLUDING, BUT NOT LIMITED TO THE IMPLIED WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE, TITLE AND NON-INFRINGEMENT. Specifically, we make no warranty that (a) the Embedl Models will meet your requirements, goals or needs, or (b) any errors or deficiencies in the Embedl Models will be corrected.
+
+4.2	Exclusion. Notwithstanding any other provision of these Terms, we will not be liable for any losses or damages, whether direct, indirect, punitive, special, incidental or consequential damages, or liable for interruption of business, diminution of value, cost of replacement, downtime, or loss of profits, revenue, use, data, or other economic advantage, in connection with, related to or arising out of the Embedl Models or these Terms, regardless of the theory of liability, whether for breach of these Terms, including breach of warranty, or in tort or otherwise, even if we have been previously advised of the possibility of such damages and even if such damages are reasonably foreseeable. Liability for damages will be so limited and excluded, regardless of the validity or efficacy of any remedy provided herein and even if any remedy fails of its essential purpose. Nothing in these Terms excludes or limits liability for death or personal injury caused by negligence, fraud, intentional misconduct, or gross negligence, or any other liability that cannot be excluded under applicable law.
+
+4.3	Indemnification. You hereby agree to defend, indemnify, release and hold harmless Embedl, and its officers, directors, employees, consultants, agents, and representatives, in full from and against any and all claims (including third party claims), losses, liability, damages, and/or costs of every kind and nature, arising out of or in any way connected with: (a) your access to or use of, the Embedl Models or conduct engendered thereby, (b) your breach of these Terms or the Third Party Terms, (c) your infringement or misappropriation of any intellectual property right or other right of any person or entity, or (d) your violations of applicable law.
+
+5.0	TERMINATION
+
+5.1	Upon Breach. The licenses granted by these Terms will terminate automatically if you are in breach of these Terms, including if you exceed the scope of the licenses granted in section 2.0. Breach of these Terms may result in pursuit of all available remedies for intellectual property rights (including intellectual property rights infringement), the availability of which you hereby acknowledge.
+
+5.2	Upon Litigation. The license rights set forth in section 2.0 will terminate automatically if you institute litigation (including a cross-claim or counterclaim in a lawsuit) against any entity alleging that the use, distribution, modification, sale, offer to sell, transfer or import of the Embedl Models, any Contribution (alone or in combination), or the outputs or results of the Embedl Models, constitutes direct, indirect, vicarious, or contributory infringement of intellectual property rights. 
+
+5.3	By Us. We may on written notice terminate rights granted under these Terms in the event the Embedl Models become subject to third party claims of infringement or violation of intellectual property or other rights or applicable law. 
+
+5.4	Effect. Upon termination of rights for any reason, all licensed rights granted under these Terms will terminate, and you will immediately cease all use, and destroy all and not retain any copies, of the Embedl Models. Our rights under this section 5.0 are in addition to any other rights and remedies permitted by law or under these Terms.
+
+6.0	GENERAL PROVISIONS
+
+6.1	Governing Law. These Terms and any disputes between us and related to or concerning these Terms or the Embedl Models (including tort as well as contract claims, and whether pre-contractual or extra-contractual), notwithstanding the choice of laws rules of any jurisdiction to the contrary, will be governed by the procedural and substantive laws of Sweden. The United Nations Convention on Contracts for the International Sale of Goods will not apply and is hereby excluded. 
+
+6.2 	Arbitration. Subject to section 6.3, any disputes between or claims brought by either party arising out of or related to the Embedl Models or these Terms shall be referred to and finally settled exclusively by binding arbitration before the Arbitration Institute of the Stockholm Chamber of Commerce (the "SCC"), and the venue for arbitration will be Stockholm, Sweden. The Rules for Expedited Arbitrations will apply, unless the SCC in its discretion determines, taking into account the complexity of the case, the amount in dispute and other circumstances, that the Arbitration Rules will apply. In the latter case, the SCC will also decide whether the Arbitral Tribunal will be composed of one or three arbitrators. The arbitration will be conducted by telephone, on-line and/or based solely upon written submissions where no in-person appearance is required. If in-person appearance is required, such hearings will be held in Stockholm, Sweden. The arbitrator will apply the law specified in section 6.1. All awards may if necessary be enforced by any court having jurisdiction. The existence of any dispute, the existence or details of the arbitration proceeding, and all related documents, materials, evidence, judgments and awards therein, shall be kept confidential. Except as required by law, no party will make any public announcements with respect to the other party or the proceeding or the award, except as required to enforce same. All disputes will be arbitrated only on an individual basis and not in a class, consolidated or representative action. The arbitrator does not have the power to vary these provisions. Subject to section 6.3, all claims between the parties must be resolved via arbitration in accordance with this section. Should either party file an action contrary to this section, the other party may recover lawyers' fees and costs associated with enforcing this section, provided that the party seeking the award has notified the other party in writing of the improperly filed claim, and the other party has failed to withdraw the claim in a timely fashion. 
+
+6.3	Injunctive Relief; Prevailing Party. You acknowledge and agree that breach of these Terms, or any unauthorized use or distribution of the Embedl Models, may cause irreparable harm to us, the extent of which would be difficult to ascertain, and that we will be entitled to seek immediate injunctive relief and specific performance (in addition to any other available remedies), and to enforce and/or seek redress for infringement of intellectual property rights, in any court of competent jurisdiction under the applicable laws thereto (and such proceeding, and our right to prosecute such a claim, is not subject to arbitration). You hereby consent to the jurisdiction of such courts for any such actions. A party prevailing in any litigation or arbitration related to these Terms or the Embedl Models will be entitled, in addition to such other relief as may be granted, to an award of reasonable attorneys' fees. This section is in addition and without prejudice to the availability of attorneys' fees as a remedy under applicable law, including laws pertaining to intellectual property rights.
+
+6.4	Important Final Terms. The Terms constitute the entire agreement between you and us and govern your use of the Embedl Models, superseding any prior agreements, understandings, communications or proposals; except that these Terms do not override, supersede, or modify the terms of any written license agreement signed by you and us. If any provision of the Terms is found by a court of competent jurisdiction to be invalid, the parties nevertheless agree that the court should endeavor to give effect to the parties' intentions as reflected in the provision, and the other provisions of the Terms will remain in full force and effect. No waiver of any provision of these Terms will be deemed a further waiver or continuing waiver or such provision or any other provision, and our failure to assert any right or provision under these Terms will not constitute a waiver of such right or provision. Nothing herein will be deemed to create an agency, partnership, joint venture, or franchisor-franchisee relationship of any kind between us and any user or other person or entity, nor do these terms extend rights to any third party. These Terms are in the English language only, which language will be controlling in all respects, and all versions of these Terms in any other language will be for accommodation only and will not be binding on you or us.
+

flash_head-0.1.0/PKG-INFO

@@ -0,0 +1,7 @@
+Metadata-Version: 2.4
+Name: flash-head
+Version: 0.1.0
+Summary: FlashHead: vLLM Plugin for Fast Language Model Head Inference
+Requires-Python: >=3.10
+License-File: LICENSE
+Dynamic: license-file

flash_head-0.1.0/README.md

@@ -0,0 +1,240 @@
+<div align="center">
+
+# FlashHead 
+
+### vLLM Plugin for Fast Language Model Head Inference
+
+
+The dense classification head accounts for up to 60% of parameters in small LLMs and roughly half of decode-step compute. FlashHead replaces it with a two-stage retrieval pipeline — **up to 2.0x model-level inference speedup** while maintaining accuracy — training-free and hardware-friendly. FlashHead integrates via vLLM's official `vllm.general_plugins` entry point: no source patches, no custom Docker image.
+
+<a href="https://python.org/">
+  <img alt="Python" src="https://img.shields.io/badge/Python-3.10+-blue.svg?logo=python" />
+</a>
+<a href="https://github.com/vllm-project/vllm">
+  <img alt="vLLM" src="https://img.shields.io/badge/vLLM-0.14.0+-green.svg" />
+</a>
+<a href="https://github.com/embedl/flash-head/LICENSE">
+    <img alt="License" src="https://img.shields.io/badge/License-Embedl Models Community-red.svg" />
+</a>
+<br>
+<a href="https://arxiv.org/abs/2603.14591">
+  <img alt="Paper" src="https://img.shields.io/badge/arXiv-2603.14591-b31b1b.svg?logo=arxiv" />
+</a>
+<a href="https://huggingface.co/collections/embedl/flashhead">
+  <img alt="Collection" src="https://img.shields.io/badge/Collection-FlashHead Models-yellow.svg?logo=huggingface" />
+</a>
+<a href="https://huggingface.co/spaces/embedl/Edge-Inference-Benchmarks">
+  <img alt="Benchmarks" src="https://img.shields.io/badge/Spaces-Edge Inference Benchmarks-yellow.svg?logo=huggingface" />
+</a>
+
+
+</div>
+
+
+
+## FlashHead: Efficient Drop-In Replacement for the Classification Head in Language Model Inference
+
+
+The standard LM head computes a dense matrix multiplication $h_t × W_{vocab}$ at every decode step, scoring all Vocabulary tokens regardless of relevance. FlashHead reframes this as a two-stage retrieval problem over clustered token embeddings: first identify which regions of vocabulary space are relevant, then score only those candidates.
+
+<p align="center">
+  <img src="docs/flash_head_flow_diagram.svg"  width="75%" />
+</p>
+
+> **⚡ Key Tradeoff** A dense head scores **128,256 tokens per step** (for a 128K vocabulary). With *c = 8,016* clusters and *p = 256* probes, FlashHead scores only **8,016 + 256 × 16 = 12,112 tokens**, a <span style="color:#22c55e; font-weight:600;">10× reduction</span> in scored tokens, while multi-probe retrieval maintains near-perfect recall of the correct next token.
+
+
+<p align="center" width="100%">
+  <img src="docs/dense_head_scoring.svg" width="30%"/>
+  <img src="docs/arrow.svg" width="4%"/>
+  <img src="docs/flash_head_scoring.svg" width="30%"/>
+</p>
+
+<strong>Note.</strong> The offline clustering step runs once per model and adds zero overhead at inference time.
+Both stages use contiguous memory access patterns for GPU and edge accelerator efficiency.
+
+
+### Four key ideas (see [paper](https://arxiv.org/abs/2603.14591))
+
+- **Equal-sized Clustering** Token embeddings grouped into **balanced clusters** for predictable memory access and stable latency. Unlike hierarchical softmax, cluster sizes stay uniform; critical for GPU and edge accelerators.
+
+- **Multi-Probe Retrieval** Instead of committing to a single cluster, FlashHead probes multiple centroids - beam search over vocabulary space. Near-perfect recall with far fewer evaluations. 
+
+- **Full Decoding Support** Supports both greedy and sampling decoding. For sampling, clusters are selected proportionally to centroid probabilities, **preserving the output distribution**. 
+
+- **Selective Quantization** Stage 1 (coarse centroid scoring) runs in low precision; Stage 2 preserves accuracy. The head's quantization weakness becomes a **structural advantage**.
+
+
+## 📦 Installation
+
+**Prerequisites:** Python 3.10+ and [vLLM](https://github.com/vllm-project/vllm) >= 0.14.0
+
+```bash
+pip install flash-head
+```
+
+That's it. The plugin is discovered automatically by vLLM at startup. ✨
+
+<details>
+<summary>Install from source</summary>
+
+```bash
+git clone https://github.com/embedl/flash-head.git
+cd flash-head
+pip install .
+```
+</details>
+
+
+## 🚀 Usage
+
+### ⌨️ CLI
+
+```bash
+# FlashHead activates automatically for compatible models
+vllm serve embedl/Cosmos-Reason2-2B-W4A16-Edge2-FlashHead \
+    --host 0.0.0.0 --port 8000 \
+    --gpu-memory-utilization 0.75 \
+    --max-model-len 8192
+
+# Disable without uninstalling
+FLASHHEAD_ENABLED=0 vllm serve ...
+```
+
+### 🐍 Python
+
+```python
+from vllm import LLM, SamplingParams
+
+llm = LLM(
+    model="embedl/Cosmos-Reason2-2B-W4A16-Edge2-FlashHead",
+    trust_remote_code=True,
+)
+outputs = llm.generate(
+    ["Explain quantum computing."],
+    SamplingParams(max_tokens=50),
+)
+print(outputs[0].outputs[0].text)
+```
+
+The model's `config.json` contains `"flash_head_cache_dir": "flash_head_assets"` which signals FlashHead to activate. Standard models without this field are completely unaffected.
+
+
+## 🔧 vLLM Plugin Integration
+
+1. **Discovery** vLLM discovers the `flash-head` plugin via the `vllm.general_plugins` entry point at startup
+2. **Patching** `register()` is called in every process, intercepting logits computation, sampling, and speculative decoding
+3. **Inference** The worker lazily constructs the FlashHead module on GPU from the model's clustering cache
+
+
+## 🛡️ Safety
+
+FlashHead models use a custom architecture name (e.g., `FlashHeadQwen3VLForConditionalGeneration`). Without the plugin installed, vLLM does not recognize the architecture and refuses to load the model. Users cannot accidentally run at reduced speed.
+
+| Scenario                                | Behavior                                        |
+|-----------------------------------------|-------------------------------------------------|
+| Plugin not installed                    | ❌ vLLM errors: architecture not supported       |
+| Plugin installed, `FLASHHEAD_ENABLED=0` | ⏸️ Clean disable, model loads without FlashHead |
+| Plugin installed, enabled               | ✅ FlashHead loads on GPU, full speedup          |
+
+### 🏗️ Supported Architectures
+
+See most recent architectures in [_FLASHHEAD_ARCHITECTURES](https://github.com/embedl/flash-head/src/flash_head/__init__.py):
+```python
+_FLASHHEAD_ARCHITECTURES = {
+    "FlashHeadLlamaForCausalLM": "vllm.model_executor.models.llama:LlamaForCausalLM",
+    "FlashHeadQwen3ForCausalLM": "vllm.model_executor.models.qwen3:Qwen3ForCausalLM",
+    "FlashHeadQwen3VLForConditionalGeneration": "vllm.model_executor.models.qwen3_vl:Qwen3VLForConditionalGeneration",
+    "FlashHeadGemma3ForCausalLM": "vllm.model_executor.models.gemma2:Gemma2ForCausalLM",
+}
+```
+
+
+## 📤 Publishing FlashHead Models
+
+For the safety check to work, FlashHead models should use this `config.json` structure:
+
+```json
+{
+  "architectures": ["FlashHeadQwen3VLForConditionalGeneration"],
+  "model_type": "qwen3_vl",
+  "flash_head_cache_dir": "flash_head_assets"
+}
+```
+
+- `architectures` uses the `FlashHead*` prefix so vLLM rejects the model without the plugin
+- `model_type` stays standard so vLLM can resolve the base model class
+- `flash_head_cache_dir` points to the clustering cache directory
+- Do NOT include `auto_map` -- the plugin handles registration
+
+
+## 🗺️ Roadmap
+
+- ✅ Core FlashHead plugin for vLLM (greedy decoding)
+- ✅ Balanced clustering with multiprobe retrieval
+- ✅ Inference-time sampling across full vocabulary
+- ✅ Quantized model support
+- 🔄 **Speculative decoding:** Full FlashHead integration with vLLM's speculative decoding pipeline *(in progress)*
+- 🔄 **EAGLE draft proposals:** FlashHead-accelerated draft generation for EAGLE speculative decoding *(in progress)*
+- ⬜ Additional model architectures
+- ⬜ Benchmarks on additional edge platforms (Qualcomm, AMD, Intel, ...)
+
+💡 Want a feature? [Open an issue](https://github.com/embedl/flash-head/issues/new)!
+
+
+## 🤝 Contributing
+
+We welcome contributions, feedback, and collaboration. Whether you're interested in adding support for new architectures, improving performance, or integrating FlashHead into your own inference stack -- we'd love to hear from you.
+
+- **Report issues** Bug reports and feature requests help us improve. [Open an issue](https://github.com/embedl/flash-head/issues/new).
+- **Submit PRs** Code contributions for new architectures, optimizations, or bug fixes.
+- **Research collaboration** Working on efficient inference, vocabulary approximation, or edge deployment? Reach out.
+- **Model contributions** Publish FlashHead-optimized models to the [HuggingFace collection](https://huggingface.co/collections/embedl/flash-head).
+- **Benchmarks** Run FlashHead on your hardware and submit results to the [Edge Inference Benchmarks](https://huggingface.co/spaces/embedl/Edge-Inference-Benchmarks) space.
+
+
+## 📂 Project Structure
+
+```
+flash-head/
+├── src/flash_head/
+│   ├── __init__.py              # Plugin entry point (register)
+│   ├── flash_head.py            # Core clustering-based head
+│   ├── loading.py               # Model/asset loading from HF Hub
+│   └── patches/                 # vLLM runtime patches
+│
+├── pyproject.toml
+└── LICENSE
+```
+
+
+## 📖 Citation
+
+If you use FlashHead in your research, please cite:
+
+```bibtex
+@article{tranheden2026flashhead,
+  title={FlashHead: Efficient Drop-In Replacement for the Classification Head in Language Model Inference},
+  author={Tranheden, Wilhelm and Ahmed, Shahnawaz and Dubhashi, Devdatt and Matthiesen, Jonna and von Essen, Hannes},
+  journal={arXiv preprint arXiv:2603.14591},
+  year={2026}
+}
+```
+
+## License
+
+Free for non-commercial use within the Embedl Community License (v.1.0).
+
+<div align="center">
+
+### Interested in FlashHead?
+
+Enterprise licensing, custom model optimization, and engineering support available.
+
+[models@embedl.com](mailto:models@embedl.com) &nbsp;&bull;&nbsp; [embedl.com](https://embedl.com)
+
+<br>
+
+<sub>&copy; 2026 Embedl AB. All rights reserved.</sub>
+
+</div>

flash_head-0.1.0/pyproject.toml

@@ -0,0 +1,21 @@
+# Copyright (C) 2026 Embedl AB
+
+[build-system]
+requires = ["setuptools>=64"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "flash-head"
+dynamic = ["version"]
+description = "FlashHead: vLLM Plugin for Fast Language Model Head Inference"
+requires-python = ">=3.10"
+dependencies = []  # torch, safetensors, huggingface_hub provided by vLLM
+
+[project.entry-points."vllm.general_plugins"]
+flash_head = "flash_head:register"
+
+[tool.setuptools.dynamic]
+version = {attr = "flash_head._version.__version__"}
+
+[tool.setuptools.packages.find]
+where = ["src"]

flash_head-0.1.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

flash_head-0.1.0/src/flash_head/__init__.py

@@ -0,0 +1,71 @@
+# Copyright (C) 2026 Embedl AB
+
+"""FlashHead: fast approximate language model head as a vLLM plugin.
+
+FlashHead replaces the standard lm_head (final vocabulary projection layer)
+with a clustering-based approximation that only evaluates logits for tokens
+in the top-k most similar clusters, delivering significant speedups on
+edge devices.
+
+This package integrates FlashHead into vLLM via the official plugin system
+(vllm.general_plugins entry point), eliminating the need for source patches
+or custom Docker images.
+"""
+
+import logging
+import os
+
+from flash_head._version import __version__
+
+logger = logging.getLogger(__name__)
+
+_patches_applied = False
+
+
+def register():
+    """vLLM plugin entry point. Called in every process before model init."""
+    global _patches_applied
+    if _patches_applied:
+        return
+    _patches_applied = True
+
+    if os.environ.get("FLASHHEAD_ENABLED", "1") == "0":
+        logger.info("[FlashHead] Disabled via FLASHHEAD_ENABLED=0")
+        return
+
+    from flash_head.patches import apply_all
+
+    apply_all()
+    _register_architectures()
+    logger.info("[FlashHead] Plugin registered")
+
+
+def _register_architectures():
+    """Register FlashHead model architectures so vLLM recognizes them.
+
+    Models published with architectures like 'FlashHeadQwen3VLForConditionalGeneration'
+    will fail to load without this plugin -- giving a clear error instead of
+    silently falling back to the slow standard lm_head path.
+    """
+    try:
+        from vllm import ModelRegistry
+
+        # Map FlashHead architecture names to their base vLLM model classes.
+        # Uses lazy string imports to avoid premature CUDA initialization.
+        # The FlashHead interception happens via the LogitsProcessor patch,
+        # not via a custom model class, so we just need vLLM to accept the
+        # architecture name and load the base model.
+        _FLASHHEAD_ARCHITECTURES = {
+            "FlashHeadLlamaForCausalLM": "vllm.model_executor.models.llama:LlamaForCausalLM",
+            "FlashHeadQwen3ForCausalLM": "vllm.model_executor.models.qwen3:Qwen3ForCausalLM",
+            "FlashHeadQwen3VLForConditionalGeneration": "vllm.model_executor.models.qwen3_vl:Qwen3VLForConditionalGeneration",
+            "FlashHeadGemma3ForCausalLM": "vllm.model_executor.models.gemma2:Gemma2ForCausalLM",
+        }
+
+        supported = ModelRegistry.get_supported_archs()
+        for fh_arch, model_cls_path in _FLASHHEAD_ARCHITECTURES.items():
+            if fh_arch not in supported:
+                ModelRegistry.register_model(fh_arch, model_cls_path)
+                logger.info("[FlashHead] Registered architecture %s", fh_arch)
+    except Exception as e:
+        logger.debug("[FlashHead] Architecture registration skipped: %s", e)

flash_head-0.1.0/src/flash_head/_version.py

@@ -0,0 +1,5 @@
+# Copyright (C) 2026 Embedl AB
+
+"""FlashHead package version. Bump this to trigger a PyPI release."""
+
+__version__ = "0.1.0"