iHore/vLLM-4x-RTX-3090
8
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
1281bf9c037a3e9d5ee816003cd2f449d6fb5915
vLLM-4x-RTX-3090/patches/0029-rebase-remove-v0.30.0-merge-artefacts-import-order.patch
T
MOVZX 777e910dd0 Initial commit 
Signed-off-by: MOVZX <movzx@yahoo.com>
2026-10-15 01:01:52 +07:00

1073 lines
44 KiB
Diff
Raw Blame History

From 81807002fb36e2b9eab05728ea3bdf54fd12c280 Mon Sep 17 00:00:00 2001
From: halt95 <halt95@users.noreply.github.com>
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
Reference in New Issue View Git Blame Copy Permalink