[Bugfix][Quantization] Load the FP8 PLE n-gram table under a ModelOpt NVFP4 checkpoint

tests/models/qwen4_exp/test_ple.py

@@ -16,6 +16,7 @@
 from vllm.model_executor.layers.quantization.fp8 import Fp8Config
 from vllm.model_executor.layers.quantization.modelopt import (
     ModelOptMixedPrecisionConfig,
+    ModelOptNvFp4Config,
 )
 from vllm.models.qwen4_exp.common.ple import (
     PLEShardOverlap,
@@ -322,6 +323,55 @@ def test_ple_fp8_embedding_supports_mixed_precision_config() -> None:
     )
 
 
+def _nvfp4_config(exclude_modules: list[str]) -> ModelOptNvFp4Config:
+    """Mirrors ``config.json``'s ``quantization_config`` in NVFP4 checkpoints."""
+    return ModelOptNvFp4Config.from_config(
+        {
+            "quant_algo": "NVFP4",
+            "quant_method": "modelopt",
+            "ignore": exclude_modules,
+            "group_size": 16,
+        }
+    )
+
+
+def test_ple_fp8_embedding_loads_under_nvfp4_checkpoint() -> None:
+    """An NVFP4 body keeps the FP8 PLE table's global scale (see #54765)."""
+    prefix = "model.language_model.layers.1.ple.ple_embedding.ngram_embedding"
+    quant_config = _nvfp4_config(["*.ple.*"])
+
+    assert isinstance(
+        _get_ple_embedding_quant_method(quant_config, prefix, "float8_e4m3fn"),
+        Qwen4ExpPLEFp8EmbeddingMethod,
+    )
+    assert isinstance(
+        _get_ple_embedding_quant_method(quant_config, prefix, torch.float8_e4m3fn),
+        Qwen4ExpPLEFp8EmbeddingMethod,
+    )
+
+
+@pytest.mark.parametrize(
+    "exclude_modules,ple_embedding_dtype",
+    [
+        # The table is excluded but stored unquantized.
+        (["*.ple.*"], None),
+        (["*.ple.*"], "bfloat16"),
+        # The table is not excluded, so NVFP4 shards are expected.
+        ([], "float8_e4m3fn"),
+    ],
+)
+def test_ple_fp8_embedding_skipped_for_non_fp8_nvfp4_tables(
+    exclude_modules: list[str], ple_embedding_dtype: object
+) -> None:
+    prefix = "model.language_model.layers.1.ple.ple_embedding.ngram_embedding"
+    quant_config = _nvfp4_config(exclude_modules)
+
+    assert (
+        _get_ple_embedding_quant_method(quant_config, prefix, ple_embedding_dtype)
+        is None
+    )
+
+
 def test_dilated_ple_spec_state_rolls_back_before_next_forward() -> None:
     conv_state_len = 6
     dilation = 2

vllm/models/qwen4_exp/nvidia/ple_layer.py

@@ -24,6 +24,7 @@
 from vllm.model_executor.layers.quantization.fp8 import Fp8Config
 from vllm.model_executor.layers.quantization.modelopt import (
     ModelOptMixedPrecisionConfig,
+    ModelOptNvFp4Config,
 )
 from vllm.model_executor.layers.quantization.utils.fp8_utils import (
     create_fp8_scale_parameter,
@@ -130,9 +131,20 @@ def embedding(self, layer: nn.Module, input_: torch.Tensor) -> torch.Tensor:
         return F.embedding(input_, layer.weight)
 
 
+def _ple_checkpoint_is_fp8(ple_embedding_dtype: object) -> bool:
+    """Whether ``ple_embedding_dtype`` declares FP8 PLE shards."""
+
+    if isinstance(ple_embedding_dtype, str):
+        ple_embedding_dtype = getattr(
+            torch, ple_embedding_dtype.rsplit(".", 1)[-1], None
+        )
+    return isinstance(ple_embedding_dtype, torch.dtype) and is_fp8(ple_embedding_dtype)
+
+
 def _get_ple_embedding_quant_method(
     quant_config: QuantizationConfig | None,
     prefix: str,
+    ple_embedding_dtype: object = None,
 ) -> QuantizeMethodBase | None:
     """Select global-scale FP8 only for quantized PLE checkpoint shards."""
 
@@ -141,6 +153,15 @@ def _get_ple_embedding_quant_method(
             return Qwen4ExpPLEFp8EmbeddingMethod()
         return None
 
+    if isinstance(quant_config, ModelOptNvFp4Config):
+        # NVFP4 checkpoints exclude the PLE table and store it as FP8 shards
+        # with one global scale, which ``ple_embedding_dtype`` records.
+        if not quant_config.is_layer_excluded(prefix):
+            return None
+        if not _ple_checkpoint_is_fp8(ple_embedding_dtype):
+            return None
+        return Qwen4ExpPLEFp8EmbeddingMethod()
+
     if not isinstance(quant_config, Fp8Config):
         return None
     if not quant_config.is_checkpoint_fp8_serialized:
@@ -326,7 +347,9 @@ def __init__(
             padding_size=divisor,
             prefix=f"{prefix}.ngram_embedding",
             quant_method=_get_ple_embedding_quant_method(
-                quant_config, f"{prefix}.ngram_embedding"
+                quant_config,
+                f"{prefix}.ngram_embedding",
+                getattr(config, "ple_embedding_dtype", None),
             ),
         )