flashmap 0.1.0__py3-none-any.whl

flashmap/vllm/speculator.py

@@ -0,0 +1,235 @@
+"""FlashMap speculator for vLLM's model runner V2.
+
+FlashMap reuses the DFlash runtime unchanged: model loading
+(``load_dflash_model``, which attaches the target's embedding and LM head), the
+context K/V precompute, the fused input-preparation kernel, the attention
+metadata and the FULL CUDA graph that covers the whole draft step. Its block has
+DSpark's layout (``sample_from_anchor``): ``block_size`` query rows per request,
+row 0 holding the anchor, every row a draft slot. What differs is the draft
+model (model.py) and sampling, which reads logits off the folded head over the
+draft vocabulary.
+
+``num_speculative_tokens`` may be smaller than the block: the draft still
+evaluates its whole block (the DiT's self-attention sees all of it) and proposes
+the first ``num_speculative_tokens`` slots, so verification, which dominates at
+high concurrency, checks fewer tokens.
+"""
+
+import torch
+from torch import nn
+from vllm.config import VllmConfig
+from vllm.logger import init_logger
+from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample
+from vllm.v1.worker.gpu.spec_decode.acceptance_estimator import (
+    OnlineAcceptanceEstimator,
+)
+from vllm.v1.worker.gpu.spec_decode.dflash.speculator import DFlashSpeculator
+from vllm.v1.worker.gpu.spec_decode.dflash.utils import load_dflash_model
+
+from flashmap.vllm.fold import split_argmax
+from flashmap.vllm.kernels import scatter_rows
+
+logger = init_logger(__name__)
+
+__all__ = ["FlashMapSpeculator"]
+
+
+class FlashMapSpeculator(DFlashSpeculator):
+    _speculator_name = "FlashMap"
+
+    def __init__(self, vllm_config: VllmConfig, device: torch.device):
+        super().__init__(vllm_config, device)
+        block_size = int(self.draft_model_config.hf_config.block_size)
+        # Drafts proposed for verification, a prefix of the block.
+        self.num_verify = self.num_speculative_steps
+        if self.num_verify > block_size:
+            raise ValueError(
+                f"num_speculative_tokens={self.num_verify} exceeds FlashMap's "
+                f"block_size={block_size}"
+            )
+        # Every one of the block's rows is a draft slot; row 0 is the anchor.
+        self.sample_from_anchor = True
+        self.num_query_per_req = block_size
+        if self.num_verify < block_size:
+            # Per-slot buffers for the whole block (the DFlash ones are sized
+            # for the verified drafts).
+            self.num_speculative_steps = block_size
+            num_slots = self.max_num_reqs * block_size
+            self.draft_tokens = torch.zeros(
+                self.max_num_reqs, block_size, dtype=torch.int64, device=device
+            )
+            self.sample_indices = torch.zeros(
+                num_slots, dtype=torch.int64, device=device
+            )
+            self.sample_pos = torch.zeros(num_slots, dtype=torch.int64, device=device)
+            self.sample_idx_mapping = torch.full(
+                (num_slots,), -1, dtype=torch.int32, device=device
+            )
+            self.sample_col = torch.arange(
+                block_size, dtype=torch.int32, device=device
+            ).repeat(self.max_num_reqs)
+
+    def draft_logits_spec(self, vllm_config: VllmConfig) -> tuple[torch.dtype, float]:
+        # A restricted draft vocabulary writes only its own columns of the
+        # cached proposal distribution; every other token has probability 0.
+        dtype, fill = super().draft_logits_spec(vllm_config)
+        hf_config = self.draft_model_config.hf_config
+        draft_vocab_size = getattr(hf_config, "draft_vocab_size", None)
+        if draft_vocab_size and draft_vocab_size < hf_config.vocab_size:
+            fill = float("-inf")
+        return dtype, fill
+
+    def load_draft_model(
+        self,
+        target_model: nn.Module,
+        target_attn_layer_names: set[str],  # noqa: ARG002
+    ) -> nn.Module:
+        model = load_dflash_model(target_model, self.vllm_config)
+        model.prepare()
+        # Key each query row's noise by its request's seed and its position.
+        model.set_noise_keys(self.seeds, self.sample_idx_mapping)
+        if self.draft_watermarker is not None and model.hot_ids is not None:
+            raise ValueError(
+                "Watermarked drafting needs FlashMap's full draft vocabulary."
+            )
+        return model
+
+    def load_model(self, target_model: nn.Module) -> None:
+        super().load_model(target_model)
+        if self.acceptance_estimator is not None:
+            # Acceptance is observed for the verified drafts only.
+            self.acceptance_estimator = OnlineAcceptanceEstimator(
+                self.max_num_reqs, self.num_verify, self.device
+            )
+
+    def set_attn(self, *args, **kwargs) -> None:
+        super().set_attn(*args, **kwargs)
+        target_attn_groups = kwargs.get("target_attn_groups", args[-1] if args else [])
+        # The cross-attention schedules itself (model.py); vLLM's ahead-of-time
+        # FlashAttention schedule for the draft group would be unused work, and
+        # it mistakes an fp8 draft cache under a bf16 target for an unsupported
+        # dtype.
+        for groups in self.attn_groups:
+            for group in groups:
+                builder = group.get_metadata_builder()
+                if getattr(builder, "aot_schedule", False):
+                    builder.aot_schedule = False  # type: ignore[attr-defined]
+        # The target runner builds metadata for the draft group too (cudagraph
+        # capture), through its own builders.
+        for gid in self.draft_kv_cache_group_ids:
+            for group in target_attn_groups[gid]:
+                builder = group.get_metadata_builder()
+                if getattr(builder, "aot_schedule", False) and getattr(
+                    builder, "kv_cache_dtype", None
+                ) not in (torch.bfloat16, torch.float16):
+                    builder.aot_schedule = False  # type: ignore[attr-defined]
+
+    def _precompute_context_kv(
+        self, start: int, end: int, dummy_run: bool = False
+    ) -> None:
+        # The context rows a draft step stores right before its forward are
+        # projected by the forward's GEMM, together with the anchors.
+        if start == 0 and not dummy_run and self._layer_group_idx is None:
+            self.model.defer_context_kv(
+                self.hidden_states[start:end],
+                self.context_positions[start:end],
+                self._context_slot_mappings[0][start:end],
+            )
+            return
+        super()._precompute_context_kv(start, end, dummy_run)
+
+    def capture(self) -> None:
+        # Record the KV caches' layout (a host-to-device copy) before capturing.
+        self.model.cache_layout()
+        super().capture()
+
+    def propose(self, *args, **kwargs) -> torch.Tensor:
+        return super().propose(*args, **kwargs)[:, : self.num_verify]
+
+    def sample_draft(
+        self,
+        hidden_states: torch.Tensor,
+        sample_src_positions: torch.Tensor,
+        idx_mapping: torch.Tensor,
+        temperature: torch.Tensor,
+        seeds: torch.Tensor,
+        draft_step: torch.Tensor,
+        draft_logits: torch.Tensor | None,
+    ) -> torch.Tensor:
+        block, verify = self.num_speculative_steps, self.num_verify
+        if verify == block:
+            return self._sample(
+                hidden_states,
+                sample_src_positions,
+                idx_mapping,
+                temperature,
+                seeds,
+                draft_step,
+                draft_logits,
+            )
+        # Sample the verified prefix of each block only; the rest stay zero.
+        num_blocks = hidden_states.shape[0] // block
+
+        def prefix(t: torch.Tensor) -> torch.Tensor:
+            return t.view(num_blocks, block, *t.shape[1:])[:, :verify].flatten(0, 1)
+
+        sampled = self._sample(
+            prefix(hidden_states),
+            prefix(sample_src_positions),
+            prefix(idx_mapping),
+            temperature,
+            seeds,
+            prefix(draft_step),
+            draft_logits,
+        )
+        tokens = sampled.new_zeros(num_blocks, block)
+        tokens[:, :verify] = sampled.view(num_blocks, verify)
+        return tokens.view(-1)
+
+    def _sample(
+        self,
+        hidden_states: torch.Tensor,
+        sample_src_positions: torch.Tensor,
+        idx_mapping: torch.Tensor,
+        temperature: torch.Tensor,
+        seeds: torch.Tensor,
+        draft_step: torch.Tensor,
+        draft_logits: torch.Tensor | None,
+    ) -> torch.Tensor:
+        logits = self.model.compute_draft_logits(hidden_states)
+        hot_ids = self.model.hot_ids
+        if draft_logits is None:
+            sampled = split_argmax(logits)
+        elif hot_ids is None:
+            sampled = gumbel_sample(
+                logits,
+                idx_mapping,
+                temperature,
+                seeds,
+                sample_src_positions,
+                apply_temperature=True,
+                is_drafting=True,
+                logits_cache=draft_logits,
+                logits_cache_col=draft_step,
+                use_fp64=self.use_fp64_gumbel,
+            )
+            if self.draft_watermarker is not None:
+                sampled = self.draft_watermarker.sample(
+                    logits, sampled, idx_mapping, temperature
+                )
+        else:
+            # Sample over the draft vocabulary and cache its logits for
+            # probabilistic verification (the other columns stay -inf).
+            sampled = gumbel_sample(
+                logits,
+                idx_mapping,
+                temperature,
+                seeds,
+                sample_src_positions,
+                apply_temperature=True,
+                is_drafting=True,
+                use_fp64=self.use_fp64_gumbel,
+            )
+            scatter_rows(logits, idx_mapping, draft_step, hot_ids, draft_logits)
+        self._maybe_predict_acceptance(logits, idx_mapping, draft_step)
+        return sampled if hot_ids is None else hot_ids[sampled]

flashmap-0.1.0.dist-info/METADATA

@@ -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.dist-info/RECORD

@@ -0,0 +1,22 @@
+flashmap/__init__.py,sha256=iLu7hWQJMNR304XnHiwdi_GWhfYljkuezvPBHdY4q2k,745
+flashmap/configuration_flashmap.py,sha256=BZl55qpiNyTcQ5thljUkcjAh2EhDjSoK7OKWMxcCOyk,4561
+flashmap/modeling_flashmap.py,sha256=KegDQeqNRc0tWtAaPTZOID3olMGceBgNYV_bROAsiqU,29630
+flashmap/sglang/__init__.py,sha256=b7817GTu7_tZP2U3ZvyqBkVa8UtBBjcKF5we9gbWZJY,3227
+flashmap/sglang/jump_bank.py,sha256=a5GFyXr-BpeJ1oqCoPEO0ZZQ5jCYUrsDjYjJ9mgSfvY,9007
+flashmap/sglang/kernels.py,sha256=k6lugxnjRba6rXW6VFQJNoIFzlpQjReuW0-UL3hjO7s,17409
+flashmap/sglang/model.py,sha256=WYxBbEix5NYFY54TQr4zyYWsVGoj1HJhIfdG-6l4Dvo,48072
+flashmap/sglang/sampler.py,sha256=RqDdDMe3oPpf3TXqS1KhAPqgEk9eLvWEt8K5TzcSZoc,8169
+flashmap/sglang/tree.py,sha256=QAfPpzTcEtU0f8Dt4wu_ldptwEozrwS-XN9A-CH_sIc,3589
+flashmap/sglang/tree_worker.py,sha256=5s2QHocg77dxFfsjMigfrX47ATaQyh5aISN78ykSREY,16276
+flashmap/sglang/verify_attention.py,sha256=uWT4sJ0QePUcP-hhXbraiDJLwxfqJYuH_D_SO_MywyE,15867
+flashmap/vllm/__init__.py,sha256=KXUpQz1PO_ivsVsLz67FfIcsTQ_08cAdiC3gO_hoI_U,10571
+flashmap/vllm/fold.py,sha256=NpYDC3zaMb6dps5OjUO-GHK4wfoOnDtx0WdaJaqQccE,5043
+flashmap/vllm/kernels.py,sha256=wYc_0CqH8_m5W989k6uCEDS0bNy39lo1MsMfUPvK2Ek,15622
+flashmap/vllm/model.py,sha256=LhdGSe5JEVeOn3XKKQNv3pu2MqDKnf4xtPE0AfRbSu4,37405
+flashmap/vllm/speculator.py,sha256=8Ms67rtgscFesi_4TTSdxG2744FizLU2FpenbHmRQG0,9937
+flashmap-0.1.0.dist-info/licenses/LICENSE,sha256=fCprpx2_9vwmNfySvgmLmFvc2h-Qholc22PsEnKX_OU,1066
+flashmap-0.1.0.dist-info/METADATA,sha256=Ws3Y4gnghJcm7c9SMXSO_bSpiOgnzE4ysSKsD66CBe0,7129
+flashmap-0.1.0.dist-info/WHEEL,sha256=YVMoNqKzERt-wjUZwJ33xBGAwnFl-4cqbYkTtWa4itE,91
+flashmap-0.1.0.dist-info/entry_points.txt,sha256=tAz7OjTvfRKJHnn6EVnM5r8uwiqD7EDmpwXlCAhIxCU,115
+flashmap-0.1.0.dist-info/top_level.txt,sha256=L1exL6jCAPP_DKhiqwOmI-PSMo16XxAQlT8K37wK6ok,9
+flashmap-0.1.0.dist-info/RECORD,,

flashmap-0.1.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (84.0.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

flashmap-0.1.0.dist-info/entry_points.txt

@@ -0,0 +1,5 @@
+[sglang.srt.plugins]
+flashmap = flashmap.sglang:register
+
+[vllm.general_plugins]
+flashmap = flashmap.vllm:register

flashmap-0.1.0.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+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
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+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.dist-info/top_level.txt

@@ -0,0 +1 @@
+flashmap