submissions/byteshark/tree-v2-align-debug-v0/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"
RUNNER_TARGET = "vllm.v1.worker.gpu_model_runner"
TOP_TOKEN_TARGET = "vllm.model_executor.models.gemma4_mtp"
PROPOSER_TARGET = "vllm.v1.spec_decode.llm_base_proposer"
REJECTION_TARGET = "vllm.v1.sample.rejection_sampler"
LOOPGRAPH_WARMUP_CALLS = int(os.environ.get("LOOPGRAPH_WARMUP_CALLS", "48"))
LOOPGRAPH_REQUIRE_CAPTURE = os.environ.get("LOOPGRAPH_REQUIRE_CAPTURE") == "1"
LOOPGRAPH_PINGPONG_SLOTS = max(1, int(os.environ.get("LOOPGRAPH_PINGPONG_SLOTS", "1")))
LOOPGRAPH_CLEAR_STALE_SLOT_EVENTS = (
    os.environ.get("LOOPGRAPH_CLEAR_STALE_SLOT_EVENTS", "1") == "1"
)
FUSED_GREEDY_REJECTION_PREP = os.environ.get("FUSED_GREEDY_REJECTION_PREP", "1") == "1"
FUSED_SPEC_DECODE_METADATA = os.environ.get("FUSED_SPEC_DECODE_METADATA", "1") == "1"
FUSED_SPEC_DECODE_METADATA_REQUIRE = (
    os.environ.get("FUSED_SPEC_DECODE_METADATA_REQUIRE") == "1"
)
FUSED_DIRECT_GREEDY_REJECTION = (
    os.environ.get("FUSED_DIRECT_GREEDY_REJECTION", "1") == "1"
)
FUSED_DIRECT_GREEDY_REJECTION_REQUIRE = (
    os.environ.get("FUSED_DIRECT_GREEDY_REJECTION_REQUIRE") == "1"
)
SPEC_ASSUME_DETERMINISTIC_REJECTION = (
    os.environ.get("SPEC_ASSUME_DETERMINISTIC_REJECTION", "0") == "1"
)
SPEC_ASSUME_ARGMAX_PROCESSORS_OK = (
    os.environ.get("SPEC_ASSUME_ARGMAX_PROCESSORS_OK", "0") == "1"
)
# onegraph (@blake-fable5-1): the Gemma4 MTP drafter is Q-only and KV-shared —
# it never writes KV and has no cross-position dependencies, so the padded
# width-(K+1) first pass (and the full-prompt-width drafter pass on the first
# decode after prefill) only ever contributes the single position selected by
# token_indices_to_sample. Width-1 is exact. With ONEGRAPH=1 the whole
# propose() becomes one CUDA-graph replay of K width-1 iterations: iteration 0
# consumes next_token_ids + gathered target hidden/position, iterations 1..K-1
# are the stock loopgraph body. Drafter-only => cannot change emitted tokens.
ONEGRAPH = os.environ.get("ONEGRAPH", "1") == "1"
FUSED_SPARSE_ARGMAX = os.environ.get("FUSED_SPARSE_ARGMAX", "1") == "1"
FUSED_SPARSE_ARGMAX_REQUIRE = os.environ.get("FUSED_SPARSE_ARGMAX_REQUIRE") == "1"
FUSED_SPARSE_ARGMAX_BLOCK = int(os.environ.get("FUSED_SPARSE_ARGMAX_BLOCK", "16"))
_FUSED_SPARSE_ARGMAX_KERNELS: Any | None = None
_FUSED_GREEDY_REJECTION_PREP_KERNEL: Any | None = None
_FUSED_SPEC_DECODE_METADATA_KERNEL: Any | None = None
_LOOPGRAPH_SLOT_EVENTS_BY_PTR: dict[int, Any] = {}
_LOOPGRAPH_SLOT_EVENT_RECORDED_BY_PTR: dict[int, bool] = {}
_PUPA_GREEDY_PREP_CACHE: dict[int, tuple[Any, Any]] = {}
SPEC_ACCEPT_HISTOGRAM = os.environ.get("SPEC_ACCEPT_HISTOGRAM", "0") == "1"
SPEC_TREE_SPEC = os.environ.get("SPEC_TREE_SPEC", "1") == "1"
SPEC_TREE_WIDTH = int(os.environ.get("SPEC_TREE_WIDTH", "8"))
SPEC_TOPK_SALVAGE_PROBE = os.environ.get("SPEC_TOPK_SALVAGE_PROBE", "0") == "1"
SPEC_TOPK_SALVAGE_K = max(1, int(os.environ.get("SPEC_TOPK_SALVAGE_K", "8")))
SPEC_TOPK_SALVAGE_MAX_STEPS = int(os.environ.get("SPEC_TOPK_SALVAGE_MAX_STEPS", "4096"))
SPEC_DRAFT_TOPK_PROBE = os.environ.get("SPEC_DRAFT_TOPK_PROBE", "0") == "1" or SPEC_TREE_SPEC
SPEC_DRAFT_TOPK_REQUIRE = os.environ.get("SPEC_DRAFT_TOPK_REQUIRE", "0") == "1"
SPEC_DRAFT_TOPK_K = max(SPEC_TREE_WIDTH, int(os.environ.get("SPEC_DRAFT_TOPK_K", "8")))
TREE_FAIL_ON_PROPOSER_PLACEHOLDER = (
    os.environ.get("TREE_FAIL_ON_PROPOSER_PLACEHOLDER", "1") == "1"
)
TREE_BRANCH_SANITY = os.environ.get("TREE_BRANCH_SANITY", "0") == "1"
TREE_BRANCH_SANITY_EVERY = int(os.environ.get("TREE_BRANCH_SANITY_EVERY", "128"))
TREE_DEBUG_STOP_AFTER_STEPS = int(os.environ.get("TREE_DEBUG_STOP_AFTER_STEPS", "0"))
TREE_SAMPLERPREP_REQUIRE_CACHE = (
    os.environ.get("TREE_SAMPLERPREP_REQUIRE_CACHE", "1") == "1"
)
_ACCEPT_HIST_STATE: dict[str, Any] = {}
_TOPK_SALVAGE_STATE: dict[str, Any] = {}
_DRAFT_TOPK_STATE: dict[str, Any] = {}
_TREE_BRANCH_SANITY_STATE: dict[str, Any] = {}
_TREE_MASK_CACHE: dict[tuple[int, int], Any] = {}


def _get_tree_mask(K: int, W: int, device: Any) -> Any:
    """Create a star tree mask: main path + W-1 branches per position."""
    key = (K, W)
    if key in _TREE_MASK_CACHE:
        return _TREE_MASK_CACHE[key]

    import torch

    # Total tokens = 1 (base) + K (main) + K*(W-1) (branches) = 1 + KW
    N = 1 + K * W
    mask = torch.zeros((N, N), dtype=torch.bool, device=device)

    # Base token (index 0) attends only to itself
    mask[0, 0] = True

    for i in range(1, K + 1):
        # Main path node for pos i: index i
        # Branches for pos i: indices (K+1) + (i-1)*(W-1) ... (K+1) + i*(W-1) - 1
        
        # All candidates for pos i attend to base and main path nodes 1...i-1
        nodes_for_pos_i = [i] + list(range(K + 1 + (i - 1) * (W - 1), K + 1 + i * (W - 1)))
        for node_idx in nodes_for_pos_i:
            mask[node_idx, 0] = True
            for prev_i in range(1, i):
                mask[node_idx, prev_i] = True
            mask[node_idx, node_idx] = True

    _TREE_MASK_CACHE[key] = mask
    return mask


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 // SPEC_TREE_WIDTH
            if SPEC_TREE_SPEC
            else self.num_speculative_tokens
        )
    state["outputs"] = [
        torch.zeros((1, token_count), dtype=torch.int64, device=device)
        for _ in range(LOOPGRAPH_PINGPONG_SLOTS)
    ]
    state["out"] = state["outputs"][0]
    if SPEC_DRAFT_TOPK_PROBE:
        topk_width = max(SPEC_TOPK_SALVAGE_K, SPEC_DRAFT_TOPK_K)
        state["topk_outputs"] = [
            torch.full(
                (token_count, topk_width),
                -1,
                dtype=torch.int64,
                device=device,
            )
            for _ in range(LOOPGRAPH_PINGPONG_SLOTS)
        ]
        state["topk_out"] = state["topk_outputs"][0]
    state["next_slot"] = 0
    state["_pupa_loopgraph_slot_events"] = [
        torch.cuda.Event(blocking=False) for _ in state["outputs"]
    ]
    for output, event in zip(
        state["outputs"], state["_pupa_loopgraph_slot_events"], strict=True
    ):
        _LOOPGRAPH_SLOT_EVENTS_BY_PTR[output.data_ptr()] = event
        _LOOPGRAPH_SLOT_EVENT_RECORDED_BY_PTR[output.data_ptr()] = False
    state["seq_lens"] = torch.zeros_like(cad.seq_lens[:1])
    state["first_input"] = torch.zeros((1,), dtype=torch.int32, device=device)
    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 // SPEC_TREE_WIDTH
            if SPEC_TREE_SPEC
            else self.num_speculative_tokens
        )
    output = state["out"]
    onegraph = state.get("onegraph", False)
    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),
    ):
        if onegraph:
            # K width-1 iterations; iteration 0 reads the target-sampled
            # token from first_input and writes out[0, 0].
            for index in range(token_count):
                source = (
                    state["first_input"]
                    if index == 0
                    else output[0, index - 1 : index]
                )
                self.input_ids[:1].copy_(source)
                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 = _get_top_tokens_maybe_record(
                    self.model,
                    last_hidden[:1],
                    state,
                    index,
                )
                output[0, index : index + 1].copy_(token)
        else:
            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 = _get_top_tokens_maybe_record(
                    self.model,
                    last_hidden[:1],
                    state,
                    index + 1,
                )
                output[0, index + 1 : index + 2].copy_(token)


def _get_top_tokens_maybe_record(
    model: Any,
    hidden_states: Any,
    state: dict[str, Any],
    step: int,
) -> Any:
    embedder = None
    if SPEC_DRAFT_TOPK_PROBE:
        topk_out = state.get("topk_out")
        embedder = getattr(model, "masked_embedding", None)
        if topk_out is not None and embedder is not None:
            embedder._ff_draft_topk_row = topk_out[step : step + 1]
    try:
        return model.get_top_tokens(hidden_states)
    finally:
        if SPEC_DRAFT_TOPK_PROBE and embedder is not None:
            embedder._ff_draft_topk_row = None


def _select_loopgraph_output_slot(state: dict[str, Any]) -> Any:
    import torch

    outputs = state.get("outputs")
    if not outputs:
        return state["out"]

    slot_index = int(state.get("next_slot", 0))
    output_slot = outputs[slot_index]
    event = _LOOPGRAPH_SLOT_EVENTS_BY_PTR.get(output_slot.data_ptr())
    event_recorded = _LOOPGRAPH_SLOT_EVENT_RECORDED_BY_PTR.get(
        output_slot.data_ptr(), False
    )
    if event is not None and event_recorded:
        torch.cuda.current_stream().wait_event(event)
        if LOOPGRAPH_CLEAR_STALE_SLOT_EVENTS:
            _LOOPGRAPH_SLOT_EVENT_RECORDED_BY_PTR[output_slot.data_ptr()] = False

    state["out"] = output_slot
    state["active_slot"] = slot_index
    state["next_slot"] = (slot_index + 1) % len(outputs)
    if SPEC_DRAFT_TOPK_PROBE and "topk_outputs" in state:
        state["topk_out"] = state["topk_outputs"][slot_index]
    return output_slot


def _prime_loopgraph_outputs(state: dict[str, Any], first_token: Any) -> None:
    for output in state.get("outputs", [state["out"]]):
        output[0, 0:1].copy_(first_token)


def _tree_guard_draft_tokens(tokens: Any, label: str) -> Any:
    if not SPEC_TREE_SPEC or not TREE_FAIL_ON_PROPOSER_PLACEHOLDER:
        return tokens
    try:
        if bool((tokens < 0).any().item()):
            head = tokens.reshape(-1)[: min(12, tokens.numel())].detach().cpu().tolist()
            raise RuntimeError(
                "[tree-debug] proposer produced negative tree draft tokens "
                f"at {label}: shape={tuple(tokens.shape)} head={head}"
            )
    except RuntimeError:
        raise
    except Exception as exc:
        if SPEC_DRAFT_TOPK_REQUIRE:
            raise RuntimeError(
                f"[tree-debug] failed to inspect tree draft tokens at {label}"
            ) from exc
    return tokens


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

    graphs = []
    outputs = state.get("outputs", [state["out"]])
    topk_outputs = state.get("topk_outputs", [])
    for index, output in enumerate(outputs):
        state["out"] = output
        if SPEC_DRAFT_TOPK_PROBE and topk_outputs:
            state["topk_out"] = topk_outputs[index]
        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)
        graphs.append(graph)
    state["graphs"] = graphs
    state["graph"] = graphs[0]
    state["out"] = state.get("outputs", [state["out"]])[0]


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 // SPEC_TREE_WIDTH
            if SPEC_TREE_SPEC
            else 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)
                _prime_loopgraph_outputs(state, 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']} "
                    f"with slots={LOOPGRAPH_PINGPONG_SLOTS} (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:
            output_slot = _select_loopgraph_output_slot(state)
            _refresh_static_buffers(self, state, cad)
            output_slot[0, 0:1].copy_(first_token)
            self.hidden_states[:1].copy_(first_hidden)
            graphs = state.get("graphs")
            graph = graphs[state["active_slot"]] if graphs else state["graph"]
            graph.replay()
            if SPEC_DRAFT_TOPK_PROBE:
                _DRAFT_TOPK_STATE["latest_topk"] = state.get("topk_out")
            
            if SPEC_TREE_SPEC:
                import torch
                topk = state.get("topk_out")            # (K, capture_width>=W)
                K = topk.shape[0]
                # emit EXACTLY SPEC_TREE_WIDTH candidates per position:
                # [main(K); branches grouped by position (K*(W-1))]
                # Use the graph body's sampled main path, not the separately
                # reconstructed top-k rank-0 row. This preserves the actual
                # drafter chain even if top-k tie order ever differs.
                main_path = output_slot[0, :K].to(topk.dtype)
                branches = topk[:, 1:SPEC_TREE_WIDTH]   # W-slice fix
                all_tokens = torch.cat([main_path, branches.flatten()])
                all_tokens = _tree_guard_draft_tokens(all_tokens, "loopgraph-tree")
                if not TREE_FAIL_ON_PROPOSER_PLACEHOLDER:
                    all_tokens = torch.clamp(all_tokens, min=0)
                return all_tokens.unsqueeze(0)          # (1, K*W)

            return output_slot

        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)

    def propose_onegraph(
        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, "onegraph": True},
        )
        if not _is_loopgraph_eligible(self, state, common_attn_metadata) or (
            self.uses_xdrope_dim > 0 and self.draft_uses_xdrope_dim > 0
        ):
            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 // SPEC_TREE_WIDTH
            if SPEC_TREE_SPEC
            else self.num_speculative_tokens
        )
        cad = common_attn_metadata

        # Index of the single position whose drafter output is consumed.
        sample_index = token_indices_to_sample
        if sample_index is None:
            sample_index = cad.query_start_loc[1:] - 1

        # Loop-invariant metadata adjustments (mirror of the stock loop setup;
        # the width-(K+1) first pass is skipped entirely).
        cad.num_actual_tokens = 1
        cad.max_query_len = 1
        cad.query_start_loc = self.arange[:2]
        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

        # Width-1 inputs: position + target hidden state of the sampled slot,
        # and the token the target just sampled. All device-side gathers.
        positions_1d = target_positions
        if self.vllm_config.model_config.uses_mrope:
            positions_1d = target_positions[0]
        self.positions[:1] = positions_1d[sample_index]
        self.hidden_states[:1].copy_(target_hidden_states[sample_index])

        if state["graph"] is None and not state["failed"]:
            try:
                _build_static_buffers(self, state, cad)
                _refresh_static_buffers(self, state, cad)
                state["first_input"].copy_(next_token_ids[:1])
                _capture_graph(self, state)
                # Warmup/capture runs consumed the hidden buffer; restore the
                # real value for the replay that serves this step.
                self.hidden_states[:1].copy_(target_hidden_states[sample_index])
                print(
                    f"[onegraph] captured K={token_count} width-1 propose graph "
                    f"at eligible call {state['calls']} "
                    f"with slots={LOOPGRAPH_PINGPONG_SLOTS} (pid {os.getpid()})",
                    file=sys.stderr,
                    flush=True,
                )
            except Exception as exc:
                state["failed"] = True
                state["graph"] = None
                print(
                    f"[onegraph] capture failed: {exc!r}",
                    file=sys.stderr,
                    flush=True,
                )
                _raise_or_fallback(exc)

        if state["graph"] is not None:
            output_slot = _select_loopgraph_output_slot(state)
            _refresh_static_buffers(self, state, cad)
            state["first_input"].copy_(next_token_ids[:1])
            graphs = state.get("graphs")
            graph = graphs[state["active_slot"]] if graphs else state["graph"]
            graph.replay()
            if SPEC_DRAFT_TOPK_PROBE:
                _DRAFT_TOPK_STATE["latest_topk"] = state.get("topk_out")
            
            if SPEC_TREE_SPEC:
                import torch
                topk = state.get("topk_out")            # (K, capture_width>=W)
                K = topk.shape[0]
                # emit EXACTLY SPEC_TREE_WIDTH candidates per position:
                # [main(K); branches grouped by position (K*(W-1))]
                # Use the graph body's sampled main path, not the separately
                # reconstructed top-k rank-0 row. This preserves the actual
                # drafter chain even if top-k tie order ever differs.
                main_path = output_slot[0, :K].to(topk.dtype)
                branches = topk[:, 1:SPEC_TREE_WIDTH]   # W-slice fix
                all_tokens = torch.cat([main_path, branches.flatten()])
                all_tokens = _tree_guard_draft_tokens(all_tokens, "onegraph-tree")
                if not TREE_FAIL_ON_PROPOSER_PLACEHOLDER:
                    all_tokens = torch.clamp(all_tokens, min=0)
                return all_tokens.unsqueeze(0)          # (1, K*W)

            return output_slot

        # Eager width-1 fallback: K iterations, token-equivalent to the
        # captured body. Re-gather hidden (a failed capture's warmup runs
        # may have overwritten the buffer).
        self.hidden_states[:1].copy_(target_hidden_states[sample_index])
        cg_mode, input_batch_size, batch_size_dp = (
            self._determine_batch_execution_and_padding(1)
        )
        _, loop_metadata = self.build_per_group_and_layer_attn_metadata(
            cad,
            draft_index=1,
        )
        draft_tokens: list[Any] = []
        for index in range(token_count):
            if index == 0:
                self.input_ids[:1] = next_token_ids[:1].int()
            else:
                self.input_ids[:1] = draft_tokens[-1].int()
            forward_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(**forward_kwargs)
            self.hidden_states[:1] = 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_onegraph if ONEGRAPH else propose
    print(
        f"[pupa-loopgraph] patched Gemma4Proposer.propose in pid {os.getpid()} "
        f"(warmup_calls={LOOPGRAPH_WARMUP_CALLS}, "
        f"require_capture={LOOPGRAPH_REQUIRE_CAPTURE}, onegraph={ONEGRAPH})",
        file=sys.stderr,
        flush=True,
    )


def _get_spec_decode_metadata_kernel() -> Any:
    global _FUSED_SPEC_DECODE_METADATA_KERNEL
    if _FUSED_SPEC_DECODE_METADATA_KERNEL is not None:
        return _FUSED_SPEC_DECODE_METADATA_KERNEL

    import triton
    import triton.language as tl

    @triton.jit
    def _pupa_batch_one_spec_metadata_kernel(
        logits_indices_ptr,
        target_logits_indices_ptr,
        cu_num_draft_tokens_ptr,
        cu_num_sampled_tokens_ptr,
        bonus_logits_indices_ptr,
        start_offset,
        draft_len,
        sampled_len,
        BLOCK: tl.constexpr,
    ) -> None:
        offsets = tl.arange(0, BLOCK)
        tl.store(
            logits_indices_ptr + offsets,
            start_offset + offsets,
            mask=offsets < sampled_len,
        )
        tl.store(
            target_logits_indices_ptr + offsets,
            offsets,
            mask=offsets < draft_len,
        )
        tl.store(cu_num_draft_tokens_ptr, draft_len)
        tl.store(cu_num_sampled_tokens_ptr, sampled_len)
        tl.store(bonus_logits_indices_ptr, draft_len)

    _FUSED_SPEC_DECODE_METADATA_KERNEL = _pupa_batch_one_spec_metadata_kernel
    return _FUSED_SPEC_DECODE_METADATA_KERNEL


def _get_spec_decode_metadata_state(self: Any, sampled_len: int) -> dict[str, Any]:
    import torch

    capacity = max(2, sampled_len)
    state = self.__dict__.setdefault("_pupa_spec_decode_metadata", {})
    if int(state.get("capacity", 0)) >= capacity:
        return state

    state["capacity"] = capacity
    state["logits_indices"] = torch.empty(
        capacity, dtype=torch.int32, device=self.device
    )
    state["target_logits_indices"] = torch.empty(
        capacity - 1, dtype=torch.int32, device=self.device
    )
    state["cu_num_draft_tokens"] = torch.empty(
        1, dtype=torch.int32, device=self.device
    )
    state["cu_num_sampled_tokens"] = torch.empty(
        1, dtype=torch.int32, device=self.device
    )
    state["bonus_logits_indices"] = torch.empty(
        1, dtype=torch.int32, device=self.device
    )
    return state


def _calc_spec_decode_metadata_fast_batch_one(
    self: Any,
    original_calc_spec_decode_metadata: Any,
    metadata_cls: Any,
    num_draft_tokens: Any,
    cu_num_scheduled_tokens: Any,
) -> Any:
    if not FUSED_SPEC_DECODE_METADATA:
        return original_calc_spec_decode_metadata(
            self, num_draft_tokens, cu_num_scheduled_tokens
        )

    try:
        if getattr(self.device, "type", None) != "cuda":
            raise RuntimeError("fast spec metadata requires CUDA")
        if len(num_draft_tokens) != 1 or len(cu_num_scheduled_tokens) != 1:
            raise RuntimeError("fast spec metadata is batch-1 only")

        draft_len = int(num_draft_tokens[0])
        # tree-v2: NO adjustment — with num_speculative_tokens=K*W the
        # scheduler already reports the expanded draft count and the upstream
        # index math is correct as-is (verified against the wheel source).
        if draft_len <= 0:
            raise RuntimeError(f"unsupported draft_len={draft_len}")
        sampled_len = draft_len + 1
        start_offset = int(cu_num_scheduled_tokens[0]) - sampled_len
        if start_offset < 0:
            raise RuntimeError(f"invalid start_offset={start_offset}")

        state = _get_spec_decode_metadata_state(self, sampled_len)
        logits_indices = state["logits_indices"][:sampled_len]
        target_logits_indices = state["target_logits_indices"][:draft_len]
        cu_num_draft_tokens = state["cu_num_draft_tokens"]
        cu_num_sampled_tokens = state["cu_num_sampled_tokens"]
        bonus_logits_indices = state["bonus_logits_indices"]

        kernel = _get_spec_decode_metadata_kernel()
        block = _next_power_of_2(sampled_len)
        kernel[(1,)](
            logits_indices,
            target_logits_indices,
            cu_num_draft_tokens,
            cu_num_sampled_tokens,
            bonus_logits_indices,
            start_offset,
            draft_len,
            sampled_len,
            BLOCK=block,
        )
        draft_token_ids = self.input_ids.gpu[logits_indices[1:sampled_len]]
        return metadata_cls(
            draft_token_ids=draft_token_ids,
            num_draft_tokens=[draft_len],
            cu_num_draft_tokens=cu_num_draft_tokens,
            cu_num_sampled_tokens=cu_num_sampled_tokens,
            target_logits_indices=target_logits_indices,
            bonus_logits_indices=bonus_logits_indices,
            logits_indices=logits_indices,
        )
    except Exception as exc:
        if FUSED_SPEC_DECODE_METADATA_REQUIRE:
            raise RuntimeError(
                "FUSED_SPEC_DECODE_METADATA_REQUIRE=1 but fusion failed"
            ) from exc
        if not getattr(self, "_pupa_spec_decode_metadata_warned", False):
            self._pupa_spec_decode_metadata_warned = True
            print(
                f"[pupa-gpumeta] falling back to upstream metadata: {exc!r}",
                file=sys.stderr,
                flush=True,
            )
        return original_calc_spec_decode_metadata(
            self, num_draft_tokens, cu_num_scheduled_tokens
        )


def _loopgraph_should_record_copy_event(
    self: Any,
    scheduler_output: Any,
    zeros_only: bool,
) -> bool:
    if zeros_only:
        return False
    if not getattr(self, "use_async_scheduling", False):
        return True
    return bool(
        scheduler_output.has_structured_output_requests
        or self.input_batch.sampling_metadata.output_token_ids
    )


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

    runner_cls = module.GPUModelRunner
    original_copy_draft_token_ids_to_cpu = runner_cls._copy_draft_token_ids_to_cpu
    original_calc_spec_decode_metadata = runner_cls._calc_spec_decode_metadata

    def _copy_draft_token_ids_to_cpu(
        self: Any,
        scheduler_output: Any,
        zeros_only: bool = False,
    ) -> Any:
        draft_token_ids = getattr(self, "_draft_token_ids", None)
        result = original_copy_draft_token_ids_to_cpu(
            self, scheduler_output, zeros_only=zeros_only
        )
        if not _loopgraph_should_record_copy_event(
            self, scheduler_output, zeros_only
        ):
            return result
        if not torch.is_tensor(draft_token_ids):
            return result

        event = _LOOPGRAPH_SLOT_EVENTS_BY_PTR.get(draft_token_ids.data_ptr())
        copy_stream = getattr(self, "draft_token_ids_copy_stream", None)
        if event is not None and copy_stream is not None:
            with torch.cuda.stream(copy_stream):
                event.record(copy_stream)
            _LOOPGRAPH_SLOT_EVENT_RECORDED_BY_PTR[draft_token_ids.data_ptr()] = True
        return result

    def _calc_spec_decode_metadata(
        self: Any,
        num_draft_tokens: Any,
        cu_num_scheduled_tokens: Any,
    ) -> Any:
        return _calc_spec_decode_metadata_fast_batch_one(
            self,
            original_calc_spec_decode_metadata,
            module.SpecDecodeMetadata,
            num_draft_tokens,
            cu_num_scheduled_tokens,
        )

    runner_cls._copy_draft_token_ids_to_cpu = _copy_draft_token_ids_to_cpu
    runner_cls._calc_spec_decode_metadata = _calc_spec_decode_metadata
    print(
        f"[pupa-loopgraph] patched GPUModelRunner draft-token copy events "
        f"and spec metadata in pid {os.getpid()} "
        f"(slots={LOOPGRAPH_PINGPONG_SLOTS}, "
        f"gpumeta={FUSED_SPEC_DECODE_METADATA})",
        file=sys.stderr,
        flush=True,
    )


def _get_fused_greedy_rejection_prep_kernel() -> Any:
    global _FUSED_GREEDY_REJECTION_PREP_KERNEL
    if _FUSED_GREEDY_REJECTION_PREP_KERNEL is not None:
        return _FUSED_GREEDY_REJECTION_PREP_KERNEL

    import triton
    import triton.language as tl

    @triton.jit(do_not_specialize=["max_spec_len"])
    def _pupa_rejection_greedy_with_prep_kernel(
        output_token_ids_ptr,
        next_token_ids_ptr,
        valid_sampled_tokens_count_ptr,
        cu_num_draft_tokens_ptr,
        draft_token_ids_ptr,
        target_argmax_ptr,
        bonus_token_ids_ptr,
        max_spec_len,
    ) -> None:
        req_idx = tl.program_id(0)
        start_idx = 0
        if req_idx != 0:
            start_idx = tl.load(cu_num_draft_tokens_ptr + req_idx - 1)
        end_idx = tl.load(cu_num_draft_tokens_ptr + req_idx)
        num_draft_tokens = end_idx - start_idx

        rejected = False
        valid_count = 0
        next_token_id = tl.load(bonus_token_ids_ptr + req_idx).to(tl.int32)
        row_offset = req_idx * (max_spec_len + 1)
        for pos in range(num_draft_tokens):
            if not rejected:
                draft_token_id = tl.load(draft_token_ids_ptr + start_idx + pos)
                target_argmax_id = tl.load(target_argmax_ptr + start_idx + pos).to(
                    tl.int32
                )
                rejected = draft_token_id != target_argmax_id
                valid_count = pos + 1
                next_token_id = target_argmax_id
                tl.store(output_token_ids_ptr + row_offset + pos, target_argmax_id)

        if not rejected:
            bonus_token_id = tl.load(bonus_token_ids_ptr + req_idx).to(tl.int32)
            valid_count = num_draft_tokens + 1
            next_token_id = bonus_token_id
            tl.store(
                output_token_ids_ptr + row_offset + num_draft_tokens,
                bonus_token_id,
            )

        tl.store(next_token_ids_ptr + req_idx, next_token_id)
        tl.store(valid_sampled_tokens_count_ptr + req_idx, valid_count)

    _FUSED_GREEDY_REJECTION_PREP_KERNEL = _pupa_rejection_greedy_with_prep_kernel
    return _FUSED_GREEDY_REJECTION_PREP_KERNEL


def _has_tracked_thinking_state(sampling_metadata: Any) -> bool:
    holder = getattr(sampling_metadata, "thinking_budget_state_holder", None)
    return holder is not None and holder.has_tracked_requests()


def _direct_rejection_is_deterministic(sampling_metadata: Any) -> bool:
    return bool(
        getattr(sampling_metadata, "all_greedy", False)
        or (
            SPEC_ASSUME_DETERMINISTIC_REJECTION
            and not getattr(sampling_metadata, "all_random", False)
        )
    )


def _safe_for_argmax_rejection(sampling_metadata: Any) -> bool:
    logitsprocs = getattr(sampling_metadata, "logitsprocs", None)
    non_argmax_processors = getattr(logitsprocs, "non_argmax_invariant", ())
    return (
        _direct_rejection_is_deterministic(sampling_metadata)
        and not getattr(sampling_metadata, "all_random", False)
        and getattr(sampling_metadata, "max_num_logprobs", None) is None
        and getattr(sampling_metadata, "no_penalties", False)
        and not getattr(sampling_metadata, "bad_words_token_ids", None)
        and getattr(sampling_metadata, "allowed_token_ids_mask", None) is None
        and (SPEC_ASSUME_ARGMAX_PROCESSORS_OK or not non_argmax_processors)
        and not _has_tracked_thinking_state(sampling_metadata)
    )


def _record_accept_histogram(valid_counts: Any) -> None:
    """Observation-only emitted-token histogram: accepted drafts + bonus token."""
    if not SPEC_ACCEPT_HISTOGRAM:
        return
    try:
        import torch

        state = _ACCEPT_HIST_STATE
        hist = state.get("hist")
        if hist is None or hist.device != valid_counts.device:
            hist = torch.zeros(64, dtype=torch.int64, device=valid_counts.device)
            state["hist"] = hist
            state["ones"] = torch.ones(1, dtype=torch.int64, device=valid_counts.device)
            state["steps"] = 0
        idx = valid_counts.long().clamp_(0, 63)
        ones = state["ones"]
        if idx.shape[0] != 1:
            ones = torch.ones_like(idx)
        hist.index_add_(0, idx, ones)
        state["steps"] += 1
        if state["steps"] in (256, 1024) or state["steps"] % 2048 == 0:
            counts = hist.tolist()
            top = max(i for i, c in enumerate(counts) if c) if any(counts) else 0
            print(
                f"[accept-hist] steps={state['steps']} "
                f"valid_counts_hist={counts[: top + 1]}",
                file=sys.stderr,
                flush=True,
            )
    except Exception:
        pass


def _record_topk_salvage(
    metadata: Any,
    draft_probs: Any,
    target_argmax: Any,
    valid_counts: Any,
) -> None:
    """Gate tree speculation by checking target-token rank at first reject."""
    if not SPEC_TOPK_SALVAGE_PROBE:
        return
    state = _TOPK_SALVAGE_STATE
    steps = int(state.get("steps", 0)) + 1
    state["steps"] = steps
    if SPEC_TOPK_SALVAGE_MAX_STEPS and steps > SPEC_TOPK_SALVAGE_MAX_STEPS:
        return
    try:
        import torch

        device = target_argmax.device
        counts = state.get("counts")
        if counts is None or counts.device != device:
            counts = torch.zeros(10, dtype=torch.int64, device=device)
            state["counts"] = counts

        batch_size = valid_counts.shape[0]
        starts = torch.zeros(batch_size, dtype=torch.long, device=device)
        if batch_size > 1:
            starts[1:] = metadata.cu_num_draft_tokens[:-1].long()
        valid_long = valid_counts.long()
        full_accept = valid_long > int(metadata.max_spec_len)
        rejected = ~full_accept
        counts[0] += batch_size
        counts[1] += full_accept.long().sum()
        counts[2] += rejected.long().sum()

        if bool(rejected.any()):
            row_ids = starts[rejected] + valid_long[rejected] - 1
            target_rows = row_ids.clamp(0, target_argmax.shape[0] - 1)
            target_ids = target_argmax[target_rows].long()
            top_ids = None
            source_index = None

            if draft_probs is not None and draft_probs.ndim == 2:
                prob_rows = row_ids.clamp(0, draft_probs.shape[0] - 1)
                k = min(SPEC_TOPK_SALVAGE_K, draft_probs.shape[-1])
                top_ids = torch.topk(draft_probs[prob_rows], k=k, dim=-1).indices.long()
                source_index = 8
            else:
                draft_topk = _DRAFT_TOPK_STATE.get("latest_topk")
                if (
                    batch_size == 1
                    and draft_topk is not None
                    and getattr(draft_topk, "ndim", 0) == 2
                    and getattr(draft_topk, "device", None) == device
                ):
                    topk_rows = valid_long[rejected] - 1
                    topk_rows = topk_rows.clamp(0, draft_topk.shape[0] - 1)
                    top_ids = draft_topk[topk_rows, :SPEC_TOPK_SALVAGE_K].long()
                    source_index = 9

            if top_ids is None:
                counts[7] += rejected.long().sum()
            else:
                valid_top_ids = top_ids[:, 0] >= 0
                counts[7] += (~valid_top_ids).long().sum()
                if source_index is not None:
                    counts[source_index] += valid_top_ids.long().sum()
                top_ids = top_ids[valid_top_ids]
                target_ids = target_ids[valid_top_ids]
            if top_ids is not None and int(top_ids.shape[0]) > 0:
                k = int(top_ids.shape[-1])
                matches = top_ids == target_ids[:, None]
                counts[3] += matches[:, : min(1, k)].any(dim=1).long().sum()
                counts[4] += matches[:, : min(2, k)].any(dim=1).long().sum()
                counts[5] += matches[:, : min(4, k)].any(dim=1).long().sum()
                counts[6] += matches.any(dim=1).long().sum()

        if steps in (256, 1024) or steps % 2048 == 0:
            vals = counts.tolist()
            total, full, rejected, top1, top2, top4, topk, missing = vals[:8]
            source_probs, source_draft_topk = vals[8:10]
            denom = max(rejected, 1)
            print(
                "[topk-salvage] "
                f"steps={steps} total={total} full_accept={full} "
                f"first_reject={rejected} top1={top1} "
                f"top2={top2} ({top2 / denom:.4f}) "
                f"top4={top4} ({top4 / denom:.4f}) "
                f"top{SPEC_TOPK_SALVAGE_K}={topk} ({topk / denom:.4f}) "
                f"missing_candidates={missing} "
                f"source_probs={source_probs} "
                f"source_draft_topk={source_draft_topk}",
                file=sys.stderr,
                flush=True,
            )
    except Exception as exc:
        if not state.get("warned"):
            state["warned"] = True
            print(
                f"[topk-salvage] disabled after error: {exc!r}",
                file=sys.stderr,
                flush=True,
            )


def _tree_branch_sanity_line(state: dict[str, Any]) -> str:
    first_reject = max(1, int(state.get("first_reject", 0)))
    steps = max(1, int(state.get("steps", 0)))
    rank2 = int(state.get("rank2", 0))
    rank3 = int(state.get("rank3", 0))
    rank4 = int(state.get("rank4", 0))
    any_hit = int(state.get("any_hit", 0))
    kernel_salv = int(state.get("kernel_salv", 0))
    valid_sum = int(state.get("valid_sum", 0))
    return (
        f"steps={state.get('steps', 0)} full={state.get('full', 0)} "
        f"first_reject={state.get('first_reject', 0)} "
        f"neg_steps={state.get('neg_steps', 0)} errors={state.get('errors', 0)} "
        f"rank2={rank2} ({rank2 / first_reject:.4f}) "
        f"rank3={rank3} ({rank3 / first_reject:.4f}) "
        f"rank4={rank4} ({rank4 / first_reject:.4f}) "
        f"any_branch={any_hit} ({any_hit / first_reject:.4f}) "
        f"kernel_salv={kernel_salv} ({kernel_salv / first_reject:.4f}) "
        f"kernel_mismatch={state.get('kernel_mismatch', 0)} "
        f"mean_valid={valid_sum / steps:.3f}"
    )


def _record_tree_branch_sanity(
    target_argmax: Any,
    draft_token_ids: Any,
    K: int,
    W: int,
    salv_node: int,
    valid_count: int,
    label: str,
) -> None:
    if not TREE_BRANCH_SANITY:
        return
    state = _TREE_BRANCH_SANITY_STATE
    state["steps"] = int(state.get("steps", 0)) + 1
    state["valid_sum"] = int(state.get("valid_sum", 0)) + int(valid_count)
    kernel_hit = bool(salv_node)
    if kernel_hit:
        state["kernel_salv"] = int(state.get("kernel_salv", 0)) + 1

    try:
        import torch

        flat = draft_token_ids.reshape(-1)[: K * W].long()
        target = target_argmax.reshape(-1).long()
        if int(flat.numel()) < K * W or int(target.numel()) < K:
            raise RuntimeError(
                f"short tensors flat={int(flat.numel())} target={int(target.numel())}"
            )
        if bool((flat < 0).any().item()):
            state["neg_steps"] = int(state.get("neg_steps", 0)) + 1
        main = flat[:K]
        branches = flat[K : K + K * (W - 1)].reshape(K, W - 1)
        target_main = target[:K]
        misses = main != target_main
        if not bool(misses.any().item()):
            state["full"] = int(state.get("full", 0)) + 1
            any_hit = False
        else:
            state["first_reject"] = int(state.get("first_reject", 0)) + 1
            reject_idx = int(torch.nonzero(misses, as_tuple=False)[0].item())
            target_id = target_main[reject_idx]
            hits = branches[reject_idx] == target_id
            any_hit = bool(hits.any().item())
            if any_hit:
                state["any_hit"] = int(state.get("any_hit", 0)) + 1
            if W >= 2 and bool(hits[0].item()):
                state["rank2"] = int(state.get("rank2", 0)) + 1
            if W >= 3 and bool(hits[1].item()):
                state["rank3"] = int(state.get("rank3", 0)) + 1
            if W >= 4 and bool(hits[2].item()):
                state["rank4"] = int(state.get("rank4", 0)) + 1
        if any_hit != kernel_hit:
            state["kernel_mismatch"] = int(state.get("kernel_mismatch", 0)) + 1
    except Exception as exc:
        state["errors"] = int(state.get("errors", 0)) + 1
        if state["errors"] <= 3:
            print(
                f"[tree-branch-sanity] error label={label}: {exc!r}",
                file=sys.stderr,
                flush=True,
            )

    steps = int(state.get("steps", 0))
    if steps == 1 or (
        TREE_BRANCH_SANITY_EVERY > 0 and steps % TREE_BRANCH_SANITY_EVERY == 0
    ):
        print(
            f"[tree-branch-sanity] label={label} {_tree_branch_sanity_line(state)}",
            file=sys.stderr,
            flush=True,
        )
    if TREE_DEBUG_STOP_AFTER_STEPS > 0 and steps >= TREE_DEBUG_STOP_AFTER_STEPS:
        raise RuntimeError(
            "[tree-debug] intentional diagnostic stop after "
            f"{steps} tree steps: {_tree_branch_sanity_line(state)}"
        )


def _can_use_direct_greedy_rejection(
    self: Any,
    metadata: Any,
    draft_probs: Any,
    logits: Any,
    sampling_metadata: Any,
) -> bool:
    return (
        FUSED_DIRECT_GREEDY_REJECTION
        # tree-v2 v3: SPEC_TREE_SPEC allowed — the star salvage walk is
        # the first branch inside _direct_greedy_rejection (run5 fix)
        and not self.synthetic_mode
        and _safe_for_argmax_rejection(sampling_metadata)
        and logits is not None
        and logits.ndim == 2
        and metadata.draft_token_ids.ndim == 1
        and metadata.cu_num_draft_tokens.ndim == 1
        and metadata.target_logits_indices.ndim == 1
        and metadata.bonus_logits_indices.ndim in (1, 2)
    )


def _debug_direct_gate_miss(
    self: Any,
    metadata: Any,
    logits: Any,
    sampling_metadata: Any,
) -> None:
    if not SPEC_ACCEPT_HISTOGRAM or getattr(self, "_ff_direct_gate_debugged", False):
        return
    self._ff_direct_gate_debugged = True
    try:
        logitsprocs = getattr(sampling_metadata, "logitsprocs", None)
        non_argmax_processors = getattr(logitsprocs, "non_argmax_invariant", ())
        print(
            "[accept-hist] direct gate miss "
            f"synthetic={getattr(self, 'synthetic_mode', None)} "
            f"all_greedy={getattr(sampling_metadata, 'all_greedy', None)} "
            f"all_random={getattr(sampling_metadata, 'all_random', None)} "
            f"assume_det={SPEC_ASSUME_DETERMINISTIC_REJECTION} "
            f"max_logprobs={getattr(sampling_metadata, 'max_num_logprobs', None)} "
            f"no_penalties={getattr(sampling_metadata, 'no_penalties', None)} "
            f"bad_words={bool(getattr(sampling_metadata, 'bad_words_token_ids', None))} "
            f"allowed_mask={getattr(sampling_metadata, 'allowed_token_ids_mask', None) is not None} "
            f"processors={bool(non_argmax_processors)} "
            f"assume_processors_ok={SPEC_ASSUME_ARGMAX_PROCESSORS_OK} "
            f"tracked_thinking={_has_tracked_thinking_state(sampling_metadata)} "
            f"logits_shape={None if logits is None else tuple(logits.shape)} "
            f"draft_ndim={getattr(metadata.draft_token_ids, 'ndim', None)} "
            f"cu_ndim={getattr(metadata.cu_num_draft_tokens, 'ndim', None)} "
            f"target_ndim={getattr(metadata.target_logits_indices, 'ndim', None)} "
            f"bonus_ndim={getattr(metadata.bonus_logits_indices, 'ndim', None)}",
            file=sys.stderr,
            flush=True,
        )
    except Exception:
        pass


def _direct_greedy_rejection(
    module: Any,
    metadata: Any,
    draft_probs: Any,
    logits: Any,
    sampling_metadata: Any,
) -> Any:
    import torch

    batch_size = len(metadata.num_draft_tokens)
    output_token_ids = torch.full(
        (batch_size, metadata.max_spec_len + 1),
        module.PLACEHOLDER_TOKEN_ID,
        dtype=torch.int32,
        device=logits.device,
    )
    next_token_ids = torch.empty(batch_size, dtype=torch.int32, device=logits.device)
    valid_counts = torch.empty(batch_size, dtype=torch.int32, device=logits.device)

    argmax_token_ids = logits.argmax(dim=-1)
    target_argmax = argmax_token_ids[metadata.target_logits_indices].contiguous()
    bonus_token_ids = argmax_token_ids[
        metadata.bonus_logits_indices.reshape(-1)
    ].contiguous()

    # tree-v2 direct path fix (cheesetaco-cdx): RejectionSampler.forward()
    # reaches this direct helper before the module-level rejection_sample()
    # wrapper below. Without this branch, SPEC_TREE_SPEC silently falls back to
    # linear chain rejection over the expanded [main; branches] draft layout,
    # so sibling salvage never runs and salvaged KV is never relocated.
    if SPEC_TREE_SPEC and batch_size == 1:
        import tree_v2_ext as _T2

        K_tree = _T2.TREE_K
        expected = K_tree * _T2.TREE_W
        if K_tree > 0 and metadata.max_spec_len == K_tree * _T2.TREE_W:
            output_token_ids, valid_counts, salv = _T2.star_reject_from_drafts(
                target_argmax,
                metadata.draft_token_ids[: metadata.max_spec_len],
                bonus_token_ids.reshape(-1)[:1],
                K_tree,
                module.PLACEHOLDER_TOKEN_ID,
                logits.device,
            )
            next_tokens = torch.gather(
                output_token_ids[0], 0, (valid_counts.to(torch.int64) - 1)
            ).to(torch.int32)
            salv_node, valid_i = int(salv[0]), int(valid_counts[0])
            _record_tree_branch_sanity(
                target_argmax,
                metadata.draft_token_ids[: metadata.max_spec_len],
                K_tree,
                _T2.TREE_W,
                salv_node,
                valid_i,
                "direct",
            )
            if salv_node:
                _T2.relocate_salvaged_kv(salv_node, K_tree)
            _T2.record_step(valid_i, salv_node, K_tree)
            _PUPA_GREEDY_PREP_CACHE[output_token_ids.data_ptr()] = (
                next_tokens,
                valid_counts,
            )
            _record_accept_histogram(valid_counts)
            return module.SamplerOutput(
                sampled_token_ids=output_token_ids,
                logprobs_tensors=None,
            )
        tree_active = bool(getattr(_T2, "_STATE", {}).get("tree_step", False))
        if _T2.TREE_REQUIRE and tree_active:
            raise RuntimeError(
                f"[tree-v2] REQUIRE: direct rejection saw max_spec_len="
                f"{metadata.max_spec_len}, expected {expected}"
            )

    kernel = _get_fused_greedy_rejection_prep_kernel()
    kernel[(batch_size,)](
        output_token_ids,
        next_token_ids,
        valid_counts,
        metadata.cu_num_draft_tokens,
        metadata.draft_token_ids,
        target_argmax,
        bonus_token_ids,
        metadata.max_spec_len,
    )
    _PUPA_GREEDY_PREP_CACHE[output_token_ids.data_ptr()] = (
        next_token_ids,
        valid_counts,
    )
    _record_accept_histogram(valid_counts)
    _record_topk_salvage(metadata, draft_probs, target_argmax, valid_counts)
    return module.SamplerOutput(
        sampled_token_ids=output_token_ids,
        logprobs_tensors=None,
    )


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

    rejection_sampler_cls = module.RejectionSampler
    original_forward = rejection_sampler_cls.forward
    original_rejection_sample = module.rejection_sample

    def forward(
        self: Any,
        metadata: Any,
        draft_probs: Any,
        logits: Any,
        sampling_metadata: Any,
    ) -> Any:
        if _can_use_direct_greedy_rejection(
            self,
            metadata,
            draft_probs,
            logits,
            sampling_metadata,
        ):
            try:
                return _direct_greedy_rejection(
                    module,
                    metadata,
                    draft_probs,
                    logits,
                    sampling_metadata,
                )
            except Exception as exc:
                if FUSED_DIRECT_GREEDY_REJECTION_REQUIRE:
                    raise RuntimeError(
                        "FUSED_DIRECT_GREEDY_REJECTION_REQUIRE=1 but fusion failed"
                    ) from exc
                if not getattr(self, "_pupa_direct_greedy_rejection_warned", False):
                    self._pupa_direct_greedy_rejection_warned = True
                    print(
                        f"[pupa-directreject] falling back to upstream forward: {exc!r}",
                        file=sys.stderr,
                        flush=True,
                    )
        else:
            _debug_direct_gate_miss(self, metadata, logits, sampling_metadata)
        return original_forward(
            self,
            metadata,
            draft_probs,
            logits,
            sampling_metadata,
        )

    def rejection_sample(
        draft_token_ids: Any,
        num_draft_tokens: list[int],
        max_spec_len: int,
        cu_num_draft_tokens: Any,
        draft_probs: Any,
        target_logits: Any,
        bonus_token_ids: Any,
        sampling_metadata: Any,
        synthetic_mode: bool = False,
        synthetic_conditional_rates: Any = None,
        use_fp64_gumbel: bool = False,
    ) -> Any:
        # Standard fast path checks
        if (
            not FUSED_GREEDY_REJECTION_PREP
            or synthetic_mode
            or not _safe_for_argmax_rejection(sampling_metadata)
            or draft_token_ids.ndim != 1
            or cu_num_draft_tokens.ndim != 1
            or target_logits.ndim != 2
            or not draft_token_ids.is_contiguous()
            or not bonus_token_ids.is_contiguous()
        ):
            # Fallback to upstream
            return original_rejection_sample(
                draft_token_ids,
                num_draft_tokens,
                max_spec_len,
                cu_num_draft_tokens,
                draft_probs,
                target_logits,
                bonus_token_ids,
                sampling_metadata,
                synthetic_mode,
                synthetic_conditional_rates,
                use_fp64_gumbel,
            )

        batch_size = len(num_draft_tokens)
        num_tokens = draft_token_ids.shape[0]
        vocab_size = target_logits.shape[-1]
        device = target_logits.device

        # Star Tree Salvage v2 (chiku-inu) — fused walk on the pinned
        # layout; topk reconstructed from the scheduled drafts themselves.
        # Validated: tree-dev/sim_tree_identity.py (240-token greedy identity)
        # + tree-v2/test_ext_walk.py (2000/2000 vs that reference).
        if SPEC_TREE_SPEC and batch_size == 1:
            import tree_v2_ext as _T2

            K_tree = _T2.TREE_K
            expected = K_tree * _T2.TREE_W
            if K_tree > 0 and max_spec_len == K_tree * _T2.TREE_W:
                target_argmax = target_logits.argmax(dim=-1)
                output_token_ids, valid_counts, salv = _T2.star_reject_from_drafts(
                    target_argmax, draft_token_ids[:max_spec_len],
                    bonus_token_ids.reshape(-1)[:1], K_tree,
                    module.PLACEHOLDER_TOKEN_ID, device,
                )
                next_tokens = torch.gather(
                    output_token_ids[0], 0, (valid_counts.to(torch.int64) - 1)
                ).to(torch.int32)
                # one small host sync per step (rides with the step's existing
                # output transfer); relocation only on salvage steps (~30%)
                salv_node, valid_i = int(salv[0]), int(valid_counts[0])
                _record_tree_branch_sanity(
                    target_argmax,
                    draft_token_ids[:max_spec_len],
                    K_tree,
                    _T2.TREE_W,
                    salv_node,
                    valid_i,
                    "module",
                )
                if salv_node:
                    _T2.relocate_salvaged_kv(salv_node, K_tree)
                _T2.record_step(valid_i, salv_node, K_tree)
                _PUPA_GREEDY_PREP_CACHE[output_token_ids.data_ptr()] = (
                    next_tokens, valid_counts,
                )
                return output_token_ids
            tree_active = bool(getattr(_T2, "_STATE", {}).get("tree_step", False))
            if _T2.TREE_REQUIRE and tree_active:
                raise RuntimeError(
                    f"[tree-v2] REQUIRE: rejection saw max_spec_len={max_spec_len}, "
                    f"expected {expected}"
                )

        # Standard fused path
        if target_logits.shape != (num_tokens, vocab_size):
            return original_rejection_sample(
                draft_token_ids,
                num_draft_tokens,
                max_spec_len,
                cu_num_draft_tokens,
                draft_probs,
                target_logits,
                bonus_token_ids,
                sampling_metadata,
                synthetic_mode,
                synthetic_conditional_rates,
                use_fp64_gumbel,
            )

        output_token_ids = torch.full(
            (batch_size, max_spec_len + 1),
            module.PLACEHOLDER_TOKEN_ID,
            dtype=torch.int32,
            device=device,
        )
        next_token_ids = torch.empty(batch_size, dtype=torch.int32, device=device)
        valid_counts = torch.empty(batch_size, dtype=torch.int32, device=device)
        target_argmax = target_logits.argmax(dim=-1)
        kernel = _get_fused_greedy_rejection_prep_kernel()
        kernel[(batch_size,)](
            output_token_ids,
            next_token_ids,
            valid_counts,
            cu_num_draft_tokens,
            draft_token_ids,
            target_argmax,
            bonus_token_ids,
            max_spec_len,
        )
        _PUPA_GREEDY_PREP_CACHE[output_token_ids.data_ptr()] = (
            next_token_ids,
            valid_counts,
        )
        _record_accept_histogram(valid_counts)
        return output_token_ids

    rejection_sampler_cls.forward = forward
    module.rejection_sample = rejection_sample
    print(
        f"[pupa-samplerprep] patched greedy rejection/prep fusion "
        f"in pid {os.getpid()} (enabled={FUSED_GREEDY_REJECTION_PREP}, "
        f"direct={FUSED_DIRECT_GREEDY_REJECTION})",
        file=sys.stderr,
        flush=True,
    )


def _apply_samplerprep_proposer_patch(module: Any) -> None:
    proposer_cls = module.SpecDecodeBaseProposer
    original_prepare_next_token_ids_padded = proposer_cls.prepare_next_token_ids_padded
    cache_stats = {"hit": 0, "miss_placeholder": 0}

    def prepare_next_token_ids_padded(
        self: Any,
        sampled_token_ids: Any,
        requests: dict[str, Any],
        gpu_input_batch: Any,
        discard_request_mask: Any,
    ) -> tuple[Any, Any]:
        cached = _PUPA_GREEDY_PREP_CACHE.pop(sampled_token_ids.data_ptr(), None)
        if (
            FUSED_GREEDY_REJECTION_PREP
            and cached is not None
            and sampled_token_ids.shape[0] == gpu_input_batch.num_reqs
            and gpu_input_batch.num_reqs == 1
        ):
            cache_stats["hit"] += 1
            return cached
        if (
            SPEC_TREE_SPEC
            and FUSED_GREEDY_REJECTION_PREP
            and TREE_SAMPLERPREP_REQUIRE_CACHE
        ):
            try:
                has_placeholder = bool((sampled_token_ids < 0).any().item())
            except Exception:
                has_placeholder = False
            if has_placeholder:
                cache_stats["miss_placeholder"] += 1
                raise RuntimeError(
                    "tree-v2 placeholder row reached upstream "
                    "prepare_next_token_ids_padded without a samplerprep cache "
                    f"hit; shape={tuple(sampled_token_ids.shape)} "
                    f"cache_stats={cache_stats}"
                )
        return original_prepare_next_token_ids_padded(
            self,
            sampled_token_ids,
            requests,
            gpu_input_batch,
            discard_request_mask,
        )

    proposer_cls.prepare_next_token_ids_padded = prepare_next_token_ids_padded
    print(
        f"[pupa-samplerprep] patched SpecDecodeBaseProposer.prepare_next_token_ids_padded "
        f"in pid {os.getpid()} (enabled={FUSED_GREEDY_REJECTION_PREP})",
        file=sys.stderr,
        flush=True,
    )


def _next_power_of_2(value: int) -> int:
    return 1 << (max(1, value) - 1).bit_length()


def _get_fused_sparse_argmax_kernels() -> Any:
    global _FUSED_SPARSE_ARGMAX_KERNELS
    if _FUSED_SPARSE_ARGMAX_KERNELS is not None:
        return _FUSED_SPARSE_ARGMAX_KERNELS

    import triton
    import triton.language as tl

    @triton.jit
    def _sparse_argmax_blocks_kernel(
        hidden_states,
        lm_head_weight,
        top_centroids,
        token_ordering,
        partial_scores,
        partial_tokens,
        hidden_stride_t,
        hidden_stride_d,
        lm_head_stride_v,
        lm_head_stride_d,
        top_stride_t,
        top_stride_k,
        partial_score_stride_t,
        partial_token_stride_t,
        VOCAB_PER_CENTROID: tl.constexpr,
        SELECTED_COUNT: tl.constexpr,
        HIDDEN_SIZE: tl.constexpr,
        BLOCK_SELECTED: tl.constexpr,
        BLOCK_D: tl.constexpr,
    ) -> None:
        token_idx = tl.program_id(0)
        selected_block = tl.program_id(1)

        selected_offsets = selected_block * BLOCK_SELECTED + tl.arange(
            0, BLOCK_SELECTED
        )
        valid_selected = selected_offsets < SELECTED_COUNT
        centroid_slots = selected_offsets // VOCAB_PER_CENTROID
        token_slots = selected_offsets - centroid_slots * VOCAB_PER_CENTROID

        centroid_ids = tl.load(
            top_centroids + token_idx * top_stride_t + centroid_slots * top_stride_k,
            mask=valid_selected,
            other=0,
        )
        vocab_ids = tl.load(
            token_ordering + centroid_ids * VOCAB_PER_CENTROID + token_slots,
            mask=valid_selected,
            other=0,
        )

        d_offsets = tl.arange(0, BLOCK_D)
        valid_d = d_offsets < HIDDEN_SIZE
        hidden = tl.load(
            hidden_states + token_idx * hidden_stride_t + d_offsets * hidden_stride_d,
            mask=valid_d,
            other=0.0,
        ).to(tl.float32)
        weights = tl.load(
            lm_head_weight
            + vocab_ids[:, None] * lm_head_stride_v
            + d_offsets[None, :] * lm_head_stride_d,
            mask=valid_selected[:, None] & valid_d[None, :],
            other=0.0,
        ).to(tl.float32)
        scores = tl.sum(weights * hidden[None, :], axis=1)
        # The PyTorch sparse path materializes bf16 logits before argmax.
        scores = scores.to(tl.bfloat16).to(tl.float32)
        scores = tl.where(valid_selected, scores, -float("inf"))
        best_score, best_local_idx = tl.max(
            scores,
            axis=0,
            return_indices=True,
            return_indices_tie_break_left=True,
        )

        best_selected = selected_block * BLOCK_SELECTED + best_local_idx
        best_centroid_slot = best_selected // VOCAB_PER_CENTROID
        best_token_slot = best_selected - best_centroid_slot * VOCAB_PER_CENTROID
        best_centroid = tl.load(
            top_centroids + token_idx * top_stride_t + best_centroid_slot * top_stride_k
        )
        best_token = tl.load(
            token_ordering + best_centroid * VOCAB_PER_CENTROID + best_token_slot
        )
        tl.store(
            partial_scores + token_idx * partial_score_stride_t + selected_block,
            best_score,
        )
        tl.store(
            partial_tokens + token_idx * partial_token_stride_t + selected_block,
            best_token,
        )

    @triton.jit
    def _sparse_argmax_reduce_kernel(
        partial_scores,
        partial_tokens,
        output_tokens,
        partial_score_stride_t,
        partial_token_stride_t,
        output_stride_t,
        NUM_BLOCKS: tl.constexpr,
        BLOCK_BLOCKS: tl.constexpr,
    ) -> None:
        token_idx = tl.program_id(0)
        block_offsets = tl.arange(0, BLOCK_BLOCKS)
        valid_blocks = block_offsets < NUM_BLOCKS
        scores = tl.load(
            partial_scores + token_idx * partial_score_stride_t + block_offsets,
            mask=valid_blocks,
            other=-float("inf"),
        )
        _, best_block = tl.max(
            scores,
            axis=0,
            return_indices=True,
            return_indices_tie_break_left=True,
        )
        token = tl.load(
            partial_tokens + token_idx * partial_token_stride_t + best_block
        )
        tl.store(output_tokens + token_idx * output_stride_t, token)

    _FUSED_SPARSE_ARGMAX_KERNELS = (
        triton,
        _sparse_argmax_blocks_kernel,
        _sparse_argmax_reduce_kernel,
    )
    return _FUSED_SPARSE_ARGMAX_KERNELS


def _fallback_sparse_argmax(
    self: Any,
    original_get_top_tokens: Any,
    hidden_states: Any,
    lm_head_weight: Any,
    reason: Exception,
) -> Any:
    if FUSED_SPARSE_ARGMAX_REQUIRE:
        raise RuntimeError(
            "FUSED_SPARSE_ARGMAX_REQUIRE=1 but fusion failed"
        ) from reason
    if not getattr(self, "_pupa_fused_sparse_argmax_warned", False):
        self._pupa_fused_sparse_argmax_warned = True
        print(
            f"[pupa-fused-sparse-argmax] falling back to PyTorch path: {reason!r}",
            file=sys.stderr,
            flush=True,
        )
    return original_get_top_tokens(self, hidden_states, lm_head_weight)


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

    embedder_cls = module.Gemma4MTPMaskedEmbedder
    original_get_top_tokens = embedder_cls.get_top_tokens

    def _select_and_score_unsorted(self: Any, hidden_states: Any, lm_head_weight: Any):
        num_tokens = hidden_states.shape[0]
        _, top_k_indices = torch.topk(
            self.centroids(hidden_states),
            k=self.centroid_intermediate_top_k,
            dim=-1,
            sorted=False,
        )
        clusters = self.token_ordering.view(
            self.num_centroids,
            self.vocab_size_per_centroid,
        )
        selected = clusters[top_k_indices]
        embeddings = lm_head_weight[selected.reshape(-1)].view(
            num_tokens,
            self.num_selected,
            self.hidden_size,
        )
        logits = torch.einsum("td,tsd->ts", hidden_states, embeddings)
        return logits, selected.view(num_tokens, -1)

    def _record_draft_topk(
        self: Any,
        hidden_states: Any,
        lm_head_weight: Any,
    ) -> None:
        if not SPEC_DRAFT_TOPK_PROBE:
            return
        row = getattr(self, "_ff_draft_topk_row", None)
        if row is None:
            return
        try:
            logits, selected = _select_and_score_unsorted(
                self,
                hidden_states,
                lm_head_weight,
            )
            k = min(SPEC_DRAFT_TOPK_K, int(logits.shape[-1]), int(row.shape[-1]))
            _, rank_ids = torch.topk(logits, k=k, dim=-1, sorted=True)
            token_ids = selected.long().gather(1, rank_ids.long())
            row[: token_ids.shape[0], :k].copy_(token_ids.to(row.dtype))
        except Exception as exc:
            if SPEC_DRAFT_TOPK_REQUIRE:
                raise RuntimeError("SPEC_DRAFT_TOPK_REQUIRE=1 but top-k failed") from exc
            if not getattr(self, "_ff_draft_topk_warned", False):
                self._ff_draft_topk_warned = True
                print(
                    f"[draft-topk] disabled after error: {exc!r}",
                    file=sys.stderr,
                    flush=True,
                )

    def get_top_tokens_fused(self: Any, hidden_states: Any, lm_head_weight: Any) -> Any:
        if not FUSED_SPARSE_ARGMAX:
            return original_get_top_tokens(self, hidden_states, lm_head_weight)
        try:
            if (
                hidden_states.device.type != "cuda"
                or lm_head_weight.device.type != "cuda"
            ):
                raise RuntimeError("fusion requires CUDA tensors")
            if (
                hidden_states.dtype != torch.bfloat16
                or lm_head_weight.dtype != torch.bfloat16
            ):
                raise RuntimeError(
                    "fusion currently preserves exact PyTorch argmax only for bf16"
                )
            hidden_size = int(self.hidden_size)
            if hidden_size <= 0 or hidden_size > 1024:
                raise RuntimeError(f"unsupported hidden_size={hidden_size}")

            _record_draft_topk(self, hidden_states, lm_head_weight)

            triton, blocks_kernel, reduce_kernel = _get_fused_sparse_argmax_kernels()
            num_tokens = int(hidden_states.shape[0])
            selected_count = int(self.num_selected)
            block_selected = _next_power_of_2(FUSED_SPARSE_ARGMAX_BLOCK)
            num_blocks = triton.cdiv(selected_count, block_selected)
            reduce_block = _next_power_of_2(num_blocks)
            block_d = _next_power_of_2(hidden_size)

            _, top_k_indices = torch.topk(
                self.centroids(hidden_states),
                k=self.centroid_intermediate_top_k,
                dim=-1,
                sorted=False,
            )
            partial_scores = torch.empty(
                (num_tokens, num_blocks),
                dtype=torch.float32,
                device=hidden_states.device,
            )
            partial_tokens = torch.empty(
                (num_tokens, num_blocks),
                dtype=torch.int64,
                device=hidden_states.device,
            )
            output_tokens = torch.empty(
                (num_tokens,),
                dtype=torch.int64,
                device=hidden_states.device,
            )

            blocks_kernel[(num_tokens, num_blocks)](
                hidden_states,
                lm_head_weight,
                top_k_indices,
                self.token_ordering,
                partial_scores,
                partial_tokens,
                hidden_states.stride(0),
                hidden_states.stride(1),
                lm_head_weight.stride(0),
                lm_head_weight.stride(1),
                top_k_indices.stride(0),
                top_k_indices.stride(1),
                partial_scores.stride(0),
                partial_tokens.stride(0),
                VOCAB_PER_CENTROID=int(self.vocab_size_per_centroid),
                SELECTED_COUNT=selected_count,
                HIDDEN_SIZE=hidden_size,
                BLOCK_SELECTED=block_selected,
                BLOCK_D=block_d,
                num_warps=8,
            )
            reduce_kernel[(num_tokens,)](
                partial_scores,
                partial_tokens,
                output_tokens,
                partial_scores.stride(0),
                partial_tokens.stride(0),
                output_tokens.stride(0),
                NUM_BLOCKS=num_blocks,
                BLOCK_BLOCKS=reduce_block,
                num_warps=8,
            )
            return output_tokens
        except Exception as exc:
            return _fallback_sparse_argmax(
                self,
                original_get_top_tokens,
                hidden_states,
                lm_head_weight,
                exc,
            )

    embedder_cls._select_and_score = _select_and_score_unsorted
    embedder_cls.get_top_tokens = get_top_tokens_fused
    print(
        f"[pupa-fused-sparse-argmax] patched Gemma4MTPMaskedEmbedder top-token path "
        f"in pid {os.getpid()} (enabled={FUSED_SPARSE_ARGMAX}, "
        f"require={FUSED_SPARSE_ARGMAX_REQUIRE}, block={FUSED_SPARSE_ARGMAX_BLOCK})",
        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


def _apply_tree_attention_patch(module: Any) -> None:
    """Monkey-patch Gemma4 attention to support Star Tree masking."""
    import torch
    import functools
    
    # Gemma4 typically uses Gemma4Attention.
    # We use signature-neutral wrapper (*args, **kwargs) to satisfy Dynamo.
    for name in dir(module):
        attr = getattr(module, name)
        if (isinstance(attr, type) and "Attention" in name 
            and hasattr(attr, "forward") and not name.startswith("_")
            and attr.__module__ == module.__name__):
            
            original_forward = attr.forward
            
            @functools.wraps(original_forward)
            def forward(self: Any, *args: Any, **kwargs: Any) -> Any:
                # attn_metadata is usually a keyword or positional argument.
                metadata = kwargs.get("attn_metadata")
                if metadata is None:
                    for arg in args:
                        if hasattr(arg, "num_prefills"):
                            metadata = arg
                            break
                
                if (SPEC_TREE_SPEC and metadata is not None 
                    and getattr(metadata, "num_prefills", 1) == 0):
                    num_tokens = getattr(metadata, "num_tokens", -1)
                    # We expect 1 (base) + K*W (tree)
                    K = int(os.environ.get("SPECULATIVE_TOKENS", "7"))
                    W = SPEC_TREE_WIDTH
                    if num_tokens == 1 + K * W:
                        mask = _get_tree_mask(K, W, device=self.device)
                        if hasattr(metadata, "attn_bias"):
                            # vLLM V1: 0 for allow, -inf for block.
                            bias = torch.zeros_like(mask, dtype=torch.float32)
                            bias.masked_fill_(~mask, float("-inf"))
                            metadata.attn_bias = bias
                
                return original_forward(self, *args, **kwargs)
            
            attr.forward = forward
            print(f"[tree-spec] patched {name}.forward for Star Tree verification", 
                  file=sys.stderr, flush=True)

sys.meta_path.insert(0, _TargetFinder(TOP_TOKEN_TARGET, _apply_fused_top_token_patch))
sys.meta_path.insert(0, _TargetFinder(PROPOSER_TARGET, _apply_samplerprep_proposer_patch))
sys.meta_path.insert(0, _TargetFinder(REJECTION_TARGET, _apply_fused_greedy_rejection_prep_patch))
sys.meta_path.insert(0, _TargetFinder(RUNNER_TARGET, _apply_loopgraph_copy_event_patch))
sys.meta_path.insert(0, _TargetFinder(LOOPGRAPH_TARGET, _apply_loopgraph_patch))
# tree-v2: attn-bias registration removed (FA has no attn_bias; star
# attention is installed by tree_v2_ext at the Attention layer instead)

# PCK-04: registers Gemma4ForCausalLM lm_head rebuild + logits scatter.
import serve_patch_pck04  # noqa: E402, F401

# tree-v2 machinery (kernels + runner/attention patches)
import tree_v2_ext  # noqa: E402, F401
import steptime_patch  # noqa: E402, F401  (credit: agent-smith)