diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 921f31466..dcf35614d 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -1718,6 +1718,10 @@ class EngineArgs: ) self.speculative_config[key] = value + if self.speculative_config is None: + self.speculative_config = _checkpoint_mtp_speculative_config( + target_model_config + ) if self.speculative_config is None: return None @@ -2692,3 +2696,29 @@ def _raise_unsupported_error(feature_name: str): f"remove {feature_name} from your config." ) raise NotImplementedError(msg) + + +# Checkpoints whose text config carries a bundled MTP drafter that vLLM can load +# from the same weights (method "mtp", draft = target model). +_CHECKPOINT_MTP_MODEL_TYPES: frozenset[str] = frozenset({"qwen3_5", "qwen3_5_moe"}) +_CHECKPOINT_MTP_NUM_SPECULATIVE_TOKENS: int = 5 + + +def _checkpoint_mtp_speculative_config(model_config: ModelConfig) -> dict | None: + """Default speculative config from the checkpoint's own MTP head, or None. + + Only applies when the user gave no speculative config at all. The drafter + weights live in the target checkpoint, so nothing extra is downloaded. + """ + hf_config = getattr(model_config, "hf_config", None) + model_type = getattr(hf_config, "model_type", None) + if model_type not in _CHECKPOINT_MTP_MODEL_TYPES: + return None + text_config = getattr(model_config, "hf_text_config", None) + n_mtp_layers = getattr(text_config, "mtp_num_hidden_layers", None) + if not isinstance(n_mtp_layers, int) or n_mtp_layers < 1: + return None + return { + "method": "mtp", + "num_speculative_tokens": _CHECKPOINT_MTP_NUM_SPECULATIVE_TOKENS, + } diff --git a/vllm/entrypoints/openai/completion/serving.py b/vllm/entrypoints/openai/completion/serving.py index fef174135..0814cdd71 100644 --- a/vllm/entrypoints/openai/completion/serving.py +++ b/vllm/entrypoints/openai/completion/serving.py @@ -395,46 +395,73 @@ class OpenAIServingCompletion(OpenAIServing): self._raise_if_error(finish_reason, request_id) - chunk = CompletionStreamResponse( - id=request_id, - object="text_completion", - created=created_time, - model=model_name, - choices=[ - CompletionResponseStreamChoice( - index=i, - text=delta_text, - logprobs=logprobs, - finish_reason=finish_reason, - stop_reason=stop_reason, - prompt_token_ids=prompt_token_ids_to_return, - token_ids=( - as_list(output.token_ids) - if request.return_token_ids - else None - ), - ) - ], - ) - # Stamp on terminal chunk only when no trailing usage chunk - # will follow (that one is the true final message). + # One choice chunk per generated token. A RequestOutput can + # carry several new tokens (speculative decoding accepts a + # step's drafts at once; the frontend merges outputs when it + # lags the engine). Emitting them as one chunk hides the + # token granularity from streaming consumers, so split the + # delta into per-token slices. Text concatenation, token + # accounting and finish_reason placement are unchanged. + delta_ids = as_list(output.token_ids) if ( - not include_usage - and self.system_fingerprint is not None - and finish_reason is not None + len(delta_ids) > 1 + and logprobs is None + and not (request.echo and prompt_token_ids_to_return) ): - chunk.system_fingerprint = self.system_fingerprint - if include_continuous_usage: - prompt_tokens = num_prompt_tokens[prompt_idx] - completion_tokens = previous_num_tokens[i] - chunk.usage = UsageInfo( - prompt_tokens=prompt_tokens, - completion_tokens=completion_tokens, - total_tokens=prompt_tokens + completion_tokens, + pieces = _split_delta_by_token( + tokenizer, delta_ids, delta_text ) + else: + pieces = [(delta_text, delta_ids)] + + tokens_before = previous_num_tokens[i] - len(delta_ids) + n_pieces = len(pieces) + for piece_idx, (piece_text, piece_ids) in enumerate(pieces): + is_last = piece_idx == n_pieces - 1 + tokens_before += len(piece_ids) + chunk = CompletionStreamResponse( + id=request_id, + object="text_completion", + created=created_time, + model=model_name, + choices=[ + CompletionResponseStreamChoice( + index=i, + text=piece_text, + logprobs=logprobs, + finish_reason=finish_reason if is_last else None, + stop_reason=stop_reason if is_last else None, + prompt_token_ids=( + prompt_token_ids_to_return + if piece_idx == 0 + else None + ), + token_ids=( + piece_ids if request.return_token_ids else None + ), + ) + ], + ) + # Stamp on terminal chunk only when no trailing usage + # chunk will follow (that one is the true final message). + if ( + is_last + and not include_usage + and self.system_fingerprint is not None + and finish_reason is not None + ): + chunk.system_fingerprint = self.system_fingerprint + if include_continuous_usage: + prompt_tokens = num_prompt_tokens[prompt_idx] + completion_tokens = tokens_before + chunk.usage = UsageInfo( + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + total_tokens=prompt_tokens + completion_tokens, + ) - response_json = chunk.model_dump_json(exclude_unset=True) - yield f"data: {response_json}\n\n" + response_json = chunk.model_dump_json(exclude_unset=True) + yield f"data: {response_json}\n\n" total_prompt_tokens = sum(num_prompt_tokens) total_completion_tokens = sum(previous_num_tokens) @@ -689,3 +716,44 @@ class OpenAIServingCompletion(OpenAIServing): tokens=out_tokens, top_logprobs=out_top_logprobs, ) + + +def _split_delta_by_token( + tokenizer: TokenizerLike | None, + token_ids: list[int], + text: str, +) -> list[tuple[str, list[int]]]: + """Slice a multi-token delta into per-token (text, [token_id]) pieces. + + Boundaries come from incrementally decoding the delta's token prefix; they + are clamped and made monotone so the pieces always concatenate to exactly + ``text``. The final piece absorbs any remainder, so nothing is lost or + duplicated even when a token boundary does not fall on a character + boundary (multi-byte sequences, leading-space normalisation). + """ + n = len(token_ids) + if n <= 1: + return [(text, list(token_ids))] + if not text: + # A multi-token delta can still decode to nothing while a multi-byte + # character is incomplete. Keep one piece per token anyway: collapsing + # them would hide tokens from consumers that count streamed chunks. + return [("", [token_id]) for token_id in token_ids] + bounds: list[int] = [] + if tokenizer is not None: + try: + for j in range(1, n): + prefix = tokenizer.decode(token_ids[:j], skip_special_tokens=True) + bounds.append(len(prefix)) + except Exception: # noqa: BLE001 - fall back to an even split + bounds = [] + if len(bounds) != n - 1: + bounds = [(len(text) * j) // n for j in range(1, n)] + pieces: list[tuple[str, list[int]]] = [] + start = 0 + for j, b in enumerate(bounds): + b = min(max(b, start), len(text)) + pieces.append((text[start:b], [token_ids[j]])) + start = b + pieces.append((text[start:], [token_ids[n - 1]])) + return pieces diff --git a/vllm/model_executor/kernels/linear/__init__.py b/vllm/model_executor/kernels/linear/__init__.py index 4ac8d49cd..bca89ea4d 100644 --- a/vllm/model_executor/kernels/linear/__init__.py +++ b/vllm/model_executor/kernels/linear/__init__.py @@ -321,8 +321,10 @@ _POSSIBLE_FP8_BLOCK_KERNELS: dict[ PlatformEnum, list[type[Fp8BlockScaledMMLinearKernel | FP8ScaledMMLinearKernel]] ] = { PlatformEnum.CUDA: [ - FlashInferFp8DeepGEMMDynamicBlockScaledKernel, + # One kernel for every batch size: no M-dependent dispatch, and no + # TensorRT-LLM nvcc JIT on the cold-start path. DeepGemmFp8BlockScaledMMKernel, + FlashInferFp8DeepGEMMDynamicBlockScaledKernel, CutlassFp8BlockScaledMMKernel, MarlinFP8ScaledMMLinearKernel, TritonFp8BlockScaledMMKernel, diff --git a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py index 06bfe5c5d..62494a37e 100644 --- a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py +++ b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py @@ -1399,6 +1399,19 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention): ) num_decode_tokens = attn_metadata.num_decode_tokens + # `a`/`b` are laid out over the whole batch, but the spec kernel walks + # them with spec_query_start_loc, i.e. as if spec rows started at 0. + # That holds only for a spec-only batch; when a prefill or a plain + # decode shares the step, the gating inputs must be gathered the same + # way mixed_qkv_spec is, or spec rows read another request's gates. + if spec_sequence_masks is not None and ( + attn_metadata.num_prefills > 0 or attn_metadata.num_decodes > 0 + ): + a_spec = a.index_select(0, spec_token_indx) + b_spec = b.index_select(0, spec_token_indx) + else: + a_spec = a + b_spec = b if attn_metadata.num_prefills > 0: assert mixed_qkv_non_spec is not None, ( "mixed_qkv_non_spec must be provided for prefill path" @@ -1456,8 +1469,8 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention): core_attn_out_spec, last_recurrent_state = ( fused_sigmoid_gating_delta_rule_update( A_log=self.A_log, - a=a, - b=b, + a=a_spec, + b=b_spec, dt_bias=self.dt_bias, q=query_spec, k=key_spec, diff --git a/vllm/v1/sample/logits_processor/__init__.py b/vllm/v1/sample/logits_processor/__init__.py index 2cb89e1ea..3501c3f47 100644 --- a/vllm/v1/sample/logits_processor/__init__.py +++ b/vllm/v1/sample/logits_processor/__init__.py @@ -18,6 +18,7 @@ from vllm.v1.sample.logits_processor.builtin import ( LogitBiasLogitsProcessor, MinPLogitsProcessor, MinTokensLogitsProcessor, + RepeatLoopBreakerLogitsProcessor, process_dict_updates, ) from vllm.v1.sample.logits_processor.interface import ( @@ -50,6 +51,7 @@ BUILTIN_LOGITS_PROCESSORS: list[type[LogitsProcessor]] = [ MinTokensLogitsProcessor, LogitBiasLogitsProcessor, MinPLogitsProcessor, + RepeatLoopBreakerLogitsProcessor, ] @@ -204,8 +206,14 @@ def build_logitsprocs( logger.warning( "min_p and logit_bias parameters won't work with speculative decoding." ) + # Speculative decoding drops the parameterised processors, but the + # loop breaker takes no sampling parameter and guards against a + # degenerate repeat, which spec decode does not make less likely. return LogitsProcessors( - [MinTokensLogitsProcessor(vllm_config, device, is_pin_memory)] + [ + MinTokensLogitsProcessor(vllm_config, device, is_pin_memory), + RepeatLoopBreakerLogitsProcessor(vllm_config, device, is_pin_memory), + ] ) custom_logitsprocs_classes = _load_custom_logitsprocs(custom_logitsprocs) @@ -346,6 +354,7 @@ __all__ = [ "LogitBiasLogitsProcessor", "MinPLogitsProcessor", "MinTokensLogitsProcessor", + "RepeatLoopBreakerLogitsProcessor", "BatchUpdate", "BatchUpdateBuilder", "MoveDirectionality", diff --git a/vllm/v1/sample/logits_processor/builtin.py b/vllm/v1/sample/logits_processor/builtin.py index d7c944438..4a9bb493a 100644 --- a/vllm/v1/sample/logits_processor/builtin.py +++ b/vllm/v1/sample/logits_processor/builtin.py @@ -7,6 +7,7 @@ import numpy as np import torch from vllm import SamplingParams +from vllm.logger import init_logger from vllm.utils.torch_utils import async_tensor_h2d from vllm.v1.sample.logits_processor.interface import ( BatchUpdate, @@ -17,6 +18,8 @@ from vllm.v1.sample.logits_processor.interface import ( if TYPE_CHECKING: from vllm.config import VllmConfig +logger = init_logger(__name__) + T = TypeVar("T") @@ -286,6 +289,147 @@ class MinTokensLogitsProcessor(LogitsProcessor): return logits +class RepeatLoopBreakerLogitsProcessor(LogitsProcessor): + """Stop a greedy continuation that has collapsed into a repeat loop. + + Tracks, per request, how long the tail has been an exact copy of an earlier + stretch of the same output. Once that repeated run is long enough to + threaten the output as a whole, the token that would extend it is masked, + so decoding takes the next-best token and leaves the loop. + + The run is found by remembering where each ``_GRAM``-token window last + occurred: a hit gives the loop's true period directly, whatever its length, + which a fixed-period search does not. Work is proportional to tokens + generated. Until a request qualifies its logits are untouched, so ordinary + output is bit-identical to running without this. + """ + + # Tokens that must match exactly before a repeat is credited at all. + _GRAM = 32 + # Fire once the repeated run reaches this many tokens, and this share of + # the output. Well under the point where a loop dominates the answer. + _MIN_RUN = 64 + _RUN_FRACTION = 8 + + def __init__( + self, vllm_config: "VllmConfig", device: torch.device, is_pin_memory: bool + ): + self.device = device + # req_index -> [output ids, tokens scanned, {gram hash: index}, + # run period, run length] + self.req_state: dict[int, list] = {} + self.neg_inf = torch.tensor( + -float("inf"), dtype=torch.float32, device=self.device + ) + + def is_argmax_invariant(self) -> bool: + """Masking the looping token is the whole point, so it is not.""" + return False + + @staticmethod + def add_request( + params: SamplingParams, _: list[int] | None, output_tok_ids: list[int] + ) -> list | None: + return [output_tok_ids, 0, {}, 0, 0] + + def update_state(self, batch_update: BatchUpdate | None): + process_dict_updates(self.req_state, batch_update, self.add_request) + + @classmethod + def _looping_token(cls, state: list) -> int | None: + """The token that would extend a repeat loop, or None if not looping.""" + out, scanned, seen, period, run = state + total = len(out) + if total < scanned: + # Should not happen, but never trust stale bookkeeping. + seen.clear() + scanned = period = run = 0 + gram = cls._GRAM + for i in range(max(scanned, gram - 1), total): + window = out[i - gram + 1 : i + 1] + key = hash(tuple(window)) + prev = seen.get(key) + seen[key] = i + if prev is None or out[prev - gram + 1 : prev + 1] != window: + # No repeat here, so any run in progress has ended. + period = run = 0 + continue + distance = i - prev + if distance == period: + run += 1 + else: + # A fresh repeat: the matching window itself is the run so far. + period, run = distance, gram + state[1], state[3], state[4] = total, period, run + if period > 0 and run >= max(cls._MIN_RUN, total // cls._RUN_FRACTION): + # The tail copies what came `period` tokens earlier, so the next + # token of the loop is the one `period` back. + return out[-period] + return None + + def _masked_requests(self) -> tuple[list[int], list[int]]: + reqs: list[int] = [] + toks: list[int] = [] + for index, state in self.req_state.items(): + try: + tok = self._looping_token(state) + except Exception: # noqa: BLE001 - never fail a step over this + logger.warning_once("Repeat-loop detector disabled after an error.") + self.req_state.clear() + return [], [] + if tok is not None: + reqs.append(index) + toks.append(tok) + return reqs, toks + + def _mask(self, logits: torch.Tensor, rows: list[int], toks: list[int]) -> None: + logits.index_put_( + ( + async_tensor_h2d(rows, device=self.device, dtype=torch.int32), + async_tensor_h2d(toks, device=self.device, dtype=torch.int32), + ), + self.neg_inf, + ) + + def apply(self, logits: torch.Tensor) -> torch.Tensor: + if not self.req_state: + return logits + reqs, toks = self._masked_requests() + if reqs: + self._mask(logits, reqs, toks) + return logits + + def apply_with_spec_decode( + self, + logits: torch.Tensor, + num_draft_tokens: list[int], + ) -> torch.Tensor: + """Spec-decode version of apply(). + + Request ``i`` owns ``num_draft_tokens[i]`` consecutive rows, so the + looping token is masked at every position that request could accept. + """ + if not self.req_state: + return logits + reqs, toks = self._masked_requests() + if not reqs: + return logits + cumsum = np.concatenate( + [[0], np.cumsum(np.array(num_draft_tokens, dtype=np.int64))] + ) + rows: list[int] = [] + row_toks: list[int] = [] + for index, tok in zip(reqs, toks): + if index + 1 >= len(cumsum): + continue + for row in range(int(cumsum[index]), int(cumsum[index + 1])): + rows.append(row) + row_toks.append(tok) + if rows: + self._mask(logits, rows, row_toks) + return logits + + def process_dict_updates( req_entries: dict[int, T], batch_update: BatchUpdate | None, diff --git a/vllm/v1/sample/rejection_sampler.py b/vllm/v1/sample/rejection_sampler.py index 8b4d8c9dc..a19e769c0 100644 --- a/vllm/v1/sample/rejection_sampler.py +++ b/vllm/v1/sample/rejection_sampler.py @@ -13,7 +13,6 @@ import torch.nn as nn from vllm.logger import init_logger from vllm.triton_utils import tl, triton from vllm.v1.outputs import LogprobsLists, LogprobsTensors, SamplerOutput -from vllm.v1.sample.logits_processor.builtin import MinTokensLogitsProcessor from vllm.v1.sample.metadata import SamplingMetadata from vllm.v1.sample.ops.bad_words import apply_bad_words_with_drafts from vllm.v1.sample.ops.penalties import apply_all_penalties @@ -333,7 +332,9 @@ class RejectionSampler(nn.Module): ) for processor in sampling_metadata.logitsprocs.non_argmax_invariant: - if isinstance(processor, MinTokensLogitsProcessor): + # Any processor that knows how to expand itself across a request's + # draft positions, not just min-tokens. + if hasattr(processor, "apply_with_spec_decode"): logits = processor.apply_with_spec_decode( logits, metadata.num_draft_tokens ) diff --git a/vllm/v1/spec_decode/llm_base_proposer.py b/vllm/v1/spec_decode/llm_base_proposer.py index 9f46cbd24..5be1ed39f 100644 --- a/vllm/v1/spec_decode/llm_base_proposer.py +++ b/vllm/v1/spec_decode/llm_base_proposer.py @@ -57,6 +57,162 @@ from vllm.v1.worker.utils import AttentionGroup logger = init_logger(__name__) +# e4m3 finite max; hardcoded so this does not depend on torch.finfo supporting +# float8 dtypes. +_FP8_E4M3_MAX = 448.0 +# Batch sizes the draft head is timed at when picking a backend. The trace +# releases requests 200ms apart against multi-second generations, so the +# running batch ramps well past one. +_DRAFT_HEAD_PROBE_BATCHES = (4, 16, 32) + + +def _quantize_lm_head_to_fp8( + weight: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Per-output-channel fp8 quantization of an ``[N, K]`` lm_head weight. + + Returns ``(qweight, scale)`` with ``qweight`` fp8 ``[N, K]`` and ``scale`` + a 1-D float32 tensor of length ``N``, such that row ``i`` of the original + weight is approximately ``scale[i] * qweight[i]``. Rows are converted in + chunks: a whole-tensor float32 upcast of a 248k x 5120 head would + transiently need several GB, and this runs before the KV cache is sized. + """ + fp8_dtype = current_platform.fp8_dtype() + n, k = weight.shape + scale = torch.empty(n, dtype=torch.float32, device=weight.device) + qweight = torch.empty((n, k), dtype=fp8_dtype, device=weight.device) + for start in range(0, n, 8192): + end = min(start + 8192, n) + # Never in-place here: the source may be the target model's own head. + rows = weight[start:end].to(dtype=torch.float32, copy=True) + row_scale = rows.abs().amax(dim=1).clamp_(min=1e-12) / _FP8_E4M3_MAX + qweight[start:end] = ( + (rows / row_scale.unsqueeze(1)) + .clamp_(-_FP8_E4M3_MAX, _FP8_E4M3_MAX) + .to(fp8_dtype) + ) + scale[start:end] = row_scale + del rows, row_scale + return qweight, scale + + +class _DraftHeadQuantMethod: + """Adapter that lets a draft head stand in for a ``VocabParallelEmbedding``. + + ``LogitsProcessor`` only ever reaches a head through + ``lm_head.quant_method.apply(lm_head, ...)``, so implementing that one call + is enough. Kept separate from the head itself so an ``nn.Module`` head does + not end up registered as its own submodule. + """ + + def apply( + self, + layer: "object", + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + return layer.compute_draft_logits(x, bias) + + +class _Fp8DraftLMHead: + """An fp8 stand-in for a speculative drafter's bf16 ``lm_head``. + + The drafter only proposes tokens; under greedy sampling the rejection + sampler emits the target model's argmax at every position, so reduced draft + precision can move the acceptance rate but never an emitted token. Halving + the bytes read per draft step is worth far more than the acceptance this + costs, because the head is re-read once per speculative token. + """ + + def __init__(self, qweight: torch.Tensor, scale: torch.Tensor) -> None: + # cutlass_scaled_mm wants a column-major B. Transposing a contiguous + # [N, K] already gives that, so do not make it contiguous: that would + # both break the kernel's stride check and cost a full extra copy. + self.b = qweight.t() + self.scale = scale.reshape(-1, 1).to(torch.float32) + self.quant_method = _DraftHeadQuantMethod() + + def compute_draft_logits( + self, x: torch.Tensor, bias: torch.Tensor | None = None + ) -> torch.Tensor: + import vllm._custom_ops as ops + + x_2d = x.reshape(-1, x.shape[-1]) + x_fp8, x_scale = ops.scaled_fp8_quant(x_2d, use_per_token_if_dynamic=True) + out = ops.cutlass_scaled_mm( + x_fp8, + self.b, + scale_a=x_scale, + scale_b=self.scale, + out_dtype=x.dtype, + bias=bias, + ) + return out.reshape(*x.shape[:-1], out.shape[-1]) + + +class _MarlinDraftLMHead(nn.Module): + """Weight-only fp8 stand-in for a drafter's ``lm_head``. + + Leaves activations alone, so it costs less draft accuracy than + :class:`_Fp8DraftLMHead`, but Marlin is aimed at GPUs without native fp8 + support and is only tried when the native path does not win. Doubles as the + layer object Marlin's setup helper repacks in place. + """ + + def __init__( + self, qweight: torch.Tensor, scale: torch.Tensor, orig_dtype: torch.dtype + ) -> None: + from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import ( + prepare_fp8_layer_for_marlin, + ) + + super().__init__() + n, k = qweight.shape + self.output_size_per_partition = n + self.input_size_per_partition = k + self.orig_dtype = orig_dtype + self.weight = torch.nn.Parameter(qweight, requires_grad=False) + self.weight_scale = torch.nn.Parameter( + scale.reshape(1, n).to(orig_dtype), requires_grad=False + ) + prepare_fp8_layer_for_marlin(self, size_k_first=False) + self.quant_method = _DraftHeadQuantMethod() + + def compute_draft_logits( + self, x: torch.Tensor, bias: torch.Tensor | None = None + ) -> torch.Tensor: + from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import ( + apply_fp8_marlin_linear, + ) + + return apply_fp8_marlin_linear( + input=x, + weight=self.weight, + weight_scale=self.weight_scale, + workspace=self.workspace, + size_n=self.output_size_per_partition, + size_k=self.input_size_per_partition, + bias=bias, + ) + + +def _time_draft_head(fn, probes: list, iters: int = 12) -> float: + """Median-ish wall time in ms for ``fn`` over each probe, summed.""" + total = 0.0 + for probe in probes: + for _ in range(3): + fn(probe) + torch.cuda.synchronize() + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(iters): + fn(probe) + end.record() + torch.cuda.synchronize() + total += start.elapsed_time(end) / iters + return total + class SpecDecodeBaseProposer: def __init__( @@ -1306,6 +1462,7 @@ class SpecDecodeBaseProposer: self._maybe_share_embeddings(target_language_model) self._maybe_share_lm_head(target_language_model) + self._maybe_quantize_draft_lm_head() if ( self.parallel_drafting @@ -1501,6 +1658,154 @@ class SpecDecodeBaseProposer: "(communication: O(2*tp_size) vs O(vocab_size))." ) + def _maybe_quantize_draft_lm_head(self) -> None: + """Give the drafter an fp8 copy of its lm_head, leaving the target's alone. + + The head dominates draft cost: it is re-read once per speculative token, + so for a large vocabulary it can outweigh the drafter's transformer work + several times over. Halving those bytes is safe because the drafter only + proposes; under greedy sampling the rejection sampler emits the target's + argmax regardless, so this trades a little acceptance for a lot of + bandwidth and cannot change an emitted token. + + A candidate is adopted only if it both reproduces the original head's + logits and is measurably faster than it. Model load is not on the + measured path, so that check is free. Any failure leaves the existing + head in place rather than failing load. + """ + head = getattr(self.model, "lm_head", None) + weight = getattr(head, "weight", None) + try: + if not current_platform.is_cuda(): + return + if self.vllm_config.parallel_config.tensor_parallel_size != 1: + # The head is vocab-sharded; scales would have to be sharded too. + return + if self.use_local_argmax_reduction: + # That path reads shard_indices off the head, which a + # substitute does not carry. + return + if not isinstance(weight, torch.Tensor) or weight.dim() != 2: + return + if weight.dtype not in (torch.bfloat16, torch.float16): + # Already quantized, or something we do not understand. + return + n, k = weight.shape + if n % 16 != 0 or k % 16 != 0: + return + + qweight, scale = _quantize_lm_head_to_fp8(weight) + + probes = [ + torch.randn(m, k, dtype=weight.dtype, device=weight.device) + for m in _DRAFT_HEAD_PROBE_BATCHES + ] + baseline_ms = _time_draft_head( + lambda t: torch.nn.functional.linear(t, weight), probes + ) + reference = torch.nn.functional.linear(probes[0], weight) + + # Native fp8 first: Marlin targets GPUs without fp8 support and is + # gated off above compute capability 89 in vLLM's own selector. + fp8_head = None + backend = "" + chosen_ms = baseline_ms + similarity = 0.0 + for name, build in ( + ("cutlass", lambda: _Fp8DraftLMHead(qweight, scale)), + ( + "marlin", + lambda: _MarlinDraftLMHead(qweight, scale, weight.dtype), + ), + ): + try: + candidate_head = build() + candidate = candidate_head.compute_draft_logits(probes[0]) + score = ( + torch.nn.functional.cosine_similarity( + reference.float(), candidate.float(), dim=-1 + ) + .mean() + .item() + ) + agreement = ( + (reference.argmax(-1) == candidate.argmax(-1)) + .float() + .mean() + .item() + ) + elapsed = _time_draft_head( + candidate_head.compute_draft_logits, probes + ) + except Exception as err: # noqa: BLE001 - try the next backend + logger.info("fp8 draft lm_head (%s) unavailable: %s", name, err) + continue + # A wrong layout or scale collapses both of these to roughly + # zero, which is what they are for. They are deliberately not + # tight: per-channel e4m3 plus per-token activation scaling + # measures about 0.9993 and 0.87 on adversarial random probes, + # and rejecting that would throw away a working backend. + if score < 0.995 or agreement < 0.6: + logger.info( + "fp8 draft lm_head (%s) rejected: cosine %.5f, " + "top-1 agreement %.3f", + name, + score, + agreement, + ) + continue + if elapsed >= chosen_ms: + logger.info( + "fp8 draft lm_head (%s) rejected: %.3f ms vs %.3f ms " + "for the original head", + name, + elapsed, + chosen_ms, + ) + continue + fp8_head, backend, chosen_ms, similarity = ( + candidate_head, + name, + elapsed, + score, + ) + break + del probes, reference + + if fp8_head is None: + logger.error( + "PARETON_FP8_DRAFT_HEAD_FALLBACK: no fp8 draft lm_head was " + "both accurate and faster than the original; running with " + "the unquantized head." + ) + return + + if hasattr(self.model, "lm_head"): + del self.model.lm_head + self.model.lm_head = fp8_head + inner = getattr(self.model, "model", None) + layers = getattr(inner, "layers", None) if inner is not None else None + if layers is not None: + items = layers.values() if isinstance(layers, nn.ModuleDict) else layers + for layer in items: + shared = getattr(layer, "shared_head", None) + if shared is not None and hasattr(shared, "head"): + del shared.head + shared.head = fp8_head + logger.info( + "Draft lm_head quantized to fp8 via %s: cosine %.5f, %.3f ms vs " + "%.3f ms unquantized. The target model's head is unchanged.", + backend, + similarity, + chosen_ms, + baseline_ms, + ) + except Exception as err: # noqa: BLE001 - never fail load over this + logger.error( + "PARETON_FP8_DRAFT_HEAD_FALLBACK: skipping fp8 draft lm_head: %s", + err, + ) + @torch.inference_mode() def dummy_run( self,