Repository navigation
perf: reduce Xavier GPU transfer padding with persistent small slabs - #5621
Conversation
There was a problem hiding this comment.
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.
ba5b8a0 to
20c0318
Compare
20c0318 to
544c6d9
Compare
rogercloud
left a comment
There was a problem hiding this comment.
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 realGPUTransfer.__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 withxo.buffer_refstubbed and the CUDA/NIXL checks patched.
Blocking: no — recommended event: APPROVE
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-diffconfirmed the original increment was unchanged by the rebase; follow-upaaea73281addresses review comments.Exact-head local Xavier regression: 175 passed, 7 skipped, 1 deselected. The known baseline
test_block_trackerintegration remains deselected. Pre-commit andgit diff --checkpassed. 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
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.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:
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.