vllm.models.qwen4_exp.nvidia.ngram_embedding

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

Classes:

• Qwen4ExpNGramEmbedding –
• Qwen4ExpPLEDeviceEmbedding –

  PLE table allocated on the active model device.

• Qwen4ExpPLEDiskOffloadEmbedding –

  PLE rows read on the host from the checkpoint shard files.

• 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.

• Qwen4ExpPLENvFp4EmbeddingMethod –

  NVFP4 PLE embedding: packed E2M1 rows, block scales, one global 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 = round_up(total_vocab_size, 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 = (
            Qwen4ExpPLEDiskOffloadEmbedding
            if engram_config is not None and engram_config.disk_offload
            else 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.
        # Disk-offload ranks hold full rows, so they only shard tokens.
        self.use_reduce_scatter = (
            self.use_sequence_parallel
            and get_etp_group().world_size == get_tp_group().world_size
            and embedding_cls is not Qwen4ExpPLEDiskOffloadEmbedding
        )
        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_config=%s, quantization_method=%s, "
            "weight_dtype=%s, weight_device=%s, pinned=%s",
            embedding_prefix,
            type(quant_config).__name__,
            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])
        elif isinstance(embedding, Qwen4ExpPLEDiskOffloadEmbedding):
            embeddings = embedding.staged_rows(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_"
        embedding = self.ngram_embedding
        method = getattr(embedding, "embedding_method", None)
        nvfp4_method = (
            method if isinstance(method, Qwen4ExpPLENvFp4EmbeddingMethod) else None
        )
        # NVFP4 tables split their block scales into shards like the rows.
        shard_parameters = ("weight", "weight_scale") if nvfp4_method else ("weight",)
        shard_size = (
            embedding.org_vocab_size + self.split_ngram_parts - 1
        ) // self.split_ngram_parts

        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):
                shard_text, _, suffix = name[len(shard_prefix) :].partition(".")
                if not shard_text.isdigit() or suffix not in shard_parameters:
                    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}"
                    )
                checkpoint_start = shard_index * shard_size
                expected_rows = max(
                    0,
                    min(shard_size, embedding.org_vocab_size - checkpoint_start),
                )
                parameter = getattr(embedding, suffix)
                expected_shape = (expected_rows, parameter.shape[1])
                if tuple(loaded_weight.shape) != expected_shape:
                    raise ValueError(
                        f"Shape mismatch for PLE embedding shard {shard_index} "
                        f"{suffix}: expected {expected_shape}, got "
                        f"{tuple(loaded_weight.shape)}"
                    )
                # A cast would silently reinterpret packed E2M1 bytes.
                if nvfp4_method is not None and loaded_weight.dtype != parameter.dtype:
                    raise ValueError(
                        f"NVFP4 PLE shard {shard_index} {suffix} requires "
                        f"{parameter.dtype}, got {loaded_weight.dtype}"
                    )
                if isinstance(embedding, Qwen4ExpPLEDiskOffloadEmbedding):
                    embedding.accept_checkpoint_shard(
                        shard_index, loaded_weight, shard_size, suffix
                    )
                else:
                    parameter.weight_loader(
                        parameter,
                        loaded_weight,
                        checkpoint_start=checkpoint_start,
                    )
                loaded.add(f"ngram_embedding.{suffix}")
                continue
            if nvfp4_method is not None and name in (
                "ngram_embedding.weight",
                "ngram_embedding.weight_scale",
            ):
                raise ValueError(
                    f"NVFP4 PLE tables must be stored as {shard_prefix}<i>.weight "
                    f"and {shard_prefix}<i>.weight_scale shards, got unsharded {name}"
                )
            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])
    elif isinstance(embedding, Qwen4ExpPLEDiskOffloadEmbedding):
        embeddings = embedding.staged_rows(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_"
    embedding = self.ngram_embedding
    method = getattr(embedding, "embedding_method", None)
    nvfp4_method = (
        method if isinstance(method, Qwen4ExpPLENvFp4EmbeddingMethod) else None
    )
    # NVFP4 tables split their block scales into shards like the rows.
    shard_parameters = ("weight", "weight_scale") if nvfp4_method else ("weight",)
    shard_size = (
        embedding.org_vocab_size + self.split_ngram_parts - 1
    ) // self.split_ngram_parts

    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):
            shard_text, _, suffix = name[len(shard_prefix) :].partition(".")
            if not shard_text.isdigit() or suffix not in shard_parameters:
                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}"
                )
            checkpoint_start = shard_index * shard_size
            expected_rows = max(
                0,
                min(shard_size, embedding.org_vocab_size - checkpoint_start),
            )
            parameter = getattr(embedding, suffix)
            expected_shape = (expected_rows, parameter.shape[1])
            if tuple(loaded_weight.shape) != expected_shape:
                raise ValueError(
                    f"Shape mismatch for PLE embedding shard {shard_index} "
                    f"{suffix}: expected {expected_shape}, got "
                    f"{tuple(loaded_weight.shape)}"
                )
            # A cast would silently reinterpret packed E2M1 bytes.
            if nvfp4_method is not None and loaded_weight.dtype != parameter.dtype:
                raise ValueError(
                    f"NVFP4 PLE shard {shard_index} {suffix} requires "
                    f"{parameter.dtype}, got {loaded_weight.dtype}"
                )
            if isinstance(embedding, Qwen4ExpPLEDiskOffloadEmbedding):
                embedding.accept_checkpoint_shard(
                    shard_index, loaded_weight, shard_size, suffix
                )
            else:
                parameter.weight_loader(
                    parameter,
                    loaded_weight,
                    checkpoint_start=checkpoint_start,
                )
            loaded.add(f"ngram_embedding.{suffix}")
            continue
        if nvfp4_method is not None and name in (
            "ngram_embedding.weight",
            "ngram_embedding.weight_scale",
        ):
            raise ValueError(
                f"NVFP4 PLE tables must be stored as {shard_prefix}<i>.weight "
                f"and {shard_prefix}<i>.weight_scale shards, got unsharded {name}"
            )
        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

Qwen4ExpPLEDiskOffloadEmbedding

Bases: Qwen4ExpPLEEmbedding

PLE rows read on the host from the checkpoint shard files.

No table is allocated. Each step, stage_checkpoint_rows reads the rows of that step's n-gram ids from the safetensors files and copies them to fixed device buffers, which the forward reads, so CUDA graphs replay it. Every ETP rank stages full rows, so no ETP reduction is needed.

Methods:

• accept_checkpoint_shard –

  Keep a checkpoint shard as a zero-copy byte view of its storage.

• allocate_embedding_weight –

  Allocate no rows; they stay in the checkpoint files.

• bind_file_shards –

  Record each shard's file and offset and release the loader's views.

• stage_rows –

  Read the rows of _host_ids[:num_tokens] and copy them to staging.

• staged_rows –

  Return this step's staged rows, decoded if they are NVFP4.

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

class Qwen4ExpPLEDiskOffloadEmbedding(Qwen4ExpPLEEmbedding):
    """PLE rows read on the host from the checkpoint shard files.

    No table is allocated. Each step, `stage_checkpoint_rows` reads the rows
    of that step's n-gram ids from the safetensors files and copies them to
    fixed device buffers, which the forward reads, so CUDA graphs replay it.
    Every ETP rank stages full rows, so no ETP reduction is needed.
    """

    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:
        super().__init__(
            num_embeddings,
            embedding_dim,
            params_dtype=params_dtype,
            padding_size=padding_size,
            prefix=prefix,
            embedding_method=embedding_method,
            data_parallel_rank=data_parallel_rank,
            reduce_results=reduce_results,
        )
        self._dummy = get_current_vllm_config().load_config.load_format == "dummy"
        # NVFP4 rows are packed E2M1 bytes plus a separate plane of block scales.
        self._nvfp4 = isinstance(embedding_method, Qwen4ExpPLENvFp4EmbeddingMethod)
        device, pin = self.weight.device, self.weight.is_cuda
        rows = (max_total_tokens, num_ngram_heads)
        self._host_ids = torch.empty(
            rows, dtype=torch.long, device="cpu", pin_memory=pin
        )
        # Per plane: device staging, pinned host rows (written only after the
        # step's stream sync, so after the previous H2D) and distinct-row bytes.
        self._planes: dict[str, tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = {}
        for name in ("weight", "weight_scale") if self._nvfp4 else ("weight",):
            plane = getattr(self, name)
            shape, dtype = (*rows, plane.shape[1]), plane.dtype
            host = torch.empty(shape, dtype=dtype, device="cpu", pin_memory=pin)
            self._planes[name] = (
                torch.zeros(shape, dtype=dtype, device=device),
                host,
                torch.empty_like(host).view(torch.uint8).flatten(0, 1),
            )
        if self._nvfp4:
            self._output = torch.zeros(
                (*rows, embedding_dim), dtype=params_dtype, device=device
            )
            self._row_ids = torch.arange(rows[0] * rows[1], device=device)
        self._shards: dict[str, dict[int, torch.Tensor]] = {n: {} for n in self._planes}
        self._shard_size = 0
        self._bound = False
        self._fds: dict[str, int] = {}
        self._sources: dict[str, tuple[torch.Tensor, torch.Tensor]] = {}

    def allocate_embedding_weight(
        self,
        num_embeddings: int,
        embedding_dim: int,
        dtype: torch.dtype,
    ) -> torch.Tensor:
        """Allocate no rows; they stay in the checkpoint files."""
        return torch.empty(0, embedding_dim, dtype=dtype)

    def accept_checkpoint_shard(
        self,
        shard_index: int,
        loaded_weight: torch.Tensor,
        shard_size: int,
        name: str,
    ) -> None:
        """Keep a checkpoint shard as a zero-copy byte view of its storage."""
        if self._bound:
            raise RuntimeError("PLE disk offload does not support weight reload")
        dtype = getattr(self, name).dtype
        if loaded_weight.dtype != dtype:
            raise ValueError(
                f"PLE shard {name} dtype {loaded_weight.dtype} must match the "
                f"embedding dtype {dtype}; disk offload serves rows as stored"
            )
        if shard_index in self._shards[name]:
            raise ValueError(f"Duplicate PLE embedding shard {shard_index} {name}")
        self._shards[name][shard_index] = loaded_weight.view(torch.uint8)
        self._shard_size = shard_size

    def bind_file_shards(self) -> None:
        """Record each shard's file and offset and release the loader's views."""
        if self._bound or self._dummy:
            self._bound = True
            return
        for name, shards in self._shards.items():
            rows = sum(shard.shape[0] for shard in shards.values())
            if rows != self.org_vocab_size:
                raise ValueError(
                    f"PLE {name} shards cover {rows} of {self.org_vocab_size} rows"
                )
            fds, bases = torch.zeros(2, len(shards)).long()
            for index, shard in shards.items():
                # The loader hands over mmap views, so find their files here.
                path, offset = _maps_entry(shard.data_ptr())
                if not path or not os.path.isfile(path):
                    raise ValueError(
                        f"PLE embedding shard {index} {name} is not a "
                        f"file-backed view ({path or 'anonymous'})"
                    )
                if path not in self._fds:
                    self._fds[path] = os.open(path, os.O_RDONLY)
                fds[index], bases[index] = self._fds[path], offset
            self._sources[name] = (fds, bases)
        self._shards.clear()
        self._bound = True

    def copy_ids_to_host(self, ids: torch.Tensor) -> None:
        self._host_ids[: ids.shape[0]].copy_(ids, non_blocking=True)

    def stage_rows(self, num_tokens: int) -> None:
        """Read the rows of `_host_ids[:num_tokens]` and copy them to staging."""
        ids = self._host_ids[:num_tokens].flatten()
        if not self._fds:
            for _, host, _ in self._planes.values():
                host[:num_tokens].view(torch.uint8).zero_()
        else:
            ids, inverse = ids.unique(return_inverse=True)
            if ids.numel() and not 0 <= ids[0] <= ids[-1] < self.org_vocab_size:
                raise IndexError(f"PLE id out of range for {self.org_vocab_size} rows")
            shard = ids // self._shard_size
            local = ids - shard * self._shard_size
            for name, (_, host, rows) in self._planes.items():
                out = host[:num_tokens].flatten(0, 1).view(torch.uint8)
                self._read_rows(name, shard, local, rows[: ids.numel()])
                torch.index_select(rows[: ids.numel()], 0, inverse, out=out)
        for staging, host, _ in self._planes.values():
            staging[:num_tokens].copy_(host[:num_tokens], non_blocking=True)

    def _read_rows(
        self,
        name: str,
        shard: torch.Tensor,
        local: torch.Tensor,
        rows: torch.Tensor,
    ) -> None:
        """Read one plane's rows for sorted distinct ids into `rows`."""
        fds, bases = self._sources[name]
        page, row_bytes = mmap.PAGESIZE, rows.shape[1]
        start = bases[shard] + local * row_bytes
        first, last = start // page, (start + row_bytes - 1) // page
        run = torch.ones_like(first, dtype=torch.bool)
        run[1:] = (first[1:] > last[:-1] + 1) | (shard[1:] != shard[:-1])
        # Queue every page run before the serial reads.
        for fd, lo, hi in zip(
            fds[shard[run]].tolist(),
            first[run].tolist(),
            last[run.roll(-1)].tolist(),
        ):
            os.posix_fadvise(
                fd, lo * page, (hi - lo + 1) * page, os.POSIX_FADV_WILLNEED
            )
        for row, fd, offset in zip(rows.numpy(), fds[shard].tolist(), start.tolist()):
            if os.preadv(fd, [row], offset) != row_bytes:
                raise ValueError(
                    f"PLE shard file {os.readlink(f'/proc/self/fd/{fd}')} is "
                    f"truncated at offset {offset}"
                )

    def staged_rows(self, num_tokens: int) -> torch.Tensor:
        """Return this step's staged rows, decoded if they are NVFP4."""
        weight, *scale = (
            staging[:num_tokens] for staging, _, _ in self._planes.values()
        )
        if not self._nvfp4:
            return weight.flatten(-2)
        num_rows = weight.shape[0] * weight.shape[1]
        output = self._output[:num_tokens]
        # Staged rows are compact, so the i-th staged row has id i.
        Qwen4ExpPLENvFp4EmbeddingMethod.decode_rows(
            weight,
            scale[0],
            self.weight_scale_2,
            self._row_ids[:num_rows],
            output,
            0,
            num_rows,
        )
        return output.flatten(-2)

    def start_prefetch(self, ngram_ids: torch.Tensor) -> None:
        return None

    def finalize_prefetch(self, total_tokens: int) -> torch.Tensor:
        raise RuntimeError("PLE disk offload does not support prefetch")

_read_rows(name, shard, local, rows)

Read one plane's rows for sorted distinct ids into rows.

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

def _read_rows(
    self,
    name: str,
    shard: torch.Tensor,
    local: torch.Tensor,
    rows: torch.Tensor,
) -> None:
    """Read one plane's rows for sorted distinct ids into `rows`."""
    fds, bases = self._sources[name]
    page, row_bytes = mmap.PAGESIZE, rows.shape[1]
    start = bases[shard] + local * row_bytes
    first, last = start // page, (start + row_bytes - 1) // page
    run = torch.ones_like(first, dtype=torch.bool)
    run[1:] = (first[1:] > last[:-1] + 1) | (shard[1:] != shard[:-1])
    # Queue every page run before the serial reads.
    for fd, lo, hi in zip(
        fds[shard[run]].tolist(),
        first[run].tolist(),
        last[run.roll(-1)].tolist(),
    ):
        os.posix_fadvise(
            fd, lo * page, (hi - lo + 1) * page, os.POSIX_FADV_WILLNEED
        )
    for row, fd, offset in zip(rows.numpy(), fds[shard].tolist(), start.tolist()):
        if os.preadv(fd, [row], offset) != row_bytes:
            raise ValueError(
                f"PLE shard file {os.readlink(f'/proc/self/fd/{fd}')} is "
                f"truncated at offset {offset}"
            )

accept_checkpoint_shard(shard_index, loaded_weight, shard_size, name)

Keep a checkpoint shard as a zero-copy byte view of its storage.

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

def accept_checkpoint_shard(
    self,
    shard_index: int,
    loaded_weight: torch.Tensor,
    shard_size: int,
    name: str,
) -> None:
    """Keep a checkpoint shard as a zero-copy byte view of its storage."""
    if self._bound:
        raise RuntimeError("PLE disk offload does not support weight reload")
    dtype = getattr(self, name).dtype
    if loaded_weight.dtype != dtype:
        raise ValueError(
            f"PLE shard {name} dtype {loaded_weight.dtype} must match the "
            f"embedding dtype {dtype}; disk offload serves rows as stored"
        )
    if shard_index in self._shards[name]:
        raise ValueError(f"Duplicate PLE embedding shard {shard_index} {name}")
    self._shards[name][shard_index] = loaded_weight.view(torch.uint8)
    self._shard_size = shard_size

allocate_embedding_weight(num_embeddings, embedding_dim, dtype)

Allocate no rows; they stay in the checkpoint files.

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

def allocate_embedding_weight(
    self,
    num_embeddings: int,
    embedding_dim: int,
    dtype: torch.dtype,
) -> torch.Tensor:
    """Allocate no rows; they stay in the checkpoint files."""
    return torch.empty(0, embedding_dim, dtype=dtype)

bind_file_shards()

Record each shard's file and offset and release the loader's views.

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

def bind_file_shards(self) -> None:
    """Record each shard's file and offset and release the loader's views."""
    if self._bound or self._dummy:
        self._bound = True
        return
    for name, shards in self._shards.items():
        rows = sum(shard.shape[0] for shard in shards.values())
        if rows != self.org_vocab_size:
            raise ValueError(
                f"PLE {name} shards cover {rows} of {self.org_vocab_size} rows"
            )
        fds, bases = torch.zeros(2, len(shards)).long()
        for index, shard in shards.items():
            # The loader hands over mmap views, so find their files here.
            path, offset = _maps_entry(shard.data_ptr())
            if not path or not os.path.isfile(path):
                raise ValueError(
                    f"PLE embedding shard {index} {name} is not a "
                    f"file-backed view ({path or 'anonymous'})"
                )
            if path not in self._fds:
                self._fds[path] = os.open(path, os.O_RDONLY)
            fds[index], bases[index] = self._fds[path], offset
        self._sources[name] = (fds, bases)
    self._shards.clear()
    self._bound = True

stage_rows(num_tokens)

Read the rows of _host_ids[:num_tokens] and copy them to staging.

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

def stage_rows(self, num_tokens: int) -> None:
    """Read the rows of `_host_ids[:num_tokens]` and copy them to staging."""
    ids = self._host_ids[:num_tokens].flatten()
    if not self._fds:
        for _, host, _ in self._planes.values():
            host[:num_tokens].view(torch.uint8).zero_()
    else:
        ids, inverse = ids.unique(return_inverse=True)
        if ids.numel() and not 0 <= ids[0] <= ids[-1] < self.org_vocab_size:
            raise IndexError(f"PLE id out of range for {self.org_vocab_size} rows")
        shard = ids // self._shard_size
        local = ids - shard * self._shard_size
        for name, (_, host, rows) in self._planes.items():
            out = host[:num_tokens].flatten(0, 1).view(torch.uint8)
            self._read_rows(name, shard, local, rows[: ids.numel()])
            torch.index_select(rows[: ids.numel()], 0, inverse, out=out)
    for staging, host, _ in self._planes.values():
        staging[:num_tokens].copy_(host[:num_tokens], non_blocking=True)

staged_rows(num_tokens)

Return this step's staged rows, decoded if they are NVFP4.

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

def staged_rows(self, num_tokens: int) -> torch.Tensor:
    """Return this step's staged rows, decoded if they are NVFP4."""
    weight, *scale = (
        staging[:num_tokens] for staging, _, _ in self._planes.values()
    )
    if not self._nvfp4:
        return weight.flatten(-2)
    num_rows = weight.shape[0] * weight.shape[1]
    output = self._output[:num_tokens]
    # Staged rows are compact, so the i-th staged row has id i.
    Qwen4ExpPLENvFp4EmbeddingMethod.decode_rows(
        weight,
        scale[0],
        self.weight_scale_2,
        self._row_ids[:num_rows],
        output,
        0,
        num_rows,
    )
    return output.flatten(-2)

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.

• bind_file_shards –

  Bind the checkpoint files of a disk-offloaded table after loading.

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

    def bind_file_shards(self) -> None:
        """Bind the checkpoint files of a disk-offloaded table after loading."""
        return None

    @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

bind_file_shards()

Bind the checkpoint files of a disk-offloaded table after loading.

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

def bind_file_shards(self) -> None:
    """Bind the checkpoint files of a disk-offloaded table after loading."""
    return None

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.

• lookup_dtype –

  Return the dtype of rows that a pinned-host lookup writes.

• lookup_from_pinned –

  Copy the ETP-owned rows for ids from pinned host memory.

• process_weights_after_loading –

  Bind disk-offloaded tables once every checkpoint shard is loaded.

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 == "nvfp4":
            return Qwen4ExpPLENvFp4EmbeddingMethod()
        if embedding_dtype == "float8_e4m3fn":
            return Qwen4ExpPLEFp8EmbeddingMethod()
        if embedding_dtype is not None:
            raise ValueError(
                "ple_embedding_dtype must be 'nvfp4' or 'float8_e4m3fn', "
                f"got {embedding_dtype!r}"
            )
        if quant_config is None:
            return Qwen4ExpPLEUnquantizedEmbeddingMethod()
        if isinstance(quant_config, ModelOptQuantConfigBase):
            if quant_config.is_layer_excluded(prefix):
                return Qwen4ExpPLEUnquantizedEmbeddingMethod()
            if isinstance(quant_config, ModelOptMixedPrecisionConfig):
                quant_algo = quant_config._resolve_quant_algo(prefix)
                if quant_algo == "FP8":
                    return Qwen4ExpPLEFp8EmbeddingMethod()
                if quant_algo == "NVFP4":
                    return Qwen4ExpPLENvFp4EmbeddingMethod()
                if quant_algo is not None:
                    raise ValueError(
                        f"Unsupported quant_algo {quant_algo!r} for the PLE embedding "
                        f"{prefix!r}; expected FP8 or NVFP4"
                    )
                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 isinstance(quant_config, Fp8Config):
            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 quant_config.is_checkpoint_fp8_serialized:
                return Qwen4ExpPLEFp8EmbeddingMethod()
            raise NotImplementedError(
                "Qwen4Exp PLE embedding only supports serialized FP8 checkpoints"
            )

        logger.warning(
            "Falling back to unquantized PLE embedding for %s "
            "with quantization config %s.",
            prefix,
            type(quant_config).__name__,
        )
        return Qwen4ExpPLEUnquantizedEmbeddingMethod()

    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)

    def lookup_dtype(self, layer: nn.Module) -> torch.dtype:
        """Return the dtype of rows that a pinned-host lookup writes."""
        return layer.weight.dtype

    def lookup_from_pinned(
        self,
        layer: "Qwen4ExpPLEPinnedHostEmbedding",
        ids: torch.Tensor,
        output: torch.Tensor,
    ) -> None:
        """Copy the ETP-owned rows for ``ids`` from pinned host memory."""
        row_bytes = layer.embedding_dim * layer.weight.element_size()
        _lookup_ple_embedding_from_pinned_kernel[(ids.numel(),)](
            layer._uva_weight,
            ids,
            output,
            row_bytes,
            layer.shard_indices.org_vocab_start_index,
            layer.shard_indices.org_vocab_end_index,
            BLOCK_D=triton.next_power_of_2(row_bytes),
        )

    def process_weights_after_loading(self, layer: nn.Module) -> None:
        """Bind disk-offloaded tables once every checkpoint shard is loaded."""
        if isinstance(layer, Qwen4ExpPLEEmbedding):
            layer.bind_file_shards()

    @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 == "nvfp4":
        return Qwen4ExpPLENvFp4EmbeddingMethod()
    if embedding_dtype == "float8_e4m3fn":
        return Qwen4ExpPLEFp8EmbeddingMethod()
    if embedding_dtype is not None:
        raise ValueError(
            "ple_embedding_dtype must be 'nvfp4' or 'float8_e4m3fn', "
            f"got {embedding_dtype!r}"
        )
    if quant_config is None:
        return Qwen4ExpPLEUnquantizedEmbeddingMethod()
    if isinstance(quant_config, ModelOptQuantConfigBase):
        if quant_config.is_layer_excluded(prefix):
            return Qwen4ExpPLEUnquantizedEmbeddingMethod()
        if isinstance(quant_config, ModelOptMixedPrecisionConfig):
            quant_algo = quant_config._resolve_quant_algo(prefix)
            if quant_algo == "FP8":
                return Qwen4ExpPLEFp8EmbeddingMethod()
            if quant_algo == "NVFP4":
                return Qwen4ExpPLENvFp4EmbeddingMethod()
            if quant_algo is not None:
                raise ValueError(
                    f"Unsupported quant_algo {quant_algo!r} for the PLE embedding "
                    f"{prefix!r}; expected FP8 or NVFP4"
                )
            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 isinstance(quant_config, Fp8Config):
        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 quant_config.is_checkpoint_fp8_serialized:
            return Qwen4ExpPLEFp8EmbeddingMethod()
        raise NotImplementedError(
            "Qwen4Exp PLE embedding only supports serialized FP8 checkpoints"
        )

    logger.warning(
        "Falling back to unquantized PLE embedding for %s "
        "with quantization config %s.",
        prefix,
        type(quant_config).__name__,
    )
    return Qwen4ExpPLEUnquantizedEmbeddingMethod()

lookup_dtype(layer)

Return the dtype of rows that a pinned-host lookup writes.

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

def lookup_dtype(self, layer: nn.Module) -> torch.dtype:
    """Return the dtype of rows that a pinned-host lookup writes."""
    return layer.weight.dtype

lookup_from_pinned(layer, ids, output)

Copy the ETP-owned rows for ids from pinned host memory.

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

def lookup_from_pinned(
    self,
    layer: "Qwen4ExpPLEPinnedHostEmbedding",
    ids: torch.Tensor,
    output: torch.Tensor,
) -> None:
    """Copy the ETP-owned rows for ``ids`` from pinned host memory."""
    row_bytes = layer.embedding_dim * layer.weight.element_size()
    _lookup_ple_embedding_from_pinned_kernel[(ids.numel(),)](
        layer._uva_weight,
        ids,
        output,
        row_bytes,
        layer.shard_indices.org_vocab_start_index,
        layer.shard_indices.org_vocab_end_index,
        BLOCK_D=triton.next_power_of_2(row_bytes),
    )

process_weights_after_loading(layer)

Bind disk-offloaded tables once every checkpoint shard is loaded.

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

def process_weights_after_loading(self, layer: nn.Module) -> None:
    """Bind disk-offloaded tables once every checkpoint shard is loaded."""
    if isinstance(layer, Qwen4ExpPLEEmbedding):
        layer.bind_file_shards()

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")
        super().process_weights_after_loading(layer)

    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")
    super().process_weights_after_loading(layer)

Qwen4ExpPLENvFp4EmbeddingMethod

Bases: Qwen4ExpPLEEmbeddingMethod

NVFP4 PLE embedding: packed E2M1 rows, block scales, one global scale.

Rows stay packed in storage. A lookup decodes only the selected rows to the activation dtype, before the ETP reduction.

Methods:

• decode_rows –

  Decode the rows of ids in [vocab_start, vocab_end) into output.

• process_weights_after_loading –

  Reject NVFP4 PLE checkpoints without a global scale.

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

class Qwen4ExpPLENvFp4EmbeddingMethod(Qwen4ExpPLEEmbeddingMethod):
    """NVFP4 PLE embedding: packed E2M1 rows, block scales, one global scale.

    Rows stay packed in storage. A lookup decodes only the selected rows to
    the activation dtype, before the ETP reduction.
    """

    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
        if input_size_per_partition % _NVFP4_BLOCK_SIZE:
            raise ValueError(
                "NVFP4 PLE embedding dimension must be divisible by "
                f"{_NVFP4_BLOCK_SIZE}, got {input_size_per_partition}"
            )
        weight_loader = extra_weight_attrs.get("weight_loader")
        num_rows = sum(output_partition_sizes)
        for name, width, dtype in (
            (
                "weight",
                input_size_per_partition // _NVFP4_VALUES_PER_BYTE,
                torch.uint8,
            ),
            (
                "weight_scale",
                input_size_per_partition // _NVFP4_BLOCK_SIZE,
                torch.float8_e4m3fn,
            ),
        ):
            parameter = ModelWeightParameter(
                data=layer.allocate_embedding_weight(num_rows, width, dtype),
                input_dim=1,
                output_dim=0,
                weight_loader=weight_loader,
            )
            layer.register_parameter(name, parameter)
        weight_scale_2 = create_fp8_scale_parameter(
            PerTensorScaleParameter,
            output_partition_sizes,
            input_size_per_partition,
            None,
            weight_loader,
            scale_dtype=torch.float32,
        )
        layer.register_parameter("weight_scale_2", weight_scale_2)
        if layer.weight_scale.is_pinned():
            self._uva_weight_scale = get_accelerator_view_from_cpu_tensor(
                layer.weight_scale
            )

    def process_weights_after_loading(self, layer: nn.Module) -> None:
        """Reject NVFP4 PLE checkpoints without a global scale."""
        sentinel = torch.finfo(torch.float32).min
        if torch.any(layer.weight_scale_2 == sentinel):
            raise ValueError("NVFP4 PLE checkpoint is missing its global scale")
        super().process_weights_after_loading(layer)

    def lookup_dtype(self, layer: nn.Module) -> torch.dtype:
        return layer.params_dtype

    def embedding(self, layer: nn.Module, input_: torch.Tensor) -> torch.Tensor:
        ids = input_.reshape(-1)
        if not input_.is_cuda:
            rows = dequantize_to_dtype(
                F.embedding(ids, layer.weight),
                F.embedding(ids, layer.weight_scale),
                layer.weight_scale_2,
                layer.params_dtype,
                swizzle=False,
            )
            return rows.reshape(*input_.shape, layer.embedding_dim)
        output = torch.empty(
            (*input_.shape, layer.embedding_dim),
            dtype=layer.params_dtype,
            device=input_.device,
        )
        # Device rows are already local, so every ID in range is owned.
        self.decode_rows(
            layer.weight,
            layer.weight_scale,
            layer.weight_scale_2,
            ids,
            output,
            0,
            layer.weight.shape[0],
        )
        return output

    @staticmethod
    def decode_rows(
        weight: torch.Tensor,
        weight_scale: torch.Tensor,
        weight_scale_2: torch.Tensor,
        ids: torch.Tensor,
        output: torch.Tensor,
        vocab_start: int,
        vocab_end: int,
    ) -> None:
        """Decode the rows of `ids` in `[vocab_start, vocab_end)` into `output`."""
        if ids.numel():
            _lookup_nvfp4_ple_embedding_kernel[(ids.numel(),)](
                weight,
                weight_scale.view(torch.uint8),
                weight_scale_2,
                ids,
                output,
                output.shape[-1],
                vocab_start,
                vocab_end,
                BLOCK_D=triton.next_power_of_2(output.shape[-1]),
            )

    def lookup_from_pinned(
        self,
        layer: "Qwen4ExpPLEPinnedHostEmbedding",
        ids: torch.Tensor,
        output: torch.Tensor,
    ) -> None:
        self.decode_rows(
            layer._uva_weight,
            self._uva_weight_scale,
            layer.weight_scale_2,
            ids,
            output,
            layer.shard_indices.org_vocab_start_index,
            layer.shard_indices.org_vocab_end_index,
        )

    def dequantize(
        self,
        layer: nn.Module,
        embeddings: torch.Tensor,
        output_dtype: torch.dtype,
    ) -> torch.Tensor:
        del layer
        # The lookup already decoded the rows to the activation dtype.
        return embeddings.to(output_dtype)

decode_rows(weight, weight_scale, weight_scale_2, ids, output, vocab_start, vocab_end) staticmethod

Decode the rows of ids in [vocab_start, vocab_end) into output.

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

@staticmethod
def decode_rows(
    weight: torch.Tensor,
    weight_scale: torch.Tensor,
    weight_scale_2: torch.Tensor,
    ids: torch.Tensor,
    output: torch.Tensor,
    vocab_start: int,
    vocab_end: int,
) -> None:
    """Decode the rows of `ids` in `[vocab_start, vocab_end)` into `output`."""
    if ids.numel():
        _lookup_nvfp4_ple_embedding_kernel[(ids.numel(),)](
            weight,
            weight_scale.view(torch.uint8),
            weight_scale_2,
            ids,
            output,
            output.shape[-1],
            vocab_start,
            vocab_end,
            BLOCK_D=triton.next_power_of_2(output.shape[-1]),
        )

process_weights_after_loading(layer)

Reject NVFP4 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 NVFP4 PLE checkpoints without a global scale."""
    sentinel = torch.finfo(torch.float32).min
    if torch.any(layer.weight_scale_2 == sentinel):
        raise ValueError("NVFP4 PLE checkpoint is missing its global scale")
    super().process_weights_after_loading(layer)

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._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 in the embedding method's lookup dtype."""
        expected_shape = (*input_ids.shape, self.embedding_dim)
        lookup_dtype = self.embedding_method.lookup_dtype(self)
        if output is None:
            output = torch.empty(
                expected_shape,
                dtype=lookup_dtype,
                device=input_ids.device,
            )
        elif (
            tuple(output.shape) != expected_shape
            or output.dtype != lookup_dtype
            or output.device != input_ids.device
        ):
            raise ValueError(
                "PLE prefetch output must match the input shape, lookup dtype, "
                "and input device"
            )

        flat_ids = input_ids.reshape(-1).long()
        if flat_ids.numel():
            self.embedding_method.lookup_from_pinned(self, flat_ids, output)
        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)
        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
        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.embedding_method.lookup_dtype(self),
                        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())
        # Prevent ID storage reuse until the prefetch stream finishes reading it.
        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 in the embedding method's lookup 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 in the embedding method's lookup dtype."""
    expected_shape = (*input_ids.shape, self.embedding_dim)
    lookup_dtype = self.embedding_method.lookup_dtype(self)
    if output is None:
        output = torch.empty(
            expected_shape,
            dtype=lookup_dtype,
            device=input_ids.device,
        )
    elif (
        tuple(output.shape) != expected_shape
        or output.dtype != lookup_dtype
        or output.device != input_ids.device
    ):
        raise ValueError(
            "PLE prefetch output must match the input shape, lookup dtype, "
            "and input device"
        )

    flat_ids = input_ids.reshape(-1).long()
    if flat_ids.numel():
        self.embedding_method.lookup_from_pinned(self, flat_ids, output)
    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.embedding_method.lookup_dtype(self),
                    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())
    # Prevent ID storage reuse until the prefetch stream finishes reading it.
    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)
    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
    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

_maps_entry(ptr)

Return the path and file offset of ptr from /proc/self/maps.

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

def _maps_entry(ptr: int) -> tuple[str, int]:
    """Return the path and file offset of `ptr` from `/proc/self/maps`."""
    with open("/proc/self/maps") as maps:
        for line in maps:
            fields = line.split(maxsplit=5)
            low, high = (int(bound, 16) for bound in fields[0].split("-"))
            if low <= ptr < high:
                path = fields[5].rstrip("\n") if len(fields) > 5 else ""
                return path, ptr - low + int(fields[2], 16)
    return "", 0

stage_checkpoint_rows(modules, input_ids, query_start_loc, ngram_context)

Stage disk-offload PLE rows for one step behind one stream sync.

The ids hash this step's tokens, which the previous step samples on the GPU, so the host read has to wait for them.

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

def stage_checkpoint_rows(
    modules: list[tuple[Qwen4ExpNGramEmbedding, Qwen4ExpPLEDiskOffloadEmbedding]],
    input_ids: torch.Tensor,
    query_start_loc: torch.Tensor,
    ngram_context: torch.Tensor,
) -> None:
    """Stage disk-offload PLE rows for one step behind one stream sync.

    The ids hash this step's tokens, which the previous step samples on the
    GPU, so the host read has to wait for them.
    """
    num_tokens = input_ids.shape[0]
    for module, embedding in modules:
        ids = module.compute_ngram_ids(input_ids, query_start_loc, ngram_context)
        embedding.copy_ids_to_host(ids)
    if input_ids.device.type != "cpu":
        with gpu_sync_allowed():
            torch.accelerator.current_stream().synchronize()
    for _, embedding in modules:
        embedding.stage_rows(num_tokens)