[5/5][WIP] feat: request-aligned DBO ubatch splitting (flag off by default)

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,

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

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

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(

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(

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,