Repository navigation
feat(inference): experimental HiSparse on P/D decode - #4023
Merged
Merged
Conversation
vLLM 0.31 captures routed experts through the AuxOutput connector and rejects it with P/D connectors: a decode instance never computes a remotely prefilled prompt, so its worker would materialize missing rows (BlockObjectStoreError), and the first forward starts mid-block. Carry our upstream AuxOutput fix as a plugin patch: - the config check skips the P/D connector rejection for NixlConnector (alone or inside MultiConnector); every other check still runs; - requests that arrive with do_remote_prefill emit rows only from their first forward (num_computed_tokens then); the floor survives preemption; - the worker does not buffer or publish the partially known first block. The pinned router drops as many leading decode rows as prefill returned, so the decode server left-pads its rows to routed_experts_prompt_start (the same payload size as 0.30's full-length decode rows). Remove the config rejection added by the 0.31 bump.
vLLM 0.31 keeps routing rows in a local store keyed by block hash. No connector moves rows with the KV, so a block filled by an offload load has no rows here, and a request that needs them fails the engine step (BlockObjectStoreError). - Load clamp: native (OffloadingConnector._max_loadable_tokens) and Mooncake (get_num_new_matched_tokens) loads stop at block_floor(routed_experts_prompt_start). vLLM never returns rows before that position, so multi-turn reloads keep their hits. Decode- side P/D requests are not limited (they emit from their first forward). - Loaded-block cut: loaded blocks then serve local prefix hits. A later request with a lower prompt start would need their rows, so its local hit is cut at the first loaded block it needs rows from. Clamping is simpler than restoring rows with offloaded blocks: the row store lives in the worker, the clamp is a scheduler-side bound, and neither connector has a hook to carry extra per-block objects. Remove validate_router_replay_without_kv_offload.
Per-role KV offload for disaggregated P/D (deployment.prefill_kv_cache_offload / decode_kv_cache_offload, inheriting inference.kv_cache_offload; None disables). Mooncake cpu.num_bytes = 0 connects a role to the shared store without a storage client; under P/D Mooncake sets save_decode_cache so decode KV is reusable by the next turn's prefill. HiSparse (experimental) runs on decode only; decode cannot offload KV with it (its host-resident layout collides with prefill's under Mooncake keys). The launcher writes one engine config per role.
…rward Anchor the routing rows at the end (start = prompt + completion - 1 - rows), so a P/D decode response no longer zero-pads its rows to the prompt start. The router cuts decode rows by position since v0.2.3 (router#61).
A request resumed after preemption emits routing rows only from num_tokens-1, which its first forward always recomputes, so it never reads stored rows.
kv_cache_offload_by_role is {'all': ...} without P/D, so the decode check missed
aggregated HiSparse + native offload, which vLLM rejects only at engine start.
…ked /generate path The start moves into _PackedOutputs (#3900), so routed_experts.py and serving_tokens.py stay as on main.
# Conflicts: # src/prime_rl/entrypoints/rl.py
samsja
pushed a commit
that referenced
this pull request
Oct 14, 2026
Router replay now works with disaggregated P/D on vLLM 0.31, together with sampling replay. It replaces main's config rejection of router replay + P/D. **Problem on 0.31.** vLLM 0.31 captures routed experts through the AuxOutput connector (V2 runner only, vllm#45635) and rejects P/D connectors (`_verify_aux_output_compatibility`). The rejection is real: a decode (D) instance never computes a remotely prefilled prompt, so its worker would try to materialize the prompt rows from its block store (`BlockObjectStoreError`), and D's first forward (token `prompt_len-1`) starts mid-block, which the routed-experts buffer rejects. **Fix (plugin patch `monkey_patch_routed_experts_with_pd_connectors`, our upstream AuxOutput patch 02):** - Config: skip only the P/D connector check when the KV transfer config has `NixlConnector` (alone or inside `MultiConnector`). Every other AuxOutput check still runs. - Scheduler: a request on D (`kv_transfer_params` with a `remote_engine_id` and no `do_remote_decode`; NIXL clears `do_remote_prefill` but keeps that) emits rows only from its first forward on D (`num_computed_tokens` at that point: normally `prompt_len-1`, lower after a failed KV load, where D recomputes those rows itself). The floor survives preemption and is dropped when the request finishes. - Worker: the partially known first block is not buffered or published; later blocks D computes in full are published as usual. **Router v0.2.3 required (router#61); the decode `start` comes from #3900 (merged).** On a decode instance the rows start at its first forward, so the response carries only the rows D computed. #3900's packed `/generate` anchors the rows at the end: `start = len(prompt) + len(completion) - 1 - len(rows)` (equal to `routed_experts_prompt_start` without P/D). This PR does not touch `routed_experts.py`/`serving_tokens.py`, so it and #3900 (and #3920/#3794 on it) land in either order; P/D router replay needs both. Router v0.2.3 cuts decode rows by position (`(prefill.start + prefill.seq_len) - decode.start`) and also handles uint16 routing (models with more than 256 experts). This PR pins `vllm-router` to v0.2.3 in `pyproject.toml`/`uv.lock`. Without #3900, the merge would drop real completion rows. Also: removes `validate_disaggregated_no_routed_experts` (both validators on main), docs. Stack: **this PR** → #3831 (router replay with KV offload) → #3832 (P/D per-role KV offload) → #4023 (draft, HiSparse on decode). ## Checks run - `pytest -m "not gpu" tests/unit/test_configs.py tests/unit/inference` on this head: 214 passed. `ruff check` / `ruff format --check`: clean. - New tests (`test_routed_experts_pd.py`): the scheduler floor (first forward, preemption, cleanup) on vLLM's real `AuxOutputSchedulerConnector`, and the worker skip on the real `RoutedExpertsBuffer`. `test_combined_replay_uses_v2_runner` is parametrized over single-node and P/D. - vLLM's own `tests/distributed/aux_output_connector/test_store.py` (v0.31.0) with this patch applied plus patch 02's two new tests: 66 passed. Without the patch the two new tests fail. ## GPU validation (2026-10-06, vLLM 0.31.0, 1x 8xH200, Qwen3-30B-A3B TP1, `VLLM_BATCH_INVARIANT=1`, router v0.2.2 with zero-padded decode rows) - 1P1D (NIXL) vs colocated, router replay + sampling masks: 10/10 identical at concurrency 1 (prompts 64-6000 tokens, 128 completion tokens) and 20/20 at concurrency 10 with a fresh `cache_salt`. Identical means tokens, prompt and completion routing rows (100%) and sampling masks; rows = prompt + completion - 1. - Forced decode preemption (D with 1300 blocks, 20 requests, 1000-4000 prompts, up to 1024 tokens, concurrency 10): 6 preemptions on D, 20/20 identical to colocated. - All engines log `Using V2 Model Runner`; no `BlockObjectStoreError`. Re-run (job 2900, 2026-10-06, same setup, launched from per-role configs written by #3832's `write_pd_configs`): 1P1D (NIXL, no offload) vs colocated, router replay + masks: 10/10 identical at concurrency 1 and 20/20 at concurrency 10 with a fresh `cache_salt`; 0 errors. Not run: injected NIXL load failure (`kv_load_failure_policy=recompute`), MTP on decode, an `rl` run with a P/D config. ## GPU validation with router v0.2.3 (2026-10-14) 1x 8xH200, vLLM 0.31.0, Qwen3-30B-A3B TP1, `VLLM_BATCH_INVARIANT=1`, router replay + sampling masks. prime-rl = #3830 (978146c) merged with #3900 (20a3225); router = `release/v0.2.3` (ff818da, this PR + version bump), wheel built locally. 1P1D (NIXL) through the router vs a colocated server. Identical means: same tokens, same routing rows for prompt and completion (100%, rows = prompt + completion - 1, merged `start` = 0), same sampling masks, same logprobs. - Concurrency 1 (prompts 64-6000 tokens, 128 completion tokens): 10/10 identical. - Concurrency 10, fresh `cache_salt` on the colocated side: 20/20 identical. - uint16: no knob exists (vLLM picks uint8 for <=256 experts), so a test-only `sitecustomize` forced vLLM's capturer to `uint16` on P, D and the colocated server. The responses carry `dtype: uint16`. Concurrency 1: 10/10 identical; concurrency 10: 20/20 identical. - Forced decode preemption (D with `num_gpu_blocks_override=1300`, 1000-4000 token prompts, up to 1024 tokens, concurrency 10): 6 preemptions on D, 20/20 identical. - Router logs: no errors or warnings during the runs (only the shutdown warning). Engines: all `Using V2 Model Runner`, no `BlockObjectStoreError`. - Negative control, router v0.2.2 with the same P/D pair: 8/10 requests fail (HTTP 500 `decode routed_experts has 128 rows, expected at least <prompt_len>`). The 2 prompts of 64 tokens return 128 rows instead of 191, and the completion rows are wrong. <!-- CURSOR_SUMMARY --> --- > [!NOTE] > **High Risk** > Changes expert-routing capture on NIXL P/D via vLLM monkey patches and removes prior safety guards; training correctness depends on router 0.2.3 merge behavior and patch edge cases (preemption, failed KV load). > > **Overview** > **Router replay** and **sampling replay** are now allowed on **disaggregated prefill/decode (P/D)** deployments, where they were previously blocked by config validation and vLLM’s AuxOutput + NIXL incompatibility. > > Config validators that rejected `enable_return_routed_experts` with `deployment.type = "disaggregated"` are removed from inference and RL configs. Docs now describe P/D behavior: prompt routing from prefill, completion routing from decode, and that sampling capture can run alongside router replay under P/D. > > A startup **vLLM monkey patch** (`monkey_patch_routed_experts_with_pd_connectors`) bypasses only the P/D connector check in `_verify_aux_output_compatibility`, raises an emit floor on decode so routed-expert rows start at the instance’s first forward, and skips buffering a partial first block after remote prefill. **`vllm-router` is pinned to v0.2.3** so the router can merge prefill/decode expert rows correctly. > > Unit tests cover scheduler floors, worker buffering, and V2 runner setup for both single-node and disaggregated configs. > > <sup>Reviewed by [Cursor Bugbot](https://cursor.com/bugbot) for commit e586ef8. Bugbot is set up for automated code reviews on this repo. Configure [here](https://www.cursor.com/dashboard/bugbot).</sup> <!-- /CURSOR_SUMMARY --> --------- Signed-off-by: faresobeid <111092724+faresobeid@users.noreply.github.com>
# Conflicts: # docs/inference.md # packages/prime-rl-configs/src/prime_rl/configs/rl.py # src/prime_rl/inference/patches.py
faresobeid
marked this pull request as ready for review
October 16, 2026 17:15
mikasenghaas
approved these changes
Oct 19, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Experimental vLLM HiSparse for DSA sparse-MLA models (e.g. GLM-5.x) on P/D decode instances.
[inference.hisparse] host_pool_gib = ...addsHiSparseConnectorto the connector chain. Off by default; no example config uses it.for_pd_role("prefill")drops it).inference.deployment.decode_kv_cache_offload = "None". A HiSparse decode keeps its KV in a host-resident layout. Native offload rejects it, and Mooncake stores it under the same keys as prefill's regular layout, so prefill fails to load it (TRANSFER_FAILon GPU).cudagraph_mode = "FULL", while P/D decode launches withFULL_DECODE_ONLY).docs/inference.md.Checks
pytest -m "not gpu" tests/unit/orchestrator tests/unit/test_configs.py tests/unit/inference: 317 passed. Ruff clean.Not validated: real GLM-5.x quality at long prompts, aggregated HiSparse against the trainer forward, preemption or failed loads with HiSparse.
Note
Medium Risk
Changes P/D decode KV connector composition and host-memory sizing; misconfiguration is blocked at validate time, but the feature is experimental and not fully validated on production GLM workloads.
Overview
Adds experimental vLLM HiSparse support for DSA sparse-MLA models (e.g. GLM-5.x) via
[inference.hisparse] host_pool_gib = ..., which injectsHiSparseConnectorinto the existingkv_transfer_config/MultiConnectorchain alongside NIXL and optional Mooncake offload.Under disaggregated P/D, HiSparse applies to decode only (
for_pd_role("prefill")clears it). Config validation rejects incompatible setups: decode cannot use KV offload while HiSparse is on (native offload is also rejected globally when HiSparse is enabled). Docs add a HiSparse section (experimental caveats, RAM budgeting, offload rules). SLURM templates only note HiSparse in the pinned-memoryulimitcomment. Unit tests cover decode connector composition and validation errors.Reviewed by Cursor Bugbot for commit 7d44741. Bugbot is set up for automated code reviews on this repo. Configure here.