Init

Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai>

tests/distributed/test_context_parallel.py

@@ -33,7 +33,6 @@ CP_TEST_MODELS = [
     # [LANGUAGE GENERATION]
     "deepseek-ai/DeepSeek-V2-Lite-Chat",
     "Qwen/Qwen2.5-1.5B-Instruct",
-    "Qwen/Qwen3.5-0.8B",  # hybrid attention model
 ]
 
 # GSM8K eval configuration
@@ -47,7 +46,6 @@ MIN_ACCURACY = {
     "deepseek-ai/DeepSeek-V2-Lite-Chat": 0.64,
     # .buildkite/lm-eval-harness/configs/Qwen2.5-1.5B-Instruct.yaml
     "Qwen/Qwen2.5-1.5B-Instruct": 0.52,
-    "Qwen/Qwen3.5-0.8B": 0.33,
 }
 
 
@@ -153,12 +151,6 @@ else:
                 cp_kv_cache_interleave_size=16, attn_backend="FLASHINFER"
             ),
         ],
-        "Qwen/Qwen3.5-0.8B": [
-            CPTestSettings.detailed(
-                cp_kv_cache_interleave_size=16,
-                attn_backend="FLASH_ATTN",
-            ),
-        ],
     }
 
 

tests/distributed/test_pynccl.py

@@ -16,7 +16,6 @@ from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator
 from vllm.distributed.device_communicators.pynccl_wrapper import NCCLLibrary
 from vllm.distributed.parallel_state import (
     ensure_model_parallel_initialized,
-    get_tp_group,
     get_world_group,
     graph_capture,
     init_distributed_environment,
@@ -200,52 +199,6 @@ def test_pynccl_all_gather():
     distributed_run(all_gather_worker_fn, 2)
 
 
-@worker_fn_wrapper
-def cuda_communicator_all_gather_dim_worker_fn():
-    with ensure_current_vllm_config():
-        ensure_model_parallel_initialized(2, 1)
-
-    tp_group = get_tp_group()
-    comm = tp_group.device_communicator
-    assert comm is not None
-
-    rank = tp_group.rank_in_group
-    world_size = tp_group.world_size
-    device = tp_group.device
-
-    shape = (2, 3, 4)
-    num_elems = 1
-    for size in shape:
-        num_elems *= size
-
-    for dim in (1, -1):
-        tensor = (
-            torch.arange(num_elems, dtype=torch.float32, device=device).reshape(shape)
-            + rank * num_elems
-        )
-        expected = torch.cat(
-            [
-                torch.arange(num_elems, dtype=torch.float32, device=device).reshape(
-                    shape
-                )
-                + r * num_elems
-                for r in range(world_size)
-            ],
-            dim=dim,
-        )
-
-        result = comm.all_gather(tensor, dim=dim)
-        torch.accelerator.synchronize()
-        torch.testing.assert_close(result, expected, rtol=1e-5, atol=1e-8)
-
-
-@pytest.mark.skipif(
-    torch.accelerator.device_count() < 2, reason="Need at least 2 GPUs to run the test."
-)
-def test_cuda_communicator_all_gather_dim_not_zero():
-    distributed_run(cuda_communicator_all_gather_dim_worker_fn, 2)
-
-
 @worker_fn_wrapper
 def all_gatherv_worker_fn():
     pynccl_comm = PyNcclCommunicator(

tests/kernels/core/test_fused_embed_norm.py

@@ -0,0 +1,75 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for the replicated-embedding fused gather/norm kernels
+(``vllm.model_executor.layers.fused_embed_norm``).
+
+The guarantee: enabling ``VLLM_REPLICATE_EMBED`` (replicated table + fused
+kernels) must not change model outputs. The gathered residual is bit-exact and
+the fused norms match the unfused reference.
+"""
+
+import pytest
+import torch
+
+from vllm.model_executor.layers.fused_embed_norm import (
+    fused_embed_eh_norm,
+    fused_embed_norm,
+)
+
+# The model-local (untouched) eh-norm the replicate path must match.
+from vllm.models.deepseek_v32.nvidia.kernels import fused_eh_norm
+from vllm.platforms import current_platform
+from vllm.utils.torch_utils import set_random_seed
+
+DTYPE = torch.bfloat16
+VOCAB, HIDDEN, NUM_TOKENS, EPS = 8192, 4096, 129, 1e-6
+
+requires_cuda = pytest.mark.skipif(
+    not current_platform.is_cuda_alike(),
+    reason="fused embed/norm Triton kernels require a CUDA/ROCm device",
+)
+
+
+def _rmsnorm(x: torch.Tensor, w: torch.Tensor, eps: float) -> torch.Tensor:
+    # Full-precision (fp32) reference RMSNorm.
+    var = x.float().pow(2).mean(dim=-1, keepdim=True)
+    return x.float() * torch.rsqrt(var + eps) * w.float()
+
+
+@requires_cuda
+@torch.inference_mode()
+def test_fused_embed_norm_matches_reference():
+    """Main-model fusion: the residual is the exact gather and the second output
+    is a correct RMSNorm. The norm matches a full-precision reference to ~2 bf16
+    ulp (rtol 1e-2) -- that gap is bf16 rounding, not the kernel."""
+    set_random_seed(13)
+    table = torch.randn(VOCAB, HIDDEN, dtype=DTYPE, device="cuda")
+    ids = torch.randint(0, VOCAB, (NUM_TOKENS,), dtype=torch.int32, device="cuda")
+    weight = torch.empty(HIDDEN, dtype=DTYPE, device="cuda").normal_(1.0, 0.1)
+
+    residual, normed = fused_embed_norm(ids, table, chain_weight=weight, eps=EPS)
+
+    embeds = table[ids.long()]
+    torch.testing.assert_close(residual, embeds, atol=0.0, rtol=0.0)
+    torch.testing.assert_close(
+        normed.float(), _rmsnorm(embeds, weight, EPS), atol=1e-3, rtol=1e-2
+    )
+
+
+@requires_cuda
+@torch.inference_mode()
+def test_fused_embed_eh_norm_matches_reference():
+    """MTP fusion (folded gather) is bit-exact vs gathering the embeds and
+    feeding the untouched model-local ``fused_eh_norm``."""
+    set_random_seed(13)
+    table = torch.randn(VOCAB, HIDDEN, dtype=DTYPE, device="cuda")
+    ids = torch.randint(0, VOCAB, (NUM_TOKENS,), dtype=torch.int32, device="cuda")
+    prev = torch.randn(NUM_TOKENS, HIDDEN, dtype=DTYPE, device="cuda")
+    enorm_w = torch.randn(HIDDEN, dtype=DTYPE, device="cuda")
+    hnorm_w = torch.randn(HIDDEN, dtype=DTYPE, device="cuda")
+    positions = torch.arange(NUM_TOKENS, device="cuda")  # includes pos 0
+
+    fused = fused_embed_eh_norm(positions, ids, table, prev, enorm_w, hnorm_w, EPS)
+    ref = fused_eh_norm(positions, table[ids.long()], prev, enorm_w, hnorm_w, EPS)
+
+    torch.testing.assert_close(fused, ref, atol=0.0, rtol=0.0)

tests/models/language/generation/test_hybrid.py

@@ -43,11 +43,6 @@ HYBRID_MODELS = [
     "tiny-random/qwen3-next-moe",
 ]
 
-HYBRID_MODELS_REQUIRING_CHUNKED_PREFILL = {
-    "LiquidAI/LFM2-1.2B",
-    "tiny-random/qwen3-next-moe",
-}
-
 FULL_CUDA_GRAPH_MODELS = [
     "ai21labs/Jamba-tiny-dev",
     "pfnet/plamo-2-1b",
@@ -97,15 +92,8 @@ def test_models(
             example_prompts, max_tokens, num_logprobs
         )
 
-    extra_kwargs = {}
-    if model in HYBRID_MODELS_REQUIRING_CHUNKED_PREFILL:
-        extra_kwargs["enable_chunked_prefill"] = True
-
     with vllm_runner(
-        model,
-        max_num_seqs=MAX_NUM_SEQS,
-        attention_backend=ATTN_BACKEND,
-        **extra_kwargs,
+        model, max_num_seqs=MAX_NUM_SEQS, attention_backend=ATTN_BACKEND
     ) as vllm_model:
         vllm_outputs = vllm_model.generate_greedy_logprobs(
             example_prompts, max_tokens, num_logprobs

tests/models/multimodal/generation/test_vit_cudagraph.py

@@ -171,7 +171,6 @@ MODEL_CONFIGS: dict[str, VitCudagraphTestConfig] = {
             "Describe this video in one sentence."
         ),
         needs_video_metadata=True,
-        vllm_runner_kwargs={"enable_chunked_prefill": True},
         marks=[pytest.mark.core_model],
     ),
     "internvl": VitCudagraphTestConfig(

tests/test_config.py

@@ -1121,14 +1121,14 @@ def test_is_chunked_prefill_supported(
         (
             "Qwen/Qwen3-Next-80B-A3B-Instruct",
             "hybrid",
-            True,
-            "Generative hybrid models support prefix caching.",  # noqa: E501
+            False,
+            "Hybrid models do not support prefix caching since the feature is still experimental.",  # noqa: E501
         ),
         (
             "ibm-granite/granite-4.0-h-small",
             "hybrid",
-            True,
-            "Generative hybrid models support prefix caching.",  # noqa: E501
+            False,
+            "Hybrid models do not support prefix caching since the feature is still experimental.",  # noqa: E501
         ),
         (
             "state-spaces/mamba-130m-hf",

tests/v1/core/test_prefix_caching.py

@@ -12,12 +12,7 @@ import torch
 
 import vllm.v1.core.kv_cache_manager as kv_cache_manager
 import vllm.v1.core.kv_cache_utils as kv_cache_utils
-from vllm.distributed.kv_events import (
-    MEDIUM_GPU,
-    AllBlocksCleared,
-    BlockRemoved,
-    BlockStored,
-)
+from vllm.distributed.kv_events import AllBlocksCleared, BlockRemoved, BlockStored
 from vllm.lora.request import LoRARequest
 from vllm.multimodal.inputs import (
     MultiModalFeatureSpec,
@@ -2459,118 +2454,6 @@ def test_block_removed_event_group_idx(group_id: int):
         assert event.group_idx == group_id
 
 
-def test_emit_cached_block_events():
-    """emit_cached_block_events emits one BlockStored for already-cached
-    (reused) prefix blocks, carrying the correct group_idx /
-    parent_block_hash / token_ids, and without mutating block state."""
-    block_size = 4
-    num_cached_blocks = 3
-    kv_cache_group_id = 1
-    num_tokens = block_size * 4  # 4 full blocks; reuse the first 3
-
-    pool = BlockPool(
-        num_gpu_blocks=8,
-        enable_caching=True,
-        hash_block_size=block_size,
-        enable_kv_cache_events=True,
-    )
-
-    req = make_request(
-        "req_emit_cached",
-        prompt_token_ids=list(range(num_tokens)),
-        block_size=block_size,
-        hash_fn=sha256,
-    )
-    assert len(req.block_hashes) >= num_cached_blocks
-
-    # Snapshot block state to prove emit_cached_block_events does not mutate it.
-    free_before = pool.get_num_free_blocks()
-    assert len(pool.cached_block_hash_to_block) == 0
-
-    pool.emit_cached_block_events(
-        request=req,
-        num_cached_blocks=num_cached_blocks,
-        block_size=block_size,
-        kv_cache_group_id=kv_cache_group_id,
-    )
-
-    # No block-state mutation: nothing allocated, nothing inserted into the
-    # prefix-cache map.
-    assert pool.get_num_free_blocks() == free_before
-    assert len(pool.cached_block_hash_to_block) == 0
-
-    events = pool.take_events()
-    assert len(events) == 1
-    event = events[0]
-    assert isinstance(event, BlockStored)
-
-    expected_hashes = [
-        kv_cache_utils.maybe_convert_block_hash(req.block_hashes[i])
-        for i in range(num_cached_blocks)
-    ]
-    assert event.block_hashes == expected_hashes
-    # Reused blocks start from block 0, so there is no parent block hash.
-    assert event.parent_block_hash is None
-    assert event.token_ids == list(req.all_token_ids[: num_cached_blocks * block_size])
-    assert event.group_idx == kv_cache_group_id
-    assert event.block_size == block_size
-    assert event.medium == MEDIUM_GPU
-    assert event.lora_id is None
-    assert event.lora_name is None
-
-
-def test_emit_cached_block_events_disabled():
-    """No events are emitted when enable_kv_cache_events is False."""
-    block_size = 4
-    pool = BlockPool(
-        num_gpu_blocks=8,
-        enable_caching=True,
-        hash_block_size=block_size,
-        enable_kv_cache_events=False,
-    )
-    req = make_request(
-        "req_emit_disabled",
-        prompt_token_ids=list(range(block_size * 4)),
-        block_size=block_size,
-        hash_fn=sha256,
-    )
-
-    pool.emit_cached_block_events(
-        request=req,
-        num_cached_blocks=3,
-        block_size=block_size,
-        kv_cache_group_id=0,
-    )
-
-    assert pool.take_events() == []
-
-
-def test_emit_cached_block_events_zero_cached():
-    """No events are emitted when num_cached_blocks == 0."""
-    block_size = 4
-    pool = BlockPool(
-        num_gpu_blocks=8,
-        enable_caching=True,
-        hash_block_size=block_size,
-        enable_kv_cache_events=True,
-    )
-    req = make_request(
-        "req_emit_zero",
-        prompt_token_ids=list(range(block_size * 4)),
-        block_size=block_size,
-        hash_fn=sha256,
-    )
-
-    pool.emit_cached_block_events(
-        request=req,
-        num_cached_blocks=0,
-        block_size=block_size,
-        kv_cache_group_id=0,
-    )
-
-    assert pool.take_events() == []
-
-
 def test_eagle_enabled_removes_last_block():
     """Verify Eagle does NOT remove blocks when request
     length is divisible by block size."""

tests/v1/streaming_input/test_gpu_model_runner_streaming.py

@@ -38,7 +38,6 @@ def mock_model_runner_with_input_batch():
         vocab_size=32000,
         block_sizes=[16],
         kernel_block_sizes=[16],
-        max_num_blocks_per_req=[64],
         logitsprocs=None,
         is_pooling_model=False,
     )

tests/v1/worker/test_cp_utils.py

@@ -1,45 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-import pytest
-import torch
-
-from vllm.v1.attention.backends.utils import get_dcp_local_seq_lens
-from vllm.v1.worker.cp_utils import should_skip_dcp_context_attention
-
-
-def test_skip_gate_only_for_zero_context():
-    assert should_skip_dcp_context_attention(torch.zeros(3, dtype=torch.int32))
-    assert not should_skip_dcp_context_attention(
-        torch.tensor([0, 5, 0], dtype=torch.int32)
-    )
-
-
-@pytest.mark.parametrize(
-    "dcp_world_size,interleave_size,context_len",
-    [(2, 16, 10), (4, 16, 10), (8, 16, 10), (4, 1, 2)],
-)
-def test_skip_gate_rank_invariant_with_divergent_local_context(
-    dcp_world_size: int, interleave_size: int, context_len: int
-):
-    """Contexts shorter than a full interleave round land entirely on a
-    subset of DCP ranks, so the per-rank local context lengths diverge:
-    some ranks hold zero local context while others hold all of it. Ranks
-    with zero local context must still take the collective (non-skip) path,
-    otherwise the query all-gather in _forward_with_dcp deadlocks across
-    ranks. The skip gate must therefore depend only on the rank-invariant
-    global context lengths, never on get_dcp_local_seq_lens output.
-    """
-    context_kv_lens = torch.tensor([context_len], dtype=torch.int32)
-    local_maxes = [
-        int(
-            get_dcp_local_seq_lens(
-                context_kv_lens, dcp_world_size, rank, interleave_size
-            ).max()
-        )
-        for rank in range(dcp_world_size)
-    ]
-    # Precondition: the local view diverges across ranks.
-    assert 0 in local_maxes
-    assert max(local_maxes) > 0
-    # The batch still has context globally, so no rank may skip.
-    assert not should_skip_dcp_context_attention(context_kv_lens)

tests/v1/worker/test_gpu_input_batch.py

@@ -238,7 +238,6 @@ def test_sampling_metadata_in_input_batch(device: str, batch_size: int):
         vocab_size=1024,
         block_sizes=[1],
         kernel_block_sizes=[1],
-        max_num_blocks_per_req=[1024],
     )
     reqs: list[CachedRequestState] = []
     req_id_reqs = {}
@@ -333,7 +332,6 @@ def test_swap_states_in_input_batch(device: str, batch_size: int, swap_list: lis
         vocab_size=1024,
         block_sizes=[1],
         kernel_block_sizes=[1],
-        max_num_blocks_per_req=[1024],
     )
     ref_input_batch: InputBatch = InputBatch(
         max_num_reqs=batch_size,
@@ -343,7 +341,6 @@ def test_swap_states_in_input_batch(device: str, batch_size: int, swap_list: lis
         vocab_size=1024,
         block_sizes=[1],
         kernel_block_sizes=[1],
-        max_num_blocks_per_req=[1024],
     )
 
     reqs: list[CachedRequestState] = []
@@ -412,7 +409,6 @@ def test_pooling_prompt_lens_not_aliased(device: str):
         vocab_size=VOCAB_SIZE,
         block_sizes=[16],
         kernel_block_sizes=[16],
-        max_num_blocks_per_req=[64],
         is_pooling_model=True,
     )
 
@@ -448,7 +444,6 @@ def test_placeholder_spec_token_ids_written_verbatim():
         vocab_size=VOCAB_SIZE,
         block_sizes=[16],
         kernel_block_sizes=[16],
-        max_num_blocks_per_req=[1],
     )
     req = CachedRequestState(
         req_id="req",
@@ -496,7 +491,6 @@ def test_pooling_metadata_token_id_buffers(
         vocab_size=VOCAB_SIZE,
         block_sizes=[16],
         kernel_block_sizes=[16],
-        max_num_blocks_per_req=[64],
         is_pooling_model=True,
     )
     req = _construct_pooling_request(0, PoolingParams(**pooling_params))

tests/v1/worker/test_gpu_model_runner.py

@@ -89,7 +89,6 @@ def initialize_kv_cache(runner: GPUModelRunner):
         kernel_block_sizes=[
             kv_cache_config.kv_cache_groups[0].kv_cache_spec.block_size
         ],
-        max_num_blocks_per_req=[NUM_BLOCKS],
     )
     runner.initialize_attn_backend(kv_cache_config)
 
@@ -1398,7 +1397,6 @@ def test_input_batch_with_kernel_block_sizes():
         vocab_size=vocab_size,
         block_sizes=block_sizes,
         kernel_block_sizes=kernel_block_sizes,
-        max_num_blocks_per_req=[16, 8],
     )
 
     # Verify that block tables were created with kernel block sizes
@@ -1459,7 +1457,6 @@ def test_hybrid_cache_integration(default_vllm_config, dist_init):
         vocab_size=runner.model_config.get_vocab_size(),
         block_sizes=[kv_cache_config.kv_cache_groups[0].kv_cache_spec.block_size],
         kernel_block_sizes=[16],
-        max_num_blocks_per_req=[NUM_BLOCKS],
     )  # Use kernel block size
 
     runner.initialize_attn_backend(kv_cache_config)

vllm/config/model.py

@@ -1881,8 +1881,11 @@ class ModelConfig:
         else:
             # for generative models
             if attn_type == "hybrid":
-                logger.debug("Generative hybrid models support prefix caching.")
-                return True
+                logger.debug(
+                    "Hybrid models do not support prefix caching since the feature "
+                    "is still experimental."
+                )
+                return False
             elif attn_type == "attention_free":
                 logger.debug(
                     "Attention free models do not support prefix caching since the "

vllm/distributed/device_communicators/cuda_communicator.py

@@ -341,30 +341,13 @@ class CudaCommunicator(DeviceCommunicatorBase):
     def all_gather(self, input_: torch.Tensor, dim: int = -1) -> torch.Tensor:
         # Route uniform dim-0 all-gathers through NVLS symmetric memory when
         # enabled (mirrors reduce_scatter); otherwise fall back to the
-        # PyNccl/base-class all-gather. Sequence parallelism's
-        # gather-before-GEMM uses dim=0 with tp-aligned (uniform) shards.
+        # base-class ring all-gather. Sequence parallelism's gather-before-GEMM
+        # uses dim=0 with tp-aligned (uniform) shards.
         if dim < 0:
             dim += input_.dim()
         if dim == 0 and should_nccl_symm_mem_ag_rs():
             return self._all_gather_symm_mem(input_.contiguous())
-
-        pynccl_comm = self.pynccl_comm
-        if pynccl_comm is None or pynccl_comm.disabled:
-            return super().all_gather(input_, dim)
-
-        input_size = input_.size()
-        output_size = (input_size[0] * self.world_size,) + input_size[1:]
-        output_tensor = torch.empty(
-            output_size, dtype=input_.dtype, device=input_.device
-        )
-        pynccl_comm.all_gather(output_tensor, input_.contiguous())
-        output_tensor = output_tensor.reshape((self.world_size,) + input_size)
-        output_tensor = output_tensor.movedim(0, dim)
-        return output_tensor.reshape(
-            input_size[:dim]
-            + (self.world_size * input_size[dim],)
-            + input_size[dim + 1 :]
-        )
+        return super().all_gather(input_, dim)
 
     def reduce_scatter(self, input_: torch.Tensor, dim: int = -1):
         world_size = self.world_size

vllm/engine/arg_utils.py

@@ -2493,11 +2493,7 @@ class EngineArgs:
         self, model_config: ModelConfig
     ) -> None:
         default_chunked_prefill = model_config.is_chunked_prefill_supported
-        # Hybrid models support prefix caching but keep it opt-in for now
-        # while the feature matures.
-        default_prefix_caching = (
-            model_config.is_prefix_caching_supported and not model_config.is_hybrid
-        )
+        default_prefix_caching = model_config.is_prefix_caching_supported
 
         if self.enable_chunked_prefill is None:
             self.enable_chunked_prefill = default_chunked_prefill

vllm/envs.py

@@ -147,6 +147,7 @@ if TYPE_CHECKING:
     VLLM_ENABLE_V1_MULTIPROCESSING: bool = True
     VLLM_LOG_BATCHSIZE_INTERVAL: float = -1
     VLLM_DISABLE_COMPILE_CACHE: bool = False
+    VLLM_REPLICATE_EMBED: bool = False
     VLLM_USE_LAYERNAME: bool = True
     Q_SCALE_CONSTANT: int = 200
     K_SCALE_CONSTANT: int = 200
@@ -583,6 +584,9 @@ environment_variables: dict[str, Callable[[], Any]] = {
     # Enable batch-invariant mode: deterministic results regardless of
     # batch composition. Requires NVIDIA GPU with compute capability >= 9.0.
     "VLLM_BATCH_INVARIANT": lambda: bool(int(os.getenv("VLLM_BATCH_INVARIANT", "0"))),
+    "VLLM_REPLICATE_EMBED": lambda: (
+        os.getenv("VLLM_REPLICATE_EMBED", "0").strip().lower() in ("1", "true")
+    ),
     # Use tensor descriptors for Q/K/V loads and output stores in the
     # Triton unified-attention kernel.  Enables HW 2D block reads on
     # Intel Xe2/Xe3; the non-TD branch is dead-code-eliminated at Triton

vllm/model_executor/layers/attention_layer_base.py

@@ -19,7 +19,6 @@ class AttentionLayerBase(ABC):
     """
 
     impl: "AttentionImpl"
-    supports_dcp: bool = True
 
     @abstractmethod
     def get_attn_backend(self) -> type[AttentionBackend]:

vllm/model_executor/layers/fused_embed_norm.py

@@ -0,0 +1,244 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Replicated input embedding + its fused gather/norm kernels.
+
+Groups the ``VLLM_REPLICATE_EMBED`` path in one place: the replicated embedding
+module, its factory (with the fallback to vocab-parallel), and the two Triton
+fusions the full on-rank table unlocks --
+
+  * ``fused_embed_norm``: gather + a chained RMSNorm (e.g. the first decoder
+    layer's ``input_layernorm``), and
+  * ``fused_embed_eh_norm``: gather + pos-0 zeroing + enorm/hnorm + cat, the
+    embed/previous-hidden input norm for a speculative (MTP/eagle) depth layer
+    (the replicated-table analogue of the model-local ``fused_eh_norm``, which
+    takes precomputed embeds).
+
+Self-contained (no model-local imports) so it can live under ``layers/``.
+"""
+
+import torch
+
+import vllm.envs as envs
+from vllm.model_executor.layers.vocab_parallel_embedding import (
+    VocabParallelEmbedding,
+)
+from vllm.triton_utils import tl, triton
+
+
+@triton.jit
+def _rms_norm(x, w, eps, HIDDEN_SIZE: tl.constexpr):
+    x = x.to(tl.float32)
+    mean_sq = tl.sum(x * x, axis=0) / HIDDEN_SIZE
+    rrms = tl.rsqrt(mean_sq + eps)
+    w = w.to(tl.float32)
+    return (x * rrms) * w
+
+
+class ReplicatedEmbedding(torch.nn.Embedding):
+    """Fully-replicated input token embedding for GLM-5.2 / DeepSeek-V32.
+
+    The full [num_embeddings, embedding_dim] table lives on every TP rank and
+    the forward is a local lookup with NO all-reduce (unlike
+    VocabParallelEmbedding, which shards the vocab and all-reduces the output).
+    The full on-rank table also enables the fused gather+norm kernels
+    (``fused_embed_norm`` / ``fused_eh_norm`` GATHER mode). The weight has no
+    ``weight_loader`` attr, so it loads via ``default_weight_loader`` (a
+    shape-checked full-tensor copy). The only addition over
+    ``torch.nn.Embedding`` is the int32->int64 index cast that ``F.embedding``
+    requires (vLLM feeds int32 input_ids).
+    """
+
+    def forward(self, input_: torch.Tensor) -> torch.Tensor:
+        return super().forward(input_.long())
+
+
+def make_input_embedding(
+    num_embeddings: int,
+    embedding_dim: int,
+    *,
+    params_dtype: torch.dtype | None = None,
+    quant_config=None,
+    prefix: str = "",
+    tie_word_embeddings: bool = False,
+):
+    """Input token embedding with an optional replicated escape hatch.
+
+    With ``VLLM_REPLICATE_EMBED=1`` use a fully-replicated ``ReplicatedEmbedding``
+    to unlock the fused gather+norm path (and, at TP>1, skip the embedding
+    all-reduce), at the cost of a full table per rank at TP>1 (no extra memory at
+    TP=1, where vocab-parallel is already unsharded). The replicated table is
+    always the raw (unquantized) ``params_dtype``; ``quant_config`` is accepted
+    only for the vocab-parallel fallback (embeddings are unquantized regardless).
+    Assumes an untied embedding -- a replicated, unsharded table cannot be tied
+    to a vocab-parallel ``ParallelLMHead``, so tied word embeddings are rejected.
+    Otherwise falls back to the default TP-sharded ``VocabParallelEmbedding``
+    (byte-identical to before).
+    """
+    if envs.VLLM_REPLICATE_EMBED:
+        assert not tie_word_embeddings, (
+            "VLLM_REPLICATE_EMBED is unsupported with tied word embeddings "
+            "(the replicated table cannot tie to a vocab-parallel lm_head)"
+        )
+        emb = ReplicatedEmbedding(
+            num_embeddings,
+            embedding_dim,
+            dtype=params_dtype or torch.get_default_dtype(),
+        )
+        emb.weight.requires_grad_(False)
+        return emb
+    return VocabParallelEmbedding(
+        num_embeddings,
+        embedding_dim,
+        params_dtype=params_dtype,
+        quant_config=quant_config,
+        prefix=prefix,
+    )
+
+
+@triton.jit
+def _fused_embed_norm_kernel(
+    ids_ptr,  # [T] token ids
+    table_ptr,  # [V, H] embedding table (full vocab, replicated on-rank)
+    table_stride_0,
+    out_ptr,  # [T, H] gathered embedding (the residual stream)
+    normed_ptr,  # [T, H] rmsnorm(out, chain_w) (HAS_NORM only)
+    chain_w_ptr,  # [H] next norm weight (HAS_NORM only)
+    eps,
+    H: tl.constexpr,
+    BLOCK: tl.constexpr,
+    HAS_NORM: tl.constexpr,
+):
+    tok = tl.program_id(0).to(tl.int64)
+    off = tl.arange(0, BLOCK)
+    mask = off < H
+    row = tl.load(ids_ptr + tok).to(tl.int64)
+    x = tl.load(table_ptr + row * table_stride_0 + off, mask=mask, other=0.0)
+    tl.store(out_ptr + tok * H + off, x, mask=mask)
+    if HAS_NORM:
+        w = tl.load(chain_w_ptr + off, mask=mask)
+        y = _rms_norm(x, w, eps, H).to(normed_ptr.dtype.element_ty)
+        tl.store(normed_ptr + tok * H + off, y, mask=mask)
+
+
+# Base model fusion
+def fused_embed_norm(
+    input_ids: torch.Tensor,
+    embed_table: torch.Tensor,
+    chain_weight: torch.Tensor | None = None,
+    eps: float = 0.0,
+) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+    """Fused embedding row gather (``embed_table[input_ids]``).
+
+    Requires the full vocab on-rank (replicated embedding). When
+    ``chain_weight`` is given, also emits ``rmsnorm(gathered, chain_weight)``
+    (the first decoder layer's ``input_layernorm``) as a second output in the
+    same launch, so the returned pair is ``(residual, normed_input)``. Bit-exact
+    vs a plain gather followed by an ``RMSNorm``.
+    """
+    assert embed_table.ndim == 2, embed_table.shape
+    ids = input_ids.view(-1)
+    (t,) = ids.shape
+    h = embed_table.shape[1]
+    if chain_weight is not None:
+        assert chain_weight.shape == (h,), (chain_weight.shape, h)
+    out = torch.empty((t, h), dtype=embed_table.dtype, device=embed_table.device)
+    normed = torch.empty_like(out) if chain_weight is not None else None
+    if t > 0:
+        block = triton.next_power_of_2(h)
+        _fused_embed_norm_kernel[(t,)](
+            ids,
+            embed_table,
+            embed_table.stride(0),
+            out,
+            normed if normed is not None else out,
+            chain_weight if chain_weight is not None else embed_table,
+            eps,
+            h,
+            block,
+            HAS_NORM=chain_weight is not None,
+            num_warps=min(32, max(4, block // 512)),
+        )
+    if normed is not None:
+        return out, normed
+    return out
+
+
+@triton.jit
+def _fused_embed_eh_norm_kernel(
+    pos_ptr,
+    ids_ptr,  # [T] token ids
+    table_ptr,  # [V, H] embedding table (full vocab, replicated on-rank)
+    table_stride,
+    prev_ptr,  # [T, H] previous-step hidden
+    prev_stride,
+    enorm_w_ptr,
+    hnorm_w_ptr,
+    eps,
+    out_ptr,  # [T, 2H]
+    out_stride,
+    H: tl.constexpr,
+    BLOCK: tl.constexpr,
+):
+    """MTP input fusion with a folded embedding gather: gather
+    ``table[ids]``, zero it at position 0, RMSNorm(embed) with enorm and
+    RMSNorm(prev_hidden) with hnorm, written side-by-side into ``out`` ([N, 2H])
+    ready for the eh_proj GEMM. Replaces embedding lookup + where + 2x RMSNorm +
+    cat. Requires the full table on-rank (replicated embedding)."""
+    tok = tl.program_id(0)
+    off = tl.arange(0, BLOCK)
+    mask = off < H
+
+    pos = tl.load(pos_ptr + tok)
+    row = tl.load(ids_ptr + tok).to(tl.int64)
+    e = tl.load(table_ptr + row * table_stride + off, mask=mask, other=0.0)
+    e = tl.where(pos == 0, 0.0, e.to(tl.float32))
+    ew = tl.load(enorm_w_ptr + off, mask=mask)
+    e_normed = _rms_norm(e, ew, eps, H)
+    tl.store(out_ptr + tok * out_stride + off, e_normed, mask=mask)
+
+    p = tl.load(prev_ptr + tok * prev_stride + off, mask=mask, other=0.0)
+    hw = tl.load(hnorm_w_ptr + off, mask=mask)
+    p_normed = _rms_norm(p, hw, eps, H)
+    tl.store(out_ptr + tok * out_stride + H + off, p_normed, mask=mask)
+
+
+# MTP fusion
+def fused_embed_eh_norm(
+    positions: torch.Tensor,
+    input_ids: torch.Tensor,
+    embed_table: torch.Tensor,
+    previous_hidden: torch.Tensor,
+    enorm_w: torch.Tensor,
+    hnorm_w: torch.Tensor,
+    eps: float,
+) -> torch.Tensor:
+    """Fused ``cat([enorm(masked embed_table[ids]), hnorm(prev_hidden)])`` -> [N, 2H].
+
+    Folds the embedding row gather into the MTP eh-norm launch; requires the full
+    table on-rank (replicated embedding). Bit-exact vs gathering ``embed_table[
+    input_ids]`` and passing it to the model-local ``fused_eh_norm``.
+    """
+    assert previous_hidden.ndim == 2 and embed_table.ndim == 2
+    n, h = previous_hidden.shape
+    assert positions.shape == (n,) and input_ids.view(-1).shape == (n,)
+    assert embed_table.shape[1] == h, (embed_table.shape, h)
+    assert enorm_w.shape == (h,) and hnorm_w.shape == (h,)
+    out = torch.empty(
+        n, 2 * h, dtype=previous_hidden.dtype, device=previous_hidden.device
+    )
+    _fused_embed_eh_norm_kernel[(n,)](
+        positions,
+        input_ids,
+        embed_table,
+        embed_table.stride(0),
+        previous_hidden,
+        previous_hidden.stride(0),
+        enorm_w,
+        hnorm_w,
+        eps,
+        out,
+        out.stride(0),
+        h,
+        triton.next_power_of_2(h),
+    )
+    return out

vllm/model_executor/layers/mamba/abstract.py

@@ -22,7 +22,6 @@ class MambaBase(AttentionLayerBase):
     # Contains the KV cache (mamba state) for the layer
     # in the shape specified by `self.get_state_shape`.
     kv_cache: tuple[torch.Tensor, ...]
-    supports_dcp: bool = False
 
     @abstractmethod
     def get_state_shape(self) -> Iterable[tuple[int, ...]]:

vllm/models/deepseek_v32/nvidia/model.py

@@ -8,13 +8,15 @@ import torch
 
 from vllm.config import VllmConfig
 from vllm.distributed import get_pp_group
+from vllm.model_executor.layers.fused_embed_norm import (
+    ReplicatedEmbedding,
+    fused_embed_norm,
+    make_input_embedding,
+)
 from vllm.model_executor.layers.fused_moe import (
     fused_moe_make_expert_params_mapping,
 )
 from vllm.model_executor.layers.layernorm import RMSNorm
-from vllm.model_executor.layers.vocab_parallel_embedding import (
-    VocabParallelEmbedding,
-)
 from vllm.model_executor.model_loader.weight_utils import (
     default_weight_loader,
     maybe_remap_kv_scale_name,
@@ -103,11 +105,16 @@ class DeepseekV32DecoderLayer(torch.nn.Module):
         positions: torch.Tensor,
         hidden_states: torch.Tensor,
         residual: torch.Tensor | None,
+        attn_in: torch.Tensor | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor]:
         if residual is None:
-            # First layer: hidden_states is the (already reduced) embedding.
+            # First layer: hidden_states is the embedding (the residual).
             residual = hidden_states
-            hidden_states = self.input_layernorm(hidden_states)
+            # ``attn_in`` is input_layernorm(embedding) already computed fused
+            # with the embedding gather; otherwise apply it here.
+            hidden_states = (
+                attn_in if attn_in is not None else self.input_layernorm(hidden_states)
+            )
         else:
             # The previous layer's MLP/MoE output is left un-reduced; fuse its
             # all-reduce into this input_layernorm.
@@ -151,14 +158,17 @@ class DeepseekV32Model(torch.nn.Module):
         )
 
         if get_pp_group().is_first_rank:
-            self.embed_tokens = VocabParallelEmbedding(
+            self.embed_tokens = make_input_embedding(
                 config.vocab_size,
                 config.hidden_size,
                 quant_config=quant_config,
                 prefix=f"{prefix}.embed_tokens",
+                tie_word_embeddings=getattr(config, "tie_word_embeddings", False),
             )
         else:
             self.embed_tokens = PPMissingLayer()
+        # The fused embed+norm gather needs the full table on-rank.
+        self.replicated_embed = isinstance(self.embed_tokens, ReplicatedEmbedding)
 
         self.start_layer, self.end_layer, self.layers = make_layers(
             config.num_hidden_layers,
@@ -193,9 +203,21 @@ class DeepseekV32Model(torch.nn.Module):
         intermediate_tensors: IntermediateTensors | None = None,
         inputs_embeds: torch.Tensor | None = None,
     ) -> torch.Tensor | IntermediateTensors:
+        attn_in = None
         if get_pp_group().is_first_rank:
             if inputs_embeds is not None:
                 hidden_states = inputs_embeds
+            elif self.replicated_embed:
+                assert input_ids is not None
+                # Full table on-rank: gather the embedding and the first layer's
+                # input_layernorm in one launch. ``attn_in`` is the pre-normed
+                # attention input; ``hidden_states`` is the (residual) embedding.
+                hidden_states, attn_in = fused_embed_norm(
+                    input_ids,
+                    self.embed_tokens.weight,
+                    chain_weight=self.layers[self.start_layer].input_layernorm.weight,
+                    eps=self.config.rms_norm_eps,
+                )
             else:
                 assert input_ids is not None
                 hidden_states = self.embed_input_ids(input_ids)
@@ -212,7 +234,8 @@ class DeepseekV32Model(torch.nn.Module):
         ):
             if idx in self.aux_hidden_state_layers:
                 aux_hidden_states.append(hidden_states + residual)
-            hidden_states, residual = layer(positions, hidden_states, residual)
+            hidden_states, residual = layer(positions, hidden_states, residual, attn_in)
+            attn_in = None
 
         if not get_pp_group().is_last_rank:
             return IntermediateTensors(

vllm/models/deepseek_v32/nvidia/mtp.py

@@ -9,14 +9,16 @@ import torch.nn as nn
 from vllm._aiter_ops import rocm_aiter_ops
 from vllm.config import VllmConfig
 from vllm.distributed import tensor_model_parallel_all_reduce
+from vllm.model_executor.layers.fused_embed_norm import (
+    ReplicatedEmbedding,
+    fused_embed_eh_norm,
+    make_input_embedding,
+)
 from vllm.model_executor.layers.fused_moe import (
     fused_moe_make_expert_params_mapping,
 )
 from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.layers.logits_processor import LogitsProcessor
-from vllm.model_executor.layers.vocab_parallel_embedding import (
-    VocabParallelEmbedding,
-)
 from vllm.model_executor.model_loader.weight_utils import (
     default_weight_loader,
     maybe_remap_kv_scale_name,
@@ -73,18 +75,33 @@ class DeepseekV32MultiTokenPredictorLayer(nn.Module):
         positions: torch.Tensor,
         previous_hidden_states: torch.Tensor,
         inputs_embeds: torch.Tensor | None = None,
+        embed_table: torch.Tensor | None = None,
         spec_step_index: int = 0,
     ) -> torch.Tensor:
-        assert inputs_embeds is not None
-        # Fused: zero pos-0 embeds + enorm(embeds) + hnorm(prev) + cat -> [N, 2H].
-        eh_input = fused_eh_norm(
-            positions,
-            inputs_embeds,
-            previous_hidden_states,
-            self.enorm.weight,
-            self.hnorm.weight,
-            self.enorm.variance_epsilon,
-        )
+        # Fused zero pos-0 + enorm(embeds) + hnorm(prev) + cat -> [N, 2H]. With a
+        # replicated table the caller passes ``embed_table`` so the embedding
+        # lookup is folded in too (fused_embed_eh_norm); otherwise the embeds are
+        # precomputed and go through the model-local fused_eh_norm.
+        if embed_table is not None:
+            eh_input = fused_embed_eh_norm(
+                positions,
+                input_ids,
+                embed_table,
+                previous_hidden_states,
+                self.enorm.weight,
+                self.hnorm.weight,
+                self.enorm.variance_epsilon,
+            )
+        else:
+            assert inputs_embeds is not None
+            eh_input = fused_eh_norm(
+                positions,
+                inputs_embeds,
+                previous_hidden_states,
+                self.enorm.weight,
+                self.hnorm.weight,
+                self.enorm.variance_epsilon,
+            )
         hidden_states = self.eh_proj(eh_input)
         hidden_states, residual = self.mtp_block(
             positions=positions, hidden_states=hidden_states, residual=None
@@ -124,11 +141,15 @@ class DeepseekV32MultiTokenPredictor(nn.Module):
                 )
             }
         )
-        self.embed_tokens = VocabParallelEmbedding(
+        self.embed_tokens = make_input_embedding(
             config.vocab_size,
             config.hidden_size,
+            quant_config=vllm_config.quant_config,
             prefix=maybe_prefix(prefix, "embed_tokens"),
+            tie_word_embeddings=getattr(config, "tie_word_embeddings", False),
         )
+        # A replicated table lets the eh_norm fusion fold in the embedding gather.
+        self.replicated_embed = isinstance(self.embed_tokens, ReplicatedEmbedding)
         self.logits_processor = LogitsProcessor(config.vocab_size)
 
     def set_skip_topk(self, skip: bool):
@@ -158,14 +179,21 @@ class DeepseekV32MultiTokenPredictor(nn.Module):
         inputs_embeds: torch.Tensor | None = None,
         spec_step_idx: int = 0,
     ) -> torch.Tensor:
+        # With a replicated table, defer the embedding gather to fused_eh_norm
+        # (folded into the enorm/hnorm/cat launch); otherwise gather it here.
+        embed_table = None
         if inputs_embeds is None:
-            inputs_embeds = self.embed_tokens(input_ids)
+            if self.replicated_embed:
+                embed_table = self.embed_tokens.weight
+            else:
+                inputs_embeds = self.embed_tokens(input_ids)
         current_step_idx = spec_step_idx % self.num_mtp_layers
         return self.layers[str(self.mtp_start_layer_idx + current_step_idx)](
             input_ids,
             positions,
             previous_hidden_states,
             inputs_embeds,
+            embed_table,
             current_step_idx,
         )
 

vllm/utils/cpu_triton_utils.py

@@ -29,17 +29,13 @@ def _compute_slot_mapping_kernel_impl(
     block_table_stride: int,  # max_num_blocks_per_req
     block_size: int,
     slot_mapping: torch.Tensor,  # [max_num_tokens], int64
-    KV_CACHE_BLOCK_SIZE: int | None = None,
-    BLOCKS_PER_KV_BLOCK: int = 1,
-    TOTAL_CP_WORLD_SIZE: int = 1,
-    TOTAL_CP_RANK: int = 0,
-    CP_KV_CACHE_INTERLEAVE_SIZE: int = 1,
-    PAD_ID: int = -1,
-    BLOCK_SIZE: int = 1024,
+    TOTAL_CP_WORLD_SIZE: int,
+    TOTAL_CP_RANK: int,
+    CP_KV_CACHE_INTERLEAVE_SIZE: int,
+    PAD_ID: int,
+    BLOCK_SIZE: int,
 ) -> None:
     assert TOTAL_CP_WORLD_SIZE == 1, "Context Parallelism is not supported on CPU."
-    if BLOCKS_PER_KV_BLOCK != 1:
-        assert block_size * BLOCKS_PER_KV_BLOCK == KV_CACHE_BLOCK_SIZE
     torch.ops._C.compute_slot_mapping_kernel_impl(
         query_start_loc,
         positions,

vllm/v1/attention/backend.py

@@ -796,8 +796,6 @@ class AttentionImplBase(ABC, Generic[T]):
 
     # Whether the attention impl supports Prefill Context Parallelism.
     supports_pcp: bool = False
-    # Whether the attention impl supports Decode Context Parallelism.
-    supports_dcp: bool = True
     # Whether the attention impl(or ops) supports MTP
     # when cp_kv_cache_interleave_size > 1
     supports_mtp_with_cp_non_trivial_interleave_size: bool = False

vllm/v1/attention/backends/flash_attn.py

@@ -56,14 +56,10 @@ from vllm.v1.attention.backend import (
     AttentionMetadataBuilder,
     CommonAttentionMetadata,
 )
-from vllm.v1.attention.backends.utils import get_kv_cache_layout
-from vllm.v1.kv_cache_interface import AttentionSpec
-from vllm.v1.worker.cp_utils import (
-    run_split_fa2_dcp_context_attention,
-    should_skip_dcp_context_attention,
-    should_split_fa2_dcp_context_attention,
-    split_dcp_context_queries,
+from vllm.v1.attention.backends.utils import (
+    get_kv_cache_layout,
 )
+from vllm.v1.kv_cache_interface import AttentionSpec
 
 logger = init_logger(__name__)
 
@@ -249,13 +245,6 @@ class FlashAttentionMetadata:
     max_dcp_context_kv_len: int | None = None
     dcp_context_kv_lens: torch.Tensor | None = None
 
-    # Split counts for FA2 DCP context attention. num_prefill_* tracks
-    # context-bearing extend rows; pure prefills do not attend to DCP context.
-    num_decode_reqs: int = 0
-    num_prefill_reqs: int = 0
-    num_decode_tokens: int = 0
-    num_prefill_tokens: int = 0
-
     # Optional aot scheduling
     scheduler_metadata: torch.Tensor | None = None
     prefix_scheduler_metadata: torch.Tensor | None = None
@@ -524,10 +513,6 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad
         use_cascade = common_prefix_len > 0
         max_dcp_context_kv_len = 0
         dcp_context_kv_lens = None
-        num_decode_reqs = 0
-        num_prefill_reqs = 0
-        num_decode_tokens = 0
-        num_prefill_tokens = 0
 
         cu_prefix_query_lens = None
         prefix_kv_lens = None
@@ -547,54 +532,23 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad
             self._dcp_context_kv_lens[num_reqs:] = 0
             dcp_context_kv_lens = self._dcp_context_kv_lens[:num_reqs]
 
-            skip_dcp_context_attention = False
-            if common_attn_metadata.seq_lens_cpu_upper_bound is not None:
-                query_lens_cpu = (
-                    common_attn_metadata.query_start_loc_cpu[1 : num_reqs + 1]
-                    - common_attn_metadata.query_start_loc_cpu[:num_reqs]
-                )
-                context_kv_lens_cpu = (
-                    common_attn_metadata.seq_lens_cpu_upper_bound[:num_reqs]
-                    - query_lens_cpu
-                )
-                skip_dcp_context_attention = should_skip_dcp_context_attention(
-                    context_kv_lens_cpu
-                )
-
-            if max_query_len > 1:
-                (
-                    num_decode_reqs,
-                    num_prefill_reqs,
-                    num_decode_tokens,
-                    num_prefill_tokens,
-                ) = split_dcp_context_queries(
-                    common_attn_metadata.query_start_loc_cpu,
-                    common_attn_metadata.seq_lens_cpu_upper_bound,
-                    max_query_len,
-                    num_actual_tokens,
-                )
-
             # After DCP distribution, the maximum number of tokens for any rank is
             # ceil(L / (N * I)) * I, where L is max_seq_len, N is dcp_world_size,
             # and I is cp_kv_cache_interleave_size.
             # This eliminates GPU->CPU sync while minimizing workspace over-allocation.
-            if skip_dcp_context_attention:
-                max_dcp_context_kv_len = 0
-                scheduler_metadata = None
-            else:
-                num_partitions = self.dcp_world_size * self.cp_kv_cache_interleave_size
-                max_dcp_context_kv_len = (
-                    (max_seq_len + num_partitions - 1) // num_partitions
-                ) * self.cp_kv_cache_interleave_size
+            num_partitions = self.dcp_world_size * self.cp_kv_cache_interleave_size
+            max_dcp_context_kv_len = (
+                (max_seq_len + num_partitions - 1) // num_partitions
+            ) * self.cp_kv_cache_interleave_size
 
-                scheduler_metadata = schedule(
-                    batch_size=num_reqs,
-                    cu_query_lens=query_start_loc,
-                    max_query_len=max_query_len,
-                    seqlens=dcp_context_kv_lens,
-                    max_seq_len=max_dcp_context_kv_len,
-                    causal=False,
-                )
+            scheduler_metadata = schedule(
+                batch_size=num_reqs,
+                cu_query_lens=query_start_loc,
+                max_query_len=max_query_len,
+                seqlens=dcp_context_kv_lens,
+                max_seq_len=max_dcp_context_kv_len,
+                causal=False,
+            )
         elif use_cascade:
             cu_prefix_query_lens = torch.tensor(
                 [0, num_actual_tokens], dtype=torch.int32, device=self.device
@@ -660,10 +614,6 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad
             slot_mapping=slot_mapping,
             max_dcp_context_kv_len=max_dcp_context_kv_len,
             dcp_context_kv_lens=dcp_context_kv_lens,
-            num_decode_reqs=num_decode_reqs,
-            num_prefill_reqs=num_prefill_reqs,
-            num_decode_tokens=num_decode_tokens,
-            num_prefill_tokens=num_prefill_tokens,
             use_cascade=use_cascade,
             common_prefix_len=common_prefix_len,
             scheduler_metadata=scheduler_metadata,
@@ -793,12 +743,8 @@ class FlashAttentionImpl(AttentionImpl):
         self.dcp_combine = dcp_a2a_lse_reduce if dcp_a2a else cp_lse_ag_out_rs
 
         self._dcp_dtype: torch.dtype | None = None
-        self._dcp_max_num_tokens: int = 0
         if vllm_config is not None and self.dcp_world_size > 1:
             self._dcp_dtype = vllm_config.model_config.dtype
-            self._dcp_max_num_tokens = (
-                vllm_config.scheduler_config.max_num_batched_tokens
-            )
 
     def forward(
         self,
@@ -1117,120 +1063,40 @@ class FlashAttentionImpl(AttentionImpl):
         block_table = attn_metadata.block_table
 
         query = query.contiguous()
-        if attn_metadata.max_dcp_context_kv_len == 0:
-            flash_attn_varlen_func(
-                q=query,
-                k=key,
-                v=value,
-                out=output,
-                cu_seqlens_q=cu_seqlens_q,
-                max_seqlen_q=max_seqlen_q,
-                cu_seqlens_k=cu_seqlens_q,
-                max_seqlen_k=max_seqlen_q,
-                softmax_scale=self.scale,
-                causal=attn_metadata.causal,
-                alibi_slopes=self.alibi_slopes,
-                window_size=list(self.sliding_window)
-                if self.sliding_window is not None
-                else None,
-                softcap=self.logits_soft_cap,
-                return_softmax_lse=True,
-                fa_version=self.vllm_flash_attn_version,
-                q_descale=q_descale,
-                k_descale=k_descale,
-                v_descale=v_descale,
-                num_splits=attn_metadata.max_num_splits,
-            )
-            return output
-
         query_across_dcp = get_dcp_group().all_gather(query, dim=1)
         sliding_window_size = (
             list(self.sliding_window) if self.sliding_window is not None else None
         )
         n = query_across_dcp.shape[0]
-        num_reqs = cu_seqlens_q.shape[0] - 1
-        num_decodes = attn_metadata.num_decode_reqs
-        num_context_prefills = attn_metadata.num_prefill_reqs
-        num_decode_tokens = attn_metadata.num_decode_tokens
-        num_context_prefill_tokens = attn_metadata.num_prefill_tokens
-        split_dcp_context = should_split_fa2_dcp_context_attention(
-            self.vllm_flash_attn_version,
-            max_seqlen_q,
-            num_reqs,
-            num_decodes,
-            num_context_prefills,
-        )
-        dcp_context_out_tokens = max(n, self._dcp_max_num_tokens)
-        dcp_context_out_spec = (
+        (dcp_context_out,) = current_workspace_manager().get_simultaneous(
             (
-                dcp_context_out_tokens,
-                self.num_heads * self.dcp_world_size,
-                self.head_size,
+                (n, self.num_heads * self.dcp_world_size, self.head_size),
+                self._dcp_dtype,
             ),
-            self._dcp_dtype,
         )
-        (dcp_context_out_workspace,) = current_workspace_manager().get_simultaneous(
-            dcp_context_out_spec,
+        context_attn_out, context_lse = flash_attn_varlen_func(
+            q=query_across_dcp,
+            k=key_cache,
+            v=value_cache,
+            out=dcp_context_out,
+            cu_seqlens_q=cu_seqlens_q,
+            max_seqlen_q=max_seqlen_q,
+            seqused_k=attn_metadata.dcp_context_kv_lens,
+            max_seqlen_k=attn_metadata.max_dcp_context_kv_len,
+            softmax_scale=self.scale,
+            causal=False,
+            alibi_slopes=self.alibi_slopes,
+            window_size=sliding_window_size,
+            block_table=block_table,
+            softcap=self.logits_soft_cap,
+            return_softmax_lse=True,
+            scheduler_metadata=attn_metadata.scheduler_metadata,
+            fa_version=self.vllm_flash_attn_version,
+            q_descale=q_descale,
+            k_descale=k_descale,
+            v_descale=v_descale,
+            num_splits=attn_metadata.max_num_splits,
         )
-        dcp_context_out = dcp_context_out_workspace[:n]
-
-        if split_dcp_context:
-            # TODO: Remove this DCP + FA2 mixed decode/prefill workaround once
-            # FA4 supports this Qwen3.5 shape.
-            assert attn_metadata.dcp_context_kv_lens is not None
-            assert attn_metadata.max_dcp_context_kv_len is not None
-            assert self.vllm_flash_attn_version is not None
-            context_attn_out, context_lse = run_split_fa2_dcp_context_attention(
-                flash_attn_varlen_func,
-                query_across_dcp,
-                key_cache,
-                value_cache,
-                dcp_context_out,
-                cu_seqlens_q,
-                max_seqlen_q,
-                attn_metadata.dcp_context_kv_lens,
-                attn_metadata.max_dcp_context_kv_len,
-                self.scale,
-                self.alibi_slopes,
-                sliding_window_size,
-                block_table,
-                self.logits_soft_cap,
-                self.vllm_flash_attn_version,
-                q_descale,
-                k_descale,
-                v_descale,
-                attn_metadata.max_num_splits,
-                self.num_heads,
-                self.dcp_world_size,
-                num_decodes,
-                num_context_prefills,
-                num_decode_tokens,
-                num_context_prefill_tokens,
-            )
-        else:
-            context_attn_out, context_lse = flash_attn_varlen_func(
-                q=query_across_dcp,
-                k=key_cache,
-                v=value_cache,
-                out=dcp_context_out,
-                cu_seqlens_q=cu_seqlens_q,
-                max_seqlen_q=max_seqlen_q,
-                seqused_k=attn_metadata.dcp_context_kv_lens,
-                max_seqlen_k=attn_metadata.max_dcp_context_kv_len,
-                softmax_scale=self.scale,
-                causal=False,
-                alibi_slopes=self.alibi_slopes,
-                window_size=sliding_window_size,
-                block_table=block_table,
-                softcap=self.logits_soft_cap,
-                return_softmax_lse=True,
-                scheduler_metadata=attn_metadata.scheduler_metadata,
-                fa_version=self.vllm_flash_attn_version,
-                q_descale=q_descale,
-                k_descale=k_descale,
-                v_descale=v_descale,
-                num_splits=attn_metadata.max_num_splits,
-            )
         # FA returns LSE in shape [ H, B ] but DCP combine wants [ B, H ]
         context_attn_out_cor, context_lse_cor = self.dcp_combine(
             context_attn_out,
@@ -1240,11 +1106,14 @@ class FlashAttentionImpl(AttentionImpl):
         )
         context_lse_cor = context_lse_cor.transpose(0, 1).contiguous()
 
+        (dcp_query_out,) = current_workspace_manager().get_simultaneous(
+            ((query.shape[0], self.num_heads, self.head_size), self._dcp_dtype),
+        )
         query_attn_out, query_lse = flash_attn_varlen_func(
             q=query,
             k=key,
             v=value,
-            out=output,
+            out=dcp_query_out,
             cu_seqlens_q=cu_seqlens_q,
             max_seqlen_q=max_seqlen_q,
             cu_seqlens_k=cu_seqlens_q,

vllm/v1/core/block_pool.py

@@ -14,6 +14,8 @@ from vllm.logger import init_logger
 from vllm.v1.core.kv_cache_metrics import KVCacheMetricsCollector
 from vllm.v1.core.kv_cache_utils import (
     BlockHash,
+    BlockHashList,
+    BlockHashListWithBlockSize,
     BlockHashWithGroupId,
     ExternalBlockHash,
     FreeKVCacheBlockQueue,
@@ -23,7 +25,6 @@ from vllm.v1.core.kv_cache_utils import (
     get_group_id,
     make_block_hash_with_group_id,
     maybe_convert_block_hash,
-    resolve_block_hashes,
 )
 from vllm.v1.request import Request
 
@@ -260,7 +261,17 @@ class BlockPool:
             return
         new_full_blocks = blocks[num_cached_blocks:num_full_blocks]
         assert block_mask is None or len(block_mask) == len(new_full_blocks)
-        block_hashes = resolve_block_hashes(request, self.hash_block_size, block_size)
+        if block_size == self.hash_block_size:
+            # Common case.
+            block_hashes: BlockHashList = request.block_hashes
+        else:
+            # block_size is a multiple of hash_block_size. This happens when
+            # different KV cache groups have different block sizes.
+            assert block_size % self.hash_block_size == 0
+            block_hashes = BlockHashListWithBlockSize(
+                request.block_hashes, self.hash_block_size, block_size
+            )
+        assert len(block_hashes) >= num_full_blocks
 
         new_block_hashes = block_hashes[num_cached_blocks:]
         new_hashes: list[ExternalBlockHash] | None = (
@@ -327,117 +338,23 @@ class BlockPool:
                 extra_keys_list.append(extra_keys)
 
             self.kv_event_queue.append(
-                self._build_block_stored_event(
-                    request,
+                BlockStored(
                     block_hashes=new_hashes,
                     parent_block_hash=parent_block_hash,
-                    start_token_idx=start_token_idx,
-                    end_token_idx=end_token_idx,
+                    token_ids=request.all_token_ids[start_token_idx:end_token_idx],
                     block_size=block_size,
-                    kv_cache_group_id=kv_cache_group_id,
-                    extra_keys_list=extra_keys_list,
+                    lora_id=request.lora_request.adapter_id
+                    if request.lora_request
+                    else None,
+                    medium=MEDIUM_GPU,
+                    lora_name=request.lora_request.name
+                    if request.lora_request
+                    else None,
+                    extra_keys=extra_keys_list if extra_keys_list else None,
+                    group_idx=kv_cache_group_id,
                 )
             )
 
-    def _build_block_stored_event(
-        self,
-        request: Request,
-        block_hashes: list[ExternalBlockHash] | None,
-        parent_block_hash: ExternalBlockHash | None,
-        start_token_idx: int,
-        end_token_idx: int,
-        block_size: int,
-        kv_cache_group_id: int,
-        extra_keys_list: list[tuple[Any, ...] | None],
-    ) -> BlockStored:
-        """Build a ``BlockStored`` KV event for ``request``.
-
-        Shared by ``cache_full_blocks`` (newly cached blocks) and
-        ``emit_cached_block_events`` (prefix-cache-reused blocks) so both emit
-        identical event shapes for downstream consumers.
-        """
-        return BlockStored(
-            block_hashes=block_hashes,
-            parent_block_hash=parent_block_hash,
-            token_ids=request.all_token_ids[start_token_idx:end_token_idx],
-            block_size=block_size,
-            lora_id=request.lora_request.adapter_id if request.lora_request else None,
-            medium=MEDIUM_GPU,
-            lora_name=request.lora_request.name if request.lora_request else None,
-            extra_keys=extra_keys_list if extra_keys_list else None,
-            group_idx=kv_cache_group_id,
-        )
-
-    def emit_cached_block_events(
-        self,
-        request: Request,
-        num_cached_blocks: int,
-        block_size: int,
-        kv_cache_group_id: int,
-    ) -> None:
-        """Generate BlockStored events for blocks reused from prefix cache.
-
-        Unlike cache_full_blocks(), this does NOT modify block state —
-        the blocks are already cached. It only generates events so that
-        external consumers (e.g. gateway) can learn about reused blocks.
-
-        Args:
-            request: The request whose prefix cache blocks were reused.
-            num_cached_blocks: Number of blocks that were cache hits.
-            block_size: Number of tokens per block.
-            kv_cache_group_id: The KV cache group ID.
-        """
-        if not self.enable_kv_cache_events or num_cached_blocks == 0:
-            return
-
-        block_hashes = resolve_block_hashes(request, self.hash_block_size, block_size)
-
-        # Collect external hashes and extra_keys for cached blocks.
-        cached_hashes: list[ExternalBlockHash] = []
-        extra_keys_list: list[tuple[Any, ...] | None] = []
-        curr_mm_idx = 0
-        for i in range(num_cached_blocks):
-            block_start = i * block_size
-            block_end = block_start + block_size
-            cached_hashes.append(maybe_convert_block_hash(block_hashes[i]))
-            extra_keys, curr_mm_idx = generate_block_hash_extra_keys(
-                request, block_start, block_end, curr_mm_idx
-            )
-            extra_keys_list.append(extra_keys)
-
-        if not cached_hashes:
-            return
-
-        # Prefix-cache hits always form a contiguous prefix starting at block 0,
-        # so the first (and thus the whole group's) parent block hash is None.
-        parent_block_hash: ExternalBlockHash | None = None
-        start_token_idx = 0
-        end_token_idx = num_cached_blocks * block_size
-
-        logger.debug(
-            "EmitCachedBlock event: block_size=%d, "
-            "num_cached_blocks=%d, parent_block_hash=%s, "
-            "token_ids_len=%d, group_idx=%s",
-            block_size,
-            num_cached_blocks,
-            parent_block_hash,
-            len(request.all_token_ids[start_token_idx:end_token_idx]),
-            kv_cache_group_id,
-        )
-
-        self.kv_event_queue.append(
-            self._build_block_stored_event(
-                request,
-                block_hashes=cached_hashes,
-                parent_block_hash=parent_block_hash,
-                start_token_idx=start_token_idx,
-                end_token_idx=end_token_idx,
-                block_size=block_size,
-                kv_cache_group_id=kv_cache_group_id,
-                extra_keys_list=extra_keys_list,
-            )
-        )
-
     def cache_partial_block(
         self,
         request: Request,

vllm/v1/core/kv_cache_coordinator.py

@@ -550,28 +550,12 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):
         # different KV cache groups have different block sizes, the actual block size
         # can be a multiple of hash_block_size.
         self.hash_block_size = hash_block_size
-        self.dcp_world_size = dcp_world_size
-        group_block_sizes = [
-            manager.block_size for manager in self.single_type_managers
-        ]
         assert all(
-            block_size % hash_block_size == 0 for block_size in group_block_sizes
-        ), (
-            "Each KV cache group's real block_size must be divisible by "
-            f"hash_block_size. block_sizes={group_block_sizes}, "
-            f"hash_block_size={hash_block_size}"
-        )
+            g.kv_cache_spec.block_size % hash_block_size == 0
+            for g in kv_cache_config.kv_cache_groups
+        ), "block_size must be divisible by hash_block_size"
+        assert dcp_world_size == 1, "DCP not support hybrid attn now."
         assert pcp_world_size == 1, "PCP not support hybrid attn now."
-        if dcp_world_size > 1:
-            # DCP shards full-attention KV across ranks and replicates Mamba
-            # state; other spec types (e.g. sliding window) have no DCP-aware
-            # handling yet, so reject them explicitly.
-            for g in kv_cache_config.kv_cache_groups:
-                assert isinstance(g.kv_cache_spec, (FullAttentionSpec, MambaSpec)), (
-                    "DCP with hybrid KV cache layouts only supports "
-                    "full-attention and Mamba groups, got: "
-                    f"{type(g.kv_cache_spec).__name__}."
-                )
         self.verify_and_split_kv_cache_groups()
 
     def verify_and_split_kv_cache_groups(self) -> None:
@@ -667,11 +651,11 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):
                 - The number of tokens of the longest cache hit.
         """
 
-        def _get_block_hashes(block_size: int) -> BlockHashList:
-            if block_size == self.hash_block_size:
+        def _get_block_hashes(kv_cache_spec: KVCacheSpec) -> BlockHashList:
+            if kv_cache_spec.block_size == self.hash_block_size:
                 return block_hashes
             return BlockHashListWithBlockSize(
-                block_hashes, self.hash_block_size, block_size
+                block_hashes, self.hash_block_size, kv_cache_spec.block_size
             )
 
         num_groups = len(self.kv_cache_config.kv_cache_groups)
@@ -696,14 +680,13 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):
             for idx, (spec, group_ids, manager_cls, use_eagle) in enumerate(
                 self.attention_groups
             ):
-                group_block_size = self.single_type_managers[group_ids[0]].block_size
                 cached_blocks = hit_blocks_by_group[group_ids[0]]
                 if isinstance(spec, FullAttentionSpec) and cached_blocks is not None:
                     # Full attention is downward-closed: we only need to look
                     # up cached blocks once; on subsequent iterations just trim
                     # to the (reduced) current hit length.
                     curr_hit_length = (
-                        curr_hit_length // group_block_size * group_block_size
+                        curr_hit_length // spec.block_size * spec.block_size
                     )
                     continue
 
@@ -713,23 +696,18 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):
                 if drop_eagle_block:
                     # Eagle needs to match one more block and then pop the last.
                     _max_length = min(
-                        curr_hit_length + group_block_size, max_cache_hit_length
+                        curr_hit_length + spec.block_size, max_cache_hit_length
                     )
                 hit_blocks = manager_cls.find_longest_cache_hit(
-                    block_hashes=_get_block_hashes(group_block_size),
+                    block_hashes=_get_block_hashes(spec),
                     max_length=_max_length,
                     kv_cache_group_ids=group_ids,
                     block_pool=self.block_pool,
                     kv_cache_spec=spec,
                     drop_eagle_block=drop_eagle_block,
                     alignment_tokens=self.scheduler_block_size,
-                    dcp_world_size=(
-                        self.dcp_world_size
-                        if isinstance(spec, FullAttentionSpec)
-                        else 1
-                    ),
                 )
-                _new_hit_length = len(hit_blocks[0]) * group_block_size
+                _new_hit_length = len(hit_blocks[0]) * spec.block_size
                 if drop_eagle_block:
                     eagle_verified.add(idx)
                 elif _new_hit_length < curr_hit_length:
@@ -750,10 +728,7 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):
         # Truncate full attention blocks to final hit_length (if present)
         first_group = self.attention_groups[0]
         if isinstance(first_group.spec, FullAttentionSpec):
-            group_block_size = self.single_type_managers[
-                first_group.group_ids[0]
-            ].block_size
-            num_blocks = hit_length // group_block_size
+            num_blocks = hit_length // first_group.spec.block_size
             for group_id in first_group.group_ids:
                 if (blks := hit_blocks_by_group[group_id]) is not None:
                     del blks[num_blocks:]

vllm/v1/core/kv_cache_manager.py

@@ -136,7 +136,6 @@ class KVCacheManager:
             max_in_flight_tokens = max_model_len
 
         self.enable_caching = enable_caching
-        self.enable_kv_cache_events = enable_kv_cache_events
         self.use_eagle = use_eagle
         self.log_stats = log_stats
         self.metrics_collector = metrics_collector
@@ -236,26 +235,6 @@ class KVCacheManager:
             )
         )
 
-        # When kv_cache_report_mode is "full", emit BlockStored events
-        # for the reused prefix cache blocks so that external consumers
-        # (e.g. gateway) can learn about them.
-        if (
-            num_new_computed_tokens > 0
-            and self.enable_kv_cache_events
-            and getattr(request, "kv_cache_report_mode", "incremental") == "full"
-        ):
-            for group_idx, group_blocks in enumerate(computed_blocks):
-                num_blocks = len(group_blocks)
-                if num_blocks > 0:
-                    group = self.kv_cache_config.kv_cache_groups[group_idx]
-                    block_size = group.kv_cache_spec.block_size
-                    self.block_pool.emit_cached_block_events(
-                        request,
-                        num_blocks,
-                        block_size,
-                        group_idx,
-                    )
-
         if self.log_stats:
             assert self.prefix_cache_stats is not None
             self.prefix_cache_stats.record(

vllm/v1/core/kv_cache_utils.py

@@ -627,8 +627,7 @@ def resolve_kv_cache_block_sizes(
     - ``scheduler_block_size`` is the token-alignment invariant used by the
       scheduler (e.g. for ``num_computed_tokens`` rounding). Single group:
       ``cache_config.block_size * dcp * pcp``. Multiple groups: LCM of every
-      group's effective block size. Attention groups are scaled by DCP/PCP;
-      Mamba groups keep their full per-rank state and are not scaled.
+      group's block size — context parallelism is not supported here.
     - ``hash_block_size`` is the granularity at which ``Request.block_hashes``
       is computed. Single group: equals scheduler block size. Multiple groups:
       ``cache_config.hash_block_size`` override if set, else the GCD of group
@@ -646,12 +645,13 @@ def resolve_kv_cache_block_sizes(
         bs = cache_config.block_size * dcp * pcp
         return bs, bs
 
-    group_block_sizes = [
-        g.kv_cache_spec.block_size * dcp * pcp
-        if isinstance(g.kv_cache_spec, AttentionSpec)
-        else g.kv_cache_spec.block_size
-        for g in groups
-    ]
+    if dcp != 1 or pcp != 1:
+        raise ValueError(
+            "Hybrid KV cache groups with multiple block sizes do not "
+            "support context parallelism (dcp_world_size/pcp_world_size > 1)."
+        )
+
+    group_block_sizes = [g.kv_cache_spec.block_size for g in groups]
     scheduler_block_size = math.lcm(*group_block_sizes)
 
     # Block hashes are only consumed by prefix caching and KV connectors
@@ -2249,22 +2249,3 @@ class BlockHashListWithBlockSize:
 
 
 BlockHashList = list[BlockHash] | BlockHashListWithBlockSize
-
-
-def resolve_block_hashes(
-    request: Request,
-    hash_block_size: int,
-    block_size: int,
-) -> BlockHashList:
-    """Resolve the block-hash view for ``request`` at ``block_size``.
-
-    When ``block_size`` equals ``hash_block_size``, reuse the request's
-    precomputed ``block_hashes`` directly; otherwise recalculate at
-    ``block_size`` granularity (``block_size`` must be a multiple of
-    ``hash_block_size``, which happens when KV cache groups differ in
-    block size).
-    """
-    if block_size == hash_block_size:
-        return request.block_hashes
-    assert block_size % hash_block_size == 0
-    return BlockHashListWithBlockSize(request.block_hashes, hash_block_size, block_size)

vllm/v1/core/single_type_kv_cache_manager.py

@@ -1031,10 +1031,6 @@ class MambaManager(SingleTypeKVCacheManager):
         self, kv_cache_spec: MambaSpec, block_pool: BlockPool, **kwargs
     ) -> None:
         super().__init__(kv_cache_spec, block_pool, **kwargs)
-        # Mamba layers use TP instead of DCP, so each rank holds the full
-        # recurrent state. Undo the DCP/PCP block_size scaling that the base
-        # class applies for attention groups whose KV cache is partitioned.
-        self.block_size = kv_cache_spec.block_size
         self.cached_blocks_this_step: set[BlockHashWithGroupId] = set()
         self.mamba_cache_mode = kv_cache_spec.mamba_cache_mode
         self.num_speculative_blocks: int = kv_cache_spec.num_speculative_blocks

vllm/v1/kv_cache_interface.py

@@ -128,18 +128,6 @@ class KVCacheSpec:
         """
         raise NotImplementedError
 
-    def max_num_blocks_per_req(self, vllm_config: VllmConfig, max_len: int) -> int:
-        """
-        The number of block table entries needed per request, i.e. the row
-        length of the worker-side block table for this cache group.
-
-        Args:
-            vllm_config: The vllm config.
-            max_len: The maximum sequence length to size for, including the
-                encoder length for encoder-decoder models.
-        """
-        return cdiv(max_len, self.block_size)
-
     def copy_with_new_block_size(self, block_size: int) -> Self:
         """
         Create a new KVCacheSpec from self but replacing the block size.
@@ -213,16 +201,6 @@ class AttentionSpec(KVCacheSpec):
             * get_dtype_size(self.dtype)
         )
 
-    def max_num_blocks_per_req(self, vllm_config: VllmConfig, max_len: int) -> int:
-        # Attention KV is token-interleaved across DCP/PCP ranks, so each rank
-        # only stores max_len // (dcp * pcp) tokens per request.
-        parallel_config = vllm_config.parallel_config
-        total_cp_size = (
-            parallel_config.decode_context_parallel_size
-            * parallel_config.prefill_context_parallel_size
-        )
-        return cdiv(max_len, self.block_size * total_cp_size)
-
 
 @dataclass(frozen=True, kw_only=True)
 class FullAttentionSpec(AttentionSpec):
@@ -721,18 +699,6 @@ class MambaSpec(KVCacheSpec):
         else:
             return self.page_size_bytes * (1 + self.num_speculative_blocks)
 
-    def max_num_blocks_per_req(self, vllm_config: VllmConfig, max_len: int) -> int:
-        # Mamba state is replicated across DCP/PCP ranks, never sharded, so
-        # no CP scaling applies.
-        if vllm_config.cache_config.mamba_cache_mode == "align":
-            # Block table rows are position-indexed over the full sequence
-            # even though only 2 + num_speculative_blocks state blocks are
-            # resident at a time (earlier states are nulled out by
-            # remove_skipped_blocks), so the row length must cover max_len
-            # rather than max_memory_usage_bytes.
-            return cdiv(max_len, self.block_size) + self.num_speculative_blocks
-        return cdiv(self.max_memory_usage_bytes(vllm_config), self.page_size_bytes)
-
     def is_uniform_with_collection(
         self, kv_cache_specs: dict[str, KVCacheSpec]
     ) -> bool:

vllm/v1/request.py

@@ -115,11 +115,6 @@ class Request:
                 self.kv_transfer_params = sampling_params.extra_args.get(
                     "kv_transfer_params"
                 )
-                self.kv_cache_report_mode = sampling_params.extra_args.get(
-                    "kv_cache_report_mode", "incremental"
-                )
-            else:
-                self.kv_cache_report_mode = "incremental"
         else:
             raise ValueError("sampling_params and pooling_params can't both be unset")
 

vllm/v1/worker/block_table.py

@@ -1,8 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-from enum import Enum
-
 import numpy as np
 import torch
 
@@ -12,15 +10,11 @@ from vllm.triton_utils import tl, triton
 from vllm.utils.math_utils import cdiv
 from vllm.v1.attention.backends.utils import PAD_SLOT_ID
 from vllm.v1.utils import CpuGpuBuffer
+from vllm.v1.worker.cp_utils import get_total_cp_world_size
 
 logger = init_logger(__name__)
 
 
-class SlotMappingMode(Enum):
-    TOKEN_TO_KV_SLOT = "token_to_kv_slot"
-    NONE = "none"
-
-
 class BlockTable:
     def __init__(
         self,
@@ -32,7 +26,6 @@ class BlockTable:
         device: torch.device,
         kernel_block_size: int,
         cp_kv_cache_interleave_size: int,
-        slot_mapping_mode: SlotMappingMode = SlotMappingMode.TOKEN_TO_KV_SLOT,
     ):
         """
         Args:
@@ -45,15 +38,11 @@ class BlockTable:
             kernel_block_size: The block_size of underlying attention kernel.
                 Will be the same as `block_size` if `block_size` is supported
                 by the attention kernel.
-            slot_mapping_mode: How this cache group maps scheduled tokens to
-                cache slots. Mamba-like state caches do not use token slot
-                mappings and should use SlotMappingMode.NONE.
         """
         self.max_num_reqs = max_num_reqs
         self.max_num_batched_tokens = max_num_batched_tokens
         self.pin_memory = pin_memory
         self.device = device
-        self.kv_cache_block_size = block_size
 
         if kernel_block_size == block_size:
             # Standard case: allocation and computation use same block size
@@ -109,7 +98,6 @@ class BlockTable:
             self.dcp_world_size = 1
             self.dcp_rank = 0
         self.cp_kv_cache_interleave_size = cp_kv_cache_interleave_size
-        self.slot_mapping_mode = slot_mapping_mode
 
     def append_row(
         self,
@@ -157,12 +145,6 @@ class BlockTable:
         positions: torch.Tensor,
     ) -> None:
         num_tokens = positions.shape[0]
-        if self.slot_mapping_mode == SlotMappingMode.NONE:
-            # Mamba/GDN groups consume the block table as recurrent state
-            # indices and do not use per-token slot mappings.
-            return
-        assert self.slot_mapping_mode == SlotMappingMode.TOKEN_TO_KV_SLOT
-
         total_cp_world_size = self.pcp_world_size * self.dcp_world_size
         total_cp_rank = self.pcp_rank * self.dcp_world_size + self.dcp_rank
         _compute_slot_mapping_kernel[(num_reqs + 1,)](
@@ -174,8 +156,6 @@ class BlockTable:
             self.block_table.gpu.stride(0),
             self.block_size,
             self.slot_mapping.gpu,
-            KV_CACHE_BLOCK_SIZE=self.kv_cache_block_size,
-            BLOCKS_PER_KV_BLOCK=self.blocks_per_kv_block,
             TOTAL_CP_WORLD_SIZE=total_cp_world_size,
             TOTAL_CP_RANK=total_cp_rank,
             CP_KV_CACHE_INTERLEAVE_SIZE=self.cp_kv_cache_interleave_size,
@@ -246,27 +226,30 @@ class MultiGroupBlockTable:
     def __init__(
         self,
         max_num_reqs: int,
+        max_model_len: int,
         max_num_batched_tokens: int,
         pin_memory: bool,
         device: torch.device,
         block_sizes: list[int],
         kernel_block_sizes: list[int],
-        max_num_blocks: list[int],
+        max_num_blocks: list[int] | None = None,
         cp_kv_cache_interleave_size: int = 1,
-        slot_mapping_modes: list[SlotMappingMode] | None = None,
     ) -> None:
         if len(kernel_block_sizes) != len(block_sizes):
             raise ValueError(
                 f"kernel_block_sizes length ({len(kernel_block_sizes)}) "
                 f"must match block_sizes length ({len(block_sizes)})"
             )
-        if slot_mapping_modes is None:
-            slot_mapping_modes = [SlotMappingMode.TOKEN_TO_KV_SLOT] * len(block_sizes)
-        if len(slot_mapping_modes) != len(block_sizes):
-            raise ValueError(
-                f"slot_mapping_modes length ({len(slot_mapping_modes)}) "
-                f"must match block_sizes length ({len(block_sizes)})"
-            )
+        if max_num_blocks is None:
+            # Note(hc): each dcp rank only store
+            # (max_model_len//dcp_world_size) tokens in kvcache,
+            # so the block_size which used for calc max_num_blocks_per_req
+            # must be multiplied by dcp_world_size.
+            total_cp_world_size = get_total_cp_world_size()
+            max_num_blocks = [
+                cdiv(max_model_len, block_size * total_cp_world_size)
+                for block_size in block_sizes
+            ]
 
         if len(max_num_blocks) != len(block_sizes):
             raise ValueError(
@@ -291,15 +274,9 @@ class MultiGroupBlockTable:
                 device,
                 kernel_block_size,
                 cp_kv_cache_interleave_size,
-                slot_mapping_mode=slot_mapping_mode,
             )
-            for (
-                block_size,
-                kernel_block_size,
-                max_num_blocks_per_req,
-                slot_mapping_mode,
-            ) in zip(
-                block_sizes, kernel_block_sizes, max_num_blocks, slot_mapping_modes
+            for block_size, kernel_block_size, max_num_blocks_per_req in zip(
+                block_sizes, kernel_block_sizes, max_num_blocks
             )
         ]
 
@@ -355,8 +332,6 @@ def _compute_slot_mapping_kernel(
     block_table_stride,  # max_num_blocks_per_req
     block_size,
     slot_mapping_ptr,  # [max_num_tokens], int64
-    KV_CACHE_BLOCK_SIZE: tl.constexpr,
-    BLOCKS_PER_KV_BLOCK: tl.constexpr,
     TOTAL_CP_WORLD_SIZE: tl.constexpr,
     TOTAL_CP_RANK: tl.constexpr,
     CP_KV_CACHE_INTERLEAVE_SIZE: tl.constexpr,
@@ -379,14 +354,18 @@ def _compute_slot_mapping_kernel(
     start_idx = tl.load(query_start_loc_ptr + req_idx).to(tl.int64)
     end_idx = tl.load(query_start_loc_ptr + req_idx + 1).to(tl.int64)
 
-    virtual_block_size = KV_CACHE_BLOCK_SIZE * TOTAL_CP_WORLD_SIZE
+    virtual_block_size = block_size * TOTAL_CP_WORLD_SIZE
     row_offset = req_idx * block_table_stride
     for i in range(start_idx, end_idx, BLOCK_SIZE):
         offsets = i + tl.arange(0, BLOCK_SIZE)
         mask = offsets < end_idx
         pos = tl.load(positions_ptr + offsets, mask=mask, other=0)
-        virtual_block_indices = pos // virtual_block_size
-        virtual_block_offsets = pos - virtual_block_indices * virtual_block_size
+        block_indices = pos // virtual_block_size
+        block_numbers = tl.load(block_table_ptr + row_offset + block_indices).to(
+            tl.int64
+        )
+
+        virtual_block_offsets = pos - block_indices * virtual_block_size
         is_local = (
             virtual_block_offsets // CP_KV_CACHE_INTERLEAVE_SIZE
         ) % TOTAL_CP_WORLD_SIZE == TOTAL_CP_RANK
@@ -396,16 +375,6 @@ def _compute_slot_mapping_kernel(
             virtual_block_offsets % CP_KV_CACHE_INTERLEAVE_SIZE
         )
 
-        block_indices = (
-            virtual_block_indices * BLOCKS_PER_KV_BLOCK
-            + local_block_offsets // block_size
-        )
-        block_numbers = tl.load(
-            block_table_ptr + row_offset + block_indices,
-            mask=mask & is_local,
-            other=0,
-        ).to(tl.int64)
-        slot_offsets = local_block_offsets % block_size
-        slot_ids = block_numbers * block_size + slot_offsets
+        slot_ids = block_numbers * block_size + local_block_offsets
         slot_ids = tl.where(is_local, slot_ids, PAD_ID)
         tl.store(slot_mapping_ptr + offsets, slot_ids, mask=mask)

vllm/v1/worker/cp_utils.py

@@ -1,23 +1,15 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-from types import SimpleNamespace
 from typing import TYPE_CHECKING, Any, cast
 
-import torch
-
 from vllm.config import VllmConfig, get_layers_from_vllm_config
 from vllm.distributed import get_dcp_group, get_pcp_group
-from vllm.logger import init_logger
-from vllm.v1.attention.backend import CommonAttentionMetadata
-from vllm.v1.attention.backends.utils import split_decodes_prefills_and_extends
 
 if TYPE_CHECKING:
     from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
 else:
     AttentionLayerBase = object
 
-logger = init_logger(__name__)
-
 
 def check_attention_cp_compatibility(vllm_config: VllmConfig) -> None:
     pcp_size = vllm_config.parallel_config.prefill_context_parallel_size
@@ -64,232 +56,3 @@ def get_total_cp_world_size():
         # DCP might not be initialized in testing
         dcp_world_size = 1
     return dcp_world_size * pcp_world_size
-
-
-def get_dcp_dummy_context_len(
-    dcp_world_size: int,
-    cp_kv_cache_interleave_size: int,
-    has_kv_cache_config: bool,
-    create_mixed_batch: bool,
-    is_graph_capturing: bool,
-    uniform_decode: bool,
-) -> int:
-    if (
-        dcp_world_size <= 1
-        or not has_kv_cache_config
-        or not (create_mixed_batch or (is_graph_capturing and uniform_decode))
-    ):
-        return 0
-    return dcp_world_size * cp_kv_cache_interleave_size
-
-
-def prepare_dcp_dummy_context_metadata(
-    *,
-    input_batch: Any,
-    kv_cache_config: Any,
-    query_pos: Any,
-    positions: torch.Tensor,
-    query_start_loc: Any,
-    num_reqs: int,
-    num_tokens_unpadded: int,
-    dcp_dummy_context_len: int,
-) -> None:
-    """Populate valid fake KV metadata for DCP CUDA graph warmup/capture."""
-    if dcp_dummy_context_len == 0:
-        return
-
-    # DCP graph warmup may exercise context attention, so block-table entries
-    # must point at allocated KV blocks.
-    assert kv_cache_config is not None
-    max_valid_block_id = kv_cache_config.num_blocks - 1
-    assert max_valid_block_id > 0
-    for blk_table in input_batch.block_table.block_tables:
-        max_row_blocks = (
-            blk_table.max_num_blocks_per_req // blk_table.blocks_per_kv_block
-        )
-        block_ids = [
-            (block_idx % max_valid_block_id) + 1 for block_idx in range(max_row_blocks)
-        ]
-        for req_idx in range(num_reqs):
-            blk_table.add_row(block_ids, req_idx)
-        blk_table.commit_block_table(num_reqs)
-
-    query_pos.copy_to_gpu(num_tokens_unpadded)
-    positions[:num_tokens_unpadded] = (
-        query_pos.gpu[:num_tokens_unpadded] + dcp_dummy_context_len
-    )
-    input_batch.block_table.compute_slot_mapping(
-        num_reqs,
-        query_start_loc.gpu[: num_reqs + 1],
-        positions[:num_tokens_unpadded],
-    )
-
-
-def should_skip_dcp_context_attention(context_kv_lens_cpu: torch.Tensor) -> bool:
-    """Whether DCP context attention can be skipped for this batch.
-
-    Must be computed from rank-invariant inputs only (the global context
-    lengths, NOT this rank's local share from get_dcp_local_seq_lens): the
-    non-skip path in _forward_with_dcp issues DCP collectives (query
-    all-gather + LSE combine), so every DCP rank must take the same branch.
-    A rank can hold zero local context tokens while other ranks still hold
-    context for the same batch.
-    """
-    return int(context_kv_lens_cpu.max().item()) == 0
-
-
-def split_dcp_context_queries(
-    query_start_loc: torch.Tensor,
-    seq_lens_cpu_upper_bound: torch.Tensor | None,
-    max_query_len: int,
-    num_actual_tokens: int,
-) -> tuple[int, int, int, int]:
-    """Split reordered DCP context queries into decode and extend regions."""
-    num_reqs = query_start_loc.shape[0] - 1
-    if max_query_len <= 1:
-        return num_reqs, 0, num_actual_tokens, 0
-    if seq_lens_cpu_upper_bound is None:
-        return 0, num_reqs, 0, num_actual_tokens
-
-    common_attn_metadata = cast(
-        CommonAttentionMetadata,
-        SimpleNamespace(
-            max_query_len=max_query_len,
-            num_reqs=num_reqs,
-            num_actual_tokens=num_actual_tokens,
-            query_start_loc_cpu=query_start_loc,
-            seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound,
-            is_prefilling=None,
-        ),
-    )
-    (
-        num_decodes,
-        num_extends,
-        _num_prefills,
-        num_decode_tokens,
-        num_extend_tokens,
-        _num_prefill_tokens,
-    ) = split_decodes_prefills_and_extends(common_attn_metadata)
-    return num_decodes, num_extends, num_decode_tokens, num_extend_tokens
-
-
-def should_split_fa2_dcp_context_attention(
-    fa_version: int | None,
-    max_query_len: int,
-    num_reqs: int,
-    num_decode_reqs: int,
-    num_context_prefill_reqs: int,
-) -> bool:
-    num_prefills = num_reqs - num_decode_reqs
-    # TODO: Remove this FA2-only DCP compatibility path once FA4 supports
-    # the Qwen3.5 head_size=256 shape on Blackwell and can be used here.
-    # FA2 paged-varlen context attention can fail for DCP mixed batches when
-    # decode rows, context-bearing extend rows, and zero-context pure prefill
-    # rows are submitted together.
-    return (
-        fa_version == 2
-        and max_query_len > 1
-        and num_prefills > 0
-        and (num_decode_reqs > 0 or num_context_prefill_reqs < num_prefills)
-    )
-
-
-def run_split_fa2_dcp_context_attention(
-    flash_attn_varlen_func: Any,
-    query_across_dcp: torch.Tensor,
-    key_cache: torch.Tensor,
-    value_cache: torch.Tensor,
-    dcp_context_out: torch.Tensor,
-    cu_seqlens_q: torch.Tensor,
-    max_seqlen_q: int,
-    dcp_context_kv_lens: torch.Tensor,
-    max_dcp_context_kv_len: int,
-    softmax_scale: float,
-    alibi_slopes: torch.Tensor | None,
-    sliding_window_size: list[int] | None,
-    block_table: torch.Tensor,
-    softcap: float,
-    fa_version: int,
-    q_descale: torch.Tensor | None,
-    k_descale: torch.Tensor | None,
-    v_descale: torch.Tensor | None,
-    max_num_splits: int,
-    num_heads: int,
-    dcp_world_size: int,
-    num_decode_reqs: int,
-    num_context_prefill_reqs: int,
-    num_decode_tokens: int,
-    num_context_prefill_tokens: int,
-) -> tuple[torch.Tensor, torch.Tensor]:
-    dcp_context_out.zero_()
-    context_lse = torch.full(
-        (num_heads * dcp_world_size, query_across_dcp.shape[0]),
-        -torch.inf,
-        dtype=torch.float32,
-        device=query_across_dcp.device,
-    )
-
-    if num_decode_tokens > 0:
-        _, decode_context_lse = flash_attn_varlen_func(
-            q=query_across_dcp[:num_decode_tokens],
-            k=key_cache,
-            v=value_cache,
-            out=dcp_context_out[:num_decode_tokens],
-            cu_seqlens_q=cu_seqlens_q[: num_decode_reqs + 1],
-            max_seqlen_q=1,
-            seqused_k=dcp_context_kv_lens[:num_decode_reqs],
-            max_seqlen_k=max_dcp_context_kv_len,
-            softmax_scale=softmax_scale,
-            causal=False,
-            alibi_slopes=alibi_slopes,
-            window_size=sliding_window_size,
-            block_table=block_table[:num_decode_reqs],
-            softcap=softcap,
-            return_softmax_lse=True,
-            scheduler_metadata=None,
-            fa_version=fa_version,
-            q_descale=q_descale[:num_decode_reqs] if q_descale is not None else None,
-            k_descale=k_descale[:num_decode_reqs] if k_descale is not None else None,
-            v_descale=v_descale[:num_decode_reqs] if v_descale is not None else None,
-            num_splits=max_num_splits,
-        )
-        context_lse[:, :num_decode_tokens] = decode_context_lse
-
-    if num_context_prefill_tokens > 0:
-        prefill_start = num_decode_tokens
-        prefill_end = prefill_start + num_context_prefill_tokens
-        prefill_query_start_loc = (
-            cu_seqlens_q[
-                num_decode_reqs : num_decode_reqs + num_context_prefill_reqs + 1
-            ]
-            - num_decode_tokens
-        )
-        prefill_req_slice = slice(
-            num_decode_reqs, num_decode_reqs + num_context_prefill_reqs
-        )
-        _, prefill_context_lse = flash_attn_varlen_func(
-            q=query_across_dcp[prefill_start:prefill_end],
-            k=key_cache,
-            v=value_cache,
-            out=dcp_context_out[prefill_start:prefill_end],
-            cu_seqlens_q=prefill_query_start_loc,
-            max_seqlen_q=max_seqlen_q,
-            seqused_k=dcp_context_kv_lens[prefill_req_slice],
-            max_seqlen_k=max_dcp_context_kv_len,
-            softmax_scale=softmax_scale,
-            causal=False,
-            alibi_slopes=alibi_slopes,
-            window_size=sliding_window_size,
-            block_table=block_table[prefill_req_slice],
-            softcap=softcap,
-            return_softmax_lse=True,
-            scheduler_metadata=None,
-            fa_version=fa_version,
-            q_descale=q_descale[prefill_req_slice] if q_descale is not None else None,
-            k_descale=k_descale[prefill_req_slice] if k_descale is not None else None,
-            v_descale=v_descale[prefill_req_slice] if v_descale is not None else None,
-            num_splits=max_num_splits,
-        )
-        context_lse[:, prefill_start:prefill_end] = prefill_context_lse
-
-    return dcp_context_out, context_lse

vllm/v1/worker/gpu_input_batch.py

@@ -28,7 +28,7 @@ from vllm.v1.sample.thinking_budget_state import (
     maybe_create_thinking_budget_state_holder,
 )
 from vllm.v1.utils import copy_slice
-from vllm.v1.worker.block_table import MultiGroupBlockTable, SlotMappingMode
+from vllm.v1.worker.block_table import MultiGroupBlockTable
 
 
 @dataclass
@@ -99,14 +99,13 @@ class InputBatch:
         vocab_size: int,
         block_sizes: list[int],  # The block_size of each kv cache group
         kernel_block_sizes: list[int],
-        max_num_blocks_per_req: list[int],
+        max_num_blocks_per_req: list[int] | None = None,
         logitsprocs: LogitsProcessors | None = None,
         logitsprocs_need_output_token_ids: bool = False,
         num_spec_tokens: int = 0,
         is_pooling_model: bool = False,
         cp_kv_cache_interleave_size: int = 1,
         reasoning_config: ReasoningConfig | None = None,
-        slot_mapping_modes: list[SlotMappingMode] | None = None,
     ):
         self.thinking_budget_state_holder = maybe_create_thinking_budget_state_holder(
             reasoning_config,
@@ -172,6 +171,7 @@ class InputBatch:
         # Block table.
         self.block_table = MultiGroupBlockTable(
             max_num_reqs=max_num_reqs,
+            max_model_len=max_model_len,
             max_num_batched_tokens=max_num_batched_tokens,
             pin_memory=PIN_MEMORY,
             device=device,
@@ -179,7 +179,6 @@ class InputBatch:
             kernel_block_sizes=kernel_block_sizes,
             max_num_blocks=max_num_blocks_per_req,
             cp_kv_cache_interleave_size=cp_kv_cache_interleave_size,
-            slot_mapping_modes=slot_mapping_modes,
         )
 
         # Sampling-related.

vllm/v1/worker/gpu_model_runner.py

@@ -153,12 +153,10 @@ from vllm.v1.kv_cache_interface import (
     KVCacheConfig,
     KVCacheGroupSpec,
     KVCacheSpec,
-    KVCacheSpecKind,
     KVQuantMode,
     MambaSpec,
     SlidingWindowSpec,
     UniformTypeKVCacheSpecs,
-    get_kv_cache_spec_kind,
 )
 from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry
 from vllm.v1.outputs import (
@@ -203,11 +201,9 @@ from vllm.v1.spec_decode.utils import update_num_computed_tokens_for_batch_chang
 from vllm.v1.structured_output.utils import apply_grammar_bitmask
 from vllm.v1.utils import CpuGpuBuffer, record_function_or_nullcontext
 from vllm.v1.worker import mamba_utils
-from vllm.v1.worker.block_table import SlotMappingMode
 from vllm.v1.worker.cp_utils import (
     check_attention_cp_compatibility,
-    get_dcp_dummy_context_len,
-    prepare_dcp_dummy_context_metadata,
+    get_total_cp_world_size,
 )
 from vllm.v1.worker.dp_utils import coordinate_batch_across_dp
 from vllm.v1.worker.ec_connector_model_runner_mixin import ECConnectorModelRunnerMixin
@@ -685,13 +681,8 @@ class GPUModelRunner(
         placeholder_block_size = (
             self.cache_config.block_size or CacheConfig.DEFAULT_BLOCK_SIZE
         )
-        placeholder_max_num_blocks = cdiv(
-            max(self.max_model_len, self.max_encoder_len), placeholder_block_size
-        )
         self._init_block_sizes = [placeholder_block_size]
         self._init_kernel_block_sizes = [placeholder_block_size]
-        self._init_max_num_blocks = [placeholder_max_num_blocks]
-        self._init_slot_mapping_modes = [SlotMappingMode.TOKEN_TO_KV_SLOT]
         self.input_batch = InputBatch(
             max_num_reqs=self.max_num_reqs,
             # We need to use the encoder length for encoder-decoder
@@ -702,7 +693,6 @@ class GPUModelRunner(
             vocab_size=self.model_config.get_vocab_size(),
             block_sizes=[placeholder_block_size],
             kernel_block_sizes=[placeholder_block_size],
-            max_num_blocks_per_req=[placeholder_max_num_blocks],
             num_spec_tokens=self.num_spec_tokens,
             logitsprocs=build_logitsprocs(
                 self.vllm_config,
@@ -5868,14 +5858,6 @@ class GPUModelRunner(
         num_reqs_padded = (
             batch_desc.num_reqs if batch_desc.num_reqs is not None else num_reqs
         )
-        dcp_dummy_context_len = get_dcp_dummy_context_len(
-            self.dcp_world_size,
-            self.parallel_config.cp_kv_cache_interleave_size,
-            hasattr(self, "kv_cache_config"),
-            create_mixed_batch,
-            is_graph_capturing,
-            uniform_decode,
-        )
         ubatch_slices, ubatch_slices_padded = maybe_create_ubatch_slices(
             should_ubatch,
             num_scheduled_tokens,
@@ -5918,19 +5900,10 @@ class GPUModelRunner(
                     # In the mixed batch mode (used for FI warmup), we use
                     # shorter sequence lengths to run faster.
                     # TODO(luka) better system for describing dummy batches
-                    if dcp_dummy_context_len > 0:
-                        seq_lens = torch.tensor(  # type: ignore[assignment]
-                            [1 + dcp_dummy_context_len] * num_decode_tokens
-                            + [num_prefill_tokens + dcp_dummy_context_len],
-                            dtype=torch.int,
-                        )
-                    else:
-                        seq_lens = torch.tensor(  # type: ignore[assignment]
-                            [1] * num_decode_tokens + [num_prefill_tokens + 1],
-                            dtype=torch.int,
-                        )
-                elif dcp_dummy_context_len > 0:
-                    seq_lens = max_query_len + dcp_dummy_context_len  # type: ignore[assignment]
+                    seq_lens = torch.tensor(  # type: ignore[assignment]
+                        [1] * num_decode_tokens + [num_prefill_tokens + 1],
+                        dtype=torch.int,
+                    )
                 else:
                     seq_lens = max_query_len  # type: ignore[assignment]
                 self.optimistic_seq_lens_cpu[:num_reqs] = seq_lens
@@ -5946,17 +5919,6 @@ class GPUModelRunner(
                 )
                 self.query_start_loc.copy_to_gpu()
 
-                prepare_dcp_dummy_context_metadata(
-                    input_batch=self.input_batch,
-                    kv_cache_config=getattr(self, "kv_cache_config", None),
-                    query_pos=self.query_pos,
-                    positions=self.positions,
-                    query_start_loc=self.query_start_loc,
-                    num_reqs=num_reqs,
-                    num_tokens_unpadded=num_tokens_unpadded,
-                    dcp_dummy_context_len=dcp_dummy_context_len,
-                )
-
                 # Sync block table CPU->GPU so cleared rows from
                 # remove_request() are visible to the attention metadata
                 # builder. Without this, stale block IDs from finished
@@ -7067,34 +7029,29 @@ class GPUModelRunner(
         """
         block_sizes = []
         max_num_blocks = []
-        slot_mapping_modes = []
         max_model_len = max(self.max_model_len, self.max_encoder_len)
         for kv_cache_group in kv_cache_config.kv_cache_groups:
-            kv_cache_spec = kv_cache_group.kv_cache_spec
-            kv_cache_spec_kind = get_kv_cache_spec_kind(kv_cache_spec)
-            if kv_cache_spec_kind == KVCacheSpecKind.ENCODER_ONLY_ATTENTION:
+            if isinstance(kv_cache_group.kv_cache_spec, EncoderOnlyAttentionSpec):
                 continue
-            block_size = kv_cache_spec.block_size
+            block_size = kv_cache_group.kv_cache_spec.block_size
             block_sizes.append(block_size)
-            if kv_cache_spec_kind == KVCacheSpecKind.MAMBA:
-                slot_mapping_modes.append(SlotMappingMode.NONE)
-            else:
-                slot_mapping_modes.append(SlotMappingMode.TOKEN_TO_KV_SLOT)
-            max_num_blocks_per_req = kv_cache_spec.max_num_blocks_per_req(
-                self.vllm_config, max_model_len
+            max_num_blocks_per_req = cdiv(
+                max_model_len, block_size * get_total_cp_world_size()
             )
+            if isinstance(kv_cache_group.kv_cache_spec, MambaSpec):
+                max_num_blocks_per_req = (
+                    max_num_blocks_per_req
+                    if self.cache_config.enable_prefix_caching
+                    else 1
+                ) + kv_cache_group.kv_cache_spec.num_speculative_blocks
             max_num_blocks.append(max_num_blocks_per_req)
 
         if (
             block_sizes != self._init_block_sizes
             or kernel_block_sizes != self._init_kernel_block_sizes
-            or max_num_blocks != self._init_max_num_blocks
-            or slot_mapping_modes != self._init_slot_mapping_modes
         ):
             self._init_block_sizes = block_sizes
             self._init_kernel_block_sizes = kernel_block_sizes
-            self._init_max_num_blocks = max_num_blocks
-            self._init_slot_mapping_modes = slot_mapping_modes
             self.input_batch = InputBatch(
                 max_num_reqs=self.max_num_reqs,
                 max_model_len=max_model_len,
@@ -7108,9 +7065,7 @@ class GPUModelRunner(
                 logitsprocs=self.input_batch.logitsprocs,
                 logitsprocs_need_output_token_ids=self.input_batch.logitsprocs_need_output_token_ids,
                 is_pooling_model=self.is_pooling_model,
-                cp_kv_cache_interleave_size=self.parallel_config.cp_kv_cache_interleave_size,
                 reasoning_config=self.vllm_config.reasoning_config,
-                slot_mapping_modes=slot_mapping_modes,
             )
 
         assert self._init_block_sizes == block_sizes, (

vllm/v1/worker/tpu_input_batch.py

@@ -29,7 +29,6 @@ class InputBatch:
         vocab_size: int,
         block_sizes: list[int],  # The block_size of each kv cache group
         kernel_block_sizes: list[int],
-        max_num_blocks_per_req: list[int],
     ):
         self.max_num_reqs = max_num_reqs
         self.max_model_len = max_model_len
@@ -65,12 +64,12 @@ class InputBatch:
         # Block table.
         self.block_table = MultiGroupBlockTable(
             max_num_reqs=max_num_reqs,
+            max_model_len=max_model_len,
             max_num_batched_tokens=max_num_batched_tokens,
             pin_memory=pin_memory,
             device=device,
             block_sizes=block_sizes,
             kernel_block_sizes=kernel_block_sizes,
-            max_num_blocks=max_num_blocks_per_req,
         )
 
         # Sampling-related.