# Modification notice: DevelopingDad, 2026-10-08. Experimental isolated FP8 draft-head integration and/or byte-preserving per-boot diagnostic archival. # SPDX-License-Identifier: Apache-2.0. Upstream copyright and SPDX notices retained below. """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: # The copied v0.31 resolver receives the outer QuantizationConfig and # reads .online_quantization_config.args. Accepting a direct # OnlineQuantizationConfig here is useful for the constructor seam, but # does not broaden the accepted spec. 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 # Pinned v0.31 QuantSpec.__str__ names the CLI shorthand as # fp8_per_tensor_static. Keep this exact; block/channel FP8 are outside # this experiment's contract. 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) # The processed FP8 weight is transposed to [hidden, vocab]. Derive the # expected logical shape from the draft HF config, then retain the layer # partition metadata for an independent TP1 cross-check. Do not bake the # current 1024x262144 dimensions into this helper. 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 ), # Embeddings are deliberately metadata-only. A whole 262144x1024 # FP32 finiteness conversion would consume about 1 GiB per tensor. "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")