From 81807002fb36e2b9eab05728ea3bdf54fd12c280 Mon Sep 17 00:00:00 2001 From: halt95 Date: Tue, 22 Sep 2026 04:28:34 -0400 Subject: [PATCH 029/103] rebase: remove v0.30.0 merge artefacts; import order - Drop the PLE-offload half of ple_layer.py that upstream #54371 deleted and a 3-way merge resurrected (it no longer imported on v0.30.0). - Drop the structured-output helper duplicate that v0.30.0 already provides, and an unused import. - Re-sort the import blocks the merges left unsorted (ruff I001). --- tests/v1/core/test_deferred_block_free.py | 3 +- vllm/models/qwen4_exp/nvidia/model.py | 7 +- vllm/models/qwen4_exp/nvidia/mtp.py | 2 +- vllm/models/qwen4_exp/nvidia/ple_layer.py | 882 ---------------------- vllm/models/qwen4_exp/nvidia/qsa.py | 3 +- vllm/v1/core/sched/scheduler.py | 4 +- vllm/v1/structured_output/utils.py | 13 - 7 files changed, 8 insertions(+), 906 deletions(-) diff --git a/tests/v1/core/test_deferred_block_free.py b/tests/v1/core/test_deferred_block_free.py index c45f5a7..e94161b 100644 --- a/tests/v1/core/test_deferred_block_free.py +++ b/tests/v1/core/test_deferred_block_free.py @@ -16,14 +16,13 @@ import time from unittest.mock import PropertyMock, patch import pytest - import torch from vllm.config import VllmConfig from vllm.v1.core.block_pool import BlockPool from vllm.v1.core.kv_cache_utils import BlockHashWithGroupId, KVCacheBlockCopy -from vllm.v1.kv_cache_interface import CircularBufferSpec from vllm.v1.core.sched.output import SchedulerOutput +from vllm.v1.kv_cache_interface import CircularBufferSpec from vllm.v1.outputs import ModelRunnerOutput from vllm.v1.request import RequestStatus diff --git a/vllm/models/qwen4_exp/nvidia/model.py b/vllm/models/qwen4_exp/nvidia/model.py index 269bbbd..816a4e2 100644 --- a/vllm/models/qwen4_exp/nvidia/model.py +++ b/vllm/models/qwen4_exp/nvidia/model.py @@ -2,16 +2,16 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Inference-only Qwen4Exp model.""" +import os from collections.abc import Iterable from itertools import islice -import os - import torch from torch import nn from vllm.config import VllmConfig from vllm.distributed import get_pp_group +from vllm.logger import init_logger from vllm.model_executor.layers.fused_moe.utils import ( is_model_fused_shared_expert_compatible, ) @@ -28,8 +28,8 @@ from vllm.model_executor.layers.mamba.mamba_utils import ( ) from vllm.model_executor.layers.quantization import QuantizationConfig from vllm.model_executor.layers.vocab_parallel_embedding import ( - PinnedHostVocabParallelEmbedding, ParallelLMHead, + PinnedHostVocabParallelEmbedding, VocabParallelEmbedding, should_use_host_embed_table, ) @@ -85,7 +85,6 @@ from .hyperconnection import GatedResidual, HyperConnectionConfig from .low_latency_gemm import enable_qwen4_exp_low_latency_gemm from .ple_layer import Qwen4ExpPLELayer from .qsa import Qwen4ExpQSAAttention, load_qsa_static_kv_scales -from vllm.logger import init_logger logger = init_logger(__name__) diff --git a/vllm/models/qwen4_exp/nvidia/mtp.py b/vllm/models/qwen4_exp/nvidia/mtp.py index 0ad96a2..2df3868 100644 --- a/vllm/models/qwen4_exp/nvidia/mtp.py +++ b/vllm/models/qwen4_exp/nvidia/mtp.py @@ -28,8 +28,8 @@ from vllm.model_executor.layers.layernorm import GemmaRMSNorm from vllm.model_executor.layers.linear import ColumnParallelLinear from vllm.model_executor.layers.logits_processor import LogitsProcessor from vllm.model_executor.layers.vocab_parallel_embedding import ( - PinnedHostVocabParallelEmbedding, ParallelLMHead, + PinnedHostVocabParallelEmbedding, VocabParallelEmbedding, should_use_host_embed_table, ) diff --git a/vllm/models/qwen4_exp/nvidia/ple_layer.py b/vllm/models/qwen4_exp/nvidia/ple_layer.py index bec35e6..2ec8790 100644 --- a/vllm/models/qwen4_exp/nvidia/ple_layer.py +++ b/vllm/models/qwen4_exp/nvidia/ple_layer.py @@ -4,8 +4,6 @@ from collections.abc import Sequence -import os - import torch from torch import nn @@ -65,886 +63,6 @@ class Qwen4ExpPLEGroupedNorm(nn.Module): return (normalized * (1.0 + self.weight.float())).to(input_dtype) -class Qwen4ExpPLEFp8EmbeddingMethod(QuantizeMethodBase): - """FP8 PLE embedding with one global checkpoint scale.""" - - def create_weights( - self, - layer: nn.Module, - input_size_per_partition: int, - output_partition_sizes: list[int], - input_size: int, - output_size: int, - params_dtype: torch.dtype, - **extra_weight_attrs, - ) -> None: - del input_size, output_size, params_dtype - weight_loader = extra_weight_attrs.get("weight_loader") - weight = create_fp8_weight_parameter( - sum(output_partition_sizes), input_size_per_partition, weight_loader - ) - layer.register_parameter("weight", weight) - - weight_scale = create_fp8_scale_parameter( - PerTensorScaleParameter, - output_partition_sizes, - input_size_per_partition, - None, - weight_loader, - scale_dtype=torch.float32, - ) - layer.register_parameter("weight_scale", weight_scale) - - def process_weights_after_loading(self, layer: nn.Module) -> None: - """Reject FP8 PLE checkpoints without a global scale.""" - sentinel = torch.finfo(torch.float32).min - if torch.any(layer.weight_scale == sentinel): - raise ValueError("FP8 PLE checkpoint is missing its global scale") - - def apply( - self, - layer: nn.Module, - x: torch.Tensor, - bias: torch.Tensor | None = None, - ) -> torch.Tensor: - raise NotImplementedError("PLE FP8 weights only support embedding lookup") - - def embedding(self, layer: nn.Module, input_: torch.Tensor) -> torch.Tensor: - return F.embedding(input_, layer.weight) - - -_NVFP4_BLOCK_SIZE = 16 - - -class Qwen4ExpPLENVFp4EmbeddingMethod(QuantizeMethodBase): - """NVFP4 PLE embedding kept packed at runtime. - - The table stays as uint8 codes with FP8 block scales (16 values per - block) and one FP32 global scale. A lookup only gathers the packed rows; - the PLE layer dequantizes them after lookup. - """ - - def create_weights( - self, - layer: nn.Module, - input_size_per_partition: int, - output_partition_sizes: list[int], - input_size: int, - output_size: int, - params_dtype: torch.dtype, - **extra_weight_attrs, - ) -> None: - del input_size, output_size, extra_weight_attrs - if input_size_per_partition % _NVFP4_BLOCK_SIZE: - raise ValueError( - "NVFP4 PLE embedding requires the embedding dim to be a " - f"multiple of {_NVFP4_BLOCK_SIZE}, got {input_size_per_partition}" - ) - self.params_dtype = params_dtype - self.head_dim = input_size_per_partition - self.packed_row_width = ( - input_size_per_partition // 2 - + input_size_per_partition // _NVFP4_BLOCK_SIZE - ) - rows = sum(output_partition_sizes) - weight = nn.Parameter( - torch.empty(rows, input_size_per_partition // 2, dtype=torch.uint8), - requires_grad=False, - ) - set_weight_attrs(weight, {"input_dim": 1, "output_dim": 0}) - layer.register_parameter("weight", weight) - - weight_scale = nn.Parameter( - torch.empty( - rows, - input_size_per_partition // _NVFP4_BLOCK_SIZE, - dtype=torch.float8_e4m3fn, - ), - requires_grad=False, - ) - set_weight_attrs(weight_scale, {"input_dim": 1, "output_dim": 0}) - layer.register_parameter("weight_scale", weight_scale) - - weight_scale_2 = nn.Parameter( - torch.empty((), dtype=torch.float32), requires_grad=False - ) - layer.register_parameter("weight_scale_2", weight_scale_2) - self._lut: torch.Tensor | None = None - - def process_weights_after_loading(self, layer: nn.Module) -> None: - # Build the lookup table before CUDA graph capture. - self._lut = torch.tensor( - _FP4_VALUES, dtype=torch.float32, device=layer.weight.device - ) - - def apply( - self, - layer: nn.Module, - x: torch.Tensor, - bias: torch.Tensor | None = None, - ) -> torch.Tensor: - raise NotImplementedError("PLE NVFP4 weights only support embedding lookup") - - def embedding(self, layer: nn.Module, input_: torch.Tensor) -> torch.Tensor: - """Gather packed rows for deferred dequantization.""" - codes = F.embedding(input_, layer.weight) - scales = F.embedding(input_, layer.weight_scale) - return torch.cat((codes, scales.view(torch.uint8)), dim=-1) - - def dequantize( - self, - packed_rows: torch.Tensor, - scale_2: torch.Tensor, - output_dtype: torch.dtype, - ) -> torch.Tensor: - """Dequantize rows containing NVFP4 codes and block scales.""" - lut = self._lut - if lut is None or lut.device != packed_rows.device: - lut = torch.tensor( - _FP4_VALUES, dtype=torch.float32, device=packed_rows.device - ) - self._lut = lut - return _dequant_nvfp4_rows( - packed_rows, self.head_dim, scale_2, output_dtype, lut - ) - - -def _ple_dtype_is_fp8(ple_embedding_dtype: object) -> bool: - """Return whether the model config declares FP8 PLE checkpoint weights.""" - - if ple_embedding_dtype is None: - return False - if isinstance(ple_embedding_dtype, torch.dtype): - return ple_embedding_dtype == torch.float8_e4m3fn - return str(ple_embedding_dtype).rsplit(".", 1)[-1] == "float8_e4m3fn" - - -def _get_ple_embedding_quant_method( - quant_config: QuantizationConfig | None, - prefix: str, - ple_embedding_dtype: object = None, -) -> QuantizeMethodBase | None: - """Select a packed PLE embedding method for quantized checkpoint shards.""" - - if isinstance(quant_config, ModelOptMixedPrecisionConfig): - if quant_config._resolve_quant_algo(prefix) == "FP8": - return Qwen4ExpPLEFp8EmbeddingMethod() - return None - - if isinstance(quant_config, Fp8Config): - if not quant_config.is_checkpoint_fp8_serialized: - return None - ignored_layers = quant_config.ignored_layers - if is_layer_skipped( - prefix, - ignored_layers, - quant_config.packed_modules_mapping, - match_mode=quant_config.ignored_layers_match_mode, - ): - return None - # PLE checkpoint shards form one runtime embedding parameter. - shard_prefix = f"{prefix}.shard_" - if any(name.startswith(shard_prefix) for name in ignored_layers): - return None - return Qwen4ExpPLEFp8EmbeddingMethod() - - if isinstance(quant_config, ModelOptNvFp4Config): - if not quant_config.is_checkpoint_nvfp4_serialized: - return None - if not quant_config.is_layer_excluded(prefix): - logger.info_once("PLE embedding %s uses the runtime NVFP4 method", prefix) - return Qwen4ExpPLENVFp4EmbeddingMethod() - if _ple_dtype_is_fp8(ple_embedding_dtype): - logger.info_once( - "Excluded ModelOpt PLE embedding %s uses the runtime FP8 method", - prefix, - ) - return Qwen4ExpPLEFp8EmbeddingMethod() - - # Checkpoints pairing quantized experts with a global-scale FP8 PLE table - # (the RadixArk Flash-Next table grafted onto a CompressedTensors checkpoint) - # list the PLE under exclude_modules, so the FP8-ness is visible only in the - # tensors' dtype / the config's ple_embedding_dtype, not in the quant config - # type dispatched above. Without this branch the offload worker treats the - # table as BF16 and materializes ~96 GiB in host RAM (OOM-killed on the - # 96 GB host). Honour either the config key (the -e1 - # checkpoint copy) or the explicit VLLM_PLE_FP8_GLOBAL_SCALE=1 opt-in the - # previous release uses. Placement after the typed dispatch is the - # position already proven in service. - if _ple_dtype_is_fp8(ple_embedding_dtype) or os.environ.get( - "VLLM_PLE_FP8_GLOBAL_SCALE" - ) == "1": - logger.info_once( - "PLE embedding %s: config declares an FP8 table with quant_config %s; " - "using the runtime FP8 method (global weight_scale)", - prefix, - type(quant_config).__name__, - ) - return Qwen4ExpPLEFp8EmbeddingMethod() - - return None - - -_FP4_VALUES = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0) - - -def _dequant_nvfp4_codes( - packed: torch.Tensor, - scale: torch.Tensor, - scale_2: torch.Tensor, - lut: torch.Tensor | None = None, -) -> torch.Tensor: - """Unpack NVFP4 codes with FP8 block scales and an FP32 global scale.""" - half = packed.shape[-1] - codes = torch.stack([packed & 0xF, (packed >> 4) & 0xF], dim=-1).reshape( - *packed.shape[:-1], half * 2 - ) - if lut is None: - lut = torch.tensor(_FP4_VALUES, dtype=torch.float32, device=packed.device) - magnitude = lut[(codes & 0x7).long()] - sign = ((codes >> 3) & 1).to(torch.float32) - fp4 = (magnitude * (1 - 2 * sign)).reshape( - *packed.shape[:-1], -1, _NVFP4_BLOCK_SIZE - ) - output = fp4 * scale.float().unsqueeze(-1) * scale_2.float() - return output.reshape(*packed.shape[:-1], half * 2) - - -def _dequant_nvfp4_rows( - packed_rows: torch.Tensor, - head_dim: int, - scale_2: torch.Tensor, - output_dtype: torch.dtype, - lut: torch.Tensor, -) -> torch.Tensor: - """Dequantize rows containing NVFP4 codes followed by block scales.""" - half = head_dim // 2 - codes = packed_rows[..., :half] - scales = packed_rows[..., half:].contiguous().view(torch.float8_e4m3fn) - return _dequant_nvfp4_codes(codes, scales, scale_2, lut).to(output_dtype) - - -def _get_shared_nvfp4_outer_scale( - outer_scales: dict[int, torch.Tensor], -) -> torch.Tensor: - """Validate and return the global scale shared by NVFP4 PLE shards.""" - first_index, reference = next(iter(outer_scales.items())) - reference = reference.reshape(()) - for shard_index, outer_scale in outer_scales.items(): - if not torch.equal(outer_scale.reshape(()), reference): - raise ValueError( - "NVFP4 PLE shards must share the same global scale, but " - f"shards {first_index} and {shard_index} differ" - ) - return reference - - -class Qwen4ExpNGramEmbedding(PleOffloadLayer): - _offload_quant_method: QuantizeMethodBase | None = None - _MASK64 = (1 << 64) - 1 - _SPLITMIX_GAMMA = 0x9E3779B97F4A7C15 - _SPLITMIX_M1 = 0xBF58476D1CE4E5B9 - _SPLITMIX_M2 = 0x94D049BB133111EB - _PLE_LAYER_PRIME = 10007 - - @classmethod - def _splitmix64(cls, value: int) -> int: - """Mix an integer into a deterministic unsigned 64-bit value.""" - value = (value + cls._SPLITMIX_GAMMA) & cls._MASK64 - value = ((value ^ (value >> 30)) * cls._SPLITMIX_M1) & cls._MASK64 - value = ((value ^ (value >> 27)) * cls._SPLITMIX_M2) & cls._MASK64 - return (value ^ (value >> 31)) & cls._MASK64 - - @staticmethod - def _is_prime_64(value: int) -> bool: - """Return whether a 64-bit integer is prime.""" - if value < 2: - return False - for prime in (2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37): - if value % prime == 0: - return value == prime - exponent = value - 1 - shifts = 0 - while exponent % 2 == 0: - exponent //= 2 - shifts += 1 - for base in (2, 325, 9375, 28178, 450775, 9780504, 1795265022): - if base % value == 0: - continue - witness = pow(base, exponent, value) - if witness in (1, value - 1): - continue - for _ in range(shifts - 1): - witness = pow(witness, 2, value) - if witness == value - 1: - break - else: - return False - return True - - @classmethod - def _nth_prime_after(cls, start: int, count: int) -> int: - """Return the ``count``-th prime strictly greater than ``start``.""" - prime = int(start) - for _ in range(count): - candidate = prime + 1 - if candidate <= 2: - prime = 2 - continue - if candidate % 2 == 0: - candidate += 1 - while not cls._is_prime_64(candidate): - candidate += 2 - prime = candidate - return prime - - @classmethod - def _make_layer_multipliers( - cls, - *, - ngram_size: int, - unigram_vocab_size: int, - seed: int, - ple_dense_layer_id: int, - ) -> list[int]: - """Build deterministic hash multipliers for one PLE layer.""" - max_multiplier = ((1 << 63) - 1) // unigram_vocab_size - half_bound = max(1, max_multiplier // 2) - base_seed = seed + cls._PLE_LAYER_PRIME * ple_dense_layer_id - multipliers = [] - for index in range(ngram_size): - value = base_seed + cls._SPLITMIX_GAMMA * (index + 1) - multipliers.append(2 * (cls._splitmix64(value) % half_bound) + 1) - return multipliers - - @classmethod - def _make_vocab_layout( - cls, - *, - ngram_vocab_size_base: int, - ngram_heads: int, - ple_dense_layer_id: int, - ) -> tuple[list[int], list[int], int]: - """Build per-head vocabulary sizes, offsets, and total row count.""" - sizes: list[int] = [] - offsets: list[int] = [] - offset = 0 - for local_head in range(ngram_heads): - global_head = ple_dense_layer_id * ngram_heads + local_head - size = cls._nth_prime_after(ngram_vocab_size_base - 1, global_head + 1) - sizes.append(size) - offsets.append(offset) - offset += size - return sizes, offsets, offset - - def __init__( - self, - config: Qwen4ExpTextConfig, - embedding_dim: int, - ple_dense_layer_id: int, - prefix: str, - layer_name: str, - quant_config: QuantizationConfig | None = None, - params_dtype: torch.dtype | None = None, - ) -> None: - super().__init__() - self.layer_name = layer_name - self.embedding_dim = embedding_dim - self.ngram_size = int(config.ngram_size) - self.heads_per_ngram = int(config.heads_per_ngram) - self.ngram_heads = (self.ngram_size - 1) * self.heads_per_ngram - if self.ngram_size < 2: - raise ValueError(f"ngram_size must be >= 2, got {self.ngram_size}") - if self.heads_per_ngram <= 0: - raise ValueError(f"heads_per_ngram must be > 0, got {self.heads_per_ngram}") - if embedding_dim % self.ngram_heads: - raise ValueError( - "ple_embed_dim must be divisible by total ngram heads: " - f"{embedding_dim} % {self.ngram_heads} != 0" - ) - self.head_dim = embedding_dim // self.ngram_heads - self.eos_token_id = int(config.eos_token_id) - self.unigram_vocab_size = int(config.vocab_size) - self.split_ngram_parts = int(getattr(config, "split_ngram_parts", 512)) - if self.split_ngram_parts <= 0: - raise ValueError("split_ngram_parts must be positive") - - multipliers = self._make_layer_multipliers( - ngram_size=self.ngram_size, - unigram_vocab_size=self.unigram_vocab_size, - seed=int(getattr(config, "seed", 1234)), - ple_dense_layer_id=ple_dense_layer_id, - ) - self.register_buffer( - "layer_multipliers", - torch.tensor(multipliers, dtype=torch.long), - persistent=True, - ) - - sizes, offsets, total_vocab_size = self._make_vocab_layout( - ngram_vocab_size_base=int(config.ngram_vocab_size_base), - ngram_heads=self.ngram_heads, - ple_dense_layer_id=ple_dense_layer_id, - ) - self.register_buffer( - "ngram_heads_vocab_sizes", - torch.tensor(sizes, dtype=torch.long), - persistent=True, - ) - self.register_buffer( - "ngram_heads_offsets", - torch.tensor(offsets, dtype=torch.long), - persistent=True, - ) - divisor = int(config.make_ngram_vocab_size_divisible_by) - padded_vocab_size = ((total_vocab_size + divisor - 1) // divisor) * divisor - self.ngram_embedding = PLEVocabParallelEmbedding( - padded_vocab_size, - self.head_dim, - params_dtype=params_dtype, - padding_size=divisor, - prefix=f"{prefix}.ngram_embedding", - quant_method=_get_ple_embedding_quant_method( - quant_config, - f"{prefix}.ngram_embedding", - getattr(config, "ple_embedding_dtype", None), - ), - ) - - @staticmethod - def _shift_precompute( - tokens: torch.Tensor, eos_token_id: int - ) -> tuple[torch.Tensor, torch.Tensor]: - if tokens.dim() != 2: - raise ValueError("tokens must be a 2D tensor") - batch_size, seq_len = tokens.shape - positions = torch.arange(seq_len, device=tokens.device, dtype=torch.int64) - eos_positions = torch.where(tokens == eos_token_id, positions, -1) - previous_eos_inclusive = torch.cummax(eos_positions, dim=1).values - previous_eos = torch.cat( - [ - eos_positions.new_full((batch_size, 1), -1), - previous_eos_inclusive[:, :-1], - ], - dim=1, - ) - return positions, positions.unsqueeze(0) - previous_eos - 1 - - @staticmethod - def _shift_apply( - tokens: torch.Tensor, - positions: torch.Tensor, - position_in_segment: torch.Tensor, - shift: int, - eos_token_id: int, - ) -> torch.Tensor: - if shift == 0: - return tokens - source = positions - shift - gather_indices = source.clamp_min(0).unsqueeze(0).expand(tokens.shape[0], -1) - shifted = tokens.gather(1, gather_indices) - valid = (source.unsqueeze(0) >= 0) & (position_in_segment >= shift) - return torch.where(valid, shifted, tokens.new_full((), eos_token_id)) - - def compute_ngram_ids( - self, - input_ids: torch.Tensor, - query_start_loc: torch.Tensor, - ngram_context: torch.Tensor, - output: torch.Tensor | None = None, - ) -> torch.Tensor: - """Compute n-gram embedding indices for the current request layout.""" - input_ids = input_ids.reshape(-1) - num_reqs = query_start_loc.numel() - 1 - num_tokens = input_ids.shape[0] - - if input_ids.is_cuda: - return ple_ngram_ids( - input_ids=input_ids, - query_start_loc=query_start_loc, - ngram_context=ngram_context, - layer_multipliers=self.layer_multipliers, - ngram_heads_vocab_sizes=self.ngram_heads_vocab_sizes, - ngram_heads_offsets=self.ngram_heads_offsets, - eos_token_id=self.eos_token_id, - heads_per_ngram=self.heads_per_ngram, - output=output, - ) - input_ids = input_ids.long() - query_start_loc = query_start_loc.long() - if num_reqs <= 0: - raise ValueError("PLE requires at least one request") - - if is_offload_process(): - max_seq_len = max( - 1, - int((query_start_loc[1:] - query_start_loc[:-1]).max().item()), - ) - # CUDA graphs can pad input_ids beyond the valid request layout. - # Keep those rows from overwriting the final real token. - num_valid_tokens = min(int(query_start_loc[-1].item()), num_tokens) - else: - max_seq_len = num_tokens - num_valid_tokens = num_tokens - - positions = torch.arange(num_tokens, device=input_ids.device, dtype=torch.int64) - packed = torch.full( - (num_reqs, max_seq_len), - self.eos_token_id, - device=input_ids.device, - dtype=torch.int64, - ) - request_indices = torch.searchsorted(query_start_loc, positions, right=True) - 1 - request_indices.clamp_(max=num_reqs - 1) - columns = (positions - query_start_loc[request_indices]).clamp( - 0, packed.shape[1] - 1 - ) - packed[request_indices[:num_valid_tokens], columns[:num_valid_tokens]] = ( - input_ids[:num_valid_tokens] - ) - ngram_context = ngram_context[:num_reqs].to( - device=input_ids.device, dtype=torch.long - ) - - context = torch.cat([ngram_context, packed], dim=-1) - positions_2d, position_in_segment = self._shift_precompute( - context, self.eos_token_id - ) - shifted = [context] - for shift in range(1, self.ngram_size): - shifted.append( - self._shift_apply( - context, - positions_2d, - position_in_segment, - shift, - self.eos_token_id, - ) - ) - adjusted_columns = columns + self.ngram_size - 1 - id_blocks = [] - for ngram in range(2, self.ngram_size + 1): - start = (ngram - 2) * self.heads_per_ngram - end = start + self.heads_per_ngram - mixed = shifted[0] * self.layer_multipliers[0] - for index in range(1, ngram): - mixed = torch.bitwise_xor( - mixed, shifted[index] * self.layer_multipliers[index] - ) - sizes = self.ngram_heads_vocab_sizes[start:end] - offsets = self.ngram_heads_offsets[start:end] - ids = torch.remainder(mixed.unsqueeze(-1), sizes) + offsets - id_blocks.append(ids[request_indices, adjusted_columns]) - return torch.cat(id_blocks, dim=-1) - - def forward_impl( # type: ignore[override] - self, - hidden_states: torch.Tensor, - input_ids: torch.Tensor, - query_start_loc: torch.Tensor, - ngram_context: torch.Tensor, - output_buffer: torch.Tensor | None = None, - ) -> torch.Tensor: - del hidden_states - if input_ids.is_cuda: - # Keep request-dependent ID generation outside piecewise graphs. - ngram_ids = input_ids.new_empty( - (input_ids.numel(), self.ngram_heads), dtype=torch.long - ) - torch.ops.vllm.qwen4_exp_compute_ple_ngram_ids( - input_ids, - query_start_loc, - ngram_context, - ngram_ids, - self.layer_name, - ) - else: - ngram_ids = self.compute_ngram_ids( - input_ids, - query_start_loc, - ngram_context, - ) - num_tokens = input_ids.reshape(-1).shape[0] - if output_buffer is not None: - quant_method = getattr(self.ngram_embedding, "quant_method", None) - if isinstance(quant_method, Qwen4ExpPLENVFp4EmbeddingMethod): - output_dim = self.get_offload_output_dim(self.embedding_dim) - output = output_buffer[:num_tokens, :output_dim] - output.copy_(self.ngram_embedding(ngram_ids).flatten(-2)) - return output - output = output_buffer[:num_tokens, : self.embedding_dim] - torch.index_select( - self.ngram_embedding.weight, - 0, - ngram_ids.reshape(-1), - out=output.reshape(-1, self.head_dim), - ) - return output - return self.ngram_embedding(ngram_ids).flatten(-2) - - def get_offload_output_dtype(self, default_dtype: torch.dtype) -> torch.dtype: - """Keep quantized lookup results in their embedding storage dtype.""" - embedding = getattr(self, "ngram_embedding", None) - weight = getattr(embedding, "weight", None) - if weight is not None: - return weight.dtype - if isinstance(self._offload_quant_method, Qwen4ExpPLENVFp4EmbeddingMethod): - return torch.uint8 - if isinstance(self._offload_quant_method, Qwen4ExpPLEFp8EmbeddingMethod): - return torch.float8_e4m3fn - return default_dtype - - def get_offload_output_dim(self, default_dim: int) -> int: - """Keep NVFP4 lookup rows packed while transferring them to the GPU.""" - quant_method = getattr( - getattr(self, "ngram_embedding", None), "quant_method", None - ) - if quant_method is None: - quant_method = self._offload_quant_method - if isinstance(quant_method, Qwen4ExpPLENVFp4EmbeddingMethod): - if default_dim % _NVFP4_BLOCK_SIZE: - raise ValueError( - "NVFP4 PLE output dim must be a multiple of " - f"{_NVFP4_BLOCK_SIZE}, got {default_dim}" - ) - return default_dim // 2 + default_dim // _NVFP4_BLOCK_SIZE - return default_dim - - def initialize_dummy_offload_metadata(self, device: torch.device) -> None: - """Initialize quantization metadata skipped by the dummy loader.""" - quant_method = self._offload_quant_method - if isinstance(quant_method, Qwen4ExpPLEFp8EmbeddingMethod): - self.register_buffer( - "_offload_weight_scale", - torch.ones((), dtype=torch.float32, device=device), - persistent=False, - ) - elif isinstance(quant_method, Qwen4ExpPLENVFp4EmbeddingMethod): - self.register_buffer( - "_offload_weight_scale_2", - torch.ones((), dtype=torch.float32, device=device), - persistent=False, - ) - self.register_buffer( - "_offload_nvfp4_lut", - torch.tensor(_FP4_VALUES, dtype=torch.float32, device=device), - persistent=False, - ) - - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - """Load hash buffers and checkpoint-split embedding rows.""" - - # Prefix grouping may invoke this method repeatedly for one module. - # Treat the global FP8 scale as incrementally loaded state. - # GPU workers retain only dequantization metadata. The CPU process owns - # the embedding table and transfers quantized lookup rows unchanged. - if envs.VLLM_PLE_CPU_OFFLOAD and not is_offload_process(): - retained: set[str] = set() - quant_method = self._offload_quant_method - if isinstance(quant_method, Qwen4ExpPLEFp8EmbeddingMethod): - for name, loaded_weight in weights: - if name != "ngram_embedding.weight_scale": - continue - self.register_buffer( - "_offload_weight_scale", - loaded_weight.to( - device=torch.accelerator.current_accelerator() - ), - persistent=False, - ) - retained.add(name) - elif isinstance(quant_method, Qwen4ExpPLENVFp4EmbeddingMethod): - outer_scales: dict[int, torch.Tensor] = {} - for name, loaded_weight in weights: - if not name.startswith( - "ngram_embedding.shard_" - ) or not name.endswith(".weight_scale_2"): - continue - shard_text = name[ - len("ngram_embedding.shard_") : -len(".weight_scale_2") - ] - shard_index = int(shard_text) - outer_scales[shard_index] = loaded_weight - retained.add(name) - if not retained: - raise ValueError( - "NVFP4 PLE offload checkpoint is missing its global scale" - ) - scale_2 = _get_shared_nvfp4_outer_scale(outer_scales).to( - device=torch.accelerator.current_accelerator() - ) - self.register_buffer( - "_offload_weight_scale_2", scale_2, persistent=False - ) - self.register_buffer( - "_offload_nvfp4_lut", - torch.tensor( - _FP4_VALUES, dtype=torch.float32, device=scale_2.device - ), - persistent=False, - ) - else: - for _ in weights: - pass - return retained - - persistent_buffers = { - "layer_multipliers": self.layer_multipliers, - "ngram_heads_offsets": self.ngram_heads_offsets, - "ngram_heads_vocab_sizes": self.ngram_heads_vocab_sizes, - } - loaded: set[str] = set() - regular_weights: list[tuple[str, torch.Tensor]] = [] - shard_prefix = "ngram_embedding.shard_" - quant_method = getattr(self.ngram_embedding, "quant_method", None) - nvfp4_runtime = isinstance(quant_method, Qwen4ExpPLENVFp4EmbeddingMethod) - packed_codes: dict[int, torch.Tensor] = {} - packed_scales: dict[int, torch.Tensor] = {} - packed_outer_scales: dict[int, torch.Tensor] = {} - - for name, loaded_weight in weights: - leaf_name = name.rsplit(".", 1)[-1] - if leaf_name.startswith("hashstats_") or leaf_name == "token_lookup": - continue - if name in persistent_buffers: - buffer = persistent_buffers[name] - if buffer.shape != loaded_weight.shape: - raise ValueError( - f"Shape mismatch for {name}: expected " - f"{tuple(buffer.shape)}, got {tuple(loaded_weight.shape)}" - ) - buffer.copy_(loaded_weight.to(device=buffer.device, dtype=buffer.dtype)) - loaded.add(name) - continue - - if nvfp4_runtime and name.startswith(shard_prefix): - suffix = name[len(shard_prefix) :] - for ending, sink in ( - (".weight_scale_2", packed_outer_scales), - (".weight_scale", packed_scales), - (".weight", packed_codes), - ): - if not suffix.endswith(ending): - continue - shard_text = suffix[: -len(ending)] - shard_index = int(shard_text) - if shard_index >= self.split_ngram_parts: - raise ValueError( - f"PLE embedding shard index {shard_index} exceeds " - f"split_ngram_parts={self.split_ngram_parts}" - ) - sink[shard_index] = loaded_weight - break - else: - regular_weights.append((name, loaded_weight)) - continue - - if ( - not nvfp4_runtime - and name.startswith(shard_prefix) - and name.endswith(".weight") - ): - shard_text = name[len(shard_prefix) : -len(".weight")] - if not shard_text.isdigit(): - regular_weights.append((name, loaded_weight)) - continue - shard_index = int(shard_text) - if shard_index >= self.split_ngram_parts: - raise ValueError( - f"PLE embedding shard index {shard_index} exceeds " - f"split_ngram_parts={self.split_ngram_parts}" - ) - embedding = self.ngram_embedding - shard_size = ( - embedding.org_vocab_size + self.split_ngram_parts - 1 - ) // self.split_ngram_parts - checkpoint_start = shard_index * shard_size - expected_rows = max( - 0, - min(shard_size, embedding.org_vocab_size - checkpoint_start), - ) - expected_shape = (expected_rows, embedding.embedding_dim) - if tuple(loaded_weight.shape) != expected_shape: - raise ValueError( - f"Shape mismatch for PLE embedding shard {shard_index}: " - f"expected {expected_shape}, got " - f"{tuple(loaded_weight.shape)}" - ) - embedding.weight.weight_loader( - embedding.weight, - loaded_weight, - checkpoint_start=checkpoint_start, - ) - loaded.add("ngram_embedding.weight") - continue - regular_weights.append((name, loaded_weight)) - - if nvfp4_runtime: - if not packed_codes: - raise ValueError("NVFP4 PLE checkpoint contains no packed shards") - for shard_index in packed_codes: - if ( - shard_index not in packed_scales - or shard_index not in packed_outer_scales - ): - raise ValueError( - f"NVFP4 PLE shard {shard_index} is missing its scale tensors" - ) - shared_outer_scale = _get_shared_nvfp4_outer_scale( - { - shard_index: packed_outer_scales[shard_index] - for shard_index in packed_codes - } - ) - - for shard_index, codes in packed_codes.items(): - embedding = self.ngram_embedding - shard_size = ( - embedding.org_vocab_size + self.split_ngram_parts - 1 - ) // self.split_ngram_parts - checkpoint_start = shard_index * shard_size - expected_rows = max( - 0, min(shard_size, embedding.org_vocab_size - checkpoint_start) - ) - expected_codes_shape = (expected_rows, embedding.embedding_dim // 2) - if tuple(codes.shape) != expected_codes_shape: - raise ValueError( - f"Shape mismatch for packed PLE shard {shard_index}: " - f"expected {expected_codes_shape}, got {tuple(codes.shape)}" - ) - tp_bounds = { - "checkpoint_start": checkpoint_start, - "tp_start": embedding.shard_indices.org_vocab_start_index, - "tp_end": embedding.shard_indices.org_vocab_end_index, - } - scales = packed_scales[shard_index] - expected_scales_shape = ( - expected_rows, - embedding.embedding_dim // _NVFP4_BLOCK_SIZE, - ) - if tuple(scales.shape) != expected_scales_shape: - raise ValueError( - f"Scale shape mismatch for packed PLE shard {shard_index}: " - f"expected {expected_scales_shape}, got {tuple(scales.shape)}" - ) - copy_ple_embedding_shard_(embedding.weight.data, codes, **tp_bounds) - copy_ple_embedding_shard_(embedding.weight_scale.data, scales, **tp_bounds) - - if nvfp4_runtime: - self.ngram_embedding.weight_scale_2.data.copy_(shared_outer_scale) - loaded.update( - { - "ngram_embedding.weight", - "ngram_embedding.weight_scale", - "ngram_embedding.weight_scale_2", - } - ) - - if regular_weights: - loaded.update(AutoWeightsLoader(self).load_weights(regular_weights)) - return loaded - - class Qwen4ExpPLELayer(nn.Module, MambaBase): def __init__( self, diff --git a/vllm/models/qwen4_exp/nvidia/qsa.py b/vllm/models/qwen4_exp/nvidia/qsa.py index 91f34de..87c4b3c 100644 --- a/vllm/models/qwen4_exp/nvidia/qsa.py +++ b/vllm/models/qwen4_exp/nvidia/qsa.py @@ -22,6 +22,7 @@ from vllm.config import VllmConfig from vllm.config.cache import CacheDType from vllm.distributed import get_tensor_model_parallel_world_size from vllm.forward_context import get_forward_context +from vllm.logger import init_logger from vllm.model_executor.layers.attention.attention import ( set_default_quant_scales, ) @@ -29,7 +30,6 @@ from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase from vllm.model_executor.layers.layernorm import GemmaRMSNorm from vllm.model_executor.layers.linear import QKVParallelLinear, RowParallelLinear from vllm.model_executor.layers.quantization import QuantizationConfig -from vllm.logger import init_logger from vllm.model_executor.layers.rotary_embedding import MRotaryEmbedding, get_rope from vllm.model_executor.models.qwen3_next import Qwen3NextAttention from vllm.platforms import current_platform @@ -63,7 +63,6 @@ from ..common.qsa_cache import QSAForwardMetadata from . import model from .indexer_qsa import QSAIndexer - logger = init_logger(__name__) # --- FP8 K/V scale calibration collector ----------------------- diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index ae9eeb6..3b942cb 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -39,6 +39,7 @@ from vllm.v1.core.encoder_cache_manager import ( from vllm.v1.core.kv_cache_manager import KVCacheBlocks, KVCacheManager from vllm.v1.core.kv_cache_metrics import KVCacheMetricsCollector from vllm.v1.core.kv_cache_utils import KVCacheBlock +from vllm.v1.core.sched import e2_alloc_log as _e2_alloc_log from vllm.v1.core.sched.interface import PauseState, SchedulerInterface from vllm.v1.core.sched.output import ( CachedRequestData, @@ -70,11 +71,10 @@ from vllm.v1.metrics.stats import ( SchedulerStats, ) from vllm.v1.outputs import DraftTokenIds, KVConnectorOutput, ModelRunnerOutput -from vllm.v1.core.sched import e2_alloc_log as _e2_alloc_log from vllm.v1.request import Request, RequestStatus, StreamingUpdate from vllm.v1.spec_decode.dynamic.utils import build_dynamic_sd_schedule_lookup from vllm.v1.spec_decode.metrics import SpecDecodingStats -from vllm.v1.structured_output import StructuredOutputGrammar, StructuredOutputManager +from vllm.v1.structured_output import StructuredOutputManager from vllm.v1.structured_output.utils import strip_speculative_padding from vllm.v1.utils import record_function_or_nullcontext diff --git a/vllm/v1/structured_output/utils.py b/vllm/v1/structured_output/utils.py index 7c4af7a..92122f8 100644 --- a/vllm/v1/structured_output/utils.py +++ b/vllm/v1/structured_output/utils.py @@ -573,16 +573,3 @@ def choice_as_grammar(choice: list[str]) -> str: escaped_choices = (escape_ebnf_string(c) for c in choice) grammar = "root ::= " + " | ".join(f'"{c}"' for c in escaped_choices) return grammar - -def strip_speculative_padding(token_ids: list[int]) -> list[int]: - """Drop speculative-decoding padding from a token block. - - ngram and other speculative backends pad rejected draft positions with a - -1 sentinel. Structured-output grammars treat every entry as a real token - id, so the sentinels (and everything after the first one) are removed here, - before tokens reach any backend, rather than inside a single backend. - """ - for i, token_id in enumerate(token_ids): - if token_id < 0: - return token_ids[:i] - return token_ids -- 2.43.0