[Model][Qwen4Exp] Split the pinned-host PLE table between device and host (#479)

[Peuqui]

Purpose

Third slice of the PP enablement offered in #479 (loader fix: #485, PLE gate on the partition: #516), and the placement we run Qwen3.8 Flash Next with.

The pinned-host PLE path (#345, #374) keeps the whole FP8 n-gram shard in pinned host memory and gathers every row over UVA. Where the card has room beside its layers that is the slow option, and on a box with little host RAM it does not boot at all: the Flash Next table is 51 GB, a TP2 rank pins 25 GB. #471 answers the memory side with a disk-backed table.

This PR adds the option in between: rows stay in device memory as long as the stage's weights, the KV cache of the requested max_model_len and a reserve for the activation peak and the graph pool still fit; only the remainder is pinned on the host. The budget is derived from the real headroom when the first checkpoint shard arrives (every other weight of the stage is placed by then), or fixed per rank with VLLM_QWEN4EXP_PLE_HOST_GIB. Both halves are gathered by the existing qwen4_exp_ple_pinned_gather kernel; a row-less half is skipped without a device-to-host sync. The loader routes each shard into the half that owns its rows. Qwen4ExpPinnedHostEmbedding therefore starts with a row-less CPU placeholder and materializes its tables lazily (first shard, or prepare_accelerator_weight() for dummy weights).

Second commit: _should_use_pinned_host_ple() asked for the capability of device 0 of the visible list. On a heterogeneous pipeline the stage that owns the PLE layers is not necessarily that device. It now asks for the worker's own device, and every pre-Ampere card takes the split placement (the generic path dequantizes the whole FP8 table to fp16 inside the compiled graph, 47.7 GiB for Flash Next).

This is how we serve Qwen3.8 Flash Next NVFP4 at TP2/PP2 on 2x RTX 8000 + 2x V100 with VLLM_QWEN4EXP_PLE_HOST_GIB=6: 51.9 tok/s with MTP k=4, coherence 8/8 (measured 2026-08-28 on our fork, the same placement is in our 1.5.0 deployment). We cannot boot current main on this hardware, so the end-to-end evidence is from the fork; the unit tests below run against this branch.

Known limits: the automatic budget estimates the activation peak and the graph pool through a reserve (8 % of the device, at most 4 GiB, VLLM_QWEN4EXP_PLE_VRAM_RESERVE_GIB); the explicit host budget is the safe choice. The gather always reads both halves.

Test Plan

Wheel venv (1Cat 1.5.0) with PYTHONPATH on this checkout, V100 (SM70):

pytest tests/models/qwen4_exp/test_ple.py
  34 passed (3 of them exact-SM70: split gather with 0, 4 and 8 host rows)
Negative proof: pre-change ple_layer.py with the new tests
  6 failed (the six pinned-host tests), 24 passed
pytest tests/models/qwen4_exp/ tests/v1/worker/test_ple_offload_worker.py tests/v1/spec_decode/test_qwen4_exp.py --ignore=tests/models/qwen4_exp/test_qsa_amd.py
  234 passed, 3 failed -- the three (test_qsa_reference.py side/circular/compressed metadata)
  fail identically on unmodified 755baae; test_qsa_amd.py cannot be collected in the same
  process as the nvidia tree (duplicate custom-op registration), also on 755baae
pre-commit run --files <4 files>          all hooks passed (incl. check-torch-cuda-call, mypy-local)
pre-commit run mypy-3.10 --hook-stage manual --files <4 files>   passed

Duplicate check / AI assistance

gh pr list --state all --search "PLE host" / "pinned host PLE" / "PLE placement" / "PLE_HOST": #345, #374 (whole shard in pinned host), #471 (disk-backed table) -- all merged, none splits the table between device and host; no open PR touches this path. gh issue view 479 --comments: the third-slice offer is in the issue.

Written with AI assistance (Claude); every line was reviewed and the tests were run by me on 2x RTX 8000 + 3x V100.

🤖 Generated with Claude Code