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

def _should_share(eagle: nn.Module, flag: str, draft, target) -> bool:
    """Share when the draft has no own copy, or its copy matches the target."""
    if not getattr(eagle, flag, False) or draft is None:
        return True
    if 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)

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
    )

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