art3_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_V1 = os.environ.get("ADAPTIVE_V1") == "1"
ADAPTIVE_WARMUP_CALLS = int(os.environ.get("ADAPTIVE_WARMUP_CALLS", "400"))
ADAPTIVE_P_STAR = float(os.environ.get("ADAPTIVE_P_STAR", "0.27"))
ADAPTIVE_BINS = int(os.environ.get("ADAPTIVE_BINS", "10"))
ADAPTIVE_LOG_EVERY = int(os.environ.get("ADAPTIVE_LOG_EVERY", "2000"))


def _apply_adaptive_v1_patch(module: Any) -> None:
    """abay adaptive-K v1: acceptance-calibrated position-1 gate + remainder graph.

    v0 lesson (results/20260610-005822-259_abay.md): margin-only pooled-quantile
    thresholds stop accepted runs. v1 calibrates P(accept_pos1 | margin_1)
    EMPIRICALLY during warmup using the runner's num_rejected_tokens_gpu
    feedback, gates only at position 1 (single host sync), and replays the
    K-1 remainder draft loop as one CUDA graph (pupa loopgraph machinery,
    verbatim). Stop rule: gate when P(accept|margin) < P* ~ tau_d/(T/E[L]).
    Exactness: drafted tokens are argmax over the same sparse logits as
    get_top_tokens; gated steps pad with token 0 which verification rejects.
    """
    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_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 _calibrate(state: dict[str, Any], token_count: int) -> None:
        pairs = sorted(state["pairs"])
        n = len(pairs)
        bins = max(4, ADAPTIVE_BINS)
        tau = float("-inf")
        report = []
        for b in range(bins):
            lo, hi = b * n // bins, (b + 1) * n // bins
            if hi <= lo:
                continue
            chunk = pairs[lo:hi]
            rate = sum(a for _, a in chunk) / len(chunk)
            report.append((round(chunk[0][0], 2), round(rate, 2)))
            if rate < ADAPTIVE_P_STAR:
                tau = chunk[-1][0]
        state["tau"] = tau
        print(
            f"[abay-adaptive-v1] calibrated tau1={tau:.3f} from {n} "
            f"(margin,accept) pairs; P*={ADAPTIVE_P_STAR}; "
            f"bin (lower_edge, acc_rate): {report}; K_max={token_count} "
            f"(pid {os.getpid()})",
            file=sys.stderr,
            flush=True,
        )

    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_v1",
            {
                "calls": 0,
                "tau": None,
                "pairs": [],
                "pending": None,
                "graph": None,
                "graph_failed": False,
                "steps": 0,
                "gated": 0,
                "rolled": 0,
            },
        )
        if not _is_eligible(self, common_attn_metadata, sampling_metadata):
            state["pending"] = None
            return _call_base_propose(base_propose, self, kwargs)

        state["calls"] += 1
        calibrating = state["tau"] is None

        if (
            calibrating
            and state["pending"] is not None
            and num_rejected_tokens_gpu is not None
        ):
            prev_m1, prev_drafted = state["pending"]
            rejected = int(num_rejected_tokens_gpu.reshape(-1)[0].item())
            accepted1 = 1 if (prev_drafted - rejected) >= 1 else 0
            state["pairs"].append((prev_m1, accepted1))
        state["pending"] = 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

        state["steps"] += 1

        if not calibrating and first_margin < state["tau"]:
            state["gated"] += 1
            out = torch.zeros(
                (1, token_count), dtype=torch.int64, device=first_token.device
            )
            out[0, 0] = first_token[0]
            if ADAPTIVE_LOG_EVERY and state["steps"] % ADAPTIVE_LOG_EVERY == 0:
                print(
                    f"[abay-adaptive-v1] steps={state['steps']} "
                    f"gated={state['gated']} rolled={state['rolled']}",
                    file=sys.stderr,
                    flush=True,
                )
            return out

        state["rolled"] += 1

        graph_state = state.setdefault("graph_state", {"failed": False})
        if (
            not calibrating
            and state["graph"] is None
            and not state["graph_failed"]
        ):
            try:
                _build_static_buffers(self, graph_state, cad)
                _refresh_static_buffers(self, graph_state, cad)
                graph_state["out"][0, 0:1].copy_(first_token)
                self.hidden_states[:1].copy_(first_hidden)
                _capture_graph(self, graph_state)
                state["graph"] = graph_state["graph"]
                print(
                    f"[abay-adaptive-v1] captured K-1={token_count - 1} remainder "
                    f"graph at call {state['calls']} (pid {os.getpid()})",
                    file=sys.stderr,
                    flush=True,
                )
            except Exception as exc:
                state["graph_failed"] = True
                state["graph"] = None
                print(
                    f"[abay-adaptive-v1] graph capture failed, eager fallback: {exc!r}",
                    file=sys.stderr,
                    flush=True,
                )

        if state["graph"] is not None:
            _refresh_static_buffers(self, graph_state, cad)
            graph_state["out"][0, 0:1].copy_(first_token)
            self.hidden_states[:1].copy_(first_hidden)
            state["graph"].replay()
            if ADAPTIVE_LOG_EVERY and state["steps"] % ADAPTIVE_LOG_EVERY == 0:
                print(
                    f"[abay-adaptive-v1] steps={state['steps']} "
                    f"gated={state['gated']} rolled={state['rolled']}",
                    file=sys.stderr,
                    flush=True,
                )
            return graph_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
            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, _ = self._sample_draft_tokens(last_hidden[:1], sampling_metadata)
            draft_tokens.append(token)

        if calibrating:
            state["pending"] = (first_margin, token_count)
            if state["calls"] >= ADAPTIVE_WARMUP_CALLS and len(state["pairs"]) >= 64:
                _calibrate(state, token_count)

        return torch.stack(draft_tokens, dim=1)

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


if ADAPTIVE_V1:
    sys.meta_path.insert(0, _TargetFinder(LOOPGRAPH_TARGET, _apply_adaptive_v1_patch))
else:
    sys.meta_path.insert(0, _TargetFinder(LOOPGRAPH_TARGET, _apply_loopgraph_patch))