Repository navigation
PERF: avoid repeated Xavier snapshot size scans - #5625
Conversation
There was a problem hiding this comment.
Code Review
This pull request introduces an experimental GPU-first Xavier cache feature that enables high-performance GPU-to-GPU KV cache transfers using CUDA IPC and xoscar NIXL, falling back to CPU snapshots when necessary. The changes span documentation, frontend configuration, and core backend components—including the implementation of a GPUTransfer manager, tiered snapshot store optimizations, and asynchronous load integration in the V1 connector. The review feedback is highly constructive, recommending minor robustness enhancements such as logging swallowed exceptions in background transfer tasks and explicitly specifying dtype=torch.long when constructing PyTorch tensors from integer lists to prevent potential type mismatches.
07e2ef9 to
c98d91a
Compare
c98d91a to
f6ecad3
Compare
rogercloud
left a comment
There was a problem hiding this comment.
Minor
xinference/model/llm/vllm/xavier/test/test_tiered_snapshot.py:337— size-accounting tests only fail the first layer copy, so the kept-block path (later layer fails,_block_sizes[key]must stay at the prior bytes) is untested; likewise no test asserts_block_sizesafter a failedGPUTransfer.stagecleanup. Add one later-layer failure assertion and one_block_sizescheck in an existing failed-stage test.
Blocking: no — recommended event: APPROVE
Staging each new KV layer recomputes a block's size by visiting every layer already stored in that block. This makes size validation quadratic in the number of layers and delays cold prefills while snapshots are staged.
Current review range after #5624 merged
Rebased only this PR onto main
b9dc5ce456ed21421f1f2107b05b03dec8b09337(the #5624 merge). Current head:8bd2b1bb07595f16f8b7d690e3ce607b0905f0d9. All preceding dependencies are merged, so this is the next independently reviewable increment.Current review diff. The normal Files changed tab now contains only incremental snapshot-size accounting and the small review/lifecycle fixes described below.
The accounting commit rebased without conflicts and is unchanged by range-diff. Merged async completion, lease cleanup, dtype, slab, packed-demotion, CPU-CI coverage and GPU-LRU fixes remain intact. Shutdown now clears the new
_block_sizesmap, with successful/failed-fence assertions. Background transfer task failures are logged with their traceback even when the awaiting caller is cancelled; a gated async regression verifies cancellation still drains the write before surfacing cancellation, and the background error remains visible. Both integer-index comments were already addressed in the merged predecessors.Current validation:
test_block_trackerexpectation mismatch).gpu_remote: 108 passed, including physical demotion, size accounting, failure cleanup and cancellation logging.git diff --checkpassed.Historical performance measurements below remain attributed to their original heads; they were not rerun as new throughput claims for this rebase. Only #5625 was pushed; #5627 and #5628 remain untouched to limit CI usage.
Latest review: partial-copy failure accounting
Added test-only assertions for a later-layer copy failure: after K is stored, a failing V copy leaves only K and preserves exactly its physical byte count, equal to the sum of retained tensors. The existing GPUTransfer.stage failure test now also checks that unpublished block size metadata is removed while the published BF16 block retains its four-byte count; this runs for both OOM and runtime failures.
Local Xavier regression: 214 passed, 10 skipped, 1 known-baseline deselection. Modified-file pre-commit and diff checks passed. Runtime code is unchanged in this follow-up; the 108-pass CUDA run above belongs to the preceding revision and was not rerun for these CPU failure-injection assertions.
Previous implementation and validation record
Track physical tensor bytes once per newly stored layer instead. Clear the count when a block is dropped and initialize it again when an evicted key is reused. Immutable hits do not increment the count; demotion preserves it. GPU-first placement, CPU overflow, leases, and full-block budget validation retain their existing behavior.
Historical measurement baseline: #5624 revision
bf1cc5a25; measured candidate:07e2ef923. The current rebased review range is linked above.Validation:
test_block_trackerexpectation mismatch.07e2ef923; no server tracebacks or profiling instrumentation in the performance run.Performance: two RTX 3090 Ti GPUs with NVLink, Qwen2.5-0.5B BF16, eager vLLM 0.21.0, torch 2.11.0, xoscar 0.11.1, NIXL 1.1.0.
Incoming TTFT median decreases 6.6%, p95 decreases 2.1%, and workload throughput increases 2.2%. The throughput gap to native NIXL changes from 18.8% to 17.0%. This is a small cold-workload improvement; the large remaining TTFT gap is not resolved.
The initial warm-cache comparison was mixed:
An independent warm-only C32 ABBA control then ran three consecutive 600-request trials per deployment (3,600 measured requests per variant):
The initial 4.6% C32 throughput drop did not reproduce; the independent control differs by +0.6%. Both sets are reported because process/run variation is material. No steady-state throughput improvement is claimed.
The isolated CUDA staging microbenchmark recorded baseline 0.970–1.196 s versus candidate 0.671–0.681 s (24 layers, ten batches with a shared prefix and 2,295 demotions). This diagnostic is separate from the end-to-end measurements.
Main method: ABBA (baseline, candidate, candidate, baseline), then native NIXL and an 8 MiB candidate. Every fresh deployment starts with the same 12 initial prompts. ABBA deployments run 600 warm requests each at C16 and C32; native and mixed deployments run 120 at each concurrency. Each deployment then runs three identical interleaved rounds: eight ongoing streams followed by 32 unique approximately 4,400-token prompts staggered by 100 ms. Incoming output is capped at 64 tokens and ongoing output at 1,024; EOS can shorten responses. Profiling is disabled, prefix caching enabled, model length 8,192, and GPU-first snapshot budget 256 MiB. Client chunk gaps are not GPU per-token timings.
Both baseline and candidate retain 1,365 GPU and 26,452 CPU snapshots after 26,452 demotions. The 8 MiB candidate exercises 105 GPU batches and 882 CPU batches, including cancellation recovery. GPU services were shut down after testing. The main run plus the independent warm control completed 13,920 successful requests and four intentional cancellations.
Evidence is saved under
artifacts/pd-size-accounting-20261003/locally and/data/xuyeqin/pd-demotion-batch-final-20261003plus/data/xuyeqin/pd-size-warm-check-20261003on the GPU host.