flashmap 0.0.1__tar.gz → 0.1.0__tar.gz

{flashmap-0.0.1 → flashmap-0.1.0}/LICENSE

@@ -1,6 +1,6 @@
 MIT License
 
-Copyright (c) 2025, pecwie
+Copyright (c) 2026 Jeff Shen
 
 Permission is hereby granted, free of charge, to any person obtaining a copy
 of this software and associated documentation files (the "Software"), to deal
@@ -19,4 +19,3 @@ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 SOFTWARE.
-

flashmap-0.1.0/PKG-INFO

@@ -0,0 +1,152 @@
+Metadata-Version: 2.4
+Name: flashmap
+Version: 0.1.0
+Summary: FlashMap: flow map speculative decoding for Hugging Face transformers
+Author: Jeff Shen, Zhejian Zhou, Ryan Lee, Jonathan May, Shirley Ho
+License-Expression: MIT
+Project-URL: Homepage, https://jshen.net/flashmap
+Project-URL: Repository, https://github.com/al-jshen/flashmap
+Project-URL: Documentation, https://github.com/al-jshen/flashmap#readme
+Project-URL: Model (Qwen3-8B), https://huggingface.co/j-shen/FlashMap-Qwen3-8B
+Project-URL: Model (Qwen3-4B), https://huggingface.co/j-shen/FlashMap-Qwen3-4B
+Keywords: speculative decoding,flow map,draft model,llm inference,transformers
+Classifier: Programming Language :: Python :: 3
+Classifier: Operating System :: OS Independent
+Classifier: Intended Audience :: Science/Research
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch>=2.1
+Requires-Dist: transformers>=4.56
+Provides-Extra: benchmark
+Requires-Dist: datasets; extra == "benchmark"
+Provides-Extra: test
+Requires-Dist: pytest; extra == "test"
+Provides-Extra: sglang
+Requires-Dist: sglang==0.5.20; extra == "sglang"
+Provides-Extra: vllm
+Requires-Dist: vllm<0.32,>=0.31; extra == "vllm"
+Dynamic: license-file
+
+# FlashMap: Flow Map Speculative Decoding
+
+**[Project page](https://jshen.net/flashmap)** · **Paper** (coming soon) · **[Models](https://huggingface.co/collections/j-shen/flashmap-6ad0feb955bdc30db59769fa)**
+
+FlashMap is a drafter for lossless speculative decoding. It proposes a whole block of K tokens in a single pass. Unlike discrete diffusion drafters, which predict each token of the block independently and ignore inter-token dependencies, FlashMap uses a flow map ([Flow Map Language Models](https://arxiv.org/abs/2602.16813)), which jumps from Gaussian noise to the whole block in a single step and can capture dependencies between its tokens. On Qwen3-8B, FlashMap decodes up to 6x faster than the target alone.
+
+This repository is a small, standalone implementation for Hugging Face `transformers`. For fast serving, use
+the [SGLang](#serve-with-sglang) or [vLLM](#serve-with-vllm) integration.
+
+## Models
+
+| Drafter | Target | Drafter parameters |
+|---|---|---|
+| [j-shen/FlashMap-Qwen3-8B](https://huggingface.co/j-shen/FlashMap-Qwen3-8B) | [Qwen/Qwen3-8B](https://huggingface.co/Qwen/Qwen3-8B) | 500M |
+| [j-shen/FlashMap-Qwen3-4B](https://huggingface.co/j-shen/FlashMap-Qwen3-4B) | [Qwen/Qwen3-4B](https://huggingface.co/Qwen/Qwen3-4B) | 475M |
+
+## Install
+
+```bash
+uv venv --python 3.12
+source .venv/bin/activate
+uv pip install flashmap
+```
+
+The package needs only `torch` and `transformers>=4.56`; serving takes the `sglang` or `vllm` extra (below). It
+is tested with torch 2.9 and 2.13 and with transformers 4.57, 5.10, 5.16 and 5.20.
+
+## Quickstart
+
+```python
+import torch
+from transformers import AutoModel, AutoModelForCausalLM, AutoTokenizer
+
+tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-8B")
+target = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3-8B", dtype=torch.bfloat16).cuda().eval()
+draft = AutoModel.from_pretrained("j-shen/FlashMap-Qwen3-8B", trust_remote_code=True, dtype=torch.bfloat16).cuda().eval()
+
+messages = [{"role": "user", "content": "How many positive whole-number divisors does 196 have?"}]
+input_ids = tokenizer.apply_chat_template(
+    messages, add_generation_prompt=True, enable_thinking=False, return_tensors="pt", return_dict=True
+)["input_ids"].cuda()
+
+output = draft.spec_generate(target, input_ids, max_new_tokens=512, temperature=0.0, num_steps=1)
+print(tokenizer.decode(output.sequences[0, input_ids.shape[1]:], skip_special_tokens=True))
+print(f"acceptance length: {output.acceptance_length:.2f}")
+```
+
+`num_steps=2` drafts with two jumps of the flow map (trading speed for accuracy). 
+The model code also ships with the weights, so `AutoModel.from_pretrained(..., trust_remote_code=True)` works without this
+package. 
+
+## Serve with SGLang
+
+The SGLang integration is a plugin in this package (the `sglang` extra), which SGLang loads automatically.
+
+```bash
+uv venv --python 3.12 .venv-sglang
+source .venv-sglang/bin/activate
+uv pip install "sglang==0.5.20" "flashmap[sglang]"
+
+python -m sglang.launch_server --model-path Qwen/Qwen3-8B \
+  --speculative-algorithm DSPARK --speculative-draft-model-path j-shen/FlashMap-Qwen3-8B
+```
+
+[Our SpecForge fork](https://github.com/al-jshen/SpecForge/tree/flashmap) trains FlashMap with SGLang's
+training library, and its `specforge export --to hf` writes a trained map as a release folder that
+`flashmap[sglang]` serves as it is.
+
+## Serve with vLLM
+
+The vLLM integration is a plugin in this package (the `vllm` extra), which vLLM loads automatically.
+
+```bash
+uv venv --python 3.12 .venv-vllm
+source .venv-vllm/bin/activate
+uv pip install "vllm==0.31.*" "flashmap[vllm]"
+
+vllm serve Qwen/Qwen3-8B --speculative-config '{"model": "j-shen/FlashMap-Qwen3-8B", "num_speculative_tokens": 16}'
+```
+
+For sampling at `temperature > 0`, add `"draft_sample_method": "probabilistic"` to the speculative config.
+
+Drafters you train yourself ([Training](#training)) are served the same way once converted to the release
+format.
+
+## Sending requests
+
+Both servers expose an OpenAI-compatible API, SGLang on port 30000 and vLLM on port 8000. 
+
+```python
+from openai import OpenAI
+
+client = OpenAI(base_url="http://localhost:30000/v1", api_key="none")  # vLLM: port 8000
+response = client.chat.completions.create(
+    model=client.models.list().data[0].id,
+    messages=[{"role": "user", "content": "How many positive whole-number divisors does 196 have?"}],
+    temperature=0,
+    extra_body={"chat_template_kwargs": {"enable_thinking": False}},
+)
+print(response.choices[0].message.content)
+```
+
+## Training
+
+FlashMap is trained in two stages: a flow over one-hot token blocks, then a flow map distilled from it. The
+recipes are in the speculators fork
+([Training FlashMap](https://github.com/al-jshen/speculators/blob/flashmap/docs/user_guide/algorithms/flashmap_training.md))
+and the SpecForge fork ([FlashMap](https://github.com/al-jshen/SpecForge/blob/flashmap/docs/sections/concepts/FlashMap.md)).
+Both write a trained map as a release folder for this package (`speculators convert --algorithm flashmap`,
+`specforge export --to hf`), and [`scripts/convert_checkpoint.py`](https://github.com/al-jshen/flashmap/blob/master/scripts/convert_checkpoint.py) converts
+speculators and SpecForge checkpoints too (`python scripts/convert_checkpoint.py --help`).
+
+## License
+
+The code is released under the [MIT license](https://github.com/al-jshen/flashmap/blob/master/LICENSE). The model weights are released under Apache-2.0,
+following Qwen3.
+
+FlashMap builds on [DFlash](https://github.com/z-lab/dflash),
+[Flow Map Language Models](https://arxiv.org/abs/2602.16813), [DeepSpec](https://github.com/deepseek-ai/DeepSpec),
+[SGLang](https://github.com/sgl-project/sglang), [vLLM](https://github.com/vllm-project/vllm),
+[speculators](https://github.com/vllm-project/speculators) and [SpecForge](https://github.com/sgl-project/SpecForge).

flashmap-0.1.0/README.md

@@ -0,0 +1,121 @@
+# FlashMap: Flow Map Speculative Decoding
+
+**[Project page](https://jshen.net/flashmap)** · **Paper** (coming soon) · **[Models](https://huggingface.co/collections/j-shen/flashmap-6ad0feb955bdc30db59769fa)**
+
+FlashMap is a drafter for lossless speculative decoding. It proposes a whole block of K tokens in a single pass. Unlike discrete diffusion drafters, which predict each token of the block independently and ignore inter-token dependencies, FlashMap uses a flow map ([Flow Map Language Models](https://arxiv.org/abs/2602.16813)), which jumps from Gaussian noise to the whole block in a single step and can capture dependencies between its tokens. On Qwen3-8B, FlashMap decodes up to 6x faster than the target alone.
+
+This repository is a small, standalone implementation for Hugging Face `transformers`. For fast serving, use
+the [SGLang](#serve-with-sglang) or [vLLM](#serve-with-vllm) integration.
+
+## Models
+
+| Drafter | Target | Drafter parameters |
+|---|---|---|
+| [j-shen/FlashMap-Qwen3-8B](https://huggingface.co/j-shen/FlashMap-Qwen3-8B) | [Qwen/Qwen3-8B](https://huggingface.co/Qwen/Qwen3-8B) | 500M |
+| [j-shen/FlashMap-Qwen3-4B](https://huggingface.co/j-shen/FlashMap-Qwen3-4B) | [Qwen/Qwen3-4B](https://huggingface.co/Qwen/Qwen3-4B) | 475M |
+
+## Install
+
+```bash
+uv venv --python 3.12
+source .venv/bin/activate
+uv pip install flashmap
+```
+
+The package needs only `torch` and `transformers>=4.56`; serving takes the `sglang` or `vllm` extra (below). It
+is tested with torch 2.9 and 2.13 and with transformers 4.57, 5.10, 5.16 and 5.20.
+
+## Quickstart
+
+```python
+import torch
+from transformers import AutoModel, AutoModelForCausalLM, AutoTokenizer
+
+tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-8B")
+target = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3-8B", dtype=torch.bfloat16).cuda().eval()
+draft = AutoModel.from_pretrained("j-shen/FlashMap-Qwen3-8B", trust_remote_code=True, dtype=torch.bfloat16).cuda().eval()
+
+messages = [{"role": "user", "content": "How many positive whole-number divisors does 196 have?"}]
+input_ids = tokenizer.apply_chat_template(
+    messages, add_generation_prompt=True, enable_thinking=False, return_tensors="pt", return_dict=True
+)["input_ids"].cuda()
+
+output = draft.spec_generate(target, input_ids, max_new_tokens=512, temperature=0.0, num_steps=1)
+print(tokenizer.decode(output.sequences[0, input_ids.shape[1]:], skip_special_tokens=True))
+print(f"acceptance length: {output.acceptance_length:.2f}")
+```
+
+`num_steps=2` drafts with two jumps of the flow map (trading speed for accuracy). 
+The model code also ships with the weights, so `AutoModel.from_pretrained(..., trust_remote_code=True)` works without this
+package. 
+
+## Serve with SGLang
+
+The SGLang integration is a plugin in this package (the `sglang` extra), which SGLang loads automatically.
+
+```bash
+uv venv --python 3.12 .venv-sglang
+source .venv-sglang/bin/activate
+uv pip install "sglang==0.5.20" "flashmap[sglang]"
+
+python -m sglang.launch_server --model-path Qwen/Qwen3-8B \
+  --speculative-algorithm DSPARK --speculative-draft-model-path j-shen/FlashMap-Qwen3-8B
+```
+
+[Our SpecForge fork](https://github.com/al-jshen/SpecForge/tree/flashmap) trains FlashMap with SGLang's
+training library, and its `specforge export --to hf` writes a trained map as a release folder that
+`flashmap[sglang]` serves as it is.
+
+## Serve with vLLM
+
+The vLLM integration is a plugin in this package (the `vllm` extra), which vLLM loads automatically.
+
+```bash
+uv venv --python 3.12 .venv-vllm
+source .venv-vllm/bin/activate
+uv pip install "vllm==0.31.*" "flashmap[vllm]"
+
+vllm serve Qwen/Qwen3-8B --speculative-config '{"model": "j-shen/FlashMap-Qwen3-8B", "num_speculative_tokens": 16}'
+```
+
+For sampling at `temperature > 0`, add `"draft_sample_method": "probabilistic"` to the speculative config.
+
+Drafters you train yourself ([Training](#training)) are served the same way once converted to the release
+format.
+
+## Sending requests
+
+Both servers expose an OpenAI-compatible API, SGLang on port 30000 and vLLM on port 8000. 
+
+```python
+from openai import OpenAI
+
+client = OpenAI(base_url="http://localhost:30000/v1", api_key="none")  # vLLM: port 8000
+response = client.chat.completions.create(
+    model=client.models.list().data[0].id,
+    messages=[{"role": "user", "content": "How many positive whole-number divisors does 196 have?"}],
+    temperature=0,
+    extra_body={"chat_template_kwargs": {"enable_thinking": False}},
+)
+print(response.choices[0].message.content)
+```
+
+## Training
+
+FlashMap is trained in two stages: a flow over one-hot token blocks, then a flow map distilled from it. The
+recipes are in the speculators fork
+([Training FlashMap](https://github.com/al-jshen/speculators/blob/flashmap/docs/user_guide/algorithms/flashmap_training.md))
+and the SpecForge fork ([FlashMap](https://github.com/al-jshen/SpecForge/blob/flashmap/docs/sections/concepts/FlashMap.md)).
+Both write a trained map as a release folder for this package (`speculators convert --algorithm flashmap`,
+`specforge export --to hf`), and [`scripts/convert_checkpoint.py`](https://github.com/al-jshen/flashmap/blob/master/scripts/convert_checkpoint.py) converts
+speculators and SpecForge checkpoints too (`python scripts/convert_checkpoint.py --help`).
+
+## License
+
+The code is released under the [MIT license](https://github.com/al-jshen/flashmap/blob/master/LICENSE). The model weights are released under Apache-2.0,
+following Qwen3.
+
+FlashMap builds on [DFlash](https://github.com/z-lab/dflash),
+[Flow Map Language Models](https://arxiv.org/abs/2602.16813), [DeepSpec](https://github.com/deepseek-ai/DeepSpec),
+[SGLang](https://github.com/sgl-project/sglang), [vLLM](https://github.com/vllm-project/vllm),
+[speculators](https://github.com/vllm-project/speculators) and [SpecForge](https://github.com/sgl-project/SpecForge).

flashmap-0.1.0/flashmap/__init__.py

@@ -0,0 +1,14 @@
+from transformers import AutoConfig, AutoModel
+
+from .configuration_flashmap import FlashMapConfig
+from .modeling_flashmap import FlashMapCache, FlashMapDraftModel, FlashMapGenerateOutput
+
+# With the package imported, AutoModel loads FlashMap folders without trust_remote_code,
+# and save_pretrained writes the model code next to the weights, so that the folder
+# loads with trust_remote_code=True alone.
+AutoConfig.register(FlashMapConfig.model_type, FlashMapConfig, exist_ok=True)
+AutoModel.register(FlashMapConfig, FlashMapDraftModel, exist_ok=True)
+FlashMapConfig.register_for_auto_class()
+FlashMapDraftModel.register_for_auto_class("AutoModel")
+
+__all__ = ["FlashMapCache", "FlashMapConfig", "FlashMapDraftModel", "FlashMapGenerateOutput"]

flashmap-0.1.0/flashmap/configuration_flashmap.py

@@ -0,0 +1,105 @@
+"""Configuration of a FlashMap draft model."""
+
+from transformers import PretrainedConfig
+
+
+# Jump grids of the "decoding_error" time schedule (Flow Map Language Models,
+# eq. 25) at the Qwen3 vocabulary size, with the endpoints set to exactly 0 and 1.
+# A draft with ``n`` steps jumps from grid[i] to grid[i + 1] for i < n.
+DEFAULT_JUMP_GRIDS = {
+    "1": [0.0, 1.0],
+    "2": [0.0, 0.8172569274902344, 1.0],
+    "3": [0.0, 0.8010796308517456, 0.8310279846191406, 1.0],
+    "4": [0.0, 0.7906050682067871, 0.8172569274902344, 0.8379518389701843, 1.0],
+}
+
+
+class FlashMapConfig(PretrainedConfig):
+    r"""A FlashMap drafter for speculative decoding.
+
+    FlashMap drafts a block of ``block_size`` tokens after the last verified one
+    (the anchor) with a flow map: a DiT that jumps from Gaussian noise over the
+    target's vocabulary to the block's per-slot token distributions, conditioned on
+    the target model's hidden states through cross-attention.
+
+    The drafter borrows the target's token embedding and lm_head, which are not part
+    of its weights.
+
+    Args:
+        target_model (`str`): Hub id of the target model the drafter was trained for.
+        target_hidden_size (`int`): Hidden size of the target model.
+        vocab_size (`int`): Vocabulary size of the target model.
+        target_layer_ids (`list[int]`): Target decoder layers whose outputs (that is,
+            `hidden_states[i + 1]` of a transformers forward) are the drafter's context.
+        block_size (`int`): Number of draft tokens per block.
+        hidden_size (`int`): Width of the DiT.
+        num_hidden_layers (`int`): Number of DiT blocks.
+        num_attention_heads (`int`): Attention heads of the DiT.
+        intermediate_size (`int`): MLP width of the DiT.
+        cond_dim (`int`): Width of the time conditioning.
+        time_embed_dim (`int`): Size of the sinusoidal time features.
+        time_scale (`float`): Times in [0, 1] are multiplied by this before their
+            sinusoidal embedding.
+        rope_theta (`float`): RoPE base. RoPE rotates interleaved pairs of channels.
+        rms_norm_eps (`float`): Epsilon of the context, anchor and state RMSNorms.
+        qk_norm_eps (`float`): Epsilon of the query/key RMSNorms.
+        layer_norm_eps (`float`): Epsilon of the adaptive LayerNorms.
+        noise_gate (`str`): How the noise term of the state embedding is weighted:
+            "jump" by `noise_gate * (t - s)` for a jump from `s` to `t`, "constant" by
+            `noise_gate` alone.
+        time_schedule (`str`): The schedule the jump grids come from.
+        num_steps (`int`): Default number of jumps per draft.
+        jump_grids (`dict[str, list[float]]`): Jump times for each number of steps.
+    """
+
+    model_type = "flashmap"
+
+    def __init__(
+        self,
+        target_model: str = "Qwen/Qwen3-8B",
+        target_hidden_size: int = 4096,
+        vocab_size: int = 151936,
+        target_layer_ids=(1, 9, 17, 25, 33),
+        block_size: int = 16,
+        hidden_size: int = 2048,
+        num_hidden_layers: int = 6,
+        num_attention_heads: int = 16,
+        intermediate_size: int = 8192,
+        cond_dim: int = 256,
+        time_embed_dim: int = 256,
+        time_scale: float = 1000.0,
+        rope_theta: float = 10000.0,
+        rms_norm_eps: float = 1e-6,
+        qk_norm_eps: float = 1e-5,
+        layer_norm_eps: float = 1e-6,
+        noise_gate: str = "jump",
+        time_schedule: str = "decoding_error",
+        num_steps: int = 1,
+        jump_grids=None,
+        **kwargs,
+    ):
+        self.target_model = target_model
+        self.target_hidden_size = target_hidden_size
+        self.vocab_size = vocab_size
+        self.target_layer_ids = list(target_layer_ids)
+        self.block_size = block_size
+        self.hidden_size = hidden_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.intermediate_size = intermediate_size
+        self.cond_dim = cond_dim
+        self.time_embed_dim = time_embed_dim
+        self.time_scale = time_scale
+        self.rope_theta = rope_theta
+        self.rms_norm_eps = rms_norm_eps
+        self.qk_norm_eps = qk_norm_eps
+        self.layer_norm_eps = layer_norm_eps
+        self.noise_gate = noise_gate
+        self.time_schedule = time_schedule
+        self.num_steps = num_steps
+        self.jump_grids = dict(DEFAULT_JUMP_GRIDS if jump_grids is None else jump_grids)
+        kwargs.setdefault("tie_word_embeddings", False)
+        super().__init__(**kwargs)
+
+
+__all__ = ["FlashMapConfig"]