vllm.v1.worker.gpu.spec_decode.eagle.utils

Functions:

• get_target_lm_head –

  The target's lm_head — from get_language_model() for

• maybe_share_target_embed –

  Share the target input embedding with the drafter when needed.

• share_target_lm_head –

  Install the target lm_head, or its quantized copy, on a drafter that

_should_share(eagle, flag, draft, target)

Share unless the draft declares its own copy that differs from the target.

A draft that declares its own copy but has no top-level one (e.g. MTP heads stored per layer) keeps it.

Source code in vllm/v1/worker/gpu/spec_decode/eagle/utils.py

def _should_share(eagle: nn.Module, flag: str, draft, target) -> bool:
    """Share unless the draft declares its own copy that differs from the target.

    A draft that declares its own copy but has no top-level one (e.g. MTP
    heads stored per layer) keeps it.
    """
    if not getattr(eagle, flag, False):
        return True
    if draft is None or target is None:
        return False
    # A quantized layer registers packed weights (`weight_packed`, `qweight`,
    # ...) instead of `weight`, so the two copies cannot be compared. Keep the
    # draft's own copy rather than assuming it matches the target.
    w = getattr(draft, "weight", None)
    target_w = getattr(target, "weight", None)
    if not isinstance(w, torch.Tensor) or not isinstance(target_w, torch.Tensor):
        return False
    # torch.equal on GPU allocates a bool mask the size of the input.
    # Use the faster GPU path when there is plenty of headroom;
    # otherwise compare on CPU.
    if w.is_cuda and torch.accelerator.get_memory_info(w.device)[0] < w.numel() * 2:
        return torch.equal(w.cpu(), target_w.cpu())
    return torch.equal(w, target_w)

get_target_lm_head(target_model, target_language_model)

The target's lm_head — from get_language_model() for *ForConditionalGeneration targets, else the top-level module.

Source code in vllm/v1/worker/gpu/spec_decode/eagle/utils.py

def get_target_lm_head(target_model: nn.Module, target_language_model: nn.Module):
    """The target's lm_head — from get_language_model() for
    *ForConditionalGeneration targets, else the top-level module."""
    return getattr(target_language_model, "lm_head", None) or getattr(
        target_model, "lm_head", None
    )

maybe_share_target_embed(draft_model, draft_inner, target_inner)

Share the target input embedding with the drafter when needed.

Source code in vllm/v1/worker/gpu/spec_decode/eagle/utils.py

def maybe_share_target_embed(
    draft_model: nn.Module, draft_inner: nn.Module, target_inner: nn.Module
) -> None:
    """Share the target input embedding with the drafter when needed."""
    target_embed = getattr(target_inner, "embed_tokens", None) or getattr(
        target_inner, "embedding", None
    )
    if isinstance(target_embed, PPMissingLayer):
        target_embed = None
    # The drafter does not use the target's LoRA adapter.
    if isinstance(target_embed, BaseLayerWithLoRA):
        target_embed = target_embed.base_layer
    draft_embed = getattr(draft_inner, "embed_tokens", None)

    if get_pp_group().world_size > 1 and not hasattr(
        draft_model, "has_own_embed_tokens"
    ):
        return

    if target_embed is None:
        if hasattr(draft_inner, "embed_tokens") and not getattr(
            draft_model, "has_own_embed_tokens", False
        ):
            raise RuntimeError(
                f"{type(draft_model).__name__} needs the target input embedding, "
                "but it is unavailable on this PP stage"
            )
        return

    if _should_share(draft_model, "has_own_embed_tokens", draft_embed, target_embed):
        if draft_embed is not None:
            del draft_inner.embed_tokens
        draft_inner.embed_tokens = target_embed

share_target_lm_head(draft_model, target_lm_head, quantization, can_share=True)

Install the target lm_head, or its quantized copy, on a drafter that shares it. Returns the installed head, or None if the drafter keeps its own.

Source code in vllm/v1/worker/gpu/spec_decode/eagle/utils.py

def share_target_lm_head(
    draft_model: nn.Module,
    target_lm_head: nn.Module | None,
    quantization: Literal["fp8", "nvfp4"] | None,
    can_share: bool = True,
) -> nn.Module | None:
    """Install the target lm_head, or its quantized copy, on a drafter that
    shares it. Returns the installed head, or None if the drafter keeps its own.
    """
    draft_lm_head = getattr(draft_model, "lm_head", None)
    if not (
        can_share
        and target_lm_head is not None
        and _should_share(draft_model, "has_own_lm_head", draft_lm_head, target_lm_head)
    ):
        if quantization is not None:
            raise ValueError(
                "draft_lm_head_quantization needs a drafter that shares the "
                "target lm_head."
            )
        return None
    head = target_lm_head
    if quantization is not None:
        head = quantized_lm_head_copy(target_lm_head, quantization)
    if draft_lm_head is not None:
        del draft_model.lm_head
    draft_model.lm_head = head
    return head