patch to fix MLA multi-group KV cache

Dockerfile

@@ -1,9 +1,8 @@
 #FROM vllm/vllm-openai:v0.19.0-cu130
-#FROM vllm/vllm-openai:deepseekv4-cu130
-FROM vllm/vllm-openai:v0.20.0-cu130
+FROM vllm/vllm-openai:cu130-nightly-x86_64
 
 # Fix the broken ass nightly build that forgot to include pandas
-#RUN pip install --no-cache-dir pandas
+RUN pip install --no-cache-dir pandas
 
 # Install LMCache for KV cache offloading / sharing across nodes
 # Build with system CUDA 13.0 for Blackwell (B200)
@@ -13,37 +12,24 @@ RUN apt-get update && apt-get install -y git \
     libcublas-dev-13-0 \
     libcurand-dev-13-0 \
     libcufft-dev-13-0 \
-    cuda-nvrtc-dev-13-0 \
     libnvjitlink-dev-13-0 && \
-    git clone https://github.com/biondizzle/LMCache.git /tmp/lmcache && \
+    git clone https://github.com/Byteflux/LMCache.git /tmp/lmcache && \
     cd /tmp/lmcache && \
-    git checkout feat/redis-ttl && \
+    git checkout mla-multi-group-kv-cache-with-redis && \
     CUDA_HOME=/usr/local/cuda \
     TORCH_CUDA_ARCH_LIST="10.0" \
     pip install --no-cache-dir --no-build-isolation . && \
-    rm -rf /tmp/lmcache && export CACHE_BUSTER=2
-
-ENV VLLM_DOCKER_BUILD_CONTEXT=1
-# Needs git to be installed before this can run
-#RUN curl -fsSL https://raw.githubusercontent.com/vllm-project/vllm/v0.20.0/tools/install_deepgemm.sh | bash
-# Had to modify this
-COPY ./install_deepgemm.sh /opt/install_deepgemm.sh
-RUN chmod +x /opt/install_deepgemm.sh
-RUN /opt/install_deepgemm.sh
+    rm -rf /tmp/lmcache && export CACHE_BUSTER=1
 
 # Copy over nemotron reasonong parser
-#COPY ./super_v3_reasoning_parser.py /opt/super_v3_reasoning_parser.py
+COPY ./super_v3_reasoning_parser.py /opt/super_v3_reasoning_parser.py
 
 # Copy over deepseek tool call parser with MTP fixes
-#COPY deepseekv32_tool_parser.py /usr/local/lib/python3.12/dist-packages/vllm/tool_parsers/deepseekv32_tool_parser.py
-#COPY deepseekv4_tool_parser.py /usr/local/lib/python3.12/dist-packages/vllm/tool_parsers/deepseekv4_tool_parser.py
+COPY deepseekv32_tool_parser.py /usr/local/lib/python3.12/dist-packages/vllm/tool_parsers/deepseekv32_tool_parser.py
 
 # Copy over minimax tool call parser with kwargs fixes
-#COPY minimax_tool_parser.py /usr/local/lib/python3.12/dist-packages/vllm/tool_parsers/minimax_tool_parser.py
+COPY minimax_tool_parser.py /usr/local/lib/python3.12/dist-packages/vllm/tool_parsers/minimax_tool_parser.py
 
 # Copy over minimax parsers with kwargs fixes
-#COPY minimax_tool_parser.py /usr/local/lib/python3.12/dist-packages/vllm/tool_parsers/minimax_tool_parser.py
-#COPY minimax_m2_parser.py /usr/local/lib/python3.12/dist-packages/vllm/parser/minimax_m2_parser.py
-
-# Fix DeepseekV4Config missing max_position_embeddings attr for RoPE standardization
-#COPY deepseek_v4.py /usr/local/lib/python3.12/dist-packages/vllm/transformers_utils/configs/deepseek_v4.py
+COPY minimax_tool_parser.py /usr/local/lib/python3.12/dist-packages/vllm/tool_parsers/minimax_tool_parser.py
+COPY minimax_m2_parser.py /usr/local/lib/python3.12/dist-packages/vllm/parser/minimax_m2_parser.py

deepseek_v4.py

@@ -1,73 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-from transformers import PretrainedConfig
-
-
-class DeepseekV4Config(PretrainedConfig):
-    model_type = "deepseek_v4"
-
-    def __init__(
-        self,
-        vocab_size: int = 129024,
-        hidden_size: int = 7168,
-        intermediate_size: int = 18432,
-        moe_intermediate_size: int = 1408,
-        num_hidden_layers: int = 61,
-        num_attention_heads: int = 128,
-        num_key_value_heads: int = 128,
-        n_routed_experts: int = 256,
-        n_shared_experts: int = 1,
-        num_experts_per_tok: int = 8,
-        head_dim: int = 256,
-        q_lora_rank: int = 1536,
-        o_lora_rank: int = 1536,
-        qk_rope_head_dim: int = 64,
-        o_groups: int = 128,
-        rms_norm_eps: float = 1e-6,
-        rope_theta: float = 10000.0,
-        rope_scaling: dict | None = None,
-        max_position_embeddings: int = 163840,
-        hidden_act: str = "silu",
-        norm_topk_prob: bool = True,
-        compress_ratios: list | None = None,
-        compress_rope_theta: float = 10000.0,
-        num_hash_layers: int = 0,
-        hc_mult: float = 1.0,
-        hc_eps: float = 1e-6,
-        hc_sinkhorn_iters: int = 8,
-        swiglu_limit: int = 0,
-        index_topk: int = 0,
-        num_nextn_predict_layers: int = 0,
-        **kwargs,
-    ):
-        self.vocab_size = vocab_size
-        self.hidden_size = hidden_size
-        self.intermediate_size = intermediate_size
-        self.moe_intermediate_size = moe_intermediate_size
-        self.num_hidden_layers = num_hidden_layers
-        self.num_attention_heads = num_attention_heads
-        self.num_key_value_heads = num_key_value_heads
-        self.n_routed_experts = n_routed_experts
-        self.n_shared_experts = n_shared_experts
-        self.num_experts_per_tok = num_experts_per_tok
-        self.head_dim = head_dim
-        self.q_lora_rank = q_lora_rank
-        self.o_lora_rank = o_lora_rank
-        self.qk_rope_head_dim = qk_rope_head_dim
-        self.o_groups = o_groups
-        self.rms_norm_eps = rms_norm_eps
-        self.rope_theta = rope_theta
-        self.rope_scaling = rope_scaling
-        self.max_position_embeddings = max_position_embeddings
-        self.hidden_act = hidden_act
-        self.norm_topk_prob = norm_topk_prob
-        self.compress_ratios = compress_ratios or [1] * num_hidden_layers
-        self.compress_rope_theta = compress_rope_theta
-        self.num_hash_layers = num_hash_layers
-        self.hc_mult = hc_mult
-        self.hc_eps = hc_eps
-        self.hc_sinkhorn_iters = hc_sinkhorn_iters
-        self.swiglu_limit = swiglu_limit
-        self.index_topk = index_topk
-        self.num_nextn_predict_layers = num_nextn_predict_layers
-        super().__init__(**kwargs)

deepseekv4_tool_parser.py

@@ -1,33 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import regex as re
-
-from vllm.tool_parsers.deepseekv32_tool_parser import DeepSeekV32ToolParser
-
-
-class DeepSeekV4ToolParser(DeepSeekV32ToolParser):
-    """
-    DeepSeek V4 DSML tool parser.
-
-    V4 keeps the V3.2 DSML invoke/parameter grammar, but wraps tool calls in
-    ``<｜DSML｜tool_calls>`` instead of ``<｜DSML｜function_calls>``.
-    """
-
-    tool_call_start_token: str = "<｜DSML｜tool_calls>"
-    tool_call_end_token: str = "</｜DSML｜tool_calls>"
-
-    def __init__(self, tokenizer, tools=None):
-        super().__init__(tokenizer, tools)
-        # The V3.2 __init__ overwrites the class-level start/end tokens
-        # with instance attributes set to "function_calls". Re-override
-        # them to the V4 "tool_calls" tags so _extract_content(),
-        # _extract_invoke_regions(), etc. match the correct tags.
-        self.tool_call_start_token = "<｜DSML｜tool_calls>"
-        self.tool_call_end_token = "</｜DSML｜tool_calls>"
-        self.tool_calls_start_token = self.tool_call_start_token
-        # Recompile regexes with the V4 tool_calls tags instead of
-        # the V3.2 function_calls tags that the base class hardcodes.
-        self.tool_call_complete_regex = re.compile(
-            r"<｜DSML｜tool_calls>(.*?)</｜DSML｜tool_calls>", re.DOTALL
-        )

install_deepgemm.sh

@@ -1,121 +0,0 @@
-#!/bin/bash
-# Script to build and/or install DeepGEMM from source
-# Default: build and install immediately
-# Optional: build wheels to a directory for later installation (useful in multi-stage builds)
-set -e
-
-# Default values
-# Keep DEEPGEMM_GIT_REF in sync with cmake/external_projects/deepgemm.cmake
-DEEPGEMM_GIT_REPO="https://github.com/deepseek-ai/DeepGEMM.git"
-DEEPGEMM_GIT_REF="891d57b4db1071624b5c8fa0d1e51cb317fa709f"
-WHEEL_DIR=""
-
-# Parse command line arguments
-while [[ $# -gt 0 ]]; do
-    case $1 in
-        --ref)
-            if [[ -z "$2" || "$2" =~ ^- ]]; then
-                echo "Error: --ref requires an argument." >&2
-                exit 1
-            fi
-            DEEPGEMM_GIT_REF="$2"
-            shift 2
-            ;;
-        --cuda-version)
-            if [[ -z "$2" || "$2" =~ ^- ]]; then
-                echo "Error: --cuda-version requires an argument." >&2
-                exit 1
-            fi
-            CUDA_VERSION="$2"
-            shift 2
-            ;;
-        --wheel-dir)
-            if [[ -z "$2" || "$2" =~ ^- ]]; then
-                echo "Error: --wheel-dir requires a directory path." >&2
-                exit 1
-            fi
-            WHEEL_DIR="$2"
-            shift 2
-            ;;
-        -h|--help)
-            echo "Usage: $0 [OPTIONS]"
-            echo "Options:"
-            echo "  --ref REF          Git reference to checkout (default: $DEEPGEMM_GIT_REF)"
-            echo "  --cuda-version VER CUDA version (auto-detected if not provided)"
-            echo "  --wheel-dir PATH   If set, build wheel into PATH but do not install"
-            echo "  -h, --help         Show this help message"
-            exit 0
-            ;;
-        *)
-            echo "Unknown option: $1" >&2
-            exit 1
-            ;;
-    esac
-done
-
-# Auto-detect CUDA version if not provided
-if [ -z "$CUDA_VERSION" ]; then
-    if command -v nvcc >/dev/null 2>&1; then
-        CUDA_VERSION=$(nvcc --version | grep "release" | sed -n 's/.*release \([0-9]\+\.[0-9]\+\).*/\1/p')
-        echo "Auto-detected CUDA version: $CUDA_VERSION"
-    else
-        echo "Warning: Could not auto-detect CUDA version. Please specify with --cuda-version"
-        exit 1
-    fi
-fi
-
-# Extract major and minor version numbers
-CUDA_MAJOR="${CUDA_VERSION%%.*}"
-CUDA_MINOR="${CUDA_VERSION#"${CUDA_MAJOR}".}"
-CUDA_MINOR="${CUDA_MINOR%%.*}"
-echo "CUDA version: $CUDA_VERSION (major: $CUDA_MAJOR, minor: $CUDA_MINOR)"
-
-# Check CUDA version requirement
-if [ "$CUDA_MAJOR" -lt 12 ] || { [ "$CUDA_MAJOR" -eq 12 ] && [ "$CUDA_MINOR" -lt 8 ]; }; then
-    echo "Skipping DeepGEMM build/installation (requires CUDA 12.8+ but got ${CUDA_VERSION})"
-    exit 0
-fi
-
-echo "Preparing DeepGEMM build..."
-echo "Repository: $DEEPGEMM_GIT_REPO"
-echo "Reference: $DEEPGEMM_GIT_REF"
-
-# Create a temporary directory for the build
-INSTALL_DIR=$(mktemp -d)
-trap 'rm -rf "$INSTALL_DIR"' EXIT
-
-# Clone the repository
-git clone --recursive --shallow-submodules "$DEEPGEMM_GIT_REPO" "$INSTALL_DIR/deepgemm"
-pushd "$INSTALL_DIR/deepgemm"
-
-# Checkout the specific reference
-git checkout "$DEEPGEMM_GIT_REF"
-
-# Clean previous build artifacts
-# (Based on https://github.com/deepseek-ai/DeepGEMM/blob/main/install.sh)
-rm -rf -- build dist *.egg-info 2>/dev/null || true
-
-# Build wheel
-echo "🏗️  Building DeepGEMM wheel..."
-python3 setup.py bdist_wheel
-
-# If --wheel-dir was specified, copy wheels there and exit
-if [ -n "$WHEEL_DIR" ]; then
-    mkdir -p "$WHEEL_DIR"
-    cp dist/*.whl "$WHEEL_DIR"/
-    echo "✅ Wheel built and copied to $WHEEL_DIR"
-    popd
-    exit 0
-fi
-
-# Default behaviour: install built wheel
-if command -v uv >/dev/null 2>&1; then
-    echo "Installing DeepGEMM wheel using uv..."
-    uv pip install --system dist/*.whl
-else
-    echo "Installing DeepGEMM wheel using pip..."
-    python3 -m pip install dist/*.whl
-fi
-
-popd
-echo "✅ DeepGEMM installation completed successfully"