diff --git a/Dockerfile b/Dockerfile index 140c3f6..5926830 100644 --- a/Dockerfile +++ b/Dockerfile @@ -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 \ No newline at end of file +#COPY lmcache-config-dsv4.yaml /opt/lmcache-config-dsv4.yaml + diff --git a/deep_gemm.py b/deep_gemm.py new file mode 100644 index 0000000..5eb548f --- /dev/null +++ b/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", +] \ No newline at end of file diff --git a/indexer.py b/indexer.py index a856fe2..b58ed68 100644 --- a/indexer.py +++ b/indexer.py @@ -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) \ No newline at end of file diff --git a/multiproc_executor.py b/multiproc_executor.py index 3be6cce..7c2e4a4 100644 --- a/multiproc_executor.py +++ b/multiproc_executor.py @@ -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: diff --git a/shm_broadcast.py b/shm_broadcast.py index d695aa6..b8e09fa 100644 --- a/shm_broadcast.py +++ b/shm_broadcast.py @@ -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