| |
| |
| """Narrow draft-only FP8 LM-head selection and runtime receipt helpers. |
| |
| This module is an experiment overlay. It deliberately does not monkey-patch |
| vLLM at import time. The copied Gemma4 MTP constructor must call |
| ``select_draft_head_method`` explicitly, and only when the experiment env flag |
| is enabled. That keeps the target model and ordinary token embeddings out of |
| the policy. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import json |
| import math |
| import os |
| from dataclasses import dataclass |
| from pathlib import Path |
| from typing import Any, Mapping |
|
|
|
|
| EXPERIMENT_ENV = "VLLM_EXPERIMENTAL_DRAFT_HEAD_FP8" |
| EXPECTED_ONLINE_SPEC = "fp8_per_tensor_static" |
| EXPECTED_METHOD = "Fp8PerTensorOnlineLinearMethod" |
|
|
|
|
| class HeadFp8EligibilityError(ValueError): |
| """Raised when the explicitly requested head experiment is unsafe.""" |
|
|
|
|
| @dataclass(frozen=True) |
| class Eligibility: |
| enabled: bool |
| eligible: bool |
| reason: str |
| role: str |
| prefix: str |
| ordered_embeddings: bool |
| tie_word_embeddings: bool |
| quant_spec: str | None |
|
|
|
|
| def env_enabled(environ: Mapping[str, str] | None = None) -> bool: |
| env = os.environ if environ is None else environ |
| return env.get(EXPERIMENT_ENV, "").strip().lower() in {"1", "true", "yes"} |
|
|
|
|
| def _online_spec(quant_config: Any) -> str | None: |
| |
| |
| |
| |
| online = getattr(quant_config, "online_quantization_config", None) |
| if online is None and hasattr(quant_config, "args"): |
| online = quant_config |
| if online is None: |
| return None |
| args = getattr(online, "args", online) |
| linear = getattr(args, "linear", None) |
| if linear is None: |
| return None |
| |
| |
| |
| return str(linear) |
|
|
|
|
| def _is_head_instance( |
| layer: Any, |
| parallel_lm_head_type: type[Any] | None, |
| parallel_lm_head_class: type[Any] | None = None, |
| ) -> bool: |
| if layer is None: |
| if parallel_lm_head_class is None: |
| return False |
| if parallel_lm_head_type is not None: |
| try: |
| return issubclass(parallel_lm_head_class, parallel_lm_head_type) |
| except TypeError: |
| return False |
| return ( |
| parallel_lm_head_class.__name__ == "ParallelLMHead" |
| and parallel_lm_head_class.__module__.endswith( |
| "vocab_parallel_embedding" |
| ) |
| ) |
| if parallel_lm_head_type is not None: |
| return isinstance(layer, parallel_lm_head_type) |
| try: |
| from vllm.model_executor.layers.vocab_parallel_embedding import ( |
| ParallelLMHead, |
| ) |
|
|
| return isinstance(layer, ParallelLMHead) |
| except ImportError: |
| return ( |
| type(layer).__name__ == "ParallelLMHead" |
| and type(layer).__module__.endswith("vocab_parallel_embedding") |
| ) |
|
|
|
|
| def eligibility( |
| *, |
| layer: Any, |
| quant_config: Any, |
| role: str, |
| prefix: str, |
| enabled: bool, |
| ordered_embeddings: bool, |
| tie_word_embeddings: bool, |
| head_dtype: Any = None, |
| model_dtype: Any = None, |
| parallel_lm_head_type: type[Any] | None = None, |
| parallel_lm_head_class: type[Any] | None = None, |
| ) -> Eligibility: |
| spec = _online_spec(quant_config) |
| common = dict( |
| role=role, |
| prefix=prefix, |
| ordered_embeddings=ordered_embeddings, |
| tie_word_embeddings=tie_word_embeddings, |
| quant_spec=spec, |
| ) |
| if not enabled: |
| return Eligibility( |
| enabled=enabled, eligible=False, reason="experiment flag is disabled", **common |
| ) |
| if role != "draft": |
| return Eligibility(enabled=enabled, eligible=False, reason="role is not draft", **common) |
| if not _is_head_instance( |
| layer, parallel_lm_head_type, parallel_lm_head_class |
| ): |
| return Eligibility( |
| enabled=enabled, eligible=False, reason="layer is not ParallelLMHead", **common |
| ) |
| if prefix.rsplit(".", 1)[-1] != "lm_head": |
| return Eligibility( |
| enabled=enabled, eligible=False, reason="prefix is not an LM head", **common |
| ) |
| if spec != EXPECTED_ONLINE_SPEC: |
| return Eligibility( |
| enabled=enabled, |
| eligible=False, |
| reason=f"online spec {spec!r} is not {EXPECTED_ONLINE_SPEC!r}", |
| **common, |
| ) |
| if ordered_embeddings: |
| return Eligibility( |
| enabled=enabled, |
| eligible=False, |
| reason="ordered embeddings bypass LogitsProcessor.quant_method.apply", |
| **common, |
| ) |
| if not tie_word_embeddings: |
| return Eligibility( |
| enabled=enabled, |
| eligible=False, |
| reason="untied draft embedding/head is outside this alias-safe prototype", |
| **common, |
| ) |
| if head_dtype is not None and model_dtype is not None: |
| if str(head_dtype) != str(model_dtype): |
| return Eligibility( |
| enabled=enabled, |
| eligible=False, |
| reason="custom FP8 head requires head dtype aligned with hidden dtype", |
| **common, |
| ) |
| return Eligibility(enabled=enabled, eligible=True, reason="eligible", **common) |
|
|
|
|
| def select_draft_head_method( |
| *, |
| layer: Any, |
| quant_config: Any, |
| role: str, |
| prefix: str, |
| enabled: bool, |
| ordered_embeddings: bool, |
| tie_word_embeddings: bool, |
| head_dtype: Any = None, |
| model_dtype: Any = None, |
| parallel_lm_head_type: type[Any] | None = None, |
| parallel_lm_head_class: type[Any] | None = None, |
| method_factory: Any | None = None, |
| ) -> tuple[Any | None, Eligibility]: |
| """Return the pinned online FP8 method only for an eligible draft head. |
| |
| The caller passes the returned method as ``quant_method`` to the copied |
| ``ParallelLMHead`` constructor. Returning ``None`` means use the normal |
| resolver, preserving target and embedding behavior. |
| """ |
|
|
| decision = eligibility( |
| layer=layer, |
| quant_config=quant_config, |
| role=role, |
| prefix=prefix, |
| enabled=enabled, |
| ordered_embeddings=ordered_embeddings, |
| tie_word_embeddings=tie_word_embeddings, |
| head_dtype=head_dtype, |
| model_dtype=model_dtype, |
| parallel_lm_head_type=parallel_lm_head_type, |
| parallel_lm_head_class=parallel_lm_head_class, |
| ) |
| if not decision.eligible: |
| if enabled and role == "draft" and _is_head_instance( |
| layer, parallel_lm_head_type, parallel_lm_head_class |
| ): |
| raise HeadFp8EligibilityError(decision.reason) |
| return None, decision |
|
|
| if method_factory is None: |
| from vllm.model_executor.layers.quantization.online.fp8 import ( |
| Fp8PerTensorOnlineLinearMethod, |
| ) |
|
|
| method_factory = Fp8PerTensorOnlineLinearMethod |
| return method_factory(), decision |
|
|
|
|
| def _bounded_finite(value: Any, chunk_elements: int = 8 * 1024 * 1024) -> bool | None: |
| """Check a tensor in bounded chunks, avoiding a full FP32 copy.""" |
|
|
| if hasattr(value, "finite"): |
| return bool(value.finite) |
| try: |
| import torch |
|
|
| flat = value.reshape(-1) |
| for start in range(0, int(flat.numel()), chunk_elements): |
| chunk = flat[start : start + chunk_elements] |
| chunk = chunk.to(torch.float32) if hasattr(chunk, "to") else chunk |
| if not bool(torch.isfinite(chunk).all().item()): |
| return False |
| return True |
| except Exception: |
| return None |
|
|
|
|
| def _positive_scalar(value: Any) -> bool | None: |
| if hasattr(value, "positive"): |
| return bool(value.positive) |
| try: |
| if int(value.numel()) != 1: |
| return False |
| scalar = float(value.detach().to("cpu", dtype=getattr(value, "dtype", None)).item()) |
| return math.isfinite(scalar) and scalar > 0.0 |
| except Exception: |
| return None |
|
|
|
|
| def _tensor_snapshot( |
| value: Any, |
| *, |
| check_finite: bool = False, |
| require_positive_scalar: bool = False, |
| ) -> dict[str, Any] | None: |
| if value is None: |
| return None |
| shape = getattr(value, "shape", None) |
| dtype = getattr(value, "dtype", None) |
| device = getattr(value, "device", None) |
| result: dict[str, Any] = { |
| "dtype": str(dtype) if dtype is not None else None, |
| "shape": list(shape) if shape is not None else None, |
| "device": str(device) if device is not None else None, |
| } |
| if hasattr(value, "numel"): |
| result["numel"] = int(value.numel()) |
| if check_finite: |
| result["finite"] = _bounded_finite(value) |
| if require_positive_scalar: |
| result["positive_scalar"] = _positive_scalar(value) |
| if hasattr(value, "data_ptr"): |
| try: |
| result["data_ptr"] = int(value.data_ptr()) |
| except Exception: |
| pass |
| return result |
|
|
|
|
| def _method_name(method: Any) -> str | None: |
| return None if method is None else type(method).__name__ |
|
|
|
|
| def _draft_text_config(draft_model: Any) -> Any: |
| """Return the draft text config without importing the runtime model class.""" |
|
|
| config = getattr(draft_model, "config", None) |
| return getattr(config, "text_config", config) |
|
|
|
|
| def _draft_dimension(draft_model: Any, name: str) -> int | None: |
| value = getattr(_draft_text_config(draft_model), name, None) |
| if value is None: |
| return None |
| try: |
| value = int(value) |
| except (TypeError, ValueError): |
| return None |
| return value if value > 0 else None |
|
|
|
|
| def make_runtime_receipt( |
| *, |
| head: Any, |
| draft_model: Any, |
| target_model: Any, |
| decision: Eligibility, |
| phase: str, |
| target_lm_head: Any | None = None, |
| ) -> dict[str, Any]: |
| """Capture metadata only; never serializes weight values.""" |
|
|
| draft_body = getattr(draft_model, "model", draft_model) |
| draft_embed = getattr(draft_body, "embed_tokens", None) |
| target_language_model = ( |
| target_model.get_language_model() |
| if hasattr(target_model, "get_language_model") |
| else target_model |
| ) |
| if target_lm_head is None: |
| target_lm_head = getattr(target_language_model, "lm_head", None) |
| target_body = getattr(target_language_model, "model", target_language_model) |
| target_embed = getattr(target_body, "embed_tokens", None) |
| method = getattr(head, "quant_method", None) |
| kernel = getattr(method, "fp8_linear", None) |
| weight = getattr(head, "weight", None) |
| scale = getattr(head, "weight_scale", None) |
| target_head_weight = getattr(target_lm_head, "weight", None) |
|
|
| head_tp_size = getattr(head, "tp_size", None) |
| |
| |
| |
| |
| draft_hidden_size = _draft_dimension(draft_model, "hidden_size") |
| draft_vocab_size = _draft_dimension(draft_model, "vocab_size") |
| head_hidden_size = getattr(head, "input_size_per_partition", None) |
| head_vocab_size = getattr(head, "output_size_per_partition", None) |
| if head_hidden_size is not None: |
| head_hidden_size = int(head_hidden_size) |
| if head_vocab_size is not None: |
| head_vocab_size = int(head_vocab_size) |
| expected_hidden_size = draft_hidden_size or head_hidden_size |
| expected_vocab_size = draft_vocab_size or head_vocab_size |
|
|
| return { |
| "schema": "draft-head-fp8-runtime-receipt/v2", |
| "phase": phase, |
| "eligible": decision.eligible, |
| "eligibility_reason": decision.reason, |
| "role": decision.role, |
| "prefix": decision.prefix, |
| "ordered_embeddings": decision.ordered_embeddings, |
| "tie_word_embeddings": decision.tie_word_embeddings, |
| "quant_spec": decision.quant_spec, |
| "head_class": f"{type(head).__module__}.{type(head).__name__}", |
| "head_method": _method_name(method), |
| "head_method_module": getattr(type(method), "__module__", None), |
| "head_kernel": None if kernel is None else type(kernel).__name__, |
| "head_tp_size": head_tp_size, |
| "draft_config_hidden_size": draft_hidden_size, |
| "draft_config_vocab_size": draft_vocab_size, |
| "head_partition_hidden_size": head_hidden_size, |
| "head_partition_vocab_size": head_vocab_size, |
| "head_expected_weight_shape": ( |
| [expected_hidden_size, expected_vocab_size] |
| if expected_hidden_size is not None and expected_vocab_size is not None |
| else None |
| ), |
| "head_weight": _tensor_snapshot(weight, check_finite=True), |
| "head_weight_scale": _tensor_snapshot( |
| scale, check_finite=True, require_positive_scalar=True |
| ), |
| |
| |
| "draft_embed": _tensor_snapshot(getattr(draft_embed, "weight", None)), |
| "target_embed": _tensor_snapshot(getattr(target_embed, "weight", None)), |
| "target_lm_head_weight": _tensor_snapshot(target_head_weight), |
| "head_weight_is_draft_embed": ( |
| weight is not None |
| and draft_embed is not None |
| and weight is getattr(draft_embed, "weight", None) |
| ), |
| "head_weight_is_target_embed": ( |
| weight is not None |
| and target_embed is not None |
| and weight is getattr(target_embed, "weight", None) |
| ), |
| "draft_embed_weight_is_target_embed": ( |
| draft_embed is not None |
| and target_embed is not None |
| and getattr(draft_embed, "weight", None) |
| is getattr(target_embed, "weight", None) |
| ), |
| "draft_embedding_method": _method_name( |
| getattr(draft_embed, "quant_method", None) |
| ), |
| "target_embedding_method": _method_name( |
| getattr(target_embed, "quant_method", None) |
| ), |
| "target_lm_head_method": _method_name( |
| getattr(target_lm_head, "quant_method", None) |
| ), |
| } |
|
|
|
|
| def validate_post_share_receipt(receipt: Mapping[str, Any]) -> None: |
| """Fail closed on the facts required before any API run.""" |
|
|
| required = ( |
| "eligible", |
| "head_class", |
| "head_method", |
| "head_kernel", |
| "head_weight", |
| "head_weight_scale", |
| "draft_embed", |
| "target_embed", |
| "head_weight_is_draft_embed", |
| "head_weight_is_target_embed", |
| "draft_embed_weight_is_target_embed", |
| "head_tp_size", |
| "head_expected_weight_shape", |
| "target_lm_head_weight", |
| "target_lm_head_method", |
| ) |
| missing = [key for key in required if key not in receipt] |
| if missing: |
| raise HeadFp8EligibilityError(f"receipt missing {missing}") |
| if receipt["eligible"] is not True: |
| raise HeadFp8EligibilityError("receipt is not eligible") |
| if receipt["head_tp_size"] != 1: |
| raise HeadFp8EligibilityError("only TP1 is supported by this prototype") |
| if receipt["head_expected_weight_shape"] is None: |
| raise HeadFp8EligibilityError("draft config did not provide head dimensions") |
| if receipt.get("draft_config_hidden_size") is None or receipt.get( |
| "draft_config_vocab_size" |
| ) is None: |
| raise HeadFp8EligibilityError("draft config dimensions are missing") |
| if receipt.get("head_partition_hidden_size") != receipt["draft_config_hidden_size"]: |
| raise HeadFp8EligibilityError("head hidden partition does not match draft config") |
| if receipt.get("head_partition_vocab_size") != receipt["draft_config_vocab_size"]: |
| raise HeadFp8EligibilityError("head vocab partition does not match draft config") |
| if not str(receipt["head_class"]).endswith(".ParallelLMHead"): |
| raise HeadFp8EligibilityError("unexpected head class") |
| if receipt["head_method"] != EXPECTED_METHOD: |
| raise HeadFp8EligibilityError("unexpected online head method") |
| if not receipt["head_kernel"]: |
| raise HeadFp8EligibilityError("kernel selection was not recorded") |
| weight = receipt["head_weight"] |
| scale = receipt["head_weight_scale"] |
| if not weight or not scale or "dtype" not in weight or "dtype" not in scale: |
| raise HeadFp8EligibilityError("head FP8 weight/scale metadata is incomplete") |
| if weight.get("finite") is not True or scale.get("finite") is not True: |
| raise HeadFp8EligibilityError("head FP8 weight/scale is not finite") |
| if scale.get("shape") != [1] or scale.get("positive_scalar") is not True: |
| raise HeadFp8EligibilityError("head scale is not one finite positive scalar") |
| if weight.get("shape") != receipt["head_expected_weight_shape"]: |
| raise HeadFp8EligibilityError( |
| "processed head shape does not match [hidden, vocab] logical shape" |
| ) |
| if "float8" not in str(weight["dtype"]).lower(): |
| raise HeadFp8EligibilityError("head weight is not an FP8 dtype") |
| if receipt["head_weight_is_draft_embed"] is not False: |
| raise HeadFp8EligibilityError("head still aliases draft input embedding") |
| if receipt["head_weight_is_target_embed"] is not False: |
| raise HeadFp8EligibilityError("head aliases target input embedding") |
| if receipt["draft_embed_weight_is_target_embed"] is not True: |
| raise HeadFp8EligibilityError("draft input embedding was not shared with target") |
| if "bfloat16" not in str(receipt["draft_embed"].get("dtype", "")).lower(): |
| raise HeadFp8EligibilityError("draft input embedding is not BF16") |
| target_method = str(receipt.get("target_lm_head_method", "")) |
| if "Fp8PerTensorOnlineLinearMethod" in target_method: |
| raise HeadFp8EligibilityError("target LM head was routed to draft FP8 method") |
| target_weight = receipt["target_lm_head_weight"] |
| if "bfloat16" not in str(target_weight.get("dtype", "")).lower(): |
| raise HeadFp8EligibilityError("target LM head is not BF16") |
| if "Unquantized" not in target_method: |
| raise HeadFp8EligibilityError("target LM head is not on an unquantized method") |
|
|
|
|
| def write_runtime_receipt(receipt: Mapping[str, Any], path: str | Path) -> None: |
| output = Path(path) |
| output.parent.mkdir(parents=True, exist_ok=True) |
| if output.exists(): |
| raise FileExistsError(output) |
| output.write_text(json.dumps(receipt, indent=2, sort_keys=True) + "\n") |
|
|
