diff --git a/tests/models/qwen4_exp/test_ple.py b/tests/models/qwen4_exp/test_ple.py index d7747ebfbdb5..7c804a239d86 100644 --- a/tests/models/qwen4_exp/test_ple.py +++ b/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 diff --git a/vllm/models/qwen4_exp/nvidia/ple_layer.py b/vllm/models/qwen4_exp/nvidia/ple_layer.py index 9d7091d7a2b3..8ce776d889f0 100644 --- a/vllm/models/qwen4_exp/nvidia/ple_layer.py +++ b/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), ), )