[Bugfix][AMD] A GPU-CPU KV transfer fault in OffloadingConnector takes the engine down instead of degrading the cache. - #52838
Open
okorzh-amd wants to merge 5 commits into
Conversation
okorzh-amd
requested review from
ApostaC,
NickLucche,
ivanium,
orozery and
xuechendi
as code owners
August 18, 2026 22:22
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
mergify
Bot
added
rocm
![depthfirst-app[bot]](https://avatars.githubusercontent.com/in/1021002?s=60&v=4){kind=small}
Member
|
/amd-ci run nightly |
|
✅ Triggered Buildkite AMD CI #12216 for commit |
1 task
Collaborator
|
Thanks @okorzh-amd ! Can you elaborate in what cases will the h2d transfers may fail? |
Etelis
reviewed
Sep 6, 2026
Comment on lines
+689
to
+690
| self._buffer_pool.append((batch_src, batch_dst, batch_sizes)) | ||
| self._submit_failures.append(job_id) |
Contributor
There was a problem hiding this comment.
Can we drain any copies already queued before reporting this job as finished? hipMemcpyBatchAsync can fail partway through a batch, so the scheduler could reuse the KV blocks while earlier copies are still running.
Contributor
Author
…e engine A GPU<->CPU transfer fault in OffloadingConnector takes the engine down. `swap_blocks_batch` raises synchronously inside `transfer_async` via STD_TORCH_CHECK on a non-success return, and a device error latched by an earlier transfer surfaces the same way at `end_event.query()`. Every call site is unguarded: the V1 mixin starts the load one line above its `try`, and the V2 path (`gpu/kv_connector.py` pre_forward/post_forward) has no `try` at all. `OffloadingConnectorWorker.get_finished` then asserts `transfer_result.success`, which is dead code because the only producer hardcodes `success=True`. This is the wrong failure mode for this connector specifically: `requires_kv_delivery` is False with the comment "a dropped save is just a future cache miss", so it declares best-effort delivery and then aborts the engine when a transfer fails. Report faults instead of raising: - `transfer_async` catches the fault, queues the job as a submit failure and still returns True, so the job completes. Jobs must always complete or the scheduler's per-job pending count never reaches zero and the request is never resumed. The stream and events are deliberately dropped rather than pooled, so a latched device error cannot leak into an unrelated transfer. - `get_finished` drains submit failures and guards `end_event.query()` / `elapsed_time`, emitting `TransferResult(success=False)`. - `wait()` no longer propagates a synchronize failure; the next poll reports it, keeping failure handling in one place. - A failed store is dropped with a warning: the chunk is simply not cached. - A failed load is not droppable. Its destination GPU blocks hold undefined data while the scheduler has already advanced num_computed_tokens past them, so the blocks are reported through the newly implemented `get_block_ids_with_load_errors()` and the scheduler recomputes them. The request is still reported via `get_finished()`, per the base contract. Block-level recovery assumes a single KV cache group (`_update_requests_with_invalid_blocks` has a standing TODO for the hybrid memory allocator), so on a hybrid model a failed load raises a clear error naming that limitation rather than silently serving undefined KV. Store faults, which dominate offload traffic, are handled on every model. Test plan --------- `pytest tests/v1/kv_connector/unit/offloading_connector/ tests/v1/kv_offload/` before: 76 failed, 655 passed, 23 skipped, 32 errors after: 76 failed, 661 passed, 23 skipped, 32 errors The +6 are the new tests; the pre-existing failures and errors are environmental (no GPU visible to the test container) and identical in both arms. Mutation check: with the source change reverted and the tests kept, 6 of the 7 new/affected tests fail, so they are not vacuous. pre-commit (ruff-check, ruff-format, mypy-3.10, mypy-3.12, all other hooks) passes on the changed files. Not a duplicate: vllm-project#50984 does this for the Mooncake connector, not this one. vllm-project#42461 would replace the per-block error API with a per-request one, but its RFC (vllm-project#35780) was auto-closed as stale on 2026-08-18 and the per-request API is not on main; if it ever lands, this connector migrates alongside nixl, mooncake, flexkv and lmcache. AI-assistance disclosure: developed with Claude (Anthropic). Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: okorzh-amd <okorzh-amd@users.noreply.github.com>
…rors Narrow the three new handlers from `except Exception` to `except RuntimeError`. torch.AcceleratorError and torch.cuda.CudaError both derive from RuntimeError, as does the STD_TORCH_CHECK in swap_blocks_batch, so every genuine device fault is still caught; an AttributeError, TypeError or ValueError now propagates instead of being silently downgraded to a degraded cache. Found by running the suite against an environment where torch.ops._C_cache_ops.swap_blocks_batch was unavailable. The broad catch turned that AttributeError into a "failed transfer", so the copy silently never happened and test_canonical_layout.py reported "reader tp=4 rank=0 bytes diverge from the tp=2 writers' ground truth" -- a missing op presenting as data corruption. With the narrower catch the AttributeError surfaces directly. Adds test_get_finished_propagates_non_device_errors to pin the boundary. Test: pytest v1/kv_offload v1/kv_connector/unit/offloading_connector origin/main 783 passed, 3 skipped this branch 790 passed, 3 skipped Run on 8xMI355X (gfx950, ROCm 7.2) with the repo's compiled extensions on sys.path; the +7 are this PR's new tests. Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: okorzh-amd <okorzh-amd@users.noreply.github.com>
Review catch on this PR: a failed store was reported to the scheduler as an ordinary completion. mark_completed() carried no failure information, so _process_worker_metadata called manager.complete_store() with the default success=True, which sets ref_cnt = 0 and marks the blocks evictable (cpu/manager.py:247-255) -- publishing destination blocks that were never written. A later lookup on those keys would load whatever the recycled blocks last held, which may be another request's KV. That is a worse failure than the engine death this PR set out to fix. Propagate the verdict instead: - OffloadingWorkerMetadata gains failed_jobs, mark_completed() takes success=True, and aggregate() unions the sets across workers. - The worker reports its failed transfers with success=False. - The scheduler records the failure on TransferJobStatus and passes success=not failed to complete_store(), taking the existing cpu/manager.py:256-262 path that removes and frees the blocks. The flag is sticky on the job rather than read from the current batch: pending_count is decremented across steps, so the step that carries the failure is often not the step that reaches zero. A per-batch lookup would silently miss exactly the multi-worker case this guards. complete_load() is deliberately left alone. It has no success parameter and needs none -- a failed load does not corrupt the CPU source blocks, and the call must still run to drop their ref count. The GPU-side damage is reported separately via get_block_ids_with_load_errors(). Test: pytest v1/kv_offload v1/kv_connector/unit/offloading_connector 790 -> 794 passed, 3 skipped, 0 failed (8xMI355X, gfx950, ROCm 7.2) New: test_failed_store_is_not_published_as_cache reports the failure in an earlier step than the one completing the job, so it fails if the verdict is read per-batch; test_successful_store_is_published keeps the guard from passing vacuously. Mutation check: forcing success=True fails the first and leaves the second passing. Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: okorzh-amd <okorzh-amd@users.noreply.github.com>
okorzh-amd
force-pushed
the
okorzh/kv-offload-dma-failure-channel
branch
from
September 8, 2026 17:12
d6a8c30 to
ff68153
Compare
A batched copy can fail partway through, so the submit-failure path reported a job whose copies were still running. `swap_blocks_batch` submits the descriptor list in chunks and checks each chunk separately (ROCm caps a call at 8192 descriptors, so it always chunks), and `hipMemcpyBatchAsync` is itself a loop of async copies that records `failIdx` and breaks on the first failure. Copies submitted before the fault are already on the stream. `transfer_async` caught the exception, queued the job on `_submit_failures` and returned, and `get_finished()` reported `success=False` on the next poll with nothing synchronized. Three consequences, all on the same lines: - The scheduler is free to reuse the destination GPU blocks while an abandoned load is still writing them, and the CPU block a failed store targeted is freed back to the allocator while the D2H write may still land in it. - `_transfer_events[job_id]` was only set on the success path, so `wait()` silently no-opped for exactly these jobs. - The job never entered `_transfers`, so the next submission's `stream.wait_event(self._transfers[-1].end_event)` chained past the orphaned stream. The pinned descriptor buffers were also returned to the pool at the same point, while the driver may still be DMA-reading them for the copies it did enqueue. Record `end_event` on the failing stream instead, and queue the transfer with `failed=True`. `get_finished()` reports `success=False` only once `end_event.query()` is true, which restores the ordering guarantees for free and removes `_submit_failures` entirely -- that queue also reported failures ahead of still-running earlier jobs, breaking in-order completion. A failed transfer's stream, events and descriptor buffers are dropped rather than pooled, in both the submit and poll failure paths. If `end_event.record()` itself raises, the in-flight copies cannot be bounded and the error propagates. Test plan --------- 8xMI355X (gfx950), ROCm, nightly-e962733e08 image with the compiled extensions overlaid on the branch: `pytest tests/v1/kv_offload/ tests/v1/kv_connector/unit/offloading_connector/` before: 925 passed, 3 skipped after: 927 passed, 3 skipped (+3 new tests, -1 replaced.) Mutation check: with the source change reverted and the tests kept, all 3 new tests fail plus the tightened poll-failure assertion, so they are not vacuous. pre-commit (ruff-check, ruff-format, mypy-3.12, all other hooks) passes on the changed files. Addresses @Etelis's review comment on gpu_worker.py. AI-assistance disclosure: developed with Claude (Anthropic). Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Oxana Korzh <okorzh@amd.com>
Per AGENTS.md: minimize comments, prefer self-documenting code, assume the reader is familiar with vLLM. The rationale that mattered is in the commit messages and the PR description; the docstrings are unchanged. No functional change. `pytest tests/v1/kv_offload/ tests/v1/kv_connector/unit/offloading_connector/` gives 927 passed, 3 skipped, identical to the parent commit. ruff-check and ruff-format clean. Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Oxana Korzh <okorzh@amd.com>
Contributor
Author
@orozery On the API I did check before using it. RFC #35780, the proposal to remove per-block error #50984 is open right now doing for Mooncake P/D exactly what this PR does. Its consumer side is being actively fixed rather than retired — #54733 is open making the scheduler's For when h2d can fail, I don't have a clean repro here. #49276 is an open load-direction report, but its shapes are a segfault and a hang, which this cannot catch. I cite it as evidence the surface is real, not as something this PR fixes. |
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
4 participants
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.
Purpose
swap_blocks_batchraises synchronously insidetransfer_async—STD_TORCH_CHECK(result == hipSuccess, ...)atcsrc/libtorch_stable/cache_kernels.cu:168on the ROCm path, and the equivalent check at:151on the CUDA path. A device error latched by earlier work on the transfer stream surfaces the same way atend_event.query().Every call site is unguarded, so that exception reaches
execute_modeland kills the worker:kv_connector_model_runner_mixin.py) callsstart_load_kvat:95, one line above thetry:at:96.v1/worker/gpu/kv_connector.py) has notryat all inpre_forward(:61) orpost_forward(:77).gpu_model_runner.handle_preemptions, which submits stores and waits on them, is also unguarded.Meanwhile the failure channel that exists is dead code.
offloading/worker.py:359carries the comment "we currently do not support job failures" aboveassert transfer_result.success, and the only producer of that result hardcodessuccess=True(v1/kv_offload/cpu/gpu_worker.py:691).OffloadingConnectoralso does not overrideget_block_ids_with_load_errors, sobase.py's empty-set default applies and the scheduler's_handle_invalid_blocksrecovery is unreachable for this connector — even though the plumbing is already wired atv1/worker/gpu/kv_connector.py:89.This is the wrong failure mode for this connector in particular.
offloading_connector.py:56returnsrequires_kv_delivery = Falsewith the comment "a dropped save is just a future cache miss, so opt out of the producer-role default". It declares best-effort delivery and then aborts the engine when a transfer fails.What this changes
Faults are reported, not raised:
transfer_asynccatches the fault, queues the job as a submit failure, and still returnsTrue. Jobs must always complete: the scheduler's per-job pending count waits forworld_sizeacks, so a job that silently disappears strands the request forever. The stream and events are deliberately not returned to their pools — a latched device error would otherwise leak into an unrelated transfer. The pinned descriptor buffers are plain host memory and are still recycled.get_finisheddrains submit failures and guardsend_event.query()/elapsed_time, emittingTransferResult(success=False).wait()no longer propagates a synchronize failure; the next poll reports it, keeping failure handling in one place.requires_kv_delivery = Falsealready promises.num_computed_tokenspast them, so resuming the request would read garbage KV. The blocks are reported through a newly implementedget_block_ids_with_load_errors()and the scheduler recomputes them. The request is still reported viaget_finished(), as the base-class contract requires.Scope of the catch (second commit). The handlers catch
RuntimeError, notException.torch.AcceleratorErrorandtorch.cuda.CudaErrorboth derive fromRuntimeError, as does theSTD_TORCH_CHECKinswap_blocks_batch, so every genuine device fault is still caught — while anAttributeError,TypeErrororValueErrorkeeps propagating instead of being silently downgraded to a degraded cache.Test Plan
Test Result
8×MI355X (gfx950), ROCm 7.2, torch 2.12.0, with the compiled extensions on
sys.path:Essential Elements of an Effective PR Description Checklist
supported_models.mdandexamplesfor a new model.