From c9d922ca26218363adcf7946c9cfbecacfc3155f Mon Sep 17 00:00:00 2001 From: specture724 Date: Thu, 10 Sep 2026 16:22:27 +0800 Subject: [PATCH] feat: overlap whole requests across DBO ubatches Upstream vLLM splits a DBO batch at an even token count, which cuts whichever request straddles that point into both ubatches. AFD's overlap is between whole requests -- one request's dispatch is in flight while the other's attention runs -- so the split has to land on a request boundary, and a batch with no interior boundary (a single prefill) must run whole rather than be divided. The splitter also prefers 16-token-aligned boundaries: a ubatch's per-token tensors are views starting at the split point, and DeepSeek-V4's CuTeDSL compressor kernel rejects anything below 64-byte alignment. It is a preference, not a rule -- uniform decode has boundaries at 1, 2, 3, ... and refusing there would disable DBO for decode entirely. Two bugs that kept DBO from working at all on this connector. The batch- execution override only consulted the rank-local decision when data_parallel_size == 1, but the cross-DP agreement is *also* skipped for connectors with no control plane, where the parent's False is a placeholder rather than a decision -- so `--enable-dbo` was accepted and then silently ignored for the whole run. And the last split point came off cu_num_tokens as a numpy.int32, which Triton refuses to specialize, killing DeepSeek-V4 in _build_c128a_topk_metadata_kernel on the last ubatch. The dispatch stage is now the DBO ubatch id, read from the thread the way vLLM tracks it; forward_context.ubatch_idx does not exist, so both halves looked like stage 0 and claimed the same window slot. Co-Authored-By: Claude Opus 5 Signed-off-by: specture724 --- afd_plugin/__init__.py | 37 ++-- afd_plugin/compat/patches/ubatch_positions.py | 62 +++++++ afd_plugin/compat/patches/ubatch_split.py | 169 ++++++++++++++++++ .../model_executor/models/deepseek_v2.py | 16 +- .../model_executor/models/deepseek_v4.py | 18 +- .../v1/worker/attention_model_runner.py | 13 +- afd_plugin/v1/worker/ubatch_wrapper.py | 14 ++ .../deepseek_v4_flash/2a2f_async.sh | 19 ++ .../compat/patches/test_ubatch_positions.py | 53 ++++++ .../unit/compat/patches/test_ubatch_split.py | 160 +++++++++++++++++ .../models/test_deepseek_v2_proxy.py | 14 +- .../models/test_deepseek_v4_proxy.py | 28 ++- .../v1/worker/test_attention_model_runner.py | 60 ++++++- tests/unit/v1/worker/test_ubatch_wrapper.py | 58 ++++++ 14 files changed, 662 insertions(+), 59 deletions(-) create mode 100644 afd_plugin/compat/patches/ubatch_positions.py create mode 100644 afd_plugin/compat/patches/ubatch_split.py create mode 100644 tests/unit/compat/patches/test_ubatch_positions.py create mode 100644 tests/unit/compat/patches/test_ubatch_split.py create mode 100644 tests/unit/v1/worker/test_ubatch_wrapper.py diff --git a/afd_plugin/__init__.py b/afd_plugin/__init__.py index 2766818a..6d48bc0a 100644 --- a/afd_plugin/__init__.py +++ b/afd_plugin/__init__.py @@ -8,6 +8,7 @@ import logging import multiprocessing import os +from importlib import import_module from importlib.metadata import PackageNotFoundError, version from pathlib import Path from types import MappingProxyType @@ -170,18 +171,30 @@ def register_afd() -> None: exc_info=True, ) - try: - import afd_plugin.compat.patches.async_dp_engine # noqa: F401 - import afd_plugin.compat.patches.async_dp_forward_context # noqa: F401 - import afd_plugin.compat.patches.config_validation # noqa: F401 - import afd_plugin.compat.patches.dp_coordinator_timeout # noqa: F401 - import afd_plugin.compat.patches.engine_core # noqa: F401 - import afd_plugin.compat.patches.ffn_local_moe_prepare # noqa: F401 - except Exception: - _logger.debug( - "AFD plugin: compatibility patches could not be applied", - exc_info=True, - ) + # One import per patch, each isolated. A single try block around the whole + # list means the first failure silently skips every patch after it: a stale + # module name here once disabled the ubatch positions and split patches, + # which surfaced two layers away as "positions is required for C128A + # metadata build" inside DeepSeek-V4's kernel warmup. Warn rather than + # debug for the same reason -- a patch that did not load is not a detail. + for _patch in ( + "async_dp_engine", + "async_dp_forward_context", + "config_validation", + "dp_coordinator_timeout", + "engine_core", + "ffn_local_moe_prepare", + "ubatch_positions", + "ubatch_split", + ): + try: + import_module(f"afd_plugin.compat.patches.{_patch}") + except Exception: + _logger.warning( + "AFD plugin: compatibility patch %r could not be applied", + _patch, + exc_info=True, + ) from afd_plugin.model_executor.routing_simulator import ( register_afd_balanced_routing_strategy, diff --git a/afd_plugin/compat/patches/ubatch_positions.py b/afd_plugin/compat/patches/ubatch_positions.py new file mode 100644 index 00000000..b1eec07b --- /dev/null +++ b/afd_plugin/compat/patches/ubatch_positions.py @@ -0,0 +1,62 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Propagate ``positions`` through vLLM's ubatch attention-metadata split. + +``CommonAttentionMetadata.positions`` is per-token and optional; upstream's +``split_attn_metadata`` rebuilds each ubatch's metadata without it, so every +ubatch sees ``positions=None``. The DeepSeek-V4 C128A metadata builder asserts +``positions`` is present, which makes any DBO-split forward fail. Re-slice the +source positions onto each rebuilt metadata after the upstream split. +""" + +from __future__ import annotations + +from vllm.v1.worker import ubatch_utils as ubatch_utils_module + +_upstream_split_attn_metadata = ubatch_utils_module.split_attn_metadata + + +# Patch reason: upstream split_attn_metadata drops CommonAttentionMetadata +# .positions, and the DeepSeek-V4 C128A metadata builder asserts on it, so +# every DBO-split forward of a V4 model dies in the attention metadata build. +# Patch functionality: after the upstream split, re-slice the source +# metadata's per-token positions onto each ubatch's metadata by that +# ubatch's token slice; None stays None. +# Signature: matches upstream; no added parameters. +# Upstream: vLLM v0.26.0, vllm/v1/worker/ubatch_utils.py +def split_attn_metadata(ubatch_slices, common_attn_metadata): + results = _upstream_split_attn_metadata(ubatch_slices, common_attn_metadata) + positions = common_attn_metadata.positions + if positions is not None: + for ubatch_slice, ubatch_metadata in zip(ubatch_slices, results, strict=False): + ubatch_metadata.positions = positions[ubatch_slice.token_slice] + return results + + +split_attn_metadata.__afd_positions_propagated = True # type: ignore[attr-defined] + + +def apply_positions_propagation() -> None: + """Install the position-propagating splitter into vLLM's module namespaces. + + The plugin loads before vLLM's worker stack is importable, so importing + ``gpu_model_runner`` here would fail on a partial import; the source + module is always patched (the runner binds the name when it is first + imported, picking the patched function up), and the runner's namespace + is only re-aliased when that module already exists in ``sys.modules``. + Idempotent via a marker attribute on the installed function. + """ + import sys + + if getattr(split_attn_metadata, "_afd_positions_propagated_installed", False): + return + split_attn_metadata._afd_positions_propagated_installed = True # type: ignore[attr-defined] + ubatch_utils_module.split_attn_metadata = split_attn_metadata + runner = sys.modules.get("vllm.v1.worker.gpu_model_runner") + if runner is not None and hasattr(runner, "split_attn_metadata"): + runner.split_attn_metadata = split_attn_metadata + + +apply_positions_propagation() + +__all__ = ["apply_positions_propagation", "split_attn_metadata"] diff --git a/afd_plugin/compat/patches/ubatch_split.py b/afd_plugin/compat/patches/ubatch_split.py new file mode 100644 index 00000000..d254e509 --- /dev/null +++ b/afd_plugin/compat/patches/ubatch_split.py @@ -0,0 +1,169 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Request-aligned ubatch splitting for AFD DBO. + +Upstream vLLM's dual-batch overlap splits a batch at an even token count, +which cuts whichever request straddles that point into both ubatches. AFD's +DBO story is overlap between whole requests: one request runs in the first +ubatch while the other runs in the second, so the split must land on a +request boundary -- and a batch that cannot be split without cutting a +request must run whole instead of being divided. +""" + +from __future__ import annotations + +import numpy as np +from vllm.v1.worker.ubatch_utils import ( + UBatchSlice, + _pad_out_ubatch_slices, +) + +_DBO_UBATCH_COUNT = 2 + +# A ubatch's per-token tensors (positions, slot_mapping) are views into the +# step's buffers starting at the split point, so the split point decides their +# data pointers' alignment. DeepSeek-V4's CuTeDSL compressor kernel rejects any +# input below 64-byte alignment ("Misaligned Tensor data on argument #2"), and +# 16 four-byte tokens is the coarsest element stride that guarantees it for +# every per-token dtype in play. It is a requirement, not a preference: a batch +# with no aligned request boundary runs whole. This was once a preference with +# an unaligned fallback, to keep DBO splitting uniform decode (boundaries 1, 2, +# 3, ...), and that fallback crashed DeepSeek-V4 at startup -- a decode capture +# bucket of <=16 requests has no aligned boundary, so FULL decode graph capture +# with DBO on died in the compressor. Decode DBO was also measured as a +# regression (2.17x slower on DeepSeek-V2-Lite), so declining costs nothing. +_UBATCH_SPLIT_TOKEN_ALIGNMENT = 16 + + +def request_aligned_split_token(num_scheduled_tokens: np.ndarray) -> int | None: + """Token index of the request boundary nearest the half-way point. + + Only boundaries that leave every ubatch's per-token views aligned to + ``_UBATCH_SPLIT_TOKEN_ALIGNMENT`` qualify, because some attention kernels + reject a misaligned view outright. + + Returns ``None`` when no request boundary qualifies -- fewer than two + requests carrying tokens, or no aligned boundary among them -- meaning the + batch runs whole. + """ + cumulative = np.cumsum(np.asarray(num_scheduled_tokens, dtype=np.int64)) + total = int(cumulative[-1]) if cumulative.size else 0 + # Interior boundaries: every request edge except the batch start and the + # batch end. A boundary at either end would empty one ubatch. + boundaries = np.unique(cumulative[:-1]) + boundaries = boundaries[(boundaries > 0) & (boundaries < total)] + if boundaries.size == 0: + return None + aligned = boundaries[boundaries % _UBATCH_SPLIT_TOKEN_ALIGNMENT == 0] + if aligned.size == 0: + return None + nearest = int(np.argmin(np.abs(aligned - total / 2))) + return int(aligned[nearest]) + + +# Patch reason: upstream maybe_create_ubatch_slices splits at an even token +# count, cutting the straddling request into both ubatches. AFD overlaps whole +# requests, so the split must fall on a request boundary, and a batch with no +# interior boundary (a single request) must not be split at all. +# Patch functionality: with no explicit split point and exactly two ubatches, +# split at the request boundary nearest the even token split; return +# (None, None) -- vLLM's no-ubatch state -- when there is no such boundary. +# Explicit split points and other ubatch counts keep upstream behavior. +# Signature: matches upstream; no added parameters. +# Upstream: vLLM v0.26.0, vllm/v1/worker/ubatch_utils.py +def maybe_create_ubatch_slices( + should_ubatch: bool, + num_scheduled_tokens: np.ndarray, + num_tokens_padded: int, + num_reqs_padded: int, + num_ubatches: int, + split_point: list[int] | int | None = None, +) -> tuple[list[UBatchSlice] | None, list[UBatchSlice] | None]: + if not should_ubatch: + return None, None + + # ### PATCH START: request-aligned ubatch split + if split_point is None and num_ubatches == _DBO_UBATCH_COUNT: + aligned = request_aligned_split_token(num_scheduled_tokens) + if aligned is None: + # No interior request boundary: dividing would cut a request in + # half. Run the batch whole; vLLM treats absent slices as the + # single-batch path. + return None, None + split_point = aligned + # ### PATCH END: request-aligned ubatch split + if split_point is None: + split_point = int(num_tokens_padded) // num_ubatches + + token_split_points = [split_point * i for i in range(1, num_ubatches)] + + # TODO(lucas): Refactor the gpu_model_runner.py so we can pass + # in cu_num_tokens directly (i.e. query_start_loc) + cu_num_tokens = np.zeros(len(num_scheduled_tokens) + 1, dtype=np.int32) + np.cumsum(num_scheduled_tokens, dtype=np.int32, out=cu_num_tokens[1:]) + + ubatch_slices = [] + start_token = 0 + + # Add the end point to the split points to make iteration easier + # ### PATCH START: keep the final split point a Python int + # Upstream appends the numpy int32 straight off cu_num_tokens, which makes + # the last ubatch's token_slice.stop -- and therefore its + # num_actual_tokens, and every token count derived from it -- a + # numpy.int32. Triton refuses to specialize a numpy scalar, so DeepSeek-V4 + # dies in _build_c128a_topk_metadata_kernel on the last ubatch. + all_points = token_split_points + [int(cu_num_tokens[-1])] + # ### PATCH END: keep the final split point a Python int + + for end_token in all_points: + token_slice = slice(start_token, end_token) + + # Determine request slices using exclusive stop semantics + # Ubatch includes requests whose tokens overlap [start_token, end_token) + + # Start at the request that contains the start_token + # or the request starting exactly at start_token (if on boundary) + req_start = int(np.searchsorted(cu_num_tokens, start_token, side="right") - 1) + + # Stop at the request that starts at or after end_token + req_stop = int(np.searchsorted(cu_num_tokens, end_token, side="left")) + + req_slice = slice(req_start, req_stop) + ubatch_slices.append(UBatchSlice(req_slice, token_slice)) + + start_token = end_token + + ubatch_slices_padded = _pad_out_ubatch_slices( + ubatch_slices, num_tokens_padded, num_reqs_padded + ) + + assert sum(s.num_tokens for s in ubatch_slices_padded) == num_tokens_padded + + return ubatch_slices, ubatch_slices_padded + + +def apply_request_aligned_ubatch_split() -> None: + """Install the request-aligned splitter into vLLM's GPU runner. + + Both the execution and dummy-run call sites resolve the function through + the ``gpu_model_runner`` namespace, so patching that alias (and the + source module for any later importer) covers every caller. Idempotent via + a marker attribute on the installed function. + """ + from vllm.v1.worker import gpu_model_runner as gpu_model_runner_module + from vllm.v1.worker import ubatch_utils as ubatch_utils_module + + if getattr(maybe_create_ubatch_slices, "_afd_request_aligned", False): + return + maybe_create_ubatch_slices._afd_request_aligned = True # type: ignore[attr-defined] + gpu_model_runner_module.maybe_create_ubatch_slices = maybe_create_ubatch_slices + ubatch_utils_module.maybe_create_ubatch_slices = maybe_create_ubatch_slices + + +apply_request_aligned_ubatch_split() + +__all__ = [ + "apply_request_aligned_ubatch_split", + "maybe_create_ubatch_slices", + "request_aligned_split_token", +] diff --git a/afd_plugin/model_executor/models/deepseek_v2.py b/afd_plugin/model_executor/models/deepseek_v2.py index 235d177d..8919c1aa 100644 --- a/afd_plugin/model_executor/models/deepseek_v2.py +++ b/afd_plugin/model_executor/models/deepseek_v2.py @@ -15,7 +15,6 @@ import torch.nn as nn from transformers import DeepseekV2Config, DeepseekV3Config, GlmMoeDsaConfig from vllm.config import ParallelConfig, VllmConfig -from vllm.forward_context import get_forward_context from vllm.logger import init_logger from vllm.model_executor.layers import fused_moe from vllm.model_executor.layers.linear import ReplicatedLinear @@ -29,7 +28,10 @@ AFDTransferMetadata, ) from afd_plugin.model_executor.models import get_afd_metadata_from_forward_context -from afd_plugin.v1.worker.dbo import maybe_apply_dbo_yield +from afd_plugin.v1.worker.dbo import ( + current_dbo_ubatch_id, + maybe_apply_dbo_yield, +) logger = init_logger(__name__) @@ -127,9 +129,13 @@ def _send_and_receive( afd_metadata = get_afd_metadata_from_forward_context() if afd_metadata is None: raise RuntimeError("RemoteFFNProxy requires AFD forward metadata") - forward_context = get_forward_context() - stage_idx = int( - getattr(forward_context, "ubatch_idx", afd_metadata.stage_idx), + # vLLM tracks the ubatch by thread, not on the forward context, so + # forward_context.ubatch_idx does not exist and reading it made both + # DBO halves look like stage 0 -- one window slot for two concurrent + # dispatches, the second overwriting the first's flag. + dbo_ubatch_id = current_dbo_ubatch_id() + stage_idx = ( + afd_metadata.stage_idx if dbo_ubatch_id is None else int(dbo_ubatch_id) ) afd_metadata.stage_idx = stage_idx metadata = AFDTransferMetadata.create_attention_metadata( diff --git a/afd_plugin/model_executor/models/deepseek_v4.py b/afd_plugin/model_executor/models/deepseek_v4.py index 18707099..19234db8 100644 --- a/afd_plugin/model_executor/models/deepseek_v4.py +++ b/afd_plugin/model_executor/models/deepseek_v4.py @@ -14,7 +14,6 @@ import torch import torch.nn as nn from vllm.config import VllmConfig -from vllm.forward_context import get_forward_context from vllm.models.deepseek_v4.nvidia import model as native from afd_plugin.config import parse_afd_config @@ -24,7 +23,10 @@ ) from afd_plugin.connectors.metadata import AFDTransferContext, AFDTransferMetadata from afd_plugin.model_executor.models import get_afd_metadata_from_forward_context -from afd_plugin.v1.worker.dbo import maybe_apply_dbo_yield +from afd_plugin.v1.worker.dbo import ( + current_dbo_ubatch_id, + maybe_apply_dbo_yield, +) _ATTENTION_ROLE = frozenset(("attention",)) _FFN_ROLE = frozenset(("ffn",)) @@ -157,9 +159,15 @@ def forward( afd_metadata = get_afd_metadata_from_forward_context() if afd_metadata is None: raise RuntimeError("RemoteDeepseekV4FFN requires AFD forward metadata") - forward_context = get_forward_context() - stage_idx = int( - getattr(forward_context, "ubatch_idx", afd_metadata.stage_idx), + # vLLM tracks the ubatch by thread, not on the forward context, so + # forward_context.ubatch_idx does not exist and reading it pinned both + # DBO halves to stage 0 -- one window slot for two concurrent + # dispatches, the second overwriting the first's flag, after which the + # first forward waits for a reply that never comes. That surfaced as + # "RPC call to sample_tokens timed out" on a V4 decode run with DBO on. + dbo_ubatch_id = current_dbo_ubatch_id() + stage_idx = ( + afd_metadata.stage_idx if dbo_ubatch_id is None else int(dbo_ubatch_id) ) afd_metadata.stage_idx = stage_idx metadata = AFDTransferMetadata.create_attention_metadata( diff --git a/afd_plugin/v1/worker/attention_model_runner.py b/afd_plugin/v1/worker/attention_model_runner.py index 2dd1f32f..e862a22d 100644 --- a/afd_plugin/v1/worker/attention_model_runner.py +++ b/afd_plugin/v1/worker/attention_model_runner.py @@ -290,9 +290,16 @@ def _determine_batch_execution_and_padding( kwargs: dict[str, Any] = {} # determine if ubatch should be activated. - # 1. For dp = 1, vLLM hardcodes `should_ubatch=False`. - # This is the extra support for dp = 1 - if self.vllm_config.parallel_config.data_parallel_size == 1: + # 1. Whenever the cross-DP agreement did not run, nobody has decided + # yet and the answer above is a hardcoded False: vLLM hardcodes it for + # dp = 1, and `_dp_batch_coordination_disabled` hardcodes it for the + # connectors that opt out of the collective. Both cases need the + # rank-local decision instead, or `--enable-dbo` is accepted and then + # silently ignored for the whole run. + if ( + self.vllm_config.parallel_config.data_parallel_size == 1 + or self.connector.control_plane is None + ): should_ubatch = self._should_ubatch_single_rank( batch_descriptor, args, diff --git a/afd_plugin/v1/worker/ubatch_wrapper.py b/afd_plugin/v1/worker/ubatch_wrapper.py index 14265f2c..b1025f0d 100644 --- a/afd_plugin/v1/worker/ubatch_wrapper.py +++ b/afd_plugin/v1/worker/ubatch_wrapper.py @@ -71,6 +71,20 @@ def __call__(self, *args, **kwargs): forward_context = get_forward_context() ubatch_slices = forward_context.ubatch_slices if ubatch_slices is None: + # ### PATCH START: uncaptured FULL without a cudagraph wrapper + # runs eagerly. The AFD wrapper owns only cooperative capture and + # never builds a cudagraph_wrapper, so a whole-batch FULL dispatch + # whose key was never captured has nowhere to replay into and + # upstream asserts. A decode bucket the splitter declines to divide + # reaches this during capture of its own key. + if ( + forward_context.cudagraph_runtime_mode is CUDAGraphMode.FULL + and self.cudagraph_wrapper is None + and forward_context.batch_descriptor is not None + and forward_context.batch_descriptor.num_tokens not in self.cudagraphs + ): + return self.runnable(*args, **kwargs) + # ### PATCH END: uncaptured FULL without a cudagraph wrapper. return super().__call__(*args, **kwargs) cudagraph_runtime_mode = forward_context.cudagraph_runtime_mode diff --git a/recipe/gpu/GpuAsyncAFDConnector/deepseek_v4_flash/2a2f_async.sh b/recipe/gpu/GpuAsyncAFDConnector/deepseek_v4_flash/2a2f_async.sh index 56c8cc02..2a9ca57c 100755 --- a/recipe/gpu/GpuAsyncAFDConnector/deepseek_v4_flash/2a2f_async.sh +++ b/recipe/gpu/GpuAsyncAFDConnector/deepseek_v4_flash/2a2f_async.sh @@ -40,6 +40,23 @@ FFN_GRAPH_ARGS=() # parity (2.123 s vs 2.128 s). Leave it eager for a prefill-only run; turn it # on for anything that decodes. ATTN_EAGER=${ATTN_EAGER:-1} +# Set to 1 for vLLM's dual batch overlap. The flags go to BOTH roles: each side +# sizes its window rings from the visible stage count, and a count only one +# role sees leaves the FFN with fewer rings than the Attention dispatches into. +# +# DBO's sign on V4 is set by rows-per-expert per ubatch: splitting halves them +# on an already-skinny grouped GEMM, and the only thing the overlap hides is +# attention. Measured pure prefill on 4x L20X: -7.1% at 2048-token steps, +# -3.7% at 4096, +0.5% at 8192 with 4096-token prompts. +ENABLE_DBO=${ENABLE_DBO:-0} +DBO_ARGS=() +if [ "$ENABLE_DBO" = 1 ]; then + DBO_ARGS=( + --enable-dbo + --dbo-decode-token-threshold "${DBO_DECODE_THRESHOLD:-2}" + --dbo-prefill-token-threshold "${DBO_PREFILL_THRESHOLD:-12}" + ) +fi LOG_DIR=${LOG_DIR:-.} mkdir -p "$LOG_DIR" export VLLM_USE_V2_MODEL_RUNNER=0 @@ -126,6 +143,7 @@ CUDA_VISIBLE_DEVICES="$ATTN_DEVICES" "${VLLM_CMD[@]}" serve "$MODEL_PATH" \ --api-server-count 1 \ --gpu-memory-utilization "$GPU_MEM_UTIL" \ "${ATTN_GRAPH_ARGS[@]}" \ + "${DBO_ARGS[@]}" \ "${EXTRA_ARGS[@]}" \ --host 127.0.0.1 \ --port "$API_PORT" \ @@ -145,6 +163,7 @@ CUDA_VISIBLE_DEVICES="$FFN_DEVICES" "${VLLM_CMD[@]}" serve "$MODEL_PATH" \ --api-server-count 1 \ --gpu-memory-utilization "$GPU_MEM_UTIL" \ "${FFN_GRAPH_ARGS[@]}" \ + "${DBO_ARGS[@]}" \ "${EXTRA_ARGS[@]}" \ --host 127.0.0.1 \ --port "$FFN_API_PORT" \ diff --git a/tests/unit/compat/patches/test_ubatch_positions.py b/tests/unit/compat/patches/test_ubatch_positions.py new file mode 100644 index 00000000..10a7e3ae --- /dev/null +++ b/tests/unit/compat/patches/test_ubatch_positions.py @@ -0,0 +1,53 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Tests for the ubatch positions-propagation patch.""" + +from __future__ import annotations + +from types import SimpleNamespace + +from afd_plugin.compat.patches.ubatch_positions import split_attn_metadata + + +class _Ub: + def __init__(self, start: int, stop: int): + self.token_slice = slice(start, stop) + + +def test_split_propagates_positions_per_token_slice() -> None: + positions = list(range(10)) + source = SimpleNamespace(positions=positions) + slices = [_Ub(0, 4), _Ub(4, 10)] + + def fake_split(ubatch_slices, cm): + return [SimpleNamespace(positions=None) for _ in ubatch_slices] + + import afd_plugin.compat.patches.ubatch_positions as module + + original = module._upstream_split_attn_metadata + module._upstream_split_attn_metadata = fake_split + try: + results = split_attn_metadata(slices, source) + finally: + module._upstream_split_attn_metadata = original + + assert [list(cm.positions) for cm in results] == [[0, 1, 2, 3], [4, 5, 6, 7, 8, 9]] + + +def test_split_keeps_none_positions() -> None: + source = SimpleNamespace(positions=None) + slices = [_Ub(0, 4), _Ub(4, 10)] + + def fake_split(ubatch_slices, cm): + return [SimpleNamespace(positions=None) for _ in ubatch_slices] + + import afd_plugin.compat.patches.ubatch_positions as module + + original = module._upstream_split_attn_metadata + module._upstream_split_attn_metadata = fake_split + try: + results = split_attn_metadata(slices, source) + finally: + module._upstream_split_attn_metadata = original + + assert all(cm.positions is None for cm in results) diff --git a/tests/unit/compat/patches/test_ubatch_split.py b/tests/unit/compat/patches/test_ubatch_split.py new file mode 100644 index 00000000..af058b8c --- /dev/null +++ b/tests/unit/compat/patches/test_ubatch_split.py @@ -0,0 +1,160 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Tests for the request-aligned ubatch split patch.""" + +from __future__ import annotations + +import numpy as np +from vllm.v1.worker import gpu_model_runner, ubatch_utils +from vllm.v1.worker.ubatch_utils import UBatchSlice + +from afd_plugin.compat.patches import ubatch_split +from afd_plugin.compat.patches.ubatch_split import ( + maybe_create_ubatch_slices, + request_aligned_split_token, +) + + +def test_request_aligned_split_token_prefers_halfway_boundary() -> None: + # Requests of 192, 96, 96 tokens: boundaries at 192 and 288; the even + # split of 192 sits exactly on the first one. + assert request_aligned_split_token(np.array([192, 96, 96])) == 192 + + +def test_request_aligned_split_token_takes_nearest_boundary() -> None: + # Aligned boundaries at 64 and 288 for a 320-token batch; the halfway + # point 160 is nearer 64 (96 away) than 288 (128 away). + assert request_aligned_split_token(np.array([64, 224, 32])) == 64 + + +def test_request_aligned_split_token_single_request_has_no_boundary() -> None: + assert request_aligned_split_token(np.array([512])) is None + + +def test_request_aligned_split_token_prefers_an_aligned_boundary() -> None: + # Boundaries at 33 and 64 for a 100-token batch. 33 is nearer the halfway + # point, but only 64 leaves both ubatches' per-token views aligned, and + # some attention kernels reject a misaligned view outright. + assert request_aligned_split_token(np.array([33, 31, 36])) == 64 + + +def test_request_aligned_split_token_refuses_when_none_aligned() -> None: + # Uniform decode of six requests: boundaries 1..5, none aligned. Splitting + # at an unaligned one crashed DeepSeek-V4's compressor during FULL decode + # capture, so the batch runs whole instead. + assert request_aligned_split_token(np.array([1] * 6)) is None + + +def test_request_aligned_split_token_ignores_batch_edges() -> None: + # A zero-token request puts a cumulative sum at 0; a boundary there would + # empty the first ubatch. + assert request_aligned_split_token(np.array([0, 100])) is None + + +def test_request_aligned_split_lands_on_request_boundary() -> None: + # Two prefills of 304 and 96 tokens: the even token split (200) would cut + # the first request; the patch must split at 304 instead. + slices, slices_padded = maybe_create_ubatch_slices( + True, + np.array([304, 96]), + num_tokens_padded=400, + num_reqs_padded=2, + num_ubatches=2, + ) + assert slices is not None + assert slices_padded is not None + assert slices[0] == UBatchSlice(slice(0, 1), slice(0, 304)) + assert slices[1] == UBatchSlice(slice(1, 2), slice(304, 400)) + assert sum(s.num_tokens for s in slices_padded) == 400 + + +def test_small_uniform_decode_runs_whole() -> None: + # Four single-token decodes have boundaries 1, 2, 3 -- none aligned -- so + # the batch is not divided. This is the shape that crashed FULL decode + # capture on DeepSeek-V4 when it was split anyway. + assert maybe_create_ubatch_slices( + True, + np.array([1, 1, 1, 1]), + num_tokens_padded=4, + num_reqs_padded=4, + num_ubatches=2, + ) == (None, None) + + +def test_large_uniform_decode_splits_on_an_aligned_boundary() -> None: + # Forty-eight single-token decodes: 16 and 32 are the aligned boundaries, + # equidistant from the halfway point, so either is a valid split. + slices, _ = maybe_create_ubatch_slices( + True, + np.array([1] * 48), + num_tokens_padded=48, + num_reqs_padded=48, + num_ubatches=2, + ) + assert slices is not None + assert slices[0].token_slice.stop % 16 == 0 + + +def test_single_request_is_not_split() -> None: + # A lone 512-token prefill has no interior request boundary: DBO must not + # cut it in half. + assert maybe_create_ubatch_slices( + True, + np.array([512]), + num_tokens_padded=512, + num_reqs_padded=1, + num_ubatches=2, + ) == (None, None) + + +def test_no_split_when_disabled() -> None: + assert maybe_create_ubatch_slices( + False, + np.array([100, 100]), + num_tokens_padded=200, + num_reqs_padded=2, + num_ubatches=2, + ) == (None, None) + + +def test_explicit_split_point_keeps_upstream_behavior() -> None: + slices, _ = maybe_create_ubatch_slices( + True, + np.array([300, 100]), + num_tokens_padded=400, + num_reqs_padded=2, + num_ubatches=2, + split_point=200, + ) + # Upstream cuts at the given token count regardless of request edges, so + # the first request straddles both ubatches. + assert slices is not None + assert slices[0] == UBatchSlice(slice(0, 1), slice(0, 200)) + assert slices[1] == UBatchSlice(slice(0, 2), slice(200, 400)) + + +def test_non_two_ubatch_count_keeps_upstream_behavior() -> None: + slices, _ = maybe_create_ubatch_slices( + True, + np.array([300, 100]), + num_tokens_padded=400, + num_reqs_padded=2, + num_ubatches=4, + ) + # Upstream's even token split, cutting through the first request. + assert slices is not None + assert [s.token_slice for s in slices] == [ + slice(0, 100), + slice(100, 200), + slice(200, 300), + slice(300, 400), + ] + + +def test_patch_installed_on_vllm_modules() -> None: + assert gpu_model_runner.maybe_create_ubatch_slices is ( + ubatch_split.maybe_create_ubatch_slices + ) + assert ubatch_utils.maybe_create_ubatch_slices is ( + ubatch_split.maybe_create_ubatch_slices + ) diff --git a/tests/unit/model_executor/models/test_deepseek_v2_proxy.py b/tests/unit/model_executor/models/test_deepseek_v2_proxy.py index 7f34ca2b..71fdd9bb 100644 --- a/tests/unit/model_executor/models/test_deepseek_v2_proxy.py +++ b/tests/unit/model_executor/models/test_deepseek_v2_proxy.py @@ -47,11 +47,10 @@ def _install_fake_forward_context(monkeypatch, events, *, stage_idx=2): "get_afd_metadata_from_forward_context", lambda: afd_metadata, ) - monkeypatch.setattr( - adapter, - "get_forward_context", - lambda: SimpleNamespace(ubatch_idx=stage_idx), - ) + # The stage is the DBO ubatch this thread is running, which vLLM tracks by + # thread rather than on the forward context. Off DBO the helper returns + # None and the metadata's own stage stands. + monkeypatch.setattr(adapter, "current_dbo_ubatch_id", lambda: stage_idx) def record_yield(hidden_states, *, role): events.append(("yield", hidden_states, role)) @@ -204,11 +203,6 @@ def test_remote_proxy_requires_forward_metadata(monkeypatch): def test_remote_proxy_exchanges_cam_during_profile(monkeypatch): events: list[tuple] = [] _install_fake_forward_context(monkeypatch, events) - monkeypatch.setattr( - adapter, - "get_forward_context", - lambda: SimpleNamespace(in_profile_run=True, ubatch_idx=0), - ) hidden_states = torch.ones(2, 4) output = adapter.RemoteFFNProxy(layer_idx=0)(hidden_states) diff --git a/tests/unit/model_executor/models/test_deepseek_v4_proxy.py b/tests/unit/model_executor/models/test_deepseek_v4_proxy.py index 9ce00912..f01ed34d 100644 --- a/tests/unit/model_executor/models/test_deepseek_v4_proxy.py +++ b/tests/unit/model_executor/models/test_deepseek_v4_proxy.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project + from __future__ import annotations from types import SimpleNamespace @@ -7,6 +10,8 @@ torch = pytest.importorskip("torch") pytest.importorskip("vllm") +from torch import nn # noqa: E402 + from afd_plugin.model_executor.models import deepseek_v4 as adapter # noqa: E402 @@ -23,7 +28,7 @@ def recv_ffn_output(self, *, ref_tensor, ubatch_idx): def test_remote_v4_ffn_sends_token_ids(monkeypatch): - events = [] + events: list[tuple] = [] connector = _FakeConnector(events) afd_metadata = SimpleNamespace( connector=connector, @@ -34,11 +39,7 @@ def test_remote_v4_ffn_sends_token_ids(monkeypatch): "get_afd_metadata_from_forward_context", lambda: afd_metadata, ) - monkeypatch.setattr( - adapter, - "get_forward_context", - lambda: SimpleNamespace(ubatch_idx=2, slot_mapping={}), - ) + monkeypatch.setattr(adapter, "current_dbo_ubatch_id", lambda: 2) def record_yield(hidden_states, *, role): events.append(("yield", hidden_states, role)) @@ -68,7 +69,7 @@ def record_yield(hidden_states, *, role): def test_remote_v4_ffn_preserves_ids_in_padding_slots(monkeypatch): - events = [] + events: list[tuple] = [] connector = _FakeConnector(events) afd_metadata = SimpleNamespace(connector=connector, stage_idx=0) monkeypatch.setattr( @@ -76,14 +77,7 @@ def test_remote_v4_ffn_preserves_ids_in_padding_slots(monkeypatch): "get_afd_metadata_from_forward_context", lambda: afd_metadata, ) - monkeypatch.setattr( - adapter, - "get_forward_context", - lambda: SimpleNamespace( - ubatch_idx=0, - slot_mapping={"model.layers.0.attn": torch.tensor([5, 6, -1, -1])}, - ), - ) + monkeypatch.setattr(adapter, "current_dbo_ubatch_id", lambda: 0) monkeypatch.setattr( adapter, "maybe_apply_dbo_yield", @@ -129,7 +123,7 @@ def test_remote_v4_ffn_validates_token_ids_before_metadata_lookup( def test_v4_decoder_forward_rejects_ffn_role(monkeypatch): - class FakeMissingLayer(torch.nn.Module): + class FakeMissingLayer(nn.Module): pass monkeypatch.setattr(adapter.native, "PPMissingLayer", FakeMissingLayer) @@ -155,7 +149,7 @@ def test_v4_ffn_compute_rejects_attention_role_before_input_ids(): def test_v4_ffn_compute_requires_input_ids(monkeypatch): - class FakeMoE(torch.nn.Module): + class FakeMoE(nn.Module): pass monkeypatch.setattr(adapter.native, "DeepseekV4MoE", FakeMoE) diff --git a/tests/unit/v1/worker/test_attention_model_runner.py b/tests/unit/v1/worker/test_attention_model_runner.py index 7a8cbe03..138a4563 100644 --- a/tests/unit/v1/worker/test_attention_model_runner.py +++ b/tests/unit/v1/worker/test_attention_model_runner.py @@ -586,23 +586,68 @@ def test_should_ubatch_single_rank( @pytest.mark.parametrize( ( "dp_size", + "control_plane", "parent_should_ubatch", "num_tokens", "padded_num_tokens", "expected", ), [ - pytest.param(1, False, 48, 64, True, id="dp1-enables-local-dbo"), - pytest.param(1, False, 2, 64, False, id="dp1-rejects-empty-last-ubatch"), - pytest.param(2, True, 2, 64, True, id="dp2-keeps-coordinated-true"), - pytest.param(2, False, 48, 64, False, id="dp2-keeps-coordinated-false"), - pytest.param(2, True, 1, 1, False, id="dp2-rejects-empty-first-ubatch"), - pytest.param(2, True, 2, 2, True, id="dp2-keeps-minimal-nonempty-split"), + pytest.param( + 1, _DUMMY_CONTROL_PLANE, False, 48, 64, True, id="dp1-enables-local-dbo" + ), + pytest.param( + 1, + _DUMMY_CONTROL_PLANE, + False, + 2, + 64, + False, + id="dp1-rejects-empty-last-ubatch", + ), + pytest.param( + 2, _DUMMY_CONTROL_PLANE, True, 2, 64, True, id="dp2-keeps-coordinated-true" + ), + pytest.param( + 2, + _DUMMY_CONTROL_PLANE, + False, + 48, + 64, + False, + id="dp2-keeps-coordinated-false", + ), + pytest.param( + 2, + _DUMMY_CONTROL_PLANE, + True, + 1, + 1, + False, + id="dp2-rejects-empty-first-ubatch", + ), + pytest.param( + 2, + _DUMMY_CONTROL_PLANE, + True, + 2, + 2, + True, + id="dp2-keeps-minimal-nonempty-split", + ), + # Without a control plane the cross-DP agreement never runs, so the + # parent's False is a hardcoded placeholder rather than a decision: + # the rank has to decide locally or DBO never activates at all. + pytest.param(2, None, False, 48, 64, True, id="dp2-no-control-plane-decides"), + pytest.param( + 2, None, False, 2, 64, False, id="dp2-no-control-plane-rejects-empty" + ), ], ) -def test_determine_batch_execution_overrides_ubatch_only_for_dp1( +def test_determine_batch_execution_overrides_ubatch_without_dp_coordination( monkeypatch, dp_size, + control_plane, parent_should_ubatch, num_tokens, padded_num_tokens, @@ -610,6 +655,7 @@ def test_determine_batch_execution_overrides_ubatch_only_for_dp1( ): runner = _ubatch_runner( True, + control_plane=control_plane, data_parallel_size=dp_size, use_ubatching=True, num_ubatches=2, diff --git a/tests/unit/v1/worker/test_ubatch_wrapper.py b/tests/unit/v1/worker/test_ubatch_wrapper.py new file mode 100644 index 00000000..01cd2d78 --- /dev/null +++ b/tests/unit/v1/worker/test_ubatch_wrapper.py @@ -0,0 +1,58 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""The AFD ubatch wrapper's handling of steps the splitter did not divide.""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest + +pytest.importorskip("torch") +pytest.importorskip("vllm") + +from vllm.config import CUDAGraphMode # noqa: E402 + +from afd_plugin.v1.worker import ubatch_wrapper # noqa: E402 +from afd_plugin.v1.worker.ubatch_wrapper import AFDUBatchWrapper # noqa: E402 + + +def _wrapper(monkeypatch, *, mode, num_tokens, captured): + wrapper = object.__new__(AFDUBatchWrapper) + wrapper.cudagraph_wrapper = None # the AFD wrapper never builds one + wrapper.cudagraphs = dict.fromkeys(captured) + wrapper.runnable = lambda *a, **k: "eager" + context = SimpleNamespace( + ubatch_slices=None, + cudagraph_runtime_mode=mode, + batch_descriptor=SimpleNamespace(num_tokens=num_tokens), + ) + monkeypatch.setattr(ubatch_wrapper, "get_forward_context", lambda: context) + return wrapper + + +def test_uncaptured_full_step_runs_eagerly_instead_of_asserting(monkeypatch): + # A decode bucket the splitter declines to divide reaches the non-ubatch + # path in FULL mode with no graph for its key. Upstream would assert on + # the cudagraph_wrapper the AFD wrapper never constructs. + wrapper = _wrapper(monkeypatch, mode=CUDAGraphMode.FULL, num_tokens=8, captured=()) + assert wrapper() == "eager" + + +def test_captured_or_non_full_steps_keep_upstream_behaviour(monkeypatch): + calls: list[str] = [] + + def upstream_call(self, *args, **kwargs): + calls.append("upstream") + return "upstream" + + monkeypatch.setattr(ubatch_wrapper.UBatchWrapper, "__call__", upstream_call) + captured = _wrapper( + monkeypatch, mode=CUDAGraphMode.FULL, num_tokens=8, captured=(8,) + ) + assert captured() == "upstream" + eager_mode = _wrapper( + monkeypatch, mode=CUDAGraphMode.NONE, num_tokens=8, captured=() + ) + assert eager_mode() == "upstream" + assert calls == ["upstream", "upstream"]