vllm - 💡(How to fix) Fix [Feature]: Polymorphic buffer management for V1 worker (CPU/GPU staged tensors, lower hot-path overhead) [1 participants]

🚀 The feature, motivation and pitch

Summary

This proposes a polymorphic buffer management path for the V1 worker so internal staged tensors can be handled consistently across CPU and GPU, while reducing hot-path overhead from transient buffer work.

Prototype branch: masterFoad/vllm:pr-polymorphic-buffers-clean
Branch URL: https://github.com/masterFoad/vllm/tree/pr-polymorphic-buffers-clean
Head commit: 31e0aa38c

Motivation

The current buffer/staging flow can incur avoidable overhead in decode and input-prep paths due to temporary allocations and host-side staging patterns.

A unified buffer abstraction should improve maintainability and reduce per-step overhead in performance-critical paths.

Proposed Change

• Introduce DeviceMemoryManager to centralize staged tensor allocation strategy.
• Use polymorphic staged-write tensors via .device_tensor instead of GPU-only assumptions.
• Keep CPU-compatible and CUDA-compatible staged write paths under one abstraction.
• Refactor V1 worker callsites to use the new staged buffer flow.

Scope in Prototype Branch

Representative files touched:

• vllm/v1/worker/gpu/buffer_utils.py
• vllm/v1/worker/gpu/model_runner.py
• vllm/v1/worker/gpu/states.py
• vllm/v1/worker/gpu/block_table.py
• vllm/v1/worker/gpu/model_states/default.py
• vllm/v1/worker/gpu/sample/bad_words.py
• vllm/v1/worker/gpu/sample/logit_bias.py
• vllm/v1/worker/gpu/sample/penalties.py
• vllm/v1/worker/gpu/sample/sampler.py
• vllm/v1/worker/gpu/sample/states.py
• vllm/v1/worker/gpu/spec_decode/rejection_sampler.py

Preliminary Local Result

Local setup:

• OS: WSL Ubuntu 24.04
• GPU: RTX 4060 Laptop GPU, 8GB
• PyTorch: 2.11.0+cu129
• vLLM: 0.7.3
• Model: Qwen2.5-1.5B-Instruct

Throughput result:

Delta: +9.4%

These are prototype/local measurements and are shared only to indicate potential.

Request for Feedback

If this direction is useful, I can split the work into a smaller reviewable sequence:

1. Memory abstraction groundwork.
2. Targeted hot-path integrations.
3. Focused validation and benchmark reporting.

Alternatives

No response

Additional context

No response

Before submitting a new issue...

• Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.