diff --git a/afd_plugin/v1/worker/cuda_graph.py b/afd_plugin/v1/worker/cuda_graph.py index 71323dd7..7b412c8f 100644 --- a/afd_plugin/v1/worker/cuda_graph.py +++ b/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,10 +112,19 @@ 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)) ) @@ -122,7 +132,8 @@ def make_ffn_graph_key( 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), @@ -130,7 +141,247 @@ def make_ffn_graph_key( 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", ] diff --git a/afd_plugin/v1/worker/ffn_model_runner.py b/afd_plugin/v1/worker/ffn_model_runner.py index 08f86ab0..29f3482c 100644 --- a/afd_plugin/v1/worker/ffn_model_runner.py +++ b/afd_plugin/v1/worker/ffn_model_runner.py @@ -5,6 +5,7 @@ from __future__ import annotations from contextlib import contextmanager +from dataclasses import dataclass from typing import TYPE_CHECKING, Any import torch @@ -13,6 +14,7 @@ from vllm.config import update_config as update_vllm_config from vllm.distributed.parallel_state import get_world_group, graph_capture from vllm.forward_context import DPMetadata, get_forward_context, set_forward_context +from vllm.logger import init_logger from vllm.model_executor.layers.rotary_embedding import _ROPE_DICT from vllm.model_executor.model_loader import get_model_loader from vllm.utils.mem_utils import DeviceMemoryProfiler @@ -39,9 +41,17 @@ fail_if_unsupported_ubatching, ) from afd_plugin.v1.worker.cuda_graph import ( + MAX_REPLAY_PADDING_RATIO, AFDGraphRunMode, graph_run_mode, make_ffn_graph_key, + pad_counts_to_shape, + padded_ffn_graph_buckets, + padded_ffn_graph_layer_fits, + padded_ffn_graph_shape, + resolve_padded_ffn_graph_reserve_bytes, + select_padded_ffn_bucket, + shared_rows_for_bucket, validate_cuda_graph_mode, ) from afd_plugin.v1.worker.ffn_metadata import ( @@ -49,12 +59,30 @@ project_ffn_token_counts_to_dp, ) +# Name the logger inside vLLM's tree so its handler picks the lines up; a bare +# afd_plugin.* logger propagates to a handler-less root and is dropped. +logger = init_logger(f"vllm.{__name__}") + if TYPE_CHECKING: from vllm.sequence import IntermediateTensors from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheSpec +# How often the padded-graph runner reports its padding overhead. One line per +# this many replays keeps a 43-layer step from writing a line per layer. +_PADDED_STATS_EVERY = 2000 + + +@dataclass(slots=True) +class _PaddedFFNGraph: + """One MoE layer's experts, captured at the padded shape.""" + + graph: torch.cuda.CUDAGraph + routed_out: torch.Tensor + shared_out: torch.Tensor | None + + class GPUFFNModelRunner(LoRAModelRunnerMixin): """FFN model runner for AFD GPU execution. @@ -86,9 +114,10 @@ def __init__(self, vllm_config: VllmConfig, device: object) -> None: # A connector without a control plane drives FFN steps from its own # receive loop instead of from broadcast DP metadata. self.is_connector_driven = self.connector.control_plane is None - # The connector-driven path never touches vLLM's graph machinery, so - # vLLM's cudagraph_mode says nothing about it -- running the policy gate - # here would reject modes that are simply irrelevant. + # The connector-driven path captures its own padded graphs and never + # touches vLLM's graph machinery, so vLLM's cudagraph_mode says nothing + # about it -- running the policy gate here would reject modes that are + # simply irrelevant. The only question is whether graphs are wanted. if self.is_connector_driven: self.afd_cudagraph_policy = None else: @@ -100,12 +129,32 @@ def __init__(self, vllm_config: VllmConfig, device: object) -> None: self.model: Any = None self.model_memory_usage = 0 self.num_layers = int(self.model_config.hf_text_config.num_hidden_layers) - self.use_cuda_graph = bool( - self.afd_cudagraph_policy is not None - and self.afd_cudagraph_policy.enable_ffn_graph_cache + self.use_cuda_graph = ( + not bool(self.model_config.enforce_eager) + if self.is_connector_driven + else bool( + self.afd_cudagraph_policy is not None + and self.afd_cudagraph_policy.enable_ffn_graph_cache + ) ) self._cuda_graphs: dict[tuple, dict[str, Any]] = {} self._graph_memory_pool: Any | None = None + # Connector-driven padded graphs, one per MoE layer. Empty unless the + # async connector runs with graphs on; see capture_padded_ffn_graphs. + self._padded_graphs: dict[tuple[int, int], _PaddedFFNGraph] = {} + self._padded_buckets: tuple[int, ...] = () + # Rows a replay really carried vs rows it was charged, so a run reports + # how much of the FFN's GPU time went to padding. Host-side ints; the + # ladder is only worth tuning against a measured distribution. + self._padded_rows_real = 0 + self._padded_rows_charged = 0 + self._padded_replays = 0 + self._padded_eager_items = 0 + self._padded_hidden: torch.Tensor | None = None + self._padded_counts: torch.Tensor | None = None + self._padded_shared: torch.Tensor | None = None + self._padded_max_routed = 0 + self._padded_max_shared = 0 self.prof = create_afd_gpu_profiler("ffn") @property @@ -280,17 +329,323 @@ def execute_connector_driven_step(self) -> None: step_afd_gpu_profiler(self.prof) self._ffn_forward_connector_driven() + # ================================================================== + # Connector-driven padded CUDA graphs + # ================================================================== + + def capture_padded_ffn_graphs(self) -> int: + """Capture the local experts once per MoE layer, at the maximum shape. + + The connector-driven path has no control plane, so nothing tells this + rank the shape of the next work item -- which is why it ran eagerly. + Padding removes the need to know: a grouped GEMM takes its grouping + from a device-side count vector rather than from its row count, so the + largest shape can be captured and every item padded up to it, with the + padding charged to the last expert and its output sliced off. Only the + counts differ between replays, and a replay re-reads them. + + The shape is the largest batch the sender can produce, + ``max_num_batched_tokens``, so prefill and decode both replay it. + + One graph per layer, because a graph records the weight pointers and + each layer has its own. Input buffers are shared across layers -- work + items are served one at a time on one stream. + + Called once, before the connector joins its process group. That + rendezvous is the only barrier between the roles, and the Attention + rank profiles -- and so dispatches -- the moment it clears; capturing + after joining would leave those dispatches landing in a slot nobody is + polling yet. + + Returns: + Bytes of device memory the graphs took. + """ + if not self.use_cuda_graph or not self.is_connector_driven: + return 0 + if self._padded_graphs: + # start_ffn_server_loop is callable more than once. + return 0 + if self.model is None: + raise RuntimeError("capture_padded_ffn_graphs needs a loaded model") + + connector: Any = self.connector + device = torch.device(self.device) + # This runs on whichever thread called start_ffn_server_loop, which is + # not the serving thread and has no device bound yet. + torch.cuda.set_device(device) + + self._padded_max_routed, self._padded_max_shared = padded_ffn_graph_shape( + num_tokens=int(self.vllm_config.scheduler_config.max_num_batched_tokens), + topk=connector.topk, + ffn_size=connector.ffn_size, + has_shared_experts=connector.has_shared_experts, + ) + self._padded_hidden = torch.zeros( + (self._padded_max_routed, connector.hidden_size), + dtype=self.dtype, + device=device, + ) + self._padded_counts = torch.zeros( + connector.expert_per_rank, + dtype=torch.int32, + device=device, + ) + self._padded_shared = ( + torch.zeros( + (self._padded_max_shared, connector.hidden_size), + dtype=self.dtype, + device=device, + ) + if self._padded_max_shared + else None + ) + # A grouping that fills the captured shape. Any grouping records the + # same kernels; a replay re-reads the counts. + self._padded_buckets = padded_ffn_graph_buckets( + self._padded_max_routed, + ffn_size=connector.ffn_size, + ) + + inner = self.model.model + # Same source as the forward loop above: the model reports which of its + # layers own experts. Reading a per-layer flag instead only works for + # the adapters that happen to define one -- DeepSeek-V4's layers do not. + experts_layer_indices = frozenset(self.model.get_experts_layer_indices()) + moe_layers = [ + layer_idx + for layer_idx in range(inner.start_layer, inner.end_layer) + if layer_idx in experts_layer_indices + ] + num_warmups = max( + 1, + int(self.vllm_config.compilation_config.cudagraph_num_of_warmups), + ) + + start_free_gpu_memory = torch.cuda.mem_get_info()[0] + if self._graph_memory_pool is None: + self._graph_memory_pool = torch.cuda.graph_pool_handle() + + # Pin the shared MoE workspace at its ceiling before capturing anything. + # It grows by freeing the old buffer and allocating a bigger one, so any + # growth after a capture leaves that graph reading freed memory -- one + # run died with an illegal memory access when an oversized item took the + # eager path and grew it. One eager call at the largest shape any path + # can ask for settles it, and costs a single forward instead of the + # captured graph per layer that reserving it through the ladder would. + with _ffn_forward_context(self.vllm_config) as warmup_context: + _set_moe_layer_index(warmup_context, moe_layers[0]) + self._fill_padded_counts(self._padded_max_routed) + self._padded_ffn_compute(moe_layers[0], self._padded_max_routed) + torch.cuda.synchronize() + + reserve_bytes = resolve_padded_ffn_graph_reserve_bytes() + per_layer_bytes = 0 + captured_layers = 0 + set_cudagraph_capturing_enabled(True) + try: + with ( + _ffn_forward_context(self.vllm_config) as forward_context, + graph_capture(device=self.device), + ): + for layer_idx in moe_layers: + free_bytes = torch.cuda.mem_get_info()[0] + if not padded_ffn_graph_layer_fits( + free_bytes=free_bytes, + per_layer_bytes=per_layer_bytes, + reserve_bytes=reserve_bytes, + ): + # Every layer from here runs eagerly: a work item whose + # (layer, bucket) has no graph already takes that path. + break + _set_moe_layer_index(forward_context, layer_idx) + # Largest bucket first, and that order is load-bearing: + # vLLM's WorkspaceManager grows the shared MoE scratch by + # freeing the old buffer and allocating a bigger one, which + # leaves every graph already captured against the old + # pointer dangling. Capturing ascending therefore made each + # small bucket's graph fault on its first replay. Starting + # at the largest sizes the workspace once, and every + # smaller capture reuses it. + for bucket in reversed(self._padded_buckets): + # A replay costs its captured row count, so each bucket + # gets its own graph and an item takes the smallest one + # that holds it. + self._fill_padded_counts(bucket) + # Warm first: the fused MoE picks a kernel on its first + # call, and that choice must settle before capture -- + # autotuning synchronizes, which capture forbids. + for _ in range(num_warmups): + self._padded_ffn_compute(layer_idx, bucket) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, pool=self._graph_memory_pool): + payload = self._padded_ffn_compute(layer_idx, bucket) + self._padded_graphs[(layer_idx, bucket)] = _PaddedFFNGraph( + graph=graph, + routed_out=payload.routed_output, + shared_out=payload.shared_output, + ) + per_layer_bytes = max( + per_layer_bytes, + free_bytes - torch.cuda.mem_get_info()[0], + ) + captured_layers += 1 + finally: + set_cudagraph_capturing_enabled(False) + + graph_bytes = start_free_gpu_memory - torch.cuda.mem_get_info()[0] + logger.info( + "AFD FFN padded graphs ready: graphs=%d layers=%d/%d buckets=%s " + "shared_rows=%d experts=%d size=%.1fMiB", + len(self._padded_graphs), + captured_layers, + len(moe_layers), + list(self._padded_buckets), + self._padded_max_shared, + connector.expert_per_rank, + graph_bytes / 2**20, + ) + return int(graph_bytes) + + def _fill_padded_counts(self, bucket: int) -> None: + """Spread ``bucket`` rows over every local expert for capture. + + Not any grouping will do, even though the row total is what sizes the + grouped GEMM. The fused MoE pads each non-empty expert's rows up to a + block, so the scratch it reserves grows with the number of experts that + have rows -- and a capture is stuck with the scratch it reserved. All + rows on one expert reserves the least, and a replay whose real routing + touches every expert then runs off the end of it. Spreading rows evenly + is both the worst case for that padding and what real routing looks + like, so the capture reserves enough for any replay. + """ + assert self._padded_counts is not None + num_experts = int(self._padded_counts.numel()) + share, remainder = divmod(bucket, num_experts) + self._padded_counts.fill_(share) + if remainder: + # One row each rather than all on the last expert, so a bucket + # smaller than the expert count still touches as many experts as + # it has rows instead of collapsing onto one. + self._padded_counts[:remainder] += 1 + + def _shared_rows_for(self, bucket: int) -> int: + """This bucket's share of the shared-expert rows.""" + return shared_rows_for_bucket( + bucket, + max_routed=self._padded_max_routed, + max_shared=self._padded_max_shared, + ) + + def _padded_ffn_compute( + self, + layer_idx: int, + bucket: int, + ) -> AFDF2ATransferPayload: + """Run one layer's experts over the first ``bucket`` padded rows.""" + assert self._padded_hidden is not None + shared = self._padded_shared + if shared is not None: + shared = shared[: self._shared_rows_for(bucket)] + return self.model.compute_ffn_output( + hidden_states=self._padded_hidden[:bucket], + layer_idx=layer_idx, + group_list=self._padded_counts, + expand_x_shared=shared, + ) + def _compute_work_item( self, work_item: Any, states: GpuAsyncTransferState, ) -> torch.Tensor | AFDF2ATransferPayload: - """Run this work item's layer over the rows that arrived.""" - return self.model.compute_ffn_output( - hidden_states=work_item.hidden_states, - layer_idx=work_item.layer_idx, - group_list=states.group_list, - expand_x_shared=states.expand_x_shared, + """Replay this layer's smallest fitting padded graph, or run it eagerly. + + A replay costs its captured row count rather than the item's real one, + so taking the smallest bucket that holds the item is what keeps a small + item cheap -- a DBO ubatch is a quarter of a full batch's rows, and + charging it the full-batch shape is what made ubatching a loss. + + Eager is the fallback for an item that does not fit the captured shape + and for the empty item a decode can produce when none of a token's + experts landed here. + """ + routed_rows = int(states.routed_tokens) + shared_rows = int(states.shared_tokens) + bucket = ( + select_padded_ffn_bucket( + self._padded_buckets, + routed_rows, + shared_rows, + max_routed=self._padded_max_routed, + max_shared=self._padded_max_shared, + ) + if routed_rows > 0 + else None + ) + if bucket is not None and bucket > routed_rows * MAX_REPLAY_PADDING_RATIO: + # The padding this replay would carry costs more than the launches + # eager pays. Take the cheaper of the two per item rather than + # charging every item to the nearest bucket above it. + bucket = None + graph = ( + self._padded_graphs.get((work_item.layer_idx, bucket)) + if bucket is not None + else None + ) + fits = ( + graph is not None + and 0 < routed_rows <= self._padded_max_routed + and shared_rows <= self._padded_max_shared + ) + if not fits: + self._padded_eager_items += 1 + return self.model.compute_ffn_output( + hidden_states=work_item.hidden_states, + layer_idx=work_item.layer_idx, + group_list=states.group_list, + expand_x_shared=states.expand_x_shared, + ) + + assert graph is not None and bucket is not None + self._padded_rows_real += routed_rows + self._padded_rows_charged += bucket + self._padded_replays += 1 + if self._padded_replays % _PADDED_STATS_EVERY == 0: + logger.info( + "AFD FFN padded replays=%d eager=%d rows_real=%d rows_charged=%d " + "padding_overhead=%.2fx", + self._padded_replays, + self._padded_eager_items, + self._padded_rows_real, + self._padded_rows_charged, + self._padded_rows_charged / max(self._padded_rows_real, 1), + ) + + assert self._padded_hidden is not None + assert self._padded_counts is not None + if not states.staged_routed: + # The receive could not use the buffer -- no graph existed when the + # item arrived, or it did not fit -- so the rows still need moving. + self._padded_hidden[:routed_rows].copy_(work_item.hidden_states) + self._padded_counts.copy_(states.group_list) + pad_counts_to_shape( + self._padded_counts, + padded_rows=bucket, + actual_rows=routed_rows, + ) + if self._padded_shared is not None and shared_rows: + self._padded_shared[:shared_rows].copy_(states.expand_x_shared) + graph.graph.replay() + # Views into the graph's own output buffers. The next replay overwrites + # them, and the reply that consumes them is queued before it on this + # stream, so the ordering holds without a copy. + return AFDF2ATransferPayload( + routed_output=graph.routed_out[:routed_rows], + shared_output=( + graph.shared_out[:shared_rows] + if graph.shared_out is not None and shared_rows + else None + ), ) def _ffn_forward_connector_driven( @@ -307,6 +662,11 @@ def _ffn_forward_connector_driven( work_item = connector.recv_ffn_work_item( # type: ignore[attr-defined] stage_idx=stage_idx, max_num_tokens=self.vllm_config.scheduler_config.max_num_batched_tokens, + # Hand the graph its input buffer so the arrival's + # gather lands there directly. The gather had to write + # somewhere either way; staging afterwards would be a + # second pass over the whole payload, per layer. + routed_out=self._padded_hidden, ) except TimeoutError: # Nothing pending; hand control back so the worker loop can diff --git a/afd_plugin/v1/worker/ffn_worker.py b/afd_plugin/v1/worker/ffn_worker.py index 66784a4b..3693ff7b 100644 --- a/afd_plugin/v1/worker/ffn_worker.py +++ b/afd_plugin/v1/worker/ffn_worker.py @@ -150,6 +150,26 @@ def start_ffn_server_loop(self) -> None: return self.raise_ffn_loop_error_if_any() + # Capture before the rendezvous, not after. The connector's process + # group is the only barrier between the two roles: the Attention rank + # leaves it and immediately profiles, which dispatches. Capturing after + # joining leaves a window -- 35 seconds for a 26-layer model -- where + # those dispatches land in a slot nobody is polling. The Attention rank + # then parks its stream on a flag that will not be stamped until this + # side starts polling, runs ahead until the launch queue fills, and + # blocks mid-write with the next dispatch's flag unstamped, which the + # poll can never see. Under vLLM's ubatching that is fatal rather than + # merely slow: the blocked thread never reaches its next yield and the + # peer ubatch waits on it forever. + # + # Capture needs the model and the connector's shape config, both ready + # before init_afd_connector, so moving it earlier costs nothing. + # + # Here rather than in initialize_from_config because an AFD FFN + # EngineCore is a daemon that reaches this by collective_rpc and never + # runs KV/scheduler setup; this is the one point both entry paths cross. + self.model_runner.capture_padded_ffn_graphs() + connector = self.model_runner.connector if not connector.is_initialized: self.model_runner.initialize_afd_connector() diff --git a/recipe/gpu/GpuAsyncAFDConnector/deepseek_v4_flash/2a2f_async.sh b/recipe/gpu/GpuAsyncAFDConnector/deepseek_v4_flash/2a2f_async.sh index a426b2d1..56c8cc02 100755 --- a/recipe/gpu/GpuAsyncAFDConnector/deepseek_v4_flash/2a2f_async.sh +++ b/recipe/gpu/GpuAsyncAFDConnector/deepseek_v4_flash/2a2f_async.sh @@ -4,6 +4,9 @@ # 2A2F DeepSeek-V4-Flash on the async GPU connector. # +# FFN_EAGER picks the FFN side's run mode; it defaults to eager, which is what +# measured fastest on this model -- see the note on it below. +# # Launch under a GPU reservation, which sets CUDA_VISIBLE_DEVICES: # gpu run --gpus 4 -- \ # bash recipe/gpu/GpuAsyncAFDConnector/deepseek_v4_flash/2a2f_async.sh @@ -17,6 +20,26 @@ MODEL_PATH=${MODEL_PATH:-/path/model_weights/deepseek-v4-flash} # How to invoke vLLM. `uv run vllm` is right from a synced checkout; override # to point at an interpreter that actually has the plugin installed. read -r -a VLLM_CMD <<< "${VLLM_CMD:-uv run vllm}" +# The FFN experts run eagerly by default. Their padded graphs hold ~19 GiB +# and buy nothing on V4: the FFN is compute-bound, so the per-item launch +# saving is negligible, while a replay costs its whole captured bucket. +# Measured pure prefill, 2A2F on 4x L20X, 128x1024 tokens: 17.9 s eager +# against 19.0 s replaying every item. Set FFN_EAGER=0 to capture them anyway. +FFN_EAGER=${FFN_EAGER:-1} +FFN_GRAPH_ARGS=() +[ "$FFN_EAGER" = 1 ] && FFN_GRAPH_ARGS=(--enforce-eager) +# Attention-side run mode: +# 1 -- eager (the default) +# 0 -- FULL_DECODE_ONLY graphs +# +# Decode is where the Attention graphs pay, and they pay a lot: measured on +# DeepSeek-V2-Lite 1A1F decode (64 req x 256 output tokens, conc 16), eager +# 36.6 s against 28.3 s with graphs -- 22.6% off the wall clock, with the +# wrapper reporting a 98% replay share. Prefill is the opposite: those graphs +# are captured and never replayed there, and the same flags measured dead +# 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} LOG_DIR=${LOG_DIR:-.} mkdir -p "$LOG_DIR" export VLLM_USE_V2_MODEL_RUNNER=0 @@ -60,6 +83,22 @@ API_PORT=${API_PORT:-18307} # whichever loses exits and takes its role down with it. Give it its own. FFN_API_PORT=${FFN_API_PORT:-$((API_PORT + 1))} +# One capture bucket per decode size 1..MAX_NUM_SEQS. A single bucket at the +# max would pad every decode batch up to it, which wastes the compute the +# graphs just saved. +ATTN_GRAPH_ARGS=(--enforce-eager) +if [ "$ATTN_EAGER" != 1 ]; then + DECODE_SIZES=() + for size in $(seq 1 "$MAX_NUM_SEQS"); do + DECODE_SIZES+=("$size") + done + ATTN_GRAPH_ARGS=( + --max-cudagraph-capture-size "$MAX_NUM_SEQS" + --cudagraph-capture-sizes "${DECODE_SIZES[@]}" + --compilation-config '{"cudagraph_mode":"FULL_DECODE_ONLY"}' + ) +fi + AFD_CONFIG_ATTN='{ "afd": { "role": "attention", @@ -86,7 +125,7 @@ CUDA_VISIBLE_DEVICES="$ATTN_DEVICES" "${VLLM_CMD[@]}" serve "$MODEL_PATH" \ --kv-cache-dtype "$KV_CACHE_DTYPE" \ --api-server-count 1 \ --gpu-memory-utilization "$GPU_MEM_UTIL" \ - --enforce-eager \ + "${ATTN_GRAPH_ARGS[@]}" \ "${EXTRA_ARGS[@]}" \ --host 127.0.0.1 \ --port "$API_PORT" \ @@ -105,7 +144,7 @@ CUDA_VISIBLE_DEVICES="$FFN_DEVICES" "${VLLM_CMD[@]}" serve "$MODEL_PATH" \ --kv-cache-dtype "$KV_CACHE_DTYPE" \ --api-server-count 1 \ --gpu-memory-utilization "$GPU_MEM_UTIL" \ - --enforce-eager \ + "${FFN_GRAPH_ARGS[@]}" \ "${EXTRA_ARGS[@]}" \ --host 127.0.0.1 \ --port "$FFN_API_PORT" \ diff --git a/tests/e2e/async_gpu_ffn_padded_graph.py b/tests/e2e/async_gpu_ffn_padded_graph.py new file mode 100644 index 00000000..329fbf03 --- /dev/null +++ b/tests/e2e/async_gpu_ffn_padded_graph.py @@ -0,0 +1,143 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Replay the FFN grouped GEMM from a graph captured at a padded shape. + +Run with one GPU:: + + python tests/e2e/async_gpu_ffn_padded_graph.py + +The FFN side of the async connector never learns the next work item's shape +ahead of time -- there is no control plane -- which is why it ran eagerly. The +claim this checks is that it does not need to: a grouped GEMM takes its +grouping from a device-side count vector rather than from its row count, so one +row count can be captured and every smaller item padded up to it. + +What makes that non-obvious is that the padding changes the *grouping*, not +just the row count: the pad rows are charged to the last expert, so the count +vector differs on every replay. A graph that had baked its expert assignment in +would return the previous item's answer. Each replay here uses a different +routing and is checked against an eager run of the same routing. +""" + +from __future__ import annotations + +import os +import sys +from types import SimpleNamespace + +import torch + +from afd_plugin.model_executor.models.gpu.deepseek_v2_attention_gate import ( + compute_attention_gate_moe_ffn, +) +from afd_plugin.v1.worker.cuda_graph import pad_counts_to_shape + +HIDDEN = 128 +INTERMEDIATE = 256 +EXPERT_PER_RANK = 4 +MAX_ROWS = 24 +SCALING = 1.7 +NUM_REPLAYS = 4 + + +def build_layer(device): + gen = torch.Generator(device="cpu").manual_seed(5) + w13 = ( + torch.randn(EXPERT_PER_RANK, 2 * INTERMEDIATE, HIDDEN, generator=gen) + / HIDDEN**0.5 + ).to(device, torch.bfloat16) + w2 = ( + torch.randn(EXPERT_PER_RANK, HIDDEN, INTERMEDIATE, generator=gen) + / INTERMEDIATE**0.5 + ).to(device, torch.bfloat16) + return SimpleNamespace( + mlp=SimpleNamespace( + experts=SimpleNamespace( + routed_experts=SimpleNamespace(w13_weight=w13, w2_weight=w2), + _shared_experts=None, + routed_scaling_factor=SCALING, + ), + ), + ) + + +def routing_for(iteration: int, device): + """A different per-expert split, and a different row count, each time.""" + gen = torch.Generator(device="cpu").manual_seed(300 + iteration) + counts = torch.randint(0, 5, (EXPERT_PER_RANK,), generator=gen) + rows = int(counts.sum()) + if rows == 0: + counts[0] = 3 + rows = 3 + hidden = torch.randn(rows, HIDDEN, generator=gen).to(device, torch.bfloat16) + return hidden, counts.to(device, torch.int32), rows + + +def main() -> None: + device = torch.device("cuda", int(os.environ.get("LOCAL_RANK", 0))) + torch.cuda.set_device(device) + layer = build_layer(device) + + static_hidden = torch.zeros(MAX_ROWS, HIDDEN, dtype=torch.bfloat16, device=device) + static_counts = torch.zeros(EXPERT_PER_RANK, dtype=torch.int32, device=device) + + def padded_compute(): + return compute_attention_gate_moe_ffn( + layer, + hidden_states=static_hidden, + group_list=static_counts, + expand_x_shared=None, + ) + + # Warm first: the fused MoE picks a kernel on its first call, and that + # choice has to be settled before capture -- autotuning synchronizes. + static_counts[-1] = MAX_ROWS + warmup = torch.cuda.Stream(device=device) + warmup.wait_stream(torch.cuda.current_stream(device)) + with torch.cuda.stream(warmup): + padded_compute() + torch.cuda.current_stream(device).wait_stream(warmup) + torch.cuda.synchronize(device) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + payload = padded_compute() + print(f"captured the grouped GEMM at {MAX_ROWS} rows", flush=True) + + for replay in range(NUM_REPLAYS): + hidden, counts, rows = routing_for(replay, device) + + expected = compute_attention_gate_moe_ffn( + layer, + hidden_states=hidden, + group_list=counts, + expand_x_shared=None, + ).routed_output + + static_hidden[:rows].copy_(hidden) + static_counts.copy_(counts) + pad_counts_to_shape(static_counts, padded_rows=MAX_ROWS, actual_rows=rows) + assert int(static_counts.sum()) == MAX_ROWS + graph.replay() + torch.cuda.synchronize(device) + + got = payload.routed_output[:rows] + torch.testing.assert_close( + got.to(torch.float32), + expected.to(torch.float32), + rtol=2e-2, + atol=2e-2, + ) + print( + f"replay {replay}: rows={rows} split={counts.tolist()} matches eager", + flush=True, + ) + + print("PASS: padded FFN grouped GEMM replays correctly", flush=True) + + +if __name__ == "__main__": + if not torch.cuda.is_available(): + raise SystemExit("this test needs a GPU") + main() + sys.exit(0) diff --git a/tests/unit/v1/worker/test_cuda_graph.py b/tests/unit/v1/worker/test_cuda_graph.py index 36ac7195..c7edb50e 100644 --- a/tests/unit/v1/worker/test_cuda_graph.py +++ b/tests/unit/v1/worker/test_cuda_graph.py @@ -1,8 +1,12 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project + from __future__ import annotations from types import SimpleNamespace import pytest +import torch from afd_plugin.v1.worker.cuda_graph import ( FULL_DECODE_ONLY, @@ -10,6 +14,11 @@ cudagraph_mode_name, graph_run_mode, make_ffn_graph_key, + pad_counts_to_shape, + padded_ffn_graph_buckets, + padded_ffn_graph_shape, + select_padded_ffn_bucket, + shared_rows_for_bucket, validate_cuda_graph_mode, ) @@ -226,3 +235,174 @@ def test_graph_run_mode_requires_attention_replaying( ) is expected ) + + +# ---------------------------------------------------------------------- +# Padded FFN graph shape +# +# A grouped GEMM reads its grouping from a device-side count vector, not from +# its row count, so one row count can be captured and smaller items padded up +# to it. These pin what "big enough" means and where the padding lands. +# ---------------------------------------------------------------------- + + +def test_padded_shape_bounds_the_largest_batch(): + # Worst case for one FFN rank: every token sends every one of its topk + # slots here. Shared rows are split contiguously, so a rank holds a share. + routed, shared = padded_ffn_graph_shape( + num_tokens=8, + topk=6, + ffn_size=2, + has_shared_experts=True, + ) + assert routed == 48 + assert shared == 4 + + +def test_padded_shape_rounds_the_shared_split_up(): + # 7 tokens over 2 ranks is 4 and 3; the buffer has to hold the larger. + _, shared = padded_ffn_graph_shape( + num_tokens=7, + topk=2, + ffn_size=2, + has_shared_experts=True, + ) + assert shared == 4 + + +def test_padded_shape_has_no_shared_rows_without_shared_experts(): + routed, shared = padded_ffn_graph_shape( + num_tokens=8, + topk=6, + ffn_size=2, + has_shared_experts=False, + ) + assert routed == 48 + assert shared == 0 + + +@pytest.mark.parametrize( + ("num_tokens", "topk", "ffn_size"), + [(0, 6, 2), (8, 0, 2), (8, 6, 0)], +) +def test_padded_shape_rejects_nonpositive_inputs(num_tokens, topk, ffn_size): + with pytest.raises(ValueError): + padded_ffn_graph_shape( + num_tokens=num_tokens, + topk=topk, + ffn_size=ffn_size, + has_shared_experts=True, + ) + + +def test_padding_lands_on_the_last_expert(): + # Real rows are grouped by expert in ascending order, so the tail of the + # row range is the last expert's either way -- charging the padding there + # leaves every real row's expert assignment untouched. + counts = torch.tensor([3, 2, 1], dtype=torch.int32) + pad_counts_to_shape(counts, padded_rows=10, actual_rows=6) + assert counts.tolist() == [3, 2, 5] + assert int(counts.sum()) == 10 + + +def test_padding_an_exact_fit_changes_nothing(): + counts = torch.tensor([3, 2, 1], dtype=torch.int32) + pad_counts_to_shape(counts, padded_rows=6, actual_rows=6) + assert counts.tolist() == [3, 2, 1] + + +def test_padding_refuses_rows_that_do_not_fit(): + counts = torch.tensor([4, 4], dtype=torch.int32) + with pytest.raises(ValueError, match="do not fit"): + pad_counts_to_shape(counts, padded_rows=6, actual_rows=8) + + +def test_bucket_ladder_puts_exact_stops_on_the_common_item_sizes(): + # The sizes items actually cluster at: a whole item is max_routed/ffn_size, + # a DBO ubatch is half of that. A boundary sitting exactly on either sends + # every above-average item a full step up, which is what made the first + # bucketed run only 7% faster. + buckets = padded_ffn_graph_buckets(12288, ffn_size=2) + assert 6144 in buckets, "whole item size needs its own bucket" + assert 3072 in buckets, "DBO ubatch size needs its own bucket" + # And an item just above either pays a small step, not a doubling. + assert select_padded_ffn_bucket(buckets, 6200) / 6200 < 1.15 + assert select_padded_ffn_bucket(buckets, 3100) / 3100 < 1.15 + + +def test_bucket_ladder_scales_with_the_step_size(): + small = padded_ffn_graph_buckets(12288, ffn_size=2) + large = padded_ffn_graph_buckets(24576, ffn_size=2) + assert len(small) == len(large) + assert max(large) == 2 * max(small) + + +def test_oversized_item_has_no_bucket_and_runs_eager(): + buckets = padded_ffn_graph_buckets(12288, ffn_size=2) + assert select_padded_ffn_bucket(buckets, max(buckets) + 1) is None + + +def test_bucket_ladder_stays_within_the_row_bound(): + # The ladder no longer has to reach max_routed -- the workspace ceiling is + # pinned by an eager warm-up in capture_padded_ffn_graphs instead, which + # costs one forward rather than the largest graph per layer -- but no bucket + # may exceed the buffers. + for max_routed, ffn_size in ((12288, 2), (24576, 2), (4096, 4)): + buckets = padded_ffn_graph_buckets(max_routed, ffn_size=ffn_size) + assert max(buckets) <= max_routed + assert min(buckets) > 0 + + +def test_shared_rows_scale_with_the_bucket(): + # Both row counts scale with the item's tokens, so a half-sized bucket + # captures half the shared rows. Capturing every bucket at max_shared made + # a DBO ubatch pay double on the shared expert, which is what kept graphs + # losing to eager under DBO after the routed rows were already bucketed. + assert shared_rows_for_bucket(12288, max_routed=24576, max_shared=4096) == 2048 + assert shared_rows_for_bucket(24576, max_routed=24576, max_shared=4096) == 4096 + # No shared experts stays zero. + assert shared_rows_for_bucket(1024, max_routed=2048, max_shared=0) == 0 + + +def test_bucket_selection_escalates_when_the_shared_slice_does_not_fit(): + buckets = (12288, 24576) + kwargs = {"max_routed": 24576, "max_shared": 4096} + # Routed fits the small bucket and so does its share of the shared rows. + assert select_padded_ffn_bucket(buckets, 12288, 2048, **kwargs) == 12288 + # Same routed rows but more shared rows than the small bucket captured. + assert select_padded_ffn_bucket(buckets, 12288, 4000, **kwargs) == 24576 + # Beyond every bucket's shared capacity: eager. + assert select_padded_ffn_bucket(buckets, 12288, 9999, **kwargs) is None + + +def test_padded_graph_capture_stops_before_the_memory_reserve(): + # At 8192 tokens the default ladder took 65 GiB per V4 FFN rank and the + # NVSHMEM heap, set up after capture, had nothing left. Capture has to stop + # before a layer that would cut into the reserve. + from afd_plugin.v1.worker.cuda_graph import ( + padded_ffn_graph_layer_fits, + resolve_padded_ffn_graph_reserve_bytes, + ) + + gib = 2**30 + reserve = resolve_padded_ffn_graph_reserve_bytes({}) + assert reserve == 8 * gib + assert resolve_padded_ffn_graph_reserve_bytes( + {"AFD_FFN_GRAPH_MEM_RESERVE_GIB": "2.5"} + ) == int(2.5 * gib) + + # The first layer is taken whenever the reserve is already there. + assert padded_ffn_graph_layer_fits( + free_bytes=9 * gib, per_layer_bytes=0, reserve_bytes=reserve + ) + # A layer that would leave less than the reserve is not. + assert not padded_ffn_graph_layer_fits( + free_bytes=9 * gib, per_layer_bytes=2 * gib, reserve_bytes=reserve + ) + # Exactly the reserve left over is fine. + assert padded_ffn_graph_layer_fits( + free_bytes=10 * gib, per_layer_bytes=2 * gib, reserve_bytes=reserve + ) + + with pytest.raises(ValueError): + resolve_padded_ffn_graph_reserve_bytes({"AFD_FFN_GRAPH_MEM_RESERVE_GIB": "-1"}) diff --git a/tests/unit/v1/worker/test_ffn_model_runner.py b/tests/unit/v1/worker/test_ffn_model_runner.py index b7c1eff7..6b650eaa 100644 --- a/tests/unit/v1/worker/test_ffn_model_runner.py +++ b/tests/unit/v1/worker/test_ffn_model_runner.py @@ -28,7 +28,10 @@ from afd_plugin.model_executor.models.deepseek_v2 import ( # noqa: E402 AFDDeepseekV2ForCausalLM, ) -from afd_plugin.v1.worker.cuda_graph import make_ffn_graph_key # noqa: E402 +from afd_plugin.v1.worker.cuda_graph import ( # noqa: E402 + make_ffn_graph_key, + padded_ffn_graph_buckets, +) from afd_plugin.v1.worker.ffn_model_runner import ( # noqa: E402 GPUFFNModelRunner, _set_moe_layer_index, @@ -794,6 +797,7 @@ def test_ffn_worker_loop_logs_unexpected_thread_errors(caplog): worker._ffn_loop_error = None worker.model_runner = SimpleNamespace( connector=SimpleNamespace(is_initialized=True), + capture_padded_ffn_graphs=lambda: 0, ) expected_error = RuntimeError("boom") @@ -813,3 +817,288 @@ def fail_loop(): with pytest.raises(RuntimeError, match="AFD FFN worker loop failed") as exc: worker.raise_ffn_loop_error_if_any() assert exc.value.__cause__ is expected_error + + +# ---------------------------------------------------------------------- +# Connector-driven padded graphs +# +# The FFN side never learns the next work item's shape ahead of time, which is +# why it ran eagerly. Padding removes the need to: the grouping is device data, +# so one captured row count serves every smaller item. +# ---------------------------------------------------------------------- + + +class _RecordingPaddedGraph: + def __init__(self, routed_out, shared_out=None): + self.routed_out = routed_out + self.shared_out = shared_out + self.replays = 0 + + @property + def graph(self): + return self + + def replay(self): + self.replays += 1 + + +class _RecordingFFNModel: + """Stands in for the model; records eager compute calls.""" + + def __init__(self): + self.eager_calls = [] + + def compute_ffn_output(self, *, hidden_states, layer_idx, group_list, **kwargs): + self.eager_calls.append((layer_idx, int(hidden_states.shape[0]))) + return torch.zeros_like(hidden_states) + + +def _padded_runner(*, max_routed=8, max_shared=2, expert_per_rank=2, buckets=None): + runner = object.__new__(GPUFFNModelRunner) + runner.model = _RecordingFFNModel() + runner._padded_max_routed = max_routed + runner._padded_max_shared = max_shared + runner._padded_hidden = torch.zeros(max_routed, 4) + runner._padded_counts = torch.zeros(expert_per_rank, dtype=torch.int32) + runner._padded_shared = torch.zeros(max_shared, 4) if max_shared else None + runner._padded_graphs = {} + runner._padded_buckets = ( + buckets if buckets is not None else padded_ffn_graph_buckets(max_routed) + ) + runner._padded_rows_real = 0 + runner._padded_rows_charged = 0 + runner._padded_replays = 0 + runner._padded_eager_items = 0 + return runner + + +def _work_item(layer_idx, routed, shared, expert_per_rank=2, staged=False): + counts = torch.zeros(expert_per_rank, dtype=torch.int32) + counts[0] = routed + states = SimpleNamespace( + routed_tokens=routed, + shared_tokens=shared, + group_list=counts, + staged_routed=staged, + expand_x_shared=torch.ones(shared, 4) if shared else None, + ) + item = SimpleNamespace(layer_idx=layer_idx, hidden_states=torch.ones(routed, 4)) + return item, states + + +def test_padded_graph_replays_and_slices_off_the_padding(): + # 1000 rows into a 1024 bucket: 2.4% padding, inside the ratio that makes a + # replay worth more than the launches eager would pay. + runner = _padded_runner(max_routed=1024, max_shared=2, buckets=(1024,)) + graph = _RecordingPaddedGraph( + routed_out=torch.arange(1024 * 4, dtype=torch.float32).reshape(1024, 4), + shared_out=torch.zeros(2, 4), + ) + runner._padded_graphs[(3, runner._padded_buckets[-1])] = graph + item, states = _work_item(3, routed=1000, shared=1) + + payload = GPUFFNModelRunner._compute_work_item(runner, item, states) + + assert graph.replays == 1 + assert runner.model.eager_calls == [] + # The reply only ever sees the real rows. + assert payload.routed_output.shape[0] == 1000 + assert payload.shared_output.shape[0] == 1 + # Counts must sum to the captured row count, with the pad on the last + # expert -- otherwise the grouping and the row count disagree. + assert runner._padded_counts.tolist() == [1000, 24] + + +def test_padded_graph_stages_the_real_rows_at_the_front(): + runner = _padded_runner(max_routed=1024, max_shared=0, buckets=(1024,)) + runner._padded_graphs[(0, 1024)] = _RecordingPaddedGraph( + routed_out=torch.zeros(1024, 4) + ) + item, states = _work_item(0, routed=1000, shared=0) + item.hidden_states = torch.full((1000, 4), 7.0) + + GPUFFNModelRunner._compute_work_item(runner, item, states) + + assert torch.equal(runner._padded_hidden[:1000], torch.full((1000, 4), 7.0)) + + +def test_rows_already_gathered_into_the_buffer_are_not_copied_again(): + # The arrival's gather can write straight into the graph's input buffer, + # and it had to write somewhere regardless. Copying afterwards would be a + # second pass over the whole payload, once per layer -- which measured as + # the reason capturing was slower than running eagerly. + runner = _padded_runner(max_routed=1024, max_shared=0, buckets=(1024,)) + runner._padded_graphs[(0, 1024)] = _RecordingPaddedGraph( + routed_out=torch.zeros(1024, 4) + ) + runner._padded_hidden[:1000] = 7.0 + item, states = _work_item(0, routed=1000, shared=0, staged=True) + # What a stale copy would put there instead. + item.hidden_states = torch.full((1000, 4), -1.0) + + GPUFFNModelRunner._compute_work_item(runner, item, states) + + assert torch.equal(runner._padded_hidden[:1000], torch.full((1000, 4), 7.0)) + + +@pytest.mark.parametrize( + ("routed", "shared", "why"), + [ + (9, 0, "more routed rows than the capture"), + (4, 5, "more shared rows than the capture"), + (0, 0, "nothing routed here at all"), + ], +) +def test_work_that_does_not_fit_the_capture_runs_eagerly(routed, shared, why): + runner = _padded_runner(max_routed=8, max_shared=2) + graph = _RecordingPaddedGraph(routed_out=torch.zeros(8, 4)) + runner._padded_graphs[(1, runner._padded_buckets[-1])] = graph + item, states = _work_item(1, routed=routed, shared=shared) + + GPUFFNModelRunner._compute_work_item(runner, item, states) + + assert graph.replays == 0, why + assert runner.model.eager_calls == [(1, routed)], why + + +def test_a_layer_without_a_captured_graph_runs_eagerly(): + runner = _padded_runner() + item, states = _work_item(7, routed=4, shared=1) + + GPUFFNModelRunner._compute_work_item(runner, item, states) + + assert runner.model.eager_calls == [(7, 4)] + + +@pytest.mark.parametrize( + ("use_cuda_graph", "connector_driven"), + [(False, True), (True, False), (False, False)], +) +def test_capture_is_skipped_unless_graphs_and_connector_driven( + use_cuda_graph, + connector_driven, +): + # The padded path is for the connector-driven connector only. With a + # control plane the runner already has a shape-keyed graph cache, and with + # graphs off nothing should allocate. + runner = object.__new__(GPUFFNModelRunner) + runner.use_cuda_graph = use_cuda_graph + runner.is_connector_driven = connector_driven + runner._padded_graphs = {} + runner.model = None # would raise if capture got past the guard + + assert GPUFFNModelRunner.capture_padded_ffn_graphs(runner) == 0 + assert runner._padded_graphs == {} + + +def test_capture_does_not_run_twice(): + # start_ffn_server_loop is callable more than once; a second capture would + # leak the first set of graphs and their pool. + runner = object.__new__(GPUFFNModelRunner) + runner.use_cuda_graph = True + runner.is_connector_driven = True + runner._padded_graphs = {0: object()} + runner.model = None # would raise if capture got past the guard + + assert GPUFFNModelRunner.capture_padded_ffn_graphs(runner) == 0 + assert list(runner._padded_graphs) == [0] + + +def _ordering_worker(order, *, initialized): + worker = object.__new__(AFDFFNWorker) + worker._ffn_thread = None + worker._ffn_shutdown_event = None + worker._ffn_loop_error = None + worker.model_runner = SimpleNamespace( + connector=SimpleNamespace(is_initialized=initialized), + capture_padded_ffn_graphs=lambda: order.append("capture") or 0, + initialize_afd_connector=lambda: order.append("rendezvous"), + ) + return worker + + +def test_graphs_are_captured_before_the_connector_joins(): + # The connector's process group is the only barrier between the roles: the + # Attention rank profiles, and so dispatches, the moment it clears. Joining + # first and capturing after leaves those dispatches landing in a slot + # nobody is polling, which wedges the Attention rank mid-write with a flag + # unstamped -- fatal under ubatching, where the blocked thread never + # reaches its next yield. + order: list[str] = [] + worker = _ordering_worker(order, initialized=False) + worker._run_ffn_server_loop = lambda: None + + worker.start_ffn_server_loop() + worker._ffn_thread.join(timeout=5) + + assert order == ["capture", "rendezvous"] + + +def test_graphs_are_captured_before_the_serving_thread_starts(): + # Capture has to happen on an idle stream, and an AFD FFN EngineCore is a + # daemon that reaches start_ffn_server_loop by collective_rpc and never + # runs initialize_from_config -- so this is the only point that sees both + # entry paths. + order: list[str] = [] + worker = _ordering_worker(order, initialized=True) + started = threading.Event() + + def serve_loop(): + order.append("serve") + started.set() + + worker._run_ffn_server_loop = serve_loop + + worker.start_ffn_server_loop() + assert started.wait(timeout=5) + worker._ffn_thread.join(timeout=5) + + assert order == ["capture", "serve"] + + +def test_item_padded_far_past_its_bucket_runs_eagerly(): + # A replay costs its whole bucket, so an item far below one is cheaper run + # eagerly. Measured on DeepSeek-V4: eager beat a 1.18x-padded replay by 6%. + runner = _padded_runner(max_routed=1024, max_shared=0, buckets=(1024,)) + runner._padded_graphs[(0, 1024)] = _RecordingPaddedGraph( + routed_out=torch.zeros(1024, 4) + ) + item, states = _work_item(0, routed=600, shared=0) + + GPUFFNModelRunner._compute_work_item(runner, item, states) + + assert runner.model.eager_calls == [(0, 600)] + + +def test_small_item_replays_the_small_bucket_not_the_largest(): + # The regression this guards: one graph at the upper bound charged every + # item the full-batch row count, which is what made a DBO ubatch -- a + # quarter of the rows -- cost the same as a whole batch. + runner = _padded_runner(max_routed=4096, max_shared=0, buckets=(1024, 2048, 4096)) + graphs = { + bucket: _RecordingPaddedGraph(routed_out=torch.zeros(bucket, 4)) + for bucket in (1024, 2048, 4096) + } + for bucket, graph in graphs.items(): + runner._padded_graphs[(0, bucket)] = graph + item, states = _work_item(0, routed=2020, shared=0) + + payload = GPUFFNModelRunner._compute_work_item(runner, item, states) + + assert graphs[2048].replays == 1, "should take the smallest bucket that fits" + assert graphs[1024].replays == 0 + assert graphs[4096].replays == 0 + assert runner.model.eager_calls == [] + assert payload.routed_output.shape[0] == 2020 + # Counts sum to the chosen bucket, not to the largest one. + assert runner._padded_counts.tolist() == [2020, 28] + + +def test_item_larger_than_every_bucket_still_runs_eagerly(): + runner = _padded_runner(max_routed=8, max_shared=0, buckets=(2, 4, 8)) + runner._padded_graphs[(0, 8)] = _RecordingPaddedGraph(routed_out=torch.zeros(8, 4)) + item, states = _work_item(0, routed=9, shared=0) + + GPUFFNModelRunner._compute_work_item(runner, item, states) + + assert runner.model.eager_calls == [(0, 9)]