biondizzle/vllm-with-lmcache
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

other things

Browse Source
This commit is contained in:
biondizzle
2026-07-06 19:26:26 +00:00
parent 57477cf0ec
commit c76ee01d57
5 changed files with 714 additions and 75 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
Dockerfile
View File

@@ -23,11 +23,12 @@ RUN apt-get update && apt-get install -y git \
# 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 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
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
@@ -36,4 +37,5 @@ RUN apt-get update && apt-get install -y git \
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
#COPY lmcache-config-dsv4.yaml /opt/lmcache-config-dsv4.yaml

587
deep_gemm.py Normal file
View File

@@ -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",
]

62
indexer.py
View File

@@ -231,8 +231,6 @@ def get_max_prefill_buffer_size(vllm_config: VllmConfig):
class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
reorder_batch_threshold: int = 1
natively_supported_next_n_fp4: list[int] = [1, 2]
# TODO (matt): integrate kernel with next_n = 4 support
@classmethod
def get_cudagraph_support(
@@ -267,15 +265,21 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
next_n = self.num_speculative_tokens + 1
self.reorder_batch_threshold += self.num_speculative_tokens
# NOTE(zyongye) fp4 indexer cache only natively supports next_n in
# natively_supported_next_n_fp4; for other next_n values we fall back
# to the flattening path. Outside the SM100 datacenter family the FP8
# paged MQA logits kernel has the same [1, 2] constraint (deepgemm
# smxx_fp8_fp4_paged_mqa_logits.hpp:233), so flatten there too.
self.use_flattening = (
self.use_fp4_indexer_cache
or not current_platform.is_device_capability_family(100)
) and next_n not in self.natively_supported_next_n_fp4
# 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
@@ -288,20 +292,14 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
dtype=torch.int32,
device=self.device,
)
if not self.use_flattening and next_n > 1:
# Native MTP: 2D buffer for per-token seq_lens.
self.decode_seq_lens_buffer = torch.zeros(
(scheduler_config.max_num_seqs, next_n),
dtype=torch.int32,
device=self.device,
)
else:
# Flattening or no MTP: 1D buffer for expanded per-token seq_lens.
self.decode_seq_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,
@@ -373,7 +371,8 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
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, next_n) for native MTP.
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:
@@ -454,16 +453,19 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
# (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() == 2
assert self.decode_seq_lens_buffer.dim() == 1
# (B, max_decode_len): token j attends to
# L - max_decode_len + j + 1 KV tokens.
self.decode_seq_lens_buffer[:num_decodes, :max_decode_len] = (
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 = self.decode_seq_lens_buffer[:num_decodes, :max_decode_len]
seq_lens = seq_lens_buffer
return seq_lens, block_table, decode_lens, num_decodes, requires_padding
def build(
@@ -774,4 +776,4 @@ def _build_prefill_chunk_metadata_kernel(
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)
tl.store(token_to_seq_ptr + seq_start + offset, batch_idx, mask=mask)

72
multiproc_executor.py
View File

@@ -15,7 +15,7 @@ from concurrent.futures import Future, InvalidStateError
from contextlib import suppress
from dataclasses import dataclass
from enum import Enum, auto
from functools import cached_property, partial
from functools import partial
from multiprocessing.connection import Connection
from multiprocessing.process import BaseProcess
from multiprocessing.synchronize import Lock as LockType
@@ -60,6 +60,7 @@ from vllm.utils.system_utils import (
)
from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
from vllm.v1.executor.abstract import Executor, FailureCallback
from vllm.v1.executor.vllm_net_devices import set_worker_net_device
from vllm.v1.outputs import AsyncModelRunnerOutput, DraftTokenIds, ModelRunnerOutput
from vllm.v1.worker.worker_base import WorkerWrapperBase
@@ -395,9 +396,7 @@ class MultiprocExecutor(Executor):
return responses[0] if output_rank is not None else responses
future = FutureWrapper(
self.futures_queue,
get_response=get_response,
aggregate=aggregate,
self.futures_queue, get_response=get_response, aggregate=aggregate
)
return future if non_block else future.result()
@@ -421,27 +420,47 @@ class MultiprocExecutor(Executor):
return False
active_procs = lambda: [proc for proc in worker_procs if proc.is_alive()]
initial_count = len(active_procs())
# Give processes time to clean themselves up properly first
logger.debug("Worker Termination: allow workers to gracefully shutdown")
if wait_for_termination(active_procs(), 4):
logger.info(
"[shutdown] Executor: waiting for worker exit count=%d",
initial_count,
)
if wait_for_termination(
active_procs(), timeout=envs.VLLM_WORKER_SHUTDOWN_TIMEOUT_SECONDS
):
logger.info_once("[shutdown] Executor: all workers exited gracefully")
return
# Send SIGTERM if still running
logger.debug("Worker Termination: workers still running sending SIGTERM")
for p in active_procs():
remaining = active_procs()
logger.warning(
"[shutdown] Executor: workers still running after grace period; "
"sending SIGTERM count=%d",
len(remaining),
)
for p in remaining:
p.terminate()
if not wait_for_termination(active_procs(), 4):
# Send SIGKILL if still running
logger.debug(
"Worker Termination: resorting to SIGKILL to take down workers"
remaining = active_procs()
logger.warning(
"[shutdown] Executor: workers still running after SIGTERM; "
"sending SIGKILL count=%d",
len(remaining),
)
for p in active_procs():
for p in remaining:
p.kill()
def shutdown(self):
"""Properly shut down the executor and its workers"""
if not getattr(self, "shutting_down", False):
logger.debug("Triggering shutdown of workers")
worker_count = len(getattr(self, "workers", None) or [])
logger.debug(
"[shutdown] Executor: start worker_count=%d",
worker_count,
)
self.shutting_down = True
# Make sure all the worker processes are terminated first.
@@ -467,16 +486,12 @@ class MultiprocExecutor(Executor):
mq.shutdown()
self.response_mqs = []
logger.debug_once("[shutdown] Executor: complete")
def check_health(self) -> None:
self.collective_rpc("check_health", timeout=10)
return
@cached_property
def max_concurrent_batches(self) -> int:
# PP requires PP-size concurrent batches to fill the pipeline.
pp_size = self.parallel_config.pipeline_parallel_size
return 2 if pp_size <= 1 and self.scheduler_config.async_scheduling else pp_size
def _get_output_rank(self) -> int:
# Only returns ModelRunnerOutput from TP rank=0 and PP rank=-1
# (the first TP worker of the last PP stage).
@@ -811,6 +826,9 @@ class WorkerProc:
signal.signal(signal.SIGTERM, signal_handler)
signal.signal(signal.SIGINT, signal_handler)
# Set net device env vars for the worker if VLLM_GPU_NIC_PCIE_MAPPING is set
set_worker_net_device(kwargs.get("local_rank", 0), kwargs["vllm_config"])
worker = None
ready_writer = kwargs.pop("ready_pipe")
death_pipe = kwargs.pop("death_pipe", None)
@@ -869,7 +887,9 @@ class WorkerProc:
if ready_writer is not None:
logger.exception("WorkerProc failed to start.")
elif shutdown_requested.is_set():
logger.info("WorkerProc shutting down.")
logger.debug_once(
"[shutdown] WorkerProc: exiting after shutdown request"
)
else:
logger.exception("WorkerProc failed.")
@@ -881,7 +901,12 @@ class WorkerProc:
except SystemExit as e:
# SystemExit is raised on SIGTERM or SIGKILL, which usually indicates that
# the graceful shutdown process did not succeed
logger.warning("WorkerProc was terminated")
if shutdown_requested.is_set():
logger.debug_once(
"[shutdown] WorkerProc: terminated by shutdown signal"
)
else:
logger.warning("WorkerProc was terminated")
# SystemExit must never be ignored
raise e
@@ -955,6 +980,9 @@ class WorkerProc:
func = partial(cloudpickle.loads(method), self.worker)
output = func(*args, **kwargs)
if output_rank is None or self.rank == output_rank:
self.handle_output(output)
except Exception as e:
# Notes have been introduced in python 3.11
if hasattr(e, "add_note"):
@@ -964,10 +992,6 @@ class WorkerProc:
# string, only for logging purpose.
if output_rank is None or self.rank == output_rank:
self.handle_output(e)
continue
if output_rank is None or self.rank == output_rank:
self.handle_output(output)
@staticmethod
def setup_proc_title_and_log_prefix(enable_ep: bool) -> None:

58
shm_broadcast.py
View File

@@ -38,6 +38,21 @@ from vllm.utils.network_utils import (
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
@@ -77,9 +92,6 @@ def to_bytes_big(value: int, size: int) -> bytes:
return value.to_bytes(size, byteorder="big")
logger = init_logger(__name__)
LONG_WAIT_TIME_LOG_MSG = (
"No available shared memory broadcast block found "
"in %d seconds. This typically happens "
@@ -540,13 +552,17 @@ class MessageQueue:
n_warning = 1
while True:
with self.buffer.get_metadata(self.current_idx) as metadata_buffer:
# Memory fence ensures we see the latest read flags from readers.
# Without this, we may read stale flags from our CPU cache and
# spin indefinitely even though readers have completed.
memory_fence()
read_count = sum(metadata_buffer[1:])
written_flag = metadata_buffer[0]
if written_flag and read_count != self.buffer.n_reader:
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,
@@ -657,13 +673,21 @@ class MessageQueue:
)
with self.buffer.get_metadata(self.current_idx) as metadata_buffer:
while True:
# Memory fence ensures we see the latest writes from the writer.
# Without this, we may read stale flags from our CPU cache
# and spin indefinitely even though writer has updated them.
memory_fence()
read_flag = metadata_buffer[self.local_reader_rank + 1]
written_flag = metadata_buffer[0]
if not written_flag or read_flag:
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