other things

Dockerfile

@@ -1,6 +1,8 @@
 #FROM vllm/vllm-openai:v0.19.0-cu130
-#FROM vllm/vllm-openai:cu130-nightly-x86_64
-FROM vllm/vllm-openai:glm51-cu130
+#vllm says version 0.20.2rc1.dev9+g01d4d1ad3
+FROM  vllm/vllm-openai:glm52
+#FROM vllm/vllm-openai:nightly
+#FROM vllm/vllm-openai:glm51-cu130
 
 # Install LMCache for KV cache offloading / sharing across nodes
 # Build with system CUDA 13.0 for Blackwell (B200)
@@ -13,30 +15,27 @@ RUN apt-get update && apt-get install -y git \
     libnvjitlink-dev-13-0 && \
     git clone https://github.com/biondizzle/LMCache.git /tmp/lmcache && \
     cd /tmp/lmcache && \
-    git checkout feat/redis-ttl && \
+    git checkout dream-build && \
     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=1
-
-# Copy over nemotron reasonong parser
-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 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 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
+    rm -rf /tmp/lmcache && export CACHE_BUSTER=5
 
 
-# Patch tool parser for GLM regex fix
-COPY glm4_moe_tool_parser.py /usr/local/lib/python3.12/dist-packages/vllm/tool_parsers/glm4_moe_tool_parser.py
-COPY utils.py /usr/local/lib/python3.12/dist-packages/vllm/tool_parsers/utils.py
+# Make sure we have patch to make MTP work on GLM (I did this based on vllm-project/vllm#40989)
+COPY indexer.py /usr/local/lib/python3.12/dist-packages/vllm/v1/attention/backends/mla/indexer.py
+COPY deep_gemm.py /usr/local/lib/python3.12/dist-packages/vllm/utils/deep_gemm.py  
+
+# These were from https://github.com/vllm-project/vllm/pull/41357/changes#diff-75b8ca6d854db6a47e75db6507afd20c15624f35229e2fd0d71642bffd70b11c
+COPY shm_broadcast.py /usr/local/lib/python3.12/dist-packages/vllm/distributed/device_communicators/shm_broadcast.py
+COPY multiproc_executor.py /usr/local/lib/python3.12/dist-packages/vllm/v1/executor/multiproc_executor.py
+
+# Make sure we have the latest up to date chat template
+#COPY glm_5.1_chat_template.jinja /opt/chat_template.jinja
+
+# GLM 5.1 LMCache config
+COPY lmcache-config-glm-52.yaml /opt/lmcache-config-glm-52.yaml
+
+# DEEPSEEK v4 LMCache config
+#COPY lmcache-config-dsv4.yaml /opt/lmcache-config-dsv4.yaml
 
-# Patch hf renderer to force string content format for GLM models
-# This fixes the issue where tool response content is dropped
-COPY hf.py /usr/local/lib/python3.12/dist-packages/vllm/renderers/hf.py

deep_gemm.py

@@ -0,0 +1,587 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Compatibility wrapper for DeepGEMM API changes.
+
+Users of vLLM should always import **only** these wrappers.
+"""
+
+import functools
+import importlib
+import os
+from collections.abc import Callable
+from enum import Enum
+from typing import Any, NoReturn
+
+import torch
+
+import vllm.envs as envs
+from vllm.logger import logger
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
+    get_fp8_min_max,
+)
+from vllm.platforms import current_platform
+from vllm.utils.import_utils import has_deep_gemm
+from vllm.utils.math_utils import cdiv
+
+_DEEPGEMM_BLACKWELL_EXCLUDED_MODEL_TYPES: set[str] = {
+    "qwen3_5_text",
+    "qwen3_5_moe_text",
+}
+
+
+def should_auto_disable_deep_gemm(model_type: str | None) -> bool:
+    """Check if DeepGemm should be auto-disabled for this model on Blackwell.
+
+    Returns True if the model is known to have accuracy degradation with
+    DeepGemm's E8M0 scale format on Blackwell GPUs (SM100+).
+    """
+    if model_type is None:
+        return False
+    if not current_platform.is_device_capability_family(100):
+        return False
+    return model_type in _DEEPGEMM_BLACKWELL_EXCLUDED_MODEL_TYPES
+
+
+class DeepGemmQuantScaleFMT(Enum):
+    # Float32 scales in Float32 tensor
+    FLOAT32 = 0
+    # Compute float32 scales and ceil the scales to UE8M0.
+    # Keep the scales in Float32 tensor.
+    FLOAT32_CEIL_UE8M0 = 1
+    # Compute float32 scales and ceil the scales to UE8M0.
+    # Pack the scales into a int32 tensor where each int32
+    # element contains 4 scale values.
+    UE8M0 = 2
+
+    @classmethod
+    def init_oracle_cache(cls) -> None:
+        """Initialize the oracle decision and store it in the class cache"""
+        cached = getattr(cls, "_oracle_cache", None)
+        if cached is not None:
+            return
+
+        use_e8m0 = (
+            envs.VLLM_USE_DEEP_GEMM_E8M0
+            and is_deep_gemm_supported()
+            and (_fp8_gemm_nt_impl is not None)
+        )
+        if not use_e8m0:
+            cls._oracle_cache = cls.FLOAT32  # type: ignore
+            return
+
+        cls._oracle_cache = (  # type: ignore
+            cls.UE8M0
+            if current_platform.is_device_capability_family(100)
+            else cls.FLOAT32_CEIL_UE8M0
+        )
+
+    @classmethod
+    def from_oracle(cls) -> "DeepGemmQuantScaleFMT":
+        """Return the pre-initialized oracle decision"""
+        cached = getattr(cls, "_oracle_cache", None)
+        assert cached is not None, "DeepGemmQuantScaleFMT oracle cache not initialized"
+        return cached
+
+
+@functools.cache
+def is_deep_gemm_supported() -> bool:
+    """Return `True` if DeepGEMM is supported on the current platform.
+    Currently, only Hopper and Blackwell GPUs are supported.
+    """
+    is_supported_arch = current_platform.support_deep_gemm()
+    return envs.VLLM_USE_DEEP_GEMM and has_deep_gemm() and is_supported_arch
+
+
+@functools.cache
+def is_deep_gemm_e8m0_used() -> bool:
+    """Return `True` if vLLM is configured to use DeepGEMM "
+    "E8M0 scale on a Hopper or Blackwell-class GPU.
+    """
+    if not is_deep_gemm_supported():
+        logger.debug_once(
+            "DeepGEMM E8M0 disabled: DeepGEMM not supported on this system."
+        )
+        return False
+
+    _lazy_init()
+
+    if _fp8_gemm_nt_impl is None:
+        logger.info_once("DeepGEMM E8M0 disabled: _fp8_gemm_nt_impl not found")
+        return False
+
+    if envs.VLLM_USE_DEEP_GEMM_E8M0:
+        logger.info_once("DeepGEMM E8M0 enabled on current platform.")
+        return True
+
+    logger.info_once("DeepGEMM E8M0 disabled on current configuration.")
+    return False
+
+
+def _missing(*_: Any, **__: Any) -> NoReturn:
+    """Placeholder for unavailable DeepGEMM backend."""
+    raise RuntimeError(
+        "DeepGEMM backend is not available or outdated. Please install or "
+        "update the `deep_gemm` to a newer version to enable FP8 kernels."
+    )
+
+
+_cublaslt_gemm_nt_impl: Callable[..., Any] | None = None
+_fp8_gemm_nt_impl: Callable[..., Any] | None = None
+_fp8_einsum_impl: Callable[..., Any] | None = None
+_grouped_impl: Callable[..., Any] | None = None
+_grouped_masked_impl: Callable[..., Any] | None = None
+_grouped_fp4_impl: Callable[..., Any] | None = None
+_fp8_fp4_mqa_logits_impl: Callable[..., Any] | None = None
+_fp8_fp4_paged_mqa_logits_impl: Callable[..., Any] | None = None
+_get_paged_mqa_logits_metadata_impl: Callable[..., Any] | None = None
+_tf32_hc_prenorm_gemm_impl: Callable[..., Any] | None = None
+_get_mn_major_tma_aligned_tensor_impl: Callable[..., Any] | None = None
+_get_mk_alignment_for_contiguous_layout_impl: Callable[..., Any] | None = None
+_transform_sf_into_required_layout_impl: Callable[..., Any] | None = None
+
+
+@functools.cache
+def _import_deep_gemm():
+    """Import the deep_gemm module.
+
+    Prefers an externally installed ``deep_gemm`` package (so users can
+    pin a specific version), then falls back to the vendored copy bundled
+    in the vLLM wheel.
+
+    Returns ``None`` when neither source is usable.
+    """
+    # 1. Try the external (pip-installed) package first.
+    try:
+        module = importlib.import_module("deep_gemm")
+        logger.debug_once("Imported deep_gemm module from site-packages")
+        return module
+    except ImportError:
+        logger.debug_once(
+            "deep_gemm not found in site-packages, "
+            "trying vendored vllm.third_party.deep_gemm"
+        )
+
+    # 2. Fall back to the vendored copy bundled in the vLLM wheel.
+    try:
+        module = importlib.import_module("vllm.third_party.deep_gemm")
+        logger.debug_once("Imported deep_gemm module from vllm.third_party.deep_gemm")
+        return module
+    except ImportError:
+        logger.debug_once("Vendored deep_gemm not found either")
+    except Exception as e:
+        # The vendored module may raise RuntimeError during _C.init()
+        # if JIT include files are missing (e.g. incomplete wheel).
+        logger.warning_once("Failed to import vendored deep_gemm: %s", e)
+
+    return None
+
+
+def _lazy_init() -> None:
+    """Import deep_gemm and resolve symbols on first use."""
+    global _cublaslt_gemm_nt_impl
+    global _fp8_gemm_nt_impl, _fp8_einsum_impl
+    global _grouped_impl, _grouped_masked_impl, _grouped_fp4_impl
+    global _fp8_fp4_mqa_logits_impl, _fp8_fp4_paged_mqa_logits_impl
+    global _get_paged_mqa_logits_metadata_impl
+    global _tf32_hc_prenorm_gemm_impl
+    global _get_mn_major_tma_aligned_tensor_impl
+    global _get_mk_alignment_for_contiguous_layout_impl
+    global _transform_sf_into_required_layout_impl
+    # fast path
+    if (
+        _cublaslt_gemm_nt_impl is not None
+        or _fp8_gemm_nt_impl is not None
+        or _fp8_einsum_impl is not None
+        or _grouped_impl is not None
+        or _grouped_masked_impl is not None
+        or _grouped_fp4_impl is not None
+        or _fp8_fp4_mqa_logits_impl is not None
+        or _fp8_fp4_paged_mqa_logits_impl is not None
+        or _get_paged_mqa_logits_metadata_impl is not None
+        or _tf32_hc_prenorm_gemm_impl is not None
+        or _get_mk_alignment_for_contiguous_layout_impl is not None
+        or _transform_sf_into_required_layout_impl is not None
+    ):
+        return
+
+    if not has_deep_gemm():
+        return
+
+    # Set up deep_gemm cache path
+    DEEP_GEMM_JIT_CACHE_ENV_NAME = "DG_JIT_CACHE_DIR"
+    if not os.environ.get(DEEP_GEMM_JIT_CACHE_ENV_NAME, None):
+        os.environ[DEEP_GEMM_JIT_CACHE_ENV_NAME] = os.path.join(
+            envs.VLLM_CACHE_ROOT, "deep_gemm"
+        )
+
+    _dg = _import_deep_gemm()
+    if _dg is None:
+        return
+
+    _cublaslt_gemm_nt_impl = getattr(_dg, "cublaslt_gemm_nt", None)
+    _fp8_gemm_nt_impl = getattr(_dg, "fp8_gemm_nt", None)
+    _fp8_einsum_impl = getattr(_dg, "fp8_einsum", None)
+    _grouped_impl = getattr(_dg, "m_grouped_fp8_gemm_nt_contiguous", None)
+    _grouped_masked_impl = getattr(_dg, "fp8_m_grouped_gemm_nt_masked", None)
+    _grouped_fp4_impl = getattr(_dg, "m_grouped_fp8_fp4_gemm_nt_contiguous", None)
+    # DeepGEMM exposes fp8_fp4_*_mqa_logits as the canonical symbols that
+    # handle both the FP8 and FP4 Q/K paths via a tuple-typed `q`.
+    _fp8_fp4_mqa_logits_impl = getattr(_dg, "fp8_fp4_mqa_logits", None)
+    _fp8_fp4_paged_mqa_logits_impl = getattr(_dg, "fp8_fp4_paged_mqa_logits", None)
+    _get_paged_mqa_logits_metadata_impl = getattr(
+        _dg, "get_paged_mqa_logits_metadata", None
+    )
+    _tf32_hc_prenorm_gemm_impl = getattr(_dg, "tf32_hc_prenorm_gemm", None)
+    _get_mn_major_tma_aligned_tensor_impl = getattr(
+        _dg, "get_mn_major_tma_aligned_tensor", None
+    )
+    _get_mk_alignment_for_contiguous_layout_impl = getattr(
+        _dg, "get_mk_alignment_for_contiguous_layout", None
+    )
+    _transform_sf_into_required_layout_impl = getattr(
+        _dg, "transform_sf_into_required_layout", None
+    )
+    DeepGemmQuantScaleFMT.init_oracle_cache()
+
+
+def get_num_sms() -> int:
+    _lazy_init()
+    dg = _import_deep_gemm()
+    if dg is None:
+        raise RuntimeError("DeepGEMM is not available")
+    return int(dg.get_num_sms())
+
+
+def set_num_sms(num_sms: int) -> None:
+    _lazy_init()
+    dg = _import_deep_gemm()
+    if dg is None:
+        raise RuntimeError("DeepGEMM is not available")
+    dg.set_num_sms(num_sms)
+
+
+@functools.cache
+def get_mk_alignment_for_contiguous_layout() -> list[int]:
+    _lazy_init()
+    if _get_mk_alignment_for_contiguous_layout_impl is None:
+        return _missing()
+    mk_align_size = _get_mk_alignment_for_contiguous_layout_impl()
+    return [mk_align_size, mk_align_size]
+
+
+def get_col_major_tma_aligned_tensor(x: torch.Tensor) -> torch.Tensor:
+    """Wrapper for DeepGEMM's get_mn_major_tma_aligned_tensor"""
+    _lazy_init()
+    if _get_mn_major_tma_aligned_tensor_impl is None:
+        return _missing()
+    return _get_mn_major_tma_aligned_tensor_impl(x)
+
+
+def cublaslt_gemm_nt(*args, **kwargs):
+    _lazy_init()
+    if _cublaslt_gemm_nt_impl is None:
+        return _missing(*args, **kwargs)
+    return _cublaslt_gemm_nt_impl(*args, **kwargs)
+
+
+def fp8_gemm_nt(*args, **kwargs):
+    _lazy_init()
+    if _fp8_gemm_nt_impl is None:
+        return _missing(*args, **kwargs)
+    if "is_deep_gemm_e8m0_used" in kwargs:
+        use_ue8m0 = kwargs["is_deep_gemm_e8m0_used"]
+        del kwargs["is_deep_gemm_e8m0_used"]
+    else:
+        use_ue8m0 = is_deep_gemm_e8m0_used()
+    return _fp8_gemm_nt_impl(*args, disable_ue8m0_cast=not use_ue8m0, **kwargs)
+
+
+def fp8_einsum(*args, **kwargs):
+    _lazy_init()
+    if _fp8_einsum_impl is None:
+        return _missing(*args, **kwargs)
+    return _fp8_einsum_impl(*args, **kwargs)
+
+
+def m_grouped_fp8_gemm_nt_contiguous(*args, **kwargs):
+    _lazy_init()
+    if _grouped_impl is None:
+        return _missing(*args, **kwargs)
+    return _grouped_impl(
+        *args, disable_ue8m0_cast=not is_deep_gemm_e8m0_used(), **kwargs
+    )
+
+
+def m_grouped_fp8_fp4_gemm_nt_contiguous(*args, **kwargs):
+    _lazy_init()
+    if _grouped_fp4_impl is None:
+        return _missing(*args, **kwargs)
+    return _grouped_fp4_impl(
+        *args, disable_ue8m0_cast=not is_deep_gemm_e8m0_used(), **kwargs
+    )
+
+
+def fp8_m_grouped_gemm_nt_masked(*args, **kwargs):
+    _lazy_init()
+    if _grouped_masked_impl is None:
+        return _missing(*args, **kwargs)
+    return _grouped_masked_impl(
+        *args, disable_ue8m0_cast=not is_deep_gemm_e8m0_used(), **kwargs
+    )
+
+
+def transform_sf_into_required_layout(*args, **kwargs):
+    _lazy_init()
+    if _transform_sf_into_required_layout_impl is None:
+        return _missing(*args, **kwargs)
+    return _transform_sf_into_required_layout_impl(
+        *args, disable_ue8m0_cast=not is_deep_gemm_e8m0_used(), **kwargs
+    )
+
+
+def fp8_fp4_mqa_logits(
+    q: tuple[torch.Tensor, torch.Tensor | None],
+    kv: tuple[torch.Tensor, torch.Tensor],
+    weights: torch.Tensor,
+    cu_seqlen_ks: torch.Tensor,
+    cu_seqlen_ke: torch.Tensor,
+    clean_logits: bool,
+) -> torch.Tensor:
+    """Compute MQA logits for a single sequence without KV paging.
+
+    Unified FP8/FP4 dispatch — the underlying DeepGEMM kernel takes
+    ``q = (values, scales_or_None)`` where ``scales`` is None for FP8 Q
+    (per-token scale is folded into ``weights``) and a packed block-scale
+    tensor for MXFP4 Q.
+
+    Args:
+        q: Tuple ``(q_values, q_scale)``. FP8 path: q_values is [M, H, D]
+            float8_e4m3fn and q_scale is None (per-token scale is folded
+            into ``weights``). FP4 path: q_values is packed uint8 and
+            q_scale is the companion block-scale tensor.
+        kv: Tuple `(k_packed, k_scales)` — FP8 layout is [N, D]
+            float8_e4m3fn plus fp32 scales [N]; FP4 layout is packed uint8.
+        weights: weights of shape [M, H], dtype `torch.float32`.
+        cu_seqlen_ks: Start indices (inclusive) for valid K per query
+            position, shape [M], dtype int32.
+        cu_seqlen_ke: End indices (exclusive) for valid K per query
+            position, shape [M], dtype int32.
+        clean_logits: Whether to clean the unfilled logits into `-inf`.
+
+    Returns:
+        Logits tensor of shape [M, N], dtype `torch.float32`.
+    """
+    _lazy_init()
+    if _fp8_fp4_mqa_logits_impl is None:
+        return _missing()
+    return _fp8_fp4_mqa_logits_impl(
+        q,
+        kv,
+        weights,
+        cu_seqlen_ks,
+        cu_seqlen_ke,
+        clean_logits=clean_logits,
+    )
+
+
+def get_paged_mqa_logits_metadata(
+    context_lens: torch.Tensor, block_size: int, num_sms: int
+) -> torch.Tensor:
+    """Build scheduling metadata for paged MQA logits.
+
+    Args:
+        context_lens: Tensor of shape [B] or [B, 1], dtype int32; effective
+            context length per batch element.
+        block_size: KV-cache block size in tokens (e.g., 64).
+        num_sms: Number of SMs available. 132 for Hopper
+
+    Returns:
+        Backend-specific tensor consumed by `fp8_fp4_paged_mqa_logits` to
+        schedule work across SMs.
+    """
+    _lazy_init()
+    if _get_paged_mqa_logits_metadata_impl is None:
+        return _missing()
+    if context_lens.dim() == 1:
+        context_lens = context_lens.unsqueeze(-1)
+    context_lens = context_lens.contiguous()
+    return _get_paged_mqa_logits_metadata_impl(context_lens, block_size, num_sms)
+
+
+def fp8_fp4_paged_mqa_logits(
+    q: tuple[torch.Tensor, torch.Tensor | None],
+    kv_cache: torch.Tensor,
+    weights: torch.Tensor,
+    context_lens: torch.Tensor,
+    block_tables: torch.Tensor,
+    schedule_metadata: torch.Tensor,
+    max_model_len: int,
+    clean_logits: bool,
+) -> torch.Tensor:
+    """Compute MQA logits using a paged KV-cache.
+
+    Unified FP8/FP4 dispatch — the underlying DeepGEMM kernel takes
+    ``q = (values, scales_or_None)``; pass ``(q_tensor, None)`` for the FP8
+    path and ``(q_values, q_scale)`` for MXFP4.
+
+    Args:
+        q: Tuple ``(q_values, q_scale)``. FP8 path: q_values is
+            [B, next_n, H, D] float8_e4m3fn and q_scale is None. FP4 path:
+            q_values is packed uint8 and q_scale is the companion
+            block-scale tensor.
+        kv_cache: Paged KV-cache. FP8 layout is [num_blocks, block_size, 1,
+            D+4], dtype `torch.uint8`, with the last 4 bytes per (block, pos)
+            storing the float dequant scale.
+        weights: Tensor of shape [B * next_n, H], dtype `torch.float32`.
+        context_lens: Tensor of shape [B], dtype int32; effective context length
+            for each batch element.
+        block_tables: Tensor of shape [B, max_blocks], dtype int32; maps logical
+            block indices to physical blocks in the paged cache.
+        schedule_metadata: Returned by `get_paged_mqa_logits_metadata`;
+            used to distribute work across SMs.
+        max_model_len: Maximum sequence length used to size the logits output.
+        clean_logits: Whether to clean the unfilled logits into `-inf`.
+
+    Returns:
+        Logits tensor of shape [B * next_n, max_model_len], dtype
+        `torch.float32`.
+    """
+    _lazy_init()
+    if _fp8_fp4_paged_mqa_logits_impl is None:
+        return _missing()
+    return _fp8_fp4_paged_mqa_logits_impl(
+        q,
+        kv_cache,
+        weights,
+        context_lens,
+        block_tables,
+        schedule_metadata,
+        max_model_len,
+        clean_logits=clean_logits,
+    )
+
+
+def tf32_hc_prenorm_gemm(
+    x: torch.Tensor,
+    fn: torch.Tensor,
+    out: torch.Tensor,
+    sqrsum: torch.Tensor,
+    num_split: int,
+) -> torch.Tensor:
+    """
+    Perform the following computation:
+        out = x.float() @ fn.T
+        sqrsum = x.float().square().sum(-1)
+
+    See the caller function for shape requirement
+    """
+    _lazy_init()
+    if _tf32_hc_prenorm_gemm_impl is None:
+        return _missing()
+    return _tf32_hc_prenorm_gemm_impl(
+        x,
+        fn,
+        out,
+        sqrsum,
+        num_split,
+    )
+
+
+def _ceil_to_ue8m0(x: torch.Tensor):
+    return torch.pow(2.0, torch.ceil(torch.log2(x.abs())))
+
+
+def _align(x: int, y: int) -> int:
+    return cdiv(x, y) * y
+
+
+# Taken from https://github.com/deepseek-ai/DeepGEMM/blob/v2.1.1/csrc/utils/math.hpp#L19
+def get_tma_aligned_size(x: int, element_size: int) -> int:
+    return _align(x, 16 // element_size)
+
+
+DEFAULT_BLOCK_SIZE = [128, 128]
+
+
+# Taken from https://github.com/deepseek-ai/DeepGEMM/blob/dd6ed14acbc7445dcef224248a77ab4d22b5f240/deep_gemm/utils/math.py#L38
+@torch.compile(dynamic=True, backend=current_platform.simple_compile_backend)
+def per_block_cast_to_fp8(
+    x: torch.Tensor, block_size: list[int] = DEFAULT_BLOCK_SIZE, use_ue8m0: bool = False
+) -> tuple[torch.Tensor, torch.Tensor]:
+    fp8_dtype = current_platform.fp8_dtype()
+    assert x.dim() == 2
+    m, n = x.shape
+    block_m, block_n = block_size
+    x_padded = torch.zeros(
+        (_align(m, block_m), _align(n, block_n)), dtype=x.dtype, device=x.device
+    )
+    x_padded[:m, :n] = x
+    x_view = x_padded.view(-1, block_m, x_padded.size(1) // block_n, block_n)
+    x_amax = x_view.abs().float().amax(dim=(1, 3), keepdim=True).clamp(1e-4)
+    _, fp8_max = get_fp8_min_max()
+    sf = x_amax / fp8_max
+    sf = _ceil_to_ue8m0(sf) if use_ue8m0 else sf
+    x_scaled = (x_view * (1.0 / sf)).to(fp8_dtype)
+    return x_scaled.view_as(x_padded)[:m, :n].contiguous(), sf.view(
+        x_view.size(0), x_view.size(2)
+    )
+
+
+def calc_diff(x: torch.Tensor, y: torch.Tensor):
+    """Return a global difference metric for unit tests.
+
+    DeepGEMM kernels on Blackwell/B200 currently exhibit noticeable per-element
+    error, causing `torch.testing.assert_close` to fail.  Instead of checking
+    every element, we compute a cosine-style similarity over the whole tensor
+    and report `1 - sim`.  Once kernel accuracy improves this helper can be
+    removed.
+    """
+
+    x, y = x.double(), y.double()
+    denominator = (x * x + y * y).sum()
+    sim = 2 * (x * y).sum() / denominator
+    return 1 - sim
+
+
+def should_use_deepgemm_for_fp8_linear(
+    output_dtype: torch.dtype,
+    weight_shape: tuple[int, int],
+    supports_deep_gemm: bool | None = None,
+):
+    if supports_deep_gemm is None:
+        supports_deep_gemm = is_deep_gemm_supported()
+
+    # Verify DeepGEMM N/K dims requirements
+    # NOTE: Also synchronized with test_w8a8_block_fp8_deep_gemm_matmul
+    # test inside kernels/quantization/test_block_fp8.py
+    N_MULTIPLE = 64
+    K_MULTIPLE = 128
+
+    return (
+        supports_deep_gemm
+        and output_dtype == torch.bfloat16
+        and weight_shape[0] % N_MULTIPLE == 0
+        and weight_shape[1] % K_MULTIPLE == 0
+    )
+
+
+__all__ = [
+    "calc_diff",
+    "DeepGemmQuantScaleFMT",
+    "fp8_gemm_nt",
+    "fp8_einsum",
+    "m_grouped_fp8_gemm_nt_contiguous",
+    "m_grouped_fp8_fp4_gemm_nt_contiguous",
+    "fp8_m_grouped_gemm_nt_masked",
+    "fp8_fp4_mqa_logits",
+    "fp8_fp4_paged_mqa_logits",
+    "get_paged_mqa_logits_metadata",
+    "per_block_cast_to_fp8",
+    "is_deep_gemm_e8m0_used",
+    "is_deep_gemm_supported",
+    "get_num_sms",
+    "set_num_sms",
+    "should_use_deepgemm_for_fp8_linear",
+    "get_col_major_tma_aligned_tensor",
+    "get_mk_alignment_for_contiguous_layout",
+]

deepseekv32_tool_parser.py

@@ -1,616 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-DeepSeek-V3.2 Tool Call Parser — re-parse-and-diff version.
-
-Adapted from the GLM-4 streaming fix to make the streaming path robust
-against multi-token deltas produced by MTP speculative decoding.
-
-Instead of maintaining incremental state that advances one token at a
-time, the streaming path re-parses the *entire* current_text on every
-call, finds all <｜DSML｜invoke> regions (complete and in-progress),
-builds a JSON arguments string for each, and diffs against what was
-previously sent.  This makes the parser agnostic to how many tokens
-arrive per step.
-
-Key changes vs. the upstream buffer-until-complete parser:
-  1. _extract_content() handles partial tag overlaps so content text
-     is never swallowed or duplicated when a tag boundary lands inside
-     a multi-token chunk.
-  2. _extract_invoke_regions() finds both complete and incomplete
-     invoke blocks, enabling streaming of partial arguments.
-  3. _build_args_json_so_far() constructs the JSON arguments string
-     incrementally from complete + partial <｜DSML｜parameter> tags.
-  4. _compute_args_diff() emits only the newly-added characters.
-
-Drop-in replacement: same class name, same interface.
-"""
-
-import json
-import uuid
-from collections.abc import Sequence
-from typing import Any
-
-import regex as re
-
-from vllm.entrypoints.openai.chat_completion.protocol import (
-    ChatCompletionRequest,
-)
-from vllm.entrypoints.openai.engine.protocol import (
-    DeltaFunctionCall,
-    DeltaMessage,
-    DeltaToolCall,
-    ExtractedToolCallInformation,
-    FunctionCall,
-    ToolCall,
-)
-from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
-from vllm.logger import init_logger
-from vllm.tokenizers import TokenizerLike
-from vllm.tool_parsers.abstract_tool_parser import (
-    Tool,
-    ToolParser,
-)
-logger = init_logger(__name__)
-
-
-def partial_tag_overlap(text: str, tag: str) -> int:
-    """Length of the longest prefix of *tag* that matches a suffix of *text*.
-
-    E.g. text ending in ``"<tool_"`` returns 6 when tag is ``"<tool_call>"``.
-    Returns 0 when there is no overlap.
-    """
-    max_check = min(len(tag) - 1, len(text))
-    for k in range(max_check, 0, -1):
-        if text.endswith(tag[:k]):
-            return k
-    return 0
-
-
-class DeepSeekV32ToolParser(ToolParser):
-    """
-    Re-parse-and-diff tool parser for DeepSeek-V3.2 DSML format.
-
-    On every streaming call the parser re-parses ``current_text`` to
-    find ``<｜DSML｜invoke>`` regions, builds the JSON arguments string
-    for each tool call, and diffs against what was previously sent to
-    emit only new content.  This is robust against multi-token deltas
-    from MTP / EAGLE speculative decoding.
-
-    Example tool call format::
-
-        <｜DSML｜function_calls>
-        <｜DSML｜invoke name="get_weather">
-        <｜DSML｜parameter name="location" string="true">杭州</｜DSML｜parameter>
-        <｜DSML｜parameter name="date" string="true">2024-01-16</｜DSML｜parameter>
-        </｜DSML｜invoke>
-        </｜DSML｜function_calls>
-    """
-
-    def __init__(self, tokenizer: TokenizerLike, tools: list[Tool] | None = None):
-        super().__init__(tokenizer, tools)
-
-        # ----- Tag constants -----
-        self.tool_call_start_token: str = "<｜DSML｜function_calls>"
-        self.tool_call_end_token: str = "</｜DSML｜function_calls>"
-        self.invoke_end_token: str = "</｜DSML｜invoke>"
-        self.param_end_token: str = "</｜DSML｜parameter>"
-
-        # Alias expected by ToolParser base / adjust_request
-        self.tool_calls_start_token = self.tool_call_start_token
-
-        # ----- Compiled regexes -----
-        # Matches a complete <｜DSML｜function_calls>…</｜DSML｜function_calls>
-        self.tool_call_complete_regex = re.compile(
-            r"<｜DSML｜function_calls>(.*?)</｜DSML｜function_calls>", re.DOTALL
-        )
-        # Opening tag of an invoke block — captures the function name.
-        self.invoke_start_regex = re.compile(
-            r'<｜DSML｜invoke\s+name="([^"]+)"\s*>', re.DOTALL
-        )
-        # Complete invoke block.
-        self.invoke_complete_regex = re.compile(
-            r'<｜DSML｜invoke\s+name="([^"]+)"\s*>(.*?)</｜DSML｜invoke>',
-            re.DOTALL,
-        )
-        # Complete parameter tag — captures (name, string_attr, value).
-        self.parameter_complete_regex = re.compile(
-            r'<｜DSML｜parameter\s+name="([^"]+)"\s+string="(true|false)"\s*>'
-            r"(.*?)"
-            r"</｜DSML｜parameter>",
-            re.DOTALL,
-        )
-        # Just the opening header of a parameter tag (for partial params).
-        self.parameter_header_regex = re.compile(
-            r'<｜DSML｜parameter\s+name="([^"]+)"\s+string="(true|false)"\s*>',
-            re.DOTALL,
-        )
-
-        # ----- Streaming state (reset per request) -----
-        self._sent_content_idx: int = 0
-        self._tool_call_ids: list[str] = []
-        self.streamed_args_for_tool: list[str] = []
-        self.prev_tool_call_arr: list[dict[str, Any]] = []
-        self.current_tool_id: int = -1
-
-        if not self.model_tokenizer:
-            raise ValueError(
-                "The model tokenizer must be passed to the ToolParser "
-                "constructor during construction."
-            )
-
-        logger.debug(
-            "Successfully initialized %s", self.__class__.__name__
-        )
-
-    # ------------------------------------------------------------------
-    # Request adjustment
-    # ------------------------------------------------------------------
-
-    def adjust_request(
-        self, request: ChatCompletionRequest | ResponsesRequest
-    ) -> ChatCompletionRequest | ResponsesRequest:
-        request = super().adjust_request(request)
-        if request.tools and request.tool_choice != "none":
-            # Ensure DSML tokens are not stripped during decoding.
-            request.skip_special_tokens = False
-        return request
-
-    # ------------------------------------------------------------------
-    # Static / utility helpers
-    # ------------------------------------------------------------------
-
-    @staticmethod
-    def _tools_enabled(request: ChatCompletionRequest) -> bool:
-        """Check whether tool calling is active for this request."""
-        try:
-            tools = getattr(request, "tools", None)
-            tool_choice = getattr(request, "tool_choice", None)
-            return bool(tools) and tool_choice != "none"
-        except Exception:
-            logger.exception("Failed to determine if tools are enabled.")
-            return False
-
-    def _generate_tool_call_id(self) -> str:
-        return f"call_{uuid.uuid4().hex[:24]}"
-
-    @staticmethod
-    def _json_escape_string_content(s: str) -> str:
-        """JSON-escape a string value (without surrounding quotes)."""
-        if not s:
-            return ""
-        return json.dumps(s, ensure_ascii=False)[1:-1]
-
-    # ------------------------------------------------------------------
-    # Type conversion helpers
-    # ------------------------------------------------------------------
-
-    def _convert_param_value_checked(self, value: str, param_type: str) -> Any:
-        """Convert a raw string value to the type indicated by *param_type*.
-
-        Raises on failure so the caller can try the next candidate type.
-        """
-        if value.lower() == "null":
-            return None
-
-        param_type = param_type.lower()
-        if param_type in ("string", "str", "text"):
-            return value
-        elif param_type in ("integer", "int"):
-            return int(value)
-        elif param_type in ("number", "float"):
-            val = float(value)
-            return val if val != int(val) else int(val)
-        elif param_type in ("boolean", "bool"):
-            normed = value.strip().lower()
-            if normed not in ("false", "0", "true", "1"):
-                raise ValueError(f"Invalid boolean value: {value!r}")
-            return normed in ("true", "1")
-        elif param_type in ("object", "array"):
-            return json.loads(value)
-        else:
-            return json.loads(value)
-
-    def _convert_param_value(self, value: str, param_type: str | list[str]) -> Any:
-        """Try each candidate type in turn; fall back to the raw string."""
-        if not isinstance(param_type, list):
-            param_type = [param_type]
-        for current_type in param_type:
-            try:
-                return self._convert_param_value_checked(value, current_type)
-            except Exception:
-                continue
-        return value
-
-    def _get_param_schema_type(
-        self, func_name: str, param_name: str
-    ) -> str | list[str]:
-        """Look up the JSON-schema type for a parameter, defaulting to
-        ``"string"``."""
-        if self.tools:
-            for tool in self.tools:
-                if (
-                    hasattr(tool, "function")
-                    and tool.function.name == func_name
-                    and hasattr(tool.function, "parameters")
-                ):
-                    schema = tool.function.parameters
-                    if isinstance(schema, dict) and "properties" in schema:
-                        prop = schema["properties"].get(param_name, {})
-                        if isinstance(prop, dict):
-                            return prop.get("type", "string")
-                    break
-        return "string"
-
-    def _convert_with_schema(
-        self, func_name: str, param_name: str, value: str
-    ) -> Any:
-        """Convert *value* using the tool schema for *func_name*.*param_name*."""
-        param_type = self._get_param_schema_type(func_name, param_name)
-        return self._convert_param_value(value, param_type)
-
-    def _is_string_type(self, func_name: str, param_name: str) -> bool:
-        """Return True if the schema says this parameter is a string."""
-        ptype = self._get_param_schema_type(func_name, param_name)
-        if isinstance(ptype, list):
-            return "string" in ptype
-        return ptype in ("string", "str", "text")
-
-    # ------------------------------------------------------------------
-    # Non-streaming extraction (unchanged logic, shared helpers)
-    # ------------------------------------------------------------------
-
-    def extract_tool_calls(
-        self,
-        model_output: str,
-        request: ChatCompletionRequest,
-    ) -> ExtractedToolCallInformation:
-        """Extract tool calls from complete model output (non-streaming)."""
-        if self.tool_call_start_token not in model_output:
-            return ExtractedToolCallInformation(
-                tools_called=False, tool_calls=[], content=model_output
-            )
-
-        try:
-            tool_calls: list[ToolCall] = []
-
-            for fc_block in self.tool_call_complete_regex.findall(model_output):
-                for invoke_name, invoke_body in self.invoke_complete_regex.findall(
-                    fc_block
-                ):
-                    # Parse all parameters in this invoke.
-                    raw_params: dict[str, str] = {}
-                    for pname, _str_attr, pval in (
-                        self.parameter_complete_regex.findall(invoke_body)
-                    ):
-                        raw_params[pname] = pval
-
-                    # Convert types via schema.
-                    converted: dict[str, Any] = {}
-                    for pname, pval in raw_params.items():
-                        converted[pname] = self._convert_with_schema(
-                            invoke_name, pname, pval
-                        )
-
-                    tool_calls.append(
-                        ToolCall(
-                            type="function",
-                            function=FunctionCall(
-                                name=invoke_name,
-                                arguments=json.dumps(
-                                    converted, ensure_ascii=False
-                                ),
-                            ),
-                        )
-                    )
-
-            if not tool_calls:
-                return ExtractedToolCallInformation(
-                    tools_called=False, tool_calls=[], content=model_output
-                )
-
-            first_idx = model_output.find(self.tool_call_start_token)
-            content = model_output[:first_idx] if first_idx > 0 else None
-
-            return ExtractedToolCallInformation(
-                tools_called=True, tool_calls=tool_calls, content=content
-            )
-
-        except Exception:
-            logger.exception("Error extracting tool calls from complete output")
-            return ExtractedToolCallInformation(
-                tools_called=False, tool_calls=[], content=model_output
-            )
-
-    # ------------------------------------------------------------------
-    # Streaming helpers — re-parse-and-diff
-    # ------------------------------------------------------------------
-
-    def _reset_streaming_state(self) -> None:
-        self._sent_content_idx = 0
-        self._tool_call_ids.clear()
-        self.streamed_args_for_tool.clear()
-        self.prev_tool_call_arr.clear()
-        self.current_tool_id = -1
-
-    def _extract_content(self, current_text: str) -> str | None:
-        """Return any non-tool-call text that hasn't been sent yet.
-
-        Walks *current_text* from ``_sent_content_idx``, collecting text
-        outside ``<｜DSML｜function_calls>`` regions.  Uses
-        ``partial_tag_overlap`` to avoid emitting bytes that might turn
-        out to be the start of the function-calls tag once the next
-        chunk arrives.
-        """
-        content_segments: list[str] = []
-        pos = self._sent_content_idx
-
-        while pos < len(current_text):
-            start = current_text.find(self.tool_call_start_token, pos)
-            if start == -1:
-                # No (more) tool-call regions — send the tail, minus
-                # any suffix that could be the beginning of the tag.
-                tail = current_text[pos:]
-                overlap = partial_tag_overlap(tail, self.tool_call_start_token)
-                sendable = tail[: len(tail) - overlap] if overlap else tail
-                if sendable:
-                    content_segments.append(sendable)
-                pos = len(current_text) - overlap
-                break
-
-            # Text between previous position and the tag start is content.
-            if start > pos:
-                content_segments.append(current_text[pos:start])
-
-            # Skip past the tool-call region.
-            end = current_text.find(self.tool_call_end_token, start)
-            if end != -1:
-                pos = end + len(self.tool_call_end_token)
-            else:
-                # Region still open — park cursor at start, stop.
-                pos = start
-                break
-
-        if content_segments:
-            self._sent_content_idx = pos
-            return "".join(content_segments)
-        if pos > self._sent_content_idx:
-            self._sent_content_idx = pos
-        return None
-
-    def _extract_invoke_regions(
-        self, text: str
-    ) -> list[tuple[str, str, bool]]:
-        """Find all invoke blocks inside the function_calls region.
-
-        Returns a list of ``(func_name, inner_text, is_complete)``
-        tuples.  *inner_text* is everything between the invoke open
-        tag and the close tag (or the end of available text for the
-        last, potentially incomplete, invoke).
-        """
-        results: list[tuple[str, str, bool]] = []
-
-        fc_start = text.find(self.tool_call_start_token)
-        if fc_start == -1:
-            return results
-
-        region_start = fc_start + len(self.tool_call_start_token)
-        fc_end = text.find(self.tool_call_end_token, region_start)
-        region = text[region_start:fc_end] if fc_end != -1 else text[region_start:]
-
-        pos = 0
-        while pos < len(region):
-            inv_match = self.invoke_start_regex.search(region, pos)
-            if not inv_match:
-                break
-
-            func_name = inv_match.group(1)
-            body_start = inv_match.end()
-
-            inv_end_pos = region.find(self.invoke_end_token, body_start)
-            if inv_end_pos != -1:
-                # Complete invoke block.
-                body = region[body_start:inv_end_pos]
-                results.append((func_name, body, True))
-                pos = inv_end_pos + len(self.invoke_end_token)
-            else:
-                # Incomplete — still being generated.
-                body = region[body_start:]
-                overlap = partial_tag_overlap(body, self.invoke_end_token)
-                if overlap:
-                    body = body[:-overlap]
-                results.append((func_name, body, False))
-                break
-
-        return results
-
-    def _build_args_json_so_far(
-        self,
-        func_name: str,
-        inner_text: str,
-        is_complete: bool,
-    ) -> str:
-        """Build a JSON arguments string from the parameters found so far.
-
-        Handles both fully-closed ``<｜DSML｜parameter>`` tags and the
-        single trailing partial parameter whose value is still being
-        streamed.
-        """
-        # ---- Collect all fully-closed parameters ----
-        complete_params = self.parameter_complete_regex.findall(inner_text)
-        parts: list[str] = []
-
-        for param_name, string_attr, param_value in complete_params:
-            key_json = json.dumps(param_name, ensure_ascii=False)
-            if string_attr == "true":
-                val_json = json.dumps(param_value, ensure_ascii=False)
-            else:
-                converted = self._convert_with_schema(
-                    func_name, param_name, param_value
-                )
-                val_json = json.dumps(converted, ensure_ascii=False)
-            parts.append(f"{key_json}: {val_json}")
-
-        # ---- Handle a trailing partial parameter ----
-        last_param_open = inner_text.rfind("<｜DSML｜parameter")
-        last_param_close = inner_text.rfind(self.param_end_token)
-        has_partial = last_param_open != -1 and (
-            last_param_close == -1 or last_param_close < last_param_open
-        )
-
-        if has_partial:
-            partial_text = inner_text[last_param_open:]
-            header_match = self.parameter_header_regex.search(partial_text)
-
-            if header_match:
-                param_name = header_match.group(1)
-                string_attr = header_match.group(2)
-                partial_value = partial_text[header_match.end():]
-
-                # Strip any bytes that might be the beginning of the
-                # closing </｜DSML｜parameter> tag.
-                overlap = partial_tag_overlap(
-                    partial_value, self.param_end_token
-                )
-                if overlap:
-                    partial_value = partial_value[:-overlap]
-
-                key_json = json.dumps(param_name, ensure_ascii=False)
-
-                if is_complete:
-                    # Invoke is closed — treat whatever we have as final.
-                    if string_attr == "true":
-                        val_json = json.dumps(
-                            partial_value, ensure_ascii=False
-                        )
-                    else:
-                        converted = self._convert_with_schema(
-                            func_name, param_name, partial_value
-                        )
-                        val_json = json.dumps(converted, ensure_ascii=False)
-                    parts.append(f"{key_json}: {val_json}")
-                elif string_attr == "true" or self._is_string_type(
-                    func_name, param_name
-                ):
-                    # Stream as an open JSON string (no closing quote).
-                    escaped = self._json_escape_string_content(partial_value)
-                    parts.append(f'{key_json}: "{escaped}')
-                else:
-                    # Non-string — emit raw partial value.
-                    parts.append(f"{key_json}: {partial_value}")
-
-        # ---- Assemble ----
-        if not parts:
-            return "{}" if is_complete else ""
-
-        joined = "{" + ", ".join(parts)
-        if is_complete:
-            joined += "}"
-        return joined
-
-    def _compute_args_diff(self, index: int, args_so_far: str) -> str | None:
-        """Return only the characters in *args_so_far* that haven't been
-        sent yet, or ``None`` if there's nothing new."""
-        prev = self.streamed_args_for_tool[index]
-        if not args_so_far or len(args_so_far) <= len(prev):
-            return None
-        diff = args_so_far[len(prev):]
-        self.streamed_args_for_tool[index] = args_so_far
-        self.prev_tool_call_arr[index]["arguments"] = args_so_far
-        return diff
-
-    def _ensure_tool_state_for(self, index: int) -> None:
-        """Grow the streaming-state arrays so *index* is valid."""
-        while len(self._tool_call_ids) <= index:
-            self._tool_call_ids.append(self._generate_tool_call_id())
-        while len(self.streamed_args_for_tool) <= index:
-            self.streamed_args_for_tool.append("")
-        while len(self.prev_tool_call_arr) <= index:
-            self.prev_tool_call_arr.append({})
-
-    # ------------------------------------------------------------------
-    # Main streaming entry point
-    # ------------------------------------------------------------------
-
-    def extract_tool_calls_streaming(
-        self,
-        previous_text: str,
-        current_text: str,
-        delta_text: str,
-        previous_token_ids: Sequence[int],
-        current_token_ids: Sequence[int],
-        delta_token_ids: Sequence[int],
-        request: ChatCompletionRequest,
-    ) -> DeltaMessage | None:
-        """Extract tool calls from streaming output using re-parse-and-diff.
-
-        On every call we:
-          1. Re-scan *current_text* for content outside tool-call regions.
-          2. Find all ``<｜DSML｜invoke>`` regions (complete + partial).
-          3. Build JSON args for each, diff against previous, emit deltas.
-
-        Because the entire text is re-parsed each time, the result is
-        correct regardless of how many tokens arrived in this step.
-        """
-        # First chunk of a new stream — reset state.
-        if not previous_text:
-            self._reset_streaming_state()
-
-        # If tools aren't enabled, just forward content.
-        if not self._tools_enabled(request):
-            return DeltaMessage(content=delta_text) if delta_text else None
-
-        # 1. Extract any content outside tool-call regions.
-        content = self._extract_content(current_text)
-
-        # 2. Find all invoke regions.
-        regions = self._extract_invoke_regions(current_text)
-        tool_call_deltas: list[DeltaToolCall] = []
-
-        for i, (func_name, inner_text, is_complete) in enumerate(regions):
-            self._ensure_tool_state_for(i)
-
-            # Emit the tool name (once per tool call).
-            if "name" not in self.prev_tool_call_arr[i]:
-                self.prev_tool_call_arr[i]["name"] = func_name
-                tool_call_deltas.append(
-                    DeltaToolCall(
-                        index=i,
-                        id=self._tool_call_ids[i],
-                        type="function",
-                        function=DeltaFunctionCall(
-                            name=func_name,
-                            arguments="",
-                        ),
-                    )
-                )
-
-            # Build the JSON args so far and emit the diff.
-            args_so_far = self._build_args_json_so_far(
-                func_name, inner_text, is_complete
-            )
-            diff = self._compute_args_diff(i, args_so_far)
-            if diff:
-                tool_call_deltas.append(
-                    DeltaToolCall(
-                        index=i,
-                        function=DeltaFunctionCall(arguments=diff),
-                    )
-                )
-
-        if regions:
-            self.current_tool_id = len(regions) - 1
-
-        # 3. Return a delta if we have content or tool-call updates.
-        if content or tool_call_deltas:
-            return DeltaMessage(
-                content=content,
-                tool_calls=tool_call_deltas,
-            )
-
-        # Empty delta with token ids means EOS or closing tag — return
-        # non-None so the serving framework can finalize finish_reason.
-        if not delta_text and delta_token_ids and self.prev_tool_call_arr:
-            return DeltaMessage(content="")
-
-        return None

glm4_moe_tool_parser.py

@@ -1,491 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-GLM-4/5 Tool Call Parser — fixed version.
-
-Fixes applied over the upstream vLLM + sweetapi patch:
-
-1. **func_detail_regex no longer requires a newline** between tool name and
-   first <arg_key>.  The model's chat template instructs:
-       <tool_call>{name}<arg_key>…</arg_key><arg_value>…</arg_value>…</tool_call>
-   with NO mandatory newline, but the original regex used ``[^\\n]*\\n`` which
-   silently failed when the model omitted it.
-
-2. **Zero-argument tool calls no longer crash** (TypeError on NoneType).
-
-3. **extract_tool_calls uses the same robust extraction helpers** as the
-   streaming path, so both paths parse identically.
-
-4. **_extract_tool_name_from_region** is more tolerant of whitespace /
-   formatting variants the model may produce.
-
-Drop this file into your vLLM install as a --tool-parser-plugin, or replace
-the built-in glm4_moe_tool_parser.py.
-"""
-
-import ast
-import json
-from collections.abc import Sequence
-from typing import Any
-
-import regex as re
-
-from vllm.entrypoints.chat_utils import make_tool_call_id
-from vllm.entrypoints.openai.chat_completion.protocol import (
-    ChatCompletionRequest,
-)
-from vllm.entrypoints.openai.engine.protocol import (
-    DeltaFunctionCall,
-    DeltaMessage,
-    DeltaToolCall,
-    ExtractedToolCallInformation,
-    FunctionCall,
-    ToolCall,
-)
-from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
-from vllm.logger import init_logger
-from vllm.tokenizers import TokenizerLike
-from vllm.tool_parsers.abstract_tool_parser import (
-    Tool,
-    ToolParser,
-)
-from vllm.tool_parsers.utils import partial_tag_overlap
-
-logger = init_logger(__name__)
-
-
-class Glm4MoeModelToolParser(ToolParser):
-    """Tool parser for GLM-4/5 models with incremental string streaming.
-
-    On every streaming call the parser re-parses ``current_text`` to find
-    ``<tool_call>`` regions, builds the JSON arguments string for each tool
-    call, and diffs against what was previously sent to emit only new content.
-    """
-
-    def __init__(self, tokenizer: TokenizerLike, tools: list[Tool] | None = None):
-        super().__init__(tokenizer, tools)
-        # Stateful streaming fields
-        self.current_tool_name_sent: bool = False
-        self.prev_tool_call_arr: list[dict[str, Any]] = []
-        self.current_tool_id: int = -1
-        self.streamed_args_for_tool: list[str] = []
-
-        self.tool_call_start_token: str = "<tool_call>"
-        self.tool_call_end_token: str = "</tool_call>"
-        self.arg_key_start: str = "<arg_key>"
-        self.arg_key_end: str = "</arg_key>"
-        self.arg_val_start: str = "<arg_value>"
-        self.arg_val_end: str = "</arg_value>"
-
-        self.tool_calls_start_token = self.tool_call_start_token
-
-        # ---- FIXED regexes ------------------------------------------------
-        # Match the whole <tool_call>…</tool_call> block (unchanged).
-        self.func_call_regex = re.compile(
-            r"<tool_call>.*?</tool_call>", re.DOTALL
-        )
-
-        # FIX 1: The original regex required a literal \n between tool name
-        # and the body.  The model often omits it.  We now accept any
-        # whitespace (including none) before the first <arg_key>, and we
-        # make the body group optional so zero-argument calls don't fail.
-        self.func_detail_regex = re.compile(
-            r"<tool_call>\s*"          # opening tag + optional whitespace
-            r"([\w.\-]+)"             # group 1: tool/function name (word chars, dots, hyphens)
-            r"\s*"                     # optional whitespace / newline
-            r"((?:<arg_key>.*)?)"      # group 2: everything from first <arg_key> onward (may be empty)
-            r"\s*</tool_call>",        # closing tag
-            re.DOTALL,
-        )
-
-        self.func_arg_regex = re.compile(
-            r"<arg_key>(.*?)</arg_key>\s*<arg_value>(.*?)</arg_value>", re.DOTALL
-        )
-
-        if not self.model_tokenizer:
-            raise ValueError(
-                "The model tokenizer must be passed to the ToolParser "
-                "constructor during construction."
-            )
-
-        self.tool_call_start_token_id = self.vocab.get(self.tool_call_start_token)
-        self.tool_call_end_token_id = self.vocab.get(self.tool_call_end_token)
-
-        # Pre-compiled pattern for finding the last <arg_key>...</arg_key>
-        # before a partial <arg_value> (used in _build_args_json_so_far).
-        self._arg_key_pattern = re.compile(
-            re.escape(self.arg_key_start) + r"(.*?)" + re.escape(self.arg_key_end),
-            re.DOTALL,
-        )
-
-        # Streaming state for re-parse-and-diff approach
-        self._sent_content_idx: int = 0
-        self._tool_call_ids: list[str] = []
-
-    # ------------------------------------------------------------------
-    # Static helpers
-    # ------------------------------------------------------------------
-
-    @staticmethod
-    def _deserialize(value: str) -> Any:
-        try:
-            return json.loads(value)
-        except json.JSONDecodeError:
-            pass
-        try:
-            return ast.literal_eval(value)
-        except (ValueError, SyntaxError):
-            pass
-        return value
-
-    @staticmethod
-    def _json_escape_string_content(s: str) -> str:
-        """JSON-escape string content (without surrounding quotes)."""
-        if not s:
-            return ""
-        return json.dumps(s, ensure_ascii=False)[1:-1]
-
-    @staticmethod
-    def _is_string_type(
-        tool_name: str,
-        arg_name: str,
-        tools: list[Tool] | None,
-    ) -> bool:
-        if tools is None:
-            return False
-        for tool in tools:
-            if tool.function.name != tool_name:
-                continue
-            if tool.function.parameters is None:
-                return False
-            arg_type = (
-                tool.function.parameters.get("properties", {})
-                .get(arg_name, {})
-                .get("type", None)
-            )
-            return arg_type == "string"
-        logger.debug("No tool named '%s'.", tool_name)
-        return False
-
-    @staticmethod
-    def _tools_enabled(request: ChatCompletionRequest) -> bool:
-        try:
-            tools = getattr(request, "tools", None)
-            tool_choice = getattr(request, "tool_choice", None)
-            return bool(tools) and tool_choice != "none"
-        except Exception:
-            logger.exception("Failed to determine if tools are enabled.")
-            return False
-
-    # ------------------------------------------------------------------
-    # Request adjustment
-    # ------------------------------------------------------------------
-
-    def adjust_request(
-        self, request: ChatCompletionRequest | ResponsesRequest
-    ) -> ChatCompletionRequest | ResponsesRequest:
-        request = super().adjust_request(request)
-        if request.tools and request.tool_choice != "none":
-            request.skip_special_tokens = False
-        return request
-
-    # ------------------------------------------------------------------
-    # Non-streaming extraction
-    # ------------------------------------------------------------------
-
-    def extract_tool_calls(
-        self,
-        model_output: str,
-        request: ChatCompletionRequest,
-    ) -> ExtractedToolCallInformation:
-        matched_tool_calls = self.func_call_regex.findall(model_output)
-        logger.debug("model_output: %s", model_output)
-
-        try:
-            tool_calls: list[ToolCall] = []
-            for match in matched_tool_calls:
-                tc_detail = self.func_detail_regex.search(match)
-                if not tc_detail:
-                    logger.warning(
-                        "Failed to parse tool call details from: %s", match
-                    )
-                    continue
-
-                tc_name = tc_detail.group(1).strip()
-                tc_args_raw = tc_detail.group(2) or ""  # FIX 2: default to ""
-                pairs = self.func_arg_regex.findall(tc_args_raw) if tc_args_raw else []
-                arg_dct: dict[str, Any] = {}
-                for key, value in pairs:
-                    arg_key = key.strip()
-                    arg_val = value.strip()
-                    if not self._is_string_type(tc_name, arg_key, self.tools):
-                        arg_val = self._deserialize(arg_val)
-                    logger.debug("arg_key = %s, arg_val = %s", arg_key, arg_val)
-                    arg_dct[arg_key] = arg_val
-                tool_calls.append(
-                    ToolCall(
-                        type="function",
-                        function=FunctionCall(
-                            name=tc_name,
-                            arguments=json.dumps(arg_dct, ensure_ascii=False),
-                        ),
-                    )
-                )
-        except Exception:
-            logger.exception("Failed to extract tool call spec")
-            return ExtractedToolCallInformation(
-                tools_called=False, tool_calls=[], content=model_output
-            )
-
-        if tool_calls:
-            content: str | None = model_output[
-                : model_output.find(self.tool_calls_start_token)
-            ]
-            if not content or not content.strip():
-                content = None
-            return ExtractedToolCallInformation(
-                tools_called=True, tool_calls=tool_calls, content=content
-            )
-        return ExtractedToolCallInformation(
-            tools_called=False, tool_calls=[], content=model_output
-        )
-
-    # ------------------------------------------------------------------
-    # Streaming helpers
-    # ------------------------------------------------------------------
-
-    def _extract_content(self, current_text: str) -> str | None:
-        content_segments: list[str] = []
-        pos = self._sent_content_idx
-
-        while pos < len(current_text):
-            start = current_text.find(self.tool_call_start_token, pos)
-            if start == -1:
-                tail = current_text[pos:]
-                overlap = partial_tag_overlap(tail, self.tool_call_start_token)
-                sendable = tail[: len(tail) - overlap] if overlap else tail
-                if sendable:
-                    content_segments.append(sendable)
-                pos = len(current_text) - overlap
-                break
-
-            if start > pos:
-                content_segments.append(current_text[pos:start])
-
-            end = current_text.find(self.tool_call_end_token, start)
-            if end != -1:
-                pos = end + len(self.tool_call_end_token)
-            else:
-                pos = start
-                break
-
-        if content_segments:
-            self._sent_content_idx = pos
-            return "".join(content_segments)
-        if pos > self._sent_content_idx:
-            self._sent_content_idx = pos
-        return None
-
-    def _extract_tool_call_regions(self, text: str) -> list[tuple[str, bool]]:
-        results: list[tuple[str, bool]] = []
-        pos = 0
-        while True:
-            start = text.find(self.tool_call_start_token, pos)
-            if start == -1:
-                break
-            inner_start = start + len(self.tool_call_start_token)
-            end = text.find(self.tool_call_end_token, inner_start)
-            if end != -1:
-                results.append((text[inner_start:end], True))
-                pos = end + len(self.tool_call_end_token)
-            else:
-                raw = text[inner_start:]
-                overlap = partial_tag_overlap(raw, self.tool_call_end_token)
-                if overlap:
-                    raw = raw[:-overlap]
-                results.append((raw, False))
-                break
-        return results
-
-    def _extract_tool_name_from_region(self, inner_text: str) -> str | None:
-        """Extract the tool name from the beginning of a tool-call region.
-
-        The name is everything before the first ``\\n``, ``<arg_key>``, or
-        ``</tool_call>``.  We also accept the name being the only content
-        (for zero-argument calls that are still in-flight).
-        """
-        # Strip leading whitespace — model may emit \n after <tool_call>
-        stripped = inner_text.lstrip()
-        if not stripped:
-            return None
-
-        nl = stripped.find("\n")
-        ak = stripped.find(self.arg_key_start)
-        candidates = [i for i in [nl, ak] if i != -1]
-        if not candidates:
-            # No delimiter yet — if the text looks like a partial name
-            # (only word chars / dots / hyphens), return None to wait.
-            # If it's a complete name with no args (zero-arg call, complete),
-            # it will be handled when is_complete is True.
-            candidate_name = stripped.strip()
-            if re.fullmatch(r'[\w.\-]+', candidate_name):
-                # Could be a complete name or still arriving — return it
-                # so zero-arg complete calls work; the caller checks is_complete.
-                return candidate_name
-            return None
-        cut = min(candidates)
-        name = stripped[:cut].strip()
-        return name if name else None
-
-    def _build_args_json_so_far(
-        self,
-        tool_name: str,
-        inner_text: str,
-        is_complete: bool,
-    ) -> str:
-        pairs = self.func_arg_regex.findall(inner_text)
-
-        parts: list[str] = []
-        for key, value in pairs:
-            key = key.strip()
-            key_json = json.dumps(key, ensure_ascii=False)
-            if self._is_string_type(tool_name, key, self.tools):
-                val_json = json.dumps(value, ensure_ascii=False)
-            else:
-                val_json = json.dumps(
-                    self._deserialize(value.strip()), ensure_ascii=False
-                )
-            parts.append(f"{key_json}: {val_json}")
-
-        # Check for a partial (incomplete) arg value
-        last_val_start = inner_text.rfind(self.arg_val_start)
-        last_val_end = inner_text.rfind(self.arg_val_end)
-        has_partial_value = last_val_start != -1 and (
-            last_val_end == -1 or last_val_end < last_val_start
-        )
-
-        if has_partial_value:
-            last_key_match = None
-            for m in self._arg_key_pattern.finditer(inner_text[:last_val_start]):
-                last_key_match = m
-
-            if last_key_match:
-                partial_key = last_key_match.group(1).strip()
-                partial_content_start = last_val_start + len(self.arg_val_start)
-                partial_content = inner_text[partial_content_start:]
-
-                overlap = partial_tag_overlap(partial_content, self.arg_val_end)
-                if overlap:
-                    partial_content = partial_content[:-overlap]
-
-                key_json = json.dumps(partial_key, ensure_ascii=False)
-                if is_complete:
-                    if self._is_string_type(tool_name, partial_key, self.tools):
-                        val_json = json.dumps(partial_content, ensure_ascii=False)
-                    else:
-                        val_json = json.dumps(
-                            self._deserialize(partial_content.strip()),
-                            ensure_ascii=False,
-                        )
-                    parts.append(f"{key_json}: {val_json}")
-                elif self._is_string_type(tool_name, partial_key, self.tools):
-                    escaped = self._json_escape_string_content(partial_content)
-                    parts.append(f'{key_json}: "{escaped}')
-                else:
-                    parts.append(f"{key_json}: {partial_content}")
-
-        if not parts:
-            return "{}" if is_complete else ""
-
-        joined = "{" + ", ".join(parts)
-        if is_complete:
-            joined += "}"
-        return joined
-
-    def _compute_args_diff(self, index: int, args_so_far: str) -> str | None:
-        if not args_so_far or len(args_so_far) <= len(
-            self.streamed_args_for_tool[index]
-        ):
-            return None
-        diff = args_so_far[len(self.streamed_args_for_tool[index]) :]
-        self.streamed_args_for_tool[index] = args_so_far
-        self.prev_tool_call_arr[index]["arguments"] = args_so_far
-        return diff
-
-    def _ensure_tool_state_for(self, index: int) -> None:
-        while len(self._tool_call_ids) <= index:
-            self._tool_call_ids.append(
-                make_tool_call_id(id_type="random", func_name=None, idx=None)
-            )
-        while len(self.streamed_args_for_tool) <= index:
-            self.streamed_args_for_tool.append("")
-        while len(self.prev_tool_call_arr) <= index:
-            self.prev_tool_call_arr.append({})
-
-    # ------------------------------------------------------------------
-    # Main streaming entry point
-    # ------------------------------------------------------------------
-
-    def extract_tool_calls_streaming(
-        self,
-        previous_text: str,
-        current_text: str,
-        delta_text: str,
-        previous_token_ids: Sequence[int],
-        current_token_ids: Sequence[int],
-        delta_token_ids: Sequence[int],
-        request: ChatCompletionRequest,
-    ) -> DeltaMessage | None:
-        if not self._tools_enabled(request):
-            return DeltaMessage(content=delta_text) if delta_text else None
-
-        content = self._extract_content(current_text)
-        regions = self._extract_tool_call_regions(current_text)
-        tool_call_deltas: list[DeltaToolCall] = []
-
-        for i, (inner_text, is_complete) in enumerate(regions):
-            self._ensure_tool_state_for(i)
-
-            tool_name = self._extract_tool_name_from_region(inner_text)
-            if not tool_name:
-                break
-
-            # Emit tool name (once per tool call)
-            if "name" not in self.prev_tool_call_arr[i]:
-                self.prev_tool_call_arr[i]["name"] = tool_name
-                tool_call_deltas.append(
-                    DeltaToolCall(
-                        index=i,
-                        id=self._tool_call_ids[i],
-                        type="function",
-                        function=DeltaFunctionCall(
-                            name=tool_name,
-                            arguments="",
-                        ).model_dump(exclude_none=True),
-                    )
-                )
-
-            # Build args JSON so far, diff, emit
-            args_so_far = self._build_args_json_so_far(
-                tool_name, inner_text, is_complete
-            )
-            diff = self._compute_args_diff(i, args_so_far)
-            if diff:
-                tool_call_deltas.append(
-                    DeltaToolCall(
-                        index=i,
-                        function=DeltaFunctionCall(arguments=diff).model_dump(
-                            exclude_none=True
-                        ),
-                    )
-                )
-
-        if regions:
-            self.current_tool_id = len(regions) - 1
-
-        if content or tool_call_deltas:
-            return DeltaMessage(
-                content=content,
-                tool_calls=tool_call_deltas,
-            )
-        return None

glm_5.1_chat_template.jinja

@@ -0,0 +1,119 @@
+[gMASK]<sop>
+{%- if tools -%}
+{%- macro tool_to_json(tool) -%}
+    {%- set ns_tool = namespace(first=true) -%}
+    {{ '{' -}}
+    {%- for k, v in tool.items() -%}
+        {%- if k != 'defer_loading' and k != 'strict' -%}
+            {%- if not ns_tool.first -%}{{- ', ' -}}{%- endif -%}
+            {%- set ns_tool.first = false -%}
+            "{{ k }}": {{ v | tojson(ensure_ascii=False) }}
+        {%- endif -%}
+    {%- endfor -%}
+    {{- '}' -}}
+{%- endmacro -%}
+<|system|>
+# Tools
+
+You may call one or more functions to assist with the user query.
+
+You are provided with function signatures within <tools></tools> XML tags:
+<tools>
+{% for tool in tools %}
+{%- if 'function' in tool -%}
+    {%- set tool = tool['function'] -%}
+{%- endif -%}
+{% if tool.defer_loading is not defined or not tool.defer_loading %}
+{{ tool_to_json(tool) }}
+{% endif %}
+{% endfor %}
+</tools>
+
+For each function call, output the function name and arguments within the following XML format:
+<tool_call>{function-name}<arg_key>{arg-key-1}</arg_key><arg_value>{arg-value-1}</arg_value><arg_key>{arg-key-2}</arg_key><arg_value>{arg-value-2}</arg_value>...</tool_call>{%- endif -%}
+{%- macro visible_text(content) -%}
+    {%- if content is string -%}
+        {{- content }}
+    {%- elif content is iterable and content is not mapping -%}
+        {%- for item in content -%}
+            {%- if item is mapping and item.type == 'text' -%}
+                {{- item.text }}
+            {%- elif item is string -%}
+                {{- item }}
+            {%- endif -%}
+        {%- endfor -%}
+    {%- else -%}
+        {{- content }}
+    {%- endif -%}
+{%- endmacro -%}
+{%- set ns = namespace(last_user_index=-1, thinking_indices='') -%}
+{%- for m in messages %}
+    {%- if m.role == 'user' %}
+        {%- set ns.last_user_index = loop.index0 -%}
+    {%- elif m.role == 'assistant' %}
+        {%- if m.reasoning_content is string %}
+            {%- set ns.thinking_indices = ns.thinking_indices ~ ',' ~ ns.last_user_index ~ ',' -%}
+        {%- endif %}
+    {%- endif %}
+{%- endfor %}
+{%- set ns.has_thinking = false -%}
+{%- for m in messages -%}
+{%- if m.role == 'user' -%}<|user|>{{ visible_text(m.content) }}{% set ns.has_thinking = (',' ~ loop.index0 ~ ',') in ns.thinking_indices -%}
+{%- elif m.role == 'assistant' -%}
+<|assistant|>
+{%- set content = visible_text(m.content) %}
+{%- if m.reasoning_content is string %}
+    {%- set reasoning_content = m.reasoning_content %}
+{%- elif '</think>' in content %}
+    {%- set reasoning_content = content.split('</think>')[0].split('<think>')[-1] %}
+    {%- set content = content.split('</think>')[-1] %}
+{%- elif loop.index0 > ns.last_user_index and not (enable_thinking is defined and not enable_thinking) %}
+    {%- set reasoning_content = '' %}
+{%- elif loop.index0 < ns.last_user_index and ns.has_thinking %}
+    {%- set reasoning_content = '' %}
+{%- endif %}
+{%- if ((clear_thinking is defined and not clear_thinking) or loop.index0 > ns.last_user_index) and reasoning_content is defined -%}
+{{ '<think>' + reasoning_content +  '</think>'}}
+{%- else -%}
+{{ '</think>' }}
+{%- endif -%}
+{%- if content.strip() -%}
+{{ content.strip() }}
+{%- endif -%}
+{% if m.tool_calls %}
+{% for tc in m.tool_calls %}
+{%- if tc.function %}
+    {%- set tc = tc.function %}
+{%- endif %}
+{{- '<tool_call>' + tc.name -}}
+{% set _args = tc.arguments %}{% for k, v in _args.items() %}<arg_key>{{ k }}</arg_key><arg_value>{{ v | tojson(ensure_ascii=False) if v is not string else v }}</arg_value>{% endfor %}</tool_call>{% endfor %}
+{% endif %}
+{%- elif m.role == 'tool' -%}
+{%- if loop.first or (messages[loop.index0 - 1].role != "tool") %}
+    {{- '<|observation|>' -}}
+{%- endif %}
+{%- if m.content is string -%}
+    {{- '<tool_response>' + m.content + '</tool_response>' -}}
+{%- elif m.content is iterable and m.content is not mapping and m.content and m.content.0.type == "tool_reference" -%}
+    {{- '<tool_response><tools>\n' -}}
+    {% for tr in m.content %}
+        {%- for tool in tools -%}
+            {%- if 'function' in tool -%}
+                {%- set tool = tool['function'] -%}
+            {%- endif -%}
+            {%- if tool.name == tr.name -%}
+                {{- tool_to_json(tool) + '\n' -}}
+            {%- endif -%}
+        {%- endfor -%}
+    {%- endfor -%}
+    {{- '</tools></tool_response>' -}}
+{%- else -%}
+    {{- '<tool_response>' + visible_text(m.content) + '</tool_response>' -}}
+{% endif -%}
+{%- elif m.role == 'system' -%}
+<|system|>{{ visible_text(m.content) }}
+{%- endif -%}
+{%- endfor -%}
+{%- if add_generation_prompt -%}
+    <|assistant|>{{- '</think>' if (enable_thinking is defined and not enable_thinking) else '<think>' -}}
+{%- endif -%}

hf.py

@@ -1,771 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-import inspect
-import itertools
-from collections import defaultdict, deque
-from collections.abc import Set
-from functools import lru_cache
-from typing import Any, Literal, cast, overload
-
-import jinja2
-import jinja2.ext
-import jinja2.meta
-import jinja2.nodes
-import jinja2.parser
-import jinja2.sandbox
-
-from vllm.config import ModelConfig, VllmConfig
-from vllm.entrypoints.chat_utils import (
-    ChatCompletionMessageParam,
-    ChatTemplateContentFormat,
-    ChatTemplateContentFormatOption,
-    ChatTemplateResolutionError,
-    ConversationMessage,
-    load_chat_template,
-    parse_chat_messages,
-    parse_chat_messages_async,
-)
-from vllm.inputs import MultiModalDataDict, MultiModalUUIDDict
-from vllm.logger import init_logger
-from vllm.tokenizers.hf import HfTokenizer
-from vllm.transformers_utils.chat_templates import get_chat_template_fallback_path
-from vllm.transformers_utils.processor import cached_get_processor
-from vllm.utils.async_utils import make_async
-from vllm.utils.func_utils import supports_kw
-
-from .base import BaseRenderer
-from .inputs import DictPrompt
-from .inputs.preprocess import parse_dec_only_prompt
-from .params import ChatParams
-
-logger = init_logger(__name__)
-
-
-_PROCESSOR_CHAT_TEMPLATES = dict[tuple[str, bool], str | None]()
-"""
-Used in `_try_get_processor_chat_template` to avoid calling
-`cached_get_processor` again if the processor fails to be loaded.
-
-This is needed because `lru_cache` does not cache when an exception happens.
-"""
-
-
-def _try_get_processor_chat_template(
-    tokenizer: HfTokenizer,
-    *,
-    trust_remote_code: bool,
-) -> str | None:
-    cache_key = (tokenizer.name_or_path, trust_remote_code)
-    if cache_key in _PROCESSOR_CHAT_TEMPLATES:
-        return _PROCESSOR_CHAT_TEMPLATES[cache_key]
-
-    from transformers import (
-        PreTrainedTokenizer,
-        PreTrainedTokenizerFast,
-        ProcessorMixin,
-    )
-
-    try:
-        processor = cached_get_processor(
-            tokenizer.name_or_path,
-            processor_cls=(
-                PreTrainedTokenizer,
-                PreTrainedTokenizerFast,
-                ProcessorMixin,
-            ),
-            trust_remote_code=trust_remote_code,
-        )
-        if (
-            isinstance(processor, ProcessorMixin)
-            and hasattr(processor, "chat_template")
-            and (chat_template := processor.chat_template) is not None
-        ):
-            _PROCESSOR_CHAT_TEMPLATES[cache_key] = chat_template
-            return chat_template
-    except Exception:
-        logger.debug(
-            "Failed to load AutoProcessor chat template for %s",
-            tokenizer.name_or_path,
-            exc_info=True,
-        )
-
-    _PROCESSOR_CHAT_TEMPLATES[cache_key] = None
-    return None
-
-
-def resolve_chat_template(
-    tokenizer: HfTokenizer,
-    chat_template: str | None,
-    tools: list[dict[str, Any]] | None,
-    *,
-    model_config: "ModelConfig",
-) -> str | None:
-    # 1st priority: The given chat template
-    if chat_template is not None:
-        # Resolve template names (e.g. "tool_use") to actual Jinja content
-        # so that downstream kwargs detection can parse template variables.
-        return tokenizer.get_chat_template(chat_template, tools=tools)
-
-    # 2nd priority: AutoProcessor chat template, unless tool calling is enabled
-    if tools is None:
-        chat_template = _try_get_processor_chat_template(
-            tokenizer,
-            trust_remote_code=model_config.trust_remote_code,
-        )
-        if chat_template is not None:
-            return chat_template
-
-    # 3rd priority: AutoTokenizer chat template
-    try:
-        return tokenizer.get_chat_template(chat_template, tools=tools)
-    except Exception:
-        logger.debug(
-            "Failed to load AutoTokenizer chat template for %s",
-            tokenizer.name_or_path,
-            exc_info=True,
-        )
-
-    # 4th priority: Predefined fallbacks
-    path = get_chat_template_fallback_path(
-        model_type=model_config.hf_config.model_type,
-        tokenizer_name_or_path=tokenizer.name_or_path,
-    )
-    if path is not None:
-        logger.info_once(
-            "Loading chat template fallback for %s as there isn't one "
-            "defined on HF Hub.",
-            tokenizer.name_or_path,
-        )
-        chat_template = load_chat_template(path)
-    else:
-        logger.debug_once(
-            "There is no chat template fallback for %s", tokenizer.name_or_path
-        )
-
-    return chat_template
-
-
-def _is_var_access(node: jinja2.nodes.Node, varname: str) -> bool:
-    if isinstance(node, jinja2.nodes.Name):
-        return node.ctx == "load" and node.name == varname
-
-    return False
-
-
-def _is_attr_access(node: jinja2.nodes.Node, varname: str, key: str) -> bool:
-    if isinstance(node, jinja2.nodes.Getitem):
-        return (
-            _is_var_access(node.node, varname)
-            and isinstance(node.arg, jinja2.nodes.Const)
-            and node.arg.value == key
-        )
-
-    if isinstance(node, jinja2.nodes.Getattr):
-        return _is_var_access(node.node, varname) and node.attr == key
-
-    return False
-
-
-def _is_var_or_elems_access(
-    node: jinja2.nodes.Node,
-    varname: str,
-    key: str | None = None,
-) -> bool:
-    if isinstance(node, jinja2.nodes.Filter):
-        return node.node is not None and _is_var_or_elems_access(
-            node.node, varname, key
-        )
-    if isinstance(node, jinja2.nodes.Test):
-        return _is_var_or_elems_access(node.node, varname, key)
-
-    if isinstance(node, jinja2.nodes.Getitem) and isinstance(
-        node.arg, jinja2.nodes.Slice
-    ):
-        return _is_var_or_elems_access(node.node, varname, key)
-
-    return _is_attr_access(node, varname, key) if key else _is_var_access(node, varname)
-
-
-def _iter_nodes_assign_var_or_elems(root: jinja2.nodes.Node, varname: str):
-    # Global variable that is implicitly defined at the root
-    yield root, varname
-
-    # Iterative BFS
-    related_varnames = deque([varname])
-    while related_varnames:
-        related_varname = related_varnames.popleft()
-
-        for assign_ast in root.find_all(jinja2.nodes.Assign):
-            lhs = assign_ast.target
-            rhs = assign_ast.node
-
-            if _is_var_or_elems_access(rhs, related_varname):
-                assert isinstance(lhs, jinja2.nodes.Name)
-                yield assign_ast, lhs.name
-
-                # Avoid infinite looping for self-assignment
-                if lhs.name != related_varname:
-                    related_varnames.append(lhs.name)
-
-
-# NOTE: The proper way to handle this is to build a CFG so that we can handle
-# the scope in which each variable is defined, but that is too complicated
-def _iter_nodes_assign_messages_item(root: jinja2.nodes.Node):
-    messages_varnames = [
-        varname for _, varname in _iter_nodes_assign_var_or_elems(root, "messages")
-    ]
-
-    # Search for {%- for message in messages -%} loops
-    for loop_ast in root.find_all(jinja2.nodes.For):
-        loop_iter = loop_ast.iter
-        loop_target = loop_ast.target
-
-        for varname in messages_varnames:
-            if _is_var_or_elems_access(loop_iter, varname):
-                assert isinstance(loop_target, jinja2.nodes.Name)
-                yield loop_ast, loop_target.name
-                break
-
-
-def _iter_nodes_assign_content_item(root: jinja2.nodes.Node):
-    message_varnames = [
-        varname for _, varname in _iter_nodes_assign_messages_item(root)
-    ]
-
-    # Search for {%- for content in message['content'] -%} loops
-    for loop_ast in root.find_all(jinja2.nodes.For):
-        loop_iter = loop_ast.iter
-        loop_target = loop_ast.target
-
-        for varname in message_varnames:
-            if _is_var_or_elems_access(loop_iter, varname, "content"):
-                assert isinstance(loop_target, jinja2.nodes.Name)
-                yield loop_ast, loop_target.name
-                break
-
-
-def _try_extract_ast(chat_template: str) -> jinja2.nodes.Template | None:
-    import transformers.utils.chat_template_utils as hf_chat_utils
-
-    try:
-        jinja_compiled = hf_chat_utils._compile_jinja_template(chat_template)
-        return jinja_compiled.environment.parse(chat_template)
-    except Exception:
-        logger.exception("Error when compiling Jinja template")
-        return None
-
-
-@lru_cache(maxsize=32)
-def _detect_content_format(
-    chat_template: str,
-    *,
-    default: ChatTemplateContentFormat,
-) -> ChatTemplateContentFormat:
-    jinja_ast = _try_extract_ast(chat_template)
-    if jinja_ast is None:
-        return default
-
-    try:
-        next(_iter_nodes_assign_content_item(jinja_ast))
-    except StopIteration:
-        return "string"
-    except Exception:
-        logger.exception("Error when parsing AST of Jinja template")
-        return default
-    else:
-        return "openai"
-
-
-def _is_glm_model(tokenizer: HfTokenizer, model_config: "ModelConfig") -> bool:
-    """Check if this is a GLM model that requires string content format.
-    
-    GLM models (GLM-4, GLM-4.5, GLM-5.x) have a chat template that incorrectly
-    triggers "openai" content format detection because they iterate over
-    m.content for tool responses. However, the template expects string content
-    for tool messages (checking `m.content is string`).
-    
-    This detection ensures we force "string" format for GLM models.
-    """
-    # Check tokenizer name/path for GLM indicators
-    name_or_path = tokenizer.name_or_path.lower()
-    glm_indicators = ["glm-4", "glm-5", "glm4", "glm5", "zai-org/glm"]
-    if any(ind in name_or_path for ind in glm_indicators):
-        return True
-    
-    # Check model type in config
-    if hasattr(model_config, "hf_config") and hasattr(model_config.hf_config, "model_type"):
-        model_type = model_config.hf_config.model_type.lower()
-        if "glm" in model_type:
-            return True
-    
-    return False
-
-
-def _resolve_chat_template_content_format(
-    chat_template: str | None,
-    tools: list[dict[str, Any]] | None,
-    tokenizer: HfTokenizer,
-    *,
-    model_config: "ModelConfig",
-) -> ChatTemplateContentFormat:
-    # GLM models require "string" content format for tool responses to work
-    # The template has `{% for tr in m.content %}` which triggers "openai"
-    # detection, but then checks `m.content is string` which fails for arrays.
-    if _is_glm_model(tokenizer, model_config):
-        logger.debug(
-            "Forcing 'string' content format for GLM model: %s",
-            tokenizer.name_or_path,
-        )
-        return "string"
-    
-    resolved_chat_template = resolve_chat_template(
-        tokenizer,
-        chat_template=chat_template,
-        tools=tools,
-        model_config=model_config,
-    )
-
-    jinja_text = (
-        resolved_chat_template
-        if isinstance(resolved_chat_template, str)
-        else load_chat_template(chat_template, is_literal=True)
-    )
-
-    detected_format = (
-        "string"
-        if jinja_text is None
-        else _detect_content_format(jinja_text, default="string")
-    )
-
-    return detected_format
-
-
-@lru_cache
-def _log_chat_template_content_format(
-    chat_template: str | None,  # For caching purposes
-    given_format: ChatTemplateContentFormatOption,
-    detected_format: ChatTemplateContentFormatOption,
-):
-    logger.info(
-        "Detected the chat template content format to be '%s'. "
-        "You can set `--chat-template-content-format` to override this.",
-        detected_format,
-    )
-
-    if given_format != "auto" and given_format != detected_format:
-        logger.warning(
-            "You specified `--chat-template-content-format %s` "
-            "which is different from the detected format '%s'. "
-            "If our automatic detection is incorrect, please consider "
-            "opening a GitHub issue so that we can improve it: "
-            "https://github.com/vllm-project/vllm/issues/new/choose",
-            given_format,
-            detected_format,
-        )
-
-
-def resolve_chat_template_content_format(
-    chat_template: str | None,
-    tools: list[dict[str, Any]] | None,
-    given_format: ChatTemplateContentFormatOption,
-    tokenizer: HfTokenizer,
-    *,
-    model_config: "ModelConfig",
-) -> ChatTemplateContentFormat:
-    if given_format != "auto":
-        return given_format
-
-    detected_format = _resolve_chat_template_content_format(
-        chat_template,
-        tools,
-        tokenizer,
-        model_config=model_config,
-    )
-
-    _log_chat_template_content_format(
-        chat_template,
-        given_format=given_format,
-        detected_format=detected_format,
-    )
-
-    return detected_format
-
-
-# adapted from https://github.com/huggingface/transformers/blob/v4.56.2/src/transformers/utils/chat_template_utils.py#L398-L412
-# only preserve the parse function used to resolve chat template kwargs
-class AssistantTracker(jinja2.ext.Extension):
-    tags = {"generation"}
-
-    def parse(self, parser: jinja2.parser.Parser) -> jinja2.nodes.Node:
-        lineno = next(parser.stream).lineno
-        body = parser.parse_statements(("name:endgeneration",), drop_needle=True)
-        call = self.call_method("_generation_support")
-        call_block = jinja2.nodes.CallBlock(call, [], [], body)
-        return call_block.set_lineno(lineno)
-
-
-def _resolve_chat_template_kwargs(chat_template: str) -> Set[str]:
-    env = jinja2.sandbox.ImmutableSandboxedEnvironment(
-        trim_blocks=True,
-        lstrip_blocks=True,
-        extensions=[AssistantTracker, jinja2.ext.loopcontrols],
-    )
-    parsed_content = env.parse(chat_template)
-    template_vars = jinja2.meta.find_undeclared_variables(parsed_content)
-    return template_vars
-
-
-_cached_resolve_chat_template_kwargs = lru_cache(_resolve_chat_template_kwargs)
-
-
-@lru_cache
-def _get_hf_base_chat_template_params() -> frozenset[str]:
-    from transformers import PreTrainedTokenizer
-
-    # Get standard parameters from HuggingFace's base tokenizer class.
-    # This dynamically extracts parameters from PreTrainedTokenizer's
-    # apply_chat_template method, ensuring compatibility with tokenizers
-    # that use **kwargs to receive standard parameters.
-
-    # Read signature from HF's base class - the single source of truth
-    base_sig = inspect.signature(PreTrainedTokenizer.apply_chat_template)
-
-    # Exclude VAR_KEYWORD (**kwargs) and VAR_POSITIONAL (*args) placeholders
-    return frozenset(
-        p.name
-        for p in base_sig.parameters.values()
-        if p.kind
-        not in (inspect.Parameter.VAR_KEYWORD, inspect.Parameter.VAR_POSITIONAL)
-    )
-
-
-def resolve_chat_template_kwargs(
-    tokenizer: HfTokenizer,
-    chat_template: str,
-    chat_template_kwargs: dict[str, Any],
-    raise_on_unexpected: bool = True,
-) -> dict[str, Any]:
-    # We exclude chat_template from kwargs here, because
-    # chat template has been already resolved at this stage
-    unexpected_vars = {"chat_template", "tokenize"}
-    if raise_on_unexpected and (
-        unexpected_in_kwargs := unexpected_vars & chat_template_kwargs.keys()
-    ):
-        raise ValueError(
-            "Found unexpected chat template kwargs from request: "
-            f"{unexpected_in_kwargs}"
-        )
-
-    fn_kw = {
-        k
-        for k in chat_template_kwargs
-        if supports_kw(tokenizer.apply_chat_template, k, allow_var_kwargs=False)
-    }
-    template_vars = _cached_resolve_chat_template_kwargs(chat_template)
-
-    # Allow standard HF parameters even if tokenizer uses **kwargs to receive them
-    hf_base_params = _get_hf_base_chat_template_params()
-
-    accept_vars = (fn_kw | template_vars | hf_base_params) - unexpected_vars
-    return {k: v for k, v in chat_template_kwargs.items() if k in accept_vars}
-
-
-@overload
-def safe_apply_chat_template(
-    model_config: "ModelConfig",
-    tokenizer: HfTokenizer,
-    conversation: list[ConversationMessage],
-    *,
-    tools: list[dict[str, Any]] | None = ...,
-    chat_template: str | None = ...,
-    tokenize: Literal[True] = ...,
-    **kwargs,
-) -> list[int]: ...
-@overload
-def safe_apply_chat_template(
-    model_config: "ModelConfig",
-    tokenizer: HfTokenizer,
-    conversation: list[ConversationMessage],
-    *,
-    tools: list[dict[str, Any]] | None = ...,
-    chat_template: str | None = ...,
-    tokenize: Literal[False] = ...,
-    **kwargs,
-) -> str: ...
-def safe_apply_chat_template(
-    model_config: "ModelConfig",
-    tokenizer: HfTokenizer,
-    conversation: list[ConversationMessage],
-    *,
-    tools: list[dict[str, Any]] | None = None,
-    chat_template: str | None = None,
-    tokenize: bool = True,
-    **kwargs,
-) -> str | list[int]:
-    chat_template = resolve_chat_template(
-        tokenizer,
-        chat_template=chat_template,
-        tools=tools,
-        model_config=model_config,
-    )
-    if chat_template is None:
-        raise ChatTemplateResolutionError(
-            "As of transformers v4.44, default chat template is no longer "
-            "allowed, so you must provide a chat template if the tokenizer "
-            "does not define one."
-        )
-
-    resolved_kwargs = resolve_chat_template_kwargs(
-        tokenizer=tokenizer,
-        chat_template=chat_template,
-        chat_template_kwargs=kwargs,
-    )
-
-    try:
-        return tokenizer.apply_chat_template(
-            conversation=conversation,  # type: ignore[arg-type]
-            tools=tools,  # type: ignore[arg-type]
-            chat_template=chat_template,
-            tokenize=tokenize,
-            **resolved_kwargs,
-        )
-    # External library exceptions can sometimes occur despite the framework's
-    # internal exception management capabilities.
-    except Exception as e:
-        # Log and report any library-related exceptions for further
-        # investigation.
-        logger.exception(
-            "An error occurred in `transformers` while applying chat template"
-        )
-        raise ValueError(str(e)) from e
-
-
-def rebuild_mm_uuids_from_mm_data(
-    mm_uuids: MultiModalUUIDDict,
-    mm_data: MultiModalDataDict,
-) -> MultiModalUUIDDict:
-    """Rebuild mm_uuids after vision_chunk processing.
-
-    When videos are split into chunks, the original UUIDs need to be updated
-    to reflect the new UUIDs generated for each chunk.
-
-    Args:
-        mm_uuids: Original UUIDs dictionary
-        mm_data: Processed multimodal data with vision_chunk items
-
-    Returns:
-        Updated UUIDs dictionary with chunk UUIDs
-    """
-    vision_chunks = mm_data.get("vision_chunk")
-    if vision_chunks is None:
-        return mm_uuids
-
-    assert all(isinstance(item, dict) for item in vision_chunks), (
-        "Expected all vision_chunk items to be dicts"
-    )
-    vision_chunks = cast(list[dict[str, Any]], vision_chunks)
-    vision_chunk_uuids = [
-        uuid_val for item in vision_chunks if (uuid_val := item.get("uuid")) is not None
-    ]
-
-    if vision_chunk_uuids:
-        mm_uuids = dict(mm_uuids)
-        mm_uuids["vision_chunk"] = vision_chunk_uuids
-
-    return mm_uuids
-
-
-def build_video_prompts_from_mm_data(
-    mm_data: MultiModalDataDict,
-) -> list[str]:
-    """Build video prompts from vision_chunk data.
-
-    Collects prompts from video chunks and groups them by video_idx.
-
-    Args:
-        mm_data: Processed multimodal data with vision_chunk items
-
-    Returns:
-        List of video prompts, one per video.
-    """
-    vision_chunks = mm_data.get("vision_chunk")
-    if vision_chunks is None:
-        return []
-
-    # Group chunks by video_idx
-    video_prompts_dict: dict[int, list[str]] = defaultdict(list)
-
-    for item in vision_chunks:
-        # vision_chunk items are always dicts (VisionChunkImage/VisionChunkVideo)
-        assert isinstance(item, dict)
-        if item.get("type") == "video_chunk":
-            video_idx = item.get("video_idx", 0)
-            prompt = item.get("prompt", "")
-            video_prompts_dict[video_idx].append(prompt)
-
-    # Build prompts in video order
-    video_prompts = [
-        "".join(video_prompts_dict[video_idx])
-        for video_idx in sorted(video_prompts_dict.keys())
-    ]
-
-    return video_prompts
-
-
-def replace_vision_chunk_video_placeholder(
-    prompt_raw: str | list[int],
-    mm_data: MultiModalDataDict,
-    video_placeholder: str | None,
-) -> str | list[int]:
-    # get video placeholder, replace it with runtime video-chunk prompts
-    if video_placeholder and isinstance(prompt_raw, str):
-        video_prompts = build_video_prompts_from_mm_data(mm_data)
-
-        # replace in order
-        prompt_raw_parts = prompt_raw.split(video_placeholder)
-        if len(prompt_raw_parts) == len(video_prompts) + 1:
-            prompt_raw = "".join(
-                itertools.chain.from_iterable(zip(prompt_raw_parts, video_prompts))
-            )
-            prompt_raw += prompt_raw_parts[-1]
-        else:
-            logger.warning(
-                "Number of video placeholders (%d) does not match "
-                "number of videos (%d) in the request.",
-                len(prompt_raw_parts) - 1,
-                len(video_prompts),
-            )
-    return prompt_raw
-
-
-class HfRenderer(BaseRenderer[HfTokenizer]):
-    def __init__(
-        self,
-        config: VllmConfig,
-        tokenizer: HfTokenizer | None,
-    ) -> None:
-        super().__init__(config, tokenizer)
-
-        self.use_unified_vision_chunk = getattr(
-            config.model_config.hf_config, "use_unified_vision_chunk", False
-        )
-
-        self._apply_chat_template_async = make_async(
-            safe_apply_chat_template, executor=self._executor
-        )
-
-    def render_messages(
-        self,
-        messages: list[ChatCompletionMessageParam],
-        params: ChatParams,
-    ) -> tuple[list[ConversationMessage], DictPrompt]:
-        model_config = self.model_config
-        tokenizer = self.get_tokenizer()
-
-        conversation, mm_data, mm_uuids = parse_chat_messages(
-            messages,
-            model_config,
-            content_format=resolve_chat_template_content_format(
-                chat_template=params.chat_template,
-                tools=params.chat_template_kwargs.get("tools"),
-                given_format=params.chat_template_content_format,
-                tokenizer=tokenizer,
-                model_config=model_config,
-            ),
-            media_io_kwargs=params.media_io_kwargs,
-            mm_processor_kwargs=params.mm_processor_kwargs,
-        )
-
-        prompt_raw = safe_apply_chat_template(
-            model_config,
-            tokenizer,
-            conversation,
-            **params.get_apply_chat_template_kwargs(),
-        )
-
-        # NOTE: use_unified_vision_chunk is currently specific to Kimi-K2.5
-        # model which uses unified vision chunks for both images and videos.
-        if (
-            self.use_unified_vision_chunk
-            and mm_uuids is not None
-            and mm_data is not None
-        ):
-            mm_uuids = rebuild_mm_uuids_from_mm_data(mm_uuids, mm_data)
-
-            # get video placeholder, replace it with runtime video-chunk prompts
-            video_placeholder = getattr(
-                model_config.hf_config, "video_placeholder", None
-            )
-            prompt_raw = cast(
-                list[int],
-                replace_vision_chunk_video_placeholder(
-                    prompt_raw,
-                    mm_data,
-                    video_placeholder,
-                ),
-            )
-
-        prompt = parse_dec_only_prompt(prompt_raw)
-        if mm_data is not None:
-            prompt["multi_modal_data"] = mm_data
-        if mm_uuids is not None:
-            prompt["multi_modal_uuids"] = mm_uuids
-
-        return conversation, prompt
-
-    async def render_messages_async(
-        self,
-        messages: list[ChatCompletionMessageParam],
-        params: ChatParams,
-    ) -> tuple[list[ConversationMessage], DictPrompt]:
-        model_config = self.model_config
-        tokenizer = self.get_tokenizer()
-
-        conversation, mm_data, mm_uuids = await parse_chat_messages_async(
-            messages,
-            model_config,
-            content_format=resolve_chat_template_content_format(
-                chat_template=params.chat_template,
-                tools=params.chat_template_kwargs.get("tools"),
-                given_format=params.chat_template_content_format,
-                tokenizer=tokenizer,
-                model_config=model_config,
-            ),
-            media_io_kwargs=params.media_io_kwargs,
-            mm_processor_kwargs=params.mm_processor_kwargs,
-        )
-
-        prompt_raw = await self._apply_chat_template_async(
-            model_config,
-            tokenizer,
-            conversation,
-            **params.get_apply_chat_template_kwargs(),
-        )
-
-        # NOTE: use_unified_vision_chunk is currently specific to Kimi-K2.5
-        # model which uses unified vision chunks for both images and videos.
-        if (
-            self.use_unified_vision_chunk
-            and mm_uuids is not None
-            and mm_data is not None
-        ):
-            # get video placeholder, replace it with runtime video-chunk prompts
-            video_placeholder = getattr(
-                model_config.hf_config, "video_placeholder", None
-            )
-            prompt_raw = cast(
-                list[int],
-                replace_vision_chunk_video_placeholder(
-                    prompt_raw,
-                    mm_data,
-                    video_placeholder,
-                ),
-            )
-
-        prompt = parse_dec_only_prompt(prompt_raw)
-        if mm_data is not None:
-            prompt["multi_modal_data"] = mm_data
-        if mm_uuids is not None:
-            prompt["multi_modal_uuids"] = mm_uuids
-
-        return conversation, prompt

indexer.py

@@ -0,0 +1,779 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from dataclasses import dataclass
+
+import torch
+
+import vllm.envs as envs
+from vllm.config import VllmConfig
+from vllm.logger import init_logger
+from vllm.platforms import current_platform
+from vllm.triton_utils import tl, triton
+from vllm.utils.deep_gemm import (
+    get_paged_mqa_logits_metadata,
+    has_deep_gemm,
+)
+from vllm.utils.math_utils import cdiv
+from vllm.utils.platform_utils import num_compute_units
+from vllm.v1.attention.backend import (
+    AttentionBackend,
+    AttentionCGSupport,
+    AttentionMetadataBuilder,
+    CommonAttentionMetadata,
+    MultipleOf,
+)
+from vllm.v1.attention.backends.mla.compressor_utils import get_compressed_slot_mapping
+from vllm.v1.attention.backends.utils import (
+    split_decodes_and_prefills,
+)
+from vllm.v1.kv_cache_interface import AttentionSpec, MLAAttentionSpec
+from vllm.v1.worker.cp_utils import get_total_cp_world_size
+
+logger = init_logger(__name__)
+
+
+@triton.jit
+def _prepare_uniform_decode_kernel(
+    seq_lens_ptr,
+    decode_seq_lens_ptr,
+    block_table_ptr,
+    block_table_stride,
+    expanded_block_table_ptr,
+    expanded_bt_stride,
+    decode_lens_ptr,
+    max_decode_len,
+    BLOCK_SIZE: tl.constexpr,
+):
+    idx = tl.program_id(0)
+    req_id = idx // max_decode_len
+    local_idx = idx % max_decode_len
+
+    # Compute number of KVs attended to by this token.
+    seq_len = tl.load(seq_lens_ptr + req_id)
+    per_token_seq_len = seq_len - max_decode_len + local_idx + 1
+    tl.store(decode_seq_lens_ptr + idx, per_token_seq_len)
+
+    # Copy block table row.
+    src = block_table_ptr + req_id * block_table_stride
+    dst = expanded_block_table_ptr + idx * expanded_bt_stride
+    for i in tl.range(0, expanded_bt_stride, BLOCK_SIZE):
+        off = i + tl.arange(0, BLOCK_SIZE)
+        mask = off < expanded_bt_stride
+        src_block = tl.load(src + off, mask=mask)
+        tl.store(dst + off, src_block, mask=mask)
+
+    # All reqs now have decode_len = 1.
+    tl.store(decode_lens_ptr + idx, 1)
+
+
+def split_indexer_prefill_chunks(
+    seq_lens_cpu: torch.Tensor,
+    query_lens_cpu: torch.Tensor,
+    workspace_size: int,
+    max_logits_bytes: int,
+    request_offset: int = 0,
+) -> list[tuple[slice, slice]]:
+    """
+    Split prefill requests into chunks for the sparse indexer, respecting:
+    - N constraint: total_seq_lens <= workspace_size (existing O(N) workspace)
+    - Logits constraint: M * N * 4 <= max_logits_bytes
+
+    When a single request-level chunk still exceeds the logits budget,
+    sub-chunks on the query dimension (M) to bound peak memory.
+
+    Returns list of (req_slice, query_slice) tuples.
+    """
+    chunks: list[tuple[slice, slice]] = []
+    n = len(seq_lens_cpu)
+    max_logits_elems = max_logits_bytes // 4
+    end = 0
+
+    while end < n:
+        start, chunk_m, chunk_n = end, 0, 0
+
+        while end < n:
+            q, s = query_lens_cpu[end].item(), seq_lens_cpu[end].item()
+            new_m, new_n = chunk_m + q, chunk_n + s
+            if new_n <= workspace_size and new_m * new_n <= max_logits_elems:
+                chunk_m, chunk_n = new_m, new_n
+                end += 1
+            else:
+                break
+
+        # A single request can exceed the budget, requiring sub-chunking
+        # on the query dimension.
+        if end == start:
+            chunk_m, chunk_n = query_lens_cpu[end].item(), seq_lens_cpu[end].item()
+            end += 1
+
+        req_slice = slice(start + request_offset, end + request_offset)
+        max_q = max(1, max_logits_elems // chunk_n) if chunk_n > 0 else chunk_m
+        for q_off in range(0, chunk_m, max_q):
+            sub_m = min(max_q, chunk_m - q_off)
+            chunks.append((req_slice, slice(q_off, q_off + sub_m)))
+
+    return chunks
+
+
+class DeepseekV32IndexerBackend(AttentionBackend):
+    @staticmethod
+    def get_name() -> str:
+        return "DEEPSEEK_V32_INDEXER"
+
+    @staticmethod
+    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+        return [1, 64] if current_platform.is_rocm() else [64]
+
+    @classmethod
+    def get_supported_head_sizes(cls) -> list[int]:
+        return [32, 64, 128]
+
+    @staticmethod
+    def get_builder_cls() -> type["DeepseekV32IndexerMetadataBuilder"]:
+        return DeepseekV32IndexerMetadataBuilder
+
+    @staticmethod
+    def get_kv_cache_shape(
+        num_blocks: int,
+        block_size: int,
+        num_kv_heads: int,
+        head_size: int,
+        cache_dtype_str: str = "auto",
+    ) -> tuple[int, ...]:
+        assert num_kv_heads == 1
+        return (num_blocks, block_size, head_size)
+
+    @staticmethod
+    def get_kv_cache_stride_order(
+        include_num_layers_dimension: bool = False,
+    ) -> tuple[int, ...]:
+        if include_num_layers_dimension:
+            # DeepseekV32Indexer kernels do not support cross-layer
+            # KV cache layout. Identity permutation keeps num_layers
+            # first, signaling incompatibility.
+            return (0, 1, 2, 3)
+        return (0, 1, 2)
+
+
+class DeepseekV4IndexerBackend(DeepseekV32IndexerBackend):
+    @staticmethod
+    def get_name() -> str:
+        return "DEEPSEEK_V4_INDEXER"
+
+    @staticmethod
+    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+        return [256]
+
+
+@dataclass
+class DeepseekV32IndexerPrefillChunkMetadata:
+    block_table: torch.Tensor
+    cu_seqlen_ks: torch.Tensor
+    cu_seqlen_ke: torch.Tensor
+    cu_seq_lens: torch.Tensor
+    token_to_seq: torch.Tensor
+    total_seq_lens: int
+    token_start: int
+    token_end: int
+    num_reqs: int
+    skip_kv_gather: bool = False
+
+
+@dataclass
+class DeepseekV32IndexerPrefillMetadata:
+    chunks: list[DeepseekV32IndexerPrefillChunkMetadata]
+
+
+@dataclass
+class DeepSeekV32IndexerDecodeMetadata:
+    block_table: torch.Tensor
+    # seq_lens: per-token effective context lengths.
+    #   - flatten path / plain decode: 1D (batch_size,)
+    #   - native MTP path: 2D (B, next_n) where [b,j] = L_b - next_n + j + 1
+    # Both fp8_fp4_paged_mqa_logits and the topk kernels accept both shapes.
+    seq_lens: torch.Tensor
+    decode_lens: torch.Tensor
+    requires_padding: bool
+    schedule_metadata: torch.Tensor
+
+
+@dataclass
+class DeepseekV32IndexerMetadata:
+    # FIXME (zyongye)
+    # hacky way to access the data now, need to be in chunked meta
+    seq_lens: torch.Tensor
+    max_seq_len: int
+    slot_mapping: torch.Tensor
+
+    # New for MLA (compared to FlashAttention)
+    # For handling prefill decode split
+    num_decodes: int
+    num_decode_tokens: int
+    num_prefills: int
+    num_prefill_tokens: int
+
+    decode: DeepSeekV32IndexerDecodeMetadata | None = None
+    prefill: DeepseekV32IndexerPrefillMetadata | None = None
+
+
+def get_max_prefill_buffer_size(vllm_config: VllmConfig):
+    max_model_len = vllm_config.model_config.max_model_len
+    # NOTE(Chen): 40 is a magic number for controlling the prefill buffer size.
+    # Each entry is 128 fp8 bytes and 4 scale bytes for a total of 132 bytes.
+    # The flashmla_sparse backend uses a workspace size of 5 * max_model_len.
+    # The memory usage of the workspace there is 576 * 2 bytes; so we size this as
+    # (576 * 2 // 132) * 5 = 40 to maximize this workspace size while still fitting
+    # within the flashmla_sparse workspace.
+    # For DeepSeek-V3.2, the max_model_len is 163840.
+    #   40 * 163840 * 132 = 865075200 bytes = 825 MB
+    return max_model_len * 40
+
+
+class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
+    reorder_batch_threshold: int = 1
+
+    @classmethod
+    def get_cudagraph_support(
+        cls,
+        vllm_config: VllmConfig,
+        kv_cache_spec: AttentionSpec,
+    ) -> AttentionCGSupport:
+        return AttentionCGSupport.UNIFORM_BATCH
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        scheduler_config = self.vllm_config.scheduler_config
+        # NOTE(Chen):an estimated max size of flattened_kv. Need to double check.
+        self.max_prefill_buffer_size = get_max_prefill_buffer_size(self.vllm_config)
+        self.num_speculative_tokens = (
+            self.vllm_config.speculative_config.num_speculative_tokens
+            if self.vllm_config.speculative_config
+            else 0
+        )
+        self.use_fp4_indexer_cache = (
+            self.vllm_config.attention_config.use_fp4_indexer_cache
+        )
+
+        assert (
+            current_platform.is_device_capability_family(100)
+            or not self.use_fp4_indexer_cache
+        ), (
+            "use_fp4_indexer_cache requires Blackwell datacenter GPUs "
+            "(sm_10x, e.g. B200/GB200); sm_120 (consumer Blackwell) and "
+            "earlier architectures are not supported."
+        )
+
+        next_n = self.num_speculative_tokens + 1
+        self.reorder_batch_threshold += self.num_speculative_tokens
+        # NOTE: SM100 datacenter GPUs support any next_n natively via the
+        # multi-atom paged MQA logits kernels (FP8 and FP4 indexer
+        # caches). Outside the SM100 family the FP8
+        # paged MQA logits kernel only supports next_n in (1, 2)
+        # (deepgemm smxx_fp8_fp4_paged_mqa_logits.hpp:233), so flatten there.
+        self.use_flattening = not current_platform.is_device_capability_family(
+            100
+        ) and next_n not in (1, 2)
+        logger.info_once(
+            "DSA indexer decode path: use_flattening=%s "
+            "(next_n=%d, use_fp4_indexer_cache=%s)",
+            self.use_flattening,
+            next_n,
+            self.use_fp4_indexer_cache,
+        )
+
+        sm_count = num_compute_units(self.device.index)
+        self.num_sms = sm_count
+
+        self.offsets_buffer = torch.arange(
+            next_n, device=self.device, dtype=torch.int32
+        )
+        self.decode_lens_buffer = torch.zeros(
+            (scheduler_config.max_num_batched_tokens,),
+            dtype=torch.int32,
+            device=self.device,
+        )
+        # Shared workspace for decode seq_lens. Native MTP views this as
+        # (B, max_decode_len) at runtime, keeping context_lens contiguous even
+        # when max_decode_len is smaller than next_n.
+        self.decode_seq_lens_buffer = torch.zeros(
+            (scheduler_config.max_num_batched_tokens,),
+            dtype=torch.int32,
+            device=self.device,
+        )
+        self.arange_buffer = torch.arange(
+            max(
+                scheduler_config.max_num_seqs * next_n,
+                scheduler_config.max_num_batched_tokens,
+            ),
+            dtype=torch.int32,
+            device=self.device,
+        )
+        max_num_blocks_per_req = cdiv(
+            self.vllm_config.model_config.max_model_len,
+            self.kv_cache_spec.block_size * get_total_cp_world_size(),
+        )
+        self.expanded_block_table_buffer = torch.zeros(
+            (
+                scheduler_config.max_num_batched_tokens,
+                max_num_blocks_per_req,
+            ),
+            dtype=torch.int32,
+            device=self.device,
+        )
+
+        # See: DeepGMM/csrc/apis/attention.hpp
+        self.scheduler_metadata_buffer = torch.empty(
+            (self.num_sms + 1, 2), dtype=torch.int32, device=self.device
+        )
+
+        # KV compression. Default to 1 for no compression.
+        self.compress_ratio = 1
+        # Get compress_ratio for DeepseekV4 support
+        if isinstance(self.kv_cache_spec, MLAAttentionSpec):
+            self.compress_ratio = self.kv_cache_spec.compress_ratio
+
+        # Pre-allocate buffers for CUDA graph compatibility when
+        if self.compress_ratio > 1:
+            # compress_ratio > 1 (DeepseekV4)
+            # Compressed slot mapping output buffer
+            self.compressed_slot_mapping_buffer = torch.zeros(
+                (scheduler_config.max_num_batched_tokens,),
+                dtype=torch.int64,
+                device=self.device,
+            )
+            # Buffer for compressed seq_lens in decode path
+            self.expanded_seq_lens_buffer = torch.zeros(
+                (scheduler_config.max_num_batched_tokens,),
+                dtype=torch.int32,
+                device=self.device,
+            )
+
+    def _prepare_decode_tensors(
+        self,
+        seq_lens: torch.Tensor,
+        block_table: torch.Tensor,
+        decode_lens: torch.Tensor,
+        decode_lens_cpu: torch.Tensor,
+        query_start_loc: torch.Tensor,
+        num_decodes: int,
+        num_decode_tokens: int,
+        use_native: bool,
+        next_n: int,
+        max_decode_len: int,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int, bool]:
+        """Expand seq_lens/block_table/decode_lens for the decode kernels.
+
+        Flatten path (not use_native, max_decode_len > 1):
+          Each multi-token decode request is expanded into individual
+          single-token entries so the kernel always sees next_n=1.
+
+        Native path (use_native or max_decode_len == 1):
+          Plain decode or spec-decode with 2D per-token context lengths.
+
+        Returns (seq_lens, block_table, decode_lens, batch_size, requires_padding).
+        seq_lens is 1D (batch_size,) for flatten/plain, 2D (B, max_decode_len)
+        for native MTP.
+        """
+        min_decode_len = int(decode_lens_cpu.min().item())
+        if not use_native and max_decode_len > 1:
+            assert self.decode_seq_lens_buffer.dim() == 1
+            if min_decode_len == max_decode_len:
+                # Uniform decode lengths.
+                num_decode_tokens = num_decodes * max_decode_len
+                _prepare_uniform_decode_kernel[(num_decode_tokens,)](
+                    seq_lens,
+                    self.decode_seq_lens_buffer,
+                    block_table,
+                    block_table.stride(0),
+                    self.expanded_block_table_buffer,
+                    self.expanded_block_table_buffer.stride(0),
+                    self.decode_lens_buffer,
+                    max_decode_len,
+                    BLOCK_SIZE=1024,
+                )
+                self.decode_seq_lens_buffer[num_decode_tokens:] = 0
+                seq_lens = self.decode_seq_lens_buffer[:num_decode_tokens]
+                block_table = self.expanded_block_table_buffer[:num_decode_tokens]
+                decode_lens = self.decode_lens_buffer[:num_decode_tokens]
+                return seq_lens, block_table, decode_lens, num_decode_tokens, False
+            else:
+                # Variable decode lengths.
+                # Assume 4 requests with seq_lens [10, 7, 12, 0] (the final req is
+                # padding) and decode_lens [3, 1, 4, 0] in the below example comments.
+                # The context lengths are therefore
+                # [10-3, 7-1, 12-4, 0-0] = [7, 6, 8, 0].
+
+                # 3 + 1 + 4 + 0 = 8
+                actual_expanded = int(decode_lens_cpu.sum().item())
+
+                # Fuse expanded_base and expanded_starts into a single
+                # repeat_interleave:
+                # seq_len_i = (context_start[b] - query_start_loc[b]) + arange[i] + 1
+                # where context_start[b] = seq_lens[b] - decode_lens[b].
+                # Example: offsets = [7-0, 6-3, 8-4, 0-8] = [7, 3, 4, -8]
+                # expanded_offsets  = [7, 7, 7, 3, 4, 4, 4, 4]
+                # result            = [8, 9, 10, 7, 9, 10, 11, 12]
+                expanded_offsets = torch.repeat_interleave(
+                    seq_lens - decode_lens - query_start_loc,
+                    decode_lens,
+                    output_size=actual_expanded,
+                )
+
+                # [8, 9, 10, 7, 9, 10, 11, 12, ...] where ... is unused buffer space
+                self.decode_seq_lens_buffer[:actual_expanded] = (
+                    expanded_offsets + self.arange_buffer[:actual_expanded] + 1
+                )
+                self.decode_seq_lens_buffer[actual_expanded:] = 0
+                seq_lens = self.decode_seq_lens_buffer[:num_decode_tokens]
+
+                # Give each of the flattened entries the same block table row as the
+                # original request.
+                self.expanded_block_table_buffer[:actual_expanded] = (
+                    torch.repeat_interleave(
+                        block_table, decode_lens, dim=0, output_size=actual_expanded
+                    )
+                )
+                if actual_expanded < num_decode_tokens:
+                    self.expanded_block_table_buffer[
+                        actual_expanded:num_decode_tokens, 0
+                    ] = 0
+                block_table = self.expanded_block_table_buffer[:num_decode_tokens]
+
+                # All reqs now have decode_len=1
+                self.decode_lens_buffer[:num_decode_tokens] = 1
+                decode_lens = self.decode_lens_buffer[:num_decode_tokens]
+                return seq_lens, block_table, decode_lens, num_decode_tokens, False
+        else:
+            # Native path: plain decode (next_n==1) or spec decode
+            # with 2D per-token context lengths (next_n > 1).
+            #
+            # When decode_lens are not truly uniform (e.g. some requests have
+            # decode_len < next_n due to padding or short prefills), the simple
+            # reshape in sparse_attn_indexer won't work. Use pack_seq_triton
+            # (requires_padding) instead.
+            requires_padding = min_decode_len != max_decode_len
+            if use_native and next_n > 1:
+                assert self.decode_seq_lens_buffer.dim() == 1
+                # (B, max_decode_len): token j attends to
+                # L - max_decode_len + j + 1 KV tokens.
+                seq_lens_buffer = self.decode_seq_lens_buffer[
+                    : num_decodes * max_decode_len
+                ].view(num_decodes, max_decode_len)
+                seq_lens_buffer[:] = (
+                    seq_lens.unsqueeze(1)
+                    - max_decode_len
+                    + 1
+                    + self.offsets_buffer[:max_decode_len]
+                )
+                seq_lens = seq_lens_buffer
+            return seq_lens, block_table, decode_lens, num_decodes, requires_padding
+
+    def build(
+        self,
+        common_prefix_len: int,
+        common_attn_metadata: CommonAttentionMetadata,
+        fast_build: bool = False,
+    ) -> DeepseekV32IndexerMetadata:
+        num_reqs = common_attn_metadata.num_reqs
+        num_tokens = common_attn_metadata.num_actual_tokens
+        query_start_loc = common_attn_metadata.query_start_loc
+        query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu
+        seq_lens = common_attn_metadata.seq_lens
+        slot_mapping = common_attn_metadata.slot_mapping
+        block_table = common_attn_metadata.block_table_tensor
+
+        num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
+            split_decodes_and_prefills(
+                common_attn_metadata,
+                decode_threshold=self.reorder_batch_threshold,
+                require_uniform=not self.use_flattening,
+            )
+        )
+
+        assert num_decodes + num_prefills == num_reqs
+        assert num_decode_tokens + num_prefill_tokens == num_tokens
+
+        compressed_slot_mapping = slot_mapping
+        compressed_seq_lens = seq_lens
+        if self.compress_ratio > 1:
+            compressed_slot_mapping = get_compressed_slot_mapping(
+                num_tokens,
+                query_start_loc,
+                seq_lens,
+                block_table,
+                self.kv_cache_spec.storage_block_size,
+                self.compress_ratio,
+                out=self.compressed_slot_mapping_buffer,
+            )
+            compressed_seq_lens = seq_lens // self.compress_ratio
+
+        prefill_metadata = None
+        if num_prefills > 0:
+            # This CPU value is an upper bound for async-spec extend rows.  It
+            # is safe for chunking/allocation because CUDA metadata below is
+            # built from exact device seq_lens and gather ignores the tail.
+            assert common_attn_metadata.seq_lens_cpu_upper_bound is not None
+            seq_lens_cpu = common_attn_metadata.seq_lens_cpu_upper_bound
+            compressed_seq_lens_cpu = (
+                seq_lens_cpu // self.compress_ratio
+                if self.compress_ratio > 1
+                else seq_lens_cpu
+            )
+            prefill_query_lens_cpu = torch.diff(
+                query_start_loc_cpu[num_decodes : num_decodes + num_prefills + 1]
+            )
+            max_logits_bytes = envs.VLLM_SPARSE_INDEXER_MAX_LOGITS_MB * 1024 * 1024
+            # Upper bound is exact for prefill rows (the `[num_decodes:]`
+            # slice below).
+            assert common_attn_metadata.seq_lens_cpu_upper_bound is not None
+            seq_lens_cpu = common_attn_metadata.seq_lens_cpu_upper_bound
+            chunk_specs = split_indexer_prefill_chunks(
+                compressed_seq_lens_cpu[num_decodes:],
+                prefill_query_lens_cpu,
+                self.max_prefill_buffer_size,
+                max_logits_bytes,
+                request_offset=num_decodes,
+            )
+
+            chunks = []
+            for req_slice, query_slice in chunk_specs:
+                metadata = build_prefill_chunk_metadata(
+                    req_slice.start,
+                    req_slice.stop,
+                    query_start_loc,
+                    query_start_loc_cpu,
+                    seq_lens,
+                    compressed_seq_lens,
+                    compressed_seq_lens_cpu,
+                    common_attn_metadata.block_table_tensor,
+                    self.compress_ratio,
+                    query_slice=query_slice,
+                    skip_kv_gather=query_slice.start > 0,
+                )
+                # Skip when total_seq_lens is 0 (i.e., no compressed token).
+                if metadata is not None:
+                    chunks.append(metadata)
+            prefill_metadata = DeepseekV32IndexerPrefillMetadata(chunks)
+
+        decode_metadata = None
+        if num_decodes > 0:
+            torch.diff(
+                common_attn_metadata.query_start_loc[: num_decodes + 1],
+                out=self.decode_lens_buffer[:num_decodes],
+            )
+            decode_lens = self.decode_lens_buffer[:num_decodes]
+            decode_lens_cpu = torch.diff(
+                common_attn_metadata.query_start_loc_cpu[: num_decodes + 1]
+            )
+
+            seq_lens = common_attn_metadata.seq_lens[:num_decodes]
+            block_table = common_attn_metadata.block_table_tensor[:num_decodes, ...]
+
+            max_decode_len = int(decode_lens_cpu.max().item())
+            next_n = 1 + self.num_speculative_tokens
+            use_native = not self.use_flattening and max_decode_len <= next_n
+
+            seq_lens, block_table, decode_lens, batch_size, requires_padding = (
+                self._prepare_decode_tensors(
+                    seq_lens=seq_lens,
+                    block_table=block_table,
+                    decode_lens=decode_lens,
+                    decode_lens_cpu=decode_lens_cpu,
+                    query_start_loc=common_attn_metadata.query_start_loc[:num_decodes],
+                    num_decodes=num_decodes,
+                    num_decode_tokens=num_decode_tokens,
+                    use_native=use_native,
+                    next_n=next_n,
+                    max_decode_len=max_decode_len,
+                )
+            )
+
+            # For DeepseekV4 (compress_ratio > 1), the indexer KV cache stores
+            # compressed tokens. Convert uncompressed seq_lens to compressed.
+            if self.compress_ratio > 1:
+                # True iff seq_lens aliases decode_seq_lens_buffer (flatten or
+                # native wrote it); False iff it aliases common_attn_metadata.
+                seq_lens_is_local_view = (use_native and next_n > 1) or (
+                    not use_native and max_decode_len > 1
+                )
+                if seq_lens_is_local_view:
+                    seq_lens //= self.compress_ratio
+                else:
+                    # Copy to avoid mutating shared state; keeps CG address stable.
+                    self.expanded_seq_lens_buffer[:num_decodes] = (
+                        seq_lens // self.compress_ratio
+                    )
+                    self.expanded_seq_lens_buffer[num_decodes:num_decode_tokens] = 0
+                    seq_lens = self.expanded_seq_lens_buffer[:num_decode_tokens]
+
+            # Non-MTP: deep_gemm paged MQA logits requires 2D context_lens
+            # (csrc/apis/attention.hpp). Unsqueeze to (B, 1) so downstream
+            # kernels see the same (B, next_n) layout as the MTP path.
+            if seq_lens.dim() == 1:
+                seq_lens = seq_lens.unsqueeze(-1)
+            seq_lens = seq_lens.contiguous()
+
+            # DeepGEMM is required for the paged MQA logits on CUDA devices
+            if current_platform.is_cuda() and has_deep_gemm():
+                self.scheduler_metadata_buffer[:] = get_paged_mqa_logits_metadata(
+                    seq_lens,
+                    self.kv_cache_spec.storage_block_size,
+                    self.num_sms,
+                )
+
+            decode_metadata = DeepSeekV32IndexerDecodeMetadata(
+                block_table=block_table,
+                seq_lens=seq_lens,
+                decode_lens=decode_lens,
+                requires_padding=requires_padding,
+                schedule_metadata=self.scheduler_metadata_buffer,
+            )
+
+        attn_metadata = DeepseekV32IndexerMetadata(
+            seq_lens=common_attn_metadata.seq_lens,
+            max_seq_len=common_attn_metadata.max_seq_len,
+            slot_mapping=compressed_slot_mapping,
+            num_decodes=num_decodes,
+            num_decode_tokens=num_decode_tokens,
+            num_prefills=num_prefills,
+            num_prefill_tokens=num_prefill_tokens,
+            prefill=prefill_metadata,
+            decode=decode_metadata,
+        )
+
+        return attn_metadata
+
+
+def build_prefill_chunk_metadata(
+    start_idx: int,
+    end_idx: int,
+    query_start_loc: torch.Tensor,
+    query_start_loc_cpu: torch.Tensor,
+    uncompressed_seq_lens: torch.Tensor,
+    compressed_seq_lens: torch.Tensor,
+    compressed_seq_lens_cpu: torch.Tensor,
+    block_table: torch.Tensor,
+    compress_ratio: int,
+    query_slice: slice | None = None,
+    skip_kv_gather: bool = False,
+) -> DeepseekV32IndexerPrefillChunkMetadata | None:
+    total_seq_lens = compressed_seq_lens_cpu[start_idx:end_idx].sum().item()
+    if total_seq_lens == 0:
+        return None
+
+    num_reqs = end_idx - start_idx
+    device = block_table.device
+    token_to_seq = torch.empty(total_seq_lens, dtype=torch.int32, device=device)
+
+    cu_seq_lens = torch.empty(num_reqs + 1, dtype=torch.int32, device=device)
+    # Assigning to slice avoids cpu sync.
+    cu_seq_lens[:1] = 0
+    torch.cumsum(compressed_seq_lens[start_idx:end_idx], dim=0, out=cu_seq_lens[1:])
+
+    query_start_loc = (
+        query_start_loc[start_idx : end_idx + 1] - query_start_loc[start_idx]
+    )
+
+    total_query_len = int(
+        (query_start_loc_cpu[end_idx] - query_start_loc_cpu[start_idx]).item()
+    )
+    if query_slice is not None:
+        qs_start = query_slice.start
+        qs_stop = query_slice.stop
+    else:
+        qs_start = 0
+        qs_stop = total_query_len
+    output_query_len = qs_stop - qs_start
+
+    cu_seq_len_ks = torch.empty(output_query_len, dtype=torch.int32, device=device)
+    cu_seq_len_ke = torch.empty(output_query_len, dtype=torch.int32, device=device)
+
+    _build_prefill_chunk_metadata_kernel[(num_reqs,)](
+        query_start_loc,
+        uncompressed_seq_lens[start_idx:end_idx],
+        cu_seq_lens,
+        token_to_seq,
+        cu_seq_len_ks,
+        cu_seq_len_ke,
+        qs_start,
+        qs_stop,
+        BLOCK_SIZE=1024,
+        COMPRESS_RATIO=compress_ratio,
+    )
+
+    token_start = query_start_loc_cpu[start_idx].item()
+    if query_slice is not None:
+        token_end = token_start + qs_stop
+        token_start = token_start + qs_start
+        skip_kv_gather = skip_kv_gather or qs_start > 0
+    else:
+        token_end = query_start_loc_cpu[end_idx].item()
+
+    return DeepseekV32IndexerPrefillChunkMetadata(
+        cu_seqlen_ks=cu_seq_len_ks,
+        cu_seqlen_ke=cu_seq_len_ke,
+        cu_seq_lens=cu_seq_lens,
+        token_to_seq=token_to_seq,
+        total_seq_lens=total_seq_lens,
+        block_table=block_table[start_idx:end_idx],
+        token_start=token_start,
+        token_end=token_end,
+        num_reqs=num_reqs,
+        skip_kv_gather=skip_kv_gather,
+    )
+
+
+@triton.jit
+def _build_prefill_chunk_metadata_kernel(
+    # Inputs
+    query_start_loc_ptr,
+    uncompressed_seq_lens_ptr,
+    cu_compressed_seq_lens_ptr,
+    # Outputs
+    token_to_seq_ptr,
+    cu_compressed_seq_len_ks_ptr,
+    cu_compressed_seq_len_ke_ptr,
+    query_slice_start,
+    query_slice_stop,
+    BLOCK_SIZE: tl.constexpr,
+    COMPRESS_RATIO: tl.constexpr,
+):
+    batch_idx = tl.program_id(0)
+
+    query_start = tl.load(query_start_loc_ptr + batch_idx)
+    query_end = tl.load(query_start_loc_ptr + batch_idx + 1)
+    query_len = query_end - query_start
+
+    seq_start = tl.load(cu_compressed_seq_lens_ptr + batch_idx)
+    seq_end = tl.load(cu_compressed_seq_lens_ptr + batch_idx + 1)
+    compressed_seq_len = seq_end - seq_start
+
+    uncompressed_seq_len = tl.load(uncompressed_seq_lens_ptr + batch_idx)
+    start_pos = uncompressed_seq_len - query_len
+
+    for i in range(0, query_len, BLOCK_SIZE):
+        offset = i + tl.arange(0, BLOCK_SIZE)
+        abs_pos = query_start + offset
+        mask = (
+            (offset < query_len)
+            & (abs_pos >= query_slice_start)
+            & (abs_pos < query_slice_stop)
+        )
+        out_pos = abs_pos - query_slice_start
+
+        # Compute cu_seq_len_ks
+        tl.store(cu_compressed_seq_len_ks_ptr + out_pos, seq_start, mask=mask)
+
+        # Compute cu_seq_len_ke
+        seq_len_per_token = (start_pos + 1 + offset) // COMPRESS_RATIO
+        tl.store(
+            cu_compressed_seq_len_ke_ptr + out_pos,
+            seq_start + seq_len_per_token,
+            mask=mask,
+        )
+
+    # Compute token_to_seq
+    for i in range(0, compressed_seq_len, BLOCK_SIZE):
+        offset = i + tl.arange(0, BLOCK_SIZE)
+        mask = offset < compressed_seq_len
+        tl.store(token_to_seq_ptr + seq_start + offset, batch_idx, mask=mask)

lmcache-config-dsv4.yaml

@@ -0,0 +1,13 @@
+chunk_size: 256
+local_cpu: true
+max_local_cpu_size: 512.0
+enable_lazy_memory_allocator: true
+lazy_memory_initial_ratio: 0.2
+remote_url: "redis://10.66.0.100:6379"
+remote_serde: naive
+remote_ttl: 1800
+save_decode_cache: true
+use_layerwise: true
+save_unfull_chunk: true
+blocking_timeout_secs: 60
+cache_policy: LRU

lmcache-config-glm-52.yaml

@@ -0,0 +1,10 @@
+chunk_size: 256
+local_cpu: true
+max_local_cpu_size: 384.0
+save_decode_cache: true
+enable_lazy_memory_allocator: true
+lazy_memory_initial_ratio: 1.0
+use_gpu_connector_v3: true
+remote_url: "redis://10.66.0.100:6379"
+remote_serde: naive
+remote_ttl: 3600

minimax_m2_parser.py

@@ -1,61 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-"""
-MiniMax M2 Parser - A unified parser for MiniMax M2 models.
-
-This parser combines the existing MiniMaxM2ReasoningParser and
-MinimaxM2ToolParser into a single unified interface by delegating
-to those implementations.
-"""
-
-from vllm.logger import init_logger
-from vllm.parser.abstract_parser import DelegatingParser
-from vllm.reasoning.minimax_m2_reasoning_parser import MiniMaxM2ReasoningParser
-from vllm.tokenizers import TokenizerLike
-from vllm.tool_parsers.abstract_tool_parser import (
-    Tool,
-)
-from vllm.tool_parsers.minimax_m2_tool_parser import MinimaxM2ToolParser
-
-logger = init_logger(__name__)
-
-
-class MiniMaxM2Parser(DelegatingParser):
-    """
-    Unified parser for MiniMax M2 models that handles both reasoning
-    extraction and tool call parsing.
-
-    This parser delegates to the existing implementations:
-    - MiniMaxM2ReasoningParser for reasoning extraction
-    - MinimaxM2ToolParser for tool call parsing
-
-    MiniMax M2 models have two special behaviors:
-    1. Reasoning: They don't generate <think> start token, only </think> end
-       token. All content before </think> is reasoning, content after is the
-       actual response.
-    2. Tool Calls: They use <minimax:tool_call>...</minimax:tool_call> tags
-       with <invoke name="...">...</invoke> and <parameter name="...">...</parameter>
-       syntax.
-    """
-
-    # Class-level parser classes for compatibility
-    reasoning_parser_cls = MiniMaxM2ReasoningParser
-    tool_parser_cls = MinimaxM2ToolParser
-
-    def __init__(
-        self,
-        tokenizer: TokenizerLike,
-        tools: list[Tool] | None = None,
-        *args,
-        **kwargs,
-    ):
-        super().__init__(tokenizer, *args, **kwargs)
-
-        # Initialize the underlying parsers
-        self._reasoning_parser = MiniMaxM2ReasoningParser(tokenizer, *args, **kwargs)
-        self._tool_parser = MinimaxM2ToolParser(tokenizer, tools)
-
-        logger.debug(
-            "vLLM Successfully initialized parser %s!", self.__class__.__name__
-        )

minimax_tool_parser.py

@@ -1,852 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import json
-from collections.abc import Sequence
-from typing import Any
-
-import regex as re
-
-from vllm.entrypoints.chat_utils import make_tool_call_id
-from vllm.entrypoints.openai.chat_completion.protocol import (
-    ChatCompletionRequest,
-)
-from vllm.entrypoints.openai.engine.protocol import (
-    DeltaFunctionCall,
-    DeltaMessage,
-    DeltaToolCall,
-    ExtractedToolCallInformation,
-    FunctionCall,
-    ToolCall,
-)
-from vllm.logger import init_logger
-from vllm.tokenizers import TokenizerLike
-from vllm.tool_parsers.abstract_tool_parser import (
-    Tool,
-    ToolParser,
-)
-from vllm.tool_parsers.utils import extract_intermediate_diff
-
-logger = init_logger(__name__)
-
-
-class MinimaxToolParser(ToolParser):
-    def __init__(self, tokenizer: TokenizerLike, tools: list[Tool] | None = None, **kwargs):
-        super().__init__(tokenizer, tools)
-
-        # Initialize streaming state for tracking tool call progress
-        self.streaming_state: dict[str, Any] = {
-            "current_tool_index": -1,  # Index of current tool being processed
-            "tool_ids": [],  # List of tool call IDs
-            "sent_tools": [],  # List of tools that have been sent
-        }
-
-        # Define tool call tokens and patterns
-        self.tool_call_start_token = "<tool_calls>"
-        self.tool_call_end_token = "</tool_calls>"
-        self.tool_call_regex = re.compile(
-            r"<tool_calls>(.*?)</tool_calls>|<tool_calls>(.*)", re.DOTALL
-        )
-        self.thinking_tag_pattern = r"<think>(.*?)</think>"
-        self.tool_name_pattern = re.compile(r'"name":\s*"([^"]+)"')
-        self.tool_args_pattern = re.compile(r'"arguments":\s*')
-
-        # Buffer for handling partial tool calls during streaming
-        self.pending_buffer = ""
-        self.in_thinking_tag = False
-
-        if not self.model_tokenizer:
-            raise ValueError(
-                "The model tokenizer must be passed to the ToolParser "
-                "constructor during construction."
-            )
-
-        # Get token IDs for tool call start/end tokens
-        self.tool_call_start_token_id = self.vocab.get(self.tool_call_start_token)
-        self.tool_call_end_token_id = self.vocab.get(self.tool_call_end_token)
-
-        if self.tool_call_start_token_id is None or self.tool_call_end_token_id is None:
-            logger.warning(
-                "Minimax Tool parser could not locate tool call start/end "
-                "tokens in the tokenizer. Falling back to string matching."
-            )
-
-    def preprocess_model_output(self, model_output: str) -> str:
-        """
-        Preprocess model output by removing tool calls from thinking tags.
-
-        Args:
-            model_output: Raw model output string
-
-        Returns:
-            Preprocessed model output with tool calls removed from thinking tags
-        """
-
-        def remove_tool_calls_from_think(match):
-            think_content = match.group(1)
-            cleaned_content = re.sub(
-                r"<tool_calls>.*?</tool_calls>", "", think_content, flags=re.DOTALL
-            )
-            return f"<think>{cleaned_content}</think>"
-
-        return re.sub(
-            self.thinking_tag_pattern,
-            remove_tool_calls_from_think,
-            model_output,
-            flags=re.DOTALL,
-        )
-
-    def _clean_duplicate_braces(self, args_text: str) -> str:
-        """
-        Clean duplicate closing braces from arguments text.
-
-        Args:
-            args_text: Raw arguments text
-
-        Returns:
-            Cleaned arguments text with proper JSON formatting
-        """
-        args_text = args_text.strip()
-        if not args_text:
-            return args_text
-
-        try:
-            json.loads(args_text)
-            return args_text
-        except json.JSONDecodeError:
-            pass
-
-        while args_text.endswith("}}"):
-            candidate = args_text[:-1]
-            try:
-                json.loads(candidate)
-                return candidate
-            except json.JSONDecodeError:
-                args_text = candidate
-
-        return args_text
-
-    def _clean_delta_braces(self, delta_text: str) -> str:
-        """
-        Clean delta text by removing excessive closing braces.
-
-        Args:
-            delta_text: Delta text to clean
-
-        Returns:
-            Cleaned delta text
-        """
-        if not delta_text:
-            return delta_text
-
-        delta_stripped = delta_text.strip()
-
-        if delta_stripped and all(c in "}\n\r\t " for c in delta_stripped):
-            brace_count = delta_stripped.count("}")
-            if brace_count > 1:
-                return "}\n" if delta_text.endswith("\n") else "}"
-
-        return delta_text
-
-    def extract_tool_calls(
-        self,
-        model_output: str,
-        request: ChatCompletionRequest,
-    ) -> ExtractedToolCallInformation:
-        """
-        Extract tool calls from model output for non-streaming mode.
-
-        Args:
-            model_output: Complete model output
-            request: Chat completion request
-
-        Returns:
-            ExtractedToolCallInformation containing tool calls and content
-        """
-        processed_output = self.preprocess_model_output(model_output)
-
-        if self.tool_call_start_token not in processed_output:
-            return ExtractedToolCallInformation(
-                tools_called=False, tool_calls=[], content=model_output
-            )
-
-        try:
-            function_call_tuples = self.tool_call_regex.findall(processed_output)
-
-            raw_function_calls = []
-            for match in function_call_tuples:
-                tool_call_content = match[0] if match[0] else match[1]
-                if tool_call_content.strip():
-                    lines = tool_call_content.strip().split("\n")
-                    for line in lines:
-                        line = line.strip()
-                        if line and line.startswith("{") and line.endswith("}"):
-                            try:
-                                parsed_call = json.loads(line)
-                                raw_function_calls.append(parsed_call)
-                            except json.JSONDecodeError:
-                                continue
-
-            tool_calls = []
-            for function_call in raw_function_calls:
-                if "name" in function_call and "arguments" in function_call:
-                    tool_calls.append(
-                        ToolCall(
-                            type="function",
-                            function=FunctionCall(
-                                name=function_call["name"],
-                                arguments=json.dumps(
-                                    function_call["arguments"], ensure_ascii=False
-                                ),
-                            ),
-                        )
-                    )
-
-            processed_pos = processed_output.find(self.tool_call_start_token)
-            if processed_pos != -1:
-                processed_content = processed_output[:processed_pos].strip()
-
-                if processed_content:
-                    lines = processed_content.split("\n")
-                    for line in reversed(lines):
-                        line = line.strip()
-                        if line:
-                            pos = model_output.find(line)
-                            if pos != -1:
-                                content = model_output[: pos + len(line)]
-                                break
-                    else:
-                        content = ""
-                else:
-                    content = ""
-            else:
-                content = model_output
-
-            return ExtractedToolCallInformation(
-                tools_called=len(tool_calls) > 0,
-                tool_calls=tool_calls,
-                content=content.strip() if content.strip() else None,
-            )
-
-        except Exception:
-            logger.exception(
-                "An unexpected error occurred during tool call extraction."
-            )
-            return ExtractedToolCallInformation(
-                tools_called=False, tool_calls=[], content=model_output
-            )
-
-    def _update_thinking_state(self, text: str) -> None:
-        """
-        Update the thinking tag state based on text content.
-
-        Args:
-            text: Text to analyze for thinking tags
-        """
-        open_count = text.count("<think>")
-        close_count = text.count("</think>")
-        self.in_thinking_tag = open_count > close_count or (
-            open_count == close_count and text.endswith("</think>")
-        )
-
-    def _is_potential_tag_start(self, text: str) -> bool:
-        """
-        Check if text might be the start of a tool call tag.
-
-        Args:
-            text: Text to check
-
-        Returns:
-            True if text could be the start of a tool call tag
-        """
-        for tag in [self.tool_call_start_token, self.tool_call_end_token]:
-            if any(
-                tag.startswith(text[-i:])
-                for i in range(1, min(len(text) + 1, len(tag)))
-            ):
-                return True
-        return False
-
-    def _should_buffer_content(self, delta_text: str) -> bool:
-        """
-        Determine if content should be buffered for later processing.
-
-        Args:
-            delta_text: Delta text to check
-
-        Returns:
-            True if content should be buffered
-        """
-        if self.in_thinking_tag:
-            return False
-        return bool(
-            self.pending_buffer
-            or self.tool_call_start_token in delta_text
-            or self.tool_call_end_token in delta_text
-            or delta_text.startswith("<")
-        )
-
-    def _split_content_for_buffering(self, delta_text: str) -> tuple[str, str]:
-        """
-        Split delta text into safe content and potential tag content.
-
-        Args:
-            delta_text: Delta text to split
-
-        Returns:
-            Tuple of (safe_content, potential_tag_content)
-        """
-        if self.in_thinking_tag:
-            return delta_text, ""
-
-        for tag in [self.tool_call_start_token, self.tool_call_end_token]:
-            for i in range(1, len(tag)):
-                tag_prefix = tag[:i]
-                pos = delta_text.rfind(tag_prefix)
-                if pos != -1 and tag.startswith(delta_text[pos:]):
-                    return delta_text[:pos], delta_text[pos:]
-        return delta_text, ""
-
-    def _process_buffer(self, new_content: str) -> str:
-        """
-        Process buffered content and return output content.
-
-        Args:
-            new_content: New content to add to buffer
-
-        Returns:
-            Processed output content
-        """
-        self.pending_buffer += new_content
-        output_content = ""
-
-        if self.in_thinking_tag:
-            output_content = self.pending_buffer
-            self.pending_buffer = ""
-            return output_content
-
-        while self.pending_buffer:
-            start_pos = self.pending_buffer.find(self.tool_call_start_token)
-            end_pos = self.pending_buffer.find(self.tool_call_end_token)
-
-            if start_pos != -1 and (end_pos == -1 or start_pos < end_pos):
-                tag_pos, tag_len = start_pos, len(self.tool_call_start_token)
-            elif end_pos != -1:
-                tag_pos, tag_len = end_pos, len(self.tool_call_end_token)
-            else:
-                if self._is_potential_tag_start(self.pending_buffer):
-                    break
-                output_content += self.pending_buffer
-                self.pending_buffer = ""
-                break
-
-            output_content += self.pending_buffer[:tag_pos]
-            self.pending_buffer = self.pending_buffer[tag_pos + tag_len :]
-
-        return output_content
-
-    def _reset_streaming_state(self) -> None:
-        """Reset the streaming state to initial values."""
-        self.streaming_state = {
-            "current_tool_index": -1,
-            "tool_ids": [],
-            "sent_tools": [],
-        }
-
-    def _advance_to_next_tool(self) -> None:
-        """Advance to the next tool in the streaming sequence."""
-        self.streaming_state["current_tool_index"] = (
-            int(self.streaming_state["current_tool_index"]) + 1
-        )
-
-    def _set_current_tool_index(self, index: int) -> None:
-        """
-        Set the current tool index.
-
-        Args:
-            index: Tool index to set
-        """
-        self.streaming_state["current_tool_index"] = index
-
-    def _get_current_tool_index(self) -> int:
-        """
-        Get the current tool index.
-
-        Returns:
-            Current tool index
-        """
-        return int(self.streaming_state["current_tool_index"])
-
-    def _get_next_unsent_tool_index(self, tool_count: int) -> int:
-        """
-        Get the index of the next unsent tool.
-
-        Args:
-            tool_count: Total number of tools
-
-        Returns:
-            Index of next unsent tool, or -1 if all tools sent
-        """
-        sent_tools = list(self.streaming_state["sent_tools"])
-        for i in range(tool_count):
-            if i < len(sent_tools):
-                if not sent_tools[i]["sent_name"]:
-                    return i
-            else:
-                return i
-        return -1
-
-    def _ensure_state_arrays(self, tool_count: int) -> None:
-        """
-        Ensure state arrays have sufficient capacity for tool_count tools.
-
-        Args:
-            tool_count: Number of tools to prepare for
-        """
-        sent_tools = list(self.streaming_state["sent_tools"])
-        tool_ids = list(self.streaming_state["tool_ids"])
-
-        while len(sent_tools) < tool_count:
-            sent_tools.append(
-                {
-                    "sent_name": False,
-                    "sent_arguments": "",
-                    "id": make_tool_call_id(),
-                }
-            )
-
-        while len(tool_ids) < tool_count:
-            tool_ids.append(None)
-
-        self.streaming_state["sent_tools"] = sent_tools
-        self.streaming_state["tool_ids"] = tool_ids
-
-    def _detect_tools_in_text(self, text: str) -> int:
-        """
-        Detect the number of tools in text by counting name patterns.
-
-        Args:
-            text: Text to analyze
-
-        Returns:
-            Number of tools detected
-        """
-        matches = self.tool_name_pattern.findall(text)
-        return len(matches)
-
-    def _find_tool_boundaries(self, text: str) -> list[tuple[int, int]]:
-        """
-        Find the boundaries of tool calls in text.
-
-        Args:
-            text: Text to analyze
-
-        Returns:
-            List of (start, end) positions for tool calls
-        """
-        boundaries = []
-        i = 0
-        while i < len(text):
-            if text[i] == "{":
-                start = i
-                depth = 0
-                has_name = False
-                has_arguments = False
-
-                while i < len(text):
-                    if text[i] == "{":
-                        depth += 1
-                    elif text[i] == "}":
-                        depth -= 1
-                        if depth == 0:
-                            end = i + 1
-                            segment = text[start:end]
-                            if '"name"' in segment and '"arguments"' in segment:
-                                boundaries.append((start, end))
-                            break
-
-                    if not has_name and '"name"' in text[start : i + 1]:
-                        has_name = True
-                    if not has_arguments and '"arguments"' in text[start : i + 1]:
-                        has_arguments = True
-
-                    i += 1
-
-                if depth > 0 and has_name:
-                    boundaries.append((start, i))
-            else:
-                i += 1
-        return boundaries
-
-    def _extract_tool_args(self, tool_content: str, args_match: re.Match[str]) -> str:
-        """
-        Extract tool arguments from tool content.
-
-        Args:
-            tool_content: Tool call content
-            args_match: Regex match for arguments pattern
-
-        Returns:
-            Extracted arguments as string
-        """
-        args_start_pos = args_match.end()
-        remaining_content = tool_content[args_start_pos:]
-
-        if remaining_content.strip().startswith("{"):
-            depth = 0
-            for i, char in enumerate(remaining_content):
-                if char == "{":
-                    depth += 1
-                elif char == "}":
-                    depth -= 1
-                    if depth == 0:
-                        return remaining_content[: i + 1]
-        else:
-            args_end = remaining_content.find("}")
-            if args_end > 0:
-                return remaining_content[:args_end].strip()
-
-        return remaining_content.rstrip("}").strip()
-
-    def _get_current_tool_content(
-        self, text: str, tool_index: int
-    ) -> tuple[str | None, str | None]:
-        """
-        Get the content of a specific tool by index.
-
-        Args:
-            text: Text containing tool calls
-            tool_index: Index of tool to extract
-
-        Returns:
-            Tuple of (tool_name, tool_arguments) or (None, None) if not found
-        """
-        boundaries = self._find_tool_boundaries(text)
-
-        if tool_index >= len(boundaries):
-            return None, None
-
-        start, end = boundaries[tool_index]
-        tool_content = text[start:end]
-
-        name_match = self.tool_name_pattern.search(tool_content)
-        name = name_match.group(1) if name_match else None
-
-        args_match = self.tool_args_pattern.search(tool_content)
-        if args_match:
-            try:
-                args_text = self._extract_tool_args(tool_content, args_match)
-                return name, args_text
-            except Exception:
-                remaining_content = tool_content[args_match.end() :]
-                args_text = remaining_content.rstrip("}").strip()
-                return name, args_text
-
-        return name, None
-
-    def _handle_tool_name_streaming(
-        self, tool_content: str, tool_count: int
-    ) -> DeltaMessage | None:
-        """
-        Handle streaming of tool names.
-
-        Args:
-            tool_content: Content containing tool calls
-            tool_count: Total number of tools
-
-        Returns:
-            DeltaMessage with tool name or None if no tool to stream
-        """
-        next_idx = self._get_next_unsent_tool_index(tool_count)
-
-        if next_idx == -1:
-            return None
-
-        boundaries = self._find_tool_boundaries(tool_content)
-        if next_idx >= len(boundaries):
-            return None
-
-        tool_name, _ = self._get_current_tool_content(tool_content, next_idx)
-        if not tool_name:
-            return None
-
-        self._set_current_tool_index(next_idx)
-        sent_tools = list(self.streaming_state["sent_tools"])
-        tool_ids = list(self.streaming_state["tool_ids"])
-
-        tool_id = sent_tools[next_idx]["id"]
-        tool_ids[next_idx] = tool_id
-        sent_tools[next_idx]["sent_name"] = True
-
-        self.streaming_state["sent_tools"] = sent_tools
-        self.streaming_state["tool_ids"] = tool_ids
-
-        return DeltaMessage(
-            tool_calls=[
-                DeltaToolCall(
-                    index=next_idx,
-                    type="function",
-                    id=tool_id,
-                    function=DeltaFunctionCall(name=tool_name).model_dump(
-                        exclude_none=True
-                    ),
-                )
-            ]
-        )
-
-    def _handle_tool_args_streaming(
-        self, tool_content: str, tool_count: int
-    ) -> DeltaMessage | None:
-        """
-        Handle streaming of tool arguments.
-
-        Args:
-            tool_content: Content containing tool calls
-            tool_count: Total number of tools
-
-        Returns:
-            DeltaMessage with tool arguments or None if no arguments to stream
-        """
-        current_idx = self._get_current_tool_index()
-
-        if current_idx < 0 or current_idx >= tool_count:
-            return None
-
-        tool_name, tool_args = self._get_current_tool_content(tool_content, current_idx)
-        if not tool_name or tool_args is None:
-            return None
-
-        sent_tools = list(self.streaming_state["sent_tools"])
-
-        if not sent_tools[current_idx]["sent_name"]:
-            return None
-
-        clean_args = self._clean_duplicate_braces(tool_args)
-        sent_args = sent_tools[current_idx]["sent_arguments"]
-
-        if clean_args != sent_args:
-            if sent_args and clean_args.startswith(sent_args):
-                args_delta = extract_intermediate_diff(clean_args, sent_args)
-                if args_delta:
-                    args_delta = self._clean_delta_braces(args_delta)
-                    sent_tools[current_idx]["sent_arguments"] = clean_args
-                    self.streaming_state["sent_tools"] = sent_tools
-
-                    if clean_args.endswith("}"):
-                        self._advance_to_next_tool()
-
-                    return DeltaMessage(
-                        tool_calls=[
-                            DeltaToolCall(
-                                index=current_idx,
-                                function=DeltaFunctionCall(
-                                    arguments=args_delta
-                                ).model_dump(exclude_none=True),
-                            )
-                        ]
-                    )
-            elif not sent_args and clean_args:
-                clean_args_delta = self._clean_delta_braces(clean_args)
-                sent_tools[current_idx]["sent_arguments"] = clean_args
-                self.streaming_state["sent_tools"] = sent_tools
-
-                if clean_args.endswith("}"):
-                    self._advance_to_next_tool()
-
-                return DeltaMessage(
-                    tool_calls=[
-                        DeltaToolCall(
-                            index=current_idx,
-                            function=DeltaFunctionCall(
-                                arguments=clean_args_delta
-                            ).model_dump(exclude_none=True),
-                        )
-                    ]
-                )
-
-        return None
-
-    def _is_end_tool_calls(self, current_text: str) -> bool:
-        if self.tool_call_end_token not in current_text:
-            return False
-
-        end_token_positions = []
-        search_start = 0
-        while True:
-            pos = current_text.find(self.tool_call_end_token, search_start)
-            if pos == -1:
-                break
-            end_token_positions.append(pos)
-            search_start = pos + 1
-
-        think_regions = []
-        for match in re.finditer(
-            self.thinking_tag_pattern, current_text, flags=re.DOTALL
-        ):
-            think_regions.append((match.start(), match.end()))
-
-        for pos in end_token_positions:
-            in_think = any(
-                pos >= t_start and pos < t_end for t_start, t_end in think_regions
-            )
-            if not in_think:
-                return True
-
-        return False
-
-    def extract_tool_calls_streaming(
-        self,
-        previous_text: str,
-        current_text: str,
-        delta_text: str,
-        previous_token_ids: Sequence[int],
-        current_token_ids: Sequence[int],
-        delta_token_ids: Sequence[int],
-        request: ChatCompletionRequest,
-    ) -> DeltaMessage | None:
-        self._update_thinking_state(current_text)
-
-        if self.in_thinking_tag:
-            return DeltaMessage(content=delta_text)
-
-        if self._should_buffer_content(delta_text):
-            buffered_output = self._process_buffer(delta_text)
-            return DeltaMessage(content=buffered_output) if buffered_output else None
-
-        if self._is_end_tool_calls(current_text):
-            return DeltaMessage(content=delta_text)
-
-        safe_content, potential_tag = self._split_content_for_buffering(delta_text)
-        if potential_tag:
-            self.pending_buffer += potential_tag
-            return DeltaMessage(content=safe_content) if safe_content else None
-
-        processed_current_text = self.preprocess_model_output(current_text)
-
-        if self.tool_call_start_token not in processed_current_text:
-            if (
-                self.tool_call_end_token in delta_text
-                and self.tool_call_start_token in current_text
-            ):
-                return None
-            if delta_text.strip() == "" and self.tool_call_start_token in current_text:
-                return None
-            if (
-                self._get_current_tool_index() != -1
-                and self.tool_call_end_token in current_text
-            ):
-                self._reset_streaming_state()
-            return DeltaMessage(content=delta_text)
-
-        if (
-            self.tool_call_start_token_id is not None
-            and self.tool_call_start_token_id in delta_token_ids
-            and len(delta_token_ids) == 1
-        ):
-            return None
-
-        original_tool_start = self._find_tool_start_outside_thinking(current_text)
-        if original_tool_start is None:
-            return None
-
-        content_before_tools = self._extract_content_before_tools(
-            current_text, delta_text, original_tool_start
-        )
-        if content_before_tools:
-            return DeltaMessage(content=content_before_tools)
-
-        try:
-            tool_content = self._extract_tool_content(current_text, original_tool_start)
-            current_tools_count = self._detect_tools_in_text(tool_content)
-
-            if current_tools_count == 0:
-                return None
-
-            if self._get_current_tool_index() == -1:
-                self._reset_streaming_state()
-
-            self._ensure_state_arrays(current_tools_count)
-
-            return self._handle_tool_name_streaming(
-                tool_content, current_tools_count
-            ) or self._handle_tool_args_streaming(tool_content, current_tools_count)
-
-        except Exception:
-            logger.exception(
-                "An unexpected error occurred ", "during streaming tool call handling."
-            )
-            return None
-
-    def _find_tool_start_outside_thinking(self, current_text: str) -> int | None:
-        """
-        Find the start position of tool calls outside of thinking tags.
-
-        Args:
-            current_text: Current text to search
-
-        Returns:
-            Position of tool call start or None if not found
-        """
-        search_start = 0
-        while True:
-            pos = current_text.find(self.tool_call_start_token, search_start)
-            if pos == -1:
-                return None
-
-            think_regions = [
-                (m.start(), m.end())
-                for m in re.finditer(
-                    r"<think>(.*?)</think>", current_text, flags=re.DOTALL
-                )
-            ]
-            in_think = any(
-                pos >= t_start and pos < t_end for t_start, t_end in think_regions
-            )
-
-            if not in_think:
-                return pos
-
-            search_start = pos + 1
-
-    def _extract_content_before_tools(
-        self, current_text: str, delta_text: str, tool_start: int
-    ) -> str | None:
-        """
-        Extract content that appears before tool calls.
-
-        Args:
-            current_text: Current text
-            delta_text: Delta text
-            tool_start: Start position of tools
-
-        Returns:
-            Content before tools or None
-        """
-        if tool_start > 0:
-            delta_start_pos = len(current_text) - len(delta_text)
-            if delta_start_pos < tool_start:
-                content_part = delta_text
-                if delta_start_pos + len(delta_text) > tool_start:
-                    content_part = delta_text[: tool_start - delta_start_pos]
-                return content_part if content_part else None
-        return None
-
-    def _extract_tool_content(self, current_text: str, tool_start: int) -> str:
-        """
-        Extract tool content from current text starting at tool_start.
-
-        Args:
-            current_text: Current text
-            tool_start: Start position of tool calls
-
-        Returns:
-            Extracted tool content
-        """
-        tool_content_start = tool_start + len(self.tool_call_start_token)
-        tool_content = current_text[tool_content_start:]
-
-        end_pos = tool_content.find(self.tool_call_end_token)
-        if end_pos != -1:
-            tool_content = tool_content[:end_pos]
-
-        return tool_content

shm_broadcast.py

@@ -0,0 +1,942 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import functools
+import pickle
+import sys
+import threading
+import time
+from contextlib import contextmanager
+from dataclasses import dataclass, field
+from multiprocessing import shared_memory
+from pickle import PickleBuffer
+from typing import TYPE_CHECKING, Any, cast
+from unittest.mock import patch
+
+import torch
+import torch.distributed as dist
+import zmq
+from torch.distributed import ProcessGroup
+from zmq import (  # type: ignore
+    IPV6,  # type: ignore
+    PUB,
+    SUB,
+    SUBSCRIBE,
+    XPUB,
+    XPUB_VERBOSE,
+    Context,
+)
+
+import vllm.envs as envs
+from vllm.distributed.utils import StatelessProcessGroup, sched_yield
+from vllm.logger import init_logger
+from vllm.platforms import current_platform
+from vllm.utils.network_utils import (
+    get_ip,
+    get_open_port,
+    get_open_zmq_inproc_path,
+    get_open_zmq_ipc_path,
+    is_valid_ipv6_address,
+)
+
+logger = init_logger(__name__)
+
+
+SPINLOOP_EXT_ENABLED = False
+if envs.VLLM_USE_SPINLOOP_EXT:
+    try:
+        from vllm.spinloop import spinloop
+
+        SPINLOOP_EXT_ENABLED = True
+    except ImportError:
+        logger.warning(
+            "spinloop extension could not be loaded, disabling VLLM_USE_SPINLOOP_EXT!"
+        )
+SPINLOOP_TIMEOUT_SECONDS = 0.1
+
+if TYPE_CHECKING:
+    from _typeshed import SizedBuffer
+
+VLLM_RINGBUFFER_WARNING_INTERVAL = envs.VLLM_RINGBUFFER_WARNING_INTERVAL
+
+from_bytes_big = functools.partial(int.from_bytes, byteorder="big")
+
+
+# Memory fence for cross-process shared memory visibility.
+# Required for correct producer-consumer synchronization when using
+# shared memory without locks.
+_memory_fence_lock = threading.Lock()
+
+
+def memory_fence():
+    """
+    Full memory barrier for shared memory synchronization.
+
+    Ensures all prior memory writes are visible to other processes before
+    any subsequent reads. This is critical for lock-free producer-consumer
+    patterns using shared memory.
+
+    Implementation acquires and immediately releases a lock. Python's
+    threading.Lock provides sequentially consistent memory barrier semantics
+    across all major platforms (POSIX, Windows). This is a lightweight
+    operation (~20ns) that guarantees:
+    - All stores before the barrier are visible to other threads/processes
+    - All loads after the barrier see the latest values
+    """
+    # Lock acquire/release provides full memory barrier semantics.
+    # Using context manager ensures lock release even on exceptions.
+    with _memory_fence_lock:
+        pass
+
+
+def to_bytes_big(value: int, size: int) -> bytes:
+    return value.to_bytes(size, byteorder="big")
+
+
+LONG_WAIT_TIME_LOG_MSG = (
+    "No available shared memory broadcast block found "
+    "in %d seconds. This typically happens "
+    "when some processes are hanging or doing some "
+    "time-consuming work (e.g. compilation, "
+    "weight/kv cache quantization)."
+)
+
+
+class SpinCondition:
+    """
+    This class implements an interface similar to a threading.Condition. It
+    allows a writer to notify readers to wake up and read from the shared memory
+    buffer. This notification is done over a zmq socket.
+
+    For optimal performance under load we don't want the readers to need to poll
+    the zmq socket for every read. So the `wait` method here will return
+    immediately when reads are frequent, and will only enter "idle mode" and
+    await a notification on the zmq socket after a period of inactivity. This
+    allows the readers to spin quickly, hence "SpinCondition".
+
+    To support clean shutdown, a separate thread in the reader's process must be
+    able to wake the reader so that it can exit. A separate cancel() method is
+    implemented with an in-process socket to allow this interruption.
+    """
+
+    def __init__(
+        self,
+        is_reader: bool,
+        context: zmq.Context,
+        notify_address: str,
+        busy_loop_s: float = 1,
+    ):
+        self.is_reader = is_reader
+
+        if is_reader:
+            # Time of last shm buffer read
+            self.last_read = time.monotonic()
+
+            # Time to keep busy-looping on the shm buffer before going idle
+            self.busy_loop_s = busy_loop_s
+
+            # Readers subscribe to write notifications
+            self.local_notify_socket: zmq.Socket = context.socket(SUB)
+            # Set zmq.CONFLATE to only keep the last message that the socket
+            # receives. This prevents us from piling up notification messages
+            # under high load when we aren't polling the socket.
+            self.local_notify_socket.setsockopt(zmq.CONFLATE, 1)
+            # Subscribe to all messages on the socket
+            self.local_notify_socket.setsockopt_string(SUBSCRIBE, "")
+            self.local_notify_socket.connect(notify_address)
+
+            # Readers require a process-local socket to poll for cancellation
+            cancel_path = get_open_zmq_inproc_path()
+            self.write_cancel_socket: zmq.Socket = context.socket(zmq.PAIR)
+            self.write_cancel_socket.bind(cancel_path)
+            self.read_cancel_socket: zmq.Socket = context.socket(zmq.PAIR)
+            self.read_cancel_socket.connect(cancel_path)
+
+            # Poller allows waiting on either `.notify()` or `.cancel()`
+            self.poller = zmq.Poller()
+            self.poller.register(self.read_cancel_socket, zmq.POLLIN)
+            self.poller.register(self.local_notify_socket, zmq.POLLIN)
+        else:
+            # Writer side publishes write notifications
+            self.local_notify_socket: zmq.Socket = context.socket(PUB)  # type: ignore
+            # Set high water mark to 1 - we don't need to send a massive amount of
+            # pings during busy operation. PUB sockets will silently drop subsequent
+            # messages after the high water mark is reached.
+            self.local_notify_socket.setsockopt(zmq.SNDHWM, 1)
+            self.local_notify_socket.bind(notify_address)
+
+            self.last_read = 0
+            self.busy_loop_s = 0
+            self.read_cancel_socket = None
+            self.write_cancel_socket = None
+            self.poller = None
+
+    def record_read(self):
+        self.last_read = time.monotonic()
+
+    def cancel(self):
+        # Sends cancellation ping that will cause the reader to wake up.
+        # This is done from a monitor thread in the same process as the reader.
+        if self.is_reader:
+            logger.debug("Canceling waiting reads on SHM Buffer")
+            self.write_cancel_socket.send(b"\x00")
+
+    def wait(self, timeout_ms: int | None = None) -> None:
+        """Wait for data on the shared memory buffer.
+
+        Yields the scheduler then returns immediately if it has been less than
+        self.busy_loop_s since the last read.
+
+        Otherwise, enters idle mode and awaits a socket ping for at most
+        `timeout_ms` milliseconds, or indefinitely if timeout_ms is None.
+        """
+        assert self.is_reader, "Only readers can wait"
+
+        current_time = time.monotonic()
+        if current_time <= self.last_read + self.busy_loop_s:
+            sched_yield()
+        else:
+            events = dict(self.poller.poll(timeout=timeout_ms))
+
+            if self.read_cancel_socket in events:
+                logger.debug("Poller received cancel event")
+            elif self.local_notify_socket in events:
+                logger.debug("Poller received notify event")
+                # Since zmq.CONFLATE is set, there will only be one notification
+                # to read from the socket
+                self.local_notify_socket.recv(flags=zmq.NOBLOCK, copy=False)
+            else:
+                logger.debug("Poller timed out")
+
+    def notify(self):
+        """Notifies all readers to wake up"""
+        assert not self.is_reader, "Only writers can notify"
+        self.local_notify_socket.send(b"\x00")
+
+
+class ShmRingBuffer:
+    def __init__(
+        self,
+        n_reader: int,
+        max_chunk_bytes: int,
+        max_chunks: int,
+        name: str | None = None,
+    ):
+        """
+        A shared memory ring buffer implementation for broadcast communication.
+        Essentially, it is a queue where only one will `enqueue` and multiple
+        will `dequeue`. The max size of each item, together with the max number
+        of items that can be stored in the buffer are known in advance.
+        In this case, we don't need to synchronize the access to
+         the buffer.
+
+        Buffer memory layout:
+                  data                                 metadata
+                    |                                      |
+                    | (current_idx)                        | (current_idx)
+                    v                                      v
+        +-------------------------------+----------------------------------------+
+        | chunk0 | chunk1 | ... | chunk | metadata0 | metadata1 | ... | metadata |
+        +-------------------------------+----------------------------------------+
+        | max_chunks x max_chunk_bytes  | max_chunks x (1 + n_reader) bytes      |
+
+        metadata memory layout: each byte is a flag, the first byte is the written
+        flag, and the rest are reader flags. The flags are set to 0 by default.
+        +--------------+--------------+--------------+-----+--------------+
+        | written_flag | reader0_flag | reader1_flag | ... | readerN_flag |
+        +--------------+--------------+--------------+-----+--------------+
+
+        The state of metadata is as follows:
+
+        (case 1) 0???...???: the block is not written yet, cannot read, can write
+        (case 2) 1000...000: the block is just written, can read, cannot write
+        (case 3) 1???...???: the block is written and read by some readers, can read if not read, cannot write
+        (case 4) 1111...111: the block is written and read by all readers, cannot read, can write
+
+        State transition for readers:
+
+        When a reader finds a block that it can read (case 2 or 3), it can yield the block for caller to read.
+        Only after the caller finishes reading the block, the reader can mark the block as read.
+        Readers only mark the block as read (from 0 to 1), the writer marks the block as ready to read (from 1 to 0).
+
+        State transition for writer:
+
+        When the writer writes to a block (case 1 or 4), it first resets the written flag to 0, converting either case
+        to case 1. Then it can yield the block for caller to write. After the caller finishes writing the block, the writer
+        can reset the reader flags to 0, and mark the block as written (from 0 to 1).
+        NOTE: the order is important here, first reset the reader flags (so that we are still in case 1), then mark the block as written. The state transition is atomic. If we do it in the reverse order, it will go through case 3 and then back to case 2, and readers might read the intermediate case 3, which is not correct.
+
+        During creation, `name` is None and the buffer is created. We can pass the
+        created object to other processes by pickling it. The other processes will
+        get the name of the shared memory and open it, so that they can access the
+        same shared memory buffer.
+        """  # noqa
+        self.n_reader = n_reader
+        self.metadata_size = 1 + n_reader
+        self.max_chunk_bytes = max_chunk_bytes
+        self.max_chunks = max_chunks
+        self.total_bytes_of_buffer = (
+            self.max_chunk_bytes + self.metadata_size
+        ) * self.max_chunks
+        self.data_offset = 0
+        self.metadata_offset = self.max_chunk_bytes * self.max_chunks
+
+        if name is None:
+            # we are creating a buffer
+            self.is_creator = True
+            self.shared_memory = shared_memory.SharedMemory(
+                create=True, size=self.total_bytes_of_buffer
+            )
+            assert self.shared_memory.buf is not None, "Buffer was not created"
+            # initialize the metadata section to 0
+            with self.shared_memory.buf[self.metadata_offset :] as metadata_buffer:
+                torch.frombuffer(metadata_buffer, dtype=torch.uint8).fill_(0)
+        else:
+            # we are opening an existing buffer
+            self.is_creator = False
+            # fix to https://stackoverflow.com/q/62748654/9191338
+            # Python incorrectly tracks shared memory even if it is not
+            # created by the process. The following patch is a workaround.
+            with patch(
+                "multiprocessing.resource_tracker.register",
+                lambda *args, **kwargs: None,
+            ):
+                try:
+                    self.shared_memory = shared_memory.SharedMemory(name=name)
+                    # See https://docs.python.org/3/library/multiprocessing.shared_memory.html # noqa
+                    # Some platforms allocate memory based on page size,
+                    # so the shared memory block size may be larger or equal
+                    # to the requested size. The size parameter is ignored
+                    # when attaching to an existing block.
+                    assert self.shared_memory.size >= self.total_bytes_of_buffer
+                except FileNotFoundError:
+                    # we might deserialize the object in a different node
+                    # in this case, this object is not used,
+                    # and we should suppress the error
+                    pass
+
+    def handle(self):
+        return (
+            self.n_reader,
+            self.max_chunk_bytes,
+            self.max_chunks,
+            self.shared_memory.name,
+        )
+
+    def __reduce__(self):
+        return (
+            self.__class__,
+            self.handle(),
+        )
+
+    def __del__(self):
+        if hasattr(self, "shared_memory"):
+            self.shared_memory.close()
+            if self.is_creator:
+                self.shared_memory.unlink()
+
+    @contextmanager
+    def get_data(self, current_idx: int):
+        start = self.data_offset + current_idx * self.max_chunk_bytes
+        end = start + self.max_chunk_bytes
+        assert self.shared_memory.buf is not None, "Buffer has been closed"
+        with self.shared_memory.buf[start:end] as buf:
+            yield buf
+
+    @contextmanager
+    def get_metadata(self, current_idx: int):
+        start = self.metadata_offset + current_idx * self.metadata_size
+        end = start + self.metadata_size
+        assert self.shared_memory.buf is not None, "Buffer has been closed"
+        with self.shared_memory.buf[start:end] as buf:
+            yield buf
+
+
+@dataclass
+class Handle:
+    local_reader_ranks: list[int] = field(default_factory=list)
+
+    buffer_handle: tuple[int, int, int, str] | None = None
+    local_subscribe_addr: str | None = None
+    local_notify_addr: str | None = None
+    remote_subscribe_addr: str | None = None
+    remote_addr_ipv6: bool = False
+
+
+class MessageQueue:
+    def __init__(
+        self,
+        n_reader,  # number of all readers
+        n_local_reader,  # number of local readers through shared memory
+        local_reader_ranks: list[int] | None = None,
+        # Default of 24MiB chosen to be large enough to accommodate grammar
+        # bitmask tensors for large batches (1024 requests).
+        max_chunk_bytes: int = 1024 * 1024 * 24,
+        max_chunks: int = 10,
+        connect_ip: str | None = None,
+    ):
+        if local_reader_ranks is None:
+            local_reader_ranks = list(range(n_local_reader))
+        else:
+            assert len(local_reader_ranks) == n_local_reader
+        self.n_local_reader = n_local_reader
+        n_remote_reader = n_reader - n_local_reader
+        self.n_remote_reader = n_remote_reader
+        self.shutting_down = False
+        context = Context()
+
+        if n_local_reader > 0:
+            # for local readers, we will:
+            # 1. create a shared memory ring buffer to communicate small data
+            # 2. create a publish-subscribe socket to communicate large data
+            self.buffer = ShmRingBuffer(n_local_reader, max_chunk_bytes, max_chunks)
+
+            # XPUB is very similar to PUB,
+            # except that it can receive subscription messages
+            # to confirm the number of subscribers
+            self.local_socket = context.socket(XPUB)
+            # set the verbose option so that we can receive every subscription
+            # message. otherwise, we will only receive the first subscription
+            # see http://api.zeromq.org/3-3:zmq-setsockopt for more details
+            self.local_socket.setsockopt(XPUB_VERBOSE, True)
+            local_subscribe_addr = get_open_zmq_ipc_path()
+            logger.debug("Binding to %s", local_subscribe_addr)
+            self.local_socket.bind(local_subscribe_addr)
+
+            self.current_idx = 0
+
+            # Create the notification side of the SpinCondition
+            local_notify_addr = get_open_zmq_ipc_path()
+            self._spin_condition = SpinCondition(
+                is_reader=False, context=context, notify_address=local_notify_addr
+            )
+        else:
+            self.buffer = None  # type: ignore
+            local_subscribe_addr = None
+            self.local_socket = None
+            self.current_idx = -1
+            local_notify_addr = None
+            self._spin_condition = None  # type: ignore
+
+        remote_addr_ipv6 = False
+        if n_remote_reader > 0:
+            # for remote readers, we will:
+            # create a publish-subscribe socket to communicate large data
+            if not connect_ip:
+                connect_ip = get_ip()
+            self.remote_socket = context.socket(XPUB)
+            self.remote_socket.setsockopt(XPUB_VERBOSE, True)
+            remote_subscribe_port = get_open_port()
+            if is_valid_ipv6_address(connect_ip):
+                self.remote_socket.setsockopt(IPV6, 1)
+                remote_addr_ipv6 = True
+                connect_ip = f"[{connect_ip}]"
+            socket_addr = f"tcp://{connect_ip}:{remote_subscribe_port}"
+            self.remote_socket.bind(socket_addr)
+            remote_subscribe_addr = f"tcp://{connect_ip}:{remote_subscribe_port}"
+        else:
+            remote_subscribe_addr = None
+            self.remote_socket = None
+
+        self._is_writer = True
+        self._is_local_reader = False
+        self.local_reader_rank = -1
+        # rank does not matter for remote readers
+        self._is_remote_reader = False
+
+        self.handle = Handle(
+            local_reader_ranks=local_reader_ranks,
+            buffer_handle=self.buffer.handle() if self.buffer is not None else None,
+            local_subscribe_addr=local_subscribe_addr,
+            local_notify_addr=local_notify_addr,
+            remote_subscribe_addr=remote_subscribe_addr,
+            remote_addr_ipv6=remote_addr_ipv6,
+        )
+
+        logger.debug("vLLM message queue communication handle: %s", self.handle)
+
+    def export_handle(self) -> Handle:
+        return self.handle
+
+    @staticmethod
+    def create_from_handle(handle: Handle, rank) -> "MessageQueue":
+        self = MessageQueue.__new__(MessageQueue)
+        self.handle = handle
+        self._is_writer = False
+
+        context = Context()
+
+        if rank in handle.local_reader_ranks:
+            assert handle.buffer_handle is not None
+            self.buffer = ShmRingBuffer(*handle.buffer_handle)
+            self.current_idx = 0
+            self.local_reader_rank = handle.local_reader_ranks.index(rank)
+            self._is_local_reader = True
+            self._is_remote_reader = False
+
+            self.local_socket = context.socket(SUB)
+            self.local_socket.setsockopt_string(SUBSCRIBE, "")
+            socket_addr = handle.local_subscribe_addr
+            logger.debug("Connecting to %s", socket_addr)
+            self.local_socket.connect(socket_addr)
+
+            self.remote_socket = None
+            assert isinstance(handle.local_notify_addr, str)
+            self._spin_condition = SpinCondition(
+                is_reader=True, context=context, notify_address=handle.local_notify_addr
+            )
+        else:
+            self.buffer = None  # type: ignore
+            self.current_idx = -1
+            self.local_reader_rank = -1
+            self._is_local_reader = False
+            self._is_remote_reader = True
+
+            self.local_socket = None
+
+            self.remote_socket = context.socket(SUB)
+            self.remote_socket.setsockopt_string(SUBSCRIBE, "")
+            if handle.remote_addr_ipv6:
+                self.remote_socket.setsockopt(IPV6, 1)
+            socket_addr = handle.remote_subscribe_addr
+            logger.debug("Connecting to %s", socket_addr)
+            self.remote_socket.connect(socket_addr)
+            self._spin_condition = None  # type: ignore
+
+        self.shutting_down = False
+        return self
+
+    def wait_until_ready(self):
+        """This is a collective operation. All processes (including the
+        readers and the writer) should call this function.
+        """
+        if self._is_writer:
+            # wait for all readers to connect
+
+            # local readers
+            for i in range(self.n_local_reader):
+                # wait for subscription messages from all local readers
+                self.local_socket.recv()
+            if self.n_local_reader > 0:
+                # send a message to all local readers
+                # to make sure the publish channel is working
+                self.local_socket.send(b"READY")
+
+            # remote readers
+            for i in range(self.n_remote_reader):
+                # wait for subscription messages from all remote readers
+                self.remote_socket.recv()
+            if self.n_remote_reader > 0:
+                # send a message to all remote readers
+                # to make sure the publish channel is working
+                self.remote_socket.send(b"READY")
+        elif self._is_local_reader:
+            # wait for the writer to send a message
+            recv = self.local_socket.recv()
+            assert recv == b"READY"
+        elif self._is_remote_reader:
+            # wait for the writer to send a message
+            recv = self.remote_socket.recv()
+            assert recv == b"READY"
+
+    def shutdown(self):
+        """If this is an idle reader, wakes it up so it can clean up and shut
+        down"""
+        self.shutting_down = True
+        if self._spin_condition is not None:
+            self._spin_condition.cancel()
+
+    @contextmanager
+    def acquire_write(self, timeout: float | None = None):
+        assert self._is_writer, "Only writers can acquire write"
+        start_time = time.monotonic()
+        n_warning = 1
+        while True:
+            with self.buffer.get_metadata(self.current_idx) as metadata_buffer:
+
+                def check():
+                    memory_fence()
+                    read_count = sum(metadata_buffer[1:])
+                    written_flag = metadata_buffer[0]
+                    return not (written_flag and read_count != self.buffer.n_reader)
+
+                if SPINLOOP_EXT_ENABLED and not check():
+                    spinloop(metadata_buffer, check, timeout=SPINLOOP_TIMEOUT_SECONDS)
+
+                if not check():
+                    # this block is written and not read by all readers
+                    # for writers, `self.current_idx` is the next block to write
+                    # if this block is not ready to write,
+                    # we need to wait until it is read by all readers
+
+                    # Release the processor to other threads
+                    sched_yield()
+
+                    # if we time out, raise an exception
+                    elapsed = time.monotonic() - start_time
+                    if timeout is not None and elapsed > timeout:
+                        raise TimeoutError
+
+                    # if we wait for a long time, log a message
+                    if elapsed > VLLM_RINGBUFFER_WARNING_INTERVAL * n_warning:
+                        logger.info(
+                            LONG_WAIT_TIME_LOG_MSG, VLLM_RINGBUFFER_WARNING_INTERVAL
+                        )
+                        n_warning += 1
+
+                    continue
+                # found a block that is either
+                # (1) not written
+                # (2) read by all readers
+
+                # mark the block as not written
+                metadata_buffer[0] = 0
+                # let caller write to the buffer
+                with self.buffer.get_data(self.current_idx) as buf:
+                    yield buf
+
+                # caller has written to the buffer
+                # NOTE: order is important here
+                # first set the read flags to 0
+                # then set the written flag to 1
+                # otherwise, the readers may think they already read the block
+                for i in range(1, self.buffer.n_reader + 1):
+                    # set read flag to 0, meaning it is not read yet
+                    metadata_buffer[i] = 0
+                # Memory fence here ensures the order of the buffer and flag
+                # writes. This guarantees that when `metadata_buffer[0] = 1` is
+                # visible to readers, `buf` can be completely ready. Without
+                # this, some CPU architectures with weak ordering may incur
+                # memory inconsistency.
+                memory_fence()
+                # mark the block as written
+                metadata_buffer[0] = 1
+                # Memory fence ensures the write is visible to readers on other cores
+                # before we proceed. Without this, readers may spin indefinitely
+                # waiting for a write that's stuck in our CPU's store buffer.
+                memory_fence()
+                self.current_idx = (self.current_idx + 1) % self.buffer.max_chunks
+                break
+
+    class ReadTimeoutWithWarnings:
+        def __init__(self, timeout: float | None, should_warn: bool) -> None:
+            self.started = time.monotonic()
+            self.deadline = sys.maxsize if timeout is None else self.started + timeout
+
+            # if should_warn, we need to wake up periodically to log
+            self.warning_wait_time_ms: int | None = (
+                VLLM_RINGBUFFER_WARNING_INTERVAL * 1000 if should_warn else None
+            )
+
+            self._should_warn = should_warn
+            self.n_warning = 1
+            self.timeout = timeout
+
+        def timeout_ms(self) -> int | None:
+            """Returns a timeout that is:
+            - min(time to deadline, time to next warning) if we're logging warnings
+            - time to deadline, if we're not logging warnings
+            - None if the timeout is None and we're not logging warnings
+            - raise TimeoutError if we are past the deadline
+            """
+            warning_wait_time = self.warning_wait_time_ms
+            if self.timeout is None:
+                return warning_wait_time
+
+            time_left_ms = int((self.deadline - time.monotonic()) * 1000)
+            if time_left_ms <= 0:
+                raise TimeoutError
+
+            if warning_wait_time and warning_wait_time < time_left_ms:
+                return warning_wait_time
+
+            return time_left_ms
+
+        def should_warn(self) -> bool:
+            """Returns true if it's time to log a warning for a timeout that is not
+            indefinite"""
+            if self._should_warn:
+                elapsed = time.monotonic() - self.started
+                if elapsed >= VLLM_RINGBUFFER_WARNING_INTERVAL * self.n_warning:
+                    self.n_warning += 1
+                    return True
+            return False
+
+    @contextmanager
+    def acquire_read(
+        self,
+        timeout: float | None = None,
+        indefinite: bool = False,
+    ):
+        assert self._is_local_reader, "Only readers can acquire read"
+        read_timeout = self.ReadTimeoutWithWarnings(
+            timeout=timeout, should_warn=not indefinite
+        )
+        with self.buffer.get_metadata(self.current_idx) as metadata_buffer:
+            while True:
+
+                def check():
+                    memory_fence()
+                    read_flag = metadata_buffer[self.local_reader_rank + 1]
+                    written_flag = metadata_buffer[0]
+                    return not (not written_flag or read_flag)
+
+                if SPINLOOP_EXT_ENABLED and not check():
+                    spinloop(
+                        metadata_buffer[0 : self.local_reader_rank + 1],
+                        check,
+                        timeout=SPINLOOP_TIMEOUT_SECONDS,
+                    )
+
+                if not check():
+                    # this block is either
+                    # (1) not written
+                    # (2) already read by this reader
+
+                    # for readers, `self.current_idx` is the next block to read
+                    # if this block is not ready,
+                    # we need to wait until it is written
+                    self._spin_condition.wait(timeout_ms=read_timeout.timeout_ms())
+
+                    if self.shutting_down:
+                        raise RuntimeError("cancelled")
+
+                    # if we wait for a long time, log a message
+                    if read_timeout.should_warn():
+                        logger.info(
+                            LONG_WAIT_TIME_LOG_MSG, VLLM_RINGBUFFER_WARNING_INTERVAL
+                        )
+
+                    continue
+                # found a block that is not read by this reader
+                # let caller read from the buffer
+                with self.buffer.get_data(self.current_idx) as buf:
+                    yield buf
+
+                # caller has read from the buffer
+                # set the read flag
+                metadata_buffer[self.local_reader_rank + 1] = 1
+                # Memory fence ensures the read flag is visible to the writer.
+                # Without this, writer may not see our read completion and
+                # could wait indefinitely for all readers to finish.
+                memory_fence()
+                self.current_idx = (self.current_idx + 1) % self.buffer.max_chunks
+
+                self._spin_condition.record_read()
+                break
+
+    def enqueue(self, obj, timeout: float | None = None):
+        """Write to message queue with optional timeout (in seconds)"""
+        assert self._is_writer, "Only writers can enqueue"
+        all_buffers: list[SizedBuffer] = [b""]
+        total_bytes = 6  # 2 bytes for oob buffer count, 4 for main buffer size
+
+        def oob_callback(buf: PickleBuffer) -> bool:
+            raw_buf = buf.raw()
+            if len(raw_buf) < 1024 * 1024:
+                # In-line buffers smaller than 1MiB.
+                return True
+            all_buffers.append(raw_buf)
+            nonlocal total_bytes
+            total_bytes += len(raw_buf) + 4
+            return False
+
+        all_buffers[0] = pickle.dumps(
+            obj, protocol=pickle.HIGHEST_PROTOCOL, buffer_callback=oob_callback
+        )
+        if self.n_local_reader > 0:
+            if total_bytes + len(all_buffers[0]) >= self.buffer.max_chunk_bytes:
+                with self.acquire_write(timeout) as buf:
+                    buf[0] = 1  # overflow
+                self.local_socket.send_multipart(all_buffers, copy=False)
+            else:
+                # Byte 0: 0
+                # Bytes 1-2: Count of buffers
+                # Then each buffer follows, preceded by 4 bytes containing its length:
+                # [4 byte int L][L bytes of buffer content] ...
+                with self.acquire_write(timeout) as buf:
+                    buf[0] = 0  # not overflow
+                    offset = 3
+                    buf[1:offset] = to_bytes_big(len(all_buffers), 2)  # oob buf count
+                    for buffer in all_buffers:
+                        buf_len = len(buffer)
+                        # prepend each buffer with 4 bytes containing its size.
+                        buf_offset = offset + 4
+                        buf[offset:buf_offset] = to_bytes_big(buf_len, 4)
+                        buf[buf_offset : (offset := buf_offset + buf_len)] = buffer
+
+            self._spin_condition.notify()
+
+        if self.n_remote_reader > 0:
+            self.remote_socket.send_multipart(all_buffers, copy=False)
+
+    def dequeue(
+        self,
+        timeout: float | None = None,
+        indefinite: bool = False,
+    ):
+        """Read from message queue with optional timeout (in seconds)"""
+        if self._is_local_reader:
+            with self.acquire_read(timeout, indefinite) as buf:
+                overflow = buf[0] == 1
+                if not overflow:
+                    offset = 3
+                    buf_count = from_bytes_big(buf[1:offset])
+                    all_buffers = []
+                    for i in range(buf_count):
+                        buf_offset = offset + 4
+                        buf_len = from_bytes_big(buf[offset:buf_offset])
+                        offset = buf_offset + buf_len
+                        all_buffers.append(buf[buf_offset:offset])
+                    obj = pickle.loads(all_buffers[0], buffers=all_buffers[1:])
+            if overflow:
+                obj = MessageQueue.recv(self.local_socket, timeout)
+        elif self._is_remote_reader:
+            obj = MessageQueue.recv(self.remote_socket, timeout)
+        else:
+            raise RuntimeError("Only readers can dequeue")
+        return obj
+
+    @staticmethod
+    def recv(socket: zmq.Socket, timeout: float | None) -> Any:
+        # Ensure non-negative timeout passed to zmq poll.
+        timeout_ms = None if timeout is None else max(0, int(timeout * 1000))
+        if not socket.poll(timeout=timeout_ms):
+            raise TimeoutError
+        recv, *recv_oob = socket.recv_multipart(copy=False)
+        return pickle.loads(recv, buffers=recv_oob)
+
+    def broadcast_object(self, obj=None):
+        if self._is_writer:
+            self.enqueue(obj)
+            return obj
+        return self.dequeue()
+
+    @staticmethod
+    def create_from_process_group_single_reader(
+        pg: ProcessGroup,
+        max_chunk_bytes,
+        max_chunks,
+        reader_rank: int = 0,
+        blocking: bool = False,
+    ) -> tuple["MessageQueue", list[Handle]]:
+        """
+        Creates a MessageQueue for a process group with a single reader.
+
+        This method is designed for scenarios where only one process (the reader)
+        will consume messages, and all other processes are writers. It sets up
+        the shared memory buffer and communication handles accordingly, and
+        gathers the handles from all processes to the reader.
+
+        Args:
+            pg (ProcessGroup): The torch distributed process group.
+            max_chunk_bytes (int): Maximum size in bytes for each chunk in the buffer.
+            max_chunks (int): Maximum number of chunks in the buffer.
+            reader_rank (int, optional): The global rank that will act as the reader.
+                Defaults to 0.
+            blocking (bool, optional): If True, blocks until all processes are ready.
+                Defaults to False.
+
+        Returns:
+            tuple[MessageQueue, list[Handle]]:
+            The MessageQueue instance for the calling process,
+            and a list of handles (only non-empty for the reader process).
+        """
+        local_size = current_platform.device_count()
+        rank = dist.get_rank()
+        same_node = rank // local_size == reader_rank // local_size
+        buffer_io = MessageQueue(
+            n_reader=1,
+            n_local_reader=1 if same_node else 0,
+            max_chunk_bytes=max_chunk_bytes,
+            max_chunks=max_chunks,
+        )
+        handle = buffer_io.export_handle()
+        handles = [None] * dist.get_world_size(pg) if rank == reader_rank else None
+        dist.gather_object(handle, handles, dst=reader_rank, group=pg)
+        if blocking:
+            buffer_io.wait_until_ready()
+        return buffer_io, cast(list[Handle], handles or [])
+
+    @staticmethod
+    def create_from_process_group(
+        pg: ProcessGroup | StatelessProcessGroup,
+        max_chunk_bytes,
+        max_chunks,
+        writer_rank: int = 0,
+        external_writer_handle=None,
+        blocking: bool = True,
+    ) -> "MessageQueue":
+        """
+        Creates a MessageQueue for a distributed process group with one writer and
+        multiple readers.
+
+        This method is designed for scenarios where one process (the writer) sends
+        messages, and all other processes (the readers) receive messages. It sets up
+        the shared memory buffer and socket communication handles accordingly, and
+        broadcasts the handle from the writer to all readers.
+
+        Args:
+            pg (ProcessGroup | StatelessProcessGroup): The torch distributed process
+                group.
+            max_chunk_bytes (int): Maximum size in bytes for each chunk in the buffer.
+            max_chunks (int): Maximum number of chunks in the buffer.
+            writer_rank (int, optional): The global rank that will act as the writer.
+                Defaults to 0.
+            external_writer_handle (Handle, optional): Used when there is a handle
+                from an external Message Queue. If provided, use this handle to init
+                PG writer message queue instead of creating a new one. Defaults to None.
+            blocking (bool, optional): If True, blocks until all processes are ready.
+                Defaults to True.
+
+        Returns:
+            MessageQueue: The MessageQueue instance for the calling process.
+
+        """
+        if isinstance(pg, ProcessGroup):
+            group_rank = dist.get_rank(pg)
+            group_world_size = dist.get_world_size(pg)
+            global_ranks = dist.get_process_group_ranks(pg)
+        else:
+            group_rank = pg.rank
+            group_world_size = pg.world_size
+            global_ranks = list(range(pg.world_size))
+        from vllm.distributed.parallel_state import in_the_same_node_as
+
+        status = in_the_same_node_as(pg, source_rank=writer_rank)
+        if group_rank == writer_rank:
+            if external_writer_handle is not None:
+                buffer_io = MessageQueue.create_from_handle(
+                    external_writer_handle, group_rank
+                )
+            else:
+                same_node_ranks = [i for i, s in enumerate(status) if s]
+                n_reader = group_world_size - 1
+                n_local_reader = len(same_node_ranks) - 1
+                local_reader_ranks = [i for i in same_node_ranks if i != writer_rank]
+                buffer_io = MessageQueue(
+                    n_reader=n_reader,
+                    n_local_reader=n_local_reader,
+                    local_reader_ranks=local_reader_ranks,
+                    max_chunk_bytes=max_chunk_bytes,
+                    max_chunks=max_chunks,
+                )
+            handle = buffer_io.export_handle()
+            if isinstance(pg, ProcessGroup):
+                dist.broadcast_object_list(
+                    [handle], src=global_ranks[writer_rank], group=pg
+                )
+            else:
+                pg.broadcast_obj(handle, writer_rank)
+        else:
+            if isinstance(pg, ProcessGroup):
+                recv = [None]
+                dist.broadcast_object_list(
+                    recv, src=global_ranks[writer_rank], group=pg
+                )
+                handle = recv[0]  # type: ignore
+            else:
+                handle = pg.broadcast_obj(None, writer_rank)
+            buffer_io = MessageQueue.create_from_handle(handle, group_rank)
+        if blocking:
+            buffer_io.wait_until_ready()
+        return buffer_io

super_v3_reasoning_parser.py

@@ -1,28 +0,0 @@
-from vllm.reasoning.abs_reasoning_parsers import ReasoningParserManager
-from vllm.reasoning.deepseek_r1_reasoning_parser import DeepSeekR1ReasoningParser
-
-
-@ReasoningParserManager.register_module("super_v3")
-class SuperV3ReasoningParser(DeepSeekR1ReasoningParser):
-    def extract_reasoning(self, model_output, request):
-        reasoning_content, final_content = super().extract_reasoning(
-            model_output, request
-        )
-        if (
-            hasattr(request, "chat_template_kwargs")
-            and request.chat_template_kwargs
-            and (
-                request.chat_template_kwargs.get("enable_thinking") is False
-                or request.chat_template_kwargs.get("force_nonempty_content") is True
-            )
-            and final_content is None
-        ):
-            """
-            The original `deepseek_r1` reasoning parser this inherits from will automatically put everything in the reasoning content when it cannot parse out reasoning. This was fine for the DeepSeek R1 model that was not intended to be used without reasoning.
-            1. Since the Nemotron 3 Nano and Super both have thinking off modes modulated by "enable_thinking=false" in the chat template kwargs, this change instead which will properly place the content in cases where there is no thinking enabled via config.
-            2. There are rare cases where the model will output only reasoning without an end-think token `</think>` (e.g. reasoning exceeds max length), which results in empty content returned. End users may want to unilaterally avoid such cases and always have a content response even if the model does not finish its reasoning.
-            """
-            # Put all nonempty content into the content, rather than return content
-            reasoning_content, final_content = None, reasoning_content
-
-        return reasoning_content, final_content

utils.py

@@ -1,438 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import ast
-import json
-from json import JSONDecodeError, JSONDecoder
-from typing import Any, TypeAlias
-
-import partial_json_parser
-from openai.types.responses import (
-    FunctionTool,
-    ToolChoiceFunction,
-)
-from openai.types.responses.tool import Tool as ResponsesTool
-from partial_json_parser.core.options import Allow
-
-from vllm.entrypoints.openai.chat_completion.protocol import (
-    ChatCompletionNamedToolChoiceParam,
-    ChatCompletionToolsParam,
-)
-from vllm.entrypoints.openai.engine.protocol import (
-    DeltaFunctionCall,
-    DeltaToolCall,
-    FunctionCall,
-    ToolCall,
-)
-from vllm.logger import init_logger
-
-Tool: TypeAlias = ChatCompletionToolsParam | ResponsesTool
-
-logger = init_logger(__name__)
-
-
-def partial_tag_overlap(text: str, tag: str) -> int:
-    """Length of the longest prefix of *tag* that matches a suffix of *text*.
-
-    E.g. text ending in ``"<tool_"`` returns 6 when tag is ``"<tool_call>"``.
-    Returns 0 when there is no overlap.
-    """
-    max_check = min(len(tag) - 1, len(text))
-    for k in range(max_check, 0, -1):
-        if text.endswith(tag[:k]):
-            return k
-    return 0
-
-
-def find_common_prefix(s1: str, s2: str) -> str:
-    """
-    Finds a common prefix that is shared between two strings, if there is one.
-    Order of arguments is NOT important.
-
-    This function is provided as a UTILITY for extracting information from JSON
-    generated by partial_json_parser, to help in ensuring that the right tokens
-    are returned in streaming, so that close-quotes, close-brackets and
-    close-braces are not returned prematurely.
-
-    e.g. find_common_prefix('{"fruit": "ap"}', '{"fruit": "apple"}') ->
-    '{"fruit": "ap'
-    """
-    prefix = ""
-    min_length = min(len(s1), len(s2))
-    for i in range(0, min_length):
-        if s1[i] == s2[i]:
-            prefix += s1[i]
-        else:
-            break
-    return prefix
-
-
-def find_common_suffix(s1: str, s2: str) -> str:
-    """
-    Finds a common suffix shared between two strings, if there is one. Order of
-    arguments is NOT important.
-    Stops when the suffix ends OR it hits an alphanumeric character
-
-    e.g. find_common_suffix('{"fruit": "ap"}', '{"fruit": "apple"}') -> '"}'
-    """
-    suffix = ""
-    min_length = min(len(s1), len(s2))
-    for i in range(1, min_length + 1):
-        if s1[-i] == s2[-i] and not s1[-i].isalnum():
-            suffix = s1[-i] + suffix
-        else:
-            break
-    return suffix
-
-
-def extract_intermediate_diff(curr: str, old: str) -> str:
-    """
-    Given two strings, extract the difference in the middle between two strings
-    that are known to have a common prefix and/or suffix.
-
-    This function is provided as a UTILITY for extracting information from JSON
-    generated by partial_json_parser, to help in ensuring that the right tokens
-    are returned in streaming, so that close-quotes, close-brackets and
-    close-braces are not returned prematurely. The order of arguments IS
-    important - the new version of the partially-parsed JSON must be the first
-    argument, and the secnod argument must be from the previous generation.
-
-    What it returns, is tokens that should be streamed to the client.
-
-    e.g. extract_intermediate_diff('{"fruit": "apple"}', '{"fruit": "ap"}')
-        -> 'ple'
-
-    """
-    suffix = find_common_suffix(curr, old)
-
-    old = old[::-1].replace(suffix[::-1], "", 1)[::-1]
-    prefix = find_common_prefix(curr, old)
-    diff = curr
-    if len(suffix):
-        diff = diff[::-1].replace(suffix[::-1], "", 1)[::-1]
-
-    if len(prefix):
-        # replace the prefix only once in case it's mirrored
-        diff = diff.replace(prefix, "", 1)
-
-    return diff
-
-
-# partial_json_parser doesn't support extra data and
-# JSONDecoder.raw_decode doesn't support partial JSON
-def partial_json_loads(input_str: str, flags: Allow) -> tuple[Any, int]:
-    try:
-        return (partial_json_parser.loads(input_str, flags), len(input_str))
-    except JSONDecodeError as e:
-        if "Extra data" in e.msg:
-            dec = JSONDecoder()
-            return dec.raw_decode(input_str)
-        raise
-
-
-def is_complete_json(input_str: str) -> bool:
-    try:
-        json.loads(input_str)
-        return True
-    except JSONDecodeError:
-        return False
-
-
-def consume_space(i: int, s: str) -> int:
-    while i < len(s) and s[i].isspace():
-        i += 1
-    return i
-
-
-def _extract_tool_info(
-    tool: Tool,
-) -> tuple[str, dict[str, Any] | None]:
-    if isinstance(tool, FunctionTool):
-        return tool.name, tool.parameters
-    elif isinstance(tool, ChatCompletionToolsParam):
-        return tool.function.name, tool.function.parameters
-    else:
-        raise TypeError(f"Unsupported tool type: {type(tool)}")
-
-
-def _get_tool_schema_from_tool(tool: Tool) -> dict:
-    name, params = _extract_tool_info(tool)
-    params = params if params else {"type": "object", "properties": {}}
-    return {
-        "properties": {
-            "name": {"type": "string", "enum": [name]},
-            "parameters": params,
-        },
-        "required": ["name", "parameters"],
-    }
-
-
-def _get_tool_schema_defs(
-    tools: list[Tool],
-) -> dict:
-    all_defs: dict[str, dict[str, Any]] = {}
-    for tool in tools:
-        _, params = _extract_tool_info(tool)
-        if params is None:
-            continue
-        defs = params.pop("$defs", {})
-        for def_name, def_schema in defs.items():
-            if def_name in all_defs and all_defs[def_name] != def_schema:
-                raise ValueError(
-                    f"Tool definition '{def_name}' has multiple schemas, "
-                    "which is not supported."
-                )
-            all_defs[def_name] = def_schema
-    return all_defs
-
-
-def _get_json_schema_from_tools(
-    tools: list[Tool],
-) -> dict:
-    json_schema = {
-        "type": "array",
-        "minItems": 1,
-        "items": {
-            "type": "object",
-            "anyOf": [_get_tool_schema_from_tool(tool) for tool in tools],
-        },
-    }
-    json_schema_defs = _get_tool_schema_defs(tools)
-    if json_schema_defs:
-        json_schema["$defs"] = json_schema_defs
-    return json_schema
-
-
-def get_json_schema_from_tools(
-    tool_choice: str | ToolChoiceFunction | ChatCompletionNamedToolChoiceParam,
-    tools: list[Tool] | None,
-) -> str | dict | None:
-    # tool_choice: "none"
-    if tool_choice in ("none", None) or tools is None:
-        return None
-    # tool_choice: Forced Function (Responses)
-    if (not isinstance(tool_choice, str)) and isinstance(
-        tool_choice, ToolChoiceFunction
-    ):
-        tool_name = tool_choice.name
-        tool_map = {tool.name: tool for tool in tools if isinstance(tool, FunctionTool)}
-        if tool_name not in tool_map:
-            raise ValueError(f"Tool '{tool_name}' has not been passed in `tools`.")
-        return tool_map[tool_name].parameters
-    # tool_choice: Forced Function (ChatCompletion)
-    if (not isinstance(tool_choice, str)) and isinstance(
-        tool_choice, ChatCompletionNamedToolChoiceParam
-    ):
-        tool_name = tool_choice.function.name
-        tool_map = {
-            tool.function.name: tool
-            for tool in tools
-            if isinstance(tool, ChatCompletionToolsParam)
-        }
-        if tool_name not in tool_map:
-            raise ValueError(f"Tool '{tool_name}' has not been passed in `tools`.")
-        return tool_map[tool_name].function.parameters
-    # tool_choice: "required"
-    if tool_choice == "required":
-        return _get_json_schema_from_tools(tools)
-    # tool_choice: "auto"
-    return None
-
-
-# ---------------------------------------------------------------------------
-# Shared utilities for pythonic-style tool call parsers
-# (PythonicToolParser, Llama4PythonicToolParser, Olmo3PythonicToolParser)
-# ---------------------------------------------------------------------------
-
-
-class UnexpectedAstError(Exception):
-    """Raised when the AST structure does not match the expected
-    pythonic tool call format."""
-
-    pass
-
-
-_JSON_NAME_LITERALS = {
-    "null": None,
-    "true": True,
-    "false": False,
-}
-
-
-def get_parameter_value(val: ast.expr) -> Any:
-    """Extract a Python literal value from an AST expression node.
-
-    Handles constants, dicts, lists, and JSON-style name literals
-    (null, true, false) that some models produce instead of Python
-    literals (None, True, False).
-
-    Raises:
-        UnexpectedAstError: If the AST node is not a supported literal type.
-    """
-    if isinstance(val, ast.Constant):
-        return val.value
-    elif isinstance(val, ast.Dict):
-        if not all(isinstance(k, ast.Constant) for k in val.keys):
-            logger.warning(
-                "Dict argument keys are not all literals: %s",
-                ast.dump(val),
-            )
-            raise UnexpectedAstError("Dict tool call arguments must have literal keys")
-        return {
-            k.value: get_parameter_value(v)  # type: ignore
-            for k, v in zip(val.keys, val.values)
-        }
-    elif isinstance(val, ast.List):
-        return [get_parameter_value(v) for v in val.elts]
-    elif isinstance(val, ast.Name) and val.id in _JSON_NAME_LITERALS:
-        return _JSON_NAME_LITERALS[val.id]
-    else:
-        logger.warning(
-            "Unsupported AST node type in tool call arguments: %s",
-            ast.dump(val),
-        )
-        raise UnexpectedAstError("Tool call arguments must be literals")
-
-
-def handle_single_tool(call: ast.Call) -> ToolCall:
-    """Convert a single AST function call node into a ToolCall object.
-
-    Raises:
-        UnexpectedAstError: If the call node does not have a simple
-            function name (e.g. it's an attribute access or subscript).
-    """
-    if not isinstance(call.func, ast.Name):
-        logger.warning(
-            "Tool call has non-simple function name: %s",
-            ast.dump(call.func),
-        )
-        raise UnexpectedAstError("Invalid tool call name")
-    function_name = call.func.id
-    arguments = {}
-    for keyword in call.keywords:
-        arguments[keyword.arg] = get_parameter_value(keyword.value)
-    return ToolCall(
-        type="function",
-        function=FunctionCall(
-            name=function_name,
-            arguments=json.dumps(arguments, ensure_ascii=False),
-        ),
-    )
-
-
-def make_valid_python(text: str) -> tuple[str, str] | None:
-    """Attempt to close all open brackets/quotes to make partial Python valid.
-
-    Used during streaming to parse incomplete tool call expressions by
-    appending the necessary closing characters.
-
-    Returns:
-        A tuple of (completed_text, added_suffix) if the text can be
-        made valid, or None if the text is too incomplete to complete
-        meaningfully (e.g. mid-parameter-name or mid-dict-key).
-
-    Raises:
-        UnexpectedAstError: If mismatched brackets or parentheses
-            are detected.
-    """
-    bracket_stack: list[str] = []
-    for index, char in enumerate(text):
-        if char in {"[", "(", "{"}:
-            bracket_stack.append(char)
-        elif char == "]":
-            if not bracket_stack or bracket_stack.pop() != "[":
-                raise UnexpectedAstError("Mismatched square brackets")
-        elif char == ")":
-            if not bracket_stack or bracket_stack.pop() != "(":
-                raise UnexpectedAstError("Mismatched parentheses")
-        elif char == "}":
-            if not bracket_stack or bracket_stack.pop() != "{":
-                raise UnexpectedAstError("Mismatched curly braces")
-        elif char in {"'", '"'}:
-            if bracket_stack and bracket_stack[-1] == char:
-                if index > 0 and text[index - 1] == "\\":
-                    pass
-                else:
-                    bracket_stack.pop()
-            elif bracket_stack and bracket_stack[-1] in {"'", '"'}:
-                pass
-            else:
-                bracket_stack.append(char)
-
-    text = text.rstrip()
-    if text.endswith("=") or text.endswith(":"):
-        return None
-    if bracket_stack and bracket_stack[-1] == "{":
-        trailing_dict_text = text[: text.rfind("{")]
-        num_keys = trailing_dict_text.count(":")
-        num_values = trailing_dict_text.count(",")
-        if num_keys <= num_values:
-            return None
-    if bracket_stack and bracket_stack[-1] == "(":
-        trailing_params_text = text[: text.rfind("(")]
-        num_full_param_names = trailing_params_text.count("=")
-        num_full_param_values = trailing_params_text.count(",")
-        if num_full_param_names <= num_full_param_values:
-            return None
-    if text.endswith(","):
-        text = text[:-1]
-    if (
-        bracket_stack
-        and bracket_stack[-1] == "["
-        and not text.endswith("[")
-        and not text.endswith(")")
-    ):
-        return None
-
-    _CLOSING = {"[": "]", "(": ")", "{": "}", "'": "'", '"': '"'}
-    added_text = ""
-    for char in reversed(bracket_stack):
-        added_text += _CLOSING[char]
-
-    return text + added_text, added_text
-
-
-def compute_tool_delta(
-    previously_sent_args: str,
-    new_call: ToolCall,
-    index: int,
-    withheld_suffix: str,
-) -> DeltaToolCall | None:
-    """Compute the incremental delta between previously streamed arguments
-    and the current tool call state.
-
-    Returns:
-        A DeltaToolCall with only the new argument characters, or None
-        if there is no difference from what was previously sent.
-    """
-    new_call_args = new_call.function.arguments
-    if withheld_suffix:
-        if not new_call_args.endswith(withheld_suffix):
-            msg = (
-                f"Tool call arguments '{new_call_args}' do not end with "
-                f"expected withheld suffix '{withheld_suffix}'"
-            )
-            logger.error(msg)
-            raise ValueError(msg)
-        new_call_args = new_call_args[: -len(withheld_suffix)]
-    if not previously_sent_args:
-        return DeltaToolCall(
-            id=new_call.id,
-            type="function",
-            index=index,
-            function=DeltaFunctionCall(
-                name=new_call.function.name,
-                arguments=new_call_args,
-            ),
-        )
-
-    arg_diff = new_call_args[len(previously_sent_args) :]
-    return (
-        DeltaToolCall(
-            id=None,
-            index=index,
-            function=DeltaFunctionCall(arguments=arg_diff),
-        )
-        if arg_diff
-        else None
-    )