[Build] Update pre-commit pip-compile hook to use cu128 torch backend

Matches the test.txt lockfile change to CUDA 12.8. Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>
[Build] Recompile test.txt lockfile with cu128 torch backend
2026-02-18 16:41:53 -05:00 · 2026-02-18 14:38:23 -05:00 · 2026-02-18 14:38:23 -05:00 · 2026-02-17 18:58:18 -08:00 · 2026-02-17 17:06:54 -08:00 · 2026-02-17 16:27:15 -08:00
16 changed files with 438 additions and 220 deletions
@@ -8,7 +8,7 @@ clean_docker_tag() {
 }

 print_usage_and_exit() {
-    echo "Usage: $0 <registry> <repo> <commit> <branch> <vllm_use_precompiled> <vllm_merge_base_commit> <cache_from> <cache_to>"
+    echo "Usage: $0 <registry> <repo> <commit> <branch> <image_tag> [<image_tag_latest>]"
    exit 1
 }

@@ -159,7 +159,7 @@ print_bake_config() {
 #################################
 print_instance_info

-if [[ $# -lt 7 ]]; then
+if [[ $# -lt 5 ]]; then
    print_usage_and_exit
 fi

@@ -168,10 +168,8 @@ REGISTRY=$1
 REPO=$2
 BUILDKITE_COMMIT=$3
 BRANCH=$4
-VLLM_USE_PRECOMPILED=0
-VLLM_MERGE_BASE_COMMIT=""
-IMAGE_TAG=$7
-IMAGE_TAG_LATEST=${8:-} # only used for main branch, optional
+IMAGE_TAG=$5
+IMAGE_TAG_LATEST=${6:-} # only used for main branch, optional

 # build config
 TARGET="test-ci"
@@ -198,8 +196,6 @@ export CACHE_FROM
 export CACHE_FROM_BASE_BRANCH
 export CACHE_FROM_MAIN
 export CACHE_TO
-export VLLM_USE_PRECOMPILED
-export VLLM_MERGE_BASE_COMMIT

 # print args
 echo "--- :mag: Arguments"
@@ -207,8 +203,6 @@ echo "REGISTRY: ${REGISTRY}"
 echo "REPO: ${REPO}"
 echo "BUILDKITE_COMMIT: ${BUILDKITE_COMMIT}"
 echo "BRANCH: ${BRANCH}"
-echo "VLLM_USE_PRECOMPILED: ${VLLM_USE_PRECOMPILED}"
-echo "VLLM_MERGE_BASE_COMMIT: ${VLLM_MERGE_BASE_COMMIT}"
 echo "IMAGE_TAG: ${IMAGE_TAG}"
 echo "IMAGE_TAG_LATEST: ${IMAGE_TAG_LATEST}"

@@ -5,8 +5,7 @@ steps:
    depends_on: []
    timeout_in_minutes: 600
    commands:
-    - if [[ "$BUILDKITE_BRANCH" != "main" ]]; then .buildkite/image_build/image_build.sh $REGISTRY $REPO $BUILDKITE_COMMIT $BRANCH $VLLM_USE_PRECOMPILED $VLLM_MERGE_BASE_COMMIT $IMAGE_TAG; fi
-    - if [[ "$BUILDKITE_BRANCH" == "main" ]]; then .buildkite/image_build/image_build.sh $REGISTRY $REPO $BUILDKITE_COMMIT $BRANCH $VLLM_USE_PRECOMPILED $VLLM_MERGE_BASE_COMMIT $IMAGE_TAG $IMAGE_TAG_LATEST; fi
+    - if [[ "$BUILDKITE_BRANCH" == "main" ]]; then .buildkite/image_build/image_build.sh $REGISTRY $REPO $BUILDKITE_COMMIT $BRANCH $IMAGE_TAG $IMAGE_TAG_LATEST; else .buildkite/image_build/image_build.sh $REGISTRY $REPO $BUILDKITE_COMMIT $BRANCH $IMAGE_TAG; fi
    retry:
      automatic:
        - exit_status: -1  # Agent was lost
@@ -31,7 +31,7 @@ steps:
        commands:
          # #NOTE: torch_cuda_arch_list is derived from upstream PyTorch build files here:
          # https://github.com/pytorch/pytorch/blob/main/.ci/aarch64_linux/aarch64_ci_build.sh#L7
-          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0' --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu22.04  --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
+          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg BUILDER_CUDA_VERSION=13.0.1 --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0' --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu22.04  --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
          - "mkdir artifacts"
          - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
          - "bash .buildkite/scripts/upload-nightly-wheels.sh manylinux_2_35"
@@ -70,7 +70,7 @@ steps:
        agents:
          queue: cpu_queue_postmerge
        commands:
-          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu22.04 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
+          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg BUILDER_CUDA_VERSION=13.0.1 --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu22.04 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
          - "mkdir artifacts"
          - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
          - "bash .buildkite/scripts/upload-nightly-wheels.sh manylinux_2_35"
@@ -123,7 +123,7 @@ steps:
          queue: cpu_queue_postmerge
        commands:
          - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
-          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu22.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130 --target vllm-openai --progress plain -f docker/Dockerfile ."
+          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg BUILDER_CUDA_VERSION=13.0.1 --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu22.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130 --target vllm-openai --progress plain -f docker/Dockerfile ."
          - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130"
          # re-tag to default image tag and push, just in case arm64 build fails
          - "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu130"
@@ -137,7 +137,7 @@ steps:
        commands:
          - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
          # compute capability 12.0 for RTX-50 series / RTX PRO 6000 Blackwell, 12.1 for DGX Spark
-          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0 12.1' --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu22.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130 --target vllm-openai --progress plain -f docker/Dockerfile ."
+          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg BUILDER_CUDA_VERSION=13.0.1 --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0 12.1' --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu22.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130 --target vllm-openai --progress plain -f docker/Dockerfile ."
          - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130"

      - block: "Build release image for x86_64 CPU"
@@ -38,7 +38,7 @@ repos:
  rev: 0.9.1
  hooks:
    - id: pip-compile
-      args: [requirements/test.in, -o, requirements/test.txt, --index-strategy, unsafe-best-match, --torch-backend, cu129, --python-platform, x86_64-manylinux_2_28, --python-version, "3.12"]
+      args: [requirements/test.in, -o, requirements/test.txt, --index-strategy, unsafe-best-match, --torch-backend, cu128, --python-platform, x86_64-manylinux_2_28, --python-version, "3.12"]
      files: ^requirements/test\.(in|txt)$
 - repo: local
  hooks:
@@ -22,7 +22,12 @@
 #   docker buildx bake -f docker/docker-bake.hcl -f docker/versions.json
 # =============================================================================

-ARG CUDA_VERSION=12.9.1
+ARG CUDA_VERSION=12.8.1
+# BUILDER_CUDA_VERSION controls the CUDA toolkit used to compile csrc/ and
+# extensions (DeepGEMM, EP kernels).  It can differ from CUDA_VERSION, which
+# is the CUDA version shipped in the final runtime image and used to select
+# the matching PyTorch wheel.
+ARG BUILDER_CUDA_VERSION=12.9.1
 ARG PYTHON_VERSION=3.12

 # By parameterizing the base images, we allow third-party to use their own
@@ -36,7 +41,7 @@ ARG PYTHON_VERSION=3.12
 # compatibility with other Linux OSes. The main reason for this is that the
 # glibc version is baked into the distro, and binaries built with one glibc
 # version are not backwards compatible with OSes that use an earlier version.
-ARG BUILD_BASE_IMAGE=nvidia/cuda:${CUDA_VERSION}-devel-ubuntu20.04
+ARG BUILD_BASE_IMAGE=nvidia/cuda:${BUILDER_CUDA_VERSION}-devel-ubuntu20.04
 # Using cuda base image with minimal dependencies necessary for JIT compilation (FlashInfer, DeepGEMM, EP kernels)
 ARG FINAL_BASE_IMAGE=nvidia/cuda:${CUDA_VERSION}-base-ubuntu22.04

@@ -92,8 +97,13 @@ ARG INSTALL_KV_CONNECTORS=false
 FROM ${BUILD_BASE_IMAGE} AS base

 ARG CUDA_VERSION
+ARG BUILDER_CUDA_VERSION
 ARG PYTHON_VERSION

+# Override the CUDA_VERSION env var inherited from the nvidia base image
+# (which equals BUILDER_CUDA_VERSION) so that $CUDA_VERSION in RUN commands
+# resolves to the runtime CUDA version used for PyTorch wheel selection.
+ENV CUDA_VERSION=${CUDA_VERSION}
 ENV DEBIAN_FRONTEND=noninteractive

 # Install system dependencies including build tools
@@ -133,7 +143,7 @@ ENV UV_LINK_MODE=copy
 RUN gcc --version

 # Ensure CUDA compatibility library is loaded
-RUN echo "/usr/local/cuda-$(echo "$CUDA_VERSION" | cut -d. -f1,2)/compat/" > /etc/ld.so.conf.d/cuda-compat.conf && ldconfig
+RUN echo "/usr/local/cuda-$(echo "$BUILDER_CUDA_VERSION" | cut -d. -f1,2)/compat/" > /etc/ld.so.conf.d/cuda-compat.conf && ldconfig

 # ============================================================
 # SLOW-CHANGING DEPENDENCIES BELOW
@@ -309,7 +319,7 @@ RUN --mount=type=cache,target=/root/.cache/ccache \
 # Build DeepGEMM, pplx-kernels, DeepEP - runs in PARALLEL with csrc-build
 # This stage is independent and doesn't affect csrc cache
 FROM base AS extensions-build
-ARG CUDA_VERSION
+ARG BUILDER_CUDA_VERSION

 # This timeout (in seconds) is necessary when installing some dependencies via uv since it's likely to time out
 ENV UV_HTTP_TIMEOUT=500
@@ -325,7 +335,7 @@ COPY tools/install_deepgemm.sh /tmp/install_deepgemm.sh
 RUN --mount=type=cache,target=/root/.cache/uv \
    mkdir -p /tmp/deepgemm/dist && \
    VLLM_DOCKER_BUILD_CONTEXT=1 TORCH_CUDA_ARCH_LIST="9.0a 10.0a" /tmp/install_deepgemm.sh \
-        --cuda-version "${CUDA_VERSION}" \
+        --cuda-version "${BUILDER_CUDA_VERSION}" \
        ${DEEPGEMM_GIT_REF:+--ref "$DEEPGEMM_GIT_REF"} \
        --wheel-dir /tmp/deepgemm/dist || \
    echo "DeepGEMM build skipped (CUDA version requirement not met)"
@@ -2,6 +2,9 @@
  "_comment": "Auto-generated from Dockerfile ARGs. Do not edit manually. Run: python tools/generate_versions_json.py",
  "variable": {
    "CUDA_VERSION": {
+      "default": "12.8.1"
+    },
+    "BUILDER_CUDA_VERSION": {
      "default": "12.9.1"
    },
    "PYTHON_VERSION": {
@@ -11,7 +14,7 @@
      "default": "nvidia/cuda:12.9.1-devel-ubuntu20.04"
    },
    "FINAL_BASE_IMAGE": {
-      "default": "nvidia/cuda:12.9.1-base-ubuntu22.04"
+      "default": "nvidia/cuda:12.8.1-base-ubuntu22.04"
    },
    "GET_PIP_URL": {
      "default": "https://bootstrap.pypa.io/get-pip.py"
@@ -1,5 +1,5 @@
 # This file was autogenerated by uv via the following command:
-#    uv pip compile requirements/test.in -o requirements/test.txt --index-strategy unsafe-best-match --torch-backend cu129 --python-platform x86_64-manylinux_2_28 --python-version 3.12
+#    uv pip compile requirements/test.in -o requirements/test.txt --index-strategy unsafe-best-match --torch-backend cu128 --python-platform x86_64-manylinux_2_28 --python-version 3.12
 absl-py==2.1.0
    # via
    #   rouge-score
@@ -574,28 +574,28 @@ numpy==2.2.6
    #   tritonclient
    #   vocos
    #   xarray
-nvidia-cublas-cu12==12.9.1.4
+nvidia-cublas-cu12==12.8.4.1
    # via
    #   nvidia-cudnn-cu12
    #   nvidia-cusolver-cu12
    #   torch
-nvidia-cuda-cupti-cu12==12.9.79
+nvidia-cuda-cupti-cu12==12.8.90
    # via torch
-nvidia-cuda-nvrtc-cu12==12.9.86
+nvidia-cuda-nvrtc-cu12==12.8.93
    # via torch
-nvidia-cuda-runtime-cu12==12.9.79
+nvidia-cuda-runtime-cu12==12.8.90
    # via torch
 nvidia-cudnn-cu12==9.10.2.21
    # via torch
-nvidia-cufft-cu12==11.4.1.4
+nvidia-cufft-cu12==11.3.3.83
    # via torch
-nvidia-cufile-cu12==1.14.1.1
+nvidia-cufile-cu12==1.13.1.3
    # via torch
-nvidia-curand-cu12==10.3.10.19
+nvidia-curand-cu12==10.3.9.90
    # via torch
-nvidia-cusolver-cu12==11.7.5.82
+nvidia-cusolver-cu12==11.7.3.90
    # via torch
-nvidia-cusparse-cu12==12.5.10.65
+nvidia-cusparse-cu12==12.5.8.93
    # via
    #   nvidia-cusolver-cu12
    #   torch
@@ -603,7 +603,7 @@ nvidia-cusparselt-cu12==0.7.1
    # via torch
 nvidia-nccl-cu12==2.27.5
    # via torch
-nvidia-nvjitlink-cu12==12.9.86
+nvidia-nvjitlink-cu12==12.8.93
    # via
    #   nvidia-cufft-cu12
    #   nvidia-cusolver-cu12
@@ -611,7 +611,7 @@ nvidia-nvjitlink-cu12==12.9.86
    #   torch
 nvidia-nvshmem-cu12==3.4.5
    # via torch
-nvidia-nvtx-cu12==12.9.79
+nvidia-nvtx-cu12==12.8.90
    # via torch
 omegaconf==2.3.0
    # via
@@ -1154,7 +1154,7 @@ tomli==2.2.1
    # via schemathesis
 tomli-w==1.2.0
    # via schemathesis
-torch==2.10.0+cu129
+torch==2.10.0+cu128
    # via
    #   -r requirements/test.in
    #   accelerate
@@ -1179,7 +1179,7 @@ torch==2.10.0+cu129
    #   torchvision
    #   vector-quantize-pytorch
    #   vocos
-torchaudio==2.10.0+cu129
+torchaudio==2.10.0+cu128
    # via
    #   -r requirements/test.in
    #   encodec
@@ -1192,7 +1192,7 @@ torchmetrics==1.7.4
    #   pytorch-lightning
    #   terratorch
    #   torchgeo
-torchvision==0.25.0+cu129
+torchvision==0.25.0+cu128
    # via
    #   -r requirements/test.in
    #   lightly
@@ -19,7 +19,6 @@ else:

 logger = init_logger(__name__)

-BlockSize = Literal[1, 8, 16, 32, 64, 128, 256]
 CacheDType = Literal[
    "auto",
    "bfloat16",
@@ -39,13 +38,11 @@ KVOffloadingBackend = Literal["native", "lmcache"]
 class CacheConfig:
    """Configuration for the KV cache."""

-    block_size: SkipValidation[BlockSize] = None  # type: ignore[assignment]
-    """Size of a contiguous cache block in number of tokens. On CUDA devices,
-    only block sizes up to 32 are supported.
+    block_size: SkipValidation[int] = None  # type: ignore[assignment]
+    """Size of a contiguous cache block in number of tokens.

-    This config has no static default. If left unspecified by the user, it will
-    be set in `Platform.check_and_update_config()` based on the current
-    platform."""
+    This is None until `Platform.check_and_update_config()` sets it based on
+    the current platform. Always an int by the time the engine starts."""
    gpu_memory_utilization: float = Field(default=0.9, gt=0, le=1)
    """The fraction of GPU memory to be used for the model executor, which can
    range from 0 to 1. For example, a value of 0.5 would imply 50% GPU memory
@@ -59,7 +59,6 @@ from vllm.config import (
    get_attr_docs,
 )
 from vllm.config.cache import (
-    BlockSize,
    CacheDType,
    KVOffloadingBackend,
    MambaCacheMode,
@@ -431,7 +430,7 @@ class EngineArgs:
    max_parallel_loading_workers: int | None = (
        ParallelConfig.max_parallel_loading_workers
    )
-    block_size: BlockSize = CacheConfig.block_size
+    block_size: int = None  # type: ignore[assignment]
    enable_prefix_caching: bool | None = None
    prefix_caching_hash_algo: PrefixCachingHashAlgo = (
        CacheConfig.prefix_caching_hash_algo
@@ -163,8 +163,6 @@ class CudaPlatformBase(Platform):

    @classmethod
    def check_and_update_config(cls, vllm_config: "VllmConfig") -> None:
-        from vllm.v1.attention.backends.registry import AttentionBackendEnum
-
        parallel_config = vllm_config.parallel_config
        model_config = vllm_config.model_config

@@ -172,112 +170,19 @@ class CudaPlatformBase(Platform):
            parallel_config.worker_cls = "vllm.v1.worker.gpu_worker.Worker"

        cache_config = vllm_config.cache_config
-        if cache_config and cache_config.block_size is None:
+        user_specified_block_size = cache_config.block_size is not None
+        if not user_specified_block_size:
            cache_config.block_size = 16

-        # TODO(lucas): handle this more gracefully
-        # Note: model_config may be None during testing
-        # Note: block_size is initialized in
-        # HybridAttentionMambaModelConfig.verify_and_update_config
-        # for models with both attention and mamba,
-        # and doesn't need to be reinitialized here
-        if (
-            model_config is not None
-            and model_config.use_mla
-            and cache_config.block_size is not None
-        ):
-            use_sparse = hasattr(vllm_config.model_config.hf_config, "index_topk")
-            # If `--attention-config.backend` is not set and we are using MLA,
-            # then we default to FlashMLA backend for non-blackwell GPUs,
-            # else we default to CutlassMLA. For each case, we force the
-            # required block_size.
-            use_flashmla = False
-            use_cutlass_mla = False
-            use_flashinfer_mla = False
-            use_flashmla_sparse = False
-            use_flashinfer_mla_sparse = False
-
-            from vllm.v1.attention.ops.flashmla import is_flashmla_dense_supported
-
-            if vllm_config.attention_config.backend is None:
-                # Default case
-                hf_text_config = model_config.hf_text_config
-                qk_nope_head_dim = getattr(hf_text_config, "qk_nope_head_dim", 1)
-                if (
-                    cls.is_device_capability_family(100)
-                    and not use_sparse
-                    and qk_nope_head_dim == 128
-                ):
-                    # Blackwell => Force FlashInfer MLA (unless sparse, i.e. DSv3.2)
-                    # and only if qk_nope_head_dim == 128 (kernel constraint)
-                    use_flashinfer_mla = True
-                    # Set the backend in AttentionConfig so it's used during
-                    # backend selection
-                    vllm_config.attention_config.backend = (
-                        AttentionBackendEnum.FLASHINFER_MLA
-                    )
-                elif cls.is_device_capability_family(100) and not use_sparse:
-                    # Fall back to CUTLASS_MLA as 2nd priority on Blackwell
-                    use_cutlass_mla = True
-                elif is_flashmla_dense_supported()[0]:
-                    # Non-Blackwell with FlashMLA support
-                    use_flashmla = True
-                else:
-                    # Fallback: will use Triton MLA or other compatible backend
-                    pass
-            else:
-                # Forced case
-                backend = vllm_config.attention_config.backend
-                use_flashmla = backend == AttentionBackendEnum.FLASHMLA
-                use_cutlass_mla = backend == AttentionBackendEnum.CUTLASS_MLA
-                use_flashinfer_mla = backend == AttentionBackendEnum.FLASHINFER_MLA
-                use_flashmla_sparse = backend == AttentionBackendEnum.FLASHMLA_SPARSE
-                use_flashinfer_mla_sparse = (
-                    backend == AttentionBackendEnum.FLASHINFER_MLA_SPARSE
-                )
-
-            if (
-                use_flashmla
-                and is_flashmla_dense_supported()[0]
-                and cache_config.block_size % 64 != 0
-            ):
-                cache_config.block_size = 64
-                logger.info("Forcing kv cache block size to 64 for FlashMLA backend.")
-
-            if use_cutlass_mla and cache_config.block_size % 128 != 0:
-                cache_config.block_size = 128
-                logger.info(
-                    "Forcing kv cache block size to 128 for CUTLASS_MLA backend."
-                )
-
-            if (
-                use_flashinfer_mla
-                and cache_config.block_size != 32
-                and cache_config.block_size % 64 != 0
-            ):
-                cache_config.block_size = 64
-                logger.info(
-                    "Forcing kv cache block size to 64 for FlashInferMLA backend."
-                )
-
-            if use_sparse:
-                if not (use_flashmla_sparse or use_flashinfer_mla_sparse):
-                    use_flashmla_sparse = True
-
-                if use_flashmla_sparse and cache_config.block_size != 64:
-                    cache_config.block_size = 64
-                    logger.info(
-                        "Forcing kv cache block size to 64 for FlashMLASparse backend."
-                    )
-                elif use_flashinfer_mla_sparse and cache_config.block_size not in (
-                    32,
-                    64,
-                ):
-                    cache_config.block_size = 64
-                    logger.info(
-                        "Forcing kv cache block size to 64 for FlashInferMLASparse "
-                        "backend."
-                    )
+        # Ensure block_size is compatible with the attention backend.
+        # Note: model_config may be None during testing.
+        # Skip hybrid (attention+mamba) models — their block_size is
+        # managed by HybridAttentionMambaModelConfig
+        if model_config is not None and not model_config.is_hybrid:
+            cls._update_block_size_for_backend(
+                vllm_config,
+                user_specified_block_size,
+            )

        scheduler_config = vllm_config.scheduler_config
        # Note: model_config may be None during testing
@@ -293,6 +198,150 @@ class CudaPlatformBase(Platform):
            )
            scheduler_config.disable_chunked_mm_input = True

+    @classmethod
+    def _update_block_size_for_backend(
+        cls,
+        vllm_config: "VllmConfig",
+        user_specified_block_size: bool,
+    ) -> None:
+        """Ensure block_size is compatible with the attention backend.
+
+        If the user specified --block-size, the selector validates/filters
+        backends by that block size (raising on incompatibility). Otherwise,
+        the backend is selected unconstrained and block_size is set to the
+        backend's preferred value.
+        """
+        from vllm.config.vllm import set_current_vllm_config
+        from vllm.v1.attention.selector import AttentionSelectorConfig
+
+        model_config = vllm_config.model_config
+        cache_config = vllm_config.cache_config
+
+        device_capability = cls.get_device_capability()
+        if device_capability is None:
+            return
+
+        use_mla = model_config.use_mla
+        attn_selector_config = AttentionSelectorConfig(
+            head_size=model_config.get_head_size(),
+            dtype=model_config.dtype,  # type: ignore[arg-type]
+            kv_cache_dtype=cache_config.cache_dtype,
+            block_size=cache_config.block_size if user_specified_block_size else None,
+            use_mla=use_mla,
+            has_sink=False,
+            use_sparse=use_mla and hasattr(model_config.hf_config, "index_topk"),
+            use_mm_prefix=model_config.is_mm_prefix_lm,
+        )
+
+        user_specified_backend = vllm_config.attention_config.backend
+        num_heads = model_config.get_num_attention_heads(
+            vllm_config.parallel_config,
+        )
+        with set_current_vllm_config(vllm_config):
+            chosen_backend = cls.select_attention_backend(
+                selected_backend=user_specified_backend,
+                attn_selector_config=attn_selector_config,
+                device_capability=device_capability,
+                # Don't raise here — we produce better errors below.
+                raise_on_invalid=False,
+                num_heads=num_heads,
+            )
+
+            # If the user's --block-size forced a non-optimal backend,
+            # warn them. Only relevant when the user didn't also specify
+            # --attention-backend (in which case the choice is explicit).
+            if (
+                chosen_backend is not None
+                and user_specified_block_size
+                and user_specified_backend is None
+            ):
+                optimal = cls.select_attention_backend(
+                    selected_backend=None,
+                    attn_selector_config=attn_selector_config._replace(
+                        block_size=None,
+                    ),
+                    device_capability=device_capability,
+                    raise_on_invalid=False,
+                    num_heads=num_heads,
+                )
+                if optimal is not None and optimal != chosen_backend:
+                    logger.warning(
+                        "--block-size %d is not supported by the preferred "
+                        "%s backend. Using %s instead, which may result "
+                        "in reduced performance. Consider removing "
+                        "--block-size to auto-select the optimal "
+                        "block size.",
+                        cache_config.block_size,
+                        optimal.name,
+                        chosen_backend.name,
+                    )
+
+            if chosen_backend is not None:
+                if user_specified_block_size:
+                    # User's block_size is compatible with the chosen
+                    # backend.
+                    return
+                # User didn't specify --block-size, so auto-select the
+                # preferred block size for the chosen backend.
+                try:
+                    backend_class = chosen_backend.get_class()
+                except ImportError:
+                    return  # Will fail later with a better error
+                preferred = backend_class.get_preferred_block_size(
+                    cache_config.block_size,
+                )
+                if cache_config.block_size != preferred:
+                    logger.info(
+                        "Setting kv cache block size to %d for %s backend.",
+                        preferred,
+                        chosen_backend.name,
+                    )
+                    cache_config.block_size = preferred
+                return
+
+            # No valid backend found. If the user didn't constrain the
+            # selection, defer the error to get_attn_backend_cls where
+            # the full config (including per-layer settings) is
+            # available.
+            if not user_specified_block_size:
+                return
+
+            if user_specified_backend is not None:
+                # User specified --block-size and --attention-backend
+                # and they are incompatible.
+                try:
+                    backend_class = user_specified_backend.get_class()
+                    supported = backend_class.get_supported_kernel_block_sizes()
+                except ImportError:
+                    supported = None
+                raise ValueError(
+                    f"User-specified --block-size "
+                    f"{cache_config.block_size} is incompatible with "
+                    f"the specified --attention-backend "
+                    f"{user_specified_backend.name} (supported kernel "
+                    f"block sizes: {supported}). Either remove "
+                    f"--block-size to auto-select, or choose a "
+                    f"compatible value."
+                )
+            else:
+                # User specified --block-size but no backend supports
+                # it.
+                _, invalid_reasons = cls.get_valid_backends(
+                    device_capability=device_capability,
+                    attn_selector_config=attn_selector_config,
+                    num_heads=num_heads,
+                )
+                reasons_str = ", ".join(
+                    f"{b.name}: [{', '.join(r)}]" for b, r in invalid_reasons.items()
+                )
+                raise ValueError(
+                    f"No valid attention backend found for "
+                    f"--block-size {cache_config.block_size}. "
+                    f"Reasons: {{{reasons_str}}}. Either remove "
+                    f"--block-size to auto-select, or choose a "
+                    f"compatible value."
+                )
+
    @classmethod
    def get_current_memory_usage(
        cls, device: torch.types.Device | None = None
@@ -335,78 +384,126 @@ class CudaPlatformBase(Platform):

        return valid_backends_priorities, invalid_reasons

+    @classmethod
+    def select_attention_backend(
+        cls,
+        selected_backend: "AttentionBackendEnum | None",
+        attn_selector_config: "AttentionSelectorConfig",
+        device_capability: "DeviceCapability",
+        raise_on_invalid: bool = True,
+        num_heads: int | None = None,
+    ) -> "AttentionBackendEnum | None":
+        """Select the best attention backend for the given configuration.
+
+        Args:
+            selected_backend: User-specified backend, or None for auto-selection
+            attn_selector_config: Configuration for attention selection
+            device_capability: Device capability info
+            raise_on_invalid: If True, raise ValueError when no valid backend
+            num_heads: Number of attention heads per GPU, used for backend
+                priority ordering on Blackwell GPUs
+
+        Returns:
+            The selected backend enum, or None if no valid backend found
+            and raise_on_invalid is False
+        """
+        # First try checking just the selected backend, if there is one.
+        if selected_backend is not None:
+            try:
+                backend_class = selected_backend.get_class()
+                validation_errors = backend_class.validate_configuration(
+                    device_capability=device_capability,
+                    **attn_selector_config._asdict(),
+                )
+            except ImportError:
+                validation_errors = ["ImportError"]
+            if validation_errors:
+                if raise_on_invalid:
+                    raise ValueError(
+                        f"Selected backend {selected_backend} is not valid for "
+                        f"this configuration. Reason: {validation_errors}"
+                    )
+                return None
+            return selected_backend
+
+        # No selected backend, so find the best valid one.
+        valid_backends_priorities, invalid_reasons = cls.get_valid_backends(
+            device_capability=device_capability,
+            attn_selector_config=attn_selector_config,
+            num_heads=num_heads,
+        )
+
+        if len(valid_backends_priorities) == 0:
+            if raise_on_invalid:
+                reasons_str = (
+                    "{"
+                    + ", ".join(
+                        f"{backend.name}: [{', '.join(reasons)}]"
+                        for backend, reasons in invalid_reasons.items()
+                    )
+                    + "}"
+                )
+                config_str = attn_selector_config.__repr__()
+                raise ValueError(
+                    f"No valid attention backend found for {cls.device_name} "
+                    f"with {config_str}. Reasons: {reasons_str}."
+                )
+            return None
+
+        # Select the one with the highest priority (lowest index).
+        sorted_backends = sorted(valid_backends_priorities, key=lambda x: x[1])
+        return sorted_backends[0][0]
+
    @classmethod
    def get_attn_backend_cls(
        cls,
-        selected_backend: "AttentionBackendEnum",
+        selected_backend: "AttentionBackendEnum | None",
        attn_selector_config: "AttentionSelectorConfig",
        num_heads: int | None = None,
    ) -> str:
        device_capability = cls.get_device_capability()
        assert device_capability is not None

-        attn_selector_config = attn_selector_config._replace(block_size=None)
-        # First try checking just the selected backend, if there is one.
-        if selected_backend is not None:
-            try:
-                backend_class = selected_backend.get_class()
-                invalid_reasons = backend_class.validate_configuration(
-                    device_capability=device_capability,
-                    **attn_selector_config._asdict(),
-                )
-            except ImportError:
-                invalid_reasons = ["ImportError"]
-            if invalid_reasons:
-                raise ValueError(
-                    f"Selected backend {selected_backend} is not valid for "
-                    f"this configuration. Reason: {invalid_reasons}"
-                )
-            else:
-                logger.info("Using %s backend.", selected_backend)
-                return selected_backend.get_path()
-
-        # No selected backend or the selected backend is invalid,
-        # so we try finding a valid backend.
-        valid_backends_priorities, invalid_reasons = cls.get_valid_backends(
-            device_capability=device_capability,
+        chosen_backend = cls.select_attention_backend(
+            selected_backend=selected_backend,
            attn_selector_config=attn_selector_config,
            num_heads=num_heads,
+            device_capability=device_capability,
+            raise_on_invalid=True,
        )
-        reasons_str = (
-            "{"
-            + ", ".join(
-                f"{backend.name}: [{', '.join(reasons)}]"
-                for backend, reasons in invalid_reasons.items()
+        assert chosen_backend is not None  # raise_on_invalid=True guarantees this
+
+        # Log the selection
+        if selected_backend is not None:
+            logger.info("Using %s backend.", chosen_backend)
+        else:
+            # Get all valid backends for logging
+            valid_backends_priorities, invalid_reasons = cls.get_valid_backends(
+                device_capability=device_capability,
+                attn_selector_config=attn_selector_config,
+                num_heads=num_heads,
            )
-            + "}"
-        )
-        config_str = attn_selector_config.__repr__()
-        logger.debug_once(
-            f"Some attention backends are not valid for {cls.device_name} with "
-            f"{config_str}. Reasons: {reasons_str}."
-        )
-        if len(valid_backends_priorities) == 0:
-            raise ValueError(
-                f"No valid attention backend found for {cls.device_name} "
-                f"with {config_str}. Reasons: {reasons_str}."
+            reasons_str = (
+                "{"
+                + ", ".join(
+                    f"{backend.name}: [{', '.join(reasons)}]"
+                    for backend, reasons in invalid_reasons.items()
+                )
+                + "}"
+            )
+            config_str = attn_selector_config.__repr__()
+            logger.debug_once(
+                f"Some attention backends are not valid for {cls.device_name} with "
+                f"{config_str}. Reasons: {reasons_str}."
+            )
+            logger.info_once(
+                "Using %s attention backend out of potential backends: %s",
+                chosen_backend.name,
+                tuple(b[0].name for b in valid_backends_priorities),
+                scope="local",
            )

-        # We have found some valid backends. Select the one with the
-        # highest priority.
-        sorted_indices = sorted(
-            range(len(valid_backends_priorities)),
-            key=lambda i: valid_backends_priorities[i][1],
-        )
-        selected_index = sorted_indices[0]
-        selected_backend = valid_backends_priorities[selected_index][0]
-        logger.info_once(
-            "Using %s attention backend out of potential backends: %s.",
-            selected_backend.name,
-            "[" + ", ".join(f"'{b[0].name}'" for b in valid_backends_priorities) + "]",
-            scope="local",
-        )
-
-        return selected_backend.get_path()
+        return chosen_backend.get_path()

    @classmethod
    def get_supported_vit_attn_backends(cls) -> list["AttentionBackendEnum"]:
@@ -4,7 +4,7 @@
 from abc import ABC, abstractmethod
 from dataclasses import dataclass, replace
 from enum import Enum
-from typing import TYPE_CHECKING, Any, ClassVar, Generic, Protocol, TypeVar, get_args
+from typing import TYPE_CHECKING, Any, ClassVar, Generic, Protocol, TypeVar

 import numpy as np
 import torch
@@ -144,15 +144,9 @@ class AttentionBackend(ABC):

    @classmethod
    def supports_block_size(cls, block_size: int | None) -> bool:
-        from vllm.config.cache import BlockSize
-
        if block_size is None:
            return True

-        valid_sizes = get_args(BlockSize)
-        if block_size not in valid_sizes:
-            return False
-
        supported_kernel_block_sizes = cls.get_supported_kernel_block_sizes()
        if not supported_kernel_block_sizes:
            return True
@@ -167,6 +161,17 @@ class AttentionBackend(ABC):
                return True
        return False

+    @classmethod
+    def get_preferred_block_size(cls, default_block_size: int = 16) -> int:
+        supported_sizes = cls.get_supported_kernel_block_sizes()
+        if not supported_sizes:
+            return default_block_size
+
+        if cls.supports_block_size(default_block_size):
+            return default_block_size
+
+        return min(s.base if isinstance(s, MultipleOf) else s for s in supported_sizes)
+
    @classmethod
    def is_mla(cls) -> bool:
        return False
@@ -12,6 +12,7 @@ from vllm.v1.attention.backend import (
    AttentionMetadataBuilder,
    CommonAttentionMetadata,
 )
+from vllm.v1.attention.backends.utils import get_dcp_local_seq_lens
 from vllm.v1.kv_cache_interface import (
    AttentionSpec,
    KVCacheConfig,
@@ -143,6 +144,28 @@ def build_slot_mappings_by_layer(
    return slot_mappings_by_layer


+def prepare_dcp_local_seq_lens(
+    dcp_local_seq_lens: torch.Tensor,
+    seq_lens: torch.Tensor,
+    num_reqs: int,
+    dcp_size: int,
+    dcp_rank: int,
+    cp_kv_cache_interleave_size: int,
+) -> None:
+    """Populate the persistent DCP local seq_lens buffer (CUDA graph safe)."""
+    if dcp_size <= 1:
+        return
+
+    local_seq_lens = get_dcp_local_seq_lens(
+        seq_lens[:num_reqs],
+        dcp_size=dcp_size,
+        dcp_rank=dcp_rank,
+        cp_kv_cache_interleave_size=cp_kv_cache_interleave_size,
+    )
+    dcp_local_seq_lens[:num_reqs].copy_(local_seq_lens, non_blocking=True)
+    dcp_local_seq_lens[num_reqs:].zero_()
+
+
 def build_attn_metadata(
    attn_metadata_builders: list[AttentionMetadataBuilder],
    num_reqs: int,
@@ -155,9 +178,13 @@ def build_attn_metadata(
    block_tables: Sequence[torch.Tensor],
    slot_mappings: torch.Tensor,
    kv_cache_config: KVCacheConfig,
+    dcp_local_seq_lens: torch.Tensor | None = None,
 ) -> dict[str, Any]:
    seq_lens = seq_lens[:num_reqs]

+    if dcp_local_seq_lens is not None:
+        dcp_local_seq_lens = dcp_local_seq_lens[:num_reqs]
+
    attn_metadata: dict[str, Any] = {}
    kv_cache_groups = kv_cache_config.kv_cache_groups
    for i, kv_cache_spec in enumerate(kv_cache_groups):
@@ -175,6 +202,7 @@ def build_attn_metadata(
            block_table_tensor=block_table,
            slot_mapping=slot_mapping,
            causal=True,
+            dcp_local_seq_lens=dcp_local_seq_lens,
        )

        attn_metadata_builder = attn_metadata_builders[i]
@@ -4,6 +4,7 @@ from collections.abc import Iterable

 import torch

+from vllm.distributed import get_dcp_group
 from vllm.triton_utils import tl, triton
 from vllm.utils.math_utils import cdiv
 from vllm.v1.attention.backends.utils import PAD_SLOT_ID
@@ -18,19 +19,36 @@ class BlockTables:
        max_num_batched_tokens: int,
        max_model_len: int,
        device: torch.device,
+        cp_kv_cache_interleave_size: int = 1,
    ):
        self.block_sizes = block_sizes
        self.max_num_reqs = max_num_reqs
        self.max_num_batched_tokens = max_num_batched_tokens
        self.max_model_len = max_model_len
        self.device = device
+        assert cp_kv_cache_interleave_size >= 1
+        self.cp_kv_cache_interleave_size = cp_kv_cache_interleave_size
+
+        try:
+            dcp = get_dcp_group()
+            self.dcp_world_size, self.dcp_rank = dcp.world_size, dcp.rank_in_group
+        except AssertionError:
+            self.dcp_world_size, self.dcp_rank = 1, 0
+        # TODO(wentao): PCP supprot
+        self.total_cp_world_size = self.dcp_world_size
+        self.total_cp_rank = self.dcp_rank

        self.num_kv_cache_groups = len(self.block_sizes)
        # num_kv_cache_groups x [max_num_reqs, max_num_blocks]
        self.block_tables: list[StagedWriteTensor] = []
        for i in range(self.num_kv_cache_groups):
            block_size = self.block_sizes[i]
-            max_num_blocks = cdiv(self.max_model_len, block_size)
+            # with DCP, a request's KV is sharded across
+            # ranks, so one physical block on this rank
+            # corresponds to `block_size * total_cp_world_size`
+            # tokens in the global (unsharded) sequence.
+            virtual_block_size = block_size * self.total_cp_world_size
+            max_num_blocks = cdiv(self.max_model_len, virtual_block_size)
            block_table = StagedWriteTensor(
                (self.max_num_reqs, max_num_blocks),
                dtype=torch.int32,
@@ -131,6 +149,9 @@ class BlockTables:
            self.block_sizes_tensor,
            self.slot_mappings,
            self.slot_mappings.stride(0),
+            TOTAL_CP_WORLD_SIZE=self.total_cp_world_size,
+            TOTAL_CP_RANK=self.total_cp_rank,
+            CP_KV_CACHE_INTERLEAVE_SIZE=self.cp_kv_cache_interleave_size,
            PAD_ID=PAD_SLOT_ID,
            TRITON_BLOCK_SIZE=1024,  # type: ignore
        )
@@ -183,6 +204,9 @@ def _compute_slot_mappings_kernel(
    block_sizes,  # [num_kv_cache_groups]
    slot_mappings_ptr,  # [num_kv_cache_groups, max_num_tokens]
    slot_mappings_stride,
+    TOTAL_CP_WORLD_SIZE: tl.constexpr,
+    TOTAL_CP_RANK: tl.constexpr,
+    CP_KV_CACHE_INTERLEAVE_SIZE: tl.constexpr,
    PAD_ID: tl.constexpr,
    TRITON_BLOCK_SIZE: tl.constexpr,
 ):
@@ -201,6 +225,7 @@ def _compute_slot_mappings_kernel(
    block_table_ptr = _load_ptr(block_table_ptrs + group_id, tl.int32)
    block_table_stride = tl.load(block_table_strides + group_id)
    block_size = tl.load(block_sizes + group_id)
+    virtual_block_size = block_size * TOTAL_CP_WORLD_SIZE

    req_state_idx = tl.load(idx_mapping + batch_idx)
    start_idx = tl.load(query_start_loc + batch_idx)
@@ -208,11 +233,26 @@ def _compute_slot_mappings_kernel(
    for i in range(start_idx, end_idx, TRITON_BLOCK_SIZE):
        offset = i + tl.arange(0, TRITON_BLOCK_SIZE)
        positions = tl.load(pos + offset, mask=offset < end_idx, other=0)
-        block_indices = positions // block_size
+        block_indices = positions // virtual_block_size
        block_numbers = tl.load(
            block_table_ptr + req_state_idx * block_table_stride + block_indices
        )
-        slot_ids = block_numbers * block_size + positions % block_size
+        virtual_block_offsets = positions - block_indices * virtual_block_size
+
+        # determine whether the token is stored on this CP rank.
+        is_local = (
+            virtual_block_offsets // CP_KV_CACHE_INTERLEAVE_SIZE
+        ) % TOTAL_CP_WORLD_SIZE == TOTAL_CP_RANK
+        # mapping virture block offsets to local block offsets.
+        local_block_offsets = (
+            virtual_block_offsets // (TOTAL_CP_WORLD_SIZE * CP_KV_CACHE_INTERLEAVE_SIZE)
+        ) * CP_KV_CACHE_INTERLEAVE_SIZE + (
+            virtual_block_offsets % CP_KV_CACHE_INTERLEAVE_SIZE
+        )
+
+        # physical slot index
+        slot_ids = block_numbers * block_size + local_block_offsets
+        slot_ids = tl.where(is_local, slot_ids, PAD_ID)
        tl.store(slot_mapping_ptr + offset, slot_ids, mask=offset < end_idx)


@@ -10,6 +10,7 @@ from tqdm import tqdm

 from vllm.config import VllmConfig
 from vllm.config.compilation import CUDAGraphMode
+from vllm.distributed import get_dcp_group
 from vllm.distributed.parallel_state import graph_capture, is_global_first_rank
 from vllm.forward_context import set_forward_context
 from vllm.v1.attention.backend import AttentionMetadataBuilder
@@ -17,6 +18,7 @@ from vllm.v1.kv_cache_interface import KVCacheConfig
 from vllm.v1.worker.gpu.attn_utils import (
    build_attn_metadata,
    build_slot_mappings_by_layer,
+    prepare_dcp_local_seq_lens,
 )
 from vllm.v1.worker.gpu.block_table import BlockTables
 from vllm.v1.worker.gpu.dp_utils import make_num_tokens_across_dp
@@ -257,6 +259,23 @@ def prepare_inputs_to_capture(
    input_buffers.seq_lens[:num_reqs] = num_tokens
    input_buffers.seq_lens[num_reqs:] = 0

+    try:
+        dcp_group = get_dcp_group()
+        dcp_world_size = dcp_group.world_size
+        dcp_rank = dcp_group.rank_in_group
+    except AssertionError:
+        dcp_world_size = 1
+        dcp_rank = 0
+    if dcp_world_size > 1:
+        prepare_dcp_local_seq_lens(
+            input_buffers.dcp_local_seq_lens,
+            input_buffers.seq_lens,
+            num_reqs,
+            dcp_size=dcp_world_size,
+            dcp_rank=dcp_rank,
+            cp_kv_cache_interleave_size=block_tables.cp_kv_cache_interleave_size,
+        )
+
    input_block_tables = [x[:num_reqs] for x in block_tables.input_block_tables]
    slot_mappings = block_tables.slot_mappings[:, :num_tokens]
    slot_mappings_by_layer = build_slot_mappings_by_layer(
@@ -275,5 +294,6 @@ def prepare_inputs_to_capture(
        block_tables=input_block_tables,
        slot_mappings=slot_mappings,
        kv_cache_config=kv_cache_config,
+        dcp_local_seq_lens=input_buffers.dcp_local_seq_lens,
    )
    return attn_metadata, slot_mappings_by_layer
@@ -27,6 +27,10 @@ class InputBuffers:
            max_num_reqs + 1, dtype=torch.int32, device=device
        )
        self.seq_lens = torch.zeros(max_num_reqs, dtype=torch.int32, device=device)
+        # DCP: per-request local seq_lens buffer
+        self.dcp_local_seq_lens = torch.zeros(
+            max_num_reqs, dtype=torch.int32, device=device
+        )


@dataclass
@@ -11,6 +11,7 @@ import torch.nn as nn
 from vllm.config import VllmConfig
 from vllm.config.compilation import CUDAGraphMode
 from vllm.distributed.parallel_state import (
+    get_dcp_group,
    get_pp_group,
    prepare_communication_buffer_for_model,
 )
@@ -24,6 +25,7 @@ from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
 from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
 from vllm.v1.kv_cache_interface import KVCacheConfig
 from vllm.v1.outputs import DraftTokenIds, ModelRunnerOutput
+from vllm.v1.worker.cp_utils import check_attention_cp_compatibility
 from vllm.v1.worker.gpu.async_utils import AsyncOutput
 from vllm.v1.worker.gpu.attn_utils import (
    build_attn_metadata,
@@ -31,6 +33,7 @@ from vllm.v1.worker.gpu.attn_utils import (
    get_kv_cache_spec,
    init_attn_backend,
    init_kv_cache,
+    prepare_dcp_local_seq_lens,
 )
 from vllm.v1.worker.gpu.block_table import BlockTables
 from vllm.v1.worker.gpu.buffer_utils import async_copy_to_gpu
@@ -248,11 +251,15 @@ class GPUModelRunner(LoRAModelRunnerMixin):
            max_num_batched_tokens=self.max_num_tokens,
            max_model_len=self.max_model_len,
            device=self.device,
+            cp_kv_cache_interleave_size=(
+                self.parallel_config.cp_kv_cache_interleave_size
+            ),
        )

        self.attn_backends, self.attn_metadata_builders = init_attn_backend(
            self.kv_cache_config, self.vllm_config, self.device
        )
+        check_attention_cp_compatibility(self.vllm_config)
        if self.do_spec_decode:
            # HACK(woosuk)
            self.speculator.set_attn(
@@ -294,6 +301,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
            block_tables=block_tables,
            slot_mappings=slot_mappings,
            kv_cache_config=self.kv_cache_config,
+            dcp_local_seq_lens=self.input_buffers.dcp_local_seq_lens,
        )
        input_batch.attn_metadata = attn_metadata
        input_batch.slot_mappings = slot_mappings_by_layer
@@ -627,6 +635,19 @@ class GPUModelRunner(LoRAModelRunnerMixin):
        )
        seq_lens = self.input_buffers.seq_lens[:num_reqs]

+        dcp_size = self.parallel_config.decode_context_parallel_size
+        if dcp_size > 1:
+            prepare_dcp_local_seq_lens(
+                self.input_buffers.dcp_local_seq_lens,
+                seq_lens,
+                num_reqs,
+                dcp_size=dcp_size,
+                dcp_rank=get_dcp_group().rank_in_group,
+                cp_kv_cache_interleave_size=(
+                    self.parallel_config.cp_kv_cache_interleave_size
+                ),
+            )
+
        # Prepare M-RoPE positions.
        if self.uses_mrope:
            self.mrope_states.prepare_mrope_positions(
@@ -674,6 +695,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
            block_tables=block_tables,
            slot_mappings=slot_mappings,
            kv_cache_config=self.kv_cache_config,
+            dcp_local_seq_lens=self.input_buffers.dcp_local_seq_lens,
        )

        input_ids = self.input_buffers.input_ids[:num_tokens_after_padding]
Author	SHA1	Message	Date
Tyler Michael Smith	9df42d2000	[Build] Update pre-commit pip-compile hook to use cu128 torch backend Matches the test.txt lockfile change to CUDA 12.8. Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>	2026-02-18 16:41:53 -05:00
Tyler Michael SmithandClaude Opus 4.6	8d151aa148	[Build] Recompile test.txt lockfile with cu128 torch backend The lockfile was compiled with --torch-backend cu129, pinning torch==2.10.0+cu129. This breaks Docker builds that use CUDA_VERSION=12.8 because the cu128 PyTorch index does not carry +cu129 wheels. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>	2026-02-18 14:38:23 -05:00
Tyler Michael SmithandClaude Opus 4.6	11a13b0fd3	[Docker] Add BUILDER_CUDA_VERSION to decouple build and runtime CUDA versions Allow compiling csrc/ and extensions (DeepGEMM, EP kernels) with a different CUDA toolkit than the one shipped in the final runtime image. BUILDER_CUDA_VERSION controls the devel base image used for compilation, while CUDA_VERSION selects the runtime base image and PyTorch wheel index. Override the CUDA_VERSION env var inherited from the nvidia base image in the build stages so PyTorch index URLs resolve to the runtime version. Update the CUDA 13.0 release pipeline entries to pass the new arg. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>	2026-02-18 14:38:23 -05:00
Amr MahdiandGitHub	df3f537a66	[CI] Remove unused precompiled wheel args from image build (#34767 ) Signed-off-by: Amr Mahdi <amrmahdi@meta.com>	2026-02-17 18:58:18 -08:00
Matthew BonanniandGitHub	7743152957	[Attention] Refactor `check_and_update_config` (#33600 ) Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>	2026-02-17 17:06:54 -08:00
Wentao YeandGitHub	ab33d2a629	[Feature] Decode Context Parallel support for GPU model runner v2 (#34179 ) Signed-off-by: yewentao256 <zhyanwentao@126.com>	2026-02-17 16:27:15 -08:00