[Feature][KV-offloading]: Relax Offloading Block Size asserts

[varun-sundar-rabindranath]

🚀 The feature, motivation and pitch

Setting offloading block size for some models like,

• deepseek-ai/DeepSeek-V4-Flash, and
• google/gemma-4-31B-it

trips the assert

vllm/vllm/v1/kv_offload/base.py

Line 547 in 1ff9429

assert len(gpu_block_sizes) == 1, (

This is the "all kv groups must have the same block size" assert.

How to run,

Deepseek V4

MODEL="deepseek-ai/DeepSeek-V4-Flash"
DSV4_EXTRA="--kv-cache-dtype fp8 --block-size 256 --tokenizer-mode deepseek_v4 --tool-call-parser deepseek_v4 --enable-auto-tool-choice --reasoning-parser deepseek_v4"
TP_SIZE=4
CPU_BYTES=26843545600

KV_TRANSFER_CONFIG=$(cat <<EOF
{
  "kv_connector": "OffloadingConnector",
  "kv_role": "kv_both",
  "kv_connector_extra_config": {
    "block_size": 512,
    "spec_name": "CPUOffloadingSpec",
    "cpu_bytes_to_use": ${CPU_BYTES},
    "eviction_policy": "lru"
  }
}
EOF
)

vllm serve "${MODEL}" \
      --tensor-parallel-size="${TP_SIZE}" \
      --kv-transfer-config "${KV_TRANSFER_CONFIG}" \
      --enable-prefix-caching \
      --no-disable-hybrid-kv-cache-manager \
      ${DSV4_EXTRA} \
      --port 8000

Gemma4

MODEL="google/gemma-4-31B-it"
TP_SIZE=2
CPU_BYTES=26843545600

KV_TRANSFER_CONFIG=$(cat <<EOF
{
  "kv_connector": "OffloadingConnector",
  "kv_role": "kv_both",
  "kv_connector_extra_config": {
    "block_size": 512,
    "spec_name": "CPUOffloadingSpec",
    "cpu_bytes_to_use": ${CPU_BYTES},
    "eviction_policy": "lru"
  }
}
EOF
)

vllm serve "${MODEL}" \
      --tensor-parallel-size="${TP_SIZE}" \
      --kv-transfer-config "${KV_TRANSFER_CONFIG}" \
      --enable-prefix-caching \
      --no-disable-hybrid-kv-cache-manager \
      --port 8000

Request:
This feature request is to relax the assertion allowing models with disparate KV group block sizes to still accommodate/align on a common offloading block size.

Motivation:
This is important for disk offloading as we very quickly hit the IOPs bottleneck with large quantities of small file read/writes. Relaxing the assert lets the user tune the offloading block size per their disk capacity.

Alternatives

The users can still set the top-level GPU block size and benefit that way. But this is not ideal and is not supported by all models. For example, setting the gpu block size to anything other than 256 for deepseek-v4-flash hits the assert,

vllm/vllm/v1/worker/gpu/attn_utils.py

Line 221 in 1ff9429

.view(permuted_kv_cache_shape)

. No issues with gemma4 when setting this.

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.

[Debasish-87]

I spent some time tracing the KV offloading pipeline to better understand this feature request. From what I can see, block_size_factor is currently treated as a single global value across the offloading stack (including base.py, scheduler.py, and cpu/gpu_worker.py). This makes the assertion in base.py seem like a safeguard for a broader architectural assumption rather than a standalone validation.

Before I start working on a patch, I'd like to clarify the intended direction. Should supporting heterogeneous GPU KV block sizes involve introducing per-KV-group block size factors throughout the offloading pipeline, or is the expectation to preserve the current single-factor abstraction while aligning on a common offloaded block size?

[dannyharnik]

The issue is that today offloading chunk size (the block_size parameter) is defined as number of tokens, which is then translated internally to block_size_factor, the parameter actually used in the code.

The problem is with models that don't have a single value of number of tokens per GPU block. In such models you can't define the block_size_factor using a fixed tokens per chunk parameters.
What we need is very simple - expose an alternative API for defining the block_size_factor directly.

Note that blocks_size_factor is undergoing a name change to blocks_per_chunk #48150 . So new API should be for blocks_per_chunk.