Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
34 changes: 34 additions & 0 deletions tests/models/qwen4_exp/test_ple.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,9 @@
import vllm.model_executor.layers.vocab_parallel_embedding as embedding_module
import vllm.model_executor.parameter as parameter_module
from vllm.model_executor.layers.quantization.fp8 import Fp8Config
from vllm.model_executor.layers.quantization.modelopt import (
ModelOptMixedPrecisionConfig,
)
from vllm.models.qwen4_exp.common.ple import (
PLEShardOverlap,
compute_ple_shard_overlap,
Expand Down Expand Up @@ -288,6 +291,37 @@ def test_ple_fp8_embedding_respects_checkpoint_shard_exclusions() -> None:
assert _get_ple_embedding_quant_method(quant_config, prefix) is None


def test_ple_fp8_embedding_supports_mixed_precision_config() -> None:
prefix = "model.language_model.layers.1.ple.ple_embedding.ngram_embedding"
quant_config = ModelOptMixedPrecisionConfig.from_config(
{
"quantization": {
"quant_algo": "MIXED_PRECISION",
"exclude_modules": [],
"group_size": 16,
"quantized_layers": {
prefix: {"quant_algo": "FP8"},
"model.language_model.layers.2.moe.gate_proj": {
"quant_algo": "NVFP4"
},
},
}
}
)

assert isinstance(
_get_ple_embedding_quant_method(quant_config, prefix),
Qwen4ExpPLEFp8EmbeddingMethod,
)
assert (
_get_ple_embedding_quant_method(
quant_config,
"model.language_model.layers.2.moe.gate_proj",
)
is None
)


def test_dilated_ple_spec_state_rolls_back_before_next_forward() -> None:
conv_state_len = 6
dilation = 2
Expand Down
8 changes: 8 additions & 0 deletions vllm/models/qwen4_exp/nvidia/ple_layer.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,9 @@
QuantizeMethodBase,
)
from vllm.model_executor.layers.quantization.fp8 import Fp8Config
from vllm.model_executor.layers.quantization.modelopt import (
ModelOptMixedPrecisionConfig,
)
from vllm.model_executor.layers.quantization.utils.fp8_utils import (
create_fp8_scale_parameter,
create_fp8_weight_parameter,
Expand Down Expand Up @@ -133,6 +136,11 @@ def _get_ple_embedding_quant_method(
) -> QuantizeMethodBase | None:
"""Select global-scale FP8 only for quantized PLE checkpoint shards."""

if isinstance(quant_config, ModelOptMixedPrecisionConfig):
if quant_config._resolve_quant_algo(prefix) == "FP8":
return Qwen4ExpPLEFp8EmbeddingMethod()
return None

if not isinstance(quant_config, Fp8Config):
return None
if not quant_config.is_checkpoint_fp8_serialized:
Expand Down
Loading