[benenzhu]

Overview

#59164 removes synchronous MoRIIO READ destinations from new_block_ids_to_zero, but only for models with Mamba layers. needs_kv_cache_zeroing is also true for mixed-precision attention caches (#47574). DeepSeek-V4 is one: its MLA cache is uint8 (fp8_ds_mla) and its compressor state cache is float32. Its READ destinations are still zeroed on the compute stream, and the zeroing can run after the RDMA READ has landed. This PR applies the same exclusion to every READ group of models without Mamba layers.

Claims

• Closes the race between zeroing and synchronous READ for attention-only models with mixed-precision KV caches (DeepSeek-V4) in MoRIIO READ mode.
• No change for Mamba hybrids: they still exclude only the aligned attention pages, as in [Bugfix][MoRIIO] Exclude synchronous READ destinations from KV zeroing #59164.
• No change for uniform-precision models: the scheduler sends no zeroing list, so the branch does not run.
• Local recompute tails, pages that are not READ destinations, and later allocations of the same pages are still zeroed.

Validation

Check	main	this PR
test_update_state_drops_decode_recompute_tail_block (the non-Mamba cases now expect the exclusion) + new test_attention_only_read_skips_zeroing_in_every_group	3 failed	pass
test_moriio_connector.py + test_moriio_kv_layout.py + test_moriio_hma_scheduler.py (ROCm, MI350X, mori installed)	183 passed, 3 failed	186 passed

• The unit tests ran in a ROCm 10 vLLM nightly container, with the branch's Python tree on PYTHONPATH. The main column runs main's code against this PR's test file.
• Draft: not validated end to end yet. Our DeepSeek-V4-Pro 1P1D MoRIIO READ setup (2x MI350X) closes the same race with a host torch.cuda.current_stream().synchronize() in start_load_kv, before the READs are posted. We did not A/B the race itself. The nodes are not available right now; I will add an E2E accuracy run (with and without this change) before marking the PR ready.

Details

Flow (decode side, one step, --async-scheduling).

1. The scheduler allocates the new request's blocks. DeepSeek-V4's cache is mixed precision, so they all go to new_block_ids_to_zero.
2. build_connector_meta keeps the list unchanged, because _has_mamba is False.
3. The V2 runner's update_requests queues kv_block_zeroer.zero_block_ids on the current stream. The stream is still running the previous step's forward, so the zeroing kernel waits behind it.
4. pre_forward -> start_load_kv posts the RDMA READs from the host. The NIC writes the pages directly and does not wait on the stream.
5. If the READs finish before the previous forward, the queued zeroing kernel then overwrites the KV it just received.

Fix. build_connector_meta subtracts the READ destinations of every group when the model has no Mamba layers. The Mamba path keeps block_ids[:1], the aligned attention pages. This is the exception that KVConnectorBase_V1.build_connector_meta already documents for synchronous READ connectors.

Why skipping the zeroing is safe here.

• MoRIIO READ copies block_len bytes per block (compute_block_transfer_offsets), so each destination page becomes a byte copy of the prefiller's page. The prefiller zeroed that page when it allocated it.
• The worker awaits every READ before the forward (_await_reads_issued_this_step; [Bugfix][KV Connector][ROCm] MoRIIO READ: await every read before the forward #60808 makes this unconditional), so no kernel reads a destination page before the READ has landed.
• _reqs_need_recv holds only the aligned destination blocks. The decode-side recompute tail is not a READ destination, so it stays in the zeroing list (covered by the existing test).

Alternative considered. A stream sync before posting READs (what our internal build does) also closes the race. It makes the host wait for the previous forward on every step that receives KV, so the READs cannot overlap it. The exclusion needs no synchronization.

Related PRs.

• [KV Connector][ROCm] MoRIIO: hybrid READ for several attention groups and shared-group speculative methods #60867 (MoRIIO hybrid READ for several attention groups) changes the same set comprehension for the Mamba path to block_ids[: len(self._attn_group_ids)] and keeps the _has_mamba condition. The two compose. Whichever lands second needs a one-line rebase.
• [ROCm/MoRIIO] Route READ-mode remote KV load via async lifecycle (fix scheduler req_id_to_index KeyError) #54301 and the async READ RFC [RFC]: Enable full CUDA graphs for MoRIIO READ mode via async KV-load gating #49643 avoid the race differently: async loads already pass skip_zeroing_group_ids ([Bugfix][Core] Exempt exactly the blocks an async KV load writes from zeroing #59504).

Duplicate check. I searched open PRs for moriio zeroing, moriio zero, new_block_ids_to_zero, READ destinations zeroing and moriio mixed precision, and checked the diffs of #60867, #60808, #61726 and #54301. Only #60867 touches this exclusion, and only for Mamba hybrids.

---

Pull Request Checklist

• I used vLLM's /pr-checklist skill. (Mandatory for agents, optional for humans).

• AI assistance was used during the creation of this PR.

• Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.

• Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.

• Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.

• Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.

AI assistance (Claude Code) was used to prepare this change.

[benenzhu]

/ci run

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

[github-actions]

✅ Triggered Buildkite CI #95808 for commit 5e5cf4f6d6af.

[whx-sjtu]

Thanks for the fix! overall lgtm.

[whx-sjtu]

Nit: the docstring says these models "zero pages in every READ group", but the test asserts the opposite (READ pages are dropped from the zeroing list).

Suggested change

	"""Mixed-precision attention-only hybrids zero pages in every READ group."""
	"""Attention-only READ skips zeroing its destinations in every group."""

[whx-sjtu]

Nit: this filter now runs on every READ-mode step for mixed-precision models, including the prefill side where _reqs_need_recv is always empty, and copies new_block_ids_to_zero for nothing. Checking _reqs_need_recv first skips that.

Suggested change

	if self.mode == MoRIIOMode.READ and scheduler_output.new_block_ids_to_zero:
	if (
	self.mode == MoRIIOMode.READ
	and self._reqs_need_recv
	and scheduler_output.new_block_ids_to_zero
	):

[simondanielsson]

Question: Is the condition now overly loose? IIUC zeroing happens either if it's mamba or if using mixed precision kv. Can we check that instead so it aligns full with needs_kv_cache_zeroing?

Ideally we would use needs_kv_cache_zeroing directly so we don't have to mirror the check as that will cause drifts in the future, not sure if that's possible though

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @benenzhu.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[benenzhu]

Close cause no PD machine now.