[KV Connector] Fix PD async scheduling race condition for hybrid attn models by arpera · Pull Request #48481 · vllm-project/vllm

arpera · 2026-07-13T09:15:35Z

Purpose

IMPORTANT: PR was fully copied from #47373

This PR fixes an accuracy regression in PD disaggregation with async scheduling for hybrid attention + SSM/GDN models.
When a decode-side request receives KV blocks from a remote prefill worker, the newly allocated attention blocks are overwritten by the incoming NIXL transfer. For hybrid models, vLLM also schedules KV cache zeroing for newly allocated attention blocks. Under async scheduling, that zeroing kernel can be queued behind a previous forward pass while the NIXL RDMA transfer writes directly into GPU memory. If the RDMA transfer completes first, the delayed zeroing kernel can erase the received attention KV, causing severe accuracy degradation.
The fix makes NIXL metadata report the local attention blocks that will be overwritten by remote KV loads. The scheduler then excludes those blocks from , while preserving zeroing for other newly allocated attention blocks.
This keeps the change scoped to connector-loaded attention blocks and avoids adding a CUDA synchronization point to the async pipeline.

Details about this fix

Initially the issue with async scheduling on PD Qwen3.5 was reported in these two bug reports:

Then these issues were addressed by this PR #45096.
But this PR #45096 introduces two different fixes (see this comment):

Allocation-step zeroing from [BUGFIX][Mamba][Qwen3.5] Zero freed SSM cache blocks on GPU #35219 wiping already-landed KV
Freed-block reuse while an over-scheduled in-flight step still rewrites the old owner's mamba state

During discussion in PR #45096 maintainers decided firstly to fix (2) race condition (comment) using another approach from #45357, so PR #45096 was closed because of this reason.

Recently fix for (1) race condition was published #47373, so this current PR #48481 finishes it's integration.

Test Result

Qwen3.5 PD with TP==1

Fixes this bug #42182
Hardware: 4xGB200
Model: Qwen3.5-35B-A3B

Details

# prefill
VLLM_SSM_CONV_STATE_LAYOUT=DS \
VLLM_NIXL_SIDE_CHANNEL_HOST="$(hostname)" \
VLLM_NIXL_SIDE_CHANNEL_PORT=6550 \
vllm serve Qwen/Qwen3.5-35B-A3B \
  --served-model-name Qwen3.5-35B-A3B \
  --host 0.0.0.0 --port 8100 \
  -tp 1 --generation-config vllm --no-enable-prefix-caching --no-disable-hybrid-kv-cache-manager \
  --async-scheduling \
  --kv-transfer-config '{"kv_connector":"NixlConnector","kv_role":"kv_both","kv_load_failure_policy":"fail"}'

# decode
VLLM_SSM_CONV_STATE_LAYOUT=DS \
VLLM_NIXL_SIDE_CHANNEL_HOST="$(hostname)" \
VLLM_NIXL_SIDE_CHANNEL_PORT=6551 \
vllm serve Qwen/Qwen3.5-35B-A3B \
  --served-model-name Qwen3.5-35B-A3B \
  --host 0.0.0.0 --port 8200 \
  -tp 1 --generation-config vllm --no-enable-prefix-caching --no-disable-hybrid-kv-cache-manager \
  --async-scheduling \
  --kv-transfer-config '{"kv_connector":"NixlConnector","kv_role":"kv_both","kv_load_failure_policy":"fail"}'

# router
python tests/v1/kv_connector/nixl_integration/toy_proxy_server.py \
  --host 0.0.0.0 --port 8000 \
  --prefiller-hosts "$PREFILL_HOST" --prefiller-ports 8100 \
  --decoder-hosts "$DECODE_HOST" --decoder-ports 8200

# gsm8k test
python tests/evals/gsm8k/gsm8k_eval.py

Results:

python tests/evals/gsm8k/gsm8k_eval.py
1319 questions, 5-shot
Evaluating: 100% | 1319/1319
Results:
Accuracy: 0.865
Invalid responses: 0.000
Total latency: 27.663 s
Questions per second: 47.680
Total output tokens: 188209
Output tokens per second: 6803.552

Qwen3.5 PD with TP>1

Fixes this bug #37285 (ANOTHER BUG!!!)
Hardware: 4xGB200
Model: Qwen3.5-397B-A17B-NVFP4

Details

# prefill
VLLM_SSM_CONV_STATE_LAYOUT=DS \
VLLM_NIXL_SIDE_CHANNEL_HOST="$(hostname)" \
VLLM_NIXL_SIDE_CHANNEL_PORT=6550 \
vllm serve nvidia/Qwen3.5-397B-A17B-NVFP4 \
  --served-model-name Qwen3.5-397B-A17B-NVFP4 \
  --host 0.0.0.0 --port 8100 \
  -tp 4 --generation-config vllm --no-enable-prefix-caching --no-disable-hybrid-kv-cache-manager \
  --async-scheduling \
  --kv-transfer-config '{"kv_connector":"NixlConnector","kv_role":"kv_both","kv_load_failure_policy":"fail"}'

# decode
VLLM_SSM_CONV_STATE_LAYOUT=DS \
VLLM_NIXL_SIDE_CHANNEL_HOST="$(hostname)" \
VLLM_NIXL_SIDE_CHANNEL_PORT=6551 \
vllm serve nvidia/Qwen3.5-397B-A17B-NVFP4 \
  --served-model-name Qwen3.5-397B-A17B-NVFP4 \
  --host 0.0.0.0 --port 8200 \
  -tp 4 --generation-config vllm --no-enable-prefix-caching --no-disable-hybrid-kv-cache-manager \
  --async-scheduling \
  --kv-transfer-config '{"kv_connector":"NixlConnector","kv_role":"kv_both","kv_load_failure_policy":"fail"}'

# router
python tests/v1/kv_connector/nixl_integration/toy_proxy_server.py \
  --host 0.0.0.0 --port 8000 \
  --prefiller-hosts "$PREFILL_HOST" --prefiller-ports 8100 \
  --decoder-hosts "$DECODE_HOST" --decoder-ports 8200

# gsm8k test
python tests/evals/gsm8k/gsm8k_eval.py

Results:

python tests/evals/gsm8k/gsm8k_eval.py
Running GSM8K evaluation: 1319 questions, 5-shot
Evaluating: 100%| 1319/1319
Results:
Accuracy: 0.886
Invalid responses: 0.010
Total latency: 36.813 s
Questions per second: 35.830
Total output tokens: 197127
Output tokens per second: 5354.831

Essential Elements of an Effective PR Description Checklist

The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
The test plan, such as providing test command.
The test results, such as pasting the results comparison before and after, or e2e results
(Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

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.

llx-08 · 2026-07-13T13:53:59Z

LGTM, thanks a lot for taking this over and doing the extra accuracy validation in PD mode — really appreciate it! Let's proceed with #48481
One small thing, would you mind adding a Co-authored-by trailer to keep the credit clear? Thank you very much

Signed-off-by: Artem Perevedentsev <aperevedents@nvidia.com> Co-authored-by: llx-08 <2596671364@qq.com>

arpera · 2026-07-13T14:08:47Z

One small thing, would you mind adding a Co-authored-by trailer to keep the credit clear?

Sure, added you as a co-author.

njhill · 2026-07-14T14:50:37Z

Thanks @llx-08 @arpera. I also encountered this bug and it's a critical one to fix.

I am a bit uneasy with the approach in this PR though for the following reasons:

It adds a nontrivial new method to the connector API which every kv-loading connector must implement independently, or else they will still be subject to the bug.
It implicitly requires any unused tail of the blocks to be zeroed by the connector (this happens to be the case for the current nixl implementation, but other connectors might work at token level and still have some exposure)
There's a missed failure path - if the transfer fails and kv_load_failure_policy=recompute, the excluded blocks will be neither zeroed or overwritten by the transfer. So this would require some additional special handling.

So it's overall quite fragile and we'd want to make sure we cover all connectors in addition to the NIXL one, including multiconnector with appropriate aggregation logic.

I have been experimenting with an alternative approach here njhill@dc75fd8, which has some of its own downsides, but does not put any burden on connectors apart from that they should perform their async loading logic (and nothing else) within start_load_kv, which is the reason it updates some existing impls.

This runs the zeroing on a separate stream and defers the connector's async loading until the zeroing is complete.

I'm actually still not sure what's best. We could repair / fill in the deficiencies in this PR for now, but it's making an interface change that I'm not sure we'll necessarily want to keep/require.

arpera · 2026-07-14T15:16:29Z

As for me, I need any kind of fix for async scheduling because otherwise I cannot publish reproducible results on PD Qwen3.5 - this is currently my priority. Async scheduling significantly improves the performance in my case, so I would prefer to have any kind of solution to this problem for now even if it does not support all the connectors.

Replace the KVConnectorMetadata.get_blocks_to_skip_kv_cache_zeroing hook (and its NIXL implementation) with a generic derivation: the external-token grant from get_num_new_matched_tokens already defines exactly which blocks an async load will overwrite, so the scheduler computes the skip set itself from the request's blocks covering [num_local_tokens, num_local_tokens + num_external_tokens). This works for every consumer connector with no connector involvement, and can't drift from the allocation arithmetic. Failure cases are handled without per-request state: everything still allocated beyond the valid prefix of an async load is by construction exactly the skipped-but-unwritten set, so when a load fails and the request will recompute locally, those blocks are re-recorded into the managers' normal new-block flow and the next step's new_block_ids_to_zero zeroes them before local recompute consumes them. Fully-failed or aborted loads release their blocks, which are re-recorded for zeroing on reallocation as usual. The only remaining state is the transient per-step skip set. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Nick Hill <nickhill123@gmail.com>

njhill · 2026-07-15T22:27:07Z

@arpera @llx-08 I have pushed an additional commit which addresses most of my concerns, PTAL.

NickLucche

Thanks for fixing this race @arpera @njhill 🐐

Makes the OffloadingConnector path functional for hybrid (attention + GDN) models on TPU, with every transfer carried by raiden direct DMA between the per-layer KV arrays and raiden's DmaMapped host pool. Canonicalization and routing: each per-layer array (attention KV and every mamba/GDN state) becomes an int8 (num_blocks, row_bytes) canonical tensor; block ids arrive ordered by KV-cache group (GPULoadStoreSpec.group_sizes) and route to only that group's arrays. The scheduler-side mamba external-KV patch extends to OffloadingConnector for hybrid models, and the CPU pool is sized from the bytes the TPU arrays physically hold per block row rather than vLLM's shared-page accounting. Raiden managers are created per (KV-cache group, row class): a manager sizes and addresses every registered layer from layer 0's slice, so tensors sharing one must share slice size and kernel-row factor. Attention arrays register as source-major views at kernel-block granularity with per-block copy_sizes runs; a registration-time probe routes any tensor raiden cannot row-address to a host-staged torch sidecar (empty for Qwen3.5 with conv-state padding enabled). Transfer dispatch is gated on per-step device completion: raiden DMA has no ordering against TensorCore programs, and any fence program referencing a KV array races the per-step buffer donation swap ('Buffer has been deleted or donated' whenever its Execute enqueue is delayed past the next step's donation). Jobs queue stamped with an engine step and dispatch - always on the runner main thread - once that step's sampled-token copy has resolved (AsyncTPUCopyState), which under serial TensorCore execution means every device program through that step has finished. The same gating supports async scheduling, but async additionally requires vLLM fixes past 0.22.1 (vllm-project/vllm#45357, vllm-project/vllm#48481 - scheduler races that corrupt KV consumers); run with --no-async-scheduling until then. KV_OFFLOAD_HANG_DUMP_S arms a watchdog that dumps all thread stacks when the engine stops stepping. Validated on Qwen3.5-35B-A3B-FP8 TP4 (tpu7x, forced-overflow 512-block HBM pool): 2-pass gsm8k holds accuracy across restore-heavy reuse runs (baseline 0.77; reuse 0.80/0.76/0.757; 13.9M external hits), and a prefill-heavy benchmark (150 x ~4.8K-token prompts, 94.5% external hits) reaches 14,318 total tok/s vs 8,900 without offload at identical accuracy.

Reverts the scheduler half of vllm-project#48481, which skipped zeroing the blocks an async KV load would overwrite in order to avoid the zeroing racing the out-of-band RDMA write. The preceding commit's KVConnectorLoadGate makes that race impossible by construction - a connector's loads never start until the zeroing event for the step's blocks has fired - so the exclusion is no longer load-bearing, and full zeroing of every freshly allocated block is restored. Removed: - Scheduler._skip_zero_block_ids and the filtering in _get_new_block_ids_to_zero (the helper itself is kept). - KVCacheManager.get_zeroing_block_ids_in_range, which only fed the skip set, and record_blocks_for_zeroing plus its call on the failed-load path, which existed solely to re-zero blocks the skip had left unwritten. With no skip, those blocks were zeroed at allocation. - SingleTypeKVCacheManager.records_new_block_ids, used only by the two removed manager methods. - The three unit tests covering the skip behavior. The NixlConnector's own receive-side zeroing of untransferred regions (post_process_device_kv_on_receive) is deliberately left in place; it is now partly redundant but is a separate revert. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Nick Hill <nickhill123@gmail.com>

… models (vllm-project#48481) Signed-off-by: Artem Perevedentsev <aperevedents@nvidia.com> Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: llx-08 <2596671364@qq.com> Co-authored-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

arpera requested review from ApostaC, NickLucche, WoosukKwon, alexm-redhat, heheda12345, ivanium, njhill, orozery, robertgshaw2-redhat, xuechendi and ywang96 as code owners July 13, 2026 09:15

claude Bot reviewed Jul 13, 2026

View reviewed changes

mergify Bot added qwen Related to Qwen models v1 kv-connector labels Jul 13, 2026

arpera mentioned this pull request Jul 13, 2026

[KV Connector] Use prefill block content to overwrite decode bad blocks #47373

Closed

arpera and others added 2 commits July 13, 2026 17:06

[KV Connector] Fix PD async scheduling for Qwen3.5 model

73a4728

Signed-off-by: Artem Perevedentsev <aperevedents@nvidia.com> Co-authored-by: llx-08 <2596671364@qq.com>

fix linter issue

bb62bf0

Signed-off-by: Artem Perevedentsev <aperevedents@nvidia.com> Co-authored-by: llx-08 <2596671364@qq.com>

arpera force-pushed the pd-async-sched-qwen branch from 6f23f3b to bb62bf0 Compare July 13, 2026 14:07

njhill changed the title ~~[KV Connector] Fix PD async scheduling for Qwen3.5 model~~ [KV Connector] Fix PD async scheduling race condition for hybrid attn models Jul 15, 2026

njhill added the ready ONLY add when PR is ready to merge/full CI is needed label Jul 15, 2026

Merge branch 'main' into pd-async-sched-qwen

c53f071

NickLucche approved these changes Jul 16, 2026

View reviewed changes

njhill merged commit 530852f into vllm-project:main Jul 16, 2026
83 checks passed

yinlin09 mentioned this pull request Jul 18, 2026

[KV Offload] All-raiden group-aware offload for hybrid models (Qwen3.5) vllm-project/vllm-torchtpu#84

Closed

NickLucche mentioned this pull request Jul 27, 2026

[Bugfix] Zero new KV blocks for quantized + sliding-window hybrid caches #47574

Merged

3 of 4 tasks

shugandhi mentioned this pull request Aug 4, 2026

fix: disable async scheduling on NVFP4 disagg decode + use the GA runtime image chw001/dynamo#2

Merged

This was referenced Aug 6, 2026

[Bug]: Qwen3.5-27B Disagg accuracy gsm8k collapses with async scheduling when TP==1 #42182

Closed

[Bug]: PD disaggregation for SSM models requires --no-async-scheduling when TP>1 #37285

Open

coderabbitai Bot mentioned this pull request Aug 7, 2026

fix(vllm): align KVBM decode KV cache layout ai-dynamo/dynamo#12821

Merged

stu-cao mentioned this pull request Aug 19, 2026

[Tracking Issue]: NIXL P/D Disaggregation for Hybrid Models #40017

Open

11 of 19 tasks

cocotyty mentioned this pull request Sep 20, 2026

kvshrink: upgrade base vLLM to v0.29.0 intel/intel-accel-for-llm#15

Merged

Navigation Menu

Repository navigation

[KV Connector] Fix PD async scheduling race condition for hybrid attn models - #48481

[KV Connector] Fix PD async scheduling race condition for hybrid attn models#48481
njhill merged 4 commits into
vllm-project:mainfrom
arpera:pd-async-sched-qwen

arpera commented Jul 13, 2026 •

edited

Loading

Uh oh!

claude Bot left a comment

Uh oh!

llx-08 commented Jul 13, 2026

Uh oh!

arpera commented Jul 13, 2026

Uh oh!

njhill commented Jul 14, 2026

Uh oh!

arpera commented Jul 14, 2026

Uh oh!

njhill commented Jul 15, 2026 •

edited

Loading

Uh oh!

NickLucche left a comment •

edited

Loading

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Conversation

arpera commented Jul 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Purpose

Details about this fix

Test Result

Qwen3.5 PD with TP==1

Qwen3.5 PD with TP>1

Uh oh!

claude Bot left a comment

Choose a reason for hiding this comment

Claude Code Review

Uh oh!

llx-08 commented Jul 13, 2026

Uh oh!

arpera commented Jul 13, 2026

Uh oh!

njhill commented Jul 14, 2026

Uh oh!

arpera commented Jul 14, 2026

Uh oh!

njhill commented Jul 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

NickLucche left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

arpera commented Jul 13, 2026 •

edited

Loading

njhill commented Jul 15, 2026 •

edited

Loading

NickLucche left a comment •

edited

Loading