perf: batch Xavier KV loads across scheduled requests

[qinxuye]

Xavier currently loads each request in a scheduler batch through a separate blocking actor call. Send the batch in one call, group reads by source rank and storage tier, and retain bounded transfer sizes. Shared KV hashes can write to multiple destination slots; identical destinations are deduplicated and conflicting writes are rejected before transfer. Small requests keep the small-slab bound from #5621 so batching does not reintroduce large padded copies.

Current review range after #5621 merged

Rebased onto main a88ac52eae9c377ad4a62c2c770fe52cd051c5cb. Current head: 50d0d97a10a99ad979b8fc2f87bf8e74bd1d0657. All earlier dependencies are merged; this is now the next independently reviewable increment. The normal Files changed tab contains the batching change and the shutdown cleanup review fix.

Current review diff.

Retained the merged schema validation, mapping lock, explicit index dtypes and per-peer slab-selection hysteresis from #5619–#5621. Batched small reads remain bounded while the first small batch uses the full slab, switching only after consecutive small batches; the batch test now verifies the inherited full/small/small sequence.

Review follow-up releases snapshot storage, metadata and full slabs on close. The closed runtime remains as an idempotent shutdown/remapping guard; regressions cover successful and failing shutdown, one-time cleanup and rejection of remapping. All six existing review threads have been addressed.

Exact-head local Xavier regression: 187 passed, 7 skipped, 1 deselected (the known baseline block-tracker integration remains deselected). Pre-commit and git diff --check passed. Historical GPU and performance measurements below remain attributed to their original heads and were not rerun for this rebase. Only #5622 was pushed; later PRs remain untouched to limit CI usage.

Previous implementation and validation record

The connector waits for the entire load and lease cleanup even if cancelled repeatedly. Every lease release is attempted after the writes finish, including on load failures. The unused single-request actor entry point was removed in the review follow-up; GPUTransfer.load remains for tests. New counters report logical requests and batch calls.

Depends on #5621. Incremental commit: 1d6612c76.

Validation

• Pre-commit passed. Local Xavier tests: 138 passed, 7 skipped, excluding test_xavier.py. Its existing block-tracker empty-set assertion fails identically on an isolated, unmodified perf: reduce Xavier GPU transfer padding with persistent small slabs #5621 checkout; neither tracker implementation nor that test is changed here.
• GPU host: 59 GPU/tiered/request-transfer tests passed, including BF16 bit preservation, multiple source ranks, shared KV with distinct destinations, transfer bounds, repeated cancellation, and partial cleanup failure.
• Real serving: 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 GPU snapshot budget. Baseline is perf: reduce Xavier GPU transfer padding with persistent small slabs #5621 at 1abec59b1.
• Fresh deployments in baseline/optimized/optimized/baseline order. Each runs 12 initial requests and then concurrency 1/8/16/32 with 48/120/600/120 requests respectively. Profiling disabled. Four stream cancellations and 600 recovery requests follow the second optimized measurement; those are excluded from performance comparisons. Native NIXL and an 8 MiB mixed-tier configuration each receive one additional reference/validation deployment.

Concurrency	Baseline token/s	Batched token/s	Baseline TTFT p95	Batched TTFT p95
1	159.4	159.4	67.9 ms	68.3 ms
8	928.7	952.2	169.1 ms	160.8 ms
16	1532.3	1552.4	261.4 ms	220.5 ms
32	1986.3	2100.5	517.8 ms	441.9 ms

Both C32 comparisons improved throughput (2048.0 -> 2205.7 and 1928.2 -> 2004.8 token/s), though these are short 120-request trials. C16 throughput is close and its pairwise direction is inconsistent; tail latency improves in both pairs. Long-prompt mean TTFT in the initial pass is 176.7 -> 200.4 ms, including one 498.2 ms optimized outlier, so no cold-path improvement is claimed.

For a matched 900-request deployment, 478 batch loads handle 900 logical requests, GPU transfers fall from 916 to 842, and useful bytes remain 406,388,736. Mixed validation exercises 229 GPU batches and 436 CPU batches with 491 loads for 900 requests. Native NIXL remains faster in the same setup: 1882.0 / 2774.7 token/s at C16/C32 (single reference run).

All 6,000 completed requests succeeded, plus four intentional cancellations. All 72 initial response texts match baseline. Warm concurrent response text varies with batching in the baseline and native NIXL too; bitwise cache correctness is covered separately by the regression tests. After cancellation, 600 requests succeeded, and transfer-actor FD counts stayed at 121/124 in the latter half. Candidate GPU and connector source hashes match the measured files.

Additional review follow-up

50d0d97a1 filters shared null-block positions using vLLM KVCacheBlocks is_null metadata before constructing destination maps; genuine conflicting writes still fail. Added two-request hybrid-group coverage, a mixed large/small batch with both slab sizes that asserts the full batch bound, and identical key/destination mappings from two ranks that assert one source read. Removed the unused load_gpu_request_v1 actor method. All nine review threads have been addressed; exact-head local suite: 187 passed, 7 skipped, 1 deselected; pre-commit passed. No additional performance claims are made for this review fix.

[gemini-code-assist]

Code Review

This pull request implements an experimental GPU-first Xavier cache feature for vLLM models, enabling high-performance KV cache transfers via CUDA IPC and xoscar NIXL. It includes changes to the supervisor, worker, and model actors to support GPU cache budget validation, subpool spawning with external addresses, and resource cleanup. The review feedback focuses on enhancing robustness and preventing race conditions. Specifically, it recommends introducing an asyncio.Lock to serialize GPU cache mapping, explicitly specifying dtype=torch.int64 for PyTorch index tensors to avoid type inference issues, clearing staged blocks during cleanup to prevent memory leaks, and wrapping resource release in a try...finally block to ensure references are cleared.

[rogercloud]

Minor

• xinference/model/llm/vllm/xavier/gpu_transfer.py:272 — the new batch-wide conflict check may fail a whole step's load if vLLM reports the shared null block for out-of-window positions (risk to confirm).
• xinference/model/llm/vllm/xavier/gpu_transfer.py:322 — no test covers a small slab together with a mixed large+small batch.
• xinference/model/llm/vllm/xavier/gpu_transfer.py:273 — no test covers the same (key, destination) arriving from two source ranks.

Simplification

L377: delete: load_gpu_request_v1 has no callers after the connector moved to load_gpu_requests_v1. Remove it (keep GPUTransfer.load, still used by tests).
net: -3 lines possible

Blocking: no — recommended event: APPROVE