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.

_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
    # 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.
    w = draft.weight
    if w.is_cuda and torch.accelerator.get_memory_info(w.device)[0] < w.numel() * 2:
        return torch.equal(w.cpu(), target.weight.cpu())
    return torch.equal(w, target.weight)

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