vllm.models.qwen4_exp.amd.ple_layer

GPU-resident Qwen4Exp position-learning enhancement layers.

Classes:

• Qwen4ExpNGramEmbedding –
• Qwen4ExpPLELayer –

Qwen4ExpNGramEmbedding

Bases: Module

Methods:

• forward –

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

• load_weights –

  Load hash buffers and checkpoint-split embedding rows.

Source code in vllm/models/qwen4_exp/amd/ple_layer.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,
        prefix: str,
        layer_name: str,
        *,
        data_parallel_rank: int = 0,
        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.layer_name = layer_name
        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)
        if params_dtype is None:
            params_dtype = torch.get_default_dtype()
        embedding_prefix = f"{prefix}.ngram_embedding"
        embedding_quant_method = Qwen4ExpPLEEmbeddingMethod.from_quant_config(
            quant_config,
            embedding_prefix,
            getattr(config, "ple_embedding_dtype", None),
        )
        if isinstance(embedding_quant_method, Qwen4ExpPLENvFp4EmbeddingMethod):
            raise NotImplementedError(
                "NVFP4 Qwen4Exp PLE tables are not supported on ROCm yet; use a "
                "checkpoint with a BF16 or FP8 PLE table."
            )
        engram_config = get_current_vllm_config().engram_config
        embedding_cls = (
            Qwen4ExpPLEPinnedHostEmbedding
            if engram_config is not None and engram_config.cpu_offload
            else Qwen4ExpPLEDeviceEmbedding
        )
        # A TP reduce-scatter covers all vocabulary owners only when ETP=TP.
        self.use_reduce_scatter = (
            self.use_sequence_parallel
            and get_etp_group().world_size == get_tp_group().world_size
        )
        self.ngram_embedding = embedding_cls(
            padded_vocab_size,
            self.head_dim,
            params_dtype=params_dtype,
            padding_size=divisor,
            prefix=embedding_prefix,
            embedding_method=embedding_quant_method,
            num_ngram_heads=self.ngram_heads,
            max_total_tokens=max_total_tokens,
            data_parallel_rank=data_parallel_rank,
            reduce_results=not self.use_reduce_scatter,
        )
        logger.info(
            "Initialized AMD 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__,
            self.ngram_embedding.weight.dtype,
            self.ngram_embedding.weight.device,
            self.ngram_embedding.weight.is_pinned(),
        )

    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,
        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."""
        input_ids = input_ids.reshape(-1)
        ngram_ids = input_ids.new_empty(
            (input_ids.shape[0], self.ngram_heads), dtype=torch.int64
        )
        torch.ops.vllm.qwen4_exp_amd_ple_ngram_ids(
            input_ids,
            query_start_loc,
            ngram_context,
            ngram_ids,
            self.layer_name,
        )
        embedding = self.ngram_embedding
        if embedding.supports_prefetch:
            output = ngram_ids.new_empty(
                (ngram_ids.shape[0], self.embedding_dim),
                dtype=embedding.weight.dtype,
            )
            torch.ops.vllm.qwen4_exp_amd_ple_ngram_embedding_pinned(
                ngram_ids,
                output,
                self.layer_name,
            )
        else:
            output = ngram_ids.new_empty(
                (ngram_ids.shape[0], self.embedding_dim),
                dtype=embedding.params_dtype,
            )
            torch.ops.vllm.qwen4_exp_amd_ple_ngram_embedding(
                ngram_ids,
                output,
                self.layer_name,
            )
        if not self.use_sequence_parallel:
            return output
        return self._sp_shard_embeddings(output)

    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
        """Load hash buffers and checkpoint-split embedding rows."""
        persistent_buffers = {
            "layer_multipliers": self.layer_multipliers,
            "ngram_heads_offsets": self.ngram_heads_offsets,
            "ngram_heads_vocab_sizes": self.ngram_heads_vocab_sizes,
        }
        loaded: set[str] = set()
        regular_weights: list[tuple[str, torch.Tensor]] = []
        shard_prefix = "ngram_embedding.shard_"

        for name, loaded_weight in weights:
            leaf_name = name.rsplit(".", 1)[-1]
            if leaf_name.startswith("hashstats_") or leaf_name == "token_lookup":
                continue
            if name in persistent_buffers:
                buffer = persistent_buffers[name]
                if buffer.shape != loaded_weight.shape:
                    raise ValueError(
                        f"Shape mismatch for {name}: expected "
                        f"{tuple(buffer.shape)}, got {tuple(loaded_weight.shape)}"
                    )
                buffer.copy_(loaded_weight.to(device=buffer.device, dtype=buffer.dtype))
                loaded.add(name)
                continue
            if name.startswith(shard_prefix) and name.endswith(".weight"):
                shard_text = name[len(shard_prefix) : -len(".weight")]
                if not shard_text.isdigit():
                    regular_weights.append((name, loaded_weight))
                    continue
                shard_index = int(shard_text)
                if shard_index >= self.split_ngram_parts:
                    raise ValueError(
                        f"PLE embedding shard index {shard_index} exceeds "
                        f"split_ngram_parts={self.split_ngram_parts}"
                    )
                embedding = self.ngram_embedding
                shard_size = (
                    embedding.org_vocab_size + self.split_ngram_parts - 1
                ) // self.split_ngram_parts
                checkpoint_start = shard_index * shard_size
                expected_rows = max(
                    0,
                    min(shard_size, embedding.org_vocab_size - checkpoint_start),
                )
                expected_shape = (expected_rows, embedding.embedding_dim)
                if tuple(loaded_weight.shape) != expected_shape:
                    raise ValueError(
                        f"Shape mismatch for PLE embedding shard {shard_index}: "
                        f"expected {expected_shape}, got "
                        f"{tuple(loaded_weight.shape)}"
                    )
                embedding.weight.weight_loader(
                    embedding.weight,
                    loaded_weight,
                    checkpoint_start=checkpoint_start,
                )
                loaded.add("ngram_embedding.weight")
                continue
            regular_weights.append((name, loaded_weight))

        if regular_weights:
            loaded.update(AutoWeightsLoader(self).load_weights(regular_weights))
        return loaded

_is_prime_64(value) staticmethod

Return whether a 64-bit integer is prime.

Source code in vllm/models/qwen4_exp/amd/ple_layer.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/amd/ple_layer.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/amd/ple_layer.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/amd/ple_layer.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/amd/ple_layer.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/amd/ple_layer.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

forward(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/amd/ple_layer.py

def forward(
    self,
    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."""
    input_ids = input_ids.reshape(-1)
    ngram_ids = input_ids.new_empty(
        (input_ids.shape[0], self.ngram_heads), dtype=torch.int64
    )
    torch.ops.vllm.qwen4_exp_amd_ple_ngram_ids(
        input_ids,
        query_start_loc,
        ngram_context,
        ngram_ids,
        self.layer_name,
    )
    embedding = self.ngram_embedding
    if embedding.supports_prefetch:
        output = ngram_ids.new_empty(
            (ngram_ids.shape[0], self.embedding_dim),
            dtype=embedding.weight.dtype,
        )
        torch.ops.vllm.qwen4_exp_amd_ple_ngram_embedding_pinned(
            ngram_ids,
            output,
            self.layer_name,
        )
    else:
        output = ngram_ids.new_empty(
            (ngram_ids.shape[0], self.embedding_dim),
            dtype=embedding.params_dtype,
        )
        torch.ops.vllm.qwen4_exp_amd_ple_ngram_embedding(
            ngram_ids,
            output,
            self.layer_name,
        )
    if not self.use_sequence_parallel:
        return output
    return self._sp_shard_embeddings(output)

load_weights(weights)

Load hash buffers and checkpoint-split embedding rows.

Source code in vllm/models/qwen4_exp/amd/ple_layer.py

def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
    """Load hash buffers and checkpoint-split embedding rows."""
    persistent_buffers = {
        "layer_multipliers": self.layer_multipliers,
        "ngram_heads_offsets": self.ngram_heads_offsets,
        "ngram_heads_vocab_sizes": self.ngram_heads_vocab_sizes,
    }
    loaded: set[str] = set()
    regular_weights: list[tuple[str, torch.Tensor]] = []
    shard_prefix = "ngram_embedding.shard_"

    for name, loaded_weight in weights:
        leaf_name = name.rsplit(".", 1)[-1]
        if leaf_name.startswith("hashstats_") or leaf_name == "token_lookup":
            continue
        if name in persistent_buffers:
            buffer = persistent_buffers[name]
            if buffer.shape != loaded_weight.shape:
                raise ValueError(
                    f"Shape mismatch for {name}: expected "
                    f"{tuple(buffer.shape)}, got {tuple(loaded_weight.shape)}"
                )
            buffer.copy_(loaded_weight.to(device=buffer.device, dtype=buffer.dtype))
            loaded.add(name)
            continue
        if name.startswith(shard_prefix) and name.endswith(".weight"):
            shard_text = name[len(shard_prefix) : -len(".weight")]
            if not shard_text.isdigit():
                regular_weights.append((name, loaded_weight))
                continue
            shard_index = int(shard_text)
            if shard_index >= self.split_ngram_parts:
                raise ValueError(
                    f"PLE embedding shard index {shard_index} exceeds "
                    f"split_ngram_parts={self.split_ngram_parts}"
                )
            embedding = self.ngram_embedding
            shard_size = (
                embedding.org_vocab_size + self.split_ngram_parts - 1
            ) // self.split_ngram_parts
            checkpoint_start = shard_index * shard_size
            expected_rows = max(
                0,
                min(shard_size, embedding.org_vocab_size - checkpoint_start),
            )
            expected_shape = (expected_rows, embedding.embedding_dim)
            if tuple(loaded_weight.shape) != expected_shape:
                raise ValueError(
                    f"Shape mismatch for PLE embedding shard {shard_index}: "
                    f"expected {expected_shape}, got "
                    f"{tuple(loaded_weight.shape)}"
                )
            embedding.weight.weight_loader(
                embedding.weight,
                loaded_weight,
                checkpoint_start=checkpoint_start,
            )
            loaded.add("ngram_embedding.weight")
            continue
        regular_weights.append((name, loaded_weight))

    if regular_weights:
        loaded.update(AutoWeightsLoader(self).load_weights(regular_weights))
    return loaded

Qwen4ExpPLELayer

Bases: Module, MambaBase

Source code in vllm/models/qwen4_exp/amd/ple_layer.py

class Qwen4ExpPLELayer(nn.Module, MambaBase):
    def __init__(
        self,
        config: Qwen4ExpTextConfig,
        vllm_config: VllmConfig,
        layer_idx: int = 0,
        ple_dense_layer_id: int | None = None,
        prefix: str = "",
        use_sequence_parallel: bool = False,
    ) -> None:
        super().__init__()
        self.use_sequence_parallel = use_sequence_parallel
        model_config = vllm_config.model_config
        cache_config = vllm_config.cache_config
        quant_config = vllm_config.quant_config
        self.model_config: ModelConfig = model_config
        self.cache_config: CacheConfig = cache_config
        self.layer_idx = layer_idx
        self.ple_dense_layer_id = (
            int(ple_dense_layer_id)
            if ple_dense_layer_id is not None
            else int(layer_idx)
        )
        self.prefix = prefix
        self.hidden_size = int(config.hidden_size)
        self.hc_count = config.hc_count
        self.hc_hidden_size = self.hidden_size * self.hc_count
        self.conv_kernel_size = int(config.ple_conv_kernel_size)
        self.short_conv_dilation = int(config.ngram_size)
        self.conv_state_len = (self.conv_kernel_size - 1) * self.short_conv_dilation
        self.num_spec_tokens = vllm_config.num_speculative_tokens
        self.activation = "silu"
        self.ple_embedding: nn.Module = Qwen4ExpNGramEmbedding(
            config,
            int(config.ple_embed_dim),
            self.ple_dense_layer_id,
            vllm_config.scheduler_config.max_num_batched_tokens,
            f"{prefix}.ple_embedding",
            prefix,
            data_parallel_rank=vllm_config.parallel_config.data_parallel_rank,
            quant_config=quant_config,
            params_dtype=model_config.dtype,
            use_sequence_parallel=self.use_sequence_parallel,
        )
        # Projection weights stay replicated even when their token rows are sharded.
        self.kv_proj = MergedColumnParallelLinear(
            int(config.ple_embed_dim),
            [self.hc_hidden_size, self.hidden_size],
            bias=False,
            params_dtype=model_config.dtype,
            quant_config=quant_config,
            prefix=f"{prefix}.kv_proj",
            disable_tp=True,
        )
        norm_args = (
            self.hc_hidden_size,
            config.rms_norm_eps,
            self.hidden_size,
            model_config.dtype,
        )
        self.norm_key = Qwen4ExpPLEGroupedNorm(*norm_args)
        self.norm_query = Qwen4ExpPLEGroupedNorm(*norm_args)
        self.norm_conv = Qwen4ExpPLEGroupedNorm(*norm_args)
        self.conv1d = nn.Conv1d(
            self.hc_hidden_size,
            self.hc_hidden_size,
            self.conv_kernel_size,
            groups=self.hc_hidden_size,
            padding=self.conv_state_len,
            dilation=self.short_conv_dilation,
            bias=False,
            dtype=model_config.dtype,
        )
        nn.init.zeros_(self.conv1d.weight)
        self.conv1d.weight._no_reinit = True
        self.kv_cache = (torch.tensor([]),)
        compilation_config = get_current_vllm_config().compilation_config
        if prefix in compilation_config.static_forward_context:
            raise ValueError(f"Duplicate layer name: {prefix}")
        compilation_config.static_forward_context[prefix] = self

    @property
    def mamba_type(self) -> MambaAttentionBackendEnum:
        return MambaAttentionBackendEnum.SHORT_CONV

    @property
    def is_kv_cache_tp_replicated(self) -> bool:
        return True

    def get_attn_backend(self) -> type[PleShortConvAttentionBackend]:
        return PleShortConvAttentionBackend

    def get_state_dtype(self) -> tuple[torch.dtype, ...]:
        return MambaStateDtypeCalculator.short_conv_state_dtype(
            self.model_config.dtype, self.cache_config.mamba_cache_dtype
        )

    def get_state_shape(self) -> Sequence[tuple[int, ...]]:
        return MambaStateShapeCalculator.short_conv_state_shape(
            tp_world_size=1,
            intermediate_size=self.hc_hidden_size,
            conv_kernel=self.conv_state_len + 1,
            num_spec=self.num_spec_tokens,
        )

    def _short_conv_dilated_dispatch(
        self,
        inputs: torch.Tensor,
        residual: torch.Tensor,
        outer_residual: torch.Tensor | None,
        metadata: PleShortConvAttentionMetadata,
        conv_state: torch.Tensor,
        conv_weights: torch.Tensor,
    ) -> None:
        num_prefills = metadata.num_prefills
        num_decodes = metadata.num_decodes
        num_decode_tokens = metadata.num_decode_tokens
        num_prefill_tokens = metadata.num_prefill_tokens
        has_prefill = num_prefills > 0
        has_decode = num_decodes > 0
        has_spec = metadata.spec_sequence_masks is not None
        has_non_spec = has_prefill or has_decode
        inputs = inputs[: metadata.num_actual_tokens]
        residual = residual[: metadata.num_actual_tokens]
        if outer_residual is not None:
            outer_residual = outer_residual[: metadata.num_actual_tokens]

        spec_token_indices = None
        non_spec_token_indices = None
        if has_spec and has_non_spec:
            assert metadata.spec_token_indx is not None
            assert metadata.non_spec_token_indx is not None
            spec_token_indices = metadata.spec_token_indx
            non_spec_token_indices = metadata.non_spec_token_indx

        if has_spec:
            assert metadata.spec_state_indices_tensor is not None
            query_start_loc = metadata.spec_query_start_loc
            num_accepted_tokens = metadata.num_accepted_tokens
            assert query_start_loc is not None
            assert num_accepted_tokens is not None
            spec_state_indices = metadata.spec_state_indices_tensor[
                : metadata.num_spec_decodes
            ]
            # Mixed batches stay in their original row order; the kernels map
            # logical spec/non-spec rows instead of materializing both groups.
            ple_conv(
                inputs=inputs,
                residual=residual,
                conv_state=conv_state,
                conv_weights=conv_weights,
                state_indices=spec_state_indices,
                outer_residual=outer_residual,
                mode="spec",
                dilation=self.short_conv_dilation,
                query_start_loc=query_start_loc,
                num_accepted_tokens=num_accepted_tokens,
                spec_query_len=metadata.spec_query_len,
                token_indices=spec_token_indices,
            )

        if not has_non_spec:
            return

        state_indices = metadata.state_indices_tensor
        assert state_indices is not None
        if has_prefill:
            state_indices_d, state_indices_p = torch.split(
                state_indices, [num_decodes, num_prefills], dim=0
            )
            if non_spec_token_indices is None:
                inputs_d, inputs_p = torch.split(
                    inputs, [num_decode_tokens, num_prefill_tokens], dim=0
                )
                residual_d, residual_p = torch.split(
                    residual, [num_decode_tokens, num_prefill_tokens], dim=0
                )
                outer_residual_d: torch.Tensor | None = None
                outer_residual_p: torch.Tensor | None = None
                if outer_residual is not None:
                    outer_residual_d, outer_residual_p = torch.split(
                        outer_residual,
                        [num_decode_tokens, num_prefill_tokens],
                        dim=0,
                    )
                token_indices_d = None
                token_indices_p = None
            else:
                inputs_d = inputs_p = inputs
                residual_d = residual_p = residual
                outer_residual_d = outer_residual_p = outer_residual
                token_indices_d, token_indices_p = torch.split(
                    non_spec_token_indices,
                    [num_decode_tokens, num_prefill_tokens],
                    dim=0,
                )

            if has_decode:
                ple_conv(
                    inputs=inputs_d,
                    residual=residual_d,
                    conv_state=conv_state,
                    conv_weights=conv_weights,
                    state_indices=state_indices_d,
                    outer_residual=outer_residual_d,
                    mode="decode",
                    dilation=self.short_conv_dilation,
                    has_initial_states=metadata.has_initial_states_d,
                    token_indices=token_indices_d,
                )

            query_start_loc = metadata.query_start_loc_p
            if query_start_loc is None:
                raise ValueError("query_start_loc is required for prefill short-conv")
            has_initial_states = metadata.has_initial_states_p
            if has_initial_states is None:
                raise ValueError(
                    "has_initial_states_p is required for prefill short-conv"
                )
            ple_conv(
                inputs=inputs_p,
                residual=residual_p,
                conv_state=conv_state,
                conv_weights=conv_weights,
                state_indices=state_indices_p,
                outer_residual=outer_residual_p,
                mode="prefill",
                dilation=self.short_conv_dilation,
                query_start_loc=query_start_loc,
                has_initial_states=has_initial_states,
                token_indices=token_indices_p,
            )
        else:
            num_decode_rows = (
                non_spec_token_indices.numel()
                if non_spec_token_indices is not None
                else inputs.size(0)
            )
            ple_conv(
                inputs=inputs,
                residual=residual,
                conv_state=conv_state,
                conv_weights=conv_weights,
                state_indices=state_indices[:num_decode_rows],
                outer_residual=outer_residual,
                mode="decode",
                dilation=self.short_conv_dilation,
                has_initial_states=metadata.has_initial_states_d,
                token_indices=non_spec_token_indices,
            )

    def _short_conv(
        self,
        inputs: torch.Tensor,
        residual: torch.Tensor,
        outer_residual: torch.Tensor | None,
    ) -> None:
        """Update convolution state, optionally fusing both residual additions."""
        forward_context = get_forward_context()
        attn_metadata = forward_context.attn_metadata
        # Profiling omits all metadata or this Mamba entry. Short convolution
        # is a no-op there, but preserve the outer residual addition.
        if attn_metadata is None:
            if outer_residual is not None:
                residual.add_(outer_residual)
            return

        if not isinstance(attn_metadata, dict):
            raise RuntimeError(
                "PLE short-conv expects per-layer attention metadata dict "
                f"during inference, got {type(attn_metadata).__name__}."
            )

        layer_attn_metadata = attn_metadata.get(self.prefix)
        if layer_attn_metadata is None:
            if outer_residual is not None:
                residual.add_(outer_residual)
            return
        if not isinstance(layer_attn_metadata, PleShortConvAttentionMetadata):
            raise TypeError(
                "Expected PleShortConvAttentionMetadata for layer "
                f"'{self.prefix}', got "
                f"{type(layer_attn_metadata).__name__}."
            )

        conv_state = self.kv_cache[0]
        # Canonicalize both backend cache layouts to [slot, channel, window].
        if not is_conv_state_dim_first():
            conv_state = conv_state.transpose(-1, -2)
        conv_weights = self.conv1d.weight.squeeze(1)

        state_capacity = self.conv_state_len + self.num_spec_tokens
        if state_capacity > 0:
            if conv_state.size(-1) < state_capacity:
                raise RuntimeError(
                    "PLE short-conv cache is smaller than expected for "
                    f"dilated convolution: got {conv_state.size(-1)}, "
                    f"expect at least {state_capacity}."
                )
            conv_state = conv_state[..., -state_capacity:]
        self._short_conv_dilated_dispatch(
            inputs=inputs,
            residual=residual,
            outer_residual=outer_residual,
            metadata=layer_attn_metadata,
            conv_state=conv_state,
            conv_weights=conv_weights.to(dtype=inputs.dtype),
        )

    def forward(
        self,
        hidden_states: torch.Tensor,
        input_ids: torch.Tensor,
        query_start_loc: torch.Tensor,
        ngram_context: torch.Tensor,
    ) -> torch.Tensor:
        input_ids = input_ids.reshape(-1)
        if (
            not self.use_sequence_parallel
            and input_ids.shape[0] != hidden_states.shape[0]
        ):
            raise ValueError(
                "PLE expects input_ids and hidden_states to have the same "
                f"token length, got {input_ids.shape[0]} and "
                f"{hidden_states.shape[0]}"
            )
        embeddings = self.ple_embedding(input_ids, query_start_loc, ngram_context)
        embeddings = self.ple_embedding.ngram_embedding.dequantize(
            embeddings, hidden_states.dtype
        )
        kv, _ = self.kv_proj(embeddings)
        key, value = kv.split(self.kv_proj.output_sizes, dim=-1)
        gated_output, conv_input = torch.ops.vllm.qwen4_exp_ple_gate(
            key,
            value,
            hidden_states,
            self.norm_key.weight,
            self.norm_query.weight,
            self.norm_conv.weight,
            self.norm_key.eps,
        )
        # The short-conv op is a piecewise-graph splitting op: state routing
        # reads runtime request metadata. It accumulates the convolution and
        # the outer residual into gated_output.
        if self.use_sequence_parallel:
            # Only convolution needs full tokens. Add residuals on local shards.
            conv_input = sp_all_gather(conv_input)[: input_ids.shape[0]]
            # Metadata may omit padding tokens or all tokens during profiling.
            conv_output = torch.zeros_like(conv_input)
            torch.ops.vllm.qwen4_exp_ple_short_conv(
                conv_input, conv_output, None, self.prefix
            )
            gated_output.add_(sp_shard(conv_output))
            gated_output.add_(hidden_states)
        else:
            torch.ops.vllm.qwen4_exp_ple_short_conv(
                conv_input,
                gated_output,
                hidden_states,
                self.prefix,
            )
        return gated_output

_short_conv(inputs, residual, outer_residual)

Update convolution state, optionally fusing both residual additions.

Source code in vllm/models/qwen4_exp/amd/ple_layer.py

def _short_conv(
    self,
    inputs: torch.Tensor,
    residual: torch.Tensor,
    outer_residual: torch.Tensor | None,
) -> None:
    """Update convolution state, optionally fusing both residual additions."""
    forward_context = get_forward_context()
    attn_metadata = forward_context.attn_metadata
    # Profiling omits all metadata or this Mamba entry. Short convolution
    # is a no-op there, but preserve the outer residual addition.
    if attn_metadata is None:
        if outer_residual is not None:
            residual.add_(outer_residual)
        return

    if not isinstance(attn_metadata, dict):
        raise RuntimeError(
            "PLE short-conv expects per-layer attention metadata dict "
            f"during inference, got {type(attn_metadata).__name__}."
        )

    layer_attn_metadata = attn_metadata.get(self.prefix)
    if layer_attn_metadata is None:
        if outer_residual is not None:
            residual.add_(outer_residual)
        return
    if not isinstance(layer_attn_metadata, PleShortConvAttentionMetadata):
        raise TypeError(
            "Expected PleShortConvAttentionMetadata for layer "
            f"'{self.prefix}', got "
            f"{type(layer_attn_metadata).__name__}."
        )

    conv_state = self.kv_cache[0]
    # Canonicalize both backend cache layouts to [slot, channel, window].
    if not is_conv_state_dim_first():
        conv_state = conv_state.transpose(-1, -2)
    conv_weights = self.conv1d.weight.squeeze(1)

    state_capacity = self.conv_state_len + self.num_spec_tokens
    if state_capacity > 0:
        if conv_state.size(-1) < state_capacity:
            raise RuntimeError(
                "PLE short-conv cache is smaller than expected for "
                f"dilated convolution: got {conv_state.size(-1)}, "
                f"expect at least {state_capacity}."
            )
        conv_state = conv_state[..., -state_capacity:]
    self._short_conv_dilated_dispatch(
        inputs=inputs,
        residual=residual,
        outer_residual=outer_residual,
        metadata=layer_attn_metadata,
        conv_state=conv_state,
        conv_weights=conv_weights.to(dtype=inputs.dtype),
    )

_ple_gate_fake(key, value, hidden, norm_key_w, norm_query_w, norm_conv_w, eps)

Return the gate output shapes for AMD graph tracing.

Source code in vllm/models/qwen4_exp/amd/ple_layer.py

def _ple_gate_fake(
    key: torch.Tensor,
    value: torch.Tensor,
    hidden: torch.Tensor,
    norm_key_w: torch.Tensor,
    norm_query_w: torch.Tensor,
    norm_conv_w: torch.Tensor,
    eps: float,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Return the gate output shapes for AMD graph tracing."""
    return torch.empty_like(hidden), torch.empty_like(hidden)

qwen4_exp_amd_ple_ngram_embedding(ngram_ids, output, layer_name)

Run the large PLE embedding lookup outside Inductor's FX graph.

Keeping the embedding weight in static_forward_context prevents AOT compile-time autotuning from materializing a synthetic copy of the weight.

Source code in vllm/models/qwen4_exp/amd/ple_layer.py

def qwen4_exp_amd_ple_ngram_embedding(
    ngram_ids: torch.Tensor,
    output: torch.Tensor,
    layer_name: str,
) -> None:
    """Run the large PLE embedding lookup outside Inductor's FX graph.

    Keeping the embedding weight in ``static_forward_context`` prevents AOT
    compile-time autotuning from materializing a synthetic copy of the weight.
    """
    layer = get_forward_context().no_compile_layers[layer_name]
    if not isinstance(layer, Qwen4ExpPLELayer):
        raise TypeError(f"{layer_name} is not a Qwen4Exp PLE owner")
    result = layer.ple_embedding.ngram_embedding(ngram_ids).flatten(-2)
    output.copy_(result)

qwen4_exp_amd_ple_ngram_embedding_pinned(ngram_ids, output, layer_name)

Run the pinned PLE UVA lookup outside Inductor's FX graph.

Same rationale as the device-path escape: keeping the large embedding weight out of the graph prevents AOT compile-time autotuning from materializing a synthetic copy of the weight.

Source code in vllm/models/qwen4_exp/amd/ple_layer.py

def qwen4_exp_amd_ple_ngram_embedding_pinned(
    ngram_ids: torch.Tensor,
    output: torch.Tensor,
    layer_name: str,
) -> None:
    """Run the pinned PLE UVA lookup outside Inductor's FX graph.

    Same rationale as the device-path escape: keeping the large embedding weight
    out of the graph prevents AOT compile-time autotuning from materializing a
    synthetic copy of the weight.
    """
    layer = get_forward_context().no_compile_layers[layer_name]
    if not isinstance(layer, Qwen4ExpPLELayer):
        raise TypeError(f"{layer_name} is not a Qwen4Exp PLE owner")
    result = layer.ple_embedding.ngram_embedding.sync_lookup(ngram_ids).flatten(-2)
    output.copy_(result)

qwen4_exp_amd_ple_ngram_ids(input_ids, query_start_loc, ngram_context, output, layer_name)

Hash the current request layout into n-gram embedding indices.

The launch sizes its request binary search from the request count, which is symbolic under Dynamo, so it runs outside Inductor's FX graph.

Source code in vllm/models/qwen4_exp/amd/ple_layer.py

def qwen4_exp_amd_ple_ngram_ids(
    input_ids: torch.Tensor,
    query_start_loc: torch.Tensor,
    ngram_context: torch.Tensor,
    output: torch.Tensor,
    layer_name: str,
) -> None:
    """Hash the current request layout into n-gram embedding indices.

    The launch sizes its request binary search from the request count, which
    is symbolic under Dynamo, so it runs outside Inductor's FX graph.
    """
    layer = get_forward_context().no_compile_layers[layer_name]
    if not isinstance(layer, Qwen4ExpPLELayer):
        raise TypeError(f"{layer_name} is not a Qwen4Exp PLE owner")
    embedding = layer.ple_embedding
    ple_ngram_ids(
        input_ids=input_ids,
        query_start_loc=query_start_loc,
        ngram_context=ngram_context,
        layer_multipliers=embedding.layer_multipliers,
        ngram_heads_vocab_sizes=embedding.ngram_heads_vocab_sizes,
        ngram_heads_offsets=embedding.ngram_heads_offsets,
        eos_token_id=embedding.eos_token_id,
        heads_per_ngram=embedding.heads_per_ngram,
        output=output,
    )