Skip to content

PERF: avoid repeated Xavier snapshot size scans - #5625

Merged
qinxuye merged 3 commits into
xorbitsai:mainfrom
qinxuye:perf/xavier-snapshot-size-accounting
Oct 4, 2026
Merged

qinxuye merged 3 commits into
xorbitsai:mainfrom
qinxuye:perf/xavier-snapshot-size-accounting

Conversation

@qinxuye

@qinxuye qinxuye commented Oct 3, 2026 •

Copy link
Copy Markdown
Contributor

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_sizes map, 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:

  • Local Xavier regression: 214 passed, 10 skipped, 1 deselected (the known baseline test_block_tracker expectation mismatch).
  • CUDA-enabled GPU-transfer, tiered-snapshot and request-transfer tests on gpu_remote: 108 passed, including physical demotion, size accounting, failure cleanup and cancellation logging.
  • Modified-file pre-commit hooks and git diff --check passed.

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:

  • Local Xavier tests: 153 passed, 9 skipped; CUDA-host Xavier tests: 160 passed, 2 skipped. Both deselect the known unrelated test_block_tracker expectation mismatch.
  • All modified-file pre-commit hooks pass.
  • Regression tests cover physical versus logical dtype sizing, immutable hits, failed copies, demotion, eviction, and key reuse.
  • The main benchmark completed 6,672 successful requests, including 600 after four intentional cancellations in mixed GPU/CPU mode. All 72 initial responses match exactly across the six deployments; concurrent text identity is not asserted because batching differs.
  • Candidate benchmark and CUDA-test runtime hashes match 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.

Interleaved workload metric #5624 baseline This change Native vLLM NIXL
Incoming long-prompt TTFT p50 2537.5 ms 2371.1 ms 356.8 ms
Incoming long-prompt TTFT p95 4135.4 ms 4050.2 ms 389.9 ms
Aggregate output tokens/s 958.7 979.7 1181.0
Ongoing-stream client chunk gap p99 15.75 ms 15.83 ms 14.91 ms

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:

Concurrency Tokens/s baseline → candidate Change TTFT p95 baseline → candidate TPOT p95 baseline → candidate
16 1591.4 → 1603.8 +0.8% 253.3 → 268.0 ms 8.01 → 7.93 ms
32 2379.3 → 2270.7 -4.6% 519.0 → 692.1 ms 9.30 → 9.93 ms

An independent warm-only C32 ABBA control then ran three consecutive 600-request trials per deployment (3,600 measured requests per variant):

Metric Baseline Candidate
Output tokens/s 2248.1 2261.4
TTFT p50 360.4 ms 344.9 ms
TTFT p95 599.3 ms 542.4 ms
TPOT p95 9.78 ms 9.89 ms

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-20261003 plus /data/xuyeqin/pd-size-warm-check-20261003 on the GPU host.

@XprobeBot XprobeBot added enhancement New feature or request gpu labels Oct 3, 2026
@XprobeBot XprobeBot added this to the v3.x milestone Oct 3, 2026

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment thread xinference/model/llm/vllm/xavier/gpu_transfer.py Outdated
Comment thread xinference/model/llm/vllm/xavier/gpu_transfer.py Outdated
Comment thread xinference/model/llm/vllm/xavier/gpu_transfer.py Outdated
@qinxuye
qinxuye force-pushed the perf/xavier-snapshot-size-accounting branch from 07e2ef9 to c98d91a Compare October 3, 2026 14:51
@qinxuye
qinxuye force-pushed the perf/xavier-snapshot-size-accounting branch from c98d91a to f6ecad3 Compare October 4, 2026 09:08
@qinxuye
qinxuye requested a review from rogercloud October 4, 2026 09:11

@rogercloud rogercloud left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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_sizes after a failed GPUTransfer.stage cleanup. Add one later-layer failure assertion and one _block_sizes check in an existing failed-stage test.

Blocking: no — recommended event: APPROVE

Comment thread xinference/model/llm/vllm/xavier/test/test_tiered_snapshot.py
@qinxuye
qinxuye merged commit 1be7706 into xorbitsai:main Oct 4, 2026
15 checks passed
@qinxuye
qinxuye deleted the perf/xavier-snapshot-size-accounting branch October 4, 2026 10:41
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

enhancement New feature or request gpu

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants