[4/5] feat: capture the FFN experts at bucketed padded shapes

afd_plugin/v1/worker/cuda_graph.py

@@ -12,9 +12,10 @@
 from collections.abc import Mapping
 from dataclasses import dataclass
 from enum import Enum
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
 if TYPE_CHECKING:
+    from torch import Tensor
     from vllm.config import VllmConfig
 
 FULL_DECODE_ONLY = "FULL_DECODE_ONLY"
@@ -111,26 +112,276 @@ def make_ffn_graph_key(
     """Extract the AFD FFN graph hashable key from DP metadata."""
 
     key_parts: list[tuple[int, tuple]] = []
+    values_tuple: tuple[Any, ...]
+    # Narrowed here rather than inside _use_ffn_aggregated_key so the sizes
+    # stay non-optional for the int() calls below.
+    aggregated = (
+        attention_size is not None
+        and ffn_size is not None
+        and _use_ffn_aggregated_key(attention_size, ffn_size)
+    )
     for stage_idx, metadata in sorted(dp_metadata_list.items()):
         values = getattr(metadata, "num_tokens_across_dp_cpu", None)
         if values is None:
-            if _use_ffn_aggregated_key(attention_size, ffn_size):
+            if aggregated:
+                assert ffn_size is not None
                 values_tuple = tuple(
                     max(1, int(fallback)) for _ in range(int(ffn_size))
                 )
             else:
                 values_tuple = (repr(metadata),)
         else:
             values_tuple = _metadata_values_tuple(values)
-            if _use_ffn_aggregated_key(attention_size, ffn_size):
+            if aggregated:
+                assert attention_size is not None and ffn_size is not None
                 values_tuple = _aggregate_ffn_values_tuple(
                     values_tuple,
                     attention_size=int(attention_size),
                     ffn_size=int(ffn_size),
                     fallback=int(fallback),
                 )
         key_parts.append((int(stage_idx), values_tuple))
-    return tuple(key_parts)
+    return tuple(key_parts)  # type: ignore[return-value]
+
+
+def padded_ffn_graph_shape(
+    *,
+    num_tokens: int,
+    topk: int,
+    ffn_size: int,
+    has_shared_experts: bool,
+) -> tuple[int, int]:
+    """Rows the captured shape has to hold.
+
+    The grouped GEMM does not care how many rows are real -- it reads its
+    grouping from a device-side count vector -- so a fixed row count can be
+    captured once and every smaller item padded up to it. What the shape has to
+    be is an upper bound for every item padded into it, which is the
+    largest batch the sender can produce.
+
+    A token contributes at most ``topk`` partials to any single FFN rank (the
+    case where every one of its experts lives there). Shared-expert rows are
+    split contiguously across the FFN ranks, so a rank holds at most
+    ``ceil(num_tokens / ffn_size)`` of them, and none at all without shared
+    experts.
+
+    Returns:
+        ``(max_routed_rows, max_shared_rows)``.
+    """
+    if num_tokens <= 0 or topk <= 0 or ffn_size <= 0:
+        raise ValueError(
+            "padded FFN graph shape needs positive num_tokens, topk and "
+            f"ffn_size; got {num_tokens}, {topk}, {ffn_size}",
+        )
+    max_routed = num_tokens * topk
+    max_shared = -(-num_tokens // ffn_size) if has_shared_experts else 0
+    return max_routed, max_shared
+
+
+# Every replay costs its captured row count, not the item's real one, so a
+# single graph at the upper bound charges the worst case to every item. The
+# bound assumes all of a token's topk partials land on one rank; with experts
+# spread over the ranks a token sends about ``topk / ffn_size`` of them, so a
+# real item occupies about ``max_routed / ffn_size`` rows -- call that the
+# expected size -- and a DBO ubatch half of that again.
+#
+# The ladder is therefore built as multiples of the expected size rather than
+# as fractions of the worst case. That distinction matters at the top of the
+# common range: routing scatter puts about half the items just above the
+# expected size, and a bucket boundary sitting exactly on it sends every one of
+# them a full step up. The multiples below cluster tightly just above 1.0 (and
+# above 0.5, where a DBO ubatch lands) so those items pay a few percent instead
+# of a quarter. Each entry costs one captured graph per MoE layer, so the
+# density is a memory trade.
+# 0.5 and 0.55 cover a DBO ubatch, which is half an item; 1.0 and 1.1 cover a
+# whole one. Each entry costs one captured graph per MoE layer, so the density
+# is spent at those two clusters rather than spread evenly.
+PADDED_FFN_GRAPH_EXPECTED_MULTIPLES: tuple[float, ...] = (
+    0.25,
+    0.5,
+    0.55,
+    0.7,
+    1.0,
+    1.1,
+    1.25,
+    1.5,
+)
+
+PADDED_FFN_GRAPH_FRACTIONS_ENV = "AFD_FFN_GRAPH_FRACTIONS"
+
+# GPU memory the padded-graph capture must leave free. The ladder is sized from
+# max_num_batched_tokens * topk, so its memory grows with the token budget: at
+# 8192 tokens the default ladder captured 65 GiB per DeepSeek-V4 FFN rank, and
+# the connector's NVSHMEM heap -- initialized after capture -- then failed with
+# "cuMemCreate failed". Capture stops at the first layer whose graphs would cut
+# into this reserve and the rest run eagerly, so the token budget can no longer
+# exhaust the card. It is a calibration knob: NVSHMEM's own heap plus the
+# window is a few GiB, and the default leaves room over that.
+PADDED_FFN_GRAPH_RESERVE_ENV = "AFD_FFN_GRAPH_MEM_RESERVE_GIB"
+PADDED_FFN_GRAPH_RESERVE_GIB_DEFAULT = 8.0
+
+# Bucket row counts are rounded up to this, so a bucket is always a whole
+# number of GEMM tiles rather than a ragged tail.
+PADDED_FFN_BUCKET_ALIGNMENT = 128
+
+# How much padding a replay may carry before running the item eagerly instead.
+# A replay pays for its whole bucket but saves the per-layer launch work; eager
+# pays launches but computes only the real rows. Measured on DeepSeek-V4 2A2F,
+# eager beat a replay padded to 1.18x by 6%, so a replay only pays for itself
+# when it is close to the item's own size. Items landing exactly on a bucket
+# still replay; the rest take the cheaper path.
+MAX_REPLAY_PADDING_RATIO = 1.05
+
+
+def resolve_padded_ffn_graph_fractions(
+    environ: Mapping[str, str] | None = None,
+) -> tuple[float, ...]:
+    """Parse the multiples override, falling back to the default set."""
+    import os
+
+    source = os.environ if environ is None else environ
+    raw = source.get(PADDED_FFN_GRAPH_FRACTIONS_ENV, "").strip()
+    if not raw:
+        return PADDED_FFN_GRAPH_EXPECTED_MULTIPLES
+    return tuple(float(part) for part in raw.replace(",", " ").split())
+
+
+def resolve_padded_ffn_graph_reserve_bytes(
+    environ: Mapping[str, str] | None = None,
+) -> int:
+    """Bytes of GPU memory padded-graph capture must leave free."""
+    import os
+
+    source = os.environ if environ is None else environ
+    raw = source.get(PADDED_FFN_GRAPH_RESERVE_ENV, "").strip()
+    gib = float(raw) if raw else PADDED_FFN_GRAPH_RESERVE_GIB_DEFAULT
+    if gib < 0:
+        raise ValueError(f"{PADDED_FFN_GRAPH_RESERVE_ENV} must be >= 0; got {raw!r}")
+    return int(gib * 2**30)
+
+
+def padded_ffn_graph_layer_fits(
+    *,
+    free_bytes: int,
+    per_layer_bytes: int,
+    reserve_bytes: int,
+) -> bool:
+    """Whether capturing one more layer keeps ``reserve_bytes`` free.
+
+    ``per_layer_bytes`` is what the layers captured so far took each; before
+    the first layer it is 0, so that layer is captured whenever the reserve is
+    already there.
+    """
+    return free_bytes - per_layer_bytes >= reserve_bytes
+
+
+def padded_ffn_graph_buckets(
+    max_routed: int,
+    *,
+    ffn_size: int = 1,
+    fractions: tuple[float, ...] | None = None,
+) -> tuple[int, ...]:
+    """Ascending distinct row counts to capture for one MoE layer.
+
+    The multiples are of the expected item size, ``max_routed / ffn_size``, not
+    of ``max_routed`` -- see the comment on the default set. Entries are rounded
+    up to ``PADDED_FFN_BUCKET_ALIGNMENT`` so a bucket boundary stays a sane GEMM
+    tile count, clamped to ``max_routed``, and deduplicated.
+
+    The ladder need not reach ``max_routed``: an item above the largest bucket
+    runs eager, the same fallback an item larger than the captured shape has
+    always taken. That is only safe because ``capture_padded_ffn_graphs`` pins
+    the shared MoE workspace at the ceiling before capturing -- see the note
+    there; without it, the first oversized item grows the workspace and
+    invalidates every captured graph.
+    """
+    if max_routed <= 0:
+        raise ValueError(f"max_routed must be positive; got {max_routed}")
+    if ffn_size <= 0:
+        raise ValueError(f"ffn_size must be positive; got {ffn_size}")
+    if fractions is None:
+        fractions = resolve_padded_ffn_graph_fractions()
+    expected = max_routed / ffn_size
+    buckets = set()
+    for multiple in fractions:
+        if multiple <= 0:
+            raise ValueError(
+                f"padded FFN graph multiples must be positive; got {multiple}",
+            )
+        rows = int(expected * multiple)
+        rows = -(-rows // PADDED_FFN_BUCKET_ALIGNMENT) * PADDED_FFN_BUCKET_ALIGNMENT
+        buckets.add(min(max(rows, PADDED_FFN_BUCKET_ALIGNMENT), max_routed))
+    return tuple(sorted(buckets))
+
+
+def shared_rows_for_bucket(
+    bucket: int,
+    *,
+    max_routed: int,
+    max_shared: int,
+) -> int:
+    """Shared-expert rows captured alongside a bucket's routed rows.
+
+    Both counts scale with the item's token count, so a bucket holding half the
+    routed rows needs half the shared rows. Capturing every bucket at
+    ``max_shared`` instead makes a half-sized item pay double on the shared
+    expert -- measured as graphs losing to eager under DBO even once the routed
+    rows were bucketed.
+    """
+    if max_shared <= 0:
+        return 0
+    rows = -(-max_shared * bucket // max_routed)
+    return min(max(rows, 1), max_shared)
+
+
+def select_padded_ffn_bucket(
+    buckets: tuple[int, ...],
+    routed_rows: int,
+    shared_rows: int = 0,
+    *,
+    max_routed: int | None = None,
+    max_shared: int = 0,
+) -> int | None:
+    """Smallest captured bucket holding both row counts, or ``None``.
+
+    ``None`` means no bucket fits and the caller runs the item eagerly, which
+    is what an item larger than the captured maximum has always done. The
+    shared rows are checked too, because each bucket's graph captures only its
+    own share of them.
+    """
+    for bucket in buckets:
+        if routed_rows > bucket:
+            continue
+        if max_routed is not None and shared_rows > shared_rows_for_bucket(
+            bucket, max_routed=max_routed, max_shared=max_shared
+        ):
+            continue
+        return bucket
+    return None
+
+
+def pad_counts_to_shape(
+    counts: Tensor,
+    *,
+    padded_rows: int,
+    actual_rows: int,
+) -> None:
+    """Grow ``counts`` in place so its entries sum to ``padded_rows``.
+
+    The padding lands on the last expert, which is where the padded rows are:
+    real rows are grouped by expert in ascending order, so the tail of the row
+    range belongs to the last expert either way. That keeps every real row's
+    expert assignment untouched.
+
+    The padded rows carry whatever the input buffer last held. Their output is
+    sliced off and discarded, and a grouped GEMM is row-independent, so their
+    content cannot reach a real row.
+    """
+    if actual_rows > padded_rows:
+        raise ValueError(
+            f"{actual_rows} rows do not fit the padded shape {padded_rows}",
+        )
+    counts[-1] += padded_rows - actual_rows
 
 
 def graph_run_mode(
@@ -150,7 +401,7 @@ def graph_run_mode(
     return AFDGraphRunMode.EAGER
 
 
-def _metadata_values_tuple(values: object) -> tuple[int, ...]:
+def _metadata_values_tuple(values: Any) -> tuple[int, ...]:
     tolist = getattr(values, "tolist", None)
     if callable(tolist):
         values = tolist()
@@ -206,5 +457,18 @@ def _aggregate_ffn_values_tuple(
     "cudagraph_mode_name",
     "graph_run_mode",
     "make_ffn_graph_key",
+    "pad_counts_to_shape",
+    "MAX_REPLAY_PADDING_RATIO",
+    "PADDED_FFN_BUCKET_ALIGNMENT",
+    "PADDED_FFN_GRAPH_EXPECTED_MULTIPLES",
+    "PADDED_FFN_GRAPH_FRACTIONS_ENV",
+    "PADDED_FFN_GRAPH_RESERVE_ENV",
+    "padded_ffn_graph_layer_fits",
+    "resolve_padded_ffn_graph_reserve_bytes",
+    "padded_ffn_graph_buckets",
+    "resolve_padded_ffn_graph_fractions",
+    "padded_ffn_graph_shape",
+    "select_padded_ffn_bucket",
+    "shared_rows_for_bucket",
     "validate_cuda_graph_mode",
 ]