[KV-Offloading] Fix tensors_per_block stride (#46888)

Signed-off-by: <>
Co-authored-by: Varun Sundar Rabindranath <varun-sundar-rabindranath@h100-01.nemg-001.lab.rdu2.dc.redhat.com>
(cherry picked from commit 798185d438)

tests/evals/gsm8k/configs/Qwen3-30B-A3B-NVFP4.yaml

@@ -2,4 +2,5 @@ model_name: "nvidia/Qwen3-30B-A3B-FP4"
 accuracy_threshold: 0.89
 num_questions: 1319
 num_fewshot: 5
+startup_max_wait_seconds: 1200
 server_args: "--enforce-eager --max-model-len 4096"

tests/evals/gsm8k/configs/Qwen3-Next-80B-A3B-NVFP4-EP2.yaml

@@ -2,6 +2,7 @@ model_name: "nm-testing/Qwen3-Next-80B-A3B-Instruct-NVFP4"
 accuracy_threshold: 0.75
 num_questions: 1319
 num_fewshot: 5
+startup_max_wait_seconds: 1200
 server_args: >-
   --enforce-eager
   --max-model-len 4096

tests/evals/gsm8k/configs/moe-refactor-dp-ep/Qwen3-30B-A3B-NvFp4-CT-fi-cutlass.yaml

@@ -2,4 +2,5 @@ model_name: "RedHatAI/Qwen3-30B-A3B-NVFP4"
 accuracy_threshold: 0.88
 num_questions: 1319
 num_fewshot: 5
+startup_max_wait_seconds: 1200
 server_args: "--enforce-eager --max-model-len 8192 --data-parallel-size 2 --enable-expert-parallel --moe-backend=flashinfer_cutlass"

tests/evals/gsm8k/configs/moe-refactor-dp-ep/Qwen3-30B-A3B-NvFp4-ModelOpt-fi-trtllm.yaml

@@ -2,4 +2,5 @@ model_name: "nvidia/Qwen3-30B-A3B-NVFP4"
 accuracy_threshold: 0.88
 num_questions: 1319
 num_fewshot: 5
+startup_max_wait_seconds: 1200
 server_args: "--enforce-eager --max-model-len 8192 --data-parallel-size 2 --enable-expert-parallel --moe-backend=flashinfer_trtllm"

tests/evals/gsm8k/configs/moe-refactor/Qwen3-30B-A3B-NvFp4-CT-fi-trtllm.yaml

@@ -2,4 +2,5 @@ model_name: "RedHatAI/Qwen3-30B-A3B-NVFP4"
 accuracy_threshold: 0.88
 num_questions: 1319
 num_fewshot: 5
+startup_max_wait_seconds: 1200
 server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2 --moe-backend=flashinfer_trtllm"

tests/evals/gsm8k/configs/moe-refactor/Qwen3-30B-A3B-NvFp4-ModelOpt-fi-trtllm.yaml

@@ -2,4 +2,5 @@ model_name: "nvidia/Qwen3-30B-A3B-NVFP4"
 accuracy_threshold: 0.88
 num_questions: 1319
 num_fewshot: 5
+startup_max_wait_seconds: 1200
 server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2 --moe-backend=flashinfer_trtllm"

tests/v1/kv_connector/unit/offloading_connector/utils.py

@@ -41,6 +41,7 @@ from vllm.v1.kv_cache_interface import (
     FullAttentionSpec,
     KVCacheConfig,
     KVCacheGroupSpec,
+    KVCacheTensor,
 )
 from vllm.v1.kv_offload.base import (
     GPULoadStoreSpec,
@@ -238,9 +239,18 @@ class RequestRunner:
                 )
             ]
 
+        kv_cache_tensors = [
+            KVCacheTensor(
+                size=group.kv_cache_spec.page_size_bytes * num_gpu_blocks,
+                shared_by=[layer_name],
+            )
+            for group in kv_cache_groups
+            for layer_name in group.layer_names
+        ]
+
         kv_cache_config = KVCacheConfig(
             num_blocks=num_gpu_blocks,
-            kv_cache_tensors=[],
+            kv_cache_tensors=kv_cache_tensors,
             kv_cache_groups=kv_cache_groups,
         )
         vllm_config.cache_config.num_gpu_blocks = num_gpu_blocks

vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py

@@ -53,6 +53,16 @@ class OffloadingConnectorWorker:
         kv_cache_config = self.spec.kv_cache_config
         num_blocks = kv_cache_config.num_blocks
 
+        # Packed layouts (e.g. DSv4) set block_stride > 0; their tensors use
+        # stride(0) as the manager-block stride (equals total_num_bytes_per_block).
+        # General (non-packed) layouts size the tensor at page_size_bytes per
+        # manager block, so page_size_bytes is the correct offloading stride.
+        layer_is_packed: dict[str, bool] = {
+            ln: bool(kv_tensor.block_stride)
+            for kv_tensor in kv_cache_config.kv_cache_tensors
+            for ln in kv_tensor.shared_by
+        }
+
         # layer_name -> (num_blocks, page_size_bytes) tensor
         tensors_per_block: dict[str, tuple[torch.Tensor, ...]] = {}
         # layer_name -> size of (un-padded) page in bytes
@@ -77,7 +87,11 @@ class OffloadingConnectorWorker:
                     page = layer_kv_cache_spec.page_size_bytes
                     elem_size = layer_kv_cache.element_size()
                     byte_offset = layer_kv_cache.storage_offset() * elem_size
-                    block_stride_bytes = layer_kv_cache.stride(0) * elem_size
+                    block_stride_bytes = (
+                        layer_kv_cache.stride(0) * elem_size
+                        if layer_is_packed[layer_name]
+                        else page
+                    )
                     tensors_per_block[layer_name] = (
                         torch.tensor(
                             [],