[Peuqui]

Purpose

With the PLE disk cascade, which is on by default (KernelConfig.ple_disk_cascade) and admitted for FP16 Qwen4Exp models with FP8 or packed GGUF PLE storage on local multiprocessing workers, the n-gram embedding is computed without its resident rows. PleOffloadLayer.forward asks offload_keeps_local_tables() on every call whether to gather the layer's own rows and merge the offload worker's rows into them. For Qwen4ExpNGramEmbedding that answer is ple_cascade_configured() or pinned_decode_active(), and both read kernel_config from get_current_vllm_config_or_none(). The GPU worker makes the config current only around worker construction, init_device, initialize_from_config and load_model. Profiling, torch.compile tracing, warmup, CUDA graph capture and execute_model run without it (vllm/models/qwen4_exp/nvidia/model.py:1164 notes the same for capture), so every forward, eager or compiled, answers "no" and returns the offload worker's raw output buffer. The worker writes only the disk-tier slots of that buffer (_remote_lookup); every slot whose row is resident keeps what the buffer held: zeros from the dummy and capture runs, or rows left over from earlier chunks and requests.

The cascade reached main in #806 together with 41faaa6, which moved ple_cascade_configured() from VLLM_QWEN4EXP_PLE_DISK to the kernel config, so on main the merge has not run at runtime since the cascade exists. How much is affected depends on the placement and the prompt. On our PP4 entry (131,045,319 of 320,001,536 rows in device memory) it was 10,592 of the 25,856 n-gram slots in the first 1,616-token chunk of a German prompt (41 %). On TP2xPP2 (125,439,175 of 160,000,768 rows per rank in device memory) most rows are resident. From the source only, not run: when a placement leaves no rows on disk, the worker zeroes the whole buffer, so the n-gram embedding is all zeros. Because the stale rows depend on earlier requests, the same prompt got different log-probabilities on every pass within one boot.

The guarded constructor already decides once, while the config is current, whether the layer builds its resident tables. This PR records that decision in the instance and lets the forward use it, the same way Qwen4ExpNGramEmbedding already keeps _cascade and _pinned_decode from its constructor, and the way #1004 and #857 moved policy reads to layer construction. Not changed: the packed GGUF pinned path, whose forward returns through the _pinned_decode branch before this check; the hybrid policy, which the cascade admission excludes and for which the recorded value equals the old runtime answer; and Qwen4ExpNGramEmbedding.load_weights, which repeats the guard but runs inside load_model's config.

The new test builds a layer while its keeps-local-tables answer is true, then flips the answer as the missing runtime config does. Its forward_impl takes resident ids itself and the remaining ids from a worker buffer in which only the worker's slots are written; on main the forward returns that buffer and the resident slot comes back as 0.

Not a duplicate: gh pr list --repo 1CatAI/1Cat-vLLM --state all --search with "offload_keeps_local_tables", "ple_disk_cascade_active", "resident PLE rows", "PLE cascade forward", "wait_offloaded_output", "Flash-Next nondeterministic" and "PLE zeros", and the same searches over issues, find no PR or issue that reports or fixes this. Related but different: #806 (merged, brought the cascade and 41faaa6), #856 (merged, cascade storage lookup under checkpoint layer names), #818 (merged, whole-table placeholders without cascade state), #1004, #857. The open PRs touching the same two files (#828, #1020, #1023, #1094, #1167, #1214, #1217) do not change this check, and this PR adds no merge conflict to any of them (git merge-tree against each PR head; #1020 and #1023 conflict with main already).

AI assistance was used for this change. I reviewed every line and ran the tests and measurements below.

Test Plan

pytest tests/v1/worker/test_ple_offload_worker.py tests/models/test_qwen4_ple_placeholder.py tests/models/test_ple_optional_cuda_import.py tests/models/qwen4_exp/test_gguf_pinned_policy.py tests/models/qwen4_exp/test_ple.py tests/compile/test_sm70_decode_graph.py tests/config/test_ple_cascade_capabilities.py tests/kernels/ple
pre-commit run --from-ref origin/main --to-ref HEAD
pre-commit run mypy-3.10 --hook-stage manual --from-ref origin/main --to-ref HEAD

Row check against the checkpoint: Qwen3.8-Flash-Next-180B-A4B NVFP4 with MTP, PP4 over two Tesla V100 and two Quadro RTX 8000, disk cascade, --enforce-eager. A forward hook saved the n-gram embedding output and the worker's n-gram ids for a 3,708-token German prompt (prefill chunks 1,616 + 1,616 + 476 tokens); each slot was compared with the row read directly from the checkpoint's FP8 table (dequantized with the table scale for the fixed path, which returns the dequantized embedding).

Teacher-forced log-probabilities: the same prompt with prompt_logprobs (3,707 scored positions), two passes in one boot, same setup.

Serving A/B on our fork (main 6f3557c plus our commits): fork commit 6bbcd40 (this change) against cf50e1a (its parent), PP4 and TP2xPP2, torch.compile with CUDA graphs (FULL_AND_PIECEWISE, capture sizes 1, 2, 4, 5, 8), MTP with 4 speculative tokens. Boots fix cold, base, fix, base, fix; the cold boot is not counted. The base arm runs the production tree, the fix arm the same entries with PYTHONPATH on the fork tree and its own compile cache; the native build is the same in both arms (this PR changes Python only). Per boot: greedy answers to three prompts, four ~29k-token prompts, a decode benchmark (prose, code, edit), eight quality prompts at temperature 0, and prompts of 184,134 and 249,992 tokens.

Path coverage

Path / provider	Format	Key conditions (shape, TP/EP, graph, capability)	Expected selection / fallback	Evidence (source / static / executed / measured)	Unverified scope
Qwen4Exp n-gram PLE, disk cascade	FP8 table in safetensors (NVFP4 checkpoint)	PP4 and TP2xPP2, V100 and Quadro RTX 8000; eager (row check) and CUDA graphs (serving)	layer merges its resident rows with the worker's disk rows	executed + measured (row check, log-probabilities, serving A/B)	TP4
Qwen4Exp n-gram PLE, disk cascade	packed GGUF, not pinned	worker writes only the disk-tier rows (ple_layer.py _remote_lookup), GPU merge in gguf_embedding.py	same as above	source	not run
Packed GGUF pinned decode	GGUF	_pinned_decode set by the constructor	unchanged: the forward returns through the pinned branch before this check	source	not run
Hybrid PLE policy	any	ple_policy().hybrid, bound at runtime; excluded from cascade admission	unchanged	source

Acceleration and benchmark contract (required for performance changes)

Not a performance change. Python only; no new switches, fields or environment variables. On main this PR puts the merge path into the compiled forward and the CUDA graph captures for the first time; on our fork it ran that way in the serving A/B below.

Test Result

• Tests on main 6f3557c plus this PR: 260 passed, 34 failed. Main with only the new test applied: 259 passed, 35 failed; the extra failure is the new test, failing on its assertion (the forward returns [[0], [5], [6]] instead of [[4], [5], [6]]). The remaining failures are identical with and without this PR and come from stale test fakes (CudaGraphManager has no capture_context, SimpleNamespace configs without kernel_config, the connector and ple_offload_layer modules without envs, Qwen4ExpNGramEmbedding fakes without layer_name); with [Test] Repair stale fakes in PLE offload and SM70 decode-graph tests #1214 applied as well: 284 passed, 10 failed, all ten in tests/kernels/ple/test_gguf_packed_embedding.py. pre-commit and mypy-3.10 pass.
• Row check, first 1,616-token chunk: base: disk-tier slots 15,264 of 15,264 equal the checkpoint row, resident slots 0 of 10,592. This PR: all 25,856 slots equal the dequantized checkpoint row (max abs difference 0). The other two chunks behave the same.
• Log-probabilities: base: mean NLL 1.4115 in pass 1 and 1.4023 in pass 2, equal log-probabilities at 1 of 3,707 positions, same top-1 token at 3,466. This PR: mean NLL 1.3609 in both passes, all 3,707 log-probabilities identical.
• Serving speed, fix (two boots) against base (two boots), no slowdown: decode step times are equal (PP4 edit 54 ms = 92.3 against 92.5 tok/s, TP2xPP2 edit 45 ms = 110.0 against 110.0 tok/s); time to first token for ~29k tokens is equal (PP4 7.5-7.6 s against 7.5-8.0 s, TP2xPP2 11.8-11.9 s against 11.8 s; the first long request after each boot not counted); decode step at 249,992 tokens PP4 66.8 and 67.4 ms against 67.0 and 68.1 ms, TP2xPP2 59.2 and 57.9 ms against 58.8 and 58.4 ms. Time to first token at 184,134 tokens was lower with the fix (PP4 71.9 and 72.2 s against 73.3 and 74.5 s, TP2xPP2 99.3 and 99.6 s against 107.6 and 111.1 s); we have no explanation for the TP2xPP2 difference and do not claim it as a speedup, the long prefills on this machine also vary with the page cache of the PLE disk tier. Prose tok/s follows the generated text through the MTP acceptance (PP4 47.1 against 45.6 tok/s, TP2xPP2 54.0 against 57.2 tok/s at acceptance 2.42 against 2.58; ~29k-token prompts PP4 37.5 against 35.5, TP2xPP2 43.8 against 41.8 tok/s).
• Outputs: the three fix boots per topology produced byte-identical greedy texts and quality answers, and the decode benchmark, which runs each prompt twice per boot, now gets the same token count and acceptance on both runs (base: 519 and 553 tokens for the same prose prompt). PP4 and TP2xPP2 now write the same greedy text for two of the three prompts and the same first 625 of 665 characters for the third; on base the two topologies share only the first 377, 211 and 66 characters. All quality answers read correct by hand (arithmetic, logic, code, summaries; a question about a non-existent person was declined without inventing a biography in both arms).

🤖 Generated with Claude Code

https://claude.ai/code/session_01YaG8VjXn5eN3Aq7eZrA3dp