[v1] Initialize InputBatch in initialize_kv_cache instead of __init__

[wenyili]

Summary

• Moves InputBatch creation from GPUModelRunner.__init__ to initialize_kv_cache (via the new initialize_input_batch method), so it is built with the final block sizes from kv_cache_config rather than a placeholder value.
• Removes the now-unnecessary _init_block_sizes / _init_kernel_block_sizes tracking fields and the conditional re-init logic in the old may_reinitialize_input_batch.
• Fixes a latent bug where cp_kv_cache_interleave_size was omitted from the re-init path.

Background

The early initialization in __init__ was a workaround for the UVA pinned-memory reuse bug reported in #18298. The three-step failure chain was:

1. CPU offloading (#15354) stored cpu_data only as a Python attribute on the parameter object, backed by a no-op C++ deleter that did not hold a reference to the CPU tensor:

   p._vllm_offloaded_cpu_data = cpu_data   # sole Python reference to pinned memory
   p.data = get_cuda_view_from_cpu_tensor(cpu_data)

2. process_weights_after_loading (GPTQ, Marlin, FP8, …) replaced parameter objects; the old PackedvLLMParameter was GC'd, releasing cpu_data back to CachingHostAllocator.
3. InputBatch, if created after load_model, allocated pinned buffers (e.g. block_table_cpu) that reused that freed memory, aliasing live GPTQ weight CUDA views.

The workaround was to create InputBatch before load_model with a placeholder block size (#17945 / #18298), then conditionally re-create it in initialize_kv_cache if block sizes differed (#18593 / #18654). The root cause was described at the time as "unknown reasons".

I confirmed the root cause by temporarily keeping cpu_data alive on the module instead of the parameter:

# vllm/model_executor/models/utils.py  (temporary patch to confirm root cause)
if not hasattr(module, '_vllm_offloaded_cpu_datas'):
    module._vllm_offloaded_cpu_datas = []
module._vllm_offloaded_cpu_datas.append(cpu_data)  # survives parameter replacement
p.data = get_cuda_view_from_cpu_tensor(cpu_data)

This eliminated the corruption in test_cpu_offload_gptq, confirming that premature release of cpu_data was the cause.

Why it is now safe

Two fixes closed the root cause:

1. C++ lambda in csrc/cuda_view.cu (landed after [Don't merge] Debug failing quantization test with input batch move #18298): the deleter now captures cpu_tensor by value, keeping it alive for the full lifetime of the UVA CUDA view:

   // before: no-op, does not hold cpu_tensor alive
   auto deleter = [](void*) {};
   // after: captures cpu_tensor, increments its C++ refcount
   [base = cpu_tensor](void*) {}

   The pinned memory cannot be returned to CachingHostAllocator as long as any UVA CUDA view is alive — regardless of Python-side GC of the parameter object.
2. device_loading_context re-offload in vllm/model_executor/model_loader/utils.py: detects parameters whose UVA offload flag was lost during process_weights_after_loading and re-creates the UVA view with the C++ lambda.

PR #36461 confirmed the fix is complete by removing the offload + quantization reinit guard added in #18654.

This is not duplicating an existing PR

• [v1] Support multiple KV cache groups in GPU model runner #17945 moved init to initialize_kv_cache but was reverted (Revert "[v1] Support multiple KV cache groups in GPU model runner (#17945) #18459) due to the then-unknown root cause.
• [v1] Redo "Support multiple KV cache groups in GPU model runner (#17945)" #18593 re-did the multi-KV-cache-group work while keeping the placeholder init in __init__.
• [Bugfix] Fix cpu-offload-gb assertion with non-default block sizes #36461 removed the offload guard but did not complete the cleanup.
• No open PR removes the placeholder init entirely.

Test commands

# Linters (all passed locally)
pre-commit run --files vllm/v1/worker/gpu_model_runner.py

# Unit tests
.venv/bin/python -m pytest tests/v1/worker/ -v
.venv/bin/python -m pytest tests/quantization/test_cpu_offload.py::test_cpu_offload_gptq -v

Results on RTX 4090: tests/v1/worker/ 100 passed, 1 skipped; test_cpu_offload_gptq passed (directly exercises the root cause scenario).

AI assistance was used (Claude) to identify the root cause and draft this change. Every changed line has been reviewed by the human submitter.

🤖 Generated with Claude Code