[KV-Offloading] Fix tensors_per_block stride

[varun-sundar-rabindranath]

Purpose

On main Offloading connector fails for Qwen/Qwen3.6-35B-A3B & nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-FP8 with the assert

(Worker_TP3 pid=3812595) ERROR 06-26 23:01:48 [multiproc_executor.py:1000]   File "/home/varun-sundar-rabindranath/code/vllm/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py", line 174, in register_kv_caches
(Worker_TP3 pid=3812595) ERROR 06-26 23:01:48 [multiproc_executor.py:1000]     len({tensors_per_block[n][0].stride() for n in tensor_layer_names}) == 1

Why:
PR #44577 introduced cross-layer layouts for DSV4. The PR introduced a change in the offloading/worker.py::register_kv_caches to set the 2D tensors_per_block stride directly from the layer_kv_cache. This is incorrect when the KV cache tensor is not layer packed.

Fix:
Special case based on if the kv cache tensor is packed.

Test Plan

vllm-serve command:

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

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

lm_eval command:

lm_eval \
  --model local-completions \
  --model_args "base_url=http://127.0.0.1:${VLLM_PORT}/v1/completions,model=${MODEL},tokenized_requests=False,num_concurrent=1000,trust_remote_code=True" \
  --tasks gsm8k \
  --seed 42 \
  --num_fewshot 50 \
  --gen_kwargs temperature=0.0

Test Result

PR
nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-FP8

|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|    50|exact_match|↑  |0.9181|±  |0.0076|
|     |       |strict-match    |    50|exact_match|↑  |0.9158|±  |0.0076|

nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-FP8 - without offloading

|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|    50|exact_match|↑  |0.9121|±  |0.0078|
|     |       |strict-match    |    50|exact_match|↑  |0.9113|±  |0.0078|

Qwen/Qwen3.6-35B-A3B

|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|    50|exact_match|↑  |0.8886|±  |0.0087|
|     |       |strict-match    |    50|exact_match|↑  |0.8870|±  |0.0087|

Qwen/Qwen3.6-35B-A3B - without offloading

|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|    50|exact_match|↑  |0.8923|±  |0.0085|
|     |       |strict-match    |    50|exact_match|↑  |0.8916|±  |0.0086|

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[varun-sundar-rabindranath]

cc @tlrmchlsmth @bnellnm

[tlrmchlsmth]

looks good to me, thank you for finding this.

I am following up with a integration tests that would have caught this.

[tlrmchlsmth]

adding integration tests here #46893