art5_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


FUSED_DRAFTER = os.environ.get("FUSED_DRAFTER") == "1"
FUSED_DEBUG_PROBES = int(os.environ.get("FUSED_DEBUG_PROBES", "0"))
FUSED_WARMUP_CALLS = int(os.environ.get("FUSED_WARMUP_CALLS", "48"))
FUSED_MIN_MATCH = float(os.environ.get("FUSED_MIN_MATCH", "0.95"))
FUSED_LOG_EVERY = int(os.environ.get("FUSED_LOG_EVERY", "4000"))


def _apply_fused_patch(module: Any) -> None:
    """abay: fused Triton drafter loop (see fused_drafter.py).

    Warmup calls run the stock eager draft loop and shadow-run the fused
    forward on identical per-iteration inputs, comparing sampled tokens.
    If agreement >= FUSED_MIN_MATCH, the K-1 loop is captured as ONE CUDA
    graph over the fused kernels; any failure falls back to the loopgraph
    propose patched underneath. Output exactness is guaranteed by greedy
    verification regardless — fused drift only affects acceptance.
    """
    import torch

    from vllm.forward_context import set_forward_context

    proposer_cls = module.Gemma4Proposer
    fallback_propose = proposer_cls.propose  # loopgraph patch, applied first

    def _log(msg: str) -> None:
        print(f"[abay-fused] {msg} (pid {os.getpid()})", file=sys.stderr, flush=True)

    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 _init_fused(self: Any, state: dict[str, Any]) -> None:
        import fused_drafter

        state["fd"] = fused_drafter.FusedDrafter(self, self.device)
        state["token_buf"] = torch.zeros(1, dtype=torch.int64, device=self.device)
        state["hid_buf"] = torch.zeros(
            1,
            state["fd"].bh,
            dtype=self.hidden_states.dtype,
            device=self.device,
        )
        _log(
            "FusedDrafter initialized: "
            + ", ".join(
                f"L{i}(hd={ly['hd']},rot={ly['rot']},win={ly['window']},"
                f"gid={ly['gid']},bs={ly['kv'].shape[2]})"
                for i, ly in enumerate(state["fd"].layers)
            )
        )
        for i, ly in enumerate(state["fd"].layers):
            _log(f"L{i} cs.shape={tuple(ly['cs'].shape)} rot={ly['rot']}")
        if FUSED_DEBUG_PROBES:
            predictor = self.model.model
            stock = state["stock_probes"] = {}

            def _mk(name, pick):
                def hook(mod, args, out):
                    stock[name] = pick(out).detach().view(-1).float().clone()
                return hook

            predictor.pre_projection.register_forward_hook(
                _mk("pre", lambda o: o[0])
            )
            for i, lyr in enumerate(predictor.layers):
                lyr.register_forward_hook(_mk(f"L{i}", lambda o: o[0]))
                lyr.self_attn.register_forward_hook(
                    _mk(f"L{i}.attnout", lambda o: o)
                )
            predictor.norm.register_forward_hook(_mk("dh", lambda o: o))
            predictor.post_projection.register_forward_hook(
                _mk("backbone", lambda o: o[0])
            )
            _log("debug probes installed")

    def _build_fused_static(self: Any, state: dict[str, Any], cad: Any) -> None:
        fd = state["fd"]
        seq = torch.zeros_like(cad.seq_lens[:1])
        state["seq_lens"] = seq
        state["seqlen_bufs"] = {}
        state["bt_bufs"] = {}
        for group in self.draft_attn_groups:
            gid = group.kv_cache_group_id
            source = self._per_group_block_tables.get(gid, 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))
            state["bt_bufs"][gid] = torch.zeros(
                width, dtype=source.dtype, device=fd.device
            )
            state["seqlen_bufs"][gid] = seq

    def _refresh_fused_static(self: Any, state: dict[str, Any], cad: Any) -> None:
        state["seq_lens"].copy_(cad.seq_lens[:1])
        for gid, buf in state["bt_bufs"].items():
            source = self._per_group_block_tables.get(gid, cad.block_table_tensor)[:1]
            width = min(source.shape[1], buf.shape[0])
            buf[:width].copy_(source[0, :width])

    def _fused_loop_body(self: Any, state: dict[str, Any], out: Any) -> None:
        fd = state["fd"]
        token_count = self.num_speculative_tokens
        for index in range(token_count - 1):
            dh = fd.forward(
                state["token_buf"],
                self.positions,
                state["hid_buf"],
                state["seqlen_bufs"],
                state["bt_bufs"],
                state["hid_buf"],
            )
            tok = self.model.get_top_tokens(dh)
            out[0, index + 1 : index + 2].copy_(tok.view(1, 1))
            state["token_buf"].copy_(tok.view(-1))

    def _capture_fused_graph(self: Any, state: dict[str, Any]) -> None:
        token_count = self.num_speculative_tokens
        state["out"] = torch.zeros(
            (1, token_count), dtype=torch.int64, device=self.device
        )
        for _ in range(2):
            _fused_loop_body(self, state, state["out"])
        torch.cuda.synchronize()
        graph = torch.cuda.CUDAGraph()
        with torch.cuda.graph(graph):
            _fused_loop_body(self, state, state["out"])
        state["graph"] = graph

    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_fused",
            {
                "calls": 0,
                "fd": None,
                "failed": False,
                "decided": False,
                "match": 0,
                "total": 0,
                "max_err": 0.0,
                "graph": None,
                "steps": 0,
            },
        )
        if state["failed"] or not _is_eligible(
            self, common_attn_metadata, sampling_metadata
        ):
            return _call_base_propose(fallback_propose, self, kwargs)

        state["calls"] += 1
        if state["fd"] is None:
            try:
                _init_fused(self, state)
            except Exception as exc:
                state["failed"] = True
                _log(f"init failed, stock fallback: {exc!r}")
                return _call_base_propose(fallback_propose, self, kwargs)

        # ---- shared first pass (stock) ----
        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 "seq_lens" not in state:
            _build_fused_static(self, state, cad)
        _refresh_fused_static(self, state, cad)

        # ---- captured fast path ----
        if state["decided"] and state["graph"] is not None:
            state["out"][0, 0:1].copy_(first_token)
            state["token_buf"].copy_(first_token.view(-1))
            state["hid_buf"].copy_(first_hidden.view(1, -1))
            state["graph"].replay()
            state["steps"] += 1
            if FUSED_LOG_EVERY and state["steps"] % FUSED_LOG_EVERY == 0:
                _log(f"fused steps={state['steps']}")
            return state["out"].clone()

        # ---- warmup: stock eager loop + per-iteration fused shadow ----
        fd = state["fd"]
        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
        try:
            for index in range(token_count - 1):
                # fused shadow on this iteration's inputs
                state["token_buf"].copy_(draft_tokens[-1].view(-1))
                state["hid_buf"].copy_(hidden_current.view(1, -1))
                fused_probes = (
                    {}
                    if FUSED_DEBUG_PROBES and state["calls"] <= 6
                    else None
                )
                dh = fd.forward(
                    state["token_buf"],
                    self.positions,
                    state["hid_buf"],
                    state["seqlen_bufs"],
                    state["bt_bufs"],
                    state["hid_buf"],
                    probes=fused_probes,
                )
                tok_fused = self.model.get_top_tokens(dh)

                # stock iteration
                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
                )

                state["total"] += 1
                if int(tok_fused.view(-1)[0].item()) == int(token.view(-1)[0].item()):
                    state["match"] += 1
                if fused_probes is not None and "stock_probes" in state:
                    stock = state["stock_probes"]
                    parts = []
                    for key, mine in fused_probes.items():
                        ref = stock.get(key)
                        if ref is None:
                            continue
                        err = (mine.view(-1) - ref).abs().max()
                        scale = ref.abs().max().clamp_min(1e-6)
                        parts.append(f"{key}={float(err / scale):.4f}")
                    _log(
                        f"probe call={state['calls']} iter={index}: "
                        + " ".join(parts)
                    )
                ref = hidden_current.view(-1).float()
                err = (state["hid_buf"].view(-1).float() - ref).abs().max()
                scale = ref.abs().max().clamp_min(1e-6)
                state["max_err"] = max(state["max_err"], float(err / scale))
                draft_tokens.append(token)

            if not state["decided"] and state["calls"] >= FUSED_WARMUP_CALLS:
                rate = state["match"] / max(1, state["total"])
                _log(
                    f"shadow done: token match {state['match']}/{state['total']}"
                    f" = {rate:.3f} (min {FUSED_MIN_MATCH}), "
                    f"max backbone relerr {state['max_err']:.4f}"
                )
                if rate >= FUSED_MIN_MATCH:
                    state["token_buf"].copy_(first_token.view(-1))
                    state["hid_buf"].copy_(first_hidden.view(1, -1))
                    _capture_fused_graph(self, state)
                    _log(f"captured fused K-1={token_count - 1} graph")
                else:
                    state["failed"] = True
                    _log("match below threshold; stock fallback from here")
                state["decided"] = True
        except Exception as exc:
            state["failed"] = True
            state["decided"] = True
            _log(f"shadow/capture failed, stock fallback: {exc!r}")

        return torch.stack(draft_tokens, dim=1)

    proposer_cls.propose = propose
    _log(
        f"patched Gemma4Proposer.propose (warmup={FUSED_WARMUP_CALLS}, "
        f"min_match={FUSED_MIN_MATCH})"
    )


def _apply_loopgraph_then_fused(module: Any) -> None:
    _apply_loopgraph_patch(module)
    _apply_fused_patch(module)


if FUSED_DRAFTER:
    sys.meta_path.insert(
        0, _TargetFinder(LOOPGRAPH_TARGET, _apply_loopgraph_then_fused)
    )
else:
    sys.meta_path.insert(0, _TargetFinder(LOOPGRAPH_TARGET, _apply_loopgraph_patch))