From c6800337427b4a68c329bf2e087fbdb371ce0f26 Mon Sep 17 00:00:00 2001 From: ronenkat <16743404+ronenkat@users.noreply.github.com> Date: Sun, 25 Oct 2026 12:46:30 +0200 Subject: [PATCH 01/32] feat(gpu): DeepEP M2N connector on the afdEP transport (SM90+) Add the DeepEP M2N AFD connector: Attention ranks route tokens and dispatch only the routed copies straight to the FFN ranks that own the chosen experts; FFN ranks run their local experts and combine sums each token's top-k contributions back on its Attention rank. - Transport: afdEP's DeepEP-derived ElasticBuffer (optional `afdep` extra, pinned to afdEP main). The former AfdEPConnector is merged into DeepEPM2NAFDConnector: `role_aware` (default true) runs afdEP's AFD role path, false keeps its Masked DeepEP protocol. No kernel sources are vendored in this repository. - Experts: `deepgemm` backend uses the separately installed upstream DeepGEMM package; `fused_experts` (vLLM Triton) is the reference path. - Plugin integration: routed-experts connector interface, FFN routed-experts model runner, connector device placement, Attention-side routing proxy for Qwen3 MoE and DeepSeek V2/V3, Attention TP split, DBO lanes and CUDA graph support. Co-Authored-By: Claude Opus 5.5 Signed-off-by: ronenkat <16743404+ronenkat@users.noreply.github.com> --- README.md | 8 +- afd_plugin/__init__.py | 2 + afd_plugin/config.py | 3 + afd_plugin/connectors/__init__.py | 8 + afd_plugin/connectors/factory.py | 50 +- .../connectors/gpu/deepep_m2n/__init__.py | 7 + .../connectors/gpu/deepep_m2n/connector.py | 1669 +++++++++++++++++ .../connectors/gpu/deepep_m2n/deepgemm.py | 92 + .../connectors/gpu/deepep_m2n/experts.py | 251 +++ .../connectors/gpu/deepep_m2n/fp8_quant.py | 162 ++ .../connectors/gpu/deepep_m2n/layout.py | 70 + .../connectors/gpu/deepep_m2n/network.py | 504 +++++ afd_plugin/connectors/metadata.py | 9 +- afd_plugin/connectors/routed_experts.py | 133 ++ afd_plugin/distributed/__init__.py | 2 + afd_plugin/distributed/topology.py | 38 +- .../model_executor/models/deepseek_v2.py | 110 +- afd_plugin/model_executor/models/qwen3_moe.py | 103 +- afd_plugin/v1/worker/__init__.py | 4 + .../v1/worker/attention_model_runner_v2.py | 22 +- afd_plugin/v1/worker/attention_worker.py | 29 +- afd_plugin/v1/worker/connector_device.py | 93 + afd_plugin/v1/worker/ffn_model_runner.py | 49 +- .../worker/ffn_routed_experts_model_runner.py | 193 ++ afd_plugin/v1/worker/ffn_worker.py | 45 +- afd_plugin/validation.py | 31 +- pyproject.toml | 5 + uv.lock | 133 +- 28 files changed, 3720 insertions(+), 105 deletions(-) create mode 100644 afd_plugin/connectors/gpu/deepep_m2n/__init__.py create mode 100644 afd_plugin/connectors/gpu/deepep_m2n/connector.py create mode 100644 afd_plugin/connectors/gpu/deepep_m2n/deepgemm.py create mode 100644 afd_plugin/connectors/gpu/deepep_m2n/experts.py create mode 100644 afd_plugin/connectors/gpu/deepep_m2n/fp8_quant.py create mode 100644 afd_plugin/connectors/gpu/deepep_m2n/layout.py create mode 100644 afd_plugin/connectors/gpu/deepep_m2n/network.py create mode 100644 afd_plugin/connectors/routed_experts.py create mode 100644 afd_plugin/v1/worker/connector_device.py create mode 100644 afd_plugin/v1/worker/ffn_routed_experts_model_runner.py diff --git a/README.md b/README.md index 95f7a9b4b..1223e42c1 100644 --- a/README.md +++ b/README.md @@ -47,7 +47,7 @@ Model support: | Model family | Registered architectures | Plugin model wrappers | Notes | | --- | --- | --- | --- | | DeepSeekV2 / DeepSeekV3 / DeepSeekV3.2 | `DeepseekForCausalLM`, `DeepseekV2ForCausalLM`, `DeepseekV3ForCausalLM`, `DeepseekV32ForCausalLM` | `AFDDeepseekForCausalLM`, `AFDDeepseekV2ForCausalLM`, `AFDDeepseekV3ForCausalLM` | DeepSeekV3.2 uses `AFDDeepseekV3ForCausalLM`. Each AFD role constructs and loads only its role-required model components, while shared embedding, normalization, and output components remain available where required by the model lifecycle. | -| Qwen3 MoE | `Qwen3MoeForCausalLM` | `AFDQwen3MoeForCausalLM` | CUDA with `compute_gate_on_attention=false`. | +| Qwen3 MoE | `Qwen3MoeForCausalLM` | `AFDQwen3MoeForCausalLM` | CUDA; DeepEP M2N routes on Attention, P2P routes on FFN. | | Qwen3.5 / Qwen3.6 MoE | `Qwen3_5MoeForConditionalGeneration` | `AFDQwen3_5MoeForConditionalGeneration` | CUDA text-only Qwen3.6-35B-A3B with `--language-model-only`, native DP4/TP1/EP4, and synchronous AFD 2A1F. Eager and graph execution remain supported; GDN DBO is temporarily withdrawn; native hybrid/multimodal validation applies. Earlier DBO experiments are historical evidence only; see the [E2E limits](docs/design/module/e2e_testing.md). | Connector support: @@ -57,6 +57,7 @@ See the [recipe index](recipe/README.md) for deployment and benchmark examples. | Connector | Platform | Recommend Stage | Sync or Async | Graph Support | Notes | | --- | --- | --- | --- | --- | --- | | `P2pNcclAFDConnector` | CUDA | Decode | Sync | `FULL_DECODE_ONLY` CUDA graph | FFN ranks are ordered before Attention ranks. `num_attention_ranks` must be greater than or equal to `num_ffn_ranks`; the Attention ranks are spread over the FFN ranks in blocks that differ in size by at most one. See the [DeepSeek V2 Lite recipe](recipe/gpu/P2pNcclAFDConnector/deepseek_v2_lite/README.md). | +| `DeepEPM2NAFDConnector` | CUDA | Decode | Sync | `FULL_DECODE_ONLY` CUDA graph | Expert-routed M2N transport: Attention ranks dispatch routed tokens straight to the FFN ranks that own the chosen experts. Dispatch and combine use [afdEP](https://github.com/ms-llmd/afdEP)'s DeepEP-derived `ElasticBuffer` (install the `afdep` extra); the default `deepgemm` expert backend needs [DeepGEMM](https://github.com/deepseek-ai/DeepGEMM) installed separately. Requires Hopper (SM90) or newer and Attention-side routing. | | `CAMP2pAFDConnector` | Ascend NPU | Decode | Sync | `FULL` or `FULL_DECODE_ONLY`; plain MLA graph DBO requires `FULL_DECODE_ONLY` | Uses HCCL/CAMP2P custom ops and native MRV2. DBO requires two ubatches and Attention DP > 1. See the [CAMP2P guide](docs/npu/CAM_P2P_CONNECTOR_USER_GUIDE.md). | | `CAMAsyncAFDConnector` | Ascend NPU | Prefill / decode | Async | FFN-only experimental `FULL` | Eager Attention supports DP+TP/SP and two native DBO microbatches, including DP1. FFN FULL requires the [W4A8 runtime contract](docs/npu/ASYNC_FFN_RUNTIME.md) and accepts eager Attention DBO; FFN does not join native DBO. PCP/DCP are unsupported. See the [Async CAM guide](docs/npu/CAM_ASYNC_CONNECTOR_USER_GUIDE.md) for configuration and validation limits. | @@ -88,8 +89,9 @@ Known gaps: and graph DBO; Async CAM enables eager DBO only on Attention, including when FFN uses FULL. FFN stays outside Attention's native DBO thread schedule. Model evidence remains specific to its workload and topology. -- Qwen3 MoE currently rejects Attention-side gate placement, sequence-parallel - MoE, EPLB, pipeline parallelism, speculative decoding, LoRA, and NPU. +- Qwen3 MoE requires Attention-side gate placement for DeepEP M2N and + FFN-side gate placement for P2P. It rejects sequence-parallel MoE, EPLB, + pipeline parallelism, speculative decoding, LoRA, and NPU. - Qwen3.5/Qwen3.6 MoE adapter-family support is limited to the CUDA text-only lane; repository hardware E2E evidence currently covers Qwen3.6-35B-A3B. GDN DBO is unsupported; upstream hybrid/multimodal DBO validation applies. diff --git a/afd_plugin/__init__.py b/afd_plugin/__init__.py index e281abaef..74f3c0cc8 100644 --- a/afd_plugin/__init__.py +++ b/afd_plugin/__init__.py @@ -22,6 +22,7 @@ def __getattr__(name: str): "AFDAttentionWorker", "AFDFFNWorker", "GPUFFNModelRunner", + "GPUFFNRoutedExpertsModelRunner", }: from afd_plugin.v1 import worker @@ -222,6 +223,7 @@ def register_afd() -> None: "AFDAttentionWorker", "AFDFFNWorker", "GPUFFNModelRunner", + "GPUFFNRoutedExpertsModelRunner", "assert_compatible_afd_stack", "parse_afd_config", "parse_optional_afd_config", diff --git a/afd_plugin/config.py b/afd_plugin/config.py index 2e004e572..ec1f93909 100644 --- a/afd_plugin/config.py +++ b/afd_plugin/config.py @@ -20,11 +20,14 @@ AFD_ADDITIONAL_CONFIG_KEY: Final[str] = "afd" AFD_ASYNC_CONNECTOR: Final[str] = "CAMAsyncAFDConnector" CAMP2P_CONNECTOR: Final[str] = "CAMP2pAFDConnector" +DEEPEP_M2N_CONNECTOR: Final[str] = "DeepEPM2NAFDConnector" +DEEPEP_M2N_CONNECTORS: Final[tuple[str, ...]] = (DEEPEP_M2N_CONNECTOR,) AFDRole = Literal["attention", "ffn"] SUPPORTED_AFD_ROLES: Final[tuple[str, ...]] = ("attention", "ffn") SUPPORTED_AFD_CONNECTORS: Final[tuple[str, ...]] = ( "P2pNcclAFDConnector", + DEEPEP_M2N_CONNECTOR, CAMP2P_CONNECTOR, AFD_ASYNC_CONNECTOR, ) diff --git a/afd_plugin/connectors/__init__.py b/afd_plugin/connectors/__init__.py index ed331c743..55fe53646 100644 --- a/afd_plugin/connectors/__init__.py +++ b/afd_plugin/connectors/__init__.py @@ -20,8 +20,14 @@ AFDTransferMetadata, AFDTransferState, ) +from afd_plugin.connectors.routed_experts import ( + AFDRoutedExpertsConnector, + FFNLocalExperts, + LocalExpertWeights, +) __all__ = [ + "AFDRoutedExpertsConnector", "AFDConnectorBase", "AFDControlPlane", "ConnectorExtraInfo", @@ -36,4 +42,6 @@ "AFDForwardContextMetadata", "AFDA2FTransferPayload", "AFDSingleDPMetadata", + "FFNLocalExperts", + "LocalExpertWeights", ] diff --git a/afd_plugin/connectors/factory.py b/afd_plugin/connectors/factory.py index 6e94e1ac1..420e09a5f 100644 --- a/afd_plugin/connectors/factory.py +++ b/afd_plugin/connectors/factory.py @@ -14,7 +14,8 @@ parse_afd_config, ) from afd_plugin.connectors.base import AFDConnectorBase, ConnectorExtraInfo -from afd_plugin.distributed import resolve_role_rank +from afd_plugin.connectors.routed_experts import AFDRoutedExpertsConnector +from afd_plugin.distributed import resolve_role_rank, tp_rank_before_distributed_init if TYPE_CHECKING: from vllm.config import VllmConfig @@ -53,12 +54,15 @@ def create_connector( local_rank: int, vllm_config: VllmConfig, afd_config: AFDConfig | None = None, + *, + tp_rank: int | None = None, ) -> AFDConnectorBase: + """Create a connector; ``tp_rank`` is passed to ``resolve_role_rank``.""" config = afd_config or parse_afd_config(vllm_config) if config.connector not in cls._registry: raise ValueError(f"unsupported AFD connector type: {config.connector}") connector_cls = cls._registry[config.connector]() - role_rank = resolve_role_rank(vllm_config, config) + role_rank = resolve_role_rank(vllm_config, config, tp_rank=tp_rank) return connector_cls( rank, local_rank, @@ -67,6 +71,43 @@ def create_connector( role_rank, ) + @classmethod + def create_bootstrapped_connector( + cls, + rank: int, + local_rank: int, + vllm_config: VllmConfig, + afd_config: AFDConfig, + ) -> AFDRoutedExpertsConnector | None: + """Create and bootstrap a routed-experts connector ahead of vLLM. + + Workers call this before vLLM's distributed init. Returns ``None`` for + connectors that are not ``AFDRoutedExpertsConnector``; their runner + creates them later, as before. vLLM's TP group does not exist yet, so + the TP part of the role rank comes from the worker ``rank`` + (``tp_rank_before_distributed_init``); workers check it with + ``check_connector_role_rank`` after vLLM's init. ``local_rank`` is the + CUDA device index the connector uses; vLLM has not selected the + worker's device yet, so callers pass the index it will select. + """ + connector_cls = cls.get_connector_class(afd_config.connector) + if not issubclass(connector_cls, AFDRoutedExpertsConnector): + return None + connector = cls.create_connector( + rank, + local_rank, + vllm_config, + afd_config, + tp_rank=tp_rank_before_distributed_init(vllm_config, rank), + ) + assert isinstance(connector, AFDRoutedExpertsConnector) + try: + connector.bootstrap() + except BaseException: + connector.close() + raise + return connector + @classmethod def get_connector_class(cls, connector_name: str) -> type[AFDConnectorBase]: if connector_name not in cls._registry: @@ -90,6 +131,11 @@ def parse_connector_extra_info( "afd_plugin.connectors.gpu.p2p", "P2pNcclAFDConnector", ) +AFDConnectorFactory.register_connector( + "DeepEPM2NAFDConnector", + "afd_plugin.connectors.gpu.deepep_m2n.connector", + "DeepEPM2NAFDConnector", +) AFDConnectorFactory.register_connector( "CAMP2pAFDConnector", "afd_plugin.connectors.npu.camp2p", diff --git a/afd_plugin/connectors/gpu/deepep_m2n/__init__.py b/afd_plugin/connectors/gpu/deepep_m2n/__init__.py new file mode 100644 index 000000000..2e3480e21 --- /dev/null +++ b/afd_plugin/connectors/gpu/deepep_m2n/__init__.py @@ -0,0 +1,7 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""DeepEP M2N connector package (afdEP transport + DeepGEMM experts). + +Import ``afd_plugin.connectors.gpu.deepep_m2n.connector`` for the connector; +this package stays import-light so registering it costs nothing. +""" diff --git a/afd_plugin/connectors/gpu/deepep_m2n/connector.py b/afd_plugin/connectors/gpu/deepep_m2n/connector.py new file mode 100644 index 000000000..85b4e099b --- /dev/null +++ b/afd_plugin/connectors/gpu/deepep_m2n/connector.py @@ -0,0 +1,1669 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""DeepEP M2N AFD connector: expert-routed Attention <-> FFN transport. + +Port of FastAFD's DeepEP M2N path (``3c71619``) to the plugin connector +contract. Routing runs on the Attention side: each Attention rank sends only +the routed copies of its tokens, straight to the FFN ranks that own the chosen +experts; FFN ranks run their local experts on what they receive and send the +weighted results back, where combine sums each token's top-k contributions. + +Topology: + One AFD world ordered ``[F0, F1, ..., A0, A1, ...]`` (as the P2P + connector). DeepEP sees it as a symmetric EP group in which FFN rank ``g`` + owns real experts ``[g * E/F, (g + 1) * E/F)`` and Attention ranks own + dummy expert slots (see ``layout.py``). Any Attention:FFN ratio works; + ``num_experts`` must be divisible by the FFN rank count. + +Data path (per MoE layer and stage, identical order on every rank): + Attention: ``send_attn_output(h, ctx, topk_ids=, topk_weights=)`` + dispatches; ``recv_ffn_output(ref, ubatch_idx)`` combines and returns + ``[tokens, hidden]``. + FFN: ``recv_attn_output(ubatch_idx)`` dispatches with no tokens and + returns the received rows; ``run_experts(payload, weights)`` runs the + local experts (the MoE FFN); ``send_ffn_output(out, ctx)`` combines. + Both calls are collectives over all ranks: every rank must issue them for + every MoE layer and stage, including ranks with no tokens. + +Attention tensor parallelism: + Every Attention TP rank is its own union rank. TP peers hold the same + tokens, so each dispatches only its slice: the model pads the DP rank's + ``N`` tokens to a multiple of the TP size, TP rank ``r`` dispatches rows + ``[r * ceil(N/TP), (r + 1) * ceil(N/TP))`` and the TP group all-gathers the + combined rows. Attention reports its TP size in the handshake; the FFN, + which runs without TP, adopts it to plan the per-rank token counts. + +Control plane: + The existing ``AFDControlPayload`` is enough. Every rank derives the same + per-stage plan from it (dispatch bucket, host sync, graph mode), so no + per-layer control traffic exists. Attention rank 0 sends each step's + payload to every FFN rank over the AFD world's NCCL group. + +Transport: + Dispatch and combine run on afdEP's DeepEP-derived ``ElasticBuffer`` + (https://github.com/ms-llmd/afdEP). By default it runs afdEP's AFD role + path: Attention ranks only send in dispatch and only receive in combine, + FFN ranks the reverse. It removes Masked DeepEP's entry barriers, union + count exchange and all-rank exits, and sizes buffers and kernels by role. + ``role_aware: false`` keeps afdEP's Masked DeepEP protocol for comparison. + +Scope: BF16 or FP8 (block-128, fp32 scales) payload; eager, ``torch.compile`` +(dispatch/combine are custom ops) and CUDA graphs; up to two lanes (DBO); +single node (multi-node needs RDMA pods). +""" + +from __future__ import annotations + +import itertools +import os +import time +from collections.abc import Mapping, Sequence +from dataclasses import asdict, dataclass, field, fields +from datetime import timedelta +from typing import TYPE_CHECKING, Any + +import torch +import torch.distributed as dist +from afdep.deepep import DispatchHandle, ElasticBuffer, ReceiveResult +from vllm.logger import init_logger +from vllm.model_executor.layers.quantization.fp8 import Fp8Config +from vllm.model_executor.layers.quantization.utils.fp8_utils import ( + per_token_group_quant_fp8, +) +from vllm.utils.torch_utils import direct_register_custom_op +from vllm.v1.worker.ubatching import dbo_current_ubatch_id, dbo_enabled + +from afd_plugin.config import AFDConfig +from afd_plugin.connectors.base import AFDControlPlane, ConnectorExtraInfo +from afd_plugin.connectors.gpu.deepep_m2n.experts import ( + EXPERT_BACKENDS, + DeepEPM2NExperts, + ExpertBackend, +) +from afd_plugin.connectors.gpu.deepep_m2n.layout import DeepEPM2NLayout +from afd_plugin.connectors.gpu.deepep_m2n.network import ( + NetworkConfig, + NetworkPlan, + RankTopology, + apply_rank_environment, + discover_local_topology, + plan_network, +) +from afd_plugin.connectors.metadata import ( + AFDA2FTransferPayload, + AFDControlPayload, + AFDDPMetadata, + AFDTransferContext, + AFDTransferMetadata, + AFDTransferState, + recv_control_payload, + send_control_payload, +) +from afd_plugin.connectors.routed_experts import ( + FP8_DTYPE, + AFDRoutedExpertsConnector, + FFNLocalExperts, + LocalExpertWeights, +) +from afd_plugin.distributed import init_afd_process_group + +if TYPE_CHECKING: + from vllm.config import VllmConfig + from vllm.model_executor.layers.quantization.base_config import ( + QuantizationConfig, + ) + +logger = init_logger(__name__) + +PROTOCOL_VERSION = "deepep_m2n/1" +BF16_PAYLOAD = "bf16" +# FP8 e4m3 rows with one fp32 scale per row and 128 channels (SM90 recipe). +FP8_PAYLOAD = "fp8_block128_fp32" +PAYLOAD_FORMATS = (BF16_PAYLOAD, FP8_PAYLOAD) +FP8_GROUP_SIZE = 128 +# Expert weight formats the FFN must load for each payload format. +BF16_WEIGHTS = "bf16" +FP8_BLOCK128_WEIGHTS = "fp8_block128" +PAYLOAD_WEIGHT_FORMATS = { + BF16_PAYLOAD: BF16_WEIGHTS, + FP8_PAYLOAD: FP8_BLOCK128_WEIGHTS, +} +# vLLM's all2all backend on the FFN. The connector already moved the rows, so +# vLLM's own MoE all-to-all never runs; this backend allocates no buffers and +# keeps vLLM's DeepEP out of the process. +IDLE_ALL2ALL_BACKEND = "allgather_reducescatter" +CONTROL_BACKENDS = ("nccl", "gloo") +# DeepEP V2 needs Hopper (SM90) or newer. +MIN_SM_MAJOR = 9 +PROCESS_GROUP_NAME = "afd_deepep_m2n" +# CPU (gloo) group for the handshake and DeepEP's communicator bootstrap, +# created before any NCCL communicator so the network plan can still set NCCL_*. +BOOTSTRAP_GROUP_NAME = "afd_deepep_m2n_bootstrap" +# FastAFD passes 120 s for both the CPU wait and the GPU spin timeouts. +DEFAULT_TIMEOUT_S = 120 +# Only Attention rank 0 sends the per-step control payload. +CONTROL_SENDER_ATTN_RANK = 0 +# Lanes = in-flight micro-batches (DBO uses two). +MAX_LANES = 2 +PER_BUFFER_COMM_STREAM_ENV = "AFD_DEEPEP_PER_BUFFER_COMM_STREAM" + + +@dataclass(frozen=True) +class DeepEPM2NConfig(ConnectorExtraInfo): + """``connector_extra_config`` for ``DeepEPM2NAFDConnector``. + + Attributes: + capacity_tokens_per_rank: Max tokens one Attention rank dispatches per + call (with Attention TP, its slice); sizes the DeepEP buffers. + ``None`` uses the larger of vLLM's ``max_num_batched_tokens`` and + its largest CUDA-graph capture size, divided by the Attention TP + size (``resolve_capacity_tokens_per_rank``). + num_lanes: In-flight micro-batches (one DeepEP buffer each). ``None`` + uses vLLM's micro-batch count: two under DBO, otherwise one + (``resolve_num_lanes``). + payload_format: Dispatch wire format. + expert_backend: ``deepgemm`` (fast) or ``fused_experts`` (reference). + network: Fabric settings (GIN, NIC choice, hybrid mode, traffic + class, extra NCCL env); see ``network.NetworkConfig`` and its + presets. + bucket_granularity: Dispatch bucket is rounded up to this multiple, + then up to the next rung of the bucket ladder (``bucket_ladder``) + or, above the ladder, to the capacity, so every step uses one of a + few DeepEP kernel sets compiled at startup. + num_sms: SMs DeepEP kernels use; ``0`` = a fixed fraction of the + device's SMs (24 on H100). + num_qps: RDMA queue pairs per DeepEP buffer; ``0`` = DeepEP's default. + expert_num_sms: SMs the FFN's DeepGEMM expert kernels use; ``0`` = all. + DeepGEMM splits its tiles statically over a persistent grid, so a + transport kernel of the other DBO lane holding even one SM delays + a share of every overlapping GEMM. Size the grid to the SMs that + transport leaves free while it waits. Must be even (DeepGEMM + launches 2-CTA clusters). + cpu_timeout_s / gpu_timeout_s: DeepEP host-wait and kernel-spin + timeouts. + cpu_sync_min_tokens: Eager steps whose bucket is below this skip the + host sync that reads receive counts; they allocate the worst-case + receive size instead. ``0`` always syncs in eager mode. Attention + ranks keep a zero buffer for the largest bucket that skips the + sync (the top rung of ``bucket_ladder``), so raising this costs + memory. + control_backend: Backend of the connector's own process group + (handshake, DeepEP communicator bootstrap, per-step control + payload). ``gloo`` keeps it off the GPUs and the RDMA NICs, so + DeepEP's communicator is the only RDMA user. + role_aware: Run afdEP's AFD role path (the default). ``False`` runs its + Masked DeepEP protocol. Every rank must use the same value; afdEP + rejects a mismatch when the buffers are created. + """ + + capacity_tokens_per_rank: int | None = None + num_lanes: int | None = None + payload_format: str = "bf16" + expert_backend: ExpertBackend = "deepgemm" + network: NetworkConfig = field(default_factory=NetworkConfig) + bucket_granularity: int = 1 + num_sms: int = 0 + num_qps: int = 0 + expert_num_sms: int = 0 + cpu_timeout_s: int = DEFAULT_TIMEOUT_S + gpu_timeout_s: int = DEFAULT_TIMEOUT_S + cpu_sync_min_tokens: int = 0 + control_backend: str = "nccl" + role_aware: bool = True + + @classmethod + def from_mapping(cls, raw: Mapping[str, Any] | None) -> DeepEPM2NConfig: + raw = dict(raw or {}) + known = {item.name for item in fields(cls)} + unknown = sorted(set(raw) - known) + if unknown: + raise ValueError(f"unknown DeepEPM2NAFDConnector options: {unknown}") + raw["network"] = NetworkConfig.from_mapping(raw.get("network")) + config = cls(**raw) + config.validate() + return config + + def validate(self) -> None: + if ( + self.capacity_tokens_per_rank is not None + and self.capacity_tokens_per_rank <= 0 + ): + raise ValueError("capacity_tokens_per_rank must be positive") + if self.num_lanes is not None and not 1 <= self.num_lanes <= MAX_LANES: + raise ValueError(f"num_lanes must be in [1, {MAX_LANES}]") + if self.payload_format not in PAYLOAD_FORMATS: + raise ValueError(f"payload_format must be one of {PAYLOAD_FORMATS}") + if self.expert_backend not in EXPERT_BACKENDS: + raise ValueError(f"expert_backend must be one of {EXPERT_BACKENDS}") + if self.bucket_granularity <= 0: + raise ValueError("bucket_granularity must be positive") + if self.num_sms < 0: + raise ValueError("num_sms must be non-negative") + if self.expert_num_sms < 0 or self.expert_num_sms % 2: + raise ValueError("expert_num_sms must be a non-negative even number") + if self.cpu_sync_min_tokens < 0: + raise ValueError("cpu_sync_min_tokens must be non-negative") + if self.num_qps < 0: + raise ValueError("num_qps must be non-negative") + if self.cpu_timeout_s <= 0 or self.gpu_timeout_s <= 0: + raise ValueError("cpu_timeout_s and gpu_timeout_s must be positive") + if self.control_backend not in CONTROL_BACKENDS: + raise ValueError(f"control_backend must be one of {CONTROL_BACKENDS}") + if type(self.role_aware) is not bool: + raise ValueError("role_aware must be a boolean") + + def to_mapping(self) -> dict[str, Any]: + return {**asdict(self), "network": self.network.to_mapping()} + + +# Attention derives the capacity from its own vLLM config and the FFN adopts +# it in the handshake (the FFN does not know the Attention TP size). The lane +# count is derived on both roles and must match, so Attention and FFN must +# share the DBO setting. +def resolve_capacity_tokens_per_rank( + config: DeepEPM2NConfig, vllm_config: VllmConfig +) -> int: + """Largest dispatch one Attention rank may issue: its TP slice of a prefill + chunk or of the largest decode graph, whichever is bigger.""" + if config.capacity_tokens_per_rank is not None: + return config.capacity_tokens_per_rank + max_tokens = max( + vllm_config.scheduler_config.max_num_batched_tokens, + vllm_config.compilation_config.max_cudagraph_capture_size, + ) + return -(-max_tokens // vllm_config.parallel_config.tensor_parallel_size) + + +def resolve_num_lanes(config: DeepEPM2NConfig, vllm_config: VllmConfig) -> int: + """One lane per vLLM micro-batch, so a DBO deployment gets two.""" + if config.num_lanes is not None: + return config.num_lanes + parallel_config = vllm_config.parallel_config + num_lanes = parallel_config.num_ubatches if parallel_config.use_ubatching else 1 + if num_lanes > MAX_LANES: + raise ValueError( + f"vLLM runs {num_lanes} micro-batches; DeepEP M2N supports at most " + f"{MAX_LANES} lanes" + ) + return num_lanes + + +@dataclass(frozen=True) +class StagePlan: + """Per-stage values every rank derives from the same control payload.""" + + tokens_per_attn_rank: tuple[int, ...] + bucket: int + do_cpu_sync: bool + graph_mode: bool + + +def tokens_per_attention_rank( + dp_metadata: AFDDPMetadata, attn_size: int, attn_tp_size: int +) -> tuple[int, ...]: + """Token count every Attention rank dispatches for one stage. + + DP metadata lists one count per DP rank. The TP peers of a DP rank hold + the same tokens and each dispatches a ``ceil(count / attn_tp_size)`` slice + (the last slices include padding rows that route nowhere). + """ + counts = [int(count) for count in dp_metadata.num_tokens_across_dp_cpu.flatten()] + if len(counts) * attn_tp_size != attn_size: + raise ValueError( + f"DP metadata has {len(counts)} token counts; with Attention TP " + f"{attn_tp_size} that does not cover {attn_size} Attention ranks" + ) + return tuple( + -(-count // attn_tp_size) for count in counts for _ in range(attn_tp_size) + ) + + +def plan_stage( + dp_metadata: AFDDPMetadata, + *, + attn_size: int, + attn_tp_size: int, + capacity_tokens_per_rank: int, + bucket_granularity: int, + graph_mode: bool, + bucket_ladder: Sequence[int], + cpu_sync_min_tokens: int = 0, +) -> StagePlan: + tokens = tokens_per_attention_rank(dp_metadata, attn_size, attn_tp_size) + max_tokens = max(1, max(tokens)) + bucket = -(-max_tokens // bucket_granularity) * bucket_granularity + if bucket > capacity_tokens_per_rank: + raise ValueError( + f"dispatch bucket {bucket} exceeds capacity_tokens_per_rank=" + f"{capacity_tokens_per_rank}" + ) + # Graph capture/replay cannot block on the host. Eager steps read the real + # receive counts so DeepEP allocates by actual rather than worst-case size, + # unless the bucket is small enough that skipping the host round trip wins. + do_cpu_sync = not graph_mode and bucket >= cpu_sync_min_tokens + # DeepEP compiles its dispatch and combine kernels per bucket (it is a + # template argument), so a bucket per token count runs nvcc on the request + # path. Buckets therefore come from a fixed set, compiled at startup + # (``precompile_buckets``). Above the ladder only synced eager steps + # remain, which size receives by the real counts: they all use the + # capacity. Smaller buckets may skip the sync (graph mode, or below + # ``cpu_sync_min_tokens``) and allocate the worst case for their bucket, + # so they round up only to the next rung: graph steps already sit on one + # (vLLM pads them to a capture size). The rule depends only on the token + # counts, because every rank must pass DeepEP the same bucket even when a + # rank plans the step in another mode (an FFN eager fallback). + if not bucket_ladder or bucket > bucket_ladder[-1]: + bucket = capacity_tokens_per_rank + else: + bucket = next(rung for rung in bucket_ladder if rung >= bucket) + return StagePlan( + tokens_per_attn_rank=tokens, + bucket=bucket, + do_cpu_sync=do_cpu_sync, + graph_mode=graph_mode, + ) + + +def no_sync_bucket_limit( + *, + capacity_tokens_per_rank: int, + max_cudagraph_capture_size: int, + attn_tp_size: int, + bucket_granularity: int, + cpu_sync_min_tokens: int, +) -> int: + """Largest bucket a stage can dispatch without the host sync. + + ``plan_stage`` skips the sync in graph mode, where a DP rank's tokens are + at most the largest CUDA-graph capture size and each Attention TP rank + dispatches its slice, and in eager mode below ``cpu_sync_min_tokens``. + Every other step syncs, and an Attention rank then receives no rows. + """ + graph_tokens = -(-max_cudagraph_capture_size // attn_tp_size) + graph_bucket = -(-graph_tokens // bucket_granularity) * bucket_granularity + return min(capacity_tokens_per_rank, max(graph_bucket, cpu_sync_min_tokens - 1)) + + +def bucket_ladder( + *, + capacity_tokens_per_rank: int, + cudagraph_capture_sizes: Sequence[int], + attn_tp_size: int, + bucket_granularity: int, + cpu_sync_min_tokens: int, +) -> tuple[int, ...]: + """Sorted buckets that steps at or below the no-sync limit round up to. + + The rungs are the per-rank CUDA-graph buckets (each capture size's + Attention TP slice, rounded to ``bucket_granularity``) plus the no-sync + limit as the top rung. Empty when no step can skip the host sync (no + graphs and ``cpu_sync_min_tokens`` 0): every eager step then uses the + capacity. + """ + limit = no_sync_bucket_limit( + capacity_tokens_per_rank=capacity_tokens_per_rank, + max_cudagraph_capture_size=max(cudagraph_capture_sizes, default=0), + attn_tp_size=attn_tp_size, + bucket_granularity=bucket_granularity, + cpu_sync_min_tokens=cpu_sync_min_tokens, + ) + if limit < 1: + return () + rungs = {limit} + for size in cudagraph_capture_sizes: + tokens = -(-size // attn_tp_size) + rung = -(-tokens // bucket_granularity) * bucket_granularity + if 1 <= rung <= limit: + rungs.add(rung) + return tuple(sorted(rungs)) + + +def precompile_buckets( + ladder: Sequence[int], capacity_tokens_per_rank: int, cpu_sync_min_tokens: int +) -> tuple[tuple[int, bool], ...]: + """``(bucket, do_cpu_sync)`` pairs ``plan_stage`` can give an eager step. + + A synced step lands on a rung at or above ``cpu_sync_min_tokens`` or on the + capacity; a step below the threshold skips the sync and rounds up to the + first rung that covers it. Graph steps also skip the sync, but CUDA-graph + capture compiles their kernels already. + """ + buckets = [(rung, True) for rung in ladder if rung >= cpu_sync_min_tokens] + buckets.append((capacity_tokens_per_rank, True)) + if cpu_sync_min_tokens > 0 and ladder: + covering = next( + (rung for rung in ladder if rung >= cpu_sync_min_tokens - 1), ladder[-1] + ) + buckets += [(rung, False) for rung in ladder if rung <= covering] + return tuple(dict.fromkeys(buckets)) + + +def worst_case_receive_rows( + *, + num_ranks: int, + bucket: int, + top_k: int, + experts_per_rank: int, + expert_alignment: int, +) -> int: + """Rows DeepEP allocates for one rank's receive when it skips the host sync. + + Mirrors DeepEP V2's worst case for expanded receives (``ElasticBuffer`` + non-cached dispatch without CPU sync): every rank's bucket times the routes + one token can send to this rank, plus alignment padding for each local + expert. + """ + rows = num_ranks * bucket * min(top_k, experts_per_rank) + rows += (expert_alignment - 1) * experts_per_rank + return -(-rows // expert_alignment) * expert_alignment + + +@dataclass(frozen=True) +class HandshakeResult: + expert_alignment: int + network_plan: NetworkPlan + # Reported by Attention; FFN ranks adopt them. + attn_tp_size: int + capacity_tokens_per_rank: int + bucket_ladder: tuple[int, ...] + + +def expert_weight_format(quant_config: QuantizationConfig | None) -> str: + """Expert weight format vLLM will load, named like ``payload_format``.""" + if quant_config is None: + return BF16_WEIGHTS + if isinstance(quant_config, Fp8Config) and quant_config.weight_block_size == [ + FP8_GROUP_SIZE, + FP8_GROUP_SIZE, + ]: + return FP8_BLOCK128_WEIGHTS + return f"unsupported ({quant_config.get_name()})" + + +def ffn_deployment_info(vllm_config: VllmConfig) -> dict[str, Any]: + """FFN-side vLLM settings that decide which experts each rank loads.""" + parallel_config = vllm_config.parallel_config + return { + "expert_parallel": parallel_config.enable_expert_parallel, + "ep_size": parallel_config.data_parallel_size + * parallel_config.tensor_parallel_size, + "tensor_parallel_size": parallel_config.tensor_parallel_size, + "expert_placement_strategy": parallel_config.expert_placement_strategy, + "eplb": parallel_config.enable_eplb, + "all2all_backend": parallel_config.all2all_backend, + "weight_format": expert_weight_format(vllm_config.quant_config), + } + + +def local_handshake_info( + *, + role: str, + role_rank: int, + world_rank: int, + config: DeepEPM2NConfig, + hidden_size: int, + dtype: torch.dtype, + num_experts: int, + top_k: int, + num_lanes: int, + capacity_tokens_per_rank: int | None, + bucket_ladder: Sequence[int] | None, + attn_tp_size: int | None, + compute_gate_on_attention: bool, + expert_alignment: int | None, + ffn_deployment: dict[str, Any] | None, + sm: tuple[int, int], + topology: RankTopology, +) -> dict[str, Any]: + return { + "protocol": PROTOCOL_VERSION, + "role": role, + "role_rank": role_rank, + "world_rank": world_rank, + "sm": list(sm), + "payload_format": config.payload_format, + "hidden_size": hidden_size, + "dtype": str(dtype), + "num_experts": num_experts, + "top_k": top_k, + "num_lanes": num_lanes, + "capacity_tokens_per_rank": capacity_tokens_per_rank, + "bucket_ladder": None if bucket_ladder is None else list(bucket_ladder), + "bucket_granularity": config.bucket_granularity, + "cpu_sync_min_tokens": config.cpu_sync_min_tokens, + "attn_tp_size": attn_tp_size, + "compute_gate_on_attention": compute_gate_on_attention, + "network": str(config.network.to_mapping()), + "expert_alignment": expert_alignment, + "ffn_deployment": ffn_deployment, + "topology": topology.to_mapping(), + } + + +_MUST_MATCH = ( + "protocol", + "payload_format", + "hidden_size", + "dtype", + "num_experts", + "top_k", + "num_lanes", + "compute_gate_on_attention", + "network", + "bucket_granularity", + "cpu_sync_min_tokens", +) + + +def _check_ffn_deployment( + world_rank: int, + deployment: dict[str, Any] | None, + layout: DeepEPM2NLayout, + payload_format: str, +) -> None: + """vLLM must place experts on FFN ranks exactly as DeepEP routes them. + + Otherwise rows reach a rank that holds other experts and the results are + silently wrong. Every rank runs this, so Attention fails with the FFN. + """ + if deployment is None: + raise RuntimeError(f"FFN world rank {world_rank} sent no deployment info") + expected = { + "expert_parallel": True, + "ep_size": layout.ffn_size, + "tensor_parallel_size": 1, + "expert_placement_strategy": "linear", + "eplb": False, + "all2all_backend": IDLE_ALL2ALL_BACKEND, + "weight_format": PAYLOAD_WEIGHT_FORMATS[payload_format], + } + mismatches = { + key: deployment[key] + for key, value in expected.items() + if deployment[key] != value + } + if mismatches: + required = {key: expected[key] for key in mismatches} + raise RuntimeError( + f"FFN world rank {world_rank} is deployed with {mismatches}; " + f"DeepEP M2N with payload_format={payload_format!r} needs {required}" + ) + + +def validate_handshake( + infos: list[dict[str, Any]], layout: DeepEPM2NLayout, network: NetworkConfig +) -> HandshakeResult: + """Fail fast, before any buffer allocation, if ranks disagree.""" + if len(infos) != layout.union_size: + raise RuntimeError( + f"handshake has {len(infos)} ranks, expected {layout.union_size}" + ) + for world_rank, info in enumerate(infos): + expected_role = "ffn" if world_rank < layout.ffn_size else "attention" + if info["world_rank"] != world_rank or info["role"] != expected_role: + raise RuntimeError( + f"world rank {world_rank} reported role={info['role']!r} " + f"world_rank={info['world_rank']}, expected {expected_role!r}" + ) + if info["sm"][0] < MIN_SM_MAJOR: + raise RuntimeError( + f"world rank {world_rank} has SM{info['sm'][0]}{info['sm'][1]}; " + f"DeepEP M2N needs SM{MIN_SM_MAJOR}0 or newer" + ) + for key in _MUST_MATCH: + values = {str(info[key]) for info in infos} + if len(values) != 1: + raise RuntimeError(f"ranks disagree on {key}: {sorted(values)}") + if not infos[0]["compute_gate_on_attention"]: + raise RuntimeError( + "DeepEP M2N routes tokens on Attention: set " + "compute_gate_on_attention=true on both roles" + ) + for world_rank, info in enumerate(infos[: layout.ffn_size]): + _check_ffn_deployment( + world_rank, info["ffn_deployment"], layout, info["payload_format"] + ) + alignments = {info["expert_alignment"] for info in infos[: layout.ffn_size]} + if len(alignments) != 1 or None in alignments: + raise RuntimeError(f"FFN ranks disagree on expert_alignment: {alignments}") + attn_infos = infos[layout.ffn_size :] + tp_sizes = {info["attn_tp_size"] for info in attn_infos} + if len(tp_sizes) != 1 or None in tp_sizes: + raise RuntimeError(f"Attention ranks disagree on their TP size: {tp_sizes}") + attn_tp_size = int(tp_sizes.pop()) + if attn_tp_size <= 0 or layout.attn_size % attn_tp_size != 0: + raise RuntimeError( + f"Attention TP size {attn_tp_size} does not divide the " + f"{layout.attn_size} Attention ranks" + ) + # Attention resolves the capacity (it knows its TP size); an FFN rank + # reports one only when it is configured explicitly, and it must agree. + capacities = {info["capacity_tokens_per_rank"] for info in attn_infos} + capacities |= { + info["capacity_tokens_per_rank"] + for info in infos[: layout.ffn_size] + if info["capacity_tokens_per_rank"] is not None + } + if len(capacities) != 1 or None in capacities: + raise RuntimeError( + "ranks disagree on capacity_tokens_per_rank: " + f"{sorted(map(str, capacities))}" + ) + # Attention resolves it from its CUDA-graph sizes; plan_stage needs the + # same ladder on every rank to pick the same bucket. + ladders = { + None if info["bucket_ladder"] is None else tuple(info["bucket_ladder"]) + for info in attn_infos + } + if len(ladders) != 1 or None in ladders: + raise RuntimeError( + f"Attention ranks disagree on bucket_ladder: {sorted(map(str, ladders))}" + ) + # Every rank evaluates the same plan on the same list. + plan = plan_network( + [RankTopology.from_mapping(info["topology"]) for info in infos], network + ) + return HandshakeResult( + expert_alignment=int(alignments.pop()), + network_plan=plan, + attn_tp_size=attn_tp_size, + capacity_tokens_per_rank=int(capacities.pop()), + bucket_ladder=ladders.pop() or (), + ) + + +def _describe_tensor(tensor: torch.Tensor | None) -> list[Any] | None: + return None if tensor is None else [str(tensor.dtype), list(tensor.shape)] + + +def local_model_info( + moe_layer_indices: Sequence[int] | None, + local_experts: Mapping[int, FFNLocalExperts] | None, +) -> dict[str, Any] | None: + """What this rank loaded, in a picklable form; ``None`` if not reported.""" + if moe_layer_indices is None: + return None + experts = None + if local_experts is not None: + experts = { + layer_idx: { + "ids": list(layer_experts.expert_ids), + "w13": _describe_tensor(layer_experts.weights.w13), + "w2": _describe_tensor(layer_experts.weights.w2), + "w13_scale": _describe_tensor(layer_experts.weights.w13_scale), + "w2_scale": _describe_tensor(layer_experts.weights.w2_scale), + } + for layer_idx, layer_experts in local_experts.items() + } + return {"moe_layers": list(moe_layer_indices), "experts": experts} + + +def _expected_expert_tensors( + num_experts: int, + hidden_size: int, + intermediate_size: int, + payload_format: str, +) -> dict[str, list[Any] | None]: + """dtype and shape of every ``LocalExpertWeights`` tensor.""" + if payload_format == BF16_PAYLOAD: + weight_dtype = str(torch.bfloat16) + w13_scale = w2_scale = None + else: + weight_dtype = str(FP8_DTYPE) + hidden_blocks = -(-hidden_size // FP8_GROUP_SIZE) + intermediate_blocks = -(-intermediate_size // FP8_GROUP_SIZE) + w13_scale = [ + str(torch.float32), + [num_experts, -(-2 * intermediate_size // FP8_GROUP_SIZE), hidden_blocks], + ] + w2_scale = [ + str(torch.float32), + [num_experts, hidden_blocks, intermediate_blocks], + ] + return { + "w13": [weight_dtype, [num_experts, 2 * intermediate_size, hidden_size]], + "w2": [weight_dtype, [num_experts, hidden_size, intermediate_size]], + "w13_scale": w13_scale, + "w2_scale": w2_scale, + } + + +def validate_model_infos( + infos: list[dict[str, Any] | None], + layout: DeepEPM2NLayout, + *, + hidden_size: int, + payload_format: str, +) -> None: + """After weights load, check what every rank actually holds. + + All ranks must run the same MoE layers (dispatch and combine are + collectives issued once per MoE layer), and every FFN rank must hold + exactly the experts DeepEP routes to it, in the payload's format. Ranks + that report nothing (``None``) are skipped; an FFN rank must report its + experts once any rank reports. + """ + reported = [ + (world_rank, info) for world_rank, info in enumerate(infos) if info is not None + ] + if not reported: + return + layer_lists = {str(info["moe_layers"]) for _, info in reported} + if len(layer_lists) != 1: + raise RuntimeError(f"ranks disagree on MoE layers: {sorted(layer_lists)}") + moe_layers = reported[0][1]["moe_layers"] + for world_rank in range(layout.ffn_size): + info = infos[world_rank] + if info is None or info["experts"] is None: + raise RuntimeError(f"FFN world rank {world_rank} reported no local experts") + experts = info["experts"] + if sorted(experts) != sorted(moe_layers): + raise RuntimeError( + f"FFN world rank {world_rank} holds experts for layers " + f"{sorted(experts)}, expected MoE layers {moe_layers}" + ) + expected_ids = list(layout.expert_range(world_rank)) + for layer_idx, described in experts.items(): + if described["ids"] != expected_ids: + raise RuntimeError( + f"FFN world rank {world_rank} layer {layer_idx} holds experts " + f"{described['ids']}, but DeepEP M2N routes experts " + f"[{expected_ids[0]}, {expected_ids[-1] + 1}) to it: use " + "linear expert placement with EP size = num_ffn_ranks" + ) + intermediate_size = described["w2"][1][2] + expected = _expected_expert_tensors( + layout.experts_per_rank, hidden_size, intermediate_size, payload_format + ) + wrong = { + name: described[name] + for name, value in expected.items() + if described[name] != value + } + if wrong: + required = {name: expected[name] for name in wrong} + raise RuntimeError( + f"FFN world rank {world_rank} layer {layer_idx} expert tensors " + f"{wrong} do not match payload_format={payload_format!r}, " + f"expected {required}" + ) + + +@dataclass(slots=True) +class DeepEPM2NTransferState(AFDTransferState): + """FFN-side state from ``recv_attn_output`` needed by experts and combine.""" + + handle: DispatchHandle + recv_topk_weights: torch.Tensor + recv_count_psum: torch.Tensor + recv_counts_host: list[int] + recv_scale: torch.Tensor | None + lane: int + # Recorded on the lane stream after dispatch; experts wait on it. + ready_event: torch.cuda.Event + # Recorded on the caller's stream after this lane's experts; the combine + # waits on it instead of on everything the caller queued. + experts_done: torch.cuda.Event | None = None + + +# Custom ops cannot take Python objects, so they find their connector by id +# (same pattern as the P2P connector's communicator registry). +_CONNECTORS: dict[int, DeepEPM2NAFDConnector] = {} +_CONNECTOR_IDS = itertools.count() +_DISPATCH_OP = "afd_deepep_m2n_dispatch" +_COMBINE_OP = "afd_deepep_m2n_combine" + + +def _live_stage_idx(traced_stage_idx: int) -> int: + """The stage of the calling micro-batch, read when the op body runs. + + ``torch.compile`` freezes the ``stage_idx`` op argument at trace time, and + vLLM's DBO micro-batch threads share one compiled graph, so every thread + would pass the same stage. Inside a DBO thread the vLLM micro-batch id is + the stage (one AFD stage per micro-batch); otherwise, as in direct eager + calls, the argument is right. + """ + return dbo_current_ubatch_id() if dbo_enabled() else traced_stage_idx + + +def _dispatch_op_impl( + hidden_states: torch.Tensor, + topk_ids: torch.Tensor, + topk_weights: torch.Tensor, + connector_id: int, + stage_idx: int, +) -> None: + _CONNECTORS[connector_id].dispatch_attention( + hidden_states, topk_ids, topk_weights, _live_stage_idx(stage_idx) + ) + + +def _dispatch_op_fake( + hidden_states: torch.Tensor, + topk_ids: torch.Tensor, + topk_weights: torch.Tensor, + connector_id: int, + stage_idx: int, +) -> None: + return None + + +def _combine_op_impl(out: torch.Tensor, connector_id: int, stage_idx: int) -> None: + _CONNECTORS[connector_id].combine_attention(out, _live_stage_idx(stage_idx)) + + +def _combine_op_fake(out: torch.Tensor, connector_id: int, stage_idx: int) -> None: + return None + + +def _register_custom_ops() -> None: + """Register the Attention-side ops once per process. + + Wrapping dispatch and combine in opaque custom ops keeps them intact under + ``torch.compile``. ``hidden_states`` / ``out`` are declared mutated so the + side-effecting ops are never dropped or reordered. + """ + if not hasattr(torch.ops.vllm, _DISPATCH_OP): + direct_register_custom_op( + op_name=_DISPATCH_OP, + op_func=_dispatch_op_impl, + mutates_args=["hidden_states"], + fake_impl=_dispatch_op_fake, + ) + if not hasattr(torch.ops.vllm, _COMBINE_OP): + direct_register_custom_op( + op_name=_COMBINE_OP, + op_func=_combine_op_impl, + mutates_args=["out"], + fake_impl=_combine_op_fake, + ) + + +def _keep_alive_on(stream: torch.cuda.Stream, *tensors: torch.Tensor | None) -> None: + """Tell the caching allocator ``tensors`` are used on ``stream``. + + Also while capturing a CUDA graph: within one capture the allocator reuses + a freed graph-pool block for the next allocation on the freeing stream, + which would overwrite data a lane still reads on replay (a micro-batch's + hidden states under its dispatch, combine output under the copy). PyTorch + tracks stream uses recorded during capture and keeps such blocks out of + reuse until the capture ends. + """ + for tensor in tensors: + if tensor is not None: + tensor.record_stream(stream) + + +class DeepEPM2NAFDConnector(AFDRoutedExpertsConnector): + """Expert-routed Attention <-> FFN connector over afdEP's DeepEP V2. + + Each stage (micro-batch) has a lane: its own DeepEP buffer, its own DeepEP + communication stream, and a lane stream on which the connector issues that + lane's dispatch and combine. Lane work joins the caller's stream only + where results are consumed (``recv_ffn_output``, ``run_experts``, + ``synchronize_lanes``), so two lanes can overlap communication with + compute (DBO order ``D0 D1 C0 C1``). With one lane everything runs on the + caller's stream, which avoids the fork/join cost in eager mode. + + CUDA graphs: capture with ``is_graph_capturing``/``is_graph_replaying`` + payloads (no host sync, fixed worst-case shapes) after one warm-up step at + the same bucket (it JIT-compiles the kernels and caches the Attention-side + zero buffer). FFN ranks call ``synchronize_lanes`` before the capture ends. + + FFN memory: a lane's received rows and its expert output live in that + lane's caching-allocator pool, without ``record_stream`` on the caller's + stream. The caller's stream uses them only in that lane's experts, which + the lane's combine waits for, and the lane's next allocation is used only + by work issued after that combine. So the lane can reuse a freed block at + once, also within one CUDA-graph capture: a capture holds about two + layers' rows per lane, not every layer's. + + Lifecycle: ``bootstrap`` (before vLLM's distributed init: handshake and + NCCL environment), then ``init_afd_connector`` (after weights load: + post-load checks, NCCL control group and DeepEP buffers). + """ + + @classmethod + def parse_extra_config(cls, raw: Mapping[str, Any] | None) -> DeepEPM2NConfig: + return DeepEPM2NConfig.from_mapping(raw) + + def __init__( + self, + rank: int, + local_rank: int, + vllm_config: VllmConfig, + afd_config: AFDConfig, + role_rank: int, + ) -> None: + super().__init__(rank, local_rank, vllm_config, afd_config, role_rank) + if not isinstance(self.extra_info, DeepEPM2NConfig): + raise TypeError("DeepEPM2NAFDConnector expects DeepEPM2NConfig extra info") + self.config = self.extra_info + text_config = vllm_config.model_config.hf_text_config + self.hidden_size: int = text_config.hidden_size + self.top_k: int = text_config.num_experts_per_tok + self.dtype: torch.dtype = vllm_config.model_config.dtype + if self.dtype != torch.bfloat16: + raise ValueError(f"DeepEP M2N needs a bfloat16 model, got {self.dtype}") + self.layout = DeepEPM2NLayout( + ffn_size=afd_config.num_ffn_ranks, + attn_size=afd_config.num_attention_ranks, + num_experts=text_config.num_experts, + ) + self.attn_size = self.layout.attn_size + self.ffn_size = self.layout.ffn_size + self.is_ffn = afd_config.role == "ffn" + self.is_fp8_payload = self.config.payload_format == FP8_PAYLOAD + self.payload_dtype = FP8_DTYPE if self.is_fp8_payload else self.dtype + if self.is_fp8_payload and self.hidden_size % FP8_GROUP_SIZE != 0: + raise ValueError(f"FP8 payload needs hidden_size % {FP8_GROUP_SIZE} == 0") + self.world_rank = self.layout.world_rank(afd_config.role, role_rank) + # The FFN runs without TP and learns the Attention TP size and the + # resulting capacity in the handshake (``bootstrap``). + self.attn_tp_size: int | None = ( + None if self.is_ffn else vllm_config.parallel_config.tensor_parallel_size + ) + self.capacity_tokens_per_rank: int | None = ( + self.config.capacity_tokens_per_rank + if self.is_ffn + else resolve_capacity_tokens_per_rank(self.config, vllm_config) + ) + # Buckets below the capacity (``plan_stage``); Attention derives them + # from its CUDA-graph sizes, the FFN adopts them in the handshake. + self.bucket_ladder: tuple[int, ...] | None = None + if not self.is_ffn: + assert self.attn_tp_size is not None + assert self.capacity_tokens_per_rank is not None + self.bucket_ladder = bucket_ladder( + capacity_tokens_per_rank=self.capacity_tokens_per_rank, + cudagraph_capture_sizes=( + vllm_config.compilation_config.cudagraph_capture_sizes or () + ), + attn_tp_size=self.attn_tp_size, + bucket_granularity=self.config.bucket_granularity, + cpu_sync_min_tokens=self.config.cpu_sync_min_tokens, + ) + self.num_lanes = resolve_num_lanes(self.config, vllm_config) + self.device = torch.device("cuda", local_rank) + # Where control payloads are staged: gloo moves CPU tensors. + self.control_device = ( + torch.device("cpu") + if self.config.control_backend == "gloo" + else self.device + ) + self._rendezvous = f"tcp://{afd_config.host}:{afd_config.port}" + self._timeout = timedelta(seconds=afd_config.afd_process_group_timeout_s) + self.group: dist.ProcessGroup | None = None + self.bootstrap_group: dist.ProcessGroup | None = None + self.network_plan: NetworkPlan | None = None + self.lanes: list[ElasticBuffer] = [] + self.lane_streams: list[torch.cuda.Stream] = [] + self.experts: DeepEPM2NExperts | None = None + self.expert_alignment = 1 + self.stage_plans: dict[int, StagePlan] = {} + # Attention side: dispatch handles waiting for their combine, by stage. + self._pending_combines: dict[int, DispatchHandle] = {} + # FFN side: combine completion per lane, joined by synchronize_lanes. + self._lane_done: dict[int, torch.cuda.Event] = {} + # FFN side: lanes join the caller's stream once per step, at an event + # recorded before any experts of the step (see _join_step). + self._step_fork: torch.cuda.Event | None = None + self._step_fork_capturing = False + self._lanes_joined: set[int] = set() + # Attention side: the zero rows combine takes from Attention ranks, + # sized for the largest receive (see _allocate_combine_zeros). + self._combine_zeros: torch.Tensor | None = None + self._op_id = next(_CONNECTOR_IDS) + self._initialized = False + self.control_plane = DeepEPM2NAFDControlPlane(self) + + # ============================== + # Lifecycle + # ============================== + + def bootstrap(self) -> None: + """Collective: join the AFD world, check compatibility, set NCCL env. + + Workers call this before vLLM's distributed init, because NCCL reads + its environment (``NCCL_IB_HCA``, the GIN switch) once per process, + when the first communicator is created. Uses gloo only, so it creates + no NCCL communicator itself. ``init_afd_connector`` calls it when the + caller has not. + """ + if self.bootstrap_group is not None: + return + torch.cuda.set_device(self.device) + self.bootstrap_group = init_afd_process_group( + backend="gloo", + init_method=self._rendezvous, + world_size=self.layout.union_size, + rank=self.world_rank, + group_name=BOOTSTRAP_GROUP_NAME, + timeout=self._timeout, + ) + local_alignment = None + if self.is_ffn: + local_alignment = ( + DeepEPM2NExperts.deepgemm_alignment() + if self.config.expert_backend == "deepgemm" + else 1 + ) + info = local_handshake_info( + role=self.afd_config.role, + role_rank=self.role_rank, + world_rank=self.world_rank, + config=self.config, + hidden_size=self.hidden_size, + dtype=self.dtype, + num_experts=self.layout.num_experts, + top_k=self.top_k, + num_lanes=self.num_lanes, + capacity_tokens_per_rank=self.capacity_tokens_per_rank, + bucket_ladder=self.bucket_ladder, + attn_tp_size=self.attn_tp_size, + compute_gate_on_attention=self.afd_config.compute_gate_on_attention, + expert_alignment=local_alignment, + ffn_deployment=( + ffn_deployment_info(self.vllm_config) if self.is_ffn else None + ), + sm=torch.cuda.get_device_capability(self.device), + topology=discover_local_topology( + self.world_rank, self.device, self.local_rank + ), + ) + infos: list[dict[str, Any] | None] = [None] * self.layout.union_size + dist.all_gather_object(infos, info, group=self.bootstrap_group) + result = validate_handshake( + [info for info in infos if info is not None], + self.layout, + self.config.network, + ) + self.expert_alignment = result.expert_alignment + self.network_plan = result.network_plan + self.attn_tp_size = result.attn_tp_size + self.capacity_tokens_per_rank = result.capacity_tokens_per_rank + self.bucket_ladder = result.bucket_ladder + if self.world_rank == 0: + logger.info("DeepEP M2N network plan: %s", self.network_plan.summary()) + # Before any NCCL communicator in this process exists, when the worker + # bootstraps ahead of vLLM's distributed init. + apply_rank_environment(self.network_plan, self.config.network, self.world_rank) + # DeepEP reads this when each buffer is created: with several lanes + # each buffer needs its own communication stream. + if self.num_lanes > 1: + os.environ[PER_BUFFER_COMM_STREAM_ENV] = "1" + + def init_afd_connector( + self, + *, + moe_layer_indices: Sequence[int] | None = None, + local_experts: Mapping[int, FFNLocalExperts] | None = None, + ) -> None: + """Collective: check what every rank loaded, then create buffers. + + Args: + moe_layer_indices: MoE layers this rank's model runs, in order. + ``None`` skips the post-load checks for this rank. + local_experts: FFN only: this rank's experts per MoE layer, as + vLLM loaded them. + """ + if self._initialized: + return + if local_experts is not None and not self.is_ffn: + raise ValueError("only FFN ranks hold local experts") + self.bootstrap() + torch.cuda.set_device(self.device) + _register_custom_ops() + assert self.bootstrap_group is not None and self.network_plan is not None + assert self.capacity_tokens_per_rank is not None + model_infos: list[dict[str, Any] | None] = [None] * self.layout.union_size + dist.all_gather_object( + model_infos, + local_model_info(moe_layer_indices, local_experts), + group=self.bootstrap_group, + ) + validate_model_infos( + model_infos, + self.layout, + hidden_size=self.hidden_size, + payload_format=self.config.payload_format, + ) + self.group = ( + self.bootstrap_group + if self.config.control_backend == "gloo" + else init_afd_process_group( + backend="nccl", + init_method=self._rendezvous, + world_size=self.layout.union_size, + rank=self.world_rank, + group_name=PROCESS_GROUP_NAME, + timeout=self._timeout, + ) + ) + try: + for _ in range(self.num_lanes): + self.lanes.append( + self._new_buffer( + group=self.bootstrap_group, + capacity_tokens_per_rank=self.capacity_tokens_per_rank, + hidden_size=self.hidden_size, + num_experts=self.layout.deepep_num_experts, + top_k=self.top_k, + allow_hybrid_mode=self.network_plan.hybrid, + use_fp8_dispatch=self.is_fp8_payload, + num_sms=self.config.num_sms, + num_qps=self.config.num_qps, + rdma_traffic_class=self.network_plan.rdma_traffic_class, + cpu_timeout_s=self.config.cpu_timeout_s, + gpu_timeout_s=self.config.gpu_timeout_s, + ) + ) + # One lane runs on the caller's stream: no fork/join cost. + if self.num_lanes > 1: + self.lane_streams.append(torch.cuda.Stream(self.device)) + if self.is_ffn: + self.experts = DeepEPM2NExperts( + self.config.expert_backend, + self.expert_alignment, + self.config.expert_num_sms, + ) + else: + self._allocate_combine_zeros() + self._precompile_dispatch_kernels() + except Exception: + self.close() + raise + _CONNECTORS[self._op_id] = self + self._initialized = True + + def _new_buffer(self, **kwargs: Any) -> ElasticBuffer: + """One lane's afdEP transport.""" + # The DeepEP M2N layout puts FFN ranks first, as afdEP's role path expects. + return ElasticBuffer( + **kwargs, role_aware=self.config.role_aware, ffn_ranks=self.ffn_size + ) + + def close(self) -> None: + """Destroy DeepEP buffers; safe after partial init and on repeat calls.""" + _CONNECTORS.pop(self._op_id, None) + for lane in self.lanes: + lane.destroy() + self.lanes = [] + self.lane_streams = [] + self._pending_combines.clear() + self._lane_done.clear() + self._reset_step_fork() + self._combine_zeros = None + self._initialized = False + + @property + def is_initialized(self) -> bool: + return self._initialized + + @property + def domain(self) -> tuple[int, int]: + """DeepEP ``(scale-out ranks, scale-up ranks)``: nodes and GPUs per node + in hybrid mode, ``(1, world)`` when everything is in one NVLink domain.""" + self._check_initialized() + return self.lanes[0].num_scaleout_ranks, self.lanes[0].num_scaleup_ranks + + # ============================== + # Attention side + # ============================== + + def send_attn_output( + self, + hidden_states: torch.Tensor, + context: AFDTransferContext, + *, + topk_ids: torch.Tensor | None = None, + topk_weights: torch.Tensor | None = None, + **kwargs: Any, + ) -> None: + """Dispatch each token to the FFN ranks owning its top-k experts. + + Args: + hidden_states: ``[tokens, hidden]`` BF16 Attention output. + context: Transfer context; ``metadata.stage_idx`` selects the lane + (inside a vLLM DBO micro-batch thread, its micro-batch id does). + topk_ids: ``[tokens, top_k]`` real expert ids. Required: routing + runs on the Attention side with this connector. + topk_weights: ``[tokens, top_k]`` routing weights. Required. + """ + self._check_initialized() + if topk_ids is None or topk_weights is None: + raise ValueError("DeepEPM2NAFDConnector needs topk_ids and topk_weights") + if not context.metadata.validate_tensor_shape(tuple(hidden_states.shape)): + raise ValueError( + f"hidden_states shape {tuple(hidden_states.shape)} does not match " + f"{context.metadata.total_tokens} tokens" + ) + torch.ops.vllm.afd_deepep_m2n_dispatch( + hidden_states, + topk_ids, + topk_weights, + self._op_id, + context.metadata.stage_idx, + ) + + def recv_ffn_output( + self, + ref_tensor: torch.Tensor, + ubatch_idx: int = 0, + **kwargs: Any, + ) -> torch.Tensor: + """Combine into ``ref_tensor``: ``sum_k weight_k * expert_k(token)``.""" + self._check_initialized() + torch.ops.vllm.afd_deepep_m2n_combine(ref_tensor, self._op_id, ubatch_idx) + return ref_tensor + + def dispatch_attention( + self, + hidden_states: torch.Tensor, + topk_ids: torch.Tensor, + topk_weights: torch.Tensor, + stage_idx: int, + ) -> None: + """Eager body of the dispatch custom op (runs on the stage's lane).""" + self._check_initialized() + if stage_idx in self._pending_combines: + raise RuntimeError(f"stage {stage_idx} dispatched twice without combine") + lane_stream = self._lane_stream(stage_idx) + self._fork_to(lane_stream) + with torch.cuda.stream(lane_stream): + x, x_scale = hidden_states, None + if self.is_fp8_payload: + x, x_scale = per_token_group_quant_fp8( + hidden_states, FP8_GROUP_SIZE, use_ue8m0=False + ) + result = self._dispatch(stage_idx, x, x_scale, topk_ids, topk_weights) + _keep_alive_on(lane_stream, hidden_states, topk_ids, topk_weights) + self._pending_combines[stage_idx] = result.handle + + def combine_attention(self, out: torch.Tensor, stage_idx: int) -> None: + """Eager body of the combine custom op: lane combine, then copy to out.""" + self._check_initialized() + handle = self._pending_combines.pop(stage_idx) + lane_stream = self._lane_stream(stage_idx) + with torch.cuda.stream(lane_stream): + # Attention ranks hold no experts; they contribute zero rows. + combined = self._lanes(stage_idx).combine( + self._combine_zero_rows(handle.recv_shape), handle + ) + if combined.shape != out.shape: + raise RuntimeError( + f"combined shape {tuple(combined.shape)} != out {tuple(out.shape)}" + ) + current = torch.cuda.current_stream(self.device) + if lane_stream != current: + current.wait_stream(lane_stream) + _keep_alive_on(current, combined) + out.copy_(combined) + + # ============================== + # FFN side + # ============================== + + def recv_attn_output( + self, + ubatch_idx: int = 0, + **kwargs: Any, + ) -> AFDA2FTransferPayload: + """Dispatch with no tokens; return the rows routed to local experts. + + The dispatch runs on the stage's lane; ``run_experts`` waits for it, so + a later lane's dispatch can overlap this lane's experts. The lane does + not wait for the caller's stream: its inputs are empty, and this + lane's previous combine, which waits for the experts that read the + lane's last rows, is ahead of it on the lane (see _join_step). The + same order lets the rows stay in the lane's allocator pool (see the + class docstring). + """ + self._check_initialized() + lane_stream = self._lane_stream(ubatch_idx) + self._join_step(ubatch_idx, lane_stream) + with torch.cuda.stream(lane_stream): + empty = torch.empty( + (0, self.hidden_size), dtype=self.payload_dtype, device=self.device + ) + empty_scale = ( + torch.empty( + (0, self.hidden_size // FP8_GROUP_SIZE), + dtype=torch.float32, + device=self.device, + ) + if self.is_fp8_payload + else None + ) + no_routes = torch.empty( + (0, self.top_k), dtype=torch.int64, device=self.device + ) + result = self._dispatch( + ubatch_idx, empty, empty_scale, no_routes, no_routes.float() + ) + ready_event = torch.cuda.Event() + ready_event.record(lane_stream) + state = DeepEPM2NTransferState( + handle=result.handle, + recv_topk_weights=result.recv_topk_weights, + recv_count_psum=result.recv_count_psum, + recv_counts_host=result.recv_counts_host, + recv_scale=result.recv_scale, + lane=ubatch_idx, + ready_event=ready_event, + ) + metadata = AFDTransferMetadata.create_ffn_metadata( + layer_idx=0, + stage_idx=ubatch_idx, + seq_lens=[result.handle.recv_shape[0]], + ) + return AFDA2FTransferPayload( + hidden_states=result.recv_x, + context=AFDTransferContext(metadata=metadata, states=state), + ) + + def run_experts( + self, payload: AFDA2FTransferPayload, weights: LocalExpertWeights + ) -> torch.Tensor: + """MoE FFN on received rows: top-k weighted local-expert outputs.""" + self._check_initialized() + if self.experts is None: + raise RuntimeError("run_experts is only available on FFN ranks") + state = self._transfer_state(payload.context) + current = torch.cuda.current_stream(self.device) + current.wait_event(state.ready_event) + # In the lane's pool, like the rows (see the class docstring). + with torch.cuda.stream(self._lane_stream(state.lane)): + out = torch.empty( + payload.hidden_states.shape, dtype=torch.bfloat16, device=self.device + ) + out = self.experts( + payload.hidden_states, + state.recv_scale, + state.recv_topk_weights, + state.recv_count_psum, + weights, + out=out, + ) + if self.lane_streams: + experts_done = torch.cuda.Event() + experts_done.record(current) + state.experts_done = experts_done + return out + + def send_ffn_output( + self, + ffn_output: torch.Tensor, + context: AFDTransferContext, + **kwargs: Any, + ) -> None: + """Combine on the lane: send each weighted row back to its token's rank. + + The lane waits only for this lane's experts (``run_experts``), not for + other lanes' experts queued after them on the caller's stream; under + DBO that wait would hold the other micro-batch's Attention. + """ + self._check_initialized() + state = self._transfer_state(context) + lane_stream = self._lane_stream(state.lane) + if state.experts_done is not None and lane_stream != torch.cuda.current_stream( + self.device + ): + lane_stream.wait_event(state.experts_done) + self._lanes_joined.add(state.lane) + else: + # Output not produced by run_experts: order after everything. + self._fork_to(lane_stream) + _keep_alive_on(lane_stream, ffn_output) + with torch.cuda.stream(lane_stream): + self._lanes(state.lane).combine(ffn_output, state.handle) + done = torch.cuda.Event() + done.record(lane_stream) + self._lane_done[state.lane] = done + + def synchronize_lanes(self) -> None: + """Make the caller's stream wait for every lane's outstanding combine. + + FFN ranks call this once per step (and before ending a CUDA-graph + capture, which requires all forked streams to rejoin). + """ + current = torch.cuda.current_stream(self.device) + for done in self._lane_done.values(): + current.wait_event(done) + self._lane_done.clear() + self._reset_step_fork() + + # ============================== + # Internals + # ============================== + + def _dispatch( + self, + stage_idx: int, + x: torch.Tensor, + x_scale: torch.Tensor | None, + topk_ids: torch.Tensor, + topk_weights: torch.Tensor, + ) -> ReceiveResult: + plan = self.stage_plans.get(stage_idx) + if plan is None: + raise RuntimeError( + f"no plan for stage {stage_idx}; call " + "control_plane.update_state_from_dp_metadata first" + ) + return self._lanes(stage_idx).dispatch( + x, + x_scale, + topk_ids, + topk_weights, + bucket=plan.bucket, + expert_alignment=self.expert_alignment, + do_cpu_sync=plan.do_cpu_sync, + ) + + def _lanes(self, stage_idx: int) -> ElasticBuffer: + if not 0 <= stage_idx < len(self.lanes): + raise RuntimeError(f"stage {stage_idx} has no lane ({len(self.lanes)})") + return self.lanes[stage_idx] + + def _lane_stream(self, stage_idx: int) -> torch.cuda.Stream: + self._lanes(stage_idx) + if not self.lane_streams: + return torch.cuda.current_stream(self.device) + return self.lane_streams[stage_idx] + + def _fork_to(self, lane_stream: torch.cuda.Stream) -> None: + """Order ``lane_stream`` after the caller's work (no-op for one lane).""" + current = torch.cuda.current_stream(self.device) + if lane_stream != current: + lane_stream.wait_stream(current) + + def _join_step(self, lane: int, lane_stream: torch.cuda.Stream) -> None: + """Order ``lane_stream`` after the caller's work from before this step. + + The first FFN lane operation of a step records one event on the + caller's stream, before any of the step's experts; each lane waits on + it once. This orders the lane after the previous step and, while + capturing, joins the lane to the capture. Later operations on the lane + are ordered by the lane itself. A change of capture mode starts over. + """ + current = torch.cuda.current_stream(self.device) + if lane_stream == current: + return + capturing = torch.cuda.is_current_stream_capturing() + if self._step_fork is not None and capturing != self._step_fork_capturing: + self._reset_step_fork() + if lane in self._lanes_joined: + return + if self._step_fork is None: + self._step_fork = torch.cuda.Event() + self._step_fork.record(current) + self._step_fork_capturing = capturing + lane_stream.wait_event(self._step_fork) + self._lanes_joined.add(lane) + + def _reset_step_fork(self) -> None: + self._step_fork = None + self._step_fork_capturing = False + self._lanes_joined.clear() + + def _precompile_dispatch_kernels(self) -> None: + """Compile every eager dispatch/combine variant before serving. + + Collective: every rank runs one empty dispatch and combine per lane + for each ``precompile_buckets`` pair, in the same order, so DeepEP's + JIT builds them now instead of on the first step of each size. Graph + buckets compile during CUDA-graph capture. + """ + assert self.bucket_ladder is not None + assert self.capacity_tokens_per_rank is not None + buckets = precompile_buckets( + self.bucket_ladder, + self.capacity_tokens_per_rank, + self.config.cpu_sync_min_tokens, + ) + x = torch.empty( + (0, self.hidden_size), dtype=self.payload_dtype, device=self.device + ) + x_scale = ( + torch.empty( + (0, self.hidden_size // FP8_GROUP_SIZE), + dtype=torch.float32, + device=self.device, + ) + if self.is_fp8_payload + else None + ) + topk_ids = torch.empty((0, self.top_k), dtype=torch.int64, device=self.device) + topk_weights = torch.empty( + (0, self.top_k), dtype=torch.float32, device=self.device + ) + start = time.perf_counter() + for lane in self.lanes: + for bucket, do_cpu_sync in buckets: + result = lane.dispatch( + x, + x_scale, + topk_ids, + topk_weights, + bucket=bucket, + expert_alignment=self.expert_alignment, + do_cpu_sync=do_cpu_sync, + ) + recv_shape = result.handle.recv_shape + expert_output = ( + torch.zeros(recv_shape, dtype=self.dtype, device=self.device) + if self.is_ffn + else self._combine_zero_rows(recv_shape) + ) + lane.combine(expert_output, result.handle) + torch.cuda.synchronize(self.device) + logger.info( + "DeepEP M2N precompiled %d dispatch buckets x %d lanes in %.1f s: %s", + len(buckets), + len(self.lanes), + time.perf_counter() - start, + buckets, + ) + + def _allocate_combine_zeros(self) -> None: + """One zero buffer for every combine on this Attention rank. + + Combine takes one zero row per row this rank received. A synced + dispatch gives an Attention rank no rows; a dispatch that skips the + sync gets DeepEP's worst case for its bucket, so the buffer covers the + largest such bucket. Both lanes share it (combine only reads it), and + captured graphs keep one stable address. + """ + assert self.bucket_ladder is not None + rows = worst_case_receive_rows( + num_ranks=self.layout.union_size, + bucket=self.bucket_ladder[-1] if self.bucket_ladder else 0, + top_k=self.top_k, + experts_per_rank=self.layout.experts_per_rank, + expert_alignment=self.expert_alignment, + ) + self._combine_zeros = torch.zeros( + (rows, self.hidden_size), dtype=self.dtype, device=self.device + ) + + def _combine_zero_rows(self, shape: tuple[int, int]) -> torch.Tensor: + """The first ``shape[0]`` rows of the zero buffer (a contiguous view).""" + zeros = self._combine_zeros + assert zeros is not None + rows = shape[0] + if rows > zeros.shape[0] or shape[1] != zeros.shape[1]: + raise RuntimeError( + f"combine needs zero rows of shape {shape} but the buffer holds " + f"{tuple(zeros.shape)}: a dispatch skipped the host sync with a " + "bucket above the ladder's top rung (no_sync_bucket_limit)" + ) + return zeros[:rows] + + @staticmethod + def _transfer_state(context: AFDTransferContext) -> DeepEPM2NTransferState: + if not isinstance(context.states, DeepEPM2NTransferState): + raise RuntimeError("context does not come from DeepEPM2N recv_attn_output") + return context.states + + def _check_initialized(self) -> None: + if not self._initialized: + raise RuntimeError("DeepEPM2NAFDConnector is not initialized") + + +class DeepEPM2NAFDControlPlane(AFDControlPlane): + """Per-step plan derivation plus payload delivery to every FFN rank.""" + + def __init__(self, connector: DeepEPM2NAFDConnector) -> None: + self.connector = connector + + def update_state_from_dp_metadata(self, payload: AFDControlPayload) -> None: + connector = self.connector + if len(payload.dp_metadata_list) > connector.num_lanes: + raise ValueError( + f"{len(payload.dp_metadata_list)} stages exceed " + f"num_lanes={connector.num_lanes}" + ) + graph_mode = payload.is_graph_capturing or payload.is_graph_replaying + if ( + connector.attn_tp_size is None + or connector.capacity_tokens_per_rank is None + or connector.bucket_ladder is None + ): + raise RuntimeError("DeepEPM2N control plane used before bootstrap") + connector.stage_plans = { + stage_idx: plan_stage( + dp_metadata, + attn_size=connector.attn_size, + attn_tp_size=connector.attn_tp_size, + capacity_tokens_per_rank=connector.capacity_tokens_per_rank, + bucket_granularity=connector.config.bucket_granularity, + graph_mode=graph_mode, + bucket_ladder=connector.bucket_ladder, + cpu_sync_min_tokens=connector.config.cpu_sync_min_tokens, + ) + for stage_idx, dp_metadata in payload.dp_metadata_list.items() + } + + def send_dp_metadata_list(self, payload: AFDControlPayload) -> None: + connector = self.connector + if connector.is_ffn or connector.role_rank != CONTROL_SENDER_ATTN_RANK: + return + if connector.group is None: + raise RuntimeError("DeepEPM2N control plane is not initialized") + send_control_payload( + payload, + dst=list(range(connector.ffn_size)), + group=connector.group, + device=connector.control_device, + ) + + def recv_dp_metadata_list(self) -> AFDControlPayload: + connector = self.connector + if connector.group is None: + raise RuntimeError("DeepEPM2N control plane is not initialized") + return recv_control_payload( + src=connector.layout.world_rank("attention", CONTROL_SENDER_ATTN_RANK), + group=connector.group, + device=connector.control_device, + ) diff --git a/afd_plugin/connectors/gpu/deepep_m2n/deepgemm.py b/afd_plugin/connectors/gpu/deepep_m2n/deepgemm.py new file mode 100644 index 000000000..ed9acd259 --- /dev/null +++ b/afd_plugin/connectors/gpu/deepep_m2n/deepgemm.py @@ -0,0 +1,92 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""DeepGEMM grouped GEMMs the M2N experts use, from the ``deep_gemm`` package. + +Uses the installed upstream DeepGEMM (https://github.com/deepseek-ai/DeepGEMM), +installed separately: its build needs torch, CUDA and git submodules. Its +kernels are JIT-compiled at runtime. Only the ``deepgemm`` expert backend +needs it. + +Only the psum ("prefix-sum contiguous") layout is used: DeepEP's expanded +dispatch already groups received rows per local expert, padded to the +DeepGEMM M alignment, and its per-expert end offsets are the group boundaries. +""" + +from __future__ import annotations + +import functools +from types import ModuleType + +import torch + + +@functools.cache +def load_deepgemm() -> ModuleType: + """DeepGEMM's native module; importing ``deep_gemm`` initializes its JIT.""" + # Deferred: the fused_experts backend runs without DeepGEMM installed. + import deep_gemm + + return deep_gemm._C + + +def contiguous_m_alignment() -> int: + """Row alignment DeepGEMM requires between experts in the contiguous layout.""" + return int(load_deepgemm().get_theoretical_mk_alignment_for_contiguous_layout()) + + +def set_contiguous_m_alignment(alignment: int) -> None: + load_deepgemm().set_mk_alignment_for_contiguous_layout(alignment) + + +def set_num_sms(num_sms: int) -> None: + """Size DeepGEMM's persistent grid (process-wide) to ``num_sms`` SMs.""" + load_deepgemm().set_num_sms(num_sms) + + +def m_grouped_bf16_gemm_nt_psum( + a: torch.Tensor, + b: torch.Tensor, + out: torch.Tensor, + psum_rows_per_expert: torch.Tensor, + expected_m: int, +) -> None: + """``out[rows of e] = a[rows of e] @ b[e].T`` for every local expert ``e``. + + ``a``: ``[M, K]`` BF16, rows grouped per expert. ``b``: ``[E, N, K]`` BF16. + ``out``: ``[M, N]`` BF16. ``psum_rows_per_expert``: ``[E]`` int32 end offset + of each expert's rows (DeepEP's expanded-dispatch prefix sum). + """ + load_deepgemm().m_grouped_bf16_gemm_nt_contiguous( + a, + b, + out, + psum_rows_per_expert, + use_psum_layout=True, + expected_m_for_psum_layout=expected_m, + ) + + +def m_grouped_fp8_gemm_nt_psum( + a: torch.Tensor, + a_scale: torch.Tensor, + b: torch.Tensor, + b_scale: torch.Tensor, + out: torch.Tensor, + psum_rows_per_expert: torch.Tensor, + expected_m: int, +) -> None: + """FP8 version of ``m_grouped_bf16_gemm_nt_psum`` (SM90 "1d2d" recipe). + + ``a``: ``[M, K]`` FP8 with ``a_scale`` ``[M, K/128]`` fp32 (one scale per row + and 128 channels; MN-major TMA-aligned avoids a transpose). ``b``: + ``[E, N, K]`` FP8 with ``b_scale`` ``[E, N/128, K/128]`` fp32 block scales + (vLLM's block-FP8 layout). ``out``: ``[M, N]`` BF16. + """ + load_deepgemm().m_grouped_fp8_gemm_nt_contiguous( + (a, a_scale), + (b, b_scale), + out, + psum_rows_per_expert, + use_psum_layout=True, + expected_m_for_psum_layout=expected_m, + ) diff --git a/afd_plugin/connectors/gpu/deepep_m2n/experts.py b/afd_plugin/connectors/gpu/deepep_m2n/experts.py new file mode 100644 index 000000000..102e89120 --- /dev/null +++ b/afd_plugin/connectors/gpu/deepep_m2n/experts.py @@ -0,0 +1,251 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""MoE FFN on DeepEP-received rows: this FFN rank's local experts only. + +DeepEP's expanded dispatch hands the FFN rank one row per (token, local +expert) pair, grouped per expert with each group starting at a multiple of +``expert_alignment``, plus one top-k weight per row. The executor returns +``weight * expert(row)`` in BF16 for every row in the received shape; DeepEP's +combine then sums each token's rows back on its Attention rank. Padding rows +between expert groups are never read by combine, so their content is +irrelevant. + +Payloads: + +- BF16 rows with BF16 weights. +- FP8 rows (per-row, per-128-channel fp32 scales) with FP8 block-128 weights + (vLLM's block-FP8 ``FusedMoE`` layout). + +Backends: + +- ``deepgemm``: FastAFD's path. DeepGEMM grouped GEMMs consume DeepEP's + per-expert prefix sum directly (no re-permutation). BF16: GEMM, vLLM + ``silu_and_mul``, GEMM, weight multiply. FP8: GEMM, fused SiLU-mul with the + top-k weight folded in and FP8 requant (``fp8_quant.py``), GEMM. Requires + ``expert_alignment`` equal to the DeepGEMM contiguous-layout M alignment. +- ``fused_experts``: vLLM's Triton ``fused_experts`` with ``top_k=1`` on the + valid rows (FP8 inputs and weights are dequantized to BF16 first). Reads the + prefix sum on the host: a correctness reference, not a fast path. +""" + +from __future__ import annotations + +from typing import Literal + +import torch +import vllm._custom_ops # noqa: F401 (registers torch.ops._C.silu_and_mul) +from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts + +from afd_plugin.connectors.gpu.deepep_m2n import deepgemm +from afd_plugin.connectors.gpu.deepep_m2n.fp8_quant import ( + dequantize_fp8_block128, + dequantize_fp8_rows, + silu_mul_fp8_quant, +) +from afd_plugin.connectors.routed_experts import FP8_DTYPE, LocalExpertWeights + +ExpertBackend = Literal["deepgemm", "fused_experts"] +EXPERT_BACKENDS: tuple[ExpertBackend, ...] = ("deepgemm", "fused_experts") + + +def align_up(value: int, alignment: int) -> int: + return (value + alignment - 1) // alignment * alignment + + +def expert_row_ranges( + recv_count_psum: torch.Tensor, expert_alignment: int +) -> list[tuple[int, int]]: + """Host-side ``[start, end)`` row range of each local expert's group.""" + ranges: list[tuple[int, int]] = [] + previous_end = 0 + for end in recv_count_psum.tolist(): + start = align_up(previous_end, expert_alignment) + ranges.append((start, int(end))) + previous_end = int(end) + return ranges + + +class DeepEPM2NExperts: + """Runs local experts on received rows; output rows are top-k weighted.""" + + def __init__( + self, backend: ExpertBackend, expert_alignment: int, num_sms: int = 0 + ) -> None: + if backend not in EXPERT_BACKENDS: + raise ValueError(f"expert backend must be one of {EXPERT_BACKENDS}") + self.backend = backend + self.expert_alignment = expert_alignment + if backend == "deepgemm": + deepgemm.set_contiguous_m_alignment(expert_alignment) + if num_sms: + device_sms = torch.cuda.get_device_properties( + torch.cuda.current_device() + ).multi_processor_count + if num_sms > device_sms: + raise ValueError( + f"expert_num_sms={num_sms} exceeds the device's " + f"{device_sms} SMs" + ) + deepgemm.set_num_sms(num_sms) + + @staticmethod + def deepgemm_alignment() -> int: + """Row alignment DeepEP must use so DeepGEMM can read its layout.""" + return deepgemm.contiguous_m_alignment() + + def __call__( + self, + recv_x: torch.Tensor, + recv_scale: torch.Tensor | None, + recv_topk_weights: torch.Tensor, + recv_count_psum: torch.Tensor, + weights: LocalExpertWeights, + out: torch.Tensor | None = None, + ) -> torch.Tensor: + """``out``, if given, receives the rows (bf16, ``recv_x``'s shape); the + caller then picks the allocator pool (stream) they live in.""" + if recv_count_psum.shape[0] != weights.num_experts: + raise ValueError( + f"received rows for {recv_count_psum.shape[0]} local experts, " + f"weights hold {weights.num_experts}" + ) + is_fp8_payload = recv_x.dtype == FP8_DTYPE + if is_fp8_payload != weights.is_fp8 or is_fp8_payload != ( + recv_scale is not None + ): + raise ValueError( + f"payload dtype {recv_x.dtype} does not match the expert weights " + f"({weights.w13.dtype}) or the presence of scales" + ) + if out is None: + out = torch.empty(recv_x.shape, dtype=torch.bfloat16, device=recv_x.device) + elif out.shape != recv_x.shape or out.dtype != torch.bfloat16: + raise ValueError( + f"expert output must be bf16 {tuple(recv_x.shape)}, got " + f"{out.dtype} {tuple(out.shape)}" + ) + if recv_x.shape[0] == 0: + return out + if self.backend == "fused_experts": + return self._run_fused_experts( + recv_x, recv_scale, recv_topk_weights, recv_count_psum, weights, out + ) + if recv_scale is not None: + return self._run_deepgemm_fp8( + recv_x, recv_scale, recv_topk_weights, recv_count_psum, weights, out + ) + return self._run_deepgemm_bf16( + recv_x, recv_topk_weights, recv_count_psum, weights, out + ) + + def _run_deepgemm_bf16( + self, + recv_x: torch.Tensor, + recv_topk_weights: torch.Tensor, + recv_count_psum: torch.Tensor, + weights: LocalExpertWeights, + out: torch.Tensor, + ) -> torch.Tensor: + num_rows = recv_x.shape[0] + intermediate_x2 = weights.w13.shape[1] + expected_m = max(1, num_rows // weights.num_experts) + gate_up = torch.empty( + (num_rows, intermediate_x2), dtype=torch.bfloat16, device=recv_x.device + ) + deepgemm.m_grouped_bf16_gemm_nt_psum( + recv_x, weights.w13, gate_up, recv_count_psum, expected_m + ) + activated = torch.empty( + (num_rows, intermediate_x2 // 2), dtype=torch.bfloat16, device=recv_x.device + ) + torch.ops._C.silu_and_mul(activated, gate_up) + deepgemm.m_grouped_bf16_gemm_nt_psum( + activated, weights.w2, out, recv_count_psum, expected_m + ) + # Combine is an unweighted sum, so apply each row's top-k weight here. + out.mul_(recv_topk_weights[:num_rows].to(torch.bfloat16).view(-1, 1)) + return out + + def _run_deepgemm_fp8( + self, + recv_x: torch.Tensor, + recv_scale: torch.Tensor, + recv_topk_weights: torch.Tensor, + recv_count_psum: torch.Tensor, + weights: LocalExpertWeights, + out: torch.Tensor, + ) -> torch.Tensor: + assert weights.w13_scale is not None and weights.w2_scale is not None + num_rows = recv_x.shape[0] + expected_m = max(1, num_rows // weights.num_experts) + gate_up = torch.empty( + (num_rows, weights.w13.shape[1]), dtype=torch.bfloat16, device=recv_x.device + ) + deepgemm.m_grouped_fp8_gemm_nt_psum( + recv_x, + recv_scale, + weights.w13, + weights.w13_scale, + gate_up, + recv_count_psum, + expected_m, + ) + # Fold the top-k weight into the second GEMM's (linear) input. + activated, activated_scale = silu_mul_fp8_quant( + gate_up, recv_count_psum, recv_topk_weights + ) + deepgemm.m_grouped_fp8_gemm_nt_psum( + activated, + activated_scale, + weights.w2, + weights.w2_scale, + out, + recv_count_psum, + expected_m, + ) + return out + + def _run_fused_experts( + self, + recv_x: torch.Tensor, + recv_scale: torch.Tensor | None, + recv_topk_weights: torch.Tensor, + recv_count_psum: torch.Tensor, + weights: LocalExpertWeights, + out: torch.Tensor, + ) -> torch.Tensor: + out.zero_() + ranges = expert_row_ranges(recv_count_psum, self.expert_alignment) + rows = [ + torch.arange(start, end, device=recv_x.device) + for start, end in ranges + if end > start + ] + if not rows: + return out + row_idx = torch.cat(rows) + local_ids = torch.cat( + [ + torch.full( + (end - start,), expert, dtype=torch.int32, device=recv_x.device + ) + for expert, (start, end) in enumerate(ranges) + if end > start + ] + ) + x = recv_x[row_idx] + w13, w2 = weights.w13, weights.w2 + if recv_scale is not None: + assert weights.w13_scale is not None and weights.w2_scale is not None + x = dequantize_fp8_rows(x, recv_scale[row_idx]).to(torch.bfloat16) + w13 = dequantize_fp8_block128(w13, weights.w13_scale).to(torch.bfloat16) + w2 = dequantize_fp8_block128(w2, weights.w2_scale).to(torch.bfloat16) + out[row_idx] = fused_experts( + x.contiguous(), + w13, + w2, + topk_weights=recv_topk_weights[row_idx].float().view(-1, 1), + topk_ids=local_ids.view(-1, 1), + global_num_experts=weights.num_experts, + ) + return out diff --git a/afd_plugin/connectors/gpu/deepep_m2n/fp8_quant.py b/afd_plugin/connectors/gpu/deepep_m2n/fp8_quant.py new file mode 100644 index 000000000..c2e9074c9 --- /dev/null +++ b/afd_plugin/connectors/gpu/deepep_m2n/fp8_quant.py @@ -0,0 +1,162 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""FP8 block-128 helpers for the DeepEP M2N expert path (SM90, fp32 scales). + +``silu_mul_fp8_quant`` is the fused step between the two FP8 grouped GEMMs: +``fp8(silu(gate) * up * topk_weight)`` with one fp32 scale per row and +128-channel group, written in the MN-major, TMA-aligned layout DeepGEMM's SM90 +kernel reads without a transpose. Folding the top-k weight in here makes the +second GEMM's output already weighted (it is linear), so no separate multiply +is needed before combine. + +Adapted from FastAFD ``kernel/deepgemm_fused_quant.py`` +(``_silu_mul_fp8_quant_psum_deep_gemm``, itself from vLLM), which only emits +SM100 UE8M0 scales. Rows are processed in parallel blocks instead of per +expert; rows past the last expert's end are skipped, and padding rows between +expert groups are quantized but never read by combine. +""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + +from afd_plugin.connectors.routed_experts import FP8_DTYPE + +FP8_MAX = torch.finfo(FP8_DTYPE).max +FP8_MIN = torch.finfo(FP8_DTYPE).min +GROUP_SIZE = 128 +# DeepGEMM's TMA needs the M stride of fp32 scales aligned to 16 bytes. +TMA_ALIGN_FP32_ELEMENTS = 4 +EPS = 1e-10 +BLOCK_M = 32 + + +def tma_aligned_rows(num_rows: int) -> int: + return -(-num_rows // TMA_ALIGN_FP32_ELEMENTS) * TMA_ALIGN_FP32_ELEMENTS + + +@triton.jit +def _silu_mul_fp8_quant_kernel( + x_ptr, + topk_weights_ptr, + psum_ptr, + y_q_ptr, + y_s_ptr, + num_rows, + num_experts, + stride_x_m, + stride_yq_m, + stride_ys_m, + stride_ys_g, + intermediate: tl.constexpr, + group_size: tl.constexpr, + apply_topk_weight: tl.constexpr, + block: tl.constexpr, + eps: tl.constexpr, + fp8_min: tl.constexpr, + fp8_max: tl.constexpr, +): + pid_m = tl.program_id(0) + group = tl.program_id(1) + valid_rows = tl.minimum(tl.load(psum_ptr + num_experts - 1), num_rows) + rows = (pid_m * block + tl.arange(0, block)).to(tl.int64) + row_mask = rows < valid_rows + cols = group * group_size + tl.arange(0, group_size) + gate_ptrs = x_ptr + rows[:, None] * stride_x_m + cols[None, :] + gate = tl.load(gate_ptrs, mask=row_mask[:, None], other=0.0).to(tl.float32) + up = tl.load(gate_ptrs + intermediate, mask=row_mask[:, None], other=0.0).to( + tl.float32 + ) + y = gate / (1.0 + tl.exp(-gate)) * up + if apply_topk_weight: + weights = tl.load(topk_weights_ptr + rows, mask=row_mask, other=0.0) + y = y * weights.to(tl.float32)[:, None] + scale = tl.maximum(tl.max(tl.abs(y), axis=1), eps) / fp8_max + y_q = tl.clamp(y / scale[:, None], fp8_min, fp8_max).to(y_q_ptr.dtype.element_ty) + tl.store( + y_q_ptr + rows[:, None] * stride_yq_m + cols[None, :], + y_q, + mask=row_mask[:, None], + ) + tl.store(y_s_ptr + rows * stride_ys_m + group * stride_ys_g, scale, mask=row_mask) + + +def silu_mul_fp8_quant( + gate_up: torch.Tensor, + recv_count_psum: torch.Tensor, + topk_weights: torch.Tensor | None, +) -> tuple[torch.Tensor, torch.Tensor]: + """``gate_up`` ``[M, 2I]`` BF16 -> FP8 ``[M, I]`` and fp32 scales ``[M, I/128]``. + + The scales are MN-major with a TMA-aligned M stride (DeepGEMM SM90 SFA). + """ + num_rows, intermediate_x2 = gate_up.shape + intermediate = intermediate_x2 // 2 + if intermediate % GROUP_SIZE != 0: + raise ValueError(f"intermediate size {intermediate} is not a multiple of 128") + num_groups = intermediate // GROUP_SIZE + y_q = torch.empty((num_rows, intermediate), dtype=FP8_DTYPE, device=gate_up.device) + y_s = torch.empty_strided( + (num_rows, num_groups), + (1, tma_aligned_rows(num_rows)), + dtype=torch.float32, + device=gate_up.device, + ) + if num_rows == 0: + return y_q, y_s + grid = (triton.cdiv(num_rows, BLOCK_M), num_groups) + _silu_mul_fp8_quant_kernel[grid]( + gate_up, + topk_weights if topk_weights is not None else gate_up, + recv_count_psum, + y_q, + y_s, + num_rows, + recv_count_psum.shape[0], + gate_up.stride(0), + y_q.stride(0), + y_s.stride(0), + y_s.stride(1), + intermediate=intermediate, + group_size=GROUP_SIZE, + apply_topk_weight=topk_weights is not None, + block=BLOCK_M, + eps=EPS, + fp8_min=FP8_MIN, + fp8_max=FP8_MAX, + ) + return y_q, y_s + + +def quantize_fp8_block128(w: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """``[E, N, K]`` weights -> FP8 and fp32 per-128x128-block scales. + + Same layout as vLLM's block-FP8 ``FusedMoE`` weights (``w13_weight`` / + ``w13_weight_scale_inv``), so real checkpoints can be passed through as is. + """ + num_experts, n, k = w.shape + if n % GROUP_SIZE or k % GROUP_SIZE: + raise ValueError(f"weight shape {tuple(w.shape)} is not 128-block aligned") + blocks = w.float().view( + num_experts, n // GROUP_SIZE, GROUP_SIZE, k // GROUP_SIZE, GROUP_SIZE + ) + scale = blocks.abs().amax(dim=(2, 4)).clamp(min=EPS) / FP8_MAX + w_q = (blocks / scale[:, :, None, :, None]).clamp(FP8_MIN, FP8_MAX).to(FP8_DTYPE) + return w_q.view(num_experts, n, k), scale + + +def dequantize_fp8_block128(w_q: torch.Tensor, scale: torch.Tensor) -> torch.Tensor: + num_experts, n, k = w_q.shape + blocks = w_q.float().view( + num_experts, n // GROUP_SIZE, GROUP_SIZE, k // GROUP_SIZE, GROUP_SIZE + ) + return (blocks * scale[:, :, None, :, None]).view(num_experts, n, k) + + +def dequantize_fp8_rows(x_q: torch.Tensor, scale: torch.Tensor) -> torch.Tensor: + """Per-row, per-128-channel FP8 activations -> fp32 (any scale strides).""" + num_rows, hidden = x_q.shape + groups = x_q.float().view(num_rows, hidden // GROUP_SIZE, GROUP_SIZE) + return (groups * scale[:, :, None]).view(num_rows, hidden) diff --git a/afd_plugin/connectors/gpu/deepep_m2n/layout.py b/afd_plugin/connectors/gpu/deepep_m2n/layout.py new file mode 100644 index 000000000..24d1bde89 --- /dev/null +++ b/afd_plugin/connectors/gpu/deepep_m2n/layout.py @@ -0,0 +1,70 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Expert placement over the DeepEP M2N union world. + +Adapted from FastAFD ``python/minisgl/moe/deepep_m2n_union.py`` (``3c71619``). + +DeepEP treats the union of Attention and FFN ranks as one symmetric EP group +in which every rank owns ``experts_per_rank`` experts. FastAFD orders the +union ``[Attention.., FFN..]`` and remaps real expert ids by an offset. The +plugin orders AFD worlds FFN first (``[F0.., A0..]``, as the P2P connector +does), so FFN rank ``g`` owns DeepEP experts ``[g * E_local, (g + 1) * E_local)``, +which are exactly the real expert ids: no remap. Attention ranks own dummy +expert slots after the real ones that routing never selects. +""" + +from __future__ import annotations + +from dataclasses import dataclass + + +@dataclass(frozen=True) +class DeepEPM2NLayout: + ffn_size: int + attn_size: int + num_experts: int + + def __post_init__(self) -> None: + if self.ffn_size <= 0 or self.attn_size <= 0: + raise ValueError( + "DeepEP M2N needs at least one FFN and one Attention rank, got " + f"ffn_size={self.ffn_size} attn_size={self.attn_size}" + ) + if self.num_experts % self.ffn_size != 0: + raise ValueError( + f"num_experts={self.num_experts} must be divisible by " + f"ffn_size={self.ffn_size}" + ) + + @property + def union_size(self) -> int: + return self.ffn_size + self.attn_size + + @property + def experts_per_rank(self) -> int: + return self.num_experts // self.ffn_size + + @property + def deepep_num_experts(self) -> int: + """Expert count DeepEP sees: real experts plus Attention dummy slots.""" + return self.union_size * self.experts_per_rank + + def world_rank(self, role: str, role_rank: int) -> int: + if role == "ffn": + _check_rank("ffn role rank", role_rank, self.ffn_size) + return role_rank + if role == "attention": + _check_rank("attention role rank", role_rank, self.attn_size) + return self.ffn_size + role_rank + raise ValueError(f"unknown AFD role {role!r}") + + def expert_range(self, ffn_rank: int) -> range: + """Real expert ids owned by FFN rank ``ffn_rank``.""" + _check_rank("ffn rank", ffn_rank, self.ffn_size) + start = ffn_rank * self.experts_per_rank + return range(start, start + self.experts_per_rank) + + +def _check_rank(name: str, rank: int, size: int) -> None: + if not 0 <= rank < size: + raise ValueError(f"{name} must be in [0, {size}), got {rank}") diff --git a/afd_plugin/connectors/gpu/deepep_m2n/network.py b/afd_plugin/connectors/gpu/deepep_m2n/network.py new file mode 100644 index 000000000..33f480182 --- /dev/null +++ b/afd_plugin/connectors/gpu/deepep_m2n/network.py @@ -0,0 +1,504 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Cluster topology discovery and the network plan for DeepEP M2N. + +DeepEP V2 moves data inside an NVLink domain with load/stores (LSA) and +between domains with NCCL GIN (RDMA). What that needs depends on the cluster: + +- one host, or an NVL72 rack: one NVLink domain, no RDMA (GIN off); +- InfiniBand fat-tree or RoCE with routed rails: GIN on, each GPU uses its + nearest NIC (NCCL's own choice); +- RoCE whose rails are separate, unrouted subnets: GIN on, and every rank must + use NICs on one shared rail, because DeepEP's GIN setup connects every pair + of ranks. + +Each rank reports a ``RankTopology`` in the connector's init handshake, and +every rank evaluates the same ``plan_network`` on the gathered list, so the +plan needs no further agreement. Whether rails are routed cannot be seen from +inside a pod; ``rails_routed`` declares it, and ``auto`` assumes they are not +when NICs sit on several rails (the choice that works everywhere, at the cost +of bandwidth). +""" + +from __future__ import annotations + +import ctypes +import ipaddress +import os +import socket +from collections.abc import Mapping +from dataclasses import asdict, dataclass, field, fields +from pathlib import Path +from typing import Any + +import psutil +import torch + +SYSFS_INFINIBAND = Path("/sys/class/infiniband") +SYSFS_PCI = Path("/sys/bus/pci/devices") +# RoCE v2 GIDs carry the port's IPv4 address as ::ffff:a.b.c.d. +IPV4_MAPPED_GID_PREFIX = "0000:0000:0000:0000:0000:ffff:" +ETHERNET = "ethernet" +INFINIBAND = "infiniband" +UNKNOWN_DISTANCE = 1 << 16 + +GIN_MODES = ("auto", "on", "off") +NIC_POLICIES = ("auto", "nccl", "shared", "explicit") +HYBRID_MODES = ("auto", "on", "off") +NCCL_ENV_PREFIXES = ("NCCL_", "EP_") +# Set from the plan (``network.gin``); a copy in ``nccl_env`` would fight it. +GIN_DISABLE_ENV = "EP_DISABLE_GIN" +# DeepEP's default RDMA service level / traffic class. +DEFAULT_RDMA_TRAFFIC_CLASS = 3 + +# Preset -> NetworkConfig defaults; explicit keys in the config override them. +PRESETS: dict[str, dict[str, Any]] = { + "auto": {}, + "single-node": {"gin": "off", "nic_policy": "nccl"}, + "roce-rail-isolated": {"nic_policy": "shared", "rails_routed": False}, + "roce-routed": {"nic_policy": "nccl", "rails_routed": True}, + "ib": {"nic_policy": "nccl", "rails_routed": True}, + "nvl72": {"nic_policy": "nccl", "rails_routed": True}, +} + + +@dataclass(frozen=True) +class NetworkConfig: + """The ``network`` section of ``connector_extra_config``. + + Attributes: + preset: Defaults for a fabric type (see ``PRESETS``). + gin: ``auto`` enables NCCL GIN only when ranks span several NVLink + domains; ``on``/``off`` force it (``off`` is rejected when ranks + span several domains: GIN is the only path between them). + nic_policy: ``auto`` picks ``nccl`` when rails are routed or all + NICs share one rail, else ``shared``; ``nccl`` keeps NCCL's + per-GPU choice; ``shared`` puts every rank on NICs of one rail + present on every host; ``explicit`` uses ``nics`` by local rank. + nics: Explicit NIC names: by local rank for ``explicit``, or the NIC + to prefer (first entry) for ``shared``. + rails_routed: Whether different rails can reach each other. + ``None`` means unknown (treated as not routed). + hybrid: DeepEP hybrid mode (rail-aligned RDMA + NVLink forwarding); + ``auto`` = on. + rdma_traffic_class: GIN service level / traffic class. + nccl_env: Extra ``NCCL_*``/``EP_*`` variables applied before any + NCCL communicator is created (process-wide). + """ + + preset: str = "auto" + gin: str = "auto" + nic_policy: str = "auto" + nics: tuple[str, ...] = () + rails_routed: bool | None = None + hybrid: str = "auto" + rdma_traffic_class: int = DEFAULT_RDMA_TRAFFIC_CLASS + nccl_env: tuple[tuple[str, str], ...] = () + + @classmethod + def from_mapping(cls, raw: Mapping[str, Any] | None) -> NetworkConfig: + raw = dict(raw or {}) + known = {item.name for item in fields(cls)} + unknown = sorted(set(raw) - known) + if unknown: + raise ValueError(f"unknown network options: {unknown}") + preset = raw.get("preset", "auto") + if preset not in PRESETS: + raise ValueError(f"network.preset must be one of {sorted(PRESETS)}") + merged = {**PRESETS[preset], **raw} + if "nics" in merged: + merged["nics"] = tuple(merged["nics"]) + if "nccl_env" in merged: + merged["nccl_env"] = tuple( + sorted((str(k), str(v)) for k, v in dict(merged["nccl_env"]).items()) + ) + config = cls(**merged) + config.validate() + return config + + def validate(self) -> None: + if self.gin not in GIN_MODES: + raise ValueError(f"network.gin must be one of {GIN_MODES}") + if self.nic_policy not in NIC_POLICIES: + raise ValueError(f"network.nic_policy must be one of {NIC_POLICIES}") + if self.hybrid not in HYBRID_MODES: + raise ValueError(f"network.hybrid must be one of {HYBRID_MODES}") + if not all(isinstance(nic, str) and nic for nic in self.nics): + raise ValueError("network.nics must be non-empty NIC names") + if self.nic_policy == "explicit" and not self.nics: + raise ValueError("network.nic_policy=explicit needs network.nics") + if self.rdma_traffic_class < 0: + raise ValueError("network.rdma_traffic_class must be non-negative") + for key, _ in self.nccl_env: + if not key.startswith(NCCL_ENV_PREFIXES): + raise ValueError( + f"network.nccl_env key {key!r} must start with {NCCL_ENV_PREFIXES}" + ) + if key == GIN_DISABLE_ENV: + raise ValueError( + f"network.nccl_env cannot set {GIN_DISABLE_ENV}; use network.gin" + ) + + def to_mapping(self) -> dict[str, Any]: + mapping = asdict(self) + mapping["nics"] = list(self.nics) + mapping["nccl_env"] = dict(self.nccl_env) + return mapping + + +@dataclass(frozen=True) +class NicInfo: + """One active RDMA port as seen from a rank.""" + + name: str + link_layer: str + # Ranks can talk RDMA only between NICs on the same rail (unless routed). + rail: str + # PCIe hops between the rank's GPU and this NIC (lower is closer). + distance: int + + +@dataclass(frozen=True) +class RankTopology: + world_rank: int + host: str + # Ranks with the same id share an NVLink (LSA) domain. + nvlink_domain: str + local_index: int + nics: tuple[NicInfo, ...] = field(default_factory=tuple) + + def to_mapping(self) -> dict[str, Any]: + return asdict(self) + + @classmethod + def from_mapping(cls, raw: Mapping[str, Any]) -> RankTopology: + return cls( + world_rank=int(raw["world_rank"]), + host=str(raw["host"]), + nvlink_domain=str(raw["nvlink_domain"]), + local_index=int(raw["local_index"]), + nics=tuple(NicInfo(**nic) for nic in raw["nics"]), + ) + + +@dataclass(frozen=True) +class NetworkPlan: + """What every rank does, derived identically from all ``RankTopology``.""" + + num_domains: int + ranks_per_domain: int + use_gin: bool + hybrid: bool + # world rank -> NIC to pin via NCCL_IB_HCA (None = NCCL's own choice). + nic_by_rank: tuple[str | None, ...] + nic_policy: str + rail: str | None + rdma_traffic_class: int + + def summary(self) -> str: + nics = sorted({nic for nic in self.nic_by_rank if nic is not None}) + return ( + f"domains={self.num_domains}x{self.ranks_per_domain} gin={self.use_gin} " + f"hybrid={self.hybrid} nic_policy={self.nic_policy} rail={self.rail} " + f"nics={nics or 'nccl'} traffic_class={self.rdma_traffic_class}" + ) + + +# --------------------------------------------------------------------------- +# Discovery (runs on each rank) +# --------------------------------------------------------------------------- + + +def discover_local_topology( + world_rank: int, + device: torch.device, + local_index: int, + *, + sysfs_infiniband: Path = SYSFS_INFINIBAND, + sysfs_pci: Path = SYSFS_PCI, + interface_networks: Mapping[str, str] | None = None, +) -> RankTopology: + host = socket.gethostname() + gpu_pci = gpu_pci_address(device) + return RankTopology( + world_rank=world_rank, + host=host, + nvlink_domain=nvlink_domain_id(gpu_pci) or f"host:{host}", + local_index=local_index, + nics=discover_nics( + gpu_pci, + sysfs_infiniband=sysfs_infiniband, + sysfs_pci=sysfs_pci, + interface_networks=( + local_interface_networks() + if interface_networks is None + else interface_networks + ), + ), + ) + + +def gpu_pci_address(device: torch.device) -> str | None: + """``dddd:bb:dd.f`` of the CUDA device (CUDA order, not NVML order).""" + props = torch.cuda.get_device_properties(device) + return ( + f"{props.pci_domain_id:04x}:{props.pci_bus_id:02x}:{props.pci_device_id:02x}.0" + ) + + +def nvlink_domain_id(gpu_pci: str | None) -> str | None: + """Multi-node NVLink (NVL72) clique id from NVML fabric info, if any. + + Returns ``None`` when the GPU is not part of an NVLink fabric (or NVML + lacks the API), so the caller falls back to "one domain per host". + """ + if gpu_pci is None: + return None + from vllm.third_party import pynvml + + info = pynvml.c_nvmlGpuFabricInfo_t() + try: + pynvml.nvmlInit() + handle = pynvml.nvmlDeviceGetHandleByPciBusId(gpu_pci.upper()) + pynvml.nvmlDeviceGetGpuFabricInfo(handle, ctypes.byref(info)) + except pynvml.NVMLError: + # Not an NVLink-fabric GPU (e.g. HGX H100): NOT_SUPPORTED. + return None + # Read the raw UUID bytes: a c_char array field stops at the first zero. + field_info = type(info).clusterUuid + cluster = ctypes.string_at( + ctypes.addressof(info) + field_info.offset, field_info.size + ).hex() + if info.state != pynvml.NVML_GPU_FABRIC_STATE_COMPLETED or not int(cluster, 16): + return None + return f"nvl:{cluster}:{info.cliqueId}" + + +def local_interface_networks() -> dict[str, str]: + """Local IPv4 address -> its network (``a.b.c.d/len``).""" + networks: dict[str, str] = {} + for addresses in psutil.net_if_addrs().values(): + for address in addresses: + if address.family == socket.AF_INET and address.netmask: + network = ipaddress.ip_network( + f"{address.address}/{address.netmask}", strict=False + ) + networks[address.address] = str(network) + return networks + + +def discover_nics( + gpu_pci: str | None, + *, + sysfs_infiniband: Path, + sysfs_pci: Path, + interface_networks: Mapping[str, str], +) -> tuple[NicInfo, ...]: + nics = [] + if not sysfs_infiniband.is_dir(): + return () + for device_dir in sorted(sysfs_infiniband.iterdir()): + port = device_dir / "ports" / "1" + state = _read(port / "state") + if state is None or "ACTIVE" not in state: + continue + link_layer = (_read(port / "link_layer") or "").strip().lower() + rail = nic_rail(port, link_layer, interface_networks) + if rail is None: + continue + nics.append( + NicInfo( + name=device_dir.name, + link_layer=link_layer, + rail=rail, + distance=pci_distance( + gpu_pci, _pci_address(device_dir), sysfs_pci=sysfs_pci + ), + ) + ) + return tuple(nics) + + +def nic_rail( + port: Path, link_layer: str, interface_networks: Mapping[str, str] +) -> str | None: + """Rail id: the RoCE port's IPv4 network, or the IB GID subnet prefix.""" + gids_dir = port / "gids" + if not gids_dir.is_dir(): + return None + gids = [_read(gid) for gid in sorted(gids_dir.iterdir(), key=_gid_index)] + if link_layer == ETHERNET: + for gid in gids: + if gid and gid.strip().startswith(IPV4_MAPPED_GID_PREFIX): + ipv4 = _ipv4_from_gid(gid.strip()) + # Without the netmask the port is its own rail (isolated). + return f"roce:{interface_networks.get(ipv4, ipv4 + '/32')}" + return None + if link_layer == INFINIBAND and gids and gids[0]: + prefix = ":".join(gids[0].strip().split(":")[:4]) + return f"ib:{prefix}" + return None + + +def pci_distance(gpu_pci: str | None, nic_pci: str | None, *, sysfs_pci: Path) -> int: + """PCIe hops between two functions via their common upstream bridge.""" + if gpu_pci is None or nic_pci is None: + return UNKNOWN_DISTANCE + try: + gpu_path = (sysfs_pci / gpu_pci).resolve(strict=True).parts + nic_path = (sysfs_pci / nic_pci).resolve(strict=True).parts + except OSError: + return UNKNOWN_DISTANCE + common = 0 + for gpu_part, nic_part in zip(gpu_path, nic_path, strict=False): + if gpu_part != nic_part: + break + common += 1 + return (len(gpu_path) - common) + (len(nic_path) - common) + + +def _pci_address(device_dir: Path) -> str | None: + try: + return (device_dir / "device").resolve(strict=True).name + except OSError: + return None + + +def _gid_index(path: Path) -> int: + return int(path.name) if path.name.isdigit() else 1 << 30 + + +def _ipv4_from_gid(gid: str) -> str: + high, low = gid[len(IPV4_MAPPED_GID_PREFIX) :].split(":") + return str(ipaddress.IPv4Address(int(high + low, 16))) + + +def _read(path: Path) -> str | None: + try: + return path.read_text(encoding="utf-8") + except OSError: + return None + + +# --------------------------------------------------------------------------- +# Plan (pure; identical on every rank) +# --------------------------------------------------------------------------- + + +def plan_network(topologies: list[RankTopology], config: NetworkConfig) -> NetworkPlan: + topologies = sorted(topologies, key=lambda topology: topology.world_rank) + domains = _domain_blocks(topologies) + multi_domain = len(domains) > 1 + if multi_domain and config.gin == "off": + # Without GIN, DeepEP reaches peers only by NVLink load/store, which + # cannot cross NVLink domains. + raise RuntimeError( + f"network.gin=off, but the ranks span {len(domains)} NVLink domains; " + "DeepEP needs NCCL GIN (RDMA) between domains" + ) + use_gin = multi_domain if config.gin == "auto" else config.gin == "on" + hybrid = config.hybrid != "off" + policy, rail, nic_by_rank = _choose_nics(topologies, config, use_gin) + return NetworkPlan( + num_domains=len(domains), + ranks_per_domain=len(domains[0]), + use_gin=use_gin, + hybrid=hybrid, + nic_by_rank=nic_by_rank, + nic_policy=policy, + rail=rail, + rdma_traffic_class=config.rdma_traffic_class, + ) + + +def _domain_blocks(topologies: list[RankTopology]) -> list[list[int]]: + """Ranks grouped by NVLink domain; DeepEP needs equal, contiguous blocks.""" + blocks: list[tuple[str, list[int]]] = [] + for topology in topologies: + if blocks and blocks[-1][0] == topology.nvlink_domain: + blocks[-1][1].append(topology.world_rank) + else: + blocks.append((topology.nvlink_domain, [topology.world_rank])) + names = [name for name, _ in blocks] + if len(set(names)) != len(names): + raise RuntimeError(f"ranks of one NVLink domain are not contiguous: {names}") + sizes = {len(ranks) for _, ranks in blocks} + if len(sizes) != 1: + raise RuntimeError( + "NVLink domains hold different rank counts: " + f"{[(name, len(ranks)) for name, ranks in blocks]}" + ) + return [ranks for _, ranks in blocks] + + +def _choose_nics( + topologies: list[RankTopology], config: NetworkConfig, use_gin: bool +) -> tuple[str, str | None, tuple[str | None, ...]]: + none = tuple(None for _ in topologies) + if not use_gin: + return "none", None, none + policy = config.nic_policy + if policy == "explicit": + nics = tuple(config.nics[t.local_index % len(config.nics)] for t in topologies) + return policy, None, nics + rails = {nic.rail for t in topologies for nic in t.nics} + if policy == "auto": + policy = "nccl" if config.rails_routed or len(rails) <= 1 else "shared" + if policy == "nccl": + return policy, None, none + rail = _shared_rail(topologies, config) + nics = tuple( + min( + (nic for nic in t.nics if nic.rail == rail), + key=lambda nic: (nic.distance, nic.name), + ).name + for t in topologies + ) + return policy, rail, nics + + +def _shared_rail(topologies: list[RankTopology], config: NetworkConfig) -> str: + """A rail every rank has a NIC on; prefer ``nics[0]``'s, then the closest.""" + common = set.intersection(*({nic.rail for nic in t.nics} for t in topologies)) + if not common: + raise RuntimeError("no RDMA rail is reachable from every rank") + if config.nics: + preferred = { + nic.rail + for t in topologies + for nic in t.nics + if nic.name == config.nics[0] and nic.rail in common + } + if preferred: + return min(preferred) + return min( + common, + key=lambda rail: ( + sum( + min(nic.distance for nic in t.nics if nic.rail == rail) + for t in topologies + ), + rail, + ), + ) + + +def apply_rank_environment( + plan: NetworkPlan, config: NetworkConfig, world_rank: int +) -> None: + """Process-wide NCCL/DeepEP settings for this rank, before any communicator. + + NCCL reads these when its first communicator in the process is created. + AFD workers therefore run the connector's ``bootstrap`` (which calls this) + before vLLM's distributed init creates its own NCCL groups; a process that + already created a communicator keeps the values it started with. + """ + for key, value in config.nccl_env: + os.environ[key] = value + nic = plan.nic_by_rank[world_rank] + if nic is not None: + os.environ["NCCL_IB_HCA"] = "=" + nic + if plan.use_gin: + os.environ.pop(GIN_DISABLE_ENV, None) + else: + os.environ[GIN_DISABLE_ENV] = "1" diff --git a/afd_plugin/connectors/metadata.py b/afd_plugin/connectors/metadata.py index f8c23bd02..728d94950 100644 --- a/afd_plugin/connectors/metadata.py +++ b/afd_plugin/connectors/metadata.py @@ -182,7 +182,10 @@ class AFDTransferMetadata: seq_lens: Per-peer or per-split token lengths. The sum of this list is the expected leading dimension for tensors validated against this metadata. For one-to-one transfers this is usually a single-item - list; for fan-in/fan-out paths it can describe split sizes. + list; for fan-in/fan-out paths it can describe split sizes. An + entry may be ``0``: with expert-routed transports (DeepEP M2N) an + FFN rank can receive no rows for a layer while still taking part + in the collective. """ layer_idx: int @@ -192,8 +195,8 @@ class AFDTransferMetadata: def __post_init__(self) -> None: if not self.seq_lens: raise ValueError("seq_lens cannot be empty") - if any(length <= 0 for length in self.seq_lens): - raise ValueError("all sequence lengths must be positive") + if any(length < 0 for length in self.seq_lens): + raise ValueError("all sequence lengths must be non-negative") @property def total_tokens(self) -> int: diff --git a/afd_plugin/connectors/routed_experts.py b/afd_plugin/connectors/routed_experts.py new file mode 100644 index 000000000..9ebd0d438 --- /dev/null +++ b/afd_plugin/connectors/routed_experts.py @@ -0,0 +1,133 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Contract for connectors whose FFN side receives routed expert rows. + +Routing (gate and top-k) runs on Attention. The FFN receives one row per +(token, local expert) pair, runs only its local experts and sends the weighted +rows back. Workers and runners use ``AFDRoutedExpertsConnector`` and the types +here instead of any specific connector, and this module imports no connector +backend, so they stay connector-agnostic. ``AFDConnectorBase`` is unchanged. +""" + +from __future__ import annotations + +from abc import abstractmethod +from collections.abc import Mapping, Sequence +from dataclasses import dataclass + +import torch + +from afd_plugin.connectors.base import AFDConnectorBase +from afd_plugin.connectors.metadata import AFDA2FTransferPayload + +FP8_DTYPE = torch.float8_e4m3fn + + +@dataclass(frozen=True) +class LocalExpertWeights: + """This FFN rank's experts, in vLLM ``FusedMoE`` layout. + + Attributes: + w13: ``[E_local, 2 * I, H]`` gate and up projections (gate first), + BF16 or FP8. + w2: ``[E_local, H, I]`` down projection, same dtype as ``w13``. + w13_scale: FP8 only: ``[E_local, 2 * I / 128, H / 128]`` fp32 block + scales (vLLM ``w13_weight_scale_inv``). + w2_scale: FP8 only: ``[E_local, H / 128, I / 128]`` fp32 block scales. + """ + + w13: torch.Tensor + w2: torch.Tensor + w13_scale: torch.Tensor | None = None + w2_scale: torch.Tensor | None = None + + def __post_init__(self) -> None: + if self.w13.dtype != self.w2.dtype: + raise ValueError("w13 and w2 must have the same dtype") + if self.is_fp8 != (self.w13_scale is not None and self.w2_scale is not None): + raise ValueError( + "FP8 expert weights need both block scales (and BF16 none)" + ) + + @property + def num_experts(self) -> int: + return int(self.w13.shape[0]) + + @property + def is_fp8(self) -> bool: + return self.w13.dtype == FP8_DTYPE + + +@dataclass(frozen=True) +class FFNLocalExperts: + """One MoE layer's experts on this FFN rank, as vLLM loaded them. + + Attributes: + expert_ids: Global ids of the experts in ``weights``, in slot order. + weights: The same experts in the executor's layout. + """ + + expert_ids: tuple[int, ...] + weights: LocalExpertWeights + + +class AFDRoutedExpertsConnector(AFDConnectorBase): + """A connector whose FFN side receives routed (token, local expert) rows. + + Lifecycle: workers create the connector and call ``bootstrap`` before + vLLM's distributed init (it may set process-wide NCCL environment, which + NCCL reads when the process creates its first communicator); runners then + call ``init_afd_connector`` after weights load, with what they loaded. + The FFN worker selects the routed-experts runner for these connectors. + """ + + @abstractmethod + def bootstrap(self) -> None: + """Collective: join the AFD world before vLLM's distributed init.""" + raise NotImplementedError + + # Extends AFDConnectorBase.init_afd_connector with optional keywords only, + # so calls through the base contract stay valid. + @abstractmethod + def init_afd_connector( + self, + *, + moe_layer_indices: Sequence[int] | None = None, + local_experts: Mapping[int, FFNLocalExperts] | None = None, + ) -> None: + """Collective: check what every rank loaded, then create resources. + + Args: + moe_layer_indices: MoE layers this rank's model runs, in order. + ``None`` skips the post-load checks for this rank. + local_experts: FFN only: this rank's experts per MoE layer, as + vLLM loaded them. + """ + raise NotImplementedError + + @abstractmethod + def run_experts( + self, + payload: AFDA2FTransferPayload, + weights: LocalExpertWeights, + ) -> torch.Tensor: + """Run this rank's experts on the rows ``recv_attn_output`` returned; + the result is what ``send_ffn_output`` sends back.""" + raise NotImplementedError + + @abstractmethod + def synchronize_lanes(self) -> None: + """Join the streams the connector forked into the caller's stream. + + FFN runners call this at the end of each step, before a CUDA-graph + capture ends. + """ + raise NotImplementedError + + +__all__ = [ + "AFDRoutedExpertsConnector", + "FFNLocalExperts", + "FP8_DTYPE", + "LocalExpertWeights", +] diff --git a/afd_plugin/distributed/__init__.py b/afd_plugin/distributed/__init__.py index b7b28853b..ba822f04f 100644 --- a/afd_plugin/distributed/__init__.py +++ b/afd_plugin/distributed/__init__.py @@ -8,6 +8,7 @@ resolve_role_rank, subgroup_attention_block, topology_from_config, + tp_rank_before_distributed_init, validate_p2p_topology, ) @@ -35,5 +36,6 @@ def __getattr__(name: str): "resolve_role_rank", "subgroup_attention_block", "topology_from_config", + "tp_rank_before_distributed_init", "validate_p2p_topology", ] diff --git a/afd_plugin/distributed/topology.py b/afd_plugin/distributed/topology.py index 20be8f64c..d2853b621 100644 --- a/afd_plugin/distributed/topology.py +++ b/afd_plugin/distributed/topology.py @@ -58,12 +58,40 @@ def validate_p2p_topology(config: AFDConfig) -> None: ) -def resolve_role_rank(vllm_config: VllmConfig, config: AFDConfig) -> int: +def tp_rank_before_distributed_init(vllm_config: VllmConfig, worker_rank: int) -> int: + """TP rank of a worker, before vLLM creates its TP group. + + vLLM orders ranks ``[ExternalDP, DP, PP, PCP, TP]`` with TP innermost + (vLLM v0.30.0 ``initialize_model_parallel``), and AFD runs with PP = PCP = + 1, so the TP rank is the worker rank modulo the TP size whether the worker + rank is DP-local or global. ``check_connector_role_rank`` compares the + result with vLLM's TP group once it exists. + """ + parallel_config = vllm_config.parallel_config + if ( + parallel_config.pipeline_parallel_size, + parallel_config.prefill_context_parallel_size, + ) != (1, 1): + raise ValueError("the TP rank before distributed init needs PP = PCP = 1") + return worker_rank % int(parallel_config.tensor_parallel_size) + + +def resolve_role_rank( + vllm_config: VllmConfig, + config: AFDConfig, + *, + tp_rank: int | None = None, +) -> int: """Resolve this worker's connector-independent AFD role rank. The resolver linearizes vLLM's global DP rank and local PCP/TP coordinates. Connectors receive the result as runtime state and map it to their own communication-world rank. + + Args: + tp_rank: This worker's TP rank, for callers that run before vLLM's + distributed init (see ``tp_rank_before_distributed_init``). + ``None`` reads it from vLLM's TP group. """ parallel_config = vllm_config.parallel_config @@ -82,12 +110,14 @@ def resolve_role_rank(vllm_config: VllmConfig, config: AFDConfig) -> int: pcp_rank = int(get_pcp_group().rank_in_group) else: pcp_rank = 0 - if tp_size > 1: + if tp_size == 1: + tp_rank = 0 + elif tp_rank is None: from vllm.distributed.parallel_state import get_tensor_model_parallel_rank tp_rank = int(get_tensor_model_parallel_rank()) - else: - tp_rank = 0 + elif not 0 <= tp_rank < tp_size: + raise ValueError(f"tp_rank must be in [0, {tp_size}), got {tp_rank}") role_rank = (dp_rank * pcp_size + pcp_rank) * tp_size + tp_rank if config.role == "attention": diff --git a/afd_plugin/model_executor/models/deepseek_v2.py b/afd_plugin/model_executor/models/deepseek_v2.py index 5e40c05b8..6433f6f92 100644 --- a/afd_plugin/model_executor/models/deepseek_v2.py +++ b/afd_plugin/model_executor/models/deepseek_v2.py @@ -10,7 +10,7 @@ from collections.abc import Iterable, Iterator from copy import copy -from typing import Any, TypeAlias +from typing import TYPE_CHECKING, Any, TypeAlias import torch import torch.nn as nn @@ -19,9 +19,15 @@ 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.fused_moe.router.fused_topk_router import fused_topk +from vllm.model_executor.layers.linear import ReplicatedLinear from vllm.model_executor.models import deepseek_v2 as native -from afd_plugin.config import AFD_ASYNC_CONNECTOR, parse_afd_config +from afd_plugin.config import ( + AFD_ASYNC_CONNECTOR, + DEEPEP_M2N_CONNECTORS, + parse_afd_config, +) from afd_plugin.connectors import ( AFDExpertRoutingSpec, AFDF2ATransferPayload, @@ -31,6 +37,9 @@ from afd_plugin.model_executor.models import get_afd_metadata_from_forward_context from afd_plugin.v1.worker.dbo import maybe_apply_dbo_yield +if TYPE_CHECKING: + from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts + logger = init_logger(__name__) _DeepseekAdapterConfig: TypeAlias = ( @@ -50,6 +59,15 @@ def _is_moe_layer(config: _DeepseekAdapterConfig, layer_idx: int) -> bool: _ATTENTION_ROLE = frozenset(("attention",)) _FFN_ROLE = frozenset(("ffn",)) _BOTH_ROLES = frozenset(("attention", "ffn")) +_NO_ROLE: frozenset[str] = frozenset() +# Connectors whose FFN builds no shared experts: Async CAM runs them on +# Attention; DeepEP M2N's FFN runs routed experts only. +_FFN_WITHOUT_SHARED_EXPERTS_CONNECTORS = ( + AFD_ASYNC_CONNECTOR, + *DEEPEP_M2N_CONNECTORS, +) +# Expert id of a padding row: a routed-experts connector sends it nowhere. +UNROUTED_EXPERT_ID = -1 def _weight_layer_path(name: str) -> tuple[int, str, tuple[str, ...]] | None: @@ -72,8 +90,14 @@ def _checkpoint_weight_roles( *, compute_gate_on_attention: bool, attention_shared_experts: bool = False, + ffn_skips_shared_experts: bool = False, ) -> frozenset[str]: - """Classify one native checkpoint path by its AFD execution owner.""" + """Classify one native checkpoint path by its AFD execution owner. + + ``ffn_skips_shared_experts`` keeps shared-expert weights off the FFN + without giving them to Attention, for connectors whose FFN runs routed + experts only while Attention does not build shared experts yet. + """ layer_path = _weight_layer_path(name) if layer_path is None: return _BOTH_ROLES @@ -87,8 +111,11 @@ def _checkpoint_weight_roles( if not _is_moe_layer(config, layer_idx): return _ATTENTION_ROLE if compute_gate_on_attention else _FFN_ROLE - if attention_shared_experts and remainder and remainder[0] == "shared_experts": - return _ATTENTION_ROLE + if remainder and remainder[0] == "shared_experts": + if attention_shared_experts: + return _ATTENTION_ROLE + if ffn_skips_shared_experts: + return _NO_ROLE is_moe_gate = bool(remainder) and remainder[0] == "gate" if is_moe_gate and compute_gate_on_attention: return _BOTH_ROLES @@ -102,6 +129,7 @@ def _iter_role_weights( config: _DeepseekAdapterConfig, compute_gate_on_attention: bool, attention_shared_experts: bool = False, + ffn_skips_shared_experts: bool = False, ) -> Iterator[tuple[str, torch.Tensor]]: """Consume a checkpoint iterator once and retain this role's paths.""" for name, loaded_weight in weights: @@ -110,19 +138,71 @@ def _iter_role_weights( config, compute_gate_on_attention=compute_gate_on_attention, attention_shared_experts=attention_shared_experts, + ffn_skips_shared_experts=ffn_skips_shared_experts, ): yield name, loaded_weight class RemoteFFNProxy(nn.Module): - """Parameter-free FFN stage executed through the AFD connector.""" + """FFN stage executed through the AFD connector. + + By default the proxy is parameter-free: it sends the hidden states and the + FFN role runs the complete MLP or MoE block. A model whose connector routes + on Attention (an M2N routed-experts connector such as DeepEP M2N) sets + ``gate``, ``top_k``, ``renormalize`` and ``tp_size`` after construction, + from its model and parallel config. The proxy then computes the gate and + vLLM's default fused top-k (``FusedTopKRouter``; EPLB is not supported) + and sends ``topk_ids`` and ``topk_weights``; combine returns the top-k + weighted sum, which replaces the native MoE block's output. + + With Attention TP > 1 every TP rank holds the same tokens after the + attention all-reduce; dispatching them from every rank would make the FFN + run each token TP times. Each rank instead pads the tokens to a multiple + of ``tp_size`` (``sequence_parallel_chunk``), routes and dispatches only + its slice, and the TP group all-gathers the combined slices, as vLLM's + sequence-parallel MoE does. Padding rows get ``UNROUTED_EXPERT_ID`` and + reach no FFN rank; they are dropped after the all-gather. + """ def __init__(self, *, layer_idx: int) -> None: super().__init__() self.layer_idx = layer_idx + self.gate: ReplicatedLinear | None = None + self.top_k = 0 + self.renormalize = False + self.tp_size = 1 def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - return self._send_and_receive(hidden_states) + if self.gate is None: + return self._send_and_receive(hidden_states) + num_tokens = hidden_states.shape[0] + if self.tp_size > 1: + hidden_states = native.sequence_parallel_chunk(hidden_states) + router_logits, _ = self.gate(hidden_states) + topk_weights, topk_ids, _ = fused_topk( + hidden_states=hidden_states, + gating_output=router_logits, + topk=self.top_k, + renormalize=self.renormalize, + ) + if self.tp_size > 1: + local_tokens = hidden_states.shape[0] + token_index = native.get_tensor_model_parallel_rank() * local_tokens + token_index = token_index + torch.arange( + local_tokens, device=topk_ids.device + ) + topk_ids = topk_ids.masked_fill( + (token_index >= num_tokens).unsqueeze(1), UNROUTED_EXPERT_ID + ) + output = self._send_and_receive( + hidden_states, + topk_ids=topk_ids, + topk_weights=topk_weights, + ) + if self.tp_size > 1: + output = native.tensor_model_parallel_all_gather(output, 0) + output = output[:num_tokens] + return output def _send_and_receive( self, @@ -528,7 +608,7 @@ def __init__( if device_type == "npu": native.FusedMoEFactory = fused_moe.FusedMoEFactory moe_config = config - if afd_config.connector == AFD_ASYNC_CONNECTOR: + if afd_config.connector in _FFN_WITHOUT_SHARED_EXPERTS_CONNECTORS: moe_config = copy(config) # HF validates this public field as int, but pinned native # MoE uses only None to omit shared construction (0 builds @@ -894,6 +974,15 @@ def get_experts_layer_indices(self) -> tuple[int, ...]: if isinstance(layer, AFDDeepseekV2DecoderLayer) and layer.is_moe_layer ) + def get_routed_experts(self, layer_idx: int) -> "RoutedExperts": + """FFN only: the vLLM module holding one MoE layer's local experts. + + Connectors that route on Attention run these experts themselves, so + ``routed_scaling_factor`` must be folded into the top-k weights that + Attention sends; the FFN never applies it. + """ + return self.layers[layer_idx].mlp.experts.routed_experts + def get_experts_routing_spec( self, layer_idx: int, @@ -955,6 +1044,9 @@ def compute_experts_output( def get_experts_layer_indices(self) -> tuple[int, ...]: return self.model.get_experts_layer_indices() + def get_routed_experts(self, layer_idx: int) -> "RoutedExperts": + return self.model.get_routed_experts(layer_idx) + def get_experts_routing_spec( self, layer_idx: int, @@ -970,6 +1062,8 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: compute_gate_on_attention=self.afd_config.compute_gate_on_attention, attention_shared_experts=self.afd_config.connector == AFD_ASYNC_CONNECTOR, + ffn_skips_shared_experts=self.afd_config.connector + in DEEPEP_M2N_CONNECTORS, ) ) diff --git a/afd_plugin/model_executor/models/qwen3_moe.py b/afd_plugin/model_executor/models/qwen3_moe.py index a7969b6ce..d3fca0b16 100644 --- a/afd_plugin/model_executor/models/qwen3_moe.py +++ b/afd_plugin/model_executor/models/qwen3_moe.py @@ -5,10 +5,13 @@ The wrapper keeps the native vLLM Qwen3 MoE lifecycle and injects a role-aware decoder layer through the native ``decoder_layer_type`` hook. Attention executes the native Attention, residual, and normalization path; -FFN executes the complete native dense MLP or MoE block. +FFN executes the complete native dense MLP or MoE block. With the DeepEP M2N +connector (which requires ``compute_gate_on_attention``), Attention also runs +the MoE gate and top-k, and FFN runs only the routed experts. """ from collections.abc import Iterable, Iterator +from typing import TYPE_CHECKING import torch import torch.nn as nn @@ -17,9 +20,16 @@ from vllm.model_executor.models import qwen3_moe as native from vllm.platforms import current_platform -from afd_plugin.config import AFDConfig, parse_afd_config +from afd_plugin.config import ( + DEEPEP_M2N_CONNECTORS, + AFDConfig, + parse_afd_config, +) from afd_plugin.model_executor.models.deepseek_v2 import RemoteFFNProxy +if TYPE_CHECKING: + from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts + _ATTENTION_ROLE = frozenset(("attention",)) _FFN_ROLE = frozenset(("ffn",)) _BOTH_ROLES = frozenset(("attention", "ffn")) @@ -34,8 +44,13 @@ def _is_moe_layer(config: Qwen3MoeConfig, layer_idx: int) -> bool: ) -def _weight_layer_path(name: str) -> tuple[int, str] | None: - """Return ``(layer index, stage)`` for a native decoder weight path.""" +def _routes_on_attention(afd_config: AFDConfig) -> bool: + """Whether the DeepEP M2N runner performs gate and top-k on Attention.""" + return afd_config.connector in DEEPEP_M2N_CONNECTORS + + +def _weight_layer_path(name: str) -> tuple[int, str, tuple[str, ...]] | None: + """Return ``(layer index, stage, remainder)`` for a native decoder path.""" parts = name.split(".") for marker_idx, part in enumerate(parts[:-2]): if part != "layers": @@ -44,18 +59,30 @@ def _weight_layer_path(name: str) -> tuple[int, str] | None: layer_idx = int(parts[marker_idx + 1]) except ValueError: continue - return layer_idx, parts[marker_idx + 2] + return layer_idx, parts[marker_idx + 2], tuple(parts[marker_idx + 3 :]) return None -def _checkpoint_weight_roles(name: str) -> frozenset[str]: - """Classify one native checkpoint path by its AFD execution owner.""" +def _checkpoint_weight_roles( + name: str, + *, + routes_on_attention: bool = False, +) -> frozenset[str]: + """Classify one native checkpoint path by its AFD execution owner. + + With ``routes_on_attention`` the MoE gate goes to both roles: + Attention routes with it, and FFN keeps the native gate parameter that + vLLM's strict loader requires but never runs it. + """ layer_path = _weight_layer_path(name) if layer_path is None: return _BOTH_ROLES - _layer_idx, stage = layer_path + _layer_idx, stage, remainder = layer_path if stage == "mlp": + is_moe_gate = bool(remainder) and remainder[0] == "gate" + if is_moe_gate and routes_on_attention: + return _BOTH_ROLES return _FFN_ROLE if stage in {"self_attn", "input_layernorm", "post_attention_layernorm"}: return _ATTENTION_ROLE @@ -66,10 +93,14 @@ def _iter_role_weights( weights: Iterable[tuple[str, torch.Tensor]], *, role: str, + routes_on_attention: bool = False, ) -> Iterator[tuple[str, torch.Tensor]]: """Consume a checkpoint iterator once and retain this role's paths.""" for name, loaded_weight in weights: - if role in _checkpoint_weight_roles(name): + if role in _checkpoint_weight_roles( + name, + routes_on_attention=routes_on_attention, + ): yield name, loaded_weight @@ -81,10 +112,34 @@ def _validate_supported_config( parallel_config = vllm_config.parallel_config if not current_platform.is_cuda(): raise RuntimeError("AFD Qwen3 MoE currently supports CUDA only") - if afd_config.compute_gate_on_attention: + # DeepEP M2N routes on Attention. Its FFN runs only routed experts and + # keeps the native gate built but unused. + # validate_gpu_model_runner_v2_config requires the Attention-side gate. + if afd_config.compute_gate_on_attention and not _routes_on_attention(afd_config): raise RuntimeError( - "AFD Qwen3 MoE requires compute_gate_on_attention=false", + "AFD Qwen3 MoE requires compute_gate_on_attention=false, except " + f"with {' or '.join(DEEPEP_M2N_CONNECTORS)}", ) + if _routes_on_attention(afd_config): + text_config = vllm_config.model_config.hf_text_config + # The routed FFN runner serves MoE layers only; a dense layer would + # have no owner. + if not all( + _is_moe_layer(text_config, layer_idx) + for layer_idx in range(text_config.num_hidden_layers) + ): + raise RuntimeError( + f"AFD Qwen3 MoE with {afd_config.connector} requires every " + "layer to be MoE (no mlp_only_layers, decoder_sparse_step=1)", + ) + # Neither the Attention proxy nor the routed FFN runner computes a + # shared expert, so its output would be dropped silently. The field is + # optional on Qwen3 configs; native vLLM reads it the same way. + if getattr(text_config, "shared_expert_intermediate_size", 0) > 0: + raise RuntimeError( + f"AFD Qwen3 MoE with {afd_config.connector} does not support " + "shared experts (shared_expert_intermediate_size > 0)", + ) if parallel_config.use_sequence_parallel_moe: raise RuntimeError( "AFD Qwen3 MoE does not support sequence-parallel MoE", @@ -155,6 +210,19 @@ def __init__( dual_chunk_attention_config=dual_chunk_attention_config, ) self.mlp = RemoteFFNProxy(layer_idx=layer_idx) + if _routes_on_attention(afd_config): + # Same gate as native Qwen3MoeSparseMoeBlock, same weight path. + self.mlp.gate = native.ReplicatedLinear( + config.hidden_size, + config.num_experts, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.mlp.gate", + ) + self.mlp.top_k = int(config.num_experts_per_tok) + self.mlp.renormalize = bool(config.norm_topk_prob) + # Attention TP peers each dispatch their slice of the tokens. + self.mlp.tp_size = vllm_config.parallel_config.tensor_parallel_size self.input_layernorm = native.RMSNorm( config.hidden_size, eps=config.rms_norm_eps, @@ -217,6 +285,10 @@ def get_experts_layer_indices(self) -> tuple[int, ...]: if _is_moe_layer(self.config, layer_idx) ) + def get_routed_experts(self, layer_idx: int) -> "RoutedExperts": + """FFN only: the vLLM module holding one MoE layer's local experts.""" + return self.layers[layer_idx].mlp.experts.routed_experts + class AFDQwen3MoeForCausalLM(native.Qwen3MoeForCausalLM): """Qwen3 MoE causal LM wrapper for AFD CUDA execution.""" @@ -316,6 +388,9 @@ def compute_ffn_output( def get_experts_layer_indices(self) -> tuple[int, ...]: return self.model.get_experts_layer_indices() + def get_routed_experts(self, layer_idx: int) -> "RoutedExperts": + return self.model.get_routed_experts(layer_idx) + # Patch reason: the native loader would materialize both execution stages # even though each AFD role constructs only one stage's decoder modules. # Patch functionality: retain only role-owned checkpoint paths, then use the @@ -325,7 +400,11 @@ def get_experts_layer_indices(self) -> tuple[int, ...]: # Commit: ced6857afa0ea7b2e3f0846a62e1394e90f15607 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: # ### PATCH START: filter the checkpoint stream by AFD execution role. - role_weights = _iter_role_weights(weights, role=self.afd_role) + role_weights = _iter_role_weights( + weights, + role=self.afd_role, + routes_on_attention=_routes_on_attention(self.afd_config), + ) # ### PATCH END: filter the checkpoint stream by AFD execution role. return super().load_weights(role_weights) diff --git a/afd_plugin/v1/worker/__init__.py b/afd_plugin/v1/worker/__init__.py index 1488f305b..2c4fe6da4 100644 --- a/afd_plugin/v1/worker/__init__.py +++ b/afd_plugin/v1/worker/__init__.py @@ -9,6 +9,9 @@ "AFDAttentionWorker": "afd_plugin.v1.worker.attention_worker", "AFDFFNWorker": "afd_plugin.v1.worker.ffn_worker", "GPUFFNModelRunner": "afd_plugin.v1.worker.ffn_model_runner", + "GPUFFNRoutedExpertsModelRunner": ( + "afd_plugin.v1.worker.ffn_routed_experts_model_runner" + ), } @@ -27,4 +30,5 @@ def __getattr__(name: str): "AFDAttentionWorker", "AFDFFNWorker", "GPUFFNModelRunner", + "GPUFFNRoutedExpertsModelRunner", ] diff --git a/afd_plugin/v1/worker/attention_model_runner_v2.py b/afd_plugin/v1/worker/attention_model_runner_v2.py index a4061f84d..4aa62b41d 100644 --- a/afd_plugin/v1/worker/attention_model_runner_v2.py +++ b/afd_plugin/v1/worker/attention_model_runner_v2.py @@ -38,6 +38,7 @@ AFDConnectorBase, AFDConnectorFactory, AFDForwardContextMetadata, + AFDRoutedExpertsConnector, ) from afd_plugin.model_executor.models.forward_context import use_afd_metadata_provider from afd_plugin.v1.worker.attention_metadata import ( @@ -446,9 +447,20 @@ def install_afd_metadata_on_forward_context( def parse_config(vllm_config: VllmConfig) -> AFDConfig: return parse_afd_config(vllm_config, expected_role="attention") + def adopt_bootstrapped_connector(self, connector: AFDConnectorBase) -> None: + """Use the connector the worker bootstrapped before vLLM's init. + + Replaces the connector created at construction, which never + communicated, so the native constructor signature stays unchanged. + """ + self.connector.close() + self.connector = connector + # Patch reason: native V2 load_model has no AFD connector lifecycle. # Patch functionality: initialize the AFD connector after native weight - # loading so Attention and FFN model loading overlap across roles. + # loading so Attention and FFN model loading overlap across roles. A + # connector that delivers routed expert rows also receives the MoE layers + # this model dispatches, which every rank checks against the FFN's. # Signature: matches vLLM v0.30.0 GPUModelRunnerV2.load_model exactly, # including the upstream *args/**kwargs pass-through. def load_model( @@ -458,7 +470,13 @@ def load_model( **kwargs, ) -> None: super().load_model(load_dummy_weights, *args, **kwargs) - if not self.connector.is_initialized: + if self.connector.is_initialized: + return + if isinstance(self.connector, AFDRoutedExpertsConnector): + self.connector.init_afd_connector( + moe_layer_indices=tuple(self.get_model().get_experts_layer_indices()), + ) + else: self.connector.init_afd_connector() # Patch reason: native KV initialization excludes DP1 from UBatchRunner diff --git a/afd_plugin/v1/worker/attention_worker.py b/afd_plugin/v1/worker/attention_worker.py index a33b9b4ed..85ebaca58 100644 --- a/afd_plugin/v1/worker/attention_worker.py +++ b/afd_plugin/v1/worker/attention_worker.py @@ -9,10 +9,17 @@ from vllm.v1.worker.gpu import model_runner as gpu_model_runner_v2 from vllm.v1.worker.gpu_worker import Worker +from afd_plugin.config import parse_afd_config +from afd_plugin.connectors import AFDConnectorFactory from afd_plugin.model_executor.models.model_utils import get_afd_model_config from afd_plugin.v1.worker.attention_model_runner_v2 import ( AFDAttentionModelRunnerV2, ) +from afd_plugin.v1.worker.connector_device import ( + check_connector_device, + check_connector_role_rank, + connector_device_index, +) from afd_plugin.validation import ( assert_compatible_afd_stack, validate_gpu_model_runner_v2_config, @@ -44,7 +51,9 @@ def __init__( # after an internal module import, with no injectable runner factory. # Patch functionality: validate the supported V2 contract, select the AFD model # identity before construction, and scope the module-class substitution to - # the delegated native initialization window. + # the delegated native initialization window. A routed-experts connector + # is created first on the device vLLM will select, and checked against + # that device once vLLM has selected it. # Signature: matches vLLM v0.30.0 Worker.init_device exactly: (self). # Upstream source: vllm/v1/worker/gpu_worker.py, Worker.init_device; # runner local-import/construct seam; commit @@ -79,12 +88,30 @@ def init_device(self): self.vllm_config.model_config = afd_model_config self.model_config = afd_model_config + # A routed-experts connector joins the AFD world here, before vLLM's + # distributed init creates NCCL groups, so its NCCL environment applies; + # the FFN worker does the same at the same point. + connector = AFDConnectorFactory.create_bootstrapped_connector( + self.rank, + connector_device_index(self.vllm_config, self.local_rank), + self.vllm_config, + parse_afd_config(self.vllm_config, expected_role="attention"), + ) native_runner_cls = gpu_model_runner_v2.GPUModelRunner gpu_model_runner_v2.GPUModelRunner = AFDAttentionModelRunnerV2 try: super().init_device() + if connector is not None: + check_connector_device(connector, self.device) + check_connector_role_rank(connector, self.vllm_config) + except BaseException: + if connector is not None: + connector.close() + raise finally: gpu_model_runner_v2.GPUModelRunner = native_runner_cls + if connector is not None: + self.model_runner.adopt_bootstrapped_connector(connector) # ### PATCH END: select and directly construct the AFD runner. torch.accelerator.empty_cache() diff --git a/afd_plugin/v1/worker/connector_device.py b/afd_plugin/v1/worker/connector_device.py new file mode 100644 index 000000000..ae8ed2b1a --- /dev/null +++ b/afd_plugin/v1/worker/connector_device.py @@ -0,0 +1,93 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""CUDA device for connectors created before vLLM's ``Worker.init_device``. + +A routed-experts connector bootstraps before vLLM's distributed init, so it is +created before vLLM selects the worker's device. ``connector_device_index`` +mirrors that selection; ``check_connector_device`` compares the result with +the device vLLM chose, so a drift after a vLLM upgrade fails at startup. +``check_connector_role_rank`` does the same for the role rank, whose TP part +is derived from the worker rank before vLLM's TP group exists. +""" + +from __future__ import annotations + +import torch +from vllm.config import VllmConfig +from vllm.platforms import current_platform +from vllm.platforms.interface import set_assigned_physical_gpu_ids + +from afd_plugin.connectors import AFDConnectorBase +from afd_plugin.distributed import resolve_role_rank + +# vLLM's GPU worker skips its local data-parallel device offset for these. +_NO_DP_DEVICE_OFFSET_EXECUTORS = ("ray", "external_launcher") + + +def connector_device_index(vllm_config: VllmConfig, local_rank: int) -> int: + """CUDA device index vLLM's GPU worker will select for ``local_rank``. + + Mirrors ``Worker.init_device`` (vLLM v0.30.0, ``vllm/v1/worker/ + gpu_worker.py``, commit ced6857): add the local data-parallel offset, + then map the logical index through ``assigned_physical_gpu_ids`` + (``--device-ids``) and ``CUDA_VISIBLE_DEVICES``. Registering the assigned + ids here is what vLLM does next; its setter accepts the same value again. + """ + parallel_config = vllm_config.parallel_config + if ( + parallel_config.distributed_executor_backend + not in _NO_DP_DEVICE_OFFSET_EXECUTORS + and parallel_config.data_parallel_backend != "ray" + and parallel_config.nnodes_within_dp == 1 + ): + dp_local_rank = parallel_config.data_parallel_rank_local + if dp_local_rank is None: + dp_local_rank = parallel_config.data_parallel_index + tp_pp_world_size = ( + parallel_config.pipeline_parallel_size + * parallel_config.tensor_parallel_size + ) + local_rank += dp_local_rank * tp_pp_world_size + assigned_physical_gpu_ids = parallel_config.assigned_physical_gpu_ids + if assigned_physical_gpu_ids is not None: + set_assigned_physical_gpu_ids(assigned_physical_gpu_ids) + return current_platform.logical_device_id_to_visible_device_id(local_rank) + + +def check_connector_device( + connector: AFDConnectorBase, + worker_device: torch.device, +) -> None: + """Fail when a connector does not use the device vLLM gave the worker.""" + if connector.local_rank != worker_device.index: + raise RuntimeError( + f"{type(connector).__name__} uses cuda:{connector.local_rank} but " + f"vLLM placed this worker on {worker_device}; the device selection " + "in afd_plugin/v1/worker/connector_device.py no longer matches " + "vLLM's Worker.init_device", + ) + + +def check_connector_role_rank( + connector: AFDConnectorBase, + vllm_config: VllmConfig, +) -> None: + """Fail when a connector's role rank differs from vLLM's distributed ranks. + + Call after vLLM's distributed init, which creates the TP group. + """ + role_rank = resolve_role_rank(vllm_config, connector.afd_config) + if connector.role_rank != role_rank: + raise RuntimeError( + f"{type(connector).__name__} has role rank {connector.role_rank} but " + f"vLLM's distributed ranks give {role_rank}; the TP rank derived " + "before distributed init (tp_rank_before_distributed_init) no longer " + "matches vLLM's rank layout", + ) + + +__all__ = [ + "check_connector_device", + "check_connector_role_rank", + "connector_device_index", +] diff --git a/afd_plugin/v1/worker/ffn_model_runner.py b/afd_plugin/v1/worker/ffn_model_runner.py index e31d6c239..4ca5b3717 100644 --- a/afd_plugin/v1/worker/ffn_model_runner.py +++ b/afd_plugin/v1/worker/ffn_model_runner.py @@ -25,6 +25,7 @@ ) from afd_plugin.config import AFDConfig, parse_afd_config from afd_plugin.connectors import ( + AFDConnectorBase, AFDConnectorFactory, AFDControlPayload, AFDControlPlane, @@ -108,6 +109,15 @@ def get_model(self) -> Any: def initialize_afd_connector(self) -> None: self.connector.init_afd_connector() + def adopt_bootstrapped_connector(self, connector: AFDConnectorBase) -> None: + """Use the connector the worker bootstrapped before vLLM's init. + + Replaces the connector created at construction, which never + communicated. + """ + self.connector.close() + self.connector = connector + def load_model(self, *, load_dummy_weights: bool = False, **kwargs: Any) -> None: """Load the vLLM model.""" @@ -176,10 +186,14 @@ def execute_model( self._cuda_graphs[graph_key]["graph"].replay() return None + # An eager fallback for a step Attention replays keeps the replay flag, + # so the connector plans the step the way Attention does (graph mode, + # no host sync). self._ffn_forward( dp_metadata_list=dp_metadata_list, is_graph_capturing=is_graph_capturing, is_warmup=is_warmup, + is_graph_replaying=is_graph_replaying, ) return None @@ -189,17 +203,15 @@ def _ffn_forward( dp_metadata_list: dict[int, DPMetadata | AFDDPMetadata], is_graph_capturing: bool = False, is_warmup: bool = False, + is_graph_replaying: bool = False, update_connector_state: bool = True, ) -> torch.Tensor | None: if update_connector_state: - cast( - AFDControlPlane, self.connector.control_plane - ).update_state_from_dp_metadata( - _make_dp_metadata_payload( - dp_metadata_list, - is_graph_capturing=is_graph_capturing, - is_warmup=is_warmup, - ), + self._update_connector_state( + dp_metadata_list, + is_graph_capturing=is_graph_capturing, + is_warmup=is_warmup, + is_graph_replaying=is_graph_replaying, ) model = self.get_model() @@ -268,6 +280,25 @@ def _ffn_forward( self.connector.send_ffn_output(rank_ffn_output, context) return rank_ffn_output + def _update_connector_state( + self, + dp_metadata_list: dict[int, DPMetadata | AFDDPMetadata], + *, + is_graph_capturing: bool, + is_warmup: bool, + is_graph_replaying: bool, + ) -> None: + cast( + AFDControlPlane, self.connector.control_plane + ).update_state_from_dp_metadata( + _make_dp_metadata_payload( + dp_metadata_list, + is_graph_capturing=is_graph_capturing, + is_warmup=is_warmup, + is_graph_replaying=is_graph_replaying, + ), + ) + def _execute_eager_mode( self, hidden_states: torch.Tensor, @@ -472,11 +503,13 @@ def _make_dp_metadata_payload( *, is_graph_capturing: bool = False, is_warmup: bool = False, + is_graph_replaying: bool = False, ) -> AFDControlPayload: return AFDControlPayload( dp_metadata_list=dp_metadata_list, is_graph_capturing=is_graph_capturing, is_warmup=is_warmup, + is_graph_replaying=is_graph_replaying, ) diff --git a/afd_plugin/v1/worker/ffn_routed_experts_model_runner.py b/afd_plugin/v1/worker/ffn_routed_experts_model_runner.py new file mode 100644 index 000000000..09a0220a4 --- /dev/null +++ b/afd_plugin/v1/worker/ffn_routed_experts_model_runner.py @@ -0,0 +1,193 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""FFN-side model runner for connectors that deliver routed expert rows.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +import torch +from vllm.config import VllmConfig +from vllm.forward_context import DPMetadata + +from afd_plugin.connectors import ( + AFDConnectorBase, + AFDDPMetadata, + AFDRoutedExpertsConnector, + AFDTransferContext, + FFNLocalExperts, + LocalExpertWeights, +) +from afd_plugin.connectors.routed_experts import FP8_DTYPE +from afd_plugin.v1.worker.ffn_model_runner import GPUFFNModelRunner + +if TYPE_CHECKING: + from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts + + +# vLLM MoE backends whose ``process_weights_after_loading`` keeps the FusedMoE +# layout the connector's executor reads: w13 as [gate; up] with block-128 +# scales. Others reorder in place without changing dtype or shape (FlashInfer +# swaps w13 to [up; gate]; TRT-LLM and Marlin also shuffle), which yields wrong +# outputs with no error. +ROUTED_MOE_BACKEND = "triton" +_LAYOUT_PRESERVING_BACKENDS = frozenset({"TRITON"}) + + +def select_routed_moe_backend(vllm_config: VllmConfig) -> None: + """Pin the FFN's vLLM MoE backend to one that keeps the expert layout. + + vLLM's MoE forward never runs on a routed-experts FFN; its backend only + decides how the expert weights are processed at load. ``auto`` picks + FlashInfer CUTLASS on Hopper with EP > 1 for block FP8, which reorders + w13, so ``auto`` becomes ``triton`` and any other explicit choice fails. + """ + kernel_config = vllm_config.kernel_config + backend = str(kernel_config.moe_backend) + if backend == "auto": + kernel_config.moe_backend = ROUTED_MOE_BACKEND + elif backend != ROUTED_MOE_BACKEND: + raise RuntimeError( + f"routed-experts AFD FFN needs --moe-backend {ROUTED_MOE_BACKEND} (or " + f"auto); moe_backend={backend!r} reorders the expert weights the " + "connector reads" + ) + + +def _check_expert_layout(routed_experts: RoutedExperts) -> None: + quant_method = routed_experts.quant_method + backend = getattr(quant_method, "fp8_backend", None) or getattr( + quant_method, "unquantized_backend", None + ) + name = getattr(backend, "name", None) + if name not in _LAYOUT_PRESERVING_BACKENDS: + raise RuntimeError( + f"vLLM processed the FFN experts for MoE backend {name or backend!r}, " + f"which may reorder them; the routed-experts connector reads them " + f"as [gate; up] FusedMoE weights (supported: " + f"{sorted(_LAYOUT_PRESERVING_BACKENDS)}; set --moe-backend " + f"{ROUTED_MOE_BACKEND})" + ) + + +def local_experts_from_vllm(routed_experts: RoutedExperts) -> FFNLocalExperts: + """This rank's experts of one vLLM MoE layer, as views of its parameters. + + Read after ``process_weights_after_loading``, which must have kept the + FusedMoE layout (``select_routed_moe_backend``); checked here because a + reordered layout keeps dtypes and shapes. + """ + _check_expert_layout(routed_experts) + w13 = routed_experts.w13_weight.data + w2 = routed_experts.w2_weight.data + w13_scale = w2_scale = None + if w13.dtype == FP8_DTYPE: + w13_scale = routed_experts.w13_weight_scale_inv.data + w2_scale = routed_experts.w2_weight_scale_inv.data + return FFNLocalExperts( + expert_ids=tuple(routed_experts.expert_map_manager.get_local_expert_ids()), + weights=LocalExpertWeights( + w13=w13, w2=w2, w13_scale=w13_scale, w2_scale=w2_scale + ), + ) + + +def _require_routed(connector: AFDConnectorBase) -> AFDRoutedExpertsConnector: + if not isinstance(connector, AFDRoutedExpertsConnector): + raise TypeError( + f"{type(connector).__name__} is not an AFDRoutedExpertsConnector; " + "use GPUFFNModelRunner", + ) + return connector + + +class GPUFFNRoutedExpertsModelRunner(GPUFFNModelRunner): + """FFN runner for connectors that deliver routed (token, expert) rows. + + Routing (gate and top-k) runs on Attention. For each MoE layer and stage, + in the same order as Attention, the FFN receives the rows routed to its + local experts, runs those experts with the connector's executor and sends + the weighted rows back just before that stage's next receive. vLLM still + builds the model and loads and places the expert weights; its MoE forward + never runs, so the P2P path's forward-context and DP metadata setup is + skipped. + + The worker selects this runner for an ``AFDRoutedExpertsConnector`` + (DeepEP M2N). + """ + + def __init__(self, vllm_config: VllmConfig, device: torch.device) -> None: + super().__init__(vllm_config, device) + _require_routed(self.connector) + self.moe_layer_indices: tuple[int, ...] = () + self.local_experts: dict[int, FFNLocalExperts] = {} + + @property + def routed_connector(self) -> AFDRoutedExpertsConnector: + """The connector; the worker may replace it with the one it + bootstrapped, so it is checked where it is used.""" + return _require_routed(self.connector) + + def load_model(self, *, load_dummy_weights: bool = False, **kwargs: Any) -> None: + """Load with the native loader, then take each MoE layer's experts. + + Runs again on ``reload_weights``, whose parameters replace the old ones. + """ + super().load_model(load_dummy_weights=load_dummy_weights, **kwargs) + model = self.get_model() + self.moe_layer_indices = tuple(model.get_experts_layer_indices()) + self.local_experts = { + layer_idx: local_experts_from_vllm(model.get_routed_experts(layer_idx)) + for layer_idx in self.moe_layer_indices + } + + def initialize_afd_connector(self) -> None: + self.routed_connector.init_afd_connector( + moe_layer_indices=self.moe_layer_indices, + local_experts=self.local_experts, + ) + + def _ffn_forward( + self, + *, + dp_metadata_list: dict[int, DPMetadata | AFDDPMetadata], + is_graph_capturing: bool = False, + is_warmup: bool = False, + is_graph_replaying: bool = False, + update_connector_state: bool = True, + ) -> torch.Tensor | None: + if update_connector_state: + self._update_connector_state( + dp_metadata_list, + is_graph_capturing=is_graph_capturing, + is_warmup=is_warmup, + is_graph_replaying=is_graph_replaying, + ) + connector = self.routed_connector + stage_ids = sorted(int(stage_idx) for stage_idx in dp_metadata_list) or [0] + ffn_output = None + # Each stage's combine waits until just before that stage's next + # dispatch, matching Attention's DBO micro-batch order + # D0(0) D1(0) C0(0) D0(1) C1(0) D1(1) ... C0(n) C1(n). Dispatch and + # combine can block the host, so a combine issued earlier would wait + # on an Attention combine that only follows the next stage's dispatch. + # On the first layer no stage has a result to send back yet. + pending_outputs: dict[int, tuple[torch.Tensor, AFDTransferContext]] = {} + for layer_idx in self.moe_layer_indices: + weights = self.local_experts[layer_idx].weights + for stage_idx in stage_ids: + if stage_idx in pending_outputs: + connector.send_ffn_output(*pending_outputs.pop(stage_idx)) + payload = connector.recv_attn_output(ubatch_idx=stage_idx) + ffn_output = connector.run_experts(payload, weights) + pending_outputs[stage_idx] = (ffn_output, payload.context) + for stage_idx in stage_ids: + if stage_idx in pending_outputs: + connector.send_ffn_output(*pending_outputs.pop(stage_idx)) + # Lane streams forked by dispatch/combine must rejoin the caller's + # stream before a CUDA-graph capture ends. + connector.synchronize_lanes() + return ffn_output + + +__all__ = ["GPUFFNRoutedExpertsModelRunner", "local_experts_from_vllm"] diff --git a/afd_plugin/v1/worker/ffn_worker.py b/afd_plugin/v1/worker/ffn_worker.py index b657da8f0..a85026353 100644 --- a/afd_plugin/v1/worker/ffn_worker.py +++ b/afd_plugin/v1/worker/ffn_worker.py @@ -14,7 +14,14 @@ from vllm.v1.worker.gpu_worker import Worker from vllm.v1.worker.worker_base import CompilationTimes +from afd_plugin.config import parse_afd_config +from afd_plugin.connectors import AFDConnectorFactory, AFDRoutedExpertsConnector from afd_plugin.model_executor.models.model_utils import get_afd_model_config +from afd_plugin.v1.worker.connector_device import ( + check_connector_device, + check_connector_role_rank, + connector_device_index, +) from afd_plugin.v1.worker.ffn_model_runner import GPUFFNModelRunner from afd_plugin.validation import ( assert_compatible_afd_stack, @@ -61,7 +68,12 @@ def __init__( # Patch reason: vLLM 0.30.0 selects its V2 runner from a module-level # symbol inside Worker.init_device and exposes no injectable runner factory. # Patch functionality: replace only the V2 construction seam with the - # connector-driven GPUFFNModelRunner, the sole GPU FFN runtime. + # connector-driven FFN runner: GPUFFNRoutedExpertsModelRunner (imported + # only then) for an AFDRoutedExpertsConnector, else GPUFFNModelRunner. A + # routed-experts connector is created and bootstrapped first, so its NCCL + # environment precedes vLLM's NCCL groups, and the runner adopts it after + # construction. The connector gets the device vLLM will select, checked + # once vLLM has selected it. # Signature: matches vLLM v0.30.0 Worker.init_device exactly: (self). # Upstream source: vllm/v1/worker/gpu_worker.py, Worker.init_device; # ced6857afa0ea7b2e3f0846a62e1394e90f15607. The 0.30.0 init_device still @@ -90,12 +102,41 @@ def init_device(self): ) self.vllm_config.model_config = afd_model_config self.model_config = afd_model_config + afd_config = parse_afd_config(self.vllm_config, expected_role="ffn") + runner_cls: type[GPUFFNModelRunner] = GPUFFNModelRunner + connector_cls = AFDConnectorFactory.get_connector_class(afd_config.connector) + if issubclass(connector_cls, AFDRoutedExpertsConnector): + # Lazy: only workers whose connector delivers routed rows load the + # routed-experts runner. + from afd_plugin.v1.worker.ffn_routed_experts_model_runner import ( + GPUFFNRoutedExpertsModelRunner, + select_routed_moe_backend, + ) + + runner_cls = GPUFFNRoutedExpertsModelRunner + # Before bootstrap: a rejected backend fails before any collective. + select_routed_moe_backend(self.vllm_config) + connector = AFDConnectorFactory.create_bootstrapped_connector( + self.rank, + connector_device_index(self.vllm_config, self.local_rank), + self.vllm_config, + afd_config, + ) native_runner_cls = gpu_model_runner_v2.GPUModelRunner - gpu_model_runner_v2.GPUModelRunner = GPUFFNModelRunner + gpu_model_runner_v2.GPUModelRunner = runner_cls try: super().init_device() + if connector is not None: + check_connector_device(connector, self.device) + check_connector_role_rank(connector, self.vllm_config) + except BaseException: + if connector is not None: + connector.close() + raise finally: gpu_model_runner_v2.GPUModelRunner = native_runner_cls + if connector is not None: + self.model_runner.adopt_bootstrapped_connector(connector) # ### PATCH END: construct the FFN runner through the V2 seam. torch.accelerator.empty_cache() diff --git a/afd_plugin/validation.py b/afd_plugin/validation.py index ad4ce5c27..89750b076 100644 --- a/afd_plugin/validation.py +++ b/afd_plugin/validation.py @@ -12,6 +12,7 @@ from afd_plugin.config import ( AFD_ASYNC_CONNECTOR, CAMP2P_CONNECTOR, + DEEPEP_M2N_CONNECTORS, AFDConfig, parse_afd_config, ) @@ -29,6 +30,9 @@ "afd_plugin.v1.worker.AFDAttentionModelRunnerV2" ) FFN_MODEL_RUNNER_FQCN: Final[str] = "afd_plugin.v1.worker.GPUFFNModelRunner" +ROUTED_FFN_MODEL_RUNNER_FQCN: Final[str] = ( + "afd_plugin.v1.worker.GPUFFNRoutedExpertsModelRunner" +) NPU_ATTENTION_WORKER_FQCN: Final[str] = "afd_plugin.v1.worker.npu.AFDNPUAttentionWorker" NPU_FFN_WORKER_FQCN: Final[str] = "afd_plugin.v1.worker.npu.AFDNPUFFNWorker" NPU_ATTENTION_MODEL_RUNNER_FQCN: Final[str] = ( @@ -40,6 +44,10 @@ VLLM_ASCEND_310P_WORKER_FQCN: Final[str] = "vllm_ascend._310p.worker_310p.NPUWorker310" VLLM_ASCEND_XLITE_WORKER_FQCN: Final[str] = "vllm_ascend.xlite.xlite_worker.XliteWorker" GPU_MRV2_NUM_UBATCHES: Final[int] = 2 +GPU_MRV2_CONNECTORS: Final[tuple[str, ...]] = ( + "P2pNcclAFDConnector", + *DEEPEP_M2N_CONNECTORS, +) CAM_MAX_RECEIVE_TOKENS = 262144 @@ -91,12 +99,29 @@ def validate_gpu_model_runner_v2_config( raise RuntimeError("AFD GPU requires ModelRunnerV2") afd_config = parse_afd_config(vllm_config, expected_role=expected_role) - if device_type != "cuda" or afd_config.connector != "P2pNcclAFDConnector": + if device_type != "cuda" or afd_config.connector not in GPU_MRV2_CONNECTORS: raise RuntimeError( - "AFD ModelRunnerV2 requires CUDA and synchronous P2pNcclAFDConnector", + "AFD ModelRunnerV2 requires CUDA and one of the synchronous " + f"connectors {GPU_MRV2_CONNECTORS}", ) parallel = vllm_config.parallel_config + # DeepEP M2N bootstraps before vLLM's distributed init and derives the TP + # part of the role rank from the worker rank. Attention may run TP > 1: + # each TP rank dispatches its own slice of the tokens. FFN experts are not + # TP-sharded; that and the other deployment checks (expert placement, + # all2all backend, weight format) run in the connector handshake so that + # every rank fails together. It routes on Attention, so it requires the + # gate there. + if afd_config.connector in DEEPEP_M2N_CONNECTORS: + if expected_role == "ffn" and parallel.tensor_parallel_size != 1: + raise RuntimeError( + f"{afd_config.connector} requires FFN tensor_parallel_size=1" + ) + if not afd_config.compute_gate_on_attention: + raise RuntimeError( + f"{afd_config.connector} requires compute_gate_on_attention=true" + ) if ( parallel.pipeline_parallel_size, parallel.prefill_context_parallel_size, @@ -374,10 +399,12 @@ def _ctx() -> str: "ATTENTION_WORKER_FQCN", "FFN_MODEL_RUNNER_FQCN", "FFN_WORKER_FQCN", + "GPU_MRV2_CONNECTORS", "NPU_ATTENTION_MODEL_RUNNER_FQCN", "NPU_ATTENTION_WORKER_FQCN", "NPU_FFN_MODEL_RUNNER_FQCN", "NPU_FFN_WORKER_FQCN", + "ROUTED_FFN_MODEL_RUNNER_FQCN", "assert_compatible_afd_stack", "expected_npu_worker_qualname", "expected_worker_qualname", diff --git a/pyproject.toml b/pyproject.toml index 750dadf23..658667885 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -35,6 +35,9 @@ dependencies = [] # Install the runtime extra only in environments with a compatible vLLM wheel. [project.optional-dependencies] vllm = ["vllm==0.30.0"] +afdep = [ + "afdep @ git+https://github.com/ms-llmd/afdEP.git@e763e3189398beda24f592d2117e6ade135b0a7f", +] [project.urls] Homepage = "https://github.com/vllm-project/afd-plugin" @@ -191,4 +194,6 @@ extend-ignore-identifiers-re = [ "DataCopyPad2DA8W4NDNew", ".*isND2NZ.*", ".*enVecND2NZ.*", + # UE8M0: the FP8 unsigned-exponent scale format (DeepEP/DeepGEMM docs). + ".*[Uu][Ee]8[Mm]0.*", ] diff --git a/uv.lock b/uv.lock index 238e01688..77a951457 100644 --- a/uv.lock +++ b/uv.lock @@ -27,6 +27,17 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/58/0a/a10b45aab35b175aded078a462dc8d0c698f5b13946e7cb0869097b78bb6/absl_py-2.5.0-py3-none-any.whl", hash = "sha256:0f17b89f2a4eaaedc4f28c622998aa690564b3012a396a4ffad0821007fe03ba", size = 137410, upload-time = "2026-07-03T10:57:46.735Z" }, ] +[[package]] +name = "afdep" +version = "1.0.0.dev0" +source = { git = "https://github.com/ms-llmd/afdEP.git?rev=e763e3189398beda24f592d2117e6ade135b0a7f#e763e3189398beda24f592d2117e6ade135b0a7f" } +dependencies = [ + { name = "ninja" }, + { name = "packaging" }, + { name = "setuptools" }, + { name = "torch" }, +] + [[package]] name = "aiohappyeyeballs" version = "2.6.1" @@ -630,7 +641,7 @@ resolution-markers = [ "python_full_version < '3.11' and sys_platform == 'darwin'", ] dependencies = [ - { name = "cuda-pathfinder", marker = "sys_platform == 'darwin'" }, + { name = "cuda-pathfinder" }, ] [[package]] @@ -650,7 +661,7 @@ resolution-markers = [ "python_full_version < '3.11' and sys_platform != 'darwin'", ] dependencies = [ - { name = "cuda-pathfinder", marker = "sys_platform != 'darwin'" }, + { name = "cuda-pathfinder" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/a9/21/8464d133752951c154feafb3b65c297e7d80f301183d220bec4c830f1441/cuda_bindings-13.3.1-cp310-cp310-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:120fcc53d57903df529c3486962c56528cba5b7d6c57c99537320ed9922c8b86", size = 6073403, upload-time = "2026-05-29T23:11:36.22Z" }, @@ -710,7 +721,7 @@ resolution-markers = [ "python_full_version < '3.11' and sys_platform == 'darwin'", ] dependencies = [ - { name = "cuda-bindings", version = "12.9.4", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform == 'darwin'" }, + { name = "cuda-bindings", version = "12.9.4", source = { registry = "https://pypi.org/simple" } }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/af/f3/6b032a554019cfb3447e671798c1bd3e79b5f1af20d10253f56cea269ef2/cuda_python-12.9.4-py3-none-any.whl", hash = "sha256:d2cacea882a69863f1e7d27ee71d75f0684f4c76910aff839067e4f89c902279", size = 7594, upload-time = "2025-10-21T14:55:12.846Z" }, @@ -733,9 +744,9 @@ resolution-markers = [ "python_full_version < '3.11' and sys_platform != 'darwin'", ] dependencies = [ - { name = "cuda-bindings", version = "13.3.1", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform != 'darwin'" }, - { name = "cuda-core", marker = "sys_platform != 'darwin'" }, - { name = "cuda-pathfinder", marker = "sys_platform != 'darwin'" }, + { name = "cuda-bindings", version = "13.3.1", source = { registry = "https://pypi.org/simple" } }, + { name = "cuda-core" }, + { name = "cuda-pathfinder" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/38/31/7ff3f7768eded7535c621abc2fecb9d181a34ea4cae3afe682feb796f242/cuda_python-13.3.1-py3-none-any.whl", hash = "sha256:280b014139ab447b6dd70a377db1596f310d6e887d9d342e6651b919ec145fb3", size = 8295, upload-time = "2026-05-29T23:28:47.012Z" }, @@ -773,43 +784,43 @@ wheels = [ [package.optional-dependencies] cublas = [ - { name = "nvidia-cublas", marker = "(python_full_version < '3.11' and platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "nvidia-cuda-nvrtc", version = "13.0.88", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.11' and platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nvidia-cublas", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, + { name = "nvidia-cuda-nvrtc", version = "13.0.88", source = { registry = "https://pypi.org/simple" }, marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, ] cudart = [ - { name = "nvidia-cuda-runtime", version = "13.0.96", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.11' and platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nvidia-cuda-runtime", version = "13.0.96", source = { registry = "https://pypi.org/simple" }, marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, ] cufft = [ - { name = "nvidia-cufft", marker = "(python_full_version < '3.11' and platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "nvidia-nvjitlink", marker = "(python_full_version < '3.11' and platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nvidia-cufft", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, + { name = "nvidia-nvjitlink", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, ] cufile = [ - { name = "nvidia-cufile", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nvidia-cufile", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, ] cupti = [ - { name = "nvidia-cuda-cupti", marker = "(python_full_version < '3.11' and platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nvidia-cuda-cupti", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, ] curand = [ - { name = "nvidia-curand", marker = "(python_full_version < '3.11' and platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nvidia-curand", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, ] cusolver = [ - { name = "nvidia-cublas", marker = "(python_full_version < '3.11' and platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "nvidia-cusolver", marker = "(python_full_version < '3.11' and platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "nvidia-cusparse", marker = "(python_full_version < '3.11' and platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "nvidia-nvjitlink", marker = "(python_full_version < '3.11' and platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nvidia-cublas", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, + { name = "nvidia-cusolver", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, + { name = "nvidia-cusparse", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, + { name = "nvidia-nvjitlink", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, ] cusparse = [ - { name = "nvidia-cusparse", marker = "(python_full_version < '3.11' and platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "nvidia-nvjitlink", marker = "(python_full_version < '3.11' and platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nvidia-cusparse", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, + { name = "nvidia-nvjitlink", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, ] nvjitlink = [ - { name = "nvidia-nvjitlink", marker = "(python_full_version < '3.11' and platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nvidia-nvjitlink", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, ] nvrtc = [ - { name = "nvidia-cuda-nvrtc", version = "13.0.88", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.11' and platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nvidia-cuda-nvrtc", version = "13.0.88", source = { registry = "https://pypi.org/simple" }, marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, ] nvtx = [ - { name = "nvidia-nvtx", marker = "(python_full_version < '3.11' and platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nvidia-nvtx", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, ] [[package]] @@ -968,7 +979,7 @@ name = "exceptiongroup" version = "1.3.1" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "typing-extensions", marker = "python_full_version < '3.11'" }, + { name = "typing-extensions" }, ] sdist = { url = "https://files.pythonhosted.org/packages/50/79/66800aadf48771f6b62f7eb014e352e5d06856655206165d775e675a02c9/exceptiongroup-1.3.1.tar.gz", hash = "sha256:8b412432c6055b0b7d14c310000ae93352ed6754f70fa8f7c34141f91c4e3219", size = 30371, upload-time = "2025-11-21T23:01:54.787Z" } wheels = [ @@ -1397,8 +1408,8 @@ name = "httpcore2" version = "2.13.1" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "h11", marker = "python_full_version < '3.11' or sys_platform != 'emscripten'" }, - { name = "truststore", marker = "python_full_version < '3.11' or sys_platform != 'emscripten'" }, + { name = "h11" }, + { name = "truststore" }, ] sdist = { url = "https://files.pythonhosted.org/packages/cb/f3/1db7aa2bc2524062192bb0e0323969492d1883152a232fe36eea65f4e35c/httpcore2-2.13.1.tar.gz", hash = "sha256:e0aa977abe17e69a3b820a24542a6fa88702676d83880b8d194dcd18408e5103", size = 68071, upload-time = "2026-09-23T07:47:22.372Z" } wheels = [ @@ -2610,7 +2621,7 @@ name = "nvidia-cublas" version = "13.1.1.3" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-cuda-nvrtc", version = "13.0.88", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32')" }, + { name = "nvidia-cuda-nvrtc", version = "13.0.88", source = { registry = "https://pypi.org/simple" } }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/a7/a1/0bd24ee8c8d03adac032fd2909426a00c88f8c57961b1277ded97f91119f/nvidia_cublas-13.1.1.3-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:b7a210458267ac818974c53038fbec2e969d5c99f305ab15c72522fa9f001dd5", size = 542848918, upload-time = "2026-04-08T18:46:22.985Z" }, @@ -2743,7 +2754,7 @@ name = "nvidia-cudnn-cu13" version = "9.20.0.48" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-cublas", marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32')" }, + { name = "nvidia-cublas" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/56/c5/83384d846b2fd17c44bd499b36c75a45ed4f095fbbb2252294e89cea5c5c/nvidia_cudnn_cu13-9.20.0.48-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:e31454ae00094b0c55319d9d15b6fa2fc50a9e1c0f5c8c80fb75258234e731e1", size = 444574296, upload-time = "2026-03-09T19:28:27.751Z" }, @@ -2774,7 +2785,7 @@ name = "nvidia-cufft" version = "12.0.0.61" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-nvjitlink", marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32')" }, + { name = "nvidia-nvjitlink" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/8b/ae/f417a75c0259e85c1d2f83ca4e960289a5f814ed0cea74d18c353d3e989d/nvidia_cufft-12.0.0.61-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:2708c852ef8cd89d1d2068bdbece0aa188813a0c934db3779b9b1faa8442e5f5", size = 214053554, upload-time = "2025-09-04T08:31:38.196Z" }, @@ -2804,9 +2815,9 @@ name = "nvidia-cusolver" version = "12.0.4.66" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-cublas", marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32')" }, - { name = "nvidia-cusparse", marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32')" }, - { name = "nvidia-nvjitlink", marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32')" }, + { name = "nvidia-cublas" }, + { name = "nvidia-cusparse" }, + { name = "nvidia-nvjitlink" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/c8/c3/b30c9e935fc01e3da443ec0116ed1b2a009bb867f5324d3f2d7e533e776b/nvidia_cusolver-12.0.4.66-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:02c2457eaa9e39de20f880f4bd8820e6a1cfb9f9a34f820eb12a155aa5bc92d2", size = 223467760, upload-time = "2025-09-04T08:33:04.222Z" }, @@ -2818,7 +2829,7 @@ name = "nvidia-cusparse" version = "12.6.3.3" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-nvjitlink", marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32')" }, + { name = "nvidia-nvjitlink" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/f8/94/5c26f33738ae35276672f12615a64bd008ed5be6d1ebcb23579285d960a9/nvidia_cusparse-12.6.3.3-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:80bcc4662f23f1054ee334a15c72b8940402975e0eab63178fc7e670aa59472c", size = 162155568, upload-time = "2025-09-04T08:33:42.864Z" }, @@ -3265,10 +3276,10 @@ resolution-markers = [ "python_full_version < '3.11' and sys_platform != 'darwin'", ] dependencies = [ - { name = "numpy", marker = "python_full_version < '3.11'" }, - { name = "python-dateutil", marker = "python_full_version < '3.11'" }, - { name = "pytz", marker = "python_full_version < '3.11'" }, - { name = "tzdata", marker = "python_full_version < '3.11'" }, + { name = "numpy" }, + { name = "python-dateutil" }, + { name = "pytz" }, + { name = "tzdata" }, ] sdist = { url = "https://files.pythonhosted.org/packages/33/01/d40b85317f86cf08d853a4f495195c73815fdf205eef3993821720274518/pandas-2.3.3.tar.gz", hash = "sha256:e05e1af93b977f7eafa636d043f9f94c7ee3ac81af99c13508215942e64c993b", size = 4495223, upload-time = "2025-09-29T23:34:51.853Z" } wheels = [ @@ -3327,9 +3338,9 @@ resolution-markers = [ "python_full_version == '3.11.*' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32'", ] dependencies = [ - { name = "numpy", marker = "python_full_version >= '3.11'" }, - { name = "python-dateutil", marker = "python_full_version >= '3.11'" }, - { name = "tzdata", marker = "(python_full_version >= '3.11' and sys_platform == 'emscripten') or (python_full_version >= '3.11' and sys_platform == 'win32')" }, + { name = "numpy" }, + { name = "python-dateutil" }, + { name = "tzdata", marker = "sys_platform == 'emscripten' or sys_platform == 'win32'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/be/4f/5f3422a2afec5ffc46308b79e53291365a93748b498ac2e58bead0197916/pandas-3.0.5.tar.gz", hash = "sha256:dca3734d6ab7c906e6730f0788b0a1dbb9f2467731f9711f77995c8e9d62d712", size = 4658219, upload-time = "2026-07-22T22:19:28.819Z" } wheels = [ @@ -4683,10 +4694,10 @@ resolution-markers = [ "python_full_version < '3.11' and sys_platform != 'darwin'", ] dependencies = [ - { name = "joblib", marker = "python_full_version < '3.11'" }, - { name = "numpy", marker = "python_full_version < '3.11'" }, - { name = "scipy", version = "1.15.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, - { name = "threadpoolctl", marker = "python_full_version < '3.11'" }, + { name = "joblib" }, + { name = "numpy" }, + { name = "scipy", version = "1.15.3", source = { registry = "https://pypi.org/simple" } }, + { name = "threadpoolctl" }, ] sdist = { url = "https://files.pythonhosted.org/packages/98/c2/a7855e41c9d285dfe86dc50b250978105dce513d6e459ea66a6aeb0e1e0c/scikit_learn-1.7.2.tar.gz", hash = "sha256:20e9e49ecd130598f1ca38a1d85090e1a600147b9c02fa6f15d69cb53d968fda", size = 7193136, upload-time = "2025-09-09T08:21:29.075Z" } wheels = [ @@ -4736,12 +4747,12 @@ resolution-markers = [ "python_full_version == '3.11.*' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32'", ] dependencies = [ - { name = "joblib", marker = "python_full_version >= '3.11'" }, - { name = "narwhals", marker = "python_full_version >= '3.11'" }, - { name = "numpy", marker = "python_full_version >= '3.11'" }, - { name = "scipy", version = "1.17.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version == '3.11.*'" }, + { name = "joblib" }, + { name = "narwhals" }, + { name = "numpy" }, + { name = "scipy", version = "1.17.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" }, { name = "scipy", version = "1.18.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" }, - { name = "threadpoolctl", marker = "python_full_version >= '3.11'" }, + { name = "threadpoolctl" }, ] sdist = { url = "https://files.pythonhosted.org/packages/d2/eb/eaf5e07fcc0da7149b0e084f24e54edd7441b9a89ce7e034032ae97fe3a0/scikit_learn-1.9.1.tar.gz", hash = "sha256:629cada3e33e2b9bf376cdc7614a47a4140b8aedc1d836579e359736fbd82977", size = 7786908, upload-time = "2026-09-10T18:34:04.679Z" } wheels = [ @@ -4774,7 +4785,7 @@ resolution-markers = [ "python_full_version < '3.11' and sys_platform != 'darwin'", ] dependencies = [ - { name = "numpy", marker = "python_full_version < '3.11'" }, + { name = "numpy" }, ] sdist = { url = "https://files.pythonhosted.org/packages/0f/37/6964b830433e654ec7485e45a00fc9a27cf868d622838f6b6d9c5ec0d532/scipy-1.15.3.tar.gz", hash = "sha256:eae3cf522bc7df64b42cad3925c876e1b0b6c35c1337c93e12c0f366f55b0eaf", size = 59419214, upload-time = "2025-05-08T16:13:05.955Z" } wheels = [ @@ -4836,7 +4847,7 @@ resolution-markers = [ "python_full_version == '3.11.*' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32'", ] dependencies = [ - { name = "numpy", marker = "python_full_version == '3.11.*'" }, + { name = "numpy" }, ] sdist = { url = "https://files.pythonhosted.org/packages/7a/97/5a3609c4f8d58b039179648e62dd220f89864f56f7357f5d4f45c29eb2cc/scipy-1.17.1.tar.gz", hash = "sha256:95d8e012d8cb8816c226aef832200b1d45109ed4464303e997c5b13122b297c0", size = 30573822, upload-time = "2026-02-23T00:26:24.851Z" } wheels = [ @@ -4897,7 +4908,7 @@ resolution-markers = [ "python_full_version == '3.12.*' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32'", ] dependencies = [ - { name = "numpy", marker = "python_full_version >= '3.12'" }, + { name = "numpy" }, ] sdist = { url = "https://files.pythonhosted.org/packages/7e/74/66de6258867beb2ef08f35f9f2ac017a52cacd5081714d239ff1a442d458/scipy-1.18.1.tar.gz", hash = "sha256:52c4b7422442aba924d03ad4019852b08a92e64ea187b933135687bfe2747307", size = 30781235, upload-time = "2026-08-21T23:28:50.599Z" } wheels = [ @@ -5306,10 +5317,10 @@ name = "tokenspeed-mla" version = "0.1.8" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "apache-tvm-ffi", marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32')" }, - { name = "nvidia-cutlass-dsl", marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32')" }, - { name = "tokenspeed-triton", marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32')" }, - { name = "torch", marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32')" }, + { name = "apache-tvm-ffi" }, + { name = "nvidia-cutlass-dsl" }, + { name = "tokenspeed-triton" }, + { name = "torch" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/1e/65/81d7e9f14472bc4c6abb576c9b1edd8e40ab01832027c4e647bfe2890749/tokenspeed_mla-0.1.8-py3-none-manylinux_2_28_aarch64.whl", hash = "sha256:952209cf4b29a54e6b6e7e088be9d4a40f24792b06fdfa330de8077b9926a7d9", size = 752341, upload-time = "2026-06-24T03:36:28.619Z" }, @@ -5841,6 +5852,9 @@ name = "vllm-afd-plugin" source = { editable = "." } [package.optional-dependencies] +afdep = [ + { name = "afdep" }, +] vllm = [ { name = "vllm" }, ] @@ -5860,8 +5874,11 @@ e2e-tests = [ ] [package.metadata] -requires-dist = [{ name = "vllm", marker = "extra == 'vllm'", specifier = "==0.30.0" }] -provides-extras = ["vllm"] +requires-dist = [ + { name = "afdep", marker = "extra == 'afdep'", git = "https://github.com/ms-llmd/afdEP.git?rev=e763e3189398beda24f592d2117e6ade135b0a7f" }, + { name = "vllm", marker = "extra == 'vllm'", specifier = "==0.30.0" }, +] +provides-extras = ["vllm", "afdep"] [package.metadata.requires-dev] dev = [ From 259bfe5be1f648684fa11138a52d154f4bf072fc Mon Sep 17 00:00:00 2001 From: ronenkat <16743404+ronenkat@users.noreply.github.com> Date: Sun, 25 Oct 2026 13:54:28 +0200 Subject: [PATCH 02/32] test(gpu): cover the DeepEP M2N connector on the afdEP transport Add the DeepEP M2N connector, routed-experts runner and model tests, and move device tests out of tests/unit. - Connector unit tests (config, stage plans, handshake, control plane, DBO lane streams) run against a stand-in afdEP module (tests/unit/fakes/afdep.py) where afdEP is not installed. They cover the merged role_aware option and that each lane builds afdEP's ElasticBuffer with the protocol and FFN rank count. - Network planner, connector device placement, routed FFN runner, worker bootstrap, validation, topology, and Qwen3 / DeepSeek V2 routing-proxy and weight-policy tests. - GPU tests move to tests/gpu/ (gpu marker), with a masked-protocol (role_aware: false) case; the M2N vLLM-config fakes they share with the unit tests live in tests/unit/fakes/vllm_config.py, the recording routed connector in tests/unit/fakes/connector.py. - Qwen3 DeepEP M2N E2E scenarios (2A2F eager, graph, graph+DBO, Attention TP2) in the E2E runner. - validation_paths entries for the new test files. Co-Authored-By: Claude Opus 5.5 Signed-off-by: ronenkat <16743404+ronenkat@users.noreply.github.com> --- docs/design/module/connector_contracts.md | 2 + docs/design/module/ffn_runtime.md | 1 + tests/e2e/models/qwen3_moe/test_qwen3_moe.py | 20 +- tests/e2e/runner.py | 29 + tests/gpu/__init__.py | 3 + tests/gpu/deepep_m2n_harness.py | 1561 ++++++++++++++++ tests/gpu/test_deepep_m2n_gpu.py | 415 +++++ tests/unit/config/test_validation.py | 15 + tests/unit/connectors/test_base_factory.py | 151 +- .../connectors/test_deepep_m2n_connector.py | 1599 +++++++++++++++++ .../connectors/test_deepep_m2n_network.py | 380 ++++ tests/unit/distributed/test_topology.py | 73 +- tests/unit/fakes/afdep.py | 35 + tests/unit/fakes/connector.py | 82 + tests/unit/fakes/vllm_config.py | 114 ++ .../models/test_deepseek_v2_proxy.py | 292 +++ .../models/test_deepseek_v2_weight_policy.py | 27 + .../models/test_qwen3_moe_construction.py | 170 +- .../models/test_qwen3_moe_weight_policy.py | 89 +- tests/unit/package/test_package.py | 1 + tests/unit/v1/worker/test_connector_device.py | 159 ++ .../test_ffn_routed_experts_model_runner.py | 312 ++++ tests/unit/v1/worker/test_model_runner_v2.py | 285 +++ .../unit/v1/worker/test_runtime_classpaths.py | 42 + 24 files changed, 5846 insertions(+), 11 deletions(-) create mode 100644 tests/gpu/__init__.py create mode 100644 tests/gpu/deepep_m2n_harness.py create mode 100644 tests/gpu/test_deepep_m2n_gpu.py create mode 100644 tests/unit/connectors/test_deepep_m2n_connector.py create mode 100644 tests/unit/connectors/test_deepep_m2n_network.py create mode 100644 tests/unit/fakes/afdep.py create mode 100644 tests/unit/fakes/connector.py create mode 100644 tests/unit/v1/worker/test_connector_device.py create mode 100644 tests/unit/v1/worker/test_ffn_routed_experts_model_runner.py diff --git a/docs/design/module/connector_contracts.md b/docs/design/module/connector_contracts.md index 80a227186..450e9080f 100644 --- a/docs/design/module/connector_contracts.md +++ b/docs/design/module/connector_contracts.md @@ -22,6 +22,8 @@ validation_paths: - "tests/e2e/models/deepseek_v2_lite/test_deepseek_v2_lite.py" - "tests/e2e/models/deepseek_v2_lite/test_async_cam_npu.py" - "tests/unit/distributed/**" + - "tests/unit/v1/worker/test_connector_device.py" + - "tests/gpu/**" upstream_refs: - "vLLM vllm.forward_context.DPMetadata" - "vLLM vllm.distributed.parallel_state" diff --git a/docs/design/module/ffn_runtime.md b/docs/design/module/ffn_runtime.md index bb28bd4b4..ae4f4df40 100644 --- a/docs/design/module/ffn_runtime.md +++ b/docs/design/module/ffn_runtime.md @@ -26,6 +26,7 @@ depends_on: validation_paths: - "tests/unit/v1/worker/test_ffn_metadata.py" - "tests/unit/v1/worker/test_ffn_model_runner.py" + - "tests/unit/v1/worker/test_ffn_routed_experts_model_runner.py" - "tests/unit/v1/worker/test_model_runner_v2.py" - "tests/unit/v1/worker/test_npu_runtime.py" - "tests/unit/v1/worker/test_async_ffn_two_stream.py" diff --git a/tests/e2e/models/qwen3_moe/test_qwen3_moe.py b/tests/e2e/models/qwen3_moe/test_qwen3_moe.py index b3b3091cd..ec02ccab9 100644 --- a/tests/e2e/models/qwen3_moe/test_qwen3_moe.py +++ b/tests/e2e/models/qwen3_moe/test_qwen3_moe.py @@ -17,6 +17,7 @@ preserve_environment_variable, run_runner, ) +from tests.e2e.runner import DEEPEP_M2N_SCENARIOS GSM8K_DATASET_ID = "openai/gsm8k" GSM8K_DATASET_CONFIG = "main" @@ -25,6 +26,7 @@ DEFAULT_DEVICE_IDS = ("0", "1", "2", "3") ATTENTION_DEVICE_COUNT = 2 AFD_FFN_DEVICE_COUNT = 1 +DEEPEP_M2N_FFN_DEVICE_COUNT = 2 BASELINE_DEVICE_COUNT = 4 SCENARIOS = ( "baseline-graph", @@ -32,6 +34,10 @@ "afd-graph-2a1f", "afd-graph-dbo-2a1f", ) +# DeepEP M2N (routing on Attention) needs EP over at least two FFN ranks, so +# its scenarios (DEEPEP_M2N_SCENARIOS) run 2A2F. CI selects test ids +# explicitly and does not run them; running the whole module (pytest -m e2e +# tests/e2e/models/qwen3_moe) does, and they need DeepEP and 4 GPUs. def _required_env(name: str) -> str: @@ -74,9 +80,14 @@ def build_runner_command(scenario: str, gsm8k_output_path: Path) -> list[str]: attention_devices = devices[:BASELINE_DEVICE_COUNT] ffn_devices = [] else: + ffn_device_count = ( + DEEPEP_M2N_FFN_DEVICE_COUNT + if scenario in DEEPEP_M2N_SCENARIOS + else AFD_FFN_DEVICE_COUNT + ) attention_devices = devices[:ATTENTION_DEVICE_COUNT] ffn_devices = devices[ - ATTENTION_DEVICE_COUNT : ATTENTION_DEVICE_COUNT + AFD_FFN_DEVICE_COUNT + ATTENTION_DEVICE_COUNT : ATTENTION_DEVICE_COUNT + ffn_device_count ] command = [ @@ -121,3 +132,10 @@ def _prepare_e2e_assets() -> Iterator[None]: def test_qwen3_moe(scenario: str, tmp_path: Path) -> None: command = build_runner_command(scenario, tmp_path / scenario) run_runner(command) + + +@pytest.mark.e2e +@pytest.mark.parametrize("scenario", DEEPEP_M2N_SCENARIOS, ids=DEEPEP_M2N_SCENARIOS) +def test_qwen3_moe_deepep_m2n(scenario: str, tmp_path: Path) -> None: + command = build_runner_command(scenario, tmp_path / scenario) + run_runner(command) diff --git a/tests/e2e/runner.py b/tests/e2e/runner.py index 6dd556021..bac0ea998 100644 --- a/tests/e2e/runner.py +++ b/tests/e2e/runner.py @@ -67,6 +67,22 @@ "afd-v2-eager-dbo-dp2", "afd-v2-graph-dbo-dp2", ) +# DeepEP M2N routes on Attention (gate and top-k there) and dispatches each +# token straight to the FFN ranks owning its experts. +DEEPEP_M2N_CONNECTOR = "DeepEPM2NAFDConnector" +DEEPEP_M2N_SCENARIOS = ( + "afd-deepep-m2n-eager-2a2f", + "afd-deepep-m2n-graph-2a2f", + "afd-deepep-m2n-graph-dbo-2a2f", + "afd-deepep-m2n-eager-tp2-2a2f", + "afd-deepep-m2n-graph-tp2-2a2f", +) +# Attention runs DP1 x TP2: each TP rank dispatches its slice of the tokens. +# FFN stays DP2 x TP1 (experts are not TP-sharded under DeepEP M2N). +DEEPEP_M2N_ATTENTION_TP_SCENARIOS = frozenset( + ("afd-deepep-m2n-eager-tp2-2a2f", "afd-deepep-m2n-graph-tp2-2a2f"), +) +DEEPEP_M2N_ATTENTION_TP_SIZE = 2 V2_DBO_COMPARISON_SCENARIOS = frozenset( ("afd-v2-eager-dp2", "afd-v2-eager-dbo-dp2", "afd-v2-graph-dbo-dp2"), ) @@ -327,6 +343,7 @@ def parse_args() -> argparse.Namespace: ASYNC_UBATCH_SCENARIO, *DSV4_SCENARIOS, *V2_SCENARIOS, + *DEEPEP_M2N_SCENARIOS, ], required=True, help="Fixed E2E scenario to run.", @@ -477,6 +494,11 @@ def configure_scenario(args: argparse.Namespace) -> None: "afd-v2-graph-tp2": (False, True, False, 2, 2), "afd-v2-eager-dbo-dp2": (False, False, True, 2, 2), "afd-v2-graph-dbo-dp2": (False, True, True, 2, 2), + "afd-deepep-m2n-eager-2a2f": (False, False, False, 2, 2), + "afd-deepep-m2n-graph-2a2f": (False, True, False, 2, 2), + "afd-deepep-m2n-graph-dbo-2a2f": (False, True, True, 2, 2), + "afd-deepep-m2n-eager-tp2-2a2f": (False, False, False, 2, 2), + "afd-deepep-m2n-graph-tp2-2a2f": (False, True, False, 2, 2), } for sync_scenario, sync_profile in DSV4_SYNC_SHAPES.items(): # The A5 script's native DBO stays off for these scenarios: a split batch @@ -509,6 +531,8 @@ def configure_scenario(args: argparse.Namespace) -> None: args.attention_tp_size = ASYNC_UBATCH_ATTENTION_TP_SIZE elif args.scenario in V2_TENSOR_PARALLEL_SCENARIOS: args.attention_tp_size = 2 + elif args.scenario in DEEPEP_M2N_ATTENTION_TP_SCENARIOS: + args.attention_tp_size = DEEPEP_M2N_ATTENTION_TP_SIZE else: args.attention_tp_size = 1 if args.scenario in V2_TENSOR_PARALLEL_SCENARIOS: @@ -556,6 +580,9 @@ def configure_scenario(args: argparse.Namespace) -> None: for arg in args.common_vllm_arg ): args.common_vllm_arg.extend(["--gpu-memory-utilization", "0.8"]) + if args.scenario in DEEPEP_M2N_SCENARIOS: + args.afd_connector = DEEPEP_M2N_CONNECTOR + args.compute_gate_on_attention = True if args.scenario == DSV4_ASYNC_CAM_SCENARIO: dsv4_config.configure_scenario(args) elif args.scenario in DSV4_SYNC_CAMP2P_SCENARIOS: @@ -607,6 +634,8 @@ def validate_topology( return if args.scenario in V2_SCENARIOS and args.device_backend != "gpu": raise ValueError("ModelRunnerV2 E2E scenarios require GPU") + if args.scenario in DEEPEP_M2N_SCENARIOS and args.device_backend != "gpu": + raise ValueError("DeepEP M2N E2E scenarios require GPU") if ( args.scenario in (ASYNC_CAM_SCENARIO, ASYNC_UBATCH_SCENARIO) and args.device_backend != "npu" diff --git a/tests/gpu/__init__.py b/tests/gpu/__init__.py new file mode 100644 index 000000000..198dc66cc --- /dev/null +++ b/tests/gpu/__init__.py @@ -0,0 +1,3 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Opt-in GPU tests that need real devices and connector transports.""" diff --git a/tests/gpu/deepep_m2n_harness.py b/tests/gpu/deepep_m2n_harness.py new file mode 100644 index 000000000..21a835b37 --- /dev/null +++ b/tests/gpu/deepep_m2n_harness.py @@ -0,0 +1,1561 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Multi-process GPU harness for ``DeepEPM2NAFDConnector`` tests and benchmarks. + +``run_ranks`` spawns one process per union rank (FFN ranks first, then +Attention ranks), each on its own GPU. Every rank builds the connector through +``AFDConnectorFactory`` from a fake ``vllm_config``, initializes it, and runs +a scenario; results come back as JSON-able dicts. + +Inputs are deterministic, so every rank can regenerate every Attention rank's +tokens and routing: FFN ranks check exactly what they received, and Attention +ranks compare the combined output against a dense MoE reference. +""" + +from __future__ import annotations + +import functools +import gc +import json +import os +import random +import socket +import time +import traceback +from collections.abc import Callable, Mapping +from contextlib import nullcontext +from dataclasses import asdict, dataclass, field, replace +from pathlib import Path +from typing import Any + +import torch +import torch.distributed as dist +import torch.multiprocessing as mp + +from tests.unit.fakes.vllm_config import ( + M2N_HIDDEN_SIZE, + M2N_INTERMEDIATE_SIZE, + M2N_NUM_EXPERTS, + M2N_TOP_K, + M2NTopology, + m2n_vllm_config, +) + +INPUT_SEED_STRIDE = 1_000_003 +WEIGHT_SEED_OFFSET = 7_777 +# Limits are about 1.3x the worst value seen on H100 (2026-10-15), so a real +# regression fails. Rows are matched exactly; the per-expert sums only differ +# by fp32 summation order. +CHECKSUM_REL_TOL = 1e-5 +BF16_TOLERANCES = { + "dispatch_rows_rel_err": CHECKSUM_REL_TOL, + "dispatch_weights_rel_err": CHECKSUM_REL_TOL, + "experts_rel_err": 9e-3, # worst seen 6.5e-3 + "experts_vs_fused_rel_err": 7e-3, # worst seen 4.8e-3 + "combine_rel_err": 8e-3, # worst seen 5.9e-3 +} +# FP8: the experts' references use the same FP8 inputs and weights, so they +# only see the requantized intermediate; combine is checked end to end against +# the BF16 model and also sees activation and weight quantization. +FP8_TOLERANCES = { + "dispatch_rows_rel_err": CHECKSUM_REL_TOL, + "dispatch_weights_rel_err": CHECKSUM_REL_TOL, + "experts_rel_err": 4e-2, # worst seen 2.84e-2 + "experts_vs_fused_rel_err": 4e-2, # worst seen 2.87e-2 + "combine_rel_err": 9e-2, # worst seen 6.87e-2 +} +RANK_TIMEOUT_S = 900 +# Rendezvous ports come from [PORT_RANGE_LOW, ephemeral range start). +PORT_RANGE_LOW = 20_000 +PORT_RANGE_HIGH = 32_768 +DEFAULT_EPHEMERAL_PORT_LOW = 32_768 +EPHEMERAL_PORT_RANGE = Path("/proc/sys/net/ipv4/ip_local_port_range") +CLOSE_FAILED = "\nclose() failed:\n" +NO_RESULT = "no result" +PEER_EXIT_CODE = 17 +ROW_KEY_SEED = 4_242 +ROW_KEY_MAX_COEFFICIENT = 1 << 31 +BENCH_WARMUP_ITERS = 5 +# Seeds reserve room for this many stages (lanes) per step. +MAX_STAGES = 2 + + +@dataclass(frozen=True) +class Case: + """One step: a control payload, then ``num_layers`` MoE layers under it.""" + + name: str + # One count per Attention DP rank (per Attention rank without TP); the + # control payload carries them as vLLM's DP metadata does. + tokens_per_attn_rank: tuple[int, ...] + # Graph-mode payloads skip the host sync (worst-case receive allocation). + graph_mode: bool = False + num_layers: int = 2 + # Second stage (lane 1) token counts for two-lane (DBO) cases. + second_stage: tuple[int, ...] | None = None + # Route every top-k choice to FFN rank 0's experts (worst-case receive). + hot_ffn_rank: bool = False + # FFN ranks plan the step as graph (True) or eager (False) regardless of + # the payload: one side replaying a graph while the other runs eagerly. + ffn_graph_mode: bool | None = None + + def stage_tokens(self) -> list[tuple[int, ...]]: + if self.second_stage is None: + return [self.tokens_per_attn_rank] + return [self.tokens_per_attn_rank, self.second_stage] + + +@dataclass(frozen=True) +class Placement: + """Which world ranks this host runs and where the ranks rendezvous. + + Single-node runs use the defaults. Multi-node runs launch one ``run_ranks`` + per node with that node's ``local_ranks`` (contiguous; local index = + GPU index), the same ``master_addr`` (node 0) and the same fixed ports. + """ + + master_addr: str = "127.0.0.1" + ports: tuple[int, int] | None = None + local_ranks: tuple[int, ...] | None = None + + +@dataclass +class RankResult: + world_rank: int + role: str + ok: bool = True + error: str = "" + metrics: dict[str, Any] = field(default_factory=dict) + + +def standard_cases(attn_size: int, capacity: int) -> list[Case]: + cases: list[Case] = [] + for graph_mode in (False, True): + tag = "nosync" if graph_mode else "sync" + uneven = (capacity,) + (0,) * (attn_size - 1) + mixed = tuple(17 + 5 * rank for rank in range(attn_size)) + cases += [ + Case(f"even64_{tag}", (64,) * attn_size, graph_mode), + Case(f"uneven_{tag}", uneven, graph_mode), + Case(f"decode1_{tag}", (1,) * attn_size, graph_mode), + Case(f"empty_{tag}", (0,) * attn_size, graph_mode), + Case(f"mixed_{tag}", mixed, graph_mode), + ] + return cases + + +def hot_cases(attn_size: int, capacity: int) -> list[Case]: + """Standard cases with all routes on FFN rank 0's experts.""" + return [ + replace(case, name=f"hot_{case.name}", hot_ffn_rank=True) + for case in standard_cases(attn_size, capacity) + ] + + +def mixed_mode_cases(attn_size: int) -> list[Case]: + """Attention and FFN plan the same step differently (graph vs eager).""" + cases: list[Case] = [] + for attention_graph in (True, False): + tag = "attn_graph_ffn_eager" if attention_graph else "attn_eager_ffn_graph" + cases += [ + Case( + f"{tag}_{name}", + tokens, + graph_mode=attention_graph, + ffn_graph_mode=not attention_graph, + ) + for name, tokens in ( + ("even64", (64,) * attn_size), + ("decode1", (1,) * attn_size), + ("mixed", tuple(17 + 5 * rank for rank in range(attn_size))), + ) + ] + return cases + + +def lane_cases(attn_size: int, capacity: int) -> list[Case]: + """Two-lane (DBO) cases: different token counts per lane, incl. empty lanes.""" + cases: list[Case] = [] + for graph_mode in (False, True): + tag = "nosync" if graph_mode else "sync" + cases += [ + Case( + f"lanes_even_{tag}", (64,) * attn_size, graph_mode, 2, (32,) * attn_size + ), + Case( + f"lanes_uneven_{tag}", + (capacity,) + (0,) * (attn_size - 1), + graph_mode, + 2, + (1,) * attn_size, + ), + Case( + f"lanes_empty_{tag}", (0,) * attn_size, graph_mode, 2, (17,) * attn_size + ), + ] + return cases + + +# --------------------------------------------------------------------------- +# Deterministic data +# --------------------------------------------------------------------------- + + +def make_inputs( + seed: int, attn_rank: int, num_tokens: int, expert_limit: int = M2N_NUM_EXPERTS +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Hidden states and Qwen3-style (softmax, renormalized) top-k routing. + + ``expert_limit`` restricts routing to experts ``[0, expert_limit)``; the + random draws are the same either way. + """ + gen = torch.Generator(device="cuda") + gen.manual_seed(seed * INPUT_SEED_STRIDE + attn_rank) + x = torch.randn(num_tokens, M2N_HIDDEN_SIZE, generator=gen, device="cuda") + logits = torch.randn(num_tokens, M2N_NUM_EXPERTS, generator=gen, device="cuda") + logits[:, expert_limit:] = float("-inf") + weights, ids = torch.topk(torch.softmax(logits, dim=-1), M2N_TOP_K, dim=-1) + weights = weights / weights.sum(dim=-1, keepdim=True) + return x.to(torch.bfloat16), ids, weights.float() + + +def attention_rank_inputs( + seed: int, + attn_rank: int, + attn_tp_size: int, + dp_tokens: tuple[int, ...], + expert_limit: int = M2N_NUM_EXPERTS, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]: + """What one Attention rank dispatches, and how many of its rows are real. + + TP peers regenerate their DP rank's tokens and each takes rows + ``[r * c, (r + 1) * c)`` of them padded to a multiple of the TP size + (``c = ceil(N / TP)``), as the Qwen3 routing proxy does. Padding rows come + last and route nowhere (expert id -1). Without TP this is ``make_inputs``. + """ + dp_rank, tp_rank = divmod(attn_rank, attn_tp_size) + num_tokens = dp_tokens[dp_rank] + x, ids, w = make_inputs(seed, dp_rank, num_tokens, expert_limit) + if attn_tp_size == 1: + return x, ids, w, num_tokens + rows = -(-num_tokens // attn_tp_size) + start = tp_rank * rows + real = max(0, min(rows, num_tokens - start)) + padding = rows - real + end = start + real + return ( + torch.cat([x[start:end], x.new_zeros(padding, M2N_HIDDEN_SIZE)]), + torch.cat([ids[start:end], ids.new_full((padding, M2N_TOP_K), -1)]), + torch.cat([w[start:end], w.new_zeros(padding, M2N_TOP_K)]), + real, + ) + + +def make_expert_weights(experts: range) -> tuple[torch.Tensor, torch.Tensor]: + w13 = torch.empty( + len(experts), + 2 * M2N_INTERMEDIATE_SIZE, + M2N_HIDDEN_SIZE, + dtype=torch.bfloat16, + device="cuda", + ) + w2 = torch.empty( + len(experts), + M2N_HIDDEN_SIZE, + M2N_INTERMEDIATE_SIZE, + dtype=torch.bfloat16, + device="cuda", + ) + for index, expert in enumerate(experts): + gen = torch.Generator(device="cuda") + gen.manual_seed(WEIGHT_SEED_OFFSET + expert) + w13[index] = ( + torch.randn( + 2 * M2N_INTERMEDIATE_SIZE, M2N_HIDDEN_SIZE, generator=gen, device="cuda" + ) + / M2N_HIDDEN_SIZE**0.5 + ) + w2[index] = ( + torch.randn( + M2N_HIDDEN_SIZE, M2N_INTERMEDIATE_SIZE, generator=gen, device="cuda" + ) + / M2N_INTERMEDIATE_SIZE**0.5 + ) + return w13, w2 + + +def expert_mlp_fp32( + x: torch.Tensor, w13: torch.Tensor, w2: torch.Tensor +) -> torch.Tensor: + gate, up = (x.float() @ w13.float().t()).chunk(2, dim=-1) + return (torch.nn.functional.silu(gate) * up) @ w2.float().t() + + +def dense_moe_reference( + x: torch.Tensor, + ids: torch.Tensor, + weights: torch.Tensor, + w13: torch.Tensor, + w2: torch.Tensor, +) -> torch.Tensor: + ref = torch.zeros(x.shape, dtype=torch.float32, device=x.device) + for expert in ids.unique().tolist(): + rows, slots = (ids == expert).nonzero(as_tuple=True) + ref.index_add_( + 0, + rows, + expert_mlp_fp32(x[rows], w13[expert], w2[expert]) + * weights[rows, slots, None], + ) + return ref + + +def routing_limit(case: Case, layout: Any) -> int: + """Routing range for ``case``: FFN rank 0's experts when it is hot.""" + return layout.experts_per_rank if case.hot_ffn_rank else M2N_NUM_EXPERTS + + +def row_keys(rows: torch.Tensor) -> torch.Tensor: + """An exact 64-bit fingerprint per fp32 row (bit pattern, wrapping sum). + + Equal rows give equal keys whatever the batch shape, so received rows can + be matched against sent rows without depending on their order. + """ + gen = torch.Generator(device=rows.device) + gen.manual_seed(ROW_KEY_SEED) + coefficients = torch.randint( + 1, ROW_KEY_MAX_COEFFICIENT, (rows.shape[-1],), generator=gen, device=rows.device + ) + return (rows.contiguous().view(torch.int32).long() * coefficients).sum(-1) + + +def sort_rows( + local_ids: torch.Tensor, keys: torch.Tensor, weights: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Order rows by (local expert, key) for an order-independent comparison.""" + by_key = torch.argsort(keys, stable=True) + order = by_key[torch.argsort(local_ids[by_key], stable=True)] + return local_ids[order], keys[order], weights[order] + + +def rel_err(actual: torch.Tensor, expected: torch.Tensor) -> float: + denom = expected.float().norm().item() + diff = (actual.float() - expected.float()).norm().item() + return diff / denom if denom > 0 else diff + + +# --------------------------------------------------------------------------- +# Process plumbing +# --------------------------------------------------------------------------- + + +def free_ports(count: int) -> tuple[int, ...]: + """Distinct free ports below the kernel's ephemeral range. + + The ports are released before the ranks bind them. A port from the + ephemeral range (what binding port 0 returns) can be taken in between by + any outgoing connection, e.g. another test's NCCL or gloo sockets, which + fails the rendezvous with EADDRINUSE. Below that range only explicit binds + compete. The sockets stay open until all are picked, so one call cannot + return the same port twice. + """ + rng = random.SystemRandom() + high = min(_ephemeral_port_low(), PORT_RANGE_HIGH) + if high <= PORT_RANGE_LOW: + high = PORT_RANGE_HIGH + sockets: list[socket.socket] = [] + try: + while len(sockets) < count: + sock = socket.socket() + try: + sock.bind(("", rng.randrange(PORT_RANGE_LOW, high))) + except OSError: + sock.close() + continue + sockets.append(sock) + return tuple(int(sock.getsockname()[1]) for sock in sockets) + finally: + for sock in sockets: + sock.close() + + +def _ephemeral_port_low() -> int: + try: + return int(EPHEMERAL_PORT_RANGE.read_text().split()[0]) + except (OSError, ValueError, IndexError): + return DEFAULT_EPHEMERAL_PORT_LOW + + +def make_connector( + topology: M2NTopology, + world_rank: int, + local_rank: int, + extra_config: dict[str, Any], + port: int, + host: str, +) -> Any: + """The connector as the runtime constructs it (not yet initialized).""" + from afd_plugin.config import afd_config_from_mapping + from afd_plugin.connectors import AFDConnectorFactory + + role, role_rank = topology.role(world_rank) + vllm_config = m2n_vllm_config(topology, role, role_rank, extra_config, port, host) + afd_config = afd_config_from_mapping(vllm_config.additional_config["afd"]) + # No vLLM TP group here: pass the TP rank, as workers do before init. + return AFDConnectorFactory.create_connector( + world_rank, + local_rank, + vllm_config, + afd_config, + tp_rank=role_rank % vllm_config.parallel_config.tensor_parallel_size, + ) + + +def control_step( + connector: Any, + tokens_by_stage: list[tuple[int, ...]], + *, + graph_mode: bool = False, + ffn_graph_mode: bool | None = None, +) -> None: + """Deliver one step's control payload like the runtime does. + + ``tokens_by_stage[s]`` holds every Attention rank's token count for stage + (lane) ``s``. ``graph_mode`` marks the step as graph capture/replay (no + host sync, fixed worst-case shapes). ``ffn_graph_mode`` makes FFN ranks + plan the step in the other mode than the payload says, as an FFN runner + does when it has no graph for a step Attention replays. + """ + from afd_plugin.connectors import AFDControlPayload, AFDDPMetadata + + if connector.is_ffn: + payload = connector.control_plane.recv_dp_metadata_list() + if ffn_graph_mode is not None: + payload = replace( + payload, is_graph_capturing=False, is_graph_replaying=ffn_graph_mode + ) + else: + payload = AFDControlPayload( + dp_metadata_list={ + stage: AFDDPMetadata(num_tokens_across_dp_cpu=torch.tensor(tokens)) + for stage, tokens in enumerate(tokens_by_stage) + }, + is_graph_capturing=False, + is_warmup=False, + is_graph_replaying=graph_mode, + ) + connector.control_plane.send_dp_metadata_list(payload) + connector.control_plane.update_state_from_dp_metadata(payload) + + +def _rank_main( + world_rank: int, + local_rank: int, + topology: M2NTopology, + scenario: Callable[[Any, M2NTopology, dict[str, Any]], dict[str, Any]], + scenario_args: dict[str, Any], + extra_config: dict[str, Any], + master_addr: str, + ports: tuple[int, int], + out_dir: str, +) -> None: + role, _ = topology.role(world_rank) + result = RankResult(world_rank=world_rank, role=role) + connector = None + try: + torch.cuda.set_device(local_rank) + # vLLM workers always have a default group; the AFD world is separate. + dist.init_process_group( + "gloo", + init_method=f"tcp://{master_addr}:{ports[0]}", + world_size=topology.world_size, + rank=world_rank, + ) + connector = make_connector( + topology, world_rank, local_rank, extra_config, ports[1], master_addr + ) + # Separate from construction so close() also runs after a failed init. + connector.init_afd_connector() + result.metrics = scenario(connector, topology, scenario_args) + except Exception: + result.ok = False + result.error = traceback.format_exc() + finally: + if connector is not None: + try: + connector.close() + except Exception: + result.ok = False + result.error += CLOSE_FAILED + traceback.format_exc() + Path(out_dir, f"rank{world_rank}.json").write_text( + json.dumps(asdict(result), default=str), encoding="utf-8" + ) + if dist.is_initialized(): + dist.destroy_process_group() + + +def run_ranks( + topology: M2NTopology, + scenario: Callable[[Any, M2NTopology, dict[str, Any]], dict[str, Any]], + scenario_args: dict[str, Any], + extra_config: dict[str, Any], + out_dir: Path, + placement: Placement | None = None, + rank_extra_config: Mapping[int, dict[str, Any]] | None = None, + timeout_s: float = RANK_TIMEOUT_S, +) -> list[RankResult]: + """Run ``scenario`` on this host's ranks; results come back in rank order. + + ``rank_extra_config`` adds per-rank config keys on top of ``extra_config`` + (to make ranks disagree). Ranks still alive after ``timeout_s`` are killed + and reported as ``no result (hung=True)``. + """ + placement = placement or Placement() + rank_extra_config = rank_extra_config or {} + ctx = mp.get_context("spawn") + ports = placement.ports or free_ports(2) + local_ranks = placement.local_ranks or tuple(range(topology.world_size)) + procs = [ + ctx.Process( + target=_rank_main, + args=( + rank, + local_index, + topology, + scenario, + scenario_args, + {**extra_config, **rank_extra_config.get(rank, {})}, + placement.master_addr, + ports, + str(out_dir), + ), + ) + for local_index, rank in enumerate(local_ranks) + ] + for proc in procs: + proc.start() + deadline = time.monotonic() + timeout_s + for proc in procs: + proc.join(timeout=max(1.0, deadline - time.monotonic())) + hung = [proc for proc in procs if proc.is_alive()] + for proc in hung: + proc.kill() + results = [] + for rank in local_ranks: + path = out_dir / f"rank{rank}.json" + if path.exists(): + results.append(RankResult(**json.loads(path.read_text(encoding="utf-8")))) + else: + role, _ = topology.role(rank) + results.append( + RankResult( + rank, role, ok=False, error=f"{NO_RESULT} (hung={bool(hung)})" + ) + ) + return results + + +# --------------------------------------------------------------------------- +# Scenarios (run inside each rank) +# --------------------------------------------------------------------------- + + +def local_weights(experts: range, payload_format: str) -> tuple[Any, Any, Any]: + """Expert weights for ``experts`` plus fp32 references ``(w13, w2)``. + + FP8 payloads use FP8 block-128 weights; the references are then the + dequantized weights, so kernel checks isolate the kernels' own error. + """ + from afd_plugin.connectors.gpu.deepep_m2n.fp8_quant import ( + dequantize_fp8_block128, + quantize_fp8_block128, + ) + from afd_plugin.connectors.routed_experts import LocalExpertWeights + + w13, w2 = make_expert_weights(experts) + if payload_format == "bf16": + return LocalExpertWeights(w13=w13, w2=w2), w13.float(), w2.float() + w13_q, w13_s = quantize_fp8_block128(w13) + w2_q, w2_s = quantize_fp8_block128(w2) + return ( + LocalExpertWeights(w13=w13_q, w2=w2_q, w13_scale=w13_s, w2_scale=w2_s), + dequantize_fp8_block128(w13_q, w13_s), + dequantize_fp8_block128(w2_q, w2_s), + ) + + +def as_sent(x: torch.Tensor, payload_format: str) -> torch.Tensor: + """What the FFN side receives for ``x`` (dequantized when FP8), in fp32.""" + from vllm.model_executor.layers.quantization.utils.fp8_utils import ( + per_token_group_quant_fp8, + ) + + from afd_plugin.connectors.gpu.deepep_m2n.fp8_quant import dequantize_fp8_rows + + if payload_format == "bf16" or x.shape[0] == 0: + return x.float() + x_q, x_s = per_token_group_quant_fp8(x, 128, use_ue8m0=False) + return dequantize_fp8_rows(x_q, x_s) + + +def stage_seed(step: int, stage: int) -> int: + return step * MAX_STAGES + stage + + +class Checker: + """Per-rank verification state: references and worst observed errors.""" + + def __init__(self, connector: Any) -> None: + from afd_plugin.connectors.gpu.deepep_m2n.experts import DeepEPM2NExperts + + self.connector = connector + self.payload_format = connector.config.payload_format + self.tolerances = ( + BF16_TOLERANCES if self.payload_format == "bf16" else FP8_TOLERANCES + ) + self.worst: dict[str, float] = {} + if connector.is_ffn: + self.experts = connector.layout.expert_range(connector.role_rank) + self.weights, self.w13_ref, self.w2_ref = local_weights( + self.experts, self.payload_format + ) + self.reference_executor = DeepEPM2NExperts( + "fused_experts", connector.expert_alignment + ) + else: + # End to end, the Attention side compares against the BF16 model. + self.w13_ref, self.w2_ref = make_expert_weights(range(M2N_NUM_EXPERTS)) + + def record(self, key: str, value: float) -> None: + self.worst[key] = max(self.worst.get(key, 0.0), value) + + def check_ffn_receive( + self, + payload: Any, + tokens_per_attn_rank: tuple[int, ...], + seed: int, + label: str, + limit: int = M2N_NUM_EXPERTS, + ) -> None: + """Exact per-expert counts, then every row and weight matched exactly. + + Rows are matched as a multiset per expert (bit-exact fingerprints and + their top-k weights), so dropped, duplicated, swapped or wrongly weighted + rows fail even when per-expert sums agree. The sums are still recorded + as error metrics. ``limit`` is the routing range used for the inputs. + ``tokens_per_attn_rank`` holds one count per Attention DP rank; with + Attention TP every TP peer's slice is expected once. + """ + from afd_plugin.connectors.gpu.deepep_m2n.experts import ( + align_up, + expert_row_ranges, + ) + from afd_plugin.connectors.gpu.deepep_m2n.fp8_quant import ( + dequantize_fp8_rows, + ) + + connector, experts = self.connector, self.experts + state = payload.context.states + ranges = expert_row_ranges( + state.recv_count_psum.to(torch.int64), connector.expert_alignment + ) + received = ( + payload.hidden_states.float() + if state.recv_scale is None + else dequantize_fp8_rows(payload.hidden_states, state.recv_scale) + ) + expected_counts = torch.zeros(len(experts), dtype=torch.int64, device="cuda") + expected_xsum = torch.zeros(len(experts), M2N_HIDDEN_SIZE, device="cuda") + expected_wsum = torch.zeros(len(experts), device="cuda") + expected_rows: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = [] + attn_tp_size = connector.attn_tp_size + for attn_rank in range(len(tokens_per_attn_rank) * attn_tp_size): + x, ids, w, _ = attention_rank_inputs( + seed, attn_rank, attn_tp_size, tokens_per_attn_rank, limit + ) + sent = as_sent(x, self.payload_format) + mask = (ids >= experts.start) & (ids < experts.stop) + rows, slots = mask.nonzero(as_tuple=True) + local = ids[rows, slots] - experts.start + expected_counts.index_add_(0, local, torch.ones_like(local)) + expected_xsum.index_add_(0, local, sent[rows]) + expected_wsum.index_add_(0, local, w[rows, slots]) + expected_rows.append((local, row_keys(sent[rows]), w[rows, slots])) + got_counts = [end - start for start, end in ranges] + if got_counts != expected_counts.tolist(): + raise AssertionError( + f"{label}: counts {got_counts} != {expected_counts.tolist()}" + ) + valid = [ + torch.arange(start, end, device="cuda") + for start, end in ranges + if end > start + ] + if valid: + index = torch.cat(valid) + got_local = torch.repeat_interleave( + torch.arange(len(experts), device="cuda"), + torch.tensor(got_counts, device="cuda"), + ) + got = sort_rows( + got_local, + row_keys(received[index]), + state.recv_topk_weights[index].float(), + ) + local, keys, weights = ( + torch.cat(parts) for parts in zip(*expected_rows, strict=True) + ) + want = sort_rows(local, keys, weights) + for name, got_part, want_part in zip( + ("expert ids", "rows", "top-k weights"), got, want, strict=True + ): + if not torch.equal(got_part, want_part): + raise AssertionError(f"{label}: received {name} differ from sent") + plan = connector.stage_plans[state.lane] + padded = [align_up(count, connector.expert_alignment) for count in got_counts] + if plan.do_cpu_sync and state.recv_counts_host != padded: + raise AssertionError( + f"{label}: host counts {state.recv_counts_host} != padded {padded}" + ) + got_xsum = torch.stack([received[s:e].sum(0) for s, e in ranges]) + got_wsum = torch.stack([state.recv_topk_weights[s:e].sum() for s, e in ranges]) + self.record("dispatch_rows_rel_err", rel_err(got_xsum, expected_xsum)) + self.record("dispatch_weights_rel_err", rel_err(got_wsum, expected_wsum)) + + def check_ffn_experts(self, payload: Any, out: torch.Tensor) -> None: + """Experts output vs fp32 math and vLLM ``fused_experts`` on the same rows.""" + from afd_plugin.connectors.gpu.deepep_m2n.experts import expert_row_ranges + from afd_plugin.connectors.gpu.deepep_m2n.fp8_quant import ( + dequantize_fp8_rows, + ) + + state = payload.context.states + ranges = expert_row_ranges( + state.recv_count_psum.to(torch.int64), self.connector.expert_alignment + ) + received = ( + payload.hidden_states.float() + if state.recv_scale is None + else dequantize_fp8_rows(payload.hidden_states, state.recv_scale) + ) + reference = self.reference_executor( + payload.hidden_states, + state.recv_scale, + state.recv_topk_weights, + state.recv_count_psum, + self.weights, + ) + for expert, (start, end) in enumerate(ranges): + if end == start: + continue + exact = ( + expert_mlp_fp32( + received[start:end], self.w13_ref[expert], self.w2_ref[expert] + ) + * state.recv_topk_weights[start:end, None] + ) + self.record("experts_rel_err", rel_err(out[start:end], exact)) + self.record( + "experts_vs_fused_rel_err", + rel_err(out[start:end], reference[start:end]), + ) + + def check_combined( + self, + combined: torch.Tensor, + x: torch.Tensor, + ids: torch.Tensor, + w: torch.Tensor, + ) -> None: + if x.shape[0]: + expected = dense_moe_reference(x, ids, w, self.w13_ref, self.w2_ref) + self.record("combine_rel_err", rel_err(combined, expected)) + + def finish(self) -> dict[str, Any]: + torch.cuda.synchronize() + for key, limit in self.tolerances.items(): + if self.worst.get(key, 0.0) > limit: + raise AssertionError(f"{key}={self.worst[key]:.3e} exceeds {limit:.0e}") + return { + **self.worst, + "expert_alignment": self.connector.expert_alignment, + "domain": list(self.connector.domain), + "network_plan": self.network_plan_summary(), + } + + def network_plan_summary(self) -> str: + plan = self.connector.network_plan + if plan is None: + return "" + nic = plan.nic_by_rank[self.connector.world_rank] + return f"{plan.summary()} this_rank_nic={nic or 'nccl'}" + + +def attention_context(layer: int, stage: int, num_tokens: int) -> Any: + from afd_plugin.connectors import AFDTransferContext, AFDTransferMetadata + + return AFDTransferContext( + metadata=AFDTransferMetadata.create_attention_metadata( + layer_idx=layer, stage_idx=stage, seq_len=num_tokens + ) + ) + + +def run_attention_step( + connector: Any, + inputs: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor]], + outs: list[torch.Tensor], + layer: int = 0, +) -> None: + """DBO order: dispatch every stage, then combine every stage.""" + for stage, (x, ids, w) in enumerate(inputs): + connector.send_attn_output( + x, attention_context(layer, stage, x.shape[0]), topk_ids=ids, topk_weights=w + ) + for stage, out in enumerate(outs): + connector.recv_ffn_output(out, stage) + + +def run_ffn_step( + connector: Any, weights: Any, num_stages: int +) -> tuple[list[Any], list[torch.Tensor]]: + """Receive every stage, then experts + combine per stage, then join lanes.""" + payloads = [connector.recv_attn_output(stage) for stage in range(num_stages)] + outs = [] + for payload in payloads: + out = connector.run_experts(payload, weights) + connector.send_ffn_output(out, payload.context) + outs.append(out) + connector.synchronize_lanes() + return payloads, outs + + +def correctness_scenario( + connector: Any, topology: M2NTopology, args: dict[str, Any] +) -> dict[str, Any]: + """Eager steps over every case; each case lists per-stage token counts. + + Like the runtime, one control payload covers a whole step: every layer of + the case reuses the plan it set up. + """ + checker = Checker(connector) + step = 0 + for case in [Case(**case) for case in args["cases"]]: + stages = case.stage_tokens() + limit = routing_limit(case, connector.layout) + control_step( + connector, + stages, + graph_mode=case.graph_mode, + ffn_graph_mode=case.ffn_graph_mode, + ) + for layer in range(case.num_layers): + step += 1 + if connector.is_ffn: + payloads, outs = run_ffn_step(connector, checker.weights, len(stages)) + for stage, (payload, out) in enumerate( + zip(payloads, outs, strict=True) + ): + label = f"{case.name} L{layer} S{stage}" + checker.check_ffn_receive( + payload, stages[stage], stage_seed(step, stage), label, limit + ) + checker.check_ffn_experts(payload, out) + else: + sliced = [ + attention_rank_inputs( + stage_seed(step, stage), + connector.role_rank, + connector.attn_tp_size, + tokens, + limit, + ) + for stage, tokens in enumerate(stages) + ] + inputs = [(x, ids, w) for x, ids, w, _ in sliced] + outs = [torch.empty_like(x) for x, _, _ in inputs] + run_attention_step(connector, inputs, outs, layer) + # Padding rows (last) route nowhere; check the real rows. + for (x, ids, w, real), out in zip(sliced, outs, strict=True): + checker.check_combined(out[:real], x[:real], ids[:real], w[:real]) + return checker.finish() + + +def compile_scenario( + connector: Any, topology: M2NTopology, args: dict[str, Any] +) -> dict[str, Any]: + """Attention side runs send + recv under ``torch.compile(fullgraph=True)``.""" + + def attention_layer( + x: torch.Tensor, ids: torch.Tensor, w: torch.Tensor, out: torch.Tensor + ) -> torch.Tensor: + connector.send_attn_output( + x, attention_context(0, 0, x.shape[0]), topk_ids=ids, topk_weights=w + ) + return connector.recv_ffn_output(out, 0) * 1.0 + + compiled = torch.compile(attention_layer, fullgraph=True) + checker = Checker(connector) + step = 0 + for case in [Case(**case) for case in args["cases"]]: + stages = case.stage_tokens() + control_step(connector, stages, graph_mode=case.graph_mode) + for _ in range(case.num_layers): + step += 1 + if connector.is_ffn: + run_ffn_step(connector, checker.weights, 1) + else: + rank = connector.role_rank + x, ids, w = make_inputs(stage_seed(step, 0), rank, stages[0][rank]) + combined = compiled(x, ids, w, torch.empty_like(x)) + checker.check_combined(combined, x, ids, w) + return checker.finish() + + +Inputs = tuple[torch.Tensor, torch.Tensor, torch.Tensor] + + +@dataclass +class CapturedStep: + """One CUDA graph for one bucket: ``num_layers`` MoE layers of one step. + + ``inputs``/``outs``/``payloads`` are indexed ``[layer][stage]``; inputs and + outs are filled on Attention ranks, payloads and outs on FFN ranks. + """ + + graph: torch.cuda.CUDAGraph + inputs: list[list[Inputs]] + outs: list[list[torch.Tensor]] + payloads: list[list[Any]] + + +def capture_step( + connector: Any, + topology: M2NTopology, + bucket: int, + weights: Any, + num_layers: int = 1, + limit: int = M2N_NUM_EXPERTS, +) -> CapturedStep: + """Warm up one eager graph-mode step, then capture the same step. + + Every Attention rank runs exactly ``bucket`` tokens per stage, so all + shapes are fixed. Graph-mode payloads skip the host sync and give + worst-case receive shapes on every rank. The graph holds ``num_layers`` + layers under one control payload, as a decode graph holds a whole step. + """ + num_stages = connector.num_lanes + stages = [(bucket,) * topology.attn_size] * num_stages + inputs: list[list[Inputs]] = [] + outs: list[list[torch.Tensor]] = [] + if not connector.is_ffn: + inputs = [ + [ + make_inputs(stage_seed(0, stage), connector.role_rank, bucket, limit) + for stage in range(num_stages) + ] + for _ in range(num_layers) + ] + outs = [[torch.empty_like(x) for x, _, _ in layer] for layer in inputs] + payloads: list[list[Any]] = [] + for capture in (False, True): + control_step(connector, stages, graph_mode=True) + graph = torch.cuda.CUDAGraph() + torch.cuda.synchronize() + with torch.cuda.graph(graph) if capture else nullcontext(): + if connector.is_ffn: + payloads, outs = [], [] + for _ in range(num_layers): + layer_payloads, layer_outs = run_ffn_step( + connector, weights, num_stages + ) + payloads.append(layer_payloads) + outs.append(layer_outs) + else: + for layer in range(num_layers): + run_attention_step(connector, inputs[layer], outs[layer], layer) + torch.cuda.synchronize() + return CapturedStep(graph=graph, inputs=inputs, outs=outs, payloads=payloads) + + +def replay_step( + connector: Any, topology: M2NTopology, captured: CapturedStep, bucket: int +) -> None: + stages = [(bucket,) * topology.attn_size] * connector.num_lanes + control_step(connector, stages, graph_mode=True) + captured.graph.replay() + + +def graph_scenario( + connector: Any, topology: M2NTopology, args: dict[str, Any] +) -> dict[str, Any]: + """Capture one graph per bucket; replay them interleaved with new inputs. + + Args: ``buckets``, ``replays``, optional ``num_layers`` (layers per graph, + default 1) and ``hot_ffn_rank`` (route everything to FFN rank 0). + """ + checker = Checker(connector) + weights = checker.weights if connector.is_ffn else None + buckets = args["buckets"] + num_layers = args.get("num_layers", 1) + case = Case("graph", (), hot_ffn_rank=args.get("hot_ffn_rank", False)) + limit = routing_limit(case, connector.layout) + captured = { + bucket: capture_step(connector, topology, bucket, weights, num_layers, limit) + for bucket in buckets + } + step = 0 + for _ in range(args["replays"]): + for bucket in buckets: + graph_step = captured[bucket] + first_step = step + 1 + step += num_layers + if not connector.is_ffn: + rank = connector.role_rank + for layer, layer_inputs in enumerate(graph_step.inputs): + for stage, static in enumerate(layer_inputs): + fresh = make_inputs( + stage_seed(first_step + layer, stage), rank, bucket, limit + ) + for static_tensor, new_tensor in zip( + static, fresh, strict=True + ): + static_tensor.copy_(new_tensor) + replay_step(connector, topology, graph_step, bucket) + torch.cuda.synchronize() + if connector.is_ffn: + tokens = (bucket,) * topology.attn_size + for layer in range(num_layers): + for stage, (payload, out) in enumerate( + zip( + graph_step.payloads[layer], + graph_step.outs[layer], + strict=True, + ) + ): + label = f"graph bucket={bucket} step={step} L{layer} S{stage}" + checker.check_ffn_receive( + payload, + tokens, + stage_seed(first_step + layer, stage), + label, + limit, + ) + checker.check_ffn_experts(payload, out) + else: + for layer_inputs, layer_outs in zip( + graph_step.inputs, graph_step.outs, strict=True + ): + for (x, ids, w), out in zip(layer_inputs, layer_outs, strict=True): + checker.check_combined(out, x, ids, w) + return checker.finish() + + +def run_attention_dbo_layers( + connector: Any, + inputs: list[list[Inputs]], + outs: list[list[torch.Tensor]], +) -> None: + """Attention side of vLLM's DBO order over two micro-batches: + D0(0) D1(0) C0(0) D0(1) C1(0) D1(1) C0(1) ... C1(last).""" + num_layers = len(inputs) + for layer in range(num_layers): + for stage in (0, 1): + if stage == 1 and layer > 0: + connector.recv_ffn_output(outs[layer - 1][1], 1) + x, ids, w = inputs[layer][stage] + connector.send_attn_output( + x, + attention_context(layer, stage, x.shape[0]), + topk_ids=ids, + topk_weights=w, + ) + connector.recv_ffn_output(outs[layer][0], 0) + connector.recv_ffn_output(outs[-1][1], 1) + + +@dataclass +class FFNPipelineStep: + payloads: list[list[Any]] + outs: list[list[torch.Tensor]] + # (layer, C1(layer-1) end, E0(layer) end) timing events when stretched. + timings: list[tuple[int, Any, Any]] + start: Any + + +def run_ffn_dbo_layers( + connector: Any, + weights: Any, + num_layers: int, + sleep_cycles: int = 0, + *, + retain: bool = True, + stretch_cycles: int = 0, +) -> FFNPipelineStep: + """FFN side in the routed runner's order: each stage's send waits until + just before that stage's next receive, then the lanes join. + + ``sleep_cycles`` stretches every E0 with a GPU sleep queued right after + it on the caller stream, and records when C1 of the previous layer and + the stretched E0 end. ``stretch_cycles`` queues the same sleep after + every experts launch without timing events (capturable). ``retain=False`` + drops each layer's rows and output once sent, as the routed runner does, + so the allocator may reuse them (nothing is returned to check). + """ + # Timing events only in the stretched eager variant (not while capturing). + start = None + if sleep_cycles: + start = torch.cuda.Event(enable_timing=True) + start.record() + payloads: list[list[Any]] = [[None, None] for _ in range(num_layers)] + outs: list[list[Any]] = [[None, None] for _ in range(num_layers)] + timings: list[tuple[int, Any, Any]] = [] + pending: dict[int, tuple[torch.Tensor, Any]] = {} + c1_end = None + for layer in range(num_layers): + for stage in (0, 1): + if stage in pending: + connector.send_ffn_output(*pending.pop(stage)) + if stage == 1 and sleep_cycles: + c1_end = torch.cuda.Event(enable_timing=True) + c1_end.record(connector.lane_streams[1]) + payload = connector.recv_attn_output(stage) + out = connector.run_experts(payload, weights) + if stretch_cycles: + torch.cuda._sleep(stretch_cycles) + if stage == 0 and sleep_cycles: + torch.cuda._sleep(sleep_cycles) + e0_end = torch.cuda.Event(enable_timing=True) + e0_end.record() + if layer > 0: + timings.append((layer, None, e0_end)) + if stage == 1 and sleep_cycles and layer > 0: + prev_layer, _, prev_e0 = timings[-1] + timings[-1] = (prev_layer, c1_end, prev_e0) + pending[stage] = (out, payload.context) + if retain: + payloads[layer][stage] = payload + outs[layer][stage] = out + del payload, out + for stage in (0, 1): + connector.send_ffn_output(*pending.pop(stage)) + connector.synchronize_lanes() + return FFNPipelineStep(payloads=payloads, outs=outs, timings=timings, start=start) + + +def dbo_pipeline_scenario( + connector: Any, topology: M2NTopology, args: dict[str, Any] +) -> dict[str, Any]: + """Several MoE layers in vLLM's two-micro-batch DBO order, with the FFN's + deferred sends, checked layer by layer; eager or captured and replayed. + + Args: ``num_layers``, ``tokens`` (per Attention rank and stage), + ``graph`` (capture once, replay ``replays`` times with fresh inputs) and, + eager only, ``sleep_cycles``: stretch every E0 on the FFN and require + C1(L-1) to finish before the stretched E0(L) does, i.e. one lane's + combine must not wait for the other lane's experts. + + Graph only, ``growth_layers=(short, long)``: also capture a short and a + long step whose FFN drops each layer's rows once sent (as the routed + runner does) and stretches every experts launch; replay both and check + Attention's outputs, and require the memory the FFN's capture reserves + not to grow with the layer count (``GRAPH_GROWTH_LIMIT``). + """ + if connector.num_lanes != 2: + raise RuntimeError("dbo_pipeline_scenario needs two lanes") + checker = Checker(connector) + weights = checker.weights if connector.is_ffn else None + num_layers = args["num_layers"] + tokens = (args["tokens"],) * topology.attn_size + stages = [tokens, tokens] + graph = args.get("graph", False) + sleep_cycles = args.get("sleep_cycles", 0) + limit = M2N_NUM_EXPERTS + overlap: list[float] = [] + + def fresh_inputs(first_step: int, layers: int = num_layers) -> list[list[Inputs]]: + return [ + [ + make_inputs( + stage_seed(first_step + layer, stage), + connector.role_rank, + tokens[connector.role_rank], + limit, + ) + for stage in (0, 1) + ] + for layer in range(layers) + ] + + def check(first_step: int, ffn_step: Any, inputs: Any, outs: Any, label: str): + if connector.is_ffn: + for layer in range(num_layers): + for stage in (0, 1): + payload = ffn_step.payloads[layer][stage] + checker.check_ffn_receive( + payload, + tokens, + stage_seed(first_step + layer, stage), + f"{label} L{layer} S{stage}", + limit, + ) + checker.check_ffn_experts(payload, ffn_step.outs[layer][stage]) + else: + for layer_inputs, layer_outs in zip(inputs, outs, strict=True): + for (x, ids, w), out in zip(layer_inputs, layer_outs, strict=True): + checker.check_combined(out, x, ids, w) + + step = 0 + if not graph: + for _ in range(args.get("steps", 2)): + control_step(connector, stages) + first_step = step + 1 + step += num_layers + inputs = outs = None + ffn_step = None + if connector.is_ffn: + ffn_step = run_ffn_dbo_layers( + connector, weights, num_layers, sleep_cycles + ) + else: + inputs = fresh_inputs(first_step) + outs = [[torch.empty_like(x) for x, _, _ in layer] for layer in inputs] + run_attention_dbo_layers(connector, inputs, outs) + torch.cuda.synchronize() + check(first_step, ffn_step, inputs, outs, "dbo eager") + if ffn_step is not None and sleep_cycles: + for layer, c1_end, e0_end in ffn_step.timings: + c1_ms = ffn_step.start.elapsed_time(c1_end) + e0_ms = ffn_step.start.elapsed_time(e0_end) + overlap.append(e0_ms - c1_ms) + if c1_ms >= e0_ms: + raise RuntimeError( + f"C1(L{layer - 1}) ended at {c1_ms:.3f} ms, after " + f"the stretched E0(L{layer}) at {e0_ms:.3f} ms: the " + "lane-1 combine waited for lane 0's experts" + ) + result = checker.finish() + if overlap: + result["min_c1_lead_ms"] = min(overlap) + return result + + # Graph: warm up uncaptured, capture once, then replay with fresh inputs. + inputs = outs = None + if not connector.is_ffn: + inputs = fresh_inputs(0) + outs = [[torch.empty_like(x) for x, _, _ in layer] for layer in inputs] + cuda_graph = torch.cuda.CUDAGraph() + ffn_step = None + for capture in (False, True): + control_step(connector, stages, graph_mode=True) + torch.cuda.synchronize() + with torch.cuda.graph(cuda_graph) if capture else nullcontext(): + if connector.is_ffn: + ffn_step = run_ffn_dbo_layers(connector, weights, num_layers) + else: + assert inputs is not None and outs is not None + run_attention_dbo_layers(connector, inputs, outs) + torch.cuda.synchronize() + for _ in range(args.get("replays", 3)): + first_step = step + 1 + step += num_layers + if not connector.is_ffn: + assert inputs is not None + for static_layer, new_layer in zip( + inputs, fresh_inputs(first_step), strict=True + ): + for static, new in zip(static_layer, new_layer, strict=True): + for static_tensor, new_tensor in zip(static, new, strict=True): + static_tensor.copy_(new_tensor) + control_step(connector, stages, graph_mode=True) + cuda_graph.replay() + torch.cuda.synchronize() + check(first_step, ffn_step, inputs, outs, "dbo graph") + result = checker.finish() + if args.get("growth_layers"): + result["capture_pool_mib"] = _dbo_capture_growth( + connector, stages, weights, fresh_inputs, args + ) + return result + + +# The memory a capture whose FFN drops sent rows reserves (its graph pool) may +# grow with the layer count by at most this factor (a capture that keeps every +# layer's rows grows linearly). +GRAPH_GROWTH_LIMIT = 1.5 + + +def _dbo_capture_growth( + connector: Any, + stages: list[tuple[int, ...]], + weights: Any, + fresh_inputs: Any, + args: dict[str, Any], +) -> dict[int, float]: + """See ``dbo_pipeline_scenario`` (``growth_layers``).""" + pools: dict[int, float] = {} + checker = None if connector.is_ffn else Checker(connector) + for layers in args["growth_layers"]: + inputs = outs = None + if not connector.is_ffn: + inputs = fresh_inputs(0, layers) + outs = [[torch.empty_like(x) for x, _, _ in layer] for layer in inputs] + control_step(connector, stages, graph_mode=True) + torch.cuda.synchronize() + # Reserved, not allocated: a block freed during capture but kept out + # of reuse leaves the allocated count at once and stays reserved. + # torch.cuda.graph empties the cache on entry; do it before the base. + gc.collect() + torch.cuda.empty_cache() + base = torch.cuda.memory_reserved() + cuda_graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(cuda_graph): + if connector.is_ffn: + run_ffn_dbo_layers( + connector, + weights, + layers, + retain=False, + stretch_cycles=args.get("stretch_cycles", 0), + ) + else: + assert inputs is not None and outs is not None + run_attention_dbo_layers(connector, inputs, outs) + torch.cuda.synchronize() + pools[layers] = (torch.cuda.memory_reserved() - base) / 2**20 + for replay in range(args.get("replays", 3)): + if not connector.is_ffn: + assert inputs is not None + for static_layer, new_layer in zip( + inputs, fresh_inputs(1000 * (replay + 1), layers), strict=True + ): + for static, new in zip(static_layer, new_layer, strict=True): + for static_tensor, new_tensor in zip(static, new, strict=True): + static_tensor.copy_(new_tensor) + control_step(connector, stages, graph_mode=True) + cuda_graph.replay() + torch.cuda.synchronize() + if checker is not None: + assert inputs is not None and outs is not None + for layer_inputs, layer_outs in zip(inputs, outs, strict=True): + for (x, ids, w), out in zip(layer_inputs, layer_outs, strict=True): + checker.check_combined(out, x, ids, w) + del cuda_graph + torch.cuda.empty_cache() + short, long = args["growth_layers"] + if connector.is_ffn and pools[long] > GRAPH_GROWTH_LIMIT * pools[short]: + raise RuntimeError( + f"FFN capture pool grew from {pools[short]:.1f} MiB ({short} layers) " + f"to {pools[long]:.1f} MiB ({long} layers): sent rows are not reused" + ) + return pools + + +def peer_exit_scenario( + connector: Any, topology: M2NTopology, args: dict[str, Any] +) -> dict[str, Any]: + """One rank dies (``os._exit``, no cleanup) after the step's control payload. + + The others then run the layer and are expected to fail through DeepEP's + timeouts rather than wait forever. Returns normally only if the layer + completed despite the missing peer, which the test treats as a failure. + """ + tokens = (args["tokens"],) * topology.attn_size + control_step(connector, [tokens]) + if connector.world_rank == args["exit_world_rank"]: + os._exit(PEER_EXIT_CODE) + start = time.monotonic() + if connector.is_ffn: + weights, _, _ = local_weights( + connector.layout.expert_range(connector.role_rank), + connector.config.payload_format, + ) + run_ffn_step(connector, weights, 1) + else: + x, ids, w = make_inputs(0, connector.role_rank, args["tokens"]) + run_attention_step(connector, [(x, ids, w)], [torch.empty_like(x)]) + torch.cuda.synchronize() + return {"completed_without_peer_s": time.monotonic() - start} + + +def _timed(fn: Any, iters: int) -> float: + """Mean wall time of ``fn`` in µs; ranks start each iteration together.""" + total = 0.0 + for iteration in range(BENCH_WARMUP_ITERS + iters): + torch.cuda.synchronize() + dist.barrier() + start = time.perf_counter() + fn() + torch.cuda.synchronize() + if iteration >= BENCH_WARMUP_ITERS: + total += time.perf_counter() - start + return total / iters * 1e6 + + +def eager_step( + connector: Any, + stages: list[tuple[int, ...]], + weights: Any, + inputs: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor]], + outs: list[torch.Tensor], +) -> None: + """One eager step over every stage (control payload included).""" + control_step(connector, stages) + if connector.is_ffn: + run_ffn_step(connector, weights, len(stages)) + else: + run_attention_step(connector, inputs, outs) + + +def benchmark_scenario( + connector: Any, topology: M2NTopology, args: dict[str, Any] +) -> dict[str, Any]: + """Eager per-layer time. One lane: dispatch/experts/combine breakdown.""" + num_stages = connector.num_lanes + if connector.is_ffn: + experts = connector.layout.expert_range(connector.role_rank) + weights, _, _ = local_weights(experts, connector.config.payload_format) + timings: dict[str, Any] = {} + for tokens in args["tokens_per_attn_rank"]: + stages = [(tokens,) * topology.attn_size] * num_stages + if not connector.is_ffn: + inputs = [ + make_inputs(stage, connector.role_rank, tokens) + for stage in range(num_stages) + ] + outs = [torch.empty_like(x) for x, _, _ in inputs] + if num_stages > 1: + step = functools.partial( + eager_step, + connector, + stages, + weights if connector.is_ffn else None, + [] if connector.is_ffn else inputs, + [] if connector.is_ffn else outs, + ) + timings[str(tokens)] = {"total_us": _timed(step, args["iters"])} + continue + per_iter: list[tuple[float, float, float]] = [] + for iteration in range(BENCH_WARMUP_ITERS + args["iters"]): + control_step(connector, stages) + torch.cuda.synchronize() + dist.barrier() + t0 = time.perf_counter() + if connector.is_ffn: + payload = connector.recv_attn_output(0) + torch.cuda.synchronize() + t1 = time.perf_counter() + out = connector.run_experts(payload, weights) + torch.cuda.synchronize() + t2 = time.perf_counter() + connector.send_ffn_output(out, payload.context) + connector.synchronize_lanes() + else: + x, ids, w = inputs[0] + if args.get("bypass_custom_ops", False): + # Measures the custom-op wrappers' own cost. + connector.dispatch_attention(x, ids, w, 0) + else: + connector.send_attn_output( + x, attention_context(0, 0, tokens), topk_ids=ids, topk_weights=w + ) + torch.cuda.synchronize() + t1 = t2 = time.perf_counter() + if args.get("bypass_custom_ops", False): + connector.combine_attention(outs[0], 0) + else: + connector.recv_ffn_output(outs[0], 0) + torch.cuda.synchronize() + t3 = time.perf_counter() + if iteration >= BENCH_WARMUP_ITERS: + per_iter.append((t1 - t0, t2 - t1, t3 - t2)) + count = len(per_iter) + timings[str(tokens)] = { + "dispatch_us": sum(item[0] for item in per_iter) / count * 1e6, + "experts_us": sum(item[1] for item in per_iter) / count * 1e6, + "combine_us": sum(item[2] for item in per_iter) / count * 1e6, + "total_us": sum(sum(item) for item in per_iter) / count * 1e6, + } + return timings + + +def graph_benchmark_scenario( + connector: Any, topology: M2NTopology, args: dict[str, Any] +) -> dict[str, Any]: + """Per-layer time of a replayed CUDA graph (all lanes) per token count. + + The control payload is delivered outside the timed region, as a real + runtime sends it once per step rather than once per layer. + """ + weights = None + if connector.is_ffn: + experts = connector.layout.expert_range(connector.role_rank) + weights, _, _ = local_weights(experts, connector.config.payload_format) + timings: dict[str, Any] = {} + for tokens in args["tokens_per_attn_rank"]: + captured = capture_step(connector, topology, tokens, weights) + stages = [(tokens,) * topology.attn_size] * connector.num_lanes + total = 0.0 + for iteration in range(BENCH_WARMUP_ITERS + args["iters"]): + control_step(connector, stages, graph_mode=True) + torch.cuda.synchronize() + dist.barrier() + start = time.perf_counter() + captured.graph.replay() + torch.cuda.synchronize() + if iteration >= BENCH_WARMUP_ITERS: + total += time.perf_counter() - start + timings[str(tokens)] = {"total_us": total / args["iters"] * 1e6} + del captured + return timings + + +def local_moe_baseline( + tokens_list: list[int], iters: int, payload_format: str +) -> dict[str, Any]: + """Single-GPU MoE: vLLM ``fused_experts``, all experts local, no communication. + + FP8 uses vLLM's block-FP8 path (128x128 weight blocks, dynamic per-token + activation groups), the same recipe the connector's FP8 payload uses. + """ + from vllm.model_executor.layers.fused_moe.config import fp8_w8a8_moe_quant_config + from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts + + weights, _, _ = local_weights(range(M2N_NUM_EXPERTS), payload_format) + quant_config = ( + fp8_w8a8_moe_quant_config( + w1_scale=weights.w13_scale, + w2_scale=weights.w2_scale, + block_shape=[128, 128], + ) + if weights.is_fp8 + else None + ) + timings: dict[str, Any] = {} + for tokens in tokens_list: + x, ids, w = make_inputs(0, 0, tokens) + call = { + "hidden_states": x, + "w1": weights.w13, + "w2": weights.w2, + "topk_weights": w, + "topk_ids": ids.to(torch.int32), + "quant_config": quant_config, + } + for _ in range(BENCH_WARMUP_ITERS): + fused_experts(**call) + torch.cuda.synchronize() + start = time.perf_counter() + for _ in range(iters): + fused_experts(**call) + torch.cuda.synchronize() + timings[str(tokens)] = {"total_us": (time.perf_counter() - start) / iters * 1e6} + return timings + + +def available_gpus() -> int: + return torch.cuda.device_count() if torch.cuda.is_available() else 0 + + +def bench_output_path() -> Path: + return Path(os.environ.get("AFD_DEEPEP_M2N_BENCH_OUT", "deepep_m2n_bench.json")) diff --git a/tests/gpu/test_deepep_m2n_gpu.py b/tests/gpu/test_deepep_m2n_gpu.py new file mode 100644 index 000000000..94d7027cf --- /dev/null +++ b/tests/gpu/test_deepep_m2n_gpu.py @@ -0,0 +1,415 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""GPU tests for ``DeepEPM2NAFDConnector``: N->M transport and expert compute. + +Each test spawns one process per AFD rank on one node (FFN ranks first). FFN +ranks check the received rows exactly (per-expert counts, then every row and +top-k weight matched bit-exactly) and the MoE FFN output against an fp32 +reference and vLLM ``fused_experts``; Attention ranks check the combined +output against a dense MoE reference. The same checks run in eager mode, with +two lanes (DBO order), under ``torch.compile`` and after CUDA-graph replay. + +Each case is one step: one control payload, then several MoE layers under the +plan it sets, as the runtime does. +""" + +from __future__ import annotations + +import time +from collections.abc import Callable +from dataclasses import asdict +from typing import Any + +import pytest + +torch = pytest.importorskip("torch") +pytest.importorskip("vllm") + +from tests.gpu.deepep_m2n_harness import ( # noqa: E402 + CLOSE_FAILED, + NO_RESULT, + Case, + RankResult, + available_gpus, + compile_scenario, + correctness_scenario, + dbo_pipeline_scenario, + graph_scenario, + hot_cases, + lane_cases, + mixed_mode_cases, + peer_exit_scenario, + run_ranks, + standard_cases, +) +from tests.unit.fakes.vllm_config import M2NTopology # noqa: E402 + +pytestmark = pytest.mark.gpu + +CAPACITY = 256 +BF16 = "bf16" +FP8 = "fp8_block128_fp32" +PAYLOADS = [BF16, FP8] +GRAPH_BUCKETS = [1, 8, 64, CAPACITY] +GRAPH_REPLAYS = 3 +# Layers per captured graph: several dispatch/combine pairs per graph, as in +# a decode graph that holds a whole step. +GRAPH_LAYERS = 3 +# DeepEP timeouts for the peer-exit test, and the wall-clock budget for the +# whole run (spawn, init, the failing layer, teardown). +PEER_EXIT_TIMEOUT_S = 10 +PEER_EXIT_BUDGET_S = 300 + + +def _skip_without_gpus(topology: M2NTopology) -> None: + if available_gpus() < topology.world_size: + pytest.skip(f"needs {topology.world_size} GPUs") + if torch.cuda.get_device_capability()[0] < 9: + pytest.skip("DeepEP M2N needs SM90+") + + +def _run( + topology: M2NTopology, + extra_config: dict, + tmp_path, + scenario: Callable[..., dict[str, Any]] = correctness_scenario, + scenario_args: dict[str, Any] | None = None, +) -> None: + _skip_without_gpus(topology) + if scenario_args is None: + scenario_args = { + "cases": _cases(standard_cases(topology.attn_dp_size, CAPACITY)) + } + results = run_ranks( + topology, + scenario, + scenario_args, + {"capacity_tokens_per_rank": CAPACITY, **extra_config}, + tmp_path, + ) + failures = [ + f"rank {r.world_rank} ({r.role}):\n{r.error}" for r in results if not r.ok + ] + assert not failures, "\n".join(failures) + for result in results: + print(f"rank {result.world_rank} {result.role}: {result.metrics}") + + +def _cases(cases: list[Case]) -> list[dict[str, Any]]: + return [asdict(case) for case in cases] + + +@pytest.mark.parametrize("payload_format", PAYLOADS, ids=["bf16", "fp8"]) +@pytest.mark.parametrize( + ("ffn_size", "attn_size"), + [(1, 1), (2, 2), (1, 3)], + ids=["1F1A", "2F2A", "1F3A"], +) +def test_transport_and_deepgemm_experts(ffn_size, attn_size, payload_format, tmp_path): + """Standard cases (even, uneven, decode, empty, mixed; with and without + host sync) through dispatch, DeepGEMM experts and combine. + + Why: the core claim, for several rank ratios and both payload formats. + """ + _run(M2NTopology(ffn_size, attn_size), {"payload_format": payload_format}, tmp_path) + + +@pytest.mark.parametrize("payload_format", PAYLOADS, ids=["bf16", "fp8"]) +@pytest.mark.parametrize( + "topology", + [M2NTopology(2, 2, attn_tp_size=2), M2NTopology(2, 4, attn_tp_size=2)], + ids=["2F-1DPx2TP", "2F-2DPx2TP"], +) +def test_attention_tp_peers_dispatch_disjoint_slices( + topology, payload_format, tmp_path +): + """With Attention TP, each TP peer dispatches its padded slice of its DP + rank's tokens (padding rows with expert id -1): the FFN receives every + token's routes exactly once and each slice combines to the dense MoE + reference. Covers single-lane and two-lane steps, including decode + batches smaller than the TP size and empty DP ranks. + + Why: TP peers hold the same tokens; dispatching all of them from every + peer would run each token's experts TP times, and a lost or doubled slice + would silently corrupt the output. + """ + cases = standard_cases(topology.attn_dp_size, CAPACITY) + lane_cases( + topology.attn_dp_size, CAPACITY + ) + _run( + topology, + {"payload_format": payload_format, "num_lanes": 2}, + tmp_path, + scenario_args={"cases": _cases(cases)}, + ) + + +@pytest.mark.parametrize("payload_format", PAYLOADS, ids=["bf16", "fp8"]) +def test_transport_with_fused_experts_backend(payload_format, tmp_path): + """Standard cases on 2F2A with vLLM's Triton ``fused_experts`` backend. + + Why: the reference backend must keep working to cross-check DeepGEMM. + """ + _run( + M2NTopology(2, 2), + {"expert_backend": "fused_experts", "payload_format": payload_format}, + tmp_path, + ) + + +@pytest.mark.parametrize("payload_format", PAYLOADS, ids=["bf16", "fp8"]) +def test_transport_with_masked_protocol(payload_format, tmp_path): + """Standard cases on 2F2A with ``role_aware: false``: afdEP's Masked DeepEP + protocol instead of its AFD role path. + + Why: the masked protocol is the comparison baseline for the role path and + must stay correct on its own. + """ + _run( + M2NTopology(2, 2), + {"role_aware": False, "payload_format": payload_format}, + tmp_path, + ) + + +@pytest.mark.parametrize("payload_format", PAYLOADS, ids=["bf16", "fp8"]) +def test_hot_ffn_rank_routing(payload_format, tmp_path): + """Standard cases with every top-k choice on FFN rank 0's experts: rank 0 + receives every token ``top_k`` times, rank 1 receives nothing. + + Why: real routing is skewed; this is the worst-case receive size the + buffers and the no-sync path must handle. + """ + _run( + M2NTopology(2, 2), + {"payload_format": payload_format}, + tmp_path, + scenario_args={"cases": _cases(hot_cases(2, CAPACITY))}, + ) + + +def test_fp8_without_host_sync_for_small_buckets(tmp_path): + """FP8 standard cases with ``cpu_sync_min_tokens`` above capacity, so no + eager step reads receive counts on the host. + + Why: the no-sync path uses worst-case receive shapes and must stay correct. + """ + _run( + M2NTopology(2, 2), + {"payload_format": FP8, "cpu_sync_min_tokens": CAPACITY + 1}, + tmp_path, + ) + + +def test_graph_and_eager_plans_mixed_across_ranks(tmp_path): + """Attention plans a step as graph (no host sync) while FFN plans it eager + (host sync), and the reverse, for even, decode and mixed token counts. + + Why: an FFN runner without a graph for a step that Attention replays runs + it eagerly; the two sides then disagree on the host sync. + """ + _run( + M2NTopology(2, 2), + {"payload_format": FP8}, + tmp_path, + scenario_args={"cases": _cases(mixed_mode_cases(2))}, + ) + + +@pytest.mark.parametrize( + ("ffn_size", "attn_size", "payload_format"), + [(2, 2, BF16), (2, 2, FP8), (1, 3, FP8)], + ids=["2F2A-bf16", "2F2A-fp8", "1F3A-fp8"], +) +def test_two_lanes_dbo_order(ffn_size, attn_size, payload_format, tmp_path): + """Two lanes with different token counts per lane (including an empty + lane); Attention dispatches both lanes, then combines both. + + Why: DBO keeps two micro-batches in flight; they must not corrupt each other. + """ + topology = M2NTopology(ffn_size, attn_size) + _run( + topology, + {"payload_format": payload_format, "num_lanes": 2}, + tmp_path, + scenario_args={"cases": _cases(lane_cases(attn_size, CAPACITY))}, + ) + + +# GPU sleep long enough (~5 ms on H100) to dwarf a combine. +DBO_SLEEP_CYCLES = 10_000_000 + + +@pytest.mark.parametrize("payload_format", PAYLOADS, ids=["bf16", "fp8"]) +def test_dbo_pipeline_eager_lanes_do_not_wait_for_each_other(payload_format, tmp_path): + """Four MoE layers in vLLM's DBO order with the FFN's deferred sends; every + E0 on the FFN is stretched by a GPU sleep, and C1(L-1) must still finish + before the stretched E0(L). + + Why: a lane's combine that waited for the other lane's experts serialized + Attention and FFN compute, so DBO overlapped nothing. + """ + _run( + M2NTopology(2, 2), + {"payload_format": payload_format, "num_lanes": 2}, + tmp_path, + scenario=dbo_pipeline_scenario, + scenario_args={ + "num_layers": 4, + "tokens": 64, + "steps": 2, + "sleep_cycles": DBO_SLEEP_CYCLES, + }, + ) + + +@pytest.mark.parametrize("payload_format", PAYLOADS, ids=["bf16", "fp8"]) +def test_dbo_pipeline_cuda_graph_capture_and_replay(payload_format, tmp_path): + """The same DBO order captured in one graph per side and replayed with fresh + inputs; every layer checked after each replay. + + Why: the per-lane waits become graph edges; a lane that never joined the + capture or a missing edge would corrupt rows only in replays. + """ + _run( + M2NTopology(2, 2), + {"payload_format": payload_format, "num_lanes": 2}, + tmp_path, + scenario=dbo_pipeline_scenario, + scenario_args={"num_layers": 4, "tokens": 64, "graph": True, "replays": 3}, + ) + + +@pytest.mark.parametrize("payload_format", PAYLOADS, ids=["bf16", "fp8"]) +def test_dbo_graph_capture_reuses_sent_rows(payload_format, tmp_path): + """DBO steps of 2 and 8 MoE layers captured with an FFN that drops each + layer's rows once sent (as the routed runner does), every experts launch + stretched; Attention's outputs are checked after each replay, and the + memory the FFN's capture reserves must not grow with the layer count. + + Why: rows kept alive across streams during capture were never reused + within it, so every layer's rows stayed allocated and Qwen3-235B + graph+DBO (94 MoE layers) ran out of FFN memory at startup. Reusing them + early must still not let a lane overwrite rows the experts read. + """ + _run( + M2NTopology(2, 2), + {"payload_format": payload_format, "num_lanes": 2}, + tmp_path, + scenario=dbo_pipeline_scenario, + scenario_args={ + "num_layers": 2, + "tokens": 64, + "graph": True, + "replays": 2, + "growth_layers": (2, 8), + "stretch_cycles": DBO_SLEEP_CYCLES // 10, + }, + ) + + +@pytest.mark.parametrize( + ("num_lanes", "payload_format", "hot"), + [(1, BF16, False), (1, FP8, False), (2, FP8, False), (1, FP8, True)], + ids=["1lane-bf16", "1lane-fp8", "2lanes-fp8", "1lane-fp8-hot"], +) +def test_cuda_graph_capture_and_replay(num_lanes, payload_format, hot, tmp_path): + """Capture one graph per bucket (1, 8, 64, 256), each holding several MoE + layers, then replay them 3 times interleaved with fresh inputs; check all + layers after every replay. ``hot`` routes everything to FFN rank 0. + + Why: decode runs from replays of whole-step graphs; stale pointers, values + baked in at capture, or buffer reuse across layers would show up only here. + """ + _run( + M2NTopology(2, 2), + {"payload_format": payload_format, "num_lanes": num_lanes}, + tmp_path, + scenario=graph_scenario, + scenario_args={ + "buckets": GRAPH_BUCKETS, + "replays": GRAPH_REPLAYS, + "num_layers": GRAPH_LAYERS, + "hot_ffn_rank": hot, + }, + ) + + +def test_torch_compile_attention_side(tmp_path): + """Attention's send and receive inside ``torch.compile(fullgraph=True)`` + for 64, 1 and 17/22 tokens per rank. + + Why: vLLM compiles the Attention model; the custom ops must not break it. + """ + topology = M2NTopology(2, 2) + cases = [ + Case("even64", (64,) * topology.attn_size), + Case("decode1", (1,) * topology.attn_size), + Case("mixed", (17, 22)), + ] + _run( + topology, + {"payload_format": FP8}, + tmp_path, + scenario=compile_scenario, + scenario_args={"cases": _cases(cases)}, + ) + + +def test_failed_init_reports_and_closes(tmp_path): + """FFN rank 0 asks for FP8 while the others use BF16: every rank's init + must fail on the handshake, and ``close()`` must then succeed. + + Why: a misconfigured deployment must fail cleanly on all ranks, and + teardown after a partial init must not raise or hang. + """ + topology = M2NTopology(2, 2) + _skip_without_gpus(topology) + results = run_ranks( + topology, + correctness_scenario, + {"cases": []}, + {"capacity_tokens_per_rank": CAPACITY}, + tmp_path, + rank_extra_config={0: {"payload_format": FP8}}, + ) + + for result in results: + assert not result.ok, f"rank {result.world_rank} initialized" + assert "disagree on payload_format" in result.error, result.error + assert CLOSE_FAILED not in result.error, result.error + + +def test_peer_exit_fails_within_timeout(tmp_path): + """FFN rank 1 exits abruptly after the step's control payload; the other + ranks then run the layer with DeepEP timeouts of 10 s. Every survivor must + stop (error or abort) well within the budget instead of waiting forever. + + Why: a crashed peer must not hang the rest of the deployment. + """ + topology = M2NTopology(2, 2) + _skip_without_gpus(topology) + start = time.monotonic() + results = run_ranks( + topology, + peer_exit_scenario, + {"tokens": 64, "exit_world_rank": 1}, + { + "capacity_tokens_per_rank": CAPACITY, + "control_backend": "gloo", + "cpu_timeout_s": PEER_EXIT_TIMEOUT_S, + "gpu_timeout_s": PEER_EXIT_TIMEOUT_S, + }, + tmp_path, + timeout_s=PEER_EXIT_BUDGET_S, + ) + elapsed = time.monotonic() - start + + print(f"peer exit: {elapsed:.0f} s") + for result in results: + print(f"rank {result.world_rank}: ok={result.ok} {result.error[-300:]}") + assert results[1].error.startswith(NO_RESULT), "the exiting rank reported" + survivors: list[RankResult] = [r for r in results if r.world_rank != 1] + assert all(not r.ok for r in survivors), "a layer completed without its peer" + assert all("hung=True" not in r.error for r in survivors), "a survivor hung" + assert elapsed < PEER_EXIT_BUDGET_S diff --git a/tests/unit/config/test_validation.py b/tests/unit/config/test_validation.py index 99a08b367..c11c77208 100644 --- a/tests/unit/config/test_validation.py +++ b/tests/unit/config/test_validation.py @@ -165,6 +165,21 @@ def mrv2_config(monkeypatch): } +@pytest.mark.parametrize("role", ["attention", "ffn"]) +def test_routed_transport_preserves_runner_and_gate_policy(mrv2_config, role): + mrv2_config.additional_config["afd"].update( + role=role, connector="DeepEPM2NAFDConnector", compute_gate_on_attention=True + ) + validate_gpu_model_runner_v2_config( + mrv2_config, expected_role=role, device_type="cuda" + ) + mrv2_config.additional_config["afd"]["compute_gate_on_attention"] = False + with pytest.raises(RuntimeError, match="requires compute_gate_on_attention=true"): + validate_gpu_model_runner_v2_config( + mrv2_config, expected_role=role, device_type="cuda" + ) + + def _symmetric_ranks(dp_size, tp_size): return { "parallel_config.data_parallel_size": dp_size, diff --git a/tests/unit/connectors/test_base_factory.py b/tests/unit/connectors/test_base_factory.py index 5dbe37d46..fa932a3b7 100644 --- a/tests/unit/connectors/test_base_factory.py +++ b/tests/unit/connectors/test_base_factory.py @@ -3,6 +3,7 @@ from __future__ import annotations +import sys from types import SimpleNamespace import pytest @@ -15,6 +16,7 @@ AFDA2FTransferPayload, AFDConnectorBase, AFDConnectorFactory, + AFDRoutedExpertsConnector, AFDTransferContext, AFDTransferMetadata, ) @@ -61,10 +63,21 @@ def test_connector_metadata_validates_sequence_lengths(): AFDTransferMetadata( layer_idx=0, stage_idx=0, - seq_lens=[0], + seq_lens=[-1], ) +def test_connector_metadata_allows_empty_receive(): + metadata = AFDTransferMetadata.create_ffn_metadata( + layer_idx=0, + stage_idx=0, + seq_lens=[0], + ) + + assert metadata.total_tokens == 0 + assert metadata.validate_tensor_shape((0, 16)) + + def test_attn_output_carries_transfer_context(): metadata = AFDTransferMetadata.create_ffn_metadata( layer_idx=1, @@ -180,3 +193,139 @@ def test_factory_resolves_role_rank_before_connector_construction(monkeypatch): ) assert connector.role_rank == 1 + + +class _BootstrappingConnector(_MinimalConnector, AFDRoutedExpertsConnector): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.events: list[str] = [] + + def bootstrap(self): + self.events.append("bootstrap") + if self.vllm_config.fail_bootstrap: + raise RuntimeError("bootstrap failed") + + def init_afd_connector(self, *, moe_layer_indices=None, local_experts=None): + return None + + def run_experts(self, payload, weights): + return payload.hidden_states + + def synchronize_lanes(self): + return None + + def close(self): + self.events.append("close") + + +def _register(monkeypatch, name, connector_cls): + monkeypatch.setitem(AFDConnectorFactory._registry, name, lambda: connector_cls) + + +def _single_rank_config( + *, fail_bootstrap=False, data_parallel_size=1, data_parallel_rank=0, tp_size=1 +): + return SimpleNamespace( + additional_config={}, + fail_bootstrap=fail_bootstrap, + parallel_config=SimpleNamespace( + data_parallel_size=data_parallel_size, + data_parallel_rank=data_parallel_rank, + pipeline_parallel_size=1, + prefill_context_parallel_size=1, + tensor_parallel_size=tp_size, + ), + ) + + +def test_create_bootstrapped_connector_skips_other_connectors(monkeypatch): + """A connector that is not an ``AFDRoutedExpertsConnector`` is not + created before vLLM's distributed init; its runner creates it later. + + Why: P2P and the NPU connectors keep their existing bring-up order. + """ + _register(monkeypatch, "MinimalConnector", _MinimalConnector) + + connector = AFDConnectorFactory.create_bootstrapped_connector( + 0, + 0, + _single_rank_config(), + AFDConfig(connector="MinimalConnector", role="ffn"), + ) + + assert connector is None + + +def test_create_bootstrapped_connector_creates_and_bootstraps(monkeypatch): + """A routed-experts connector is created with the worker's ranks and + bootstrapped. + + Why: workers call this before vLLM's distributed init so the connector's + NCCL environment precedes vLLM's NCCL groups. + """ + _register(monkeypatch, "BootstrappingConnector", _BootstrappingConnector) + + connector = AFDConnectorFactory.create_bootstrapped_connector( + 3, + 1, + _single_rank_config(), + AFDConfig(connector="BootstrappingConnector", role="ffn"), + ) + + assert isinstance(connector, _BootstrappingConnector) + assert (connector.rank, connector.local_rank, connector.role_rank) == (3, 1, 0) + assert connector.events == ["bootstrap"] + + +def test_create_bootstrapped_connector_derives_tp_rank_from_worker_rank( + monkeypatch, +): + """With TP > 1 the role rank's TP part comes from the worker rank. + + Why: vLLM's TP group does not exist before its distributed init; reading + it there would fail. TP is vLLM's innermost rank dimension, so the worker + rank modulo the TP size is the TP rank. + """ + _register(monkeypatch, "BootstrappingConnector", _BootstrappingConnector) + monkeypatch.setitem(sys.modules, "vllm.distributed.parallel_state", None) + + connector = AFDConnectorFactory.create_bootstrapped_connector( + 1, + 1, + _single_rank_config(data_parallel_size=2, data_parallel_rank=1, tp_size=2), + AFDConfig( + connector="BootstrappingConnector", + role="attention", + num_attention_ranks=4, + ), + ) + + assert isinstance(connector, _BootstrappingConnector) + # DP rank 1, TP rank 1 of 2. + assert connector.role_rank == 3 + + +def test_create_bootstrapped_connector_closes_on_bootstrap_failure(monkeypatch): + """A failed bootstrap closes the connector before the error propagates. + + Why: the worker never receives the connector, so nobody else closes it. + """ + _register(monkeypatch, "BootstrappingConnector", _BootstrappingConnector) + created: list[_BootstrappingConnector] = [] + original_init = _BootstrappingConnector.__init__ + + def recording_init(self, *args, **kwargs): + original_init(self, *args, **kwargs) + created.append(self) + + monkeypatch.setattr(_BootstrappingConnector, "__init__", recording_init) + + with pytest.raises(RuntimeError, match="bootstrap failed"): + AFDConnectorFactory.create_bootstrapped_connector( + 0, + 0, + _single_rank_config(fail_bootstrap=True), + AFDConfig(connector="BootstrappingConnector", role="ffn"), + ) + + assert created[0].events == ["bootstrap", "close"] diff --git a/tests/unit/connectors/test_deepep_m2n_connector.py b/tests/unit/connectors/test_deepep_m2n_connector.py new file mode 100644 index 000000000..9abc365be --- /dev/null +++ b/tests/unit/connectors/test_deepep_m2n_connector.py @@ -0,0 +1,1599 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""CPU-only tests for ``DeepEPM2NAFDConnector``: config, layout, plans, handshake.""" + +from __future__ import annotations + +import contextlib +import itertools +from types import SimpleNamespace + +import pytest + +torch = pytest.importorskip("torch") +pytest.importorskip("vllm") + +from tests.unit.fakes.afdep import FakeElasticBuffer, install_fake_afdep # noqa: E402 + +# The connector module imports afdEP's transport when it loads. +install_fake_afdep() + +from afd_plugin.config import afd_config_from_mapping # noqa: E402 +from afd_plugin.connectors import ( # noqa: E402 + AFDConnectorFactory, + AFDControlPayload, + AFDDPMetadata, + AFDTransferContext, + AFDTransferMetadata, +) +from afd_plugin.connectors.gpu.deepep_m2n import ( # noqa: E402 + connector as connector_module, +) +from afd_plugin.connectors.gpu.deepep_m2n.connector import ( # noqa: E402 + DeepEPM2NConfig, + bucket_ladder, + expert_weight_format, + local_handshake_info, + local_model_info, + no_sync_bucket_limit, + plan_stage, + precompile_buckets, + resolve_capacity_tokens_per_rank, + resolve_num_lanes, + tokens_per_attention_rank, + validate_handshake, + validate_model_infos, + worst_case_receive_rows, +) +from afd_plugin.connectors.gpu.deepep_m2n.experts import ( # noqa: E402 + expert_row_ranges, +) +from afd_plugin.connectors.gpu.deepep_m2n.fp8_quant import ( # noqa: E402 + dequantize_fp8_block128, + quantize_fp8_block128, +) +from afd_plugin.connectors.gpu.deepep_m2n.layout import DeepEPM2NLayout # noqa: E402 +from afd_plugin.connectors.gpu.deepep_m2n.network import ( # noqa: E402 + NetworkConfig, + RankTopology, +) +from afd_plugin.connectors.routed_experts import ( # noqa: E402 + FFNLocalExperts, + LocalExpertWeights, +) +from tests.unit.fakes.vllm_config import ( # noqa: E402 + M2N_HIDDEN_SIZE, + M2NTopology, + m2n_vllm_config, +) + +CONNECTOR = "DeepEPM2NAFDConnector" +HOPPER = (9, 0) + + +def _dp(*tokens: int) -> AFDDPMetadata: + return AFDDPMetadata(num_tokens_across_dp_cpu=torch.tensor(tokens)) + + +def _connector( + role: str, role_rank: int, topology: M2NTopology, extra: dict | None = None +): + vllm_config = m2n_vllm_config(topology, role, role_rank, extra or {}, port=12345) + afd_config = afd_config_from_mapping(vllm_config.additional_config["afd"]) + return AFDConnectorFactory.create_connector(0, 0, vllm_config, afd_config) + + +def _ffn_deployment(layout: DeepEPM2NLayout, weight_format: str = "bf16") -> dict: + """A valid FFN deployment: linear EP over the FFN ranks, idle all2all.""" + return { + "expert_parallel": True, + "ep_size": layout.ffn_size, + "tensor_parallel_size": 1, + "expert_placement_strategy": "linear", + "eplb": False, + "all2all_backend": "allgather_reducescatter", + "weight_format": weight_format, + } + + +def _handshake( + layout: DeepEPM2NLayout, + overrides: dict | None = None, + hosts: str | None = None, + attn_tp_size: int = 1, +) -> list[dict]: + """Handshake infos; rank r runs on host hosts[r] (one NVLink domain per host). + + Attention ranks report their TP size, a capacity of 256 and a bucket + ladder (8, 16, 64); FFN ranks report none, as without an explicit capacity. + """ + hosts = hosts or "a" * layout.union_size + infos = [] + for world_rank in range(layout.union_size): + is_ffn = world_rank < layout.ffn_size + info = local_handshake_info( + role="ffn" if is_ffn else "attention", + role_rank=world_rank if is_ffn else world_rank - layout.ffn_size, + world_rank=world_rank, + config=DeepEPM2NConfig(), + hidden_size=M2N_HIDDEN_SIZE, + dtype=torch.bfloat16, + num_experts=layout.num_experts, + top_k=8, + num_lanes=1, + capacity_tokens_per_rank=None if is_ffn else 256, + bucket_ladder=None if is_ffn else (8, 16, 64), + attn_tp_size=None if is_ffn else attn_tp_size, + compute_gate_on_attention=True, + expert_alignment=128 if is_ffn else None, + ffn_deployment=_ffn_deployment(layout) if is_ffn else None, + sm=HOPPER, + topology=RankTopology( + world_rank=world_rank, + host=hosts[world_rank], + nvlink_domain=f"host:{hosts[world_rank]}", + local_index=hosts[: world_rank + 1].count(hosts[world_rank]) - 1, + ), + ) + infos.append(info) + for (world_rank, key), value in (overrides or {}).items(): + infos[world_rank][key] = value + return infos + + +def test_connector_is_registered(): + """The factory resolves ``DeepEPM2NAFDConnector`` to this module's class. + + Why: deployments select the connector by name in the AFD config. + """ + cls = AFDConnectorFactory.get_connector_class(CONNECTOR) + + assert cls.__name__ == CONNECTOR + assert cls.__module__ == "afd_plugin.connectors.gpu.deepep_m2n.connector" + + +def test_config_defaults_and_round_trip(): + """An empty config gives DeepGEMM experts, a BF16 payload and afdEP's role + path, and the config survives to_mapping/from_mapping unchanged. + + Why: defaults are what deployments get, and the handshake compares + serialized configs. + """ + config = DeepEPM2NConfig.from_mapping(None) + masked = DeepEPM2NConfig.from_mapping({"role_aware": False}) + + assert config.expert_backend == "deepgemm" + assert config.payload_format == "bf16" + assert config.role_aware is True + assert masked.role_aware is False + assert DeepEPM2NConfig.from_mapping(config.to_mapping()) == config + assert DeepEPM2NConfig.from_mapping(masked.to_mapping()) == masked + + +@pytest.mark.parametrize( + ("raw", "message"), + [ + ({"bogus": 1}, "unknown DeepEPM2NAFDConnector options"), + ({"expert_backend": "cutlass"}, "expert_backend"), + ({"payload_format": "fp4"}, "payload_format"), + # Removed in Phase 4b (now under ``network``): rejected, not ignored. + ({"disable_gin": "auto"}, "unknown DeepEPM2NAFDConnector options.*disable_gin"), + ({"num_lanes": 3}, "num_lanes"), + ({"capacity_tokens_per_rank": 0}, "capacity_tokens_per_rank"), + ({"bucket_granularity": 0}, "bucket_granularity"), + ({"cpu_sync_min_tokens": -1}, "cpu_sync_min_tokens"), + ({"control_backend": "mpi"}, "control_backend"), + ({"network": {"preset": "token-ring"}}, "network.preset"), + ({"num_qps": -1}, "num_qps"), + ({"expert_num_sms": -2}, "expert_num_sms"), + ({"expert_num_sms": 131}, "expert_num_sms"), + ({"cpu_timeout_s": 0}, "cpu_timeout_s"), + ({"role_aware": "true"}, "role_aware must be a boolean"), + ({"role_aware": 1}, "role_aware must be a boolean"), + ], +) +def test_config_rejects_invalid_options(raw, message): + """Each bad or stale option (unknown or removed key, bad enum, zero or + negative size, bad network preset, bad timeout, non-boolean protocol + switch) raises a ValueError naming the field. + + Why: typos and removed keys must fail at startup, not be ignored. + """ + with pytest.raises(ValueError, match=message): + DeepEPM2NConfig.from_mapping(raw) + + +def test_config_accepts_fp8_payload_and_sync_threshold(): + """The FP8 payload format and ``cpu_sync_min_tokens`` parse as given. + + Why: the Phase 2 options must reach the connector intact. + """ + config = DeepEPM2NConfig.from_mapping( + {"payload_format": "fp8_block128_fp32", "cpu_sync_min_tokens": 512} + ) + + assert config.payload_format == "fp8_block128_fp32" + assert config.cpu_sync_min_tokens == 512 + + +def test_expert_num_sms_sizes_the_deepgemm_grid(monkeypatch): + """``expert_num_sms`` defaults to all SMs and, when set, sizes DeepGEMM's + persistent grid once, at expert construction; above the device's SM count + it is rejected. + + Why: under DBO a grid that needs every SM straggles behind the other + lane's waiting transport kernel. + """ + from afd_plugin.connectors.gpu.deepep_m2n import experts as experts_module + + calls: list[int] = [] + monkeypatch.setattr( + experts_module.deepgemm, "set_contiguous_m_alignment", lambda a: None + ) + monkeypatch.setattr(experts_module.deepgemm, "set_num_sms", calls.append) + monkeypatch.setattr(torch.cuda, "current_device", lambda: 0) + monkeypatch.setattr( + torch.cuda, + "get_device_properties", + lambda device: SimpleNamespace(multi_processor_count=132), + ) + + assert DeepEPM2NConfig.from_mapping({}).expert_num_sms == 0 + experts_module.DeepEPM2NExperts("deepgemm", 128) + assert calls == [] + experts_module.DeepEPM2NExperts("deepgemm", 128, 130) + assert calls == [130] + with pytest.raises(ValueError, match="exceeds the device"): + experts_module.DeepEPM2NExperts("deepgemm", 128, 134) + + +def test_plan_stage_skips_host_sync_below_threshold(): + """With a threshold of 256, an eager 255-token step skips the host sync + and a 256-token step keeps it. + + Why: pins the boundary of the small-batch no-sync shortcut. + """ + + def plan(tokens: int): + return plan_stage( + _dp(tokens), + attn_size=1, + attn_tp_size=1, + capacity_tokens_per_rank=1024, + bucket_granularity=1, + graph_mode=False, + bucket_ladder=(255,), + cpu_sync_min_tokens=256, + ) + + assert not plan(255).do_cpu_sync + assert plan(256).do_cpu_sync + + +def test_fp8_block128_round_trip(): + """Block-128 FP8 quantization gives e4m3 values and one scale per + 128x128 block, and dequantizes back within 5% relative error. + + Why: test weights and the reference math rely on this helper. + """ + w = torch.randn(2, 256, 384) + + w_q, scale = quantize_fp8_block128(w) + + assert w_q.dtype == torch.float8_e4m3fn + assert scale.shape == (2, 2, 3) + error = (dequantize_fp8_block128(w_q, scale) - w).norm() / w.norm() + assert error < 5e-2 + + +def test_expert_weights_require_matching_fp8_scales(): + """FP8 weights with scales are FP8, BF16 weights without scales are not, + and FP8 weights without scales are rejected. + + Why: a half-configured FP8 layer must fail loudly, not compute garbage. + """ + w13, w2 = torch.zeros(1, 256, 128), torch.zeros(1, 128, 128) + w13_q, w13_s = quantize_fp8_block128(w13) + w2_q, w2_s = quantize_fp8_block128(w2) + + assert LocalExpertWeights(w13_q, w2_q, w13_s, w2_s).is_fp8 + assert not LocalExpertWeights(w13, w2).is_fp8 + with pytest.raises(ValueError, match="block scales"): + LocalExpertWeights(w13_q, w2_q) + + +def test_layout_is_ffn_first_with_identity_expert_ids(): + """2F3A with 128 experts: 64 experts per FFN rank, 5x64 DeepEP expert + slots, FFN ranks first in world order, FFN 1 owning experts 64-127. + + Why: real expert ids must equal DeepEP ids in the plugin's rank order. + """ + layout = DeepEPM2NLayout(ffn_size=2, attn_size=3, num_experts=128) + + assert layout.experts_per_rank == 64 + assert layout.deepep_num_experts == 5 * 64 + assert layout.world_rank("ffn", 1) == 1 + assert layout.world_rank("attention", 0) == 2 + assert layout.expert_range(1) == range(64, 128) + + +def test_layout_rejects_indivisible_experts(): + """128 experts over 3 FFN ranks is rejected. + + Why: uneven expert splits are unsupported and must be refused. + """ + with pytest.raises(ValueError, match="divisible"): + DeepEPM2NLayout(ffn_size=3, attn_size=1, num_experts=128) + + +def test_tokens_per_attention_rank_splits_each_dp_rank_over_its_tp_peers(): + """Without TP each Attention rank dispatches its DP rank's count; with TP + each TP peer dispatches a ceil(count / TP) slice; a count list that does + not cover the Attention ranks raises. + + Why: TP peers hold the same tokens; dispatching them all from every peer + would run each token's experts TP times. + """ + assert tokens_per_attention_rank(_dp(3, 5), attn_size=2, attn_tp_size=1) == ( + 3, + 5, + ) + assert tokens_per_attention_rank(_dp(3, 5), attn_size=4, attn_tp_size=2) == ( + 2, + 2, + 3, + 3, + ) + assert tokens_per_attention_rank(_dp(1, 0), attn_size=8, attn_tp_size=4) == ( + 1, + 1, + 1, + 1, + 0, + 0, + 0, + 0, + ) + with pytest.raises(ValueError, match="token counts"): + tokens_per_attention_rank(_dp(3, 5, 7), attn_size=2, attn_tp_size=1) + # Two counts with TP 1 do not cover four ranks: no TP is guessed. + with pytest.raises(ValueError, match="token counts"): + tokens_per_attention_rank(_dp(3, 5), attn_size=4, attn_tp_size=1) + + +def test_plan_stage_sizes_the_bucket_for_the_tp_slice(): + """A 17-token DP rank with Attention TP 2 dispatches 9-token slices. + + Why: the bucket and capacity check apply to what one rank dispatches. + """ + plan = plan_stage( + _dp(17), + attn_size=2, + attn_tp_size=2, + capacity_tokens_per_rank=9, + bucket_granularity=1, + graph_mode=False, + bucket_ladder=(9,), + ) + + assert plan.tokens_per_attn_rank == (9, 9) + assert plan.bucket == 9 + + +def test_plan_stage_rounds_graph_buckets_up_the_ladder_without_host_sync(): + """With ladder (8, 16, 32, 64), graph-mode steps of (17, 4), (16, 16) and + (0, 0) tokens get buckets 32, 16 and 8 and never a host sync. + + Why: graph steps sit on a rung (vLLM pads them to a capture size) and must + keep it; anything else rounds up to the next rung, never to 0, and graph + mode must never block on the host. + """ + + def plan(*tokens: int): + return plan_stage( + _dp(*tokens), + attn_size=2, + attn_tp_size=1, + capacity_tokens_per_rank=64, + bucket_granularity=1, + graph_mode=True, + bucket_ladder=(8, 16, 32, 64), + ) + + assert (plan(17, 4).bucket, plan(17, 4).do_cpu_sync) == (32, False) + assert plan(17, 4).tokens_per_attn_rank == (17, 4) + assert plan(16, 16).bucket == 16 + assert plan(0, 0).bucket == 8 + + +def test_plan_stage_gives_buckets_above_the_no_sync_limit_the_capacity(): + """With ladder (8, 16) and capacity 64, eager steps of 17, 40 and 64 + tokens all get bucket 64, while a 5-token step rounds up to rung 8. + + Why: DeepEP compiles its kernels per bucket; a bucket per eager token + count ran nvcc on the request path (seconds per new prefill size). + """ + + def plan(tokens: int): + return plan_stage( + _dp(tokens), + attn_size=1, + attn_tp_size=1, + capacity_tokens_per_rank=64, + bucket_granularity=1, + graph_mode=False, + bucket_ladder=(8, 16), + ) + + assert [plan(tokens).bucket for tokens in (17, 40, 64)] == [64] * 3 + assert (plan(5).bucket, plan(5).do_cpu_sync) == (8, True) + + +@pytest.mark.parametrize("tokens", [8, 40]) +def test_plan_stage_bucket_does_not_depend_on_the_rank_s_mode(tokens): + """A step planned as graph on one rank and eager on another gets the same + bucket, both on the ladder (8 tokens) and above it (40). + + Why: an FFN eager fallback can plan a step Attention replays in another + mode, and DeepEP needs the same bucket on every rank; a capacity bucket + chosen by the sync decision broke test_graph_and_eager_plans_mixed_across_ranks. + """ + buckets = { + plan_stage( + _dp(tokens), + attn_size=1, + attn_tp_size=1, + capacity_tokens_per_rank=64, + bucket_granularity=1, + graph_mode=graph_mode, + bucket_ladder=(8, 16), + cpu_sync_min_tokens=sync_min, + ).bucket + for graph_mode in (False, True) + for sync_min in (0, 10) + } + + assert len(buckets) == 1 + + +def test_plan_stage_keeps_a_ladder_bucket_for_eager_steps_without_host_sync(): + """With a threshold of 256 and ladder (32, 128, 255), an eager 20-token + step skips the sync with rung 32, while a 300-token step syncs with the + capacity bucket. + + Why: without the sync DeepEP allocates the worst case for the bucket, so + a capacity bucket there would allocate for the capacity on every step. + """ + + def plan(tokens: int): + return plan_stage( + _dp(tokens), + attn_size=1, + attn_tp_size=1, + capacity_tokens_per_rank=1024, + bucket_granularity=1, + graph_mode=False, + bucket_ladder=(32, 128, 255), + cpu_sync_min_tokens=256, + ) + + assert (plan(20).bucket, plan(20).do_cpu_sync) == (32, False) + assert (plan(300).bucket, plan(300).do_cpu_sync) == (1024, True) + + +@pytest.mark.parametrize( + ("sizes", "tp", "granularity", "sync_min", "capacity", "expected"), + [ + ((1, 2, 4, 8, 16, 24, 32), 1, 1, 0, 1024, (1, 2, 4, 8, 16, 24, 32)), + ((1, 2, 4, 8, 16, 24, 32), 2, 1, 0, 1024, (1, 2, 4, 8, 12, 16)), + ((1, 2, 4, 8, 16, 24, 32), 1, 8, 0, 1024, (8, 16, 24, 32)), + ((1, 2, 4, 8, 16, 24, 32), 1, 1, 100, 1024, (1, 2, 4, 8, 16, 24, 32, 99)), + ((1, 2, 4, 8, 16, 24, 32), 1, 1, 0, 20, (1, 2, 4, 8, 16, 20)), + ((), 1, 1, 0, 1024, ()), + ], + ids=["plain", "tp2", "granularity8", "sync-threshold", "capped", "no-graphs"], +) +def test_bucket_ladder(sizes, tp, granularity, sync_min, capacity, expected): + """The ladder holds each capture size's TP slice, rounded to the + granularity, plus the no-sync limit as its top rung, capped at capacity; + without graphs or a sync threshold it is empty. + + Why: graph steps must land on a rung, the top rung sizes the combine zero + buffer, and an empty ladder sends every eager step to the capacity. + """ + ladder = bucket_ladder( + capacity_tokens_per_rank=capacity, + cudagraph_capture_sizes=sizes, + attn_tp_size=tp, + bucket_granularity=granularity, + cpu_sync_min_tokens=sync_min, + ) + + assert ladder == expected + + +@pytest.mark.parametrize( + ("ladder", "sync_min", "expected"), + [ + ((8, 16, 32), 0, ((8, True), (16, True), (32, True), (64, True))), + ( + (8, 16, 32), + 20, + ((32, True), (64, True), (8, False), (16, False), (32, False)), + ), + ((), 0, ((64, True),)), + ], + ids=["all-synced", "sync-threshold", "no-ladder"], +) +def test_precompile_buckets(ladder, sync_min, expected): + """Synced eager steps can land on rungs at or above the sync threshold + and on the capacity (64); steps below the threshold skip the sync on any + rung up to the first one covering the threshold. + + Why: every pair plan_stage can give an eager step must be compiled at + startup, or nvcc runs on the request path again. + """ + assert precompile_buckets(ladder, 64, sync_min) == expected + + +class _RecordingLane: + """A DeepEP buffer stand-in that records dispatch and combine calls.""" + + def __init__(self) -> None: + self.calls: list[tuple] = [] + + def dispatch(self, x, x_scale, topk_ids, topk_weights, **kwargs): + self.calls.append( + ("dispatch", kwargs["bucket"], kwargs["do_cpu_sync"], x.shape[0]) + ) + return SimpleNamespace(handle=SimpleNamespace(recv_shape=(0, M2N_HIDDEN_SIZE))) + + def combine(self, expert_output, handle): + self.calls.append(("combine", tuple(expert_output.shape))) + + +def test_precompile_runs_the_same_empty_steps_on_attention_and_ffn(monkeypatch): + """Attention and FFN each run one empty dispatch and combine per lane for + every precompile bucket, in the same order. + + Why: dispatch and combine are collectives; ranks that warm up different + buckets would hang at startup. + """ + monkeypatch.setattr(torch.cuda, "synchronize", lambda device=None: None) + topology = M2NTopology(ffn_size=2, attn_size=2) + calls = {} + for role in ("attention", "ffn"): + connector = _connector(role, 0, topology, {"capacity_tokens_per_rank": 64}) + connector.device = torch.device("cpu") + connector.capacity_tokens_per_rank = 64 + connector.bucket_ladder = (8, 16, 64) + connector.expert_alignment = DEEPGEMM_ALIGNMENT + connector.lanes = [_RecordingLane(), _RecordingLane()] + if role == "attention": + connector._allocate_combine_zeros() + connector._precompile_dispatch_kernels() + calls[role] = [lane.calls for lane in connector.lanes] + + expected = [ + call + for bucket in (8, 16, 64) + for call in (("dispatch", bucket, True, 0), ("combine", (0, M2N_HIDDEN_SIZE))) + ] + assert calls["attention"] == calls["ffn"] == [expected, expected] + + +def test_plan_stage_rejects_bucket_above_capacity(): + """65 tokens against a capacity of 64 raises. + + Why: an overflow must be an error, not memory corruption in DeepEP. + """ + with pytest.raises(ValueError, match="exceeds capacity"): + plan_stage( + _dp(65), + attn_size=1, + attn_tp_size=1, + capacity_tokens_per_rank=64, + bucket_granularity=1, + graph_mode=False, + bucket_ladder=(64,), + ) + + +def test_handshake_accepts_consistent_ranks(): + """Four consistent handshake records pass, giving expert alignment 128 + and no GIN (one host). + + Why: the normal path must pass and yield the right settings. + """ + layout = DeepEPM2NLayout(ffn_size=2, attn_size=2, num_experts=128) + + result = validate_handshake(_handshake(layout), layout, NetworkConfig()) + + assert result.expert_alignment == 128 + assert not result.network_plan.use_gin + assert (result.attn_tp_size, result.capacity_tokens_per_rank) == (1, 256) + + +def test_handshake_gives_ffn_the_attention_tp_size_capacity_and_ladder(): + """FFN ranks report no TP size, capacity or bucket ladder; the result + carries Attention's. + + Why: the FFN runs without TP or Attention's CUDA graphs, so only + Attention knows its TP size, the capacity of a TP slice and its graph + buckets; both roles must size buffers and pick buckets alike. + """ + layout = DeepEPM2NLayout(ffn_size=2, attn_size=4, num_experts=128) + + result = validate_handshake( + _handshake(layout, attn_tp_size=2), layout, NetworkConfig() + ) + + assert ( + result.attn_tp_size, + result.capacity_tokens_per_rank, + result.bucket_ladder, + ) == (2, 256, (8, 16, 64)) + + +@pytest.mark.parametrize( + ("overrides", "message"), + [ + ({(3, "attn_tp_size"): 2}, "disagree on their TP size"), + ({(2, "attn_tp_size"): 3, (3, "attn_tp_size"): 3}, "does not divide"), + ({(3, "capacity_tokens_per_rank"): 128}, "capacity_tokens_per_rank"), + ({(0, "capacity_tokens_per_rank"): 128}, "capacity_tokens_per_rank"), + ({(3, "bucket_ladder"): [8, 32]}, "bucket_ladder"), + ({(0, "cpu_sync_min_tokens"): 128}, "disagree on cpu_sync_min_tokens"), + ({(0, "bucket_granularity"): 16}, "disagree on bucket_granularity"), + ], +) +def test_handshake_rejects_attention_tp_and_capacity_mismatches(overrides, message): + """Attention ranks that disagree on TP, a TP size that does not divide + the Attention ranks, or a capacity some rank disagrees with (an FFN rank + reports one only when configured explicitly) fail the handshake. + + Why: ranks that plan different slices or buffer sizes would hang or + overflow DeepEP buffers. + """ + layout = DeepEPM2NLayout(ffn_size=2, attn_size=2, num_experts=128) + + with pytest.raises(RuntimeError, match=message): + validate_handshake(_handshake(layout, overrides), layout, NetworkConfig()) + + +@pytest.mark.parametrize( + ("overrides", "message"), + [ + ({(3, "hidden_size"): 4096}, "disagree on hidden_size"), + ({(0, "payload_format"): "fp8"}, "disagree on payload_format"), + ({(1, "expert_alignment"): 64}, "expert_alignment"), + ({(0, "role"): "attention"}, "expected 'ffn'"), + ({(2, "sm"): [8, 0]}, "SM90 or newer"), + ], +) +def test_handshake_rejects_mismatches(overrides, message): + """One rank's record differs (hidden size, payload format, expert + alignment, role order, or a pre-Hopper GPU): the handshake raises. + + Why: disagreeing ranks would corrupt data or hang; fail before buffers. + """ + layout = DeepEPM2NLayout(ffn_size=2, attn_size=2, num_experts=128) + + with pytest.raises(RuntimeError, match=message): + validate_handshake(_handshake(layout, overrides), layout, NetworkConfig()) + + +def test_handshake_plans_gin_across_nvlink_domains(): + """Ranks on hosts ``aabb`` give a plan of 2 domains of 2 ranks with GIN. + + Why: the handshake must hand the network plan the right domains. + """ + layout = DeepEPM2NLayout(ffn_size=2, attn_size=2, num_experts=128) + + result = validate_handshake( + _handshake(layout, hosts="aabb"), layout, NetworkConfig() + ) + + assert result.network_plan.use_gin + assert (result.network_plan.num_domains, result.network_plan.ranks_per_domain) == ( + 2, + 2, + ) + + +@pytest.mark.parametrize( + ("hosts", "message"), + [("aaab", "different rank counts"), ("abab", "not contiguous")], +) +def test_handshake_rejects_unsupported_domain_layouts(hosts, message): + """Hosts ``aaab`` (unequal counts) and ``abab`` (not contiguous) raise. + + Why: DeepEP's hybrid mode needs equal, contiguous ranks per domain. + """ + layout = DeepEPM2NLayout(ffn_size=2, attn_size=2, num_experts=128) + + with pytest.raises(RuntimeError, match=message): + validate_handshake(_handshake(layout, hosts=hosts), layout, NetworkConfig()) + + +def test_handshake_rejects_mismatched_network_config(): + """One rank with a different ``network`` config fails the handshake. + + Why: every rank must derive the same NIC and GIN plan. + """ + layout = DeepEPM2NLayout(ffn_size=1, attn_size=1, num_experts=128) + infos = _handshake(layout) + infos[1]["network"] = "something else" + + with pytest.raises(RuntimeError, match="disagree on network"): + validate_handshake(infos, layout, NetworkConfig()) + + +def test_gloo_control_backend_stages_payloads_on_cpu(): + """With ``control_backend: gloo`` control payloads are staged on CPU. + + Why: gloo cannot send CUDA tensors. + """ + connector = _connector( + "attention", + 0, + M2NTopology(ffn_size=1, attn_size=1), + {"control_backend": "gloo"}, + ) + + assert connector.control_device == torch.device("cpu") + + +def test_config_parses_network_section_and_tunables(): + """The ``network`` section (preset plus NIC), ``num_qps`` and timeouts + parse and survive a round trip. + + Why: the Phase 4b options must reach the connector intact. + """ + config = DeepEPM2NConfig.from_mapping( + { + "network": {"preset": "roce-rail-isolated", "nics": ["mlx5_0"]}, + "num_qps": 33, + "gpu_timeout_s": 60, + } + ) + + assert config.network.nic_policy == "shared" + assert config.network.nics == ("mlx5_0",) + assert (config.num_qps, config.gpu_timeout_s) == (33, 60) + assert DeepEPM2NConfig.from_mapping(config.to_mapping()) == config + + +def test_expert_row_ranges_follow_alignment(): + """Per-expert row ranges from the prefix sum: contiguous with alignment 1, + each group starting on a multiple of 128 with alignment 128. + + Why: every receive check reads DeepEP's padded per-expert layout. + """ + psum = torch.tensor([3, 128 + 5, 256]) + + assert expert_row_ranges(psum, 1) == [(0, 3), (3, 133), (133, 256)] + assert expert_row_ranges(psum, 128) == [(0, 3), (128, 133), (256, 256)] + + +def test_connector_constructs_without_initialization(): + """Constructing Attention and FFN connectors computes world ranks without + initializing. Attention knows its TP size and capacity; the FFN learns + both in the handshake unless the capacity is explicit. + + Why: construction must not communicate, and rank math must be right. + """ + topology = M2NTopology(ffn_size=2, attn_size=3) + attention = _connector("attention", 1, topology, {"capacity_tokens_per_rank": 64}) + default_attention = _connector("attention", 0, topology) + ffn = _connector("ffn", 1, topology) + explicit_ffn = _connector("ffn", 0, topology, {"capacity_tokens_per_rank": 64}) + + assert attention.is_initialized is False + assert attention.world_rank == 3 + assert (attention.attn_tp_size, attention.capacity_tokens_per_rank) == (1, 64) + # scheduler max_num_batched_tokens + assert default_attention.capacity_tokens_per_rank == 8192 + assert ffn.world_rank == 1 + assert (ffn.attn_tp_size, ffn.capacity_tokens_per_rank) == (None, None) + assert explicit_ffn.capacity_tokens_per_rank == 64 + + +def test_control_plane_requires_the_handshake_on_ffn(): + """An FFN connector cannot plan a step before the handshake. + + Why: it does not know the Attention TP size until then. + """ + connector = _connector("ffn", 0, M2NTopology(ffn_size=1, attn_size=1)) + payload = AFDControlPayload( + dp_metadata_list={0: _dp(1)}, + is_graph_capturing=False, + is_warmup=False, + ) + + with pytest.raises(RuntimeError, match="before bootstrap"): + connector.control_plane.update_state_from_dp_metadata(payload) + + +def test_connector_rejects_non_bf16_model(): + """An FP16 model is rejected at construction. + + Why: only BF16 hidden states are supported. + """ + topology = M2NTopology(ffn_size=1, attn_size=1) + vllm_config = m2n_vllm_config(topology, "ffn", 0, {}, port=12345) + vllm_config.model_config.dtype = torch.float16 + afd_config = afd_config_from_mapping(vllm_config.additional_config["afd"]) + + with pytest.raises(ValueError, match="bfloat16"): + AFDConnectorFactory.create_connector(0, 0, vllm_config, afd_config) + + +def test_control_plane_derives_one_plan_per_stage(): + """An eager control payload with counts (5, 9) leaves stage 0 with a + host sync and bucket 16, the next rung of the ladder Attention derives + from its capture sizes (1, 2, 4, 8, 16, ... capped at capacity 64). + + Why: the runtime path from control payload to plan must work. + """ + connector = _connector( + "attention", + 0, + M2NTopology(ffn_size=1, attn_size=2), + {"capacity_tokens_per_rank": 64}, + ) + payload = AFDControlPayload( + dp_metadata_list={0: _dp(5, 9)}, + is_graph_capturing=False, + is_warmup=False, + ) + + connector.control_plane.update_state_from_dp_metadata(payload) + + assert connector.stage_plans[0].bucket == 16 + assert connector.stage_plans[0].do_cpu_sync + + +def test_control_plane_rejects_more_stages_than_lanes(): + """Two stages on a one-lane connector raise. + + Why: each in-flight micro-batch needs its own buffer. + """ + connector = _connector("ffn", 0, M2NTopology(ffn_size=1, attn_size=1)) + payload = AFDControlPayload( + dp_metadata_list={0: _dp(1), 1: _dp(1)}, + is_graph_capturing=False, + is_warmup=False, + ) + + with pytest.raises(ValueError, match="num_lanes"): + connector.control_plane.update_state_from_dp_metadata(payload) + + +def test_close_without_init_is_safe(): + """``close()`` on a never-initialized connector succeeds, twice. + + Why: runtimes call close() on every exit path, including before init. + """ + connector = _connector("ffn", 0, M2NTopology(ffn_size=1, attn_size=1)) + + connector.close() + connector.close() + + assert connector.is_initialized is False + + +def test_data_path_requires_initialization(): + """``send_attn_output`` before init raises "not initialized". + + Why: misuse must give a clear Python error, not a native crash. + """ + connector = _connector("attention", 0, M2NTopology(ffn_size=1, attn_size=1)) + context = AFDTransferContext( + metadata=AFDTransferMetadata.create_attention_metadata( + layer_idx=0, stage_idx=0, seq_len=1 + ) + ) + + with pytest.raises(RuntimeError, match="not initialized"): + connector.send_attn_output( + torch.zeros(1, M2N_HIDDEN_SIZE), + context, + topk_ids=torch.zeros(1, 8, dtype=torch.int64), + topk_weights=torch.zeros(1, 8), + ) + + +# --------------------------------------------------------------------------- +# Bring-up: defaults, FFN deployment in the handshake, post-load checks +# --------------------------------------------------------------------------- + +SMALL_HIDDEN = 256 +SMALL_INTERMEDIATE = 128 +FP8_BLOCK = 128 +SMALL_LAYOUT = DeepEPM2NLayout(ffn_size=2, attn_size=1, num_experts=8) +MOE_LAYERS = [1, 3] + + +def _ffn_vllm_config(): + return m2n_vllm_config(M2NTopology(ffn_size=1, attn_size=1), "ffn", 0, {}, port=1) + + +@pytest.mark.parametrize( + ("capture_size", "tp_size", "explicit", "expected"), + [ + (512, 1, None, 8192), + (16384, 1, None, 16384), + (16384, 1, 64, 64), + (512, 2, None, 4096), + (512, 3, None, 2731), + (512, 2, 64, 64), + ], +) +def test_capacity_covers_prefill_budget_and_largest_graph( + capture_size, tp_size, explicit, expected +): + """Default capacity is the larger of max_num_batched_tokens (8192) and + the largest CUDA-graph capture size, divided (rounded up) by the + Attention TP size; an explicit value wins. + + Why: a decode graph larger than the prefill budget must still fit, and + with TP each rank dispatches only its slice. + """ + vllm_config = _ffn_vllm_config() + vllm_config.compilation_config.max_cudagraph_capture_size = capture_size + vllm_config.parallel_config.tensor_parallel_size = tp_size + config = DeepEPM2NConfig(capacity_tokens_per_rank=explicit) + + assert resolve_capacity_tokens_per_rank(config, vllm_config) == expected + + +@pytest.mark.parametrize( + ("use_ubatching", "num_ubatches", "explicit", "expected"), + [(False, 1, None, 1), (True, 2, None, 2), (True, 2, 1, 1)], +) +def test_num_lanes_follows_vllm_micro_batches( + use_ubatching, num_ubatches, explicit, expected +): + """One lane without micro-batching, two under DBO; explicit wins. + + Why: a DBO deployment needs a buffer per in-flight micro-batch. + """ + vllm_config = _ffn_vllm_config() + vllm_config.parallel_config.use_ubatching = use_ubatching + vllm_config.parallel_config.num_ubatches = num_ubatches + + lanes = resolve_num_lanes(DeepEPM2NConfig(num_lanes=explicit), vllm_config) + + assert lanes == expected + + +def test_num_lanes_rejects_more_micro_batches_than_lanes(): + vllm_config = _ffn_vllm_config() + vllm_config.parallel_config.use_ubatching = True + vllm_config.parallel_config.num_ubatches = 3 + + with pytest.raises(ValueError, match="at most 2 lanes"): + resolve_num_lanes(DeepEPM2NConfig(), vllm_config) + + +def test_expert_weight_format_names_what_vllm_loads(): + """BF16 checkpoints, block-128 FP8 checkpoints, and anything else. + + Why: the handshake compares this with the payload format before loading. + """ + from vllm.model_executor.layers.quantization.fp8 import Fp8Config + + block_fp8 = Fp8Config( + is_checkpoint_fp8_serialized=True, weight_block_size=[FP8_BLOCK, FP8_BLOCK] + ) + per_tensor_fp8 = Fp8Config(is_checkpoint_fp8_serialized=True) + + assert expert_weight_format(None) == "bf16" + assert expert_weight_format(block_fp8) == "fp8_block128" + assert expert_weight_format(per_tensor_fp8) == "unsupported (fp8)" + + +@pytest.mark.parametrize( + ("key", "value"), + [ + ("expert_parallel", False), + ("ep_size", 4), + ("tensor_parallel_size", 2), + ("expert_placement_strategy", "round_robin"), + ("eplb", True), + ("all2all_backend", "deepep_low_latency"), + ("weight_format", "fp8_block128"), + ], +) +def test_handshake_rejects_ffn_deployments_that_misplace_experts(key, value): + """An FFN rank whose vLLM deployment would hold other experts, run + vLLM's own all-to-all, or load weights in another format fails the + handshake on every rank. + + Why: rows would reach the wrong experts with no error later. + """ + layout = DeepEPM2NLayout(ffn_size=2, attn_size=2, num_experts=128) + infos = _handshake(layout) + infos[1]["ffn_deployment"] = {**infos[1]["ffn_deployment"], key: value} + + with pytest.raises( + RuntimeError, match=f"FFN world rank 1 is deployed with .*{key}" + ): + validate_handshake(infos, layout, NetworkConfig()) + + +def test_handshake_accepts_fp8_block_weights_for_fp8_payload(): + layout = DeepEPM2NLayout(ffn_size=2, attn_size=2, num_experts=128) + infos = _handshake( + layout, + {(rank, "payload_format"): "fp8_block128_fp32" for rank in range(4)}, + ) + for rank in range(layout.ffn_size): + infos[rank]["ffn_deployment"] = _ffn_deployment(layout, "fp8_block128") + + validate_handshake(infos, layout, NetworkConfig()) + + +def test_handshake_requires_gate_on_attention(): + """Both roles must route on Attention; a mix, or false everywhere, fails. + + Why: DeepEP M2N dispatches by top-k ids computed on Attention. + """ + layout = DeepEPM2NLayout(ffn_size=1, attn_size=1, num_experts=128) + mixed = _handshake(layout, {(0, "compute_gate_on_attention"): False}) + neither = _handshake( + layout, + {(rank, "compute_gate_on_attention"): False for rank in range(2)}, + ) + + with pytest.raises(RuntimeError, match="disagree on compute_gate_on_attention"): + validate_handshake(mixed, layout, NetworkConfig()) + with pytest.raises(RuntimeError, match="compute_gate_on_attention=true"): + validate_handshake(neither, layout, NetworkConfig()) + + +def _local_experts( + expert_ids: range, *, fp8: bool = False, w2_scale_blocks: int | None = None +) -> FFNLocalExperts: + num_experts = len(expert_ids) + w13_shape = (num_experts, 2 * SMALL_INTERMEDIATE, SMALL_HIDDEN) + w2_shape = (num_experts, SMALL_HIDDEN, SMALL_INTERMEDIATE) + if fp8: + weights = LocalExpertWeights( + w13=torch.zeros(w13_shape, dtype=torch.float8_e4m3fn), + w2=torch.zeros(w2_shape, dtype=torch.float8_e4m3fn), + w13_scale=torch.ones( + num_experts, + 2 * SMALL_INTERMEDIATE // FP8_BLOCK, + SMALL_HIDDEN // FP8_BLOCK, + ), + w2_scale=torch.ones( + num_experts, + SMALL_HIDDEN // FP8_BLOCK, + w2_scale_blocks or SMALL_INTERMEDIATE // FP8_BLOCK, + ), + ) + else: + weights = LocalExpertWeights( + w13=torch.zeros(w13_shape, dtype=torch.bfloat16), + w2=torch.zeros(w2_shape, dtype=torch.bfloat16), + ) + return FFNLocalExperts(expert_ids=tuple(expert_ids), weights=weights) + + +def _model_infos( + *, + ffn_ranges: tuple[range, range] = (range(0, 4), range(4, 8)), + attention_layers: list[int] = MOE_LAYERS, + fp8: bool = False, + w2_scale_blocks: int | None = None, +) -> list: + infos = [ + local_model_info( + MOE_LAYERS, + { + layer: _local_experts(experts, fp8=fp8, w2_scale_blocks=w2_scale_blocks) + for layer in MOE_LAYERS + }, + ) + for experts in ffn_ranges + ] + infos.append(local_model_info(attention_layers, None)) + return infos + + +def test_model_check_accepts_matching_layers_and_expert_placement(): + """Two FFN ranks hold experts [0, 4) and [4, 8) of layers 1 and 3, and + Attention dispatches the same layers: the check passes, BF16 and FP8. + + Why: this is the deployment DeepEP M2N routes for. + """ + validate_model_infos( + _model_infos(), SMALL_LAYOUT, hidden_size=SMALL_HIDDEN, payload_format="bf16" + ) + validate_model_infos( + _model_infos(fp8=True), + SMALL_LAYOUT, + hidden_size=SMALL_HIDDEN, + payload_format="fp8_block128_fp32", + ) + + +def test_model_check_skips_when_no_rank_reports(): + """Ranks that report nothing (the connector harness) skip the check.""" + validate_model_infos( + [None, None, None], + SMALL_LAYOUT, + hidden_size=SMALL_HIDDEN, + payload_format="bf16", + ) + + +@pytest.mark.parametrize( + ("infos_kwargs", "payload_format", "message"), + [ + ({"attention_layers": [1]}, "bf16", "disagree on MoE layers"), + ( + {"ffn_ranges": (range(0, 4), range(0, 4))}, + "bf16", + r"routes experts \[4, 8\) to it", + ), + ({}, "fp8_block128_fp32", "do not match payload_format"), + ({"fp8": True, "w2_scale_blocks": 2}, "fp8_block128_fp32", "w2_scale"), + ], + ids=["layers", "placement", "format", "scale-shape"], +) +def test_model_check_rejects_what_would_compute_wrong_rows( + infos_kwargs, payload_format, message +): + """Ranks running different MoE layers, an FFN rank holding experts DeepEP + does not route to it, weights in another format, or scales of the wrong + shape fail on every rank after loading. + + Why: each of these hangs the collectives or silently corrupts outputs. + """ + with pytest.raises(RuntimeError, match=message): + validate_model_infos( + _model_infos(**infos_kwargs), + SMALL_LAYOUT, + hidden_size=SMALL_HIDDEN, + payload_format=payload_format, + ) + + +def test_model_check_requires_ffn_experts_once_anyone_reports(): + infos = _model_infos() + infos[1] = None + + with pytest.raises(RuntimeError, match="FFN world rank 1 reported no local"): + validate_model_infos( + infos, SMALL_LAYOUT, hidden_size=SMALL_HIDDEN, payload_format="bf16" + ) + + +def test_attention_init_rejects_local_experts_before_communicating(): + """Passing experts on an Attention rank raises before any rendezvous. + + Why: a wiring bug must not leave the peers waiting in a collective. + """ + connector = _connector("attention", 0, M2NTopology(ffn_size=1, attn_size=1)) + + with pytest.raises(ValueError, match="only FFN ranks hold local experts"): + connector.init_afd_connector(local_experts={}) + + assert connector.bootstrap_group is None + + +@pytest.mark.parametrize(("role", "role_rank"), [("attention", 2), ("ffn", 1)]) +@pytest.mark.parametrize("role_aware", [True, False]) +def test_lane_transport_is_afdep_with_the_protocol_and_ffn_rank_count( + monkeypatch, role, role_rank, role_aware +): + """Each lane's transport is afdEP's ``ElasticBuffer``, built with the + selected protocol and the FFN rank count; the DeepEP M2N arguments pass + through unchanged. + + Why: afdEP's role path needs the FFN-first rank split, and every rank must + request the same protocol. + """ + monkeypatch.setattr(connector_module, "ElasticBuffer", FakeElasticBuffer) + connector = _connector( + role, + role_rank, + M2NTopology(ffn_size=2, attn_size=3), + {"role_aware": role_aware}, + ) + + buffer = connector._new_buffer(capacity_tokens_per_rank=64, top_k=4) + + assert isinstance(buffer, FakeElasticBuffer) + assert buffer.kwargs == { + "capacity_tokens_per_rank": 64, + "top_k": 4, + "role_aware": role_aware, + "ffn_ranks": 2, + } + + +def test_attention_ops_take_the_stage_from_the_live_dbo_micro_batch(monkeypatch): + """Under DBO the dispatch/combine op bodies use the calling micro-batch's + id, not the ``stage_idx`` argument; outside DBO they use the argument. + + Why: torch.compile froze ``stage_idx=0`` into the graph both DBO + micro-batch threads share, so micro-batch 1 dispatched "stage 0" again + and raised, and every rank hung at the DeepEP barrier. + """ + from afd_plugin.connectors.gpu.deepep_m2n import connector as connector_module + + class _RecordingConnector: + def __init__(self): + self.calls = [] + + def dispatch_attention(self, hidden_states, topk_ids, topk_weights, stage): + self.calls.append(("dispatch", stage)) + + def combine_attention(self, out, stage): + self.calls.append(("combine", stage)) + + recorder = _RecordingConnector() + connector_id = -1 + monkeypatch.setitem(connector_module._CONNECTORS, connector_id, recorder) + tensor = torch.zeros(1) + traced_stage_idx = 0 + + monkeypatch.setattr(connector_module, "dbo_enabled", lambda: True) + monkeypatch.setattr(connector_module, "dbo_current_ubatch_id", lambda: 1) + connector_module._dispatch_op_impl( + tensor, tensor, tensor, connector_id, traced_stage_idx + ) + connector_module._combine_op_impl(tensor, connector_id, traced_stage_idx) + + monkeypatch.setattr(connector_module, "dbo_enabled", lambda: False) + connector_module._dispatch_op_impl(tensor, tensor, tensor, connector_id, 1) + + assert recorder.calls == [("dispatch", 1), ("combine", 1), ("dispatch", 1)] + + +# --------------------------------------------------------------------------- +# Attention-side combine zeros: sized for the largest receive +# --------------------------------------------------------------------------- + +DEEPGEMM_ALIGNMENT = 128 + + +@pytest.mark.parametrize( + ("num_ranks", "bucket", "top_k", "experts_per_rank", "alignment", "expected"), + [ + # Qwen3-30B 2A2F at the default capacity, and at a 512 bucket. + (4, 8192, 8, 64, DEEPGEMM_ALIGNMENT, 270_336), + (4, 512, 8, 64, DEEPGEMM_ALIGNMENT, 24_576), + # Qwen3-235B 4A2F at a 512 bucket. + (6, 512, 8, 64, DEEPGEMM_ALIGNMENT, 32_768), + # No alignment (fused_experts backend): exactly ranks x bucket x top-k. + (4, 10, 8, 64, 1, 320), + # A token sends at most one route per local expert. + (4, 10, 8, 4, 1, 160), + ], +) +def test_worst_case_receive_rows_mirrors_deepep( + num_ranks, bucket, top_k, experts_per_rank, alignment, expected +): + """Rows match DeepEP's worst-case receive without the host sync. + + Why: the Attention-side zero buffer is sized with this; too few rows + would fail combine, and DeepEP's formula is the contract. + """ + rows = worst_case_receive_rows( + num_ranks=num_ranks, + bucket=bucket, + top_k=top_k, + experts_per_rank=experts_per_rank, + expert_alignment=alignment, + ) + + assert rows == expected + + +@pytest.mark.parametrize( + ("capture", "tp", "granularity", "sync_min", "capacity", "expected"), + [ + (512, 1, 1, 0, 8192, 512), # graph buckets only + (512, 2, 1, 0, 4096, 256), # each TP rank dispatches its slice + (100, 1, 16, 0, 8192, 112), # graph bucket rounded to granularity + (512, 1, 1, 1000, 8192, 999), # eager buckets below the sync threshold + (512, 1, 1, 0, 256, 256), # never above capacity + (0, 1, 1, 0, 8192, 0), # eager-only and always syncing + ], +) +def test_no_sync_bucket_limit(capture, tp, granularity, sync_min, capacity, expected): + """The largest bucket that skips the host sync: the largest graph bucket or + the largest eager bucket below cpu_sync_min_tokens, capped by capacity. + + Why: only those dispatches give an Attention rank a worst-case receive; + sizing the zero buffer for capacity instead would waste ~1 GiB per GPU. + """ + limit = no_sync_bucket_limit( + capacity_tokens_per_rank=capacity, + max_cudagraph_capture_size=capture, + attn_tp_size=tp, + bucket_granularity=granularity, + cpu_sync_min_tokens=sync_min, + ) + + assert limit == expected + + +def test_attention_combine_zeros_cover_the_largest_no_sync_receive(): + """Attention allocates one zero buffer for its largest no-sync receive and + hands combine a contiguous view of its first rows. + + Why: a cache keyed by receive shape grew without bound when eager buckets + skipped the sync; a fixed buffer bounds memory and keeps one address for + captured graphs. + """ + connector = _connector( + "attention", + 0, + M2NTopology(ffn_size=2, attn_size=2), + {"capacity_tokens_per_rank": 16}, + ) + connector.device = torch.device("cpu") + connector.expert_alignment = DEEPGEMM_ALIGNMENT + + connector._allocate_combine_zeros() + + # Bucket min(16, 512 capture) = 16: 4 ranks x 16 x top-k 8 + 127 x 64. + expected_rows = worst_case_receive_rows( + num_ranks=4, + bucket=16, + top_k=8, + experts_per_rank=64, + expert_alignment=DEEPGEMM_ALIGNMENT, + ) + zeros = connector._combine_zeros + assert zeros is not None + assert zeros.shape == (expected_rows, M2N_HIDDEN_SIZE) + view = connector._combine_zero_rows((5, M2N_HIDDEN_SIZE)) + assert view.shape == (5, M2N_HIDDEN_SIZE) and view.is_contiguous() + assert view.data_ptr() == zeros.data_ptr() + assert connector._combine_zero_rows((0, M2N_HIDDEN_SIZE)).shape == ( + 0, + M2N_HIDDEN_SIZE, + ) + with pytest.raises(RuntimeError, match="no_sync_bucket_limit"): + connector._combine_zero_rows((expected_rows + 1, M2N_HIDDEN_SIZE)) + + +class _LoggedEvent: + """CPU stand-in for ``torch.cuda.Event``: logs where it was recorded.""" + + _ids = itertools.count() + + def __init__(self, log: list) -> None: + self._log = log + self.name = f"ev{next(self._ids)}" + + def record(self, stream) -> None: + self._log.append(("record", self.name, stream.name)) + + +class _LoggedStream: + def __init__(self, name: str, log: list) -> None: + self.name = name + self._log = log + + def wait_event(self, event) -> None: + self._log.append(("wait_event", self.name, event.name)) + + def wait_stream(self, other) -> None: + self._log.append(("wait_stream", self.name, other.name)) + + +class _FFNLane: + def dispatch(self, x, x_scale, topk_ids, topk_weights, **kwargs): + rows = torch.zeros((0, M2N_HIDDEN_SIZE)) + return SimpleNamespace( + recv_x=rows, + recv_scale=None, + recv_topk_weights=torch.zeros(0), + recv_count_psum=torch.zeros(0, dtype=torch.int32), + recv_counts_host=[], + handle=SimpleNamespace(recv_shape=(0, M2N_HIDDEN_SIZE)), + ) + + def combine(self, expert_output, handle): + return expert_output + + +def _logged_ffn_connector(monkeypatch, num_lanes: int): + """FFN connector whose streams and events log every record and wait, and + every allocation (with the stream it is made under) and keep-alive.""" + from afd_plugin.connectors.gpu.deepep_m2n import connector as m2n + + log: list = [] + main = _LoggedStream("main", log) + active: list[str] = [] + + @contextlib.contextmanager + def stream_context(stream): + active.append(stream.name) + try: + yield + finally: + active.pop() + + real_empty = torch.empty + + def logged_empty(*args, **kwargs): + log.append(("alloc", active[-1] if active else "main")) + return real_empty(*args, **kwargs) + + monkeypatch.setattr(torch.cuda, "current_stream", lambda device=None: main) + monkeypatch.setattr(torch.cuda, "Event", lambda *a, **k: _LoggedEvent(log)) + monkeypatch.setattr(torch.cuda, "stream", stream_context) + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) + monkeypatch.setattr(torch, "empty", logged_empty) + monkeypatch.setattr( + m2n, + "_keep_alive_on", + lambda stream, *tensors: log.append(("keep_alive", stream.name)), + ) + + connector = _connector("ffn", 0, M2NTopology(ffn_size=2, attn_size=2)) + connector.device = torch.device("cpu") + connector._initialized = True + connector.lanes = [_FFNLane() for _ in range(num_lanes)] + connector.lane_streams = ( + [_LoggedStream(f"lane{i}", log) for i in range(num_lanes)] + if num_lanes > 1 + else [] + ) + connector.stage_plans = { + lane: SimpleNamespace(bucket=8, do_cpu_sync=False) for lane in range(num_lanes) + } + + def experts(x, scale, weights, psum, local, out=None): + log.append(("experts", "main", None)) + return out + + connector.experts = experts + return connector, log + + +def _ffn_dbo_step(connector, num_layers: int, num_stages: int) -> None: + """The routed FFN runner's order: deferred sends, then join lanes.""" + pending: dict[int, tuple] = {} + for _layer in range(num_layers): + for stage in range(num_stages): + if stage in pending: + connector.send_ffn_output(*pending.pop(stage)) + payload = connector.recv_attn_output(stage) + out = connector.run_experts(payload, None) + pending[stage] = (out, payload.context) + for stage in range(num_stages): + connector.send_ffn_output(*pending.pop(stage)) + connector.synchronize_lanes() + + +def test_ffn_lane_combine_waits_only_for_its_own_experts(monkeypatch): + """Under DBO each lane's combine waits on the event recorded right after + that lane's experts, and no lane waits on the whole caller stream. + + Why: waiting on the whole stream also waited for the other micro-batch's + experts, which serialized Attention and FFN compute (no DBO overlap). + """ + connector, log = _logged_ffn_connector(monkeypatch, num_lanes=2) + + _ffn_dbo_step(connector, num_layers=3, num_stages=2) + + assert not [e for e in log if e[0] == "wait_stream"] + # The event each experts launch is followed by on the caller stream. + experts_events = [] + for i, entry in enumerate(log): + if entry[0] == "experts": + record = log[i + 1] + assert record[0] == "record" and record[2] == "main" + experts_events.append(record[1]) + # Experts run E0(L0) E1(L0) E0(L1) E1(L1) ...; combine C_s(L) waits E_s(L). + lane_waits = { + lane: [e[2] for e in log if e[0] == "wait_event" and e[1] == lane] + for lane in ("lane0", "lane1") + } + step_fork = lane_waits["lane0"][0] + assert lane_waits["lane1"][0] == step_fork + assert lane_waits["lane0"][1:] == experts_events[0::2] + assert lane_waits["lane1"][1:] == experts_events[1::2] + # The step fork is recorded before any experts of the step. + fork_at = log.index(("record", step_fork, "main")) + assert fork_at < next(i for i, e in enumerate(log) if e[0] == "experts") + + +def test_ffn_lanes_join_each_step_once(monkeypatch): + """Each lane joins the caller stream once per step; synchronize_lanes + starts a new step (a new fork event). + + Why: the per-step join orders a lane after the previous step and, while + capturing a CUDA graph, joins the lane to the capture. + """ + connector, log = _logged_ffn_connector(monkeypatch, num_lanes=2) + + _ffn_dbo_step(connector, num_layers=1, num_stages=2) + first = len(log) + _ffn_dbo_step(connector, num_layers=1, num_stages=2) + + def forks(entries): + return [ + e + for e in entries + if e[0] == "wait_event" + and e[1].startswith("lane") + and e[2] not in experts_ids + ] + + experts_ids = {log[i + 1][1] for i, e in enumerate(log) if e[0] == "experts"} + first_forks, second_forks = forks(log[:first]), forks(log[first:]) + assert [e[1] for e in first_forks] == ["lane0", "lane1"] + assert [e[1] for e in second_forks] == ["lane0", "lane1"] + assert first_forks[0][2] != second_forks[0][2] + + +def test_ffn_rows_and_expert_output_stay_in_the_lane_pool(monkeypatch): + """Under DBO every FFN allocation of a stage (its dispatch inputs and the + expert output) is made on that stage's lane, and nothing is kept alive on + the caller stream. + + Why: rows allocated on a lane and recorded on the caller stream (or an + output allocated on the caller stream and recorded on the lane) cannot be + reused within one CUDA-graph capture; every MoE layer's rows stayed + allocated, and Qwen3-235B graph+DBO ran out of FFN memory at startup. + """ + connector, log = _logged_ffn_connector(monkeypatch, num_lanes=2) + + _ffn_dbo_step(connector, num_layers=3, num_stages=2) + + assert {e[1] for e in log if e[0] == "keep_alive"} <= {"lane0", "lane1"} + assert {e[1] for e in log if e[0] == "alloc"} == {"lane0", "lane1"} + # The expert output is the last allocation before each experts launch. + output_lanes = [] + for i, entry in enumerate(log): + if entry[0] == "experts": + output_lanes.append(next(e for e in reversed(log[:i]) if e[0] == "alloc")) + assert [e[1] for e in output_lanes] == ["lane0", "lane1"] * 3 + + +def test_single_lane_ffn_has_no_cross_stream_waits(monkeypatch): + """One lane runs on the caller stream: no events, no waits. + + Why: the single-lane path must keep its zero fork/join cost. + """ + connector, log = _logged_ffn_connector(monkeypatch, num_lanes=1) + + _ffn_dbo_step(connector, num_layers=2, num_stages=1) + + assert not [e for e in log if e[0] == "wait_stream"] + assert not [e for e in log if e[0] == "wait_event" and e[1] != "main"] + # One ready event per receive and one done event per send; no experts event. + assert len([e for e in log if e[0] == "record"]) == 2 + 2 + + +def test_keep_alive_records_stream_uses_while_capturing(monkeypatch): + """Cross-stream uses are recorded also during CUDA graph capture. + + Why: skipping them let the allocator reuse a freed graph-pool block on the + freeing stream while a DeepEP lane still read it on replay (a micro-batch's + hidden states under its dispatch); Qwen3-235B graph+DBO died with an + illegal instruction on the first concurrent requests. + """ + from afd_plugin.connectors.gpu.deepep_m2n import connector as m2n + + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True) + recorded = [] + + class _Tensor: + def record_stream(self, stream): + recorded.append(stream) + + m2n._keep_alive_on("lane0", _Tensor(), None, _Tensor()) + + assert recorded == ["lane0", "lane0"] diff --git a/tests/unit/connectors/test_deepep_m2n_network.py b/tests/unit/connectors/test_deepep_m2n_network.py new file mode 100644 index 000000000..7d0cc83d8 --- /dev/null +++ b/tests/unit/connectors/test_deepep_m2n_network.py @@ -0,0 +1,380 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""CPU tests for DeepEP M2N topology discovery and network planning.""" + +from __future__ import annotations + +import os + +import pytest + +pytest.importorskip("torch") +pytest.importorskip("vllm") + +from afd_plugin.connectors.gpu.deepep_m2n.network import ( # noqa: E402 + NetworkConfig, + NicInfo, + RankTopology, + apply_rank_environment, + discover_nics, + plan_network, +) + +RAIL_A = "roce:172.22.176.0/21" +RAIL_B = "roce:172.22.136.0/21" +RAIL_C = "roce:172.22.200.0/21" +IB_RAIL = "ib:fe80:0000:0000:0000" + + +def _rank(rank: int, host: str, local: int, nics=(), domain: str | None = None): + return RankTopology( + world_rank=rank, + host=host, + nvlink_domain=domain or f"host:{host}", + local_index=local, + nics=tuple(nics), + ) + + +def _nic(name: str, rail: str, distance: int) -> NicInfo: + return NicInfo(name=name, link_layer="ethernet", rail=rail, distance=distance) + + +def _isolated_rail_two_nodes() -> list[RankTopology]: + """Every NIC on its own rail; the nearest NIC differs per rank and host.""" + node0 = [ + [ + _nic("mlx5_0", RAIL_A, 2), + _nic("mlx5_5", RAIL_B, 6), + _nic("mlx5_3", RAIL_C, 6), + ], + [ + _nic("mlx5_0", RAIL_A, 6), + _nic("mlx5_5", RAIL_B, 6), + _nic("mlx5_3", RAIL_C, 2), + ], + ] + node1 = [ + [ + _nic("mlx5_0", RAIL_A, 2), + _nic("mlx5_5", RAIL_B, 6), + _nic("mlx5_3", RAIL_C, 6), + ], + [ + _nic("mlx5_0", RAIL_A, 6), + _nic("mlx5_5", RAIL_B, 2), + _nic("mlx5_3", RAIL_C, 6), + ], + ] + return [ + _rank(0, "n0", 0, node0[0]), + _rank(1, "n0", 1, node0[1]), + _rank(2, "n1", 0, node1[0]), + _rank(3, "n1", 1, node1[1]), + ] + + +def test_presets_expand_and_explicit_keys_override(): + """Presets fill in defaults, explicit keys override them, and the config + survives a round trip. + + Why: presets are defaults, not locks. + """ + rail_isolated = NetworkConfig.from_mapping({"preset": "roce-rail-isolated"}) + overridden = NetworkConfig.from_mapping({"preset": "ib", "nic_policy": "shared"}) + + assert (rail_isolated.nic_policy, rail_isolated.rails_routed) == ("shared", False) + assert (overridden.nic_policy, overridden.rails_routed) == ("shared", True) + assert NetworkConfig.from_mapping(rail_isolated.to_mapping()) == rail_isolated + + +@pytest.mark.parametrize( + ("raw", "message"), + [ + ({"gin": "maybe"}, "network.gin"), + ({"nic_policy": "explicit"}, "needs network.nics"), + ({"nccl_env": {"PATH": "/tmp"}}, "nccl_env key"), + ({"nccl_env": {"EP_DISABLE_GIN": "1"}}, "use network.gin"), + ({"bogus": 1}, "unknown network options"), + ], +) +def test_network_config_rejects_invalid_values(raw, message): + """Bad ``gin``, ``explicit`` without NICs, non-NCCL env keys, + ``EP_DISABLE_GIN`` in ``nccl_env`` and unknown keys are rejected. + + Why: bad or conflicting network settings must fail at parse time. + """ + with pytest.raises(ValueError, match=message): + NetworkConfig.from_mapping(raw) + + +def test_single_host_needs_no_rdma(): + """Four ranks on one host: one domain, GIN off, no NIC pins. + + Why: single-node runs must never touch RDMA. + """ + plan = plan_network( + [_rank(r, "n0", r, [_nic("mlx5_0", RAIL_A, 2)]) for r in range(4)], + NetworkConfig(), + ) + + assert (plan.num_domains, plan.use_gin) == (1, False) + assert plan.nic_by_rank == (None,) * 4 + + +def test_isolated_rails_put_every_rank_on_one_shared_rail(): + """Hosts whose NICs each sit on their own subnet get the shared policy, + one rail, and the same NIC for every rank. + + Why: encodes the Phase 4 finding as the automatic choice. + """ + plan = plan_network(_isolated_rail_two_nodes(), NetworkConfig()) + + assert (plan.use_gin, plan.nic_policy, plan.rail) == (True, "shared", RAIL_A) + assert plan.nic_by_rank == ("mlx5_0",) * 4 + + +def test_shared_rail_uses_each_hosts_own_nic_name(): + """The shared subnet is ``mlx5_0`` on one host and ``mlx5_7`` on the + other; each rank gets its own host's name. + + Why: NIC names are not consistent across hosts. + """ + ranks = [ + _rank(0, "n0", 0, [_nic("mlx5_0", RAIL_A, 2)]), + _rank(1, "n1", 0, [_nic("mlx5_7", RAIL_A, 2), _nic("mlx5_0", RAIL_B, 2)]), + ] + + plan = plan_network(ranks, NetworkConfig()) + + assert plan.nic_by_rank == ("mlx5_0", "mlx5_7") + + +def test_shared_rail_prefers_the_configured_nic(): + """``nics: [mlx5_5]`` selects that NIC's subnet as the shared rail. + + Why: operators can steer which NIC is shared. + """ + plan = plan_network( + _isolated_rail_two_nodes(), NetworkConfig(nic_policy="shared", nics=("mlx5_5",)) + ) + + assert (plan.rail, plan.nic_by_rank) == (RAIL_B, ("mlx5_5",) * 4) + + +@pytest.mark.parametrize( + "config", + [NetworkConfig(rails_routed=True), NetworkConfig.from_mapping({"preset": "ib"})], + ids=["routed-roce", "ib-preset"], +) +def test_routed_fabrics_keep_nccls_nearest_nic(config): + """``rails_routed`` and the ``ib`` preset keep NCCL's per-GPU NIC choice + (no pins) with GIN on. + + Why: routed fabrics should use every NIC, not throttle to one. + """ + plan = plan_network(_isolated_rail_two_nodes(), config) + + assert (plan.use_gin, plan.nic_policy) == (True, "nccl") + assert plan.nic_by_rank == (None,) * 4 + + +def test_ib_fat_tree_is_one_rail(): + """IB NICs on one subnet prefix get the NCCL policy without pins. + + Why: one rail needs no pinning. + """ + ib_nic = NicInfo(name="mlx5_0", link_layer="infiniband", rail=IB_RAIL, distance=2) + plan = plan_network( + [_rank(r, f"n{r // 2}", r % 2, [ib_nic]) for r in range(4)], NetworkConfig() + ) + + assert (plan.use_gin, plan.nic_policy) == (True, "nccl") + + +def test_nvl72_is_one_domain_even_with_uneven_hosts(): + # One NVLink fabric spans the hosts, so per-host counts need not match. + """Four ranks on two trays (3+1) in one NVLink fabric: one domain of 4, + GIN off. + + Why: NVL72 is one domain; per-host layout rules must not apply. + """ + ranks = [ + _rank(r, f"tray{0 if r < 3 else 1}", r, domain="nvl:abc:1") for r in range(4) + ] + + plan = plan_network(ranks, NetworkConfig()) + + assert (plan.num_domains, plan.ranks_per_domain, plan.use_gin) == (1, 4, False) + + +def test_explicit_nics_follow_local_rank(): + """``explicit`` assigns ``nics`` by local rank on every host. + + Why: a manual override for fabrics the planner cannot model. + """ + plan = plan_network( + _isolated_rail_two_nodes(), + NetworkConfig(nic_policy="explicit", nics=("mlx5_0", "mlx5_5")), + ) + + assert plan.nic_by_rank == ("mlx5_0", "mlx5_5", "mlx5_0", "mlx5_5") + + +def test_no_common_rail_is_an_error(): + """Hosts with no subnet in common raise "no RDMA rail". + + Why: fail at init instead of a QP timeout minutes later. + """ + ranks = [ + _rank(0, "n0", 0, [_nic("mlx5_0", RAIL_A, 2)]), + _rank(1, "n1", 0, [_nic("mlx5_0", RAIL_B, 2)]), + ] + + with pytest.raises(RuntimeError, match="no RDMA rail"): + plan_network(ranks, NetworkConfig()) + + +@pytest.mark.parametrize( + "config", + [NetworkConfig(gin="off"), NetworkConfig.from_mapping({"preset": "single-node"})], + ids=["gin-off", "single-node-preset"], +) +def test_gin_off_across_nvlink_domains_is_an_error(config): + # NVLink load/store cannot reach another domain; GIN is the only path. + """``gin: off`` (or the ``single-node`` preset) with ranks on two hosts + raises. + + Why: without GIN DeepEP cannot reach another NVLink domain. + """ + with pytest.raises(RuntimeError, match="span 2 NVLink domains"): + plan_network(_isolated_rail_two_nodes(), config) + + +def test_gin_off_on_one_host_is_allowed(): + """``gin: off`` with all ranks on one host plans normally. + + Why: the multi-domain check must not block single-node use. + """ + plan = plan_network( + [_rank(r, "n0", r, [_nic("mlx5_0", RAIL_A, 2)]) for r in range(2)], + NetworkConfig(gin="off"), + ) + + assert not plan.use_gin + + +def test_forced_gin_on_one_domain_still_plans_nics(): + """``gin: on`` on one host keeps GIN on and pins no NICs. + + Why: a forced setting is honored without inventing NIC pins. + """ + plan = plan_network( + [_rank(r, "n0", r, [_nic("mlx5_0", RAIL_A, 2)]) for r in range(2)], + NetworkConfig(gin="on"), + ) + + assert plan.use_gin + assert plan.nic_by_rank == (None, None) + + +def test_rank_environment_pins_nic_and_gin(monkeypatch): + """Applying a plan sets ``NCCL_IB_HCA``, ``nccl_env`` values and + ``EP_DISABLE_GIN`` (absent with GIN, 1 without). + + Why: the plan only works if these reach NCCL and DeepEP. + """ + for key in ("NCCL_IB_HCA", "EP_DISABLE_GIN", "NCCL_IB_TC"): + monkeypatch.delenv(key, raising=False) + config = NetworkConfig(nccl_env=(("NCCL_IB_TC", "106"),)) + plan = plan_network(_isolated_rail_two_nodes(), config) + + apply_rank_environment(plan, config, world_rank=3) + + assert os.environ["NCCL_IB_HCA"] == "=mlx5_0" + assert os.environ["NCCL_IB_TC"] == "106" + assert "EP_DISABLE_GIN" not in os.environ + + single = plan_network([_rank(0, "n0", 0)], NetworkConfig()) + apply_rank_environment(single, NetworkConfig(), world_rank=0) + assert os.environ["EP_DISABLE_GIN"] == "1" + + +def _fake_sysfs(tmp_path): + """Two RoCE ports (near and far from the GPU) and one IB port.""" + pci_root = tmp_path / "devices" + near_bridge = pci_root / "pci0000:00" / "0000:00:01.0" + far_bridge = pci_root / "pci0000:80" / "0000:80:01.0" + functions = { + "0000:18:00.0": near_bridge / "0000:18:00.0", + "0000:19:00.0": near_bridge / "0000:19:00.0", + "0000:9b:00.0": far_bridge / "0000:9b:00.0", + "0000:aa:00.0": far_bridge / "0000:aa:00.0", + } + sysfs_pci = tmp_path / "bus_pci" + sysfs_pci.mkdir() + for bdf, path in functions.items(): + path.mkdir(parents=True) + (sysfs_pci / bdf).symlink_to(path) + sysfs_ib = tmp_path / "infiniband" + ports = { + "mlx5_0": ( + "0000:19:00.0", + "Ethernet", + "0000:0000:0000:0000:0000:ffff:ac16:b021", + ), + "mlx5_10": ( + "0000:9b:00.0", + "Ethernet", + "0000:0000:0000:0000:0000:ffff:ac16:a021", + ), + "mlx5_12": ( + "0000:aa:00.0", + "InfiniBand", + "fe80:0000:0000:0000:0e42:a1ff:fe00:0001", + ), + } + for name, (bdf, link_layer, gid) in ports.items(): + port = sysfs_ib / name / "ports" / "1" + (port / "gids").mkdir(parents=True) + (port / "state").write_text("4: ACTIVE\n") + (port / "link_layer").write_text(f"{link_layer}\n") + (port / "gids" / "0").write_text( + "fe80:0000:0000:0000:0000:0000:0000:0000\n" + if link_layer == "Ethernet" + else gid + ) + if link_layer == "Ethernet": + (port / "gids" / "3").write_text(f"{gid}\n") + (sysfs_ib / name / "device").symlink_to(functions[bdf]) + down = sysfs_ib / "mlx5_6" / "ports" / "1" + (down / "gids").mkdir(parents=True) + (down / "state").write_text("1: DOWN\n") + return sysfs_ib, sysfs_pci + + +def test_discovery_reads_rails_and_pcie_distance(tmp_path): + """A fake sysfs tree: active RoCE and IB ports found (down port skipped), + rails from interface subnets, GID prefix or /32, and PCIe distance + ordering near before far. + + Why: topology discovery is tested without RDMA hardware. + """ + sysfs_ib, sysfs_pci = _fake_sysfs(tmp_path) + + nics = { + nic.name: nic + for nic in discover_nics( + "0000:18:00.0", + sysfs_infiniband=sysfs_ib, + sysfs_pci=sysfs_pci, + interface_networks={"172.22.176.33": "172.22.176.0/21"}, + ) + } + + assert set(nics) == {"mlx5_0", "mlx5_10", "mlx5_12"} + assert nics["mlx5_0"].rail == "roce:172.22.176.0/21" + # No netmask known for this address: it is its own (isolated) rail. + assert nics["mlx5_10"].rail == "roce:172.22.160.33/32" + assert nics["mlx5_12"].rail == IB_RAIL + assert nics["mlx5_0"].distance < nics["mlx5_10"].distance diff --git a/tests/unit/distributed/test_topology.py b/tests/unit/distributed/test_topology.py index 9970282a0..0e6b7aa55 100644 --- a/tests/unit/distributed/test_topology.py +++ b/tests/unit/distributed/test_topology.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project + from __future__ import annotations import sys @@ -6,7 +9,7 @@ import pytest from afd_plugin.config import AFDConfig -from afd_plugin.distributed import resolve_role_rank +from afd_plugin.distributed import resolve_role_rank, tp_rank_before_distributed_init def _vllm_config( @@ -50,10 +53,10 @@ def _install_parallel_state( distributed_module = ModuleType("vllm.distributed") distributed_module.__path__ = [] parallel_state_module = ModuleType("vllm.distributed.parallel_state") - parallel_state_module.get_pcp_group = lambda: SimpleNamespace( - rank_in_group=pcp_rank, + vars(parallel_state_module).update( + get_pcp_group=lambda: SimpleNamespace(rank_in_group=pcp_rank), + get_tensor_model_parallel_rank=lambda: tp_rank, ) - parallel_state_module.get_tensor_model_parallel_rank = lambda: tp_rank monkeypatch.setitem(sys.modules, "vllm", vllm_module) monkeypatch.setitem(sys.modules, "vllm.distributed", distributed_module) monkeypatch.setitem( @@ -107,3 +110,65 @@ def test_resolve_role_rank_rejects_derived_rank_outside_role_size(monkeypatch): _vllm_config(dp_size=2, dp_rank=1, tp_size=2), _afd_config(num_attention_ranks=2), ) + + +def test_resolve_role_rank_uses_given_tp_rank_without_vllm_groups(monkeypatch): + """A TP rank passed by the caller replaces vLLM's TP group lookup. + + Why: routed-experts connectors resolve their role rank before vLLM's + distributed init, when the TP group does not exist yet. + """ + monkeypatch.setitem(sys.modules, "vllm.distributed.parallel_state", None) + + role_rank = resolve_role_rank( + _vllm_config(dp_size=2, dp_rank=1, tp_size=2), + _afd_config(num_attention_ranks=4), + tp_rank=1, + ) + + assert role_rank == 3 + + +def test_resolve_role_rank_rejects_given_tp_rank_outside_tp_size(): + with pytest.raises(ValueError, match="tp_rank must be in"): + resolve_role_rank( + _vllm_config(tp_size=2), + _afd_config(num_attention_ranks=2), + tp_rank=2, + ) + + +def _pre_init_config(*, tp_size: int, pp_size: int = 1, pcp_size: int = 1): + return SimpleNamespace( + parallel_config=SimpleNamespace( + tensor_parallel_size=tp_size, + pipeline_parallel_size=pp_size, + prefill_context_parallel_size=pcp_size, + ), + ) + + +@pytest.mark.parametrize( + ("worker_rank", "tp_size", "expected"), + [(0, 1, 0), (3, 1, 0), (0, 2, 0), (1, 2, 1), (5, 4, 1)], +) +def test_tp_rank_before_distributed_init_is_worker_rank_mod_tp( + worker_rank, tp_size, expected +): + """TP is vLLM's innermost rank dimension, so the worker rank gives it. + + Why: this holds for DP-local and global worker ranks alike; + ``check_connector_role_rank`` re-checks it after vLLM's init. + """ + assert ( + tp_rank_before_distributed_init(_pre_init_config(tp_size=tp_size), worker_rank) + == expected + ) + + +@pytest.mark.parametrize(("pp_size", "pcp_size"), [(2, 1), (1, 2)]) +def test_tp_rank_before_distributed_init_rejects_pp_and_pcp(pp_size, pcp_size): + with pytest.raises(ValueError, match="PP = PCP = 1"): + tp_rank_before_distributed_init( + _pre_init_config(tp_size=2, pp_size=pp_size, pcp_size=pcp_size), 1 + ) diff --git a/tests/unit/fakes/afdep.py b/tests/unit/fakes/afdep.py new file mode 100644 index 000000000..23697e6f0 --- /dev/null +++ b/tests/unit/fakes/afdep.py @@ -0,0 +1,35 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""A stand-in for the afdEP package where it is not installed. + +``DeepEPM2NAFDConnector`` imports ``afdep.deepep`` when its module loads. The +stand-in provides the names it imports, so connector tests run on hosts +without afdEP; tests that build a lane substitute ``FakeElasticBuffer``. +""" + +from __future__ import annotations + +import sys +from importlib.util import find_spec +from types import ModuleType + + +class FakeElasticBuffer: + """Records the arguments one lane's transport is built with.""" + + def __init__(self, **kwargs: object) -> None: + self.kwargs = kwargs + + +def install_fake_afdep() -> None: + """Register the stand-in ``afdep.deepep`` unless afdEP is importable.""" + if "afdep" in sys.modules or find_spec("afdep") is not None: + return + package = ModuleType("afdep") + deepep = ModuleType("afdep.deepep") + deepep.ElasticBuffer = FakeElasticBuffer # type: ignore[attr-defined] + deepep.DispatchHandle = type("DispatchHandle", (), {}) # type: ignore[attr-defined] + deepep.ReceiveResult = type("ReceiveResult", (), {}) # type: ignore[attr-defined] + package.deepep = deepep # type: ignore[attr-defined] + sys.modules["afdep"] = package + sys.modules["afdep.deepep"] = deepep diff --git a/tests/unit/fakes/connector.py b/tests/unit/fakes/connector.py new file mode 100644 index 000000000..92b9c8813 --- /dev/null +++ b/tests/unit/fakes/connector.py @@ -0,0 +1,82 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""A recording ``AFDRoutedExpertsConnector`` for worker and runner tests. + +Constructed without ``AFDConnectorBase.__init__`` (no vLLM config needed). +``events`` records the lifecycle and data-path calls in order. +""" + +from __future__ import annotations + +from typing import Any + +from afd_plugin.connectors import ( + AFDA2FTransferPayload, + AFDRoutedExpertsConnector, + AFDTransferContext, + AFDTransferMetadata, +) + + +class FakeRoutedConnector(AFDRoutedExpertsConnector): + def __init__(self) -> None: + self.events: list[tuple] = [] + self.control_updates: list[Any] = [] + self.init_kwargs: dict[str, Any] | None = None + self.initialized = False + self.control_plane = self # type: ignore[assignment] + + @classmethod + def parse_extra_config(cls, raw): + return None + + @property + def is_initialized(self) -> bool: + return self.initialized + + def close(self) -> None: + self.events.append(("close",)) + + def bootstrap(self) -> None: + self.events.append(("bootstrap",)) + + def init_afd_connector(self, *, moe_layer_indices=None, local_experts=None): + self.init_kwargs = { + "moe_layer_indices": moe_layer_indices, + "local_experts": local_experts, + } + self.initialized = True + + # Control plane: the fake serves as its own. + def update_state_from_dp_metadata(self, payload) -> None: + self.control_updates.append(payload) + + def send_attn_output(self, hidden_states, context, **kwargs) -> None: + self.events.append(("send_attn",)) + + def recv_ffn_output(self, ref_tensor, ubatch_idx=0, **kwargs): + return ref_tensor + + def recv_attn_output(self, ubatch_idx=0, **kwargs): + assert not kwargs, "the routed path passes no P2P receive options" + self.events.append(("recv", ubatch_idx)) + return AFDA2FTransferPayload( + hidden_states=f"rows{ubatch_idx}", + context=AFDTransferContext( + metadata=AFDTransferMetadata.create_ffn_metadata( + layer_idx=0, + stage_idx=ubatch_idx, + seq_lens=[1], + ), + ), + ) + + def run_experts(self, payload, weights): + self.events.append(("run", payload.hidden_states, weights)) + return f"out({payload.hidden_states})" + + def send_ffn_output(self, ffn_output, context, **kwargs) -> None: + self.events.append(("send", ffn_output, context.metadata.stage_idx)) + + def synchronize_lanes(self) -> None: + self.events.append(("sync",)) diff --git a/tests/unit/fakes/vllm_config.py b/tests/unit/fakes/vllm_config.py index 1ef730fc5..0f0f467f7 100644 --- a/tests/unit/fakes/vllm_config.py +++ b/tests/unit/fakes/vllm_config.py @@ -1,7 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +from dataclasses import dataclass from types import SimpleNamespace +from typing import Any def camp2p_config( @@ -132,3 +134,115 @@ def async_moe_config( **extra_config, }, ) + + +# DeepEP M2N connector tests: Qwen3-30B-A3B MoE shapes. +M2N_HIDDEN_SIZE = 2048 +M2N_NUM_EXPERTS = 128 +M2N_TOP_K = 8 +M2N_INTERMEDIATE_SIZE = 768 +# Fake vLLM limits: capacity defaults to the larger of the two. +M2N_MAX_NUM_BATCHED_TOKENS = 8192 +M2N_MAX_CUDAGRAPH_CAPTURE_SIZE = 512 +# Fake CUDA-graph capture sizes: 1, 2, 4, then every 8 up to the maximum above. +M2N_CUDAGRAPH_CAPTURE_SIZES = ( + 1, + 2, + 4, + *range(8, M2N_MAX_CUDAGRAPH_CAPTURE_SIZE + 1, 8), +) + + +@dataclass(frozen=True) +class M2NTopology: + ffn_size: int + attn_size: int + # Attention TP peers (consecutive Attention ranks) share one DP rank's + # tokens; each dispatches its slice. + attn_tp_size: int = 1 + + @property + def world_size(self) -> int: + return self.ffn_size + self.attn_size + + @property + def attn_dp_size(self) -> int: + return self.attn_size // self.attn_tp_size + + def role(self, world_rank: int) -> tuple[str, int]: + if world_rank < self.ffn_size: + return "ffn", world_rank + return "attention", world_rank - self.ffn_size + + +def m2n_vllm_config( + topology: M2NTopology, + role: str, + role_rank: int, + extra_config: dict[str, Any], + port: int, + host: str = "127.0.0.1", +) -> SimpleNamespace: + """A ``vllm_config`` for one DeepEP M2N rank of ``topology``.""" + # Deferred: this fakes module must import on CPU-only hosts without torch. + import torch + from vllm.model_executor.layers.quantization.fp8 import Fp8Config + + text_config = SimpleNamespace( + hidden_size=M2N_HIDDEN_SIZE, + num_hidden_layers=2, + num_experts=M2N_NUM_EXPERTS, + num_experts_per_tok=M2N_TOP_K, + moe_intermediate_size=M2N_INTERMEDIATE_SIZE, + ) + # vLLM's view: FFN ranks are DP ranks; Attention ranks are DP x TP. + tp_size = 1 if role == "ffn" else topology.attn_tp_size + dp_size = topology.ffn_size if role == "ffn" else topology.attn_dp_size + # The FFN deployment the handshake requires: linear EP over the FFN ranks, + # with checkpoint weights in the payload's format. + quant_config = ( + Fp8Config(is_checkpoint_fp8_serialized=True, weight_block_size=[128, 128]) + if extra_config.get("payload_format") == "fp8_block128_fp32" + else None + ) + return SimpleNamespace( + additional_config={ + "afd": { + "connector": "DeepEPM2NAFDConnector", + "role": role, + "host": host, + "port": port, + "num_attention_ranks": topology.attn_size, + "num_ffn_ranks": topology.ffn_size, + "compute_gate_on_attention": True, + "connector_extra_config": extra_config, + } + }, + quant_config=quant_config, + compilation_config=SimpleNamespace( + max_cudagraph_capture_size=M2N_MAX_CUDAGRAPH_CAPTURE_SIZE, + cudagraph_capture_sizes=list(M2N_CUDAGRAPH_CAPTURE_SIZES), + ), + model_config=SimpleNamespace( + dtype=torch.bfloat16, + enforce_eager=True, + hf_config=text_config, + hf_text_config=text_config, + ), + parallel_config=SimpleNamespace( + data_parallel_size=dp_size, + data_parallel_rank=role_rank // tp_size, + pipeline_parallel_size=1, + prefill_context_parallel_size=1, + tensor_parallel_size=tp_size, + enable_expert_parallel=True, + expert_placement_strategy="linear", + enable_eplb=False, + all2all_backend="allgather_reducescatter", + use_ubatching=False, + num_ubatches=1, + ), + scheduler_config=SimpleNamespace( + max_num_batched_tokens=M2N_MAX_NUM_BATCHED_TOKENS + ), + ) 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 91210d1cc..6ede9f77b 100644 --- a/tests/unit/model_executor/models/test_deepseek_v2_proxy.py +++ b/tests/unit/model_executor/models/test_deepseek_v2_proxy.py @@ -10,10 +10,17 @@ torch = pytest.importorskip("torch") pytest.importorskip("vllm") from torch import nn # noqa: E402 +from vllm.model_executor.layers.fused_moe.router import ( # noqa: E402 + fused_topk_router, +) from afd_plugin.config import AFD_ASYNC_CONNECTOR, AFDConfig # noqa: E402 from afd_plugin.model_executor.models import deepseek_v2 as adapter # noqa: E402 +ROUTING_HIDDEN_SIZE = 16 +ROUTING_NUM_EXPERTS = 8 +ROUTING_TOP_K = 2 + class _FakeConnector: def __init__(self, events: list[tuple]) -> None: @@ -360,3 +367,288 @@ def forward(self, hidden_states): torch.testing.assert_close(actual_ids, expected_ids.to(torch.int32)) torch.testing.assert_close(actual_weights, weights) assert actual_logits is logits + + +class _RecordingGate(nn.Module): + """Linear gate returning ``(logits, bias)`` like vLLM's ReplicatedLinear.""" + + def __init__(self, calls: list) -> None: + super().__init__() + self.calls = calls + self.weight = nn.Parameter( + torch.randn(ROUTING_NUM_EXPERTS, ROUTING_HIDDEN_SIZE) + ) + + def forward(self, hidden_states): + self.calls.append(hidden_states) + return hidden_states @ self.weight.t(), None + + +def _routing_proxy(gate: nn.Module, *, renormalize: bool = True): + """A proxy that routes on Attention, configured as a model sets it up.""" + proxy = adapter.RemoteFFNProxy(layer_idx=3) + proxy.gate = gate + proxy.top_k = ROUTING_TOP_K + proxy.renormalize = renormalize + return proxy + + +def test_default_proxy_is_parameter_free_and_sends_no_routes(monkeypatch): + """Without a gate the proxy sends only hidden states. + + Why: connectors other than M2N run the whole MoE block on FFN; adding the + routing settings to the base proxy must not change their payload. + """ + events: list[tuple] = [] + _install_fake_forward_context(monkeypatch, events) + proxy = adapter.RemoteFFNProxy(layer_idx=0) + + proxy(torch.ones(2, 4)) + + assert list(proxy.parameters()) == [] + assert proxy.tp_size == 1 + assert events[0][3] == {} + + +def test_routing_proxy_routes_on_attention_and_sends_topk(monkeypatch): + """Gate and top-k run once on Attention; only the routes go to the connector. + + Why: DeepEP M2N dispatches by topk_ids/topk_weights and its combine returns + the weighted sum, so the proxy must not send logits or apply weights. + """ + events: list[tuple] = [] + _install_fake_forward_context(monkeypatch, events, stage_idx=1) + gate_calls: list = [] + topk_calls: list[dict] = [] + topk_weights = torch.tensor([[0.75, 0.25]]) + topk_ids = torch.tensor([[1, 3]], dtype=torch.int32) + + def fake_fused_topk(**kwargs): + topk_calls.append(kwargs) + return topk_weights, topk_ids, torch.zeros_like(topk_ids) + + monkeypatch.setattr(adapter, "fused_topk", fake_fused_topk) + proxy = _routing_proxy(_RecordingGate(gate_calls), renormalize=False) + hidden_states = torch.ones(1, ROUTING_HIDDEN_SIZE) + + output = proxy(hidden_states) + + assert gate_calls == [hidden_states] + assert len(topk_calls) == 1 + call = topk_calls[0] + assert call["hidden_states"] is hidden_states + assert torch.equal(call["gating_output"], hidden_states @ proxy.gate.weight.t()) + assert call["topk"] == ROUTING_TOP_K + assert call["renormalize"] is False + assert [event[0] for event in events] == ["send", "yield", "recv"] + context = events[0][2] + assert context.metadata.layer_idx == 3 + assert context.metadata.stage_idx == 1 + assert events[0][3] == {"topk_ids": topk_ids, "topk_weights": topk_weights} + assert events[2][2] == 1 + assert torch.equal(output, hidden_states * 0.25) + + +def test_routing_proxy_owns_gate_at_native_checkpoint_path(): + """The gate is the proxy's only parameter, under ``gate``. + + Why: the proxy replaces ``mlp``, so ``model.layers.N.mlp.gate.weight`` + from the checkpoint loads into it with no name mapping. + """ + proxy = _routing_proxy(_RecordingGate([])) + + assert set(dict(proxy.named_parameters())) == {"gate.weight"} + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="fused top-k needs CUDA") +@pytest.mark.parametrize("renormalize", [True, False]) +def test_routing_proxy_matches_native_router(monkeypatch, renormalize): + """On GPU the proxy's routes equal vLLM's native fused top-k router. + + Why: the FFN applies the weights the Attention side computes; any drift + from native routing changes the model's output. + """ + torch.manual_seed(0) + events: list[tuple] = [] + _install_fake_forward_context(monkeypatch, events, stage_idx=0) + gate = _RecordingGate([]).cuda().to(torch.bfloat16) + proxy = _routing_proxy(gate, renormalize=renormalize) + hidden_states = torch.randn( + 33, ROUTING_HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16 + ) + + proxy(hidden_states) + + sent = events[0][3] + router_logits, _ = gate(hidden_states) + native_weights, native_ids = fused_topk_router.FusedTopKRouter( + top_k=ROUTING_TOP_K, + global_num_experts=ROUTING_NUM_EXPERTS, + renormalize=renormalize, + )._compute_routing(hidden_states, router_logits, indices_type=None) + assert torch.equal(sent["topk_ids"], native_ids) + assert torch.equal(sent["topk_weights"], native_weights) + + reference = torch.softmax(router_logits.float(), dim=-1) + reference_weights, reference_ids = reference.topk(ROUTING_TOP_K, dim=-1) + if renormalize: + reference_weights = reference_weights / reference_weights.sum(-1, keepdim=True) + assert torch.equal( + sent["topk_ids"].long().sort(-1).values, reference_ids.sort(-1).values + ) + torch.testing.assert_close( + sent["topk_weights"].sort(-1).values, + reference_weights.sort(-1).values, + rtol=1e-3, + atol=1e-3, + ) + + +# Fake combine output scale of ``_FakeConnector.recv_ffn_output``. +COMBINE_SCALE = 0.25 +# Added to the peer rank's rows by the fake all-gather, to tell them apart. +PEER_OFFSET = 100.0 + + +class _TokenGate(nn.Module): + """Logits that make expert ``token % ROUTING_NUM_EXPERTS`` win per row.""" + + def forward(self, hidden_states): + logits = torch.zeros(hidden_states.shape[0], ROUTING_NUM_EXPERTS) + experts = hidden_states[:, 0].long() % ROUTING_NUM_EXPERTS + logits[torch.arange(hidden_states.shape[0]), experts] = 1.0 + return logits, None + + +def _torch_topk(*, hidden_states, gating_output, topk, renormalize): + weights, ids = torch.topk(gating_output.softmax(dim=-1), topk, dim=-1) + return weights, ids.to(torch.int32), None + + +def _chunk_for_rank(tp_rank: int, tp_size: int): + """vLLM's ``sequence_parallel_chunk`` for one rank: pad, then slice.""" + + def chunk(x): + padded = torch.nn.functional.pad(x, (0, 0, 0, (-x.shape[0]) % tp_size)) + rows = padded.shape[0] // tp_size + return padded[tp_rank * rows : (tp_rank + 1) * rows].clone() + + return chunk + + +def _tokens(num_tokens: int): + """Token ``i`` has value ``i`` in every channel.""" + return torch.arange(num_tokens, dtype=torch.float32)[:, None].repeat( + 1, ROUTING_HIDDEN_SIZE + ) + + +def _tp_routing_proxy(monkeypatch, tp_size: int): + monkeypatch.setattr(adapter, "fused_topk", _torch_topk) + proxy = _routing_proxy(_TokenGate()) + proxy.tp_size = tp_size + return proxy + + +def test_routing_proxy_tp1_dispatches_every_token(monkeypatch): + """Without Attention TP the proxy dispatches the whole batch unsplit. + + Why: TP = 1 must keep the existing data path. + """ + + def no_split(x): + raise AssertionError("TP = 1 must not split") + + events: list[tuple] = [] + _install_fake_forward_context(monkeypatch, events) + monkeypatch.setattr(adapter.native, "sequence_parallel_chunk", no_split) + hidden_states = _tokens(3) + + output = _tp_routing_proxy(monkeypatch, tp_size=1)(hidden_states) + + sends = [event for event in events if event[0] == "send"] + assert len(sends) == 1 + assert torch.equal(sends[0][1], hidden_states) + assert torch.equal(output, hidden_states * COMBINE_SCALE) + + +@pytest.mark.parametrize( + ("tp_rank", "sent_tokens", "padding_rows"), + [(0, [0, 1], []), (1, [2, 0], [1])], +) +def test_routing_proxy_tp2_rank_dispatches_its_slice_and_gathers( + monkeypatch, tp_rank, sent_tokens, padding_rows +): + """With TP 2 and 3 tokens, rank 0 dispatches tokens 0-1 and rank 1 token 2 + plus one padding row whose experts are all ``UNROUTED_EXPERT_ID``; the + combined slices are all-gathered and trimmed back to 3 tokens. + + Why: TP peers hold the same tokens; each must dispatch a disjoint slice + so every token reaches its experts exactly once. + """ + tp_size = 2 + events: list[tuple] = [] + _install_fake_forward_context(monkeypatch, events) + monkeypatch.setattr( + adapter.native, "sequence_parallel_chunk", _chunk_for_rank(tp_rank, tp_size) + ) + monkeypatch.setattr( + adapter.native, "get_tensor_model_parallel_rank", lambda: tp_rank + ) + gathered_inputs = [] + + def all_gather(local_output, dim): + gathered_inputs.append(local_output.clone()) + return torch.cat([local_output, local_output + PEER_OFFSET], dim=dim) + + monkeypatch.setattr(adapter.native, "tensor_model_parallel_all_gather", all_gather) + + output = _tp_routing_proxy(monkeypatch, tp_size)(_tokens(3)) + + (send,) = [event for event in events if event[0] == "send"] + sent_states, sent_kwargs = send[1], send[3] + assert sent_states[:, 0].tolist() == sent_tokens + topk_ids = sent_kwargs["topk_ids"] + unrouted = (topk_ids == adapter.UNROUTED_EXPERT_ID).all(dim=1) + assert unrouted.nonzero().flatten().tolist() == padding_rows + assert not (topk_ids[~unrouted] == adapter.UNROUTED_EXPERT_ID).any() + (local_output,) = gathered_inputs + assert torch.equal(local_output, sent_states * COMBINE_SCALE) + expected = torch.cat([local_output, local_output + PEER_OFFSET])[:3] + assert torch.equal(output, expected) + + +def test_routing_proxy_tp4_with_one_token_routes_only_rank0(monkeypatch): + """One decode token over TP 4: rank 0 routes it, ranks 1-3 send one + unrouted padding row each. + + Why: decode batches can be smaller than the TP size; padding must never + reach an expert. + """ + tp_size = 4 + events: list[tuple] = [] + _install_fake_forward_context(monkeypatch, events) + monkeypatch.setattr( + adapter.native, + "tensor_model_parallel_all_gather", + lambda local, dim: torch.cat([local] * tp_size, dim=dim), + ) + for tp_rank in range(tp_size): + monkeypatch.setattr( + adapter.native, "sequence_parallel_chunk", _chunk_for_rank(tp_rank, tp_size) + ) + monkeypatch.setattr( + adapter.native, + "get_tensor_model_parallel_rank", + lambda rank=tp_rank: rank, + ) + + output = _tp_routing_proxy(monkeypatch, tp_size)(_tokens(1)) + + assert output.shape == (1, ROUTING_HIDDEN_SIZE) + unrouted = [ + bool((event[3]["topk_ids"] == adapter.UNROUTED_EXPERT_ID).all()) + for event in events + if event[0] == "send" + ] + assert unrouted == [False, True, True, True] diff --git a/tests/unit/model_executor/models/test_deepseek_v2_weight_policy.py b/tests/unit/model_executor/models/test_deepseek_v2_weight_policy.py index c09812733..0b6d822d4 100644 --- a/tests/unit/model_executor/models/test_deepseek_v2_weight_policy.py +++ b/tests/unit/model_executor/models/test_deepseek_v2_weight_policy.py @@ -244,3 +244,30 @@ def test_async_shared_weights_belong_to_attention(suffix): ) == FFN_ROLES ) + + +@pytest.mark.parametrize( + ("name", "expected"), + [ + ("model.layers.3.mlp.shared_experts.gate_proj.weight", frozenset()), + ("model.layers.3.mlp.experts.0.down_proj.weight", FFN_ROLES), + ("model.layers.3.mlp.gate.weight", BOTH_ROLES), + ], + ids=["shared-experts", "routed-experts", "gate"], +) +def test_deepep_m2n_ffn_loads_routed_experts_only(name, expected): + """With ``ffn_skips_shared_experts`` the FFN loads no shared experts and + Attention does not take them either, until it builds them; routed + experts and the gate keep their gate-on-Attention owners. + + Why: the DeepEP M2N FFN runs routed experts only. + """ + assert ( + _checkpoint_weight_roles( + name, + _config(), + compute_gate_on_attention=True, + ffn_skips_shared_experts=True, + ) + == expected + ) diff --git a/tests/unit/model_executor/models/test_qwen3_moe_construction.py b/tests/unit/model_executor/models/test_qwen3_moe_construction.py index dee09f2e9..3b0a7964a 100644 --- a/tests/unit/model_executor/models/test_qwen3_moe_construction.py +++ b/tests/unit/model_executor/models/test_qwen3_moe_construction.py @@ -4,16 +4,22 @@ from __future__ import annotations from types import SimpleNamespace +from typing import TYPE_CHECKING import pytest -torch = pytest.importorskip("torch") +if TYPE_CHECKING: + import torch + from torch import nn +else: + torch = pytest.importorskip("torch") + nn = torch.nn pytest.importorskip("vllm") from torch import nn # noqa: E402 from vllm.compilation.wrapper import TorchCompileWithNoGuardsWrapper # noqa: E402 from vllm.config import CompilationMode # noqa: E402 -from afd_plugin.config import AFDConfig # noqa: E402 +from afd_plugin.config import DEEPEP_M2N_CONNECTOR, AFDConfig # noqa: E402 from afd_plugin.model_executor.models import qwen3_moe as adapter # noqa: E402 from tests.unit.fakes.model_layers import stage_type as _stage_type # noqa: E402 @@ -27,6 +33,7 @@ def construction_env(monkeypatch: pytest.MonkeyPatch): calls: dict[str, list[str]] = { "attention": [], "dense": [], + "gate": [], "moe": [], "norm": [], } @@ -50,6 +57,11 @@ def bind(stage_type): bind(_stage_type("moe")), ) monkeypatch.setattr(adapter.native, "RMSNorm", bind(_stage_type("norm"))) + monkeypatch.setattr( + adapter.native, + "ReplicatedLinear", + bind(_stage_type("gate")), + ) monkeypatch.setattr(adapter.native, "PPMissingLayer", _FakeMissingLayer) return calls @@ -69,6 +81,7 @@ def _model_config( intermediate_size=32, max_position_embeddings=4096, mlp_only_layers=mlp_only_layers or [], + norm_topk_prob=True, num_attention_heads=2, num_experts=8, num_experts_per_tok=2, @@ -86,6 +99,7 @@ def _model_config( parallel_config=SimpleNamespace( enable_eplb=False, pipeline_parallel_size=1, + tensor_parallel_size=1, use_sequence_parallel_moe=False, ), quant_config=None, @@ -103,11 +117,14 @@ def _make_layer( role: str, layer_idx: int, vllm_config=None, + afd_config: AFDConfig | None = None, ): + if afd_config is None: + afd_config = AFDConfig(role=role) monkeypatch.setattr( adapter, "parse_afd_config", - lambda *_args, **_kwargs: AFDConfig(role=role), + lambda *_args, **_kwargs: afd_config, ) if vllm_config is None: vllm_config = _model_config() @@ -127,7 +144,8 @@ def test_attention_role_constructs_only_attention_and_norms( assert construction_env["norm"] == ["", ""] assert construction_env["dense"] == [] assert construction_env["moe"] == [] - assert isinstance(layer.mlp, adapter.RemoteFFNProxy) + assert construction_env["gate"] == [] + assert type(layer.mlp) is adapter.RemoteFFNProxy assert layer.mlp.layer_idx == 0 assert set(dict(layer.named_parameters())) == { "self_attn.weight", @@ -136,6 +154,70 @@ def test_attention_role_constructs_only_attention_and_norms( } +def test_attention_role_routes_with_native_gate_under_deepep_m2n( + monkeypatch: pytest.MonkeyPatch, + construction_env, +) -> None: + """With DeepEP M2N, Attention sets the proxy's gate and top-k settings. + + Why: DeepEP M2N dispatches by top-k routes computed on Attention; the gate + keeps the native ``mlp.gate`` path so the checkpoint loads unchanged. + """ + layer = _make_layer( + monkeypatch, + role="attention", + layer_idx=5, + afd_config=AFDConfig( + role="attention", + connector=DEEPEP_M2N_CONNECTOR, + compute_gate_on_attention=True, + ), + ) + + assert construction_env["gate"] == ["model.layers.5.mlp.gate"] + assert construction_env["moe"] == [] + assert type(layer.mlp) is adapter.RemoteFFNProxy + assert layer.mlp.gate is not None + assert layer.mlp.layer_idx == 5 + assert layer.mlp.top_k == 2 + assert layer.mlp.renormalize is True + assert layer.mlp.tp_size == 1 + assert set(dict(layer.named_parameters())) == { + "self_attn.weight", + "mlp.gate.weight", + "input_layernorm.weight", + "post_attention_layernorm.weight", + } + + +def test_attention_role_gives_the_proxy_the_attention_tp_size( + monkeypatch: pytest.MonkeyPatch, + construction_env, +) -> None: + """Under DeepEP M2N the proxy gets the Attention TP size. + + Why: with Attention TP > 1 each TP rank must dispatch only its slice of + the tokens; without the TP size the proxy would dispatch every token from + every TP rank. + """ + vllm_config = _model_config() + vllm_config.parallel_config.tensor_parallel_size = 2 + layer = _make_layer( + monkeypatch, + role="attention", + layer_idx=5, + vllm_config=vllm_config, + afd_config=AFDConfig( + role="attention", + connector=DEEPEP_M2N_CONNECTOR, + compute_gate_on_attention=True, + ), + ) + + assert type(layer.mlp) is adapter.RemoteFFNProxy + assert layer.mlp.tp_size == 2 + + def test_ffn_role_constructs_only_native_moe( monkeypatch: pytest.MonkeyPatch, construction_env, @@ -314,6 +396,86 @@ def test_unsupported_modes_fail_closed( adapter._validate_supported_config(vllm_config, afd_config) +def test_gate_on_attention_is_accepted_with_deepep_m2n( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """DeepEP M2N routes on Attention, so Qwen3 accepts the gate there. + + Why: its FFN runs only the routed experts and never the native gate. + """ + monkeypatch.setattr(adapter, "current_platform", _platform(is_cuda=True)) + + adapter._validate_supported_config( + _model_config(), + AFDConfig(connector="DeepEPM2NAFDConnector", compute_gate_on_attention=True), + ) + + +def test_attention_role_selects_routing_proxy_by_connector_only( + monkeypatch: pytest.MonkeyPatch, + construction_env, +) -> None: + """Without DeepEP M2N, Attention builds the plain proxy even if the + gate-on-Attention flag is set. + + Why: only DeepEP M2N dispatches by Attention-side routes; for any other + connector the flag is rejected by model validation, and the layer itself + must not pick the routing proxy from the flag alone. + """ + layer = _make_layer( + monkeypatch, + role="attention", + layer_idx=0, + afd_config=AFDConfig(role="attention", compute_gate_on_attention=True), + ) + + assert type(layer.mlp) is adapter.RemoteFFNProxy + assert layer.mlp.gate is None + assert construction_env["gate"] == [] + + +def test_gate_on_attention_rejects_dense_layers( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Gate-on-Attention requires every layer to be MoE. + + Why: the routed FFN runner serves MoE layers only, so a dense layer's MLP + would run on neither role. + """ + monkeypatch.setattr(adapter, "current_platform", _platform(is_cuda=True)) + + with pytest.raises(RuntimeError, match="every layer to be MoE"): + adapter._validate_supported_config( + _model_config(layer_count=4, mlp_only_layers=[3]), + AFDConfig( + connector=DEEPEP_M2N_CONNECTOR, + compute_gate_on_attention=True, + ), + ) + + +def test_deepep_m2n_rejects_shared_experts( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """DeepEP M2N fails closed on a Qwen3 checkpoint with shared experts. + + Why: neither the Attention proxy nor the routed FFN runner computes a + shared expert, so its output would be dropped without an error. + """ + monkeypatch.setattr(adapter, "current_platform", _platform(is_cuda=True)) + vllm_config = _model_config() + vllm_config.model_config.hf_text_config.shared_expert_intermediate_size = 64 + + with pytest.raises(RuntimeError, match="does not support shared experts"): + adapter._validate_supported_config( + vllm_config, + AFDConfig( + connector=DEEPEP_M2N_CONNECTOR, + compute_gate_on_attention=True, + ), + ) + + def test_rocm_platform_fails_closed(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr( adapter, diff --git a/tests/unit/model_executor/models/test_qwen3_moe_weight_policy.py b/tests/unit/model_executor/models/test_qwen3_moe_weight_policy.py index f1222dc73..20938151d 100644 --- a/tests/unit/model_executor/models/test_qwen3_moe_weight_policy.py +++ b/tests/unit/model_executor/models/test_qwen3_moe_weight_policy.py @@ -14,6 +14,7 @@ torch = pytest.importorskip("torch") pytest.importorskip("vllm") +from afd_plugin.config import DEEPEP_M2N_CONNECTOR, AFDConfig # noqa: E402 from afd_plugin.model_executor.models import qwen3_moe as adapter # noqa: E402 from tests.unit.fakes.weights import OneShotWeights as _OneShotWeights # noqa: E402 @@ -73,6 +74,33 @@ ) +# With gate-on-Attention (DeepEP M2N) only the MoE gate changes owner. +GATE_ON_ATTENTION_ROLE_CASES = ( + ("moe-gate", "model.layers.0.mlp.gate.weight", BOTH_ROLES), + ( + "moe-expert-projection", + "model.layers.0.mlp.experts.0.gate_proj.weight", + FFN_ROLE, + ), + ( + "shared-expert-gate", + "model.layers.0.mlp.shared_expert_gate.weight", + FFN_ROLE, + ), + ( + "dense-gate-projection", + "model.layers.3.mlp.gate_proj.weight", + FFN_ROLE, + ), + ( + "attention-q-projection", + "model.layers.0.self_attn.q_proj.weight", + ATTENTION_ROLE, + ), + ("embedding", "model.embed_tokens.weight", BOTH_ROLES), +) + + @pytest.mark.parametrize( ("checkpoint_name", "roles"), [case[1:] for case in WEIGHT_ROLE_CASES], @@ -86,19 +114,65 @@ def test_weight_role_policy( @pytest.mark.parametrize( - ("role", "expected_names"), + ("checkpoint_name", "roles"), + [case[1:] for case in GATE_ON_ATTENTION_ROLE_CASES], + ids=[case[0] for case in GATE_ON_ATTENTION_ROLE_CASES], +) +def test_gate_on_attention_weight_role_policy( + checkpoint_name: str, + roles: frozenset[str], +) -> None: + """The MoE gate loads on both roles; everything else keeps its owner. + + Why: Attention routes with the gate, and the FFN's native MoE block still + owns a gate parameter that vLLM's strict loader requires to be loaded. + """ + assert ( + adapter._checkpoint_weight_roles( + checkpoint_name, + routes_on_attention=True, + ) + == roles + ) + + +@pytest.mark.parametrize( + ("role", "uses_deepep_m2n", "expected_names"), [ ( "attention", + False, + [ + "model.embed_tokens.weight", + "model.layers.0.self_attn.q_proj.weight", + "model.layers.0.input_layernorm.weight", + "lm_head.weight", + ], + ), + ( + "ffn", + False, + [ + "model.embed_tokens.weight", + "model.layers.0.mlp.gate.weight", + "model.layers.0.mlp.experts.0.down_proj.weight", + "lm_head.weight", + ], + ), + ( + "attention", + True, [ "model.embed_tokens.weight", "model.layers.0.self_attn.q_proj.weight", "model.layers.0.input_layernorm.weight", + "model.layers.0.mlp.gate.weight", "lm_head.weight", ], ), ( "ffn", + True, [ "model.embed_tokens.weight", "model.layers.0.mlp.gate.weight", @@ -111,6 +185,7 @@ def test_weight_role_policy( def test_load_weights_filters_once_and_delegates_to_native_loader( monkeypatch: pytest.MonkeyPatch, role: str, + uses_deepep_m2n: bool, expected_names: list[str], ) -> None: names = [ @@ -137,6 +212,18 @@ def fake_native_loader(self, filtered_weights): ) model = object.__new__(adapter.AFDQwen3MoeForCausalLM) object.__setattr__(model, "afd_role", role) + # The gate follows the connector: only DeepEP M2N routes on Attention. + object.__setattr__( + model, + "afd_config", + AFDConfig( + role=role, + connector=DEEPEP_M2N_CONNECTOR, + compute_gate_on_attention=True, + ) + if uses_deepep_m2n + else AFDConfig(role=role), + ) result = model.load_weights(weights) diff --git a/tests/unit/package/test_package.py b/tests/unit/package/test_package.py index 0bfebfadd..c745522ef 100644 --- a/tests/unit/package/test_package.py +++ b/tests/unit/package/test_package.py @@ -20,6 +20,7 @@ def test_gpu_runtime_exports_only_mrv2_attention_and_connector_ffn(): "AFDAttentionModelRunnerV2", "AFDFFNWorker", "GPUFFNModelRunner", + "GPUFFNRoutedExpertsModelRunner", } diff --git a/tests/unit/v1/worker/test_connector_device.py b/tests/unit/v1/worker/test_connector_device.py new file mode 100644 index 000000000..76ee793b9 --- /dev/null +++ b/tests/unit/v1/worker/test_connector_device.py @@ -0,0 +1,159 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""A connector created before ``Worker.init_device`` gets vLLM's device.""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest + +torch = pytest.importorskip("torch") +pytest.importorskip("vllm") + +from vllm.platforms import interface as platform_interface # noqa: E402 + +from afd_plugin.config import AFDConfig # noqa: E402 +from afd_plugin.v1.worker import connector_device # noqa: E402 + + +class _CudaLikePlatform(platform_interface.Platform): + device_control_env_var = "CUDA_VISIBLE_DEVICES" + + +@pytest.fixture(autouse=True) +def _cuda_like_platform(monkeypatch): + monkeypatch.setattr(connector_device, "current_platform", _CudaLikePlatform) + # vLLM keeps the assigned ids process-wide; start every case unset. + monkeypatch.setattr(platform_interface, "_assigned_physical_gpu_ids", None) + monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False) + + +def _config( + *, + data_parallel_rank_local=0, + distributed_executor_backend="mp", + assigned_physical_gpu_ids=None, +): + return SimpleNamespace( + parallel_config=SimpleNamespace( + distributed_executor_backend=distributed_executor_backend, + data_parallel_backend="mp", + nnodes_within_dp=1, + data_parallel_rank_local=data_parallel_rank_local, + data_parallel_index=data_parallel_rank_local, + pipeline_parallel_size=1, + tensor_parallel_size=1, + assigned_physical_gpu_ids=assigned_physical_gpu_ids, + ), + ) + + +@pytest.mark.parametrize( + ("executor_backend", "expected_index"), + [("mp", 1), ("ray", 0), ("external_launcher", 0)], +) +def test_adds_local_dp_offset_like_vllm(executor_backend, expected_index): + """The local DP rank offsets the device, except where vLLM skips it. + + Why: vLLM adds the offset only inside init_device, after the connector + bootstrap; without it every DP rank's connector used GPU 0. + """ + config = _config( + data_parallel_rank_local=1, + distributed_executor_backend=executor_backend, + ) + + assert connector_device.connector_device_index(config, 0) == expected_index + + +def test_maps_assigned_device_ids_without_visible_devices(): + """``--device-ids 4,5`` without CUDA_VISIBLE_DEVICES places rank 0 on cuda:4. + + Why: vLLM maps its logical index through the assigned physical ids, so a + connector using the logical index would run on another rank's GPU. + """ + config = _config(assigned_physical_gpu_ids=[4, 5]) + + assert connector_device.connector_device_index(config, 0) == 4 + assert connector_device.connector_device_index(config, 1) == 5 + + +def test_maps_assigned_device_ids_through_visible_devices(monkeypatch): + """With CUDA_VISIBLE_DEVICES the assigned ids map to visible ordinals.""" + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "4,5") + config = _config(assigned_physical_gpu_ids=[5, 4]) + + assert connector_device.connector_device_index(config, 0) == 1 + assert connector_device.connector_device_index(config, 1) == 0 + + +def test_without_assigned_ids_the_index_is_the_local_rank(monkeypatch): + """The default mp launch keeps the logical index as the device index.""" + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "2,3") + + assert connector_device.connector_device_index(_config(), 1) == 1 + + +def test_check_accepts_the_worker_device(): + connector_device.check_connector_device( + SimpleNamespace(local_rank=1), + torch.device("cuda", 1), + ) + + +def test_check_rejects_another_device(): + """A connector on another GPU than the worker's fails with both devices. + + Why: a vLLM upgrade that changes device selection must fail at startup + instead of moving traffic onto another rank's GPU. + """ + with pytest.raises(RuntimeError, match="uses cuda:0 but vLLM placed"): + connector_device.check_connector_device( + SimpleNamespace(local_rank=0), + torch.device("cuda", 4), + ) + + +def _role_rank_config(*, tp_size: int): + return SimpleNamespace( + parallel_config=SimpleNamespace( + data_parallel_size=1, + data_parallel_rank=0, + prefill_context_parallel_size=1, + tensor_parallel_size=tp_size, + ), + ) + + +def test_role_rank_check_accepts_the_distributed_role_rank(monkeypatch): + monkeypatch.setattr( + "vllm.distributed.parallel_state.get_tensor_model_parallel_rank", lambda: 1 + ) + connector = SimpleNamespace( + role_rank=1, + afd_config=AFDConfig(role="attention", num_attention_ranks=2), + ) + + connector_device.check_connector_role_rank(connector, _role_rank_config(tp_size=2)) + + +def test_role_rank_check_rejects_a_different_role_rank(monkeypatch): + """A connector whose early role rank differs from vLLM's fails startup. + + Why: the TP part of the role rank is derived from the worker rank before + vLLM's TP group exists; a vLLM change to its rank order must fail here + instead of routing another rank's tokens. + """ + monkeypatch.setattr( + "vllm.distributed.parallel_state.get_tensor_model_parallel_rank", lambda: 0 + ) + connector = SimpleNamespace( + role_rank=1, + afd_config=AFDConfig(role="attention", num_attention_ranks=2), + ) + + with pytest.raises(RuntimeError, match="has role rank 1 but vLLM"): + connector_device.check_connector_role_rank( + connector, _role_rank_config(tp_size=2) + ) diff --git a/tests/unit/v1/worker/test_ffn_routed_experts_model_runner.py b/tests/unit/v1/worker/test_ffn_routed_experts_model_runner.py new file mode 100644 index 000000000..a11c95be5 --- /dev/null +++ b/tests/unit/v1/worker/test_ffn_routed_experts_model_runner.py @@ -0,0 +1,312 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""CPU tests for ``GPUFFNRoutedExpertsModelRunner`` with a fake connector.""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest + +torch = pytest.importorskip("torch") +pytest.importorskip("vllm") + +from afd_plugin.v1.worker import ( # noqa: E402 + ffn_routed_experts_model_runner as routed_module, +) +from afd_plugin.v1.worker.ffn_model_runner import GPUFFNModelRunner # noqa: E402 +from afd_plugin.v1.worker.ffn_routed_experts_model_runner import ( # noqa: E402 + GPUFFNRoutedExpertsModelRunner, + local_experts_from_vllm, + select_routed_moe_backend, +) +from tests.unit.fakes.connector import ( # noqa: E402 + FakeRoutedConnector, +) + +MOE_LAYERS = (1, 3) +NUM_LOCAL_EXPERTS = 2 +HIDDEN = 256 +INTERMEDIATE = 128 +FP8_BLOCK = 128 + + +def _runner(connector=None): + runner = object.__new__(GPUFFNRoutedExpertsModelRunner) + runner.connector = connector or FakeRoutedConnector() + runner.moe_layer_indices = MOE_LAYERS + runner.local_experts = { + layer: SimpleNamespace(weights=f"w{layer}") for layer in MOE_LAYERS + } + runner.use_cuda_graph = False + runner._cuda_graphs = {} + runner.prof = None + return runner + + +def _dp(*tokens): + return SimpleNamespace(num_tokens_across_dp_cpu=torch.tensor(tokens)) + + +def test_routed_runner_reuses_the_connector_driven_ffn_runner(): + """The routed runner is the FFN runner with a different forward. + + Why: graph capture, replay, the loop contract and shutdown are shared. + """ + assert issubclass(GPUFFNRoutedExpertsModelRunner, GPUFFNModelRunner) + + +def _recv_and_run(stage, layer): + return [("recv", stage), ("run", f"rows{stage}", f"w{layer}")] + + +def _send(stage): + return [("send", f"out(rows{stage})", stage)] + + +def test_forward_follows_attention_dbo_order_then_joins_lanes(): + """For layers (1, 3) and stages (0, 1) the FFN issues Attention's DBO + micro-batch order: both stages dispatch layer 1, then each stage sends + its layer-1 result just before dispatching layer 3, and the last results + go back at the end. Each recv runs that layer's experts; the lanes rejoin + once at the end. + + Why: Attention's micro-batch threads issue D0 D1 C0 D0' C1 D1' and + dispatch/combine can block the host, so sending stage 0 right after its + experts (D0 C0 D1) hung every rank in the graph+DBO warmup. A CUDA-graph + capture also fails if forked lane streams are not joined. + """ + runner = _runner() + + output = runner._ffn_forward(dp_metadata_list={1: _dp(3), 0: _dp(5)}) + + first, last = MOE_LAYERS + expected = [ + *_recv_and_run(0, first), + *_recv_and_run(1, first), + *_send(0), + *_recv_and_run(0, last), + *_send(1), + *_recv_and_run(1, last), + *_send(0), + *_send(1), + ("sync",), + ] + assert runner.connector.events == expected + assert output == "out(rows1)" + assert len(runner.connector.control_updates) == 1 + + +def test_forward_with_one_stage_sends_each_layer_before_the_next_recv(): + """With one stage the order is recv, run, send per MoE layer. + + Why: without DBO, Attention combines a layer before it dispatches the next. + """ + runner = _runner() + + runner._ffn_forward(dp_metadata_list={0: _dp(5)}) + + first, last = MOE_LAYERS + assert runner.connector.events == [ + *_recv_and_run(0, first), + *_send(0), + *_recv_and_run(0, last), + *_send(0), + ("sync",), + ] + + +def test_forward_skips_control_update_inside_capture(): + runner = _runner() + + runner._ffn_forward(dp_metadata_list={0: _dp(1)}, update_connector_state=False) + + assert runner.connector.control_updates == [] + assert runner.connector.events[-1] == ("sync",) + + +def test_eager_fallback_keeps_attention_replay_flag(): + """A step Attention replays but the FFN has no graph for runs eagerly + with ``is_graph_replaying`` still set in the control payload. + + Why: the connector then plans the step like Attention (no host sync). + """ + runner = _runner() + runner.use_cuda_graph = True + + runner.execute_model( + dp_metadata_list={0: _dp(4)}, + is_graph_replaying=True, + ) + + (payload,) = runner.connector.control_updates + assert payload.is_graph_replaying is True + assert payload.is_graph_capturing is False + + +def test_load_model_takes_each_moe_layers_local_experts(monkeypatch): + """After the native load, the runner asks the model for its MoE layers and + converts each layer's vLLM experts once. + + Why: weights are read after vLLM's post-load processing, not per step. + """ + routed_experts = {layer: object() for layer in MOE_LAYERS} + model = SimpleNamespace( + get_experts_layer_indices=lambda: list(MOE_LAYERS), + get_routed_experts=routed_experts.__getitem__, + ) + + def native_load(self, *, load_dummy_weights=False, **kwargs): + self.model = model + + monkeypatch.setattr(GPUFFNModelRunner, "load_model", native_load) + monkeypatch.setattr( + routed_module, + "local_experts_from_vllm", + lambda experts: ("converted", experts), + ) + runner = _runner() + + runner.load_model() + + assert runner.moe_layer_indices == MOE_LAYERS + assert runner.local_experts == { + layer: ("converted", routed_experts[layer]) for layer in MOE_LAYERS + } + + +def test_initialize_passes_moe_layers_and_experts_to_the_connector(): + runner = _runner() + + runner.initialize_afd_connector() + + assert runner.connector.init_kwargs == { + "moe_layer_indices": MOE_LAYERS, + "local_experts": runner.local_experts, + } + + +def test_constructor_rejects_connectors_without_routed_rows(monkeypatch): + def base_init(self, vllm_config, device): + self.connector = SimpleNamespace(close=lambda: None) + + monkeypatch.setattr(GPUFFNModelRunner, "__init__", base_init) + + with pytest.raises(TypeError, match="is not an AFDRoutedExpertsConnector"): + GPUFFNRoutedExpertsModelRunner(SimpleNamespace(), torch.device("cpu")) + + +def test_ffn_runner_adopts_bootstrapped_connector_and_closes_placeholder(): + """The worker's bootstrapped connector replaces the one built at + construction, which is closed. + + Why: bootstrap must precede vLLM's distributed init, construction cannot. + """ + events: list[str] = [] + runner = _runner(SimpleNamespace(close=lambda: events.append("close"))) + bootstrapped = FakeRoutedConnector() + + runner.adopt_bootstrapped_connector(bootstrapped) + + assert runner.connector is bootstrapped + assert events == ["close"] + + +def _routed_experts(*, fp8: bool, expert_ids=(4, 5), backend: str = "TRITON"): + num_experts = len(expert_ids) + dtype = torch.float8_e4m3fn if fp8 else torch.bfloat16 + selected = SimpleNamespace(name=backend) + experts = SimpleNamespace( + quant_method=( + SimpleNamespace(fp8_backend=selected) + if fp8 + else SimpleNamespace(unquantized_backend=selected) + ), + w13_weight=torch.nn.Parameter( + torch.zeros(num_experts, 2 * INTERMEDIATE, HIDDEN, dtype=dtype), + requires_grad=False, + ), + w2_weight=torch.nn.Parameter( + torch.zeros(num_experts, HIDDEN, INTERMEDIATE, dtype=dtype), + requires_grad=False, + ), + expert_map_manager=SimpleNamespace( + get_local_expert_ids=lambda: list(expert_ids), + ), + ) + if fp8: + experts.w13_weight_scale_inv = torch.nn.Parameter( + torch.ones(num_experts, 2 * INTERMEDIATE // FP8_BLOCK, HIDDEN // FP8_BLOCK), + requires_grad=False, + ) + experts.w2_weight_scale_inv = torch.nn.Parameter( + torch.ones(num_experts, HIDDEN // FP8_BLOCK, INTERMEDIATE // FP8_BLOCK), + requires_grad=False, + ) + return experts + + +@pytest.mark.parametrize("fp8", [False, True], ids=["bf16", "fp8"]) +def test_local_experts_are_views_of_vllm_parameters(fp8): + """The executor's weights share storage with vLLM's parameters, take the + block scales only for FP8, and carry vLLM's local expert ids. + + Why: no second copy of the experts, and the ids feed the placement check. + """ + experts = _routed_experts(fp8=fp8) + + local = local_experts_from_vllm(experts) + + assert local.expert_ids == (4, 5) + assert local.weights.w13.data_ptr() == experts.w13_weight.data_ptr() + assert local.weights.w2.data_ptr() == experts.w2_weight.data_ptr() + assert local.weights.is_fp8 is fp8 + if fp8: + assert ( + local.weights.w13_scale.data_ptr() + == experts.w13_weight_scale_inv.data_ptr() + ) + else: + assert local.weights.w13_scale is None + + +@pytest.mark.parametrize("fp8", [False, True], ids=["bf16", "fp8"]) +@pytest.mark.parametrize( + "backend", ["FLASHINFER_CUTLASS", "FLASHINFER_TRTLLM", "MARLIN"] +) +def test_local_experts_reject_backends_that_reorder_weights(fp8, backend): + """Experts processed for a backend that reorders them in place are refused. + + Why: FlashInfer swaps w13 to [up; gate] (TRT-LLM and Marlin also shuffle) + without changing dtype or shape; the executor would read [gate; up] and + every MoE layer would compute silu(up) * gate, with no error. + """ + with pytest.raises(RuntimeError, match="moe-backend triton"): + local_experts_from_vllm(_routed_experts(fp8=fp8, backend=backend)) + + +def _kernel_config(moe_backend: str) -> SimpleNamespace: + return SimpleNamespace(kernel_config=SimpleNamespace(moe_backend=moe_backend)) + + +def test_routed_ffn_pins_auto_moe_backend_to_triton(): + """``auto`` becomes ``triton``; an explicit ``triton`` stays. + + Why: auto picks FlashInfer CUTLASS on Hopper with EP > 1 for block FP8, + which reorders the experts; the FFN never runs vLLM's MoE forward, so the + backend only decides the weight layout. + """ + for configured in ("auto", "triton"): + vllm_config = _kernel_config(configured) + select_routed_moe_backend(vllm_config) + assert vllm_config.kernel_config.moe_backend == "triton" + + +@pytest.mark.parametrize("configured", ["flashinfer_cutlass", "deep_gemm", "marlin"]) +def test_routed_ffn_rejects_other_explicit_moe_backends(configured): + """An explicit backend other than triton fails before the weights load. + + Why: the failure must be loud at startup, not wrong tokens at serving. + """ + with pytest.raises(RuntimeError, match="moe-backend triton"): + select_routed_moe_backend(_kernel_config(configured)) diff --git a/tests/unit/v1/worker/test_model_runner_v2.py b/tests/unit/v1/worker/test_model_runner_v2.py index a6c68d0e1..bc1a72af5 100644 --- a/tests/unit/v1/worker/test_model_runner_v2.py +++ b/tests/unit/v1/worker/test_model_runner_v2.py @@ -40,6 +40,9 @@ attention_worker, ffn_worker, ) +from afd_plugin.v1.worker import ( # noqa: E402 + ffn_routed_experts_model_runner as routed_runner_module, +) from afd_plugin.v1.worker.attention_model_runner_v2 import ( # noqa: E402 AFDAttentionModelRunnerV2, ) @@ -48,6 +51,9 @@ validate_gpu_model_runner_v2_config, validate_npu_model_runner_v2_config, ) +from tests.unit.fakes.connector import ( # noqa: E402 + FakeRoutedConnector, +) from tests.unit.fakes.profiler import StepProfiler as _StepProfiler # noqa: E402 TEST_MODEL_MAX_LEN = 163840 @@ -129,6 +135,12 @@ def _v2_config( parallel_config = SimpleNamespace( data_parallel_size=data_parallel_size, data_parallel_rank=0, + data_parallel_rank_local=0, + data_parallel_index=0, + data_parallel_backend="mp", + distributed_executor_backend="mp", + nnodes_within_dp=1, + assigned_physical_gpu_ids=None, tensor_parallel_size=tensor_parallel_size, pipeline_parallel_size=1, prefill_context_parallel_size=prefill_context_parallel_size, @@ -168,6 +180,7 @@ def _v2_config( parallel_config=parallel_config, compilation_config=compilation_config, cache_config=SimpleNamespace(enable_prefix_caching=False), + kernel_config=SimpleNamespace(moe_backend="auto"), scheduler_config=SimpleNamespace( enable_chunked_prefill=False, async_scheduling=False, @@ -248,6 +261,8 @@ def _construct_like_native_worker(worker): def _prepare_worker(worker_cls, config, *, use_v2): worker = object.__new__(worker_cls) + worker.rank = 0 + worker.local_rank = 0 worker.vllm_config = config worker.model_config = config.model_config worker.use_v2_model_runner = use_v2 @@ -2029,6 +2044,47 @@ def native_load(self, load_dummy_weights=False, *args, **kwargs): assert runner.connector.is_initialized +def test_v2_load_model_reports_moe_layers_to_routed_connector(monkeypatch): + """An ``AFDRoutedExpertsConnector`` receives the MoE layers Attention + dispatches, which every rank checks against the FFN's. + + Why: dispatch and combine run once per MoE layer on every rank. + """ + connector = FakeRoutedConnector() + runner = object.__new__(AFDAttentionModelRunnerV2) + runner.connector = connector + runner.model = SimpleNamespace(get_experts_layer_indices=lambda: [1, 3]) + monkeypatch.setattr( + native_v2.GPUModelRunner, "load_model", lambda self, *args, **kwargs: None + ) + monkeypatch.setattr(AFDAttentionModelRunnerV2, "get_model", lambda self: self.model) + + runner.load_model() + + assert connector.init_kwargs == { + "moe_layer_indices": (1, 3), + "local_experts": None, + } + + +def test_v2_runner_adopts_bootstrapped_connector_and_closes_placeholder(): + """The worker's bootstrapped connector replaces the one built at + construction, which is closed. + + Why: the native constructor signature stays unchanged. + """ + events: list[str] = [] + placeholder = SimpleNamespace(close=lambda: events.append("placeholder_close")) + bootstrapped = SimpleNamespace() + runner = object.__new__(AFDAttentionModelRunnerV2) + runner.connector = placeholder + + runner.adopt_bootstrapped_connector(bootstrapped) + + assert runner.connector is bootstrapped + assert events == ["placeholder_close"] + + def test_v2_constructor_rejects_missing_control_plane_and_cleans_native( monkeypatch, ): @@ -2237,6 +2293,235 @@ def raising_worker_init(current_worker): assert _NativeRunnerSentinel.instances == 1 +class _ConnectorRunnerRecorder: + adopted: list[object] = [] + + def __init__(self, vllm_config, device): + pass + + def adopt_bootstrapped_connector(self, connector): + type(self).adopted.append(connector) + + +_ROUTED_WORKER_CASES = [ + ( + attention_worker, + attention_worker.AFDAttentionWorker, + attention_worker, + "AFDAttentionModelRunnerV2", + ), + # The FFN worker imports the routed runner lazily, from its own module. + ( + ffn_worker, + AFDFFNWorker, + routed_runner_module, + "GPUFFNRoutedExpertsModelRunner", + ), +] + + +@pytest.mark.parametrize( + ("worker_module", "worker_cls", "runner_module", "runner_name"), + _ROUTED_WORKER_CASES, +) +def test_workers_bootstrap_connector_before_vllm_and_hand_it_to_runner( + monkeypatch, worker_module, worker_cls, runner_module, runner_name +): + """An ``AFDRoutedExpertsConnector`` is created and bootstrapped before + vLLM's distributed init and adopted by the runner; the FFN picks the + routed-experts runner. + + Why: NCCL reads its environment when vLLM creates the first communicator. + """ + config = _v2_config(role="ffn" if worker_cls is AFDFFNWorker else "attention") + worker = _prepare_worker(worker_cls, config, use_v2=True) + events: list[str] = [] + connector = SimpleNamespace(close=lambda: events.append("close"), local_rank=0) + _ConnectorRunnerRecorder.adopted = [] + + def create_bootstrapped_connector(rank, local_rank, vllm_config, afd_config): + events.append("bootstrap") + return connector + + def vllm_init(current_worker): + events.append("vllm_init") + _fake_worker_init(current_worker) + current_worker.device = torch.device("cuda", 0) + + factory = worker_module.AFDConnectorFactory + monkeypatch.setattr( + factory, "create_bootstrapped_connector", create_bootstrapped_connector + ) + monkeypatch.setattr( + factory, "get_connector_class", lambda name: FakeRoutedConnector + ) + monkeypatch.setattr(runner_module, runner_name, _ConnectorRunnerRecorder) + monkeypatch.setattr(native_v2, "GPUModelRunner", _NativeRunnerSentinel) + monkeypatch.setattr(worker_module.Worker, "init_device", vllm_init) + monkeypatch.setattr( + worker_module, + "assert_compatible_afd_stack", + lambda *args, **kwargs: None, + ) + monkeypatch.setattr(worker_module.torch.accelerator, "empty_cache", lambda: None) + monkeypatch.setattr( + worker_module, + "check_connector_role_rank", + lambda checked, vllm_config: events.append("role_rank_check"), + ) + + worker.init_device() + + assert events == ["bootstrap", "vllm_init", "role_rank_check"] + assert _ConnectorRunnerRecorder.adopted == [connector] + assert native_v2.GPUModelRunner is _NativeRunnerSentinel + # The routed FFN pins vLLM's MoE backend to one that keeps the expert layout. + expected_backend = "triton" if worker_cls is AFDFFNWorker else "auto" + assert config.kernel_config.moe_backend == expected_backend + + +@pytest.mark.parametrize( + ("worker_module", "worker_cls", "runner_module", "runner_name"), + _ROUTED_WORKER_CASES, +) +def test_workers_close_bootstrapped_connector_when_vllm_init_fails( + monkeypatch, worker_module, worker_cls, runner_module, runner_name +): + config = _v2_config(role="ffn" if worker_cls is AFDFFNWorker else "attention") + worker = _prepare_worker(worker_cls, config, use_v2=True) + events: list[str] = [] + connector = SimpleNamespace(close=lambda: events.append("close")) + + def failing_vllm_init(current_worker): + raise RuntimeError("vLLM init failed") + + factory = worker_module.AFDConnectorFactory + monkeypatch.setattr( + factory, "create_bootstrapped_connector", lambda *args: connector + ) + monkeypatch.setattr( + factory, "get_connector_class", lambda name: FakeRoutedConnector + ) + monkeypatch.setattr(runner_module, runner_name, _ConnectorRunnerRecorder) + monkeypatch.setattr(native_v2, "GPUModelRunner", _NativeRunnerSentinel) + monkeypatch.setattr(worker_module.Worker, "init_device", failing_vllm_init) + monkeypatch.setattr( + worker_module, + "assert_compatible_afd_stack", + lambda *args, **kwargs: None, + ) + + with pytest.raises(RuntimeError, match="vLLM init failed"): + worker.init_device() + + assert events == ["close"] + assert native_v2.GPUModelRunner is _NativeRunnerSentinel + + +@pytest.mark.parametrize( + ("worker_module", "worker_cls", "runner_module", "runner_name"), + _ROUTED_WORKER_CASES, +) +def test_workers_reject_bootstrapped_connector_on_another_device( + monkeypatch, worker_module, worker_cls, runner_module, runner_name +): + """A connector bootstrapped on a different GPU than vLLM's fails startup. + + Why: the connector picks its device before vLLM does, by mirroring vLLM's + selection; a vLLM upgrade that changes the selection must fail here + rather than run the connector on another rank's GPU. + """ + config = _v2_config(role="ffn" if worker_cls is AFDFFNWorker else "attention") + worker = _prepare_worker(worker_cls, config, use_v2=True) + events: list[str] = [] + connector = SimpleNamespace(close=lambda: events.append("close"), local_rank=0) + _ConnectorRunnerRecorder.adopted = [] + + def vllm_init(current_worker): + _fake_worker_init(current_worker) + current_worker.device = torch.device("cuda", 4) + + factory = worker_module.AFDConnectorFactory + monkeypatch.setattr( + factory, "create_bootstrapped_connector", lambda *args: connector + ) + monkeypatch.setattr( + factory, "get_connector_class", lambda name: FakeRoutedConnector + ) + monkeypatch.setattr(runner_module, runner_name, _ConnectorRunnerRecorder) + monkeypatch.setattr(native_v2, "GPUModelRunner", _NativeRunnerSentinel) + monkeypatch.setattr(worker_module.Worker, "init_device", vllm_init) + monkeypatch.setattr( + worker_module, + "assert_compatible_afd_stack", + lambda *args, **kwargs: None, + ) + + with pytest.raises(RuntimeError, match="uses cuda:0 but vLLM placed"): + worker.init_device() + + assert events == ["close"] + assert _ConnectorRunnerRecorder.adopted == [] + assert native_v2.GPUModelRunner is _NativeRunnerSentinel + + +def _deepep_m2n_v2_config(**kwargs): + config = _v2_config(**kwargs) + config.additional_config["afd"]["connector"] = "DeepEPM2NAFDConnector" + config.additional_config["afd"]["compute_gate_on_attention"] = True + return config + + +@pytest.mark.parametrize("tp_size", [1, 2]) +def test_v2_validator_accepts_deepep_m2n_attention_tp(tp_size): + """DeepEP M2N accepts Attention TP > 1. + + Why: each Attention TP rank dispatches its own slice of the tokens, and + the bootstrap derives the TP rank from the worker rank before vLLM's + process groups exist. + """ + validate_gpu_model_runner_v2_config( + _deepep_m2n_v2_config( + tensor_parallel_size=tp_size, num_attention_ranks=tp_size + ), + expected_role="attention", + device_type="cuda", + ) + + +def test_v2_validator_rejects_deepep_m2n_ffn_tp(): + """DeepEP M2N FFN ranks must run TP = 1. + + Why: DeepEP routes each token to the one FFN rank that owns the whole + expert; FFN experts are not TP-sharded. + """ + config = _deepep_m2n_v2_config(role="ffn", tensor_parallel_size=2, num_ffn_ranks=2) + with pytest.raises(RuntimeError, match="requires FFN tensor_parallel_size=1"): + validate_gpu_model_runner_v2_config( + config, + expected_role="ffn", + device_type="cuda", + ) + + +@pytest.mark.parametrize("role", ["attention", "ffn"]) +def test_v2_validator_requires_gate_on_attention_for_deepep_m2n(role): + """DeepEP M2N without compute_gate_on_attention fails at validation. + + Why: the connector routes on Attention, so a run without the gate there + cannot dispatch; failing here names the setting instead of failing later + in the connector handshake. + """ + config = _deepep_m2n_v2_config(role=role) + config.additional_config["afd"]["compute_gate_on_attention"] = False + with pytest.raises(RuntimeError, match="requires compute_gate_on_attention=true"): + validate_gpu_model_runner_v2_config( + config, + expected_role=role, + device_type="cuda", + ) + + @pytest.mark.parametrize("use_v2", [False, True]) def test_non_afd_model_config_and_native_runner_identity_are_unchanged( monkeypatch, diff --git a/tests/unit/v1/worker/test_runtime_classpaths.py b/tests/unit/v1/worker/test_runtime_classpaths.py index 2fe1541a0..24fce919a 100644 --- a/tests/unit/v1/worker/test_runtime_classpaths.py +++ b/tests/unit/v1/worker/test_runtime_classpaths.py @@ -89,6 +89,48 @@ def test_gpu_runtime_class_paths_resolve_when_vllm_is_available(qualname): assert cls.__module__.startswith("afd_plugin.v1.worker") +_GPU_WORKER_IMPORT_SCRIPT = """ +import sys + +import afd_plugin.v1.worker.attention_model_runner_v2 +import afd_plugin.v1.worker.attention_worker +import afd_plugin.v1.worker.ffn_worker + +loaded = sorted( + name + for name in sys.modules + if name.startswith("afd_plugin.connectors.gpu.deepep_m2n") + or name == "afd_plugin.v1.worker.ffn_routed_experts_model_runner" +) +assert not loaded, loaded +""" + + +@pytest.mark.vllm_runtime +def test_gpu_workers_and_runners_import_no_connector_backend(): + """Importing the GPU workers and the Attention runner loads no DeepEP M2N + module and not the routed-experts runner. + + Why: workers and runners stay connector-agnostic; P2P deployments must + not pay for, or depend on, another connector's backend. + """ + pytest.importorskip("torch") + pytest.importorskip("vllm") + + result = subprocess.run( + [sys.executable, "-c", _GPU_WORKER_IMPORT_SCRIPT], + capture_output=True, + text=True, + check=False, + ) + + assert result.returncode == 0, ( + "GPU worker import loaded a connector backend:\n" + f"stdout:\n{result.stdout}\n" + f"stderr:\n{result.stderr}" + ) + + @pytest.mark.vllm_runtime def test_npu_runtime_class_paths_resolve_when_vllm_ascend_is_available(): pytest.importorskip("torch") From f2afe7518312d63baa09dabfbbbeb5a4ff957e03 Mon Sep 17 00:00:00 2001 From: Itay Etelis <92247226+Etelis@users.noreply.github.com> Date: Sun, 25 Oct 2026 16:50:56 +0200 Subject: [PATCH 03/32] fix(gpu): run the DeepEP M2N connector on afdEP-comm The afdep extra pinned ms-llmd/afdEP e763e318, a 2026-10-19 development tree under the old distribution name. It rejects the role path across NVLink domains and has none of the cross-node work: the sparse directional hybrid, FP8 combine replies, role-sized inboxes. The library now lives at afdEP/afdEP-comm as the afdep-comm distribution (import name still afdep); pin its main at 2b74d57 and relock. afdEP-comm main also removed ElasticBuffer's role_aware keyword (afdEP-comm #17): ffn_ranks alone selects the AFD transport. Drop the connector's role_aware option with its check and docstring, and build each lane with ElasticBuffer(**kwargs, ffn_ranks=ffn_size). The two changes go together: on the old pin, role_aware defaulted to false, so dropping it there would silently fall back to the unspecialized protocol. A config that still sets role_aware is rejected as an unknown option. The 2F2A GPU case for the removed protocol goes. The fake ElasticBuffer now rejects keywords the real one lacks: the old fake accepted anything, so the unit tests passed while every real rank would fail at lane creation. The module docstring and README name afdEP-comm and describe the baseline as FastAFD's DeepEP M2N. Signed-off-by: Itay Etelis <92247226+Etelis@users.noreply.github.com> --- README.md | 2 +- .../connectors/gpu/deepep_m2n/connector.py | 26 +++++-------- pyproject.toml | 2 +- tests/gpu/test_deepep_m2n_gpu.py | 15 ------- .../connectors/test_deepep_m2n_connector.py | 39 +++++++------------ tests/unit/fakes/afdep.py | 26 ++++++++++++- uv.lock | 10 ++--- 7 files changed, 54 insertions(+), 66 deletions(-) diff --git a/README.md b/README.md index 1223e42c1..432f2bfc7 100644 --- a/README.md +++ b/README.md @@ -57,7 +57,7 @@ See the [recipe index](recipe/README.md) for deployment and benchmark examples. | Connector | Platform | Recommend Stage | Sync or Async | Graph Support | Notes | | --- | --- | --- | --- | --- | --- | | `P2pNcclAFDConnector` | CUDA | Decode | Sync | `FULL_DECODE_ONLY` CUDA graph | FFN ranks are ordered before Attention ranks. `num_attention_ranks` must be greater than or equal to `num_ffn_ranks`; the Attention ranks are spread over the FFN ranks in blocks that differ in size by at most one. See the [DeepSeek V2 Lite recipe](recipe/gpu/P2pNcclAFDConnector/deepseek_v2_lite/README.md). | -| `DeepEPM2NAFDConnector` | CUDA | Decode | Sync | `FULL_DECODE_ONLY` CUDA graph | Expert-routed M2N transport: Attention ranks dispatch routed tokens straight to the FFN ranks that own the chosen experts. Dispatch and combine use [afdEP](https://github.com/ms-llmd/afdEP)'s DeepEP-derived `ElasticBuffer` (install the `afdep` extra); the default `deepgemm` expert backend needs [DeepGEMM](https://github.com/deepseek-ai/DeepGEMM) installed separately. Requires Hopper (SM90) or newer and Attention-side routing. | +| `DeepEPM2NAFDConnector` | CUDA | Decode | Sync | `FULL_DECODE_ONLY` CUDA graph | Expert-routed M2N transport: Attention ranks dispatch routed tokens straight to the FFN ranks that own the chosen experts. Dispatch and combine use [afdEP-comm](https://github.com/afdEP/afdEP-comm)'s DeepEP-derived `ElasticBuffer` (install the `afdep` extra, which installs the `afdep-comm` package; uninstall an older `afdep` package first, since both provide the `afdep` module); the default `deepgemm` expert backend needs [DeepGEMM](https://github.com/deepseek-ai/DeepGEMM) installed separately. Requires Hopper (SM90) or newer and Attention-side routing. | | `CAMP2pAFDConnector` | Ascend NPU | Decode | Sync | `FULL` or `FULL_DECODE_ONLY`; plain MLA graph DBO requires `FULL_DECODE_ONLY` | Uses HCCL/CAMP2P custom ops and native MRV2. DBO requires two ubatches and Attention DP > 1. See the [CAMP2P guide](docs/npu/CAM_P2P_CONNECTOR_USER_GUIDE.md). | | `CAMAsyncAFDConnector` | Ascend NPU | Prefill / decode | Async | FFN-only experimental `FULL` | Eager Attention supports DP+TP/SP and two native DBO microbatches, including DP1. FFN FULL requires the [W4A8 runtime contract](docs/npu/ASYNC_FFN_RUNTIME.md) and accepts eager Attention DBO; FFN does not join native DBO. PCP/DCP are unsupported. See the [Async CAM guide](docs/npu/CAM_ASYNC_CONNECTOR_USER_GUIDE.md) for configuration and validation limits. | diff --git a/afd_plugin/connectors/gpu/deepep_m2n/connector.py b/afd_plugin/connectors/gpu/deepep_m2n/connector.py index 85b4e099b..dca2f0ba7 100644 --- a/afd_plugin/connectors/gpu/deepep_m2n/connector.py +++ b/afd_plugin/connectors/gpu/deepep_m2n/connector.py @@ -40,12 +40,12 @@ payload to every FFN rank over the AFD world's NCCL group. Transport: - Dispatch and combine run on afdEP's DeepEP-derived ``ElasticBuffer`` - (https://github.com/ms-llmd/afdEP). By default it runs afdEP's AFD role - path: Attention ranks only send in dispatch and only receive in combine, - FFN ranks the reverse. It removes Masked DeepEP's entry barriers, union - count exchange and all-rank exits, and sizes buffers and kernels by role. - ``role_aware: false`` keeps afdEP's Masked DeepEP protocol for comparison. + Dispatch and combine run on afdEP-comm's DeepEP-derived ``ElasticBuffer`` + (https://github.com/afdEP/afdEP-comm): Attention ranks only send in + dispatch and only receive in combine, FFN ranks the reverse. Compared with + FastAFD's DeepEP M2N (DeepEP V2 over one symmetric group with dummy expert + slots on Attention ranks), it drops the entry barriers, the all-rank count + exchange and exits, and sizes buffers and kernels by role. Scope: BF16 or FP8 (block-128, fp32 scales) payload; eager, ``torch.compile`` (dispatch/combine are custom ops) and CUDA graphs; up to two lanes (DBO); @@ -191,9 +191,6 @@ class DeepEPM2NConfig(ConnectorExtraInfo): (handshake, DeepEP communicator bootstrap, per-step control payload). ``gloo`` keeps it off the GPUs and the RDMA NICs, so DeepEP's communicator is the only RDMA user. - role_aware: Run afdEP's AFD role path (the default). ``False`` runs its - Masked DeepEP protocol. Every rank must use the same value; afdEP - rejects a mismatch when the buffers are created. """ capacity_tokens_per_rank: int | None = None @@ -209,7 +206,6 @@ class DeepEPM2NConfig(ConnectorExtraInfo): gpu_timeout_s: int = DEFAULT_TIMEOUT_S cpu_sync_min_tokens: int = 0 control_backend: str = "nccl" - role_aware: bool = True @classmethod def from_mapping(cls, raw: Mapping[str, Any] | None) -> DeepEPM2NConfig: @@ -249,8 +245,6 @@ def validate(self) -> None: raise ValueError("cpu_timeout_s and gpu_timeout_s must be positive") if self.control_backend not in CONTROL_BACKENDS: raise ValueError(f"control_backend must be one of {CONTROL_BACKENDS}") - if type(self.role_aware) is not bool: - raise ValueError("role_aware must be a boolean") def to_mapping(self) -> dict[str, Any]: return {**asdict(self), "network": self.network.to_mapping()} @@ -1177,11 +1171,9 @@ def init_afd_connector( self._initialized = True def _new_buffer(self, **kwargs: Any) -> ElasticBuffer: - """One lane's afdEP transport.""" - # The DeepEP M2N layout puts FFN ranks first, as afdEP's role path expects. - return ElasticBuffer( - **kwargs, role_aware=self.config.role_aware, ffn_ranks=self.ffn_size - ) + """One lane's afdEP-comm transport.""" + # The DeepEP M2N layout puts FFN ranks first, as afdEP-comm expects. + return ElasticBuffer(**kwargs, ffn_ranks=self.ffn_size) def close(self) -> None: """Destroy DeepEP buffers; safe after partial init and on repeat calls.""" diff --git a/pyproject.toml b/pyproject.toml index 658667885..a2da9103b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -36,7 +36,7 @@ dependencies = [] [project.optional-dependencies] vllm = ["vllm==0.30.0"] afdep = [ - "afdep @ git+https://github.com/ms-llmd/afdEP.git@e763e3189398beda24f592d2117e6ade135b0a7f", + "afdep-comm @ git+https://github.com/afdEP/afdEP-comm.git@2b74d574e9d25cc794de8d2359f815fcfb48955a", ] [project.urls] diff --git a/tests/gpu/test_deepep_m2n_gpu.py b/tests/gpu/test_deepep_m2n_gpu.py index 94d7027cf..28b99bac4 100644 --- a/tests/gpu/test_deepep_m2n_gpu.py +++ b/tests/gpu/test_deepep_m2n_gpu.py @@ -157,21 +157,6 @@ def test_transport_with_fused_experts_backend(payload_format, tmp_path): ) -@pytest.mark.parametrize("payload_format", PAYLOADS, ids=["bf16", "fp8"]) -def test_transport_with_masked_protocol(payload_format, tmp_path): - """Standard cases on 2F2A with ``role_aware: false``: afdEP's Masked DeepEP - protocol instead of its AFD role path. - - Why: the masked protocol is the comparison baseline for the role path and - must stay correct on its own. - """ - _run( - M2NTopology(2, 2), - {"role_aware": False, "payload_format": payload_format}, - tmp_path, - ) - - @pytest.mark.parametrize("payload_format", PAYLOADS, ids=["bf16", "fp8"]) def test_hot_ffn_rank_routing(payload_format, tmp_path): """Standard cases with every top-k choice on FFN rank 0's experts: rank 0 diff --git a/tests/unit/connectors/test_deepep_m2n_connector.py b/tests/unit/connectors/test_deepep_m2n_connector.py index 9abc365be..59d32b62b 100644 --- a/tests/unit/connectors/test_deepep_m2n_connector.py +++ b/tests/unit/connectors/test_deepep_m2n_connector.py @@ -153,21 +153,17 @@ def test_connector_is_registered(): def test_config_defaults_and_round_trip(): - """An empty config gives DeepGEMM experts, a BF16 payload and afdEP's role - path, and the config survives to_mapping/from_mapping unchanged. + """An empty config gives DeepGEMM experts and a BF16 payload, and the config + survives to_mapping/from_mapping unchanged. Why: defaults are what deployments get, and the handshake compares serialized configs. """ config = DeepEPM2NConfig.from_mapping(None) - masked = DeepEPM2NConfig.from_mapping({"role_aware": False}) assert config.expert_backend == "deepgemm" assert config.payload_format == "bf16" - assert config.role_aware is True - assert masked.role_aware is False assert DeepEPM2NConfig.from_mapping(config.to_mapping()) == config - assert DeepEPM2NConfig.from_mapping(masked.to_mapping()) == masked @pytest.mark.parametrize( @@ -188,14 +184,14 @@ def test_config_defaults_and_round_trip(): ({"expert_num_sms": -2}, "expert_num_sms"), ({"expert_num_sms": 131}, "expert_num_sms"), ({"cpu_timeout_s": 0}, "cpu_timeout_s"), - ({"role_aware": "true"}, "role_aware must be a boolean"), - ({"role_aware": 1}, "role_aware must be a boolean"), + # afdEP-comm removed the switch: the transport is always role-specialized. + ({"role_aware": False}, "unknown DeepEPM2NAFDConnector options.*role_aware"), ], ) def test_config_rejects_invalid_options(raw, message): """Each bad or stale option (unknown or removed key, bad enum, zero or - negative size, bad network preset, bad timeout, non-boolean protocol - switch) raises a ValueError naming the field. + negative size, bad network preset, bad timeout) raises a ValueError + naming the field. Why: typos and removed keys must fail at startup, not be ignored. """ @@ -1186,24 +1182,17 @@ def test_attention_init_rejects_local_experts_before_communicating(): @pytest.mark.parametrize(("role", "role_rank"), [("attention", 2), ("ffn", 1)]) -@pytest.mark.parametrize("role_aware", [True, False]) -def test_lane_transport_is_afdep_with_the_protocol_and_ffn_rank_count( - monkeypatch, role, role_rank, role_aware +def test_lane_transport_is_afdep_comm_with_the_ffn_rank_count( + monkeypatch, role, role_rank ): - """Each lane's transport is afdEP's ``ElasticBuffer``, built with the - selected protocol and the FFN rank count; the DeepEP M2N arguments pass - through unchanged. + """Each lane's transport is afdEP-comm's ``ElasticBuffer``, built with the + FFN rank count; the DeepEP M2N arguments pass through unchanged. - Why: afdEP's role path needs the FFN-first rank split, and every rank must - request the same protocol. + Why: ``ffn_ranks`` selects afdEP-comm's AFD transport (without it the + buffer is unspecialized DeepEP V2), and it needs the FFN-first rank split. """ monkeypatch.setattr(connector_module, "ElasticBuffer", FakeElasticBuffer) - connector = _connector( - role, - role_rank, - M2NTopology(ffn_size=2, attn_size=3), - {"role_aware": role_aware}, - ) + connector = _connector(role, role_rank, M2NTopology(ffn_size=2, attn_size=3), {}) buffer = connector._new_buffer(capacity_tokens_per_rank=64, top_k=4) @@ -1211,11 +1200,11 @@ def test_lane_transport_is_afdep_with_the_protocol_and_ffn_rank_count( assert buffer.kwargs == { "capacity_tokens_per_rank": 64, "top_k": 4, - "role_aware": role_aware, "ffn_ranks": 2, } + def test_attention_ops_take_the_stage_from_the_live_dbo_micro_batch(monkeypatch): """Under DBO the dispatch/combine op bodies use the calling micro-batch's id, not the ``stage_idx`` argument; outside DBO they use the argument. diff --git a/tests/unit/fakes/afdep.py b/tests/unit/fakes/afdep.py index 23697e6f0..14ffa495c 100644 --- a/tests/unit/fakes/afdep.py +++ b/tests/unit/fakes/afdep.py @@ -1,10 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project -"""A stand-in for the afdEP package where it is not installed. +"""A stand-in for afdEP-comm's ``afdep`` package where it is not installed. ``DeepEPM2NAFDConnector`` imports ``afdep.deepep`` when its module loads. The stand-in provides the names it imports, so connector tests run on hosts -without afdEP; tests that build a lane substitute ``FakeElasticBuffer``. +without afdEP-comm; tests that build a lane substitute ``FakeElasticBuffer``. """ from __future__ import annotations @@ -13,11 +13,33 @@ from importlib.util import find_spec from types import ModuleType +# Keyword parameters of afdep.deepep.ElasticBuffer at the pinned afdEP-comm +# (afdEP/afdEP-comm 2b74d57). The stand-in rejects any other keyword, so a lane +# built with one the real transport lacks fails here as it would on a GPU. +ELASTIC_BUFFER_PARAMETERS = ( + "group", + "capacity_tokens_per_rank", + "hidden_size", + "num_experts", + "top_k", + "allow_hybrid_mode", + "rdma_traffic_class", + "cpu_timeout_s", + "gpu_timeout_s", + "use_fp8_dispatch", + "num_sms", + "num_qps", + "ffn_ranks", +) + class FakeElasticBuffer: """Records the arguments one lane's transport is built with.""" def __init__(self, **kwargs: object) -> None: + unknown = sorted(set(kwargs) - set(ELASTIC_BUFFER_PARAMETERS)) + if unknown: + raise TypeError(f"ElasticBuffer() got unexpected keywords {unknown}") self.kwargs = kwargs diff --git a/uv.lock b/uv.lock index 77a951457..4965a963b 100644 --- a/uv.lock +++ b/uv.lock @@ -28,9 +28,9 @@ wheels = [ ] [[package]] -name = "afdep" -version = "1.0.0.dev0" -source = { git = "https://github.com/ms-llmd/afdEP.git?rev=e763e3189398beda24f592d2117e6ade135b0a7f#e763e3189398beda24f592d2117e6ade135b0a7f" } +name = "afdep-comm" +version = "1.1.0" +source = { git = "https://github.com/afdEP/afdEP-comm.git?rev=2b74d574e9d25cc794de8d2359f815fcfb48955a#2b74d574e9d25cc794de8d2359f815fcfb48955a" } dependencies = [ { name = "ninja" }, { name = "packaging" }, @@ -5853,7 +5853,7 @@ source = { editable = "." } [package.optional-dependencies] afdep = [ - { name = "afdep" }, + { name = "afdep-comm" }, ] vllm = [ { name = "vllm" }, @@ -5875,7 +5875,7 @@ e2e-tests = [ [package.metadata] requires-dist = [ - { name = "afdep", marker = "extra == 'afdep'", git = "https://github.com/ms-llmd/afdEP.git?rev=e763e3189398beda24f592d2117e6ade135b0a7f" }, + { name = "afdep-comm", marker = "extra == 'afdep'", git = "https://github.com/afdEP/afdEP-comm.git?rev=2b74d574e9d25cc794de8d2359f815fcfb48955a" }, { name = "vllm", marker = "extra == 'vllm'", specifier = "==0.30.0" }, ] provides-extras = ["vllm", "afdep"] From 8a79b86c69ba3d1b0881e0a3ea2a65052a465a68 Mon Sep 17 00:00:00 2001 From: Itay Etelis <92247226+Etelis@users.noreply.github.com> Date: Sun, 25 Oct 2026 16:57:12 +0200 Subject: [PATCH 04/32] fix(gpu): import Triton through vllm.triton_utils in the FP8 quant kernel fp8_quant.py imported triton directly, so importing the DeepEP M2N connector failed on hosts without Triton and its unit tests could not be collected on a CPU-only host, which AGENTS.md requires of tests/unit. vllm.triton_utils provides Triton when it is installed and placeholders otherwise, as vLLM's own kernels do; the kernel is unchanged on GPUs. Signed-off-by: Itay Etelis <92247226+Etelis@users.noreply.github.com> --- afd_plugin/connectors/gpu/deepep_m2n/fp8_quant.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/afd_plugin/connectors/gpu/deepep_m2n/fp8_quant.py b/afd_plugin/connectors/gpu/deepep_m2n/fp8_quant.py index c2e9074c9..0badf760a 100644 --- a/afd_plugin/connectors/gpu/deepep_m2n/fp8_quant.py +++ b/afd_plugin/connectors/gpu/deepep_m2n/fp8_quant.py @@ -19,8 +19,7 @@ from __future__ import annotations import torch -import triton -import triton.language as tl +from vllm.triton_utils import tl, triton from afd_plugin.connectors.routed_experts import FP8_DTYPE From 8f1a73f90e92fdec89289310dd0977f834b8d59b Mon Sep 17 00:00:00 2001 From: Itay Etelis <92247226+Etelis@users.noreply.github.com> Date: Sun, 25 Oct 2026 16:57:12 +0200 Subject: [PATCH 05/32] perf(gpu): give Attention ranks an empty combine input Attention ranks allocated a zero buffer sized for DeepEP's worst-case no-sync receive at the top bucket rung (union ranks x bucket x top-k, 96 MiB to 1 GiB per Attention GPU), before vLLM sizes the KV cache. afdEP-comm gives Attention ranks no receive rows on any dispatch, with or without the host sync, so combine only ever read zero rows of it. Allocate a (0, hidden) tensor once instead, keeping one stable address for captured graphs, and fail loudly if the transport ever reports receive rows for an Attention rank. worst_case_receive_rows and its test go with the buffer they sized; the precompile test already checks that every Attention combine gets an empty input. Signed-off-by: Itay Etelis <92247226+Etelis@users.noreply.github.com> --- .../connectors/gpu/deepep_m2n/connector.py | 55 +++--------- .../connectors/test_deepep_m2n_connector.py | 89 +++++-------------- 2 files changed, 35 insertions(+), 109 deletions(-) diff --git a/afd_plugin/connectors/gpu/deepep_m2n/connector.py b/afd_plugin/connectors/gpu/deepep_m2n/connector.py index dca2f0ba7..701e12e0b 100644 --- a/afd_plugin/connectors/gpu/deepep_m2n/connector.py +++ b/afd_plugin/connectors/gpu/deepep_m2n/connector.py @@ -182,11 +182,8 @@ class DeepEPM2NConfig(ConnectorExtraInfo): cpu_timeout_s / gpu_timeout_s: DeepEP host-wait and kernel-spin timeouts. cpu_sync_min_tokens: Eager steps whose bucket is below this skip the - host sync that reads receive counts; they allocate the worst-case - receive size instead. ``0`` always syncs in eager mode. Attention - ranks keep a zero buffer for the largest bucket that skips the - sync (the top rung of ``bucket_ladder``), so raising this costs - memory. + host sync that reads receive counts; FFN ranks then allocate the + worst-case receive size instead. ``0`` always syncs in eager mode. control_backend: Backend of the connector's own process group (handshake, DeepEP communicator bootstrap, per-step control payload). ``gloo`` keeps it off the GPUs and the RDMA NICs, so @@ -432,26 +429,6 @@ def precompile_buckets( return tuple(dict.fromkeys(buckets)) -def worst_case_receive_rows( - *, - num_ranks: int, - bucket: int, - top_k: int, - experts_per_rank: int, - expert_alignment: int, -) -> int: - """Rows DeepEP allocates for one rank's receive when it skips the host sync. - - Mirrors DeepEP V2's worst case for expanded receives (``ElasticBuffer`` - non-cached dispatch without CPU sync): every rank's bucket times the routes - one token can send to this rank, plus alignment padding for each local - expert. - """ - rows = num_ranks * bucket * min(top_k, experts_per_rank) - rows += (expert_alignment - 1) * experts_per_rank - return -(-rows // expert_alignment) * expert_alignment - - @dataclass(frozen=True) class HandshakeResult: expert_alignment: int @@ -1004,8 +981,7 @@ def __init__( self._step_fork: torch.cuda.Event | None = None self._step_fork_capturing = False self._lanes_joined: set[int] = set() - # Attention side: the zero rows combine takes from Attention ranks, - # sized for the largest receive (see _allocate_combine_zeros). + # Attention side: the empty combine input (see _allocate_combine_zeros). self._combine_zeros: torch.Tensor | None = None self._op_id = next(_CONNECTOR_IDS) self._initialized = False @@ -1559,24 +1535,15 @@ def _precompile_dispatch_kernels(self) -> None: ) def _allocate_combine_zeros(self) -> None: - """One zero buffer for every combine on this Attention rank. + """The empty input every combine on this Attention rank takes. - Combine takes one zero row per row this rank received. A synced - dispatch gives an Attention rank no rows; a dispatch that skips the - sync gets DeepEP's worst case for its bucket, so the buffer covers the - largest such bucket. Both lanes share it (combine only reads it), and - captured graphs keep one stable address. + Combine takes one zero row per row this rank received, and afdEP-comm + gives Attention ranks no receive rows on any dispatch, with or without + the host sync. One zero-row tensor serves both lanes and keeps a + stable address for captured graphs. """ - assert self.bucket_ladder is not None - rows = worst_case_receive_rows( - num_ranks=self.layout.union_size, - bucket=self.bucket_ladder[-1] if self.bucket_ladder else 0, - top_k=self.top_k, - experts_per_rank=self.layout.experts_per_rank, - expert_alignment=self.expert_alignment, - ) self._combine_zeros = torch.zeros( - (rows, self.hidden_size), dtype=self.dtype, device=self.device + (0, self.hidden_size), dtype=self.dtype, device=self.device ) def _combine_zero_rows(self, shape: tuple[int, int]) -> torch.Tensor: @@ -1587,8 +1554,8 @@ def _combine_zero_rows(self, shape: tuple[int, int]) -> torch.Tensor: if rows > zeros.shape[0] or shape[1] != zeros.shape[1]: raise RuntimeError( f"combine needs zero rows of shape {shape} but the buffer holds " - f"{tuple(zeros.shape)}: a dispatch skipped the host sync with a " - "bucket above the ladder's top rung (no_sync_bucket_limit)" + f"{tuple(zeros.shape)}: afdEP-comm gave an Attention rank receive " + "rows, which its AFD transport never does" ) return zeros[:rows] diff --git a/tests/unit/connectors/test_deepep_m2n_connector.py b/tests/unit/connectors/test_deepep_m2n_connector.py index 59d32b62b..85b44c179 100644 --- a/tests/unit/connectors/test_deepep_m2n_connector.py +++ b/tests/unit/connectors/test_deepep_m2n_connector.py @@ -43,7 +43,6 @@ tokens_per_attention_rank, validate_handshake, validate_model_infos, - worst_case_receive_rows, ) from afd_plugin.connectors.gpu.deepep_m2n.experts import ( # noqa: E402 expert_row_ranges, @@ -1204,7 +1203,6 @@ def test_lane_transport_is_afdep_comm_with_the_ffn_rank_count( } - def test_attention_ops_take_the_stage_from_the_live_dbo_micro_batch(monkeypatch): """Under DBO the dispatch/combine op bodies use the calling micro-batch's id, not the ``stage_idx`` argument; outside DBO they use the argument. @@ -1251,39 +1249,6 @@ def combine_attention(self, out, stage): DEEPGEMM_ALIGNMENT = 128 -@pytest.mark.parametrize( - ("num_ranks", "bucket", "top_k", "experts_per_rank", "alignment", "expected"), - [ - # Qwen3-30B 2A2F at the default capacity, and at a 512 bucket. - (4, 8192, 8, 64, DEEPGEMM_ALIGNMENT, 270_336), - (4, 512, 8, 64, DEEPGEMM_ALIGNMENT, 24_576), - # Qwen3-235B 4A2F at a 512 bucket. - (6, 512, 8, 64, DEEPGEMM_ALIGNMENT, 32_768), - # No alignment (fused_experts backend): exactly ranks x bucket x top-k. - (4, 10, 8, 64, 1, 320), - # A token sends at most one route per local expert. - (4, 10, 8, 4, 1, 160), - ], -) -def test_worst_case_receive_rows_mirrors_deepep( - num_ranks, bucket, top_k, experts_per_rank, alignment, expected -): - """Rows match DeepEP's worst-case receive without the host sync. - - Why: the Attention-side zero buffer is sized with this; too few rows - would fail combine, and DeepEP's formula is the contract. - """ - rows = worst_case_receive_rows( - num_ranks=num_ranks, - bucket=bucket, - top_k=top_k, - experts_per_rank=experts_per_rank, - expert_alignment=alignment, - ) - - assert rows == expected - - @pytest.mark.parametrize( ("capture", "tp", "granularity", "sync_min", "capacity", "expected"), [ @@ -1299,8 +1264,8 @@ def test_no_sync_bucket_limit(capture, tp, granularity, sync_min, capacity, expe """The largest bucket that skips the host sync: the largest graph bucket or the largest eager bucket below cpu_sync_min_tokens, capped by capacity. - Why: only those dispatches give an Attention rank a worst-case receive; - sizing the zero buffer for capacity instead would waste ~1 GiB per GPU. + Why: it is the ladder's top rung, and dispatches up to it give FFN ranks + DeepEP's worst-case receive instead of a synced count. """ limit = no_sync_bucket_limit( capacity_tokens_per_rank=capacity, @@ -1313,45 +1278,39 @@ def test_no_sync_bucket_limit(capture, tp, granularity, sync_min, capacity, expe assert limit == expected -def test_attention_combine_zeros_cover_the_largest_no_sync_receive(): - """Attention allocates one zero buffer for its largest no-sync receive and - hands combine a contiguous view of its first rows. +class _ReceivingLane(_RecordingLane): + """A transport that reports receive rows, which afdEP-comm never gives an + Attention rank.""" - Why: a cache keyed by receive shape grew without bound when eager buckets - skipped the sync; a fixed buffer bounds memory and keeps one address for - captured graphs. + def dispatch(self, x, x_scale, topk_ids, topk_weights, **kwargs): + super().dispatch(x, x_scale, topk_ids, topk_weights, **kwargs) + return SimpleNamespace(handle=SimpleNamespace(recv_shape=(5, M2N_HIDDEN_SIZE))) + + +def test_attention_combine_rejects_receive_rows(monkeypatch): + """An Attention rank's combine input is empty, so a transport that reports + receive rows for it raises instead of combining. + + Why: afdEP-comm gives Attention ranks no receive rows, with or without the + host sync. The worst-case zero buffer this replaces took up to ~1 GiB per + Attention GPU before KV-cache sizing and was never read. """ + monkeypatch.setattr(torch.cuda, "synchronize", lambda device=None: None) connector = _connector( "attention", 0, M2NTopology(ffn_size=2, attn_size=2), - {"capacity_tokens_per_rank": 16}, + {"capacity_tokens_per_rank": 64}, ) connector.device = torch.device("cpu") + connector.capacity_tokens_per_rank = 64 + connector.bucket_ladder = (8,) connector.expert_alignment = DEEPGEMM_ALIGNMENT - + connector.lanes = [_ReceivingLane()] connector._allocate_combine_zeros() - # Bucket min(16, 512 capture) = 16: 4 ranks x 16 x top-k 8 + 127 x 64. - expected_rows = worst_case_receive_rows( - num_ranks=4, - bucket=16, - top_k=8, - experts_per_rank=64, - expert_alignment=DEEPGEMM_ALIGNMENT, - ) - zeros = connector._combine_zeros - assert zeros is not None - assert zeros.shape == (expected_rows, M2N_HIDDEN_SIZE) - view = connector._combine_zero_rows((5, M2N_HIDDEN_SIZE)) - assert view.shape == (5, M2N_HIDDEN_SIZE) and view.is_contiguous() - assert view.data_ptr() == zeros.data_ptr() - assert connector._combine_zero_rows((0, M2N_HIDDEN_SIZE)).shape == ( - 0, - M2N_HIDDEN_SIZE, - ) - with pytest.raises(RuntimeError, match="no_sync_bucket_limit"): - connector._combine_zero_rows((expected_rows + 1, M2N_HIDDEN_SIZE)) + with pytest.raises(RuntimeError, match="receive rows"): + connector._precompile_dispatch_kernels() class _LoggedEvent: From ca97ef5a75093901b4372e2bb75cd8a22b64b5b8 Mon Sep 17 00:00:00 2001 From: Itay Etelis <92247226+Etelis@users.noreply.github.com> Date: Sun, 25 Oct 2026 18:22:50 +0200 Subject: [PATCH 06/32] test(gpu): build each lane in one method so the fake checks real arguments The fake ElasticBuffer rejected unknown keywords but accepted calls that left out the real constructor's required ones, and the lane test built a buffer from two hand-picked arguments. A connector change that stopped passing, say, cpu_timeout_s would still pass while every GPU rank failed at lane creation. Lanes are now built by _new_lane, which assembles the full argument set the init path used to pass inline. The fake requires every required keyword of afdEP-comm's ElasticBuffer and rejects any other, and the lane test goes through _new_lane. Signed-off-by: Itay Etelis <92247226+Etelis@users.noreply.github.com> --- .../connectors/gpu/deepep_m2n/connector.py | 39 ++++++++++--------- .../connectors/test_deepep_m2n_connector.py | 20 +++++----- tests/unit/fakes/afdep.py | 25 ++++++------ 3 files changed, 45 insertions(+), 39 deletions(-) diff --git a/afd_plugin/connectors/gpu/deepep_m2n/connector.py b/afd_plugin/connectors/gpu/deepep_m2n/connector.py index 701e12e0b..d974a3772 100644 --- a/afd_plugin/connectors/gpu/deepep_m2n/connector.py +++ b/afd_plugin/connectors/gpu/deepep_m2n/connector.py @@ -1112,22 +1112,7 @@ def init_afd_connector( ) try: for _ in range(self.num_lanes): - self.lanes.append( - self._new_buffer( - group=self.bootstrap_group, - capacity_tokens_per_rank=self.capacity_tokens_per_rank, - hidden_size=self.hidden_size, - num_experts=self.layout.deepep_num_experts, - top_k=self.top_k, - allow_hybrid_mode=self.network_plan.hybrid, - use_fp8_dispatch=self.is_fp8_payload, - num_sms=self.config.num_sms, - num_qps=self.config.num_qps, - rdma_traffic_class=self.network_plan.rdma_traffic_class, - cpu_timeout_s=self.config.cpu_timeout_s, - gpu_timeout_s=self.config.gpu_timeout_s, - ) - ) + self.lanes.append(self._new_lane()) # One lane runs on the caller's stream: no fork/join cost. if self.num_lanes > 1: self.lane_streams.append(torch.cuda.Stream(self.device)) @@ -1146,10 +1131,26 @@ def init_afd_connector( _CONNECTORS[self._op_id] = self self._initialized = True - def _new_buffer(self, **kwargs: Any) -> ElasticBuffer: + def _new_lane(self) -> ElasticBuffer: """One lane's afdEP-comm transport.""" - # The DeepEP M2N layout puts FFN ranks first, as afdEP-comm expects. - return ElasticBuffer(**kwargs, ffn_ranks=self.ffn_size) + assert self.bootstrap_group is not None and self.network_plan is not None + assert self.capacity_tokens_per_rank is not None + return ElasticBuffer( + group=self.bootstrap_group, + capacity_tokens_per_rank=self.capacity_tokens_per_rank, + hidden_size=self.hidden_size, + num_experts=self.layout.deepep_num_experts, + top_k=self.top_k, + allow_hybrid_mode=self.network_plan.hybrid, + use_fp8_dispatch=self.is_fp8_payload, + num_sms=self.config.num_sms, + num_qps=self.config.num_qps, + rdma_traffic_class=self.network_plan.rdma_traffic_class, + cpu_timeout_s=self.config.cpu_timeout_s, + gpu_timeout_s=self.config.gpu_timeout_s, + # The DeepEP M2N layout puts FFN ranks first, as afdEP-comm expects. + ffn_ranks=self.ffn_size, + ) def close(self) -> None: """Destroy DeepEP buffers; safe after partial init and on repeat calls.""" diff --git a/tests/unit/connectors/test_deepep_m2n_connector.py b/tests/unit/connectors/test_deepep_m2n_connector.py index 85b44c179..cffada8c9 100644 --- a/tests/unit/connectors/test_deepep_m2n_connector.py +++ b/tests/unit/connectors/test_deepep_m2n_connector.py @@ -1184,23 +1184,25 @@ def test_attention_init_rejects_local_experts_before_communicating(): def test_lane_transport_is_afdep_comm_with_the_ffn_rank_count( monkeypatch, role, role_rank ): - """Each lane's transport is afdEP-comm's ``ElasticBuffer``, built with the - FFN rank count; the DeepEP M2N arguments pass through unchanged. + """Each lane's transport is afdEP-comm's ``ElasticBuffer``, built with every + argument its constructor requires and the FFN rank count. Why: ``ffn_ranks`` selects afdEP-comm's AFD transport (without it the - buffer is unspecialized DeepEP V2), and it needs the FFN-first rank split. + buffer is unspecialized DeepEP V2), and the stand-in fails on any argument + set the real constructor would refuse. """ monkeypatch.setattr(connector_module, "ElasticBuffer", FakeElasticBuffer) connector = _connector(role, role_rank, M2NTopology(ffn_size=2, attn_size=3), {}) + connector.bootstrap_group = SimpleNamespace() + connector.network_plan = SimpleNamespace(hybrid=True, rdma_traffic_class=0) + connector.capacity_tokens_per_rank = 64 - buffer = connector._new_buffer(capacity_tokens_per_rank=64, top_k=4) + buffer = connector._new_lane() assert isinstance(buffer, FakeElasticBuffer) - assert buffer.kwargs == { - "capacity_tokens_per_rank": 64, - "top_k": 4, - "ffn_ranks": 2, - } + assert buffer.kwargs["ffn_ranks"] == 2 + assert buffer.kwargs["capacity_tokens_per_rank"] == 64 + assert buffer.kwargs["allow_hybrid_mode"] is True def test_attention_ops_take_the_stage_from_the_live_dbo_micro_batch(monkeypatch): diff --git a/tests/unit/fakes/afdep.py b/tests/unit/fakes/afdep.py index 14ffa495c..aaf564877 100644 --- a/tests/unit/fakes/afdep.py +++ b/tests/unit/fakes/afdep.py @@ -13,10 +13,11 @@ from importlib.util import find_spec from types import ModuleType -# Keyword parameters of afdep.deepep.ElasticBuffer at the pinned afdEP-comm -# (afdEP/afdEP-comm 2b74d57). The stand-in rejects any other keyword, so a lane -# built with one the real transport lacks fails here as it would on a GPU. -ELASTIC_BUFFER_PARAMETERS = ( +# Keyword-only parameters of afdep.deepep.ElasticBuffer at the pinned afdEP-comm +# (afdEP/afdEP-comm 2b74d57). The stand-in requires the required ones and +# rejects any other keyword, so a lane built with arguments the real transport +# would refuse fails here as it would on a GPU. +ELASTIC_BUFFER_REQUIRED = ( "group", "capacity_tokens_per_rank", "hidden_size", @@ -26,20 +27,22 @@ "rdma_traffic_class", "cpu_timeout_s", "gpu_timeout_s", - "use_fp8_dispatch", - "num_sms", - "num_qps", - "ffn_ranks", ) +ELASTIC_BUFFER_OPTIONAL = ("use_fp8_dispatch", "num_sms", "num_qps", "ffn_ranks") class FakeElasticBuffer: """Records the arguments one lane's transport is built with.""" def __init__(self, **kwargs: object) -> None: - unknown = sorted(set(kwargs) - set(ELASTIC_BUFFER_PARAMETERS)) - if unknown: - raise TypeError(f"ElasticBuffer() got unexpected keywords {unknown}") + missing = sorted(set(ELASTIC_BUFFER_REQUIRED) - set(kwargs)) + unknown = sorted( + set(kwargs) - set(ELASTIC_BUFFER_REQUIRED) - set(ELASTIC_BUFFER_OPTIONAL) + ) + if missing or unknown: + raise TypeError( + f"ElasticBuffer() missing keywords {missing}, unexpected {unknown}" + ) self.kwargs = kwargs From 3d6ed6b0482ac55f79092595f3f47a9b99e90933 Mon Sep 17 00:00:00 2001 From: ronenkat <16743404+ronenkat@users.noreply.github.com> Date: Mon, 26 Oct 2026 10:09:41 +0200 Subject: [PATCH 07/32] fix(gpu): reject DeepEP M2N for models that do not route on Attention DeepEP M2N needs the Attention side to send topk_ids and topk_weights with the hidden states. Only Qwen3 MoE does: on CUDA, DeepSeek V2/V3 still builds the remote-experts MoE shell, which sends router logits only, so the first MoE forward failed in send_attn_output. DeepSeek would also need its grouped top-k, routed scaling and shared experts on Attention, none of which the M2N path implements yet. Validate the model architecture with the other M2N deployment checks. Both workers run this validation before the connector bootstrap, so every rank fails at startup, naming the supported models, before any collective. Co-Authored-By: Claude Opus 5.5 Signed-off-by: ronenkat <16743404+ronenkat@users.noreply.github.com> --- afd_plugin/validation.py | 18 ++++++++- tests/unit/config/test_validation.py | 40 +++++++++++++++++++- tests/unit/v1/worker/test_model_runner_v2.py | 2 +- 3 files changed, 57 insertions(+), 3 deletions(-) diff --git a/afd_plugin/validation.py b/afd_plugin/validation.py index 89750b076..eab1ccced 100644 --- a/afd_plugin/validation.py +++ b/afd_plugin/validation.py @@ -48,6 +48,9 @@ "P2pNcclAFDConnector", *DEEPEP_M2N_CONNECTORS, ) +# Models whose Attention side computes the MoE gate and top-k and sends them +# with the hidden states, as a routed-experts connector (DeepEP M2N) needs. +ATTENTION_ROUTING_MODEL_ARCHITECTURES: Final[tuple[str, ...]] = ("Qwen3MoeForCausalLM",) CAM_MAX_RECEIVE_TOKENS = 262144 @@ -112,7 +115,9 @@ def validate_gpu_model_runner_v2_config( # TP-sharded; that and the other deployment checks (expert placement, # all2all backend, weight format) run in the connector handshake so that # every rank fails together. It routes on Attention, so it requires the - # gate there. + # gate there, and a model whose Attention side sends the top-k. Workers + # validate before the connector bootstrap, so both roles fail before any + # collective instead of one role timing out. if afd_config.connector in DEEPEP_M2N_CONNECTORS: if expected_role == "ffn" and parallel.tensor_parallel_size != 1: raise RuntimeError( @@ -122,6 +127,16 @@ def validate_gpu_model_runner_v2_config( raise RuntimeError( f"{afd_config.connector} requires compute_gate_on_attention=true" ) + # Runners validate the AFD-prefixed architecture of the wrapped model. + architectures = vllm_config.model_config.hf_config.architectures + if not any( + architecture.removeprefix("AFD") in ATTENTION_ROUTING_MODEL_ARCHITECTURES + for architecture in architectures + ): + raise RuntimeError( + f"{afd_config.connector} supports only " + f"{ATTENTION_ROUTING_MODEL_ARCHITECTURES}, got {architectures}" + ) if ( parallel.pipeline_parallel_size, parallel.prefill_context_parallel_size, @@ -396,6 +411,7 @@ def _ctx() -> str: __all__ = [ "ATTENTION_MODEL_RUNNER_FQCN", + "ATTENTION_ROUTING_MODEL_ARCHITECTURES", "ATTENTION_WORKER_FQCN", "FFN_MODEL_RUNNER_FQCN", "FFN_WORKER_FQCN", diff --git a/tests/unit/config/test_validation.py b/tests/unit/config/test_validation.py index c11c77208..de1ee4e22 100644 --- a/tests/unit/config/test_validation.py +++ b/tests/unit/config/test_validation.py @@ -145,7 +145,10 @@ def mrv2_config(monkeypatch): use_ubatching=True, num_ubatches=2, ), - model_config=SimpleNamespace(enforce_eager=True), + model_config=SimpleNamespace( + enforce_eager=True, + hf_config=SimpleNamespace(architectures=["Qwen3MoeForCausalLM"]), + ), compilation_config=SimpleNamespace( cudagraph_mode="FULL_DECODE_ONLY", pass_config=SimpleNamespace(enable_sp=False), @@ -180,6 +183,41 @@ def test_routed_transport_preserves_runner_and_gate_policy(mrv2_config, role): ) +@pytest.mark.parametrize("role", ["attention", "ffn"]) +@pytest.mark.parametrize( + "architecture, accepted", + [ + pytest.param("Qwen3MoeForCausalLM", True, id="qwen3-moe"), + # Runners validate the wrapped model's AFD-prefixed architecture. + pytest.param("AFDQwen3MoeForCausalLM", True, id="afd-qwen3-moe"), + pytest.param("DeepseekV2ForCausalLM", False, id="deepseek-v2"), + pytest.param("DeepseekV3ForCausalLM", False, id="deepseek-v3"), + pytest.param("AFDDeepseekV3ForCausalLM", False, id="afd-deepseek-v3"), + pytest.param("Qwen3_5MoeForConditionalGeneration", False, id="qwen3-5-moe"), + ], +) +def test_routed_transport_requires_model_that_routes_on_attention( + monkeypatch, mrv2_config, role, architecture, accepted +): + mrv2_config.additional_config["afd"].update( + role=role, connector="DeepEPM2NAFDConnector", compute_gate_on_attention=True + ) + mrv2_config.model_config.hf_config.architectures = [architecture] + + if accepted: + validate_gpu_model_runner_v2_config( + mrv2_config, expected_role=role, device_type="cuda" + ) + return + # Workers validate before the connector bootstrap; rejecting a model must + # not need the model runtime. + monkeypatch.setitem(sys.modules, MODEL_UTILS_MODULE, None) + with pytest.raises(RuntimeError, match="DeepEPM2NAFDConnector supports only"): + validate_gpu_model_runner_v2_config( + mrv2_config, expected_role=role, device_type="cuda" + ) + + def _symmetric_ranks(dp_size, tp_size): return { "parallel_config.data_parallel_size": dp_size, diff --git a/tests/unit/v1/worker/test_model_runner_v2.py b/tests/unit/v1/worker/test_model_runner_v2.py index bc1a72af5..1578aa3df 100644 --- a/tests/unit/v1/worker/test_model_runner_v2.py +++ b/tests/unit/v1/worker/test_model_runner_v2.py @@ -2466,7 +2466,7 @@ def vllm_init(current_worker): def _deepep_m2n_v2_config(**kwargs): - config = _v2_config(**kwargs) + config = _v2_config(architecture="Qwen3MoeForCausalLM", **kwargs) config.additional_config["afd"]["connector"] = "DeepEPM2NAFDConnector" config.additional_config["afd"]["compute_gate_on_attention"] = True return config From 72f743101cd3abc3de2413a37ed87564612b673c Mon Sep 17 00:00:00 2001 From: ronenkat <16743404+ronenkat@users.noreply.github.com> Date: Mon, 26 Oct 2026 10:09:49 +0200 Subject: [PATCH 08/32] fix(gpu): compare control_backend in the DeepEP M2N handshake The handshake did not advertise control_backend. With one role on gloo and the other on nccl it succeeded; the NCCL ranks then created a communicator the gloo ranks never joined, and startup hung until the timeout. Advertise the backend and require every rank to match, so the mismatch fails before any communicator or buffer is created. Co-Authored-By: Claude Opus 5.5 Signed-off-by: ronenkat <16743404+ronenkat@users.noreply.github.com> --- .../connectors/gpu/deepep_m2n/connector.py | 3 +++ .../connectors/test_deepep_m2n_connector.py | 18 +++++++++++++++++- 2 files changed, 20 insertions(+), 1 deletion(-) diff --git a/afd_plugin/connectors/gpu/deepep_m2n/connector.py b/afd_plugin/connectors/gpu/deepep_m2n/connector.py index d974a3772..9669301b3 100644 --- a/afd_plugin/connectors/gpu/deepep_m2n/connector.py +++ b/afd_plugin/connectors/gpu/deepep_m2n/connector.py @@ -504,6 +504,7 @@ def local_handshake_info( "cpu_sync_min_tokens": config.cpu_sync_min_tokens, "attn_tp_size": attn_tp_size, "compute_gate_on_attention": compute_gate_on_attention, + "control_backend": config.control_backend, "network": str(config.network.to_mapping()), "expert_alignment": expert_alignment, "ffn_deployment": ffn_deployment, @@ -520,6 +521,8 @@ def local_handshake_info( "top_k", "num_lanes", "compute_gate_on_attention", + # NCCL ranks create a communicator that gloo ranks would never join. + "control_backend", "network", "bucket_granularity", "cpu_sync_min_tokens", diff --git a/tests/unit/connectors/test_deepep_m2n_connector.py b/tests/unit/connectors/test_deepep_m2n_connector.py index cffada8c9..980f82a1e 100644 --- a/tests/unit/connectors/test_deepep_m2n_connector.py +++ b/tests/unit/connectors/test_deepep_m2n_connector.py @@ -100,6 +100,7 @@ def _handshake( overrides: dict | None = None, hosts: str | None = None, attn_tp_size: int = 1, + config: DeepEPM2NConfig | None = None, ) -> list[dict]: """Handshake infos; rank r runs on host hosts[r] (one NVLink domain per host). @@ -114,7 +115,7 @@ def _handshake( role="ffn" if is_ffn else "attention", role_rank=world_rank if is_ffn else world_rank - layout.ffn_size, world_rank=world_rank, - config=DeepEPM2NConfig(), + config=config or DeepEPM2NConfig(), hidden_size=M2N_HIDDEN_SIZE, dtype=torch.bfloat16, num_experts=layout.num_experts, @@ -724,6 +725,21 @@ def test_handshake_rejects_mismatched_network_config(): validate_handshake(infos, layout, NetworkConfig()) +def test_handshake_rejects_mixed_control_backends(): + """FFN on ``nccl`` and Attention on ``gloo``: the handshake raises. + + Why: NCCL ranks would create a communicator that the gloo ranks never + join, so startup would hang until the timeout instead of failing. + """ + layout = DeepEPM2NLayout(ffn_size=1, attn_size=1, num_experts=128) + nccl = _handshake(layout) + gloo = _handshake(layout, config=DeepEPM2NConfig(control_backend="gloo")) + + validate_handshake(gloo, layout, NetworkConfig()) + with pytest.raises(RuntimeError, match="disagree on control_backend"): + validate_handshake([nccl[0], gloo[1]], layout, NetworkConfig()) + + def test_gloo_control_backend_stages_payloads_on_cpu(): """With ``control_backend: gloo`` control payloads are staged on CPU. From 84dd9bb486628d90994e71d5aa6e285d34263b1c Mon Sep 17 00:00:00 2001 From: ronenkat <16743404+ronenkat@users.noreply.github.com> Date: Mon, 26 Oct 2026 10:09:50 +0200 Subject: [PATCH 09/32] test(gpu): fail the DeepEP M2N checks on non-finite errors The harness kept the worst error per metric with max(), which ignores NaN, so NaN expert or combine outputs kept the previous worst value and passed the tolerance checks. Reject a non-finite error when it is recorded, and add a negative case for NaN and infinite outputs. Co-Authored-By: Claude Opus 5.5 Signed-off-by: ronenkat <16743404+ronenkat@users.noreply.github.com> --- tests/gpu/deepep_m2n_harness.py | 14 +++++++++++++- tests/gpu/test_deepep_m2n_gpu.py | 20 ++++++++++++++++++++ 2 files changed, 33 insertions(+), 1 deletion(-) diff --git a/tests/gpu/deepep_m2n_harness.py b/tests/gpu/deepep_m2n_harness.py index 21a835b37..9afc18264 100644 --- a/tests/gpu/deepep_m2n_harness.py +++ b/tests/gpu/deepep_m2n_harness.py @@ -17,6 +17,7 @@ import functools import gc import json +import math import os import random import socket @@ -344,6 +345,17 @@ def rel_err(actual: torch.Tensor, expected: torch.Tensor) -> float: return diff / denom if denom > 0 else diff +def update_worst_error(worst: float, key: str, value: float) -> float: + """The larger of ``worst`` and ``value``; a non-finite error fails at once. + + ``max`` ignores NaN (every comparison with it is false), so a NaN output + would otherwise keep the previous worst and pass the tolerance check. + """ + if not math.isfinite(value): + raise AssertionError(f"{key}={value} is not finite") + return max(worst, value) + + # --------------------------------------------------------------------------- # Process plumbing # --------------------------------------------------------------------------- @@ -627,7 +639,7 @@ def __init__(self, connector: Any) -> None: self.w13_ref, self.w2_ref = make_expert_weights(range(M2N_NUM_EXPERTS)) def record(self, key: str, value: float) -> None: - self.worst[key] = max(self.worst.get(key, 0.0), value) + self.worst[key] = update_worst_error(self.worst.get(key, 0.0), key, value) def check_ffn_receive( self, diff --git a/tests/gpu/test_deepep_m2n_gpu.py b/tests/gpu/test_deepep_m2n_gpu.py index 28b99bac4..33a835626 100644 --- a/tests/gpu/test_deepep_m2n_gpu.py +++ b/tests/gpu/test_deepep_m2n_gpu.py @@ -39,8 +39,10 @@ lane_cases, mixed_mode_cases, peer_exit_scenario, + rel_err, run_ranks, standard_cases, + update_worst_error, ) from tests.unit.fakes.vllm_config import M2NTopology # noqa: E402 @@ -99,6 +101,24 @@ def _cases(cases: list[Case]) -> list[dict[str, Any]]: return [asdict(case) for case in cases] +@pytest.mark.parametrize("bad_value", [float("nan"), float("inf")], ids=["nan", "inf"]) +def test_non_finite_output_fails_the_numerical_checks(bad_value): + """An all-NaN or infinite output fails, even after a finite worst error. + + Why: ``max`` ignores NaN, so without this check a broken expert or + combine output would keep the previous worst error and pass the suite. + """ + expected = torch.ones(2, 4) + worst = update_worst_error(0.0, "combine_rel_err", 1e-3) + + with pytest.raises(AssertionError, match="combine_rel_err=.* is not finite"): + update_worst_error( + worst, + "combine_rel_err", + rel_err(torch.full_like(expected, bad_value), expected), + ) + + @pytest.mark.parametrize("payload_format", PAYLOADS, ids=["bf16", "fp8"]) @pytest.mark.parametrize( ("ffn_size", "attn_size"), From e3460fecedf454cca4c711e7f67303f8ab541ba6 Mon Sep 17 00:00:00 2001 From: ronenkat <16743404+ronenkat@users.noreply.github.com> Date: Mon, 26 Oct 2026 12:51:27 +0200 Subject: [PATCH 10/32] refactor(gpu): one connector instance and lifecycle for pre-init connectors A routed-experts connector must join the AFD world and set its NCCL environment before vLLM's distributed init, which vLLM runs before it builds the runner. That requirement was spread across a factory method (create_bootstrapped_connector), both workers, and adopt_bootstrapped_connector in both runners, which closed the connector their constructor made and swapped in the worker's. Give the worker layer one owner for it. connector_lifecycle.py creates the connector with the normal factory call and runs its pre-init step (prepare_before_distributed_init, renamed from bootstrap), and checks the connector's device and role rank against vLLM's after init. The workers pass that same instance to the runner through the existing runner seam, as a keyword-only connector argument; runners create the connector themselves only when none is passed, so P2P is unchanged. Runners no longer create a throwaway connector, the factory only creates connectors again, and init_afd_connector after model load stays the common second step. Move worker device resolution to afd_plugin/distributed/device.py as resolve_worker_device_index, with vLLM imported inside the functions so the package stays CPU-safe to import. Registering vLLM's assigned physical GPU ids is now an explicit call (register_assigned_physical_gpu_ids), and resolving without it raises instead of mapping to another rank's GPU. Co-Authored-By: Claude Opus 5.5 Signed-off-by: ronenkat <16743404+ronenkat@users.noreply.github.com> --- afd_plugin/connectors/factory.py | 40 +--- .../connectors/gpu/deepep_m2n/connector.py | 19 +- .../connectors/gpu/deepep_m2n/network.py | 5 +- afd_plugin/connectors/routed_experts.py | 14 +- afd_plugin/distributed/__init__.py | 6 + afd_plugin/distributed/device.py | 80 ++++++++ afd_plugin/distributed/topology.py | 4 +- .../v1/worker/attention_model_runner_v2.py | 34 ++- afd_plugin/v1/worker/attention_worker.py | 39 ++-- afd_plugin/v1/worker/connector_device.py | 93 --------- afd_plugin/v1/worker/connector_lifecycle.py | 113 ++++++++++ afd_plugin/v1/worker/ffn_model_runner.py | 35 ++-- .../worker/ffn_routed_experts_model_runner.py | 13 +- afd_plugin/v1/worker/ffn_worker.py | 36 ++-- afd_plugin/validation.py | 6 +- docs/design/module/connector_contracts.md | 2 +- tests/unit/config/test_validation.py | 4 +- tests/unit/connectors/test_base_factory.py | 138 ------------- .../connectors/test_deepep_m2n_connector.py | 2 +- tests/unit/distributed/test_device.py | 85 ++++++++ tests/unit/distributed/test_topology.py | 2 +- tests/unit/fakes/connector.py | 16 +- tests/unit/fakes/platform.py | 22 ++ tests/unit/fakes/vllm_config.py | 23 +++ tests/unit/v1/worker/test_connector_device.py | 159 --------------- .../v1/worker/test_connector_lifecycle.py | 193 ++++++++++++++++++ tests/unit/v1/worker/test_ffn_model_runner.py | 19 +- .../test_ffn_routed_experts_model_runner.py | 32 +-- tests/unit/v1/worker/test_model_runner_v2.py | 175 +++++++++------- 29 files changed, 785 insertions(+), 624 deletions(-) create mode 100644 afd_plugin/distributed/device.py delete mode 100644 afd_plugin/v1/worker/connector_device.py create mode 100644 afd_plugin/v1/worker/connector_lifecycle.py create mode 100644 tests/unit/distributed/test_device.py create mode 100644 tests/unit/fakes/platform.py delete mode 100644 tests/unit/v1/worker/test_connector_device.py create mode 100644 tests/unit/v1/worker/test_connector_lifecycle.py diff --git a/afd_plugin/connectors/factory.py b/afd_plugin/connectors/factory.py index 420e09a5f..ab778f942 100644 --- a/afd_plugin/connectors/factory.py +++ b/afd_plugin/connectors/factory.py @@ -14,8 +14,7 @@ parse_afd_config, ) from afd_plugin.connectors.base import AFDConnectorBase, ConnectorExtraInfo -from afd_plugin.connectors.routed_experts import AFDRoutedExpertsConnector -from afd_plugin.distributed import resolve_role_rank, tp_rank_before_distributed_init +from afd_plugin.distributed import resolve_role_rank if TYPE_CHECKING: from vllm.config import VllmConfig @@ -71,43 +70,6 @@ def create_connector( role_rank, ) - @classmethod - def create_bootstrapped_connector( - cls, - rank: int, - local_rank: int, - vllm_config: VllmConfig, - afd_config: AFDConfig, - ) -> AFDRoutedExpertsConnector | None: - """Create and bootstrap a routed-experts connector ahead of vLLM. - - Workers call this before vLLM's distributed init. Returns ``None`` for - connectors that are not ``AFDRoutedExpertsConnector``; their runner - creates them later, as before. vLLM's TP group does not exist yet, so - the TP part of the role rank comes from the worker ``rank`` - (``tp_rank_before_distributed_init``); workers check it with - ``check_connector_role_rank`` after vLLM's init. ``local_rank`` is the - CUDA device index the connector uses; vLLM has not selected the - worker's device yet, so callers pass the index it will select. - """ - connector_cls = cls.get_connector_class(afd_config.connector) - if not issubclass(connector_cls, AFDRoutedExpertsConnector): - return None - connector = cls.create_connector( - rank, - local_rank, - vllm_config, - afd_config, - tp_rank=tp_rank_before_distributed_init(vllm_config, rank), - ) - assert isinstance(connector, AFDRoutedExpertsConnector) - try: - connector.bootstrap() - except BaseException: - connector.close() - raise - return connector - @classmethod def get_connector_class(cls, connector_name: str) -> type[AFDConnectorBase]: if connector_name not in cls._registry: diff --git a/afd_plugin/connectors/gpu/deepep_m2n/connector.py b/afd_plugin/connectors/gpu/deepep_m2n/connector.py index 9669301b3..5c475ce37 100644 --- a/afd_plugin/connectors/gpu/deepep_m2n/connector.py +++ b/afd_plugin/connectors/gpu/deepep_m2n/connector.py @@ -892,9 +892,9 @@ class DeepEPM2NAFDConnector(AFDRoutedExpertsConnector): once, also within one CUDA-graph capture: a capture holds about two layers' rows per lane, not every layer's. - Lifecycle: ``bootstrap`` (before vLLM's distributed init: handshake and - NCCL environment), then ``init_afd_connector`` (after weights load: - post-load checks, NCCL control group and DeepEP buffers). + Lifecycle: ``prepare_before_distributed_init`` (before vLLM's distributed + init: handshake and NCCL environment), then ``init_afd_connector`` (after + weights load: post-load checks, NCCL control group and DeepEP buffers). """ @classmethod @@ -933,7 +933,8 @@ def __init__( raise ValueError(f"FP8 payload needs hidden_size % {FP8_GROUP_SIZE} == 0") self.world_rank = self.layout.world_rank(afd_config.role, role_rank) # The FFN runs without TP and learns the Attention TP size and the - # resulting capacity in the handshake (``bootstrap``). + # resulting capacity in the handshake + # (``prepare_before_distributed_init``). self.attn_tp_size: int | None = ( None if self.is_ffn else vllm_config.parallel_config.tensor_parallel_size ) @@ -994,7 +995,7 @@ def __init__( # Lifecycle # ============================== - def bootstrap(self) -> None: + def prepare_before_distributed_init(self) -> None: """Collective: join the AFD world, check compatibility, set NCCL env. Workers call this before vLLM's distributed init, because NCCL reads @@ -1059,7 +1060,7 @@ def bootstrap(self) -> None: if self.world_rank == 0: logger.info("DeepEP M2N network plan: %s", self.network_plan.summary()) # Before any NCCL communicator in this process exists, when the worker - # bootstraps ahead of vLLM's distributed init. + # calls this ahead of vLLM's distributed init. apply_rank_environment(self.network_plan, self.config.network, self.world_rank) # DeepEP reads this when each buffer is created: with several lanes # each buffer needs its own communication stream. @@ -1084,7 +1085,7 @@ def init_afd_connector( return if local_experts is not None and not self.is_ffn: raise ValueError("only FFN ranks hold local experts") - self.bootstrap() + self.prepare_before_distributed_init() torch.cuda.set_device(self.device) _register_custom_ops() assert self.bootstrap_group is not None and self.network_plan is not None @@ -1593,7 +1594,9 @@ def update_state_from_dp_metadata(self, payload: AFDControlPayload) -> None: or connector.capacity_tokens_per_rank is None or connector.bucket_ladder is None ): - raise RuntimeError("DeepEPM2N control plane used before bootstrap") + raise RuntimeError( + "DeepEPM2N control plane used before prepare_before_distributed_init" + ) connector.stage_plans = { stage_idx: plan_stage( dp_metadata, diff --git a/afd_plugin/connectors/gpu/deepep_m2n/network.py b/afd_plugin/connectors/gpu/deepep_m2n/network.py index 33f480182..133bd77c1 100644 --- a/afd_plugin/connectors/gpu/deepep_m2n/network.py +++ b/afd_plugin/connectors/gpu/deepep_m2n/network.py @@ -489,8 +489,9 @@ def apply_rank_environment( """Process-wide NCCL/DeepEP settings for this rank, before any communicator. NCCL reads these when its first communicator in the process is created. - AFD workers therefore run the connector's ``bootstrap`` (which calls this) - before vLLM's distributed init creates its own NCCL groups; a process that + AFD workers therefore run the connector's + ``prepare_before_distributed_init`` (which calls this) before vLLM's + distributed init creates its own NCCL groups; a process that already created a communicator keeps the values it started with. """ for key, value in config.nccl_env: diff --git a/afd_plugin/connectors/routed_experts.py b/afd_plugin/connectors/routed_experts.py index 9ebd0d438..c6a8a99ac 100644 --- a/afd_plugin/connectors/routed_experts.py +++ b/afd_plugin/connectors/routed_experts.py @@ -74,15 +74,17 @@ class FFNLocalExperts: class AFDRoutedExpertsConnector(AFDConnectorBase): """A connector whose FFN side receives routed (token, local expert) rows. - Lifecycle: workers create the connector and call ``bootstrap`` before - vLLM's distributed init (it may set process-wide NCCL environment, which - NCCL reads when the process creates its first communicator); runners then - call ``init_afd_connector`` after weights load, with what they loaded. - The FFN worker selects the routed-experts runner for these connectors. + Lifecycle: the worker creates the connector and calls + ``prepare_before_distributed_init`` before vLLM's distributed init (it may + set process-wide NCCL environment, which NCCL reads when the process + creates its first communicator), then hands the same instance to its + runner, which calls ``init_afd_connector`` after weights load, with what + it loaded (``afd_plugin/v1/worker/connector_lifecycle.py``). The FFN + worker selects the routed-experts runner for these connectors. """ @abstractmethod - def bootstrap(self) -> None: + def prepare_before_distributed_init(self) -> None: """Collective: join the AFD world before vLLM's distributed init.""" raise NotImplementedError diff --git a/afd_plugin/distributed/__init__.py b/afd_plugin/distributed/__init__.py index ba822f04f..c0777a76e 100644 --- a/afd_plugin/distributed/__init__.py +++ b/afd_plugin/distributed/__init__.py @@ -2,6 +2,10 @@ # SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project """AFD distributed helper namespace.""" +from afd_plugin.distributed.device import ( + register_assigned_physical_gpu_ids, + resolve_worker_device_index, +) from afd_plugin.distributed.topology import ( AFDRankMapping, build_rank_mapping, @@ -33,7 +37,9 @@ def __getattr__(name: str): "build_rank_mapping", "create_hccl_process_group_options", "init_afd_process_group", + "register_assigned_physical_gpu_ids", "resolve_role_rank", + "resolve_worker_device_index", "subgroup_attention_block", "topology_from_config", "tp_rank_before_distributed_init", diff --git a/afd_plugin/distributed/device.py b/afd_plugin/distributed/device.py new file mode 100644 index 000000000..6adab044b --- /dev/null +++ b/afd_plugin/distributed/device.py @@ -0,0 +1,80 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Worker device resolution before vLLM's ``Worker.init_device``. + +A connector that joins the AFD world before vLLM's distributed init is created +before vLLM selects the worker's device. These helpers mirror that selection +(vLLM v0.30.0 ``vllm/v1/worker/gpu_worker.py``, ``Worker.init_device``, commit +ced6857); the worker checks the result against the device vLLM then selects. +vLLM is imported inside the functions, so importing this module stays +CPU-safe. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from vllm.config import VllmConfig + +# vLLM's GPU worker skips its local data-parallel device offset for these. +_NO_DP_DEVICE_OFFSET_EXECUTORS = ("ray", "external_launcher") + + +def register_assigned_physical_gpu_ids(vllm_config: VllmConfig) -> None: + """Register ``--device-ids`` with vLLM's platform, as ``init_device`` does. + + Side effect: sets vLLM's process-wide assigned physical GPU ids, through + which vLLM maps logical device indices. Call it before + ``resolve_worker_device_index``. vLLM's setter accepts the same ids again + when ``init_device`` registers them later; different ids raise. + """ + assigned_physical_gpu_ids = vllm_config.parallel_config.assigned_physical_gpu_ids + if assigned_physical_gpu_ids is None: + return + from vllm.platforms.interface import set_assigned_physical_gpu_ids + + set_assigned_physical_gpu_ids(assigned_physical_gpu_ids) + + +def resolve_worker_device_index(vllm_config: VllmConfig, local_rank: int) -> int: + """CUDA device index vLLM's GPU worker will select for ``local_rank``. + + Adds the local data-parallel offset, then maps the logical index through + the assigned physical ids (``--device-ids``, registered by + ``register_assigned_physical_gpu_ids``) and ``CUDA_VISIBLE_DEVICES``. + """ + from vllm.platforms import current_platform + from vllm.platforms.interface import get_assigned_physical_gpu_ids + + parallel_config = vllm_config.parallel_config + assigned_physical_gpu_ids = parallel_config.assigned_physical_gpu_ids + if ( + assigned_physical_gpu_ids is not None + and get_assigned_physical_gpu_ids() != assigned_physical_gpu_ids + ): + raise RuntimeError( + "call register_assigned_physical_gpu_ids before resolving the " + "worker device: vLLM maps device indices through those ids" + ) + if ( + parallel_config.distributed_executor_backend + not in _NO_DP_DEVICE_OFFSET_EXECUTORS + and parallel_config.data_parallel_backend != "ray" + and parallel_config.nnodes_within_dp == 1 + ): + dp_local_rank = parallel_config.data_parallel_rank_local + if dp_local_rank is None: + dp_local_rank = parallel_config.data_parallel_index + tp_pp_world_size = ( + parallel_config.pipeline_parallel_size + * parallel_config.tensor_parallel_size + ) + local_rank += dp_local_rank * tp_pp_world_size + return current_platform.logical_device_id_to_visible_device_id(local_rank) + + +__all__ = [ + "register_assigned_physical_gpu_ids", + "resolve_worker_device_index", +] diff --git a/afd_plugin/distributed/topology.py b/afd_plugin/distributed/topology.py index d2853b621..223672b5e 100644 --- a/afd_plugin/distributed/topology.py +++ b/afd_plugin/distributed/topology.py @@ -64,8 +64,8 @@ def tp_rank_before_distributed_init(vllm_config: VllmConfig, worker_rank: int) - vLLM orders ranks ``[ExternalDP, DP, PP, PCP, TP]`` with TP innermost (vLLM v0.30.0 ``initialize_model_parallel``), and AFD runs with PP = PCP = 1, so the TP rank is the worker rank modulo the TP size whether the worker - rank is DP-local or global. ``check_connector_role_rank`` compares the - result with vLLM's TP group once it exists. + rank is DP-local or global. The worker compares the resulting role rank + with vLLM's TP group once it exists (``check_worker_connector``). """ parallel_config = vllm_config.parallel_config if ( diff --git a/afd_plugin/v1/worker/attention_model_runner_v2.py b/afd_plugin/v1/worker/attention_model_runner_v2.py index 4aa62b41d..2ce70c4a6 100644 --- a/afd_plugin/v1/worker/attention_model_runner_v2.py +++ b/afd_plugin/v1/worker/attention_model_runner_v2.py @@ -367,10 +367,16 @@ class AFDAttentionModelRunnerV2(AFDMetadataProviderMixin, GPUModelRunnerV2): afd_expected_role = "attention" + # Signature: native (vllm_config, device) plus a keyword-only + # ``connector``: the instance the worker created before vLLM's distributed + # init (see connector_lifecycle.py), which the worker keeps owning if + # construction fails. ``None`` creates the connector here. def __init__( self, vllm_config: VllmConfig, device: torch.device, + *, + connector: AFDConnectorBase | None = None, ): validate_gpu_model_runner_v2_config( vllm_config, @@ -378,16 +384,17 @@ def __init__( device_type=device.type, ) super().__init__(vllm_config, device) - connector: AFDConnectorBase | None = None + owns_connector = connector is None try: self.afd_config = self.parse_config(self.vllm_config) - rank, local_rank = _resolve_world_ranks() - connector = AFDConnectorFactory.create_connector( - rank, - local_rank, - self.vllm_config, - self.afd_config, - ) + if connector is None: + rank, local_rank = _resolve_world_ranks() + connector = AFDConnectorFactory.create_connector( + rank, + local_rank, + self.vllm_config, + self.afd_config, + ) self.connector = connector # The connector rendezvous is deferred to the end of ``load_model()`` # so Attention and FFN weight loading overlap, matching the V1 @@ -405,7 +412,7 @@ def __init__( self.prof = create_afd_gpu_profiler("attention") except BaseException: try: - if connector is not None: + if owns_connector and connector is not None: connector.close() finally: super().shutdown() @@ -447,15 +454,6 @@ def install_afd_metadata_on_forward_context( def parse_config(vllm_config: VllmConfig) -> AFDConfig: return parse_afd_config(vllm_config, expected_role="attention") - def adopt_bootstrapped_connector(self, connector: AFDConnectorBase) -> None: - """Use the connector the worker bootstrapped before vLLM's init. - - Replaces the connector created at construction, which never - communicated, so the native constructor signature stays unchanged. - """ - self.connector.close() - self.connector = connector - # Patch reason: native V2 load_model has no AFD connector lifecycle. # Patch functionality: initialize the AFD connector after native weight # loading so Attention and FFN model loading overlap across roles. A diff --git a/afd_plugin/v1/worker/attention_worker.py b/afd_plugin/v1/worker/attention_worker.py index 85ebaca58..e108e7d96 100644 --- a/afd_plugin/v1/worker/attention_worker.py +++ b/afd_plugin/v1/worker/attention_worker.py @@ -4,21 +4,21 @@ from __future__ import annotations +from functools import partial + import torch from vllm.config import VllmConfig from vllm.v1.worker.gpu import model_runner as gpu_model_runner_v2 from vllm.v1.worker.gpu_worker import Worker from afd_plugin.config import parse_afd_config -from afd_plugin.connectors import AFDConnectorFactory from afd_plugin.model_executor.models.model_utils import get_afd_model_config from afd_plugin.v1.worker.attention_model_runner_v2 import ( AFDAttentionModelRunnerV2, ) -from afd_plugin.v1.worker.connector_device import ( - check_connector_device, - check_connector_role_rank, - connector_device_index, +from afd_plugin.v1.worker.connector_lifecycle import ( + check_worker_connector, + create_worker_connector, ) from afd_plugin.validation import ( assert_compatible_afd_stack, @@ -51,9 +51,10 @@ def __init__( # after an internal module import, with no injectable runner factory. # Patch functionality: validate the supported V2 contract, select the AFD model # identity before construction, and scope the module-class substitution to - # the delegated native initialization window. A routed-experts connector - # is created first on the device vLLM will select, and checked against - # that device once vLLM has selected it. + # the delegated native initialization window. A connector that must join + # the AFD world before vLLM's distributed init is created first + # (connector_lifecycle.py), handed to the runner through the seam, and + # checked against vLLM's device and ranks once vLLM has chosen them. # Signature: matches vLLM v0.30.0 Worker.init_device exactly: (self). # Upstream source: vllm/v1/worker/gpu_worker.py, Worker.init_device; # runner local-import/construct seam; commit @@ -88,30 +89,30 @@ def init_device(self): self.vllm_config.model_config = afd_model_config self.model_config = afd_model_config - # A routed-experts connector joins the AFD world here, before vLLM's - # distributed init creates NCCL groups, so its NCCL environment applies; - # the FFN worker does the same at the same point. - connector = AFDConnectorFactory.create_bootstrapped_connector( - self.rank, - connector_device_index(self.vllm_config, self.local_rank), + # Before vLLM's distributed init creates NCCL groups, so the + # connector's NCCL environment applies; the FFN worker does the same. + connector = create_worker_connector( self.vllm_config, parse_afd_config(self.vllm_config, expected_role="attention"), + rank=self.rank, + local_rank=self.local_rank, ) native_runner_cls = gpu_model_runner_v2.GPUModelRunner - gpu_model_runner_v2.GPUModelRunner = AFDAttentionModelRunnerV2 + gpu_model_runner_v2.GPUModelRunner = ( + AFDAttentionModelRunnerV2 + if connector is None + else partial(AFDAttentionModelRunnerV2, connector=connector) + ) try: super().init_device() if connector is not None: - check_connector_device(connector, self.device) - check_connector_role_rank(connector, self.vllm_config) + check_worker_connector(connector, self.vllm_config, self.device) except BaseException: if connector is not None: connector.close() raise finally: gpu_model_runner_v2.GPUModelRunner = native_runner_cls - if connector is not None: - self.model_runner.adopt_bootstrapped_connector(connector) # ### PATCH END: select and directly construct the AFD runner. torch.accelerator.empty_cache() diff --git a/afd_plugin/v1/worker/connector_device.py b/afd_plugin/v1/worker/connector_device.py deleted file mode 100644 index ae8ed2b1a..000000000 --- a/afd_plugin/v1/worker/connector_device.py +++ /dev/null @@ -1,93 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project -"""CUDA device for connectors created before vLLM's ``Worker.init_device``. - -A routed-experts connector bootstraps before vLLM's distributed init, so it is -created before vLLM selects the worker's device. ``connector_device_index`` -mirrors that selection; ``check_connector_device`` compares the result with -the device vLLM chose, so a drift after a vLLM upgrade fails at startup. -``check_connector_role_rank`` does the same for the role rank, whose TP part -is derived from the worker rank before vLLM's TP group exists. -""" - -from __future__ import annotations - -import torch -from vllm.config import VllmConfig -from vllm.platforms import current_platform -from vllm.platforms.interface import set_assigned_physical_gpu_ids - -from afd_plugin.connectors import AFDConnectorBase -from afd_plugin.distributed import resolve_role_rank - -# vLLM's GPU worker skips its local data-parallel device offset for these. -_NO_DP_DEVICE_OFFSET_EXECUTORS = ("ray", "external_launcher") - - -def connector_device_index(vllm_config: VllmConfig, local_rank: int) -> int: - """CUDA device index vLLM's GPU worker will select for ``local_rank``. - - Mirrors ``Worker.init_device`` (vLLM v0.30.0, ``vllm/v1/worker/ - gpu_worker.py``, commit ced6857): add the local data-parallel offset, - then map the logical index through ``assigned_physical_gpu_ids`` - (``--device-ids``) and ``CUDA_VISIBLE_DEVICES``. Registering the assigned - ids here is what vLLM does next; its setter accepts the same value again. - """ - parallel_config = vllm_config.parallel_config - if ( - parallel_config.distributed_executor_backend - not in _NO_DP_DEVICE_OFFSET_EXECUTORS - and parallel_config.data_parallel_backend != "ray" - and parallel_config.nnodes_within_dp == 1 - ): - dp_local_rank = parallel_config.data_parallel_rank_local - if dp_local_rank is None: - dp_local_rank = parallel_config.data_parallel_index - tp_pp_world_size = ( - parallel_config.pipeline_parallel_size - * parallel_config.tensor_parallel_size - ) - local_rank += dp_local_rank * tp_pp_world_size - assigned_physical_gpu_ids = parallel_config.assigned_physical_gpu_ids - if assigned_physical_gpu_ids is not None: - set_assigned_physical_gpu_ids(assigned_physical_gpu_ids) - return current_platform.logical_device_id_to_visible_device_id(local_rank) - - -def check_connector_device( - connector: AFDConnectorBase, - worker_device: torch.device, -) -> None: - """Fail when a connector does not use the device vLLM gave the worker.""" - if connector.local_rank != worker_device.index: - raise RuntimeError( - f"{type(connector).__name__} uses cuda:{connector.local_rank} but " - f"vLLM placed this worker on {worker_device}; the device selection " - "in afd_plugin/v1/worker/connector_device.py no longer matches " - "vLLM's Worker.init_device", - ) - - -def check_connector_role_rank( - connector: AFDConnectorBase, - vllm_config: VllmConfig, -) -> None: - """Fail when a connector's role rank differs from vLLM's distributed ranks. - - Call after vLLM's distributed init, which creates the TP group. - """ - role_rank = resolve_role_rank(vllm_config, connector.afd_config) - if connector.role_rank != role_rank: - raise RuntimeError( - f"{type(connector).__name__} has role rank {connector.role_rank} but " - f"vLLM's distributed ranks give {role_rank}; the TP rank derived " - "before distributed init (tp_rank_before_distributed_init) no longer " - "matches vLLM's rank layout", - ) - - -__all__ = [ - "check_connector_device", - "check_connector_role_rank", - "connector_device_index", -] diff --git a/afd_plugin/v1/worker/connector_lifecycle.py b/afd_plugin/v1/worker/connector_lifecycle.py new file mode 100644 index 000000000..9ec01b935 --- /dev/null +++ b/afd_plugin/v1/worker/connector_lifecycle.py @@ -0,0 +1,113 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Connector lifecycle for GPU workers whose connector joins before vLLM. + +Every connector follows one lifecycle: create, then ``init_afd_connector`` +once the runner has loaded the model. A routed-experts connector +(``AFDRoutedExpertsConnector``) also has a step that must run before vLLM's +distributed init: it joins the AFD world and applies its NCCL environment, +which NCCL reads once per process, when the first communicator is created. +vLLM creates its communicators in ``Worker.init_device`` before it builds the +runner, so the worker owns that step: + +1. ``create_worker_connector`` creates the connector and runs + ``prepare_before_distributed_init``. +2. The worker hands that same instance to the runner it constructs. +3. ``check_worker_connector`` compares what was derived before vLLM's init + (device, TP rank) with what vLLM then chose. +4. The runner calls ``init_afd_connector`` after loading the model. + +Other connectors are created by their runner, as before. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from afd_plugin.config import AFDConfig +from afd_plugin.connectors import ( + AFDConnectorBase, + AFDConnectorFactory, + AFDRoutedExpertsConnector, +) +from afd_plugin.distributed import ( + register_assigned_physical_gpu_ids, + resolve_role_rank, + resolve_worker_device_index, + tp_rank_before_distributed_init, +) + +if TYPE_CHECKING: + import torch + from vllm.config import VllmConfig + + +def create_worker_connector( + vllm_config: VllmConfig, + afd_config: AFDConfig, + *, + rank: int, + local_rank: int, +) -> AFDRoutedExpertsConnector | None: + """Create a connector that must join the AFD world before vLLM's init. + + Call before vLLM's distributed init. Returns ``None`` for connectors that + are not ``AFDRoutedExpertsConnector``; their runner creates them. + Otherwise the connector gets the device vLLM will select and a role rank + whose TP part comes from the worker ``rank`` (vLLM's TP group does not + exist yet), and has run ``prepare_before_distributed_init``. The caller + owns it until it hands it to the runner. + + Side effect: registers vLLM's assigned physical GPU ids + (``register_assigned_physical_gpu_ids``). + """ + connector_cls = AFDConnectorFactory.get_connector_class(afd_config.connector) + if not issubclass(connector_cls, AFDRoutedExpertsConnector): + return None + register_assigned_physical_gpu_ids(vllm_config) + connector = AFDConnectorFactory.create_connector( + rank, + resolve_worker_device_index(vllm_config, local_rank), + vllm_config, + afd_config, + tp_rank=tp_rank_before_distributed_init(vllm_config, rank), + ) + assert isinstance(connector, AFDRoutedExpertsConnector) + try: + connector.prepare_before_distributed_init() + except BaseException: + connector.close() + raise + return connector + + +def check_worker_connector( + connector: AFDConnectorBase, + vllm_config: VllmConfig, + worker_device: torch.device, +) -> None: + """Fail when a connector created before vLLM's init disagrees with vLLM. + + Call after vLLM's distributed init. The connector must use the device vLLM + gave the worker, and its role rank must match vLLM's distributed ranks, so + a vLLM upgrade that changes device selection or rank order fails at + startup instead of moving traffic onto another rank's GPU or tokens. + """ + if connector.local_rank != worker_device.index: + raise RuntimeError( + f"{type(connector).__name__} uses cuda:{connector.local_rank} but " + f"vLLM placed this worker on {worker_device}; the device selection " + "in afd_plugin/distributed/device.py no longer matches vLLM's " + "Worker.init_device", + ) + role_rank = resolve_role_rank(vllm_config, connector.afd_config) + if connector.role_rank != role_rank: + raise RuntimeError( + f"{type(connector).__name__} has role rank {connector.role_rank} but " + f"vLLM's distributed ranks give {role_rank}; the TP rank derived " + "before distributed init (tp_rank_before_distributed_init) no longer " + "matches vLLM's rank layout", + ) + + +__all__ = ["check_worker_connector", "create_worker_connector"] diff --git a/afd_plugin/v1/worker/ffn_model_runner.py b/afd_plugin/v1/worker/ffn_model_runner.py index 4ca5b3717..918fb746c 100644 --- a/afd_plugin/v1/worker/ffn_model_runner.py +++ b/afd_plugin/v1/worker/ffn_model_runner.py @@ -61,7 +61,15 @@ class GPUFFNModelRunner: afd_expected_role = "ffn" - def __init__(self, vllm_config: VllmConfig, device: torch.device) -> None: + # ``connector``: the instance the worker created before vLLM's distributed + # init (see connector_lifecycle.py); ``None`` creates the connector here. + def __init__( + self, + vllm_config: VllmConfig, + device: torch.device, + *, + connector: AFDConnectorBase | None = None, + ) -> None: self.vllm_config = vllm_config self.model_config = vllm_config.model_config self.load_config = vllm_config.load_config @@ -77,13 +85,15 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device) -> None: vllm_config, role="ffn", ) - rank, local_rank = _resolve_world_ranks() - self.connector = AFDConnectorFactory.create_connector( - rank, - local_rank, - vllm_config, - self.afd_config, - ) + if connector is None: + rank, local_rank = _resolve_world_ranks() + connector = AFDConnectorFactory.create_connector( + rank, + local_rank, + vllm_config, + self.afd_config, + ) + self.connector = connector # TODO: Async GPU connector will be supported in the future assert self.connector.control_plane is not None, ( "GPU model runner only supports control-plane-driven connectors" @@ -109,15 +119,6 @@ def get_model(self) -> Any: def initialize_afd_connector(self) -> None: self.connector.init_afd_connector() - def adopt_bootstrapped_connector(self, connector: AFDConnectorBase) -> None: - """Use the connector the worker bootstrapped before vLLM's init. - - Replaces the connector created at construction, which never - communicated. - """ - self.connector.close() - self.connector = connector - def load_model(self, *, load_dummy_weights: bool = False, **kwargs: Any) -> None: """Load the vLLM model.""" diff --git a/afd_plugin/v1/worker/ffn_routed_experts_model_runner.py b/afd_plugin/v1/worker/ffn_routed_experts_model_runner.py index 09a0220a4..87c8a58aa 100644 --- a/afd_plugin/v1/worker/ffn_routed_experts_model_runner.py +++ b/afd_plugin/v1/worker/ffn_routed_experts_model_runner.py @@ -116,16 +116,21 @@ class GPUFFNRoutedExpertsModelRunner(GPUFFNModelRunner): (DeepEP M2N). """ - def __init__(self, vllm_config: VllmConfig, device: torch.device) -> None: - super().__init__(vllm_config, device) + def __init__( + self, + vllm_config: VllmConfig, + device: torch.device, + *, + connector: AFDConnectorBase | None = None, + ) -> None: + super().__init__(vllm_config, device, connector=connector) _require_routed(self.connector) self.moe_layer_indices: tuple[int, ...] = () self.local_experts: dict[int, FFNLocalExperts] = {} @property def routed_connector(self) -> AFDRoutedExpertsConnector: - """The connector; the worker may replace it with the one it - bootstrapped, so it is checked where it is used.""" + """The connector, typed as the routed-experts contract it must meet.""" return _require_routed(self.connector) def load_model(self, *, load_dummy_weights: bool = False, **kwargs: Any) -> None: diff --git a/afd_plugin/v1/worker/ffn_worker.py b/afd_plugin/v1/worker/ffn_worker.py index a85026353..4cd249bc1 100644 --- a/afd_plugin/v1/worker/ffn_worker.py +++ b/afd_plugin/v1/worker/ffn_worker.py @@ -6,6 +6,7 @@ import logging import threading +from functools import partial from typing import TYPE_CHECKING import torch @@ -17,10 +18,9 @@ from afd_plugin.config import parse_afd_config from afd_plugin.connectors import AFDConnectorFactory, AFDRoutedExpertsConnector from afd_plugin.model_executor.models.model_utils import get_afd_model_config -from afd_plugin.v1.worker.connector_device import ( - check_connector_device, - check_connector_role_rank, - connector_device_index, +from afd_plugin.v1.worker.connector_lifecycle import ( + check_worker_connector, + create_worker_connector, ) from afd_plugin.v1.worker.ffn_model_runner import GPUFFNModelRunner from afd_plugin.validation import ( @@ -70,10 +70,10 @@ def __init__( # Patch functionality: replace only the V2 construction seam with the # connector-driven FFN runner: GPUFFNRoutedExpertsModelRunner (imported # only then) for an AFDRoutedExpertsConnector, else GPUFFNModelRunner. A - # routed-experts connector is created and bootstrapped first, so its NCCL - # environment precedes vLLM's NCCL groups, and the runner adopts it after - # construction. The connector gets the device vLLM will select, checked - # once vLLM has selected it. + # connector that must join the AFD world before vLLM's distributed init is + # created first (connector_lifecycle.py), so its NCCL environment precedes + # vLLM's NCCL groups, handed to the runner through the seam, and checked + # against vLLM's device and ranks once vLLM has chosen them. # Signature: matches vLLM v0.30.0 Worker.init_device exactly: (self). # Upstream source: vllm/v1/worker/gpu_worker.py, Worker.init_device; # ced6857afa0ea7b2e3f0846a62e1394e90f15607. The 0.30.0 init_device still @@ -114,29 +114,31 @@ def init_device(self): ) runner_cls = GPUFFNRoutedExpertsModelRunner - # Before bootstrap: a rejected backend fails before any collective. + # Before the connector joins the AFD world: a rejected backend + # fails before any collective. select_routed_moe_backend(self.vllm_config) - connector = AFDConnectorFactory.create_bootstrapped_connector( - self.rank, - connector_device_index(self.vllm_config, self.local_rank), + connector = create_worker_connector( self.vllm_config, afd_config, + rank=self.rank, + local_rank=self.local_rank, ) native_runner_cls = gpu_model_runner_v2.GPUModelRunner - gpu_model_runner_v2.GPUModelRunner = runner_cls + gpu_model_runner_v2.GPUModelRunner = ( + runner_cls + if connector is None + else partial(runner_cls, connector=connector) + ) try: super().init_device() if connector is not None: - check_connector_device(connector, self.device) - check_connector_role_rank(connector, self.vllm_config) + check_worker_connector(connector, self.vllm_config, self.device) except BaseException: if connector is not None: connector.close() raise finally: gpu_model_runner_v2.GPUModelRunner = native_runner_cls - if connector is not None: - self.model_runner.adopt_bootstrapped_connector(connector) # ### PATCH END: construct the FFN runner through the V2 seam. torch.accelerator.empty_cache() diff --git a/afd_plugin/validation.py b/afd_plugin/validation.py index eab1ccced..27ba28b5a 100644 --- a/afd_plugin/validation.py +++ b/afd_plugin/validation.py @@ -109,14 +109,14 @@ def validate_gpu_model_runner_v2_config( ) parallel = vllm_config.parallel_config - # DeepEP M2N bootstraps before vLLM's distributed init and derives the TP - # part of the role rank from the worker rank. Attention may run TP > 1: + # DeepEP M2N joins the AFD world before vLLM's distributed init and derives + # the TP part of the role rank from the worker rank. Attention may run TP > 1: # each TP rank dispatches its own slice of the tokens. FFN experts are not # TP-sharded; that and the other deployment checks (expert placement, # all2all backend, weight format) run in the connector handshake so that # every rank fails together. It routes on Attention, so it requires the # gate there, and a model whose Attention side sends the top-k. Workers - # validate before the connector bootstrap, so both roles fail before any + # validate before the connector joins the AFD world, so both roles fail before any # collective instead of one role timing out. if afd_config.connector in DEEPEP_M2N_CONNECTORS: if expected_role == "ffn" and parallel.tensor_parallel_size != 1: diff --git a/docs/design/module/connector_contracts.md b/docs/design/module/connector_contracts.md index 450e9080f..27c5c809b 100644 --- a/docs/design/module/connector_contracts.md +++ b/docs/design/module/connector_contracts.md @@ -22,7 +22,7 @@ validation_paths: - "tests/e2e/models/deepseek_v2_lite/test_deepseek_v2_lite.py" - "tests/e2e/models/deepseek_v2_lite/test_async_cam_npu.py" - "tests/unit/distributed/**" - - "tests/unit/v1/worker/test_connector_device.py" + - "tests/unit/v1/worker/test_connector_lifecycle.py" - "tests/gpu/**" upstream_refs: - "vLLM vllm.forward_context.DPMetadata" diff --git a/tests/unit/config/test_validation.py b/tests/unit/config/test_validation.py index de1ee4e22..665930679 100644 --- a/tests/unit/config/test_validation.py +++ b/tests/unit/config/test_validation.py @@ -209,8 +209,8 @@ def test_routed_transport_requires_model_that_routes_on_attention( mrv2_config, expected_role=role, device_type="cuda" ) return - # Workers validate before the connector bootstrap; rejecting a model must - # not need the model runtime. + # Workers validate before the connector joins the AFD world; rejecting a + # model must not need the model runtime. monkeypatch.setitem(sys.modules, MODEL_UTILS_MODULE, None) with pytest.raises(RuntimeError, match="DeepEPM2NAFDConnector supports only"): validate_gpu_model_runner_v2_config( diff --git a/tests/unit/connectors/test_base_factory.py b/tests/unit/connectors/test_base_factory.py index fa932a3b7..cc12ee3e4 100644 --- a/tests/unit/connectors/test_base_factory.py +++ b/tests/unit/connectors/test_base_factory.py @@ -3,7 +3,6 @@ from __future__ import annotations -import sys from types import SimpleNamespace import pytest @@ -16,7 +15,6 @@ AFDA2FTransferPayload, AFDConnectorBase, AFDConnectorFactory, - AFDRoutedExpertsConnector, AFDTransferContext, AFDTransferMetadata, ) @@ -193,139 +191,3 @@ def test_factory_resolves_role_rank_before_connector_construction(monkeypatch): ) assert connector.role_rank == 1 - - -class _BootstrappingConnector(_MinimalConnector, AFDRoutedExpertsConnector): - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) - self.events: list[str] = [] - - def bootstrap(self): - self.events.append("bootstrap") - if self.vllm_config.fail_bootstrap: - raise RuntimeError("bootstrap failed") - - def init_afd_connector(self, *, moe_layer_indices=None, local_experts=None): - return None - - def run_experts(self, payload, weights): - return payload.hidden_states - - def synchronize_lanes(self): - return None - - def close(self): - self.events.append("close") - - -def _register(monkeypatch, name, connector_cls): - monkeypatch.setitem(AFDConnectorFactory._registry, name, lambda: connector_cls) - - -def _single_rank_config( - *, fail_bootstrap=False, data_parallel_size=1, data_parallel_rank=0, tp_size=1 -): - return SimpleNamespace( - additional_config={}, - fail_bootstrap=fail_bootstrap, - parallel_config=SimpleNamespace( - data_parallel_size=data_parallel_size, - data_parallel_rank=data_parallel_rank, - pipeline_parallel_size=1, - prefill_context_parallel_size=1, - tensor_parallel_size=tp_size, - ), - ) - - -def test_create_bootstrapped_connector_skips_other_connectors(monkeypatch): - """A connector that is not an ``AFDRoutedExpertsConnector`` is not - created before vLLM's distributed init; its runner creates it later. - - Why: P2P and the NPU connectors keep their existing bring-up order. - """ - _register(monkeypatch, "MinimalConnector", _MinimalConnector) - - connector = AFDConnectorFactory.create_bootstrapped_connector( - 0, - 0, - _single_rank_config(), - AFDConfig(connector="MinimalConnector", role="ffn"), - ) - - assert connector is None - - -def test_create_bootstrapped_connector_creates_and_bootstraps(monkeypatch): - """A routed-experts connector is created with the worker's ranks and - bootstrapped. - - Why: workers call this before vLLM's distributed init so the connector's - NCCL environment precedes vLLM's NCCL groups. - """ - _register(monkeypatch, "BootstrappingConnector", _BootstrappingConnector) - - connector = AFDConnectorFactory.create_bootstrapped_connector( - 3, - 1, - _single_rank_config(), - AFDConfig(connector="BootstrappingConnector", role="ffn"), - ) - - assert isinstance(connector, _BootstrappingConnector) - assert (connector.rank, connector.local_rank, connector.role_rank) == (3, 1, 0) - assert connector.events == ["bootstrap"] - - -def test_create_bootstrapped_connector_derives_tp_rank_from_worker_rank( - monkeypatch, -): - """With TP > 1 the role rank's TP part comes from the worker rank. - - Why: vLLM's TP group does not exist before its distributed init; reading - it there would fail. TP is vLLM's innermost rank dimension, so the worker - rank modulo the TP size is the TP rank. - """ - _register(monkeypatch, "BootstrappingConnector", _BootstrappingConnector) - monkeypatch.setitem(sys.modules, "vllm.distributed.parallel_state", None) - - connector = AFDConnectorFactory.create_bootstrapped_connector( - 1, - 1, - _single_rank_config(data_parallel_size=2, data_parallel_rank=1, tp_size=2), - AFDConfig( - connector="BootstrappingConnector", - role="attention", - num_attention_ranks=4, - ), - ) - - assert isinstance(connector, _BootstrappingConnector) - # DP rank 1, TP rank 1 of 2. - assert connector.role_rank == 3 - - -def test_create_bootstrapped_connector_closes_on_bootstrap_failure(monkeypatch): - """A failed bootstrap closes the connector before the error propagates. - - Why: the worker never receives the connector, so nobody else closes it. - """ - _register(monkeypatch, "BootstrappingConnector", _BootstrappingConnector) - created: list[_BootstrappingConnector] = [] - original_init = _BootstrappingConnector.__init__ - - def recording_init(self, *args, **kwargs): - original_init(self, *args, **kwargs) - created.append(self) - - monkeypatch.setattr(_BootstrappingConnector, "__init__", recording_init) - - with pytest.raises(RuntimeError, match="bootstrap failed"): - AFDConnectorFactory.create_bootstrapped_connector( - 0, - 0, - _single_rank_config(fail_bootstrap=True), - AFDConfig(connector="BootstrappingConnector", role="ffn"), - ) - - assert created[0].events == ["bootstrap", "close"] diff --git a/tests/unit/connectors/test_deepep_m2n_connector.py b/tests/unit/connectors/test_deepep_m2n_connector.py index 980f82a1e..7444ba8b5 100644 --- a/tests/unit/connectors/test_deepep_m2n_connector.py +++ b/tests/unit/connectors/test_deepep_m2n_connector.py @@ -822,7 +822,7 @@ def test_control_plane_requires_the_handshake_on_ffn(): is_warmup=False, ) - with pytest.raises(RuntimeError, match="before bootstrap"): + with pytest.raises(RuntimeError, match="before prepare_before_distributed_init"): connector.control_plane.update_state_from_dp_metadata(payload) diff --git a/tests/unit/distributed/test_device.py b/tests/unit/distributed/test_device.py new file mode 100644 index 000000000..26200d2f7 --- /dev/null +++ b/tests/unit/distributed/test_device.py @@ -0,0 +1,85 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Worker device resolution matches ``Worker.init_device`` before it runs.""" + +from __future__ import annotations + +import pytest + +pytest.importorskip("vllm") + +from afd_plugin.distributed import ( # noqa: E402 + register_assigned_physical_gpu_ids, + resolve_worker_device_index, +) +from tests.unit.fakes.platform import use_cuda_like_platform # noqa: E402 +from tests.unit.fakes.vllm_config import gpu_worker_vllm_config # noqa: E402 + + +@pytest.fixture(autouse=True) +def _cuda_like_platform(monkeypatch): + use_cuda_like_platform(monkeypatch) + + +def _resolve(config, local_rank): + register_assigned_physical_gpu_ids(config) + return resolve_worker_device_index(config, local_rank) + + +@pytest.mark.parametrize( + ("executor_backend", "expected_index"), + [("mp", 1), ("ray", 0), ("external_launcher", 0)], +) +def test_adds_local_dp_offset_like_vllm(executor_backend, expected_index): + """The local DP rank offsets the device, except where vLLM skips it. + + Why: vLLM adds the offset only inside init_device, after the connector + joined the AFD world; without it every DP rank's connector used GPU 0. + """ + config = gpu_worker_vllm_config( + data_parallel_rank_local=1, + data_parallel_index=1, + distributed_executor_backend=executor_backend, + ) + + assert _resolve(config, 0) == expected_index + + +def test_maps_assigned_device_ids_without_visible_devices(): + """``--device-ids 4,5`` without CUDA_VISIBLE_DEVICES places rank 0 on cuda:4. + + Why: vLLM maps its logical index through the assigned physical ids, so a + connector using the logical index would run on another rank's GPU. + """ + config = gpu_worker_vllm_config(assigned_physical_gpu_ids=[4, 5]) + + assert _resolve(config, 0) == 4 + assert _resolve(config, 1) == 5 + + +def test_maps_assigned_device_ids_through_visible_devices(monkeypatch): + """With CUDA_VISIBLE_DEVICES the assigned ids map to visible ordinals.""" + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "4,5") + config = gpu_worker_vllm_config(assigned_physical_gpu_ids=[5, 4]) + + assert _resolve(config, 0) == 1 + assert _resolve(config, 1) == 0 + + +def test_without_assigned_ids_the_index_is_the_local_rank(monkeypatch): + """The default mp launch keeps the logical index as the device index.""" + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "2,3") + + assert _resolve(gpu_worker_vllm_config(), 1) == 1 + + +def test_resolving_before_registering_assigned_ids_fails(): + """Resolving with ``--device-ids`` set but not registered raises. + + Why: vLLM would map the index without those ids and place the connector + on another rank's GPU; the side effect must be explicit and come first. + """ + config = gpu_worker_vllm_config(assigned_physical_gpu_ids=[4, 5]) + + with pytest.raises(RuntimeError, match="register_assigned_physical_gpu_ids"): + resolve_worker_device_index(config, 0) diff --git a/tests/unit/distributed/test_topology.py b/tests/unit/distributed/test_topology.py index 0e6b7aa55..bdfcfca89 100644 --- a/tests/unit/distributed/test_topology.py +++ b/tests/unit/distributed/test_topology.py @@ -158,7 +158,7 @@ def test_tp_rank_before_distributed_init_is_worker_rank_mod_tp( """TP is vLLM's innermost rank dimension, so the worker rank gives it. Why: this holds for DP-local and global worker ranks alike; - ``check_connector_role_rank`` re-checks it after vLLM's init. + ``check_worker_connector`` re-checks it after vLLM's init. """ assert ( tp_rank_before_distributed_init(_pre_init_config(tp_size=tp_size), worker_rank) diff --git a/tests/unit/fakes/connector.py b/tests/unit/fakes/connector.py index 92b9c8813..485e7de43 100644 --- a/tests/unit/fakes/connector.py +++ b/tests/unit/fakes/connector.py @@ -2,7 +2,8 @@ # SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project """A recording ``AFDRoutedExpertsConnector`` for worker and runner tests. -Constructed without ``AFDConnectorBase.__init__`` (no vLLM config needed). +Constructed without arguments it skips ``AFDConnectorBase.__init__`` (no vLLM +config needed); the factory constructs it with the base arguments. ``events`` records the lifecycle and data-path calls in order. """ @@ -19,7 +20,12 @@ class FakeRoutedConnector(AFDRoutedExpertsConnector): - def __init__(self) -> None: + # Raised by prepare_before_distributed_init when set. + prepare_error: BaseException | None = None + + def __init__(self, *base_args: Any) -> None: + if base_args: + super().__init__(*base_args) self.events: list[tuple] = [] self.control_updates: list[Any] = [] self.init_kwargs: dict[str, Any] | None = None @@ -37,8 +43,10 @@ def is_initialized(self) -> bool: def close(self) -> None: self.events.append(("close",)) - def bootstrap(self) -> None: - self.events.append(("bootstrap",)) + def prepare_before_distributed_init(self) -> None: + self.events.append(("prepare",)) + if self.prepare_error is not None: + raise self.prepare_error def init_afd_connector(self, *, moe_layer_indices=None, local_experts=None): self.init_kwargs = { diff --git a/tests/unit/fakes/platform.py b/tests/unit/fakes/platform.py new file mode 100644 index 000000000..4ad40fe35 --- /dev/null +++ b/tests/unit/fakes/platform.py @@ -0,0 +1,22 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""A CUDA-like vLLM platform for device-selection tests on hosts without GPUs.""" + +from __future__ import annotations + +import pytest + + +def use_cuda_like_platform(monkeypatch: pytest.MonkeyPatch) -> None: + """Map device ids like vLLM's CUDA platform: no assigned physical ids and + no ``CUDA_VISIBLE_DEVICES`` until a test sets them.""" + # Deferred: this fakes module must import on CPU-only hosts without vLLM. + from vllm.platforms import interface + + class CudaLikePlatform(interface.Platform): + device_control_env_var = "CUDA_VISIBLE_DEVICES" + + monkeypatch.setattr("vllm.platforms.current_platform", CudaLikePlatform) + # vLLM keeps the assigned ids process-wide; start every case unset. + monkeypatch.setattr(interface, "_assigned_physical_gpu_ids", None) + monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False) diff --git a/tests/unit/fakes/vllm_config.py b/tests/unit/fakes/vllm_config.py index 0f0f467f7..66f2e92b3 100644 --- a/tests/unit/fakes/vllm_config.py +++ b/tests/unit/fakes/vllm_config.py @@ -54,6 +54,29 @@ def npu_parallel_config(**overrides): return SimpleNamespace(**values) +def gpu_worker_vllm_config(**overrides): + """The ``parallel_config`` fields that worker device and role-rank + resolution read, for one GPU worker (mp executor, DP1 x TP1 by default). + """ + values = { + "distributed_executor_backend": "mp", + "data_parallel_backend": "mp", + "nnodes_within_dp": 1, + "data_parallel_size": 1, + "data_parallel_rank": 0, + "data_parallel_rank_local": 0, + "data_parallel_index": 0, + "pipeline_parallel_size": 1, + "prefill_context_parallel_size": 1, + "tensor_parallel_size": 1, + "assigned_physical_gpu_ids": None, + } + values.update(overrides) + return SimpleNamespace( + additional_config={}, parallel_config=SimpleNamespace(**values) + ) + + def npu_vllm_config( *, role="attention", diff --git a/tests/unit/v1/worker/test_connector_device.py b/tests/unit/v1/worker/test_connector_device.py deleted file mode 100644 index 76ee793b9..000000000 --- a/tests/unit/v1/worker/test_connector_device.py +++ /dev/null @@ -1,159 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project -"""A connector created before ``Worker.init_device`` gets vLLM's device.""" - -from __future__ import annotations - -from types import SimpleNamespace - -import pytest - -torch = pytest.importorskip("torch") -pytest.importorskip("vllm") - -from vllm.platforms import interface as platform_interface # noqa: E402 - -from afd_plugin.config import AFDConfig # noqa: E402 -from afd_plugin.v1.worker import connector_device # noqa: E402 - - -class _CudaLikePlatform(platform_interface.Platform): - device_control_env_var = "CUDA_VISIBLE_DEVICES" - - -@pytest.fixture(autouse=True) -def _cuda_like_platform(monkeypatch): - monkeypatch.setattr(connector_device, "current_platform", _CudaLikePlatform) - # vLLM keeps the assigned ids process-wide; start every case unset. - monkeypatch.setattr(platform_interface, "_assigned_physical_gpu_ids", None) - monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False) - - -def _config( - *, - data_parallel_rank_local=0, - distributed_executor_backend="mp", - assigned_physical_gpu_ids=None, -): - return SimpleNamespace( - parallel_config=SimpleNamespace( - distributed_executor_backend=distributed_executor_backend, - data_parallel_backend="mp", - nnodes_within_dp=1, - data_parallel_rank_local=data_parallel_rank_local, - data_parallel_index=data_parallel_rank_local, - pipeline_parallel_size=1, - tensor_parallel_size=1, - assigned_physical_gpu_ids=assigned_physical_gpu_ids, - ), - ) - - -@pytest.mark.parametrize( - ("executor_backend", "expected_index"), - [("mp", 1), ("ray", 0), ("external_launcher", 0)], -) -def test_adds_local_dp_offset_like_vllm(executor_backend, expected_index): - """The local DP rank offsets the device, except where vLLM skips it. - - Why: vLLM adds the offset only inside init_device, after the connector - bootstrap; without it every DP rank's connector used GPU 0. - """ - config = _config( - data_parallel_rank_local=1, - distributed_executor_backend=executor_backend, - ) - - assert connector_device.connector_device_index(config, 0) == expected_index - - -def test_maps_assigned_device_ids_without_visible_devices(): - """``--device-ids 4,5`` without CUDA_VISIBLE_DEVICES places rank 0 on cuda:4. - - Why: vLLM maps its logical index through the assigned physical ids, so a - connector using the logical index would run on another rank's GPU. - """ - config = _config(assigned_physical_gpu_ids=[4, 5]) - - assert connector_device.connector_device_index(config, 0) == 4 - assert connector_device.connector_device_index(config, 1) == 5 - - -def test_maps_assigned_device_ids_through_visible_devices(monkeypatch): - """With CUDA_VISIBLE_DEVICES the assigned ids map to visible ordinals.""" - monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "4,5") - config = _config(assigned_physical_gpu_ids=[5, 4]) - - assert connector_device.connector_device_index(config, 0) == 1 - assert connector_device.connector_device_index(config, 1) == 0 - - -def test_without_assigned_ids_the_index_is_the_local_rank(monkeypatch): - """The default mp launch keeps the logical index as the device index.""" - monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "2,3") - - assert connector_device.connector_device_index(_config(), 1) == 1 - - -def test_check_accepts_the_worker_device(): - connector_device.check_connector_device( - SimpleNamespace(local_rank=1), - torch.device("cuda", 1), - ) - - -def test_check_rejects_another_device(): - """A connector on another GPU than the worker's fails with both devices. - - Why: a vLLM upgrade that changes device selection must fail at startup - instead of moving traffic onto another rank's GPU. - """ - with pytest.raises(RuntimeError, match="uses cuda:0 but vLLM placed"): - connector_device.check_connector_device( - SimpleNamespace(local_rank=0), - torch.device("cuda", 4), - ) - - -def _role_rank_config(*, tp_size: int): - return SimpleNamespace( - parallel_config=SimpleNamespace( - data_parallel_size=1, - data_parallel_rank=0, - prefill_context_parallel_size=1, - tensor_parallel_size=tp_size, - ), - ) - - -def test_role_rank_check_accepts_the_distributed_role_rank(monkeypatch): - monkeypatch.setattr( - "vllm.distributed.parallel_state.get_tensor_model_parallel_rank", lambda: 1 - ) - connector = SimpleNamespace( - role_rank=1, - afd_config=AFDConfig(role="attention", num_attention_ranks=2), - ) - - connector_device.check_connector_role_rank(connector, _role_rank_config(tp_size=2)) - - -def test_role_rank_check_rejects_a_different_role_rank(monkeypatch): - """A connector whose early role rank differs from vLLM's fails startup. - - Why: the TP part of the role rank is derived from the worker rank before - vLLM's TP group exists; a vLLM change to its rank order must fail here - instead of routing another rank's tokens. - """ - monkeypatch.setattr( - "vllm.distributed.parallel_state.get_tensor_model_parallel_rank", lambda: 0 - ) - connector = SimpleNamespace( - role_rank=1, - afd_config=AFDConfig(role="attention", num_attention_ranks=2), - ) - - with pytest.raises(RuntimeError, match="has role rank 1 but vLLM"): - connector_device.check_connector_role_rank( - connector, _role_rank_config(tp_size=2) - ) diff --git a/tests/unit/v1/worker/test_connector_lifecycle.py b/tests/unit/v1/worker/test_connector_lifecycle.py new file mode 100644 index 000000000..568a0d357 --- /dev/null +++ b/tests/unit/v1/worker/test_connector_lifecycle.py @@ -0,0 +1,193 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Workers create a routed-experts connector before vLLM's distributed init.""" + +from __future__ import annotations + +import sys +from types import SimpleNamespace + +import pytest + +torch = pytest.importorskip("torch") +pytest.importorskip("vllm") + +from afd_plugin.config import AFDConfig # noqa: E402 +from afd_plugin.connectors import AFDConnectorFactory # noqa: E402 +from afd_plugin.v1.worker.connector_lifecycle import ( # noqa: E402 + check_worker_connector, + create_worker_connector, +) +from tests.unit.fakes.connector import FakeRoutedConnector # noqa: E402 +from tests.unit.fakes.platform import use_cuda_like_platform # noqa: E402 +from tests.unit.fakes.vllm_config import gpu_worker_vllm_config # noqa: E402 + +ROUTED = "FakeRoutedConnector" + + +@pytest.fixture(autouse=True) +def _routed_connector_and_cuda_like_platform(monkeypatch): + monkeypatch.setitem( + AFDConnectorFactory._registry, ROUTED, lambda: FakeRoutedConnector + ) + use_cuda_like_platform(monkeypatch) + + +def test_other_connectors_are_left_to_their_runner(): + """A connector that is not an ``AFDRoutedExpertsConnector`` is not created + before vLLM's distributed init. + + Why: P2P and the NPU connectors keep their runner-owned bring-up. + """ + connector = create_worker_connector( + gpu_worker_vllm_config(), + AFDConfig(connector="P2pNcclAFDConnector", role="ffn"), + rank=0, + local_rank=0, + ) + + assert connector is None + + +def test_routed_connector_is_created_on_vllms_device_and_prepared(): + """The connector gets the worker's ranks, the device vLLM will select and + has joined the AFD world. + + Why: its NCCL environment must precede vLLM's first communicator, and its + device must be the one vLLM gives the worker. + """ + config = gpu_worker_vllm_config(data_parallel_rank_local=1) + + connector = create_worker_connector( + config, AFDConfig(connector=ROUTED, role="ffn"), rank=3, local_rank=0 + ) + + assert isinstance(connector, FakeRoutedConnector) + assert (connector.rank, connector.local_rank, connector.role_rank) == (3, 1, 0) + assert connector.events == [("prepare",)] + + +def test_assigned_device_ids_select_the_connector_device(): + """With ``--device-ids 4,5`` worker 0's connector runs on cuda:4. + + Why: vLLM maps device indices through the assigned ids, so they must be + registered before the connector's device is resolved. + """ + config = gpu_worker_vllm_config(assigned_physical_gpu_ids=[4, 5]) + + connector = create_worker_connector( + config, AFDConfig(connector=ROUTED, role="ffn"), rank=0, local_rank=0 + ) + + assert connector is not None + assert connector.local_rank == 4 + + +def test_tp_part_of_the_role_rank_comes_from_the_worker_rank(monkeypatch): + """With TP > 1 the role rank's TP part comes from the worker rank. + + Why: vLLM's TP group does not exist before its distributed init; reading + it there would fail. TP is vLLM's innermost rank dimension, so the worker + rank modulo the TP size is the TP rank. + """ + monkeypatch.setitem(sys.modules, "vllm.distributed.parallel_state", None) + config = gpu_worker_vllm_config( + data_parallel_size=2, data_parallel_rank=1, tensor_parallel_size=2 + ) + + connector = create_worker_connector( + config, + AFDConfig(connector=ROUTED, role="attention", num_attention_ranks=4), + rank=1, + local_rank=1, + ) + + assert connector is not None + # DP rank 1, TP rank 1 of 2. + assert connector.role_rank == 3 + + +def test_connector_is_closed_when_it_fails_to_join(monkeypatch): + """A failed ``prepare_before_distributed_init`` closes the connector. + + Why: the worker never receives the connector, so nobody else closes it. + """ + created: list[FakeRoutedConnector] = [] + original_init = FakeRoutedConnector.__init__ + + def recording_init(self, *args): + original_init(self, *args) + created.append(self) + + monkeypatch.setattr(FakeRoutedConnector, "__init__", recording_init) + monkeypatch.setattr( + FakeRoutedConnector, "prepare_error", RuntimeError("join failed") + ) + + with pytest.raises(RuntimeError, match="join failed"): + create_worker_connector( + gpu_worker_vllm_config(), + AFDConfig(connector=ROUTED, role="ffn"), + rank=0, + local_rank=0, + ) + + assert created[0].events == [("prepare",), ("close",)] + + +def _connector(*, local_rank=0, role_rank=0): + return SimpleNamespace( + local_rank=local_rank, + role_rank=role_rank, + afd_config=AFDConfig(role="attention", num_attention_ranks=2), + ) + + +def _tp_rank(monkeypatch, tp_rank): + monkeypatch.setattr( + "vllm.distributed.parallel_state.get_tensor_model_parallel_rank", + lambda: tp_rank, + ) + + +def test_check_accepts_vllms_device_and_role_rank(monkeypatch): + _tp_rank(monkeypatch, 1) + + check_worker_connector( + _connector(local_rank=1, role_rank=1), + gpu_worker_vllm_config(tensor_parallel_size=2), + torch.device("cuda", 1), + ) + + +def test_check_rejects_another_device(monkeypatch): + """A connector on another GPU than the worker's fails with both devices. + + Why: a vLLM upgrade that changes device selection must fail at startup + instead of moving traffic onto another rank's GPU. + """ + _tp_rank(monkeypatch, 0) + + with pytest.raises(RuntimeError, match="uses cuda:0 but vLLM placed"): + check_worker_connector( + _connector(local_rank=0), + gpu_worker_vllm_config(tensor_parallel_size=2), + torch.device("cuda", 4), + ) + + +def test_check_rejects_a_different_role_rank(monkeypatch): + """A connector whose early role rank differs from vLLM's fails startup. + + Why: the TP part of the role rank is derived from the worker rank before + vLLM's TP group exists; a vLLM change to its rank order must fail here + instead of routing another rank's tokens. + """ + _tp_rank(monkeypatch, 0) + + with pytest.raises(RuntimeError, match="has role rank 1 but vLLM"): + check_worker_connector( + _connector(role_rank=1), + gpu_worker_vllm_config(tensor_parallel_size=2), + torch.device("cuda", 0), + ) diff --git a/tests/unit/v1/worker/test_ffn_model_runner.py b/tests/unit/v1/worker/test_ffn_model_runner.py index eea8a4aa2..7114c2a79 100644 --- a/tests/unit/v1/worker/test_ffn_model_runner.py +++ b/tests/unit/v1/worker/test_ffn_model_runner.py @@ -179,7 +179,16 @@ def test_ffn_runner_owns_connector_execution_without_native_runner_inheritance() assert all(base.__name__ != "GPUModelRunner" for base in GPUFFNModelRunner.__mro__) -def test_ffn_constructor_defers_connector_init_and_owns_model_and_graphs(monkeypatch): +# A connector the worker created before vLLM's init is the runner's one +# instance: the runner then creates none. +@pytest.mark.parametrize( + ("workers_connector", "expected_events"), + [(False, ["validate", "connector"]), (True, ["validate"])], + ids=["runner-creates", "worker-passes"], +) +def test_ffn_constructor_defers_connector_init_and_owns_model_and_graphs( + monkeypatch, workers_connector, expected_events +): config = SimpleNamespace( use_v2_model_runner=True, model_config=SimpleNamespace( @@ -220,9 +229,13 @@ def create_connector(*_args): ffn_model_runner_module, "create_afd_gpu_profiler", lambda _role: None ) - runner = GPUFFNModelRunner(config, torch.device("cpu")) + runner = GPUFFNModelRunner( + config, + torch.device("cpu"), + connector=connector if workers_connector else None, + ) - assert events == ["validate", "connector"] + assert events == expected_events assert runner.connector is connector assert runner.model is None assert runner.model_memory_usage == 0 diff --git a/tests/unit/v1/worker/test_ffn_routed_experts_model_runner.py b/tests/unit/v1/worker/test_ffn_routed_experts_model_runner.py index a11c95be5..1c1ce752c 100644 --- a/tests/unit/v1/worker/test_ffn_routed_experts_model_runner.py +++ b/tests/unit/v1/worker/test_ffn_routed_experts_model_runner.py @@ -186,30 +186,36 @@ def test_initialize_passes_moe_layers_and_experts_to_the_connector(): } -def test_constructor_rejects_connectors_without_routed_rows(monkeypatch): - def base_init(self, vllm_config, device): - self.connector = SimpleNamespace(close=lambda: None) +def _base_init_with(monkeypatch): + def base_init(self, vllm_config, device, *, connector=None): + self.connector = connector or SimpleNamespace(close=lambda: None) monkeypatch.setattr(GPUFFNModelRunner, "__init__", base_init) + +def test_constructor_rejects_connectors_without_routed_rows(monkeypatch): + _base_init_with(monkeypatch) + with pytest.raises(TypeError, match="is not an AFDRoutedExpertsConnector"): GPUFFNRoutedExpertsModelRunner(SimpleNamespace(), torch.device("cpu")) -def test_ffn_runner_adopts_bootstrapped_connector_and_closes_placeholder(): - """The worker's bootstrapped connector replaces the one built at - construction, which is closed. +def test_constructor_uses_the_workers_connector(monkeypatch): + """The routed runner runs on the connector the worker created before + vLLM's distributed init. - Why: bootstrap must precede vLLM's distributed init, construction cannot. + Why: that connector already joined the AFD world; a second instance + would never be initialized by its peers. """ - events: list[str] = [] - runner = _runner(SimpleNamespace(close=lambda: events.append("close"))) - bootstrapped = FakeRoutedConnector() + connector = FakeRoutedConnector() + _base_init_with(monkeypatch) - runner.adopt_bootstrapped_connector(bootstrapped) + runner = GPUFFNRoutedExpertsModelRunner( + SimpleNamespace(), torch.device("cpu"), connector=connector + ) - assert runner.connector is bootstrapped - assert events == ["close"] + assert runner.connector is connector + assert runner.routed_connector is connector def _routed_experts(*, fp8: bool, expert_ids=(4, 5), backend: str = "TRITON"): diff --git a/tests/unit/v1/worker/test_model_runner_v2.py b/tests/unit/v1/worker/test_model_runner_v2.py index 1578aa3df..5ba3fd190 100644 --- a/tests/unit/v1/worker/test_model_runner_v2.py +++ b/tests/unit/v1/worker/test_model_runner_v2.py @@ -2067,22 +2067,67 @@ def test_v2_load_model_reports_moe_layers_to_routed_connector(monkeypatch): } -def test_v2_runner_adopts_bootstrapped_connector_and_closes_placeholder(): - """The worker's bootstrapped connector replaces the one built at - construction, which is closed. +def _patch_constructor_with_workers_connector(monkeypatch, events): + def native_init(self, vllm_config, device): + self.vllm_config = vllm_config + + def unexpected_create(*args, **kwargs): + pytest.fail("the runner created a second connector") + + monkeypatch.setattr(native_v2.GPUModelRunner, "__init__", native_init) + monkeypatch.setattr( + native_v2.GPUModelRunner, + "shutdown", + lambda self: events.append("native_shutdown"), + ) + monkeypatch.setattr( + "afd_plugin.v1.worker.attention_model_runner_v2.create_afd_gpu_profiler", + lambda role: _StepProfiler(), + ) + config = _patch_v2_constructor_dependencies(monkeypatch, None) + monkeypatch.setattr( + v2_runner_module.AFDConnectorFactory, "create_connector", unexpected_create + ) + return config + - Why: the native constructor signature stays unchanged. +def test_v2_constructor_uses_the_workers_connector_without_creating_one( + monkeypatch, +): + """A connector the worker created before vLLM's init is the runner's one + instance; the runner creates none. + + Why: one connector instance per worker, with a single create, then + ``init_afd_connector`` lifecycle. """ events: list[str] = [] - placeholder = SimpleNamespace(close=lambda: events.append("placeholder_close")) - bootstrapped = SimpleNamespace() - runner = object.__new__(AFDAttentionModelRunnerV2) - runner.connector = placeholder + connector = SimpleNamespace( + control_plane=object(), close=lambda: events.append("close") + ) + config = _patch_constructor_with_workers_connector(monkeypatch, events) + + runner = AFDAttentionModelRunnerV2(config, torch.device("cpu"), connector=connector) + + assert runner.connector is connector + assert events == [] - runner.adopt_bootstrapped_connector(bootstrapped) - assert runner.connector is bootstrapped - assert events == ["placeholder_close"] +def test_v2_constructor_failure_leaves_the_workers_connector_open(monkeypatch): + """A failed construction closes native state but not the worker's + connector. + + Why: the worker still owns that connector and closes it itself. + """ + events: list[str] = [] + connector = SimpleNamespace( + control_plane=None, close=lambda: events.append("close") + ) + config = _patch_constructor_with_workers_connector(monkeypatch, events) + + with pytest.raises(RuntimeError, match="control-plane"): + AFDAttentionModelRunnerV2(config, torch.device("cpu"), connector=connector) + + assert events == ["native_shutdown"] def test_v2_constructor_rejects_missing_control_plane_and_cleans_native( @@ -2294,13 +2339,10 @@ def raising_worker_init(current_worker): class _ConnectorRunnerRecorder: - adopted: list[object] = [] + connectors: list[object] = [] - def __init__(self, vllm_config, device): - pass - - def adopt_bootstrapped_connector(self, connector): - type(self).adopted.append(connector) + def __init__(self, vllm_config, device, *, connector=None): + type(self).connectors.append(connector) _ROUTED_WORKER_CASES = [ @@ -2320,60 +2362,67 @@ def adopt_bootstrapped_connector(self, connector): ] +def _patch_routed_worker(monkeypatch, worker_module, runner_module, runner_name): + """Select a routed-experts connector and record the runner construction.""" + _ConnectorRunnerRecorder.connectors = [] + # The FFN worker picks its runner from the connector class. + monkeypatch.setattr( + ffn_worker.AFDConnectorFactory, + "get_connector_class", + lambda name: FakeRoutedConnector, + ) + monkeypatch.setattr(runner_module, runner_name, _ConnectorRunnerRecorder) + monkeypatch.setattr(native_v2, "GPUModelRunner", _NativeRunnerSentinel) + monkeypatch.setattr( + worker_module, + "assert_compatible_afd_stack", + lambda *args, **kwargs: None, + ) + monkeypatch.setattr(worker_module.torch.accelerator, "empty_cache", lambda: None) + + @pytest.mark.parametrize( ("worker_module", "worker_cls", "runner_module", "runner_name"), _ROUTED_WORKER_CASES, ) -def test_workers_bootstrap_connector_before_vllm_and_hand_it_to_runner( +def test_workers_create_connector_before_vllm_and_construct_runner_with_it( monkeypatch, worker_module, worker_cls, runner_module, runner_name ): - """An ``AFDRoutedExpertsConnector`` is created and bootstrapped before - vLLM's distributed init and adopted by the runner; the FFN picks the - routed-experts runner. + """A connector that joins the AFD world before vLLM's distributed init is + created first and passed to the runner vLLM constructs, which uses that + one instance; the FFN picks the routed-experts runner. - Why: NCCL reads its environment when vLLM creates the first communicator. + Why: NCCL reads its environment when vLLM creates the first communicator, + and vLLM builds the runner only after that. """ config = _v2_config(role="ffn" if worker_cls is AFDFFNWorker else "attention") worker = _prepare_worker(worker_cls, config, use_v2=True) events: list[str] = [] - connector = SimpleNamespace(close=lambda: events.append("close"), local_rank=0) - _ConnectorRunnerRecorder.adopted = [] + connector = SimpleNamespace(close=lambda: events.append("close")) - def create_bootstrapped_connector(rank, local_rank, vllm_config, afd_config): - events.append("bootstrap") + def create_worker_connector(vllm_config, afd_config, *, rank, local_rank): + events.append("create") return connector def vllm_init(current_worker): events.append("vllm_init") _fake_worker_init(current_worker) - current_worker.device = torch.device("cuda", 0) - factory = worker_module.AFDConnectorFactory - monkeypatch.setattr( - factory, "create_bootstrapped_connector", create_bootstrapped_connector - ) + _patch_routed_worker(monkeypatch, worker_module, runner_module, runner_name) monkeypatch.setattr( - factory, "get_connector_class", lambda name: FakeRoutedConnector + worker_module, "create_worker_connector", create_worker_connector ) - monkeypatch.setattr(runner_module, runner_name, _ConnectorRunnerRecorder) - monkeypatch.setattr(native_v2, "GPUModelRunner", _NativeRunnerSentinel) monkeypatch.setattr(worker_module.Worker, "init_device", vllm_init) monkeypatch.setattr( worker_module, - "assert_compatible_afd_stack", - lambda *args, **kwargs: None, - ) - monkeypatch.setattr(worker_module.torch.accelerator, "empty_cache", lambda: None) - monkeypatch.setattr( - worker_module, - "check_connector_role_rank", - lambda checked, vllm_config: events.append("role_rank_check"), + "check_worker_connector", + lambda checked, vllm_config, device: events.append("check"), ) worker.init_device() - assert events == ["bootstrap", "vllm_init", "role_rank_check"] - assert _ConnectorRunnerRecorder.adopted == [connector] + assert events == ["create", "vllm_init", "check"] + assert _ConnectorRunnerRecorder.connectors == [connector] assert native_v2.GPUModelRunner is _NativeRunnerSentinel # The routed FFN pins vLLM's MoE backend to one that keeps the expert layout. expected_backend = "triton" if worker_cls is AFDFFNWorker else "auto" @@ -2384,7 +2433,7 @@ def vllm_init(current_worker): ("worker_module", "worker_cls", "runner_module", "runner_name"), _ROUTED_WORKER_CASES, ) -def test_workers_close_bootstrapped_connector_when_vllm_init_fails( +def test_workers_close_their_connector_when_vllm_init_fails( monkeypatch, worker_module, worker_cls, runner_module, runner_name ): config = _v2_config(role="ffn" if worker_cls is AFDFFNWorker else "attention") @@ -2395,21 +2444,11 @@ def test_workers_close_bootstrapped_connector_when_vllm_init_fails( def failing_vllm_init(current_worker): raise RuntimeError("vLLM init failed") - factory = worker_module.AFDConnectorFactory - monkeypatch.setattr( - factory, "create_bootstrapped_connector", lambda *args: connector - ) + _patch_routed_worker(monkeypatch, worker_module, runner_module, runner_name) monkeypatch.setattr( - factory, "get_connector_class", lambda name: FakeRoutedConnector + worker_module, "create_worker_connector", lambda *args, **kwargs: connector ) - monkeypatch.setattr(runner_module, runner_name, _ConnectorRunnerRecorder) - monkeypatch.setattr(native_v2, "GPUModelRunner", _NativeRunnerSentinel) monkeypatch.setattr(worker_module.Worker, "init_device", failing_vllm_init) - monkeypatch.setattr( - worker_module, - "assert_compatible_afd_stack", - lambda *args, **kwargs: None, - ) with pytest.raises(RuntimeError, match="vLLM init failed"): worker.init_device() @@ -2422,10 +2461,10 @@ def failing_vllm_init(current_worker): ("worker_module", "worker_cls", "runner_module", "runner_name"), _ROUTED_WORKER_CASES, ) -def test_workers_reject_bootstrapped_connector_on_another_device( +def test_workers_reject_connector_on_another_device( monkeypatch, worker_module, worker_cls, runner_module, runner_name ): - """A connector bootstrapped on a different GPU than vLLM's fails startup. + """A connector created on a different GPU than vLLM's fails startup. Why: the connector picks its device before vLLM does, by mirroring vLLM's selection; a vLLM upgrade that changes the selection must fail here @@ -2435,33 +2474,21 @@ def test_workers_reject_bootstrapped_connector_on_another_device( worker = _prepare_worker(worker_cls, config, use_v2=True) events: list[str] = [] connector = SimpleNamespace(close=lambda: events.append("close"), local_rank=0) - _ConnectorRunnerRecorder.adopted = [] def vllm_init(current_worker): _fake_worker_init(current_worker) current_worker.device = torch.device("cuda", 4) - factory = worker_module.AFDConnectorFactory + _patch_routed_worker(monkeypatch, worker_module, runner_module, runner_name) monkeypatch.setattr( - factory, "create_bootstrapped_connector", lambda *args: connector + worker_module, "create_worker_connector", lambda *args, **kwargs: connector ) - monkeypatch.setattr( - factory, "get_connector_class", lambda name: FakeRoutedConnector - ) - monkeypatch.setattr(runner_module, runner_name, _ConnectorRunnerRecorder) - monkeypatch.setattr(native_v2, "GPUModelRunner", _NativeRunnerSentinel) monkeypatch.setattr(worker_module.Worker, "init_device", vllm_init) - monkeypatch.setattr( - worker_module, - "assert_compatible_afd_stack", - lambda *args, **kwargs: None, - ) with pytest.raises(RuntimeError, match="uses cuda:0 but vLLM placed"): worker.init_device() assert events == ["close"] - assert _ConnectorRunnerRecorder.adopted == [] assert native_v2.GPUModelRunner is _NativeRunnerSentinel @@ -2477,7 +2504,7 @@ def test_v2_validator_accepts_deepep_m2n_attention_tp(tp_size): """DeepEP M2N accepts Attention TP > 1. Why: each Attention TP rank dispatches its own slice of the tokens, and - the bootstrap derives the TP rank from the worker rank before vLLM's + the worker derives the TP rank from the worker rank before vLLM's process groups exist. """ validate_gpu_model_runner_v2_config( From 9fe72136c2fc87c50cb01bc82233c0b4716d680c Mon Sep 17 00:00:00 2001 From: ronenkat <16743404+ronenkat@users.noreply.github.com> Date: Mon, 26 Oct 2026 14:47:33 +0200 Subject: [PATCH 11/32] refactor(gpu): drop unrouted-rail NIC planning and the pre-init lifecycle The DeepEP M2N connector planned NICs for RoCE fabrics whose rails are separate, unrouted subnets: it pinned every rank to one shared rail through NCCL_IB_HCA. NCCL reads that variable once per process, when the first communicator is created, so the connector had to join the AFD world before vLLM's distributed init. That requirement was the only reason for the worker lifecycle (connector_lifecycle.py), the pre-init device mirroring (distributed/device.py), the TP rank derived from the worker rank, and the connector injection into both runners. Routed fabrics (InfiniBand, routed RoCE, NVL72, one host) need no NIC planning: NCCL's own per-GPU choice works there. Drop NIC and rail discovery, the shared/explicit NIC policies, rails_routed, nics, nccl_env and the presets; keep NVLink-domain discovery, the GIN decision, hybrid mode and the traffic class. The only variables the connector still sets, EP_DISABLE_GIN and the per-buffer comm stream switch, are read by afdEP when it creates its buffers, so the handshake moves into init_afd_connector and the connector follows the same lifecycle as P2P: the runner creates it and initializes it after the model loads. On unrouted-rail RoCE, set NCCL_IB_HCA to one shared rail in the launch environment. Co-Authored-By: Claude Opus 5.5 Signed-off-by: ronenkat <16743404+ronenkat@users.noreply.github.com> --- afd_plugin/connectors/factory.py | 5 +- .../connectors/gpu/deepep_m2n/connector.py | 48 +-- .../connectors/gpu/deepep_m2n/network.py | 328 ++--------------- afd_plugin/connectors/routed_experts.py | 15 +- afd_plugin/distributed/__init__.py | 8 - afd_plugin/distributed/device.py | 80 ----- afd_plugin/distributed/topology.py | 38 +- .../v1/worker/attention_model_runner_v2.py | 25 +- afd_plugin/v1/worker/attention_worker.py | 32 +- afd_plugin/v1/worker/connector_lifecycle.py | 113 ------ afd_plugin/v1/worker/ffn_model_runner.py | 27 +- .../worker/ffn_routed_experts_model_runner.py | 10 +- afd_plugin/v1/worker/ffn_worker.py | 32 +- afd_plugin/validation.py | 17 +- docs/design/module/connector_contracts.md | 1 - tests/gpu/deepep_m2n_harness.py | 19 +- .../connectors/test_deepep_m2n_connector.py | 21 +- .../connectors/test_deepep_m2n_network.py | 339 +++--------------- tests/unit/distributed/test_device.py | 85 ----- tests/unit/distributed/test_topology.py | 64 +--- tests/unit/fakes/connector.py | 15 +- tests/unit/fakes/platform.py | 22 -- tests/unit/fakes/vllm_config.py | 23 -- .../v1/worker/test_connector_lifecycle.py | 193 ---------- tests/unit/v1/worker/test_ffn_model_runner.py | 19 +- .../test_ffn_routed_experts_model_runner.py | 28 +- tests/unit/v1/worker/test_model_runner_v2.py | 216 +---------- 27 files changed, 163 insertions(+), 1660 deletions(-) delete mode 100644 afd_plugin/distributed/device.py delete mode 100644 afd_plugin/v1/worker/connector_lifecycle.py delete mode 100644 tests/unit/distributed/test_device.py delete mode 100644 tests/unit/fakes/platform.py delete mode 100644 tests/unit/v1/worker/test_connector_lifecycle.py diff --git a/afd_plugin/connectors/factory.py b/afd_plugin/connectors/factory.py index ab778f942..4cc19400d 100644 --- a/afd_plugin/connectors/factory.py +++ b/afd_plugin/connectors/factory.py @@ -53,15 +53,12 @@ def create_connector( local_rank: int, vllm_config: VllmConfig, afd_config: AFDConfig | None = None, - *, - tp_rank: int | None = None, ) -> AFDConnectorBase: - """Create a connector; ``tp_rank`` is passed to ``resolve_role_rank``.""" config = afd_config or parse_afd_config(vllm_config) if config.connector not in cls._registry: raise ValueError(f"unsupported AFD connector type: {config.connector}") connector_cls = cls._registry[config.connector]() - role_rank = resolve_role_rank(vllm_config, config, tp_rank=tp_rank) + role_rank = resolve_role_rank(vllm_config, config) return connector_cls( rank, local_rank, diff --git a/afd_plugin/connectors/gpu/deepep_m2n/connector.py b/afd_plugin/connectors/gpu/deepep_m2n/connector.py index 5c475ce37..e688c8dd9 100644 --- a/afd_plugin/connectors/gpu/deepep_m2n/connector.py +++ b/afd_plugin/connectors/gpu/deepep_m2n/connector.py @@ -85,7 +85,7 @@ NetworkConfig, NetworkPlan, RankTopology, - apply_rank_environment, + apply_gin_environment, discover_local_topology, plan_network, ) @@ -163,9 +163,8 @@ class DeepEPM2NConfig(ConnectorExtraInfo): (``resolve_num_lanes``). payload_format: Dispatch wire format. expert_backend: ``deepgemm`` (fast) or ``fused_experts`` (reference). - network: Fabric settings (GIN, NIC choice, hybrid mode, traffic - class, extra NCCL env); see ``network.NetworkConfig`` and its - presets. + network: Fabric settings (GIN, hybrid mode, traffic class); see + ``network.NetworkConfig``. NICs are NCCL's own per-GPU choice. bucket_granularity: Dispatch bucket is rounded up to this multiple, then up to the next rung of the bucket ladder (``bucket_ladder``) or, above the ladder, to the capacity, so every step uses one of a @@ -892,9 +891,9 @@ class DeepEPM2NAFDConnector(AFDRoutedExpertsConnector): once, also within one CUDA-graph capture: a capture holds about two layers' rows per lane, not every layer's. - Lifecycle: ``prepare_before_distributed_init`` (before vLLM's distributed - init: handshake and NCCL environment), then ``init_afd_connector`` (after - weights load: post-load checks, NCCL control group and DeepEP buffers). + Lifecycle: the runner creates the connector and calls + ``init_afd_connector`` after weights load: handshake and network plan, + post-load checks, NCCL control group, then DeepEP buffers. """ @classmethod @@ -933,8 +932,7 @@ def __init__( raise ValueError(f"FP8 payload needs hidden_size % {FP8_GROUP_SIZE} == 0") self.world_rank = self.layout.world_rank(afd_config.role, role_rank) # The FFN runs without TP and learns the Attention TP size and the - # resulting capacity in the handshake - # (``prepare_before_distributed_init``). + # resulting capacity in the handshake (``_handshake``). self.attn_tp_size: int | None = ( None if self.is_ffn else vllm_config.parallel_config.tensor_parallel_size ) @@ -995,18 +993,12 @@ def __init__( # Lifecycle # ============================== - def prepare_before_distributed_init(self) -> None: - """Collective: join the AFD world, check compatibility, set NCCL env. + def _handshake(self) -> None: + """Collective: join the AFD world, check compatibility, plan the network. - Workers call this before vLLM's distributed init, because NCCL reads - its environment (``NCCL_IB_HCA``, the GIN switch) once per process, - when the first communicator is created. Uses gloo only, so it creates - no NCCL communicator itself. ``init_afd_connector`` calls it when the - caller has not. + Uses gloo only. Sets the DeepEP environment the buffers read when + ``init_afd_connector`` creates them. """ - if self.bootstrap_group is not None: - return - torch.cuda.set_device(self.device) self.bootstrap_group = init_afd_process_group( backend="gloo", init_method=self._rendezvous, @@ -1041,9 +1033,7 @@ def prepare_before_distributed_init(self) -> None: ffn_deployment_info(self.vllm_config) if self.is_ffn else None ), sm=torch.cuda.get_device_capability(self.device), - topology=discover_local_topology( - self.world_rank, self.device, self.local_rank - ), + topology=discover_local_topology(self.world_rank, self.device), ) infos: list[dict[str, Any] | None] = [None] * self.layout.union_size dist.all_gather_object(infos, info, group=self.bootstrap_group) @@ -1059,11 +1049,9 @@ def prepare_before_distributed_init(self) -> None: self.bucket_ladder = result.bucket_ladder if self.world_rank == 0: logger.info("DeepEP M2N network plan: %s", self.network_plan.summary()) - # Before any NCCL communicator in this process exists, when the worker - # calls this ahead of vLLM's distributed init. - apply_rank_environment(self.network_plan, self.config.network, self.world_rank) - # DeepEP reads this when each buffer is created: with several lanes - # each buffer needs its own communication stream. + # DeepEP reads both when each buffer is created: the GIN switch, and + # with several lanes a communication stream per buffer. + apply_gin_environment(self.network_plan) if self.num_lanes > 1: os.environ[PER_BUFFER_COMM_STREAM_ENV] = "1" @@ -1085,8 +1073,8 @@ def init_afd_connector( return if local_experts is not None and not self.is_ffn: raise ValueError("only FFN ranks hold local experts") - self.prepare_before_distributed_init() torch.cuda.set_device(self.device) + self._handshake() _register_custom_ops() assert self.bootstrap_group is not None and self.network_plan is not None assert self.capacity_tokens_per_rank is not None @@ -1594,9 +1582,7 @@ def update_state_from_dp_metadata(self, payload: AFDControlPayload) -> None: or connector.capacity_tokens_per_rank is None or connector.bucket_ladder is None ): - raise RuntimeError( - "DeepEPM2N control plane used before prepare_before_distributed_init" - ) + raise RuntimeError("DeepEPM2N control plane used before the handshake") connector.stage_plans = { stage_idx: plan_stage( dp_metadata, diff --git a/afd_plugin/connectors/gpu/deepep_m2n/network.py b/afd_plugin/connectors/gpu/deepep_m2n/network.py index 133bd77c1..715a49ed6 100644 --- a/afd_plugin/connectors/gpu/deepep_m2n/network.py +++ b/afd_plugin/connectors/gpu/deepep_m2n/network.py @@ -3,97 +3,57 @@ """Cluster topology discovery and the network plan for DeepEP M2N. DeepEP V2 moves data inside an NVLink domain with load/stores (LSA) and -between domains with NCCL GIN (RDMA). What that needs depends on the cluster: +between domains with NCCL GIN (RDMA): - one host, or an NVL72 rack: one NVLink domain, no RDMA (GIN off); -- InfiniBand fat-tree or RoCE with routed rails: GIN on, each GPU uses its - nearest NIC (NCCL's own choice); -- RoCE whose rails are separate, unrouted subnets: GIN on, and every rank must - use NICs on one shared rail, because DeepEP's GIN setup connects every pair - of ranks. +- several NVLink domains (InfiniBand, or RoCE with routed rails): GIN on, and + each GPU uses NCCL's own NIC choice. + +DeepEP's GIN setup connects every pair of ranks, so the fabric must route +between the NICs NCCL picks. On RoCE whose rails are separate, unrouted +subnets, pin every rank to NICs on one shared rail through ``NCCL_IB_HCA`` in +the launch environment; the connector does not plan NICs. Each rank reports a ``RankTopology`` in the connector's init handshake, and every rank evaluates the same ``plan_network`` on the gathered list, so the -plan needs no further agreement. Whether rails are routed cannot be seen from -inside a pod; ``rails_routed`` declares it, and ``auto`` assumes they are not -when NICs sit on several rails (the choice that works everywhere, at the cost -of bandwidth). +plan needs no further agreement. """ from __future__ import annotations import ctypes -import ipaddress import os import socket from collections.abc import Mapping -from dataclasses import asdict, dataclass, field, fields -from pathlib import Path +from dataclasses import asdict, dataclass, fields from typing import Any -import psutil import torch -SYSFS_INFINIBAND = Path("/sys/class/infiniband") -SYSFS_PCI = Path("/sys/bus/pci/devices") -# RoCE v2 GIDs carry the port's IPv4 address as ::ffff:a.b.c.d. -IPV4_MAPPED_GID_PREFIX = "0000:0000:0000:0000:0000:ffff:" -ETHERNET = "ethernet" -INFINIBAND = "infiniband" -UNKNOWN_DISTANCE = 1 << 16 - GIN_MODES = ("auto", "on", "off") -NIC_POLICIES = ("auto", "nccl", "shared", "explicit") HYBRID_MODES = ("auto", "on", "off") -NCCL_ENV_PREFIXES = ("NCCL_", "EP_") -# Set from the plan (``network.gin``); a copy in ``nccl_env`` would fight it. +# DeepEP reads it when it creates each buffer; set from the plan. GIN_DISABLE_ENV = "EP_DISABLE_GIN" # DeepEP's default RDMA service level / traffic class. DEFAULT_RDMA_TRAFFIC_CLASS = 3 -# Preset -> NetworkConfig defaults; explicit keys in the config override them. -PRESETS: dict[str, dict[str, Any]] = { - "auto": {}, - "single-node": {"gin": "off", "nic_policy": "nccl"}, - "roce-rail-isolated": {"nic_policy": "shared", "rails_routed": False}, - "roce-routed": {"nic_policy": "nccl", "rails_routed": True}, - "ib": {"nic_policy": "nccl", "rails_routed": True}, - "nvl72": {"nic_policy": "nccl", "rails_routed": True}, -} - @dataclass(frozen=True) class NetworkConfig: """The ``network`` section of ``connector_extra_config``. Attributes: - preset: Defaults for a fabric type (see ``PRESETS``). gin: ``auto`` enables NCCL GIN only when ranks span several NVLink domains; ``on``/``off`` force it (``off`` is rejected when ranks span several domains: GIN is the only path between them). - nic_policy: ``auto`` picks ``nccl`` when rails are routed or all - NICs share one rail, else ``shared``; ``nccl`` keeps NCCL's - per-GPU choice; ``shared`` puts every rank on NICs of one rail - present on every host; ``explicit`` uses ``nics`` by local rank. - nics: Explicit NIC names: by local rank for ``explicit``, or the NIC - to prefer (first entry) for ``shared``. - rails_routed: Whether different rails can reach each other. - ``None`` means unknown (treated as not routed). hybrid: DeepEP hybrid mode (rail-aligned RDMA + NVLink forwarding); ``auto`` = on. rdma_traffic_class: GIN service level / traffic class. - nccl_env: Extra ``NCCL_*``/``EP_*`` variables applied before any - NCCL communicator is created (process-wide). """ - preset: str = "auto" gin: str = "auto" - nic_policy: str = "auto" - nics: tuple[str, ...] = () - rails_routed: bool | None = None hybrid: str = "auto" rdma_traffic_class: int = DEFAULT_RDMA_TRAFFIC_CLASS - nccl_env: tuple[tuple[str, str], ...] = () @classmethod def from_mapping(cls, raw: Mapping[str, Any] | None) -> NetworkConfig: @@ -102,60 +62,20 @@ def from_mapping(cls, raw: Mapping[str, Any] | None) -> NetworkConfig: unknown = sorted(set(raw) - known) if unknown: raise ValueError(f"unknown network options: {unknown}") - preset = raw.get("preset", "auto") - if preset not in PRESETS: - raise ValueError(f"network.preset must be one of {sorted(PRESETS)}") - merged = {**PRESETS[preset], **raw} - if "nics" in merged: - merged["nics"] = tuple(merged["nics"]) - if "nccl_env" in merged: - merged["nccl_env"] = tuple( - sorted((str(k), str(v)) for k, v in dict(merged["nccl_env"]).items()) - ) - config = cls(**merged) + config = cls(**raw) config.validate() return config def validate(self) -> None: if self.gin not in GIN_MODES: raise ValueError(f"network.gin must be one of {GIN_MODES}") - if self.nic_policy not in NIC_POLICIES: - raise ValueError(f"network.nic_policy must be one of {NIC_POLICIES}") if self.hybrid not in HYBRID_MODES: raise ValueError(f"network.hybrid must be one of {HYBRID_MODES}") - if not all(isinstance(nic, str) and nic for nic in self.nics): - raise ValueError("network.nics must be non-empty NIC names") - if self.nic_policy == "explicit" and not self.nics: - raise ValueError("network.nic_policy=explicit needs network.nics") if self.rdma_traffic_class < 0: raise ValueError("network.rdma_traffic_class must be non-negative") - for key, _ in self.nccl_env: - if not key.startswith(NCCL_ENV_PREFIXES): - raise ValueError( - f"network.nccl_env key {key!r} must start with {NCCL_ENV_PREFIXES}" - ) - if key == GIN_DISABLE_ENV: - raise ValueError( - f"network.nccl_env cannot set {GIN_DISABLE_ENV}; use network.gin" - ) def to_mapping(self) -> dict[str, Any]: - mapping = asdict(self) - mapping["nics"] = list(self.nics) - mapping["nccl_env"] = dict(self.nccl_env) - return mapping - - -@dataclass(frozen=True) -class NicInfo: - """One active RDMA port as seen from a rank.""" - - name: str - link_layer: str - # Ranks can talk RDMA only between NICs on the same rail (unless routed). - rail: str - # PCIe hops between the rank's GPU and this NIC (lower is closer). - distance: int + return asdict(self) @dataclass(frozen=True) @@ -164,8 +84,6 @@ class RankTopology: host: str # Ranks with the same id share an NVLink (LSA) domain. nvlink_domain: str - local_index: int - nics: tuple[NicInfo, ...] = field(default_factory=tuple) def to_mapping(self) -> dict[str, Any]: return asdict(self) @@ -176,8 +94,6 @@ def from_mapping(cls, raw: Mapping[str, Any]) -> RankTopology: world_rank=int(raw["world_rank"]), host=str(raw["host"]), nvlink_domain=str(raw["nvlink_domain"]), - local_index=int(raw["local_index"]), - nics=tuple(NicInfo(**nic) for nic in raw["nics"]), ) @@ -189,18 +105,12 @@ class NetworkPlan: ranks_per_domain: int use_gin: bool hybrid: bool - # world rank -> NIC to pin via NCCL_IB_HCA (None = NCCL's own choice). - nic_by_rank: tuple[str | None, ...] - nic_policy: str - rail: str | None rdma_traffic_class: int def summary(self) -> str: - nics = sorted({nic for nic in self.nic_by_rank if nic is not None}) return ( f"domains={self.num_domains}x{self.ranks_per_domain} gin={self.use_gin} " - f"hybrid={self.hybrid} nic_policy={self.nic_policy} rail={self.rail} " - f"nics={nics or 'nccl'} traffic_class={self.rdma_traffic_class}" + f"hybrid={self.hybrid} traffic_class={self.rdma_traffic_class}" ) @@ -209,32 +119,12 @@ def summary(self) -> str: # --------------------------------------------------------------------------- -def discover_local_topology( - world_rank: int, - device: torch.device, - local_index: int, - *, - sysfs_infiniband: Path = SYSFS_INFINIBAND, - sysfs_pci: Path = SYSFS_PCI, - interface_networks: Mapping[str, str] | None = None, -) -> RankTopology: +def discover_local_topology(world_rank: int, device: torch.device) -> RankTopology: host = socket.gethostname() - gpu_pci = gpu_pci_address(device) return RankTopology( world_rank=world_rank, host=host, - nvlink_domain=nvlink_domain_id(gpu_pci) or f"host:{host}", - local_index=local_index, - nics=discover_nics( - gpu_pci, - sysfs_infiniband=sysfs_infiniband, - sysfs_pci=sysfs_pci, - interface_networks=( - local_interface_networks() - if interface_networks is None - else interface_networks - ), - ), + nvlink_domain=nvlink_domain_id(gpu_pci_address(device)) or f"host:{host}", ) @@ -274,112 +164,6 @@ def nvlink_domain_id(gpu_pci: str | None) -> str | None: return f"nvl:{cluster}:{info.cliqueId}" -def local_interface_networks() -> dict[str, str]: - """Local IPv4 address -> its network (``a.b.c.d/len``).""" - networks: dict[str, str] = {} - for addresses in psutil.net_if_addrs().values(): - for address in addresses: - if address.family == socket.AF_INET and address.netmask: - network = ipaddress.ip_network( - f"{address.address}/{address.netmask}", strict=False - ) - networks[address.address] = str(network) - return networks - - -def discover_nics( - gpu_pci: str | None, - *, - sysfs_infiniband: Path, - sysfs_pci: Path, - interface_networks: Mapping[str, str], -) -> tuple[NicInfo, ...]: - nics = [] - if not sysfs_infiniband.is_dir(): - return () - for device_dir in sorted(sysfs_infiniband.iterdir()): - port = device_dir / "ports" / "1" - state = _read(port / "state") - if state is None or "ACTIVE" not in state: - continue - link_layer = (_read(port / "link_layer") or "").strip().lower() - rail = nic_rail(port, link_layer, interface_networks) - if rail is None: - continue - nics.append( - NicInfo( - name=device_dir.name, - link_layer=link_layer, - rail=rail, - distance=pci_distance( - gpu_pci, _pci_address(device_dir), sysfs_pci=sysfs_pci - ), - ) - ) - return tuple(nics) - - -def nic_rail( - port: Path, link_layer: str, interface_networks: Mapping[str, str] -) -> str | None: - """Rail id: the RoCE port's IPv4 network, or the IB GID subnet prefix.""" - gids_dir = port / "gids" - if not gids_dir.is_dir(): - return None - gids = [_read(gid) for gid in sorted(gids_dir.iterdir(), key=_gid_index)] - if link_layer == ETHERNET: - for gid in gids: - if gid and gid.strip().startswith(IPV4_MAPPED_GID_PREFIX): - ipv4 = _ipv4_from_gid(gid.strip()) - # Without the netmask the port is its own rail (isolated). - return f"roce:{interface_networks.get(ipv4, ipv4 + '/32')}" - return None - if link_layer == INFINIBAND and gids and gids[0]: - prefix = ":".join(gids[0].strip().split(":")[:4]) - return f"ib:{prefix}" - return None - - -def pci_distance(gpu_pci: str | None, nic_pci: str | None, *, sysfs_pci: Path) -> int: - """PCIe hops between two functions via their common upstream bridge.""" - if gpu_pci is None or nic_pci is None: - return UNKNOWN_DISTANCE - try: - gpu_path = (sysfs_pci / gpu_pci).resolve(strict=True).parts - nic_path = (sysfs_pci / nic_pci).resolve(strict=True).parts - except OSError: - return UNKNOWN_DISTANCE - common = 0 - for gpu_part, nic_part in zip(gpu_path, nic_path, strict=False): - if gpu_part != nic_part: - break - common += 1 - return (len(gpu_path) - common) + (len(nic_path) - common) - - -def _pci_address(device_dir: Path) -> str | None: - try: - return (device_dir / "device").resolve(strict=True).name - except OSError: - return None - - -def _gid_index(path: Path) -> int: - return int(path.name) if path.name.isdigit() else 1 << 30 - - -def _ipv4_from_gid(gid: str) -> str: - high, low = gid[len(IPV4_MAPPED_GID_PREFIX) :].split(":") - return str(ipaddress.IPv4Address(int(high + low, 16))) - - -def _read(path: Path) -> str | None: - try: - return path.read_text(encoding="utf-8") - except OSError: - return None - - # --------------------------------------------------------------------------- # Plan (pure; identical on every rank) # --------------------------------------------------------------------------- @@ -396,17 +180,11 @@ def plan_network(topologies: list[RankTopology], config: NetworkConfig) -> Netwo f"network.gin=off, but the ranks span {len(domains)} NVLink domains; " "DeepEP needs NCCL GIN (RDMA) between domains" ) - use_gin = multi_domain if config.gin == "auto" else config.gin == "on" - hybrid = config.hybrid != "off" - policy, rail, nic_by_rank = _choose_nics(topologies, config, use_gin) return NetworkPlan( num_domains=len(domains), ranks_per_domain=len(domains[0]), - use_gin=use_gin, - hybrid=hybrid, - nic_by_rank=nic_by_rank, - nic_policy=policy, - rail=rail, + use_gin=multi_domain if config.gin == "auto" else config.gin == "on", + hybrid=config.hybrid != "off", rdma_traffic_class=config.rdma_traffic_class, ) @@ -431,74 +209,12 @@ def _domain_blocks(topologies: list[RankTopology]) -> list[list[int]]: return [ranks for _, ranks in blocks] -def _choose_nics( - topologies: list[RankTopology], config: NetworkConfig, use_gin: bool -) -> tuple[str, str | None, tuple[str | None, ...]]: - none = tuple(None for _ in topologies) - if not use_gin: - return "none", None, none - policy = config.nic_policy - if policy == "explicit": - nics = tuple(config.nics[t.local_index % len(config.nics)] for t in topologies) - return policy, None, nics - rails = {nic.rail for t in topologies for nic in t.nics} - if policy == "auto": - policy = "nccl" if config.rails_routed or len(rails) <= 1 else "shared" - if policy == "nccl": - return policy, None, none - rail = _shared_rail(topologies, config) - nics = tuple( - min( - (nic for nic in t.nics if nic.rail == rail), - key=lambda nic: (nic.distance, nic.name), - ).name - for t in topologies - ) - return policy, rail, nics - - -def _shared_rail(topologies: list[RankTopology], config: NetworkConfig) -> str: - """A rail every rank has a NIC on; prefer ``nics[0]``'s, then the closest.""" - common = set.intersection(*({nic.rail for nic in t.nics} for t in topologies)) - if not common: - raise RuntimeError("no RDMA rail is reachable from every rank") - if config.nics: - preferred = { - nic.rail - for t in topologies - for nic in t.nics - if nic.name == config.nics[0] and nic.rail in common - } - if preferred: - return min(preferred) - return min( - common, - key=lambda rail: ( - sum( - min(nic.distance for nic in t.nics if nic.rail == rail) - for t in topologies - ), - rail, - ), - ) - - -def apply_rank_environment( - plan: NetworkPlan, config: NetworkConfig, world_rank: int -) -> None: - """Process-wide NCCL/DeepEP settings for this rank, before any communicator. +def apply_gin_environment(plan: NetworkPlan) -> None: + """Set DeepEP's GIN switch for this process from the plan. - NCCL reads these when its first communicator in the process is created. - AFD workers therefore run the connector's - ``prepare_before_distributed_init`` (which calls this) before vLLM's - distributed init creates its own NCCL groups; a process that - already created a communicator keeps the values it started with. + DeepEP reads it when it creates each buffer, so call this before the + connector creates its buffers. Every rank sets the same value. """ - for key, value in config.nccl_env: - os.environ[key] = value - nic = plan.nic_by_rank[world_rank] - if nic is not None: - os.environ["NCCL_IB_HCA"] = "=" + nic if plan.use_gin: os.environ.pop(GIN_DISABLE_ENV, None) else: diff --git a/afd_plugin/connectors/routed_experts.py b/afd_plugin/connectors/routed_experts.py index c6a8a99ac..67d6f6264 100644 --- a/afd_plugin/connectors/routed_experts.py +++ b/afd_plugin/connectors/routed_experts.py @@ -74,20 +74,11 @@ class FFNLocalExperts: class AFDRoutedExpertsConnector(AFDConnectorBase): """A connector whose FFN side receives routed (token, local expert) rows. - Lifecycle: the worker creates the connector and calls - ``prepare_before_distributed_init`` before vLLM's distributed init (it may - set process-wide NCCL environment, which NCCL reads when the process - creates its first communicator), then hands the same instance to its - runner, which calls ``init_afd_connector`` after weights load, with what - it loaded (``afd_plugin/v1/worker/connector_lifecycle.py``). The FFN - worker selects the routed-experts runner for these connectors. + Lifecycle: as every connector, the runner creates it and calls + ``init_afd_connector`` after weights load, here with what it loaded. The + FFN worker selects the routed-experts runner for these connectors. """ - @abstractmethod - def prepare_before_distributed_init(self) -> None: - """Collective: join the AFD world before vLLM's distributed init.""" - raise NotImplementedError - # Extends AFDConnectorBase.init_afd_connector with optional keywords only, # so calls through the base contract stay valid. @abstractmethod diff --git a/afd_plugin/distributed/__init__.py b/afd_plugin/distributed/__init__.py index c0777a76e..b7b28853b 100644 --- a/afd_plugin/distributed/__init__.py +++ b/afd_plugin/distributed/__init__.py @@ -2,17 +2,12 @@ # SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project """AFD distributed helper namespace.""" -from afd_plugin.distributed.device import ( - register_assigned_physical_gpu_ids, - resolve_worker_device_index, -) from afd_plugin.distributed.topology import ( AFDRankMapping, build_rank_mapping, resolve_role_rank, subgroup_attention_block, topology_from_config, - tp_rank_before_distributed_init, validate_p2p_topology, ) @@ -37,11 +32,8 @@ def __getattr__(name: str): "build_rank_mapping", "create_hccl_process_group_options", "init_afd_process_group", - "register_assigned_physical_gpu_ids", "resolve_role_rank", - "resolve_worker_device_index", "subgroup_attention_block", "topology_from_config", - "tp_rank_before_distributed_init", "validate_p2p_topology", ] diff --git a/afd_plugin/distributed/device.py b/afd_plugin/distributed/device.py deleted file mode 100644 index 6adab044b..000000000 --- a/afd_plugin/distributed/device.py +++ /dev/null @@ -1,80 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project -"""Worker device resolution before vLLM's ``Worker.init_device``. - -A connector that joins the AFD world before vLLM's distributed init is created -before vLLM selects the worker's device. These helpers mirror that selection -(vLLM v0.30.0 ``vllm/v1/worker/gpu_worker.py``, ``Worker.init_device``, commit -ced6857); the worker checks the result against the device vLLM then selects. -vLLM is imported inside the functions, so importing this module stays -CPU-safe. -""" - -from __future__ import annotations - -from typing import TYPE_CHECKING - -if TYPE_CHECKING: - from vllm.config import VllmConfig - -# vLLM's GPU worker skips its local data-parallel device offset for these. -_NO_DP_DEVICE_OFFSET_EXECUTORS = ("ray", "external_launcher") - - -def register_assigned_physical_gpu_ids(vllm_config: VllmConfig) -> None: - """Register ``--device-ids`` with vLLM's platform, as ``init_device`` does. - - Side effect: sets vLLM's process-wide assigned physical GPU ids, through - which vLLM maps logical device indices. Call it before - ``resolve_worker_device_index``. vLLM's setter accepts the same ids again - when ``init_device`` registers them later; different ids raise. - """ - assigned_physical_gpu_ids = vllm_config.parallel_config.assigned_physical_gpu_ids - if assigned_physical_gpu_ids is None: - return - from vllm.platforms.interface import set_assigned_physical_gpu_ids - - set_assigned_physical_gpu_ids(assigned_physical_gpu_ids) - - -def resolve_worker_device_index(vllm_config: VllmConfig, local_rank: int) -> int: - """CUDA device index vLLM's GPU worker will select for ``local_rank``. - - Adds the local data-parallel offset, then maps the logical index through - the assigned physical ids (``--device-ids``, registered by - ``register_assigned_physical_gpu_ids``) and ``CUDA_VISIBLE_DEVICES``. - """ - from vllm.platforms import current_platform - from vllm.platforms.interface import get_assigned_physical_gpu_ids - - parallel_config = vllm_config.parallel_config - assigned_physical_gpu_ids = parallel_config.assigned_physical_gpu_ids - if ( - assigned_physical_gpu_ids is not None - and get_assigned_physical_gpu_ids() != assigned_physical_gpu_ids - ): - raise RuntimeError( - "call register_assigned_physical_gpu_ids before resolving the " - "worker device: vLLM maps device indices through those ids" - ) - if ( - parallel_config.distributed_executor_backend - not in _NO_DP_DEVICE_OFFSET_EXECUTORS - and parallel_config.data_parallel_backend != "ray" - and parallel_config.nnodes_within_dp == 1 - ): - dp_local_rank = parallel_config.data_parallel_rank_local - if dp_local_rank is None: - dp_local_rank = parallel_config.data_parallel_index - tp_pp_world_size = ( - parallel_config.pipeline_parallel_size - * parallel_config.tensor_parallel_size - ) - local_rank += dp_local_rank * tp_pp_world_size - return current_platform.logical_device_id_to_visible_device_id(local_rank) - - -__all__ = [ - "register_assigned_physical_gpu_ids", - "resolve_worker_device_index", -] diff --git a/afd_plugin/distributed/topology.py b/afd_plugin/distributed/topology.py index 223672b5e..20be8f64c 100644 --- a/afd_plugin/distributed/topology.py +++ b/afd_plugin/distributed/topology.py @@ -58,40 +58,12 @@ def validate_p2p_topology(config: AFDConfig) -> None: ) -def tp_rank_before_distributed_init(vllm_config: VllmConfig, worker_rank: int) -> int: - """TP rank of a worker, before vLLM creates its TP group. - - vLLM orders ranks ``[ExternalDP, DP, PP, PCP, TP]`` with TP innermost - (vLLM v0.30.0 ``initialize_model_parallel``), and AFD runs with PP = PCP = - 1, so the TP rank is the worker rank modulo the TP size whether the worker - rank is DP-local or global. The worker compares the resulting role rank - with vLLM's TP group once it exists (``check_worker_connector``). - """ - parallel_config = vllm_config.parallel_config - if ( - parallel_config.pipeline_parallel_size, - parallel_config.prefill_context_parallel_size, - ) != (1, 1): - raise ValueError("the TP rank before distributed init needs PP = PCP = 1") - return worker_rank % int(parallel_config.tensor_parallel_size) - - -def resolve_role_rank( - vllm_config: VllmConfig, - config: AFDConfig, - *, - tp_rank: int | None = None, -) -> int: +def resolve_role_rank(vllm_config: VllmConfig, config: AFDConfig) -> int: """Resolve this worker's connector-independent AFD role rank. The resolver linearizes vLLM's global DP rank and local PCP/TP coordinates. Connectors receive the result as runtime state and map it to their own communication-world rank. - - Args: - tp_rank: This worker's TP rank, for callers that run before vLLM's - distributed init (see ``tp_rank_before_distributed_init``). - ``None`` reads it from vLLM's TP group. """ parallel_config = vllm_config.parallel_config @@ -110,14 +82,12 @@ def resolve_role_rank( pcp_rank = int(get_pcp_group().rank_in_group) else: pcp_rank = 0 - if tp_size == 1: - tp_rank = 0 - elif tp_rank is None: + if tp_size > 1: from vllm.distributed.parallel_state import get_tensor_model_parallel_rank tp_rank = int(get_tensor_model_parallel_rank()) - elif not 0 <= tp_rank < tp_size: - raise ValueError(f"tp_rank must be in [0, {tp_size}), got {tp_rank}") + else: + tp_rank = 0 role_rank = (dp_rank * pcp_size + pcp_rank) * tp_size + tp_rank if config.role == "attention": diff --git a/afd_plugin/v1/worker/attention_model_runner_v2.py b/afd_plugin/v1/worker/attention_model_runner_v2.py index 2ce70c4a6..610c656a2 100644 --- a/afd_plugin/v1/worker/attention_model_runner_v2.py +++ b/afd_plugin/v1/worker/attention_model_runner_v2.py @@ -367,16 +367,10 @@ class AFDAttentionModelRunnerV2(AFDMetadataProviderMixin, GPUModelRunnerV2): afd_expected_role = "attention" - # Signature: native (vllm_config, device) plus a keyword-only - # ``connector``: the instance the worker created before vLLM's distributed - # init (see connector_lifecycle.py), which the worker keeps owning if - # construction fails. ``None`` creates the connector here. def __init__( self, vllm_config: VllmConfig, device: torch.device, - *, - connector: AFDConnectorBase | None = None, ): validate_gpu_model_runner_v2_config( vllm_config, @@ -384,17 +378,16 @@ def __init__( device_type=device.type, ) super().__init__(vllm_config, device) - owns_connector = connector is None + connector: AFDConnectorBase | None = None try: self.afd_config = self.parse_config(self.vllm_config) - if connector is None: - rank, local_rank = _resolve_world_ranks() - connector = AFDConnectorFactory.create_connector( - rank, - local_rank, - self.vllm_config, - self.afd_config, - ) + rank, local_rank = _resolve_world_ranks() + connector = AFDConnectorFactory.create_connector( + rank, + local_rank, + self.vllm_config, + self.afd_config, + ) self.connector = connector # The connector rendezvous is deferred to the end of ``load_model()`` # so Attention and FFN weight loading overlap, matching the V1 @@ -412,7 +405,7 @@ def __init__( self.prof = create_afd_gpu_profiler("attention") except BaseException: try: - if owns_connector and connector is not None: + if connector is not None: connector.close() finally: super().shutdown() diff --git a/afd_plugin/v1/worker/attention_worker.py b/afd_plugin/v1/worker/attention_worker.py index e108e7d96..a33b9b4ed 100644 --- a/afd_plugin/v1/worker/attention_worker.py +++ b/afd_plugin/v1/worker/attention_worker.py @@ -4,22 +4,15 @@ from __future__ import annotations -from functools import partial - import torch from vllm.config import VllmConfig from vllm.v1.worker.gpu import model_runner as gpu_model_runner_v2 from vllm.v1.worker.gpu_worker import Worker -from afd_plugin.config import parse_afd_config from afd_plugin.model_executor.models.model_utils import get_afd_model_config from afd_plugin.v1.worker.attention_model_runner_v2 import ( AFDAttentionModelRunnerV2, ) -from afd_plugin.v1.worker.connector_lifecycle import ( - check_worker_connector, - create_worker_connector, -) from afd_plugin.validation import ( assert_compatible_afd_stack, validate_gpu_model_runner_v2_config, @@ -51,10 +44,7 @@ def __init__( # after an internal module import, with no injectable runner factory. # Patch functionality: validate the supported V2 contract, select the AFD model # identity before construction, and scope the module-class substitution to - # the delegated native initialization window. A connector that must join - # the AFD world before vLLM's distributed init is created first - # (connector_lifecycle.py), handed to the runner through the seam, and - # checked against vLLM's device and ranks once vLLM has chosen them. + # the delegated native initialization window. # Signature: matches vLLM v0.30.0 Worker.init_device exactly: (self). # Upstream source: vllm/v1/worker/gpu_worker.py, Worker.init_device; # runner local-import/construct seam; commit @@ -89,28 +79,10 @@ def init_device(self): self.vllm_config.model_config = afd_model_config self.model_config = afd_model_config - # Before vLLM's distributed init creates NCCL groups, so the - # connector's NCCL environment applies; the FFN worker does the same. - connector = create_worker_connector( - self.vllm_config, - parse_afd_config(self.vllm_config, expected_role="attention"), - rank=self.rank, - local_rank=self.local_rank, - ) native_runner_cls = gpu_model_runner_v2.GPUModelRunner - gpu_model_runner_v2.GPUModelRunner = ( - AFDAttentionModelRunnerV2 - if connector is None - else partial(AFDAttentionModelRunnerV2, connector=connector) - ) + gpu_model_runner_v2.GPUModelRunner = AFDAttentionModelRunnerV2 try: super().init_device() - if connector is not None: - check_worker_connector(connector, self.vllm_config, self.device) - except BaseException: - if connector is not None: - connector.close() - raise finally: gpu_model_runner_v2.GPUModelRunner = native_runner_cls # ### PATCH END: select and directly construct the AFD runner. diff --git a/afd_plugin/v1/worker/connector_lifecycle.py b/afd_plugin/v1/worker/connector_lifecycle.py deleted file mode 100644 index 9ec01b935..000000000 --- a/afd_plugin/v1/worker/connector_lifecycle.py +++ /dev/null @@ -1,113 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project -"""Connector lifecycle for GPU workers whose connector joins before vLLM. - -Every connector follows one lifecycle: create, then ``init_afd_connector`` -once the runner has loaded the model. A routed-experts connector -(``AFDRoutedExpertsConnector``) also has a step that must run before vLLM's -distributed init: it joins the AFD world and applies its NCCL environment, -which NCCL reads once per process, when the first communicator is created. -vLLM creates its communicators in ``Worker.init_device`` before it builds the -runner, so the worker owns that step: - -1. ``create_worker_connector`` creates the connector and runs - ``prepare_before_distributed_init``. -2. The worker hands that same instance to the runner it constructs. -3. ``check_worker_connector`` compares what was derived before vLLM's init - (device, TP rank) with what vLLM then chose. -4. The runner calls ``init_afd_connector`` after loading the model. - -Other connectors are created by their runner, as before. -""" - -from __future__ import annotations - -from typing import TYPE_CHECKING - -from afd_plugin.config import AFDConfig -from afd_plugin.connectors import ( - AFDConnectorBase, - AFDConnectorFactory, - AFDRoutedExpertsConnector, -) -from afd_plugin.distributed import ( - register_assigned_physical_gpu_ids, - resolve_role_rank, - resolve_worker_device_index, - tp_rank_before_distributed_init, -) - -if TYPE_CHECKING: - import torch - from vllm.config import VllmConfig - - -def create_worker_connector( - vllm_config: VllmConfig, - afd_config: AFDConfig, - *, - rank: int, - local_rank: int, -) -> AFDRoutedExpertsConnector | None: - """Create a connector that must join the AFD world before vLLM's init. - - Call before vLLM's distributed init. Returns ``None`` for connectors that - are not ``AFDRoutedExpertsConnector``; their runner creates them. - Otherwise the connector gets the device vLLM will select and a role rank - whose TP part comes from the worker ``rank`` (vLLM's TP group does not - exist yet), and has run ``prepare_before_distributed_init``. The caller - owns it until it hands it to the runner. - - Side effect: registers vLLM's assigned physical GPU ids - (``register_assigned_physical_gpu_ids``). - """ - connector_cls = AFDConnectorFactory.get_connector_class(afd_config.connector) - if not issubclass(connector_cls, AFDRoutedExpertsConnector): - return None - register_assigned_physical_gpu_ids(vllm_config) - connector = AFDConnectorFactory.create_connector( - rank, - resolve_worker_device_index(vllm_config, local_rank), - vllm_config, - afd_config, - tp_rank=tp_rank_before_distributed_init(vllm_config, rank), - ) - assert isinstance(connector, AFDRoutedExpertsConnector) - try: - connector.prepare_before_distributed_init() - except BaseException: - connector.close() - raise - return connector - - -def check_worker_connector( - connector: AFDConnectorBase, - vllm_config: VllmConfig, - worker_device: torch.device, -) -> None: - """Fail when a connector created before vLLM's init disagrees with vLLM. - - Call after vLLM's distributed init. The connector must use the device vLLM - gave the worker, and its role rank must match vLLM's distributed ranks, so - a vLLM upgrade that changes device selection or rank order fails at - startup instead of moving traffic onto another rank's GPU or tokens. - """ - if connector.local_rank != worker_device.index: - raise RuntimeError( - f"{type(connector).__name__} uses cuda:{connector.local_rank} but " - f"vLLM placed this worker on {worker_device}; the device selection " - "in afd_plugin/distributed/device.py no longer matches vLLM's " - "Worker.init_device", - ) - role_rank = resolve_role_rank(vllm_config, connector.afd_config) - if connector.role_rank != role_rank: - raise RuntimeError( - f"{type(connector).__name__} has role rank {connector.role_rank} but " - f"vLLM's distributed ranks give {role_rank}; the TP rank derived " - "before distributed init (tp_rank_before_distributed_init) no longer " - "matches vLLM's rank layout", - ) - - -__all__ = ["check_worker_connector", "create_worker_connector"] diff --git a/afd_plugin/v1/worker/ffn_model_runner.py b/afd_plugin/v1/worker/ffn_model_runner.py index 918fb746c..9f6478d09 100644 --- a/afd_plugin/v1/worker/ffn_model_runner.py +++ b/afd_plugin/v1/worker/ffn_model_runner.py @@ -25,7 +25,6 @@ ) from afd_plugin.config import AFDConfig, parse_afd_config from afd_plugin.connectors import ( - AFDConnectorBase, AFDConnectorFactory, AFDControlPayload, AFDControlPlane, @@ -61,15 +60,7 @@ class GPUFFNModelRunner: afd_expected_role = "ffn" - # ``connector``: the instance the worker created before vLLM's distributed - # init (see connector_lifecycle.py); ``None`` creates the connector here. - def __init__( - self, - vllm_config: VllmConfig, - device: torch.device, - *, - connector: AFDConnectorBase | None = None, - ) -> None: + def __init__(self, vllm_config: VllmConfig, device: torch.device) -> None: self.vllm_config = vllm_config self.model_config = vllm_config.model_config self.load_config = vllm_config.load_config @@ -85,15 +76,13 @@ def __init__( vllm_config, role="ffn", ) - if connector is None: - rank, local_rank = _resolve_world_ranks() - connector = AFDConnectorFactory.create_connector( - rank, - local_rank, - vllm_config, - self.afd_config, - ) - self.connector = connector + rank, local_rank = _resolve_world_ranks() + self.connector = AFDConnectorFactory.create_connector( + rank, + local_rank, + vllm_config, + self.afd_config, + ) # TODO: Async GPU connector will be supported in the future assert self.connector.control_plane is not None, ( "GPU model runner only supports control-plane-driven connectors" diff --git a/afd_plugin/v1/worker/ffn_routed_experts_model_runner.py b/afd_plugin/v1/worker/ffn_routed_experts_model_runner.py index 87c8a58aa..85a848c8d 100644 --- a/afd_plugin/v1/worker/ffn_routed_experts_model_runner.py +++ b/afd_plugin/v1/worker/ffn_routed_experts_model_runner.py @@ -116,14 +116,8 @@ class GPUFFNRoutedExpertsModelRunner(GPUFFNModelRunner): (DeepEP M2N). """ - def __init__( - self, - vllm_config: VllmConfig, - device: torch.device, - *, - connector: AFDConnectorBase | None = None, - ) -> None: - super().__init__(vllm_config, device, connector=connector) + def __init__(self, vllm_config: VllmConfig, device: torch.device) -> None: + super().__init__(vllm_config, device) _require_routed(self.connector) self.moe_layer_indices: tuple[int, ...] = () self.local_experts: dict[int, FFNLocalExperts] = {} diff --git a/afd_plugin/v1/worker/ffn_worker.py b/afd_plugin/v1/worker/ffn_worker.py index 4cd249bc1..f09778a96 100644 --- a/afd_plugin/v1/worker/ffn_worker.py +++ b/afd_plugin/v1/worker/ffn_worker.py @@ -6,7 +6,6 @@ import logging import threading -from functools import partial from typing import TYPE_CHECKING import torch @@ -18,10 +17,6 @@ from afd_plugin.config import parse_afd_config from afd_plugin.connectors import AFDConnectorFactory, AFDRoutedExpertsConnector from afd_plugin.model_executor.models.model_utils import get_afd_model_config -from afd_plugin.v1.worker.connector_lifecycle import ( - check_worker_connector, - create_worker_connector, -) from afd_plugin.v1.worker.ffn_model_runner import GPUFFNModelRunner from afd_plugin.validation import ( assert_compatible_afd_stack, @@ -69,11 +64,7 @@ def __init__( # symbol inside Worker.init_device and exposes no injectable runner factory. # Patch functionality: replace only the V2 construction seam with the # connector-driven FFN runner: GPUFFNRoutedExpertsModelRunner (imported - # only then) for an AFDRoutedExpertsConnector, else GPUFFNModelRunner. A - # connector that must join the AFD world before vLLM's distributed init is - # created first (connector_lifecycle.py), so its NCCL environment precedes - # vLLM's NCCL groups, handed to the runner through the seam, and checked - # against vLLM's device and ranks once vLLM has chosen them. + # only then) for an AFDRoutedExpertsConnector, else GPUFFNModelRunner. # Signature: matches vLLM v0.30.0 Worker.init_device exactly: (self). # Upstream source: vllm/v1/worker/gpu_worker.py, Worker.init_device; # ced6857afa0ea7b2e3f0846a62e1394e90f15607. The 0.30.0 init_device still @@ -114,29 +105,12 @@ def init_device(self): ) runner_cls = GPUFFNRoutedExpertsModelRunner - # Before the connector joins the AFD world: a rejected backend - # fails before any collective. + # Before vLLM builds the model, which reads the selected backend. select_routed_moe_backend(self.vllm_config) - connector = create_worker_connector( - self.vllm_config, - afd_config, - rank=self.rank, - local_rank=self.local_rank, - ) native_runner_cls = gpu_model_runner_v2.GPUModelRunner - gpu_model_runner_v2.GPUModelRunner = ( - runner_cls - if connector is None - else partial(runner_cls, connector=connector) - ) + gpu_model_runner_v2.GPUModelRunner = runner_cls try: super().init_device() - if connector is not None: - check_worker_connector(connector, self.vllm_config, self.device) - except BaseException: - if connector is not None: - connector.close() - raise finally: gpu_model_runner_v2.GPUModelRunner = native_runner_cls # ### PATCH END: construct the FFN runner through the V2 seam. diff --git a/afd_plugin/validation.py b/afd_plugin/validation.py index 27ba28b5a..bca1fa743 100644 --- a/afd_plugin/validation.py +++ b/afd_plugin/validation.py @@ -109,15 +109,14 @@ def validate_gpu_model_runner_v2_config( ) parallel = vllm_config.parallel_config - # DeepEP M2N joins the AFD world before vLLM's distributed init and derives - # the TP part of the role rank from the worker rank. Attention may run TP > 1: - # each TP rank dispatches its own slice of the tokens. FFN experts are not - # TP-sharded; that and the other deployment checks (expert placement, - # all2all backend, weight format) run in the connector handshake so that - # every rank fails together. It routes on Attention, so it requires the - # gate there, and a model whose Attention side sends the top-k. Workers - # validate before the connector joins the AFD world, so both roles fail before any - # collective instead of one role timing out. + # DeepEP M2N: Attention may run TP > 1, each TP rank dispatching its own + # slice of the tokens. FFN experts are not TP-sharded; that and the other + # deployment checks (expert placement, all2all backend, weight format) run + # in the connector handshake so that every rank fails together. It routes + # on Attention, so it requires the gate there, and a model whose Attention + # side sends the top-k. Workers validate before the connector joins the AFD + # world, so both roles fail before any collective instead of one role + # timing out. if afd_config.connector in DEEPEP_M2N_CONNECTORS: if expected_role == "ffn" and parallel.tensor_parallel_size != 1: raise RuntimeError( diff --git a/docs/design/module/connector_contracts.md b/docs/design/module/connector_contracts.md index 27c5c809b..f3c5275ed 100644 --- a/docs/design/module/connector_contracts.md +++ b/docs/design/module/connector_contracts.md @@ -22,7 +22,6 @@ validation_paths: - "tests/e2e/models/deepseek_v2_lite/test_deepseek_v2_lite.py" - "tests/e2e/models/deepseek_v2_lite/test_async_cam_npu.py" - "tests/unit/distributed/**" - - "tests/unit/v1/worker/test_connector_lifecycle.py" - "tests/gpu/**" upstream_refs: - "vLLM vllm.forward_context.DPMetadata" diff --git a/tests/gpu/deepep_m2n_harness.py b/tests/gpu/deepep_m2n_harness.py index 9afc18264..b193f5428 100644 --- a/tests/gpu/deepep_m2n_harness.py +++ b/tests/gpu/deepep_m2n_harness.py @@ -406,21 +406,17 @@ def make_connector( port: int, host: str, ) -> Any: - """The connector as the runtime constructs it (not yet initialized).""" + """The connector class the factory resolves, not yet initialized.""" from afd_plugin.config import afd_config_from_mapping from afd_plugin.connectors import AFDConnectorFactory role, role_rank = topology.role(world_rank) vllm_config = m2n_vllm_config(topology, role, role_rank, extra_config, port, host) afd_config = afd_config_from_mapping(vllm_config.additional_config["afd"]) - # No vLLM TP group here: pass the TP rank, as workers do before init. - return AFDConnectorFactory.create_connector( - world_rank, - local_rank, - vllm_config, - afd_config, - tp_rank=role_rank % vllm_config.parallel_config.tensor_parallel_size, - ) + # No vLLM process groups here, so pass the role rank the topology gives + # instead of resolving it from vLLM's ranks as the factory does. + connector_cls = AFDConnectorFactory.get_connector_class(afd_config.connector) + return connector_cls(world_rank, local_rank, vllm_config, afd_config, role_rank) def control_step( @@ -797,10 +793,7 @@ def finish(self) -> dict[str, Any]: def network_plan_summary(self) -> str: plan = self.connector.network_plan - if plan is None: - return "" - nic = plan.nic_by_rank[self.connector.world_rank] - return f"{plan.summary()} this_rank_nic={nic or 'nccl'}" + return "" if plan is None else plan.summary() def attention_context(layer: int, stage: int, num_tokens: int) -> Any: diff --git a/tests/unit/connectors/test_deepep_m2n_connector.py b/tests/unit/connectors/test_deepep_m2n_connector.py index 7444ba8b5..6d9ae4d23 100644 --- a/tests/unit/connectors/test_deepep_m2n_connector.py +++ b/tests/unit/connectors/test_deepep_m2n_connector.py @@ -132,7 +132,6 @@ def _handshake( world_rank=world_rank, host=hosts[world_rank], nvlink_domain=f"host:{hosts[world_rank]}", - local_index=hosts[: world_rank + 1].count(hosts[world_rank]) - 1, ), ) infos.append(info) @@ -179,7 +178,9 @@ def test_config_defaults_and_round_trip(): ({"bucket_granularity": 0}, "bucket_granularity"), ({"cpu_sync_min_tokens": -1}, "cpu_sync_min_tokens"), ({"control_backend": "mpi"}, "control_backend"), - ({"network": {"preset": "token-ring"}}, "network.preset"), + ({"network": {"gin": "maybe"}}, "network.gin"), + # The connector no longer plans NICs: NCCL picks them. + ({"network": {"nic_policy": "shared"}}, "unknown network options"), ({"num_qps": -1}, "num_qps"), ({"expert_num_sms": -2}, "expert_num_sms"), ({"expert_num_sms": 131}, "expert_num_sms"), @@ -190,7 +191,7 @@ def test_config_defaults_and_round_trip(): ) def test_config_rejects_invalid_options(raw, message): """Each bad or stale option (unknown or removed key, bad enum, zero or - negative size, bad network preset, bad timeout) raises a ValueError + negative size, bad network option, bad timeout) raises a ValueError naming the field. Why: typos and removed keys must fail at startup, not be ignored. @@ -756,21 +757,21 @@ def test_gloo_control_backend_stages_payloads_on_cpu(): def test_config_parses_network_section_and_tunables(): - """The ``network`` section (preset plus NIC), ``num_qps`` and timeouts - parse and survive a round trip. + """The ``network`` section, ``num_qps`` and timeouts parse and survive a + round trip. - Why: the Phase 4b options must reach the connector intact. + Why: the options must reach the connector intact. """ config = DeepEPM2NConfig.from_mapping( { - "network": {"preset": "roce-rail-isolated", "nics": ["mlx5_0"]}, + "network": {"gin": "on", "hybrid": "off", "rdma_traffic_class": 5}, "num_qps": 33, "gpu_timeout_s": 60, } ) - assert config.network.nic_policy == "shared" - assert config.network.nics == ("mlx5_0",) + assert (config.network.gin, config.network.hybrid) == ("on", "off") + assert config.network.rdma_traffic_class == 5 assert (config.num_qps, config.gpu_timeout_s) == (33, 60) assert DeepEPM2NConfig.from_mapping(config.to_mapping()) == config @@ -822,7 +823,7 @@ def test_control_plane_requires_the_handshake_on_ffn(): is_warmup=False, ) - with pytest.raises(RuntimeError, match="before prepare_before_distributed_init"): + with pytest.raises(RuntimeError, match="before the handshake"): connector.control_plane.update_state_from_dp_metadata(payload) diff --git a/tests/unit/connectors/test_deepep_m2n_network.py b/tests/unit/connectors/test_deepep_m2n_network.py index 7d0cc83d8..c438bd94c 100644 --- a/tests/unit/connectors/test_deepep_m2n_network.py +++ b/tests/unit/connectors/test_deepep_m2n_network.py @@ -13,183 +13,69 @@ from afd_plugin.connectors.gpu.deepep_m2n.network import ( # noqa: E402 NetworkConfig, - NicInfo, RankTopology, - apply_rank_environment, - discover_nics, + apply_gin_environment, plan_network, ) -RAIL_A = "roce:172.22.176.0/21" -RAIL_B = "roce:172.22.136.0/21" -RAIL_C = "roce:172.22.200.0/21" -IB_RAIL = "ib:fe80:0000:0000:0000" - -def _rank(rank: int, host: str, local: int, nics=(), domain: str | None = None): +def _rank(rank: int, host: str, domain: str | None = None) -> RankTopology: return RankTopology( - world_rank=rank, - host=host, - nvlink_domain=domain or f"host:{host}", - local_index=local, - nics=tuple(nics), + world_rank=rank, host=host, nvlink_domain=domain or f"host:{host}" ) -def _nic(name: str, rail: str, distance: int) -> NicInfo: - return NicInfo(name=name, link_layer="ethernet", rail=rail, distance=distance) - - -def _isolated_rail_two_nodes() -> list[RankTopology]: - """Every NIC on its own rail; the nearest NIC differs per rank and host.""" - node0 = [ - [ - _nic("mlx5_0", RAIL_A, 2), - _nic("mlx5_5", RAIL_B, 6), - _nic("mlx5_3", RAIL_C, 6), - ], - [ - _nic("mlx5_0", RAIL_A, 6), - _nic("mlx5_5", RAIL_B, 6), - _nic("mlx5_3", RAIL_C, 2), - ], - ] - node1 = [ - [ - _nic("mlx5_0", RAIL_A, 2), - _nic("mlx5_5", RAIL_B, 6), - _nic("mlx5_3", RAIL_C, 6), - ], - [ - _nic("mlx5_0", RAIL_A, 6), - _nic("mlx5_5", RAIL_B, 2), - _nic("mlx5_3", RAIL_C, 6), - ], - ] - return [ - _rank(0, "n0", 0, node0[0]), - _rank(1, "n0", 1, node0[1]), - _rank(2, "n1", 0, node1[0]), - _rank(3, "n1", 1, node1[1]), - ] - - -def test_presets_expand_and_explicit_keys_override(): - """Presets fill in defaults, explicit keys override them, and the config - survives a round trip. - - Why: presets are defaults, not locks. - """ - rail_isolated = NetworkConfig.from_mapping({"preset": "roce-rail-isolated"}) - overridden = NetworkConfig.from_mapping({"preset": "ib", "nic_policy": "shared"}) +def _two_hosts() -> list[RankTopology]: + return [_rank(r, f"n{r // 2}") for r in range(4)] + + +def test_network_config_survives_a_round_trip(): + config = NetworkConfig.from_mapping( + {"gin": "on", "hybrid": "off", "rdma_traffic_class": 5} + ) - assert (rail_isolated.nic_policy, rail_isolated.rails_routed) == ("shared", False) - assert (overridden.nic_policy, overridden.rails_routed) == ("shared", True) - assert NetworkConfig.from_mapping(rail_isolated.to_mapping()) == rail_isolated + assert NetworkConfig.from_mapping(config.to_mapping()) == config @pytest.mark.parametrize( ("raw", "message"), [ ({"gin": "maybe"}, "network.gin"), - ({"nic_policy": "explicit"}, "needs network.nics"), - ({"nccl_env": {"PATH": "/tmp"}}, "nccl_env key"), - ({"nccl_env": {"EP_DISABLE_GIN": "1"}}, "use network.gin"), + ({"hybrid": "maybe"}, "network.hybrid"), + ({"rdma_traffic_class": -1}, "rdma_traffic_class"), + # NICs are NCCL's own choice; the connector plans none. + ({"nic_policy": "shared"}, "unknown network options"), + ({"nccl_env": {"NCCL_IB_HCA": "mlx5_0"}}, "unknown network options"), ({"bogus": 1}, "unknown network options"), ], ) def test_network_config_rejects_invalid_values(raw, message): - """Bad ``gin``, ``explicit`` without NICs, non-NCCL env keys, - ``EP_DISABLE_GIN`` in ``nccl_env`` and unknown keys are rejected. + """Bad ``gin``/``hybrid``/traffic class and unknown keys are rejected. - Why: bad or conflicting network settings must fail at parse time. + Why: bad or removed network settings must fail at parse time. """ with pytest.raises(ValueError, match=message): NetworkConfig.from_mapping(raw) def test_single_host_needs_no_rdma(): - """Four ranks on one host: one domain, GIN off, no NIC pins. + """Four ranks on one host: one domain, GIN off. Why: single-node runs must never touch RDMA. """ - plan = plan_network( - [_rank(r, "n0", r, [_nic("mlx5_0", RAIL_A, 2)]) for r in range(4)], - NetworkConfig(), - ) - - assert (plan.num_domains, plan.use_gin) == (1, False) - assert plan.nic_by_rank == (None,) * 4 - - -def test_isolated_rails_put_every_rank_on_one_shared_rail(): - """Hosts whose NICs each sit on their own subnet get the shared policy, - one rail, and the same NIC for every rank. - - Why: encodes the Phase 4 finding as the automatic choice. - """ - plan = plan_network(_isolated_rail_two_nodes(), NetworkConfig()) - - assert (plan.use_gin, plan.nic_policy, plan.rail) == (True, "shared", RAIL_A) - assert plan.nic_by_rank == ("mlx5_0",) * 4 - - -def test_shared_rail_uses_each_hosts_own_nic_name(): - """The shared subnet is ``mlx5_0`` on one host and ``mlx5_7`` on the - other; each rank gets its own host's name. - - Why: NIC names are not consistent across hosts. - """ - ranks = [ - _rank(0, "n0", 0, [_nic("mlx5_0", RAIL_A, 2)]), - _rank(1, "n1", 0, [_nic("mlx5_7", RAIL_A, 2), _nic("mlx5_0", RAIL_B, 2)]), - ] - - plan = plan_network(ranks, NetworkConfig()) - - assert plan.nic_by_rank == ("mlx5_0", "mlx5_7") + plan = plan_network([_rank(r, "n0") for r in range(4)], NetworkConfig()) - -def test_shared_rail_prefers_the_configured_nic(): - """``nics: [mlx5_5]`` selects that NIC's subnet as the shared rail. - - Why: operators can steer which NIC is shared. - """ - plan = plan_network( - _isolated_rail_two_nodes(), NetworkConfig(nic_policy="shared", nics=("mlx5_5",)) - ) - - assert (plan.rail, plan.nic_by_rank) == (RAIL_B, ("mlx5_5",) * 4) - - -@pytest.mark.parametrize( - "config", - [NetworkConfig(rails_routed=True), NetworkConfig.from_mapping({"preset": "ib"})], - ids=["routed-roce", "ib-preset"], -) -def test_routed_fabrics_keep_nccls_nearest_nic(config): - """``rails_routed`` and the ``ib`` preset keep NCCL's per-GPU NIC choice - (no pins) with GIN on. - - Why: routed fabrics should use every NIC, not throttle to one. - """ - plan = plan_network(_isolated_rail_two_nodes(), config) - - assert (plan.use_gin, plan.nic_policy) == (True, "nccl") - assert plan.nic_by_rank == (None,) * 4 + assert (plan.num_domains, plan.ranks_per_domain, plan.use_gin) == (1, 4, False) -def test_ib_fat_tree_is_one_rail(): - """IB NICs on one subnet prefix get the NCCL policy without pins. +def test_ranks_on_several_hosts_use_gin(): + """Two hosts of two ranks: two domains, GIN on. - Why: one rail needs no pinning. + Why: GIN (RDMA) is DeepEP's only path between NVLink domains. """ - ib_nic = NicInfo(name="mlx5_0", link_layer="infiniband", rail=IB_RAIL, distance=2) - plan = plan_network( - [_rank(r, f"n{r // 2}", r % 2, [ib_nic]) for r in range(4)], NetworkConfig() - ) + plan = plan_network(_two_hosts(), NetworkConfig()) - assert (plan.use_gin, plan.nic_policy) == (True, "nccl") + assert (plan.num_domains, plan.ranks_per_domain, plan.use_gin) == (2, 2, True) def test_nvl72_is_one_domain_even_with_uneven_hosts(): @@ -199,182 +85,43 @@ def test_nvl72_is_one_domain_even_with_uneven_hosts(): Why: NVL72 is one domain; per-host layout rules must not apply. """ - ranks = [ - _rank(r, f"tray{0 if r < 3 else 1}", r, domain="nvl:abc:1") for r in range(4) - ] + ranks = [_rank(r, f"tray{0 if r < 3 else 1}", domain="nvl:abc:1") for r in range(4)] plan = plan_network(ranks, NetworkConfig()) assert (plan.num_domains, plan.ranks_per_domain, plan.use_gin) == (1, 4, False) -def test_explicit_nics_follow_local_rank(): - """``explicit`` assigns ``nics`` by local rank on every host. - - Why: a manual override for fabrics the planner cannot model. - """ - plan = plan_network( - _isolated_rail_two_nodes(), - NetworkConfig(nic_policy="explicit", nics=("mlx5_0", "mlx5_5")), - ) - - assert plan.nic_by_rank == ("mlx5_0", "mlx5_5", "mlx5_0", "mlx5_5") - - -def test_no_common_rail_is_an_error(): - """Hosts with no subnet in common raise "no RDMA rail". - - Why: fail at init instead of a QP timeout minutes later. - """ - ranks = [ - _rank(0, "n0", 0, [_nic("mlx5_0", RAIL_A, 2)]), - _rank(1, "n1", 0, [_nic("mlx5_0", RAIL_B, 2)]), - ] - - with pytest.raises(RuntimeError, match="no RDMA rail"): - plan_network(ranks, NetworkConfig()) - - -@pytest.mark.parametrize( - "config", - [NetworkConfig(gin="off"), NetworkConfig.from_mapping({"preset": "single-node"})], - ids=["gin-off", "single-node-preset"], -) -def test_gin_off_across_nvlink_domains_is_an_error(config): +def test_gin_off_across_nvlink_domains_is_an_error(): # NVLink load/store cannot reach another domain; GIN is the only path. - """``gin: off`` (or the ``single-node`` preset) with ranks on two hosts - raises. + """``gin: off`` with ranks on two hosts raises. Why: without GIN DeepEP cannot reach another NVLink domain. """ with pytest.raises(RuntimeError, match="span 2 NVLink domains"): - plan_network(_isolated_rail_two_nodes(), config) + plan_network(_two_hosts(), NetworkConfig(gin="off")) -def test_gin_off_on_one_host_is_allowed(): - """``gin: off`` with all ranks on one host plans normally. +@pytest.mark.parametrize(("gin", "use_gin"), [("off", False), ("on", True)]) +def test_forced_gin_on_one_host_is_honored(gin, use_gin): + """``gin: off``/``on`` with all ranks on one host plans as forced. Why: the multi-domain check must not block single-node use. """ - plan = plan_network( - [_rank(r, "n0", r, [_nic("mlx5_0", RAIL_A, 2)]) for r in range(2)], - NetworkConfig(gin="off"), - ) - - assert not plan.use_gin - - -def test_forced_gin_on_one_domain_still_plans_nics(): - """``gin: on`` on one host keeps GIN on and pins no NICs. - - Why: a forced setting is honored without inventing NIC pins. - """ - plan = plan_network( - [_rank(r, "n0", r, [_nic("mlx5_0", RAIL_A, 2)]) for r in range(2)], - NetworkConfig(gin="on"), - ) + plan = plan_network([_rank(r, "n0") for r in range(2)], NetworkConfig(gin=gin)) - assert plan.use_gin - assert plan.nic_by_rank == (None, None) + assert plan.use_gin is use_gin -def test_rank_environment_pins_nic_and_gin(monkeypatch): - """Applying a plan sets ``NCCL_IB_HCA``, ``nccl_env`` values and - ``EP_DISABLE_GIN`` (absent with GIN, 1 without). +def test_gin_environment_follows_the_plan(monkeypatch): + """``EP_DISABLE_GIN`` is 1 without GIN and absent with it. - Why: the plan only works if these reach NCCL and DeepEP. + Why: DeepEP reads it when the connector creates its buffers. """ - for key in ("NCCL_IB_HCA", "EP_DISABLE_GIN", "NCCL_IB_TC"): - monkeypatch.delenv(key, raising=False) - config = NetworkConfig(nccl_env=(("NCCL_IB_TC", "106"),)) - plan = plan_network(_isolated_rail_two_nodes(), config) - - apply_rank_environment(plan, config, world_rank=3) - - assert os.environ["NCCL_IB_HCA"] == "=mlx5_0" - assert os.environ["NCCL_IB_TC"] == "106" - assert "EP_DISABLE_GIN" not in os.environ + monkeypatch.delenv("EP_DISABLE_GIN", raising=False) - single = plan_network([_rank(0, "n0", 0)], NetworkConfig()) - apply_rank_environment(single, NetworkConfig(), world_rank=0) + apply_gin_environment(plan_network([_rank(0, "n0")], NetworkConfig())) assert os.environ["EP_DISABLE_GIN"] == "1" - -def _fake_sysfs(tmp_path): - """Two RoCE ports (near and far from the GPU) and one IB port.""" - pci_root = tmp_path / "devices" - near_bridge = pci_root / "pci0000:00" / "0000:00:01.0" - far_bridge = pci_root / "pci0000:80" / "0000:80:01.0" - functions = { - "0000:18:00.0": near_bridge / "0000:18:00.0", - "0000:19:00.0": near_bridge / "0000:19:00.0", - "0000:9b:00.0": far_bridge / "0000:9b:00.0", - "0000:aa:00.0": far_bridge / "0000:aa:00.0", - } - sysfs_pci = tmp_path / "bus_pci" - sysfs_pci.mkdir() - for bdf, path in functions.items(): - path.mkdir(parents=True) - (sysfs_pci / bdf).symlink_to(path) - sysfs_ib = tmp_path / "infiniband" - ports = { - "mlx5_0": ( - "0000:19:00.0", - "Ethernet", - "0000:0000:0000:0000:0000:ffff:ac16:b021", - ), - "mlx5_10": ( - "0000:9b:00.0", - "Ethernet", - "0000:0000:0000:0000:0000:ffff:ac16:a021", - ), - "mlx5_12": ( - "0000:aa:00.0", - "InfiniBand", - "fe80:0000:0000:0000:0e42:a1ff:fe00:0001", - ), - } - for name, (bdf, link_layer, gid) in ports.items(): - port = sysfs_ib / name / "ports" / "1" - (port / "gids").mkdir(parents=True) - (port / "state").write_text("4: ACTIVE\n") - (port / "link_layer").write_text(f"{link_layer}\n") - (port / "gids" / "0").write_text( - "fe80:0000:0000:0000:0000:0000:0000:0000\n" - if link_layer == "Ethernet" - else gid - ) - if link_layer == "Ethernet": - (port / "gids" / "3").write_text(f"{gid}\n") - (sysfs_ib / name / "device").symlink_to(functions[bdf]) - down = sysfs_ib / "mlx5_6" / "ports" / "1" - (down / "gids").mkdir(parents=True) - (down / "state").write_text("1: DOWN\n") - return sysfs_ib, sysfs_pci - - -def test_discovery_reads_rails_and_pcie_distance(tmp_path): - """A fake sysfs tree: active RoCE and IB ports found (down port skipped), - rails from interface subnets, GID prefix or /32, and PCIe distance - ordering near before far. - - Why: topology discovery is tested without RDMA hardware. - """ - sysfs_ib, sysfs_pci = _fake_sysfs(tmp_path) - - nics = { - nic.name: nic - for nic in discover_nics( - "0000:18:00.0", - sysfs_infiniband=sysfs_ib, - sysfs_pci=sysfs_pci, - interface_networks={"172.22.176.33": "172.22.176.0/21"}, - ) - } - - assert set(nics) == {"mlx5_0", "mlx5_10", "mlx5_12"} - assert nics["mlx5_0"].rail == "roce:172.22.176.0/21" - # No netmask known for this address: it is its own (isolated) rail. - assert nics["mlx5_10"].rail == "roce:172.22.160.33/32" - assert nics["mlx5_12"].rail == IB_RAIL - assert nics["mlx5_0"].distance < nics["mlx5_10"].distance + apply_gin_environment(plan_network(_two_hosts(), NetworkConfig())) + assert "EP_DISABLE_GIN" not in os.environ diff --git a/tests/unit/distributed/test_device.py b/tests/unit/distributed/test_device.py deleted file mode 100644 index 26200d2f7..000000000 --- a/tests/unit/distributed/test_device.py +++ /dev/null @@ -1,85 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project -"""Worker device resolution matches ``Worker.init_device`` before it runs.""" - -from __future__ import annotations - -import pytest - -pytest.importorskip("vllm") - -from afd_plugin.distributed import ( # noqa: E402 - register_assigned_physical_gpu_ids, - resolve_worker_device_index, -) -from tests.unit.fakes.platform import use_cuda_like_platform # noqa: E402 -from tests.unit.fakes.vllm_config import gpu_worker_vllm_config # noqa: E402 - - -@pytest.fixture(autouse=True) -def _cuda_like_platform(monkeypatch): - use_cuda_like_platform(monkeypatch) - - -def _resolve(config, local_rank): - register_assigned_physical_gpu_ids(config) - return resolve_worker_device_index(config, local_rank) - - -@pytest.mark.parametrize( - ("executor_backend", "expected_index"), - [("mp", 1), ("ray", 0), ("external_launcher", 0)], -) -def test_adds_local_dp_offset_like_vllm(executor_backend, expected_index): - """The local DP rank offsets the device, except where vLLM skips it. - - Why: vLLM adds the offset only inside init_device, after the connector - joined the AFD world; without it every DP rank's connector used GPU 0. - """ - config = gpu_worker_vllm_config( - data_parallel_rank_local=1, - data_parallel_index=1, - distributed_executor_backend=executor_backend, - ) - - assert _resolve(config, 0) == expected_index - - -def test_maps_assigned_device_ids_without_visible_devices(): - """``--device-ids 4,5`` without CUDA_VISIBLE_DEVICES places rank 0 on cuda:4. - - Why: vLLM maps its logical index through the assigned physical ids, so a - connector using the logical index would run on another rank's GPU. - """ - config = gpu_worker_vllm_config(assigned_physical_gpu_ids=[4, 5]) - - assert _resolve(config, 0) == 4 - assert _resolve(config, 1) == 5 - - -def test_maps_assigned_device_ids_through_visible_devices(monkeypatch): - """With CUDA_VISIBLE_DEVICES the assigned ids map to visible ordinals.""" - monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "4,5") - config = gpu_worker_vllm_config(assigned_physical_gpu_ids=[5, 4]) - - assert _resolve(config, 0) == 1 - assert _resolve(config, 1) == 0 - - -def test_without_assigned_ids_the_index_is_the_local_rank(monkeypatch): - """The default mp launch keeps the logical index as the device index.""" - monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "2,3") - - assert _resolve(gpu_worker_vllm_config(), 1) == 1 - - -def test_resolving_before_registering_assigned_ids_fails(): - """Resolving with ``--device-ids`` set but not registered raises. - - Why: vLLM would map the index without those ids and place the connector - on another rank's GPU; the side effect must be explicit and come first. - """ - config = gpu_worker_vllm_config(assigned_physical_gpu_ids=[4, 5]) - - with pytest.raises(RuntimeError, match="register_assigned_physical_gpu_ids"): - resolve_worker_device_index(config, 0) diff --git a/tests/unit/distributed/test_topology.py b/tests/unit/distributed/test_topology.py index bdfcfca89..fbece0ca6 100644 --- a/tests/unit/distributed/test_topology.py +++ b/tests/unit/distributed/test_topology.py @@ -9,7 +9,7 @@ import pytest from afd_plugin.config import AFDConfig -from afd_plugin.distributed import resolve_role_rank, tp_rank_before_distributed_init +from afd_plugin.distributed import resolve_role_rank def _vllm_config( @@ -110,65 +110,3 @@ def test_resolve_role_rank_rejects_derived_rank_outside_role_size(monkeypatch): _vllm_config(dp_size=2, dp_rank=1, tp_size=2), _afd_config(num_attention_ranks=2), ) - - -def test_resolve_role_rank_uses_given_tp_rank_without_vllm_groups(monkeypatch): - """A TP rank passed by the caller replaces vLLM's TP group lookup. - - Why: routed-experts connectors resolve their role rank before vLLM's - distributed init, when the TP group does not exist yet. - """ - monkeypatch.setitem(sys.modules, "vllm.distributed.parallel_state", None) - - role_rank = resolve_role_rank( - _vllm_config(dp_size=2, dp_rank=1, tp_size=2), - _afd_config(num_attention_ranks=4), - tp_rank=1, - ) - - assert role_rank == 3 - - -def test_resolve_role_rank_rejects_given_tp_rank_outside_tp_size(): - with pytest.raises(ValueError, match="tp_rank must be in"): - resolve_role_rank( - _vllm_config(tp_size=2), - _afd_config(num_attention_ranks=2), - tp_rank=2, - ) - - -def _pre_init_config(*, tp_size: int, pp_size: int = 1, pcp_size: int = 1): - return SimpleNamespace( - parallel_config=SimpleNamespace( - tensor_parallel_size=tp_size, - pipeline_parallel_size=pp_size, - prefill_context_parallel_size=pcp_size, - ), - ) - - -@pytest.mark.parametrize( - ("worker_rank", "tp_size", "expected"), - [(0, 1, 0), (3, 1, 0), (0, 2, 0), (1, 2, 1), (5, 4, 1)], -) -def test_tp_rank_before_distributed_init_is_worker_rank_mod_tp( - worker_rank, tp_size, expected -): - """TP is vLLM's innermost rank dimension, so the worker rank gives it. - - Why: this holds for DP-local and global worker ranks alike; - ``check_worker_connector`` re-checks it after vLLM's init. - """ - assert ( - tp_rank_before_distributed_init(_pre_init_config(tp_size=tp_size), worker_rank) - == expected - ) - - -@pytest.mark.parametrize(("pp_size", "pcp_size"), [(2, 1), (1, 2)]) -def test_tp_rank_before_distributed_init_rejects_pp_and_pcp(pp_size, pcp_size): - with pytest.raises(ValueError, match="PP = PCP = 1"): - tp_rank_before_distributed_init( - _pre_init_config(tp_size=2, pp_size=pp_size, pcp_size=pcp_size), 1 - ) diff --git a/tests/unit/fakes/connector.py b/tests/unit/fakes/connector.py index 485e7de43..5795633fd 100644 --- a/tests/unit/fakes/connector.py +++ b/tests/unit/fakes/connector.py @@ -2,8 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project """A recording ``AFDRoutedExpertsConnector`` for worker and runner tests. -Constructed without arguments it skips ``AFDConnectorBase.__init__`` (no vLLM -config needed); the factory constructs it with the base arguments. +Constructed without ``AFDConnectorBase.__init__`` (no vLLM config needed). ``events`` records the lifecycle and data-path calls in order. """ @@ -20,12 +19,7 @@ class FakeRoutedConnector(AFDRoutedExpertsConnector): - # Raised by prepare_before_distributed_init when set. - prepare_error: BaseException | None = None - - def __init__(self, *base_args: Any) -> None: - if base_args: - super().__init__(*base_args) + def __init__(self) -> None: self.events: list[tuple] = [] self.control_updates: list[Any] = [] self.init_kwargs: dict[str, Any] | None = None @@ -43,11 +37,6 @@ def is_initialized(self) -> bool: def close(self) -> None: self.events.append(("close",)) - def prepare_before_distributed_init(self) -> None: - self.events.append(("prepare",)) - if self.prepare_error is not None: - raise self.prepare_error - def init_afd_connector(self, *, moe_layer_indices=None, local_experts=None): self.init_kwargs = { "moe_layer_indices": moe_layer_indices, diff --git a/tests/unit/fakes/platform.py b/tests/unit/fakes/platform.py deleted file mode 100644 index 4ad40fe35..000000000 --- a/tests/unit/fakes/platform.py +++ /dev/null @@ -1,22 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project -"""A CUDA-like vLLM platform for device-selection tests on hosts without GPUs.""" - -from __future__ import annotations - -import pytest - - -def use_cuda_like_platform(monkeypatch: pytest.MonkeyPatch) -> None: - """Map device ids like vLLM's CUDA platform: no assigned physical ids and - no ``CUDA_VISIBLE_DEVICES`` until a test sets them.""" - # Deferred: this fakes module must import on CPU-only hosts without vLLM. - from vllm.platforms import interface - - class CudaLikePlatform(interface.Platform): - device_control_env_var = "CUDA_VISIBLE_DEVICES" - - monkeypatch.setattr("vllm.platforms.current_platform", CudaLikePlatform) - # vLLM keeps the assigned ids process-wide; start every case unset. - monkeypatch.setattr(interface, "_assigned_physical_gpu_ids", None) - monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False) diff --git a/tests/unit/fakes/vllm_config.py b/tests/unit/fakes/vllm_config.py index 66f2e92b3..0f0f467f7 100644 --- a/tests/unit/fakes/vllm_config.py +++ b/tests/unit/fakes/vllm_config.py @@ -54,29 +54,6 @@ def npu_parallel_config(**overrides): return SimpleNamespace(**values) -def gpu_worker_vllm_config(**overrides): - """The ``parallel_config`` fields that worker device and role-rank - resolution read, for one GPU worker (mp executor, DP1 x TP1 by default). - """ - values = { - "distributed_executor_backend": "mp", - "data_parallel_backend": "mp", - "nnodes_within_dp": 1, - "data_parallel_size": 1, - "data_parallel_rank": 0, - "data_parallel_rank_local": 0, - "data_parallel_index": 0, - "pipeline_parallel_size": 1, - "prefill_context_parallel_size": 1, - "tensor_parallel_size": 1, - "assigned_physical_gpu_ids": None, - } - values.update(overrides) - return SimpleNamespace( - additional_config={}, parallel_config=SimpleNamespace(**values) - ) - - def npu_vllm_config( *, role="attention", diff --git a/tests/unit/v1/worker/test_connector_lifecycle.py b/tests/unit/v1/worker/test_connector_lifecycle.py deleted file mode 100644 index 568a0d357..000000000 --- a/tests/unit/v1/worker/test_connector_lifecycle.py +++ /dev/null @@ -1,193 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project -"""Workers create a routed-experts connector before vLLM's distributed init.""" - -from __future__ import annotations - -import sys -from types import SimpleNamespace - -import pytest - -torch = pytest.importorskip("torch") -pytest.importorskip("vllm") - -from afd_plugin.config import AFDConfig # noqa: E402 -from afd_plugin.connectors import AFDConnectorFactory # noqa: E402 -from afd_plugin.v1.worker.connector_lifecycle import ( # noqa: E402 - check_worker_connector, - create_worker_connector, -) -from tests.unit.fakes.connector import FakeRoutedConnector # noqa: E402 -from tests.unit.fakes.platform import use_cuda_like_platform # noqa: E402 -from tests.unit.fakes.vllm_config import gpu_worker_vllm_config # noqa: E402 - -ROUTED = "FakeRoutedConnector" - - -@pytest.fixture(autouse=True) -def _routed_connector_and_cuda_like_platform(monkeypatch): - monkeypatch.setitem( - AFDConnectorFactory._registry, ROUTED, lambda: FakeRoutedConnector - ) - use_cuda_like_platform(monkeypatch) - - -def test_other_connectors_are_left_to_their_runner(): - """A connector that is not an ``AFDRoutedExpertsConnector`` is not created - before vLLM's distributed init. - - Why: P2P and the NPU connectors keep their runner-owned bring-up. - """ - connector = create_worker_connector( - gpu_worker_vllm_config(), - AFDConfig(connector="P2pNcclAFDConnector", role="ffn"), - rank=0, - local_rank=0, - ) - - assert connector is None - - -def test_routed_connector_is_created_on_vllms_device_and_prepared(): - """The connector gets the worker's ranks, the device vLLM will select and - has joined the AFD world. - - Why: its NCCL environment must precede vLLM's first communicator, and its - device must be the one vLLM gives the worker. - """ - config = gpu_worker_vllm_config(data_parallel_rank_local=1) - - connector = create_worker_connector( - config, AFDConfig(connector=ROUTED, role="ffn"), rank=3, local_rank=0 - ) - - assert isinstance(connector, FakeRoutedConnector) - assert (connector.rank, connector.local_rank, connector.role_rank) == (3, 1, 0) - assert connector.events == [("prepare",)] - - -def test_assigned_device_ids_select_the_connector_device(): - """With ``--device-ids 4,5`` worker 0's connector runs on cuda:4. - - Why: vLLM maps device indices through the assigned ids, so they must be - registered before the connector's device is resolved. - """ - config = gpu_worker_vllm_config(assigned_physical_gpu_ids=[4, 5]) - - connector = create_worker_connector( - config, AFDConfig(connector=ROUTED, role="ffn"), rank=0, local_rank=0 - ) - - assert connector is not None - assert connector.local_rank == 4 - - -def test_tp_part_of_the_role_rank_comes_from_the_worker_rank(monkeypatch): - """With TP > 1 the role rank's TP part comes from the worker rank. - - Why: vLLM's TP group does not exist before its distributed init; reading - it there would fail. TP is vLLM's innermost rank dimension, so the worker - rank modulo the TP size is the TP rank. - """ - monkeypatch.setitem(sys.modules, "vllm.distributed.parallel_state", None) - config = gpu_worker_vllm_config( - data_parallel_size=2, data_parallel_rank=1, tensor_parallel_size=2 - ) - - connector = create_worker_connector( - config, - AFDConfig(connector=ROUTED, role="attention", num_attention_ranks=4), - rank=1, - local_rank=1, - ) - - assert connector is not None - # DP rank 1, TP rank 1 of 2. - assert connector.role_rank == 3 - - -def test_connector_is_closed_when_it_fails_to_join(monkeypatch): - """A failed ``prepare_before_distributed_init`` closes the connector. - - Why: the worker never receives the connector, so nobody else closes it. - """ - created: list[FakeRoutedConnector] = [] - original_init = FakeRoutedConnector.__init__ - - def recording_init(self, *args): - original_init(self, *args) - created.append(self) - - monkeypatch.setattr(FakeRoutedConnector, "__init__", recording_init) - monkeypatch.setattr( - FakeRoutedConnector, "prepare_error", RuntimeError("join failed") - ) - - with pytest.raises(RuntimeError, match="join failed"): - create_worker_connector( - gpu_worker_vllm_config(), - AFDConfig(connector=ROUTED, role="ffn"), - rank=0, - local_rank=0, - ) - - assert created[0].events == [("prepare",), ("close",)] - - -def _connector(*, local_rank=0, role_rank=0): - return SimpleNamespace( - local_rank=local_rank, - role_rank=role_rank, - afd_config=AFDConfig(role="attention", num_attention_ranks=2), - ) - - -def _tp_rank(monkeypatch, tp_rank): - monkeypatch.setattr( - "vllm.distributed.parallel_state.get_tensor_model_parallel_rank", - lambda: tp_rank, - ) - - -def test_check_accepts_vllms_device_and_role_rank(monkeypatch): - _tp_rank(monkeypatch, 1) - - check_worker_connector( - _connector(local_rank=1, role_rank=1), - gpu_worker_vllm_config(tensor_parallel_size=2), - torch.device("cuda", 1), - ) - - -def test_check_rejects_another_device(monkeypatch): - """A connector on another GPU than the worker's fails with both devices. - - Why: a vLLM upgrade that changes device selection must fail at startup - instead of moving traffic onto another rank's GPU. - """ - _tp_rank(monkeypatch, 0) - - with pytest.raises(RuntimeError, match="uses cuda:0 but vLLM placed"): - check_worker_connector( - _connector(local_rank=0), - gpu_worker_vllm_config(tensor_parallel_size=2), - torch.device("cuda", 4), - ) - - -def test_check_rejects_a_different_role_rank(monkeypatch): - """A connector whose early role rank differs from vLLM's fails startup. - - Why: the TP part of the role rank is derived from the worker rank before - vLLM's TP group exists; a vLLM change to its rank order must fail here - instead of routing another rank's tokens. - """ - _tp_rank(monkeypatch, 0) - - with pytest.raises(RuntimeError, match="has role rank 1 but vLLM"): - check_worker_connector( - _connector(role_rank=1), - gpu_worker_vllm_config(tensor_parallel_size=2), - torch.device("cuda", 0), - ) diff --git a/tests/unit/v1/worker/test_ffn_model_runner.py b/tests/unit/v1/worker/test_ffn_model_runner.py index 7114c2a79..eea8a4aa2 100644 --- a/tests/unit/v1/worker/test_ffn_model_runner.py +++ b/tests/unit/v1/worker/test_ffn_model_runner.py @@ -179,16 +179,7 @@ def test_ffn_runner_owns_connector_execution_without_native_runner_inheritance() assert all(base.__name__ != "GPUModelRunner" for base in GPUFFNModelRunner.__mro__) -# A connector the worker created before vLLM's init is the runner's one -# instance: the runner then creates none. -@pytest.mark.parametrize( - ("workers_connector", "expected_events"), - [(False, ["validate", "connector"]), (True, ["validate"])], - ids=["runner-creates", "worker-passes"], -) -def test_ffn_constructor_defers_connector_init_and_owns_model_and_graphs( - monkeypatch, workers_connector, expected_events -): +def test_ffn_constructor_defers_connector_init_and_owns_model_and_graphs(monkeypatch): config = SimpleNamespace( use_v2_model_runner=True, model_config=SimpleNamespace( @@ -229,13 +220,9 @@ def create_connector(*_args): ffn_model_runner_module, "create_afd_gpu_profiler", lambda _role: None ) - runner = GPUFFNModelRunner( - config, - torch.device("cpu"), - connector=connector if workers_connector else None, - ) + runner = GPUFFNModelRunner(config, torch.device("cpu")) - assert events == expected_events + assert events == ["validate", "connector"] assert runner.connector is connector assert runner.model is None assert runner.model_memory_usage == 0 diff --git a/tests/unit/v1/worker/test_ffn_routed_experts_model_runner.py b/tests/unit/v1/worker/test_ffn_routed_experts_model_runner.py index 1c1ce752c..6bb7afe4e 100644 --- a/tests/unit/v1/worker/test_ffn_routed_experts_model_runner.py +++ b/tests/unit/v1/worker/test_ffn_routed_experts_model_runner.py @@ -186,38 +186,16 @@ def test_initialize_passes_moe_layers_and_experts_to_the_connector(): } -def _base_init_with(monkeypatch): - def base_init(self, vllm_config, device, *, connector=None): - self.connector = connector or SimpleNamespace(close=lambda: None) +def test_constructor_rejects_connectors_without_routed_rows(monkeypatch): + def base_init(self, vllm_config, device): + self.connector = SimpleNamespace(close=lambda: None) monkeypatch.setattr(GPUFFNModelRunner, "__init__", base_init) - -def test_constructor_rejects_connectors_without_routed_rows(monkeypatch): - _base_init_with(monkeypatch) - with pytest.raises(TypeError, match="is not an AFDRoutedExpertsConnector"): GPUFFNRoutedExpertsModelRunner(SimpleNamespace(), torch.device("cpu")) -def test_constructor_uses_the_workers_connector(monkeypatch): - """The routed runner runs on the connector the worker created before - vLLM's distributed init. - - Why: that connector already joined the AFD world; a second instance - would never be initialized by its peers. - """ - connector = FakeRoutedConnector() - _base_init_with(monkeypatch) - - runner = GPUFFNRoutedExpertsModelRunner( - SimpleNamespace(), torch.device("cpu"), connector=connector - ) - - assert runner.connector is connector - assert runner.routed_connector is connector - - def _routed_experts(*, fp8: bool, expert_ids=(4, 5), backend: str = "TRITON"): num_experts = len(expert_ids) dtype = torch.float8_e4m3fn if fp8 else torch.bfloat16 diff --git a/tests/unit/v1/worker/test_model_runner_v2.py b/tests/unit/v1/worker/test_model_runner_v2.py index 5ba3fd190..f1b8d42cf 100644 --- a/tests/unit/v1/worker/test_model_runner_v2.py +++ b/tests/unit/v1/worker/test_model_runner_v2.py @@ -2067,69 +2067,6 @@ def test_v2_load_model_reports_moe_layers_to_routed_connector(monkeypatch): } -def _patch_constructor_with_workers_connector(monkeypatch, events): - def native_init(self, vllm_config, device): - self.vllm_config = vllm_config - - def unexpected_create(*args, **kwargs): - pytest.fail("the runner created a second connector") - - monkeypatch.setattr(native_v2.GPUModelRunner, "__init__", native_init) - monkeypatch.setattr( - native_v2.GPUModelRunner, - "shutdown", - lambda self: events.append("native_shutdown"), - ) - monkeypatch.setattr( - "afd_plugin.v1.worker.attention_model_runner_v2.create_afd_gpu_profiler", - lambda role: _StepProfiler(), - ) - config = _patch_v2_constructor_dependencies(monkeypatch, None) - monkeypatch.setattr( - v2_runner_module.AFDConnectorFactory, "create_connector", unexpected_create - ) - return config - - -def test_v2_constructor_uses_the_workers_connector_without_creating_one( - monkeypatch, -): - """A connector the worker created before vLLM's init is the runner's one - instance; the runner creates none. - - Why: one connector instance per worker, with a single create, then - ``init_afd_connector`` lifecycle. - """ - events: list[str] = [] - connector = SimpleNamespace( - control_plane=object(), close=lambda: events.append("close") - ) - config = _patch_constructor_with_workers_connector(monkeypatch, events) - - runner = AFDAttentionModelRunnerV2(config, torch.device("cpu"), connector=connector) - - assert runner.connector is connector - assert events == [] - - -def test_v2_constructor_failure_leaves_the_workers_connector_open(monkeypatch): - """A failed construction closes native state but not the worker's - connector. - - Why: the worker still owns that connector and closes it itself. - """ - events: list[str] = [] - connector = SimpleNamespace( - control_plane=None, close=lambda: events.append("close") - ) - config = _patch_constructor_with_workers_connector(monkeypatch, events) - - with pytest.raises(RuntimeError, match="control-plane"): - AFDAttentionModelRunnerV2(config, torch.device("cpu"), connector=connector) - - assert events == ["native_shutdown"] - - def test_v2_constructor_rejects_missing_control_plane_and_cleans_native( monkeypatch, ): @@ -2338,158 +2275,35 @@ def raising_worker_init(current_worker): assert _NativeRunnerSentinel.instances == 1 -class _ConnectorRunnerRecorder: - connectors: list[object] = [] - - def __init__(self, vllm_config, device, *, connector=None): - type(self).connectors.append(connector) - - -_ROUTED_WORKER_CASES = [ - ( - attention_worker, - attention_worker.AFDAttentionWorker, - attention_worker, - "AFDAttentionModelRunnerV2", - ), - # The FFN worker imports the routed runner lazily, from its own module. - ( - ffn_worker, - AFDFFNWorker, - routed_runner_module, - "GPUFFNRoutedExpertsModelRunner", - ), -] +def test_ffn_worker_selects_routed_runner_for_routed_connector(monkeypatch): + """An ``AFDRoutedExpertsConnector`` gets the routed-experts runner, and + vLLM's MoE backend is pinned to one that keeps the expert layout. - -def _patch_routed_worker(monkeypatch, worker_module, runner_module, runner_name): - """Select a routed-experts connector and record the runner construction.""" - _ConnectorRunnerRecorder.connectors = [] - # The FFN worker picks its runner from the connector class. + Why: the routed runner runs the experts itself, from vLLM's layout. + """ + config = _v2_config(role="ffn") + worker = _prepare_worker(AFDFFNWorker, config, use_v2=True) + _RunnerRecorder.instances.clear() monkeypatch.setattr( ffn_worker.AFDConnectorFactory, "get_connector_class", lambda name: FakeRoutedConnector, ) - monkeypatch.setattr(runner_module, runner_name, _ConnectorRunnerRecorder) - monkeypatch.setattr(native_v2, "GPUModelRunner", _NativeRunnerSentinel) - monkeypatch.setattr( - worker_module, - "assert_compatible_afd_stack", - lambda *args, **kwargs: None, - ) - monkeypatch.setattr(worker_module.torch.accelerator, "empty_cache", lambda: None) - - -@pytest.mark.parametrize( - ("worker_module", "worker_cls", "runner_module", "runner_name"), - _ROUTED_WORKER_CASES, -) -def test_workers_create_connector_before_vllm_and_construct_runner_with_it( - monkeypatch, worker_module, worker_cls, runner_module, runner_name -): - """A connector that joins the AFD world before vLLM's distributed init is - created first and passed to the runner vLLM constructs, which uses that - one instance; the FFN picks the routed-experts runner. - - Why: NCCL reads its environment when vLLM creates the first communicator, - and vLLM builds the runner only after that. - """ - config = _v2_config(role="ffn" if worker_cls is AFDFFNWorker else "attention") - worker = _prepare_worker(worker_cls, config, use_v2=True) - events: list[str] = [] - connector = SimpleNamespace(close=lambda: events.append("close")) - - def create_worker_connector(vllm_config, afd_config, *, rank, local_rank): - events.append("create") - return connector - - def vllm_init(current_worker): - events.append("vllm_init") - _fake_worker_init(current_worker) - - _patch_routed_worker(monkeypatch, worker_module, runner_module, runner_name) monkeypatch.setattr( - worker_module, "create_worker_connector", create_worker_connector + routed_runner_module, "GPUFFNRoutedExpertsModelRunner", _RunnerRecorder ) - monkeypatch.setattr(worker_module.Worker, "init_device", vllm_init) + monkeypatch.setattr(native_v2, "GPUModelRunner", _NativeRunnerSentinel) + monkeypatch.setattr(ffn_worker.Worker, "init_device", _fake_worker_init) monkeypatch.setattr( - worker_module, - "check_worker_connector", - lambda checked, vllm_config, device: events.append("check"), + ffn_worker, "assert_compatible_afd_stack", lambda *args, **kwargs: None ) + monkeypatch.setattr(ffn_worker.torch.accelerator, "empty_cache", lambda: None) worker.init_device() - assert events == ["create", "vllm_init", "check"] - assert _ConnectorRunnerRecorder.connectors == [connector] - assert native_v2.GPUModelRunner is _NativeRunnerSentinel - # The routed FFN pins vLLM's MoE backend to one that keeps the expert layout. - expected_backend = "triton" if worker_cls is AFDFFNWorker else "auto" - assert config.kernel_config.moe_backend == expected_backend - - -@pytest.mark.parametrize( - ("worker_module", "worker_cls", "runner_module", "runner_name"), - _ROUTED_WORKER_CASES, -) -def test_workers_close_their_connector_when_vllm_init_fails( - monkeypatch, worker_module, worker_cls, runner_module, runner_name -): - config = _v2_config(role="ffn" if worker_cls is AFDFFNWorker else "attention") - worker = _prepare_worker(worker_cls, config, use_v2=True) - events: list[str] = [] - connector = SimpleNamespace(close=lambda: events.append("close")) - - def failing_vllm_init(current_worker): - raise RuntimeError("vLLM init failed") - - _patch_routed_worker(monkeypatch, worker_module, runner_module, runner_name) - monkeypatch.setattr( - worker_module, "create_worker_connector", lambda *args, **kwargs: connector - ) - monkeypatch.setattr(worker_module.Worker, "init_device", failing_vllm_init) - - with pytest.raises(RuntimeError, match="vLLM init failed"): - worker.init_device() - - assert events == ["close"] - assert native_v2.GPUModelRunner is _NativeRunnerSentinel - - -@pytest.mark.parametrize( - ("worker_module", "worker_cls", "runner_module", "runner_name"), - _ROUTED_WORKER_CASES, -) -def test_workers_reject_connector_on_another_device( - monkeypatch, worker_module, worker_cls, runner_module, runner_name -): - """A connector created on a different GPU than vLLM's fails startup. - - Why: the connector picks its device before vLLM does, by mirroring vLLM's - selection; a vLLM upgrade that changes the selection must fail here - rather than run the connector on another rank's GPU. - """ - config = _v2_config(role="ffn" if worker_cls is AFDFFNWorker else "attention") - worker = _prepare_worker(worker_cls, config, use_v2=True) - events: list[str] = [] - connector = SimpleNamespace(close=lambda: events.append("close"), local_rank=0) - - def vllm_init(current_worker): - _fake_worker_init(current_worker) - current_worker.device = torch.device("cuda", 4) - - _patch_routed_worker(monkeypatch, worker_module, runner_module, runner_name) - monkeypatch.setattr( - worker_module, "create_worker_connector", lambda *args, **kwargs: connector - ) - monkeypatch.setattr(worker_module.Worker, "init_device", vllm_init) - - with pytest.raises(RuntimeError, match="uses cuda:0 but vLLM placed"): - worker.init_device() - - assert events == ["close"] + assert worker.model_runner.__class__ is _RunnerRecorder assert native_v2.GPUModelRunner is _NativeRunnerSentinel + assert config.kernel_config.moe_backend == "triton" def _deepep_m2n_v2_config(**kwargs): From 2afd130ec044d15927d6709ef2ae9c3766fd692c Mon Sep 17 00:00:00 2001 From: ronenkat <16743404+ronenkat@users.noreply.github.com> Date: Mon, 26 Oct 2026 17:28:09 +0200 Subject: [PATCH 12/32] test(e2e): multi-node Qwen3-30B-A3B cases on the DeepEP M2N connector Run the DeepEP M2N eager, graph and graph+DBO scenarios across two pods, Attention DP2 on one and FFN DP2 on the other (layout 2A0F,0A2F), with GSM8K at 100 samples by default. M2N needs FFN ranks first (afdEP-comm) and each NVLink domain to hold an equal, contiguous block of world ranks (DeepEP), so a two-pod run splits by role with equal rank counts; dispatch and combine then cross pods over NCCL GIN. The multi-pod runner now knows the connector's rendezvous role: world rank 0, which hosts its gloo store, is the first FFN rank. Co-Authored-By: Claude Opus 5.5 Signed-off-by: ronenkat <16743404+ronenkat@users.noreply.github.com> --- docs/design/module/e2e_testing.md | 1 + .../qwen3_moe/test_qwen3_moe_multi_pod.py | 118 ++++++++++++++++++ tests/e2e/multi_pod/layout.py | 4 +- 3 files changed, 122 insertions(+), 1 deletion(-) create mode 100644 tests/e2e/models/qwen3_moe/test_qwen3_moe_multi_pod.py diff --git a/docs/design/module/e2e_testing.md b/docs/design/module/e2e_testing.md index d68ededbb..c2a15bd6e 100644 --- a/docs/design/module/e2e_testing.md +++ b/docs/design/module/e2e_testing.md @@ -37,6 +37,7 @@ validation_paths: - "tests/e2e/models/deepseek_v4_flash/config.py" - "tests/e2e/models/deepseek_v4_flash/completions.py" - "tests/e2e/models/qwen3_moe/test_qwen3_moe.py" + - "tests/e2e/models/qwen3_moe/test_qwen3_moe_multi_pod.py" - "tests/e2e/models/qwen3_6/test_qwen3_6.py" upstream_refs: - "vLLM 0.30.0 serving and shutdown interfaces" diff --git a/tests/e2e/models/qwen3_moe/test_qwen3_moe_multi_pod.py b/tests/e2e/models/qwen3_moe/test_qwen3_moe_multi_pod.py new file mode 100644 index 000000000..8b7511900 --- /dev/null +++ b/tests/e2e/models/qwen3_moe/test_qwen3_moe_multi_pod.py @@ -0,0 +1,118 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Qwen3 MoE multi-pod AFD E2E cases on the DeepEP M2N connector. + +This test runs as *one pod's slice* of an already-provisioned multi-pod +deployment: the k8s pods or Docker containers must already exist -- brought +up by hand, following `.agents/skills/run-e2e/resources/k8-multi-pod.md` -- +before this test is invoked once inside each of them. Every test decision -- +launch order, readiness, evaluation, teardown -- is made by the in-pod +runner itself; this only supplies what the runner cannot infer from its own +pod's environment. + +DeepEP M2N needs FFN ranks first (afdEP-comm) and every NVLink domain to hold +an equal, contiguous block of world ranks (DeepEP), so the two pods split by +role with equal rank counts: Attention on one pod, FFN on the other. The +ranks then span two NVLink domains, so dispatch and combine cross pods over +NCCL GIN (RDMA). +""" + +from __future__ import annotations + +import os +import sys +from collections.abc import Iterator + +import pytest + +from tests.conftest import ( + download_dataset, + download_model, + preserve_environment_variable, + run_runner, +) +from tests.e2e.runner import GSM8K_LIMIT_ENV + +GSM8K_DATASET_ID = "openai/gsm8k" +GSM8K_DATASET_CONFIG = "main" +QWEN3_MOE_REPO_ID = "Qwen/Qwen3-30B-A3B" +QWEN3_MOE_MAX_MODEL_LEN = 4096 +# The default sample limit of these cases; AFD_GSM8K_LIMIT overrides it. +GSM8K_SAMPLE_LIMIT = 100 +# The layout is a second axis, orthogonal to the scenario id, so accuracy +# evidence stays comparable with the single-host rows. +POD_LAYOUTS = { + "2pod-role-split": "2A0F,0A2F", +} +MULTI_POD_CASES = [ + ("afd-deepep-m2n-eager-2a2f", "2pod-role-split"), + ("afd-deepep-m2n-graph-2a2f", "2pod-role-split"), + ("afd-deepep-m2n-graph-dbo-2a2f", "2pod-role-split"), +] + + +def _required_env(name: str) -> str: + value = os.environ.get(name) + if not value: + raise RuntimeError(f"{name} must be set") + return value + + +@pytest.fixture(scope="module", autouse=True) +def _prepare_e2e_assets() -> Iterator[None]: + """Make GSM8K and Qwen3-30B-A3B available in this pod for the runner.""" + with preserve_environment_variable("AFD_GPU_E2E_MODEL"): + download_dataset(GSM8K_DATASET_ID, GSM8K_DATASET_CONFIG) + if not os.environ.get("AFD_GPU_E2E_MODEL"): + os.environ["AFD_GPU_E2E_MODEL"] = str(download_model(QWEN3_MOE_REPO_ID)) + yield + + +def build_runner_command(scenario: str, layout_name: str) -> list[str]: + """Build this pod's in-pod runner argv for one case. + + Assumes this process is itself running inside one pod of an + already-provisioned deployment. Pod identity (AFD_E2E_POD_INDEX / + JOB_COMPLETION_INDEX / HOSTNAME) and peer addresses + (AFD_E2E_POD_ADDRESSES, or the rendezvous store's own address-exchange + barrier) are resolved by the runner itself from its environment, so this + only supplies what it cannot infer: the scenario, the shared rendezvous + store host, and where to write results. + """ + command = [ + sys.executable, + "-m", + "tests.e2e.multi_pod.runner", + "--scenario", + scenario, + "--pod-layout", + POD_LAYOUTS[layout_name], + "--run-id", + _required_env("AFD_E2E_RUN_ID"), + "--model", + _required_env("AFD_GPU_E2E_MODEL"), + "--gsm8k-output-path", + _required_env("AFD_E2E_GSM8K_OUTPUT"), + "--store-host", + _required_env("AFD_E2E_STORE_HOST"), + f"--common-vllm-arg=--max-model-len={QWEN3_MOE_MAX_MODEL_LEN}", + ] + store_port = os.environ.get("AFD_E2E_STORE_PORT") + if store_port: + command.extend(["--store-port", store_port]) + for value in os.environ.get("AFD_E2E_POD_ENV", "").split(";"): + if value.strip(): + command.extend(["--pod-env", value.strip()]) + return command + + +@pytest.mark.e2e +@pytest.mark.parametrize( + ("scenario", "layout_name"), + MULTI_POD_CASES, + ids=[f"{scenario}-{layout}" for scenario, layout in MULTI_POD_CASES], +) +def test_multi_pod_deepep_m2n(scenario: str, layout_name: str) -> None: + env = os.environ.copy() + env.setdefault(GSM8K_LIMIT_ENV, str(GSM8K_SAMPLE_LIMIT)) + run_runner(build_runner_command(scenario, layout_name), env=env) diff --git a/tests/e2e/multi_pod/layout.py b/tests/e2e/multi_pod/layout.py index 40c91ed18..0282ec923 100644 --- a/tests/e2e/multi_pod/layout.py +++ b/tests/e2e/multi_pod/layout.py @@ -22,9 +22,11 @@ ASYNC_AFD_CONNECTOR = "CAMAsyncAFDConnector" # docs/gpu/NCCL_P2P_CONNECTOR_USER_GUIDE.md:88 - the synchronous connectors # rendezvous at the first FFN rank. The async CAM connector rendezvous at the -# Attention side instead (tools/itask/launch_dsv4_afd_cross_node.sh). +# Attention side instead (tools/itask/launch_dsv4_afd_cross_node.sh). DeepEP M2N +# puts FFN ranks first, so its world rank 0 (the gloo store) is the first FFN rank. RENDEZVOUS_ROLE_BY_CONNECTOR = { "P2pNcclAFDConnector": FFN_ROLE, + "DeepEPM2NAFDConnector": FFN_ROLE, "CAMP2pAFDConnector": FFN_ROLE, ASYNC_AFD_CONNECTOR: ATTENTION_ROLE, } From 7eff267e35e2042abe011f41b38865eccf7b32ab Mon Sep 17 00:00:00 2001 From: ronenkat <16743404+ronenkat@users.noreply.github.com> Date: Tue, 27 Oct 2026 12:54:24 +0200 Subject: [PATCH 13/32] fix(gpu): run the DeepEP M2N connector on the worker's visible CUDA device The runners pass vLLM's world-group local rank, which is a logical index. vLLM places the worker on that index mapped through the assigned physical GPU ids (--device-ids) and CUDA_VISIBLE_DEVICES, so with --device-ids 4,5 the connector chose cuda:0/cuda:1 while the model ran on cuda:4/cuda:5, and init_afd_connector created its transport buffers on another worker's GPU. The pre-init device mirroring that handled this left with the bootstrap lifecycle. Map the rank with current_platform.logical_device_id_to_visible_device_id, as vLLM's Worker.init_device does; vLLM registers the assigned ids before it builds the runner. P2P and the runners are unchanged. Cover the identity, CUDA_VISIBLE_DEVICES, assigned-ids and reordered cases with a CUDA-like platform fake. Co-Authored-By: Claude Opus 5.5 Signed-off-by: ronenkat <16743404+ronenkat@users.noreply.github.com> --- .../connectors/gpu/deepep_m2n/connector.py | 9 +++- .../connectors/test_deepep_m2n_connector.py | 47 ++++++++++++++++++- tests/unit/fakes/platform.py | 22 +++++++++ 3 files changed, 75 insertions(+), 3 deletions(-) create mode 100644 tests/unit/fakes/platform.py diff --git a/afd_plugin/connectors/gpu/deepep_m2n/connector.py b/afd_plugin/connectors/gpu/deepep_m2n/connector.py index e688c8dd9..1fcf70f75 100644 --- a/afd_plugin/connectors/gpu/deepep_m2n/connector.py +++ b/afd_plugin/connectors/gpu/deepep_m2n/connector.py @@ -70,6 +70,7 @@ from vllm.model_executor.layers.quantization.utils.fp8_utils import ( per_token_group_quant_fp8, ) +from vllm.platforms import current_platform from vllm.utils.torch_utils import direct_register_custom_op from vllm.v1.worker.ubatching import dbo_current_ubatch_id, dbo_enabled @@ -957,7 +958,13 @@ def __init__( cpu_sync_min_tokens=self.config.cpu_sync_min_tokens, ) self.num_lanes = resolve_num_lanes(self.config, vllm_config) - self.device = torch.device("cuda", local_rank) + # ``local_rank`` is vLLM's logical local rank (the world group's). The + # worker runs on its visible ordinal, which differs under + # ``--device-ids`` or a reordered CUDA_VISIBLE_DEVICES; vLLM has + # registered the assigned ids by the time the runner builds us. + self.device = torch.device( + "cuda", current_platform.logical_device_id_to_visible_device_id(local_rank) + ) # Where control payloads are staged: gloo moves CPU tensors. self.control_device = ( torch.device("cpu") diff --git a/tests/unit/connectors/test_deepep_m2n_connector.py b/tests/unit/connectors/test_deepep_m2n_connector.py index 6d9ae4d23..87c21ee03 100644 --- a/tests/unit/connectors/test_deepep_m2n_connector.py +++ b/tests/unit/connectors/test_deepep_m2n_connector.py @@ -13,6 +13,8 @@ torch = pytest.importorskip("torch") pytest.importorskip("vllm") +from vllm.platforms.interface import set_assigned_physical_gpu_ids # noqa: E402 + from tests.unit.fakes.afdep import FakeElasticBuffer, install_fake_afdep # noqa: E402 # The connector module imports afdEP's transport when it loads. @@ -60,6 +62,7 @@ FFNLocalExperts, LocalExpertWeights, ) +from tests.unit.fakes.platform import cuda_like_platform # noqa: E402 from tests.unit.fakes.vllm_config import ( # noqa: E402 M2N_HIDDEN_SIZE, M2NTopology, @@ -75,11 +78,15 @@ def _dp(*tokens: int) -> AFDDPMetadata: def _connector( - role: str, role_rank: int, topology: M2NTopology, extra: dict | None = None + role: str, + role_rank: int, + topology: M2NTopology, + extra: dict | None = None, + local_rank: int = 0, ): vllm_config = m2n_vllm_config(topology, role, role_rank, extra or {}, port=12345) afd_config = afd_config_from_mapping(vllm_config.additional_config["afd"]) - return AFDConnectorFactory.create_connector(0, 0, vllm_config, afd_config) + return AFDConnectorFactory.create_connector(0, local_rank, vllm_config, afd_config) def _ffn_deployment(layout: DeepEPM2NLayout, weight_format: str = "bf16") -> dict: @@ -811,6 +818,42 @@ def test_connector_constructs_without_initialization(): assert explicit_ffn.capacity_tokens_per_rank == 64 +@pytest.mark.parametrize( + ("visible_devices", "assigned_ids", "expected_indices"), + [ + (None, None, (0, 1)), + ("2,3", None, (0, 1)), + (None, [4, 5], (4, 5)), + ("4,5", [5, 4], (1, 0)), + ], +) +def test_connector_uses_the_device_vllm_gives_the_worker( + monkeypatch, visible_devices, assigned_ids, expected_indices +): + """The runner passes vLLM's logical local rank; the connector runs on the + visible CUDA ordinal vLLM's ``Worker.init_device`` maps it to. + + Why: with ``--device-ids 4,5`` (registered by vLLM before the runner is + built) or a reordered CUDA_VISIBLE_DEVICES, the logical rank names another + worker's GPU, so the buffers would sit apart from the model's tensors. + """ + monkeypatch.setattr( + connector_module, "current_platform", cuda_like_platform(monkeypatch) + ) + if visible_devices is not None: + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", visible_devices) + if assigned_ids is not None: + set_assigned_physical_gpu_ids(assigned_ids) + topology = M2NTopology(ffn_size=1, attn_size=2) + + devices = tuple( + _connector("attention", rank, topology, local_rank=rank).device + for rank in range(2) + ) + + assert devices == tuple(torch.device("cuda", i) for i in expected_indices) + + def test_control_plane_requires_the_handshake_on_ffn(): """An FFN connector cannot plan a step before the handshake. diff --git a/tests/unit/fakes/platform.py b/tests/unit/fakes/platform.py new file mode 100644 index 000000000..3562eb69a --- /dev/null +++ b/tests/unit/fakes/platform.py @@ -0,0 +1,22 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""A CUDA-like vLLM platform for device-selection tests on hosts without GPUs.""" + +from __future__ import annotations + +import pytest + + +def cuda_like_platform(monkeypatch: pytest.MonkeyPatch) -> type: + """A platform that maps device ids like vLLM's CUDA platform, starting with + no assigned physical ids and no ``CUDA_VISIBLE_DEVICES``.""" + # Deferred: this fakes module must import on CPU-only hosts without vLLM. + from vllm.platforms import interface + + class CudaLikePlatform(interface.Platform): + device_control_env_var = "CUDA_VISIBLE_DEVICES" + + # vLLM keeps the assigned ids process-wide; start every case unset. + monkeypatch.setattr(interface, "_assigned_physical_gpu_ids", None) + monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False) + return CudaLikePlatform From a4a667d3c146bfc8e9aa23079020ab7b68bbbcb5 Mon Sep 17 00:00:00 2001 From: ronenkat <16743404+ronenkat@users.noreply.github.com> Date: Tue, 27 Oct 2026 13:20:25 +0200 Subject: [PATCH 14/32] refactor(connectors): keep the routed-experts contract with the base contract AFDRoutedExpertsConnector is an abstract interface, not an expert-compute implementation, and routed_experts.py was easy to mistake for one next to vLLM's RoutedExperts. Move it into connectors/base.py next to AFDConnectorBase, still as its own subclass so P2P and CAM connectors implement none of its methods. Move LocalExpertWeights, FFNLocalExperts and FP8_DTYPE into connectors/metadata.py with the other shared payload types; base.py already imports metadata.py, so imports stay acyclic. routed_experts.py has no remaining users and is removed. The package-level names in afd_plugin.connectors are unchanged. Co-Authored-By: Claude Opus 5.5 Signed-off-by: ronenkat <16743404+ronenkat@users.noreply.github.com> --- afd_plugin/connectors/__init__.py | 4 +- afd_plugin/connectors/base.py | 66 ++++++++- .../connectors/gpu/deepep_m2n/connector.py | 15 ++- .../connectors/gpu/deepep_m2n/experts.py | 2 +- .../connectors/gpu/deepep_m2n/fp8_quant.py | 2 +- afd_plugin/connectors/metadata.py | 53 ++++++++ afd_plugin/connectors/routed_experts.py | 126 ------------------ .../worker/ffn_routed_experts_model_runner.py | 2 +- tests/gpu/deepep_m2n_harness.py | 2 +- .../connectors/test_deepep_m2n_connector.py | 2 +- 10 files changed, 131 insertions(+), 143 deletions(-) delete mode 100644 afd_plugin/connectors/routed_experts.py diff --git a/afd_plugin/connectors/__init__.py b/afd_plugin/connectors/__init__.py index 55fe53646..2b270cc6e 100644 --- a/afd_plugin/connectors/__init__.py +++ b/afd_plugin/connectors/__init__.py @@ -5,6 +5,7 @@ from afd_plugin.connectors.base import ( AFDConnectorBase, AFDControlPlane, + AFDRoutedExpertsConnector, ConnectorExtraInfo, ) from afd_plugin.connectors.factory import AFDConnectorFactory @@ -19,9 +20,6 @@ AFDTransferContext, AFDTransferMetadata, AFDTransferState, -) -from afd_plugin.connectors.routed_experts import ( - AFDRoutedExpertsConnector, FFNLocalExperts, LocalExpertWeights, ) diff --git a/afd_plugin/connectors/base.py b/afd_plugin/connectors/base.py index 95201b838..d0bb01770 100644 --- a/afd_plugin/connectors/base.py +++ b/afd_plugin/connectors/base.py @@ -5,7 +5,7 @@ from __future__ import annotations from abc import ABC, abstractmethod -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from dataclasses import dataclass from typing import TYPE_CHECKING, Any @@ -16,6 +16,8 @@ AFDA2FTransferPayload, AFDControlPayload, AFDTransferContext, + FFNLocalExperts, + LocalExpertWeights, ) if TYPE_CHECKING: @@ -255,6 +257,61 @@ def send_ffn_output( raise NotImplementedError +class AFDRoutedExpertsConnector(AFDConnectorBase): + """A connector whose FFN side receives routed (token, local expert) rows. + + Routing (gate and top-k) runs on Attention. The FFN receives one row per + (token, local expert) pair, runs only its local experts and sends the + weighted rows back. Workers and runners use this interface and the + ``FFNLocalExperts``/``LocalExpertWeights`` types from ``metadata.py`` + instead of any specific connector, so they stay connector-agnostic. It + extends ``AFDConnectorBase`` rather than adding to it, so connectors + without routed experts implement none of this. + + Lifecycle: as every connector, the runner creates it and calls + ``init_afd_connector`` after weights load, here with what it loaded. The + FFN worker selects the routed-experts runner for these connectors. + """ + + # Extends AFDConnectorBase.init_afd_connector with optional keywords only, + # so calls through the base contract stay valid. + @abstractmethod + def init_afd_connector( + self, + *, + moe_layer_indices: Sequence[int] | None = None, + local_experts: Mapping[int, FFNLocalExperts] | None = None, + ) -> None: + """Collective: check what every rank loaded, then create resources. + + Args: + moe_layer_indices: MoE layers this rank's model runs, in order. + ``None`` skips the post-load checks for this rank. + local_experts: FFN only: this rank's experts per MoE layer, as + vLLM loaded them. + """ + raise NotImplementedError + + @abstractmethod + def run_experts( + self, + payload: AFDA2FTransferPayload, + weights: LocalExpertWeights, + ) -> torch.Tensor: + """Run this rank's experts on the rows ``recv_attn_output`` returned; + the result is what ``send_ffn_output`` sends back.""" + raise NotImplementedError + + @abstractmethod + def synchronize_lanes(self) -> None: + """Join the streams the connector forked into the caller's stream. + + FFN runners call this at the end of each step, before a CUDA-graph + capture ends. + """ + raise NotImplementedError + + class AFDControlPlane(ABC): """DP metadata control plane""" @@ -307,4 +364,9 @@ def recv_dp_metadata_list(self) -> AFDControlPayload: raise NotImplementedError -__all__ = ["AFDConnectorBase", "AFDControlPlane", "ConnectorExtraInfo"] +__all__ = [ + "AFDConnectorBase", + "AFDControlPlane", + "AFDRoutedExpertsConnector", + "ConnectorExtraInfo", +] diff --git a/afd_plugin/connectors/gpu/deepep_m2n/connector.py b/afd_plugin/connectors/gpu/deepep_m2n/connector.py index 1fcf70f75..8a7e86280 100644 --- a/afd_plugin/connectors/gpu/deepep_m2n/connector.py +++ b/afd_plugin/connectors/gpu/deepep_m2n/connector.py @@ -75,7 +75,11 @@ from vllm.v1.worker.ubatching import dbo_current_ubatch_id, dbo_enabled from afd_plugin.config import AFDConfig -from afd_plugin.connectors.base import AFDControlPlane, ConnectorExtraInfo +from afd_plugin.connectors.base import ( + AFDControlPlane, + AFDRoutedExpertsConnector, + ConnectorExtraInfo, +) from afd_plugin.connectors.gpu.deepep_m2n.experts import ( EXPERT_BACKENDS, DeepEPM2NExperts, @@ -91,20 +95,17 @@ plan_network, ) from afd_plugin.connectors.metadata import ( + FP8_DTYPE, AFDA2FTransferPayload, AFDControlPayload, AFDDPMetadata, AFDTransferContext, AFDTransferMetadata, AFDTransferState, - recv_control_payload, - send_control_payload, -) -from afd_plugin.connectors.routed_experts import ( - FP8_DTYPE, - AFDRoutedExpertsConnector, FFNLocalExperts, LocalExpertWeights, + recv_control_payload, + send_control_payload, ) from afd_plugin.distributed import init_afd_process_group diff --git a/afd_plugin/connectors/gpu/deepep_m2n/experts.py b/afd_plugin/connectors/gpu/deepep_m2n/experts.py index 102e89120..2b5fe0042 100644 --- a/afd_plugin/connectors/gpu/deepep_m2n/experts.py +++ b/afd_plugin/connectors/gpu/deepep_m2n/experts.py @@ -42,7 +42,7 @@ dequantize_fp8_rows, silu_mul_fp8_quant, ) -from afd_plugin.connectors.routed_experts import FP8_DTYPE, LocalExpertWeights +from afd_plugin.connectors.metadata import FP8_DTYPE, LocalExpertWeights ExpertBackend = Literal["deepgemm", "fused_experts"] EXPERT_BACKENDS: tuple[ExpertBackend, ...] = ("deepgemm", "fused_experts") diff --git a/afd_plugin/connectors/gpu/deepep_m2n/fp8_quant.py b/afd_plugin/connectors/gpu/deepep_m2n/fp8_quant.py index 0badf760a..0f4125870 100644 --- a/afd_plugin/connectors/gpu/deepep_m2n/fp8_quant.py +++ b/afd_plugin/connectors/gpu/deepep_m2n/fp8_quant.py @@ -21,7 +21,7 @@ import torch from vllm.triton_utils import tl, triton -from afd_plugin.connectors.routed_experts import FP8_DTYPE +from afd_plugin.connectors.metadata import FP8_DTYPE FP8_MAX = torch.finfo(FP8_DTYPE).max FP8_MIN = torch.finfo(FP8_DTYPE).min diff --git a/afd_plugin/connectors/metadata.py b/afd_plugin/connectors/metadata.py index 728d94950..30c63db0f 100644 --- a/afd_plugin/connectors/metadata.py +++ b/afd_plugin/connectors/metadata.py @@ -19,6 +19,8 @@ from afd_plugin.connectors.base import AFDConnectorBase +FP8_DTYPE = torch.float8_e4m3fn + @dataclass(slots=True) class AFDDPMetadata: @@ -315,6 +317,54 @@ def clone(self) -> AFDForwardContextMetadata: return cloned +@dataclass(frozen=True) +class LocalExpertWeights: + """This FFN rank's experts, in vLLM ``FusedMoE`` layout. + + Attributes: + w13: ``[E_local, 2 * I, H]`` gate and up projections (gate first), + BF16 or FP8. + w2: ``[E_local, H, I]`` down projection, same dtype as ``w13``. + w13_scale: FP8 only: ``[E_local, 2 * I / 128, H / 128]`` fp32 block + scales (vLLM ``w13_weight_scale_inv``). + w2_scale: FP8 only: ``[E_local, H / 128, I / 128]`` fp32 block scales. + """ + + w13: torch.Tensor + w2: torch.Tensor + w13_scale: torch.Tensor | None = None + w2_scale: torch.Tensor | None = None + + def __post_init__(self) -> None: + if self.w13.dtype != self.w2.dtype: + raise ValueError("w13 and w2 must have the same dtype") + if self.is_fp8 != (self.w13_scale is not None and self.w2_scale is not None): + raise ValueError( + "FP8 expert weights need both block scales (and BF16 none)" + ) + + @property + def num_experts(self) -> int: + return int(self.w13.shape[0]) + + @property + def is_fp8(self) -> bool: + return self.w13.dtype == FP8_DTYPE + + +@dataclass(frozen=True) +class FFNLocalExperts: + """One MoE layer's experts on this FFN rank, as vLLM loaded them. + + Attributes: + expert_ids: Global ids of the experts in ``weights``, in slot order. + weights: The same experts in the executor's layout. + """ + + expert_ids: tuple[int, ...] + weights: LocalExpertWeights + + def _to_int(value: object) -> int: item = getattr(value, "item", None) return int(cast(Any, item() if callable(item) else value)) @@ -434,6 +484,9 @@ def recv_control_payload( "AFDForwardContextMetadata", "AFDA2FTransferPayload", "AFDSingleDPMetadata", + "FFNLocalExperts", + "FP8_DTYPE", + "LocalExpertWeights", "decode_control_payload", "encode_control_payload", "recv_control_payload", diff --git a/afd_plugin/connectors/routed_experts.py b/afd_plugin/connectors/routed_experts.py deleted file mode 100644 index 67d6f6264..000000000 --- a/afd_plugin/connectors/routed_experts.py +++ /dev/null @@ -1,126 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project -"""Contract for connectors whose FFN side receives routed expert rows. - -Routing (gate and top-k) runs on Attention. The FFN receives one row per -(token, local expert) pair, runs only its local experts and sends the weighted -rows back. Workers and runners use ``AFDRoutedExpertsConnector`` and the types -here instead of any specific connector, and this module imports no connector -backend, so they stay connector-agnostic. ``AFDConnectorBase`` is unchanged. -""" - -from __future__ import annotations - -from abc import abstractmethod -from collections.abc import Mapping, Sequence -from dataclasses import dataclass - -import torch - -from afd_plugin.connectors.base import AFDConnectorBase -from afd_plugin.connectors.metadata import AFDA2FTransferPayload - -FP8_DTYPE = torch.float8_e4m3fn - - -@dataclass(frozen=True) -class LocalExpertWeights: - """This FFN rank's experts, in vLLM ``FusedMoE`` layout. - - Attributes: - w13: ``[E_local, 2 * I, H]`` gate and up projections (gate first), - BF16 or FP8. - w2: ``[E_local, H, I]`` down projection, same dtype as ``w13``. - w13_scale: FP8 only: ``[E_local, 2 * I / 128, H / 128]`` fp32 block - scales (vLLM ``w13_weight_scale_inv``). - w2_scale: FP8 only: ``[E_local, H / 128, I / 128]`` fp32 block scales. - """ - - w13: torch.Tensor - w2: torch.Tensor - w13_scale: torch.Tensor | None = None - w2_scale: torch.Tensor | None = None - - def __post_init__(self) -> None: - if self.w13.dtype != self.w2.dtype: - raise ValueError("w13 and w2 must have the same dtype") - if self.is_fp8 != (self.w13_scale is not None and self.w2_scale is not None): - raise ValueError( - "FP8 expert weights need both block scales (and BF16 none)" - ) - - @property - def num_experts(self) -> int: - return int(self.w13.shape[0]) - - @property - def is_fp8(self) -> bool: - return self.w13.dtype == FP8_DTYPE - - -@dataclass(frozen=True) -class FFNLocalExperts: - """One MoE layer's experts on this FFN rank, as vLLM loaded them. - - Attributes: - expert_ids: Global ids of the experts in ``weights``, in slot order. - weights: The same experts in the executor's layout. - """ - - expert_ids: tuple[int, ...] - weights: LocalExpertWeights - - -class AFDRoutedExpertsConnector(AFDConnectorBase): - """A connector whose FFN side receives routed (token, local expert) rows. - - Lifecycle: as every connector, the runner creates it and calls - ``init_afd_connector`` after weights load, here with what it loaded. The - FFN worker selects the routed-experts runner for these connectors. - """ - - # Extends AFDConnectorBase.init_afd_connector with optional keywords only, - # so calls through the base contract stay valid. - @abstractmethod - def init_afd_connector( - self, - *, - moe_layer_indices: Sequence[int] | None = None, - local_experts: Mapping[int, FFNLocalExperts] | None = None, - ) -> None: - """Collective: check what every rank loaded, then create resources. - - Args: - moe_layer_indices: MoE layers this rank's model runs, in order. - ``None`` skips the post-load checks for this rank. - local_experts: FFN only: this rank's experts per MoE layer, as - vLLM loaded them. - """ - raise NotImplementedError - - @abstractmethod - def run_experts( - self, - payload: AFDA2FTransferPayload, - weights: LocalExpertWeights, - ) -> torch.Tensor: - """Run this rank's experts on the rows ``recv_attn_output`` returned; - the result is what ``send_ffn_output`` sends back.""" - raise NotImplementedError - - @abstractmethod - def synchronize_lanes(self) -> None: - """Join the streams the connector forked into the caller's stream. - - FFN runners call this at the end of each step, before a CUDA-graph - capture ends. - """ - raise NotImplementedError - - -__all__ = [ - "AFDRoutedExpertsConnector", - "FFNLocalExperts", - "FP8_DTYPE", - "LocalExpertWeights", -] diff --git a/afd_plugin/v1/worker/ffn_routed_experts_model_runner.py b/afd_plugin/v1/worker/ffn_routed_experts_model_runner.py index 85a848c8d..55c26f41d 100644 --- a/afd_plugin/v1/worker/ffn_routed_experts_model_runner.py +++ b/afd_plugin/v1/worker/ffn_routed_experts_model_runner.py @@ -18,7 +18,7 @@ FFNLocalExperts, LocalExpertWeights, ) -from afd_plugin.connectors.routed_experts import FP8_DTYPE +from afd_plugin.connectors.metadata import FP8_DTYPE from afd_plugin.v1.worker.ffn_model_runner import GPUFFNModelRunner if TYPE_CHECKING: diff --git a/tests/gpu/deepep_m2n_harness.py b/tests/gpu/deepep_m2n_harness.py index b193f5428..2a228ce84 100644 --- a/tests/gpu/deepep_m2n_harness.py +++ b/tests/gpu/deepep_m2n_harness.py @@ -578,7 +578,7 @@ def local_weights(experts: range, payload_format: str) -> tuple[Any, Any, Any]: dequantize_fp8_block128, quantize_fp8_block128, ) - from afd_plugin.connectors.routed_experts import LocalExpertWeights + from afd_plugin.connectors.metadata import LocalExpertWeights w13, w2 = make_expert_weights(experts) if payload_format == "bf16": diff --git a/tests/unit/connectors/test_deepep_m2n_connector.py b/tests/unit/connectors/test_deepep_m2n_connector.py index 87c21ee03..42f0baf9c 100644 --- a/tests/unit/connectors/test_deepep_m2n_connector.py +++ b/tests/unit/connectors/test_deepep_m2n_connector.py @@ -58,7 +58,7 @@ NetworkConfig, RankTopology, ) -from afd_plugin.connectors.routed_experts import ( # noqa: E402 +from afd_plugin.connectors.metadata import ( # noqa: E402 FFNLocalExperts, LocalExpertWeights, ) From 9bd2bfc9cab1df488326987aa7235dc4af7c7e90 Mon Sep 17 00:00:00 2001 From: Shlomit Koyfman Date: Tue, 27 Oct 2026 13:15:00 +0200 Subject: [PATCH 15/32] Provide a skill for benchmarking AFD recipes (#465) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Purpose Add run-vllm-bench-k8s skill for benchmarking AFD recipes on Kubernetes. Code the k8s benchmarking workflow as a skill so that measuring throughput/TTFT/TPOT/ITL for a deployed AFD GPU recipe is repeatable and self-documenting. ## Issue - Related issue(s): #464 - Closing keyword, only if fully resolved: Closes #464 ## Scope - In scope: GPU colocation recipes only - Out of scope: - baseline recipes - local benchmarking (run request_generator.sh against decode_bench_server.sh directly) - NPU recipes - correctness E2E (run-e2e). ## Implementation Notes .agents/skills/run-vllm-bench-k8s/SKILL.md — the main workflow: - Workload selection (e.g., balanced, max-throughput, prefix-heavy, long-context). - Recipe deployment: delegates to deploy-afd-k8s. - Sweep execution: single Pod spec running the repo's tools/benchmarks/request_generator.sh from the afd-plugin image (generated using the existing docker file). - Result collection (copies from the pod to local storage). - Teardown: cleanup of vllm-bench, vllm-pod, and vllm-service without waiting to be asked; the reports PVC is kept until all runs of interest are copied out. References: - references/comparing-runs.md — comparability checks and per-point deltas for two-run comparisons; recipes to compare must be given explicitly by the user. - references/fetching-past-reports.md — recovering reports from the vllm-bench-reports PVC after the pod is gone. - agents/openai.yaml — agent-facing registration. ## Test Plan - User prompt: "Follow the `run-vllm-bench-k8s` skill to compare two recipes". State, for example, 2a2f_eager_dbo_dp1tp2.sh and 2a2f_graph_dbo_dp1tp2.sh from recipe/gpu/P2pNcclAFDConnector/deepseek_v2_lite/prefill_decode_colocation. - Provide afd-plugin image and model PVC if exists. - Select the workload - e.g., balanced. ## Test Result Expected: a table with runs' comparison and a summary with findings. Example: ``` Reference: 2a2f_eager_dbo_dp1tp2-20261018-111201 (--enforce-eager) Compared: 2a2f_graph_dbo_dp1tp2-20261018-114152 (FULL_DECODE_ONLY CUDA Graph, capture size 64) Model /models/DeepSeek-V2-Lite, 4 GPUs, workload balanced (identical across runs: random ISL 1024 / OSL 256, 512 prompts/point, sweep 2:64 4:64 8:64 16:64) ``` | Rate | Conc | Metric | Eager | Graph | Δ | % | |------|------|--------|-------|-------|-----|---| | 2 | 64 | Output tok/s | 167.9 | 138.5 | −29.4 | −17.5% | | | | TTFT mean / p99 (ms) | 70301 / 76996 | 86034 / 95290 | +15733 / +18294 | +22% | | | | TPOT mean / p99 (ms) | 74.5 / 78.9 | 90.7 / 93.9 | +16.2 / +15.0 | +22% / +19% | | | | E2EL mean (ms) | 89306 | 109173 | +19867 | +22% | | 4 | 64 | Output tok/s | 723.1 | 873.2 | +150.1 | **+20.8%** | | | | TTFT mean / p99 (ms) | 1095 / 1720 | 198 / 323 | −898 / −1397 | **−82% / −81%** | | | | TPOT mean / p99 (ms) | 75.5 / 77.8 | 33.1 / 84.9 | −42.4 / +7.1 | **−56%** / +9% | | | | E2EL mean (ms) | 20345 | 8645 | −11700 | −58% | | 8 | 64 | Output tok/s | 750.7 | 772.7 | +22.0 | +2.9% | | | | TTFT mean / p99 (ms) | 1094 / 1752 | 1354 / 2272 | +260 / +520 | +24% / +30% | | | | TPOT mean / p99 (ms) | 74.7 / 76.1 | 71.1 / 83.1 | −3.6 / +6.9 | −5% / +9% | | | | E2EL mean (ms) | 20141 | 19477 | −664 | −3% | | 16 | 64 | Output tok/s | 794.6 | 972.1 | +177.4 | **+22.3%** | | | | TTFT mean / p99 (ms) | 1065 / 1941 | 1401 / 2560 | +336 / +619 | +32% / +32% | | | | TPOT mean / p99 (ms) | 71.9 / 75.5 | 55.8 / 67.5 | −16.1 / −8.0 | **−22% / −11%** | | | | E2EL mean (ms) | 19389 | 15624 | −3765 | −19% | Completed = 512/512 at every point in both runs; no failures. ### Findings 1. **Saturation**: both runs are saturated across the *entire* sweep — even at 2 req/s, mean TTFT is 70–86 s (pure queue wait), and steady-state output throughput plateaus at ~750 (eager) / ~770–970 (graph) tok/s regardless of offered rate. With `--max-num-batched-tokens 64`, each 1024-token prompt needs 16 chunked-prefill steps, so the offered load exceeds capacity at every point; this sweep measured **queueing behavior under overload, not headroom**. 2. **What moved**: exactly what FULL_DECODE_ONLY graphs should move — **decode**. TPOT improved −56% at 4 req/s and −22% at 16 req/s; TTFT moved the *other* way at 8–16 req/s (+24–32%), consistent with prefill running eager in both cases while graph capture adds memory pressure/interference with chunked prefill. At 4 req/s graph improved everything simultaneously because it drained the queue faster. 3. **Divergence**: the 2 req/s point is inverted (graph −17.5% throughput) — but at that point both runs are fully backlogged, so it reflects deep-queue dynamics rather than steady serving; treat it as noise-of-overload. 4. **Headline**: under sustained load the graph recipe sustains **~22% higher output throughput** (972 vs 795 tok/s at 16 req/s) and **22% lower TPOT**, at the cost of ~32% higher TTFT — i.e., better decode cost and capacity, slightly worse prefill/queue latency. 5. The 8 req/s point is nearly within noise (+2.9% throughput); the meaningful capacity difference shows at 16 req/s. ### Artifacts - `./reports/2a2f_eager_dbo_dp1tp2-20261018-111201/` (one JSON per sweep point + `run.json`) - `./reports/2a2f_graph_dbo_dp1tp2-20261018-114152/` (one JSON per sweep point + `run.json`) ## Docs Impact - Files updated: .agents/skills/run-vllm-bench-k8s/ ---
Essential PR Checklist - [x] Purpose is clear and linked to public context when possible. - [x] Scope is bounded. - [x] Compatibility with vLLM v0.30.0 is considered. - [x] No changes are made to the vLLM source checkout. - [x] Plugin-owned classes or explicit dotted class paths are preferred over monkey patches. - [x] Any compat shim or monkey patch is isolated, idempotent, version-guarded, documented, and tested. - [x] Imports remain CPU-safe; CUDA-heavy work is delayed or GPU-gated. - [x] Validation evidence is included, including skipped GPU tests when applicable. - [x] Documentation impact is stated.
--------- Signed-off-by: Shlomit Koyfman --- .agents/skills/run-vllm-bench-k8s/SKILL.md | 412 ++++++++++++++++++ .../run-vllm-bench-k8s/agents/openai.yaml | 4 + .../references/comparing-runs.md | 98 +++++ .../references/fetching-past-reports.md | 84 ++++ 4 files changed, 598 insertions(+) create mode 100644 .agents/skills/run-vllm-bench-k8s/SKILL.md create mode 100644 .agents/skills/run-vllm-bench-k8s/agents/openai.yaml create mode 100644 .agents/skills/run-vllm-bench-k8s/references/comparing-runs.md create mode 100644 .agents/skills/run-vllm-bench-k8s/references/fetching-past-reports.md diff --git a/.agents/skills/run-vllm-bench-k8s/SKILL.md b/.agents/skills/run-vllm-bench-k8s/SKILL.md new file mode 100644 index 000000000..e418a05fa --- /dev/null +++ b/.agents/skills/run-vllm-bench-k8s/SKILL.md @@ -0,0 +1,412 @@ +--- +name: run-vllm-bench-k8s +description: Use when the user asks to benchmark an AFD GPU recipe on Kubernetes/OpenShift with vllm bench serve (or the repo's tools/benchmarks/request_generator.sh) - measuring throughput, TTFT, TPOT, or ITL for a deployed recipe and producing a report for that run. Deploys via deploy-afd-k8s. Do not use for local (non-k8s) benchmarking, NPU recipes, or E2E correctness testing (see run-e2e). +--- + +# Benchmark an AFD recipe with `vllm bench serve` + +Deploy a recipe, drive load against it with the repo's +[tools/benchmarks/request_generator.sh](../../../tools/benchmarks/request_generator.sh) +(a thin wrapper around `vllm bench serve`), and produce a self-contained +report for **that one run**. + +One invocation = one run = one `RUN_ID` = one report. + +## Scope + +- GPU colocation recipes only -- whatever `deploy-afd-k8s` supports. +- Not for local benchmarking (run `request_generator.sh` directly against + `decode_bench_server.sh` for that) or correctness E2E (`run-e2e`). +- Requires everything `deploy-afd-k8s` requires. No extra image: the bench + pod runs `AFD_PLUGIN_IMAGE`, which already carries the repo at + `/opt/afd-plugin` and vLLM's `bench` CLI. + +## Workflow + +### 1. Establish the run identity + +```bash +RECIPE_SCRIPT_PATH= +RUN_ID="$(basename "${RECIPE_SCRIPT_PATH}" .sh)-$(date +%Y%m%d-%H%M%S)" +LOCAL_DIR="./reports/${RUN_ID}" +echo "RUN_ID=${RUN_ID}" +``` + +Every artifact is keyed on `RUN_ID`, which is what keeps one run's numbers +from being confused with another's. + +If the recipe is not provided, ask which recipe to use. Use the `./recipe` +folder as the list of possible recipes. + +### 2. Deploy the recipe + +Follow **`deploy-afd-k8s`** with `RECIPE_SCRIPT_PATH`. It returns +`MODEL_ID`, `PVC_NAME`, `GPU_COUNT`, the serving node, the recipe's +`--max-num-batched-tokens`, and an endpoint at `http://vllm-service:18305`. + +`vllm bench serve` is pointed at that endpoint with `HOST=vllm-service` +`PORT=18305`, and `MODEL_PATH=${MODEL_ID}` -- it must name exactly what the +pod serves, since the model id is sent in each request and used to load the +tokenizer. + +### 3. Choose the workload + +**If the user did not say what to run, ask before deploying anything.** Put +the question as a choice between the named workloads below, and say which +one you will use if they have no preference. Do not silently pick one -- the +workload *is* the result: a throughput number without its prompt shape and +offered load is not quotable. + +Ask in this form: + +> Which workload should I benchmark `` with? +> 1. **balanced** (default) -- random 1024-token prompts, 256-token +> outputs, Poisson arrivals swept 2 -> 16 req/s. General serving +> behaviour and where latency starts to climb. +> 2. **max-throughput** -- same prompts, all requests issued at once, +> concurrency swept 16 -> 128. The peak tokens/s this recipe can reach. +> 3. **prefix-heavy** -- 10 shared 256-token prefixes reused across +> requests. Prefix-cache behaviour. +> 4. **long-context** -- 8192-token prompts, 256-token outputs. Prefill +> cost on long inputs. +> 5. **smoke** -- 64 short requests, ~1 minute. Checks the wiring only; +> never a published number. +> Or give me your own ISL/OSL, rate and concurrency. + +#### Built-in defaults + +Each workload is a complete set of values -- adopt them as-is unless the +user overrides a specific knob. `SWEEP` is a list of `rate:concurrency` +points; `vllm bench serve` is **single-shot**, so each point is one separate +invocation of the wrapper (see step 4). + +| Workload | `DATASET_NAME` | `INPUT_LEN` | `OUTPUT_LEN` | `NUM_PROMPTS` | `SWEEP` (`rate:conc`) | Extra flags | +|---|---|---|---|---|---|---| +| **balanced** (default) | `random` | 1024 | 256 | 512 | `2:64 4:64 8:64 16:64` | -- | +| **max-throughput** | `random` | 1024 | 256 | 512 | `inf:16 inf:32 inf:64 inf:128` | -- | +| **prefix-heavy** | `prefix_repetition` | (unused) | (unused) | 512 | `2:64 4:64 8:64` | `--prefix-repetition-prefix-len 256 --prefix-repetition-suffix-len 256 --prefix-repetition-num-prefixes 10 --prefix-repetition-output-len 256` | +| **long-context** | `random` | 8192 | 256 | 256 | `1:32 2:32 4:32` | -- | +| **smoke** | `random` | 256 | 128 | 64 | `inf:8` | -- | + +`REQUEST_RATE=inf` means "issue every request at time 0", so with +`MAX_CONCURRENCY` it is a closed-loop test measuring peak throughput. A +finite rate uses Poisson arrivals and measures behaviour under an offered +load, which is what latency SLOs care about. The two answer different +questions -- do not mix them in a single sweep. + +`prefix_repetition` ignores `INPUT_LEN`/`OUTPUT_LEN` and takes its shape +from its own flags, which go in `EXTRA_ARGS`. It must **not** be given +`--dataset-path` (nor may `random`); vLLM rejects that combination. + +#### Other datasets + +`--dataset-name` also accepts `sharegpt`, `burstgpt`, `sonnet`, `hf`, +`custom`, and several multimodal variants. These need a real corpus, so pass +`--dataset-path ` through `EXTRA_ARGS` and mount it into the +bench pod. Use one when the user asks for realistic traffic rather than +synthetic prompts; otherwise prefer `random`, whose ISL/OSL are exactly +controllable and therefore reproducible. + +#### Sizing the sweep against the recipe + +Match offered load to the recipe's `--max-num-batched-tokens` from step 2. A +recipe pinned at 64 absorbs a 1024-token prompt in ~16 chunked-prefill +steps, so a high rate puts it in permanent backpressure and the report +measures queueing rather than the recipe. Long-context makes this sharper: +an 8192-token prompt at that budget needs ~128 steps. + +Sweep until at least one point saturates -- TTFT climbing, or `completed` +below `NUM_PROMPTS`. A sweep where nothing saturates measured headroom, not +capacity, and the report has to say so. + +### 4. Run the sweep + +The pod runs the AFD image, so the wrapper is already on disk at +`/opt/afd-plugin/tools/benchmarks/request_generator.sh`. Two details: + +- The wrapper calls `uv run vllm bench serve`. The image's base Python + environment is already correct, so a `uv` shim on `PATH` strips `uv run` -- + resolved before the shim is installed, since afterwards `command -v uv` + would find the shim itself and recurse. +- `EXTRA_ARGS` adds `e2el` and p90. `vllm bench serve` defaults + `--percentile-metrics` to `ttft,tpot,itl` (no `e2el`) and + `--metric-percentiles` to `99` only, so without this the report cannot + show end-to-end latency or p90. + +Export the workload chosen in step 3 -- these are the `balanced` defaults. +`envsubst` only sees **exported** variables -- an unexported name +substitutes as empty with no error, so `export` every name in the list +below, including `TEMPLATE_IMAGE`/`MODEL_ID`/`RUN_ID`/`PVC_NAME` carried +over from step 2, and `VLLM_NODE` too if you will pin the pod (see below): + +```bash +export DATASET_NAME=random +export INPUT_LEN=1024 +export OUTPUT_LEN=256 +export NUM_PROMPTS=512 +export SWEEP="2:64 4:64 8:64 16:64" # rate:concurrency +export WORKLOAD_ARGS="" # e.g. the --prefix-repetition-* flags + +kubectl delete pod vllm-bench --ignore-not-found + +envsubst '${TEMPLATE_IMAGE} ${MODEL_ID} ${RUN_ID} ${SWEEP} ${DATASET_NAME} ${INPUT_LEN} ${OUTPUT_LEN} ${NUM_PROMPTS} ${WORKLOAD_ARGS} ${PVC_NAME} ${VLLM_NODE}' <<'EOF' | kubectl apply -f - +apiVersion: v1 +kind: PersistentVolumeClaim +metadata: + name: vllm-bench-reports +spec: + accessModes: + - ReadWriteOnce + resources: + requests: + storage: 5Gi +--- +apiVersion: v1 +kind: Pod +metadata: + name: vllm-bench + labels: + app: vllm-bench +spec: + restartPolicy: Never + activeDeadlineSeconds: 14400 + containers: + - name: bench + image: ${TEMPLATE_IMAGE} + imagePullPolicy: Always + command: ["/bin/bash", "-c"] + args: + - | + set -uo pipefail + mkdir -p /work/bin + REAL_UV="$(command -v uv || true)" + printf '#!/bin/bash\nif [ "$1" = "run" ]; then\n shift\n exec "$@"\nfi\nREAL_UV="%s"\n[ -n "$REAL_UV" ] || { echo "uv not found in image" >&2; exit 127; }\nexec "$REAL_UV" "$@"\n' "$REAL_UV" > /work/bin/uv + chmod +x /work/bin/uv + export PATH="/work/bin:$PATH" + + GEN=/opt/afd-plugin/tools/benchmarks/request_generator.sh + [ -x "$GEN" ] || { echo "ERROR: $GEN missing from image"; sleep infinity; } + + OUT="/reports/${RUN_ID}" + mkdir -p "$OUT" + rc=0 + for point in ${SWEEP}; do + rate="${point%%:*}" + conc="${point##*:}" + echo "=== sweep point rate=${rate} concurrency=${conc} ===" + MODEL_PATH="${MODEL_ID}" \ + HOST=vllm-service PORT=18305 \ + DATASET_NAME="${DATASET_NAME}" \ + INPUT_LEN="${INPUT_LEN}" OUTPUT_LEN="${OUTPUT_LEN}" \ + NUM_PROMPTS="${NUM_PROMPTS}" \ + REQUEST_RATE="${rate}" MAX_CONCURRENCY="${conc}" \ + RESULT_DIR="$OUT" \ + RESULT_FILENAME="rate_${rate}_conc_${conc}.json" \ + EXTRA_ARGS="--percentile-metrics ttft,tpot,itl,e2el --metric-percentiles 90,99 --label rate_${rate}_conc_${conc} ${WORKLOAD_ARGS}" \ + bash "$GEN" || { echo "sweep point rate=${rate} conc=${conc} FAILED"; rc=1; } + done + + echo "$rc" > "$OUT/.exit_code" + echo "=== sweep finished, exit=$rc; holding pod so reports can be copied ===" + sleep infinity + env: + - name: USER + value: "vllm" + - name: HOME + value: "/work/home" + - name: HF_HOME + value: "/models/.hf_home" + - name: HF_TOKEN + valueFrom: + secretKeyRef: + name: hf-token-secret + key: token + optional: true + resources: + requests: + cpu: "4" + memory: 16Gi + limits: + cpu: "8" + memory: 32Gi + volumeMounts: + - name: work + mountPath: /work + - name: reports + mountPath: /reports + - name: model-storage + mountPath: /models + readOnly: true + volumes: + - name: work + emptyDir: {} + - name: reports + persistentVolumeClaim: + claimName: vllm-bench-reports + - name: model-storage + persistentVolumeClaim: + claimName: ${PVC_NAME} + readOnly: true +EOF +``` + +The loop does not `set -e`: one failing sweep point records `rc=1` but the +remaining points still run, so a single bad rate does not discard the whole +sweep. The sentinel is written either way. + +**The model-storage mount above is required whenever `MODEL_ID` is an +in-container path** rather than a HF repo id (step 2) -- which is the common +case, since every recipe in this repo defaults `MODEL_PATH` to one. Without +it, `vllm bench serve` resolves `--model` as a HF repo id and fails with +`HFValidationError: Repo id must be in the form ...`, because it cannot load +the tokenizer from a path that doesn't exist in the bench pod. If `MODEL_ID` +is a genuine HF repo id instead, the mount is harmless to leave in place. + +**Pin the pod to the serving node only if the model PVC is `ReadWriteOnce`**: + +```bash +kubectl get pvc "${PVC_NAME}" -o jsonpath='{.spec.accessModes[0]}' +``` + +If that prints `ReadWriteOnce`, a second pod can only attach the PVC from +the node that already holds it. `export VLLM_NODE=`, add `${VLLM_NODE}` to the `envsubst` variable list above, and add +`nodeName: ${VLLM_NODE}` at the pod spec level (sibling of `restartPolicy`) +-- otherwise the pod leaves `Pending` with `Multi-Attach error for volume +...`. If it prints `ReadWriteMany`, skip `nodeName` entirely -- no pinning +needed, any node can attach it. + +### 5. Wait for completion, then collect + +The pod holds after the sweep, so completion is the sentinel file, not a +`Succeeded` phase -- which also preserves the exit code: + +```bash +kubectl logs -f pod/vllm-bench 2>/dev/null || true + +deadline=$(( $(date +%s) + 7200 )) +RC="" +while true; do + phase="$(kubectl get pod vllm-bench -o jsonpath='{.status.phase}' 2>/dev/null || true)" + case "$phase" in + "") echo "vllm-bench pod is gone"; exit 1 ;; + Failed) echo "vllm-bench Failed (OOM, or activeDeadlineSeconds hit)" + kubectl describe pod vllm-bench | tail -20 + kubectl logs vllm-bench --tail=100; exit 1 ;; + esac + RC="$(kubectl exec vllm-bench -- cat "/reports/${RUN_ID}/.exit_code" 2>/dev/null || true)" + [ -n "$RC" ] && break + if [ "$(date +%s)" -ge "$deadline" ]; then + echo "timed out after 2h waiting for the sweep (phase=${phase})" + kubectl logs vllm-bench --tail=50; exit 1 + fi + sleep 10 +done + +[ "$RC" = "0" ] || echo "at least one sweep point failed; the report must say which" +echo "=== sweep finished, exit=${RC} ===" + +mkdir -p "${LOCAL_DIR}" +kubectl cp "vllm-bench:/reports/${RUN_ID}/." "${LOCAL_DIR}" +ls "${LOCAL_DIR}" + +kubectl delete pod vllm-bench +``` + +Delete only **after** confirming the files are local -- that is the whole +reason the pod holds. + +**If that pod is already gone** -- an interrupted run, or an earlier +`RUN_ID` -- reports remain on the `vllm-bench-reports` PVC. See +[references/fetching-past-reports.md](references/fetching-past-reports.md). + +### 6. Record the run manifest + +```bash +cat > "${LOCAL_DIR}/run.json" <, + "workload": "", + "dataset": "${DATASET_NAME}", + "workload_args": "${WORKLOAD_ARGS}", + "input_len": ${INPUT_LEN}, + "output_len": ${OUTPUT_LEN}, + "num_prompts": ${NUM_PROMPTS}, + "sweep": "${SWEEP}", + "node": "${VLLM_NODE}", + "finished_utc": "$(date -u +%Y-%m-%dT%H:%M:%SZ)" +} +EOF +``` + +Recording the workload and sweep next to the numbers is what keeps the +report interpretable after the cluster is torn down, and lets +`comparing-runs.md` check that two runs are actually comparable. + +### 7. Report this run + +`${LOCAL_DIR}` holds one JSON per sweep point +(`rate__conc_.json`). Each is a flat object -- no nesting -- with +the run envelope (`date`, `model_id`, `num_prompts`, `request_rate`, +`max_concurrency`, `label`) merged with the metrics. + +**TTFT and TPOT are required in every report**: they separate prefill cost +from decode cost, and an AFD topology can move one without moving the other, +so end-to-end numbers alone hide the effect. + +| Report as | Key in the result JSON | Means | +|---|---|---| +| **TTFT** | `mean_ttft_ms`, `median_ttft_ms`, `p90_ttft_ms`, `p99_ttft_ms` | Prefill + queue wait. Rises first under backpressure -- the earliest saturation signal. | +| **TPOT** | `mean_tpot_ms`, `p90_/p99_tpot_ms` | Decode cost per output token, excluding the first. | +| ITL | `mean_itl_ms`, `p90_/p99_itl_ms` | Measured gap between consecutive streamed tokens. Related to TPOT, not identical -- do not substitute one for the other. | +| E2EL | `mean_e2el_ms`, `p90_/p99_e2el_ms` | End-to-end per request. Present only because step 4 passes `e2el` in `--percentile-metrics`. | +| Throughput | `output_throughput`, `total_token_throughput`, `request_throughput` | tokens/s and req/s. | +| Completion | `completed` vs `num_prompts`, plus `failed` | Error/drop rate. | + +All latency keys are **milliseconds** (the `_ms` suffix is literal); vLLM +reports throughput in tokens/s. State units once and stay consistent. + +``` +Run: +Harness: vllm bench serve (random ISL / OSL , prompts/point) +Recipe: (GPUs: , max-num-batched-tokens: ) +Model: +Profile: -- sweep + +| Rate | Conc | Output tok/s | TTFT mean | TTFT p99 | TPOT mean | TPOT p99 | E2EL mean | Completed | +|------|------|--------------|-----------|----------|-----------|----------|-----------|-----------| +``` + +Close with the saturation point -- the first sweep point where TTFT departs +from flat (it moves before end-to-end latency) or `completed` drops below +`num_prompts` -- and name any point that failed outright. If nothing +saturated, say so: the sweep measured headroom and the ceiling is above what +was offered. + +### 8. Comparing several runs + +See [references/comparing-runs.md](references/comparing-runs.md) -- +confirming two runs are comparable and reporting per-point deltas. Both +recipes being compared (e.g. an AFD recipe and its non-AFD baseline) must be +given explicitly by the user; never derive one from the other. + +### 9. Tear down the cluster + +`deploy-afd-k8s` leaves `vllm-pod`/`vllm-service` running by default so a +*follow-up* run can skip the model download and graph capture, but this +skill's own run is done once the report is produced -- tear both down at the +end, every time, without waiting to be asked: + +```bash +kubectl delete pod vllm-bench --ignore-not-found # already deleted in step 5 on the happy path +kubectl delete pod vllm-pod --ignore-not-found +kubectl delete svc vllm-service --ignore-not-found +``` + +Keep `vllm-bench-reports` until every run of interest has been copied out. diff --git a/.agents/skills/run-vllm-bench-k8s/agents/openai.yaml b/.agents/skills/run-vllm-bench-k8s/agents/openai.yaml new file mode 100644 index 000000000..3c9f1f053 --- /dev/null +++ b/.agents/skills/run-vllm-bench-k8s/agents/openai.yaml @@ -0,0 +1,4 @@ +interface: + display_name: "Benchmark AFD with vllm bench" + short_description: "Run a vllm bench serve sweep against an AFD recipe on k8s" + default_prompt: "Use $run-vllm-bench-k8s to benchmark an AFD GPU recipe on Kubernetes with vllm bench serve and report the results." diff --git a/.agents/skills/run-vllm-bench-k8s/references/comparing-runs.md b/.agents/skills/run-vllm-bench-k8s/references/comparing-runs.md new file mode 100644 index 000000000..a6c15985e --- /dev/null +++ b/.agents/skills/run-vllm-bench-k8s/references/comparing-runs.md @@ -0,0 +1,98 @@ +# Appendix: comparing multiple runs + +The main skill produces one report per `RUN_ID`. This appendix covers +comparing several of them -- most often an AFD recipe against its +GPU-equivalent non-AFD baseline, but equally two AFD configurations +(eager vs graph, DP2TP1 vs DP1TP2) or the same recipe across images. + +**Both recipes must be named explicitly by the user up front.** Never +derive, construct, or guess a counterpart recipe (e.g. a non-AFD baseline) +on the AFD recipe's behalf -- if the user asks for a comparison but gives +only one recipe, stop and ask them for the second `RECIPE_SCRIPT_PATH` +rather than building one. + +## 1. Confirm the runs are comparable + +Read every `run.json` first. A comparison is only meaningful when the runs +agree on **model, GPU count, and workload** (same dataset, ISL/OSL, +`NUM_PROMPTS`, and the same `SWEEP` points). Anything else being equal is a +bonus, not a requirement. + +```bash +for d in ./reports/*/; do + echo "--- $d"; cat "$d/run.json" +done +``` + +If a run you want is missing locally -- it was never copied out, or predates +the current session -- fetch it from the reports PVC first, per +[fetching-past-reports.md](fetching-past-reports.md). + +Refuse to produce a delta if the workloads differ -- rerun instead. Two runs +at different offered rates or prompt shapes are not a comparison, and a +percentage between them is worse than no number at all. If the images or +nodes differ, the comparison is still valid but say so in the report: node +class and image are confounds worth naming. + +Also confirm each run's report directory is genuinely its own -- one +`RUN_ID`, one directory. If two runs somehow share a path, discard both: the +sweep writes one file per point, so a shorter sweep leaves the earlier run's +files in place and the delta is computed against stale data. + +## 2. Run the set + +Run the main skill once per recipe, **sequentially** -- each run tears down +`vllm-pod`/`vllm-service` in its own step 9, so a second run cannot start +until the first has finished and the previous recipe's pod/service are gone. +Keep `MODEL_ID`, `PVC_NAME`, `GPU_COUNT`, the image, and the whole workload +(dataset, ISL/OSL, `NUM_PROMPTS`, `SWEEP`) fixed; change only +`RECIPE_SCRIPT_PATH` (and therefore `RUN_ID`). + +Reuse the same model PVC throughout so weights download once. Prefer running +back to back on the same node: a node change between runs is a confound. + +## 3. Compare + +Compare **per sweep point**. At low offered rate variants usually look +alike; the interesting result is the rate at which one starts queueing and +the other doesn't. + +For each sweep point present in all runs -- matched by `rate:concurrency`, +never across different sweeps -- tabulate the same metrics the single-run +report uses. **TTFT (`mean_ttft_ms`, `p99_ttft_ms`) and TPOT +(`mean_tpot_ms`, `p99_tpot_ms`) are required**, alongside `mean_e2el_ms`, +`mean_itl_ms`, `output_throughput`, and `completed` against `num_prompts`. + +TTFT and TPOT are what make an AFD-vs-baseline delta interpretable: splitting +attention from FFN changes prefill and decode by different amounts, so a +single end-to-end number can stay flat while TTFT improves and TPOT +regresses (or the reverse). Report the two separately and say which moved. + +Give both absolute delta and percentage, always naming which run is the +reference: + +``` +Reference: +Compared: +Model , GPUs, workload (identical across runs) + +| Rate | Conc | Metric | Reference | Compared | Delta | % | +|------|------|--------|-----------|----------|-------|---| +``` + +Call out, in this order: + +1. **Saturation point per run** -- the first sweep point where TTFT departs + from flat (it moves before end-to-end latency) or `completed` drops below + `num_prompts`. A recipe that saturates two points later is the headline + result, more than any single-point percentage. +2. **Which of TTFT / TPOT moved**, and in which direction. "Faster" with no + split between prefill and decode is not a usable finding. +3. **Sweep points where the runs diverge**, with the direction named. +4. **Sweep points where they are within noise** -- say so explicitly rather than + reporting a small percentage as a finding. +5. **Any run where no point saturated** -- the ramp measured headroom, so + the comparison bounds the difference from below and nothing more. + +A single averaged percentage as the sole result is not an acceptable +report: it depends entirely on how much of the sweep sat past saturation. diff --git a/.agents/skills/run-vllm-bench-k8s/references/fetching-past-reports.md b/.agents/skills/run-vllm-bench-k8s/references/fetching-past-reports.md new file mode 100644 index 000000000..e58f359f4 --- /dev/null +++ b/.agents/skills/run-vllm-bench-k8s/references/fetching-past-reports.md @@ -0,0 +1,84 @@ +# Fetching reports whose pod is already gone + +The main skill's step 5 copies reports straight out of the `vllm-bench` +pod, which holds open until the copy is done. Use this instead when that pod +no longer exists: + +- the run was interrupted, or the pod was deleted, evicted, or hit + `activeDeadlineSeconds` +- you want an **earlier** `RUN_ID` from a previous run, e.g. to complete a + comparison (see [comparing-runs.md](comparing-runs.md)) +- you want to see which `RUN_ID`s are on the PVC at all + +Reports live on a reports PVC, independently of any pod, so every past run +is still retrievable as long as that PVC exists. Mount it from a short-lived +busybox helper, pinned to the serving node -- the PVC is `ReadWriteOnce`, so +a second pod can only attach it from the node where it is already mounted, +otherwise `Multi-Attach error for volume ...` leaves the helper `Pending` +indefinitely. + +Set `RUN_ID` to the run you want and `LOCAL_DIR` to where it should land: + +```bash +RUN_ID= +LOCAL_DIR="./reports/${RUN_ID}" +``` + +```bash +HELPER=vllm-bench-reports-copy +VLLM_NODE="$(kubectl get pod vllm-pod -o jsonpath='{.spec.nodeName}')" + +kubectl delete pod "${HELPER}" --ignore-not-found +envsubst '${HELPER} ${VLLM_NODE}' <<'EOF' | kubectl apply -f - +apiVersion: v1 +kind: Pod +metadata: + name: ${HELPER} + labels: + app: vllm-bench + role: reports-copy +spec: + restartPolicy: Never + nodeName: ${VLLM_NODE} + containers: + - name: copy + image: busybox:1.36 + command: ["sleep", "3600"] + volumeMounts: + - name: reports + mountPath: /reports + readOnly: true + volumes: + - name: reports + persistentVolumeClaim: + claimName: vllm-bench-reports + readOnly: true +EOF + +deadline=$(( $(date +%s) + 300 )) +until [ "$(kubectl get pod "${HELPER}" -o jsonpath='{.status.phase}' 2>/dev/null)" = "Running" ]; do + phase="$(kubectl get pod "${HELPER}" -o jsonpath='{.status.phase}' 2>/dev/null || true)" + if [ "$phase" = "Failed" ] || [ "$(date +%s)" -ge "$deadline" ]; then + echo "helper pod not Running (phase=${phase:-})" + kubectl describe pod "${HELPER}" | tail -30 + kubectl delete pod "${HELPER}" --ignore-not-found + exit 1 + fi + sleep 5 +done + +kubectl exec "${HELPER}" -- ls /reports # which RUN_IDs are on the PVC +mkdir -p "${LOCAL_DIR}" +kubectl cp "${HELPER}:/reports/${RUN_ID}/." "${LOCAL_DIR}" +kubectl delete pod "${HELPER}" --ignore-not-found +``` + +`kubectl exec ... -- ls /reports` is the cheap way to see what is on the PVC +before choosing a `RUN_ID`. + +The helper mounts the PVC `readOnly`, so it cannot disturb a run in +progress. It is still safe to skip the node pin only if no other pod has the +PVC attached -- pinning is harmless either way, so keep it. + +If the PVC itself is gone, the reports are gone: nothing here recovers them. +Copy reports out promptly rather than relying on the PVC as archival storage. From bfaed8ba8239dcb4b9d974ee952756b3cefae3cd Mon Sep 17 00:00:00 2001 From: ronenkat <16743404+ronenkat@users.noreply.github.com> Date: Wed, 28 Oct 2026 11:32:23 +0200 Subject: [PATCH 16/32] feat(benchmarks): decode-only benchmark over a pinned real-KV prompt pool Add a prompt-pool mode to AFDDecodeBenchConnector (num_prompts > 0 in kv_connector_extra_config). Each Attention DP rank turns its first num_prompts requests into a seeded random prompt pool, identical on every rank, prefills it for real, and pins the finished warmup requests' KV blocks through the connector's delayed-free contract so LRU eviction can never drop a pool entry. Every later request is remapped round-robin onto a pool prompt regardless of its content; the prompt length is rounded to one token past a block boundary, so the prefix cache covers all but the last token and each measured request prefills exactly one token before decoding over real KV. A startup check fails fast when the pool and the running requests cannot both fit the KV cache, and reports the largest pool that fits. num_prompts = 0 keeps the dummy-KV mode unchanged. Add decode_bench_client.py: warms every DP rank with num_prompts requests pinned by the X-data-parallel-rank header, then sends the measured requests through vLLM's load balancer and reports TTFT/TPOT/ITL, throughput, and prefilled tokens per request from the prefix-cache counters (1.00 = decode only). Add a Qwen3-30B-A3B 2A2F single-node recipe with selectable AFD connector (DeepEPM2NAFDConnector or P2pNcclAFDConnector) and KV connector (none, decode_bench, or a raw KV_TRANSFER_CONFIG), usable for plain serving as well; the P2pNcclAFDConnector directory carries a thin entry point. It replaces decode_bench_server.sh. Name the connector logger under the vllm tree: vLLM only configures handlers for "vllm", so INFO lines from other logger names are dropped. Validated on 4xH200 with both connectors: plain serving smoke, and the decode-only benchmark with 1000 x 1025-token prompts per rank, 2000 measured requests, 0 prefix-cache misses, 1.00 prefilled tokens/request. Co-Authored-By: Claude Opus 5.5 Signed-off-by: ronenkat <16743404+ronenkat@users.noreply.github.com> --- .agents/skills/deploy-afd-k8s/SKILL.md | 3 +- .agents/skills/run-vllm-bench-k8s/SKILL.md | 5 +- recipe/README.md | 6 +- .../qwen3_30b_a3b/README.md | 121 +++++ .../2a2f_graph_dp2tp1.sh | 153 ++++++ .../qwen3_30b_a3b/README.md | 11 + .../2a2f_graph_dp2tp1.sh | 12 + tools/benchmarks/decode_bench.py | 336 ++++++++++++- tools/benchmarks/decode_bench_client.py | 455 ++++++++++++++++++ tools/benchmarks/decode_bench_server.sh | 104 ---- tools/benchmarks/request_generator.sh | 7 +- 11 files changed, 1089 insertions(+), 124 deletions(-) create mode 100644 recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/README.md create mode 100755 recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/2a2f_graph_dp2tp1.sh create mode 100644 recipe/gpu/P2pNcclAFDConnector/qwen3_30b_a3b/README.md create mode 100755 recipe/gpu/P2pNcclAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/2a2f_graph_dp2tp1.sh create mode 100644 tools/benchmarks/decode_bench_client.py delete mode 100755 tools/benchmarks/decode_bench_server.sh diff --git a/.agents/skills/deploy-afd-k8s/SKILL.md b/.agents/skills/deploy-afd-k8s/SKILL.md index f51ae4190..123de61bf 100644 --- a/.agents/skills/deploy-afd-k8s/SKILL.md +++ b/.agents/skills/deploy-afd-k8s/SKILL.md @@ -44,7 +44,8 @@ than one pod) -- callers never talk to these directly. See `# For pod N: POD=...` line per pod, but the comment isn't guaranteed to use that exact grammar going forward. See [resolve-recipe.md](references/resolve-recipe.md). -- Not for `tools/benchmarks/decode_bench_server.sh` (local, non-k8s). +- Not for the local decode-only benchmark (`tools/benchmarks/decode_bench_client.py` + against `recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/`, non-k8s). ## Requirements diff --git a/.agents/skills/run-vllm-bench-k8s/SKILL.md b/.agents/skills/run-vllm-bench-k8s/SKILL.md index e418a05fa..ac8c50087 100644 --- a/.agents/skills/run-vllm-bench-k8s/SKILL.md +++ b/.agents/skills/run-vllm-bench-k8s/SKILL.md @@ -15,8 +15,9 @@ One invocation = one run = one `RUN_ID` = one report. ## Scope - GPU colocation recipes only -- whatever `deploy-afd-k8s` supports. -- Not for local benchmarking (run `request_generator.sh` directly against - `decode_bench_server.sh` for that) or correctness E2E (`run-e2e`). +- Not for local benchmarking (run `request_generator.sh` or + `decode_bench_client.py` directly against a local recipe for that) or + correctness E2E (`run-e2e`). - Requires everything `deploy-afd-k8s` requires. No extra image: the bench pod runs `AFD_PLUGIN_IMAGE`, which already carries the repo at `/opt/afd-plugin` and vLLM's `bench` CLI. diff --git a/recipe/README.md b/recipe/README.md index c53a8b355..237a6f624 100644 --- a/recipe/README.md +++ b/recipe/README.md @@ -12,8 +12,11 @@ Recipes are organized by hardware backend, connector, and model: ```text recipe/ ├── gpu/ +│ ├── DeepEPM2NAFDConnector/ +│ │ └── qwen3_30b_a3b/ │ └── P2pNcclAFDConnector/ -│ └── deepseek_v2_lite/ +│ ├── deepseek_v2_lite/ +│ └── qwen3_30b_a3b/ └── npu/ ├── CAMAsyncAFDConnector/ │ └── deepseek_v3_2/ @@ -34,6 +37,7 @@ Directory names follow these conventions: | Hardware | Connector | Model | Recommended stage | v0.30 status | Recipe | | --- | --- | --- | --- | --- | --- | | GPU | `P2pNcclAFDConnector` | DeepSeek-V2-Lite | Decode | Legacy 2A2F E2E validated; examples inspected | [Launch examples](gpu/P2pNcclAFDConnector/deepseek_v2_lite/README.md) | +| GPU | `DeepEPM2NAFDConnector` or `P2pNcclAFDConnector` | Qwen3-30B-A3B | Prefill / decode; decode-only benchmark | 4xH200: serving smoke and decode-only benchmark validated on both connectors | [2A2F single node](gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/README.md) | | Ascend NPU | `CAMP2pAFDConnector` | DeepSeek-V3.2 | Decode | Integrated v0.30 hardware validation pending; linked recipe is historical | [Synchronous decode](npu/CAMP2pAFDConnector/deepseek_v3_2/README.md) | | Ascend NPU | `CAMAsyncAFDConnector` | DeepSeek-V3.2 | Prefill / decode | Integrated v0.30 hardware validation pending; historical v0.26 DP2TP8+EP16 token split reached `0.9522` strict match; legacy PCP8 results are v0.19-only | [Async CAM](npu/CAMAsyncAFDConnector/deepseek_v3_2/README.md) | diff --git a/recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/README.md b/recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/README.md new file mode 100644 index 000000000..30f7c34d2 --- /dev/null +++ b/recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/README.md @@ -0,0 +1,121 @@ +# Qwen3-30B-A3B AFD Recipe + +One launch script for Qwen3-30B-A3B in a 2A2F prefill/decode colocation on a +single 4-GPU node, with a selectable AFD connector and an optional KV +connector. With the decode-bench KV connector it doubles as a decode-only +benchmark over real KV (see [Decode-only benchmark](#decode-only-benchmark)). + +## Prerequisites + +- 4 GPUs on one node. `DeepEPM2NAFDConnector` needs SM90 or newer (H100, H200, + GB200) and DeepEP/DeepGEMM installed; `P2pNcclAFDConnector` runs on any + A/H-class GPU. +- vLLM `v0.30.0` and the `afd-plugin` package in the same environment. +- Qwen3-30B-A3B weights. The script defaults to + `/path/model_weights/Qwen3-30B-A3B`; override with `MODEL_PATH=...`. +- Free ports: `6269` for the AFD connector, `18305`/`18306` for the Attention + and FFN HTTP servers. + +## Topology + +| GPUs | Role | Parallelism | Port | +|------|-----------|----------------------|-------| +| 0, 1 | Attention | DP2 / TP1 | 18305 | +| 2, 3 | FFN | DP2 / TP1, EP2 | 18306 | + +Both roles run `FULL_DECODE_ONLY` CUDA graphs; prefill runs eager in chunks of +`MAX_NUM_BATCHED_TOKENS`. Send traffic to the Attention port. + +## Running + +Run from the repository root: + +```bash +export MODEL_PATH=/path/model_weights/Qwen3-30B-A3B +# DeepEP M2N (default) +bash recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/2a2f_graph_dp2tp1.sh +# NCCL P2P: same script with AFD_CONNECTOR=P2pNcclAFDConnector, or its entry point +bash recipe/gpu/P2pNcclAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/2a2f_graph_dp2tp1.sh +``` + +The script backgrounds both roles, logs to `attn.log` and `ffn.log` in the +current directory, and waits. Wait for `Application startup complete` in +`attn.log` before sending traffic. + +Every knob is an environment variable documented at the top of the script: +connector (`AFD_CONNECTOR`), KV connector (`KV_CONNECTOR`, or a raw +`KV_TRANSFER_CONFIG` for any other connector), prefix caching, DBO, +`MAX_NUM_SEQS`, `MAX_NUM_BATCHED_TOKENS`, graph capture sizes, devices, ports, +and `EXTRA_ARGS` for anything else. + +## Decode-only benchmark + +`KV_CONNECTOR=decode_bench` attaches `AFDDecodeBenchConnector` +(`tools/benchmarks/decode_bench.py`) to Attention in its prompt-pool mode: + +1. **Warmup.** Each Attention DP rank turns its first `DECODE_BENCH_NUM_PROMPTS` + requests into a pool of seeded random prompts of `DECODE_BENCH_PROMPT_LEN` + tokens (the same pool on every rank, whatever the requests carried) and + prefills them for real, one output token each. +2. **Pinning.** When a warmup request finishes, the connector keeps its KV + blocks allocated for the server's lifetime (it reports the blocks as still + being sent and never completes the send). Pinned blocks cannot be evicted + and stay in the prefix cache. +3. **Measured requests.** Every later request is remapped round-robin onto a + pool prompt, whatever it asked for. The prefix cache covers every prompt + token but the last (the prompt length is rounded to one token past a block + boundary), so the request prefills one token and decodes over real KV, on + whichever rank serves it. + +Start the server, then drive it with the two-phase client: + +```bash +export MODEL_PATH=/path/model_weights/Qwen3-30B-A3B +KV_CONNECTOR=decode_bench \ +DECODE_BENCH_NUM_PROMPTS=1000 \ +DECODE_BENCH_PROMPT_LEN=1025 \ +DECODE_BENCH_MAX_OUTPUT_LEN=256 \ + bash recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/2a2f_graph_dp2tp1.sh + +python tools/benchmarks/decode_bench_client.py \ + --port 18305 --dp-size 2 --num-prompts 1000 \ + --num-requests 4000 --output-len 256 --max-concurrency 64 \ + --result-file /tmp/results/decode_bench.json +``` + +The client warms every rank with `--num-prompts` requests pinned by the +`X-data-parallel-rank` header, then sends the measured requests through vLLM's +load balancer and reports TTFT, TPOT, ITL, and throughput. It also reports +`prefilled tokens/request` from the server's prefix-cache counters: `1.00` +means every measured request was decode-only. `--num-prompts` must match +`DECODE_BENCH_NUM_PROMPTS`, and `--output-len` must not exceed +`DECODE_BENCH_MAX_OUTPUT_LEN`. + +### Sizing the pool + +The pool must fit each Attention rank's KV cache next to the running requests. +At startup the connector checks, per rank, + +```text +num_prompts * ceil((prompt_len + 1) / block_size) + + max_num_seqs * ceil((max_output_len + 1) / block_size) <= free KV blocks +``` + +and fails with the largest `num_prompts` that fits. Qwen3-30B-A3B holds about +96 KiB of BF16 KV per token, so 1000 prompts of 1025 tokens need about 94 GiB +per rank: an H200 fits it, an H100 needs fewer or shorter prompts. Keep +`num_prompts` at or above `MAX_NUM_SEQS` (the connector warns otherwise), or +concurrent requests share physical KV blocks and attention reads look cheaper +than they are. + +### Limitations + +- Generated text is meaningless: prompts are random tokens, and the client's + prompt is discarded. +- Measured requests decode at a context of `prompt_len` regardless of the + client's prompt length. +- Pinned warmup requests keep each Attention engine loop stepping while idle + (one busy CPU core per DP rank, no GPU work). Restart the server to free the + pool. +- `DECODE_BENCH_NUM_PROMPTS=0` keeps the older dummy-KV mode, which writes no + KV at all. diff --git a/recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/2a2f_graph_dp2tp1.sh b/recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/2a2f_graph_dp2tp1.sh new file mode 100755 index 000000000..3d3766404 --- /dev/null +++ b/recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/2a2f_graph_dp2tp1.sh @@ -0,0 +1,153 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +# +# Qwen3-30B-A3B 2A2F prefill/decode colocation on one 4-GPU node: Attention +# DP2/TP1 and FFN DP2/TP1 (expert parallel), FULL_DECODE_ONLY CUDA graphs. +# Every knob is an environment variable (default in parentheses): +# +# MODEL_PATH weights dir or HF id (/path/model_weights/Qwen3-30B-A3B) +# AFD_CONNECTOR DeepEPM2NAFDConnector | P2pNcclAFDConnector +# (DeepEPM2NAFDConnector) +# KV_CONNECTOR none | decode_bench (none); see the README +# KV_TRANSFER_CONFIG raw --kv-transfer-config JSON for Attention; +# overrides KV_CONNECTOR (unset) +# ENABLE_PREFIX_CACHING 1 | 0 (1) +# ENABLE_DBO 1 | 0 (0) +# DBO_DECODE_TOKEN_THRESHOLD / DBO_PREFILL_TOKEN_THRESHOLD (2 / 12) +# MAX_NUM_SEQS (64) +# MAX_NUM_BATCHED_TOKENS prefill chunk per step (2048) +# CUDAGRAPH_CAPTURE_SIZES comma-separated decode graph sizes (MAX_NUM_SEQS) +# MAX_MODEL_LEN (unset: model default) +# GPU_MEMORY_UTILIZATION (unset: vLLM default) +# ATTN_DEVICES / FFN_DEVICES (0,1 / 2,3) +# AFD_HOST / AFD_PORT AFD connector rendezvous (127.0.0.1 / 6269) +# ATTN_PORT / FFN_PORT HTTP ports (18305 / 18306) +# EXTRA_ARGS extra `vllm serve` args for both roles (empty) +# +# KV_CONNECTOR=decode_bench only (tools/benchmarks/decode_bench.py): +# DECODE_BENCH_NUM_PROMPTS pool prompts per Attention rank; 0 = dummy KV (1000) +# DECODE_BENCH_PROMPT_LEN pool prompt length (1025) +# DECODE_BENCH_MAX_OUTPUT_LEN longest measured output; sizes the KV check (1024) +# DECODE_BENCH_SEED pool token seed (0) +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +REPO_ROOT="$(cd "${SCRIPT_DIR}/../../../../.." && pwd)" + +MODEL_PATH=${MODEL_PATH:-/path/model_weights/Qwen3-30B-A3B} +AFD_CONNECTOR=${AFD_CONNECTOR:-DeepEPM2NAFDConnector} +KV_CONNECTOR=${KV_CONNECTOR:-none} +KV_TRANSFER_CONFIG=${KV_TRANSFER_CONFIG:-} +ENABLE_PREFIX_CACHING=${ENABLE_PREFIX_CACHING:-1} +ENABLE_DBO=${ENABLE_DBO:-0} +DBO_DECODE_TOKEN_THRESHOLD=${DBO_DECODE_TOKEN_THRESHOLD:-2} +DBO_PREFILL_TOKEN_THRESHOLD=${DBO_PREFILL_TOKEN_THRESHOLD:-12} +MAX_NUM_SEQS=${MAX_NUM_SEQS:-64} +MAX_NUM_BATCHED_TOKENS=${MAX_NUM_BATCHED_TOKENS:-2048} +CUDAGRAPH_CAPTURE_SIZES=${CUDAGRAPH_CAPTURE_SIZES:-$MAX_NUM_SEQS} +MAX_MODEL_LEN=${MAX_MODEL_LEN:-} +GPU_MEMORY_UTILIZATION=${GPU_MEMORY_UTILIZATION:-} +ATTN_DEVICES=${ATTN_DEVICES:-0,1} +FFN_DEVICES=${FFN_DEVICES:-2,3} +AFD_HOST=${AFD_HOST:-127.0.0.1} +AFD_PORT=${AFD_PORT:-6269} +ATTN_PORT=${ATTN_PORT:-18305} +FFN_PORT=${FFN_PORT:-18306} +EXTRA_ARGS=${EXTRA_ARGS:-} +DECODE_BENCH_NUM_PROMPTS=${DECODE_BENCH_NUM_PROMPTS:-1000} +DECODE_BENCH_PROMPT_LEN=${DECODE_BENCH_PROMPT_LEN:-1025} +DECODE_BENCH_MAX_OUTPUT_LEN=${DECODE_BENCH_MAX_OUTPUT_LEN:-1024} +DECODE_BENCH_SEED=${DECODE_BENCH_SEED:-0} + +export VLLM_USE_V2_MODEL_RUNNER=1 + +case "$AFD_CONNECTOR" in + # DeepEP M2N routes on Attention: the gate and top-k run there. + DeepEPM2NAFDConnector) COMPUTE_GATE_ON_ATTENTION=true ;; + P2pNcclAFDConnector) COMPUTE_GATE_ON_ATTENTION=false ;; + *) + echo "unsupported AFD_CONNECTOR=${AFD_CONNECTOR}" >&2 + exit 1 + ;; +esac + +if [[ -z "$KV_TRANSFER_CONFIG" ]]; then + case "$KV_CONNECTOR" in + none) ;; + decode_bench) + KV_TRANSFER_CONFIG=$(cat <&2 + exit 1 + ;; + esac +fi +if [[ "$KV_TRANSFER_CONFIG" == *tools.benchmarks.* ]]; then + # tools/benchmarks is not in the wheel; the vLLM scheduler and worker + # subprocesses import it from the repository root. + export PYTHONPATH="${REPO_ROOT}${PYTHONPATH:+:${PYTHONPATH}}" +fi + +afd_config() { + cat < attn.log 2>&1 & + +CUDA_VISIBLE_DEVICES=$FFN_DEVICES uv run vllm serve "$MODEL_PATH" \ + --additional-config "$(afd_config ffn)" \ + "${COMMON_ARGS[@]}" \ + --host 127.0.0.1 \ + --port "$FFN_PORT" > ffn.log 2>&1 & + +wait diff --git a/recipe/gpu/P2pNcclAFDConnector/qwen3_30b_a3b/README.md b/recipe/gpu/P2pNcclAFDConnector/qwen3_30b_a3b/README.md new file mode 100644 index 000000000..cb715b288 --- /dev/null +++ b/recipe/gpu/P2pNcclAFDConnector/qwen3_30b_a3b/README.md @@ -0,0 +1,11 @@ +# Qwen3-30B-A3B AFD Recipe (P2pNcclAFDConnector) + +`prefill_decode_colocation/2a2f_graph_dp2tp1.sh` runs the shared Qwen3-30B-A3B +2A2F recipe with `AFD_CONNECTOR=P2pNcclAFDConnector`. Topology, environment +knobs, and the decode-only benchmark are documented with the script itself in +[`../../DeepEPM2NAFDConnector/qwen3_30b_a3b/README.md`](../../DeepEPM2NAFDConnector/qwen3_30b_a3b/README.md). + +```bash +export MODEL_PATH=/path/model_weights/Qwen3-30B-A3B +bash recipe/gpu/P2pNcclAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/2a2f_graph_dp2tp1.sh +``` diff --git a/recipe/gpu/P2pNcclAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/2a2f_graph_dp2tp1.sh b/recipe/gpu/P2pNcclAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/2a2f_graph_dp2tp1.sh new file mode 100755 index 000000000..18ba6bff1 --- /dev/null +++ b/recipe/gpu/P2pNcclAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/2a2f_graph_dp2tp1.sh @@ -0,0 +1,12 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +# +# P2pNcclAFDConnector entry point of the shared Qwen3-30B-A3B 2A2F recipe. +# The topology and every environment knob are documented in +# recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/2a2f_graph_dp2tp1.sh. +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +AFD_CONNECTOR=P2pNcclAFDConnector exec bash \ + "${SCRIPT_DIR}/../../../DeepEPM2NAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/2a2f_graph_dp2tp1.sh" diff --git a/tools/benchmarks/decode_bench.py b/tools/benchmarks/decode_bench.py index 6ddd2a1be..9412f82ea 100644 --- a/tools/benchmarks/decode_bench.py +++ b/tools/benchmarks/decode_bench.py @@ -1,9 +1,22 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project -"""Decode benchmark KV connector for AFD decode-side pressure tests.""" +"""Decode benchmark KV connector for AFD decode-side pressure tests. + +Two modes, chosen by ``kv_connector_extra_config``: + +- Dummy KV (``num_prompts`` unset or 0): every prompt token but the last is + reported as externally computed and never written, so requests skip prefill. +- Prompt pool (``num_prompts`` > 0): each Attention DP rank prefills a pool of + ``num_prompts`` seeded random prompts for real, from its first + ``num_prompts`` requests (the warmup), and pins their KV blocks in the prefix + cache. Every later request is remapped, whatever its content, onto a pool + prompt round-robin, hits the prefix cache for all but the last prompt token, + and goes straight to decode over real KV. +""" from __future__ import annotations +from collections import deque from dataclasses import dataclass from typing import TYPE_CHECKING, Any @@ -21,15 +34,51 @@ if TYPE_CHECKING: from vllm.config import VllmConfig from vllm.forward_context import ForwardContext - from vllm.v1.core.kv_cache_manager import KVCacheBlocks + from vllm.v1.core.kv_cache_manager import KVCacheBlocks, KVCacheManager + from vllm.v1.core.kv_cache_utils import BlockHash from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.kv_cache_interface import KVCacheConfig from vllm.v1.request import Request -logger = init_logger(__name__) +# vLLM configures handlers only for the "vllm" logger tree; a logger outside it +# falls back to the root WARNING level and drops the prompt-pool INFO lines. +logger = init_logger(f"vllm.{__name__}") KVCacheLayer = torch.Tensor | tuple[torch.Tensor, ...] +# A warmup request only has to prefill its pool prompt; one output token also +# fixes how many blocks each pinned pool entry holds. +WARMUP_OUTPUT_TOKENS = 1 +# Finished measured requests between two prompt-pool hit/miss log lines. +POOL_STATS_LOG_INTERVAL = 1000 +# Warmup requests that end this way have fully prefilled their pool prompt. +PINNABLE_WARMUP_STATUSES = ( + RequestStatus.FINISHED_STOPPED, + RequestStatus.FINISHED_LENGTH_CAPPED, +) + + +@dataclass(frozen=True) +class PromptPoolConfig: + """Prompt-pool keys of ``kv_connector_extra_config``. + + Attributes: + num_prompts: Pool entries per Attention DP rank; 0 keeps the dummy-KV + mode. + prompt_len: Requested pool prompt length. It is rounded down to one + token past a block boundary, where a prefix-cache hit leaves exactly + the last token to prefill. + seed: Seed of the pool tokens, shared by every rank so all ranks hold + the same pool. + max_output_len: Longest output a measured request decodes; sizes the KV + capacity check. ``None`` uses ``max_model_len - prompt_len``. + """ + + num_prompts: int = 0 + prompt_len: int = 1025 + seed: int = 0 + max_output_len: int | None = None + @dataclass class AFDDecodeBenchConnectorMetadata(KVConnectorMetadata): @@ -42,9 +91,10 @@ class AFDDecodeBenchConnectorMetadata(KVConnectorMetadata): class AFDDecodeBenchConnector(KVConnectorBase_V1, SupportsHMA): """KV connector used to benchmark AFD decode instances with long ISL. - The connector emulates remote prefill by marking all prompt tokens except - the final decode token as externally provided, then fills the allocated KV - cache pages with deterministic or random non-zero values on the worker. + In the dummy-KV mode the connector emulates remote prefill by marking all + prompt tokens except the final decode token as externally provided. In the + prompt-pool mode only a prefix-cache miss is emulated that way; see the + module docstring. """ def __init__( @@ -59,7 +109,10 @@ def __init__( self.connector_worker: AFDDecodeBenchConnectorWorker | None = None if role == KVConnectorRole.SCHEDULER: - self.connector_scheduler = AFDDecodeBenchConnectorScheduler(vllm_config) + self.connector_scheduler = AFDDecodeBenchConnectorScheduler( + vllm_config, + kv_cache_config, + ) elif role == KVConnectorRole.WORKER: self.connector_worker = AFDDecodeBenchConnectorWorker( vllm_config, @@ -91,6 +144,15 @@ def save_kv_layer( def wait_for_save(self) -> None: pass + def bind_kv_cache_manager(self, kv_cache_manager: KVCacheManager) -> None: + super().bind_kv_cache_manager(kv_cache_manager) + assert self.connector_scheduler is not None + self.connector_scheduler.bind_kv_cache_manager(kv_cache_manager) + + def on_new_request(self, request: Request) -> None: + assert self.connector_scheduler is not None + self.connector_scheduler.on_new_request(request) + def get_num_new_matched_tokens( self, request: Request, @@ -128,8 +190,11 @@ def request_finished( block_ids: list[int], ) -> tuple[bool, dict[str, Any] | None]: assert self.connector_scheduler is not None - self.connector_scheduler.request_finished(request) - return False, None + delay_free_blocks = self.connector_scheduler.request_finished( + request, + len(block_ids), + ) + return delay_free_blocks, None def request_finished_all_groups( self, @@ -137,24 +202,70 @@ def request_finished_all_groups( block_ids: tuple[list[int], ...], ) -> tuple[bool, dict[str, Any] | None]: assert self.connector_scheduler is not None - self.connector_scheduler.request_finished(request) - return False, None + delay_free_blocks = self.connector_scheduler.request_finished( + request, + sum(len(group_block_ids) for group_block_ids in block_ids), + ) + return delay_free_blocks, None class AFDDecodeBenchConnectorScheduler: """Scheduler-side implementation for AFDDecodeBenchConnector.""" - def __init__(self, vllm_config: VllmConfig) -> None: + def __init__( + self, + vllm_config: VllmConfig, + kv_cache_config: KVCacheConfig | None = None, + ) -> None: self.vllm_config = vllm_config self.block_size = vllm_config.cache_config.block_size self._filled_requests: set[str] = set() self._pending_fills: dict[str, tuple[tuple[list[int], ...], int]] = {} + kv_transfer_config = vllm_config.kv_transfer_config + assert kv_transfer_config is not None + pool_config = PromptPoolConfig( + num_prompts=kv_transfer_config.get_from_extra_config( + "num_prompts", PromptPoolConfig.num_prompts + ), + prompt_len=kv_transfer_config.get_from_extra_config( + "prompt_len", PromptPoolConfig.prompt_len + ), + seed=kv_transfer_config.get_from_extra_config( + "seed", PromptPoolConfig.seed + ), + max_output_len=kv_transfer_config.get_from_extra_config( + "max_output_len", PromptPoolConfig.max_output_len + ), + ) + self.prompt_pool = ( + PromptPool(pool_config, vllm_config, kv_cache_config) + if pool_config.num_prompts > 0 + else None + ) + + def bind_kv_cache_manager(self, kv_cache_manager: KVCacheManager) -> None: + if self.prompt_pool is not None: + self.prompt_pool.check_capacity( + kv_cache_manager.block_pool.get_num_free_blocks() + ) + + def on_new_request(self, request: Request) -> None: + if self.prompt_pool is not None: + self.prompt_pool.assign(request) + def get_num_new_matched_tokens( self, request: Request, num_computed_tokens: int, ) -> tuple[int, bool]: + if self.prompt_pool is not None: + num_tokens_to_fill = self.prompt_pool.num_tokens_to_fill( + request, + num_computed_tokens, + ) + return num_tokens_to_fill, False + req_id = request.request_id allow_refill_after_preempt = ( request.status == RequestStatus.PREEMPTED @@ -208,8 +319,205 @@ def build_connector_meta( self._pending_fills.clear() return meta - def request_finished(self, request: Request) -> None: + def request_finished(self, request: Request, num_blocks: int) -> bool: + """Return True to keep the request's blocks allocated (pinned).""" self._filled_requests.discard(request.request_id) + if self.prompt_pool is None: + return False + return self.prompt_pool.request_finished(request, num_blocks) + + +class PromptPool: + """Per-rank pool of real-KV prompts that measured requests decode over. + + Pinning: a warmup request that prefilled its pool prompt reports + ``delay_free_blocks`` from ``request_finished`` and is never returned from + ``get_finished``, so the scheduler keeps its blocks allocated for the + server's lifetime. Pinned blocks never enter the free queue, so LRU + eviction cannot drop a pool entry, and they stay in the prefix cache, so a + measured request that maps to the entry shares them like any prefix hit. + """ + + def __init__( + self, + config: PromptPoolConfig, + vllm_config: VllmConfig, + kv_cache_config: KVCacheConfig | None, + ) -> None: + if not vllm_config.cache_config.enable_prefix_caching: + raise ValueError( + "AFDDecodeBenchConnector prompt pool requires prefix caching; " + "launch with --enable-prefix-caching" + ) + if kv_cache_config is not None and len(kv_cache_config.kv_cache_groups) != 1: + raise ValueError( + "AFDDecodeBenchConnector prompt pool supports a single KV cache " + f"group, got {len(kv_cache_config.kv_cache_groups)}" + ) + + self.block_size = vllm_config.cache_config.block_size + self.max_num_seqs = vllm_config.scheduler_config.max_num_seqs + max_model_len = vllm_config.model_config.max_model_len + self.num_prompts = config.num_prompts + # A prefix-cache hit covers full blocks only and stops one token short + # of the prompt, so a prompt one token past a block boundary hits every + # token but the last. + self.prompt_len = ( + config.prompt_len - 1 + ) // self.block_size * self.block_size + 1 + if self.prompt_len == 1: + raise ValueError( + f"prompt_len ({config.prompt_len}) must exceed the KV block size " + f"({self.block_size})" + ) + self.max_output_len = ( + config.max_output_len + if config.max_output_len is not None + else max_model_len - self.prompt_len + ) + if self.max_output_len < 1 or ( + self.prompt_len + self.max_output_len > max_model_len + ): + raise ValueError( + f"prompt_len ({self.prompt_len}) + max_output_len " + f"({self.max_output_len}) must fit max_model_len ({max_model_len})" + ) + + # A CPU generator gives every rank the same pool for the same seed. + generator = torch.Generator().manual_seed(config.seed) + self.prompt_tokens = torch.randint( + vllm_config.model_config.get_vocab_size(), + (self.num_prompts, self.prompt_len), + generator=generator, + dtype=torch.int32, + ) + # Hashing a long prompt per request is avoidable work on the scheduler: + # an entry's hashes are computed once and copied after that. + self._entry_block_hashes: dict[int, list[BlockHash]] = {} + self._unwarmed_entries = deque(range(self.num_prompts)) + self._warming_entries: dict[str, int] = {} + self._num_pinned_entries = 0 + self._num_pinned_blocks = 0 + self._num_assigned_measured = 0 + # request_id -> prompt tokens the prefix cache missed (0 on a hit). + self._measured_misses: dict[str, int] = {} + self._num_finished_measured = 0 + self._num_missed_requests = 0 + self._num_missed_tokens = 0 + + def check_capacity(self, num_free_blocks: int) -> None: + """Fail fast when the pool and the running requests cannot both fit.""" + blocks_per_entry = cdiv( + self.prompt_len + WARMUP_OUTPUT_TOKENS, + self.block_size, + ) + # A measured request shares its entry's full blocks and allocates only + # the last prompt token and its output. + blocks_per_seq = cdiv(1 + self.max_output_len, self.block_size) + num_reserved_blocks = self.max_num_seqs * blocks_per_seq + num_pool_blocks = self.num_prompts * blocks_per_entry + if num_pool_blocks + num_reserved_blocks > num_free_blocks: + max_num_prompts = max( + 0, + (num_free_blocks - num_reserved_blocks) // blocks_per_entry, + ) + raise ValueError( + f"AFDDecodeBenchConnector prompt pool needs {num_pool_blocks} KV " + f"blocks for {self.num_prompts} prompts of {self.prompt_len} " + f"tokens plus {num_reserved_blocks} for {self.max_num_seqs} " + f"sequences of up to {self.max_output_len} output tokens, but " + f"only {num_free_blocks} are free; use num_prompts <= " + f"{max_num_prompts} or shorten prompt_len/max_output_len" + ) + if self.num_prompts < self.max_num_seqs: + logger.warning( + "AFDDecodeBenchConnector: %d pool prompts are fewer than " + "max_num_seqs (%d); concurrent requests will share KV blocks", + self.num_prompts, + self.max_num_seqs, + ) + logger.info( + "AFDDecodeBenchConnector prompt pool: %d prompts x %d tokens use " + "%d of %d free KV blocks; %d reserved for running requests", + self.num_prompts, + self.prompt_len, + num_pool_blocks, + num_free_blocks, + num_reserved_blocks, + ) + + def assign(self, request: Request) -> None: + """Rewrite a new request's prompt to a pool entry's, before scheduling. + + The first ``num_prompts`` requests warm one unwarmed entry each; every + later request maps round-robin onto the pool, whatever it asked for. + """ + if self._unwarmed_entries: + entry = self._unwarmed_entries.popleft() + self._warming_entries[request.request_id] = entry + request.max_tokens = WARMUP_OUTPUT_TOKENS + else: + entry = self._num_assigned_measured % self.num_prompts + self._num_assigned_measured += 1 + + prompt_token_ids = self.prompt_tokens[entry].tolist() + request.prompt_token_ids = prompt_token_ids + request.num_prompt_tokens = len(prompt_token_ids) + # all_token_ids is a read-only view of _all_token_ids: replace in place. + request._all_token_ids[:] = prompt_token_ids + request.block_hashes.clear() + entry_block_hashes = self._entry_block_hashes.get(entry) + if entry_block_hashes is None: + request.update_block_hashes() + self._entry_block_hashes[entry] = list(request.block_hashes) + else: + request.block_hashes.extend(entry_block_hashes) + + def num_tokens_to_fill(self, request: Request, num_computed_tokens: int) -> int: + """Prompt tokens to fake as external: 0 on a hit or during warmup.""" + if request.request_id in self._warming_entries: + return 0 + # A pinned entry hits every prompt token but the last. Anything the + # cache misses is faked so the step stays decode-only, and counted. + num_missed_tokens = max(0, request.num_tokens - 1 - num_computed_tokens) + self._measured_misses[request.request_id] = num_missed_tokens + return num_missed_tokens + + def request_finished(self, request: Request, num_blocks: int) -> bool: + """Pin a finished warmup request's blocks; tally measured hits.""" + entry = self._warming_entries.pop(request.request_id, None) + if entry is not None: + if request.status not in PINNABLE_WARMUP_STATUSES: + # Aborted before its prompt was cached: warm it again later. + self._unwarmed_entries.appendleft(entry) + return False + self._num_pinned_entries += 1 + self._num_pinned_blocks += num_blocks + if self._num_pinned_entries == self.num_prompts: + logger.info( + "AFDDecodeBenchConnector prompt pool warm: %d prompts " + "pinned in %d KV blocks", + self._num_pinned_entries, + self._num_pinned_blocks, + ) + return True + + num_missed_tokens = self._measured_misses.pop(request.request_id, None) + if num_missed_tokens is None: + return False + self._num_finished_measured += 1 + if num_missed_tokens > 0: + self._num_missed_requests += 1 + self._num_missed_tokens += num_missed_tokens + if self._num_finished_measured % POOL_STATS_LOG_INTERVAL == 0: + logger.info( + "AFDDecodeBenchConnector prompt pool: %d measured requests, " + "%d prefix-cache misses (%d prompt tokens faked)", + self._num_finished_measured, + self._num_missed_requests, + self._num_missed_tokens, + ) + return False class AFDDecodeBenchConnectorWorker: @@ -339,4 +647,6 @@ def _fill_blocks(self, group_idx: int, block_ids: list[int]) -> None: "AFDDecodeBenchConnectorMetadata", "AFDDecodeBenchConnectorScheduler", "AFDDecodeBenchConnectorWorker", + "PromptPool", + "PromptPoolConfig", ] diff --git a/tools/benchmarks/decode_bench_client.py b/tools/benchmarks/decode_bench_client.py new file mode 100644 index 000000000..ca71b3c0c --- /dev/null +++ b/tools/benchmarks/decode_bench_client.py @@ -0,0 +1,455 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Two-phase load generator for the decode benchmark's prompt-pool mode. + +Run it against an Attention instance whose ``AFDDecodeBenchConnector`` has +``num_prompts`` > 0 (see ``tools/benchmarks/decode_bench.py``): + +1. Warmup: sends ``--num-prompts`` requests to every Attention DP rank, pinned + with the ``X-data-parallel-rank`` header, so each rank prefills and pins its + whole prompt pool. Their metrics are discarded. +2. Measured: sends ``--num-requests`` requests through vLLM's load balancer. + The server remaps each onto a pool prompt, so the prompt sent here is only a + placeholder and every request decodes over real, prefix-cached KV. + +The server's prefix-cache counters before and after the measured phase give the +prompt tokens prefilled per measured request; a clean decode-only run prefills +exactly one. + +Usage: + python tools/benchmarks/decode_bench_client.py --num-prompts 1000 \ + --num-requests 4000 --output-len 256 --max-concurrency 64 +""" + +from __future__ import annotations + +import argparse +import asyncio +import json +import math +import random +import sys +import time +from collections.abc import Sequence +from dataclasses import dataclass, field +from pathlib import Path + +import aiohttp +import numpy as np + +DATA_PARALLEL_RANK_HEADER = "X-data-parallel-rank" +COMPLETIONS_PATH = "/v1/completions" +MODELS_PATH = "/v1/models" +METRICS_PATH = "/metrics" +PREFIX_CACHE_QUERIES_METRIC = "vllm:prefix_cache_queries" +PREFIX_CACHE_HITS_METRIC = "vllm:prefix_cache_hits" +SSE_DATA_PREFIX = b"data: " +SSE_DONE = b"[DONE]" +HTTP_OK = 200 +# Placeholder prompt token ids stay below this, which every vocabulary covers. +PLACEHOLDER_TOKEN_ID_LIMIT = 1000 +# The server forces warmup requests to one output token as well. +WARMUP_OUTPUT_TOKENS = 1 +READY_POLL_INTERVAL_S = 5.0 +PERCENTILES = (50, 90, 99) +MS_PER_S = 1000.0 + + +@dataclass +class RequestResult: + """Timing of one streamed completion.""" + + ok: bool + ttft_s: float = 0.0 + latency_s: float = 0.0 + output_tokens: int = 0 + itl_s: list[float] = field(default_factory=list) + error: str = "" + + +@dataclass(frozen=True) +class PlannedRequest: + payload: dict[str, object] + headers: dict[str, str] + + +def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__.split("\n\n")[0]) + parser.add_argument("--host", default="127.0.0.1") + parser.add_argument("--port", type=int, default=18305) + parser.add_argument( + "--model", + default=None, + help="Served model name; defaults to the first one in /v1/models.", + ) + parser.add_argument( + "--dp-size", + type=int, + default=2, + help="Attention data-parallel ranks; each is warmed separately.", + ) + parser.add_argument( + "--num-prompts", + type=int, + required=True, + help="Pool prompts per Attention DP rank; must match the server's " + "kv_connector_extra_config num_prompts.", + ) + parser.add_argument( + "--num-requests", + type=int, + default=1000, + help="Measured requests.", + ) + parser.add_argument("--output-len", type=int, default=128) + parser.add_argument("--max-concurrency", type=int, default=64) + parser.add_argument( + "--request-rate", + type=float, + default=math.inf, + help="Measured requests per second (Poisson); inf sends them as fast " + "as --max-concurrency allows.", + ) + parser.add_argument( + "--warmup-concurrency", + type=int, + default=64, + help="In-flight warmup requests across all ranks.", + ) + parser.add_argument( + "--placeholder-prompt-len", + type=int, + default=8, + help="Tokens in each request's placeholder prompt; the server replaces " + "them with a pool prompt.", + ) + parser.add_argument( + "--pin-measured-ranks", + action="store_true", + help="Spread measured requests round-robin over DP ranks with the " + f"{DATA_PARALLEL_RANK_HEADER} header instead of vLLM's load balancer.", + ) + parser.add_argument( + "--skip-warmup", + action="store_true", + help="Measure only, against a server whose pool is already warm.", + ) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--ready-timeout-s", type=float, default=1800.0) + parser.add_argument("--result-file", type=Path, default=None) + return parser.parse_args(argv) + + +async def wait_for_model( + session: aiohttp.ClientSession, + base_url: str, + timeout_s: float, +) -> str: + """Poll until the server answers /v1/models; return the first model id.""" + deadline = time.monotonic() + timeout_s + while True: + try: + async with session.get(f"{base_url}{MODELS_PATH}") as response: + if response.status == HTTP_OK: + models = await response.json() + return models["data"][0]["id"] + except aiohttp.ClientError: + pass + if time.monotonic() > deadline: + raise TimeoutError(f"{base_url} not ready after {timeout_s:.0f}s") + await asyncio.sleep(READY_POLL_INTERVAL_S) + + +async def read_prefix_cache_counters( + session: aiohttp.ClientSession, + base_url: str, +) -> tuple[float, float]: + """Return the (queries, hits) prefix-cache token counters, summed.""" + async with session.get(f"{base_url}{METRICS_PATH}") as response: + text = await response.text() + totals = {PREFIX_CACHE_QUERIES_METRIC: 0.0, PREFIX_CACHE_HITS_METRIC: 0.0} + for line in text.splitlines(): + if not line or line.startswith("#"): + continue + sample, _, value = line.rpartition(" ") + metric = sample.split("{", 1)[0].removesuffix("_total") + if metric in totals: + totals[metric] += float(value) + return totals[PREFIX_CACHE_QUERIES_METRIC], totals[PREFIX_CACHE_HITS_METRIC] + + +async def send_request( + session: aiohttp.ClientSession, + url: str, + request: PlannedRequest, +) -> RequestResult: + start = time.perf_counter() + first_chunk_time: float | None = None + last_chunk_time = start + itl_s: list[float] = [] + output_tokens = 0 + try: + async with session.post( + url, json=request.payload, headers=request.headers + ) as response: + if response.status != HTTP_OK: + body = await response.text() + return RequestResult(ok=False, error=f"HTTP {response.status}: {body}") + async for raw_line in response.content: + line = raw_line.strip() + if not line.startswith(SSE_DATA_PREFIX): + continue + data = line.removeprefix(SSE_DATA_PREFIX) + if data == SSE_DONE: + break + chunk = json.loads(data) + usage = chunk.get("usage") + if usage: + output_tokens = usage["completion_tokens"] + if not chunk.get("choices"): + continue + now = time.perf_counter() + if first_chunk_time is None: + first_chunk_time = now + else: + itl_s.append(now - last_chunk_time) + last_chunk_time = now + except aiohttp.ClientError as error: + return RequestResult(ok=False, error=repr(error)) + if first_chunk_time is None: + return RequestResult(ok=False, error="stream ended without output") + return RequestResult( + ok=True, + ttft_s=first_chunk_time - start, + latency_s=last_chunk_time - start, + output_tokens=output_tokens, + itl_s=itl_s, + ) + + +async def run_phase( + session: aiohttp.ClientSession, + url: str, + requests: list[PlannedRequest], + max_concurrency: int, + request_rate: float, + rng: random.Random, +) -> tuple[list[RequestResult], float]: + """Send ``requests``; return their results and the phase duration.""" + semaphore = asyncio.Semaphore(max_concurrency) + + async def limited(request: PlannedRequest) -> RequestResult: + async with semaphore: + return await send_request(session, url, request) + + start = time.perf_counter() + tasks = [] + for request in requests: + tasks.append(asyncio.create_task(limited(request))) + if not math.isinf(request_rate): + await asyncio.sleep(rng.expovariate(request_rate)) + results = await asyncio.gather(*tasks) + return list(results), time.perf_counter() - start + + +def build_payload( + model: str, + prompt_token_ids: list[int], + max_tokens: int, +) -> dict[str, object]: + return { + "model": model, + "prompt": prompt_token_ids, + "max_tokens": max_tokens, + "ignore_eos": True, + "temperature": 0.0, + "stream": True, + # Per-chunk usage gives the exact output token count even when one + # chunk carries several tokens. + "stream_options": {"include_usage": True, "continuous_usage_stats": True}, + } + + +def placeholder_prompt(rng: random.Random, length: int) -> list[int]: + return [rng.randrange(PLACEHOLDER_TOKEN_ID_LIMIT) for _ in range(length)] + + +def plan_warmup( + args: argparse.Namespace, model: str, rng: random.Random +) -> list[PlannedRequest]: + # Interleave ranks so every rank warms its pool concurrently. + return [ + PlannedRequest( + payload=build_payload( + model, + placeholder_prompt(rng, args.placeholder_prompt_len), + WARMUP_OUTPUT_TOKENS, + ), + headers={DATA_PARALLEL_RANK_HEADER: str(rank)}, + ) + for _ in range(args.num_prompts) + for rank in range(args.dp_size) + ] + + +def plan_measured( + args: argparse.Namespace, model: str, rng: random.Random +) -> list[PlannedRequest]: + return [ + PlannedRequest( + payload=build_payload( + model, + placeholder_prompt(rng, args.placeholder_prompt_len), + args.output_len, + ), + headers=( + {DATA_PARALLEL_RANK_HEADER: str(index % args.dp_size)} + if args.pin_measured_ranks + else {} + ), + ) + for index in range(args.num_requests) + ] + + +def distribution_ms(values_s: list[float]) -> dict[str, float]: + if not values_s: + return {} + values_ms = np.asarray(values_s) * MS_PER_S + summary = {"mean": float(values_ms.mean())} + for percentile in PERCENTILES: + summary[f"p{percentile}"] = float(np.percentile(values_ms, percentile)) + return summary + + +def summarize(results: list[RequestResult], duration_s: float) -> dict[str, object]: + completed = [result for result in results if result.ok] + output_tokens = sum(result.output_tokens for result in completed) + tpot_s = [ + (result.latency_s - result.ttft_s) / (result.output_tokens - 1) + for result in completed + if result.output_tokens > 1 + ] + return { + "completed": len(completed), + "failed": len(results) - len(completed), + "duration_s": duration_s, + "request_throughput": len(completed) / duration_s, + "output_token_throughput": output_tokens / duration_s, + "total_output_tokens": output_tokens, + "ttft_ms": distribution_ms([result.ttft_s for result in completed]), + "tpot_ms": distribution_ms(tpot_s), + "itl_ms": distribution_ms( + [itl for result in completed for itl in result.itl_s] + ), + "e2el_ms": distribution_ms([result.latency_s for result in completed]), + } + + +def first_error(results: list[RequestResult]) -> str: + return next((result.error for result in results if not result.ok), "") + + +async def run(args: argparse.Namespace) -> dict[str, object]: + base_url = f"http://{args.host}:{args.port}" + url = f"{base_url}{COMPLETIONS_PATH}" + rng = random.Random(args.seed) + timeout = aiohttp.ClientTimeout(total=None) + connector = aiohttp.TCPConnector(limit=0) + async with aiohttp.ClientSession(timeout=timeout, connector=connector) as session: + served_model = await wait_for_model(session, base_url, args.ready_timeout_s) + model = args.model or served_model + report: dict[str, object] = {"config": vars(args) | {"model": model}} + + if not args.skip_warmup: + warmup = plan_warmup(args, model, rng) + print( + f"Warmup: {args.num_prompts} prompts x {args.dp_size} ranks", + flush=True, + ) + warmup_results, warmup_duration_s = await run_phase( + session, url, warmup, args.warmup_concurrency, math.inf, rng + ) + num_failed = sum(not result.ok for result in warmup_results) + if num_failed: + raise RuntimeError( + f"{num_failed} warmup requests failed; first error: " + f"{first_error(warmup_results)}" + ) + report["warmup"] = { + "requests": len(warmup_results), + "duration_s": warmup_duration_s, + } + print(f"Warmup done in {warmup_duration_s:.1f}s", flush=True) + + queries_before, hits_before = await read_prefix_cache_counters( + session, base_url + ) + print(f"Measured: {args.num_requests} requests", flush=True) + measured_results, measured_duration_s = await run_phase( + session, + url, + plan_measured(args, model, rng), + args.max_concurrency, + args.request_rate, + rng, + ) + queries_after, hits_after = await read_prefix_cache_counters(session, base_url) + + measured = summarize(measured_results, measured_duration_s) + if measured["failed"]: + measured["first_error"] = first_error(measured_results) + num_completed = sum(result.ok for result in measured_results) + prompt_tokens = queries_after - queries_before + hit_tokens = hits_after - hits_before + measured["prefix_cache"] = { + "queried_tokens": prompt_tokens, + "hit_tokens": hit_tokens, + "hit_rate": hit_tokens / prompt_tokens if prompt_tokens else 0.0, + # 1.0 when every measured request prefilled only its last token. + "prefilled_tokens_per_request": ( + (prompt_tokens - hit_tokens) / num_completed if num_completed else 0.0 + ), + } + report["measured"] = measured + return report + + +def print_summary(measured: dict[str, object]) -> None: + prefix_cache = measured["prefix_cache"] + assert isinstance(prefix_cache, dict) + print( + f"completed={measured['completed']} failed={measured['failed']} " + f"duration={measured['duration_s']:.1f}s" + ) + print( + f"request throughput={measured['request_throughput']:.2f} req/s " + f"output throughput={measured['output_token_throughput']:.1f} tok/s" + ) + for metric in ("ttft_ms", "tpot_ms", "itl_ms", "e2el_ms"): + values = measured[metric] + assert isinstance(values, dict) + formatted = " ".join(f"{key}={value:.2f}" for key, value in values.items()) + print(f"{metric}: {formatted}") + print( + f"prefix cache hit rate={prefix_cache['hit_rate']:.4f} " + f"prefilled tokens/request={prefix_cache['prefilled_tokens_per_request']:.2f} " + "(1.00 = decode only)" + ) + + +def main(argv: Sequence[str] | None = None) -> int: + args = parse_args(argv) + report = asyncio.run(run(args)) + measured = report["measured"] + assert isinstance(measured, dict) + print_summary(measured) + if args.result_file is not None: + args.result_file.parent.mkdir(parents=True, exist_ok=True) + args.result_file.write_text(json.dumps(report, indent=2, default=str)) + print(f"Saved {args.result_file}") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tools/benchmarks/decode_bench_server.sh b/tools/benchmarks/decode_bench_server.sh deleted file mode 100755 index 683b0aaab..000000000 --- a/tools/benchmarks/decode_bench_server.sh +++ /dev/null @@ -1,104 +0,0 @@ -#!/usr/bin/env bash -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project -# -# AF-disaggregated (2A2F) decode-only benchmark. -# -# This is the DECODE side of prefill-decode (PD) disaggregation with the prefill -# instance faked out. Compare recipe/gpu/P2pNcclAFDConnector/deepseek_v2_lite/ -# prefill_decode_disaggregation/2p1a1f_graph_dbo.sh, where the decode side -# (attention + ffn) carries an LMCacheConnectorV1 kv_consumer and pulls real KV -# from a separate 1P prefill instance via a proxy. -# -# Here the AFDDecodeBenchConnector REPLACES that consumer connector on the -# ATTENTION instance and fabricates the KV: every prompt token except the last -# is reported as externally computed, and the KV cache is filled with dummy -# values. So requests skip prefill and go straight into decode -- letting you -# stress the decode path with arbitrary ISL, with NO prefill instance, no -# LMCache producer, and no proxy. Throughput/latency are meaningful; generated -# text is garbage. The connector is attached to attention only (the ffn side -# needs no kv-transfer-config), matching tests/e2e/runner.py. -# -# The connector lives in tools/benchmarks/ and is NOT shipped in the wheel, so -# it is loaded purely via kv_connector_module_path and needs the repo root on -# PYTHONPATH (inherited by the vLLM scheduler/worker subprocesses). -# -# Usage: -# MODEL_PATH=/path/to/DeepSeek-V2-Lite tools/benchmarks/2a2f_graph_dbo_dp2tp1_decode_bench.sh -# then, once both instances are ready, drive load against the attention port: -# MODEL_PATH=/path/to/DeepSeek-V2-Lite tools/benchmarks/vllm_bench.sh -set -euo pipefail - -# --- make tools.benchmarks.decode_bench importable in vLLM subprocesses ------- -SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" -REPO_ROOT="$(cd "${SCRIPT_DIR}/../.." && pwd)" -export PYTHONPATH="${REPO_ROOT}${PYTHONPATH:+:${PYTHONPATH}}" - -MODEL_PATH=${MODEL_PATH:-/path/model_weights/DeepSeek-V2-Lite} - -# dummy-KV fill params for the decode-bench connector -FILL_MEAN=${FILL_MEAN:-0.015} -FILL_STD=${FILL_STD:-0.0} -DECODE_BENCH_KV_CONFIG=$(cat < attn.log 2>&1 & - -# --- ffn instance (no decode-bench connector) --------------------------------- -CUDA_VISIBLE_DEVICES=2,3 uv run vllm serve "$MODEL_PATH" \ - --data-parallel-size 2 \ - --tensor-parallel-size 1 \ - --enable-expert-parallel \ - --additional-config '{ - "afd": { - "role": "ffn", - "connector": "P2pNcclAFDConnector", - "host": "127.0.0.1", - "port": 6269, - "num_attention_ranks": 2, - "num_ffn_ranks": 2 - } - }' \ - --max-num-seqs 64 \ - --enable-dbo \ - --dbo-decode-token-threshold 2 \ - --dbo-prefill-token-threshold 12 \ - --max-num-batched-tokens 64 \ - --max-cudagraph-capture-size 64 \ - --compilation-config '{ - "cudagraph_mode": "FULL_DECODE_ONLY", "cudagraph_capture_sizes":[64] - }' \ - --host 127.0.0.1 \ - --port 18305 \ - --trust-remote-code > ffn.log 2>&1 & - -wait diff --git a/tools/benchmarks/request_generator.sh b/tools/benchmarks/request_generator.sh index a60dbbf7b..0a7340bce 100755 --- a/tools/benchmarks/request_generator.sh +++ b/tools/benchmarks/request_generator.sh @@ -2,9 +2,10 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project # -# Drive load against a running AFD instance (e.g. the attention endpoint started -# by decode_bench_server.sh). Thin wrapper around `vllm bench serve`; every knob -# is overridable via an environment variable. +# Drive load against a running AFD instance (e.g. the attention endpoint of a +# recipe under recipe/gpu/). Thin wrapper around `vllm bench serve`; every knob +# is overridable via an environment variable. For a decode-only run over a +# warmed prompt pool, use decode_bench_client.py instead. # # Usage: # MODEL_PATH=/path/to/weights tools/benchmarks/request_generator.sh From 4fa8cf25c3bb7e5ded41642ce16e742de38f739a Mon Sep 17 00:00:00 2001 From: ronenkat <16743404+ronenkat@users.noreply.github.com> Date: Wed, 28 Oct 2026 21:35:48 +0200 Subject: [PATCH 17/32] feat(benchmarks): run the decode-only benchmark phases with vllm bench decode_bench_client.py now drives every phase with `vllm bench serve` and keeps only the orchestration and the decode-only checks: - Warmup: one vllm bench run per Attention DP rank, in parallel, pinned with the X-data-parallel-rank header, sends --pool-prompts requests of --prompt-len tokens with one output token. --warmup-concurrency and --warmup-bench-args configure these runs. - Pre-check: the server's prefix-cache counters must show --prompt-len prompt tokens per warmup request, so the server's pool prompts are as long as the measured prompts, and every rank must have received a full pool of requests. - Measured: vllm bench serve --dataset-name random --random-input-len --random-range-ratio 0, plus every argument after `--` (request rate, concurrency, number of requests, output length, ...). Flags the client owns are rejected after `--`. - Post-check: over the measured run the server must have prefilled one token per finished request, with no failed requests and no preemptions; otherwise the client exits non-zero. vllm bench's result JSON for each run and decode_bench_report.json with the checks land in --result-dir. The client's own HTTP load generator, and its aiohttp and numpy dependencies, are gone. --num-prompts is renamed --pool-prompts so it does not collide with vllm bench's --num-prompts. Validated on 4xH200 with both AFD connectors (1000 x 1025-token pool per rank, 2000 measured requests, 256 output tokens, concurrency 64): pool length pre-check 1025.00 tokens/request, 1.0000 prefilled tokens per measured request, 0 misses, 0 preemptions; a mismatched --prompt-len is rejected before the measured run. Co-Authored-By: Claude Opus 5.5 Signed-off-by: ronenkat <16743404+ronenkat@users.noreply.github.com> --- .../qwen3_30b_a3b/README.md | 39 +- tools/benchmarks/decode_bench_client.py | 712 +++++++++--------- tools/benchmarks/request_generator.sh | 2 +- 3 files changed, 400 insertions(+), 353 deletions(-) diff --git a/recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/README.md b/recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/README.md index 30f7c34d2..885f61ac2 100644 --- a/recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/README.md +++ b/recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/README.md @@ -67,7 +67,8 @@ and `EXTRA_ARGS` for anything else. boundary), so the request prefills one token and decodes over real KV, on whichever rank serves it. -Start the server, then drive it with the two-phase client: +Start the server, then drive it with `tools/benchmarks/decode_bench_client.py`, +which runs every phase with `vllm bench serve`: ```bash export MODEL_PATH=/path/model_weights/Qwen3-30B-A3B @@ -78,18 +79,34 @@ DECODE_BENCH_MAX_OUTPUT_LEN=256 \ bash recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/2a2f_graph_dp2tp1.sh python tools/benchmarks/decode_bench_client.py \ - --port 18305 --dp-size 2 --num-prompts 1000 \ - --num-requests 4000 --output-len 256 --max-concurrency 64 \ - --result-file /tmp/results/decode_bench.json + --port 18305 --dp-size 2 --pool-prompts 1000 --prompt-len 1025 \ + --result-dir /tmp/results/decode_bench \ + -- --num-prompts 4000 --random-output-len 256 \ + --max-concurrency 64 --request-rate 20 ``` -The client warms every rank with `--num-prompts` requests pinned by the -`X-data-parallel-rank` header, then sends the measured requests through vLLM's -load balancer and reports TTFT, TPOT, ITL, and throughput. It also reports -`prefilled tokens/request` from the server's prefix-cache counters: `1.00` -means every measured request was decode-only. `--num-prompts` must match -`DECODE_BENCH_NUM_PROMPTS`, and `--output-len` must not exceed -`DECODE_BENCH_MAX_OUTPUT_LEN`. +1. **Warmup.** One `vllm bench serve` run per Attention rank, pinned with the + `X-data-parallel-rank` header, sends `--pool-prompts` requests to that rank. + `--warmup-concurrency` sets their concurrency; `--warmup-bench-args` (last + before `--`) passes anything else to these runs. +2. **Pre-check.** From the server's prefix-cache counters, every warmup + request must have looked up `--prompt-len` prompt tokens, so the pool + prompts are as long as the measured prompts. +3. **Measured run.** `vllm bench serve --dataset-name random --random-input-len + ` plus every argument after `--`: request rate, concurrency, + number of requests, output length, burstiness, percentiles, and so on. The + client sets the dataset, input length, endpoint, and result file itself and + rejects them after `--`. +4. **Post-check.** Over the measured run the server must have prefilled one + token per finished request (decode only), with no failed requests and no + preemptions. Otherwise the client exits non-zero. + +`--result-dir` receives vllm bench's own result JSON for each run +(`warmup-rank.json`, `measured.json`) plus `decode_bench_report.json` with +the checks. `--pool-prompts` must match `DECODE_BENCH_NUM_PROMPTS`, +`--prompt-len` must match `DECODE_BENCH_PROMPT_LEN`, and `--random-output-len` +must not exceed `DECODE_BENCH_MAX_OUTPUT_LEN`. The client needs vLLM and the +model's tokenizer installed, as vllm bench does. ### Sizing the pool diff --git a/tools/benchmarks/decode_bench_client.py b/tools/benchmarks/decode_bench_client.py index ca71b3c0c..ba93a6881 100644 --- a/tools/benchmarks/decode_bench_client.py +++ b/tools/benchmarks/decode_bench_client.py @@ -1,80 +1,146 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project -"""Two-phase load generator for the decode benchmark's prompt-pool mode. +"""Decode-only benchmark over the decode-bench prompt pool, run with vllm bench. -Run it against an Attention instance whose ``AFDDecodeBenchConnector`` has +Drives an Attention instance whose ``AFDDecodeBenchConnector`` has ``num_prompts`` > 0 (see ``tools/benchmarks/decode_bench.py``): -1. Warmup: sends ``--num-prompts`` requests to every Attention DP rank, pinned - with the ``X-data-parallel-rank`` header, so each rank prefills and pins its - whole prompt pool. Their metrics are discarded. -2. Measured: sends ``--num-requests`` requests through vLLM's load balancer. - The server remaps each onto a pool prompt, so the prompt sent here is only a - placeholder and every request decodes over real, prefix-cached KV. - -The server's prefix-cache counters before and after the measured phase give the -prompt tokens prefilled per measured request; a clean decode-only run prefills -exactly one. +1. Warmup: one ``vllm bench serve`` run per Attention DP rank, each pinned to + its rank with the ``X-data-parallel-rank`` header, sends ``--pool-prompts`` + requests, so every rank prefills and pins its whole prompt pool. +2. Check: the server's prefix-cache counters must show each warmup request + looked up ``--prompt-len`` prompt tokens, i.e. the server's pool prompts are + as long as the measured prompts below. +3. Measured: one ``vllm bench serve`` run with ``--random-input-len + --prompt-len``. Every argument after ``--`` goes to it unchanged (request + rate, concurrency, number of requests, output length, ...). The server remaps + each request onto a pool prompt, so the prompt text is only a placeholder. +4. Check: over the measured run the server must have prefilled exactly one + token per finished request (the last prompt token), with no failed requests + and no preemptions; otherwise this script exits non-zero. + +vllm bench writes its usual result JSON for each run into ``--result-dir``; +this script adds ``decode_bench_report.json`` with the checks. Usage: - python tools/benchmarks/decode_bench_client.py --num-prompts 1000 \ - --num-requests 4000 --output-len 256 --max-concurrency 64 + python tools/benchmarks/decode_bench_client.py --pool-prompts 1000 \ + --prompt-len 1025 -- --num-prompts 4000 --random-output-len 256 \ + --max-concurrency 64 --request-rate 20 """ from __future__ import annotations import argparse -import asyncio import json -import math -import random +import subprocess import sys import time +import urllib.error +import urllib.request from collections.abc import Sequence -from dataclasses import dataclass, field +from dataclasses import asdict, dataclass from pathlib import Path -import aiohttp -import numpy as np - DATA_PARALLEL_RANK_HEADER = "X-data-parallel-rank" -COMPLETIONS_PATH = "/v1/completions" MODELS_PATH = "/v1/models" METRICS_PATH = "/metrics" PREFIX_CACHE_QUERIES_METRIC = "vllm:prefix_cache_queries" PREFIX_CACHE_HITS_METRIC = "vllm:prefix_cache_hits" -SSE_DATA_PREFIX = b"data: " -SSE_DONE = b"[DONE]" +REQUEST_SUCCESS_METRIC = "vllm:request_success" +PREEMPTIONS_METRIC = "vllm:num_preemptions" HTTP_OK = 200 -# Placeholder prompt token ids stay below this, which every vocabulary covers. -PLACEHOLDER_TOKEN_ID_LIMIT = 1000 +READY_POLL_INTERVAL_S = 5.0 # The server forces warmup requests to one output token as well. WARMUP_OUTPUT_TOKENS = 1 -READY_POLL_INTERVAL_S = 5.0 -PERCENTILES = (50, 90, 99) -MS_PER_S = 1000.0 +# A decode-only request prefills only its last prompt token. +DECODE_ONLY_PREFILLED_TOKENS = 1.0 +PREFILLED_TOKENS_TOLERANCE = 0.01 +PROMPT_LEN_TOLERANCE = 0.5 +# vllm bench flags this script sets for every run; overriding them would break +# the prompt-length contract or lose the result file. +CLIENT_OWNED_BENCH_FLAGS = frozenset( + { + "--host", + "--port", + "--base-url", + "--model", + "--tokenizer", + "--dataset-name", + "--random-input-len", + "--random-range-ratio", + "--random-prefix-len", + "--save-result", + "--result-dir", + "--result-filename", + } +) +# The warmup also owns how many requests each rank gets and where they go. +WARMUP_OWNED_BENCH_FLAGS = CLIENT_OWNED_BENCH_FLAGS | { + "--num-prompts", + "--random-output-len", + "--header", + "--max-concurrency", +} +WARMUP_RESULT_FILENAME = "warmup-rank{rank}.json" +WARMUP_LOG_FILENAME = "warmup-rank{rank}.log" +MEASURED_RESULT_FILENAME = "measured.json" +REPORT_FILENAME = "decode_bench_report.json" +MEASURED_SUMMARY_KEYS = ( + "completed", + "failed", + "duration", + "request_throughput", + "output_throughput", + "total_input_tokens", + "total_output_tokens", + "median_ttft_ms", + "p99_ttft_ms", + "median_tpot_ms", + "p99_tpot_ms", + "median_itl_ms", + "p99_itl_ms", +) -@dataclass -class RequestResult: - """Timing of one streamed completion.""" +@dataclass(frozen=True) +class ServerCounters: + """Server-wide totals of the counters the checks need, summed over engines.""" + + prefix_cache_queries: float = 0.0 + prefix_cache_hits: float = 0.0 + finished_requests: float = 0.0 + preemptions: float = 0.0 + + def __sub__(self, other: ServerCounters) -> ServerCounters: + return ServerCounters( + prefix_cache_queries=self.prefix_cache_queries - other.prefix_cache_queries, + prefix_cache_hits=self.prefix_cache_hits - other.prefix_cache_hits, + finished_requests=self.finished_requests - other.finished_requests, + preemptions=self.preemptions - other.preemptions, + ) - ok: bool - ttft_s: float = 0.0 - latency_s: float = 0.0 - output_tokens: int = 0 - itl_s: list[float] = field(default_factory=list) - error: str = "" + @property + def prefilled_tokens(self) -> float: + return self.prefix_cache_queries - self.prefix_cache_hits @dataclass(frozen=True) -class PlannedRequest: - payload: dict[str, object] - headers: dict[str, str] +class Check: + name: str + passed: bool + detail: str + +def parse_args(argv: Sequence[str]) -> tuple[argparse.Namespace, list[str]]: + """Split ``argv`` at ``--``: client flags before, measured vllm bench after.""" + argv = list(argv) + split = argv.index("--") if "--" in argv else len(argv) + client_argv, bench_argv = argv[:split], argv[split + 1 :] -def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: - parser = argparse.ArgumentParser(description=__doc__.split("\n\n")[0]) + parser = argparse.ArgumentParser( + description=__doc__.split("\n\n")[0], + epilog="Arguments after `--` are passed to the measured `vllm bench serve`.", + ) parser.add_argument("--host", default="127.0.0.1") parser.add_argument("--port", type=int, default=18305) parser.add_argument( @@ -82,92 +148,97 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: default=None, help="Served model name; defaults to the first one in /v1/models.", ) + parser.add_argument( + "--tokenizer", + default=None, + help="Tokenizer vllm bench uses to build prompts; defaults to --model.", + ) parser.add_argument( "--dp-size", type=int, default=2, - help="Attention data-parallel ranks; each is warmed separately.", + help="Attention data-parallel ranks; each gets its own warmup run.", ) parser.add_argument( - "--num-prompts", + "--pool-prompts", type=int, required=True, help="Pool prompts per Attention DP rank; must match the server's " "kv_connector_extra_config num_prompts.", ) parser.add_argument( - "--num-requests", + "--prompt-len", type=int, - default=1000, - help="Measured requests.", - ) - parser.add_argument("--output-len", type=int, default=128) - parser.add_argument("--max-concurrency", type=int, default=64) - parser.add_argument( - "--request-rate", - type=float, - default=math.inf, - help="Measured requests per second (Poisson); inf sends them as fast " - "as --max-concurrency allows.", + default=1025, + help="Prompt length of the warmup and measured runs; must equal the " + "server's pool prompt length (checked after the warmup).", ) parser.add_argument( "--warmup-concurrency", type=int, default=64, - help="In-flight warmup requests across all ranks.", - ) - parser.add_argument( - "--placeholder-prompt-len", - type=int, - default=8, - help="Tokens in each request's placeholder prompt; the server replaces " - "them with a pool prompt.", + help="In-flight requests of each rank's warmup run.", ) parser.add_argument( - "--pin-measured-ranks", - action="store_true", - help="Spread measured requests round-robin over DP ranks with the " - f"{DATA_PARALLEL_RANK_HEADER} header instead of vLLM's load balancer.", + "--warmup-bench-args", + nargs=argparse.REMAINDER, + default=[], + help="Extra vllm bench arguments for the warmup runs; must come last " + "before `--`.", ) parser.add_argument( "--skip-warmup", action="store_true", help="Measure only, against a server whose pool is already warm.", ) - parser.add_argument("--seed", type=int, default=0) + parser.add_argument( + "--result-dir", + type=Path, + default=Path("decode_bench_results"), + ) parser.add_argument("--ready-timeout-s", type=float, default=1800.0) - parser.add_argument("--result-file", type=Path, default=None) - return parser.parse_args(argv) + args = parser.parse_args(client_argv) + reject_owned_flags(args.warmup_bench_args, WARMUP_OWNED_BENCH_FLAGS, "warmup") + reject_owned_flags(bench_argv, CLIENT_OWNED_BENCH_FLAGS, "measured") + return args, bench_argv + + +def reject_owned_flags( + bench_args: Sequence[str], owned_flags: frozenset[str], phase: str +) -> None: + for bench_arg in bench_args: + flag = bench_arg.split("=", 1)[0] + if flag in owned_flags: + raise SystemExit( + f"{flag} is set by decode_bench_client.py for the {phase} vllm " + f"bench run; drop it (owned: {', '.join(sorted(owned_flags))})" + ) -async def wait_for_model( - session: aiohttp.ClientSession, - base_url: str, - timeout_s: float, -) -> str: +def wait_for_model(base_url: str, timeout_s: float) -> str: """Poll until the server answers /v1/models; return the first model id.""" deadline = time.monotonic() + timeout_s while True: try: - async with session.get(f"{base_url}{MODELS_PATH}") as response: + with urllib.request.urlopen(f"{base_url}{MODELS_PATH}") as response: if response.status == HTTP_OK: - models = await response.json() - return models["data"][0]["id"] - except aiohttp.ClientError: + return json.loads(response.read())["data"][0]["id"] + except (urllib.error.URLError, ConnectionError): pass if time.monotonic() > deadline: raise TimeoutError(f"{base_url} not ready after {timeout_s:.0f}s") - await asyncio.sleep(READY_POLL_INTERVAL_S) + time.sleep(READY_POLL_INTERVAL_S) -async def read_prefix_cache_counters( - session: aiohttp.ClientSession, - base_url: str, -) -> tuple[float, float]: - """Return the (queries, hits) prefix-cache token counters, summed.""" - async with session.get(f"{base_url}{METRICS_PATH}") as response: - text = await response.text() - totals = {PREFIX_CACHE_QUERIES_METRIC: 0.0, PREFIX_CACHE_HITS_METRIC: 0.0} +def read_counters(base_url: str) -> ServerCounters: + with urllib.request.urlopen(f"{base_url}{METRICS_PATH}") as response: + text = response.read().decode() + totals = { + PREFIX_CACHE_QUERIES_METRIC: 0.0, + PREFIX_CACHE_HITS_METRIC: 0.0, + REQUEST_SUCCESS_METRIC: 0.0, + PREEMPTIONS_METRIC: 0.0, + } for line in text.splitlines(): if not line or line.startswith("#"): continue @@ -175,280 +246,239 @@ async def read_prefix_cache_counters( metric = sample.split("{", 1)[0].removesuffix("_total") if metric in totals: totals[metric] += float(value) - return totals[PREFIX_CACHE_QUERIES_METRIC], totals[PREFIX_CACHE_HITS_METRIC] - - -async def send_request( - session: aiohttp.ClientSession, - url: str, - request: PlannedRequest, -) -> RequestResult: - start = time.perf_counter() - first_chunk_time: float | None = None - last_chunk_time = start - itl_s: list[float] = [] - output_tokens = 0 - try: - async with session.post( - url, json=request.payload, headers=request.headers - ) as response: - if response.status != HTTP_OK: - body = await response.text() - return RequestResult(ok=False, error=f"HTTP {response.status}: {body}") - async for raw_line in response.content: - line = raw_line.strip() - if not line.startswith(SSE_DATA_PREFIX): - continue - data = line.removeprefix(SSE_DATA_PREFIX) - if data == SSE_DONE: - break - chunk = json.loads(data) - usage = chunk.get("usage") - if usage: - output_tokens = usage["completion_tokens"] - if not chunk.get("choices"): - continue - now = time.perf_counter() - if first_chunk_time is None: - first_chunk_time = now - else: - itl_s.append(now - last_chunk_time) - last_chunk_time = now - except aiohttp.ClientError as error: - return RequestResult(ok=False, error=repr(error)) - if first_chunk_time is None: - return RequestResult(ok=False, error="stream ended without output") - return RequestResult( - ok=True, - ttft_s=first_chunk_time - start, - latency_s=last_chunk_time - start, - output_tokens=output_tokens, - itl_s=itl_s, + return ServerCounters( + prefix_cache_queries=totals[PREFIX_CACHE_QUERIES_METRIC], + prefix_cache_hits=totals[PREFIX_CACHE_HITS_METRIC], + finished_requests=totals[REQUEST_SUCCESS_METRIC], + preemptions=totals[PREEMPTIONS_METRIC], ) -async def run_phase( - session: aiohttp.ClientSession, - url: str, - requests: list[PlannedRequest], - max_concurrency: int, - request_rate: float, - rng: random.Random, -) -> tuple[list[RequestResult], float]: - """Send ``requests``; return their results and the phase duration.""" - semaphore = asyncio.Semaphore(max_concurrency) - - async def limited(request: PlannedRequest) -> RequestResult: - async with semaphore: - return await send_request(session, url, request) - - start = time.perf_counter() - tasks = [] - for request in requests: - tasks.append(asyncio.create_task(limited(request))) - if not math.isinf(request_rate): - await asyncio.sleep(rng.expovariate(request_rate)) - results = await asyncio.gather(*tasks) - return list(results), time.perf_counter() - start - - -def build_payload( +def bench_command( + args: argparse.Namespace, model: str, - prompt_token_ids: list[int], - max_tokens: int, -) -> dict[str, object]: - return { - "model": model, - "prompt": prompt_token_ids, - "max_tokens": max_tokens, - "ignore_eos": True, - "temperature": 0.0, - "stream": True, - # Per-chunk usage gives the exact output token count even when one - # chunk carries several tokens. - "stream_options": {"include_usage": True, "continuous_usage_stats": True}, - } - - -def placeholder_prompt(rng: random.Random, length: int) -> list[int]: - return [rng.randrange(PLACEHOLDER_TOKEN_ID_LIMIT) for _ in range(length)] - - -def plan_warmup( - args: argparse.Namespace, model: str, rng: random.Random -) -> list[PlannedRequest]: - # Interleave ranks so every rank warms its pool concurrently. + result_filename: str, + bench_args: Sequence[str], +) -> list[str]: + """A ``vllm bench serve`` command over --prompt-len random prompts.""" return [ - PlannedRequest( - payload=build_payload( - model, - placeholder_prompt(rng, args.placeholder_prompt_len), - WARMUP_OUTPUT_TOKENS, - ), - headers={DATA_PARALLEL_RANK_HEADER: str(rank)}, - ) - for _ in range(args.num_prompts) - for rank in range(args.dp_size) + sys.executable, + "-m", + "vllm.entrypoints.cli.main", + "bench", + "serve", + "--host", + args.host, + "--port", + str(args.port), + "--model", + model, + "--tokenizer", + args.tokenizer or model, + "--dataset-name", + "random", + "--random-input-len", + str(args.prompt_len), + "--random-range-ratio", + "0", + "--random-prefix-len", + "0", + "--ignore-eos", + "--save-result", + "--result-dir", + str(args.result_dir), + "--result-filename", + result_filename, + *bench_args, ] -def plan_measured( - args: argparse.Namespace, model: str, rng: random.Random -) -> list[PlannedRequest]: - return [ - PlannedRequest( - payload=build_payload( - model, - placeholder_prompt(rng, args.placeholder_prompt_len), - args.output_len, - ), - headers=( - {DATA_PARALLEL_RANK_HEADER: str(index % args.dp_size)} - if args.pin_measured_ranks - else {} - ), +def read_bench_result(args: argparse.Namespace, filename: str) -> dict[str, object]: + return json.loads((args.result_dir / filename).read_text()) + + +def run_warmup(args: argparse.Namespace, model: str) -> list[dict[str, object]]: + """Run one pinned vllm bench per rank in parallel; return their results.""" + processes = [] + for rank in range(args.dp_size): + command = bench_command( + args, + model, + WARMUP_RESULT_FILENAME.format(rank=rank), + [ + "--num-prompts", + str(args.pool_prompts), + "--random-output-len", + str(WARMUP_OUTPUT_TOKENS), + "--max-concurrency", + str(args.warmup_concurrency), + "--header", + f"{DATA_PARALLEL_RANK_HEADER}={rank}", + *args.warmup_bench_args, + ], + ) + log_path = args.result_dir / WARMUP_LOG_FILENAME.format(rank=rank) + with log_path.open("w") as log: + processes.append( + (rank, log_path, subprocess.Popen(command, stdout=log, stderr=log)) + ) + results = [] + for rank, log_path, process in processes: + if process.wait() != 0: + raise RuntimeError( + f"warmup vllm bench for rank {rank} exited {process.returncode}; " + f"see {log_path}" + ) + results.append( + read_bench_result(args, WARMUP_RESULT_FILENAME.format(rank=rank)) + ) + return results + + +def check_warmup( + args: argparse.Namespace, + results: list[dict[str, object]], + delta: ServerCounters, +) -> list[Check]: + num_failed = sum(int(str(result["failed"])) for result in results) + checks = [ + Check( + "warmup requests succeeded", + num_failed == 0, + f"{num_failed} failed; a failed warmup request leaves its pool entry " + "for a later request to prefill", ) - for index in range(args.num_requests) ] + min_requests = args.dp_size * args.pool_prompts + checks.append( + Check( + "every rank received a full pool of warmup requests", + delta.finished_requests >= min_requests, + f"{delta.finished_requests:.0f} finished, need >= {min_requests}", + ) + ) + tokens_per_request = ( + delta.prefix_cache_queries / delta.finished_requests + if delta.finished_requests + else 0.0 + ) + prompt_len_matches = ( + abs(tokens_per_request - args.prompt_len) <= PROMPT_LEN_TOLERANCE + ) + checks.append( + Check( + "server pool prompts match --prompt-len", + prompt_len_matches, + f"{tokens_per_request:.2f} prompt tokens per warmup request, " + f"--prompt-len {args.prompt_len}" + + ( + "" + if prompt_len_matches + else "; set the server's prompt_len to one token past a block " + "boundary and pass the same value here" + ), + ) + ) + return checks -def distribution_ms(values_s: list[float]) -> dict[str, float]: - if not values_s: - return {} - values_ms = np.asarray(values_s) * MS_PER_S - summary = {"mean": float(values_ms.mean())} - for percentile in PERCENTILES: - summary[f"p{percentile}"] = float(np.percentile(values_ms, percentile)) - return summary - - -def summarize(results: list[RequestResult], duration_s: float) -> dict[str, object]: - completed = [result for result in results if result.ok] - output_tokens = sum(result.output_tokens for result in completed) - tpot_s = [ - (result.latency_s - result.ttft_s) / (result.output_tokens - 1) - for result in completed - if result.output_tokens > 1 - ] - return { - "completed": len(completed), - "failed": len(results) - len(completed), - "duration_s": duration_s, - "request_throughput": len(completed) / duration_s, - "output_token_throughput": output_tokens / duration_s, - "total_output_tokens": output_tokens, - "ttft_ms": distribution_ms([result.ttft_s for result in completed]), - "tpot_ms": distribution_ms(tpot_s), - "itl_ms": distribution_ms( - [itl for result in completed for itl in result.itl_s] +def check_measured(result: dict[str, object], delta: ServerCounters) -> list[Check]: + prefilled_per_request = ( + delta.prefilled_tokens / delta.finished_requests + if delta.finished_requests + else 0.0 + ) + return [ + Check( + "measured requests succeeded", + int(str(result["failed"])) == 0, + f"{result['failed']} failed of {result['completed']} completed", ), - "e2el_ms": distribution_ms([result.latency_s for result in completed]), - } + Check( + "decode only: one prefilled token per request", + abs(prefilled_per_request - DECODE_ONLY_PREFILLED_TOKENS) + <= PREFILLED_TOKENS_TOLERANCE, + f"{prefilled_per_request:.4f} prefilled tokens per finished request " + f"({delta.prefilled_tokens:.0f} tokens, " + f"{delta.finished_requests:.0f} requests)", + ), + Check( + "no preemptions", + delta.preemptions == 0, + f"{delta.preemptions:.0f} preemptions; a preempted request recomputes " + "its context", + ), + ] -def first_error(results: list[RequestResult]) -> str: - return next((result.error for result in results if not result.ok), "") +def print_checks(phase: str, checks: list[Check]) -> None: + for check in checks: + verdict = "PASS" if check.passed else "FAIL" + print(f"[{phase}] {verdict} {check.name}: {check.detail}", flush=True) -async def run(args: argparse.Namespace) -> dict[str, object]: +def main(argv: Sequence[str] | None = None) -> int: + args, measured_bench_args = parse_args(sys.argv[1:] if argv is None else argv) + args.result_dir.mkdir(parents=True, exist_ok=True) base_url = f"http://{args.host}:{args.port}" - url = f"{base_url}{COMPLETIONS_PATH}" - rng = random.Random(args.seed) - timeout = aiohttp.ClientTimeout(total=None) - connector = aiohttp.TCPConnector(limit=0) - async with aiohttp.ClientSession(timeout=timeout, connector=connector) as session: - served_model = await wait_for_model(session, base_url, args.ready_timeout_s) - model = args.model or served_model - report: dict[str, object] = {"config": vars(args) | {"model": model}} - - if not args.skip_warmup: - warmup = plan_warmup(args, model, rng) - print( - f"Warmup: {args.num_prompts} prompts x {args.dp_size} ranks", - flush=True, - ) - warmup_results, warmup_duration_s = await run_phase( - session, url, warmup, args.warmup_concurrency, math.inf, rng - ) - num_failed = sum(not result.ok for result in warmup_results) - if num_failed: - raise RuntimeError( - f"{num_failed} warmup requests failed; first error: " - f"{first_error(warmup_results)}" - ) - report["warmup"] = { - "requests": len(warmup_results), - "duration_s": warmup_duration_s, - } - print(f"Warmup done in {warmup_duration_s:.1f}s", flush=True) - - queries_before, hits_before = await read_prefix_cache_counters( - session, base_url - ) - print(f"Measured: {args.num_requests} requests", flush=True) - measured_results, measured_duration_s = await run_phase( - session, - url, - plan_measured(args, model, rng), - args.max_concurrency, - args.request_rate, - rng, + served_model = wait_for_model(base_url, args.ready_timeout_s) + model = args.model or served_model + report: dict[str, object] = { + "config": {key: str(value) for key, value in vars(args).items()} + | {"model": model, "measured_bench_args": measured_bench_args}, + } + checks: list[Check] = [] + + if not args.skip_warmup: + print( + f"Warmup: vllm bench x {args.dp_size} ranks, {args.pool_prompts} " + f"prompts of {args.prompt_len} tokens each", + flush=True, ) - queries_after, hits_after = await read_prefix_cache_counters(session, base_url) - - measured = summarize(measured_results, measured_duration_s) - if measured["failed"]: - measured["first_error"] = first_error(measured_results) - num_completed = sum(result.ok for result in measured_results) - prompt_tokens = queries_after - queries_before - hit_tokens = hits_after - hits_before - measured["prefix_cache"] = { - "queried_tokens": prompt_tokens, - "hit_tokens": hit_tokens, - "hit_rate": hit_tokens / prompt_tokens if prompt_tokens else 0.0, - # 1.0 when every measured request prefilled only its last token. - "prefilled_tokens_per_request": ( - (prompt_tokens - hit_tokens) / num_completed if num_completed else 0.0 - ), + before = read_counters(base_url) + warmup_results = run_warmup(args, model) + warmup_delta = read_counters(base_url) - before + warmup_checks = check_warmup(args, warmup_results, warmup_delta) + print_checks("warmup", warmup_checks) + checks += warmup_checks + report["warmup"] = { + "server_counters": asdict(warmup_delta), + "checks": [asdict(check) for check in warmup_checks], } - report["measured"] = measured - return report - - -def print_summary(measured: dict[str, object]) -> None: - prefix_cache = measured["prefix_cache"] - assert isinstance(prefix_cache, dict) - print( - f"completed={measured['completed']} failed={measured['failed']} " - f"duration={measured['duration_s']:.1f}s" - ) - print( - f"request throughput={measured['request_throughput']:.2f} req/s " - f"output throughput={measured['output_token_throughput']:.1f} tok/s" - ) - for metric in ("ttft_ms", "tpot_ms", "itl_ms", "e2el_ms"): - values = measured[metric] - assert isinstance(values, dict) - formatted = " ".join(f"{key}={value:.2f}" for key, value in values.items()) - print(f"{metric}: {formatted}") - print( - f"prefix cache hit rate={prefix_cache['hit_rate']:.4f} " - f"prefilled tokens/request={prefix_cache['prefilled_tokens_per_request']:.2f} " - "(1.00 = decode only)" + if not all(check.passed for check in warmup_checks): + write_report(args, report, checks) + return 1 + + print(f"Measured: vllm bench serve {' '.join(measured_bench_args)}", flush=True) + before = read_counters(base_url) + completed = subprocess.run( + bench_command(args, model, MEASURED_RESULT_FILENAME, measured_bench_args) ) + if completed.returncode != 0: + print(f"measured vllm bench exited {completed.returncode}", flush=True) + return completed.returncode + measured_delta = read_counters(base_url) - before + measured_result = read_bench_result(args, MEASURED_RESULT_FILENAME) + measured_checks = check_measured(measured_result, measured_delta) + print_checks("measured", measured_checks) + checks += measured_checks + report["measured"] = { + "vllm_bench_result": str(args.result_dir / MEASURED_RESULT_FILENAME), + "summary": {key: measured_result.get(key) for key in MEASURED_SUMMARY_KEYS}, + "server_counters": asdict(measured_delta), + "checks": [asdict(check) for check in measured_checks], + } + write_report(args, report, checks) + return 0 if all(check.passed for check in checks) else 1 -def main(argv: Sequence[str] | None = None) -> int: - args = parse_args(argv) - report = asyncio.run(run(args)) - measured = report["measured"] - assert isinstance(measured, dict) - print_summary(measured) - if args.result_file is not None: - args.result_file.parent.mkdir(parents=True, exist_ok=True) - args.result_file.write_text(json.dumps(report, indent=2, default=str)) - print(f"Saved {args.result_file}") - return 0 +def write_report( + args: argparse.Namespace, report: dict[str, object], checks: list[Check] +) -> None: + report["passed"] = all(check.passed for check in checks) + path = args.result_dir / REPORT_FILENAME + path.write_text(json.dumps(report, indent=2)) + print(f"{'PASS' if report['passed'] else 'FAIL'}: report saved to {path}") if __name__ == "__main__": diff --git a/tools/benchmarks/request_generator.sh b/tools/benchmarks/request_generator.sh index 0a7340bce..c0240c277 100755 --- a/tools/benchmarks/request_generator.sh +++ b/tools/benchmarks/request_generator.sh @@ -5,7 +5,7 @@ # Drive load against a running AFD instance (e.g. the attention endpoint of a # recipe under recipe/gpu/). Thin wrapper around `vllm bench serve`; every knob # is overridable via an environment variable. For a decode-only run over a -# warmed prompt pool, use decode_bench_client.py instead. +# warmed prompt pool, use decode_bench_client.py, which wraps vllm bench too. # # Usage: # MODEL_PATH=/path/to/weights tools/benchmarks/request_generator.sh From b7aca181e29137a4d1a55108ae5e016a14c64527 Mon Sep 17 00:00:00 2001 From: ronenkat <16743404+ronenkat@users.noreply.github.com> Date: Thu, 29 Oct 2026 15:03:04 +0200 Subject: [PATCH 18/32] feat(recipe): Qwen3-235B-A22B-FP8 4A2F recipe with decode-only benchmark One launch script for Qwen3-235B-A22B-FP8 in a 4A2F prefill/decode colocation on a single 6-GPU node: Attention DP4/TP1 on GPUs 0-3 and FFN DP2/TP1 with expert parallelism on GPUs 4-5, FULL_DECODE_ONLY CUDA graphs (8-128). The connector is selectable (DeepEPM2NAFDConnector by default, P2pNcclAFDConnector through its own entry point), and KV_CONNECTOR=decode_bench turns it into the decode-only benchmark over a pinned real-KV prompt pool, driven by decode_bench_client.py with --dp-size 4. Differences from the Qwen3-30B-A3B recipe: - DeepEPM2NAFDConnector dispatches FP8 (M2N_PAYLOAD_FORMAT, default fp8_block128_fp32); M2N_EXTRA_CONFIG adds connector_extra_config members. - Per-role ATTN_EXTRA_ARGS / FFN_EXTRA_ARGS, e.g. --api-server-count 1 --cudagraph-metrics on Attention. - The AFD join timeout and vLLM's CPU-group timeout default to 2 h, since a cold 235B load can finish minutes apart on the two roles. - Defaults sized for the benchmark: 128 sequences per Attention rank, MBT 3072, max-model-len 4096, a 500 x 1025-token pool per rank (about 555K of the ~640K KV tokens an H200 Attention rank has). - The README warns that P2pNcclAFDConnector's FFN hangs with PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True. Validated on 6xH200 with both connectors (M2N + DBO at concurrency 64 and 128): decode-only sweeps at concurrency 64/128/256 and request rates 1-50 (rate x 300 requests), every step one prefilled token per request and no preemptions. Co-Authored-By: Claude Opus 5.5 Signed-off-by: ronenkat <16743404+ronenkat@users.noreply.github.com> --- recipe/README.md | 3 + .../qwen3_235b_a22b_fp8/README.md | 90 +++++++++ .../4a2f_graph_dp4tp1.sh | 182 ++++++++++++++++++ .../qwen3_235b_a22b_fp8/README.md | 15 ++ .../4a2f_graph_dp4tp1.sh | 12 ++ 5 files changed, 302 insertions(+) create mode 100644 recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/README.md create mode 100755 recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/4a2f_graph_dp4tp1.sh create mode 100644 recipe/gpu/P2pNcclAFDConnector/qwen3_235b_a22b_fp8/README.md create mode 100755 recipe/gpu/P2pNcclAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/4a2f_graph_dp4tp1.sh diff --git a/recipe/README.md b/recipe/README.md index 237a6f624..509566b89 100644 --- a/recipe/README.md +++ b/recipe/README.md @@ -13,9 +13,11 @@ Recipes are organized by hardware backend, connector, and model: recipe/ ├── gpu/ │ ├── DeepEPM2NAFDConnector/ +│ │ ├── qwen3_235b_a22b_fp8/ │ │ └── qwen3_30b_a3b/ │ └── P2pNcclAFDConnector/ │ ├── deepseek_v2_lite/ +│ ├── qwen3_235b_a22b_fp8/ │ └── qwen3_30b_a3b/ └── npu/ ├── CAMAsyncAFDConnector/ @@ -38,6 +40,7 @@ Directory names follow these conventions: | --- | --- | --- | --- | --- | --- | | GPU | `P2pNcclAFDConnector` | DeepSeek-V2-Lite | Decode | Legacy 2A2F E2E validated; examples inspected | [Launch examples](gpu/P2pNcclAFDConnector/deepseek_v2_lite/README.md) | | GPU | `DeepEPM2NAFDConnector` or `P2pNcclAFDConnector` | Qwen3-30B-A3B | Prefill / decode; decode-only benchmark | 4xH200: serving smoke and decode-only benchmark validated on both connectors | [2A2F single node](gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/README.md) | +| GPU | `DeepEPM2NAFDConnector` or `P2pNcclAFDConnector` | Qwen3-235B-A22B-FP8 | Prefill / decode; decode-only benchmark | 6xH200 decode-only benchmark | [4A2F single node](gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/README.md) | | Ascend NPU | `CAMP2pAFDConnector` | DeepSeek-V3.2 | Decode | Integrated v0.30 hardware validation pending; linked recipe is historical | [Synchronous decode](npu/CAMP2pAFDConnector/deepseek_v3_2/README.md) | | Ascend NPU | `CAMAsyncAFDConnector` | DeepSeek-V3.2 | Prefill / decode | Integrated v0.30 hardware validation pending; historical v0.26 DP2TP8+EP16 token split reached `0.9522` strict match; legacy PCP8 results are v0.19-only | [Async CAM](npu/CAMAsyncAFDConnector/deepseek_v3_2/README.md) | diff --git a/recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/README.md b/recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/README.md new file mode 100644 index 000000000..75f70577c --- /dev/null +++ b/recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/README.md @@ -0,0 +1,90 @@ +# Qwen3-235B-A22B-FP8 AFD Recipe + +One launch script for Qwen3-235B-A22B-FP8 in a 4A2F prefill/decode colocation +on a single node with 6 GPUs, with a selectable AFD connector and an optional +KV connector. With the decode-bench KV connector it doubles as a decode-only +benchmark over real KV, exactly like the +[Qwen3-30B-A3B recipe](../qwen3_30b_a3b/README.md#decode-only-benchmark). + +## Prerequisites + +- 6 GPUs with about 140 GB each on one node (H200). `DeepEPM2NAFDConnector` + needs SM90 or newer and DeepEP/DeepGEMM installed; `P2pNcclAFDConnector` + runs on any A/H-class GPU with enough memory. +- vLLM `v0.30.0` and the `afd-plugin` package in the same environment. +- Qwen3-235B-A22B-FP8 weights (block-128 FP8). The script defaults to + `/path/model_weights/Qwen3-235B-A22B-FP8`; override with `MODEL_PATH=...`. +- Free ports: `6269` for the AFD connector, `18305`/`18306` for the Attention + and FFN HTTP servers. + +## Topology + +| GPUs | Role | Parallelism | Port | +|------------|-----------|----------------------|-------| +| 0, 1, 2, 3 | Attention | DP4 / TP1 | 18305 | +| 4, 5 | FFN | DP2 / TP1, EP2 | 18306 | + +Both roles run `FULL_DECODE_ONLY` CUDA graphs (sizes 8 to 128 by default); +prefill runs eager in chunks of `MAX_NUM_BATCHED_TOKENS`. Send traffic to the +Attention port. `DeepEPM2NAFDConnector` routes on Attention and dispatches FP8 +tokens (`M2N_PAYLOAD_FORMAT=fp8_block128_fp32`); `P2pNcclAFDConnector` sends +BF16 hidden states and routes on the FFN. + +## Running + +Run from the repository root: + +```bash +export MODEL_PATH=/path/model_weights/Qwen3-235B-A22B-FP8 +# DeepEP M2N (default) +bash recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/4a2f_graph_dp4tp1.sh +# NCCL P2P: same script with AFD_CONNECTOR=P2pNcclAFDConnector, or its entry point +bash recipe/gpu/P2pNcclAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/4a2f_graph_dp4tp1.sh +``` + +The script backgrounds both roles, logs to `attn.log` and `ffn.log` in the +current directory, and waits. Wait for `Application startup complete` in +`attn.log` before sending traffic. A cold weight load can take tens of minutes +and the two roles may finish minutes apart, so the recipe raises the AFD join +timeout (`AFD_PROCESS_GROUP_TIMEOUT_S`) and vLLM's CPU-group timeout to 2 h. + +Every knob is an environment variable documented at the top of the script, +including DBO (`ENABLE_DBO=1`), `MAX_NUM_SEQS` per Attention rank, graph +capture sizes, devices, ports, `EXTRA_ARGS` for both roles, and +`ATTN_EXTRA_ARGS` / `FFN_EXTRA_ARGS` for one role. For example, +`ATTN_EXTRA_ARGS="--api-server-count 1 --cudagraph-metrics"` logs how many +Attention steps ran in each CUDA graph and how many ran eager (vLLM disables +the stats logger with more than one API server, and DP4 starts four). + +Do not set `PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True` for the +`P2pNcclAFDConnector` FFN: it hangs on its first step. + +## Decode-only benchmark + +`KV_CONNECTOR=decode_bench` attaches `AFDDecodeBenchConnector` to Attention; +the pool, pinning, and remapping work as described for +[Qwen3-30B-A3B](../qwen3_30b_a3b/README.md#decode-only-benchmark). Drive it +with `--dp-size 4`: + +```bash +export MODEL_PATH=/path/model_weights/Qwen3-235B-A22B-FP8 +KV_CONNECTOR=decode_bench \ +DECODE_BENCH_NUM_PROMPTS=500 \ +DECODE_BENCH_PROMPT_LEN=1025 \ +DECODE_BENCH_MAX_OUTPUT_LEN=256 \ + bash recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/4a2f_graph_dp4tp1.sh + +python tools/benchmarks/decode_bench_client.py \ + --port 18305 --dp-size 4 --pool-prompts 500 --prompt-len 1025 \ + --result-dir /tmp/results/decode_bench \ + -- --num-prompts 3000 --random-output-len 256 \ + --max-concurrency 128 --request-rate 10 +``` + +### Sizing the pool + +Qwen3-235B holds 188 KiB of BF16 KV per token (94 layers, 4 KV heads of 128). +With the defaults, each Attention rank needs +`500 * ceil(1026 / 16) + 128 * ceil(257 / 16)` blocks, about 555K tokens or +97 GiB; an H200 Attention rank has room for about 640K tokens. Keep +`DECODE_BENCH_NUM_PROMPTS` at or above `MAX_NUM_SEQS`. diff --git a/recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/4a2f_graph_dp4tp1.sh b/recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/4a2f_graph_dp4tp1.sh new file mode 100755 index 000000000..0d8b87e7c --- /dev/null +++ b/recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/4a2f_graph_dp4tp1.sh @@ -0,0 +1,182 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +# +# Qwen3-235B-A22B-FP8 4A2F prefill/decode colocation on one node with 6 GPUs: +# Attention DP4/TP1 and FFN DP2/TP1 (expert parallel), FULL_DECODE_ONLY CUDA +# graphs. Every knob is an environment variable (default in parentheses): +# +# MODEL_PATH weights dir or HF id (/path/model_weights/Qwen3-235B-A22B-FP8) +# AFD_CONNECTOR DeepEPM2NAFDConnector | P2pNcclAFDConnector +# (DeepEPM2NAFDConnector) +# M2N_PAYLOAD_FORMAT DeepEPM2NAFDConnector dispatch format (fp8_block128_fp32) +# M2N_EXTRA_CONFIG extra connector_extra_config JSON members for +# DeepEPM2NAFDConnector, e.g. "expert_num_sms": 130 (empty) +# KV_CONNECTOR none | decode_bench (none); see the README +# KV_TRANSFER_CONFIG raw --kv-transfer-config JSON for Attention; +# overrides KV_CONNECTOR (unset) +# ENABLE_PREFIX_CACHING 1 | 0 (1) +# ENABLE_DBO 1 | 0 (0) +# DBO_DECODE_TOKEN_THRESHOLD / DBO_PREFILL_TOKEN_THRESHOLD (2 / 8) +# MAX_NUM_SEQS per Attention rank (128) +# MAX_NUM_BATCHED_TOKENS prefill chunk per step (3072) +# CUDAGRAPH_CAPTURE_SIZES comma-separated decode graph sizes (8,16,32,64,96,128) +# MAX_MODEL_LEN (4096) +# GPU_MEMORY_UTILIZATION (unset: vLLM default) +# ATTN_DEVICES / FFN_DEVICES (0,1,2,3 / 4,5) +# AFD_HOST / AFD_PORT AFD connector rendezvous (127.0.0.1 / 6269) +# AFD_PROCESS_GROUP_TIMEOUT_S AFD join timeout; the roles can finish +# loading minutes apart (7200) +# ATTN_PORT / FFN_PORT HTTP ports (18305 / 18306) +# EXTRA_ARGS extra `vllm serve` args for both roles (empty) +# ATTN_EXTRA_ARGS / FFN_EXTRA_ARGS extra `vllm serve` args for one role (empty) +# +# KV_CONNECTOR=decode_bench only (tools/benchmarks/decode_bench.py): +# DECODE_BENCH_NUM_PROMPTS pool prompts per Attention rank; 0 = dummy KV (500) +# DECODE_BENCH_PROMPT_LEN pool prompt length (1025) +# DECODE_BENCH_MAX_OUTPUT_LEN longest measured output; sizes the KV check (256) +# DECODE_BENCH_SEED pool token seed (0) +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +REPO_ROOT="$(cd "${SCRIPT_DIR}/../../../../.." && pwd)" + +MODEL_PATH=${MODEL_PATH:-/path/model_weights/Qwen3-235B-A22B-FP8} +AFD_CONNECTOR=${AFD_CONNECTOR:-DeepEPM2NAFDConnector} +M2N_PAYLOAD_FORMAT=${M2N_PAYLOAD_FORMAT:-fp8_block128_fp32} +M2N_EXTRA_CONFIG=${M2N_EXTRA_CONFIG:-} +KV_CONNECTOR=${KV_CONNECTOR:-none} +KV_TRANSFER_CONFIG=${KV_TRANSFER_CONFIG:-} +ENABLE_PREFIX_CACHING=${ENABLE_PREFIX_CACHING:-1} +ENABLE_DBO=${ENABLE_DBO:-0} +DBO_DECODE_TOKEN_THRESHOLD=${DBO_DECODE_TOKEN_THRESHOLD:-2} +DBO_PREFILL_TOKEN_THRESHOLD=${DBO_PREFILL_TOKEN_THRESHOLD:-8} +MAX_NUM_SEQS=${MAX_NUM_SEQS:-128} +MAX_NUM_BATCHED_TOKENS=${MAX_NUM_BATCHED_TOKENS:-3072} +CUDAGRAPH_CAPTURE_SIZES=${CUDAGRAPH_CAPTURE_SIZES:-8,16,32,64,96,128} +MAX_MODEL_LEN=${MAX_MODEL_LEN:-4096} +GPU_MEMORY_UTILIZATION=${GPU_MEMORY_UTILIZATION:-} +ATTN_DEVICES=${ATTN_DEVICES:-0,1,2,3} +FFN_DEVICES=${FFN_DEVICES:-4,5} +AFD_HOST=${AFD_HOST:-127.0.0.1} +AFD_PORT=${AFD_PORT:-6269} +AFD_PROCESS_GROUP_TIMEOUT_S=${AFD_PROCESS_GROUP_TIMEOUT_S:-7200} +ATTN_PORT=${ATTN_PORT:-18305} +FFN_PORT=${FFN_PORT:-18306} +EXTRA_ARGS=${EXTRA_ARGS:-} +ATTN_EXTRA_ARGS=${ATTN_EXTRA_ARGS:-} +FFN_EXTRA_ARGS=${FFN_EXTRA_ARGS:-} +DECODE_BENCH_NUM_PROMPTS=${DECODE_BENCH_NUM_PROMPTS:-500} +DECODE_BENCH_PROMPT_LEN=${DECODE_BENCH_PROMPT_LEN:-1025} +DECODE_BENCH_MAX_OUTPUT_LEN=${DECODE_BENCH_MAX_OUTPUT_LEN:-256} +DECODE_BENCH_SEED=${DECODE_BENCH_SEED:-0} + +export VLLM_USE_V2_MODEL_RUNNER=1 + +NUM_ATTN_RANKS=$(tr ',' '\n' <<<"$ATTN_DEVICES" | wc -l | tr -d ' ') +NUM_FFN_RANKS=$(tr ',' '\n' <<<"$FFN_DEVICES" | wc -l | tr -d ' ') + +case "$AFD_CONNECTOR" in + # DeepEP M2N routes on Attention and sends FP8 tokens to the FFN experts. + DeepEPM2NAFDConnector) + COMPUTE_GATE_ON_ATTENTION=true + CONNECTOR_EXTRA_CONFIG=", \"connector_extra_config\": {\"payload_format\": \"${M2N_PAYLOAD_FORMAT}\"${M2N_EXTRA_CONFIG:+, ${M2N_EXTRA_CONFIG}}}" + ;; + # P2P routes on the FFN and sends BF16 hidden states. Do not run its FFN + # with PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True: the first step hangs. + P2pNcclAFDConnector) + COMPUTE_GATE_ON_ATTENTION=false + CONNECTOR_EXTRA_CONFIG="" + ;; + *) + echo "unsupported AFD_CONNECTOR=${AFD_CONNECTOR}" >&2 + exit 1 + ;; +esac + +if [[ -z "$KV_TRANSFER_CONFIG" ]]; then + case "$KV_CONNECTOR" in + none) ;; + decode_bench) + KV_TRANSFER_CONFIG=$(cat <&2 + exit 1 + ;; + esac +fi +if [[ "$KV_TRANSFER_CONFIG" == *tools.benchmarks.* ]]; then + # tools/benchmarks is not in the wheel; the vLLM scheduler and worker + # subprocesses import it from the repository root. + export PYTHONPATH="${REPO_ROOT}${PYTHONPATH:+:${PYTHONPATH}}" +fi + +afd_config() { + cat < ffn.log 2>&1 & + +CUDA_VISIBLE_DEVICES=$ATTN_DEVICES uv run vllm serve "$MODEL_PATH" \ + --additional-config "$(afd_config attention)" \ + "${COMMON_ARGS[@]}" \ + "${ATTN_ARGS[@]}" \ + --host 127.0.0.1 \ + --port "$ATTN_PORT" > attn.log 2>&1 & + +wait diff --git a/recipe/gpu/P2pNcclAFDConnector/qwen3_235b_a22b_fp8/README.md b/recipe/gpu/P2pNcclAFDConnector/qwen3_235b_a22b_fp8/README.md new file mode 100644 index 000000000..087aa0ad8 --- /dev/null +++ b/recipe/gpu/P2pNcclAFDConnector/qwen3_235b_a22b_fp8/README.md @@ -0,0 +1,15 @@ +# Qwen3-235B-A22B-FP8 AFD Recipe (P2pNcclAFDConnector) + +`prefill_decode_colocation/4a2f_graph_dp4tp1.sh` runs the shared +Qwen3-235B-A22B-FP8 4A2F recipe with `AFD_CONNECTOR=P2pNcclAFDConnector`. +Topology, environment knobs, and the decode-only benchmark are documented with +the script itself in +[`../../DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/README.md`](../../DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/README.md). + +```bash +export MODEL_PATH=/path/model_weights/Qwen3-235B-A22B-FP8 +bash recipe/gpu/P2pNcclAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/4a2f_graph_dp4tp1.sh +``` + +Do not set `PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True` for this +connector's FFN: the FFN hangs on its first step. diff --git a/recipe/gpu/P2pNcclAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/4a2f_graph_dp4tp1.sh b/recipe/gpu/P2pNcclAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/4a2f_graph_dp4tp1.sh new file mode 100755 index 000000000..f5e183f4c --- /dev/null +++ b/recipe/gpu/P2pNcclAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/4a2f_graph_dp4tp1.sh @@ -0,0 +1,12 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +# +# P2pNcclAFDConnector entry point of the shared Qwen3-235B-A22B-FP8 4A2F recipe. +# The topology and every environment knob are documented in +# recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/4a2f_graph_dp4tp1.sh. +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +AFD_CONNECTOR=P2pNcclAFDConnector exec bash \ + "${SCRIPT_DIR}/../../../DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/4a2f_graph_dp4tp1.sh" From d39d2c84adc510f5ec2198f73edfdb5058678de1 Mon Sep 17 00:00:00 2001 From: ronenkat <16743404+ronenkat@users.noreply.github.com> Date: Thu, 29 Oct 2026 17:07:00 +0200 Subject: [PATCH 19/32] fix(dbo): replay the microbatched CUDA graph for underfilled DP ranks vLLM 0.30 runs a microbatched (DBO) step eager whenever the smallest DP rank holds at most half of the CUDA graph's tokens (`2 * min < N`). A FULL microbatched graph is captured with a fixed request/token split, and native slicing re-derives the split from the real requests, so an underfilled rank would replay shifted metadata. With DP imbalance this fires on most decode steps and an eager microbatched step costs several times a graph step. For FULL one-token decodes, keep the captured split instead: each microbatch takes the request and token range the graph was captured with, and a trailing microbatch without real requests keeps its padded shape (one-token query bound, nonzero host-side sequence bound). The step then replays the microbatched graph on every rank. Configs this is not validated for (speculative decoding, inner-state models, MLA, PCP/DCP) replay the regular graph instead of running eager. Other steps keep native slicing and coordination. - gpu_mrv2_ubatching: instance-bound UBatchRunner.prepare copy with the captured split, a copied sync_cudagraph_and_dp_padding for DP>1 without the eager guard, and the same rule in the DP1 branch. - Unit tests cover the captured split against native vLLM slicing, the padding-only microbatch, and DP1/DP2 graph selection. Co-Authored-By: Claude Opus 5.5 Signed-off-by: ronenkat <16743404+ronenkat@users.noreply.github.com> --- .../compat/patches/gpu_mrv2_ubatching.py | 401 +++++++++++++++++- .../module/compatibility_and_patches.md | 1 + .../compat/patches/test_gpu_mrv2_ubatching.py | 377 +++++++++++++++- 3 files changed, 755 insertions(+), 24 deletions(-) diff --git a/afd_plugin/compat/patches/gpu_mrv2_ubatching.py b/afd_plugin/compat/patches/gpu_mrv2_ubatching.py index 4715b2d92..867048386 100644 --- a/afd_plugin/compat/patches/gpu_mrv2_ubatching.py +++ b/afd_plugin/compat/patches/gpu_mrv2_ubatching.py @@ -3,7 +3,9 @@ """Scoped DBO seams for the pinned native GPU ModelRunnerV2. The native UBatchRunner owns input slicing, Attention preparation, threads and -CUDA graphs. AFD extends DP1 eligibility, dispatch and child-context metadata. +CUDA graphs. AFD extends DP1 eligibility, dispatch and child-context metadata, +and keeps FULL microbatched decodes on the captured request split so an +underfilled DP rank replays the graph instead of running eager. Bindings are limited to AFD Attention initialization/execution and restored on exit; worker initialization and dispatch run synchronously in each process. """ @@ -13,28 +15,116 @@ from collections.abc import Iterator from contextlib import contextmanager from dataclasses import replace +from itertools import pairwise from types import MethodType from typing import Any import torch +import torch.distributed as dist from vllm.config import CUDAGraphMode, ParallelConfig, VllmConfig +from vllm.distributed.parallel_state import get_dp_group from vllm.forward_context import DPMetadata, ForwardContext, create_forward_context from vllm.v1.kv_cache_interface import KVCacheConfig +from vllm.v1.worker.dp_utils import should_skip_dp_coordination from vllm.v1.worker.gpu import model_runner +from vllm.v1.worker.gpu.attn_utils import build_slot_mappings_by_layer from vllm.v1.worker.gpu.cudagraph_utils import ( BatchExecutionDescriptor, CudaGraphManager, ) -from vllm.v1.worker.gpu.dp_utils import DPSyncState, sync_cudagraph_and_dp_padding +from vllm.v1.worker.gpu.dp_utils import DPSyncState +from vllm.v1.worker.gpu.input_batch import InputBatch from vllm.v1.worker.gpu.model_states.interface import ModelState -from vllm.v1.worker.gpu.ubatch_utils import UBatchRunner, logger +from vllm.v1.worker.gpu.ubatch_utils import ( + UBatchRunner, + UBatchState, + _slice_input_batch, + create_ubatch_slices, + logger, +) from vllm.v1.worker.ubatch_utils import ( + UBatchSlice, UBatchSlices, check_ubatch_thresholds, get_num_ubatches, ) from vllm.v1.worker.utils import AttentionGroup +# A FULL microbatched decode graph is captured with one token per request. +DECODE_QUERY_LEN = 1 +# Host-side sequence bound for a microbatch holding only padding requests; the +# graph's kernels were captured with max_model_len and ignore it. +PADDING_ONLY_SEQ_LEN_BOUND = 1 + + +def replays_captured_ubatch_split(vllm_config: VllmConfig) -> bool: + """Whether FULL microbatched decodes keep the captured request split. + + The split keeps every microbatch's captured request and token range, so a + rank holding at most half of the graph's requests replays its trailing + microbatch on padding rows only. Limited to one-token decodes over + attention without inner state or context parallelism; MLA is not yet + validated with a padding-only microbatch. + """ + model_config = vllm_config.model_config + parallel_config = vllm_config.parallel_config + return ( + vllm_config.speculative_config is None + and not model_config.has_inner_state + and not model_config.use_mla + and parallel_config.prefill_context_parallel_size == 1 + and parallel_config.decode_context_parallel_size == 1 + ) + + +def needs_regular_graph_fallback( + cudagraph_manager: CudaGraphManager, + num_reqs: int, + min_num_tokens: int, + max_num_tokens: int, + uniform_token_count: int | None, + num_active_loras: int, + num_ubatches: int, +) -> bool: + """Whether a microbatched graph would replay a shifted request split. + + Native slicing re-derives the split from the real requests, so a graph only + replays correctly when the smallest rank fills more than its leading + microbatches. Without the captured split, such a step replays the regular + graph instead of running eager. Every rank passes the same agreed token + bounds, so every rank returns the same answer. + """ + if replays_captured_ubatch_split(cudagraph_manager.vllm_config): + return False + ubatch_desc = cudagraph_manager.dispatch( + num_reqs, + max_num_tokens, + uniform_token_count, + num_active_loras=num_active_loras, + num_ubatches=num_ubatches, + ) + return ( + ubatch_desc.cg_mode == CUDAGraphMode.FULL + and num_ubatches * min_num_tokens <= ubatch_desc.num_tokens + ) + + +def captured_ubatch_slices(input_batch: InputBatch, num_ubatches: int) -> UBatchSlices: + """Slices a FULL one-token-per-request graph was captured with. + + Capture splits a dummy batch whose request and token counts are equal at + `num_tokens // num_ubatches`, the last microbatch taking the remainder. + Native slicing yields the same slices whenever real requests reach the + last microbatch. + """ + num_tokens = input_batch.num_tokens_after_padding + split_point = num_tokens // num_ubatches + bounds = [split_point * i for i in range(num_ubatches)] + [num_tokens] + return [ + UBatchSlice(slice(start, stop), slice(start, stop)) + for start, stop in pairwise(bounds) + ] + # Patch reason: the native factory excludes DP1, although AFD schedules its two # stages locally and sends their complete metadata to the FFN control plane. @@ -70,9 +160,103 @@ def maybe_build_ubatch_runner( if parallel_config.data_parallel_size == 1: runner._make_forward_contexts = MethodType(_make_forward_contexts, runner) # ### PATCH END: bind native DP1 child metadata on this instance. + # ### PATCH START: bind the captured-split prepare on this instance. + if replays_captured_ubatch_split(vllm_config): + runner.prepare = MethodType(prepare, runner) + # ### PATCH END: bind the captured-split prepare on this instance. return runner +# Patch reason: native prepare re-derives the microbatch split from the real +# requests and clamps an empty trailing microbatch onto the last real request. +# A FULL microbatched graph replays fixed request/token offsets, so a rank with +# at most half of the graph's requests would replay with shifted metadata, and +# native dispatch runs such steps eager instead (`2 * min < N`). +# Patch functionality: slice FULL one-token decodes at the captured split. A +# trailing microbatch without real requests keeps its padded shape, with a +# one-token query bound and a nonzero host-side sequence bound. +# Signature: matches UBatchRunner.prepare exactly at the pinned commit. +# Upstream: vllm/v1/worker/gpu/ubatch_utils.py, UBatchRunner.prepare; +# ced6857afa0ea7b2e3f0846a62e1394e90f15607 (vLLM 0.30.0). +# Removal plan: remove when native FULL microbatching keeps the captured split. +def prepare( + self, + input_batch: InputBatch, + block_tables: tuple[torch.Tensor, ...], + slot_mappings: torch.Tensor, + cg_mode: CUDAGraphMode = CUDAGraphMode.NONE, + for_capture: bool = False, +) -> UBatchState: + """Split the batch into the microbatches the step will run on. + + Attention metadata is built per microbatch and carried in the forward + contexts the threads install, not in the caller's. + + FULL uses padded sizes; `for_capture` refreshes metadata from dummy + block tables. + """ + # ### PATCH START: slice FULL one-token decodes at the captured split. + keeps_captured_split = ( + cg_mode == CUDAGraphMode.FULL + and input_batch.num_tokens == input_batch.num_reqs * DECODE_QUERY_LEN + and input_batch.num_tokens_after_padding + == input_batch.num_reqs_after_padding * DECODE_QUERY_LEN + ) + if keeps_captured_split: + # A microbatch without real requests has no query length to derive. + input_batch = replace(input_batch, max_query_len=DECODE_QUERY_LEN) + ubatch_slices = captured_ubatch_slices(input_batch, self.num_ubatches) + else: + ubatch_slices = create_ubatch_slices(input_batch, self.num_ubatches) + # ### PATCH END: slice FULL one-token decodes at the captured split. + + attn_metadata = [] + slot_mappings_by_layer = [] + is_padding = [] + for i, ubatch_slice in enumerate(ubatch_slices): + ubatch = _slice_input_batch( + input_batch, + ubatch_slice, + self.ubatch_query_start_loc[i], + self.ubatch_seq_lens[i], + self.ubatch_dcp_local_seq_lens[i], + dcp_size=self.dcp_size, + dcp_rank=self.dcp_rank, + cp_interleave=self.cp_interleave, + ) + # ### PATCH START: keep a padding-only microbatch's sequence bound nonzero. + if keeps_captured_split and ubatch.num_reqs == 0: + ubatch.seq_lens_cpu_upper_bound[0] = PADDING_ONLY_SEQ_LEN_BOUND + # ### PATCH END: keep a padding-only microbatch's sequence bound nonzero. + ubatch_slot_mappings = slot_mappings[:, ubatch_slice.token_slice] + ubatch_block_tables = tuple( + block_table[ubatch_slice.request_slice] for block_table in block_tables + ) + attn_metadata.append( + self.model_state.prepare_attn( + ubatch, + cg_mode, + ubatch_block_tables, + ubatch_slot_mappings, + self.attn_groups, + self.kv_cache_config, + ubatch_idx=i, + for_capture=for_capture, + ) + ) + slot_mappings_by_layer.append( + build_slot_mappings_by_layer(ubatch_slot_mappings, self.kv_cache_config) + ) + is_padding.append(ubatch.is_padding) + + return UBatchState( + slices=ubatch_slices, + forward_contexts=self._make_forward_contexts( + ubatch_slices, attn_metadata, slot_mappings_by_layer, is_padding + ), + ) + + # Patch reason: DPMetadata.make asserts DP>1 or SP; AFD DP1 uses neither. # Patch functionality: build the same native DPMetadata directly for this DP1 # runner instance. Native context fields, token counts and SP methods are kept. @@ -115,9 +299,203 @@ def _make_forward_contexts( return forward_contexts +# Patch reason: native DP>1 microbatching runs eager whenever the smallest rank +# holds at most half of the graph (`2 * min < N`), because native slicing would +# replay a shifted request split. +# Patch functionality: keep the microbatched FULL graph when AFD's prepare +# replays the captured split; otherwise replay the regular graph, never eager. +# Every other branch is the native coordination unchanged. +# Signature: matches sync_cudagraph_and_dp_padding exactly at the pinned commit. +# Upstream: vllm/v1/worker/gpu/dp_utils.py, sync_cudagraph_and_dp_padding; +# ced6857afa0ea7b2e3f0846a62e1394e90f15607 (vLLM 0.30.0). +# Removal plan: remove when native FULL microbatching keeps the captured split. +def sync_cudagraph_and_dp_padding( + cudagraph_manager: CudaGraphManager | None, + desired_batch_desc: BatchExecutionDescriptor, + num_tokens: int, + num_reqs: int, + uniform_token_count: int | None, + dp_size: int, + dp_rank: int, + max_query_len: int | None = None, + num_active_loras: int = 0, + parallel_config: ParallelConfig | None = None, + allow_ubatching: bool = False, + uniform_decode: bool = False, +) -> tuple[BatchExecutionDescriptor, DPSyncState | None]: + """ + Coordinates the batch descriptor and DP padding across all ranks. + + `parallel_config` is only needed to decide whether to microbatch, so callers + that never do (`allow_ubatching=False`) can leave it out. + + Returns (synced_batch_desc, sync). `sync` is None when no rank has work. + """ + assert dp_size > 1, "DP size must be greater than 1" + group = get_dp_group().cpu_group + tensor = torch.zeros(6, dp_size, dtype=torch.int32, device="cpu") + tensor[0][dp_rank] = num_tokens + tensor[1][dp_rank] = desired_batch_desc.cg_mode.value + tensor[2][dp_rank] = uniform_token_count or 0 # (0 means None) + tensor[3][dp_rank] = max_query_len or -1 # (-1 means None) + tensor[4][dp_rank] = int(allow_ubatching) + tensor[5][dp_rank] = num_reqs + if should_skip_dp_coordination(): + tensor[:] = tensor[:, dp_rank, None].clone() + else: + dist.all_reduce(tensor, group=group) + + num_tokens_across_dp = tensor[0] + cg_mode_across_dp = tensor[1] + uniform_token_counts_across_dp = tensor[2] + max_query_lens_across_dp = tensor[3] + allow_ubatching_across_dp = tensor[4] + num_reqs_across_dp = tensor[5] + + # If ranks disagree on the uniform token count, or its 0 (means None) set to None + synced_uniform_token_count: int | None = int(uniform_token_counts_across_dp[0]) + if synced_uniform_token_count == 0 or not torch.all( + uniform_token_counts_across_dp == synced_uniform_token_count + ): + synced_uniform_token_count = None + + if torch.all(num_tokens_across_dp == 0).item(): + synced_desc = BatchExecutionDescriptor( + cg_mode=CUDAGraphMode.NONE, num_tokens=0, num_reqs=0 + ) + return synced_desc, None + + if torch.all(allow_ubatching_across_dp == 1): + # This rank voted too, so its caller passed a config. + assert parallel_config is not None + # A uniform decode only if every rank runs the same uniform query + # length, so every rank reaches the same answer here. + uniform_decode_across_dp = uniform_decode and torch.all( + uniform_token_counts_across_dp == (uniform_token_count or 0) + ) + # ### PATCH START: regular graph when the captured split cannot serve. + if check_ubatch_thresholds( + parallel_config, + # Thresholds only grow with the token count, so holding the smallest + # rank to them holds every rank to them. + int(num_tokens_across_dp.min()), + uniform_decode=uniform_decode_across_dp, + ) and not ( + cudagraph_manager is not None + and needs_regular_graph_fallback( + cudagraph_manager, + num_reqs, + int(num_tokens_across_dp.min()), + int(num_tokens_across_dp.max()), + synced_uniform_token_count, + num_active_loras, + get_num_ubatches(parallel_config), + ) + ): + # ### PATCH END: regular graph when the captured split cannot serve. + # Expert all-to-all requires every rank to split and pad equally. + # Empty microbatches run as dummy batches. + ubatch_num_tokens = int(num_tokens_across_dp.max()) + num_ubatches = get_num_ubatches(parallel_config) + ubatch_desc = None + if cudagraph_manager is not None: + # Match a FULL microbatched CUDA graph, falling back to an eager + # descriptor if no graph matches. + ubatch_desc = cudagraph_manager.dispatch( + num_reqs, + ubatch_num_tokens, + synced_uniform_token_count, + num_active_loras=num_active_loras, + num_ubatches=num_ubatches, + ) + # ### PATCH START: no eager fallback for an underfilled rank. + # The captured split replays it; otherwise the threshold check + # above already routed the step to the regular graph. + # ### PATCH END: no eager fallback for an underfilled rank. + if ubatch_desc is None: + ubatch_desc = BatchExecutionDescriptor( + cg_mode=CUDAGraphMode.NONE, + num_tokens=ubatch_num_tokens, + num_reqs=num_reqs, + num_ubatches=num_ubatches, + ) + # Refresh the token count to include CUDA graph padding. + ubatch_num_tokens = ubatch_desc.num_tokens + num_reqs = int(num_reqs_across_dp.max()) + if ubatch_desc.cg_mode == CUDAGraphMode.FULL: + assert ubatch_desc.num_reqs is not None + num_reqs = ubatch_desc.num_reqs + return ubatch_desc, DPSyncState( + num_tokens_across_dp=torch.full_like( + num_tokens_across_dp, ubatch_num_tokens + ), + uniform_token_count=synced_uniform_token_count, + eager=ubatch_desc.cg_mode == CUDAGraphMode.NONE, + num_reqs=num_reqs, + ) + + synced_cg_mode = CUDAGraphMode(int(cg_mode_across_dp.min().item())) + + # If any rank wants to run eager, all ranks run eager + if synced_cg_mode == CUDAGraphMode.NONE: + return ( + BatchExecutionDescriptor( + cg_mode=CUDAGraphMode.NONE, + num_tokens=num_tokens, + num_reqs=num_reqs, + num_active_loras=desired_batch_desc.num_active_loras, + ), + DPSyncState( + num_tokens_across_dp=num_tokens_across_dp, + uniform_token_count=synced_uniform_token_count, + eager=True, + num_reqs=int(num_reqs_across_dp.max()), + ), + ) + + assert cudagraph_manager is not None, ( + "cudagraph_manager should only be None during profile run, " + "where synced_cg_mode must be NONE across all DP ranks" + ) + synced_num_tokens = int(num_tokens_across_dp.max().item()) + + # Varlen decode graphs are selected by the query-length bound, so ranks must agree + # on it or they pad to different token counts below. + synced_max_query_len: int | None = None + if bool(torch.all(max_query_lens_across_dp != -1).item()): + synced_max_query_len = int(max_query_lens_across_dp.max().item()) + + # Dispatch for the final synced values, use num_reqs instead of synced_num_reqs + # so we don't perform request padding for PIECEWISE graphs. + # num_active_loras is per-rank and doesn't need cross-rank agreement. + synced_desc = cudagraph_manager.dispatch( + num_reqs, + synced_num_tokens, + synced_uniform_token_count, + num_active_loras=num_active_loras, + max_query_len=synced_max_query_len, + ) + + # Update num_tokens_across_dp to reflect padded size. + num_tokens_across_dp[:] = synced_desc.num_tokens + + return synced_desc, DPSyncState( + num_tokens_across_dp=num_tokens_across_dp, + uniform_token_count=synced_uniform_token_count, + eager=False, + num_reqs=( + synced_desc.num_reqs + if synced_desc.cg_mode == CUDAGraphMode.FULL + and synced_desc.num_reqs is not None + else int(num_reqs_across_dp.max()) + ), + ) + + # Patch reason: native DP1 dispatch returns before its DBO threshold decision. -# Patch functionality: apply the native threshold/graph-padding policy locally; -# retain native DP>1 coordination and native descriptor selection. +# Patch functionality: apply the native threshold/graph-padding policy locally, +# replaying the regular graph when the captured split cannot serve the step; +# route DP>1 coordination through AFD's sync_cudagraph_and_dp_padding. # Signature: matches dispatch_cg_and_sync_dp exactly at the pinned commit. # Upstream: vllm/v1/worker/gpu/dp_utils.py, dispatch_cg_and_sync_dp; # ced6857afa0ea7b2e3f0846a62e1394e90f15607 (vLLM 0.30.0). @@ -203,6 +581,16 @@ def dispatch_cg_and_sync_dp( ubatch_desc = None if not need_eager and not reuse_eager: assert cudagraph_manager is not None + if needs_regular_graph_fallback( + cudagraph_manager, + num_reqs, + num_tokens, + num_tokens, + uniform_token_count, + num_active_loras, + num_ubatches, + ): + return batch_desc, None ubatch_desc = cudagraph_manager.dispatch( num_reqs, num_tokens, @@ -210,9 +598,6 @@ def dispatch_cg_and_sync_dp( num_active_loras=num_active_loras, num_ubatches=num_ubatches, ) - if num_ubatches * num_tokens < ubatch_desc.num_tokens: - # Match the native DP>1 guard against an empty trailing stage. - ubatch_desc = None if ubatch_desc is None: ubatch_desc = BatchExecutionDescriptor( cg_mode=CUDAGraphMode.NONE, diff --git a/docs/design/module/compatibility_and_patches.md b/docs/design/module/compatibility_and_patches.md index 64c13ae4b..4eae54022 100644 --- a/docs/design/module/compatibility_and_patches.md +++ b/docs/design/module/compatibility_and_patches.md @@ -125,6 +125,7 @@ not the package dependency policy. | [`async_dp_forward_context.py`](../../../afd_plugin/compat/patches/async_dp_forward_context.py): `vllm.forward_context.set_forward_context` plus already-imported worker aliases | Skips native `DPMetadata` construction/coordination only for AFD async-DP; otherwise uses the copied upstream flow. | Imported by `register_afd`; same target/dev/unknown guard. Rebinds known already-imported aliases so callers do not retain the old function. | [`test_async_dp_forward_context.py`](../../../tests/unit/compat/patches/test_async_dp_forward_context.py) covers async skip and non-async coordination. | Remove when vLLM supports a per-engine-role opt-out from native MoE DP metadata coordination. | | [`config_validation.py`](../../../afd_plugin/compat/patches/config_validation.py): `EngineArgs.create_engine_config`, `VllmConfig.__post_init__` | For AFD-owned ubatching with a non-DeepEP backend, temporarily presents `deepep_low_latency` during upstream validation and restores the configured backend. After upstream platform normalization, maps an initial `worker_cls="auto"` to the role-specific CUDA or standard Ascend AFD worker. Native hybrid/multimodal DBO rejection remains intact, including Qwen3.5/3.6 GDN models. | Imported by `register_afd`; accepts the target version, development versions, or missing version metadata. Saves originals on upstream modules under AFD-specific attributes before installing wrappers. Explicit worker paths and non-AFD configs are not remapped. | [`test_config_validation.py`](../../../tests/unit/compat/patches/test_config_validation.py) covers backend relaxation, four role/platform mappings, explicit and non-AFD preservation, repeated validation, unsupported platforms, and dev versions. | Remove the backend branch when upstream validation distinguishes plugin-owned ubatching; remove worker mapping when vLLM offers plugin-owned role-aware worker selection. | | [`engine_core.py`](../../../afd_plugin/compat/patches/engine_core.py): `EngineCore.__init__`, `_initialize_kv_caches`, `shutdown`; `EngineCoreProc.run_busy_loop`; `DPEngineCoreProc.run_busy_loop` | AFD FFN becomes a connector daemon: construct executor, skip scheduler/KV setup, return an empty KV-shaped result on late paths, start/monitor/stop the FFN worker loop, and use FFN-safe shutdown. Non-FFN branches copy pinned upstream behavior. | Imported by `register_afd`; Installs under the target/dev compatibility check (also permitting missing version metadata); no saved-original sentinel. Direct class assignment remains tied to the pinned upstream implementation. | [`test_engine_core.py`](../../../tests/unit/compat/patches/test_engine_core.py) covers FFN initialization, non-FFN behavior, and daemon start/stop; role runtime tests cover error propagation. | Remove when vLLM offers a headless connector-daemon engine lifecycle or an executor mode that does not require scheduler/KV ownership. | +| [`gpu_mrv2_ubatching.py`](../../../afd_plugin/compat/patches/gpu_mrv2_ubatching.py): `vllm.v1.worker.gpu.model_runner.maybe_build_ubatch_runner` and `dispatch_cg_and_sync_dp` aliases, `UBatchRunner.prepare` and `_make_forward_contexts` (instance-bound), `sync_cudagraph_and_dp_padding` (copied) | Builds the native `UBatchRunner` at DP1 with DP1 child DP metadata. FULL one-token decode microbatches keep the captured request/token split, so a DP rank holding at most half of the graph's requests replays the microbatched graph instead of running eager; a trailing microbatch without real requests keeps its padded shape. Configs without that split (speculative decoding, inner-state models, MLA, PCP/DCP) replay the regular graph instead. Other steps use native slicing and coordination. | Scoped to AFD Attention KV initialization and each `execute_model`; aliases are restored on exit. The prepare binding is per runner instance, chosen once from the config. | [`test_gpu_mrv2_ubatching.py`](../../../tests/unit/compat/patches/test_gpu_mrv2_ubatching.py) covers the DP1 factory and metadata, captured-split slicing against native vLLM slicing, padding-only microbatches, and DP1 and DP2 graph selection for underfilled ranks. | Remove the split and dispatch patches when native FULL microbatching keeps the captured split; the factory when native GPU DBO supports DP1. | | [`npu/ascend_platform.py`](../../../afd_plugin/compat/patches/npu/ascend_platform.py): `NPUPlatform.check_and_update_config` | Snapshots AFD DBO state, runs upstream normalization, and restores configured `enable_dbo`, `ubatch_size`, and `all2all_backend` in `finally`; non-AFD behavior is unchanged. | Installed through the config facade during AFD NPU normalization and verified again by the worker runtime facade; no version guard. Saves the original on the class and uses a class sentinel. The runtime facade caches success only after the wrapper is installed, so an early missing vLLM-Ascend import remains retryable. | [`test_runtime.py`](../../../tests/unit/compat/test_runtime.py) and [`test_npu_runtime.py`](../../../tests/unit/v1/worker/test_npu_runtime.py). | Remove when vLLM-Ascend recognizes plugin-owned DBO workers or no longer clears these fields. | | [`npu/ascend_config.py`](../../../afd_plugin/compat/patches/npu/ascend_config.py): `vllm_ascend.ascend_config.init_ascend_config`, loaded factory aliases, and the AFD DP child entry point | Excludes the AFD namespace from native constructor arguments; preserves config identity, validation, singleton publication, and cache refresh. DP children install the factory before native Ascend initialization. | Installed by the NPU config facade; rebinds already-loaded aliases. The DP entry point is selected only for AFD synchronous Attention or FFN. | [`test_ascend_config_lifecycle.py`](../../../tests/unit/compat/npu/test_ascend_config_lifecycle.py), [`test_ascend_engine_core_binding.py`](../../../tests/unit/compat/npu/test_ascend_engine_core_binding.py). | Remove when Ascend accepts plugin namespaces and loads general plugins before config validation in every DP child. | | [`npu/mla_graph.py`](../../../afd_plugin/compat/patches/npu/mla_graph.py): `vllm_ascend.attention.mla_v1.get_graph_params` | Resolves an AFD ubatch-owned `GraphParams` registry from the active forward context and otherwise delegates to the saved upstream process-global resolver. | Called through `apply_afd_ascend_patches_if_needed` during NPU worker startup; no version guard. Saves the original on the upstream module and uses a module sentinel. | [`test_runtime.py`](../../../tests/unit/compat/test_runtime.py) covers AFD-context resolution, upstream fallback, and idempotence; [`test_npu_mla_graph.py`](../../../tests/unit/v1/worker/test_npu_mla_graph.py) covers the owning graph lifecycle. | Remove when vLLM-Ascend accepts a forward-context-local MLA graph registry or exposes an equivalent resolver hook. | diff --git a/tests/unit/compat/patches/test_gpu_mrv2_ubatching.py b/tests/unit/compat/patches/test_gpu_mrv2_ubatching.py index 83c027a57..2c7f929f1 100644 --- a/tests/unit/compat/patches/test_gpu_mrv2_ubatching.py +++ b/tests/unit/compat/patches/test_gpu_mrv2_ubatching.py @@ -8,6 +8,7 @@ from contextlib import contextmanager from dataclasses import dataclass, replace from enum import Enum +from itertools import pairwise from pathlib import Path from types import MethodType, SimpleNamespace @@ -34,9 +35,42 @@ class _BatchDescriptor: num_ubatches: int = 1 +def _vllm_config( + *, + speculative=False, + inner_state=False, + mla=False, + pcp_size=1, + dcp_size=1, + parallel=None, +): + parallel = SimpleNamespace( + **{ + **vars(parallel or _parallel()), + "prefill_context_parallel_size": pcp_size, + "decode_context_parallel_size": dcp_size, + } + ) + return SimpleNamespace( + speculative_config=object() if speculative else None, + model_config=SimpleNamespace(has_inner_state=inner_state, use_mla=mla), + parallel_config=parallel, + ) + + +UNSUPPORTED_SPLIT_CONFIGS: dict[str, dict[str, bool | int]] = { + "speculative": {"speculative": True}, + "inner-state": {"inner_state": True}, + "mla": {"mla": True}, + "pcp": {"pcp_size": 2}, + "dcp": {"dcp_size": 2}, +} + + class _GraphManager: - def __init__(self, capture_tokens: int | None = None): + def __init__(self, capture_tokens: int | None = None, vllm_config=None): self.capture_tokens = capture_tokens + self.vllm_config = vllm_config or _vllm_config() self.calls: list[int] = [] def dispatch( @@ -75,16 +109,94 @@ def _parallel(*, dp_size=1, use_ubatching=True): num_ubatches=2, dbo_decode_token_threshold=2, dbo_prefill_token_threshold=12, + prefill_context_parallel_size=1, + decode_context_parallel_size=1, ) +@dataclass(frozen=True) +class _UBatchSlice: + request_slice: slice + token_slice: slice + + +def _bounds(slices): + return [ + ( + s.request_slice.start, + s.request_slice.stop, + s.token_slice.start, + s.token_slice.stop, + ) + for s in slices + ] + + +@dataclass(frozen=True) +class _DecodeBatch: + """The InputBatch fields prepare reads for a one-token-per-request batch.""" + + num_reqs: int + num_tokens: int + num_reqs_after_padding: int + num_tokens_after_padding: int + max_query_len: int | None = None + + +def _native_slices(input_batch, num_ubatches): + # Native vLLM 0.30 slicing clamps a trailing microbatch without real + # requests onto the last real request; only that outcome matters here. + split = input_batch.num_tokens_after_padding // num_ubatches + last_start = min(split, input_batch.num_reqs - 1) + return [ + _UBatchSlice(slice(0, split), slice(0, split)), + _UBatchSlice( + slice(last_start, input_batch.num_reqs_after_padding), + slice(split, input_batch.num_tokens_after_padding), + ), + ] + + +def _slice_decode_batch(input_batch, ubatch_slice, *args, **kwargs): + start, stop = ubatch_slice.request_slice.start, ubatch_slice.request_slice.stop + return SimpleNamespace( + num_reqs=max(0, min(stop, input_batch.num_reqs) - start), + max_query_len=input_batch.max_query_len, + seq_lens_cpu_upper_bound=[0] * (stop - start), + is_padding=ubatch_slice.token_slice, + ) + + +class _ModelState: + def prepare_attn(self, ubatch, cg_mode, block_tables, slot_mappings, *args, **kw): + return { + "max_query_len": ubatch.max_query_len, + "max_seq_len": max(ubatch.seq_lens_cpu_upper_bound), + "block_tables": block_tables, + } + + +class _SlotMappings: + def __getitem__(self, key): + return key[1] + + @pytest.fixture def patches(): # Execute the patch's functions with small native API doubles. This keeps # the real eligibility/control policy under test without importing CUDA. tree = ast.parse(PATCH_PATH.read_text()) functions: list[ast.stmt] = [ - node for node in tree.body if isinstance(node, ast.FunctionDef) + node + for node in tree.body + if isinstance(node, ast.FunctionDef) + or ( + isinstance(node, ast.Assign) + and all( + isinstance(target, ast.Name) and target.id.isupper() + for target in node.targets + ) + ) ] class NativeUBatchRunner: @@ -153,6 +265,14 @@ def check_thresholds(config, num_tokens, uniform_decode): "get_num_ubatches": lambda config: ( config.num_ubatches if config.use_ubatching else 1 ), + "pairwise": pairwise, + "UBatchSlice": _UBatchSlice, + "UBatchState": SimpleNamespace, + "create_ubatch_slices": _native_slices, + "_slice_input_batch": _slice_decode_batch, + "build_slot_mappings_by_layer": lambda slot_mappings, kv_cache_config: ( + slot_mappings + ), } future = ast.parse("from __future__ import annotations").body module = ast.Module(body=future + functions, type_ignores=[]) @@ -163,7 +283,7 @@ def check_thresholds(config, num_tokens, uniform_decode): @pytest.mark.parametrize("dp_size", [1, 2]) def test_factory_keeps_native_runner_and_changes_only_dp1_instance(patches, dp_size): parallel = _parallel(dp_size=dp_size) - config = SimpleNamespace(parallel_config=parallel) + config = _vllm_config(parallel=parallel) native_context_factory = patches.UBatchRunner._make_forward_contexts inputs = ("device", "model-state", [["attn-group"]], "kv-cache", 32) @@ -175,13 +295,25 @@ def test_factory_keeps_native_runner_and_changes_only_dp1_instance(patches, dp_s assert ("_make_forward_contexts" in vars(runner)) is (dp_size == 1) +@pytest.mark.parametrize("unsupported", [None, *UNSUPPORTED_SPLIT_CONFIGS]) +def test_factory_binds_captured_split_prepare_only_where_supported( + patches, unsupported +): + overrides = UNSUPPORTED_SPLIT_CONFIGS.get(unsupported, {}) + config = _vllm_config(parallel=_parallel(dp_size=2), **overrides) + + runner = patches.maybe_build_ubatch_runner(config, None, None, [], None, 32) + + assert ("prepare" in vars(runner)) is (unsupported is None) + + def test_disabled_dbo_does_not_construct_runner(patches): - config = SimpleNamespace(parallel_config=_parallel(use_ubatching=False)) + config = _vllm_config(parallel=_parallel(use_ubatching=False)) assert patches.maybe_build_ubatch_runner(config, None, None, [], None, 32) is None def test_dp1_native_child_metadata_keeps_stage_counts_and_context_fields(patches): - config = SimpleNamespace(parallel_config=_parallel()) + config = _vllm_config(parallel=_parallel()) runner = patches.maybe_build_ubatch_runner(config, None, None, [], None, 32) slices = [SimpleNamespace(num_tokens=3), SimpleNamespace(num_tokens=5)] attention = [{"layer": object()}, {"layer": object()}] @@ -237,17 +369,24 @@ def test_dp1_thresholds_select_eager_microbatches_without_handshake( @pytest.mark.parametrize( - "num_tokens, capture_tokens, expected_mode, expected_tokens", + "num_tokens, capture_tokens, split_supported, expected", [ - (8, 16, _CUDAGraphMode.FULL, 16), - (3, 16, _CUDAGraphMode.NONE, 3), - (8, None, _CUDAGraphMode.NONE, 8), + # (cg_mode, padded tokens, microbatches) + (9, 16, False, (_CUDAGraphMode.FULL, 16, 2)), + # At most half of the graph: the captured split replays it ... + (3, 16, True, (_CUDAGraphMode.FULL, 16, 2)), + (8, 16, True, (_CUDAGraphMode.FULL, 16, 2)), + # ... and without it the regular graph does, never eager. + (3, 16, False, (_CUDAGraphMode.FULL, 16, 1)), + (8, 16, False, (_CUDAGraphMode.FULL, 16, 1)), + (8, None, True, (_CUDAGraphMode.NONE, 8, 2)), ], ) -def test_dp1_uses_microbatch_graph_bucket_or_native_eager_padding_fallback( - patches, num_tokens, capture_tokens, expected_mode, expected_tokens +def test_dp1_replays_a_graph_for_an_underfilled_microbatched_step( + patches, num_tokens, capture_tokens, split_supported, expected ): - manager = _GraphManager(capture_tokens) + vllm_config = _vllm_config() if split_supported else _vllm_config(mla=True) + manager = _GraphManager(capture_tokens, vllm_config) desc, sync = patches.dispatch_cg_and_sync_dp( manager, num_tokens, @@ -260,10 +399,101 @@ def test_dp1_uses_microbatch_graph_bucket_or_native_eager_padding_fallback( uniform_decode=True, ) assert sync is None - assert desc.cg_mode is expected_mode - assert desc.num_tokens == expected_tokens - assert desc.num_ubatches == 2 - assert manager.calls == [1, 2] + assert (desc.cg_mode, desc.num_tokens, desc.num_ubatches) == expected + + +@pytest.mark.parametrize( + "graph_tokens, expected", + [ + (8, [(0, 4, 0, 4), (4, 8, 4, 8)]), + (32, [(0, 16, 0, 16), (16, 32, 16, 32)]), + (96, [(0, 48, 0, 48), (48, 96, 48, 96)]), + ], +) +def test_captured_split_halves_requests_and_tokens_alike( + patches, graph_tokens, expected +): + batch = _DecodeBatch(3, 3, graph_tokens, graph_tokens) + assert _bounds(patches.captured_ubatch_slices(batch, 2)) == expected + + +@pytest.mark.parametrize( + "unsupported, min_tokens, max_tokens, capture_tokens, falls_back", + [ + (None, 15, 17, 32, False), + ("mla", 15, 17, 32, True), + ("mla", 16, 17, 32, True), + ("mla", 17, 17, 32, False), + ("speculative", 2, 3, 8, True), + ("mla", 15, 17, None, False), + ], +) +def test_regular_graph_fallback_only_without_the_captured_split( + patches, unsupported, min_tokens, max_tokens, capture_tokens, falls_back +): + config = _vllm_config(**UNSUPPORTED_SPLIT_CONFIGS.get(unsupported, {})) + manager = _GraphManager(capture_tokens, config) + assert ( + patches.needs_regular_graph_fallback( + manager, max_tokens, min_tokens, max_tokens, 1, 0, 2 + ) + is falls_back + ) + + +def _prepare(patches, input_batch, cg_mode): + runner = SimpleNamespace( + num_ubatches=2, + ubatch_query_start_loc=[None, None], + ubatch_seq_lens=[None, None], + ubatch_dcp_local_seq_lens=[None, None], + dcp_size=1, + dcp_rank=0, + cp_interleave=1, + model_state=_ModelState(), + attn_groups=[], + kv_cache_config=None, + _make_forward_contexts=lambda slices, attn, slot_mappings, padding: attn, + ) + block_tables = (list(range(input_batch.num_reqs_after_padding)),) + return patches.prepare(runner, input_batch, block_tables, _SlotMappings(), cg_mode) + + +def test_underfilled_full_decode_replays_the_captured_split(patches): + # 15 real requests padded into a 32 graph: native slicing would start the + # trailing microbatch at request 14 instead of the captured 16. + state = _prepare(patches, _DecodeBatch(15, 15, 32, 32), _CUDAGraphMode.FULL) + + assert _bounds(state.slices) == [(0, 16, 0, 16), (16, 32, 16, 32)] + trailing = state.forward_contexts[1] + assert trailing["block_tables"] == (list(range(16, 32)),) + assert trailing["max_query_len"] == 1 + assert trailing["max_seq_len"] == 1 + + +@pytest.mark.parametrize("num_reqs", [17, 24, 32]) +def test_filled_full_decode_split_matches_native_slicing(patches, num_reqs): + input_batch = _DecodeBatch(num_reqs, num_reqs, 32, 32) + state = _prepare(patches, input_batch, _CUDAGraphMode.FULL) + assert state.slices == _native_slices(input_batch, 2) + + +@pytest.mark.parametrize( + "input_batch, cg_mode", + [ + (_DecodeBatch(15, 15, 32, 32), _CUDAGraphMode.NONE), + (_DecodeBatch(15, 20, 32, 32), _CUDAGraphMode.FULL), + ], + ids=["eager", "multi-token-request"], +) +def test_other_microbatched_steps_keep_native_slicing(patches, input_batch, cg_mode): + state = _prepare(patches, input_batch, cg_mode) + + assert state.slices == _native_slices(input_batch, 2) + assert [context["max_query_len"] for context in state.forward_contexts] == [ + None, + None, + ] def test_profile_dispatch_does_not_select_a_graph(patches): @@ -399,3 +629,118 @@ def initialize_kv_cache(self, *args, **kwargs): ) ] assert patches.model_runner.maybe_build_ubatch_runner is original + + +@pytest.fixture +def native_patch(): + pytest.importorskip("torch") + pytest.importorskip("vllm") + from afd_plugin.compat.patches import gpu_mrv2_ubatching + + return gpu_mrv2_ubatching + + +@pytest.mark.vllm_runtime +@pytest.mark.parametrize("graph_tokens", [8, 16, 32, 64, 96, 128]) +def test_captured_split_is_native_slicing_once_real_requests_reach_it( + native_patch, graph_tokens +): + import numpy as np + from vllm.v1.worker.gpu.ubatch_utils import create_ubatch_slices + + for num_reqs in range(1, graph_tokens + 1): + input_batch = SimpleNamespace( + num_reqs=num_reqs, + num_scheduled_tokens=np.ones(num_reqs, dtype=np.int32), + num_tokens_after_padding=graph_tokens, + num_reqs_after_padding=graph_tokens, + ) + native = _bounds(create_ubatch_slices(input_batch, 2)) + captured = _bounds(native_patch.captured_ubatch_slices(input_batch, 2)) + assert (native == captured) is (2 * num_reqs > graph_tokens), num_reqs + + +class _NativeGraphManager: + CAPTURE_SIZES = (8, 16, 32, 64) + + def __init__(self, vllm_config): + self.vllm_config = vllm_config + + def dispatch( + self, + num_reqs, + num_tokens, + uniform_token_count, + *, + num_active_loras=0, + max_query_len=None, + num_ubatches=1, + ): + from vllm.config import CUDAGraphMode + from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor + + size = next(size for size in self.CAPTURE_SIZES if size >= num_tokens) + return BatchExecutionDescriptor( + cg_mode=CUDAGraphMode.FULL, + num_tokens=size, + num_reqs=size, + uniform_token_count=uniform_token_count, + num_ubatches=num_ubatches, + ) + + +@pytest.mark.vllm_runtime +@pytest.mark.parametrize( + "tokens_across_dp, split_supported, expected", + [ + # (padded tokens, microbatches); every case replays a FULL graph. + ((15, 17), True, (32, 2)), + ((15, 17), False, (32, 1)), + ((16, 32), True, (32, 2)), + ((16, 32), False, (32, 1)), + ((16, 16), False, (16, 2)), + ], +) +def test_dp2_underfilled_rank_replays_a_graph_on_every_rank( + native_patch, monkeypatch, tokens_across_dp, split_supported, expected +): + import torch + from vllm.config import CUDAGraphMode + from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor + + def all_reduce(tensor, group=None): + # Peers vote like the calling rank except for their own batch size. + local_vote = tensor.sum(dim=1) + for rank, tokens in enumerate(tokens_across_dp): + tensor[:, rank] = local_vote + tensor[0, rank] = tokens + tensor[5, rank] = tokens + + monkeypatch.setattr(native_patch.dist, "all_reduce", all_reduce) + monkeypatch.setattr( + native_patch, "get_dp_group", lambda: SimpleNamespace(cpu_group=None) + ) + vllm_config = _vllm_config() if split_supported else _vllm_config(mla=True) + manager = _NativeGraphManager(vllm_config) + + results = [] + for rank, tokens in enumerate(tokens_across_dp): + desc, sync = native_patch.sync_cudagraph_and_dp_padding( + manager, + BatchExecutionDescriptor(CUDAGraphMode.FULL, tokens, tokens, 1), + tokens, + tokens, + 1, + len(tokens_across_dp), + rank, + parallel_config=_parallel(dp_size=len(tokens_across_dp)), + allow_ubatching=True, + uniform_decode=True, + ) + assert desc.cg_mode == CUDAGraphMode.FULL + assert sync is not None and not sync.eager + assert torch.equal( + sync.num_tokens_across_dp, torch.full((2,), expected[0], dtype=torch.int32) + ) + results.append((desc.num_tokens, desc.num_ubatches)) + assert results == [expected] * len(tokens_across_dp) From 3a187ee6b8b2d17fcb0d86e7118eed01b2679287 Mon Sep 17 00:00:00 2001 From: ronenkat <16743404+ronenkat@users.noreply.github.com> Date: Thu, 29 Oct 2026 18:29:57 +0200 Subject: [PATCH 20/32] feat(recipe): Attention TP knob for the Qwen3-235B 4A2F recipe ATTN_TP sets the Attention tensor-parallel size; Attention DP becomes the number of ATTN_DEVICES divided by it (default 1, the DP4/TP1 layout as before). ATTN_TP=4 runs Attention as one DP1/TP4 engine on the same four GPUs; num_attention_ranks stays 4 and the FFN keeps DP2/TP1. The script rejects a TP size that does not divide the Attention device count, and the README notes that one engine then needs MAX_NUM_SEQS and the capture sizes raised to hold every running request. Validated on 6xH200 with both AFD connectors: decode-only sweeps with ATTN_TP=4, MAX_NUM_SEQS=256, capture sizes up to 256, concurrency 256 and request rates 1-50, one prefilled token per request and no preemptions in every step. Co-Authored-By: Claude Opus 5.5 Signed-off-by: ronenkat <16743404+ronenkat@users.noreply.github.com> --- .../qwen3_235b_a22b_fp8/README.md | 4 ++++ .../4a2f_graph_dp4tp1.sh | 20 +++++++++++++------ 2 files changed, 18 insertions(+), 6 deletions(-) diff --git a/recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/README.md b/recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/README.md index 75f70577c..503029fc9 100644 --- a/recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/README.md +++ b/recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/README.md @@ -24,6 +24,10 @@ benchmark over real KV, exactly like the | 0, 1, 2, 3 | Attention | DP4 / TP1 | 18305 | | 4, 5 | FFN | DP2 / TP1, EP2 | 18306 | +`ATTN_TP=4` runs Attention as one DP1/TP4 engine on the same four GPUs instead +(`num_attention_ranks` stays 4); raise `MAX_NUM_SEQS` and the capture sizes +with it, since one engine then holds every running request. + Both roles run `FULL_DECODE_ONLY` CUDA graphs (sizes 8 to 128 by default); prefill runs eager in chunks of `MAX_NUM_BATCHED_TOKENS`. Send traffic to the Attention port. `DeepEPM2NAFDConnector` routes on Attention and dispatches FP8 diff --git a/recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/4a2f_graph_dp4tp1.sh b/recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/4a2f_graph_dp4tp1.sh index 0d8b87e7c..334afebed 100755 --- a/recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/4a2f_graph_dp4tp1.sh +++ b/recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/4a2f_graph_dp4tp1.sh @@ -3,8 +3,9 @@ # SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project # # Qwen3-235B-A22B-FP8 4A2F prefill/decode colocation on one node with 6 GPUs: -# Attention DP4/TP1 and FFN DP2/TP1 (expert parallel), FULL_DECODE_ONLY CUDA -# graphs. Every knob is an environment variable (default in parentheses): +# Attention DP4/TP1 (or DP1/TP4 with ATTN_TP=4) and FFN DP2/TP1 (expert +# parallel), FULL_DECODE_ONLY CUDA graphs. Every knob is an environment +# variable (default in parentheses): # # MODEL_PATH weights dir or HF id (/path/model_weights/Qwen3-235B-A22B-FP8) # AFD_CONNECTOR DeepEPM2NAFDConnector | P2pNcclAFDConnector @@ -18,12 +19,14 @@ # ENABLE_PREFIX_CACHING 1 | 0 (1) # ENABLE_DBO 1 | 0 (0) # DBO_DECODE_TOKEN_THRESHOLD / DBO_PREFILL_TOKEN_THRESHOLD (2 / 8) -# MAX_NUM_SEQS per Attention rank (128) +# MAX_NUM_SEQS per Attention DP rank (128) # MAX_NUM_BATCHED_TOKENS prefill chunk per step (3072) # CUDAGRAPH_CAPTURE_SIZES comma-separated decode graph sizes (8,16,32,64,96,128) # MAX_MODEL_LEN (4096) # GPU_MEMORY_UTILIZATION (unset: vLLM default) # ATTN_DEVICES / FFN_DEVICES (0,1,2,3 / 4,5) +# ATTN_TP Attention tensor-parallel size; Attention DP is the +# number of ATTN_DEVICES divided by it (1) # AFD_HOST / AFD_PORT AFD connector rendezvous (127.0.0.1 / 6269) # AFD_PROCESS_GROUP_TIMEOUT_S AFD join timeout; the roles can finish # loading minutes apart (7200) @@ -58,6 +61,7 @@ MAX_MODEL_LEN=${MAX_MODEL_LEN:-4096} GPU_MEMORY_UTILIZATION=${GPU_MEMORY_UTILIZATION:-} ATTN_DEVICES=${ATTN_DEVICES:-0,1,2,3} FFN_DEVICES=${FFN_DEVICES:-4,5} +ATTN_TP=${ATTN_TP:-1} AFD_HOST=${AFD_HOST:-127.0.0.1} AFD_PORT=${AFD_PORT:-6269} AFD_PROCESS_GROUP_TIMEOUT_S=${AFD_PROCESS_GROUP_TIMEOUT_S:-7200} @@ -75,6 +79,11 @@ export VLLM_USE_V2_MODEL_RUNNER=1 NUM_ATTN_RANKS=$(tr ',' '\n' <<<"$ATTN_DEVICES" | wc -l | tr -d ' ') NUM_FFN_RANKS=$(tr ',' '\n' <<<"$FFN_DEVICES" | wc -l | tr -d ' ') +if ((NUM_ATTN_RANKS % ATTN_TP)); then + echo "ATTN_TP=${ATTN_TP} does not divide the ${NUM_ATTN_RANKS} Attention devices" >&2 + exit 1 +fi +ATTN_DP=$((NUM_ATTN_RANKS / ATTN_TP)) case "$AFD_CONNECTOR" in # DeepEP M2N routes on Attention and sends FP8 tokens to the FFN experts. @@ -123,7 +132,6 @@ JSON MAX_CUDAGRAPH_CAPTURE_SIZE=$(tr ',' '\n' <<<"$CUDAGRAPH_CAPTURE_SIZES" | sort -n | tail -n 1) COMMON_ARGS=( - --tensor-parallel-size 1 --enable-expert-parallel --enable-ep-weight-filter --max-num-seqs "$MAX_NUM_SEQS" @@ -149,7 +157,7 @@ fi # shellcheck disable=SC2206 COMMON_ARGS+=($EXTRA_ARGS) -ATTN_ARGS=(--data-parallel-size "$NUM_ATTN_RANKS") +ATTN_ARGS=(--data-parallel-size "$ATTN_DP" --tensor-parallel-size "$ATTN_TP") if [[ "$ENABLE_PREFIX_CACHING" == 1 ]]; then ATTN_ARGS+=(--enable-prefix-caching) else @@ -161,7 +169,7 @@ fi # shellcheck disable=SC2206 ATTN_ARGS+=($ATTN_EXTRA_ARGS) -FFN_ARGS=(--data-parallel-size "$NUM_FFN_RANKS") +FFN_ARGS=(--data-parallel-size "$NUM_FFN_RANKS" --tensor-parallel-size 1) # shellcheck disable=SC2206 FFN_ARGS+=($FFN_EXTRA_ARGS) From ce34f526766236d23ae38a9d8f4c437366b4bae1 Mon Sep 17 00:00:00 2001 From: ronenkat <16743404+ronenkat@users.noreply.github.com> Date: Sat, 31 Oct 2026 07:23:20 +0200 Subject: [PATCH 21/32] update run config Signed-off-by: ronenkat <16743404+ronenkat@users.noreply.github.com> --- .../prefill_decode_colocation/4a2f_graph_dp4tp1.sh | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/4a2f_graph_dp4tp1.sh b/recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/4a2f_graph_dp4tp1.sh index 334afebed..8f25b016e 100755 --- a/recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/4a2f_graph_dp4tp1.sh +++ b/recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/4a2f_graph_dp4tp1.sh @@ -18,10 +18,11 @@ # overrides KV_CONNECTOR (unset) # ENABLE_PREFIX_CACHING 1 | 0 (1) # ENABLE_DBO 1 | 0 (0) -# DBO_DECODE_TOKEN_THRESHOLD / DBO_PREFILL_TOKEN_THRESHOLD (2 / 8) +# DBO_DECODE_TOKEN_THRESHOLD / DBO_PREFILL_TOKEN_THRESHOLD (2 / 2) # MAX_NUM_SEQS per Attention DP rank (128) # MAX_NUM_BATCHED_TOKENS prefill chunk per step (3072) -# CUDAGRAPH_CAPTURE_SIZES comma-separated decode graph sizes (8,16,32,64,96,128) +# CUDAGRAPH_CAPTURE_SIZES comma-separated decode graph sizes +# (8,16,32,48,64,96,112,128) # MAX_MODEL_LEN (4096) # GPU_MEMORY_UTILIZATION (unset: vLLM default) # ATTN_DEVICES / FFN_DEVICES (0,1,2,3 / 4,5) @@ -53,10 +54,10 @@ KV_TRANSFER_CONFIG=${KV_TRANSFER_CONFIG:-} ENABLE_PREFIX_CACHING=${ENABLE_PREFIX_CACHING:-1} ENABLE_DBO=${ENABLE_DBO:-0} DBO_DECODE_TOKEN_THRESHOLD=${DBO_DECODE_TOKEN_THRESHOLD:-2} -DBO_PREFILL_TOKEN_THRESHOLD=${DBO_PREFILL_TOKEN_THRESHOLD:-8} +DBO_PREFILL_TOKEN_THRESHOLD=${DBO_PREFILL_TOKEN_THRESHOLD:-2} MAX_NUM_SEQS=${MAX_NUM_SEQS:-128} MAX_NUM_BATCHED_TOKENS=${MAX_NUM_BATCHED_TOKENS:-3072} -CUDAGRAPH_CAPTURE_SIZES=${CUDAGRAPH_CAPTURE_SIZES:-8,16,32,64,96,128} +CUDAGRAPH_CAPTURE_SIZES=${CUDAGRAPH_CAPTURE_SIZES:-8,16,32,48,64,96,112,128} MAX_MODEL_LEN=${MAX_MODEL_LEN:-4096} GPU_MEMORY_UTILIZATION=${GPU_MEMORY_UTILIZATION:-} ATTN_DEVICES=${ATTN_DEVICES:-0,1,2,3} From 9ad2999649c960d1166c2ddecf0edfca173387f7 Mon Sep 17 00:00:00 2001 From: ronenkat <16743404+ronenkat@users.noreply.github.com> Date: Sat, 31 Oct 2026 18:04:39 +0200 Subject: [PATCH 22/32] add baseline vllm (vanilla) and decode benchmark readme Signed-off-by: ronenkat <16743404+ronenkat@users.noreply.github.com> --- .../qwen3_235b_a22b_fp8/vanilla_tp6.sh | 89 +++++++++++ tools/benchmarks/README.md | 149 ++++++++++++++++++ 2 files changed, 238 insertions(+) create mode 100755 recipe/baseline/qwen3_235b_a22b_fp8/vanilla_tp6.sh create mode 100644 tools/benchmarks/README.md diff --git a/recipe/baseline/qwen3_235b_a22b_fp8/vanilla_tp6.sh b/recipe/baseline/qwen3_235b_a22b_fp8/vanilla_tp6.sh new file mode 100755 index 000000000..ea372120d --- /dev/null +++ b/recipe/baseline/qwen3_235b_a22b_fp8/vanilla_tp6.sh @@ -0,0 +1,89 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +# +# Qwen3-235B-A22B-FP8 vanilla vLLM baseline on one 6-GPU node. +# +# This is the no-AFD counterpart to +# recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/4a2f_graph_dp4tp1.sh. +# It launches a single vLLM process with tensor parallelism 6 and no AFD +# connector, AFD additional config, or separate Attention/FFN roles. +# +# Every knob is an environment variable: +# +# MODEL_PATH weights dir or HF id (/path/model_weights/Qwen3-235B-A22B-FP8) +# TENSOR_PARALLEL_SIZE tensor-parallel size (6) +# CUDA_VISIBLE_DEVICES GPUs to use (0,1,2,3,4,5) +# HOST / PORT HTTP bind address and port (127.0.0.1 / 18000) +# MAX_NUM_SEQS maximum sequences (128) +# MAX_NUM_BATCHED_TOKENS prefill chunk per step (3072) +# MAX_MODEL_LEN maximum model length (34816) +# CUDAGRAPH_CAPTURE_SIZES comma-separated decode graph sizes +# (8,16,32,48,64,96,112,128) +# ENABLE_PREFIX_CACHING 1 | 0 (1) +# ENABLE_DBO 1 | 0 (0) +# DBO_DECODE_TOKEN_THRESHOLD / DBO_PREFILL_TOKEN_THRESHOLD (2 / 2) +# GPU_MEMORY_UTILIZATION optional GPU memory utilization fraction (unset) +# EXTRA_ARGS extra `vllm serve` arguments (empty) + +set -euo pipefail + +MODEL_PATH=${MODEL_PATH:-/path/model_weights/Qwen3-235B-A22B-FP8} +TENSOR_PARALLEL_SIZE=${TENSOR_PARALLEL_SIZE:-6} +CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES:-0,1,2,3,4,5} +HOST=${HOST:-127.0.0.1} +PORT=${PORT:-18000} +MAX_NUM_SEQS=${MAX_NUM_SEQS:-128} +MAX_NUM_BATCHED_TOKENS=${MAX_NUM_BATCHED_TOKENS:-3072} +MAX_MODEL_LEN=${MAX_MODEL_LEN:-34816} +CUDAGRAPH_CAPTURE_SIZES=${CUDAGRAPH_CAPTURE_SIZES:-8,16,32,48,64,96,112,128} +ENABLE_PREFIX_CACHING=${ENABLE_PREFIX_CACHING:-1} +ENABLE_DBO=${ENABLE_DBO:-0} +DBO_DECODE_TOKEN_THRESHOLD=${DBO_DECODE_TOKEN_THRESHOLD:-2} +DBO_PREFILL_TOKEN_THRESHOLD=${DBO_PREFILL_TOKEN_THRESHOLD:-2} +GPU_MEMORY_UTILIZATION=${GPU_MEMORY_UTILIZATION:-} +EXTRA_ARGS=${EXTRA_ARGS:-} + +export VLLM_USE_V2_MODEL_RUNNER=1 +export CUDA_VISIBLE_DEVICES + +MAX_CUDAGRAPH_CAPTURE_SIZE=$(tr ',' '\n' <<<"$CUDAGRAPH_CAPTURE_SIZES" | sort -n | tail -n 1) + +ARGS=( + --tensor-parallel-size "$TENSOR_PARALLEL_SIZE" + --enable-expert-parallel + --enable-ep-weight-filter + --max-num-seqs "$MAX_NUM_SEQS" + --max-num-batched-tokens "$MAX_NUM_BATCHED_TOKENS" + --max-model-len "$MAX_MODEL_LEN" + --max-cudagraph-capture-size "$MAX_CUDAGRAPH_CAPTURE_SIZE" + --compilation-config "{\"cudagraph_mode\": \"FULL_DECODE_ONLY\", \"cudagraph_capture_sizes\": [${CUDAGRAPH_CAPTURE_SIZES}]}" + --cpu-distributed-timeout-seconds 3600 + --trust-remote-code +) + +if [[ "$ENABLE_PREFIX_CACHING" == 1 ]]; then + ARGS+=(--enable-prefix-caching) +else + ARGS+=(--no-enable-prefix-caching) +fi + +if [[ "$ENABLE_DBO" == 1 ]]; then + ARGS+=( + --enable-dbo + --dbo-decode-token-threshold "$DBO_DECODE_TOKEN_THRESHOLD" + --dbo-prefill-token-threshold "$DBO_PREFILL_TOKEN_THRESHOLD" + ) +fi + +if [[ -n "$GPU_MEMORY_UTILIZATION" ]]; then + ARGS+=(--gpu-memory-utilization "$GPU_MEMORY_UTILIZATION") +fi + +# shellcheck disable=SC2206 +ARGS+=($EXTRA_ARGS) + +exec vllm serve "$MODEL_PATH" \ + "${ARGS[@]}" \ + --host "$HOST" \ + --port "$PORT" diff --git a/tools/benchmarks/README.md b/tools/benchmarks/README.md new file mode 100644 index 000000000..d3d220ab2 --- /dev/null +++ b/tools/benchmarks/README.md @@ -0,0 +1,149 @@ +# Decode-only benchmark + +This directory contains the decode-only benchmark used to measure decode +throughput and latency over long, real KV contexts while avoiding large +prefill work during the measured phase. + +- `decode_bench.py` implements the server-side KV connector. +- `decode_bench_client.py` wraps `vllm bench serve`, performs warmup, and writes + the benchmark checks. + +## How it works + +The benchmark starts by bringing up a vLLM serving instance using one of the +recipes under `recipe/gpu/`. The recipe attaches the decode-bench KV connector +to the Attention instance and starts the server. + +After the server is ready, the benchmark has two phases: + +1. **Warmup.** The server-side connector creates a pool of seeded prompts of + `--prompt-len` tokens and prefills them once. Their KV blocks are pinned in + the prefix cache. +2. **Measured.** The client sends requests that the connector remaps onto the + pool prompts. Each measured request prefills only its last prompt token and + then decodes for the requested number of output tokens. + +This makes the benchmark decode-only while still exercising real KV-cache +attention over the full prompt length. + +## Running with an AFD recipe + +The AFD recipes already support the connector through environment variables. +For example, to run the Qwen3-235B-A22B-FP8 4A2F recipe: + +```bash +export MODEL_PATH=/path/to/Qwen3-235B-A22B-FP8 + +KV_CONNECTOR=decode_bench \ +DECODE_BENCH_NUM_PROMPTS=500 \ +DECODE_BENCH_PROMPT_LEN=1025 \ +DECODE_BENCH_MAX_OUTPUT_LEN=256 \ + bash recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/4a2f_graph_dp4tp1.sh +``` + +Then run the client from the repository root: + +```bash +python tools/benchmarks/decode_bench_client.py \ + --port 18000 \ + --dp-size 4 \ + --pool-prompts 500 \ + --prompt-len 1025 \ + --result-dir /tmp/results/decode_bench \ + -- --num-prompts 3000 \ + --random-output-len 256 \ + --max-concurrency 128 \ + --request-rate inf +``` + +Use the Attention port, not the FFN port. + +## Running with vanilla vLLM + +The client is not AFD-specific. It works with any vLLM server that has the +decode-bench KV connector attached. + +For a single vanilla vLLM process on six GPUs, start the server with tensor +parallelism 6 and the connector’s `--kv-transfer-config`: + +```bash +export PYTHONPATH=$PWD +export MODEL_PATH=/path/to/model + +vllm serve "$MODEL_PATH" \ + --tensor-parallel-size 6 \ + --trust-remote-code \ + --kv-transfer-config '{"kv_connector":"AFDDecodeBenchConnector","kv_connector_module_path":"tools.benchmarks.decode_bench","kv_role":"kv_both","kv_connector_extra_config":{"num_prompts":500,"prompt_len":1025,"seed":0,"max_output_len":256}}' +``` + +Then run the benchmark client with `--dp-size 1`, because a vanilla TP6 server +has one logical engine: + +```bash +python tools/benchmarks/decode_bench_client.py \ + --port 8000 \ + --dp-size 1 \ + --pool-prompts 500 \ + --prompt-len 1025 \ + --result-dir /tmp/results/decode_bench_vanilla \ + -- --num-prompts 3000 \ + --random-output-len 256 \ + --max-concurrency 256 \ + --request-rate inf +``` + +To repeat the run at concurrency 512, rerun the same command with +`--max-concurrency 512` and a different `--result-dir`. + +## Pool sizing + +The connector checks that the prompt pool and running requests fit in the KV +cache: + +```text +num_prompts * ceil((prompt_len + 1) / block_size) + + max_num_seqs * ceil((max_output_len + 1) / block_size) + <= free KV blocks +``` + +Increase `num_prompts` only if the server still has enough free KV blocks after +reserving space for the configured `max_num_seqs` and `max_output_len`. + +## Client options + +- `--dp-size` is the number of Attention data-parallel ranks. Use `4` for the + AFD 4A2F recipe and `1` for a single vanilla vLLM engine. +- `--pool-prompts` must match the server’s `num_prompts`. +- `--prompt-len` must match the server’s `prompt_len`. +- `--result-dir` receives the warmup and measured benchmark JSON. +- `--skip-warmup` measures only, against an already warmed pool. +- Every argument after `--` is passed unchanged to `vllm bench serve`. + +Useful `vllm bench serve` arguments include: + +- `--num-prompts` +- `--random-output-len` +- `--max-concurrency` +- `--request-rate` +- `--percentile-metrics` +- `--metric-percentiles` +- `--save-detailed` + +## Output files + +Each run writes: + +- `warmup-rank.json` and `warmup-rank.log` +- `measured.json` +- `decode_bench_report.json` + +`decode_bench_report.json` contains the connector’s checks, including whether +every measured request prefilled exactly one token and whether any requests +were preempted. + +## Notes + +- The generated text is meaningless because the prompts are random. +- Restart the server to release the pinned pool. +- The connector’s name includes `AF`, but it does not require the AFD plugin; + it uses vLLM’s standard KV-transfer interface. From 666d24951bb35e493ab7f428869b575625702c49 Mon Sep 17 00:00:00 2001 From: ronenkat <16743404+ronenkat@users.noreply.github.com> Date: Sat, 31 Oct 2026 18:16:33 +0200 Subject: [PATCH 23/32] move baseline vllm script location Signed-off-by: ronenkat <16743404+ronenkat@users.noreply.github.com> --- recipe/{ => gpu}/baseline/qwen3_235b_a22b_fp8/vanilla_tp6.sh | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename recipe/{ => gpu}/baseline/qwen3_235b_a22b_fp8/vanilla_tp6.sh (100%) diff --git a/recipe/baseline/qwen3_235b_a22b_fp8/vanilla_tp6.sh b/recipe/gpu/baseline/qwen3_235b_a22b_fp8/vanilla_tp6.sh similarity index 100% rename from recipe/baseline/qwen3_235b_a22b_fp8/vanilla_tp6.sh rename to recipe/gpu/baseline/qwen3_235b_a22b_fp8/vanilla_tp6.sh From 51fcaf37bfdf42894f9caad574c4518ec1ebbd86 Mon Sep 17 00:00:00 2001 From: ronenkat <16743404+ronenkat@users.noreply.github.com> Date: Sat, 31 Oct 2026 20:51:57 +0200 Subject: [PATCH 24/32] update the baseline to 4 GPUs Signed-off-by: ronenkat <16743404+ronenkat@users.noreply.github.com> --- .../{vanilla_tp6.sh => vanilla_tp4.sh} | 16 ++++++++++------ tools/benchmarks/README.md | 8 ++++---- 2 files changed, 14 insertions(+), 10 deletions(-) rename recipe/gpu/baseline/qwen3_235b_a22b_fp8/{vanilla_tp6.sh => vanilla_tp4.sh} (83%) diff --git a/recipe/gpu/baseline/qwen3_235b_a22b_fp8/vanilla_tp6.sh b/recipe/gpu/baseline/qwen3_235b_a22b_fp8/vanilla_tp4.sh similarity index 83% rename from recipe/gpu/baseline/qwen3_235b_a22b_fp8/vanilla_tp6.sh rename to recipe/gpu/baseline/qwen3_235b_a22b_fp8/vanilla_tp4.sh index ea372120d..7adba7f33 100755 --- a/recipe/gpu/baseline/qwen3_235b_a22b_fp8/vanilla_tp6.sh +++ b/recipe/gpu/baseline/qwen3_235b_a22b_fp8/vanilla_tp4.sh @@ -2,18 +2,18 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project # -# Qwen3-235B-A22B-FP8 vanilla vLLM baseline on one 6-GPU node. +# Qwen3-235B-A22B-FP8 vanilla vLLM baseline on one 4-GPU node. # # This is the no-AFD counterpart to # recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/4a2f_graph_dp4tp1.sh. -# It launches a single vLLM process with tensor parallelism 6 and no AFD +# It launches a single vLLM process with tensor parallelism 4 and no AFD # connector, AFD additional config, or separate Attention/FFN roles. # # Every knob is an environment variable: # # MODEL_PATH weights dir or HF id (/path/model_weights/Qwen3-235B-A22B-FP8) -# TENSOR_PARALLEL_SIZE tensor-parallel size (6) -# CUDA_VISIBLE_DEVICES GPUs to use (0,1,2,3,4,5) +# TENSOR_PARALLEL_SIZE tensor-parallel size (4) +# CUDA_VISIBLE_DEVICES GPUs to use (0,1,2,3) # HOST / PORT HTTP bind address and port (127.0.0.1 / 18000) # MAX_NUM_SEQS maximum sequences (128) # MAX_NUM_BATCHED_TOKENS prefill chunk per step (3072) @@ -23,14 +23,16 @@ # ENABLE_PREFIX_CACHING 1 | 0 (1) # ENABLE_DBO 1 | 0 (0) # DBO_DECODE_TOKEN_THRESHOLD / DBO_PREFILL_TOKEN_THRESHOLD (2 / 2) +# ALL2ALL_BACKEND MoE all-to-all backend (allgather_reducescatter); +# DBO requires DeepEP or NIXL instead # GPU_MEMORY_UTILIZATION optional GPU memory utilization fraction (unset) # EXTRA_ARGS extra `vllm serve` arguments (empty) set -euo pipefail MODEL_PATH=${MODEL_PATH:-/path/model_weights/Qwen3-235B-A22B-FP8} -TENSOR_PARALLEL_SIZE=${TENSOR_PARALLEL_SIZE:-6} -CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES:-0,1,2,3,4,5} +TENSOR_PARALLEL_SIZE=${TENSOR_PARALLEL_SIZE:-4} +CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES:-0,1,2,3} HOST=${HOST:-127.0.0.1} PORT=${PORT:-18000} MAX_NUM_SEQS=${MAX_NUM_SEQS:-128} @@ -41,6 +43,7 @@ ENABLE_PREFIX_CACHING=${ENABLE_PREFIX_CACHING:-1} ENABLE_DBO=${ENABLE_DBO:-0} DBO_DECODE_TOKEN_THRESHOLD=${DBO_DECODE_TOKEN_THRESHOLD:-2} DBO_PREFILL_TOKEN_THRESHOLD=${DBO_PREFILL_TOKEN_THRESHOLD:-2} +ALL2ALL_BACKEND=${ALL2ALL_BACKEND:-allgather_reducescatter} GPU_MEMORY_UTILIZATION=${GPU_MEMORY_UTILIZATION:-} EXTRA_ARGS=${EXTRA_ARGS:-} @@ -53,6 +56,7 @@ ARGS=( --tensor-parallel-size "$TENSOR_PARALLEL_SIZE" --enable-expert-parallel --enable-ep-weight-filter + --all2all-backend "$ALL2ALL_BACKEND" --max-num-seqs "$MAX_NUM_SEQS" --max-num-batched-tokens "$MAX_NUM_BATCHED_TOKENS" --max-model-len "$MAX_MODEL_LEN" diff --git a/tools/benchmarks/README.md b/tools/benchmarks/README.md index d3d220ab2..f08d7cfca 100644 --- a/tools/benchmarks/README.md +++ b/tools/benchmarks/README.md @@ -63,20 +63,20 @@ Use the Attention port, not the FFN port. The client is not AFD-specific. It works with any vLLM server that has the decode-bench KV connector attached. -For a single vanilla vLLM process on six GPUs, start the server with tensor -parallelism 6 and the connector’s `--kv-transfer-config`: +For a single vanilla vLLM process on four GPUs, start the server with tensor +parallelism 4 and the connector’s `--kv-transfer-config`: ```bash export PYTHONPATH=$PWD export MODEL_PATH=/path/to/model vllm serve "$MODEL_PATH" \ - --tensor-parallel-size 6 \ + --tensor-parallel-size 4 \ --trust-remote-code \ --kv-transfer-config '{"kv_connector":"AFDDecodeBenchConnector","kv_connector_module_path":"tools.benchmarks.decode_bench","kv_role":"kv_both","kv_connector_extra_config":{"num_prompts":500,"prompt_len":1025,"seed":0,"max_output_len":256}}' ``` -Then run the benchmark client with `--dp-size 1`, because a vanilla TP6 server +Then run the benchmark client with `--dp-size 1`, because a vanilla TP4 server has one logical engine: ```bash From 2cd098bc1c66ecb0bbdc867d21631c4e4f50637f Mon Sep 17 00:00:00 2001 From: ronenkat <16743404+ronenkat@users.noreply.github.com> Date: Sun, 1 Nov 2026 17:47:43 +0200 Subject: [PATCH 25/32] fix(benchmarks): pin the prompt pool with block references, map only onto pinned entries Pin the prompt pool by holding one reference on each entry's full prompt blocks (BlockPool.touch) instead of reporting delay_free_blocks for warmup requests and never completing them. Delayed-free warmup requests stayed in the scheduler for the server's lifetime, so has_requests() never turned False: the engine loop spun on empty steps whenever no request was running, and shutdown and pause could not see the engine as idle. Warmup requests are now freed as usual, and only the full prompt blocks a prefix hit shares are pinned, one block fewer per entry. Map measured requests round-robin over pinned entries only. An entry that is still warming, or whose warmup was aborted, is not in the prefix cache yet, so a request mapped onto it decoded over unwritten blocks and could leave them cached under the entry's hashes. A request that arrives before any entry is pinned falls back to the whole pool; such early requests are counted, logged, and caught by the client's decode-only check. Co-Authored-By: Claude Opus 5.5 Signed-off-by: ronenkat <16743404+ronenkat@users.noreply.github.com> --- .../qwen3_30b_a3b/README.md | 18 +- tools/benchmarks/decode_bench.py | 168 ++++++++++++------ 2 files changed, 125 insertions(+), 61 deletions(-) diff --git a/recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/README.md b/recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/README.md index 6ab09c6ee..cc843b89c 100644 --- a/recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/README.md +++ b/recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/README.md @@ -99,15 +99,17 @@ For an accuracy check, run the matching E2E case, for example requests into a pool of seeded random prompts of `DECODE_BENCH_PROMPT_LEN` tokens (the same pool on every rank, whatever the requests carried) and prefills them for real, one output token each. -2. **Pinning.** When a warmup request finishes, the connector keeps its KV - blocks allocated for the server's lifetime (it reports the blocks as still - being sent and never completes the send). Pinned blocks cannot be evicted - and stay in the prefix cache. +2. **Pinning.** When a warmup request finishes, the connector takes a + reference of its own on the request's full prompt blocks and holds it for + the server's lifetime, while the request itself is freed as usual. Pinned + blocks cannot be evicted and stay in the prefix cache. 3. **Measured requests.** Every later request is remapped round-robin onto a - pool prompt, whatever it asked for. The prefix cache covers every prompt - token but the last (the prompt length is rounded to one token past a block - boundary), so the request prefills one token and decodes over real KV, on - whichever rank serves it. + pinned pool prompt, whatever it asked for. The prefix cache covers every + prompt token but the last (the prompt length is rounded to one token past a + block boundary), so the request prefills one token and decodes over real KV, + on whichever rank serves it. A request that arrives before any prompt is + pinned may miss the cache; the connector logs it and the client's checks + fail the run. Start the server, then drive it with `tools/benchmarks/decode_bench_client.py`, which runs every phase with `vllm bench serve`: diff --git a/tools/benchmarks/decode_bench.py b/tools/benchmarks/decode_bench.py index 3ff08a261..3f00d839f 100644 --- a/tools/benchmarks/decode_bench.py +++ b/tools/benchmarks/decode_bench.py @@ -9,9 +9,9 @@ - Prompt pool (``num_prompts`` > 0): each Attention DP rank prefills a pool of ``num_prompts`` seeded random prompts for real, from its first ``num_prompts`` requests (the warmup), and pins their KV blocks in the prefix - cache. Every later request is remapped, whatever its content, onto a pool - prompt round-robin, hits the prefix cache for all but the last prompt token, - and goes straight to decode over real KV. + cache. Every later request is remapped, whatever its content, onto a pinned + pool prompt round-robin, hits the prefix cache for all but the last prompt + token, and goes straight to decode over real KV. """ from __future__ import annotations @@ -34,6 +34,7 @@ if TYPE_CHECKING: from vllm.config import VllmConfig from vllm.forward_context import ForwardContext + from vllm.v1.core.block_pool import BlockPool from vllm.v1.core.kv_cache_manager import KVCacheBlocks, KVCacheManager from vllm.v1.core.kv_cache_utils import BlockHash from vllm.v1.core.sched.output import SchedulerOutput @@ -46,8 +47,7 @@ KVCacheLayer: TypeAlias = torch.Tensor | tuple[torch.Tensor, ...] -# A warmup request only has to prefill its pool prompt; one output token also -# fixes how many blocks each pinned pool entry holds. +# A warmup request only has to prefill its pool prompt. WARMUP_OUTPUT_TOKENS = 1 # Finished measured requests between two prompt-pool hit/miss log lines. POOL_STATS_LOG_INTERVAL = 1000 @@ -190,11 +190,10 @@ def request_finished( block_ids: list[int], ) -> tuple[bool, dict[str, Any] | None]: assert self.connector_scheduler is not None - delay_free_blocks = self.connector_scheduler.request_finished( - request, - len(block_ids), - ) - return delay_free_blocks, None + self.connector_scheduler.request_finished(request, block_ids) + # Never delay the free: the prompt pool pins its blocks with references + # of its own, so every request is released as soon as it finishes. + return False, None def request_finished_all_groups( self, @@ -202,11 +201,12 @@ def request_finished_all_groups( block_ids: tuple[list[int], ...], ) -> tuple[bool, dict[str, Any] | None]: assert self.connector_scheduler is not None - delay_free_blocks = self.connector_scheduler.request_finished( + # The prompt pool supports a single KV cache group (checked at startup). + self.connector_scheduler.request_finished( request, - sum(len(group_block_ids) for group_block_ids in block_ids), + block_ids[0] if block_ids else [], ) - return delay_free_blocks, None + return False, None class AFDDecodeBenchConnectorScheduler: @@ -246,9 +246,7 @@ def __init__( def bind_kv_cache_manager(self, kv_cache_manager: KVCacheManager) -> None: if self.prompt_pool is not None: - self.prompt_pool.check_capacity( - kv_cache_manager.block_pool.get_num_free_blocks() - ) + self.prompt_pool.bind_block_pool(kv_cache_manager.block_pool) def on_new_request(self, request: Request) -> None: if self.prompt_pool is not None: @@ -319,23 +317,28 @@ def build_connector_meta( self._pending_fills.clear() return meta - def request_finished(self, request: Request, num_blocks: int) -> bool: - """Return True to keep the request's blocks allocated (pinned).""" + def request_finished(self, request: Request, block_ids: list[int]) -> None: self._filled_requests.discard(request.request_id) - if self.prompt_pool is None: - return False - return self.prompt_pool.request_finished(request, num_blocks) + if self.prompt_pool is not None: + self.prompt_pool.request_finished(request, block_ids) class PromptPool: """Per-rank pool of real-KV prompts that measured requests decode over. - Pinning: a warmup request that prefilled its pool prompt reports - ``delay_free_blocks`` from ``request_finished`` and is never returned from - ``get_finished``, so the scheduler keeps its blocks allocated for the - server's lifetime. Pinned blocks never enter the free queue, so LRU - eviction cannot drop a pool entry, and they stay in the prefix cache, so a - measured request that maps to the entry shares them like any prefix hit. + Pinning: when a warmup request that prefilled its pool prompt finishes, + the pool takes one reference of its own on the entry's full prompt blocks + (``BlockPool.touch``) and holds it for the server's lifetime; the scheduler + then frees the request as usual. The extra reference keeps the blocks out + of the free queue, so LRU eviction cannot drop a pool entry, and they stay + in the prefix cache, so a measured request that maps to the entry shares + them like any prefix hit. Unlike a delayed free, this leaves no finished + request behind in the scheduler, so the engine still goes idle and shuts + down normally. + + Measured requests map only onto pinned entries. One that arrives before + any entry is pinned falls back to the whole pool; it may miss the prefix + cache, so it is counted and logged, and the client's checks catch it. """ def __init__( @@ -357,6 +360,7 @@ def __init__( self.block_size = vllm_config.cache_config.block_size self.max_num_seqs = vllm_config.scheduler_config.max_num_seqs + self._block_pool: BlockPool | None = None max_model_len = vllm_config.model_config.max_model_len self.num_prompts = config.num_prompts # A prefix-cache hit covers full blocks only and stops one token short @@ -370,6 +374,9 @@ def __init__( f"prompt_len ({config.prompt_len}) must exceed the KV block size " f"({self.block_size})" ) + # The full prompt blocks a hit shares, and so the blocks a pinned entry + # holds; the block with the last prompt token is never shared. + self.num_entry_blocks = self.prompt_len // self.block_size self.max_output_len = ( config.max_output_len if config.max_output_len is not None @@ -396,21 +403,26 @@ def __init__( self._entry_block_hashes: dict[int, list[BlockHash]] = {} self._unwarmed_entries = deque(range(self.num_prompts)) self._warming_entries: dict[str, int] = {} - self._num_pinned_entries = 0 - self._num_pinned_blocks = 0 + # Pinned entries in pin order, and the block ids each one holds. + self._pinned_entries: list[int] = [] + self._pinned_block_ids: dict[int, list[int]] = {} self._num_assigned_measured = 0 + # Measured requests assigned while part of the pool was still unpinned. + self._num_early_measured = 0 # request_id -> prompt tokens the prefix cache missed (0 on a hit). self._measured_misses: dict[str, int] = {} self._num_finished_measured = 0 self._num_missed_requests = 0 self._num_missed_tokens = 0 + def bind_block_pool(self, block_pool: BlockPool) -> None: + """Bind the scheduler's block pool, whose blocks the pool pins.""" + self.check_capacity(block_pool.get_num_free_blocks()) + self._block_pool = block_pool + def check_capacity(self, num_free_blocks: int) -> None: """Fail fast when the pool and the running requests cannot both fit.""" - blocks_per_entry = cdiv( - self.prompt_len + WARMUP_OUTPUT_TOKENS, - self.block_size, - ) + blocks_per_entry = self.num_entry_blocks # A measured request shares its entry's full blocks and allocates only # the last prompt token and its output. blocks_per_seq = cdiv(1 + self.max_output_len, self.block_size) @@ -450,15 +462,15 @@ def assign(self, request: Request) -> None: """Rewrite a new request's prompt to a pool entry's, before scheduling. The first ``num_prompts`` requests warm one unwarmed entry each; every - later request maps round-robin onto the pool, whatever it asked for. + later request maps round-robin onto the pinned entries, whatever it + asked for. """ if self._unwarmed_entries: entry = self._unwarmed_entries.popleft() self._warming_entries[request.request_id] = entry request.max_tokens = WARMUP_OUTPUT_TOKENS else: - entry = self._num_assigned_measured % self.num_prompts - self._num_assigned_measured += 1 + entry = self._next_measured_entry() prompt_token_ids = self.prompt_tokens[entry].tolist() request.prompt_token_ids = prompt_token_ids @@ -473,6 +485,31 @@ def assign(self, request: Request) -> None: else: request.block_hashes.extend(entry_block_hashes) + def _next_measured_entry(self) -> int: + """Pick a measured request's entry: round-robin over pinned entries. + + An entry that is still warming, or whose warmup was aborted, is not + in the prefix cache yet: a request mapped onto it would decode over + unwritten blocks and could leave them cached under the entry's hashes. + """ + index = self._num_assigned_measured + self._num_assigned_measured += 1 + num_pinned = len(self._pinned_entries) + if num_pinned < self.num_prompts: + if self._num_early_measured == 0: + logger.warning( + "AFDDecodeBenchConnector: a measured request arrived with " + "%d of %d pool prompts pinned; finish the warmup before " + "measuring", + num_pinned, + self.num_prompts, + ) + self._num_early_measured += 1 + if num_pinned == 0: + # Nothing is pinned yet: any entry may miss, and the miss is faked. + return index % self.num_prompts + return self._pinned_entries[index % num_pinned] + def num_tokens_to_fill(self, request: Request, num_computed_tokens: int) -> int: """Prompt tokens to fake as external: 0 on a hit or during warmup.""" if request.request_id in self._warming_entries: @@ -483,28 +520,24 @@ def num_tokens_to_fill(self, request: Request, num_computed_tokens: int) -> int: self._measured_misses[request.request_id] = num_missed_tokens return num_missed_tokens - def request_finished(self, request: Request, num_blocks: int) -> bool: - """Pin a finished warmup request's blocks; tally measured hits.""" + def request_finished(self, request: Request, block_ids: list[int]) -> None: + """Pin a finished warmup request's blocks; tally measured hits. + + Called before the scheduler frees the request, so a pinned block's + reference count never drops to zero. + """ entry = self._warming_entries.pop(request.request_id, None) if entry is not None: - if request.status not in PINNABLE_WARMUP_STATUSES: + if request.status not in PINNABLE_WARMUP_STATUSES or not self._pin( + entry, block_ids + ): # Aborted before its prompt was cached: warm it again later. self._unwarmed_entries.appendleft(entry) - return False - self._num_pinned_entries += 1 - self._num_pinned_blocks += num_blocks - if self._num_pinned_entries == self.num_prompts: - logger.info( - "AFDDecodeBenchConnector prompt pool warm: %d prompts " - "pinned in %d KV blocks", - self._num_pinned_entries, - self._num_pinned_blocks, - ) - return True + return num_missed_tokens = self._measured_misses.pop(request.request_id, None) if num_missed_tokens is None: - return False + return self._num_finished_measured += 1 if num_missed_tokens > 0: self._num_missed_requests += 1 @@ -512,12 +545,41 @@ def request_finished(self, request: Request, num_blocks: int) -> bool: if self._num_finished_measured % POOL_STATS_LOG_INTERVAL == 0: logger.info( "AFDDecodeBenchConnector prompt pool: %d measured requests, " - "%d prefix-cache misses (%d prompt tokens faked)", + "%d prefix-cache misses (%d prompt tokens faked), %d assigned " + "before the pool was warm", self._num_finished_measured, self._num_missed_requests, self._num_missed_tokens, + self._num_early_measured, + ) + + def _pin(self, entry: int, block_ids: list[int]) -> bool: + """Hold one reference on an entry's cached full prompt blocks.""" + assert self._block_pool is not None, "the KV cache manager is not bound" + entry_block_ids = block_ids[: self.num_entry_blocks] + blocks = [self._block_pool.blocks[block_id] for block_id in entry_block_ids] + if len(blocks) < self.num_entry_blocks or any( + block.is_null or block.block_hash is None for block in blocks + ): + logger.warning( + "AFDDecodeBenchConnector: pool prompt %d finished warming with " + "%d of %d prompt blocks cached; warming it again", + entry, + sum(block.block_hash is not None for block in blocks), + self.num_entry_blocks, + ) + return False + self._block_pool.touch(blocks) + self._pinned_entries.append(entry) + self._pinned_block_ids[entry] = entry_block_ids + if len(self._pinned_entries) == self.num_prompts: + logger.info( + "AFDDecodeBenchConnector prompt pool warm: %d prompts " + "pinned in %d KV blocks", + len(self._pinned_entries), + self.num_prompts * self.num_entry_blocks, ) - return False + return True class AFDDecodeBenchConnectorWorker: From cfe4175c3c86868818c9fe5d56d0fca362ea8045 Mon Sep 17 00:00:00 2001 From: ronenkat <16743404+ronenkat@users.noreply.github.com> Date: Mon, 2 Nov 2026 15:35:58 +0200 Subject: [PATCH 26/32] [Feat] DeepEP M2N: flexible Attention/FFN placement across pods The M2N connector accepted only placements whose NVLink domains held FFN-first, contiguous blocks of the role world, so pods that each hold FFN and Attention ranks (2 x 4A1F) were rejected. The handshake now plans a transport order from the ranks' NVLink domains: domain by domain, FFN ranks first in each. When the FFN ranks fill the first domains the order is the role world and nothing changes. When every domain holds the same number of FFN ranks, the lanes use a gloo group in transport order and afdEP-comm's per-domain mode (ffn_ranks_per_domain). The role world still drives the handshake, model checks, control plane and expert ownership. - layout.py: RankPlacement, TransportLayout, plan_transport_layout - connector.py: transport layout in HandshakeResult, transport group, ffn_ranks_per_domain on each lane, network plan in transport order - pin afdEP-comm to the ffn-domains-hybrid head (fdff967) - recipe: Qwen3-235B-A22B-FP8 2 x 4A1F multi-pod recipe - E2E: Qwen3 multi-pod 1A1F,1A1F M2N eager/graph/DBO cases - docs and deploy-afd-k8s skill: placement rule, POD=AFD_ Co-Authored-By: Claude Opus 5.5 Signed-off-by: ronenkat <16743404+ronenkat@users.noreply.github.com> --- .agents/skills/deploy-afd-k8s/SKILL.md | 4 +- .../deploy-afd-k8s/references/deploy.md | 15 ++ .../references/resolve-recipe.md | 22 ++- .../deploy-afd-k8s/references/services.md | 4 + .../connectors/gpu/deepep_m2n/connector.py | 81 ++++++++-- .../connectors/gpu/deepep_m2n/layout.py | 146 +++++++++++++++++- .../connectors/gpu/deepep_m2n/network.py | 7 +- docs/design/module/connector_contracts.md | 2 +- docs/gpu/DEEPEP_M2N_CONNECTOR_USER_GUIDE.md | 84 +++++++++- pyproject.toml | 2 +- recipe/gpu/DeepEPM2NAFDConnector/README.md | 13 +- .../qwen3_235b_a22b_fp8/README.md | 41 ++++- .../multipod_4a1f_4a1f_graph.sh | 82 ++++++++++ tests/e2e/README.md | 6 +- .../qwen3_moe/test_qwen3_moe_multi_pod.py | 22 ++- .../connectors/test_deepep_m2n_connector.py | 97 ++++++++++-- .../unit/connectors/test_deepep_m2n_layout.py | 136 ++++++++++++++++ tests/unit/fakes/afdep.py | 10 +- uv.lock | 4 +- 19 files changed, 715 insertions(+), 63 deletions(-) create mode 100644 recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/multipod_4a1f_4a1f_graph.sh create mode 100644 tests/unit/connectors/test_deepep_m2n_layout.py diff --git a/.agents/skills/deploy-afd-k8s/SKILL.md b/.agents/skills/deploy-afd-k8s/SKILL.md index f51ae4190..e33170e90 100644 --- a/.agents/skills/deploy-afd-k8s/SKILL.md +++ b/.agents/skills/deploy-afd-k8s/SKILL.md @@ -29,7 +29,9 @@ than one pod) -- callers never talk to these directly. See ## Scope -- Recipes under `recipe/gpu/P2pNcclAFDConnector/**` (GPU only, not NPU), +- Recipes under `recipe/gpu/P2pNcclAFDConnector/**` or + `recipe/gpu/DeepEPM2NAFDConnector/**` (GPU only, not NPU; the M2N image + must contain afdEP-comm and DeepGEMM), **`prefill_decode_colocation` only** (attention + FFN workers, no prefill split). `prefill_decode_disaggregation` recipes are **out of scope**: they need a diff --git a/.agents/skills/deploy-afd-k8s/references/deploy.md b/.agents/skills/deploy-afd-k8s/references/deploy.md index eccd5423a..1957efadd 100644 --- a/.agents/skills/deploy-afd-k8s/references/deploy.md +++ b/.agents/skills/deploy-afd-k8s/references/deploy.md @@ -28,6 +28,14 @@ devices="$(role_block "$RECIPE_SCRIPT_PATH" "$ROLE" \ GPU_COUNT=$(( $(grep -o ',' <<<"$devices" | wc -l) + 1 )) ``` +`POD=AFD_` entries (both roles per pod, see +[resolve-recipe.md](resolve-recipe.md)): every pod has the same shape, + +```bash +default() { grep -oE "$1:-[0-9]+" "$RECIPE_SCRIPT_PATH" | grep -oE '[0-9]+$'; } +GPU_COUNT=$(( ${ATTN_PER_POD:-$(default ATTN_PER_POD)} + ${FFN_PER_POD:-$(default FFN_PER_POD)} )) +``` + ## 3. Confirm before touching the cluster State and confirm: target cluster/namespace, `AFD_PLUGIN_IMAGE`, `MODEL_ID`, @@ -168,6 +176,13 @@ deploy_pod vllm-pod-attention-1 2 worker none "POD=ATTENTION_1 ATTENTION_HEADLES deploy_pod vllm-pod-ffn-0 2 none head "POD=FFN_0" ``` +`POD=AFD_` recipe, e.g. the DeepEP M2N 2 x 4A1F plan (5 GPUs per pod): + +```bash +deploy_pod vllm-pod-afd-0 5 head head "POD=AFD_0" +deploy_pod vllm-pod-afd-1 5 worker worker "POD=AFD_1" +``` + ## 4f. Wait for Running, then for each pod's own readiness line Bounded 15-minute wait per pod for `phase=Running` (unschedulable pods never diff --git a/.agents/skills/deploy-afd-k8s/references/resolve-recipe.md b/.agents/skills/deploy-afd-k8s/references/resolve-recipe.md index 84e4ac2a9..cc27b69d9 100644 --- a/.agents/skills/deploy-afd-k8s/references/resolve-recipe.md +++ b/.agents/skills/deploy-afd-k8s/references/resolve-recipe.md @@ -1,7 +1,8 @@ # 1. Resolve the recipe, its model, and its placement Take the recipe path from the user, or ask. It must be a colocation recipe -under `recipe/gpu/P2pNcclAFDConnector/**/prefill_decode_colocation/` (GPU +under `recipe/gpu/P2pNcclAFDConnector/**/prefill_decode_colocation/` or +`recipe/gpu/DeepEPM2NAFDConnector/**/prefill_decode_colocation/` (GPU only; prefill/decode disaggregation recipes are out of scope -- see `SKILL.md` Scope). @@ -88,3 +89,22 @@ For each entry, derive (used by [deploy.md](deploy.md) and which case `worker` if that line's `ATTENTION_HEADLESS` override is `1`, else `head`. - **`afd-ffn-node-role`**: symmetric, `none` unless `POD` starts with `FFN_`. + +### Pods that run both roles: `POD=AFD_` + +DeepEP M2N multi-pod recipes (e.g. +`DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/multipod_4a1f_4a1f_graph.sh`) +list `POD=AFD_0`, `POD=AFD_1`, ... with no other per-pod overrides. Every +such pod runs one FFN and one Attention server and derives its DP flags +(headless, start rank) from its own index, so: + +- **`afd-attn-node-role`** and **`afd-ffn-node-role`**: both `head` for + `AFD_0`, both `worker` for every other `AFD_`. +- The pod shape comes from the script's `NUM_PODS`, `ATTN_PER_POD` and + `FFN_PER_POD` defaults (`NUM_PODS` must equal the number of entries). If + the user overrides them, pass the same values to every pod as container + env. +- Ports are literals in the script, not `*_PORT:-N` defaults: the AFD + rendezvous port is the `"port": N` of the `afd` JSON, and the FFN and + Attention DP RPC ports are the last argument of the `role_dp "$F" ...` and + `role_dp "$A" ...` lines. diff --git a/.agents/skills/deploy-afd-k8s/references/services.md b/.agents/skills/deploy-afd-k8s/references/services.md index 3aa423a8d..761d23e65 100644 --- a/.agents/skills/deploy-afd-k8s/references/services.md +++ b/.agents/skills/deploy-afd-k8s/references/services.md @@ -39,6 +39,10 @@ AFD_CONNECTOR_PORT="$AFD_CONNECTOR_PORT" envsubst '${AFD_CONNECTOR_PORT}' \ Only when **more than one** plan entry's `POD` carries the same role (i.e. that role is split across pods, not just placed in its own single pod). +Every `POD=AFD_` entry carries both roles, so a plan of two or more +such entries needs both Services; read their ports (and the rendezvous port) +as [resolve-recipe.md](resolve-recipe.md) describes for those recipes rather +than with the `*_PORT:-N` greps below. Carries that role's DP-RPC coordination traffic to its head pod. One per split role, generic over role via `templates/service-dp-coord.yaml`: diff --git a/afd_plugin/connectors/gpu/deepep_m2n/connector.py b/afd_plugin/connectors/gpu/deepep_m2n/connector.py index badab047a..34c65b748 100644 --- a/afd_plugin/connectors/gpu/deepep_m2n/connector.py +++ b/afd_plugin/connectors/gpu/deepep_m2n/connector.py @@ -9,12 +9,22 @@ weighted results back, where combine sums each token's top-k contributions. Topology: - One AFD world ordered ``[F0, F1, ..., A0, A1, ...]`` (as the P2P - connector). DeepEP sees it as a symmetric EP group in which FFN rank ``g`` + One AFD role world ordered ``[F0, F1, ..., A0, A1, ...]`` (as the P2P + connector) for the handshake, the model checks and the control plane. + DeepEP sees the ranks as a symmetric EP group in which FFN rank ``g`` owns real experts ``[g * E/F, (g + 1) * E/F)`` and Attention ranks own dummy expert slots (see ``layout.py``). Any Attention:FFN ratio works; ``num_experts`` must be divisible by the FFN rank count. +Placement: + Every NVLink domain (on HGX, every pod) holds the same number of ranks, + and either the FFN ranks fill the first domains or every domain holds the + same number of FFN ranks, with FFN DP ranks rising in pod order (2 x 4A1F + is two pods of one FFN and four Attention ranks). The handshake derives + the transport's physical rank order from the ranks' NVLink domains + (``plan_transport_layout``) and picks afdEP-comm's FFN-prefix or + per-domain mode; nothing is configured. + Data path (per MoE layer and stage, identical order on every rank): Attention: ``send_attn_output(h, ctx, topk_ids=, topk_weights=)`` dispatches; ``recv_ffn_output(ref, ubatch_idx)`` combines and returns @@ -88,7 +98,12 @@ DeepEPM2NExperts, ExpertBackend, ) -from afd_plugin.connectors.gpu.deepep_m2n.layout import DeepEPM2NLayout +from afd_plugin.connectors.gpu.deepep_m2n.layout import ( + DeepEPM2NLayout, + RankPlacement, + TransportLayout, + plan_transport_layout, +) from afd_plugin.connectors.gpu.deepep_m2n.network import ( NetworkConfig, NetworkPlan, @@ -143,6 +158,9 @@ # CPU (gloo) group for the handshake and DeepEP's communicator bootstrap, # created before any NCCL communicator so the network plan can still set NCCL_*. BOOTSTRAP_GROUP_NAME = "afd_deepep_m2n_bootstrap" +# CPU (gloo) group in transport order for the lanes, created only when that +# order differs from the role world (per-domain FFN placement). +TRANSPORT_GROUP_NAME = "afd_deepep_m2n_transport" # FastAFD passes 120 s for both the CPU wait and the GPU spin timeouts. DEFAULT_TIMEOUT_S = 120 # Only Attention rank 0 sends the per-step control payload. @@ -433,6 +451,7 @@ class HandshakeResult: attn_tp_size: int capacity_tokens_per_rank: int bucket_ladder: tuple[int, ...] + transport: TransportLayout def expert_weight_format(quant_config: QuantizationConfig | None) -> str: @@ -626,9 +645,26 @@ def validate_handshake( raise RuntimeError( f"Attention ranks disagree on bucket_ladder: {sorted(map(str, ladders))}" ) - # Every rank evaluates the same plan on the same list. + # Every rank evaluates the same plans on the same list. The network plan + # sees the ranks in the transport order afdEP-comm will see. + topologies = [RankTopology.from_mapping(info["topology"]) for info in infos] + transport = plan_transport_layout( + [ + RankPlacement(info["role"], info["role_rank"], topology.nvlink_domain) + for info, topology in zip(infos, topologies, strict=True) + ], + layout.ffn_size, + ) plan = plan_network( - [RankTopology.from_mapping(info["topology"]) for info in infos], network + [ + RankTopology( + world_rank=transport.transport_rank(info["role"], info["role_rank"]), + host=topology.host, + nvlink_domain=topology.nvlink_domain, + ) + for info, topology in zip(infos, topologies, strict=True) + ], + network, ) return HandshakeResult( expert_alignment=int(alignments.pop()), @@ -636,6 +672,7 @@ def validate_handshake( attn_tp_size=attn_tp_size, capacity_tokens_per_rank=int(capacities.pop()), bucket_ladder=ladders.pop() or (), + transport=transport, ) @@ -962,6 +999,9 @@ def __init__( self._timeout = timedelta(seconds=afd_config.afd_process_group_timeout_s) self.group: dist.ProcessGroup | None = None self.bootstrap_group: dist.ProcessGroup | None = None + # The lanes' group: the bootstrap group, or one in transport order. + self.transport_group: dist.ProcessGroup | None = None + self.transport: TransportLayout | None = None self.network_plan: NetworkPlan | None = None self.lanes: list[ElasticBuffer] = [] self.lane_streams: list[torch.cuda.Stream] = [] @@ -1041,8 +1081,25 @@ def _handshake(self) -> None: self.attn_tp_size = result.attn_tp_size self.capacity_tokens_per_rank = result.capacity_tokens_per_rank self.bucket_ladder = result.bucket_ladder + self.transport = result.transport + if self.transport.is_identity: + self.transport_group = self.bootstrap_group + else: + # Transport rank 0 is FFN rank 0, the process that hosts the + # rendezvous store, so this group joins the same store. + self.transport_group = init_afd_process_group( + backend="gloo", + init_method=self._rendezvous, + world_size=self.layout.union_size, + rank=self.transport.transport_rank( + self.afd_config.role, self.role_rank + ), + group_name=TRANSPORT_GROUP_NAME, + timeout=self._timeout, + ) if self.world_rank == 0: logger.info("DeepEP M2N network plan: %s", self.network_plan.summary()) + logger.info("DeepEP M2N transport: %s", self.transport.summary()) # DeepEP reads both when each buffer is created: the GIN switch, and # with several lanes a communication stream per buffer. apply_gin_environment(self.network_plan) @@ -1117,10 +1174,11 @@ def init_afd_connector( def _new_lane(self) -> ElasticBuffer: """One lane's afdEP-comm transport.""" - assert self.bootstrap_group is not None and self.network_plan is not None + assert self.transport_group is not None and self.transport is not None + assert self.network_plan is not None assert self.capacity_tokens_per_rank is not None return ElasticBuffer( - group=self.bootstrap_group, + group=self.transport_group, capacity_tokens_per_rank=self.capacity_tokens_per_rank, hidden_size=self.hidden_size, num_experts=self.layout.deepep_num_experts, @@ -1132,8 +1190,10 @@ def _new_lane(self) -> ElasticBuffer: rdma_traffic_class=self.network_plan.rdma_traffic_class, cpu_timeout_s=self.config.cpu_timeout_s, gpu_timeout_s=self.config.gpu_timeout_s, - # The DeepEP M2N layout puts FFN ranks first, as afdEP-comm expects. + # FFN ranks are transport ranks [0, F) (FFN prefix), or the first + # ``ffn_ranks_per_domain`` of every NVLink domain (per-domain mode). ffn_ranks=self.ffn_size, + ffn_ranks_per_domain=self.transport.ffn_ranks_per_domain, ) def close(self) -> None: @@ -1157,8 +1217,9 @@ def is_initialized(self) -> bool: @property def domain(self) -> tuple[int, int]: - """DeepEP ``(scale-out ranks, scale-up ranks)``: nodes and GPUs per node - in hybrid mode, ``(1, world)`` when everything is in one NVLink domain.""" + """DeepEP ``(scale-out ranks, scale-up ranks)``: NVLink domains and + ranks per domain in hybrid mode (2 x 4A1F: ``(2, 5)``), ``(1, world)`` + when everything is in one NVLink domain.""" self._check_initialized() return self.lanes[0].num_scaleout_ranks, self.lanes[0].num_scaleup_ranks diff --git a/afd_plugin/connectors/gpu/deepep_m2n/layout.py b/afd_plugin/connectors/gpu/deepep_m2n/layout.py index 24d1bde89..d8bb225d2 100644 --- a/afd_plugin/connectors/gpu/deepep_m2n/layout.py +++ b/afd_plugin/connectors/gpu/deepep_m2n/layout.py @@ -7,16 +7,28 @@ DeepEP treats the union of Attention and FFN ranks as one symmetric EP group in which every rank owns ``experts_per_rank`` experts. FastAFD orders the union ``[Attention.., FFN..]`` and remaps real expert ids by an offset. The -plugin orders AFD worlds FFN first (``[F0.., A0..]``, as the P2P connector -does), so FFN rank ``g`` owns DeepEP experts ``[g * E_local, (g + 1) * E_local)``, -which are exactly the real expert ids: no remap. Attention ranks own dummy -expert slots after the real ones that routing never selects. +plugin orders the AFD role world FFN first (``[F0.., A0..]``, as the P2P +connector does), so FFN rank ``g`` owns real experts +``[g * E_local, (g + 1) * E_local)``. Attention ranks own dummy expert slots +that routing never selects. + +The transport world (``plan_transport_layout``) holds the same ranks in +physical order: NVLink domain by domain, as DeepEP numbers ranks, with each +domain's FFN ranks first. When the FFN ranks fill the first domains, it equals +the role world and DeepEP expert ids are the real ones (afdEP-comm's FFN +prefix). When every domain holds ``f`` FFN ranks, afdEP-comm's per-domain mode +places FFN rank ``j`` at transport rank ``(j // f) * U + j % f`` (``U`` ranks +per domain) and translates the real expert ids itself. """ from __future__ import annotations +from collections.abc import Sequence from dataclasses import dataclass +FFN_ROLE = "ffn" +ATTENTION_ROLE = "attention" + @dataclass(frozen=True) class DeepEPM2NLayout: @@ -50,10 +62,11 @@ def deepep_num_experts(self) -> int: return self.union_size * self.experts_per_rank def world_rank(self, role: str, role_rank: int) -> int: - if role == "ffn": + """Rank in the role world ``[F0.., A0..]``.""" + if role == FFN_ROLE: _check_rank("ffn role rank", role_rank, self.ffn_size) return role_rank - if role == "attention": + if role == ATTENTION_ROLE: _check_rank("attention role rank", role_rank, self.attn_size) return self.ffn_size + role_rank raise ValueError(f"unknown AFD role {role!r}") @@ -68,3 +81,124 @@ def expert_range(self, ffn_rank: int) -> range: def _check_rank(name: str, rank: int, size: int) -> None: if not 0 <= rank < size: raise ValueError(f"{name} must be in [0, {size}), got {rank}") + + +@dataclass(frozen=True) +class RankPlacement: + """Where one rank runs, as it reports in the handshake.""" + + role: str + role_rank: int + nvlink_domain: str + + +@dataclass(frozen=True) +class TransportLayout: + """The physical rank order afdEP-comm's transport sees. + + Attributes: + order: ``(role, role_rank)`` of each transport rank. + ffn_ranks_per_domain: afdEP-comm's ``ffn_ranks_per_domain``: 0 when the + FFN ranks are transport ranks ``[0, F)`` (FFN prefix), else the FFN + rank count of every NVLink domain (per-domain mode). + num_domains: NVLink domains. + domain_size: ranks per NVLink domain. + shape: ``AF`` per domain in transport order, for logs and errors. + """ + + order: tuple[tuple[str, int], ...] + ffn_ranks_per_domain: int + num_domains: int + domain_size: int + shape: str + + def transport_rank(self, role: str, role_rank: int) -> int: + return self.order.index((role, role_rank)) + + @property + def is_identity(self) -> bool: + """The transport order is the role world ``[F0.., A0..]``.""" + ffn_size = sum(role == FFN_ROLE for role, _ in self.order) + return all( + rank == (role_rank if role == FFN_ROLE else ffn_size + role_rank) + for rank, (role, role_rank) in enumerate(self.order) + ) + + @property + def mode(self) -> str: + return "per_domain" if self.ffn_ranks_per_domain else "ffn_prefix" + + def summary(self) -> str: + per_domain = ( + f" f={self.ffn_ranks_per_domain}" if self.ffn_ranks_per_domain else "" + ) + return ( + f"{self.mode}{per_domain} " + f"domains={self.num_domains}x{self.domain_size} placement={self.shape}" + ) + + +def plan_transport_layout( + placements: Sequence[RankPlacement], ffn_size: int +) -> TransportLayout: + """Order the ranks for afdEP-comm; every rank computes the same layout. + + Each NVLink domain is a contiguous block (DeepEP numbers ranks domain by + domain) with its FFN ranks first. Domains follow their lowest role-world + rank, so the domains holding FFN ranks come first. afdEP-comm takes two + shapes: the FFN ranks fill the first domains (FFN prefix: one domain, or + role-split pods), or every domain holds the same number of FFN ranks, in + pod order (per-domain mode). Every domain must hold the same number of + ranks. + """ + + def role_world_rank(placement: RankPlacement) -> int: + if placement.role == FFN_ROLE: + return placement.role_rank + return ffn_size + placement.role_rank + + domains: dict[str, list[RankPlacement]] = {} + for placement in placements: + domains.setdefault(placement.nvlink_domain, []).append(placement) + ordered = sorted( + domains.values(), key=lambda ranks: min(map(role_world_rank, ranks)) + ) + order: list[tuple[str, int]] = [] + shapes: list[str] = [] + ffn_counts: set[int] = set() + for ranks in ordered: + ffn = sorted(p.role_rank for p in ranks if p.role == FFN_ROLE) + attention = sorted(p.role_rank for p in ranks if p.role != FFN_ROLE) + order += [(FFN_ROLE, rank) for rank in ffn] + order += [(ATTENTION_ROLE, rank) for rank in attention] + shapes.append(f"{len(attention)}A{len(ffn)}F") + ffn_counts.add(len(ffn)) + shape = ",".join(shapes) + domain_sizes = {len(ranks) for ranks in ordered} + if len(domain_sizes) != 1: + raise RuntimeError( + f"NVLink domains hold {shape}: DeepEP needs the same number of ranks " + "in every NVLink domain (pod)" + ) + ffn_order = [rank for role, rank in order if role == FFN_ROLE] + if ffn_order != list(range(ffn_size)): + raise RuntimeError( + f"NVLink domains ({shape}) hold FFN ranks in the order {ffn_order}: " + "FFN DP ranks must rise in pod order (--data-parallel-start-rank)" + ) + if order[:ffn_size] == [(FFN_ROLE, rank) for rank in range(ffn_size)]: + ffn_ranks_per_domain = 0 + elif len(ffn_counts) == 1: + ffn_ranks_per_domain = ffn_counts.pop() + else: + raise RuntimeError( + f"NVLink domains hold {shape}: every domain (pod) must hold the same " + "number of FFN ranks, or the FFN ranks must fill the first domains" + ) + return TransportLayout( + order=tuple(order), + ffn_ranks_per_domain=ffn_ranks_per_domain, + num_domains=len(ordered), + domain_size=domain_sizes.pop(), + shape=shape, + ) diff --git a/afd_plugin/connectors/gpu/deepep_m2n/network.py b/afd_plugin/connectors/gpu/deepep_m2n/network.py index 715a49ed6..321f7d811 100644 --- a/afd_plugin/connectors/gpu/deepep_m2n/network.py +++ b/afd_plugin/connectors/gpu/deepep_m2n/network.py @@ -190,7 +190,12 @@ def plan_network(topologies: list[RankTopology], config: NetworkConfig) -> Netwo def _domain_blocks(topologies: list[RankTopology]) -> list[list[int]]: - """Ranks grouped by NVLink domain; DeepEP needs equal, contiguous blocks.""" + """Ranks grouped by NVLink domain; DeepEP needs equal, contiguous blocks. + + The connector passes the ranks in transport order + (``plan_transport_layout``), which makes the blocks contiguous; this + checks that order. + """ blocks: list[tuple[str, list[int]]] = [] for topology in topologies: if blocks and blocks[-1][0] == topology.nvlink_domain: diff --git a/docs/design/module/connector_contracts.md b/docs/design/module/connector_contracts.md index 3923f7bbe..28c48e76d 100644 --- a/docs/design/module/connector_contracts.md +++ b/docs/design/module/connector_contracts.md @@ -311,7 +311,7 @@ metadata is present. | Connector | Connector-owned resources | Topology constraints and mapping | | --- | --- | --- | | CUDA P2P | AFD process group, PyNccl data communicators, separate NCCL metadata group, compiled custom-op communicator registry, and graph-oriented receive buffers/state. | Requires `A >= F`. One FFN rank is grouped with a consecutive block of Attention ranks in FFN-first ordering; the blocks differ in size by at most one, and hold `A/F` ranks each when `F` divides `A`. | -| CUDA DeepEP M2N | AFD process group, a CPU bootstrap group for the handshake, the connector's NCCL control group, afdEP-comm transport buffers (one per lane), and compiled dispatch/combine kernels per bucket. | FFN-first ordering, any `A/F`, `num_experts % F == 0`, FFN TP 1. Across nodes every NVLink domain must hold an equal, contiguous block of world ranks. The startup handshake checks payload format, model shapes, lanes, network settings, SM90+, and the FFN expert deployment on every rank before buffers are allocated. | +| CUDA DeepEP M2N | AFD process group, a CPU bootstrap group for the handshake, a CPU transport group when the physical rank order differs from the role order, the connector's NCCL control group, afdEP-comm transport buffers (one per lane), and compiled dispatch/combine kernels per bucket. | FFN-first role ordering, any `A/F`, `num_experts % F == 0`, FFN TP 1. Every NVLink domain must hold the same number of ranks, and either the FFN ranks fill the first domains (afdEP-comm FFN prefix) or every domain holds the same number of FFN ranks (afdEP-comm per-domain mode); the transport sees the ranks domain by domain, FFN first in each. The startup handshake checks payload format, model shapes, lanes, network settings, SM90+, and the FFN expert deployment on every rank before buffers are allocated. | | Ascend CAMP2P | AFD process group, one HCCL communication group per ubatch, FFN HCCL state, Gloo metadata group, custom-op state and transfer handles. | FFN-first ordering and `A >= F`; group construction derives each FFN/Attention mapping. | | Ascend CAM async | Attention-first HCCL group, plugin-owned routed-only CAM operator state, per-stage pending Attention payload queues, and connector work-item state. | Role ranks map into a combined Attention-first world; Compact metadata determines the actual layer and expert interval; counts.sum() gives this chunk’s routed token count. | diff --git a/docs/gpu/DEEPEP_M2N_CONNECTOR_USER_GUIDE.md b/docs/gpu/DEEPEP_M2N_CONNECTOR_USER_GUIDE.md index bd2a67d4a..48d7c2039 100644 --- a/docs/gpu/DEEPEP_M2N_CONNECTOR_USER_GUIDE.md +++ b/docs/gpu/DEEPEP_M2N_CONNECTOR_USER_GUIDE.md @@ -42,6 +42,16 @@ owns real experts `[g * E/F, (g + 1) * E/F)`, so `num_experts` must be divisible by `F`; Attention ranks own dummy expert slots that routing never selects. Any Attention:FFN ratio works. +The handshake, the model checks and the control plane use this order. The +afdEP-comm transport sees the same ranks in physical order instead: NVLink +domain by domain (DeepEP numbers ranks that way), with each domain's FFN ranks +first. When the FFN ranks fill the first domains (one node, or FFN and +Attention on separate pods) the two orders are equal. When every domain holds +the same number of FFN ranks (`4A1F,4A1F`), the connector runs afdEP-comm's +per-domain mode, which finds the FFN ranks at the start of every domain. The +connector derives the order from where the ranks run; see +[Supported layouts](#supported-layouts). + ### Data path For every MoE layer and stage, in the same order on every rank: @@ -69,8 +79,11 @@ together if they disagree. The handshake checks the payload format, model shapes, lane count (DBO), `network`, bucket settings, GPU architecture (SM90 or newer), and the FFN deployment (expert parallel, EP size, FFN TP=1, linear expert placement, no EPLB, all2all backend and expert -weight format). The FFN adopts the Attention TP size, dispatch capacity and -bucket ladder reported by the Attention ranks. +weight format), and the placement of the ranks over NVLink domains. The FFN +adopts the Attention TP size, dispatch capacity and bucket ladder reported by +the Attention ranks. FFN rank 0 logs the network plan and the transport it +picked, for example +`DeepEP M2N transport: per_domain f=1 domains=2x5 placement=4A1F,4A1F`. ## Installation @@ -219,7 +232,64 @@ that sets `capacity_tokens_per_rank` explicitly must agree with Attention. | Attention DP x TP | Supported, including TP > 1: each TP rank dispatches its slice of the DP rank's tokens and the TP group all-gathers the combined rows. | | FFN DP x TP1 | Required. FFN experts are not TP-sharded. | | One node | Supported; ranks reach each other over NVLink. | -| Several nodes | Supported when every node (NVLink domain) holds an equal, contiguous block of world ranks, with FFN ranks first. Writing each node's ranks as `AF`, the layouts `2A0F,0A2F` (Attention on one node, FFN on another) and `4A0F,0A4F` are valid; `4A1F,4A1F` (FFN ranks not contiguous) and `4A0F,0A2F` (unequal blocks) are not. The fabric must route between the NICs NCCL picks. | +| Several nodes or pods | Supported for the placements below. The fabric must route between the NICs NCCL picks. | + +Write each pod's ranks as `AF`, in pod order (on HGX every pod is one +NVLink domain; on NVL72, pods in one NVLink clique form one domain). A +placement is valid when: + +1. every pod holds the same number of GPUs; +2. either the FFN ranks fill the first pods, or every pod holds the same + number of FFN ranks; +3. FFN DP ranks rise in pod order, and so do Attention DP ranks; +4. `num_experts % F == 0`, and Attention TP stays inside a pod. + +Inside a pod, put the FFN ranks on the same local GPUs in every pod +(the recipes use the first ones): DeepEP's hybrid mode pairs ranks across pods +by their position in the pod. + +| Placement | Accepted | afdEP-comm mode | +| --- | --- | --- | +| `4A2F` (one pod) | yes | FFN prefix | +| `2A0F,0A2F`, `0A4F,4A0F` | yes | FFN prefix | +| `3A2F,5A0F` | yes | FFN prefix (one mixed pod at the boundary) | +| `0A1F,0A1F,1A0F,1A0F` | yes | FFN prefix, pods of one GPU | +| `4A1F,4A1F`, `1A1F,1A1F`, `6A2F,6A2F` | yes | per-domain, `f` = FFN ranks per pod | +| `4A1F,3A2F` | no | the FFN count varies from pod to pod | +| `0A1F,0A1F,2A0F,2A0F`, `0A2F,4A0F` | no | pods of different sizes | + +The rejected placements usually have an equivalent with equal pods: for +example `0A2F,2A0F,2A0F` instead of `0A2F,4A0F`. There is nothing to +configure; the handshake rejects other placements with the pods' shapes in +the message. + +## Multi-pod deployment + +Each pod starts one `vllm serve` per role it holds, with the AFD config +(`num_attention_ranks = A`, `num_ffn_ranks = F`, `host:port` = the pod that +holds FFN rank 0) the same on every rank. For each role, with `S` = the role's +DP size, `L` = this pod's local DP size and `start` = the sum of the local +sizes of the pods before it: + +| Flag | First pod holding the role | Other pods holding the role | +| --- | --- | --- | +| `--data-parallel-size` | `S` | `S` | +| `--data-parallel-size-local` | `L` | `L` | +| `--data-parallel-address` / `--data-parallel-rpc-port` | its own address | the first pod's address | +| `--data-parallel-start-rank` | omitted | `start` | +| `--headless` | no | yes | + +When one pod holds the whole role, pass only `--data-parallel-size`. Do not +pass `--data-parallel-start-rank 0` on the first pod: vLLM reads an explicit +start rank on a non-headless engine as hybrid load balancing and then rejects +the headless pods. The FFN DP rank is the FFN's EP rank, so it decides which +experts a pod loads; `--enable-ep-weight-filter` keeps each FFN from reading +the other experts' weights. + +The recipe +[`multipod_4a1f_4a1f_graph.sh`](../../recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/README.md#multi-pod-2-x-4a1f) +derives all of these flags for symmetric placements `N x AF` from the +pod index (`POD=AFD_`). ## Minimal launch shape @@ -266,13 +336,15 @@ and FP8 examples, see the - CPU unit tests: [`tests/unit/connectors/test_deepep_m2n_connector.py`](../../tests/unit/connectors/test_deepep_m2n_connector.py) - runs against a fake afdEP module when afdEP-comm is not installed. + runs against a fake afdEP module when afdEP-comm is not installed; + [`test_deepep_m2n_layout.py`](../../tests/unit/connectors/test_deepep_m2n_layout.py) + checks the placement table above. - GPU connector tests: [`tests/gpu`](../../tests/gpu/README.md) checks dispatch, DeepGEMM experts and combine against references on one node. - Model E2E: the Qwen3 MoE cases `afd-deepep-m2n-*-2a2f` in [`tests/e2e`](../../tests/e2e/README.md#deepep-m2n-scenarios) (eager, graph, - graph + DBO, and Attention TP2) and the two-pod cases in - `tests/e2e/models/qwen3_moe/test_qwen3_moe_multi_pod.py`. + graph + DBO, and Attention TP2) and the two-pod cases (`2A0F,0A2F` and + `1A1F,1A1F`) in `tests/e2e/models/qwen3_moe/test_qwen3_moe_multi_pod.py`. ## Requirements and limitations diff --git a/pyproject.toml b/pyproject.toml index 4f1d7518a..ae1ef3394 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -35,7 +35,7 @@ dependencies = [] [project.optional-dependencies] vllm = ["vllm==0.32.0"] afdep = [ - "afdep-comm @ git+https://github.com/afdEP/afdEP-comm.git@2b74d574e9d25cc794de8d2359f815fcfb48955a", + "afdep-comm @ git+https://github.com/afdEP/afdEP-comm.git@fdff967f6f0a29c154175e0d5865955d347ec69a", ] [project.urls] diff --git a/recipe/gpu/DeepEPM2NAFDConnector/README.md b/recipe/gpu/DeepEPM2NAFDConnector/README.md index da6ffbcbb..4b0e05871 100644 --- a/recipe/gpu/DeepEPM2NAFDConnector/README.md +++ b/recipe/gpu/DeepEPM2NAFDConnector/README.md @@ -16,7 +16,8 @@ supported layouts. │ └── prefill_decode_colocation/ # 2A2F: eager, graph, graph + DBO, Attention TP2 └── qwen3_235b_a22b_fp8/ ├── README.md - └── prefill_decode_colocation/ # 4A2F FP8 payload: graph, graph + DBO + └── prefill_decode_colocation/ # 4A2F FP8 payload: graph, graph + DBO; + # 2 pods x 4A1F (multipod_4a1f_4a1f_graph.sh) ``` Each model directory holds its own scripts; open its `README.md` before @@ -61,6 +62,14 @@ slower than later starts. ## Deploying on Kubernetes -The recipes have no placement header, so the +The single-node recipes have no placement header, so the **`deploy-afd-k8s`** skill (`.agents/skills/deploy-afd-k8s/SKILL.md`) deploys each one as a single pod. The image must contain afdEP-comm and DeepGEMM. + +`multipod_4a1f_4a1f_graph.sh` carries a `POD=AFD_` header, so the skill +deploys it as one pod per entry, each running both roles. Every pod gets the +same `NUM_PODS`, `ATTN_PER_POD` and `FFN_PER_POD` (2, 4 and 1 by default), +which select any symmetric layout `N x AF`; the FFN uses each pod's +first `f` GPUs. See +[Multi-pod deployment](../../../docs/gpu/DEEPEP_M2N_CONNECTOR_USER_GUIDE.md#multi-pod-deployment) +for the placements the connector accepts. diff --git a/recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/README.md b/recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/README.md index 5b5ec1e74..0e64727a1 100644 --- a/recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/README.md +++ b/recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/README.md @@ -5,8 +5,9 @@ Launch scripts for Qwen3-235B-A22B-FP8 on `DeepEPM2NAFDConnector` with vLLM ## Prerequisites -- 6 SM90 or newer GPUs on one node with enough memory for the model. This - configuration was run on H200. +- 6 SM90 or newer GPUs on one node with enough memory for the model, or two + pods of 5 such GPUs on an RDMA fabric for the multi-pod recipe. These + configurations were run on H200. - The environment described in the [connector recipes README](../README.md#prerequisites). - Qwen3-235B-A22B-FP8 weights on disk (block-128 FP8). The scripts default to @@ -21,10 +22,42 @@ owns half of the 128 experts. ```text prefill_decode_colocation/ -├── 4a2f_graph_dp4tp1.sh # 256 sequences per Attention rank, FULL_DECODE_ONLY graph -└── 4a2f_graph_dbo_dp4tp1.sh # 512 sequences per Attention rank, graph + DBO +├── 4a2f_graph_dp4tp1.sh # 256 sequences per Attention rank, FULL_DECODE_ONLY graph +├── 4a2f_graph_dbo_dp4tp1.sh # 512 sequences per Attention rank, graph + DBO +└── multipod_4a1f_4a1f_graph.sh # 2 pods x 4A1F, FULL_DECODE_ONLY graph (DBO optional) ``` +### Multi-pod: 2 x 4A1F + +`multipod_4a1f_4a1f_graph.sh` runs on two pods of 5 GPUs. Each pod runs one +FFN rank on GPU 0 and four Attention ranks on GPUs 1-4 (Attention DP8 x TP1, +FFN DP2/EP2 in total), so every FFN GPU holds half of the experts, as in 4A2F. +Pod 0 leads both DP groups and hosts the AFD rendezvous; pod 1 starts both +servers headless and serves no requests. Send traffic to pod 0's port `18000`. + +Only `POD` differs between pods: + +```bash +# Pod 0 +POD=AFD_0 bash recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/multipod_4a1f_4a1f_graph.sh +# Pod 1 +POD=AFD_1 bash recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/multipod_4a1f_4a1f_graph.sh +``` + +On Kubernetes the script reaches pod 0 through the Services that the +`deploy-afd-k8s` skill creates. Without Kubernetes, set `HEAD_HOST` to pod 0's +address on both pods; pod 0 must accept ports `6239`, `13345` (Attention DP) +and `13355` (FFN DP). Pass the fabric's NCCL settings (for example +`NCCL_IB_HCA`) through the environment. + +`NUM_PODS`, `ATTN_PER_POD` and `FFN_PER_POD` select another symmetric layout +(`2 x 6A2F`, ...); the tuning values (`MAX_NUM_SEQS`, `MAX_NUM_BATCHED_TOKENS`, +`MAX_MODEL_LEN`, `CAPTURE_SIZES`, `PAYLOAD_FORMAT`) and DBO (`ENABLE_DBO=1`, +`DBO_DECODE_TOKEN_THRESHOLD`, `DBO_PREFILL_TOKEN_THRESHOLD`) are environment +variables with the 4A2F values as defaults. The connector logs the placement +it found and the afdEP-comm mode it picked, for example +`DeepEP M2N transport: per_domain f=1 domains=2x5 placement=4A1F,4A1F`. + ## Configuration notes - `connector_extra_config` sets `payload_format: "fp8_block128_fp32"`. The diff --git a/recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/multipod_4a1f_4a1f_graph.sh b/recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/multipod_4a1f_4a1f_graph.sh new file mode 100644 index 000000000..ed48f9b2a --- /dev/null +++ b/recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/multipod_4a1f_4a1f_graph.sh @@ -0,0 +1,82 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project + +# Qwen3-235B-A22B-FP8 with DeepEPM2NAFDConnector, 2 x 4A1F, FULL_DECODE_ONLY +# CUDA graphs: 2 pods of 5 GPUs, each with FFN on GPU 0 and Attention DP4 x TP1 +# on GPUs 1-4 (Attention DP8, FFN DP2/EP2 in total). Send traffic to port +# 18000 of the POD=AFD_0 pod. +# +# Use .agents/skills/deploy-afd-k8s/SKILL.md to deploy this recipe +# on 2 pods using the setup below. +# For pod 1: POD=AFD_0 +# For pod 2: POD=AFD_1 +# +# NUM_PODS x (ATTN_PER_POD Attention + FFN_PER_POD FFN) sets any symmetric +# layout (2 x 1A1F, 3 x 6A2F, ...); give every pod the same values. HEAD_HOST +# points the AFD rendezvous and both DP leaders at pod 0 (runs without k8s). + +MODEL_PATH=${MODEL_PATH:-/path/model_weights/Qwen3-235B-A22B-FP8} +NUM_PODS=${NUM_PODS:-2} +ATTN_PER_POD=${ATTN_PER_POD:-4} +FFN_PER_POD=${FFN_PER_POD:-1} +PAYLOAD_FORMAT=${PAYLOAD_FORMAT:-fp8_block128_fp32} +MAX_NUM_SEQS=${MAX_NUM_SEQS:-256} # per Attention DP rank +MAX_NUM_BATCHED_TOKENS=${MAX_NUM_BATCHED_TOKENS:-3072} +MAX_MODEL_LEN=${MAX_MODEL_LEN:-4096} +CAPTURE_SIZES=${CAPTURE_SIZES:-8,16,32,64,96,128,192,256} # ascending +ENABLE_DBO=${ENABLE_DBO:-0} +DBO_DECODE_TOKEN_THRESHOLD=${DBO_DECODE_TOKEN_THRESHOLD:-2} +DBO_PREFILL_TOKEN_THRESHOLD=${DBO_PREFILL_TOKEN_THRESHOLD:-2} +export VLLM_USE_V2_MODEL_RUNNER=1 + +K=${POD#AFD_} A=$ATTN_PER_POD F=$FFN_PER_POD + +# role_dp PER_POD ADDRESS RPC_PORT -> DP: this pod's DP flags for one role. +# Pod 0 leads and omits its start rank: vLLM reads an explicit start rank on a +# non-headless engine as hybrid load balancing and rejects the headless peers. +role_dp() { + DP=(--data-parallel-size $((NUM_PODS * $1)) --data-parallel-size-local "$1" + --data-parallel-address "$2" --data-parallel-rpc-port "$3") + if ((K > 0)); then DP+=(--headless --data-parallel-start-rank $((K * $1))); fi +} + +afd() { + cat < ffn.log 2>&1 & + +role_dp "$A" "${HEAD_HOST:-vllm-attn-dp-service}" 13345 +CUDA_VISIBLE_DEVICES=$(seq "$F" $((F + A - 1)) | paste -sd, -) vllm serve "$MODEL_PATH" \ + --additional-config "$(afd attention)" "${COMMON[@]}" "${DP[@]}" \ + --port 18000 > attn.log 2>&1 & + +wait diff --git a/tests/e2e/README.md b/tests/e2e/README.md index 35966476b..46ad6747a 100644 --- a/tests/e2e/README.md +++ b/tests/e2e/README.md @@ -173,8 +173,10 @@ python -m pytest -q -s \ ``` `tests/e2e/models/qwen3_moe/test_qwen3_moe_multi_pod.py` runs the eager, -graph and graph + DBO cases with Attention and FFN in two pods, one role per -pod, so dispatch and combine cross nodes over NCCL GIN. The pods must be +graph and graph + DBO cases on two pods in two layouts: one role per pod +(`2A0F,0A2F`) and one Attention and one FFN rank per pod (`1A1F,1A1F`, which +runs afdEP-comm's per-domain mode). Dispatch and combine cross nodes over NCCL +GIN. The pods must be provisioned first; see [`.agents/skills/run-e2e/resources/k8-multi-pod.md`](../../.agents/skills/run-e2e/resources/k8-multi-pod.md) and the [multi-pod runner walkthrough](multi_pod/README.md). diff --git a/tests/e2e/models/qwen3_moe/test_qwen3_moe_multi_pod.py b/tests/e2e/models/qwen3_moe/test_qwen3_moe_multi_pod.py index 8b7511900..5918c0669 100644 --- a/tests/e2e/models/qwen3_moe/test_qwen3_moe_multi_pod.py +++ b/tests/e2e/models/qwen3_moe/test_qwen3_moe_multi_pod.py @@ -10,11 +10,12 @@ runner itself; this only supplies what the runner cannot infer from its own pod's environment. -DeepEP M2N needs FFN ranks first (afdEP-comm) and every NVLink domain to hold -an equal, contiguous block of world ranks (DeepEP), so the two pods split by -role with equal rank counts: Attention on one pod, FFN on the other. The -ranks then span two NVLink domains, so dispatch and combine cross pods over -NCCL GIN (RDMA). +DeepEP M2N needs every NVLink domain (pod) to hold the same number of ranks, +with the FFN ranks either filling the first pods or the same count in every +pod. The cases cover both afdEP-comm modes: the role split (Attention on one +pod, FFN on the other: FFN prefix) and the interleaved layout (one Attention +and one FFN rank per pod: per-domain mode). The ranks span two NVLink domains, +so dispatch and combine cross pods over NCCL GIN (RDMA). """ from __future__ import annotations @@ -43,11 +44,16 @@ # evidence stays comparable with the single-host rows. POD_LAYOUTS = { "2pod-role-split": "2A0F,0A2F", + "2pod-interleaved": "1A1F,1A1F", } MULTI_POD_CASES = [ - ("afd-deepep-m2n-eager-2a2f", "2pod-role-split"), - ("afd-deepep-m2n-graph-2a2f", "2pod-role-split"), - ("afd-deepep-m2n-graph-dbo-2a2f", "2pod-role-split"), + (scenario, layout) + for layout in POD_LAYOUTS + for scenario in ( + "afd-deepep-m2n-eager-2a2f", + "afd-deepep-m2n-graph-2a2f", + "afd-deepep-m2n-graph-dbo-2a2f", + ) ] diff --git a/tests/unit/connectors/test_deepep_m2n_connector.py b/tests/unit/connectors/test_deepep_m2n_connector.py index 551ec0735..9717bff09 100644 --- a/tests/unit/connectors/test_deepep_m2n_connector.py +++ b/tests/unit/connectors/test_deepep_m2n_connector.py @@ -54,7 +54,11 @@ dequantize_fp8_block128, quantize_fp8_block128, ) -from afd_plugin.connectors.gpu.deepep_m2n.layout import DeepEPM2NLayout # noqa: E402 +from afd_plugin.connectors.gpu.deepep_m2n.layout import ( # noqa: E402 + DeepEPM2NLayout, + RankPlacement, + plan_transport_layout, +) from afd_plugin.connectors.gpu.deepep_m2n.network import ( # noqa: E402 NetworkConfig, RankTopology, @@ -621,6 +625,8 @@ def test_handshake_accepts_consistent_ranks(): assert result.expert_alignment == 128 assert not result.network_plan.use_gin assert (result.attn_tp_size, result.capacity_tokens_per_rank) == (1, 256) + assert result.transport.is_identity + assert result.transport.ffn_ranks_per_domain == 0 def test_handshake_gives_ffn_the_attention_tp_size_capacity_and_ladder(): @@ -696,9 +702,11 @@ def test_handshake_rejects_mismatches(overrides, message): def test_handshake_plans_gin_across_nvlink_domains(): - """Ranks on hosts ``aabb`` give a plan of 2 domains of 2 ranks with GIN. + """Ranks on hosts ``aabb`` (role-split pods) give a plan of 2 domains of 2 + ranks with GIN, and keep the role world as the transport order. - Why: the handshake must hand the network plan the right domains. + Why: the handshake must hand the network plan the right domains, and the + role-split layouts that run today must take the unchanged transport. """ layout = DeepEPM2NLayout(ffn_size=2, attn_size=2, num_experts=128) @@ -711,18 +719,54 @@ def test_handshake_plans_gin_across_nvlink_domains(): 2, 2, ) + assert result.transport.is_identity + + +def test_handshake_gives_interleaved_pods_the_per_domain_transport(): + """FFN 0 and FFN 1 on hosts ``a`` and ``b``, each with one Attention rank + (``1A1F,1A1F``), give a per-domain transport ``[F0, A0, F1, A1]`` over 2 + domains of 2 ranks. + + Why: afdEP-comm's lanes must see the ranks domain by domain with each + domain's FFN rank first; the network plan is checked in that order. + """ + layout = DeepEPM2NLayout(ffn_size=2, attn_size=2, num_experts=128) + + result = validate_handshake( + _handshake(layout, hosts="abab"), layout, NetworkConfig() + ) + + assert result.transport.order == ( + ("ffn", 0), + ("attention", 0), + ("ffn", 1), + ("attention", 1), + ) + assert result.transport.ffn_ranks_per_domain == 1 + assert (result.network_plan.num_domains, result.network_plan.ranks_per_domain) == ( + 2, + 2, + ) @pytest.mark.parametrize( - ("hosts", "message"), - [("aaab", "different rank counts"), ("abab", "not contiguous")], + ("ffn_size", "attn_size", "hosts", "message"), + [ + (2, 2, "aaab", "same number of ranks"), + # Host a holds F0, A0, A1 and host b F1, F2, A2. + (3, 3, "abbaab", "same number of FFN ranks"), + ], ) -def test_handshake_rejects_unsupported_domain_layouts(hosts, message): - """Hosts ``aaab`` (unequal counts) and ``abab`` (not contiguous) raise. +def test_handshake_rejects_unsupported_domain_layouts( + ffn_size, attn_size, hosts, message +): + """Hosts with unequal rank counts, or with FFN counts that differ while + the FFN ranks do not fill the first hosts, raise. - Why: DeepEP's hybrid mode needs equal, contiguous ranks per domain. + Why: DeepEP needs equal NVLink domains, and afdEP-comm finds FFN ranks + only as a prefix or as the same count in every domain. """ - layout = DeepEPM2NLayout(ffn_size=2, attn_size=2, num_experts=128) + layout = DeepEPM2NLayout(ffn_size=ffn_size, attn_size=attn_size, num_experts=96) with pytest.raises(RuntimeError, match=message): validate_handshake(_handshake(layout, hosts=hosts), layout, NetworkConfig()) @@ -1285,27 +1329,48 @@ def test_attention_init_rejects_local_experts_before_communicating(): assert connector.bootstrap_group is None -@pytest.mark.parametrize(("role", "role_rank"), [("attention", 2), ("ffn", 1)]) +@pytest.mark.parametrize( + ("role", "role_rank", "hosts", "ffn_ranks_per_domain"), + [ + ("attention", 2, "aaaaaa", 0), + ("ffn", 1, "aaaaaa", 0), + # Hosts of F0 F1 A0 A1 A2 A3: 2A1F,2A1F. + ("ffn", 1, "ababab", 1), + ("attention", 3, "ababab", 1), + ], +) def test_lane_transport_is_afdep_comm_with_the_ffn_rank_count( - monkeypatch, role, role_rank + monkeypatch, role, role_rank, hosts, ffn_ranks_per_domain ): """Each lane's transport is afdEP-comm's ``ElasticBuffer``, built with every - argument its constructor requires and the FFN rank count. + argument its constructor requires, the FFN rank count and the transport's + ``ffn_ranks_per_domain``, on the transport group. Why: ``ffn_ranks`` selects afdEP-comm's AFD transport (without it the - buffer is unspecialized DeepEP V2), and the stand-in fails on any argument - set the real constructor would refuse. + buffer is unspecialized DeepEP V2), ``ffn_ranks_per_domain`` tells it + where the FFN ranks sit, and the stand-in fails on any argument set the + real constructor would refuse. """ monkeypatch.setattr(connector_module, "ElasticBuffer", FakeElasticBuffer) - connector = _connector(role, role_rank, M2NTopology(ffn_size=2, attn_size=3), {}) - connector.bootstrap_group = SimpleNamespace() + topology = M2NTopology(ffn_size=2, attn_size=4) + connector = _connector(role, role_rank, topology, {}) + connector.transport_group = SimpleNamespace() + connector.transport = plan_transport_layout( + [ + RankPlacement(*topology.role(rank), f"host:{host}") + for rank, host in enumerate(hosts) + ], + ffn_size=topology.ffn_size, + ) connector.network_plan = SimpleNamespace(hybrid=True, rdma_traffic_class=0) connector.capacity_tokens_per_rank = 64 buffer = connector._new_lane() assert isinstance(buffer, FakeElasticBuffer) + assert buffer.kwargs["group"] is connector.transport_group assert buffer.kwargs["ffn_ranks"] == 2 + assert buffer.kwargs["ffn_ranks_per_domain"] == ffn_ranks_per_domain assert buffer.kwargs["capacity_tokens_per_rank"] == 64 assert buffer.kwargs["allow_hybrid_mode"] is True diff --git a/tests/unit/connectors/test_deepep_m2n_layout.py b/tests/unit/connectors/test_deepep_m2n_layout.py new file mode 100644 index 000000000..328f1177a --- /dev/null +++ b/tests/unit/connectors/test_deepep_m2n_layout.py @@ -0,0 +1,136 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""CPU tests for the DeepEP M2N transport planner (``plan_transport_layout``).""" + +from __future__ import annotations + +import pytest + +pytest.importorskip("torch") +pytest.importorskip("vllm") + +from afd_plugin.connectors.gpu.deepep_m2n.layout import ( # noqa: E402 + RankPlacement, + TransportLayout, + plan_transport_layout, +) + + +def _pods(placement: str) -> tuple[list[RankPlacement], int]: + """Ranks of ``AF,...`` with DP ranks rising in pod order. + + Each pod is one NVLink domain and lists its Attention ranks first, the way + a pod that starts Attention on its first GPUs reports them. + """ + ranks: list[RankPlacement] = [] + next_rank = {"attention": 0, "ffn": 0} + for pod, entry in enumerate(placement.split(",")): + attention, ffn = (int(n) for n in entry.removesuffix("F").split("A")) + for role, count in (("attention", attention), ("ffn", ffn)): + for _ in range(count): + ranks.append(RankPlacement(role, next_rank[role], f"host:pod{pod}")) + next_rank[role] += 1 + return ranks, next_rank["ffn"] + + +def _plan(placement: str) -> TransportLayout: + ranks, ffn_size = _pods(placement) + return plan_transport_layout(ranks, ffn_size) + + +def _order(layout: TransportLayout) -> str: + return " ".join(f"{role[0].upper()}{rank}" for role, rank in layout.order) + + +@pytest.mark.parametrize( + ("placement", "order", "ffn_ranks_per_domain", "is_identity"), + [ + ("4A2F", "F0 F1 A0 A1 A2 A3", 0, True), + ("2A0F,0A2F", "F0 F1 A0 A1", 0, True), + ("3A2F,5A0F", "F0 F1 A0 A1 A2 A3 A4 A5 A6 A7", 0, True), + ("0A1F,0A1F,1A0F,1A0F", "F0 F1 A0 A1", 0, True), + ("1A1F,1A1F", "F0 A0 F1 A1", 1, False), + ("4A1F,4A1F", "F0 A0 A1 A2 A3 F1 A4 A5 A6 A7", 1, False), + ("2A2F,2A2F,2A2F", "F0 F1 A0 A1 F2 F3 A2 A3 F4 F5 A4 A5", 2, False), + ], +) +def test_transport_orders_domains_ffn_first_and_picks_the_afdep_mode( + placement, order, ffn_ranks_per_domain, is_identity +): + """Each pod becomes one block with its FFN ranks first. FFN ranks that + fill the first pods keep the role world (FFN prefix); the same FFN count + in every pod gives afdEP-comm's per-domain mode. + + Why: DeepEP numbers ranks domain by domain, and afdEP-comm finds the FFN + ranks either at [0, F) or at the first f ranks of every domain; any other + order routes rows to the wrong ranks. Layouts that run today must keep the + role world so they take the unchanged path. + """ + layout = _plan(placement) + + assert _order(layout) == order + assert layout.ffn_ranks_per_domain == ffn_ranks_per_domain + assert layout.is_identity is is_identity + + +def test_per_domain_ffn_ranks_sit_where_afdep_expects_them(): + """In 3 x 6A2F, FFN rank ``j`` is transport rank ``(j // 2) * 8 + j % 2``. + + Why: an FFN rank loads the experts of its DP (EP) rank, and afdEP-comm + routes those experts to this transport rank; a mismatch silently computes + rows with the wrong experts. + """ + layout = _plan("6A2F,6A2F,6A2F") + + assert (layout.num_domains, layout.domain_size) == (3, 8) + assert [layout.transport_rank("ffn", j) for j in range(6)] == [ + (j // 2) * 8 + j % 2 for j in range(6) + ] + + +def test_role_split_pods_in_any_order_keep_the_ffn_prefix(): + """Attention pods listed before the FFN pod still give FFN-first order. + + Why: domains follow their lowest role-world rank, so the pod order does + not change the transport the role-split layouts run today. + """ + layout = _plan("4A0F,0A4F") + + assert _order(layout) == "F0 F1 F2 F3 A0 A1 A2 A3" + assert layout.is_identity + + +@pytest.mark.parametrize( + ("ranks", "ffn_size", "message"), + [ + # 4A1F,3A2F: the FFN count varies from pod to pod. + (_pods("4A1F,3A2F")[0], 3, "same number of FFN ranks"), + # Pods of 1, 1, 2 and 2 GPUs. + (_pods("0A1F,0A1F,2A0F,2A0F")[0], 2, "same number of ranks"), + (_pods("0A2F,4A0F")[0], 2, "same number of ranks"), + # 2A2F,2A2F with FFN DP ranks 0, 2 on pod 0 and 1, 3 on pod 1. + ( + [ + RankPlacement(role, rank, f"host:pod{pod}") + for pod, roles in enumerate( + [ + [("ffn", 0), ("ffn", 2), ("attention", 0), ("attention", 1)], + [("ffn", 1), ("ffn", 3), ("attention", 2), ("attention", 3)], + ] + ) + for role, rank in roles + ], + 4, + "FFN DP ranks must rise in pod order", + ), + ], +) +def test_transport_rejects_placements_afdep_cannot_carry(ranks, ffn_size, message): + """Unequal pods, an FFN count that varies per pod, and FFN DP ranks out + of pod order raise, naming the placement. + + Why: DeepEP needs equal NVLink domains and afdEP-comm has no mode for the + rest; failing in the handshake beats a hang in the first dispatch. + """ + with pytest.raises(RuntimeError, match=message): + plan_transport_layout(ranks, ffn_size) diff --git a/tests/unit/fakes/afdep.py b/tests/unit/fakes/afdep.py index aaf564877..4aaf5a98d 100644 --- a/tests/unit/fakes/afdep.py +++ b/tests/unit/fakes/afdep.py @@ -14,7 +14,7 @@ from types import ModuleType # Keyword-only parameters of afdep.deepep.ElasticBuffer at the pinned afdEP-comm -# (afdEP/afdEP-comm 2b74d57). The stand-in requires the required ones and +# (afdEP/afdEP-comm fdff967). The stand-in requires the required ones and # rejects any other keyword, so a lane built with arguments the real transport # would refuse fails here as it would on a GPU. ELASTIC_BUFFER_REQUIRED = ( @@ -28,7 +28,13 @@ "cpu_timeout_s", "gpu_timeout_s", ) -ELASTIC_BUFFER_OPTIONAL = ("use_fp8_dispatch", "num_sms", "num_qps", "ffn_ranks") +ELASTIC_BUFFER_OPTIONAL = ( + "use_fp8_dispatch", + "num_sms", + "num_qps", + "ffn_ranks", + "ffn_ranks_per_domain", +) class FakeElasticBuffer: diff --git a/uv.lock b/uv.lock index 3c3f6e353..3394d9826 100644 --- a/uv.lock +++ b/uv.lock @@ -31,7 +31,7 @@ resolution-markers = [ [[package]] name = "afdep-comm" version = "1.1.0" -source = { git = "https://github.com/afdEP/afdEP-comm.git?rev=2b74d574e9d25cc794de8d2359f815fcfb48955a#2b74d574e9d25cc794de8d2359f815fcfb48955a" } +source = { git = "https://github.com/afdEP/afdEP-comm.git?rev=fdff967f6f0a29c154175e0d5865955d347ec69a#fdff967f6f0a29c154175e0d5865955d347ec69a" } dependencies = [ { name = "ninja" }, { name = "packaging" }, @@ -5584,7 +5584,7 @@ e2e-tests = [ [package.metadata] requires-dist = [ - { name = "afdep-comm", marker = "extra == 'afdep'", git = "https://github.com/afdEP/afdEP-comm.git?rev=2b74d574e9d25cc794de8d2359f815fcfb48955a" }, + { name = "afdep-comm", marker = "extra == 'afdep'", git = "https://github.com/afdEP/afdEP-comm.git?rev=fdff967f6f0a29c154175e0d5865955d347ec69a" }, { name = "vllm", marker = "extra == 'vllm'", specifier = "==0.32.0" }, ] provides-extras = ["afdep", "vllm"] From a8f953659946e0a67bd8d373773929379d1d75c7 Mon Sep 17 00:00:00 2001 From: Ronen Kat <16743404+ronenkat@users.noreply.github.com> Date: Mon, 2 Nov 2026 19:05:08 +0200 Subject: [PATCH 27/32] Remove captured-split DBO changes --- .../compat/patches/gpu_mrv2_ubatching.py | 400 +----------------- .../module/compatibility_and_patches.md | 1 - .../compat/patches/test_gpu_mrv2_ubatching.py | 377 +---------------- 3 files changed, 24 insertions(+), 754 deletions(-) diff --git a/afd_plugin/compat/patches/gpu_mrv2_ubatching.py b/afd_plugin/compat/patches/gpu_mrv2_ubatching.py index 1c577a0ba..c264fa0e7 100644 --- a/afd_plugin/compat/patches/gpu_mrv2_ubatching.py +++ b/afd_plugin/compat/patches/gpu_mrv2_ubatching.py @@ -3,9 +3,7 @@ """Scoped DBO seams for the pinned native GPU ModelRunnerV2. The native UBatchRunner owns input slicing, Attention preparation, threads and -CUDA graphs. AFD extends DP1 eligibility, dispatch and child-context metadata, -and keeps FULL microbatched decodes on the captured request split so an -underfilled DP rank replays the graph instead of running eager. +CUDA graphs. AFD extends DP1 eligibility, dispatch and child-context metadata. Bindings are limited to AFD Attention initialization/execution and restored on exit; worker initialization and dispatch run synchronously in each process. """ @@ -15,116 +13,28 @@ from collections.abc import Iterator from contextlib import contextmanager from dataclasses import replace -from itertools import pairwise from types import MethodType from typing import Any import torch -import torch.distributed as dist from vllm.config import CUDAGraphMode, ParallelConfig, VllmConfig -from vllm.distributed.parallel_state import get_dp_group from vllm.forward_context import DPMetadata, ForwardContext, create_forward_context from vllm.v1.kv_cache_interface import KVCacheConfig -from vllm.v1.worker.dp_utils import should_skip_dp_coordination from vllm.v1.worker.gpu import model_runner -from vllm.v1.worker.gpu.attn_utils import build_slot_mappings_by_layer from vllm.v1.worker.gpu.cudagraph_utils import ( BatchExecutionDescriptor, CudaGraphManager, ) -from vllm.v1.worker.gpu.dp_utils import DPSyncState -from vllm.v1.worker.gpu.input_batch import InputBatch +from vllm.v1.worker.gpu.dp_utils import DPSyncState, sync_cudagraph_and_dp_padding from vllm.v1.worker.gpu.model_states.interface import ModelState -from vllm.v1.worker.gpu.ubatch_utils import ( - UBatchRunner, - UBatchState, - _slice_input_batch, - create_ubatch_slices, - logger, -) +from vllm.v1.worker.gpu.ubatch_utils import UBatchRunner, logger from vllm.v1.worker.ubatch_utils import ( - UBatchSlice, UBatchSlices, check_ubatch_thresholds, get_num_ubatches, ) from vllm.v1.worker.utils import AttentionGroup -# A FULL microbatched decode graph is captured with one token per request. -DECODE_QUERY_LEN = 1 -# Host-side sequence bound for a microbatch holding only padding requests; the -# graph's kernels were captured with max_model_len and ignore it. -PADDING_ONLY_SEQ_LEN_BOUND = 1 - - -def replays_captured_ubatch_split(vllm_config: VllmConfig) -> bool: - """Whether FULL microbatched decodes keep the captured request split. - - The split keeps every microbatch's captured request and token range, so a - rank holding at most half of the graph's requests replays its trailing - microbatch on padding rows only. Limited to one-token decodes over - attention without inner state or context parallelism; MLA is not yet - validated with a padding-only microbatch. - """ - model_config = vllm_config.model_config - parallel_config = vllm_config.parallel_config - return ( - vllm_config.speculative_config is None - and not model_config.has_inner_state - and not model_config.use_mla - and parallel_config.prefill_context_parallel_size == 1 - and parallel_config.decode_context_parallel_size == 1 - ) - - -def needs_regular_graph_fallback( - cudagraph_manager: CudaGraphManager, - num_reqs: int, - min_num_tokens: int, - max_num_tokens: int, - uniform_token_count: int | None, - num_active_loras: int, - num_ubatches: int, -) -> bool: - """Whether a microbatched graph would replay a shifted request split. - - Native slicing re-derives the split from the real requests, so a graph only - replays correctly when the smallest rank fills more than its leading - microbatches. Without the captured split, such a step replays the regular - graph instead of running eager. Every rank passes the same agreed token - bounds, so every rank returns the same answer. - """ - if replays_captured_ubatch_split(cudagraph_manager.vllm_config): - return False - ubatch_desc = cudagraph_manager.dispatch( - num_reqs, - max_num_tokens, - uniform_token_count, - num_active_loras=num_active_loras, - num_ubatches=num_ubatches, - ) - return ( - ubatch_desc.cg_mode == CUDAGraphMode.FULL - and num_ubatches * min_num_tokens <= ubatch_desc.num_tokens - ) - - -def captured_ubatch_slices(input_batch: InputBatch, num_ubatches: int) -> UBatchSlices: - """Slices a FULL one-token-per-request graph was captured with. - - Capture splits a dummy batch whose request and token counts are equal at - `num_tokens // num_ubatches`, the last microbatch taking the remainder. - Native slicing yields the same slices whenever real requests reach the - last microbatch. - """ - num_tokens = input_batch.num_tokens_after_padding - split_point = num_tokens // num_ubatches - bounds = [split_point * i for i in range(num_ubatches)] + [num_tokens] - return [ - UBatchSlice(slice(start, stop), slice(start, stop)) - for start, stop in pairwise(bounds) - ] - # Patch reason: the native factory excludes DP1, although AFD schedules its two # stages locally and sends their complete metadata to the FFN control plane. @@ -160,103 +70,9 @@ def maybe_build_ubatch_runner( if parallel_config.data_parallel_size == 1: runner._make_forward_contexts = MethodType(_make_forward_contexts, runner) # ### PATCH END: bind native DP1 child metadata on this instance. - # ### PATCH START: bind the captured-split prepare on this instance. - if replays_captured_ubatch_split(vllm_config): - runner.prepare = MethodType(prepare, runner) - # ### PATCH END: bind the captured-split prepare on this instance. return runner -# Patch reason: native prepare re-derives the microbatch split from the real -# requests and clamps an empty trailing microbatch onto the last real request. -# A FULL microbatched graph replays fixed request/token offsets, so a rank with -# at most half of the graph's requests would replay with shifted metadata, and -# native dispatch runs such steps eager instead (`2 * min < N`). -# Patch functionality: slice FULL one-token decodes at the captured split. A -# trailing microbatch without real requests keeps its padded shape, with a -# one-token query bound and a nonzero host-side sequence bound. -# Signature: matches UBatchRunner.prepare exactly at the pinned commit. -# Upstream: vllm/v1/worker/gpu/ubatch_utils.py, UBatchRunner.prepare; -# 5df63a7203cdb9328d824a27b682e77c7b9928be (vLLM 0.32.0). -# Removal plan: remove when native FULL microbatching keeps the captured split. -def prepare( - self, - input_batch: InputBatch, - block_tables: tuple[torch.Tensor, ...], - slot_mappings: torch.Tensor, - cg_mode: CUDAGraphMode = CUDAGraphMode.NONE, - for_capture: bool = False, -) -> UBatchState: - """Split the batch into the microbatches the step will run on. - - Attention metadata is built per microbatch and carried in the forward - contexts the threads install, not in the caller's. - - FULL uses padded sizes; `for_capture` refreshes metadata from dummy - block tables. - """ - # ### PATCH START: slice FULL one-token decodes at the captured split. - keeps_captured_split = ( - cg_mode == CUDAGraphMode.FULL - and input_batch.num_tokens == input_batch.num_reqs * DECODE_QUERY_LEN - and input_batch.num_tokens_after_padding - == input_batch.num_reqs_after_padding * DECODE_QUERY_LEN - ) - if keeps_captured_split: - # A microbatch without real requests has no query length to derive. - input_batch = replace(input_batch, max_query_len=DECODE_QUERY_LEN) - ubatch_slices = captured_ubatch_slices(input_batch, self.num_ubatches) - else: - ubatch_slices = create_ubatch_slices(input_batch, self.num_ubatches) - # ### PATCH END: slice FULL one-token decodes at the captured split. - - attn_metadata = [] - slot_mappings_by_layer = [] - is_padding = [] - for i, ubatch_slice in enumerate(ubatch_slices): - ubatch = _slice_input_batch( - input_batch, - ubatch_slice, - self.ubatch_query_start_loc[i], - self.ubatch_seq_lens[i], - self.ubatch_dcp_local_seq_lens[i], - dcp_size=self.dcp_size, - dcp_rank=self.dcp_rank, - cp_interleave=self.cp_interleave, - ) - # ### PATCH START: keep a padding-only microbatch's sequence bound nonzero. - if keeps_captured_split and ubatch.num_reqs == 0: - ubatch.seq_lens_cpu_upper_bound[0] = PADDING_ONLY_SEQ_LEN_BOUND - # ### PATCH END: keep a padding-only microbatch's sequence bound nonzero. - ubatch_slot_mappings = slot_mappings[:, ubatch_slice.token_slice] - ubatch_block_tables = tuple( - block_table[ubatch_slice.request_slice] for block_table in block_tables - ) - attn_metadata.append( - self.model_state.prepare_attn( - ubatch, - cg_mode, - ubatch_block_tables, - ubatch_slot_mappings, - self.attn_groups, - self.kv_cache_config, - ubatch_idx=i, - for_capture=for_capture, - ) - ) - slot_mappings_by_layer.append( - build_slot_mappings_by_layer(ubatch_slot_mappings, self.kv_cache_config) - ) - is_padding.append(ubatch.is_padding) - - return UBatchState( - slices=ubatch_slices, - forward_contexts=self._make_forward_contexts( - ubatch_slices, attn_metadata, slot_mappings_by_layer, is_padding - ), - ) - - # Patch reason: DPMetadata.make asserts DP>1 or SP; AFD DP1 uses neither. # Patch functionality: build the same native DPMetadata directly for this DP1 # runner instance. Native context fields, token counts and SP methods are kept. @@ -299,202 +115,9 @@ def _make_forward_contexts( return forward_contexts -# Patch reason: native DP>1 microbatching runs eager whenever the smallest rank -# holds at most half of the graph (`2 * min < N`), because native slicing would -# replay a shifted request split. -# Patch functionality: keep the microbatched FULL graph when AFD's prepare -# replays the captured split; otherwise replay the regular graph, never eager. -# Every other branch is the native coordination unchanged. -# Signature: matches sync_cudagraph_and_dp_padding exactly at the pinned commit. -# Upstream: vllm/v1/worker/gpu/dp_utils.py, sync_cudagraph_and_dp_padding; -# 5df63a7203cdb9328d824a27b682e77c7b9928be (vLLM 0.32.0). -# Removal plan: remove when native FULL microbatching keeps the captured split. -def sync_cudagraph_and_dp_padding( - cudagraph_manager: CudaGraphManager | None, - desired_batch_desc: BatchExecutionDescriptor, - num_tokens: int, - num_reqs: int, - uniform_token_count: int | None, - dp_size: int, - dp_rank: int, - max_query_len: int | None = None, - num_active_loras: int = 0, - parallel_config: ParallelConfig | None = None, - allow_ubatching: bool = False, - uniform_decode: bool = False, -) -> tuple[BatchExecutionDescriptor, DPSyncState | None]: - """Coordinates the batch descriptor and DP padding across all ranks. - - `parallel_config` is only needed to decide whether to microbatch, so callers - that never do (`allow_ubatching=False`) can leave it out. - - Returns (synced_batch_desc, sync). `sync` is None when no rank has work. - """ - assert dp_size > 1, "DP size must be greater than 1" - group = get_dp_group().cpu_group - tensor = torch.zeros(6, dp_size, dtype=torch.int32, device="cpu") - tensor[0][dp_rank] = num_tokens - tensor[1][dp_rank] = desired_batch_desc.cg_mode.value - tensor[2][dp_rank] = uniform_token_count or 0 # (0 means None) - tensor[3][dp_rank] = max_query_len or -1 # (-1 means None) - tensor[4][dp_rank] = int(allow_ubatching) - tensor[5][dp_rank] = num_reqs - if should_skip_dp_coordination(): - tensor[:] = tensor[:, dp_rank, None].clone() - else: - dist.all_reduce(tensor, group=group) - - num_tokens_across_dp = tensor[0] - cg_mode_across_dp = tensor[1] - uniform_token_counts_across_dp = tensor[2] - max_query_lens_across_dp = tensor[3] - allow_ubatching_across_dp = tensor[4] - num_reqs_across_dp = tensor[5] - - # If ranks disagree on the uniform token count, or its 0 (means None) set to None - synced_uniform_token_count: int | None = int(uniform_token_counts_across_dp[0]) - if synced_uniform_token_count == 0 or not torch.all( - uniform_token_counts_across_dp == synced_uniform_token_count - ): - synced_uniform_token_count = None - - if torch.all(num_tokens_across_dp == 0).item(): - synced_desc = BatchExecutionDescriptor( - cg_mode=CUDAGraphMode.NONE, num_tokens=0, num_reqs=0 - ) - return synced_desc, None - - if torch.all(allow_ubatching_across_dp == 1): - # This rank voted too, so its caller passed a config. - assert parallel_config is not None - # A uniform decode only if every rank runs the same uniform query - # length, so every rank reaches the same answer here. - uniform_decode_across_dp = uniform_decode and torch.all( - uniform_token_counts_across_dp == (uniform_token_count or 0) - ) - # ### PATCH START: regular graph when the captured split cannot serve. - if check_ubatch_thresholds( - parallel_config, - # Thresholds only grow with the token count, so holding the smallest - # rank to them holds every rank to them. - int(num_tokens_across_dp.min()), - uniform_decode=uniform_decode_across_dp, - ) and not ( - cudagraph_manager is not None - and needs_regular_graph_fallback( - cudagraph_manager, - num_reqs, - int(num_tokens_across_dp.min()), - int(num_tokens_across_dp.max()), - synced_uniform_token_count, - num_active_loras, - get_num_ubatches(parallel_config), - ) - ): - # ### PATCH END: regular graph when the captured split cannot serve. - # Expert all-to-all requires every rank to split and pad equally. - # Empty microbatches run as dummy batches. - ubatch_num_tokens = int(num_tokens_across_dp.max()) - num_ubatches = get_num_ubatches(parallel_config) - ubatch_desc = None - if cudagraph_manager is not None: - # Match a FULL microbatched CUDA graph, falling back to an eager - # descriptor if no graph matches. - ubatch_desc = cudagraph_manager.dispatch( - num_reqs, - ubatch_num_tokens, - synced_uniform_token_count, - num_active_loras=num_active_loras, - num_ubatches=num_ubatches, - ) - # ### PATCH START: no eager fallback for an underfilled rank. - # The captured split replays it; otherwise the threshold check - # above already routed the step to the regular graph. - # ### PATCH END: no eager fallback for an underfilled rank. - if ubatch_desc is None: - ubatch_desc = BatchExecutionDescriptor( - cg_mode=CUDAGraphMode.NONE, - num_tokens=ubatch_num_tokens, - num_reqs=num_reqs, - num_ubatches=num_ubatches, - ) - # Refresh the token count to include CUDA graph padding. - ubatch_num_tokens = ubatch_desc.num_tokens - num_reqs = int(num_reqs_across_dp.max()) - if ubatch_desc.cg_mode == CUDAGraphMode.FULL: - assert ubatch_desc.num_reqs is not None - num_reqs = ubatch_desc.num_reqs - return ubatch_desc, DPSyncState( - num_tokens_across_dp=torch.full_like( - num_tokens_across_dp, ubatch_num_tokens - ), - uniform_token_count=synced_uniform_token_count, - eager=ubatch_desc.cg_mode == CUDAGraphMode.NONE, - num_reqs=num_reqs, - ) - - synced_cg_mode = CUDAGraphMode(int(cg_mode_across_dp.min())) - - # If any rank wants to run eager, all ranks run eager - if synced_cg_mode == CUDAGraphMode.NONE: - return ( - BatchExecutionDescriptor( - cg_mode=CUDAGraphMode.NONE, - num_tokens=num_tokens, - num_reqs=num_reqs, - num_active_loras=desired_batch_desc.num_active_loras, - ), - DPSyncState( - num_tokens_across_dp=num_tokens_across_dp, - uniform_token_count=synced_uniform_token_count, - eager=True, - num_reqs=int(num_reqs_across_dp.max()), - ), - ) - - assert cudagraph_manager is not None, ( - "cudagraph_manager should only be None during profile run, " - "where synced_cg_mode must be NONE across all DP ranks" - ) - synced_num_tokens = int(num_tokens_across_dp.max()) - - # Varlen decode graphs are selected by the query-length bound, so ranks must agree - # on it or they pad to different token counts below. - synced_max_query_len: int | None = None - if bool(torch.all(max_query_lens_across_dp != -1)): - synced_max_query_len = int(max_query_lens_across_dp.max()) - - # Dispatch for the final synced values, use num_reqs instead of synced_num_reqs - # so we don't perform request padding for PIECEWISE graphs. - # num_active_loras is per-rank and doesn't need cross-rank agreement. - synced_desc = cudagraph_manager.dispatch( - num_reqs, - synced_num_tokens, - synced_uniform_token_count, - num_active_loras=num_active_loras, - max_query_len=synced_max_query_len, - ) - - # Update num_tokens_across_dp to reflect padded size. - num_tokens_across_dp[:] = synced_desc.num_tokens - - return synced_desc, DPSyncState( - num_tokens_across_dp=num_tokens_across_dp, - uniform_token_count=synced_uniform_token_count, - eager=False, - num_reqs=( - synced_desc.num_reqs - if synced_desc.cg_mode == CUDAGraphMode.FULL - and synced_desc.num_reqs is not None - else int(num_reqs_across_dp.max()) - ), - ) - - # Patch reason: native DP1 dispatch returns before its DBO threshold decision. -# Patch functionality: apply the native threshold/graph-padding policy locally, -# replaying the regular graph when the captured split cannot serve the step; -# route DP>1 coordination through AFD's sync_cudagraph_and_dp_padding. +# Patch functionality: apply the native threshold/graph-padding policy locally; +# retain native DP>1 coordination and native descriptor selection. # Signature: matches dispatch_cg_and_sync_dp exactly at the pinned commit. # Upstream: vllm/v1/worker/gpu/dp_utils.py, dispatch_cg_and_sync_dp; # 5df63a7203cdb9328d824a27b682e77c7b9928be (vLLM 0.32.0). @@ -588,16 +211,6 @@ def dispatch_cg_and_sync_dp( ubatch_desc = None if not need_eager and not reuse_eager: assert cudagraph_manager is not None - if needs_regular_graph_fallback( - cudagraph_manager, - num_reqs, - num_tokens, - num_tokens, - uniform_token_count, - num_active_loras, - num_ubatches, - ): - return batch_desc, None ubatch_desc = cudagraph_manager.dispatch( num_reqs, num_tokens, @@ -605,6 +218,9 @@ def dispatch_cg_and_sync_dp( num_active_loras=num_active_loras, num_ubatches=num_ubatches, ) + if num_ubatches * num_tokens < ubatch_desc.num_tokens: + # Match the native DP>1 guard against an empty trailing stage. + ubatch_desc = None if ubatch_desc is None: ubatch_desc = BatchExecutionDescriptor( cg_mode=CUDAGraphMode.NONE, diff --git a/docs/design/module/compatibility_and_patches.md b/docs/design/module/compatibility_and_patches.md index 37e0bfe0e..86b107c68 100644 --- a/docs/design/module/compatibility_and_patches.md +++ b/docs/design/module/compatibility_and_patches.md @@ -132,7 +132,6 @@ not the package dependency policy. | [`async_dp_forward_context.py`](../../../afd_plugin/compat/patches/async_dp_forward_context.py): `vllm.forward_context.set_forward_context` plus already-imported worker aliases | Skips native `DPMetadata` construction/coordination only for AFD async-DP; otherwise uses the copied upstream flow. | Imported by `register_afd`; same target/dev/unknown guard. Rebinds known already-imported aliases so callers do not retain the old function. | [`test_async_dp_forward_context.py`](../../../tests/unit/compat/patches/test_async_dp_forward_context.py) covers async skip and non-async coordination. | Remove when vLLM supports a per-engine-role opt-out from native MoE DP metadata coordination. | | [`config_validation.py`](../../../afd_plugin/compat/patches/config_validation.py): `EngineArgs.create_engine_config`, `VllmConfig.__post_init__` | For AFD-owned ubatching with a non-DeepEP backend, temporarily presents `deepep_low_latency` during upstream validation and restores the configured backend. After upstream platform normalization, maps an initial `worker_cls="auto"` to the role-specific CUDA or standard Ascend AFD worker. Native hybrid/multimodal DBO rejection remains intact, including Qwen3.5/3.6 GDN models. | Imported by `register_afd`; accepts the target version, development versions, or missing version metadata. Saves originals on upstream modules under AFD-specific attributes before installing wrappers. Explicit worker paths and non-AFD configs are not remapped. | [`test_config_validation.py`](../../../tests/unit/compat/patches/test_config_validation.py) covers backend relaxation, four role/platform mappings, explicit and non-AFD preservation, repeated validation, unsupported platforms, and dev versions. | Remove the backend branch when upstream validation distinguishes plugin-owned ubatching; remove worker mapping when vLLM offers plugin-owned role-aware worker selection. | | [`engine_core.py`](../../../afd_plugin/compat/patches/engine_core.py): `EngineCore.__init__`, `_initialize_kv_caches`, `shutdown`; `EngineCoreProc.run_busy_loop`; `DPEngineCoreProc.run_busy_loop` | AFD FFN becomes a connector daemon: construct executor, skip scheduler/KV setup, return an empty KV-shaped result on late paths, start/monitor/stop the FFN worker loop, and use FFN-safe shutdown. Non-FFN branches copy pinned upstream behavior. | Imported by `register_afd`; Installs under the target/dev compatibility check (also permitting missing version metadata); no saved-original sentinel. Direct class assignment remains tied to the pinned upstream implementation. | [`test_engine_core.py`](../../../tests/unit/compat/patches/test_engine_core.py) covers FFN initialization, non-FFN behavior, and daemon start/stop; role runtime tests cover error propagation. | Remove when vLLM offers a headless connector-daemon engine lifecycle or an executor mode that does not require scheduler/KV ownership. | -| [`gpu_mrv2_ubatching.py`](../../../afd_plugin/compat/patches/gpu_mrv2_ubatching.py): `vllm.v1.worker.gpu.model_runner.maybe_build_ubatch_runner` and `dispatch_cg_and_sync_dp` aliases, `UBatchRunner.prepare` and `_make_forward_contexts` (instance-bound), `sync_cudagraph_and_dp_padding` (copied) | Builds the native `UBatchRunner` at DP1 with DP1 child DP metadata. FULL one-token decode microbatches keep the captured request/token split, so a DP rank holding at most half of the graph's requests replays the microbatched graph instead of running eager; a trailing microbatch without real requests keeps its padded shape. Configs without that split (speculative decoding, inner-state models, MLA, PCP/DCP) replay the regular graph instead. Other steps use native slicing and coordination. | Scoped to AFD Attention KV initialization and each `execute_model`; aliases are restored on exit. The prepare binding is per runner instance, chosen once from the config. | [`test_gpu_mrv2_ubatching.py`](../../../tests/unit/compat/patches/test_gpu_mrv2_ubatching.py) covers the DP1 factory and metadata, captured-split slicing against native vLLM slicing, padding-only microbatches, and DP1 and DP2 graph selection for underfilled ranks. | Remove the split and dispatch patches when native FULL microbatching keeps the captured split; the factory when native GPU DBO supports DP1. | | [`npu/ascend_platform.py`](../../../afd_plugin/compat/patches/npu/ascend_platform.py): `NPUPlatform.check_and_update_config` | Snapshots AFD DBO state, runs upstream normalization, and restores configured `enable_dbo`, `ubatch_size`, and `all2all_backend` in `finally`; non-AFD behavior is unchanged. | Installed through the config facade during AFD NPU normalization and verified again by the worker runtime facade; no version guard. Saves the original on the class and uses a class sentinel. The runtime facade caches success only after the wrapper is installed, so an early missing vLLM-Ascend import remains retryable. | [`test_runtime.py`](../../../tests/unit/compat/test_runtime.py) and [`test_npu_runtime.py`](../../../tests/unit/v1/worker/test_npu_runtime.py). | Remove when vLLM-Ascend recognizes plugin-owned DBO workers or no longer clears these fields. | | [`npu/ascend_config.py`](../../../afd_plugin/compat/patches/npu/ascend_config.py): `vllm_ascend.ascend_config.init_ascend_config`, loaded factory aliases, and the AFD DP child entry point | Excludes the AFD namespace from native constructor arguments; preserves config identity, validation, singleton publication, and cache refresh. DP children install the factory before native Ascend initialization. | Installed by the NPU config facade; rebinds already-loaded aliases. The DP entry point is selected only for AFD synchronous Attention or FFN. | [`test_ascend_config_lifecycle.py`](../../../tests/unit/compat/npu/test_ascend_config_lifecycle.py), [`test_ascend_engine_core_binding.py`](../../../tests/unit/compat/npu/test_ascend_engine_core_binding.py). | Remove when Ascend accepts plugin namespaces and loads general plugins before config validation in every DP child. | | [`npu/mla_graph.py`](../../../afd_plugin/compat/patches/npu/mla_graph.py): `vllm_ascend.attention.mla_v1.get_graph_params` | Resolves an AFD ubatch-owned `GraphParams` registry from the active forward context and otherwise delegates to the saved upstream process-global resolver. | Called through `apply_afd_ascend_patches_if_needed` during NPU worker startup; no version guard. Saves the original on the upstream module and uses a module sentinel. | [`test_runtime.py`](../../../tests/unit/compat/test_runtime.py) covers AFD-context resolution, upstream fallback, and idempotence; [`test_npu_mla_graph.py`](../../../tests/unit/v1/worker/test_npu_mla_graph.py) covers the owning graph lifecycle. | Remove when vLLM-Ascend accepts a forward-context-local MLA graph registry or exposes an equivalent resolver hook. | diff --git a/tests/unit/compat/patches/test_gpu_mrv2_ubatching.py b/tests/unit/compat/patches/test_gpu_mrv2_ubatching.py index 6a115f2eb..deeb1f7bf 100644 --- a/tests/unit/compat/patches/test_gpu_mrv2_ubatching.py +++ b/tests/unit/compat/patches/test_gpu_mrv2_ubatching.py @@ -8,7 +8,6 @@ from contextlib import contextmanager from dataclasses import dataclass, replace from enum import Enum -from itertools import pairwise from pathlib import Path from types import MethodType, SimpleNamespace @@ -35,38 +34,6 @@ class _BatchDescriptor: num_ubatches: int = 1 -def _vllm_config( - *, - speculative=False, - inner_state=False, - mla=False, - pcp_size=1, - dcp_size=1, - parallel=None, -): - parallel = SimpleNamespace( - **{ - **vars(parallel or _parallel()), - "prefill_context_parallel_size": pcp_size, - "decode_context_parallel_size": dcp_size, - } - ) - return SimpleNamespace( - speculative_config=object() if speculative else None, - model_config=SimpleNamespace(has_inner_state=inner_state, use_mla=mla), - parallel_config=parallel, - ) - - -UNSUPPORTED_SPLIT_CONFIGS: dict[str, dict[str, bool | int]] = { - "speculative": {"speculative": True}, - "inner-state": {"inner_state": True}, - "mla": {"mla": True}, - "pcp": {"pcp_size": 2}, - "dcp": {"dcp_size": 2}, -} - - @dataclass(frozen=True) class _DPSyncState: num_tokens_across_dp: list[int] @@ -76,9 +43,8 @@ class _DPSyncState: class _GraphManager: - def __init__(self, capture_tokens: int | None = None, vllm_config=None): + def __init__(self, capture_tokens: int | None = None): self.capture_tokens = capture_tokens - self.vllm_config = vllm_config or _vllm_config() self.calls: list[int] = [] def dispatch( @@ -117,94 +83,16 @@ def _parallel(*, dp_size=1, use_ubatching=True): num_ubatches=2, dbo_decode_token_threshold=2, dbo_prefill_token_threshold=12, - prefill_context_parallel_size=1, - decode_context_parallel_size=1, ) -@dataclass(frozen=True) -class _UBatchSlice: - request_slice: slice - token_slice: slice - - -def _bounds(slices): - return [ - ( - s.request_slice.start, - s.request_slice.stop, - s.token_slice.start, - s.token_slice.stop, - ) - for s in slices - ] - - -@dataclass(frozen=True) -class _DecodeBatch: - """The InputBatch fields prepare reads for a one-token-per-request batch.""" - - num_reqs: int - num_tokens: int - num_reqs_after_padding: int - num_tokens_after_padding: int - max_query_len: int | None = None - - -def _native_slices(input_batch, num_ubatches): - # Native vLLM 0.30 slicing clamps a trailing microbatch without real - # requests onto the last real request; only that outcome matters here. - split = input_batch.num_tokens_after_padding // num_ubatches - last_start = min(split, input_batch.num_reqs - 1) - return [ - _UBatchSlice(slice(0, split), slice(0, split)), - _UBatchSlice( - slice(last_start, input_batch.num_reqs_after_padding), - slice(split, input_batch.num_tokens_after_padding), - ), - ] - - -def _slice_decode_batch(input_batch, ubatch_slice, *args, **kwargs): - start, stop = ubatch_slice.request_slice.start, ubatch_slice.request_slice.stop - return SimpleNamespace( - num_reqs=max(0, min(stop, input_batch.num_reqs) - start), - max_query_len=input_batch.max_query_len, - seq_lens_cpu_upper_bound=[0] * (stop - start), - is_padding=ubatch_slice.token_slice, - ) - - -class _ModelState: - def prepare_attn(self, ubatch, cg_mode, block_tables, slot_mappings, *args, **kw): - return { - "max_query_len": ubatch.max_query_len, - "max_seq_len": max(ubatch.seq_lens_cpu_upper_bound), - "block_tables": block_tables, - } - - -class _SlotMappings: - def __getitem__(self, key): - return key[1] - - @pytest.fixture def patches(): # Execute the patch's functions with small native API doubles. This keeps # the real eligibility/control policy under test without importing CUDA. tree = ast.parse(PATCH_PATH.read_text()) functions: list[ast.stmt] = [ - node - for node in tree.body - if isinstance(node, ast.FunctionDef) - or ( - isinstance(node, ast.Assign) - and all( - isinstance(target, ast.Name) and target.id.isupper() - for target in node.targets - ) - ) + node for node in tree.body if isinstance(node, ast.FunctionDef) ] class NativeUBatchRunner: @@ -273,14 +161,6 @@ def check_thresholds(config, num_tokens, uniform_decode): "get_num_ubatches": lambda config: ( config.num_ubatches if config.use_ubatching else 1 ), - "pairwise": pairwise, - "UBatchSlice": _UBatchSlice, - "UBatchState": SimpleNamespace, - "create_ubatch_slices": _native_slices, - "_slice_input_batch": _slice_decode_batch, - "build_slot_mappings_by_layer": lambda slot_mappings, kv_cache_config: ( - slot_mappings - ), } future = ast.parse("from __future__ import annotations").body module = ast.Module(body=future + functions, type_ignores=[]) @@ -291,7 +171,7 @@ def check_thresholds(config, num_tokens, uniform_decode): @pytest.mark.parametrize("dp_size", [1, 2]) def test_factory_keeps_native_runner_and_changes_only_dp1_instance(patches, dp_size): parallel = _parallel(dp_size=dp_size) - config = _vllm_config(parallel=parallel) + config = SimpleNamespace(parallel_config=parallel) native_context_factory = patches.UBatchRunner._make_forward_contexts inputs = ("device", "model-state", [["attn-group"]], "kv-cache", 32) @@ -303,25 +183,13 @@ def test_factory_keeps_native_runner_and_changes_only_dp1_instance(patches, dp_s assert ("_make_forward_contexts" in vars(runner)) is (dp_size == 1) -@pytest.mark.parametrize("unsupported", [None, *UNSUPPORTED_SPLIT_CONFIGS]) -def test_factory_binds_captured_split_prepare_only_where_supported( - patches, unsupported -): - overrides = UNSUPPORTED_SPLIT_CONFIGS.get(unsupported, {}) - config = _vllm_config(parallel=_parallel(dp_size=2), **overrides) - - runner = patches.maybe_build_ubatch_runner(config, None, None, [], None, 32) - - assert ("prepare" in vars(runner)) is (unsupported is None) - - def test_disabled_dbo_does_not_construct_runner(patches): - config = _vllm_config(parallel=_parallel(use_ubatching=False)) + config = SimpleNamespace(parallel_config=_parallel(use_ubatching=False)) assert patches.maybe_build_ubatch_runner(config, None, None, [], None, 32) is None def test_dp1_native_child_metadata_keeps_stage_counts_and_context_fields(patches): - config = _vllm_config(parallel=_parallel()) + config = SimpleNamespace(parallel_config=_parallel()) runner = patches.maybe_build_ubatch_runner(config, None, None, [], None, 32) slices = [SimpleNamespace(num_tokens=3), SimpleNamespace(num_tokens=5)] attention = [{"layer": object()}, {"layer": object()}] @@ -377,24 +245,17 @@ def test_dp1_thresholds_select_eager_microbatches_without_handshake( @pytest.mark.parametrize( - "num_tokens, capture_tokens, split_supported, expected", + "num_tokens, capture_tokens, expected_mode, expected_tokens", [ - # (cg_mode, padded tokens, microbatches) - (9, 16, False, (_CUDAGraphMode.FULL, 16, 2)), - # At most half of the graph: the captured split replays it ... - (3, 16, True, (_CUDAGraphMode.FULL, 16, 2)), - (8, 16, True, (_CUDAGraphMode.FULL, 16, 2)), - # ... and without it the regular graph does, never eager. - (3, 16, False, (_CUDAGraphMode.FULL, 16, 1)), - (8, 16, False, (_CUDAGraphMode.FULL, 16, 1)), - (8, None, True, (_CUDAGraphMode.NONE, 8, 2)), + (8, 16, _CUDAGraphMode.FULL, 16), + (3, 16, _CUDAGraphMode.NONE, 3), + (8, None, _CUDAGraphMode.NONE, 8), ], ) -def test_dp1_replays_a_graph_for_an_underfilled_microbatched_step( - patches, num_tokens, capture_tokens, split_supported, expected +def test_dp1_uses_microbatch_graph_bucket_or_native_eager_padding_fallback( + patches, num_tokens, capture_tokens, expected_mode, expected_tokens ): - vllm_config = _vllm_config() if split_supported else _vllm_config(mla=True) - manager = _GraphManager(capture_tokens, vllm_config) + manager = _GraphManager(capture_tokens) desc, sync = patches.dispatch_cg_and_sync_dp( manager, num_tokens, @@ -407,101 +268,10 @@ def test_dp1_replays_a_graph_for_an_underfilled_microbatched_step( uniform_decode=True, ) assert sync is None - assert (desc.cg_mode, desc.num_tokens, desc.num_ubatches) == expected - - -@pytest.mark.parametrize( - "graph_tokens, expected", - [ - (8, [(0, 4, 0, 4), (4, 8, 4, 8)]), - (32, [(0, 16, 0, 16), (16, 32, 16, 32)]), - (96, [(0, 48, 0, 48), (48, 96, 48, 96)]), - ], -) -def test_captured_split_halves_requests_and_tokens_alike( - patches, graph_tokens, expected -): - batch = _DecodeBatch(3, 3, graph_tokens, graph_tokens) - assert _bounds(patches.captured_ubatch_slices(batch, 2)) == expected - - -@pytest.mark.parametrize( - "unsupported, min_tokens, max_tokens, capture_tokens, falls_back", - [ - (None, 15, 17, 32, False), - ("mla", 15, 17, 32, True), - ("mla", 16, 17, 32, True), - ("mla", 17, 17, 32, False), - ("speculative", 2, 3, 8, True), - ("mla", 15, 17, None, False), - ], -) -def test_regular_graph_fallback_only_without_the_captured_split( - patches, unsupported, min_tokens, max_tokens, capture_tokens, falls_back -): - config = _vllm_config(**UNSUPPORTED_SPLIT_CONFIGS.get(unsupported, {})) - manager = _GraphManager(capture_tokens, config) - assert ( - patches.needs_regular_graph_fallback( - manager, max_tokens, min_tokens, max_tokens, 1, 0, 2 - ) - is falls_back - ) - - -def _prepare(patches, input_batch, cg_mode): - runner = SimpleNamespace( - num_ubatches=2, - ubatch_query_start_loc=[None, None], - ubatch_seq_lens=[None, None], - ubatch_dcp_local_seq_lens=[None, None], - dcp_size=1, - dcp_rank=0, - cp_interleave=1, - model_state=_ModelState(), - attn_groups=[], - kv_cache_config=None, - _make_forward_contexts=lambda slices, attn, slot_mappings, padding: attn, - ) - block_tables = (list(range(input_batch.num_reqs_after_padding)),) - return patches.prepare(runner, input_batch, block_tables, _SlotMappings(), cg_mode) - - -def test_underfilled_full_decode_replays_the_captured_split(patches): - # 15 real requests padded into a 32 graph: native slicing would start the - # trailing microbatch at request 14 instead of the captured 16. - state = _prepare(patches, _DecodeBatch(15, 15, 32, 32), _CUDAGraphMode.FULL) - - assert _bounds(state.slices) == [(0, 16, 0, 16), (16, 32, 16, 32)] - trailing = state.forward_contexts[1] - assert trailing["block_tables"] == (list(range(16, 32)),) - assert trailing["max_query_len"] == 1 - assert trailing["max_seq_len"] == 1 - - -@pytest.mark.parametrize("num_reqs", [17, 24, 32]) -def test_filled_full_decode_split_matches_native_slicing(patches, num_reqs): - input_batch = _DecodeBatch(num_reqs, num_reqs, 32, 32) - state = _prepare(patches, input_batch, _CUDAGraphMode.FULL) - assert state.slices == _native_slices(input_batch, 2) - - -@pytest.mark.parametrize( - "input_batch, cg_mode", - [ - (_DecodeBatch(15, 15, 32, 32), _CUDAGraphMode.NONE), - (_DecodeBatch(15, 20, 32, 32), _CUDAGraphMode.FULL), - ], - ids=["eager", "multi-token-request"], -) -def test_other_microbatched_steps_keep_native_slicing(patches, input_batch, cg_mode): - state = _prepare(patches, input_batch, cg_mode) - - assert state.slices == _native_slices(input_batch, 2) - assert [context["max_query_len"] for context in state.forward_contexts] == [ - None, - None, - ] + assert desc.cg_mode is expected_mode + assert desc.num_tokens == expected_tokens + assert desc.num_ubatches == 2 + assert manager.calls == [1, 2] def test_profile_dispatch_does_not_select_a_graph(patches): @@ -678,118 +448,3 @@ def initialize_kv_cache(self, *args, **kwargs): ) ] assert patches.model_runner.maybe_build_ubatch_runner is original - - -@pytest.fixture -def native_patch(): - pytest.importorskip("torch") - pytest.importorskip("vllm") - from afd_plugin.compat.patches import gpu_mrv2_ubatching - - return gpu_mrv2_ubatching - - -@pytest.mark.vllm_runtime -@pytest.mark.parametrize("graph_tokens", [8, 16, 32, 64, 96, 128]) -def test_captured_split_is_native_slicing_once_real_requests_reach_it( - native_patch, graph_tokens -): - import numpy as np - from vllm.v1.worker.gpu.ubatch_utils import create_ubatch_slices - - for num_reqs in range(1, graph_tokens + 1): - input_batch = SimpleNamespace( - num_reqs=num_reqs, - num_scheduled_tokens=np.ones(num_reqs, dtype=np.int32), - num_tokens_after_padding=graph_tokens, - num_reqs_after_padding=graph_tokens, - ) - native = _bounds(create_ubatch_slices(input_batch, 2)) - captured = _bounds(native_patch.captured_ubatch_slices(input_batch, 2)) - assert (native == captured) is (2 * num_reqs > graph_tokens), num_reqs - - -class _NativeGraphManager: - CAPTURE_SIZES = (8, 16, 32, 64) - - def __init__(self, vllm_config): - self.vllm_config = vllm_config - - def dispatch( - self, - num_reqs, - num_tokens, - uniform_token_count, - *, - num_active_loras=0, - max_query_len=None, - num_ubatches=1, - ): - from vllm.config import CUDAGraphMode - from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor - - size = next(size for size in self.CAPTURE_SIZES if size >= num_tokens) - return BatchExecutionDescriptor( - cg_mode=CUDAGraphMode.FULL, - num_tokens=size, - num_reqs=size, - uniform_token_count=uniform_token_count, - num_ubatches=num_ubatches, - ) - - -@pytest.mark.vllm_runtime -@pytest.mark.parametrize( - "tokens_across_dp, split_supported, expected", - [ - # (padded tokens, microbatches); every case replays a FULL graph. - ((15, 17), True, (32, 2)), - ((15, 17), False, (32, 1)), - ((16, 32), True, (32, 2)), - ((16, 32), False, (32, 1)), - ((16, 16), False, (16, 2)), - ], -) -def test_dp2_underfilled_rank_replays_a_graph_on_every_rank( - native_patch, monkeypatch, tokens_across_dp, split_supported, expected -): - import torch - from vllm.config import CUDAGraphMode - from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor - - def all_reduce(tensor, group=None): - # Peers vote like the calling rank except for their own batch size. - local_vote = tensor.sum(dim=1) - for rank, tokens in enumerate(tokens_across_dp): - tensor[:, rank] = local_vote - tensor[0, rank] = tokens - tensor[5, rank] = tokens - - monkeypatch.setattr(native_patch.dist, "all_reduce", all_reduce) - monkeypatch.setattr( - native_patch, "get_dp_group", lambda: SimpleNamespace(cpu_group=None) - ) - vllm_config = _vllm_config() if split_supported else _vllm_config(mla=True) - manager = _NativeGraphManager(vllm_config) - - results = [] - for rank, tokens in enumerate(tokens_across_dp): - desc, sync = native_patch.sync_cudagraph_and_dp_padding( - manager, - BatchExecutionDescriptor(CUDAGraphMode.FULL, tokens, tokens, 1), - tokens, - tokens, - 1, - len(tokens_across_dp), - rank, - parallel_config=_parallel(dp_size=len(tokens_across_dp)), - allow_ubatching=True, - uniform_decode=True, - ) - assert desc.cg_mode == CUDAGraphMode.FULL - assert sync is not None and not sync.eager - assert torch.equal( - sync.num_tokens_across_dp, torch.full((2,), expected[0], dtype=torch.int32) - ) - results.append((desc.num_tokens, desc.num_ubatches)) - assert results == [expected] * len(tokens_across_dp) From 2a49124c482c3aac2c0257fd8dfaf9f63b2dba1e Mon Sep 17 00:00:00 2001 From: Ronen Kat <16743404+ronenkat@users.noreply.github.com> Date: Mon, 2 Nov 2026 19:17:43 +0200 Subject: [PATCH 28/32] Update Qwen3 graph capture sizes --- .../prefill_decode_colocation/4a2f_graph_dp4tp1.sh | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/4a2f_graph_dp4tp1.sh b/recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/4a2f_graph_dp4tp1.sh index 19e86f2cc..960b7998d 100755 --- a/recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/4a2f_graph_dp4tp1.sh +++ b/recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/4a2f_graph_dp4tp1.sh @@ -22,7 +22,7 @@ # MAX_NUM_SEQS per Attention DP rank (256) # MAX_NUM_BATCHED_TOKENS prefill chunk per step (3072) # CUDAGRAPH_CAPTURE_SIZES comma-separated decode graph sizes -# (8,16,32,48,64,96,112,128,192,256) +# (8,12,16,24,32,48,64,96,128,192,256) # MAX_MODEL_LEN (4096) # GPU_MEMORY_UTILIZATION (unset: vLLM default) # ATTN_DEVICES / FFN_DEVICES (0,1,2,3 / 4,5) @@ -57,7 +57,7 @@ DBO_DECODE_TOKEN_THRESHOLD=${DBO_DECODE_TOKEN_THRESHOLD:-2} DBO_PREFILL_TOKEN_THRESHOLD=${DBO_PREFILL_TOKEN_THRESHOLD:-2} MAX_NUM_SEQS=${MAX_NUM_SEQS:-256} MAX_NUM_BATCHED_TOKENS=${MAX_NUM_BATCHED_TOKENS:-3072} -CUDAGRAPH_CAPTURE_SIZES=${CUDAGRAPH_CAPTURE_SIZES:-8,16,32,48,64,96,112,128,192,256} +CUDAGRAPH_CAPTURE_SIZES=${CUDAGRAPH_CAPTURE_SIZES:-8,12,16,24,32,48,64,96,128,192,256} MAX_MODEL_LEN=${MAX_MODEL_LEN:-4096} GPU_MEMORY_UTILIZATION=${GPU_MEMORY_UTILIZATION:-} ATTN_DEVICES=${ATTN_DEVICES:-0,1,2,3} From 5165ba336fae297ac489bc8dce919d496e1c48cc Mon Sep 17 00:00:00 2001 From: ronenkat <16743404+ronenkat@users.noreply.github.com> Date: Mon, 2 Nov 2026 20:43:43 +0200 Subject: [PATCH 29/32] [Doc] DeepEP M2N: drop the transport log line from the placement docs The plugin's INFO lines are not printed under vLLM's logging setup, so the docs should not point users at the transport log line. A placement the connector cannot carry still fails at startup with the pods' shapes. Co-Authored-By: Claude Opus 5.5 Signed-off-by: ronenkat <16743404+ronenkat@users.noreply.github.com> --- docs/gpu/DEEPEP_M2N_CONNECTOR_USER_GUIDE.md | 4 +--- .../gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/README.md | 5 ++--- 2 files changed, 3 insertions(+), 6 deletions(-) diff --git a/docs/gpu/DEEPEP_M2N_CONNECTOR_USER_GUIDE.md b/docs/gpu/DEEPEP_M2N_CONNECTOR_USER_GUIDE.md index 48d7c2039..8056bab7e 100644 --- a/docs/gpu/DEEPEP_M2N_CONNECTOR_USER_GUIDE.md +++ b/docs/gpu/DEEPEP_M2N_CONNECTOR_USER_GUIDE.md @@ -81,9 +81,7 @@ architecture (SM90 or newer), and the FFN deployment (expert parallel, EP size, FFN TP=1, linear expert placement, no EPLB, all2all backend and expert weight format), and the placement of the ranks over NVLink domains. The FFN adopts the Attention TP size, dispatch capacity and bucket ladder reported by -the Attention ranks. FFN rank 0 logs the network plan and the transport it -picked, for example -`DeepEP M2N transport: per_domain f=1 domains=2x5 placement=4A1F,4A1F`. +the Attention ranks. ## Installation diff --git a/recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/README.md b/recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/README.md index 0e64727a1..c9bd1a761 100644 --- a/recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/README.md +++ b/recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/README.md @@ -54,9 +54,8 @@ and `13355` (FFN DP). Pass the fabric's NCCL settings (for example (`2 x 6A2F`, ...); the tuning values (`MAX_NUM_SEQS`, `MAX_NUM_BATCHED_TOKENS`, `MAX_MODEL_LEN`, `CAPTURE_SIZES`, `PAYLOAD_FORMAT`) and DBO (`ENABLE_DBO=1`, `DBO_DECODE_TOKEN_THRESHOLD`, `DBO_PREFILL_TOKEN_THRESHOLD`) are environment -variables with the 4A2F values as defaults. The connector logs the placement -it found and the afdEP-comm mode it picked, for example -`DeepEP M2N transport: per_domain f=1 domains=2x5 placement=4A1F,4A1F`. +variables with the 4A2F values as defaults. A placement the connector cannot +carry fails at startup with the pods' shapes in the error. ## Configuration notes From f3d55a7a2cc01604d6c4cefc6ecdebe7d6d92b6f Mon Sep 17 00:00:00 2001 From: ronenkat <16743404+ronenkat@users.noreply.github.com> Date: Wed, 4 Nov 2026 10:03:59 +0200 Subject: [PATCH 30/32] recipe(gpu): one LAYOUT-driven graph/eager/dbo recipe per M2N model Replace the eleven DeepEPM2NAFDConnector recipe scripts (2a2f_*, 4a2f_*, multipod_4a1f_4a1f_graph.sh) with graph.sh, eager.sh and dbo.sh per model (DeepSeek-V2-Lite, Qwen3-30B-A3B, Qwen3-235B-A22B-FP8), about 95 lines each instead of up to 191. - LAYOUT (default 2A2F) lists one AF entry per pod; POD=AFD_ selects the entry. The script derives devices (FFN first), rank totals, DP size/local/start rank, headless and the leader of each role, so one script runs single-host, symmetric multi-pod (1A1F,1A1F, 4A1F,4A1F) and role-split (2A0F,0A2F) placements. ATTN_TP replaces the dp1tp2 scripts. - Run-specific flags move out of the recipe: VLLM_ARGS, ATTN_ARGS, FFN_ARGS and script arguments are appended last, so they override the defaults. DECODE_BENCH takes the decode-only benchmark's kv_connector_extra_config JSON and replaces KV_CONNECTOR and the DECODE_BENCH_* knobs. - Each script stays self-contained (the deploy skill mounts one file); the block below its model settings is identical in every recipe. - The P2P Qwen3 entry points now run graph.sh (235B keeps LAYOUT=4A2F). Docs: the connector recipes README holds the shared layout, settings, multi-pod and benchmark instructions; model READMEs keep only their defaults and notes; tools/benchmarks/README.md holds the decode-only benchmark mechanics. The deploy-afd-k8s skill reads placement and GPU counts from LAYOUT and passes env values with spaces or JSON as quoted YAML. Validated with Qwen3-30B-A3B on H100 (vLLM 0.32.0): 2A2F graph, eager, DBO, ATTN_TP=2, flag overrides and the P2P entry point serve correct completions; decode-only benchmark (500 x 1025-token pool, 2000 requests, OSL 256) passes its checks on 2A2F (13.6 req/s, TPOT p50 18.0 ms) and on two-node 2A0F,0A2F and 1A1F,1A1F (11.0 req/s, TPOT p50 22 ms), plus smoke on both multi-pod layouts. Signed-off-by: ronenkat Co-Authored-By: Claude Opus 5.5 --- .agents/skills/deploy-afd-k8s/SKILL.md | 8 +- .../deploy-afd-k8s/references/deploy.md | 47 ++-- .../references/resolve-recipe.md | 51 ++-- .../deploy-afd-k8s/references/services.md | 8 +- docs/gpu/DEEPEP_M2N_CONNECTOR_USER_GUIDE.md | 7 +- recipe/README.md | 4 +- recipe/gpu/DeepEPM2NAFDConnector/README.md | 230 +++++++++++++++--- .../deepseek_v2_lite/README.md | 59 ++--- .../2a2f_eager_dp2tp1.sh | 55 ----- .../2a2f_graph_dbo_dp2tp1.sh | 64 ----- .../2a2f_graph_dp1tp2.sh | 58 ----- .../2a2f_graph_dp2tp1.sh | 57 ----- .../prefill_decode_colocation/dbo.sh | 98 ++++++++ .../prefill_decode_colocation/eager.sh | 95 ++++++++ .../prefill_decode_colocation/graph.sh | 97 ++++++++ .../qwen3_235b_a22b_fp8/README.md | 206 ++++------------ .../4a2f_graph_dbo_dp4tp1.sh | 68 ------ .../4a2f_graph_dp4tp1.sh | 191 --------------- .../prefill_decode_colocation/dbo.sh | 98 ++++++++ .../prefill_decode_colocation/eager.sh | 95 ++++++++ .../prefill_decode_colocation/graph.sh | 97 ++++++++ .../multipod_4a1f_4a1f_graph.sh | 82 ------- .../qwen3_30b_a3b/README.md | 201 +++------------ .../2a2f_eager_dp2tp1.sh | 53 ---- .../2a2f_graph_dbo_dp2tp1.sh | 62 ----- .../2a2f_graph_dp1tp2.sh | 56 ----- .../2a2f_graph_dp2tp1.sh | 153 ------------ .../prefill_decode_colocation/dbo.sh | 98 ++++++++ .../prefill_decode_colocation/eager.sh | 95 ++++++++ .../prefill_decode_colocation/graph.sh | 97 ++++++++ .../qwen3_235b_a22b_fp8/README.md | 11 +- .../4a2f_graph_dp4tp1.sh | 11 +- .../qwen3_30b_a3b/README.md | 8 +- .../2a2f_graph_dp2tp1.sh | 8 +- .../qwen3_235b_a22b_fp8/vanilla_tp4.sh | 2 +- tools/benchmarks/README.md | 199 +++++++-------- 36 files changed, 1364 insertions(+), 1465 deletions(-) delete mode 100644 recipe/gpu/DeepEPM2NAFDConnector/deepseek_v2_lite/prefill_decode_colocation/2a2f_eager_dp2tp1.sh delete mode 100644 recipe/gpu/DeepEPM2NAFDConnector/deepseek_v2_lite/prefill_decode_colocation/2a2f_graph_dbo_dp2tp1.sh delete mode 100644 recipe/gpu/DeepEPM2NAFDConnector/deepseek_v2_lite/prefill_decode_colocation/2a2f_graph_dp1tp2.sh delete mode 100644 recipe/gpu/DeepEPM2NAFDConnector/deepseek_v2_lite/prefill_decode_colocation/2a2f_graph_dp2tp1.sh create mode 100644 recipe/gpu/DeepEPM2NAFDConnector/deepseek_v2_lite/prefill_decode_colocation/dbo.sh create mode 100644 recipe/gpu/DeepEPM2NAFDConnector/deepseek_v2_lite/prefill_decode_colocation/eager.sh create mode 100644 recipe/gpu/DeepEPM2NAFDConnector/deepseek_v2_lite/prefill_decode_colocation/graph.sh delete mode 100644 recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/4a2f_graph_dbo_dp4tp1.sh delete mode 100755 recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/4a2f_graph_dp4tp1.sh create mode 100644 recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/dbo.sh create mode 100644 recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/eager.sh create mode 100644 recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/graph.sh delete mode 100644 recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/multipod_4a1f_4a1f_graph.sh delete mode 100644 recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/2a2f_eager_dp2tp1.sh delete mode 100644 recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/2a2f_graph_dbo_dp2tp1.sh delete mode 100644 recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/2a2f_graph_dp1tp2.sh delete mode 100755 recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/2a2f_graph_dp2tp1.sh create mode 100644 recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/dbo.sh create mode 100644 recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/eager.sh create mode 100644 recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/graph.sh diff --git a/.agents/skills/deploy-afd-k8s/SKILL.md b/.agents/skills/deploy-afd-k8s/SKILL.md index ff84023c5..5468f24eb 100644 --- a/.agents/skills/deploy-afd-k8s/SKILL.md +++ b/.agents/skills/deploy-afd-k8s/SKILL.md @@ -46,8 +46,12 @@ than one pod) -- callers never talk to these directly. See `# For pod N: POD=...` line per pod, but the comment isn't guaranteed to use that exact grammar going forward. See [resolve-recipe.md](references/resolve-recipe.md). -- Not for the local decode-only benchmark (`tools/benchmarks/decode_bench_client.py` - against `recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/`, non-k8s). +- `DeepEPM2NAFDConnector` recipes take their placement from `LAYOUT` (one + pod per `AF` entry) and run-specific settings as env (`VLLM_ARGS`, + `DECODE_BENCH`, ...); see [resolve-recipe.md](references/resolve-recipe.md). + This skill only serves; the decode-only benchmark's client + (`tools/benchmarks/decode_bench_client.py`) is run separately against the + endpoint. ## Requirements diff --git a/.agents/skills/deploy-afd-k8s/references/deploy.md b/.agents/skills/deploy-afd-k8s/references/deploy.md index 1957efadd..c351725e7 100644 --- a/.agents/skills/deploy-afd-k8s/references/deploy.md +++ b/.agents/skills/deploy-afd-k8s/references/deploy.md @@ -28,12 +28,13 @@ devices="$(role_block "$RECIPE_SCRIPT_PATH" "$ROLE" \ GPU_COUNT=$(( $(grep -o ',' <<<"$devices" | wc -l) + 1 )) ``` -`POD=AFD_` entries (both roles per pod, see -[resolve-recipe.md](resolve-recipe.md)): every pod has the same shape, +`LAYOUT` recipes (see [resolve-recipe.md](resolve-recipe.md)): each entry's +GPU count is its `a + f`: ```bash -default() { grep -oE "$1:-[0-9]+" "$RECIPE_SCRIPT_PATH" | grep -oE '[0-9]+$'; } -GPU_COUNT=$(( ${ATTN_PER_POD:-$(default ATTN_PER_POD)} + ${FFN_PER_POD:-$(default FFN_PER_POD)} )) +LAYOUT=${LAYOUT:-$(grep -oE 'LAYOUT:-[0-9AF,]+' "$RECIPE_SCRIPT_PATH" | cut -d- -f2)} +IFS=, read -ra ENTRIES <<<"$LAYOUT" +for e in "${ENTRIES[@]}"; do f=${e#*A}; echo "$e -> $(( ${e%%A*} + ${f%F} )) GPUs"; done ``` ## 3. Confirm before touching the cluster @@ -138,17 +139,19 @@ recipe) against the one shared `templates/pod.yaml`: ```bash deploy_pod() { local pod_name="$1" gpu="$2" attn_role="$3" ffn_role="$4" extra_env_pairs="$5" - # extra_env_pairs: space-separated VAR=value tokens straight off this - # plan entry's header line (including POD=... itself), empty for a plain - # recipe. Build the spliced env-entries block for templates/pod.yaml. - # Leading (not trailing) newlines keep the `# ${TEMPLATE_EXTRA_ENV}` - # placeholder line a harmless comment when this is empty, and valid - # sibling list items when it isn't. - local extra_env="" - for pair in $extra_env_pairs; do + # extra_env_pairs: one VAR=value per line -- this plan entry's overrides + # (including POD=... itself) plus the caller's settings -- empty for a plain + # recipe. Values may hold spaces and JSON (VLLM_ARGS, DECODE_BENCH); they + # become single-quoted YAML strings. Leading (not trailing) newlines keep + # the `# ${TEMPLATE_EXTRA_ENV}` placeholder line a harmless comment when + # this is empty, and valid sibling list items when it isn't. + local extra_env="" pair val q="'" + while IFS= read -r pair; do + [ -n "$pair" ] || continue + val=${pair#*=} extra_env="$extra_env - - {name: ${pair%%=*}, value: \"${pair#*=}\"}" - done + - {name: ${pair%%=*}, value: '${val//$q/$q$q}'}" + done <<<"$extra_env_pairs" kubectl delete pod "$pod_name" --ignore-not-found @@ -172,17 +175,25 @@ Placement-driven recipe, one call per entry parsed in ```bash deploy_pod vllm-pod-attention-0 2 head none "POD=ATTENTION_0" -deploy_pod vllm-pod-attention-1 2 worker none "POD=ATTENTION_1 ATTENTION_HEADLESS=1 ATTENTION_DP_START_RANK=2" +deploy_pod vllm-pod-attention-1 2 worker none $'POD=ATTENTION_1\nATTENTION_HEADLESS=1\nATTENTION_DP_START_RANK=2' deploy_pod vllm-pod-ffn-0 2 none head "POD=FFN_0" ``` -`POD=AFD_` recipe, e.g. the DeepEP M2N 2 x 4A1F plan (5 GPUs per pod): +`LAYOUT` recipe, e.g. `LAYOUT=2A0F,0A2F` (2 GPUs per pod) with the +decode-only benchmark, or `LAYOUT=4A1F,4A1F` (5 GPUs per pod): ```bash -deploy_pod vllm-pod-afd-0 5 head head "POD=AFD_0" -deploy_pod vllm-pod-afd-1 5 worker worker "POD=AFD_1" +SETTINGS=$'LAYOUT=2A0F,0A2F\nDECODE_BENCH={"num_prompts": 500, "max_output_len": 256}\nAFD_REPO=/opt/afd-plugin' +deploy_pod vllm-pod-afd-0 2 head none $'POD=AFD_0\n'"$SETTINGS" +deploy_pod vllm-pod-afd-1 2 none head $'POD=AFD_1\n'"$SETTINGS" + +deploy_pod vllm-pod-afd-0 5 head head $'POD=AFD_0\nLAYOUT=4A1F,4A1F' +deploy_pod vllm-pod-afd-1 5 worker worker $'POD=AFD_1\nLAYOUT=4A1F,4A1F' ``` +A one-entry `LAYOUT` with settings is a plain pod with env, e.g. +`deploy_pod vllm-pod 4 head head $'VLLM_ARGS=--max-num-seqs 32'`. + ## 4f. Wait for Running, then for each pod's own readiness line Bounded 15-minute wait per pod for `phase=Running` (unschedulable pods never diff --git a/.agents/skills/deploy-afd-k8s/references/resolve-recipe.md b/.agents/skills/deploy-afd-k8s/references/resolve-recipe.md index cc27b69d9..d5eeae63f 100644 --- a/.agents/skills/deploy-afd-k8s/references/resolve-recipe.md +++ b/.agents/skills/deploy-afd-k8s/references/resolve-recipe.md @@ -90,21 +90,36 @@ For each entry, derive (used by [deploy.md](deploy.md) and else `head`. - **`afd-ffn-node-role`**: symmetric, `none` unless `POD` starts with `FFN_`. -### Pods that run both roles: `POD=AFD_` - -DeepEP M2N multi-pod recipes (e.g. -`DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/multipod_4a1f_4a1f_graph.sh`) -list `POD=AFD_0`, `POD=AFD_1`, ... with no other per-pod overrides. Every -such pod runs one FFN and one Attention server and derives its DP flags -(headless, start rank) from its own index, so: - -- **`afd-attn-node-role`** and **`afd-ffn-node-role`**: both `head` for - `AFD_0`, both `worker` for every other `AFD_`. -- The pod shape comes from the script's `NUM_PODS`, `ATTN_PER_POD` and - `FFN_PER_POD` defaults (`NUM_PODS` must equal the number of entries). If - the user overrides them, pass the same values to every pod as container - env. -- Ports are literals in the script, not `*_PORT:-N` defaults: the AFD - rendezvous port is the `"port": N` of the `afd` JSON, and the FFN and - Attention DP RPC ports are the last argument of the `role_dp "$F" ...` and - `role_dp "$A" ...` lines. +### `LAYOUT` recipes: `POD=AFD_` + +The `DeepEPM2NAFDConnector` recipes (`graph.sh`, `eager.sh`, `dbo.sh` under +`recipe/gpu/DeepEPM2NAFDConnector//prefill_decode_colocation/`) state +their placement through `LAYOUT`: one pod per comma-separated `AF` +entry, in order, pod `k` running `POD=AFD_`, every pod with the same +`LAYOUT`. Take `LAYOUT` from the caller's settings, else from the script's +`LAYOUT=${LAYOUT:-...}` default; never ask. Each pod derives its own DP flags +(headless, start rank), devices and roles from `LAYOUT` and `POD`, so for +entry `k` with `a` Attention and `f` FFN ranks: + +- **k8s Pod name**: `vllm-pod-afd-`. A one-entry `LAYOUT` (the default + `2A2F`) is a plain single pod: `vllm-pod`, no `POD` needed. +- **GPU count**: `a + f`. The recipe computes `CUDA_VISIBLE_DEVICES` (FFN on + the pod's first `f` GPUs), so there is no literal to read. +- **`afd-attn-node-role`**: `none` when `a` is 0, `head` on the first entry + with `a > 0`, `worker` on later ones. **`afd-ffn-node-role`**: the same + with `f`. For example `2A0F,0A2F` is `head none` and `none head`, and + `1A1F,1A1F` is `head head` and `worker worker`. +- **Container env**: `POD=AFD_` and `LAYOUT=`, plus every setting + the caller gives (`MODEL_PATH`, `ATTN_TP`, `VLLM_ARGS`, `ATTN_ARGS`, + `FFN_ARGS`, `DECODE_BENCH`, `M2N_EXTRA_CONFIG`, NCCL variables), the same on + every pod. Leave `HEAD_HOST`, `ATTN_HEAD_HOST` and `FFN_HEAD_HOST` unset: + the defaults name the Services from [services.md](services.md). With + `DECODE_BENCH`, also set `AFD_REPO=/opt/afd-plugin` (the repository in + `docker/Dockerfile.ci` images), since the connector lives in + `tools/benchmarks` and the recipe runs from `/work`. +- **Ports** are literals: `--port 18000` on the Attention `vllm serve` + (`CLIENT_PORT`), the AFD rendezvous `"port": 6239` in the `afd` JSON, and + the DP RPC ports as the third argument of the two `role_dp` calls (`13355` + FFN, `13345` Attention). +- Whether the connector accepts the placement is checked at startup, not + here; see the connector guide's supported layouts. diff --git a/.agents/skills/deploy-afd-k8s/references/services.md b/.agents/skills/deploy-afd-k8s/references/services.md index 761d23e65..a3bc18648 100644 --- a/.agents/skills/deploy-afd-k8s/references/services.md +++ b/.agents/skills/deploy-afd-k8s/references/services.md @@ -39,10 +39,10 @@ AFD_CONNECTOR_PORT="$AFD_CONNECTOR_PORT" envsubst '${AFD_CONNECTOR_PORT}' \ Only when **more than one** plan entry's `POD` carries the same role (i.e. that role is split across pods, not just placed in its own single pod). -Every `POD=AFD_` entry carries both roles, so a plan of two or more -such entries needs both Services; read their ports (and the rendezvous port) -as [resolve-recipe.md](resolve-recipe.md) describes for those recipes rather -than with the `*_PORT:-N` greps below. +For a `LAYOUT` recipe, a role is split when more than one entry has ranks of +it (both roles in `1A1F,1A1F`, neither in `2A0F,0A2F`); read its ports (and +the rendezvous port) as [resolve-recipe.md](resolve-recipe.md) describes for +those recipes rather than with the `*_PORT:-N` greps below. Carries that role's DP-RPC coordination traffic to its head pod. One per split role, generic over role via `templates/service-dp-coord.yaml`: diff --git a/docs/gpu/DEEPEP_M2N_CONNECTOR_USER_GUIDE.md b/docs/gpu/DEEPEP_M2N_CONNECTOR_USER_GUIDE.md index 20515c0ff..442d62e1b 100644 --- a/docs/gpu/DEEPEP_M2N_CONNECTOR_USER_GUIDE.md +++ b/docs/gpu/DEEPEP_M2N_CONNECTOR_USER_GUIDE.md @@ -291,10 +291,9 @@ the headless pods. The FFN DP rank is the FFN's EP rank, so it decides which experts a pod loads; `--enable-ep-weight-filter` keeps each FFN from reading the other experts' weights. -The recipe -[`multipod_4a1f_4a1f_graph.sh`](../../recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/README.md#multi-pod-2-x-4a1f) -derives all of these flags for symmetric placements `N x AF` from the -pod index (`POD=AFD_`). +The [DeepEP M2N recipes](../../recipe/gpu/DeepEPM2NAFDConnector/README.md#layout) +derive all of these flags from a `LAYOUT` such as `2A0F,0A2F` or `4A1F,4A1F` +and the pod index (`POD=AFD_`). ## Minimal launch shape diff --git a/recipe/README.md b/recipe/README.md index 65a6a0502..ffdc157e4 100644 --- a/recipe/README.md +++ b/recipe/README.md @@ -41,8 +41,8 @@ Directory names follow these conventions: | Hardware | Connector | Model | Recommended stage | Status | Recipe | | --- | --- | --- | --- | --- | --- | | GPU | `P2pNcclAFDConnector` | DeepSeek-V2-Lite | Decode | Deployment examples; see the linked README for validation scope | [Launch examples](gpu/P2pNcclAFDConnector/deepseek_v2_lite/README.md) | -| GPU (SM90+ for `DeepEPM2NAFDConnector`) | `DeepEPM2NAFDConnector` or `P2pNcclAFDConnector` | Qwen3-30B-A3B | Prefill / decode; decode-only benchmark | Same 2A2F layouts as the Qwen3 MoE DeepEP M2N E2E cases; serving smoke and decode-only benchmark on both connectors | [2A2F single node](gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/README.md) | -| GPU (SM90+ for `DeepEPM2NAFDConnector`) | `DeepEPM2NAFDConnector` or `P2pNcclAFDConnector` | Qwen3-235B-A22B-FP8 | Prefill / decode; decode-only benchmark | 4A2F FP8-payload configuration and decode-only benchmark run on H200 | [4A2F single node](gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/README.md) | +| GPU (SM90+ for `DeepEPM2NAFDConnector`) | `DeepEPM2NAFDConnector` or `P2pNcclAFDConnector` | Qwen3-30B-A3B | Prefill / decode; decode-only benchmark | Graph, eager and DBO scripts for any `LAYOUT` (default 2A2F; multi-pod such as `2A0F,0A2F` or `1A1F,1A1F`), the layouts of the Qwen3 MoE DeepEP M2N E2E cases | [Graph, eager, DBO](gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/README.md) | +| GPU (SM90+ for `DeepEPM2NAFDConnector`) | `DeepEPM2NAFDConnector` or `P2pNcclAFDConnector` | Qwen3-235B-A22B-FP8 | Prefill / decode; decode-only benchmark | FP8-payload graph, eager and DBO scripts for any `LAYOUT` (2A2F, 4A2F, 2 x 4A1F, ...); H200-class FFN GPUs | [Graph, eager, DBO](gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/README.md) | | Ascend NPU | `CAMP2pAFDConnector` | DeepSeek-V3.2 | Decode | Historical v0.19.1rc1 recipe; not qualified for the 0.32 target | [Synchronous decode](npu/CAMP2pAFDConnector/deepseek_v3_2/README.md) | | Ascend NPU | `CAMAsyncAFDConnector` | DeepSeek-V3.2 | Prefill / decode | Historical v0.19.1rc1/v0.26 recipes, not qualified for 0.32; v0.26 DP2TP8+EP16 token split reached `0.9522` strict match; legacy PCP8 results are v0.19-only | [Async CAM](npu/CAMAsyncAFDConnector/deepseek_v3_2/README.md) | | Ascend NPU | `CAMAsyncAFDConnector` | DeepSeek-V4-Flash W4A8 DSpark | Separate P / decode, or colocated | v0.32, single-host eager; native full-draft P TP4, A DP1/DP2 TP1, F EP4 with legacy or layered W4A8 | [Eager DSpark](npu/CAMAsyncAFDConnector/deepseek_v4_flash/README.md) | diff --git a/recipe/gpu/DeepEPM2NAFDConnector/README.md b/recipe/gpu/DeepEPM2NAFDConnector/README.md index bfbb024dd..07b4c254c 100644 --- a/recipe/gpu/DeepEPM2NAFDConnector/README.md +++ b/recipe/gpu/DeepEPM2NAFDConnector/README.md @@ -2,34 +2,42 @@ Launch scripts for running Attention-FFN Disaggregation (AFD) with `DeepEPM2NAFDConnector`, the expert-routed M2N connector for Hopper (SM90) and -newer GPUs. See +newer GPUs, with vLLM `v0.32.0`. See [`docs/gpu/DEEPEP_M2N_CONNECTOR_USER_GUIDE.md`](../../../docs/gpu/DEEPEP_M2N_CONNECTOR_USER_GUIDE.md) for installation, the configuration contract, the required vLLM settings, and -supported layouts. +the placements the connector accepts. -## Directory layout +This README covers what every recipe here shares; each model's `README.md` +lists only its defaults and model-specific notes. + +## Recipes + +Every model directory holds the same three scripts in +`prefill_decode_colocation/`: + +| Script | Mode | +| --- | --- | +| `graph.sh` | `FULL_DECODE_ONLY` CUDA graphs; prefill runs eager | +| `eager.sh` | Eager (`--enforce-eager`) | +| `dbo.sh` | `FULL_DECODE_ONLY` CUDA graphs and dual-batch overlap on both roles (`--enable-dbo`, decode/prefill token thresholds 2 and 8) | ```text . -├── deepseek_v2_lite/ -│ ├── README.md -│ └── prefill_decode_colocation/ # 2A2F: eager, graph, graph + DBO -├── qwen3_30b_a3b/ -│ ├── README.md -│ └── prefill_decode_colocation/ # 2A2F: eager, graph, graph + DBO, Attention TP2 -└── qwen3_235b_a22b_fp8/ - ├── README.md - └── prefill_decode_colocation/ # 4A2F FP8 payload: graph, graph + DBO; - # 2 pods x 4A1F (multipod_4a1f_4a1f_graph.sh) +├── deepseek_v2_lite/ # BF16 payload +├── qwen3_30b_a3b/ # BF16 payload +└── qwen3_235b_a22b_fp8/ # FP8 payload (fp8_block128_fp32) ``` -Each model directory holds its own scripts; open its `README.md` before -running anything in it. +Each script is self-contained, so it can be copied or mounted on its own (the +`deploy-afd-k8s` skill mounts the single file). Above the line +`The rest is the same in every DeepEPM2NAFDConnector recipe` a script sets its +model path, layout default, tuning flags (`TUNING`), connector extras and +mode flags (`MODE`); everything below that line is identical in all of them. ## Prerequisites -- SM90 or newer GPUs on one node: 4 for the DeepSeek-V2-Lite and - Qwen3-30B-A3B recipes, 6 for the Qwen3-235B-A22B-FP8 recipes. +- SM90 or newer GPUs. The default `2A2F` layout uses 4 GPUs on one host; see + the model README for memory needs. - vLLM `v0.32.0`, the plugin, afdEP-comm (the `afdep` extra) and DeepGEMM installed in one environment, with the elfutils headers available to DeepGEMM's JIT. See @@ -42,37 +50,183 @@ running anything in it. export AFDEP_NCCL_ROOT=$NCCL LD_LIBRARY_PATH=$NCCL/lib:${LD_LIBRARY_PATH:-} ``` +- Free ports `6239` (AFD rendezvous), `18000` (Attention API) and `18001` + (FFN, serves no requests); across pods also `13345` (Attention DP RPC) and + `13355` (FFN DP RPC). + The scripts call `vllm serve` directly rather than through `uv run`, so they use the environment that holds afdEP-comm and DeepGEMM. -## Running on a local host +## Layout + +`LAYOUT` lists the pods in order, one `AF` entry each: `a` Attention +ranks and `f` FFN ranks. Inside a pod, the FFN ranks take the first `f` GPUs +and the Attention ranks the next `a` GPUs; every pod runs the same script with +the same `LAYOUT` and its own `POD=AFD_` (`k` counts from 0; one pod needs +no `POD`). + +| `LAYOUT` | Pods | Per pod | +| --- | --- | --- | +| `2A2F` (default) | 1 | FFN DP2 on GPUs 0,1, Attention DP2 on GPUs 2,3 | +| `4A2F` | 1 | FFN on GPUs 0,1, Attention DP4 on GPUs 2-5 | +| `2A0F,0A2F` | 2 | pod 0: Attention DP2; pod 1: FFN DP2 | +| `1A1F,1A1F` | 2 | FFN on GPU 0, Attention on GPU 1 | +| `4A1F,4A1F` | 2 | FFN on GPU 0, Attention on GPUs 1-4 | + +`ATTN_TP` splits each pod's Attention GPUs into TP groups (it must divide each +pod's `a`), so `ATTN_TP=2` runs `2A2F`'s Attention as one DP1 x TP2 engine; +`num_attention_ranks` still counts GPUs. The FFN always runs TP 1, and its DP +rank (its EP rank) decides which experts it loads. + +The script derives every DP flag from `LAYOUT` and `POD`. A role held by one +pod gets only `--data-parallel-size`; a role spread over pods is led by its +first pod, and the others start headless from their start rank. The +connector checks the placement at startup and names the pods' shapes when it +rejects one; the +[supported layouts](../../../docs/gpu/DEEPEP_M2N_CONNECTOR_USER_GUIDE.md#supported-layouts) +table lists which placements it accepts (for example `4A0F,0A2F` is not, its +pods differ in size). + +## Recipe settings + +Every setting is an environment variable; the defaults reproduce the +published recipe. + +| Variable | Default | Meaning | +| --- | --- | --- | +| `MODEL_PATH` | `/path/model_weights/` | Weights directory or HF id. | +| `LAYOUT` | `2A2F` | Pods and their ranks; see [Layout](#layout). | +| `POD` | `AFD_0` | This pod's entry in `LAYOUT`, `AFD_`. | +| `ATTN_TP` | `1` | Attention TP size. | +| `CUDAGRAPH_CAPTURE_SIZES` | per model | Ascending decode graph sizes (`graph.sh`, `dbo.sh`); the largest is the maximum capture size. | +| `VLLM_ARGS` | empty | Extra `vllm serve` flags for both roles. | +| `ATTN_ARGS` / `FFN_ARGS` | empty | Extra `vllm serve` flags for one role. | +| `DECODE_BENCH` | empty | Turns on the decode-only benchmark; see [Decode-only benchmark](#decode-only-benchmark). | +| `AFD_REPO` | current directory | Repository root that holds `tools/benchmarks` (with `DECODE_BENCH`). | +| `M2N_EXTRA_CONFIG` | empty | Extra `connector_extra_config` members, for example `'"expert_num_sms": 130'`. | +| `HEAD_HOST` | `127.0.0.1` on one pod | Address of pod 0 when it leads both roles (multi-pod, no Kubernetes). | +| `ATTN_HEAD_HOST` / `FFN_HEAD_HOST` | `HEAD_HOST` | Address of the first pod with Attention / FFN ranks, when they differ. The FFN head also hosts the AFD rendezvous. | +| `AFD_CONNECTOR` | `DeepEPM2NAFDConnector` | `P2pNcclAFDConnector` runs the same layout on NCCL P2P (used by the P2P recipe entry points). | + +`VLLM_ARGS`, `ATTN_ARGS`, `FFN_ARGS` and any arguments after the script name +come after the recipe's own flags, so a repeated flag overrides its default +(argparse keeps the last value). They are split on whitespace without shell +quoting, so pass JSON values without spaces. For example: + +```bash +# 512 sequences per Attention rank, no prefix caching on Attention +VLLM_ARGS="--max-num-seqs 512" ATTN_ARGS="--no-enable-prefix-caching" \ + bash recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/graph.sh +# Count graph and eager Attention steps (vLLM disables the stats logger with +# more than one API server) +ATTN_ARGS="--api-server-count 1 --cudagraph-metrics" bash .../graph.sh +``` + +Raise `CUDAGRAPH_CAPTURE_SIZES` together with `--max-num-seqs`, so decode +batches stay in a captured size. + +## Running on one host -Run any script from the repository root, for example: +Run a script from the repository root: ```bash export MODEL_PATH=/path/model_weights/Qwen3-30B-A3B -bash recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/2a2f_graph_dp2tp1.sh +bash recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/graph.sh ``` -Each script starts the FFN server and then the Attention server in the -background and writes `ffn.log` and `attn.log` into the current directory. -Wait for `Application startup complete` in both logs, then send traffic to the -Attention server on port `18000`. The AFD rendezvous uses port `6239`; the FFN -server listens on `18001` but serves no requests. +The script starts the FFN server and then the Attention server in the +background, writes `ffn.log` and `attn.log` into the current directory, and +waits. Wait for `Application startup complete` in `attn.log`, then send +traffic to the Attention server on port `18000`. The first start on a host +compiles transport and DeepGEMM kernels, so it is slower than later starts. + +## Running on several pods + +Start the same script on every pod with the same settings and the pod's own +`POD`. Without Kubernetes, point the pods at the leaders' addresses: -The first start on a node compiles transport and DeepGEMM kernels, so it is -slower than later starts. +```bash +# 2 x 1A1F: pod 0 leads both roles +LAYOUT=1A1F,1A1F POD=AFD_0 HEAD_HOST= bash .../graph.sh # pod 0 +LAYOUT=1A1F,1A1F POD=AFD_1 HEAD_HOST= bash .../graph.sh # pod 1 + +# 2A0F,0A2F: Attention on pod 0, FFN on pod 1 +LAYOUT=2A0F,0A2F POD=AFD_0 ATTN_HEAD_HOST= FFN_HEAD_HOST= bash .../graph.sh +LAYOUT=2A0F,0A2F POD=AFD_1 ATTN_HEAD_HOST= FFN_HEAD_HOST= bash .../graph.sh +``` + +Send traffic to port `18000` of the first pod with Attention ranks; the other +pods serve no requests. Pass the fabric's NCCL settings (for example +`NCCL_IB_HCA`) through the environment. The roles can finish loading minutes +apart, so the recipes raise `afd_process_group_timeout_s` and +`--cpu-distributed-timeout-seconds` to 3600 s. ## Deploying on Kubernetes -The single-node recipes have no placement header, so the -**`deploy-afd-k8s`** skill (`.agents/skills/deploy-afd-k8s/SKILL.md`) deploys -each one as a single pod. The image must contain afdEP-comm and DeepGEMM. - -`multipod_4a1f_4a1f_graph.sh` carries a `POD=AFD_` header, so the skill -deploys it as one pod per entry, each running both roles. Every pod gets the -same `NUM_PODS`, `ATTN_PER_POD` and `FFN_PER_POD` (2, 4 and 1 by default), -which select any symmetric layout `N x AF`; the FFN uses each pod's -first `f` GPUs. See -[Multi-pod deployment](../../../docs/gpu/DEEPEP_M2N_CONNECTOR_USER_GUIDE.md#multi-pod-deployment) -for the placements the connector accepts. +The **`deploy-afd-k8s`** skill (`.agents/skills/deploy-afd-k8s/SKILL.md`) +deploys a recipe as one pod per `LAYOUT` entry, sizes each pod's GPUs from its +entry, and creates the Services the recipe's default host names point at +(`vllm-ffn-p2p-service`, `vllm-attn-dp-service`, `vllm-ffn-dp-service`). The +image must contain afdEP-comm and DeepGEMM. + +## Checking a deployment + +```bash +curl -s http://127.0.0.1:18000/v1/completions \ + -H 'Content-Type: application/json' \ + -d '{"model":"'"$MODEL_PATH"'","prompt":"The capital of France is","max_tokens":12,"temperature":0}' +``` + +The answer should name Paris. For an accuracy check, run the matching E2E case +(`afd-deepep-m2n-*`); see +[`tests/e2e`](../../../tests/e2e/README.md#deepep-m2n-scenarios). + +## Decode-only benchmark + +`DECODE_BENCH` attaches the decode-only benchmark's KV connector +(`AFDDecodeBenchConnector`) to Attention and takes its +`kv_connector_extra_config` as JSON: every Attention DP rank prefills a pool +of `num_prompts` random prompts of `prompt_len` tokens, pins their KV, and +remaps every later request onto the pool, so measured requests decode over +real KV without prefill. Set `max_output_len` to the longest output you +measure; it sizes the startup KV check. The connector lives in +`tools/benchmarks`, which is not in the wheel: run from the repository root or +set `AFD_REPO`. + +```bash +export MODEL_PATH=/path/model_weights/Qwen3-30B-A3B +DECODE_BENCH='{"num_prompts": 500, "prompt_len": 1025, "max_output_len": 256}' \ + bash recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/graph.sh + +python tools/benchmarks/decode_bench_client.py \ + --port 18000 --dp-size 2 --pool-prompts 500 --prompt-len 1025 \ + --result-dir /tmp/results/decode_bench \ + -- --num-prompts 2000 --random-output-len 256 \ + --max-concurrency 64 --request-rate inf +``` + +`--dp-size` is the Attention DP size over all pods (total Attention ranks +divided by `ATTN_TP`; 2 for `2A2F`, `2A0F,0A2F` and `1A1F,1A1F`), and +`--pool-prompts` and `--prompt-len` must match `DECODE_BENCH`. The benchmark +works in every layout and mode. How the pool, the client and its checks work, +and how to size the pool, are described in +[`tools/benchmarks/README.md`](../../../tools/benchmarks/README.md); the model +READMEs give each model's KV size per token. + +## Configuration notes + +- Both roles set `compute_gate_on_attention: true` and run with + `--enable-expert-parallel` and `--enable-ep-weight-filter`, so each FFN rank + reads only its own experts. +- The FFN keeps vLLM's default `--moe-backend auto`, which the plugin pins to + `triton` so the expert weights stay in the layout the connector reads. +- The connector derives its dispatch buckets from the CUDA graph capture + sizes and compiles one transport kernel set per bucket at startup. The + dispatch capacity is the larger of `--max-num-batched-tokens` and the + largest capture size, so lowering `--max-num-batched-tokens` shrinks the + transport buffers. +- With DBO the connector uses two lanes; DBO must be on for both roles. +- `expert_num_sms` and `num_sms` split SMs between the DeepGEMM experts and + the transport kernels; pass them with `M2N_EXTRA_CONFIG`. See + [`connector_extra_config`](../../../docs/gpu/DEEPEP_M2N_CONNECTOR_USER_GUIDE.md#connector_extra_config). +- `FULL_DECODE_ONLY` is the only CUDA graph mode. diff --git a/recipe/gpu/DeepEPM2NAFDConnector/deepseek_v2_lite/README.md b/recipe/gpu/DeepEPM2NAFDConnector/deepseek_v2_lite/README.md index f2ff98ed8..4e014b0d9 100644 --- a/recipe/gpu/DeepEPM2NAFDConnector/deepseek_v2_lite/README.md +++ b/recipe/gpu/DeepEPM2NAFDConnector/deepseek_v2_lite/README.md @@ -1,46 +1,41 @@ # DeepSeek-V2-Lite with DeepEPM2NAFDConnector -Launch scripts for DeepSeek-V2-Lite (BF16) on `DeepEPM2NAFDConnector` with -vLLM `v0.32.0`. They use the same layouts as the DeepSeek-V2-Lite DeepEP M2N -E2E cases (`test_deepseek_v2_lite_deepep_m2n` in -[`tests/e2e`](../../../../tests/e2e/README.md#deepep-m2n-scenarios)). - -## Prerequisites - -- 4 SM90 or newer GPUs on one node. -- The environment described in the - [connector recipes README](../README.md#prerequisites). -- DeepSeek-V2-Lite weights on disk. The scripts default to - `/path/model_weights/DeepSeek-V2-Lite`; override with `MODEL_PATH=...`. -- Free ports `6239` (AFD rendezvous), `18000` (Attention API) and `18001` - (FFN). - -## Layout - -`2A2F`: Attention on GPUs 0,1 (DP2, or DP1 x TP2) and FFN DP2 on GPUs 2,3. -FFN rank 0 owns routed experts 0-31 and FFN rank 1 owns routed experts 32-63. +DeepSeek-V2-Lite (BF16) recipes. The scripts, layouts, settings, multi-pod +runs and the decode-only benchmark are described once in the +[connector recipes README](../README.md); this page lists what is specific to +the model. ```text prefill_decode_colocation/ -├── 2a2f_eager_dp2tp1.sh # Attention DP2 x TP1, eager -├── 2a2f_graph_dp2tp1.sh # Attention DP2 x TP1, FULL_DECODE_ONLY graph -├── 2a2f_graph_dbo_dp2tp1.sh # Attention DP2 x TP1, graph + DBO -└── 2a2f_graph_dp1tp2.sh # Attention DP1 x TP2, FULL_DECODE_ONLY graph +├── graph.sh # FULL_DECODE_ONLY CUDA graphs +├── eager.sh # eager +└── dbo.sh # graph + DBO ``` -## Configuration notes +## Defaults + +| Setting | Value | +| --- | --- | +| `LAYOUT` | `2A2F`: FFN DP2 on GPUs 0,1, Attention DP2 on GPUs 2,3 | +| `--max-num-seqs` | 64 per Attention DP rank | +| `--max-num-batched-tokens` | 4096 | +| `--max-model-len` | 4096 | +| `CUDAGRAPH_CAPTURE_SIZES` | 8,16,32,64 | +| Payload | BF16 (no `connector_extra_config`); the FFN runs DeepGEMM experts | + +These match the DeepSeek-V2-Lite DeepEP M2N E2E cases +(`test_deepseek_v2_lite_deepep_m2n` in +[`tests/e2e`](../../../../tests/e2e/README.md#deepep-m2n-scenarios)); +`ATTN_TP=2` matches the Attention TP2 case. + +## Notes -- Both roles set `compute_gate_on_attention: true`, use the same rank counts, - and run with `--enable-expert-parallel`. The FFN runs TP 1. - Attention runs the gate, the router, the dense first layer and the two - shared experts; the FFN runs only the routed experts. -- With Attention TP 2 each TP rank routes and dispatches its slice of the + shared experts; the FFN runs only the routed experts. With two FFN ranks, + FFN rank 0 owns routed experts 0-31 and FFN rank 1 owns 32-63. +- With Attention TP each TP rank routes and dispatches its slice of the tokens and runs the shared experts on that slice with replicated weights; the TP group all-gathers the result. -- `connector_extra_config` is left empty, so the payload is BF16 and the FFN - runs DeepGEMM experts. -- The graph and DBO settings match the - [Qwen3-30B-A3B recipes](../qwen3_30b_a3b/README.md#configuration-notes). ## Limitations diff --git a/recipe/gpu/DeepEPM2NAFDConnector/deepseek_v2_lite/prefill_decode_colocation/2a2f_eager_dp2tp1.sh b/recipe/gpu/DeepEPM2NAFDConnector/deepseek_v2_lite/prefill_decode_colocation/2a2f_eager_dp2tp1.sh deleted file mode 100644 index 3da8f033f..000000000 --- a/recipe/gpu/DeepEPM2NAFDConnector/deepseek_v2_lite/prefill_decode_colocation/2a2f_eager_dp2tp1.sh +++ /dev/null @@ -1,55 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project - -# 2A2F eager: Attention DP2 x TP1 on GPUs 0,1; FFN DP2 on GPUs 2,3. - -MODEL_PATH=${MODEL_PATH:-/path/model_weights/DeepSeek-V2-Lite} -export VLLM_USE_V2_MODEL_RUNNER=1 - -CUDA_VISIBLE_DEVICES=2,3 vllm serve "$MODEL_PATH" \ - --data-parallel-size 2 \ - --tensor-parallel-size 1 \ - --enable-expert-parallel \ - --additional-config '{ - "afd": { - "role": "ffn", - "connector": "DeepEPM2NAFDConnector", - "host": "127.0.0.1", - "port": 6239, - "num_attention_ranks": 2, - "num_ffn_ranks": 2, - "compute_gate_on_attention": true - } - }' \ - --max-num-seqs 64 \ - --max-num-batched-tokens 4096 \ - --max-model-len 4096 \ - --enforce-eager \ - --host 127.0.0.1 \ - --port 18001 \ - --trust-remote-code > ffn.log 2>&1 & - -CUDA_VISIBLE_DEVICES=0,1 vllm serve "$MODEL_PATH" \ - --data-parallel-size 2 \ - --tensor-parallel-size 1 \ - --enable-expert-parallel \ - --additional-config '{ - "afd": { - "role": "attention", - "connector": "DeepEPM2NAFDConnector", - "host": "127.0.0.1", - "port": 6239, - "num_attention_ranks": 2, - "num_ffn_ranks": 2, - "compute_gate_on_attention": true - } - }' \ - --max-num-seqs 64 \ - --max-num-batched-tokens 4096 \ - --max-model-len 4096 \ - --enforce-eager \ - --host 127.0.0.1 \ - --port 18000 \ - --trust-remote-code > attn.log 2>&1 & - -wait diff --git a/recipe/gpu/DeepEPM2NAFDConnector/deepseek_v2_lite/prefill_decode_colocation/2a2f_graph_dbo_dp2tp1.sh b/recipe/gpu/DeepEPM2NAFDConnector/deepseek_v2_lite/prefill_decode_colocation/2a2f_graph_dbo_dp2tp1.sh deleted file mode 100644 index 0ea5a336e..000000000 --- a/recipe/gpu/DeepEPM2NAFDConnector/deepseek_v2_lite/prefill_decode_colocation/2a2f_graph_dbo_dp2tp1.sh +++ /dev/null @@ -1,64 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project - -# 2A2F graph + DBO: Attention DP2 x TP1 on GPUs 0,1; FFN DP2 on GPUs 2,3. -# DBO must be enabled on both roles. - -MODEL_PATH=${MODEL_PATH:-/path/model_weights/DeepSeek-V2-Lite} -export VLLM_USE_V2_MODEL_RUNNER=1 - -CUDA_VISIBLE_DEVICES=2,3 vllm serve "$MODEL_PATH" \ - --data-parallel-size 2 \ - --tensor-parallel-size 1 \ - --enable-expert-parallel \ - --additional-config '{ - "afd": { - "role": "ffn", - "connector": "DeepEPM2NAFDConnector", - "host": "127.0.0.1", - "port": 6239, - "num_attention_ranks": 2, - "num_ffn_ranks": 2, - "compute_gate_on_attention": true - } - }' \ - --max-num-seqs 64 \ - --max-num-batched-tokens 4096 \ - --max-model-len 4096 \ - --max-cudagraph-capture-size 64 \ - --compilation-config '{"cudagraph_mode": "FULL_DECODE_ONLY", "cudagraph_capture_sizes": [8, 16, 32, 64]}' \ - --enable-dbo \ - --dbo-decode-token-threshold 2 \ - --dbo-prefill-token-threshold 8 \ - --host 127.0.0.1 \ - --port 18001 \ - --trust-remote-code > ffn.log 2>&1 & - -CUDA_VISIBLE_DEVICES=0,1 vllm serve "$MODEL_PATH" \ - --data-parallel-size 2 \ - --tensor-parallel-size 1 \ - --enable-expert-parallel \ - --additional-config '{ - "afd": { - "role": "attention", - "connector": "DeepEPM2NAFDConnector", - "host": "127.0.0.1", - "port": 6239, - "num_attention_ranks": 2, - "num_ffn_ranks": 2, - "compute_gate_on_attention": true - } - }' \ - --max-num-seqs 64 \ - --max-num-batched-tokens 4096 \ - --max-model-len 4096 \ - --max-cudagraph-capture-size 64 \ - --compilation-config '{"cudagraph_mode": "FULL_DECODE_ONLY", "cudagraph_capture_sizes": [8, 16, 32, 64]}' \ - --enable-dbo \ - --dbo-decode-token-threshold 2 \ - --dbo-prefill-token-threshold 8 \ - --host 127.0.0.1 \ - --port 18000 \ - --trust-remote-code > attn.log 2>&1 & - -wait diff --git a/recipe/gpu/DeepEPM2NAFDConnector/deepseek_v2_lite/prefill_decode_colocation/2a2f_graph_dp1tp2.sh b/recipe/gpu/DeepEPM2NAFDConnector/deepseek_v2_lite/prefill_decode_colocation/2a2f_graph_dp1tp2.sh deleted file mode 100644 index e08e1d695..000000000 --- a/recipe/gpu/DeepEPM2NAFDConnector/deepseek_v2_lite/prefill_decode_colocation/2a2f_graph_dp1tp2.sh +++ /dev/null @@ -1,58 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project - -# 2A2F graph with Attention TP: Attention DP1 x TP2 on GPUs 0,1; FFN DP2 x TP1 on -# GPUs 2,3. Each Attention TP rank dispatches its slice of the tokens. - -MODEL_PATH=${MODEL_PATH:-/path/model_weights/DeepSeek-V2-Lite} -export VLLM_USE_V2_MODEL_RUNNER=1 - -CUDA_VISIBLE_DEVICES=2,3 vllm serve "$MODEL_PATH" \ - --data-parallel-size 2 \ - --tensor-parallel-size 1 \ - --enable-expert-parallel \ - --additional-config '{ - "afd": { - "role": "ffn", - "connector": "DeepEPM2NAFDConnector", - "host": "127.0.0.1", - "port": 6239, - "num_attention_ranks": 2, - "num_ffn_ranks": 2, - "compute_gate_on_attention": true - } - }' \ - --max-num-seqs 64 \ - --max-num-batched-tokens 4096 \ - --max-model-len 4096 \ - --max-cudagraph-capture-size 64 \ - --compilation-config '{"cudagraph_mode": "FULL_DECODE_ONLY", "cudagraph_capture_sizes": [8, 16, 32, 64]}' \ - --host 127.0.0.1 \ - --port 18001 \ - --trust-remote-code > ffn.log 2>&1 & - -CUDA_VISIBLE_DEVICES=0,1 vllm serve "$MODEL_PATH" \ - --data-parallel-size 1 \ - --tensor-parallel-size 2 \ - --enable-expert-parallel \ - --additional-config '{ - "afd": { - "role": "attention", - "connector": "DeepEPM2NAFDConnector", - "host": "127.0.0.1", - "port": 6239, - "num_attention_ranks": 2, - "num_ffn_ranks": 2, - "compute_gate_on_attention": true - } - }' \ - --max-num-seqs 64 \ - --max-num-batched-tokens 4096 \ - --max-model-len 4096 \ - --max-cudagraph-capture-size 64 \ - --compilation-config '{"cudagraph_mode": "FULL_DECODE_ONLY", "cudagraph_capture_sizes": [8, 16, 32, 64]}' \ - --host 127.0.0.1 \ - --port 18000 \ - --trust-remote-code > attn.log 2>&1 & - -wait diff --git a/recipe/gpu/DeepEPM2NAFDConnector/deepseek_v2_lite/prefill_decode_colocation/2a2f_graph_dp2tp1.sh b/recipe/gpu/DeepEPM2NAFDConnector/deepseek_v2_lite/prefill_decode_colocation/2a2f_graph_dp2tp1.sh deleted file mode 100644 index c7f1801ad..000000000 --- a/recipe/gpu/DeepEPM2NAFDConnector/deepseek_v2_lite/prefill_decode_colocation/2a2f_graph_dp2tp1.sh +++ /dev/null @@ -1,57 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project - -# 2A2F FULL_DECODE_ONLY graph: Attention DP2 x TP1 on GPUs 0,1; FFN DP2 on GPUs 2,3. - -MODEL_PATH=${MODEL_PATH:-/path/model_weights/DeepSeek-V2-Lite} -export VLLM_USE_V2_MODEL_RUNNER=1 - -CUDA_VISIBLE_DEVICES=2,3 vllm serve "$MODEL_PATH" \ - --data-parallel-size 2 \ - --tensor-parallel-size 1 \ - --enable-expert-parallel \ - --additional-config '{ - "afd": { - "role": "ffn", - "connector": "DeepEPM2NAFDConnector", - "host": "127.0.0.1", - "port": 6239, - "num_attention_ranks": 2, - "num_ffn_ranks": 2, - "compute_gate_on_attention": true - } - }' \ - --max-num-seqs 64 \ - --max-num-batched-tokens 4096 \ - --max-model-len 4096 \ - --max-cudagraph-capture-size 64 \ - --compilation-config '{"cudagraph_mode": "FULL_DECODE_ONLY", "cudagraph_capture_sizes": [8, 16, 32, 64]}' \ - --host 127.0.0.1 \ - --port 18001 \ - --trust-remote-code > ffn.log 2>&1 & - -CUDA_VISIBLE_DEVICES=0,1 vllm serve "$MODEL_PATH" \ - --data-parallel-size 2 \ - --tensor-parallel-size 1 \ - --enable-expert-parallel \ - --additional-config '{ - "afd": { - "role": "attention", - "connector": "DeepEPM2NAFDConnector", - "host": "127.0.0.1", - "port": 6239, - "num_attention_ranks": 2, - "num_ffn_ranks": 2, - "compute_gate_on_attention": true - } - }' \ - --max-num-seqs 64 \ - --max-num-batched-tokens 4096 \ - --max-model-len 4096 \ - --max-cudagraph-capture-size 64 \ - --compilation-config '{"cudagraph_mode": "FULL_DECODE_ONLY", "cudagraph_capture_sizes": [8, 16, 32, 64]}' \ - --host 127.0.0.1 \ - --port 18000 \ - --trust-remote-code > attn.log 2>&1 & - -wait diff --git a/recipe/gpu/DeepEPM2NAFDConnector/deepseek_v2_lite/prefill_decode_colocation/dbo.sh b/recipe/gpu/DeepEPM2NAFDConnector/deepseek_v2_lite/prefill_decode_colocation/dbo.sh new file mode 100644 index 000000000..c7e17b515 --- /dev/null +++ b/recipe/gpu/DeepEPM2NAFDConnector/deepseek_v2_lite/prefill_decode_colocation/dbo.sh @@ -0,0 +1,98 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +# +# DeepSeek-V2-Lite (BF16) on DeepEPM2NAFDConnector, FULL_DECODE_ONLY graphs + DBO. +# Placement: one pod per comma-separated LAYOUT entry (default 2A2F, one pod of +# 4 GPUs). AF puts f FFN ranks on the pod's first f GPUs and a Attention +# ranks on the next a GPUs, e.g. LAYOUT=2A0F,0A2F or LAYOUT=1A1F,1A1F. +# For pod k (k = 0, 1, ...): POD=AFD_, with the same LAYOUT on every pod. +# Settings: recipe/gpu/DeepEPM2NAFDConnector/README.md#recipe-settings + +MODEL_PATH=${MODEL_PATH:-/path/model_weights/DeepSeek-V2-Lite} +LAYOUT=${LAYOUT:-2A2F} +TUNING=(--max-num-seqs 64 --max-num-batched-tokens 4096 --max-model-len 4096) +CONNECTOR_CONFIG=${M2N_EXTRA_CONFIG:-} # connector_extra_config members; BF16 payload +CUDAGRAPH_CAPTURE_SIZES=${CUDAGRAPH_CAPTURE_SIZES:-8,16,32,64} # ascending +MODE=(--max-cudagraph-capture-size "${CUDAGRAPH_CAPTURE_SIZES##*,}" + --compilation-config "{\"cudagraph_mode\": \"FULL_DECODE_ONLY\", \"cudagraph_capture_sizes\": [${CUDAGRAPH_CAPTURE_SIZES}]}" + --enable-dbo --dbo-decode-token-threshold 2 --dbo-prefill-token-threshold 8) + +# ---- The rest is the same in every DeepEPM2NAFDConnector recipe. ---- +export VLLM_USE_V2_MODEL_RUNNER=1 +ATTN_TP=${ATTN_TP:-1} +AFD_CONNECTOR=${AFD_CONNECTOR:-DeepEPM2NAFDConnector} GATE_ON_ATTENTION=true +# P2P routes on the FFN and takes no connector_extra_config. +[[ $AFD_CONNECTOR == P2pNcclAFDConnector ]] && GATE_ON_ATTENTION=false CONNECTOR_CONFIG= +[[ -n $CONNECTOR_CONFIG ]] && CONNECTOR_CONFIG=", \"connector_extra_config\": {${CONNECTOR_CONFIG}}" +IFS=, read -ra PODS <<<"$LAYOUT" +K=${POD:-AFD_0} && K=${K#AFD_} +if [[ ! $LAYOUT =~ ^[0-9]+A[0-9]+F(,[0-9]+A[0-9]+F)*$ || ! $K =~ ^[0-9]+$ ]] || ((K >= ${#PODS[@]})); then + echo "need LAYOUT=AF[,...] and POD=AFD_ in it: LAYOUT=$LAYOUT POD=${POD:-}" >&2 + exit 1 +fi +NUM_A=() NUM_F=() +for p in "${PODS[@]}"; do NUM_A+=("${p%%A*}") && p=${p#*A} && NUM_F+=("${p%F}"); done +((${#PODS[@]} > 1)) || HEAD_HOST=${HEAD_HOST:-127.0.0.1} + +# Extra `vllm serve` flags go last, so they override the defaults. +read -ra ARGS <<<"${VLLM_ARGS:-}" && ARGS+=("$@") +read -ra ATTN_EXTRA <<<"${ATTN_ARGS:-}" +read -ra FFN_EXTRA <<<"${FFN_ARGS:-}" +if [[ -n ${DECODE_BENCH:-} ]]; then # tools/benchmarks is imported from the repository + ATTN_EXTRA=(--kv-transfer-config "{\"kv_connector\": \"AFDDecodeBenchConnector\", \"kv_connector_module_path\": \"tools.benchmarks.decode_bench\", \"kv_role\": \"kv_both\", \"kv_connector_extra_config\": ${DECODE_BENCH}}" "${ATTN_EXTRA[@]}") + export PYTHONPATH=${AFD_REPO:-$PWD}${PYTHONPATH:+:$PYTHONPATH} +fi + +# role_dp TP DP_HOST DP_RPC_PORT ... -> TOTAL and this pod's DP +# flags. A role spread over pods is led by its first pod; the others run +# headless from their start rank (the leader must omit it, or vLLM selects +# hybrid load balancing and rejects the headless peers). +role_dp() { + local tp=$1 host=$2 port=$3 i pods=0 start=0 lead=-1 + shift 3 + local n=("$@") + TOTAL=0 + for i in "${!n[@]}"; do + ((n[i])) || continue + ((n[i] % tp)) && { echo "TP $tp does not divide ${n[i]} ranks on pod $i" >&2; exit 1; } + ((lead < 0)) && lead=$i + ((i < K)) && start=$((start + n[i])) + pods=$((pods + 1)) TOTAL=$((TOTAL + n[i])) + done + DP=(--tensor-parallel-size "$tp" --data-parallel-size $((TOTAL / tp))) + ((pods > 1)) && DP+=(--data-parallel-size-local $((n[K] / tp)) + --data-parallel-address "$host" --data-parallel-rpc-port "$port") + ((pods > 1 && K != lead)) && DP+=(--headless --data-parallel-start-rank $((start / tp))) +} +role_dp 1 "${FFN_HEAD_HOST:-${HEAD_HOST:-vllm-ffn-dp-service}}" 13355 "${NUM_F[@]}" +FFN_DP=("${DP[@]}") NUM_FFN=$TOTAL +role_dp "$ATTN_TP" "${ATTN_HEAD_HOST:-${HEAD_HOST:-vllm-attn-dp-service}}" 13345 "${NUM_A[@]}" +ATTN_DP=("${DP[@]}") NUM_ATTN=$TOTAL + +afd() { # The AFD rendezvous runs on the pod that holds FFN rank 0. + cat < ffn.log 2>&1 & +fi +if ((A)); then + CUDA_VISIBLE_DEVICES=$(seq "$F" $((F + A - 1)) | paste -sd, -) vllm serve "$MODEL_PATH" \ + --additional-config "$(afd attention)" "${COMMON[@]}" "${ATTN_DP[@]}" \ + "${ARGS[@]}" "${ATTN_EXTRA[@]}" --port 18000 > attn.log 2>&1 & +fi +wait diff --git a/recipe/gpu/DeepEPM2NAFDConnector/deepseek_v2_lite/prefill_decode_colocation/eager.sh b/recipe/gpu/DeepEPM2NAFDConnector/deepseek_v2_lite/prefill_decode_colocation/eager.sh new file mode 100644 index 000000000..82fc56903 --- /dev/null +++ b/recipe/gpu/DeepEPM2NAFDConnector/deepseek_v2_lite/prefill_decode_colocation/eager.sh @@ -0,0 +1,95 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +# +# DeepSeek-V2-Lite (BF16) on DeepEPM2NAFDConnector, eager mode. +# Placement: one pod per comma-separated LAYOUT entry (default 2A2F, one pod of +# 4 GPUs). AF puts f FFN ranks on the pod's first f GPUs and a Attention +# ranks on the next a GPUs, e.g. LAYOUT=2A0F,0A2F or LAYOUT=1A1F,1A1F. +# For pod k (k = 0, 1, ...): POD=AFD_, with the same LAYOUT on every pod. +# Settings: recipe/gpu/DeepEPM2NAFDConnector/README.md#recipe-settings + +MODEL_PATH=${MODEL_PATH:-/path/model_weights/DeepSeek-V2-Lite} +LAYOUT=${LAYOUT:-2A2F} +TUNING=(--max-num-seqs 64 --max-num-batched-tokens 4096 --max-model-len 4096) +CONNECTOR_CONFIG=${M2N_EXTRA_CONFIG:-} # connector_extra_config members; BF16 payload +MODE=(--enforce-eager) + +# ---- The rest is the same in every DeepEPM2NAFDConnector recipe. ---- +export VLLM_USE_V2_MODEL_RUNNER=1 +ATTN_TP=${ATTN_TP:-1} +AFD_CONNECTOR=${AFD_CONNECTOR:-DeepEPM2NAFDConnector} GATE_ON_ATTENTION=true +# P2P routes on the FFN and takes no connector_extra_config. +[[ $AFD_CONNECTOR == P2pNcclAFDConnector ]] && GATE_ON_ATTENTION=false CONNECTOR_CONFIG= +[[ -n $CONNECTOR_CONFIG ]] && CONNECTOR_CONFIG=", \"connector_extra_config\": {${CONNECTOR_CONFIG}}" +IFS=, read -ra PODS <<<"$LAYOUT" +K=${POD:-AFD_0} && K=${K#AFD_} +if [[ ! $LAYOUT =~ ^[0-9]+A[0-9]+F(,[0-9]+A[0-9]+F)*$ || ! $K =~ ^[0-9]+$ ]] || ((K >= ${#PODS[@]})); then + echo "need LAYOUT=AF[,...] and POD=AFD_ in it: LAYOUT=$LAYOUT POD=${POD:-}" >&2 + exit 1 +fi +NUM_A=() NUM_F=() +for p in "${PODS[@]}"; do NUM_A+=("${p%%A*}") && p=${p#*A} && NUM_F+=("${p%F}"); done +((${#PODS[@]} > 1)) || HEAD_HOST=${HEAD_HOST:-127.0.0.1} + +# Extra `vllm serve` flags go last, so they override the defaults. +read -ra ARGS <<<"${VLLM_ARGS:-}" && ARGS+=("$@") +read -ra ATTN_EXTRA <<<"${ATTN_ARGS:-}" +read -ra FFN_EXTRA <<<"${FFN_ARGS:-}" +if [[ -n ${DECODE_BENCH:-} ]]; then # tools/benchmarks is imported from the repository + ATTN_EXTRA=(--kv-transfer-config "{\"kv_connector\": \"AFDDecodeBenchConnector\", \"kv_connector_module_path\": \"tools.benchmarks.decode_bench\", \"kv_role\": \"kv_both\", \"kv_connector_extra_config\": ${DECODE_BENCH}}" "${ATTN_EXTRA[@]}") + export PYTHONPATH=${AFD_REPO:-$PWD}${PYTHONPATH:+:$PYTHONPATH} +fi + +# role_dp TP DP_HOST DP_RPC_PORT ... -> TOTAL and this pod's DP +# flags. A role spread over pods is led by its first pod; the others run +# headless from their start rank (the leader must omit it, or vLLM selects +# hybrid load balancing and rejects the headless peers). +role_dp() { + local tp=$1 host=$2 port=$3 i pods=0 start=0 lead=-1 + shift 3 + local n=("$@") + TOTAL=0 + for i in "${!n[@]}"; do + ((n[i])) || continue + ((n[i] % tp)) && { echo "TP $tp does not divide ${n[i]} ranks on pod $i" >&2; exit 1; } + ((lead < 0)) && lead=$i + ((i < K)) && start=$((start + n[i])) + pods=$((pods + 1)) TOTAL=$((TOTAL + n[i])) + done + DP=(--tensor-parallel-size "$tp" --data-parallel-size $((TOTAL / tp))) + ((pods > 1)) && DP+=(--data-parallel-size-local $((n[K] / tp)) + --data-parallel-address "$host" --data-parallel-rpc-port "$port") + ((pods > 1 && K != lead)) && DP+=(--headless --data-parallel-start-rank $((start / tp))) +} +role_dp 1 "${FFN_HEAD_HOST:-${HEAD_HOST:-vllm-ffn-dp-service}}" 13355 "${NUM_F[@]}" +FFN_DP=("${DP[@]}") NUM_FFN=$TOTAL +role_dp "$ATTN_TP" "${ATTN_HEAD_HOST:-${HEAD_HOST:-vllm-attn-dp-service}}" 13345 "${NUM_A[@]}" +ATTN_DP=("${DP[@]}") NUM_ATTN=$TOTAL + +afd() { # The AFD rendezvous runs on the pod that holds FFN rank 0. + cat < ffn.log 2>&1 & +fi +if ((A)); then + CUDA_VISIBLE_DEVICES=$(seq "$F" $((F + A - 1)) | paste -sd, -) vllm serve "$MODEL_PATH" \ + --additional-config "$(afd attention)" "${COMMON[@]}" "${ATTN_DP[@]}" \ + "${ARGS[@]}" "${ATTN_EXTRA[@]}" --port 18000 > attn.log 2>&1 & +fi +wait diff --git a/recipe/gpu/DeepEPM2NAFDConnector/deepseek_v2_lite/prefill_decode_colocation/graph.sh b/recipe/gpu/DeepEPM2NAFDConnector/deepseek_v2_lite/prefill_decode_colocation/graph.sh new file mode 100644 index 000000000..0121b1d0e --- /dev/null +++ b/recipe/gpu/DeepEPM2NAFDConnector/deepseek_v2_lite/prefill_decode_colocation/graph.sh @@ -0,0 +1,97 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +# +# DeepSeek-V2-Lite (BF16) on DeepEPM2NAFDConnector, FULL_DECODE_ONLY CUDA graphs. +# Placement: one pod per comma-separated LAYOUT entry (default 2A2F, one pod of +# 4 GPUs). AF puts f FFN ranks on the pod's first f GPUs and a Attention +# ranks on the next a GPUs, e.g. LAYOUT=2A0F,0A2F or LAYOUT=1A1F,1A1F. +# For pod k (k = 0, 1, ...): POD=AFD_, with the same LAYOUT on every pod. +# Settings: recipe/gpu/DeepEPM2NAFDConnector/README.md#recipe-settings + +MODEL_PATH=${MODEL_PATH:-/path/model_weights/DeepSeek-V2-Lite} +LAYOUT=${LAYOUT:-2A2F} +TUNING=(--max-num-seqs 64 --max-num-batched-tokens 4096 --max-model-len 4096) +CONNECTOR_CONFIG=${M2N_EXTRA_CONFIG:-} # connector_extra_config members; BF16 payload +CUDAGRAPH_CAPTURE_SIZES=${CUDAGRAPH_CAPTURE_SIZES:-8,16,32,64} # ascending +MODE=(--max-cudagraph-capture-size "${CUDAGRAPH_CAPTURE_SIZES##*,}" + --compilation-config "{\"cudagraph_mode\": \"FULL_DECODE_ONLY\", \"cudagraph_capture_sizes\": [${CUDAGRAPH_CAPTURE_SIZES}]}") + +# ---- The rest is the same in every DeepEPM2NAFDConnector recipe. ---- +export VLLM_USE_V2_MODEL_RUNNER=1 +ATTN_TP=${ATTN_TP:-1} +AFD_CONNECTOR=${AFD_CONNECTOR:-DeepEPM2NAFDConnector} GATE_ON_ATTENTION=true +# P2P routes on the FFN and takes no connector_extra_config. +[[ $AFD_CONNECTOR == P2pNcclAFDConnector ]] && GATE_ON_ATTENTION=false CONNECTOR_CONFIG= +[[ -n $CONNECTOR_CONFIG ]] && CONNECTOR_CONFIG=", \"connector_extra_config\": {${CONNECTOR_CONFIG}}" +IFS=, read -ra PODS <<<"$LAYOUT" +K=${POD:-AFD_0} && K=${K#AFD_} +if [[ ! $LAYOUT =~ ^[0-9]+A[0-9]+F(,[0-9]+A[0-9]+F)*$ || ! $K =~ ^[0-9]+$ ]] || ((K >= ${#PODS[@]})); then + echo "need LAYOUT=AF[,...] and POD=AFD_ in it: LAYOUT=$LAYOUT POD=${POD:-}" >&2 + exit 1 +fi +NUM_A=() NUM_F=() +for p in "${PODS[@]}"; do NUM_A+=("${p%%A*}") && p=${p#*A} && NUM_F+=("${p%F}"); done +((${#PODS[@]} > 1)) || HEAD_HOST=${HEAD_HOST:-127.0.0.1} + +# Extra `vllm serve` flags go last, so they override the defaults. +read -ra ARGS <<<"${VLLM_ARGS:-}" && ARGS+=("$@") +read -ra ATTN_EXTRA <<<"${ATTN_ARGS:-}" +read -ra FFN_EXTRA <<<"${FFN_ARGS:-}" +if [[ -n ${DECODE_BENCH:-} ]]; then # tools/benchmarks is imported from the repository + ATTN_EXTRA=(--kv-transfer-config "{\"kv_connector\": \"AFDDecodeBenchConnector\", \"kv_connector_module_path\": \"tools.benchmarks.decode_bench\", \"kv_role\": \"kv_both\", \"kv_connector_extra_config\": ${DECODE_BENCH}}" "${ATTN_EXTRA[@]}") + export PYTHONPATH=${AFD_REPO:-$PWD}${PYTHONPATH:+:$PYTHONPATH} +fi + +# role_dp TP DP_HOST DP_RPC_PORT ... -> TOTAL and this pod's DP +# flags. A role spread over pods is led by its first pod; the others run +# headless from their start rank (the leader must omit it, or vLLM selects +# hybrid load balancing and rejects the headless peers). +role_dp() { + local tp=$1 host=$2 port=$3 i pods=0 start=0 lead=-1 + shift 3 + local n=("$@") + TOTAL=0 + for i in "${!n[@]}"; do + ((n[i])) || continue + ((n[i] % tp)) && { echo "TP $tp does not divide ${n[i]} ranks on pod $i" >&2; exit 1; } + ((lead < 0)) && lead=$i + ((i < K)) && start=$((start + n[i])) + pods=$((pods + 1)) TOTAL=$((TOTAL + n[i])) + done + DP=(--tensor-parallel-size "$tp" --data-parallel-size $((TOTAL / tp))) + ((pods > 1)) && DP+=(--data-parallel-size-local $((n[K] / tp)) + --data-parallel-address "$host" --data-parallel-rpc-port "$port") + ((pods > 1 && K != lead)) && DP+=(--headless --data-parallel-start-rank $((start / tp))) +} +role_dp 1 "${FFN_HEAD_HOST:-${HEAD_HOST:-vllm-ffn-dp-service}}" 13355 "${NUM_F[@]}" +FFN_DP=("${DP[@]}") NUM_FFN=$TOTAL +role_dp "$ATTN_TP" "${ATTN_HEAD_HOST:-${HEAD_HOST:-vllm-attn-dp-service}}" 13345 "${NUM_A[@]}" +ATTN_DP=("${DP[@]}") NUM_ATTN=$TOTAL + +afd() { # The AFD rendezvous runs on the pod that holds FFN rank 0. + cat < ffn.log 2>&1 & +fi +if ((A)); then + CUDA_VISIBLE_DEVICES=$(seq "$F" $((F + A - 1)) | paste -sd, -) vllm serve "$MODEL_PATH" \ + --additional-config "$(afd attention)" "${COMMON[@]}" "${ATTN_DP[@]}" \ + "${ARGS[@]}" "${ATTN_EXTRA[@]}" --port 18000 > attn.log 2>&1 & +fi +wait diff --git a/recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/README.md b/recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/README.md index 8bef91e70..38255d164 100644 --- a/recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/README.md +++ b/recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/README.md @@ -1,170 +1,56 @@ -# Qwen3-235B-A22B-FP8 AFD Recipe +# Qwen3-235B-A22B-FP8 with DeepEPM2NAFDConnector -Launch scripts for Qwen3-235B-A22B-FP8 in a 4A2F prefill/decode colocation on -a single node with 6 GPUs, with vLLM `v0.32.0`. `4a2f_graph_dp4tp1.sh` also -takes a selectable AFD connector and an optional KV connector; with the -decode-bench KV connector it doubles as a decode-only benchmark over real KV, -exactly like the -[Qwen3-30B-A3B recipe](../qwen3_30b_a3b/README.md#decode-only-benchmark). - -## Prerequisites - -- 6 GPUs with about 140 GB each on one node (H200), or two pods of 5 such - GPUs on an RDMA fabric for the multi-pod recipe. `DeepEPM2NAFDConnector` - needs SM90 or newer and the environment described in the - [connector recipes README](../README.md#prerequisites); - `P2pNcclAFDConnector` runs on any A/H-class GPU with enough memory. -- Qwen3-235B-A22B-FP8 weights on disk (block-128 FP8). The scripts default to - `/path/model_weights/Qwen3-235B-A22B-FP8`; override with `MODEL_PATH=...`. -- Free ports `6239` (AFD rendezvous), `18000` (Attention API) and `18001` - (FFN). - -## Layout - -`4A2F`: Attention DP4 x TP1 on GPUs 0-3 and FFN DP2 on GPUs 4,5. Each FFN rank -owns half of the 128 experts. - -| GPUs | Role | Parallelism | Port | -|------------|-----------|----------------------|-------| -| 0, 1, 2, 3 | Attention | DP4 / TP1 | 18000 | -| 4, 5 | FFN | DP2 / TP1, EP2 | 18001 | +Qwen3-235B-A22B-FP8 recipes with the FP8 dispatch payload. The scripts, +layouts, settings, multi-pod runs and the decode-only benchmark are described +once in the [connector recipes README](../README.md); this page lists what is +specific to the model. ```text prefill_decode_colocation/ -├── 4a2f_graph_dp4tp1.sh # 256 sequences per Attention rank, FULL_DECODE_ONLY graph; every knob an env var -├── 4a2f_graph_dbo_dp4tp1.sh # 512 sequences per Attention rank, graph + DBO -└── multipod_4a1f_4a1f_graph.sh # 2 pods x 4A1F, FULL_DECODE_ONLY graph (DBO optional) +├── graph.sh # FULL_DECODE_ONLY CUDA graphs, 256 sequences per Attention rank +├── eager.sh # eager, 256 sequences per Attention rank +└── dbo.sh # graph + DBO, 512 sequences per Attention rank ``` -`ATTN_TP=4` runs `4a2f_graph_dp4tp1.sh`'s Attention as one DP1/TP4 engine on -the same four GPUs instead (`num_attention_ranks` stays 4); raise -`MAX_NUM_SEQS` and the capture sizes with it, since one engine then holds every -running request. - -Both roles run `FULL_DECODE_ONLY` CUDA graphs (sizes 8 to 256 by default, -including 48, 112 and 192); prefill runs eager in chunks of -`MAX_NUM_BATCHED_TOKENS`. Send traffic to the Attention port. -`DeepEPM2NAFDConnector` routes on Attention and dispatches FP8 tokens -(`M2N_PAYLOAD_FORMAT=fp8_block128_fp32`); `P2pNcclAFDConnector` sends BF16 -hidden states and routes on the FFN. - -### Multi-pod: 2 x 4A1F - -`multipod_4a1f_4a1f_graph.sh` runs on two pods of 5 GPUs. Each pod runs one -FFN rank on GPU 0 and four Attention ranks on GPUs 1-4 (Attention DP8 x TP1, -FFN DP2/EP2 in total), so every FFN GPU holds half of the experts, as in 4A2F. -Pod 0 leads both DP groups and hosts the AFD rendezvous; pod 1 starts both -servers headless and serves no requests. Send traffic to pod 0's port `18000`. - -Only `POD` differs between pods: - -```bash -# Pod 0 -POD=AFD_0 bash recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/multipod_4a1f_4a1f_graph.sh -# Pod 1 -POD=AFD_1 bash recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/multipod_4a1f_4a1f_graph.sh -``` - -On Kubernetes the script reaches pod 0 through the Services that the -`deploy-afd-k8s` skill creates. Without Kubernetes, set `HEAD_HOST` to pod 0's -address on both pods; pod 0 must accept ports `6239`, `13345` (Attention DP) -and `13355` (FFN DP). Pass the fabric's NCCL settings (for example -`NCCL_IB_HCA`) through the environment. - -`NUM_PODS`, `ATTN_PER_POD` and `FFN_PER_POD` select another symmetric layout -(`2 x 6A2F`, ...); the tuning values (`MAX_NUM_SEQS`, `MAX_NUM_BATCHED_TOKENS`, -`MAX_MODEL_LEN`, `CUDAGRAPH_CAPTURE_SIZES`, `PAYLOAD_FORMAT`) and DBO -(`ENABLE_DBO=1`, `DBO_DECODE_TOKEN_THRESHOLD`, `DBO_PREFILL_TOKEN_THRESHOLD`) -are environment variables with the 4A2F values as defaults. A placement the connector cannot -carry fails at startup with the pods' shapes in the error. - -## Configuration notes - -- `connector_extra_config` sets `payload_format: "fp8_block128_fp32"`. The - FFN must load the block-128 FP8 checkpoint; the startup handshake rejects a - mismatch between the payload format and the loaded expert weights. -- The FFN keeps vLLM's default `--moe-backend auto`, which the plugin pins to - `triton` so the expert weights stay in the layout the connector reads. -- `afd_process_group_timeout_s` and `--cpu-distributed-timeout-seconds` are - raised to 3600 s because loading the weights can take longer than the - default 120 s join timeout, especially from cold storage, and the two roles - may finish minutes apart. -- The capture sizes cover `--max-num-seqs`. The connector derives its dispatch - buckets from them and compiles one transport kernel set per bucket at - startup. -- `expert_num_sms` (default `0`, all SMs) and `num_sms` tune how SMs are split - between DeepGEMM experts and the transport kernels; see the connector guide. - Pass them through `M2N_EXTRA_CONFIG`, for example `M2N_EXTRA_CONFIG='"expert_num_sms": 130'`. - -## Running - -Run from the repository root: - -```bash -export MODEL_PATH=/path/model_weights/Qwen3-235B-A22B-FP8 -# DeepEP M2N (default) -bash recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/4a2f_graph_dp4tp1.sh -# NCCL P2P: same script with AFD_CONNECTOR=P2pNcclAFDConnector, or its entry point -bash recipe/gpu/P2pNcclAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/4a2f_graph_dp4tp1.sh -``` - -The script backgrounds both roles, logs to `attn.log` and `ffn.log` in the -current directory, and waits. Wait for `Application startup complete` in -`attn.log` before sending traffic. - -Every knob of `4a2f_graph_dp4tp1.sh` is an environment variable documented at -the top of the script, including DBO (`ENABLE_DBO=1`), `MAX_NUM_SEQS` per -Attention rank, graph capture sizes, devices, ports, `EXTRA_ARGS` for both -roles, and `ATTN_EXTRA_ARGS` / `FFN_EXTRA_ARGS` for one role. For example, -`ATTN_EXTRA_ARGS="--api-server-count 1 --cudagraph-metrics"` logs how many -Attention steps ran in each CUDA graph and how many ran eager (vLLM disables -the stats logger with more than one API server, and DP4 starts four). - -Do not set `PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True` for the -`P2pNcclAFDConnector` FFN: it hangs on its first step. - -## Checking a deployment - -```bash -curl -s http://127.0.0.1:18000/v1/completions \ - -H 'Content-Type: application/json' \ - -d '{"model":"'"$MODEL_PATH"'","prompt":"The capital of France is","max_tokens":12,"temperature":0}' -``` - -The answer should name Paris. An FFN log line `Using TRITON Fp8 MoE backend` -confirms the expected expert-weight layout. - -## Decode-only benchmark - -`KV_CONNECTOR=decode_bench` attaches `AFDDecodeBenchConnector` to Attention; -the pool, pinning, and remapping work as described for -[Qwen3-30B-A3B](../qwen3_30b_a3b/README.md#decode-only-benchmark). Drive it -with `--dp-size 4`: - -```bash -export MODEL_PATH=/path/model_weights/Qwen3-235B-A22B-FP8 -KV_CONNECTOR=decode_bench \ -DECODE_BENCH_NUM_PROMPTS=500 \ -DECODE_BENCH_PROMPT_LEN=1025 \ -DECODE_BENCH_MAX_OUTPUT_LEN=256 \ - bash recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/4a2f_graph_dp4tp1.sh - -python tools/benchmarks/decode_bench_client.py \ - --port 18000 --dp-size 4 --pool-prompts 500 --prompt-len 1025 \ - --result-dir /tmp/results/decode_bench \ - -- --num-prompts 3000 --random-output-len 256 \ - --max-concurrency 128 --request-rate 10 -``` - -### Sizing the pool - -Qwen3-235B holds 188 KiB of BF16 KV per token (94 layers, 4 KV heads of 128). -With the defaults, each Attention rank needs -`500 * ceil(1026 / 16) + 256 * ceil(257 / 16)` blocks, about 590K tokens or -106 GiB; an H200 Attention rank has room for about 640K tokens. Keep -`DECODE_BENCH_NUM_PROMPTS` at or above `MAX_NUM_SEQS`. +## Defaults + +| Setting | `graph.sh`, `eager.sh` | `dbo.sh` | +| --- | --- | --- | +| `LAYOUT` | `2A2F` | `2A2F` | +| `--max-num-seqs` | 256 per Attention DP rank | 512 | +| `--max-num-batched-tokens` | 3072 | 3072 | +| `--max-model-len` | 4096 | 4096 | +| `CUDAGRAPH_CAPTURE_SIZES` | 8,12,16,24,32,48,64,96,128,192,256 | 8,16,32,64,96,128,192,256,384,512 | +| `connector_extra_config` | `"payload_format": "fp8_block128_fp32"` | same | + +Common layouts: + +| `LAYOUT` | GPUs | Notes | +| --- | --- | --- | +| `2A2F` | 4 on one host | each FFN GPU holds half of the 128 experts | +| `4A2F` | 6 on one host | Attention DP4 on GPUs 2-5 | +| `4A1F,4A1F` | 2 pods of 5 | Attention DP8, FFN DP2/EP2 | + +## Notes + +- Every layout with two FFN ranks puts about 110 GB of FP8 experts on each + FFN GPU, so it needs H200-class GPUs (about 140 GB); four or more FFN ranks + fit smaller GPUs. Attention ranks hold no experts + (`--enable-ep-weight-filter`). +- The FFN must load the block-128 FP8 checkpoint; the startup handshake + rejects a mismatch between the payload format and the loaded expert weights. +- `ATTN_TP=4` with `LAYOUT=4A2F` runs Attention as one DP1/TP4 engine; raise + `--max-num-seqs` and the capture sizes with it, since one engine then holds + every running request. +- The baseline without AFD is + [`recipe/gpu/baseline/qwen3_235b_a22b_fp8/vanilla_tp4.sh`](../../baseline/qwen3_235b_a22b_fp8/vanilla_tp4.sh). +- Decode-only benchmark pool sizing: Qwen3-235B holds 188 KiB of BF16 KV per + token (94 layers, 4 KV heads of 128). With + `{"num_prompts": 500, "prompt_len": 1025, "max_output_len": 256}` and 256 + sequences, each Attention rank needs + `500 * ceil(1026 / 16) + 256 * ceil(257 / 16)` blocks, about 590K tokens or + 106 GiB; an H200 Attention rank has room for about 640K tokens. ## Limitations - Qwen3 MoE only; see the connector guide for the full list of requirements. -- `FULL_DECODE_ONLY` is the only CUDA graph mode. diff --git a/recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/4a2f_graph_dbo_dp4tp1.sh b/recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/4a2f_graph_dbo_dp4tp1.sh deleted file mode 100644 index 37aef9314..000000000 --- a/recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/4a2f_graph_dbo_dp4tp1.sh +++ /dev/null @@ -1,68 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project - -# 4A2F graph + DBO, FP8 payload: Attention DP4 x TP1 on GPUs 0-3; FFN DP2 on -# GPUs 4,5. DBO must be enabled on both roles. - -MODEL_PATH=${MODEL_PATH:-/path/model_weights/Qwen3-235B-A22B-FP8} -export VLLM_USE_V2_MODEL_RUNNER=1 - -CUDA_VISIBLE_DEVICES=4,5 vllm serve "$MODEL_PATH" \ - --data-parallel-size 2 \ - --tensor-parallel-size 1 \ - --enable-expert-parallel \ - --additional-config '{ - "afd": { - "role": "ffn", - "connector": "DeepEPM2NAFDConnector", - "host": "127.0.0.1", - "port": 6239, - "num_attention_ranks": 4, - "num_ffn_ranks": 2, - "compute_gate_on_attention": true, - "afd_process_group_timeout_s": 3600, - "connector_extra_config": {"payload_format": "fp8_block128_fp32"} - } - }' \ - --max-num-seqs 512 \ - --max-num-batched-tokens 3072 \ - --max-model-len 4096 \ - --max-cudagraph-capture-size 512 \ - --compilation-config '{"cudagraph_mode": "FULL_DECODE_ONLY", "cudagraph_capture_sizes": [8, 16, 32, 64, 96, 128, 192, 256, 384, 512]}' \ - --cpu-distributed-timeout-seconds 3600 \ - --enable-dbo \ - --dbo-decode-token-threshold 2 \ - --dbo-prefill-token-threshold 8 \ - --host 127.0.0.1 \ - --port 18001 > ffn.log 2>&1 & - -CUDA_VISIBLE_DEVICES=0,1,2,3 vllm serve "$MODEL_PATH" \ - --data-parallel-size 4 \ - --tensor-parallel-size 1 \ - --enable-expert-parallel \ - --additional-config '{ - "afd": { - "role": "attention", - "connector": "DeepEPM2NAFDConnector", - "host": "127.0.0.1", - "port": 6239, - "num_attention_ranks": 4, - "num_ffn_ranks": 2, - "compute_gate_on_attention": true, - "afd_process_group_timeout_s": 3600, - "connector_extra_config": {"payload_format": "fp8_block128_fp32"} - } - }' \ - --max-num-seqs 512 \ - --max-num-batched-tokens 3072 \ - --max-model-len 4096 \ - --max-cudagraph-capture-size 512 \ - --compilation-config '{"cudagraph_mode": "FULL_DECODE_ONLY", "cudagraph_capture_sizes": [8, 16, 32, 64, 96, 128, 192, 256, 384, 512]}' \ - --cpu-distributed-timeout-seconds 3600 \ - --enable-dbo \ - --dbo-decode-token-threshold 2 \ - --dbo-prefill-token-threshold 8 \ - --host 127.0.0.1 \ - --port 18000 > attn.log 2>&1 & - -wait diff --git a/recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/4a2f_graph_dp4tp1.sh b/recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/4a2f_graph_dp4tp1.sh deleted file mode 100755 index 960b7998d..000000000 --- a/recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/4a2f_graph_dp4tp1.sh +++ /dev/null @@ -1,191 +0,0 @@ -#!/usr/bin/env bash -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project -# -# Qwen3-235B-A22B-FP8 4A2F prefill/decode colocation on one node with 6 GPUs: -# Attention DP4/TP1 (or DP1/TP4 with ATTN_TP=4) and FFN DP2/TP1 (expert -# parallel), FULL_DECODE_ONLY CUDA graphs. Every knob is an environment -# variable (default in parentheses): -# -# MODEL_PATH weights dir or HF id (/path/model_weights/Qwen3-235B-A22B-FP8) -# AFD_CONNECTOR DeepEPM2NAFDConnector | P2pNcclAFDConnector -# (DeepEPM2NAFDConnector) -# M2N_PAYLOAD_FORMAT DeepEPM2NAFDConnector dispatch format (fp8_block128_fp32) -# M2N_EXTRA_CONFIG extra connector_extra_config JSON members for -# DeepEPM2NAFDConnector, e.g. "expert_num_sms": 130 (empty) -# KV_CONNECTOR none | decode_bench (none); see the README -# KV_TRANSFER_CONFIG raw --kv-transfer-config JSON for Attention; -# overrides KV_CONNECTOR (unset) -# ENABLE_PREFIX_CACHING 1 | 0 (1) -# ENABLE_DBO 1 | 0 (0) -# DBO_DECODE_TOKEN_THRESHOLD / DBO_PREFILL_TOKEN_THRESHOLD (2 / 2) -# MAX_NUM_SEQS per Attention DP rank (256) -# MAX_NUM_BATCHED_TOKENS prefill chunk per step (3072) -# CUDAGRAPH_CAPTURE_SIZES comma-separated decode graph sizes -# (8,12,16,24,32,48,64,96,128,192,256) -# MAX_MODEL_LEN (4096) -# GPU_MEMORY_UTILIZATION (unset: vLLM default) -# ATTN_DEVICES / FFN_DEVICES (0,1,2,3 / 4,5) -# ATTN_TP Attention tensor-parallel size; Attention DP is the -# number of ATTN_DEVICES divided by it (1) -# AFD_HOST / AFD_PORT AFD connector rendezvous (127.0.0.1 / 6239) -# AFD_PROCESS_GROUP_TIMEOUT_S AFD join timeout; the roles can finish -# loading minutes apart (3600) -# ATTN_PORT / FFN_PORT HTTP ports (18000 / 18001) -# EXTRA_ARGS extra `vllm serve` args for both roles (empty) -# ATTN_EXTRA_ARGS / FFN_EXTRA_ARGS extra `vllm serve` args for one role (empty) -# -# KV_CONNECTOR=decode_bench only (tools/benchmarks/decode_bench.py): -# DECODE_BENCH_NUM_PROMPTS pool prompts per Attention rank; 0 = dummy KV (500) -# DECODE_BENCH_PROMPT_LEN pool prompt length (1025) -# DECODE_BENCH_MAX_OUTPUT_LEN longest measured output; sizes the KV check (256) -# DECODE_BENCH_SEED pool token seed (0) -set -euo pipefail - -SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" -REPO_ROOT="$(cd "${SCRIPT_DIR}/../../../../.." && pwd)" - -MODEL_PATH=${MODEL_PATH:-/path/model_weights/Qwen3-235B-A22B-FP8} -AFD_CONNECTOR=${AFD_CONNECTOR:-DeepEPM2NAFDConnector} -M2N_PAYLOAD_FORMAT=${M2N_PAYLOAD_FORMAT:-fp8_block128_fp32} -M2N_EXTRA_CONFIG=${M2N_EXTRA_CONFIG:-} -KV_CONNECTOR=${KV_CONNECTOR:-none} -KV_TRANSFER_CONFIG=${KV_TRANSFER_CONFIG:-} -ENABLE_PREFIX_CACHING=${ENABLE_PREFIX_CACHING:-1} -ENABLE_DBO=${ENABLE_DBO:-0} -DBO_DECODE_TOKEN_THRESHOLD=${DBO_DECODE_TOKEN_THRESHOLD:-2} -DBO_PREFILL_TOKEN_THRESHOLD=${DBO_PREFILL_TOKEN_THRESHOLD:-2} -MAX_NUM_SEQS=${MAX_NUM_SEQS:-256} -MAX_NUM_BATCHED_TOKENS=${MAX_NUM_BATCHED_TOKENS:-3072} -CUDAGRAPH_CAPTURE_SIZES=${CUDAGRAPH_CAPTURE_SIZES:-8,12,16,24,32,48,64,96,128,192,256} -MAX_MODEL_LEN=${MAX_MODEL_LEN:-4096} -GPU_MEMORY_UTILIZATION=${GPU_MEMORY_UTILIZATION:-} -ATTN_DEVICES=${ATTN_DEVICES:-0,1,2,3} -FFN_DEVICES=${FFN_DEVICES:-4,5} -ATTN_TP=${ATTN_TP:-1} -AFD_HOST=${AFD_HOST:-127.0.0.1} -AFD_PORT=${AFD_PORT:-6239} -AFD_PROCESS_GROUP_TIMEOUT_S=${AFD_PROCESS_GROUP_TIMEOUT_S:-3600} -ATTN_PORT=${ATTN_PORT:-18000} -FFN_PORT=${FFN_PORT:-18001} -EXTRA_ARGS=${EXTRA_ARGS:-} -ATTN_EXTRA_ARGS=${ATTN_EXTRA_ARGS:-} -FFN_EXTRA_ARGS=${FFN_EXTRA_ARGS:-} -DECODE_BENCH_NUM_PROMPTS=${DECODE_BENCH_NUM_PROMPTS:-500} -DECODE_BENCH_PROMPT_LEN=${DECODE_BENCH_PROMPT_LEN:-1025} -DECODE_BENCH_MAX_OUTPUT_LEN=${DECODE_BENCH_MAX_OUTPUT_LEN:-256} -DECODE_BENCH_SEED=${DECODE_BENCH_SEED:-0} - -export VLLM_USE_V2_MODEL_RUNNER=1 - -NUM_ATTN_RANKS=$(tr ',' '\n' <<<"$ATTN_DEVICES" | wc -l | tr -d ' ') -NUM_FFN_RANKS=$(tr ',' '\n' <<<"$FFN_DEVICES" | wc -l | tr -d ' ') -if ((NUM_ATTN_RANKS % ATTN_TP)); then - echo "ATTN_TP=${ATTN_TP} does not divide the ${NUM_ATTN_RANKS} Attention devices" >&2 - exit 1 -fi -ATTN_DP=$((NUM_ATTN_RANKS / ATTN_TP)) - -case "$AFD_CONNECTOR" in - # DeepEP M2N routes on Attention and sends FP8 tokens to the FFN experts. - DeepEPM2NAFDConnector) - COMPUTE_GATE_ON_ATTENTION=true - CONNECTOR_EXTRA_CONFIG=", \"connector_extra_config\": {\"payload_format\": \"${M2N_PAYLOAD_FORMAT}\"${M2N_EXTRA_CONFIG:+, ${M2N_EXTRA_CONFIG}}}" - ;; - # P2P routes on the FFN and sends BF16 hidden states. Do not run its FFN - # with PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True: the first step hangs. - P2pNcclAFDConnector) - COMPUTE_GATE_ON_ATTENTION=false - CONNECTOR_EXTRA_CONFIG="" - ;; - *) - echo "unsupported AFD_CONNECTOR=${AFD_CONNECTOR}" >&2 - exit 1 - ;; -esac - -if [[ -z "$KV_TRANSFER_CONFIG" ]]; then - case "$KV_CONNECTOR" in - none) ;; - decode_bench) - KV_TRANSFER_CONFIG=$(cat <&2 - exit 1 - ;; - esac -fi -if [[ "$KV_TRANSFER_CONFIG" == *tools.benchmarks.* ]]; then - # tools/benchmarks is not in the wheel; the vLLM scheduler and worker - # subprocesses import it from the repository root. - export PYTHONPATH="${REPO_ROOT}${PYTHONPATH:+:${PYTHONPATH}}" -fi - -afd_config() { - cat < ffn.log 2>&1 & - -CUDA_VISIBLE_DEVICES=$ATTN_DEVICES vllm serve "$MODEL_PATH" \ - --additional-config "$(afd_config attention)" \ - "${COMMON_ARGS[@]}" \ - "${ATTN_ARGS[@]}" \ - --host 127.0.0.1 \ - --port "$ATTN_PORT" > attn.log 2>&1 & - -wait diff --git a/recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/dbo.sh b/recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/dbo.sh new file mode 100644 index 000000000..6530cac95 --- /dev/null +++ b/recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/dbo.sh @@ -0,0 +1,98 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +# +# Qwen3-235B-A22B-FP8 on DeepEPM2NAFDConnector, FULL_DECODE_ONLY graphs + DBO. +# Placement: one pod per comma-separated LAYOUT entry (default 2A2F, one pod of +# 4 GPUs). AF puts f FFN ranks on the pod's first f GPUs and a Attention +# ranks on the next a GPUs, e.g. LAYOUT=2A0F,0A2F or LAYOUT=1A1F,1A1F. +# For pod k (k = 0, 1, ...): POD=AFD_, with the same LAYOUT on every pod. +# Settings: recipe/gpu/DeepEPM2NAFDConnector/README.md#recipe-settings + +MODEL_PATH=${MODEL_PATH:-/path/model_weights/Qwen3-235B-A22B-FP8} +LAYOUT=${LAYOUT:-2A2F} +TUNING=(--max-num-seqs 512 --max-num-batched-tokens 3072 --max-model-len 4096) +CONNECTOR_CONFIG="\"payload_format\": \"fp8_block128_fp32\"${M2N_EXTRA_CONFIG:+, $M2N_EXTRA_CONFIG}" +CUDAGRAPH_CAPTURE_SIZES=${CUDAGRAPH_CAPTURE_SIZES:-8,16,32,64,96,128,192,256,384,512} # ascending +MODE=(--max-cudagraph-capture-size "${CUDAGRAPH_CAPTURE_SIZES##*,}" + --compilation-config "{\"cudagraph_mode\": \"FULL_DECODE_ONLY\", \"cudagraph_capture_sizes\": [${CUDAGRAPH_CAPTURE_SIZES}]}" + --enable-dbo --dbo-decode-token-threshold 2 --dbo-prefill-token-threshold 8) + +# ---- The rest is the same in every DeepEPM2NAFDConnector recipe. ---- +export VLLM_USE_V2_MODEL_RUNNER=1 +ATTN_TP=${ATTN_TP:-1} +AFD_CONNECTOR=${AFD_CONNECTOR:-DeepEPM2NAFDConnector} GATE_ON_ATTENTION=true +# P2P routes on the FFN and takes no connector_extra_config. +[[ $AFD_CONNECTOR == P2pNcclAFDConnector ]] && GATE_ON_ATTENTION=false CONNECTOR_CONFIG= +[[ -n $CONNECTOR_CONFIG ]] && CONNECTOR_CONFIG=", \"connector_extra_config\": {${CONNECTOR_CONFIG}}" +IFS=, read -ra PODS <<<"$LAYOUT" +K=${POD:-AFD_0} && K=${K#AFD_} +if [[ ! $LAYOUT =~ ^[0-9]+A[0-9]+F(,[0-9]+A[0-9]+F)*$ || ! $K =~ ^[0-9]+$ ]] || ((K >= ${#PODS[@]})); then + echo "need LAYOUT=AF[,...] and POD=AFD_ in it: LAYOUT=$LAYOUT POD=${POD:-}" >&2 + exit 1 +fi +NUM_A=() NUM_F=() +for p in "${PODS[@]}"; do NUM_A+=("${p%%A*}") && p=${p#*A} && NUM_F+=("${p%F}"); done +((${#PODS[@]} > 1)) || HEAD_HOST=${HEAD_HOST:-127.0.0.1} + +# Extra `vllm serve` flags go last, so they override the defaults. +read -ra ARGS <<<"${VLLM_ARGS:-}" && ARGS+=("$@") +read -ra ATTN_EXTRA <<<"${ATTN_ARGS:-}" +read -ra FFN_EXTRA <<<"${FFN_ARGS:-}" +if [[ -n ${DECODE_BENCH:-} ]]; then # tools/benchmarks is imported from the repository + ATTN_EXTRA=(--kv-transfer-config "{\"kv_connector\": \"AFDDecodeBenchConnector\", \"kv_connector_module_path\": \"tools.benchmarks.decode_bench\", \"kv_role\": \"kv_both\", \"kv_connector_extra_config\": ${DECODE_BENCH}}" "${ATTN_EXTRA[@]}") + export PYTHONPATH=${AFD_REPO:-$PWD}${PYTHONPATH:+:$PYTHONPATH} +fi + +# role_dp TP DP_HOST DP_RPC_PORT ... -> TOTAL and this pod's DP +# flags. A role spread over pods is led by its first pod; the others run +# headless from their start rank (the leader must omit it, or vLLM selects +# hybrid load balancing and rejects the headless peers). +role_dp() { + local tp=$1 host=$2 port=$3 i pods=0 start=0 lead=-1 + shift 3 + local n=("$@") + TOTAL=0 + for i in "${!n[@]}"; do + ((n[i])) || continue + ((n[i] % tp)) && { echo "TP $tp does not divide ${n[i]} ranks on pod $i" >&2; exit 1; } + ((lead < 0)) && lead=$i + ((i < K)) && start=$((start + n[i])) + pods=$((pods + 1)) TOTAL=$((TOTAL + n[i])) + done + DP=(--tensor-parallel-size "$tp" --data-parallel-size $((TOTAL / tp))) + ((pods > 1)) && DP+=(--data-parallel-size-local $((n[K] / tp)) + --data-parallel-address "$host" --data-parallel-rpc-port "$port") + ((pods > 1 && K != lead)) && DP+=(--headless --data-parallel-start-rank $((start / tp))) +} +role_dp 1 "${FFN_HEAD_HOST:-${HEAD_HOST:-vllm-ffn-dp-service}}" 13355 "${NUM_F[@]}" +FFN_DP=("${DP[@]}") NUM_FFN=$TOTAL +role_dp "$ATTN_TP" "${ATTN_HEAD_HOST:-${HEAD_HOST:-vllm-attn-dp-service}}" 13345 "${NUM_A[@]}" +ATTN_DP=("${DP[@]}") NUM_ATTN=$TOTAL + +afd() { # The AFD rendezvous runs on the pod that holds FFN rank 0. + cat < ffn.log 2>&1 & +fi +if ((A)); then + CUDA_VISIBLE_DEVICES=$(seq "$F" $((F + A - 1)) | paste -sd, -) vllm serve "$MODEL_PATH" \ + --additional-config "$(afd attention)" "${COMMON[@]}" "${ATTN_DP[@]}" \ + "${ARGS[@]}" "${ATTN_EXTRA[@]}" --port 18000 > attn.log 2>&1 & +fi +wait diff --git a/recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/eager.sh b/recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/eager.sh new file mode 100644 index 000000000..3ae1de33c --- /dev/null +++ b/recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/eager.sh @@ -0,0 +1,95 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +# +# Qwen3-235B-A22B-FP8 on DeepEPM2NAFDConnector, eager mode. +# Placement: one pod per comma-separated LAYOUT entry (default 2A2F, one pod of +# 4 GPUs). AF puts f FFN ranks on the pod's first f GPUs and a Attention +# ranks on the next a GPUs, e.g. LAYOUT=2A0F,0A2F or LAYOUT=1A1F,1A1F. +# For pod k (k = 0, 1, ...): POD=AFD_, with the same LAYOUT on every pod. +# Settings: recipe/gpu/DeepEPM2NAFDConnector/README.md#recipe-settings + +MODEL_PATH=${MODEL_PATH:-/path/model_weights/Qwen3-235B-A22B-FP8} +LAYOUT=${LAYOUT:-2A2F} +TUNING=(--max-num-seqs 256 --max-num-batched-tokens 3072 --max-model-len 4096) +CONNECTOR_CONFIG="\"payload_format\": \"fp8_block128_fp32\"${M2N_EXTRA_CONFIG:+, $M2N_EXTRA_CONFIG}" +MODE=(--enforce-eager) + +# ---- The rest is the same in every DeepEPM2NAFDConnector recipe. ---- +export VLLM_USE_V2_MODEL_RUNNER=1 +ATTN_TP=${ATTN_TP:-1} +AFD_CONNECTOR=${AFD_CONNECTOR:-DeepEPM2NAFDConnector} GATE_ON_ATTENTION=true +# P2P routes on the FFN and takes no connector_extra_config. +[[ $AFD_CONNECTOR == P2pNcclAFDConnector ]] && GATE_ON_ATTENTION=false CONNECTOR_CONFIG= +[[ -n $CONNECTOR_CONFIG ]] && CONNECTOR_CONFIG=", \"connector_extra_config\": {${CONNECTOR_CONFIG}}" +IFS=, read -ra PODS <<<"$LAYOUT" +K=${POD:-AFD_0} && K=${K#AFD_} +if [[ ! $LAYOUT =~ ^[0-9]+A[0-9]+F(,[0-9]+A[0-9]+F)*$ || ! $K =~ ^[0-9]+$ ]] || ((K >= ${#PODS[@]})); then + echo "need LAYOUT=AF[,...] and POD=AFD_ in it: LAYOUT=$LAYOUT POD=${POD:-}" >&2 + exit 1 +fi +NUM_A=() NUM_F=() +for p in "${PODS[@]}"; do NUM_A+=("${p%%A*}") && p=${p#*A} && NUM_F+=("${p%F}"); done +((${#PODS[@]} > 1)) || HEAD_HOST=${HEAD_HOST:-127.0.0.1} + +# Extra `vllm serve` flags go last, so they override the defaults. +read -ra ARGS <<<"${VLLM_ARGS:-}" && ARGS+=("$@") +read -ra ATTN_EXTRA <<<"${ATTN_ARGS:-}" +read -ra FFN_EXTRA <<<"${FFN_ARGS:-}" +if [[ -n ${DECODE_BENCH:-} ]]; then # tools/benchmarks is imported from the repository + ATTN_EXTRA=(--kv-transfer-config "{\"kv_connector\": \"AFDDecodeBenchConnector\", \"kv_connector_module_path\": \"tools.benchmarks.decode_bench\", \"kv_role\": \"kv_both\", \"kv_connector_extra_config\": ${DECODE_BENCH}}" "${ATTN_EXTRA[@]}") + export PYTHONPATH=${AFD_REPO:-$PWD}${PYTHONPATH:+:$PYTHONPATH} +fi + +# role_dp TP DP_HOST DP_RPC_PORT ... -> TOTAL and this pod's DP +# flags. A role spread over pods is led by its first pod; the others run +# headless from their start rank (the leader must omit it, or vLLM selects +# hybrid load balancing and rejects the headless peers). +role_dp() { + local tp=$1 host=$2 port=$3 i pods=0 start=0 lead=-1 + shift 3 + local n=("$@") + TOTAL=0 + for i in "${!n[@]}"; do + ((n[i])) || continue + ((n[i] % tp)) && { echo "TP $tp does not divide ${n[i]} ranks on pod $i" >&2; exit 1; } + ((lead < 0)) && lead=$i + ((i < K)) && start=$((start + n[i])) + pods=$((pods + 1)) TOTAL=$((TOTAL + n[i])) + done + DP=(--tensor-parallel-size "$tp" --data-parallel-size $((TOTAL / tp))) + ((pods > 1)) && DP+=(--data-parallel-size-local $((n[K] / tp)) + --data-parallel-address "$host" --data-parallel-rpc-port "$port") + ((pods > 1 && K != lead)) && DP+=(--headless --data-parallel-start-rank $((start / tp))) +} +role_dp 1 "${FFN_HEAD_HOST:-${HEAD_HOST:-vllm-ffn-dp-service}}" 13355 "${NUM_F[@]}" +FFN_DP=("${DP[@]}") NUM_FFN=$TOTAL +role_dp "$ATTN_TP" "${ATTN_HEAD_HOST:-${HEAD_HOST:-vllm-attn-dp-service}}" 13345 "${NUM_A[@]}" +ATTN_DP=("${DP[@]}") NUM_ATTN=$TOTAL + +afd() { # The AFD rendezvous runs on the pod that holds FFN rank 0. + cat < ffn.log 2>&1 & +fi +if ((A)); then + CUDA_VISIBLE_DEVICES=$(seq "$F" $((F + A - 1)) | paste -sd, -) vllm serve "$MODEL_PATH" \ + --additional-config "$(afd attention)" "${COMMON[@]}" "${ATTN_DP[@]}" \ + "${ARGS[@]}" "${ATTN_EXTRA[@]}" --port 18000 > attn.log 2>&1 & +fi +wait diff --git a/recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/graph.sh b/recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/graph.sh new file mode 100644 index 000000000..126ddf8a5 --- /dev/null +++ b/recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/graph.sh @@ -0,0 +1,97 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +# +# Qwen3-235B-A22B-FP8 on DeepEPM2NAFDConnector, FULL_DECODE_ONLY CUDA graphs. +# Placement: one pod per comma-separated LAYOUT entry (default 2A2F, one pod of +# 4 GPUs). AF puts f FFN ranks on the pod's first f GPUs and a Attention +# ranks on the next a GPUs, e.g. LAYOUT=2A0F,0A2F or LAYOUT=1A1F,1A1F. +# For pod k (k = 0, 1, ...): POD=AFD_, with the same LAYOUT on every pod. +# Settings: recipe/gpu/DeepEPM2NAFDConnector/README.md#recipe-settings + +MODEL_PATH=${MODEL_PATH:-/path/model_weights/Qwen3-235B-A22B-FP8} +LAYOUT=${LAYOUT:-2A2F} +TUNING=(--max-num-seqs 256 --max-num-batched-tokens 3072 --max-model-len 4096) +CONNECTOR_CONFIG="\"payload_format\": \"fp8_block128_fp32\"${M2N_EXTRA_CONFIG:+, $M2N_EXTRA_CONFIG}" +CUDAGRAPH_CAPTURE_SIZES=${CUDAGRAPH_CAPTURE_SIZES:-8,12,16,24,32,48,64,96,128,192,256} # ascending +MODE=(--max-cudagraph-capture-size "${CUDAGRAPH_CAPTURE_SIZES##*,}" + --compilation-config "{\"cudagraph_mode\": \"FULL_DECODE_ONLY\", \"cudagraph_capture_sizes\": [${CUDAGRAPH_CAPTURE_SIZES}]}") + +# ---- The rest is the same in every DeepEPM2NAFDConnector recipe. ---- +export VLLM_USE_V2_MODEL_RUNNER=1 +ATTN_TP=${ATTN_TP:-1} +AFD_CONNECTOR=${AFD_CONNECTOR:-DeepEPM2NAFDConnector} GATE_ON_ATTENTION=true +# P2P routes on the FFN and takes no connector_extra_config. +[[ $AFD_CONNECTOR == P2pNcclAFDConnector ]] && GATE_ON_ATTENTION=false CONNECTOR_CONFIG= +[[ -n $CONNECTOR_CONFIG ]] && CONNECTOR_CONFIG=", \"connector_extra_config\": {${CONNECTOR_CONFIG}}" +IFS=, read -ra PODS <<<"$LAYOUT" +K=${POD:-AFD_0} && K=${K#AFD_} +if [[ ! $LAYOUT =~ ^[0-9]+A[0-9]+F(,[0-9]+A[0-9]+F)*$ || ! $K =~ ^[0-9]+$ ]] || ((K >= ${#PODS[@]})); then + echo "need LAYOUT=AF[,...] and POD=AFD_ in it: LAYOUT=$LAYOUT POD=${POD:-}" >&2 + exit 1 +fi +NUM_A=() NUM_F=() +for p in "${PODS[@]}"; do NUM_A+=("${p%%A*}") && p=${p#*A} && NUM_F+=("${p%F}"); done +((${#PODS[@]} > 1)) || HEAD_HOST=${HEAD_HOST:-127.0.0.1} + +# Extra `vllm serve` flags go last, so they override the defaults. +read -ra ARGS <<<"${VLLM_ARGS:-}" && ARGS+=("$@") +read -ra ATTN_EXTRA <<<"${ATTN_ARGS:-}" +read -ra FFN_EXTRA <<<"${FFN_ARGS:-}" +if [[ -n ${DECODE_BENCH:-} ]]; then # tools/benchmarks is imported from the repository + ATTN_EXTRA=(--kv-transfer-config "{\"kv_connector\": \"AFDDecodeBenchConnector\", \"kv_connector_module_path\": \"tools.benchmarks.decode_bench\", \"kv_role\": \"kv_both\", \"kv_connector_extra_config\": ${DECODE_BENCH}}" "${ATTN_EXTRA[@]}") + export PYTHONPATH=${AFD_REPO:-$PWD}${PYTHONPATH:+:$PYTHONPATH} +fi + +# role_dp TP DP_HOST DP_RPC_PORT ... -> TOTAL and this pod's DP +# flags. A role spread over pods is led by its first pod; the others run +# headless from their start rank (the leader must omit it, or vLLM selects +# hybrid load balancing and rejects the headless peers). +role_dp() { + local tp=$1 host=$2 port=$3 i pods=0 start=0 lead=-1 + shift 3 + local n=("$@") + TOTAL=0 + for i in "${!n[@]}"; do + ((n[i])) || continue + ((n[i] % tp)) && { echo "TP $tp does not divide ${n[i]} ranks on pod $i" >&2; exit 1; } + ((lead < 0)) && lead=$i + ((i < K)) && start=$((start + n[i])) + pods=$((pods + 1)) TOTAL=$((TOTAL + n[i])) + done + DP=(--tensor-parallel-size "$tp" --data-parallel-size $((TOTAL / tp))) + ((pods > 1)) && DP+=(--data-parallel-size-local $((n[K] / tp)) + --data-parallel-address "$host" --data-parallel-rpc-port "$port") + ((pods > 1 && K != lead)) && DP+=(--headless --data-parallel-start-rank $((start / tp))) +} +role_dp 1 "${FFN_HEAD_HOST:-${HEAD_HOST:-vllm-ffn-dp-service}}" 13355 "${NUM_F[@]}" +FFN_DP=("${DP[@]}") NUM_FFN=$TOTAL +role_dp "$ATTN_TP" "${ATTN_HEAD_HOST:-${HEAD_HOST:-vllm-attn-dp-service}}" 13345 "${NUM_A[@]}" +ATTN_DP=("${DP[@]}") NUM_ATTN=$TOTAL + +afd() { # The AFD rendezvous runs on the pod that holds FFN rank 0. + cat < ffn.log 2>&1 & +fi +if ((A)); then + CUDA_VISIBLE_DEVICES=$(seq "$F" $((F + A - 1)) | paste -sd, -) vllm serve "$MODEL_PATH" \ + --additional-config "$(afd attention)" "${COMMON[@]}" "${ATTN_DP[@]}" \ + "${ARGS[@]}" "${ATTN_EXTRA[@]}" --port 18000 > attn.log 2>&1 & +fi +wait diff --git a/recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/multipod_4a1f_4a1f_graph.sh b/recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/multipod_4a1f_4a1f_graph.sh deleted file mode 100644 index 6ee50cd7b..000000000 --- a/recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/multipod_4a1f_4a1f_graph.sh +++ /dev/null @@ -1,82 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project - -# Qwen3-235B-A22B-FP8 with DeepEPM2NAFDConnector, 2 x 4A1F, FULL_DECODE_ONLY -# CUDA graphs: 2 pods of 5 GPUs, each with FFN on GPU 0 and Attention DP4 x TP1 -# on GPUs 1-4 (Attention DP8, FFN DP2/EP2 in total). Send traffic to port -# 18000 of the POD=AFD_0 pod. -# -# Use .agents/skills/deploy-afd-k8s/SKILL.md to deploy this recipe -# on 2 pods using the setup below. -# For pod 1: POD=AFD_0 -# For pod 2: POD=AFD_1 -# -# NUM_PODS x (ATTN_PER_POD Attention + FFN_PER_POD FFN) sets any symmetric -# layout (2 x 1A1F, 3 x 6A2F, ...); give every pod the same values. HEAD_HOST -# points the AFD rendezvous and both DP leaders at pod 0 (runs without k8s). - -MODEL_PATH=${MODEL_PATH:-/path/model_weights/Qwen3-235B-A22B-FP8} -NUM_PODS=${NUM_PODS:-2} -ATTN_PER_POD=${ATTN_PER_POD:-4} -FFN_PER_POD=${FFN_PER_POD:-1} -PAYLOAD_FORMAT=${PAYLOAD_FORMAT:-fp8_block128_fp32} -MAX_NUM_SEQS=${MAX_NUM_SEQS:-256} # per Attention DP rank -MAX_NUM_BATCHED_TOKENS=${MAX_NUM_BATCHED_TOKENS:-3072} -MAX_MODEL_LEN=${MAX_MODEL_LEN:-4096} -CUDAGRAPH_CAPTURE_SIZES=${CUDAGRAPH_CAPTURE_SIZES:-8,12,16,24,32,48,64,96,128,192,256} # ascending -ENABLE_DBO=${ENABLE_DBO:-0} -DBO_DECODE_TOKEN_THRESHOLD=${DBO_DECODE_TOKEN_THRESHOLD:-2} -DBO_PREFILL_TOKEN_THRESHOLD=${DBO_PREFILL_TOKEN_THRESHOLD:-2} -export VLLM_USE_V2_MODEL_RUNNER=1 - -K=${POD#AFD_} A=$ATTN_PER_POD F=$FFN_PER_POD - -# role_dp PER_POD ADDRESS RPC_PORT -> DP: this pod's DP flags for one role. -# Pod 0 leads and omits its start rank: vLLM reads an explicit start rank on a -# non-headless engine as hybrid load balancing and rejects the headless peers. -role_dp() { - DP=(--data-parallel-size $((NUM_PODS * $1)) --data-parallel-size-local "$1" - --data-parallel-address "$2" --data-parallel-rpc-port "$3") - if ((K > 0)); then DP+=(--headless --data-parallel-start-rank $((K * $1))); fi -} - -afd() { - cat < ffn.log 2>&1 & - -role_dp "$A" "${HEAD_HOST:-vllm-attn-dp-service}" 13345 -CUDA_VISIBLE_DEVICES=$(seq "$F" $((F + A - 1)) | paste -sd, -) vllm serve "$MODEL_PATH" \ - --additional-config "$(afd attention)" "${COMMON[@]}" "${DP[@]}" \ - --port 18000 > attn.log 2>&1 & - -wait diff --git a/recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/README.md b/recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/README.md index cc843b89c..9e9c7bbec 100644 --- a/recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/README.md +++ b/recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/README.md @@ -1,187 +1,42 @@ -# Qwen3-30B-A3B AFD Recipe +# Qwen3-30B-A3B with DeepEPM2NAFDConnector -Launch scripts for Qwen3-30B-A3B (BF16) in a 2A2F prefill/decode colocation on -a single 4-GPU node with vLLM `v0.32.0`. They use the same layouts as the Qwen3 -MoE DeepEP M2N E2E cases (`afd-deepep-m2n-*-2a2f` in -[`tests/e2e`](../../../../tests/e2e/README.md#deepep-m2n-scenarios)). -`2a2f_graph_dp2tp1.sh` also takes a selectable AFD connector and an optional KV -connector; with the decode-bench KV connector it doubles as a decode-only -benchmark over real KV (see [Decode-only benchmark](#decode-only-benchmark)). - -## Prerequisites - -- 4 GPUs on one node. `DeepEPM2NAFDConnector` needs SM90 or newer (H100, H200, - GB200) and the environment described in the - [connector recipes README](../README.md#prerequisites); - `P2pNcclAFDConnector` runs on any A/H-class GPU. -- Qwen3-30B-A3B weights on disk. The scripts default to - `/path/model_weights/Qwen3-30B-A3B`; override with `MODEL_PATH=...`. -- Free ports `6239` (AFD rendezvous), `18000` (Attention API) and `18001` - (FFN). - -## Layout - -`2A2F`: Attention on GPUs 0,1 and FFN DP2 on GPUs 2,3. FFN rank 0 owns experts -0-63 and FFN rank 1 owns experts 64-127. - -| GPUs | Role | Parallelism | Port | -|------|-----------|----------------------|-------| -| 0, 1 | Attention | DP2 / TP1 | 18000 | -| 2, 3 | FFN | DP2 / TP1, EP2 | 18001 | +Qwen3-30B-A3B (BF16) recipes. The scripts, layouts, settings, multi-pod runs +and the decode-only benchmark are described once in the +[connector recipes README](../README.md); this page lists what is specific to +the model. ```text prefill_decode_colocation/ -├── 2a2f_eager_dp2tp1.sh # Attention DP2 x TP1, eager -├── 2a2f_graph_dp2tp1.sh # Attention DP2 x TP1, FULL_DECODE_ONLY graph; every knob an env var -├── 2a2f_graph_dbo_dp2tp1.sh # Attention DP2 x TP1, graph + DBO -└── 2a2f_graph_dp1tp2.sh # Attention DP1 x TP2, FULL_DECODE_ONLY graph -``` - -The graph scripts run `FULL_DECODE_ONLY` CUDA graphs; prefill runs eager in -chunks of `--max-num-batched-tokens`. Send traffic to the Attention port. - -## Configuration notes - -- Both roles set `compute_gate_on_attention: true` for `DeepEPM2NAFDConnector` - (`false` for `P2pNcclAFDConnector`), use the same rank counts, and run with - `--enable-expert-parallel`. The FFN runs TP 1. -- `connector_extra_config` is left empty, so the payload is BF16 and the FFN - runs DeepGEMM experts. -- The graph scripts capture batch sizes 8, 16, 32 and 64. The connector - derives its dispatch buckets from the capture sizes and compiles one - transport kernel set per bucket at startup. -- In the DBO script, both roles pass `--enable-dbo` with the same thresholds; - the connector then uses two lanes. `2a2f_graph_dp2tp1.sh` does the same with - `ENABLE_DBO=1`. -- The dispatch capacity is derived from `--max-num-batched-tokens` and the - largest capture size, so lowering `--max-num-batched-tokens` shrinks the - transport buffers. - -## Running - -Run from the repository root: - -```bash -export MODEL_PATH=/path/model_weights/Qwen3-30B-A3B -# DeepEP M2N (default) -bash recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/2a2f_graph_dp2tp1.sh -# NCCL P2P: same script with AFD_CONNECTOR=P2pNcclAFDConnector, or its entry point -bash recipe/gpu/P2pNcclAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/2a2f_graph_dp2tp1.sh -``` - -The script backgrounds both roles, logs to `attn.log` and `ffn.log` in the -current directory, and waits. Wait for `Application startup complete` in -`attn.log` before sending traffic. - -Every knob of `2a2f_graph_dp2tp1.sh` is an environment variable documented at -the top of the script: connector (`AFD_CONNECTOR`), KV connector -(`KV_CONNECTOR`, or a raw `KV_TRANSFER_CONFIG` for any other connector), prefix -caching, DBO, `MAX_NUM_SEQS`, `MAX_NUM_BATCHED_TOKENS`, graph capture sizes, -devices, ports, and `EXTRA_ARGS` for anything else. - -## Checking a deployment - -```bash -curl -s http://127.0.0.1:18000/v1/completions \ - -H 'Content-Type: application/json' \ - -d '{"model":"'"$MODEL_PATH"'","prompt":"The capital of France is","max_tokens":12,"temperature":0}' +├── graph.sh # FULL_DECODE_ONLY CUDA graphs +├── eager.sh # eager +└── dbo.sh # graph + DBO ``` -For an accuracy check, run the matching E2E case, for example -`afd-deepep-m2n-graph-2a2f`; see [`tests/e2e`](../../../../tests/e2e/README.md#deepep-m2n-scenarios). - -## Decode-only benchmark - -`KV_CONNECTOR=decode_bench` attaches `AFDDecodeBenchConnector` -(`tools/benchmarks/decode_bench.py`) to Attention in its prompt-pool mode: +## Defaults -1. **Warmup.** Each Attention DP rank turns its first `DECODE_BENCH_NUM_PROMPTS` - requests into a pool of seeded random prompts of `DECODE_BENCH_PROMPT_LEN` - tokens (the same pool on every rank, whatever the requests carried) and - prefills them for real, one output token each. -2. **Pinning.** When a warmup request finishes, the connector takes a - reference of its own on the request's full prompt blocks and holds it for - the server's lifetime, while the request itself is freed as usual. Pinned - blocks cannot be evicted and stay in the prefix cache. -3. **Measured requests.** Every later request is remapped round-robin onto a - pinned pool prompt, whatever it asked for. The prefix cache covers every - prompt token but the last (the prompt length is rounded to one token past a - block boundary), so the request prefills one token and decodes over real KV, - on whichever rank serves it. A request that arrives before any prompt is - pinned may miss the cache; the connector logs it and the client's checks - fail the run. - -Start the server, then drive it with `tools/benchmarks/decode_bench_client.py`, -which runs every phase with `vllm bench serve`: - -```bash -export MODEL_PATH=/path/model_weights/Qwen3-30B-A3B -KV_CONNECTOR=decode_bench \ -DECODE_BENCH_NUM_PROMPTS=1000 \ -DECODE_BENCH_PROMPT_LEN=1025 \ -DECODE_BENCH_MAX_OUTPUT_LEN=256 \ - bash recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/2a2f_graph_dp2tp1.sh - -python tools/benchmarks/decode_bench_client.py \ - --port 18000 --dp-size 2 --pool-prompts 1000 --prompt-len 1025 \ - --result-dir /tmp/results/decode_bench \ - -- --num-prompts 4000 --random-output-len 256 \ - --max-concurrency 64 --request-rate 20 -``` - -1. **Warmup.** One `vllm bench serve` run per Attention rank, pinned with the - `X-data-parallel-rank` header, sends `--pool-prompts` requests to that rank. - `--warmup-concurrency` sets their concurrency; `--warmup-bench-args` (last - before `--`) passes anything else to these runs. -2. **Pre-check.** From the server's prefix-cache counters, every warmup - request must have looked up `--prompt-len` prompt tokens, so the pool - prompts are as long as the measured prompts. -3. **Measured run.** `vllm bench serve --dataset-name random --random-input-len - ` plus every argument after `--`: request rate, concurrency, - number of requests, output length, burstiness, percentiles, and so on. The - client sets the dataset, input length, endpoint, and result file itself and - rejects them after `--`. -4. **Post-check.** Over the measured run the server must have prefilled one - token per finished request (decode only), with no failed requests and no - preemptions. Otherwise the client exits non-zero. - -`--result-dir` receives vllm bench's own result JSON for each run -(`warmup-rank.json`, `measured.json`) plus `decode_bench_report.json` with -the checks. `--pool-prompts` must match `DECODE_BENCH_NUM_PROMPTS`, -`--prompt-len` must match `DECODE_BENCH_PROMPT_LEN`, and `--random-output-len` -must not exceed `DECODE_BENCH_MAX_OUTPUT_LEN`. The client needs vLLM and the -model's tokenizer installed, as vllm bench does. - -### Sizing the pool - -The pool must fit each Attention rank's KV cache next to the running requests. -At startup the connector checks, per rank, - -```text -num_prompts * ceil((prompt_len + 1) / block_size) - + max_num_seqs * ceil((max_output_len + 1) / block_size) <= free KV blocks -``` +| Setting | Value | +| --- | --- | +| `LAYOUT` | `2A2F`: FFN DP2 on GPUs 0,1, Attention DP2 on GPUs 2,3 | +| `--max-num-seqs` | 64 per Attention DP rank | +| `--max-num-batched-tokens` | 4096 | +| `--max-model-len` | 4096 | +| `CUDAGRAPH_CAPTURE_SIZES` | 8,16,32,64 | +| Payload | BF16 (no `connector_extra_config`); the FFN runs DeepGEMM experts | -and fails with the largest `num_prompts` that fits. Qwen3-30B-A3B holds about -96 KiB of BF16 KV per token, so 1000 prompts of 1025 tokens need about 94 GiB -per rank: an H200 fits it, an H100 needs fewer or shorter prompts. Keep -`num_prompts` at or above `MAX_NUM_SEQS` (the connector warns otherwise), or -concurrent requests share physical KV blocks and attention reads look cheaper -than they are. +These match the Qwen3 MoE DeepEP M2N E2E cases (`afd-deepep-m2n-*-2a2f` and +the two-pod `2A0F,0A2F` and `1A1F,1A1F` cases in +[`tests/e2e`](../../../../tests/e2e/README.md#deepep-m2n-scenarios)); +`ATTN_TP=2` matches the Attention TP2 cases. -### Limitations +## Notes -- Generated text is meaningless: prompts are random tokens, and the client's - prompt is discarded. -- Measured requests decode at a context of `prompt_len` regardless of the - client's prompt length. -- Pinned warmup requests keep each Attention engine loop stepping while idle - (one busy CPU core per DP rank, no GPU work). Restart the server to free the - pool. -- `DECODE_BENCH_NUM_PROMPTS=0` keeps the older dummy-KV mode, which writes no - KV at all. +- 4 GPUs of 80 GB (H100) or more fit `2A2F`. With two FFN ranks, FFN rank 0 + owns experts 0-63 and FFN rank 1 owns experts 64-127; the FFN count must + divide the 128 experts. +- Decode-only benchmark pool sizing: Qwen3-30B-A3B holds about 96 KiB of BF16 + KV per token, so 500 prompts of 1025 tokens need about 47 GiB per Attention + rank and fit an H100; 1000 prompts need about 94 GiB and need an H200. ## Limitations - Qwen3 MoE only; see the connector guide for the full list of requirements. -- `FULL_DECODE_ONLY` is the only CUDA graph mode. diff --git a/recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/2a2f_eager_dp2tp1.sh b/recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/2a2f_eager_dp2tp1.sh deleted file mode 100644 index 18b37e29a..000000000 --- a/recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/2a2f_eager_dp2tp1.sh +++ /dev/null @@ -1,53 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project - -# 2A2F eager: Attention DP2 x TP1 on GPUs 0,1; FFN DP2 on GPUs 2,3. - -MODEL_PATH=${MODEL_PATH:-/path/model_weights/Qwen3-30B-A3B} -export VLLM_USE_V2_MODEL_RUNNER=1 - -CUDA_VISIBLE_DEVICES=2,3 vllm serve "$MODEL_PATH" \ - --data-parallel-size 2 \ - --tensor-parallel-size 1 \ - --enable-expert-parallel \ - --additional-config '{ - "afd": { - "role": "ffn", - "connector": "DeepEPM2NAFDConnector", - "host": "127.0.0.1", - "port": 6239, - "num_attention_ranks": 2, - "num_ffn_ranks": 2, - "compute_gate_on_attention": true - } - }' \ - --max-num-seqs 64 \ - --max-num-batched-tokens 4096 \ - --max-model-len 4096 \ - --enforce-eager \ - --host 127.0.0.1 \ - --port 18001 > ffn.log 2>&1 & - -CUDA_VISIBLE_DEVICES=0,1 vllm serve "$MODEL_PATH" \ - --data-parallel-size 2 \ - --tensor-parallel-size 1 \ - --enable-expert-parallel \ - --additional-config '{ - "afd": { - "role": "attention", - "connector": "DeepEPM2NAFDConnector", - "host": "127.0.0.1", - "port": 6239, - "num_attention_ranks": 2, - "num_ffn_ranks": 2, - "compute_gate_on_attention": true - } - }' \ - --max-num-seqs 64 \ - --max-num-batched-tokens 4096 \ - --max-model-len 4096 \ - --enforce-eager \ - --host 127.0.0.1 \ - --port 18000 > attn.log 2>&1 & - -wait diff --git a/recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/2a2f_graph_dbo_dp2tp1.sh b/recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/2a2f_graph_dbo_dp2tp1.sh deleted file mode 100644 index 1736b681c..000000000 --- a/recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/2a2f_graph_dbo_dp2tp1.sh +++ /dev/null @@ -1,62 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project - -# 2A2F graph + DBO: Attention DP2 x TP1 on GPUs 0,1; FFN DP2 on GPUs 2,3. -# DBO must be enabled on both roles. - -MODEL_PATH=${MODEL_PATH:-/path/model_weights/Qwen3-30B-A3B} -export VLLM_USE_V2_MODEL_RUNNER=1 - -CUDA_VISIBLE_DEVICES=2,3 vllm serve "$MODEL_PATH" \ - --data-parallel-size 2 \ - --tensor-parallel-size 1 \ - --enable-expert-parallel \ - --additional-config '{ - "afd": { - "role": "ffn", - "connector": "DeepEPM2NAFDConnector", - "host": "127.0.0.1", - "port": 6239, - "num_attention_ranks": 2, - "num_ffn_ranks": 2, - "compute_gate_on_attention": true - } - }' \ - --max-num-seqs 64 \ - --max-num-batched-tokens 4096 \ - --max-model-len 4096 \ - --max-cudagraph-capture-size 64 \ - --compilation-config '{"cudagraph_mode": "FULL_DECODE_ONLY", "cudagraph_capture_sizes": [8, 16, 32, 64]}' \ - --enable-dbo \ - --dbo-decode-token-threshold 2 \ - --dbo-prefill-token-threshold 8 \ - --host 127.0.0.1 \ - --port 18001 > ffn.log 2>&1 & - -CUDA_VISIBLE_DEVICES=0,1 vllm serve "$MODEL_PATH" \ - --data-parallel-size 2 \ - --tensor-parallel-size 1 \ - --enable-expert-parallel \ - --additional-config '{ - "afd": { - "role": "attention", - "connector": "DeepEPM2NAFDConnector", - "host": "127.0.0.1", - "port": 6239, - "num_attention_ranks": 2, - "num_ffn_ranks": 2, - "compute_gate_on_attention": true - } - }' \ - --max-num-seqs 64 \ - --max-num-batched-tokens 4096 \ - --max-model-len 4096 \ - --max-cudagraph-capture-size 64 \ - --compilation-config '{"cudagraph_mode": "FULL_DECODE_ONLY", "cudagraph_capture_sizes": [8, 16, 32, 64]}' \ - --enable-dbo \ - --dbo-decode-token-threshold 2 \ - --dbo-prefill-token-threshold 8 \ - --host 127.0.0.1 \ - --port 18000 > attn.log 2>&1 & - -wait diff --git a/recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/2a2f_graph_dp1tp2.sh b/recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/2a2f_graph_dp1tp2.sh deleted file mode 100644 index 447ace10f..000000000 --- a/recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/2a2f_graph_dp1tp2.sh +++ /dev/null @@ -1,56 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project - -# 2A2F graph with Attention TP: Attention DP1 x TP2 on GPUs 0,1; FFN DP2 x TP1 on -# GPUs 2,3. Each Attention TP rank dispatches its slice of the tokens. - -MODEL_PATH=${MODEL_PATH:-/path/model_weights/Qwen3-30B-A3B} -export VLLM_USE_V2_MODEL_RUNNER=1 - -CUDA_VISIBLE_DEVICES=2,3 vllm serve "$MODEL_PATH" \ - --data-parallel-size 2 \ - --tensor-parallel-size 1 \ - --enable-expert-parallel \ - --additional-config '{ - "afd": { - "role": "ffn", - "connector": "DeepEPM2NAFDConnector", - "host": "127.0.0.1", - "port": 6239, - "num_attention_ranks": 2, - "num_ffn_ranks": 2, - "compute_gate_on_attention": true - } - }' \ - --max-num-seqs 64 \ - --max-num-batched-tokens 4096 \ - --max-model-len 4096 \ - --max-cudagraph-capture-size 64 \ - --compilation-config '{"cudagraph_mode": "FULL_DECODE_ONLY", "cudagraph_capture_sizes": [8, 16, 32, 64]}' \ - --host 127.0.0.1 \ - --port 18001 > ffn.log 2>&1 & - -CUDA_VISIBLE_DEVICES=0,1 vllm serve "$MODEL_PATH" \ - --data-parallel-size 1 \ - --tensor-parallel-size 2 \ - --enable-expert-parallel \ - --additional-config '{ - "afd": { - "role": "attention", - "connector": "DeepEPM2NAFDConnector", - "host": "127.0.0.1", - "port": 6239, - "num_attention_ranks": 2, - "num_ffn_ranks": 2, - "compute_gate_on_attention": true - } - }' \ - --max-num-seqs 64 \ - --max-num-batched-tokens 4096 \ - --max-model-len 4096 \ - --max-cudagraph-capture-size 64 \ - --compilation-config '{"cudagraph_mode": "FULL_DECODE_ONLY", "cudagraph_capture_sizes": [8, 16, 32, 64]}' \ - --host 127.0.0.1 \ - --port 18000 > attn.log 2>&1 & - -wait diff --git a/recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/2a2f_graph_dp2tp1.sh b/recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/2a2f_graph_dp2tp1.sh deleted file mode 100755 index fa58b4953..000000000 --- a/recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/2a2f_graph_dp2tp1.sh +++ /dev/null @@ -1,153 +0,0 @@ -#!/usr/bin/env bash -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project -# -# Qwen3-30B-A3B 2A2F prefill/decode colocation on one 4-GPU node: Attention -# DP2/TP1 and FFN DP2/TP1 (expert parallel), FULL_DECODE_ONLY CUDA graphs. -# Every knob is an environment variable (default in parentheses): -# -# MODEL_PATH weights dir or HF id (/path/model_weights/Qwen3-30B-A3B) -# AFD_CONNECTOR DeepEPM2NAFDConnector | P2pNcclAFDConnector -# (DeepEPM2NAFDConnector) -# KV_CONNECTOR none | decode_bench (none); see the README -# KV_TRANSFER_CONFIG raw --kv-transfer-config JSON for Attention; -# overrides KV_CONNECTOR (unset) -# ENABLE_PREFIX_CACHING 1 | 0 (1) -# ENABLE_DBO 1 | 0 (0) -# DBO_DECODE_TOKEN_THRESHOLD / DBO_PREFILL_TOKEN_THRESHOLD (2 / 12) -# MAX_NUM_SEQS (64) -# MAX_NUM_BATCHED_TOKENS prefill chunk per step (4096) -# CUDAGRAPH_CAPTURE_SIZES comma-separated decode graph sizes (8,16,32,64) -# MAX_MODEL_LEN (4096) -# GPU_MEMORY_UTILIZATION (unset: vLLM default) -# ATTN_DEVICES / FFN_DEVICES (0,1 / 2,3) -# AFD_HOST / AFD_PORT AFD connector rendezvous (127.0.0.1 / 6239) -# ATTN_PORT / FFN_PORT HTTP ports (18000 / 18001) -# EXTRA_ARGS extra `vllm serve` args for both roles (empty) -# -# KV_CONNECTOR=decode_bench only (tools/benchmarks/decode_bench.py): -# DECODE_BENCH_NUM_PROMPTS pool prompts per Attention rank; 0 = dummy KV (1000) -# DECODE_BENCH_PROMPT_LEN pool prompt length (1025) -# DECODE_BENCH_MAX_OUTPUT_LEN longest measured output; sizes the KV check (1024) -# DECODE_BENCH_SEED pool token seed (0) -set -euo pipefail - -SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" -REPO_ROOT="$(cd "${SCRIPT_DIR}/../../../../.." && pwd)" - -MODEL_PATH=${MODEL_PATH:-/path/model_weights/Qwen3-30B-A3B} -AFD_CONNECTOR=${AFD_CONNECTOR:-DeepEPM2NAFDConnector} -KV_CONNECTOR=${KV_CONNECTOR:-none} -KV_TRANSFER_CONFIG=${KV_TRANSFER_CONFIG:-} -ENABLE_PREFIX_CACHING=${ENABLE_PREFIX_CACHING:-1} -ENABLE_DBO=${ENABLE_DBO:-0} -DBO_DECODE_TOKEN_THRESHOLD=${DBO_DECODE_TOKEN_THRESHOLD:-2} -DBO_PREFILL_TOKEN_THRESHOLD=${DBO_PREFILL_TOKEN_THRESHOLD:-12} -MAX_NUM_SEQS=${MAX_NUM_SEQS:-64} -MAX_NUM_BATCHED_TOKENS=${MAX_NUM_BATCHED_TOKENS:-4096} -CUDAGRAPH_CAPTURE_SIZES=${CUDAGRAPH_CAPTURE_SIZES:-8,16,32,64} -MAX_MODEL_LEN=${MAX_MODEL_LEN:-4096} -GPU_MEMORY_UTILIZATION=${GPU_MEMORY_UTILIZATION:-} -ATTN_DEVICES=${ATTN_DEVICES:-0,1} -FFN_DEVICES=${FFN_DEVICES:-2,3} -AFD_HOST=${AFD_HOST:-127.0.0.1} -AFD_PORT=${AFD_PORT:-6239} -ATTN_PORT=${ATTN_PORT:-18000} -FFN_PORT=${FFN_PORT:-18001} -EXTRA_ARGS=${EXTRA_ARGS:-} -DECODE_BENCH_NUM_PROMPTS=${DECODE_BENCH_NUM_PROMPTS:-1000} -DECODE_BENCH_PROMPT_LEN=${DECODE_BENCH_PROMPT_LEN:-1025} -DECODE_BENCH_MAX_OUTPUT_LEN=${DECODE_BENCH_MAX_OUTPUT_LEN:-1024} -DECODE_BENCH_SEED=${DECODE_BENCH_SEED:-0} - -export VLLM_USE_V2_MODEL_RUNNER=1 - -case "$AFD_CONNECTOR" in - # DeepEP M2N routes on Attention: the gate and top-k run there. - DeepEPM2NAFDConnector) COMPUTE_GATE_ON_ATTENTION=true ;; - P2pNcclAFDConnector) COMPUTE_GATE_ON_ATTENTION=false ;; - *) - echo "unsupported AFD_CONNECTOR=${AFD_CONNECTOR}" >&2 - exit 1 - ;; -esac - -if [[ -z "$KV_TRANSFER_CONFIG" ]]; then - case "$KV_CONNECTOR" in - none) ;; - decode_bench) - KV_TRANSFER_CONFIG=$(cat <&2 - exit 1 - ;; - esac -fi -if [[ "$KV_TRANSFER_CONFIG" == *tools.benchmarks.* ]]; then - # tools/benchmarks is not in the wheel; the vLLM scheduler and worker - # subprocesses import it from the repository root. - export PYTHONPATH="${REPO_ROOT}${PYTHONPATH:+:${PYTHONPATH}}" -fi - -afd_config() { - cat < ffn.log 2>&1 & - -CUDA_VISIBLE_DEVICES=$ATTN_DEVICES vllm serve "$MODEL_PATH" \ - --additional-config "$(afd_config attention)" \ - "${COMMON_ARGS[@]}" \ - "${ATTN_ARGS[@]}" \ - --host 127.0.0.1 \ - --port "$ATTN_PORT" > attn.log 2>&1 & - -wait diff --git a/recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/dbo.sh b/recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/dbo.sh new file mode 100644 index 000000000..8a611278a --- /dev/null +++ b/recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/dbo.sh @@ -0,0 +1,98 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +# +# Qwen3-30B-A3B (BF16) on DeepEPM2NAFDConnector, FULL_DECODE_ONLY graphs + DBO. +# Placement: one pod per comma-separated LAYOUT entry (default 2A2F, one pod of +# 4 GPUs). AF puts f FFN ranks on the pod's first f GPUs and a Attention +# ranks on the next a GPUs, e.g. LAYOUT=2A0F,0A2F or LAYOUT=1A1F,1A1F. +# For pod k (k = 0, 1, ...): POD=AFD_, with the same LAYOUT on every pod. +# Settings: recipe/gpu/DeepEPM2NAFDConnector/README.md#recipe-settings + +MODEL_PATH=${MODEL_PATH:-/path/model_weights/Qwen3-30B-A3B} +LAYOUT=${LAYOUT:-2A2F} +TUNING=(--max-num-seqs 64 --max-num-batched-tokens 4096 --max-model-len 4096) +CONNECTOR_CONFIG=${M2N_EXTRA_CONFIG:-} # connector_extra_config members; BF16 payload +CUDAGRAPH_CAPTURE_SIZES=${CUDAGRAPH_CAPTURE_SIZES:-8,16,32,64} # ascending +MODE=(--max-cudagraph-capture-size "${CUDAGRAPH_CAPTURE_SIZES##*,}" + --compilation-config "{\"cudagraph_mode\": \"FULL_DECODE_ONLY\", \"cudagraph_capture_sizes\": [${CUDAGRAPH_CAPTURE_SIZES}]}" + --enable-dbo --dbo-decode-token-threshold 2 --dbo-prefill-token-threshold 8) + +# ---- The rest is the same in every DeepEPM2NAFDConnector recipe. ---- +export VLLM_USE_V2_MODEL_RUNNER=1 +ATTN_TP=${ATTN_TP:-1} +AFD_CONNECTOR=${AFD_CONNECTOR:-DeepEPM2NAFDConnector} GATE_ON_ATTENTION=true +# P2P routes on the FFN and takes no connector_extra_config. +[[ $AFD_CONNECTOR == P2pNcclAFDConnector ]] && GATE_ON_ATTENTION=false CONNECTOR_CONFIG= +[[ -n $CONNECTOR_CONFIG ]] && CONNECTOR_CONFIG=", \"connector_extra_config\": {${CONNECTOR_CONFIG}}" +IFS=, read -ra PODS <<<"$LAYOUT" +K=${POD:-AFD_0} && K=${K#AFD_} +if [[ ! $LAYOUT =~ ^[0-9]+A[0-9]+F(,[0-9]+A[0-9]+F)*$ || ! $K =~ ^[0-9]+$ ]] || ((K >= ${#PODS[@]})); then + echo "need LAYOUT=AF[,...] and POD=AFD_ in it: LAYOUT=$LAYOUT POD=${POD:-}" >&2 + exit 1 +fi +NUM_A=() NUM_F=() +for p in "${PODS[@]}"; do NUM_A+=("${p%%A*}") && p=${p#*A} && NUM_F+=("${p%F}"); done +((${#PODS[@]} > 1)) || HEAD_HOST=${HEAD_HOST:-127.0.0.1} + +# Extra `vllm serve` flags go last, so they override the defaults. +read -ra ARGS <<<"${VLLM_ARGS:-}" && ARGS+=("$@") +read -ra ATTN_EXTRA <<<"${ATTN_ARGS:-}" +read -ra FFN_EXTRA <<<"${FFN_ARGS:-}" +if [[ -n ${DECODE_BENCH:-} ]]; then # tools/benchmarks is imported from the repository + ATTN_EXTRA=(--kv-transfer-config "{\"kv_connector\": \"AFDDecodeBenchConnector\", \"kv_connector_module_path\": \"tools.benchmarks.decode_bench\", \"kv_role\": \"kv_both\", \"kv_connector_extra_config\": ${DECODE_BENCH}}" "${ATTN_EXTRA[@]}") + export PYTHONPATH=${AFD_REPO:-$PWD}${PYTHONPATH:+:$PYTHONPATH} +fi + +# role_dp TP DP_HOST DP_RPC_PORT ... -> TOTAL and this pod's DP +# flags. A role spread over pods is led by its first pod; the others run +# headless from their start rank (the leader must omit it, or vLLM selects +# hybrid load balancing and rejects the headless peers). +role_dp() { + local tp=$1 host=$2 port=$3 i pods=0 start=0 lead=-1 + shift 3 + local n=("$@") + TOTAL=0 + for i in "${!n[@]}"; do + ((n[i])) || continue + ((n[i] % tp)) && { echo "TP $tp does not divide ${n[i]} ranks on pod $i" >&2; exit 1; } + ((lead < 0)) && lead=$i + ((i < K)) && start=$((start + n[i])) + pods=$((pods + 1)) TOTAL=$((TOTAL + n[i])) + done + DP=(--tensor-parallel-size "$tp" --data-parallel-size $((TOTAL / tp))) + ((pods > 1)) && DP+=(--data-parallel-size-local $((n[K] / tp)) + --data-parallel-address "$host" --data-parallel-rpc-port "$port") + ((pods > 1 && K != lead)) && DP+=(--headless --data-parallel-start-rank $((start / tp))) +} +role_dp 1 "${FFN_HEAD_HOST:-${HEAD_HOST:-vllm-ffn-dp-service}}" 13355 "${NUM_F[@]}" +FFN_DP=("${DP[@]}") NUM_FFN=$TOTAL +role_dp "$ATTN_TP" "${ATTN_HEAD_HOST:-${HEAD_HOST:-vllm-attn-dp-service}}" 13345 "${NUM_A[@]}" +ATTN_DP=("${DP[@]}") NUM_ATTN=$TOTAL + +afd() { # The AFD rendezvous runs on the pod that holds FFN rank 0. + cat < ffn.log 2>&1 & +fi +if ((A)); then + CUDA_VISIBLE_DEVICES=$(seq "$F" $((F + A - 1)) | paste -sd, -) vllm serve "$MODEL_PATH" \ + --additional-config "$(afd attention)" "${COMMON[@]}" "${ATTN_DP[@]}" \ + "${ARGS[@]}" "${ATTN_EXTRA[@]}" --port 18000 > attn.log 2>&1 & +fi +wait diff --git a/recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/eager.sh b/recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/eager.sh new file mode 100644 index 000000000..656764492 --- /dev/null +++ b/recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/eager.sh @@ -0,0 +1,95 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +# +# Qwen3-30B-A3B (BF16) on DeepEPM2NAFDConnector, eager mode. +# Placement: one pod per comma-separated LAYOUT entry (default 2A2F, one pod of +# 4 GPUs). AF puts f FFN ranks on the pod's first f GPUs and a Attention +# ranks on the next a GPUs, e.g. LAYOUT=2A0F,0A2F or LAYOUT=1A1F,1A1F. +# For pod k (k = 0, 1, ...): POD=AFD_, with the same LAYOUT on every pod. +# Settings: recipe/gpu/DeepEPM2NAFDConnector/README.md#recipe-settings + +MODEL_PATH=${MODEL_PATH:-/path/model_weights/Qwen3-30B-A3B} +LAYOUT=${LAYOUT:-2A2F} +TUNING=(--max-num-seqs 64 --max-num-batched-tokens 4096 --max-model-len 4096) +CONNECTOR_CONFIG=${M2N_EXTRA_CONFIG:-} # connector_extra_config members; BF16 payload +MODE=(--enforce-eager) + +# ---- The rest is the same in every DeepEPM2NAFDConnector recipe. ---- +export VLLM_USE_V2_MODEL_RUNNER=1 +ATTN_TP=${ATTN_TP:-1} +AFD_CONNECTOR=${AFD_CONNECTOR:-DeepEPM2NAFDConnector} GATE_ON_ATTENTION=true +# P2P routes on the FFN and takes no connector_extra_config. +[[ $AFD_CONNECTOR == P2pNcclAFDConnector ]] && GATE_ON_ATTENTION=false CONNECTOR_CONFIG= +[[ -n $CONNECTOR_CONFIG ]] && CONNECTOR_CONFIG=", \"connector_extra_config\": {${CONNECTOR_CONFIG}}" +IFS=, read -ra PODS <<<"$LAYOUT" +K=${POD:-AFD_0} && K=${K#AFD_} +if [[ ! $LAYOUT =~ ^[0-9]+A[0-9]+F(,[0-9]+A[0-9]+F)*$ || ! $K =~ ^[0-9]+$ ]] || ((K >= ${#PODS[@]})); then + echo "need LAYOUT=AF[,...] and POD=AFD_ in it: LAYOUT=$LAYOUT POD=${POD:-}" >&2 + exit 1 +fi +NUM_A=() NUM_F=() +for p in "${PODS[@]}"; do NUM_A+=("${p%%A*}") && p=${p#*A} && NUM_F+=("${p%F}"); done +((${#PODS[@]} > 1)) || HEAD_HOST=${HEAD_HOST:-127.0.0.1} + +# Extra `vllm serve` flags go last, so they override the defaults. +read -ra ARGS <<<"${VLLM_ARGS:-}" && ARGS+=("$@") +read -ra ATTN_EXTRA <<<"${ATTN_ARGS:-}" +read -ra FFN_EXTRA <<<"${FFN_ARGS:-}" +if [[ -n ${DECODE_BENCH:-} ]]; then # tools/benchmarks is imported from the repository + ATTN_EXTRA=(--kv-transfer-config "{\"kv_connector\": \"AFDDecodeBenchConnector\", \"kv_connector_module_path\": \"tools.benchmarks.decode_bench\", \"kv_role\": \"kv_both\", \"kv_connector_extra_config\": ${DECODE_BENCH}}" "${ATTN_EXTRA[@]}") + export PYTHONPATH=${AFD_REPO:-$PWD}${PYTHONPATH:+:$PYTHONPATH} +fi + +# role_dp TP DP_HOST DP_RPC_PORT ... -> TOTAL and this pod's DP +# flags. A role spread over pods is led by its first pod; the others run +# headless from their start rank (the leader must omit it, or vLLM selects +# hybrid load balancing and rejects the headless peers). +role_dp() { + local tp=$1 host=$2 port=$3 i pods=0 start=0 lead=-1 + shift 3 + local n=("$@") + TOTAL=0 + for i in "${!n[@]}"; do + ((n[i])) || continue + ((n[i] % tp)) && { echo "TP $tp does not divide ${n[i]} ranks on pod $i" >&2; exit 1; } + ((lead < 0)) && lead=$i + ((i < K)) && start=$((start + n[i])) + pods=$((pods + 1)) TOTAL=$((TOTAL + n[i])) + done + DP=(--tensor-parallel-size "$tp" --data-parallel-size $((TOTAL / tp))) + ((pods > 1)) && DP+=(--data-parallel-size-local $((n[K] / tp)) + --data-parallel-address "$host" --data-parallel-rpc-port "$port") + ((pods > 1 && K != lead)) && DP+=(--headless --data-parallel-start-rank $((start / tp))) +} +role_dp 1 "${FFN_HEAD_HOST:-${HEAD_HOST:-vllm-ffn-dp-service}}" 13355 "${NUM_F[@]}" +FFN_DP=("${DP[@]}") NUM_FFN=$TOTAL +role_dp "$ATTN_TP" "${ATTN_HEAD_HOST:-${HEAD_HOST:-vllm-attn-dp-service}}" 13345 "${NUM_A[@]}" +ATTN_DP=("${DP[@]}") NUM_ATTN=$TOTAL + +afd() { # The AFD rendezvous runs on the pod that holds FFN rank 0. + cat < ffn.log 2>&1 & +fi +if ((A)); then + CUDA_VISIBLE_DEVICES=$(seq "$F" $((F + A - 1)) | paste -sd, -) vllm serve "$MODEL_PATH" \ + --additional-config "$(afd attention)" "${COMMON[@]}" "${ATTN_DP[@]}" \ + "${ARGS[@]}" "${ATTN_EXTRA[@]}" --port 18000 > attn.log 2>&1 & +fi +wait diff --git a/recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/graph.sh b/recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/graph.sh new file mode 100644 index 000000000..d8ea0feed --- /dev/null +++ b/recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/graph.sh @@ -0,0 +1,97 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +# +# Qwen3-30B-A3B (BF16) on DeepEPM2NAFDConnector, FULL_DECODE_ONLY CUDA graphs. +# Placement: one pod per comma-separated LAYOUT entry (default 2A2F, one pod of +# 4 GPUs). AF puts f FFN ranks on the pod's first f GPUs and a Attention +# ranks on the next a GPUs, e.g. LAYOUT=2A0F,0A2F or LAYOUT=1A1F,1A1F. +# For pod k (k = 0, 1, ...): POD=AFD_, with the same LAYOUT on every pod. +# Settings: recipe/gpu/DeepEPM2NAFDConnector/README.md#recipe-settings + +MODEL_PATH=${MODEL_PATH:-/path/model_weights/Qwen3-30B-A3B} +LAYOUT=${LAYOUT:-2A2F} +TUNING=(--max-num-seqs 64 --max-num-batched-tokens 4096 --max-model-len 4096) +CONNECTOR_CONFIG=${M2N_EXTRA_CONFIG:-} # connector_extra_config members; BF16 payload +CUDAGRAPH_CAPTURE_SIZES=${CUDAGRAPH_CAPTURE_SIZES:-8,16,32,64} # ascending +MODE=(--max-cudagraph-capture-size "${CUDAGRAPH_CAPTURE_SIZES##*,}" + --compilation-config "{\"cudagraph_mode\": \"FULL_DECODE_ONLY\", \"cudagraph_capture_sizes\": [${CUDAGRAPH_CAPTURE_SIZES}]}") + +# ---- The rest is the same in every DeepEPM2NAFDConnector recipe. ---- +export VLLM_USE_V2_MODEL_RUNNER=1 +ATTN_TP=${ATTN_TP:-1} +AFD_CONNECTOR=${AFD_CONNECTOR:-DeepEPM2NAFDConnector} GATE_ON_ATTENTION=true +# P2P routes on the FFN and takes no connector_extra_config. +[[ $AFD_CONNECTOR == P2pNcclAFDConnector ]] && GATE_ON_ATTENTION=false CONNECTOR_CONFIG= +[[ -n $CONNECTOR_CONFIG ]] && CONNECTOR_CONFIG=", \"connector_extra_config\": {${CONNECTOR_CONFIG}}" +IFS=, read -ra PODS <<<"$LAYOUT" +K=${POD:-AFD_0} && K=${K#AFD_} +if [[ ! $LAYOUT =~ ^[0-9]+A[0-9]+F(,[0-9]+A[0-9]+F)*$ || ! $K =~ ^[0-9]+$ ]] || ((K >= ${#PODS[@]})); then + echo "need LAYOUT=AF[,...] and POD=AFD_ in it: LAYOUT=$LAYOUT POD=${POD:-}" >&2 + exit 1 +fi +NUM_A=() NUM_F=() +for p in "${PODS[@]}"; do NUM_A+=("${p%%A*}") && p=${p#*A} && NUM_F+=("${p%F}"); done +((${#PODS[@]} > 1)) || HEAD_HOST=${HEAD_HOST:-127.0.0.1} + +# Extra `vllm serve` flags go last, so they override the defaults. +read -ra ARGS <<<"${VLLM_ARGS:-}" && ARGS+=("$@") +read -ra ATTN_EXTRA <<<"${ATTN_ARGS:-}" +read -ra FFN_EXTRA <<<"${FFN_ARGS:-}" +if [[ -n ${DECODE_BENCH:-} ]]; then # tools/benchmarks is imported from the repository + ATTN_EXTRA=(--kv-transfer-config "{\"kv_connector\": \"AFDDecodeBenchConnector\", \"kv_connector_module_path\": \"tools.benchmarks.decode_bench\", \"kv_role\": \"kv_both\", \"kv_connector_extra_config\": ${DECODE_BENCH}}" "${ATTN_EXTRA[@]}") + export PYTHONPATH=${AFD_REPO:-$PWD}${PYTHONPATH:+:$PYTHONPATH} +fi + +# role_dp TP DP_HOST DP_RPC_PORT ... -> TOTAL and this pod's DP +# flags. A role spread over pods is led by its first pod; the others run +# headless from their start rank (the leader must omit it, or vLLM selects +# hybrid load balancing and rejects the headless peers). +role_dp() { + local tp=$1 host=$2 port=$3 i pods=0 start=0 lead=-1 + shift 3 + local n=("$@") + TOTAL=0 + for i in "${!n[@]}"; do + ((n[i])) || continue + ((n[i] % tp)) && { echo "TP $tp does not divide ${n[i]} ranks on pod $i" >&2; exit 1; } + ((lead < 0)) && lead=$i + ((i < K)) && start=$((start + n[i])) + pods=$((pods + 1)) TOTAL=$((TOTAL + n[i])) + done + DP=(--tensor-parallel-size "$tp" --data-parallel-size $((TOTAL / tp))) + ((pods > 1)) && DP+=(--data-parallel-size-local $((n[K] / tp)) + --data-parallel-address "$host" --data-parallel-rpc-port "$port") + ((pods > 1 && K != lead)) && DP+=(--headless --data-parallel-start-rank $((start / tp))) +} +role_dp 1 "${FFN_HEAD_HOST:-${HEAD_HOST:-vllm-ffn-dp-service}}" 13355 "${NUM_F[@]}" +FFN_DP=("${DP[@]}") NUM_FFN=$TOTAL +role_dp "$ATTN_TP" "${ATTN_HEAD_HOST:-${HEAD_HOST:-vllm-attn-dp-service}}" 13345 "${NUM_A[@]}" +ATTN_DP=("${DP[@]}") NUM_ATTN=$TOTAL + +afd() { # The AFD rendezvous runs on the pod that holds FFN rank 0. + cat < ffn.log 2>&1 & +fi +if ((A)); then + CUDA_VISIBLE_DEVICES=$(seq "$F" $((F + A - 1)) | paste -sd, -) vllm serve "$MODEL_PATH" \ + --additional-config "$(afd attention)" "${COMMON[@]}" "${ATTN_DP[@]}" \ + "${ARGS[@]}" "${ATTN_EXTRA[@]}" --port 18000 > attn.log 2>&1 & +fi +wait diff --git a/recipe/gpu/P2pNcclAFDConnector/qwen3_235b_a22b_fp8/README.md b/recipe/gpu/P2pNcclAFDConnector/qwen3_235b_a22b_fp8/README.md index 087aa0ad8..c89af5fd7 100644 --- a/recipe/gpu/P2pNcclAFDConnector/qwen3_235b_a22b_fp8/README.md +++ b/recipe/gpu/P2pNcclAFDConnector/qwen3_235b_a22b_fp8/README.md @@ -1,10 +1,13 @@ # Qwen3-235B-A22B-FP8 AFD Recipe (P2pNcclAFDConnector) `prefill_decode_colocation/4a2f_graph_dp4tp1.sh` runs the shared -Qwen3-235B-A22B-FP8 4A2F recipe with `AFD_CONNECTOR=P2pNcclAFDConnector`. -Topology, environment knobs, and the decode-only benchmark are documented with -the script itself in -[`../../DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/README.md`](../../DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/README.md). +Qwen3-235B-A22B-FP8 graph recipe +(`DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/graph.sh`) +in the 4A2F layout (`LAYOUT=4A2F`, one 6-GPU node) with +`AFD_CONNECTOR=P2pNcclAFDConnector`. Its layouts, settings and the +decode-only benchmark are documented in the +[`DeepEPM2NAFDConnector` recipes README](../../DeepEPM2NAFDConnector/README.md) +and the [model README](../../DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/README.md). ```bash export MODEL_PATH=/path/model_weights/Qwen3-235B-A22B-FP8 diff --git a/recipe/gpu/P2pNcclAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/4a2f_graph_dp4tp1.sh b/recipe/gpu/P2pNcclAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/4a2f_graph_dp4tp1.sh index f5e183f4c..a494f68f1 100755 --- a/recipe/gpu/P2pNcclAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/4a2f_graph_dp4tp1.sh +++ b/recipe/gpu/P2pNcclAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/4a2f_graph_dp4tp1.sh @@ -2,11 +2,12 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project # -# P2pNcclAFDConnector entry point of the shared Qwen3-235B-A22B-FP8 4A2F recipe. -# The topology and every environment knob are documented in -# recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/4a2f_graph_dp4tp1.sh. +# P2pNcclAFDConnector entry point of the shared Qwen3-235B-A22B-FP8 graph +# recipe in the 4A2F layout (one 6-GPU node), +# recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/graph.sh, +# which documents the layout and settings. set -euo pipefail SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" -AFD_CONNECTOR=P2pNcclAFDConnector exec bash \ - "${SCRIPT_DIR}/../../../DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/4a2f_graph_dp4tp1.sh" +LAYOUT=${LAYOUT:-4A2F} AFD_CONNECTOR=P2pNcclAFDConnector exec bash \ + "${SCRIPT_DIR}/../../../DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/graph.sh" "$@" diff --git a/recipe/gpu/P2pNcclAFDConnector/qwen3_30b_a3b/README.md b/recipe/gpu/P2pNcclAFDConnector/qwen3_30b_a3b/README.md index cb715b288..41ea965bb 100644 --- a/recipe/gpu/P2pNcclAFDConnector/qwen3_30b_a3b/README.md +++ b/recipe/gpu/P2pNcclAFDConnector/qwen3_30b_a3b/README.md @@ -1,9 +1,11 @@ # Qwen3-30B-A3B AFD Recipe (P2pNcclAFDConnector) `prefill_decode_colocation/2a2f_graph_dp2tp1.sh` runs the shared Qwen3-30B-A3B -2A2F recipe with `AFD_CONNECTOR=P2pNcclAFDConnector`. Topology, environment -knobs, and the decode-only benchmark are documented with the script itself in -[`../../DeepEPM2NAFDConnector/qwen3_30b_a3b/README.md`](../../DeepEPM2NAFDConnector/qwen3_30b_a3b/README.md). +graph recipe (`DeepEPM2NAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/graph.sh`, +default layout 2A2F) with `AFD_CONNECTOR=P2pNcclAFDConnector`. Its layouts, +settings and the decode-only benchmark are documented in the +[`DeepEPM2NAFDConnector` recipes README](../../DeepEPM2NAFDConnector/README.md) +and the [model README](../../DeepEPM2NAFDConnector/qwen3_30b_a3b/README.md). ```bash export MODEL_PATH=/path/model_weights/Qwen3-30B-A3B diff --git a/recipe/gpu/P2pNcclAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/2a2f_graph_dp2tp1.sh b/recipe/gpu/P2pNcclAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/2a2f_graph_dp2tp1.sh index 18ba6bff1..f7e63fa87 100755 --- a/recipe/gpu/P2pNcclAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/2a2f_graph_dp2tp1.sh +++ b/recipe/gpu/P2pNcclAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/2a2f_graph_dp2tp1.sh @@ -2,11 +2,11 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project # -# P2pNcclAFDConnector entry point of the shared Qwen3-30B-A3B 2A2F recipe. -# The topology and every environment knob are documented in -# recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/2a2f_graph_dp2tp1.sh. +# P2pNcclAFDConnector entry point of the shared Qwen3-30B-A3B graph recipe, +# recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/graph.sh, +# which documents the layout and settings. set -euo pipefail SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" AFD_CONNECTOR=P2pNcclAFDConnector exec bash \ - "${SCRIPT_DIR}/../../../DeepEPM2NAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/2a2f_graph_dp2tp1.sh" + "${SCRIPT_DIR}/../../../DeepEPM2NAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/graph.sh" "$@" diff --git a/recipe/gpu/baseline/qwen3_235b_a22b_fp8/vanilla_tp4.sh b/recipe/gpu/baseline/qwen3_235b_a22b_fp8/vanilla_tp4.sh index 7adba7f33..deda834a9 100755 --- a/recipe/gpu/baseline/qwen3_235b_a22b_fp8/vanilla_tp4.sh +++ b/recipe/gpu/baseline/qwen3_235b_a22b_fp8/vanilla_tp4.sh @@ -5,7 +5,7 @@ # Qwen3-235B-A22B-FP8 vanilla vLLM baseline on one 4-GPU node. # # This is the no-AFD counterpart to -# recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/4a2f_graph_dp4tp1.sh. +# recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/graph.sh. # It launches a single vLLM process with tensor parallelism 4 and no AFD # connector, AFD additional config, or separate Attention/FFN roles. # diff --git a/tools/benchmarks/README.md b/tools/benchmarks/README.md index f08d7cfca..c217ea3e1 100644 --- a/tools/benchmarks/README.md +++ b/tools/benchmarks/README.md @@ -1,70 +1,67 @@ # Decode-only benchmark This directory contains the decode-only benchmark used to measure decode -throughput and latency over long, real KV contexts while avoiding large -prefill work during the measured phase. +throughput and latency over long, real KV contexts while avoiding prefill +work during the measured phase. -- `decode_bench.py` implements the server-side KV connector. -- `decode_bench_client.py` wraps `vllm bench serve`, performs warmup, and writes - the benchmark checks. +- `decode_bench.py` implements the server-side KV connector, + `AFDDecodeBenchConnector`. +- `decode_bench_client.py` wraps `vllm bench serve`, performs warmup, and + checks that the measured run was decode-only. ## How it works -The benchmark starts by bringing up a vLLM serving instance using one of the -recipes under `recipe/gpu/`. The recipe attaches the decode-bench KV connector -to the Attention instance and starts the server. - -After the server is ready, the benchmark has two phases: - -1. **Warmup.** The server-side connector creates a pool of seeded prompts of - `--prompt-len` tokens and prefills them once. Their KV blocks are pinned in - the prefix cache. -2. **Measured.** The client sends requests that the connector remaps onto the - pool prompts. Each measured request prefills only its last prompt token and - then decodes for the requested number of output tokens. - -This makes the benchmark decode-only while still exercising real KV-cache -attention over the full prompt length. +The connector runs in its prompt-pool mode when `kv_connector_extra_config` +sets `num_prompts > 0`: + +1. **Warmup.** Each Attention DP rank turns its first `num_prompts` requests + into a pool of seeded random prompts of `prompt_len` tokens (the same pool + on every rank, whatever the requests carried) and prefills them for real, + one output token each. +2. **Pinning.** When a warmup request finishes, the connector takes a + reference of its own on the request's full prompt blocks and holds it for + the server's lifetime, while the request itself is freed as usual. Pinned + blocks cannot be evicted and stay in the prefix cache. +3. **Measured requests.** Every later request is remapped round-robin onto a + pinned pool prompt, whatever it asked for. The prefix cache covers every + prompt token but the last (the prompt length is rounded to one token past a + block boundary), so the request prefills one token and decodes over real + KV, on whichever rank serves it. A request that arrives before any prompt + is pinned may miss the cache; the connector logs it and the client's checks + fail the run. + +`kv_connector_extra_config` keys: `num_prompts` (pool size per Attention DP +rank, `0` keeps the dummy-KV mode below), `prompt_len` (default 1025), `seed` +(default 0) and `max_output_len` (the longest measured output; it sizes the +KV check and defaults to `max_model_len - prompt_len`). ## Running with an AFD recipe -The AFD recipes already support the connector through environment variables. -For example, to run the Qwen3-235B-A22B-FP8 4A2F recipe: +The `DeepEPM2NAFDConnector` recipes attach the connector to Attention when +`DECODE_BENCH` holds its `kv_connector_extra_config`; see the +[recipes README](../../recipe/gpu/DeepEPM2NAFDConnector/README.md#decode-only-benchmark). +Every layout works, including multi-pod ones. From the repository root: ```bash -export MODEL_PATH=/path/to/Qwen3-235B-A22B-FP8 - -KV_CONNECTOR=decode_bench \ -DECODE_BENCH_NUM_PROMPTS=500 \ -DECODE_BENCH_PROMPT_LEN=1025 \ -DECODE_BENCH_MAX_OUTPUT_LEN=256 \ - bash recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/4a2f_graph_dp4tp1.sh -``` - -Then run the client from the repository root: +export MODEL_PATH=/path/to/Qwen3-30B-A3B +DECODE_BENCH='{"num_prompts": 500, "prompt_len": 1025, "max_output_len": 256}' \ + bash recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/graph.sh -```bash python tools/benchmarks/decode_bench_client.py \ - --port 18000 \ - --dp-size 4 \ - --pool-prompts 500 \ - --prompt-len 1025 \ + --port 18000 --dp-size 2 --pool-prompts 500 --prompt-len 1025 \ --result-dir /tmp/results/decode_bench \ - -- --num-prompts 3000 \ - --random-output-len 256 \ - --max-concurrency 128 \ - --request-rate inf + -- --num-prompts 2000 --random-output-len 256 \ + --max-concurrency 64 --request-rate inf ``` -Use the Attention port, not the FFN port. +Use the Attention port, not the FFN port. `--dp-size` is the Attention DP +size over all pods. ## Running with vanilla vLLM The client is not AFD-specific. It works with any vLLM server that has the -decode-bench KV connector attached. - -For a single vanilla vLLM process on four GPUs, start the server with tensor -parallelism 4 and the connector’s `--kv-transfer-config`: +connector attached. For a single vanilla vLLM process on four GPUs, start the +server with tensor parallelism 4 and the connector's `--kv-transfer-config`: ```bash export PYTHONPATH=$PWD @@ -76,74 +73,80 @@ vllm serve "$MODEL_PATH" \ --kv-transfer-config '{"kv_connector":"AFDDecodeBenchConnector","kv_connector_module_path":"tools.benchmarks.decode_bench","kv_role":"kv_both","kv_connector_extra_config":{"num_prompts":500,"prompt_len":1025,"seed":0,"max_output_len":256}}' ``` -Then run the benchmark client with `--dp-size 1`, because a vanilla TP4 server -has one logical engine: +Then run the client with `--dp-size 1`, because a vanilla TP4 server has one +logical engine: ```bash python tools/benchmarks/decode_bench_client.py \ - --port 8000 \ - --dp-size 1 \ - --pool-prompts 500 \ - --prompt-len 1025 \ + --port 8000 --dp-size 1 --pool-prompts 500 --prompt-len 1025 \ --result-dir /tmp/results/decode_bench_vanilla \ - -- --num-prompts 3000 \ - --random-output-len 256 \ - --max-concurrency 256 \ - --request-rate inf + -- --num-prompts 3000 --random-output-len 256 \ + --max-concurrency 256 --request-rate inf ``` -To repeat the run at concurrency 512, rerun the same command with -`--max-concurrency 512` and a different `--result-dir`. - -## Pool sizing - -The connector checks that the prompt pool and running requests fit in the KV -cache: - -```text -num_prompts * ceil((prompt_len + 1) / block_size) - + max_num_seqs * ceil((max_output_len + 1) / block_size) - <= free KV blocks -``` - -Increase `num_prompts` only if the server still has enough free KV blocks after -reserving space for the configured `max_num_seqs` and `max_output_len`. +## Client phases + +1. **Warmup.** One `vllm bench serve` run per Attention rank, pinned with the + `X-data-parallel-rank` header, sends `--pool-prompts` requests to that rank. + `--warmup-concurrency` sets their concurrency; `--warmup-bench-args` (last + before `--`) passes anything else to these runs. +2. **Pre-check.** From the server's prefix-cache counters, every warmup + request must have looked up `--prompt-len` prompt tokens, so the pool + prompts are as long as the measured prompts. +3. **Measured run.** `vllm bench serve --dataset-name random --random-input-len + ` plus every argument after `--`: request rate, concurrency, + number of requests, output length, burstiness, percentiles, and so on. The + client sets the dataset, input length, endpoint, and result file itself and + rejects them after `--`. +4. **Post-check.** Over the measured run the server must have prefilled one + token per finished request (decode only), with no failed requests and no + preemptions. Otherwise the client exits non-zero. ## Client options -- `--dp-size` is the number of Attention data-parallel ranks. Use `4` for the - AFD 4A2F recipe and `1` for a single vanilla vLLM engine. -- `--pool-prompts` must match the server’s `num_prompts`. -- `--prompt-len` must match the server’s `prompt_len`. +- `--dp-size` is the number of Attention data-parallel ranks. +- `--pool-prompts` must match the server's `num_prompts`, and `--prompt-len` + its `prompt_len`; `--random-output-len` must not exceed `max_output_len`. - `--result-dir` receives the warmup and measured benchmark JSON. - `--skip-warmup` measures only, against an already warmed pool. -- Every argument after `--` is passed unchanged to `vllm bench serve`. +- Every argument after `--` is passed unchanged to `vllm bench serve`, for + example `--num-prompts`, `--random-output-len`, `--max-concurrency`, + `--request-rate`, `--percentile-metrics`, `--metric-percentiles` and + `--save-detailed`. -Useful `vllm bench serve` arguments include: - -- `--num-prompts` -- `--random-output-len` -- `--max-concurrency` -- `--request-rate` -- `--percentile-metrics` -- `--metric-percentiles` -- `--save-detailed` +The client needs vLLM and the model's tokenizer installed, as vllm bench does. ## Output files -Each run writes: +Each run writes `warmup-rank.json` and `warmup-rank.log`, +`measured.json`, and `decode_bench_report.json`, which holds the client's +checks, including whether every measured request prefilled exactly one token +and whether any requests were preempted. -- `warmup-rank.json` and `warmup-rank.log` -- `measured.json` -- `decode_bench_report.json` +## Pool sizing -`decode_bench_report.json` contains the connector’s checks, including whether -every measured request prefilled exactly one token and whether any requests -were preempted. +The pool must fit each Attention rank's KV cache next to the running +requests. At startup the connector checks, per rank, -## Notes +```text +num_prompts * ceil((prompt_len + 1) / block_size) + + max_num_seqs * ceil((max_output_len + 1) / block_size) <= free KV blocks +``` -- The generated text is meaningless because the prompts are random. -- Restart the server to release the pinned pool. -- The connector’s name includes `AF`, but it does not require the AFD plugin; - it uses vLLM’s standard KV-transfer interface. +and fails with the largest `num_prompts` that fits. Keep `num_prompts` at or +above `--max-num-seqs` (the connector warns otherwise), or concurrent requests +share physical KV blocks and attention reads look cheaper than they are. The +recipe READMEs give each model's KV size per token. + +## Limitations + +- Generated text is meaningless: prompts are random tokens, and the client's + prompt is discarded. +- Measured requests decode at a context of `prompt_len` regardless of the + client's prompt length. +- The pool stays pinned for the server's lifetime; restart the server to free + it. +- `num_prompts: 0` keeps the older dummy-KV mode, which writes no KV at all: + every prompt token but the last is reported as externally computed. +- The connector's name includes `AF`, but it does not require the AFD plugin; + it uses vLLM's standard KV-transfer interface. From 974f7805f766db9fdba239b5b5757ba4fc79fe63 Mon Sep 17 00:00:00 2001 From: ronenkat <16743404+ronenkat@users.noreply.github.com> Date: Tue, 10 Nov 2026 14:41:10 +0200 Subject: [PATCH 31/32] revert(benchmarks): restore the dummy-KV decode-bench connector Drop the prompt-pool mode from AFDDecodeBenchConnector and return tools/benchmarks/decode_bench.py to main's behavior: every prompt token but the last is reported as externally computed, so requests skip prefill. The only difference from main is an explicit TypeAlias on KVCacheLayer so mypy passes without torch installed. Remove decode_bench_client.py, whose warmup and prefix-cache checks only apply to the prompt pool. Decode-only runs now start a recipe with DECODE_BENCH='{}' and drive it with request_generator.sh, whose default port follows the recipes' Attention port (18000). Rewrite tools/benchmarks/README.md for the dummy-KV mode and update the recipe READMEs and the deploy/bench skills that referenced the pool. Co-Authored-By: Claude Opus 5.5 Signed-off-by: ronenkat <16743404+ronenkat@users.noreply.github.com> --- .agents/skills/deploy-afd-k8s/SKILL.md | 4 +- .../skills/deploy-afd-k8s/resources/deploy.md | 2 +- .agents/skills/run-vllm-bench-k8s/SKILL.md | 5 +- recipe/gpu/DeepEPM2NAFDConnector/README.md | 31 +- .../qwen3_235b_a22b_fp8/README.md | 8 +- tools/benchmarks/README.md | 162 ++---- tools/benchmarks/decode_bench.py | 394 +------------- tools/benchmarks/decode_bench_client.py | 485 ------------------ tools/benchmarks/request_generator.sh | 12 +- 9 files changed, 81 insertions(+), 1022 deletions(-) delete mode 100644 tools/benchmarks/decode_bench_client.py diff --git a/.agents/skills/deploy-afd-k8s/SKILL.md b/.agents/skills/deploy-afd-k8s/SKILL.md index 4d5df550a..e7ea5fe0c 100644 --- a/.agents/skills/deploy-afd-k8s/SKILL.md +++ b/.agents/skills/deploy-afd-k8s/SKILL.md @@ -49,8 +49,8 @@ than one pod) -- callers never talk to these directly. See - `DeepEPM2NAFDConnector` recipes take their placement from `LAYOUT` (one pod per `AF` entry) and run-specific settings as env (`VLLM_ARGS`, `DECODE_BENCH`, ...); see [resolve-recipe.md](resources/resolve-recipe.md). - This skill only serves; the decode-only benchmark's client - (`tools/benchmarks/decode_bench_client.py`) is run separately against the + This skill only serves; load for the decode-only benchmark + (`tools/benchmarks/request_generator.sh`) is run separately against the endpoint. ## Requirements diff --git a/.agents/skills/deploy-afd-k8s/resources/deploy.md b/.agents/skills/deploy-afd-k8s/resources/deploy.md index 83ad1fd6c..58643cbbe 100644 --- a/.agents/skills/deploy-afd-k8s/resources/deploy.md +++ b/.agents/skills/deploy-afd-k8s/resources/deploy.md @@ -161,7 +161,7 @@ deploy_pod vllm-pod-ffn-0 2 none head "POD=FFN_0" decode-only benchmark, or `LAYOUT=4A1F,4A1F` (5 GPUs per pod): ```bash -SETTINGS=$'LAYOUT=2A0F,0A2F\nDECODE_BENCH={"num_prompts": 500, "max_output_len": 256}\nAFD_REPO=/opt/afd-plugin' +SETTINGS=$'LAYOUT=2A0F,0A2F\nDECODE_BENCH={}\nAFD_REPO=/opt/afd-plugin' deploy_pod vllm-pod-afd-0 2 head none $'POD=AFD_0\n'"$SETTINGS" deploy_pod vllm-pod-afd-1 2 none head $'POD=AFD_1\n'"$SETTINGS" diff --git a/.agents/skills/run-vllm-bench-k8s/SKILL.md b/.agents/skills/run-vllm-bench-k8s/SKILL.md index ac8c50087..b61c4f99f 100644 --- a/.agents/skills/run-vllm-bench-k8s/SKILL.md +++ b/.agents/skills/run-vllm-bench-k8s/SKILL.md @@ -15,9 +15,8 @@ One invocation = one run = one `RUN_ID` = one report. ## Scope - GPU colocation recipes only -- whatever `deploy-afd-k8s` supports. -- Not for local benchmarking (run `request_generator.sh` or - `decode_bench_client.py` directly against a local recipe for that) or - correctness E2E (`run-e2e`). +- Not for local benchmarking (run `request_generator.sh` directly against + a local recipe for that) or correctness E2E (`run-e2e`). - Requires everything `deploy-afd-k8s` requires. No extra image: the bench pod runs `AFD_PLUGIN_IMAGE`, which already carries the repo at `/opt/afd-plugin` and vLLM's `bench` CLI. diff --git a/recipe/gpu/DeepEPM2NAFDConnector/README.md b/recipe/gpu/DeepEPM2NAFDConnector/README.md index 07b4c254c..29f8ef848 100644 --- a/recipe/gpu/DeepEPM2NAFDConnector/README.md +++ b/recipe/gpu/DeepEPM2NAFDConnector/README.md @@ -185,33 +185,24 @@ The answer should name Paris. For an accuracy check, run the matching E2E case `DECODE_BENCH` attaches the decode-only benchmark's KV connector (`AFDDecodeBenchConnector`) to Attention and takes its -`kv_connector_extra_config` as JSON: every Attention DP rank prefills a pool -of `num_prompts` random prompts of `prompt_len` tokens, pins their KV, and -remaps every later request onto the pool, so measured requests decode over -real KV without prefill. Set `max_output_len` to the longest output you -measure; it sizes the startup KV check. The connector lives in -`tools/benchmarks`, which is not in the wheel: run from the repository root or -set `AFD_REPO`. +`kv_connector_extra_config` as JSON (`{}` is enough). The connector reports +every prompt token but the last as externally computed, so requests skip +prefill and decode at a context of the full prompt length over unwritten KV. +The connector lives in `tools/benchmarks`, which is not in the wheel: run from +the repository root or set `AFD_REPO`. ```bash export MODEL_PATH=/path/model_weights/Qwen3-30B-A3B -DECODE_BENCH='{"num_prompts": 500, "prompt_len": 1025, "max_output_len": 256}' \ +DECODE_BENCH='{}' \ bash recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/graph.sh -python tools/benchmarks/decode_bench_client.py \ - --port 18000 --dp-size 2 --pool-prompts 500 --prompt-len 1025 \ - --result-dir /tmp/results/decode_bench \ - -- --num-prompts 2000 --random-output-len 256 \ - --max-concurrency 64 --request-rate inf +INPUT_LEN=1024 OUTPUT_LEN=256 NUM_PROMPTS=2000 \ + MAX_CONCURRENCY=64 REQUEST_RATE=inf RESULT_DIR=/tmp/results/decode_bench \ + tools/benchmarks/request_generator.sh ``` -`--dp-size` is the Attention DP size over all pods (total Attention ranks -divided by `ATTN_TP`; 2 for `2A2F`, `2A0F,0A2F` and `1A1F,1A1F`), and -`--pool-prompts` and `--prompt-len` must match `DECODE_BENCH`. The benchmark -works in every layout and mode. How the pool, the client and its checks work, -and how to size the pool, are described in -[`tools/benchmarks/README.md`](../../../tools/benchmarks/README.md); the model -READMEs give each model's KV size per token. +Drive the Attention port (18000). The benchmark works in every layout and +mode; see [`tools/benchmarks/README.md`](../../../tools/benchmarks/README.md). ## Configuration notes diff --git a/recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/README.md b/recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/README.md index 38255d164..4c94ac5ac 100644 --- a/recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/README.md +++ b/recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/README.md @@ -44,12 +44,8 @@ Common layouts: every running request. - The baseline without AFD is [`recipe/gpu/baseline/qwen3_235b_a22b_fp8/vanilla_tp4.sh`](../../baseline/qwen3_235b_a22b_fp8/vanilla_tp4.sh). -- Decode-only benchmark pool sizing: Qwen3-235B holds 188 KiB of BF16 KV per - token (94 layers, 4 KV heads of 128). With - `{"num_prompts": 500, "prompt_len": 1025, "max_output_len": 256}` and 256 - sequences, each Attention rank needs - `500 * ceil(1026 / 16) + 256 * ceil(257 / 16)` blocks, about 590K tokens or - 106 GiB; an H200 Attention rank has room for about 640K tokens. +- Qwen3-235B holds 188 KiB of BF16 KV per token (94 layers, 4 KV heads of + 128); an H200 Attention rank has room for about 640K tokens. ## Limitations diff --git a/tools/benchmarks/README.md b/tools/benchmarks/README.md index c217ea3e1..c5f8d5627 100644 --- a/tools/benchmarks/README.md +++ b/tools/benchmarks/README.md @@ -1,152 +1,82 @@ # Decode-only benchmark -This directory contains the decode-only benchmark used to measure decode -throughput and latency over long, real KV contexts while avoiding prefill -work during the measured phase. +This directory contains the decode-only benchmark, which measures decode +throughput and latency at a chosen input length without prefill work. - `decode_bench.py` implements the server-side KV connector, `AFDDecodeBenchConnector`. -- `decode_bench_client.py` wraps `vllm bench serve`, performs warmup, and - checks that the measured run was decode-only. +- `request_generator.sh` wraps `vllm bench serve` to drive load against the + server. ## How it works -The connector runs in its prompt-pool mode when `kv_connector_extra_config` -sets `num_prompts > 0`: - -1. **Warmup.** Each Attention DP rank turns its first `num_prompts` requests - into a pool of seeded random prompts of `prompt_len` tokens (the same pool - on every rank, whatever the requests carried) and prefills them for real, - one output token each. -2. **Pinning.** When a warmup request finishes, the connector takes a - reference of its own on the request's full prompt blocks and holds it for - the server's lifetime, while the request itself is freed as usual. Pinned - blocks cannot be evicted and stay in the prefix cache. -3. **Measured requests.** Every later request is remapped round-robin onto a - pinned pool prompt, whatever it asked for. The prefix cache covers every - prompt token but the last (the prompt length is rounded to one token past a - block boundary), so the request prefills one token and decodes over real - KV, on whichever rank serves it. A request that arrives before any prompt - is pinned may miss the cache; the connector logs it and the client's checks - fail the run. - -`kv_connector_extra_config` keys: `num_prompts` (pool size per Attention DP -rank, `0` keeps the dummy-KV mode below), `prompt_len` (default 1025), `seed` -(default 0) and `max_output_len` (the longest measured output; it sizes the -KV check and defaults to `max_model_len - prompt_len`). +The connector emulates the decode side of prefill-decode disaggregation with +the prefill instance faked out. For every new request, the scheduler side +reports every prompt token except the last as externally computed, so the +request prefills one token and goes straight to decode at a context of the +full prompt length. No prefill instance, KV producer or proxy is needed. + +The connector writes no KV values: the blocks allocated for the "external" +tokens keep whatever the KV cache held before. Throughput and latency are +meaningful; generated text is not. A preempted request that restarts from +scratch is reported as external again. + +Attach the connector to Attention only; the FFN needs no +`--kv-transfer-config`. The connector lives in `tools/benchmarks`, which is +not in the wheel, so it is loaded through `kv_connector_module_path` and needs +the repository root on `PYTHONPATH`. ## Running with an AFD recipe -The `DeepEPM2NAFDConnector` recipes attach the connector to Attention when -`DECODE_BENCH` holds its `kv_connector_extra_config`; see the +The `DeepEPM2NAFDConnector` recipes, and the `P2pNcclAFDConnector` entry +points that reuse them, attach the connector to Attention when `DECODE_BENCH` +holds its `kv_connector_extra_config` as JSON (`{}` is enough) and add the +repository root (`AFD_REPO`, default the current directory) to `PYTHONPATH`; +see the [recipes README](../../recipe/gpu/DeepEPM2NAFDConnector/README.md#decode-only-benchmark). Every layout works, including multi-pod ones. From the repository root: ```bash export MODEL_PATH=/path/to/Qwen3-30B-A3B -DECODE_BENCH='{"num_prompts": 500, "prompt_len": 1025, "max_output_len": 256}' \ +DECODE_BENCH='{}' \ bash recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/graph.sh -python tools/benchmarks/decode_bench_client.py \ - --port 18000 --dp-size 2 --pool-prompts 500 --prompt-len 1025 \ - --result-dir /tmp/results/decode_bench \ - -- --num-prompts 2000 --random-output-len 256 \ - --max-concurrency 64 --request-rate inf +# once Attention answers on port 18000: +INPUT_LEN=1024 OUTPUT_LEN=256 NUM_PROMPTS=2000 MAX_CONCURRENCY=64 \ + REQUEST_RATE=inf RESULT_DIR=/tmp/results/decode_bench \ + tools/benchmarks/request_generator.sh ``` -Use the Attention port, not the FFN port. `--dp-size` is the Attention DP -size over all pods. +Use the Attention port, not the FFN port. `INPUT_LEN` plus `OUTPUT_LEN` must +fit the recipe's `--max-model-len`. ## Running with vanilla vLLM -The client is not AFD-specific. It works with any vLLM server that has the -connector attached. For a single vanilla vLLM process on four GPUs, start the -server with tensor parallelism 4 and the connector's `--kv-transfer-config`: +The connector is not AFD-specific. For a single vanilla vLLM process, pass +the same `--kv-transfer-config`: ```bash export PYTHONPATH=$PWD -export MODEL_PATH=/path/to/model - vllm serve "$MODEL_PATH" \ --tensor-parallel-size 4 \ - --trust-remote-code \ - --kv-transfer-config '{"kv_connector":"AFDDecodeBenchConnector","kv_connector_module_path":"tools.benchmarks.decode_bench","kv_role":"kv_both","kv_connector_extra_config":{"num_prompts":500,"prompt_len":1025,"seed":0,"max_output_len":256}}' + --kv-transfer-config '{"kv_connector":"AFDDecodeBenchConnector","kv_connector_module_path":"tools.benchmarks.decode_bench","kv_role":"kv_both","kv_connector_extra_config":{}}' ``` -Then run the client with `--dp-size 1`, because a vanilla TP4 server has one -logical engine: +and run `request_generator.sh` with `PORT=8000`. -```bash -python tools/benchmarks/decode_bench_client.py \ - --port 8000 --dp-size 1 --pool-prompts 500 --prompt-len 1025 \ - --result-dir /tmp/results/decode_bench_vanilla \ - -- --num-prompts 3000 --random-output-len 256 \ - --max-concurrency 256 --request-rate inf -``` - -## Client phases - -1. **Warmup.** One `vllm bench serve` run per Attention rank, pinned with the - `X-data-parallel-rank` header, sends `--pool-prompts` requests to that rank. - `--warmup-concurrency` sets their concurrency; `--warmup-bench-args` (last - before `--`) passes anything else to these runs. -2. **Pre-check.** From the server's prefix-cache counters, every warmup - request must have looked up `--prompt-len` prompt tokens, so the pool - prompts are as long as the measured prompts. -3. **Measured run.** `vllm bench serve --dataset-name random --random-input-len - ` plus every argument after `--`: request rate, concurrency, - number of requests, output length, burstiness, percentiles, and so on. The - client sets the dataset, input length, endpoint, and result file itself and - rejects them after `--`. -4. **Post-check.** Over the measured run the server must have prefilled one - token per finished request (decode only), with no failed requests and no - preemptions. Otherwise the client exits non-zero. - -## Client options - -- `--dp-size` is the number of Attention data-parallel ranks. -- `--pool-prompts` must match the server's `num_prompts`, and `--prompt-len` - its `prompt_len`; `--random-output-len` must not exceed `max_output_len`. -- `--result-dir` receives the warmup and measured benchmark JSON. -- `--skip-warmup` measures only, against an already warmed pool. -- Every argument after `--` is passed unchanged to `vllm bench serve`, for - example `--num-prompts`, `--random-output-len`, `--max-concurrency`, - `--request-rate`, `--percentile-metrics`, `--metric-percentiles` and - `--save-detailed`. - -The client needs vLLM and the model's tokenizer installed, as vllm bench does. - -## Output files - -Each run writes `warmup-rank.json` and `warmup-rank.log`, -`measured.json`, and `decode_bench_report.json`, which holds the client's -checks, including whether every measured request prefilled exactly one token -and whether any requests were preempted. - -## Pool sizing - -The pool must fit each Attention rank's KV cache next to the running -requests. At startup the connector checks, per rank, - -```text -num_prompts * ceil((prompt_len + 1) / block_size) - + max_num_seqs * ceil((max_output_len + 1) / block_size) <= free KV blocks -``` +## Checking that a run was decode-only -and fails with the largest `num_prompts` that fits. Keep `num_prompts` at or -above `--max-num-seqs` (the connector warns otherwise), or concurrent requests -share physical KV blocks and attention reads look cheaper than they are. The -recipe READMEs give each model's KV size per token. +The server's `/metrics` split prompt tokens by source. Over a decode-only +run, `vllm:prompt_tokens_by_source{source="local_compute"}` grows by exactly +one per finished request (`vllm:request_success`). The rest of each prompt is +counted under `source="external_kv_transfer"`, or `source="local_cache_hit"` +for the few prompt blocks the prefix cache already holds. ## Limitations -- Generated text is meaningless: prompts are random tokens, and the client's - prompt is discarded. -- Measured requests decode at a context of `prompt_len` regardless of the - client's prompt length. -- The pool stays pinned for the server's lifetime; restart the server to free - it. -- `num_prompts: 0` keeps the older dummy-KV mode, which writes no KV at all: - every prompt token but the last is reported as externally computed. -- The connector's name includes `AF`, but it does not require the AFD plugin; +- KV contents are uninitialized, so attention reads arbitrary values and the + generated text is meaningless. +- `kv_connector_extra_config` accepts `fill_mean` and `fill_std`, but the + worker never writes them into the KV cache (`start_load_kv` is a no-op). +- The connector's name includes `AFD`, but it does not require the AFD plugin; it uses vLLM's standard KV-transfer interface. diff --git a/tools/benchmarks/decode_bench.py b/tools/benchmarks/decode_bench.py index 3f00d839f..f2e18c91f 100644 --- a/tools/benchmarks/decode_bench.py +++ b/tools/benchmarks/decode_bench.py @@ -1,22 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project -"""Decode benchmark KV connector for AFD decode-side pressure tests. - -Two modes, chosen by ``kv_connector_extra_config``: - -- Dummy KV (``num_prompts`` unset or 0): every prompt token but the last is - reported as externally computed and never written, so requests skip prefill. -- Prompt pool (``num_prompts`` > 0): each Attention DP rank prefills a pool of - ``num_prompts`` seeded random prompts for real, from its first - ``num_prompts`` requests (the warmup), and pins their KV blocks in the prefix - cache. Every later request is remapped, whatever its content, onto a pinned - pool prompt round-robin, hits the prefix cache for all but the last prompt - token, and goes straight to decode over real KV. -""" +"""Decode benchmark KV connector for AFD decode-side pressure tests.""" from __future__ import annotations -from collections import deque from dataclasses import dataclass from typing import TYPE_CHECKING, Any, TypeAlias @@ -34,51 +21,15 @@ if TYPE_CHECKING: from vllm.config import VllmConfig from vllm.forward_context import ForwardContext - from vllm.v1.core.block_pool import BlockPool - from vllm.v1.core.kv_cache_manager import KVCacheBlocks, KVCacheManager - from vllm.v1.core.kv_cache_utils import BlockHash + from vllm.v1.core.kv_cache_manager import KVCacheBlocks from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.kv_cache_interface import KVCacheConfig from vllm.v1.request import Request -# vLLM configures handlers only for the "vllm" logger tree; a logger outside it -# falls back to the root WARNING level and drops the prompt-pool INFO lines. -logger = init_logger(f"vllm.{__name__}") +logger = init_logger(__name__) KVCacheLayer: TypeAlias = torch.Tensor | tuple[torch.Tensor, ...] -# A warmup request only has to prefill its pool prompt. -WARMUP_OUTPUT_TOKENS = 1 -# Finished measured requests between two prompt-pool hit/miss log lines. -POOL_STATS_LOG_INTERVAL = 1000 -# Warmup requests that end this way have fully prefilled their pool prompt. -PINNABLE_WARMUP_STATUSES = ( - RequestStatus.FINISHED_STOPPED, - RequestStatus.FINISHED_LENGTH_CAPPED, -) - - -@dataclass(frozen=True) -class PromptPoolConfig: - """Prompt-pool keys of ``kv_connector_extra_config``. - - Attributes: - num_prompts: Pool entries per Attention DP rank; 0 keeps the dummy-KV - mode. - prompt_len: Requested pool prompt length. It is rounded down to one - token past a block boundary, where a prefix-cache hit leaves exactly - the last token to prefill. - seed: Seed of the pool tokens, shared by every rank so all ranks hold - the same pool. - max_output_len: Longest output a measured request decodes; sizes the KV - capacity check. ``None`` uses ``max_model_len - prompt_len``. - """ - - num_prompts: int = 0 - prompt_len: int = 1025 - seed: int = 0 - max_output_len: int | None = None - @dataclass class AFDDecodeBenchConnectorMetadata(KVConnectorMetadata): @@ -91,10 +42,9 @@ class AFDDecodeBenchConnectorMetadata(KVConnectorMetadata): class AFDDecodeBenchConnector(KVConnectorBase_V1, SupportsHMA): """KV connector used to benchmark AFD decode instances with long ISL. - In the dummy-KV mode the connector emulates remote prefill by marking all - prompt tokens except the final decode token as externally provided. In the - prompt-pool mode only a prefix-cache miss is emulated that way; see the - module docstring. + The connector emulates remote prefill by marking all prompt tokens except + the final decode token as externally provided, then fills the allocated KV + cache pages with deterministic or random non-zero values on the worker. """ def __init__( @@ -109,10 +59,7 @@ def __init__( self.connector_worker: AFDDecodeBenchConnectorWorker | None = None if role == KVConnectorRole.SCHEDULER: - self.connector_scheduler = AFDDecodeBenchConnectorScheduler( - vllm_config, - kv_cache_config, - ) + self.connector_scheduler = AFDDecodeBenchConnectorScheduler(vllm_config) elif role == KVConnectorRole.WORKER: self.connector_worker = AFDDecodeBenchConnectorWorker( vllm_config, @@ -144,15 +91,6 @@ def save_kv_layer( def wait_for_save(self) -> None: pass - def bind_kv_cache_manager(self, kv_cache_manager: KVCacheManager) -> None: - super().bind_kv_cache_manager(kv_cache_manager) - assert self.connector_scheduler is not None - self.connector_scheduler.bind_kv_cache_manager(kv_cache_manager) - - def on_new_request(self, request: Request) -> None: - assert self.connector_scheduler is not None - self.connector_scheduler.on_new_request(request) - def get_num_new_matched_tokens( self, request: Request, @@ -190,9 +128,7 @@ def request_finished( block_ids: list[int], ) -> tuple[bool, dict[str, Any] | None]: assert self.connector_scheduler is not None - self.connector_scheduler.request_finished(request, block_ids) - # Never delay the free: the prompt pool pins its blocks with references - # of its own, so every request is released as soon as it finishes. + self.connector_scheduler.request_finished(request) return False, None def request_finished_all_groups( @@ -201,69 +137,24 @@ def request_finished_all_groups( block_ids: tuple[list[int], ...], ) -> tuple[bool, dict[str, Any] | None]: assert self.connector_scheduler is not None - # The prompt pool supports a single KV cache group (checked at startup). - self.connector_scheduler.request_finished( - request, - block_ids[0] if block_ids else [], - ) + self.connector_scheduler.request_finished(request) return False, None class AFDDecodeBenchConnectorScheduler: """Scheduler-side implementation for AFDDecodeBenchConnector.""" - def __init__( - self, - vllm_config: VllmConfig, - kv_cache_config: KVCacheConfig | None = None, - ) -> None: + def __init__(self, vllm_config: VllmConfig) -> None: self.vllm_config = vllm_config self.block_size = vllm_config.cache_config.block_size self._filled_requests: set[str] = set() self._pending_fills: dict[str, tuple[tuple[list[int], ...], int]] = {} - kv_transfer_config = vllm_config.kv_transfer_config - assert kv_transfer_config is not None - pool_config = PromptPoolConfig( - num_prompts=kv_transfer_config.get_from_extra_config( - "num_prompts", PromptPoolConfig.num_prompts - ), - prompt_len=kv_transfer_config.get_from_extra_config( - "prompt_len", PromptPoolConfig.prompt_len - ), - seed=kv_transfer_config.get_from_extra_config( - "seed", PromptPoolConfig.seed - ), - max_output_len=kv_transfer_config.get_from_extra_config( - "max_output_len", PromptPoolConfig.max_output_len - ), - ) - self.prompt_pool = ( - PromptPool(pool_config, vllm_config, kv_cache_config) - if pool_config.num_prompts > 0 - else None - ) - - def bind_kv_cache_manager(self, kv_cache_manager: KVCacheManager) -> None: - if self.prompt_pool is not None: - self.prompt_pool.bind_block_pool(kv_cache_manager.block_pool) - - def on_new_request(self, request: Request) -> None: - if self.prompt_pool is not None: - self.prompt_pool.assign(request) - def get_num_new_matched_tokens( self, request: Request, num_computed_tokens: int, ) -> tuple[int, bool]: - if self.prompt_pool is not None: - num_tokens_to_fill = self.prompt_pool.num_tokens_to_fill( - request, - num_computed_tokens, - ) - return num_tokens_to_fill, False - req_id = request.request_id allow_refill_after_preempt = ( request.status == RequestStatus.PREEMPTED @@ -317,269 +208,8 @@ def build_connector_meta( self._pending_fills.clear() return meta - def request_finished(self, request: Request, block_ids: list[int]) -> None: + def request_finished(self, request: Request) -> None: self._filled_requests.discard(request.request_id) - if self.prompt_pool is not None: - self.prompt_pool.request_finished(request, block_ids) - - -class PromptPool: - """Per-rank pool of real-KV prompts that measured requests decode over. - - Pinning: when a warmup request that prefilled its pool prompt finishes, - the pool takes one reference of its own on the entry's full prompt blocks - (``BlockPool.touch``) and holds it for the server's lifetime; the scheduler - then frees the request as usual. The extra reference keeps the blocks out - of the free queue, so LRU eviction cannot drop a pool entry, and they stay - in the prefix cache, so a measured request that maps to the entry shares - them like any prefix hit. Unlike a delayed free, this leaves no finished - request behind in the scheduler, so the engine still goes idle and shuts - down normally. - - Measured requests map only onto pinned entries. One that arrives before - any entry is pinned falls back to the whole pool; it may miss the prefix - cache, so it is counted and logged, and the client's checks catch it. - """ - - def __init__( - self, - config: PromptPoolConfig, - vllm_config: VllmConfig, - kv_cache_config: KVCacheConfig | None, - ) -> None: - if not vllm_config.cache_config.enable_prefix_caching: - raise ValueError( - "AFDDecodeBenchConnector prompt pool requires prefix caching; " - "launch with --enable-prefix-caching" - ) - if kv_cache_config is not None and len(kv_cache_config.kv_cache_groups) != 1: - raise ValueError( - "AFDDecodeBenchConnector prompt pool supports a single KV cache " - f"group, got {len(kv_cache_config.kv_cache_groups)}" - ) - - self.block_size = vllm_config.cache_config.block_size - self.max_num_seqs = vllm_config.scheduler_config.max_num_seqs - self._block_pool: BlockPool | None = None - max_model_len = vllm_config.model_config.max_model_len - self.num_prompts = config.num_prompts - # A prefix-cache hit covers full blocks only and stops one token short - # of the prompt, so a prompt one token past a block boundary hits every - # token but the last. - self.prompt_len = ( - config.prompt_len - 1 - ) // self.block_size * self.block_size + 1 - if self.prompt_len == 1: - raise ValueError( - f"prompt_len ({config.prompt_len}) must exceed the KV block size " - f"({self.block_size})" - ) - # The full prompt blocks a hit shares, and so the blocks a pinned entry - # holds; the block with the last prompt token is never shared. - self.num_entry_blocks = self.prompt_len // self.block_size - self.max_output_len = ( - config.max_output_len - if config.max_output_len is not None - else max_model_len - self.prompt_len - ) - if self.max_output_len < 1 or ( - self.prompt_len + self.max_output_len > max_model_len - ): - raise ValueError( - f"prompt_len ({self.prompt_len}) + max_output_len " - f"({self.max_output_len}) must fit max_model_len ({max_model_len})" - ) - - # A CPU generator gives every rank the same pool for the same seed. - generator = torch.Generator().manual_seed(config.seed) - self.prompt_tokens = torch.randint( - vllm_config.model_config.get_vocab_size(), - (self.num_prompts, self.prompt_len), - generator=generator, - dtype=torch.int32, - ) - # Hashing a long prompt per request is avoidable work on the scheduler: - # an entry's hashes are computed once and copied after that. - self._entry_block_hashes: dict[int, list[BlockHash]] = {} - self._unwarmed_entries = deque(range(self.num_prompts)) - self._warming_entries: dict[str, int] = {} - # Pinned entries in pin order, and the block ids each one holds. - self._pinned_entries: list[int] = [] - self._pinned_block_ids: dict[int, list[int]] = {} - self._num_assigned_measured = 0 - # Measured requests assigned while part of the pool was still unpinned. - self._num_early_measured = 0 - # request_id -> prompt tokens the prefix cache missed (0 on a hit). - self._measured_misses: dict[str, int] = {} - self._num_finished_measured = 0 - self._num_missed_requests = 0 - self._num_missed_tokens = 0 - - def bind_block_pool(self, block_pool: BlockPool) -> None: - """Bind the scheduler's block pool, whose blocks the pool pins.""" - self.check_capacity(block_pool.get_num_free_blocks()) - self._block_pool = block_pool - - def check_capacity(self, num_free_blocks: int) -> None: - """Fail fast when the pool and the running requests cannot both fit.""" - blocks_per_entry = self.num_entry_blocks - # A measured request shares its entry's full blocks and allocates only - # the last prompt token and its output. - blocks_per_seq = cdiv(1 + self.max_output_len, self.block_size) - num_reserved_blocks = self.max_num_seqs * blocks_per_seq - num_pool_blocks = self.num_prompts * blocks_per_entry - if num_pool_blocks + num_reserved_blocks > num_free_blocks: - max_num_prompts = max( - 0, - (num_free_blocks - num_reserved_blocks) // blocks_per_entry, - ) - raise ValueError( - f"AFDDecodeBenchConnector prompt pool needs {num_pool_blocks} KV " - f"blocks for {self.num_prompts} prompts of {self.prompt_len} " - f"tokens plus {num_reserved_blocks} for {self.max_num_seqs} " - f"sequences of up to {self.max_output_len} output tokens, but " - f"only {num_free_blocks} are free; use num_prompts <= " - f"{max_num_prompts} or shorten prompt_len/max_output_len" - ) - if self.num_prompts < self.max_num_seqs: - logger.warning( - "AFDDecodeBenchConnector: %d pool prompts are fewer than " - "max_num_seqs (%d); concurrent requests will share KV blocks", - self.num_prompts, - self.max_num_seqs, - ) - logger.info( - "AFDDecodeBenchConnector prompt pool: %d prompts x %d tokens use " - "%d of %d free KV blocks; %d reserved for running requests", - self.num_prompts, - self.prompt_len, - num_pool_blocks, - num_free_blocks, - num_reserved_blocks, - ) - - def assign(self, request: Request) -> None: - """Rewrite a new request's prompt to a pool entry's, before scheduling. - - The first ``num_prompts`` requests warm one unwarmed entry each; every - later request maps round-robin onto the pinned entries, whatever it - asked for. - """ - if self._unwarmed_entries: - entry = self._unwarmed_entries.popleft() - self._warming_entries[request.request_id] = entry - request.max_tokens = WARMUP_OUTPUT_TOKENS - else: - entry = self._next_measured_entry() - - prompt_token_ids = self.prompt_tokens[entry].tolist() - request.prompt_token_ids = prompt_token_ids - request.num_prompt_tokens = len(prompt_token_ids) - # all_token_ids is a read-only view of _all_token_ids: replace in place. - request._all_token_ids[:] = prompt_token_ids - request.block_hashes.clear() - entry_block_hashes = self._entry_block_hashes.get(entry) - if entry_block_hashes is None: - request.update_block_hashes() - self._entry_block_hashes[entry] = list(request.block_hashes) - else: - request.block_hashes.extend(entry_block_hashes) - - def _next_measured_entry(self) -> int: - """Pick a measured request's entry: round-robin over pinned entries. - - An entry that is still warming, or whose warmup was aborted, is not - in the prefix cache yet: a request mapped onto it would decode over - unwritten blocks and could leave them cached under the entry's hashes. - """ - index = self._num_assigned_measured - self._num_assigned_measured += 1 - num_pinned = len(self._pinned_entries) - if num_pinned < self.num_prompts: - if self._num_early_measured == 0: - logger.warning( - "AFDDecodeBenchConnector: a measured request arrived with " - "%d of %d pool prompts pinned; finish the warmup before " - "measuring", - num_pinned, - self.num_prompts, - ) - self._num_early_measured += 1 - if num_pinned == 0: - # Nothing is pinned yet: any entry may miss, and the miss is faked. - return index % self.num_prompts - return self._pinned_entries[index % num_pinned] - - def num_tokens_to_fill(self, request: Request, num_computed_tokens: int) -> int: - """Prompt tokens to fake as external: 0 on a hit or during warmup.""" - if request.request_id in self._warming_entries: - return 0 - # A pinned entry hits every prompt token but the last. Anything the - # cache misses is faked so the step stays decode-only, and counted. - num_missed_tokens = max(0, request.num_tokens - 1 - num_computed_tokens) - self._measured_misses[request.request_id] = num_missed_tokens - return num_missed_tokens - - def request_finished(self, request: Request, block_ids: list[int]) -> None: - """Pin a finished warmup request's blocks; tally measured hits. - - Called before the scheduler frees the request, so a pinned block's - reference count never drops to zero. - """ - entry = self._warming_entries.pop(request.request_id, None) - if entry is not None: - if request.status not in PINNABLE_WARMUP_STATUSES or not self._pin( - entry, block_ids - ): - # Aborted before its prompt was cached: warm it again later. - self._unwarmed_entries.appendleft(entry) - return - - num_missed_tokens = self._measured_misses.pop(request.request_id, None) - if num_missed_tokens is None: - return - self._num_finished_measured += 1 - if num_missed_tokens > 0: - self._num_missed_requests += 1 - self._num_missed_tokens += num_missed_tokens - if self._num_finished_measured % POOL_STATS_LOG_INTERVAL == 0: - logger.info( - "AFDDecodeBenchConnector prompt pool: %d measured requests, " - "%d prefix-cache misses (%d prompt tokens faked), %d assigned " - "before the pool was warm", - self._num_finished_measured, - self._num_missed_requests, - self._num_missed_tokens, - self._num_early_measured, - ) - - def _pin(self, entry: int, block_ids: list[int]) -> bool: - """Hold one reference on an entry's cached full prompt blocks.""" - assert self._block_pool is not None, "the KV cache manager is not bound" - entry_block_ids = block_ids[: self.num_entry_blocks] - blocks = [self._block_pool.blocks[block_id] for block_id in entry_block_ids] - if len(blocks) < self.num_entry_blocks or any( - block.is_null or block.block_hash is None for block in blocks - ): - logger.warning( - "AFDDecodeBenchConnector: pool prompt %d finished warming with " - "%d of %d prompt blocks cached; warming it again", - entry, - sum(block.block_hash is not None for block in blocks), - self.num_entry_blocks, - ) - return False - self._block_pool.touch(blocks) - self._pinned_entries.append(entry) - self._pinned_block_ids[entry] = entry_block_ids - if len(self._pinned_entries) == self.num_prompts: - logger.info( - "AFDDecodeBenchConnector prompt pool warm: %d prompts " - "pinned in %d KV blocks", - len(self._pinned_entries), - self.num_prompts * self.num_entry_blocks, - ) - return True class AFDDecodeBenchConnectorWorker: @@ -709,6 +339,4 @@ def _fill_blocks(self, group_idx: int, block_ids: list[int]) -> None: "AFDDecodeBenchConnectorMetadata", "AFDDecodeBenchConnectorScheduler", "AFDDecodeBenchConnectorWorker", - "PromptPool", - "PromptPoolConfig", ] diff --git a/tools/benchmarks/decode_bench_client.py b/tools/benchmarks/decode_bench_client.py deleted file mode 100644 index e3e9ddf20..000000000 --- a/tools/benchmarks/decode_bench_client.py +++ /dev/null @@ -1,485 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project -"""Decode-only benchmark over the decode-bench prompt pool, run with vllm bench. - -Drives an Attention instance whose ``AFDDecodeBenchConnector`` has -``num_prompts`` > 0 (see ``tools/benchmarks/decode_bench.py``): - -1. Warmup: one ``vllm bench serve`` run per Attention DP rank, each pinned to - its rank with the ``X-data-parallel-rank`` header, sends ``--pool-prompts`` - requests, so every rank prefills and pins its whole prompt pool. -2. Check: the server's prefix-cache counters must show each warmup request - looked up ``--prompt-len`` prompt tokens, i.e. the server's pool prompts are - as long as the measured prompts below. -3. Measured: one ``vllm bench serve`` run with ``--random-input-len - --prompt-len``. Every argument after ``--`` goes to it unchanged (request - rate, concurrency, number of requests, output length, ...). The server remaps - each request onto a pool prompt, so the prompt text is only a placeholder. -4. Check: over the measured run the server must have prefilled exactly one - token per finished request (the last prompt token), with no failed requests - and no preemptions; otherwise this script exits non-zero. - -vllm bench writes its usual result JSON for each run into ``--result-dir``; -this script adds ``decode_bench_report.json`` with the checks. - -Usage: - python tools/benchmarks/decode_bench_client.py --pool-prompts 1000 \ - --prompt-len 1025 -- --num-prompts 4000 --random-output-len 256 \ - --max-concurrency 64 --request-rate 20 -""" - -from __future__ import annotations - -import argparse -import json -import subprocess -import sys -import time -import urllib.error -import urllib.request -from collections.abc import Sequence -from dataclasses import asdict, dataclass -from pathlib import Path - -DATA_PARALLEL_RANK_HEADER = "X-data-parallel-rank" -MODELS_PATH = "/v1/models" -METRICS_PATH = "/metrics" -PREFIX_CACHE_QUERIES_METRIC = "vllm:prefix_cache_queries" -PREFIX_CACHE_HITS_METRIC = "vllm:prefix_cache_hits" -REQUEST_SUCCESS_METRIC = "vllm:request_success" -PREEMPTIONS_METRIC = "vllm:num_preemptions" -HTTP_OK = 200 -READY_POLL_INTERVAL_S = 5.0 -# The server forces warmup requests to one output token as well. -WARMUP_OUTPUT_TOKENS = 1 -# A decode-only request prefills only its last prompt token. -DECODE_ONLY_PREFILLED_TOKENS = 1.0 -PREFILLED_TOKENS_TOLERANCE = 0.01 -PROMPT_LEN_TOLERANCE = 0.5 -# vllm bench flags this script sets for every run; overriding them would break -# the prompt-length contract or lose the result file. -CLIENT_OWNED_BENCH_FLAGS = frozenset( - { - "--host", - "--port", - "--base-url", - "--model", - "--tokenizer", - "--dataset-name", - "--random-input-len", - "--random-range-ratio", - "--random-prefix-len", - "--save-result", - "--result-dir", - "--result-filename", - } -) -# The warmup also owns how many requests each rank gets and where they go. -WARMUP_OWNED_BENCH_FLAGS = CLIENT_OWNED_BENCH_FLAGS | { - "--num-prompts", - "--random-output-len", - "--header", - "--max-concurrency", -} -WARMUP_RESULT_FILENAME = "warmup-rank{rank}.json" -WARMUP_LOG_FILENAME = "warmup-rank{rank}.log" -MEASURED_RESULT_FILENAME = "measured.json" -REPORT_FILENAME = "decode_bench_report.json" -MEASURED_SUMMARY_KEYS = ( - "completed", - "failed", - "duration", - "request_throughput", - "output_throughput", - "total_input_tokens", - "total_output_tokens", - "median_ttft_ms", - "p99_ttft_ms", - "median_tpot_ms", - "p99_tpot_ms", - "median_itl_ms", - "p99_itl_ms", -) - - -@dataclass(frozen=True) -class ServerCounters: - """Server-wide totals of the counters the checks need, summed over engines.""" - - prefix_cache_queries: float = 0.0 - prefix_cache_hits: float = 0.0 - finished_requests: float = 0.0 - preemptions: float = 0.0 - - def __sub__(self, other: ServerCounters) -> ServerCounters: - return ServerCounters( - prefix_cache_queries=self.prefix_cache_queries - other.prefix_cache_queries, - prefix_cache_hits=self.prefix_cache_hits - other.prefix_cache_hits, - finished_requests=self.finished_requests - other.finished_requests, - preemptions=self.preemptions - other.preemptions, - ) - - @property - def prefilled_tokens(self) -> float: - return self.prefix_cache_queries - self.prefix_cache_hits - - -@dataclass(frozen=True) -class Check: - name: str - passed: bool - detail: str - - -def parse_args(argv: Sequence[str]) -> tuple[argparse.Namespace, list[str]]: - """Split ``argv`` at ``--``: client flags before, measured vllm bench after.""" - argv = list(argv) - split = argv.index("--") if "--" in argv else len(argv) - client_argv, bench_argv = argv[:split], argv[split + 1 :] - - parser = argparse.ArgumentParser( - description=__doc__.split("\n\n")[0], - epilog="Arguments after `--` are passed to the measured `vllm bench serve`.", - ) - parser.add_argument("--host", default="127.0.0.1") - parser.add_argument("--port", type=int, default=18000) - parser.add_argument( - "--model", - default=None, - help="Served model name; defaults to the first one in /v1/models.", - ) - parser.add_argument( - "--tokenizer", - default=None, - help="Tokenizer vllm bench uses to build prompts; defaults to --model.", - ) - parser.add_argument( - "--dp-size", - type=int, - default=2, - help="Attention data-parallel ranks; each gets its own warmup run.", - ) - parser.add_argument( - "--pool-prompts", - type=int, - required=True, - help="Pool prompts per Attention DP rank; must match the server's " - "kv_connector_extra_config num_prompts.", - ) - parser.add_argument( - "--prompt-len", - type=int, - default=1025, - help="Prompt length of the warmup and measured runs; must equal the " - "server's pool prompt length (checked after the warmup).", - ) - parser.add_argument( - "--warmup-concurrency", - type=int, - default=64, - help="In-flight requests of each rank's warmup run.", - ) - parser.add_argument( - "--warmup-bench-args", - nargs=argparse.REMAINDER, - default=[], - help="Extra vllm bench arguments for the warmup runs; must come last " - "before `--`.", - ) - parser.add_argument( - "--skip-warmup", - action="store_true", - help="Measure only, against a server whose pool is already warm.", - ) - parser.add_argument( - "--result-dir", - type=Path, - default=Path("decode_bench_results"), - ) - parser.add_argument("--ready-timeout-s", type=float, default=1800.0) - args = parser.parse_args(client_argv) - reject_owned_flags(args.warmup_bench_args, WARMUP_OWNED_BENCH_FLAGS, "warmup") - reject_owned_flags(bench_argv, CLIENT_OWNED_BENCH_FLAGS, "measured") - return args, bench_argv - - -def reject_owned_flags( - bench_args: Sequence[str], owned_flags: frozenset[str], phase: str -) -> None: - for bench_arg in bench_args: - flag = bench_arg.split("=", 1)[0] - if flag in owned_flags: - raise SystemExit( - f"{flag} is set by decode_bench_client.py for the {phase} vllm " - f"bench run; drop it (owned: {', '.join(sorted(owned_flags))})" - ) - - -def wait_for_model(base_url: str, timeout_s: float) -> str: - """Poll until the server answers /v1/models; return the first model id.""" - deadline = time.monotonic() + timeout_s - while True: - try: - with urllib.request.urlopen(f"{base_url}{MODELS_PATH}") as response: - if response.status == HTTP_OK: - return json.loads(response.read())["data"][0]["id"] - except (urllib.error.URLError, ConnectionError): - pass - if time.monotonic() > deadline: - raise TimeoutError(f"{base_url} not ready after {timeout_s:.0f}s") - time.sleep(READY_POLL_INTERVAL_S) - - -def read_counters(base_url: str) -> ServerCounters: - with urllib.request.urlopen(f"{base_url}{METRICS_PATH}") as response: - text = response.read().decode() - totals = { - PREFIX_CACHE_QUERIES_METRIC: 0.0, - PREFIX_CACHE_HITS_METRIC: 0.0, - REQUEST_SUCCESS_METRIC: 0.0, - PREEMPTIONS_METRIC: 0.0, - } - for line in text.splitlines(): - if not line or line.startswith("#"): - continue - sample, _, value = line.rpartition(" ") - metric = sample.split("{", 1)[0].removesuffix("_total") - if metric in totals: - totals[metric] += float(value) - return ServerCounters( - prefix_cache_queries=totals[PREFIX_CACHE_QUERIES_METRIC], - prefix_cache_hits=totals[PREFIX_CACHE_HITS_METRIC], - finished_requests=totals[REQUEST_SUCCESS_METRIC], - preemptions=totals[PREEMPTIONS_METRIC], - ) - - -def bench_command( - args: argparse.Namespace, - model: str, - result_filename: str, - bench_args: Sequence[str], -) -> list[str]: - """A ``vllm bench serve`` command over --prompt-len random prompts.""" - return [ - sys.executable, - "-m", - "vllm.entrypoints.cli.main", - "bench", - "serve", - "--host", - args.host, - "--port", - str(args.port), - "--model", - model, - "--tokenizer", - args.tokenizer or model, - "--dataset-name", - "random", - "--random-input-len", - str(args.prompt_len), - "--random-range-ratio", - "0", - "--random-prefix-len", - "0", - "--ignore-eos", - "--save-result", - "--result-dir", - str(args.result_dir), - "--result-filename", - result_filename, - *bench_args, - ] - - -def read_bench_result(args: argparse.Namespace, filename: str) -> dict[str, object]: - return json.loads((args.result_dir / filename).read_text()) - - -def run_warmup(args: argparse.Namespace, model: str) -> list[dict[str, object]]: - """Run one pinned vllm bench per rank in parallel; return their results.""" - processes = [] - for rank in range(args.dp_size): - command = bench_command( - args, - model, - WARMUP_RESULT_FILENAME.format(rank=rank), - [ - "--num-prompts", - str(args.pool_prompts), - "--random-output-len", - str(WARMUP_OUTPUT_TOKENS), - "--max-concurrency", - str(args.warmup_concurrency), - "--header", - f"{DATA_PARALLEL_RANK_HEADER}={rank}", - *args.warmup_bench_args, - ], - ) - log_path = args.result_dir / WARMUP_LOG_FILENAME.format(rank=rank) - with log_path.open("w") as log: - processes.append( - (rank, log_path, subprocess.Popen(command, stdout=log, stderr=log)) - ) - results = [] - for rank, log_path, process in processes: - if process.wait() != 0: - raise RuntimeError( - f"warmup vllm bench for rank {rank} exited {process.returncode}; " - f"see {log_path}" - ) - results.append( - read_bench_result(args, WARMUP_RESULT_FILENAME.format(rank=rank)) - ) - return results - - -def check_warmup( - args: argparse.Namespace, - results: list[dict[str, object]], - delta: ServerCounters, -) -> list[Check]: - num_failed = sum(int(str(result["failed"])) for result in results) - checks = [ - Check( - "warmup requests succeeded", - num_failed == 0, - f"{num_failed} failed; a failed warmup request leaves its pool entry " - "for a later request to prefill", - ) - ] - min_requests = args.dp_size * args.pool_prompts - checks.append( - Check( - "every rank received a full pool of warmup requests", - delta.finished_requests >= min_requests, - f"{delta.finished_requests:.0f} finished, need >= {min_requests}", - ) - ) - tokens_per_request = ( - delta.prefix_cache_queries / delta.finished_requests - if delta.finished_requests - else 0.0 - ) - prompt_len_matches = ( - abs(tokens_per_request - args.prompt_len) <= PROMPT_LEN_TOLERANCE - ) - checks.append( - Check( - "server pool prompts match --prompt-len", - prompt_len_matches, - f"{tokens_per_request:.2f} prompt tokens per warmup request, " - f"--prompt-len {args.prompt_len}" - + ( - "" - if prompt_len_matches - else "; set the server's prompt_len to one token past a block " - "boundary and pass the same value here" - ), - ) - ) - return checks - - -def check_measured(result: dict[str, object], delta: ServerCounters) -> list[Check]: - prefilled_per_request = ( - delta.prefilled_tokens / delta.finished_requests - if delta.finished_requests - else 0.0 - ) - return [ - Check( - "measured requests succeeded", - int(str(result["failed"])) == 0, - f"{result['failed']} failed of {result['completed']} completed", - ), - Check( - "decode only: one prefilled token per request", - abs(prefilled_per_request - DECODE_ONLY_PREFILLED_TOKENS) - <= PREFILLED_TOKENS_TOLERANCE, - f"{prefilled_per_request:.4f} prefilled tokens per finished request " - f"({delta.prefilled_tokens:.0f} tokens, " - f"{delta.finished_requests:.0f} requests)", - ), - Check( - "no preemptions", - delta.preemptions == 0, - f"{delta.preemptions:.0f} preemptions; a preempted request recomputes " - "its context", - ), - ] - - -def print_checks(phase: str, checks: list[Check]) -> None: - for check in checks: - verdict = "PASS" if check.passed else "FAIL" - print(f"[{phase}] {verdict} {check.name}: {check.detail}", flush=True) - - -def main(argv: Sequence[str] | None = None) -> int: - args, measured_bench_args = parse_args(sys.argv[1:] if argv is None else argv) - args.result_dir.mkdir(parents=True, exist_ok=True) - base_url = f"http://{args.host}:{args.port}" - served_model = wait_for_model(base_url, args.ready_timeout_s) - model = args.model or served_model - report: dict[str, object] = { - "config": {key: str(value) for key, value in vars(args).items()} - | {"model": model, "measured_bench_args": measured_bench_args}, - } - checks: list[Check] = [] - - if not args.skip_warmup: - print( - f"Warmup: vllm bench x {args.dp_size} ranks, {args.pool_prompts} " - f"prompts of {args.prompt_len} tokens each", - flush=True, - ) - before = read_counters(base_url) - warmup_results = run_warmup(args, model) - warmup_delta = read_counters(base_url) - before - warmup_checks = check_warmup(args, warmup_results, warmup_delta) - print_checks("warmup", warmup_checks) - checks += warmup_checks - report["warmup"] = { - "server_counters": asdict(warmup_delta), - "checks": [asdict(check) for check in warmup_checks], - } - if not all(check.passed for check in warmup_checks): - write_report(args, report, checks) - return 1 - - print(f"Measured: vllm bench serve {' '.join(measured_bench_args)}", flush=True) - before = read_counters(base_url) - completed = subprocess.run( - bench_command(args, model, MEASURED_RESULT_FILENAME, measured_bench_args) - ) - if completed.returncode != 0: - print(f"measured vllm bench exited {completed.returncode}", flush=True) - return completed.returncode - measured_delta = read_counters(base_url) - before - measured_result = read_bench_result(args, MEASURED_RESULT_FILENAME) - measured_checks = check_measured(measured_result, measured_delta) - print_checks("measured", measured_checks) - checks += measured_checks - report["measured"] = { - "vllm_bench_result": str(args.result_dir / MEASURED_RESULT_FILENAME), - "summary": {key: measured_result.get(key) for key in MEASURED_SUMMARY_KEYS}, - "server_counters": asdict(measured_delta), - "checks": [asdict(check) for check in measured_checks], - } - write_report(args, report, checks) - return 0 if all(check.passed for check in checks) else 1 - - -def write_report( - args: argparse.Namespace, report: dict[str, object], checks: list[Check] -) -> None: - report["passed"] = all(check.passed for check in checks) - path = args.result_dir / REPORT_FILENAME - path.write_text(json.dumps(report, indent=2)) - print(f"{'PASS' if report['passed'] else 'FAIL'}: report saved to {path}") - - -if __name__ == "__main__": - sys.exit(main()) diff --git a/tools/benchmarks/request_generator.sh b/tools/benchmarks/request_generator.sh index c0240c277..c5e3bf091 100755 --- a/tools/benchmarks/request_generator.sh +++ b/tools/benchmarks/request_generator.sh @@ -2,20 +2,20 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project # -# Drive load against a running AFD instance (e.g. the attention endpoint of a +# Drive load against a running AFD instance (e.g. the Attention endpoint of a # recipe under recipe/gpu/). Thin wrapper around `vllm bench serve`; every knob -# is overridable via an environment variable. For a decode-only run over a -# warmed prompt pool, use decode_bench_client.py, which wraps vllm bench too. +# is overridable via an environment variable. For a decode-only run, start the +# recipe with DECODE_BENCH set; see tools/benchmarks/README.md. # # Usage: # MODEL_PATH=/path/to/weights tools/benchmarks/request_generator.sh -# HOST=127.0.0.1 PORT=18305 INPUT_LEN=8192 OUTPUT_LEN=256 MAX_CONCURRENCY=64 \ +# HOST=127.0.0.1 PORT=18000 INPUT_LEN=8192 OUTPUT_LEN=256 MAX_CONCURRENCY=64 \ # MODEL_PATH=/path/to/weights tools/benchmarks/request_generator.sh # # Variables (default): # MODEL_PATH model weights dir/name (required) # HOST server host (127.0.0.1) -# PORT server port (18305) +# PORT server port (18000) # DATASET_NAME vllm bench dataset (random) # INPUT_LEN random input length (ISL) (1024) # OUTPUT_LEN random output length (OSL) (128) @@ -29,7 +29,7 @@ set -euo pipefail MODEL_PATH=${MODEL_PATH:?set MODEL_PATH to the model weights dir/name} HOST=${HOST:-127.0.0.1} -PORT=${PORT:-18305} +PORT=${PORT:-18000} DATASET_NAME=${DATASET_NAME:-random} INPUT_LEN=${INPUT_LEN:-1024} OUTPUT_LEN=${OUTPUT_LEN:-128} From 891a232a091c9f9d0720ba653916f139c888c7e3 Mon Sep 17 00:00:00 2001 From: ronenkat <16743404+ronenkat@users.noreply.github.com> Date: Tue, 10 Nov 2026 18:49:00 +0200 Subject: [PATCH 32/32] recipe(gpu): one shared launcher for the GPU colocation recipes The GPU prefill_decode_colocation recipes came in two styles: nine LAYOUT-driven M2N scripts that each carried an identical ~75-line launch block, and hand-written P2P scripts that spelled out both vllm serve commands per DP/TP/pod variant. The P2P Qwen3 entry points exec'd the M2N scripts through a script-relative path that breaks when a single recipe file is mounted into a pod. Move the shared block into recipe/gpu/common/afd_colocation.sh (LAYOUT/POD placement, DP flags, the afd config, the decode-bench hook, the override variables and the two servers, plus FFN_TP, SERVE_HOST and an AFD_DRY_RUN print mode). Every recipe now states only its model, connector, default layout, mode and serve flags, and sources the launcher from AFD_RECIPE_LIB or its repository location. The P2P recipes move to the same style: DeepSeek-V2-Lite dbo.sh and eager_dbo.sh replace the seven 2A2F/4A4F/multi-pod variants (selected with LAYOUT, ATTN_TP and FFN_TP), Qwen3.5-122B graph.sh replaces the three-pod script (default LAYOUT=2A0F,2A0F,0A2F), and the Qwen3 entry points become standalone graph.sh recipes. The P2P recipes now use the launcher's ports (18000/18001/6239), POD=AFD_, FFN-first placement and 3600 s timeouts. deploy-afd-k8s ships the launcher in the recipe ConfigMap, sets AFD_RECIPE_LIB and SERVE_HOST in the pod instead of patching the script, and reads placement and ports from LAYOUT and the launcher only. Co-Authored-By: Claude Opus 5.5 Signed-off-by: ronenkat <16743404+ronenkat@users.noreply.github.com> --- .agents/skills/deploy-afd-k8s/SKILL.md | 35 ++-- .../skills/deploy-afd-k8s/resources/deploy.md | 46 +++-- .../resources/resolve-recipe.md | 178 ++++++------------ .../deploy-afd-k8s/resources/services.md | 30 ++- .../skills/deploy-afd-k8s/templates/pod.yaml | 28 +-- .../features/disaggregation-topologies.md | 7 +- .agents/skills/run-vllm-bench-k8s/SKILL.md | 6 +- recipe/README.md | 6 + recipe/gpu/DeepEPM2NAFDConnector/README.md | 127 +++---------- .../prefill_decode_colocation/dbo.sh | 101 ++-------- .../prefill_decode_colocation/eager.sh | 97 +--------- .../prefill_decode_colocation/graph.sh | 99 +--------- .../prefill_decode_colocation/dbo.sh | 102 ++-------- .../prefill_decode_colocation/eager.sh | 98 ++-------- .../prefill_decode_colocation/graph.sh | 100 ++-------- .../prefill_decode_colocation/dbo.sh | 101 ++-------- .../prefill_decode_colocation/eager.sh | 97 +--------- .../prefill_decode_colocation/graph.sh | 99 +--------- recipe/gpu/P2pNcclAFDConnector/README.md | 47 ++--- .../deepseek_v2_lite/README.md | 96 ++++------ .../2a2f_eager_dbo_dp1tp2.sh | 55 ------ .../2a2f_eager_dbo_dp2tp1.sh | 55 ------ .../2a2f_graph_dbo_dp1tp2.sh | 61 ------ .../2a2f_graph_dbo_dp2tp1.sh | 61 ------ .../4a4f_eager_dbo_dp2tp2.sh | 55 ------ .../4a4f_graph_dbo_dp2tp2.sh | 61 ------ .../prefill_decode_colocation/dbo.sh | 21 +++ .../prefill_decode_colocation/eager_dbo.sh | 21 +++ .../multipod_2a_2f_graph_dbo_dp1tp2.sh | 73 ------- .../qwen3_235b_a22b_fp8/README.md | 17 +- .../4a2f_graph_dp4tp1.sh | 13 -- .../prefill_decode_colocation/graph.sh | 20 ++ .../qwen3_30b_a3b/README.md | 14 +- .../2a2f_graph_dp2tp1.sh | 12 -- .../prefill_decode_colocation/graph.sh | 18 ++ .../prefill_decode_colocation/graph.sh | 21 +++ .../multipod_2a_2a_2f_graph.sh | 100 ---------- recipe/gpu/common/README.md | 156 +++++++++++++++ recipe/gpu/common/afd_colocation.sh | 116 ++++++++++++ 39 files changed, 692 insertions(+), 1758 deletions(-) delete mode 100644 recipe/gpu/P2pNcclAFDConnector/deepseek_v2_lite/prefill_decode_colocation/2a2f_eager_dbo_dp1tp2.sh delete mode 100644 recipe/gpu/P2pNcclAFDConnector/deepseek_v2_lite/prefill_decode_colocation/2a2f_eager_dbo_dp2tp1.sh delete mode 100644 recipe/gpu/P2pNcclAFDConnector/deepseek_v2_lite/prefill_decode_colocation/2a2f_graph_dbo_dp1tp2.sh delete mode 100644 recipe/gpu/P2pNcclAFDConnector/deepseek_v2_lite/prefill_decode_colocation/2a2f_graph_dbo_dp2tp1.sh delete mode 100644 recipe/gpu/P2pNcclAFDConnector/deepseek_v2_lite/prefill_decode_colocation/4a4f_eager_dbo_dp2tp2.sh delete mode 100644 recipe/gpu/P2pNcclAFDConnector/deepseek_v2_lite/prefill_decode_colocation/4a4f_graph_dbo_dp2tp2.sh create mode 100644 recipe/gpu/P2pNcclAFDConnector/deepseek_v2_lite/prefill_decode_colocation/dbo.sh create mode 100644 recipe/gpu/P2pNcclAFDConnector/deepseek_v2_lite/prefill_decode_colocation/eager_dbo.sh delete mode 100644 recipe/gpu/P2pNcclAFDConnector/deepseek_v2_lite/prefill_decode_colocation/multipod_2a_2f_graph_dbo_dp1tp2.sh delete mode 100755 recipe/gpu/P2pNcclAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/4a2f_graph_dp4tp1.sh create mode 100644 recipe/gpu/P2pNcclAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/graph.sh delete mode 100755 recipe/gpu/P2pNcclAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/2a2f_graph_dp2tp1.sh create mode 100644 recipe/gpu/P2pNcclAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/graph.sh create mode 100644 recipe/gpu/P2pNcclAFDConnector/qwen3_5_122b_a10b_fp8/prefill_decode_colocation/graph.sh delete mode 100644 recipe/gpu/P2pNcclAFDConnector/qwen3_5_122b_a10b_fp8/prefill_decode_colocation/multipod_2a_2a_2f_graph.sh create mode 100644 recipe/gpu/common/README.md create mode 100644 recipe/gpu/common/afd_colocation.sh diff --git a/.agents/skills/deploy-afd-k8s/SKILL.md b/.agents/skills/deploy-afd-k8s/SKILL.md index e7ea5fe0c..a7ef9b68e 100644 --- a/.agents/skills/deploy-afd-k8s/SKILL.md +++ b/.agents/skills/deploy-afd-k8s/SKILL.md @@ -1,6 +1,6 @@ --- name: deploy-afd-k8s -description: Use when the user asks to deploy, serve, or stand up an AFD GPU recipe on Kubernetes/OpenShift from a recipe .sh script - creating the model PVC, recipe ConfigMap, serve pod(s), and Service(s), and waiting until the endpoint answers. Placement across pods is never asked for or assumed -- it is read directly off the recipe's own top-of-file comment block (however that block describes it, strict or prose): a recipe whose comment describes a per-pod placement deploys exactly that many pods with exactly that per-pod setup; a recipe without one deploys as a single pod. Do not use for local (non-k8s) serving, NPU recipes, prefill-decode disaggregation recipes, or E2E correctness testing (see run-e2e). +description: Use when the user asks to deploy, serve, or stand up an AFD GPU recipe on Kubernetes/OpenShift from a recipe .sh script - creating the model PVC, recipe ConfigMap, serve pod(s), and Service(s), and waiting until the endpoint answers. Placement across pods is never asked for or assumed -- it is the recipe's LAYOUT (the caller's setting, else the script's default): one pod per LAYOUT entry, a one-entry LAYOUT being a single pod. Do not use for local (non-k8s) serving, NPU recipes, prefill-decode disaggregation recipes, or E2E correctness testing (see run-e2e). --- # Deploy an AFD GPU recipe on Kubernetes @@ -15,13 +15,13 @@ On success this skill guarantees, in the target namespace: | | | |---|---| -| Endpoint | `http://vllm-service:${CLIENT_PORT}` (OpenAI-compatible, in-namespace; `CLIENT_PORT` defaults to `18305`) | +| Endpoint | `http://vllm-service:${CLIENT_PORT}` (OpenAI-compatible, in-namespace; `CLIENT_PORT` is the launcher's Attention port, `18000`) | | Serving | the model named by `MODEL_ID` | -| Pod(s) / Service | `vllm-pod`, or one pod per pod the recipe's header comment describes (named `vllm-pod-`); Service `vllm-service`, labels `app=afd-recipe,role=serve` | +| Pod(s) / Service | `vllm-pod`, or one pod per `LAYOUT` entry (named `vllm-pod-afd-`); Service `vllm-service`, labels `app=afd-recipe,role=serve` | | Model cache | `PVC_NAME`, mounted at `/models`, `HF_HOME=/models/.hf_home` | | Left running | yes, deliberately -- so weights stay warm for follow-up runs; step 5 asks the caller when to delete the pods | -Whenever a recipe's header spans more than one pod, two more internal-only +Whenever a recipe's `LAYOUT` spans more than one pod, two more internal-only Services may also exist (`vllm-ffn-p2p-service` always; `vllm-attn-dp-service` and/or `vllm-ffn-dp-service` only if that role is itself split across more than one pod) -- callers never talk to these directly. See @@ -38,18 +38,13 @@ than one pod) -- callers never talk to these directly. See NIXL-enabled image plus a proxy, and resolve a `SCRIPT_DIR`-relative path that this skill's flat script mount (step 4b) breaks. Ask the user for a different workflow for disaggregation. -- Recipes are hand-written shell scripts, not JSON configs -- there is no - generator step and no manifest to read. -- A recipe is **placement-driven** (more than one pod) exactly when its - top-of-file comment block describes a per-pod placement -- read that - block, don't infer shape from flags. Today's recipes spell it as one - `# For pod N: POD=...` line per pod, but the comment isn't guaranteed to - use that exact grammar going forward. See - [resolve-recipe.md](resources/resolve-recipe.md). -- `DeepEPM2NAFDConnector` recipes take their placement from `LAYOUT` (one - pod per `AF` entry) and run-specific settings as env (`VLLM_ARGS`, - `DECODE_BENCH`, ...); see [resolve-recipe.md](resources/resolve-recipe.md). - This skill only serves; load for the decode-only benchmark +- Recipes are short shell scripts on the shared launcher + `recipe/gpu/common/afd_colocation.sh` (see `recipe/gpu/common/README.md`); + the recipe ConfigMap ships both files (step 4b). +- Placement is `LAYOUT` (one pod per `AF` entry) and run-specific + settings are env (`ATTN_TP`, `FFN_TP`, `VLLM_ARGS`, `DECODE_BENCH`, ...); + see [resolve-recipe.md](resources/resolve-recipe.md). This skill only + serves; load for the decode-only benchmark (`tools/benchmarks/request_generator.sh`) is run separately against the endpoint. @@ -68,9 +63,9 @@ than one pod) -- callers never talk to these directly. See docker push "$IMAGE" ``` - The image supplies only the `afd-plugin` install. The recipe script is - mounted fresh from local disk on every run (step 4b), so editing or adding - a recipe never requires a rebuild. + The image supplies only the `afd-plugin` install. The recipe script and + the shared launcher are mounted fresh from local disk on every run (step + 4b), so editing or adding a recipe never requires a rebuild. - An `hf-token-secret` Secret with a `token` key, needed when `MODEL_ID` is a HF Hub repo id (not needed when weights are already staged on the PVC): ```bash @@ -86,7 +81,7 @@ exist **before** any Pod is deployed -- see | Resource | Covers | |---|---| -| [resolve-recipe.md](resources/resolve-recipe.md) | Step 1 (the only step that parses the recipe): read the recipe's serve-block fields and model id; read its placement (or lack of one) from its header comment -- never ask; record `GPU_COUNT` per pod. | +| [resolve-recipe.md](resources/resolve-recipe.md) | Step 1 (the only step that parses the recipe): model id, the `LAYOUT` plan (never ask), `GPU_COUNT` per pod, and the launcher's ports. | | [deploy.md](resources/deploy.md) | Steps 2-4f: take step 1's `GPU_COUNT`, confirm, PVC, ConfigMap, fsGroup, Services (step 4d -- which of the four, see [services.md](resources/services.md)), pod(s) via `templates/pod.yaml`, wait for Running and then for each pod's `Ready` condition. | | [report-and-teardown.md](resources/report-and-teardown.md) | Step 5: report the endpoint/model/nodes and ask when to delete the running pods. Step 6: teardown when the caller says so (pods/Services only -- never the PVC unless asked). | diff --git a/.agents/skills/deploy-afd-k8s/resources/deploy.md b/.agents/skills/deploy-afd-k8s/resources/deploy.md index 58643cbbe..0a971ec66 100644 --- a/.agents/skills/deploy-afd-k8s/resources/deploy.md +++ b/.agents/skills/deploy-afd-k8s/resources/deploy.md @@ -3,11 +3,8 @@ ## 2. GPU_COUNT Use the `GPU_COUNT` recorded by [resolve-recipe.md](resolve-recipe.md) -- -one value for a plain recipe, one per plan entry. Don't re-read the script -here; all recipe parsing lives in step 1. - -`LAYOUT` recipes (see [resolve-recipe.md](resolve-recipe.md)): each entry's -GPU count is its `a + f`: +`a + f` per `LAYOUT` entry. Don't re-read the script here; all recipe +parsing lives in step 1: ```bash LAYOUT=${LAYOUT:-$(grep -oE 'LAYOUT:-[0-9AF,]+' "$RECIPE_SCRIPT_PATH" | cut -d- -f2)} @@ -28,7 +25,7 @@ existing Pod(s)** with the same name(s) about to be deployed. AFD_PLUGIN_IMAGE= MODEL_ID= PVC_NAME= -CLIENT_PORT= # from the script's --port; default 18305 +CLIENT_PORT= # the launcher's ATTN_PORT (18000), see resolve-recipe.md RECIPE_SCRIPT_PATH= ``` @@ -81,11 +78,14 @@ rather than re-downloading. ## 4b. Recipe ConfigMap Lets the pod run any local recipe -- edited, uncommitted, or new -- without -rebuilding the image. Run from wherever `RECIPE_SCRIPT_PATH` resolves: +rebuilding the image. The recipe sources the shared launcher, so ship both; +`templates/pod.yaml` points the recipe at `/recipe/afd_colocation.sh` +(`AFD_RECIPE_LIB`). Run from the repository root: ```bash kubectl create configmap afd-recipe-script \ --from-file=recipe.sh="${RECIPE_SCRIPT_PATH}" \ + --from-file=afd_colocation.sh=recipe/gpu/common/afd_colocation.sh \ --dry-run=client -o yaml | kubectl apply -f - ``` @@ -142,23 +142,18 @@ deploy_pod() { } ``` -Plain recipe, one pod, both roles head, no overrides: +One-entry `LAYOUT` (e.g. the default `2A2F`), one pod, both roles head, +with or without the caller's settings: ```bash deploy_pod vllm-pod "$GPU_COUNT" head head "" +deploy_pod vllm-pod 4 head head $'VLLM_ARGS=--max-num-seqs 32' ``` -Placement-driven recipe, one call per entry parsed in -[resolve-recipe.md](resolve-recipe.md), e.g. the qwen3.5 3-pod plan: - -```bash -deploy_pod vllm-pod-attention-0 2 head none "POD=ATTENTION_0" -deploy_pod vllm-pod-attention-1 2 worker none $'POD=ATTENTION_1\nATTENTION_HEADLESS=1\nATTENTION_DP_START_RANK=2' -deploy_pod vllm-pod-ffn-0 2 none head "POD=FFN_0" -``` - -`LAYOUT` recipe, e.g. `LAYOUT=2A0F,0A2F` (2 GPUs per pod) with the -decode-only benchmark, or `LAYOUT=4A1F,4A1F` (5 GPUs per pod): +Several entries, one call per entry parsed in +[resolve-recipe.md](resolve-recipe.md), e.g. `LAYOUT=2A0F,0A2F` (2 GPUs per +pod) with the decode-only benchmark, `LAYOUT=4A1F,4A1F` (5 GPUs per pod), +or the Qwen3.5 recipe's default `2A0F,2A0F,0A2F`: ```bash SETTINGS=$'LAYOUT=2A0F,0A2F\nDECODE_BENCH={}\nAFD_REPO=/opt/afd-plugin' @@ -167,10 +162,11 @@ deploy_pod vllm-pod-afd-1 2 none head $'POD=AFD_1\n'"$SETTINGS" deploy_pod vllm-pod-afd-0 5 head head $'POD=AFD_0\nLAYOUT=4A1F,4A1F' deploy_pod vllm-pod-afd-1 5 worker worker $'POD=AFD_1\nLAYOUT=4A1F,4A1F' -``` -A one-entry `LAYOUT` with settings is a plain pod with env, e.g. -`deploy_pod vllm-pod 4 head head $'VLLM_ARGS=--max-num-seqs 32'`. +deploy_pod vllm-pod-afd-0 2 head none 'POD=AFD_0' +deploy_pod vllm-pod-afd-1 2 worker none 'POD=AFD_1' +deploy_pod vllm-pod-afd-2 2 none head 'POD=AFD_2' +``` ## 4f. Wait for Running, then for each pod's Ready condition @@ -217,6 +213,6 @@ done Only report the endpoint ready once every pod in the plan is `Ready` -- one pod becoming ready doesn't imply the others did. -Startup markers are keyed on the fixed log names every colocation recipe -uses today -- see `templates/pod.yaml`'s `ready_marker()` for the exact -lines. A recipe using different log names needs that template updated. +Startup markers are keyed on the launcher's fixed log names (`attn.log`, +`ffn.log`) -- see `templates/pod.yaml`'s `ready_marker()` for the exact +lines. diff --git a/.agents/skills/deploy-afd-k8s/resources/resolve-recipe.md b/.agents/skills/deploy-afd-k8s/resources/resolve-recipe.md index e550fe8f9..3c244867e 100644 --- a/.agents/skills/deploy-afd-k8s/resources/resolve-recipe.md +++ b/.agents/skills/deploy-afd-k8s/resources/resolve-recipe.md @@ -6,140 +6,74 @@ under `recipe/gpu/P2pNcclAFDConnector/**/prefill_decode_colocation/` or only; prefill/decode disaggregation recipes are out of scope -- see `SKILL.md` Scope). -Read the script and note, per `vllm serve` block (each backgrounded process -ending `> .log 2>&1 &`): +Every such recipe is a short script on the shared launcher +`recipe/gpu/common/afd_colocation.sh` (see `recipe/gpu/common/README.md`): +the recipe sets `MODEL_PATH`, `AFD_CONNECTOR`, `LAYOUT`, `MODE` and its +`SERVE_ARGS`; the launcher derives devices, DP flags and ports. Never parse +the `vllm serve` command out of the script -- print it: -- `CUDA_VISIBLE_DEVICES`, `--data-parallel-size`, `--tensor-parallel-size` -- the `"afd": {"role": ...}` block in `--additional-config` -- eager vs graph, and every other flag (`--enable-expert-parallel`, - `--max-num-seqs`, `--max-num-batched-tokens`, `--max-model-len`, - `--trust-remote-code`, host/port) - -Record `--max-num-batched-tokens` and report it to the caller later -- it -caps how much prefill work a single step can absorb. - -Resolve `MODEL_ID`: read the script's `MODEL_PATH` default. It is either a -HF Hub id (map the recipe's model directory name to one, e.g. -`deepseek_v2_lite` -> `deepseek-ai/DeepSeek-V2-Lite`; ask if the mapping -isn't obvious) or a **placeholder** path of the form -`/path/model_weights/` -- every colocation recipe uses exactly this -literal `/path/` stand-in today, regardless of what's actually staged on any -given cluster's PVC. Never pass this through verbatim: it names only the -model directory (``), not a real in-container path. Derive the real -`MODEL_ID` as `/`, where `` is the PVC mount -point fixed by `templates/pod.yaml` (`/models`) -- e.g. -`/path/model_weights/Qwen3.5-122B-A10B-FP8` resolves to -`/models/Qwen3.5-122B-A10B-FP8`. Before deploying, confirm that directory -actually exists on the target `PVC_NAME` (if it already exists, `kubectl -exec` into any pod already mounting it, e.g. `ls /models`, or spin up a -short-lived busybox pod mounting the PVC) -- don't rely on the recipe's -claim; a mismatch here fails fast inside the pod with a confusing -HF-repo-id validation error instead of a clear message. - -Tell the caller which form it is, and the resolved value if derived. - -Set `CLIENT_PORT` to the `--port` value on the attention (or sole, for a -single-pod recipe) `vllm serve` block -- `templates/pod.yaml` rebinds this -port off loopback and `vllm-service` (step 4d, see -[services.md](services.md)) exposes it, whatever value is actually there. - -## Read the placement plan -- never ask - -Look at the top-of-file comment block (above the first `vllm serve`/env-var -line). If it says nothing about per-pod placement, this is a plain -single-pod recipe: one pod, `POD` unset, both roles run unconditionally -- -continue to [deploy.md](deploy.md) with no plan. - -If it does describe a placement across multiple pods, read and interpret -that description directly. - -Example of a placement description: -``` -# For pod 1: POD=ATTENTION_0 -# For pod 2: POD=ATTENTION_1, ATTENTION_HEADLESS=1, ATTENTION_DP_START_RANK=2 -# For pod 3: POD=FFN_0 -``` - -Or -``` -"runs -on 3 pods: the first two run ATTENTION_0 and ATTENTION_1 -- the second one -headless, starting at data-parallel rank 2 -- and the third runs FFN_0" +```bash +RECIPE_LIB_PATH=recipe/gpu/common/afd_colocation.sh +AFD_DRY_RUN=1 LAYOUT= POD=AFD_ bash "$RECIPE_SCRIPT_PATH" ``` -Whatever the phrasing, read it and extract, per pod, in the order the -comment presents them: - -- that pod's `POD=` identity -- any further container-env overrides for that pod only (e.g. - `ATTENTION_HEADLESS=1`, `ATTENTION_DP_START_RANK=2`) - -The recipe already states ranks/ordering per pod -- there is nothing to -derive or ask about, regardless of how it's phrased. (A quick -`grep -E '^# For pod [0-9]+: '` is a fine way to spot today's convention at -a glance, but its absence doesn't mean "no plan" -- check the comment block -itself before concluding this is a plain single-pod recipe.) - -For each entry, derive (used by [deploy.md](deploy.md) and -[services.md](services.md)): - -- **k8s Pod name**: `vllm-pod-$(echo "$POD" | tr 'A-Z_' 'a-z-')`, e.g. - `vllm-pod-attention-0`. -- **`afd-attn-node-role`**: `none` unless `POD` starts with `ATTENTION_`, in - which case `worker` if that line's `ATTENTION_HEADLESS` override is `1`, - else `head`. -- **`afd-ffn-node-role`**: symmetric, `none` unless `POD` starts with `FFN_`. -- **`GPU_COUNT`**: the number of devices in the literal - `CUDA_VISIBLE_DEVICES=` of the `vllm serve` block this pod runs -- the one - inside the recipe's branch that tests `$POD` for this pod's role - (`ATTENTION` or `FFN`). Every pod sharing a role shares that literal, so - read it off the script; don't derive it from ranks or DP sizes. - -### `LAYOUT` recipes: `POD=AFD_` - -The `DeepEPM2NAFDConnector` recipes (`graph.sh`, `eager.sh`, `dbo.sh` under -`recipe/gpu/DeepEPM2NAFDConnector//prefill_decode_colocation/`) state -their placement through `LAYOUT`: one pod per comma-separated `AF` -entry, in order, pod `k` running `POD=AFD_`, every pod with the same -`LAYOUT`. Take `LAYOUT` from the caller's settings, else from the script's -`LAYOUT=${LAYOUT:-...}` default; never ask. Each pod derives its own DP flags -(headless, start rank), devices and roles from `LAYOUT` and `POD`, so for -entry `k` with `a` Attention and `f` FFN ranks: - -- **k8s Pod name**: `vllm-pod-afd-`. A one-entry `LAYOUT` (the default - `2A2F`) is a plain single pod: `vllm-pod`, no `POD` needed. -- **`GPU_COUNT`**: `a + f`. The recipe computes `CUDA_VISIBLE_DEVICES` (FFN on - the pod's first `f` GPUs), so there is no literal to read. +prints, per role this pod runs, the `CUDA_VISIBLE_DEVICES=... vllm serve ...` +line. Record `--max-num-batched-tokens` and report it to the caller later -- +it caps how much prefill work a single step can absorb. + +## Model + +Resolve `MODEL_ID` from the script's `MODEL_PATH` default. It is either a +HF Hub id or a **placeholder** path of the form `/path/model_weights/` +-- every colocation recipe uses exactly this literal `/path/` stand-in. +Never pass this through verbatim: derive the real `MODEL_ID` as +`/models/` (`/models` is the PVC mount point fixed by +`templates/pod.yaml`), e.g. `/path/model_weights/Qwen3.5-122B-A10B-FP8` +resolves to `/models/Qwen3.5-122B-A10B-FP8`. Before deploying, confirm that +directory exists on the target `PVC_NAME` (`kubectl exec` into any pod +already mounting it, e.g. `ls /models`, or a short-lived busybox pod +mounting the PVC); a mismatch fails inside the pod with a confusing +HF-repo-id validation error. Tell the caller which form it is, and the +resolved value if derived. + +## Placement plan -- never ask + +The placement is `LAYOUT`: one pod per comma-separated `AF` entry, in +order, pod `k` running `POD=AFD_`, every pod with the same `LAYOUT`. Take +`LAYOUT` from the caller's settings, else from the script's +`LAYOUT=${LAYOUT:-...}` default; never ask. For entry `k` with `a` +Attention and `f` FFN GPUs: + +- **k8s Pod name**: `vllm-pod-afd-`. A one-entry `LAYOUT` (e.g. `2A2F`) + is a plain single pod: `vllm-pod`, no `POD` needed. +- **`GPU_COUNT`**: `a + f` (FFN on the pod's first `f` GPUs). - **`afd-attn-node-role`**: `none` when `a` is 0, `head` on the first entry with `a > 0`, `worker` on later ones. **`afd-ffn-node-role`**: the same - with `f`. For example `2A0F,0A2F` is `head none` and `none head`, and - `1A1F,1A1F` is `head head` and `worker worker`. + with `f`. For example `2A0F,0A2F` is `head none` and `none head`, + `1A1F,1A1F` is `head head` and `worker worker`, and `2A0F,2A0F,0A2F` is + `head none`, `worker none` and `none head`. - **Container env**: `POD=AFD_` and `LAYOUT=`, plus every setting - the caller gives (`MODEL_PATH`, `ATTN_TP`, `VLLM_ARGS`, `ATTN_ARGS`, - `FFN_ARGS`, `DECODE_BENCH`, `M2N_EXTRA_CONFIG`, NCCL variables), the same on - every pod. Leave `HEAD_HOST`, `ATTN_HEAD_HOST` and `FFN_HEAD_HOST` unset: - the defaults name the Services from [services.md](services.md). With - `DECODE_BENCH`, also set `AFD_REPO=/opt/afd-plugin` (the repository in - `docker/Dockerfile.ci` images), since the connector lives in - `tools/benchmarks` and the recipe runs from `/work`. -- **Ports** are literals: `--port 18000` on the Attention `vllm serve` - (`CLIENT_PORT`), the AFD rendezvous `"port": 6239` in the `afd` JSON, and - the DP RPC ports as the third argument of the two `role_dp` calls (`13355` - FFN, `13345` Attention). + the caller gives (`MODEL_PATH`, `ATTN_TP`, `FFN_TP`, `VLLM_ARGS`, + `ATTN_ARGS`, `FFN_ARGS`, `DECODE_BENCH`, `M2N_EXTRA_CONFIG`, NCCL + variables), the same on every pod. Leave `HEAD_HOST`, `ATTN_HEAD_HOST` and + `FFN_HEAD_HOST` unset: the defaults name the Services from + [services.md](services.md). With `DECODE_BENCH`, also set + `AFD_REPO=/opt/afd-plugin` (the repository in `docker/Dockerfile.ci` + images), since the connector lives in `tools/benchmarks` and the recipe + runs from `/work`. - Whether the connector accepts the placement is checked at startup, not here; see the connector guide's supported layouts. -## GPU_COUNT for a plain single-pod recipe +## Ports -The union of every literal `CUDA_VISIBLE_DEVICES=` in the script: +The launcher fixes them on one line (`ATTN_PORT=18000 FFN_PORT=18001 +AFD_PORT=6239 ATTN_DP_RPC_PORT=13345 FFN_DP_RPC_PORT=13355`); read them from +it rather than assuming: ```bash -GPU_COUNT="$(grep -oE 'CUDA_VISIBLE_DEVICES=[0-9,]+' "$RECIPE_SCRIPT_PATH" \ - | cut -d= -f2 | tr ',' '\n' | sort -u | wc -l)" +lib_port() { grep -oE "\b$1=[0-9]+" "$RECIPE_LIB_PATH" | head -1 | cut -d= -f2; } +CLIENT_PORT=$(lib_port ATTN_PORT) # the Attention API, exposed by vllm-service ``` -A one-entry `LAYOUT` recipe has no literal to read: its `GPU_COUNT` is -`a + f` of that entry, as above. - This step is the only place that parses the recipe: -[deploy.md](deploy.md) uses the `GPU_COUNT` (and plan) recorded here as-is. +[deploy.md](deploy.md) uses the `GPU_COUNT` and plan recorded here as-is. diff --git a/.agents/skills/deploy-afd-k8s/resources/services.md b/.agents/skills/deploy-afd-k8s/resources/services.md index a3bc18648..5f5302dfc 100644 --- a/.agents/skills/deploy-afd-k8s/resources/services.md +++ b/.agents/skills/deploy-afd-k8s/resources/services.md @@ -24,32 +24,28 @@ CLIENT_PORT="$CLIENT_PORT" envsubst '${CLIENT_PORT}' \ Internal only, carries AFD rendezvous traffic to whichever pod holds `afd-ffn-node-role: head`. Must exist before any multi-pod deploy, since the -recipe's own `AFD_CONNECTOR_HOST` default now points at this name directly. -Traffic to a headless Service goes straight to the pod IP, and there's one -rendezvous port, so a single static port entry is enough -- read -`AFD_CONNECTOR_PORT` off the recipe's own shell default: +launcher's default rendezvous host points at this name directly. Traffic to +a headless Service goes straight to the pod IP, and there's one rendezvous +port, so a single static port entry is enough (`lib_port` from +[resolve-recipe.md](resolve-recipe.md#ports)): ```bash -AFD_CONNECTOR_PORT="$(grep -oE 'AFD_CONNECTOR_PORT:-[0-9]+' "$RECIPE_SCRIPT_PATH" | grep -oE '[0-9]+')" -AFD_CONNECTOR_PORT="$AFD_CONNECTOR_PORT" envsubst '${AFD_CONNECTOR_PORT}' \ +AFD_CONNECTOR_PORT="$(lib_port AFD_PORT)" envsubst '${AFD_CONNECTOR_PORT}' \ < templates/service-ffn-p2p.yaml | kubectl apply -f - ``` ## `vllm-attn-dp-service` / `vllm-ffn-dp-service` -- iff that role is split -Only when **more than one** plan entry's `POD` carries the same role (i.e. -that role is split across pods, not just placed in its own single pod). -For a `LAYOUT` recipe, a role is split when more than one entry has ranks of -it (both roles in `1A1F,1A1F`, neither in `2A0F,0A2F`); read its ports (and -the rendezvous port) as [resolve-recipe.md](resolve-recipe.md) describes for -those recipes rather than with the `*_PORT:-N` greps below. -Carries that role's DP-RPC coordination traffic to its head pod. One per -split role, generic over role via `templates/service-dp-coord.yaml`: +Only when **more than one** `LAYOUT` entry has GPUs of that role (both +roles in `1A1F,1A1F`, Attention in `2A0F,2A0F,0A2F`, neither in +`2A0F,0A2F`). Carries that role's DP-RPC coordination traffic to its head +pod. One per split role, generic over role via +`templates/service-dp-coord.yaml`: ```bash -TEMPLATE_ROLE=attn -TEMPLATE_PORT="$(grep -oE 'ATTENTION_DP_RPC_PORT:-[0-9]+' "$RECIPE_SCRIPT_PATH" | grep -oE '[0-9]+')" -TEMPLATE_ROLE="$TEMPLATE_ROLE" TEMPLATE_PORT="$TEMPLATE_PORT" \ +TEMPLATE_ROLE=attn TEMPLATE_PORT="$(lib_port ATTN_DP_RPC_PORT)" \ + envsubst '${TEMPLATE_ROLE} ${TEMPLATE_PORT}' < templates/service-dp-coord.yaml | kubectl apply -f - +TEMPLATE_ROLE=ffn TEMPLATE_PORT="$(lib_port FFN_DP_RPC_PORT)" \ envsubst '${TEMPLATE_ROLE} ${TEMPLATE_PORT}' < templates/service-dp-coord.yaml | kubectl apply -f - ``` diff --git a/.agents/skills/deploy-afd-k8s/templates/pod.yaml b/.agents/skills/deploy-afd-k8s/templates/pod.yaml index 561f16f18..f43a0694e 100644 --- a/.agents/skills/deploy-afd-k8s/templates/pod.yaml +++ b/.agents/skills/deploy-afd-k8s/templates/pod.yaml @@ -33,33 +33,19 @@ spec: - -c - | set -euo pipefail - RECIPE_SCRIPT=/recipe/recipe.sh - mkdir -p /work/bin - # Recipes only ever call `uv run `; drop the `run` and exec - # the rest directly instead of invoking uv at all. - printf '#!/bin/bash\nshift\nexec "$@"\n' > /work/bin/uv - chmod +x /work/bin/uv - export PATH="/work/bin:$PATH" - - echo "=== launching $(basename "$RECIPE_SCRIPT") (POD=${POD:-} attn=${AFD_ATTN_NODE_ROLE} ffn=${AFD_FFN_NODE_ROLE}) ===" - - # Rebind the client-facing port off loopback: every `--host - # 127.0.0.1` on the branch is an API-server flag (the recipe binds - # :${TEMPLATE_CLIENT_PORT} on 127.0.0.1), but the Service needs - # 0.0.0.0. - RECIPE_BASENAME="$(basename "$RECIPE_SCRIPT" .sh)" - PATCHED_SCRIPT="/work/${RECIPE_BASENAME}.patched.sh" - sed 's/--host 127\.0\.0\.1/--host 0.0.0.0/' "$RECIPE_SCRIPT" > "$PATCHED_SCRIPT" + echo "=== launching recipe.sh (POD=${POD:-} attn=${AFD_ATTN_NODE_ROLE} ffn=${AFD_FFN_NODE_ROLE}) ===" + # The recipe sources the launcher from AFD_RECIPE_LIB and binds its + # API server to SERVE_HOST (0.0.0.0, so the Service can reach it). cd /work - bash "$PATCHED_SCRIPT" & + bash /recipe/recipe.sh & RECIPE_PID=$! echo "$RECIPE_PID" > /work/recipe.pid # Which log(s) this pod will actually produce, and this pod's own # readiness marker per log -- both keyed on this pod's node-role # labels, not re-derived from $POD. Fixed to the attn.log/ffn.log - # names every colocation recipe uses today. + # names of the colocation launcher. LOGS="" [ "$AFD_ATTN_NODE_ROLE" != "none" ] && LOGS="$LOGS attn.log" [ "$AFD_FFN_NODE_ROLE" != "none" ] && LOGS="$LOGS ffn.log" @@ -148,6 +134,10 @@ spec: value: "${TEMPLATE_ATTN_ROLE}" - name: AFD_FFN_NODE_ROLE value: "${TEMPLATE_FFN_ROLE}" + - name: AFD_RECIPE_LIB + value: "/recipe/afd_colocation.sh" + - name: SERVE_HOST + value: "0.0.0.0" # ${TEMPLATE_EXTRA_ENV} - name: HF_TOKEN valueFrom: diff --git a/.agents/skills/review-pr/references/features/disaggregation-topologies.md b/.agents/skills/review-pr/references/features/disaggregation-topologies.md index b5d920306..c23a93b38 100644 --- a/.agents/skills/review-pr/references/features/disaggregation-topologies.md +++ b/.agents/skills/review-pr/references/features/disaggregation-topologies.md @@ -14,10 +14,11 @@ recipes under `recipe/gpu/` and `recipe/npu/`. | Topology | Where it lives | | --- | --- | | Baseline (native DP/TP, no AFD) | `baseline-graph` E2E scenarios; native DP4/TP1/EP4 | -| CUDA NCCL P2P, colocation | `recipe/gpu/P2pNcclAFDConnector/deepseek_v2_lite/prefill_decode_colocation/` (2A2F and 4A4F, eager/graph, DBO variants, plus `multipod_2a_2f_graph_dbo_dp1tp2.sh`) | +| CUDA NCCL P2P, colocation | `recipe/gpu/P2pNcclAFDConnector/deepseek_v2_lite/prefill_decode_colocation/` (`dbo.sh`, `eager_dbo.sh`; any `LAYOUT`, `ATTN_TP`, `FFN_TP`, e.g. 2A2F, 4A4F, `2A0F,0A2F`) | | CUDA NCCL P2P, disaggregation | `recipe/gpu/P2pNcclAFDConnector/deepseek_v2_lite/prefill_decode_disaggregation/` (2P1A1F eager/graph DBO) | -| CUDA NCCL P2P, Qwen3.5 multi-pod | `recipe/gpu/P2pNcclAFDConnector/qwen3_5_122b_a10b_fp8/prefill_decode_colocation/multipod_2a_2a_2f_graph.sh` | -| CUDA DeepEP M2N (routed experts) | `recipe/gpu/DeepEPM2NAFDConnector/{deepseek_v2_lite,qwen3_30b_a3b,qwen3_235b_a22b_fp8}/` (2A2F, and 4A2F for the 235B FP8 case) | +| CUDA NCCL P2P, Qwen3.5 multi-pod | `recipe/gpu/P2pNcclAFDConnector/qwen3_5_122b_a10b_fp8/prefill_decode_colocation/graph.sh` (default `LAYOUT=2A0F,2A0F,0A2F`) | +| CUDA DeepEP M2N (routed experts) | `recipe/gpu/DeepEPM2NAFDConnector/{deepseek_v2_lite,qwen3_30b_a3b,qwen3_235b_a22b_fp8}/` (graph/eager/dbo, any `LAYOUT`) | +| GPU colocation launcher | `recipe/gpu/common/afd_colocation.sh`, shared by every GPU colocation recipe | | NPU CAM P2P disaggregation | `recipe/npu/CAMP2pAFDConnector/deepseek_v3_2/` (attention/FFN/baseline scripts) | | NPU async CAM | `recipe/npu/CAMAsyncAFDConnector/deepseek_v3_2/v0_26_accuracy/` and `deepseek_v4_flash/` (DSpark A/F) | | Gate matrices | `tests/e2e/models/*/` scenario lists; `docs/design/module/e2e_testing.md` | diff --git a/.agents/skills/run-vllm-bench-k8s/SKILL.md b/.agents/skills/run-vllm-bench-k8s/SKILL.md index b61c4f99f..d0a3a66fc 100644 --- a/.agents/skills/run-vllm-bench-k8s/SKILL.md +++ b/.agents/skills/run-vllm-bench-k8s/SKILL.md @@ -42,10 +42,10 @@ folder as the list of possible recipes. Follow **`deploy-afd-k8s`** with `RECIPE_SCRIPT_PATH`. It returns `MODEL_ID`, `PVC_NAME`, `GPU_COUNT`, the serving node, the recipe's -`--max-num-batched-tokens`, and an endpoint at `http://vllm-service:18305`. +`--max-num-batched-tokens`, and an endpoint at `http://vllm-service:${CLIENT_PORT}` (18000). `vllm bench serve` is pointed at that endpoint with `HOST=vllm-service` -`PORT=18305`, and `MODEL_PATH=${MODEL_ID}` -- it must name exactly what the +`PORT=18000` (the launcher's Attention port), and `MODEL_PATH=${MODEL_ID}` -- it must name exactly what the pod serves, since the model id is sent in each request and used to load the tokenizer. @@ -195,7 +195,7 @@ spec: conc="${point##*:}" echo "=== sweep point rate=${rate} concurrency=${conc} ===" MODEL_PATH="${MODEL_ID}" \ - HOST=vllm-service PORT=18305 \ + HOST=vllm-service PORT=18000 \ DATASET_NAME="${DATASET_NAME}" \ INPUT_LEN="${INPUT_LEN}" OUTPUT_LEN="${OUTPUT_LEN}" \ NUM_PROMPTS="${NUM_PROMPTS}" \ diff --git a/recipe/README.md b/recipe/README.md index 2c88fdd88..c98bea1cf 100644 --- a/recipe/README.md +++ b/recipe/README.md @@ -25,6 +25,7 @@ Recipes are organized by hardware backend, connector, and model: | GPU (SM90+) | `DeepEPM2NAFDConnector` | DeepSeek-V2-Lite | Decode | Deployment examples matching the checked-in E2E layouts; hardware qualification is separate | [Launch examples](gpu/DeepEPM2NAFDConnector/deepseek_v2_lite/README.md) | | GPU (SM90+ for `DeepEPM2NAFDConnector`) | `DeepEPM2NAFDConnector` or `P2pNcclAFDConnector` | Qwen3-30B-A3B | Prefill / decode; decode-only benchmark | Graph, eager and DBO scripts for any `LAYOUT` (default 2A2F; multi-pod such as `2A0F,0A2F` or `1A1F,1A1F`), the layouts of the Qwen3 MoE DeepEP M2N E2E cases | [Graph, eager, DBO](gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/README.md) | | GPU (SM90+ for `DeepEPM2NAFDConnector`) | `DeepEPM2NAFDConnector` or `P2pNcclAFDConnector` | Qwen3-235B-A22B-FP8 | Prefill / decode; decode-only benchmark | FP8-payload graph, eager and DBO scripts for any `LAYOUT` (2A2F, 4A2F, 2 x 4A1F, ...); H200-class FFN GPUs | [Graph, eager, DBO](gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/README.md) | +| GPU | `P2pNcclAFDConnector` | Qwen3.5-122B-A10B-FP8 | Decode | Graph script for any `LAYOUT` (default 3 pods, `2A0F,2A0F,0A2F`) | [Graph](gpu/P2pNcclAFDConnector/qwen3_5_122b_a10b_fp8/prefill_decode_colocation/graph.sh) | | Ascend NPU | `CAMP2pAFDConnector` | DeepSeek-V3.2 | Decode | Historical v0.19.1rc1 recipe; not qualified for the 0.32 target | [Synchronous decode](npu/CAMP2pAFDConnector/deepseek_v3_2/README.md) | | Ascend NPU | `CAMAsyncAFDConnector` | DeepSeek-V3.2 | Prefill / decode | Historical v0.19.1rc1/v0.26 recipes; standard V3.2 is unsupported by the current Async CAM path | [Async CAM](npu/CAMAsyncAFDConnector/deepseek_v3_2/README.md) | | Ascend NPU | `CAMAsyncAFDConnector` | DeepSeek-V4-Flash W4A8/W8A8 DSpark | Separate P / decode, or colocated | v0.32; eager A TP1, F layered eager/FULL. W4A8 PD: native full-draft P TP4, A DP2, F EP4. W8A8 A/F: A DP2, F EP8; PD not qualified | [DSpark A/F](npu/CAMAsyncAFDConnector/deepseek_v4_flash/README.md) | @@ -41,3 +42,8 @@ directory should contain a `README.md` describing prerequisites, topology, launch commands, validation steps, and limitations. Keep model-specific launch scripts, configuration files, and result images in the same directory so the recipe can be moved or linked as one unit. + +GPU `prefill_decode_colocation` recipes are short scripts on the shared +[colocation launcher](gpu/common/README.md): a recipe sets its model, +connector, default `LAYOUT` and `vllm serve` flags, and the launcher handles +placement across GPUs and pods. Start a new one from an existing recipe. diff --git a/recipe/gpu/DeepEPM2NAFDConnector/README.md b/recipe/gpu/DeepEPM2NAFDConnector/README.md index 29f8ef848..01b411de6 100644 --- a/recipe/gpu/DeepEPM2NAFDConnector/README.md +++ b/recipe/gpu/DeepEPM2NAFDConnector/README.md @@ -28,11 +28,9 @@ Every model directory holds the same three scripts in └── qwen3_235b_a22b_fp8/ # FP8 payload (fp8_block128_fp32) ``` -Each script is self-contained, so it can be copied or mounted on its own (the -`deploy-afd-k8s` skill mounts the single file). Above the line -`The rest is the same in every DeepEPM2NAFDConnector recipe` a script sets its -model path, layout default, tuning flags (`TUNING`), connector extras and -mode flags (`MODE`); everything below that line is identical in all of them. +Each script sets the model, its tuning and mode flags, and runs the shared +[colocation launcher](../common/README.md), which documents the layout, +every setting, and how to run a recipe directly or in a pod. ## Prerequisites @@ -59,115 +57,40 @@ use the environment that holds afdEP-comm and DeepGEMM. ## Layout -`LAYOUT` lists the pods in order, one `AF` entry each: `a` Attention -ranks and `f` FFN ranks. Inside a pod, the FFN ranks take the first `f` GPUs -and the Attention ranks the next `a` GPUs; every pod runs the same script with -the same `LAYOUT` and its own `POD=AFD_` (`k` counts from 0; one pod needs -no `POD`). - -| `LAYOUT` | Pods | Per pod | -| --- | --- | --- | -| `2A2F` (default) | 1 | FFN DP2 on GPUs 0,1, Attention DP2 on GPUs 2,3 | -| `4A2F` | 1 | FFN on GPUs 0,1, Attention DP4 on GPUs 2-5 | -| `2A0F,0A2F` | 2 | pod 0: Attention DP2; pod 1: FFN DP2 | -| `1A1F,1A1F` | 2 | FFN on GPU 0, Attention on GPU 1 | -| `4A1F,4A1F` | 2 | FFN on GPU 0, Attention on GPUs 1-4 | - -`ATTN_TP` splits each pod's Attention GPUs into TP groups (it must divide each -pod's `a`), so `ATTN_TP=2` runs `2A2F`'s Attention as one DP1 x TP2 engine; -`num_attention_ranks` still counts GPUs. The FFN always runs TP 1, and its DP -rank (its EP rank) decides which experts it loads. - -The script derives every DP flag from `LAYOUT` and `POD`. A role held by one -pod gets only `--data-parallel-size`; a role spread over pods is led by its -first pod, and the others start headless from their start rank. The -connector checks the placement at startup and names the pods' shapes when it -rejects one; the +`LAYOUT` (default `2A2F`, one pod of 4 GPUs) lists the pods, one `AF` +entry each; see [Layout](../common/README.md#layout) and +[Settings](../common/README.md#settings). Examples for this connector: +`4A2F` (one 6-GPU node), `2A0F,0A2F` and `1A1F,1A1F` (two pods), `4A1F,4A1F` +(two 5-GPU pods). `ATTN_TP=2` runs `2A2F`'s Attention as one DP1 x TP2 engine; +the FFN runs TP 1, and its DP rank (its EP rank) decides which experts it +loads. The connector checks the placement at startup and names the pods' +shapes when it rejects one; the [supported layouts](../../../docs/gpu/DEEPEP_M2N_CONNECTOR_USER_GUIDE.md#supported-layouts) table lists which placements it accepts (for example `4A0F,0A2F` is not, its pods differ in size). -## Recipe settings - -Every setting is an environment variable; the defaults reproduce the -published recipe. - -| Variable | Default | Meaning | -| --- | --- | --- | -| `MODEL_PATH` | `/path/model_weights/` | Weights directory or HF id. | -| `LAYOUT` | `2A2F` | Pods and their ranks; see [Layout](#layout). | -| `POD` | `AFD_0` | This pod's entry in `LAYOUT`, `AFD_`. | -| `ATTN_TP` | `1` | Attention TP size. | -| `CUDAGRAPH_CAPTURE_SIZES` | per model | Ascending decode graph sizes (`graph.sh`, `dbo.sh`); the largest is the maximum capture size. | -| `VLLM_ARGS` | empty | Extra `vllm serve` flags for both roles. | -| `ATTN_ARGS` / `FFN_ARGS` | empty | Extra `vllm serve` flags for one role. | -| `DECODE_BENCH` | empty | Turns on the decode-only benchmark; see [Decode-only benchmark](#decode-only-benchmark). | -| `AFD_REPO` | current directory | Repository root that holds `tools/benchmarks` (with `DECODE_BENCH`). | -| `M2N_EXTRA_CONFIG` | empty | Extra `connector_extra_config` members, for example `'"expert_num_sms": 130'`. | -| `HEAD_HOST` | `127.0.0.1` on one pod | Address of pod 0 when it leads both roles (multi-pod, no Kubernetes). | -| `ATTN_HEAD_HOST` / `FFN_HEAD_HOST` | `HEAD_HOST` | Address of the first pod with Attention / FFN ranks, when they differ. The FFN head also hosts the AFD rendezvous. | -| `AFD_CONNECTOR` | `DeepEPM2NAFDConnector` | `P2pNcclAFDConnector` runs the same layout on NCCL P2P (used by the P2P recipe entry points). | - -`VLLM_ARGS`, `ATTN_ARGS`, `FFN_ARGS` and any arguments after the script name -come after the recipe's own flags, so a repeated flag overrides its default -(argparse keeps the last value). They are split on whitespace without shell -quoting, so pass JSON values without spaces. For example: - -```bash -# 512 sequences per Attention rank, no prefix caching on Attention -VLLM_ARGS="--max-num-seqs 512" ATTN_ARGS="--no-enable-prefix-caching" \ - bash recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/graph.sh -# Count graph and eager Attention steps (vLLM disables the stats logger with -# more than one API server) -ATTN_ARGS="--api-server-count 1 --cudagraph-metrics" bash .../graph.sh -``` - -Raise `CUDAGRAPH_CAPTURE_SIZES` together with `--max-num-seqs`, so decode -batches stay in a captured size. - -## Running on one host - -Run a script from the repository root: +## Running ```bash export MODEL_PATH=/path/model_weights/Qwen3-30B-A3B bash recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/graph.sh ``` -The script starts the FFN server and then the Attention server in the -background, writes `ffn.log` and `attn.log` into the current directory, and -waits. Wait for `Application startup complete` in `attn.log`, then send -traffic to the Attention server on port `18000`. The first start on a host -compiles transport and DeepGEMM kernels, so it is slower than later starts. - -## Running on several pods - -Start the same script on every pod with the same settings and the pod's own -`POD`. Without Kubernetes, point the pods at the leaders' addresses: - -```bash -# 2 x 1A1F: pod 0 leads both roles -LAYOUT=1A1F,1A1F POD=AFD_0 HEAD_HOST= bash .../graph.sh # pod 0 -LAYOUT=1A1F,1A1F POD=AFD_1 HEAD_HOST= bash .../graph.sh # pod 1 - -# 2A0F,0A2F: Attention on pod 0, FFN on pod 1 -LAYOUT=2A0F,0A2F POD=AFD_0 ATTN_HEAD_HOST= FFN_HEAD_HOST= bash .../graph.sh -LAYOUT=2A0F,0A2F POD=AFD_1 ATTN_HEAD_HOST= FFN_HEAD_HOST= bash .../graph.sh -``` - -Send traffic to port `18000` of the first pod with Attention ranks; the other -pods serve no requests. Pass the fabric's NCCL settings (for example -`NCCL_IB_HCA`) through the environment. The roles can finish loading minutes -apart, so the recipes raise `afd_process_group_timeout_s` and -`--cpu-distributed-timeout-seconds` to 3600 s. +Send traffic to the Attention server on port `18000` once `attn.log` prints +`Application startup complete`. The first start on a host compiles transport +and DeepGEMM kernels, so it is slower than later starts. Raise +`CUDAGRAPH_CAPTURE_SIZES` together with `--max-num-seqs`, so decode batches +stay in a captured size. Across pods, pass the fabric's NCCL settings (for +example `NCCL_IB_HCA`) through the environment; see +[Running directly](../common/README.md#running-directly) for hosts without +Kubernetes. ## Deploying on Kubernetes The **`deploy-afd-k8s`** skill (`.agents/skills/deploy-afd-k8s/SKILL.md`) -deploys a recipe as one pod per `LAYOUT` entry, sizes each pod's GPUs from its -entry, and creates the Services the recipe's default host names point at -(`vllm-ffn-p2p-service`, `vllm-attn-dp-service`, `vllm-ffn-dp-service`). The -image must contain afdEP-comm and DeepGEMM. +deploys a recipe as one pod per `LAYOUT` entry; see +[Running in a pod](../common/README.md#running-in-a-pod). The image must +contain afdEP-comm and DeepGEMM. ## Checking a deployment @@ -207,8 +130,8 @@ mode; see [`tools/benchmarks/README.md`](../../../tools/benchmarks/README.md). ## Configuration notes - Both roles set `compute_gate_on_attention: true` and run with - `--enable-expert-parallel` and `--enable-ep-weight-filter`, so each FFN rank - reads only its own experts. + `--enable-expert-parallel` and `--enable-ep-weight-filter` (in each + recipe's `SERVE_ARGS`), so each FFN rank reads only its own experts. - The FFN keeps vLLM's default `--moe-backend auto`, which the plugin pins to `triton` so the expert weights stay in the layout the connector reads. - The connector derives its dispatch buckets from the CUDA graph capture diff --git a/recipe/gpu/DeepEPM2NAFDConnector/deepseek_v2_lite/prefill_decode_colocation/dbo.sh b/recipe/gpu/DeepEPM2NAFDConnector/deepseek_v2_lite/prefill_decode_colocation/dbo.sh index c7e17b515..dbb24388f 100644 --- a/recipe/gpu/DeepEPM2NAFDConnector/deepseek_v2_lite/prefill_decode_colocation/dbo.sh +++ b/recipe/gpu/DeepEPM2NAFDConnector/deepseek_v2_lite/prefill_decode_colocation/dbo.sh @@ -3,96 +3,17 @@ # SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project # # DeepSeek-V2-Lite (BF16) on DeepEPM2NAFDConnector, FULL_DECODE_ONLY graphs + DBO. -# Placement: one pod per comma-separated LAYOUT entry (default 2A2F, one pod of -# 4 GPUs). AF puts f FFN ranks on the pod's first f GPUs and a Attention -# ranks on the next a GPUs, e.g. LAYOUT=2A0F,0A2F or LAYOUT=1A1F,1A1F. -# For pod k (k = 0, 1, ...): POD=AFD_, with the same LAYOUT on every pod. -# Settings: recipe/gpu/DeepEPM2NAFDConnector/README.md#recipe-settings +# Placement: one pod per comma-separated LAYOUT entry (default 2A2F), pod k +# running POD=AFD_ with the same LAYOUT; FFN on each pod's first GPUs. +# Settings and running in a pod: recipe/gpu/common/README.md +# shellcheck disable=SC1090,SC2034 # the sourced launcher reads these variables -MODEL_PATH=${MODEL_PATH:-/path/model_weights/DeepSeek-V2-Lite} +AFD_CONNECTOR=${AFD_CONNECTOR:-DeepEPM2NAFDConnector} LAYOUT=${LAYOUT:-2A2F} -TUNING=(--max-num-seqs 64 --max-num-batched-tokens 4096 --max-model-len 4096) -CONNECTOR_CONFIG=${M2N_EXTRA_CONFIG:-} # connector_extra_config members; BF16 payload -CUDAGRAPH_CAPTURE_SIZES=${CUDAGRAPH_CAPTURE_SIZES:-8,16,32,64} # ascending -MODE=(--max-cudagraph-capture-size "${CUDAGRAPH_CAPTURE_SIZES##*,}" - --compilation-config "{\"cudagraph_mode\": \"FULL_DECODE_ONLY\", \"cudagraph_capture_sizes\": [${CUDAGRAPH_CAPTURE_SIZES}]}" - --enable-dbo --dbo-decode-token-threshold 2 --dbo-prefill-token-threshold 8) - -# ---- The rest is the same in every DeepEPM2NAFDConnector recipe. ---- -export VLLM_USE_V2_MODEL_RUNNER=1 -ATTN_TP=${ATTN_TP:-1} -AFD_CONNECTOR=${AFD_CONNECTOR:-DeepEPM2NAFDConnector} GATE_ON_ATTENTION=true -# P2P routes on the FFN and takes no connector_extra_config. -[[ $AFD_CONNECTOR == P2pNcclAFDConnector ]] && GATE_ON_ATTENTION=false CONNECTOR_CONFIG= -[[ -n $CONNECTOR_CONFIG ]] && CONNECTOR_CONFIG=", \"connector_extra_config\": {${CONNECTOR_CONFIG}}" -IFS=, read -ra PODS <<<"$LAYOUT" -K=${POD:-AFD_0} && K=${K#AFD_} -if [[ ! $LAYOUT =~ ^[0-9]+A[0-9]+F(,[0-9]+A[0-9]+F)*$ || ! $K =~ ^[0-9]+$ ]] || ((K >= ${#PODS[@]})); then - echo "need LAYOUT=AF[,...] and POD=AFD_ in it: LAYOUT=$LAYOUT POD=${POD:-}" >&2 - exit 1 -fi -NUM_A=() NUM_F=() -for p in "${PODS[@]}"; do NUM_A+=("${p%%A*}") && p=${p#*A} && NUM_F+=("${p%F}"); done -((${#PODS[@]} > 1)) || HEAD_HOST=${HEAD_HOST:-127.0.0.1} - -# Extra `vllm serve` flags go last, so they override the defaults. -read -ra ARGS <<<"${VLLM_ARGS:-}" && ARGS+=("$@") -read -ra ATTN_EXTRA <<<"${ATTN_ARGS:-}" -read -ra FFN_EXTRA <<<"${FFN_ARGS:-}" -if [[ -n ${DECODE_BENCH:-} ]]; then # tools/benchmarks is imported from the repository - ATTN_EXTRA=(--kv-transfer-config "{\"kv_connector\": \"AFDDecodeBenchConnector\", \"kv_connector_module_path\": \"tools.benchmarks.decode_bench\", \"kv_role\": \"kv_both\", \"kv_connector_extra_config\": ${DECODE_BENCH}}" "${ATTN_EXTRA[@]}") - export PYTHONPATH=${AFD_REPO:-$PWD}${PYTHONPATH:+:$PYTHONPATH} -fi - -# role_dp TP DP_HOST DP_RPC_PORT ... -> TOTAL and this pod's DP -# flags. A role spread over pods is led by its first pod; the others run -# headless from their start rank (the leader must omit it, or vLLM selects -# hybrid load balancing and rejects the headless peers). -role_dp() { - local tp=$1 host=$2 port=$3 i pods=0 start=0 lead=-1 - shift 3 - local n=("$@") - TOTAL=0 - for i in "${!n[@]}"; do - ((n[i])) || continue - ((n[i] % tp)) && { echo "TP $tp does not divide ${n[i]} ranks on pod $i" >&2; exit 1; } - ((lead < 0)) && lead=$i - ((i < K)) && start=$((start + n[i])) - pods=$((pods + 1)) TOTAL=$((TOTAL + n[i])) - done - DP=(--tensor-parallel-size "$tp" --data-parallel-size $((TOTAL / tp))) - ((pods > 1)) && DP+=(--data-parallel-size-local $((n[K] / tp)) - --data-parallel-address "$host" --data-parallel-rpc-port "$port") - ((pods > 1 && K != lead)) && DP+=(--headless --data-parallel-start-rank $((start / tp))) -} -role_dp 1 "${FFN_HEAD_HOST:-${HEAD_HOST:-vllm-ffn-dp-service}}" 13355 "${NUM_F[@]}" -FFN_DP=("${DP[@]}") NUM_FFN=$TOTAL -role_dp "$ATTN_TP" "${ATTN_HEAD_HOST:-${HEAD_HOST:-vllm-attn-dp-service}}" 13345 "${NUM_A[@]}" -ATTN_DP=("${DP[@]}") NUM_ATTN=$TOTAL - -afd() { # The AFD rendezvous runs on the pod that holds FFN rank 0. - cat < ffn.log 2>&1 & -fi -if ((A)); then - CUDA_VISIBLE_DEVICES=$(seq "$F" $((F + A - 1)) | paste -sd, -) vllm serve "$MODEL_PATH" \ - --additional-config "$(afd attention)" "${COMMON[@]}" "${ATTN_DP[@]}" \ - "${ARGS[@]}" "${ATTN_EXTRA[@]}" --port 18000 > attn.log 2>&1 & -fi -wait +source "${AFD_RECIPE_LIB:-$(dirname "${BASH_SOURCE[0]}")/../../../common/afd_colocation.sh}" +afd_colocation_serve "$@" diff --git a/recipe/gpu/DeepEPM2NAFDConnector/deepseek_v2_lite/prefill_decode_colocation/eager.sh b/recipe/gpu/DeepEPM2NAFDConnector/deepseek_v2_lite/prefill_decode_colocation/eager.sh index 82fc56903..53c3bc605 100644 --- a/recipe/gpu/DeepEPM2NAFDConnector/deepseek_v2_lite/prefill_decode_colocation/eager.sh +++ b/recipe/gpu/DeepEPM2NAFDConnector/deepseek_v2_lite/prefill_decode_colocation/eager.sh @@ -3,93 +3,16 @@ # SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project # # DeepSeek-V2-Lite (BF16) on DeepEPM2NAFDConnector, eager mode. -# Placement: one pod per comma-separated LAYOUT entry (default 2A2F, one pod of -# 4 GPUs). AF puts f FFN ranks on the pod's first f GPUs and a Attention -# ranks on the next a GPUs, e.g. LAYOUT=2A0F,0A2F or LAYOUT=1A1F,1A1F. -# For pod k (k = 0, 1, ...): POD=AFD_, with the same LAYOUT on every pod. -# Settings: recipe/gpu/DeepEPM2NAFDConnector/README.md#recipe-settings +# Placement: one pod per comma-separated LAYOUT entry (default 2A2F), pod k +# running POD=AFD_ with the same LAYOUT; FFN on each pod's first GPUs. +# Settings and running in a pod: recipe/gpu/common/README.md +# shellcheck disable=SC1090,SC2034 # the sourced launcher reads these variables -MODEL_PATH=${MODEL_PATH:-/path/model_weights/DeepSeek-V2-Lite} +AFD_CONNECTOR=${AFD_CONNECTOR:-DeepEPM2NAFDConnector} LAYOUT=${LAYOUT:-2A2F} -TUNING=(--max-num-seqs 64 --max-num-batched-tokens 4096 --max-model-len 4096) -CONNECTOR_CONFIG=${M2N_EXTRA_CONFIG:-} # connector_extra_config members; BF16 payload -MODE=(--enforce-eager) - -# ---- The rest is the same in every DeepEPM2NAFDConnector recipe. ---- -export VLLM_USE_V2_MODEL_RUNNER=1 -ATTN_TP=${ATTN_TP:-1} -AFD_CONNECTOR=${AFD_CONNECTOR:-DeepEPM2NAFDConnector} GATE_ON_ATTENTION=true -# P2P routes on the FFN and takes no connector_extra_config. -[[ $AFD_CONNECTOR == P2pNcclAFDConnector ]] && GATE_ON_ATTENTION=false CONNECTOR_CONFIG= -[[ -n $CONNECTOR_CONFIG ]] && CONNECTOR_CONFIG=", \"connector_extra_config\": {${CONNECTOR_CONFIG}}" -IFS=, read -ra PODS <<<"$LAYOUT" -K=${POD:-AFD_0} && K=${K#AFD_} -if [[ ! $LAYOUT =~ ^[0-9]+A[0-9]+F(,[0-9]+A[0-9]+F)*$ || ! $K =~ ^[0-9]+$ ]] || ((K >= ${#PODS[@]})); then - echo "need LAYOUT=AF[,...] and POD=AFD_ in it: LAYOUT=$LAYOUT POD=${POD:-}" >&2 - exit 1 -fi -NUM_A=() NUM_F=() -for p in "${PODS[@]}"; do NUM_A+=("${p%%A*}") && p=${p#*A} && NUM_F+=("${p%F}"); done -((${#PODS[@]} > 1)) || HEAD_HOST=${HEAD_HOST:-127.0.0.1} - -# Extra `vllm serve` flags go last, so they override the defaults. -read -ra ARGS <<<"${VLLM_ARGS:-}" && ARGS+=("$@") -read -ra ATTN_EXTRA <<<"${ATTN_ARGS:-}" -read -ra FFN_EXTRA <<<"${FFN_ARGS:-}" -if [[ -n ${DECODE_BENCH:-} ]]; then # tools/benchmarks is imported from the repository - ATTN_EXTRA=(--kv-transfer-config "{\"kv_connector\": \"AFDDecodeBenchConnector\", \"kv_connector_module_path\": \"tools.benchmarks.decode_bench\", \"kv_role\": \"kv_both\", \"kv_connector_extra_config\": ${DECODE_BENCH}}" "${ATTN_EXTRA[@]}") - export PYTHONPATH=${AFD_REPO:-$PWD}${PYTHONPATH:+:$PYTHONPATH} -fi - -# role_dp TP DP_HOST DP_RPC_PORT ... -> TOTAL and this pod's DP -# flags. A role spread over pods is led by its first pod; the others run -# headless from their start rank (the leader must omit it, or vLLM selects -# hybrid load balancing and rejects the headless peers). -role_dp() { - local tp=$1 host=$2 port=$3 i pods=0 start=0 lead=-1 - shift 3 - local n=("$@") - TOTAL=0 - for i in "${!n[@]}"; do - ((n[i])) || continue - ((n[i] % tp)) && { echo "TP $tp does not divide ${n[i]} ranks on pod $i" >&2; exit 1; } - ((lead < 0)) && lead=$i - ((i < K)) && start=$((start + n[i])) - pods=$((pods + 1)) TOTAL=$((TOTAL + n[i])) - done - DP=(--tensor-parallel-size "$tp" --data-parallel-size $((TOTAL / tp))) - ((pods > 1)) && DP+=(--data-parallel-size-local $((n[K] / tp)) - --data-parallel-address "$host" --data-parallel-rpc-port "$port") - ((pods > 1 && K != lead)) && DP+=(--headless --data-parallel-start-rank $((start / tp))) -} -role_dp 1 "${FFN_HEAD_HOST:-${HEAD_HOST:-vllm-ffn-dp-service}}" 13355 "${NUM_F[@]}" -FFN_DP=("${DP[@]}") NUM_FFN=$TOTAL -role_dp "$ATTN_TP" "${ATTN_HEAD_HOST:-${HEAD_HOST:-vllm-attn-dp-service}}" 13345 "${NUM_A[@]}" -ATTN_DP=("${DP[@]}") NUM_ATTN=$TOTAL - -afd() { # The AFD rendezvous runs on the pod that holds FFN rank 0. - cat < ffn.log 2>&1 & -fi -if ((A)); then - CUDA_VISIBLE_DEVICES=$(seq "$F" $((F + A - 1)) | paste -sd, -) vllm serve "$MODEL_PATH" \ - --additional-config "$(afd attention)" "${COMMON[@]}" "${ATTN_DP[@]}" \ - "${ARGS[@]}" "${ATTN_EXTRA[@]}" --port 18000 > attn.log 2>&1 & -fi -wait +source "${AFD_RECIPE_LIB:-$(dirname "${BASH_SOURCE[0]}")/../../../common/afd_colocation.sh}" +afd_colocation_serve "$@" diff --git a/recipe/gpu/DeepEPM2NAFDConnector/deepseek_v2_lite/prefill_decode_colocation/graph.sh b/recipe/gpu/DeepEPM2NAFDConnector/deepseek_v2_lite/prefill_decode_colocation/graph.sh index 0121b1d0e..1e267ee28 100644 --- a/recipe/gpu/DeepEPM2NAFDConnector/deepseek_v2_lite/prefill_decode_colocation/graph.sh +++ b/recipe/gpu/DeepEPM2NAFDConnector/deepseek_v2_lite/prefill_decode_colocation/graph.sh @@ -3,95 +3,16 @@ # SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project # # DeepSeek-V2-Lite (BF16) on DeepEPM2NAFDConnector, FULL_DECODE_ONLY CUDA graphs. -# Placement: one pod per comma-separated LAYOUT entry (default 2A2F, one pod of -# 4 GPUs). AF puts f FFN ranks on the pod's first f GPUs and a Attention -# ranks on the next a GPUs, e.g. LAYOUT=2A0F,0A2F or LAYOUT=1A1F,1A1F. -# For pod k (k = 0, 1, ...): POD=AFD_, with the same LAYOUT on every pod. -# Settings: recipe/gpu/DeepEPM2NAFDConnector/README.md#recipe-settings +# Placement: one pod per comma-separated LAYOUT entry (default 2A2F), pod k +# running POD=AFD_ with the same LAYOUT; FFN on each pod's first GPUs. +# Settings and running in a pod: recipe/gpu/common/README.md +# shellcheck disable=SC1090,SC2034 # the sourced launcher reads these variables -MODEL_PATH=${MODEL_PATH:-/path/model_weights/DeepSeek-V2-Lite} +AFD_CONNECTOR=${AFD_CONNECTOR:-DeepEPM2NAFDConnector} LAYOUT=${LAYOUT:-2A2F} -TUNING=(--max-num-seqs 64 --max-num-batched-tokens 4096 --max-model-len 4096) -CONNECTOR_CONFIG=${M2N_EXTRA_CONFIG:-} # connector_extra_config members; BF16 payload -CUDAGRAPH_CAPTURE_SIZES=${CUDAGRAPH_CAPTURE_SIZES:-8,16,32,64} # ascending -MODE=(--max-cudagraph-capture-size "${CUDAGRAPH_CAPTURE_SIZES##*,}" - --compilation-config "{\"cudagraph_mode\": \"FULL_DECODE_ONLY\", \"cudagraph_capture_sizes\": [${CUDAGRAPH_CAPTURE_SIZES}]}") - -# ---- The rest is the same in every DeepEPM2NAFDConnector recipe. ---- -export VLLM_USE_V2_MODEL_RUNNER=1 -ATTN_TP=${ATTN_TP:-1} -AFD_CONNECTOR=${AFD_CONNECTOR:-DeepEPM2NAFDConnector} GATE_ON_ATTENTION=true -# P2P routes on the FFN and takes no connector_extra_config. -[[ $AFD_CONNECTOR == P2pNcclAFDConnector ]] && GATE_ON_ATTENTION=false CONNECTOR_CONFIG= -[[ -n $CONNECTOR_CONFIG ]] && CONNECTOR_CONFIG=", \"connector_extra_config\": {${CONNECTOR_CONFIG}}" -IFS=, read -ra PODS <<<"$LAYOUT" -K=${POD:-AFD_0} && K=${K#AFD_} -if [[ ! $LAYOUT =~ ^[0-9]+A[0-9]+F(,[0-9]+A[0-9]+F)*$ || ! $K =~ ^[0-9]+$ ]] || ((K >= ${#PODS[@]})); then - echo "need LAYOUT=AF[,...] and POD=AFD_ in it: LAYOUT=$LAYOUT POD=${POD:-}" >&2 - exit 1 -fi -NUM_A=() NUM_F=() -for p in "${PODS[@]}"; do NUM_A+=("${p%%A*}") && p=${p#*A} && NUM_F+=("${p%F}"); done -((${#PODS[@]} > 1)) || HEAD_HOST=${HEAD_HOST:-127.0.0.1} - -# Extra `vllm serve` flags go last, so they override the defaults. -read -ra ARGS <<<"${VLLM_ARGS:-}" && ARGS+=("$@") -read -ra ATTN_EXTRA <<<"${ATTN_ARGS:-}" -read -ra FFN_EXTRA <<<"${FFN_ARGS:-}" -if [[ -n ${DECODE_BENCH:-} ]]; then # tools/benchmarks is imported from the repository - ATTN_EXTRA=(--kv-transfer-config "{\"kv_connector\": \"AFDDecodeBenchConnector\", \"kv_connector_module_path\": \"tools.benchmarks.decode_bench\", \"kv_role\": \"kv_both\", \"kv_connector_extra_config\": ${DECODE_BENCH}}" "${ATTN_EXTRA[@]}") - export PYTHONPATH=${AFD_REPO:-$PWD}${PYTHONPATH:+:$PYTHONPATH} -fi - -# role_dp TP DP_HOST DP_RPC_PORT ... -> TOTAL and this pod's DP -# flags. A role spread over pods is led by its first pod; the others run -# headless from their start rank (the leader must omit it, or vLLM selects -# hybrid load balancing and rejects the headless peers). -role_dp() { - local tp=$1 host=$2 port=$3 i pods=0 start=0 lead=-1 - shift 3 - local n=("$@") - TOTAL=0 - for i in "${!n[@]}"; do - ((n[i])) || continue - ((n[i] % tp)) && { echo "TP $tp does not divide ${n[i]} ranks on pod $i" >&2; exit 1; } - ((lead < 0)) && lead=$i - ((i < K)) && start=$((start + n[i])) - pods=$((pods + 1)) TOTAL=$((TOTAL + n[i])) - done - DP=(--tensor-parallel-size "$tp" --data-parallel-size $((TOTAL / tp))) - ((pods > 1)) && DP+=(--data-parallel-size-local $((n[K] / tp)) - --data-parallel-address "$host" --data-parallel-rpc-port "$port") - ((pods > 1 && K != lead)) && DP+=(--headless --data-parallel-start-rank $((start / tp))) -} -role_dp 1 "${FFN_HEAD_HOST:-${HEAD_HOST:-vllm-ffn-dp-service}}" 13355 "${NUM_F[@]}" -FFN_DP=("${DP[@]}") NUM_FFN=$TOTAL -role_dp "$ATTN_TP" "${ATTN_HEAD_HOST:-${HEAD_HOST:-vllm-attn-dp-service}}" 13345 "${NUM_A[@]}" -ATTN_DP=("${DP[@]}") NUM_ATTN=$TOTAL - -afd() { # The AFD rendezvous runs on the pod that holds FFN rank 0. - cat < ffn.log 2>&1 & -fi -if ((A)); then - CUDA_VISIBLE_DEVICES=$(seq "$F" $((F + A - 1)) | paste -sd, -) vllm serve "$MODEL_PATH" \ - --additional-config "$(afd attention)" "${COMMON[@]}" "${ATTN_DP[@]}" \ - "${ARGS[@]}" "${ATTN_EXTRA[@]}" --port 18000 > attn.log 2>&1 & -fi -wait +source "${AFD_RECIPE_LIB:-$(dirname "${BASH_SOURCE[0]}")/../../../common/afd_colocation.sh}" +afd_colocation_serve "$@" diff --git a/recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/dbo.sh b/recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/dbo.sh index 6530cac95..bbbbb0cc3 100644 --- a/recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/dbo.sh +++ b/recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/dbo.sh @@ -3,96 +3,18 @@ # SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project # # Qwen3-235B-A22B-FP8 on DeepEPM2NAFDConnector, FULL_DECODE_ONLY graphs + DBO. -# Placement: one pod per comma-separated LAYOUT entry (default 2A2F, one pod of -# 4 GPUs). AF puts f FFN ranks on the pod's first f GPUs and a Attention -# ranks on the next a GPUs, e.g. LAYOUT=2A0F,0A2F or LAYOUT=1A1F,1A1F. -# For pod k (k = 0, 1, ...): POD=AFD_, with the same LAYOUT on every pod. -# Settings: recipe/gpu/DeepEPM2NAFDConnector/README.md#recipe-settings +# Placement: one pod per comma-separated LAYOUT entry (default 2A2F), pod k +# running POD=AFD_ with the same LAYOUT; FFN on each pod's first GPUs. +# Settings and running in a pod: recipe/gpu/common/README.md +# shellcheck disable=SC1090,SC2034 # the sourced launcher reads these variables -MODEL_PATH=${MODEL_PATH:-/path/model_weights/Qwen3-235B-A22B-FP8} +AFD_CONNECTOR=${AFD_CONNECTOR:-DeepEPM2NAFDConnector} LAYOUT=${LAYOUT:-2A2F} -TUNING=(--max-num-seqs 512 --max-num-batched-tokens 3072 --max-model-len 4096) -CONNECTOR_CONFIG="\"payload_format\": \"fp8_block128_fp32\"${M2N_EXTRA_CONFIG:+, $M2N_EXTRA_CONFIG}" -CUDAGRAPH_CAPTURE_SIZES=${CUDAGRAPH_CAPTURE_SIZES:-8,16,32,64,96,128,192,256,384,512} # ascending -MODE=(--max-cudagraph-capture-size "${CUDAGRAPH_CAPTURE_SIZES##*,}" - --compilation-config "{\"cudagraph_mode\": \"FULL_DECODE_ONLY\", \"cudagraph_capture_sizes\": [${CUDAGRAPH_CAPTURE_SIZES}]}" - --enable-dbo --dbo-decode-token-threshold 2 --dbo-prefill-token-threshold 8) - -# ---- The rest is the same in every DeepEPM2NAFDConnector recipe. ---- -export VLLM_USE_V2_MODEL_RUNNER=1 -ATTN_TP=${ATTN_TP:-1} -AFD_CONNECTOR=${AFD_CONNECTOR:-DeepEPM2NAFDConnector} GATE_ON_ATTENTION=true -# P2P routes on the FFN and takes no connector_extra_config. -[[ $AFD_CONNECTOR == P2pNcclAFDConnector ]] && GATE_ON_ATTENTION=false CONNECTOR_CONFIG= -[[ -n $CONNECTOR_CONFIG ]] && CONNECTOR_CONFIG=", \"connector_extra_config\": {${CONNECTOR_CONFIG}}" -IFS=, read -ra PODS <<<"$LAYOUT" -K=${POD:-AFD_0} && K=${K#AFD_} -if [[ ! $LAYOUT =~ ^[0-9]+A[0-9]+F(,[0-9]+A[0-9]+F)*$ || ! $K =~ ^[0-9]+$ ]] || ((K >= ${#PODS[@]})); then - echo "need LAYOUT=AF[,...] and POD=AFD_ in it: LAYOUT=$LAYOUT POD=${POD:-}" >&2 - exit 1 -fi -NUM_A=() NUM_F=() -for p in "${PODS[@]}"; do NUM_A+=("${p%%A*}") && p=${p#*A} && NUM_F+=("${p%F}"); done -((${#PODS[@]} > 1)) || HEAD_HOST=${HEAD_HOST:-127.0.0.1} - -# Extra `vllm serve` flags go last, so they override the defaults. -read -ra ARGS <<<"${VLLM_ARGS:-}" && ARGS+=("$@") -read -ra ATTN_EXTRA <<<"${ATTN_ARGS:-}" -read -ra FFN_EXTRA <<<"${FFN_ARGS:-}" -if [[ -n ${DECODE_BENCH:-} ]]; then # tools/benchmarks is imported from the repository - ATTN_EXTRA=(--kv-transfer-config "{\"kv_connector\": \"AFDDecodeBenchConnector\", \"kv_connector_module_path\": \"tools.benchmarks.decode_bench\", \"kv_role\": \"kv_both\", \"kv_connector_extra_config\": ${DECODE_BENCH}}" "${ATTN_EXTRA[@]}") - export PYTHONPATH=${AFD_REPO:-$PWD}${PYTHONPATH:+:$PYTHONPATH} -fi - -# role_dp TP DP_HOST DP_RPC_PORT ... -> TOTAL and this pod's DP -# flags. A role spread over pods is led by its first pod; the others run -# headless from their start rank (the leader must omit it, or vLLM selects -# hybrid load balancing and rejects the headless peers). -role_dp() { - local tp=$1 host=$2 port=$3 i pods=0 start=0 lead=-1 - shift 3 - local n=("$@") - TOTAL=0 - for i in "${!n[@]}"; do - ((n[i])) || continue - ((n[i] % tp)) && { echo "TP $tp does not divide ${n[i]} ranks on pod $i" >&2; exit 1; } - ((lead < 0)) && lead=$i - ((i < K)) && start=$((start + n[i])) - pods=$((pods + 1)) TOTAL=$((TOTAL + n[i])) - done - DP=(--tensor-parallel-size "$tp" --data-parallel-size $((TOTAL / tp))) - ((pods > 1)) && DP+=(--data-parallel-size-local $((n[K] / tp)) - --data-parallel-address "$host" --data-parallel-rpc-port "$port") - ((pods > 1 && K != lead)) && DP+=(--headless --data-parallel-start-rank $((start / tp))) -} -role_dp 1 "${FFN_HEAD_HOST:-${HEAD_HOST:-vllm-ffn-dp-service}}" 13355 "${NUM_F[@]}" -FFN_DP=("${DP[@]}") NUM_FFN=$TOTAL -role_dp "$ATTN_TP" "${ATTN_HEAD_HOST:-${HEAD_HOST:-vllm-attn-dp-service}}" 13345 "${NUM_A[@]}" -ATTN_DP=("${DP[@]}") NUM_ATTN=$TOTAL - -afd() { # The AFD rendezvous runs on the pod that holds FFN rank 0. - cat < ffn.log 2>&1 & -fi -if ((A)); then - CUDA_VISIBLE_DEVICES=$(seq "$F" $((F + A - 1)) | paste -sd, -) vllm serve "$MODEL_PATH" \ - --additional-config "$(afd attention)" "${COMMON[@]}" "${ATTN_DP[@]}" \ - "${ARGS[@]}" "${ATTN_EXTRA[@]}" --port 18000 > attn.log 2>&1 & -fi -wait +source "${AFD_RECIPE_LIB:-$(dirname "${BASH_SOURCE[0]}")/../../../common/afd_colocation.sh}" +afd_colocation_serve "$@" diff --git a/recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/eager.sh b/recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/eager.sh index 3ae1de33c..96d1b0cc6 100644 --- a/recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/eager.sh +++ b/recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/eager.sh @@ -3,93 +3,17 @@ # SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project # # Qwen3-235B-A22B-FP8 on DeepEPM2NAFDConnector, eager mode. -# Placement: one pod per comma-separated LAYOUT entry (default 2A2F, one pod of -# 4 GPUs). AF puts f FFN ranks on the pod's first f GPUs and a Attention -# ranks on the next a GPUs, e.g. LAYOUT=2A0F,0A2F or LAYOUT=1A1F,1A1F. -# For pod k (k = 0, 1, ...): POD=AFD_, with the same LAYOUT on every pod. -# Settings: recipe/gpu/DeepEPM2NAFDConnector/README.md#recipe-settings +# Placement: one pod per comma-separated LAYOUT entry (default 2A2F), pod k +# running POD=AFD_ with the same LAYOUT; FFN on each pod's first GPUs. +# Settings and running in a pod: recipe/gpu/common/README.md +# shellcheck disable=SC1090,SC2034 # the sourced launcher reads these variables -MODEL_PATH=${MODEL_PATH:-/path/model_weights/Qwen3-235B-A22B-FP8} +AFD_CONNECTOR=${AFD_CONNECTOR:-DeepEPM2NAFDConnector} LAYOUT=${LAYOUT:-2A2F} -TUNING=(--max-num-seqs 256 --max-num-batched-tokens 3072 --max-model-len 4096) -CONNECTOR_CONFIG="\"payload_format\": \"fp8_block128_fp32\"${M2N_EXTRA_CONFIG:+, $M2N_EXTRA_CONFIG}" -MODE=(--enforce-eager) - -# ---- The rest is the same in every DeepEPM2NAFDConnector recipe. ---- -export VLLM_USE_V2_MODEL_RUNNER=1 -ATTN_TP=${ATTN_TP:-1} -AFD_CONNECTOR=${AFD_CONNECTOR:-DeepEPM2NAFDConnector} GATE_ON_ATTENTION=true -# P2P routes on the FFN and takes no connector_extra_config. -[[ $AFD_CONNECTOR == P2pNcclAFDConnector ]] && GATE_ON_ATTENTION=false CONNECTOR_CONFIG= -[[ -n $CONNECTOR_CONFIG ]] && CONNECTOR_CONFIG=", \"connector_extra_config\": {${CONNECTOR_CONFIG}}" -IFS=, read -ra PODS <<<"$LAYOUT" -K=${POD:-AFD_0} && K=${K#AFD_} -if [[ ! $LAYOUT =~ ^[0-9]+A[0-9]+F(,[0-9]+A[0-9]+F)*$ || ! $K =~ ^[0-9]+$ ]] || ((K >= ${#PODS[@]})); then - echo "need LAYOUT=AF[,...] and POD=AFD_ in it: LAYOUT=$LAYOUT POD=${POD:-}" >&2 - exit 1 -fi -NUM_A=() NUM_F=() -for p in "${PODS[@]}"; do NUM_A+=("${p%%A*}") && p=${p#*A} && NUM_F+=("${p%F}"); done -((${#PODS[@]} > 1)) || HEAD_HOST=${HEAD_HOST:-127.0.0.1} - -# Extra `vllm serve` flags go last, so they override the defaults. -read -ra ARGS <<<"${VLLM_ARGS:-}" && ARGS+=("$@") -read -ra ATTN_EXTRA <<<"${ATTN_ARGS:-}" -read -ra FFN_EXTRA <<<"${FFN_ARGS:-}" -if [[ -n ${DECODE_BENCH:-} ]]; then # tools/benchmarks is imported from the repository - ATTN_EXTRA=(--kv-transfer-config "{\"kv_connector\": \"AFDDecodeBenchConnector\", \"kv_connector_module_path\": \"tools.benchmarks.decode_bench\", \"kv_role\": \"kv_both\", \"kv_connector_extra_config\": ${DECODE_BENCH}}" "${ATTN_EXTRA[@]}") - export PYTHONPATH=${AFD_REPO:-$PWD}${PYTHONPATH:+:$PYTHONPATH} -fi - -# role_dp TP DP_HOST DP_RPC_PORT ... -> TOTAL and this pod's DP -# flags. A role spread over pods is led by its first pod; the others run -# headless from their start rank (the leader must omit it, or vLLM selects -# hybrid load balancing and rejects the headless peers). -role_dp() { - local tp=$1 host=$2 port=$3 i pods=0 start=0 lead=-1 - shift 3 - local n=("$@") - TOTAL=0 - for i in "${!n[@]}"; do - ((n[i])) || continue - ((n[i] % tp)) && { echo "TP $tp does not divide ${n[i]} ranks on pod $i" >&2; exit 1; } - ((lead < 0)) && lead=$i - ((i < K)) && start=$((start + n[i])) - pods=$((pods + 1)) TOTAL=$((TOTAL + n[i])) - done - DP=(--tensor-parallel-size "$tp" --data-parallel-size $((TOTAL / tp))) - ((pods > 1)) && DP+=(--data-parallel-size-local $((n[K] / tp)) - --data-parallel-address "$host" --data-parallel-rpc-port "$port") - ((pods > 1 && K != lead)) && DP+=(--headless --data-parallel-start-rank $((start / tp))) -} -role_dp 1 "${FFN_HEAD_HOST:-${HEAD_HOST:-vllm-ffn-dp-service}}" 13355 "${NUM_F[@]}" -FFN_DP=("${DP[@]}") NUM_FFN=$TOTAL -role_dp "$ATTN_TP" "${ATTN_HEAD_HOST:-${HEAD_HOST:-vllm-attn-dp-service}}" 13345 "${NUM_A[@]}" -ATTN_DP=("${DP[@]}") NUM_ATTN=$TOTAL - -afd() { # The AFD rendezvous runs on the pod that holds FFN rank 0. - cat < ffn.log 2>&1 & -fi -if ((A)); then - CUDA_VISIBLE_DEVICES=$(seq "$F" $((F + A - 1)) | paste -sd, -) vllm serve "$MODEL_PATH" \ - --additional-config "$(afd attention)" "${COMMON[@]}" "${ATTN_DP[@]}" \ - "${ARGS[@]}" "${ATTN_EXTRA[@]}" --port 18000 > attn.log 2>&1 & -fi -wait +source "${AFD_RECIPE_LIB:-$(dirname "${BASH_SOURCE[0]}")/../../../common/afd_colocation.sh}" +afd_colocation_serve "$@" diff --git a/recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/graph.sh b/recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/graph.sh index 126ddf8a5..354fbe141 100644 --- a/recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/graph.sh +++ b/recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/graph.sh @@ -3,95 +3,17 @@ # SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project # # Qwen3-235B-A22B-FP8 on DeepEPM2NAFDConnector, FULL_DECODE_ONLY CUDA graphs. -# Placement: one pod per comma-separated LAYOUT entry (default 2A2F, one pod of -# 4 GPUs). AF puts f FFN ranks on the pod's first f GPUs and a Attention -# ranks on the next a GPUs, e.g. LAYOUT=2A0F,0A2F or LAYOUT=1A1F,1A1F. -# For pod k (k = 0, 1, ...): POD=AFD_, with the same LAYOUT on every pod. -# Settings: recipe/gpu/DeepEPM2NAFDConnector/README.md#recipe-settings +# Placement: one pod per comma-separated LAYOUT entry (default 2A2F), pod k +# running POD=AFD_ with the same LAYOUT; FFN on each pod's first GPUs. +# Settings and running in a pod: recipe/gpu/common/README.md +# shellcheck disable=SC1090,SC2034 # the sourced launcher reads these variables -MODEL_PATH=${MODEL_PATH:-/path/model_weights/Qwen3-235B-A22B-FP8} +AFD_CONNECTOR=${AFD_CONNECTOR:-DeepEPM2NAFDConnector} LAYOUT=${LAYOUT:-2A2F} -TUNING=(--max-num-seqs 256 --max-num-batched-tokens 3072 --max-model-len 4096) -CONNECTOR_CONFIG="\"payload_format\": \"fp8_block128_fp32\"${M2N_EXTRA_CONFIG:+, $M2N_EXTRA_CONFIG}" -CUDAGRAPH_CAPTURE_SIZES=${CUDAGRAPH_CAPTURE_SIZES:-8,12,16,24,32,48,64,96,128,192,256} # ascending -MODE=(--max-cudagraph-capture-size "${CUDAGRAPH_CAPTURE_SIZES##*,}" - --compilation-config "{\"cudagraph_mode\": \"FULL_DECODE_ONLY\", \"cudagraph_capture_sizes\": [${CUDAGRAPH_CAPTURE_SIZES}]}") - -# ---- The rest is the same in every DeepEPM2NAFDConnector recipe. ---- -export VLLM_USE_V2_MODEL_RUNNER=1 -ATTN_TP=${ATTN_TP:-1} -AFD_CONNECTOR=${AFD_CONNECTOR:-DeepEPM2NAFDConnector} GATE_ON_ATTENTION=true -# P2P routes on the FFN and takes no connector_extra_config. -[[ $AFD_CONNECTOR == P2pNcclAFDConnector ]] && GATE_ON_ATTENTION=false CONNECTOR_CONFIG= -[[ -n $CONNECTOR_CONFIG ]] && CONNECTOR_CONFIG=", \"connector_extra_config\": {${CONNECTOR_CONFIG}}" -IFS=, read -ra PODS <<<"$LAYOUT" -K=${POD:-AFD_0} && K=${K#AFD_} -if [[ ! $LAYOUT =~ ^[0-9]+A[0-9]+F(,[0-9]+A[0-9]+F)*$ || ! $K =~ ^[0-9]+$ ]] || ((K >= ${#PODS[@]})); then - echo "need LAYOUT=AF[,...] and POD=AFD_ in it: LAYOUT=$LAYOUT POD=${POD:-}" >&2 - exit 1 -fi -NUM_A=() NUM_F=() -for p in "${PODS[@]}"; do NUM_A+=("${p%%A*}") && p=${p#*A} && NUM_F+=("${p%F}"); done -((${#PODS[@]} > 1)) || HEAD_HOST=${HEAD_HOST:-127.0.0.1} - -# Extra `vllm serve` flags go last, so they override the defaults. -read -ra ARGS <<<"${VLLM_ARGS:-}" && ARGS+=("$@") -read -ra ATTN_EXTRA <<<"${ATTN_ARGS:-}" -read -ra FFN_EXTRA <<<"${FFN_ARGS:-}" -if [[ -n ${DECODE_BENCH:-} ]]; then # tools/benchmarks is imported from the repository - ATTN_EXTRA=(--kv-transfer-config "{\"kv_connector\": \"AFDDecodeBenchConnector\", \"kv_connector_module_path\": \"tools.benchmarks.decode_bench\", \"kv_role\": \"kv_both\", \"kv_connector_extra_config\": ${DECODE_BENCH}}" "${ATTN_EXTRA[@]}") - export PYTHONPATH=${AFD_REPO:-$PWD}${PYTHONPATH:+:$PYTHONPATH} -fi - -# role_dp TP DP_HOST DP_RPC_PORT ... -> TOTAL and this pod's DP -# flags. A role spread over pods is led by its first pod; the others run -# headless from their start rank (the leader must omit it, or vLLM selects -# hybrid load balancing and rejects the headless peers). -role_dp() { - local tp=$1 host=$2 port=$3 i pods=0 start=0 lead=-1 - shift 3 - local n=("$@") - TOTAL=0 - for i in "${!n[@]}"; do - ((n[i])) || continue - ((n[i] % tp)) && { echo "TP $tp does not divide ${n[i]} ranks on pod $i" >&2; exit 1; } - ((lead < 0)) && lead=$i - ((i < K)) && start=$((start + n[i])) - pods=$((pods + 1)) TOTAL=$((TOTAL + n[i])) - done - DP=(--tensor-parallel-size "$tp" --data-parallel-size $((TOTAL / tp))) - ((pods > 1)) && DP+=(--data-parallel-size-local $((n[K] / tp)) - --data-parallel-address "$host" --data-parallel-rpc-port "$port") - ((pods > 1 && K != lead)) && DP+=(--headless --data-parallel-start-rank $((start / tp))) -} -role_dp 1 "${FFN_HEAD_HOST:-${HEAD_HOST:-vllm-ffn-dp-service}}" 13355 "${NUM_F[@]}" -FFN_DP=("${DP[@]}") NUM_FFN=$TOTAL -role_dp "$ATTN_TP" "${ATTN_HEAD_HOST:-${HEAD_HOST:-vllm-attn-dp-service}}" 13345 "${NUM_A[@]}" -ATTN_DP=("${DP[@]}") NUM_ATTN=$TOTAL - -afd() { # The AFD rendezvous runs on the pod that holds FFN rank 0. - cat < ffn.log 2>&1 & -fi -if ((A)); then - CUDA_VISIBLE_DEVICES=$(seq "$F" $((F + A - 1)) | paste -sd, -) vllm serve "$MODEL_PATH" \ - --additional-config "$(afd attention)" "${COMMON[@]}" "${ATTN_DP[@]}" \ - "${ARGS[@]}" "${ATTN_EXTRA[@]}" --port 18000 > attn.log 2>&1 & -fi -wait +source "${AFD_RECIPE_LIB:-$(dirname "${BASH_SOURCE[0]}")/../../../common/afd_colocation.sh}" +afd_colocation_serve "$@" diff --git a/recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/dbo.sh b/recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/dbo.sh index 8a611278a..845b912d8 100644 --- a/recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/dbo.sh +++ b/recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/dbo.sh @@ -3,96 +3,17 @@ # SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project # # Qwen3-30B-A3B (BF16) on DeepEPM2NAFDConnector, FULL_DECODE_ONLY graphs + DBO. -# Placement: one pod per comma-separated LAYOUT entry (default 2A2F, one pod of -# 4 GPUs). AF puts f FFN ranks on the pod's first f GPUs and a Attention -# ranks on the next a GPUs, e.g. LAYOUT=2A0F,0A2F or LAYOUT=1A1F,1A1F. -# For pod k (k = 0, 1, ...): POD=AFD_, with the same LAYOUT on every pod. -# Settings: recipe/gpu/DeepEPM2NAFDConnector/README.md#recipe-settings +# Placement: one pod per comma-separated LAYOUT entry (default 2A2F), pod k +# running POD=AFD_ with the same LAYOUT; FFN on each pod's first GPUs. +# Settings and running in a pod: recipe/gpu/common/README.md +# shellcheck disable=SC1090,SC2034 # the sourced launcher reads these variables -MODEL_PATH=${MODEL_PATH:-/path/model_weights/Qwen3-30B-A3B} +AFD_CONNECTOR=${AFD_CONNECTOR:-DeepEPM2NAFDConnector} LAYOUT=${LAYOUT:-2A2F} -TUNING=(--max-num-seqs 64 --max-num-batched-tokens 4096 --max-model-len 4096) -CONNECTOR_CONFIG=${M2N_EXTRA_CONFIG:-} # connector_extra_config members; BF16 payload -CUDAGRAPH_CAPTURE_SIZES=${CUDAGRAPH_CAPTURE_SIZES:-8,16,32,64} # ascending -MODE=(--max-cudagraph-capture-size "${CUDAGRAPH_CAPTURE_SIZES##*,}" - --compilation-config "{\"cudagraph_mode\": \"FULL_DECODE_ONLY\", \"cudagraph_capture_sizes\": [${CUDAGRAPH_CAPTURE_SIZES}]}" - --enable-dbo --dbo-decode-token-threshold 2 --dbo-prefill-token-threshold 8) - -# ---- The rest is the same in every DeepEPM2NAFDConnector recipe. ---- -export VLLM_USE_V2_MODEL_RUNNER=1 -ATTN_TP=${ATTN_TP:-1} -AFD_CONNECTOR=${AFD_CONNECTOR:-DeepEPM2NAFDConnector} GATE_ON_ATTENTION=true -# P2P routes on the FFN and takes no connector_extra_config. -[[ $AFD_CONNECTOR == P2pNcclAFDConnector ]] && GATE_ON_ATTENTION=false CONNECTOR_CONFIG= -[[ -n $CONNECTOR_CONFIG ]] && CONNECTOR_CONFIG=", \"connector_extra_config\": {${CONNECTOR_CONFIG}}" -IFS=, read -ra PODS <<<"$LAYOUT" -K=${POD:-AFD_0} && K=${K#AFD_} -if [[ ! $LAYOUT =~ ^[0-9]+A[0-9]+F(,[0-9]+A[0-9]+F)*$ || ! $K =~ ^[0-9]+$ ]] || ((K >= ${#PODS[@]})); then - echo "need LAYOUT=AF[,...] and POD=AFD_ in it: LAYOUT=$LAYOUT POD=${POD:-}" >&2 - exit 1 -fi -NUM_A=() NUM_F=() -for p in "${PODS[@]}"; do NUM_A+=("${p%%A*}") && p=${p#*A} && NUM_F+=("${p%F}"); done -((${#PODS[@]} > 1)) || HEAD_HOST=${HEAD_HOST:-127.0.0.1} - -# Extra `vllm serve` flags go last, so they override the defaults. -read -ra ARGS <<<"${VLLM_ARGS:-}" && ARGS+=("$@") -read -ra ATTN_EXTRA <<<"${ATTN_ARGS:-}" -read -ra FFN_EXTRA <<<"${FFN_ARGS:-}" -if [[ -n ${DECODE_BENCH:-} ]]; then # tools/benchmarks is imported from the repository - ATTN_EXTRA=(--kv-transfer-config "{\"kv_connector\": \"AFDDecodeBenchConnector\", \"kv_connector_module_path\": \"tools.benchmarks.decode_bench\", \"kv_role\": \"kv_both\", \"kv_connector_extra_config\": ${DECODE_BENCH}}" "${ATTN_EXTRA[@]}") - export PYTHONPATH=${AFD_REPO:-$PWD}${PYTHONPATH:+:$PYTHONPATH} -fi - -# role_dp TP DP_HOST DP_RPC_PORT ... -> TOTAL and this pod's DP -# flags. A role spread over pods is led by its first pod; the others run -# headless from their start rank (the leader must omit it, or vLLM selects -# hybrid load balancing and rejects the headless peers). -role_dp() { - local tp=$1 host=$2 port=$3 i pods=0 start=0 lead=-1 - shift 3 - local n=("$@") - TOTAL=0 - for i in "${!n[@]}"; do - ((n[i])) || continue - ((n[i] % tp)) && { echo "TP $tp does not divide ${n[i]} ranks on pod $i" >&2; exit 1; } - ((lead < 0)) && lead=$i - ((i < K)) && start=$((start + n[i])) - pods=$((pods + 1)) TOTAL=$((TOTAL + n[i])) - done - DP=(--tensor-parallel-size "$tp" --data-parallel-size $((TOTAL / tp))) - ((pods > 1)) && DP+=(--data-parallel-size-local $((n[K] / tp)) - --data-parallel-address "$host" --data-parallel-rpc-port "$port") - ((pods > 1 && K != lead)) && DP+=(--headless --data-parallel-start-rank $((start / tp))) -} -role_dp 1 "${FFN_HEAD_HOST:-${HEAD_HOST:-vllm-ffn-dp-service}}" 13355 "${NUM_F[@]}" -FFN_DP=("${DP[@]}") NUM_FFN=$TOTAL -role_dp "$ATTN_TP" "${ATTN_HEAD_HOST:-${HEAD_HOST:-vllm-attn-dp-service}}" 13345 "${NUM_A[@]}" -ATTN_DP=("${DP[@]}") NUM_ATTN=$TOTAL - -afd() { # The AFD rendezvous runs on the pod that holds FFN rank 0. - cat < ffn.log 2>&1 & -fi -if ((A)); then - CUDA_VISIBLE_DEVICES=$(seq "$F" $((F + A - 1)) | paste -sd, -) vllm serve "$MODEL_PATH" \ - --additional-config "$(afd attention)" "${COMMON[@]}" "${ATTN_DP[@]}" \ - "${ARGS[@]}" "${ATTN_EXTRA[@]}" --port 18000 > attn.log 2>&1 & -fi -wait +source "${AFD_RECIPE_LIB:-$(dirname "${BASH_SOURCE[0]}")/../../../common/afd_colocation.sh}" +afd_colocation_serve "$@" diff --git a/recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/eager.sh b/recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/eager.sh index 656764492..1df79d3cf 100644 --- a/recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/eager.sh +++ b/recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/eager.sh @@ -3,93 +3,16 @@ # SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project # # Qwen3-30B-A3B (BF16) on DeepEPM2NAFDConnector, eager mode. -# Placement: one pod per comma-separated LAYOUT entry (default 2A2F, one pod of -# 4 GPUs). AF puts f FFN ranks on the pod's first f GPUs and a Attention -# ranks on the next a GPUs, e.g. LAYOUT=2A0F,0A2F or LAYOUT=1A1F,1A1F. -# For pod k (k = 0, 1, ...): POD=AFD_, with the same LAYOUT on every pod. -# Settings: recipe/gpu/DeepEPM2NAFDConnector/README.md#recipe-settings +# Placement: one pod per comma-separated LAYOUT entry (default 2A2F), pod k +# running POD=AFD_ with the same LAYOUT; FFN on each pod's first GPUs. +# Settings and running in a pod: recipe/gpu/common/README.md +# shellcheck disable=SC1090,SC2034 # the sourced launcher reads these variables -MODEL_PATH=${MODEL_PATH:-/path/model_weights/Qwen3-30B-A3B} +AFD_CONNECTOR=${AFD_CONNECTOR:-DeepEPM2NAFDConnector} LAYOUT=${LAYOUT:-2A2F} -TUNING=(--max-num-seqs 64 --max-num-batched-tokens 4096 --max-model-len 4096) -CONNECTOR_CONFIG=${M2N_EXTRA_CONFIG:-} # connector_extra_config members; BF16 payload -MODE=(--enforce-eager) - -# ---- The rest is the same in every DeepEPM2NAFDConnector recipe. ---- -export VLLM_USE_V2_MODEL_RUNNER=1 -ATTN_TP=${ATTN_TP:-1} -AFD_CONNECTOR=${AFD_CONNECTOR:-DeepEPM2NAFDConnector} GATE_ON_ATTENTION=true -# P2P routes on the FFN and takes no connector_extra_config. -[[ $AFD_CONNECTOR == P2pNcclAFDConnector ]] && GATE_ON_ATTENTION=false CONNECTOR_CONFIG= -[[ -n $CONNECTOR_CONFIG ]] && CONNECTOR_CONFIG=", \"connector_extra_config\": {${CONNECTOR_CONFIG}}" -IFS=, read -ra PODS <<<"$LAYOUT" -K=${POD:-AFD_0} && K=${K#AFD_} -if [[ ! $LAYOUT =~ ^[0-9]+A[0-9]+F(,[0-9]+A[0-9]+F)*$ || ! $K =~ ^[0-9]+$ ]] || ((K >= ${#PODS[@]})); then - echo "need LAYOUT=AF[,...] and POD=AFD_ in it: LAYOUT=$LAYOUT POD=${POD:-}" >&2 - exit 1 -fi -NUM_A=() NUM_F=() -for p in "${PODS[@]}"; do NUM_A+=("${p%%A*}") && p=${p#*A} && NUM_F+=("${p%F}"); done -((${#PODS[@]} > 1)) || HEAD_HOST=${HEAD_HOST:-127.0.0.1} - -# Extra `vllm serve` flags go last, so they override the defaults. -read -ra ARGS <<<"${VLLM_ARGS:-}" && ARGS+=("$@") -read -ra ATTN_EXTRA <<<"${ATTN_ARGS:-}" -read -ra FFN_EXTRA <<<"${FFN_ARGS:-}" -if [[ -n ${DECODE_BENCH:-} ]]; then # tools/benchmarks is imported from the repository - ATTN_EXTRA=(--kv-transfer-config "{\"kv_connector\": \"AFDDecodeBenchConnector\", \"kv_connector_module_path\": \"tools.benchmarks.decode_bench\", \"kv_role\": \"kv_both\", \"kv_connector_extra_config\": ${DECODE_BENCH}}" "${ATTN_EXTRA[@]}") - export PYTHONPATH=${AFD_REPO:-$PWD}${PYTHONPATH:+:$PYTHONPATH} -fi - -# role_dp TP DP_HOST DP_RPC_PORT ... -> TOTAL and this pod's DP -# flags. A role spread over pods is led by its first pod; the others run -# headless from their start rank (the leader must omit it, or vLLM selects -# hybrid load balancing and rejects the headless peers). -role_dp() { - local tp=$1 host=$2 port=$3 i pods=0 start=0 lead=-1 - shift 3 - local n=("$@") - TOTAL=0 - for i in "${!n[@]}"; do - ((n[i])) || continue - ((n[i] % tp)) && { echo "TP $tp does not divide ${n[i]} ranks on pod $i" >&2; exit 1; } - ((lead < 0)) && lead=$i - ((i < K)) && start=$((start + n[i])) - pods=$((pods + 1)) TOTAL=$((TOTAL + n[i])) - done - DP=(--tensor-parallel-size "$tp" --data-parallel-size $((TOTAL / tp))) - ((pods > 1)) && DP+=(--data-parallel-size-local $((n[K] / tp)) - --data-parallel-address "$host" --data-parallel-rpc-port "$port") - ((pods > 1 && K != lead)) && DP+=(--headless --data-parallel-start-rank $((start / tp))) -} -role_dp 1 "${FFN_HEAD_HOST:-${HEAD_HOST:-vllm-ffn-dp-service}}" 13355 "${NUM_F[@]}" -FFN_DP=("${DP[@]}") NUM_FFN=$TOTAL -role_dp "$ATTN_TP" "${ATTN_HEAD_HOST:-${HEAD_HOST:-vllm-attn-dp-service}}" 13345 "${NUM_A[@]}" -ATTN_DP=("${DP[@]}") NUM_ATTN=$TOTAL - -afd() { # The AFD rendezvous runs on the pod that holds FFN rank 0. - cat < ffn.log 2>&1 & -fi -if ((A)); then - CUDA_VISIBLE_DEVICES=$(seq "$F" $((F + A - 1)) | paste -sd, -) vllm serve "$MODEL_PATH" \ - --additional-config "$(afd attention)" "${COMMON[@]}" "${ATTN_DP[@]}" \ - "${ARGS[@]}" "${ATTN_EXTRA[@]}" --port 18000 > attn.log 2>&1 & -fi -wait +source "${AFD_RECIPE_LIB:-$(dirname "${BASH_SOURCE[0]}")/../../../common/afd_colocation.sh}" +afd_colocation_serve "$@" diff --git a/recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/graph.sh b/recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/graph.sh index d8ea0feed..9859a3c61 100644 --- a/recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/graph.sh +++ b/recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/graph.sh @@ -3,95 +3,16 @@ # SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project # # Qwen3-30B-A3B (BF16) on DeepEPM2NAFDConnector, FULL_DECODE_ONLY CUDA graphs. -# Placement: one pod per comma-separated LAYOUT entry (default 2A2F, one pod of -# 4 GPUs). AF puts f FFN ranks on the pod's first f GPUs and a Attention -# ranks on the next a GPUs, e.g. LAYOUT=2A0F,0A2F or LAYOUT=1A1F,1A1F. -# For pod k (k = 0, 1, ...): POD=AFD_, with the same LAYOUT on every pod. -# Settings: recipe/gpu/DeepEPM2NAFDConnector/README.md#recipe-settings +# Placement: one pod per comma-separated LAYOUT entry (default 2A2F), pod k +# running POD=AFD_ with the same LAYOUT; FFN on each pod's first GPUs. +# Settings and running in a pod: recipe/gpu/common/README.md +# shellcheck disable=SC1090,SC2034 # the sourced launcher reads these variables -MODEL_PATH=${MODEL_PATH:-/path/model_weights/Qwen3-30B-A3B} +AFD_CONNECTOR=${AFD_CONNECTOR:-DeepEPM2NAFDConnector} LAYOUT=${LAYOUT:-2A2F} -TUNING=(--max-num-seqs 64 --max-num-batched-tokens 4096 --max-model-len 4096) -CONNECTOR_CONFIG=${M2N_EXTRA_CONFIG:-} # connector_extra_config members; BF16 payload -CUDAGRAPH_CAPTURE_SIZES=${CUDAGRAPH_CAPTURE_SIZES:-8,16,32,64} # ascending -MODE=(--max-cudagraph-capture-size "${CUDAGRAPH_CAPTURE_SIZES##*,}" - --compilation-config "{\"cudagraph_mode\": \"FULL_DECODE_ONLY\", \"cudagraph_capture_sizes\": [${CUDAGRAPH_CAPTURE_SIZES}]}") - -# ---- The rest is the same in every DeepEPM2NAFDConnector recipe. ---- -export VLLM_USE_V2_MODEL_RUNNER=1 -ATTN_TP=${ATTN_TP:-1} -AFD_CONNECTOR=${AFD_CONNECTOR:-DeepEPM2NAFDConnector} GATE_ON_ATTENTION=true -# P2P routes on the FFN and takes no connector_extra_config. -[[ $AFD_CONNECTOR == P2pNcclAFDConnector ]] && GATE_ON_ATTENTION=false CONNECTOR_CONFIG= -[[ -n $CONNECTOR_CONFIG ]] && CONNECTOR_CONFIG=", \"connector_extra_config\": {${CONNECTOR_CONFIG}}" -IFS=, read -ra PODS <<<"$LAYOUT" -K=${POD:-AFD_0} && K=${K#AFD_} -if [[ ! $LAYOUT =~ ^[0-9]+A[0-9]+F(,[0-9]+A[0-9]+F)*$ || ! $K =~ ^[0-9]+$ ]] || ((K >= ${#PODS[@]})); then - echo "need LAYOUT=AF[,...] and POD=AFD_ in it: LAYOUT=$LAYOUT POD=${POD:-}" >&2 - exit 1 -fi -NUM_A=() NUM_F=() -for p in "${PODS[@]}"; do NUM_A+=("${p%%A*}") && p=${p#*A} && NUM_F+=("${p%F}"); done -((${#PODS[@]} > 1)) || HEAD_HOST=${HEAD_HOST:-127.0.0.1} - -# Extra `vllm serve` flags go last, so they override the defaults. -read -ra ARGS <<<"${VLLM_ARGS:-}" && ARGS+=("$@") -read -ra ATTN_EXTRA <<<"${ATTN_ARGS:-}" -read -ra FFN_EXTRA <<<"${FFN_ARGS:-}" -if [[ -n ${DECODE_BENCH:-} ]]; then # tools/benchmarks is imported from the repository - ATTN_EXTRA=(--kv-transfer-config "{\"kv_connector\": \"AFDDecodeBenchConnector\", \"kv_connector_module_path\": \"tools.benchmarks.decode_bench\", \"kv_role\": \"kv_both\", \"kv_connector_extra_config\": ${DECODE_BENCH}}" "${ATTN_EXTRA[@]}") - export PYTHONPATH=${AFD_REPO:-$PWD}${PYTHONPATH:+:$PYTHONPATH} -fi - -# role_dp TP DP_HOST DP_RPC_PORT ... -> TOTAL and this pod's DP -# flags. A role spread over pods is led by its first pod; the others run -# headless from their start rank (the leader must omit it, or vLLM selects -# hybrid load balancing and rejects the headless peers). -role_dp() { - local tp=$1 host=$2 port=$3 i pods=0 start=0 lead=-1 - shift 3 - local n=("$@") - TOTAL=0 - for i in "${!n[@]}"; do - ((n[i])) || continue - ((n[i] % tp)) && { echo "TP $tp does not divide ${n[i]} ranks on pod $i" >&2; exit 1; } - ((lead < 0)) && lead=$i - ((i < K)) && start=$((start + n[i])) - pods=$((pods + 1)) TOTAL=$((TOTAL + n[i])) - done - DP=(--tensor-parallel-size "$tp" --data-parallel-size $((TOTAL / tp))) - ((pods > 1)) && DP+=(--data-parallel-size-local $((n[K] / tp)) - --data-parallel-address "$host" --data-parallel-rpc-port "$port") - ((pods > 1 && K != lead)) && DP+=(--headless --data-parallel-start-rank $((start / tp))) -} -role_dp 1 "${FFN_HEAD_HOST:-${HEAD_HOST:-vllm-ffn-dp-service}}" 13355 "${NUM_F[@]}" -FFN_DP=("${DP[@]}") NUM_FFN=$TOTAL -role_dp "$ATTN_TP" "${ATTN_HEAD_HOST:-${HEAD_HOST:-vllm-attn-dp-service}}" 13345 "${NUM_A[@]}" -ATTN_DP=("${DP[@]}") NUM_ATTN=$TOTAL - -afd() { # The AFD rendezvous runs on the pod that holds FFN rank 0. - cat < ffn.log 2>&1 & -fi -if ((A)); then - CUDA_VISIBLE_DEVICES=$(seq "$F" $((F + A - 1)) | paste -sd, -) vllm serve "$MODEL_PATH" \ - --additional-config "$(afd attention)" "${COMMON[@]}" "${ATTN_DP[@]}" \ - "${ARGS[@]}" "${ATTN_EXTRA[@]}" --port 18000 > attn.log 2>&1 & -fi -wait +source "${AFD_RECIPE_LIB:-$(dirname "${BASH_SOURCE[0]}")/../../../common/afd_colocation.sh}" +afd_colocation_serve "$@" diff --git a/recipe/gpu/P2pNcclAFDConnector/README.md b/recipe/gpu/P2pNcclAFDConnector/README.md index 58862a6f0..de503ede9 100644 --- a/recipe/gpu/P2pNcclAFDConnector/README.md +++ b/recipe/gpu/P2pNcclAFDConnector/README.md @@ -11,36 +11,33 @@ for the connector's rank mapping and configuration contract. ```text . ├── deepseek_v2_lite/ -│ ├── README.md -│ ├── prefill_decode_colocation/ # 2A2F, DP/TP variants, eager/graph +│ ├── prefill_decode_colocation/ # dbo.sh, eager_dbo.sh (any LAYOUT, ATTN_TP, FFN_TP) │ └── prefill_decode_disaggregation/ # 2P1A1F -└── qwen3_5_122b_a10b_fp8/ - └── prefill_decode_colocation/ # multipod_2a_2a_2f_graph.sh: 4A2F split across 2 attention pods, graph +├── qwen3_30b_a3b/prefill_decode_colocation/ # graph.sh, default 2A2F +├── qwen3_235b_a22b_fp8/prefill_decode_colocation/ # graph.sh, default 4A2F +└── qwen3_5_122b_a10b_fp8/prefill_decode_colocation/ # graph.sh, default 2A0F,2A0F,0A2F (3 pods) ``` Each model directory holds its own recipe scripts; open its `README.md` -(where present) before running anything in it -- it documents prerequisites, -topology, ports, and known limitations for that model. +(where present) before running anything in it. -Recipes under `prefill_decode_colocation/` background an attention worker and -an FFN worker (no prefill split). Recipes under `prefill_decode_disaggregation/` -add a separate prefill stage and a proxy in front of it -- see the caveat in -"Deploying on Kubernetes" below. +The `prefill_decode_colocation/` recipes run an Attention server and an FFN +server (no prefill split) through the shared +[colocation launcher](../common/README.md), which documents the layout, +settings, ports, and how to run a recipe directly or in a pod. Recipes under +`prefill_decode_disaggregation/` add a separate prefill stage and a proxy in +front of it -- see the caveat in "Deploying on Kubernetes" below. ## Running on a local host -Run any script directly from the repository root, e.g.: - ```bash export MODEL_PATH=/path/model_weights/DeepSeek-V2-Lite -bash recipe/gpu/P2pNcclAFDConnector/deepseek_v2_lite/prefill_decode_colocation/2a2f_graph_dbo_dp1tp2.sh +bash recipe/gpu/P2pNcclAFDConnector/deepseek_v2_lite/prefill_decode_colocation/dbo.sh ``` -Each script backgrounds its workers and writes per-worker logs (`attn.log`, -`ffn.log`, and for disaggregation `afd_prefill*.log`) into the current -directory. Wait for `Application startup complete` in each log before -sending traffic. See the model README for exact prerequisites (GPU count, -weights, ports) and the benchmark command for that model. +A colocation recipe writes `attn.log` and `ffn.log` into the current +directory; send traffic to port `18000` once `attn.log` prints +`Application startup complete`. ## Deploying on Kubernetes @@ -49,15 +46,11 @@ to stand up a recipe from this folder on a Kubernetes/OpenShift cluster -- see that skill for the full contract (PVC/ConfigMap/pod/Service lifecycle, scope, and requirements). -The skill requires the recipe path (e.g. -`recipe/gpu/P2pNcclAFDConnector/deepseek_v2_lite/prefill_decode_colocation/2a2f_graph_dbo_dp1tp2.sh`) -plus the image, model id, and PVC name it needs. Placement is never asked -for: a recipe without a top-of-file placement comment deploys as a single -pod; a recipe whose header describes a per-pod placement (e.g. -`qwen3_5_122b_a10b_fp8/prefill_decode_colocation/multipod_2a_2a_2f_graph.sh`, -`deepseek_v2_lite/prefill_decode_colocation/multipod_2a_2f_graph_dbo_dp1tp2.sh`) -deploys exactly that many pods with exactly that per-pod setup -- see the -skill's `resolve-recipe.md` for the header convention. +The skill takes the recipe path (e.g. +`recipe/gpu/P2pNcclAFDConnector/deepseek_v2_lite/prefill_decode_colocation/dbo.sh`) +plus the image, model id, and PVC name, and deploys one pod per `LAYOUT` +entry (for example `LAYOUT=2A0F,0A2F ATTN_TP=2 FFN_TP=2` on the DeepSeek +recipe, or the Qwen3.5 recipe's default three pods). `prefill_decode_disaggregation` recipes are out of scope for the skill; run those locally instead. diff --git a/recipe/gpu/P2pNcclAFDConnector/deepseek_v2_lite/README.md b/recipe/gpu/P2pNcclAFDConnector/deepseek_v2_lite/README.md index 2f18e18a9..6e8d2acf8 100644 --- a/recipe/gpu/P2pNcclAFDConnector/deepseek_v2_lite/README.md +++ b/recipe/gpu/P2pNcclAFDConnector/deepseek_v2_lite/README.md @@ -14,18 +14,18 @@ for tested E2E configurations. ## Prerequisites - Install [NIXL](https://github.com/ai-dynamo/nixl). -- Four GPUs for the `2a2f_*` and `2p1a1f_*` examples; the `4a4f_*` - examples require eight GPUs. The multipod example uses two GPUs per pod - across two pods. +- Four GPUs for the default `2A2F` colocation layout and the `2p1a1f_*` + examples; `LAYOUT=4A4F` needs eight GPUs, and `LAYOUT=2A0F,0A2F` two GPUs + on each of two pods. - vLLM `v0.32.0` and the `afd-plugin` package installed in the same environment (see repository root `AGENTS.md`). - DeepSeek-V2-Lite weights on disk. Colocation scripts default to `/path/model_weights/DeepSeek-V2-Lite`; disaggregation scripts default to `/home/models/hub/models--deepseek-ai--DeepSeek-V2-Lite/snapshots/604d5664dddd88a0433dbae533b7fe9472482de0`. Set `MODEL_PATH=...` to the local weights when launching. -- A free TCP port `6269` on `127.0.0.1` for the AFD p2p connector, and - HTTP ports `18305` and `18306` for local colocation, or - `18301`–`18305` for prefill/decode disaggregation. +- Free ports `6239` (AFD rendezvous), `18000` and `18001` (Attention and + FFN HTTP) for colocation, or `18301`–`18305` and `6269` for + prefill/decode disaggregation. ## Directory layout @@ -34,14 +34,9 @@ for tested E2E configurations. ├── prefill_decode_disaggregation/ # prefill_decode_disaggregation, 2P1A1F topology │ ├── 2p1a1f_eager_dbo.sh │ └── 2p1a1f_graph_dbo.sh -└── prefill_decode_colocation/ # prefill_decode_colocation, 2A2F/4A4F topology - ├── 2a2f_eager_dbo_dp1tp2.sh - ├── 2a2f_eager_dbo_dp2tp1.sh - ├── 2a2f_graph_dbo_dp1tp2.sh - ├── 2a2f_graph_dbo_dp2tp1.sh - ├── 4a4f_eager_dbo_dp2tp2.sh - ├── 4a4f_graph_dbo_dp2tp2.sh - └── multipod_2a_2f_graph_dbo_dp1tp2.sh # 2a2f_graph_dbo_dp1tp2, split across 2 k8s pods +└── prefill_decode_colocation/ # any LAYOUT, ATTN_TP, FFN_TP + ├── dbo.sh # FULL_DECODE_ONLY graphs + DBO + └── eager_dbo.sh # eager + DBO ``` ### 1. Prefill/Decode Disaggregation — `1a1f` @@ -58,31 +53,17 @@ for tested E2E configurations. ### 2. Prefill/Decode Colocation — `2a2f` -2 processes, two GPUs each: - -| GPUs | Role | Port | -|------|-----------|-------| -| 0, 1 | Attention | 18305 | -| 2, 3 | FFN | 18306 | - -The four `2a2f_*` variants cover the TP/DP cross product; `4a4f_*` doubles -both attention and FFN ranks (DP=2 on each side); `multipod_2a_2f_graph_dbo_dp1tp2.sh` -is the same `2a2f_graph_dbo_dp1tp2` config, but gated on a `POD` env var -(`ATTENTION_0` or `FFN_0`) for the `deploy-afd-k8s` skill to deploy it -across 2 pods (see repo-root `recipe/gpu/P2pNcclAFDConnector/README.md`); -its Attention and FFN HTTP ports can both be 18305 because they are in -separate pods -- -unlike the scripts above, it is **not** meant to be run directly: with -`POD` unset, neither `vllm serve` block executes, and its -`AFD_CONNECTOR_HOST` default (`vllm-ffn-p2p-service`) only resolves inside -the cluster the skill deploys it to: - -| File | DP | TP | -|---------------------------------------|----|----| -| `2a2f_*_dp1tp2.sh` | 1 | 2 | -| `2a2f_*_dp2tp1.sh` | 2 | 1 | -| `4a4f_*_dp2tp2.sh` | 2 | 2 | -| `multipod_2a_2f_graph_dbo_dp1tp2.sh` | 1 | 2 | +The colocation recipes run on the shared +[colocation launcher](../../common/README.md): `LAYOUT` places the ranks +(default `2A2F`: FFN DP2 on GPUs 0,1, Attention DP2 on GPUs 2,3, Attention +HTTP port 18000), and `ATTN_TP`/`FFN_TP` set each role's TP size: + +| Settings | Attention / FFN | GPUs | +| ---------- | ----------------- | ------ | +| (defaults) | DP2 x TP1 each | 4 | +| `ATTN_TP=2 FFN_TP=2` | DP1 x TP2 each | 4 | +| `LAYOUT=4A4F ATTN_TP=2 FFN_TP=2` | DP2 x TP2 each | 8 | +| `LAYOUT=2A0F,0A2F ATTN_TP=2 FFN_TP=2` | DP1 x TP2 each, one pod per role (`POD=AFD_0`, `POD=AFD_1`) | 2 + 2 | ## Running @@ -96,8 +77,8 @@ before sending traffic. ```bash export MODEL_PATH=/path/model_weights/DeepSeek-V2-Lite -export VLLM_USE_V2_MODEL_RUNNER=1 -bash recipe/gpu/P2pNcclAFDConnector/deepseek_v2_lite/prefill_decode_colocation/2a2f_graph_dbo_dp1tp2.sh +ATTN_TP=2 FFN_TP=2 \ + bash recipe/gpu/P2pNcclAFDConnector/deepseek_v2_lite/prefill_decode_colocation/dbo.sh ``` ### prefill_decode_disaggregation @@ -115,14 +96,15 @@ Confirm the serving stack is ready before running: ```bash export MODEL_PATH=/path/model_weights/DeepSeek-V2-Lite export RESULT_DIR=/tmp/results -export RESULT_FILENAME=2a2f_graph_dbo_dp1tp2.json +export RESULT_FILENAME=2a2f_dbo_dp1tp2.json bash tools/benchmarks/request_generator.sh ``` The script runs `vllm bench serve`; it does not wait for readiness or send a separate smoke request. By default it sends 1024 random requests (1024 input tokens / 128 output tokens) at request rate 5 with `--max-concurrency 32` -against `127.0.0.1:18305`, and saves the JSON result to +against `127.0.0.1:18000` (the colocation Attention port; set `PORT=18305` +for the disaggregation proxy), and saves the JSON result to `$RESULT_DIR/$RESULT_FILENAME`. Override `HOST`, `PORT`, `MODEL_PATH`, `NUM_PROMPTS`, `REQUEST_RATE`, `MAX_CONCURRENCY`, `INPUT_LEN`, and `OUTPUT_LEN` for smaller smoke runs or larger throughput sweeps. @@ -130,7 +112,9 @@ for smaller smoke runs or larger throughput sweeps. ## Common AFD configuration Every AFD worker is wired through `--additional-config` with the same -shape; `role` differs between attention and FFN: +shape; `role` differs between attention and FFN. The colocation launcher +builds it from `LAYOUT` (the disaggregation scripts write it inline, on port +6269); print it with `AFD_DRY_RUN=1`: ```jsonc { @@ -138,9 +122,11 @@ shape; `role` differs between attention and FFN: "role": "attention", // or "ffn" "connector": "P2pNcclAFDConnector", "host": "127.0.0.1", - "port": 6269, - "num_attention_ranks": 1, // 2 in 2A2F - "num_ffn_ranks": 1 // 2 in 2A2F + "port": 6239, + "num_attention_ranks": 2, // Attention GPUs in LAYOUT + "num_ffn_ranks": 2, // FFN GPUs in LAYOUT + "compute_gate_on_attention": false, + "afd_process_group_timeout_s": 3600 } } ``` @@ -150,16 +136,6 @@ DBO (Dual Batch Overlap) is turned on for all examples with These threshold values are example defaults; tune them for the actual workload and deployment configuration. -### Switching eager → graph - -Graph mode replaces `--enforce-eager` with: - -```bash ---max-cudagraph-capture-size 64 ---compilation-config '{"cudagraph_mode": "FULL_DECODE_ONLY", - "cudagraph_capture_sizes":[64]}' -``` - -The capture size of `64` is an example value; adjust -`--max-cudagraph-capture-size` and `cudagraph_capture_sizes` together based on -the actual batch sizes and available GPU memory. +`dbo.sh` captures `FULL_DECODE_ONLY` CUDA graphs of size `64` +(`CUDAGRAPH_CAPTURE_SIZES=64`); `eager_dbo.sh` runs `--enforce-eager`. Adjust +`CUDAGRAPH_CAPTURE_SIZES` to the actual batch sizes and available GPU memory. diff --git a/recipe/gpu/P2pNcclAFDConnector/deepseek_v2_lite/prefill_decode_colocation/2a2f_eager_dbo_dp1tp2.sh b/recipe/gpu/P2pNcclAFDConnector/deepseek_v2_lite/prefill_decode_colocation/2a2f_eager_dbo_dp1tp2.sh deleted file mode 100644 index 1cac638bc..000000000 --- a/recipe/gpu/P2pNcclAFDConnector/deepseek_v2_lite/prefill_decode_colocation/2a2f_eager_dbo_dp1tp2.sh +++ /dev/null @@ -1,55 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project - -MODEL_PATH=${MODEL_PATH:-/path/model_weights/DeepSeek-V2-Lite} -export VLLM_USE_V2_MODEL_RUNNER=1 - -CUDA_VISIBLE_DEVICES=0,1 uv run vllm serve "$MODEL_PATH" \ - --data-parallel-size 1 \ - --tensor-parallel-size 2 \ - --enable-expert-parallel \ - --additional-config '{ - "afd": { - "role": "attention", - "connector": "P2pNcclAFDConnector", - "host": "127.0.0.1", - "port": 6269, - "num_attention_ranks": 2, - "num_ffn_ranks": 2 - } - }' \ - --max-num-seqs 64 \ - --max-num-batched-tokens 64 \ - --enable-dbo \ - --dbo-decode-token-threshold 2 \ - --dbo-prefill-token-threshold 12 \ - --enforce-eager \ - --host 127.0.0.1 \ - --port 18305 \ - --trust-remote-code > attn.log 2>&1 & - -CUDA_VISIBLE_DEVICES=2,3 uv run vllm serve "$MODEL_PATH" \ - --data-parallel-size 1 \ - --tensor-parallel-size 2 \ - --enable-expert-parallel \ - --additional-config '{ - "afd": { - "role": "ffn", - "connector": "P2pNcclAFDConnector", - "host": "127.0.0.1", - "port": 6269, - "num_attention_ranks": 2, - "num_ffn_ranks": 2 - } - }' \ - --max-num-seqs 64 \ - --enable-dbo \ - --dbo-decode-token-threshold 2 \ - --dbo-prefill-token-threshold 12 \ - --max-num-batched-tokens 64 \ - --enforce-eager \ - --host 127.0.0.1 \ - --port 18306 \ - --trust-remote-code > ffn.log 2>&1 & - -wait diff --git a/recipe/gpu/P2pNcclAFDConnector/deepseek_v2_lite/prefill_decode_colocation/2a2f_eager_dbo_dp2tp1.sh b/recipe/gpu/P2pNcclAFDConnector/deepseek_v2_lite/prefill_decode_colocation/2a2f_eager_dbo_dp2tp1.sh deleted file mode 100644 index edbab3812..000000000 --- a/recipe/gpu/P2pNcclAFDConnector/deepseek_v2_lite/prefill_decode_colocation/2a2f_eager_dbo_dp2tp1.sh +++ /dev/null @@ -1,55 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project - -MODEL_PATH=${MODEL_PATH:-/path/model_weights/DeepSeek-V2-Lite} -export VLLM_USE_V2_MODEL_RUNNER=1 - -CUDA_VISIBLE_DEVICES=0,1 uv run vllm serve "$MODEL_PATH" \ - --data-parallel-size 2 \ - --tensor-parallel-size 1 \ - --enable-expert-parallel \ - --additional-config '{ - "afd": { - "role": "attention", - "connector": "P2pNcclAFDConnector", - "host": "127.0.0.1", - "port": 6269, - "num_attention_ranks": 2, - "num_ffn_ranks": 2 - } - }' \ - --max-num-seqs 64 \ - --max-num-batched-tokens 64 \ - --enable-dbo \ - --dbo-decode-token-threshold 2 \ - --dbo-prefill-token-threshold 12 \ - --enforce-eager \ - --host 127.0.0.1 \ - --port 18305 \ - --trust-remote-code > attn.log 2>&1 & - -CUDA_VISIBLE_DEVICES=2,3 uv run vllm serve "$MODEL_PATH" \ - --data-parallel-size 2 \ - --tensor-parallel-size 1 \ - --enable-expert-parallel \ - --additional-config '{ - "afd": { - "role": "ffn", - "connector": "P2pNcclAFDConnector", - "host": "127.0.0.1", - "port": 6269, - "num_attention_ranks": 2, - "num_ffn_ranks": 2 - } - }' \ - --max-num-seqs 64 \ - --enable-dbo \ - --dbo-decode-token-threshold 2 \ - --dbo-prefill-token-threshold 12 \ - --max-num-batched-tokens 64 \ - --enforce-eager \ - --host 127.0.0.1 \ - --port 18306 \ - --trust-remote-code > ffn.log 2>&1 & - -wait diff --git a/recipe/gpu/P2pNcclAFDConnector/deepseek_v2_lite/prefill_decode_colocation/2a2f_graph_dbo_dp1tp2.sh b/recipe/gpu/P2pNcclAFDConnector/deepseek_v2_lite/prefill_decode_colocation/2a2f_graph_dbo_dp1tp2.sh deleted file mode 100644 index e72893d22..000000000 --- a/recipe/gpu/P2pNcclAFDConnector/deepseek_v2_lite/prefill_decode_colocation/2a2f_graph_dbo_dp1tp2.sh +++ /dev/null @@ -1,61 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project - -MODEL_PATH=${MODEL_PATH:-/path/model_weights/DeepSeek-V2-Lite} -export VLLM_USE_V2_MODEL_RUNNER=1 - -CUDA_VISIBLE_DEVICES=0,1 uv run vllm serve "$MODEL_PATH" \ - --data-parallel-size 1 \ - --tensor-parallel-size 2 \ - --enable-expert-parallel \ - --additional-config '{ - "afd": { - "role": "attention", - "connector": "P2pNcclAFDConnector", - "host": "127.0.0.1", - "port": 6269, - "num_attention_ranks": 2, - "num_ffn_ranks": 2 - } - }' \ - --max-num-seqs 64 \ - --max-num-batched-tokens 64 \ - --enable-dbo \ - --dbo-decode-token-threshold 2 \ - --dbo-prefill-token-threshold 12 \ - --max-cudagraph-capture-size 64 \ - --compilation-config '{ - "cudagraph_mode": "FULL_DECODE_ONLY", "cudagraph_capture_sizes":[64] - }' \ - --host 127.0.0.1 \ - --port 18305 \ - --trust-remote-code > attn.log 2>&1 & - -CUDA_VISIBLE_DEVICES=2,3 uv run vllm serve "$MODEL_PATH" \ - --data-parallel-size 1 \ - --tensor-parallel-size 2 \ - --enable-expert-parallel \ - --additional-config '{ - "afd": { - "role": "ffn", - "connector": "P2pNcclAFDConnector", - "host": "127.0.0.1", - "port": 6269, - "num_attention_ranks": 2, - "num_ffn_ranks": 2 - } - }' \ - --max-num-seqs 64 \ - --enable-dbo \ - --dbo-decode-token-threshold 2 \ - --dbo-prefill-token-threshold 12 \ - --max-num-batched-tokens 64 \ - --max-cudagraph-capture-size 64 \ - --compilation-config '{ - "cudagraph_mode": "FULL_DECODE_ONLY", "cudagraph_capture_sizes":[64] - }' \ - --host 127.0.0.1 \ - --port 18306 \ - --trust-remote-code > ffn.log 2>&1 & - -wait diff --git a/recipe/gpu/P2pNcclAFDConnector/deepseek_v2_lite/prefill_decode_colocation/2a2f_graph_dbo_dp2tp1.sh b/recipe/gpu/P2pNcclAFDConnector/deepseek_v2_lite/prefill_decode_colocation/2a2f_graph_dbo_dp2tp1.sh deleted file mode 100644 index c775780db..000000000 --- a/recipe/gpu/P2pNcclAFDConnector/deepseek_v2_lite/prefill_decode_colocation/2a2f_graph_dbo_dp2tp1.sh +++ /dev/null @@ -1,61 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project - -MODEL_PATH=${MODEL_PATH:-/path/model_weights/DeepSeek-V2-Lite} -export VLLM_USE_V2_MODEL_RUNNER=1 - -CUDA_VISIBLE_DEVICES=0,1 uv run vllm serve "$MODEL_PATH" \ - --data-parallel-size 2 \ - --tensor-parallel-size 1 \ - --enable-expert-parallel \ - --additional-config '{ - "afd": { - "role": "attention", - "connector": "P2pNcclAFDConnector", - "host": "127.0.0.1", - "port": 6269, - "num_attention_ranks": 2, - "num_ffn_ranks": 2 - } - }' \ - --max-num-seqs 64 \ - --max-num-batched-tokens 64 \ - --enable-dbo \ - --dbo-decode-token-threshold 2 \ - --dbo-prefill-token-threshold 12 \ - --max-cudagraph-capture-size 64 \ - --compilation-config '{ - "cudagraph_mode": "FULL_DECODE_ONLY", "cudagraph_capture_sizes":[64] - }' \ - --host 127.0.0.1 \ - --port 18305 \ - --trust-remote-code > attn.log 2>&1 & - -CUDA_VISIBLE_DEVICES=2,3 uv run vllm serve "$MODEL_PATH" \ - --data-parallel-size 2 \ - --tensor-parallel-size 1 \ - --enable-expert-parallel \ - --additional-config '{ - "afd": { - "role": "ffn", - "connector": "P2pNcclAFDConnector", - "host": "127.0.0.1", - "port": 6269, - "num_attention_ranks": 2, - "num_ffn_ranks": 2 - } - }' \ - --max-num-seqs 64 \ - --enable-dbo \ - --dbo-decode-token-threshold 2 \ - --dbo-prefill-token-threshold 12 \ - --max-num-batched-tokens 64 \ - --max-cudagraph-capture-size 64 \ - --compilation-config '{ - "cudagraph_mode": "FULL_DECODE_ONLY", "cudagraph_capture_sizes":[64] - }' \ - --host 127.0.0.1 \ - --port 18306 \ - --trust-remote-code > ffn.log 2>&1 & - -wait diff --git a/recipe/gpu/P2pNcclAFDConnector/deepseek_v2_lite/prefill_decode_colocation/4a4f_eager_dbo_dp2tp2.sh b/recipe/gpu/P2pNcclAFDConnector/deepseek_v2_lite/prefill_decode_colocation/4a4f_eager_dbo_dp2tp2.sh deleted file mode 100644 index e2cbf9e71..000000000 --- a/recipe/gpu/P2pNcclAFDConnector/deepseek_v2_lite/prefill_decode_colocation/4a4f_eager_dbo_dp2tp2.sh +++ /dev/null @@ -1,55 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project - -MODEL_PATH=${MODEL_PATH:-/path/model_weights/DeepSeek-V2-Lite} -export VLLM_USE_V2_MODEL_RUNNER=1 - -CUDA_VISIBLE_DEVICES=0,1,2,3 uv run vllm serve "$MODEL_PATH" \ - --data-parallel-size 2 \ - --tensor-parallel-size 2 \ - --enable-expert-parallel \ - --additional-config '{ - "afd": { - "role": "attention", - "connector": "P2pNcclAFDConnector", - "host": "127.0.0.1", - "port": 6269, - "num_attention_ranks": 4, - "num_ffn_ranks": 4 - } - }' \ - --max-num-seqs 64 \ - --max-num-batched-tokens 64 \ - --enable-dbo \ - --dbo-decode-token-threshold 2 \ - --dbo-prefill-token-threshold 12 \ - --enforce-eager \ - --host 127.0.0.1 \ - --port 18305 \ - --trust-remote-code > attn.log 2>&1 & - -CUDA_VISIBLE_DEVICES=4,5,6,7 uv run vllm serve "$MODEL_PATH" \ - --data-parallel-size 2 \ - --tensor-parallel-size 2 \ - --enable-expert-parallel \ - --additional-config '{ - "afd": { - "role": "ffn", - "connector": "P2pNcclAFDConnector", - "host": "127.0.0.1", - "port": 6269, - "num_attention_ranks": 4, - "num_ffn_ranks": 4 - } - }' \ - --max-num-seqs 64 \ - --enable-dbo \ - --dbo-decode-token-threshold 2 \ - --dbo-prefill-token-threshold 12 \ - --max-num-batched-tokens 64 \ - --enforce-eager \ - --host 127.0.0.1 \ - --port 18306 \ - --trust-remote-code > ffn.log 2>&1 & - -wait diff --git a/recipe/gpu/P2pNcclAFDConnector/deepseek_v2_lite/prefill_decode_colocation/4a4f_graph_dbo_dp2tp2.sh b/recipe/gpu/P2pNcclAFDConnector/deepseek_v2_lite/prefill_decode_colocation/4a4f_graph_dbo_dp2tp2.sh deleted file mode 100644 index cd87ccbb4..000000000 --- a/recipe/gpu/P2pNcclAFDConnector/deepseek_v2_lite/prefill_decode_colocation/4a4f_graph_dbo_dp2tp2.sh +++ /dev/null @@ -1,61 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project - -MODEL_PATH=${MODEL_PATH:-/path/model_weights/DeepSeek-V2-Lite} -export VLLM_USE_V2_MODEL_RUNNER=1 - -CUDA_VISIBLE_DEVICES=0,1,2,3 uv run vllm serve "$MODEL_PATH" \ - --data-parallel-size 2 \ - --tensor-parallel-size 2 \ - --enable-expert-parallel \ - --additional-config '{ - "afd": { - "role": "attention", - "connector": "P2pNcclAFDConnector", - "host": "127.0.0.1", - "port": 6269, - "num_attention_ranks": 4, - "num_ffn_ranks": 4 - } - }' \ - --max-num-seqs 64 \ - --max-num-batched-tokens 64 \ - --enable-dbo \ - --dbo-decode-token-threshold 2 \ - --dbo-prefill-token-threshold 12 \ - --max-cudagraph-capture-size 64 \ - --compilation-config '{ - "cudagraph_mode": "FULL_DECODE_ONLY", "cudagraph_capture_sizes":[64] - }' \ - --host 127.0.0.1 \ - --port 18305 \ - --trust-remote-code > attn.log 2>&1 & - -CUDA_VISIBLE_DEVICES=4,5,6,7 uv run vllm serve "$MODEL_PATH" \ - --data-parallel-size 2 \ - --tensor-parallel-size 2 \ - --enable-expert-parallel \ - --additional-config '{ - "afd": { - "role": "ffn", - "connector": "P2pNcclAFDConnector", - "host": "127.0.0.1", - "port": 6269, - "num_attention_ranks": 4, - "num_ffn_ranks": 4 - } - }' \ - --max-num-seqs 64 \ - --enable-dbo \ - --dbo-decode-token-threshold 2 \ - --dbo-prefill-token-threshold 12 \ - --max-num-batched-tokens 64 \ - --max-cudagraph-capture-size 64 \ - --compilation-config '{ - "cudagraph_mode": "FULL_DECODE_ONLY", "cudagraph_capture_sizes":[64] - }' \ - --host 127.0.0.1 \ - --port 18306 \ - --trust-remote-code > ffn.log 2>&1 & - -wait diff --git a/recipe/gpu/P2pNcclAFDConnector/deepseek_v2_lite/prefill_decode_colocation/dbo.sh b/recipe/gpu/P2pNcclAFDConnector/deepseek_v2_lite/prefill_decode_colocation/dbo.sh new file mode 100644 index 000000000..f2af47086 --- /dev/null +++ b/recipe/gpu/P2pNcclAFDConnector/deepseek_v2_lite/prefill_decode_colocation/dbo.sh @@ -0,0 +1,21 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +# +# DeepSeek-V2-Lite (BF16) on P2pNcclAFDConnector, FULL_DECODE_ONLY graphs + DBO. +# Placement: one pod per comma-separated LAYOUT entry (default 2A2F), pod k +# running POD=AFD_ with the same LAYOUT; ATTN_TP/FFN_TP set the TP size +# of each role, e.g. ATTN_TP=2 FFN_TP=2 for DP1 x TP2, LAYOUT=4A4F ATTN_TP=2 +# FFN_TP=2 on 8 GPUs, or LAYOUT=2A0F,0A2F ATTN_TP=2 FFN_TP=2 on two pods. +# Settings and running in a pod: recipe/gpu/common/README.md +# shellcheck disable=SC1090,SC2034 # the sourced launcher reads these variables + +AFD_CONNECTOR=${AFD_CONNECTOR:-P2pNcclAFDConnector} +LAYOUT=${LAYOUT:-2A2F} +MODEL_PATH=${MODEL_PATH:-/path/model_weights/DeepSeek-V2-Lite} +MODE=graph CUDAGRAPH_CAPTURE_SIZES=${CUDAGRAPH_CAPTURE_SIZES:-64} # ascending +SERVE_ARGS=(--max-num-seqs 64 --max-num-batched-tokens 64 + --enable-dbo --dbo-decode-token-threshold 2 --dbo-prefill-token-threshold 12) + +source "${AFD_RECIPE_LIB:-$(dirname "${BASH_SOURCE[0]}")/../../../common/afd_colocation.sh}" +afd_colocation_serve "$@" diff --git a/recipe/gpu/P2pNcclAFDConnector/deepseek_v2_lite/prefill_decode_colocation/eager_dbo.sh b/recipe/gpu/P2pNcclAFDConnector/deepseek_v2_lite/prefill_decode_colocation/eager_dbo.sh new file mode 100644 index 000000000..5d803ec70 --- /dev/null +++ b/recipe/gpu/P2pNcclAFDConnector/deepseek_v2_lite/prefill_decode_colocation/eager_dbo.sh @@ -0,0 +1,21 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +# +# DeepSeek-V2-Lite (BF16) on P2pNcclAFDConnector, eager mode + DBO. +# Placement: one pod per comma-separated LAYOUT entry (default 2A2F), pod k +# running POD=AFD_ with the same LAYOUT; ATTN_TP/FFN_TP set the TP size +# of each role, e.g. ATTN_TP=2 FFN_TP=2 for DP1 x TP2, LAYOUT=4A4F ATTN_TP=2 +# FFN_TP=2 on 8 GPUs, or LAYOUT=2A0F,0A2F ATTN_TP=2 FFN_TP=2 on two pods. +# Settings and running in a pod: recipe/gpu/common/README.md +# shellcheck disable=SC1090,SC2034 # the sourced launcher reads these variables + +AFD_CONNECTOR=${AFD_CONNECTOR:-P2pNcclAFDConnector} +LAYOUT=${LAYOUT:-2A2F} +MODEL_PATH=${MODEL_PATH:-/path/model_weights/DeepSeek-V2-Lite} +MODE=eager +SERVE_ARGS=(--max-num-seqs 64 --max-num-batched-tokens 64 + --enable-dbo --dbo-decode-token-threshold 2 --dbo-prefill-token-threshold 12) + +source "${AFD_RECIPE_LIB:-$(dirname "${BASH_SOURCE[0]}")/../../../common/afd_colocation.sh}" +afd_colocation_serve "$@" diff --git a/recipe/gpu/P2pNcclAFDConnector/deepseek_v2_lite/prefill_decode_colocation/multipod_2a_2f_graph_dbo_dp1tp2.sh b/recipe/gpu/P2pNcclAFDConnector/deepseek_v2_lite/prefill_decode_colocation/multipod_2a_2f_graph_dbo_dp1tp2.sh deleted file mode 100644 index b1a1f0b14..000000000 --- a/recipe/gpu/P2pNcclAFDConnector/deepseek_v2_lite/prefill_decode_colocation/multipod_2a_2f_graph_dbo_dp1tp2.sh +++ /dev/null @@ -1,73 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project - -# Use .agents/skills/deploy-afd-k8s/SKILL.md to deploy this recipe -# on 2 pods using the setup below. -# For pod 1: POD=ATTENTION_0 -# For pod 2: POD=FFN_0 - -MODEL_PATH=${MODEL_PATH:-/path/model_weights/DeepSeek-V2-Lite} -export VLLM_USE_V2_MODEL_RUNNER=1 - -AFD_CONNECTOR_HOST=${AFD_CONNECTOR_HOST:-vllm-ffn-p2p-service} -AFD_CONNECTOR_PORT=${AFD_CONNECTOR_PORT:-6269} - -if [ "$POD" == "ATTENTION_0" ]; then - CUDA_VISIBLE_DEVICES=0,1 uv run vllm serve "$MODEL_PATH" \ - --data-parallel-size 1 \ - --tensor-parallel-size 2 \ - --enable-expert-parallel \ - --additional-config '{ - "afd": { - "role": "attention", - "connector": "P2pNcclAFDConnector", - "host": "'"${AFD_CONNECTOR_HOST}"'", - "port": '"${AFD_CONNECTOR_PORT}"', - "num_attention_ranks": 2, - "num_ffn_ranks": 2 - } - }' \ - --max-num-seqs 64 \ - --max-num-batched-tokens 64 \ - --enable-dbo \ - --dbo-decode-token-threshold 2 \ - --dbo-prefill-token-threshold 12 \ - --max-cudagraph-capture-size 64 \ - --compilation-config '{ - "cudagraph_mode": "FULL_DECODE_ONLY", "cudagraph_capture_sizes":[64] - }' \ - --host 127.0.0.1 \ - --port 18305 \ - --trust-remote-code > attn.log 2>&1 & -fi - -if [ "$POD" == "FFN_0" ]; then - CUDA_VISIBLE_DEVICES=0,1 uv run vllm serve "$MODEL_PATH" \ - --data-parallel-size 1 \ - --tensor-parallel-size 2 \ - --enable-expert-parallel \ - --additional-config '{ - "afd": { - "role": "ffn", - "connector": "P2pNcclAFDConnector", - "host": "'"${AFD_CONNECTOR_HOST}"'", - "port": '"${AFD_CONNECTOR_PORT}"', - "num_attention_ranks": 2, - "num_ffn_ranks": 2 - } - }' \ - --max-num-seqs 64 \ - --enable-dbo \ - --dbo-decode-token-threshold 2 \ - --dbo-prefill-token-threshold 12 \ - --max-num-batched-tokens 64 \ - --max-cudagraph-capture-size 64 \ - --compilation-config '{ - "cudagraph_mode": "FULL_DECODE_ONLY", "cudagraph_capture_sizes":[64] - }' \ - --host 127.0.0.1 \ - --port 18305 \ - --trust-remote-code > ffn.log 2>&1 & -fi - -wait diff --git a/recipe/gpu/P2pNcclAFDConnector/qwen3_235b_a22b_fp8/README.md b/recipe/gpu/P2pNcclAFDConnector/qwen3_235b_a22b_fp8/README.md index c89af5fd7..9cc63e56b 100644 --- a/recipe/gpu/P2pNcclAFDConnector/qwen3_235b_a22b_fp8/README.md +++ b/recipe/gpu/P2pNcclAFDConnector/qwen3_235b_a22b_fp8/README.md @@ -1,17 +1,16 @@ # Qwen3-235B-A22B-FP8 AFD Recipe (P2pNcclAFDConnector) -`prefill_decode_colocation/4a2f_graph_dp4tp1.sh` runs the shared -Qwen3-235B-A22B-FP8 graph recipe -(`DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/graph.sh`) -in the 4A2F layout (`LAYOUT=4A2F`, one 6-GPU node) with -`AFD_CONNECTOR=P2pNcclAFDConnector`. Its layouts, settings and the -decode-only benchmark are documented in the -[`DeepEPM2NAFDConnector` recipes README](../../DeepEPM2NAFDConnector/README.md) -and the [model README](../../DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/README.md). +`prefill_decode_colocation/graph.sh` runs Qwen3-235B-A22B-FP8 with +`FULL_DECODE_ONLY` CUDA graphs on `P2pNcclAFDConnector`, with the same +settings as the +[`DeepEPM2NAFDConnector` graph recipe](../../DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/README.md), +in the 4A2F layout by default (one 6-GPU node). The layout, settings and +running in a pod are documented in the +[colocation launcher README](../../common/README.md). ```bash export MODEL_PATH=/path/model_weights/Qwen3-235B-A22B-FP8 -bash recipe/gpu/P2pNcclAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/4a2f_graph_dp4tp1.sh +bash recipe/gpu/P2pNcclAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/graph.sh ``` Do not set `PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True` for this diff --git a/recipe/gpu/P2pNcclAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/4a2f_graph_dp4tp1.sh b/recipe/gpu/P2pNcclAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/4a2f_graph_dp4tp1.sh deleted file mode 100755 index a494f68f1..000000000 --- a/recipe/gpu/P2pNcclAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/4a2f_graph_dp4tp1.sh +++ /dev/null @@ -1,13 +0,0 @@ -#!/usr/bin/env bash -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project -# -# P2pNcclAFDConnector entry point of the shared Qwen3-235B-A22B-FP8 graph -# recipe in the 4A2F layout (one 6-GPU node), -# recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/graph.sh, -# which documents the layout and settings. -set -euo pipefail - -SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" -LAYOUT=${LAYOUT:-4A2F} AFD_CONNECTOR=P2pNcclAFDConnector exec bash \ - "${SCRIPT_DIR}/../../../DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/graph.sh" "$@" diff --git a/recipe/gpu/P2pNcclAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/graph.sh b/recipe/gpu/P2pNcclAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/graph.sh new file mode 100644 index 000000000..cd24d2468 --- /dev/null +++ b/recipe/gpu/P2pNcclAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/graph.sh @@ -0,0 +1,20 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +# +# Qwen3-235B-A22B-FP8 on P2pNcclAFDConnector, FULL_DECODE_ONLY CUDA graphs. +# Placement: one pod per comma-separated LAYOUT entry (default 4A2F), pod k +# running POD=AFD_ with the same LAYOUT; the default is one +# 6-GPU node. Do not set PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True: the +# P2P FFN hangs on its first step. +# Settings and running in a pod: recipe/gpu/common/README.md +# shellcheck disable=SC1090,SC2034 # the sourced launcher reads these variables + +AFD_CONNECTOR=${AFD_CONNECTOR:-P2pNcclAFDConnector} +LAYOUT=${LAYOUT:-4A2F} +MODEL_PATH=${MODEL_PATH:-/path/model_weights/Qwen3-235B-A22B-FP8} +MODE=graph CUDAGRAPH_CAPTURE_SIZES=${CUDAGRAPH_CAPTURE_SIZES:-8,12,16,24,32,48,64,96,128,192,256} # ascending +SERVE_ARGS=(--enable-ep-weight-filter --max-num-seqs 256 --max-num-batched-tokens 3072 --max-model-len 4096) + +source "${AFD_RECIPE_LIB:-$(dirname "${BASH_SOURCE[0]}")/../../../common/afd_colocation.sh}" +afd_colocation_serve "$@" diff --git a/recipe/gpu/P2pNcclAFDConnector/qwen3_30b_a3b/README.md b/recipe/gpu/P2pNcclAFDConnector/qwen3_30b_a3b/README.md index 41ea965bb..d3692de3a 100644 --- a/recipe/gpu/P2pNcclAFDConnector/qwen3_30b_a3b/README.md +++ b/recipe/gpu/P2pNcclAFDConnector/qwen3_30b_a3b/README.md @@ -1,13 +1,13 @@ # Qwen3-30B-A3B AFD Recipe (P2pNcclAFDConnector) -`prefill_decode_colocation/2a2f_graph_dp2tp1.sh` runs the shared Qwen3-30B-A3B -graph recipe (`DeepEPM2NAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/graph.sh`, -default layout 2A2F) with `AFD_CONNECTOR=P2pNcclAFDConnector`. Its layouts, -settings and the decode-only benchmark are documented in the -[`DeepEPM2NAFDConnector` recipes README](../../DeepEPM2NAFDConnector/README.md) -and the [model README](../../DeepEPM2NAFDConnector/qwen3_30b_a3b/README.md). +`prefill_decode_colocation/graph.sh` runs Qwen3-30B-A3B with +`FULL_DECODE_ONLY` CUDA graphs on `P2pNcclAFDConnector`, with the same +settings as the +[`DeepEPM2NAFDConnector` graph recipe](../../DeepEPM2NAFDConnector/qwen3_30b_a3b/README.md) +(default layout 2A2F). The layout, settings and running in a pod are +documented in the [colocation launcher README](../../common/README.md). ```bash export MODEL_PATH=/path/model_weights/Qwen3-30B-A3B -bash recipe/gpu/P2pNcclAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/2a2f_graph_dp2tp1.sh +bash recipe/gpu/P2pNcclAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/graph.sh ``` diff --git a/recipe/gpu/P2pNcclAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/2a2f_graph_dp2tp1.sh b/recipe/gpu/P2pNcclAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/2a2f_graph_dp2tp1.sh deleted file mode 100755 index f7e63fa87..000000000 --- a/recipe/gpu/P2pNcclAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/2a2f_graph_dp2tp1.sh +++ /dev/null @@ -1,12 +0,0 @@ -#!/usr/bin/env bash -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project -# -# P2pNcclAFDConnector entry point of the shared Qwen3-30B-A3B graph recipe, -# recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/graph.sh, -# which documents the layout and settings. -set -euo pipefail - -SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" -AFD_CONNECTOR=P2pNcclAFDConnector exec bash \ - "${SCRIPT_DIR}/../../../DeepEPM2NAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/graph.sh" "$@" diff --git a/recipe/gpu/P2pNcclAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/graph.sh b/recipe/gpu/P2pNcclAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/graph.sh new file mode 100644 index 000000000..9b1dd03c5 --- /dev/null +++ b/recipe/gpu/P2pNcclAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/graph.sh @@ -0,0 +1,18 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +# +# Qwen3-30B-A3B (BF16) on P2pNcclAFDConnector, FULL_DECODE_ONLY CUDA graphs. +# Placement: one pod per comma-separated LAYOUT entry (default 2A2F), pod k +# running POD=AFD_ with the same LAYOUT; FFN on each pod's first GPUs. +# Settings and running in a pod: recipe/gpu/common/README.md +# shellcheck disable=SC1090,SC2034 # the sourced launcher reads these variables + +AFD_CONNECTOR=${AFD_CONNECTOR:-P2pNcclAFDConnector} +LAYOUT=${LAYOUT:-2A2F} +MODEL_PATH=${MODEL_PATH:-/path/model_weights/Qwen3-30B-A3B} +MODE=graph CUDAGRAPH_CAPTURE_SIZES=${CUDAGRAPH_CAPTURE_SIZES:-8,16,32,64} # ascending +SERVE_ARGS=(--enable-ep-weight-filter --max-num-seqs 64 --max-num-batched-tokens 4096 --max-model-len 4096) + +source "${AFD_RECIPE_LIB:-$(dirname "${BASH_SOURCE[0]}")/../../../common/afd_colocation.sh}" +afd_colocation_serve "$@" diff --git a/recipe/gpu/P2pNcclAFDConnector/qwen3_5_122b_a10b_fp8/prefill_decode_colocation/graph.sh b/recipe/gpu/P2pNcclAFDConnector/qwen3_5_122b_a10b_fp8/prefill_decode_colocation/graph.sh new file mode 100644 index 000000000..5582c73ee --- /dev/null +++ b/recipe/gpu/P2pNcclAFDConnector/qwen3_5_122b_a10b_fp8/prefill_decode_colocation/graph.sh @@ -0,0 +1,21 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +# +# Qwen3.5-122B-A10B-FP8 on P2pNcclAFDConnector, FULL_DECODE_ONLY CUDA graphs. +# Placement: one pod per comma-separated LAYOUT entry (default 2A0F,2A0F,0A2F), pod k +# running POD=AFD_ with the same LAYOUT; the default runs +# 4A2F on three 2-GPU pods (Attention DP4 over pods 0-1, FFN DP2 on pod 2); +# LAYOUT=4A2F runs it on one 6-GPU node. +# Settings and running in a pod: recipe/gpu/common/README.md +# shellcheck disable=SC1090,SC2034 # the sourced launcher reads these variables + +AFD_CONNECTOR=${AFD_CONNECTOR:-P2pNcclAFDConnector} +LAYOUT=${LAYOUT:-2A0F,2A0F,0A2F} +MODEL_PATH=${MODEL_PATH:-/path/model_weights/Qwen3.5-122B-A10B-FP8} +MODE=graph CUDAGRAPH_CAPTURE_SIZES=${CUDAGRAPH_CAPTURE_SIZES:-32} # ascending +SERVE_ARGS=(--dtype bfloat16 --language-model-only --max-model-len 114688 --mamba-cache-mode align + --all2all-backend allgather_reducescatter --seed 0 --max-num-seqs 32 --max-num-batched-tokens 8192) + +source "${AFD_RECIPE_LIB:-$(dirname "${BASH_SOURCE[0]}")/../../../common/afd_colocation.sh}" +afd_colocation_serve "$@" diff --git a/recipe/gpu/P2pNcclAFDConnector/qwen3_5_122b_a10b_fp8/prefill_decode_colocation/multipod_2a_2a_2f_graph.sh b/recipe/gpu/P2pNcclAFDConnector/qwen3_5_122b_a10b_fp8/prefill_decode_colocation/multipod_2a_2a_2f_graph.sh deleted file mode 100644 index ec63badbd..000000000 --- a/recipe/gpu/P2pNcclAFDConnector/qwen3_5_122b_a10b_fp8/prefill_decode_colocation/multipod_2a_2a_2f_graph.sh +++ /dev/null @@ -1,100 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project - -# Use .agents/skills/deploy-afd-k8s/SKILL.md to deploy this recipe -# on 3 pods using the setup below. -# For pod 1: POD=ATTENTION_0 -# For pod 2: POD=ATTENTION_1, ATTENTION_HEADLESS=1, ATTENTION_DP_START_RANK=2 -# For pod 3: POD=FFN_0 - -MODEL_PATH=${MODEL_PATH:-/path/model_weights/Qwen3.5-122B-A10B-FP8} -export VLLM_USE_V2_MODEL_RUNNER=1 - -AFD_CONNECTOR_HOST=${AFD_CONNECTOR_HOST:-vllm-ffn-p2p-service} -AFD_CONNECTOR_PORT=${AFD_CONNECTOR_PORT:-1239} - -ATTENTION_DP_START_RANK=${ATTENTION_DP_START_RANK:-0} -ATTENTION_HEADLESS=${ATTENTION_HEADLESS:-0} -ATTENTION_DP_ADDRESS=${ATTENTION_DP_ADDRESS:-vllm-attn-dp-service} -ATTENTION_DP_RPC_PORT=${ATTENTION_DP_RPC_PORT:-13345} - -HEADLESS_ARGS=() -if [[ "$ATTENTION_HEADLESS" == "1" ]]; then - HEADLESS_ARGS=(--headless) -fi - -DP_START_RANK_ARGS=() -if [[ "$ATTENTION_DP_START_RANK" != "0" ]]; then - DP_START_RANK_ARGS=(--data-parallel-start-rank "$ATTENTION_DP_START_RANK") -fi - -if [[ "$POD" == ATTENTION_* ]]; then - CUDA_VISIBLE_DEVICES=0,1 uv run vllm serve "$MODEL_PATH" \ - --data-parallel-size 4 \ - --data-parallel-size-local 2 \ - "${DP_START_RANK_ARGS[@]}" \ - --data-parallel-address "$ATTENTION_DP_ADDRESS" \ - --data-parallel-rpc-port "$ATTENTION_DP_RPC_PORT" \ - --tensor-parallel-size 1 \ - --enable-expert-parallel \ - --dtype bfloat16 \ - --language-model-only \ - --max-model-len 114688 \ - --mamba-cache-mode align \ - --all2all-backend allgather_reducescatter \ - --seed 0 \ - --additional-config '{ - "afd": { - "role": "attention", - "connector": "P2pNcclAFDConnector", - "afd_process_group_timeout_s": 120, - "host": "'"${AFD_CONNECTOR_HOST}"'", - "port": '"${AFD_CONNECTOR_PORT}"', - "num_attention_ranks": 4, - "num_ffn_ranks": 2 - } - }' \ - --max-num-seqs 32 \ - --max-num-batched-tokens 8192 \ - --max-cudagraph-capture-size 32 \ - --compilation-config '{ - "cudagraph_mode": "FULL_DECODE_ONLY", "cudagraph_capture_sizes":[32] - }' \ - "${HEADLESS_ARGS[@]}" \ - --host 127.0.0.1 \ - --port 18305 > attn.log 2>&1 & -fi - -if [[ "$POD" == FFN_* ]]; then - CUDA_VISIBLE_DEVICES=0,1 uv run vllm serve "$MODEL_PATH" \ - --data-parallel-size 2 \ - --tensor-parallel-size 1 \ - --enable-expert-parallel \ - --dtype bfloat16 \ - --language-model-only \ - --max-model-len 114688 \ - --mamba-cache-mode align \ - --all2all-backend allgather_reducescatter \ - --seed 0 \ - --additional-config '{ - "afd": { - "role": "ffn", - "connector": "P2pNcclAFDConnector", - "afd_process_group_timeout_s": 120, - "host": "'"${AFD_CONNECTOR_HOST}"'", - "port": '"${AFD_CONNECTOR_PORT}"', - "num_attention_ranks": 4, - "num_ffn_ranks": 2 - } - }' \ - --max-num-seqs 32 \ - --max-num-batched-tokens 8192 \ - --max-cudagraph-capture-size 32 \ - --compilation-config '{ - "cudagraph_mode": "FULL_DECODE_ONLY", "cudagraph_capture_sizes":[32] - }' \ - --host 127.0.0.1 \ - --port 18305 > ffn.log 2>&1 & -fi - -wait diff --git a/recipe/gpu/common/README.md b/recipe/gpu/common/README.md new file mode 100644 index 000000000..f0c3cf5c1 --- /dev/null +++ b/recipe/gpu/common/README.md @@ -0,0 +1,156 @@ +# Shared launcher for the GPU colocation recipes + +Every GPU `prefill_decode_colocation` recipe (`DeepEPM2NAFDConnector` and +`P2pNcclAFDConnector`) is a short script that states a model, a connector, a +default layout and its `vllm serve` flags, and then hands over to +[`afd_colocation.sh`](afd_colocation.sh). The launcher owns everything the +recipes share: the placement (`LAYOUT` and `POD` to GPUs and data-parallel +flags), the `afd` additional config, the decode-only benchmark hook, the +override variables, and the two `vllm serve` processes. + +## Writing a recipe + +```bash +#!/usr/bin/env bash +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +# +# Qwen3-30B-A3B (BF16) on DeepEPM2NAFDConnector, FULL_DECODE_ONLY graphs + DBO. +# Placement: one pod per comma-separated LAYOUT entry (default 2A2F), pod k +# running POD=AFD_ with the same LAYOUT; FFN on each pod's first GPUs. +# Settings and running in a pod: recipe/gpu/common/README.md +# shellcheck disable=SC1090,SC2034 # the sourced launcher reads these variables + +AFD_CONNECTOR=${AFD_CONNECTOR:-DeepEPM2NAFDConnector} +LAYOUT=${LAYOUT:-2A2F} +MODEL_PATH=${MODEL_PATH:-/path/model_weights/Qwen3-30B-A3B} +MODE=graph CUDAGRAPH_CAPTURE_SIZES=${CUDAGRAPH_CAPTURE_SIZES:-8,16,32,64} # ascending +SERVE_ARGS=(--enable-ep-weight-filter --max-num-seqs 64 --max-num-batched-tokens 4096 --max-model-len 4096 + --enable-dbo --dbo-decode-token-threshold 2 --dbo-prefill-token-threshold 8) + +source "${AFD_RECIPE_LIB:-$(dirname "${BASH_SOURCE[0]}")/../../../common/afd_colocation.sh}" +afd_colocation_serve "$@" +``` + +| Recipe variable | Meaning | +| --- | --- | +| `AFD_CONNECTOR` | Connector class. `P2pNcclAFDConnector` routes on the FFN and drops `CONNECTOR_CONFIG`; any other connector computes the gate on Attention. | +| `LAYOUT` | Default placement; see [Layout](#layout). | +| `MODEL_PATH` | Default weights directory or HF id. | +| `MODE` | `graph` (`FULL_DECODE_ONLY` CUDA graphs of `CUDAGRAPH_CAPTURE_SIZES`, the largest being the maximum capture size) or `eager` (`--enforce-eager`). | +| `SERVE_ARGS` | Array of model and tuning flags passed to both roles (DBO, `--enable-ep-weight-filter`, batch limits, dtype, ...). | +| `CONNECTOR_CONFIG` | Optional `connector_extra_config` members, for example `'"payload_format": "fp8_block128_fp32"'`. | + +The launcher always adds `--enable-expert-parallel`, `--trust-remote-code`, +`--cpu-distributed-timeout-seconds 3600` and `afd_process_group_timeout_s: +3600` (the roles can finish loading minutes apart). Its ports are fixed: +`18000` Attention API, `18001` FFN (serves no requests), `6239` AFD +rendezvous, `13345` and `13355` the Attention and FFN DP RPC ports across +pods. The default launcher path assumes the recipe sits at +`recipe/gpu////`. + +## Layout + +`LAYOUT` lists the pods in order, one `AF` entry each: `a` Attention +GPUs and `f` FFN GPUs. Inside a pod, the FFN takes the first `f` GPUs and the +Attention the next `a`; every pod runs the same script with the same +`LAYOUT` and its own `POD=AFD_` (`k` counts from 0; one pod needs no +`POD`). + +| `LAYOUT` | Pods | Per pod | +| --- | --- | --- | +| `2A2F` | 1 | FFN DP2 on GPUs 0,1, Attention DP2 on GPUs 2,3 | +| `4A2F` | 1 | FFN on GPUs 0,1, Attention DP4 on GPUs 2-5 | +| `2A0F,0A2F` | 2 | pod 0: Attention DP2; pod 1: FFN DP2 | +| `1A1F,1A1F` | 2 | FFN on GPU 0, Attention on GPU 1 | +| `2A0F,2A0F,0A2F` | 3 | pods 0-1: Attention DP4 (pod 1 headless); pod 2: FFN DP2 | + +`ATTN_TP` and `FFN_TP` split each pod's GPUs of that role into TP groups (they +must divide each pod's `a` or `f`); `num_attention_ranks` and `num_ffn_ranks` +still count GPUs. A role held by one pod gets only `--data-parallel-size`; a +role spread over pods is led by its first pod, and the others start headless +from their start rank. Each connector checks the placement at startup; see +its user guide for the layouts it accepts. + +## Settings + +Every setting is an environment variable; the defaults reproduce the recipe. + +| Variable | Default | Meaning | +| --- | --- | --- | +| `MODEL_PATH` | recipe | Weights directory or HF id. | +| `LAYOUT` | recipe | Pods and their GPUs; see [Layout](#layout). | +| `POD` | `AFD_0` | This pod's entry in `LAYOUT`, `AFD_`. | +| `ATTN_TP` / `FFN_TP` | `1` | TP size of the Attention / FFN role. | +| `CUDAGRAPH_CAPTURE_SIZES` | recipe | Ascending decode graph sizes (graph recipes). | +| `AFD_CONNECTOR` | recipe | Connector class, for example `P2pNcclAFDConnector` to run an M2N recipe on NCCL P2P. | +| `VLLM_ARGS` | empty | Extra `vllm serve` flags for both roles. | +| `ATTN_ARGS` / `FFN_ARGS` | empty | Extra `vllm serve` flags for one role. | +| `M2N_EXTRA_CONFIG` | empty | Extra `connector_extra_config` members, for example `'"expert_num_sms": 130'`. | +| `DECODE_BENCH` | empty | `kv_connector_extra_config` JSON of the decode-only benchmark connector on Attention (`{}` is enough); see [`tools/benchmarks`](../../../tools/benchmarks/). | +| `AFD_REPO` | current directory | Repository root that holds `tools/benchmarks` (with `DECODE_BENCH`). | +| `HEAD_HOST` | `127.0.0.1` on one pod | Address of pod 0 when it leads both roles (several pods, no Kubernetes). | +| `ATTN_HEAD_HOST` / `FFN_HEAD_HOST` | `HEAD_HOST`, else the Kubernetes Service names | Address of the first pod with Attention / FFN GPUs. The FFN head also hosts the AFD rendezvous. | +| `SERVE_HOST` | `127.0.0.1` | API server bind address (`0.0.0.0` to serve outside the pod). | +| `AFD_RECIPE_LIB` | next to the recipe in the repository | Path of `afd_colocation.sh`, for a recipe copied out of the repository. | +| `AFD_DRY_RUN` | empty | Print this pod's `vllm serve` commands instead of running them. | + +`VLLM_ARGS`, `ATTN_ARGS`, `FFN_ARGS` and any arguments after the script name +come after the recipe's flags, so a repeated flag overrides its default. They +are split on whitespace without shell quoting, so pass JSON without spaces: + +```bash +VLLM_ARGS="--max-num-seqs 512" ATTN_ARGS="--no-enable-prefix-caching" \ + bash recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/graph.sh +``` + +## Running directly + +From any directory (logs go to the current directory): + +```bash +export MODEL_PATH=/path/model_weights/Qwen3-30B-A3B +bash recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/graph.sh +``` + +The script starts the FFN and Attention servers in the background, writes +`ffn.log` and `attn.log`, and waits. Send traffic to port `18000` once +`attn.log` prints `Application startup complete`. To see what a pod would +run without starting anything: + +```bash +AFD_DRY_RUN=1 LAYOUT=2A0F,0A2F POD=AFD_1 bash .../graph.sh +``` + +Several hosts without Kubernetes run the same script with their own `POD` +and the leaders' addresses: + +```bash +LAYOUT=2A0F,0A2F POD=AFD_0 ATTN_HEAD_HOST= FFN_HEAD_HOST= bash .../graph.sh +LAYOUT=2A0F,0A2F POD=AFD_1 ATTN_HEAD_HOST= FFN_HEAD_HOST= bash .../graph.sh +``` + +## Running in a pod + +A recipe needs `afd_colocation.sh` next to it. Ship both files and point +`AFD_RECIPE_LIB` at the launcher: + +- **`deploy-afd-k8s` skill** (`.agents/skills/deploy-afd-k8s/`): puts both + files into the `afd-recipe-script` ConfigMap (`recipe.sh` and + `afd_colocation.sh`), mounts it at `/recipe`, and sets + `AFD_RECIPE_LIB=/recipe/afd_colocation.sh` and `SERVE_HOST=0.0.0.0`. It + creates one pod per `LAYOUT` entry and the Services that the default head + host names resolve to. Edited recipes need no image rebuild. +- **A pod you manage**: copy both files and run the recipe with the launcher + path: + + ```bash + kubectl cp recipe/gpu/common/afd_colocation.sh :/work/afd_colocation.sh + kubectl cp .sh :/work/recipe.sh + kubectl exec -- bash -c 'cd /work && SERVE_HOST=0.0.0.0 \ + AFD_RECIPE_LIB=/work/afd_colocation.sh MODEL_PATH=/models/ bash recipe.sh' + ``` + +- **An image with this repository** (`docker/Dockerfile.ci` puts it in + `/opt/afd-plugin`): run the recipe in place, as the image's copy finds its + launcher: `bash /opt/afd-plugin/recipe/gpu///prefill_decode_colocation/.sh`. diff --git a/recipe/gpu/common/afd_colocation.sh b/recipe/gpu/common/afd_colocation.sh new file mode 100644 index 000000000..2abf78226 --- /dev/null +++ b/recipe/gpu/common/afd_colocation.sh @@ -0,0 +1,116 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +# +# Shared launcher of the GPU prefill_decode_colocation recipes: placement +# (LAYOUT/POD -> devices and DP flags), the AFD config and the two `vllm serve` +# processes. A recipe sets the variables below, sources this file and calls +# afd_colocation_serve "$@". See recipe/gpu/common/README.md. +# +# Set by the recipe: +# MODEL_PATH, AFD_CONNECTOR, LAYOUT model, connector class, default layout +# MODE graph (FULL_DECODE_ONLY CUDA graphs of +# CUDAGRAPH_CAPTURE_SIZES) or eager +# SERVE_ARGS array of model/tuning flags, both roles +# CONNECTOR_CONFIG optional connector_extra_config members +# Run settings (environment): POD, ATTN_TP, FFN_TP, VLLM_ARGS, ATTN_ARGS, +# FFN_ARGS, M2N_EXTRA_CONFIG, DECODE_BENCH, AFD_REPO, HEAD_HOST, +# ATTN_HEAD_HOST, FFN_HEAD_HOST, SERVE_HOST, AFD_DRY_RUN. + +ATTN_PORT=18000 FFN_PORT=18001 AFD_PORT=6239 ATTN_DP_RPC_PORT=13345 FFN_DP_RPC_PORT=13355 + +# role_dp TP DP_HOST DP_RPC_PORT ... -> TOTAL and this pod's DP +# flags in DP. A role spread over pods is led by its first pod; the others run +# headless from their start rank (the leader must omit it, or vLLM selects +# hybrid load balancing and rejects the headless peers). +role_dp() { + local tp=$1 host=$2 port=$3 i pods=0 start=0 lead=-1 + shift 3 + local n=("$@") + TOTAL=0 + for i in "${!n[@]}"; do + ((n[i])) || continue + ((n[i] % tp)) && { echo "TP $tp does not divide ${n[i]} ranks on pod $i" >&2; exit 1; } + ((lead < 0)) && lead=$i + ((i < K)) && start=$((start + n[i])) + pods=$((pods + 1)) TOTAL=$((TOTAL + n[i])) + done + DP=(--tensor-parallel-size "$tp" --data-parallel-size $((TOTAL / tp))) + ((pods > 1)) && DP+=(--data-parallel-size-local $((n[K] / tp)) + --data-parallel-address "$host" --data-parallel-rpc-port "$port") + ((pods > 1 && K != lead)) && DP+=(--headless --data-parallel-start-rank $((start / tp))) + return 0 +} + +afd_json() { # The AFD rendezvous runs on the pod that holds FFN rank 0. + cat <... +launch() { + local log=$1 devices + devices=$(seq "$2" $(($2 + $3 - 1)) | paste -sd, -) + shift 3 + if [[ -n ${AFD_DRY_RUN:-} ]]; then + echo "CUDA_VISIBLE_DEVICES=$devices vllm serve $(printf '%q ' "$@")> $log" + else + CUDA_VISIBLE_DEVICES=$devices vllm serve "$@" > "$log" 2>&1 & + fi +} + +afd_colocation_serve() { + export VLLM_USE_V2_MODEL_RUNNER=1 + local p + case ${MODE:-} in + graph) MODE_ARGS=(--max-cudagraph-capture-size "${CUDAGRAPH_CAPTURE_SIZES##*,}" + --compilation-config "{\"cudagraph_mode\": \"FULL_DECODE_ONLY\", \"cudagraph_capture_sizes\": [${CUDAGRAPH_CAPTURE_SIZES}]}") ;; + eager) MODE_ARGS=(--enforce-eager) ;; + *) echo "MODE must be graph or eager, got '${MODE:-}'" >&2; exit 1 ;; + esac + + GATE_ON_ATTENTION=true + CONNECTOR_CONFIG=${CONNECTOR_CONFIG:-}${M2N_EXTRA_CONFIG:+${CONNECTOR_CONFIG:+, }$M2N_EXTRA_CONFIG} + # P2P routes on the FFN and takes no connector_extra_config. + [[ $AFD_CONNECTOR == P2pNcclAFDConnector ]] && GATE_ON_ATTENTION=false CONNECTOR_CONFIG= + [[ -n $CONNECTOR_CONFIG ]] && CONNECTOR_CONFIG=", \"connector_extra_config\": {${CONNECTOR_CONFIG}}" + + IFS=, read -ra PODS <<<"$LAYOUT" + K=${POD:-AFD_0} && K=${K#AFD_} + if [[ ! $LAYOUT =~ ^[0-9]+A[0-9]+F(,[0-9]+A[0-9]+F)*$ || ! $K =~ ^[0-9]+$ ]] || ((K >= ${#PODS[@]})); then + echo "need LAYOUT=AF[,...] and POD=AFD_ in it: LAYOUT=$LAYOUT POD=${POD:-}" >&2 + exit 1 + fi + NUM_A=() NUM_F=() + for p in "${PODS[@]}"; do NUM_A+=("${p%%A*}") && p=${p#*A} && NUM_F+=("${p%F}"); done + ((${#PODS[@]} > 1)) || HEAD_HOST=${HEAD_HOST:-127.0.0.1} + + # Extra `vllm serve` flags go last, so they override the defaults. + read -ra ARGS <<<"${VLLM_ARGS:-}" && ARGS+=("$@") + read -ra ATTN_EXTRA <<<"${ATTN_ARGS:-}" + read -ra FFN_EXTRA <<<"${FFN_ARGS:-}" + if [[ -n ${DECODE_BENCH:-} ]]; then # tools/benchmarks is imported from the repository + ATTN_EXTRA=(--kv-transfer-config "{\"kv_connector\": \"AFDDecodeBenchConnector\", \"kv_connector_module_path\": \"tools.benchmarks.decode_bench\", \"kv_role\": \"kv_both\", \"kv_connector_extra_config\": ${DECODE_BENCH}}" "${ATTN_EXTRA[@]}") + export PYTHONPATH=${AFD_REPO:-$PWD}${PYTHONPATH:+:$PYTHONPATH} + fi + + role_dp "${FFN_TP:-1}" "${FFN_HEAD_HOST:-${HEAD_HOST:-vllm-ffn-dp-service}}" "$FFN_DP_RPC_PORT" "${NUM_F[@]}" + FFN_DP=("${DP[@]}") NUM_FFN=$TOTAL + role_dp "${ATTN_TP:-1}" "${ATTN_HEAD_HOST:-${HEAD_HOST:-vllm-attn-dp-service}}" "$ATTN_DP_RPC_PORT" "${NUM_A[@]}" + ATTN_DP=("${DP[@]}") NUM_ATTN=$TOTAL + + # The roles can finish loading minutes apart, hence the 3600 s timeouts. + COMMON=(--enable-expert-parallel --trust-remote-code --cpu-distributed-timeout-seconds 3600 + "${SERVE_ARGS[@]}" "${MODE_ARGS[@]}" --host "${SERVE_HOST:-127.0.0.1}") + + # FFN on this pod's GPUs [0, F), Attention on [F, F + A). + local F=${NUM_F[K]} A=${NUM_A[K]} + ((F)) && launch ffn.log 0 "$F" "$MODEL_PATH" --additional-config "$(afd_json ffn)" \ + "${COMMON[@]}" "${FFN_DP[@]}" "${ARGS[@]}" "${FFN_EXTRA[@]}" --port "$FFN_PORT" + ((A)) && launch attn.log "$F" "$A" "$MODEL_PATH" --additional-config "$(afd_json attention)" \ + "${COMMON[@]}" "${ATTN_DP[@]}" "${ARGS[@]}" "${ATTN_EXTRA[@]}" --port "$ATTN_PORT" + wait +}