ASD: Approximate Speculative Decoding custom generation method

.gitignore

@@ -0,0 +1,27 @@
+__pycache__/
+*.py[cod]
+*$py.class
+*.so
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+.pytest_cache/
+.ipynb_checkpoints
+*.log
+*.pt
+*.safetensors
+checkpoints/
+models/

README.md

@@ -0,0 +1,99 @@
+---
+tags:
+  - custom_generate
+library_name: transformers
+---
+
+# ASD — Approximate Speculative Decoding
+
+A training-free, verifier-side acceptance policy for **greedy assisted decoding** in Hugging Face Transformers.
+
+Standard greedy speculative decoding discards the entire draft suffix at the first position where the draft token differs from the target argmax. ASD relaxes this strictly deterministic check with three explicit, bounded rules so that near-optimal draft tokens can still be accepted.
+
+> Paper: [ASD: Approximate Speculative Decoding](https://arxiv.org/abs/2608.03447)
+> Reference implementation: https://github.com/Kissmetothemoon/ASD
+
+## Usage
+
+```python
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+model_id = "Qwen/Qwen3-14B"
+assistant_id = "your-draft-model"
+
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto")
+assistant_model = AutoModelForCausalLM.from_pretrained(assistant_id, device_map="auto")
+
+inputs = tokenizer("The quick brown", return_tensors="pt").to(model.device)
+
+outputs = model.generate(
+    **inputs,
+    max_new_tokens=64,
+    assistant_model=assistant_model,
+    custom_generate="ynFeng/approximate-speculative-decoding",
+    trust_remote_code=True,
+    assistant_asd_budget=2.0,
+    assistant_asd_local_ratio=0.25,
+    assistant_asd_max_mismatches=2,
+)
+```
+
+## Parameters
+
+| Name | Type | Default | Description |
+|------|------|---------|-------------|
+| `assistant_asd_budget` | `float` | `None` | Request-level cumulative regret budget `B`. `None` = unlimited (only ratio / mismatch caps apply). `0.0` forces strict argmax verification. |
+| `assistant_asd_local_ratio` | `float` | `None` | Max allowed `r_i / (K - i)` for a relaxed mismatch at draft position `i`, where `K` is the candidate length. `None` = unlimited. |
+| `assistant_asd_max_mismatches` | `int` | `None` | Max number of relaxed mismatches accepted per block. `None` = unlimited. `0` forces strict verification. |
+
+If any parameter is `0`, the policy **exactly reduces to strict greedy verification** — this is the built-in correctness baseline.
+
+## Rules
+
+For each drafted token `x_i` with target logits `z_i`:
+
+1. `r_i = max_v z_i(v) - z_i(x_i)` (`r_i = 0` iff `x_i` is the target argmax).
+2. Accept the draft prefix while **all** hold:
+   - request-level cumulative regret `<= assistant_asd_budget`
+   - local regret `r_i / (K - i) <= assistant_asd_local_ratio`
+   - relaxed mismatches in the current block `<= assistant_asd_max_mismatches`
+3. On the first infeasible position, commit the target argmax as the bonus token.
+
+## Empirical results (paper)
+
+Qwen3-14B + DSpark block-7 draft, 8×L20, greedy:
+
+| Setting | Throughput | GSM8K |
+|---------|-----------|-------|
+| Strict verification | 66.21 tok/s | 79.682 % |
+| ASD (`B=2.0625`) | 75.04 tok/s (+13.34 %) | 79.454 % (-0.227 pp) |
+| ASD (`B=0`) | token-identical to strict on all 1319 requests | identical |
+
+## Requirements
+
+```text
+transformers>=5.19.0,<5.20.0
+torch>=2.0.0
+```
+
+The acceptance loop mirrors `GenerationMixin._assisted_decoding` internals, so the
+transformers minor version is pinned. Verified on `transformers==5.19.0`, `torch==2.6.0`.
+
+## Validation
+
+Checked with Qwen2.5-1.5B-Instruct (target) + Qwen2.5-0.5B-Instruct (draft) on an
+NVIDIA L20 (`scripts/example.py` reproduces):
+
+- `assistant_asd_budget=0` and no-knob calls are **token-identical** to native
+  greedy assisted decoding on every tested prompt (also with `prompt_lookup_num_tokens`);
+- relaxed runs are deterministic across repeated calls (no RNG anywhere in the policy);
+- `do_sample=True`, combining with `assistant_ensemble_weight`, ASD knobs without a
+  draft source, and negative budgets all raise `ValueError`.
+
+ASD also works with prompt-lookup drafting (`prompt_lookup_num_tokens`) — unlike
+`assistant_ensemble_weight`, the regret rule needs target logits only.
+
+## License
+
+Apache-2.0

custom_generate/generate.py

@@ -0,0 +1,600 @@
+"""Approximate Speculative Decoding (ASD) as a Transformers custom generation method.
+
+ASD is a training-free, deterministic acceptance policy for *greedy* assisted
+decoding. Instead of discarding the whole draft suffix at the first position
+where the draft token differs from the target argmax, it accepts a draft prefix
+while ALL of the following hold:
+
+1. request-level cumulative regret ``sum(r_j) <= assistant_asd_budget``;
+2. local regret ``r_i / (K - i) <= assistant_asd_local_ratio``
+   (suffix-value weighting: later draft positions get less slack);
+3. relaxed mismatches in the current block ``<= assistant_asd_max_mismatches``.
+
+Here ``r_i = max_v z_i(v) - z_i(x_i)`` is the local regret of draft token
+``x_i`` under target logits ``z_i`` (``r_i = 0`` iff ``x_i`` is the target
+argmax) and ``K`` is the candidate length. The policy uses *target logits only*
+(no draft logits, no RNG) and exactly reduces to strict greedy verification
+when the budget, the local ratio, or the mismatch cap is ``0``.
+
+Paper: "ASD: Approximate Speculative Decoding", https://arxiv.org/abs/2608.03447
+
+Implementation notes
+--------------------
+`generate` below mirrors the signature of `GenerationMixin.generate` (this is
+the contract for Hub custom generation methods). It validates the ASD knobs and
+then temporarily swaps `type(model)._assisted_decoding` for the ASD loop in
+this file, so the standard `generate()` pipeline (logits processors, stopping
+criteria, cache preparation, streamer, assistant handling, ...) runs unchanged.
+The swap is reverted in a `finally` block.
+
+`_asd_assisted_decoding` is a byte-faithful copy of
+`transformers.generation.utils.GenerationMixin._assisted_decoding`
+(transformers v5.19.0) with exactly three changes, marked `ASD-CHANGE`:
+  1. request-level risk state is initialised before the decode loop;
+  2. the strict greedy prefix match (upstream "Case 2") is replaced by the ASD
+     acceptance rule (the sampling path is rejected in `generate`, so upstream
+     "Case 1" is unreachable here);
+  3. the risk account is charged after the EOS/length crop, only for draft
+     tokens that were actually committed.
+"""
+
+from __future__ import annotations
+
+import copy
+import math
+from typing import TYPE_CHECKING, Any, NamedTuple
+
+import torch
+from transformers import GenerationMixin
+from transformers.cache_utils import StaticCache
+from transformers.generation.candidate_generator import (
+    AssistedCandidateGenerator,
+    MTPCandidateGenerator,
+    _prepare_attention_mask,
+    _prepare_position_ids,
+    _prepare_token_type_ids,
+)
+from transformers.generation.utils import (
+    ALL_CACHE_NAMES,
+    GenerateDecoderOnlyOutput,
+    GenerateEncoderDecoderOutput,
+    _split_model_outputs,
+)
+
+if TYPE_CHECKING:
+    from transformers import (
+        GenerationConfig,
+        LogitsProcessorList,
+        PreTrainedModel,
+        PreTrainedTokenizerBase,
+        StoppingCriteriaList,
+    )
+    from transformers.generation import BaseStreamer
+    from transformers.generation.utils import GenerateNonBeamOutput
+
+
+class _ASDVerification(NamedTuple):
+    """Resolved ASD acceptance parameters for one request."""
+
+    budget: float  # request-level cumulative regret budget B (inf = uncapped)
+    local_ratio: float  # per-token regret / suffix-value cap g (inf = uncapped)
+    max_mismatches: int  # per-block relaxed mismatch cap m
+    strict_identity: bool  # True => exactly upstream strict argmax verification
+
+
+_ASD_CONFIG_ATTR = "_asd_verification"
+
+
+def _resolve_asd_params(
+    budget: Any,
+    local_ratio: Any,
+    max_mismatches: Any,
+) -> _ASDVerification:
+    """Validate the user-facing knobs and return the resolved acceptance config.
+
+    All three knobs ``None`` means "ASD disabled": the policy is strict argmax
+    verification, i.e. token-identical to native assisted decoding. When any
+    knob is set, the unset ones default to their permissive limit. A value of
+    ``0`` for any knob forces strict verification (the paper's correctness
+    baseline).
+    """
+    all_none = budget is None and local_ratio is None and max_mismatches is None
+
+    if budget is not None:
+        budget = float(budget)
+        if not math.isfinite(budget) or budget < 0:
+            raise ValueError(f"`assistant_asd_budget` must be a finite float >= 0, got {budget}.")
+    else:
+        budget = math.inf
+
+    if local_ratio is not None:
+        local_ratio = float(local_ratio)
+        if not math.isfinite(local_ratio) or local_ratio < 0:
+            raise ValueError(
+                f"`assistant_asd_local_ratio` must be a finite float >= 0, got {local_ratio}."
+            )
+    else:
+        local_ratio = math.inf
+
+    if max_mismatches is not None:
+        max_mismatches = int(max_mismatches)
+        if max_mismatches < 0:
+            raise ValueError(f"`assistant_asd_max_mismatches` must be an int >= 0, got {max_mismatches}.")
+    else:
+        max_mismatches = 2**31 - 1
+
+    strict_identity = all_none or budget == 0.0 or local_ratio == 0.0 or max_mismatches == 0
+    return _ASDVerification(budget, local_ratio, max_mismatches, strict_identity)
+
+
+def _effective_generation_value(name: str, kwargs: dict, generation_config, model):
+    """Resolve a generation flag the way `generate()` would: kwargs > config arg > model config."""
+    if name in kwargs:
+        return kwargs[name]
+    if generation_config is not None and getattr(generation_config, name, None) is not None:
+        return getattr(generation_config, name)
+    return getattr(model.generation_config, name, None)
+
+
+def _choose_asd_prefix(
+    candidate_new_tokens: torch.LongTensor,
+    new_logits: torch.FloatTensor,
+    target_argmax: torch.LongTensor,
+    config: _ASDVerification,
+) -> tuple[int, list[float], list[bool]]:
+    """Return ``(accepted, regrets, mismatches)`` for one draft block.
+
+    ``candidate_new_tokens`` is ``(1, K)``, ``new_logits`` is ``(1, K, V)`` (the
+    processed target logits at each drafted position) and ``target_argmax`` is
+    ``(1, K)``: the target argmax token per position, computed with the exact
+    same ``argmax`` op upstream strict verification uses, so the strict path is
+    token-identical by construction. ``accepted`` is the longest prefix
+    satisfying all ASD constraints; ``regrets``/``mismatches`` are the
+    per-position values so the caller can charge the budget after the
+    EOS/length crop, only for tokens that were actually committed.
+    """
+    candidate_length = candidate_new_tokens.shape[1]
+    if candidate_length == 0:
+        return 0, [], []
+
+    # r_i = max_v z_i(v) - z_i(x_i); mismatch_i = (x_i != argmax_v z_i(v))
+    top_logits = new_logits[0].max(dim=-1).values
+    draft_logits = new_logits[0].gather(-1, candidate_new_tokens[0].unsqueeze(-1)).squeeze(-1)
+    regrets = (top_logits - draft_logits).clamp_min(0.0).tolist()
+    mismatches = (candidate_new_tokens[0] != target_argmax[0]).tolist()
+
+    # `config.budget` is the *remaining* request budget for this block: the caller
+    # (`_asd_assisted_decoding`) subtracts the running spend before calling.
+    remaining = max(0.0, config.budget)
+    cumulative_risk = 0.0
+    cumulative_relaxed = 0
+    accepted = 0
+    for index in range(candidate_length):
+        mismatch = mismatches[index]
+        regret = regrets[index]
+        value = float(candidate_length - index)  # suffix-value weighting
+
+        if mismatch:
+            if config.strict_identity:
+                break
+            cumulative_risk += regret
+            cumulative_relaxed += 1
+
+        worthwhile = (not mismatch) or (regret / value <= config.local_ratio)
+        affordable = cumulative_risk <= remaining + 1e-9
+        within_cap = cumulative_relaxed <= config.max_mismatches
+
+        if not (worthwhile and affordable and within_cap):
+            break
+        accepted = index + 1
+
+    return accepted, regrets, mismatches
+
+
+def _asd_assisted_decoding(
+    self: "PreTrainedModel",
+    input_ids: torch.LongTensor,
+    logits_processor: "LogitsProcessorList",
+    stopping_criteria: "StoppingCriteriaList",
+    generation_config: "GenerationConfig",
+    synced_gpus: bool = False,
+    streamer: "BaseStreamer | None" = None,
+    inputs_tensor: torch.FloatTensor | None = None,
+    assistant_model: "PreTrainedModel | None" = None,
+    assistant_tokenizer: "PreTrainedTokenizerBase | None" = None,
+    tokenizer: "PreTrainedTokenizerBase | None" = None,
+    **model_kwargs,
+) -> "GenerateNonBeamOutput | torch.LongTensor":
+    """Drop-in replacement for `GenerationMixin._assisted_decoding` (greedy only).
+
+    Identical to upstream except for the ASD changes marked `ASD-CHANGE`.
+    """
+    asd: _ASDVerification | None = getattr(self, _ASD_CONFIG_ATTR, None)
+    if asd is None:  # defensive: behave exactly like upstream strict verification
+        asd = _ASDVerification(math.inf, math.inf, 2**31 - 1, True)
+
+    # The cache must be dynamic for assisted generation, and the check must happen AFTER preparing cache
+    if not model_kwargs["use_cache"]:
+        raise ValueError("assisted generate requires `use_cache=True`")
+    if (
+        generation_config.cache_implementation in ["static", "hybrid", "sliding_window"]
+        or type(model_kwargs.get("past_key_values")) is StaticCache
+    ):
+        raise ValueError("assisted generate is not supported with Static cache classes`")
+
+    # Same tensor the stopping criteria are built from
+    eos_token_id = getattr(generation_config, "_eos_token_tensor", None)
+
+    # Make sure we can record past on the cache
+    cache = model_kwargs.get("past_key_values")
+    if cache is None:
+        raise RuntimeError("assisted decoding requires a cache")
+    cache.activate_past_recording()
+
+    # Get the candidate generator, given the parameterization
+    candidate_generator = self._get_candidate_generator(
+        generation_config=generation_config,
+        input_ids=input_ids,
+        inputs_tensor=inputs_tensor,
+        assistant_model=assistant_model,
+        logits_processor=logits_processor,
+        target_tokenizer=tokenizer,
+        assistant_tokenizer=assistant_tokenizer,
+        model_kwargs=model_kwargs,
+    )
+    # init values
+    output_attentions = generation_config.output_attentions
+    output_hidden_states = generation_config.output_hidden_states
+    output_scores = generation_config.output_scores
+    output_logits = generation_config.output_logits
+    return_dict_in_generate = generation_config.return_dict_in_generate
+
+    # init attention / hidden states / scores tuples
+    scores = () if (return_dict_in_generate and output_scores) else None
+    raw_logits = () if (return_dict_in_generate and output_logits) else None
+    decoder_attentions = () if (return_dict_in_generate and output_attentions) else None
+    cross_attentions = () if (return_dict_in_generate and output_attentions) else None
+    decoder_hidden_states = () if (return_dict_in_generate and output_hidden_states) else None
+
+    # if model is an encoder-decoder, retrieve encoder attention weights and hidden states
+    if return_dict_in_generate and self.config.is_encoder_decoder:
+        encoder_attentions = model_kwargs["encoder_outputs"].get("attentions") if output_attentions else None
+        encoder_hidden_states = (
+            model_kwargs["encoder_outputs"].get("hidden_states") if output_hidden_states else None
+        )
+
+    # keep track of which sequences are already finished
+    batch_size, cur_len = input_ids.shape[:2]
+    if batch_size > 1:
+        raise ValueError("assisted generate is only supported for batch_size = 1")
+    unfinished_sequences = torch.ones(batch_size, dtype=torch.long, device=input_ids.device)
+
+    this_peer_finished = False
+    is_first_iteration = True  # to preserve the same API in the output as other generation methods
+    outputs = None
+    n_matches = 0
+    total_risk_spent = 0.0  # ASD-CHANGE (1): request-level cumulative regret account
+    while self._has_unfinished_sequences(this_peer_finished, synced_gpus, device=input_ids.device):
+        cur_len = input_ids.shape[1]
+
+        #  1. Fetch candidate sequences from a `CandidateGenerator` and move to the correct device
+        if candidate_generator.requires_model_outputs:
+            candidate_input_ids, candidate_logits = candidate_generator.get_candidates(
+                input_ids,
+                model_kwargs=model_kwargs,
+                model_outputs=outputs,
+                is_first_iteration=is_first_iteration,
+                n_last_matches=n_matches,
+            )
+        else:
+            candidate_input_ids, candidate_logits = candidate_generator.get_candidates(input_ids)
+        candidate_input_ids = candidate_input_ids.to(self.device)
+        if candidate_logits is not None:
+            candidate_logits = candidate_logits.to(self.device)
+
+        candidate_length = candidate_input_ids.shape[1] - input_ids.shape[1]
+
+        # 2. Use the original model to obtain the next token logits given the candidate sequence. We obtain
+        # `candidate_length + 1` relevant logits from this process: in the event that all candidates are correct,
+        # we use this forward pass to also pick the subsequent logits in the original model.
+
+        # 2.1. Prepare the model inputs
+        candidate_kwargs = copy.copy(model_kwargs)
+        candidate_kwargs = _prepare_attention_mask(
+            candidate_kwargs, candidate_input_ids.shape[1], self.config.is_encoder_decoder
+        )
+        candidate_kwargs = _prepare_token_type_ids(candidate_kwargs, candidate_input_ids.shape[1])
+        if (position_ids := candidate_kwargs.get("position_ids")) is not None and candidate_length > 0:
+            new_length = candidate_length + position_ids.shape[-1]
+            candidate_kwargs = _prepare_position_ids(candidate_kwargs, new_length, self.config.is_encoder_decoder)
+
+        next_sequence_length = candidate_length + 1 if not is_first_iteration else None
+        model_inputs = self.prepare_inputs_for_generation(
+            candidate_input_ids,
+            next_sequence_length=next_sequence_length,
+            is_first_iteration=is_first_iteration,
+            **candidate_kwargs,
+        )
+
+        if "logits_to_keep" in model_inputs:
+            model_inputs["logits_to_keep"] = candidate_length + 1
+
+        # We usually need to force returning hidden_states and some others if the candidate_generator requires the model's outputs
+        if candidate_generator.requires_model_outputs:
+            model_inputs |= candidate_generator.model_kwargs_overrides
+
+        # 2.2. Run a forward pass on the candidate sequence
+        outputs = self(**model_inputs)
+
+        # 2.3. Process the new logits
+        # .float() is needed to retain precision for later logits manipulations
+        new_logits = outputs.logits[:, -candidate_length - 1 :].to(
+            dtype=torch.float32, device=input_ids.device
+        )  # excludes the input prompt if present
+        next_token_logits = new_logits.clone()
+        if len(logits_processor) > 0:
+            for i in range(candidate_length + 1):
+                new_logits[:, i, :] = logits_processor(candidate_input_ids[:, : cur_len + i], new_logits[:, i, :])
+
+        # 3. Select the accepted tokens.
+        # ASD-CHANGE (2): the greedy verification below replaces upstream's strict argmax
+        # prefix match with the ASD acceptance rule. The sampling path ("Case 1") and the
+        # `assistant_ensemble_weight` variant are unreachable: `generate()` rejects
+        # `do_sample=True` and `assistant_ensemble_weight` before installing this method.
+        candidate_new_tokens = candidate_input_ids[:, cur_len:]
+        # Same op as upstream strict verification, so strict mode matches by construction.
+        target_argmax = new_logits.argmax(dim=-1)  # (1, candidate_length + 1)
+        block_config = _ASDVerification(
+            budget=asd.budget - total_risk_spent,  # remaining request budget for this block
+            local_ratio=asd.local_ratio,
+            max_mismatches=asd.max_mismatches,
+            strict_identity=asd.strict_identity,
+        )
+        accepted, regrets, mismatches = _choose_asd_prefix(
+            candidate_new_tokens, new_logits[:, :-1, :], target_argmax[:, :-1], block_config
+        )
+
+        # Commit the accepted draft prefix, then the target argmax at the first position
+        # after it as the bonus token (deterministic, exactly as upstream: upstream's
+        # `selected_tokens[:, :n_matches + 1]` is the argmax at position `n_matches`).
+        bonus_token = target_argmax[:, accepted : accepted + 1]
+        if accepted > 0:
+            valid_tokens = torch.cat([candidate_new_tokens[:, :accepted], bonus_token], dim=-1)
+        else:
+            valid_tokens = bonus_token
+        n_matches = accepted
+
+        # Whenever we are drafting several tokens at once with the candidate without capping its draft (e.g. MTP), we need to
+        # make sure that we did not just validate tokens outside the max length, or outside an eos token
+        # Note that we should technically crop based on other stopping criteria as well in all generality, not only EOS and Length
+        tokens_budget = generation_config.max_length - input_ids.shape[1]
+        # This is for the max length
+        if valid_tokens.shape[1] > tokens_budget:
+            valid_tokens = valid_tokens[:, :tokens_budget]
+        # This is for eos tokens
+        if eos_token_id is not None:
+            # We are restricted to batch_size == 1, so we can squeeze the batch dim to simplify
+            eos_positions = torch.isin(valid_tokens.squeeze(0), eos_token_id.to(valid_tokens.device)).nonzero()
+            if eos_positions.numel() > 0:
+                num_drafted = eos_positions[0].item() + 1
+                valid_tokens = valid_tokens[:, :num_drafted]
+        # Recompute how many matches we accepted if we just cropped above due to length/eos
+        n_matches = valid_tokens.shape[1] - 1
+
+        # ASD-CHANGE (3): charge the request-level regret account after the crop, only
+        # for relaxed (mismatched) draft tokens that were actually committed.
+        for index in range(max(0, min(n_matches, accepted))):
+            if mismatches[index]:
+                total_risk_spent += regrets[index]
+
+        # 4. Update variables according to the number of matching assistant tokens. Remember: the token generated
+        # by the model after the last candidate match is also valid, as it is generated from a correct sequence.
+        # Because of this last token, assisted generation search reduces to a normal greedy search/sample if there
+        # is no match.
+
+        # 4.1. Get the valid continuation, after the matching tokens
+        input_ids = torch.cat((input_ids, valid_tokens), dim=-1)
+        if streamer is not None:
+            streamer.put(valid_tokens.cpu())
+        new_cur_len = input_ids.shape[1]
+
+        # 4.2. Discard past key values relative to unused assistant tokens if any - if `number_of_tokens_to_crop == 0`, `crop` will
+        # still take care of shrinking back sliding window or linear attention cache layers back to their max size, so it's important
+        # to call it still
+        number_of_tokens_to_crop = candidate_length - n_matches
+        outputs.past_key_values.crop(-number_of_tokens_to_crop)
+
+        # 5. Update the candidate generation strategy if needed
+        candidate_generator.update_candidate_strategy(input_ids, new_logits, n_matches)
+
+        # synced_gpus: don't waste resources running the code we don't need; kwargs must be updated before skipping
+        model_kwargs = self._update_model_kwargs_for_generation(
+            outputs,
+            model_kwargs,
+            is_encoder_decoder=self.config.is_encoder_decoder,
+            num_new_tokens=n_matches + 1,
+        )
+        if synced_gpus and this_peer_finished:
+            continue
+
+        # Store scores, attentions and hidden_states when required
+        # Assistant: modified to append one tuple element per token, as in the other generation methods.
+        if return_dict_in_generate:
+            newly_added_length = n_matches + 1
+            if output_scores:
+                scores += tuple(new_logits[:, i, :] for i in range(newly_added_length))
+            if output_logits:
+                raw_logits += tuple(next_token_logits[:, i, :] for i in range(newly_added_length))
+
+            newly_added_length = new_cur_len if is_first_iteration else newly_added_length
+            if output_attentions:
+                if self.config.is_encoder_decoder:
+                    cross_attentions = _split_model_outputs(
+                        cross_attentions, outputs.cross_attentions, cur_len, newly_added_length
+                    )
+                    decoder_attentions = _split_model_outputs(
+                        decoder_attentions,
+                        outputs.decoder_attentions,
+                        cur_len,
+                        newly_added_length,
+                        is_decoder_attention=True,
+                    )
+                # some (V)LLMs have hard requirement on SDPA and thus never return attn
+                elif outputs.attentions[0] is not None:
+                    decoder_attentions = _split_model_outputs(
+                        decoder_attentions,
+                        outputs.attentions,
+                        cur_len,
+                        newly_added_length,
+                        is_decoder_attention=True,
+                    )
+            if output_hidden_states:
+                if self.config.is_encoder_decoder:
+                    decoder_hidden_states = _split_model_outputs(
+                        decoder_hidden_states, outputs.decoder_hidden_states, cur_len, newly_added_length
+                    )
+                else:
+                    decoder_hidden_states = _split_model_outputs(
+                        decoder_hidden_states, outputs.hidden_states, cur_len, newly_added_length
+                    )
+
+        unfinished_sequences = unfinished_sequences & ~stopping_criteria(input_ids, scores)
+        this_peer_finished = unfinished_sequences.max() == 0
+        is_first_iteration = False
+
+    if streamer is not None:
+        streamer.end()
+
+    if (
+        isinstance(candidate_generator, AssistedCandidateGenerator)
+        and not isinstance(candidate_generator, MTPCandidateGenerator)
+        and candidate_generator.assistant_model.generation_config.num_assistant_tokens_schedule == "heuristic"
+    ):
+        candidate_generator.assistant_model.generation_config.num_assistant_tokens = (
+            candidate_generator.num_assistant_tokens
+        )
+    if return_dict_in_generate:
+        cache = None
+        if any(cache_key in model_kwargs for cache_key in ALL_CACHE_NAMES):
+            cache_key = next(cache_key for cache_key in ALL_CACHE_NAMES if cache_key in model_kwargs)
+            cache = model_kwargs[cache_key]
+        if self.config.is_encoder_decoder:
+            return GenerateEncoderDecoderOutput(
+                sequences=input_ids,
+                scores=scores,
+                logits=raw_logits,
+                encoder_attentions=encoder_attentions,
+                encoder_hidden_states=encoder_hidden_states,
+                decoder_attentions=decoder_attentions,
+                cross_attentions=cross_attentions,
+                decoder_hidden_states=decoder_hidden_states,
+                past_key_values=cache,
+            )
+        else:
+            return GenerateDecoderOnlyOutput(
+                sequences=input_ids,
+                scores=scores,
+                logits=raw_logits,
+                attentions=decoder_attentions,
+                hidden_states=decoder_hidden_states,
+                past_key_values=cache,
+            )
+    else:
+        return input_ids
+
+
+def generate(
+    model: "PreTrainedModel",
+    inputs: torch.Tensor | None = None,
+    generation_config: "GenerationConfig | None" = None,
+    logits_processor: "LogitsProcessorList | None" = None,
+    stopping_criteria: "StoppingCriteriaList | None" = None,
+    prefix_allowed_tokens_fn=None,
+    synced_gpus: bool | None = None,
+    assistant_model: "PreTrainedModel | None" = None,
+    streamer: "BaseStreamer | None" = None,
+    negative_prompt_ids: torch.Tensor | None = None,
+    negative_prompt_attention_mask: torch.Tensor | None = None,
+    assistant_asd_budget: float | None = None,
+    assistant_asd_local_ratio: float | None = None,
+    assistant_asd_max_mismatches: int | None = None,
+    **kwargs,
+):
+    """Custom generation entry point loaded from the Hub.
+
+    Accepts every argument of `GenerationMixin.generate` plus the three ASD
+    knobs below, then runs the standard `generate()` pipeline with the ASD
+    assisted-decoding loop installed. With no ASD knob set (or any knob equal
+    to ``0``), the result is token-identical to native greedy assisted decoding.
+
+    Parameters
+    ----------
+    assistant_asd_budget:
+        Request-level cumulative regret budget ``B >= 0``. ``None`` (default)
+        leaves the budget uncapped; ``0.0`` forces strict verification.
+    assistant_asd_local_ratio:
+        Per-token cap ``g >= 0`` on ``r_i / (K - i)``. ``None`` leaves it
+        uncapped; ``0.0`` forces strict verification.
+    assistant_asd_max_mismatches:
+        Per-block cap ``m >= 0`` on relaxed mismatches. ``None`` leaves it
+        uncapped; ``0`` forces strict verification.
+    """
+    asd = _resolve_asd_params(assistant_asd_budget, assistant_asd_local_ratio, assistant_asd_max_mismatches)
+    any_asd_knob = any(
+        knob is not None
+        for knob in (assistant_asd_budget, assistant_asd_local_ratio, assistant_asd_max_mismatches)
+    )
+
+    if _effective_generation_value("do_sample", kwargs, generation_config, model) is True:
+        raise ValueError(
+            "ASD is a deterministic verifier for greedy decoding (`do_sample=False`). "
+            "For the sampling path, use native assisted decoding with `assistant_ensemble_weight` instead."
+        )
+    if _effective_generation_value("assistant_ensemble_weight", kwargs, generation_config, model) is not None:
+        raise ValueError(
+            "`assistant_ensemble_weight` and ASD are mutually exclusive verification relaxations. "
+            "Please use only one of them."
+        )
+    if (
+        any_asd_knob
+        and assistant_model is None
+        and _effective_generation_value("prompt_lookup_num_tokens", kwargs, generation_config, model) is None
+    ):
+        raise ValueError(
+            "ASD knobs require a draft source: pass `assistant_model` "
+            "(or set `prompt_lookup_num_tokens` for prompt-lookup drafting)."
+        )
+
+    # Install the ASD loop as the model class's assisted-decoding method for the
+    # duration of this call. `generate()` dispatches assisted generation through
+    # `getattr(type(self), "_assisted_decoding")`, so a class-level swap is the
+    # supported seam; it is reverted in `finally` below.
+    cls = type(model)
+    was_inherited = "_assisted_decoding" not in cls.__dict__
+    original = cls._assisted_decoding
+    cls._assisted_decoding = _asd_assisted_decoding
+    setattr(model, _ASD_CONFIG_ATTR, asd)
+    try:
+        return GenerationMixin.generate(
+            model,
+            inputs=inputs,
+            generation_config=generation_config,
+            logits_processor=logits_processor,
+            stopping_criteria=stopping_criteria,
+            prefix_allowed_tokens_fn=prefix_allowed_tokens_fn,
+            synced_gpus=synced_gpus,
+            assistant_model=assistant_model,
+            streamer=streamer,
+            negative_prompt_ids=negative_prompt_ids,
+            negative_prompt_attention_mask=negative_prompt_attention_mask,
+            **kwargs,
+        )
+    finally:
+        if was_inherited:
+            del cls._assisted_decoding
+        else:
+            cls._assisted_decoding = original
+        try:
+            delattr(model, _ASD_CONFIG_ATTR)
+        except AttributeError:
+            pass

custom_generate/requirements.txt

@@ -0,0 +1,4 @@
+# The acceptance loop mirrors `GenerationMixin._assisted_decoding` internals,
+# so the minor version is pinned. Verified on transformers 5.19.0.
+transformers>=5.19.0,<5.20.0
+torch>=2.0.0

scripts/example.py

@@ -0,0 +1,89 @@
+"""Sanity-check script for the ASD custom generation method.
+
+Run from anywhere:
+    python scripts/example.py \
+        --model Qwen/Qwen2.5-1.5B-Instruct \
+        --assistant-model Qwen/Qwen2.5-0.5B-Instruct \
+        --budget 2.0 --local-ratio 0.25 --max-mismatches 2
+
+Checks:
+1. ASD runs and produces text.
+2. `assistant_asd_budget=0` is token-identical to native strict assisted decoding
+   (the gold-standard correctness check).
+3. No ASD knob set is also token-identical to native (lossless default).
+"""
+
+import argparse
+import os
+
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model", required=True)
+    parser.add_argument("--assistant-model", required=True)
+    parser.add_argument("--budget", type=float, default=2.0)
+    parser.add_argument("--local-ratio", type=float, default=0.25)
+    parser.add_argument("--max-mismatches", type=int, default=2)
+    parser.add_argument("--max-new-tokens", type=int, default=64)
+    parser.add_argument("--prompt", default="Give me a short introduction to large language models.")
+    parser.add_argument("--repo-path", default=REPO_ROOT,
+                        help="Directory containing custom_generate/ (default: repo root)")
+    args = parser.parse_args()
+
+    tokenizer = AutoTokenizer.from_pretrained(args.model)
+    model = AutoModelForCausalLM.from_pretrained(args.model, dtype="auto", device_map="auto")
+    assistant_model = AutoModelForCausalLM.from_pretrained(
+        args.assistant_model, dtype="auto", device_map="auto"
+    )
+
+    inputs = tokenizer(args.prompt, return_tensors="pt").to(model.device)
+    gen_kwargs = dict(max_new_tokens=args.max_new_tokens, do_sample=False)
+
+    # 1. Native strict assisted decoding baseline.
+    strict_out = model.generate(**inputs, assistant_model=assistant_model, **gen_kwargs)
+    print("Strict      :", tokenizer.decode(strict_out[0], skip_special_tokens=True))
+
+    # 2. ASD relaxed.
+    asd_out = model.generate(
+        **inputs,
+        assistant_model=assistant_model,
+        custom_generate=args.repo_path,
+        trust_remote_code=True,
+        assistant_asd_budget=args.budget,
+        assistant_asd_local_ratio=args.local_ratio,
+        assistant_asd_max_mismatches=args.max_mismatches,
+        **gen_kwargs,
+    )
+    print("ASD         :", tokenizer.decode(asd_out[0], skip_special_tokens=True))
+
+    # 3. Gold check: budget=0 must be token-identical to native strict decoding.
+    asd_b0_out = model.generate(
+        **inputs,
+        assistant_model=assistant_model,
+        custom_generate=args.repo_path,
+        trust_remote_code=True,
+        assistant_asd_budget=0.0,
+        **gen_kwargs,
+    )
+    assert torch.equal(strict_out, asd_b0_out), "ASD(budget=0) differs from strict decoding!"
+    print("OK: assistant_asd_budget=0 is token-identical to strict verification.")
+
+    # 4. Lossless default: no ASD knobs must also be token-identical.
+    asd_default_out = model.generate(
+        **inputs,
+        assistant_model=assistant_model,
+        custom_generate=args.repo_path,
+        trust_remote_code=True,
+        **gen_kwargs,
+    )
+    assert torch.equal(strict_out, asd_default_out), "ASD (no knobs) differs from strict decoding!"
+    print("OK: no ASD knobs is token-identical to strict verification.")
+
+
+if __name__ == "__main__":
+    main()

scripts/validate.py

@@ -0,0 +1,180 @@
+"""Comprehensive GPU validation for the ASD custom generation method.
+
+Covers:
+1. Native strict baseline vs ASD(B=0)        -> must be token-identical (gold check)
+2. Native strict vs custom (no knobs)        -> must be token-identical (lossless default)
+3. ASD(B=2.0) relaxed                        -> runs, deterministic, >= strict acceptance
+4. Determinism                               -> two identical relaxed runs, token-identical
+5. Error handling                            -> do_sample / ensemble_weight / no draft source
+6. Throughput + acceptance comparison        -> strict vs relaxed
+"""
+
+import os
+import sys
+import time
+
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+REPO = os.environ.get("ASD_REPO", os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+TARGET = os.environ.get("ASD_TARGET", "Qwen/Qwen2.5-1.5B-Instruct")
+DRAFT = os.environ.get("ASD_DRAFT", "Qwen/Qwen2.5-0.5B-Instruct")
+MAX_NEW = 128
+PROMPTS = [
+    "Give me a short introduction to large language models.",
+    "What is the capital of France? Answer in one sentence.",
+    "Write a Python function that computes the nth Fibonacci number.",
+    "Explain the difference between TCP and UDP briefly.",
+    "List three prime numbers greater than 10 and justify one of them.",
+]
+
+failures = []
+
+
+def check(name, cond, detail=""):
+    status = "PASS" if cond else "FAIL"
+    print(f"[{status}] {name}" + (f" -- {detail}" if detail else ""))
+    if not cond:
+        failures.append(name)
+
+
+def count_fwd(tok, out, prompt_len):
+    return out.shape[1] - prompt_len
+
+
+def main():
+    tokenizer = AutoTokenizer.from_pretrained(TARGET)
+    model = AutoModelForCausalLM.from_pretrained(TARGET, dtype="auto", device_map="auto")
+    assistant = AutoModelForCausalLM.from_pretrained(DRAFT, dtype="auto", device_map="auto")
+    model.eval()
+    assistant.eval()
+
+    def run(prompt, **extra):
+        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+        t0 = time.perf_counter()
+        out = model.generate(
+            **inputs, max_new_tokens=MAX_NEW, do_sample=False, **extra
+        )
+        dt = time.perf_counter() - t0
+        n = out.shape[1] - inputs["input_ids"].shape[1]
+        return out, n / dt
+
+    # ---------- 1&2. gold checks ----------
+    for prompt in PROMPTS:
+        strict, _ = run(prompt, assistant_model=assistant)
+        b0, _ = run(
+            prompt,
+            assistant_model=assistant,
+            custom_generate=REPO,
+            trust_remote_code=True,
+            assistant_asd_budget=0.0,
+        )
+        noknobs, _ = run(
+            prompt,
+            assistant_model=assistant,
+            custom_generate=REPO,
+            trust_remote_code=True,
+        )
+        check(f"gold B=0 identical ({prompt[:25]!r})", torch.equal(strict, b0))
+        check(f"gold no-knobs identical ({prompt[:25]!r})", torch.equal(strict, noknobs))
+
+    # ---------- 3&4. relaxed run + determinism ----------
+    for prompt in PROMPTS:
+        relaxed1, tps1 = run(
+            prompt,
+            assistant_model=assistant,
+            custom_generate=REPO,
+            trust_remote_code=True,
+            assistant_asd_budget=2.0,
+            assistant_asd_local_ratio=0.25,
+            assistant_asd_max_mismatches=2,
+        )
+        relaxed2, _ = run(
+            prompt,
+            assistant_model=assistant,
+            custom_generate=REPO,
+            trust_remote_code=True,
+            assistant_asd_budget=2.0,
+            assistant_asd_local_ratio=0.25,
+            assistant_asd_max_mismatches=2,
+        )
+        check(f"determinism ({prompt[:25]!r})", torch.equal(relaxed1, relaxed2))
+        check(f"relaxed non-empty ({prompt[:25]!r})", relaxed1.shape[1] > 0)
+
+    # ---------- 5. error handling ----------
+    inputs = tokenizer(PROMPTS[0], return_tensors="pt").to(model.device)
+
+    def expect_raise(name, **kw):
+        try:
+            model.generate(**inputs, max_new_tokens=8, **kw)
+        except ValueError:
+            check(f"raises {name}", True)
+        except Exception as e:  # noqa: BLE001
+            check(f"raises {name}", False, f"wrong exception: {type(e).__name__}: {e}")
+        else:
+            check(f"raises {name}", False, "no exception")
+
+    expect_raise(
+        "do_sample=True",
+        assistant_model=assistant,
+        custom_generate=REPO,
+        trust_remote_code=True,
+        do_sample=True,
+        assistant_asd_budget=2.0,
+    )
+    expect_raise(
+        "ensemble_weight conflict",
+        assistant_model=assistant,
+        custom_generate=REPO,
+        trust_remote_code=True,
+        assistant_ensemble_weight=0.5,
+        assistant_asd_budget=2.0,
+    )
+    expect_raise(
+        "knobs without draft source",
+        custom_generate=REPO,
+        trust_remote_code=True,
+        assistant_asd_budget=2.0,
+    )
+    expect_raise(
+        "negative budget",
+        assistant_model=assistant,
+        custom_generate=REPO,
+        trust_remote_code=True,
+        assistant_asd_budget=-1.0,
+    )
+
+    # ---------- 6. throughput + acceptance ----------
+    print("\n=== throughput comparison ===")
+    tot_strict_tps = tot_relaxed_tps = 0.0
+    for prompt in PROMPTS:
+        strict, s_tps = run(prompt, assistant_model=assistant)
+        relaxed, r_tps = run(
+            prompt,
+            assistant_model=assistant,
+            custom_generate=REPO,
+            trust_remote_code=True,
+            assistant_asd_budget=2.0,
+            assistant_asd_local_ratio=0.25,
+            assistant_asd_max_mismatches=2,
+        )
+        tot_strict_tps += s_tps
+        tot_relaxed_tps += r_tps
+        overlap = sum(
+            1 for a, b in zip(strict[0].tolist(), relaxed[0].tolist()) if a == b
+        ) / max(len(strict[0]), len(relaxed[0]))
+        print(
+            f"  strict {s_tps:6.2f} tok/s | relaxed {r_tps:6.2f} tok/s "
+            f"({(r_tps / s_tps - 1) * 100:+5.1f}%) | prefix overlap {overlap:.2%}"
+        )
+    print(
+        f"  MEAN strict {tot_strict_tps / len(PROMPTS):.2f} tok/s | "
+        f"relaxed {tot_relaxed_tps / len(PROMPTS):.2f} tok/s"
+    )
+
+    print("\n" + ("ALL PASS" if not failures else f"FAILURES: {failures}"))
+    sys.exit(1 if failures else 0)
+
+
+if __name__ == "__main__":
+    main()