[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)

CMakeLists.txt

@@ -869,9 +869,9 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
   endif()
 
   if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
-    cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0f" "${CUDA_ARCHS}")
+    cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0f;11.0f" "${CUDA_ARCHS}")
   else()
-    cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0a;10.3a" "${CUDA_ARCHS}")
+    cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0a;10.1a;10.3a" "${CUDA_ARCHS}")
   endif()
   if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND SCALED_MM_ARCHS)
     set(CUTLASS_MOE_SM100_SRCS "csrc/libtorch_stable/quantization/w8a8/cutlass/moe/grouped_mm_c3x_sm100.cu")

csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_entry.cu

@@ -180,7 +180,7 @@ bool cutlass_group_gemm_supported(int64_t cuda_device_capability) {
 
 #if defined CUDA_VERSION
   #if defined ENABLE_CUTLASS_MOE_SM100 && ENABLE_CUTLASS_MOE_SM100
-  if (cuda_device_capability >= 100 && cuda_device_capability < 110) {
+  if (cuda_device_capability >= 100 && cuda_device_capability < 120) {
     return CUDA_VERSION >= 12080;
   }
   #endif

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/core/test_contiguous_kv_packing.py

@@ -7,9 +7,8 @@ from unittest.mock import MagicMock
 import pytest
 import torch
 
-from vllm import envs
 from vllm.v1.core.kv_cache_utils import (
-    _get_kv_cache_config_deepseek_v4,
+    _get_kv_cache_config_packed,
     get_kv_cache_config_from_groups,
 )
 from vllm.v1.kv_cache_interface import (
@@ -84,15 +83,19 @@ def _make_groups(n_c4, n_c128, n_swa):
     return [mla_group, swa_group]
 
 
-def _mock_vllm_config():
+def _mock_vllm_config(kv_connector_extra_config: dict[str, str] | None = None):
     config = MagicMock()
     config.cache_config.num_gpu_blocks_override = None
+    config.kv_transfer_config = None
+    if kv_connector_extra_config is not None:
+        config.kv_transfer_config = MagicMock()
+        config.kv_transfer_config.kv_connector_extra_config = kv_connector_extra_config
     return config
 
 
 def _run(n_c4=3, n_c128=2, n_swa=5, mem=100 * 1024 * 1024):
     groups = _make_groups(n_c4, n_c128, n_swa)
-    return _get_kv_cache_config_deepseek_v4(_mock_vllm_config(), groups, mem)
+    return _get_kv_cache_config_packed(_mock_vllm_config(), groups, mem)
 
 
 def _page_sizes_by_layer(
@@ -135,7 +138,7 @@ class TestInterleavedPacking:
     def test_strided_views_are_independent(self):
         groups = _make_groups(n_c4=3, n_c128=2, n_swa=5)
         page_sizes = _page_sizes_by_layer(groups)
-        num_blocks, tensors = _get_kv_cache_config_deepseek_v4(
+        num_blocks, tensors = _get_kv_cache_config_packed(
             _mock_vllm_config(), groups, 100 * 1024 * 1024
         )
         backing = torch.zeros(tensors[0].size, dtype=torch.uint8)
@@ -156,8 +159,7 @@ class TestInterleavedPacking:
         for i, v in enumerate(views):
             assert (v == i + 1).all(), f"View {i} was corrupted"
 
-    def test_hma_attention_groups_keep_default_backing(self, monkeypatch):
-        monkeypatch.setattr(envs, "VLLM_USE_PACKED_HMA_KV_CACHE", False, raising=False)
+    def test_hma_attention_groups_keep_default_backing(self):
         full = _make_full_spec()
         sw = _make_sw_spec()
         page_size = full.page_size_bytes
@@ -178,8 +180,7 @@ class TestInterleavedPacking:
             KVCacheTensor(size=page_size * 32, shared_by=["full.1", "sw.2", "sw.3"]),
         ]
 
-    def test_hma_attention_groups_use_packed_backing_with_flag(self, monkeypatch):
-        monkeypatch.setattr(envs, "VLLM_USE_PACKED_HMA_KV_CACHE", True, raising=False)
+    def test_hma_attention_groups_use_packed_backing_with_enable_cross_layers(self):
         full = _make_full_spec()
         sw = _make_sw_spec()
         page_size = full.page_size_bytes
@@ -190,7 +191,9 @@ class TestInterleavedPacking:
         ]
 
         config = get_kv_cache_config_from_groups(
-            _mock_vllm_config(), groups, available_memory=page_size * 2 * 32
+            _mock_vllm_config({"enable_cross_layers_blocks": "True"}),
+            groups,
+            available_memory=page_size * 2 * 32,
         )
 
         assert config.num_blocks == 32

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(
                             [],
@@ -124,13 +138,14 @@ class OffloadingConnectorWorker:
                     raise NotImplementedError
 
         packed_kv_cache_tensor = next(
-            (t for t in kv_cache_config.kv_cache_tensors if t.block_stride), None
+            (
+                t
+                for t in kv_cache_config.kv_cache_tensors
+                if t.block_stride and t.shared_by
+            ),
+            None,
         )
-        is_dsv4 = all(
-            isinstance(group.kv_cache_spec, UniformTypeKVCacheSpecs)
-            for group in kv_cache_config.kv_cache_groups
-        )
-        if packed_kv_cache_tensor is not None and not is_dsv4:
+        if packed_kv_cache_tensor is not None:
             (tensor,) = tensors_per_block[packed_kv_cache_tensor.shared_by[0]]
             block_stride = tensor.stride(0)
             packed_tensor = tensor.as_strided(
@@ -153,7 +168,7 @@ class OffloadingConnectorWorker:
         block_data_refs: dict[str, list[CanonicalKVCacheRef]] = defaultdict(list)
         for kv_cache_tensor in kv_cache_config.kv_cache_tensors:
             # Filter to layers that were actually processed above.
-            # _get_kv_cache_config_deepseek_v4 emits KVCacheTensor entries for
+            # Packed KV allocation emits KVCacheTensor entries for
             # every (tuple_idx, page_size) slot; slots where no group has a
             # layer at that index produce an empty shared_by (reserved memory
             # with no corresponding model layer).

vllm/envs.py

@@ -209,7 +209,6 @@ if TYPE_CHECKING:
     VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS: int = 300
     VLLM_WORKER_SHUTDOWN_TIMEOUT_SECONDS: int = 5
     VLLM_KV_CACHE_LAYOUT: Literal["NHD", "HND"] | None = None
-    VLLM_USE_PACKED_HMA_KV_CACHE: bool = False
     VLLM_SSM_CONV_STATE_LAYOUT: Literal["SD", "DS"] | None = None
     VLLM_COMPUTE_NANS_IN_LOGITS: bool = False
     VLLM_ROCM_QUICK_REDUCE_QUANTIZATION: Literal[
@@ -1608,11 +1607,6 @@ environment_variables: dict[str, Callable[[], Any]] = {
     "VLLM_KV_CACHE_LAYOUT": env_with_choices(
         "VLLM_KV_CACHE_LAYOUT", None, ["NHD", "HND"]
     ),
-    # Opt into packed per-block KV cache allocation for multi-group
-    # attention-only HMA models (e.g. gpt-oss, Gemma 3/4).
-    "VLLM_USE_PACKED_HMA_KV_CACHE": lambda: bool(
-        int(os.getenv("VLLM_USE_PACKED_HMA_KV_CACHE", "0"))
-    ),
     # SSM conv state layout used for Mamba models.
     # - SD: (state_len, dim) — dim contiguous (default)
     # - DS: (dim, state_len) — TP-sharded dim on dim1,

vllm/model_executor/layers/fused_moe/config.py

@@ -861,6 +861,7 @@ def nvfp4_w4a16_moe_quant_config(
     g2_alphas: torch.Tensor,
     w1_scale: torch.Tensor,
     w2_scale: torch.Tensor,
+    gemm1_clamp_limit: float | None = None,
 ) -> FusedMoEQuantConfig:
     """
     Construct a quant config for 16-but activations and nvp4 weights.
@@ -872,6 +873,7 @@ def nvfp4_w4a16_moe_quant_config(
         g1_alphas=g1_alphas,
         g2_alphas=g2_alphas,
         weight_dtype="nvfp4",
+        gemm1_clamp_limit=gemm1_clamp_limit,
     )
 
 

vllm/model_executor/layers/fused_moe/oracle/nvfp4.py

@@ -425,6 +425,7 @@ def make_nvfp4_moe_quant_config(
             g2_alphas=w2_scale_2,
             w1_scale=w13_scale,
             w2_scale=w2_scale,
+            gemm1_clamp_limit=swiglu_limit,
         )
     elif backend == NvFp4MoeBackend.EMULATION:
         return nvfp4_moe_quant_config(

vllm/platforms/rocm.py

@@ -119,6 +119,14 @@ def _sync_hip_cuda_env_vars():
     hip_val = os.environ.get("HIP_VISIBLE_DEVICES") or None
     cuda_val = os.environ.get("CUDA_VISIBLE_DEVICES") or None
 
+    if cuda_val is not None:
+        logger.warning_once(
+            "Using CUDA_VISIBLE_DEVICES on ROCm is deprecated and support "
+            "will be removed in vLLM v0.26.0. Please use HIP_VISIBLE_DEVICES "
+            "instead.",
+            scope="process",
+        )
+
     if hip_val is not None and cuda_val is not None:
         if hip_val != cuda_val:
             raise ValueError(

vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm120.py

@@ -81,11 +81,13 @@ class FlashInferMLASparseSM120Impl(SparseMLAAttentionImpl[FlashInferMLASparseMet
             )
         self.kv_scale_format = _kv_scale_format_for_model(model_type)
 
-        assert indexer is not None, (
-            "FLASHINFER_MLA_SPARSE_SM120 requires a sparse-MLA indexer "
-            "(model with index_topk in its config)."
+        # Skip-topk layers are built with indexer=None and get the shared
+        # buffer via mla_args instead (cf. FLASHMLA_SPARSE).
+        self.topk_indices_buffer: torch.Tensor | None = (
+            indexer.topk_indices_buffer
+            if indexer is not None
+            else mla_args.get("topk_indices_buffer")
         )
-        self.topk_indices_buffer: torch.Tensor | None = indexer.topk_indices_buffer
         from vllm.utils.flashinfer import has_flashinfer_sparse_mla_sm120
 
         if not has_flashinfer_sparse_mla_sm120():

vllm/v1/core/kv_cache_utils.py

@@ -947,7 +947,9 @@ def may_override_num_blocks(vllm_config: VllmConfig, num_blocks: int) -> int:
     return num_blocks
 
 
-def _pool_bytes_per_block(kv_cache_groups: list[KVCacheGroupSpec]) -> int:
+def _pool_bytes_per_block(
+    vllm_config: VllmConfig, kv_cache_groups: list[KVCacheGroupSpec]
+) -> int:
     """
     Bytes consumed by one block in the worker's shared KV cache pool, mirroring
     the divisor used by `get_kv_cache_config_from_groups` to convert
@@ -958,7 +960,7 @@ def _pool_bytes_per_block(kv_cache_groups: list[KVCacheGroupSpec]) -> int:
         kv_cache_groups[0].kv_cache_spec, UniformTypeKVCacheSpecs
     ):
         return kv_cache_groups[0].kv_cache_spec.page_size_bytes
-    if _use_packed_kv_cache_groups(kv_cache_groups):
+    if _use_packed_kv_cache_config(vllm_config, kv_cache_groups):
         # buckets = {page_size: [[layer_names], [layer_names], ...]}
         buckets = _bucket_layers_by_page_size(kv_cache_groups)
         return sum(ps * len(slots) for ps, slots in buckets.items())
@@ -1250,16 +1252,26 @@ def _bucket_layers_by_page_size(
     return buckets
 
 
-def _use_packed_kv_cache_groups(
+def _use_packed_kv_cache_config(
+    vllm_config: VllmConfig,
     kv_cache_groups: list[KVCacheGroupSpec],
 ) -> bool:
     is_dsv4 = all(
         isinstance(group.kv_cache_spec, UniformTypeKVCacheSpecs)
         for group in kv_cache_groups
     )
-    return is_dsv4 or (
-        bool(envs.VLLM_USE_PACKED_HMA_KV_CACHE) and len(kv_cache_groups) > 1
+    kv_transfer_config = vllm_config.kv_transfer_config
+    extra_config = (
+        kv_transfer_config.kv_connector_extra_config
+        if kv_transfer_config is not None
+        else {}
     )
+    # NOTE: enable_cross_layers_blocks is an experimental API and subject to change with
+    # https://github.com/vllm-project/vllm/issues/42082
+    enable_cross_layers = (
+        str(extra_config.get("enable_cross_layers_blocks", "False")).lower() == "true"
+    )
+    return is_dsv4 or (enable_cross_layers and len(kv_cache_groups) > 1)
 
 
 def _get_kv_cache_config_packed(
@@ -1347,10 +1359,9 @@ def get_kv_cache_config_from_groups(
             )
             for layer_name in kv_cache_groups[0].layer_names
         ]
-    elif _use_packed_kv_cache_groups(kv_cache_groups):
-        # DeepSeek V4 keeps the existing packed layout. Other multi-group
-        # attention-only HMA layouts can opt in with
-        # VLLM_USE_PACKED_HMA_KV_CACHE=1.
+    elif _use_packed_kv_cache_config(vllm_config, kv_cache_groups):
+        # DeepSeek V4 uses the packed layout by default. Other multi-group
+        # layouts can opt in with --enable-cross-layers.
         num_blocks, kv_cache_tensors = _get_kv_cache_config_packed(
             vllm_config, kv_cache_groups, available_memory
         )
@@ -2069,7 +2080,7 @@ def get_kv_cache_configs(
             if not groups:
                 adjusted_memory.append(avail_mem)
                 continue
-            bytes_per_block = _pool_bytes_per_block(groups)
+            bytes_per_block = _pool_bytes_per_block(vllm_config, groups)
             logger.info(
                 "Overriding num_gpu_blocks=%d with num_gpu_blocks_override=%d",
                 avail_mem // bytes_per_block,

vllm/v1/engine/core.py

@@ -809,8 +809,10 @@ class EngineCore:
         if tags is None or tags:
             self.model_executor.wake_up(tags)
 
-        # Resume scheduling (applies to all levels)
-        self.resume_scheduler()
+        # Partial wakes intentionally keep the remaining allocations asleep.
+        # Resume scheduling only once all executor memory is resident again.
+        if not self.model_executor.is_sleeping:
+            self.resume_scheduler()
 
     def is_sleeping(self) -> bool:
         """Check if engine is sleeping at any level."""
@@ -1173,10 +1175,10 @@ class EngineCoreProc(EngineCore):
                 numa_utils.log_current_affinity_state(process_title)
 
             if data_parallel and vllm_config.kv_transfer_config is not None:
-                # modify the engine_id and append the local_dp_rank to it to ensure
+                # modify the engine_id and append the dp_rank to it to ensure
                 # that the kv_transfer_config is unique for each DP rank.
                 vllm_config.kv_transfer_config.engine_id = (
-                    f"{vllm_config.kv_transfer_config.engine_id}_dp{local_dp_rank}"
+                    f"{vllm_config.kv_transfer_config.engine_id}_dp{dp_rank}"
                 )
                 logger.debug(
                     "Setting kv_transfer_config.engine_id to %s",
@@ -1947,10 +1949,11 @@ class DPEngineCoreProc(EngineCoreProc):
                     # All engines are idle.
                     continue
 
-                # We are in a running state and so must execute a dummy pass
-                # if the model didn't execute any ready requests.
-                with self.log_iteration_details(None):
-                    self.execute_dummy_batch()
+                # Execute a dummy pass when no ready requests ran, unless the
+                # engine is sleeping.
+                elif not self.model_executor.is_sleeping:
+                    with self.log_iteration_details(None):
+                        self.execute_dummy_batch()
 
             # 3) All-reduce operation to determine global unfinished reqs.
             self.engines_running = self._has_global_unfinished_reqs(