FlashMap-Qwen3-4B

Project page · Code · Paper · Models

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), 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 is the FlashMap drafter for Qwen/Qwen3-4B.

Install

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).

Quickstart

import torch
from transformers import AutoModel, AutoModelForCausalLM, AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-4B")
target = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3-4B", dtype=torch.bfloat16).cuda().eval()
draft = AutoModel.from_pretrained("j-shen/FlashMap-Qwen3-4B", 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 the flashmap package.

Serve with SGLang

The SGLang integration is a plugin in the flashmap package (the sglang extra), which SGLang loads automatically.

uv venv --python 3.12 .venv-sglang
source .venv-sglang/bin/activate
uv pip install "sglang==0.5.20" "flashmap[sglang]"

sglang serve --model-path Qwen/Qwen3-4B \
  --speculative-algorithm DSPARK --speculative-draft-model-path j-shen/FlashMap-Qwen3-4B

Serve with vLLM

The vLLM integration is a plugin in the flashmap package (the vllm extra), which vLLM loads automatically.

uv venv --python 3.12 .venv-vllm
source .venv-vllm/bin/activate
uv pip install "vllm==0.31.*" "flashmap[vllm]"

vllm serve Qwen/Qwen3-4B --speculative-config '{"model": "j-shen/FlashMap-Qwen3-4B", "num_speculative_tokens": 16}'

For sampling at temperature > 0, add "draft_sample_method": "probabilistic" to the speculative config.

Sending requests

Both servers expose an OpenAI-compatible API, SGLang on port 30000 and vLLM on port 8000.

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)

License

The model weights are released under Apache-2.0, following Qwen3. The model code (*.py, from the GitHub repository) is released under the MIT license.

Citation

@misc{flashmap2026,
  title         = {{FlashMap}: Flow Map Speculative Decoding},
  author        = {Shen, Jeff and Zhou, Zhejian and Lee, Ryan and May, Jonathan and Ho, Shirley},
  year          = {2026},
  eprint        = {2610.24365},
  archivePrefix = {arXiv},
  primaryClass  = {cs.CL},
  url           = {https://arxiv.org/abs/2610.24365}
}

j-shen/FlashMap-Qwen3-4B