[Quantization] Support packed NVFP4 Qwen4Exp PLE embeddings

[stecasta]

Purpose

Enable the local-inference-lab/Qwen3.8-Flash-Next-NVFP4 checkpoint to run on 1x DGX Spark without CPU or disk PLE offloading. Its packed NVFP4 ngram table occupies 26.822 GiB and remains resident in CUDA memory. The PR's resident path was validated on one GB10 with a 4,096-token context, 2 GiB KV cache, and MTP disabled; model loading reported 97.47 GiB.

Fixes #56272.

Support ple_embedding_dtype="nvfp4" in Qwen4Exp ngram embeddings. The local-inference-lab checkpoint stores each logical 160-value row as 80 packed E2M1 bytes and 10 E4M3 block scales, with one FP32 global scale. Upstream currently selects an unquantized table for this explicit storage format inside a ModelOpt mixed-precision checkpoint.

The change adds an embedding method to the existing device/pinned-host storage interface. Both paths keep packed weights and block scales in their original storage and decode only the requested rows. Lookup reuses the existing E2M1 decoder and produces activation-dtype rows before ETP reduction. The loader copies each checkpoint shard's overlap with the local ETP vocabulary range and validates coverage after all weight streams finish, along with the global scale, tensor shapes, and storage dtypes.

Related work was checked on September 10, 2026. #54882 and #55334 concern an FP8 PLE table inside a model with NVFP4 body weights. #53899 contains an older NVFP4 PLE implementation within offload work, but its mixed-precision dispatch does not select this checkpoint's explicit NVFP4 table. #42791 adds generic quantized vocabulary embeddings. This change targets the serialized split PLE layout through the storage interface merged in #54371.

AI assistance was used to develop and test this change.

Test Plan

Extend the existing PLE suite with streamed packed-weight/block-scale loading across ETP boundaries, missing row coverage, malformed storage/global scales, independent numerical decoding, compiled GPU lookup, CUDA graph replay with changing IDs, empty/padded inputs, and pinned-host prefetch. Add the NVFP4 cases to the existing H200 Qwen4Exp CI lane.

Test base: b28c3e1568bfae930f61d4b24940e47528c85d4a, with the matching upstream precompiled CUDA wheel, Torch 2.13.0+cu130, and FlashInfer 0.6.18.post1. GPUs: one isolated RTX PRO 6000 Blackwell Max-Q and one DGX Spark / GB10. The Spark run used the published PR commit e251861aab69d3eb438d1a76ddec036ae45cfd4c.

.venv/bin/python -m pytest -v --tb=short --timeout=300 \
  tests/models/qwen4_exp/test_ple.py tests/models/qwen4_exp/test_config.py
.venv/bin/python -m pytest -v --tb=short --timeout=600 \
  tests/compile/h100/test_startup.py::test_moe_startup
.venv/bin/python -m pre_commit run --files \
  vllm/models/qwen4_exp/nvidia/ngram_embedding.py \
  tests/models/qwen4_exp/test_ple.py .buildkite/test_areas/models_basic.yaml
.venv/bin/python -m pre_commit run mypy-3.12 --hook-stage manual --files \
  vllm/models/qwen4_exp/nvidia/ngram_embedding.py \
  tests/models/qwen4_exp/test_ple.py

Full-model text-generation validation used the pinned checkpoint above (image/video inputs disabled), a 4,096-token context, 512-token prefill chunks, four sequences, 2 GiB of KV cache, and full/piecewise CUDA graph capture sizes [1, 2, 4]. The offline harness ran with PYTHONPATH=/opt/ngram/vllm:/run/ngram /opt/ngram/.venv/bin/python /run/ngram/model_smoke.py on Max-Q and model_smoke_spark.py on Spark. Spark used engram_config={"cpu_offload":false} and cpu_offload_gb=0; Max-Q used pinned-host PLE storage. It checked selected rows from the complete table against an independent decoder, then four chat prompts and the first 32 canonical GSM8K test examples (temperature 0, thinking disabled, 768 output tokens).

Test Result

• PLE and configuration suites: 76 passed on Max-Q and 76 passed on DGX Spark.
• Compilation startup canary on Max-Q: 2 passed.
• Actual checkpoint sample: 8 rows of width 160 matched an independent E2M1/FP8/global-scale reference exactly after BF16 conversion, through both device lookup and pinned-host prefetch with simulated ETP ownership.
• CPU PLE/configuration suites: 43 passed, 33 skipped because CUDA was unavailable.
• Pre-commit, including Python 3.10 type checks, and manual Python 3.12 type checks: passed.
• Full checkpoint loaded with 26.822 GiB of packed PLE weights/scales in CUDA memory on Spark and pinned host memory on Max-Q. On both devices, selected rows from the full table, including addresses beyond 32-bit offsets, matched the independent BF16 reference exactly.
• Full-model startup captured all six CUDA graphs on each device. All four generation checks passed on each, including a prompt requiring chunked prefill.
• GSM8K smoke evaluation: 31/32 exact numeric answers (96.875%) on each GPU. On each device, the remaining response hit the 768-token limit. This is a small smoke evaluation, not a full benchmark or comparison with the checkpoint publisher's runtime.

The checkpoint uses preview architecture names. Validation normalizes them with hf_overrides={"model_type":"qwen4_exp","architectures":["Qwen4ExpForConditionalGeneration"]}. PLE offload is disabled on Spark and enabled on Max-Q. It selects kernel_config={"linear_backend":"marlin"} because the automatic MXFP8 linear backend rejects the checkpoint's 96-output GDN projection. Real multi-GPU collectives and MTP have not been validated with this PR.

Single-Spark reference measurements

Separate end-to-end measurements on September 10, 2026 compared the local-inference-lab checkpoint using its b12x runtime with the NVIDIA checkpoint using upstream vLLM plus disk PLE offload patches. These measurements provide context for fitting the checkpoints on one DGX Spark. They were collected independently of this PR and do not measure its performance impact.

Workload: 32,768 shared prefix tokens, 2,048 suffix tokens, and 256 output tokens; streaming chat, temperature 0, three warmups per cell, and concurrency 1/4/8 with 10/40/80 measured requests. Both runtimes used one GB10, TP1/PP1, a 40,960-token context limit, prefix caching, chunked prefill, and MTP with four draft tokens. Latency and TTFT are means; decode throughput is aggregate active decode throughput from AIPerf.

Checkpoint / runtime	Concurrency	Requests	Latency (s)	TTFT (s)	Decode (tok/s)	Acceptance length
local-inference-lab / b12x	1	10	11.27	3.40	32.41	2.40
local-inference-lab / b12x	4	40	28.19	5.44	36.31	2.13
local-inference-lab / b12x	8	80	45.99	6.75	43.44	2.20
NVIDIA / vLLM with disk PLE	1	10	14.30	3.45	23.51	2.64
NVIDIA / vLLM with disk PLE	4	40	29.60	5.09	34.25	2.39
NVIDIA / vLLM with disk PLE	8	80	44.34	7.12	45.92	2.59

Each checkpoint completed all 130 measured requests with zero errors and zero preemptions. Every measured response contained 256 output tokens. Acceptance length is 1 + accepted_tokens / draft_steps, calculated from server counters.

The local checkpoint's packed PLE payload is 26.822 GiB, held in b12x's resident managed CUDA weight storage with no explicit CPU or disk PLE offload; KV cache was 8 GiB FP8. The NVIDIA checkpoint's FP8 PLE payload is 47.684 GiB, served through disk-backed mmap and the OS page cache; KV cache was 10 GiB BF16. Quantization, backends, and prefix reuse also differ (27,216 versus 30,704 reused tokens per request), so this is a comparison of complete configurations. Each cell is a single run, and accuracy equivalence was not evaluated. Host swap activity was observed, so no swap-free claim is made.

Reference runtime revisions

• Local runtime: vLLM aba94d396, b12x 75ffee637.
• NVIDIA runtime: upstream 73029d424, with #54129 at 50a061f79, #54882 at d4d703caf, #55513 at 60ad959b6, and #56026 at 4a365dee4.
• AIPerf 00ba1c5db; Torch 2.13.0+cu130, FlashInfer 0.6.18, NVIDIA driver 580.95.05.

[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.