[Qwen4Exp] Clean up HC and PLE implementations for Qwen3.8-Flash-Next (#61454)

Signed-off-by: huanghaoyan.hhy <huanghaoyan.hhy@alibaba-inc.com>
Co-authored-by: Thien Tran <gau.nernst@yahoo.com.sg>

tests/models/qwen4_exp/test_ple.py

@@ -1,13 +1,14 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+import logging
 from dataclasses import dataclass
-from functools import partial
 from itertools import accumulate
 from types import SimpleNamespace
 
 import pytest
 import torch
+from compressed_tensors.quantization import QuantizationArgs
 from torch import nn
 from torch.nn import functional as F
 from transformers import Qwen4ExpTextConfig
@@ -19,6 +20,10 @@ from tests.v1.attention.utils import BatchSpec, create_common_attn_metadata
 from vllm.config import SpeculativeConfig, VllmConfig
 from vllm.config.compilation import CUDAGraphMode
 from vllm.config.quantization import QuantizationConfigArgs
+from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
+from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors import (  # noqa: E501
+    CompressedTensorsConfig,
+)
 from vllm.model_executor.layers.quantization.fp8 import Fp8Config
 from vllm.model_executor.layers.quantization.inc import INCConfig
 from vllm.model_executor.layers.quantization.modelopt import (
@@ -33,8 +38,7 @@ from vllm.models.qwen4_exp.amd.ple_layer import (
 )
 from vllm.models.qwen4_exp.common.ple import (
     PLEShardOverlap,
-    compute_ple_shard_overlap,
-    copy_ple_embedding_shard_,
+    PLEVocabParallelEmbedding,
 )
 from vllm.models.qwen4_exp.nvidia.ngram_embedding import (
     Qwen4ExpNGramEmbedding,
@@ -210,12 +214,23 @@ def _make_ngram_embedding_for_load_test() -> Qwen4ExpNGramEmbedding:
     return module
 
 
+def _make_test_ple_embedding(tp_start: int, tp_end: int) -> PLEVocabParallelEmbedding:
+    """Create an embedding with only the state needed for shard loading."""
+    embedding = PLEVocabParallelEmbedding.__new__(PLEVocabParallelEmbedding)
+    nn.Module.__init__(embedding)
+    embedding.shard_indices = SimpleNamespace(
+        org_vocab_start_index=tp_start,
+        org_vocab_end_index=tp_end,
+    )
+    return embedding
+
+
 def _set_test_embedding_weight_loader(embedding) -> None:
-    embedding.weight.weight_loader = partial(
-        copy_ple_embedding_shard_,
+    loader = _make_test_ple_embedding(
         tp_start=embedding.shard_indices.org_vocab_start_index,
         tp_end=embedding.shard_indices.org_vocab_end_index,
     )
+    embedding.weight.weight_loader = loader.weight_loader
 
 
 def _make_fp8_ngram_embedding_for_load_test() -> Qwen4ExpNGramEmbedding:
@@ -244,19 +259,16 @@ def _make_fp8_ngram_embedding_for_load_test() -> Qwen4ExpNGramEmbedding:
 
 
 def test_ple_shard_overlap_and_copy() -> None:
-    overlap = compute_ple_shard_overlap(
-        checkpoint_start=2, checkpoint_rows=5, tp_start=4, tp_end=8
-    )
+    embedding = _make_test_ple_embedding(tp_start=4, tp_end=8)
+    overlap = embedding._compute_shard_overlap(checkpoint_start=2, checkpoint_rows=5)
     assert overlap == PLEShardOverlap(source_start=2, destination_start=0, row_count=3)
 
     destination = torch.full((4, 2), -1.0)
     loaded = torch.arange(10, dtype=torch.float64).reshape(5, 2)
-    copied = copy_ple_embedding_shard_(
+    copied = embedding._copy_embedding_shard_(
         destination,
         loaded,
         checkpoint_start=2,
-        tp_start=4,
-        tp_end=8,
     )
 
     assert copied == 3
@@ -265,13 +277,12 @@ def test_ple_shard_overlap_and_copy() -> None:
 
 
 def test_ple_shard_copy_is_a_noop_without_overlap() -> None:
+    embedding = _make_test_ple_embedding(tp_start=4, tp_end=8)
     destination = torch.ones(4, 2)
-    copied = copy_ple_embedding_shard_(
+    copied = embedding._copy_embedding_shard_(
         destination,
         torch.zeros(2, 2),
         checkpoint_start=10,
-        tp_start=4,
-        tp_end=8,
     )
 
     assert copied == 0
@@ -704,17 +715,44 @@ def test_ple_fp8_embedding_respects_checkpoint_shard_exclusions() -> None:
     )
 
 
-def test_ple_embedding_rejects_unsupported_quantization_configs() -> None:
+@pytest.mark.parametrize(
+    "quant_config",
+    [
+        pytest.param(ModelOptNvFp4Config(exclude_modules=[]), id="modelopt-nvfp4"),
+        pytest.param(
+            OnlineQuantizationConfig(QuantizationConfigArgs(linear="fp8_per_tensor")),
+            id="online-fp8",
+        ),
+        pytest.param(
+            CompressedTensorsConfig(
+                target_scheme_map={
+                    "re:.*ngram_embedding": {
+                        "weights": QuantizationArgs(num_bits=8, type="int"),
+                    },
+                },
+                ignore=[],
+                quant_format="int-quantized",
+            ),
+            id="compressed-tensors",
+        ),
+    ],
+)
+def test_ple_embedding_falls_back_for_unsupported_quantization_configs(
+    quant_config: QuantizationConfig,
+    caplog: pytest.LogCaptureFixture,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """Warn when global quantization falls back to an unquantized PLE table."""
     prefix = "model.layers.1.ple.ple_embedding.ngram_embedding"
-    nvfp4_config = ModelOptNvFp4Config(exclude_modules=[])
-    with pytest.raises(NotImplementedError, match="ModelOptNvFp4Config"):
-        Qwen4ExpPLEEmbeddingMethod.from_quant_config(nvfp4_config, prefix)
-
-    online_fp8_config = OnlineQuantizationConfig(
-        QuantizationConfigArgs(linear="fp8_per_tensor")
+    test_logger = logging.getLogger("test.ple_fallback")
+    monkeypatch.setattr(ngram_embedding_module, "logger", test_logger)
+    caplog.set_level(logging.WARNING, logger=test_logger.name)
+    assert isinstance(
+        Qwen4ExpPLEEmbeddingMethod.from_quant_config(quant_config, prefix),
+        Qwen4ExpPLEUnquantizedEmbeddingMethod,
     )
-    with pytest.raises(NotImplementedError, match="OnlineQuantizationConfig"):
-        Qwen4ExpPLEEmbeddingMethod.from_quant_config(online_fp8_config, prefix)
+    assert caplog.records[0].levelno == logging.WARNING
+    assert "Falling back to unquantized PLE" in caplog.text
 
 
 def test_ple_embedding_respects_modelopt_exclusion() -> None:
@@ -741,8 +779,10 @@ def test_ple_embedding_respects_inc_layer_config() -> None:
     )
 
     quant_config.extra_config = None
-    with pytest.raises(NotImplementedError, match="INCConfig"):
-        Qwen4ExpPLEEmbeddingMethod.from_quant_config(quant_config, prefix)
+    assert isinstance(
+        Qwen4ExpPLEEmbeddingMethod.from_quant_config(quant_config, prefix),
+        Qwen4ExpPLEUnquantizedEmbeddingMethod,
+    )
 
 
 def test_ple_embedding_is_unquantized_under_quark() -> None:
@@ -926,12 +966,10 @@ def test_ple_pinned_embedding_loads_on_cpu_and_looks_up_through_uva(
     loaded_weight = (
         torch.arange(12, dtype=torch.float32).reshape(4, 3).to(storage_dtype)
     )
-    copied = copy_ple_embedding_shard_(
+    copied = embedding._copy_embedding_shard_(
         embedding.weight,
         loaded_weight,
         checkpoint_start=0,
-        tp_start=0,
-        tp_end=4,
     )
     input_ids = torch.tensor([[3, 0], [1, 2]], device="cuda:0")
     output = embedding._lookup(input_ids)
@@ -988,12 +1026,10 @@ def test_pinned_embedding_start_prefetch_allocates_lazily(
 
     loaded_weight = torch.arange(12, dtype=torch.float32).reshape(4, 3)
     loaded_weight = loaded_weight.to(torch.bfloat16)
-    copy_ple_embedding_shard_(
+    embedding._copy_embedding_shard_(
         embedding.weight,
         loaded_weight,
         checkpoint_start=0,
-        tp_start=0,
-        tp_end=4,
     )
 
     ngram_ids = torch.tensor([[3, 0], [1, 2]], device="cuda:0")
@@ -1012,8 +1048,18 @@ def test_pinned_embedding_start_prefetch_allocates_lazily(
     torch.testing.assert_close(output, expected, rtol=0, atol=0)
 
 
-def test_ple_fp8_embedding_supports_mixed_precision_config() -> None:
+@pytest.mark.parametrize(
+    "exclude_pattern",
+    [
+        "model.language_model.layers.1.ple.ple_embedding.ngram_embedding",
+        "*.ngram_embedding",
+    ],
+)
+def test_ple_fp8_embedding_supports_mixed_precision_config(
+    exclude_pattern: str,
+) -> None:
     prefix = "model.language_model.layers.1.ple.ple_embedding.ngram_embedding"
+    nvfp4_prefix = "model.language_model.layers.2.ple.ple_embedding.ngram_embedding"
     quant_config = ModelOptMixedPrecisionConfig.from_config(
         {
             "quantization": {
@@ -1022,6 +1068,7 @@ def test_ple_fp8_embedding_supports_mixed_precision_config() -> None:
                 "group_size": 16,
                 "quantized_layers": {
                     prefix: {"quant_algo": "FP8"},
+                    nvfp4_prefix: {"quant_algo": "NVFP4"},
                     "model.language_model.layers.2.moe.gate_proj": {
                         "quant_algo": "MXFP8"
                     },
@@ -1034,13 +1081,30 @@ def test_ple_fp8_embedding_supports_mixed_precision_config() -> None:
         Qwen4ExpPLEEmbeddingMethod.from_quant_config(quant_config, prefix),
         Qwen4ExpPLEFp8EmbeddingMethod,
     )
-    # A quant_algo the PLE table cannot load must not fall back to unquantized.
+    assert isinstance(
+        Qwen4ExpPLEEmbeddingMethod.from_quant_config(
+            quant_config,
+            "model.language_model.layers.3.ple.ple_embedding.ngram_embedding",
+        ),
+        Qwen4ExpPLEUnquantizedEmbeddingMethod,
+    )
+    assert isinstance(
+        Qwen4ExpPLEEmbeddingMethod.from_quant_config(quant_config, nvfp4_prefix),
+        Qwen4ExpPLENvFp4EmbeddingMethod,
+    )
+    # An explicitly unsupported PLE format must not fall back to unquantized.
     with pytest.raises(ValueError, match="Unsupported quant_algo 'MXFP8'"):
         Qwen4ExpPLEEmbeddingMethod.from_quant_config(
             quant_config,
             "model.language_model.layers.2.moe.gate_proj",
         )
 
+    quant_config.exclude_modules = [exclude_pattern]
+    assert isinstance(
+        Qwen4ExpPLEEmbeddingMethod.from_quant_config(quant_config, prefix),
+        Qwen4ExpPLEUnquantizedEmbeddingMethod,
+    )
+
 
 def test_dilated_ple_spec_state_rolls_back_before_next_forward() -> None:
     conv_state_len = 6

vllm/models/qwen4_exp/amd/hyperconnection.py

@@ -70,7 +70,7 @@ class GatedResidual(nn.Module):
     ) -> None:
         super().__init__()
         self.config = config
-        self.lora_rank = config.hc_lowrank
+        self.hc_lowrank = config.hc_lowrank
         self.hc_count = config.hc_count
         self.hidden_size = config.hidden_size
         self.use_combine = use_combine
@@ -91,11 +91,11 @@ class GatedResidual(nn.Module):
         # -- vLLM Linear weights --------------------------------------------
         # The merged skinny-GEMM shape is physically padded to 16 rows for
         # alignment and efficient backend dispatch.
-        self.pad_size = (-(self.lora_rank + self.hc_count)) % 16 if use_combine else 0
+        self.pad_size = (-(self.hc_lowrank + self.hc_count)) % 16 if use_combine else 0
         if use_combine:
             self.input_mix_weight_down_block_inject = MergedColumnParallelLinear(
                 self.hyper_hidden_size,
-                [self.lora_rank, self.hc_count]
+                [self.hc_lowrank, self.hc_count]
                 + ([self.pad_size] if self.pad_size else []),
                 bias=False,
                 params_dtype=config.params_dtype,
@@ -107,7 +107,7 @@ class GatedResidual(nn.Module):
         else:
             self.input_mix_weight_down = ReplicatedLinear(
                 self.hyper_hidden_size,
-                self.lora_rank,
+                self.hc_lowrank,
                 bias=False,
                 params_dtype=config.params_dtype,
                 quant_config=None,
@@ -115,7 +115,7 @@ class GatedResidual(nn.Module):
                 return_bias=False,
             )
         self.input_mix_weight_up = ReplicatedLinear(
-            self.lora_rank,
+            self.hc_lowrank,
             self.hyper_hidden_size,
             bias=False,
             params_dtype=config.params_dtype,
@@ -136,15 +136,15 @@ class GatedResidual(nn.Module):
 
         if self.use_combine:
             # produce injection logits for combine
-            split_sizes = [self.lora_rank, self.hc_count, self.pad_size]
+            split_sizes = [self.hc_lowrank, self.hc_count, self.pad_size]
             down_and_injection = self.input_mix_weight_down_block_inject(xn)
-            lora, injection, _ = down_and_injection.split(split_sizes, dim=-1)
+            lowrank, injection, _ = down_and_injection.split(split_sizes, dim=-1)
         else:
-            lora = self.input_mix_weight_down(xn)
+            lowrank = self.input_mix_weight_down(xn)
             injection = None
 
-        lora = hc_silu(lora, self.hc_count)
-        gate = self.input_mix_weight_up(lora)  # [M, D]
+        lowrank = hc_silu(lowrank, self.hc_count)
+        gate = self.input_mix_weight_up(lowrank)  # [M, D]
         block_input = hc_gate_mix(xn, gate, self.hc_count)
 
         return hidden_states, block_input, injection
@@ -172,15 +172,15 @@ class GatedResidual(nn.Module):
 
         if self.use_combine:
             # produce injection logits for combine
-            split_sizes = [self.lora_rank, self.hc_count, self.pad_size]
+            split_sizes = [self.hc_lowrank, self.hc_count, self.pad_size]
             down_and_injection = self.input_mix_weight_down_block_inject(xn)
-            lora, injection, _ = down_and_injection.split(split_sizes, dim=-1)
+            lowrank, injection, _ = down_and_injection.split(split_sizes, dim=-1)
         else:
-            lora = self.input_mix_weight_down(xn)
+            lowrank = self.input_mix_weight_down(xn)
             injection = None
 
-        lora = hc_silu(lora, self.hc_count)
-        gate = self.input_mix_weight_up(lora)  # [M, D]
+        lowrank = hc_silu(lowrank, self.hc_count)
+        gate = self.input_mix_weight_up(lowrank)  # [M, D]
         block_input = hc_gate_mix(xn, gate, self.hc_count)
 
         return hidden_states, block_input, injection

vllm/models/qwen4_exp/amd/ple_layer.py

@@ -29,6 +29,7 @@ from vllm.models.common.ops.sequence_parallel import (
     sp_reduce_scatter,
     sp_shard,
 )
+from vllm.utils.math_utils import round_up
 from vllm.utils.torch_utils import direct_register_custom_op
 from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum
 from vllm.v1.attention.backends.short_conv_attn import (
@@ -242,7 +243,7 @@ class Qwen4ExpNGramEmbedding(nn.Module):
             persistent=True,
         )
         divisor = int(config.make_ngram_vocab_size_divisible_by)
-        padded_vocab_size = ((total_vocab_size + divisor - 1) // divisor) * divisor
+        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"
@@ -280,9 +281,11 @@ class Qwen4ExpNGramEmbedding(nn.Module):
             reduce_results=not self.use_reduce_scatter,
         )
         logger.info(
-            "Initialized AMD PLE embedding %s: quantization_method=%s, "
+            "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,
@@ -742,7 +745,7 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase):
         )
         kv, _ = self.kv_proj(embeddings)
         key, value = kv.split(self.kv_proj.output_sizes, dim=-1)
-        gated_output, conv_input = ple_gate(
+        gated_output, conv_input = torch.ops.vllm.qwen4_exp_ple_gate(
             key,
             value,
             hidden_states,
@@ -848,6 +851,28 @@ def qwen4_exp_ple_short_conv(
     layer._short_conv(inputs, residual, outer_residual)
 
 
+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)
+
+
+# Keep the AMD gate opaque to Dynamo while reusing the shared Triton kernel.
+direct_register_custom_op(
+    op_name="qwen4_exp_ple_gate",
+    op_func=ple_gate,
+    mutates_args=[],
+    fake_impl=_ple_gate_fake,
+)
+
+
 direct_register_custom_op(
     op_name="qwen4_exp_amd_ple_ngram_ids",
     op_func=qwen4_exp_amd_ple_ngram_ids,

vllm/models/qwen4_exp/common/ngram_embedding.py

@@ -194,22 +194,21 @@ class Qwen4ExpPLEEmbeddingMethod(QuantizeMethodBase):
             )
         if quant_config is None:
             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, ModelOptQuantConfigBase
-        ) and quant_config.is_layer_excluded(prefix):
-            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
@@ -223,29 +222,32 @@ class Qwen4ExpPLEEmbeddingMethod(QuantizeMethodBase):
         # Quark quantizes only Linear and MoE layers; PLE tables stay BF16.
         if isinstance(quant_config, QuarkConfig):
             return Qwen4ExpPLEUnquantizedEmbeddingMethod()
-        if not isinstance(quant_config, Fp8Config):
-            raise NotImplementedError(
-                "Qwen4Exp PLE embedding does not support quantization config "
-                f"{type(quant_config).__name__}"
-            )
-
-        ignored_layers = quant_config.ignored_layers
-        if is_layer_skipped(
-            prefix,
-            ignored_layers,
-            quant_config.packed_modules_mapping,
-            match_mode=quant_config.ignored_layers_match_mode,
-        ):
-            return Qwen4ExpPLEUnquantizedEmbeddingMethod()
-        # PLE checkpoint shards form one runtime embedding parameter.
-        shard_prefix = f"{prefix}.shard_"
-        if any(name.startswith(shard_prefix) for name in ignored_layers):
-            return Qwen4ExpPLEUnquantizedEmbeddingMethod()
-        if not quant_config.is_checkpoint_fp8_serialized:
+        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"
             )
-        return Qwen4ExpPLEFp8EmbeddingMethod()
+
+        logger.warning(
+            "Falling back to unquantized PLE embedding for %s "
+            "with quantization config %s.",
+            prefix,
+            type(quant_config).__name__,
+        )
+        return Qwen4ExpPLEUnquantizedEmbeddingMethod()
 
     def apply(
         self,
@@ -773,6 +775,7 @@ class Qwen4ExpPLEPinnedHostEmbedding(Qwen4ExpPLEEmbedding):
             )
         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)

vllm/models/qwen4_exp/common/ops/ple.py

@@ -16,7 +16,6 @@ import torch
 
 from vllm.platforms import current_platform
 from vllm.triton_utils import tl, triton
-from vllm.utils.torch_utils import direct_register_custom_op
 from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
 
 
@@ -240,7 +239,7 @@ def _ple_gate_kernel(
     tl.store(normed_ptr + t * HC * H + offs, normed, mask=mask)
 
 
-def _ple_gate(
+def ple_gate(
     key: torch.Tensor,
     value: torch.Tensor,
     hidden: torch.Tensor,
@@ -281,40 +280,6 @@ def _ple_gate(
     return gated, normed
 
 
-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 torch.empty_like(hidden), torch.empty_like(hidden)
-
-
-direct_register_custom_op(
-    op_name="qwen4_exp_ple_gate",
-    op_func=_ple_gate,
-    mutates_args=[],
-    fake_impl=_ple_gate_fake,
-)
-
-
-def ple_gate(
-    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 torch.ops.vllm.qwen4_exp_ple_gate(
-        key, value, hidden, norm_key_w, norm_query_w, norm_conv_w, eps
-    )
-
-
 # Dilated short convolution
 @triton.jit(do_not_specialize=["num_reqs", "bs_iters", "has_token_map"])
 def _ple_conv_kernel(

vllm/models/qwen4_exp/common/ple.py

@@ -20,66 +20,64 @@ class PLEShardOverlap:
     row_count: int
 
 
-def compute_ple_shard_overlap(
-    *,
-    checkpoint_start: int,
-    checkpoint_rows: int,
-    tp_start: int,
-    tp_end: int,
-) -> PLEShardOverlap | None:
-    """Compute the overlap of a checkpoint shard and one TP vocabulary range."""
-    if checkpoint_start < 0 or checkpoint_rows < 0:
-        raise ValueError("checkpoint shard bounds must be non-negative")
-    if tp_start < 0 or tp_end < tp_start:
-        raise ValueError("invalid TP vocabulary range")
-    checkpoint_end = checkpoint_start + checkpoint_rows
-    overlap_start = max(checkpoint_start, tp_start)
-    overlap_end = min(checkpoint_end, tp_end)
-    if overlap_start >= overlap_end:
-        return None
-    return PLEShardOverlap(
-        source_start=overlap_start - checkpoint_start,
-        destination_start=overlap_start - tp_start,
-        row_count=overlap_end - overlap_start,
-    )
-
-
-def copy_ple_embedding_shard_(
-    destination: torch.Tensor,
-    loaded_weight: torch.Tensor,
-    *,
-    checkpoint_start: int,
-    tp_start: int,
-    tp_end: int,
-) -> int:
-    """Copy the overlapping rows of a PLE checkpoint shard into a TP table."""
-    if destination.ndim == 0 or loaded_weight.ndim != destination.ndim:
-        raise ValueError("destination and loaded weight must have matching ranks")
-    if destination.shape[1:] != loaded_weight.shape[1:]:
-        raise ValueError(
-            "embedding shard dimensions do not match: "
-            f"{tuple(destination.shape[1:])} != {tuple(loaded_weight.shape[1:])}"
-        )
-    if destination.shape[0] < tp_end - tp_start:
-        raise ValueError("destination does not cover the requested TP range")
-    overlap = compute_ple_shard_overlap(
-        checkpoint_start=checkpoint_start,
-        checkpoint_rows=loaded_weight.shape[0],
-        tp_start=tp_start,
-        tp_end=tp_end,
-    )
-    if overlap is None:
-        return 0
-    source = loaded_weight.narrow(0, overlap.source_start, overlap.row_count)
-    target = destination.narrow(0, overlap.destination_start, overlap.row_count)
-    with torch.no_grad():
-        target.copy_(source.to(device=target.device, dtype=target.dtype))
-    return overlap.row_count
-
-
 class PLEVocabParallelEmbedding(VocabParallelEmbedding):
     """Vocab-parallel embedding that accepts checkpoint row shards."""
 
+    def _compute_shard_overlap(
+        self,
+        *,
+        checkpoint_start: int,
+        checkpoint_rows: int,
+    ) -> PLEShardOverlap | None:
+        """Compute the overlap of a checkpoint shard and one TP vocabulary range."""
+        tp_start = self.shard_indices.org_vocab_start_index
+        tp_end = self.shard_indices.org_vocab_end_index
+        if checkpoint_start < 0 or checkpoint_rows < 0:
+            raise ValueError("checkpoint shard bounds must be non-negative")
+        if tp_start < 0 or tp_end < tp_start:
+            raise ValueError("invalid TP vocabulary range")
+        checkpoint_end = checkpoint_start + checkpoint_rows
+        overlap_start = max(checkpoint_start, tp_start)
+        overlap_end = min(checkpoint_end, tp_end)
+        if overlap_start >= overlap_end:
+            return None
+        return PLEShardOverlap(
+            source_start=overlap_start - checkpoint_start,
+            destination_start=overlap_start - tp_start,
+            row_count=overlap_end - overlap_start,
+        )
+
+    def _copy_embedding_shard_(
+        self,
+        destination: torch.Tensor,
+        loaded_weight: torch.Tensor,
+        *,
+        checkpoint_start: int,
+    ) -> int:
+        """Copy the overlapping rows of a PLE checkpoint shard into a TP table."""
+        tp_start = self.shard_indices.org_vocab_start_index
+        tp_end = self.shard_indices.org_vocab_end_index
+        if destination.ndim == 0 or loaded_weight.ndim != destination.ndim:
+            raise ValueError("destination and loaded weight must have matching ranks")
+        if destination.shape[1:] != loaded_weight.shape[1:]:
+            raise ValueError(
+                "embedding shard dimensions do not match: "
+                f"{tuple(destination.shape[1:])} != {tuple(loaded_weight.shape[1:])}"
+            )
+        if destination.shape[0] < tp_end - tp_start:
+            raise ValueError("destination does not cover the requested TP range")
+        overlap = self._compute_shard_overlap(
+            checkpoint_start=checkpoint_start,
+            checkpoint_rows=loaded_weight.shape[0],
+        )
+        if overlap is None:
+            return 0
+        source = loaded_weight.narrow(0, overlap.source_start, overlap.row_count)
+        target = destination.narrow(0, overlap.destination_start, overlap.row_count)
+        with torch.no_grad():
+            target.copy_(source.to(device=target.device, dtype=target.dtype))
+        return overlap.row_count
+
     def weight_loader(
         self,
         param: torch.Tensor,
@@ -89,10 +87,8 @@ class PLEVocabParallelEmbedding(VocabParallelEmbedding):
         if checkpoint_start is None:
             super().weight_loader(param, loaded_weight)
             return
-        copy_ple_embedding_shard_(
+        self._copy_embedding_shard_(
             param,
             loaded_weight,
             checkpoint_start=checkpoint_start,
-            tp_start=self.shard_indices.org_vocab_start_index,
-            tp_end=self.shard_indices.org_vocab_end_index,
         )

vllm/models/qwen4_exp/common/qsa_cache.py

@@ -24,7 +24,7 @@ from vllm.config.cache import CacheDType
 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
 from vllm.platforms import current_platform
 from vllm.triton_utils import HAS_TRITON, tl, triton
-from vllm.utils.math_utils import cdiv
+from vllm.utils.math_utils import cdiv, round_up
 from vllm.v1.attention.backend import (
     AttentionBackend,
     AttentionCGSupport,
@@ -818,9 +818,7 @@ class QSAKeyStateCache(_QSAStateCache):
         key_head_size = int(kwargs.pop("head_size"))
         self.key_head_size = key_head_size
         self.cache_rope_positions = bool(cache_rope_positions)
-        self.rope_position_offset = (
-            (key_head_size + self._BF16_PER_INT64 - 1) // self._BF16_PER_INT64
-        ) * self._BF16_PER_INT64
+        self.rope_position_offset = round_up(key_head_size, self._BF16_PER_INT64)
         storage_head_size = key_head_size
         if self.cache_rope_positions:
             storage_head_size = self.rope_position_offset + (
@@ -848,7 +846,7 @@ class QSAKeyStateCache(_QSAStateCache):
         # the scheduler block size). Anything narrower lets a rejected draft row
         # overwrite a committed key the next step needs to close the group.
         span = self.compress_ratio + vllm_config.num_speculative_tokens
-        capacity = self.compress_ratio * cdiv(span, self.compress_ratio)
+        capacity = round_up(span, self.compress_ratio)
         assert self.cache_config.block_size % capacity == 0, (
             f"QSA ring capacity {capacity} must divide the attention block "
             f"size {self.cache_config.block_size}"

vllm/models/qwen4_exp/nvidia/hyperconnection.py

@@ -74,7 +74,7 @@ class GatedResidual(nn.Module):
     ) -> None:
         super().__init__()
         self.config = config
-        self.lora_rank = config.hc_lowrank
+        self.hc_lowrank = config.hc_lowrank
         self.hc_count = config.hc_count
         self.hidden_size = config.hidden_size
         self.use_combine = use_combine
@@ -95,11 +95,11 @@ class GatedResidual(nn.Module):
         # -- vLLM Linear weights --------------------------------------------
         # The merged skinny-GEMM shape is physically padded to 16 rows to ensure
         # good alignment and performant implementation chosen by CuBLAS heuristics.
-        self.pad_size = (-(self.lora_rank + self.hc_count)) % 16 if use_combine else 0
+        self.pad_size = (-(self.hc_lowrank + self.hc_count)) % 16 if use_combine else 0
         if use_combine:
             self.input_mix_weight_down_block_inject = MergedColumnParallelLinear(
                 self.hyper_hidden_size,
-                [self.lora_rank, self.hc_count]
+                [self.hc_lowrank, self.hc_count]
                 + ([self.pad_size] if self.pad_size else []),
                 bias=False,
                 params_dtype=config.params_dtype,
@@ -117,7 +117,7 @@ class GatedResidual(nn.Module):
         else:
             self.input_mix_weight_down = ReplicatedLinear(
                 self.hyper_hidden_size,
-                self.lora_rank,
+                self.hc_lowrank,
                 bias=False,
                 params_dtype=config.params_dtype,
                 quant_config=None,
@@ -125,7 +125,7 @@ class GatedResidual(nn.Module):
                 return_bias=False,
             )
         self.input_mix_weight_up = ReplicatedLinear(
-            self.lora_rank,
+            self.hc_lowrank,
             self.hyper_hidden_size,
             bias=False,
             params_dtype=config.params_dtype,
@@ -150,13 +150,13 @@ class GatedResidual(nn.Module):
             return hc_down_silu(
                 xn,
                 self.input_mix_weight_down_block_inject.weight,
-                self.lora_rank,
+                self.hc_lowrank,
                 self.hc_count,
             )
         down_and_injection = self.input_mix_weight_down_block_inject(xn)
-        split_sizes = [self.lora_rank, self.hc_count, self.pad_size]
-        lora, injection, _ = down_and_injection.split(split_sizes, dim=-1)
-        return hc_silu(lora, self.hc_count), injection
+        split_sizes = [self.hc_lowrank, self.hc_count, self.pad_size]
+        lowrank, injection, _ = down_and_injection.split(split_sizes, dim=-1)
+        return hc_silu(lowrank, self.hc_count), injection
 
     def mix(
         self, hidden_states: torch.Tensor
@@ -168,8 +168,8 @@ class GatedResidual(nn.Module):
             self.hc_count,
         )
 
-        lora, injection = self._down_and_inject(xn)
-        gate = self.input_mix_weight_up(lora)  # [M, D]
+        lowrank, injection = self._down_and_inject(xn)
+        gate = self.input_mix_weight_up(lowrank)  # [M, D]
         block_input = hc_gate_mix(xn, gate, self.hc_count)
 
         return hidden_states, block_input, injection
@@ -196,8 +196,8 @@ class GatedResidual(nn.Module):
             self.hc_count,
         )
 
-        lora, injection = self._down_and_inject(xn)
-        gate = self.input_mix_weight_up(lora)  # [M, D]
+        lowrank, injection = self._down_and_inject(xn)
+        gate = self.input_mix_weight_up(lowrank)  # [M, D]
         block_input = hc_gate_mix(xn, gate, self.hc_count)
 
         return hidden_states, block_input, injection

vllm/models/qwen4_exp/nvidia/ngram_embedding.py

@@ -17,6 +17,7 @@ from vllm.model_executor.layers.quantization.base_config import (
 from vllm.model_executor.layers.quantization.utils.fp8_utils import is_fp8
 from vllm.model_executor.models.utils import AutoWeightsLoader
 from vllm.models.common.ops.sequence_parallel import sp_reduce_scatter, sp_shard
+from vllm.utils.math_utils import round_up
 
 from ..common.ngram_embedding import (
     Qwen4ExpPLEDeviceEmbedding,
@@ -203,7 +204,7 @@ class Qwen4ExpNGramEmbedding(nn.Module):
             persistent=True,
         )
         divisor = int(config.make_ngram_vocab_size_divisible_by)
-        padded_vocab_size = ((total_vocab_size + divisor - 1) // divisor) * divisor
+        padded_vocab_size = round_up(total_vocab_size, divisor)
         embedding_prefix = f"{prefix}.ngram_embedding"
         embedding_quant_method = Qwen4ExpPLEEmbeddingMethod.from_quant_config(
             quant_config,
@@ -243,9 +244,11 @@ class Qwen4ExpNGramEmbedding(nn.Module):
             )
         weight = self.ngram_embedding.weight
         logger.info(
-            "Initialized PLE embedding %s: quantization_method=%s, "
+            "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,

vllm/models/qwen4_exp/nvidia/ops/hc.py

@@ -6,7 +6,6 @@ import torch
 
 from vllm.platforms import current_platform
 from vllm.triton_utils import tl, triton
-from vllm.utils.torch_utils import direct_register_custom_op
 
 
 @triton.jit
@@ -52,7 +51,7 @@ def _grouped_gemma_rmsnorm_kernel(
     tl.store(y_ptr + row * stride_y + offsets, y, mask)
 
 
-def _grouped_gemma_rmsnorm(
+def grouped_gemma_rmsnorm(
     x: torch.Tensor, weight: torch.Tensor, eps: float, num_groups: int
 ) -> torch.Tensor:
     N, DIM = x.shape
@@ -105,7 +104,7 @@ def _hc_silu_kernel(
     tl.store(y_ptr + row * stride_y + offs, y, mask)
 
 
-def _hc_silu(x: torch.Tensor, hc_count: int) -> torch.Tensor:
+def hc_silu(x: torch.Tensor, hc_count: int) -> torch.Tensor:
     num_tokens, DIM = x.shape
     assert x.stride(1) == 1
 
@@ -160,7 +159,7 @@ def _hc_gate_mix_kernel(
     tl.store(y_ptr + row * stride_y + offs_inner, acc, mask)
 
 
-def _hc_gate_mix(x: torch.Tensor, gate: torch.Tensor, hc_count: int) -> torch.Tensor:
+def hc_gate_mix(x: torch.Tensor, gate: torch.Tensor, hc_count: int) -> torch.Tensor:
     N, DIM = gate.shape
     assert x.shape == gate.shape
     assert DIM % hc_count == 0
@@ -232,7 +231,7 @@ def _hc_combine_kernel(
     tl.store(out_ptr + row * stride_out + offs, out, mask=mask)
 
 
-def _hc_combine(
+def hc_combine(
     residual: torch.Tensor,
     block_output: torch.Tensor,
     injection_logits: torch.Tensor | None,
@@ -345,7 +344,7 @@ def _hc_combine_norm_kernel(
     tl.store(y_ptr + row * stride_y + offs, y, mask_inner)
 
 
-def _hc_combine_norm(
+def hc_combine_norm(
     residual: torch.Tensor,
     block_output: torch.Tensor,
     injection_logits: torch.Tensor | None,
@@ -392,109 +391,6 @@ def _hc_combine_norm(
     return out, y
 
 
-def _same_shape_fake(x: torch.Tensor, *args) -> torch.Tensor:
-    return x.new_empty(x.shape)
-
-
-def _hc_gate_mix_fake(
-    x: torch.Tensor, gate: torch.Tensor, hc_count: int
-) -> torch.Tensor:
-    del gate
-    return x.new_empty((x.shape[0], x.shape[1] // hc_count))
-
-
-def _hc_combine_fake(
-    residual: torch.Tensor,
-    block_output: torch.Tensor,
-    injection_logits: torch.Tensor | None,
-    hc_count: int,
-) -> torch.Tensor:
-    del block_output, injection_logits, hc_count
-    return residual.new_empty(residual.shape)
-
-
-def _hc_combine_norm_fake(
-    residual: torch.Tensor,
-    block_output: torch.Tensor,
-    injection_logits: torch.Tensor | None,
-    norm_weight: torch.Tensor,
-    eps: float,
-    hc_count: int,
-) -> tuple[torch.Tensor, torch.Tensor]:
-    del block_output, injection_logits, norm_weight, eps, hc_count
-    return residual.new_empty(residual.shape), residual.new_empty(residual.shape)
-
-
-direct_register_custom_op(
-    op_name="qwen4_exp_grouped_gemma_rmsnorm",
-    op_func=_grouped_gemma_rmsnorm,
-    fake_impl=_same_shape_fake,
-)
-direct_register_custom_op(
-    op_name="qwen4_exp_hc_silu",
-    op_func=_hc_silu,
-    fake_impl=_same_shape_fake,
-)
-direct_register_custom_op(
-    op_name="qwen4_exp_hc_gate_mix",
-    op_func=_hc_gate_mix,
-    fake_impl=_hc_gate_mix_fake,
-)
-direct_register_custom_op(
-    op_name="qwen4_exp_hc_combine",
-    op_func=_hc_combine,
-    fake_impl=_hc_combine_fake,
-)
-direct_register_custom_op(
-    op_name="qwen4_exp_hc_combine_norm",
-    op_func=_hc_combine_norm,
-    fake_impl=_hc_combine_norm_fake,
-)
-
-
-def grouped_gemma_rmsnorm(
-    x: torch.Tensor, weight: torch.Tensor, eps: float, num_groups: int
-) -> torch.Tensor:
-    return torch.ops.vllm.qwen4_exp_grouped_gemma_rmsnorm(x, weight, eps, num_groups)
-
-
-def hc_silu(x: torch.Tensor, hc_count: int) -> torch.Tensor:
-    return torch.ops.vllm.qwen4_exp_hc_silu(x, hc_count)
-
-
-def hc_gate_mix(x: torch.Tensor, gate: torch.Tensor, hc_count: int) -> torch.Tensor:
-    return torch.ops.vllm.qwen4_exp_hc_gate_mix(x, gate, hc_count)
-
-
-def hc_combine(
-    residual: torch.Tensor,
-    block_output: torch.Tensor,
-    injection_logits: torch.Tensor | None,
-    hc_count: int,
-) -> torch.Tensor:
-    return torch.ops.vllm.qwen4_exp_hc_combine(
-        residual, block_output, injection_logits, hc_count
-    )
-
-
-def hc_combine_norm(
-    residual: torch.Tensor,
-    block_output: torch.Tensor,
-    injection_logits: torch.Tensor | None,
-    norm_weight: torch.Tensor,
-    eps: float,
-    hc_count: int,
-) -> tuple[torch.Tensor, torch.Tensor]:
-    return torch.ops.vllm.qwen4_exp_hc_combine_norm(
-        residual,
-        block_output,
-        injection_logits,
-        norm_weight,
-        eps,
-        hc_count,
-    )
-
-
 __all__ = [
     "grouped_gemma_rmsnorm",
     "hc_combine",

vllm/models/qwen4_exp/nvidia/ops/qsa_indexer.py

@@ -10,6 +10,7 @@ from vllm.model_executor.warmup.jit_warmup_triton_helper import (
 )
 from vllm.platforms import current_platform
 from vllm.triton_utils import tl, triton
+from vllm.utils.math_utils import round_up
 
 _TOPK_WORKSPACE_BYTES = 1024 * 1024
 _DECODE_BLOCK_N = 64
@@ -611,7 +612,7 @@ def qsa_select_paged_prefill(
     # No row scores beyond cdiv(max_seq_len, compress_ratio) compressed
     # columns. Round up to 64 to keep the logits row stride
     # cooperative_topk-compatible.
-    logits_width = triton.cdiv(triton.cdiv(max_seq_len, compress_ratio), 64) * 64
+    logits_width = round_up(triton.cdiv(max_seq_len, compress_ratio), 64)
     logits_width = min(max(64, logits_width), page_table.shape[1] * k_cache.shape[1])
 
     # chunk the inputs to keep temp logits below VLLM_SPARSE_INDEXER_MAX_LOGITS_MB