Skip to content

perf: reduce Xavier GPU transfer padding with persistent small slabs - #5621

Merged
qinxuye merged 3 commits into
xorbitsai:mainfrom
qinxuye:perf/xavier-sized-gpu-slabs
Oct 4, 2026
Merged

qinxuye merged 3 commits into
xorbitsai:mainfrom
qinxuye:perf/xavier-sized-gpu-slabs

Conversation

@qinxuye

@qinxuye qinxuye commented Oct 3, 2026 •

Copy link
Copy Markdown
Contributor

Small warm-cache transfers currently copy the entire persistent GPU slab (12 MiB in the measured configuration). Keep a persistent 256 KiB view alongside the full slab and use it after two consecutive small GPU batches from the same source rank. A large batch immediately selects the full slab and resets that streak. Source and destination lengths match; both views share the original allocation and transfer locks. Repeated same-size transfers reuse buffer identities, although switching sizes can still update xoscar registrations.

Current review range after #5619 and #5620 merged

Rebased onto main bb03fb9e58d33df4ba328ef99f61a7615979ac7b. Current head: c72606288427e2a0c9a14c218a805fa8a8e803d6.

Both #5619 and #5620 are merged. This PR is now the first independently reviewable increment in the performance stack. The normal Files changed tab contains the persistent-small-slab change, GPU-mapping synchronization and explicit index dtypes from review, and their tests.

Current review diff. git range-diff confirmed the original increment was unchanged by the rebase; follow-up aaea73281 addresses review comments.

Exact-head local Xavier regression: 175 passed, 7 skipped, 1 deselected. The known baseline test_block_tracker integration remains deselected. Pre-commit and git diff --check passed. Historical GPU/performance measurements below remain attributed to their original heads; they were not rerun for this rebase.

Only this PR was rebased/pushed in this update to limit CI usage. Later stack PRs remain unchanged and can be updated in dependency order after this one merges. The four existing review threads have been addressed: inherited logical-dtype initialization was verified, mapping registration is serialized with concurrent/failure-retry regressions, and both GPU indexing sites explicitly use torch.long. Thread resolution does not imply review approval.

Previous implementation and validation record

Depends on #5620, which follows #5616. Incremental commit: 1abec59b1. The peer-reference and writeback-index experiments are not included. No new dependency or configuration option.

Validation

  • Local GPU/tiered/request-transfer tests: 50 passed, 2 CUDA skips; pre-commit passed.
  • Real GPU-transfer tests: 24 passed, including all 65,536 BF16 bit patterns.
  • Two RTX 3090 Ti GPUs with NVLink; xoscar 0.11.1/NIXL 1.1; vLLM 0.21; Qwen2.5-0.5B BF16 eager; 256 MiB snapshot budget; concurrency 16.
  • Baseline is perf: avoid gathering already published Xavier GPU snapshots #5620 at 9e563b00a. Fresh models in baseline/optimized/optimized/baseline order, each with 12 cold and 600 warm requests, profiling disabled. An additional cancellation/recovery workload followed the second optimized measurement and is excluded from the comparison.
Metric Baseline Optimized
Warm throughput, 1,200 requests/version 1,479.7 token/s 1,512.8 token/s (+2.2%)
Warm TTFT p50 164.6 ms 149.6 ms
Warm TTFT p95 272.5 ms 254.1 ms
Cold long-prompt mean TTFT 174.3 ms 179.6 ms
Copy bytes, matched 612-request deployments 7,902,068,736 472,645,632 (-94.0%)

Matched deployments both moved 349,765,632 useful bytes in 628 batches. Both long comparisons improved throughput (1476.9 -> 1512.1 and 1482.4 -> 1513.5). A shorter comparison was noisy; cold long-prompt latency did not improve. These are modest warm-path gains on this setup, not claims about other models or cross-host performance.

All 4,668 completed requests across short/long comparisons, mixed-tier validation and cancellation recovery succeeded, plus four intentional stream cancellations. Cold response texts matched baseline. Mixed validation exercised 160 GPU batches and 468 CPU batches. After cancellation, 600 requests succeeded; transfer-actor FD counts stayed at 121/124 in the latter half. GPUs returned to 554 MiB / 0% utilization after cleanup.

Mixed-size registration review follow-up

Confirmed the churn and fixed it in c726062. On one persistent channel between the two RTX 3090 Ti GPUs (NVLink; xoscar 0.11.1/NIXL 1.1.0), alternating 256 KiB/12 MiB copies was slower than always copying 12 MiB. We now require two consecutive small GPU batches per source rank before selecting the small view; a large batch immediately resets the streak. Alternating full/small tails therefore remain on the full slab, while sustained small batches reuse the small view.

Paired forward/reverse-order microbenchmarks, 20 warmup copies plus 1,000 measured copies per case, two trials per case:

  • Always full: 1.4745 ms/copy, zero registration changes after warmup.
  • Alternating, old selection: 1.8486 ms/copy, 1,000 registration changes on each endpoint per trial.
  • Alternating, real updated GPUTransfer selector: 1.4613 ms/copy, zero registration changes on either endpoint.
  • Sustained small, old/new selection: 0.4894 / 0.4834 ms/copy, zero registration changes after warmup.

The target is cleared before each trial and copied contents verified. This isolates copy_to channel behavior, including registration/protocol overhead; it is not a new end-to-end model-serving throughput claim. Added regressions for alternating batches, per-peer independent streaks, and stable small-view reuse. Full focused Xavier suite: 175 passed, 7 skipped, 1 deselected; pre-commit passed.

The real constructor is also covered below/at/above the 256 KiB threshold, including exact buffer/ref keys and shared-storage prefix-view identity. The existing end-to-end measurements above are historical measurements from before this selection-policy follow-up.

@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 implements an experimental GPU-first Xavier cache feature for vLLM models, allowing local CUDA IPC caches and persistent xoscar NIXL transfer buffers to optimize KV cache transfers. It includes frontend option exposure, supervisor/worker launch integrations, and comprehensive test coverage. Feedback on these changes highlights a few critical issues: the new 'logical_dtypes' attribute in 'TieredKVSnapshotStore' must be initialized in 'init' to prevent an AttributeError; the asynchronous '_ensure_gpu_cache_mapping' method in 'v1_connector.py' requires synchronization (such as an 'asyncio.Lock') to avoid concurrent registration crashes; and index tensors in 'gpu_transfer.py' should explicitly specify 'dtype=torch.long' to prevent runtime errors when index lists are empty.

Comment thread xinference/model/llm/vllm/xavier/tiered_snapshot.py
Comment thread xinference/model/llm/vllm/xavier/v1_connector.py
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-sized-gpu-slabs branch 2 times, most recently from ba5b8a0 to 20c0318 Compare October 3, 2026 19:53
@qinxuye
qinxuye force-pushed the perf/xavier-sized-gpu-slabs branch from 20c0318 to 544c6d9 Compare October 3, 2026 19:55
@qinxuye
qinxuye requested a review from rogercloud October 3, 2026 19:58

@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/gpu_transfer.py:260 — Mixed-size workloads may churn NIXL registrations: xoscar reuses a registration only for the identical buffer object, so alternating <=256 KiB / >256 KiB batches on one channel re-register GPU memory on both ends and re-exchange agent metadata at every switch. The benchmark only covers a uniform warm workload; please benchmark an alternating small/full batch case (or confirm churn is negligible). If it is costly, consider keeping the slab choice sticky per channel.
  • xinference/model/llm/vllm/xavier/gpu_transfer.py:73 — The real GPUTransfer.__init__ small-slab setup is not tested (the fixture builds via __new__ with hand-filled buffers), so a wrong constant or a missing/extra small view would not fail any test. Add a constructor test with xo.buffer_ref stubbed and the CUDA/NIXL checks patched.

Blocking: no — recommended event: APPROVE

Comment thread xinference/model/llm/vllm/xavier/gpu_transfer.py Outdated
Comment thread xinference/model/llm/vllm/xavier/gpu_transfer.py
@qinxuye
qinxuye merged commit a88ac52 into xorbitsai:main Oct 4, 2026
15 checks passed
@qinxuye
qinxuye deleted the perf/xavier-sized-gpu-slabs branch October 4, 2026 00:51
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