art2_src/sitecustomize.py

"""Patch vLLM Gemma4 MTP drafting with a CUDA graph loop replay.

This file is loaded by the vLLM child process through PYTHONPATH. It intentionally
does not patch PLE. Pupa's serve.py patches PLE textfast and scale-folds through
the installed vLLM source so the fold can be verified fail-closed at load time.
"""

from __future__ import annotations

import importlib.abc
import importlib.util
import os
import sys
from copy import copy
from typing import Any


LOOPGRAPH_TARGET = "vllm.v1.spec_decode.gemma4"
LOOPGRAPH_WARMUP_CALLS = int(os.environ.get("LOOPGRAPH_WARMUP_CALLS", "48"))
LOOPGRAPH_REQUIRE_CAPTURE = os.environ.get("LOOPGRAPH_REQUIRE_CAPTURE") == "1"


def _call_base_propose(base_propose: Any, self: Any, kwargs: dict[str, Any]) -> Any:
    return base_propose(self, **kwargs)


def _build_static_buffers(self: Any, state: dict[str, Any], cad: Any) -> None:
    import torch

    device = self.device
    token_count = self.num_speculative_tokens
    state["out"] = torch.zeros((1, token_count), dtype=torch.int64, device=device)
    state["seq_lens"] = torch.zeros_like(cad.seq_lens[:1])
    state["block_tables"] = {}

    static_cad = copy(cad)
    static_cad.seq_lens = state["seq_lens"]
    static_cad.num_actual_tokens = 1
    static_cad.max_query_len = 1
    static_cad.max_seq_len = self.max_model_len
    static_cad.slot_mapping = self._slot_mapping_buffer[:1]
    static_cad.query_start_loc = self.arange[:2]

    per_layer_metadata = {}
    for group in self.draft_attn_groups:
        group_id = group.kv_cache_group_id
        source = self._per_group_block_tables.get(group_id, cad.block_table_tensor)[:1]
        block_size = group.get_metadata_builder().kv_cache_spec.block_size
        width = max(source.shape[1], -(-self.max_model_len // block_size))
        static_block_table = torch.zeros(
            (1, width), dtype=source.dtype, device=device
        )
        state["block_tables"][group_id] = static_block_table

        group_cad = copy(static_cad)
        group_cad.block_table_tensor = static_block_table
        metadata = group.get_metadata_builder().build_for_drafting(
            common_attn_metadata=group_cad,
            draft_index=1,
        )
        for layer_name in group.layer_names:
            per_layer_metadata[layer_name] = metadata
    state["metadata"] = per_layer_metadata


def _refresh_static_buffers(self: Any, state: dict[str, Any], cad: Any) -> None:
    state["seq_lens"].copy_(cad.seq_lens[:1])
    for group_id, static_block_table in state["block_tables"].items():
        source = self._per_group_block_tables.get(group_id, cad.block_table_tensor)[:1]
        width = min(source.shape[1], static_block_table.shape[1])
        static_block_table[:, :width].copy_(source[:, :width])


def _run_graph_body(self: Any, state: dict[str, Any]) -> None:
    from vllm.config import CUDAGraphMode
    from vllm.forward_context import set_forward_context

    token_count = self.num_speculative_tokens
    output = state["out"]
    with set_forward_context(
        state["metadata"],
        self.vllm_config,
        num_tokens=1,
        num_tokens_across_dp=None,
        cudagraph_runtime_mode=CUDAGraphMode.NONE,
        slot_mapping=self._get_slot_mapping(1),
    ):
        for index in range(token_count - 1):
            self.input_ids[:1].copy_(output[0, index : index + 1])
            last_hidden, backbone_hidden = self.model(
                input_ids=self.input_ids[:1],
                positions=self._get_positions(1),
                inputs_embeds=None,
                hidden_states=self.hidden_states[:1],
            )
            self.hidden_states[:1].copy_(backbone_hidden[:1])
            token = self.model.get_top_tokens(last_hidden[:1])
            output[0, index + 1 : index + 2].copy_(token)


def _capture_graph(self: Any, state: dict[str, Any]) -> None:
    import torch

    for _ in range(2):
        _run_graph_body(self, state)
    torch.cuda.synchronize()
    graph = torch.cuda.CUDAGraph()
    with torch.cuda.graph(graph):
        _run_graph_body(self, state)
    state["graph"] = graph


def _is_loopgraph_eligible(self: Any, state: dict[str, Any], cad: Any) -> bool:
    return (
        not state["failed"]
        and self.num_speculative_tokens > 1
        and not self.parallel_drafting
        and not self._enable_probabilistic_draft_probs
        and not self.supports_mm_inputs
        and not self.uses_mrope
        and self.constant_draft_positions
        and cad.batch_size() == 1
    )


def _raise_or_fallback(exc: Exception) -> None:
    if LOOPGRAPH_REQUIRE_CAPTURE:
        raise RuntimeError("LOOPGRAPH_REQUIRE_CAPTURE=1 but capture failed") from exc


def _apply_loopgraph_patch(module: Any) -> None:
    import torch

    from vllm.forward_context import set_forward_context

    proposer_cls = module.Gemma4Proposer
    base_propose = proposer_cls.propose

    def propose(
        self: Any,
        target_token_ids: Any,
        target_positions: Any,
        target_hidden_states: Any,
        next_token_ids: Any,
        token_indices_to_sample: Any,
        common_attn_metadata: Any,
        sampling_metadata: Any,
        mm_embed_inputs: Any = None,
        num_rejected_tokens_gpu: Any = None,
        slot_mappings: Any = None,
    ) -> Any:
        kwargs = {
            "target_token_ids": target_token_ids,
            "target_positions": target_positions,
            "target_hidden_states": target_hidden_states,
            "next_token_ids": next_token_ids,
            "token_indices_to_sample": token_indices_to_sample,
            "common_attn_metadata": common_attn_metadata,
            "sampling_metadata": sampling_metadata,
            "mm_embed_inputs": mm_embed_inputs,
            "num_rejected_tokens_gpu": num_rejected_tokens_gpu,
            "slot_mappings": slot_mappings,
        }
        state = self.__dict__.setdefault(
            "_pupa_loopgraph",
            {"calls": 0, "graph": None, "failed": False},
        )
        if not _is_loopgraph_eligible(self, state, common_attn_metadata):
            return _call_base_propose(base_propose, self, kwargs)

        state["calls"] += 1
        if state["graph"] is None and state["calls"] <= LOOPGRAPH_WARMUP_CALLS:
            return _call_base_propose(base_propose, self, kwargs)

        self._last_draft_probs = None
        token_count = self.num_speculative_tokens
        num_tokens, token_indices_to_sample, cad = self.set_inputs_first_pass(
            target_token_ids=target_token_ids,
            next_token_ids=next_token_ids,
            target_positions=target_positions,
            target_hidden_states=target_hidden_states,
            token_indices_to_sample=token_indices_to_sample,
            cad=common_attn_metadata,
            num_rejected_tokens_gpu=num_rejected_tokens_gpu,
        )
        _, per_layer_metadata = self.build_per_group_and_layer_attn_metadata(cad)
        cg_mode, num_input_tokens, num_tokens_across_dp = (
            self._determine_batch_execution_and_padding(num_tokens)
        )
        model_kwargs, slot_map_size = self.build_model_inputs_first_pass(
            num_tokens,
            num_input_tokens,
            mm_embed_inputs,
        )
        with set_forward_context(
            per_layer_metadata,
            self.vllm_config,
            num_tokens=num_input_tokens,
            num_tokens_across_dp=num_tokens_across_dp,
            cudagraph_runtime_mode=cg_mode,
            slot_mapping=self._get_slot_mapping(slot_map_size, cad.slot_mapping),
        ):
            last_hidden, hidden = self.model(**model_kwargs)

        sample_hidden = last_hidden[token_indices_to_sample]
        positions = self.positions[token_indices_to_sample]
        first_hidden = hidden[token_indices_to_sample]
        self.positions[:1] = positions
        first_token, _ = self._sample_draft_tokens(sample_hidden, sampling_metadata)

        cad.num_actual_tokens = 1
        cad.max_query_len = 1
        cad.query_start_loc = self.arange[:2]
        cad.query_start_loc_cpu = torch.from_numpy(self.token_arange_np[:2]).clone()
        if num_rejected_tokens_gpu is not None:
            cad.seq_lens -= num_rejected_tokens_gpu
            cad._seq_lens_cpu = None
            cad._num_computed_tokens_cpu = None

        if state["graph"] is None and not state["failed"]:
            try:
                _build_static_buffers(self, state, cad)
                _refresh_static_buffers(self, state, cad)
                state["out"][0, 0:1].copy_(first_token)
                self.hidden_states[:1].copy_(first_hidden)
                _capture_graph(self, state)
                print(
                    f"[pupa-loopgraph] captured K-1={token_count - 1} graph "
                    f"at eligible call {state['calls']} (pid {os.getpid()})",
                    file=sys.stderr,
                    flush=True,
                )
            except Exception as exc:
                state["failed"] = True
                state["graph"] = None
                print(
                    f"[pupa-loopgraph] capture failed: {exc!r}",
                    file=sys.stderr,
                    flush=True,
                )
                _raise_or_fallback(exc)

        if state["graph"] is not None:
            _refresh_static_buffers(self, state, cad)
            state["out"][0, 0:1].copy_(first_token)
            self.hidden_states[:1].copy_(first_hidden)
            state["graph"].replay()
            return state["out"].clone()

        cg_mode, input_batch_size, batch_size_dp = (
            self._determine_batch_execution_and_padding(1)
        )
        draft_tokens = [first_token]
        hidden_current = first_hidden
        loop_metadata = None
        for index in range(token_count - 1):
            input_ids = draft_tokens[-1].int()
            if index == 0:
                _, loop_metadata = self.build_per_group_and_layer_attn_metadata(
                    cad,
                    draft_index=1,
                )
            self.input_ids[:1] = input_ids
            self.hidden_states[:1] = hidden_current
            kwargs = {
                "input_ids": self.input_ids[:input_batch_size],
                "positions": self._get_positions(input_batch_size),
                "inputs_embeds": None,
                "hidden_states": self.hidden_states[:input_batch_size],
            }
            with set_forward_context(
                loop_metadata,
                self.vllm_config,
                num_tokens=input_batch_size,
                num_tokens_across_dp=batch_size_dp,
                cudagraph_runtime_mode=cg_mode,
                slot_mapping=self._get_slot_mapping(input_batch_size),
            ):
                last_hidden, hidden = self.model(**kwargs)
            hidden_current = hidden[:1]
            token, _ = self._sample_draft_tokens(last_hidden[:1], sampling_metadata)
            draft_tokens.append(token)
        return torch.stack(draft_tokens, dim=1)

    proposer_cls.propose = propose
    print(
        f"[pupa-loopgraph] patched Gemma4Proposer.propose in pid {os.getpid()} "
        f"(warmup_calls={LOOPGRAPH_WARMUP_CALLS}, "
        f"require_capture={LOOPGRAPH_REQUIRE_CAPTURE})",
        file=sys.stderr,
        flush=True,
    )


class _PatchingLoader(importlib.abc.Loader):
    def __init__(self, inner: importlib.abc.Loader, patch_fn: Any) -> None:
        self._inner = inner
        self._patch_fn = patch_fn

    def create_module(self, spec: Any) -> Any:
        return self._inner.create_module(spec)

    def exec_module(self, module: Any) -> None:
        self._inner.exec_module(module)
        self._patch_fn(module)


class _TargetFinder(importlib.abc.MetaPathFinder):
    def __init__(self, target: str, patch_fn: Any) -> None:
        self._target = target
        self._patch_fn = patch_fn
        self._busy = False

    def find_spec(self, fullname: str, path: Any = None, target: Any = None) -> Any:
        if fullname != self._target or self._busy:
            return None
        self._busy = True
        try:
            spec = importlib.util.find_spec(fullname)
        finally:
            self._busy = False
        if spec is None or spec.loader is None:
            return None
        spec.loader = _PatchingLoader(spec.loader, self._patch_fn)
        return spec


ADAPTIVE_SPEC = os.environ.get("ADAPTIVE_SPEC") == "1"
ADAPTIVE_WARMUP_CALLS = int(os.environ.get("ADAPTIVE_WARMUP_CALLS", "96"))
ADAPTIVE_STOP_QUANTILE = float(os.environ.get("ADAPTIVE_STOP_QUANTILE", "0.30"))
ADAPTIVE_LOG_EVERY = int(os.environ.get("ADAPTIVE_LOG_EVERY", "2000"))


def _apply_adaptive_patch(module: Any) -> None:
    """abay: margin-gated adaptive speculation depth.

    Derivation (board msg 20260610-002612): conditional acceptance RISES with
    depth (0.68 -> 0.81), so steps are bimodal easy/hard. Fixed K both
    truncates easy runs and burns full drafts on hard steps. Policy: roll the
    AR drafter up to K_max=num_speculative_tokens, stop as soon as the
    drafter's sparse-logit top1-top2 margin drops below a threshold
    calibrated from its own warmup margin distribution. Token selection is
    argmax over the SAME sparse logits as get_top_tokens, so proposals are
    unchanged where drafted; early-stopped positions are padded with token 0,
    which exact-match verification simply rejects. Output tokens cannot
    change by construction; only drafting cost does.
    """
    import torch

    from vllm.forward_context import set_forward_context

    proposer_cls = module.Gemma4Proposer
    base_propose = proposer_cls.propose

    def _margin_sample(self: Any, hidden: Any) -> tuple[Any, float]:
        me = self.model.masked_embedding
        weight = self.model._get_full_lm_head_weight()
        logits, indices = me._select_and_score(hidden, weight)
        top2 = logits.topk(2, dim=-1)
        token = indices.gather(-1, top2.indices[:, :1]).squeeze(-1)
        margin = float(top2.values[0, 0] - top2.values[0, 1])
        return token, margin

    def _is_adaptive_eligible(self: Any, cad: Any, sampling_metadata: Any) -> bool:
        return (
            self.num_speculative_tokens > 1
            and not self.parallel_drafting
            and not self._enable_probabilistic_draft_probs
            and not self.supports_mm_inputs
            and not self.uses_mrope
            and self.constant_draft_positions
            and sampling_metadata.all_greedy
            and getattr(self.model, "masked_embedding", None) is not None
            and cad.batch_size() == 1
        )

    def propose(
        self: Any,
        target_token_ids: Any,
        target_positions: Any,
        target_hidden_states: Any,
        next_token_ids: Any,
        token_indices_to_sample: Any,
        common_attn_metadata: Any,
        sampling_metadata: Any,
        mm_embed_inputs: Any = None,
        num_rejected_tokens_gpu: Any = None,
        slot_mappings: Any = None,
    ) -> Any:
        kwargs = {
            "target_token_ids": target_token_ids,
            "target_positions": target_positions,
            "target_hidden_states": target_hidden_states,
            "next_token_ids": next_token_ids,
            "token_indices_to_sample": token_indices_to_sample,
            "common_attn_metadata": common_attn_metadata,
            "sampling_metadata": sampling_metadata,
            "mm_embed_inputs": mm_embed_inputs,
            "num_rejected_tokens_gpu": num_rejected_tokens_gpu,
            "slot_mappings": slot_mappings,
        }
        state = self.__dict__.setdefault(
            "_abay_adaptive",
            {
                "calls": 0,
                "tau": None,
                "pool": [],
                "steps": 0,
                "forwards": 0,
                "stops": [0] * 32,
            },
        )
        if not _is_adaptive_eligible(self, common_attn_metadata, sampling_metadata):
            return _call_base_propose(base_propose, self, kwargs)

        state["calls"] += 1
        calibrated = state["tau"] is not None

        self._last_draft_probs = None
        token_count = self.num_speculative_tokens
        num_tokens, token_indices_to_sample, cad = self.set_inputs_first_pass(
            target_token_ids=target_token_ids,
            next_token_ids=next_token_ids,
            target_positions=target_positions,
            target_hidden_states=target_hidden_states,
            token_indices_to_sample=token_indices_to_sample,
            cad=common_attn_metadata,
            num_rejected_tokens_gpu=num_rejected_tokens_gpu,
        )
        _, per_layer_metadata = self.build_per_group_and_layer_attn_metadata(cad)
        cg_mode, num_input_tokens, num_tokens_across_dp = (
            self._determine_batch_execution_and_padding(num_tokens)
        )
        model_kwargs, slot_map_size = self.build_model_inputs_first_pass(
            num_tokens,
            num_input_tokens,
            mm_embed_inputs,
        )
        with set_forward_context(
            per_layer_metadata,
            self.vllm_config,
            num_tokens=num_input_tokens,
            num_tokens_across_dp=num_tokens_across_dp,
            cudagraph_runtime_mode=cg_mode,
            slot_mapping=self._get_slot_mapping(slot_map_size, cad.slot_mapping),
        ):
            last_hidden, hidden = self.model(**model_kwargs)

        sample_hidden = last_hidden[token_indices_to_sample]
        positions = self.positions[token_indices_to_sample]
        first_hidden = hidden[token_indices_to_sample]
        self.positions[:1] = positions
        first_token, first_margin = _margin_sample(self, sample_hidden)

        cad.num_actual_tokens = 1
        cad.max_query_len = 1
        cad.query_start_loc = self.arange[:2]
        cad.query_start_loc_cpu = torch.from_numpy(self.token_arange_np[:2]).clone()
        if num_rejected_tokens_gpu is not None:
            cad.seq_lens -= num_rejected_tokens_gpu
            cad._seq_lens_cpu = None
            cad._num_computed_tokens_cpu = None

        margins = [first_margin]
        draft_tokens = [first_token]
        state["steps"] += 1
        state["forwards"] += 1

        stop_now = calibrated and first_margin < state["tau"]
        if not stop_now:
            cg_mode, input_batch_size, batch_size_dp = (
                self._determine_batch_execution_and_padding(1)
            )
            hidden_current = first_hidden
            loop_metadata = None
            for index in range(token_count - 1):
                input_ids = draft_tokens[-1].int()
                if index == 0:
                    _, loop_metadata = self.build_per_group_and_layer_attn_metadata(
                        cad,
                        draft_index=1,
                    )
                self.input_ids[:1] = input_ids
                self.hidden_states[:1] = hidden_current
                fwd_kwargs = {
                    "input_ids": self.input_ids[:input_batch_size],
                    "positions": self._get_positions(input_batch_size),
                    "inputs_embeds": None,
                    "hidden_states": self.hidden_states[:input_batch_size],
                }
                with set_forward_context(
                    loop_metadata,
                    self.vllm_config,
                    num_tokens=input_batch_size,
                    num_tokens_across_dp=batch_size_dp,
                    cudagraph_runtime_mode=cg_mode,
                    slot_mapping=self._get_slot_mapping(input_batch_size),
                ):
                    last_hidden, hidden = self.model(**fwd_kwargs)
                hidden_current = hidden[:1]
                token, margin = _margin_sample(self, last_hidden[:1])
                draft_tokens.append(token)
                margins.append(margin)
                state["forwards"] += 1
                if calibrated and margin < state["tau"]:
                    break

        n_drafted = len(draft_tokens)
        state["stops"][min(n_drafted, 31)] += 1

        if not calibrated:
            state["pool"].extend(margins)
            if state["calls"] >= ADAPTIVE_WARMUP_CALLS:
                pool = sorted(state["pool"])
                idx = max(0, min(len(pool) - 1, int(ADAPTIVE_STOP_QUANTILE * len(pool))))
                state["tau"] = pool[idx]
                qs = {
                    q: pool[int(q * (len(pool) - 1))]
                    for q in (0.1, 0.3, 0.5, 0.7, 0.9)
                }
                print(
                    f"[abay-adaptive] calibrated tau={state['tau']:.3f} "
                    f"(quantile {ADAPTIVE_STOP_QUANTILE}, pool={len(pool)}, "
                    f"quantiles={ {k: round(v, 2) for k, v in qs.items()} }, "
                    f"K_max={token_count}, pid {os.getpid()})",
                    file=sys.stderr,
                    flush=True,
                )

        if ADAPTIVE_LOG_EVERY and state["steps"] % ADAPTIVE_LOG_EVERY == 0:
            mean_fwd = state["forwards"] / state["steps"]
            print(
                f"[abay-adaptive] steps={state['steps']} mean_drafts={mean_fwd:.2f} "
                f"stop_hist={state['stops'][1:token_count + 1]}",
                file=sys.stderr,
                flush=True,
            )

        stacked = torch.stack(draft_tokens, dim=1)
        if n_drafted < token_count:
            out = torch.zeros(
                (1, token_count), dtype=stacked.dtype, device=stacked.device
            )
            out[:, :n_drafted] = stacked
            return out
        return stacked

    proposer_cls.propose = propose
    print(
        f"[abay-adaptive] patched Gemma4Proposer.propose in pid {os.getpid()} "
        f"(warmup_calls={ADAPTIVE_WARMUP_CALLS}, "
        f"stop_quantile={ADAPTIVE_STOP_QUANTILE})",
        file=sys.stderr,
        flush=True,
    )


if ADAPTIVE_SPEC:
    sys.meta_path.insert(0, _TargetFinder(LOOPGRAPH_TARGET, _apply_adaptive_patch))
else:
    sys.meta_path.insert(0, _TargetFinder(LOOPGRAPH_TARGET, _apply_loopgraph_patch))