megamoe

Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>

.buildkite/test-amd.yaml

@@ -1594,9 +1594,10 @@ steps:
 #----------------------------------------------------------  mi300 · kernels  ----------------------------------------------------------#
 
 - label: Kernels Attention Test %N # TBD
-  timeout_in_minutes: 180
+  timeout_in_minutes: 55
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
   agent_pool: mi300_1
+  optional: true
   parallelism: 2
   working_dir: "/vllm-workspace/tests"
   source_file_dependencies:
@@ -1627,10 +1628,11 @@ steps:
   - pytest -v -s kernels/core --ignore=kernels/core/test_minimax_reduce_rms.py  kernels/test_concat_mla_q.py kernels/test_top_k_per_row.py
 
 - label: Kernels MoE Test %N # TBD
-  timeout_in_minutes: 180
+  timeout_in_minutes: 50
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
   agent_pool: mi300_1
-  parallelism: 4
+  optional: true
+  parallelism: 5
   working_dir: "/vllm-workspace/tests"
   source_file_dependencies:
   - csrc/quantization/cutlass_w8a8/moe/
@@ -2120,9 +2122,10 @@ steps:
   - pytest -v -s v1/e2e/spec_decode -k "draft_model or no_sync or batch_inference"
 
 - label: Spec Decode Eagle # TBD
-  timeout_in_minutes: 180
+  timeout_in_minutes: 45
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
   agent_pool: mi300_1
+  optional: true
   working_dir: "/vllm-workspace/tests"
   source_file_dependencies:
   - vllm/v1/spec_decode/
@@ -3040,7 +3043,7 @@ steps:
 #----------------------------------------------------------  mi355 · kernels  ----------------------------------------------------------#
 
 - label: Kernels (B200-MI355) # TBD
-  timeout_in_minutes: 180
+  timeout_in_minutes: 15
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
   agent_pool: mi355_1
   working_dir: "/vllm-workspace/"
@@ -3064,11 +3067,10 @@ steps:
   - pytest -v -s tests/kernels/attention/test_attention_selector.py
 
 - label: Kernels Attention Test %N # TBD
-  timeout_in_minutes: 180
+  timeout_in_minutes: 60
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
   agent_pool: mi355_1
   parallelism: 2
-  optional: true
   working_dir: "/vllm-workspace/tests"
   source_file_dependencies:
   - csrc/attention/
@@ -3082,10 +3084,10 @@ steps:
   - pytest -v -s kernels/attention --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
 
 - label: Kernels MoE Test %N # TBD
-  timeout_in_minutes: 180
+  timeout_in_minutes: 50
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
   agent_pool: mi355_1
-  parallelism: 4
+  parallelism: 5
   working_dir: "/vllm-workspace/tests"
   source_file_dependencies:
   - csrc/quantization/cutlass_w8a8/moe/

.buildkite/test_areas/kernels.yaml

@@ -74,6 +74,20 @@ steps:
   commands:
     - pytest -v -s kernels/attention --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
   parallelism: 2
+  mirror:
+    amd:
+      device: mi325_1
+      timeout_in_minutes: 55
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - csrc/attention/
+      - vllm/v1/attention
+      - vllm/model_executor/layers/attention
+      - tests/kernels/attention
+      - vllm/_aiter_ops.py
+      - vllm/envs.py
+      - vllm/platforms/rocm.py
 
 - label: Kernels Attention DiffKV Test (H100)
   key: kernels-attention-diffkv-test-h100
@@ -128,6 +142,22 @@ steps:
     - pytest -v -s kernels/moe --ignore=kernels/moe/test_modular_oai_triton_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
     - pytest -v -s kernels/moe/test_modular_oai_triton_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
   parallelism: 5
+  mirror:
+    amd:
+      device: mi325_1
+      timeout_in_minutes: 50
+      source_file_dependencies:
+      - csrc/quantization/cutlass_w8a8/moe/
+      - csrc/moe/
+      - tests/kernels/moe
+      - vllm/model_executor/layers/fused_moe/
+      - vllm/distributed/device_communicators/
+      - vllm/envs.py
+      - vllm/config
+      - vllm/_aiter_ops.py
+      - vllm/platforms/rocm.py
+      depends_on:
+      - image-build-amd
 
 - label: Kernels Mamba Test
   key: kernels-mamba-test

.buildkite/test_areas/pytorch.yaml

@@ -107,6 +107,12 @@ steps:
     - tests/compile/passes
   commands:
   - pytest -s -v compile/passes --ignore compile/passes/distributed
+  mirror:
+    amd:
+      device: mi300_1
+      timeout_in_minutes: 180
+      depends_on:
+      - image-build-amd
 
 - label: PyTorch Fullgraph Smoke Test
   key: pytorch-fullgraph-smoke-test

.buildkite/test_areas/spec_decode.yaml

@@ -12,6 +12,20 @@ steps:
     - tests/v1/e2e/spec_decode/
   commands:
     - pytest -v -s v1/e2e/spec_decode -k "eagle_correctness"
+  mirror:
+    amd:
+      device: mi325_1
+      timeout_in_minutes: 45
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - vllm/v1/spec_decode/
+      - vllm/v1/worker/gpu/spec_decode/
+      - vllm/model_executor/model_loader/
+      - vllm/v1/sample/
+      - vllm/model_executor/layers/
+      - tests/v1/e2e/spec_decode/
+      - vllm/platforms/rocm.py
 
 - label: Spec Decode Eagle Nightly B200
   key: spec-decode-eagle-nightly-b200

tests/compile/passes/test_double_aiter_rms_quant_fusion.py

@@ -22,7 +22,7 @@ import torch
 
 import vllm.config
 from tests.compile.backend import TestBackend
-from vllm._aiter_ops import is_aiter_found_and_supported, rocm_aiter_ops
+from vllm._aiter_ops import rocm_aiter_ops
 from vllm.compilation.passes.utility.noop_elimination import NoOpEliminationPass
 from vllm.compilation.passes.utility.post_cleanup import PostCleanupPass
 from vllm.config import (
@@ -83,9 +83,8 @@ class _ViewDoubleQuantModel(torch.nn.Module):
     [_NoViewDoubleQuantModel, _ViewDoubleQuantModel],
     ids=["no_view", "with_view"],
 )
-@pytest.mark.skipif(
-    not is_aiter_found_and_supported(),
-    reason="Only test on ROCm with AITER installed and supported",
+@pytest.mark.skip(
+    reason="Skipping for now because pytorch compiler removes one the two quant ops"
 )
 def test_double_aiter_rms_fp8_group_quant_fusion(
     model_cls: type[torch.nn.Module],

tests/distributed/test_pipeline_parallel.py

@@ -175,7 +175,7 @@ MULTIMODAL_MODELS = {
     "facebook/chameleon-7b": PPTestSettings.fast(),
     "adept/fuyu-8b": PPTestSettings.fast(),
     "zai-org/glm-4v-9b": PPTestSettings.fast(),
-    "OpenGVLab/InternVL2-1B": PPTestSettings.fast(),
+    "OpenGVLab/InternVL3-1B": PPTestSettings.fast(),
     "llava-hf/llava-1.5-7b-hf": PPTestSettings.fast(),
     "llava-hf/llava-v1.6-mistral-7b-hf": PPTestSettings.fast(),
     "llava-hf/LLaVA-NeXT-Video-7B-hf": PPTestSettings.fast(),
@@ -203,7 +203,7 @@ TEST_MODELS = [
     "intfloat/e5-mistral-7b-instruct",
     "BAAI/bge-multilingual-gemma2",
     # [MULTIMODAL GENERATION]
-    "OpenGVLab/InternVL2-1B",
+    "OpenGVLab/InternVL3-1B",
     "microsoft/Phi-3.5-vision-instruct",
     "fixie-ai/ultravox-v0_5-llama-3_2-1b",
     # [LANGUAGE GENERATION - HYBRID ARCH]

tests/entrypoints/pooling/classify/test_online_vision.py

@@ -25,7 +25,7 @@ def server():
         "--runner",
         "pooling",
         "--max-model-len",
-        "5000",
+        "16384",
         "--enforce-eager",
         "--limit-mm-per-prompt",
         json.dumps({"video": MAXIMUM_VIDEOS}),
@@ -143,4 +143,4 @@ def test_chat_video_url_request(server: RemoteOpenAIServer, model_name: str):
     assert output.model == model_name
     assert len(output.data) == 1
     assert len(output.data[0].probs) == 2
-    assert output.usage.prompt_tokens == 4807
+    assert output.usage.prompt_tokens == 8993

tests/kernels/attention/test_attention_selector.py

@@ -15,16 +15,14 @@ from vllm.config import (
 from vllm.platforms import current_platform
 from vllm.platforms.cpu import CpuPlatform
 
-# CudaPlatform and RocmPlatform import their respective compiled C extensions
-# at module level, raising ModuleNotFoundError on incompatible builds.
-try:
+if current_platform.is_cuda():
     from vllm.platforms.cuda import CudaPlatform
-except (ImportError, ModuleNotFoundError):
+else:
     CudaPlatform = None
 
-try:
+if current_platform.is_rocm():
     from vllm.platforms.rocm import RocmPlatform
-except (ImportError, ModuleNotFoundError):
+else:
     RocmPlatform = None
 
 from vllm.v1.attention.backends.registry import AttentionBackendEnum
@@ -434,9 +432,15 @@ def test_per_head_quant_scales_backend_selection(
     [
         ("FLASH_ATTN", True, True),  # FlashAttn supports non-causal
         ("FLASH_ATTN", False, True),  # FlashAttn also works with causal
-        ("FLASHINFER", True, False),  # FlashInfer does not support non-causal
-        ("FLASHINFER", False, True),  # FlashInfer works with causal
-    ],
+    ]
+    + (
+        [
+            ("FLASHINFER", True, False),  # FlashInfer does not support non-causal
+            ("FLASHINFER", False, True),  # FlashInfer works with causal
+        ]
+        if CudaPlatform is not None
+        else []
+    ),
 )
 def test_non_causal_backend_selection(
     backend_name: str, use_non_causal: bool, should_succeed: bool
@@ -459,11 +463,12 @@ def test_non_causal_backend_selection(
         attention_config=attention_config, cache_config=cache_config
     )
 
-    if CudaPlatform is None:
-        pytest.skip("CudaPlatform not available")
+    platform = CudaPlatform or RocmPlatform
+    if platform is None:
+        pytest.skip("CudaPlatform and RocmPlatform are not available")
     with (
         set_current_vllm_config(vllm_config),
-        patch("vllm.platforms.current_platform", CudaPlatform()),
+        patch("vllm.platforms.current_platform", platform()),
     ):
         if should_succeed:
             backend = get_attn_backend(

tests/kernels/attention/test_prefix_prefill.py

@@ -5,10 +5,12 @@ import math
 import random
 import time
 from collections.abc import Callable
+from contextlib import nullcontext
 
 import pytest
 import torch
 import torch.nn.functional as F
+from torch.nn.attention import SDPBackend, sdpa_kernel
 
 from vllm.platforms import current_platform
 from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE, set_random_seed
@@ -557,15 +559,21 @@ def test_contexted_kv_attention_alibi(
             query_len, seq_len, alibi_slopes, device, dtype
         )
 
-        # Compute attention
-        out = F.scaled_dot_product_attention(
-            q_sdpa,
-            k_sdpa,
-            v_sdpa,
-            attn_mask=alibi_mask,
-            dropout_p=0.0,
-            scale=scale,
-        )
+        # Compute attention. On ROCm we force use of the Math SDPA backend rather than
+        # the Flash or Mem-Efficient backends for increased numerical accuracy
+        if current_platform.is_rocm():
+            sdpa_context = sdpa_kernel(SDPBackend.MATH)
+        else:
+            sdpa_context = nullcontext()
+        with sdpa_context:
+            out = F.scaled_dot_product_attention(
+                q_sdpa,
+                k_sdpa,
+                v_sdpa,
+                attn_mask=alibi_mask,
+                dropout_p=0.0,
+                scale=scale,
+            )
 
         # Reshape output back to [query_len, num_heads, head_size]
         out = out.view(num_heads, query_len, head_size).permute(1, 0, 2)

tests/kernels/attention/test_rocm_triton_attn_dsv4.py

@@ -90,7 +90,9 @@ def _ref_sparse_prefill_ragged(
     return out.to(torch.bfloat16)
 
 
-def _pack_fp8_ds_mla_cache(kv: torch.Tensor, block_size: int) -> torch.Tensor:
+def _pack_fp8_ds_mla_cache(
+    kv: torch.Tensor, block_size: int, is_extra: bool = False
+) -> torch.Tensor:
     assert kv.shape[-1] == HEAD_DIM
     num_tokens = kv.shape[0]
     num_blocks = (num_tokens + block_size - 1) // block_size
@@ -101,7 +103,9 @@ def _pack_fp8_ds_mla_cache(kv: torch.Tensor, block_size: int) -> torch.Tensor:
     )
     cache_flat = cache.view(torch.uint8).flatten()
     kv_nope_fp8 = (
-        kv[:, :NOPE_HEAD_DIM].to(current_platform.fp8_dtype()).view(torch.uint8)
+        kv[:, :NOPE_HEAD_DIM]
+        .to(torch.float8_e4m3fn if is_extra else current_platform.fp8_dtype())
+        .view(torch.uint8)
     )
     kv_rope_u8 = kv[:, NOPE_HEAD_DIM:].contiguous().view(torch.uint8)
 
@@ -120,7 +124,7 @@ def _pack_fp8_ds_mla_cache(kv: torch.Tensor, block_size: int) -> torch.Tensor:
 
 
 def _read_fp8_ds_mla_cache(
-    cache: torch.Tensor, slot: int, block_size: int
+    cache: torch.Tensor, slot: int, block_size: int, is_extra: bool = False
 ) -> torch.Tensor:
     cache_flat = cache.view(torch.uint8).flatten()
     block_idx = slot // block_size
@@ -129,7 +133,9 @@ def _read_fp8_ds_mla_cache(
     token_base = block_base + pos * 576
 
     nope_u8 = cache_flat[token_base : token_base + NOPE_HEAD_DIM]
-    nope = nope_u8.view(current_platform.fp8_dtype()).to(torch.float32)
+    nope = nope_u8.view(
+        torch.float8_e4m3fn if is_extra else current_platform.fp8_dtype()
+    ).to(torch.float32)
     rope_u8 = cache_flat[
         token_base + NOPE_HEAD_DIM : token_base + NOPE_HEAD_DIM + ROPE_HEAD_DIM * 2
     ]
@@ -157,7 +163,9 @@ def _ref_sparse_decode_ragged(
         ]
         if extra_cache is not None and extra_rows is not None:
             row_kv.extend(
-                _read_fp8_ds_mla_cache(extra_cache, int(slot), block_size)
+                _read_fp8_ds_mla_cache(
+                    extra_cache, int(slot), block_size, is_extra=True
+                )
                 for slot in extra_rows[query_idx]
             )
 
@@ -326,7 +334,7 @@ def test_sparse_attn_decode_ragged_kernel() -> None:
     main_kv = torch.randn(6, HEAD_DIM, dtype=torch.bfloat16, device=device) * 0.125
     extra_kv = torch.randn(5, HEAD_DIM, dtype=torch.bfloat16, device=device) * 0.125
     main_cache = _pack_fp8_ds_mla_cache(main_kv, block_size)
-    extra_cache = _pack_fp8_ds_mla_cache(extra_kv, block_size)
+    extra_cache = _pack_fp8_ds_mla_cache(extra_kv, block_size, is_extra=True)
     main_indices = torch.tensor([0, 2, 4, 1], dtype=torch.int32, device=device)
     main_indptr = torch.tensor([0, 2, 4], dtype=torch.int32, device=device)
     extra_indices = torch.tensor([1, 3, 0], dtype=torch.int32, device=device)
@@ -477,7 +485,7 @@ def test_sparse_attn_decode_split_k_kernel(
         rows = [[1, 3, 0, 5, 2, 4], [3, 0, 6]]
         extra_kv = torch.randn(7, HEAD_DIM, dtype=torch.bfloat16, device=device) * 0.125
         extra_rows = rows
-        extra_cache = _pack_fp8_ds_mla_cache(extra_kv, block_size)
+        extra_cache = _pack_fp8_ds_mla_cache(extra_kv, block_size, is_extra=True)
         extra_indices, extra_indptr = _ragged_from_rows(rows, device)
 
     attn_sink = (

tests/kernels/attention/test_triton_unified_attention.py

@@ -18,11 +18,7 @@ HEAD_SIZES = [128, 256]
 BLOCK_SIZES = [16]
 
 DTYPES = [torch.bfloat16]
-QDTYPES = (
-    [None, torch.float8_e4m3fn]
-    if not current_platform.is_rocm()
-    else [None, torch.float8_e4m3fnuz]
-)
+QDTYPES = [None, current_platform.fp8_dtype()]
 FP8_DTYPE = current_platform.fp8_dtype()
 
 # one value large enough to test overflow in index calculation.

tests/kernels/core/test_minimax_reduce_rms.py

@@ -10,8 +10,12 @@ from torch.multiprocessing import spawn
 from tests.kernels.utils import opcheck
 from tests.utils import ensure_current_vllm_config, init_test_distributed_environment
 from vllm.distributed import cleanup_dist_env_and_memory
-from vllm.model_executor.layers.minimax_rms_norm import MiniMaxText01RMSNormTP
+from vllm.model_executor.layers.minimax_rms_norm import (
+    MiniMaxText01RMSNormTP,
+    rms_norm_tp,
+)
 from vllm.platforms import current_platform
+from vllm.triton_utils import HAS_TRITON
 from vllm.utils.network_utils import get_open_port
 from vllm.utils.torch_utils import set_random_seed
 
@@ -54,8 +58,19 @@ def _worker_forward_qk(
     torch.manual_seed(seed + 1000 + local_rank)
     qkv = torch.randn(num_tokens, hq + hk + hk, dtype=dtype, device="cuda")
 
-    q_ref, k_ref, v_ref = qkv.clone().split([hq, hk, hk], dim=-1)
-    ref_q, ref_k = MiniMaxText01RMSNormTP.forward_qk(q_norm, k_norm, q_ref, k_ref)
+    # Reference: eager all-reduce path. ``forward_qk`` no longer all-reduces
+    # the variance (it is the tp==1 / already-reduced building block), so the
+    # multi-rank reference must use the eager path that performs the global
+    # variance all-reduce, matching the fused kernel below.
+    ref_q, ref_k = rms_norm_tp._minimax_qk_norm_tp_eager(
+        qkv.clone(),
+        q_norm.weight,
+        k_norm.weight,
+        hq,
+        hk,
+        world_size,
+        eps,
+    )
 
     # Set up Lamport workspace.
     from vllm.distributed.parallel_state import get_tp_group
@@ -150,3 +165,44 @@ def test_minimax_reduce_rms_qk(
         nprocs=world_size,
         join=True,
     )
+
+
+@pytest.mark.skipif(
+    not current_platform.is_cuda() or not HAS_TRITON,
+    reason="CUDA and Triton required",
+)
+@pytest.mark.parametrize("num_tokens", [1, 7, 128, 333, 2049])
+@pytest.mark.parametrize("hidden_dims", [(3072, 512), (768, 256), (3000, 500)])
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
+@pytest.mark.parametrize("tp_world", [1, 4, 8])
+@pytest.mark.parametrize("eps", [1e-6])
+@pytest.mark.parametrize("seed", [42])
+def test_minimax_qk_norm_triton_fallback(
+    monkeypatch, num_tokens, hidden_dims, dtype, tp_world, eps, seed
+):
+    """Single-GPU check: Triton fallback kernels vs the pure-torch reference.
+
+    The all-reduce is a TP communication barrier, so it is monkeypatched to
+    identity here; both the Triton path and the reference see the same
+    (patched) reduction. This validates the kernel math and the folded
+    ``/ tp_world`` scaling without needing multiple ranks -- ``hidden_dims``
+    are the per-rank q/k segment widths.
+    """
+    monkeypatch.setattr(rms_norm_tp, "_all_reduce_variance", lambda v: v)
+
+    q_size, kv_size = hidden_dims
+    device = "cuda"
+    torch.manual_seed(seed)
+    qkv = torch.randn(num_tokens, q_size + 2 * kv_size, dtype=dtype, device=device)
+    q_weight = torch.randn(q_size, dtype=dtype, device=device)
+    k_weight = torch.randn(kv_size, dtype=dtype, device=device)
+
+    q_triton, k_triton = rms_norm_tp._minimax_qk_norm_tp_fallback(
+        qkv, q_weight, k_weight, q_size, kv_size, 0, tp_world, eps
+    )
+    q_ref, k_ref = rms_norm_tp._minimax_qk_norm_tp_eager(
+        qkv, q_weight, k_weight, q_size, kv_size, tp_world, eps
+    )
+
+    torch.testing.assert_close(q_triton, q_ref, atol=3e-2, rtol=3e-2)
+    torch.testing.assert_close(k_triton, k_ref, atol=3e-2, rtol=3e-2)

tests/kernels/moe/test_ocp_mx_moe.py

@@ -9,6 +9,7 @@ import pytest
 import torch
 from packaging import version
 
+from vllm._aiter_ops import is_aiter_found
 from vllm.platforms import current_platform
 from vllm.utils.flashinfer import has_flashinfer
 
@@ -31,17 +32,15 @@ HOPPER_MXFP4_BF16_AVAILABLE = (
 # ROCm platform and dependencies
 ROCM_AVAILABLE = current_platform.is_rocm()
 ROCM_TRITON_KERNELS_AVAILABLE = False
-ROCM_AITER_AVAILABLE = False
+ROCM_AITER_AVAILABLE = is_aiter_found()
 ROCM_GFX950 = False
 
 if ROCM_AVAILABLE:
-    from vllm._aiter_ops import rocm_aiter_ops
     from vllm.platforms.rocm import on_gfx950
     from vllm.utils.import_utils import has_triton_kernels
 
     ROCM_TRITON_KERNELS_AVAILABLE = has_triton_kernels()
     ROCM_GFX950 = on_gfx950()
-    ROCM_AITER_AVAILABLE = rocm_aiter_ops.is_enabled()
 
     if ROCM_AITER_AVAILABLE:
         from aiter.ops.triton.moe.quant_moe import upcast_from_mxfp
@@ -83,7 +82,7 @@ def enable_pickle(monkeypatch):
     [
         ModelCase("fxmarty/qwen_1.5-moe-a2.7b-mxfp4", tp=2),
         ModelCase("fxmarty/deepseek_r1_3_layers_mxfp4", tp=8),
-        ModelCase("fxmarty/Llama-4-Scout-17B-16E-Instruct-2-layers-mxfp4", tp=1),
+        ModelCase("mawong-amd/Llama-4-Scout-17B-16E-Instruct-2-layers-mxfp4", tp=1),
         ModelCase("fxmarty/Llama-3.1-70B-Instruct-2-layers-mxfp6", tp=1),
         ModelCase("fxmarty/Llama-3.1-70B-Instruct-2-layers-mxfp6", tp=4),
     ],
@@ -102,6 +101,7 @@ def test_mxfp4_loading_and_execution_moe(vllm_runner, model_case: ModelCase):
         tensor_parallel_size=model_case.tp,
         load_format="dummy",
         compilation_config={"cudagraph_capture_sizes": [16]},
+        gpu_memory_utilization=0.8,  # mxfp6 models use more scratch space
     ) as llm:
         # Disabled as check_model is broken: https://github.com/vllm-project/vllm/pull/18465#issuecomment-3329880562
         # def check_model(model):
@@ -1267,7 +1267,7 @@ def test_rocm_mxfp4_moe_oracle(
 
     This test validates that the oracle functions work end-to-end:
     - select_mxfp4_moe_backend() selects a valid backend
-    - convert_to_mxfp4_moe_kernel_format() converts weights without error
+    - convert_gpt_oss_weight_to_mxfp4_moe_kernel_format() converts weights without error
     - make_mxfp4_moe_quant_config() builds a valid quant config
     - make_mxfp4_moe_kernel() creates a kernel that runs without error
     - The kernel output is within accuracy tolerance of reference
@@ -1287,7 +1287,7 @@ def test_rocm_mxfp4_moe_oracle(
     from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import (
         Mxfp4MoeBackend,
         backend_to_kernel_cls,
-        convert_to_mxfp4_moe_kernel_format,
+        convert_gpt_oss_weight_to_mxfp4_moe_kernel_format,
         make_mxfp4_moe_kernel,
         make_mxfp4_moe_quant_config,
     )
@@ -1387,7 +1387,7 @@ def test_rocm_mxfp4_moe_oracle(
 
     # Convert weights using oracle
     w13_conv, w2_conv, w13_scale_conv, w2_scale_conv, w13_bias_conv, w2_bias_conv = (
-        convert_to_mxfp4_moe_kernel_format(
+        convert_gpt_oss_weight_to_mxfp4_moe_kernel_format(
             mxfp4_backend=backend,
             layer=layer,  # type: ignore[arg-type]
             w13_weight=w13_quant,
@@ -1423,7 +1423,7 @@ def test_rocm_mxfp4_moe_oracle(
             mxfp4_backend=backend,
             experts_cls=experts_cls,
             routing_tables=None,
-            shared_experts=None,
+            layer=None,
         )
 
         # Create inputs

tests/models/multimodal/processing/test_qwen3_vl.py

@@ -138,3 +138,49 @@ def test_processor_multi_video(
         assert video_phs[i].offset >= prev_end, (
             f"Placeholder {i} overlaps with placeholder {i - 1}"
         )
+
+
+@pytest.mark.parametrize("model_id", [MODEL_ID])
+@pytest.mark.parametrize(
+    "hf_mm_kwargs",
+    [{"num_frames": [8, 16]}, {"fps": [2.0, 4.0]}],
+)
+def test_processor_multi_video_list_kwargs(
+    model_id: str,
+    hf_mm_kwargs: dict[str, Any],
+) -> None:
+    """Regression test: a multi-video request with list-valued per-video
+    ``mm_processor_kwargs`` (one ``fps``/``num_frames`` per video) must not
+    crash.
+
+    Before the fix, ``_call_hf_processor`` copied the whole kwargs to every
+    video without slicing, so ``_get_video_second_idx`` received the list
+    where a scalar was expected and raised ``TypeError``.
+    """
+    ctx = build_model_context(
+        model_id,
+        limit_mm_per_prompt={"image": 0, "video": 2},
+    )
+    processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config)
+
+    prompt = (
+        "<|vision_start|><|video_pad|><|vision_end|>"
+        "<|vision_start|><|video_pad|><|vision_end|>"
+    )
+    mm_data = {
+        "video": [
+            _build_video_mm_data(num_frames=16)["video"][0],
+            _build_video_mm_data(num_frames=32)["video"][0],
+        ]
+    }
+
+    processed = processor(
+        prompt,
+        mm_items=processor.info.parse_mm_data(mm_data),
+        hf_processor_mm_kwargs=hf_mm_kwargs,
+    )
+
+    video_phs = processed["mm_placeholders"].get("video", [])
+    assert len(video_phs) == 2, (
+        f"Expected exactly 2 video placeholders, got {len(video_phs)}"
+    )

tests/models/registry.py

@@ -1530,6 +1530,16 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {
         "Qwen/Qwen3-VL-8B-Instruct",
         speculative_model="taobao-mnn/Qwen3-VL-8B-Instruct-Eagle3",
     ),
+    "Eagle3Qwen3ForCausalLM": _HfExamplesInfo(
+        "Qwen/Qwen3-8B",
+        trust_remote_code=True,
+        speculative_model=(
+            "inference-optimization/"
+            "Qwen3-8B-from-Qwen3-8B_regen-speculators.eagle3-qwen3arch-ckpt1"
+        ),
+        tokenizer="Qwen/Qwen3-8B",
+        use_original_num_layers=True,
+    ),
     # [PEagle]
     "PEagleDraftModel": _HfExamplesInfo(
         "Qwen/Qwen3-8B",
@@ -1545,6 +1555,15 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {
         tokenizer="Qwen/Qwen3-8B",
         use_original_num_layers=True,
     ),
+    "PeagleQwen3ForCausalLM": _HfExamplesInfo(
+        "Qwen/Qwen3-8B",
+        trust_remote_code=True,
+        speculative_model=(
+            "inference-optimization/Qwen3-8B-speculators.peagle-qwen3arch-ckpt4"
+        ),
+        tokenizer="Qwen/Qwen3-8B",
+        use_original_num_layers=True,
+    ),
     # [MTP]
     "DeepSeekMTPModel": _HfExamplesInfo(
         "luccafong/deepseek_mtp_main_random",

tests/models/test_deepseek_v4_mega_moe.py

@@ -46,7 +46,8 @@ def test_deepseek_v4_mega_moe_ue8m0_uint8_to_float():
 
 def test_deepseek_v4_mega_moe_weight_loader_uses_ep_expert_ownership():
     vllm_config = SimpleNamespace(
-        scheduler_config=SimpleNamespace(max_num_batched_tokens=4)
+        scheduler_config=SimpleNamespace(max_num_batched_tokens=4),
+        compilation_config=SimpleNamespace(static_forward_context={}),
     )
     experts = DeepseekV4MegaMoEExperts(
         vllm_config,
@@ -182,3 +183,81 @@ def test_deepseek_v4_mega_moe_fused_input_staging_is_bitwise_exact():
         fused_topk_weights.view(torch.uint8),
         ref_topk_weights.view(torch.uint8),
     )
+
+
+@pytest.mark.skipif(
+    not torch.cuda.is_available(),
+    reason="DeepSeek V4 MegaMoE fused input staging requires CUDA.",
+)
+def test_deepseek_v4_mega_moe_fused_input_staging_masks_padding():
+    from vllm.third_party.deep_gemm.utils import per_token_cast_to_fp8
+
+    device = torch.device("cuda")
+    num_tokens = 7
+    hidden_size = 256
+    top_k = 8
+
+    generator = torch.Generator(device=device)
+    generator.manual_seed(1)
+    hidden_states = torch.randn(
+        num_tokens,
+        hidden_size,
+        device=device,
+        dtype=torch.bfloat16,
+        generator=generator,
+    )
+    topk_ids = torch.randint(
+        0,
+        256,
+        (num_tokens, top_k),
+        device=device,
+        dtype=torch.int32,
+        generator=generator,
+    )
+    topk_weights = torch.randn(
+        num_tokens,
+        top_k,
+        device=device,
+        dtype=torch.float32,
+        generator=generator,
+    )
+    is_padding = torch.tensor(
+        [False, True, False, False, True, False, True],
+        device=device,
+    )
+
+    ref_x, ref_x_sf = per_token_cast_to_fp8(
+        hidden_states,
+        use_ue8m0=True,
+        gran_k=32,
+        use_packed_ue8m0=True,
+    )
+    ref_topk_idx = topk_ids.to(torch.int64)
+    ref_topk_idx[is_padding] = -1
+    ref_topk_weights = topk_weights.clone()
+    ref_topk_weights[is_padding] = 0.0
+
+    fused_x = torch.empty_like(ref_x)
+    fused_x_sf = torch.empty_like(ref_x_sf)
+    fused_topk_idx = torch.empty_like(ref_topk_idx)
+    fused_topk_weights = torch.empty_like(ref_topk_weights)
+
+    prepare_megamoe_inputs(
+        hidden_states,
+        topk_weights,
+        topk_ids,
+        fused_x,
+        fused_x_sf,
+        fused_topk_idx,
+        fused_topk_weights,
+        is_padding=is_padding,
+    )
+    torch.accelerator.synchronize()
+
+    assert torch.equal(fused_x.view(torch.uint8), ref_x.view(torch.uint8))
+    assert torch.equal(fused_x_sf, ref_x_sf)
+    assert torch.equal(fused_topk_idx, ref_topk_idx)
+    assert torch.equal(
+        fused_topk_weights.view(torch.uint8),
+        ref_topk_weights.view(torch.uint8),
+    )

tests/multimodal/test_video.py

@@ -15,6 +15,7 @@ from vllm.multimodal.video import (
     DynamicVideoBackend,
     GLM46VVideoBackend,
     Molmo2VideoBackend,
+    Qwen2VLVideoBackend,
     Qwen3VLVideoBackend,
     VideoLoader,
     VideoSourceMetadata,
@@ -70,11 +71,12 @@ def test_video_loader_type_doesnt_exist():
 
 
 @pytest.mark.parametrize(
-    "model_repo, expected_loader_cls",
+    "model_repo, expected_loader_cls, hf_sample_kwargs",
     [
         pytest.param(
             "allenai/Molmo2-4B",
             Molmo2VideoBackend,
+            None,
             marks=pytest.mark.skip(
                 reason="Video processor not aligned, investigate later.",
             ),
@@ -83,23 +85,44 @@ def test_video_loader_type_doesnt_exist():
         pytest.param(
             "zai-org/GLM-4.1V-9B-Thinking",
             DynamicVideoBackend,
+            None,
             id="glm4v",
         ),
         pytest.param(
             "zai-org/GLM-4.6V-Flash",
             GLM46VVideoBackend,
+            None,
             id="glm46v",
         ),
         pytest.param(
             "Qwen/Qwen3-VL-4B-Instruct",
             Qwen3VLVideoBackend,
+            None,
             id="qwen3vl",
         ),
+        # Qwen2-VL/Qwen2.5-VL ship no ``video_processor_type`` in their
+        # preprocessor config, so resolution relies on the model_type ->
+        # video processor fallback in get_video_processor_cls_name_from_config.
+        # They also ship no default fps/num_frames, so the HF sampler needs an
+        # explicit target rate; pass fps=2 to match the loader default.
+        pytest.param(
+            "Qwen/Qwen2-VL-7B-Instruct",
+            Qwen2VLVideoBackend,
+            {"fps": 2},
+            id="qwen2vl",
+        ),
+        pytest.param(
+            "Qwen/Qwen2.5-VL-7B-Instruct",
+            Qwen2VLVideoBackend,
+            {"fps": 2},
+            id="qwen2_5_vl",
+        ),
     ],
 )
 def test_video_processor_from_model_repo(
     model_repo: str,
     expected_loader_cls: type,
+    hf_sample_kwargs: dict[str, int | float] | None,
 ):
     """Test that a model repo resolves to the correct video loader backend.
 
@@ -143,7 +166,7 @@ def test_video_processor_from_model_repo(
             fps=vllm_meta["fps"],
             duration=vllm_meta["duration"],
         )
-        hf_indices = processor.sample_frames(hf_metadata)
+        hf_indices = processor.sample_frames(hf_metadata, **(hf_sample_kwargs or {}))
         vllm_indices = np.array(vllm_meta["frames_indices"])
         np.testing.assert_array_equal(
             hf_indices,

tests/quantization/test_per_token_kv_cache.py

@@ -61,8 +61,8 @@ class QuantConfig:
     quant_max: float
     quant_min: float
     kv_quant_mode: KVQuantMode
-    # INT8 Triton stores truncate; FP8 hardware casts round.
-    uses_trunc: bool
+    # INT8 rounds explicitly; FP8 relies on dtype cast rounding.
+    rounds_before_store: bool
 
 
 INT8_CONFIG = QuantConfig(
@@ -71,7 +71,7 @@ INT8_CONFIG = QuantConfig(
     quant_max=127.0,
     quant_min=-128.0,
     kv_quant_mode=KVQuantMode.INT8_PER_TOKEN_HEAD,
-    uses_trunc=True,
+    rounds_before_store=True,
 )
 FP8_CONFIG = QuantConfig(
     cache_dtype=FP8_DTYPE,
@@ -79,7 +79,7 @@ FP8_CONFIG = QuantConfig(
     quant_max=FP8_MAX,
     quant_min=FP8_MIN,
     kv_quant_mode=KVQuantMode.FP8_PER_TOKEN_HEAD,
-    uses_trunc=False,
+    rounds_before_store=False,
 )
 
 QUANT_CONFIGS = [INT8_CONFIG, FP8_CONFIG]
@@ -104,7 +104,7 @@ def _quantize_per_token_head_ref(
     absmax = data.float().abs().amax(dim=2)  # [num_tokens, num_heads]
     scales = (absmax / cfg.quant_max).clamp(min=1e-6)
     scaled = data.float() * (1.0 / scales[:, :, None])
-    if cfg.uses_trunc:
+    if cfg.rounds_before_store:
         q = scaled.round().clamp(cfg.quant_min, cfg.quant_max).to(cfg.cache_dtype)
     else:
         q = scaled.clamp(cfg.quant_min, cfg.quant_max).to(cfg.cache_dtype)
@@ -255,7 +255,7 @@ def test_per_token_head_round_trip_accuracy(
 ):
     """Verify per-token-head round-trip: kernel dequant matches reference.
 
-    INT8: Triton truncates on float->int8 store.
+    INT8: round-to-nearest before int8 store.
     FP8: hardware cast (clamp then cast).
     """
     from vllm.v1.attention.ops.triton_reshape_and_cache_flash import (
@@ -315,6 +315,52 @@ def test_per_token_head_round_trip_accuracy(
                 )
 
 
+@torch.inference_mode()
+def test_int8_per_token_head_raw_cache_matches_round_reference():
+    """INT8 cache writes should match round-to-nearest quantization exactly."""
+    from vllm.v1.attention.ops.triton_reshape_and_cache_flash import (
+        triton_reshape_and_cache_flash_per_token_head_quant,
+    )
+
+    torch.set_default_device(DEVICE_TYPE)
+
+    head_size = 8
+    block_size = 4
+
+    key = torch.tensor(
+        [[[-127.0, -2.6, -2.4, -1.6, -1.4, -0.6, -0.4, 127.0]]],
+        dtype=torch.bfloat16,
+    )
+    value = -key
+
+    key_cache = torch.zeros(1, block_size, 1, head_size, dtype=torch.int8)
+    value_cache = torch.zeros_like(key_cache)
+    k_scale_cache = torch.ones(1, block_size, 1, dtype=torch.float32)
+    v_scale_cache = torch.ones_like(k_scale_cache)
+    slot_mapping = torch.tensor([2], dtype=torch.long)
+
+    triton_reshape_and_cache_flash_per_token_head_quant(
+        key,
+        value,
+        key_cache,
+        value_cache,
+        k_scale_cache,
+        v_scale_cache,
+        slot_mapping,
+    )
+
+    ref_k_quant, ref_k_scales = _quantize_per_token_head_ref(key, INT8_CONFIG)
+    ref_v_quant, ref_v_scales = _quantize_per_token_head_ref(value, INT8_CONFIG)
+
+    slot = slot_mapping.item()
+    blk = slot // block_size
+    off = slot % block_size
+    assert torch.equal(key_cache[blk, off], ref_k_quant[0])
+    assert torch.equal(value_cache[blk, off], ref_v_quant[0])
+    torch.testing.assert_close(k_scale_cache[blk, off], ref_k_scales[0])
+    torch.testing.assert_close(v_scale_cache[blk, off], ref_v_scales[0])
+
+
 # ===========================================================================
 # 4. Negative slot mapping (padding tokens should be skipped)
 # ===========================================================================
@@ -461,7 +507,7 @@ def test_triton_unified_attention_per_token_head_scale(
 
     scaled_k = key_cache_bf16.float() / k_scale_cache[:, :, :, None]
     scaled_v = value_cache_bf16.float() / v_scale_cache[:, :, :, None]
-    if qcfg.uses_trunc:
+    if qcfg.rounds_before_store:
         key_cache_q = (
             scaled_k.round().clamp(qcfg.quant_min, qcfg.quant_max).to(qcfg.cache_dtype)
         )

tests/v1/core/test_kv_cache_utils.py

@@ -117,6 +117,7 @@ def new_kv_cache_spec(
     page_size_padded=None,
     sliding_window=None,
     attention_chunk_size=None,
+    indexes_kv_by_block_stride=False,
 ):
     return FullAttentionSpec(
         block_size=block_size,
@@ -126,6 +127,7 @@ def new_kv_cache_spec(
         page_size_padded=page_size_padded,
         sliding_window=sliding_window,
         attention_chunk_size=attention_chunk_size,
+        indexes_kv_by_block_stride=indexes_kv_by_block_stride,
     )
 
 
@@ -136,6 +138,7 @@ def new_sliding_window_spec(
     dtype=torch.float32,
     page_size_padded=None,
     sliding_window=1,
+    indexes_kv_by_block_stride=False,
 ):
     return SlidingWindowSpec(
         block_size=block_size,
@@ -144,6 +147,7 @@ def new_sliding_window_spec(
         dtype=dtype,
         page_size_padded=page_size_padded,
         sliding_window=sliding_window,
+        indexes_kv_by_block_stride=indexes_kv_by_block_stride,
     )
 
 
@@ -1799,16 +1803,38 @@ def test_get_kv_cache_config_one_worker():
         ],
     )
 
-    # different hidden size that cannot be aligned by using different block size
+    # different hidden size that cannot be aligned by using different block size,
+    # but can be aligned by padding the smaller physical page.
+    swa_spec = new_sliding_window_spec(head_size=96, indexes_kv_by_block_stride=True)
     kv_cache_specs_hybrid = {
-        "layer_1": new_kv_cache_spec(head_size=64),
-        "layer_2": new_sliding_window_spec(head_size=96),
+        "layer_1": new_kv_cache_spec(head_size=64, indexes_kv_by_block_stride=True),
+        "layer_2": swa_spec,
     }
 
-    with pytest.raises(NotImplementedError):
-        get_kv_cache_configs(
-            vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 2 * 32]
-        )[0]
+    kv_cache_config_hybrid = get_kv_cache_configs(
+        vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 2 * 32]
+    )[0]
+    padded_page_size = swa_spec.page_size_bytes
+    assert kv_cache_config_hybrid == KVCacheConfig(
+        num_blocks=42,
+        kv_cache_tensors=[
+            KVCacheTensor(size=padded_page_size * 42, shared_by=["layer_1", "layer_2"]),
+        ],
+        kv_cache_groups=[
+            KVCacheGroupSpec(
+                ["layer_1"],
+                new_kv_cache_spec(
+                    head_size=64,
+                    page_size_padded=padded_page_size,
+                    indexes_kv_by_block_stride=True,
+                ),
+            ),
+            KVCacheGroupSpec(
+                ["layer_2"],
+                new_sliding_window_spec(head_size=96, indexes_kv_by_block_stride=True),
+            ),
+        ],
+    )
 
     # Test num_gpu_blocks_override
     vllm_config.cache_config.num_gpu_blocks_override = 16
@@ -2322,6 +2348,75 @@ def test_check_enough_kv_cache_memory_respects_num_gpu_blocks_override():
         get_kv_cache_configs(vllm_config, [kv_cache_specs], [large_available_memory])
 
 
+def test_unify_kv_cache_page_size_uses_padding_for_non_divisible_sizes():
+    """DFlash drafters can have a smaller head size than the target model.
+
+    For example, MiMo uses 192-dim target KV heads while its DFlash draft uses
+    128-dim KV heads. The resulting page sizes are 3:2 rather than an integer
+    block-size multiple, so the smaller page must be padded instead.
+    """
+    # Both layers' backends opt into the padded-page strided view (e.g.
+    # FlashAttention / its DiffKV subclass), so padding is allowed.
+    target_spec = new_kv_cache_spec(
+        block_size=16,
+        num_kv_heads=1,
+        head_size=192,
+        dtype=torch.bfloat16,
+        indexes_kv_by_block_stride=True,
+    )
+    draft_spec = new_sliding_window_spec(
+        block_size=16,
+        num_kv_heads=1,
+        head_size=128,
+        dtype=torch.bfloat16,
+        sliding_window=1024,
+        indexes_kv_by_block_stride=True,
+    )
+
+    unified_specs = kv_cache_utils.unify_kv_cache_spec_page_size(
+        {
+            "target_attn": target_spec,
+            "draft_attn": draft_spec,
+        }
+    )
+
+    assert unified_specs["target_attn"] == target_spec
+    unified_draft_spec = unified_specs["draft_attn"]
+    assert unified_draft_spec.block_size == draft_spec.block_size
+    assert unified_draft_spec.real_page_size_bytes == draft_spec.real_page_size_bytes
+    assert unified_draft_spec.page_size_padded == target_spec.page_size_bytes
+    assert unified_draft_spec.page_size_bytes == target_spec.page_size_bytes
+
+
+def test_unify_kv_cache_page_size_padding_requires_backend_support():
+    """Padding is gated on the backend declaring ``indexes_kv_by_block_stride``.
+
+    A backend that does not support the strided padded-page view must raise
+    rather than silently padding (and misreading KV at runtime).
+    """
+    target_spec = new_kv_cache_spec(
+        block_size=16,
+        num_kv_heads=1,
+        head_size=192,
+        dtype=torch.bfloat16,
+        indexes_kv_by_block_stride=True,
+    )
+    # The non-divisible draft layer needs padding but its backend does not
+    # support the strided padded-page view -> must raise, not silently pad.
+    draft_spec = new_sliding_window_spec(
+        block_size=16,
+        num_kv_heads=1,
+        head_size=128,
+        dtype=torch.bfloat16,
+        sliding_window=1024,
+        indexes_kv_by_block_stride=False,
+    )
+    specs = {"target_attn": target_spec, "draft_attn": draft_spec}
+
+    with pytest.raises(NotImplementedError):
+        kv_cache_utils.unify_kv_cache_spec_page_size(specs)
+
+
 def test_unify_hybrid_kv_cache_specs():
     # 1. has_full_attention and has_sliding_window
     before_spec_1 = new_kv_cache_spec()

tests/v1/e2e/spec_decode/test_spec_decode.py

@@ -425,7 +425,7 @@ def _run_eagle_correctness(
             if "deepseek" in model_setup[1].lower():
                 m.setenv("VLLM_ROCM_USE_AITER", "1")
                 m.delenv("VLLM_MLA_DISABLE", raising=False)
-                attention_config = {"backend": "TRITON_MLA"}
+                attention_config = {"backend": "ROCM_AITER_MLA"}
             else:
                 m.setenv("VLLM_ROCM_USE_AITER", "1")
 

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

@@ -22,7 +22,7 @@ from vllm.v1.kv_offload.base import (
     OffloadingGaugeMetadata,
     OffloadingHistogramMetadata,
 )
-from vllm.v1.kv_offload.cpu.spec import CPUOffloadingSpec
+from vllm.v1.kv_offload.factory import OffloadingSpecFactory
 
 LOAD_BYTES = _TransferMetricName.LOAD_BYTES
 LOAD_TIME = _TransferMetricName.LOAD_TIME
@@ -33,6 +33,8 @@ STORE_SIZE = _TransferMetricName.STORE_SIZE
 STORES_SKIPPED = "vllm:kv_offload_stores_skipped"
 PENDING_STORES = "vllm:kv_offload_pending_stores"
 LOOKUP_LATENCY = "vllm:kv_offload_lookup_latency_seconds"
+MY_COUNTER = "my_counter"
+MY_LABEL = "my_label"
 
 
 class _FakeMetric:
@@ -67,6 +69,20 @@ class _FakeVllmConfig:
         )
 
 
+def _spec_cls_with_metric_definitions(
+    metric_definitions: dict[str, Any],
+) -> type:
+    """Build a fake offloading spec class reporting the given metric
+    definitions, so tests don't need to patch the real CPU spec."""
+
+    class _FakeOffloadingSpec:
+        @staticmethod
+        def build_metric_definitions(extra_config):
+            return metric_definitions
+
+    return _FakeOffloadingSpec
+
+
 def _metric_metadata():
     return {
         LOAD_BYTES: OffloadingCounterMetadata(
@@ -96,9 +112,17 @@ def _metric_metadata():
         LOOKUP_LATENCY: OffloadingHistogramMetadata(
             documentation="lookup latency",
         ),
+        MY_COUNTER: OffloadingCounterMetadata(
+            documentation="counter with a label",
+            labelnames=(MY_LABEL,),
+        ),
     }
 
 
+def _unlabeled(values: dict[str, Any], metric_name: str) -> Any:
+    return values[metric_name][()]
+
+
 def test_build_kv_connector_stats_with_none():
     """Test that build_kv_connector_stats returns empty stats when given None."""
     stats = OffloadingConnector.build_kv_connector_stats(data=None)
@@ -131,13 +155,13 @@ def test_build_kv_connector_stats_reconstructs_offload_stats():
             STORES_SKIPPED: _MetricType.COUNTER,
         },
         _StatsKey.DATA: {
-            LOAD_BYTES: 24,
-            LOAD_TIME: 1.5,
-            LOAD_SIZE: [16, 8],
-            STORE_BYTES: 3,
-            STORE_TIME: 0.3,
-            STORE_SIZE: [1, 2],
-            STORES_SKIPPED: 5,
+            LOAD_BYTES: {(): 24},
+            LOAD_TIME: {(): 1.5},
+            LOAD_SIZE: {(): [16, 8]},
+            STORE_BYTES: {(): 3},
+            STORE_TIME: {(): 0.3},
+            STORE_SIZE: {(): [1, 2]},
+            STORES_SKIPPED: {(): 5},
         },
     }
 
@@ -145,22 +169,28 @@ def test_build_kv_connector_stats_reconstructs_offload_stats():
 
     assert isinstance(stats, OffloadingConnectorStats)
     values = stats.data[_StatsKey.DATA]
-    assert values[LOAD_BYTES] == 24
-    assert values[LOAD_TIME] == 1.5
-    assert values[LOAD_SIZE] == [16, 8]
-    assert values[STORE_BYTES] == 3
-    assert values[STORE_TIME] == 0.3
-    assert values[STORE_SIZE] == [1, 2]
-    assert values[STORES_SKIPPED] == 5
+    assert _unlabeled(values, LOAD_BYTES) == 24
+    assert _unlabeled(values, LOAD_TIME) == 1.5
+    assert _unlabeled(values, LOAD_SIZE) == [16, 8]
+    assert _unlabeled(values, STORE_BYTES) == 3
+    assert _unlabeled(values, STORE_TIME) == 0.3
+    assert _unlabeled(values, STORE_SIZE) == [1, 2]
+    assert _unlabeled(values, STORES_SKIPPED) == 5
 
 
 def _make_stats_data(
     metric_data: dict[str, Any],
     metric_metadata: dict[str, Any],
 ) -> dict[str, Any]:
-    """Build a structured data dict from flat metric data and metadata."""
+    """Build a structured data dict from flat metric data and metadata.
+
+    Values for unlabeled metrics may be passed flat (wrapped here under the
+    empty label tuple); values for labeled metrics must already be passed as
+    a ``{labelvalues: value}`` map.
+    """
     metric_types = {}
-    for key in metric_data:
+    data = {}
+    for key, value in metric_data.items():
         md = metric_metadata[key]
         if isinstance(md, OffloadingCounterMetadata):
             metric_types[key] = _MetricType.COUNTER
@@ -168,9 +198,10 @@ def _make_stats_data(
             metric_types[key] = _MetricType.GAUGE
         elif isinstance(md, OffloadingHistogramMetadata):
             metric_types[key] = _MetricType.HISTOGRAM
+        data[key] = value if md.labelnames else {(): value}
     return {
         _StatsKey.TYPES: metric_types,
-        _StatsKey.DATA: metric_data,
+        _StatsKey.DATA: data,
     }
 
 
@@ -215,34 +246,106 @@ def test_aggregate_same_connector():
 
     assert result is stats1  # Should return self
     values = result.data[_StatsKey.DATA]
-    assert values[LOAD_BYTES] == 34
-    assert values[LOAD_TIME] == 2.6
-    assert values[LOAD_SIZE] == [16, 8, 3, 7]
-    assert values[STORE_BYTES] == 19
-    assert values[STORE_TIME] == 2.3
-    assert values[STORE_SIZE] == [1, 2, 16]
-    assert values[STORES_SKIPPED] == 4
-    assert values[PENDING_STORES] == 1
-    assert values[LOOKUP_LATENCY] == [0.1, 0.2, 0.3]
+    assert _unlabeled(values, LOAD_BYTES) == 34
+    assert _unlabeled(values, LOAD_TIME) == 2.6
+    assert _unlabeled(values, LOAD_SIZE) == [16, 8, 3, 7]
+    assert _unlabeled(values, STORE_BYTES) == 19
+    assert _unlabeled(values, STORE_TIME) == 2.3
+    assert _unlabeled(values, STORE_SIZE) == [1, 2, 16]
+    assert _unlabeled(values, STORES_SKIPPED) == 4
+    assert _unlabeled(values, PENDING_STORES) == 1
+    assert _unlabeled(values, LOOKUP_LATENCY) == [0.1, 0.2, 0.3]
+
+
+def test_aggregate_labeled_metrics():
+    metadata = _metric_metadata()
+    stats1 = OffloadingConnectorStats(
+        data=_make_stats_data(
+            {
+                MY_COUNTER: {
+                    ("a",): 10,
+                    ("b",): 3,
+                },
+            },
+            metadata,
+        ),
+    )
+    stats2 = OffloadingConnectorStats(
+        data=_make_stats_data(
+            {
+                MY_COUNTER: {
+                    ("a",): 7,
+                    ("c",): 5,
+                },
+            },
+            metadata,
+        ),
+    )
+
+    stats1.aggregate(stats2)
+
+    values = stats1.data[_StatsKey.DATA][MY_COUNTER]
+    assert values[("a",)] == 17
+    assert values[("b",)] == 3
+    assert values[("c",)] == 5
+
+
+def test_aggregate_labeled_metric_missing_from_self():
+    """Aggregating a labeled metric that self doesn't have at all yet."""
+    metadata = _metric_metadata()
+    stats1 = OffloadingConnectorStats()
+    stats2 = OffloadingConnectorStats(
+        data=_make_stats_data(
+            {
+                MY_COUNTER: {
+                    ("a",): 7,
+                    ("b",): 5,
+                },
+            },
+            metadata,
+        ),
+    )
+
+    stats1.aggregate(stats2)
+
+    values = stats1.data[_StatsKey.DATA][MY_COUNTER]
+    assert values[("a",)] == 7
+    assert values[("b",)] == 5
+    assert stats1.data[_StatsKey.TYPES][MY_COUNTER] == _MetricType.COUNTER
+
+
+def test_helper_methods_accept_labeled_metrics():
+    stats = OffloadingConnectorStats()
+
+    stats.increase_counter(MY_COUNTER, 3, ("a",))
+    stats.increase_counter(MY_COUNTER, 4, ("a",))
+    stats.set_gauge(PENDING_STORES, 2, ("b",))
+    stats.observe_histogram(LOOKUP_LATENCY, 0.1, ("b",))
+    stats.observe_histogram(LOOKUP_LATENCY, 0.2, ("b",))
+
+    values = stats.data[_StatsKey.DATA]
+    assert values[MY_COUNTER][("a",)] == 7
+    assert values[PENDING_STORES][("b",)] == 2
+    assert values[LOOKUP_LATENCY][("b",)] == [0.1, 0.2]
 
 
 def test_aggregate_merges_types():
     stats1 = OffloadingConnectorStats(
         data={
             _StatsKey.TYPES: {LOAD_BYTES: _MetricType.COUNTER},
-            _StatsKey.DATA: {LOAD_BYTES: 1},
+            _StatsKey.DATA: {LOAD_BYTES: {(): 1}},
         },
     )
     stats2 = OffloadingConnectorStats(
         data={
             _StatsKey.TYPES: {PENDING_STORES: _MetricType.GAUGE},
-            _StatsKey.DATA: {PENDING_STORES: 2},
+            _StatsKey.DATA: {PENDING_STORES: {(): 2}},
         },
     )
 
     result = stats1.aggregate(stats2)
 
-    assert result.data[_StatsKey.DATA][PENDING_STORES] == 2
+    assert _unlabeled(result.data[_StatsKey.DATA], PENDING_STORES) == 2
     assert result.data[_StatsKey.TYPES][PENDING_STORES] == _MetricType.GAUGE
 
 
@@ -283,6 +386,26 @@ def test_reduce():
     assert reduced[f"{LOOKUP_LATENCY}_sum"] == sum([0.1, 0.2, 0.3])
 
 
+def test_reduce_labeled_metrics():
+    metadata = _metric_metadata()
+    stats = OffloadingConnectorStats(
+        data=_make_stats_data(
+            {
+                MY_COUNTER: {
+                    ("a",): 17,
+                    ("b",): 3,
+                },
+            },
+            metadata,
+        ),
+    )
+
+    reduced = stats.reduce()
+
+    assert reduced[f"{MY_COUNTER}:{('a',)}"] == 17
+    assert reduced[f"{MY_COUNTER}:{('b',)}"] == 3
+
+
 def test_reset():
     """Test that reset() resets all connector stats."""
     metadata = _metric_metadata()
@@ -326,11 +449,11 @@ def test_prom_metrics_observes_manager_counter():
     prom_metrics.observe(
         {
             _StatsKey.TYPES: {STORES_SKIPPED: _MetricType.COUNTER},
-            _StatsKey.DATA: {STORES_SKIPPED: 7},
+            _StatsKey.DATA: {STORES_SKIPPED: {(): 7}},
         }
     )
 
-    counter = prom_metrics.offloading_metrics[(0, STORES_SKIPPED)]
+    counter = prom_metrics.offloading_metrics[(0, STORES_SKIPPED, ())]
     assert counter.increments == [7]
     counter_def = prom_metrics._offloading_metric_defs[STORES_SKIPPED]
     assert counter_def.kwargs["name"] == "vllm:kv_offload_stores_skipped"
@@ -360,22 +483,22 @@ def test_prom_metrics_observes_flat_transfer_metrics_and_legacy_metrics():
                 STORE_SIZE: _MetricType.HISTOGRAM,
             },
             _StatsKey.DATA: {
-                LOAD_BYTES: 24,
-                LOAD_TIME: 1.5,
-                LOAD_SIZE: [16, 8],
-                STORE_BYTES: 3,
-                STORE_TIME: 0.3,
-                STORE_SIZE: [1, 2],
+                LOAD_BYTES: {(): 24},
+                LOAD_TIME: {(): 1.5},
+                LOAD_SIZE: {(): [16, 8]},
+                STORE_BYTES: {(): 3},
+                STORE_TIME: {(): 0.3},
+                STORE_SIZE: {(): [1, 2]},
             },
         }
     )
 
-    assert prom_metrics.offloading_metrics[(0, LOAD_BYTES)].increments == [24]
-    assert prom_metrics.offloading_metrics[(0, LOAD_TIME)].increments == [1.5]
-    assert prom_metrics.offloading_metrics[(0, LOAD_SIZE)].observed == [16, 8]
-    assert prom_metrics.offloading_metrics[(0, STORE_BYTES)].increments == [3]
-    assert prom_metrics.offloading_metrics[(0, STORE_TIME)].increments == [0.3]
-    assert prom_metrics.offloading_metrics[(0, STORE_SIZE)].observed == [1, 2]
+    assert prom_metrics.offloading_metrics[(0, LOAD_BYTES, ())].increments == [24]
+    assert prom_metrics.offloading_metrics[(0, LOAD_TIME, ())].increments == [1.5]
+    assert prom_metrics.offloading_metrics[(0, LOAD_SIZE, ())].observed == [16, 8]
+    assert prom_metrics.offloading_metrics[(0, STORE_BYTES, ())].increments == [3]
+    assert prom_metrics.offloading_metrics[(0, STORE_TIME, ())].increments == [0.3]
+    assert prom_metrics.offloading_metrics[(0, STORE_SIZE, ())].observed == [1, 2]
 
     assert prom_metrics.counter_kv_bytes[(0, "CPU_to_GPU")].increments == [24]
     assert prom_metrics.counter_kv_transfer_time[(0, "CPU_to_GPU")].increments == [1.5]
@@ -396,7 +519,9 @@ def test_prom_metrics_observes_manager_gauge_and_histogram():
         ),
     }
     with patch.object(
-        CPUOffloadingSpec, "build_metric_definitions", return_value=metric_definitions
+        OffloadingSpecFactory,
+        "get_spec_cls",
+        return_value=_spec_cls_with_metric_definitions(metric_definitions),
     ):
         prom_metrics = OffloadPromMetrics(
             vllm_config=_FakeVllmConfig(store_threshold=0),  # type: ignore[arg-type]
@@ -416,20 +541,91 @@ def test_prom_metrics_observes_manager_gauge_and_histogram():
                 LOOKUP_LATENCY: _MetricType.HISTOGRAM,
             },
             _StatsKey.DATA: {
-                PENDING_STORES: 5,
-                LOOKUP_LATENCY: [0.2, 0.4],
+                PENDING_STORES: {(): 5},
+                LOOKUP_LATENCY: {(): [0.2, 0.4]},
             },
         }
     )
 
-    gauge = prom_metrics.offloading_metrics[(0, PENDING_STORES)]
-    histogram = prom_metrics.offloading_metrics[(0, LOOKUP_LATENCY)]
+    gauge = prom_metrics.offloading_metrics[(0, PENDING_STORES, ())]
+    histogram = prom_metrics.offloading_metrics[(0, LOOKUP_LATENCY, ())]
     assert gauge.set_values == [5]
     assert histogram.observed == [0.2, 0.4]
     histogram_def = prom_metrics._offloading_metric_defs[LOOKUP_LATENCY]
     assert histogram_def.kwargs["buckets"] == (0.1, 1.0)
 
 
+def test_prom_metrics_lazily_observes_labeled_metric():
+    metric_definitions = {
+        MY_COUNTER: OffloadingCounterMetadata(
+            documentation="counter with a label",
+            labelnames=(MY_LABEL,),
+        ),
+    }
+    with patch.object(
+        OffloadingSpecFactory,
+        "get_spec_cls",
+        return_value=_spec_cls_with_metric_definitions(metric_definitions),
+    ):
+        prom_metrics = OffloadPromMetrics(
+            vllm_config=_FakeVllmConfig(store_threshold=0),  # type: ignore[arg-type]
+            metric_types={
+                Gauge: _FakeMetric,
+                Counter: _FakeMetric,
+                Histogram: _FakeMetric,
+            },
+            labelnames=["model_name", "engine"],
+            per_engine_labelvalues={0: ["model", "0"]},
+        )
+
+    assert (0, MY_COUNTER, ("a",)) not in prom_metrics.offloading_metrics
+
+    prom_metrics.observe(
+        {
+            _StatsKey.TYPES: {MY_COUNTER: _MetricType.COUNTER},
+            _StatsKey.DATA: {MY_COUNTER: {("a",): 7}},
+        }
+    )
+
+    counter = prom_metrics.offloading_metrics[(0, MY_COUNTER, ("a",))]
+    assert counter.increments == [7]
+    assert counter.labelvalues == ("model", "0", "a")
+    counter_def = prom_metrics._offloading_metric_defs[MY_COUNTER]
+    assert counter_def.kwargs["labelnames"] == ["model_name", "engine", MY_LABEL]
+
+
+def test_prom_metrics_rejects_wrong_label_count():
+    metric_definitions = {
+        MY_COUNTER: OffloadingCounterMetadata(
+            documentation="counter with a label",
+            labelnames=(MY_LABEL,),
+        ),
+    }
+    with patch.object(
+        OffloadingSpecFactory,
+        "get_spec_cls",
+        return_value=_spec_cls_with_metric_definitions(metric_definitions),
+    ):
+        prom_metrics = OffloadPromMetrics(
+            vllm_config=_FakeVllmConfig(store_threshold=0),  # type: ignore[arg-type]
+            metric_types={
+                Gauge: _FakeMetric,
+                Counter: _FakeMetric,
+                Histogram: _FakeMetric,
+            },
+            labelnames=["model_name", "engine"],
+            per_engine_labelvalues={0: ["model", "0"]},
+        )
+
+    with pytest.raises(AssertionError, match="expects 1 labels"):
+        prom_metrics.observe(
+            {
+                _StatsKey.TYPES: {MY_COUNTER: _MetricType.COUNTER},
+                _StatsKey.DATA: {MY_COUNTER: {("a", "extra"): 7}},
+            }
+        )
+
+
 def test_prom_metrics_uses_configured_manager_metrics():
     prom_metrics = OffloadPromMetrics(
         vllm_config=_FakeVllmConfig(store_threshold=0),  # type: ignore[arg-type]
@@ -458,9 +654,9 @@ def test_aggregate_into_empty_stats():
                 PENDING_STORES: _MetricType.GAUGE,
             },
             _StatsKey.DATA: {
-                LOAD_BYTES: 42,
-                LOAD_SIZE: [10, 20],
-                PENDING_STORES: 3,
+                LOAD_BYTES: {(): 42},
+                LOAD_SIZE: {(): [10, 20]},
+                PENDING_STORES: {(): 3},
             },
         },
     )
@@ -469,9 +665,9 @@ def test_aggregate_into_empty_stats():
 
     assert result is empty
     values = result.data[_StatsKey.DATA]
-    assert values[LOAD_BYTES] == 42
-    assert values[LOAD_SIZE] == [10, 20]
-    assert values[PENDING_STORES] == 3
+    assert _unlabeled(values, LOAD_BYTES) == 42
+    assert _unlabeled(values, LOAD_SIZE) == [10, 20]
+    assert _unlabeled(values, PENDING_STORES) == 3
 
 
 def test_prom_metrics_multi_engine_routing():
@@ -490,14 +686,13 @@ def test_prom_metrics_multi_engine_routing():
     prom_metrics.observe(
         {
             _StatsKey.TYPES: {LOAD_BYTES: _MetricType.COUNTER},
-            _StatsKey.DATA: {LOAD_BYTES: 100},
+            _StatsKey.DATA: {LOAD_BYTES: {(): 100}},
         },
         engine_idx=1,
     )
 
-    engine0 = prom_metrics.offloading_metrics[(0, LOAD_BYTES)]
-    engine1 = prom_metrics.offloading_metrics[(1, LOAD_BYTES)]
-    assert engine0.increments == []
+    assert (0, LOAD_BYTES, ()) not in prom_metrics.offloading_metrics
+    engine1 = prom_metrics.offloading_metrics[(1, LOAD_BYTES, ())]
     assert engine1.increments == [100]
 
 
@@ -518,6 +713,6 @@ def test_prom_metrics_rejects_undeclared_metric():
         prom_metrics.observe(
             {
                 _StatsKey.TYPES: {"unknown:metric": _MetricType.COUNTER},
-                _StatsKey.DATA: {"unknown:metric": 1},
+                _StatsKey.DATA: {"unknown:metric": {(): 1}},
             }
         )

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

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from collections.abc import Iterable
+from types import SimpleNamespace
 from unittest.mock import MagicMock
 
 import pytest
@@ -1278,11 +1279,11 @@ def test_reset_cache_finalizes_finished_request_with_pending_store(
     )
 
     finalized: list[str] = []
-    runner.manager.on_request_finished.side_effect = (
-        lambda req_context: finalized.append(req_context.req_id)
+    runner.manager.on_request_finished.side_effect = lambda req_context: (
+        finalized.append(req_context.req_id)
     )
-    runner.manager.prepare_store.side_effect = (
-        lambda keys, req_context: generate_store_output(keys)
+    runner.manager.prepare_store.side_effect = lambda keys, req_context: (
+        generate_store_output(keys)
     )
 
     # Decode a couple of blocks and keep every transfer in flight, so the
@@ -1314,6 +1315,100 @@ def test_reset_cache_finalizes_finished_request_with_pending_store(
     assert req_id not in cs._req_status
 
 
+def test_pending_transfer_defers_prefix_lookup():
+    """A request with an in-flight store must not issue a load on re-admission.
+
+    With async scheduling, a preempted request's store can be flushed by the
+    worker before the scheduler consumes its completion. If the request is
+    re-admitted in that window, the connector should defer it instead of
+    looking up offloaded blocks and later asserting when a load is queued while
+    the store job is still tracked.
+    """
+    scheduler = object.__new__(OffloadingConnectorScheduler)
+    scheduler.manager = MagicMock(spec=OffloadingManager)
+
+    request = SimpleNamespace(request_id="req-0")
+    group_state = SimpleNamespace(block_ids=[1, 2, 3])
+    req_status = SimpleNamespace(
+        group_states=[group_state],
+        transfer_jobs={123},
+    )
+    scheduler._req_status = {request.request_id: req_status}
+
+    matched_tokens, is_async = scheduler.get_num_new_matched_tokens(
+        request,
+        num_computed_tokens=0,
+    )
+
+    assert matched_tokens is None
+    assert is_async is False
+    assert group_state.block_ids == []
+    scheduler.manager.lookup.assert_not_called()
+
+
+def test_async_preempt_readmit_before_transfer_output_is_deferred(request_runner):
+    """A preempted request can be scheduled again before flush output is read.
+
+    EngineCore.step_with_batch_queue() may schedule a new batch while a prior
+    preemption batch is still queued. The store completion from jobs_to_flush is
+    only cleared when that queued output reaches update_from_output(), so the
+    re-admission path must defer while the scheduler still tracks the store.
+    """
+    block_size = 4
+    block_size_factor = 3
+    offloaded_block_size = block_size * block_size_factor
+
+    runner = request_runner(
+        block_size=block_size,
+        num_gpu_blocks=100,
+        async_scheduling=True,
+        block_size_factor=block_size_factor,
+    )
+    free_block_queue = runner.scheduler.kv_cache_manager.block_pool.free_block_queue
+    num_free_blocks_empty = free_block_queue.num_free_blocks
+
+    req_id = "0"
+    runner.new_request(token_ids=[0] * offloaded_block_size * 2)
+    runner.manager.prepare_store.side_effect = lambda keys, req_context: (
+        generate_store_output(keys)
+    )
+
+    runner.run(decoded_tokens=[0], complete_transfers=False)
+    runner.run(
+        decoded_tokens=[0] * (2 * offloaded_block_size - block_size),
+        complete_transfers=False,
+    )
+
+    req_status = runner.connector_scheduler._req_status[req_id]
+    pending_store_jobs = set(req_status.transfer_jobs)
+    assert pending_store_jobs
+    assert all(
+        runner.connector_scheduler._jobs[jid].is_store for jid in pending_store_jobs
+    )
+
+    free_block_queue.num_free_blocks = 0
+    preempt_output = runner.scheduler.schedule()
+    assert preempt_output.preempted_req_ids == {req_id}
+    assert preempt_output.kv_connector_metadata is not None
+    assert pending_store_jobs <= preempt_output.kv_connector_metadata.jobs_to_flush
+    assert req_status.transfer_jobs == pending_store_jobs
+
+    # Simulate the async batch-queue window: schedule again before the
+    # preemption batch's ModelRunnerOutput is consumed by update_from_output().
+    free_block_queue.num_free_blocks = num_free_blocks_empty
+    assert runner.scheduler.reset_prefix_cache()
+    runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: len(
+        key
+    )
+
+    readmit_output = runner.scheduler.schedule()
+
+    assert readmit_output.num_scheduled_tokens == {}
+    assert readmit_output.kv_connector_metadata is not None
+    assert readmit_output.kv_connector_metadata.load_jobs == {}
+    assert req_status.transfer_jobs == pending_store_jobs
+
+
 @pytest.mark.parametrize("async_scheduling", [True, False])
 def test_swa_alignment_skip(request_runner, async_scheduling: bool):
     """SWA blocks unreachable by the load path are skipped during store.

tests/v1/kv_connector/unit/test_moriio_connector.py

@@ -36,13 +36,33 @@ from vllm.utils.network_utils import (
     get_ip,
     make_zmq_path,
 )
-from vllm.v1.kv_cache_interface import KVCacheConfig
+from vllm.v1.kv_cache_interface import (
+    FullAttentionSpec,
+    KVCacheConfig,
+    KVCacheGroupSpec,
+    KVCacheTensor,
+)
 
 from .utils import create_request, create_scheduler
 
 
 def _make_test_kv_cache_config() -> KVCacheConfig:
-    return KVCacheConfig(num_blocks=0, kv_cache_tensors=[], kv_cache_groups=[])
+    layer_names = ["layer0", "layer1", "layer2"]
+    return KVCacheConfig(
+        num_blocks=2,
+        kv_cache_tensors=[KVCacheTensor(size=0, shared_by=layer_names)],
+        kv_cache_groups=[
+            KVCacheGroupSpec(
+                layer_names=layer_names,
+                kv_cache_spec=FullAttentionSpec(
+                    block_size=16,
+                    num_kv_heads=4,
+                    head_size=64,
+                    dtype=torch.float16,
+                ),
+            )
+        ],
+    )
 
 
 aiter_available = importlib.util.find_spec("aiter") is not None
@@ -175,9 +195,18 @@ class FakeMoRIIOConnectorWorker(MoRIIOConnectorWorker):
     REMOTE_ENGINE_ID = "remote_engine"
 
     def __init__(
-        self, *args, hand_shake_latency: float = 1.8, kv_cache_layout="HND", **kwargs
+        self,
+        vllm_config,
+        engine_id,
+        *args,
+        hand_shake_latency: float = 1.8,
+        kv_cache_layout="HND",
+        kv_cache_config=None,
+        **kwargs,
     ):
-        super().__init__(*args, **kwargs)
+        super().__init__(
+            vllm_config, engine_id, kv_cache_config or _make_test_kv_cache_config()
+        )
 
 
 def create_vllm_config(

tests/v1/kv_connector/unit/test_moriio_kv_layout.py

@@ -0,0 +1,228 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import importlib.util
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm.platforms import current_platform
+from vllm.v1.kv_cache_interface import FullAttentionSpec, MLAAttentionSpec
+
+aiter_available = importlib.util.find_spec("aiter") is not None
+mori_available = importlib.util.find_spec("mori") is not None
+
+if not (current_platform.is_rocm() and mori_available):
+    pytest.skip(
+        "MoRIIOs are only available on ROCm with mori package installed",
+        allow_module_level=True,
+    )
+
+moriio_layout = importlib.import_module(
+    "vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_layout"
+)
+
+
+def _full_spec(block_size: int = 4) -> FullAttentionSpec:
+    return FullAttentionSpec(
+        block_size=block_size,
+        num_kv_heads=2,
+        head_size=3,
+        dtype=torch.bfloat16,
+    )
+
+
+def _mla_spec(block_size: int = 4) -> MLAAttentionSpec:
+    return MLAAttentionSpec(
+        block_size=block_size,
+        num_kv_heads=1,
+        head_size=3,
+        dtype=torch.bfloat16,
+    )
+
+
+def _worker(
+    kv_caches: dict[str, torch.Tensor],
+    layer_to_spec: dict[str, object],
+    num_blocks: int = 8,
+) -> SimpleNamespace:
+    return SimpleNamespace(
+        kv_caches=kv_caches,
+        layer_to_spec=layer_to_spec,
+        num_blocks=num_blocks,
+        block_size=4,
+    )
+
+
+def _remote_meta(num_blocks: int = 16) -> SimpleNamespace:
+    return SimpleNamespace(num_blocks=num_blocks)
+
+
+def test_separated_kv_layout_uses_kv_axis_zero_and_block_axis_one():
+    cache = torch.empty((2, 8, 4, 2, 3), dtype=torch.bfloat16)
+    worker = _worker({"layer": cache}, {"layer": _full_spec()})
+
+    geometry = moriio_layout.get_layer_transfer_geometry(
+        "layer", cache, worker.layer_to_spec, remote_num_blocks=16
+    )
+    assert geometry.block_stride == 24
+    assert geometry.local_kv_stride == 192
+    assert geometry.remote_kv_stride == 384
+    assert geometry.split_kv_regions
+
+    assert moriio_layout.compute_block_transfer_offsets(
+        "layer", cache, worker.layer_to_spec, [1, 3], [4, 5], _remote_meta().num_blocks
+    ) == ([48, 144, 432, 528], [192, 240, 960, 1008], [48, 48, 48, 48])
+
+
+def test_interleaved_kv_layout_uses_block_axis_zero_and_kv_axis_one():
+    cache = torch.empty((8, 2, 4, 2, 3), dtype=torch.bfloat16)
+    worker = _worker({"layer": cache}, {"layer": _full_spec()})
+
+    geometry = moriio_layout.get_layer_transfer_geometry(
+        "layer", cache, worker.layer_to_spec, remote_num_blocks=16
+    )
+    assert geometry.block_stride == 48
+    assert geometry.local_kv_stride == 24
+    assert geometry.remote_kv_stride == 24
+    assert not geometry.split_kv_regions
+
+    assert moriio_layout.compute_block_transfer_offsets(
+        "layer", cache, worker.layer_to_spec, [1, 3], [4, 5], _remote_meta().num_blocks
+    ) == ([96, 288], [384, 480], [96, 96])
+
+
+def test_mla_key_only_layout_transfers_one_slab_per_block():
+    cache = torch.empty((8, 4, 3), dtype=torch.bfloat16)
+    worker = _worker({"layer": cache}, {"layer": _mla_spec()})
+
+    geometry = moriio_layout.get_layer_transfer_geometry(
+        "layer", cache, worker.layer_to_spec, remote_num_blocks=16
+    )
+    assert geometry.block_stride == 12
+    assert geometry.local_kv_stride is None
+    assert geometry.remote_kv_stride is None
+    assert geometry.transfers_per_block == 1
+
+    assert moriio_layout.compute_block_transfer_offsets(
+        "layer", cache, worker.layer_to_spec, [1, 3], [4, 5], _remote_meta().num_blocks
+    ) == ([24, 72], [96, 120], [24, 24])
+
+
+def test_mixed_layers_compute_distinct_offsets_per_layer():
+    kv_caches = {
+        "separated": torch.empty((2, 8, 4, 2, 3), dtype=torch.bfloat16),
+        "interleaved": torch.empty((8, 2, 4, 2, 3), dtype=torch.bfloat16),
+        "indexer": torch.empty((8, 4, 3), dtype=torch.bfloat16),
+    }
+    worker = _worker(
+        kv_caches,
+        {
+            "separated": _full_spec(),
+            "interleaved": _full_spec(),
+            "indexer": _mla_spec(),
+        },
+    )
+
+    separated = moriio_layout.compute_block_transfer_offsets(
+        "separated",
+        kv_caches["separated"],
+        worker.layer_to_spec,
+        [1, 3],
+        [4, 5],
+        _remote_meta().num_blocks,
+    )
+    interleaved = moriio_layout.compute_block_transfer_offsets(
+        "interleaved",
+        kv_caches["interleaved"],
+        worker.layer_to_spec,
+        [1, 3],
+        [4, 5],
+        _remote_meta().num_blocks,
+    )
+    indexer = moriio_layout.compute_block_transfer_offsets(
+        "indexer",
+        kv_caches["indexer"],
+        worker.layer_to_spec,
+        [1, 3],
+        [4, 5],
+        _remote_meta().num_blocks,
+    )
+
+    assert separated != interleaved
+    assert separated != indexer
+    assert interleaved != indexer
+
+
+def test_block_id_length_mismatch_raises_value_error():
+    cache = torch.empty((8, 2, 4, 2, 3), dtype=torch.bfloat16)
+    worker = _worker({"layer": cache}, {"layer": _full_spec()})
+
+    with pytest.raises(ValueError, match="must have the same length"):
+        moriio_layout.compute_block_transfer_offsets(
+            "layer", cache, worker.layer_to_spec, [1, 3], [4], _remote_meta().num_blocks
+        )
+
+
+def test_registration_regions_do_not_split_interleaved_or_mla_cache():
+    separated = torch.empty((2, 8, 4, 2, 3), dtype=torch.bfloat16)
+    interleaved = torch.empty((8, 2, 4, 2, 3), dtype=torch.bfloat16)
+    indexer = torch.empty((8, 4, 3), dtype=torch.bfloat16)
+    worker = _worker(
+        {
+            "separated": separated,
+            "interleaved": interleaved,
+            "indexer": indexer,
+        },
+        {
+            "separated": _full_spec(),
+            "interleaved": _full_spec(),
+            "indexer": _mla_spec(),
+        },
+    )
+
+    separated_regions = moriio_layout.iter_layer_registration_regions(
+        "separated", separated, worker.layer_to_spec
+    )
+    interleaved_regions = moriio_layout.iter_layer_registration_regions(
+        "interleaved", interleaved, worker.layer_to_spec
+    )
+    indexer_regions = moriio_layout.iter_layer_registration_regions(
+        "indexer", indexer, worker.layer_to_spec
+    )
+
+    assert [region[0].data_ptr() for region in separated_regions] == [
+        separated[0].data_ptr(),
+        separated[1].data_ptr(),
+    ]
+    assert separated_regions[0][1] == 8 * 48
+    assert separated_regions[1][1] == 8 * 48
+
+    assert len(interleaved_regions) == 1
+    assert interleaved_regions[0][0].data_ptr() == interleaved.data_ptr()
+    assert interleaved_regions[0][1] == 8 * 2 * 48
+
+    assert len(indexer_regions) == 1
+    assert indexer_regions[0][0].data_ptr() == indexer.data_ptr()
+    assert indexer_regions[0][1] == 8 * 24
+
+
+def test_registration_regions_use_layer_num_blocks():
+    cache = torch.empty((4, 2, 4, 2, 3), dtype=torch.bfloat16)
+    worker = _worker({"layer": cache}, {"layer": _full_spec()}, num_blocks=8)
+
+    regions = moriio_layout.iter_layer_registration_regions(
+        "layer", cache, worker.layer_to_spec
+    )
+
+    assert len(regions) == 1
+    assert regions[0][1] == 4 * 2 * 48
+
+
+def test_unsupported_shape_raises_value_error():
+    cache = torch.empty((8, 4, 2, 3), dtype=torch.bfloat16)
+    worker = _worker({"layer": cache}, {"layer": _full_spec()})
+
+    with pytest.raises(ValueError, match="Unsupported MoRIIO K/V cache shape"):
+        moriio_layout.get_layer_transfer_geometry("layer", cache, worker.layer_to_spec)

tests/v1/spec_decode/test_speculators_correctness.py

@@ -53,9 +53,39 @@ PEAGLE_CONFIG = SpeculatorTestConfig(
     parallel_drafting=True,
 )
 
+QWEN3_EAGLE3_CONFIG = SpeculatorTestConfig(
+    model_path=(
+        "inference-optimization/"
+        "Qwen3-8B-from-Qwen3-8B_regen-speculators.eagle3-qwen3arch-ckpt1"
+    ),
+    method="eagle3",
+    display_name="Qwen3 Eagle3",
+    expected_gsm8k_accuracy=0.88,
+    accuracy_rtol=0.05,
+    expected_acceptance_len=2.67,
+    acceptance_len_rtol=0.10,
+    expected_per_pos_acceptance_rates=(0.76, 0.55, 0.36),
+    per_pos_rtol=0.10,
+)
+
+QWEN3_PEAGLE_CONFIG = SpeculatorTestConfig(
+    model_path="inference-optimization/Qwen3-8B-speculators.peagle-qwen3arch-ckpt4",
+    method="eagle3",
+    display_name="Qwen3 PEagle",
+    expected_gsm8k_accuracy=0.88,
+    accuracy_rtol=0.05,
+    expected_acceptance_len=3.42,
+    acceptance_len_rtol=0.15,
+    expected_per_pos_acceptance_rates=(0.78, 0.59, 0.43, 0.29, 0.18, 0.10, 0.05),
+    per_pos_rtol=0.10,
+    parallel_drafting=True,
+)
+
 SPECULATOR_CONFIGS = [
     pytest.param(DFLASH_CONFIG, id="dflash"),
     pytest.param(PEAGLE_CONFIG, id="peagle"),
+    pytest.param(QWEN3_EAGLE3_CONFIG, id="qwen3arch_eagle3"),
+    pytest.param(QWEN3_PEAGLE_CONFIG, id="qwen3arch_peagle"),
 ]
 
 
@@ -176,6 +206,7 @@ def test_speculators_correctness(monkeypatch, config):
 
     results = evaluate_gsm8k_offline(spec_llm)
     accuracy = results["accuracy"]
+    print(f"GSM8K Accuracy: {accuracy:.4f}")
     accuracy_threshold = config.expected_gsm8k_accuracy * (1 - config.accuracy_rtol)
     assert accuracy >= accuracy_threshold, (
         f"Expected GSM8K accuracy >= {accuracy_threshold:.3f}, got {accuracy:.3f}"

tests/v1/worker/test_attn_utils.py

@@ -0,0 +1,242 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import torch
+
+from vllm.v1.kv_cache_interface import FullAttentionSpec, KVQuantMode
+from vllm.v1.worker.gpu.attn_utils import _reshape_kv_cache
+from vllm.v1.worker.utils import AttentionGroup
+
+
+class FakeFlashAttentionBackend:
+    @staticmethod
+    def get_kv_cache_shape(
+        num_blocks: int,
+        block_size: int,
+        num_kv_heads: int,
+        head_size: int,
+        cache_dtype_str: str = "auto",
+    ) -> tuple[int, ...]:
+        return (num_blocks, 2, block_size, num_kv_heads, head_size)
+
+    @staticmethod
+    def get_kv_cache_stride_order(
+        include_num_layers_dimension: bool = False,
+    ) -> tuple[int, ...]:
+        assert not include_num_layers_dimension
+        return (0, 1, 2, 3, 4)
+
+
+class FakeHNDFlashAttentionBackend(FakeFlashAttentionBackend):
+    @staticmethod
+    def get_kv_cache_stride_order(
+        include_num_layers_dimension: bool = False,
+    ) -> tuple[int, ...]:
+        assert not include_num_layers_dimension
+        return (0, 1, 3, 2, 4)
+
+
+def test_reshape_padded_flash_attention_kv_cache_strides_by_page():
+    num_blocks = 3
+    spec = FullAttentionSpec(
+        block_size=16,
+        num_kv_heads=1,
+        head_size=2,
+        dtype=torch.float32,
+        page_size_padded=384,
+    )
+    assert spec.real_page_size_bytes == 256
+
+    raw_tensors = {
+        "layer": torch.zeros(spec.page_size_bytes * num_blocks, dtype=torch.int8)
+    }
+    attn_groups = [
+        AttentionGroup(
+            backend=FakeFlashAttentionBackend,
+            layer_names=["layer"],
+            kv_cache_spec=spec,
+            kv_cache_group_id=0,
+        )
+    ]
+
+    kv_cache = _reshape_kv_cache(
+        attn_groups,
+        raw_tensors,
+        "auto",
+        [spec.block_size],
+        {},
+    )["layer"]
+
+    assert kv_cache.shape == (num_blocks, 2, 16, 1, 2)
+    assert kv_cache.stride(0) == spec.page_size_bytes // 4
+    assert kv_cache.stride(1) == spec.real_page_size_bytes // 2 // 4
+    assert kv_cache[1, 0].storage_offset() == spec.page_size_bytes // 4
+    assert (
+        kv_cache[1, 1].storage_offset()
+        == (spec.page_size_bytes + spec.real_page_size_bytes // 2) // 4
+    )
+
+
+def test_reshape_padded_hnd_flash_attention_kv_cache_strides_by_page():
+    num_blocks = 3
+    spec = FullAttentionSpec(
+        block_size=16,
+        num_kv_heads=3,
+        head_size=2,
+        dtype=torch.float32,
+        page_size_padded=1024,
+    )
+    assert spec.real_page_size_bytes == 768
+
+    raw_tensors = {
+        "layer": torch.zeros(spec.page_size_bytes * num_blocks, dtype=torch.int8)
+    }
+    attn_groups = [
+        AttentionGroup(
+            backend=FakeHNDFlashAttentionBackend,
+            layer_names=["layer"],
+            kv_cache_spec=spec,
+            kv_cache_group_id=0,
+        )
+    ]
+
+    kv_cache = _reshape_kv_cache(
+        attn_groups,
+        raw_tensors,
+        "auto",
+        [spec.block_size],
+        {},
+    )["layer"]
+
+    assert kv_cache.shape == (num_blocks, 2, 16, 3, 2)
+    assert kv_cache.stride(0) == spec.page_size_bytes // 4
+    assert kv_cache.stride(1) == spec.real_page_size_bytes // 2 // 4
+    assert kv_cache.stride(2) == 2
+    assert kv_cache.stride(3) == spec.block_size * spec.head_size
+    assert kv_cache[1, 0].storage_offset() == spec.page_size_bytes // 4
+    assert (
+        kv_cache[1, 1].storage_offset()
+        == (spec.page_size_bytes + spec.real_page_size_bytes // 2) // 4
+    )
+    assert (
+        kv_cache[1, 1, 3, 2].storage_offset()
+        == (
+            spec.page_size_bytes
+            + spec.real_page_size_bytes // 2
+            + 3 * spec.head_size * 4
+            + 2 * spec.block_size * spec.head_size * 4
+        )
+        // 4
+    )
+
+
+class FakeDiffKVBackend:
+    @staticmethod
+    def get_kv_cache_shape(
+        num_blocks: int,
+        block_size: int,
+        num_kv_heads: int,
+        head_size: int,
+        cache_dtype_str: str = "auto",
+    ) -> tuple[int, ...]:
+        return (num_blocks, block_size, num_kv_heads, head_size * 2)
+
+    @staticmethod
+    def get_kv_cache_stride_order(
+        include_num_layers_dimension: bool = False,
+    ) -> tuple[int, ...]:
+        assert not include_num_layers_dimension
+        return (0, 1, 2, 3)
+
+
+def test_reshape_padded_diff_kv_cache_does_not_infer_kv_dim():
+    num_blocks = 3
+    spec = FullAttentionSpec(
+        block_size=16,
+        num_kv_heads=1,
+        head_size=2,
+        dtype=torch.float32,
+        page_size_padded=384,
+    )
+
+    raw_tensors = {
+        "layer": torch.zeros(spec.page_size_bytes * num_blocks, dtype=torch.int8)
+    }
+    attn_groups = [
+        AttentionGroup(
+            backend=FakeDiffKVBackend,
+            layer_names=["layer"],
+            kv_cache_spec=spec,
+            kv_cache_group_id=0,
+        )
+    ]
+
+    kv_cache = _reshape_kv_cache(
+        attn_groups,
+        raw_tensors,
+        "auto",
+        [spec.block_size],
+        {},
+    )["layer"]
+
+    assert kv_cache.shape == (num_blocks, 16, 1, 4)
+    assert kv_cache.stride(0) == spec.page_size_bytes // 4
+    assert kv_cache.stride(1) == 4
+
+
+class FakePerTokenScaleBackend:
+    @staticmethod
+    def get_kv_cache_shape(
+        num_blocks: int,
+        block_size: int,
+        num_kv_heads: int,
+        head_size: int,
+        cache_dtype_str: str = "auto",
+    ) -> tuple[int, ...]:
+        return (num_blocks, 2, block_size, num_kv_heads, head_size + 4)
+
+    @staticmethod
+    def get_kv_cache_stride_order(
+        include_num_layers_dimension: bool = False,
+    ) -> tuple[int, ...]:
+        assert not include_num_layers_dimension
+        return (0, 1, 2, 3, 4)
+
+
+def test_reshape_padded_quantized_kv_cache_preserves_scale_stride():
+    num_blocks = 3
+    spec = FullAttentionSpec(
+        block_size=16,
+        num_kv_heads=1,
+        head_size=4,
+        dtype=torch.int8,
+        kv_quant_mode=KVQuantMode.INT8_PER_TOKEN_HEAD,
+        page_size_padded=384,
+    )
+    assert spec.real_page_size_bytes == 128
+    assert spec.page_size_bytes == 384
+
+    raw_tensors = {
+        "layer": torch.zeros(spec.page_size_bytes * num_blocks, dtype=torch.int8)
+    }
+    attn_groups = [
+        AttentionGroup(
+            backend=FakePerTokenScaleBackend,
+            layer_names=["layer"],
+            kv_cache_spec=spec,
+            kv_cache_group_id=0,
+        )
+    ]
+
+    kv_cache = _reshape_kv_cache(
+        attn_groups,
+        raw_tensors,
+        "int8_per_token_head",
+        [spec.block_size],
+        {},
+    )["layer"]
+
+    assert kv_cache.shape == (num_blocks, 2, 16, 1, 8)
+    assert kv_cache.stride(0) == spec.page_size_bytes
+    assert kv_cache.stride(1) == 16 * 1 * 8
+    assert kv_cache[1, 1].storage_offset() == spec.page_size_bytes + 16 * 1 * 8

vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py

@@ -47,6 +47,14 @@ from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_engine import (
     MoRIIOWrapper,
     MoRIIOWriter,
 )
+from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_layout import (
+    LayerTransferGeometry,
+    build_layer_to_spec,
+    compute_block_transfer_offsets,
+    get_layer_transfer_geometry,
+    is_mla_cache_layer,
+    iter_layer_registration_regions,
+)
 from vllm.distributed.parallel_state import (
     get_tensor_model_parallel_world_size,
     get_tp_group,
@@ -71,6 +79,7 @@ if TYPE_CHECKING:
 
 logger = init_logger(__name__)
 
+
 try:
     from mori.io import (
         BackendType,
@@ -117,7 +126,9 @@ class MoRIIOConnector(KVConnectorBase_V1):
             self.connector_worker: MoRIIOConnectorWorker | None = None
         elif role == KVConnectorRole.WORKER:
             self.connector_scheduler = None
-            self.connector_worker = MoRIIOConnectorWorker(vllm_config, self.engine_id)
+            self.connector_worker = MoRIIOConnectorWorker(
+                vllm_config, self.engine_id, kv_cache_config
+            )
         logger.info(
             "Initialized MoRIIO Connector,engine_id:%s,role: %s",
             self.engine_id,
@@ -683,7 +694,12 @@ class MoRIIOConnectorScheduler:
 class MoRIIOConnectorWorker:
     """Implementation of Worker side methods"""
 
-    def __init__(self, vllm_config: VllmConfig, engine_id: str):
+    def __init__(
+        self,
+        vllm_config: VllmConfig,
+        engine_id: str,
+        kv_cache_config: "KVCacheConfig",
+    ):
         if not is_moriio_available():
             raise RuntimeError(
                 "MoRIIO is not available. Please ensure the 'mori' package "
@@ -707,6 +723,7 @@ class MoRIIOConnectorWorker:
         )
         self.kv_transfer_config = vllm_config.kv_transfer_config
         self.is_producer = self.kv_transfer_config.is_kv_producer
+        self.layer_to_spec = build_layer_to_spec(kv_cache_config)
 
         if self.is_producer:
             set_role(ROLE.PRODUCER)
@@ -809,6 +826,8 @@ class MoRIIOConnectorWorker:
         self.kv_cache_shape = None
         self.block_shape = None
         self.kv_element_size = 0
+        self.kv_cache_shapes: dict[str, torch.Size] = {}
+        self.block_lens: dict[str, int] = {}
 
         # Map of engine_id -> {agent_name0, agent_name1..}.
         self._remote_agents: dict[EngineId, set[str]] = {}
@@ -1218,51 +1237,86 @@ class MoRIIOConnectorWorker:
         all_done_future = self._handshake_initiation_executor.submit(wait_all_dp)
         all_done_future.add_done_callback(request_ready)
 
+    def _is_mla_cache_layer(self, layer_name: str) -> bool:
+        return is_mla_cache_layer(self.layer_to_spec, layer_name)
+
+    def _get_layer_transfer_geometry(
+        self, layer_name: str, remote_num_blocks: int | None = None
+    ) -> LayerTransferGeometry:
+        return get_layer_transfer_geometry(
+            layer_name,
+            self.kv_caches[layer_name],
+            self.layer_to_spec,
+            remote_num_blocks,
+        )
+
+    def _iter_layer_registration_regions(
+        self, layer_name: str
+    ) -> list[tuple[torch.Tensor, int]]:
+        return iter_layer_registration_regions(
+            layer_name,
+            self.kv_caches[layer_name],
+            self.layer_to_spec,
+        )
+
     def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
         """Register the KV Cache data in moriio."""
 
-        _, first_kv_cache = next(iter(kv_caches.items()))
+        self.kv_caches = kv_caches  # layer name to kv cache
+        self.kv_cache_shapes = {
+            layer_name: kv_cache.shape for layer_name, kv_cache in kv_caches.items()
+        }
+
+        first_layer_name, first_kv_cache = next(
+            (
+                (layer_name, kv_cache)
+                for layer_name, kv_cache in kv_caches.items()
+                if (
+                    not self._is_mla_cache_layer(layer_name)
+                    and len(kv_cache.shape) == 5
+                    and (kv_cache.shape[0] == 2 or kv_cache.shape[1] == 2)
+                )
+            ),
+            next(iter(kv_caches.items())),
+        )
         kv_elem_size = first_kv_cache.element_size()
 
-        use_mla = len(first_kv_cache.shape) == 3
-        assert use_mla == self.use_mla
+        use_mla = self._is_mla_cache_layer(first_layer_name)
+        first_geometry = self._get_layer_transfer_geometry(first_layer_name)
 
         if use_mla:
             # MLA case.
-            self.num_blocks = first_kv_cache.shape[0]
             block_rank = 2  # [block_size, latent_dim]
             block_shape = first_kv_cache.shape[-block_rank:]
-            block_size, kv_latent_dim = block_shape
-            self.slot_size_bytes = kv_elem_size * kv_latent_dim
         else:
-            # [2 (k and v), num_blocks, ...]
-            self.num_blocks = first_kv_cache.shape[1]
+            # [2, num_blocks, ...] or [num_blocks, 2, ...]
             block_rank = 3  # [block_size, kv_heads, head_dim]
             block_shape = first_kv_cache.shape[-block_rank:]
-            block_size, n_kv_heads, head_dim = block_shape[-3:]
-            # head size in bytes.
-            self.slot_size_bytes = (
-                kv_elem_size * n_kv_heads * head_dim
-            )  # 1 token 1 layer size , slot size
-        assert block_size == self.block_size
+        self.num_blocks = first_geometry.num_blocks
+        self.slot_size_bytes = first_geometry.slot_size_bytes
+        assert first_geometry.block_size == self.block_size
         # TODO(tms): self.block_len needs to be per-layer for sliding window,
         # hybrid attn, etc
         # block size in bytes
-        self.block_len = kv_elem_size * math.prod(block_shape)
+        self.block_len = first_geometry.block_len
         self.kv_cache_shape = first_kv_cache.shape
         self.block_shape = block_shape
         self.kv_element_size = kv_elem_size
 
         self.dst_num_blocks[self.engine_id] = self.num_blocks
-        self.kv_caches = kv_caches  # layer name to kv cache
         kv_caches_base_addr = []
         caches_data = []
 
-        for cache_or_caches in kv_caches.values():
-            cache_list = [cache_or_caches] if use_mla else cache_or_caches
-            for cache in cache_list:
+        for layer_name in kv_caches:
+            geometry = self._get_layer_transfer_geometry(layer_name)
+            if geometry.block_size != self.block_size:
+                raise ValueError(
+                    "MoRIIO KV cache block size mismatch for layer "
+                    f"{layer_name}: {geometry.block_size} != {self.block_size}"
+                )
+            self.block_lens[layer_name] = geometry.block_len
+            for cache, region_len in self._iter_layer_registration_regions(layer_name):
                 base_addr = cache.data_ptr()
-                region_len = self.num_blocks * self.block_len
                 caches_data.append((base_addr, region_len, cache.device.index, ""))
                 kv_caches_base_addr.append(base_addr)
 
@@ -1275,7 +1329,9 @@ class MoRIIOConnectorWorker:
                 moriio_mem_metadata
             )
 
-            self.local_kv_cache_size.append(cache.nelement() * cache.element_size())
+            self.local_kv_cache_size.append(
+                kv_cache.nelement() * kv_cache.element_size()
+            )
 
         self.kv_caches_base_addr[self.engine_id] = kv_caches_base_addr
         self.num_regions = len(caches_data)
@@ -1666,47 +1722,17 @@ class MoRIIOConnectorWorker:
         Returns:
             Tuple of (local_offsets, remote_offsets, transfer_sizes)
         """
-        assert self.kv_cache_shape is not None, "KV caches shape not initialized"
-        is_mla = len(self.kv_cache_shape) == 3
-        stride = self.kv_caches[layer_name].stride()
-        sz = self.kv_caches[layer_name].element_size()
-        if is_mla:
-            blknum, blksize, hs = self.kv_cache_shape
-            hn = 1
-            block_stride = stride[0]
-        else:
-            _, blknum, blksize, hn, hs = self.kv_cache_shape
-            local_ktov_stride = stride[0]
-            block_stride = stride[1]
-            remote_ktov_stride = block_stride * remote_moriio_meta.num_blocks
-
-        transfer_size_byte = blksize * hn * hs * sz
-        per_block = 1 if is_mla else 2
-        total = len(local_block_ids) * per_block
-        offset_local = [0] * total
-        offset_remote = [0] * total
-        sizes = [transfer_size_byte] * total
-
-        w = 0
-        for i, lb in enumerate(local_block_ids):
-            rb = remote_block_ids[i]
-            # K
-            offset_local[w] = sz * (lb * block_stride)
-            offset_remote[w] = sz * (rb * block_stride)
-            w += 1
-            if not is_mla:
-                # V
-                # Handle num_block variations originating from PD (different kv strides)
-                # TODO: address block_sz differences in heterogeneous TP scenarios
-                # In MLA, we don't need to consider these two cases.
-                offset_local[w] = sz * (1 * local_ktov_stride + lb * block_stride)
-                offset_remote[w] = sz * (1 * remote_ktov_stride + rb * block_stride)
-                w += 1
-
-        merged_l, merged_r, merged_s = self.merge_contiguous_blocks(
-            offset_local, offset_remote, sizes, assume_sorted=False
+        return compute_block_transfer_offsets(
+            layer_name=layer_name,
+            kv_cache=self.kv_caches[layer_name],
+            layer_to_spec=self.layer_to_spec,
+            local_block_ids=local_block_ids,
+            remote_block_ids=remote_block_ids,
+            remote_num_blocks=remote_moriio_meta.num_blocks,
+            merge_fn=lambda local, remote, sizes: self.merge_contiguous_blocks(
+                local, remote, sizes, assume_sorted=False
+            ),
         )
-        return merged_l, merged_r, merged_s
 
     def _read_blocks(
         self,
@@ -1724,15 +1750,13 @@ class MoRIIOConnectorWorker:
         dp0_engine_id = self.get_engine_name_with_dp(dst_engine_id, 0)
         sessions, remote_moriio_meta = self._get_built_session(dp0_engine_id)
 
-        first_layer = list(self.layer_name_to_local_kv_cache_metadata.keys())[0]
-        offs = self._compute_block_transfer_offsets(
-            first_layer, local_block_ids, remote_block_ids, remote_moriio_meta
-        )
-
         for layer_name in self.layer_name_to_local_kv_cache_metadata:
             sess_idx = list(self.layer_name_to_local_kv_cache_metadata.keys()).index(
                 layer_name
             )
+            offs = self._compute_block_transfer_offsets(
+                layer_name, local_block_ids, remote_block_ids, remote_moriio_meta
+            )
             # TODO : apply multi-session batch-read when moriio support it
             transfer_status = self.moriio_wrapper.read_remote_data(
                 offs[2], offs[0], offs[1], sessions[sess_idx]

vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_layout.py

@@ -0,0 +1,213 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from collections.abc import Callable, Mapping
+from typing import NamedTuple
+
+import torch
+
+from vllm.v1.kv_cache_interface import (
+    KVCacheConfig,
+    KVCacheSpec,
+    MLAAttentionSpec,
+    SlidingWindowMLASpec,
+    UniformTypeKVCacheSpecs,
+)
+
+
+class LayerTransferGeometry(NamedTuple):
+    num_blocks: int
+    block_size: int
+    block_len: int
+    slot_size_bytes: int
+    block_stride: int
+    local_kv_stride: int | None
+    remote_kv_stride: int | None
+    transfers_per_block: int
+    regions_per_block: int
+    split_kv_regions: bool
+
+
+def build_layer_to_spec(kv_cache_config: KVCacheConfig) -> dict[str, KVCacheSpec]:
+    layer_to_spec: dict[str, KVCacheSpec] = {}
+    for group in kv_cache_config.kv_cache_groups:
+        group_spec = group.kv_cache_spec
+        if isinstance(group_spec, UniformTypeKVCacheSpecs):
+            layer_to_spec.update(
+                {
+                    layer_name: group_spec.kv_cache_specs[layer_name]
+                    for layer_name in group.layer_names
+                }
+            )
+        else:
+            layer_to_spec.update(
+                {layer_name: group_spec for layer_name in group.layer_names}
+            )
+    return layer_to_spec
+
+
+def is_mla_cache_layer(
+    layer_to_spec: Mapping[str, KVCacheSpec], layer_name: str
+) -> bool:
+    try:
+        spec = layer_to_spec[layer_name]
+    except KeyError as e:
+        raise ValueError(f"Missing KV cache spec for layer {layer_name}") from e
+    return isinstance(spec, (MLAAttentionSpec, SlidingWindowMLASpec))
+
+
+def get_layer_transfer_geometry(
+    layer_name: str,
+    kv_cache: torch.Tensor,
+    layer_to_spec: Mapping[str, KVCacheSpec],
+    remote_num_blocks: int | None = None,
+) -> LayerTransferGeometry:
+    shape = kv_cache.shape
+    stride = kv_cache.stride()
+    element_size = kv_cache.element_size()
+    is_mla_cache = is_mla_cache_layer(layer_to_spec, layer_name)
+
+    if is_mla_cache and len(shape) == 3:
+        num_blocks, block_size, latent_dim = shape
+        slot_size_bytes = latent_dim * element_size
+        block_len = block_size * slot_size_bytes
+        return LayerTransferGeometry(
+            num_blocks=num_blocks,
+            block_size=block_size,
+            block_len=block_len,
+            slot_size_bytes=slot_size_bytes,
+            block_stride=stride[0],
+            local_kv_stride=None,
+            remote_kv_stride=None,
+            transfers_per_block=1,
+            regions_per_block=1,
+            split_kv_regions=False,
+        )
+
+    if not is_mla_cache and len(shape) == 5 and shape[0] == 2:
+        _, num_blocks, block_size, num_kv_heads, head_dim = shape
+        slot_size_bytes = num_kv_heads * head_dim * element_size
+        block_len = block_size * slot_size_bytes
+        remote_kv_stride = stride[1] * (remote_num_blocks or num_blocks)
+        return LayerTransferGeometry(
+            num_blocks=num_blocks,
+            block_size=block_size,
+            block_len=block_len,
+            slot_size_bytes=slot_size_bytes,
+            block_stride=stride[1],
+            local_kv_stride=stride[0],
+            remote_kv_stride=remote_kv_stride,
+            transfers_per_block=2,
+            regions_per_block=1,
+            split_kv_regions=True,
+        )
+
+    if not is_mla_cache and len(shape) == 5 and shape[1] == 2:
+        num_blocks, _, block_size, num_kv_heads, head_dim = shape
+        slot_size_bytes = num_kv_heads * head_dim * element_size
+        block_len = block_size * slot_size_bytes
+        return LayerTransferGeometry(
+            num_blocks=num_blocks,
+            block_size=block_size,
+            block_len=block_len,
+            slot_size_bytes=slot_size_bytes,
+            block_stride=stride[0],
+            local_kv_stride=stride[1],
+            remote_kv_stride=stride[1],
+            transfers_per_block=2,
+            regions_per_block=2,
+            split_kv_regions=False,
+        )
+
+    cache_kind = "MLA" if is_mla_cache else "K/V"
+    raise ValueError(
+        f"Unsupported MoRIIO {cache_kind} cache shape for layer "
+        f"{layer_name}: {tuple(shape)}"
+    )
+
+
+def iter_layer_registration_regions(
+    layer_name: str,
+    kv_cache: torch.Tensor,
+    layer_to_spec: Mapping[str, KVCacheSpec],
+) -> list[tuple[torch.Tensor, int]]:
+    geometry = get_layer_transfer_geometry(layer_name, kv_cache, layer_to_spec)
+    region_len = geometry.num_blocks * geometry.regions_per_block * geometry.block_len
+    if geometry.split_kv_regions:
+        return [(cache, region_len) for cache in kv_cache]
+    return [(kv_cache, region_len)]
+
+
+def merge_contiguous_offsets(
+    offsets_local: list[int],
+    offsets_remote: list[int],
+    sizes: list[int],
+) -> tuple[list[int], list[int], list[int]]:
+    if not offsets_local:
+        return [], [], []
+    if not (len(offsets_local) == len(offsets_remote) == len(sizes)):
+        raise ValueError("Input list lengths mismatch")
+
+    rows = sorted(zip(offsets_local, offsets_remote, sizes), key=lambda row: row[0])
+    merged: list[list[int]] = []
+    for local, remote, size in rows:
+        if (
+            merged
+            and local == merged[-1][0] + merged[-1][2]
+            and remote == merged[-1][1] + merged[-1][2]
+        ):
+            merged[-1][2] += size
+        else:
+            merged.append([local, remote, size])
+
+    return (
+        [row[0] for row in merged],
+        [row[1] for row in merged],
+        [row[2] for row in merged],
+    )
+
+
+def compute_block_transfer_offsets(
+    layer_name: str,
+    kv_cache: torch.Tensor,
+    layer_to_spec: Mapping[str, KVCacheSpec],
+    local_block_ids: list[int],
+    remote_block_ids: list[int],
+    remote_num_blocks: int,
+    merge_fn: Callable[
+        [list[int], list[int], list[int]], tuple[list[int], list[int], list[int]]
+    ] = merge_contiguous_offsets,
+) -> tuple[list[int], list[int], list[int]]:
+    if len(local_block_ids) != len(remote_block_ids):
+        raise ValueError(
+            "local_block_ids and remote_block_ids must have the same length: "
+            f"{len(local_block_ids)} != {len(remote_block_ids)}"
+        )
+    geometry = get_layer_transfer_geometry(
+        layer_name, kv_cache, layer_to_spec, remote_num_blocks
+    )
+    element_size = kv_cache.element_size()
+    transfer_size_byte = geometry.block_len
+    per_block = geometry.transfers_per_block
+    total = len(local_block_ids) * per_block
+    offset_local = [0] * total
+    offset_remote = [0] * total
+    sizes = [transfer_size_byte] * total
+
+    w = 0
+    for lb, rb in zip(local_block_ids, remote_block_ids):
+        offset_local[w] = element_size * (lb * geometry.block_stride)
+        offset_remote[w] = element_size * (rb * geometry.block_stride)
+        w += 1
+        if per_block == 2:
+            assert geometry.local_kv_stride is not None
+            assert geometry.remote_kv_stride is not None
+            offset_local[w] = element_size * (
+                geometry.local_kv_stride + lb * geometry.block_stride
+            )
+            offset_remote[w] = element_size * (
+                geometry.remote_kv_stride + rb * geometry.block_stride
+            )
+            w += 1
+
+    return merge_fn(offset_local, offset_remote, sizes)

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

@@ -112,7 +112,7 @@ class _StatsKey:
 
     # Maps metric name -> _MetricType value
     TYPES = "types"
-    # Maps metric name -> observed value (number or list)
+    # Maps metric name -> {label values tuple -> observed value (number or list)}
     DATA = "data"
 
 
@@ -125,15 +125,17 @@ class OffloadingConnectorStats(KVConnectorStats):
 
         {
             _StatsKey.TYPES: {name: _MetricType.*, ...},
-            _StatsKey.DATA:  {name: value, ...},
+            _StatsKey.DATA:  {name: {labelvalues: value, ...}, ...},
         }
 
     This structure is self-describing: it survives IPC serialization
     without needing the full ``OffloadingMetricMetadata`` objects on the
     receiving side.
 
-    Counter values are aggregated by summing, gauge values use the latest
-    snapshot, and histogram values are lists of observed samples.
+    Counter values are aggregated by summing per-label-tuple, gauge values
+    use the latest snapshot per-label-tuple, and histogram values are lists of
+    observed samples per-label-tuple. Unlabeled metrics use ``()`` as their
+    labelvalues tuple.
     """
 
     def __post_init__(self):
@@ -160,26 +162,32 @@ class OffloadingConnectorStats(KVConnectorStats):
         assert isinstance(other, OffloadingConnectorStats)
         other_types = other._types
         other_values = other._values
-        for key, value in other_values.items():
+        for key, other_label_values in other_values.items():
             type_str = other_types.get(key)
             if type_str is None:
                 raise AssertionError(f"Unknown offloading stats key: {key}")
             self._types.setdefault(key, type_str)
-            if type_str == _MetricType.HISTOGRAM:
-                assert isinstance(value, list)
-                if key not in self._values:
-                    self._values[key] = value
+            current_label_values = self._values.setdefault(key, {})
+            for labelvalues, value in other_label_values.items():
+                if type_str == _MetricType.HISTOGRAM:
+                    assert isinstance(value, list)
+                    if labelvalues not in current_label_values:
+                        current_label_values[labelvalues] = list(value)
+                    else:
+                        assert isinstance(current_label_values[labelvalues], list)
+                        current_label_values[labelvalues].extend(value)
+                elif type_str == _MetricType.COUNTER:
+                    assert isinstance(value, int | float)
+                    current_label_values[labelvalues] = (
+                        current_label_values.get(labelvalues, 0) + value
+                    )
+                elif type_str == _MetricType.GAUGE:
+                    assert isinstance(value, int | float)
+                    current_label_values[labelvalues] = value
                 else:
-                    assert isinstance(self._values[key], list)
-                    self._values[key].extend(value)
-            elif type_str == _MetricType.COUNTER:
-                assert isinstance(value, int | float)
-                self._values[key] = self._values.get(key, 0) + value
-            elif type_str == _MetricType.GAUGE:
-                assert isinstance(value, int | float)
-                self._values[key] = value
-            else:
-                raise AssertionError(f"Unknown metric type '{type_str}' for key: {key}")
+                    raise AssertionError(
+                        f"Unknown metric type '{type_str}' for key: {key}"
+                    )
         return self
 
     def reduce(self) -> dict[str, int | float]:
@@ -190,44 +198,62 @@ class OffloadingConnectorStats(KVConnectorStats):
         stats for the last time interval.
         """
         return_dict: dict[str, int | float] = {}
-        for key, value in self._values.items():
+        for key, label_value_map in self._values.items():
             type_str = self._types.get(key)
             if type_str is None:
                 raise AssertionError(f"Unknown offloading stats key: {key}")
-            if type_str == _MetricType.HISTOGRAM:
-                assert isinstance(value, list)
-                return_dict[f"{key}_count"] = len(value)
-                return_dict[f"{key}_sum"] = sum(value)
-            elif type_str in (_MetricType.COUNTER, _MetricType.GAUGE):
-                assert isinstance(value, int | float)
-                return_dict[key] = value
-            else:
-                raise AssertionError(f"Unknown metric type '{type_str}' for key: {key}")
+            for labelvalues, value in label_value_map.items():
+                key_with_labels = f"{key}:{labelvalues}" if labelvalues else key
+                if type_str == _MetricType.HISTOGRAM:
+                    assert isinstance(value, list)
+                    return_dict[f"{key_with_labels}_count"] = len(value)
+                    return_dict[f"{key_with_labels}_sum"] = sum(value)
+                elif type_str in (_MetricType.COUNTER, _MetricType.GAUGE):
+                    assert isinstance(value, int | float)
+                    return_dict[key_with_labels] = value
+                else:
+                    raise AssertionError(
+                        f"Unknown metric type '{type_str}' for key: {key}"
+                    )
         return return_dict
 
     def is_empty(self) -> bool:
         return not self.data.get(_StatsKey.DATA)
 
     def increase_counter(
-        self, counter_name: str, counter_increase_value: int | float
+        self,
+        counter_name: str,
+        counter_increase_value: int | float,
+        labelvalues: tuple[str, ...] = (),
     ) -> None:
         """Increase a counter on the stats payload."""
         self._types.setdefault(counter_name, _MetricType.COUNTER)
-        self._values[counter_name] = (
-            self._values.get(counter_name, 0) + counter_increase_value
+        counter_values = self._values.setdefault(counter_name, {})
+        counter_values[labelvalues] = (
+            counter_values.get(labelvalues, 0) + counter_increase_value
         )
 
-    def set_gauge(self, gauge_name: str, gauge_value: int | float) -> None:
+    def set_gauge(
+        self,
+        gauge_name: str,
+        gauge_value: int | float,
+        labelvalues: tuple[str, ...] = (),
+    ) -> None:
         """Set a gauge snapshot on the stats payload."""
         self._types.setdefault(gauge_name, _MetricType.GAUGE)
-        self._values[gauge_name] = gauge_value
+        gauge_values = self._values.setdefault(gauge_name, {})
+        gauge_values[labelvalues] = gauge_value
 
     def observe_histogram(
-        self, histogram_name: str, histogram_value: int | float
+        self,
+        histogram_name: str,
+        histogram_value: int | float,
+        labelvalues: tuple[str, ...] = (),
     ) -> None:
         """Record a histogram observation on the stats payload."""
         self._types.setdefault(histogram_name, _MetricType.HISTOGRAM)
-        self._values.setdefault(histogram_name, []).append(histogram_value)
+        histogram_values = self._values.setdefault(histogram_name, {})
+        histogram_values.setdefault(labelvalues, []).append(histogram_value)
 
 
 class OffloadPromMetrics(KVConnectorPromMetrics):
@@ -255,7 +281,10 @@ class OffloadPromMetrics(KVConnectorPromMetrics):
 
         self._observe_deprecated_metrics = issubclass(spec_cls, CPUOffloadingSpec)
         self._offloading_metric_defs: dict[str, PromMetricT] = {}
-        self.offloading_metrics: dict[tuple[int, str], PromMetricT] = {}
+        # (engine_idx, metric_name, labelvalues) -> metric with bound labels
+        self.offloading_metrics: dict[
+            tuple[int, str, tuple[str, ...]], PromMetricT
+        ] = {}
 
         self._counter_kv_bytes = self._counter_cls(
             name=_DEPRECATED_TOTAL_BYTES,
@@ -301,10 +330,6 @@ class OffloadPromMetrics(KVConnectorPromMetrics):
             self._offloading_metric_defs[metric_name] = self._create_metric(
                 metric_name, metadata
             )
-            for engine_idx, labelvalues in per_engine_labelvalues.items():
-                self.offloading_metrics[(engine_idx, metric_name)] = (
-                    self._offloading_metric_defs[metric_name].labels(*labelvalues)
-                )
 
     def _create_metric(
         self, metric_name: str, metadata: OffloadingMetricMetadata
@@ -312,7 +337,7 @@ class OffloadPromMetrics(KVConnectorPromMetrics):
         kwargs: dict[str, Any] = {
             "name": metric_name,
             "documentation": metadata.documentation,
-            "labelnames": self._labelnames,
+            "labelnames": self._labelnames + list(metadata.labelnames),
         }
         if isinstance(metadata, OffloadingCounterMetadata):
             metric_cls = self._counter_cls
@@ -326,11 +351,37 @@ class OffloadPromMetrics(KVConnectorPromMetrics):
             raise AssertionError(f"Unknown offloading metric metadata: {metadata}")
         return metric_cls(**kwargs)
 
+    def _get_prometheus_metric(
+        self,
+        metric_name: str,
+        labelvalues: tuple[str, ...],
+        engine_idx: int,
+    ) -> PromMetric:
+        metadata = self._offloading_metric_metadata[metric_name]
+        if len(labelvalues) != len(metadata.labelnames):
+            raise AssertionError(
+                f"Metric {metric_name} expects {len(metadata.labelnames)} labels, "
+                f"got {len(labelvalues)}"
+            )
+        key = (engine_idx, metric_name, labelvalues)
+        prom_metric = self.offloading_metrics.get(key)
+        if prom_metric is None:
+            engine_labelvalues = self.per_engine_labelvalues[engine_idx]
+            prom_metric = self._offloading_metric_defs[metric_name].labels(
+                *(engine_labelvalues + list(labelvalues))
+            )
+            self.offloading_metrics[key] = prom_metric
+        return prom_metric
+
     def _increase_counter(
-        self, metric_name: str, value: int | float, engine_idx: int
+        self,
+        metric_name: str,
+        value: int | float,
+        labelvalues: tuple[str, ...],
+        engine_idx: int,
     ) -> None:
-        self.offloading_metrics[(engine_idx, metric_name)].inc(value)
-        if not self._observe_deprecated_metrics:
+        self._get_prometheus_metric(metric_name, labelvalues, engine_idx).inc(value)
+        if labelvalues or not self._observe_deprecated_metrics:
             return
         # Keep deprecated CPU offload transfer metrics updated during the
         # transition to flat metric names.
@@ -343,15 +394,26 @@ class OffloadPromMetrics(KVConnectorPromMetrics):
         elif metric_name == _TransferMetricName.STORE_TIME:
             self.counter_kv_transfer_time[(engine_idx, _TransferType.STORE)].inc(value)
 
-    def _set_gauge(self, metric_name: str, value: int | float, engine_idx: int) -> None:
-        self.offloading_metrics[(engine_idx, metric_name)].set(value)
+    def _set_gauge(
+        self,
+        metric_name: str,
+        value: int | float,
+        labelvalues: tuple[str, ...],
+        engine_idx: int,
+    ) -> None:
+        self._get_prometheus_metric(metric_name, labelvalues, engine_idx).set(value)
 
     def _observe_histogram(
-        self, metric_name: str, value: list[int | float], engine_idx: int
+        self,
+        metric_name: str,
+        value: list[int | float],
+        labelvalues: tuple[str, ...],
+        engine_idx: int,
     ) -> None:
+        prom_metric = self._get_prometheus_metric(metric_name, labelvalues, engine_idx)
         for observation in value:
-            self.offloading_metrics[(engine_idx, metric_name)].observe(observation)
-            if not self._observe_deprecated_metrics:
+            prom_metric.observe(observation)
+            if labelvalues or not self._observe_deprecated_metrics:
                 continue
             # Keep deprecated CPU offload transfer metrics updated during the
             # transition to flat metric names.
@@ -368,20 +430,23 @@ class OffloadPromMetrics(KVConnectorPromMetrics):
         """Observe transfer statistics."""
         metric_types = transfer_stats_data.get(_StatsKey.TYPES, {})
         metric_data = transfer_stats_data.get(_StatsKey.DATA, {})
-        for key, value in metric_data.items():
+        for key, label_value_map in metric_data.items():
             type_str = metric_types.get(key)
             if type_str is None:
                 raise AssertionError(f"Unknown offloading stats key: {key}")
             assert key in self._offloading_metric_defs
-            if type_str == _MetricType.COUNTER:
-                assert isinstance(value, int | float)
-                self._increase_counter(key, value, engine_idx)
-            elif type_str == _MetricType.GAUGE:
-                assert isinstance(value, int | float)
-                self._set_gauge(key, value, engine_idx)
-            elif type_str == _MetricType.HISTOGRAM:
-                assert isinstance(value, list)
-                assert all(isinstance(v, int | float) for v in value)
-                self._observe_histogram(key, value, engine_idx)
-            else:
-                raise AssertionError(f"Unknown metric type '{type_str}' for key: {key}")
+            for labelvalues, value in label_value_map.items():
+                if type_str == _MetricType.COUNTER:
+                    assert isinstance(value, int | float)
+                    self._increase_counter(key, value, labelvalues, engine_idx)
+                elif type_str == _MetricType.GAUGE:
+                    assert isinstance(value, int | float)
+                    self._set_gauge(key, value, labelvalues, engine_idx)
+                elif type_str == _MetricType.HISTOGRAM:
+                    assert isinstance(value, list)
+                    assert all(isinstance(v, int | float) for v in value)
+                    self._observe_histogram(key, value, labelvalues, engine_idx)
+                else:
+                    raise AssertionError(
+                        f"Unknown metric type '{type_str}' for key: {key}"
+                    )

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

@@ -647,6 +647,13 @@ class OffloadingConnectorScheduler:
         for group_state in req_status.group_states:
             group_state.block_ids.clear()
 
+        if req_status.transfer_jobs:
+            logger.debug(
+                "Delaying request %s since it still has in-flight transfers",
+                request.request_id,
+            )
+            return None, False
+
         req_status.update_offload_keys()
         req_status.num_locally_computed_tokens = num_computed_tokens
 

vllm/entrypoints/serve/disagg/protocol.py

@@ -181,6 +181,7 @@ class GenerateResponseStreamChoice(BaseModel):
     logprobs: ChatCompletionLogProbs | None = None
     finish_reason: str | None = None
     token_ids: list[int] | None = None
+    routed_experts: str | None = None
 
 
 class GenerateStreamResponse(BaseModel):

vllm/entrypoints/serve/disagg/serving.py

@@ -400,6 +400,14 @@ class ServingTokens(OpenAIServing):
                     else:
                         logprobs = None
 
+                    routed_experts_b64 = None
+                    if output.routed_experts is not None:
+                        buf = io.BytesIO()
+                        np.save(buf, output.routed_experts)
+                        routed_experts_b64 = base64.b64encode(buf.getvalue()).decode(
+                            "ascii"
+                        )
+
                     chunk = GenerateStreamResponse(
                         request_id=request_id,
                         choices=[
@@ -408,6 +416,7 @@ class ServingTokens(OpenAIServing):
                                 logprobs=logprobs,
                                 finish_reason=finish_reason,
                                 token_ids=as_list(delta_token_ids),
+                                routed_experts=routed_experts_b64,
                             )
                         ],
                     )

vllm/forward_context.py

@@ -147,6 +147,12 @@ class ForwardContext:
 
     ubatch_slices: UBatchSlices | None = None
 
+    # Boolean mask over the (cudagraph-padded) token axis: True for trailing
+    # padding rows that are not real tokens. Consumers (e.g. DeepEP V2 dispatch)
+    # use it to skip a2a communication / MoE compute for padded tokens. None
+    # when the producer (currently the V2 model runner) does not set it.
+    is_padding: torch.Tensor | None = None
+
     # If True, bypass the compiled model call, e.g. by using .forward() directly
     skip_compiled: bool = False
 
@@ -211,6 +217,7 @@ def create_forward_context(
     slot_mapping: dict[str, torch.Tensor] | list[dict[str, torch.Tensor]] | None = None,
     additional_kwargs: dict[str, Any] | None = None,
     skip_compiled: bool = False,
+    is_padding: torch.Tensor | None = None,
 ):
     if vllm_config.compilation_config.fast_moe_cold_start:
         all_moe_layers = vllm_config.compilation_config.static_all_moe_layers
@@ -228,6 +235,7 @@ def create_forward_context(
         ubatch_slices=ubatch_slices,
         skip_compiled=skip_compiled,
         additional_kwargs=additional_kwargs or {},
+        is_padding=is_padding,
     )
 
 
@@ -257,6 +265,7 @@ def set_forward_context(
     ubatch_slices: UBatchSlices | None = None,
     slot_mapping: dict[str, torch.Tensor] | list[dict[str, torch.Tensor]] | None = None,
     skip_compiled: bool = False,
+    is_padding: torch.Tensor | None = None,
 ):
     """A context manager that stores the current forward context,
     can be attention metadata, etc.
@@ -316,6 +325,7 @@ def set_forward_context(
         slot_mapping,
         additional_kwargs,
         skip_compiled,
+        is_padding=is_padding,
     )
 
     try:

vllm/model_executor/layers/fused_moe/experts/gpt_oss_triton_kernels_moe.py

@@ -4,7 +4,6 @@
 import torch
 
 import vllm.model_executor.layers.fused_moe.modular_kernel as mk
-from vllm import _custom_ops as ops
 from vllm.logger import init_logger
 from vllm.model_executor.layers.fused_moe.activation import MoEActivation
 from vllm.model_executor.layers.fused_moe.config import (
@@ -220,6 +219,241 @@ def _patch_make_bitmatrix_metadata() -> None:
     _bm.make_bitmatrix_metadata = _make_bitmatrix_metadata_pow2_safe
 
 
+def _patch_legacy_routing_for_nonpow2_topk() -> None:
+    """Monkey-patch the legacy (v3.5.1) triton_kernels routing path to support
+    non-power-of-2 top_k (e.g. DeepSeek-V4 top_k=6).
+
+    The bundled ``_routing_compute_indx`` does ``tl.arange(0, N_EXPTS_ACT *
+    BLOCK_M)``, which fails to compile when ``N_EXPTS_ACT`` (top_k) is not a
+    power of 2 (6 * 32 = 192). This installs a pow2-safe variant that pads the
+    ``tl.arange`` to the next power of 2, strides by the real per-block size,
+    and masks the padded tail so it neither loads the next block's gates nor
+    writes any output. For power-of-2 top_k it is identical to the original.
+
+    A matching ``sort_tokens`` is installed that threads the padded size into
+    the patched kernel. Only needed on the legacy path; the v3.6+ SparseMatrix
+    path is handled by ``_patch_make_bitmatrix_metadata``.
+    """
+    import triton
+    import triton.language as tl
+
+    # Import via the `triton_kernels` alias (set up by has_triton_kernels) so
+    # we patch the SAME module object that `make_routing_data` consumes. The
+    # `vllm.third_party.triton_kernels.routing` path is a *different* module
+    # object under the import alias, so patching it would have no effect.
+    try:
+        import triton_kernels.routing as _routing
+        from triton_kernels.routing_details import _routing_compute as _rc
+    except ImportError:
+        return
+
+    _keyed_add = _rc._keyed_add
+    _expt_data_compute = _rc._expt_data_compute
+
+    @triton.jit
+    def _routing_compute_indx_pow2(
+        pid_m,
+        GatherIndx,
+        ScatterIndx,
+        GateScal,
+        ExptScal,
+        ExptIndx,
+        PartialOffs,
+        stride_pm,
+        stride_pn,
+        TokensStart,
+        n_tokens,
+        BLOCK_M: tl.constexpr,
+        N_EXPTS_ACT: tl.constexpr,
+        BLOCK_SIZE_PADDED: tl.constexpr,
+    ):
+        if isinstance(n_tokens, tl.tensor) and n_tokens.dtype.is_ptr():
+            n_tokens = tl.load(n_tokens)
+        n_gates = n_tokens * N_EXPTS_ACT
+        BLOCK_SIZE: tl.constexpr = N_EXPTS_ACT * BLOCK_M
+        tl.static_assert(BLOCK_SIZE_PADDED <= 32768)
+        local_offs = tl.arange(0, BLOCK_SIZE_PADDED)
+        offs = pid_m * BLOCK_SIZE + local_offs
+        expert = tl.load(
+            ExptIndx + offs,
+            mask=(local_offs < BLOCK_SIZE) & (offs < n_gates),
+            other=-1,
+        ).to(tl.uint32)
+        kv_pairs = ((expert << 16) | local_offs).to(tl.uint32)
+        kv_pairs = tl.sort(kv_pairs, 0)
+        expert = kv_pairs >> 16
+        offs = pid_m * BLOCK_SIZE + (kv_pairs & 0xFFFF)
+        mask = expert != 0xFFFF
+        gate_scal = tl.load(ExptScal + offs, mask=mask)
+        x = kv_pairs & 0xFFFF0000 | 0x00000001
+        run_lengths = tl.associative_scan(x, 0, _keyed_add)
+        exclusive_run_lengths = (run_lengths - 1) & 0xFFFF
+        gates = tl.load(PartialOffs + pid_m * stride_pm + expert * stride_pn, mask=mask)
+        gates += tl.load(TokensStart + expert, mask=mask)
+        gates += exclusive_run_lengths
+        tl.store(ScatterIndx + offs, gates, mask=mask)
+        tl.store(GatherIndx + gates, offs, mask=mask)
+        tl.store(GateScal + gates, gate_scal, mask=mask)
+
+    @triton.jit
+    def _combined_routing_compute_pow2(
+        GatherIndx,
+        ScatterIndx,
+        GateScal,
+        ExptScal,
+        ExptIndx,
+        PartialOffs,
+        stride_pm,
+        stride_pn,
+        TokensStart,
+        n_tokens,
+        BLOCK_M: tl.constexpr,
+        N_EXPTS_ACT: tl.constexpr,
+        Hist,
+        MDTileStarts,
+        tile_starts_stridem,
+        MDTileInfo,
+        tile_info_stridem,
+        first_tile_dim_log2,
+        SIZES: tl.constexpr,
+        BLOCK: tl.constexpr,
+        blocks2a,
+        BLOCK_SIZE_PADDED: tl.constexpr,
+    ):
+        pid = tl.program_id(0)
+        if pid < blocks2a:
+            _expt_data_compute(
+                Hist,
+                MDTileStarts,
+                tile_starts_stridem,
+                MDTileInfo,
+                tile_info_stridem,
+                first_tile_dim_log2,
+                SIZES,
+                BLOCK,
+            )
+        else:
+            pid -= blocks2a
+            _routing_compute_indx_pow2(
+                pid,
+                GatherIndx,
+                ScatterIndx,
+                GateScal,
+                ExptScal,
+                ExptIndx,
+                PartialOffs,
+                stride_pm,
+                stride_pn,
+                TokensStart,
+                n_tokens,
+                BLOCK_M,
+                N_EXPTS_ACT,
+                BLOCK_SIZE_PADDED,
+            )
+
+    def _sort_tokens_pow2(expt_scal, expt_indx, n_expts_tot, bitmatrix):
+        import torch
+
+        HIST_BLOCK_M = 32
+        INDX_OFFS_BLOCK_M = 512
+        MEMSET_BLOCK = 1024
+        cdiv = triton.cdiv
+        device = expt_scal.device
+        dtype = expt_scal.dtype
+        n_tokens_raw, _ = bitmatrix.shape
+        n_tokens_pad, n_expts_act = expt_scal.shape
+        n_gates_pad = n_tokens_pad * n_expts_act
+        # pad per-block gate count (HIST_BLOCK_M * top_k) up to a pow2.
+        block_size_padded = triton.next_power_of_2(HIST_BLOCK_M * n_expts_act)
+
+        hist, partial_hist = bitmatrix.sum(partials_block_size=HIST_BLOCK_M)
+        hist = hist[:n_expts_tot]
+        expt_offs = torch.empty(n_expts_tot, dtype=torch.int32, device=device)
+        combined_indx = torch.empty(n_gates_pad * 2, dtype=torch.int32, device=device)
+        topk_indx = combined_indx[:n_gates_pad]
+        gate_indx = combined_indx[n_gates_pad:]
+        gate_scal = torch.empty(n_gates_pad, dtype=dtype, device=device)
+
+        (
+            token_offs_combined,
+            token_offs_raw,
+            token_offs_pad,
+            block_pid_map,
+            blocks1a,
+            blocks2a,
+            MEMSET_BLOCK_A,
+            HIST2_BLOCK_M,
+            block_m_log2_start,
+            block_m_num,
+        ) = _routing._compute_expt_data_internal(hist, n_expts_tot, n_gates_pad)
+
+        blocks1b = cdiv(n_gates_pad * 2, MEMSET_BLOCK) + n_expts_tot + 1
+        blocks2b = cdiv(n_tokens_pad, HIST_BLOCK_M)
+
+        _rc._combined_routing_memset[(blocks1a + blocks1b,)](
+            combined_indx,
+            n_gates_pad * 2,
+            -1,
+            MEMSET_BLOCK,
+            hist,
+            expt_offs,
+            hist.shape[0],
+            n_expts_tot,
+            partial_hist,
+            partial_hist.shape[0],
+            partial_hist.stride(0),
+            partial_hist.stride(1),
+            token_offs_combined,
+            token_offs_combined.stride(0),
+            blocks1a,
+            block_pid_map,
+            block_m_log2_start,
+            SIZES=block_m_num,
+            BLOCK_A=MEMSET_BLOCK_A,
+            BLOCK_N=512,
+            BLOCK_M=INDX_OFFS_BLOCK_M,
+        )
+
+        indx_offs = partial_hist
+        _combined_routing_compute_pow2[(blocks2a + blocks2b,)](
+            topk_indx,
+            gate_indx,
+            gate_scal,
+            expt_scal,
+            expt_indx,
+            indx_offs,
+            indx_offs.stride(0),
+            indx_offs.stride(1),
+            expt_offs,
+            n_tokens_raw,
+            HIST_BLOCK_M,
+            n_expts_act,
+            hist,
+            token_offs_pad,
+            token_offs_pad.stride(0),
+            block_pid_map,
+            block_pid_map.stride(0),
+            block_m_log2_start,
+            block_m_num,
+            HIST2_BLOCK_M,
+            blocks2a,
+            block_size_padded,
+        )
+        return (
+            hist,
+            topk_indx,
+            gate_indx,
+            gate_scal,
+            token_offs_raw,
+            token_offs_pad,
+            block_pid_map,
+        )
+
+    # `routing_from_bitmatrix` looks up `sort_tokens` via the routing module
+    # global, so replacing it here redirects the legacy path to the pow2 kernel.
+    _routing.sort_tokens = _sort_tokens_pow2
+
+
 # Two API generations of triton_kernels are supported:
 #   - v3.5.1 (the version bundled with vLLM): exposes `routing()` and
 #     `routing_from_bitmatrix()` in triton_kernels.routing; the `Bitmatrix`
@@ -260,6 +494,9 @@ if has_triton_kernels():
             use_legacy_triton_kernels = True
         if not use_legacy_triton_kernels:
             _patch_make_bitmatrix_metadata()
+        else:
+            # Legacy routing fails to compile for non-pow2 top_k (DeepSeek-V4).
+            _patch_legacy_routing_for_nonpow2_topk()
     except (AttributeError, ImportError) as e:
         logger.error(
             "Failed to import Triton kernels. Please make sure your triton "
@@ -577,6 +814,80 @@ def make_routing_data(
     return routing_data, gather_indx, scatter_indx
 
 
+@triton.jit
+def _masked_topk_sum_kernel(
+    inp_ptr,  # (M, topk, K) contiguous
+    topk_ids_ptr,  # (M, topk) int: -1 marks an invalid / non-local slot
+    out_ptr,  # (M, K), same dtype as inp
+    K,
+    topk: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+):
+    pid_m = tl.program_id(0).to(tl.int64)
+    k = tl.program_id(1) * BLOCK_K + tl.arange(0, BLOCK_K)
+    k_mask = k < K
+    base = pid_m * topk
+    acc = tl.zeros((BLOCK_K,), dtype=tl.float32)
+    for j in tl.static_range(topk):
+        eid = tl.load(topk_ids_ptr + base + j)
+        # NOTE: This is NaN-safe because the invalid slots are skipped.
+        if eid >= 0:
+            x = tl.load(inp_ptr + (base + j) * K + k, mask=k_mask)
+            acc += x.to(tl.float32)
+    tl.store(out_ptr + pid_m * K + k, acc.to(out_ptr.dtype.element_ty), mask=k_mask)
+
+
+def masked_moe_sum(
+    intermediate: torch.Tensor,  # (M, topk, K)
+    topk_ids: torch.Tensor,  # (M, topk) int, -1 = invalid / non-local slot
+    output: torch.Tensor,  # (M, K)
+) -> None:
+    M, topk, K = intermediate.shape
+    BLOCK_K = 1024
+    grid = (M, triton.cdiv(K, BLOCK_K))
+    _masked_topk_sum_kernel[grid](
+        intermediate, topk_ids, output, K, topk=topk, BLOCK_K=BLOCK_K
+    )
+
+
+@triton.jit
+def _apply_expert_map_kernel(
+    topk_ids_ptr,  # [n] global expert IDs (-1 = invalid)
+    expert_map_ptr,  # [num_experts] global->local (-1 for non-local)
+    out_ptr,  # [n] int64 local expert IDs (-1 for invalid/non-local)
+    n_elements,
+    BLOCK: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    offs = pid * BLOCK + tl.arange(0, BLOCK)
+    mask = offs < n_elements
+    tid = tl.load(topk_ids_ptr + offs, mask=mask, other=-1)
+    # Gather expert_map[tid] for valid (tid >= 0); clamp the index so invalid
+    # rows don't read OOB, then select -1 for them. Matches
+    # torch.where(tid >= 0, expert_map[clamp(tid, 0)], -1) -- preserving -1 (a
+    # plain expert_map[-1] would wrap to a valid local id and misroute).
+    valid = tid >= 0
+    idx = tl.where(valid, tid, 0)
+    local = tl.load(expert_map_ptr + idx, mask=mask, other=-1)
+    out = tl.where(valid, local.to(tl.int64), -1)
+    tl.store(out_ptr + offs, out, mask=mask)
+
+
+def apply_expert_map(topk_ids: torch.Tensor, expert_map: torch.Tensor) -> torch.Tensor:
+    """Fused global->local expert-id mapping preserving -1.
+
+    Replaces ``torch.where(topk_ids >= 0, expert_map[topk_ids.clamp(min=0)], -1)``
+    with one kernel. Returns a NEW int64 tensor -- the caller keeps the original
+    ``topk_ids`` as ``global_topk_ids``, so this must not write in place.
+    """
+    out = torch.empty_like(topk_ids, dtype=torch.int64)
+    n = topk_ids.numel()
+    BLOCK = 1024
+    grid = (triton.cdiv(n, BLOCK),)
+    _apply_expert_map_kernel[grid](topk_ids, expert_map, out, n, BLOCK=BLOCK)
+    return out
+
+
 class BaseOAITritonExperts(mk.FusedMoEExpertsModular):
     @property
     def expects_unquantized_inputs(self) -> bool:
@@ -708,7 +1019,9 @@ class OAITritonExperts(BaseOAITritonExperts):
             self.quant_config: FusedMoEQuantConfig = FUSED_MOE_UNQUANTIZED_CONFIG
 
         if expert_map is not None:
-            topk_ids = expert_map[topk_ids]
+            # Preserve -1 (invalid / non-local slots, e.g. from EP dispatch):
+            # make_routing_data treats -1 as the skip sentinel.
+            topk_ids = apply_expert_map(topk_ids, expert_map)
 
         local_num_experts = w1.shape[0]
         if global_num_experts == -1:
@@ -780,9 +1093,6 @@ class UnfusedOAITritonExperts(LoRAExpertsMixin, BaseOAITritonExperts):
         output = (M, K)
         return (workspace1, workspace2, output)
 
-    def moe_sum(self, input: torch.Tensor, output: torch.Tensor):
-        ops.moe_sum(input, output)
-
     def activation(
         self,
         activation: MoEActivation,
@@ -853,7 +1163,9 @@ class UnfusedOAITritonExperts(LoRAExpertsMixin, BaseOAITritonExperts):
 
         global_topk_ids = topk_ids
         if expert_map is not None:
-            topk_ids = expert_map[topk_ids]
+            # Preserve -1 (invalid / non-local slots, e.g. from EP dispatch):
+            # make_routing_data treats -1 as the skip sentinel.
+            topk_ids = apply_expert_map(topk_ids, expert_map)
 
         local_num_experts = w1.shape[0]
         if global_num_experts == -1:
@@ -976,7 +1288,9 @@ class UnfusedOAITritonExperts(LoRAExpertsMixin, BaseOAITritonExperts):
                 top_k_num=topk,
             )
 
-        self.moe_sum(intermediate_cache3.view(-1, topk, K), output)
+        # matmul_ogs leaves invalid (-1 / non-local EP) slots unwritten.
+        # Reduce over topk skipping those slots.
+        masked_moe_sum(intermediate_cache3.view(-1, topk, K), topk_ids, output)
 
 
 class OAITritonMxfp4ExpertsMonolithic(mk.FusedMoEExpertsMonolithic):

vllm/model_executor/layers/fused_moe/modular_kernel.py

@@ -10,6 +10,7 @@ from typing import final
 import torch
 
 import vllm.envs as envs
+from vllm.forward_context import get_forward_context
 from vllm.logger import init_logger
 from vllm.model_executor.layers.fused_moe.activation import (
     MoEActivation,
@@ -1132,6 +1133,17 @@ class FusedMoEKernelModularImpl:
         The _prepare method is a wrapper around self.prepare_finalize.prepare
         that handles DBO and async.
         """
+        # Skip cudagraph/DP padding tokens uniformly across all a2a backends:
+        # forcing padded rows' expert ids to -1 makes every prepare_finalize drop
+        # them (not dispatched / not computed by the experts). The V2 model runner
+        # marks them in forward_context.is_padding; it is None for runners that do
+        # not populate it, leaving topk_ids unchanged.
+        is_padding = get_forward_context().is_padding
+        if is_padding is not None:
+            n = topk_ids.shape[0]
+            # TODO: Properly support DBO (padding lives at the batch tail).
+            topk_ids = torch.where(is_padding[:n].unsqueeze(1), -1, topk_ids)
+
         if not self.prepare_finalize.supports_async():
             # We shouldn't be running an a2a kernel that doesn't
             # support async prepare/finalize

vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_v2.py

@@ -6,12 +6,14 @@ import deep_ep
 import torch
 
 import vllm.model_executor.layers.fused_moe.modular_kernel as mk
+from vllm.forward_context import get_forward_context
 from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
 from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
     TopKWeightAndReduceContiguous,
     TopKWeightAndReduceDelegate,
 )
 from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input
+from vllm.triton_utils import tl, triton
 from vllm.utils.math_utils import round_up
 from vllm.v1.worker.ubatching import (
     dbo_current_ubatch_id,
@@ -116,6 +118,28 @@ class DeepEPV2PrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):
         do_expand = not self.use_cudagraph
         do_cpu_sync = not self.use_cudagraph
 
+        # In do_expand=False mode, the recv buffer is the worst case
+        # R * num_max_tokens_per_rank. Defaulting to the buffer's init value
+        # (= max_num_batched_tokens) makes the experts process ~R*8192 rows even
+        # for a handful of decode tokens. Bound it to the actual DP-padded batch
+        # size (uniform across ranks): max(num_tokens_across_dp).
+        #
+        # DeepEP JIT-compiles a separate dispatch kernel per distinct
+        # num_max_tokens_per_rank, so feeding it the raw per-step size would make
+        # it recompile for every batch size (a cicc storm that starves the GPU at
+        # high concurrency). Round up to a power of 2 instead: this bounds the
+        # set to ~log2(max_num_batched_tokens) values (compiled once, then
+        # cached) while staying small for decode (e.g. 1 token -> 1) and capped
+        # at the buffer's init capacity for prefill.
+        num_max_tokens_per_rank = None
+        if not do_expand:
+            dp_meta = get_forward_context().dp_metadata
+            if dp_meta is not None:
+                n = int(dp_meta.num_tokens_across_dp_cpu.max())
+            else:
+                n = tokens.shape[0]
+            num_max_tokens_per_rank = 1 << max(n - 1, 0).bit_length()
+
         (
             recv_x,
             recv_topk_idx,
@@ -127,6 +151,7 @@ class DeepEPV2PrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):
             topk_idx=rank_topk_ids,
             topk_weights=rank_topk_weights,
             num_experts=num_experts,
+            num_max_tokens_per_rank=num_max_tokens_per_rank,
             do_expand=do_expand,
             do_cpu_sync=do_cpu_sync,
             async_with_compute_stream=False,
@@ -196,17 +221,22 @@ class DeepEPV2PrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):
                 )
             recv_topk_idx = recv_topk_idx.unsqueeze(1)
         else:
-            # do_expand=False (decode/cudagraph mode): recv_topk_idx has
-            # LOCAL expert IDs (-1 for non-local and padding rows).
-            # Convert valid local IDs to global. Rows with -1 are
-            # skipped by expert kernels (TrtLLM tile-level skipping,
-            # DeepGemm is_computation_valid), so no need to zero
-            # hidden states, scales, or weights for padding rows.
-            valid_mask = recv_topk_idx >= 0
-            recv_topk_idx = torch.where(
-                valid_mask,
-                recv_topk_idx + self.rank_expert_offset,
+            # do_expand=False (decode/cudagraph mode): the dispatch only writes
+            # rows [0, num_recv_tokens); the rest of the worst-case-allocated
+            # buffer is left UNINITIALIZED. For valid rows, recv_topk_idx holds
+            # LOCAL expert IDs (-1 for non-local slots). Convert valid local IDs
+            # to global and force everything else to -1:
+            #   * non-local / out-of-range expert slots, and
+            #   * every row >= num_recv_tokens (uninitialized padding): its
+            #     stale contents can alias valid expert IDs and would otherwise
+            #     be treated as real routed tokens by experts that build routing
+            #     over *all* rows (e.g. triton MoE backend's make_routing_data),
+            #     polluting the per-expert token lists and corrupting real tokens.
+            recv_topk_idx = _globalize_recv_topk_idx(
                 recv_topk_idx,
+                psum_recv_per_rank,
+                self.rank_expert_offset,
+                self.num_experts,
             )
 
         # Reshape recv_topk_weights to match recv_topk_idx shape [N, 1]
@@ -392,3 +422,51 @@ class DeepEPV2PrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):
             weight_and_reduce_impl,
             False,
         )
+
+
+@triton.jit
+def _globalize_recv_topk_idx_kernel(
+    topk_idx_ptr,  # [N*topk] local expert IDs (-1 = non-local), modified in place
+    psum_ptr,  # [P] per-scaleup-rank recv prefix sum; num_recv = psum[P-1]
+    P,
+    rank_expert_offset,
+    num_experts,
+    n_elements,  # N * topk
+    topk: tl.constexpr,
+    BLOCK: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    offs = pid * BLOCK + tl.arange(0, BLOCK)
+    mask = offs < n_elements
+    # num_recv_tokens read on-device (no host sync) -> cudagraph-safe.
+    num_recv = tl.load(psum_ptr + P - 1)
+    val = tl.load(topk_idx_ptr + offs, mask=mask, other=-1)
+    g = val + rank_expert_offset
+    row = offs // topk
+    # Keep a slot iff: it is a local expert (val >= 0), its global id is in
+    # range, and its row is a real received token (< num_recv). Otherwise -1.
+    valid = (val >= 0) & (g < num_experts) & (row < num_recv)
+    tl.store(topk_idx_ptr + offs, tl.where(valid, g, -1), mask=mask)
+
+
+def _globalize_recv_topk_idx(
+    recv_topk_idx: torch.Tensor,  # [N, topk] local expert IDs, -1 = non-local
+    psum_recv_per_rank: torch.Tensor,
+    rank_expert_offset: int,
+    num_experts: int,
+) -> torch.Tensor:
+    N, topk = recv_topk_idx.shape
+    n = N * topk
+    BLOCK = 1024
+    grid = (triton.cdiv(n, BLOCK),)
+    _globalize_recv_topk_idx_kernel[grid](
+        recv_topk_idx,
+        psum_recv_per_rank,
+        psum_recv_per_rank.shape[0],
+        rank_expert_offset,
+        num_experts,
+        n,
+        topk=topk,
+        BLOCK=BLOCK,
+    )
+    return recv_topk_idx

vllm/model_executor/layers/minimax_rms_norm/rms_norm_tp.py

@@ -14,7 +14,7 @@ from vllm.distributed.parallel_state import (
 )
 from vllm.logger import init_logger
 from vllm.model_executor.custom_op import CustomOp
-from vllm.platforms import current_platform
+from vllm.triton_utils import HAS_TRITON, tl, triton
 from vllm.utils.torch_utils import direct_register_custom_op
 
 logger = init_logger(__name__)
@@ -40,8 +40,116 @@ def _all_reduce_variance(var: torch.Tensor) -> torch.Tensor:
     return tensor_model_parallel_all_reduce(var.flatten()).view_as(var)
 
 
-@torch.compile(backend=current_platform.simple_compile_backend, dynamic=True)
-def _minimax_qk_norm_fallback(
+@triton.jit
+def _minimax_qk_var_kernel(
+    qkv_ptr,  # [num_tokens, hidden], 16-bit activations
+    var_ptr,  # [num_tokens, 2], fp32
+    row_stride,  # element stride between tokens in qkv
+    q_size: tl.constexpr,  # constant per deployment -> loops unroll, mask elides
+    kv_size: tl.constexpr,
+    BLOCK: tl.constexpr,
+):
+    """TP-pre stage: per-token mean-of-squares for the q and k segments.
+
+    Accumulates in fp32 while reading the 16-bit qkv in place, so no fp32
+    copy of q/k is materialized. ``var[:, 0]`` is the q variance and
+    ``var[:, 1]`` the k variance; both are the local-shard means, ready for
+    the all-reduce that follows.
+    """
+    token = tl.program_id(0)
+    base = qkv_ptr + token * row_stride
+
+    q_acc = 0.0
+    for off in range(0, q_size, BLOCK):
+        idx = off + tl.arange(0, BLOCK)
+        mask = idx < q_size
+        x = tl.load(base + idx, mask=mask, other=0.0).to(tl.float32)
+        q_acc += tl.sum(x * x, axis=0)
+
+    k_acc = 0.0
+    for off in range(0, kv_size, BLOCK):
+        idx = off + tl.arange(0, BLOCK)
+        mask = idx < kv_size
+        x = tl.load(base + q_size + idx, mask=mask, other=0.0).to(tl.float32)
+        k_acc += tl.sum(x * x, axis=0)
+
+    tl.store(var_ptr + token * 2 + 0, q_acc / q_size)
+    tl.store(var_ptr + token * 2 + 1, k_acc / kv_size)
+
+
+@triton.jit
+def _minimax_rms_apply_kernel(
+    qkv_ptr,  # [num_tokens, hidden]
+    var_ptr,  # [num_tokens, 2], fp32, all-reduced sum of per-shard means
+    q_w_ptr,  # [q_size], q per-channel weight
+    k_w_ptr,  # [kv_size], k per-channel weight
+    q_out_ptr,  # [num_tokens, q_size], contiguous
+    k_out_ptr,  # [num_tokens, kv_size], contiguous
+    row_stride,  # element stride between tokens in qkv
+    q_size: tl.constexpr,  # constant per deployment -> loops unroll, mask elides
+    kv_size: tl.constexpr,
+    tp_world: tl.constexpr,  # folds the post-all-reduce /tp_world into rsqrt
+    eps: tl.constexpr,
+    BLOCK: tl.constexpr,
+):
+    """TP-post stage: ``x * rsqrt(var / tp_world + eps) * weight``.
+
+    A single program normalizes both the q and k segments of one token, so q
+    and k share one launch instead of two. The all-reduce yields the sum of
+    per-shard means, so the ``/ tp_world`` that recovers the global
+    mean-of-squares is folded into the ``rsqrt`` here rather than run as a
+    separate elementwise pass over the ``[num_tokens, 2]`` variance tensor.
+    """
+    token = tl.program_id(0)
+    base = qkv_ptr + token * row_stride
+
+    q_inv = tl.rsqrt(tl.load(var_ptr + token * 2 + 0) / tp_world + eps)
+    q_out_row = q_out_ptr + token * q_size
+    for off in range(0, q_size, BLOCK):
+        idx = off + tl.arange(0, BLOCK)
+        mask = idx < q_size
+        x = tl.load(base + idx, mask=mask, other=0.0).to(tl.float32)
+        w = tl.load(q_w_ptr + idx, mask=mask, other=0.0).to(tl.float32)
+        y = x * q_inv * w
+        tl.store(q_out_row + idx, y.to(q_out_ptr.dtype.element_ty), mask=mask)
+
+    k_inv = tl.rsqrt(tl.load(var_ptr + token * 2 + 1) / tp_world + eps)
+    k_out_row = k_out_ptr + token * kv_size
+    for off in range(0, kv_size, BLOCK):
+        idx = off + tl.arange(0, BLOCK)
+        mask = idx < kv_size
+        x = tl.load(base + q_size + idx, mask=mask, other=0.0).to(tl.float32)
+        w = tl.load(k_w_ptr + idx, mask=mask, other=0.0).to(tl.float32)
+        y = x * k_inv * w
+        tl.store(k_out_row + idx, y.to(k_out_ptr.dtype.element_ty), mask=mask)
+
+
+def _minimax_qk_norm_tp_eager(
+    qkv: torch.Tensor,
+    q_weight: torch.Tensor,
+    k_weight: torch.Tensor,
+    q_size: int,
+    kv_size: int,
+    tp_world: int,
+    eps: float,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Pure-torch reference path used when Triton is unavailable."""
+    q, k, _ = qkv.split([q_size, kv_size, kv_size], dim=-1)
+    orig_dtype = q.dtype
+    q = q.to(torch.float32)
+    k = k.to(torch.float32)
+    q_var = q.pow(2).mean(dim=-1, keepdim=True)
+    k_var = k.pow(2).mean(dim=-1, keepdim=True)
+
+    qk_var = torch.cat([q_var, k_var], dim=-1)
+    qk_var = _all_reduce_variance(qk_var) / tp_world
+    q_var, k_var = qk_var.chunk(2, dim=-1)
+    q = q * torch.rsqrt(q_var + eps) * q_weight
+    k = k * torch.rsqrt(k_var + eps) * k_weight
+    return q.to(orig_dtype), k.to(orig_dtype)
+
+
+def _minimax_qk_norm_tp_fallback(
     qkv: torch.Tensor,
     q_weight: torch.Tensor,
     k_weight: torch.Tensor,
@@ -51,19 +159,50 @@ def _minimax_qk_norm_fallback(
     tp_world: int,
     eps: float,
 ) -> tuple[torch.Tensor, torch.Tensor]:
-    q, k, _ = qkv.split([q_size, kv_size, kv_size], dim=-1)
-    orig_dtype = q.dtype
-    q = q.to(torch.float32)
-    k = k.to(torch.float32)
-    q_var = q.pow(2).mean(dim=-1, keepdim=True)
-    k_var = k.pow(2).mean(dim=-1, keepdim=True)
-    if tp_world > 1:
-        qk_var = torch.cat([q_var, k_var], dim=-1)
-        qk_var = _all_reduce_variance(qk_var) / tp_world
-        q_var, k_var = qk_var.chunk(2, dim=-1)
-    q = q * torch.rsqrt(q_var + eps) * q_weight
-    k = k * torch.rsqrt(k_var + eps) * k_weight
-    return q.to(orig_dtype), k.to(orig_dtype)
+    """All-reduce + QK RMSNorm without the Lamport fused kernel.
+
+    The all-reduce is a TP communication barrier and cannot live inside a
+    single kernel, so the eager-torch path is split into two Triton kernels
+    around it: a variance reduction before the all-reduce and a normalize
+    after. Compared to the ``torch.compile`` path this avoids materializing
+    fp32 copies of q/k and the ``cat``/``chunk`` temporaries.
+    """
+    if not HAS_TRITON:
+        return _minimax_qk_norm_tp_eager(
+            qkv, q_weight, k_weight, q_size, kv_size, tp_world, eps
+        )
+
+    num_tokens = qkv.shape[0]
+    row_stride = qkv.stride(0)
+    BLOCK = 1024
+    grid = (num_tokens,)
+
+    qk_var = torch.empty(num_tokens, 2, dtype=torch.float32, device=qkv.device)
+    _minimax_qk_var_kernel[grid](
+        qkv, qk_var, row_stride, q_size=q_size, kv_size=kv_size, BLOCK=BLOCK
+    )
+
+    # All-reduce sums the per-shard means; the /tp_world that turns this back
+    # into the global mean is folded into the apply kernel's rsqrt below.
+    qk_var = _all_reduce_variance(qk_var)
+
+    q_out = torch.empty(num_tokens, q_size, dtype=qkv.dtype, device=qkv.device)
+    k_out = torch.empty(num_tokens, kv_size, dtype=qkv.dtype, device=qkv.device)
+    _minimax_rms_apply_kernel[grid](
+        qkv,
+        qk_var,
+        q_weight,
+        k_weight,
+        q_out,
+        k_out,
+        row_stride,
+        q_size=q_size,
+        kv_size=kv_size,
+        tp_world=tp_world,
+        eps=eps,
+        BLOCK=BLOCK,
+    )
+    return q_out, k_out
 
 
 def _minimax_qk_norm_fusion(
@@ -96,7 +235,7 @@ def _minimax_qk_norm_fusion(
             tp_world,
             eps,
         )
-    return _minimax_qk_norm_fallback(
+    return _minimax_qk_norm_tp_fallback(
         qkv, q_weight, k_weight, q_size, kv_size, tp_rank, tp_world, eps
     )
 
@@ -231,10 +370,7 @@ class MiniMaxText01RMSNormTP(CustomOp):
         k = k.to(torch.float32)
         q_var = q.pow(2).mean(dim=-1, keepdim=True)
         k_var = k.pow(2).mean(dim=-1, keepdim=True)
-        if q_norm.tp_world > 1:
-            qk_var = torch.cat([q_var, k_var], dim=-1)
-            qk_var = _all_reduce_variance(qk_var) / q_norm.tp_world
-            q_var, k_var = qk_var.chunk(2, dim=-1)
+
         q = q * torch.rsqrt(q_var + q_norm.variance_epsilon) * q_norm.weight
         k = k * torch.rsqrt(k_var + k_norm.variance_epsilon) * k_norm.weight
         q = q.to(orig_dtype)
@@ -250,7 +386,14 @@ class MiniMaxText01RMSNormTP(CustomOp):
         kv_size: int,
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         assert qkv.ndim == 2
+
         assert q_norm.variance_epsilon == k_norm.variance_epsilon
+        # Case 0： tp_size=1
+        if get_tensor_model_parallel_world_size() == 1:
+            q, k, v = qkv.split([q_size, kv_size, kv_size], dim=-1)
+            q, k = MiniMaxText01RMSNormTP.forward_qk(q_norm, k_norm, q, k)
+            return q, k, v
+        # Case : tp_size>1
         q, k = torch.ops.vllm.minimax_qk_norm_fusion(
             qkv,
             q_norm.weight,

vllm/model_executor/models/qwen3_eagle3.py

@@ -0,0 +1,453 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from collections.abc import Iterable
+
+import torch
+import torch.nn as nn
+from transformers import Qwen3Config
+
+from vllm.compilation.decorators import support_torch_compile
+from vllm.config import VllmConfig, get_current_vllm_config
+from vllm.logger import init_logger
+from vllm.model_executor.layers.layernorm import RMSNorm
+from vllm.model_executor.layers.linear import QKVParallelLinear, ReplicatedLinear
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
+from vllm.model_executor.layers.vocab_parallel_embedding import (
+    ParallelLMHead,
+    VocabParallelEmbedding,
+)
+from vllm.model_executor.model_loader.weight_utils import (
+    default_weight_loader,
+    maybe_remap_kv_scale_name,
+)
+from vllm.model_executor.models.qwen3 import Qwen3DecoderLayer, Qwen3ForCausalLM
+from vllm.multimodal.inputs import NestedTensors
+
+from .utils import (
+    AutoWeightsLoader,
+    get_draft_quant_config,
+    maybe_prefix,
+    process_eagle_weight,
+)
+
+logger = init_logger(__name__)
+
+
+class Qwen3Eagle3DecoderLayer(Qwen3DecoderLayer):
+    def __init__(
+        self,
+        vllm_config: VllmConfig,
+        prefix: str = "",
+        config: Qwen3Config | None = None,
+        layer_idx: int = 0,
+    ) -> None:
+        config = config or vllm_config.model_config.hf_config
+        cache_config = vllm_config.cache_config
+        quant_config = get_draft_quant_config(vllm_config)
+
+        super().__init__(
+            config=config,
+            cache_config=cache_config,
+            quant_config=quant_config,
+            prefix=prefix,
+        )
+
+        # First layer uses 2*hidden_size (embeds + hidden_states concatenated)
+        # Subsequent layers use hidden_size (only hidden_states, no embeds)
+        qkv_input_size = 2 * self.hidden_size if layer_idx == 0 else self.hidden_size
+
+        # Parallel drafting checkpoints may have attention bias enabled
+        qkv_bias = getattr(config, "attention_bias", False)
+
+        # Override qkv_proj with correct input size and bias setting
+        self.self_attn.qkv_proj = QKVParallelLinear(
+            qkv_input_size,
+            self.self_attn.head_dim,
+            self.self_attn.total_num_heads,
+            self.self_attn.total_num_kv_heads,
+            bias=qkv_bias,
+            quant_config=quant_config,
+            prefix=maybe_prefix(prefix, "qkv_proj"),
+        )
+
+        self.hidden_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.layer_idx = layer_idx
+
+        if getattr(config, "norm_before_residual", False):
+            self._residual_norm = self._norm_before_residual
+        else:
+            self._residual_norm = self._norm_after_residual
+
+    def _norm_before_residual(
+        self, hidden_states: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        hidden_states = self.hidden_norm(hidden_states)
+        residual = hidden_states
+        return hidden_states, residual
+
+    def _norm_after_residual(
+        self, hidden_states: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        residual = hidden_states
+        hidden_states = self.hidden_norm(hidden_states)
+        return hidden_states, residual
+
+    def forward(
+        self,
+        positions: torch.Tensor,
+        embeds: torch.Tensor,
+        hidden_states: torch.Tensor,
+        residual: torch.Tensor | None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if self.layer_idx == 0:
+            # First layer: concatenate embeds with hidden_states
+            embeds = self.input_layernorm(embeds)
+            hidden_states, residual = self._residual_norm(hidden_states=hidden_states)
+            hidden_states = torch.cat([embeds, hidden_states], dim=-1)
+        else:
+            # Subsequent layers: process hidden_states and residuals only
+            hidden_states, residual = self.input_layernorm(hidden_states, residual)
+
+        # Self Attention
+        hidden_states = self.self_attn(
+            positions=positions,
+            hidden_states=hidden_states,
+        )
+
+        hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
+
+        # Fully Connected
+        hidden_states = self.mlp(hidden_states)
+
+        return hidden_states, residual
+
+
+@support_torch_compile(
+    dynamic_arg_dims={
+        "input_ids": 0,
+        "positions": -1,
+        "hidden_states": 0,
+        "input_embeds": 0,
+    }
+)
+class Qwen3Eagle3Model(nn.Module):
+    def __init__(
+        self,
+        *,
+        vllm_config: VllmConfig,
+        start_layer_id: int = 0,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.config = vllm_config.speculative_config.draft_model_config.hf_config
+        self.vocab_size = self.config.vocab_size
+
+        # Get drafter's quantization config
+        self.quant_config = get_draft_quant_config(vllm_config)
+
+        eagle_config = getattr(self.config, "eagle_config", None) or {}
+        if "use_aux_hidden_state" in eagle_config:
+            self.use_aux_hidden_state = eagle_config["use_aux_hidden_state"]
+        else:
+            self.use_aux_hidden_state = True
+        self.norm_before_fc = bool(
+            eagle_config.get(
+                "norm_before_fc", getattr(self.config, "norm_before_fc", False)
+            )
+        )
+        self.fc_input_size = self.config.hidden_size
+
+        current_vllm_config = get_current_vllm_config()
+
+        self.embed_tokens = VocabParallelEmbedding(
+            self.config.vocab_size,
+            self.config.hidden_size,
+            prefix=maybe_prefix(prefix, "embed_tokens"),
+        )
+
+        self.layers = nn.ModuleList(
+            [
+                Qwen3Eagle3DecoderLayer(
+                    current_vllm_config,
+                    prefix=maybe_prefix(prefix, f"layers.{layer_idx + start_layer_id}"),
+                    config=self.config,
+                    layer_idx=layer_idx,
+                )
+                for layer_idx in range(self.config.num_hidden_layers)
+            ]
+        )
+        if self.use_aux_hidden_state:
+            num_aux_features = getattr(self.config, "num_aux_layers", None)
+            if num_aux_features is None:
+                num_aux_features = getattr(self.config, "num_aux_hidden_states", None)
+            if num_aux_features is None:
+                aux_ids = getattr(
+                    self.config, "eagle_aux_hidden_state_layer_ids", None
+                ) or eagle_config.get("eagle_aux_hidden_state_layer_ids")
+                num_aux_features = len(aux_ids) if aux_ids else 3
+            self.num_aux_layers = num_aux_features
+            target_hidden_size = getattr(
+                self.config, "target_hidden_size", self.config.hidden_size
+            )
+            self.fc_input_size = target_hidden_size * num_aux_features
+            if self.norm_before_fc:
+                self.input_norm = RMSNorm(
+                    self.fc_input_size,
+                    eps=self.config.rms_norm_eps,
+                )
+            else:
+                self.input_norm = None
+
+            use_fc_norm = getattr(self.config, "fc_norm", False)
+            if use_fc_norm:
+                self.fc_norm = nn.ModuleList(
+                    [
+                        RMSNorm(target_hidden_size, eps=self.config.rms_norm_eps)
+                        for _ in range(num_aux_features)
+                    ]
+                )
+            else:
+                self.fc_norm = None
+
+            self.fc = ReplicatedLinear(
+                input_size=self.fc_input_size,
+                output_size=self.config.hidden_size,
+                bias=False,
+                params_dtype=vllm_config.model_config.dtype,
+                quant_config=self.quant_config,
+                prefix=maybe_prefix(prefix, "fc"),
+                return_bias=False,
+            )
+
+        self.norm_output = getattr(self.config, "norm_output", False)
+        self.norm = RMSNorm(
+            self.config.hidden_size,
+            eps=self.config.rms_norm_eps,
+        )
+
+    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
+        return self.embed_tokens(input_ids)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+        input_embeds: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if input_embeds is None:
+            input_embeds = self.embed_input_ids(input_ids)
+        assert hidden_states.shape[-1] == input_embeds.shape[-1]
+
+        residual = None
+        for layer in self.layers:
+            hidden_states, residual = layer(
+                positions=positions,
+                embeds=input_embeds,
+                hidden_states=hidden_states,
+                residual=residual,
+            )
+        hidden_states, hidden_prenorm = self.norm(hidden_states, residual)
+
+        # norm_output variant uses the post-norm hidden states.
+        aux_output = hidden_states if self.norm_output else hidden_prenorm
+
+        return hidden_states, aux_output
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        stacked_params_mapping = [
+            # (param_name, shard_name, shard_id)
+            (".qkv_proj", ".q_proj", "q"),
+            (".qkv_proj", ".k_proj", "k"),
+            (".qkv_proj", ".v_proj", "v"),
+            (".gate_up_proj", ".gate_proj", 0),
+            (".gate_up_proj", ".up_proj", 1),
+        ]
+        params_dict = dict(self.named_parameters())
+        loaded_params: set[str] = set()
+        for name, loaded_weight in weights:
+            if "midlayer." in name:
+                name = name.replace("midlayer.", "layers.0.")
+            # Remapping the name FP8 kv-scale or zero point.
+            if "scale" in name or "zero_point" in name:
+                name = maybe_remap_kv_scale_name(name, params_dict)
+                if name is None:
+                    continue
+            for param_name, weight_name, shard_id in stacked_params_mapping:
+                if weight_name not in name:
+                    continue
+                name = name.replace(weight_name, param_name)
+                param = params_dict[name]
+                weight_loader = param.weight_loader
+                weight_loader(param, loaded_weight, shard_id)
+                break
+            else:
+                param = params_dict[name]
+                weight_loader = getattr(param, "weight_loader", default_weight_loader)
+                weight_loader(param, loaded_weight)
+            loaded_params.add(name)
+        return loaded_params
+
+
+class Eagle3Qwen3ForCausalLM(Qwen3ForCausalLM):
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        nn.Module.__init__(self)
+        self.config = vllm_config.speculative_config.draft_model_config.hf_config
+        # Ensure draft_vocab_size is set
+        # default to the base vocab size when absent
+        if getattr(self.config, "draft_vocab_size", None) is None:
+            base_vocab_size = getattr(self.config, "vocab_size", None)
+            self.config.draft_vocab_size = base_vocab_size
+        target_layer_num = vllm_config.model_config.get_num_layers(
+            vllm_config.parallel_config
+        )
+
+        # Store target layer count in draft config for
+        # proper layer_types indexing in draft models
+        self.config.target_layer_count = target_layer_num
+        self.model = Qwen3Eagle3Model(
+            vllm_config=vllm_config,
+            prefix=maybe_prefix(prefix, "model"),
+            start_layer_id=target_layer_num,
+        )
+
+        logit_scale = getattr(self.config, "logit_scale", 1.0)
+        self.lm_head = ParallelLMHead(
+            self.config.draft_vocab_size,
+            self.config.hidden_size,
+            quant_config=get_draft_quant_config(vllm_config),
+            prefix=maybe_prefix(prefix, "lm_head"),
+        )
+        self.logits_processor = LogitsProcessor(
+            self.config.draft_vocab_size, scale=logit_scale
+        )
+        self.draft_id_to_target_id = nn.Parameter(
+            torch.zeros(self.config.draft_vocab_size, dtype=torch.long),
+            requires_grad=False,
+        )
+
+        self.use_parallel_drafting = vllm_config.speculative_config.parallel_drafting
+
+        if self.use_parallel_drafting:
+            self.register_buffer(
+                "mask_hidden",
+                torch.zeros(1, self.model.fc_input_size),
+                persistent=False,
+            )
+
+    def embed_input_ids(
+        self,
+        input_ids: torch.Tensor,
+        multimodal_embeddings: NestedTensors | None = None,
+        is_multimodal: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        return self.model.embed_input_ids(input_ids)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+        inputs_embeds: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        return self.model(input_ids, positions, hidden_states, inputs_embeds)
+
+    def compute_logits(
+        self,
+        hidden_states: torch.Tensor,
+    ) -> torch.Tensor | None:
+        logits = self.logits_processor(self.lm_head, hidden_states)
+        if self.draft_id_to_target_id is None:
+            assert logits.shape[1] == self.config.vocab_size, (
+                "Expected logits to have shape "
+                f"(*, {self.config.vocab_size}), but got {logits.shape}"
+            )
+            return logits
+
+        base = torch.arange(self.config.draft_vocab_size, device=logits.device)
+        targets = base + self.draft_id_to_target_id
+        logits_new = logits.new_full(
+            (
+                logits.shape[0],
+                self.config.vocab_size,
+            ),
+            float("-inf"),
+        )
+        logits_new[:, targets] = logits
+        return logits_new
+
+    def combine_hidden_states(
+        self,
+        hidden_states: torch.Tensor,
+    ) -> torch.Tensor:
+        if not self.model.use_aux_hidden_state:
+            return hidden_states
+        # combine multiple auxiliary hidden states returned by eagle3
+
+        if self.model.norm_before_fc:
+            hidden_states = self.model.input_norm(hidden_states)
+
+        # `norm_before_fc` adds a single RMSNorm before the FC layer, whereas `fc_norm`
+        # applies separate RMSNorms to each chunk of the hidden states.
+        if self.model.fc_norm is not None:
+            chunks = hidden_states.chunk(self.model.num_aux_layers, dim=-1)
+            hidden_states = torch.cat(
+                [norm(chunk) for norm, chunk in zip(self.model.fc_norm, chunks)],
+                dim=-1,
+            )
+
+        return self.model.fc(hidden_states)
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
+        model_weights = {}
+        includes_draft_id_mapping = False
+        includes_embed_tokens = False
+        includes_mask_hidden = False
+        for name, loaded_weight in weights:
+            if "t2d" in name:
+                continue
+            if "d2t" in name:
+                name = name.replace("d2t", "draft_id_to_target_id")
+                includes_draft_id_mapping = True
+            elif "mask_hidden" in name:
+                # Load mask_hidden directly into buffer
+                if not self.use_parallel_drafting:
+                    logger.warning(
+                        "mask_hidden found in weights but "
+                        "model is not configured for parallel drafting. "
+                        "Skipping loading mask_hidden."
+                    )
+                    continue
+                self.mask_hidden.copy_(loaded_weight.view(1, -1))
+                includes_mask_hidden = True
+                continue
+            elif "lm_head" not in name:
+                name = "model." + name
+            if "embed_tokens" in name:
+                includes_embed_tokens = True
+            model_weights[name] = loaded_weight
+            process_eagle_weight(self, name)
+
+        if not includes_mask_hidden and self.use_parallel_drafting:
+            raise ValueError(
+                "mask_hidden not found in weights but "
+                "model is configured for parallel drafting. "
+                "Please provide mask_hidden in the weights."
+            )
+
+        skip_substrs = ["mask_hidden"]
+        if not includes_draft_id_mapping:
+            skip_substrs.append("draft_id_to_target_id")
+        if not includes_embed_tokens:
+            skip_substrs.append("embed_tokens")
+        if not self.model.use_aux_hidden_state:
+            skip_substrs.append("fc.")
+        if not self.model.norm_before_fc:
+            skip_substrs.append("input_norm.")
+        loader = AutoWeightsLoader(
+            self,
+            skip_prefixes=None,
+            skip_substrs=skip_substrs,
+        )
+        loader.load_weights(model_weights.items())

vllm/model_executor/models/qwen3_vl.py

@@ -1271,7 +1271,7 @@ class Qwen3VLMultiModalProcessor(BaseMultiModalProcessor[Qwen3VLProcessingInfo])
             vision_end_token_id = hf_config.vision_end_token_id
             video_token_id = hf_config.video_token_id
 
-            for item in videos:
+            for item_idx, item in enumerate(videos):
                 video_array, metadata = item
 
                 # NOTE: @JJJYmmm new attr metadata.frames_indices indicates
@@ -1282,6 +1282,12 @@ class Qwen3VLMultiModalProcessor(BaseMultiModalProcessor[Qwen3VLProcessingInfo])
                 # NOTE: a copy of is created to update do_sample_frames,
                 # otherwise mm_hash for the object will be incorrect.
                 video_mm_kwargs = dict(**mm_kwargs)
+                sampled_fps = video_mm_kwargs.get("fps")
+                if is_list_of(sampled_fps, float):
+                    video_mm_kwargs["fps"] = sampled_fps[item_idx]
+                sampled_num_frames = video_mm_kwargs.get("num_frames")
+                if is_list_of(sampled_num_frames, int):
+                    video_mm_kwargs["num_frames"] = sampled_num_frames[item_idx]
                 if "do_sample_frames" not in video_mm_kwargs:
                     # qwen_vl_utils already has "do_sample_frames" in
                     # mm_kwargs, don't overwrite it.
@@ -1363,10 +1369,15 @@ class Qwen3VLMultiModalProcessor(BaseMultiModalProcessor[Qwen3VLProcessingInfo])
         else:
             video_outputs = dict()
 
+        # fps/num_frames are video-only kwargs already consumed by the loop;
+        # exclude them so the text/image processor call below never gets a list.
+        non_video_mm_kwargs = {
+            k: v for k, v in mm_kwargs.items() if k not in ("fps", "num_frames")
+        }
         processed_outputs = super()._call_hf_processor(
             prompt=prompt,
             mm_data=mm_data,
-            mm_kwargs=mm_kwargs,
+            mm_kwargs=non_video_mm_kwargs,
             tok_kwargs=tok_kwargs,
         )
 

vllm/model_executor/models/registry.py

@@ -617,6 +617,8 @@ _SPECULATIVE_DECODING_MODELS = {
     "LlamaForCausalLMEagle3": ("llama_eagle3", "Eagle3LlamaForCausalLM"),
     "Eagle3Qwen2_5vlForCausalLM": ("llama_eagle3", "Eagle3LlamaForCausalLM"),
     "Eagle3Qwen3vlForCausalLM": ("llama_eagle3", "Eagle3LlamaForCausalLM"),
+    "Eagle3Qwen3ForCausalLM": ("qwen3_eagle3", "Eagle3Qwen3ForCausalLM"),
+    "PeagleQwen3ForCausalLM": ("qwen3_eagle3", "Eagle3Qwen3ForCausalLM"),
     "EagleMistralForCausalLM": ("mistral_eagle", "EagleMistralForCausalLM"),
     "EagleMistralLarge3ForCausalLM": (
         "mistral_large_3_eagle",

vllm/models/deepseek_v4/nvidia/model.py

@@ -16,6 +16,7 @@ from vllm.distributed import (
     get_tensor_model_parallel_world_size,
 )
 from vllm.distributed.eplb.eplb_state import EplbLayerState
+from vllm.forward_context import get_forward_context, is_forward_context_available
 from vllm.model_executor.kernels.mhc.tilelang import (
     hc_head_fused_kernel_tilelang,
     mhc_fused_post_pre_tilelang,
@@ -440,6 +441,11 @@ class DeepseekV4MegaMoEExperts(nn.Module):
 
         symm_buffer = self.get_symm_buffer()
         num_tokens = hidden_states.shape[0]
+        is_padding = None
+        if is_forward_context_available():
+            is_padding = get_forward_context().is_padding
+            if is_padding is not None:
+                is_padding = is_padding[:num_tokens]
 
         # EPLB: map logical expert IDs to physical replicas and record load.
         eplb_state = self.eplb_state
@@ -447,6 +453,8 @@ class DeepseekV4MegaMoEExperts(nn.Module):
             assert eplb_state.expert_load_view is not None
             assert eplb_state.logical_replica_count is not None
             assert eplb_state.should_record_tensor is not None
+            if is_padding is not None:
+                topk_ids = torch.where(is_padding.unsqueeze(1), -1, topk_ids)
             topk_ids = eplb_map_to_physical_and_record(
                 topk_ids=topk_ids,
                 expert_load_view=eplb_state.expert_load_view,
@@ -463,6 +471,7 @@ class DeepseekV4MegaMoEExperts(nn.Module):
             symm_buffer.x_sf[:num_tokens],
             symm_buffer.topk_idx[:num_tokens],
             symm_buffer.topk_weights[:num_tokens],
+            is_padding=is_padding,
         )
 
         # This method must have been already called during the weight loading phase.

vllm/models/deepseek_v4/nvidia/ops/prepare_megamoe.py

@@ -19,6 +19,7 @@ def _prepare_megamoe_inputs_kernel(
     x_sf,
     topk_ids,
     topk_weights,
+    is_padding,
     topk_idx_out,
     topk_weights_out,
     hidden_stride_m: tl.constexpr,
@@ -31,6 +32,7 @@ def _prepare_megamoe_inputs_kernel(
     topk_ids_stride_k: tl.constexpr,
     topk_weights_stride_m: tl.constexpr,
     topk_weights_stride_k: tl.constexpr,
+    is_padding_stride_m: tl.constexpr,
     topk_idx_stride_m: tl.constexpr,
     topk_idx_stride_k: tl.constexpr,
     topk_weights_out_stride_m: tl.constexpr,
@@ -85,12 +87,16 @@ def _prepare_megamoe_inputs_kernel(
     if k_block_id == 0:
         topk_offsets = tl.arange(0, BLOCK_TOPK)
         topk_mask = topk_offsets < top_k
+        token_is_padding = False
+        if is_padding is not None:
+            token_is_padding = tl.load(is_padding + token_id * is_padding_stride_m)
 
         ids = tl.load(
             topk_ids + token_id * topk_ids_stride_m + topk_offsets * topk_ids_stride_k,
             mask=topk_mask,
             other=0,
         ).to(tl.int64)
+        ids = tl.where(token_is_padding, -1, ids)
         tl.store(
             topk_idx_out
             + token_id * topk_idx_stride_m
@@ -106,6 +112,7 @@ def _prepare_megamoe_inputs_kernel(
             mask=topk_mask,
             other=0.0,
         )
+        weights = tl.where(token_is_padding, 0.0, weights)
         tl.store(
             topk_weights_out
             + token_id * topk_weights_out_stride_m
@@ -123,6 +130,7 @@ def prepare_megamoe_inputs(
     x_sf: torch.Tensor,
     topk_idx_out: torch.Tensor,
     topk_weights_out: torch.Tensor,
+    is_padding: torch.Tensor | None = None,
 ) -> None:
     num_tokens, hidden_size = hidden_states.shape
     if num_tokens == 0:
@@ -142,12 +150,14 @@ def prepare_megamoe_inputs(
     block_k = 128
     grid = (num_tokens, triton.cdiv(hidden_size, block_k))
     block_topk = triton.next_power_of_2(top_k)
+    padding_stride_m = is_padding.stride(0) if is_padding is not None else 0
     _prepare_megamoe_inputs_kernel[grid](
         hidden_states,
         x_fp8,
         x_sf,
         topk_ids,
         topk_weights,
+        is_padding,
         topk_idx_out,
         topk_weights_out,
         hidden_states.stride(0),
@@ -160,6 +170,7 @@ def prepare_megamoe_inputs(
         topk_ids.stride(1),
         topk_weights.stride(0),
         topk_weights.stride(1),
+        padding_stride_m,
         topk_idx_out.stride(0),
         topk_idx_out.stride(1),
         topk_weights_out.stride(0),

vllm/multimodal/video.py

@@ -7,6 +7,7 @@ from typing import Any, ClassVar, Literal, NamedTuple, cast
 
 import numpy as np
 import numpy.typing as npt
+import torch
 
 from vllm.logger import init_logger
 from vllm.utils.import_utils import PlaceholderModule
@@ -653,6 +654,85 @@ class Qwen3VLVideoBackend(VideoBackend):
         )
 
 
+@VIDEO_LOADER_REGISTRY.register(
+    "qwen2_vl",
+    video_processor="Qwen2VLVideoProcessor",
+)
+class Qwen2VLVideoBackend(VideoBackend):
+    """Qwen2-VL / Qwen2.5-VL fps-based video backend.
+
+    Ports transformers' ``Qwen2VLVideoProcessor.sample_frames`` (fps mode),
+    shared by Qwen2-VL and Qwen2.5-VL (the latter has no video processor of its
+    own): sample ``total / original_fps * fps`` frames, clamp to
+    ``[min_frames, max_frames]`` (4 and 768), floor to a multiple of
+    ``temporal_patch_size`` (2), and take indices with the exact
+    ``torch.arange(0, total, total / n)`` call so they match HF byte-for-byte.
+
+    ``num_frames`` is ignored (fps-driven, like the Qwen3-VL loader). The
+    float32 step can emit an out-of-range tail index (e.g. 451 for a 451-frame
+    clip); it is clamped to the last valid frame.
+    """
+
+    @classmethod
+    def compute_frames_index_to_sample(
+        cls,
+        source: VideoSourceMetadata,
+        target: VideoTargetMetadata,
+        **kwargs,
+    ) -> list[int]:
+        # Refer to:
+        # https://github.com/huggingface/transformers/blob/v5.7.0/src/transformers/models/qwen2_vl/video_processing_qwen2_vl.py#L122-L190
+        total_frames_num = source.total_frames_num
+        original_fps = source.original_fps
+        temporal_patch_size = kwargs.get("temporal_patch_size", 2)
+        min_frames = kwargs.get("min_frames", 4)
+        max_frames = kwargs.get("max_frames", 768)
+
+        # vLLM reports original_fps == 0 for clips with unknown/variable fps
+        # (VFR, malformed, streaming); fail loudly instead of dividing by zero.
+        if original_fps <= 0:
+            raise ValueError(
+                "Qwen2-VL video sampling needs a known source fps, but the "
+                "container reported 0 (variable or unknown frame rate)."
+            )
+
+        max_frames = (
+            math.floor(min(max_frames, total_frames_num) / temporal_patch_size)
+            * temporal_patch_size
+        )
+        n = total_frames_num / original_fps * target.fps
+        n = min(max(n, min_frames), max_frames, total_frames_num)
+        n = math.floor(n / temporal_patch_size) * temporal_patch_size
+
+        # ``torch.arange`` matches transformers' float32 index math exactly
+        # (numpy's float64 diverges by a frame on some inputs); clamp the tail
+        # because that step can emit an index == total_frames_num.
+        indices = torch.arange(0, total_frames_num, total_frames_num / n).int()
+        return torch.clamp(indices, max=total_frames_num - 1).tolist()
+
+    @classmethod
+    def load_bytes(
+        cls,
+        data: bytes,
+        num_frames: int = -1,
+        fps: int = 2,
+        max_duration: int = 300,
+        frame_recovery: bool = False,
+        *,
+        backend: Literal["opencv", "pyav"] = "opencv",
+        **kwargs,
+    ) -> tuple[npt.NDArray, dict[str, Any]]:
+        return super().load_bytes(
+            data,
+            num_frames=num_frames,
+            fps=fps,
+            max_duration=max_duration,
+            frame_recovery=frame_recovery,
+            backend=backend,
+            **kwargs,
+        )
+
+
 @VIDEO_LOADER_REGISTRY.register(
     "opencv_dynamic",
     video_processor="Glm4vVideoProcessor",

vllm/transformers_utils/configs/speculators/algos.py

@@ -36,7 +36,14 @@ def update_eagle3(config_dict: dict, pre_trained_config: dict) -> None:
         "norm_before_residual", True
     )
     pre_trained_config["norm_before_fc"] = config_dict.get("norm_before_fc", False)
-    pre_trained_config["architectures"] = ["Eagle3LlamaForCausalLM"]
+    eagle3_arch_map = {
+        "qwen3": "Eagle3Qwen3ForCausalLM",
+        "llama": "Eagle3LlamaForCausalLM",
+    }
+    model_type = pre_trained_config.get("model_type", "llama")
+    if model_type not in eagle3_arch_map:
+        raise ValueError(f"Unsupported model_type {model_type} for Eagle3 speculator")
+    pre_trained_config["architectures"] = [eagle3_arch_map[model_type]]
     if config_dict.get("eagle_aux_hidden_state_layer_ids"):
         pre_trained_config["eagle_aux_hidden_state_layer_ids"] = config_dict[
             "eagle_aux_hidden_state_layer_ids"
@@ -59,7 +66,6 @@ def update_peagle(config_dict: dict, pre_trained_config: dict) -> None:
     - eagle_aux_hidden_state_layer_ids: Layer indices from the target model
         whose intermediate hidden states are used as auxiliary inputs
     """
-    pre_trained_config["architectures"] = ["PeagleLlamaForCausalLM"]
     pre_trained_config["draft_vocab_size"] = config_dict.get("draft_vocab_size")
     if config_dict.get("target_hidden_size") is not None:
         pre_trained_config["target_hidden_size"] = config_dict["target_hidden_size"]
@@ -67,6 +73,14 @@ def update_peagle(config_dict: dict, pre_trained_config: dict) -> None:
         "norm_before_residual", False
     )
     pre_trained_config["norm_before_fc"] = config_dict.get("norm_before_fc", False)
+    peagle_arch_map = {
+        "qwen3": "PeagleQwen3ForCausalLM",
+        "llama": "PeagleLlamaForCausalLM",
+    }
+    model_type = pre_trained_config.get("model_type", "llama")
+    if model_type not in peagle_arch_map:
+        raise ValueError(f"Unsupported model_type {model_type} for PEagle speculator")
+    pre_trained_config["architectures"] = [peagle_arch_map[model_type]]
     pre_trained_config["pard_token"] = config_dict["mask_token_id"]
     if config_dict.get("eagle_aux_hidden_state_layer_ids"):
         pre_trained_config["eagle_aux_hidden_state_layer_ids"] = config_dict[

vllm/transformers_utils/processor.py

@@ -18,6 +18,7 @@ from transformers.audio_utils import AudioInput
 from transformers.feature_extraction_utils import FeatureExtractionMixin
 from transformers.image_processing_utils import BaseImageProcessor
 from transformers.image_utils import ImageInput
+from transformers.models.auto.video_processing_auto import VIDEO_PROCESSOR_MAPPING_NAMES
 from transformers.processing_utils import ProcessorMixin
 from transformers.video_processing_utils import BaseVideoProcessor
 from transformers.video_utils import VideoInput
@@ -169,6 +170,15 @@ def get_video_processor_cls_name_from_config(
         config = get_hf_file_to_dict(file, processor_name, revision=revision)
         if config and "video_processor_type" in config:
             return config["video_processor_type"]
+
+    # Some models ship no explicit ``video_processor_type`` in their
+    # preprocessor config. Fall back to transformers' ``model_type`` -> video
+    # processor mapping so these still resolve to their registered loader
+    # instead of the generic opencv fallback. The mapping is ``None`` for a
+    # given type when torchvision is unavailable; callers then use opencv.
+    model_config = get_hf_file_to_dict("config.json", processor_name, revision=revision)
+    if model_config and "model_type" in model_config:
+        return VIDEO_PROCESSOR_MAPPING_NAMES.get(model_config["model_type"])
     return None
 
 

vllm/v1/attention/backend.py

@@ -201,6 +201,38 @@ class AttentionBackend(ABC):
 
         return min(s.base if isinstance(s, MultipleOf) else s for s in supported_sizes)
 
+    @classmethod
+    def indexes_kv_by_block_stride(cls) -> bool:
+        """Whether the backend reads KV pages by the runtime block stride.
+
+        True when ``num_blocks`` is the outermost physical dimension of the KV
+        cache, so the backend tolerates a non-contiguous block dim. This gates
+        page size padding and cross-layer uniform KV layout.
+
+        Returns:
+            True if the backend's physical KV layout is num-blocks-first. False
+            otherwise, including when the backend does not define a layered
+            stride order.
+        """
+        try:
+            kv_cache_stride_order = cls.get_kv_cache_stride_order(
+                include_num_layers_dimension=False
+            )
+            layered_kv_cache_stride_order = cls.get_kv_cache_stride_order(
+                include_num_layers_dimension=True
+            )
+        except (AttributeError, NotImplementedError):
+            return False
+
+        # Check that attention backend includes a layers dimension.
+        if len(layered_kv_cache_stride_order) != len(kv_cache_stride_order) + 1:
+            return False
+
+        # stride_order[0] == 0 means num_layers stays first in physical
+        # layout (identity permutation), so indexing by block stride is
+        # not supported.
+        return layered_kv_cache_stride_order[0] != 0
+
     @classmethod
     def is_mla(cls) -> bool:
         return False

vllm/v1/attention/ops/triton_reshape_and_cache_flash.py

@@ -181,6 +181,7 @@ def _reshape_cache_per_token_head(
     HEAD_SIZE_PADDED: tl.constexpr,  # next_power_of_2(max(head_size, head_size_v))
     QUANT_MAX: tl.constexpr = 127.0,
     QUANT_MIN: tl.constexpr = -128.0,
+    IS_INT_QUANT: tl.constexpr = False,
 ):
     tok = tl.program_id(0)
     head = tl.program_id(1)
@@ -211,7 +212,11 @@ def _reshape_cache_per_token_head(
         k_scale,
     )
 
-    k_q = tl.clamp(k_h * (1.0 / k_scale), QUANT_MIN, QUANT_MAX)
+    k_q = k_h * (1.0 / k_scale)
+    if IS_INT_QUANT:
+        # Round half away from zero before the int8 store truncates.
+        k_q = tl.where(k_q >= 0, k_q + 0.5, k_q - 0.5)
+    k_q = tl.clamp(k_q, QUANT_MIN, QUANT_MAX)
     tl.store(
         key_cache_ptr
         + blk * stride_kc_blk
@@ -239,7 +244,11 @@ def _reshape_cache_per_token_head(
         v_scale,
     )
 
-    v_q = tl.clamp(v_h * (1.0 / v_scale), QUANT_MIN, QUANT_MAX)
+    v_q = v_h * (1.0 / v_scale)
+    if IS_INT_QUANT:
+        # Round half away from zero before the int8 store truncates.
+        v_q = tl.where(v_q >= 0, v_q + 0.5, v_q - 0.5)
+    v_q = tl.clamp(v_q, QUANT_MIN, QUANT_MAX)
     tl.store(
         value_cache_ptr
         + blk * stride_vc_blk
@@ -327,6 +336,7 @@ def triton_reshape_and_cache_flash_per_token_head_quant(
         HEAD_SIZE_PADDED=head_size_padded,
         QUANT_MAX=quant_max,
         QUANT_MIN=quant_min,
+        IS_INT_QUANT=cache_dtype == torch.int8,
         num_warps=num_warps,
     )
 

vllm/v1/core/kv_cache_utils.py

@@ -20,6 +20,7 @@ from vllm.utils.math_utils import cdiv, round_up
 from vllm.utils.mem_utils import format_gib
 from vllm.utils.torch_utils import get_dtype_size
 from vllm.v1.kv_cache_interface import (
+    AttentionSpec,
     ChunkedLocalAttentionSpec,
     FullAttentionSpec,
     HiddenStateCacheSpec,
@@ -1029,9 +1030,14 @@ def unify_kv_cache_spec_page_size(
 ) -> dict[str, KVCacheSpec]:
     """
     Unify the page size of the given KVCacheSpec. If the page size of all layers
-    are the same, return the original KVCacheSpec. If not same, unify the page
-    size by increasing the block size of layers with smaller page size. Raise
-    NotImplementedError if failed to unify the page size.
+    are the same, return the original KVCacheSpec. If not same, first try to
+    unify page size by increasing the block size of layers with smaller page
+    size. If a smaller attention page does not evenly divide the maximum page
+    size, keep its logical block size and pad its physical page instead --- but
+    only for attention layers whose backend opts in via
+    ``AttentionSpec.indexes_kv_by_block_stride`` (the padded page is read through
+    a strided view, which not every backend handles). Raise NotImplementedError
+    if failed to unify the page size.
 
     Args:
         kv_cache_spec: The KVCacheSpec of each attention layer in the model
@@ -1051,14 +1057,23 @@ def unify_kv_cache_spec_page_size(
             new_kv_cache_spec[layer_name] = layer_spec
         else:
             layer_page_size = layer_spec.page_size_bytes
-            if max_page_size % layer_page_size != 0:
+            if max_page_size % layer_page_size == 0:
+                ratio = max_page_size // layer_page_size
+                new_block_size = layer_spec.block_size * ratio
+                new_spec = replace(layer_spec, block_size=new_block_size)
+            elif (
+                isinstance(layer_spec, AttentionSpec)
+                and layer_spec.indexes_kv_by_block_stride
+            ):
+                new_spec = replace(layer_spec, page_size_padded=max_page_size)
+            else:
                 raise NotImplementedError(
-                    "The page size of the layer is not divisible by the "
-                    "maximum page size. Cannot unify by adjusting block_size."
+                    f"Layer {layer_name}: page size is not divisible by the "
+                    "maximum page size and cannot be padded. Padding is only "
+                    "supported for attention layers whose backend indexes KV "
+                    "pages by the block stride (indexes_kv_by_block_stride is "
+                    "True)."
                 )
-            ratio = max_page_size // layer_page_size
-            new_block_size = layer_spec.block_size * ratio
-            new_spec = replace(layer_spec, block_size=new_block_size)
             assert new_spec.page_size_bytes == max_page_size
             new_kv_cache_spec[layer_name] = new_spec
     return new_kv_cache_spec

vllm/v1/kv_cache_interface.py

@@ -163,6 +163,7 @@ class AttentionSpec(KVCacheSpec):
     dtype: torch.dtype
     kv_quant_mode: KVQuantMode = KVQuantMode.NONE
     page_size_padded: int | None = None
+    indexes_kv_by_block_stride: bool = False
 
     @property
     def page_size_bytes(self) -> int:
@@ -283,6 +284,7 @@ class FullAttentionSpec(AttentionSpec):
             dtype=specs[0].dtype,
             kv_quant_mode=specs[0].kv_quant_mode,
             page_size_padded=specs[0].page_size_padded,
+            indexes_kv_by_block_stride=specs[0].indexes_kv_by_block_stride,
             sliding_window=cls.merge_window_sizes(sliding_window),
             attention_chunk_size=cls.merge_window_sizes(attention_chunk_size),
             # If any layer in the group is non-causal, treat the group as
@@ -403,13 +405,16 @@ class MLAAttentionSpec(FullAttentionSpec):
         cache_dtype_str_set = set(spec.cache_dtype_str for spec in specs)
         compress_ratio_set = set(spec.compress_ratio for spec in specs)
         model_version_set = set(spec.model_version for spec in specs)
+        block_stride_set = set(spec.indexes_kv_by_block_stride for spec in specs)
         assert (
             len(cache_dtype_str_set) == 1
             and len(compress_ratio_set) == 1
             and len(model_version_set) == 1
+            and len(block_stride_set) == 1
         ), (
             "All attention layers in the same KV cache group must use the same "
-            "quantization method, compress ratio, and model version."
+            "quantization method, compress ratio, model version, and KV block "
+            "stride indexing."
         )
         return cls(
             block_size=specs[0].block_size,
@@ -418,6 +423,7 @@ class MLAAttentionSpec(FullAttentionSpec):
             dtype=specs[0].dtype,
             kv_quant_mode=specs[0].kv_quant_mode,
             page_size_padded=specs[0].page_size_padded,
+            indexes_kv_by_block_stride=block_stride_set.pop(),
             cache_dtype_str=cache_dtype_str_set.pop(),
             compress_ratio=compress_ratio_set.pop(),
             model_version=model_version_set.pop(),
@@ -584,15 +590,17 @@ class SlidingWindowMLASpec(SlidingWindowSpec):
         compress_ratio_set = set(spec.compress_ratio for spec in specs)
         model_version_set = set(spec.model_version for spec in specs)
         sliding_window_set = set(spec.sliding_window for spec in specs)
+        block_stride_set = set(spec.indexes_kv_by_block_stride for spec in specs)
         assert (
             len(cache_dtype_str_set) == 1
             and len(compress_ratio_set) == 1
             and len(model_version_set) == 1
             and len(sliding_window_set) == 1
+            and len(block_stride_set) == 1
         ), (
             "All attention layers in the same KV cache group must use the same "
-            "quantization method, compress ratio, model version and sliding "
-            "window size."
+            "quantization method, compress ratio, model version, sliding "
+            "window size, and KV block stride indexing."
         )
         return cls(
             block_size=specs[0].block_size,
@@ -600,6 +608,7 @@ class SlidingWindowMLASpec(SlidingWindowSpec):
             head_size=specs[0].head_size,
             dtype=specs[0].dtype,
             page_size_padded=specs[0].page_size_padded,
+            indexes_kv_by_block_stride=block_stride_set.pop(),
             sliding_window=sliding_window_set.pop(),
             cache_dtype_str=cache_dtype_str_set.pop(),
             compress_ratio=compress_ratio_set.pop(),
@@ -711,6 +720,7 @@ class SinkFullAttentionSpec(FullAttentionSpec):
             dtype=specs[0].dtype,
             kv_quant_mode=specs[0].kv_quant_mode,
             page_size_padded=specs[0].page_size_padded,
+            indexes_kv_by_block_stride=specs[0].indexes_kv_by_block_stride,
             sliding_window=cls.merge_window_sizes(sliding_window),
             attention_chunk_size=cls.merge_window_sizes(attention_chunk_size),
             non_causal=any(spec.non_causal for spec in specs),

vllm/v1/kv_offload/base.py

@@ -129,6 +129,7 @@ The class provides the following primitives:
 @dataclass(frozen=True)
 class OffloadingMetricMetadata:
     documentation: str
+    labelnames: tuple[str, ...] = ()
 
 
 @dataclass(frozen=True)

vllm/v1/spec_decode/llm_base_proposer.py

@@ -24,6 +24,7 @@ from vllm.model_executor.models.deepseek_eagle3 import Eagle3DeepseekV2ForCausal
 from vllm.model_executor.models.interfaces import SupportsMultiModal
 from vllm.model_executor.models.llama_eagle3 import Eagle3LlamaForCausalLM
 from vllm.model_executor.models.qwen3_dflash import DFlashQwen3ForCausalLM
+from vllm.model_executor.models.qwen3_eagle3 import Eagle3Qwen3ForCausalLM
 from vllm.multimodal import MULTIMODAL_REGISTRY
 from vllm.platforms import current_platform
 from vllm.utils.torch_utils import PIN_MEMORY, async_tensor_h2d
@@ -473,6 +474,7 @@ class SpecDecodeBaseProposer:
                     Eagle3LlamaForCausalLM,
                     Eagle3DeepseekV2ForCausalLM,
                     DFlashQwen3ForCausalLM,
+                    Eagle3Qwen3ForCausalLM,
                 ),
             )
             target_hidden_states = self.model.combine_hidden_states(

vllm/v1/worker/gpu/attn_utils.py

@@ -1,13 +1,17 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from collections.abc import Iterable, Sequence
-from dataclasses import dataclass
+from dataclasses import dataclass, replace
 from math import prod
 from typing import Any, cast
 
 import torch
 
-from vllm.config import VllmConfig, get_layers_from_vllm_config
+from vllm.config import (
+    VllmConfig,
+    get_layers_from_vllm_config,
+    set_current_vllm_config,
+)
 from vllm.model_executor.layers.attention import Attention
 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
 from vllm.utils.torch_utils import get_dtype_size
@@ -47,6 +51,13 @@ def get_kv_cache_spec(vllm_config: VllmConfig) -> dict[str, KVCacheSpec]:
             continue
         # Skip modules that don't need KV cache (eg encoder-only attention)
         if spec := attn_module.get_kv_cache_spec(vllm_config):
+            if isinstance(spec, AttentionSpec):
+                backend = attn_module.get_attn_backend()
+                # indexes_kv_by_block_stride() -> get_kv_cache_stride_order() ->
+                # get_kv_cache_layout() needs the current vLLM config.
+                with set_current_vllm_config(vllm_config):
+                    indexes = backend.indexes_kv_by_block_stride()
+                spec = replace(spec, indexes_kv_by_block_stride=indexes)
             kv_cache_spec[layer_name] = spec
     return kv_cache_spec
 
@@ -180,6 +191,62 @@ def _allocate_kv_cache(
     return kv_cache_raw_tensors
 
 
+def _reshape_attention_kv_cache(
+    kv_raw_tensor: torch.Tensor,
+    kv_cache_spec: AttentionSpec,
+    kv_cache_shape: tuple[int, ...],
+    kv_cache_stride_order: tuple[int, ...],
+    num_blocks: int,
+    packing: tuple[int, int] | None,
+) -> torch.Tensor:
+    permuted_kv_cache_shape = tuple(kv_cache_shape[i] for i in kv_cache_stride_order)
+    inv_order = [
+        kv_cache_stride_order.index(i) for i in range(len(kv_cache_stride_order))
+    ]
+    dtype = kv_cache_spec.dtype
+
+    if packing is not None:
+        offset, block_stride = packing
+        assert inv_order[0] == 0
+        page_bytes = prod(kv_cache_shape[1:]) * get_dtype_size(dtype)
+        kv_cache = (
+            kv_raw_tensor.view(-1, block_stride)[:, offset : offset + page_bytes]
+            .view(dtype)
+            .view(kv_cache_shape)
+        )
+    elif kv_cache_spec.page_size_padded is not None:
+        # Use a strided view to skip the padding between physical pages.
+        #
+        # Only num-blocks-first layouts are supported (the block dimension is
+        # dim 0 of the unpermuted shape). kv-first layouts such as ROCm's
+        # ``(2, num_blocks, ...)`` are intentionally not supported here. For a
+        # num-blocks-first layout the only stride that must change is the block
+        # stride: every other (contiguous) stride already steps within the
+        # unpadded region of a page, so no further adjustment is needed.
+        assert kv_cache_shape[0] == num_blocks, (
+            "Padded KV pages require a num-blocks-first KV cache layout (got "
+            f"shape {kv_cache_shape} with num_blocks={num_blocks}); "
+            "kv-first layouts are not supported."
+        )
+        dtype_size = get_dtype_size(kv_cache_spec.dtype)
+        page_stride = kv_cache_spec.page_size_bytes // dtype_size
+
+        num_blocks_dim = inv_order[0]
+        strides = list(torch.empty(permuted_kv_cache_shape).stride())
+        strides[num_blocks_dim] = page_stride
+
+        kv_cache = torch.as_strided(
+            kv_raw_tensor.view(dtype),
+            size=permuted_kv_cache_shape,
+            stride=tuple(strides),
+        )
+    else:
+        # No padding — safe to use a contiguous view.
+        kv_cache = kv_raw_tensor.view(dtype).view(permuted_kv_cache_shape)
+
+    return kv_cache.permute(*inv_order)
+
+
 def _reshape_kv_cache(
     attn_groups: Sequence[AttentionGroup],
     kv_cache_raw_tensors: dict[str, torch.Tensor],
@@ -248,45 +315,14 @@ def _reshape_kv_cache(
                 except (AttributeError, NotImplementedError):
                     kv_cache_stride_order = tuple(range(len(kv_cache_shape)))
 
-                kv_cache_shape = tuple(kv_cache_shape[i] for i in kv_cache_stride_order)
-                inv_order = [
-                    kv_cache_stride_order.index(i)
-                    for i in range(len(kv_cache_stride_order))
-                ]
-
-                dtype = kv_cache_spec.dtype
-                if packing is not None:
-                    offset, block_stride = packing
-                    assert inv_order[0] == 0
-                    page_bytes = prod(kv_cache_shape[1:]) * get_dtype_size(dtype)
-                    kv_cache = (
-                        kv_raw_tensor.view(-1, block_stride)[
-                            :, offset : offset + page_bytes
-                        ]
-                        .view(dtype)
-                        .view(kv_cache_shape)
-                    )
-                elif kv_cache_spec.page_size_padded is not None:
-                    # Use strided view to handle page_size_bytes that
-                    # include padding. This follows the same pattern as
-                    # MambaSpec handling in gpu_model_runner.py.
-                    # NOTE: This assumes kv_cache_shape[0] == num_blocks
-                    # (i.e. the first physical dimension is the block
-                    # index), which holds for all current backends
-                    # (MLA, FlashAttention, TritonAttention, etc.).
-                    dtype_size = get_dtype_size(dtype)
-                    page_stride = kv_cache_spec.page_size_bytes // dtype_size
-                    strides = list(torch.empty(kv_cache_shape).stride())
-                    strides[inv_order[0]] = page_stride
-                    kv_cache = torch.as_strided(
-                        kv_raw_tensor.view(dtype),
-                        size=kv_cache_shape,
-                        stride=tuple(strides),
-                    )
-                else:
-                    # No padding — safe to use a contiguous view.
-                    kv_cache = kv_raw_tensor.view(dtype).view(kv_cache_shape)
-                kv_caches[layer_name] = kv_cache.permute(*inv_order)
+                kv_caches[layer_name] = _reshape_attention_kv_cache(
+                    kv_raw_tensor,
+                    kv_cache_spec,
+                    kv_cache_shape,
+                    kv_cache_stride_order,
+                    kernel_num_blocks,
+                    packing,
+                )
 
             elif isinstance(kv_cache_spec, MambaSpec):
                 has_mamba = True

vllm/v1/worker/gpu/cudagraph_utils.py

@@ -471,6 +471,9 @@ class ModelCudaGraphManager(CudaGraphManager):
                 skip_attn=(desc.cg_mode == CUDAGraphMode.PIECEWISE),
             )
 
+            # Capture with all-padding padding mask.
+            input_buffers.is_padding.fill_(True)
+
             def forward_fn(cg_mode: CUDAGraphMode) -> None:
                 batch_descriptor = None
                 if cg_mode == CUDAGraphMode.PIECEWISE:
@@ -488,6 +491,7 @@ class ModelCudaGraphManager(CudaGraphManager):
                     num_tokens_across_dp=num_tokens_across_dp,
                     slot_mapping=slot_mappings,
                     batch_descriptor=batch_descriptor,
+                    is_padding=input_buffers.is_padding[:num_tokens],
                 ):
                     if cg_mode == CUDAGraphMode.PIECEWISE:
                         # PIECEWISE graph (compiled PW or breakable, chosen inside

vllm/v1/worker/gpu/input_batch.py

@@ -22,6 +22,7 @@ class InputBuffers:
 
         self.input_ids = torch.zeros(max_num_tokens, dtype=torch.int32, device=device)
         self.positions = torch.zeros(max_num_tokens, dtype=torch.int64, device=device)
+        self.is_padding = torch.zeros(max_num_tokens, dtype=torch.bool, device=device)
         self.query_start_loc = torch.zeros(
             max_num_reqs + 1, dtype=torch.int32, device=device
         )
@@ -83,6 +84,8 @@ class InputBatch:
     input_ids: torch.Tensor
     # [num_tokens_after_padding]
     positions: torch.Tensor
+    # [num_tokens_after_padding]
+    is_padding: torch.Tensor
 
     # [total_num_logits]
     logits_indices: torch.Tensor
@@ -99,9 +102,14 @@ class InputBatch:
         num_reqs: int,
         num_tokens: int,
         input_buffers: InputBuffers,
+        num_actual_tokens: int | None = None,
     ) -> "InputBatch":
         assert 0 < num_reqs <= num_tokens
         device = input_buffers.device
+        # Rows [num_actual_tokens, num_tokens) are treated as padding.
+        # Default (None) means all real.
+        if num_actual_tokens is None:
+            num_actual_tokens = num_tokens
 
         req_ids = [f"req_{i}_{random_uuid()}" for i in range(num_reqs)]
         idx_mapping_np = np.arange(num_reqs, dtype=np.int32)
@@ -134,6 +142,10 @@ class InputBatch:
         input_ids = input_buffers.input_ids[:num_tokens].zero_()
         positions = input_buffers.positions[:num_tokens].zero_()
 
+        input_buffers.is_padding[:num_actual_tokens].fill_(False)
+        input_buffers.is_padding[num_actual_tokens:num_tokens].fill_(True)
+        is_padding = input_buffers.is_padding[:num_tokens]
+
         logits_indices = query_start_loc[1:] - 1
         cu_num_logits = torch.arange(num_reqs + 1, device=device, dtype=torch.int32)
         cu_num_logits_np = np.arange(num_reqs + 1, dtype=np.int32)
@@ -164,6 +176,7 @@ class InputBatch:
             max_seq_len_np=None,
             input_ids=input_ids,
             positions=positions,
+            is_padding=is_padding,
             logits_indices=logits_indices,
             cu_num_logits=cu_num_logits,
             cu_num_logits_np=cu_num_logits_np,

vllm/v1/worker/gpu/model_runner.py

@@ -847,6 +847,10 @@ class GPUModelRunner(LoRAModelRunnerMixin):
         num_tokens = scheduler_output.total_num_scheduled_tokens
         num_tokens_after_padding = batch_desc.num_tokens
         assert num_tokens > 0
+        # Mark trailing cudagraph-padding rows so the MoE can optionally
+        # skip a2a + compute for them.
+        self.input_buffers.is_padding[:num_tokens].fill_(False)
+        self.input_buffers.is_padding[num_tokens:num_tokens_after_padding].fill_(True)
         num_tokens_per_req = scheduler_output.num_scheduled_tokens
         num_reqs = len(num_tokens_per_req)
 
@@ -1001,6 +1005,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
             max_seq_len_np=max_seq_len_np,
             input_ids=self.input_buffers.input_ids[:num_tokens_after_padding],
             positions=self.input_buffers.positions[:num_tokens_after_padding],
+            is_padding=self.input_buffers.is_padding[:num_tokens_after_padding],
             logits_indices=logits_indices,
             cu_num_logits=cu_num_logits,
             cu_num_logits_np=cu_num_logits_np,
@@ -1175,6 +1180,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
                 batch_desc.num_reqs or num_reqs,
                 batch_desc.num_tokens,
                 self.input_buffers,
+                num_actual_tokens=0,
             )
             if not skip_attn_for_dummy_run:
                 block_tables, slot_mappings = self.prepare_dummy_attn(input_batch)
@@ -1277,6 +1283,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
                 batch_descriptor=batch_descriptor,
                 slot_mapping=slot_mappings_by_layer,
                 skip_compiled=skip_compiled,
+                is_padding=input_batch.is_padding,
             ):
                 self.kv_connector.pre_forward(scheduler_output)
                 if batch_desc.cg_mode == CUDAGraphMode.PIECEWISE:

vllm/v1/worker/gpu_model_runner.py

@@ -12,7 +12,6 @@ from contextlib import contextmanager
 from copy import copy, deepcopy
 from dataclasses import dataclass, replace
 from functools import reduce
-from math import prod
 from typing import TYPE_CHECKING, Any, NamedTuple, TypeAlias, cast
 
 import numpy as np
@@ -203,6 +202,7 @@ from vllm.v1.worker.cp_utils import (
 )
 from vllm.v1.worker.dp_utils import coordinate_batch_across_dp
 from vllm.v1.worker.ec_connector_model_runner_mixin import ECConnectorModelRunnerMixin
+from vllm.v1.worker.gpu.attn_utils import _reshape_attention_kv_cache
 from vllm.v1.worker.gpu.pool.late_interaction_runner import LateInteractionRunner
 from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch
 from vllm.v1.worker.gpu_ubatch_wrapper import UBatchWrapper
@@ -7125,62 +7125,20 @@ class GPUModelRunner(
                         kv_cache_spec.head_size,
                         cache_dtype_str=self.cache_config.cache_dtype,
                     )
-                    dtype = kv_cache_spec.dtype
                     try:
                         kv_cache_stride_order = attn_backend.get_kv_cache_stride_order()
                         assert len(kv_cache_stride_order) == len(kv_cache_shape)
                     except (AttributeError, NotImplementedError):
                         kv_cache_stride_order = tuple(range(len(kv_cache_shape)))
-                    # The allocation respects the backend-defined stride order
-                    # to ensure the semantic remains consistent for each
-                    # backend. We first obtain the generic kv cache shape and
-                    # then permute it according to the stride order which could
-                    # result in a non-contiguous tensor.
-                    kv_cache_shape = tuple(
-                        kv_cache_shape[i] for i in kv_cache_stride_order
+                    raw_tensor = kv_cache_raw_tensors[layer_name]
+                    kv_caches[layer_name] = _reshape_attention_kv_cache(
+                        raw_tensor,
+                        kv_cache_spec,
+                        kv_cache_shape,
+                        kv_cache_stride_order,
+                        kernel_num_blocks,
+                        packing,
                     )
-                    # Maintain original KV shape view.
-                    inv_order = [
-                        kv_cache_stride_order.index(i)
-                        for i in range(len(kv_cache_stride_order))
-                    ]
-
-                    if packing is not None:
-                        offset, block_stride = packing
-                        assert inv_order[0] == 0
-                        page_bytes = prod(kv_cache_shape[1:]) * get_dtype_size(dtype)
-                        kv_cache = (
-                            kv_cache_raw_tensors[layer_name]
-                            .view(-1, block_stride)[:, offset : offset + page_bytes]
-                            .view(dtype)
-                            .view(kv_cache_shape)
-                        )
-                    elif kv_cache_spec.page_size_padded is not None:
-                        # Use strided view to handle page_size_bytes that
-                        # include padding. This follows
-                        # the same pattern as MambaSpec handling below.
-                        # NOTE: This assumes kv_cache_shape[0] == num_blocks
-                        # (i.e. the first physical dimension is the block
-                        # index), which holds for MLA backends but NOT for
-                        # standard attention backends whose shape starts with
-                        # a K/V dimension of size 2.
-                        dtype_size = get_dtype_size(dtype)
-                        page_stride = kv_cache_spec.page_size_bytes // dtype_size
-                        strides = list(torch.empty(kv_cache_shape).stride())
-                        strides[inv_order[0]] = page_stride
-                        kv_cache = torch.as_strided(
-                            kv_cache_raw_tensors[layer_name].view(dtype),
-                            size=kv_cache_shape,
-                            stride=tuple(strides),
-                        )
-                    else:
-                        # No padding — safe to use a contiguous view.
-                        kv_cache = (
-                            kv_cache_raw_tensors[layer_name]
-                            .view(dtype)
-                            .view(kv_cache_shape)
-                        )
-                    kv_caches[layer_name] = kv_cache.permute(*inv_order)
 
                 elif isinstance(kv_cache_spec, MambaSpec):
                     has_mamba = True
@@ -7265,7 +7223,7 @@ class GPUModelRunner(
 
         # Try creating KV caches optimized for kv-connector transfers
         cache_dtype = self.cache_config.cache_dtype
-        if self.use_uniform_kv_cache(self.attn_groups, cache_dtype):
+        if self.use_uniform_kv_cache(self.attn_groups):
             kv_caches, cross_layers_kv_cache, attn_backend = (
                 self.allocate_uniform_kv_caches(
                     kv_cache_config,
@@ -7515,6 +7473,13 @@ class GPUModelRunner(
                 continue
             # Skip modules that don't need KV cache (eg encoder-only attention)
             if spec := attn_module.get_kv_cache_spec(self.vllm_config):
+                if isinstance(spec, AttentionSpec):
+                    backend = attn_module.get_attn_backend()
+                    # indexes_kv_by_block_stride() -> get_kv_cache_stride_order()
+                    # -> get_kv_cache_layout() needs the current vLLM config.
+                    with set_current_vllm_config(self.vllm_config):
+                        indexes = backend.indexes_kv_by_block_stride()
+                    spec = replace(spec, indexes_kv_by_block_stride=indexes)
                 kv_cache_spec[layer_name] = spec
 
         return kv_cache_spec

vllm/v1/worker/kv_connector_model_runner_mixin.py

@@ -114,7 +114,6 @@ class KVConnectorModelRunnerMixin:
     @staticmethod
     def use_uniform_kv_cache(
         attn_groups: list[list[AttentionGroup]],
-        cache_dtype: CacheDType,
     ) -> bool:
         """
         Determines whether a uniform KV layout should be used.
@@ -128,9 +127,9 @@ class KVConnectorModelRunnerMixin:
             have the same page size.
         2. A KV connector is configured, and the KV connector instance prefers
             to use this layout (prefer_cross_layer_blocks() returns True)
-        2. The flash attention backend supports this layout
-            (get_kv_cache_stride_order(True) includes a placement for a
-            num_layers dimension)
+        3. The attention backend indexes KV by the block stride
+            (kv_cache_spec.indexes_kv_by_block_stride), i.e. num_blocks is the
+            outermost physical dim so per-block all-layers data is contiguous.
 
         Note that the actual placement of the num_layers dimensions
         in the unified layers tensors will be determined by the attention
@@ -140,7 +139,6 @@ class KVConnectorModelRunnerMixin:
 
         Args:
             attn_groups: The list of attention groups for this model
-            cache_dtype: The KV cache dtype
         Returns:
             True if we should use a uniform KV cache layout.
         """
@@ -157,30 +155,7 @@ class KVConnectorModelRunnerMixin:
         kv_cache_spec = attn_group.kv_cache_spec
         if not isinstance(kv_cache_spec, AttentionSpec):
             return False
-
-        attn_backend = attn_group.backend
-        kv_cache_shape = attn_backend.get_kv_cache_shape(
-            1234,
-            kv_cache_spec.block_size,
-            kv_cache_spec.num_kv_heads,
-            kv_cache_spec.head_size,
-            cache_dtype_str=cache_dtype,
-        )
-
-        try:
-            kv_cache_stride_order = attn_backend.get_kv_cache_stride_order(
-                include_num_layers_dimension=True
-            )
-        except (AttributeError, NotImplementedError):
-            return False
-
-        # check that attention backend includes a layers dimension
-        if len(kv_cache_stride_order) != len(kv_cache_shape) + 1:
-            return False
-
-        # stride_order[0] == 0 means num_layers stays first in physical
-        # layout (identity permutation), so cross-layer is unsupported.
-        return kv_cache_stride_order[0] != 0
+        return kv_cache_spec.indexes_kv_by_block_stride
 
     @staticmethod
     def allocate_uniform_kv_caches(

vllm/v1/worker/mamba_utils.py

@@ -257,9 +257,10 @@ class MambaCopyBuffers:
             for gid in mamba_group_ids
         ) * len(copy_funcs)
         n = max_num_reqs * entries_per_req
+
         return cls(
-            src_ptrs=make_buffer(n, dtype=torch.int64),
-            dst_ptrs=make_buffer(n, dtype=torch.int64),
+            src_ptrs=make_buffer(n, dtype=torch.uint64),
+            dst_ptrs=make_buffer(n, dtype=torch.uint64),
             sizes=make_buffer(n, dtype=torch.int32),
             mamba_group_ids=mamba_group_ids,
             mamba_spec=mamba_spec,