FlashMap-Qwen3-8B / modeling_flashmap.py
FlashMap release
ac003c0 verified 6 days ago
29.6 kB
"""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"]