vllm.models.qwen4_exp.nvidia.ngram_embedding

Qwen4Exp n-gram embeddings with device and pinned-host storage.

Classes:

• Qwen4ExpNGramEmbedding –
• Qwen4ExpPLEDeviceEmbedding –

  PLE table allocated on the active model device.

• Qwen4ExpPLEEmbedding –

  ETP-sharded PLE table shared by device and pinned-host backends.

• Qwen4ExpPLEEmbeddingMethod –

  Quantization interface shared by resident and pinned PLE tables.

• Qwen4ExpPLEFp8EmbeddingMethod –

  FP8 PLE embedding with one global checkpoint scale.

• Qwen4ExpPLEPinnedHostEmbedding –

  PLE table loaded into pinned CPU memory and looked up through UVA.

• Qwen4ExpPLEUnquantizedEmbeddingMethod –

  Unquantized PLE embedding storage and lookup semantics.

Qwen4ExpNGramEmbedding

Bases: Module

Methods:

• compute_ngram_ids –

  Compute n-gram embedding indices for the current request layout.

• forward –

  Return full embeddings or a token shard after combining vocabulary owners.

• load_weights –

  Load hash buffers and checkpoint-split embedding rows.

• start_prefetch –

  Start the pinned lookup while the preceding decoder layer runs.

Source code in vllm/models/qwen4_exp/nvidia/ngram_embedding.py

class Qwen4ExpNGramEmbedding(nn.Module):
    _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,
        max_total_tokens: int,
        *,
        data_parallel_rank: int,
        prefix: str,
        quant_config: QuantizationConfig | None = None,
        params_dtype: torch.dtype | None = None,
        use_sequence_parallel: bool = False,
    ) -> None:
        super().__init__()
        self.use_sequence_parallel = use_sequence_parallel
        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
        embedding_prefix = f"{prefix}.ngram_embedding"
        embedding_quant_method = Qwen4ExpPLEEmbeddingMethod.from_quant_config(
            quant_config,
            embedding_prefix,
            getattr(config, "ple_embedding_dtype", None),
        )
        if params_dtype is None:
            params_dtype = torch.get_default_dtype()
        engram_config = get_current_vllm_config().engram_config
        embedding_cls = (
            Qwen4ExpPLEPinnedHostEmbedding
            if engram_config is not None and engram_config.cpu_offload
            else Qwen4ExpPLEDeviceEmbedding
        )
        # A TP reduce-scatter covers all vocabulary owners only when ETP=TP.
        self.use_reduce_scatter = (
            self.use_sequence_parallel
            and get_etp_group().world_size == get_tp_group().world_size
        )
        self.ngram_embedding = embedding_cls(
            padded_vocab_size,
            self.head_dim,
            params_dtype=params_dtype,
            padding_size=divisor,
            prefix=embedding_prefix,
            embedding_method=embedding_quant_method,
            num_ngram_heads=self.ngram_heads,
            max_total_tokens=max_total_tokens,
            data_parallel_rank=data_parallel_rank,
            reduce_results=not self.use_reduce_scatter,
        )
        if self.ngram_embedding.supports_prefetch:
            # The side-stream lookup outlives eager-break args, whose
            # graph-pool storage later segments may reuse.
            self._prefetch_ids = torch.empty(
                max_total_tokens, self.ngram_heads, dtype=torch.long
            )
        weight = self.ngram_embedding.weight
        logger.info(
            "Initialized PLE embedding %s: quantization_method=%s, "
            "weight_dtype=%s, weight_device=%s, pinned=%s",
            embedding_prefix,
            type(embedding_quant_method).__name__,
            weight.dtype,
            weight.device,
            weight.is_pinned(),
        )

    @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()
        positions = torch.arange(num_tokens, device=input_ids.device, dtype=torch.int64)
        packed = torch.full(
            (num_reqs, num_tokens),
            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, columns] = input_ids
        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 _sp_shard_embeddings(self, embeddings: torch.Tensor) -> torch.Tensor:
        """Shard embedding token rows, optionally reducing vocabulary contributions."""
        # Each embedding element has one vocabulary owner, including FP8 bytes.
        output_dtype = embeddings.dtype
        if is_fp8(embeddings):
            embeddings = embeddings.view(torch.int8)
        if self.use_reduce_scatter:
            embeddings = sp_reduce_scatter(embeddings)
        else:
            embeddings = sp_shard(embeddings)
        return embeddings.view(output_dtype)

    def forward(
        self,
        hidden_states: torch.Tensor,
        input_ids: torch.Tensor,
        query_start_loc: torch.Tensor,
        ngram_context: torch.Tensor,
    ) -> torch.Tensor:
        """Return full embeddings or a token shard after combining vocabulary owners."""
        embedding = self.ngram_embedding
        if embedding.supports_prefetch:
            # Prefetch contains full tokens even when hidden states are sharded.
            embeddings = embedding.finalize_prefetch(total_tokens=input_ids.shape[0])
        else:
            ngram_ids = self.compute_ngram_ids(
                input_ids, query_start_loc, ngram_context
            )
            embeddings = embedding(ngram_ids).flatten(-2)
        if not self.use_sequence_parallel:
            return embeddings
        return self._sp_shard_embeddings(embeddings)

    def start_prefetch(
        self,
        input_ids: torch.Tensor,
        query_start_loc: torch.Tensor,
        ngram_context: torch.Tensor,
    ) -> None:
        """Start the pinned lookup while the preceding decoder layer runs."""
        embedding = self.ngram_embedding
        if not embedding.supports_prefetch:
            return
        ngram_ids = self.compute_ngram_ids(
            input_ids,
            query_start_loc,
            ngram_context,
            output=self._prefetch_ids[: input_ids.numel()],
        )
        embedding.start_prefetch(ngram_ids)

    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
        """Load hash buffers and checkpoint-split embedding rows."""
        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_"

        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 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 regular_weights:
            loaded.update(AutoWeightsLoader(self).load_weights(regular_weights))
        return loaded

_is_prime_64(value) staticmethod

Return whether a 64-bit integer is prime.

Source code in vllm/models/qwen4_exp/nvidia/ngram_embedding.py

@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

_make_layer_multipliers(*, ngram_size, unigram_vocab_size, seed, ple_dense_layer_id) classmethod

Build deterministic hash multipliers for one PLE layer.

Source code in vllm/models/qwen4_exp/nvidia/ngram_embedding.py

@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

_make_vocab_layout(*, ngram_vocab_size_base, ngram_heads, ple_dense_layer_id) classmethod

Build per-head vocabulary sizes, offsets, and total row count.

Source code in vllm/models/qwen4_exp/nvidia/ngram_embedding.py

@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

_nth_prime_after(start, count) classmethod

Return the count-th prime strictly greater than start.

Source code in vllm/models/qwen4_exp/nvidia/ngram_embedding.py

@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

_sp_shard_embeddings(embeddings)

Shard embedding token rows, optionally reducing vocabulary contributions.

Source code in vllm/models/qwen4_exp/nvidia/ngram_embedding.py

def _sp_shard_embeddings(self, embeddings: torch.Tensor) -> torch.Tensor:
    """Shard embedding token rows, optionally reducing vocabulary contributions."""
    # Each embedding element has one vocabulary owner, including FP8 bytes.
    output_dtype = embeddings.dtype
    if is_fp8(embeddings):
        embeddings = embeddings.view(torch.int8)
    if self.use_reduce_scatter:
        embeddings = sp_reduce_scatter(embeddings)
    else:
        embeddings = sp_shard(embeddings)
    return embeddings.view(output_dtype)

_splitmix64(value) classmethod

Mix an integer into a deterministic unsigned 64-bit value.

Source code in vllm/models/qwen4_exp/nvidia/ngram_embedding.py

@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

compute_ngram_ids(input_ids, query_start_loc, ngram_context, output=None)

Compute n-gram embedding indices for the current request layout.

Source code in vllm/models/qwen4_exp/nvidia/ngram_embedding.py

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()
    positions = torch.arange(num_tokens, device=input_ids.device, dtype=torch.int64)
    packed = torch.full(
        (num_reqs, num_tokens),
        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, columns] = input_ids
    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)

forward(hidden_states, input_ids, query_start_loc, ngram_context)

Return full embeddings or a token shard after combining vocabulary owners.

Source code in vllm/models/qwen4_exp/nvidia/ngram_embedding.py

def forward(
    self,
    hidden_states: torch.Tensor,
    input_ids: torch.Tensor,
    query_start_loc: torch.Tensor,
    ngram_context: torch.Tensor,
) -> torch.Tensor:
    """Return full embeddings or a token shard after combining vocabulary owners."""
    embedding = self.ngram_embedding
    if embedding.supports_prefetch:
        # Prefetch contains full tokens even when hidden states are sharded.
        embeddings = embedding.finalize_prefetch(total_tokens=input_ids.shape[0])
    else:
        ngram_ids = self.compute_ngram_ids(
            input_ids, query_start_loc, ngram_context
        )
        embeddings = embedding(ngram_ids).flatten(-2)
    if not self.use_sequence_parallel:
        return embeddings
    return self._sp_shard_embeddings(embeddings)

load_weights(weights)

Load hash buffers and checkpoint-split embedding rows.

Source code in vllm/models/qwen4_exp/nvidia/ngram_embedding.py

def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
    """Load hash buffers and checkpoint-split embedding rows."""
    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_"

    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 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 regular_weights:
        loaded.update(AutoWeightsLoader(self).load_weights(regular_weights))
    return loaded

start_prefetch(input_ids, query_start_loc, ngram_context)

Start the pinned lookup while the preceding decoder layer runs.

Source code in vllm/models/qwen4_exp/nvidia/ngram_embedding.py

def start_prefetch(
    self,
    input_ids: torch.Tensor,
    query_start_loc: torch.Tensor,
    ngram_context: torch.Tensor,
) -> None:
    """Start the pinned lookup while the preceding decoder layer runs."""
    embedding = self.ngram_embedding
    if not embedding.supports_prefetch:
        return
    ngram_ids = self.compute_ngram_ids(
        input_ids,
        query_start_loc,
        ngram_context,
        output=self._prefetch_ids[: input_ids.numel()],
    )
    embedding.start_prefetch(ngram_ids)

Qwen4ExpPLEDeviceEmbedding

Bases: Qwen4ExpPLEEmbedding

PLE table allocated on the active model device.

Methods:

• allocate_embedding_weight –

  Allocate the complete PLE weight on the active device.

• finalize_prefetch –

  Reject prefetch finalization for resident embeddings.

• forward –

  Gather ETP IDs, optionally reduce lookup results, and select local rows.

• start_prefetch –

  Resident embedding prefetch is a no-op.

Source code in vllm/models/qwen4_exp/common/ngram_embedding.py

class Qwen4ExpPLEDeviceEmbedding(Qwen4ExpPLEEmbedding):
    """PLE table allocated on the active model device."""

    def allocate_embedding_weight(
        self,
        num_embeddings: int,
        embedding_dim: int,
        dtype: torch.dtype,
    ) -> torch.Tensor:
        """Allocate the complete PLE weight on the active device."""
        return torch.empty(num_embeddings, embedding_dim, dtype=dtype)

    def start_prefetch(self, ngram_ids: torch.Tensor) -> None:
        """Resident embedding prefetch is a no-op."""
        return None

    def finalize_prefetch(self, total_tokens: int) -> torch.Tensor:
        """Reject prefetch finalization for resident embeddings."""
        raise RuntimeError("Device PLE embedding does not support prefetch")

    def forward(self, ngram_ids: torch.Tensor) -> torch.Tensor:
        """Gather ETP IDs, optionally reduce lookup results, and select local rows."""
        slot_size, slot_offset = self._get_dp_gather_slot(ngram_ids.shape[0])
        gathered_ids = self._gather_dp_ids(ngram_ids, slot_size)
        embeddings = super().forward(gathered_ids)
        return self._select_embeddings(
            embeddings,
            ngram_ids.shape[0],
            slot_offset,
        )

allocate_embedding_weight(num_embeddings, embedding_dim, dtype)

Allocate the complete PLE weight on the active device.

Source code in vllm/models/qwen4_exp/common/ngram_embedding.py

def allocate_embedding_weight(
    self,
    num_embeddings: int,
    embedding_dim: int,
    dtype: torch.dtype,
) -> torch.Tensor:
    """Allocate the complete PLE weight on the active device."""
    return torch.empty(num_embeddings, embedding_dim, dtype=dtype)

finalize_prefetch(total_tokens)

Reject prefetch finalization for resident embeddings.

Source code in vllm/models/qwen4_exp/common/ngram_embedding.py

def finalize_prefetch(self, total_tokens: int) -> torch.Tensor:
    """Reject prefetch finalization for resident embeddings."""
    raise RuntimeError("Device PLE embedding does not support prefetch")

forward(ngram_ids)

Gather ETP IDs, optionally reduce lookup results, and select local rows.

Source code in vllm/models/qwen4_exp/common/ngram_embedding.py

def forward(self, ngram_ids: torch.Tensor) -> torch.Tensor:
    """Gather ETP IDs, optionally reduce lookup results, and select local rows."""
    slot_size, slot_offset = self._get_dp_gather_slot(ngram_ids.shape[0])
    gathered_ids = self._gather_dp_ids(ngram_ids, slot_size)
    embeddings = super().forward(gathered_ids)
    return self._select_embeddings(
        embeddings,
        ngram_ids.shape[0],
        slot_offset,
    )

start_prefetch(ngram_ids)

Resident embedding prefetch is a no-op.

Source code in vllm/models/qwen4_exp/common/ngram_embedding.py

def start_prefetch(self, ngram_ids: torch.Tensor) -> None:
    """Resident embedding prefetch is a no-op."""
    return None

Qwen4ExpPLEEmbedding

Bases: PLEVocabParallelEmbedding, ABC

ETP-sharded PLE table shared by device and pinned-host backends.

Methods:

• allocate_embedding_weight –

  Allocate storage for the complete embedding weight.

• dequantize –

  Delegate storage-format conversion to the embedding method.

• finalize_prefetch –

  Return prefetched embeddings for the full input token set.

• start_prefetch –

  Start an asynchronous lookup when supported.

Source code in vllm/models/qwen4_exp/common/ngram_embedding.py

class Qwen4ExpPLEEmbedding(PLEVocabParallelEmbedding, ABC):
    """ETP-sharded PLE table shared by device and pinned-host backends."""

    supports_prefetch: ClassVar[bool] = False

    def __init__(
        self,
        num_embeddings: int,
        embedding_dim: int,
        *,
        params_dtype: torch.dtype,
        padding_size: int,
        prefix: str,
        embedding_method: "Qwen4ExpPLEEmbeddingMethod",
        num_ngram_heads: int = 1,
        max_total_tokens: int = 0,
        data_parallel_rank: int = 0,
        reduce_results: bool = True,
    ) -> None:
        del num_ngram_heads, max_total_tokens
        super().__init__(
            num_embeddings,
            embedding_dim,
            params_dtype=params_dtype,
            padding_size=padding_size,
            prefix=prefix,
            quant_method=embedding_method,
            parallel_group=get_etp_group(),
            reduce_results=reduce_results,
        )
        self.embedding_method = embedding_method
        self.data_parallel_rank = data_parallel_rank
        tp_size = get_tp_group().world_size
        if self.tp_size % tp_size:
            raise ValueError(
                "ETP size must be divisible by TP size, but got "
                f"ETP={self.tp_size} and TP={tp_size}"
            )
        self.etp_data_parallel_size = self.tp_size // tp_size

    @abstractmethod
    def allocate_embedding_weight(
        self,
        num_embeddings: int,
        embedding_dim: int,
        dtype: torch.dtype,
    ) -> torch.Tensor:
        """Allocate storage for the complete embedding weight."""
        raise NotImplementedError

    def dequantize(
        self,
        embeddings: torch.Tensor,
        output_dtype: torch.dtype,
    ) -> torch.Tensor:
        """Delegate storage-format conversion to the embedding method."""
        return self.embedding_method.dequantize(self, embeddings, output_dtype)

    def _get_dp_gather_slot(self, local_num_tokens: int) -> tuple[int, int]:
        """Return the per-DP slot size and this rank's slot offset."""
        if self.etp_data_parallel_size == 1:
            return local_num_tokens, 0
        dp_metadata: DPMetadata | None = get_forward_context().dp_metadata
        if dp_metadata is None:
            raise RuntimeError("ETP spanning DP requires DP token metadata")
        group_start = (self.data_parallel_rank // self.etp_data_parallel_size) * (
            self.etp_data_parallel_size
        )
        group_end = group_start + self.etp_data_parallel_size
        token_counts = dp_metadata.num_tokens_across_dp_cpu.tolist()
        group_counts = token_counts[group_start:group_end]
        slot_size = max(group_counts)
        dp_rank = get_dp_group().rank_in_group
        return slot_size, dp_rank * slot_size

    def _gather_dp_ids(
        self,
        ngram_ids: torch.Tensor,
        slot_size: int,
    ) -> torch.Tensor:
        """Gather DP-local IDs that share one ETP-sharded PLE table."""
        if self.etp_data_parallel_size == 1:
            return ngram_ids
        if ngram_ids.shape[0] < slot_size:
            padding = ngram_ids.new_zeros(
                slot_size - ngram_ids.shape[0], ngram_ids.shape[1]
            )
            ngram_ids = torch.cat((ngram_ids, padding), dim=0)
        return get_dp_group().all_gather(ngram_ids, dim=0)

    def _select_embeddings(
        self,
        embeddings: torch.Tensor,
        local_num_tokens: int,
        slot_offset: int,
    ) -> torch.Tensor:
        """Select this DP rank's rows from the ETP-reduced embeddings."""
        if self.etp_data_parallel_size == 1:
            return embeddings
        return embeddings.narrow(0, slot_offset, local_num_tokens)

    @abstractmethod
    def start_prefetch(self, ngram_ids: torch.Tensor) -> None:
        """Start an asynchronous lookup when supported."""
        raise NotImplementedError

    @abstractmethod
    def finalize_prefetch(self, total_tokens: int) -> torch.Tensor:
        """Return prefetched embeddings for the full input token set."""
        raise NotImplementedError

_gather_dp_ids(ngram_ids, slot_size)

Gather DP-local IDs that share one ETP-sharded PLE table.

Source code in vllm/models/qwen4_exp/common/ngram_embedding.py

def _gather_dp_ids(
    self,
    ngram_ids: torch.Tensor,
    slot_size: int,
) -> torch.Tensor:
    """Gather DP-local IDs that share one ETP-sharded PLE table."""
    if self.etp_data_parallel_size == 1:
        return ngram_ids
    if ngram_ids.shape[0] < slot_size:
        padding = ngram_ids.new_zeros(
            slot_size - ngram_ids.shape[0], ngram_ids.shape[1]
        )
        ngram_ids = torch.cat((ngram_ids, padding), dim=0)
    return get_dp_group().all_gather(ngram_ids, dim=0)

_get_dp_gather_slot(local_num_tokens)

Return the per-DP slot size and this rank's slot offset.

Source code in vllm/models/qwen4_exp/common/ngram_embedding.py

def _get_dp_gather_slot(self, local_num_tokens: int) -> tuple[int, int]:
    """Return the per-DP slot size and this rank's slot offset."""
    if self.etp_data_parallel_size == 1:
        return local_num_tokens, 0
    dp_metadata: DPMetadata | None = get_forward_context().dp_metadata
    if dp_metadata is None:
        raise RuntimeError("ETP spanning DP requires DP token metadata")
    group_start = (self.data_parallel_rank // self.etp_data_parallel_size) * (
        self.etp_data_parallel_size
    )
    group_end = group_start + self.etp_data_parallel_size
    token_counts = dp_metadata.num_tokens_across_dp_cpu.tolist()
    group_counts = token_counts[group_start:group_end]
    slot_size = max(group_counts)
    dp_rank = get_dp_group().rank_in_group
    return slot_size, dp_rank * slot_size

_select_embeddings(embeddings, local_num_tokens, slot_offset)

Select this DP rank's rows from the ETP-reduced embeddings.

Source code in vllm/models/qwen4_exp/common/ngram_embedding.py

def _select_embeddings(
    self,
    embeddings: torch.Tensor,
    local_num_tokens: int,
    slot_offset: int,
) -> torch.Tensor:
    """Select this DP rank's rows from the ETP-reduced embeddings."""
    if self.etp_data_parallel_size == 1:
        return embeddings
    return embeddings.narrow(0, slot_offset, local_num_tokens)

allocate_embedding_weight(num_embeddings, embedding_dim, dtype) abstractmethod

Allocate storage for the complete embedding weight.

Source code in vllm/models/qwen4_exp/common/ngram_embedding.py

@abstractmethod
def allocate_embedding_weight(
    self,
    num_embeddings: int,
    embedding_dim: int,
    dtype: torch.dtype,
) -> torch.Tensor:
    """Allocate storage for the complete embedding weight."""
    raise NotImplementedError

dequantize(embeddings, output_dtype)

Delegate storage-format conversion to the embedding method.

Source code in vllm/models/qwen4_exp/common/ngram_embedding.py

def dequantize(
    self,
    embeddings: torch.Tensor,
    output_dtype: torch.dtype,
) -> torch.Tensor:
    """Delegate storage-format conversion to the embedding method."""
    return self.embedding_method.dequantize(self, embeddings, output_dtype)

finalize_prefetch(total_tokens) abstractmethod

Return prefetched embeddings for the full input token set.

Source code in vllm/models/qwen4_exp/common/ngram_embedding.py

@abstractmethod
def finalize_prefetch(self, total_tokens: int) -> torch.Tensor:
    """Return prefetched embeddings for the full input token set."""
    raise NotImplementedError

start_prefetch(ngram_ids) abstractmethod

Start an asynchronous lookup when supported.

Source code in vllm/models/qwen4_exp/common/ngram_embedding.py

@abstractmethod
def start_prefetch(self, ngram_ids: torch.Tensor) -> None:
    """Start an asynchronous lookup when supported."""
    raise NotImplementedError

Qwen4ExpPLEEmbeddingMethod

Bases: QuantizeMethodBase

Quantization interface shared by resident and pinned PLE tables.

Methods:

• dequantize –

  Convert looked-up PLE rows to the activation dtype.

• from_quant_config –

  Select the concrete PLE embedding format for a layer.

Source code in vllm/models/qwen4_exp/common/ngram_embedding.py

class Qwen4ExpPLEEmbeddingMethod(QuantizeMethodBase):
    """Quantization interface shared by resident and pinned PLE tables."""

    # PLE post-load processing only validates scales in their current storage.
    requires_device_loading: bool = False

    @staticmethod
    def from_quant_config(
        quant_config: QuantizationConfig | None,
        prefix: str,
        embedding_dtype: str | None = None,
    ) -> "Qwen4ExpPLEEmbeddingMethod":
        """Select the concrete PLE embedding format for a layer."""
        if embedding_dtype == "float8_e4m3fn":
            return Qwen4ExpPLEFp8EmbeddingMethod()
        if quant_config is None:
            return Qwen4ExpPLEUnquantizedEmbeddingMethod()
        if isinstance(quant_config, ModelOptMixedPrecisionConfig):
            if quant_config._resolve_quant_algo(prefix) == "FP8":
                return Qwen4ExpPLEFp8EmbeddingMethod()
            return Qwen4ExpPLEUnquantizedEmbeddingMethod()
        if isinstance(
            quant_config, ModelOptQuantConfigBase
        ) and quant_config.is_layer_excluded(prefix):
            return Qwen4ExpPLEUnquantizedEmbeddingMethod()
        if (
            isinstance(quant_config, CompressedTensorsConfig)
            and quant_config.get_scheme_dict(None, layer_name=prefix) is None
        ):
            return Qwen4ExpPLEUnquantizedEmbeddingMethod()
        if (
            isinstance(quant_config, INCConfig)
            and not quant_config.config_parser.resolve(None, prefix).quantized
        ):
            return Qwen4ExpPLEUnquantizedEmbeddingMethod()
        # Quark quantizes only Linear and MoE layers; PLE tables stay BF16.
        if isinstance(quant_config, QuarkConfig):
            return Qwen4ExpPLEUnquantizedEmbeddingMethod()
        if not isinstance(quant_config, Fp8Config):
            raise NotImplementedError(
                "Qwen4Exp PLE embedding does not support quantization config "
                f"{type(quant_config).__name__}"
            )

        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 Qwen4ExpPLEUnquantizedEmbeddingMethod()
        # 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 Qwen4ExpPLEUnquantizedEmbeddingMethod()
        if not quant_config.is_checkpoint_fp8_serialized:
            raise NotImplementedError(
                "Qwen4Exp PLE embedding only supports serialized FP8 checkpoints"
            )
        return Qwen4ExpPLEFp8EmbeddingMethod()

    def apply(
        self,
        layer: nn.Module,
        x: torch.Tensor,
        bias: torch.Tensor | None = None,
    ) -> torch.Tensor:
        raise NotImplementedError("PLE weights only support embedding lookup")

    def embedding(self, layer: nn.Module, input_: torch.Tensor) -> torch.Tensor:
        return F.embedding(input_, layer.weight)

    @abstractmethod
    def dequantize(
        self,
        layer: nn.Module,
        embeddings: torch.Tensor,
        output_dtype: torch.dtype,
    ) -> torch.Tensor:
        """Convert looked-up PLE rows to the activation dtype."""
        raise NotImplementedError

dequantize(layer, embeddings, output_dtype) abstractmethod

Convert looked-up PLE rows to the activation dtype.

Source code in vllm/models/qwen4_exp/common/ngram_embedding.py

@abstractmethod
def dequantize(
    self,
    layer: nn.Module,
    embeddings: torch.Tensor,
    output_dtype: torch.dtype,
) -> torch.Tensor:
    """Convert looked-up PLE rows to the activation dtype."""
    raise NotImplementedError

from_quant_config(quant_config, prefix, embedding_dtype=None) staticmethod

Select the concrete PLE embedding format for a layer.

Source code in vllm/models/qwen4_exp/common/ngram_embedding.py

@staticmethod
def from_quant_config(
    quant_config: QuantizationConfig | None,
    prefix: str,
    embedding_dtype: str | None = None,
) -> "Qwen4ExpPLEEmbeddingMethod":
    """Select the concrete PLE embedding format for a layer."""
    if embedding_dtype == "float8_e4m3fn":
        return Qwen4ExpPLEFp8EmbeddingMethod()
    if quant_config is None:
        return Qwen4ExpPLEUnquantizedEmbeddingMethod()
    if isinstance(quant_config, ModelOptMixedPrecisionConfig):
        if quant_config._resolve_quant_algo(prefix) == "FP8":
            return Qwen4ExpPLEFp8EmbeddingMethod()
        return Qwen4ExpPLEUnquantizedEmbeddingMethod()
    if isinstance(
        quant_config, ModelOptQuantConfigBase
    ) and quant_config.is_layer_excluded(prefix):
        return Qwen4ExpPLEUnquantizedEmbeddingMethod()
    if (
        isinstance(quant_config, CompressedTensorsConfig)
        and quant_config.get_scheme_dict(None, layer_name=prefix) is None
    ):
        return Qwen4ExpPLEUnquantizedEmbeddingMethod()
    if (
        isinstance(quant_config, INCConfig)
        and not quant_config.config_parser.resolve(None, prefix).quantized
    ):
        return Qwen4ExpPLEUnquantizedEmbeddingMethod()
    # Quark quantizes only Linear and MoE layers; PLE tables stay BF16.
    if isinstance(quant_config, QuarkConfig):
        return Qwen4ExpPLEUnquantizedEmbeddingMethod()
    if not isinstance(quant_config, Fp8Config):
        raise NotImplementedError(
            "Qwen4Exp PLE embedding does not support quantization config "
            f"{type(quant_config).__name__}"
        )

    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 Qwen4ExpPLEUnquantizedEmbeddingMethod()
    # 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 Qwen4ExpPLEUnquantizedEmbeddingMethod()
    if not quant_config.is_checkpoint_fp8_serialized:
        raise NotImplementedError(
            "Qwen4Exp PLE embedding only supports serialized FP8 checkpoints"
        )
    return Qwen4ExpPLEFp8EmbeddingMethod()

Qwen4ExpPLEFp8EmbeddingMethod

Bases: Qwen4ExpPLEEmbeddingMethod

FP8 PLE embedding with one global checkpoint scale.

Methods:

• process_weights_after_loading –

  Reject FP8 PLE checkpoints without a global scale.

Source code in vllm/models/qwen4_exp/common/ngram_embedding.py

class Qwen4ExpPLEFp8EmbeddingMethod(Qwen4ExpPLEEmbeddingMethod):
    """FP8 PLE embedding with one global checkpoint scale."""

    def create_weights(
        self,
        layer: Qwen4ExpPLEEmbedding,
        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 = ModelWeightParameter(
            data=layer.allocate_embedding_weight(
                sum(output_partition_sizes),
                input_size_per_partition,
                torch.float8_e4m3fn,
            ),
            input_dim=1,
            output_dim=0,
            weight_loader=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 dequantize(
        self,
        layer: nn.Module,
        embeddings: torch.Tensor,
        output_dtype: torch.dtype,
    ) -> torch.Tensor:
        weight_scale = getattr(layer, "weight_scale", None)
        if weight_scale is None:
            raise RuntimeError("FP8 PLE embedding is missing its global scale")
        if weight_scale.device != embeddings.device:
            raise RuntimeError("FP8 PLE embedding scale must be on the output device")
        return embeddings.to(output_dtype) * weight_scale.to(output_dtype)

process_weights_after_loading(layer)

Reject FP8 PLE checkpoints without a global scale.

Source code in vllm/models/qwen4_exp/common/ngram_embedding.py

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")

Qwen4ExpPLEPinnedHostEmbedding

Bases: Qwen4ExpPLEEmbedding

PLE table loaded into pinned CPU memory and looked up through UVA.

Methods:

• allocate_embedding_weight –

  Allocate the complete PLE weight directly in pinned CPU memory.

• finalize_prefetch –

  Finalize the full-token pinned lookup with optional ETP reduction.

• start_prefetch –

  Gather ETP IDs and launch their UVA lookup on the side stream.

• sync_lookup –

  Synchronous UVA lookup for platforms without prefetch wiring.

Source code in vllm/models/qwen4_exp/common/ngram_embedding.py

class Qwen4ExpPLEPinnedHostEmbedding(Qwen4ExpPLEEmbedding):
    """PLE table loaded into pinned CPU memory and looked up through UVA."""

    supports_prefetch: ClassVar[bool] = True

    def __init__(
        self,
        num_embeddings: int,
        embedding_dim: int,
        *,
        params_dtype: torch.dtype,
        padding_size: int,
        prefix: str,
        embedding_method: Qwen4ExpPLEEmbeddingMethod,
        num_ngram_heads: int = 1,
        max_total_tokens: int = 0,
        data_parallel_rank: int = 0,
        reduce_results: bool = True,
    ) -> None:
        if not is_uva_available():
            raise RuntimeError("Engram CPU offload requires UVA support")
        super().__init__(
            num_embeddings,
            embedding_dim,
            params_dtype=params_dtype,
            padding_size=padding_size,
            prefix=prefix,
            embedding_method=embedding_method,
            num_ngram_heads=num_ngram_heads,
            max_total_tokens=max_total_tokens,
            data_parallel_rank=data_parallel_rank,
            reduce_results=reduce_results,
        )
        self._uva_weight = get_accelerator_view_from_cpu_tensor(self.weight)
        self._row_bytes = self.embedding_dim * self.weight.element_size()
        self._block_d = triton.next_power_of_2(self._row_bytes)
        self._prefetch_stream: torch.cuda.Stream | None = None
        self._prefetch_buffer: torch.Tensor | None = None
        self._prefetch_alloc_lock = threading.Lock()
        self._prefetch_rows = max_total_tokens * self.etp_data_parallel_size
        self._num_ngram_heads = num_ngram_heads
        self._output_dim = num_ngram_heads * self.embedding_dim

    def allocate_embedding_weight(
        self,
        num_embeddings: int,
        embedding_dim: int,
        dtype: torch.dtype,
    ) -> torch.Tensor:
        """Allocate the complete PLE weight directly in pinned CPU memory."""
        return torch.empty(
            num_embeddings,
            embedding_dim,
            dtype=dtype,
            device="cpu",
            pin_memory=True,
        )

    def _lookup(
        self,
        input_ids: torch.Tensor,
        output: torch.Tensor | None = None,
    ) -> torch.Tensor:
        """Look up local ETP rows while preserving the weight storage dtype."""
        expected_shape = (*input_ids.shape, self.embedding_dim)
        if output is None:
            output = torch.empty(
                expected_shape,
                dtype=self.weight.dtype,
                device=input_ids.device,
            )
        elif (
            tuple(output.shape) != expected_shape
            or output.dtype != self.weight.dtype
            or output.device != input_ids.device
        ):
            raise ValueError(
                "PLE prefetch output must match the input shape, weight dtype, "
                "and input device"
            )

        flat_ids = input_ids.reshape(-1).long()
        if flat_ids.numel():
            _lookup_ple_embedding_from_pinned_kernel[(flat_ids.numel(),)](
                self._uva_weight,
                flat_ids,
                output,
                self._row_bytes,
                self.shard_indices.org_vocab_start_index,
                self.shard_indices.org_vocab_end_index,
                BLOCK_D=self._block_d,
            )
        return output

    def sync_lookup(self, ngram_ids: torch.Tensor) -> torch.Tensor:
        """Synchronous UVA lookup for platforms without prefetch wiring."""
        slot_size, slot_offset = self._get_dp_gather_slot(ngram_ids.shape[0])
        gathered_ids = self._gather_dp_ids(ngram_ids, slot_size)
        embeddings = self._lookup(gathered_ids)
        embeddings = self._reduce_etp_embeddings(embeddings)
        return self._select_embeddings(
            embeddings,
            ngram_ids.shape[0],
            slot_offset,
        )

    def _reduce_etp_embeddings(self, embeddings: torch.Tensor) -> torch.Tensor:
        """Combine pinned lookup results owned by different ETP ranks."""
        if self.tp_size == 1:
            return embeddings
        assert self.parallel_group is not None
        if embeddings.dtype in (torch.float8_e4m3fn, torch.float8_e5m2):
            # Each vocabulary row has one owner, so reduce the raw FP8 bytes.
            reduced = self.parallel_group.all_reduce(embeddings.view(torch.int8))
            return reduced.view(embeddings.dtype)
        return self.parallel_group.all_reduce(embeddings)

    @eager_break_during_capture
    def start_prefetch(self, ngram_ids: torch.Tensor) -> None:
        """Gather ETP IDs and launch their UVA lookup on the side stream."""
        buffer = self._prefetch_buffer
        if buffer is None:
            # First use allocates. The eager profile run always precedes
            # cudagraph capture, so allocation never happens mid-capture;
            # the lock keeps concurrent first callers from tearing the
            # stream/buffer pair.
            with self._prefetch_alloc_lock:
                buffer = self._prefetch_buffer
                if buffer is None:
                    if torch.cuda.is_current_stream_capturing():
                        raise RuntimeError(
                            "pinned PLE prefetch buffer must be allocated "
                            "eagerly, before cudagraph capture"
                        )
                    self._prefetch_stream = torch.cuda.Stream(
                        device=self._uva_weight.device
                    )
                    buffer = torch.empty(
                        self._prefetch_rows,
                        self._num_ngram_heads,
                        self.embedding_dim,
                        dtype=self.weight.dtype,
                        device=self._uva_weight.device,
                    )
                    self._prefetch_buffer = buffer
        prefetch_stream = self._prefetch_stream
        if prefetch_stream is None:
            raise RuntimeError("pinned PLE prefetch stream was not allocated")
        slot_size, _ = self._get_dp_gather_slot(ngram_ids.shape[0])
        gathered_ids = self._gather_dp_ids(ngram_ids, slot_size)
        if gathered_ids.shape[0] > buffer.shape[0]:
            raise ValueError(
                f"pinned PLE prefetch buffer holds {buffer.shape[0]} rows, "
                f"but the batch needs {gathered_ids.shape[0]}"
            )
        active_output = buffer[: gathered_ids.shape[0]]
        prefetch_stream.wait_stream(torch.cuda.current_stream())
        gathered_ids.record_stream(prefetch_stream)
        with torch.cuda.stream(prefetch_stream):
            self._lookup(gathered_ids, output=active_output)

    @eager_break_during_capture
    def _finalize_prefetch(
        self,
        prefetch_output: torch.Tensor,
        output: torch.Tensor,
    ) -> None:
        """Join the side stream, optionally reduce ETP shards, and select rows."""
        prefetch_stream = self._prefetch_stream
        if prefetch_stream is None:
            raise RuntimeError("pinned PLE finalize requires a prior start_prefetch")
        torch.cuda.current_stream().wait_stream(prefetch_stream)
        slot_size, slot_offset = self._get_dp_gather_slot(output.shape[0])
        embeddings = prefetch_output[: slot_size * self.etp_data_parallel_size]
        if self.reduce_results:
            embeddings = self._reduce_etp_embeddings(embeddings)
        else:
            # The outer TP reduce-scatter requires ETP=TP.
            assert self.etp_data_parallel_size == 1
        embeddings = self._select_embeddings(
            embeddings,
            output.shape[0],
            slot_offset,
        )
        output.copy_(embeddings.flatten(-2))

    def finalize_prefetch(self, total_tokens: int) -> torch.Tensor:
        """Finalize the full-token pinned lookup with optional ETP reduction."""
        buffer = self._prefetch_buffer
        if buffer is None:
            raise RuntimeError("pinned PLE lookup requires a prior start_prefetch")
        # total_tokens counts this DP rank's tokens before SP sharding.
        output = buffer.new_empty((total_tokens, self._output_dim))
        self._finalize_prefetch(buffer, output)
        return output

    def forward(self, input_: torch.Tensor) -> torch.Tensor:
        raise RuntimeError(
            "Pinned-host PLE embedding requires "
            "start_prefetch() and finalize_prefetch()."
        )

_finalize_prefetch(prefetch_output, output)

Join the side stream, optionally reduce ETP shards, and select rows.

Source code in vllm/models/qwen4_exp/common/ngram_embedding.py

@eager_break_during_capture
def _finalize_prefetch(
    self,
    prefetch_output: torch.Tensor,
    output: torch.Tensor,
) -> None:
    """Join the side stream, optionally reduce ETP shards, and select rows."""
    prefetch_stream = self._prefetch_stream
    if prefetch_stream is None:
        raise RuntimeError("pinned PLE finalize requires a prior start_prefetch")
    torch.cuda.current_stream().wait_stream(prefetch_stream)
    slot_size, slot_offset = self._get_dp_gather_slot(output.shape[0])
    embeddings = prefetch_output[: slot_size * self.etp_data_parallel_size]
    if self.reduce_results:
        embeddings = self._reduce_etp_embeddings(embeddings)
    else:
        # The outer TP reduce-scatter requires ETP=TP.
        assert self.etp_data_parallel_size == 1
    embeddings = self._select_embeddings(
        embeddings,
        output.shape[0],
        slot_offset,
    )
    output.copy_(embeddings.flatten(-2))

_lookup(input_ids, output=None)

Look up local ETP rows while preserving the weight storage dtype.

Source code in vllm/models/qwen4_exp/common/ngram_embedding.py

def _lookup(
    self,
    input_ids: torch.Tensor,
    output: torch.Tensor | None = None,
) -> torch.Tensor:
    """Look up local ETP rows while preserving the weight storage dtype."""
    expected_shape = (*input_ids.shape, self.embedding_dim)
    if output is None:
        output = torch.empty(
            expected_shape,
            dtype=self.weight.dtype,
            device=input_ids.device,
        )
    elif (
        tuple(output.shape) != expected_shape
        or output.dtype != self.weight.dtype
        or output.device != input_ids.device
    ):
        raise ValueError(
            "PLE prefetch output must match the input shape, weight dtype, "
            "and input device"
        )

    flat_ids = input_ids.reshape(-1).long()
    if flat_ids.numel():
        _lookup_ple_embedding_from_pinned_kernel[(flat_ids.numel(),)](
            self._uva_weight,
            flat_ids,
            output,
            self._row_bytes,
            self.shard_indices.org_vocab_start_index,
            self.shard_indices.org_vocab_end_index,
            BLOCK_D=self._block_d,
        )
    return output

_reduce_etp_embeddings(embeddings)

Combine pinned lookup results owned by different ETP ranks.

Source code in vllm/models/qwen4_exp/common/ngram_embedding.py

def _reduce_etp_embeddings(self, embeddings: torch.Tensor) -> torch.Tensor:
    """Combine pinned lookup results owned by different ETP ranks."""
    if self.tp_size == 1:
        return embeddings
    assert self.parallel_group is not None
    if embeddings.dtype in (torch.float8_e4m3fn, torch.float8_e5m2):
        # Each vocabulary row has one owner, so reduce the raw FP8 bytes.
        reduced = self.parallel_group.all_reduce(embeddings.view(torch.int8))
        return reduced.view(embeddings.dtype)
    return self.parallel_group.all_reduce(embeddings)

allocate_embedding_weight(num_embeddings, embedding_dim, dtype)

Allocate the complete PLE weight directly in pinned CPU memory.

Source code in vllm/models/qwen4_exp/common/ngram_embedding.py

def allocate_embedding_weight(
    self,
    num_embeddings: int,
    embedding_dim: int,
    dtype: torch.dtype,
) -> torch.Tensor:
    """Allocate the complete PLE weight directly in pinned CPU memory."""
    return torch.empty(
        num_embeddings,
        embedding_dim,
        dtype=dtype,
        device="cpu",
        pin_memory=True,
    )

finalize_prefetch(total_tokens)

Finalize the full-token pinned lookup with optional ETP reduction.

Source code in vllm/models/qwen4_exp/common/ngram_embedding.py

def finalize_prefetch(self, total_tokens: int) -> torch.Tensor:
    """Finalize the full-token pinned lookup with optional ETP reduction."""
    buffer = self._prefetch_buffer
    if buffer is None:
        raise RuntimeError("pinned PLE lookup requires a prior start_prefetch")
    # total_tokens counts this DP rank's tokens before SP sharding.
    output = buffer.new_empty((total_tokens, self._output_dim))
    self._finalize_prefetch(buffer, output)
    return output

start_prefetch(ngram_ids)

Gather ETP IDs and launch their UVA lookup on the side stream.

Source code in vllm/models/qwen4_exp/common/ngram_embedding.py

@eager_break_during_capture
def start_prefetch(self, ngram_ids: torch.Tensor) -> None:
    """Gather ETP IDs and launch their UVA lookup on the side stream."""
    buffer = self._prefetch_buffer
    if buffer is None:
        # First use allocates. The eager profile run always precedes
        # cudagraph capture, so allocation never happens mid-capture;
        # the lock keeps concurrent first callers from tearing the
        # stream/buffer pair.
        with self._prefetch_alloc_lock:
            buffer = self._prefetch_buffer
            if buffer is None:
                if torch.cuda.is_current_stream_capturing():
                    raise RuntimeError(
                        "pinned PLE prefetch buffer must be allocated "
                        "eagerly, before cudagraph capture"
                    )
                self._prefetch_stream = torch.cuda.Stream(
                    device=self._uva_weight.device
                )
                buffer = torch.empty(
                    self._prefetch_rows,
                    self._num_ngram_heads,
                    self.embedding_dim,
                    dtype=self.weight.dtype,
                    device=self._uva_weight.device,
                )
                self._prefetch_buffer = buffer
    prefetch_stream = self._prefetch_stream
    if prefetch_stream is None:
        raise RuntimeError("pinned PLE prefetch stream was not allocated")
    slot_size, _ = self._get_dp_gather_slot(ngram_ids.shape[0])
    gathered_ids = self._gather_dp_ids(ngram_ids, slot_size)
    if gathered_ids.shape[0] > buffer.shape[0]:
        raise ValueError(
            f"pinned PLE prefetch buffer holds {buffer.shape[0]} rows, "
            f"but the batch needs {gathered_ids.shape[0]}"
        )
    active_output = buffer[: gathered_ids.shape[0]]
    prefetch_stream.wait_stream(torch.cuda.current_stream())
    gathered_ids.record_stream(prefetch_stream)
    with torch.cuda.stream(prefetch_stream):
        self._lookup(gathered_ids, output=active_output)

sync_lookup(ngram_ids)

Synchronous UVA lookup for platforms without prefetch wiring.

Source code in vllm/models/qwen4_exp/common/ngram_embedding.py

def sync_lookup(self, ngram_ids: torch.Tensor) -> torch.Tensor:
    """Synchronous UVA lookup for platforms without prefetch wiring."""
    slot_size, slot_offset = self._get_dp_gather_slot(ngram_ids.shape[0])
    gathered_ids = self._gather_dp_ids(ngram_ids, slot_size)
    embeddings = self._lookup(gathered_ids)
    embeddings = self._reduce_etp_embeddings(embeddings)
    return self._select_embeddings(
        embeddings,
        ngram_ids.shape[0],
        slot_offset,
    )

Qwen4ExpPLEUnquantizedEmbeddingMethod

Bases: Qwen4ExpPLEEmbeddingMethod

Unquantized PLE embedding storage and lookup semantics.

Source code in vllm/models/qwen4_exp/common/ngram_embedding.py

class Qwen4ExpPLEUnquantizedEmbeddingMethod(Qwen4ExpPLEEmbeddingMethod):
    """Unquantized PLE embedding storage and lookup semantics."""

    def create_weights(
        self,
        layer: Qwen4ExpPLEEmbedding,
        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
        weight = nn.Parameter(
            layer.allocate_embedding_weight(
                sum(output_partition_sizes),
                input_size_per_partition,
                params_dtype,
            ),
            requires_grad=False,
        )
        set_weight_attrs(weight, {"input_dim": 1, "output_dim": 0})
        set_weight_attrs(weight, extra_weight_attrs)
        layer.register_parameter("weight", weight)

    def dequantize(
        self,
        layer: nn.Module,
        embeddings: torch.Tensor,
        output_dtype: torch.dtype,
    ) -> torch.Tensor:
        del layer, output_dtype
        return embeddings