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

[stecasta]

Purpose

Fixes #54765.

ModelOpt NVFP4 checkpoints of Qwen3.8-Flash-Next exclude *.ple.* from NVFP4 and ship the 51.2B-parameter n-gram table separately as FP8 shards with one global scale. That is exactly the layout Qwen4ExpPLEFp8EmbeddingMethod implements, but _get_ple_embedding_quant_method only selects it for Fp8Config or ModelOptMixedPrecisionConfig. Under a plain ModelOptNvFp4Config the unquantized embedding is built, no weight_scale parameter is registered, and loading fails:

ValueError: There is no module or parameter named 'ngram_embedding.weight_scale' in
Qwen4ExpNGramEmbedding. The available parameters belonging to ngram_embedding
(PLEVocabParallelEmbedding) are: {'ngram_embedding.weight'}

#54882 fixed this for quant_algo: MIXED_PRECISION. Checkpoints that declare quant_algo: NVFP4 still fail, which is the layout of the most widely used NVFP4 build of this model.

Fix

Select the FP8 embedding method under ModelOptNvFp4Config when both hold:

1. the PLE layer is excluded from NVFP4 (is_layer_excluded), and
2. the config declares the table as FP8 via ple_embedding_dtype.

Requiring both means a checkpoint that excludes the PLE but keeps it in bf16 still takes the unquantized path. ple_embedding_dtype is the same signal #53899 uses for this.

Validation

Unit level, plus checkpoint metadata inspection. I did not run a full serving load of a 135 GB checkpoint.

pytest tests/models/qwen4_exp/test_ple.py -k ple_fp8_embedding -v   # 9 passed
pytest tests/models/qwen4_exp/test_ple.py                           # 14 passed, 23 skipped
pre-commit run --files <changed files>                              # all hooks passed

Run inside vllm/vllm-openai:nightly-8a728663c1 (0.28.1rc1.dev388+g8a728663c) on both an sm_121 and an sm_120 host, unpatched then patched. The new tests fail against the unpatched selector and pass with the fix; the five pre-existing PLE tests pass in both.

Note on the red-before-green: the unpatched failures are TypeError from the added parameter, not assertion failures. test_ple_fp8_embedding_loads_under_nvfp4_checkpoint is the one that encodes the bug, since the unpatched selector returns None for this checkpoint. The three skipped_for_non_fp8 cases assert behaviour the unpatched code also has, so they are regression guards rather than proof of the fix.

Table dtypes read from the safetensors headers of published checkpoints:

checkpoint	ple_embedding_dtype	table on disk
RadixArk/Qwen3.8-Flash-Next-NVFP4	float8_e4m3fn	F8_E4M3 + weight_scale
dealignai/Qwen3.8-Flash-Next-ABLITERATED-NVFP4	float8_e4m3fn	F8_E4M3 + weight_scale
Inferact/Qwen3.8-Flash-Next-NVFP4, nota-ai/Qwen3.8-Flash-Next-Nota-NVFP4	absent	BF16

Note on ple_embedding_dtype

primitive-ai/Qwen3.8-Flash-Next-NVFP4 and primitive-ai/Qwen3.8-Flash-Next-mixed-NVFP4-FP8 declare ple_embedding_dtype: float8_e4m3fn but ship bf16 tables with no weight_scale. They load today and would fail at startup after this change, with the clear scale-sentinel error added in #54722. I think the config is wrong rather than the gate, but flagging it since it is a visible behaviour change for those two repos. Happy to key on something else if a maintainer prefers.

Why this is not a duplicate

I ran the checks in AGENTS.md. No open PR references #54765. The same selector is touched by the PLE offload PRs (#53899, #54371, #54070, #54129), which are large unmerged features; this is the standalone load fix and does not conflict with the direction they take.

Model evaluation

Not applicable. Affected checkpoints currently fail to load, so there is no baseline to regress against. The change only decides which method reads an embedding table that the checkpoint already stores in FP8; unaffected checkpoints keep their existing path, which the second test pins.

AI assistance

AI assistance was used to author this change. I have reviewed every changed line, ran the tests and linters above myself, and can defend the design.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Repository UI

Review profile: CHILL

Plan: Team

Run ID: cdca1949-c626-4fb1-a937-82a5bbe7d58a

📥 Commits

Reviewing files that changed from the base of the PR and between 8340fe1 and 5d2003b.

📒 Files selected for processing (2)

• tests/models/qwen4_exp/test_ple.py
• vllm/models/qwen4_exp/nvidia/ple_layer.py

Included review availability: Your plan provides up to 10 included reviews per hour; 9 remain after this review.

---

📝 Summary

Summary by CodeRabbit

• New Features

  • Added support for NVFP4 checkpoints containing FP8 PLE embeddings.
  • FP8 PLE embeddings remain available when PLE modules are excluded from NVFP4 quantization.

• Bug Fixes

  • Improved validation of PLE embedding configurations, rejecting unsupported non-FP8 or non-excluded NVFP4 tables.

Walkthrough

The PLE embedding selector now supports FP8 PLE tables in ModelOpt NVFP4 checkpoints. The constructor passes the checkpoint PLE dtype to the selector. Tests cover excluded and non-excluded tables with FP8 and BF16 configurations.

Changes

NVFP4 PLE support

Layer / File(s)	Summary
NVFP4 PLE method selection vllm/models/qwen4_exp/nvidia/ple_layer.py	The selector identifies FP8 PLE checkpoint dtypes and selects Qwen4ExpPLEFp8EmbeddingMethod when ModelOpt NVFP4 excludes the PLE layer.
PLE wiring and selection tests vllm/models/qwen4_exp/nvidia/ple_layer.py, tests/models/qwen4_exp/test_ple.py	Qwen4ExpNGramEmbedding passes ple_embedding_dtype to the selector. Tests cover FP8, BF16, excluded, and non-excluded NVFP4 configurations.

Estimated code review effort: 2 (Simple) | ~10 minutes

Merge Risk: ⚪ Minimal · up to 5d200

This enables loading supported NVFP4 checkpoints with separately stored FP8 PLE tables while preserving existing behavior for BF16 and non-excluded PLE tables. The covered selection paths indicate the change is ready to merge.

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 44.44% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 9 functions across 2 files.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly identifies the bug fix: loading FP8 PLE n-gram tables from ModelOpt NVFP4 checkpoints.
Description check	✅ Passed	The description directly explains the loading failure, the conditional fix, affected checkpoint layouts, and validation results.
Linked Issues check	✅ Passed	The implementation satisfies issue #54765 by selecting the FP8 PLE embedding method only when the PLE layer is excluded from NVFP4 and declares an FP8 dtype. Non-FP8 PLE tables retain the unquantized …
Out of Scope Changes check	✅ Passed	The changes are limited to the PLE quantization selector and focused tests for ModelOpt NVFP4 checkpoint loading. No unrelated code changes are present.

• Fix all pre-merge checks with AI

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @stecasta.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork