flashmap 0.1.0__py3-none-any.whl

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/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"]

flashmap/modeling_flashmap.py

@@ -0,0 +1,556 @@
+"""FlashMap: speculative decoding with a flow-map drafter.
+
+FlashMap drafts the ``K`` tokens after the last verified token (the anchor) in one
+network evaluation, like DFlash, but from a generative model that can couple the
+slots of a block. A block is a point ``x`` in ``R^{K x V}`` (``V`` the target's
+vocabulary) moving along the path ``x_s = (1 - s) x0 + s onehot(y)`` from Gaussian
+noise ``x0`` at ``s = 0`` to the one-hot tokens ``y`` at ``s = 1``. The drafter is a
+*flow map* (Flow Map Language Models, arXiv:2602.16813): a network that jumps along
+this path from time ``s`` to time ``t``, distilled from a flow trained on it. Its
+output is a distribution over the vocabulary for every slot. Drafting with one step
+is a single jump from ``s = 0`` to ``t = 1``:
+
+    x0 ~ N(0, I)  ->  logits = map(x0; 0 -> 1)  ->  y_k ~ softmax(logits_k) per slot.
+
+Given the noise the slots are independent draws, so ``softmax(logits)`` is an exact
+proposal distribution for lossless speculative (rejection) sampling. Across noise
+draws the slots can be correlated, since the same noise conditions all of them; one
+jump captures this joint structure only partly, and more jumps capture more. More
+steps jump over a grid ``0 = t_0 < t_1 < ... < t_n = 1``, moving the state to
+``x <- ((1 - t) x + (t - s) softmax(logits)) / (1 - s)`` between jumps.
+
+The network is a small DiT. It reads the block's state through the target's token
+embedding ``E`` (the noise enters as ``x0 @ E``), conditions on the jump's start and
+end times through adaptive LayerNorms, self-attends across the block's slots, and
+cross-attends to the target's hidden states at earlier positions (from a few of its
+layers) and to the anchor's embedding. Its output goes through the target's lm_head.
+The target's embedding and lm_head are not part of the drafter's weights.
+
+``FlashMapDraftModel.spec_generate`` runs lossless speculative decoding with a
+transformers target model. This module is inference-only.
+"""
+
+import math
+import warnings
+from contextlib import contextmanager
+from dataclasses import dataclass, field
+from typing import Optional
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+from transformers import DynamicCache, PreTrainedModel
+
+from .configuration_flashmap import FlashMapConfig
+
+
+def rope_tables(positions: torch.Tensor, head_dim: int, theta: float) -> tuple[torch.Tensor, torch.Tensor]:
+    """``cos`` and ``sin`` of the rotary angles at ``positions (B, T)``, each
+    ``(B, T, 1, head_dim // 2)`` in float32; computed once and shared by all layers."""
+    inv_freq = 1.0 / theta ** (torch.arange(0, head_dim, 2, device=positions.device).float() / head_dim)
+    angles = positions.float()[..., None, None] * inv_freq
+    return angles.cos(), angles.sin()
+
+
+def apply_rope(x: torch.Tensor, rope: tuple[torch.Tensor, torch.Tensor]) -> torch.Tensor:
+    """Rotate the interleaved channel pairs ``(x[2i], x[2i + 1])`` of ``x (B, T, heads,
+    head_dim)`` by the angles of ``rope_tables``, in float32."""
+    cos, sin = rope
+    x1, x2 = x.float().unflatten(-1, (-1, 2)).unbind(-1)
+    return torch.stack([x1 * cos - x2 * sin, x2 * cos + x1 * sin], dim=-1).flatten(-2).type_as(x)
+
+
+class FlashMapRMSNorm(nn.Module):
+    """RMSNorm whose scale ``weight / rms(x)`` is rounded to the input's dtype before
+    it multiplies the input (the order the released weights were trained with)."""
+
+    def __init__(self, dim: int, eps: float):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(dim))
+        self.eps = eps
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        rms = torch.rsqrt(x.float().pow(2).mean(dim=-1, keepdim=True) + self.eps)
+        return x * (self.weight.float() * rms).to(x.dtype)
+
+
+class FlashMapTimeEmbedding(nn.Module):
+    """Sinusoidal features of a time in [0, 1], then an MLP."""
+
+    def __init__(self, config: FlashMapConfig):
+        super().__init__()
+        self.dim, self.scale = config.time_embed_dim, config.time_scale
+        self.linear_1 = nn.Linear(config.time_embed_dim, config.cond_dim)
+        self.linear_2 = nn.Linear(config.cond_dim, config.cond_dim)
+
+    def forward(self, t: torch.Tensor) -> torch.Tensor:
+        half = self.dim // 2
+        freqs = torch.exp(-math.log(10000.0) * torch.arange(half, device=t.device).float() / half)
+        args = (t.float() * self.scale)[..., None] * freqs
+        emb = torch.cat([args.cos(), args.sin()], dim=-1).to(self.linear_1.weight.dtype)
+        return self.linear_2(F.silu(self.linear_1(emb)))
+
+
+class FlashMapAdaLayerNorm(nn.Module):
+    """LayerNorm modulated by the time conditioning (AdaLN-zero)."""
+
+    def __init__(self, config: FlashMapConfig):
+        super().__init__()
+        self.linear = nn.Linear(config.cond_dim, 3 * config.hidden_size)
+        self.eps = config.layer_norm_eps
+
+    def modulation(self, cond: torch.Tensor) -> tuple[torch.Tensor, ...]:
+        """``(scale, shift, gate)`` for the time conditioning ``cond``."""
+        return self.linear(F.silu(cond)).chunk(3, dim=-1)
+
+    def forward(self, x: torch.Tensor, modulation) -> tuple[torch.Tensor, torch.Tensor]:
+        """``norm(x) * (1 + scale) + shift``, and the gate of the residual branch."""
+        scale, shift, gate = modulation
+        return F.layer_norm(x, x.shape[-1:], eps=self.eps) * (1 + scale) + shift, gate
+
+
+class FlashMapAttention(nn.Module):
+    """Multi-head attention with RMS-normed queries and keys and rotary positions.
+
+    Keys and values come from ``project_kv`` so that cross-attention can project its
+    context once and cache it."""
+
+    def __init__(self, config: FlashMapConfig):
+        super().__init__()
+        dim = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        self.head_dim = dim // self.num_heads
+        self.q_proj = nn.Linear(dim, dim, bias=False)
+        self.k_proj = nn.Linear(dim, dim, bias=False)
+        self.v_proj = nn.Linear(dim, dim, bias=False)
+        self.o_proj = nn.Linear(dim, dim, bias=False)
+        self.q_norm = FlashMapRMSNorm(self.head_dim, config.qk_norm_eps)
+        self.k_norm = FlashMapRMSNorm(self.head_dim, config.qk_norm_eps)
+
+    def _heads(self, x: torch.Tensor) -> torch.Tensor:
+        return x.view(*x.shape[:2], self.num_heads, self.head_dim)
+
+    def project_kv(self, x: torch.Tensor, rope) -> tuple[torch.Tensor, torch.Tensor]:
+        """Keys and values ``(B, heads, T, head_dim)`` of ``x (B, T, dim)``."""
+        keys = apply_rope(self.k_norm(self._heads(self.k_proj(x))), rope)
+        return keys.transpose(1, 2), self._heads(self.v_proj(x)).transpose(1, 2)
+
+    def forward(self, x, rope, keys, values) -> torch.Tensor:
+        queries = apply_rope(self.q_norm(self._heads(self.q_proj(x))), rope)
+        out = F.scaled_dot_product_attention(queries.transpose(1, 2), keys, values)
+        return self.o_proj(out.transpose(1, 2).flatten(2))
+
+
+class FlashMapMLP(nn.Module):
+    def __init__(self, config: FlashMapConfig):
+        super().__init__()
+        self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size)
+        self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.fc2(F.gelu(self.fc1(x)))
+
+
+class FlashMapLayer(nn.Module):
+    """A DiT block: self-attention across the block's slots, cross-attention to the
+    context, and an MLP, each behind a time-modulated LayerNorm and a gated residual."""
+
+    def __init__(self, config: FlashMapConfig):
+        super().__init__()
+        self.self_attn_adaln = FlashMapAdaLayerNorm(config)
+        self.self_attn = FlashMapAttention(config)
+        self.cross_attn_adaln = FlashMapAdaLayerNorm(config)
+        self.cross_attn = FlashMapAttention(config)
+        self.mlp_adaln = FlashMapAdaLayerNorm(config)
+        self.mlp = FlashMapMLP(config)
+
+    def forward(self, x, modulations, rope, context: tuple[torch.Tensor, torch.Tensor]) -> torch.Tensor:
+        self_mod, cross_mod, mlp_mod = modulations
+        h, gate = self.self_attn_adaln(x, self_mod)
+        x = x + gate * self.self_attn(h, rope, *self.self_attn.project_kv(h, rope))
+        h, gate = self.cross_attn_adaln(x, cross_mod)
+        x = x + gate * self.cross_attn(h, rope, *context)
+        h, gate = self.mlp_adaln(x, mlp_mod)
+        return x + gate * self.mlp(h)
+
+
+class FlashMapCache:
+    """Cross-attention keys and values of the drafter's context, for every DiT layer.
+
+    Positions ``[0, length)`` hold the target's hidden states at the tokens verified
+    so far. Position ``length`` belongs to the anchor, the last token, whose hidden
+    state the target has not computed yet: drafting writes the anchor's embedding
+    there, and the next ``extend_context`` overwrites it with its hidden state.
+    Create it on the drafter's device and dtype.
+    """
+
+    def __init__(self, config: FlashMapConfig, batch_size: int, max_length: int, device=None, dtype=None):
+        head_dim = config.hidden_size // config.num_attention_heads
+        shape = (config.num_hidden_layers, 2, batch_size, config.num_attention_heads, max_length, head_dim)
+        self.kv = torch.zeros(shape, device=device, dtype=dtype)
+        self.length = 0
+
+    def write(self, layer: int, start: int, keys: torch.Tensor, values: torch.Tensor) -> None:
+        end = start + keys.shape[2]
+        if end > self.kv.shape[4]:
+            raise ValueError(f"FlashMapCache holds {self.kv.shape[4]} positions, writing up to {end}")
+        self.kv[layer, 0, :, :, start:end] = keys
+        self.kv[layer, 1, :, :, start:end] = values
+
+    def context(self, end: int) -> list[tuple[torch.Tensor, torch.Tensor]]:
+        return [(kv[0, :, :, :end], kv[1, :, :, :end]) for kv in self.kv]
+
+
+@dataclass
+class FlashMapGenerateOutput:
+    """``sequences``: prompt and generated tokens, ``(1, length)``.
+    ``acceptance_lengths``: tokens committed by each target forward pass (accepted
+    draft tokens plus the target's own token). The last pass counts in full even
+    where it runs past ``max_new_tokens``, as in the evaluation of the paper."""
+
+    sequences: torch.LongTensor
+    num_new_tokens: int
+    acceptance_lengths: list = field(default_factory=list)
+
+    @property
+    def acceptance_length(self) -> float:
+        """Mean tokens per target forward pass (the paper's gamma)."""
+        return sum(self.acceptance_lengths) / max(len(self.acceptance_lengths), 1)
+
+
+class FlashMapDraftModel(PreTrainedModel):
+    config_class = FlashMapConfig
+    base_model_prefix = "flashmap"
+    _no_split_modules = ["FlashMapLayer"]
+    _supports_sdpa = True
+
+    def __init__(self, config: FlashMapConfig):
+        super().__init__(config)
+        dim, target_dim = config.hidden_size, config.target_hidden_size
+        self.head_dim = dim // config.num_attention_heads
+        # Context: the target's hidden states at target_layer_ids, and the anchor.
+        self.fc = nn.Linear(len(config.target_layer_ids) * target_dim, dim, bias=False)
+        self.hidden_norm = FlashMapRMSNorm(dim, config.rms_norm_eps)
+        self.anchor_proj = nn.Linear(target_dim, dim)
+        self.anchor_norm = FlashMapRMSNorm(dim, config.rms_norm_eps)
+        # State featurization (see embed_state).
+        self.state_norm = FlashMapRMSNorm(target_dim, config.rms_norm_eps)
+        self.noise_norm = FlashMapRMSNorm(target_dim, config.rms_norm_eps)
+        self.noise_gate = nn.Parameter(torch.ones(1))
+        # The DiT.
+        self.x_embed = nn.Linear(target_dim, dim)
+        self.time_embed = FlashMapTimeEmbedding(config)  # the jump's start time s
+        self.jump_embed = FlashMapTimeEmbedding(config)  # the jump's end time t
+        self.layers = nn.ModuleList(FlashMapLayer(config) for _ in range(config.num_hidden_layers))
+        self.final_adaln = FlashMapAdaLayerNorm(config)
+        self.head = nn.Linear(dim, target_dim)
+        self._clear_caches()
+        self.post_init()
+
+    # Values computed from the weights once and reused by every draft: the AdaLN
+    # modulations of each jump, and constants of the target's embedding. Moving or
+    # reloading the weights clears them.
+
+    def _clear_caches(self) -> None:
+        self._time_cache, self._target_cache = {}, {}
+
+    def _apply(self, fn, *args, **kwargs):
+        self._clear_caches()
+        return super()._apply(fn, *args, **kwargs)
+
+    def load_state_dict(self, *args, **kwargs):
+        self._clear_caches()
+        return super().load_state_dict(*args, **kwargs)
+
+    @torch.no_grad()
+    def time_conditioning(self, s: float, t: float):
+        """The AdaLN modulations of the jump from ``s`` to ``t``, for every norm of the
+        DiT: per layer (self-attention, cross-attention, MLP), then the final norm. They
+        depend on the two times alone, so each jump computes them once."""
+        if (s, t) not in self._time_cache:
+            times = torch.tensor([s, t], device=self.device)
+            cond = (self.time_embed(times[:1]) + self.jump_embed(times[1:]))[None]
+            layers = [
+                [norm.modulation(cond) for norm in (block.self_attn_adaln, block.cross_attn_adaln, block.mlp_adaln)]
+                for block in self.layers
+            ]
+            self._time_cache[(s, t)] = (layers, self.final_adaln.modulation(cond))
+        return self._time_cache[(s, t)]
+
+    def forward(self, inputs_embeds, s: float, t: float, position_ids, context) -> torch.Tensor:
+        """The jump from time ``s`` to ``t``: the featurized state ``inputs_embeds (B, K,
+        target_hidden)`` of blocks at ``position_ids (B, K)`` -> hidden states ``(B, K,
+        target_hidden)`` whose target lm_head logits are the distribution it lands on.
+        ``context`` holds each layer's cross-attention keys and values."""
+        layer_modulations, final_modulation = self.time_conditioning(s, t)
+        rope = rope_tables(position_ids, self.head_dim, self.config.rope_theta)
+        h = self.x_embed(inputs_embeds)
+        for layer, modulations, layer_context in zip(self.layers, layer_modulations, context):
+            h = layer(h, modulations, rope, layer_context)
+        h, _ = self.final_adaln(h, final_modulation)
+        return self.head(h)
+
+    # ---- context ------------------------------------------------------------------
+
+    def select_target_hidden(self, hidden_states) -> torch.Tensor:
+        """The drafter's features from a target forward's ``hidden_states`` tuple: the
+        outputs of layers ``target_layer_ids``, concatenated."""
+        return torch.cat([hidden_states[i + 1] for i in self.config.target_layer_ids], dim=-1)
+
+    @contextmanager
+    def target_features(self, target: PreTrainedModel):
+        """Forward hooks that keep the outputs of the target's layers ``target_layer_ids``.
+        Yields a function returning them, concatenated, for the target's last forward.
+        Unlike ``output_hidden_states``, this does not keep every layer's output."""
+        layers, outputs = target.get_decoder().layers, {}
+
+        def keep(i):
+            def hook(module, args, output):
+                outputs[i] = output[0] if isinstance(output, tuple) else output
+
+            return hook
+
+        hooks = [layers[i].register_forward_hook(keep(i)) for i in self.config.target_layer_ids]
+        try:
+            yield lambda: torch.cat([outputs[i] for i in self.config.target_layer_ids], dim=-1)
+        finally:
+            for hook in hooks:
+                hook.remove()
+
+    @torch.no_grad()
+    def extend_context(self, cache: FlashMapCache, target_hidden: torch.Tensor) -> None:
+        """Append the target's features ``(B, n, ...)`` at the next ``n`` verified positions."""
+        start, n = cache.length, target_hidden.shape[1]
+        context = self.hidden_norm(self.fc(target_hidden.to(self.dtype)))
+        positions = torch.arange(start, start + n, device=context.device).expand(context.shape[0], -1)
+        rope = rope_tables(positions, self.head_dim, self.config.rope_theta)
+        for i, layer in enumerate(self.layers):
+            cache.write(i, start, *layer.cross_attn.project_kv(context, rope))
+        cache.length += n
+
+    # ---- drafting -----------------------------------------------------------------
+
+    def embed_state(self, state, s: float, t: float, embedding: torch.Tensor) -> torch.Tensor:
+        """DiT input for a vocabulary-sized state ``(B, K, V)`` at time ``s`` jumping to ``t``.
+
+        The state is read through the target's token embedding twice: as the posterior
+        over clean tokens it implies, ``softmax(s x / (1 - s)^2)``, and linearly, weighted
+        by the jump. At ``s = 0`` the posterior is uniform whatever the noise, so the
+        linear term is how a jump from noise sees which noise it starts from."""
+        if s == 0:
+            posterior_embeds = self.target_constants(embedding)[0]
+        else:
+            posterior = (state.float() * (s / max(1.0 - s, 1e-3) ** 2)).softmax(dim=-1)
+            posterior_embeds = posterior.to(embedding.dtype) @ embedding
+        return self.mix_state(posterior_embeds, state.to(embedding.dtype) @ embedding, s, t)
+
+    def mix_state(self, posterior_embeds, state_embeds, s: float, t: float) -> torch.Tensor:
+        """``embed_state`` from the state's two projections onto the embedding."""
+        # the noise term's weight: noise_gate times the jump ("jump") or alone ("constant")
+        jump = 1.0 if self.config.noise_gate == "constant" else t - s
+        gate = (self.noise_gate.float() * jump).to(self.dtype)
+        posterior_term = self.state_norm(posterior_embeds.to(self.dtype))
+        return posterior_term + gate * self.noise_norm(state_embeds.to(self.dtype))
+
+    @torch.no_grad()
+    def target_constants(self, embedding: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        """What drafting needs of the target's embedding ``E (V, H)``: the embedded
+        uniform posterior (the posterior of any state at ``s = 0``), and ``L`` with
+        ``L L^T = E^T E``.
+
+        A one-step draft reads its noise only as ``x0 @ E`` with ``x0 ~ N(0, I_V)``,
+        a Gaussian with covariance ``E^T E``; ``z @ L^T`` with ``z ~ N(0, I_H)`` has the
+        same distribution and needs no vocabulary-sized noise."""
+        key = (embedding.data_ptr(), tuple(embedding.shape), embedding.device, embedding.dtype)
+        if key not in self._target_cache:
+            vocab_size = embedding.shape[0]
+            uniform = torch.full((1, vocab_size), 1.0 / vocab_size, device=embedding.device)
+            # E^T E in float32 where E is; its (H, H) factor in float64 on the CPU, as
+            # many accelerators lack fast (or any) float64.
+            gram = sum(rows.float().T @ rows.float() for rows in embedding.split(16384))
+            factor = torch.linalg.cholesky(gram.cpu().double())
+            uniform_embeds = uniform.to(embedding.dtype) @ embedding
+            self._target_cache = {key: (uniform_embeds, factor.to(self.device, self.dtype))}
+        return self._target_cache[key]
+
+    @torch.no_grad()
+    def draft(
+        self,
+        anchor_ids: torch.LongTensor,
+        cache: FlashMapCache,
+        embed_tokens: nn.Module,
+        lm_head: nn.Module,
+        num_steps: Optional[int] = None,
+        noise: Optional[torch.Tensor] = None,
+        generator: Optional[torch.Generator] = None,
+    ) -> torch.Tensor:
+        """Draft logits ``(B, block_size, V)`` for the block after ``anchor_ids (B,)``.
+
+        ``cache`` holds the target's features at every position before the anchor.
+        ``embed_tokens`` and ``lm_head`` are the target's. ``noise (B, block_size, V)``
+        is the starting state ``x0``; by default it is drawn, and a one-step draft draws
+        only its embedding (see ``target_constants``). Given the noise, the block's
+        tokens are independent draws from ``softmax(logits)``, one per slot."""
+        num_steps = num_steps or self.config.num_steps
+        if str(num_steps) not in self.config.jump_grids:
+            supported = ", ".join(sorted(self.config.jump_grids, key=int))
+            raise ValueError(f"num_steps must be one of {supported} (the config's jump_grids), got {num_steps}")
+        grid = self.config.jump_grids[str(num_steps)]
+        batch_size, block_size = anchor_ids.shape[0], self.config.block_size
+        embedding, device = embed_tokens.weight, anchor_ids.device
+
+        # The anchor's embedding enters the context at the anchor's position.
+        anchor = cache.length
+        anchor_context = self.anchor_norm(self.anchor_proj(embed_tokens(anchor_ids[:, None]).to(self.dtype)))
+        anchor_position = torch.full((batch_size, 1), anchor, device=device)
+        anchor_rope = rope_tables(anchor_position, self.head_dim, self.config.rope_theta)
+        for i, layer in enumerate(self.layers):
+            cache.write(i, anchor, *layer.cross_attn.project_kv(anchor_context, anchor_rope))
+        context = cache.context(anchor + 1)
+        positions = anchor + 1 + torch.arange(block_size, device=device).expand(batch_size, -1)
+
+        state = noise
+        if state is None and num_steps == 1:
+            uniform_embeds, factor = self.target_constants(embedding)
+            shape = (batch_size, block_size, factor.shape[0])
+            z = torch.randn(shape, device=device, dtype=factor.dtype, generator=generator)
+            inputs = self.mix_state(uniform_embeds, z @ factor.T, 0.0, grid[1])
+        elif state is None:
+            shape = (batch_size, block_size, embedding.shape[0])
+            state = torch.randn(shape, device=device, generator=generator)
+
+        for s, t in zip(grid[:-1], grid[1:]):
+            if state is not None:
+                inputs = self.embed_state(state, s, t, embedding)
+            logits = lm_head(self(inputs, s, t, positions, context).to(lm_head.weight.dtype))
+            if t < 1.0:
+                # move to where the jump lands (eq. 19 of arXiv:2602.16813)
+                state = ((1.0 - t) * state + (t - s) * logits.float().softmax(dim=-1)) / (1.0 - s)
+        return logits
+
+    # ---- speculative decoding -----------------------------------------------------
+
+    def check_target(self, target: PreTrainedModel) -> None:
+        """Raise if ``target`` cannot serve as this drafter's target; warn if it is not
+        the model the drafter was trained for (the output stays lossless, but drafts
+        from another model's hidden states are poor)."""
+        config = target.config.get_text_config()
+        embedding, lm_head = target.get_input_embeddings().weight, target.get_output_embeddings().weight
+        expected = (self.config.vocab_size, self.config.target_hidden_size)
+        if tuple(embedding.shape) != expected or tuple(lm_head.shape) != expected:
+            raise ValueError(
+                f"this drafter needs a target with vocabulary {expected[0]} and hidden size {expected[1]} "
+                f"({self.config.target_model}), got an embedding of shape {tuple(embedding.shape)}"
+            )
+        if max(self.config.target_layer_ids) >= config.num_hidden_layers - 1:
+            layers = config.num_hidden_layers
+            raise ValueError(f"target_layer_ids {self.config.target_layer_ids} need more than {layers} target layers")
+        name = (getattr(config, "_name_or_path", None) or "").rstrip("/")
+        if name and name.split("/")[-1] != self.config.target_model.split("/")[-1]:
+            warnings.warn(f"this drafter was trained for {self.config.target_model}, not {name}")
+        devices = {self.device, embedding.device, lm_head.device}
+        if len(devices) > 1:
+            raise ValueError(f"the drafter and the target's embedding and lm_head must share a device, got {devices}")
+
+    @torch.no_grad()
+    def spec_generate(
+        self,
+        target: PreTrainedModel,
+        input_ids: torch.LongTensor,
+        max_new_tokens: int = 256,
+        temperature: float = 0.0,
+        num_steps: Optional[int] = None,
+        stop_token_ids: Optional[list] = None,
+        generator: Optional[torch.Generator] = None,
+    ) -> FlashMapGenerateOutput:
+        """Lossless speculative decoding of ``target`` (a transformers causal LM) from
+        the prompt ``input_ids (1, P)``.
+
+        Every round drafts a block after the last token, has the target score it in one
+        forward pass, and keeps the longest prefix the target accepts plus one token of
+        the target's own. At ``temperature`` 0 the output is the target's greedy
+        decoding; above it, tokens are distributed exactly as the target samples them
+        (rejection sampling against the draft's per-slot probabilities). Generation stops
+        after ``max_new_tokens`` or at a token in ``stop_token_ids`` (default: the
+        target's ``eos_token_id``). The drafter, the target and ``generator`` (which
+        makes sampling reproducible) must be on one device."""
+        if input_ids.shape[0] != 1:
+            raise ValueError("spec_generate decodes one sequence at a time")
+        self.check_target(target)
+        if generator is not None:  # a "cuda" generator (no index) is on the current device
+            device = generator.device
+            if device.type == "cuda" and device.index is None:
+                device = torch.device("cuda", torch.cuda.current_device())
+            if (device.type, device.index or 0) != (self.device.type, self.device.index or 0):
+                raise ValueError(f"generator must be on {self.device}, got one on {generator.device}")
+        if stop_token_ids is None:
+            eos = target.generation_config.eos_token_id
+            stop_token_ids = [] if eos is None else [eos] if isinstance(eos, int) else list(eos)
+        input_ids = input_ids.to(self.device)
+        stop = torch.tensor(stop_token_ids, device=self.device, dtype=torch.long)
+        embed_tokens, lm_head = target.get_input_embeddings(), target.get_output_embeddings()
+        block_size, prompt_length = self.config.block_size, input_ids.shape[1]
+
+        def sample(logits):  # (1, n, V) -> tokens (1, n) and their distributions
+            if temperature <= 0:
+                return logits.argmax(dim=-1), None
+            probs = (logits.float() / temperature).softmax(dim=-1)
+            return torch.multinomial(probs[0], 1, generator=generator).T, probs
+
+        with self.target_features(target) as target_features:
+            target_cache = DynamicCache()
+            out = target(input_ids, past_key_values=target_cache, use_cache=True, logits_to_keep=1)
+            tokens = sample(out.logits[:, -1:])[0]
+            max_length = prompt_length + max_new_tokens + block_size + 1  # the last round may overshoot
+            cache = FlashMapCache(self.config, 1, max_length, device=self.device, dtype=self.dtype)
+            self.extend_context(cache, target_features())
+            acceptance_lengths = []
+
+            while tokens.shape[1] < max_new_tokens and not torch.isin(tokens[0, -1], stop):
+                anchor = tokens[:, -1]
+                draft_logits = self.draft(anchor, cache, embed_tokens, lm_head, num_steps, generator=generator)
+                draft, draft_probs = sample(draft_logits)
+                position = prompt_length + tokens.shape[1] - 1  # the anchor's
+                out = target(
+                    torch.cat([anchor[:, None], draft], dim=1),
+                    position_ids=torch.arange(position, position + block_size + 1, device=draft.device)[None],
+                    past_key_values=target_cache,
+                    use_cache=True,
+                )
+                if temperature <= 0:
+                    target_tokens = out.logits.argmax(dim=-1)
+                    accepted = int((draft == target_tokens[:, :-1]).cumprod(dim=1).sum())
+                    bonus = target_tokens[:, accepted]
+                else:
+                    target_probs = (out.logits.float() / temperature).softmax(dim=-1)[0]
+                    p = target_probs[:-1].gather(-1, draft[0, :, None])[:, 0]
+                    q = draft_probs[0].gather(-1, draft[0, :, None])[:, 0]
+                    u = torch.rand(q.shape, device=q.device, generator=generator)
+                    accepted = int((u * q < p).cumprod(dim=0).sum())
+                    residual = target_probs[accepted]
+                    if accepted < block_size:  # resample the rejected slot from max(p - q, 0)
+                        residual = (residual - draft_probs[0, accepted]).clamp_min(0)
+                        if residual.sum() <= 0:
+                            residual = target_probs[accepted]
+                    bonus = torch.multinomial(residual, 1, generator=generator)
+                new_tokens = torch.cat([draft[:, :accepted], bonus.view(1, 1)], dim=1)
+                hits = torch.isin(new_tokens[0], stop).nonzero()
+                if len(hits):
+                    new_tokens = new_tokens[:, : int(hits[0]) + 1]
+                tokens = torch.cat([tokens, new_tokens], dim=1)
+                acceptance_lengths.append(new_tokens.shape[1])
+                # keep the anchor and the accepted drafts; the bonus token is the next anchor
+                excess = target_cache.get_seq_length() - (position + accepted + 1)
+                if excess > 0:  # crop(-n) drops n tokens in every transformers version
+                    target_cache.crop(-excess)
+                self.extend_context(cache, target_features()[:, : accepted + 1])
+
+        tokens = tokens[:, :max_new_tokens]
+        return FlashMapGenerateOutput(
+            sequences=torch.cat([input_ids, tokens], dim=1),
+            num_new_tokens=tokens.shape[1],
+            acceptance_lengths=acceptance_lengths,
+        )
+
+
+__all__ = ["FlashMapCache", "FlashMapDraftModel", "FlashMapGenerateOutput"]