FlashMap: Flow Map Speculative Decoding

TL;DR FlashMap is a speculative-decoding drafter that samples a whole block of tokens jointly, in one forward pass of a flow map, so the target model decodes much faster with exactly the same outputs.

How it works

In each round, FlashMap reads hidden states from a few layers of the target model and turns Gaussian noise into the next K tokens with a single evaluation of a flow map. The target then checks all K in one forward pass and keeps the longest prefix it agrees with, so the output is exactly what the target would have generated.

Diagram of one speculative decoding round: the target LLM passes hidden states to the FlashMap drafter, a DiT with self-attention, cross-attention and MLP blocks modulated by the jump times, which maps noise to logits for all K slots in one pass; the target verifies the drafted tokens in one pass. — **(a)** One round of speculative decoding with FlashMap. The drafter is a DiT that cross-attends to the target's hidden states and drafts all K slots in one pass, using the target's frozen embedding and LM head.

Autoregressive drafters are accurate and their predicted tokens are self-consistent, but they require many function evaluations and are slow. On the other hand, parallel drafters based on discrete diffusion like DFlash predict a whole block in a single pass, but each position of the block is predicted independently, and inter-token dependencies are ignored. FlashMap is both parallel and self-consistent: we draw the whole block in a single pass while modelling joint inter-token dependencies.

Two-slot example after 'we landed in': DFlash samples each slot independently, so half its drafts mix 'new diego' or 'san york' and are rejected; FlashMap maps each noise sample to one mode, so every draft is 'new york' or 'san diego'. — **(b)** Drafting two correlated slots. DFlash samples each slot independently and half of its drafts mix two modes; the flow map can represent the joint distribution over the slots.

Quickstart

Drafters for Qwen3-8B and Qwen3-4B are on Hugging Face. Use Transformers for a small, readable reference, or vLLM or SGLang for fast serving.

Qwen3-8B Qwen3-4B

Transformers

Install

uv venv --python 3.12 && source .venv/bin/activate
uv pip install flashmap

Generate

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

prompt = "How many positive whole-number divisors does 196 have?"
input_ids = tokenizer.apply_chat_template(
    [{"role": "user", "content": prompt}],
    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
)
new_tokens = output.sequences[0, input_ids.shape[1]:]
print(tokenizer.decode(new_tokens, skip_special_tokens=True))
print(f"acceptance length: {output.acceptance_length:.2f}")

Run the drafter and the target on one GPU. At temperature > 0 the output follows the target's sampling distribution exactly; num_steps=2 drafts with two jumps of the flow map.

vLLM

Install

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

Start the server

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

Send a request

curl OpenAI

curl http://localhost:8000/v1/chat/completions \
  -H "Content-Type: application/json" \
  -d '{
    "model": "Qwen/Qwen3-8B",
    "messages": [{
      "role": "user",
      "content": "How many positive whole-number divisors does 196 have?"
    }],
    "temperature": 0,
    "chat_template_kwargs": {"enable_thinking": false}
  }'

from openai import OpenAI

client = OpenAI(base_url="http://localhost:8000/v1", api_key="none")
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)

SGLang

Install

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

Start the server

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

Send a request

curl OpenAI

curl http://localhost:30000/v1/chat/completions \
  -H "Content-Type: application/json" \
  -d '{
    "model": "Qwen/Qwen3-8B",
    "messages": [{
      "role": "user",
      "content": "How many positive whole-number divisors does 196 have?"
    }],
    "temperature": 0,
    "chat_template_kwargs": {"enable_thinking": false}
  }'

from openai import OpenAI

client = OpenAI(base_url="http://localhost:30000/v1", api_key="none")
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)

Citation

BibTeX

@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},
  note   = {Preprint, coming soon}
}