diff --git a/.agents/skills/deploy-afd-k8s/SKILL.md b/.agents/skills/deploy-afd-k8s/SKILL.md index f51ae4190..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,36 +15,38 @@ 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 | +| 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 -[services.md](references/services.md). +[services.md](resources/services.md). ## 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 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](references/resolve-recipe.md). -- Not for `tools/benchmarks/decode_bench_server.sh` (local, non-k8s). +- 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. ## Requirements @@ -61,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 @@ -75,13 +77,13 @@ than one pod) -- callers never talk to these directly. See Load these in order; each covers one phase and nothing else. Services must exist **before** any Pod is deployed -- see -[services.md](references/services.md) for why: +[services.md](resources/services.md) for why: -| Reference | Covers | +| Resource | Covers | |---|---| -| [resolve-recipe.md](references/resolve-recipe.md) | Step 1: read the recipe's serve-block fields and model id; read its placement (or lack of one) from its header comment -- never ask. | -| [deploy.md](references/deploy.md) | Steps 2-4f: GPU_COUNT, confirm, PVC, ConfigMap, fsGroup, Services (step 4d -- which of the four, see [services.md](references/services.md)), pod(s) via `templates/pod.yaml`, wait for Running and for each pod's own readiness line. | -| [report-and-teardown.md](references/report-and-teardown.md) | Step 5: report the endpoint/model/nodes. Step 6: teardown (pods/Services only -- never the PVC unless asked). | +| [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). | All k8s manifests live in `templates/*.yaml`, applied via `envsubst | kubectl apply -f -`; no manifest is inlined in these docs. diff --git a/.agents/skills/deploy-afd-k8s/references/resolve-recipe.md b/.agents/skills/deploy-afd-k8s/references/resolve-recipe.md deleted file mode 100644 index 84e4ac2a9..000000000 --- a/.agents/skills/deploy-afd-k8s/references/resolve-recipe.md +++ /dev/null @@ -1,90 +0,0 @@ -# 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 -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 &`): - -- `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" -``` - -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_`. diff --git a/.agents/skills/deploy-afd-k8s/references/deploy.md b/.agents/skills/deploy-afd-k8s/resources/deploy.md similarity index 55% rename from .agents/skills/deploy-afd-k8s/references/deploy.md rename to .agents/skills/deploy-afd-k8s/resources/deploy.md index eccd5423a..0a971ec66 100644 --- a/.agents/skills/deploy-afd-k8s/references/deploy.md +++ b/.agents/skills/deploy-afd-k8s/resources/deploy.md @@ -1,31 +1,15 @@ -# 2-4. Compute GPU_COUNT, confirm, and deploy +# 2-4. Confirm and deploy ## 2. GPU_COUNT -Plain single-pod recipe (no plan from [resolve-recipe.md](resolve-recipe.md)): -union of every literal `CUDA_VISIBLE_DEVICES=` in the script. +Use the `GPU_COUNT` recorded by [resolve-recipe.md](resolve-recipe.md) -- +`a + f` per `LAYOUT` entry. Don't re-read the script here; all recipe +parsing lives in step 1: ```bash -GPU_COUNT="$(grep -oE 'CUDA_VISIBLE_DEVICES=[0-9,]+' "$RECIPE_SCRIPT_PATH" \ - | cut -d= -f2 | tr ',' '\n' | sort -u | wc -l)" -``` - -Per plan entry: count the devices in the literal `CUDA_VISIBLE_DEVICES=` inside -the recipe's `if` block that tests this pod's role (`ATTENTION` or `FFN` -- -whichever prefix `POD` starts with). Every pod sharing a role shares that -literal, so this is a plain text extraction, not a per-pod computation: - -```bash -role_block() { # $1=script $2=ATTENTION|FFN - awk -v role="$2" ' - /if.*\$POD/ && $0 ~ role { inblock=1; next } - inblock && /^fi/ { inblock=0 } - inblock { print } - ' "$1" -} -devices="$(role_block "$RECIPE_SCRIPT_PATH" "$ROLE" \ - | grep -oE 'CUDA_VISIBLE_DEVICES=[0-9,]+' | head -1 | cut -d= -f2)" -GPU_COUNT=$(( $(grep -o ',' <<<"$devices" | wc -l) + 1 )) +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 @@ -41,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= ``` @@ -94,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 - ``` @@ -130,17 +117,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 @@ -153,27 +142,51 @@ 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: +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 -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-ffn-0 2 none head "POD=FFN_0" +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" + +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' + +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 own readiness line +## 4f. Wait for Running, then for each pod's Ready condition Bounded 15-minute wait per pod for `phase=Running` (unschedulable pods never -reach `Failed`, so an unbounded loop would spin forever), then poll logs for -`=== pod READY ===` (budget 60 minutes -- model download and graph -capture dominate): +reach `Failed`, so an unbounded wait would spin forever), then a 60-minute +wait per pod for its `Ready` condition (model download and graph capture +dominate). + +`Ready` is the single readiness signal. The pod's `readinessProbe` only +turns it on after the in-pod startup script has seen every log marker, got +`/health` from the attention head, and printed `=== pod READY ===`; +the same condition is what puts the attention head behind `vllm-service`. +If the in-pod script gives up, it prints `ERROR:` with log tails and exits, +the pod goes `Failed`, and the loop below stops at once. + +Timeouts are nested on purpose: the in-pod startup deadline +(`STARTUP_TIMEOUT_S` in `templates/pod.yaml`, 55 minutes, one budget shared +by every log and `/health`) is shorter than the 60-minute wait here, so a +slow pod always fails with its own `ERROR:` reason instead of an outer +timeout. Change the two together. ```bash for pod_name in ; do @@ -184,18 +197,22 @@ done for pod_name in ; do deadline=$(( $(date +%s) + 3600 )) while true; do - logs="$(kubectl logs "pod/$pod_name" 2>/dev/null || true)" - echo "$logs" | grep -q "pod ${pod_name} READY" && break - echo "$logs" | grep -q "ERROR:" && { kubectl logs "$pod_name" --tail=200; exit 1; } - [ "$(date +%s)" -ge "$deadline" ] && { kubectl logs "$pod_name" --tail=200; exit 1; } + read -r phase ready < <(kubectl get pod "$pod_name" \ + -o jsonpath='{.status.phase} {.status.conditions[?(@.type=="Ready")].status}' || true) + [ "$ready" = "True" ] && break + if [ "$phase" = "Failed" ] || [ "$(date +%s)" -ge "$deadline" ]; then + echo "pod $pod_name not Ready (phase=${phase:-unknown})" + kubectl logs "$pod_name" --tail=200 + exit 1 + fi sleep 10 done done ``` -Only report the endpoint ready once every pod in the plan has printed its -own `READY` line -- one pod becoming ready doesn't imply the others did. +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/references/report-and-teardown.md b/.agents/skills/deploy-afd-k8s/resources/report-and-teardown.md similarity index 69% rename from .agents/skills/deploy-afd-k8s/references/report-and-teardown.md rename to .agents/skills/deploy-afd-k8s/resources/report-and-teardown.md index 70fbb7e03..7949c2fbb 100644 --- a/.agents/skills/deploy-afd-k8s/references/report-and-teardown.md +++ b/.agents/skills/deploy-afd-k8s/resources/report-and-teardown.md @@ -12,6 +12,13 @@ kubectl get pods -l app=afd-recipe -o wide Every node matters if a follow-up pod needs to attach the same `ReadWriteOnce` PVC -- it can only attach from that node. +**Ask when to delete the running pods.** End the report by asking the +caller -- the user, or the task that invoked this skill (e.g. +`run-vllm-bench-k8s`) -- when the pods should be deleted: right after its +follow-up work, at a later point it names, or never (keep them warm). The +caller owns that decision and triggers step 6 when the time comes; this +skill never tears down on its own. + # 6. Teardown Left running by design, so weights stay warm -- report what's up and how to diff --git a/.agents/skills/deploy-afd-k8s/resources/resolve-recipe.md b/.agents/skills/deploy-afd-k8s/resources/resolve-recipe.md new file mode 100644 index 000000000..7c2995ffa --- /dev/null +++ b/.agents/skills/deploy-afd-k8s/resources/resolve-recipe.md @@ -0,0 +1,74 @@ +# 1. Resolve the recipe, model, and pod resources + +Take the recipe path from the user, or ask. It must be a colocation recipe +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). + +Every such recipe states its `MODEL_PATH`, default `LAYOUT`, mode, and serve +flags, then sources the shared launcher, +`recipe/gpu/common/afd_colocation.sh`. Treat the launcher as the execution +contract: do not reconstruct or parse its `vllm serve` commands to determine +Kubernetes resources. + +## Model + +Resolve `MODEL_ID` from the caller's `MODEL_PATH` if given, else the recipe's +default. It is either a HF Hub id or the placeholder +`/path/model_weights/` used by every colocation recipe. +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. + +## Pod resource plan -- never ask + +Take `LAYOUT` from the caller's settings, else the recipe's default; never +ask. The selected `LAYOUT`, not a recipe filename or descriptive comment, +controls the plan. Split it by comma into one `AF` entry per pod. For +entry `k`, `a` is that pod's Attention GPU count and `f` its FFN GPU count: + +- **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`, passed to `templates/pod.yaml` as `TEMPLATE_GPU` + so the container requests that many `nvidia.com/gpu` devices. +- **`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` gives `head none` then `none head`. +- **Container env**: every pod gets the selected `LAYOUT`, its own + `POD=AFD_`, and the caller's other recipe overrides unchanged. Leave + `HEAD_HOST`, `ATTN_HEAD_HOST`, and `FFN_HEAD_HOST` unset; their defaults + name the Services in [services.md](services.md). With `DECODE_BENCH`, also + set `AFD_REPO=/opt/afd-plugin`. + +The launcher owns role-to-device ordering, DP/TP command construction, and +startup validation of the placement. Record the pod/resource table for +[deploy.md](deploy.md), which uses it as-is. + +## Launcher ports + +Read the client and internal Service ports from the launcher's fixed assignments: + +```bash +RECIPE_LIB_PATH=recipe/gpu/common/afd_colocation.sh +lib_port() { + grep -oE "\b$1=[0-9]+" "$RECIPE_LIB_PATH" | head -1 | cut -d= -f2 +} + +CLIENT_PORT="$(lib_port ATTN_PORT)" +AFD_CONNECTOR_PORT="$(lib_port AFD_PORT)" +ATTN_DP_RPC_PORT="$(lib_port ATTN_DP_RPC_PORT)" +FFN_DP_RPC_PORT="$(lib_port FFN_DP_RPC_PORT)" +``` + +## Other deployment values + +Record the effective `--max-num-batched-tokens` -- the caller's `VLLM_ARGS` +value if overridden, else the recipe's `SERVE_ARGS` value -- for the final +report. diff --git a/.agents/skills/deploy-afd-k8s/references/services.md b/.agents/skills/deploy-afd-k8s/resources/services.md similarity index 63% rename from .agents/skills/deploy-afd-k8s/references/services.md rename to .agents/skills/deploy-afd-k8s/resources/services.md index 3aa423a8d..94173686c 100644 --- a/.agents/skills/deploy-afd-k8s/references/services.md +++ b/.agents/skills/deploy-afd-k8s/resources/services.md @@ -24,28 +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 (`AFD_CONNECTOR_PORT` from +[resolve-recipe.md](resolve-recipe.md#launcher-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="${AFD_CONNECTOR_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). -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="${ATTN_DP_RPC_PORT}" \ + envsubst '${TEMPLATE_ROLE} ${TEMPLATE_PORT}' < templates/service-dp-coord.yaml | kubectl apply -f - +TEMPLATE_ROLE=ffn TEMPLATE_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 225524486..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" @@ -92,35 +78,33 @@ spec: esac } - echo "=== wait for every worker to report startup ===" + # One startup budget shared by every log marker and /health. Keep + # it below the 60-minute Ready wait in resources/deploy.md (step + # 4f), so a slow pod fails here with its own ERROR: reason first. + STARTUP_TIMEOUT_S=3300 + STARTUP_DEADLINE=$(( $(date +%s) + STARTUP_TIMEOUT_S )) + past_deadline() { [ "$(date +%s)" -ge "$STARTUP_DEADLINE" ]; } + + echo "=== wait for every worker to report startup (budget ${STARTUP_TIMEOUT_S}s) ===" for log in $LOGS; do marker="$(ready_marker "$log")" echo "--- waiting on $log (marker: $marker) ---" - ready=0 - for i in $(seq 1 240); do - if [ -f "/work/$log" ] && grep -q "$marker" "/work/$log"; then - echo "$log: ready" - ready=1 - break - fi + until [ -f "/work/$log" ] && grep -q "$marker" "/work/$log"; do kill -0 "$RECIPE_PID" 2>/dev/null || fail "recipe process exited before $log came up" + past_deadline && fail "timed out after ${STARTUP_TIMEOUT_S}s waiting for $log" sleep 10 done - [ "$ready" = "1" ] || fail "timed out waiting for $log" + echo "$log: ready" done if [ "$AFD_ATTN_NODE_ROLE" = "head" ]; then echo "--- waiting on attention server /health (127.0.0.1:${TEMPLATE_CLIENT_PORT}) ---" - endpoint_ready=0 - for i in $(seq 1 60); do - if curl -fsS "http://127.0.0.1:${TEMPLATE_CLIENT_PORT}/health" >/dev/null 2>&1; then - echo "attention server: ready" - endpoint_ready=1 - break - fi + until curl -fsS "http://127.0.0.1:${TEMPLATE_CLIENT_PORT}/health" >/dev/null 2>&1; do + kill -0 "$RECIPE_PID" 2>/dev/null || fail "recipe process exited before /health answered" + past_deadline && fail "attention server did not answer /health within ${STARTUP_TIMEOUT_S}s" sleep 5 done - [ "$endpoint_ready" = "1" ] || fail "attention server did not answer /health" + echo "attention server: ready" fi echo "=== pod ${TEMPLATE_POD} READY ===" @@ -150,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/deploy-dsv4-flash-async/references/afd-prefill.md b/.agents/skills/deploy-dsv4-flash-async/references/afd-prefill.md index cec0d2c01..5d6e98f29 100644 --- a/.agents/skills/deploy-dsv4-flash-async/references/afd-prefill.md +++ b/.agents/skills/deploy-dsv4-flash-async/references/afd-prefill.md @@ -68,4 +68,4 @@ bash "${DSV4_SCRIPT_DIR}/run_prefill.sh" `attn_ranks_per_dp=A TP`，跨节点 local DP 分段不重叠。 真实性能运行保持 `AFD_FORCE_BALANCED_TOPK_IDS=0`、`enable_force_eplb=false`，关闭 IO 调试日志。 验证路径、数学语义和停止边界见 [验证](validation.md) 与 -[主线 FFN runtime](../../../../docs/npu/ASYNC_FFN_RUNTIME.md)。 +[主线 FFN runtime](../../../../docs/design/npu/async_ffn_runtime.md)。 diff --git a/.agents/skills/deploy-dsv4-flash-async/references/environment.md b/.agents/skills/deploy-dsv4-flash-async/references/environment.md index 1d9dc86bf..dbea6c998 100644 --- a/.agents/skills/deploy-dsv4-flash-async/references/environment.md +++ b/.agents/skills/deploy-dsv4-flash-async/references/environment.md @@ -36,7 +36,7 @@ export PID_DIR="${LOG_DIR}/pids" | vLLM 0.30.0 | `ced6857afa0ea7b2e3f0846a62e1394e90f15607` | | vLLM-Ascend | `8d4409d6256d8a6729140ddcc0d1889e3f96cdd6` | -环境与历史实测见 [安装说明](../../../../README.md#ascend-npu-installation)。 +环境与历史实测见 [安装说明](../../../../docs/installation.md#ascend-npu-installation)。 记录实际 CANN、torch、torch-npu 和构建产物；已有 A3 记录不是任意版本组合的兼容承诺。 源码路径默认使用当前 AFD checkout 和同级 vLLM/Ascend，挂载或安装包不同则显式覆盖。 @@ -87,7 +87,7 @@ PYTHON 仅 `import afd_plugin` 或 loader 成功不证明旧安装包含这四个算子；注册检查也不代替实际通信验收。缺少算子时检查安装日志和实际导入路径，并重新安装匹配源码。P 的 Attention/FFN 节点均需匹配的算子版本。 -仓库内构建细节见 `csrc/npu/README.md`，通信算子说明见 `docs/npu/CAM_ASYNC_ROUTED_OPS.md`。 +仓库内构建细节见 `csrc/npu/README.md`，通信算子说明见 `docs/design/npu/cam_async_routed_ops.md`。 | 用途 | 变量 / 默认值 | | --- | --- | diff --git a/.agents/skills/deploy-dsv4-flash-async/references/validation.md b/.agents/skills/deploy-dsv4-flash-async/references/validation.md index 0bcdb2726..d6dff1bea 100644 --- a/.agents/skills/deploy-dsv4-flash-async/references/validation.md +++ b/.agents/skills/deploy-dsv4-flash-async/references/validation.md @@ -3,8 +3,8 @@ 此 skill 的命令和配置已迁移到固定 0.30 源码契约，硬件资格按具体运行记录认定。 AFD 原始模型、DBO 和 FFN FULL 的已有证据入口为 [PR #462](https://github.com/vllm-project/afd-plugin/pull/462)、 -[原生 DBO](../../../../docs/npu/TESTING.md#native-dbo-contracts-and-device-execution) 和 -[FFN FULL](../../../../docs/npu/ASYNC_FFN_RUNTIME.md#validation)。 +[原生 DBO](../../../../tests/npu/README.md#native-dbo-contracts-and-device-execution) 和 +[FFN FULL](../../../../docs/design/npu/async_ffn_runtime.md#validation)。 PR 中复用的旧运行、维护者确认和新运行各有范围，不能改写为本次 skill 的新硬件结果。 ## 按任务选择覆盖 diff --git a/.agents/skills/review-pr/SKILL.md b/.agents/skills/review-pr/SKILL.md index 81d94c95d..c15faf101 100644 --- a/.agents/skills/review-pr/SKILL.md +++ b/.agents/skills/review-pr/SKILL.md @@ -1,6 +1,6 @@ --- name: review-pr -description: Review pull requests and local branches for vllm-project/afd-plugin with a frozen snapshot, module-design ownership, feature overlays, patch-contract enforcement, targeted validation, and concise evidence-backed findings. Use for default, detailed, or repeat reviews; checking correctness, NPU/GPU compatibility, vLLM 0.26.0 patch discipline, connectors, distributed topology, workers and model runners, tests, benchmarks, or model-family additions; and identifying or explicitly requesting code-owner reviewers. Use run-e2e instead to execute E2E suites, and AGENTS.md for authoring rules. +description: Review pull requests and local branches for vllm-project/afd-plugin with a frozen snapshot, module-design ownership, feature overlays, patch-contract enforcement, targeted validation, and concise evidence-backed findings. Use for default, detailed, or repeat reviews; checking correctness, NPU/GPU compatibility, pinned vLLM/vLLM-Ascend patch discipline, ModelRunnerV2-only execution, connectors and routed-experts transports, distributed topology, workers and model runners, E2E gate structure, tests, benchmarks, or model-family additions; and identifying or explicitly requesting code-owner reviewers. Use run-e2e instead to execute E2E suites, and AGENTS.md plus docs/UT_MAINTENANCE.md for authoring and test rules. --- # Review AFD Plugin Pull Requests @@ -21,8 +21,8 @@ Make every finding: - **Calibrated:** match severity to user and maintainer impact. Do not report unrelated backlog, style already enforced by pre-commit (ruff, -mypy, SPDX headers, markdownlint, DCO), or a missing test that would not -protect changed behavior. +mypy, typos, SPDX headers, markdownlint, shellcheck, DCO), or a missing test +that would not protect changed behavior. ## Select the input and depth @@ -38,6 +38,25 @@ checkouts; use another skill for unrelated repositories. Default to maintainer brevity. A detailed or audit request expands coverage and lists `path:line` findings, but keeps the same confidence and severity bar. +## Pins are data, not constants + +Version pins move with every upstream upgrade. Read them from the frozen head, +never from this skill and never from memory: + +- `afd_plugin/compat/vllm.py` — `TARGET_VLLM_VERSION` (single pin shared by + both backends, matched exactly, not by range). +- `docs/design/module/index.md` — `upstream_refs` (vLLM commit and vLLM-Ascend + commit) and `last_reviewed`. +- `pyproject.toml` optional `vllm` extra, `docker/Dockerfile.ci` + (`VLLM_BASE_TAG` plus its version assertion), and `docs/installation.md`'s + runtime table. + +The values quoted in this skill are the `main` values as of its last refresh +(vLLM `0.32.0` at `5df63a72`, vLLM-Ascend `bd42eea0`); a mismatch between the +skill and the reviewed head is expected and the head wins. A PR that moves a pin +without moving the gate, the CI image tag, the docs, and the design pages +together is a finding. + ## Reference guide Load references in review order: process, one primary module contract, matching @@ -56,7 +75,7 @@ drift. | --- | --- | | [review-execution.md](references/process/review-execution.md) | Every review; freeze inputs, inspect safely, and deliver against the same snapshot. | | [general-checks.md](references/process/general-checks.md) | Every review; apply repository-wide correctness, patch-contract, and evidence rules. | -| [design-contracts.md](references/process/design-contracts.md) | Every production review; resolve branch-local module design status. | +| [design-contracts.md](references/process/design-contracts.md) | Every production review; resolve branch-local module design status and ownership. | | [review-routing.md](references/process/review-routing.md) | After the diff census; select one primary module and conditional overlays. | ### Primary module contract @@ -64,22 +83,24 @@ drift. | Reference | Read when | | --- | --- | | [plugin-config.md](references/modules/plugin-config.md) | Plugin registration, `AFDConfig` schema, parse/validation, extra-config, or env flags change. | -| [compat-patches.md](references/modules/compat-patches.md) | A monkey patch is added or changed under `afd_plugin/compat/patches/`, or the vLLM/vLLM-Ascend version gate moves. | -| [npu-compat.md](references/modules/npu-compat.md) | `afd_plugin/compat/npu/` runtime init, deferred NPU patch application, feature validation, or Ascend ops loading changes. | -| [connectors.md](references/modules/connectors.md) | An attention↔FFN connector backend, its control plane, or transfer metadata changes. | +| [compat-patches.md](references/modules/compat-patches.md) | A monkey patch is added or changed under `afd_plugin/compat/patches/`, `afd_plugin/v1/worker/`, or the vLLM/vLLM-Ascend version gate moves. | +| [npu-compat.md](references/modules/npu-compat.md) | `afd_plugin/compat/npu/` runtime init, deferred NPU patch application, feature validation, routing ids, or Ascend ops loading changes. | +| [connectors.md](references/modules/connectors.md) | An attention↔FFN connector backend, routed-experts transport, its control plane, or transfer metadata changes. | | [distributed-topology.md](references/modules/distributed-topology.md) | AFD process groups, rank mapping, or topology validation changes. | | [engine-orchestration.md](references/modules/engine-orchestration.md) | Async-DP engine behavior, cross-role request routing, DBO yield, or ubatching orchestration changes. | -| [workers-runners.md](references/modules/workers-runners.md) | Role workers, model runners (V1/V2), graph capture, or attention/FFN metadata change. | +| [workers-runners.md](references/modules/workers-runners.md) | Role workers, ModelRunnerV2 runners, routed-experts/async FFN runners, graph capture, or attention/FFN metadata change. | | [execution-platforms.md](references/modules/execution-platforms.md) | NPU/GPU selection, platform checks, Ascend op build, or `csrc/` kernels change. | | [model-integration.md](references/modules/model-integration.md) | Model registration, role-split wrappers, weight filtering, or a model family changes. | +| [e2e-harness.md](references/modules/e2e-harness.md) | `tests/e2e/**`, case IDs, gate selection, sample counts, device layout, or `run-e2e` skill behavior change. | ### Feature-design overlays | Reference | Read when | | --- | --- | -| [disaggregation-topologies.md](references/features/disaggregation-topologies.md) | 2A2F/2A1F, colocation vs disaggregation placement, or recipe launch scripts change. | -| [async-cam-ubatching.md](references/features/async-cam-ubatching.md) | The async CAM connector, two-stage MoE ubatching, or DSV3.2/DSV4 async CAM forward paths change. | +| [disaggregation-topologies.md](references/features/disaggregation-topologies.md) | 2A2F/2A1F/1A1F/4A2F/4A4F, colocation vs disaggregation, multi-pod placement, or recipe launch scripts change. | +| [async-cam-ubatching.md](references/features/async-cam-ubatching.md) | The async CAM connector, layered W4A8/W8A8 FFN execution, FFN two-stream, or DSV3.2/DSV4 async CAM forward paths change. | | [graphs-and-compilation.md](references/features/graphs-and-compilation.md) | CUDA graph, ACL graph, capture metadata, or graph/eager fallback behavior changes. | +| [dspark-draft.md](references/features/dspark-draft.md) | DSpark draft-model speculation, its speculator patch, or target/draft layered FFN wiring changes. | ### Evidence and quality checks @@ -102,30 +123,31 @@ drift. ### 1. Freeze and report the snapshot -Pin the base and head before reading source or running validation. Within 60 seconds, -report the pinned head, CI, mergeability, and preliminary findings in the host -conversation. Do not wait for CI or post this update to GitHub. +Pin the base and head before reading source or running validation. Within 60 +seconds, report the pinned head, CI, mergeability, and preliminary findings in +the host conversation. Do not wait for CI or post this update to GitHub. -If the target changes while fetching, discard the evidence and retry once. If -it changes again, report the churn and wait for a stable target. +If the target changes while fetching, discard the evidence and retry once. If it +changes again, report the churn and wait for a stable target. For a trusted PR head, materialize the pinned head in an isolated detached worktree. A worktree freezes identity but is not a security sandbox. Treat fork heads as untrusted unless the user and environment policy explicitly establish -otherwise: execute them only in a disposable, secret-free sandbox with restricted -filesystem, network, and resources; without one, use static SHA-addressed reads -and CI evidence only. For a local review, freeze the committed, index, worktree, -and NUL-safe in-scope untracked contents. Follow +otherwise: execute them only in a disposable, secret-free sandbox with +restricted filesystem, network, and resources; without one, use static +SHA-addressed reads and CI evidence only. For a local review, freeze the +committed, index, worktree, and NUL-safe in-scope untracked contents. Follow [review-execution.md](references/process/review-execution.md) for trust gates, state fingerprints, and byte-for-byte staleness checks. ### 2. Build the diff census -Group files into production code, tests, docs, configuration, build/CI, and -vendored artifacts (for example `afd_plugin/connectors/npu/bin/`). Map each -changed production file and test group to the PR goal. Compare the title/body -claims with the actual diff; use linked issues only when they define the -contract or reproduction. +Group files into production code, tests, docs, configuration, build/CI, agent +skills, and vendored artifacts (for example +`afd_plugin/connectors/npu/bin/*.run|*.whl` and the packaged +`afd_plugin/_cann_ops_custom` tree). Map each changed production file and test +group to the PR goal. Compare the title/body claims with the actual diff; use +linked issues only when they define the contract or reproduction. Mark unrelated scope and unexplained generated artifacts. Do not infer behavior from the PR description without tracing the live code. @@ -139,7 +161,8 @@ module contract, a second only for a real documented cross-boundary call path, and every matching feature and evidence overlay. Treat titles and paths as hints; live behavior and the frozen head's current design metadata are authoritative. For docs-, tests-, or CI-only changes, route to the production -contract they protect or use only the applicable evidence checks. +contract they protect or use only the applicable evidence checks; E2E harness +changes route to [e2e-harness.md](references/modules/e2e-harness.md). ### 4. Run the blocker scan @@ -156,24 +179,40 @@ vLLM engine/config ingress (--additional-config {"afd": ...}) -> sampling/output in the owning role -> shutdown/teardown ``` -Cover every applicable GPU/NPU, graph/eager, DBO on/off, 2A2F/2A1F, -runner V1/V2, colocation/disaggregation, and AFD-on/AFD-off path. Patched code -must stay transparent when AFD is inactive. Search bounded callers and sibling -implementations rather than assuming the changed hunk is the only path. +Cover every applicable GPU/NPU, graph/eager, DBO on/off, 2A2F/2A1F/1A1F, +colocation/disaggregation, and AFD-on/AFD-off path, and the per-family variants +that actually exist: + +- connector backend: CUDA P2P, DeepEP M2N (routed experts), Ascend CAMP2P, + Ascend CAM async; +- FFN execution: control-plane-driven, routed-experts, async layered GMM, + experimental FFN two-stream; +- draft model on/off (DSpark) where the suite supports it. + +There is no V1 fallback on either platform: both roles require +`VLLM_USE_V2_MODEL_RUNNER=1`, Attention runs ModelRunnerV2, and FFN stays +connector-driven. Patched code must stay transparent when AFD is inactive +(`is_afd_active` false). Search bounded callers and sibling implementations +rather than assuming the changed hunk is the only path. ### 5. Apply module, feature, and patch contracts Apply the reference set selected in step 3 and any matching repo-local skill. Read the exact module and feature pages in the frozen head, including status, -ownership, `primary_code_paths`, dependencies, and `last_reviewed`. Candidate -or draft rules are questions, not blockers, unless current code, tests, or -AGENTS.md enforce them. Inspect both sides of every config, registration, -connector, rank-mapping, and patch boundary. - -For every change under `afd_plugin/compat/`, enforce the AGENTS.md patch -contract: pinned upstream refs, full upstream copy with marked AFD differences, -identical signatures, `Patch reason`/`Patch functionality` comments, and -CPU-safe test coverage. +owners, `primary_code_paths`, `depends_on`, `validation_paths`, and +`last_reviewed`. Only `plugin_boundary.md` and `e2e_testing.md` are currently +`normative`; every other page is `draft`, so its rules are questions, not +blockers, unless current code, tests, or AGENTS.md enforce them. Inspect both +sides of every config, registration, connector, rank-mapping, and patch +boundary. + +For every patch change — under `afd_plugin/compat/` and the patched methods in +`afd_plugin/v1/worker/**` — enforce the AGENTS.md patch contract: pinned +upstream refs, full upstream copy with marked AFD differences, identical +signatures, `Patch reason:`/`Patch functionality:` comments, `# Upstream +source:` pointers, and CPU-safe test coverage. The single parametrized table in +`tests/unit/compat/patches/test_signatures.py` owns upstream signature +alignment; a new patched function without a table row is a finding. When a diff adds or expands a helper, class, fallback, compatibility branch, or public behavior, run a subtraction pass: remove out-of-scope behavior and check @@ -190,9 +229,20 @@ imports, tests, builds, hooks, or repo-configurable tooling from an untrusted head on the reviewer host. - Treat CI as status evidence; inspect only the first overlapping failure. - Buildkite `test-ready` runs unit tests plus the DeepSeek-V2-Lite E2E gate; - NPU has no in-repo CI, so NPU evidence is author-provided or manually run. -- For docs-only changes, use diff hygiene, links checks, and bounded live +- Buildkite `test-ready` (L2) and `test-merge` (L3) run the CPU unit selection + — `tests/unit` with `-m "not gpu and not vllm_runtime"` plus the four + CPU-only `tests/npu` files the pipeline names — then + `tests/gpu/test_attention_routing.py`, on `l4_1`; the DeepSeek-V2-Lite GPU + seven-case E2E gate runs on `l4_4`. `test-weekly` adds Qwen3 MoE, Qwen3.6 + MoE, and the DeepSeek-V2-Lite 2A1F Graph + DBO case on `h100_4`. +- CI mirrors `docker/Dockerfile.ci`, which builds on + `vllm/vllm-openai:` and asserts the installed vLLM pin. +- `skip_ci.py` suppresses all default CI when the diff is docs-only + (`docs/**`, any `*.md`, `mkdocs.yml`) or changes only pytest skip marks. For + such a PR, a green pipeline is not evidence about the change. +- NPU has no in-repo CI (`.buildkite/npu/` does not exist yet); NPU evidence is + author-provided or manually run. +- For docs-only changes, use diff hygiene, link checks, and bounded live contract verification instead of dependency setup or pytest. - For hardware-dependent paths, run available static/CPU checks (`pytest -m "not gpu and not vllm_runtime"`) and name the exact GPU/NPU gap. Never @@ -238,9 +288,16 @@ permission from a request to review the code. ## Related skills -- `run-e2e` — execute the DeepSeek-V2-Lite / Qwen3 MoE / Qwen3.6 E2E gates; use - when the review needs active GPU or NPU evidence. +- `run-e2e` — execute the DeepSeek-V2-Lite / Qwen3 MoE / Qwen3.6 MoE and + DeepSeek-V4 Flash E2E gates; use when the review needs active GPU or NPU + evidence. - `adapt-model` — author-side workflow for adapting a new model family; suggest to contributors alongside the model-addition checklist. - `upgrade-npu` / `upgrade-gpu-version` — the pinned-ref upgrade workflow; a PR that moves vLLM or vLLM-Ascend pins should follow them. +- `upgrade-layered-gmm-ops` — replaying the A3 layered GMM selection onto a new + ops-transformer revision. +- `deploy-afd-k8s` / `run-vllm-bench-k8s` / `deploy-dsv4-flash-async` — + deployment and benchmark workflows; a PR that changes their inputs, recipes, + or skill scripts needs the matching owner. +- `maintain-vllm-main` — incremental vLLM-main tracking between pin upgrades. diff --git a/.agents/skills/review-pr/references/checks/model-addition-checklist.md b/.agents/skills/review-pr/references/checks/model-addition-checklist.md index 55affad46..4d595a839 100644 --- a/.agents/skills/review-pr/references/checks/model-addition-checklist.md +++ b/.agents/skills/review-pr/references/checks/model-addition-checklist.md @@ -12,7 +12,10 @@ skill; this checklist is the reviewer's evidence bar. only). Each role constructs only its role-required components. - **Registration.** Entries in the `register_afd()` registration maps with the correct upstream architecture names; NPU/GPU variant selection - (`torch_npu` branch) resolved at registration, not in forward code. + (`torch_npu` branch) resolved at registration, not in forward code. The maps + are frozen (`MappingProxyType`) and the AFD alias is only installed when + `is_afd_active` is true — non-AFD workers must keep the checkpoint + architecture and resolve to the native vLLM class. - **Weight-role filtering.** `_checkpoint_weight_roles` coverage so no role loads weights it never reads and none drops weights it later needs; verify against the checkpoint's full weight list, including gate/expert and @@ -20,20 +23,29 @@ skill; this checklist is the reviewer's evidence bar. - **Upstream fidelity.** Copied model code follows the patch contract: marked AFD differences, `# Upstream source:` pointers, identical signatures. - **Config compatibility.** Direct attribute access to upstream config; - documented assumptions about quantization, MLA/attention backend, and - MoE settings the family requires. + documented assumptions about quantization, MLA/attention backend, and MoE + settings the family requires. +- **Connector/execution fit.** State which connectors and topologies the family + actually supports, and make the unsupported combinations fail at startup + rather than silently degrade. Family restrictions are real: Qwen3.5/3.6 GDN + DBO is unsupported; DeepEP M2N requires an all-MoE family with no shared + expert; Async CAM is DeepSeek-V4 W4A8/W8A8 only. ## Evidence required in the PR - A unit suite under `tests/unit/model_executor/` (CPU-safe, checkpoint- or - fixture-driven) covering role construction and weight filtering. -- An E2E suite or scenario under `tests/e2e/models/` with the four gate - scenarios (or the suite's stated gate set) and GSM8K-7 accuracy evidence. + fixture-driven) covering role construction, proxy fidelity, and weight + filtering — the existing naming is `test__construction.py`, + `test__proxy.py`, `test__weight_policy.py`. +- An E2E suite or scenario under `tests/e2e/models//` with the suite's + gate scenarios and GSM8K accuracy evidence, respecting the four-device bound + (`E2E-INV-002`) and the no-skip rule (`E2E-INV-004`). - A recipe under `recipe/` for at least one working topology, with launch scripts consistent with the user guides. -- Docs: the relevant `docs/design/module/` page mention (model integration) - and, for user-facing support, a user guide or guide section. -- For NPU support: manual validation evidence per `docs/npu/TESTING.md`; +- Docs: the relevant `docs/design/module/` page (model integration) and, for + user-facing support, a user guide or guide section; the design-page + registration table must gain the new row. +- For NPU support: manual validation evidence per `tests/npu/README.md`; there is no NPU CI. A family missing any required piece is a P1 finding; state which piece and diff --git a/.agents/skills/review-pr/references/checks/perf-verification.md b/.agents/skills/review-pr/references/checks/perf-verification.md index afa4159a6..3537351ef 100644 --- a/.agents/skills/review-pr/references/checks/perf-verification.md +++ b/.agents/skills/review-pr/references/checks/perf-verification.md @@ -1,29 +1,32 @@ # Performance and Accuracy Verification Use when a PR claims latency, throughput, memory, or accuracy improvement — -including "[PERF]" titles, ubatching/CAM control-path optimizations, graph -changes, and routing changes. +including `[PERF]` titles, layered GMM / CAM control-path optimizations, graph +changes, TP/SP scheduling changes, and routing changes. ## Comparable A/B contract Accept a claim only with base and head runs that hold constant: - Hardware pool and placement (CI presets: `l4_1`–`l4_4`, `h100_4`; NPU: - documented manual environment per `docs/npu/TESTING.md`). -- Container image and dependency versions (CI image pins - `VLLM_BASE_TAG=v0.26.0`). -- Workload and datasets (`tools/benchmarks/decode_bench.py` or the named E2E - scenario; GSM8K-7 for accuracy gates). -- Graph/eager mode, DBO on/off, ubatching mode, and topology (2A2F/2A1F, - colocation/disaggregation). + documented manual environment per `tests/npu/README.md`). +- Container image and dependency versions (`docker/Dockerfile.ci` pins + `VLLM_BASE_TAG` and asserts the installed vLLM version). +- Workload and datasets (`tools/benchmarks/decode_bench.py`, + `tools/benchmarks/request_generator.sh`, or the named E2E scenario; GSM8K + for accuracy gates). +- Graph/eager mode, DBO on/off, ubatching mode, FFN execution path, and + topology (1A1F/2A1F/2A2F/4A2F/4A4F, colocation/disaggregation, multi-pod). - Warmup, repetitions, and the metric definition (mean vs median, percentiles, - tokenization boundaries). + tokenization boundaries). The NPU benchmark harnesses deliberately do not + pool or interpolate samples — preserve that when comparing. ## Isolate per claim One claim, one comparison. A PR claiming several improvements (transfer -overlap, MoE split, control-path overhead) needs evidence per mechanism or a -staged ablation; a single end-to-end delta cannot attribute multiple causes. +overlap, MoE split, layered scheduling, control-path overhead) needs evidence +per mechanism or a staged ablation; a single end-to-end delta cannot attribute +multiple causes. ## Evidence table @@ -32,18 +35,21 @@ staged ablation; a single end-to-end delta cannot attribute multiple causes. | Latency | Decode bench per-step or E2E per-request timing, same batch mix | | Throughput | Tokens/s at stated concurrency, steady-state window | | Memory | Peak allocator/device memory, same graph/capture settings | -| Accuracy | GSM8K-7 gate plus the model suite's accuracy scenario; state pass criteria | +| Accuracy | GSM8K gate plus the model suite's accuracy scenario; state pass criteria and sample count | ## Reviewer verification ladder 1. Check the A/B contract above against the PR's Test Result section; missing constants make the claim unproven, not merely weak. 2. When a runnable GPU path exists, offer to rerun the named benchmark or E2E - scenario via the `run-e2e` skill; otherwise mark the claim unverified and - name the exact missing run. + scenario via the `run-e2e` skill (or `run-vllm-bench-k8s` for a deployed + recipe); otherwise mark the claim unverified and name the exact missing run. 3. For NPU claims, require the manual-run evidence (environment, command, numbers) in the PR; there is no NPU CI, and unit tests never substitute for device measurements. +4. Watch for numbers carried over from a previous pin: evidence measured on an + older vLLM/vLLM-Ascend pair does not support a claim about the current pair. + The design pages mark such rows historical for the same reason. An unproven perf claim on a hot path is a P1 finding (missing blocking evidence), stated as the specific missing comparison — not as distrust. diff --git a/.agents/skills/review-pr/references/checks/test-quality-evaluation.md b/.agents/skills/review-pr/references/checks/test-quality-evaluation.md index 9ac0a7fb7..a10b34664 100644 --- a/.agents/skills/review-pr/references/checks/test-quality-evaluation.md +++ b/.agents/skills/review-pr/references/checks/test-quality-evaluation.md @@ -1,13 +1,14 @@ # Test Quality Evaluation Use when tests change in the diff, are absent for risky changed code, or may -pass without exercising production behavior. +pass without exercising production behavior. Standards live in AGENTS.md +("Test Conventions") and `docs/UT_MAINTENANCE.md`; this file is the reviewer's +application of them. ## CPU-safe imports first -The unit suite must run on a CPU-only host without torch. The recurring -regression class (fixed in `f5fa9fc`) is a test that imports a worker or -connector module before guarding its runtime dependencies. Enforce: +The unit suite must run on a CPU-only host without torch, torch_npu, or vLLM. +Enforce: - Runtime dependencies (`vllm`, `vllm_ascend`, `torch_npu`, `torch`) are gated with `pytest.importorskip` (or the `_require_npu_runtime()` helper in @@ -16,10 +17,13 @@ connector module before guarding its runtime dependencies. Enforce: import, never after. - Module-level imports in test files stay CPU-safe; a top-level `from afd_plugin.v1.worker... import ...` in a test file breaks collection - on CPU hosts even if every test body is guarded. + on CPU hosts even if every test body is guarded. This is the recurring + regression class; check it before anything else. - New tests pass in the CI selection: `pytest tests/unit -m "not gpu and not - vllm_runtime"`. Run (or reason through) the selection for the changed - directory. + vllm_runtime"`, plus the four CPU-only `tests/npu` modules CI names + (`test_async_cam_harness.py`, `test_async_cam_performance.py`, + `test_async_cam_golden.py`, `test_gmm_layered_performance.py`). Run (or + reason through) the selection for the changed directory. ## Markers and selection @@ -31,16 +35,62 @@ connector module before guarding its runtime dependencies. Enforce: | `eval` | Datasets/accuracy | Missing dataset guard | | `slow` | > 120 s | Unmarked long tests that inflate CI time | +Note that `tests/npu` is not uniformly device-gated: several modules there are +CPU host-harness checks that CI runs directly. Do not demand an `npu` marker on +those, and do not assume a `tests/npu` path implies hardware. + ## Quality bar +- Assert the seam, not the plumbing: return values, raised errors and their + messages, protocol messages sent to a peer, final file or environment state, + config-validation outcomes. Private attributes, internal-collaborator + arguments, and internal helper call counts are forbidden. Call order is a + valid assertion only when the order is itself the contract (connector + handshakes, service announce, DBO stage order). +- The refactor test: if the implementation is restructured without changing a + protocol, the test must stay green. A test that cannot fail when behavior + regresses is not coverage — flag it and ask for deletion rather than + retention. - Tests exercise production behavior, not a mock echo: patched functions get AFD-on and AFD-off assertions; state machines get abort/timeout paths; fail-fast contracts assert the error message, not just `pytest.raises`. -- Device-independent logic (planners, rank mapping, config validation) is - tested CPU-side with `parametrize`; device-only behavior is honestly - delegated to E2E scenarios rather than fake-simulated in unit tests. -- A risky change without a test that would protect it is a P2 finding (P1 if - the untested path already regressed before); name the specific behavior to - cover and the cheapest harness for it. -- Do not demand coverage the repo's structure already provides elsewhere - (E2E gates, package tests, upstream CI). +- Shared fakes live in `tests/unit/fakes/`, one module per upstream concept; a + builder defined in two test modules is a layout bug. A test patching more + than ~4 attributes of the module under test is testing its own mocks. +- Setup over ~15 lines belongs in a fixture or factory; a test body over ~50 + lines needs `# long-test: `; the third test with the same prefix and + setup becomes a `parametrize` table. +- Deterministic: no sleeps, no wall-clock assertions, no reliance on set + iteration order. Concurrency is coordinated with events or bounded + conditions. Numerical tolerances are justified. +- Text artifacts (build manifests, patch markers, docs links) may be asserted + as text. Production Python modules may not: no test reads + `afd_plugin/**/*.py` and asserts on a substring. +- File placement mirrors the source tree under `tests/unit//`, + and a new test file must fall under an existing `validation_paths` entry in + `docs/design/module/`. A test file owned by no module is a finding — or, if + it is a full-stack check, it belongs in `tests/e2e`. +- Upstream signature alignment stays in the single table in + `tests/unit/compat/patches/test_signatures.py`; inheritance alignment is + mypy's job, not pytest's. +- Device-independent logic (planners, rank mapping, config validation, layered + layout math) is tested CPU-side with `parametrize`; device-only behavior is + honestly delegated to E2E or `tests/gpu`/`tests/npu` rather than + fake-simulated in unit tests. + +## Lifecycle + +- A PR that removes behavior deletes its tests and now-unused helpers in the + same PR. Skipping is not retirement; a skip needs an issue link. +- Merging duplicate tests is acceptable only when the retained test still + detects the relevant failures; similar bodies or equal expected values are + not sufficient justification. +- Never make CI green with retries, weaker checks, relaxed tolerances, or + deleted necessary coverage. A gate case turned into a skip is a P1. + +## Finding severity + +A risky change without a test that would protect it is a P2 finding (P1 if the +untested path already regressed before); name the specific behavior to cover +and the cheapest harness for it. Do not demand coverage the repo's structure +already provides elsewhere (E2E gates, package tests, upstream CI). diff --git a/.agents/skills/review-pr/references/checks/tests-docs-checklist.md b/.agents/skills/review-pr/references/checks/tests-docs-checklist.md index 969acd35d..f1e83832b 100644 --- a/.agents/skills/review-pr/references/checks/tests-docs-checklist.md +++ b/.agents/skills/review-pr/references/checks/tests-docs-checklist.md @@ -5,15 +5,33 @@ technical review, before delivery. ## CI coverage -- GitHub Actions runs pre-commit only (ruff, format, mypy/ty, typos, - markdownlint, actionlint, shellcheck, SPDX headers, Buildkite YAML checks, - DCO). Treat lint findings as gate status, not review findings. -- Buildkite: `test-ready` (L2 premerge) runs unit tests (`l4_1`) plus the - DeepSeek-V2-Lite E2E gate (`l4_4`); `test-merge` (L3) repeats post-merge; - `test-weekly` (main or `weekly-test` label) adds Qwen3 MoE, Qwen3.6, and - DSV2-Lite 2A1F DBO on `h100_4`. -- A change touching an area no pipeline covers (NPU paths, non-gate model - suites, `csrc/` kernels) needs author-provided evidence; name the gap. +- GitHub Actions runs pre-commit only (check-yaml, merge-conflict, end-of-file, + mixed-line-ending, trailing-whitespace, ruff check/format, typos, + actionlint, markdownlint, mypy on changed files for 3.11 with 3.12/3.13 + manual-only, shellcheck, SPDX headers, Buildkite YAML validation, DCO). Treat + lint findings as gate status, not review findings. +- Pre-commit only inspects files changed in the push/PR (`--from-ref/--to-ref`); + a violation in an untouched file is not this PR's finding. +- Buildkite image: `docker/Dockerfile.ci` builds on + `vllm/vllm-openai:${VLLM_BASE_TAG}` (currently `v0.32.0`) and asserts the + installed vLLM version equals the pin. A PR that changes the pin without + changing that assertion leaves CI unable to catch drift. +- `test-ready` (L2) and `test-merge` (L3) run on `l4_1`: the CPU unit + selection (`pytest -m "not gpu and not vllm_runtime" tests/unit` plus + `tests/npu/test_async_cam_harness.py`, `test_async_cam_performance.py`, + `test_async_cam_golden.py`, `test_gmm_layered_performance.py`) and + `pytest tests/gpu/test_attention_routing.py`. The DeepSeek-V2-Lite + seven-case GPU E2E gate runs on `l4_4`. +- `test-weekly` (main + `WEEKLY=1`, or PR label `weekly-test`) adds Qwen3 MoE + (four cases), Qwen3.6 MoE (three cases), and DeepSeek-V2-Lite + `afd-graph-dbo-2a1f` on `h100_4`. +- `skip_ci.py` suppresses the whole default pipeline for docs-only + (`docs/**`, any `*.md`, `mkdocs.yml`) or skip-mark-only diffs; label-only CI + YAML edits gate per platform and level. Say what a green run does and does + not cover. +- A change touching an area no pipeline covers (NPU paths, `csrc/` kernels, + non-gate model suites, `tests/gpu/` device tests) needs author-provided + evidence; name the gap. ## PR template and checklist @@ -21,12 +39,14 @@ The template requires Purpose / Issue / Scope / Implementation Notes / Test Plan / Test Result / Docs Impact, plus the Essential PR Checklist. Verify the high-signal items rather than restating them: -- vLLM 0.26.0 compatibility considered; no vLLM source-checkout changes. -- Plugin-owned classes or dotted paths preferred over monkey patches; any - shim isolated, idempotent, version-guarded, documented, tested. -- Imports remain CPU-safe; CUDA-heavy work delayed or GPU-gated. +- vLLM and, for NPU changes, the pinned vLLM-Ascend revision compatibility + considered; no vLLM source-checkout changes. +- Plugin-owned classes or dotted paths preferred over monkey patches; any shim + isolated, idempotent, version-guarded, documented, and tested. +- Imports remain CPU-safe; CUDA/NPU-heavy work delayed or device-gated. - Validation evidence included, including skipped GPU tests and NPU manual - evidence where applicable. + evidence, with exact upstream revisions and backend-scoped results where + native builds or ABIs matter. - DCO `Signed-off-by` present (the pre-commit hook auto-adds it; a missing sign-off blocks the DCO gate). @@ -37,8 +57,13 @@ high-signal items rather than restating them: [design-contracts.md](../process/design-contracts.md)). - User-facing changes (config fields, connector usage, launch flows) update the relevant user guide (`docs/gpu/`, `docs/npu/`) or state why not. -- Topology/launch changes update `recipe/` scripts consistently; E2E gate - structure changes update `docs/design/module/e2e_testing.md`. +- Topology/launch changes update `recipe/` scripts and their READMEs + consistently; E2E gate structure changes update + `docs/design/module/e2e_testing.md`. +- `.agents/skills/**` are executable instructions, not prose: a change to + commands, case lists, pins, or workflows that leaves a SKILL.md stating the + old behavior is a P2 finding. Markdownlint excludes `.agents/`, so nothing + else catches it. - New source files carry SPDX headers; user guides keep the documented env vars and defaults in sync with `AFDConfig`. diff --git a/.agents/skills/review-pr/references/checks/verification.md b/.agents/skills/review-pr/references/checks/verification.md index e2708bfd0..882b73892 100644 --- a/.agents/skills/review-pr/references/checks/verification.md +++ b/.agents/skills/review-pr/references/checks/verification.md @@ -8,7 +8,8 @@ and the review needs evidence beyond static reading and CI. | Available environment | Verification | | --- | --- | | GPU host with devices free | Narrowest changed unit tests, then the single most relevant E2E scenario via `run-e2e` | -| NPU host | Manual per `docs/npu/TESTING.md`; only what the user explicitly asked to run | +| GPU host for connector device tests | `tests/gpu/test_deepep_m2n_gpu.py` on SM90+ with afdEP-comm and DeepGEMM; most cases need 4 GPUs, the Attention-TP cases need 6 | +| NPU host | Manual per `tests/npu/README.md`; only what the user explicitly asked to run | | CPU-only trusted checkout | Import preflight + `pytest tests/unit/ -m "not gpu and not vllm_runtime"` | | Static only (untrusted head, no sandbox) | SHA-addressed reads + existing CI evidence; report execution as a gap | @@ -19,13 +20,16 @@ and the review needs evidence beyond static reading and CI. 2. Start from the narrowest command that can change a conclusion; a full E2E suite is never the first run. 3. E2E runs go through the `run-e2e` skill with its prerequisite checks - (devices, model path, `HF_HOME`, mirror reachability); a missing - prerequisite fails the plan — it never becomes a silent skip. + (devices, model path, `HF_HOME`, mirror reachability); a missing prerequisite + fails the plan — it never becomes a silent skip. 4. Record every run: head SHA, snapshot fingerprint, command, result, Python - and platform, dependency fingerprint. + and platform, dependency fingerprint, and the vLLM/vLLM-Ascend revisions. 5. Classify failures before drawing conclusions: environment/setup failures are not product evidence; a product failure on the changed path is a finding; an unrelated pre-existing failure is context, cited not fixed. + Known open failures in the current tree include DeepSeek-V2-Lite DBO + shutdown and NPU graceful exit — do not present them as this PR's regression + without a base-commit comparison. Never simulate device evidence: no fabricated NPU numbers, no "should pass" substitutions for hardware gates. If the gap cannot be closed here, the diff --git a/.agents/skills/review-pr/references/delivery/delivery-style.md b/.agents/skills/review-pr/references/delivery/delivery-style.md index 2715a5af8..fdf4a7afc 100644 --- a/.agents/skills/review-pr/references/delivery/delivery-style.md +++ b/.agents/skills/review-pr/references/delivery/delivery-style.md @@ -7,24 +7,31 @@ direct engineering tone of its merged history — into comment craft. ## Recurring review behavior - Start from the changed contract, not style: pre-commit already owns ruff, - format, SPDX, markdownlint, and DCO; never restate them. + format, typos, SPDX, markdownlint, shellcheck, and DCO; never restate them. - Protect the boundaries that make the plugin maintainable: patch fidelity - against pinned refs, CPU-safe imports, role split, platform twins, AFD-off - transparency. These are the findings this repo's maintainers need most. + against pinned refs, CPU-safe imports, role split, platform twins, the + no-V1/ModelRunnerV2-only contract, AFD-off transparency, and E2E gate + integrity. These are the findings this repo's maintainers need most. - Challenge scope when a PR smuggles in an unexplained abstraction, a new - global, or a second mechanism for an existing one; ask for the smallest - version. -- Demand evidence exactly where it is load-bearing: perf claims, NPU - behavior, E2E gate impact — and nowhere else. -- Address named owners (`@hsliuustc0106`, `@jiangkuaixue123`) for - architectural questions, not a generic "PTAL". + global, a second mechanism for an existing one, or a permissive fallback + where the design says "fail at startup"; ask for the smallest version. +- Demand evidence exactly where it is load-bearing: perf claims, NPU behavior, + E2E gate impact — and nowhere else. +- Address the owner who actually holds the contract. CODEOWNERS is no longer a + single global rule: GPU paths add `@ronenkat`, `/csrc/npu/` adds + `@francisleely`, E2E paths add `@ronenkat`, and `e2e_testing.md` is owned by + `@yujuancao07`. Use the design page's `owners` plus the path-matched + CODEOWNERS entry, not a generic "PTAL". - Defer cleanup to follow-ups; do not block on debt the PR did not create. +- Known open failures are context, not ammunition: cite DeepSeek-V2-Lite DBO + shutdown or NPU graceful-exit limitations only to explain why the PR's + evidence differs from a clean pass. ## Write the finding Lead with the issue; no preamble, no praise sandwich. One sentence when the fix is obvious. Prefer root-cause comments ("this breaks patch comparability -at the next upgrade") over restating the diff. Hedge only when evidence is +at the next pin move") over restating the diff. Hedge only when evidence is genuinely incomplete, and say what would complete it. Format (severity tiers from @@ -35,19 +42,22 @@ Format (severity tiers from . . . ``` -Patch-contract findings name the exact missing element (marker, signature -parity, upstream source pointer, pinned ref) because that is what the next -upgrade will trip over. +Patch-contract findings name the exact missing element (marker, `Patch reason:` +/ `Patch functionality:` comment, signature parity, upstream source pointer, +pinned ref) because that is what the next upgrade will trip over. Gate findings +name the invariant (`E2E-INV-004`) and the case ID. Calibrate to roughly 1–5 short comments; treat that as calibration, not a -quota. Zero comments is a valid outcome — say so and name material -validation gaps. +quota. Zero comments is a valid outcome — say so and name material validation +gaps. ## Anti-patterns - Generic praise, "Nit:" prefixes, decorative severity labels. -- Audit-template artifacts: rule IDs, checklist restatements, comment-count - narration. -- Style findings pre-commit owns, or AGENTS.md wording quotes with no - concrete failure. +- Audit-template artifacts: rule IDs as decoration, checklist restatements, + comment-count narration. +- Style findings pre-commit owns, or AGENTS.md wording quotes with no concrete + failure. - Batching many small findings that share one root cause; merge them. +- Reporting a stale pin as a finding when the head deliberately moved it in an + upgrade PR that also moved the gate, image, docs, and design pages. diff --git a/.agents/skills/review-pr/references/delivery/review-requests.md b/.agents/skills/review-pr/references/delivery/review-requests.md index e6ee29b05..30203bf9b 100644 --- a/.agents/skills/review-pr/references/delivery/review-requests.md +++ b/.agents/skills/review-pr/references/delivery/review-requests.md @@ -7,24 +7,28 @@ authorization for this PR or session. ## Owners -`CODEOWNERS` is a single global rule: `* @hsliuustc0106 @jiangkuaixue123`. -The design pages under `docs/design/module/` declare the same owners. There -are no path-specific owners, so route by contract expertise instead: +`.github/CODEOWNERS` has a global catch-all +(`* @hsliuustc0106 @jiangkuaixue123`) plus path-specific entries: -| Changed contract | Ask for | +| Path | Adds | | --- | --- | -| Patches, version pins, platform selection | Owner with upstream vLLM/vLLM-Ascend upgrade context | -| Connectors, topology, async CAM | Owner with the connector/NPU runtime context | -| Model families, E2E gates | Owner with the model-suite context | -| CI / Buildkite pipelines | Owner with CI pool context | +| `/afd_plugin/connectors/gpu/`, `/csrc/gpu/`, `/recipe/gpu/`, `/tests/gpu/`, `/docs/gpu/`, `/.buildkite/cuda/` | `@ronenkat` | +| `/csrc/npu/` | `@francisleely` | +| `/tests/e2e/`, `/.agents/skills/run-e2e/`, `/.agents/skills/deploy-afd-k8s/` | `@ronenkat` | + +Design pages declare their own `owners` under `docs/design/module/`; most list +`@hsliuustc0106` and `@jiangkuaixue123`, while `e2e_testing.md` is owned by +`@yujuancao07` alone. Prefer the design page's owners for the *contract*, and +CODEOWNERS for the *paths*; when they disagree, say which one you are +following. ## Procedure 1. From the frozen head, map the diff census to the primary module contract and feature overlays. 2. Propose one to three focused reviewers with an explicit rationale naming - the contract ("patch contract + async CAM ubatching"), never a bare - "PTAL". + the contract ("patch contract + async CAM layered FFN", "E2E gate + selection"), never a bare "PTAL". 3. Check existing reviews and requests first; do not duplicate an owner already engaged on the PR threads. 4. When authorized to post, recheck the head SHA, then post at most one diff --git a/.agents/skills/review-pr/references/features/async-cam-ubatching.md b/.agents/skills/review-pr/references/features/async-cam-ubatching.md index 341d7e57c..a8f861270 100644 --- a/.agents/skills/review-pr/references/features/async-cam-ubatching.md +++ b/.agents/skills/review-pr/references/features/async-cam-ubatching.md @@ -1,35 +1,55 @@ -# Async CAM and Ubatching +# Async CAM, Layered FFN, and Ubatching -Load with the primary module contract for changes to the experimental async -CAM data path: `CAMAsyncAFDConnector`, AFD-managed two-stage MoE ubatching, -the NPU ubatch wrappers/metadata, DBO yield interplay, and the DSV3.2/DSV4 -async CAM forward modules under `afd_plugin/model_executor/models/npu/`. +Load with the primary module contract for changes to the async CAM data path: +`CAMAsyncAFDConnector` and `AFDAsyncTransferState`, AFD-managed native +ubatching, the NPU async FFN runner and its two-stream mode, DBO yield +interplay, the layered GMM executor, and the DSV3.2/DSV4 async CAM forward +modules under `afd_plugin/model_executor/models/npu/`. Designs: `docs/npu/CAM_ASYNC_CONNECTOR_USER_GUIDE.md`, -`docs/design/module/connector_contracts.md`; recent work concentrates here -(DSV4 CAM FFN control path, `v0.26.0_camasync_dsv4`). +`docs/design/npu/async_ffn_runtime.md`, `docs/design/npu/cam_async_routed_ops.md`, +`docs/design/module/connector_contracts.md`, +`docs/design/module/ffn_runtime.md`. ## Feature checks -- Respect experimental status: async CAM must not become a default or silently - required path; the opt-in remains explicit in config/recipes, and the sync - CAM P2P path keeps working. -- Preserve MoE semantics across the two-stage split: gate computation - placement, shared-expert handling, and topk routing inputs must match the - single-stage semantics the model defines; the execution planner - (`model_executor/npu/async_cam_ubatching.py`) is pure planning — no device - side effects. +- Respect the layer's status. `CAMAsyncAFDConnector` itself is a supported + product path for DeepSeek-V4 W4A8/W8A8; the *experimental* parts are FFN + two-stream (`ffn_two_stream`, default off) and any DP layout beyond the + documented ones. Experimental flags must not become defaults or silently + required paths; the synchronous `CAMP2pAFDConnector` path keeps working, and + DSpark must not acquire DBO or two-stream. +- Layered GMM is the only Async CAM FFN compute path. Incompatible weights + fail startup rather than selecting another FFN path — a PR that adds a + fallback contradicts `ffn_runtime.md`. `dynamicQuant=1` is required; + `w13_scale_bias`/`w2_scale_bias` are required for W4A8; W8A8 uses INT8 NZ + tensors with a single shared zero bias; `swiglu_limit` is applied in both + eager and FULL while W8A8 keeps its unclipped computation. +- Keep `AsyncCAMLayeredExecutor` and `LayeredFFNWeights` pure about device + side effects where the design says so: planning (`async_cam_layout.py`) is + host-side and testable on CPU, device execution belongs to the executor. - Prove state-machine safety: `AFDAsyncTransferState`, async FFN work items, - and ubatch contexts handle abort, timeout, and shutdown at every pending - state; a dropped work item or a lost completion signal hangs a rank. + ubatch contexts, and two-stream slots handle abort, timeout, and shutdown at + every pending state; a dropped work item or a lost completion signal hangs a + rank. FULL retires a slot only after its send is observed complete; eager + slot reuse requires the same guarantee. - Keep hidden-state handoff correct at ubatch boundaries: residual streams, - layernorm state, and metadata keys (`AscendUbatchMetadata`, graph metadata) - must be consistent between attention and FFN sides and across graph capture. + layernorm state, and graph/ubatch metadata keys must be consistent between + attention and FFN sides and across capture. +- Both roles select `VLLM_USE_V2_MODEL_RUNNER=1`; Attention stays eager under + Async CAM (`--enforce-eager`) and may enable native DBO with exactly two + microbatches, while FFN never joins Attention's DBO thread schedule. Native + DBO on FFN is rejected. - Require evidence proportional to risk: control-path or scheduling changes - here carry perf implications (the DSV4 CAM FFN work was a `[PERF][NPU]` - change) — demand comparable base/head NPU measurements or an explicit + here carry perf implications and have historically shipped as `[PERF][NPU]` + changes — demand comparable base/head NPU measurements or an explicit no-impact rationale; unit tests alone do not prove device behavior. +- Cleanup is part of the contract: a pending CAM receive cannot be interrupted + cleanly, so the E2E harness uses a marker-based SIGKILL sweep with a 120 s + survivor window. Changes that touch receiver teardown must keep that path + working; a "fix" that makes teardown look clean by suppressing survivors is + a regression. Name the NPU evidence (who ran what, on which environment per -`docs/npu/TESTING.md`) for any behavioral claim; CPU tests cover planning and +`tests/npu/README.md`) for any behavioral claim; CPU tests cover planning and state transitions only. diff --git a/.agents/skills/review-pr/references/features/disaggregation-topologies.md b/.agents/skills/review-pr/references/features/disaggregation-topologies.md index 210985c86..af959a048 100644 --- a/.agents/skills/review-pr/references/features/disaggregation-topologies.md +++ b/.agents/skills/review-pr/references/features/disaggregation-topologies.md @@ -2,38 +2,52 @@ Load with the primary module contract whenever a review touches role placement, rank counts, or launch scripts: attention/FFN colocation vs -prefill-decode disaggregation, 2A2F vs 2A1F rank counts, and the recipe -scripts that encode them. +prefill-decode disaggregation, the A/F rank ratios, and the recipe scripts that +encode them. Designs: `docs/design/module/attention_runtime.md`, -`docs/design/module/ffn_runtime.md`; recipes under `recipe/gpu/` and -`recipe/npu/`. +`docs/design/module/ffn_runtime.md`, `docs/design/module/e2e_testing.md`; +recipes under `recipe/gpu/` and `recipe/npu/`. ## Topology map | Topology | Where it lives | | --- | --- | | Baseline (native DP/TP, no AFD) | `baseline-graph` E2E scenarios; native DP4/TP1/EP4 | -| GPU NCCL P2P, colocation + disaggregation | `recipe/gpu/P2pNcclAFDConnector/deepseek_v2_lite/` | -| NPU CAM P2P disaggregation | `recipe/npu/CAMP2pAFDConnector/deepseek_v3_2/` | -| NPU async CAM | `recipe/npu/CAMAsyncAFDConnector/deepseek_v3_2/` | -| 2A2F / 2A1F gates | `tests/e2e/models/*/` scenario matrix | +| CUDA NCCL P2P, colocation | `recipe/gpu/P2pNcclAFDConnector/deepseek_v2_lite/prefill_decode_colocation/` (explicit 2A2F, 4A4F and two-pod graph/eager DBO recipes) | +| 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` (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` | ## Feature checks - Validate `num_attention_ranks` / `num_ffn_ranks` against the actual world - size on every rank before collectives; a 2A1F misconfiguration must fail at + size on every rank before collectives; a misconfiguration must fail at startup with a named mismatch, not deadlock mid-run. +- Keep the per-suite topology contract: DeepSeek-V2-Lite gate cases default to + the V2 DP2/TP2 families on GPU and 2A2F on NPU; Qwen3 MoE and Qwen3.6 MoE use + 2A1F; DeepEP M2N requires at least two FFN ranks for EP, so its DeepSeek and + Qwen3 cases run 2A2F (4A2F for 235B FP8). A new case that breaks the suite's + default needs an explicit reason in the design page. +- Keep the four-device bound (`E2E-INV-002`): a PR case must not exceed four + unique devices without changing that normative invariant. - Treat a launch-script change in `recipe/` as a user contract: env vars, - device counts, and model paths in recipes must stay consistent with the E2E - scenarios and user guides that document them. + device counts, and model paths must stay consistent with the E2E scenarios + and user guides that document them. `VLLM_USE_V2_MODEL_RUNNER=1` is required + on both roles in every current recipe — a recipe that drops it will fail at + startup validation. - Keep baseline comparability: a topology or launch change that affects E2E gates must state how `baseline-graph` and AFD scenarios remain comparable. - Cover the matrix edge the change implies (graph/eager, DBO on/off, DBO + - 2A1F) in at least one named scenario or an explicit follow-up. -- Cross-node setups additionally - name host/port and rank assumptions; flag silent assumptions about network - interfaces or device ordering. + 2A1F, TP variants) in at least one named scenario or an explicit follow-up. +- Cross-node setups additionally name host/port and rank assumptions; multi-pod + runners live in `tests/e2e/multi_pod/` and + `.agents/skills/run-e2e/resources/k8-multi-pod.md`. Flag silent assumptions + about network interfaces, NVLink domains, or device ordering. Name the scenario or recipe that exercises the changed topology; for a new topology, require the E2E case and design-page update together. diff --git a/.agents/skills/review-pr/references/features/dspark-draft.md b/.agents/skills/review-pr/references/features/dspark-draft.md new file mode 100644 index 000000000..40cfd3729 --- /dev/null +++ b/.agents/skills/review-pr/references/features/dspark-draft.md @@ -0,0 +1,46 @@ +# DSpark Draft-Model Speculation + +Load with the primary module contract when a change touches DSpark: the +draft-model registration (`DSparkDraftModel` → +`AFDDSparkDeepseekV4ForCausalLM`), the speculator seam +(`afd_plugin/v1/worker/npu/dspark_speculator.py`), the target/draft layered FFN +wiring (`models/npu/deepseek_v4_dspark.py`), or the DSpark E2E evidence +harnesses under `tests/e2e/models/deepseek_v4_flash/`. + +Designs: `docs/design/module/model_integration.md`, +`docs/design/module/execution_platforms.md`, +`docs/design/module/attention_runtime.md`; +guide: `recipe/npu/CAMAsyncAFDConnector/deepseek_v4_flash/README.md`. + +## Scope + +DSpark is an Ascend-only draft model paired with a DeepSeek-V4 target under +`CAMAsyncAFDConnector`. The attention role keeps native eager proposal; the FFN +role may independently use layered eager or `FULL` graph execution. Supported +combinations are listed in the [support matrix](../../../../../docs/support.md#tested-runtime-matrix); everything else is +excluded rather than best-effort. + +## Feature checks + +- Respect the excluded set: DSpark does not support DBO or FFN two-stream, and + it does not extend to model families outside DeepSeek-V4. A PR that enables + either — even behind a flag — needs a design-page change and device evidence + first. +- The speculator patch (`AFDDSparkSpeculator.set_attn`, `.propose`) must follow + the patch contract against the pinned vLLM-Ascend ref, including signature + parity and `Patch reason:`/`Patch functionality:` comments. +- Target and draft must agree on layered FFN weight extraction + (`_extract_async_cam_layered_weights`) and on the layer set each side runs; + a mismatch fails at startup or, worse, silently runs a partial draft. +- The prefill/decode split used by the DSpark recipes is a launch contract: + `MODEL`, `DSPARK_CONFIG`, TP/DP/EP counts, and the upstream worker choice on + P must stay consistent between the recipe, the README, and the E2E harness. +- Evidence for a behavioral claim comes from the DSpark evidence harnesses + (`dspark_eager.py`, `dspark_evidence.DSparkEagerEvidence`, + `dspark_ffn_graph_evidence.DSparkFFNGraphEvidence`), which exercise SWA + boundaries and chunked prompts. A unit test alone cannot qualify draft + execution. + +Test the draft/target wiring CPU-side in +`tests/unit/model_executor/models/npu/test_deepseek_v4_dspark.py` and name the +NPU run for anything device-dependent. diff --git a/.agents/skills/review-pr/references/features/graphs-and-compilation.md b/.agents/skills/review-pr/references/features/graphs-and-compilation.md index 33dd8c361..f31d8d451 100644 --- a/.agents/skills/review-pr/references/features/graphs-and-compilation.md +++ b/.agents/skills/review-pr/references/features/graphs-and-compilation.md @@ -1,9 +1,11 @@ # Graphs and Compilation Load with the primary module contract for changes to graph capture or -compilation paths: FULL_DECODE_ONLY CUDA graph (`cuda_graph.py`, GPU), -ACL graph via the `mla_graph` patches (NPU), `_AFDCaptureEventTracker`, graph -metadata preparation, and graph/eager fallback. +compilation paths: `FULL_DECODE_ONLY` / `FULL` CUDA graph +(`v1/worker/cuda_graph.py`, GPU), ACL graph via the `mla_graph` module and +patches (`v1/worker/npu/mla_graph.py`, `compat/patches/npu/mla_graph.py`, NPU), +`_AFDCaptureEventTracker`, graph metadata preparation, and graph/eager +fallback. Designs: `docs/design/module/attention_runtime.md`, `docs/design/module/execution_platforms.md`. @@ -11,22 +13,37 @@ Designs: `docs/design/module/attention_runtime.md`, ## Feature checks - Keep capture/replay parity: every input to a captured graph (metadata, - buffers, connector state, ubatch keys) must be identical between capture - and replay; flag any capture-time-only branch that replay can also reach. -- Respect decode-only scoping: CUDA graph paths must not trigger for prefill - or mixed batches; the FULL_DECODE_ONLY gate stays enforced where the - connector contract declares it. -- Keep ACL graph patches pinned: `mla_graph` patches - `vllm_ascend/attention/mla_v1.py` at commit `80d8c194f` — apply the patch - contract (markers, upstream source, signature parity) and re-verify the copy - against that ref. + buffers, connector state, ubatch keys, async CAM graph payload) must be + identical between capture and replay; flag any capture-time-only branch that + replay can also reach. `_AFDCaptureEventTracker` asserts that capture input + preparation ran exactly as expected — do not accept weakening or deleting + those assertions without a replacement guarantee. +- Respect mode scoping. GPU uses `FULL_DECODE_ONLY` where the connector + contract says so; Ascend CAM uses `FULL`, rejects piecewise and decode-only + modes, and requires plain MLA graph + DBO to use `FULL_DECODE_ONLY`. A change + that lets an unsupported mode through to capture is a P1. +- Keep ACL graph patches pinned: `compat/patches/npu/mla_graph.py` patches + `vllm_ascend/attention/mla_v1.py` at the pinned vLLM-Ascend ref — apply the + patch contract (markers, upstream source, signature parity) and re-verify the + copy against that ref. - Preserve the eager escape hatch: every graph path has a tested eager fallback, and fallback selection is explicit config or documented platform - behavior, not an accident of ordering. -- State resource implications: graph size, memory pools, and capture-count - changes belong in the PR description and the relevant user guide; a capture - change that shifts memory is a perf claim needing evidence. + behavior, not an accident of ordering. Async CAM explicitly has *no* + automatic eager fallback: unsupported FULL configurations and mismatched + transport parameters fail startup. +- Graph-owned resources are bounded: async CAM FULL owns four private graph + pools and four graphs; FFN two-stream FULL owns its own pools per slot. State + resource implications (graph size, memory pools, capture counts) in the PR + description and the relevant user guide; a capture change that shifts memory + is a perf claim needing evidence. +- Async CAM graph startup is a host-only handshake + (`afd_plugin/async_cam_startup.py`, `FFNGraphPlan`): the requested + eager/full mode, executor error, and signature must agree between roles + before HCCL creation, with `GRAPH_STARTUP_TIMEOUT_SECONDS` bounding the wait. + A change to those keys is a cross-role protocol change. Test graph gating and metadata parity CPU-side where possible -(`tests/unit/v1/worker/`); device behavior claims need the named GPU scenario -(`afd-graph-*`, `afd-graph-dbo-*`) or NPU run. +(`tests/unit/v1/worker/test_cuda_graph.py`, `test_npu_mla_graph.py`, +`test_async_ffn_graph.py`, `tests/unit/connectors/test_async_cam_startup.py`); +device behavior claims need the named GPU scenario (`afd-v2-graph-*`, +`afd-v2-graph-dbo-*`) or the NPU run recorded per `tests/npu/README.md`. diff --git a/.agents/skills/review-pr/references/modules/compat-patches.md b/.agents/skills/review-pr/references/modules/compat-patches.md index b47fd0360..d6fa2b984 100644 --- a/.agents/skills/review-pr/references/modules/compat-patches.md +++ b/.agents/skills/review-pr/references/modules/compat-patches.md @@ -5,9 +5,14 @@ reviewed head. Authoring rules: `AGENTS.md` ("Patching Requirement"). Use for monkey patches under `afd_plugin/compat/patches/` (generic engine/GPU: `async_dp_engine.py`, `async_dp_forward_context.py`, -`config_validation.py`, `engine_core.py`) and the version gate in -`compat/vllm.py`. Does not own the deferred NPU runtime layer -([npu-compat.md](npu-compat.md)) or plugin-owned classes. +`config_validation.py`, `engine_core.py`, `gpu_mrv2_ubatching.py`, +`gpu_deepseek_v4_dbo.py`), the vLLM-Ascend patches under +`compat/patches/npu/` (`ascend_config.py`, `ascend_platform.py`, `mla_graph.py`, +`mrv2_ubatching.py`), the version gate in `compat/vllm.py`, and the patched +methods in `afd_plugin/v1/worker/**` (`attention_worker.py`, `ffn_worker.py`, +`ffn_model_runner.py`, `attention_model_runner_v2.py`, and the NPU twins). +Does not own the deferred NPU runtime layer ([npu-compat.md](npu-compat.md)) or +plugin-owned classes. ## Contract checks @@ -15,25 +20,41 @@ engine/GPU: `async_dp_engine.py`, `async_dp_forward_context.py`, new monkey patch; a new patch needs an architectural-review rationale: correct upstream target, minimal scope, understood performance impact, and a long-term upstream-or-removal plan. -- Copy the upstream function wholesale from the pinned ref (vLLM 0.26.0; - vLLM-Ascend commit `80d8c194f`) and mark only AFD-specific differences with - `# ### PATCH START:` / `# ### PATCH END:`; keep the marker text short and - specific. +- Copy the upstream function wholesale from the pinned ref and mark only + AFD-specific differences with `# ### PATCH START:` / `# ### PATCH END:`; + keep the marker text short and specific. - Keep a patched function's signature and return type identical to upstream; any added parameter must be documented in the comments immediately above the patch function. -- Require `Patch reason:` / `Patch functionality:` comments immediately above - every patch function, and keep `# Upstream source:` pointers accurate. +- Require, immediately above every patch function: `# Upstream source:` + (module, symbol, and the pinned ref — a bare `vLLM 0.32.0` label is weaker + than the commit but acceptable when it names the file and symbol), + `Patch reason:`, and `Patch functionality:`. `# Removal plan:` and + `# Expansion exception:` are expected where the patch delegates to + `_original_*` or expands a large upstream function; their absence there is a + finding. - Treat `_original_*` delegation as the exception, not the default non-AFD path; require the exception to be called out in the patch comments. - Keep patches idempotent and version-aware (guarded by - `TARGET_VLLM_VERSION`); re-applying or importing twice must not corrupt - state. + `TARGET_VLLM_VERSION` / `is_target_vllm_compatible`); re-applying or + importing twice must not corrupt state. - Do not flag upstream style inside patch copies — ruff ignores E501/N806/N807 - there so diffs stay comparable to upstream; do flag upstream drift (the copy - no longer matches the pinned ref). + under `afd_plugin/compat/patches/**` so diffs stay comparable to upstream; do + flag upstream drift (the copy no longer matches the pinned ref). - Patched behavior must be transparent when AFD is inactive: non-AFD requests execute the copied upstream logic unchanged. -Test every patch in `tests/unit/compat/patches/` with CPU-safe tests covering -idempotence, version guarding, and both AFD-on and AFD-off behavior. +## Signature alignment + +`tests/unit/compat/patches/test_signatures.py` is the single parametrized table +that owns upstream signature checks for every scoped runtime patch, including +the ones that live in `afd_plugin/v1/worker/**` (capture inputs, fullgraph +replay, workspace slot, ubatch prepare, CAMP2P dispatch) and the NPU +`vllm_ascend` installers. It reads the installed upstream definition from +source rather than comparing a patch to itself. A newly patched function with +no table row, or a row whose target symbol no longer exists upstream, is a P1 +finding: the next pin move will fail or, worse, silently stop being checked. + +Test every patch in `tests/unit/compat/patches/` (and +`tests/unit/compat/patches/npu/`) with CPU-safe tests covering idempotence, +version guarding, and both AFD-on and AFD-off behavior. diff --git a/.agents/skills/review-pr/references/modules/connectors.md b/.agents/skills/review-pr/references/modules/connectors.md index 954d541a2..3e9647452 100644 --- a/.agents/skills/review-pr/references/modules/connectors.md +++ b/.agents/skills/review-pr/references/modules/connectors.md @@ -1,21 +1,28 @@ # Connectors Primary design: `docs/design/module/connector_contracts.md` in the reviewed -head; backend guides in `docs/gpu/NCCL_P2P_CONNECTOR_USER_GUIDE.md` and -`docs/npu/CAM_P2P_CONNECTOR_USER_GUIDE.md` / `docs/npu/CAM_ASYNC_CONNECTOR_USER_GUIDE.md`. +head; backend guides in `docs/gpu/NCCL_P2P_CONNECTOR_USER_GUIDE.md`, +`docs/gpu/DEEPEP_M2N_CONNECTOR_USER_GUIDE.md`, +`docs/npu/CAM_P2P_CONNECTOR_USER_GUIDE.md`, and +`docs/npu/CAM_ASYNC_CONNECTOR_USER_GUIDE.md`. -Use for `afd_plugin/connectors/`: `AFDConnectorBase`, `AFDControlPlane`, -`ConnectorExtraInfo`, `AFDConnectorFactory`, and the concrete backends — -`gpu/p2p.py` (NCCL P2P), `npu/camp2p.py` (HCCL/CAM P2P), `npu/async_cam.py` -(experimental async CAM). Connectors do not own process-group creation or rank -mapping ([distributed-topology.md](distributed-topology.md)), model-side -proxies, or role workers. +Use for `afd_plugin/connectors/`: `AFDConnectorBase`, +`AFDRoutedExpertsConnector`, `AFDControlPlane`, `ConnectorExtraInfo`, +`AFDConnectorFactory`, and the concrete backends — `gpu/p2p.py` (NCCL P2P), +`gpu/deepep_m2n/` (DeepEP routed-experts M2N), `npu/camp2p.py` +(HCCL/CAM P2P), and `npu/async_cam.py` (async CAM). Connectors do not own +process-group creation or rank mapping ([distributed-topology.md](distributed-topology.md)), +model-side proxies, or role workers. ## Contract checks - Register backends through `AFDConnectorFactory`; role workers, runners, and models must stay connector-agnostic and never branch on a concrete - connector type. + connector type. The factory stores lazy loaders, so importing it must not + import CUDA or Ascend implementations. Registration alone does not admit an + external connector: `AFDConfig` validates `connector` against its own + allow-list in `afd_plugin/config.py`, and the design page records that gap as + draft. A PR claiming "extensible connectors" must close both. - Change both sides of the transfer boundary together: `ConnectorExtraInfo` handshake, tensor shapes/dtypes/layout, and completion signaling must agree between attention and FFN implementations, including metadata versioning @@ -24,14 +31,42 @@ proxies, or role workers. cancellation, and close must surface as loud errors, not a silent hang of both roles; every state machine handles abort before completion. - Declare graph-capture compatibility per connector (for example - FULL_DECODE_ONLY CUDA graph for P2pNccl, ACL graph constraints for CAM) and - encode unsupported combinations in `feature_validation`, not runtime - surprises. -- Treat additions under `connectors/npu/bin/` (vendored `.run`/`.whl` - artifacts) as design-level changes: require source, version, and target - CANN/Ascend runtime justification; they are pre-commit-excluded by design, - which is not a license to skip review. + `FULL_DECODE_ONLY` CUDA graph constraints for the CUDA connectors, ACL graph + constraints for CAM) and encode unsupported combinations in + `feature_validation`, not runtime surprises. +- Keep `connector_extra_config` connector-owned: unknown keys must fail in the + selected connector's parser, not be silently ignored. Adding a key to one + backend is a user-contract change that needs the guide and design-page table + updated together. + +## Routed-experts transports + +`DeepEPM2NAFDConnector` extends the contract through +`AFDRoutedExpertsConnector`: routing runs on Attention, FFN ranks receive one +row per (token, local expert) pair, run only local experts via `run_experts()`, +and send weighted rows back. `synchronize_lanes()` joins connector-forked +streams at the end of each step. When reviewing changes here: + +- The FFN worker selects `GPUFFNRoutedExpertsModelRunner` from the connector + class, not from the connector name; keep that selection type-based. +- The runner pins vLLM's MoE backend to `triton` because other backends + reorder expert weights in place without changing dtype or shape, producing + wrong outputs with no error. A new backend needs a layout proof, not an + allow-list entry. +- Dispatch/combine are collectives over all ranks, including ranks with no + tokens for a layer; a zero-length split is valid, an empty list is not. +- Weight loading must report what each rank actually loaded through + `init_afd_connector(moe_layer_indices=, local_experts=)`; skipping the + post-load check silently admits mismatched expert placement. + +## Vendored Ascend artifacts + +Additions under `connectors/npu/bin/` (`CAM_ascend*_openEuler_*.run`, +`umdk_cam_op_lib-*.whl`) are design-level changes: require source, version, and +target CANN/Ascend runtime justification. They are pre-commit-excluded by +design, which is not a license to skip review. Test each backend in `tests/unit/connectors/` (state machines, handshake, -failure paths, CPU-safe) and name the E2E gate scenario that exercises the -real transport. +failure paths, CPU-safe), name the device suite (`tests/gpu/` for M2N, +`tests/npu/` for CAM) and the E2E gate scenario that exercises the real +transport. diff --git a/.agents/skills/review-pr/references/modules/distributed-topology.md b/.agents/skills/review-pr/references/modules/distributed-topology.md index 451ff398b..d5e4f9ff4 100644 --- a/.agents/skills/review-pr/references/modules/distributed-topology.md +++ b/.agents/skills/review-pr/references/modules/distributed-topology.md @@ -4,11 +4,12 @@ Primary design: role lifecycle pages `docs/design/module/attention_runtime.md` and `docs/design/module/ffn_runtime.md`; connector handoff in `docs/design/module/connector_contracts.md`. -Use for `afd_plugin/distributed/`: `init_afd_process_group` and the -`DefaultProcessGroupSwitcher` (`afd_process_group.py`), `AFDRankMapping`, -`validate_p2p_topology`, `resolve_role_rank`, and `build_rank_mapping` -(`topology.py`). Does not own connector transfer semantics or worker startup -ordering. +Use for `afd_plugin/distributed/`: `init_afd_process_group`, +`create_hccl_process_group_options`, and `DefaultProcessGroupSwitcher` +(`afd_process_group.py`), plus `AFDRankMapping`, `topology_from_config`, +`validate_p2p_topology`, `resolve_role_rank`, `subgroup_attention_block`, and +`build_rank_mapping` (`topology.py`). Does not own connector transfer +semantics or worker startup ordering. ## Contract checks @@ -16,15 +17,26 @@ ordering. connectors, and models must not create overlapping groups or bypass the switcher. - Keep rank mappings total and bijective across the attention and FFN worlds; - `validate_p2p_topology` must reject mismatched world sizes (2A2F/2A1F) - before any collective runs. + `validate_p2p_topology` must reject mismatched world sizes before any + collective runs. - Derive backend choice (HCCL vs NCCL) from the detected platform; no silent fallback to a different backend. - Resolve roles and ranks deterministically from the same inputs on every - rank; any divergence is a startup hang or collective mismatch. + rank; any divergence is a startup hang or collective mismatch. Per-process + role ranks are runtime state resolved from vLLM's parallel placement and are + deliberately absent from `AFDConfig` — a PR that stores one there is a + contract change. +- Topology assumptions differ per connector and must stay coherent with them: + P2P and CAM connectors use FFN-then-Attention world ordering with role ranks + equal to DP x TP; DeepEP M2N treats the AFD world as one symmetric EP group + with attention ranks holding dummy expert slots, requires `num_experts` to + be divisible by the FFN rank count, and lets attention TP ranks act as + separate union ranks. `afd_process_group_timeout_s` (default 120 s) is the + rendezvous knob; `hccl_buffer_size` sizes connector-owned HCCL domains only. - Release AFD groups and restored defaults on shutdown and engine restart; no leaked process groups or stale switcher state. Test mapping construction, validation rejections, and teardown in -`tests/unit/distributed/`; name the E2E topology (2A2F or 2A1F) that covers -real collectives. +`tests/unit/distributed/` (`test_topology.py`, `test_afd_process_group.py`); +name the E2E topology (1A1F, 2A1F, 2A2F, 4A2F, 4A4F, 2P1A1F, or multi-pod) +that covers real collectives. diff --git a/.agents/skills/review-pr/references/modules/e2e-harness.md b/.agents/skills/review-pr/references/modules/e2e-harness.md new file mode 100644 index 000000000..eec7ee331 --- /dev/null +++ b/.agents/skills/review-pr/references/modules/e2e-harness.md @@ -0,0 +1,64 @@ +# E2E Harness and Gate Selection + +Primary design: `docs/design/module/e2e_testing.md` in the reviewed head. That +page is `normative`; its `E2E-INV-001` … `E2E-INV-008` invariants are +blockers when broken. Author-side runner: the `run-e2e` skill. + +Use for `tests/e2e/**`, the Buildkite case selection in +`.buildkite/cuda/*.yml`, `.buildkite/common/scripts/skip_ci.py` as it affects +what runs, and the `run-e2e` / `deploy-dsv4-flash-async` skill text that +describes cases. + +## Contract checks + +- Case IDs are stable lower-kebab-case and unique (`E2E-INV-001`). Renaming a + case silently breaks every CI node ID and every recipe/README that names it; + require the rename to land with the CI, docs, and skill updates. +- Gate cases never skip (`E2E-INV-004`): no `skip`, `xfail`, or + success-on-empty. A new `pytest.importorskip` or a conditional `return` in a + gate case is a P1. Missing setup must fail. +- CI selects cases by pytest node ID and must not lower sample counts, + thresholds, or the required set (`E2E-INV-007`). Compare + `.buildkite/cuda/test-ready.yml` / `test-merge.yml` / `test-weekly.yml` + against the module's `SCENARIOS` tuple: a case added to the module but not + selected (or selected but renamed) is drift worth reporting. +- The current gate set, which a change must keep coherent: + - DeepSeek-V2-Lite GPU: seven cases — `baseline-graph`, `afd-v2-eager-dp2`, + `afd-v2-graph-dp2`, `afd-v2-eager-tp2`, `afd-v2-graph-tp2`, + `afd-v2-eager-dbo-dp2`, `afd-v2-graph-dbo-dp2`; the module also holds + local-only `afd-v2-{eager,graph}-1a1f`, `afd-v2-graph-dbo-tp2`, and the + `afd-{eager,graph,graph-dbo}-2a1f` cases. + - DeepSeek-V2-Lite NPU: `baseline-graph` plus `afd-eager-2a2f`, + `afd-graph-2a2f`, `afd-graph-dbo-2a2f` (selected by + `AFD_E2E_BACKEND=npu`). + - Qwen3 MoE GPU: `baseline-graph`, `afd-eager-2a1f`, `afd-graph-2a1f`, + `afd-graph-dbo-2a1f`. The module also collects the DeepEP M2N 2A2F + scenarios (`afd-deepep-m2n-{eager,graph,graph-dbo}-2a2f` and the two + Attention-TP variants) through `test_qwen3_moe_deepep_m2n`; CI selects node + IDs explicitly and does not run them. + - Qwen3.6 MoE GPU: `baseline-graph`, `afd-eager-2a1f`, `afd-graph-2a1f` + (GDN DBO unsupported). + - DeepSeek-V4 Flash: separate `test_gpu.py`, `test_async_cam_npu.py`, and + `test_sync_camp2p_npu.py` suites; DSV4 Flash W4A8 async CAM is a local-only + 16-NPU case. +- Accuracy plumbing: ordinary cases evaluate the first 7 GSM8K samples, DBO + cases require a 24-sample minimum, and the GPU DP2 eager/eager-DBO/graph-DBO + comparison keeps 128 samples, concurrency 12, `max_num_seqs=8`, and the 0.27 + accuracy threshold. A change here is a gate change. +- Execution evidence: GPU and NPU graph/DBO cases load a test-only worker + extension and must report completed replays or two-stage execution on every + participating rank; dummy, profile, capture, and empty-second-microbatch work + does not qualify. Pure eager and baseline cases must not load the extension, + and performance comparisons must omit it. +- Process hygiene: child processes use owned process groups; cancellation is + SIGTERM, bounded grace, then reaping. The NPU Async CAM marker-based SIGKILL + sweep and its 120 s survivor window are a documented test-scoped exception, + not a general strategy. +- Harness code is not unit-tested (`tests/unit` covers `afd_plugin` only); + validate harness changes by running the affected E2E cases or + `tests/e2e/test_runner.py` / `test_deploy_dsv4_flash_async_skill.py`. + Production code must never import the harness. + +A PR that changes a case, sample count, device layout, or gate selection +without updating `docs/design/module/e2e_testing.md` and the `run-e2e` skill +narrative is a P2 (P1 when the change weakens a blocking gate). diff --git a/.agents/skills/review-pr/references/modules/engine-orchestration.md b/.agents/skills/review-pr/references/modules/engine-orchestration.md index db4982984..80cb3ad11 100644 --- a/.agents/skills/review-pr/references/modules/engine-orchestration.md +++ b/.agents/skills/review-pr/references/modules/engine-orchestration.md @@ -5,10 +5,13 @@ Primary design: `docs/design/module/attention_runtime.md` and head. Use for cross-role request flow semantics: async-DP engine behavior from the -patched `EngineCore`/`DPEngineCoreProc` busy loops, AFD request routing -between attention and FFN, the DBO yield custom op, ubatching orchestration, -and the `AFDBalancedRoutingStrategy` routing simulator. Does not own the patch -mechanics ([compat-patches.md](compat-patches.md)) or worker internals +patched `EngineCore` / `EngineCoreProc` / `DPEngineCoreProc` busy loops, +`DPAsyncMPClient.add_request_async`, AFD request routing between attention and +FFN, the DBO yield custom op (`v1/worker/dbo.py` → +`maybe_apply_dbo_yield`), ubatching orchestration +(`v1/worker/ubatch_wrapper.py`), and the `AFDBalancedRoutingStrategy` routing +simulator. Does not own the patch mechanics +([compat-patches.md](compat-patches.md)) or worker internals ([workers-runners.md](workers-runners.md)). ## Contract checks @@ -16,19 +19,31 @@ mechanics ([compat-patches.md](compat-patches.md)) or worker internals - Keep one owner for attention→FFN routing; models and workers must not independently select a route or a peer. - Keep the async-DP patches externally transparent: non-AFD requests observe - upstream vLLM 0.26.0 scheduling, output order, and finish semantics - unchanged, on both `EngineCoreProc` and `DPEngineCoreProc` paths. + upstream scheduling, output order, and finish semantics unchanged, on both + `EngineCoreProc` and `DPEngineCoreProc` paths. - Preserve request identity and ordering across the two roles; correlate every connector result, ubatch work item, and DBO yield to the originating request, and make terminal state monotonic (no forwarding after finish, abort, or failure). - Define shutdown ordering between the worlds: FFN shutdown must not hang attention and vice versa; cancellation propagates to in-flight transfers and - queued ubatch items. -- Keep the balanced routing strategy benchmark-only and explicitly registered; - it must never become the default routing behavior silently. + queued ubatch items. DeepSeek-V2-Lite has known DBO shutdown failures — check + whether a PR claims to fix that honestly, with evidence, rather than + suppressing the symptom. +- Native DBO is a vLLM capability that AFD reuses: exactly two microbatches, + enabled on Attention only. Async CAM FFN consumes work items outside + Attention's DBO thread schedule. GPU DBO accepts Attention DP1 or DP > 1; + NPU graph DBO with plain MLA requires `FULL_DECODE_ONLY`. A change that adds + a second scheduling path or relaxes the two-ubatch bound is a design change. +- Keep the balanced routing strategy benchmark-only and explicitly registered + (gated by `AFD_FORCE_BALANCED_TOPK_IDS`); it must never become the default + routing behavior silently. Test engine-loop patch behavior (AFD on/off), ubatch state transitions, and -shutdown/cancellation in `tests/unit/compat/patches/` and -`tests/unit/model_executor/`; name the DBO E2E scenario that exercises the -real path. +shutdown/cancellation in `tests/unit/compat/patches/` +(`test_engine_core.py`, `test_async_dp_engine.py`, +`test_async_dp_forward_context.py`, `test_gpu_mrv2_ubatching.py`), +`tests/unit/v1/worker/test_dbo.py`, +`tests/unit/model_executor/test_async_cam_dbo.py`, and +`tests/unit/model_executor/test_routing_simulator.py`; name the DBO E2E +scenario that exercises the real path. diff --git a/.agents/skills/review-pr/references/modules/execution-platforms.md b/.agents/skills/review-pr/references/modules/execution-platforms.md index 00aaa7d89..d3f42c909 100644 --- a/.agents/skills/review-pr/references/modules/execution-platforms.md +++ b/.agents/skills/review-pr/references/modules/execution-platforms.md @@ -4,12 +4,14 @@ Primary design: `docs/design/module/execution_platforms.md` in the reviewed head. Use for NPU/GPU selection and platform mechanics: `validation.py` -(`assert_compatible_afd_stack`, -`afd_worker_qualname_for_platform_default`), `setup.py` Ascend-op build -gating, and `csrc/npu/` (`ascend_kernels/` a2e/e2a ACLNN operators under the -`npu_op_*` build system, `pybind/` → `afd_plugin._C_ascend`, and the build -drivers in `scripts/`). Does not own connector behavior or the deferred NPU -patch application ([npu-compat.md](npu-compat.md)). +(`assert_compatible_afd_stack`, `afd_worker_qualname_for_platform_default`, +`validate_gpu_model_runner_v2_config`, +`validate_npu_model_runner_v2_config`), `setup.py` Ascend-op build gating, and +`csrc/npu/` (`ascend_kernels/` a2e/e2a and `grouped_matmul_layered` ACLNN +operators under the `npu_op_*` build system, `pybind/` → `afd_plugin._C_ascend`, +`build_aclnn.sh` installing into `afd_plugin/_cann_ops_custom`, and the drivers +in `scripts/`). Does not own connector behavior or the deferred NPU patch +application ([npu-compat.md](npu-compat.md)). ## Contract checks @@ -17,18 +19,30 @@ patch application ([npu-compat.md](npu-compat.md)). default worker classes), never hardcoded device strings; reject unsupported platforms (Ascend 310P/xlite) explicitly with an actionable error. - Keep the build contract: compile Ascend ops only in a detected Ascend - environment or with `AFD_BUILD_ASCEND_OPS=1`; honor - `AFD_SKIP_ACLNN_BUILD=1`; document `SOC_VERSION` defaults and any new build - knob in the install docs. -- Pair `csrc/` kernel changes (a2e/e2a hosts and kernels) with the torch - adapter/extension updates they require, and state accuracy and performance - impact; kernel behavior changes need NPU evidence, never simulation. -- Keep version pins coherent: vLLM 0.26.0 (`TARGET_VLLM_VERSION`) and - vLLM-Ascend `80d8c194f` move together across the compat gate, docker base - image, docs, and design-page `upstream_refs`, or the PR states why not. -- Keep GPU path parity visible: `csrc/gpu/` is reserved — a PR adding GPU - native ops must state where the GPU counterpart lives. + environment or with `AFD_BUILD_ASCEND_OPS=1` (unset/empty means auto-detect; + an invalid value must raise, not default); honor `AFD_SKIP_ACLNN_BUILD=1`; + document `SOC_VERSION` defaults and any new build knob in the install docs. +- Pair `csrc/` kernel changes (a2e/e2a and grouped-matmul-layered hosts and + kernels) with the torch adapter/extension updates they require, and state + accuracy and performance impact; kernel behavior changes need NPU evidence, + never simulation. Layered GMM selection changes belong to the + `upgrade-layered-gmm-ops` workflow. +- Keep version pins coherent: `TARGET_VLLM_VERSION` in `compat/vllm.py`, the + `vllm` extra in `pyproject.toml`, `VLLM_BASE_TAG` and its version assertion + in `docker/Dockerfile.ci`, installation/support references, and design-page `upstream_refs` move + together, or the PR states why not. Read the actual pin; do not trust the + value quoted in a review note or in this skill. +- Keep GPU path parity visible: `csrc/gpu/` contains only a README reserving + the directory — the plugin builds no GPU native extension. A PR adding GPU + native ops must state where the GPU counterpart lives and why it is not + plugin-owned today. +- Vendored `.run`/`.whl` artifacts under `afd_plugin/connectors/npu/bin/` and + the generated `_cann_ops_custom` tree are packaging, not source; a change + there needs version/CANN justification (see + [connectors.md](connectors.md#vendored-ascend-artifacts)). -Test platform selection and build gating in `tests/unit/v1/worker/` (runtime -classpath tests) and `tests/unit/package/`; kernel changes additionally name -the NPU run that produced their evidence. +Test platform selection and build gating in `tests/unit/v1/worker/` +(`test_runtime_classpaths.py`), `tests/unit/config/test_validation.py`, and +`tests/unit/package/` (`test_package.py`, `test_ascend_build_files.py`, +`test_ascend_op_host_build.py`, `test_cam_async_build.py`); kernel changes +additionally name the NPU run that produced their evidence. diff --git a/.agents/skills/review-pr/references/modules/model-integration.md b/.agents/skills/review-pr/references/modules/model-integration.md index 2db06fa89..450684a50 100644 --- a/.agents/skills/review-pr/references/modules/model-integration.md +++ b/.agents/skills/review-pr/references/modules/model-integration.md @@ -7,16 +7,19 @@ Use for `afd_plugin/model_executor/`: the registration maps in `afd_plugin/__init__.py` (`_DEEPSEEK_MODEL_REGISTRATIONS`, `_QWEN_MODEL_REGISTRATIONS`, `_QWEN3_5_MODEL_REGISTRATIONS`), role-split model wrappers (`RemoteFFNProxy`, `AFDAttentionFusedMoE`, `GateOnlyRemoteMoE`, the -`AFDDeepseekV2/V3/V4` and `AFDQwen3Moe`/`AFDQwen3_5` families), checkpoint -weight-role filtering, and the NPU-specific forward/gate modules under -`models/npu/`. Does not own runner execution, connector transport, or the -routing simulator's registration mechanics. +`AFDDeepseekV2/V3/V4` and `AFDQwen3Moe`/`AFDQwen3_5` families, plus +`AFDGlmMoeDsaForCausalLM` and `AFDDSparkDeepseekV4ForCausalLM`), checkpoint +weight-role filtering, the NPU-specific forward/gate modules under +`models/npu/`, and the layered executor under `model_executor/npu/`. Does not +own runner execution, connector transport, or the routing simulator's +registration mechanics. ## Contract checks - Keep registrations explicit in the `register_afd()` maps; select NPU vs GPU variants (for example `DeepseekV4ForCausalLM` branching on `torch_npu`) at - registration time, never deep inside forward paths. + registration time, never deep inside forward paths. The maps are frozen with + `MappingProxyType` — a change that mutates them at runtime is a finding. - Preserve the role construction contract: each role builds only role-required components, and weight filtering via `_checkpoint_weight_roles` must not silently drop weights that role later @@ -28,11 +31,19 @@ routing simulator's registration mechanics. `getattr`/`hasattr` guard around vLLM model fields hides the exact incompatibility the pinned-version contract exists to surface. - Mark copied upstream model code with the patch contract (markers, upstream - source pointers); the qwen3_5 and deepseek_v4 files are among the heaviest - patch surfaces. + source pointers); `qwen3_5.py`, `deepseek_v4.py`, and + `models/npu/deepseek_v4.py` are among the heaviest patch surfaces. +- Family-specific constraints are contracts, not suggestions: Qwen3.5/3.6 GDN + DBO is unsupported; DeepEP M2N covers Qwen3 MoE (all-MoE layers, no shared + expert); Async CAM covers DeepSeek-V4 W4A8/W8A8 with static EP and + Attention-local SP, and no longer supports DeepSeek-V2/V2-Lite. A PR that + widens a family's supported matrix without E2E and design-page evidence is a + finding. - Register new model families against the [model-addition-checklist.md](../checks/model-addition-checklist.md): both role wrappers, registration, recipe, E2E suite, and docs. Test wrappers and weight filtering in `tests/unit/model_executor/` -(CPU-safe, checkpoint-driven); name the E2E model suite covering the family. +(`models/test_*_construction.py`, `models/test_*_weight_policy.py`, +`models/test_*_proxy.py`, `models/npu/`, and `test_async_cam_*.py`); name the +E2E model suite covering the family. diff --git a/.agents/skills/review-pr/references/modules/npu-compat.md b/.agents/skills/review-pr/references/modules/npu-compat.md index 60976b510..a41d6698b 100644 --- a/.agents/skills/review-pr/references/modules/npu-compat.md +++ b/.agents/skills/review-pr/references/modules/npu-compat.md @@ -6,10 +6,11 @@ sections) in the reviewed head; platform init ordering lives in Use for `afd_plugin/compat/npu/`: deferred patch application (`runtime.py::apply_afd_ascend_patches_if_needed`), runtime config, forward -context, feature validation, Ascend ops loading (`ops.py`), and the NPU -profiler. Does not own the vLLM-Ascend patch files under -`compat/patches/npu/` ([compat-patches.md](compat-patches.md)) or the NPU -workers ([workers-runners.md](workers-runners.md)). +context, feature validation, Ascend ops loading (`ops.py`), the opaque +routing-id balancing boundary (`routing_ids.py`), and the NPU profiler. Does +not own the vLLM-Ascend patch files under `compat/patches/npu/` +([compat-patches.md](compat-patches.md)) or the NPU workers +([workers-runners.md](workers-runners.md)). ## Contract checks @@ -22,9 +23,19 @@ workers ([workers-runners.md](workers-runners.md)). silently continue when ops are missing. - Guard every `torch_npu`/`vllm_ascend` import so module import stays CPU-safe; runtime checks happen inside guarded functions. +- `routing_ids.py` is an opaque boundary registered during worker + initialization; keep profile routing decisions inside the op rather than + exporting them as Python-visible state. Standalone scripts import the module + after device init. Its fixture lives in + `tests/unit/fakes/routing_ids_op.py` and is registered from + `tests/unit/conftest.py`; a new consumer that bypasses the fixture will break + CPU collection. - Keep NPU behavior observable: a deferred patch or runtime tweak that changes execution must be loggable and attributable when troubleshooting per `docs/npu/TROUBLESHOOTING.md`. -Test in `tests/unit/compat/npu/`: application ordering, idempotence, feature -rejections, and op-presence failure paths — all CPU-safe via guarded imports. +Test in `tests/unit/compat/npu/` (`test_runtime_config.py`, +`test_feature_validation.py`, `test_ascend_config_lifecycle.py`, +`test_ascend_engine_core_binding.py`, `test_routing_ids.py`, +`test_profiler.py`): application ordering, idempotence, feature rejections, and +op-presence failure paths — all CPU-safe via guarded imports. diff --git a/.agents/skills/review-pr/references/modules/plugin-config.md b/.agents/skills/review-pr/references/modules/plugin-config.md index f1891c355..b0c0eaa17 100644 --- a/.agents/skills/review-pr/references/modules/plugin-config.md +++ b/.agents/skills/review-pr/references/modules/plugin-config.md @@ -1,27 +1,50 @@ # Plugin Entry and Config Primary design: `docs/design/module/plugin_boundary.md` in the reviewed head. +This page is `normative`; its `ENTRY-INV-*` and `CFG-INV-*` invariants are +blockers when broken. Use for plugin registration (`register_afd()`), the `AFDConfig` dataclass and -schema, `parse_afd_config`/`validate_afd_config`, extra-config coercion, and -`afd_plugin/envs.py` flags. Plugin config does not own platform/worker -selection, patch mechanics, or worker construction. +schema, `parse_afd_config` / `parse_optional_afd_config` / +`validate_afd_config`, extra-config coercion, `afd_plugin/envs.py` flags, and +`afd_plugin/async_cam_startup.py` (the host-only Async CAM graph startup +protocol). Plugin config does not own platform/worker selection, patch +mechanics, or worker construction. ## Contract checks - Keep `register_afd()` import-time work minimal, idempotent, and version-guarded; importing `afd_plugin` must stay CPU-safe (no torch, - torch_npu, or vLLM runtime imports at module level). + torch_npu, or vLLM runtime imports at module level). The package uses + module-level `__getattr__` for the runtime classes and + `_force_spawn_multiprocessing_if_requested()` at import time — that call is + opt-in behind `AFD_FORCE_SPAWN_MULTIPROCESSING=1` and must stay that way. - Parse AFD settings only from the `afd` key of `--additional-config`; route every derived consumption through validated `AFDConfig` fields, not raw dicts, and fail fast on unsupported combinations with an actionable message. + Short keys are canonical (`connector`, `role`, `port`, `host`, + `num_attention_ranks`, `num_ffn_ranks`, `compute_gate_on_attention`, + `async`→`async_dp`, `afd_process_group_timeout_s`, + `connector_extra_config`); the `afd_*` names survive only as read-only + compatibility aliases, and unknown keys belong under + `connector_extra_config`, not at the top level. - Give every new `AFDConfig` field a default, validation, a stated role and platform implication (attention vs FFN, GPU vs NPU), and a docs update; changing a default or meaning of an existing field is a user-contract change. + `compute_hash()` covers only graph-affecting fields — a new field that + affects graph capture must join it or capture will be reused wrongly. - Read `envs.py` flags at use time, keep them opt-in and documented, and never let an env flag silently override validated config in production paths. -- Keep the plugin entry point (`vllm.general_plugins`) the only installation - path; no side effects from importing submodules. + Today `envs.py` holds a single diagnostic flag + (`AFD_FORCE_BALANCED_TOPK_IDS`); a new runtime-behavior flag needs a design + and guide entry, not just a constant. +- Keep the plugin entry point (`vllm.general_plugins` → `afd_plugin:register_afd`) + the only installation path; no side effects from importing submodules. +- `async_cam_startup.py` is host-only and must not import device runtimes: it + is imported on both roles before and after HCCL setup. Test registration idempotence, config parse/validate rejections, defaults, and -CPU-safe import in `tests/unit/config/` and `tests/unit/package/`. +CPU-safe import in `tests/unit/config/` (`test_config.py`, +`test_validation.py`), `tests/unit/test_envs.py`, +`tests/unit/test_async_cam_startup_wiring.py`, and +`tests/unit/package/test_package.py`. diff --git a/.agents/skills/review-pr/references/modules/workers-runners.md b/.agents/skills/review-pr/references/modules/workers-runners.md index 2e2439ef9..8f50d50c2 100644 --- a/.agents/skills/review-pr/references/modules/workers-runners.md +++ b/.agents/skills/review-pr/references/modules/workers-runners.md @@ -4,33 +4,56 @@ Primary design: `docs/design/module/attention_runtime.md` and `docs/design/module/ffn_runtime.md` in the reviewed head. Use for `afd_plugin/v1/worker/` on both platforms: `AFDAttentionWorker` / -`AFDFFNWorker` and `AFDNPU*Worker`, the model runners (V1 and V2 variants), -attention/FFN metadata builders, `cuda_graph.py`, `dbo.py`, and the ubatch -wrappers. Does not own platform selection -([execution-platforms.md](execution-platforms.md)), connector transport, or -model wrappers ([model-integration.md](model-integration.md)). +`AFDFFNWorker` and `AFDNPUAttentionWorker` / `AFDNPUFFNWorker`, the +Attention ModelRunnerV2 variants (`AFDAttentionModelRunnerV2`, +`AFDNPUAttentionModelRunnerV2`), the FFN runners (`GPUFFNModelRunner`, +`GPUFFNRoutedExpertsModelRunner`, `AFDNPUFFNModelRunner`, +`AsyncNPUFFNModelRunner`), attention/FFN metadata builders, `cuda_graph.py`, +`dbo.py`, `ubatch_wrapper.py`, and the NPU `ubatch_runner.py`, +`mla_graph.py`, `async_ffn_two_stream.py`, and `dspark_speculator.py`. Does +not own platform selection ([execution-platforms.md](execution-platforms.md)), +connector transport, or model wrappers ([model-integration.md](model-integration.md)). ## Contract checks - Preserve the role split: FFN-side `execute_model()` is connector-driven and fails fast on scheduler-driven calls; attention-side runners never construct FFN components and vice versa. -- Keep runner V1/V2 in pairs: a change to shared runner behavior lands in both - variants or states the exclusion; GPU keeps V2 gated off - (`VLLM_USE_V2_MODEL_RUNNER=0` required) while NPU V2 exists — flag any change - that silently moves that gate. -- Enforce the patch contract on every patched runner method (the NPU - attention model runner is the heaviest patch site): marked AFD differences, - upstream-copy fidelity against the pinned ref, identical signatures. +- There is no V1 runner and no fallback. Both roles require + `VLLM_USE_V2_MODEL_RUNNER=1`; Attention runs native ModelRunnerV2 (GPU) or + the NPU V2 runner, while FFN keeps a connector-driven runner that is *not* + a V2 runner. `validate_gpu_model_runner_v2_config` and + `validate_npu_model_runner_v2_config` enforce the flag; do not accept a + change that turns the requirement into a warning or a silent default. +- Keep the three FFN execution paths distinct and selected by contract, not by + convenience: + | Path | Selected by | Runner | + | --- | --- | --- | + | Control-plane-driven | `connector.control_plane is not None` | `GPUFFNModelRunner` / `AFDNPUFFNModelRunner` | + | Routed experts | `issubclass(connector_cls, AFDRoutedExpertsConnector)` | `GPUFFNRoutedExpertsModelRunner` | + | Async CAM (no control plane) | `CAMAsyncAFDConnector` on NPU | `AsyncNPUFFNModelRunner` | + A change that collapses two of these, or routes by connector *name*, is a + finding. +- Enforce the patch contract on every patched worker/runner method — the + Attention runners are the heaviest patch sites on both platforms: marked AFD + differences, upstream-copy fidelity against the pinned ref, identical + signatures, `Patch reason:`/`Patch functionality:` comments, and a row in + `tests/unit/compat/patches/test_signatures.py`. - Keep graph capture sound: metadata must be prepared identically for capture and replay; capture-only state must not leak into eager paths; DBO yield and ubatch wrappers behave the same under graph and eager. -- Keep the lazy export map in `v1/worker/__init__.py` CPU-safe: dotted-path - resolution must not import runtime modules eagerly. + `_AFDCaptureEventTracker` exists to detect capture input-preparation drift — + a change that silences those assertions is worse than the drift. +- Keep the lazy export maps CPU-safe: `afd_plugin/v1/worker/__init__.py` and + `afd_plugin/v1/worker/npu/__init__.py` resolve dotted paths through + `__getattr__`; new runtime classes must be added there rather than imported + eagerly, or plugin import on a CPU host breaks. - Update the GPU and NPU twins together or prove platform neutrality; the NPU - runners (largest patch surface) and GPU runners drift independently. + runners (largest patch surface) and GPU runners drift independently, and the + NPU sync/async FFN runners diverge by design. Test role contracts, metadata building, and gating in -`tests/unit/v1/worker/` using the CPU-safe runtime-guard convention (see +`tests/unit/v1/worker/` and its `npu/` subtree using the CPU-safe runtime-guard +convention (see [test-quality-evaluation.md](../checks/test-quality-evaluation.md)); name the GPU/NPU gap for anything not run. diff --git a/.agents/skills/review-pr/references/process/design-contracts.md b/.agents/skills/review-pr/references/process/design-contracts.md index a9215355a..19a69f6a8 100644 --- a/.agents/skills/review-pr/references/process/design-contracts.md +++ b/.agents/skills/review-pr/references/process/design-contracts.md @@ -7,15 +7,25 @@ the ownership and status source of truth for the frozen head. ## Resolution order 1. Open `docs/design/module/index.md` in the reviewed checkout. It declares - the reading order and the dependency direction between module designs. + the reading order, the dependency direction, and a "Production path routing" + table that assigns every shipped path, native build path, and packaging + file to exactly one primary document. Tests, tools, recipes, generated + files, user guides, and design docs are outside that inventory; the E2E + harness is owned separately by `e2e_testing.md`. 2. Match the changed production paths against each design page's `primary_code_paths` (and `related_code_paths`) frontmatter to identify the owning design document; a path may appear in `related_code_paths` of several - pages while having one primary owner. -3. Read the matched page's `status`, `depends_on`, `validation_paths`, and - `last_reviewed` fields. `status: draft` means its boundaries and invariants - are candidates, not enforced contracts — treat them as questions, and raise - a blocking finding only when current code, tests, or AGENTS.md enforce them. + pages while having one primary owner. A new shipped file that appears in no + page is an ownership gap — that is a finding, not a reason to route by + guesswork. +3. Read the matched page's `status`, `owners`, `depends_on`, + `validation_paths`, `upstream_refs`, and `last_reviewed`. Today only + `plugin_boundary.md` and `e2e_testing.md` are `normative`; every other page + is `draft`. `status: draft` means its boundaries and invariants are + candidates, not enforced contracts — treat them as questions, and raise a + blocking finding only when current code, tests, or AGENTS.md enforce them. + A `normative` page's `MUST`/`MUST NOT` invariants (`ENTRY-INV-*`, + `CFG-INV-*`, `E2E-INV-*`) are blockers. 4. Prefer live code and tests over stale prose when they disagree; report the drift against the design page as a finding when the PR touches that area. 5. Honor the declared dependency direction: role modules (attention/FFN/model @@ -23,7 +33,13 @@ the ownership and status source of truth for the frozen head. boundary, execution platforms, compatibility). A PR that adds a dependency from a shared boundary back into a role implementation contradicts the documented architecture; question it. -6. Never import a newer design page into an older branch's review; release +6. Compare `upstream_refs` and `last_reviewed` against the pin in + `afd_plugin/compat/vllm.py` and `docs/design/module/index.md`. A page whose + `upstream_refs` still name the previous pair while the code moved is drift + worth reporting when the PR touches that area — but do not report a page as + stale merely because it records historical evidence; those rows are + explicitly historical. +7. Never import a newer design page into an older branch's review; release branches (for example `release/v0.19.1rc1`, `update/v0.28.0`) are judged against their own head's pages. @@ -36,8 +52,25 @@ A PR changes a documented contract, and must update the matching - changes connector, worker, platform, or compatibility behavior that a design page describes; - changes the E2E gate structure or pass criteria (`e2e_testing.md`); -- moves the pinned vLLM or vLLM-Ascend refs recorded in `upstream_refs`. +- moves the pinned vLLM or vLLM-Ascend refs recorded in `upstream_refs`, or + changes the `TARGET_VLLM_VERSION` gate in `afd_plugin/compat/vllm.py`; +- adds, retires, or renames an E2E case, or changes a gate's sample count, + device count, or required set (`e2e_testing.md`); +- changes behavior that an `.agents/skills/**` SKILL.md describes as a command, + case list, or pin. Doc-only updates to a design page route back to the production contract the page claims; verify the claim against code before accepting the page edit as the review's whole surface. + +## Ownership + +CODEOWNERS is no longer a single global rule. The global catch-all maps to +`@hsliuustc0106 @jiangkuaixue123`, with path-specific owners on top: +`/afd_plugin/connectors/gpu/`, `/csrc/gpu/`, `/recipe/gpu/`, `/tests/gpu/`, +`/docs/gpu/`, and `/.buildkite/cuda/` add `@ronenkat`; `/csrc/npu/` adds +`@francisleely`; `/tests/e2e/`, `/.agents/skills/run-e2e/`, and +`/.agents/skills/deploy-afd-k8s/` add `@ronenkat`. Design pages name their own +`owners`; `e2e_testing.md` is owned by `@yujuancao07` alone. When a design page +and CODEOWNERS disagree about who owns a contract, report the disagreement +rather than choosing one silently. diff --git a/.agents/skills/review-pr/references/process/general-checks.md b/.agents/skills/review-pr/references/process/general-checks.md index 8533a9e24..161d30c40 100644 --- a/.agents/skills/review-pr/references/process/general-checks.md +++ b/.agents/skills/review-pr/references/process/general-checks.md @@ -16,33 +16,56 @@ vLLM engine/config ingress (--additional-config {"afd": ...}) ``` - Prove the path in the frozen head, not from the PR description. -- Cover the matrix the change can affect: GPU/NPU, graph/eager, DBO on/off, - 2A2F/2A1F, runner V1/V2, colocation/disaggregation, AFD-on/AFD-off. +- Cover the matrix the change can affect: + + | Axis | Values in the current tree | + | --- | --- | + | Platform | CUDA, Ascend NPU | + | Execution | eager, `FULL`, `FULL_DECODE_ONLY` graph | + | Ubatching | DBO off, two native microbatches | + | Topology | 1A1F, 2A1F, 2A2F, 4A2F, 4A4F, 2P1A1F, multi-pod | + | Connector | `P2pNcclAFDConnector`, `DeepEPM2NAFDConnector`, `CAMP2pAFDConnector`, `CAMAsyncAFDConnector` | + | FFN execution | control-plane-driven, routed-experts, async layered GMM, FFN two-stream | + | Draft model | target-only, DSpark target + draft | + | AFD | on, off (transparent passthrough) | + + Not every cell exists; claim only the combinations the code and design pages + enumerate, and flag a change that silently widens the supported set. + +- There is no V1 model-runner fallback. Both roles require + `VLLM_USE_V2_MODEL_RUNNER=1`; Attention runs ModelRunnerV2 and FFN stays + connector-driven. `validate_gpu_model_runner_v2_config` / + `validate_npu_model_runner_v2_config` raise on a missing flag — a change that + relaxes or bypasses that check is a contract change, not a convenience. - Patched or wrapper code must be transparent when AFD is inactive - (`is_afd_active` false): a non-AFD request must keep upstream vLLM 0.26.0 - behavior exactly. + (`is_afd_active` false): a non-AFD request must keep upstream vLLM behavior + exactly, at the pinned ref. - Search bounded callers and sibling implementations (the GPU vs NPU twin of - every worker/model file) rather than assuming the changed hunk is the only - path; the two platforms drift independently. + every worker/model file, and the sync vs async NPU FFN runner pair) rather + than assuming the changed hunk is the only path; the platforms drift + independently. ## Module, API, and class design -These are the repo's own normative rules (AGENTS.md); enforce them as review -rules, not style nits: +These are the repo's own normative rules (AGENTS.md, docs/UT_MAINTENANCE.md); +enforce them as review rules, not style nits: 1. Access vLLM / vLLM-Ascend attributes directly; `getattr`/`hasattr` guards and proactive custom exceptions hide upstream-compatibility breakage that upgrades must surface. Flag them unless the change is genuinely optional. 2. Keep parameter types concrete; `Any`/`object` parameters hide contract - violations from mypy/pyright. + violations from mypy/ty. 3. No new magic numbers; name constants (`MAX_CONTEXT_LENGTH`, not `2048`). -4. No new mutable global state; constants and immutable config objects only. +4. No new mutable global state; constants and immutable config objects only + (`MappingProxyType` for the frozen registration maps is the house pattern). New mutable globals require approval and must be justified in the PR. 5. Do not split simple functions into helper layers; extract only for real complexity, real duplication, or an established local pattern. 6. Keep imports CPU-safe at module level: importing `afd_plugin` (or any unit - test module) must not import torch, torch_npu, or vLLM runtime modules. - Defer or `importorskip` them inside functions/tests. + test module) must not import torch, torch_npu, or vLLM runtime modules. The + package uses `__getattr__` lazy export maps in `afd_plugin/__init__.py`, + `afd_plugin/v1/worker/__init__.py`, and `afd_plugin/v1/worker/npu/__init__.py`; + keep new runtime classes behind those maps, not as eager imports. 7. New public behavior needs a named owner for its design page and a docs impact statement in the PR template. 8. Performance-affecting changes need a comparable A/B claim or an explicit @@ -53,16 +76,22 @@ rules, not style nits: | Risk | Prove before reporting | | --- | --- | | Correctness | A reachable input reaches the changed code and produces wrong output, a hang, or a crash. | -| Patch-contract drift | A patched function's signature, return type, or upstream copy diverges from the pinned ref, or missing/incorrect `# ### PATCH` markers, `Patch reason`/`Patch functionality` comments, or `# Upstream source:` pointers break upgrade comparability. | -| Version/compat | The change breaks the vLLM 0.26.0 / vLLM-Ascend `80d8c194f` contract, or moves a pin without updating the version gate, docker base, docs, and design pages together. | -| Platform divergence | The GPU or NPU twin (worker/runner/model variant) is not updated and the change is not provably platform-neutral. | +| Patch-contract drift | A patched function's signature, return type, or upstream copy diverges from the pinned ref, or missing/incorrect `# ### PATCH START/END` markers, `Patch reason:` / `Patch functionality:` / `# Upstream source:` comments break upgrade comparability. Patched methods now live in `afd_plugin/v1/worker/**` as well as `afd_plugin/compat/`, and `tests/unit/compat/patches/test_signatures.py` owns the signature table for both. | +| Version/compat | The change breaks the current `TARGET_VLLM_VERSION` / vLLM-Ascend pin contract, or moves a pin without updating the version gate, `docker/Dockerfile.ci` base tag and its assertion, `pyproject.toml`, installation/support references, and design-page `upstream_refs` together. Read the pin from `afd_plugin/compat/vllm.py`, not from memory. | +| Platform divergence | The GPU or NPU twin (worker/runner/model variant, sync vs async NPU FFN runner) is not updated and the change is not provably platform-neutral. | | CPU-safe imports | A new module-level import of torch/torch_npu/vLLM runtime breaks `pytest -m "not gpu and not vllm_runtime"` or plugin import on a CPU host. | -| Concurrency/async | Async-DP busy loops, ubatching state machines, DBO yield, or connector transfer states can deadlock, drop, reorder, or leak work items on abort/timeout. | +| Concurrency/async | Async-DP busy loops, ubatching state machines, DBO yield, async CAM work items, or two-stream slots can deadlock, drop, reorder, or leak work items on abort/timeout. | | Distributed/topology | Rank mapping is not bijective/validated, process groups are created outside the AFD owner, or teardown leaks groups across restarts. | -| Lifecycle | Graph capture state leaks into eager paths, shutdown ordering between attention and FFN worlds hangs either side, or connector close leaves a transfer half-open. | -| Security/validation | Untrusted `--additional-config`, env flags, or connector metadata are consumed without validation; host/port/rank inputs reach device code unchecked. | -| Evidence | Perf/accuracy claims lack comparable base/head runs; NPU behavior claims lack manual evidence per `docs/npu/TESTING.md`. | -| User contract | `AFDConfig` fields, defaults, or recipe launch scripts change meaning without a docs update and migration note. | +| Lifecycle | Graph capture state leaks into eager paths, shutdown ordering between attention and FFN worlds hangs either side, or connector close leaves a transfer half-open. Known open failures (for example DeepSeek-V2-Lite DBO shutdown, NPU graceful exit) must be named as pre-existing context, not silently repaired by assertion changes. | +| Security/validation | Untrusted `--additional-config`, `connector_extra_config`, env flags, or connector metadata are consumed without validation; host/port/rank inputs reach device code unchecked. | +| Routed-expert contract | A routed-experts connector (`DeepEPM2NAFDConnector`) change assumes a weight layout the pinned vLLM MoE backend does not preserve. `ROUTED_MOE_BACKEND = "triton"` is pinned deliberately; a new backend needs an explicit layout proof, not a silent allow-list addition. | +| Quantization/layer contract | Async CAM layered FFN rejects incompatible weights at startup (`dynamicQuant=1`, homogeneous W4A8/W8A8, layering). A change that falls back to another FFN path instead of failing startup contradicts the design page. | +| E2E gate integrity | Case IDs, sample counts, device counts, or the required case set change without a design-page update; a gate case gains a `skip`/`xfail`; CI stops selecting a node ID it used to select. `docs/design/module/e2e_testing.md` is `normative`, so its `E2E-INV-*` items are blockers. | +| CI suppression | The diff is docs-only or skip-mark-only, so `skip_ci.py` suppresses all default CI. A green pipeline then proves nothing about the code; say so instead of treating it as validation. | +| Vendored artifacts | New or changed files under `afd_plugin/connectors/npu/bin/` (`.run`/`.whl`) are pre-commit-excluded by design; they still need source, version, and target CANN justification. | +| Agent/skill drift | A change to production paths, CI, or E2E gates that leaves `.agents/skills/**` stating stale behavior (pins, case counts, commands) is a P2 finding; those skills are executed instructions. | +| Evidence | Perf/accuracy claims lack comparable base/head runs; NPU behavior claims lack manual evidence per `tests/npu/README.md`. | +| User contract | `AFDConfig` fields, defaults, `connector_extra_config` keys, or recipe launch scripts change meaning without a docs update and migration note. | ## Finding bar @@ -70,7 +99,8 @@ Report a finding only when you can name the changed `path:line`, the trigger or call path, the current behavior, the user or maintainer impact, and the smallest safe fix direction. Do not report: -- Style pre-commit already enforces (ruff, format, SPDX, markdownlint, DCO). +- Style pre-commit already enforces (ruff, format, typos, SPDX, markdownlint, + shellcheck, DCO). - Hypothetical issues with no reachable trigger in the supported matrix. - Missing tests that would not protect the changed behavior. - Debt the PR only modifies or removes. diff --git a/.agents/skills/review-pr/references/process/review-execution.md b/.agents/skills/review-pr/references/process/review-execution.md index 5436b37f7..b96d58d7c 100644 --- a/.agents/skills/review-pr/references/process/review-execution.md +++ b/.agents/skills/review-pr/references/process/review-execution.md @@ -113,6 +113,12 @@ Mark early findings as preliminary and continue. This is not a GitHub comment. blocks the verdict. Start with the first error, not the last cascade. - Treat the PR body as navigation. Commands, benchmark tables, and claimed tests become evidence only after their provenance and relevance are checked. +- Check what the pipeline actually selects before trusting a green run. + `.buildkite/common/scripts/skip_ci.py` suppresses the whole default pipeline + when every changed path is docs (`docs/**`, any `*.md`, `mkdocs.yml`) or a + pytest skip-mark-only test edit. On such a diff, "CI passed" is not evidence + about the code. Label-only CI files gate per platform and level; there is no + `.buildkite/npu/` pipeline yet. ## Run bounded validation @@ -129,14 +135,24 @@ paths have supported findings or explicit no-issue conclusions: 1. Import preflight on a trusted head: `python -c "import afd_plugin"` and the changed modules import without torch/torch_npu present. 2. Narrowest unit tests for the changed area, CPU-safe by default: - `pytest tests/unit/ -m "not gpu and not vllm_runtime"`. -3. Targeted GPU-gated tests only when a GPU and a trusted head exist. + `pytest tests/unit/ -m "not gpu and not vllm_runtime"`. Mirror the CI + selection when the change is broad: + `pytest -m "not gpu and not vllm_runtime" tests/unit tests/npu/test_async_cam_harness.py tests/npu/test_async_cam_performance.py tests/npu/test_async_cam_golden.py tests/npu/test_gmm_layered_performance.py`. +3. Targeted device tests only when the device and a trusted head exist: + `pytest tests/gpu/test_attention_routing.py` on CUDA (this one runs in CI), + `tests/gpu/test_deepep_m2n_gpu.py` on SM90+ with afdEP-comm and DeepGEMM + installed, and the `tests/npu/` device suites on Ascend. 4. E2E gate scenarios only through the `run-e2e` skill, only when the review genuinely needs active evidence and devices are free; otherwise request - author evidence. + author evidence. Gate cases must never be turned into skips to make a run + pass. 5. NPU paths have no in-repo CI. Require author-provided evidence per - `docs/npu/TESTING.md`, or a manually executed run the user explicitly asks + `tests/npu/README.md`, or a manually executed run the user explicitly asks for. Never simulate device evidence. +6. When a pin moves, verify the coherent set: `TARGET_VLLM_VERSION` in + `afd_plugin/compat/vllm.py`, the `vllm` extra in `pyproject.toml`, + `VLLM_BASE_TAG` and its assertion in `docker/Dockerfile.ci`, installation/support references, + and every design page's `upstream_refs`. ## Deliver maintainer-style findings diff --git a/.agents/skills/review-pr/references/process/review-routing.md b/.agents/skills/review-pr/references/process/review-routing.md index 5cf7f349a..658493264 100644 --- a/.agents/skills/review-pr/references/process/review-routing.md +++ b/.agents/skills/review-pr/references/process/review-routing.md @@ -12,14 +12,15 @@ path (for example a connector change that also moves rank mapping). | Module | Live behavior signals | Read | | --- | --- | --- | | Plugin entry & config | `register_afd()`, `AFDConfig`, parse/validate, extra-config coercion, env flags | [plugin-config.md](../modules/plugin-config.md) | -| Compatibility patches | New or changed monkey patch under `afd_plugin/compat/patches/`, version gate, upstream copy | [compat-patches.md](../modules/compat-patches.md) | -| NPU runtime compat | `compat/npu/` deferred patch application, feature validation, Ascend ops loading | [npu-compat.md](../modules/npu-compat.md) | -| Connectors | Connector backend, control plane, transfer/metadata handshake | [connectors.md](../modules/connectors.md) | +| Compatibility patches | New or changed monkey patch under `afd_plugin/compat/patches/` or a patched method under `afd_plugin/v1/worker/`, version gate, upstream copy | [compat-patches.md](../modules/compat-patches.md) | +| NPU runtime compat | `compat/npu/` deferred patch application, feature validation, routing ids, Ascend ops loading | [npu-compat.md](../modules/npu-compat.md) | +| Connectors | Connector backend, routed-experts transport, control plane, transfer/metadata handshake | [connectors.md](../modules/connectors.md) | | Distributed topology | AFD process groups, `AFDRankMapping`, topology validation | [distributed-topology.md](../modules/distributed-topology.md) | | Engine orchestration | Async-DP engine loops, cross-role routing, DBO yield op, ubatch orchestration | [engine-orchestration.md](../modules/engine-orchestration.md) | -| Workers & runners | Role workers, model runners V1/V2, graph capture, attention/FFN metadata | [workers-runners.md](../modules/workers-runners.md) | +| Workers & runners | Role workers, ModelRunnerV2, routed-experts and async FFN runners, graph capture, attention/FFN metadata | [workers-runners.md](../modules/workers-runners.md) | | Execution platforms | NPU/GPU worker selection, Ascend op build, `csrc/` kernels | [execution-platforms.md](../modules/execution-platforms.md) | | Model integration | Model registrations, role-split wrappers, weight-role filtering | [model-integration.md](../modules/model-integration.md) | +| E2E harness | `tests/e2e/**`, case IDs, gate selection, device layout, accuracy plumbing | [e2e-harness.md](../modules/e2e-harness.md) | Docs-, tests-, or CI-only changes route to the production contract they protect; if none applies, use only the applicable evidence checks. @@ -28,9 +29,10 @@ protect; if none applies, use only the applicable evidence checks. | Overlay | Signals | Read | | --- | --- | --- | -| Disaggregation topologies | 2A2F/2A1F, colocation vs disaggregation, `recipe/` launch scripts | [disaggregation-topologies.md](../features/disaggregation-topologies.md) | -| Async CAM & ubatching | `CAMAsyncAFDConnector`, two-stage MoE ubatching, DSV3.2/DSV4 async CAM forwards | [async-cam-ubatching.md](../features/async-cam-ubatching.md) | +| Disaggregation topologies | 1A1F/2A1F/2A2F/4A2F/4A4F/2P1A1F, colocation vs disaggregation, multi-pod, `recipe/` launch scripts | [disaggregation-topologies.md](../features/disaggregation-topologies.md) | +| Async CAM & layered FFN | `CAMAsyncAFDConnector`, layered W4A8/W8A8 GMM, FFN two-stream, DSV3.2/DSV4 async CAM forwards | [async-cam-ubatching.md](../features/async-cam-ubatching.md) | | Graphs & compilation | CUDA graph, ACL graph, capture/replay metadata, eager fallback | [graphs-and-compilation.md](../features/graphs-and-compilation.md) | +| DSpark draft | `AFDDSparkSpeculator`, `DSparkDraftModel`, target/draft layered FFN | [dspark-draft.md](../features/dspark-draft.md) | ## Evidence and change overlays @@ -42,18 +44,23 @@ protect; if none applies, use only the applicable evidence checks. | CI, docs, PR evidence, checklist | [tests-docs-checklist.md](../checks/tests-docs-checklist.md) | — | | Hardware or runnable path available | [verification.md](../checks/verification.md) | `run-e2e` | | vLLM/vLLM-Ascend pin moves | [compat-patches.md](../modules/compat-patches.md) + [execution-platforms.md](../modules/execution-platforms.md) | `upgrade-npu` / `upgrade-gpu-version` | +| Layered GMM ops pinned to a new ops-transformer revision | [async-cam-ubatching.md](../features/async-cam-ubatching.md) + [execution-platforms.md](../modules/execution-platforms.md) | `upgrade-layered-gmm-ops` | +| Recipe, deployment, or benchmark flow changes | [disaggregation-topologies.md](../features/disaggregation-topologies.md) | `deploy-afd-k8s` / `run-vllm-bench-k8s` / `deploy-dsv4-flash-async` | +| `.agents/skills/**` changes | [tests-docs-checklist.md](../checks/tests-docs-checklist.md) | — | ## Calibrate findings - **P0:** security exposure, data corruption, an unusable plugin or engine start, a patch that silently changes non-AFD upstream behavior. Block merge. - **P1:** reachable runtime failure, wrong output, platform divergence, - patch-contract violation that breaks upgrade comparability, missing evidence - for a blocking claim. Block merge until fixed or explicitly waived by an - owner. + patch-contract violation that breaks upgrade comparability, a `normative` + design invariant broken (`plugin_boundary.md`, `e2e_testing.md`), missing + evidence for a blocking claim. Block merge until fixed or explicitly waived + by an owner. - **P2:** real defect with a concrete future failure mode, contract drift the - code does not yet enforce, missing docs/test for changed behavior. - Non-blocking; do not inflate to block. + code does not yet enforce, missing docs/test for changed behavior, stale + agent-skill instructions that would mislead the next run. Non-blocking; do + not inflate to block. Everything else is a question for the author, not a finding. One root cause yields one finding even if it surfaces at several `path:line` sites. diff --git a/.agents/skills/run-e2e/SKILL.md b/.agents/skills/run-e2e/SKILL.md index 22bb0c03a..b98eea221 100644 --- a/.agents/skills/run-e2e/SKILL.md +++ b/.agents/skills/run-e2e/SKILL.md @@ -18,7 +18,6 @@ To split a scenario's ranks across multiple pods/nodes on a Kubernetes cluster instead of running single-process, see `resources/k8-multi-pod.md`. -Each suite contains four gate scenarios: DeepSeek-V2-Lite on GPU has seven gate scenarios: - `baseline-graph` @@ -154,7 +153,7 @@ For Qwen3 MoE on GPU, run: ~~~bash python -m pytest -q -s \ - tests/e2e/models/qwen3_moe/test_qwen3_moe.py + tests/e2e/models/qwen3_moe/test_qwen3_moe.py::test_qwen3_moe ~~~ For Qwen3.6 MoE on CUDA, run: @@ -192,8 +191,13 @@ On cancellation, forward SIGTERM and allow over 90 seconds for cleanup. Success means all selected cases pass with zero skips: seven for the GPU DeepSeek gate; four for the NPU DeepSeek gate or GPU Qwen3 MoE suite; three for the GPU Qwen3.6 MoE suite. -The entire DeepSeek module includes local cases and collects 12 on GPU or -seven on NPU. Report the actual selected count, failed scenario, first +The DeepSeek P2P entry collects 13 cases on GPU or seven on NPU. Its separate +DeepEP M2N entry adds five GPU-only IDs even to an unfiltered NPU collection; +select the intended entry or node IDs rather than running an unfiltered module. +GPU multi-pod entries now invoke the model module's `multi-pod` subcommand; +follow the linked Kubernetes walkthrough for local checkpoint paths, per-pod +acceptance records and the existing workload-only evidence scope. +Report the actual selected count, failed scenario, first actionable error, and cleanup status. Any skip is a gate failure. ## Environment reference diff --git a/.agents/skills/run-e2e/resources/k8-multi-pod.md b/.agents/skills/run-e2e/resources/k8-multi-pod.md index 596f5e308..976ea0620 100644 --- a/.agents/skills/run-e2e/resources/k8-multi-pod.md +++ b/.agents/skills/run-e2e/resources/k8-multi-pod.md @@ -5,9 +5,9 @@ layout splits Attention and FFN ranks across more than one pod — optionally spread across nodes, for a genuine cross-node fabric test — rather than running everything in one process on one machine. -This walkthrough assumes the model under test is -`deepseek-ai/DeepSeek-V2-Lite`; substitute `MODEL`, `MODEL_PVC`, and the -scenario's topology if you're running a different model. +This walkthrough uses the model-owned DeepSeek-V2-Lite GPU entry. Select the +appropriate model module, local checkpoint and PVC for another existing suite; +do not infer model selection from the shared coordinator. ## How it works @@ -19,8 +19,8 @@ Attention/FFN ranks to launch) from its Kubernetes-assigned completion index, then rendezvous with its peers over a shared store before serving and evaluating GSM8K. No process outside the pods holds test state or drives the run: you apply the Service and Job once, and every pod decides -its own role, coordinates with its peers, and reports its own pass/fail -through its container exit code. +its own role, coordinates with its peers, and reports execution through its +container exit code and acceptance through its run record. ## Prerequisites @@ -75,10 +75,9 @@ both FFN ranks; that string becomes the runner's `--pod-layout` argument below, and its entry count becomes `NUM_PODS`, computed automatically from `POD_LAYOUT` in the Deploy script — you don't set it yourself. -The ranks across all entries must sum to the scenario's topology (a -`2a2f` scenario needs 2 Attention and 2 FFN ranks in total), and no single -pod's rank count for a role may exceed that role's TP size — a TP group -cannot span pods. +The ranks across all entries must sum to the scenario's topology (a `2a2f` +scenario needs 2 Attention and 2 FFN ranks in total). Each pod's rank count +for a role must be divisible by that role's TP size: a TP group cannot span pods. This assumes every pod requests the same number of GPUs (`GPUS_PER_POD` below): the layout can vary how many Attention/FFN ranks each pod carries, @@ -99,7 +98,7 @@ SCENARIO=afd-graph-2a2f POD_LAYOUT=2A0F,0A2F NUM_PODS=$(($(tr -cd ',' <<<"$POD_LAYOUT" | wc -c) + 1)) # entry count of POD_LAYOUT RUN_ID=$(date +%s) -MODEL=deepseek-ai/DeepSeek-V2-Lite +MODEL=/models # local checkpoint root on the pre-warmed model PVC GSM8K_OUTPUT_PATH=/work/gsm8k-results MODEL_PVC=deepseek-v2-lite-weights GPUS_PER_POD=2 @@ -161,7 +160,8 @@ spec: args: - python - -m - - tests.e2e.multi_pod.runner + - tests.e2e.models.deepseek_v2_lite.cases + - multi-pod - --scenario - ${SCENARIO} - --pod-layout @@ -170,7 +170,9 @@ spec: - ${RUN_ID} - --model - ${MODEL} - - --gsm8k-output-path + - --device-backend + - gpu + - --output-path - ${GSM8K_OUTPUT_PATH} - --store-host - ${JOB_NAME}-0.${JOB_NAME} @@ -247,15 +249,26 @@ kubectl get pods -n ${NAMESPACE} -l job-name=${JOB_NAME} \ kubectl logs -n ${NAMESPACE} ``` -The run passed only if every pod's exit code is `0`; a pod with no terminal -state yet (still running when the wait timed out) is not a pass. +Acceptance requires every pod's exit code to be `0` and every pod's `run.json` +to have `acceptance_passed=true`. A pod with no terminal state yet is not a pass. +The model CLI allows debug workloads, so CLI exit code alone is insufficient. +The pytest entries enforce that requirement through `run_case_entry()`. +`GSM8K_OUTPUT_PATH` is the run root: role logs and `run.json` are written there, +with evaluator artifacts under `evaluation/`. + +The ten existing GPU multi-pod entries retain workload, liveness and cleanup +checks. They do not collect cross-pod graph/DBO/Hash execution evidence. Qwen3 +DeepEP M2N uses `tests.e2e.models.qwen3_moe.cases multi-pod`, its declared +role-split/interleaved layouts and 100-sample workload; its connector runtime must already +be available. See [execution and maintenance](../../../../tests/e2e/multi_pod/README.md). ## DSV4 Flash on Ascend NPU `afd-dsv4-flash-async-cam-dp2tp4-ep8` splits its fixed 16-NPU deployment (Attention DP2/TP4, FFN DP8/TP1/EP8, async CAM) across two 8-NPU pods. Its -pass criterion is the ten concurrent chat completions, not GSM8K. This path -has unit coverage only; it has not yet run on NPU hardware. +pass criterion is the ten concurrent chat completions, not GSM8K. This retained +NPU path requires separate hardware evidence; GPU/host validation does not +qualify it. Start from the Deploy template above and change: diff --git a/.agents/skills/run-vllm-bench-k8s/SKILL.md b/.agents/skills/run-vllm-bench-k8s/SKILL.md index e418a05fa..c52b400a0 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/.agents/skills/upgrade-gpu-version/SKILL.md b/.agents/skills/upgrade-gpu-version/SKILL.md index 129109f3b..853cccb49 100644 --- a/.agents/skills/upgrade-gpu-version/SKILL.md +++ b/.agents/skills/upgrade-gpu-version/SKILL.md @@ -175,7 +175,7 @@ and any detached read-only worktrees. Discover the current vLLM ref in this order: 1. `pyproject.toml`, lockfile, and package metadata; -2. compatibility constants, patch annotations, design matrices, README, GPU +2. compatibility constants, patch annotations, installation/support references, GPU guides, recipes, CI, and active upgrade RFC; 3. installed package versions and adjacent checkout state as corroboration only. diff --git a/.agents/skills/upgrade-gpu-version/references/20261028-v0.30.0-to-v0.32.0-lessons.md b/.agents/skills/upgrade-gpu-version/references/20261028-v0.30.0-to-v0.32.0-lessons.md index c9d9c2918..100971660 100644 --- a/.agents/skills/upgrade-gpu-version/references/20261028-v0.30.0-to-v0.32.0-lessons.md +++ b/.agents/skills/upgrade-gpu-version/references/20261028-v0.30.0-to-v0.32.0-lessons.md @@ -64,7 +64,7 @@ not a substitute for checking the next upstream version. ## Documentation and process lessons - Keep hardware, CPU-only, historical and excluded claims distinct. The - [runtime matrix](../../../../docs/design/module/execution_platforms.md#tested-runtime-matrix) + [runtime matrix](../../../../docs/support.md#tested-runtime-matrix) records the measured cells, score limits and shutdown caveats. Checked-in recipes were inspected but not all were executed as recipes. - Freeze source and runtime identity per test, preserve raw failed results, diff --git a/.agents/skills/upgrade-gpu-version/references/documentation-refresh.md b/.agents/skills/upgrade-gpu-version/references/documentation-refresh.md index 60c1c1744..33a29763f 100644 --- a/.agents/skills/upgrade-gpu-version/references/documentation-refresh.md +++ b/.agents/skills/upgrade-gpu-version/references/documentation-refresh.md @@ -8,7 +8,7 @@ must describe the exact target revision and the evidence actually collected. 1. Evidence ledger 2. Claim levels -3. Root documentation +3. Installation, support, and configuration 4. GPU connector and runtime guides 5. Recipes and launch scripts 6. Design and contributor documentation @@ -56,9 +56,9 @@ Do not turn CPU-only evidence into a GPU support claim. Do not describe an option as validated merely because it was accepted by the command line: confirm that logs, counters, traces, or execution behavior show the path was used. -## 3. Root Documentation +## 3. Installation, Support, and Configuration -Review the root `README.md` and any linked installation page for: +Review `docs/installation.md`, the [support matrix](../../../../docs/support.md#tested-runtime-matrix), and `docs/configuration.md` for: - the exact supported vLLM revision or range; - the target Python, PyTorch, CUDA, driver, compiler, and NCCL requirements; @@ -66,8 +66,7 @@ Review the root `README.md` and any linked installation page for: - installation commands and package extras; - supported models, connectors, parallel modes, and execution modes; - CUDA Graph, DBO or uBatch, TP, profiling, and isolation support; -- a minimal launch example with current arguments and environment variables; -- links to current recipes, tests, design notes, and troubleshooting pages; +- links from `README.md` and `docs/README.md` to current guides and recipes; - explicit unsupported or unvalidated combinations. Avoid broad phrases such as “GPU fully supported” when only selected models or @@ -77,7 +76,7 @@ record. ## 4. GPU Connector and Runtime Guides Review `afd_plugin/connectors/README.md` first. Keep its connector status matrix -consistent with the root README, GPU guide, recipes, design runtime matrix, and +consistent with `docs/support.md`, GPU guides, recipes, and E2E evidence. State eager, CUDA Graph, DBO/uBatch, TP, asymmetric topology, and profiling evidence separately where their support differs. @@ -126,7 +125,7 @@ appear related to the changed code: |---|---| | `docs/design/module/index.md` | target runtime, upstream refs, validation entry points, ownership, and links | | `docs/design/module/compatibility_and_patches.md` | exact patch symbols/signatures, reasons, tests, dispositions, and removal/upstream plans | -| `docs/design/module/execution_platforms.md` | GPU worker/runner initialization, device selection, CUDA Graph, DBO/uBatch, profiling, and runtime matrix | +| `docs/design/module/execution_platforms.md` | GPU worker/runner initialization, device selection, CUDA Graph, DBO/uBatch, and profiling | | `docs/design/module/connector_contracts.md` | NCCL/P2P configuration, payload/metadata ownership, ordering, topology, lifecycle, and cleanup | | `docs/design/module/attention_runtime.md` | Attention lifecycle, KV/scheduler ownership, forward context, connector handoff, graph, and ubatching | | `docs/design/module/ffn_runtime.md` | FFN daemon/EngineCore lifecycle, MoE execution, graph dispatch, failure propagation, and shutdown | @@ -213,7 +212,7 @@ Classify each hit as one of: - generated and must be regenerated; - unrelated to the GPU upgrade. -Then check that README, connector README, GPU guide, recipes, design modules, +Then check that installation/support references, connector README, GPU guides, recipes, design modules, tests, CI, package/build metadata, contributor templates, and release notes agree on the same target identity, native-build status, and support scope. diff --git a/.agents/skills/upgrade-npu/references/20261029-v0.30.0-to-v0.32.0-lessons.md b/.agents/skills/upgrade-npu/references/20261029-v0.30.0-to-v0.32.0-lessons.md index 39c04f9d3..6dd2643a3 100644 --- a/.agents/skills/upgrade-npu/references/20261029-v0.30.0-to-v0.32.0-lessons.md +++ b/.agents/skills/upgrade-npu/references/20261029-v0.30.0-to-v0.32.0-lessons.md @@ -9,7 +9,7 @@ ## Upstream changes that mattered -The architecture gate classified the release changes as local API/lifecycle adaptations rather than an authorized redesign. The shared engine startup/coordinator contract changed; native MRV2 execution, DP dispatch, Ascend configuration, graph metadata, SP ownership and model construction changed at their respective seams. AFD must preserve Attention request/KV ownership, connector-driven FFN lifecycle, control/transport ordering and platform separation. See [compatibility and patches](../../../../docs/design/module/compatibility_and_patches.md), [runtime matrix](../../../../docs/design/module/execution_platforms.md), and commits `d29db75`, `2e30083`, `49dc295`, `3eec8b4`, `d85159b` and `e6ced6c`. +The architecture gate classified the release changes as local API/lifecycle adaptations rather than an authorized redesign. The shared engine startup/coordinator contract changed; native MRV2 execution, DP dispatch, Ascend configuration, graph metadata, SP ownership and model construction changed at their respective seams. AFD must preserve Attention request/KV ownership, connector-driven FFN lifecycle, control/transport ordering and platform separation. See [compatibility and patches](../../../../docs/design/module/compatibility_and_patches.md), [runtime matrix](../../../../docs/support.md#tested-runtime-matrix), and commits `d29db75`, `2e30083`, `49dc295`, `3eec8b4`, `d85159b` and `e6ced6c`. ## Adaptations and patch lifecycle @@ -28,7 +28,7 @@ Check GPU fixes per backend before reusing them. NPU shares control-plane close/ ## Validation and environment lessons -The target `tests/unit` run recorded 1843 passed, six skipped, 180 subtests passed. Four device invocations had separate hardware coverage; two mutually exclusive absence cases skipped because the extension was present. The hardware and functional gates precede accuracy. The eight agreed Lite cases used identical first 100 GSM8K documents and prompts, eight-shot, max generation 512, strict match >= 0.27; all passed their request/path and accuracy gates. Seven AFD graceful exits failed, while native control's strict exit passed; physical cleanup and scheduler release were independently verified. See the [validation record](https://github.com/vllm-project/afd-plugin/pull/507) and [NPU tests](../../../../docs/npu/TESTING.md). +The target `tests/unit` run recorded 1843 passed, six skipped, 180 subtests passed. Four device invocations had separate hardware coverage; two mutually exclusive absence cases skipped because the extension was present. The hardware and functional gates precede accuracy. The eight agreed Lite cases used identical first 100 GSM8K documents and prompts, eight-shot, max generation 512, strict match >= 0.27; all passed their request/path and accuracy gates. Seven AFD graceful exits failed, while native control's strict exit passed; physical cleanup and scheduler release were independently verified. See the [validation record](https://github.com/vllm-project/afd-plugin/pull/507) and [NPU tests](../../../../tests/npu/README.md). The V4 W4A8 native/layered-off/layered-on/DBO first100 entries scored 96/95/94/95 for both strict and flexible filters, using the same documents, rendered prompts and targets, eight-shot, max generation 512 and concurrency 12. The eager Attention DBO entry used DP2/TP4 with FFN DP8/TP1 ordinary GMM; each DP0 Attention TP rank recorded 42 completed two-nonempty-microbatch steps and each DP1 TP rank recorded 44 during evaluation. This qualifies the agreed request-window cell, not per-question stage splitting, A/F transaction freshness, padding-only or other V4 boundary tests. V4 native DBO was outside the pinned Ascend runtime's support. These are the explicitly agreed 100-sample results, not 1,319-question qualification. diff --git a/.agents/skills/upgrade-npu/references/documentation-refresh.md b/.agents/skills/upgrade-npu/references/documentation-refresh.md index 4172e67a4..94faff293 100644 --- a/.agents/skills/upgrade-npu/references/documentation-refresh.md +++ b/.agents/skills/upgrade-npu/references/documentation-refresh.md @@ -51,21 +51,20 @@ backend supports the new vLLM revision. Distinguish these evidence levels: ## User-facing runtime documentation -### Root README +### Installation, support, and configuration -Review `README.md` as a complete public contract, including: +Review `docs/installation.md`, the [support matrix](../../../../docs/support.md#tested-runtime-matrix), and `docs/configuration.md` for: - target runtime wording and whether the version applies to NPU, GPU, or both; -- Current Status, supported models, connector matrix, and Known Gaps; +- supported models, connector matrix, and validation limits; - NPU hardware and full software-stack prerequisites; - vLLM and vLLM-Ascend installation pins and resolved source revisions; - build flags, `SOC_VERSION`, custom operator packages, and container guidance; - plugin installation and import or environment verification commands; -- NPU Attention and FFN launch commands, worker selection, and required - environment variables; +- worker selection and required environment variables; - AFD configuration keys, connector names and aliases, topology fields, graph, DBO, uBatch, PCP, quantization, and operator settings; -- links to the current connector guides, recipes, design docs, and E2E workflow. +- links from `README.md` and `docs/README.md` to current guides and recipes. Remove stale commands and unsupported examples. If a shortened launch example is not a complete runnable Attention/FFN deployment, label it as a fragment and @@ -74,7 +73,7 @@ link to the complete recipe. ### Connector package overview Review `afd_plugin/connectors/README.md`. Keep its connector status matrix -consistent with the root README, NPU guides, recipes, design runtime matrix, and +consistent with `docs/support.md`, NPU guides, recipes, and test evidence. State validation separately for CAMP2p and CAMAsync and for each relevant eager, graph, DBO, PCP, or topology cell. @@ -139,7 +138,7 @@ Refresh the internal source-of-truth documents when their contracts changed: | --- | --- | | `docs/design/module/index.md` | Runtime pair, upstream references, validation entry points, code ownership, and links | | `docs/design/module/compatibility_and_patches.md` | Exact target patch inventory, source symbols and signatures, patch reasons, tests, and upstream/removal plans | -| `docs/design/module/execution_platforms.md` | NPU Worker/ModelRunner initialization, ForwardContext, graph, DBO/uBatch, operator loading, and tested runtime matrix | +| `docs/design/module/execution_platforms.md` | NPU Worker/ModelRunner initialization, ForwardContext, graph, DBO/uBatch, and operator loading | | `docs/design/module/connector_contracts.md` | Configuration schema, payload/work-item ownership, control plane, topology, lifecycle, and execution order | | `docs/design/module/attention_runtime.md` | Attention worker/runner selection, metadata, forward context, connector handoff, graph, and ubatching | | `docs/design/module/ffn_runtime.md` | FFN daemon/EngineCore lifecycle, MoE execution, connector steps, graph dispatch, and shutdown | @@ -194,11 +193,11 @@ Compare every public support claim across these sources: | Claim | Sources that must agree | | --- | --- | -| Runtime versions | README, dependency pins, guides, recipes, design index, templates | -| Connector support | README, connector README, both NPU guides, recipe index, execution-platform matrix, E2E evidence | -| Models and quantization | README, recipes, model integration design, model E2E evidence | -| DBO/graph/PCP/topology | README, connector guides, recipes/scripts, execution and connector design docs, E2E evidence | -| Installation stack | README, guides, recipes, container/operator configuration, native-control evidence | +| Runtime versions | Installation and support references, dependency pins, guides, recipes, design index, templates | +| Connector support | Support reference, connector README, both NPU guides, recipe index, E2E evidence | +| Models and quantization | Support reference, recipes, model integration design, model E2E evidence | +| DBO/graph/PCP/topology | Support reference, connector guides, recipes/scripts, execution and connector design docs, E2E evidence | +| Installation stack | Installation reference, guides, recipes, container/operator configuration, native-control evidence | Search the whole documentation surface for stale versions and conflicting support language. Adapt these commands to the current and target refs: diff --git a/.github/ISSUE_TEMPLATE/500-ci-failure.yml b/.github/ISSUE_TEMPLATE/500-ci-failure.yml index 5f6737c46..2364621d5 100644 --- a/.github/ISSUE_TEMPLATE/500-ci-failure.yml +++ b/.github/ISSUE_TEMPLATE/500-ci-failure.yml @@ -1,7 +1,7 @@ name: CI failure report description: Report a failing CI job, flaky test, or validation regression. title: "[CI Failure]: " -labels: ["ci-failure"] +labels: ["CI-failure"] body: - type: input @@ -51,4 +51,3 @@ body: attributes: label: CC list description: People who should be aware of this failure. - diff --git a/.github/labeler.yml b/.github/labeler.yml new file mode 100644 index 000000000..b545cb044 --- /dev/null +++ b/.github/labeler.yml @@ -0,0 +1,74 @@ +# Path labels applied to pull requests by actions/labeler (v5). +# +# The workflow sets sync-labels: false, so this file only adds labels. +# A label that a person applied stays even when paths no longer match. +# Labels that are not keys in this file are never added or removed. +# +# Intentionally not auto-applied: +# ready, merge-test, weekly-test — Buildkite reads these CI triggers +# high priority — maintainer signal +# bug, feature request, RFC, CI-failure — GitHub Issue Forms +# +# There is no .buildkite/npu tree. Globs are case-sensitive, matching +# minimatch in actions/labeler@v5. dot: true in the workflow lets these +# patterns match .agents and .buildkite. + +Ascend: + - changed-files: + - any-glob-to-any-file: + # npu directories, plus names such as upgrade-npu and npu-compat.md. + - "**/npu/**" + - "**/*npu*/**" + - "**/*npu*" + - "**/*ascend*" + - "**/*ascend*/**" + # CAM and CAMP2P are Ascend connectors. Uppercase recipe and doc + # names live under npu/ and match the directory globs above. + - "**/*cam*" + - "**/*cam*/**" + # DSpark draft execution is implemented only on NPU. + - "**/*dspark*" + - "**/*dspark*/**" + - ".agents/skills/upgrade-layered-gmm-ops/**" + - ".agents/skills/deploy-dsv4-flash-async/**" + - "tests/e2e/test_deploy_dsv4_flash_async_skill.py" + - "tests/unit/v1/worker/test_async_ffn_*.py" + - "tests/unit/package/test_mega_activation_patch.py" + # DSV4 Flash E2E is mixed. These files are the NPU-only ones whose + # names do not contain npu, cam, ascend, or dspark. + - "tests/e2e/models/deepseek_v4_flash/config.py" + - "tests/e2e/models/deepseek_v4_flash/ffn_graph.py" + - "tests/e2e/models/deepseek_v4_flash/ffn_graph_evidence.py" + - "tests/e2e/models/deepseek_v4_flash/test_deepseek_v4_flash_multi_pod.py" + +NVIDIA: + - changed-files: + - any-glob-to-any-file: + # gpu directories, plus names such as upgrade-gpu-version and gpu_*.py. + - "**/gpu/**" + - "**/*gpu*/**" + - "**/*gpu*" + - "**/*cuda*" + - "**/*cuda*/**" + - "**/*deepep*" + - "**/*deepep*/**" + - ".buildkite/cuda/**" + # CUDA CI image (vllm-openai base). docker/ is not labeled as a whole. + - "docker/Dockerfile.ci" + # NCCL P2P unit tests. A broad *p2p* glob would also match CAMP2P + # and docs/npu/CAM_P2P, which are Ascend. + - "tests/unit/connectors/test_p2p_*.py" + # These E2E suites reject a non-GPU backend. + - "tests/e2e/models/qwen3_moe/**" + - "tests/e2e/models/qwen3_6/**" + - ".agents/skills/deploy-afd-k8s/**" + - ".agents/skills/run-vllm-bench-k8s/**" + +documentation: + - changed-files: + - any-glob-to-any-file: + # User-facing docs only. Module READMEs under code, tests, csrc, + # and .agents are left unlabeled so code PRs do not pick this up. + - "docs/**" + - "README.md" + - "recipe/**/*.md" diff --git a/.github/workflows/labeler.yml b/.github/workflows/labeler.yml new file mode 100644 index 000000000..daaed09e6 --- /dev/null +++ b/.github/workflows/labeler.yml @@ -0,0 +1,34 @@ +# SPDX-License-Identifier: Apache-2.0 +name: labeler + +# Fork pull_request events only receive a read-only token, so they cannot +# apply labels. pull_request_target runs this file from the base branch +# with a token that can label. Do not check out the pull request head and +# do not run code from it: untrusted PR content must not execute with this +# token. actions/labeler reads the changed-file list and .github/labeler.yml +# from the base repository through the GitHub API. +on: + pull_request_target: + types: [opened, synchronize, reopened, ready_for_review] + +permissions: + contents: read + +concurrency: + group: labeler-${{ github.event.pull_request.number }} + cancel-in-progress: true + +jobs: + labeler: + name: label pull request + runs-on: ubuntu-latest + permissions: + contents: read + pull-requests: write + steps: + - uses: actions/labeler@v5 + with: + repo-token: ${{ secrets.GITHUB_TOKEN }} + configuration-path: .github/labeler.yml + sync-labels: false + dot: true diff --git a/README.md b/README.md index 2cace4383..c5644bab3 100644 --- a/README.md +++ b/README.md @@ -3,7 +3,7 @@

-| Documentation | Recipe | Slides | DeepWiki | User Forum | Developer Slack | WeChat | +| Documentation | Recipe | Slides | DeepWiki | User Forum | Developer Slack | WeChat |

## Overview @@ -18,342 +18,46 @@ tests for GPU and Ascend NPU deployments. > This project is still experimental and needs more large-scale testing across > different hardware backends. -The GPU and Ascend target is **vLLM `v0.32.0`** at commit `5df63a72`. -The Ascend integration also targets vLLM-Ascend `bd42eea0`. The plugin does not modify the vLLM -source tree. AFD behavior is installed through the `vllm.general_plugins` entry -point, `--additional-config`, automatically selected role workers, plugin-owned -model wrappers, and narrow version-scoped compatibility shims. +The plugin extends vLLM through its general-plugin entry point without +modifying the vLLM source tree. ## Architecture ![afd-plugin architecture](docs/assets/vllm-afd-plugin-architecture.svg) -## Current Status +## Getting started -Core runtime support: +1. [Install AFD](docs/installation.md) for CUDA or Ascend NPU. +2. Choose a [connector and supported model](docs/support.md), then follow its + [deployment recipe](recipe/README.md). +3. Pass role/topology settings through + [`--additional-config`](docs/configuration.md). Send requests to Attention; + FFN serves connector work. -- vLLM plugin registration, AFD configuration, and runtime validation. -- Attention/FFN workers, model runners, model wrappers, and connector-driven - execution for CUDA and Ascend NPU. -- Eager and `FULL_DECODE_ONLY` graph execution, plus backend-specific profiling - support. -- Native DBO with exactly two ubatches on CUDA V2. Ascend constraints are - recorded in the existing [NPU runtime matrix](docs/design/module/execution_platforms.md#tested-runtime-matrix). -- DeepSeek MoE handoff at the remote-experts boundary on CUDA, with the gate - placed on either Attention or FFN. -- Expert-routed M:N exchange on CUDA (`DeepEPM2NAFDConnector`) for Qwen3 MoE: - Attention routes each token and dispatches it only to the FFN ranks that own - its experts. +## Support matrix -Model support: +See the [support matrix](docs/support.md) for models, connectors, tested +configurations, and known accuracy/shutdown limits. -| 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. | -| DeepSeek V4 Flash | `DeepseekV4ForCausalLM` | `AFDDeepseekV4ForCausalLM` | CUDA synchronous and Ascend CAM paths; see the [runtime matrix](docs/design/module/execution_platforms.md#tested-runtime-matrix) for backend-specific validation limits. | -| 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). | +## Documentation -Connector support: - -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 (SM90+) | Decode | Sync | `FULL_DECODE_ONLY` CUDA graph | Expert-routed M2N transport: Attention computes the gate and top-k and dispatches each token only to the FFN ranks that own its experts; the FFN runs only routed experts. Requires `compute_gate_on_attention=true`, FFN TP 1, and the same DBO setting on both roles; supports two-microbatch DBO, Attention TP, BF16 or block-128 FP8 payloads, and one or several nodes. Qwen3 MoE only. Dispatch and combine use [afdEP-comm](https://github.com/afdEP/afdEP-comm) (the `afdep` extra); the default expert backend needs [DeepGEMM](https://github.com/deepseek-ai/DeepGEMM) installed separately. See the [DeepEP M2N guide](docs/gpu/DEEPEP_M2N_CONNECTOR_USER_GUIDE.md) and the [Qwen3 recipes](recipe/gpu/DeepEPM2NAFDConnector/README.md). | -| `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. | - -Connector implementations are grouped by backend package: -`afd_plugin.connectors.gpu` for GPU-only connectors, -`afd_plugin.connectors.npu` for NPU-only connectors. - -Known gaps: - -- The GPU target is exactly vLLM `0.32.0`; other vLLM versions are not - claimed as supported on GPU. -- NPU hardware evidence covers the specific single-Pod DeepSeek-V2-Lite and - DSV4 configurations in the [0.32 validation record](https://github.com/vllm-project/afd-plugin/pull/507). - Other models and topologies remain unverified. AFD post-request shutdown - failures remain open; successful inference does not imply graceful shutdown. -- Synchronous DSV4 E2E profiles are smoke cases; A5 concurrent answers remain - incorrect. See the [profile limits](tests/e2e/README.md#dsv4-flash-sync-camp2p-concurrent-requests-local-4-or-16-npus). -- GPU Attention requires native ModelRunnerV2; there is no GPU V1 fallback. - FFN remains connector-driven. Synchronous `P2pNcclAFDConnector` supports - two-microbatch DBO and `FULL_DECODE_ONLY`, with model-specific gate placement. - See the [runtime matrix](docs/design/module/execution_platforms.md#tested-runtime-matrix) - for tested configurations and known DeepSeek-V2-Lite DBO shutdown failures. -- `DeepEPM2NAFDConnector` requires SM90 or newer GPUs, Attention-side - routing, and FFN TP 1, and supports only Qwen3 MoE models whose layers are - all MoE and have no shared expert. Its GPU device tests - ([`tests/gpu`](tests/gpu/README.md)) and E2E cases need SM90+ hardware and - run outside CI. -- GPU and NPU E2E tests are opt-in and require real hardware plus model weights. -- GPU CUDA graph support is limited to `FULL_DECODE_ONLY`. -- Native DBO is limited to exactly two ubatches. NPU CAMP2P supports eager - 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 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 covers Qwen3.6-35B-A3B. - GDN DBO is unsupported; upstream hybrid/multimodal DBO validation applies. - NPU and multimodal execution are unsupported/unverified; - `compute_gate_on_attention=true`, pipeline parallelism, asynchronous and - multi-node execution are unsupported; quantization is unverified; no - performance claim is made. - -## Install - -Requires Python **3.11-3.13** and [uv](https://docs.astral.sh/uv/). - -```bash -git clone https://github.com/vllm-project/afd-plugin.git -cd afd-plugin -uv sync --group dev -``` - -`vllm` is an optional runtime extra so CPU-only or macOS development -environments can still run import/config tests without a CUDA wheel. - -### GPU installation - -For Linux / CUDA-capable environments only; Ascend users should skip this -command: - -```bash -uv sync --group dev --extra vllm -``` - -The optional GPU extra pins `vllm==0.32.0`. - -`DeepEPM2NAFDConnector` also needs the `afdep` extra (afdEP-comm) and -DeepGEMM, which is built from source; see -[Installation](docs/gpu/DEEPEP_M2N_CONNECTOR_USER_GUIDE.md#installation) in -the connector guide. - -```bash -uv sync --group dev --extra vllm --extra afdep -``` - -### Ascend NPU installation - -The Ascend integration uses the source pair below. See the -[0.32 NPU validation record](https://github.com/vllm-project/afd-plugin/pull/507) for the tested -scope and retained shutdown limitations. Earlier -[0.30 results](https://github.com/vllm-project/afd-plugin/pull/425#issuecomment-6063910923) -retain their historical environment and sample counts. -Install a compatible driver and firmware, and confirm the devices with -`npu-smi info`. - -| Component | Version | -| --- | --- | -| Python | Package range `3.11–3.13`; tested `3.12.13` | -| vLLM | `0.32.0`, commit `5df63a7203cdb9328d824a27b682e77c7b9928be` | -| vLLM-Ascend | commit [`bd42eea0`](https://github.com/vllm-project/vllm-ascend/commit/bd42eea01d8312c4df469fabaed45f77a12df4d3) | -| Tested CANN / torch / torch-npu | `9.2` / `2.10.0+cpu` / `2.10.0.post7` | -| Tested Triton Ascend / CAM | `3.6` / `1.3` | - -> Hardware evidence is scoped to the recorded models, configurations, and -> sample counts; it does not cover every NPU deployment or guarantee clean -> device shutdown. - -#### Environment - -The recorded A3 environment uses the pinned pair and stack above. The -repository does not claim a released container tag for the Ascend target. Use the -[installation guide at that source snapshot](https://github.com/vllm-project/vllm-ascend/blob/bd42eea01d8312c4df469fabaed45f77a12df4d3/docs/source/getting_started/installation.md) -to prepare a matching A3/openEuler environment, then install AFD from the -repository root. - -#### Install AFD - -From the AFD repository root: - -```bash -AFD_BUILD_ASCEND_OPS=1 \ -SOC_VERSION=ascend910_9391 \ -python -m pip install -v --no-build-isolation --no-deps -e . -``` - -`--no-deps` preserves the matched NPU runtime, and `--no-build-isolation` uses -its CANN/torch-npu toolchain. The command above forces the Ascend op build for -reproducibility. Otherwise, leave `AFD_BUILD_ASCEND_OPS` unset to use the -auto-detection described in [Development](#development), set it to `1` if -detection misses an Ascend environment, or set it to `0` to skip the ops. - -#### Verify - -```bash -python - <<'PY' -import torch -import torch_npu -import vllm_ascend - -from afd_plugin.compat.npu import ensure_afd_ascend_ops_loaded - -assert torch.npu.is_available(), "torch-npu cannot see an Ascend device" -ensure_afd_ascend_ops_loaded() -print("AFD_OPS_OK") -PY -``` - -After `AFD_OPS_OK`, the environment is ready to run the NPU examples and E2E -tests. See the -[NPU testing guide](docs/npu/TESTING.md). -For implementation details, see the -[Attention runtime design](docs/design/module/attention_runtime.md) and -[FFN runtime design](docs/design/module/ffn_runtime.md). - -## Using the Plugin - -Install or sync the distribution as `vllm-afd-plugin`. Python imports use the -`afd_plugin` package name. - -AFD is configured through vLLM `--additional-config`. There is no separate -`--afd-config` flag. When AFD is configured and `--worker-cls` is omitted, the -plugin automatically selects the Attention or FFN worker for the active CUDA -or standard Ascend NPU platform. Explicit AFD worker paths remain accepted for -compatibility with existing commands, but are not required or stable launch -interfaces. - -GPU and NPU launches require `VLLM_USE_V2_MODEL_RUNNER=1` on both roles. -Attention uses native ModelRunnerV2; FFN keeps its connector-driven runner. -Neither backend has a V1 fallback. GPU DBO requires two microbatches and accepts Attention -DP1 or DP > 1. -See the [dedicated E2E scenarios](tests/e2e/README.md#gpu-mrv2-dbo-development-comparison). - -```bash -export VLLM_USE_V2_MODEL_RUNNER=1 -``` - -GPU Attention-side shape: - -```bash -vllm serve /path/to/DeepSeek-V2-Lite \ - --served-model-name deepseek-v2-lite-afd-attention \ - --data-parallel-size 1 \ - --tensor-parallel-size 1 \ - --enable-expert-parallel \ - --enforce-eager \ - --host 127.0.0.1 \ - --port 18000 \ - --additional-config '{"afd":{"role":"attention","connector":"P2pNcclAFDConnector","host":"127.0.0.1","port":6239,"num_attention_ranks":1,"num_ffn_ranks":1}}' -``` - -GPU FFN-side shape: - -```bash -vllm serve /path/to/DeepSeek-V2-Lite \ - --served-model-name deepseek-v2-lite-afd-ffn \ - --data-parallel-size 1 \ - --tensor-parallel-size 1 \ - --enable-expert-parallel \ - --enforce-eager \ - --host 127.0.0.1 \ - --port 18001 \ - --additional-config '{"afd":{"role":"ffn","connector":"P2pNcclAFDConnector","host":"127.0.0.1","port":6239,"num_attention_ranks":1,"num_ffn_ranks":1}}' -``` - -NPU uses the same config channel with `CAMP2pAFDConnector`; the plugin selects -the NPU worker automatically. Attention uses native Ascend ModelRunnerV2: - -```bash -VLLM_USE_V2_MODEL_RUNNER=1 vllm serve /path/to/DeepSeek-V2-Lite \ - --served-model-name deepseek-v2-lite-afd-attention \ - --data-parallel-size 1 \ - --tensor-parallel-size 1 \ - --enable-expert-parallel \ - --enforce-eager \ - --host 127.0.0.1 \ - --port 18000 \ - --additional-config '{"afd":{"role":"attention","connector":"CAMP2pAFDConnector","host":"127.0.0.1","port":6239,"num_attention_ranks":1,"num_ffn_ranks":1}}' -``` - -Attention and FFN may be started in either order. Send requests only to the -Attention API server. FFN workers are connector-driven; scheduler-driven FFN -`execute_model()` calls fail fast. - -For repeatable local smoke testing, prefer the bundled runner: - -```bash -export HF_HOME=/path/to/huggingface -python -c 'from datasets import load_dataset; load_dataset("openai/gsm8k", "main")' - -uv run python tests/e2e/runner.py \ - --model /path/to/DeepSeek-V2-Lite \ - --device-backend gpu \ - --scenario afd-eager-2a1f \ - --attention-devices 0,1 \ - --ffn-devices 2 \ - --gsm8k-output-path /tmp/afd-e2e-gsm8k \ - --api-port-base 18000 \ - --afd-port 6239 \ - --common-vllm-arg=--trust-remote-code -``` - -For NPU, use `--device-backend npu`; the runner maps the same device arguments -to `ASCEND_RT_VISIBLE_DEVICES` and selects `CAMP2pAFDConnector`. - -## AFD Config - -The canonical config shape is: - -```json -{ - "afd": { - "role": "attention", - "connector": "P2pNcclAFDConnector", - "host": "127.0.0.1", - "port": 1239, - "num_attention_ranks": 2, - "num_ffn_ranks": 1, - "compute_gate_on_attention": false, - "connector_extra_config": {} - } -} -``` - -`role` must be `attention` or `ffn`. `connector` must be `P2pNcclAFDConnector`, -`DeepEPM2NAFDConnector`, `CAMP2pAFDConnector`, or `CAMAsyncAFDConnector`. -AFD is active when `additional_config["afd"]` is present and passes common AFD config validation; -omit `additional_config["afd"]` to disable AFD. Connector-owned -`connector_extra_config` is strictly validated by the selected connector parser -when the connector/runtime is constructed. The plugin also accepts selected -compatibility aliases such as `afd_role`, `afd_connector`, `afd_host`, and -`afd_port`. +Browse the [documentation index](docs/README.md) for user guides, module design, +validation and troubleshooting. ## Development -Before adding, changing, moving, or removing tests, or changing the behavior they -cover, read and follow the [test maintenance guide](docs/UT_MAINTENANCE.md). - -Run the default CPU-safe checks: +For CPU-only or macOS development, follow [source setup](docs/installation.md#source-setup), +then run: ```bash uv run pytest uv run ruff check . ``` -The NPU-dependent unit suite must run in an Ascend development environment; -see [Running NPU unit tests](docs/npu/TESTING.md). - -Native C/C++ sources are grouped by backend under `csrc/`: Ascend/CANN sources -live in `csrc/npu`, including the `a2e` and `e2a` ACLNN operators, and -`csrc/gpu` is reserved for GPU native sources. - -Ascend custom ops are built automatically only when the build environment looks -like an Ascend NPU platform, for example when `torch_npu`, CANN environment -variables, or the default Ascend toolkit path are present. GPU builds skip -Ascend ops by default. Set `AFD_BUILD_ASCEND_OPS=1` or -`AFD_BUILD_ASCEND_OPS=0` to override the auto-detection. - -## E2E Test - -To run E2E tests, use the [`run-e2e` skill](.agents/skills/run-e2e/SKILL.md). +Read the [contributor rules](AGENTS.md) and [test maintenance guide](docs/UT_MAINTENANCE.md) +before changing code or tests. See [module design](docs/design/module/index.md) +for ownership, [native builds](csrc/README.md) for backend operators, and the +[E2E guide](tests/e2e/README.md) for hardware validation. ## License diff --git a/afd_plugin/__init__.py b/afd_plugin/__init__.py index eeff9c467..a586088c5 100644 --- a/afd_plugin/__init__.py +++ b/afd_plugin/__init__.py @@ -125,6 +125,9 @@ def get_spawn_context(method: str | None = None): "Qwen3MoeForCausalLM": ( "afd_plugin.model_executor.models.qwen3_moe:AFDQwen3MoeForCausalLM" ), + "Qwen3MoeRoutedExperts": ( + "afd_plugin.model_executor.models.qwen3_moe:AFDQwen3MoeRoutedExperts" + ), } _QWEN3_5_MODEL_REGISTRATIONS = MappingProxyType( diff --git a/afd_plugin/compat/npu/feature_validation.py b/afd_plugin/compat/npu/feature_validation.py index cf151fac1..80f8de270 100644 --- a/afd_plugin/compat/npu/feature_validation.py +++ b/afd_plugin/compat/npu/feature_validation.py @@ -14,7 +14,11 @@ is_afd_async_dp, parse_afd_config, ) -from afd_plugin.validation import cam_receive_capacity, requested_ffn_mode +from afd_plugin.validation import ( + cam_receive_capacity, + requested_ffn_mode, + validate_async_cam_attention_graph, +) if TYPE_CHECKING: from vllm.config import VllmConfig @@ -56,8 +60,15 @@ def fail_if_unsupported_npu_afd_features( return if afd_config.compute_gate_on_attention: - raise RuntimeError( - "AFD NPU runtime does not support compute_gate_on_attention=true yet", + # The sync Attention-side gate prototype is deliberately narrow: one + # connector, the target model and BF16 transport. Widening it without the + # corresponding transport work would produce wrong numerics instead of an + # error, so anything outside that shape stays rejected. + _fail_if_unsupported_sync_attention_gate( + vllm_config, + afd_config, + extra_info, + is_dsv4=is_dsv4, ) if afd_config.connector == CAMP2P_CONNECTOR: from afd_plugin.connectors.npu.camp2p import CAMP2PExtraInfo @@ -151,6 +162,53 @@ def _fail_if_unsupported_dsv4_connector(afd_config: AFDConfig) -> None: ) +def _fail_if_unsupported_sync_attention_gate( + vllm_config: VllmConfig, + afd_config: AFDConfig, + extra_info: object, + *, + is_dsv4: bool, +) -> None: + """Restrict the sync Attention-side gate prototype to what it supports today. + + DBO and ubatching are rejected rather than merely unused: the sync shell does not + call the DBO yield, so enabling DBO here would interleave microbatches against a + transport that never yields. Not calling the yield is not the same as forbidding + the feature. + """ + + parallel = vllm_config.parallel_config + # The routing travels on each rank's own activation, and only the TP=1 shape has + # been verified end to end; a wider TP would split the activation without a + # matching routing split on the Attention side. + if int(parallel.tensor_parallel_size) != 1: + raise RuntimeError( + "Attention-side gate is only supported at tensor-parallel size 1; got " + f"TP={int(parallel.tensor_parallel_size)}", + ) + if bool(parallel.enable_dbo) or bool(parallel.use_ubatching): + raise RuntimeError( + "Attention-side gate with the synchronous connector does not support " + "DBO or ubatching yet", + ) + if afd_config.connector != CAMP2P_CONNECTOR: + raise RuntimeError( + "Attention-side gate is only supported with " + f"{CAMP2P_CONNECTOR}, got {afd_config.connector!r}", + ) + if not is_dsv4: + raise RuntimeError( + "Attention-side gate is only supported for the DeepSeek V4 target", + ) + from afd_plugin.connectors.npu.camp2p import CAMP2PExtraInfo + + if not isinstance(extra_info, CAMP2PExtraInfo): + raise TypeError( + "CAMP2pAFDConnector requires CAMP2PExtraInfo, got " + f"{type(extra_info).__name__}", + ) + + def _fail_if_unsupported_npu_afd_async_features( vllm_config: VllmConfig, afd_config: AFDConfig, @@ -170,22 +228,6 @@ def _fail_if_unsupported_npu_afd_async_features( ) parallel_config = vllm_config.parallel_config - if extra_info.ffn_two_stream: - if not afd_config.is_ffn_server: - raise RuntimeError("ffn_two_stream may only be enabled on the FFN role") - if ( - not vllm_config.use_v2_model_runner - or not afd_config.compute_gate_on_attention - or extra_info.dynamic_quant != 1 - or parallel_config.tensor_parallel_size != 1 - or parallel_config.use_sequence_parallel_moe - or os.environ.get("VLLM_ASCEND_ENABLE_FLASHCOMM1", "0").strip().lower() - in ("1", "true") - ): - raise RuntimeError( - "FFN two-stream requires MRV2 layered GMM, dynamicQuant=1, " - "Attention-side gate, F TP=1 and SP disabled" - ) if not is_afd_async_dp(vllm_config): raise RuntimeError( "CAMAsyncAFDConnector requires additional_config['afd'] " @@ -193,14 +235,13 @@ def _fail_if_unsupported_npu_afd_async_features( ) if not bool(vllm_config.model_config.enforce_eager): if afd_config.is_attention_server: - raise RuntimeError( - "CAMAsyncAFDConnector supports only eager Attention execution" - ) - requested_ffn_mode(vllm_config) - if not vllm_config.use_v2_model_runner: - raise RuntimeError("Async FFN FULL requires MRV2") - if parallel_config.tensor_parallel_size != 1: - raise RuntimeError("Async FFN FULL requires F TP=1") + validate_async_cam_attention_graph(vllm_config) + else: + requested_ffn_mode(vllm_config) + if not vllm_config.use_v2_model_runner: + raise RuntimeError("Async FFN FULL requires MRV2") + if parallel_config.tensor_parallel_size != 1: + raise RuntimeError("Async FFN FULL requires F TP=1") if bool(parallel_config.enable_dbo) or bool(parallel_config.use_ubatching): if not afd_config.is_attention_server: raise RuntimeError("Async CAM DBO/ubatching requires the Attention role") @@ -222,30 +263,34 @@ def _fail_if_unsupported_npu_afd_async_features( if speculative_config.method == "dspark": if vllm_config.model_config.hf_text_config.model_type != "deepseek_v4": raise RuntimeError("AFD DSpark supports only DeepSeek-V4 on NPU") - if extra_info.ffn_two_stream: - raise RuntimeError("AFD DSpark does not support FFN two-stream") # The shared MRV2 validator owns PP/CP and static EP constraints. # Keep only the additional target/draft restrictions here. - # F may replay its expert transaction in FULL; Attention and native - # draft proposal execution stay eager. + # Both graphs use the native decode path. Draft capture is owned + # by native target capture, so draft-only graph mode is unsupported. + # F keeps its independent TP1 expert execution. if ( - ( - afd_config.is_attention_server - and not vllm_config.model_config.enforce_eager + speculative_config.draft_parallel_config.tensor_parallel_size + != parallel_config.tensor_parallel_size + or ( + afd_config.is_ffn_server + and parallel_config.tensor_parallel_size != 1 ) - or not speculative_config.enforce_eager - or parallel_config.tensor_parallel_size != 1 - or speculative_config.draft_parallel_config.tensor_parallel_size != 1 ): raise RuntimeError( - "AFD DSpark requires eager Attention/draft and target/draft TP=1" + "AFD DSpark requires matching target/draft TP and F TP=1" ) + if ( + afd_config.is_attention_server + and not speculative_config.enforce_eager + and vllm_config.model_config.enforce_eager + ): + raise RuntimeError("DSpark draft Graph requires target Attention graph") elif afd_config.is_attention_server and ( speculative_config.use_eagle() or speculative_config.uses_draft_model() ): raise RuntimeError( "CAMAsyncAFDConnector does not support native draft models or MTP; " - "only DSV4 DSpark eager is supported", + "only DSV4 DSpark is supported", ) factor_value = os.environ.get("BATCH_SIZE_FACTOR", "1") @@ -273,9 +318,10 @@ def _fail_if_unsupported_npu_afd_async_features( "AFD DSpark draft query capacity exceeds max_num_batched_tokens: " f"{draft_rows} > {max_num_batched_tokens}" ) - # whole-expert chunks need a full TP group's batch; smaller - # buffers require native sub-expert chunking. - if receive_capacity < max_num_batched_tokens: + if receive_capacity < 1: + raise RuntimeError("BATCH_SIZE_FACTOR must allocate at least one receive row") + # Push chunks whole experts. Pull can split one expert across receive rounds. + if not extra_info.ffn_mega_moe and receive_capacity < max_num_batched_tokens: raise RuntimeError( "CAMAsyncAFDConnector requires BATCH_SIZE_FACTOR to allocate at least " f"max_num_batched_tokens={max_num_batched_tokens} " diff --git a/afd_plugin/connectors/README.md b/afd_plugin/connectors/README.md index e77e20ede..2c5ed3bb8 100644 --- a/afd_plugin/connectors/README.md +++ b/afd_plugin/connectors/README.md @@ -23,7 +23,7 @@ AFD connector implementations are grouped by backend: Both backends target vLLM `0.32.0`; Ascend additionally targets vLLM-Ascend `bd42eea0`. See the -[runtime matrix](../../docs/design/module/execution_platforms.md#tested-runtime-matrix) +[runtime matrix](../../docs/support.md#tested-runtime-matrix) for tested configurations and backend-specific validation limits. Shared connector contracts, metadata containers, factory registration, and diff --git a/afd_plugin/connectors/__init__.py b/afd_plugin/connectors/__init__.py index 2b270cc6e..00384908f 100644 --- a/afd_plugin/connectors/__init__.py +++ b/afd_plugin/connectors/__init__.py @@ -2,27 +2,49 @@ # SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project """AFD connector namespace.""" -from afd_plugin.connectors.base import ( - AFDConnectorBase, - AFDControlPlane, - AFDRoutedExpertsConnector, - ConnectorExtraInfo, -) -from afd_plugin.connectors.factory import AFDConnectorFactory -from afd_plugin.connectors.metadata import ( - AFDA2FTransferPayload, - AFDControlPayload, - AFDDPMetadata, - AFDExpertRoutingSpec, - AFDF2ATransferPayload, - AFDForwardContextMetadata, - AFDSingleDPMetadata, - AFDTransferContext, - AFDTransferMetadata, - AFDTransferState, - FFNLocalExperts, - LocalExpertWeights, -) +from importlib import import_module +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from afd_plugin.connectors.base import ( + AFDConnectorBase, + AFDControlPlane, + AFDRoutedExpertsConnector, + ConnectorExtraInfo, + ) + from afd_plugin.connectors.factory import AFDConnectorFactory + from afd_plugin.connectors.metadata import ( + AFDA2FTransferPayload, + AFDControlPayload, + AFDDPMetadata, + AFDExpertRoutingSpec, + AFDF2ATransferPayload, + AFDForwardContextMetadata, + AFDSingleDPMetadata, + AFDTransferContext, + AFDTransferMetadata, + AFDTransferState, + FFNLocalExperts, + LocalExpertWeights, + ) + + +def __getattr__(name: str): + if name in { + "AFDConnectorBase", + "AFDControlPlane", + "AFDRoutedExpertsConnector", + "ConnectorExtraInfo", + }: + module = "base" + elif name == "AFDConnectorFactory": + module = "factory" + elif name in __all__: + module = "metadata" + else: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + return getattr(import_module(f"{__name__}.{module}"), name) + __all__ = [ "AFDRoutedExpertsConnector", diff --git a/afd_plugin/connectors/base.py b/afd_plugin/connectors/base.py index e0fd3528d..4bf86a7e2 100644 --- a/afd_plugin/connectors/base.py +++ b/afd_plugin/connectors/base.py @@ -5,24 +5,25 @@ from __future__ import annotations from abc import ABC, abstractmethod -from collections.abc import Mapping, Sequence +from collections.abc import Callable, Mapping, Sequence +from contextlib import AbstractContextManager, nullcontext from dataclasses import dataclass from typing import TYPE_CHECKING, Any -import torch - from afd_plugin.config import AFDConfig, connector_extra_config_from_source -from afd_plugin.connectors.metadata import ( - AFDA2FTransferPayload, - AFDControlPayload, - AFDTransferContext, - FFNLocalExperts, - LocalExpertWeights, -) if TYPE_CHECKING: + import torch from vllm.config import VllmConfig + from afd_plugin.connectors.metadata import ( + AFDA2FTransferPayload, + AFDControlPayload, + AFDTransferContext, + FFNLocalExperts, + LocalExpertWeights, + ) + @dataclass(frozen=True) class ConnectorExtraInfo: @@ -36,6 +37,37 @@ def to_mapping(self) -> dict[str, Any]: return {} +class FFNExecutionSession: + """Backend lifecycle used by FFN runners and workers. + + By default, serving stops on local shutdown and the remaining hooks do + nothing. Backends with remote producers override the hooks and set + ``requires_drain``; pending work and the receiver handshake must finish + before communication or model resources are released. + """ + + requires_drain = False + + def keep_serving(self, shutdown_requested: bool) -> bool: + """Continue until local shutdown; remote producers may extend serving.""" + return not shutdown_requested + + def before_first_step(self) -> None: + """Prepare backend state once before the first FFN step.""" + + def after_step(self) -> None: + """Check backend status after a completed FFN step.""" + + def drain(self, run_one: Callable[[], None]) -> None: + """Drain remote work when overridden; by default do not invoke run_one.""" + + def finish_receiver(self) -> None: + """Finish the receiver handshake after pending work has drained.""" + + def validate_shutdown(self) -> None: + """Reject release when backend resources are still in use.""" + + class AFDConnectorBase(ABC): """Base class for plugin-owned AFD Attention/FFN connectors. @@ -112,6 +144,17 @@ def __init__( # Lifecycle methods # ============================== + def attention_step(self) -> AbstractContextManager[None]: + """Return a context for ``with connector.attention_step():`` around forward. + + Backends may defer transport status checks until the Attention model + step completes. The default context performs no additional work. + """ + return nullcontext() + + def ffn_session(self) -> FFNExecutionSession: + return FFNExecutionSession() + @abstractmethod def close(self) -> None: """Release connector-owned communication resources. @@ -378,4 +421,5 @@ def recv_dp_metadata_list(self) -> AFDControlPayload: "AFDControlPlaneClosedError", "AFDRoutedExpertsConnector", "ConnectorExtraInfo", + "FFNExecutionSession", ] diff --git a/afd_plugin/connectors/factory.py b/afd_plugin/connectors/factory.py index 4cc19400d..20b0a08c7 100644 --- a/afd_plugin/connectors/factory.py +++ b/afd_plugin/connectors/factory.py @@ -37,7 +37,10 @@ def register_connector( def loader() -> type[AFDConnectorBase]: module = importlib.import_module(module_path) - connector_cls = vars(module)[class_name] + # ``getattr`` rather than ``vars(module)``: packages in this tree resolve + # their names through a module-level ``__getattr__``, so the class is not + # necessarily present in the module dictionary. + connector_cls = getattr(module, class_name) if not issubclass(connector_cls, AFDConnectorBase): raise TypeError( f"{module_path}.{class_name} is not an AFDConnectorBase", diff --git a/afd_plugin/connectors/gpu/deepep_m2n/connector.py b/afd_plugin/connectors/gpu/deepep_m2n/connector.py index c0d2fa917..4eb11dca3 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: @@ -1143,6 +1203,12 @@ def close(self) -> None: _CONNECTORS.pop(self._op_id, None) for lane in self.lanes: lane.destroy() + if ( + self.transport_group is not None + and self.transport_group is not self.bootstrap_group + ): + dist.destroy_process_group(self.transport_group) + self.transport_group = None self.lanes = [] self.lane_streams = [] self._pending_combines.clear() @@ -1157,8 +1223,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/afd_plugin/connectors/metadata.py b/afd_plugin/connectors/metadata.py index 02351e7c8..4b5cf6bc9 100644 --- a/afd_plugin/connectors/metadata.py +++ b/afd_plugin/connectors/metadata.py @@ -288,12 +288,20 @@ class AFDA2FTransferPayload: router_logits: Optional Attention-side routing tensor. input_ids: Optional token-aligned identifiers required by the remote FFN implementation. + topk_ids: Optional expert indices selected by an Attention-side gate, + shaped ``(tokens, topk)``. Present instead of ``input_ids`` when the + Attention role transported its gate result over the ids channel. + topk_weights: Optional routing weights paired with ``topk_ids``, same + shape. The FFN side consumes them as the gate weights, so any + model-level routed scaling is already folded in. """ hidden_states: torch.Tensor context: AFDTransferContext router_logits: torch.Tensor | None = None input_ids: torch.Tensor | None = None + topk_ids: torch.Tensor | None = None + topk_weights: torch.Tensor | None = None @dataclass(slots=True) diff --git a/afd_plugin/connectors/npu/__init__.py b/afd_plugin/connectors/npu/__init__.py index 14e61344e..8abb26e20 100644 --- a/afd_plugin/connectors/npu/__init__.py +++ b/afd_plugin/connectors/npu/__init__.py @@ -2,12 +2,24 @@ # SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project """NPU-specific AFD connector implementations.""" -from afd_plugin.connectors.npu.camp2p import ( - CAMP2pAFDConnector, - CAMP2PExtraInfo, - CAMP2PTransferState, - build_camp2p_topology, -) +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from afd_plugin.connectors.npu.camp2p import ( + CAMP2pAFDConnector, + CAMP2PExtraInfo, + CAMP2PTransferState, + build_camp2p_topology, + ) + + +def __getattr__(name: str): + if name in __all__: + from afd_plugin.connectors.npu import camp2p + + return getattr(camp2p, name) + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + __all__ = [ "CAMP2pAFDConnector", diff --git a/afd_plugin/connectors/npu/async_cam.py b/afd_plugin/connectors/npu/async_cam.py index f981fe153..60a504189 100644 --- a/afd_plugin/connectors/npu/async_cam.py +++ b/afd_plugin/connectors/npu/async_cam.py @@ -1,28 +1,15 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project -"""Ascend CAM asynchronous connector for Attention/FFN disaggregation. - -``CAMAsyncAFDConnector`` supports eager Attention and eager or FULL FFN. -Attention ranks run MoE routing, submit activations with CAM async -dispatch-send, and receive combined expert output with combine-recv. FFN ranks -receive routed-expert activations with dispatch-recv, execute their -local experts, and return the results with combine-send. - -The connector creates one HCCL world ordered as -``[A0, A1, ..., F0, F1, ...]``. Attention world ranks therefore equal their -role ranks, while FFN world ranks start at ``num_attention_ranks``. Unlike the -synchronous connectors, CAM async carries routing and token-count metadata in -the CAM operator payload and does not create a separate Gloo DP-metadata -control plane. - -The supported deployment requires ``async=true``, Ascend CAM -operator packages, and matching topology/configuration on every rank. Regular -prefill and autoregressive decode steps use the same connector. Attention -microbatching uses the native DBO runtime with Ascend adaptations. FFN FULL -requires MRV2 and accepts Attention-side DBO. Native DBO owns microbatch -splitting and scheduling. See ``docs/npu/CAM_ASYNC_CONNECTOR_USER_GUIDE.md`` for -configuration, rank -derivation, launch guidance, and the full limitations. +"""Ascend CAM connector for Attention/FFN disaggregation. + +The connector validates payloads and manages the Attention-first HCCL world +``[A0, A1, ..., F0, F1, ...]``. Push/Pull sessions own transfer state and +protocol rules; the FFN ModelRunner owns eager/FULL execution. Attention +microbatching uses native DBO. Native operators carry routing/count metadata +without a separate Gloo control plane. + +See ``docs/npu/CAM_ASYNC_CONNECTOR_USER_GUIDE.md`` for configuration, rank +derivation and supported paths. """ from __future__ import annotations @@ -30,6 +17,7 @@ import os import time from collections.abc import Mapping +from contextlib import AbstractContextManager from dataclasses import dataclass from datetime import timedelta from typing import TYPE_CHECKING, Any, Final, cast @@ -38,11 +26,6 @@ from torch import Tensor from vllm.logger import init_logger -from afd_plugin.async_cam_startup import ( - GRAPH_STARTUP_TIMEOUT_SECONDS, - POLL_SECONDS, - AsyncCAMGraphStartup, -) from afd_plugin.compat.npu.ops import ensure_cam_async_ops_available from afd_plugin.config import AFDConfig from afd_plugin.config_utils import ( @@ -54,6 +37,7 @@ from afd_plugin.connectors.base import ( AFDConnectorBase, ConnectorExtraInfo, + FFNExecutionSession, ) from afd_plugin.connectors.metadata import ( AFDA2FTransferPayload, @@ -61,6 +45,16 @@ AFDTransferMetadata, AFDTransferState, ) +from afd_plugin.connectors.npu.async_cam_runtime.async_cam_session import ( + CAM_COMM_ID, + PullCAMSession, + PushCAMSession, +) +from afd_plugin.connectors.npu.async_cam_runtime.async_cam_startup import ( + GRAPH_STARTUP_TIMEOUT_SECONDS, + POLL_SECONDS, + AsyncCAMGraphStartup, +) from afd_plugin.distributed import ( create_hccl_process_group_options, init_afd_process_group, @@ -73,24 +67,20 @@ from vllm.config import VllmConfig from vllm.model_executor.layers.fused_moe import FusedMoERouter +logger = init_logger(__name__) + AFD_ASYNC_CAM_GROUP_NAME = "afd_async_cam" -CAM_COMM_ID = 0 ATTN_RANKS_PER_DP_CONFIG_KEY = "attn_ranks_per_dp" _AFD_ASYNC_EXTRA_CONFIG_FIELDS: Final[frozenset[str]] = frozenset( { "dynamicQuant", + "ffn_mega_moe", "attn_ranks_per_dp", "hccl_buffer_size", - "ffn_two_stream", - "ffn_compute_aic_num", - "ffn_compute_aiv_num", - "ffn_combine_aiv_num", }, ) -logger = init_logger(__name__) - def select_cam_experts( *, @@ -149,15 +139,13 @@ class AFDAsyncExtraInfo(ConnectorExtraInfo): dynamic_quant: Dynamic quantization mode accepted by CAM operators. attn_ranks_per_dp: Number of Attention ranks in each data-parallel group. hccl_buffer_size: Optional buffer size in MB for the CAM HCCL domain. + ffn_mega_moe: Enable Mega FFN and its pull protocol on every A/F rank. """ + ffn_mega_moe: bool = False dynamic_quant: int = 0 attn_ranks_per_dp: int = 1 hccl_buffer_size: int | None = None - ffn_two_stream: bool = False - ffn_compute_aic_num: int = 12 - ffn_compute_aiv_num: int = 24 - ffn_combine_aiv_num: int = 24 @classmethod def from_mapping(cls, raw: Mapping[str, Any] | None) -> AFDAsyncExtraInfo: @@ -177,7 +165,10 @@ def from_mapping(cls, raw: Mapping[str, Any] | None) -> AFDAsyncExtraInfo: + ", ".join(unknown), ) - return cls( + result = cls( + ffn_mega_moe=coerce_extra_bool( + raw.get("ffn_mega_moe", False), field_name="ffn_mega_moe" + ), dynamic_quant=coerce_extra_int( raw.get("dynamicQuant", 0), field_name="dynamicQuant", @@ -190,55 +181,24 @@ def from_mapping(cls, raw: Mapping[str, Any] | None) -> AFDAsyncExtraInfo: raw.get("hccl_buffer_size"), field_name="hccl_buffer_size", ), - ffn_two_stream=coerce_extra_bool( - raw.get("ffn_two_stream", False), field_name="ffn_two_stream" - ), - ffn_compute_aic_num=coerce_extra_positive_int( - raw.get("ffn_compute_aic_num", 12), field_name="ffn_compute_aic_num" - ), - ffn_compute_aiv_num=coerce_extra_positive_int( - raw.get("ffn_compute_aiv_num", 24), field_name="ffn_compute_aiv_num" - ), - ffn_combine_aiv_num=coerce_extra_positive_int( - raw.get("ffn_combine_aiv_num", 24), field_name="ffn_combine_aiv_num" - ), ) + if result.ffn_mega_moe and result.dynamic_quant != 1: + raise ValueError("ffn_mega_moe requires dynamicQuant=1") + return result + def to_mapping(self) -> dict[str, Any]: result: dict[str, Any] = { "dynamicQuant": self.dynamic_quant, "attn_ranks_per_dp": self.attn_ranks_per_dp, } + if self.ffn_mega_moe: + result["ffn_mega_moe"] = True if self.hccl_buffer_size is not None: result["hccl_buffer_size"] = self.hccl_buffer_size - if self.ffn_two_stream or ( - self.ffn_compute_aic_num, - self.ffn_compute_aiv_num, - self.ffn_combine_aiv_num, - ) != (12, 24, 24): - result.update( - ffn_two_stream=self.ffn_two_stream, - ffn_compute_aic_num=self.ffn_compute_aic_num, - ffn_compute_aiv_num=self.ffn_compute_aiv_num, - ffn_combine_aiv_num=self.ffn_combine_aiv_num, - ) return result -@dataclass(slots=True) -class AFDAsyncTransferState(AFDTransferState): - """Routed activations, counts and untouched compact combine metadata.""" - - batch_size: int = 1 - hidden_size: int = 1 - topk: int = 1 - layer_idx: int = 0 - token_nums_rankid_layeridx: Tensor | None = None - group_list: Tensor | None = None - dynamic_scales: Tensor | None = None - send_completion: Tensor | None = None - - @dataclass(frozen=True, slots=True) class AFDAsyncTopology: """Role-local and HCCL-world rank information for one CAM participant.""" @@ -256,6 +216,20 @@ def world_size(self) -> int: return self.attn_size + self.ffn_size +@dataclass(slots=True) +class AFDAsyncTransferState(AFDTransferState): + """Routed activations, counts and untouched compact combine metadata.""" + + batch_size: int = 1 + hidden_size: int = 1 + topk: int = 1 + layer_idx: int = 0 + token_nums_rankid_layeridx: Tensor | None = None + group_list: Tensor | None = None + dynamic_scales: Tensor | None = None + send_completion: Tensor | None = None + + class CAMAsyncAFDConnector(AFDConnectorBase): """CAM-backed asynchronous connector for Ascend NPU AFD. @@ -305,7 +279,6 @@ def __init__( self.max_num_batched_tokens = ( vllm_config.scheduler_config.max_num_batched_tokens ) - self.comm_id = CAM_COMM_ID self.tp_size = extra_info.attn_ranks_per_dp self.cam_pg: ProcessGroup | None = None self.topology = build_async_topology( @@ -317,15 +290,35 @@ def __init__( self.attn_size = self.topology.attn_size self.ffn_size = self.topology.ffn_size self.expert_per_rank = self.topology.expert_per_rank - self.comm_args: Tensor | None = None - self._placeholder: Tensor | None = None - self._pending_attention_payloads: dict[ - int, - list[tuple[AFDTransferContext, Tensor, Tensor]], - ] = {} self.graph_startup = AsyncCAMGraphStartup( self.world_rank, self.attn_size, self.topology.world_size ) + self._session: PushCAMSession | PullCAMSession = ( + PullCAMSession( + self.graph_startup, + self.topology, + hidden_size=self.hidden_size, + topk=self.topk, + max_num_batched_tokens=self.max_num_batched_tokens, + tp_size=self.tp_size, + receive_capacity=cam_receive_capacity(), + ) + if extra_info.ffn_mega_moe + else PushCAMSession(self.activation_dtype) + ) + + @property + def comm_args(self) -> Tensor | None: + """Expose the selected protocol's live window to Mega compute.""" + return self._session.comm_args + + @property + def comm_id(self) -> int: + return self._session.comm_id + + @property + def has_pending_attention(self) -> bool: + return self._session.has_pending_attention @property def is_initialized(self) -> bool: @@ -357,21 +350,16 @@ def init_afd_connector(self) -> None: ) backend = self.cam_pg._get_backend(torch.device("npu")) self.group_name = str(backend.get_hccl_comm_name(self.world_rank)) - device = f"npu:{self.local_rank}" - self.comm_args = torch.empty((1,), dtype=torch.float16, device=device) - self._placeholder = torch.empty( - (1,), - dtype=self.activation_dtype, - device=device, - ) + try: + self._session.initialize(f"npu:{self.local_rank}", self.group_name) + except Exception as exc: + self.graph_startup.fail(exc) + raise self._initialized = True if self.graph_startup.mode == "full" and self.world_rank < self.attn_size: self._run_attention_graph_startup() def _initialize_graph_startup(self, store: Store) -> None: - self.graph_startup.validate_attention_streams = ( - self._validate_two_stream_attention - ) signature = repr( ( self.attn_size, @@ -385,37 +373,23 @@ def _initialize_graph_startup(self, store: Store) -> None: self.max_num_batched_tokens, ) ) + signature = self._session.signature_prefix + signature self.graph_startup.initialize(store, signature) if self.graph_startup.mode == "full": try: if not self.vllm_config.use_v2_model_runner: raise ValueError("Async FFN FULL requires MRV2 on every rank") validate_graph_capacity( - self.max_num_batched_tokens, self.tp_size, cam_receive_capacity() + self.max_num_batched_tokens, + self.tp_size, + cam_receive_capacity(), + allow_row_slicing=self._session.allow_row_slicing, ) self.graph_startup.phase("validated") except Exception as exc: self.graph_startup.fail(exc) raise - def _validate_two_stream_attention(self) -> None: - config = self.vllm_config - parallel = config.parallel_config - if ( - not config.use_v2_model_runner - or not config.model_config.enforce_eager - or parallel.pipeline_parallel_size != 1 - or parallel.prefill_context_parallel_size != 1 - or parallel.decode_context_parallel_size != 1 - or parallel.use_sequence_parallel_moe - or os.environ.get("VLLM_ASCEND_ENABLE_FLASHCOMM1", "0").strip().lower() - in ("1", "true") - or config.speculative_config is not None - ): - raise ValueError( - "FFN two-stream requires eager MRV2 Attention without PP/CP/SP" - ) - def wait_graph_device(self, event: torch.npu.Event) -> None: deadline = time.monotonic() + GRAPH_STARTUP_TIMEOUT_SECONDS while not event.query(): @@ -426,6 +400,9 @@ def wait_graph_device(self, event: torch.npu.Event) -> None: def graph_startup_chunks(self) -> int: """Every A TP rank sends one row to experts [0, topk).""" + if isinstance(self._session, PullCAMSession): + # A fused launch drains every receive round of the selected source. + return 1 local_experts = max( 0, min( @@ -462,16 +439,21 @@ def _run_attention_graph_startup(self) -> None: self.graph_startup.fail(exc) raise + def ffn_session(self) -> FFNExecutionSession: + return self._session + + def attention_step(self) -> AbstractContextManager[None]: + return self._session.attention_step() + def close(self) -> None: """Destroy the HCCL process group and clear pending transfer states.""" + self._session.prepare_close(self._initialized) if self.cam_pg is not None: import torch.distributed as dist dist.destroy_process_group(self.cam_pg) self.cam_pg = None - self.comm_args = None - self._placeholder = None - self._pending_attention_payloads.clear() + self._session.release() self._initialized = False def select_experts(self, **kwargs: Any) -> tuple[Tensor, Tensor]: @@ -484,12 +466,7 @@ def send_attn_output( context: AFDTransferContext, **kwargs: Any, ) -> None: - """Dispatch Attention activations and top-k routing to FFN ranks. - - The matching context and routing tensors are queued per async stage - for ``recv_ffn_output``. The input token count and top-k tensor shapes - must match the transfer metadata and model configuration. - """ + """Validate routed activations, then submit through the selected protocol.""" self._require_initialized() metadata = context.metadata if not metadata.validate_tensor_shape(tuple(hidden_states.shape)): @@ -523,54 +500,18 @@ def send_attn_output( "Async CAM batch exceeds max_num_batched_tokens / TP capacity" ) hidden_states = hidden_states.contiguous() - - _log_cam_op_values( - "async_dispatch_send", - "inputs", - hidden_states=hidden_states, - topk_ids=topk_ids, - comm_args=self.comm_args, - comm_id=self.comm_id, + args = self._cam_args(states) | dict( + x=hidden_states, + expert_ids=topk_ids, max_seq_len=self.max_num_batched_tokens, batch_size=states.batch_size, - hidden_size=states.hidden_size, - topk=states.topk, - ffn_size=self.ffn_size, - attn_size=self.attn_size, - expert_per_rank=self.expert_per_rank, - rank=self.world_rank, - world_size=self.topology.world_size, - layer_idx=states.layer_idx, + attn_rank_id=self.world_rank, + layer_index=states.layer_idx, tp_size=self.tp_size, dynamic_quant=self.dynamic_quant, - group_name=self.group_name, - ) - torch.ops.afd_ascend.afd_async_dispatch_send( - hidden_states, - topk_ids, - self.comm_args, - self.comm_id, - self.max_num_batched_tokens, - states.batch_size, - states.hidden_size, - states.topk, - self.ffn_size, - self.attn_size, - self.expert_per_rank, - self.world_rank, - self.topology.world_size, - states.layer_idx, - self.tp_size, - self.dynamic_quant, - self.group_name, - ) - self._pending_attention_payloads.setdefault( - context.metadata.stage_idx, [] - ).append( - (context, topk_ids, topk_weights), ) - - return None + _log_cam_op_values("async_dispatch_send", "inputs", **args) + self._session.dispatch((context, topk_ids, topk_weights), **args) def recv_ffn_output( self, @@ -578,36 +519,21 @@ def recv_ffn_output( ubatch_idx: int = 0, **kwargs: Any, ) -> Tensor: - """Receive and combine expert output on an Attention rank. - - The optional context must match the oldest pending dispatch for the - stage. Routing tensors always come from that dispatch. ``ref_tensor`` - selects the output device and dtype for the CAM receive placeholder. - - Args: - ref_tensor: Tensor from which device/dtype are taken for the CAM - combine-recv placeholder. - ubatch_idx: Stage/microbatch index of the pending Attention - dispatch payload. Defaults to ``0``. - **kwargs: Optional ``context`` from the matching dispatch. Uses - FIFO order when omitted. - """ + """Combine the oldest matching Push dispatch or the exclusive Pull window.""" self._require_initialized() context = kwargs.get("context") - payloads = self._pending_attention_payloads.get(ubatch_idx) - if not payloads: + payload = self._session.pending_for(ubatch_idx) + if payload is None: raise RuntimeError( "CAMAsyncAFDConnector recv_ffn_output is missing pending " "Attention metadata", ) - if context is not None and context is not payloads[0][0]: + if context is not None and context is not payload[0]: raise RuntimeError( "CAMAsyncAFDConnector recv_ffn_output context does not match " "the next pending Attention dispatch", ) - context, topk_ids, topk_weights = payloads.pop(0) - if not payloads: - self._pending_attention_payloads.pop(ubatch_idx, None) + context, topk_ids, topk_weights = payload states = _require_async_transfer_state(context) _validate_topk_payload( @@ -617,40 +543,15 @@ def recv_ffn_output( topk=states.topk, ) placeholder = ref_tensor.new_empty((1,)) - _log_cam_op_values( - "async_combine_recv", - "inputs", - placeholder=placeholder, - topk_ids=topk_ids, - topk_weights=topk_weights, - comm_args=self.comm_args, - comm_id=self.comm_id, + args = self._cam_args(states) | dict( + expand_x=placeholder, + expert_ids=topk_ids, + expert_scales=topk_weights, batch_size=states.batch_size, - hidden_size=states.hidden_size, - topk=states.topk, - ffn_size=self.ffn_size, - attn_size=self.attn_size, - expert_per_rank=self.expert_per_rank, - rank=self.world_rank, - world_size=self.topology.world_size, - group_name=self.group_name, - ) - output = torch.ops.afd_ascend.afd_async_combine_recv( - placeholder, - topk_ids, - topk_weights, - self.comm_args, - self.comm_id, - states.batch_size, - states.hidden_size, - states.topk, - self.ffn_size, - self.attn_size, - self.expert_per_rank, - self.world_rank, - self.topology.world_size, - self.group_name, + attn_rank_id=self.world_rank, ) + _log_cam_op_values("async_combine_recv", "inputs", **args) + output = self._session.combine(ubatch_idx, **args) _log_cam_op_values("async_combine_recv", "outputs", output=output) return output @@ -659,8 +560,8 @@ def recv_attn_output( ubatch_idx: int = 0, **kwargs: Any, ) -> AFDA2FTransferPayload: - """Receive routed activations and preserve their compact metadata.""" - self._require_initialized() + """Receive a Push work item and preserve its compact return metadata.""" + session = self._require_push_session() batch_size = int(kwargs.get("batch_size", self.max_num_batched_tokens) or 1) layer_idx = int(kwargs.get("layer_idx", 0) or 0) metadata = AFDTransferMetadata.create_ffn_metadata( @@ -678,42 +579,16 @@ def recv_attn_output( metadata=metadata, states=states, ) - placeholder = kwargs.get("placeholder", self._placeholder) - _log_cam_op_values( - "async_dispatch_recv", - "inputs", - placeholder=placeholder, - comm_args=self.comm_args, - comm_id=self.comm_id, + placeholder = kwargs.get("placeholder", session.placeholder) + args = self._cam_args(states) | dict( + x=placeholder, max_seq_len=self.max_num_batched_tokens, - batch_size=states.batch_size, - hidden_size=states.hidden_size, - topk=states.topk, - ffn_size=self.ffn_size, - attn_size=self.attn_size, - expert_per_rank=self.expert_per_rank, - rank=self.world_rank, - world_size=self.topology.world_size, + moe_rank_id=self.world_rank, tp_size=self.tp_size, dynamic_quant=self.dynamic_quant, - group_name=self.group_name, - ) - outputs = torch.ops.afd_ascend.afd_async_dispatch_recv( - placeholder, - self.comm_args, - self.comm_id, - self.max_num_batched_tokens, - states.hidden_size, - states.topk, - self.ffn_size, - self.attn_size, - self.expert_per_rank, - self.world_rank, - self.topology.world_size, - self.tp_size, - self.dynamic_quant, - self.group_name, ) + _log_cam_op_values("async_dispatch_recv", "inputs", **args) + outputs = torch.ops.afd_ascend.afd_async_dispatch_recv(**args) hidden_states, dynamic_scales, batch_info, expert_token_nums = outputs _log_cam_op_values( "async_dispatch_recv", @@ -737,12 +612,8 @@ def send_ffn_output( context: AFDTransferContext, **kwargs: Any, ) -> None: - """Send routed expert outputs for CAM combination. - - ``TokenNums_Rankid_Layeridx`` from the matching dispatch-recv is - mandatory because CAM uses it to return results to Attention ranks. - """ - self._require_initialized() + """Return a Push work item with its original dispatch metadata.""" + self._require_push_session() states = _require_async_transfer_state(context) token_nums_rankid_layeridx = states.token_nums_rankid_layeridx if token_nums_rankid_layeridx is None: @@ -752,95 +623,45 @@ def send_ffn_output( "AFD async CAM combine send requires " "TokenNums_Rankid_Layeridx from async_dispatch_recv", ) - _log_cam_op_values( - "async_combine_send", - "inputs", - ffn_output=ffn_output, + args = self._cam_args(states) | dict( + expand_x=ffn_output, + batch_info=token_nums_rankid_layeridx, + max_seq_len=self.max_num_batched_tokens, + moe_rank_id=self.world_rank, + tp_size=self.tp_size, + ) + _log_cam_op_values("async_combine_send", "inputs", **args) + states.send_completion = torch.ops.afd_ascend.afd_async_combine_send(**args) + + def _require_push_session(self) -> PushCAMSession: + self._require_initialized() + if not isinstance(self._session, PushCAMSession): + raise RuntimeError( + "Pull FFN receive/compute/send belongs to the Mega executor" + ) + return self._session + + def _cam_args( + self, states: AFDAsyncTransferState + ) -> dict[str, Tensor | int | str | None]: + """Fields shared by the four registered CAM operator schemas.""" + return dict( comm_args=self.comm_args, - token_nums_rankid_layeridx=token_nums_rankid_layeridx, - metadata_values=False, comm_id=self.comm_id, - max_seq_len=self.max_num_batched_tokens, - batch_size=states.batch_size, hidden_size=states.hidden_size, - topk=states.topk, - ffn_size=self.ffn_size, - attn_size=self.attn_size, - expert_per_rank=self.expert_per_rank, - rank=self.world_rank, + top_k=states.topk, + moe_rank_num=self.ffn_size, + attn_rank_num=self.attn_size, + route_expert_num_per_moe=self.expert_per_rank, world_size=self.topology.world_size, - tp_size=self.tp_size, group_name=self.group_name, ) - states.send_completion = torch.ops.afd_ascend.afd_async_combine_send( - ffn_output, - self.comm_args, - token_nums_rankid_layeridx, - self.comm_id, - self.max_num_batched_tokens, - states.hidden_size, - states.topk, - self.ffn_size, - self.attn_size, - self.expert_per_rank, - self.world_rank, - self.topology.world_size, - self.tp_size, - self.group_name, - ) def _require_initialized(self) -> None: if not self._initialized: raise RuntimeError("CAMAsyncAFDConnector is not initialized") -def _require_async_transfer_state( - context: AFDTransferContext, -) -> AFDAsyncTransferState: - states = context.states - if not isinstance(states, AFDAsyncTransferState): - raise RuntimeError( - "CAMAsyncAFDConnector requires AFDAsyncTransferState in the " - "transfer context", - ) - return states - - -_CAM_LOG_SKIPPED_ARGS = frozenset({"comm_args", "comm_id", "group_name"}) -_CAM_OP_IO_LOG_ENV = "AFD_CAM_OP_IO_LOG" - - -def _log_cam_op_values( - op_name: str, label: str, *, metadata_values: bool = True, **kwargs: object -) -> None: - if os.environ.get(_CAM_OP_IO_LOG_ENV, "").lower() not in { - "1", - "true", - "yes", - "on", - }: - return - lines: list[str] = [] - for name, value in kwargs.items(): - if name in _CAM_LOG_SKIPPED_ARGS: - continue - if isinstance(value, Tensor): - description = ( - f"Tensor(dtype={value.dtype}, shape={tuple(value.shape)}, " - f"device={value.device})" - ) - if metadata_values and name == "token_nums_rankid_layeridx": - try: - first5: object = value.detach().flatten()[:5].cpu().tolist() - except Exception as exc: # pragma: no cover - defensive logging helper - first5 = f"" - description = f"{description}, first5={first5!r}" - else: - description = repr(value) - lines.append(f" {name}={description}") - logger.warning("AFD CAM %s %s:\n%s", op_name, label, "\n".join(lines)) - - def build_async_topology( afd_config: AFDConfig, role_rank: int, @@ -892,20 +713,57 @@ def build_async_topology( ) +def _require_async_transfer_state( + context: AFDTransferContext, +) -> AFDAsyncTransferState: + states = context.states + if not isinstance(states, AFDAsyncTransferState): + raise RuntimeError( + "CAMAsyncAFDConnector requires AFDAsyncTransferState in the " + "transfer context", + ) + return states + + +_CAM_LOG_SKIPPED_ARGS = frozenset({"comm_args", "comm_id", "group_name"}) +_CAM_OP_IO_LOG_ENV = "AFD_CAM_OP_IO_LOG" + + +def _log_cam_op_values(op_name: str, label: str, **kwargs: object) -> None: + if os.environ.get(_CAM_OP_IO_LOG_ENV, "").lower() not in { + "1", + "true", + "yes", + "on", + }: + return + lines: list[str] = [] + for name, value in kwargs.items(): + if name in _CAM_LOG_SKIPPED_ARGS: + continue + if isinstance(value, Tensor): + description = ( + f"Tensor(dtype={value.dtype}, shape={tuple(value.shape)}, " + f"device={value.device})" + ) + + else: + description = repr(value) + lines.append(f" {name}={description}") + logger.warning("AFD CAM %s %s:\n%s", op_name, label, "\n".join(lines)) + + def _validate_topk_payload( topk_ids: Tensor, topk_weights: Tensor | None, *, batch_size: int, topk: int, - require_weights: bool = True, ) -> None: if tuple(topk_ids.shape) != (batch_size, topk): raise ValueError( f"topk_ids shape must match ({batch_size}, {topk}), got {topk_ids.shape!r}", ) - if not require_weights and topk_weights is None: - return if topk_weights is None: raise ValueError("topk_weights is required") if tuple(topk_weights.shape) != (batch_size, topk): diff --git a/afd_plugin/connectors/npu/async_cam_runtime/__init__.py b/afd_plugin/connectors/npu/async_cam_runtime/__init__.py new file mode 100644 index 000000000..281bcb14b --- /dev/null +++ b/afd_plugin/connectors/npu/async_cam_runtime/__init__.py @@ -0,0 +1,3 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Async CAM startup and Push/Pull communication sessions.""" diff --git a/afd_plugin/connectors/npu/async_cam_runtime/async_cam_session.py b/afd_plugin/connectors/npu/async_cam_runtime/async_cam_session.py new file mode 100644 index 000000000..7405077af --- /dev/null +++ b/afd_plugin/connectors/npu/async_cam_runtime/async_cam_session.py @@ -0,0 +1,277 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Push/Pull CAM transactions, communication state and FFN lifecycle.""" + +from __future__ import annotations + +from collections.abc import Callable, Iterator +from contextlib import AbstractContextManager, contextmanager, nullcontext +from typing import TYPE_CHECKING + +from afd_plugin.connectors.base import FFNExecutionSession + +if TYPE_CHECKING: + import torch + from torch import Tensor + + from afd_plugin.connectors.metadata import AFDTransferContext + from afd_plugin.connectors.npu.async_cam import AFDAsyncTopology + from afd_plugin.connectors.npu.async_cam_runtime.async_cam_startup import ( + AsyncCAMGraphStartup, + ) + +CAM_COMM_ID = 0 +PULL_COMM_ID = 0x41464401 +PULL_FAILURE_INDEX = 24 +PULL_SOURCE_INDEX = 32 +PULL_IDLE_BUDGET_INDEX = 64 +PULL_IDLE_SCANS = 8 + + +class PushCAMSession(FFNExecutionSession): + """Push owns per-stage routing queues; FFN uses separate receive/send calls.""" + + comm_id = CAM_COMM_ID + signature_prefix = "" + allow_row_slicing = False + + def __init__(self, dtype: torch.dtype) -> None: + self.dtype = dtype + self.comm_args: Tensor | None = None + self.placeholder: Tensor | None = None + self.pending: dict[int, list[tuple[AFDTransferContext, Tensor, Tensor]]] = {} + + @property + def has_pending_attention(self) -> bool: + return bool(self.pending) + + def initialize(self, device: str, group_name: str) -> None: + import torch + + self.comm_args = torch.empty((1,), dtype=torch.float16, device=device) + self.placeholder = torch.empty((1,), dtype=self.dtype, device=device) + + def pending_for( + self, stage: int + ) -> tuple[AFDTransferContext, Tensor, Tensor] | None: + queue = self.pending.get(stage) + return queue[0] if queue else None + + def dispatch( + self, + payload: tuple[AFDTransferContext, Tensor, Tensor], + **args: Tensor | int | str | None, + ) -> None: + import torch + + torch.ops.afd_ascend.afd_async_dispatch_send(**args) + self.pending.setdefault(payload[0].metadata.stage_idx, []).append(payload) + + def combine(self, stage: int, **args: Tensor | int | str | None) -> Tensor: + import torch + + output = torch.ops.afd_ascend.afd_async_combine_recv(**args) + queue = self.pending[stage] + queue.pop(0) + if not queue: + del self.pending[stage] + return output + + def attention_step(self) -> AbstractContextManager[None]: + return nullcontext() + + def prepare_close(self, initialized: bool) -> None: + pass + + def release(self) -> None: + self.comm_args = self.placeholder = None + self.pending.clear() + + +class PullCAMSession(FFNExecutionSession): + """One in-flight Attention window; FFN runs the native fused transaction.""" + + requires_drain = True + comm_id = PULL_COMM_ID + signature_prefix = "pull_v1:" + allow_row_slicing = True + + def __init__( + self, + startup: AsyncCAMGraphStartup, + topology: AFDAsyncTopology, + *, + hidden_size: int, + topk: int, + max_num_batched_tokens: int, + tp_size: int, + receive_capacity: int, + ) -> None: + self.startup = startup + self.max_num_batched_tokens = max_num_batched_tokens + self.topk = topk + self.tp_size = tp_size + self.receive_capacity = receive_capacity + self._init_config = [ + max_num_batched_tokens // tp_size, + hidden_size, + topk, + topology.ffn_size, + topology.attn_size, + topology.expert_per_rank, + topology.world_rank, + tp_size, + ] + self.comm_args: Tensor | None = None + self._pending: tuple[AFDTransferContext, Tensor, Tensor] | None = None + self._defer_status = False + self.receiver_finished = False + self._idle_enabled = False + self._drained = False + + @property + def has_pending_attention(self) -> bool: + return self._pending is not None + + def initialize(self, device: str, group_name: str) -> None: + import torch + + self.comm_args = torch.ops.afd_ascend.afd_mega_pull_init( + torch.empty(1, dtype=torch.bfloat16, device=device), + self._init_config, + list(self.startup.plan.layer_ids), + self.startup.session_epoch, + group_name, + ) + torch.npu.synchronize() + self.check_pull_status() + self.startup.phase("pull-initialized") + + def pending_for( + self, stage: int + ) -> tuple[AFDTransferContext, Tensor, Tensor] | None: + if self._pending is not None and self._pending[0].metadata.stage_idx == stage: + return self._pending + return None + + def dispatch( + self, + payload: tuple[AFDTransferContext, Tensor, Tensor], + **args: Tensor | int | str | None, + ) -> None: + import torch + + if self._pending is not None: + raise RuntimeError("Pull window is pending; finish Combine before Dispatch") + # A launch can enqueue work before raising; retain ownership on failure. + self._pending = payload + torch.ops.afd_ascend.afd_mega_dispatch_send(**args) + + def combine(self, stage: int, **args: Tensor | int | str | None) -> Tensor: + import torch + + output = torch.ops.afd_ascend.afd_mega_combine_recv( + **args, max_seq_len=self.max_num_batched_tokens, tp_size=self.tp_size + ) + if not self._defer_status: + self.check_pull_status() + self._pending = None + return output + + @contextmanager + def attention_step(self) -> Iterator[None]: + previous = self._defer_status + self._defer_status = True + try: + yield + if not previous: + self.check_pull_status() + finally: + self._defer_status = previous + + def check_pull_status(self) -> None: + import torch + + assert self.comm_args is not None + # The model runner batches this fence outside the DBO receive hook. + status = int(self.comm_args.view(torch.int64)[PULL_FAILURE_INDEX].item()) + if status: + raise RuntimeError(f"AFD pull transaction failed with status {status}") + + def enable_pull_idle_returns(self) -> None: + import torch + + assert self.comm_args is not None + self.comm_args.view(torch.int64)[PULL_IDLE_BUDGET_INDEX].fill_(PULL_IDLE_SCANS) + + def pull_has_active_source(self) -> bool: + import torch + + assert self.comm_args is not None + return int(self.comm_args.view(torch.int64)[PULL_SOURCE_INDEX].item()) >= 0 + + def prepare_close(self, initialized: bool) -> None: + if self.has_pending_attention: + raise RuntimeError( + "Pull dispatch is pending; retain communication resources" + ) + if initialized: + if self.startup.world_rank < self.startup.attn_size: + self.close_attention() + else: + self.validate_shutdown() + + def release(self) -> None: + self.comm_args = None + self._pending = None + + def keep_serving(self, shutdown_requested: bool) -> bool: + startup = self.startup + assert startup._store is not None + startup.check_failure() + return not startup._store.check( + startup._keys("pull_quiet", range(startup.attn_size)) + ) + + def before_first_step(self) -> None: + if not self._idle_enabled: + self.enable_pull_idle_returns() + self._idle_enabled = True + + def after_step(self) -> None: + self.check_pull_status() + + def drain(self, run_one: Callable[[], None]) -> None: + rounds = ( + self.max_num_batched_tokens * self.topk + self.receive_capacity - 1 + ) // self.receive_capacity + for _ in range(rounds): + if not self.pull_has_active_source(): + self._drained = True + return + run_one() + if self.pull_has_active_source(): + raise RuntimeError("Pull source did not drain; retain resources") + self._drained = True + + def finish_receiver(self) -> None: + if self.receiver_finished: + return + if not self._drained: + raise RuntimeError("Pull receiver has not drained; retain resources") + startup = self.startup + startup._wait(startup._keys("pull_quiet", range(startup.attn_size))) + startup._publish("pull_stopped") + self.receiver_finished = True + + def close_attention(self) -> None: + startup = self.startup + startup._publish("pull_quiet") + startup._wait( + startup._keys("pull_stopped", range(startup.attn_size, startup.world_size)) + ) + + def validate_shutdown(self) -> None: + if not self.receiver_finished or self.pull_has_active_source(): + raise RuntimeError("Pull receiver has not drained; retain resources") + self.check_pull_status() diff --git a/afd_plugin/async_cam_startup.py b/afd_plugin/connectors/npu/async_cam_runtime/async_cam_startup.py similarity index 83% rename from afd_plugin/async_cam_startup.py rename to afd_plugin/connectors/npu/async_cam_runtime/async_cam_startup.py index f2c48f14f..97b12f7cd 100644 --- a/afd_plugin/async_cam_startup.py +++ b/afd_plugin/connectors/npu/async_cam_runtime/async_cam_startup.py @@ -7,7 +7,6 @@ import json import time import uuid -from collections.abc import Callable from dataclasses import asdict, dataclass from typing import TYPE_CHECKING, Literal @@ -24,18 +23,11 @@ class FFNGraphPlan: executor_error: str = "" signature: str = "" layer_ids: tuple[int, ...] = () - stream_signature: str = "" - stream_error: str = "" def resolve_ffn_mode(plans: list[FFNGraphPlan]) -> str: if not plans or len({plan.requested for plan in plans}) != 1: raise ValueError("All FFN ranks must request the same execution mode") - if len({plan.stream_signature for plan in plans}) != 1: - raise ValueError("FFN two-stream execution signatures differ across ranks") - errors = [plan.stream_error for plan in plans if plan.stream_error] - if errors: - raise ValueError("Async FFN two-stream is unsupported: " + "; ".join(errors)) errors = [plan.executor_error for plan in plans if plan.executor_error] if errors: raise ValueError( @@ -46,8 +38,6 @@ def resolve_ffn_mode(plans: list[FFNGraphPlan]) -> str: if mode not in ("eager", "full"): raise ValueError("Unknown Async FFN execution mode") execution = "FULL" if mode == "full" else "eager" - if mode == "eager" and plans[0].stream_signature: - execution = "two-stream" if len({plan.signature for plan in plans}) > 1: raise ValueError( f"Async FFN {execution} is unsupported: " @@ -70,15 +60,18 @@ def __init__(self, world_rank: int, attn_size: int, world_size: int) -> None: self.plan = FFNGraphPlan() self.mode = "eager" self.warmup_layer_id: int | None = None + self.session_epoch = 0 self._store: Store | None = None self._prefix = "" - self.validate_attention_streams: Callable[[], None] | None = None + self._shared_errors = False def initialize(self, store: Store, transport_signature: str) -> None: self._store = store if self.world_rank == 0: store.set("ffn_graph/nonce", uuid.uuid4().hex) - self._prefix = f"ffn_graph/{store.get('ffn_graph/nonce').decode()}/" + nonce = store.get("ffn_graph/nonce").decode() + self.session_epoch = int(nonce[:15], 16) or 1 + self._prefix = f"ffn_graph/{nonce}/" self._negotiate( self.plan if self.world_rank >= self.attn_size else None, transport_signature, @@ -89,10 +82,19 @@ def _keys(self, stage: str, ranks: range) -> list[str]: def _publish(self, stage: str, value: str = "ready") -> None: assert self._store is not None + if stage == "error": + # Publish the shared signal first so new peers still fail if the + # compatibility publication below loses its Store connection. + self._store.set(f"{self._prefix}error", value) self._store.set(f"{self._prefix}{stage}/{self.world_rank}", value) def check_failure(self) -> None: assert self._store is not None + signal = f"{self._prefix}error" + if self._store.check([signal]): + raise RuntimeError(self._store.get(signal).decode()) + if self._shared_errors: + return for key in self._keys("error", range(self.world_size)): if self._store.check([key]): raise RuntimeError(self._store.get(key).decode()) @@ -126,6 +128,7 @@ def _negotiate(self, plan: FFNGraphPlan | None, transport_signature: str) -> Non json.dumps( { "transport": transport_signature, + "shared_errors": True, "ffn": asdict(plan) if plan else None, } ), @@ -133,6 +136,10 @@ def _negotiate(self, plan: FFNGraphPlan | None, transport_signature: str) -> Non keys = self._keys("plan", range(self.world_size)) self._wait(keys) entries = [json.loads(self._store.get(key)) for key in keys] + # Old peers publish only per-rank errors; keep their original check path. + self._shared_errors = all( + entry.get("shared_errors", False) for entry in entries + ) plans = [] for entry in entries[self.attn_size :]: fields = entry["ffn"] @@ -140,20 +147,14 @@ def _negotiate(self, plan: FFNGraphPlan | None, transport_signature: str) -> Non plans.append(FFNGraphPlan(**fields)) self.mode = resolve_ffn_mode(plans) ffn_plan = plans[0] - two_stream = bool(ffn_plan.stream_signature) - validate_transport = self.mode == "full" or two_stream + validate_transport = self.mode == "full" or any( + entry["transport"].startswith("pull_v1:") for entry in entries + ) if ( validate_transport and len({entry["transport"] for entry in entries}) != 1 ): raise ValueError("Async CAM graph transport signatures differ") - if two_stream: - if ( - self.world_rank < self.attn_size - and self.validate_attention_streams is not None - ): - self.validate_attention_streams() - self.phase("streams-validated") if self.mode == "full": self.warmup_layer_id = ffn_plan.layer_ids[0] except Exception as exc: diff --git a/afd_plugin/connectors/npu/camp2p/__init__.py b/afd_plugin/connectors/npu/camp2p/__init__.py new file mode 100644 index 000000000..77752d2f9 --- /dev/null +++ b/afd_plugin/connectors/npu/camp2p/__init__.py @@ -0,0 +1,43 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Synchronous CAMP2P connector package. + +The connector implementation lives in :mod:`~.connector` and the Attention-side routing +payload in :mod:`~.external_routing`. Names are resolved lazily so importing the payload +does not pull in the connector's vLLM and Ascend dependencies. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from afd_plugin.connectors.npu.camp2p.connector import ( + CAMP2pAFDConnector, + CAMP2PAFDConnectorData, + CAMP2pAFDControlPlane, + CAMP2PExtraInfo, + CAMP2PTransferState, + build_camp2p_topology, + prepare_token_id_transfer, + received_token_ids, + ) + +__all__ = [ + "CAMP2pAFDConnector", + "CAMP2pAFDControlPlane", + "CAMP2PAFDConnectorData", + "CAMP2PExtraInfo", + "CAMP2PTransferState", + "build_camp2p_topology", + "prepare_token_id_transfer", + "received_token_ids", +] + + +def __getattr__(name: str) -> Any: + if name in __all__: + from afd_plugin.connectors.npu.camp2p import connector + + return getattr(connector, name) + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/afd_plugin/connectors/npu/camp2p.py b/afd_plugin/connectors/npu/camp2p/connector.py similarity index 85% rename from afd_plugin/connectors/npu/camp2p.py rename to afd_plugin/connectors/npu/camp2p/connector.py index 99403a9ce..40d93aa35 100644 --- a/afd_plugin/connectors/npu/camp2p.py +++ b/afd_plugin/connectors/npu/camp2p/connector.py @@ -146,10 +146,10 @@ def aiv_num_for_role(self, role: str) -> int: return self.core_num def validate_supported(self) -> None: - if self.compute_gate_on_attention: - raise RuntimeError( - "AFD NPU runtime does not support compute_gate_on_attention=true yet", - ) + # The gate switch is read from ``additional_config["afd"]``; the copy in + # ``connector_extra_config`` is parsed for compatibility and deliberately has no + # say here. feature_validation owns the supported shapes, so a second opinion in + # this class could only disagree with it. if self.quant_mode != 0: raise RuntimeError("AFD NPU runtime currently supports only quant_mode=0") @@ -308,6 +308,112 @@ def received_token_ids( return sim_expert_ids[:expected_tokens, 0].contiguous() +def prepare_routing_transfer( + topk_ids: torch.Tensor, + topk_weights: torch.Tensor, + *, + expected_tokens: int, + top_k: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """Pack an Attention-side gate result for the A2E routing channel. + + The A2E ids channel transports an ``int32`` ids tensor and a ``float32`` + scales tensor of shape ``(batch, topk)``. It was opened to carry token + identity for a token-keyed router; a gate that ran on the Attention ranks + fills the same two slots with the routing it already computed, so the FFN + rank selects experts without running the gate a second time. + + Args: + topk_ids: Expert indices chosen by the Attention-side gate, shaped + ``(tokens, topk)``. + topk_weights: Routing weights for those indices, same shape. + expected_tokens: Token count the transfer metadata declares. + top_k: Experts per token the operator is configured with. The wire + declares this column count, so a routing tensor of another width + would be read with the wrong stride. + + Returns: + The ``(ids, scales)`` pair to hand to the operator, already in the + channel's ``int32`` and ``float32`` layouts. + + Raises: + ValueError: If the routing tensors are not two-dimensional, disagree on + shape, do not describe exactly ``expected_tokens`` tokens, or their + column count differs from ``top_k``. + """ + + if topk_ids.dim() != 2 or topk_weights.dim() != 2: + raise ValueError( + "A2E routing transfer needs two-dimensional (tokens, topk) tensors, " + f"got shapes {tuple(topk_ids.shape)!r} and {tuple(topk_weights.shape)!r}", + ) + if topk_ids.shape != topk_weights.shape: + raise ValueError( + "A2E routing transfer needs ids and weights of the same shape, got " + f"{tuple(topk_ids.shape)!r} and {tuple(topk_weights.shape)!r}", + ) + if int(topk_ids.shape[1]) != top_k: + raise ValueError( + f"A2E routing transfer carries {int(topk_ids.shape[1])} experts per " + f"token but the operator is configured with k={top_k}", + ) + num_tokens = int(topk_ids.shape[0]) + if num_tokens != expected_tokens: + raise ValueError( + f"A2E routing transfer covers {num_tokens} tokens but the AFD " + f"transfer declares {expected_tokens}", + ) + return ( + topk_ids.to(dtype=torch.int32).contiguous(), + topk_weights.to(dtype=torch.float32).contiguous(), + ) + + +def received_routing( + sim_expert_ids: torch.Tensor, + sim_expert_scales: torch.Tensor, + *, + expected_tokens: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """Trim the received routing channel to the tokens this FFN rank computes on. + + The operator sizes the ids and scales outputs to the capacity it was given, + so rows past the FFN rank's token count are padding the MoE must not route. + Both tensors are sliced together so ids and weights stay aligned. + + Args: + sim_expert_ids: The operator's ids output, shaped ``(N, topk)``. + sim_expert_scales: The operator's scales output, shaped ``(N, topk)``. + The channel reuses the scales slot for the routing weights. + expected_tokens: Token count derived from the FFN rank's DP metadata, + which is what the FFN compute will actually run on. + + Returns: + The ``(topk_ids, topk_weights)`` pair restricted to ``expected_tokens``. + + Raises: + ValueError: If either tensor arrives with fewer rows than the FFN rank + computes on, which would misroute tokens. + """ + + received_tokens = int(sim_expert_ids.shape[0]) + if received_tokens < expected_tokens: + raise ValueError( + f"received {received_tokens} routing rows but the FFN rank computes " + f"on {expected_tokens} tokens; the routing channel is not aligned " + "with the FFN token layout", + ) + if int(sim_expert_scales.shape[0]) < expected_tokens: + raise ValueError( + f"received {int(sim_expert_scales.shape[0])} routing scale rows but " + f"the FFN rank computes on {expected_tokens} tokens", + ) + return ( + sim_expert_ids[:expected_tokens].contiguous(), + sim_expert_scales[:expected_tokens].contiguous(), + ) + + class CAMP2pAFDConnector(AFDConnectorBase): """Move model data between Attention and FFN workers on Ascend NPU. @@ -513,17 +619,20 @@ def send_attn_output( hidden_states: Model data with shape ``(tokens, hidden_size)``. context: Transfer context whose ``metadata`` supplies the layer number, ubatch number, and token count for this transfer. - **kwargs: An optional token-aligned ``input_ids`` tensor. When it is - supplied, the transfer runs with ``compute_gate=1`` so the ids - reach the FFN rank through the operator's ids channel, and the - matching ``recv_attn_output(recv_input_ids=True)`` returns them - on the payload's ``input_ids`` field. + **kwargs: Either an optional token-aligned ``input_ids`` tensor or an + optional ``routing`` pair ``(topk_ids, topk_weights)`` produced + by an Attention-side gate; at most one may be supplied because + they share the operator's ids channel. Either one makes the + transfer run with ``compute_gate=1``, and the matching + ``recv_attn_output(recv_input_ids=True)`` or + ``recv_attn_output(recv_routing=True)`` reads the slot back. Raises: RuntimeError: If the communication groups are not ready. ValueError: If the number of tokens in ``hidden_states`` does not - match ``context.metadata`` outside a ``torch.compile`` trace, or - if a supplied ``input_ids`` tensor is malformed. + match ``context.metadata`` outside a ``torch.compile`` trace, if + both wire-slot payloads were supplied, or if a supplied + ``input_ids``/``routing`` value is malformed. """ if not self._initialized: raise RuntimeError("CAMP2P connector is not initialized") @@ -536,6 +645,15 @@ def send_attn_output( f"CAMP2P metadata token count {metadata.total_tokens}", ) input_ids = cast(torch.Tensor | None, kwargs.get("input_ids")) + routing = cast( + tuple[torch.Tensor, torch.Tensor] | None, + kwargs.get("routing"), + ) + if input_ids is not None and routing is not None: + raise ValueError( + "CAMP2P send cannot put both token ids and Attention-side " + "routing on the operator's ids channel", + ) expert_ids: torch.Tensor | None = None expert_scales: torch.Tensor | None = None compute_gate = 0 @@ -546,6 +664,14 @@ def send_attn_output( expected_tokens=metadata.total_tokens, ) compute_gate = 1 + elif routing is not None: + expert_ids, expert_scales = prepare_routing_transfer( + routing[0], + routing[1], + expected_tokens=metadata.total_tokens, + top_k=self.num_experts_per_tok, + ) + compute_gate = 1 transfer_state = CAMP2PTransferState( aiv_num=self.aiv_num, batch_size=metadata.total_tokens, @@ -631,27 +757,36 @@ def recv_attn_output( number needed to create it. ``recv_input_ids`` states that this FFN rank expects token ids on the transfer, which selects the operator's ids mode and makes the connector validate and expose - the operator's ids slot on the payload. The sending rank has to - select the same mode, so only request ids for a run whose - Attention role transports them. + the operator's ids slot on the payload. ``recv_routing`` states + that the same slot carries an Attention-side gate result + instead, exposed as ``topk_ids``/``topk_weights``. The two are + mutually exclusive because they share one wire slot, and the + sending rank has to select the matching mode. Returns: The received hidden states, the information FFN needs to process them - and send the result back, and the transported ``input_ids`` when the - ids mode was selected. + and send the result back, and whichever of the transported + ``input_ids`` or gate result the selected mode carries. Raises: RuntimeError: If communication is not ready, transfer information is missing, or the requested ubatch group does not exist. - ValueError: If ids were requested but do not align with the FFN - rank's token layout. + ValueError: If both wire-slot modes were requested, or if the + transported slot does not align with the FFN rank's token + layout. """ if not self._initialized: raise RuntimeError("CAMP2P connector is not initialized") layer_idx: int = kwargs.get("layer_idx", 0) max_num_tokens: int = kwargs.get("max_num_tokens", 0) recv_input_ids: bool = bool(kwargs.get("recv_input_ids", False)) - compute_gate_mode = 1 if recv_input_ids else 0 + recv_routing: bool = bool(kwargs.get("recv_routing", False)) + if recv_input_ids and recv_routing: + raise ValueError( + "CAMP2P receive cannot read the wire slot as both token ids and " + "Attention-side routing; request exactly one", + ) + compute_gate_mode = 1 if (recv_input_ids or recv_routing) else 0 batch_size = _num_tokens_for_ffn_rank( self.dp_metadata_list, ubatch_idx, @@ -707,15 +842,25 @@ def recv_attn_output( # the other mode would hand the model uninitialised device memory as token # ids, which a token-keyed router turns into an out-of-range table read. received_ids: torch.Tensor | None = None - if compute_gate_mode == 1: + received_topk_ids: torch.Tensor | None = None + received_topk_weights: torch.Tensor | None = None + if recv_input_ids: received_ids = received_token_ids( outputs[1], expected_tokens=batch_size, ) + elif recv_routing: + received_topk_ids, received_topk_weights = received_routing( + outputs[1], + outputs[2], + expected_tokens=batch_size, + ) return AFDA2FTransferPayload( hidden_states=outputs[0], context=context, input_ids=received_ids, + topk_ids=received_topk_ids, + topk_weights=received_topk_weights, ) def send_ffn_output( diff --git a/afd_plugin/connectors/npu/camp2p/external_routing.py b/afd_plugin/connectors/npu/camp2p/external_routing.py new file mode 100644 index 000000000..f519e8a00 --- /dev/null +++ b/afd_plugin/connectors/npu/camp2p/external_routing.py @@ -0,0 +1,159 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Routing inputs the Attention side computed and the FFN side must reuse. + +The FFN side runs the MoE with routing that arrived over the A2E transport +instead of computing it locally. That changes what ``prepare`` has to do: the +activation and the routing must be padded and gathered by exactly the same rule, +or a token row ends up weighted by another row's routing. + +Two shapes of work exist, and this module owns the part they share: + +- ALLGATHER pads to ``max_tokens_across_dp`` and then gathers across the DP + group, so every rank sees every token. +- MC2 pads to ``padded_num_tokens`` and does not gather; the dispatcher moves + tokens to their expert ranks. + +``padded``, ``gathered`` and ``split_rows`` are row-aligned transforms. They keep +the caller's row order and are not required to share storage with the input, so +callers compare contents, not buffers. + +``gathered`` takes the collective as an argument so the DP group, a fake, and a +unit test all use the same code path. +""" + +from __future__ import annotations + +from dataclasses import dataclass, replace +from typing import Final + +import torch + +# Expert indices reach the MoE kernels as integers, and gather/combine treat them as +# integers. Listing the accepted dtypes explicitly also rejects bool and complex, +# which ``is_floating_point()`` alone would let through. +SUPPORTED_INDEX_DTYPES: Final[tuple[torch.dtype, ...]] = ( + torch.int32, + torch.int64, +) + + +@dataclass(frozen=True, eq=False) +class AFDExternalRoutingInputs: + """Routing computed on the Attention side, in the FFN rank's local row order. + + Attributes: + topk_ids: Expert ids per token, ``[tokens, top_k]``. + topk_weights: Gate weights per token, ``[tokens, top_k]``. Any model-level + routed scaling is already folded into these, mirroring the native + router, so the FFN side must not apply it again. + """ + + topk_ids: torch.Tensor + topk_weights: torch.Tensor + + def __post_init__(self) -> None: + # Shapes and dtypes are checked here rather than at the padding site: the + # expert ids and the weights are consumed by kernels that do not validate + # them, so a wrong dtype is only noticed as wrong results. + _require_2d(self.topk_ids, "topk_ids") + _require_2d(self.topk_weights, "topk_weights") + if self.topk_ids.shape != self.topk_weights.shape: + raise ValueError( + "topk_ids and topk_weights must have the same shape, got " + f"{tuple(self.topk_ids.shape)} and {tuple(self.topk_weights.shape)}", + ) + if self.topk_ids.dtype not in SUPPORTED_INDEX_DTYPES: + raise ValueError( + "topk_ids must hold expert indices in one of " + f"{[str(dtype) for dtype in SUPPORTED_INDEX_DTYPES]}, got " + f"{self.topk_ids.dtype}", + ) + if not self.topk_weights.is_floating_point(): + raise ValueError( + "topk_weights must hold gate weights, got dtype " + f"{self.topk_weights.dtype}", + ) + + @property + def token_count(self) -> int: + """Rows before any padding.""" + + return int(self.topk_ids.shape[0]) + + def padded(self, pad_size: int) -> AFDExternalRoutingInputs: + """Append ``pad_size`` inert rows so later collectives stay aligned. + + Padding rows carry weight 0, so whichever expert they name contributes + nothing. That keeps the padding value irrelevant to the result and lets + the ids pad stay at 0. + """ + + if pad_size < 0: + raise ValueError(f"pad_size must be non-negative, got {pad_size}") + if pad_size == 0: + return self + + return replace( + self, + topk_ids=_pad_rows(self.topk_ids, pad_size, 0), + topk_weights=_pad_rows(self.topk_weights, pad_size, 0), + ) + + def gathered(self, all_gather) -> AFDExternalRoutingInputs: + """Concatenate the routing of every DP rank, in rank order. + + ``all_gather`` must be the same collective that gathers the activation, + invoked the same way, so the row order matches token for token. + + The gathered row count is ``dp_size * max(local_rows)``, not the sum of + the local row counts: every rank pads to the DP-wide maximum before the + gather, so each rank's padding rows survive into the result. + """ + + return replace( + self, + topk_ids=all_gather(self.topk_ids), + topk_weights=all_gather(self.topk_weights), + ) + + def split_rows(self, split) -> AFDExternalRoutingInputs: + """Keep one rank's slice, exactly as the activation is sliced. + + MC2 shards the token dimension across TP ranks instead of gathering it, so + the routing has to take the same slice. ``split`` must be the split the + caller applies to the activation. + """ + + return replace( + self, + topk_ids=split(self.topk_ids), + topk_weights=split(self.topk_weights), + ) + + +def _require_2d(tensor: torch.Tensor, name: str) -> None: + """Reject anything that is not a token-by-something matrix.""" + + if tensor.dim() != 2: + raise ValueError( + f"{name} must be 2-D (tokens by something), got shape " + f"{tuple(tensor.shape)}", + ) + + +def _pad_rows(tensor: torch.Tensor, pad_size: int, value: int) -> torch.Tensor: + """Append ``pad_size`` rows filled with ``value``.""" + + pad = torch.full( + (pad_size, *tensor.shape[1:]), + value, + dtype=tensor.dtype, + device=tensor.device, + ) + return torch.cat([tensor, pad], dim=0) + + +__all__ = [ + "AFDExternalRoutingInputs", +] diff --git a/afd_plugin/ffn_faas/__init__.py b/afd_plugin/ffn_faas/__init__.py index aaac2bd74..49e5f0c5a 100644 --- a/afd_plugin/ffn_faas/__init__.py +++ b/afd_plugin/ffn_faas/__init__.py @@ -1,3 +1,3 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project -"""Standalone FFN process using the existing Async CAM backend.""" +"""Standalone FFN process: Async CAM on NPU, DeepEP M2N on GPU.""" diff --git a/afd_plugin/ffn_faas/__main__.py b/afd_plugin/ffn_faas/__main__.py index 6108dc201..9aedb9936 100644 --- a/afd_plugin/ffn_faas/__main__.py +++ b/afd_plugin/ffn_faas/__main__.py @@ -13,6 +13,7 @@ from contextlib import suppress from pathlib import Path +from afd_plugin.ffn_faas.interfaces import FFNPlatformAdapter from afd_plugin.ffn_faas.platforms import SUPPORTED_BACKENDS, get_platform_adapter from afd_plugin.ffn_faas.runner import FFNModelRunner from afd_plugin.ffn_faas.service import FFNService @@ -24,6 +25,24 @@ def _report_exception() -> None: traceback.print_exc() +def create_parser(adapter: FFNPlatformAdapter) -> argparse.ArgumentParser: + parser = adapter.create_parser() + parser.description = __doc__ + parser.add_argument("command", choices=["serve"]) + parser.add_argument("model_tag", nargs="?", help="local model checkpoint") + parser.add_argument("--backend", choices=SUPPORTED_BACKENDS, default="ascend-cam") + parser.add_argument( + "--receiver-shutdown-timeout", + type=float, + default=5.0, + help="seconds to join the standalone receiver before process isolation", + ) + parser.add_argument("--status-host", default="127.0.0.1") + parser.add_argument("--status-port", type=int, default=29700) + parser.add_argument("--status-file", type=Path) + return parser + + def main() -> int: selection = argparse.ArgumentParser(add_help=False) selection.add_argument( @@ -31,15 +50,7 @@ def main() -> int: ) selected, _ = selection.parse_known_args() adapter = get_platform_adapter(selected.backend) - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("command", choices=["serve"]) - parser.add_argument("model", type=Path) - parser.add_argument("--backend", choices=SUPPORTED_BACKENDS, default="ascend-cam") - parser.add_argument("--shutdown-timeout", type=float, default=5.0) - parser.add_argument("--status-host", default="127.0.0.1") - parser.add_argument("--status-port", type=int, default=29700) - parser.add_argument("--status-file", type=Path) - adapter.add_arguments(parser) + parser = create_parser(adapter) args = parser.parse_args() try: launch = adapter.configure(args) diff --git a/afd_plugin/ffn_faas/interfaces.py b/afd_plugin/ffn_faas/interfaces.py index 91f419d7b..9ba73bae6 100644 --- a/afd_plugin/ffn_faas/interfaces.py +++ b/afd_plugin/ffn_faas/interfaces.py @@ -21,7 +21,24 @@ def prepare(self, service: FFNService) -> None: ... def start_receiver(self) -> None: ... - def execute_step(self) -> None: ... + def execute_step(self) -> bool: + """Run one connector-driven step on the receiver thread. + + Returns ``True`` after a completed step. Returns ``False`` only when a + ``cancel_receive()`` request ended an idle receive, so the host loop + exits cleanly; every other failure raises. + """ + ... + + def cancel_receive(self) -> None: + """Unblock an idle receive after a host stop request. + + Called on the host thread while the receiver thread may be blocked in + ``execute_step()``. A backend that cannot cancel a pending receive + implements this as a no-op; the host then reports the pending receive + and retains every owner. + """ + ... def receiver_exited(self) -> None: ... @@ -42,8 +59,8 @@ class FFNLaunch: class FFNPlatformAdapter(Protocol): - """Register platform options and validate a CPU-only launch description.""" + """Build the selected CLI and describe a launch without constructing a backend.""" - def add_arguments(self, parser: ArgumentParser) -> None: ... + def create_parser(self) -> ArgumentParser: ... def configure(self, args: Namespace) -> FFNLaunch: ... diff --git a/afd_plugin/ffn_faas/model_utils.py b/afd_plugin/ffn_faas/model_utils.py new file mode 100644 index 000000000..47a0fa295 --- /dev/null +++ b/afd_plugin/ffn_faas/model_utils.py @@ -0,0 +1,30 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Shared utilities for standalone FFN model owners.""" + +from __future__ import annotations + +from collections.abc import Iterator +from typing import TYPE_CHECKING, Protocol + +if TYPE_CHECKING: + from torch import nn + + +class ParameterStorageOwner(Protocol): + """A module-like owner whose distinct parameter storages are counted.""" + + def parameters(self) -> Iterator[nn.Parameter]: ... + + +def parameter_storage_bytes(model: ParameterStorageOwner) -> int: + """Bytes of the distinct storages behind ``model``'s parameters. + + Views that share storage (for example the connector's expert views) are + counted once. + """ + storage_sizes: dict[tuple[str, int], int] = {} + for parameter in model.parameters(): + storage = parameter.untyped_storage() + storage_sizes[(str(parameter.device), storage.data_ptr())] = storage.nbytes() + return sum(storage_sizes.values()) diff --git a/afd_plugin/ffn_faas/platforms/__init__.py b/afd_plugin/ffn_faas/platforms/__init__.py index fa764be69..b1db5054c 100644 --- a/afd_plugin/ffn_faas/platforms/__init__.py +++ b/afd_plugin/ffn_faas/platforms/__init__.py @@ -6,15 +6,27 @@ from typing import TYPE_CHECKING +from afd_plugin.ffn_faas.platforms.gpu import GPU_DEEPEP_M2N_BACKEND + if TYPE_CHECKING: from afd_plugin.ffn_faas.interfaces import FFNPlatformAdapter -SUPPORTED_BACKENDS = ("ascend-cam",) +ASCEND_CAM_BACKEND = "ascend-cam" +SUPPORTED_BACKENDS = (ASCEND_CAM_BACKEND, GPU_DEEPEP_M2N_BACKEND) def get_platform_adapter(name: str) -> FFNPlatformAdapter: - if name == "ascend-cam": + """Import and construct the CPU-only adapter of one backend by name. + + Each adapter is imported only when selected, so choosing one platform + never imports another platform's modules. + """ + if name == ASCEND_CAM_BACKEND: from afd_plugin.ffn_faas.platforms.npu.adapter import AscendCAMAdapter return AscendCAMAdapter() + if name == GPU_DEEPEP_M2N_BACKEND: + from afd_plugin.ffn_faas.platforms.gpu.adapter import GPUDeepEPM2NAdapter + + return GPUDeepEPM2NAdapter() raise ValueError(f"unsupported standalone FFN backend: {name}") diff --git a/afd_plugin/ffn_faas/platforms/environment.py b/afd_plugin/ffn_faas/platforms/environment.py new file mode 100644 index 000000000..cb5a8ea3c --- /dev/null +++ b/afd_plugin/ffn_faas/platforms/environment.py @@ -0,0 +1,124 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Shared invariants for standalone FFN device and compute environments.""" + +from __future__ import annotations + +import os +from abc import ABC, abstractmethod +from typing import TYPE_CHECKING, Generic, Protocol, TypeVar + +if TYPE_CHECKING: + import torch + from vllm.distributed.parallel_state import GroupCoordinator + + +class FFNEnvironmentSettings(Protocol): + """The standalone-environment subset shared by GPU and NPU settings.""" + + @property + def ffn_ranks(self) -> int: ... + + @property + def rank(self) -> int: ... + + +FFNEnvironmentSettingsT = TypeVar( + "FFNEnvironmentSettingsT", bound=FFNEnvironmentSettings +) + + +class FFNEnvironmentBase(ABC, Generic[FFNEnvironmentSettingsT]): + """Own common F-only environment state without owning native setup.""" + + def __init__( + self, settings: FFNEnvironmentSettingsT, *, platform_name: str + ) -> None: + self.settings = settings + self._platform_name = platform_name + self._device: torch.device | None = None + self._ep_group: GroupCoordinator | None = None + self._groups_owned = False + + @property + def device(self) -> torch.device: + """This rank's device, available after ``initialize_device``.""" + if self._device is None: + raise RuntimeError( + f"{self._platform_name} FFN device has not been initialized" + ) + return self._device + + @property + def ep_group(self) -> GroupCoordinator: + """The F-only expert-parallel group, after ``initialize_compute``.""" + if self._ep_group is None: + raise RuntimeError( + f"{self._platform_name} FFN expert group has not been initialized" + ) + return self._ep_group + + def require_uninitialized(self) -> None: + """Reject a second device initialization.""" + if self._groups_owned or self._device is not None: + raise RuntimeError( + f"{self._platform_name} FFN environment is already initialized" + ) + + def require_launcher_environment(self) -> None: + """Require the launcher world and rank to still match the settings.""" + try: + world_size = int(os.environ["WORLD_SIZE"]) + rank = int(os.environ["RANK"]) + except (KeyError, ValueError) as error: + raise ValueError( + f"{self._platform_name} FFN requires an external F-only launcher" + ) from error + if world_size != self.settings.ffn_ranks or rank != self.settings.rank: + raise ValueError( + f"{self._platform_name} FFN environment differs from the " + "F-only launcher" + ) + + def require_no_borrowed_native_groups(self) -> None: + """Refuse process groups this standalone F process does not own.""" + import torch + from vllm.distributed.parallel_state import model_parallel_is_initialized + + if torch.distributed.is_initialized() or model_parallel_is_initialized(): + raise RuntimeError( + f"{self._platform_name} FFN cannot borrow existing native groups" + ) + + def require_compute_initializable(self) -> None: + """Require device selection and reject repeated compute initialization.""" + if self._device is None or self._groups_owned: + raise RuntimeError( + f"{self._platform_name} FFN compute initialization requires its device" + ) + + def mark_groups_owned(self) -> None: + """Record ownership before the first collective can partially fail.""" + self._groups_owned = True + + def validate_ep_group(self, *, error_message: str) -> None: + """Require the expert group to span exactly the F-only world.""" + ep_group = self.ep_group + if ( + ep_group.world_size != self.settings.ffn_ranks + or ep_group.rank_in_group != self.settings.rank + ): + raise ValueError(error_message) + + def close(self) -> None: + """Destroy owned groups, then clear environment ownership state.""" + if not self._groups_owned: + return + self._destroy_groups() + self._groups_owned = False + self._ep_group = None + self._device = None + + @abstractmethod + def _destroy_groups(self) -> None: + """Destroy this platform's model-parallel and process groups.""" diff --git a/afd_plugin/ffn_faas/platforms/ffn_model/__init__.py b/afd_plugin/ffn_faas/platforms/ffn_model/__init__.py new file mode 100644 index 000000000..2a445863e --- /dev/null +++ b/afd_plugin/ffn_faas/platforms/ffn_model/__init__.py @@ -0,0 +1,3 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Model-specific FFN implementations grouped by model and backend.""" diff --git a/afd_plugin/ffn_faas/platforms/ffn_model/deepseek_v4/__init__.py b/afd_plugin/ffn_faas/platforms/ffn_model/deepseek_v4/__init__.py new file mode 100644 index 000000000..8e63e14ac --- /dev/null +++ b/afd_plugin/ffn_faas/platforms/ffn_model/deepseek_v4/__init__.py @@ -0,0 +1,3 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""DeepSeek-V4 routed-expert implementations.""" diff --git a/afd_plugin/ffn_faas/platforms/ffn_model/deepseek_v4/npu/__init__.py b/afd_plugin/ffn_faas/platforms/ffn_model/deepseek_v4/npu/__init__.py new file mode 100644 index 000000000..ec934b79e --- /dev/null +++ b/afd_plugin/ffn_faas/platforms/ffn_model/deepseek_v4/npu/__init__.py @@ -0,0 +1,3 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Ascend DeepSeek-V4 routed-expert parameter loading.""" diff --git a/afd_plugin/ffn_faas/platforms/npu/v4_model.py b/afd_plugin/ffn_faas/platforms/ffn_model/deepseek_v4/npu/v4_model.py similarity index 98% rename from afd_plugin/ffn_faas/platforms/npu/v4_model.py rename to afd_plugin/ffn_faas/platforms/ffn_model/deepseek_v4/npu/v4_model.py index 0f6cc078d..fcfa2a6ff 100644 --- a/afd_plugin/ffn_faas/platforms/npu/v4_model.py +++ b/afd_plugin/ffn_faas/platforms/ffn_model/deepseek_v4/npu/v4_model.py @@ -214,6 +214,8 @@ def local_sources( return selected def create_owner(self, config: VllmConfig) -> RoutedWeightOwner: - from afd_plugin.ffn_faas.platforms.npu.v4_weights import V4RoutedWeightOwner + from afd_plugin.ffn_faas.platforms.ffn_model.deepseek_v4.npu.v4_weights import ( + V4RoutedWeightOwner, + ) return V4RoutedWeightOwner(self.checkpoint, config) diff --git a/afd_plugin/ffn_faas/platforms/npu/v4_weights.py b/afd_plugin/ffn_faas/platforms/ffn_model/deepseek_v4/npu/v4_weights.py similarity index 98% rename from afd_plugin/ffn_faas/platforms/npu/v4_weights.py rename to afd_plugin/ffn_faas/platforms/ffn_model/deepseek_v4/npu/v4_weights.py index 5e7f9ef52..45325c4e1 100644 --- a/afd_plugin/ffn_faas/platforms/npu/v4_weights.py +++ b/afd_plugin/ffn_faas/platforms/ffn_model/deepseek_v4/npu/v4_weights.py @@ -20,8 +20,10 @@ from vllm.model_executor.layers.fused_moe.router.fused_moe_router import FusedMoERouter from vllm_ascend.models.deepseek_v4 import model as native +from afd_plugin.ffn_faas.platforms.ffn_model.deepseek_v4.npu.v4_model import ( + V4CheckpointDescription, +) from afd_plugin.ffn_faas.platforms.npu.settings import validate_static_experts -from afd_plugin.ffn_faas.platforms.npu.v4_model import V4CheckpointDescription from afd_plugin.model_executor.npu.async_cam_layered import ( AsyncCAMLayeredExecutor, LayeredFFNWeights, diff --git a/afd_plugin/ffn_faas/platforms/gpu/__init__.py b/afd_plugin/ffn_faas/platforms/gpu/__init__.py new file mode 100644 index 000000000..8d763486d --- /dev/null +++ b/afd_plugin/ffn_faas/platforms/gpu/__init__.py @@ -0,0 +1,10 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""GPU standalone FFN backend for routed experts over DeepEP M2N. + +The adapter and settings are CPU-only. The environment, FFN-only model and +backend import torch, vLLM and afdEP lazily, after the adapter has validated +the launch. +""" + +GPU_DEEPEP_M2N_BACKEND = "gpu-deepep-m2n" diff --git a/afd_plugin/ffn_faas/platforms/gpu/adapter.py b/afd_plugin/ffn_faas/platforms/gpu/adapter.py new file mode 100644 index 000000000..e8e587734 --- /dev/null +++ b/afd_plugin/ffn_faas/platforms/gpu/adapter.py @@ -0,0 +1,135 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""GPU DeepEP M2N CLI adaptation; importing this module remains CPU-only. + +Imports neither torch, vLLM nor afdEP at module import: ``create_parser`` +loads vLLM's CLI layer for the selected backend, and every check in +``configure`` runs before a device or process group exists, so a bad launch +exits with a usage error and allocates nothing. Native imports happen in the +backend factory. +""" + +from __future__ import annotations + +from argparse import ArgumentParser, Namespace +from functools import partial +from pathlib import Path +from types import MappingProxyType +from typing import TYPE_CHECKING + +from afd_plugin.config import DEEPEP_M2N_CONNECTOR +from afd_plugin.ffn_faas.interfaces import FFNLaunch +from afd_plugin.ffn_faas.platforms.gpu.settings import ( + BF16_PAYLOAD, + FFN_GRAPH_MODE, + GPUFFNConfig, +) +from afd_plugin.ffn_faas.platforms.launch import ( + configure_ffn_engine_args, + create_ffn_engine_parser, +) + +if TYPE_CHECKING: + from afd_plugin.ffn_faas.interfaces import FFNBackend + + +def create_backend(settings: GPUFFNConfig, args: Namespace) -> FFNBackend: + """Import the GPU backend only when the host loop constructs it. + + Called by ``FFNModelRunner.run`` after the management endpoint is up, so + torch, vLLM and afdEP load inside the launched F process and an import + failure is reported as a FAILED service rather than a usage error. + """ + from afd_plugin.ffn_faas.platforms.gpu.backend import DeepEPM2NFFNBackend + + return DeepEPM2NFFNBackend(settings, args) + + +def uses_cuda_graphs(args: Namespace) -> bool: + """Return the qualified execution mode selected by the native engine CLI. + + ``--enforce-eager`` selects eager; otherwise the compilation config must + name ``FULL_DECODE_ONLY``, the only graph mode the FFN graph cache + serves. Any other mode, including vLLM's unset default, is a usage error. + """ + if args.enforce_eager: + return False + compilation = args.compilation_config + mode = ( + compilation.get("cudagraph_mode") + if isinstance(compilation, dict) + else getattr(compilation, "cudagraph_mode", None) + ) + if getattr(mode, "name", mode) != FFN_GRAPH_MODE: + raise ValueError( + "GPU standalone FFN runs --enforce-eager or --compilation-config " + f'\'{{"cudagraph_mode":"{FFN_GRAPH_MODE}"}}\'' + ) + return True + + +class GPUDeepEPM2NAdapter: + """Keep GPU, topology and connector options out of the host loop. + + The native engine parser and the F-only launcher checks come from the + shared launch module. This adapter adds the GPU constraints: TP 1 with + expert parallelism, the execution mode and the DeepEP lane count. + Attention TP, capacity and the bucket ladder are not launch options; the + connector adopts them from the M2N handshake with Attention. + """ + + def create_parser(self) -> ArgumentParser: + """Build vLLM's engine CLI with the F-only service's fixed defaults. + + F reads no request text, so the tokenizer is skipped, and non-local + expert tensors are skipped before they are read, as the Worker-path + recipes do. + """ + parser = create_ffn_engine_parser() + parser.set_defaults(skip_tokenizer_init=True, enable_ep_weight_filter=True) + return parser + + def configure(self, args: Namespace) -> FFNLaunch: + """Validate the launch on the CPU and describe the backend to build. + + ``additional_config["afd"]`` must select ``DeepEPM2NAFDConnector`` in + the ``ffn`` role, and the ``torchrun`` ``WORLD_SIZE``, + ``--data-parallel-size`` and ``afd.num_ffn_ranks`` must agree. TP is 1 + with expert parallelism, so vLLM's EP group is the F world. The model + path is resolved without interpreting its architecture; vLLM and the + AFD model registry own model selection. Raises ``ValueError`` or + ``OSError``, which the entry point reports as a usage error. + """ + afd, launcher = configure_ffn_engine_args(args, connector=DEEPEP_M2N_CONNECTOR) + if args.tensor_parallel_size != 1 or not args.enable_expert_parallel: + raise ValueError( + "GPU standalone FFN requires --tensor-parallel-size 1 and " + "--enable-expert-parallel" + ) + model_path = Path(args.model).resolve(strict=True) + args.model = str(model_path) + settings = GPUFFNConfig( + model=model_path, + ffn_ranks=launcher.world_size, + rank=launcher.rank, + local_rank=launcher.local_rank, + graph=uses_cuda_graphs(args), + dbo=bool(args.enable_dbo), + connector_extra_config=MappingProxyType(dict(afd.connector_extra_config)), + shutdown_timeout=args.receiver_shutdown_timeout, + ) + return FFNLaunch( + rank=launcher.rank, + local_rank=launcher.local_rank, + model_path=str(model_path), + backend_name=DEEPEP_M2N_CONNECTOR, + shutdown_timeout=settings.shutdown_timeout, + backend_factory=partial(create_backend, settings, args), + initial_status={ + "payload_format": BF16_PAYLOAD, + "num_lanes": settings.num_lanes, + "graph_mode": settings.graph_mode, + "dbo": settings.dbo, + "m2n_paired": False, + }, + ) diff --git a/afd_plugin/ffn_faas/platforms/gpu/backend.py b/afd_plugin/ffn_faas/platforms/gpu/backend.py new file mode 100644 index 000000000..ba9bf830f --- /dev/null +++ b/afd_plugin/ffn_faas/platforms/gpu/backend.py @@ -0,0 +1,186 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""FFN-only routed experts composed with the existing GPU M2N runtime. + +Imported inside the externally launched F process, after the adapter has +validated the launch. No vLLM Worker, upstream ModelRunner, EngineCore, +scheduler, KV cache or complete CausalLM is constructed. Step, graph, DBO and +transport logic are not re-implemented: they belong to +``GPUFFNRoutedExpertsModelRunner`` and ``DeepEPM2NAFDConnector``. +""" + +from __future__ import annotations + +import threading +from argparse import Namespace + +import torch +from vllm.config import set_current_vllm_config + +from afd_plugin.connectors.gpu.deepep_m2n.connector import DeepEPM2NAFDConnector +from afd_plugin.ffn_faas.model_utils import parameter_storage_bytes +from afd_plugin.ffn_faas.platforms.gpu.environment import GPUFFNEnvironment +from afd_plugin.ffn_faas.platforms.gpu.settings import GPUFFNConfig +from afd_plugin.ffn_faas.service import FFNService, ServicePhase +from afd_plugin.v1.worker.ffn_routed_experts_model_runner import ( + GPUFFNRoutedExpertsModelRunner, + RoutedExpertsModel, +) + + +class DeepEPM2NFFNBackend: + """Keep weights and communication owners until receiver exit is proven. + + Implements the ``FFNBackend`` lifecycle over DeepEP M2N: ``prepare`` runs + on the host thread; ``start_receiver``, ``execute_step`` and + ``receiver_exited`` run on the receiver thread; ``close`` then releases + the runtime, the compute groups and the weights, in that order. + ``cancel_receive`` can end an idle receive because the M2N control plane + owns an abortable NCCL group. + """ + + def __init__(self, settings: GPUFFNConfig, args: Namespace) -> None: + """Hold the CPU-side launch facts; native work starts in ``prepare``.""" + self.settings = settings + self.args = args + self.environment = GPUFFNEnvironment(settings) + self.runner: GPUFFNRoutedExpertsModelRunner | None = None + self.model: RoutedExpertsModel | None = None + self._cancel_requested = threading.Event() + + def prepare(self, service: FFNService) -> None: + """Bring this F rank from launch to paired with Attention. + + Runs on the host thread, in the order that lets a bad deployment fail + before the expert weights are read: + + 1. build the F-only ``VllmConfig`` (AFD architecture, Triton MoE pin), + which resolves ``--device-ids``; no device or weight is touched; + 2. select and check the GPU through that assignment, and the + afdEP/DeepGEMM stack; + 3. create the F-only NCCL world and vLLM's model-parallel groups; + 4. construct the routed-experts runner, which runs the ModelRunnerV2 + and CUDA-graph validation and creates the connector (no + collective yet); + 5. load the FFN-only routed-experts shell through vLLM's native loader; + 6. start the connector: the collective M2N handshake with Attention, + the model-info check, the NCCL control group, the lanes' DeepEP + buffers and the dispatch-kernel precompile. + + Phases: ``loading`` through step 5 (with ``loaded_experts`` and + ``total_experts`` after the load), ``connecting`` from step 5 on; the + host marks ``ready`` afterwards. CUDA graphs are not captured here: they + are captured lazily, driven by Attention's warmup and capture + payloads, exactly as on the Worker path. Every owner is assigned + before the next native step, so a failure leaves it reachable. + """ + environment = self.environment + config = environment.build_config(self.args) + environment.initialize_device(config) + environment.initialize_compute(config) + with set_current_vllm_config(config): + self.runner = GPUFFNRoutedExpertsModelRunner(config, environment.device) + connector = self.runner.connector + if not isinstance(connector, DeepEPM2NAFDConnector): + raise TypeError( + f"standalone GPU FFN needs DeepEPM2NAFDConnector, got " + f"{type(connector).__name__}" + ) + self.runner.load_model() + model: RoutedExpertsModel = self.runner.get_model() + self.model = model + # vLLM's loader reports no per-expert progress; this rank's experts + # are all loaded once it returns. + total_experts = sum( + len(self.runner.local_experts[layer_idx].expert_ids) + for layer_idx in self.runner.moe_layer_indices + ) + service.loading_progress(total_experts, total_experts) + if not isinstance(model, torch.nn.Module): + raise TypeError( + f"FFN-only model must be a torch module, got {type(model).__name__}" + ) + parameter_bytes = parameter_storage_bytes(model) + local_expert_ids = self.runner.local_experts[ + self.runner.moe_layer_indices[0] + ].expert_ids + service.phase( + ServicePhase.CONNECTING, + model_loaded=True, + parameter_bytes=parameter_bytes, + local_expert_first=local_expert_ids[0], + local_expert_last=local_expert_ids[-1], + ) + self.runner.initialize_afd_connector() + assert connector.attn_tp_size is not None + assert connector.capacity_tokens_per_rank is not None + assert connector.bucket_ladder is not None + assert connector.network_plan is not None + service.phase( + ServicePhase.CONNECTING, + m2n_paired=True, + attn_tp_size=connector.attn_tp_size, + capacity_tokens_per_rank=connector.capacity_tokens_per_rank, + num_buckets=len(connector.bucket_ladder), + network_hybrid=connector.network_plan.hybrid, + compute_ep_world_size=environment.ep_group.world_size, + ) + + def start_receiver(self) -> None: + """Select the GPU once on the receiving thread, as the Worker does.""" + self.environment.activate_thread() + + def execute_step(self) -> bool: + """Run one control payload's step through the shared GPU runner method. + + The runner ends every step with a device synchronize, so this method + adds no fence of its own. + + Returns ``False`` only when ``cancel_receive`` closed the control + plane while the receive was idle. A closed control plane without a + cancel request, and any other error (including a DeepEP timeout after + a peer failed), raises and fails the service with owners retained. + """ + assert self.runner is not None + return self.runner.execute_connector_driven_step( + cancel_requested=self._cancel_requested.is_set + ) + + def cancel_receive(self) -> None: + """End an idle control receive so a host stop can exit cleanly. + + Records the request first, then closes the connector's control plane, + which aborts its dedicated NCCL control group; a blocked + ``recv_dp_metadata_list`` then raises ``AFDControlPlaneClosedError`` + and ``execute_step`` returns ``False``. The DeepEP lane buffers keep + their own group, so a step already in flight still completes. Same + mechanism as ``AFDFFNWorker.stop_ffn_server_loop``. + """ + self._cancel_requested.set() + if self.runner is not None: + control_plane = self.runner.connector.control_plane + if control_plane is not None: + control_plane.close() + + def receiver_exited(self) -> None: + """Record on the receiver thread that its loop ended cleanly.""" + assert self.runner is not None + self.runner.mark_receiver_exited() + + def close(self) -> None: + """Release CUDA graphs, DeepEP buffers, compute groups and weights. + + Only after the receiver exit was proven. ``runner.shutdown`` resets + the captured graphs, drops their pool, clears vLLM's forward-context + registrations and the model reference, and closes the connector + (control group and lanes); then the environment destroys the F-only + groups, and the last reference to the FFN-only model is dropped. The + connector's gloo bootstrap group outlives this, which only matters + for an in-process restart, which is not supported. + """ + assert self.runner is not None + if not self.runner.receiver_exited: + raise RuntimeError("GPU FFN receiver has not exited; retain resources") + self.runner.shutdown() + self.environment.close() + self.model = None diff --git a/afd_plugin/ffn_faas/platforms/gpu/environment.py b/afd_plugin/ffn_faas/platforms/gpu/environment.py new file mode 100644 index 000000000..845776d3c --- /dev/null +++ b/afd_plugin/ffn_faas/platforms/gpu/environment.py @@ -0,0 +1,204 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Own the F-local GPU, the F-only vLLM config and its compute groups. + +The M2N world (bootstrap, control and lane groups) belongs to the connector. +The vLLM world, DP and EP groups created here span only the ``torchrun`` F +ranks; they exist so vLLM can construct FusedMoE and place experts, and carry +no per-step traffic. +""" + +from __future__ import annotations + +import importlib +import os +from typing import TYPE_CHECKING + +from afd_plugin.ffn_faas.platforms.environment import FFNEnvironmentBase +from afd_plugin.ffn_faas.platforms.gpu.settings import GPUFFNConfig +from afd_plugin.ffn_faas.platforms.launch import reassert_external_launcher + +if TYPE_CHECKING: + from argparse import Namespace + + from vllm.config import VllmConfig + +# afdEP-comm transport and the DeepGEMM grouped GEMMs the connector runs; both +# are imported lazily by the connector, so a missing package would otherwise +# fail only after 27 GiB of experts were read. +AFDEP_TRANSPORT_MODULE = "afdep.deepep" +DEEPGEMM_MODULE = "deep_gemm" +FUSED_EXPERTS_BACKEND = "fused_experts" +# Set by torchrun: rendezvous in its workers join the agent's TCP store. +TORCHELASTIC_AGENT_STORE_ENV = "TORCHELASTIC_USE_AGENT_STORE" +# vLLM's GPU Worker gives DBO two workspace slots. +DBO_WORKSPACE_UBATCHES = 2 + + +class GPUFFNEnvironment(FFNEnvironmentBase[GPUFFNConfig]): + """Never borrow unrelated groups or tear them down on a failed native seam. + + ``FFNEnvironmentBase`` owns the state, launcher checks, ownership flags, + the EP-group check and close-state handling. This class adds the GPU + device, the F-only ``VllmConfig`` (AFD routed-experts architecture and + pinned MoE backend) and vLLM's GPU Worker distributed initializer. + """ + + def __init__(self, settings: GPUFFNConfig) -> None: + """Record the settings; nothing native happens until ``initialize_device``.""" + super().__init__(settings, platform_name="GPU") + + def initialize_device(self, config: VllmConfig) -> None: + """Select this rank's GPU after checking the launcher and native stack. + + Requires the ``torchrun`` environment to still match the settings and + no pre-initialized ``torch.distributed`` or vLLM model-parallel state, + so this process never adopts groups it does not own. Imports afdEP's + DeepEP transport and, for the DeepGEMM expert backend, DeepGEMM, so a + missing package fails before any weight is read. Maps ``LOCAL_RANK`` + to the visible GPU the way vLLM's GPU Worker does: first registers the + ``--device-ids`` assignment that ``config`` resolved, then maps through + it (and ``CUDA_VISIBLE_DEVICES``). Requires SM90 or newer. + """ + self.require_uninitialized() + self.require_launcher_environment() + self.require_no_borrowed_native_groups() + settings = self.settings + import torch + from vllm.platforms import current_platform + from vllm.platforms.interface import set_assigned_physical_gpu_ids + + importlib.import_module(AFDEP_TRANSPORT_MODULE) + expert_backend = settings.connector_extra_config.get("expert_backend") + if expert_backend != FUSED_EXPERTS_BACKEND: + importlib.import_module(DEEPGEMM_MODULE) + # Importable only after the afdEP preflight above. + from afd_plugin.connectors.gpu.deepep_m2n.connector import MIN_SM_MAJOR + + # Same step as vLLM's GPU Worker.init_device: register --device-ids + # before mapping the local rank to a device. + assigned = config.parallel_config.assigned_physical_gpu_ids + if assigned is not None: + if settings.local_rank >= len(assigned): + raise ValueError("LOCAL_RANK is outside the --device-ids list") + set_assigned_physical_gpu_ids(assigned) + visible_index = current_platform.logical_device_id_to_visible_device_id( + settings.local_rank + ) + device = torch.device("cuda", visible_index) + major, minor = torch.cuda.get_device_capability(device) + if major < MIN_SM_MAJOR: + raise RuntimeError( + f"DeepEP M2N needs SM{MIN_SM_MAJOR}0 or newer, got SM{major}{minor}" + ) + torch.accelerator.set_device_index(device) + self._device = device + + def build_config(self, args: Namespace) -> VllmConfig: + """Build the F-only ``VllmConfig`` exactly as the GPU FFN Worker sees it. + + Creates the config from the native engine CLI arguments through + ``EngineArgs.from_cli_args`` and ``create_engine_config`` (with the + AFD plugin's config patches, which accept DBO without a DeepEP + all2all backend), then + re-asserts the ``external_launcher`` executor, swaps in the + FFN-only routed-experts architecture that ``AFDFFNWorker.init_device`` + selects for this connector, and pins vLLM's MoE backend to Triton as + ``AFDFFNWorker._make_model_runner`` does. The pin must happen before + any FusedMoE is constructed: FusedMoE picks its weight layout at + construction, and a reordered layout would be read silently wrong. + """ + from vllm.engine.arg_utils import EngineArgs + + from afd_plugin.model_executor.models.model_utils import get_afd_model_config + from afd_plugin.v1.worker.ffn_routed_experts_model_runner import ( + select_routed_moe_backend, + ) + + config = EngineArgs.from_cli_args(args).create_engine_config() + # The F world is the already launched torchrun world, whatever executor + # platform defaults would pick for an ordinary engine. + reassert_external_launcher(config) + config.model_config = get_afd_model_config( + config.model_config, + device_type="cuda", + routed_experts_connector=True, + ) + select_routed_moe_backend(config) + return config + + def initialize_compute(self, config: VllmConfig) -> None: + """Create the F-only NCCL world, model-parallel groups and workspace. + + Uses ``init_worker_distributed_environment``, the same function vLLM's + GPU Worker calls (batch invariance, custom all-reduce setting, world, + TP/PP/DP/EP groups), with the ``env://`` rendezvous ``torchrun`` + provides. With ``external_launcher`` vLLM keeps the launcher's world + size and rank, so EP spans exactly the F ranks; this is re-checked + because the connector's handshake assumes contiguous expert ranges + over that world. Ownership is recorded before the first collective so + a partial initialization is still torn down by ``close``. Then stops + later rendezvous from joining the torchrun agent's store, so the M2N + world can host its own, and initializes vLLM's workspace manager with + the Worker's micro-batch count (two under DBO). + """ + self.require_compute_initializable() + from vllm.config import set_current_vllm_config + from vllm.distributed.parallel_state import get_ep_group + from vllm.platforms import current_platform + from vllm.v1.worker.gpu_worker import init_worker_distributed_environment + from vllm.v1.worker.workspace import init_workspace_manager + + settings = self.settings + with set_current_vllm_config(config): + self.mark_groups_owned() + init_worker_distributed_environment( + config, + settings.rank, + "env://", + settings.local_rank, + current_platform.dist_backend, + ) + self._ep_group = get_ep_group() + self.validate_ep_group( + error_message="vLLM expert group differs from the F-only world" + ) + # torchrun asks every ``tcp://``/``env://`` rendezvous in its workers + # to join the agent's store. That is right for the F-only world above, + # but the connector's M2N rendezvous at ``--host:--port`` must be + # hosted by M2N world rank 0 (F rank 0), as in a ``vllm serve`` FFN; + # with the flag set, every rank would wait as a client and nobody + # would listen. vLLM's groups are already created, so drop it. + os.environ.pop(TORCHELASTIC_AGENT_STORE_ENV, None) + # Same slot count as Worker._init_workspace_and_model_runner. + init_workspace_manager( + self._device, + DBO_WORKSPACE_UBATCHES if config.parallel_config.enable_dbo else 1, + ) + + def activate_thread(self) -> None: + """Select this rank's GPU on the calling (receiver) thread. + + CUDA's current device is per thread; graph capture and the connector's + lane streams run on the receiver thread. + """ + import torch + + torch.cuda.set_device(self.device) + + def _destroy_groups(self) -> None: + """Destroy the vLLM model-parallel groups and the F-only world. + + Only after the runner released CUDA graphs and DeepEP buffers; a failed + native seam never reaches here, so groups stay owned until process + isolation. + """ + from vllm.distributed.parallel_state import ( + destroy_distributed_environment, + destroy_model_parallel, + ) + from vllm.v1.worker.workspace import reset_workspace_manager + + reset_workspace_manager() + destroy_model_parallel() + destroy_distributed_environment() diff --git a/afd_plugin/ffn_faas/platforms/gpu/settings.py b/afd_plugin/ffn_faas/platforms/gpu/settings.py new file mode 100644 index 000000000..feb713e60 --- /dev/null +++ b/afd_plugin/ffn_faas/platforms/gpu/settings.py @@ -0,0 +1,83 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""CPU-only launch settings for the standalone GPU DeepEP M2N FFN process. + +vLLM's native engine CLI and ``additional_config["afd"]`` own the engine, AFD +topology and connector options; ``GPUFFNConfig`` keeps only the launcher and +shutdown facts plus the options the FFN service itself fixes (payload format, +execution mode and lane count). Connector option values stay with +``DeepEPM2NConfig``. +""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass, field +from pathlib import Path +from types import MappingProxyType +from typing import Any + +from afd_plugin.ffn_faas.platforms.launch import validate_ffn_rank_settings + +# The only qualified payload: BF16 rows against BF16 expert weights. +BF16_PAYLOAD = "bf16" +SINGLE_LANE = 1 +# DBO runs two micro-batches, one DeepEP lane each. +DBO_LANES = 2 +FFN_GRAPH_MODE = "FULL_DECODE_ONLY" +EAGER_MODE = "eager" + + +@dataclass(frozen=True) +class GPUFFNConfig: + """Frozen settings of one F rank in a standalone DeepEP M2N deployment.""" + + model: Path + ffn_ranks: int + rank: int = 0 + local_rank: int = 0 + graph: bool = False + dbo: bool = False + connector_extra_config: Mapping[str, Any] = field( + default_factory=lambda: MappingProxyType({}) + ) + shutdown_timeout: float = 5.0 + + def __post_init__(self) -> None: + """Reject launches no Attention peer could pair with, on the CPU. + + Checks this rank, the shutdown timeout, and the two connector options + the service fixes: ``payload_format`` must be ``bf16`` and an explicit + ``num_lanes`` must equal the lane count ``--enable-dbo`` implies (1 + without, 2 with). Every other connector option is validated by + ``DeepEPM2NConfig.from_mapping`` once the backend runs. + """ + validate_ffn_rank_settings( + ffn_ranks=self.ffn_ranks, + rank=self.rank, + local_rank=self.local_rank, + shutdown_timeout=self.shutdown_timeout, + ) + payload_format = self.connector_extra_config.get("payload_format", BF16_PAYLOAD) + if payload_format != BF16_PAYLOAD: + raise ValueError( + f"GPU standalone FFN supports only payload_format={BF16_PAYLOAD!r}, " + f"got {payload_format!r}" + ) + num_lanes = self.connector_extra_config.get("num_lanes") + if num_lanes is not None and num_lanes != self.num_lanes: + raise ValueError( + f"connector num_lanes={num_lanes!r} disagrees with " + f"{'--enable-dbo' if self.dbo else 'no --enable-dbo'} " + f"({self.num_lanes} lane(s))" + ) + + @property + def num_lanes(self) -> int: + """DeepEP lanes this rank creates: two under DBO, otherwise one.""" + return DBO_LANES if self.dbo else SINGLE_LANE + + @property + def graph_mode(self) -> str: + """Status name of the FFN execution mode.""" + return FFN_GRAPH_MODE if self.graph else EAGER_MODE diff --git a/afd_plugin/ffn_faas/platforms/launch.py b/afd_plugin/ffn_faas/platforms/launch.py new file mode 100644 index 000000000..3bb01c510 --- /dev/null +++ b/afd_plugin/ffn_faas/platforms/launch.py @@ -0,0 +1,114 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Shared launch-time CLI and environment checks for standalone FFN backends. + +Importing this module stays CPU-only; ``create_ffn_engine_parser`` imports the +selected vLLM CLI layer only when the platform adapter builds its parser. +""" + +from __future__ import annotations + +import math +import os +from argparse import ArgumentParser, Namespace +from dataclasses import dataclass +from typing import TYPE_CHECKING + +from afd_plugin.config import AFDConfig, parse_afd_config + +EXTERNAL_LAUNCHER = "external_launcher" + +if TYPE_CHECKING: + from vllm.config import VllmConfig + + +@dataclass(frozen=True) +class FFNLauncherEnvironment: + """One rank's externally launched F-only world coordinates.""" + + rank: int + local_rank: int + world_size: int + + +def read_ffn_launcher_environment() -> FFNLauncherEnvironment: + """Read the ``torchrun`` coordinates of this F-only rank.""" + rank, local_rank, world_size = ( + int(os.environ.get(key, default)) + for key, default in ( + ("RANK", "0"), + ("LOCAL_RANK", "0"), + ("WORLD_SIZE", "1"), + ) + ) + return FFNLauncherEnvironment( + rank=rank, local_rank=local_rank, world_size=world_size + ) + + +def validate_ffn_rank_settings( + *, ffn_ranks: int, rank: int, local_rank: int, shutdown_timeout: float +) -> None: + """Validate the launcher and host-shutdown facts every backend keeps. + + Engine and connector definitions are validated by vLLM and AFD config. + """ + if ffn_ranks <= 0 or not 0 <= rank < ffn_ranks: + raise ValueError("invalid F-only rank or world size") + if local_rank < 0: + raise ValueError("local rank must be nonnegative") + if not math.isfinite(shutdown_timeout) or shutdown_timeout <= 0: + raise ValueError("shutdown timeout must be finite and positive") + + +def create_ffn_engine_parser() -> ArgumentParser: + """Build vLLM's native engine CLI for an externally launched F-only rank.""" + # Selected-platform imports: the Service/host entry stays CPU-only. + from vllm.engine.arg_utils import EngineArgs + from vllm.plugins import load_general_plugins + from vllm.utils.argparse_utils import FlexibleArgumentParser + + load_general_plugins() + parser = EngineArgs.add_cli_args(FlexibleArgumentParser()) + # This process is already launched by torchrun, not a vLLM executor. + parser.set_defaults(distributed_executor_backend=EXTERNAL_LAUNCHER) + return parser + + +def configure_ffn_engine_args( + args: Namespace, *, connector: str +) -> tuple[AFDConfig, FFNLauncherEnvironment]: + """Validate the native engine arguments against the F-only launcher. + + Applies vLLM serve's positional-model precedence, parses the FFN-role AFD + config, requires ``connector``, and requires the launcher ``WORLD_SIZE``, + ``--data-parallel-size`` and ``afd.num_ffn_ranks`` to agree. An omitted + ``--data-parallel-rank`` is taken from ``RANK``. Mutates ``args`` so + ``EngineArgs.from_cli_args`` sees the resolved model and DP rank. + """ + if args.model_tag is not None: + # Match vLLM serve's positional-model precedence. + args.model = args.model_tag + afd = parse_afd_config(args.additional_config, expected_role="ffn") + if afd.connector != connector: + raise ValueError(f"this standalone FFN backend requires {connector}") + launcher = read_ffn_launcher_environment() + if ( + launcher.world_size != afd.num_ffn_ranks + or launcher.world_size != args.data_parallel_size + ): + raise ValueError( + "F-only WORLD_SIZE, --data-parallel-size and " + "additional_config.afd.num_ffn_ranks must agree" + ) + if args.data_parallel_rank is not None and args.data_parallel_rank != launcher.rank: + raise ValueError("--data-parallel-rank must match the F-only launcher RANK") + args.data_parallel_rank = launcher.rank + if args.distributed_executor_backend != EXTERNAL_LAUNCHER: + raise ValueError("standalone FFN requires the external_launcher backend") + return afd, launcher + + +def reassert_external_launcher(config: VllmConfig) -> None: + """Keep an already launched F-only world from borrowing an executor.""" + config.parallel_config.distributed_executor_backend = EXTERNAL_LAUNCHER diff --git a/afd_plugin/ffn_faas/platforms/npu/__init__.py b/afd_plugin/ffn_faas/platforms/npu/__init__.py index c401ffbe2..2dfaa7c31 100644 --- a/afd_plugin/ffn_faas/platforms/npu/__init__.py +++ b/afd_plugin/ffn_faas/platforms/npu/__init__.py @@ -17,11 +17,11 @@ def get_model(model_path: Path, ffn_ranks: int, ffn_rank: int) -> CAMRoutedModel This is a loading boundary, not a model registry or multi-model support claim. Native CAM still requires its existing HF configuration and W4A8 contract. """ - from afd_plugin.ffn_faas.platforms.npu.model import CAMRoutedModel - from afd_plugin.ffn_faas.platforms.npu.v4_model import ( + from afd_plugin.ffn_faas.platforms.ffn_model.deepseek_v4.npu.v4_model import ( V4CheckpointDescription, V4ModelAdapter, ) + from afd_plugin.ffn_faas.platforms.npu.model import CAMRoutedModel adapter = V4ModelAdapter(V4CheckpointDescription.read(model_path)) return CAMRoutedModel(adapter, ffn_ranks, ffn_rank) diff --git a/afd_plugin/ffn_faas/platforms/npu/adapter.py b/afd_plugin/ffn_faas/platforms/npu/adapter.py index c37346ff1..152db9968 100644 --- a/afd_plugin/ffn_faas/platforms/npu/adapter.py +++ b/afd_plugin/ffn_faas/platforms/npu/adapter.py @@ -1,15 +1,20 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project -"""CPU-only Ascend/V4 launch adaptation; native imports happen on preparation.""" +"""Ascend/V4 CLI adaptation; importing this module remains CPU-only.""" from __future__ import annotations -import os from argparse import ArgumentParser, Namespace from functools import partial +from pathlib import Path from typing import TYPE_CHECKING +from afd_plugin.config import AFD_ASYNC_CONNECTOR from afd_plugin.ffn_faas.interfaces import FFNLaunch +from afd_plugin.ffn_faas.platforms.launch import ( + configure_ffn_engine_args, + create_ffn_engine_parser, +) from afd_plugin.ffn_faas.platforms.npu import get_model from afd_plugin.ffn_faas.platforms.npu.settings import FFNConfig @@ -18,88 +23,39 @@ from afd_plugin.ffn_faas.platforms.npu.model import CAMRoutedModel -def create_backend(model: CAMRoutedModel, settings: FFNConfig) -> FFNBackend: +def create_backend( + model: CAMRoutedModel, settings: FFNConfig, args: Namespace +) -> FFNBackend: from vllm_ascend.utils import bootstrap_custom_op_env bootstrap_custom_op_env() from afd_plugin.ffn_faas.platforms.npu.backend import CAMFFNBackend - return CAMFFNBackend(model, settings) + return CAMFFNBackend(model, settings, args) class AscendCAMAdapter: """Keep model, quantization, topology and device options out of the host loop.""" - def add_arguments(self, parser: ArgumentParser) -> None: - parser.add_argument("--host", default="127.0.0.1") - parser.add_argument("--port", type=int, default=29600) - parser.add_argument("--attention-ranks", type=int, required=True) - parser.add_argument("--attention-tp", type=int, required=True) - parser.add_argument("--ffn-ranks", type=int, required=True) - parser.add_argument("--devices", help="allocated logical NPU IDs on this node") - parser.add_argument("--max-num-batched-tokens", type=int, default=2048) - parser.add_argument("--max-num-seqs", type=int, default=256) - parser.add_argument("--full-graph", action="store_true") - parser.add_argument("--ffn-two-stream", action="store_true") - parser.add_argument("--ffn-compute-aic-num", type=int) - parser.add_argument("--ffn-compute-aiv-num", type=int) - parser.add_argument("--ffn-combine-aiv-num", type=int) - parser.add_argument("--hccl-buffer-size", type=int) + def create_parser(self) -> ArgumentParser: + return create_ffn_engine_parser() def configure(self, args: Namespace) -> FFNLaunch: - rank, local_rank, world = ( - int(os.environ.get(key, default)) - for key, default in ( - ("RANK", "0"), - ("LOCAL_RANK", "0"), - ("WORLD_SIZE", "1"), - ) - ) - if world != args.ffn_ranks: - raise ValueError("F-only launcher WORLD_SIZE must equal --ffn-ranks") - local_world = int(os.environ.get("LOCAL_WORLD_SIZE", str(world))) - devices = ( - tuple(map(int, args.devices.split(","))) - if args.devices - else tuple(range(local_world)) - ) - if ( - not 0 <= local_rank < len(devices) - or len(set(devices)) != len(devices) - or min(devices) < 0 - ): - raise ValueError("--devices must supply distinct local NPU IDs") - if os.environ.get("VLLM_USE_V2_MODEL_RUNNER") != "1": - raise ValueError("set VLLM_USE_V2_MODEL_RUNNER=1 on both A and F") - if os.environ.get("AFD_ASYNC_CAM_LAYERED_GMM") != "1": - raise ValueError("standalone V4 requires AFD_ASYNC_CAM_LAYERED_GMM=1") + _, launcher = configure_ffn_engine_args(args, connector=AFD_ASYNC_CONNECTOR) settings = FFNConfig( - model=args.model, - attention_ranks=args.attention_ranks, - attention_tp=args.attention_tp, - ffn_ranks=args.ffn_ranks, - rank=rank, - local_rank=local_rank, - device=devices[local_rank], - host=args.host, - port=args.port, - max_num_batched_tokens=args.max_num_batched_tokens, - max_num_seqs=args.max_num_seqs, - full_graph=args.full_graph, - ffn_two_stream=args.ffn_two_stream, - ffn_compute_aic_num=args.ffn_compute_aic_num, - ffn_compute_aiv_num=args.ffn_compute_aiv_num, - ffn_combine_aiv_num=args.ffn_combine_aiv_num, - hccl_buffer_size=args.hccl_buffer_size, - shutdown_timeout=args.shutdown_timeout, + model=Path(args.model), + ffn_ranks=launcher.world_size, + rank=launcher.rank, + local_rank=launcher.local_rank, + shutdown_timeout=args.receiver_shutdown_timeout, ) - model = get_model(args.model, settings.ffn_ranks, settings.rank) + model = get_model(settings.model, settings.ffn_ranks, settings.rank) return FFNLaunch( - rank=rank, - local_rank=local_rank, + rank=launcher.rank, + local_rank=launcher.local_rank, model_path=str(model.root), backend_name="CAMAsyncAFDConnector", shutdown_timeout=settings.shutdown_timeout, - backend_factory=partial(create_backend, model, settings), + backend_factory=partial(create_backend, model, settings, args), initial_status={"cam_paired": False}, ) diff --git a/afd_plugin/ffn_faas/platforms/npu/backend.py b/afd_plugin/ffn_faas/platforms/npu/backend.py index 8f1e0a5b9..d1180e835 100644 --- a/afd_plugin/ffn_faas/platforms/npu/backend.py +++ b/afd_plugin/ffn_faas/platforms/npu/backend.py @@ -8,58 +8,54 @@ from __future__ import annotations +from argparse import Namespace + import torch import torch_npu # noqa: F401 from vllm.config import set_current_vllm_config from vllm.engine.arg_utils import EngineArgs -from afd_plugin.compat.npu import fix_all2all_backend_for_afd +from afd_plugin.compat.npu import ( + fail_if_unsupported_npu_afd_features, + fix_all2all_backend_for_afd, +) from afd_plugin.config import parse_afd_config from afd_plugin.connectors.npu.async_cam import CAMAsyncAFDConnector +from afd_plugin.ffn_faas.platforms.launch import reassert_external_launcher from afd_plugin.ffn_faas.platforms.npu.environment import AscendFFNEnvironment from afd_plugin.ffn_faas.platforms.npu.model import CAMRoutedModel from afd_plugin.ffn_faas.platforms.npu.settings import FFNConfig from afd_plugin.ffn_faas.service import FFNService, ServicePhase from afd_plugin.v1.worker.npu.async_ffn_model_runner import AsyncNPUFFNModelRunner +from afd_plugin.validation import validate_npu_model_runner_v2_config class CAMFFNBackend: """Keep weights and communication owners until receiver exit is proven.""" - def __init__(self, model: CAMRoutedModel, settings: FFNConfig) -> None: + def __init__( + self, model: CAMRoutedModel, settings: FFNConfig, args: Namespace + ) -> None: self.settings = settings + self.args = args self.environment = AscendFFNEnvironment(settings) self.model = model self.runtime: AsyncNPUFFNModelRunner | None = None self.connector: CAMAsyncAFDConnector | None = None def prepare(self, service: FFNService) -> None: - self.environment.initialize_device() - args = EngineArgs( - model=str(self.settings.model), - trust_remote_code=True, - dtype="bfloat16", - quantization="ascend", - enforce_eager=not self.settings.full_graph, - compilation_config={"cudagraph_mode": "FULL"} - if self.settings.full_graph - else {}, - tensor_parallel_size=1, - data_parallel_size=self.settings.ffn_ranks, - data_parallel_rank=self.settings.rank, - enable_expert_parallel=True, - distributed_executor_backend="external_launcher", - # F owns a dispatch envelope, not a request-context/KV capacity. - max_model_len=self.settings.max_num_batched_tokens, - max_num_batched_tokens=self.settings.max_num_batched_tokens, - max_num_seqs=self.settings.max_num_seqs, - additional_config=self.settings.additional_config(), - ) + args = EngineArgs.from_cli_args(self.args) config = args.create_engine_config() - # Platform configuration can choose an executor for ordinary Workers; - # this loader belongs to the already launched F-only torchrun world. - config.parallel_config.distributed_executor_backend = "external_launcher" + # Keep the externally launched F-only ownership after platform defaults. + reassert_external_launcher(config) fix_all2all_backend_for_afd(config) + validate_npu_model_runner_v2_config( + config, expected_role="ffn", device_type="npu" + ) + fail_if_unsupported_npu_afd_features(config) + if config.speculative_config is not None: + raise ValueError("standalone V4 loads main-model routed layers only") + self.environment.initialize_device(config) self.environment.initialize_compute(config) with set_current_vllm_config(config): self.model.load( @@ -72,7 +68,7 @@ def prepare(self, service: FFNService) -> None: ) self.connector = CAMAsyncAFDConnector( rank=self.settings.rank, - local_rank=self.settings.device, + local_rank=self.environment.device.index, vllm_config=config, afd_config=parse_afd_config(config, expected_role="ffn"), role_rank=self.settings.rank, @@ -83,8 +79,8 @@ def prepare(self, service: FFNService) -> None: self.runtime.attach_model( self.model, parameter_bytes=self.model.parameter_bytes ) - # Reuse paired signatures, warmup, capture, Graph pools and two-slot - # ready/done/FIFO ownership from the existing CAM implementation. + # Reuse paired signatures, warmup, capture and Graph pools from + # the existing CAM implementation. self.runtime.initialize_afd_connector() self.environment.bind_host_threads(config) service.phase( @@ -97,12 +93,25 @@ def start_receiver(self) -> None: """Select the device once on the receiving thread, as the Worker does.""" self.environment.activate_thread() - def execute_step(self) -> None: + def execute_step(self) -> bool: + """Run one CAM submission batch; a CAM receive never reports a cancel. + + The return value follows ``FFNBackend.execute_step``: CAM cannot end a + pending receive, so a completed batch always returns ``True``. + """ assert self.runtime is not None self.runtime.execute_connector_driven_step() - if not self.settings.full_graph and not self.settings.ffn_two_stream: + if self.runtime.effective_mode == "eager": # Preserve the mature eager loop's completion fence. torch.npu.synchronize() + return True + + def cancel_receive(self) -> None: + """No-op: a pending CAM receive cannot be cancelled. + + The host join then times out, reports the pending receive and retains + every owner, exactly as before ``cancel_receive`` existed. + """ def receiver_exited(self) -> None: assert self.runtime is not None diff --git a/afd_plugin/ffn_faas/platforms/npu/environment.py b/afd_plugin/ffn_faas/platforms/npu/environment.py index 952e59acf..b7e832310 100644 --- a/afd_plugin/ffn_faas/platforms/npu/environment.py +++ b/afd_plugin/ffn_faas/platforms/npu/environment.py @@ -9,71 +9,60 @@ from __future__ import annotations import logging -import os from datetime import timedelta from typing import TYPE_CHECKING +from afd_plugin.ffn_faas.platforms.environment import FFNEnvironmentBase from afd_plugin.ffn_faas.platforms.npu.settings import ( FFNConfig, validate_static_experts, ) if TYPE_CHECKING: - import torch from vllm.config import VllmConfig - from vllm.distributed.parallel_state import GroupCoordinator COMPUTE_GROUP_TIMEOUT_SECONDS = 180 -class AscendFFNEnvironment: - """Never borrow unrelated groups or tear them down on a failed native seam.""" +class AscendFFNEnvironment(FFNEnvironmentBase[FFNConfig]): + """Never borrow unrelated groups or tear them down on a failed native seam. + + ``FFNEnvironmentBase`` owns the state, launcher checks, ownership flags, + the EP-group check and close-state handling. This class owns Ascend + device selection, custom operators, HCCL groups and CPU binding. + """ def __init__(self, settings: FFNConfig) -> None: - self.settings = settings - self._device: torch.device | None = None - self._ep_group: GroupCoordinator | None = None - self._groups_owned = False - - @property - def device(self) -> torch.device: - if self._device is None: - raise RuntimeError("Ascend FFN device has not been initialized") - return self._device - - @property - def ep_group(self) -> GroupCoordinator: - if self._ep_group is None: - raise RuntimeError("Ascend FFN expert group has not been initialized") - return self._ep_group - - def initialize_device(self) -> None: - if self._groups_owned or self._device is not None: - raise RuntimeError("Ascend FFN environment is already initialized") + super().__init__(settings, platform_name="Ascend") + + def initialize_device(self, config: VllmConfig) -> None: + self.require_uninitialized() + self.require_launcher_environment() settings = self.settings - try: - world, rank = int(os.environ["WORLD_SIZE"]), int(os.environ["RANK"]) - except (KeyError, ValueError) as error: - raise ValueError( - "Ascend FFN requires an external F-only launcher" - ) from error - if world != settings.ffn_ranks or rank != settings.rank: - raise ValueError("Ascend FFN environment differs from the F-only launcher") from vllm_ascend.utils import bootstrap_custom_op_env bootstrap_custom_op_env() import torch import torch_npu # noqa: F401 - from vllm.distributed.parallel_state import model_parallel_is_initialized + from vllm.platforms import current_platform + from vllm.platforms.interface import set_assigned_physical_gpu_ids from vllm_ascend.utils import enable_custom_op from afd_plugin.compat.npu.ops import ensure_cam_async_ops_available - if torch.distributed.is_initialized() or model_parallel_is_initialized(): - raise RuntimeError("Ascend FFN cannot borrow existing native groups") + self.require_no_borrowed_native_groups() ensure_cam_async_ops_available() - torch.npu.set_device(settings.device) - self._device = torch.device(f"npu:{settings.device}") + # Reuse the native Worker's mapping for --device-ids and visibility. + assigned = config.parallel_config.assigned_physical_gpu_ids + if assigned is not None: + if settings.local_rank >= len(assigned): + raise ValueError("LOCAL_RANK is outside the assigned device list") + set_assigned_physical_gpu_ids(assigned) + index = current_platform.logical_device_id_to_visible_device_id( + settings.local_rank + ) + self._device = torch.device(f"npu:{index}") + torch.npu.set_device(self._device) # Native NPUModelRunner v1 installs this packing policy at import time. # This service does not import that runner; native W4A8 still needs NZ. torch.npu.config.allow_internal_format = True @@ -81,8 +70,7 @@ def initialize_device(self) -> None: raise RuntimeError("native V4 W4A8 custom operators are unavailable") def initialize_compute(self, config: VllmConfig) -> None: - if self._device is None or self._groups_owned: - raise RuntimeError("Ascend FFN compute initialization requires its device") + self.require_compute_initializable() settings = self.settings parallel = config.parallel_config validate_static_experts( @@ -106,11 +94,11 @@ def initialize_compute(self, config: VllmConfig) -> None: register_ascend_customop(config) init_ascend_config(config) # Ownership starts before a collective can partially initialize. - self._groups_owned = True + self.mark_groups_owned() init_distributed_environment( world_size=settings.ffn_ranks, rank=settings.rank, - local_rank=settings.device, + local_rank=self.device.index, distributed_init_method="env://", backend="hccl", timeout=timedelta(seconds=COMPUTE_GROUP_TIMEOUT_SECONDS), @@ -122,11 +110,9 @@ def initialize_compute(self, config: VllmConfig) -> None: ) init_ascend_model_parallel(parallel) self._ep_group = get_ep_group() - if ( - self._ep_group.world_size != settings.ffn_ranks - or self._ep_group.rank_in_group != settings.rank - ): - raise ValueError("native V4 expert group differs from the F world") + self.validate_ep_group( + error_message="native V4 expert group differs from the F world" + ) def activate_thread(self) -> None: import torch @@ -151,7 +137,7 @@ def bind_host_threads(self, config: VllmConfig) -> None: bind_cpus( self.settings.local_rank, npu_id=current_platform.device_id_to_physical_device_id( - self.settings.device + self.settings.local_rank ), migrate_memory=not (engram is not None and engram.cpu_offload), ) @@ -162,10 +148,8 @@ def bind_host_threads(self, config: VllmConfig) -> None: error, ) - def close(self) -> None: + def _destroy_groups(self) -> None: """Only after CAM/executor completion; failures retain parameter owners.""" - if not self._groups_owned: - return from vllm.distributed.parallel_state import ( destroy_distributed_environment, destroy_model_parallel, @@ -175,6 +159,3 @@ def close(self) -> None: destroy_ascend_model_parallel() destroy_model_parallel() destroy_distributed_environment() - self._groups_owned = False - self._ep_group = None - self._device = None diff --git a/afd_plugin/ffn_faas/platforms/npu/model.py b/afd_plugin/ffn_faas/platforms/npu/model.py index 8209002d6..efbb91bb2 100644 --- a/afd_plugin/ffn_faas/platforms/npu/model.py +++ b/afd_plugin/ffn_faas/platforms/npu/model.py @@ -14,6 +14,7 @@ from pathlib import Path from typing import TYPE_CHECKING, Protocol +from afd_plugin.ffn_faas.model_utils import parameter_storage_bytes from afd_plugin.ffn_faas.platforms.npu.settings import validate_native_expert_map if TYPE_CHECKING: @@ -166,12 +167,7 @@ def _load_native( if tuple(view.layer_idx for view in views) != adapter.layer_indices: raise ValueError("native views do not match the routed model layers") self._views = views - storage_sizes: dict[tuple[str, int], int] = {} - for parameter in owner.parameters(): - storage = parameter.untyped_storage() - key = (str(parameter.device), storage.data_ptr()) - storage_sizes[key] = storage.nbytes() - self._parameter_bytes = sum(storage_sizes.values()) + self._parameter_bytes = parameter_storage_bytes(owner) self._loaded = True def get_experts_layer_indices(self) -> tuple[int, ...]: diff --git a/afd_plugin/ffn_faas/platforms/npu/settings.py b/afd_plugin/ffn_faas/platforms/npu/settings.py index cb78d996b..2316f6e51 100644 --- a/afd_plugin/ffn_faas/platforms/npu/settings.py +++ b/afd_plugin/ffn_faas/platforms/npu/settings.py @@ -4,98 +4,29 @@ from __future__ import annotations -import math from dataclasses import dataclass from pathlib import Path -from afd_plugin.config import AFD_ASYNC_CONNECTOR +from afd_plugin.ffn_faas.platforms.launch import validate_ffn_rank_settings @dataclass(frozen=True) class FFNConfig: model: Path - attention_ranks: int - attention_tp: int ffn_ranks: int rank: int = 0 local_rank: int = 0 - device: int = 0 - host: str = "127.0.0.1" - port: int = 29600 - max_num_batched_tokens: int = 2048 - max_num_seqs: int = 256 - full_graph: bool = False - ffn_two_stream: bool = False - ffn_compute_aic_num: int | None = None - ffn_compute_aiv_num: int | None = None - ffn_combine_aiv_num: int | None = None - hccl_buffer_size: int | None = None shutdown_timeout: float = 5.0 def __post_init__(self) -> None: - positive = ( - self.attention_ranks, - self.attention_tp, - self.ffn_ranks, - self.max_num_batched_tokens, - self.max_num_seqs, + # Engine/connector definitions and validation belong to vLLM/AFD. + # Only external-launcher and host-shutdown facts live here. + validate_ffn_rank_settings( + ffn_ranks=self.ffn_ranks, + rank=self.rank, + local_rank=self.local_rank, + shutdown_timeout=self.shutdown_timeout, ) - if any(value <= 0 for value in positive): - raise ValueError( - "rank counts, TP and scheduler capacities must be positive" - ) - if self.attention_ranks % self.attention_tp: - raise ValueError("Attention ranks must be divisible by Attention TP") - if self.max_num_batched_tokens % self.attention_tp: - raise ValueError( - "the CAM dispatch envelope must be divisible by Attention TP" - ) - if not 0 <= self.rank < self.ffn_ranks or min(self.local_rank, self.device) < 0: - raise ValueError("invalid F rank or local device") - if not self.host or not 0 < self.port < 65536: - raise ValueError("CAM rendezvous requires a host and a valid port") - if not math.isfinite(self.shutdown_timeout) or self.shutdown_timeout <= 0: - raise ValueError("shutdown timeout must be finite and positive") - optional_positive = ( - self.ffn_compute_aic_num, - self.ffn_compute_aiv_num, - self.ffn_combine_aiv_num, - self.hccl_buffer_size, - ) - if any(value is not None and value <= 0 for value in optional_positive): - raise ValueError("explicit CAM core and HCCL budgets must be positive") - - def additional_config(self) -> dict: - extra = { - "dynamicQuant": 1, - "attn_ranks_per_dp": self.attention_tp, - "ffn_two_stream": self.ffn_two_stream, - } - for name, value in ( - ("ffn_compute_aic_num", self.ffn_compute_aic_num), - ("ffn_compute_aiv_num", self.ffn_compute_aiv_num), - ("ffn_combine_aiv_num", self.ffn_combine_aiv_num), - ("hccl_buffer_size", self.hccl_buffer_size), - ): - if value is not None: - extra[name] = value - return { - "afd": { - "connector": AFD_ASYNC_CONNECTOR, - "role": "ffn", - "async_dp": True, - "host": self.host, - "port": self.port, - "num_attention_ranks": self.attention_ranks, - "num_ffn_ranks": self.ffn_ranks, - "compute_gate_on_attention": True, - "connector_extra_config": extra, - }, - "enable_prefill_mc2": False, - "enable_force_eplb": False, - "enable_cpu_binding": True, - "ascend_fusion_config": {"fusion_ops_gmmswigluquant": True}, - } def validate_static_experts( diff --git a/afd_plugin/ffn_faas/runner.py b/afd_plugin/ffn_faas/runner.py index 199293ff5..b7b59aba7 100644 --- a/afd_plugin/ffn_faas/runner.py +++ b/afd_plugin/ffn_faas/runner.py @@ -35,11 +35,19 @@ def __init__( self._failure: BaseException | None = None def _serve(self) -> None: + """Receiver-thread loop: one backend step per control payload. + + The loop ends when the host requests a stop between steps, or when + ``execute_step()`` reports that ``cancel_receive()`` ended an idle + receive. Only then does the backend record its receiver exit, which is + the proof ``close()`` requires before releasing native resources. + """ assert self.backend is not None try: self.backend.start_receiver() while not self.service.stop_requested.is_set(): - self.backend.execute_step() + if not self.backend.execute_step(): + break self.service.completed_batches() self.backend.receiver_exited() except BaseException as error: @@ -48,6 +56,16 @@ def _serve(self) -> None: self._done.set() def run(self) -> None: + """Prepare the backend, serve until stopped, then release or retain. + + The backend is constructed before native initialization so a failed + preparation keeps its owners reachable. After READY the receiver runs + on a daemon thread; on a host stop the backend is asked to cancel an + idle receive, then the thread is joined with ``shutdown_timeout``. + Resources are closed only after a clean receiver exit; any failure, + including a receive that does not end, marks the service FAILED and + re-raises with every owner retained for process isolation. + """ try: if self.service.stop_requested.is_set(): self.service.phase(ServicePhase.STOPPED, resources_released=True) @@ -68,11 +86,17 @@ def run(self) -> None: while not self.service.stop_requested.wait(HOST_STATUS_POLL_SECONDS): if self._done.is_set(): break + if self.service.stop_requested.is_set() and not self._done.is_set(): + # The receiver may be blocked in an idle receive that only + # the backend can end; a backend without cancellation keeps + # it pending and the join below reports it. + self.backend.cancel_receive() self._thread.join(timeout=self.shutdown_timeout) if self._thread.is_alive(): raise RuntimeError( "FFN receive is pending; retain owners and isolate the " - "deployment. A host stop cannot cancel backend device work." + "deployment. The backend did not end its receive or " + "in-flight device work within the shutdown timeout." ) if self._failure is not None: raise RuntimeError( diff --git a/afd_plugin/model_executor/checkpoint_paths.py b/afd_plugin/model_executor/checkpoint_paths.py new file mode 100644 index 000000000..f433b526e --- /dev/null +++ b/afd_plugin/model_executor/checkpoint_paths.py @@ -0,0 +1,19 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""CPU-safe checkpoint path parsing shared by AFD model wrappers.""" + +from __future__ import annotations + + +def weight_layer_path(name: str) -> tuple[int, str, tuple[str, ...]] | None: + """Return ``(layer index, stage, remainder)`` for a decoder path.""" + parts = name.split(".") + for marker_idx, part in enumerate(parts[:-2]): + if part != "layers": + continue + try: + layer_idx = int(parts[marker_idx + 1]) + except ValueError: + continue + return layer_idx, parts[marker_idx + 2], tuple(parts[marker_idx + 3 :]) + return None diff --git a/afd_plugin/model_executor/models/model_utils.py b/afd_plugin/model_executor/models/model_utils.py index 5a0518475..226c262b9 100644 --- a/afd_plugin/model_executor/models/model_utils.py +++ b/afd_plugin/model_executor/models/model_utils.py @@ -5,7 +5,7 @@ from __future__ import annotations from copy import deepcopy -from typing import TYPE_CHECKING, Literal +from typing import TYPE_CHECKING, Final, Literal from afd_plugin import _MODEL_REGISTRATIONS, _QWEN3_5_MODEL_REGISTRATIONS from afd_plugin.config import AFD_ASYNC_CONNECTOR, parse_afd_config @@ -14,6 +14,11 @@ from vllm.config import ModelConfig, VllmConfig +_ROUTED_EXPERTS_MODEL_ARCHITECTURES: Final[dict[str, str]] = { + "Qwen3MoeForCausalLM": "AFDQwen3MoeRoutedExperts", +} + + def has_afd_model_registration(model_config: ModelConfig) -> bool: """Return whether the model resolves to a registered AFD implementation.""" @@ -27,6 +32,7 @@ def get_afd_model_config( model_config: ModelConfig, *, device_type: Literal["cuda", "npu"], + routed_experts_connector: bool = False, ) -> ModelConfig: """Return a model config that resolves to an AFD model implementation.""" @@ -44,7 +50,14 @@ def get_afd_model_config( # automatically. vLLM Ascend uses that identity to distinguish # text models from multimodal models. afd_model_config = deepcopy(model_config) - afd_model_config.hf_config.architectures = [f"AFD{model_arch}"] + routed_architecture = ( + _ROUTED_EXPERTS_MODEL_ARCHITECTURES.get(model_arch) + if routed_experts_connector and device_type == "cuda" + else None + ) + afd_model_config.hf_config.architectures = [ + routed_architecture or f"AFD{model_arch}" + ] return afd_model_config return model_config diff --git a/afd_plugin/model_executor/models/npu/deepseek_v4.py b/afd_plugin/model_executor/models/npu/deepseek_v4.py index a025717bd..a7d45536d 100644 --- a/afd_plugin/model_executor/models/npu/deepseek_v4.py +++ b/afd_plugin/model_executor/models/npu/deepseek_v4.py @@ -24,7 +24,11 @@ from vllm.model_executor.layers.linear import ReplicatedLinear from vllm.sequence import IntermediateTensors -from afd_plugin.config import AFD_ASYNC_CONNECTOR, parse_afd_config +from afd_plugin.config import ( + AFD_ASYNC_CONNECTOR, + CAMP2P_CONNECTOR, + parse_afd_config, +) from afd_plugin.connectors import ( AFDExpertRoutingSpec, AFDTransferContext, @@ -311,6 +315,72 @@ def forward( return output +class AFDDeepseekV4SyncAttentionGateRemoteMoE(AFDDeepseekV4AttentionGateRemoteMoE): + """DSV4 gate shell that keeps the router on Attention for the sync connector. + + Inheritance is deliberate: the gate construction, the FP32 precast weight, the + hash-layer table and the shared-expert module are identical to the async shell, + and only the transport differs. Duplicating that construction would let the two + shells drift on routing parameters. + + The forward differs in three ways, all forced by the sync contract: + + - routing travels on the A2E wire slot instead of a CAM dispatch payload, so the + send and the receive are separate calls; + - the shared experts run *between* those calls, which is the whole point of + moving them here: their work overlaps the FFN side's routed work; + - no DBO yield, because the first prototype keeps the sync path free of the + scheduling interaction the async shell needs. + + Reusing ``RemoteFFNProxy._send_and_receive`` is not possible here: it issues the + send and the receive back to back, leaving no window for the shared experts. + """ + + def forward( + self, + hidden_states: torch.Tensor, + input_ids: torch.Tensor | None = None, + hidden_states_fp32: torch.Tensor | None = None, + ) -> torch.Tensor: + from afd_plugin.model_executor.models.npu import deepseek_v4_attention_gate + + topk_weights, topk_ids = deepseek_v4_attention_gate.compute_attention_gate_topk( + self, + hidden_states, + input_ids=input_ids, + hidden_states_fp32=hidden_states_fp32, + ) + afd_metadata = get_afd_metadata_from_forward_context() + if afd_metadata is None: + raise RuntimeError("Sync Attention-side gate requires AFD forward metadata") + context = AFDTransferContext( + metadata=AFDTransferMetadata.create_attention_metadata( + layer_idx=self.layer_idx, + stage_idx=afd_metadata.stage_idx, + seq_len=int(hidden_states.shape[0]), + ) + ) + # Routing shares the wire slot with token ids, so this send carries routing + # only. The FFN side is configured with recv_routing for the same reason. + afd_metadata.connector.send_attn_output( + hidden_states, + context, + routing=(topk_ids, topk_weights), + ) + shared_output = ( + self.shared_experts(hidden_states) + if self.shared_experts is not None + else None + ) + output = afd_metadata.connector.recv_ffn_output( + ref_tensor=hidden_states, + ubatch_idx=afd_metadata.stage_idx, + ) + if shared_output is not None: + output = output + shared_output + return output + + class AFDDeepseekV4DecoderLayer(native.DeepseekV4DecoderLayer): """Role-local DSV4 decoder layer. @@ -347,6 +417,13 @@ def __init__( self.config = config self.afd_role = afd_config.role self.compute_gate_on_attention = bool(afd_config.compute_gate_on_attention) + # Which wire protocol this layer expects, decided from the configuration. + # The payload only validates that expectation; it must not pick the path, + # because a missing or mixed payload would otherwise silently reroute. + self.uses_external_routing = bool( + afd_config.compute_gate_on_attention + and afd_config.connector == CAMP2P_CONNECTOR + ) self.is_moe_layer = True # ### PATCH END: AFD decoder metadata cache_config = vllm_config.cache_config @@ -397,7 +474,14 @@ def __init__( need_gather_q_kv=self.use_sequence_parallel_moe and self.enable_dsa_cp, ) if self.compute_gate_on_attention: - self.mlp = AFDDeepseekV4AttentionGateRemoteMoE( + # Both shells keep the router on Attention; only the transport + # differs. Choosing by connector keeps the async CAM path unchanged. + gate_shell = ( + AFDDeepseekV4AttentionGateRemoteMoE + if afd_config.connector == AFD_ASYNC_CONNECTOR + else AFDDeepseekV4SyncAttentionGateRemoteMoE + ) + self.mlp = gate_shell( config=config, layer_idx=self.layer_idx, prefix=f"{prefix}.mlp", @@ -410,7 +494,10 @@ def __init__( self.self_attn = native.PPMissingLayer() _refresh_ascend_fused_moe() moe_config = config - if afd_config.connector == AFD_ASYNC_CONNECTOR: + if self.compute_gate_on_attention: + # The shared experts live wherever the router lives, because they + # run in the same window as the gate. That is true for either + # connector, not just for the async one. moe_config = copy(config) # HF validates this public field as int, but pinned native # MoE uses only None to omit shared construction (0 builds a @@ -427,6 +514,23 @@ def __init__( raise ValueError(f"Unsupported AFD role: {afd_config.role!r}") # ### PATCH END: role-local Attention and remote experts + # The external-routing component is built here, while a vLLM config context + # exists, so the first forward onto MC2 never constructs a comm. + self.afd_external_routing_moe = None + # Synchronous external routing only: the async CAM connector keeps its own + # FFN path, even though it also computes the gate on the Attention side. + if ( + self.uses_external_routing + and afd_config.role == _FFN_ROLE + and isinstance(self.mlp, native.DeepseekV4MoE) + ): + from afd_plugin.model_executor.models.npu import ( + deepseek_v4_external_routing_moe as external_routing_moe, + ) + + self.afd_external_routing_moe = external_routing_moe.AFDExternalRoutingMoE( + self.mlp.experts.routed_experts, + ) self.input_layernorm = native.RMSNorm(config.hidden_size, eps=self.norm_eps) self.post_attention_layernorm = native.RMSNorm( @@ -448,17 +552,90 @@ def __init__( def compute_ffn_output( self, hidden_states: torch.Tensor, + *, + group_list: torch.Tensor | None = None, + topk_ids: torch.Tensor | None = None, + topk_weights: torch.Tensor | None = None, **_: Any, ) -> torch.Tensor: if not isinstance(self.mlp, native.DeepseekV4MoE): raise RuntimeError( "DSV4 compute_ffn_output requires the FFN role to own the MoE" ) + # The configuration decides which protocol this layer speaks; the payload + # only has to satisfy it. Letting the payload choose would turn a missing or + # mixed payload into a silent reroute to a different computation. + routing_payload_present = topk_ids is not None or topk_weights is not None + + if self.uses_external_routing: + # CAM dispatch data belongs to the async connector. Keep rejecting it + # instead of letting ``**_`` absorb it silently: a layer speaks one + # protocol, and a mixed payload would otherwise run half-initialized. + if group_list is not None: + raise RuntimeError( + "DSV4 sync external routing received CAM dispatch data as well; " + "a layer speaks one protocol", + ) + if topk_ids is None or topk_weights is None: + raise RuntimeError( + "DSV4 sync external routing requires both topk_ids and " + f"topk_weights; got ids={topk_ids is not None}, " + f"weights={topk_weights is not None}", + ) + return self._compute_external_routing_ffn( + hidden_states, + topk_ids=topk_ids, + topk_weights=topk_weights, + ) + + # The async CAM FFN runs through the layered executor upstream; this method + # only serves the synchronous external-routing path and the target-hash MoE. if self.compute_gate_on_attention: raise RuntimeError("Async CAM FFN requires the layered executor") + if routing_payload_present: + raise RuntimeError( + "DSV4 gate-on-FFN received Attention-side routing; the FFN side owns " + "the gate for this configuration and must compute it itself", + ) # Target Hash MoE takes explicit token IDs from the CAM transfer. return self.mlp(hidden_states, input_ids=get_forward_context().input_ids) + def _compute_external_routing_ffn( + self, + hidden_states: torch.Tensor, + *, + topk_ids: torch.Tensor, + topk_weights: torch.Tensor, + ) -> torch.Tensor: + """Run the FFN on routing the Attention side computed. + + Mirrors ``AscendRoutedExperts.forward`` from the pinned vLLM-Ascend tag, with + two differences: the top-k comes from the transport instead of a local gate, + and both the padding/gathering of the routing and the scale it carries are + supplied by the AFD prepare stage. + + Routed scaling needs no separate step here. The Attention-side selector folds + ``routed_scaling_factor`` into ``topk_weights``, and ``apply`` consumes those + weights directly in the token combine, so the factor is applied exactly once. + Applying it again here would scale the routed output a second time. + + Both ALLGATHER and MC2 are wired. They share this flow because both AFD prepare + implementations pad the routing, and MC2's additionally splits it for the TP + sharding before returning the aligned routing. What differs is only whether + finalize reduces across ranks. + + Any other communication method is rejected rather than silently taking one of + these routes. + """ + + component = self.afd_external_routing_moe + assert component is not None, "external routing component was not built at init" + return component.forward( + hidden_states, + topk_ids=topk_ids, + topk_weights=topk_weights, + ) + @native.support_torch_compile class AFDDeepseekV4Model(native.DeepseekV4Model): @@ -662,7 +839,8 @@ def compute_experts_output( router_logits: torch.Tensor, ) -> torch.Tensor: raise NotImplementedError( - "DSV4 AFD does not support compute_gate_on_attention yet" + "DSV4 AFD never returns gate outputs through this method; the " + "Attention-side gate keeps its top-k inside the gate shell" ) @@ -676,9 +854,20 @@ class AFDDeepseekV4ForCausalLM(native.AscendDeepseekV4ForCausalLM): # installs them in the forward context before the FFN compute. afd_requires_input_ids = True + # The wire slot carries either token ids or Attention-side routing, never both. + # When the router lives on Attention with the sync connector, the FFN side needs + # the routing instead of the ids, and the runner must ask for it by name. + afd_requires_attention_routing = False + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: self.afd_config = parse_afd_config(vllm_config, validate=False) self.afd_role = self.afd_config.role + # Per-instance, because the same model class serves both connectors and the + # async CAM shell carries its routing inside its own dispatch payload. + self.afd_requires_attention_routing = bool( + self.afd_config.compute_gate_on_attention + and self.afd_config.connector == CAMP2P_CONNECTOR + ) super().__init__(vllm_config=vllm_config, prefix=prefix) # vLLM 5df63a7 model_loader/utils.py invokes this model hook after quant @@ -733,11 +922,14 @@ def compute_experts_output( def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: # Only the gate-on-Attention configuration gives Attention a router; with # the gate on FFN its MoE slot is a parameter-free transfer shell. + # The shared experts follow the router, so their weights load on Attention + # for either connector; keying this on the async connector alone would + # construct the modules on the sync path without ever loading them. role_weights = _iter_role_weights( weights, role=self.afd_role, attn_owns_gate=bool(self.afd_config.compute_gate_on_attention), - attention_shared_experts=self.afd_config.connector == AFD_ASYNC_CONNECTOR, + attention_shared_experts=bool(self.afd_config.compute_gate_on_attention), ) loaded = super().load_weights(role_weights) return loaded diff --git a/afd_plugin/model_executor/models/npu/deepseek_v4_external_routing_moe.py b/afd_plugin/model_executor/models/npu/deepseek_v4_external_routing_moe.py new file mode 100644 index 000000000..c47a9eb28 --- /dev/null +++ b/afd_plugin/model_executor/models/npu/deepseek_v4_external_routing_moe.py @@ -0,0 +1,350 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""AFD-owned MoE execution for the external-routing FFN path. + +Attention computes the gate when ``compute_gate_on_attention`` is on, so the FFN +side receives routing for its own input rows instead of computing it. Upstream +prepares the routing inside the prepare stage because upstream owns the gate; +here the routing arrives as an explicit input, is aligned by the same padding, +gathering and sharding rules as the activation, and comes back as an explicit +result. + +The AFD prepare/finalize subclasses are selected through AFD comm implementations +that only override the factory method, so ordinary execution keeps the upstream +classes. The MoE component owns the per-token state that expert execution reads +back, and clears it on both success and failure. +""" + +from __future__ import annotations + +import torch +import torch.nn as nn +from vllm.distributed.parallel_state import get_pcp_group +from vllm_ascend.ops.fused_moe.dataclass.prepare_finalize import MoEPrepareOutput +from vllm_ascend.ops.fused_moe.moe_comm_method import ( + AllGatherCommImpl, + MC2CommImpl, +) +from vllm_ascend.ops.fused_moe.prepare_finalize import ( + PrepareAndFinalizeWithAllGather, + PrepareAndFinalizeWithMC2, +) +from vllm_ascend.quantization.quant_type import QuantType + +from afd_plugin.connectors.npu.camp2p.external_routing import ( + AFDExternalRoutingInputs, +) + + +class AFDPrepareAndFinalizeWithAllGather(PrepareAndFinalizeWithAllGather): + """ALLGATHER prepare that accepts routing from the Attention side.""" + + def prepare_external( + self, + hidden_states: torch.Tensor, + routing: AFDExternalRoutingInputs, + *, + router_logits: torch.Tensor | None = None, + replace_allreduce: bool = False, + quant_type: QuantType = QuantType.NONE, + ) -> tuple[MoEPrepareOutput, AFDExternalRoutingInputs]: + """Pad and gather the routing exactly like the activation. + + The native ``prepare`` stays untouched for callers that own the gate; this + entry returns the aligned routing instead of publishing it on ``self``. + """ + from vllm_ascend.ascend_forward_context import _EXTRA_CTX + + external = routing + if external.token_count != hidden_states.shape[0]: + raise ValueError( + "external routing covers a different token count than the activation: " + f"{external.token_count} routing rows for " + f"{hidden_states.shape[0]} activation rows", + ) + + if self.moe_config.dp_size > 1: + max_tokens_across_dp = _EXTRA_CTX.max_tokens_across_dp + + self.num_tokens = hidden_states.shape[0] + pad_size = max_tokens_across_dp - self.num_tokens + if pad_size > 0: + hidden_states = nn.functional.pad(hidden_states, (0, 0, 0, pad_size)) + if router_logits is not None: + router_logits = nn.functional.pad( + router_logits, (0, 0, 0, pad_size) + ) + external = external.padded(pad_size) + + hidden_states = self.moe_config.dp_group.all_gather(hidden_states, 0) + if router_logits is not None: + router_logits = self.moe_config.dp_group.all_gather(router_logits, 0) + external = external.gathered( + lambda tensor: self.moe_config.dp_group.all_gather(tensor, 0), + ) + + if self.moe_config.pcp_size > 1: + max_tokens_across_pcp = _EXTRA_CTX.max_tokens_across_pcp + + self.num_tokens_pcp = hidden_states.shape[0] + pad_size = max_tokens_across_pcp - self.num_tokens_pcp + if pad_size > 0: + hidden_states = nn.functional.pad(hidden_states, (0, 0, 0, pad_size)) + if router_logits is not None: + router_logits = nn.functional.pad( + router_logits, (0, 0, 0, pad_size) + ) + external = external.padded(pad_size) + + hidden_states = get_pcp_group().all_gather(hidden_states, dim=0) + if router_logits is not None: + router_logits = get_pcp_group().all_gather(router_logits, dim=0) + external = external.gathered( + lambda tensor: get_pcp_group().all_gather(tensor, dim=0), + ) + + return ( + MoEPrepareOutput( + hidden_states=hidden_states, + router_logits=router_logits, + mc2_mask=None, + padded_hidden_states_shape=None, + pertoken_scale=None, + ), + external, + ) + + +class AFDPrepareAndFinalizeWithMC2(PrepareAndFinalizeWithMC2): + """MC2 prepare that accepts routing from the Attention side.""" + + def prepare_external( + self, + hidden_states: torch.Tensor, + routing: AFDExternalRoutingInputs, + *, + router_logits: torch.Tensor | None = None, + replace_allreduce: bool = False, + quant_type: QuantType = QuantType.NONE, + ) -> tuple[MoEPrepareOutput, AFDExternalRoutingInputs]: + """Pad and TP-shard the routing exactly like the activation.""" + from vllm_ascend.ascend_forward_context import _EXTRA_CTX + from vllm_ascend.ops.fused_moe.moe_utils import _pad_tokens_with_cat + + external = routing + if external.token_count != hidden_states.shape[0]: + raise ValueError( + "external routing covers a different token count than the activation: " + f"{external.token_count} routing rows for " + f"{hidden_states.shape[0]} activation rows", + ) + + self.replace_allreduce = replace_allreduce + mc2_mask = _EXTRA_CTX.mc2_mask + if self.replace_allreduce: + self.num_tokens = hidden_states.shape[0] + start = self.tp_rank * self.num_tokens + mc2_mask = mc2_mask[start : start + self.num_tokens] + pad_size = _EXTRA_CTX.padded_num_tokens // self.tp_size - self.num_tokens + if pad_size > 0: + hidden_states = nn.functional.pad(hidden_states, (0, 0, 0, pad_size)) + if router_logits is not None: + router_logits = nn.functional.pad( + router_logits, (0, 0, 0, pad_size) + ) + external = external.padded(pad_size) + mc2_mask = nn.functional.pad(mc2_mask, (0, pad_size), value=False) + elif self.tp_size > 1: + split_mc2_mask = torch.tensor_split(mc2_mask, self.tp_size, dim=0) + mc2_mask = split_mc2_mask[self.tp_rank] + + padded_hidden_states_shape = hidden_states.shape + if not self.replace_allreduce: + self.num_tokens, _ = hidden_states.shape + target_pad_length = _EXTRA_CTX.padded_num_tokens + pad_size = target_pad_length - self.num_tokens + + if pad_size > 0: + hidden_states = _pad_tokens_with_cat(hidden_states, target_pad_length) + if router_logits is not None: + router_logits = _pad_tokens_with_cat( + router_logits, target_pad_length + ) + external = external.padded(pad_size) + padded_hidden_states_shape = hidden_states.shape + + if self.tp_size > 1: + split_hidden_states = torch.tensor_split( + hidden_states, self.tp_size, dim=0 + ) + if router_logits is not None: + split_router_logits = torch.tensor_split( + router_logits, self.tp_size, dim=0 + ) + router_logits = split_router_logits[self.tp_rank] + hidden_states = split_hidden_states[self.tp_rank] + external = external.split_rows( + lambda tensor: torch.tensor_split(tensor, self.tp_size, dim=0)[ + self.tp_rank + ], + ) + + return ( + MoEPrepareOutput( + hidden_states=hidden_states, + router_logits=router_logits, + mc2_mask=mc2_mask, + padded_hidden_states_shape=padded_hidden_states_shape, + pertoken_scale=None, + ), + external, + ) + + +class AFDAllGatherCommImpl(AllGatherCommImpl): + """ALLGATHER comm that prepares external routing instead of computing it.""" + + def _get_prepare_finalize(self) -> AFDPrepareAndFinalizeWithAllGather: + return AFDPrepareAndFinalizeWithAllGather(self.moe_config) + + +class AFDMC2CommImpl(MC2CommImpl): + """MC2 comm that prepares external routing instead of computing it.""" + + def _get_prepare_finalize(self) -> AFDPrepareAndFinalizeWithMC2: + return AFDPrepareAndFinalizeWithMC2(self.moe_config) + + +# Native comm types the external-routing FFN path supports, and the AFD +# counterparts that own the external-routing prepare/finalize. +_AFD_COMM_BY_NATIVE: dict[type, type] = { + AllGatherCommImpl: AFDAllGatherCommImpl, + MC2CommImpl: AFDMC2CommImpl, +} +_NATIVE_COMM_BY_AFD: dict[type, type] = { + afd: native for native, afd in _AFD_COMM_BY_NATIVE.items() +} + + +def _afd_types_by_comm_type() -> dict[object, type]: + """Map the runtime comm type to the AFD class that owns its prepare stage.""" + from vllm_ascend.ascend_forward_context import MoECommType + + return { + MoECommType.ALLGATHER: AFDAllGatherCommImpl, + MoECommType.MC2: AFDMC2CommImpl, + } + + +def _prebuilt_comms(moe_config) -> dict[type, object]: + """Build the AFD comm instances the native selector can pick from. + + The native selector is keyed by ``MoECommType`` and returns ``None`` when this + configuration was not initialised with that method, so a missing candidate means + the type is simply absent here. Construction happens in the caller's config + context, which is why it runs at component creation and not per forward. + """ + from vllm_ascend.ascend_forward_context import MoECommType + from vllm_ascend.ops.fused_moe.moe_comm_method import get_moe_comm_method + + candidates: dict[type, object] = {} + for comm_type, afd_type in ( + (MoECommType.ALLGATHER, AFDAllGatherCommImpl), + (MoECommType.MC2, AFDMC2CommImpl), + ): + native = get_moe_comm_method(comm_type, moe_config) + if native is None: + continue + candidates[afd_type] = afd_type(native.moe_config) + return candidates + + +class AFDExternalRoutingMoE: + """Run the routed experts on routing supplied by the Attention side. + + The AFD comm instances are built while the component is constructed, in a config + context. A forward only looks up the instance matching the native selection; it + never builds a missing one. + """ + + def __init__(self, routed_experts): + self.routed_experts = routed_experts + self._candidates = _prebuilt_comms(routed_experts.moe_config) + if not self._candidates: + raise RuntimeError( + "DSV4 external routing found no AFD comm candidate for this " + "configuration; the ALLGATHER and MC2 paths are the supported ones", + ) + + def _comm_for(self, comm_type, native_comm): + if native_comm is None: + # The MRV2 hook resolves the method without a config, which misses the + # (type, expert-count) registry key and leaves the instance unset. The + # comm type is always present, so fall back to it. + afd_type = _afd_types_by_comm_type().get(comm_type) + afd_comm = self._candidates.get(afd_type) if afd_type is not None else None + if afd_comm is not None: + return afd_comm + raise RuntimeError( + "DSV4 external routing could not resolve the MoE comm method for " + f"{comm_type} with this layer's config; no AFD candidate was built", + ) + for afd_type, afd_comm in self._candidates.items(): + if isinstance(native_comm, _NATIVE_COMM_BY_AFD[afd_type]): + return afd_comm + raise RuntimeError( + "DSV4 external routing is wired for the ALLGATHER and MC2 MoE paths; " + f"got {type(native_comm).__name__}. The all-to-all and fused-MC2 paths " + "shard differently and are not supported.", + ) + + def forward( + self, + hidden_states: torch.Tensor, + topk_ids: torch.Tensor, + topk_weights: torch.Tensor, + ) -> torch.Tensor: + from vllm_ascend.ascend_forward_context import _EXTRA_CTX + + experts = self.routed_experts + previous_comm = _EXTRA_CTX.moe_comm_method + comm = self._comm_for(_EXTRA_CTX.moe_comm_type, previous_comm) + prepare_finalize = comm.prepare_finalize + + try: + # Scope switch, not a registry change: expert execution reads the comm + # from this context attribute, so both sides use the same object. + _EXTRA_CTX.moe_comm_method = comm + prepare_output, aligned = prepare_finalize.prepare_external( + hidden_states, + AFDExternalRoutingInputs( + topk_ids=topk_ids, + topk_weights=topk_weights, + ), + router_logits=None, + replace_allreduce=experts.moe_config.is_sequence_parallel, + quant_type=experts.quant_type, + ) + experts.ascend_pertoken_scale = prepare_output.pertoken_scale + experts.ascend_mc2_mask = prepare_output.mc2_mask + try: + fused = experts.quant_method.apply( + layer=experts, + x=prepare_output.hidden_states, + topk_weights=aligned.topk_weights, + topk_ids=aligned.topk_ids, + shared_experts=None, + shared_experts_input=None, + ) + finally: + experts.ascend_pertoken_scale = None + experts.ascend_mc2_mask = None + return comm.finalize( + hidden_states=fused.routed_out, + reduce_results=isinstance(comm, AllGatherCommImpl), + padded_hidden_states_shape=prepare_output.padded_hidden_states_shape, + ) + finally: + experts.ascend_pertoken_scale = None + experts.ascend_mc2_mask = None + _EXTRA_CTX.moe_comm_method = previous_comm diff --git a/afd_plugin/model_executor/models/qwen3_moe.py b/afd_plugin/model_executor/models/qwen3_moe.py index 54dc0f406..2e8717073 100644 --- a/afd_plugin/model_executor/models/qwen3_moe.py +++ b/afd_plugin/model_executor/models/qwen3_moe.py @@ -18,6 +18,7 @@ from transformers import Qwen3MoeConfig from vllm.config import VllmConfig from vllm.model_executor.models import qwen3_moe as native +from vllm.model_executor.models.utils import AutoWeightsLoader, maybe_prefix from vllm.platforms import current_platform from afd_plugin.config import ( @@ -25,6 +26,7 @@ AFDConfig, parse_afd_config, ) +from afd_plugin.model_executor.checkpoint_paths import weight_layer_path from afd_plugin.model_executor.models.deepseek_v2 import RemoteFFNProxy if TYPE_CHECKING: @@ -33,6 +35,7 @@ _ATTENTION_ROLE = frozenset(("attention",)) _FFN_ROLE = frozenset(("ffn",)) _BOTH_ROLES = frozenset(("attention", "ffn")) +ROUTED_STAGE = "mlp" def _is_moe_layer(config: Qwen3MoeConfig, layer_idx: int) -> bool: @@ -49,18 +52,10 @@ def _routes_on_attention(afd_config: AFDConfig) -> bool: 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": - continue - try: - layer_idx = int(parts[marker_idx + 1]) - except ValueError: - continue - return layer_idx, parts[marker_idx + 2], tuple(parts[marker_idx + 3 :]) - return None +def _is_routed_mlp_path(name: str) -> bool: + """Whether a checkpoint tensor belongs to a decoder layer's MoE block.""" + layer_path = weight_layer_path(name) + return layer_path is not None and layer_path[1] == ROUTED_STAGE def _checkpoint_weight_roles( @@ -74,7 +69,7 @@ def _checkpoint_weight_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) + layer_path = weight_layer_path(name) if layer_path is None: return _BOTH_ROLES @@ -292,6 +287,80 @@ def get_routed_experts(self, layer_idx: int) -> "RoutedExperts": return self.layers[layer_idx].mlp.experts.routed_experts +class _Qwen3RoutedExpertsLayer(nn.Module): + """One decoder layer reduced to its native MoE block at ``.mlp``.""" + + def __init__(self, vllm_config: VllmConfig, prefix: str) -> None: + super().__init__() + self.mlp = native.Qwen3MoeSparseMoeBlock( + vllm_config=vllm_config, prefix=f"{prefix}.{ROUTED_STAGE}" + ) + + +class _Qwen3RoutedExpertsStack(nn.Module): + """The ``model`` submodule: only ``layers``, each holding its MoE block.""" + + def __init__(self, vllm_config: VllmConfig, prefix: str) -> None: + super().__init__() + num_layers = vllm_config.model_config.hf_text_config.num_hidden_layers + self.layers = nn.ModuleList( + _Qwen3RoutedExpertsLayer(vllm_config, f"{prefix}.layers.{layer_idx}") + for layer_idx in range(num_layers) + ) + + +class AFDQwen3MoeRoutedExperts(nn.Module): + """FFN-only Qwen3 MoE model for routed-experts connectors. + + Builds only the native ``Qwen3MoeSparseMoeBlock`` tree at + ``model.layers..mlp``. The MoE forward never runs here; the connector + reads the local expert parameter views directly. + """ + + fall_back_to_pt_during_load = False + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + afd_config = parse_afd_config(vllm_config, validate=False) + _validate_supported_config(vllm_config, afd_config) + super().__init__() + self.afd_config = afd_config + self.model = _Qwen3RoutedExpertsStack( + vllm_config, maybe_prefix(prefix, "model") + ) + + def is_unused_checkpoint_weight(self, name: str) -> bool: + """Whether vLLM can skip a shard containing only this checkpoint name.""" + return not _is_routed_mlp_path(name) + + # Upstream source: vLLM v0.32.0 + # (5df63a7203cdb9328d824a27b682e77c7b9928be), + # vllm/model_executor/models/qwen3_moe.py, Qwen3MoeModel.hf_to_vllm_mapper. + # AFD difference: filter to routed MoE paths and apply the native mapper at + # this shell's outer AutoWeightsLoader. + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + """Drop non-MoE checkpoint tensors, then load like the native model.""" + # ### PATCH START: keep only routed MoE checkpoint paths. + weights = ( + (name, tensor) + for name, tensor in weights + if not self.is_unused_checkpoint_weight(name) + ) + # ### PATCH END: keep only routed MoE checkpoint paths. + loader = AutoWeightsLoader(self) + return loader.load_weights( + weights, + mapper=native.Qwen3MoeModel.hf_to_vllm_mapper, + ) + + def get_experts_layer_indices(self) -> tuple[int, ...]: + """Every decoder layer is an MoE layer for a supported checkpoint.""" + return tuple(range(len(self.model.layers))) + + def get_routed_experts(self, layer_idx: int) -> "RoutedExperts": + """The vLLM module holding one MoE layer's local experts.""" + return self.model.layers[layer_idx].mlp.experts.routed_experts + + class AFDQwen3MoeForCausalLM(native.Qwen3MoeForCausalLM): """Qwen3 MoE causal LM wrapper for AFD CUDA execution.""" @@ -411,4 +480,4 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: return super().load_weights(role_weights) -__all__ = ["AFDQwen3MoeForCausalLM"] +__all__ = ["AFDQwen3MoeForCausalLM", "AFDQwen3MoeRoutedExperts"] diff --git a/afd_plugin/model_executor/npu/async_cam_layered.py b/afd_plugin/model_executor/npu/async_cam_layered.py index 6079ed274..06962e7a4 100644 --- a/afd_plugin/model_executor/npu/async_cam_layered.py +++ b/afd_plugin/model_executor/npu/async_cam_layered.py @@ -108,6 +108,20 @@ def __init__(self, layers: list[LayeredFFNWeights]) -> None: self.w13_op = torch.ops.afd_ascend.gmm_swiglu_quant_v2_layered self.w2_op = torch.ops.afd_ascend.grouped_matmul_layered + def geometry(self) -> tuple[int, int, int]: + """Return (experts, hidden, intermediate) from the loaded weight layout. + + INT4 weights pack eight nibbles per INT32, so the stored last dimension + is a quarter of the gate/up width; INT8 stores that width directly. + Callers that need only the logical geometry must not re-derive it. + """ + experts, hidden, last = self.w13[0].shape + if self.quantization == "W8A8": + if last % 2: + raise ValueError("layered GMM requires an even W13 gate/up width") + return experts, hidden, last // 2 + return experts, hidden, last * INT4_PER_INT32 // 2 + def graph_signature(self, capacity: int) -> str: """Compare static layouts across ranks without local device addresses.""" tensors = ( diff --git a/afd_plugin/model_executor/npu/async_cam_mega.py b/afd_plugin/model_executor/npu/async_cam_mega.py new file mode 100644 index 000000000..1370346de --- /dev/null +++ b/afd_plugin/model_executor/npu/async_cam_mega.py @@ -0,0 +1,92 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""One native launch drains an exclusive quantized transaction per layer.""" + +from __future__ import annotations + +import torch + +from afd_plugin.model_executor.npu.async_cam_layered import AsyncCAMLayeredExecutor +from afd_plugin.validation import cam_receive_capacity + +# Mirror Pull::MetadataWords in afd_pull_layout.h (all fields are int64_t): +# [total, source, layer, round_begin, round_end], tp source prefixes, +# tp * experts original counts, then [session, sequence]. +METADATA_WORD_BYTES = torch.int64.itemsize +METADATA_FIXED_WORDS = 7 +METADATA_PREFIX_WORDS_PER_SOURCE = 1 +# NativeMegaWorkspace in afd_native_mega_workspace.h aligns scratch regions to +# 512 bytes. Match its counts offset, including padding after batch metadata. +WORKSPACE_ALIGNMENT_BYTES = 512 + + +class AsyncCAMMegaExecutor: + """Pull transport with a compiled W4A8 or W8A8 execution plan.""" + + def __init__(self, weights: AsyncCAMLayeredExecutor, connector) -> None: + if not weights.per_channel: + raise ValueError("native Mega requires per-channel W4A8/W8A8 weights") + self.weights = weights + self.connector = connector + self.capacity = cam_receive_capacity() + experts, hidden, intermediate = weights.geometry() + self.config = [ + connector.max_num_batched_tokens, + hidden, + intermediate, + connector.topk, + connector.ffn_size, + connector.attn_size, + experts, + connector.world_rank, + connector.topology.world_size, + connector.tp_size, + self.capacity, + 1 if weights.quantization == "W4A8" else 2, + 1, + len(weights.layer_ids), + ] + device = weights.w13[0].device + size = torch.ops.afd_ascend.afd_mega_moe_pull_workspace_size( + self.capacity, hidden, intermediate, experts, connector.tp_size + ) + self.scratch = torch.empty(size, dtype=torch.uint8, device=device) + metadata_words = METADATA_FIXED_WORDS + connector.tp_size * ( + METADATA_PREFIX_WORDS_PER_SOURCE + experts + ) + metadata_bytes = METADATA_WORD_BYTES * metadata_words + count_offset = ( + (metadata_bytes + WORKSPACE_ALIGNMENT_BYTES - 1) + // WORKSPACE_ALIGNMENT_BYTES + * WORKSPACE_ALIGNMENT_BYTES + ) + count_bytes = METADATA_WORD_BYTES * experts + self.batch_info = self.scratch[:metadata_bytes].view(torch.int64) + self.group_list = self.scratch[count_offset : count_offset + count_bytes].view( + torch.int64 + ) + self.anchor = torch.empty(1, dtype=torch.bfloat16, device=device) + mapping = [-1] * (weights.layer_ids[-1] + 1) + for slot, layer_id in enumerate(weights.layer_ids): + mapping[layer_id] = slot + self.layer_map = torch.tensor(mapping, dtype=torch.int64, device=device) + self.op = torch.ops.afd_ascend.afd_mega_moe_pull_native + + def __call__(self) -> torch.Tensor: + c, w = self.connector, self.weights + return self.op( + self.anchor, + c.comm_args, + self.layer_map, + self.scratch, + w.w13, + w.w2, + w.w13_scale, + w.w2_scale, + w.w13_bias if w.quantization == "W4A8" else [], + w.w2_bias if w.quantization == "W4A8" else [], + self.config, + c.group_name, + w.routed_scaling_factor, + w.applied_swiglu_limit, + ) diff --git a/afd_plugin/v1/worker/attention_model_runner_v2.py b/afd_plugin/v1/worker/attention_model_runner_v2.py index 82eefc156..3ff495685 100644 --- a/afd_plugin/v1/worker/attention_model_runner_v2.py +++ b/afd_plugin/v1/worker/attention_model_runner_v2.py @@ -197,8 +197,8 @@ def _use_afd_capture_input_preparation( # Patch reason: native ModelCudaGraphManager.capture calls this exact # module symbol once before each warmup/formal-capture forward. - # Patch functionality: preserve native preparation and publish the - # matching AFD event before the forward starts. + # Patch functionality: preserve native preparation and stage the matching + # AFD metadata before the forward; synchronous connectors also send control. # Signature: matches vLLM v0.32.0 prepare_inputs_to_capture exactly, # including the trailing pcp_manager parameter. # Upstream source: vllm/v1/worker/gpu/cudagraph_utils.py, @@ -237,10 +237,11 @@ def prepare_inputs_to_capture( None, int(desc.num_tokens), ) - runner.send_dp_metadata( - runner.build_capture_dp_metadata(int(desc.num_tokens)), - None, - ) + if runner.connector.control_plane is not None: + runner.send_dp_metadata( + runner.build_capture_dp_metadata(int(desc.num_tokens)), + None, + ) runner._afd_suppress_metadata_send = True return attention_state # ### PATCH END: stage one exact AFD capture event. diff --git a/afd_plugin/v1/worker/ffn_model_runner.py b/afd_plugin/v1/worker/ffn_model_runner.py index bab6ac42a..2f67a6132 100644 --- a/afd_plugin/v1/worker/ffn_model_runner.py +++ b/afd_plugin/v1/worker/ffn_model_runner.py @@ -4,6 +4,7 @@ from __future__ import annotations +from collections.abc import Callable from contextlib import contextmanager from typing import TYPE_CHECKING, Any, cast @@ -30,6 +31,7 @@ AFDControlPlane, AFDDPMetadata, ) +from afd_plugin.connectors.base import AFDControlPlaneClosedError from afd_plugin.v1.worker.cuda_graph import ( AFDGraphRunMode, graph_run_mode, @@ -150,6 +152,65 @@ def get_kv_cache_spec(self) -> dict[str, KVCacheSpec]: def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None: return None + def execute_connector_driven_step( + self, *, cancel_requested: Callable[[], bool] + ) -> bool: + """Receive one control payload and run the FFN step it describes. + + Blocks in the control plane until Attention rank 0 sends the next + payload. Warmup and capture payloads go to ``capture_model`` when FFN + CUDA graphs are enabled; every other payload (eager, replay, or an + eager fallback for a replay without a captured key) goes to + ``execute_model``. The step ends with a device synchronize so the + caller only sees completed work. + + ``cancel_requested`` tells a cancelled receive apart from a broken + control plane: when the control plane was closed on purpose (the + caller set its stop flag before closing it) the method returns + ``False``; any other closed or failed receive raises unchanged, so a + peer or transport failure is never reported as a clean stop. + + Shared by ``AFDFFNWorker._run_ffn_server_loop`` and the standalone + FFN service backend, so both drive identical capture, replay and + synchronization order. + """ + control_plane = self.connector.control_plane + if control_plane is None: + raise NotImplementedError( + "GPU FFN only supports control-plane-driven connectors; " + "connectors without a control plane (control_plane is None) " + "are not supported.", + ) + + try: + payload = control_plane.recv_dp_metadata_list() + except AFDControlPlaneClosedError: + if not cancel_requested(): + raise + return False + dp_metadata_list = payload.dp_metadata_list + is_attn_graph_capturing = payload.is_graph_capturing + is_warmup = payload.is_warmup + is_graph_replaying = payload.is_graph_replaying + + if self.use_cuda_graph and (is_warmup or is_attn_graph_capturing): + self.capture_model( + dp_metadata_list=dp_metadata_list, + is_warmup=is_warmup, + is_attn_graph_capturing=is_attn_graph_capturing, + ) + else: + self.execute_model( + dp_metadata_list=dp_metadata_list, + is_graph_capturing=is_attn_graph_capturing, + is_warmup=is_warmup, + is_graph_replaying=is_graph_replaying, + ) + + if self.device.type == "cuda": + torch.cuda.synchronize() + return True + def execute_model( self, scheduler_output: SchedulerOutput | None = None, @@ -220,7 +281,13 @@ def _ffn_forward( stage_idx: self._make_ffn_dp_metadata(dp_metadata_list[stage_idx]) for stage_idx in stage_ids } - recv_input_ids = getattr(model, "afd_requires_input_ids", False) + # The wire slot carries one payload: either token ids or Attention-side + # routing. Routing wins when the model declares it, because a sync + # gate-on-Attention model does not need ids on the FFN side at all. + recv_routing = getattr(model, "afd_requires_attention_routing", False) + recv_input_ids = ( + False if recv_routing else getattr(model, "afd_requires_input_ids", False) + ) with _ffn_forward_context(self.vllm_config) as forward_context: for layer_idx in layer_indices: uses_remote_experts = layer_idx in experts_layer_indices @@ -233,7 +300,9 @@ def _ffn_forward( recv_kwargs: dict[str, Any] = {} if routing_spec is not None: recv_kwargs["routing_spec"] = routing_spec - if recv_input_ids: + if recv_routing: + recv_kwargs["recv_routing"] = True + elif recv_input_ids: recv_kwargs["recv_input_ids"] = True payload = self.connector.recv_attn_output( ubatch_idx=stage_idx, 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 55c26f41d..7286263eb 100644 --- a/afd_plugin/v1/worker/ffn_routed_experts_model_runner.py +++ b/afd_plugin/v1/worker/ffn_routed_experts_model_runner.py @@ -4,7 +4,7 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, Protocol import torch from vllm.config import VllmConfig @@ -92,6 +92,18 @@ def local_experts_from_vllm(routed_experts: RoutedExperts) -> FFNLocalExperts: ) +class RoutedExpertsModel(Protocol): + """What the routed runner reads from a model: MoE layers and their experts. + + Implemented by the AFD CausalLM models the Worker path loads and by the + standalone FFN service's FFN-only model shells. + """ + + def get_experts_layer_indices(self) -> tuple[int, ...]: ... + + def get_routed_experts(self, layer_idx: int) -> RoutedExperts: ... + + def _require_routed(connector: AFDConnectorBase) -> AFDRoutedExpertsConnector: if not isinstance(connector, AFDRoutedExpertsConnector): raise TypeError( @@ -121,25 +133,47 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device) -> None: _require_routed(self.connector) self.moe_layer_indices: tuple[int, ...] = () self.local_experts: dict[int, FFNLocalExperts] = {} + self._receiver_exited = False @property def routed_connector(self) -> AFDRoutedExpertsConnector: """The connector, typed as the routed-experts contract it must meet.""" return _require_routed(self.connector) + @property + def receiver_exited(self) -> bool: + """Whether ``mark_receiver_exited`` proved the receive loop finished.""" + return self._receiver_exited + 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._take_local_experts(self.get_model()) + + def _take_local_experts(self, model: RoutedExpertsModel) -> None: + """Record the model's MoE layers and this rank's expert views per layer.""" 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 mark_receiver_exited(self) -> None: + """Record that the receive loop ended cleanly and its device work is done. + + Called on the receiver thread after its last step. Each step already + ends with a synchronize; the extra fence here also covers a step that + returned early because a stop cancelled its receive. The standalone FFN + backend releases weights, CUDA graphs and DeepEP buffers only after + this proof. + """ + if self.device.type == "cuda": + torch.cuda.synchronize(self.device) + self._receiver_exited = True + def initialize_afd_connector(self) -> None: self.routed_connector.init_afd_connector( moe_layer_indices=self.moe_layer_indices, @@ -189,4 +223,8 @@ def _ffn_forward( return ffn_output -__all__ = ["GPUFFNRoutedExpertsModelRunner", "local_experts_from_vllm"] +__all__ = [ + "GPUFFNRoutedExpertsModelRunner", + "RoutedExpertsModel", + "local_experts_from_vllm", +] diff --git a/afd_plugin/v1/worker/ffn_worker.py b/afd_plugin/v1/worker/ffn_worker.py index 4a998cc34..99702d130 100644 --- a/afd_plugin/v1/worker/ffn_worker.py +++ b/afd_plugin/v1/worker/ffn_worker.py @@ -13,9 +13,8 @@ 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.config import AFDConfig, parse_afd_config from afd_plugin.connectors import AFDConnectorFactory, AFDRoutedExpertsConnector -from afd_plugin.connectors.base import AFDControlPlaneClosedError from afd_plugin.model_executor.models.model_utils import get_afd_model_config from afd_plugin.v1.worker.ffn_model_runner import GPUFFNModelRunner from afd_plugin.validation import ( @@ -33,6 +32,11 @@ FFN_LOOP_JOIN_TIMEOUT_S = 5 +def _uses_routed_experts_connector(afd_config: AFDConfig) -> bool: + connector_cls = AFDConnectorFactory.get_connector_class(afd_config.connector) + return issubclass(connector_cls, AFDRoutedExpertsConnector) + + class AFDFFNWorker(Worker): """FFN worker that owns the AFD daemon loop. @@ -76,7 +80,7 @@ def init_device(self): """Initialize the native GPU worker with the connector-driven runner.""" # ### PATCH START: validate and select the AFD model. - assert_compatible_afd_stack( + afd_config = assert_compatible_afd_stack( self.vllm_config, caller="AFDFFNWorker.init_device", expected_role="ffn", @@ -89,6 +93,7 @@ def init_device(self): afd_model_config = get_afd_model_config( self.vllm_config.model_config, device_type="cuda", + routed_experts_connector=_uses_routed_experts_connector(afd_config), ) self.vllm_config.model_config = afd_model_config self.model_config = afd_model_config @@ -104,8 +109,7 @@ def init_device(self): def _make_model_runner(self) -> GPUModelRunner: """Construct the connector-driven runner via the native factory hook.""" afd_config = parse_afd_config(self.vllm_config, expected_role="ffn") - connector_cls = AFDConnectorFactory.get_connector_class(afd_config.connector) - if issubclass(connector_cls, AFDRoutedExpertsConnector): + if _uses_routed_experts_connector(afd_config): from afd_plugin.v1.worker.ffn_routed_experts_model_runner import ( GPUFFNRoutedExpertsModelRunner, select_routed_moe_backend, @@ -174,6 +178,13 @@ def ffn_worker_loop() -> None: self._ffn_thread.start() def _run_ffn_server_loop(self) -> None: + """Drive connector steps until the shutdown event stops the loop. + + Each iteration delegates one receive-and-execute step to the runner's + ``execute_connector_driven_step``. The loop ends when the event is set + between steps, or when ``stop_ffn_server_loop`` set the event and then + closed the control plane, which ends a receive blocked while idle. + """ event = self._ffn_shutdown_event if event is None: return @@ -182,43 +193,10 @@ def _run_ffn_server_loop(self) -> None: torch.cuda.set_device(self.device) while not event.is_set(): - control_plane = self.model_runner.connector.control_plane - if control_plane is None: - raise NotImplementedError( - "GPU FFN only supports control-plane-driven connectors; " - "connectors without a control plane (control_plane is None) " - "are not supported.", - ) - - try: - payload = control_plane.recv_dp_metadata_list() - except AFDControlPlaneClosedError: - if not event.is_set(): - raise - break - dp_metadata_list = payload.dp_metadata_list - is_attn_graph_capturing = payload.is_graph_capturing - is_warmup = payload.is_warmup - is_graph_replaying = payload.is_graph_replaying - - if self.model_runner.use_cuda_graph and ( - is_warmup or is_attn_graph_capturing + if not self.model_runner.execute_connector_driven_step( + cancel_requested=event.is_set ): - self.model_runner.capture_model( - dp_metadata_list=dp_metadata_list, - is_warmup=is_warmup, - is_attn_graph_capturing=is_attn_graph_capturing, - ) - else: - self.model_runner.execute_model( - dp_metadata_list=dp_metadata_list, - is_graph_capturing=is_attn_graph_capturing, - is_warmup=is_warmup, - is_graph_replaying=is_graph_replaying, - ) - - if self.device.type == "cuda": - torch.cuda.synchronize() + break def raise_ffn_loop_error_if_any(self) -> None: error = self._ffn_loop_error diff --git a/afd_plugin/v1/worker/npu/async_ffn_model_runner.py b/afd_plugin/v1/worker/npu/async_ffn_model_runner.py index 89b8d2345..8ce1e175a 100644 --- a/afd_plugin/v1/worker/npu/async_ffn_model_runner.py +++ b/afd_plugin/v1/worker/npu/async_ffn_model_runner.py @@ -12,11 +12,6 @@ from vllm.logger import init_logger from vllm_ascend.utils import AscendDeviceType, get_ascend_device_type -from afd_plugin.async_cam_startup import ( - GRAPH_STARTUP_TIMEOUT_SECONDS, - POLL_SECONDS, - FFNGraphPlan, -) from afd_plugin.compat.npu import fail_if_unsupported_npu_afd_features from afd_plugin.compat.npu.ops import ensure_cam_async_ops_available from afd_plugin.compat.npu.profiler import ( @@ -25,15 +20,21 @@ stop_afd_npu_profiler, ) from afd_plugin.config import parse_afd_config +from afd_plugin.connectors.metadata import AFDA2FTransferPayload from afd_plugin.connectors.npu.async_cam import ( AFDAsyncTransferState, CAMAsyncAFDConnector, ) +from afd_plugin.connectors.npu.async_cam_runtime.async_cam_startup import ( + GRAPH_STARTUP_TIMEOUT_SECONDS, + POLL_SECONDS, + FFNGraphPlan, +) from afd_plugin.model_executor.npu.async_cam_layered import ( AsyncCAMLayeredExecutor, LayeredFFNWeights, ) -from afd_plugin.v1.worker.npu.async_ffn_two_stream import AsyncFFNTwoStream +from afd_plugin.model_executor.npu.async_cam_mega import AsyncCAMMegaExecutor from afd_plugin.validation import ( cam_receive_capacity, requested_ffn_mode, @@ -56,7 +57,6 @@ def get_async_cam_layered_weights(self) -> tuple[list[LayeredFFNWeights], str]: class AsyncNPUFFNModelRunner: afd_expected_role = "ffn" - _two_stream: AsyncFFNTwoStream | None = None def __init__( self, @@ -76,6 +76,9 @@ def __init__( self.model_memory_usage = 0 self.layer_ids: tuple[int, ...] = () self._layered_executor: AsyncCAMLayeredExecutor | None = None + self._mega_executor: AsyncCAMMegaExecutor | None = None + self.session = self.connector.ffn_session() + self._first_step = True self.requested_mode = requested_ffn_mode(vllm_config) self.effective_mode = "eager" self._graph: torch.npu.NPUGraph | None = None @@ -165,7 +168,7 @@ def initialize_afd_connector(self) -> None: if self.connector.is_initialized: return assert self.model is not None - stream_signature = stream_error = executor_error = signature = "" + executor_error = signature = "" config = self.connector.extra_info try: models = (self.model,) + ( @@ -180,45 +183,23 @@ def initialize_afd_connector(self) -> None: self._initialize_layered_executor() assert self._layered_executor is not None signature = self._layered_executor.graph_signature(cam_receive_capacity()) + if config.ffn_mega_moe: + self._mega_executor = AsyncCAMMegaExecutor( + self._layered_executor, self.connector + ) + signature = "native_mc2:" + signature except Exception as exc: executor_error = str(exc) or type(exc).__name__ executor = self._layered_executor - if config.ffn_two_stream: - stream_signature = repr( - ( - 1, - 2, - config.ffn_compute_aic_num, - config.ffn_compute_aiv_num, - config.ffn_combine_aiv_num, - ) - ) - try: - if not executor_error: - assert executor is not None - self._two_stream = AsyncFFNTwoStream( - self.device, - config, - self.connector, - executor, - ) - except Exception as exc: - stream_error = str(exc) or type(exc).__name__ self.connector.graph_startup.plan = FFNGraphPlan( requested=self.requested_mode, executor_error=executor_error, signature=signature, layer_ids=self.layer_ids, - stream_signature=stream_signature, - stream_error=stream_error, ) try: self.connector.init_afd_connector() self.effective_mode = self.connector.graph_startup.mode - if self._two_stream is not None: - self._two_stream.compute.wait_stream( - torch.npu.current_stream(self.device) - ) if self.effective_mode == "full": assert executor is not None chunks = self.connector.graph_startup_chunks() * ( @@ -230,10 +211,9 @@ def initialize_afd_connector(self) -> None: done.record() self._wait_graph_device(done) self.connector.graph_startup.phase("warmup") - if self._two_stream is None: - stream = torch.npu.Stream(device=self.device) - stream.wait_stream(torch.npu.current_stream(self.device)) - self._initialize_graph(stream) + stream = torch.npu.Stream(device=self.device) + stream.wait_stream(torch.npu.current_stream(self.device)) + self._initialize_graph(stream) self._capture_graph() self.connector.graph_startup.phase("captured") except Exception as exc: @@ -258,6 +238,7 @@ def _initialize_graph(self, stream: torch.npu.Stream) -> None: self._graph_stream = stream self._graph_pool = torch.npu.graph_pool_handle() self._graph_tensors: list[torch.Tensor] = [] + self._graph_payload: AFDA2FTransferPayload | None = None def _check_graph_live(self) -> None: if self._graph_failure is not None: @@ -270,10 +251,6 @@ def _check_graph_live(self) -> None: @torch.inference_mode() def _capture_graph(self) -> None: self._check_graph_live() - if self._two_stream is not None: - self._two_stream.capture() - self._graph_capture_count = self._two_stream.capture_count - return if self._graph is not None: raise RuntimeError("Async FFN may only capture once") assert self._layered_executor is not None @@ -282,21 +259,27 @@ def _capture_graph(self) -> None: with torch.npu.graph( self._graph, pool=self._graph_pool, stream=self._graph_stream ): - self._graph_payload = self.connector.recv_attn_output() - state = cast(AFDAsyncTransferState, self._graph_payload.context.states) - assert state.token_nums_rankid_layeridx is not None - assert state.group_list is not None - assert state.dynamic_scales is not None - self._graph_output = self._layered_executor( - self._graph_payload.hidden_states, - state.dynamic_scales, - state.group_list, - state.token_nums_rankid_layeridx, - retained_tensors=self._graph_tensors, - ) - self.connector.send_ffn_output( - self._graph_output, self._graph_payload.context - ) + if self._mega_executor is not None: + self._graph_payload = None + self._graph_output = self._mega_executor() + else: + self._graph_payload = self.connector.recv_attn_output() + state = cast( + AFDAsyncTransferState, self._graph_payload.context.states + ) + assert state.token_nums_rankid_layeridx is not None + assert state.group_list is not None + assert state.dynamic_scales is not None + self._graph_output = self._layered_executor( + self._graph_payload.hidden_states, + state.dynamic_scales, + state.group_list, + state.token_nums_rankid_layeridx, + retained_tensors=self._graph_tensors, + ) + self.connector.send_ffn_output( + self._graph_output, self._graph_payload.context + ) self._graph_capture_count = 1 except Exception as exc: self._graph_failure = exc @@ -305,13 +288,6 @@ def _capture_graph(self) -> None: @torch.inference_mode() def _submit_graph(self, count: int = 1) -> None: self._check_graph_live() - if self._two_stream is not None: - if self._graph_pending or self._graph_receiver_exited: - raise RuntimeError("Async FFN graph is not ready for replay") - self._graph_pending = True - self._two_stream.submit(count) - self._graph_replay_count = self._two_stream.sequence - return if ( self._graph is None or self._graph_capture_count != 1 @@ -332,16 +308,13 @@ def _submit_graph(self, count: int = 1) -> None: def _complete_graph(self) -> bool: self._check_graph_live() - if self._two_stream is not None: - result = self._two_stream.complete() - self._graph_pending = False - return result if not self._graph_pending: raise RuntimeError("Async FFN has no pending replay") try: # One fence per bounded batch. The original DR cannot be cancelled; # a blocked batch retains its graph, model and communication buffers. torch.npu.synchronize() + self.session.after_step() except Exception as exc: self._graph_failure = exc raise @@ -352,10 +325,6 @@ def _close_graph(self) -> None: if self._graph_closed: return self._check_graph_live() - if self._two_stream is not None: - self._two_stream.close(receiver_exited=self._graph_receiver_exited) - self._graph_closed = True - return if self._graph_pending or not self._graph_receiver_exited: raise RuntimeError( "Async graph resources require receiver exit and device completion" @@ -400,11 +369,19 @@ def _initialize_layered_executor(self) -> None: ) @torch.inference_mode() - def execute_connector_driven_step(self) -> bool: + def execute_connector_driven_step(self, shutdown_requested: bool = False) -> bool: step_afd_npu_profiler(self.prof) - if self._graph is None and self._two_stream is None: + if not self.session.keep_serving(shutdown_requested): + return False + if self._first_step: + self.session.before_first_step() + if self._graph is not None: + self._graph_stream.wait_stream(torch.npu.current_stream(self.device)) + self._first_step = False + if self._graph is None: # Eager and graph modes use the same layered executor. self._ffn_forward_connector_driven() + self.session.after_step() return True if not self._service_announced: self.connector.graph_startup.service_ready() @@ -414,16 +391,28 @@ def execute_connector_driven_step(self) -> bool: self._submit_graph(len(self.layer_ids)) return self._complete_graph() + def drain(self) -> None: + self.session.drain(self._run_one_transaction) + + def _run_one_transaction(self) -> None: + if self._graph is not None: + self._submit_graph(1) + self._complete_graph() + else: + self._ffn_forward_connector_driven(single=True) + torch.npu.synchronize() + self.session.after_step() + def _ffn_forward_connector_driven( self, *, single: bool = False, ) -> torch.Tensor | None: assert self.model is not None - if self._two_stream is not None: - self._two_stream.submit(1 if single else len(self.layer_ids)) - self._two_stream.complete(wait=self._wait_graph_device if single else None) - return None + if self._mega_executor is not None: + for _ in self.layer_ids[:1] if single else self.layer_ids: + rank_ffn_output = self._mega_executor() + return rank_ffn_output rank_ffn_output = None assert self._layered_executor is not None for _ in self.layer_ids[:1] if single else self.layer_ids: @@ -460,15 +449,16 @@ def shutdown(self) -> None: raise RuntimeError( "Failed graph startup retains resources" ) from self._failure + if self.session.requires_drain and not self._graph_receiver_exited: + raise RuntimeError("Receiver has not exited; retain resources") + self.session.validate_shutdown() from vllm.v1.worker.gpu.shutdown import free_before_shutdown - if self._two_stream is not None: - self._two_stream.close(receiver_exited=self._graph_receiver_exited) - self._two_stream = None - elif self._graph is not None: + if self._graph is not None: self._close_graph() stop_afd_npu_profiler(self.prof) self.connector.close() + self._mega_executor = None self._layered_executor = None free_before_shutdown(self.vllm_config) self.model = None diff --git a/afd_plugin/v1/worker/npu/async_ffn_two_stream.py b/afd_plugin/v1/worker/npu/async_ffn_two_stream.py deleted file mode 100644 index f7b1310fe..000000000 --- a/afd_plugin/v1/worker/npu/async_ffn_two_stream.py +++ /dev/null @@ -1,203 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project -"""Two-slot Async CAM receive/compute and FIFO-send execution.""" - -from __future__ import annotations - -from collections.abc import Callable -from dataclasses import dataclass, field -from typing import TYPE_CHECKING, cast - -import torch - -from afd_plugin.connectors.npu.async_cam import AFDAsyncTransferState - -if TYPE_CHECKING: - from afd_plugin.connectors.metadata import AFDA2FTransferPayload - from afd_plugin.connectors.npu.async_cam import ( - AFDAsyncExtraInfo, - CAMAsyncAFDConnector, - ) - from afd_plugin.model_executor.npu.async_cam_layered import AsyncCAMLayeredExecutor - - -@dataclass -class _Slot: - ready: torch.npu.Event - done: torch.npu.Event - submitted: bool = False - payload: AFDA2FTransferPayload | None = None - output: torch.Tensor | None = None - tensors: list[torch.Tensor] = field(default_factory=list) - compute_graph: torch.npu.NPUGraph | None = None - send_graph: torch.npu.NPUGraph | None = None - pools: tuple | None = None - - def release_transaction(self) -> None: - self.payload = self.output = None - self.tensors.clear() - - -class AsyncFFNTwoStream: - def __init__( - self, - device: torch.device, - config: AFDAsyncExtraInfo, - connector: CAMAsyncAFDConnector, - executor: AsyncCAMLayeredExecutor, - ) -> None: - self.connector, self.executor = connector, executor - self.sequence = self.capture_count = self.replay_count = 0 - self.pending = self.closed = False - self.failure: Exception | None = None - self.compute = torch.npu.Stream(device=device) - self.send = torch.npu.Stream(device=device) - limits = torch.npu.get_stream_limit(self.compute) - if ( - config.ffn_compute_aic_num > limits["cube_core_num"] - or 2 * config.ffn_compute_aic_num > config.ffn_compute_aiv_num - or config.ffn_compute_aiv_num + config.ffn_combine_aiv_num - > limits["vector_core_num"] - ): - raise ValueError( - "FFN two-stream core limits exceed the available 1:2 budget" - ) - torch.npu.set_stream_limit( - self.compute, - cube_num=config.ffn_compute_aic_num, - vector_num=config.ffn_compute_aiv_num, - ) - torch.npu.set_stream_limit(self.send, vector_num=config.ffn_combine_aiv_num) - self.compute.wait_stream(torch.npu.current_stream(device)) - self.slots = [_Slot(torch.npu.Event(), torch.npu.Event()) for _ in range(2)] - - def _check_live(self) -> None: - if self.failure is not None: - raise RuntimeError( - "Async FFN two-stream failed; retain resources" - ) from self.failure - if self.closed: - raise RuntimeError("Async FFN two-stream is closed") - - def _compute(self, slot: _Slot) -> None: - slot.payload = self.connector.recv_attn_output() - state = cast(AFDAsyncTransferState, slot.payload.context.states) - slot.output = self.executor( - slot.payload.hidden_states, - state.dynamic_scales, - state.group_list, - state.token_nums_rankid_layeridx, - retained_tensors=slot.tensors, - ) - - def _send(self, slot: _Slot) -> None: - assert slot.payload is not None and slot.output is not None - self.connector.send_ffn_output(slot.output, slot.payload.context) - - def capture(self) -> None: - self._check_live() - if self.pending or self.capture_count: - raise RuntimeError("Async FFN two-stream is not ready for capture") - try: - for slot in self.slots: - slot.compute_graph, slot.send_graph = ( - torch.npu.NPUGraph(), - torch.npu.NPUGraph(), - ) - slot.pools = ( - torch.npu.graph_pool_handle(), - torch.npu.graph_pool_handle(), - ) - with torch.npu.graph( - slot.compute_graph, pool=slot.pools[0], stream=self.compute - ): - self._compute(slot) - with torch.npu.graph( - slot.send_graph, pool=slot.pools[1], stream=self.send - ): - self._send(slot) - self.capture_count += 2 - except Exception as exc: - self.failure = exc - raise - self.sequence = self.replay_count = 0 - - def submit(self, count: int) -> None: - self._check_live() - if self.pending or count < 1: - raise RuntimeError("Async FFN two-stream is not ready for submission") - self.pending = True # Even a failed call can leave queued native work. - try: - for _ in range(count): - slot = self.slots[self.sequence % 2] - captured = slot.compute_graph is not None - if slot.submitted and not captured: - # An idle DR has no startup deadline; retain until CS completes. - slot.done.synchronize() - slot.release_transaction() - with torch.npu.stream(self.compute): - if slot.submitted: - self.compute.wait_event(slot.done) - if captured: - assert slot.compute_graph is not None - slot.compute_graph.replay() - self.replay_count += 1 - else: - self._compute(slot) - slot.ready.record(self.compute) - with torch.npu.stream(self.send): - self.send.wait_event(slot.ready) - if captured: - assert slot.send_graph is not None - slot.send_graph.replay() - self.replay_count += 1 - else: - self._send(slot) - slot.done.record(self.send) - slot.submitted = True - self.sequence += 1 - except Exception as exc: - self.failure = exc - raise - - def complete( - self, *, wait: Callable[[torch.npu.Event], None] | None = None - ) -> bool: - self._check_live() - if not self.pending: - raise RuntimeError("Async FFN two-stream has no pending submission") - try: - if wait is None: - # Match the existing bounded-batch completion policy. - torch.npu.synchronize() - else: - # FIFO send completion also covers all preceding compute work. - wait(self.slots[(self.sequence - 1) % len(self.slots)].done) - except Exception as exc: - self.failure = exc - raise - self.pending = False - if not self.capture_count: - for slot in self.slots: - slot.release_transaction() - return True - - def close(self, *, receiver_exited: bool) -> None: - if self.closed: - return - self._check_live() - if self.pending or not receiver_exited: - raise RuntimeError( - "Two-stream resources require receiver exit and device completion" - ) - for slot in self.slots: - if slot.compute_graph is not None: - slot.compute_graph.reset() - if slot.send_graph is not None: - slot.send_graph.reset() - slot.release_transaction() - slot.compute_graph = slot.send_graph = None - slot.pools = None - self.slots.clear() - del self.compute, self.send, self.connector, self.executor - self.closed = True diff --git a/afd_plugin/v1/worker/npu/attention_model_runner_v2.py b/afd_plugin/v1/worker/npu/attention_model_runner_v2.py index 246079a65..8f791cca7 100644 --- a/afd_plugin/v1/worker/npu/attention_model_runner_v2.py +++ b/afd_plugin/v1/worker/npu/attention_model_runner_v2.py @@ -82,10 +82,10 @@ def _use_afd_fullgraph_replay_hook( original_run_fullgraph = manager.run_fullgraph manager_state[_AFD_FULLGRAPH_HOOK_MARKER] = True - # Patch reason: native FULL replay bypasses ForwardContext creation, so the - # execute-scoped AFD provider cannot publish runtime control. - # Patch functionality: wrap only this manager instance and publish one - # ordinary padded control payload immediately before each native replay. + # Patch reason: native FULL replay bypasses the ordinary forward, so its + # attention-parameter update needs fresh AFD real/padded token metadata. + # Patch functionality: stage that metadata on this manager's replay scope; + # synchronous connectors also publish one padded control payload. # Signature: matches vLLM v0.32.0 ModelCudaGraphManager.run_fullgraph exactly. # Upstream source: vllm/v1/worker/gpu/cudagraph_utils.py, # ModelCudaGraphManager.run_fullgraph; commit @@ -119,10 +119,11 @@ def run_fullgraph( runner._is_warmup = False runner._afd_is_graph_capturing = False runner._afd_is_graph_replaying = True - runner.send_dp_metadata( - runner.build_capture_dp_metadata(padded_tokens), - None, - ) + if runner.connector.control_plane is not None: + runner.send_dp_metadata( + runner.build_capture_dp_metadata(padded_tokens), + None, + ) result = original_run_fullgraph(desc) finally: runner._afd_is_graph_replaying = previous_is_graph_replaying @@ -269,6 +270,7 @@ def __init__( self.vllm_config, self.device, afd_context=self._afd_draft_context, + afd_runner=self, ) self.speculator.update_stream = self.update_stream except BaseException: @@ -488,6 +490,9 @@ def execute_model( ) try: with ( + self.connector.attention_step() + if self.connector is not None + else nullcontext(), skip_dp_coordination() if self.afd_config.async_dp else nullcontext(), _use_camp2p_dp_padding( self.afd_config.num_attention_ranks, self.afd_config.num_ffn_ranks diff --git a/afd_plugin/v1/worker/npu/dspark_graph.py b/afd_plugin/v1/worker/npu/dspark_graph.py new file mode 100644 index 000000000..e17e1990e --- /dev/null +++ b/afd_plugin/v1/worker/npu/dspark_graph.py @@ -0,0 +1,149 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Install AFD metadata around native draft graph capture and replay.""" + +from __future__ import annotations + +from collections.abc import Iterator, Mapping +from contextlib import contextmanager +from types import MethodType +from typing import TYPE_CHECKING + +import torch +from vllm.config import CUDAGraphMode +from vllm.v1.kv_cache_interface import KVCacheConfig +from vllm.v1.worker.gpu.block_table import BlockTables +from vllm.v1.worker.gpu.cudagraph_utils import ( + AttentionState, + BatchExecutionDescriptor, +) +from vllm.v1.worker.gpu.input_batch import InputBuffers +from vllm.v1.worker.gpu.spec_decode.dflash import cudagraph as dflash_cudagraph +from vllm.v1.worker.utils import AttentionGroup +from vllm_ascend.worker.v2.spec_decode.dflash.aclgraph import DFlashAclGraphManager + +from afd_plugin.v1.worker.attention_model_runner_v2 import _AFDCaptureEventTracker + +if TYPE_CHECKING: + from afd_plugin.v1.worker.npu.attention_model_runner_v2 import ( + AFDNPUAttentionModelRunnerV2, + ) + + +@contextmanager +def use_afd_draft_capture( + runner: AFDNPUAttentionModelRunnerV2, + manager: DFlashAclGraphManager, +) -> Iterator[None]: + """Associate every native draft capture forward with its own AFD batch.""" + descs = manager._capture_descs.get(CUDAGraphMode.FULL, [])[ + : manager._max_full_descs_to_capture + ] + if not runner.speculative_config.enforce_eager and not descs: + raise RuntimeError("AFD DSpark draft FULL has no native capture descriptors") + tracker = _AFDCaptureEventTracker(descs) + original_prepare = dflash_cudagraph._prepare_dflash_inputs_to_capture + + # Upstream: vLLM 5df63a72, dflash/cudagraph.py; Ascend bd42eea0 delegates + # draft capture here. Signature matches the native preparation function. + # Native input/KV preparation remains in the saved function; only AFD + # metadata is added. Remove when draft capture exposes a plugin callback. + def prepare_inputs( + num_reqs: int, + num_tokens: int, + input_buffers: InputBuffers, + block_tables: BlockTables, + attn_groups: list[list[AttentionGroup]], + kv_cache_config: KVCacheConfig, + max_model_len: int, + skip_attn: bool, + causal: bool | Mapping[int, bool], + ) -> AttentionState: + # ### PATCH START: stage descriptor-local draft metadata. + state = original_prepare( + num_reqs, + num_tokens, + input_buffers, + block_tables, + attn_groups, + kv_cache_config, + max_model_len, + skip_attn, + causal, + ) + desc, is_warmup = tracker.consume(num_reqs, num_tokens) + runner._afd_pending_metadata = runner.build_afd_metadata(None, desc.num_tokens) + runner._afd_suppress_metadata_send = True + runner._is_warmup = is_warmup + runner._afd_is_graph_capturing = not is_warmup + runner._afd_is_graph_replaying = False + return state + # ### PATCH END: stage descriptor-local draft metadata. + + previous_capturing = runner._afd_is_graph_capturing + previous_replaying = runner._afd_is_graph_replaying + dflash_cudagraph._prepare_dflash_inputs_to_capture = prepare_inputs + try: + yield + tracker.assert_complete() + finally: + dflash_cudagraph._prepare_dflash_inputs_to_capture = original_prepare + runner._afd_is_graph_capturing = previous_capturing + runner._afd_is_graph_replaying = previous_replaying + + +@contextmanager +def use_afd_draft_replay( + runner: AFDNPUAttentionModelRunnerV2, + manager: DFlashAclGraphManager, + real_tokens: int, +) -> Iterator[None]: + """Preserve native replay while staging the draft's real and padded rows.""" + state = manager.__dict__ + had_override = "run_fullgraph" in state + previous_override = state.get("run_fullgraph") + original_replay = manager.run_fullgraph + + # Upstream: Ascend bd42eea0 DFlashAclGraphManager.run_fullgraph. + # Signature matches upstream. Delegate graph replay and attention updates; + # add only AFD's model-local metadata. Remove when the native manager has + # an instance-local pre-replay callback. + def run_fullgraph( + self, + desc: BatchExecutionDescriptor, + ) -> torch.Tensor | tuple[torch.Tensor, list[torch.Tensor]]: + # ### PATCH START: install draft-local replay metadata. + previous = ( + runner._afd_pending_metadata, + runner._afd_suppress_metadata_send, + runner._is_warmup, + runner._afd_is_graph_capturing, + runner._afd_is_graph_replaying, + ) + try: + metadata = runner.build_afd_metadata(None, real_tokens) + metadata.tokens_lens = [desc.num_tokens] + runner._afd_pending_metadata = metadata + runner._afd_suppress_metadata_send = True + runner._is_warmup = False + runner._afd_is_graph_capturing = False + runner._afd_is_graph_replaying = True + return original_replay(desc) + finally: + ( + runner._afd_pending_metadata, + runner._afd_suppress_metadata_send, + runner._is_warmup, + runner._afd_is_graph_capturing, + runner._afd_is_graph_replaying, + ) = previous + # ### PATCH END: install draft-local replay metadata. + + try: + manager.run_fullgraph = MethodType(run_fullgraph, manager) + yield + finally: + if had_override: + manager.run_fullgraph = previous_override + else: + del manager.run_fullgraph diff --git a/afd_plugin/v1/worker/npu/dspark_speculator.py b/afd_plugin/v1/worker/npu/dspark_speculator.py index 0de54a59b..41940c87e 100644 --- a/afd_plugin/v1/worker/npu/dspark_speculator.py +++ b/afd_plugin/v1/worker/npu/dspark_speculator.py @@ -4,19 +4,30 @@ from __future__ import annotations -from collections.abc import Callable -from contextlib import AbstractContextManager -from typing import Any +from collections.abc import Callable, Iterator +from contextlib import AbstractContextManager, contextmanager, nullcontext +from typing import TYPE_CHECKING, Any, cast import torch from vllm.config import VllmConfig from vllm.v1.worker.gpu.input_batch import InputBatch from vllm.v1.worker.gpu.spec_decode.dflash import speculator as dflash_speculator +from vllm_ascend.worker.v2.spec_decode.dflash.aclgraph import DFlashAclGraphManager from vllm_ascend.worker.v2.spec_decode.dflash.speculator import ( prepare_dflash_inputs_factory, ) from vllm_ascend.worker.v2.spec_decode.dspark.speculator import AscendDSparkSpeculator +from afd_plugin.v1.worker.npu.dspark_graph import ( + use_afd_draft_capture, + use_afd_draft_replay, +) + +if TYPE_CHECKING: + from afd_plugin.v1.worker.npu.attention_model_runner_v2 import ( + AFDNPUAttentionModelRunnerV2, + ) + class AFDDSparkSpeculator(AscendDSparkSpeculator): def __init__( @@ -25,9 +36,37 @@ def __init__( device: torch.device, *, afd_context: Callable[[bool, bool], AbstractContextManager[None]], + afd_runner: AFDNPUAttentionModelRunnerV2, ) -> None: super().__init__(vllm_config, device) self.afd_context = afd_context + self.afd_runner = afd_runner + + @contextmanager + def _physical_input_preparation(self, *, enabled: bool = True) -> Iterator[None]: + original_prepare = dflash_speculator.prepare_dflash_inputs + if enabled: + dflash_speculator.prepare_dflash_inputs = self._afd_prepare_inputs + try: + yield + finally: + dflash_speculator.prepare_dflash_inputs = original_prepare + + # Upstream: vLLM 5df63a72 DFlashSpeculator.capture, inherited by Ascend + # bd42eea0 DSpark. Signature matches. Delegate native context/KV and Markov + # capture, adding draft-local AFD metadata and the physical block size. + # Remove when native capture exposes metadata and physical-size hooks. + def capture(self) -> None: + # ### PATCH START: scope native draft capture. + assert self.query_cudagraph_manager is not None + manager = cast(DFlashAclGraphManager, self.query_cudagraph_manager) + with ( + self.afd_context(True, False), + self._physical_input_preparation(), + use_afd_draft_capture(self.afd_runner, manager), + ): + super().capture() + # ### PATCH END: scope native draft capture. # Upstream: Ascend bd42eea0, AscendDSparkSpeculator.set_attn. # Its input factory captures the global cache_config.block_size, which can @@ -88,30 +127,41 @@ def propose( num_speculative_tokens: int | None = None, ) -> torch.Tensor: # ### PATCH START: scope only the native draft proposal - prepare_inputs = dflash_speculator.prepare_dflash_inputs - if not (dummy_run and skip_attn_for_dummy_run): - dflash_speculator.prepare_dflash_inputs = self._afd_prepare_inputs - try: - with self.afd_context(dummy_run, is_profile): - return super().propose( - input_batch, - attn_metadata, - slot_mappings, - last_hidden_states, - aux_hidden_states, - num_sampled, - num_rejected, - last_sampled, - next_prefill_tokens, - temperature, - seeds, - None, - dummy_run, - skip_attn_for_dummy_run, - mm_inputs, - is_profile=is_profile, - num_speculative_tokens=num_speculative_tokens, - ) - finally: - dflash_speculator.prepare_dflash_inputs = prepare_inputs + manager = self.query_cudagraph_manager + # Native memory profiling runs before KV/graph-manager initialization. + replay_scope = ( + use_afd_draft_replay( + self.afd_runner, + cast(DFlashAclGraphManager, manager), + input_batch.num_reqs * self.num_query_per_req, + ) + if manager is not None + else nullcontext() + ) + with ( + self.afd_context(dummy_run, is_profile), + self._physical_input_preparation( + enabled=not (dummy_run and skip_attn_for_dummy_run) + ), + replay_scope, + ): + return super().propose( + input_batch, + attn_metadata, + slot_mappings, + last_hidden_states, + aux_hidden_states, + num_sampled, + num_rejected, + last_sampled, + next_prefill_tokens, + temperature, + seeds, + None, + dummy_run, + skip_attn_for_dummy_run, + mm_inputs, + is_profile=is_profile, + num_speculative_tokens=num_speculative_tokens, + ) # ### PATCH END: scope only the native draft proposal diff --git a/afd_plugin/v1/worker/npu/ffn_model_runner.py b/afd_plugin/v1/worker/npu/ffn_model_runner.py index 4cb11f326..72b7764f6 100644 --- a/afd_plugin/v1/worker/npu/ffn_model_runner.py +++ b/afd_plugin/v1/worker/npu/ffn_model_runner.py @@ -284,7 +284,13 @@ def _ffn_forward( # A model whose router is keyed by token identity declares that the FFN # role needs the tokens' ids. Resolve it once per forward, outside the # layer loop, so a per-layer re-read cannot drift within a step. - recv_input_ids = getattr(self.model, "afd_requires_input_ids", False) + # The wire slot carries one payload: either ids or Attention-side routing. + recv_routing = getattr(self.model, "afd_requires_attention_routing", False) + recv_input_ids = ( + False + if recv_routing + else getattr(self.model, "afd_requires_input_ids", False) + ) for layer_idx in _ffn_layer_indices(self): for stage_idx in stage_ids: @@ -315,6 +321,7 @@ def _ffn_forward( layer_idx=layer_idx, max_num_tokens=self.max_num_tokens, recv_input_ids=recv_input_ids, + recv_routing=recv_routing, ) context = payload.context metadata = context.metadata @@ -337,9 +344,18 @@ def _ffn_forward( assert states, "Context.states must not be None" _set_moe_layer_index(forward_context, layer_idx) + # Only a model that declared Attention-side routing gets the + # extra keywords. Passing them unconditionally changes the call + # contract for every other model wrapper, including ones that + # accept only (hidden_states, layer_idx). + ffn_kwargs: dict[str, Any] = {} + if recv_routing: + ffn_kwargs["topk_ids"] = payload.topk_ids + ffn_kwargs["topk_weights"] = payload.topk_weights rank_ffn_output = self.model.compute_ffn_output( hidden_states=hidden_states, layer_idx=layer_idx, + **ffn_kwargs, ) _send_ffn_output( self.connector, diff --git a/afd_plugin/v1/worker/npu/ffn_worker.py b/afd_plugin/v1/worker/npu/ffn_worker.py index e56a7f930..ae7b84dee 100644 --- a/afd_plugin/v1/worker/npu/ffn_worker.py +++ b/afd_plugin/v1/worker/npu/ffn_worker.py @@ -48,6 +48,8 @@ logger = logging.getLogger(__name__) FFN_SHUTDOWN_TIMEOUT_SECONDS = 5 +# Pull receivers must wait for all Attention ranks to stop publishing. +ASYNC_FFN_DRAIN_TIMEOUT_SECONDS = 60 def create_npu_ffn_model_runner( @@ -201,16 +203,17 @@ def _run_ffn_server_loop(self) -> None: return torch.npu.set_device(self.device) - while not event.is_set(): - if isinstance(self.model_runner, AsyncNPUFFNModelRunner): - if not self.model_runner.execute_connector_driven_step(): - return - if ( - self.model_runner.effective_mode == "eager" - and self.model_runner._two_stream is None - ): + if isinstance(self.model_runner, AsyncNPUFFNModelRunner): + # The runner consults the session once per step; a second + # predicate here would add Store round trips to the Pull hot path. + while self.model_runner.execute_connector_driven_step(event.is_set()): + if self.model_runner.effective_mode == "eager": torch.npu.synchronize() - continue + self.model_runner.drain() + self.model_runner.session.finish_receiver() + return + + while not event.is_set(): control_plane = self.model_runner.connector.control_plane assert control_plane is not None payload = control_plane.recv_dp_metadata_list() @@ -232,7 +235,7 @@ def _run_ffn_server_loop(self) -> None: def _requires_async_drain(self) -> bool: return isinstance(self.model_runner, AsyncNPUFFNModelRunner) and ( self.model_runner.effective_mode == "full" - or self.model_runner._two_stream is not None + or self.model_runner.session.requires_drain ) def raise_ffn_loop_error_if_any(self) -> None: @@ -249,7 +252,13 @@ def stop_ffn_server_loop(self) -> None: event.set() thread = self._ffn_thread if thread is not None: - thread.join(timeout=FFN_SHUTDOWN_TIMEOUT_SECONDS) + timeout = ( + ASYNC_FFN_DRAIN_TIMEOUT_SECONDS + if isinstance(self.model_runner, AsyncNPUFFNModelRunner) + and self.model_runner.session.requires_drain + else FFN_SHUTDOWN_TIMEOUT_SECONDS + ) + thread.join(timeout=timeout) if thread.is_alive(): raise RuntimeError( "Async graph receive is pending; retain resources" diff --git a/afd_plugin/validation.py b/afd_plugin/validation.py index 87dcb3255..d399407e0 100644 --- a/afd_plugin/validation.py +++ b/afd_plugin/validation.py @@ -52,6 +52,7 @@ # with the hidden states, as a routed-experts connector (DeepEP M2N) needs. ATTENTION_ROUTING_MODEL_ARCHITECTURES: Final[tuple[str, ...]] = ( "Qwen3MoeForCausalLM", + "Qwen3MoeRoutedExperts", "DeepseekV2ForCausalLM", ) @@ -76,10 +77,12 @@ def requested_ffn_mode(config: VllmConfig) -> Literal["eager", "full"]: return "full" -def validate_graph_capacity(max_tokens: int, tp_size: int, capacity: int) -> None: +def validate_graph_capacity( + max_tokens: int, tp_size: int, capacity: int, *, allow_row_slicing: bool = False +) -> None: if min(max_tokens, tp_size, capacity) <= 0 or max_tokens % tp_size: raise ValueError("Graph dispatch envelope must be positive and divisible by TP") - if capacity < max_tokens: + if capacity < max_tokens and not allow_row_slicing: raise ValueError( "Async FFN graph capacity is smaller than the dispatch envelope" ) @@ -87,6 +90,25 @@ def validate_graph_capacity(max_tokens: int, tp_size: int, capacity: int) -> Non raise ValueError("Async CAM graph source rows exceed the uint16 protocol bound") +def validate_async_cam_attention_graph(vllm_config: VllmConfig) -> None: + """Validate AFD dependencies for native Async CAM Attention graphs.""" + model = vllm_config.model_config + parallel = vllm_config.parallel_config + if ( + not vllm_config.use_v2_model_runner + or model.hf_text_config.model_type != "deepseek_v4" + ): + raise RuntimeError("Async CAM Attention graph requires MRV2 DeepSeek-V4") + if parallel.enable_dbo or parallel.use_ubatching: + raise RuntimeError("Async CAM Attention graph does not support DBO/ubatching") + if ( + parallel.pipeline_parallel_size, + parallel.prefill_context_parallel_size, + parallel.decode_context_parallel_size, + ) != (1, 1, 1): + raise RuntimeError("Async CAM Attention graph does not support PP or CP") + + def validate_gpu_model_runner_v2_config( vllm_config: VllmConfig, *, @@ -205,18 +227,16 @@ def validate_npu_model_runner_v2_config( ) if not vllm_config.model_config.enforce_eager: if expected_role == "attention": - raise RuntimeError( - "AFD NPU ModelRunnerV2 Async CAM requires eager execution " - "on Attention" - ) - requested_ffn_mode(vllm_config) - elif ( - afd_config.connector != CAMP2P_CONNECTOR or afd_config.compute_gate_on_attention - ): + validate_async_cam_attention_graph(vllm_config) + else: + requested_ffn_mode(vllm_config) + elif afd_config.connector != CAMP2P_CONNECTOR: raise RuntimeError( - "AFD NPU ModelRunnerV2 requires synchronous CAMP2pAFDConnector " - "with compute_gate_on_attention=false, or Async CAM", + "AFD NPU ModelRunnerV2 requires either synchronous CAMP2pAFDConnector " + "or Async CAM", ) + # The sync Attention-side gate limits (connector, target, TP, DBO/ubatching) + # live in feature_validation and run from both workers before this point. parallel = vllm_config.parallel_config if ( diff --git a/csrc/npu/README.md b/csrc/npu/README.md index 535610ee6..0f1c8718c 100644 --- a/csrc/npu/README.md +++ b/csrc/npu/README.md @@ -9,8 +9,8 @@ Attention-to-Expert and Expert-to-Attention transfers: The 910C build also contains four experimental routed-only CAM operators: `afd_async_dispatch_send`, `afd_async_dispatch_recv`, `afd_async_combine_send`, and `afd_async_combine_recv`, under the same `torch.ops.afd_ascend` namespace. -Their compact protocol and future connector migration are documented in -[the routed-only operator guide](../../docs/npu/CAM_ASYNC_ROUTED_OPS.md). +Their compact protocol is documented in +[the routed-only operator guide](../../docs/design/npu/cam_async_routed_ops.md). The NPU native sources live in this directory: @@ -42,6 +42,42 @@ The supported SOC generations and the operators compiled for each are declared in `ascend_kernels/operator_registry.json`; `scripts/select_ops.py` resolves the operator list from that registry for a given SOC. +## Pull Mega MC2 dependency + +Set `connector_extra_config.ffn_mega_moe=true` on every Attention and FFN +rank to enable MegaMoE and its dedicated pull protocol. With the switch off, +CAM retains its original push Dispatch/Combine and FFN execution path. +Pull is not a standalone transport option: there is no `dispatch_mode` +setting or separate pull DispatchRecv/CombineSend operator. MegaMoE requires +`dynamicQuant=1` and per-channel W4A8 or W8A8 weights. +Per-group weights are not supported by MegaMoE. + +The A3 `AfdMegaMoePullNative` kernel directly includes the installed CANN MC2 +W4A8 `BlockMmad`. The build supports the CANN 9.1 MegaMoE header layout and +the CANN 9.2 layout with Catlass under `ascendc/common/catlass/include`. Set +`ASCEND_CANN_PACKAGE_PATH` or `ASCEND_HOME_PATH` to the CANN installation. +When this operator is selected, configuration fails if its MC2 headers are +missing; builds that exclude it do not need this dependency. The dependency +tree is not copied into this repository. + +One native launch selects a layer and drains one source transaction. Receive +capacity follows ordinary CAM: `int(262144 * BATCH_SIZE_FACTOR)`, default 262144 +rows. Only transactions exceeding this capacity use multiple device rounds; +capacity is independent of the MC2 matrix tile size. Callers must not launch +once per receive round. The workspace allocation ABI and Attention Combine +interface are unchanged. Both quantization plans use SwiGLU. W4A8 applies the +model `swiglu_limit` in both ACT paths as +`SiLU(min(gate, L)) * clamp(up, -L, L)`; zero, the default, disables clipping. +W8A8 does not apply `swiglu_limit`, matching the Layered W8A8 path. + +For W4A8 intermediate widths up to 3072, ACT directly uses the installed MC2 +W4A8 two-stage UB epilogue and expert-group synchronization. Wider supported +W4A8 rows and W8A8 use bounded-column activation. Return-address caching is +bounded independently of receive capacity. Large capacity reserves HBM scratch +without computing padded rows; for H4096/I2048 the default capacity needs about +12.5 GiB per FFN executor, excluding weights and communication windows. This +includes a 2 GiB buffer for BF16 output staging. + ## Default Behavior Ascend custom ops build by default. Set `AFD_BUILD_ASCEND_OPS=0` only when @@ -141,17 +177,17 @@ paired vendor setup. ## Current Scope -A2E/E2A continue to serve the existing CAMP2P connector. With Attention-side -gating, `CAMAsyncAFDConnector` selects the four 910C routed-only CAM operators. -The standalone FFN entry point reuses this same protocol and connector; see +A2E/E2A serve `CAMP2pAFDConnector`. `CAMAsyncAFDConnector` uses the plugin's +910C routed-only CAM operators; see the [operator contract](../../docs/design/npu/cam_async_routed_ops.md). +The standalone FFN entry point reuses this protocol and connector; see [the standalone service guide](../../docs/npu/STANDALONE_FFN_SERVICE.md). `grouped_matmul_swiglu_quant_v2_layered` is a fused grouped-matmul + SiLU + dynamic-quantization CAM operator (A8W8 / A4W4 / A8W4-MSD variants) built for `ascend910_93`. It is registered as `torch.ops.afd_ascend.gmm_swiglu_quant_v2_layered` (inference + Meta; no -autograd grad) and is excluded from `ascend950` builds. The Async CAM layered -W4A8 executor selects it with `AFD_ASYNC_CAM_LAYERED_GMM=1`. +autograd grad) and is excluded from `ascend950` builds. The default Async CAM +layered executor uses it for W4A8 and W8A8. The trailing `swiglu_limit` argument defaults to `0.0` (no clipping). A finite positive `L` computes `SiLU(min(gate, L)) * clamp(up, -L, L)` before output diff --git a/csrc/npu/ascend_kernels/afd_mega_moe_comm/op_api/aclnn_afd_mega_moe_comm.cpp b/csrc/npu/ascend_kernels/afd_mega_moe_comm/op_api/aclnn_afd_mega_moe_comm.cpp new file mode 100644 index 000000000..c552024f7 --- /dev/null +++ b/csrc/npu/ascend_kernels/afd_mega_moe_comm/op_api/aclnn_afd_mega_moe_comm.cpp @@ -0,0 +1,18 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +#include "aclnn_afd_mega_moe_comm.h" +#include "aclnnInner_afd_mega_moe_comm.h" +extern "C" void __attribute__((weak)) NnopbaseSetHcclServerType(void *, int32_t); +extern "C" aclnnStatus aclnnAfdMegaMoeCommGetWorkspaceSize( + const aclTensor *x, const aclTensor *expertIds, const aclTensor *expertScales, + const aclTensor *state, const aclIntArray *config, char *hcclGroupName, + const aclTensor *combined, const aclTensor *completion, + uint64_t *workspaceSize, aclOpExecutor **executor) { + return aclnnInnerAfdMegaMoeCommGetWorkspaceSize(x, expertIds, expertScales, state, + config, hcclGroupName, combined, completion, workspaceSize, executor); +} +extern "C" aclnnStatus aclnnAfdMegaMoeComm(void *workspace, uint64_t workspaceSize, + aclOpExecutor *executor, aclrtStream stream) { + if (NnopbaseSetHcclServerType) NnopbaseSetHcclServerType(executor, 1); + return aclnnInnerAfdMegaMoeComm(workspace, workspaceSize, executor, stream); +} diff --git a/csrc/npu/ascend_kernels/afd_mega_moe_comm/op_api/aclnn_afd_mega_moe_comm.h b/csrc/npu/ascend_kernels/afd_mega_moe_comm/op_api/aclnn_afd_mega_moe_comm.h new file mode 100644 index 000000000..0a036ee8d --- /dev/null +++ b/csrc/npu/ascend_kernels/afd_mega_moe_comm/op_api/aclnn_afd_mega_moe_comm.h @@ -0,0 +1,13 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +#pragma once +#include "aclnn/aclnn_base.h" +extern "C" { +__attribute__((visibility("default"))) aclnnStatus aclnnAfdMegaMoeCommGetWorkspaceSize( + const aclTensor *x, const aclTensor *expertIds, const aclTensor *expertScales, + const aclTensor *state, const aclIntArray *config, char *hcclGroupName, + const aclTensor *combined, const aclTensor *completion, + uint64_t *workspaceSize, aclOpExecutor **executor); +__attribute__((visibility("default"))) aclnnStatus aclnnAfdMegaMoeComm( + void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream); +} diff --git a/csrc/npu/ascend_kernels/afd_mega_moe_comm/op_host/afd_mega_moe_comm.cpp b/csrc/npu/ascend_kernels/afd_mega_moe_comm/op_host/afd_mega_moe_comm.cpp new file mode 100644 index 000000000..f529558c4 --- /dev/null +++ b/csrc/npu/ascend_kernels/afd_mega_moe_comm/op_host/afd_mega_moe_comm.cpp @@ -0,0 +1,73 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +#include +#include +#include "register/op_def_registry.h" +#include "tiling/platform/platform_ascendc.h" +#include "tiling/hccl/hccl_tiling.h" +#include "../op_kernel/afd_mega_moe_comm_tiling.h" +#include "../op_kernel/afd_pull_layout.h" +#include "ops_log.h" +namespace optiling { +static ge::graphStatus MegaCommTiling(gert::TilingContext *ctx) { + const auto attrs = ctx->GetAttrs(); + const auto config = attrs->GetListInt(0); + if (!config || config->GetSize() != 11) return ge::GRAPH_FAILED; + const auto c = config->GetData(); + if ((c[0] != 0 && c[0] != 1 && c[0] != 4) || c[8] <= 0 || c[1] % c[8] || c[9] < 0 || + c[7] < 0 || c[7] >= c[4] + c[5] || c[10] < 0 || c[10] >= 65536 || + ((c[0] == 1 || c[0] == 4) && (c[7] >= c[5] || c[9] > c[1] / c[8]))) return ge::GRAPH_FAILED; + Cam::Pull::Shape shape{uint64_t(c[1] / c[8]), uint64_t(c[2]), uint64_t(c[3]), + uint64_t(c[4]), uint64_t(c[5]), uint64_t(c[6]), uint64_t(c[8])}; + Cam::Pull::Layout layout{}; + auto platform = platform_ascendc::PlatformAscendC(ctx->GetPlatformInfo()); + const auto aiv = platform.GetCoreNumAiv(); + uint64_t ub = 0; + platform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ub); + if (!aiv || aiv > Cam::Pull::MAX_AIVS || !Cam::Pull::MakeLayout(shape, layout) || + layout.bytes > GetMaxWindowSize() || Cam::Pull::TransportUbBytes(shape) > ub || + (c[0] == 1 && Cam::Pull::PublisherUbBytes(shape, aiv) + Cam::Pull::TransportUbBytes(shape) > ub) || + (c[0] == 4 && Cam::Pull::CombineUbBytes(shape, aiv) > ub)) + return ge::GRAPH_FAILED; + auto data = ctx->GetTilingData(); + std::memset(data, 0, sizeof(*data)); + std::copy(c, c + 11, data->config); + data->aivNum = aiv; data->ubSize = ub; data->windowSize = GetMaxWindowSize(); + AscendC::Mc2CcTilingConfig comm(attrs->GetAttrPointer(1), 8U, + "AlltoAll=level0:fullmesh;level1:pairwise"); + comm.GetTiling(data->mc2InitTiling); comm.GetTiling(data->mc2CcTiling1); + ctx->SetBlockDim(aiv); ctx->SetTilingKey(1); ctx->SetScheduleMode(1); + ctx->GetWorkspaceSizes(1)[0] = 16 * 1024 * 1024; + return ge::GRAPH_SUCCESS; +} +} +namespace ops { +class AfdMegaMoeComm : public OpDef { +public: + explicit AfdMegaMoeComm(const char *name) : OpDef(name) { + Input("x").ParamType(REQUIRED).DataType({ge::DT_BF16}).Format({ge::FORMAT_ND}); + Input("expertIds").ParamType(REQUIRED).DataType({ge::DT_INT32}).Format({ge::FORMAT_ND}); + Input("expertScales").ParamType(REQUIRED).DataType({ge::DT_FLOAT}).Format({ge::FORMAT_ND}); + Input("state").ParamType(REQUIRED).DataType({ge::DT_FLOAT16}).Format({ge::FORMAT_ND}); + Output("combined").ParamType(REQUIRED).DataType({ge::DT_BF16}).Format({ge::FORMAT_ND}); + Output("completion").ParamType(REQUIRED).DataType({ge::DT_INT64}).Format({ge::FORMAT_ND}); + Attr("config").ListInt(); Attr("hccl_group_name").String(); + SetInferShape([](gert::InferShapeContext *ctx) { + const auto c = ctx->GetAttrs()->GetListInt(0)->GetData(); + const int64_t sizes[] = {c[0] == 4 ? std::max(1, c[9]) * c[2] : 1, 1}; + for (uint32_t j = 0; j < 2; ++j) { + auto shape = ctx->GetOutputShape(j); shape->SetDimNum(1); shape->SetDim(0, sizes[j]); + } + return ge::GRAPH_SUCCESS; + }); + SetInferDataType([](gert::InferDataTypeContext *ctx) { + const ge::DataType types[] = {ge::DT_BF16,ge::DT_INT64}; + for (uint32_t j = 0; j < 2; ++j) ctx->SetOutputDataType(j, types[j]); + return ge::GRAPH_SUCCESS; + }); + AICore().AddConfig("ascend910_93"); AICore().SetTiling(optiling::MegaCommTiling); + MC2().HcclGroup({"hccl_group_name"}); + } +}; +OP_ADD(AfdMegaMoeComm); +} diff --git a/csrc/npu/ascend_kernels/afd_mega_moe_comm/op_kernel/afd_mega_moe_comm.cpp b/csrc/npu/ascend_kernels/afd_mega_moe_comm/op_kernel/afd_mega_moe_comm.cpp new file mode 100644 index 000000000..3ef0b4bc3 --- /dev/null +++ b/csrc/npu/ascend_kernels/afd_mega_moe_comm/op_kernel/afd_mega_moe_comm.cpp @@ -0,0 +1,306 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +#include "kernel_operator.h" +#include "afd_mega_moe_comm_tiling.h" +#include "afd_pull_protocol.h" + +namespace AfdMegaComm { +using namespace AscendC; +using Cam::Pull::Fence; + +// Publish only to the local A window. Statistics retain the existing Combine prefix layout. +__aicore__ inline void Publish(GM_ADDR x, GM_ADDR ids, GM_ADDR state, + const Cam::AfdMegaMoeCommTilingData &t, TPipe &pipe, + Cam::Pull::Transport &control, const Cam::Pull::Layout &layout) { + const auto c = t.config; + const uint32_t id = GetBlockIdx(), cores = t.aivNum, h = c[2], k = c[3], experts = c[4] * c[6]; + const uint32_t columns = c[4] + experts, rows = c[9]; + const uint32_t begin = rows / cores * id + (id < rows % cores ? id : rows % cores); + const uint32_t end = begin + rows / cores + (id < rows % cores ? 1 : 0); + const auto local = Cam::Pull::Window(c[7], c[7]); + TBuf statsBuf, offsetsBuf, idsBuf, mathBuf; + pipe.InitBuffer(statsBuf, Cam::Pull::PublisherStatsBytes( + {uint64_t(c[1] / c[8]), h, k, uint64_t(c[4]), uint64_t(c[5]), uint64_t(c[6]), uint64_t(c[8])}, cores)); + pipe.InitBuffer(offsetsBuf, 2 * Cam::Pull::Align(4ULL * experts)); + pipe.InitBuffer(idsBuf, Cam::Pull::Align(4ULL * k)); + pipe.InitBuffer(mathBuf, 11ULL * h + 128); + auto stats = statsBuf.Get(); + auto starts = offsetsBuf.Get(); + auto prefixes = offsetsBuf.GetWithOffset(experts, Cam::Pull::Align(4ULL * experts)); + auto routes = idsBuf.Get(); + auto input = mathBuf.Get(); + auto value = mathBuf.GetWithOffset(h, 2 * h); + auto temp = mathBuf.GetWithOffset(h, 6 * h); + auto quant = mathBuf.GetWithOffset(h + 32, 10 * h); + auto maximum = mathBuf.GetWithOffset(8, 11 * h + 32); + GlobalTensor globalStats; + GlobalTensor globalIds; + GlobalTensor globalX; + GlobalTensor payload; + globalStats.SetGlobalBuffer(reinterpret_cast<__gm__ uint32_t *>(local)); + globalIds.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t *>(ids)); + globalX.SetGlobalBuffer(reinterpret_cast<__gm__ bfloat16_t *>(x)); + payload.SetGlobalBuffer(reinterpret_cast<__gm__ int8_t *>(local + layout.payload)); + Duplicate(stats, uint32_t(0), columns); + Fence(); + int64_t error = 0; + for (uint32_t token = begin; token < end; ++token) { + DataCopyPad(routes, globalIds[token * k], {1U, 4U * k, 0U, 0U, 0U}, {false, 0U, 0U, 0U}); + Fence(); + for (uint32_t j = 0; j < k; ++j) { + const auto expert = routes.GetValue(j); + if (expert < 0 || uint32_t(expert) >= experts) error = -1; + else stats.SetValue(c[4] + expert, stats.GetValue(c[4] + expert) + 1); + } + Fence(); + } + const auto lanes = Cam::Pull::LaneErrorsOffset(c[5] / c[8], control.ReadState(Cam::Pull::MAP_SIZE)); + control.Write(state + (lanes + id * Cam::Pull::WORDS_PER_LINE) * 8, error); + Fence(); + DataCopyPad(globalStats[id * columns], stats, {1U, 4U * columns, 0U, 0U, 0U}); + Fence(); + SyncAll(); + // Each lane owns adjacent expert columns; UB rows are padded independently of GM rows. + const uint32_t firstExpert = id * (experts / cores) + (id < experts % cores ? id : experts % cores); + const uint32_t expertCount = experts / cores + (id < experts % cores); + if (expertCount) { + const uint32_t stride = Cam::Pull::Align32(4U * expertCount) / 4; + DataCopyPad(stats, globalStats[c[4] + firstExpert], + {uint16_t(cores), 4U * expertCount, 4U * (columns - expertCount), 0U, 0U}, + {false, 0U, 0U, 0U}); + Fence(); + for (uint32_t lane = 1; lane < cores; ++lane) + for (uint32_t e = 0; e < expertCount; ++e) { + const auto pos = lane * stride + e; + stats.SetValue(pos, stats.GetValue(pos) + stats.GetValue(pos - stride)); + } + Fence(); + DataCopyPad(globalStats[c[4] + firstExpert], stats, + {uint16_t(cores), 4U * expertCount, 0U, 4U * (columns - expertCount), 0U}); + Fence(); + } + if (id == 0) { + GlobalTensor errors; + errors.SetGlobalBuffer(reinterpret_cast<__gm__ int64_t *>(state) + lanes); + DataCopy(stats.ReinterpretCast(), errors, cores * Cam::Pull::WORDS_PER_LINE); + Fence(); + for (uint32_t lane = 0; lane < cores; ++lane) { + const auto failure = stats.ReinterpretCast().GetValue(lane * Cam::Pull::WORDS_PER_LINE); + if (!error && failure) error = failure; + } + if (error) control.WriteState(Cam::Pull::FAILURE, error); + } + SyncAll(); + error = control.ReadState(Cam::Pull::FAILURE); + DataCopyPad(stats, globalStats[(cores - 1) * columns + c[4]], + {1U, 4U * experts, 0U, 0U, 0U}, {false, 0U, 0U, 0U}); + if (id) DataCopyPad(starts, globalStats[(id - 1) * columns + c[4]], + {1U, 4U * experts, 0U, 0U, 0U}, {false, 0U, 0U, 0U}); + Fence(); + uint32_t prefix = 0; + for (uint32_t e = 0; e < experts; ++e) { + prefixes.SetValue(e, prefix); + if (!id) starts.SetValue(e, 0); + prefix += stats.GetValue(e); + } + for (uint32_t token = begin; !error && token < end; ++token) { + DataCopyPad(routes, globalIds[token * k], {1U, 4U * k, 0U, 0U, 0U}, {false, 0U, 0U, 0U}); + DataCopy(input, globalX[token * h], h); + Fence(); + // Preserve the original Dispatch's FP32 -> INT32 -> FP16 -> INT8 rounding. + Cast(value, input, RoundMode::CAST_NONE, h); + PipeBarrier(); + Abs(temp, value, h); + PipeBarrier(); + ReduceMax(maximum, temp, temp, h, false); + PipeBarrier(); + Duplicate(temp, 127.0f, 1); + PipeBarrier(); + Div(temp, temp, maximum, 1); + Fence(); + const float multiplier = temp.GetValue(0); + Fence(); + Muls(value, value, multiplier, h); + PipeBarrier(); + Cast(value.ReinterpretCast(), value, RoundMode::CAST_RINT, h); + PipeBarrier(); + SetDeqScale(static_cast(1.0f)); + PipeBarrier(); + Cast(value.ReinterpretCast(), value.ReinterpretCast(), RoundMode::CAST_ROUND, h); + PipeBarrier(); + Cast(quant, value.ReinterpretCast(), RoundMode::CAST_TRUNC, h); + Fence(); + quant.ReinterpretCast().SetValue(h / sizeof(float), 1.0f / multiplier); + Fence(); + Fence(); + for (uint32_t j = 0; j < k; ++j) { + const uint32_t expert = routes.GetValue(j); + const uint32_t row = prefixes.GetValue(expert) + starts.GetValue(expert); + DataCopy(payload[row * layout.rowStride], quant, h + 32); + Fence(); + starts.SetValue(expert, starts.GetValue(expert) + 1); + } + Fence(); + Fence(); + } + SyncAll(); + if (id == 0) { + const auto seq = control.ReadState(Cam::Pull::ACTIVE_SEQUENCE) + 1; + control.WriteState(Cam::Pull::ACTIVE_SEQUENCE, seq); + auto counts = offsetsBuf.Get(); + for (uint32_t e = 0; e < experts; ++e) counts.SetValue(e, stats.GetValue(e)); + GlobalTensor publishedCounts; + publishedCounts.SetGlobalBuffer(reinterpret_cast<__gm__ int64_t *>(local + layout.counts)); + Fence(); + DataCopyPad(publishedCounts, counts, {1U, 8U * experts, 0U, 0U, 0U}); + Fence(); + const int64_t values[] = {seq, c[10], rows, prefix, seq, error}; + for (uint32_t field = Cam::Pull::SEQUENCE; field <= Cam::Pull::ERROR; ++field) + if (field != Cam::Pull::COMMIT) + control.Write(local + layout.header + field * Cam::Pull::ALIGNMENT, values[field - 1]); + control.Write(local + layout.header + Cam::Pull::COMMIT * Cam::Pull::ALIGNMENT, seq); + } +} + +__aicore__ inline bool WaitCombine(GM_ADDR state, const Cam::AfdMegaMoeCommTilingData &t, + const Cam::Pull::Layout &layout, Cam::Pull::Transport &control) { + const auto c = t.config; + const auto local = Cam::Pull::Window(c[7], c[7]); + if (GetBlockIdx() == 0) while (true) { + auto error = control.ReadState(Cam::Pull::FAILURE); + if (!error) error = control.Read(local + layout.header + Cam::Pull::ERROR * Cam::Pull::ALIGNMENT); + bool ready = true; + for (uint32_t f = 0; f < c[4]; ++f) { + if (!error) error = control.Read(local + layout.errors + f * Cam::Pull::ALIGNMENT); + if (!uint32_t(control.Read(local + layout.flags + 4ULL * f))) ready = false; + } + if (error) { + control.WriteState(Cam::Pull::FAILURE, error); + break; + } + if (ready) break; + } + SyncAll(); + return !control.ReadState(Cam::Pull::FAILURE); +} + +// Publish and Combine partition tokens identically; only the last and preceding +// statistics rows are needed to recover each route's returned-row position. +__aicore__ inline void Combine(GM_ADDR ids, GM_ADDR probabilities, GM_ADDR combined, + const Cam::AfdMegaMoeCommTilingData &t, + const Cam::Pull::Layout &layout, TPipe &pipe) { + const auto c = t.config; + const uint32_t id = GetBlockIdx(), cores = t.aivNum, h = c[2], k = c[3], rows = c[9]; + const uint32_t experts = c[4] * c[6], columns = c[4] + experts; + const uint32_t begin = rows / cores * id + (id < rows % cores ? id : rows % cores); + const uint32_t end = begin + rows / cores + (id < rows % cores); + const auto local = Cam::Pull::Window(c[7], c[7]); + pipe.Reset(); + TBuf offsetsBuf, routesBuf, mathBuf; + pipe.InitBuffer(offsetsBuf, 2 * Cam::Pull::Align(4ULL * experts)); + pipe.InitBuffer(routesBuf, 2 * Cam::Pull::Align(4ULL * k)); + pipe.InitBuffer(mathBuf, 14ULL * h); + auto offsets = offsetsBuf.Get(); + auto counts = offsetsBuf.GetWithOffset(experts, Cam::Pull::Align(4ULL * experts)); + auto routes = routesBuf.Get(); + auto weights = routesBuf.GetWithOffset(k, Cam::Pull::Align(4ULL * k)); + auto input = mathBuf.Get(); + auto value = mathBuf.GetWithOffset(h, 2 * h); + auto product = mathBuf.GetWithOffset(h, 6 * h); + auto sum = mathBuf.GetWithOffset(h, 10 * h); + GlobalTensor stats; + GlobalTensor globalIds; + GlobalTensor globalWeights; + GlobalTensor returned, output; + stats.SetGlobalBuffer(reinterpret_cast<__gm__ uint32_t *>(local)); + globalIds.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t *>(ids)); + globalWeights.SetGlobalBuffer(reinterpret_cast<__gm__ float *>(probabilities)); + returned.SetGlobalBuffer(reinterpret_cast<__gm__ bfloat16_t *>(local + layout.returned)); + output.SetGlobalBuffer(reinterpret_cast<__gm__ bfloat16_t *>(combined)); + if (begin < end) { + DataCopyPad(counts, stats[(cores - 1) * columns + c[4]], + {1U, 4U * experts, 0U, 0U, 0U}, {false, 0U, 0U, 0U}); + if (id) DataCopyPad(offsets, stats[(id - 1) * columns + c[4]], + {1U, 4U * experts, 0U, 0U, 0U}, {false, 0U, 0U, 0U}); + Fence(); + uint32_t prefix = 0; + for (uint32_t e = 0; e < experts; ++e) { + offsets.SetValue(e, prefix + (id ? offsets.GetValue(e) : 0)); + prefix += counts.GetValue(e); + } + for (uint32_t token = begin; token < end; ++token) { + DataCopyPad(routes, globalIds[token * k], + {1U, 4U * k, 0U, 0U, 0U}, {false, 0U, 0U, 0U}); + DataCopyPad(weights, globalWeights[token * k], + {1U, 4U * k, 0U, 0U, 0U}, {false, 0U, 0U, 0U}); + Duplicate(sum, 0.0f, h); + Fence(); + for (uint32_t j = 0; j < k; ++j) { + const uint32_t expert = routes.GetValue(j), row = offsets.GetValue(expert); + const float weight = weights.GetValue(j); + DataCopy(input, returned[uint64_t(row) * h], h); + Fence(); + // Keep the existing FP32 multiply/add and top-k accumulation order. + Cast(value, input, RoundMode::CAST_NONE, h); + PipeBarrier(); + Muls(product, value, weight, h); + PipeBarrier(); + Add(sum, sum, product, h); + Fence(); + offsets.SetValue(expert, row + 1); + } + PipeBarrier(); + Cast(input, sum, RoundMode::CAST_RINT, h); + Fence(); + DataCopy(output[uint64_t(token) * h], input, h); + Fence(); + Fence(); + Fence(); + } + } + // All returned rows have been consumed before the single source window is reusable. + SyncAll(); + if (id == 0) { + GlobalTensor flags; + flags.SetGlobalBuffer(reinterpret_cast<__gm__ uint32_t *>(local + layout.flags)); + Duplicate(offsets, uint32_t(0), uint32_t(c[4])); + Fence(); + DataCopyPad(flags, offsets, {1U, uint32_t(4 * c[4]), 0U, 0U, 0U}); + Fence(); + } +} +} // namespace AfdMegaComm + +extern "C" __global__ __aicore__ void afd_mega_moe_comm( + GM_ADDR x, GM_ADDR ids, GM_ADDR probabilities, GM_ADDR state, GM_ADDR combined, + GM_ADDR completion, GM_ADDR workspace, GM_ADDR tiling) { + using namespace AfdMegaComm; + REGISTER_TILING_DEFAULT(Cam::AfdMegaMoeCommTilingData); + REGISTER_TILING_FOR_TILINGKEY("TILING_KEY_VAR == 1", Cam::AfdMegaMoeCommTilingData); + if (TILING_KEY_IS(1)) { + GET_TILING_DATA_WITH_STRUCT(Cam::AfdMegaMoeCommTilingData, t, tiling); + const auto c = t.config; + const Cam::Pull::Shape shape{uint64_t(c[1] / c[8]), uint64_t(c[2]), uint64_t(c[3]), + uint64_t(c[4]), uint64_t(c[5]), uint64_t(c[6]), uint64_t(c[8])}; + Cam::Pull::Layout layout{}; + Cam::Pull::MakeLayout(shape, layout); + TPipe pipe; + { + Cam::Pull::Transport transport(state, shape, c[7], t.aivNum, pipe); + SyncAll(); + if (transport.ReadState(Cam::Pull::FAILURE)) {} + else if (c[0] == 0) transport.Initialize(); + else if (c[0] == 1) Publish(x, ids, state, t, pipe, transport, layout); + else if (c[0] == 4 && WaitCombine(state, t, layout, transport)) { + Combine(ids, probabilities, combined, t, layout, pipe); + } + } + SyncAll(); + // Release phase buffers before allocating the final status transport. + pipe.Reset(); + if (GetBlockIdx() == 0) { + Cam::Pull::Transport final(state, shape, c[7], t.aivNum, pipe); + final.Write(completion, final.ReadState(Cam::Pull::FAILURE)); + } + } +} diff --git a/csrc/npu/ascend_kernels/afd_mega_moe_comm/op_kernel/afd_mega_moe_comm_tiling.h b/csrc/npu/ascend_kernels/afd_mega_moe_comm/op_kernel/afd_mega_moe_comm_tiling.h new file mode 100644 index 000000000..ec0304778 --- /dev/null +++ b/csrc/npu/ascend_kernels/afd_mega_moe_comm/op_kernel/afd_mega_moe_comm_tiling.h @@ -0,0 +1,15 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +#pragma once +#include "kernel_tiling/kernel_tiling.h" +namespace Cam { +// mode,maxSeqLen,H,K,F,A,E,rank,TP,rows,layer. +struct AfdMegaMoeCommTilingData { + Mc2InitTiling mc2InitTiling; + Mc2CcTiling mc2CcTiling1; + Mc2CcTiling mc2CcTiling2; + int64_t config[11]; + uint32_t aivNum; + uint64_t ubSize, windowSize; +}; +} diff --git a/csrc/npu/ascend_kernels/afd_mega_moe_pull_native/op_api/aclnn_afd_mega_moe_pull_native.cpp b/csrc/npu/ascend_kernels/afd_mega_moe_pull_native/op_api/aclnn_afd_mega_moe_pull_native.cpp new file mode 100644 index 000000000..216c798e2 --- /dev/null +++ b/csrc/npu/ascend_kernels/afd_mega_moe_pull_native/op_api/aclnn_afd_mega_moe_pull_native.cpp @@ -0,0 +1,22 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +#include "aclnn_afd_mega_moe_pull_native.h" +#include "aclnnInner_afd_mega_moe_pull_native.h" +extern "C" void __attribute__((weak)) NnopbaseSetHcclServerType(void *, int32_t); +extern "C" aclnnStatus aclnnAfdMegaMoePullNativeGetWorkspaceSize( + const aclTensor *anchor, const aclTensor *commArgs, const aclTensor *layerMap, + const aclTensor *scratch, const aclTensorList *w13, const aclTensorList *w2, + const aclTensorList *s13, const aclTensorList *s2, const aclTensorList *b13, + const aclTensorList *b2, const aclIntArray *config, char *hcclGroupName, double routedScale, + double swigluLimit, const aclTensor *completion, uint64_t *workspaceSize, + aclOpExecutor **executor) { + return aclnnInnerAfdMegaMoePullNativeGetWorkspaceSize( + anchor, commArgs, layerMap, scratch, w13, w2, s13, s2, b13, b2, config, hcclGroupName, + routedScale, swigluLimit, completion, workspaceSize, executor); +} +extern "C" aclnnStatus aclnnAfdMegaMoePullNative(void *workspace, uint64_t workspaceSize, + aclOpExecutor *executor, aclrtStream stream) { + if (NnopbaseSetHcclServerType) + NnopbaseSetHcclServerType(executor, 1); + return aclnnInnerAfdMegaMoePullNative(workspace, workspaceSize, executor, stream); +} diff --git a/csrc/npu/ascend_kernels/afd_mega_moe_pull_native/op_api/aclnn_afd_mega_moe_pull_native.h b/csrc/npu/ascend_kernels/afd_mega_moe_pull_native/op_api/aclnn_afd_mega_moe_pull_native.h new file mode 100644 index 000000000..7174ab8ce --- /dev/null +++ b/csrc/npu/ascend_kernels/afd_mega_moe_pull_native/op_api/aclnn_afd_mega_moe_pull_native.h @@ -0,0 +1,21 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +#pragma once +#include "aclnn/aclnn_base.h" +#ifdef __cplusplus +extern "C" { +#endif +__attribute__((visibility("default"))) aclnnStatus aclnnAfdMegaMoePullNativeGetWorkspaceSize( + const aclTensor *anchor, const aclTensor *commArgs, const aclTensor *layerMap, + const aclTensor *scratch, const aclTensorList *w13, const aclTensorList *w2, + const aclTensorList *s13, const aclTensorList *s2, const aclTensorList *b13, + const aclTensorList *b2, const aclIntArray *config, char *hcclGroupName, double routedScale, + double swigluLimit, const aclTensor *completion, uint64_t *workspaceSize, + aclOpExecutor **executor); +__attribute__((visibility("default"))) aclnnStatus aclnnAfdMegaMoePullNative(void *workspace, + uint64_t workspaceSize, + aclOpExecutor *executor, + aclrtStream stream); +#ifdef __cplusplus +} +#endif diff --git a/csrc/npu/ascend_kernels/afd_mega_moe_pull_native/op_host/afd_mega_moe_pull_native.cpp b/csrc/npu/ascend_kernels/afd_mega_moe_pull_native/op_host/afd_mega_moe_pull_native.cpp new file mode 100644 index 000000000..669ae6c7d --- /dev/null +++ b/csrc/npu/ascend_kernels/afd_mega_moe_pull_native/op_host/afd_mega_moe_pull_native.cpp @@ -0,0 +1,106 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +#include +#include +#include +#include "register/op_def_registry.h" +#include "tiling/platform/platform_ascendc.h" +#include "tiling/hccl/hccl_tiling.h" +#include "../op_kernel/afd_mega_moe_pull_native_tiling.h" +#include "../op_kernel/afd_pull_layout.h" +#include "ops_log.h" + +namespace optiling { +static ge::graphStatus NativeMegaTiling(gert::TilingContext *ctx) { + const auto attrs = ctx->GetAttrs(); + const auto config = attrs->GetListInt(0); + if (config == nullptr || config->GetSize() != 14) return ge::GRAPH_FAILED; + const auto c = config->GetData(); + for (uint32_t j = 0; j < 14; ++j) + if (c[j] <= 0 || c[j] > 262144) return ge::GRAPH_FAILED; + if (c[8] != c[4] + c[5] || c[7] < c[5] || c[7] >= c[8] || + c[9] == 0 || c[0] % c[9] || c[5] % c[9] || + c[2] % 256 || c[2] > 8192 || c[6] > 128 || + c[10] > Cam::NATIVE_MAX_ROWS || (c[11] != 1 && c[11] != 2) || c[12] != 1 || + !std::isfinite(*attrs->GetFloat(2))) return ge::GRAPH_FAILED; + // Zero keeps clipping disabled, matching the layered W4A8 operator. + const float *swigluLimitPtr = attrs->GetFloat(3); + const float swigluLimit = swigluLimitPtr != nullptr ? *swigluLimitPtr : 0.0f; + if (!std::isfinite(swigluLimit) || swigluLimit < 0.0f) return ge::GRAPH_FAILED; + Cam::Pull::Shape shape{uint64_t(c[0] / c[9]), uint64_t(c[1]), uint64_t(c[3]), + uint64_t(c[4]), uint64_t(c[5]), uint64_t(c[6]), uint64_t(c[9])}; + Cam::Pull::Layout layout{}; + if (!Cam::Pull::MakeLayout(shape, layout) || layout.bytes > GetMaxWindowSize()) + return ge::GRAPH_FAILED; + auto platform = platform_ascendc::PlatformAscendC(ctx->GetPlatformInfo()); + const auto aic = platform.GetCoreNumAic(), aiv = platform.GetCoreNumAiv(); + uint64_t ub = 0; + platform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ub); + if (aic == 0 || aiv != 2 * aic || aiv > Cam::NATIVE_MAX_AIVS || + Cam::Pull::TransportUbBytes(shape) > ub || + Cam::Pull::Align(8 * Cam::Pull::MetadataWords(shape.tp, shape.experts)) + + 10ULL * c[1] + 1280 > ub || + 184ULL * 1024 + 8ULL * Cam::NATIVE_METADATA_ROWS > ub) return ge::GRAPH_FAILED; + // MC2 full-row ACT uses 54 bytes per intermediate column plus per-core scales. + if (c[11] == 1 && c[2] <= 3072 && 54ULL * c[2] + 256 + 4ULL * ((c[10] + aiv - 1) / aiv) > ub) + return ge::GRAPH_FAILED; + const Cam::NativeMegaWorkspace w(c[10], c[1], c[2], c[6], c[9]); + if (ctx->GetInputShape(3)->GetStorageShape().GetDim(0) != int64_t(w.bytes)) + return ge::GRAPH_FAILED; + auto data = ctx->GetTilingData(); + std::memset(data, 0, sizeof(*data)); + std::copy(c, c + 14, data->config); + data->aicNum = aic; + data->aivNum = aiv; + // All AIVs take part in MC2 dispatch/MSD, activation and paired output. + data->commAivNum = 0; + data->layerMapSize = ctx->GetInputShape(2)->GetStorageShape().GetDim(0); + if (data->layerMapSize == 0 || data->layerMapSize > 65536) return ge::GRAPH_FAILED; + data->routedScale = *attrs->GetFloat(2); + data->swigluLimit = swigluLimit; + AscendC::Mc2CcTilingConfig comm(attrs->GetAttrPointer(1), 8U, + "AlltoAll=level0:fullmesh;level1:pairwise"); + comm.GetTiling(data->mc2InitTiling); + comm.GetTiling(data->mc2CcTiling1); + ctx->SetBlockDim(aic); + ctx->SetTilingKey(c[11]); + ctx->SetScheduleMode(1); + ctx->GetWorkspaceSizes(1)[0] = 16 * 1024 * 1024; + return ge::GRAPH_SUCCESS; +} +} + +namespace ops { +class AfdMegaMoePullNative : public OpDef { +public: + explicit AfdMegaMoePullNative(const char *name) : OpDef(name) { + Input("anchor").ParamType(REQUIRED).DataType({ge::DT_BF16, ge::DT_BF16}).Format({ge::FORMAT_ND, ge::FORMAT_ND}); + Input("commArgs").ParamType(REQUIRED).DataType({ge::DT_FLOAT16, ge::DT_FLOAT16}).Format({ge::FORMAT_ND, ge::FORMAT_ND}); + Input("layerMap").ParamType(REQUIRED).DataType({ge::DT_INT64, ge::DT_INT64}).Format({ge::FORMAT_ND, ge::FORMAT_ND}); + Input("scratch").ParamType(REQUIRED).DataType({ge::DT_UINT8, ge::DT_UINT8}).Format({ge::FORMAT_ND, ge::FORMAT_ND}); + for (const char *w : {"w13", "w2"}) + Input(w).ParamType(DYNAMIC).DataType({ge::DT_INT32, ge::DT_INT8}).Format({ge::FORMAT_FRACTAL_NZ, ge::FORMAT_FRACTAL_NZ}); + Input("s13").ParamType(DYNAMIC).DataType({ge::DT_INT64, ge::DT_FLOAT}).Format({ge::FORMAT_ND, ge::FORMAT_ND}); + Input("s2").ParamType(DYNAMIC).DataType({ge::DT_INT64, ge::DT_BF16}).Format({ge::FORMAT_ND, ge::FORMAT_ND}); + for (const char *b : {"b13", "b2"}) + Input(b).ParamType(DYNAMIC).DataType({ge::DT_FLOAT, ge::DT_FLOAT}).Format({ge::FORMAT_ND, ge::FORMAT_ND}); + Output("completion").ParamType(REQUIRED).DataType({ge::DT_INT64, ge::DT_INT64}).Format({ge::FORMAT_ND, ge::FORMAT_ND}); + Attr("config").ListInt(); + Attr("hccl_group_name").String(); + Attr("routed_scale").Float(); + Attr("swiglu_limit").Float(); + SetInferShape([](gert::InferShapeContext *ctx) { + auto shape = ctx->GetOutputShape(0); + shape->SetDimNum(1); shape->SetDim(0, 1); + return ge::GRAPH_SUCCESS; + }); + SetInferDataType([](gert::InferDataTypeContext *ctx) { + ctx->SetOutputDataType(0, ge::DT_INT64); return ge::GRAPH_SUCCESS; + }); + AICore().AddConfig("ascend910_93"); + AICore().SetTiling(optiling::NativeMegaTiling); + MC2().HcclGroup({"hccl_group_name"}); + } +}; +OP_ADD(AfdMegaMoePullNative); +} diff --git a/csrc/npu/ascend_kernels/afd_mega_moe_pull_native/op_kernel/afd_mega_moe_patch_activation.py b/csrc/npu/ascend_kernels/afd_mega_moe_pull_native/op_kernel/afd_mega_moe_patch_activation.py new file mode 100644 index 000000000..b6eec2a72 --- /dev/null +++ b/csrc/npu/ascend_kernels/afd_mega_moe_pull_native/op_kernel/afd_mega_moe_patch_activation.py @@ -0,0 +1,49 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Generate a Layered-compatible MC2 activation from the installed CANN header.""" + +import re +import sys +from pathlib import Path + + +def patched_activation(source: str) -> str: + source = source.replace( + '"../template_linear_algebra_v2/', '"template_linear_algebra_v2/' + ) + for header in ("get_tensor_addr.hpp", "gated_activation.hpp"): + source = source.replace(f'"{header}"', f'"utils/{header}"') + + source, multipliers = re.subn( + r"127\.0f / rowAbsMax", + "1.0f / (rowAbsMax / 127.0f)", + source, + ) + source, legacy_multipliers = re.subn( + r"127\.f / GMubDequantScale", + "(GMubDequantScale > 0.0f ? 1.0f / (GMubDequantScale / 127.f) : 0.0f)", + source, + ) + if multipliers + legacy_multipliers != 1: + raise ValueError("unsupported CANN activation scale calculation") + + pattern = ( + r"AscendC::Cast\(ubQuantS32, (\w+), AscendC::RoundMode::CAST_RINT, (\w+)\);" + r"\s*AscendC::PipeBarrier\(\);" + r"\s*AscendC::SetDeqScale\(static_cast\(1(?:\.0)?\)\);" + r"\s*AscendC::Cast\(ubQuantF16, ubQuantS32, " + r"AscendC::RoundMode::CAST_RINT, \2\);" + ) + source, casts = re.subn( + pattern, + r"AscendC::Cast(ubQuantF16, \1, AscendC::RoundMode::CAST_RINT, \2);", + source, + ) + if casts != 1: + raise ValueError("unsupported CANN activation quantization sequence") + source = re.sub(r"^\s*auto ubQuantS32 = .*?;\n", "\n", source, flags=re.MULTILINE) + return source + + +if __name__ == "__main__": + Path(sys.argv[2]).write_text(patched_activation(Path(sys.argv[1]).read_text())) diff --git a/csrc/npu/ascend_kernels/afd_mega_moe_pull_native/op_kernel/afd_mega_moe_plan_w4a8.h b/csrc/npu/ascend_kernels/afd_mega_moe_pull_native/op_kernel/afd_mega_moe_plan_w4a8.h new file mode 100644 index 000000000..f6efc8554 --- /dev/null +++ b/csrc/npu/ascend_kernels/afd_mega_moe_pull_native/op_kernel/afd_mega_moe_plan_w4a8.h @@ -0,0 +1,55 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +/** + * W4A8 execution plan for the native MegaMoE pull operator. + * + * Everything the pipeline needs that depends on the weight/activation dtype + * lives here, so afd_mega_moe_pull_pipeline.h stays dtype-agnostic and a new + * quantization only adds a sibling plan header. + * + * Include order: this header assumes kernel_operator.h and the installed MC2 + * Catlass headers are already included, plus the shared AfdNativeMega aliases + * (Arch, RowMajor, Nz, VectorLayout) in afd_mega_moe_pull_native.cpp. + */ +#ifndef AFD_MEGA_MOE_PLAN_W4A8_H +#define AFD_MEGA_MOE_PLAN_W4A8_H + +namespace AfdNativeMega { + +struct W4A8Plan { + static constexpr bool kPackedInt4 = true; + static constexpr uint32_t kRows = 2; + using OutputElement = half; + using OutputScale = float; + // Packed INT4 weights. An INT8 activation row of h elements is MSD-split + // into 2*h INT4 halves, so GMM1 rows are twice the logical row count. + using Element = AscendC::int4b_t; + // Per-channel weight scales arrive as INT64-encoded floats. + using WeightScale = uint64_t; + + using A = Catlass::Gemm::GemmType; + using B = Catlass::Gemm::GemmType; + using C = Catlass::Gemm::GemmType; + + using Copy = Catlass::Gemm::Tile::QuantTileCopy; + using Policy = Catlass::Gemm::MmadAtlasA2W4A4MatmulPerChannelDequant< + 1, 2, 2, 2, 1, false, true>; + using Mmad = Catlass::Gemm::Block::BlockMmad, + Catlass::GemmShape<128, 256, 256>, A, B, C, void, Copy>; + using Scheduler = Catlass::Gemm::Block::GemmIdentityBlockSwizzle<9, 1>; + + // Installed MC2 W4A8 epilogue, patched to match Layered's rounding. + using ActOutput = Catlass::Gemm::GemmType; + using ActScale = Catlass::Gemm::GemmType; + using ActCopy = Catlass::Epilogue::Tile::TileCopy, ActScale, ActOutput>; + using ActMuls = Catlass::Epilogue::Tile::TileElemWiseMuls, 0>; + using NativeActivation = Catlass::Epilogue::Block::BlockEpilogue< + Catlass::Epilogue::EpilogueAtlasA2W4A8PostPerTokenDequantSwigluQuant<2>, + C, ActScale, ActOutput, ActMuls, ActCopy>; +}; + +} // namespace AfdNativeMega +#endif // AFD_MEGA_MOE_PLAN_W4A8_H diff --git a/csrc/npu/ascend_kernels/afd_mega_moe_pull_native/op_kernel/afd_mega_moe_plan_w8a8.h b/csrc/npu/ascend_kernels/afd_mega_moe_pull_native/op_kernel/afd_mega_moe_plan_w8a8.h new file mode 100644 index 000000000..4a9370498 --- /dev/null +++ b/csrc/npu/ascend_kernels/afd_mega_moe_pull_native/op_kernel/afd_mega_moe_plan_w8a8.h @@ -0,0 +1,25 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +#ifndef AFD_MEGA_MOE_PLAN_W8A8_H +#define AFD_MEGA_MOE_PLAN_W8A8_H + +namespace AfdNativeMega { +struct W8A8Plan { + static constexpr bool kPackedInt4 = false; + static constexpr uint32_t kRows = 1; + using Element = int8_t; + using OutputElement = int32_t; + using WeightScale = float; + using OutputScale = bfloat16_t; + using A = Catlass::Gemm::GemmType; + using B = Catlass::Gemm::GemmType; + using C = Catlass::Gemm::GemmType; + // Keep the integer accumulator until Vector applies Layered's FP32 scales. + using Policy = Catlass::Gemm::MmadAtlasA2PreloadAsyncWithCallback< + 1, 2, 2, 2, 1, false, true>; + using Mmad = Catlass::Gemm::Block::BlockMmad, + Catlass::GemmShape<128, 256, 128>, A, B, C>; + using Scheduler = Catlass::Gemm::Block::GemmIdentityBlockSwizzle<9, 1>; +}; +} // namespace AfdNativeMega +#endif diff --git a/csrc/npu/ascend_kernels/afd_mega_moe_pull_native/op_kernel/afd_mega_moe_pull_native.cpp b/csrc/npu/ascend_kernels/afd_mega_moe_pull_native/op_kernel/afd_mega_moe_pull_native.cpp new file mode 100644 index 000000000..c27d93502 --- /dev/null +++ b/csrc/npu/ascend_kernels/afd_mega_moe_pull_native/op_kernel/afd_mega_moe_pull_native.cpp @@ -0,0 +1,257 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +#include "kernel_operator.h" +#include "lib/matmul_intf.h" // CANN generated MIX wrapper calls clearWorkspace. +#include "afd_mega_moe_pull_native_tiling.h" +#include "afd_pull_protocol.h" +#include "layered/layered_tensor_addr.h" + +// Use the installed MC2 implementation; do not vendor its template dependency tree. +#if defined(AFD_CANN_CATLASS_92) +#include "catlass/arch/resource.hpp" +#include "catlass/gemm/block/block_mmad.hpp" +#include "catlass/gemm/block/block_swizzle.hpp" +#include "catlass/epilogue/tile/tile_copy.hpp" +#include "catlass/epilogue/tile/tile_elemwise_muls.hpp" +#include "catlass/epilogue/block/block_epilogue.hpp" +#else +#include "template_linear_algebra_v2/arch/resource.hpp" +#include "template_linear_algebra_v2/gemm/block/block_mmad.hpp" +#include "template_linear_algebra_v2/gemm/block/block_swizzle.hpp" +#include "template_linear_algebra_v2/epilogue/tile/tile_copy.hpp" +#include "template_linear_algebra_v2/epilogue/tile/tile_elemwise_muls.hpp" +#include "template_linear_algebra_v2/epilogue/block/block_epilogue.hpp" +#endif +#include "utils/block_mmad_w4a4.hpp" +// The installed MC2 epilogue uses unqualified AscendC names inside this namespace. +namespace Catlass::Epilogue::Block { using namespace AscendC; } +#include "afd_mega_moe_activation.h" +#include "afd_swiglu_limit.h" + +namespace AfdNativeMega { +using namespace AscendC; +// Shared Catlass aliases. The per-quantization plan header below consumes these. +using Arch = Catlass::Arch::AtlasA2; +using Resource = Catlass::Arch::Resource; +using RowMajor = Catlass::layout::RowMajor; +using Nz = Catlass::layout::zN; +using VectorLayout = Catlass::layout::VectorLayout; +using Coord = Catlass::GemmCoord; +} // namespace AfdNativeMega + +#include "afd_mega_moe_plan_w4a8.h" +#include "afd_mega_moe_plan_w8a8.h" + +namespace AfdNativeMega { +constexpr uint32_t CUBE_M = 128, CUBE_N = 256; +constexpr uint32_t FLAG_SETS = 15, GMM1_READY = 10, ACT_READY = 11; +constexpr uint32_t MAX_EXPERTS = 128; +constexpr uint32_t ROUTE_UB = 184 * 1024; + +template __aicore__ inline void Fence() { + constexpr auto id = EVENT_ID7; + SetFlag(id); + WaitFlag(id); +} +__aicore__ inline int64_t Load(GM_ADDR address, uint32_t index = 0) { + GlobalTensor gm; + gm.SetGlobalBuffer(reinterpret_cast<__gm__ int64_t *>(address) + index); + __asm__ __volatile__("" ::: "memory"); + DataCacheCleanAndInvalid(gm); + __asm__ __volatile__("" ::: "memory"); + return *(reinterpret_cast<__gm__ volatile int64_t *>(address) + index); +} +__aicore__ inline Cam::Pull::Shape PullShape(const Cam::AfdMegaMoePullNativeTilingData &t) { + const auto c = t.config; + return {uint64_t(c[0] / c[9]), uint64_t(c[1]), uint64_t(c[3]), uint64_t(c[4]), + uint64_t(c[5]), uint64_t(c[6]), uint64_t(c[9])}; +} +// MC2 W4A8 emits activation groups at the exponentially shrinking expert tail. +__aicore__ inline bool SyncExpert(uint32_t expert, uint32_t experts) { + const uint32_t tail = experts - expert; + return (tail & (tail - 1)) == 0; +} +struct ReceiveRound { + uint32_t begin[MAX_EXPERTS + 1]; + uint32_t active[MAX_EXPERTS]; + uint32_t activeCount = 0, expertCount; + uint32_t source, offset, total; + __aicore__ ReceiveRound(GM_ADDR scratch, const Cam::NativeMegaWorkspace &w, uint32_t experts) { + expertCount = experts; + source = Load(scratch + w.batch, 1); + offset = Load(scratch + w.batch, 3); + total = Load(scratch + w.batch, 4) - offset; + begin[0] = 0; + auto counts = reinterpret_cast<__gm__ volatile int64_t *>(scratch + w.counts); + for (uint32_t e = 0; e < experts; ++e) { + // Invalidate once per cache line after the preceding receive round reused it. + if (e % 8 == 0) Load(scratch + w.counts, e); + begin[e + 1] = begin[e] + counts[e]; + activeCount += begin[e + 1] != begin[e]; + } + // Keep dense rounds on the original expert order without writing an index list. + if (activeCount != experts) { + uint32_t index = 0; + for (uint32_t e = 0; e < experts; ++e) { + if (Rows(e)) active[index++] = e; + } + } + } + __aicore__ uint32_t Rows(uint32_t e) const { return begin[e + 1] - begin[e]; } + __aicore__ uint32_t Expert(uint32_t index) const { + return activeCount == expertCount ? index : active[index]; + } +}; + +// Same Catlass block implementation, preload depth, K shuffle and expert swizzle as MC2. +// The plan supplies the operand element type, GEMM tiling and scheduler; the +// receive/round/flag structure below is identical for every quantization. +template +__aicore__ inline void Compute(GM_ADDR scratch, const Cam::NativeMegaWorkspace &w, + const Cam::AfdMegaMoePullNativeTilingData &t, const ReceiveRound &round, GM_ADDR weights, + GM_ADDR scales, int64_t layer, bool second, Resource &resource) { + using Element = typename Plan::Element; + const uint32_t k = second ? t.config[2] : t.config[1]; + const uint32_t n = second ? t.config[1] : 2 * t.config[2]; + GlobalTensor a, b; + GlobalTensor scale; + GlobalTensor out; + a.SetGlobalBuffer(reinterpret_cast<__gm__ Element *>(scratch + (second ? w.a2 : w.a1))); + b.SetGlobalBuffer(GetLayerTensorAddr(layer, t.config[13], weights)); + scale.SetGlobalBuffer( + GetLayerTensorAddr(layer, t.config[13], scales)); + out.SetGlobalBuffer(reinterpret_cast<__gm__ typename Plan::OutputElement *>(scratch + (second ? w.c2 : w.c1))); + const auto layoutB = Nz::MakeLayout(k, n); + typename Plan::Mmad mm(resource); + typename Plan::Scheduler scheduler; + uint32_t startCore = 0, syncGroup = 0; + for (uint32_t index = 0; index < round.activeCount; ++index) { + const uint32_t e = round.Expert(index); + const uint32_t rows = Plan::kRows * round.Rows(e); + if (second) { + if (index == 0 || SyncExpert(index - 1, round.activeCount)) CrossCoreWaitFlag<2>(ACT_READY); + } else { + CrossCoreWaitFlag<2>(index / FLAG_SETS); + } + scheduler.Update(Coord{rows, n, k}, Catlass::MatrixCoord{CUBE_M, CUBE_N}); + const auto jobs = scheduler.GetCoreLoops(); + const auto first = (get_block_idx() + t.aicNum - startCore) % t.aicNum; + const RowMajor layoutA{rows, k}, layoutC{rows, n}; + auto expertB = b[uint64_t(e) * k * n]; + expertB.SetL2CacheHint(rows <= CUBE_M ? CacheMode::CACHE_MODE_DISABLE : CacheMode::CACHE_MODE_NORMAL); + for (uint32_t job = first; job < jobs; job += t.aicNum) { + const auto coord = scheduler.GetBlockCoord(job); + const auto actual = scheduler.GetActualBlockShape(coord); + const uint32_t m = coord.m() * CUBE_M, col = coord.n() * CUBE_N; + const int32_t done = second && job + t.aicNum >= jobs ? int32_t(index) : -1; + const uint64_t row = uint64_t(Plan::kRows) * round.begin[e] + m; + if constexpr (Plan::kPackedInt4) { + mm(scale[e * n + col], VectorLayout{n}, a[row * k], layoutA, + expertB[layoutB.GetOffset(Catlass::MatrixCoord{0, col})], layoutB, + out[row * n + col], layoutC, actual, ::Callback{}, ::Callback{}, done, 0); + } else { + mm(a[row * k], layoutA, + expertB[layoutB.GetOffset(Catlass::MatrixCoord{0, col})], layoutB, + out[row * n + col], layoutC, actual, ::Callback{}, ::Callback{}); + } + } + if (!second && SyncExpert(index, round.activeCount)) { + mm.SynchronizeBlock(); + if constexpr (Plan::kPackedInt4) mm.Finalize(syncGroup++, GMM1_READY); + else CrossCoreSetFlag<2, PIPE_FIX>(GMM1_READY); + } + if constexpr (!Plan::kPackedInt4) { + if (second) { + mm.SynchronizeBlock(); + CrossCoreSetFlag<2, PIPE_FIX>(index / FLAG_SETS); + } + } + startCore = (startCore + jobs) % t.aicNum; + } + mm.SynchronizeBlock(); + if constexpr (Plan::kPackedInt4) { + if (second) mm.Finalize(int32_t(round.activeCount) - 1, 0); + } +} +} // namespace AfdNativeMega + +#include "afd_mega_moe_pull_pipeline.h" + +namespace AfdNativeMega { +template +__aicore__ inline void Process(GM_ADDR state, GM_ADDR layerMap, GM_ADDR scratch, + GM_ADDR w13, GM_ADDR w2, GM_ADDR s13, GM_ADDR s2, GM_ADDR b13, GM_ADDR b2, + GM_ADDR completion, const Cam::AfdMegaMoePullNativeTilingData &t) { + const auto c = t.config; + const Cam::NativeMegaWorkspace w(c[10], c[1], c[2], c[6], c[9]); + Resource resource; + TPipe pipe; + // TPipe seeds M_MTE1 0..2; MC2 must own these events exclusively on AIC. + if ASCEND_IS_AIC { pipe.Destroy(); } + const uint32_t id = GetBlockIdx(); + // A fixed workspace drains one selected A transaction, including all receive rounds. + do { + if ASCEND_IS_AIV { + if (id == 0) { + Cam::Pull::Transport rx(state, PullShape(t), c[7], t.aivNum, pipe, 0); + rx.PrepareRound(scratch + w.batch, scratch + w.counts, c[10]); + const auto modelLayer = rx.Read(scratch + w.batch + 16); + int64_t slot = -1; + if (rx.ReadState(Cam::Pull::SOURCE) >= 0 && !rx.ReadState(Cam::Pull::FAILURE)) { + if (modelLayer >= 0 && modelLayer < t.layerMapSize) + slot = rx.Read(layerMap + modelLayer * 8); + if (slot < 0 || slot >= c[13]) { + const auto source = rx.ReadState(Cam::Pull::SOURCE); + rx.Fail(source >= 0 && source + c[9] <= c[5] ? source : 0, -4); + slot = -1; + } + } + rx.Write(scratch + w.layer, slot); + } + pipe.Reset(); + } + SyncAll(); + const auto layer = Load(state, Cam::Pull::FAILURE) ? -1 : Load(scratch + w.layer); + if (layer >= 0 && layer < c[13]) { + const ReceiveRound round(scratch, w, c[6]); + if (round.total) { + if ASCEND_IS_AIC { + Compute(scratch, w, t, round, w13, s13, layer, false, resource); + Compute(scratch, w, t, round, w2, s2, layer, true, resource); + } + if ASCEND_IS_AIV { + PullAndPrepare(scratch, w, t, round, resource); + Activate(scratch, w, t, round, Plan::kPackedInt4 ? b13 : s13, layer, resource); + Combine(scratch, w, t, round, Plan::kPackedInt4 ? b2 : s2, layer, resource); + } + } + } + SyncAll(); + if ASCEND_IS_AIV { pipe.Reset(); } + if ASCEND_IS_AIV { + if (id == 0) { + Cam::Pull::Transport tx(state, PullShape(t), c[7], t.aivNum, pipe, 0); + if (layer >= 0 && layer < c[13]) tx.CompleteRound(scratch + w.batch, c[10]); + tx.Write(completion, tx.ReadState(Cam::Pull::FAILURE)); + } + } + // All GMM/epilogue events have drained before scratch and flag slots are reused. + SyncAll(); + if ASCEND_IS_AIV { pipe.Reset(); } + } while (!Load(state, Cam::Pull::FAILURE) && Load(state, Cam::Pull::SOURCE) >= 0); +} +} // namespace AfdNativeMega + +extern "C" __global__ __aicore__ void afd_mega_moe_pull_native( + GM_ADDR anchor, GM_ADDR state, GM_ADDR layerMap, GM_ADDR scratch, GM_ADDR w13, GM_ADDR w2, + GM_ADDR s13, GM_ADDR s2, GM_ADDR b13, GM_ADDR b2, GM_ADDR completion, + GM_ADDR workspace, GM_ADDR tiling) { + using namespace AfdNativeMega; + KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_MIX_AIC_1_2); + REGISTER_TILING_DEFAULT(Cam::AfdMegaMoePullNativeTilingData); + REGISTER_TILING_FOR_TILINGKEY("TILING_KEY_VAR == 1", Cam::AfdMegaMoePullNativeTilingData); + REGISTER_TILING_FOR_TILINGKEY("TILING_KEY_VAR == 2", Cam::AfdMegaMoePullNativeTilingData); + GET_TILING_DATA_WITH_STRUCT(Cam::AfdMegaMoePullNativeTilingData, t, tiling); + if (TILING_KEY_IS(1)) Process(state, layerMap, scratch, w13, w2, s13, s2, b13, b2, completion, t); + if (TILING_KEY_IS(2)) Process(state, layerMap, scratch, w13, w2, s13, s2, b13, b2, completion, t); +} diff --git a/csrc/npu/ascend_kernels/afd_mega_moe_pull_native/op_kernel/afd_mega_moe_pull_native_tiling.h b/csrc/npu/ascend_kernels/afd_mega_moe_pull_native/op_kernel/afd_mega_moe_pull_native_tiling.h new file mode 100644 index 000000000..a737f8e14 --- /dev/null +++ b/csrc/npu/ascend_kernels/afd_mega_moe_pull_native/op_kernel/afd_mega_moe_pull_native_tiling.h @@ -0,0 +1,21 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +#ifndef AFD_MEGA_MOE_PULL_NATIVE_TILING_H +#define AFD_MEGA_MOE_PULL_NATIVE_TILING_H +#include +#include "kernel_tiling/kernel_tiling.h" +#include "afd_native_mega_workspace.h" +namespace Cam { +struct AfdMegaMoePullNativeTilingData { + Mc2InitTiling mc2InitTiling; + Mc2CcTiling mc2CcTiling1; + Mc2CcTiling mc2CcTiling2; + int64_t config[14]; + uint32_t aicNum, aivNum, commAivNum, layerMapSize; + alignas(8) uint8_t cube[1024]; + float routedScale; + // Model swiglu_limit; zero disables clipping, matching the layered W4A8 op. + float swigluLimit; +}; +} +#endif diff --git a/csrc/npu/ascend_kernels/afd_mega_moe_pull_native/op_kernel/afd_mega_moe_pull_pipeline.h b/csrc/npu/ascend_kernels/afd_mega_moe_pull_native/op_kernel/afd_mega_moe_pull_pipeline.h new file mode 100644 index 000000000..ac9842065 --- /dev/null +++ b/csrc/npu/ascend_kernels/afd_mega_moe_pull_native/op_kernel/afd_mega_moe_pull_pipeline.h @@ -0,0 +1,699 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + * Modifications Copyright contributors to the AFD plugin project. + */ +#ifndef AFD_MEGA_MOE_PULL_PIPELINE_H +#define AFD_MEGA_MOE_PULL_PIPELINE_H + +// Port of MC2 arch22/a2_a3_copy_helpers.hpp and W4A8 block epilogues. +// AFD changes: frozen source routes, layer descriptors, direct A return windows, +// and the existing BF16 rounding before routed_scale. No all-to-all prologue. +namespace AfdNativeMega { +__aicore__ inline uint32_t Min(uint32_t a, uint32_t b) { return a < b ? a : b; } + +// MC2's signed high nibble and low nibble minus eight; bias compensates the latter. +__aicore__ inline void Msd(LocalTensor q, LocalTensor work, + LocalTensor packed, uint32_t n) { + auto high = work, lowBits = work[n], low = work[2 * n]; + auto mask = work[3 * n].ReinterpretCast(); + Duplicate(mask, static_cast(0x0f0f), 128); + Cast(high, q, RoundMode::CAST_NONE, n); + PipeBarrier(); + Muls(high, high, static_cast(0.0625f), n); + PipeBarrier(); + Cast(packed, high, RoundMode::CAST_FLOOR, n); + And(lowBits.ReinterpretCast(), q.ReinterpretCast(), mask, + 128, n / 256, {1, 1, 1, 8, 8, 0}); + PipeBarrier(); + Cast(low, lowBits.ReinterpretCast(), RoundMode::CAST_NONE, n); + PipeBarrier(); + Adds(low, low, static_cast(-8), n); + PipeBarrier(); + Cast(packed[n], low, RoundMode::CAST_NONE, n); +} + +__aicore__ inline void FetchRow(LocalTensor input, GlobalTensor remote, + uint32_t hidden, uint32_t stage) { + WaitFlag(stage); + WaitFlag(stage); + DataCopyPad(input, remote, {1U, hidden + uint32_t(sizeof(float)), 0U, 0U, 0U}, {false, 0U, 0U, 0U}); + SetFlag(stage); +} + +template +__aicore__ inline void PullAndPrepare(GM_ADDR scratch, const Cam::NativeMegaWorkspace &w, + const Cam::AfdMegaMoePullNativeTilingData &t, const ReceiveRound &round, Resource &resource) { + using Element = typename Plan::Element; + const uint32_t h = t.config[1], experts = t.config[6], tp = t.config[9]; + const auto shape = PullShape(t); + Cam::Pull::Layout layout{}; + Cam::Pull::MakeLayout(shape, layout); + const uint32_t metadataBytes = 8 * Cam::Pull::MetadataWords(tp, experts); + auto meta = resource.ubBuf.GetBufferByByte(0); + GlobalTensor batch; + batch.SetGlobalBuffer(reinterpret_cast<__gm__ int64_t *>(scratch + w.batch)); + DataCopyPad(meta, batch, {1U, metadataBytes, 0U, 0U, 0U}, {false, 0U, 0U, 0U}); + Fence(); + uint32_t ub = Cam::Pull::Align(metadataBytes); + LocalTensor input[2]; + LocalTensor packed[2]; + for (uint32_t b = 0; b < 2; ++b) { + input[b] = resource.ubBuf.GetBufferByByte(ub); + ub += h + Cam::Pull::ROW_PADDING; + packed[b] = resource.ubBuf.GetBufferByByte(ub); + ub += h; + SetFlag(b); + SetFlag(b); + SetFlag(b); + SetFlag(b); + } + auto work = resource.ubBuf.GetBufferByByte(ub); + GlobalTensor out; + GlobalTensor scales; + GlobalTensor routes; + out.SetGlobalBuffer(reinterpret_cast<__gm__ int8_t *>(scratch + w.a1)); + scales.SetGlobalBuffer(reinterpret_cast<__gm__ float *>(scratch + w.scale1)); + // The old INT8 staging allocation now holds the per-row return address cache. + routes.SetGlobalBuffer(reinterpret_cast<__gm__ uint64_t *>(scratch + w.x)); + uint64_t block = 0; + uint32_t active = 0; + for (uint32_t e = 0; e < experts; ++e) { + const uint32_t rows = round.Rows(e), id = GetBlockIdx(); + const uint32_t perCore = rows / t.aivNum, tail = rows % t.aivNum; + const uint32_t first = round.begin[e] + id * perCore + Min(id, tail); + const uint32_t count = perCore + (id < tail); + for (uint32_t s = 0; s < tp; ++s) { + const uint32_t index = 5 + tp + s * experts + e; + const uint64_t sourceRows = meta.GetValue(index); + const uint64_t prefix = meta.GetValue(5 + s); + // Even an empty slice may precede this round's offset in the source window. + meta.SetValue(5 + s, prefix + sourceRows); + block += sourceRows; + if (!rows) continue; + const auto slice = Cam::Pull::Intersect(block - sourceRows, sourceRows, round.offset + first, count); + if (slice.count) { + GlobalTensor remote; + const auto window = Cam::Pull::Window(round.source + s, t.config[7]); + remote.SetGlobalBuffer(reinterpret_cast<__gm__ int8_t *>(window + layout.payload)); + FetchRow(input[0], remote[(prefix + slice.skip) * layout.rowStride], h, 0); + for (uint32_t r = 0; r < slice.count; ++r) { + const uint32_t stage = r % 2; + if (r + 1 < slice.count) + FetchRow(input[1 - stage], remote[(prefix + slice.skip + r + 1) * layout.rowStride], h, 1 - stage); + WaitFlag(stage); + WaitFlag(stage); + if constexpr (Plan::kPackedInt4) Msd(input[stage], work, packed[stage], h); + SetFlag(stage); + SetFlag(stage); + WaitFlag(stage); + const uint32_t row = first + slice.local + r; + if constexpr (Plan::kPackedInt4) + DataCopy(out[uint64_t(row) * h], packed[stage].template ReinterpretCast(), h); + else DataCopy(out[uint64_t(row) * h], input[stage], h); + SetFlag(stage); + DataCopyPad(scales[row], input[stage][h].ReinterpretCast(), {1U, 4U, 0U, 0U, 0U}); + auto address = input[stage][h + 32].ReinterpretCast(); + // The next source may reuse this scalar address before its DMA drains. + WaitFlag(stage); + address.SetValue(0, reinterpret_cast(window + layout.returned) + + (prefix + slice.skip + r) * h * sizeof(bfloat16_t)); + Fence(); + DataCopyPad(routes[row], address, {1U, 8U, 0U, 0U, 0U}); + SetFlag(stage); + SetFlag(stage); + } + } + } + if (!rows) continue; + SyncAll(); + CrossCoreSetFlag<2, PIPE_MTE3>(active++ / FLAG_SETS); + } + for (uint32_t b = 0; b < 2; ++b) { + WaitFlag(b); + WaitFlag(b); + WaitFlag(b); + WaitFlag(b); + } +} + +// The current MC2 epilogue's bounded column tiles also support I > UB capacity. +// For a multi-tile row, pass one reduces a global maximum and pass two quantizes. +template +class Activation { +public: + using Element = typename Plan::Element; + using Raw = typename Plan::OutputElement; + static constexpr uint32_t TILE = 3072, SCALE_BATCH = 256; + __aicore__ Activation(Resource &resource, uint32_t intermediate, float limit) + : i_(intermediate), limit_(limit) { + uint32_t ub = 0; + for (uint32_t b = 0; b < 2; ++b) { + raw_[b] = resource.ubBuf.GetBufferByByte(ub); ub += 8 * TILE; + bias_[b] = resource.ubBuf.GetBufferByByte(ub); ub += 8 * TILE; + packed_[b] = resource.ubBuf.GetBufferByByte(ub); ub += TILE; + SetFlag(b); + SetFlag(b); + } + gate_ = resource.ubBuf.GetBufferByByte(ub); ub += 4 * TILE; + up_ = resource.ubBuf.GetBufferByByte(ub); ub += 4 * TILE; + value_ = resource.ubBuf.GetBufferByByte(ub); ub += 4 * TILE; + quant_ = resource.ubBuf.GetBufferByByte(ub); ub += TILE; + // Msd reuses gate/up/value after the activation has been quantized. + msdWork_ = gate_.ReinterpretCast(); + maximum_ = resource.ubBuf.GetBufferByByte(ub); ub += 32; + scales_ = resource.ubBuf.GetBufferByByte(ROUTE_UB); + for (uint32_t b = 0; b < 2; ++b) { + scaleOut_[b] = resource.ubBuf.GetBufferByByte(ub); ub += SCALE_BATCH * 4; + SetFlag(b); + } + } + __aicore__ void LoadScales(GlobalTensor source, uint32_t rows) { + if (rows) DataCopyPad(scales_, source, {1U, 4 * rows, 0U, 0U, 0U}, {false, 0U, 0U, 0U}); + Fence(); + } + __aicore__ void Row(GlobalTensor raw, GlobalTensor bias, + GlobalTensor out, LocalTensor scaleOut, + uint32_t scaleIndex) { + const float scale = scales_.GetValue(scaleIndex); + Duplicate(maximum_, 0.0f, 8); + PipeBarrier(); + Fetch(raw, bias, 0, 0); + for (uint32_t col = 0, stage = 0; col < i_; col += TILE, stage ^= 1) { + const auto length = Min(TILE, i_ - col); + if (col + TILE < i_) Fetch(raw, bias, col + TILE, stage ^ 1); + Evaluate(stage, length, scale); + Abs(gate_, value_, length); + PipeBarrier(); + ReduceMax(up_, gate_, up_, length, false); + PipeBarrier(); + Max(maximum_, maximum_, up_, 1); + PipeBarrier(); + } + Fence(); + const float absMax = maximum_.GetValue(0); + const float quantScale = Plan::kPackedInt4 ? absMax / 127.0f : absMax * (1.0f / 127.0f); + scaleOut.SetValue(scaleIndex, quantScale); + const float multiplier = Plan::kPackedInt4 ? (absMax > 0 ? 1.0f / quantScale : 0.0f) : quantScale; + Fence(); + if (i_ <= TILE) { + Quantize(out, 0, i_, 0, multiplier); + } else { + Fetch(raw, bias, 0, 0); + for (uint32_t col = 0, stage = 0; col < i_; col += TILE, stage ^= 1) { + const auto length = Min(TILE, i_ - col); + if (col + TILE < i_) Fetch(raw, bias, col + TILE, stage ^ 1); + Evaluate(stage, length, scale); + Quantize(out, col, length, stage, multiplier); + } + } + } + __aicore__ LocalTensor BeginScale(uint32_t stage) { + WaitFlag(stage); + return scaleOut_[stage]; + } + __aicore__ void EndScale(GlobalTensor output, uint32_t stage, uint32_t count) { + Fence(); + if (count) DataCopyPad(output, scaleOut_[stage], {1U, 4 * count, 0U, 0U, 0U}); + SetFlag(stage); + } + __aicore__ void Finalize() { + for (uint32_t b = 0; b < 2; ++b) { + WaitFlag(b); + WaitFlag(b); + WaitFlag(b); + } + } +private: + __aicore__ void Fetch(GlobalTensor raw, GlobalTensor bias, uint32_t col, uint32_t b) { + WaitFlag(b); + const uint32_t n = Min(TILE, i_ - col); + for (uint32_t part = 0; part < 2 * Plan::kRows; ++part) + DataCopy(raw_[b][part * TILE], raw[part * i_ + col], n); + DataCopy(bias_[b], bias[col], n); + DataCopy(bias_[b][TILE], bias[i_ + col], n); + SetFlag(b); + } + __aicore__ void Evaluate(uint32_t b, uint32_t n, float scale) { + WaitFlag(b); + Cast(gate_, raw_[b], RoundMode::CAST_NONE, n); + if constexpr (Plan::kPackedInt4) { + Cast(value_, raw_[b][2 * TILE], RoundMode::CAST_NONE, n); + PipeBarrier(); + Muls(gate_, gate_, 16.0f, n); + PipeBarrier(); + Add(gate_, gate_, value_, n); + PipeBarrier(); + Cast(up_, raw_[b][TILE], RoundMode::CAST_NONE, n); + Cast(value_, raw_[b][3 * TILE], RoundMode::CAST_NONE, n); + PipeBarrier(); + Muls(up_, up_, 16.0f, n); + PipeBarrier(); + Add(up_, up_, value_, n); + PipeBarrier(); + Add(gate_, gate_, bias_[b], n); + Add(up_, up_, bias_[b][TILE], n); + } else { + Cast(up_, raw_[b][TILE], RoundMode::CAST_NONE, n); + PipeBarrier(); + Mul(gate_, gate_, bias_[b], n); + Mul(up_, up_, bias_[b][TILE], n); + } + SetFlag(b); + PipeBarrier(); + Muls(gate_, gate_, scale, n); + Muls(up_, up_, scale, n); + PipeBarrier(); + // Match Layered: only W4A8 applies the model limit. + if constexpr (Plan::kPackedInt4) ClampSwiGLUInputs(gate_, up_, limit_, n); + Muls(value_, gate_, -1.0f, n); + PipeBarrier(); + Exp(value_, value_, n); + PipeBarrier(); + Adds(value_, value_, 1.0f, n); + PipeBarrier(); + Div(value_, gate_, value_, n); + PipeBarrier(); + Mul(value_, value_, up_, n); + PipeBarrier(); + } + __aicore__ void Quantize(GlobalTensor out, uint32_t col, uint32_t n, + uint32_t b, float multiplier) { + if constexpr (Plan::kPackedInt4) { + Muls(gate_, value_, multiplier, n); + PipeBarrier(); + auto fp16 = gate_.ReinterpretCast(); + Cast(fp16, gate_, RoundMode::CAST_RINT, n); + PipeBarrier(); + Cast(quant_, fp16, RoundMode::CAST_RINT, n); + PipeBarrier(); + WaitFlag(b); + Msd(quant_, msdWork_, packed_[b], n); + SetFlag(b); + WaitFlag(b); + DataCopy(out[col / 2], packed_[b].template ReinterpretCast(), n / 2); + DataCopy(out[i_ / 2 + col / 2], packed_[b][n].template ReinterpretCast(), n / 2); + } else { + Duplicate(gate_, multiplier, n); + PipeBarrier(); + Div(gate_, value_, gate_, n); + PipeBarrier(); + auto integer = gate_.template ReinterpretCast(); + Cast(integer, gate_, RoundMode::CAST_RINT, n); + PipeBarrier(); + SetDeqScale(half(1.0f)); + auto fp16 = up_.template ReinterpretCast(); + Cast(fp16, integer, RoundMode::CAST_ROUND, n); + PipeBarrier(); + WaitFlag(b); + Cast(packed_[b], fp16, RoundMode::CAST_TRUNC, n); + PipeBarrier(); + SetFlag(b); + WaitFlag(b); + DataCopy(out[col], packed_[b], n); + } + SetFlag(b); + } + uint32_t i_; + float limit_; + LocalTensor raw_[2]; + LocalTensor msdWork_; + LocalTensor packed_[2]; + LocalTensor quant_; + LocalTensor bias_[2], gate_, up_, value_, maximum_, scales_, scaleOut_[2]; +}; + +template +__aicore__ inline void ActivateNative(GM_ADDR scratch, const Cam::NativeMegaWorkspace &w, + const Cam::AfdMegaMoePullNativeTilingData &t, const ReceiveRound &round, GM_ADDR biases, + int64_t layer, Resource &resource) { + const uint32_t i = t.config[2], experts = t.config[6], id = GetBlockIdx(); + // Adapt the selected layer to MC2's one-tensor list and int32 expert counts. + auto metadata = resource.ubBuf.GetBufferByByte(0); + metadata.SetValue(0, 8); + metadata.SetValue(1, reinterpret_cast(GetLayerTensorAddr(layer, t.config[13], biases))); + GM_ADDR descriptor = scratch + w.flags + id * 64; + GlobalTensor desc; + desc.SetGlobalBuffer(reinterpret_cast<__gm__ uint64_t *>(descriptor)); + Fence(); + DataCopy(desc, metadata, 4); + GlobalTensor counts; + counts.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t *>(scratch + w.tileCounts)); + if (id == 0) { + auto localCounts = resource.ubBuf.GetBufferByByte(64); + for (uint32_t e = 0; e < experts; ++e) localCounts.SetValue(e, round.Rows(e)); + Fence(); + DataCopyPad(counts, localCounts, {1U, 4 * experts, 0U, 0U, 0U}); + } + Fence(); + Load(descriptor); + GlobalTensor raw; + GlobalTensor s1, s2; + GlobalTensor out; + raw.SetGlobalBuffer(reinterpret_cast<__gm__ half *>(scratch + w.c1)); + s1.SetGlobalBuffer(reinterpret_cast<__gm__ float *>(scratch + w.scale1)); + s2.SetGlobalBuffer(reinterpret_cast<__gm__ float *>(scratch + w.scale2)); + out.SetGlobalBuffer(reinterpret_cast<__gm__ int8_t *>(scratch + w.a2)); + typename Plan::NativeActivation activation(resource, 2 * i); + uint32_t begin = 0; + for (uint32_t index = 0; index < round.activeCount; ++index) { + if (!SyncExpert(index, round.activeCount)) continue; + const uint32_t e = round.Expert(index); + CrossCoreWaitFlag<2>(GMM1_READY); + SyncAll(); + if (begin == 0) { + for (uint32_t j = 0; j < experts; j += 16) Load(scratch + w.tileCounts + 4 * j); + } + const uint32_t rows = round.begin[e + 1] - begin; + if (rows) activation(raw[4ULL * begin * i], Catlass::MatrixCoord{rows, 2 * i}, + s1[begin], reinterpret_cast<__gm__ float *>(descriptor), out[uint64_t(begin) * i], + counts, begin, s2[begin], experts, 1, 0, 1, resource, t.aivNum, t.swigluLimit, +#if defined(AFD_CANN_CATLASS_92) + static_cast(Catlass::Epilogue::GatedActivationMode::SWIGLU), + Catlass::Epilogue::SwigluOaiActivation::DEFAULT_ALPHA, + Catlass::Epilogue::SituActivation::DEFAULT_BETA, +#endif + 2 * i); + SyncAll(); + CrossCoreSetFlag<2, PIPE_MTE3>(ACT_READY); + begin = round.begin[e + 1]; + } + activation.Finalize(); +} + +template +__aicore__ inline void Activate(GM_ADDR scratch, const Cam::NativeMegaWorkspace &w, + const Cam::AfdMegaMoePullNativeTilingData &t, const ReceiveRound &round, GM_ADDR biases, + int64_t layer, Resource &resource) { + if constexpr (Plan::kPackedInt4) { + if (t.config[2] <= Activation::TILE) { + ActivateNative(scratch, w, t, round, biases, layer, resource); + return; + } + } + const uint32_t i = t.config[2], experts = t.config[6], id = GetBlockIdx(); + GlobalTensor raw; + GlobalTensor bias, s1, s2; + GlobalTensor out; + raw.SetGlobalBuffer(reinterpret_cast<__gm__ typename Plan::OutputElement *>(scratch + w.c1)); + bias.SetGlobalBuffer(GetLayerTensorAddr(layer, t.config[13], biases)); + s1.SetGlobalBuffer(reinterpret_cast<__gm__ float *>(scratch + w.scale1)); + s2.SetGlobalBuffer(reinterpret_cast<__gm__ float *>(scratch + w.scale2)); + out.SetGlobalBuffer(reinterpret_cast<__gm__ int8_t *>(scratch + w.a2)); + Activation activation(resource, i, t.swigluLimit); + uint32_t groupBegin = 0, expert = 0, group = 0; + for (uint32_t index = 0; index < round.activeCount; ++index) { + if (!SyncExpert(index, round.activeCount)) continue; + const uint32_t e = round.Expert(index); + CrossCoreWaitFlag<2>(GMM1_READY); + SyncAll(); + const uint32_t rows = round.begin[e + 1] - groupBegin; + const uint32_t each = rows / t.aivNum, tail = rows % t.aivNum; + const uint32_t first = groupBegin + id * each + Min(id, tail); + const uint32_t count = each + (id < tail); + // Bound UB usage even when one expert consumes the entire receive capacity. + for (uint32_t offset = 0; offset < count; offset += Activation::SCALE_BATCH) { + const uint32_t length = Min(Activation::SCALE_BATCH, count - offset); + const uint32_t stage = group++ % 2; + auto scaleOut = activation.BeginScale(stage); + activation.LoadScales(s1[first + offset], length); + for (uint32_t r = 0; r < length; ++r) { + const uint32_t row = first + offset + r; + while (expert < e && row >= round.begin[expert + 1]) ++expert; + activation.Row(raw[2ULL * Plan::kRows * row * i], bias[2ULL * expert * i], + out[uint64_t(row) * i], scaleOut, r); + } + activation.EndScale(s2[first + offset], stage, length); + } + SyncAll(); + CrossCoreSetFlag<2, PIPE_MTE3>(ACT_READY); + groupBegin = round.begin[e + 1]; + } + activation.Finalize(); +} + +template +class Output { +public: + using Raw = typename Plan::OutputElement; + using Scale = typename Plan::OutputScale; + static constexpr uint32_t ROWS = 16, COLS = CUBE_N, ELEMENTS = ROWS * COLS; + __aicore__ Output(Resource &resource, GM_ADDR routeBase, uint32_t rows, GM_ADDR staged) : totalRows_(rows) { + uint32_t ub = 0; + for (uint32_t b = 0; b < 2; ++b) { + hi_[b] = resource.ubBuf.GetBufferByByte(ub); ub += sizeof(Raw) * ELEMENTS; + if constexpr (Plan::kPackedInt4) { + lo_[b] = resource.ubBuf.GetBufferByByte(ub); ub += 2 * ELEMENTS; + } + bias_[b] = resource.ubBuf.GetBufferByByte(ub); ub += 4 * COLS; + y_[b] = resource.ubBuf.GetBufferByByte(ub); ub += 2 * ELEMENTS; + scale_[b] = resource.ubBuf.GetBufferByByte(ub); ub += 4 * ROWS; + scaleRow_[b] = UINT32_MAX; + SetFlag(b); + SetFlag(b + 2); + SetFlag(b); + } + value_ = resource.ubBuf.GetBufferByByte(ub); ub += 4 * ELEMENTS; + temp_ = resource.ubBuf.GetBufferByByte(ub); + routes_ = resource.ubBuf.GetBufferByByte(ROUTE_UB); + routeGm_.SetGlobalBuffer(reinterpret_cast<__gm__ uint64_t *>(routeBase)); + staged_ = reinterpret_cast(staged); + } + __aicore__ void Tile(GlobalTensor raw, GlobalTensor bias, GlobalTensor scale, + uint32_t row, uint32_t rows, uint32_t col, uint32_t hidden, float routedScale) { + const uint32_t b = stage_; + WaitFlag(b); + DataCopyPad(hi_[b], raw[uint64_t(Plan::kRows) * row * hidden + col], + {uint16_t(rows), COLS * uint32_t(sizeof(Raw)), + (Plan::kRows * hidden - COLS) * uint32_t(sizeof(Raw)), 0U, 0U}, + {false, 0U, 0U, 0U}); + if constexpr (Plan::kPackedInt4) { + DataCopyPad(lo_[b], raw[(2ULL * row + 1) * hidden + col], + {uint16_t(rows), COLS * 2, (2 * hidden - COLS) * 2, 0U, 0U}, {false, 0U, 0U, 0U}); + DataCopy(bias_[b], bias[col], COLS); + } else { + DataCopy(bias_[b].template ReinterpretCast(), bias[col], COLS); + } + SetFlag(b); + WaitFlag(b + 2); + if (scaleRow_[b] != row || scaleRows_[b] != rows) { + DataCopyPad(scale_[b], scale[row], {1U, 4 * rows, 0U, 0U, 0U}, {false, 0U, 0U, 0U}); + scaleRow_[b] = row; + scaleRows_[b] = rows; + } + if (row < routeBegin_ || row + rows > routeBegin_ + routeCount_) { + routeBegin_ = row; + routeCount_ = Min(Cam::NATIVE_METADATA_ROWS, totalRows_ - row); + DataCopyPad(routes_, routeGm_[row], {1U, 8 * routeCount_, 0U, 0U, 0U}, + {false, 0U, 0U, 0U}); + } + Fence(); + WaitFlag(b); + const uint32_t length = rows * COLS; + Cast(value_, hi_[b], RoundMode::CAST_NONE, length); + if constexpr (Plan::kPackedInt4) { + Cast(temp_, lo_[b], RoundMode::CAST_NONE, length); + PipeBarrier(); + Muls(value_, value_, 16.0f, length); + PipeBarrier(); + Add(value_, value_, temp_, length); + PipeBarrier(); + for (uint32_t r = 0; r < rows; ++r) Add(value_[r * COLS], value_[r * COLS], bias_[b], COLS); + } else { + Cast(temp_, bias_[b].template ReinterpretCast(), RoundMode::CAST_NONE, COLS); + PipeBarrier(); + for (uint32_t r = 0; r < rows; ++r) Mul(value_[r * COLS], value_[r * COLS], temp_, COLS); + } + SetFlag(b); + PipeBarrier(); + for (uint32_t r = 0; r < rows; ++r) Muls(value_[r * COLS], value_[r * COLS], scale_[b].GetValue(r), COLS); + SetFlag(b + 2); + PipeBarrier(); + WaitFlag(b); + Cast(y_[b], value_, RoundMode::CAST_RINT, length); + if (routedScale != 1.0f) { + PipeBarrier(); + Cast(value_, y_[b], RoundMode::CAST_NONE, length); + PipeBarrier(); + Muls(value_, value_, routedScale, length); + PipeBarrier(); + Cast(y_[b], value_, RoundMode::CAST_RINT, length); + } + SetFlag(b); + WaitFlag(b); + for (uint32_t r = 0; r < rows;) { + const uint64_t address = staged_ + uint64_t(row + r) * hidden * sizeof(bfloat16_t); + uint32_t contiguous = 1; + while (r + contiguous < rows) ++contiguous; + GlobalTensor remote; + remote.SetGlobalBuffer(reinterpret_cast<__gm__ bfloat16_t *>(address)); + DataCopyPad(remote[col], y_[b][r * COLS], + {uint16_t(contiguous), COLS * 2, 0U, (hidden - COLS) * 2, 0U}); + r += contiguous; + } + SetFlag(b); + stage_ ^= 1; + } + __aicore__ void DrainLocal() { + // Observe every local MTE3 write; keep tokens seeded for next expert/finalize. + for (uint32_t b=0;b<2;++b) { + WaitFlag(b); + SetFlag(b); + } + } + __aicore__ void Finalize() { + for (uint32_t b = 0; b < 2; ++b) { + WaitFlag(b); + WaitFlag(b + 2); + WaitFlag(b); + } + } +private: + uint32_t stage_ = 0, scaleRow_[2], scaleRows_[2]; + uint64_t staged_; + uint32_t totalRows_, routeBegin_ = 0, routeCount_ = 0; + LocalTensor hi_[2]; + LocalTensor lo_[2]; + LocalTensor bias_[2], scale_[2], value_, temp_; + LocalTensor y_[2]; + LocalTensor routes_; + GlobalTensor routeGm_; +}; + + +// Independent UB and events keep TX in flight while subsequent output is produced. +class ExpertRowTx { +public: + static constexpr uint32_t BATCH=1, META_ROWS=256, PUMP_ROWS=1, READY_ROWS=512; + static constexpr uint32_t QUEUE_UB=84*1024, TX_UB=96*1024, EVENT_BASE=4; + __aicore__ ExpertRowTx(Resource &resource, GM_ADDR routeBase, GM_ADDR staged, + uint32_t hidden, uint32_t cores) : h_(hidden),cores_(cores) { + local_.SetGlobalBuffer(reinterpret_cast<__gm__ bfloat16_t *>(staged)); + routeGm_.SetGlobalBuffer(reinterpret_cast<__gm__ uint64_t *>(routeBase)); + queue_=resource.ubBuf.GetBufferByByte(QUEUE_UB); + for (uint32_t b=0;b<2;++b) { + value_[b]=resource.ubBuf.GetBufferByByte(TX_UB+b*BATCH*h_*2); + SetFlag(b+EVENT_BASE); + } + routes_=resource.ubBuf.GetBufferByByte(TX_UB+2*BATCH*h_*2); + } + __aicore__ void Enqueue(uint32_t base,uint32_t rows,uint32_t expert) { + const uint32_t id=GetBlockIdx(),per=rows/cores_,rem=rows%cores_; + const uint32_t count=per+(id=cacheEnd_) { + cacheBegin_=cursor_;cacheEnd_=cursor_+Min(META_ROWS,count_-cursor_); + DataCopyPad(routes_,routeGm_[first_+cursor_],{1U,8*(cacheEnd_-cursor_),0U,0U,0U},{false,0U,0U,0U}); + Fence(); + } + uint32_t n=1; + const uint32_t limit=Min(BATCH,Min(quota,Min(count_-cursor_,cacheEnd_-cursor_))); + const uint64_t address=routes_.GetValue(cursor_-cacheBegin_); + // Never merge across a destination discontinuity or metadata-window boundary. + while (n(b+EVENT_BASE); + WaitFlag(b+EVENT_BASE); + GlobalTensor remote; + remote.SetGlobalBuffer(reinterpret_cast<__gm__ bfloat16_t *>(address)); + DataCopyPad(remote,value_[b],{1U,2*h_*n,0U,0U,0U}); + SetFlag(b+EVENT_BASE); + cursor_+=n;quota-=n; + } + } + __aicore__ void Drain() { + Pump(0x7fffffff); + for (uint32_t b=0;b<2;++b) {WaitSlot(b);SetFlag(b+EVENT_BASE);} + } + __aicore__ void Finalize() { + Pump(0x7fffffff); + for (uint32_t b=0;b<2;++b) WaitSlot(b); + } +private: + __aicore__ void WaitSlot(uint32_t b) {WaitFlag(b+EVENT_BASE);} + uint32_t h_,cores_,stage_=0,head_=0,tail_=0,loaded_=UINT32_MAX; + uint32_t first_=0,count_=0,cursor_=0,cacheBegin_=0,cacheEnd_=0; + LocalTensor value_[2]; + LocalTensor routes_; + LocalTensor queue_; + GlobalTensor local_; + GlobalTensor routeGm_; +}; + +template +__aicore__ inline void Combine(GM_ADDR scratch, const Cam::NativeMegaWorkspace &w, + const Cam::AfdMegaMoePullNativeTilingData &t, const ReceiveRound &round, GM_ADDR biases, + int64_t layer, Resource &resource) { + const uint32_t h = t.config[1]; + GlobalTensor raw; + GlobalTensor bias; + GlobalTensor scales; + raw.SetGlobalBuffer(reinterpret_cast<__gm__ typename Plan::OutputElement *>(scratch + w.c2)); + bias.SetGlobalBuffer(GetLayerTensorAddr(layer, t.config[13], biases)); + scales.SetGlobalBuffer(reinterpret_cast<__gm__ float *>(scratch + w.scale2)); + Output output(resource, scratch + w.x, round.total, scratch + w.output); + ExpertRowTx rowTx(resource, scratch + w.x, scratch + w.output, h, t.aivNum); + typename Plan::Scheduler scheduler; + uint32_t startCore = 0, sync = 0, groupBegin = 0; + for (uint32_t index = 0; index < round.activeCount; ++index) { + const uint32_t e = round.Expert(index); + scheduler.Update(Coord{Plan::kRows * round.Rows(e), h, uint32_t(t.config[2])}, + Catlass::MatrixCoord{CUBE_M, CUBE_N}); + const auto jobs = scheduler.GetCoreLoops(); + const auto first = (get_block_idx() + t.aicNum - startCore) % t.aicNum; + for (uint32_t job = first; job < jobs; job += t.aicNum) { + const auto coord = scheduler.GetBlockCoord(job); + const auto actual = scheduler.GetActualBlockShape(coord); + for (; sync <= index; ++sync) CrossCoreWaitFlag<2>(sync / FLAG_SETS); + // Each paired AIV handles bounded groups of logical output rows. + constexpr uint32_t physicalRows = Plan::kRows * Output::ROWS; + const uint32_t blocks = (actual.m() + physicalRows - 1) / physicalRows; + const uint32_t half = blocks / 2; + const uint32_t begin = get_subblockid() == 0 ? 0 : half; + const uint32_t end = get_subblockid() == 0 ? half : blocks; + for (uint32_t b = begin; b < end; ++b) { + const uint32_t rows = Min(physicalRows, actual.m() - b * physicalRows) / Plan::kRows; + const uint32_t row = round.begin[e] + coord.m() * (CUBE_M / Plan::kRows) + b * Output::ROWS; + output.Tile(raw, bias[e * h], scales, row, rows, coord.n() * CUBE_N, h, t.routedScale); + rowTx.Pump(ExpertRowTx::PUMP_ROWS); + } + } + const uint32_t groupRows=round.begin[e+1]-groupBegin; + if (groupRows>=ExpertRowTx::READY_ROWS || index+1==round.activeCount) { + output.DrainLocal(); + SyncAll(); + rowTx.Enqueue(groupBegin,groupRows,e); + rowTx.Pump(ExpertRowTx::PUMP_ROWS); + groupBegin=round.begin[e+1]; + } + startCore = (startCore + jobs) % t.aicNum; + } + for (; sync < round.activeCount; ++sync) CrossCoreWaitFlag<2>(sync / FLAG_SETS); + output.Finalize(); + rowTx.Finalize(); + +} +} // namespace AfdNativeMega +#endif diff --git a/csrc/npu/ascend_kernels/cmake_files/op_kernel/CMakeLists.txt b/csrc/npu/ascend_kernels/cmake_files/op_kernel/CMakeLists.txt index 1aaadef07..994be3ba2 100644 --- a/csrc/npu/ascend_kernels/cmake_files/op_kernel/CMakeLists.txt +++ b/csrc/npu/ascend_kernels/cmake_files/op_kernel/CMakeLists.txt @@ -14,6 +14,43 @@ set(KERNEL_COMPILE_OPTS) # (resolves the shared utils/op_kernel headers copied next to the kernels). list(APPEND KERNEL_COMPILE_OPTS -I${CMAKE_CURRENT_SOURCE_DIR} -DASCENDC_DUMP=0 --cce-auto-sync=off -Wno-deprecated-declarations -Werror) +# Mega uses CANN's W4A8 BlockMmad directly; ordinary operator builds do not need it. +if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/afd_mega_moe_pull_native.cpp") + set(AFD_TRANSFORMER_DIR "opp/built-in/op_impl/ai_core/tbe/impl/ops_transformer/ascendc") + find_path(AFD_MC2_INCLUDE_DIR + NAMES template_linear_algebra_v2/catlass.hpp + HINTS "${ASCEND_CANN_PACKAGE_PATH}" "$ENV{ASCEND_HOME_PATH}" + PATH_SUFFIXES ${AFD_TRANSFORMER_DIR}/mega_moe/arch22 + NO_DEFAULT_PATH) + if(AFD_MC2_INCLUDE_DIR AND EXISTS "${AFD_MC2_INCLUDE_DIR}/utils/block_mmad_w4a4.hpp") + list(APPEND KERNEL_COMPILE_OPTS -I${AFD_MC2_INCLUDE_DIR}) + set(AFD_MC2_ACTIVATION_DIR "${AFD_MC2_INCLUDE_DIR}") + else() + find_path(AFD_MC2_92_DIR NAMES utils/block_mmad_w4a4.hpp + HINTS "${ASCEND_CANN_PACKAGE_PATH}" "$ENV{ASCEND_HOME_PATH}" + PATH_SUFFIXES ${AFD_TRANSFORMER_DIR}/mega_moe/op_kernel/arch22 NO_DEFAULT_PATH) + find_path(AFD_CATLASS_92_DIR NAMES catlass/catlass.hpp + HINTS "${ASCEND_CANN_PACKAGE_PATH}" "$ENV{ASCEND_HOME_PATH}" + PATH_SUFFIXES ${AFD_TRANSFORMER_DIR}/common/catlass/include NO_DEFAULT_PATH) + if(NOT AFD_MC2_92_DIR OR NOT AFD_CATLASS_92_DIR) + message(FATAL_ERROR "AfdMegaMoePullNative requires CANN MC2 W4A8 and Catlass headers") + endif() + list(APPEND KERNEL_COMPILE_OPTS -DAFD_CANN_CATLASS_92=1 -DCATLASS_ARCH=2201 + -Wno-ignored-attributes -I${AFD_MC2_92_DIR} -I${AFD_CATLASS_92_DIR}) + set(AFD_MC2_ACTIVATION_DIR "${AFD_MC2_92_DIR}") + endif() + # Keep the SDK pipeline while matching Layered's activation rounding. + execute_process( + COMMAND python3 "${CMAKE_CURRENT_SOURCE_DIR}/afd_mega_moe_patch_activation.py" + "${AFD_MC2_ACTIVATION_DIR}/utils/block_epilogue_w4a8post_pertoken_swiglu.hpp" + "${CMAKE_CURRENT_SOURCE_DIR}/afd_mega_moe_activation.h" + RESULT_VARIABLE AFD_ACTIVATION_PATCH_RESULT + ERROR_VARIABLE AFD_ACTIVATION_PATCH_ERROR) + if(NOT AFD_ACTIVATION_PATCH_RESULT EQUAL 0) + message(FATAL_ERROR "Unable to adapt CANN activation: ${AFD_ACTIVATION_PATCH_ERROR}") + endif() +endif() + # set custom compile options (single call to avoid overwriting) if ("${CMAKE_BUILD_TYPE}x" STREQUAL "Debugx") npu_op_kernel_options("" ALL OPTIONS -g -O0 ${KERNEL_COMPILE_OPTS}) diff --git a/csrc/npu/ascend_kernels/grouped_matmul_swiglu_quant_v2_layered/op_kernel/grouped_matmul_swiglu_quant_v2_layered_utils.h b/csrc/npu/ascend_kernels/grouped_matmul_swiglu_quant_v2_layered/op_kernel/grouped_matmul_swiglu_quant_v2_layered_utils.h index defff5118..ce3ca6248 100644 --- a/csrc/npu/ascend_kernels/grouped_matmul_swiglu_quant_v2_layered/op_kernel/grouped_matmul_swiglu_quant_v2_layered_utils.h +++ b/csrc/npu/ascend_kernels/grouped_matmul_swiglu_quant_v2_layered/op_kernel/grouped_matmul_swiglu_quant_v2_layered_utils.h @@ -26,6 +26,7 @@ #define OP_KERNEL_GROUPED_MATMUL_DEQUANT_SWIGLU_QUANT_V2_LAYERED_UTILS_H #include "layered/layered_tensor_addr.h" +#include "afd_swiglu_limit.h" // A8W4 MSD scenario #if defined(ORIG_DTYPE_X) && defined(DT_INT8) && ORIG_DTYPE_X == DT_INT8 && defined(ORIG_DTYPE_ALL_WEIGHT) && \ @@ -71,22 +72,6 @@ constexpr int64_t SIZE_OF_HALF_2 = 2; constexpr uint8_t NUM_8 = 8; constexpr float QUANT_SCALE_INT8 = 127.0f; -// ### PATCH START: DSV4 clipped SwiGLU -// Patch reason: the model limit applies to the dequantized branches before SiLU. -// Patch functionality: clamp in UB without another kernel or GM round trip. -__aicore__ inline void ClampSwiGLUInputs(LocalTensor gate, LocalTensor up, - float limit, uint32_t count) -{ - if (limit > 0.0f) { - Mins(gate, gate, limit, count); - Mins(up, up, limit, count); - PipeBarrier(); - Maxs(up, up, -limit, count); - PipeBarrier(); - } -} -// ### PATCH END: DSV4 clipped SwiGLU - constexpr MatmulConfig matmulCFGUnitFlag{false, false, true, 0, 0, 0, false, false, false, false, false, 0, 0, 0, 0, 0, 0, 0, false}; constexpr MatmulConfig NZ_CFG_MDL = GetMDLConfig(false, false, 0, true, false, false, false); diff --git a/csrc/npu/ascend_kernels/operator_registry.json b/csrc/npu/ascend_kernels/operator_registry.json index 2e152dd13..03375e3a2 100644 --- a/csrc/npu/ascend_kernels/operator_registry.json +++ b/csrc/npu/ascend_kernels/operator_registry.json @@ -9,7 +9,9 @@ "afd_async_combine_send", "afd_async_combine_recv", "grouped_matmul_swiglu_quant_v2_layered", - "grouped_matmul_layered" + "grouped_matmul_layered", + "afd_mega_moe_pull_native", + "afd_mega_moe_comm" ], "ascend950": [ "a2e", @@ -38,8 +40,14 @@ "grouped_matmul_swiglu_quant_v2_layered": { "requires_shmem": false }, + "afd_mega_moe_pull_native": { + "requires_shmem": false + }, "grouped_matmul_layered": { "requires_shmem": false + }, + "afd_mega_moe_comm": { + "requires_shmem": false } } } diff --git a/csrc/npu/ascend_kernels/utils/op_kernel/afd_native_mega_workspace.h b/csrc/npu/ascend_kernels/utils/op_kernel/afd_native_mega_workspace.h new file mode 100644 index 000000000..abdc5c2cd --- /dev/null +++ b/csrc/npu/ascend_kernels/utils/op_kernel/afd_native_mega_workspace.h @@ -0,0 +1,46 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +#ifndef AFD_NATIVE_MEGA_WORKSPACE_H +#define AFD_NATIVE_MEGA_WORKSPACE_H +#include +#include "afd_pull_layout.h" + +#ifdef __CCE_AICORE__ +#define AFD_NATIVE_DEVICE __aicore__ +#else +#define AFD_NATIVE_DEVICE +#endif + +namespace Cam { +constexpr uint32_t NATIVE_MAX_AIVS = 48; +constexpr uint32_t NATIVE_STAGES = 5; +constexpr uint32_t NATIVE_MAX_ROWS = Pull::MAX_CAPACITY; +constexpr uint32_t NATIVE_METADATA_ROWS = 512; +constexpr uint32_t NATIVE_TILE_ROWS = 64; +AFD_NATIVE_DEVICE constexpr uint64_t NativeAlign(uint64_t n) { return (n + 511) / 512 * 512; } + +struct NativeMegaWorkspace { + uint64_t batch, counts, tileCounts, layer, flags, scale1, scale2; + uint64_t x, a1, a2, c1, c2, output, bytes; + AFD_NATIVE_DEVICE NativeMegaWorkspace(uint64_t rows, uint64_t h, uint64_t i, + uint64_t experts, uint64_t tp) { + batch = 0; + counts = NativeAlign(8 * Pull::MetadataWords(tp, experts)); + tileCounts = counts + NativeAlign(8 * experts); + layer = tileCounts + NativeAlign(8 * experts * ((rows + NATIVE_TILE_ROWS - 1) / NATIVE_TILE_ROWS)); + flags = layer + 512; + scale1 = flags + NativeAlign(NATIVE_STAGES * (NATIVE_MAX_AIVS + 1) * 64); + scale2 = scale1 + NativeAlign(4 * rows); + x = scale2 + NativeAlign(4 * rows); + a1 = x + NativeAlign(rows * h); + a2 = a1 + NativeAlign(rows * h); + c1 = a2 + NativeAlign(rows * i); + c2 = c1 + NativeAlign(8 * rows * i); + output = c2 + NativeAlign(4 * rows * h); + // Full BF16 rows are published locally before asynchronous remote TX. + bytes = output + NativeAlign(2 * rows * h); + } +}; +} // namespace Cam +#undef AFD_NATIVE_DEVICE +#endif diff --git a/csrc/npu/ascend_kernels/utils/op_kernel/afd_pull_layout.h b/csrc/npu/ascend_kernels/utils/op_kernel/afd_pull_layout.h new file mode 100644 index 000000000..8c51b1685 --- /dev/null +++ b/csrc/npu/ascend_kernels/utils/op_kernel/afd_pull_layout.h @@ -0,0 +1,135 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +#ifndef AFD_PULL_LAYOUT_H +#define AFD_PULL_LAYOUT_H + +#include + +#ifdef __CCE_AICORE__ +#define AFD_PULL_DEVICE __aicore__ +#else +#define AFD_PULL_DEVICE +#endif + +namespace Cam { namespace Pull { +constexpr uint32_t VERSION = 1; +constexpr uint32_t ALIGNMENT = 64; +constexpr uint32_t MAX_AIVS = 48; +constexpr uint32_t ROW_PADDING = 512; +constexpr uint64_t MAX_SOURCE_ASSIGNMENTS = 65535; +constexpr uint64_t MAX_CAPACITY = 262144; +constexpr int64_t MAGIC = 0x41464401; +constexpr int64_t INIT_MAGIC = 0x41464402; +constexpr uint32_t WORDS_PER_LINE = ALIGNMENT / sizeof(int64_t); +enum State : uint32_t { + SESSION = 0, SOURCE_CAPACITY = 8, TP_SIZE = 9, MAP_SIZE = 16, FAILURE = 24, + SOURCE = 32, ACTIVE_SEQUENCE = 40, ROUND_BEGIN = 48, ROUND_END = 56, + IDLE_BUDGET = 64, ACTIVE_LAYER = 72, CURSOR = 80, NEXT_SOURCE = 88, + COMPLETED = 96, STOP = 104, LAST_SEQUENCE = 112 +}; +AFD_PULL_DEVICE constexpr uint64_t LayerMapOffset(uint64_t sources) { + return LAST_SEQUENCE + sources * WORDS_PER_LINE; +} +AFD_PULL_DEVICE constexpr uint64_t LaneErrorsOffset(uint64_t sources, uint64_t mapSize) { + return LayerMapOffset(sources) + (mapSize + WORDS_PER_LINE - 1) / WORDS_PER_LINE * WORDS_PER_LINE; +} +AFD_PULL_DEVICE constexpr uint64_t StateWords(uint64_t sources, uint64_t mapSize) { + return LaneErrorsOffset(sources, mapSize) + MAX_AIVS * WORDS_PER_LINE; +} + +// One control line has one writer; payload/counts precede the commit line. +enum Header : uint32_t { EPOCH, SEQUENCE, LAYER, TOKENS, ROWS, COMMIT, ERROR, HEADER_LINES }; + +struct Shape { + uint64_t sourceCapacity, hidden, topK, ffnRanks, attnRanks, experts, tp; +}; + +AFD_PULL_DEVICE constexpr uint64_t Align(uint64_t n) { + return (n + ALIGNMENT - 1) / ALIGNMENT * ALIGNMENT; +} + +struct Layout { + uint64_t stats, flags, returned, header, errors, counts, payload, rowStride, bytes; +}; + +// The old Combine reads 32B-aligned statistics and flags at these exact offsets. +AFD_PULL_DEVICE constexpr uint64_t Align32(uint64_t n) { return (n + 31) / 32 * 32; } + +AFD_PULL_DEVICE inline bool MakeLayout(const Shape &s, Layout &out) { + if (s.sourceCapacity == 0 || s.hidden == 0 || s.hidden > 8192 || s.hidden % 256 || + s.topK == 0 || s.ffnRanks == 0 || s.attnRanks == 0 || s.experts == 0 || + s.ffnRanks > 384 || s.attnRanks > 384 || s.experts > 256 || + s.ffnRanks + s.attnRanks > 384 || s.tp == 0 || s.attnRanks % s.tp || + s.topK > s.ffnRanks * s.experts || + s.sourceCapacity > MAX_SOURCE_ASSIGNMENTS / s.topK || + s.sourceCapacity > MAX_CAPACITY / s.tp) return false; + out.stats = 0; + out.flags = Align32(4 * (s.ffnRanks + s.ffnRanks * s.experts) * MAX_AIVS); + out.returned = out.flags + Align32(4 * s.ffnRanks); + out.header = Align(out.returned + 2 * s.sourceCapacity * s.topK * s.hidden); + out.errors = out.header + HEADER_LINES * ALIGNMENT; + out.counts = Align(out.errors + s.ffnRanks * ALIGNMENT); + out.payload = Align(out.counts + 8 * s.ffnRanks * s.experts); + out.rowStride = s.hidden + ROW_PADDING; + out.bytes = out.payload + s.sourceCapacity * s.topK * out.rowStride; + return true; +} + +AFD_PULL_DEVICE constexpr uint64_t MetadataWords(uint64_t tp, uint64_t experts) { + return 7 + tp + tp * experts; +} +AFD_PULL_DEVICE constexpr uint64_t TransportUbBytes(const Shape &s) { + return ALIGNMENT + Align(8 * MetadataWords(s.tp, s.experts)) + + Align(8 * s.ffnRanks * s.experts) + Align(8 * s.experts) + + Align(2 * s.hidden + ROW_PADDING); +} + +// Publish retains these buffers alongside Transport for the entire phase. +AFD_PULL_DEVICE inline uint64_t PublisherStatsBytes(const Shape &s, uint64_t cores) { + const auto stats = Align(4 * (s.ffnRanks + s.ffnRanks * s.experts) * cores); + // Reuse this buffer for the per-lane error lines before loading token offsets. + return stats > cores * ALIGNMENT ? stats : cores * ALIGNMENT; +} +AFD_PULL_DEVICE inline uint64_t PublisherUbBytes(const Shape &s, uint64_t cores) { + const auto e = s.ffnRanks * s.experts; + return PublisherStatsBytes(s, cores) + 2 * Align(4 * e) + + Align(4 * s.topK) + Align(11 * s.hidden + 128); +} + +// Retain the former admission bound while accounting for Mega's per-token routes. +AFD_PULL_DEVICE inline uint64_t CombineUbBytes(const Shape &s, uint64_t cores) { + const auto e = s.ffnRanks * s.experts; + const auto base = 65536 + Align32(2 * e); + const auto preprocess = base + Align32(4 * (s.ffnRanks + e) * cores); + const auto retained = base + Align32(4 * e) + Align32(2 * e) + + 3 * Align32(4 * s.hidden) + Align32(2 * s.hidden) + Align32(4 * s.ffnRanks); + const auto legacy = preprocess > retained ? preprocess : retained; + const auto mega = 2 * Align(4 * e) + 2 * Align(4 * s.topK) + Align(14 * s.hidden); + return legacy > mega ? legacy : mega; +} + +struct Slice { + uint64_t count, skip, local; +}; + +// Empty intersections carry no usable source offset. +AFD_PULL_DEVICE inline Slice Intersect(uint64_t block, uint64_t count, + uint64_t begin, uint64_t capacity) { + const uint64_t lo = block > begin ? block : begin; + const uint64_t end = begin + capacity; + const uint64_t hi = block + count < end ? block + count : end; + return hi > lo ? Slice{hi - lo, lo - block, lo - begin} : Slice{0, 0, 0}; +} + +AFD_PULL_DEVICE constexpr uint64_t Rounds(uint64_t rows, uint64_t capacity) { + return rows == 0 ? 1 : 1 + (rows - 1) / capacity; +} + +AFD_PULL_DEVICE constexpr uint64_t ReturnRow(uint64_t rankPrefix, uint64_t expertPrefix, + uint64_t skip, uint64_t row) { + return rankPrefix + expertPrefix + skip + row; +} +} } // namespace Cam::Pull + +#undef AFD_PULL_DEVICE +#endif diff --git a/csrc/npu/ascend_kernels/utils/op_kernel/afd_pull_protocol.h b/csrc/npu/ascend_kernels/utils/op_kernel/afd_pull_protocol.h new file mode 100644 index 000000000..dd1033e3f --- /dev/null +++ b/csrc/npu/ascend_kernels/utils/op_kernel/afd_pull_protocol.h @@ -0,0 +1,352 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +#ifndef AFD_PULL_PROTOCOL_H +#define AFD_PULL_PROTOCOL_H + +#include "kernel_operator.h" +#include "cam_async_comm_group.h" +#include "afd_pull_layout.h" + +namespace Cam::Pull { +template __aicore__ inline void Fence() { + const auto id = GetTPipePtr()->FetchEventID(event); + AscendC::SetFlag(id); + AscendC::WaitFlag(id); +} + +// Uses the SDK's complete context; ordinary tensor addresses are never peer addresses. +__aicore__ inline GM_ADDR Window(uint32_t rank, uint32_t self) { + auto ctx = reinterpret_cast<__gm__ HcclOpResParam *>(AscendC::GetHcclContext()); + if (rank == self) return reinterpret_cast(ctx->localWindowsIn); + return reinterpret_cast(reinterpret_cast<__gm__ HcclRankRelationResV2 *>( + ctx->remoteRes[rank].nextDevicePtr)->windowsIn); +} + +class Transport { +public: + __aicore__ Transport(GM_ADDR state, Shape shape, uint32_t rank, uint32_t cores, + AscendC::TPipe &pipe, uint32_t id = AscendC::GetBlockIdx()) + : state_(state), shape_(shape), rank_(rank), cores_(cores), id_(id), pipe_(pipe) { + pipe_.Reset(); + pipe_.InitBuffer(control_, ALIGNMENT); + MakeLayout(shape_, layout_); + const auto ctx = reinterpret_cast<__gm__ HcclOpResParam *>( + AscendC::GetHcclContext()); + if (ctx->winSize < layout_.bytes || ctx->rankSize != shape_.attnRanks + shape_.ffnRanks || + ctx->localUsrRankId != rank_) { + if (id_ == 0) WriteState(FAILURE, -7); + } + if (!ReadState(FAILURE) && + (ReadState(SOURCE_CAPACITY) != int64_t(shape_.sourceCapacity) || + ReadState(TP_SIZE) != int64_t(shape_.tp))) { + if (id_ == 0) Fail(0, -6); + } + const uint64_t metaBytes = Align(8 * MetadataWords(shape_.tp, shape_.experts)); + pipe_.InitBuffer(metadata_, metaBytes); + pipe_.InitBuffer(counts_, Align(8 * shape_.ffnRanks * shape_.experts)); + pipe_.InitBuffer(groups_, Align(8 * shape_.experts)); + pipe_.InitBuffer(row_, Align(2 * shape_.hidden + ROW_PADDING)); + } + + __aicore__ int64_t Read(GM_ADDR address) { + AscendC::GlobalTensor gm; + gm.SetGlobalBuffer(reinterpret_cast<__gm__ int64_t *>(address)); + auto value = control_.Get(); + AscendC::DataCopyPad(value, gm, {1U, 8U, 0U, 0U, 0U}, {false, 0U, 0U, 0U}); + Fence(); + const int64_t result = value.GetValue(0); + Fence(); + return result; + } + + __aicore__ void Write(GM_ADDR address, int64_t value) { + auto local = control_.Get(); + local.SetValue(0, value); + AscendC::GlobalTensor gm; + gm.SetGlobalBuffer(reinterpret_cast<__gm__ int64_t *>(address)); + Fence(); + AscendC::DataCopyPad(gm, local, {1U, 8U, 0U, 0U, 0U}); + Fence(); + } + + __aicore__ int64_t ReadState(uint32_t field) { return Read(state_ + field * 8ULL); } + __aicore__ void WriteState(uint32_t field, int64_t value) { Write(state_ + field * 8ULL, value); } + + __aicore__ void Initialize() { + if (rank_ >= shape_.attnRanks) return; + AscendC::SyncAll(); + if (ReadState(FAILURE)) return; + auto zero = row_.Get(); + AscendC::Duplicate(row_.Get(), uint32_t(0), ALIGNMENT / sizeof(uint32_t)); + Fence(); + // Keep the payload untouched; only protocol state requires initialization. + for (uint64_t offset = id_ * ALIGNMENT; + offset < layout_.counts; offset += cores_ * ALIGNMENT) { + if (offset >= layout_.returned && offset < layout_.header) continue; + AscendC::GlobalTensor dst; + dst.SetGlobalBuffer(reinterpret_cast<__gm__ uint8_t *>(Peer(rank_) + offset)); + AscendC::DataCopy(dst, zero, ALIGNMENT); + } + Fence(); + AscendC::SyncAll(); + if (id_ == 0) + Write(Peer(rank_) + layout_.header + EPOCH * ALIGNMENT, ReadState(SESSION)); + } + + __aicore__ void Fail(uint32_t source, int64_t error) { + const auto first = ReadState(FAILURE); + if (!first) WriteState(FAILURE, error); + else error = first; + if (rank_ >= shape_.attnRanks) { + for (uint32_t s = 0; s < shape_.tp; ++s) + Write(Peer(source + s) + layout_.errors + + (rank_ - shape_.attnRanks) * ALIGNMENT, error); + } + } + + __aicore__ bool Select() { + if (ReadState(FAILURE)) return false; + if (ReadState(SOURCE) >= 0) return true; + if (ReadState(STOP)) return false; + const auto sources = shape_.attnRanks / shape_.tp; + const auto first = uint32_t(ReadState(NEXT_SOURCE)); + const auto budget = ReadState(IDLE_BUDGET); + const auto session = ReadState(SESSION); + uint64_t scans = 0; + while (!ReadState(STOP)) { + for (uint32_t n = 0; n < sources; ++n) { + const uint32_t dp = (first + n) % sources; + const uint32_t source = dp * shape_.tp; + const auto expected = ReadState(LAST_SEQUENCE + dp * WORDS_PER_LINE) + 1; + bool ready = true; + int64_t layer = -1; + for (uint32_t s = 0; s < shape_.tp; ++s) { + const auto header = Peer(source + s) + layout_.header; + const auto commit = Read(header + COMMIT * ALIGNMENT); + if (commit < expected) { ready = false; break; } + AscendC::GlobalTensor gm; + gm.SetGlobalBuffer(reinterpret_cast<__gm__ int64_t *>(header)); + auto values = row_.Get(); + AscendC::DataCopy(values, gm, HEADER_LINES * WORDS_PER_LINE); + Fence(); + if (commit != expected || values.GetValue(EPOCH * WORDS_PER_LINE) != session) { + Fail(source, -2); return false; + } + const auto currentLayer = values.GetValue(LAYER * WORDS_PER_LINE); + if (values.GetValue(ERROR * WORDS_PER_LINE) || + values.GetValue(SEQUENCE * WORDS_PER_LINE) != commit || + values.GetValue(COMMIT * WORDS_PER_LINE) != commit || + (layer >= 0 && layer != currentLayer)) { + Fail(source, -3); return false; + } + layer = currentLayer; + if (Read(header + COMMIT * ALIGNMENT) != commit) { ready = false; break; } + } + if (!ready) continue; + const auto mapSize = ReadState(MAP_SIZE); + if (layer < 0 || layer >= mapSize || + Read(state_ + 8 * (LayerMapOffset(sources) + layer)) < 0) { + Fail(source, -4); return false; + } + WriteState(SOURCE, source); + WriteState(ACTIVE_SEQUENCE, expected); + WriteState(ACTIVE_LAYER, layer); + WriteState(CURSOR, 0); + return true; + } + if (budget > 0 && ++scans >= uint64_t(budget)) return false; + } + return false; + } + + // Load immutable original counts; slice counts are separate from return coordinates. + __aicore__ uint64_t LoadMetadata() { + auto meta = metadata_.Get(); + const auto source = uint32_t(ReadState(SOURCE)); + const uint32_t ffn = rank_ - shape_.attnRanks; + uint64_t total = 0; + for (uint32_t s = 0; s < shape_.tp; ++s) { + AscendC::GlobalTensor gm; + gm.SetGlobalBuffer(reinterpret_cast<__gm__ int64_t *>(Peer(source + s) + layout_.counts)); + auto counts = counts_.Get(); + AscendC::DataCopyPad(counts, gm, + {1U, uint32_t(8 * shape_.ffnRanks * shape_.experts), 0U, 0U, 0U}, + {false, 0U, 0U, 0U}); + Fence(); + uint64_t prefix = 0, sourceTotal = 0; + for (uint32_t i = 0; i < shape_.ffnRanks * shape_.experts; ++i) { + const auto count = counts.GetValue(i); + if (count < 0 || uint64_t(count) > shape_.sourceCapacity * shape_.topK) { + if (id_ == 0) Fail(source, -5); + return 0; + } + if (i < ffn * shape_.experts) prefix += count; + sourceTotal += count; + if (i / shape_.experts == ffn) { + const uint32_t index = 5 + shape_.tp + s * shape_.experts + i % shape_.experts; + meta.SetValue(index, count); + total += count; + } + } + if (sourceTotal > shape_.sourceCapacity * shape_.topK || + uint64_t(Read(Peer(source + s) + layout_.header + ROWS * ALIGNMENT)) != sourceTotal) { + if (id_ == 0) Fail(source, -5); + return 0; + } + meta.SetValue(5 + s, prefix); + Fence(); + } + meta.SetValue(0, total); + meta.SetValue(1, source); + meta.SetValue(2, ReadState(ACTIVE_LAYER)); + meta.SetValue(MetadataWords(shape_.tp, shape_.experts) - 2, ReadState(SESSION)); + meta.SetValue(MetadataWords(shape_.tp, shape_.experts) - 1, ReadState(ACTIVE_SEQUENCE)); + return total; + } + + // Mega's fused pipeline owns payload DMA; this prepares only control metadata. + __aicore__ void PrepareRound(GM_ADDR batch, GM_ADDR groupCounts, uint64_t capacity) { + const bool continuing = ReadState(SOURCE) >= 0; + Select(); + auto meta = metadata_.Get(); + auto groups = groups_.Get(); + if (ReadState(FAILURE) || ReadState(SOURCE) < 0) { + ClearMetadata(); + SaveMetadata(batch, groupCounts); + return; + } + // A retains the original counts until the last round has returned its output. + if (continuing) LoadSnapshot(batch); + else LoadMetadata(); + const uint64_t total = meta.GetValue(0); + const auto source = ReadState(SOURCE); + const auto sequence = ReadState(ACTIVE_SEQUENCE); + if (!ReadState(FAILURE)) { + for (uint32_t s = 0; s < shape_.tp; ++s) { + if (Read(Peer(source + s) + layout_.header + COMMIT * ALIGNMENT) != sequence) { + Fail(source, -9); break; + } + } + } + if (ReadState(FAILURE)) { + ClearMetadata(); + SaveMetadata(batch, groupCounts); + return; + } + const uint64_t begin = ReadState(CURSOR); + if (begin > total) { + Fail(source, -9); + ClearMetadata(); + SaveMetadata(batch, groupCounts); + return; + } + const uint64_t end = begin + capacity < total ? begin + capacity : total; + meta.SetValue(3, begin); + meta.SetValue(4, end); + if (id_ == 0) { + WriteState(ROUND_BEGIN, begin); + WriteState(ROUND_END, end); + } + uint64_t block = 0; + for (uint32_t e = 0; e < shape_.experts; ++e) { + uint64_t group = 0; + for (uint32_t s = 0; s < shape_.tp; ++s) { + const uint64_t count = meta.GetValue(5 + shape_.tp + s * shape_.experts + e); + const auto slice = Intersect(block, count, begin, capacity); + group += slice.count; + block += count; + } + groups.SetValue(e, group); + } + SaveMetadata(batch, groupCounts); + Fence(); + } + + // Called only by the control owner after the mixed-core output drain. + __aicore__ bool ValidateReturn(GM_ADDR batch, uint64_t inputRows) { + if (ReadState(FAILURE) || ReadState(SOURCE) < 0) return false; + LoadSnapshot(batch); + auto meta = metadata_.Get(); + const auto source = ReadState(SOURCE); + if (source < 0 || uint64_t(source) + shape_.tp > shape_.attnRanks || + source % shape_.tp || meta.GetValue(1) != source || + meta.GetValue(2) != ReadState(ACTIVE_LAYER) || + meta.GetValue(MetadataWords(shape_.tp, shape_.experts) - 2) != ReadState(SESSION) || + meta.GetValue(MetadataWords(shape_.tp, shape_.experts) - 1) != ReadState(ACTIVE_SEQUENCE) || + meta.GetValue(3) != ReadState(CURSOR) || + meta.GetValue(3) != ReadState(ROUND_BEGIN) || + meta.GetValue(4) != ReadState(ROUND_END) || + meta.GetValue(3) < 0 || meta.GetValue(4) < meta.GetValue(3) || + meta.GetValue(4) > meta.GetValue(0) || + uint64_t(meta.GetValue(4) - meta.GetValue(3)) > inputRows) { + Fail(source >= 0 && uint64_t(source) + shape_.tp <= shape_.attnRanks ? source : 0, -9); + } + if (ReadState(FAILURE)) return false; + return true; + } + + __aicore__ void CompleteRound(GM_ADDR batch, uint64_t inputRows) { + if (!ValidateReturn(batch, inputRows)) return; + auto meta = metadata_.Get(); + const uint64_t end = meta.GetValue(4); + const uint32_t source = meta.GetValue(1); + if (end < uint64_t(meta.GetValue(0))) { WriteState(CURSOR, end); return; } + const auto sequence = ReadState(ACTIVE_SEQUENCE); + WriteState(LAST_SEQUENCE + source / shape_.tp * WORDS_PER_LINE, sequence); + WriteState(SOURCE, -1); + WriteState(CURSOR, 0); + WriteState(NEXT_SOURCE, (source / shape_.tp + 1) % (shape_.attnRanks / shape_.tp)); + WriteState(COMPLETED, ReadState(COMPLETED) + 1); + auto flag = control_.Get(); + for (uint32_t s = 0; s < shape_.tp; ++s) { + flag.SetValue(0, meta.GetValue(5 + s) + 1); + AscendC::GlobalTensor dst; + dst.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t *>(Peer(source + s) + layout_.flags) + + rank_ - shape_.attnRanks); + Fence(); + AscendC::DataCopyPad(dst, flag, {1U, 4U, 0U, 0U, 0U}); + Fence(); + } + } + +private: + __aicore__ GM_ADDR Peer(uint32_t rank) { return Window(rank, rank_); } + __aicore__ void LoadSnapshot(GM_ADDR batch) { + AscendC::GlobalTensor gm; + gm.SetGlobalBuffer(reinterpret_cast<__gm__ int64_t *>(batch)); + AscendC::DataCopyPad(metadata_.Get(), gm, + {1U, uint32_t(8 * MetadataWords(shape_.tp, shape_.experts)), 0U, 0U, 0U}, + {false, 0U, 0U, 0U}); + Fence(); + } + __aicore__ void ClearMetadata() { + auto meta = metadata_.Get(); + auto groups = groups_.Get(); + for (uint32_t i = 0; i < MetadataWords(shape_.tp, shape_.experts); ++i) meta.SetValue(i, 0); + for (uint32_t e = 0; e < shape_.experts; ++e) groups.SetValue(e, 0); + const auto offset = LayerMapOffset(shape_.attnRanks / shape_.tp); + for (int64_t layer = 0; layer < ReadState(MAP_SIZE); ++layer) { + if (Read(state_ + 8 * (offset + layer)) >= 0) { meta.SetValue(2, layer); break; } + } + } + __aicore__ void SaveMetadata(GM_ADDR batch, GM_ADDR counts) { + if (id_ != 0) return; + AscendC::GlobalTensor dstBatch, dstCounts; + dstBatch.SetGlobalBuffer(reinterpret_cast<__gm__ int64_t *>(batch)); + dstCounts.SetGlobalBuffer(reinterpret_cast<__gm__ int64_t *>(counts)); + Fence(); + AscendC::DataCopyPad(dstBatch, metadata_.Get(), + {1U, uint32_t(8 * MetadataWords(shape_.tp, shape_.experts)), 0U, 0U, 0U}); + AscendC::DataCopyPad(dstCounts, groups_.Get(), + {1U, uint32_t(8 * shape_.experts), 0U, 0U, 0U}); + Fence(); + } + GM_ADDR state_; + Shape shape_; + Layout layout_{}; + uint32_t rank_, cores_, id_; + AscendC::TPipe &pipe_; + AscendC::TBuf control_, metadata_, counts_, groups_, row_; +}; +} // namespace Cam::Pull +#endif diff --git a/csrc/npu/ascend_kernels/utils/op_kernel/afd_swiglu_limit.h b/csrc/npu/ascend_kernels/utils/op_kernel/afd_swiglu_limit.h new file mode 100644 index 000000000..c0db806f2 --- /dev/null +++ b/csrc/npu/ascend_kernels/utils/op_kernel/afd_swiglu_limit.h @@ -0,0 +1,32 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +/** + * Shared DSV4 clipped-SwiGLU clamp used by every AFD GMM/epilogue path. + * + * The model's swiglu_limit applies to the dequantized gate/up branches before + * SiLU: gate is upper-clamped at L, up is clamped to [-L, L]. Zero disables + * clipping, matching the ACLNN and PyTorch surfaces that default to 0. + * + * This lives in utils/op_kernel so the layered W4A8/A4W4 postprocessors and + * the native MegaMoE activation share one definition instead of each carrying + * its own Mins/Maxs sequence. + */ +#ifndef AFD_SWIGLU_LIMIT_H +#define AFD_SWIGLU_LIMIT_H + +#include "kernel_operator.h" + +__aicore__ inline void ClampSwiGLUInputs(AscendC::LocalTensor gate, + AscendC::LocalTensor up, + float limit, uint32_t count) +{ + if (limit > 0.0f) { + AscendC::Mins(gate, gate, limit, count); + AscendC::Mins(up, up, limit, count); + AscendC::PipeBarrier(); + AscendC::Maxs(up, up, -limit, count); + AscendC::PipeBarrier(); + } +} + +#endif // AFD_SWIGLU_LIMIT_H diff --git a/csrc/npu/pybind/CMakeLists.txt b/csrc/npu/pybind/CMakeLists.txt index f9877545a..6b6f16291 100644 --- a/csrc/npu/pybind/CMakeLists.txt +++ b/csrc/npu/pybind/CMakeLists.txt @@ -61,7 +61,9 @@ if(AFD_SOC_VERSION STREQUAL "") endif() if(AFD_SOC_VERSION STREQUAL "910c" OR AFD_SOC_VERSION MATCHES "^ascend910_93") target_sources(afd_plugin_ascend_C PRIVATE - "${CMAKE_CURRENT_LIST_DIR}/torch_binding_cam_async.cpp") + "${CMAKE_CURRENT_LIST_DIR}/torch_binding_cam_async.cpp" + "${CMAKE_CURRENT_LIST_DIR}/torch_binding_mega_moe_pull.cpp" + "${CMAKE_CURRENT_LIST_DIR}/torch_binding_mega_moe_comm.cpp") elseif(AFD_SOC_VERSION STREQUAL "950" OR AFD_SOC_VERSION MATCHES "^[Aa]scend950") message(STATUS "CAM asynchronous bindings disabled for Ascend 950") else() diff --git a/csrc/npu/pybind/torch_binding_mega_moe_comm.cpp b/csrc/npu/pybind/torch_binding_mega_moe_comm.cpp new file mode 100644 index 000000000..b1dfab444 --- /dev/null +++ b/csrc/npu/pybind/torch_binding_mega_moe_comm.cpp @@ -0,0 +1,249 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) Huawei Technologies Co., Ltd. 2025-2025. All rights reserved. +// Copyright contributors to the AFD plugin project +// +// Adapted from cam_async/src/comm_operator/pybind/moe_*_async.cpp. +// Dedicated Mega pull adapters; the original CAM entry points remain unchanged. + +#include +#include +#include + +#include +#include + +#include "pytorch_extension/op_api_common.h" +#include "utils/op_kernel/afd_pull_layout.h" +#include "afd_mega_moe_comm/op_api/aclnn_afd_mega_moe_comm.h" + +namespace afd_plugin::mega_comm { +namespace { + +constexpr int64_t MAX_SEQUENCE_LENGTH = 1024 * 256; + +void check_tensor(const at::Tensor &tensor, const at::Tensor &anchor, + const char *name) { + TORCH_CHECK(!tensor.requires_grad(), name, + ": CAM asynchronous operators support inference only"); + TORCH_CHECK(tensor.device() == anchor.device(), name, + " must be on the same device as the input"); + TORCH_CHECK(tensor.is_contiguous(), name, " must be contiguous"); +} + +void check_topology(int64_t hidden_size, int64_t top_k, int64_t moe_rank_num, + int64_t attn_rank_num, int64_t route_expert_num_per_moe, + int64_t world_size) { + TORCH_CHECK(hidden_size > 0, "hidden_size must be positive"); + TORCH_CHECK(moe_rank_num > 0 && attn_rank_num > 0 && + route_expert_num_per_moe > 0, + "rank counts and route_expert_num_per_moe must be positive"); + TORCH_CHECK(world_size == moe_rank_num + attn_rank_num, + "world_size must equal attn_rank_num + moe_rank_num"); + TORCH_CHECK(top_k > 0 && top_k <= moe_rank_num * route_expert_num_per_moe, + "top_k must be in [1, moe_rank_num * route_expert_num_per_moe]"); +} + +void check_tp(int64_t max_seq_len, int64_t tp_size, int64_t attn_rank_num) { + TORCH_CHECK(tp_size > 0 && attn_rank_num % tp_size == 0, + "tp_size must be positive and divide attn_rank_num"); + TORCH_CHECK(max_seq_len >= tp_size && max_seq_len <= MAX_SEQUENCE_LENGTH && max_seq_len % tp_size == 0, + "max_seq_len must be in [tp_size, 262144] and divisible by tp_size"); +} + +void check_pull_state(const at::Tensor &state, int64_t attn_ranks, + int64_t tp_size, const at::Tensor &anchor) { + TORCH_CHECK(tp_size > 0 && attn_ranks % tp_size == 0, "pull TP topology is invalid"); + TORCH_CHECK(anchor.scalar_type() == at::kBFloat16, "pull requires BF16"); + TORCH_CHECK(state.scalar_type() == at::kHalf && state.dim() == 1 && + state.numel() >= int64_t(4 * Cam::Pull::StateWords(attn_ranks / tp_size, 0)), + "pull state is missing; initialize with afd_mega_pull_init"); +} + +void check_float_input(const at::Tensor &tensor) { + check_tensor(tensor, tensor, "input"); + TORCH_CHECK(tensor.scalar_type() == at::kHalf || + tensor.scalar_type() == at::kBFloat16, + "input must have dtype float16 or bfloat16"); +} + +} // namespace + +template +std::vector run_comm(const at::Tensor &x, const at::Tensor &ids, + const at::Tensor &probabilities, const at::Tensor &state, + at::IntArrayRef c, c10::string_view group_name) { + Cam::Pull::Layout layout{}; + TORCH_CHECK(Cam::Pull::MakeLayout({uint64_t(c[1] / c[8]), uint64_t(c[2]), uint64_t(c[3]), + uint64_t(c[4]), uint64_t(c[5]), uint64_t(c[6]), uint64_t(c[8])}, layout), "invalid Mega communication layout"); + std::vector out{ + at::empty({c[0] == 4 ? std::max(1, c[9]) * c[2] : 1}, x.options()), + at::empty({1}, x.options().dtype(at::kLong))}; + if constexpr (EXECUTE_NPU) { + std::string group(group_name.data(), group_name.size()); + auto group_ptr = group.data(); + auto &combined = out[0], &completion = out[1]; + EXEC_NPU_CMD(aclnnAfdMegaMoeComm, x, ids, probabilities, state, c, group_ptr, + combined, completion); + } + return out; +} + +// Source: moe_dispatch_send_async.cpp::cam_dispatch_send_async_impl_npu/meta. +// AFD adapter: shared NPU/Meta validation and allocation; no autograd wrapper. +// Signature retains the source arguments; the template chooses execution only. +template +at::Tensor dispatch_send( + const at::Tensor &x, const at::Tensor &expert_ids, + const at::Tensor &comm_args, const int64_t comm_id, + const int64_t max_seq_len, const int64_t batch_size, + const int64_t hidden_size, const int64_t top_k, + const int64_t moe_rank_num, const int64_t attn_rank_num, + const int64_t route_expert_num_per_moe, const int64_t attn_rank_id, + const int64_t world_size, const int64_t layer_index, + const int64_t tp_size, const int64_t dynamic_quant, + c10::string_view group_name) { + TORCH_CHECK(comm_id == Cam::Pull::MAGIC, "Mega communication requires pull_v1"); + check_float_input(x); + check_pull_state(comm_args, attn_rank_num, tp_size, x); + check_tensor(expert_ids, x, "expert_ids"); + check_tensor(comm_args, x, "comm_args"); + check_topology(hidden_size, top_k, moe_rank_num, attn_rank_num, + route_expert_num_per_moe, world_size); + check_tp(max_seq_len, tp_size, attn_rank_num); + TORCH_CHECK(batch_size >= 0 && + batch_size <= max_seq_len / tp_size, + "batch_size is outside the dispatch envelope"); + TORCH_CHECK(attn_rank_id >= 0 && attn_rank_id < attn_rank_num, + "attn_rank_id must identify an Attention rank"); + TORCH_CHECK(layer_index >= 0, "layer_index must be nonnegative"); + TORCH_CHECK(dynamic_quant == 1, "pull requires dynamic_quant=1"); + TORCH_CHECK(x.dim() == 2 && x.size(0) == batch_size && x.size(1) == hidden_size, + "x must have shape [batch_size, hidden_size]"); + TORCH_CHECK(expert_ids.scalar_type() == at::kInt && expert_ids.dim() == 2 && + expert_ids.size(0) == batch_size && expert_ids.size(1) == top_k, + "expert_ids must be int32[batch_size, top_k]"); + TORCH_CHECK(comm_args.scalar_type() == at::kHalf, + "comm_args must have dtype float16"); + auto probs = at::empty({1}, x.options().dtype(at::kFloat)); + return run_comm(x, expert_ids, probs, comm_args, + {1,max_seq_len,hidden_size,top_k,moe_rank_num,attn_rank_num,route_expert_num_per_moe, + attn_rank_id,tp_size,batch_size,layer_index}, group_name)[1]; +} + +// Source: moe_combine_recv_async.cpp::cam_combine_recv_async_impl_npu/meta. +// Empty transactions still launch a physical output row to consume completion flags. +// Signature retains the source arguments; the template chooses execution only. +template +at::Tensor combine_recv( + const at::Tensor &expand_x, const at::Tensor &expert_ids, + const at::Tensor &expert_scales, const at::Tensor &comm_args, + const int64_t comm_id, const int64_t batch_size, const int64_t hidden_size, + const int64_t top_k, const int64_t moe_rank_num, const int64_t attn_rank_num, + const int64_t route_expert_num_per_moe, const int64_t attn_rank_id, + const int64_t world_size, c10::string_view group_name, + int64_t max_seq_len, int64_t tp_size) { + TORCH_CHECK(comm_id == Cam::Pull::MAGIC, "Mega communication requires pull_v1"); + check_float_input(expand_x); + check_tp(max_seq_len, tp_size, attn_rank_num); + check_pull_state(comm_args, attn_rank_num, tp_size, expand_x); + check_tensor(expert_ids, expand_x, "expert_ids"); + check_tensor(expert_scales, expand_x, "expert_scales"); + check_tensor(comm_args, expand_x, "comm_args"); + check_topology(hidden_size, top_k, moe_rank_num, attn_rank_num, + route_expert_num_per_moe, world_size); + TORCH_CHECK(batch_size >= 0 && + batch_size <= MAX_SEQUENCE_LENGTH, + "batch_size is outside the combine envelope"); + TORCH_CHECK(attn_rank_id >= 0 && attn_rank_id < attn_rank_num, + "attn_rank_id must identify an Attention rank"); + TORCH_CHECK(expand_x.dim() == 1 && expand_x.size(0) == 1, + "expand_x must be a one-element anchor"); + TORCH_CHECK(expert_ids.scalar_type() == at::kInt && expert_ids.dim() == 2 && + expert_ids.size(0) == batch_size && expert_ids.size(1) == top_k, + "expert_ids must be int32[batch_size, top_k]"); + TORCH_CHECK(expert_scales.scalar_type() == at::kFloat && + expert_scales.dim() == 2 && expert_scales.size(0) == batch_size && + expert_scales.size(1) == top_k, + "expert_scales must be float32[batch_size, top_k]"); + TORCH_CHECK(comm_args.scalar_type() == at::kHalf, + "comm_args must have dtype float16"); + TORCH_CHECK(batch_size <= max_seq_len / tp_size, "batch_size exceeds Mega source capacity"); + auto out = run_comm(expand_x, expert_ids, expert_scales, comm_args, + {4,max_seq_len,hidden_size,top_k,moe_rank_num,attn_rank_num,route_expert_num_per_moe, + attn_rank_id,tp_size,batch_size,0}, group_name)[0]; + return out.view({std::max(1,batch_size),hidden_size}).narrow(0,0,batch_size); +} + +at::Tensor pull_init(const at::Tensor &anchor, at::IntArrayRef config, + at::IntArrayRef layer_ids, int64_t epoch, c10::string_view group_name) { + TORCH_CHECK(config.size() == 8, "pull config must be [Ncap,H,K,F,A,E,rank,TP]"); + const auto c = config.data(); + Cam::Pull::Layout layout{}; + TORCH_CHECK(Cam::Pull::MakeLayout({uint64_t(c[0]), uint64_t(c[1]), uint64_t(c[2]), + uint64_t(c[3]), uint64_t(c[4]), uint64_t(c[5]), uint64_t(c[7])}, layout), + "invalid pull layout or uint16 Combine bound"); + TORCH_CHECK(c[6] >= 0 && c[6] < c[3] + c[4] && epoch > 0, + "pull rank/epoch is invalid"); + TORCH_CHECK(anchor.scalar_type() == at::kBFloat16, "pull requires BF16"); + int64_t map_size = 0; + for (auto layer : layer_ids) { + TORCH_CHECK(layer >= 0 && layer < 65536, "pull layer map is out of bounds"); + map_size = std::max(map_size, layer + 1); + } + TORCH_CHECK(c[6] < c[4] || !layer_ids.empty(), "FFN pull requires a layer table"); + std::vector values(Cam::Pull::StateWords(c[4] / c[7], map_size), 0); + values[Cam::Pull::SESSION] = epoch; + values[Cam::Pull::SOURCE_CAPACITY] = c[0]; + values[Cam::Pull::TP_SIZE] = c[7]; + values[Cam::Pull::MAP_SIZE] = map_size; + values[Cam::Pull::SOURCE] = -1; + const auto offset = Cam::Pull::LayerMapOffset(c[4] / c[7]); + std::fill(values.begin() + offset, values.begin() + offset + map_size, -1); + for (int64_t slot = 0; slot < int64_t(layer_ids.size()); ++slot) { + const auto layer = layer_ids[slot]; + TORCH_CHECK(layer >= 0 && values[offset + layer] == -1, "pull layer IDs must be unique and nonnegative"); + values[offset + layer] = slot; + } + auto host = at::from_blob(values.data(), {int64_t(values.size())}, at::TensorOptions().dtype(at::kLong)); + auto state = host.to(anchor.device(), at::kLong, false, true).view(at::kHalf); + if (anchor.is_meta()) return state; + auto ids = at::empty({1}, anchor.options().dtype(at::kInt)); + auto probs = at::empty({1}, anchor.options().dtype(at::kFloat)); + run_comm(anchor, ids, probs, state, + {0,c[0]*c[7],c[1],c[2],c[3],c[4],c[5],c[6],c[7],0,0}, group_name); + return state; +} + +} // namespace afd_plugin::mega_comm + +TORCH_LIBRARY_FRAGMENT(afd_ascend, ops) { + ops.def("afd_mega_pull_init(Tensor anchor, int[] config, int[] layer_ids, int epoch, str group_name) -> Tensor"); + ops.def("afd_mega_dispatch_send(Tensor x, Tensor expert_ids, Tensor(a!) comm_args, " + "int comm_id, int max_seq_len, int batch_size, int hidden_size, int top_k, " + "int moe_rank_num, int attn_rank_num, int route_expert_num_per_moe, " + "int attn_rank_id, int world_size, int layer_index, int tp_size, " + "int dynamic_quant, str group_name) -> Tensor"); + ops.def("afd_mega_combine_recv(Tensor expand_x, Tensor expert_ids, " + "Tensor expert_scales, Tensor(a!) comm_args, int comm_id, int batch_size, " + "int hidden_size, int top_k, int moe_rank_num, int attn_rank_num, " + "int route_expert_num_per_moe, int attn_rank_id, int world_size, " + "str group_name, int max_seq_len, int tp_size) -> Tensor"); +} + +TORCH_LIBRARY_IMPL(afd_ascend, PrivateUse1, ops) { + ops.impl("afd_mega_pull_init", &afd_plugin::mega_comm::pull_init); + ops.impl("afd_mega_dispatch_send", &afd_plugin::mega_comm::dispatch_send); + ops.impl("afd_mega_combine_recv", &afd_plugin::mega_comm::combine_recv); +} + +// Explicitly enforce the inference-only contract, including when grad mode is on. +TORCH_LIBRARY_IMPL(afd_ascend, AutogradPrivateUse1, ops) { + ops.impl("afd_mega_dispatch_send", &afd_plugin::mega_comm::dispatch_send); + ops.impl("afd_mega_combine_recv", &afd_plugin::mega_comm::combine_recv); +} + +TORCH_LIBRARY_IMPL(afd_ascend, Meta, ops) { + ops.impl("afd_mega_pull_init", &afd_plugin::mega_comm::pull_init); + ops.impl("afd_mega_dispatch_send", &afd_plugin::mega_comm::dispatch_send); + ops.impl("afd_mega_combine_recv", &afd_plugin::mega_comm::combine_recv); +} diff --git a/csrc/npu/pybind/torch_binding_mega_moe_pull.cpp b/csrc/npu/pybind/torch_binding_mega_moe_pull.cpp new file mode 100644 index 000000000..ef0a17d15 --- /dev/null +++ b/csrc/npu/pybind/torch_binding_mega_moe_pull.cpp @@ -0,0 +1,115 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +#include +#include +#include +#include +#include "pytorch_extension/op_api_common.h" +#include "pytorch_extension/NPUBridge.h" +#include "utils/op_kernel/afd_native_mega_workspace.h" +#include "utils/op_kernel/afd_pull_layout.h" +#include "afd_mega_moe_pull_native/op_api/aclnn_afd_mega_moe_pull_native.h" + +namespace { +int64_t workspace_size(int64_t rows, int64_t h, int64_t i, int64_t experts, int64_t tp) { + TORCH_CHECK(rows > 0 && rows <= Cam::NATIVE_MAX_ROWS && h > 0 && h <= 8192 && + i > 0 && i <= 8192 && h % 256 == 0 && i % 256 == 0 && + experts > 0 && experts <= 128 && tp > 0 && tp <= 384, + "native Mega workspace dimensions are unsupported"); + return Cam::NativeMegaWorkspace(rows, h, i, experts, tp).bytes; +} + +template +at::Tensor native_mega(const at::Tensor &anchor, const at::Tensor &state, + const at::Tensor &layer_map, const at::Tensor &scratch, + at::TensorList w13, at::TensorList w2, at::TensorList s13, at::TensorList s2, + at::TensorList b13, at::TensorList b2, at::IntArrayRef c, + c10::string_view group_name, double routed_scale, double swiglu_limit) { + TORCH_CHECK(c.size() == 14, "native Mega config requires 14 fields"); + for (auto n : c) TORCH_CHECK(n > 0 && n <= 262144, "native Mega config out of range"); + TORCH_CHECK(c[8] == c[4] + c[5] && c[7] >= c[5] && c[7] < c[8] && + c[0] % c[9] == 0 && c[5] % c[9] == 0 && (c[11] == 1 || c[11] == 2) && c[12] == 1 && + std::isfinite(routed_scale), "native Mega requires a valid F rank and per-channel weights"); + TORCH_CHECK(std::isfinite(swiglu_limit) && swiglu_limit >= 0.0 && + swiglu_limit <= std::numeric_limits::max(), + "native Mega swiglu_limit must be finite, nonnegative and representable as float32"); + Cam::Pull::Shape shape{uint64_t(c[0] / c[9]), uint64_t(c[1]), uint64_t(c[3]), + uint64_t(c[4]), uint64_t(c[5]), uint64_t(c[6]), uint64_t(c[9])}; + Cam::Pull::Layout layout{}; + TORCH_CHECK(Cam::Pull::MakeLayout(shape, layout), "native Mega pull shape unsupported"); + auto check = [&](const at::Tensor &x, at::ScalarType dtype) { + TORCH_CHECK(x.device() == anchor.device() && x.scalar_type() == dtype && + x.is_contiguous() && !x.requires_grad(), "native Mega tensor has wrong device, dtype or layout"); + }; + check(anchor, at::kBFloat16); + TORCH_CHECK(anchor.numel() == 1, "native Mega requires a one-element anchor"); + check(state, at::kHalf); + check(layer_map, at::kLong); + TORCH_CHECK(layer_map.dim() == 1 && layer_map.numel() > 0 && layer_map.numel() <= 65536, + "native Mega layer map must cover model layer IDs"); + TORCH_CHECK(state.dim() == 1 && state.numel() >= int64_t(4 * + Cam::Pull::StateWords(c[5] / c[9], layer_map.numel())), "native Mega pull state is missing"); + check(scratch, at::kByte); + TORCH_CHECK(scratch.dim() == 1 && scratch.numel() == workspace_size(c[10], c[1], c[2], c[6], c[9]), + "native Mega scratch has the wrong size"); + const bool packed = c[11] == 1; + TORCH_CHECK(packed || (b13.empty() && b2.empty()), "native Mega W8A8 requires empty compensation lists"); + const std::vector> lists{ + {w13, packed ? at::kInt : at::kChar}, {w2, packed ? at::kInt : at::kChar}, + {s13, packed ? at::kLong : at::kFloat}, {s2, packed ? at::kLong : at::kBFloat16}}; + for (const auto &list : lists) { + TORCH_CHECK(int64_t(list.first.size()) == c[13], "native Mega layer list length mismatch"); + for (const auto &tensor : list.first) check(tensor, list.second); + } + if (packed) { + for (const auto list : {b13, b2}) { + TORCH_CHECK(int64_t(list.size()) == c[13], "native Mega layer list length mismatch"); + for (const auto &tensor : list) check(tensor, at::kFloat); + } + } + for (int64_t l = 0; l < c[13]; ++l) { + if constexpr (EXECUTE) { + TORCH_CHECK(vllm_ascend::NPUBridge::GetNpuStorageImplDesc(w13[l]).npu_format_ == ACL_FORMAT_FRACTAL_NZ && + vllm_ascend::NPUBridge::GetNpuStorageImplDesc(w2[l]).npu_format_ == ACL_FORMAT_FRACTAL_NZ, + "native Mega requires physically NZ packed weights"); + } + TORCH_CHECK(w13[l].sizes() == at::IntArrayRef({c[6], c[1], packed ? c[2] / 4 : 2 * c[2]}) && + w2[l].sizes() == at::IntArrayRef({c[6], c[2], packed ? c[1] / 8 : c[1]}), + "native Mega weights mismatch"); + TORCH_CHECK(s13[l].numel() == c[6] * 2 * c[2] && s2[l].numel() == c[6] * c[1], + "native Mega scale dimensions mismatch"); + if (packed) TORCH_CHECK(b13[l].numel() == c[6] * 2 * c[2] && b2[l].numel() == c[6] * c[1], + "native Mega compensation dimensions mismatch"); + } + auto completion = at::empty({1}, anchor.options().dtype(at::kLong)); + if constexpr (EXECUTE) { + std::string group(group_name.data(), group_name.size()); + auto group_ptr = group.data(); + // CANN matches dynamic inputs by dtype even when the W8A8 plan never reads them. + std::vector empty; + if (!packed) empty.push_back(s13[0].slice(0, 0, 0)); + const at::TensorList compensation13 = packed ? b13 : at::TensorList(empty); + const at::TensorList compensation2 = packed ? b2 : at::TensorList(empty); + EXEC_NPU_CMD(aclnnAfdMegaMoePullNative, anchor, state, layer_map, scratch, + w13, w2, s13, s2, compensation13, compensation2, c, group_ptr, routed_scale, swiglu_limit, completion); + } + return completion; +} +} // namespace + +TORCH_LIBRARY_FRAGMENT(afd_ascend, m) { + m.def("afd_mega_moe_pull_workspace_size(int rows, int hidden, int intermediate, int experts, int tp) -> int", + &workspace_size); + m.def("afd_mega_moe_pull_native(Tensor anchor, Tensor(a!) state, Tensor layer_map, Tensor(b!) scratch, " + "Tensor[] w13, Tensor[] w2, Tensor[] s13, Tensor[] s2, Tensor[] b13, Tensor[] b2, " + "int[] config, str group_name, float routed_scale, float swiglu_limit=0.0) -> Tensor"); +} +TORCH_LIBRARY_IMPL(afd_ascend, PrivateUse1, m) { + m.impl("afd_mega_moe_pull_native", &native_mega); +} +TORCH_LIBRARY_IMPL(afd_ascend, AutogradPrivateUse1, m) { + m.impl("afd_mega_moe_pull_native", &native_mega); +} +TORCH_LIBRARY_IMPL(afd_ascend, Meta, m) { + m.impl("afd_mega_moe_pull_native", &native_mega); +} diff --git a/docs/README.md b/docs/README.md new file mode 100644 index 000000000..50d717def --- /dev/null +++ b/docs/README.md @@ -0,0 +1,31 @@ +# AFD documentation + +## Deploy and use + +1. [Install](installation.md) the pinned CUDA or Ascend environment. +2. Check [supported configurations and validation limits](support.md). +3. Choose a [deployment recipe](../recipe/README.md) and configure the + [Attention/FFN roles](configuration.md). + +| Backend | Connector guides | +| --- | --- | +| CUDA | [NCCL P2P](gpu/NCCL_P2P_CONNECTOR_USER_GUIDE.md), [DeepEP M2N](gpu/DEEPEP_M2N_CONNECTOR_USER_GUIDE.md) | +| Ascend NPU | [CAMP2P](npu/CAM_P2P_CONNECTOR_USER_GUIDE.md), [CAM async](npu/CAM_ASYNC_CONNECTOR_USER_GUIDE.md) | + +## Develop + +- [Module design and code ownership](design/module/index.md) +- [Contribution rules](../AGENTS.md) and [test maintenance](UT_MAINTENANCE.md) +- [Native operator builds](../csrc/README.md) +- NPU design: [Async FFN runtime](design/npu/async_ffn_runtime.md) and + [routed CAM operators](design/npu/cam_async_routed_ops.md) + +## Validate and troubleshoot + +- [Model serving and accuracy E2E](../tests/e2e/README.md), including + [multi-pod deployments](../tests/e2e/multi_pod/README.md) +- [GPU connector tests](../tests/gpu/README.md) +- [NPU runtime and operator tests](../tests/npu/README.md) +- [NPU troubleshooting](npu/TROUBLESHOOTING.md) +- Historical measurement: [Layered W13/W2 graph validation](reports/layered_w2_mdl_validation.md) + (environment, source revisions, and results retained with the report) diff --git a/docs/assets/WeChat.png b/docs/assets/WeChat.png index 8fd55cdb4..9805c873d 100644 Binary files a/docs/assets/WeChat.png and b/docs/assets/WeChat.png differ diff --git a/docs/configuration.md b/docs/configuration.md new file mode 100644 index 000000000..a073f0b0a --- /dev/null +++ b/docs/configuration.md @@ -0,0 +1,69 @@ +# AFD configuration + +Pass settings through vLLM `--additional-config` under the `afd` key. Omit +`afd` to disable the plugin. When `--worker-cls` is omitted, AFD selects the +Attention or FFN worker for CUDA or standard Ascend NPU. Explicit worker paths +are accepted for compatibility, but are not stable launch interfaces. + +Both roles require `VLLM_USE_V2_MODEL_RUNNER=1`. Attention uses native +ModelRunnerV2; FFN is connector-driven. Send requests only to the Attention API. + +## Common fields + +```json +{ + "afd": { + "role": "attention", + "connector": "P2pNcclAFDConnector", + "async_dp": false, + "host": "127.0.0.1", + "port": 1239, + "afd_process_group_timeout_s": 120, + "num_attention_ranks": 2, + "num_ffn_ranks": 1, + "compute_gate_on_attention": false, + "connector_extra_config": {} + } +} +``` + +| Field | Default | Meaning and constraint | +| --- | --- | --- | +| `role` | `attention` | `attention` or `ffn`. | +| `connector` | `P2pNcclAFDConnector` | One of the four connectors linked below. | +| `async_dp` | `false` | Must be `true` for CAM async and `false` for synchronous connectors. | +| `host` | `127.0.0.1` | Shared, reachable rendezvous host; placement depends on the connector. | +| `port` | `1239` | Rendezvous port in `1..65535`, separate from the HTTP service ports. | +| `afd_process_group_timeout_s` | `120` | Positive rendezvous timeout in seconds; also bounds the NPU native DBO startup barrier, not device-kernel execution. | +| `num_attention_ranks` | `1` | Positive total Attention worker count, including DP/TP ranks. | +| `num_ffn_ranks` | `1` | Positive total FFN worker count, including DP/TP ranks. | +| `compute_gate_on_attention` | `false` | Gate placement; allowed values depend on the connector and model. | +| `connector_extra_config` | `{}` | Connector-specific mapping; unknown keys fail in the selected connector parser. | + +The example describes Attention in a 2A1F deployment. Set `role` to `ffn` on +FFN; all ranks must agree on rendezvous and topology. Role rank is derived at +runtime and is not a configuration field. Increase the rendezvous timeout when +model loading or rank placement delays group initialization. + +## Compatibility aliases + +`afd_connector`, `afd_role`, `afd_host`, `afd_port`, and `async` map to +`connector`, `role`, `host`, `port`, and `async_dp`. Supplying an alias and its +canonical key together is rejected. Unknown top-level AFD keys are rejected; +put connector-specific settings under `connector_extra_config`. + +## Connector-specific settings + +Each guide owns its extra fields, topology rules and launch examples: + +- [NCCL P2P](gpu/NCCL_P2P_CONNECTOR_USER_GUIDE.md#configuration): synchronous CUDA; + extra config must be empty. +- [DeepEP M2N](gpu/DEEPEP_M2N_CONNECTOR_USER_GUIDE.md#configuration): synchronous CUDA; + Attention-side routing and FFN TP1. +- [CAMP2P](npu/CAM_P2P_CONNECTOR_USER_GUIDE.md#afd-configuration): synchronous NPU; + both common and connector-local gate settings must be `false`. +- [CAM async](npu/CAM_ASYNC_CONNECTOR_USER_GUIDE.md#afd-configuration): NPU; + requires `async_dp=true`, `compute_gate_on_attention=true`, and `dynamicQuant=1`. + +See [supported configurations](support.md) before combining graph, DBO, +parallelism, or quantization settings. diff --git a/docs/design/module/attention_runtime.md b/docs/design/module/attention_runtime.md index 31700b82d..57188eeda 100644 --- a/docs/design/module/attention_runtime.md +++ b/docs/design/module/attention_runtime.md @@ -13,6 +13,7 @@ primary_code_paths: - "afd_plugin/v1/worker/npu/attention_model_runner_v2.py" - "afd_plugin/v1/worker/npu/attention_worker.py" - "afd_plugin/v1/worker/npu/dspark_speculator.py" + - "afd_plugin/v1/worker/npu/dspark_graph.py" related_code_paths: - "afd_plugin/connectors/**" - "afd_plugin/model_executor/**" @@ -24,6 +25,7 @@ depends_on: - "execution_platforms.md" - "compatibility_and_patches.md" validation_paths: + - "tests/unit/v1/worker/npu/test_dspark_graph.py" - "tests/unit/v1/worker/test_attention_model_runner.py" - "tests/unit/v1/worker/test_model_runner_v2.py" - "tests/unit/v1/worker/test_npu_runtime.py" @@ -53,7 +55,7 @@ related_issues: - "#88" - "#105" - "#129" -last_reviewed: 2026-10-13 +last_reviewed: 2026-11-09 --- # Attention runtime @@ -160,21 +162,20 @@ Both platforms require V2. The paired deployment constraints are: | Constraint | CUDA V2 | Ascend V2 synchronous | Ascend V2 Async CAM | | --- | --- | --- | --- | -| Connector | Synchronous `P2pNcclAFDConnector` | Synchronous `CAMP2pAFDConnector` | `CAMAsyncAFDConnector`; `async=true` | -| Gate placement | Model-specific; Qwen requires `compute_gate_on_attention=false` | `compute_gate_on_attention=false` | `compute_gate_on_attention=true` | +| Connector | Synchronous `P2pNcclAFDConnector` or `DeepEPM2NAFDConnector` | Synchronous `CAMP2pAFDConnector` | `CAMAsyncAFDConnector`; `async_dp=true` | +| Gate placement | P2P: model-specific, Qwen requires `false`; DeepEP M2N: `true` | `compute_gate_on_attention=false` | `compute_gate_on_attention=true` | | Parallelism | PP=PCP=DCP=1; configured role ranks equal DP x TP; static EP enabled | Same | Same; `attn_ranks_per_dp` equals Attention TP | -| Excluded features | Elastic EP, EPLB, sequence-parallel MoE, and compile SP | Elastic EP and EPLB; model-specific SP restrictions | Elastic EP, EPLB, native draft models/MTP, finegrained TP, and shared-expert DP | +| Excluded features | Elastic EP, EPLB, sequence-parallel MoE, and compile SP | Elastic EP and EPLB; model-specific SP restrictions | Elastic EP, EPLB, native draft models/MTP except DSV4 DSpark, finegrained TP, and explicit shared-expert DP | | Sequence parallelism | Unsupported | Model-specific; synchronous DeepSeek-V2/V3 SP is rejected | Attention-local FlashComm1/SP supported; disabled on FFN; no PCP/DCP | | Native DBO | Exactly two microbatches, Attention DP1 or DP > 1 | Eager or graph; exactly two microbatches, Attention DP > 1 | Eager only; exactly two microbatches at Attention DP1 or DP > 1 | -| AFD async MoE ubatching | Not applicable | Not applicable | Legacy alias for native DBO; split selector is ignored | -| Graph execution | Eager or `FULL_DECODE_ONLY` | Eager, `FULL`, or `FULL_DECODE_ONLY`; plain MLA graph DBO requires `FULL_DECODE_ONLY` | Eager Attention; independent FFN also supports the narrow W4A8/W8A8 FULL contract | -| Model | Must resolve to a registered AFD architecture | Same | DeepSeek-V4 with layered W4A8/W8A8 FFN experts and `dynamicQuant=1` | +| Graph execution | Eager or `FULL_DECODE_ONLY` | Eager, `FULL`, or `FULL_DECODE_ONLY`; plain MLA graph DBO requires `FULL_DECODE_ONLY` | Eager Attention or DSV4 target native graph without DBO; native `FULL`/`FULL_DECODE_ONLY` resolution; Attention-local SP and native PD may remain enabled; DSpark draft may be eager or native Graph; independent FFN supports the narrow W4A8/W8A8 FULL contract | +| Model | Registered AFD architecture; DeepEP M2N is limited to Qwen3 MoE or DeepSeek-V2 | Registered AFD architecture | DeepSeek-V4 with layered W4A8/W8A8 FFN experts and `dynamicQuant=1` | Both roles validate the paired V2 topology, but only Attention uses the native V2 model runner. FFN remains connector-driven and uses its existing AFD runner surface. Qwen3.5/3.6 GDN DBO is unsupported; native hybrid/multimodal DBO validation applies. Qwen3 MoE DBO remains supported. See the -[runtime matrix](execution_platforms.md#tested-runtime-matrix) for tested +[runtime matrix](../../support.md#tested-runtime-matrix) for tested configurations and known shutdown limitations. The generic V2 E2E harness remains CUDA-only. @@ -247,12 +248,9 @@ plugin-owned model code. These fields describe the current implementation; the open metadata issues prevent the object shape and live connector reference from becoming a long-term API. -Ascend V1 creates pending AFD metadata while native Attention metadata is -built and installs it immediately before model forward. V2 temporarily installs a -provider at the native `ForwardContext` creation seam and attaches the same -sidecar to the context created by the upstream runner. The shared metadata -mixin supplies the ordinary single-stage and graph-control behavior. Ascend V1 -native ubatching still gives each child context stage-local metadata. +The runner temporarily installs a provider at the native `ForwardContext` +creation seam and attaches the sidecar to the upstream context. The shared +metadata mixin supplies the single-stage and graph-control behavior. Ascend V2 Async CAM uses `NPUUBatchRunner` to adapt native DBO child inputs to Ascend attention builders. Native slicing, execution threads and output merge @@ -325,25 +323,12 @@ hook. FlashComm1/SP is Attention-local and FFN topology validation is independent. See the CAM async user guide for the supported topology and environment matrix. -CAM async requires eager execution. Its -feature limits are recorded in the -[platform matrix](execution_platforms.md#tested-runtime-matrix). +CAM async requires eager Attention; FFN can use eager or FULL execution. +Feature limits are recorded in the +[platform matrix](../../support.md#tested-runtime-matrix). ## Ubatching and graph orchestration -On Ascend V1, when native ubatching is enabled, AFD accepts exactly two -ubatches. The Attention runner chooses or normalizes slices, creates -stage-local AFD/DP metadata, and installs the platform wrapper during model -load. The role-level contract is that control-plane side effects complete -before replayable model work enters formal graph capture. - -During graph warmup/capture, the Attention runner marks control payloads so the -FFN daemon can warm or capture the matching shape. A non-ubatched capture sends -the control payload explicitly before formal capture; an ubatched capture lets -the platform wrapper send the exact per-stage shape. CUDA Graph, ACL Graph, -stream, and wrapper implementation details are owned by -[execution platforms](execution_platforms.md). - CUDA V2 accepts exactly two native microbatches with Attention DP1 or DP > 1. The CUDA-only [V4 compatibility patch](../../../afd_plugin/compat/patches/gpu_deepseek_v4_dbo.py) owns aligned, persistent stage inputs, Indexer SM planning and pre-profile @@ -352,12 +337,14 @@ shutdown; Ascend and other models keep native preparation and profiling. Ascend V2 supports exactly two native microbatches: synchronous CAMP2P requires Attention DP > 1 and supports graph DBO, while Async CAM uses eager -Attention and also supports Attention DP1. Plain MLA graph DBO requires +Attention for DBO and also supports Attention DP1. Its separate target +DSV4 decode graph profile excludes DBO. Plain MLA graph DBO requires `FULL_DECODE_ONLY`. -For full-graph capture, the V2 runner -wraps native input preparation and publishes one warmup and one capture -payload per descriptor. CUDA two-stage preparation installs separate AFD +For full-graph capture, the V2 runner wraps native input preparation and +stages descriptor-matched AFD metadata. Connectors with a control plane also +publish one warmup and one capture payload per descriptor; Async CAM does not. +CUDA two-stage preparation installs separate AFD metadata on both native stage contexts and sends their complete layout before eager execution or FULL replay. Single-stage FULL replay publishes its padded control shape through an execute-scoped manager hook; two-stage replay reuses @@ -401,21 +388,8 @@ the matching platform and connector tests as well. ## Limitations and open issues -Current shared limits are the supported vLLM release and registered role-aware -model integrations. Historical Ascend V1 accepted exactly two ubatches. V2 -requires static EP and no PP/CP. CUDA V2 requires a synchronous control-plane -connector and supports the two-microbatch path above. Ascend V2 supports eager -native DBO with exactly two microbatches: synchronous CAMP2P requires Attention -DP > 1 and supports graph DBO; Async CAM supports eager Attention at DP1 or -DP > 1. Hardware evidence remains limited to the recorded model and -configuration cells; see the [0.32 record](https://github.com/vllm-project/afd-plugin/pull/507). -Platform/connector limits are intentionally centralized in -[execution platforms](execution_platforms.md#tested-runtime-matrix). - -Async CAM DBO has a Lite two-stage hardware accuracy result; broader coverage needs -DP1/DP > 1, plain TP/model SP, uneven -prefill lengths, and requests crossing a child boundary, with live two-child -execution evidence and baseline output comparison. +Platform/connector limits and hardware evidence are centralized in +[support matrix](../../support.md#tested-runtime-matrix). This document does not decide the runtime refactor, connector metadata ownership, or transfer metadata/state shape. See diff --git a/docs/design/module/compatibility_and_patches.md b/docs/design/module/compatibility_and_patches.md index 815f23948..86a936d58 100644 --- a/docs/design/module/compatibility_and_patches.md +++ b/docs/design/module/compatibility_and_patches.md @@ -10,6 +10,7 @@ primary_code_paths: - "afd_plugin/compat/vllm.py" - "afd_plugin/compat/npu/__init__.py" - "afd_plugin/compat/npu/feature_validation.py" + - "afd_plugin/compat/npu/routing_ids.py" - "afd_plugin/compat/npu/runtime.py" - "afd_plugin/compat/npu/runtime_config.py" - "afd_plugin/compat/patches/**/*.py" @@ -40,7 +41,7 @@ verified_platform_refs: related_issues: - "#86" - "#129" -last_reviewed: 2026-10-13 +last_reviewed: 2026-11-09 --- # Compatibility and patches @@ -67,7 +68,7 @@ Direct strict calls raise for a missing or different vLLM; plugin registration calls the check with `strict=False`, so it warns and continues. This warning policy does not make another vLLM release supported. Shared EngineCore and coordinator patches use the 0.32 contract. See the -[runtime matrix](execution_platforms.md#tested-runtime-matrix) for hardware +[runtime matrix](../../support.md#tested-runtime-matrix) for hardware validation scope. The NPU implementation targets vLLM `0.32.0` (`5df63a72`) and @@ -85,6 +86,7 @@ compatibility evidence rather than a released package or container tag. | Core vLLM patches | [`compat/patches/`](../../../afd_plugin/compat/patches) | [`tests/unit/compat/patches/`](../../../tests/unit/compat/patches) | | CUDA benchmark routing | [`model_executor/routing_simulator.py`](../../../afd_plugin/model_executor/routing_simulator.py) | [`test_routing_simulator.py`](../../../tests/unit/model_executor/test_routing_simulator.py) | | NPU adapters | [`compat/npu/`](../../../afd_plugin/compat/npu) | [`test_runtime.py`](../../../tests/unit/compat/test_runtime.py), [`test_ascend_ops.py`](../../../tests/unit/compat/test_ascend_ops.py) | +| NPU routing-ID balancing | [`compat/npu/routing_ids.py`](../../../afd_plugin/compat/npu/routing_ids.py) | [`test_routing_ids.py`](../../../tests/unit/compat/npu/test_routing_ids.py), using the shared [`routing_ids_op.py` fixture](../../../tests/unit/fakes/routing_ids_op.py) | | NPU patch paths | [`compat/patches/npu/`](../../../afd_plugin/compat/patches/npu) | [`test_npu_runtime.py`](../../../tests/unit/v1/worker/test_npu_runtime.py) | | CUDA native V2 DBO preparation | [`gpu_mrv2_ubatching.py`](../../../afd_plugin/compat/patches/gpu_mrv2_ubatching.py) | [`test_gpu_mrv2_ubatching.py`](../../../tests/unit/compat/patches/test_gpu_mrv2_ubatching.py) | @@ -156,7 +158,8 @@ These modules adapt upstream behavior without replacing a global symbol: | [`model_executor/routing_simulator.py`](../../../afd_plugin/model_executor/routing_simulator.py) | Registers the opt-in `afd_balanced` strategy through vLLM's public `RoutingSimulator` interface. Select it with `VLLM_MOE_ROUTING_SIMULATION_STRATEGY=afd_balanced`; `AFD_BENCHMARK_FORCE_LB_TOPN_PER_RANK` optionally limits each rank's expert pool. It provides deterministic source-rank phases and normalized weights for controlled CUDA performance tests without replacing router internals. EPLB remains outside the validated scope. | | [`compat/npu/runtime_config.py`](../../../afd_plugin/compat/npu/runtime_config.py) | Preserves model SP while applying the non-SP backend fallback; normalizes Async CAM native ubatching to Attention DBO, checks thresholds, and derives model/PP tensor layout and active ubatch count. | | [`compat/npu/feature_validation.py`](../../../afd_plugin/compat/npu/feature_validation.py) | Parses connector-owned typed extra information through the factory and fails before execution for unsupported NPU connector, quantization, graph, DBO, gate, or async MoE combinations. | -| [`compat/npu/forward_context.py`](../../../afd_plugin/compat/npu/forward_context.py) | Enters the pinned Ascend forward context for connector-driven FFN compute and installs AFD metadata in `additional_kwargs`. V1 uses `set_ascend_forward_context`; a V2-paired FFN uses native `set_forward_context` plus the Ascend MRV2 profile override so platform state follows the V2 proxy layout. | +| [`compat/npu/routing_ids.py`](../../../afd_plugin/compat/npu/routing_ids.py) | Registers the opaque `afd_balance_routing_ids_` custom op. It reads Ascend's live profile flag inside the op and balances logical expert IDs in place during profiling or when forced. NPU worker construction registers it before the first Attention-side gate; standalone scripts must import the module after device initialization and before using the op. | +| [`compat/npu/forward_context.py`](../../../afd_plugin/compat/npu/forward_context.py) | Uses native `set_forward_context` and the Ascend MRV2 profile override for connector-driven FFN compute, installing AFD metadata and the model instance in `additional_kwargs`. | | [`compat/npu/ops.py`](../../../afd_plugin/compat/npu/ops.py) | Discovers plugin-owned CAMP2P and routed-only async CAM operators with explicit missing-runtime errors; operator build/runtime ownership is documented in [execution platforms](execution_platforms.md). | Several scoped upstream adaptations live with their semantic owners rather diff --git a/docs/design/module/connector_contracts.md b/docs/design/module/connector_contracts.md index 407bf4673..8403e52b5 100644 --- a/docs/design/module/connector_contracts.md +++ b/docs/design/module/connector_contracts.md @@ -42,7 +42,7 @@ related_issues: - "#105" - "#107" - "#129" -last_reviewed: 2026-10-30 +last_reviewed: 2026-11-09 --- # Connector contracts @@ -67,8 +67,8 @@ depend on role worker implementations. | Payloads and control-plane codec | [`metadata.py`](../../../afd_plugin/connectors/metadata.py) | Connector unit tests and [`test_forward_context.py`](../../../tests/unit/model_executor/models/test_forward_context.py) | | CUDA P2P | [`gpu/p2p.py`](../../../afd_plugin/connectors/gpu/p2p.py), [`topology.py`](../../../afd_plugin/distributed/topology.py) | [`test_p2p_connector.py`](../../../tests/unit/connectors/test_p2p_connector.py), [DeepSeek-V2-Lite E2E](../../../tests/e2e/models/deepseek_v2_lite/test_deepseek_v2_lite.py) | | CUDA DeepEP M2N | [`gpu/deepep_m2n/`](../../../afd_plugin/connectors/gpu/deepep_m2n/), `AFDRoutedExpertsConnector` in [`base.py`](../../../afd_plugin/connectors/base.py) | [`test_deepep_m2n_connector.py`](../../../tests/unit/connectors/test_deepep_m2n_connector.py), [GPU device tests](../../../tests/gpu/README.md), [Qwen3 MoE E2E](../../../tests/e2e/models/qwen3_moe/test_qwen3_moe.py), [DeepSeek-V2-Lite E2E](../../../tests/e2e/models/deepseek_v2_lite/test_deepseek_v2_lite.py) (`test_deepseek_v2_lite_deepep_m2n`) | -| Ascend CAMP2P | [`npu/camp2p.py`](../../../afd_plugin/connectors/npu/camp2p.py) | [`test_camp2p_connector.py`](../../../tests/unit/connectors/test_camp2p_connector.py), [DeepSeek-V2-Lite E2E](../../../tests/e2e/models/deepseek_v2_lite/test_deepseek_v2_lite.py) | -| Ascend CAM async | [`npu/async_cam.py`](../../../afd_plugin/connectors/npu/async_cam.py) | [Connector unit tests](../../../tests/unit/connectors/test_async_cam_connector.py), [fused routing device test](../../../tests/npu/test_async_cam_connector.py), [`test_async_cam_npu.py`](../../../tests/e2e/models/deepseek_v4_flash/test_async_cam_npu.py) | +| Ascend CAMP2P | [`npu/camp2p/connector.py`](../../../afd_plugin/connectors/npu/camp2p/connector.py) | [`test_camp2p_connector.py`](../../../tests/unit/connectors/test_camp2p_connector.py), [`test_camp2p_token_ids.py`](../../../tests/unit/connectors/test_camp2p_token_ids.py), [`test_external_routing.py`](../../../tests/unit/connectors/test_external_routing.py), [DeepSeek-V2-Lite E2E](../../../tests/e2e/models/deepseek_v2_lite/test_deepseek_v2_lite.py) | +| Ascend CAM async | [`npu/async_cam.py`](../../../afd_plugin/connectors/npu/async_cam.py), [`Push/Pull sessions`](../../../afd_plugin/connectors/npu/async_cam_runtime/async_cam_session.py) | [Connector unit tests](../../../tests/unit/connectors/test_async_cam_connector.py), [CAM host contracts](../../../tests/unit/connectors/npu/test_async_cam.py), [fused routing device test](../../../tests/npu/test_async_cam_connector.py), [`test_async_cam_npu.py`](../../../tests/e2e/models/deepseek_v4_flash/test_async_cam_npu.py) | | Process-group construction | [`afd_process_group.py`](../../../afd_plugin/distributed/afd_process_group.py) | Connector initialization tests plus platform E2E paths | ## Factory and construction @@ -95,15 +95,15 @@ which lets feature validation use the same connector-owned schema. ## Connector-owned configuration -`additional_config["afd"]["connector_extra_config"]` is an envelope rather -than an `AFDConfig` field. Unknown fields fail in the selected connector parser. +`AFDConfig.connector_extra_config` stores the connector-specific mapping. +Unknown fields fail in the selected connector parser. | Connector | Accepted connector-specific fields | | --- | --- | | `P2pNcclAFDConnector` | None; the mapping must be empty. | | `DeepEPM2NAFDConnector` | `payload_format`, `expert_backend`, `capacity_tokens_per_rank`, `num_lanes`, `bucket_granularity`, `cpu_sync_min_tokens`, `num_sms`, `num_qps`, `expert_num_sms`, `cpu_timeout_s`, `gpu_timeout_s`, and a nested `network` mapping (`gin`, `hybrid`, `rdma_traffic_class`). All are optional; see the [DeepEP M2N guide](../../gpu/DEEPEP_M2N_CONNECTOR_USER_GUIDE.md#connector_extra_config) for defaults. | -| `CAMP2pAFDConnector` | `core_num`, optional `attn_core_num` / `ffn_core_num`, `compute_gate_on_attention`, `quant_mode`, and optional positive `hccl_buffer_size`. Core counts must be positive; the current runtime rejects gate-on-Attention and any nonzero quantization mode. | -| `CAMAsyncAFDConnector` | `dynamicQuant`, `attn_ranks_per_dp`, optional positive `hccl_buffer_size`, and FFN two-stream options. Attention microbatching uses native DBO configuration. | +| `CAMP2pAFDConnector` | `core_num`, optional `attn_core_num` / `ffn_core_num`, `compute_gate_on_attention`, `quant_mode`, and optional positive `hccl_buffer_size`. Core counts must be positive. Gate placement follows the common `afd.compute_gate_on_attention` field; the connector-local copy is kept for compatibility and does not select the mode. `true` requires DeepSeek-V4 at TP=1 without DBO or ubatching. Any nonzero quantization mode is rejected today. | +| `CAMAsyncAFDConnector` | `dynamicQuant`, `ffn_mega_moe`, `attn_ranks_per_dp`, and optional positive `hccl_buffer_size`. Attention microbatching uses native DBO configuration. | The common `afd_process_group_timeout_s` defaults to 120 seconds and controls AFD rendezvous. `hccl_buffer_size` sizes connector-owned HCCL domains in MB; @@ -113,10 +113,12 @@ environment rather than querying the allocated domain. The common `compute_gate_on_attention` field remains on `AFDConfig` and is the model-routing selector. CUDA P2P supports both values; DeepEP M2N requires -`true`. CAMP2P also parses a -connector-local field with that name for its operator contract, but the current -synchronous NPU runtime requires both common and connector-local values to be -`false`. CAM async requires the common field to be `true`. +`true`. CAMP2P also parses a connector-local field with that name, retained for +backwards compatibility with older configurations and ignored for routing: the +synchronous NPU runtime selects Attention-side routing from the common field, so +deployments leave the connector-local copy unset. The supported synchronous +shape is `true` on both roles for DeepSeek-V4 at TP=1 with no DBO or ubatching. +CAM async requires the common field to be `true`. ## Current connector modes @@ -207,7 +209,7 @@ The current plugin-owned objects are: | `AFDTransferMetadata` | Layer, stage, non-negative split lengths, and total-token validation; backend state is paired separately in `AFDTransferContext`. A split length may be `0`: with an expert-routed transport (DeepEP M2N) an FFN rank can receive no rows for a layer while still taking part in the collective. The list itself must not be empty. | **Draft** under [#88](https://github.com/JiusiServe/afd-plugin/issues/88). | | `AFDTransferContext` | Pairs common `metadata` with the backend `states` retained from receive through the matching send. | **Draft** under [#88](https://github.com/JiusiServe/afd-plugin/issues/88) and [#105](https://github.com/JiusiServe/afd-plugin/issues/105). | | `AFDTransferState` | Empty base for backend-specific state such as CAMP2P handles or async CAM state. | **Draft** under [#88](https://github.com/JiusiServe/afd-plugin/issues/88). | -| `AFDA2FTransferPayload` | Attention-to-FFN hidden states, common metadata, and optional router logits or token-aligned input IDs. | **Draft** and scheduled for state splitting in [#105](https://github.com/JiusiServe/afd-plugin/issues/105). | +| `AFDA2FTransferPayload` | Attention-to-FFN hidden states, common metadata, and optional router logits, token-aligned input IDs, or the routed top-k (`topk_ids` with `topk_weights`). | **Draft** and scheduled for state splitting in [#105](https://github.com/JiusiServe/afd-plugin/issues/105). | | `AFDF2ATransferPayload` | Structured routed/shared FFN outputs where the backend needs both. | **Draft** with the transfer-state work. | | `LocalExpertWeights` | One FFN rank's experts in vLLM `FusedMoE` layout (`w13` gate-first, `w2`), BF16 or block-128 FP8 with both block scales. | **Draft**; used only by routed-experts connectors. | | `FFNLocalExperts` | One MoE layer's global expert ids and their `LocalExpertWeights`, as vLLM loaded them. | **Draft**; used only by routed-experts connectors. | @@ -233,7 +235,12 @@ state. P2P can send optional router logits and one-dimensional, token-aligned the A2E operator's ids channel, which the receiving FFN rank selects with `recv_input_ids`; both roles derive that mode from the model's `afd_requires_input_ids` declaration, because the operator only writes the ids -slot in the mode the sender chose. FFN requests only the fields its +slot in the mode the sender chose. The same slot carries the one payload the +layer speaks: either token-aligned input IDs or the routed top-k, never both, and +the connector rejects a request for both. External routing keeps the expert IDs in +that ids slot and puts the gate weights in the paired scales slot; the receiving +rank selects it with `recv_routing`, derived from the model's +`afd_requires_attention_routing` declaration. FFN requests only the fields its model declares, concatenates them with the same per-peer token order as hidden states, and returns them in `AFDA2FTransferPayload`. A CUDA P2P FFN rank posts the receives for all of its Attention peers, each peer's tensors in wire order, @@ -312,10 +319,40 @@ 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. | +`CAMAsyncAFDConnector` owns configuration, topology, the HCCL group, startup, +payload validation and final cleanup decisions. It selects one session at +construction: `PushCAMSession` or `PullCAMSession`, both in +`async_cam_runtime/async_cam_session.py`. Push keeps a FIFO per Attention +stage and separate FFN receive/send calls. Pull owns one in-flight window across +stages, native publish/combine, status/idle/source state and quiet/drain/stopped. +Mega's fused FFN operation owns its receive/compute/send transaction; the public +FFN receive/send methods reject Pull use. + +The four public transfer methods still enter through the connector. The existing +`AFDAsyncTransferState` definition stays at its original import path; +`comm_args` and `comm_id` expose the selected session's live resources through read-only +properties for the current Mega compute path. `has_pending_attention` provides +the diagnostic pending-state read without exposing the queue. + +Pull Session receives Startup, topology and required dimensions/capacity +explicitly, without retaining the connector. Sessions allocate their buffers +after Store negotiation and HCCL setup; initialization failures reach Startup +before the connector is marked ready. Push records routing after a successful +launch; Pull retains window ownership before launch. Combine failures retain +routing/window references. DBO defers the sticky status check to the Attention +step boundary. Graph capture, device completion and service threads remain with +the runner/worker. + +`FFNExecutionSession` in `base.py` defines the lifecycle hooks and provides +local-stop/no-drain defaults in one concrete class. Push/Pull sessions inherit +those defaults and override their protocol-specific behavior. The base module +keeps torch and transfer metadata imports under type checking so session imports +remain safe before worker device bootstrap. + The common process-group helper creates a plugin-owned group while temporarily switching PyTorch's default group so vLLM parallel-state helpers initialize against it. It currently calls private PyTorch/vLLM symbols and is therefore an diff --git a/docs/design/module/e2e_testing.md b/docs/design/module/e2e_testing.md index e6ff87788..e3be92e16 100644 --- a/docs/design/module/e2e_testing.md +++ b/docs/design/module/e2e_testing.md @@ -5,6 +5,10 @@ status: normative owners: - "@yujuancao07" primary_code_paths: + - "tests/e2e/case.py" + - "tests/e2e/presets.py" + - "tests/e2e/deployment.py" + - "tests/e2e/checks/**" - "tests/e2e/runner.py" - "tests/e2e/cancellation.py" - "tests/e2e/multi_pod/**" @@ -51,7 +55,7 @@ verified_platform_refs: - "CUDA Qwen3 MoE" - "CUDA Qwen3.6 MoE" related_issues: [] -last_reviewed: 2026-10-13 +last_reviewed: 2026-11-10 --- # E2E testing @@ -67,13 +71,120 @@ operator tests, benchmarks, and performance tests are outside this scope. ## Structure +Single-deployment pytest entries use model-local `cases.py` builders for all +48 existing definitions: 38 CUDA and ten NPU. Stable node IDs and the 15 +existing CI selections stay unchanged. Collection reads IDs only. Builders +resolve checkpoint-dependent profile settings inside the runner subprocess; +the common `case_main` receives each model builder explicitly. + +`case.py` defines the common types; `presets.py` provides reusable synchronous +eager/decode-graph deployment construction across models and backends. V2 on +GPU/NPU, Qwen3 MoE and Qwen3.6 compose these presets with model-owned DP/TP +layouts, connector, capture size and DBO settings. V2 retains its explicit +comparison sequence-budget override. Presets select neither model IDs nor +scenarios and own no workload, threshold, evidence or cleanup requirements. +DSV4's specialized deployment configuration remains model-local. Shared +configuration does not generate a common scenario matrix across all models. + +`deployment.py` owns topology/resource binding, commands, connector launch +order, readiness, pre/post-evaluation liveness and cleanup. `runner.run_case` +executes declared phases through the serializable Task interface and writes +`run.json` after cleanup. GSM8K is shared; DSV4's concurrent answer/plumbing +semantics remain model-local. Its two concurrent rounds and GSM8K are phases, +including their existing sample, repetition and concurrency requirements. + +`checks/mrv2.py` prepares one worker probe and validates graph, DBO and Hash +inputs. The runner drives one reset/start/end per phase; checks consume a +shared immutable snapshot. Hash freshness compares declared windows. Every +required rank must report even in observation-only windows. Cases without +execution checks do not install the probe. + +The shared capabilities needed by existing DSV4 profiles are explicit: +per-role arguments/additional configuration and native DBO, connector extra +configuration, a pinned Attention rendezvous address, exact compilation +configuration, and `cleanup_profile`. Sync connectors start FFN then Attention; +async CAM starts Attention then FFN. The NPU async profile alone keeps its +scoped FFN process-ownership cleanup exception; other profiles are strict. +`Deployment.ffn_service` replaces the FFN role's `vllm serve` with the +standalone FFN service under `torchrun` on one host. Readiness adds every +rank's `/ready`; a passing case ends with `/stop` on every rank, then +Attention, and requires a clean launcher exit and released resources per rank. +Its torchrun workers sit outside the launcher's process group, so cleanup also +sweeps FFN processes by run environment. +`Deployment.dbo` supplies a default, and an explicit `Role.dbo` overrides it. +`Role.dbo=None` inherits the default; `DboConfig()` keeps native thresholds. +`Deployment.dbo_for_role()` resolves this rule for commands, probe validation +and role logging. A DBO check requires it on every named role, so Attention-only +native DBO does not imply FFN DBO. +NPU async CAM declares 60-second worker, 90-second service and 120-second runner +shutdown budgets. The worker budget is an environment default, independent of +the service CLI timeout. Existing worker environment overrides are retained and +recorded in the effective configuration. Other profiles keep their declared +shutdown budgets and do not gain this worker default. + +The ten existing GPU multi-pod entries use model-local distributed Cases. +`multi_pod` derives local RoleSlots and coordinates addresses, launch barriers, +verdict and departure; `deployment.py` builds commands and owns local processes. +The original distributed workload, liveness and cleanup checks are retained; +no cross-pod MRV2/Hash evidence is added. Unsupported execution checks fail +before launch rather than being silently dropped. Every pod records the actual +evaluator configuration/results, its local outcome and cleanup. Qualification +requires every pod's acceptance entry to pass. + +The two NPU multi-pod entries and dedicated FFN/DSpark callers retain their +legacy path. Remaining CLI scenario configuration is NPU-only; GPU-specific +tables and evaluation/Hash branches have been removed. No cases are added or +removed. Implementation/collection does not establish device qualification: +report host, same-host distributed and cross-node/device execution separately. + +On migrated Case paths (all single-deployment entries and ten GPU multi-pod +entries), an existing-capability combination changes model declarations and +IDs only. +Changing Case fields, task/check behavior or deployment protocols is a shared +capability extension. A new connector must define launch order in deployment; +Cases do not override that order independently. Configuration conflicts in +deployment-owned flags fail before service startup. +Changing a preset affects its existing consumers and requires configuration +and command comparison plus affected integration checks. A new combination +using existing presets must leave the presets unchanged. +Retained NPU multi-pod and dedicated executors are outside this extension +contract until migrated. Case execution uses deployment-owned cleanup; +retained legacy execution uses its existing cleanup, never both for one run. + +Each Case run records canonical and effective settings, documented overrides, +resource bindings, role commands, phase results, evidence and cleanup outcome. +Configuration eligibility and execution success are independent. Weakened +sample counts or accuracy thresholds produce a debug run and cannot qualify a +fixed scenario. Permitted larger workloads preserve the model's existing +threshold rule. The model-module CLI can execute debug runs; the stable pytest +acceptance node additionally requires `acceptance_passed=true`. + +The GSM8K schema records the existing small-workload threshold separately: +V2 comparison debug runs still floor at 24 samples and use the declared smoke +threshold, while their canonical acceptance remains 128 samples at 0.27. +The host CLI integration covers this policy and rejects its qualification; +this field preserves an existing override rule rather than adding coverage. + +For migrated entries, the operator keeps CPU harness integration results, +actual device results, skipped and unexecuted entries separate. `--describe`, +collection and command comparisons do not establish device acceptance. A +retained old entry exits only after its workload, evidence, cleanup and active +callers have a validated replacement. + ```text tests/e2e/ -├── models//test_*.py # Test entries and case lists -├── accuracy/.py # Accuracy tool and result parsing -├── environment.py # Backend and role launch environment -├── runner.py # Service startup, evaluation, and cleanup -└── process_utils.py # Process-group termination and reaping +├── models//test_*.py # Stable pytest entries; launch a child +├── models//cases.py # Model-local declarations for migrated entries +├── case.py # Immutable scenario and deployment types +├── presets.py # Shared execution configuration; no scenario selection +├── runner.py # Case phase execution; remaining legacy entries +├── deployment.py # Case commands, service startup and cleanup +├── checks/mrv2.py # Host probe preparation and phase windows +├── mrv2_evidence.py # Worker-side probe +├── accuracy/.py # Shared workload and result validation +├── environment.py # Resource bindings and CLI environment +├── multi_pod/ # Peer coordination and placement; NPU legacy entry +└── process_utils.py # Process-group termination and reaping ``` New model files do not copy service startup, lm-eval, signal handling, or @@ -102,8 +213,8 @@ cleanup. Production code does not depend on the E2E harness. - `E2E-INV-007` — CI **MUST NOT** lower the sample count, accuracy threshold, or required case set. It **MUST** select cases by pytest node ID. - `E2E-INV-008` — A shared case **MUST** pass without skips on every platform - where CI selects it. New behavior **MUST** have focused unit tests and real - hardware evidence. + where CI selects it. New harness behavior **MUST** have affected integration + and real hardware evidence. ## Required coverage @@ -124,11 +235,15 @@ CUDA ready/merge CI selects seven DeepSeek-V2-Lite cases: the native baseline and the six DP2/TP2/DBO V2 cases. The V2 DP2 cases replace the three duplicate GPU 2A2F cases; the eager/DBO comparison retains 128 samples, concurrency 12 and the 0.27 accuracy threshold. The backend-neutral 2A2F IDs remain available -for NPU. The GPU module collects 13 cases, including 1A1F and 2A1F coverage -and the local-only `afd-v2-graph-dbo-tp2`; the NPU module collects seven. -Across the three GPU model suites, there are 20 cases: DeepSeek 13, Qwen3 -MoE four, and Qwen3.6 MoE three. The local-only TP2 Graph + DBO case is -outside the seven-case ready/merge CI gate. +for NPU. The `test_deepseek_v2_lite` entry collects 13 GPU cases, including +1A1F, 2A1F and the local-only `afd-v2-graph-dbo-tp2`; its NPU selection +collects seven. The P2P entries across DeepSeek, Qwen3 MoE and Qwen3.6 MoE +contain 20 CUDA cases (13 + 4 + 3). Separate DeepEP M2N entries add five GPU-only +cases each for DeepSeek and Qwen3 MoE, and a Qwen3 MoE standalone FFN service +entry adds four. Unfiltered DeepSeek module collection +includes its five DeepEP IDs even on NPU; NPU runs select the P2P entry or +explicit node IDs. The local-only TP2 Graph + DBO case is outside the seven-case +ready/merge CI gate. GPU cases gate accuracy, service liveness and cleanup. AFD GPU graph/DBO cases additionally use a test-only worker extension to collect bounded per-rank diff --git a/docs/design/module/execution_platforms.md b/docs/design/module/execution_platforms.md index 9ae4ed2e5..1992482ca 100644 --- a/docs/design/module/execution_platforms.md +++ b/docs/design/module/execution_platforms.md @@ -24,6 +24,7 @@ related_code_paths: depends_on: - "plugin_boundary.md" validation_paths: + - "tests/npu/check_pull_layout.cpp" - "tests/unit/compat/test_profiler.py" - "tests/unit/compat/test_ascend_ops.py" - "tests/unit/compat/npu/test_profiler.py" @@ -43,7 +44,7 @@ upstream_refs: - "GPU: vLLM 0.32.0 5df63a7203cdb9328d824a27b682e77c7b9928be" - "NPU target: vLLM 0.32.0 5df63a7203cdb9328d824a27b682e77c7b9928be" - "vLLM vllm.compilation and vllm.v1.worker.gpu V2 graph/ubatching APIs" - - "vLLM-Ascend ACL graph, forward-context, and model-runner V1/V2 APIs in the tested environment" + - "vLLM-Ascend ACL graph, forward-context, and ModelRunnerV2 APIs in the tested environment" - "PyTorch CUDA, torch_npu, CMake, and Ascend CANN build interfaces used by the repository" verified_platform_refs: - "CUDA eager, graph, DBO, and ModelRunnerV2 E2E paths; see runtime matrix for tested scope" @@ -51,7 +52,7 @@ verified_platform_refs: related_issues: - "#86" - "#129" -last_reviewed: 2026-10-13 +last_reviewed: 2026-11-09 --- # Execution platforms @@ -60,8 +61,9 @@ last_reviewed: 2026-10-13 This document is the primary design for CUDA and NPU mechanisms: runtime class strategy, device graphs, native DBO/ubatching, streams, forward-context -adaptation, profilers, native operators, packaging, and the tested runtime -matrix. [Attention](attention_runtime.md) and [FFN](ffn_runtime.md) retain role +adaptation, profilers, native operators, and packaging. User-facing execution +limits and validation evidence live in the [support reference](../../support.md). +[Attention](attention_runtime.md) and [FFN](ffn_runtime.md) retain role lifecycle orchestration; [connector contracts](connector_contracts.md) retain transport and topology semantics. @@ -94,6 +96,20 @@ the top-level package, common configuration/validation, version checks, and the graph policy helper; importing a CUDA or Ascend runtime module requires the corresponding runtime stack. +## Backend execution session + +Connectors expose `attention_step()` and a backend-owned `FFNExecutionSession`. +The Async FFN runner and worker use its service predicate, first-step hook, +completion hook, drain callback, receiver completion and pre-release validation; +they do not inspect Pull state. Push uses the no-op default session. Pull owns +idle enablement, sticky status checks and the `quiet → drain → stopped` handshake +in `connectors/npu/async_cam_runtime/async_cam_session.py`. Shutdown validates the +session before releasing graph, profiler, model or communication resources; +failures retain those resources. Attention batches the same device-status fence +at step exit. +Generic graph startup negotiation remains backend-owned and is imported without +loading torch, vLLM or NPU modules through the connector package namespaces. + ## Platform mechanism map | Mechanism | CUDA owner | NPU owner | @@ -130,14 +146,13 @@ flowchart TB ### Worker and device setup -Both CUDA workers delegate device and distributed setup to native -`Worker.init_device()`. During that synchronous construction window, they -temporarily replace the native V2 runner symbol with the AFD runner -and restore it in `finally`. Both roles require ModelRunnerV2. Attention -selects `AFDAttentionModelRunnerV2`; FFN selects the connector-driven -`GPUFFNModelRunner`. `torch.accelerator.empty_cache()` releases -startup allocations. Attention retains the matching upstream request, -KV-cache, sampling, and output behavior; FFN remains connector-driven. +Both CUDA workers validate the AFD configuration and select the role's model +before delegating device and distributed setup to native `Worker.init_device()`. +Their `_make_model_runner()` overrides construct `AFDAttentionModelRunnerV2` +or the connector-driven FFN runner (`GPUFFNModelRunner` for P2P, +`GPUFFNRoutedExpertsModelRunner` for DeepEP M2N). Both roles require +ModelRunnerV2. Attention retains upstream request, KV-cache, sampling, and +output ownership. ### CUDA Graph policy @@ -183,9 +198,7 @@ Scoped factory and dispatch adapters enable the native path at Attention DP1 as well as DP > 1. AFD accepts exactly two native ubatches. Qwen3.5/3.6 GDN DBO is unsupported; native hybrid/multimodal DBO validation -applies. Qwen3 MoE DBO remains supported. Earlier Qwen GDN Graph + DBO and -six-cell recurrent-state/CUDA Graph parity results describe the removed adapter -and are historical evidence only. See the [runtime matrix](#tested-runtime-matrix) +applies. Qwen3 MoE DBO remains supported. See the [runtime matrix](../../support.md#tested-runtime-matrix) for tested configurations and known DeepSeek-V2-Lite DBO shutdown failures. ### CUDA profiling @@ -206,12 +219,10 @@ upstream construction. During device initialization they: 1. validate Ascend-specific feature combinations; 2. apply the non-sequence-parallel all-to-all backend correction when needed, including legacy explicit-worker launches; -3. validate the synchronous CAMP2P or eager Async CAM ModelRunnerV2 subset - when V2 is selected; +3. require ModelRunnerV2 and validate the connector's role-specific execution mode; 4. call `NPUWorker._init_device()`; 5. initialize the vLLM workspace manager for one or two ubatches; -6. construct the matching V1 or V2 Attention runner, or the connector-driven - FFN runner. +6. construct `AFDNPUAttentionModelRunnerV2` or the connector-driven FFN runner. The all-to-all correction selects `flashinfer_all2allv` when sequence parallelism is disabled. Automatic worker selection receives the matching @@ -226,11 +237,10 @@ in `ForwardContext.additional_kwargs`. FFN uses connector-driven MoE compute, including token counts and ACL graph runtime mode when applicable. -When ModelRunnerV2 is selected, `ascend_forward_context()` uses native +`ascend_forward_context()` uses native `set_forward_context()` plus vLLM-Ascend's MRV2 profile override. This places Ascend-specific state, `model_instance`, and AFD metadata in -`additional_kwargs`, matching the proxy layout read by the V2 runtime. V1 -continues to use `set_ascend_forward_context()`. +`additional_kwargs`, matching the proxy layout read by the V2 runtime. For native ubatching, `create_ascend_forward_context()` creates one context per stage with stage attention/DP metadata, batch descriptor, and graph mode. @@ -242,9 +252,7 @@ reassembled to match the upstream Ascend layout. `NPUUBatchRunner` extends native MRV2 `UBatchRunner`. Native code owns child slicing, threads, handoff, capture/replay and output merge. AFD supplies Ascend child metadata builders, NPU stream/context scopes, model-specific transport -coordinates and nested auxiliary output handling. The old NPU MRV1 runner, -wrapper, custom registry and request-boundary splitting implementation are -removed; there is no execution fallback. +coordinates and nested auxiliary output handling. `v1/worker/dbo.py` registers the model-side yield operation and dispatches to the platform DBO implementation. Eager Async Attention uses native DBO at every @@ -277,11 +285,12 @@ groups and also keeps Attention TP peers aligned. ### ACL Graph and NPU Graph NPU Attention follows the upstream ACL graph dispatcher while adding AFD -metadata and control-plane coordination. Native MRV2 owns single/two-stage -graph descriptors, warmup/capture, persistent outputs and replay. AFD retains -formal per-stage MLA registries by the complete native descriptor, including -LoRA specialization, and publishes one matching FFN control payload before -each warmup, capture or live replay. +metadata and, for synchronous CAMP2P, control-plane coordination. Native MRV2 +owns single/two-stage graph descriptors, warmup/capture, persistent outputs and +replay. AFD retains formal per-stage MLA registries by the complete native descriptor, including +LoRA specialization. Synchronous CAMP2P publishes one matching FFN control +payload before each warmup, capture or live replay; Async CAM stages local +AFD metadata without sending DP control messages. For MLA DBO full graphs, each stage records its own upstream `GraphParams`. `merge_mla_graph_params()` validates identical layer order and record counts, @@ -297,12 +306,16 @@ register no FIA workspace, so those models, DeepSeek V4 among them, take the generic two-stage path with no MLA registries. With synchronous CAMP2P, the NPU V2 runner supports eager, `FULL`, and -`FULL_DECODE_ONLY`; Async CAM is eager-only. Like CUDA V2, -it publishes descriptor-matched warmup/capture control outside formal graph -capture and installs an instance-scoped pre-replay hook because native full -replay creates no `ForwardContext`. V2 does not use `AscendUBatchWrapper` and -uses native `NPUUBatchRunner` preparation, execution and merge for eager and -graph DBO. Plain MLA graph DBO requires `FULL_DECODE_ONLY`; SFA/DSA +`FULL_DECODE_ONLY`; Async CAM additionally admits DSV4 target native +`FULL` or `FULL_DECODE_ONLY`, with native mode resolution and no DBO. DSpark +draft may be eager or native Graph. Attention-local FlashComm1/SP and native PD +can remain enabled. Device qualification uses compilation mode 0. +Synchronous CAMP2P publishes descriptor-matched warmup/capture control outside +formal graph capture. Async CAM has no control plane and only stages local +AFD metadata. Both install an instance-scoped pre-replay hook because native +full replay creates no `ForwardContext`. Native `NPUUBatchRunner` owns +preparation, execution and merge for eager and graph DBO. Plain MLA graph DBO +requires `FULL_DECODE_ONLY`; SFA/DSA accepted modes still require model-specific qualification. Native two-stage capture shares the ordinary descriptor's FIA workspace; replay merges captured records with current metadata for Ascend updates. Each child carries its actual stage ID for the opaque CAMP2P HCCL @@ -310,12 +323,13 @@ group lookup. Async CAM enables eager native DBO at DP1; it has no separate MoE stage pipeline. Async FFN FULL accepts eager Attention DBO; FFN stays outside Attention's native DBO thread schedule. -The FFN runner owns a separate ACL graph cache keyed by stage token counts and -A/F topology. Warmup runs the eager FFN path. Formal capture updates connector -state through `connector.control_plane` before entering `torch.npu.graph(...)`, +The synchronous FFN runner owns a separate ACL graph cache keyed by stage token +counts and A/F topology. Warmup runs the eager FFN path. Formal capture updates +connector state through `connector.control_plane` before entering `torch.npu.graph(...)`, so replay contains only model and data-plane operations. Replay additionally requires the control payload's -`is_graph_replaying=true`; an unknown key or eager payload executes eagerly. CAM async does not enter this path because validation -requires eager execution and `connector.control_plane` is `None`. +`is_graph_replaying=true`; an unknown key or eager payload executes eagerly. +Async CAM has no control plane and uses the separate complete-transaction +graph described in the [Async FFN runtime guide](../npu/async_ffn_runtime.md). ### Ascend native operators and packaging @@ -337,7 +351,7 @@ NPU extension, but the CAMP2P data path cannot run without it. all four `torch.ops.afd_ascend.afd_async_*` operators. The 910C package uses compact routed-only metadata; Attention owns shared experts and FFN owns routed experts. The 950 package does not register these async operators. -See [the routed-only guide](../../npu/CAM_ASYNC_ROUTED_OPS.md) for metadata +See [the routed-only guide](../npu/cam_async_routed_ops.md) for metadata and window lifetime contracts. Device validation remains necessary for each target topology. @@ -356,78 +370,6 @@ experimental output, role-specific defaults, optional stacks/modules, and a TensorBoard trace handler. Runners step the profiler on execution and stop it during shutdown. -## Tested runtime matrix - -This table records current validation gates and repository evidence. The -historical V1 rows are not current launch options or fallbacks. Evidence -does not expand the supported runtime contract. - -| Platform/path | Execution | Ubatching | Routing/quantization limits | Evidence | -| --- | --- | --- | --- | --- | -| Historical CUDA V1 (removed) + `P2pNcclAFDConnector` | Eager or `FULL_DECODE_ONLY` CUDA Graph | Native DBO, exactly two ubatches | Registered CUDA model boundaries; Attention-side or FFN-side gate where the model supports it | DeepSeek-V2-Lite eager/graph/DBO accuracy E2E; model, graph, connector, and profiler unit tests | -| CUDA V2 + `P2pNcclAFDConnector` | Eager or `FULL_DECODE_ONLY` native V2 CUDA Graph | Exactly two native microbatches; Attention DP1 or DP > 1 | Model-specific gate placement; PP/CP, elastic EP, EPLB, SP MoE, and compile SP rejected; role ranks equal DP x TP | Source-specific GPU validation is linked below; coverage is limited to the recorded model/topology combinations. DeepSeek-V2-Lite has known DBO shutdown failures. Qwen3.5/3.6 GDN DBO is unsupported. | -| CUDA V2 + `DeepEPM2NAFDConnector` | Eager or `FULL_DECODE_ONLY` native V2 CUDA Graph; SM90 or newer | Exactly two native microbatches, enabled on both roles; Attention DP1 or DP > 1, TP1 or TP2 | Qwen3 MoE only (all layers MoE, no shared expert); `compute_gate_on_attention=true`; FFN TP1 with EP over all FFN ranks, linear placement, `allgather_reducescatter` all2all and the Triton MoE weight layout; BF16 or block-128 FP8 payload; EPLB, PP/CP and SP MoE rejected | Historical vLLM 0.30 Qwen3-30B-A3B 2A2F GSM8K E2E for eager, graph, graph + DBO, and Attention TP2 eager/graph on H100 and H200, and for eager, graph and graph + DBO on GB200; two-node 2A2F eager, graph and graph + DBO with Attention and FFN on separate nodes; [GPU connector device tests](../../../tests/gpu/README.md) | -| Historical Ascend V1 (removed) + `CAMP2pAFDConnector` | Eager or current ACL Graph path | Native DBO, exactly two ubatches | Common and connector-local `compute_gate_on_attention=false`; `connector_extra_config.quant_mode=0`; plugin CANN ops required | V2-Lite 2A2F graph 300-question comparison, 2A1F/2A2F functional checks, and Legacy DBO regressions; runtime, graph, ops, connector, and profiler unit tests | -| Ascend V2 + `CAMP2pAFDConnector` | Eager, `FULL`, or `FULL_DECODE_ONLY` native V2 ACL Graph | Eager or graph native DBO with two ubatches and Attention DP > 1; plain MLA graph DBO requires `FULL_DECODE_ONLY` | `compute_gate_on_attention=false`; PP/CP, elastic EP and EPLB rejected; model-specific SP policy; role ranks equal DP x TP | Historical 0.30 V2-Lite 2A2F FULL 300-question comparison and 2A1F functional checks; current 0.32 Lite evidence is in the [validation record](https://github.com/vllm-project/afd-plugin/pull/507) | -| Ascend V2 + `CAMAsyncAFDConnector` | Eager Attention; layered FFN eager or FULL | Two native DBO microbatches on Attention, including DP1 and with FFN FULL; FFN stays outside the native DBO schedule | Experimental DeepSeek-V4 W4A8/W8A8 on A3/910C; `async=true`, `compute_gate_on_attention=true`, and `dynamicQuant=1`; PP/CP unsupported; plugin-owned async CAM ops required | Historical DSV4 layered comparisons and limitations are in the [validation record](https://github.com/vllm-project/afd-plugin/pull/507); they predate removal of the legacy path | - -The 300-question Ascend evidence in this table is historical. The 0.32 -[validation record](https://github.com/vllm-project/afd-plugin/pull/507) separately lists the -eight Lite 100-question cells, W4A8 fixture, and observed shutdown failures. - -The Async CAM row supports DeepSeek-V4 with static EP, PP=PCP=DCP=1, -and Attention-local SP. DeepSeek-V2/V2-Lite are no longer supported by -Async CAM; standard DeepSeek-V3.2 was already rejected by model admission. -Finegrained TP, explicit shared-expert DP configuration, and native -MTP and draft-model paths other than DeepSeek-V4 DSpark are excluded. DSpark -keeps Attention/native proposal eager and permits independent layered FFN FULL; -it does not support DBO or FFN two-stream. -FFN remains connector-driven. The row's upgrade evidence retains its original -scope; MRV2 deployment settings are documented in the -[CAM async guide](../../npu/CAM_ASYNC_CONNECTOR_USER_GUIDE.md#modelrunnerv2-deployment). - -DSV2/V3 Async CAM Attention with model SP selects Ascend shared-expert DP before -constructing the shared MLP. Each rank holds complete shared weights and computes -its local token rows. Attention TP/SP and routed-expert parallelism are unchanged. -Without model SP, the user's shared-expert DP setting is preserved. MRV2's -existing rejection of explicit shared-expert DP configuration is unchanged. - -Both backends target vLLM `0.32.0` (`5df63a72`); Ascend additionally targets -vLLM-Ascend `bd42eea0`. Source-specific -H20 model, topology, execution-mode and accuracy results are recorded in -the [GPU upgrade validation matrix](https://github.com/vllm-project/afd-plugin/pull/501) -and [runtime-fix validation matrix](https://github.com/vllm-project/afd-plugin/pull/502). -Those records identify the exact tested revisions. Earlier GPU results on -vLLM 0.30 do not establish vLLM 0.32 hardware qualification. - -Validation limits: - -- GPU accuracy evidence covers selected first-300 GSM8K comparisons, not - full-dataset accuracy, numerical equivalence, or performance improvement. -- DeepSeek-V2-Lite TP2 Graph + DBO and DP2 eager DBO can require force - killing during shutdown. Inference success does not imply graceful exit. -- Native V4 DBO requires a compatible upstream communication backend. The - recorded native V4 accuracy reference uses Graph without DBO. -- NPU 0.32 evidence is limited to the configurations and checks in the - [validation record](https://github.com/vllm-project/afd-plugin/pull/507); the seven AFD Lite - cases failed graceful exit. Historical NPU results do not qualify 0.32. - -The historical [NPU validation record](https://github.com/vllm-project/afd-plugin/pull/425#issuecomment-6063910923) -covers five representative configurations with the same first 300 GSM8K -questions: V2-Lite MRV1 graph, MRV2 FULL, ordinary Async CAM, and DSV4 Flash -W4A8 layered off/on. The [Legacy DBO repair](https://github.com/vllm-project/afd-plugin/pull/430#issuecomment-6063840311) -adds 2A1F/2A2F content regressions and a same-question 128-sample comparison; -it does not add a new 300-question result. These records do not establish all -model/topology support, numerical equivalence, performance, or clean device -shutdown. Synchronous DSV4 profiles remain smoke-only, with known incorrect -A5 concurrent answers; see their [E2E limits](../../../tests/e2e/README.md#dsv4-flash-sync-camp2p-concurrent-requests-local-4-or-16-npus). -GPU/NPU rank topology and connector resource rules remain owned by -[connector contracts](connector_contracts.md). - -The repository does not record a canonical CUDA container or assert an -official tag for the Ascend source pin `bd42eea0`; environment evidence is -not an authoritative package tag. - ## Failure and cleanup boundaries Unsupported graph, ubatch, async, gate, or quantization combinations fail in diff --git a/docs/design/module/ffn_runtime.md b/docs/design/module/ffn_runtime.md index e5301c1d3..5af8800fc 100644 --- a/docs/design/module/ffn_runtime.md +++ b/docs/design/module/ffn_runtime.md @@ -13,9 +13,10 @@ primary_code_paths: - "afd_plugin/v1/worker/ffn_worker.py" - "afd_plugin/v1/worker/npu/ffn_model_runner.py" - "afd_plugin/v1/worker/npu/async_ffn_model_runner.py" - - "afd_plugin/v1/worker/npu/async_ffn_two_stream.py" - "afd_plugin/v1/worker/npu/ffn_worker.py" related_code_paths: + - "afd_plugin/connectors/npu/async_cam_runtime/async_cam_startup.py" + - "afd_plugin/connectors/npu/async_cam_runtime/async_cam_session.py" - "afd_plugin/connectors/**" - "afd_plugin/model_executor/**" - "afd_plugin/compat/patches/engine_core.py" @@ -27,6 +28,7 @@ depends_on: - "compatibility_and_patches.md" validation_paths: - "tests/unit/ffn_faas/test_*.py" + - "tests/gpu/test_ffn_faas_gpu.py" - "tests/unit/model_executor/test_async_cam_layered.py" - "tests/npu/test_async_cam_layered_w8a8.py" - "tests/unit/v1/worker/test_ffn_metadata.py" @@ -34,10 +36,12 @@ validation_paths: - "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" - "tests/unit/v1/worker/test_async_ffn_runtime.py" - - "tests/unit/test_async_cam_startup_wiring.py" - - "tests/unit/connectors/test_async_cam_startup.py" + - "tests/unit/connectors/npu/test_async_cam_startup_wiring.py" + - "tests/unit/connectors/npu/test_async_cam_startup.py" + - "tests/unit/connectors/test_base.py" + - "tests/unit/connectors/npu/test_async_cam_session.py" + - "tests/unit/connectors/npu/test_async_cam.py" - "tests/npu/async_cam_runtime.py" - "tests/npu/standalone_cam_attention.py" - "tests/e2e/models/deepseek_v4_flash/ffn_graph.py" @@ -61,7 +65,7 @@ related_issues: - "#105" - "#107" - "#129" -last_reviewed: 2026-10-26 +last_reviewed: 2026-11-09 --- # FFN runtime @@ -81,6 +85,17 @@ FFN consumes plugin configuration, connector transport, model-side compute, platform mechanisms, and EngineCore compatibility. No shared module may depend on the FFN worker or runner implementation. +Async CAM's connector owns HCCL, payload validation and initialization/cleanup +coordination. It selects a Push or Pull session at construction. Both inherit +the concrete `FFNExecutionSession` lifecycle defaults from `connectors/base.py`. Push owns +per-stage routing queues; Pull owns its single window, publish/combine, device +status and quiet/drain/stopped lifecycle. Pull takes Startup, topology and the +dimensions/capacity used for initialization and bounded drain; it does not +retain the connector. The FFN runner continues to own eager/FULL execution, +graph references and completion; the worker still owns the receiver thread. +The current Mega executor reads the same live window through the connector's +read-only `comm_args` property. + ## Runtime selection The traditional FFN is launched as a `vllm serve` process, and AFD config normalization selects @@ -96,10 +111,16 @@ a vLLM Worker, upstream ModelRunner, EngineCore, scheduler or KV cache. Attention retains its existing model, connector and DBO path. Configuration, parameter loading and quantization remain vLLM/vLLM-Ascend dependencies. +The standalone Ascend CLI reuses vLLM engine argument definitions, parsing and +configuration construction. Existing AFD options remain in `additional_config`; +only host management and shutdown options are standalone-specific. The host +entry delegates parser creation to its selected platform, so CPU-only service +and host-loop imports remain independent of vLLM. + The CPU-only host `FFNModelRunner` accepts an `FFNBackend` factory, retains the backend before preparation, and owns the receiving thread and failure boundary. `FFNService` manages host status and stop requests. Platform options -and native composition live under `ffn_faas/platforms/npu/`: +and CAM runtime composition live under `ffn_faas/platforms/npu/`. - `AscendFFNEnvironment` owns device initialization, F-only compute/EP groups and native post-warmup CPU affinity. These groups coexist with the separate @@ -108,17 +129,18 @@ and native composition live under `ffn_faas/platforms/npu/`: accounting and cached-view lifetime. Startup selects a `RoutedModelAdapter` to provide checkpoint metadata and native owner construction. The backend uses the common resource without model-specific branching. -- `v4_model.py` contains V4 checkpoint metadata and `V4ModelAdapter`; - `v4_weights.py` contains the lazily imported native parameter owner. Native - routed modules, expert-map managers and quantizers remain resident. The - loader adapter delegates to the native V4 loader; dependency upgrades must - preserve its documented receiver and factory contracts. +- Model-specific code follows `ffn_faas/platforms/ffn_model///`. + Under `deepseek_v4/npu/`, `v4_model.py` contains V4 checkpoint metadata and + `V4ModelAdapter`. `v4_weights.py` contains the lazily imported native parameter + owner. Native routed modules, expert-map managers and quantizers remain + resident. The loader adapter delegates to the native V4 loader. Dependency + upgrades must preserve its documented receiver and factory contracts. - Platform settings enforce static CAM placement and reject EPLB, elastic EP and redundant experts. Retaining native expert state does not enable these algorithms. `FFNBackend` is a runtime lifecycle interface. Its `execute_step()` delegates -communication, computation, Graph and two-stream scheduling to the Async CAM +communication, computation and Graph execution to the Async CAM runtime. A shared cross-platform, cross-connector Step Loop and a pure `compute(item)` backend are not implemented. The runtime consumes the existing `LayeredFFNWeights` interface with W4A8 weights, without per-step model selection; @@ -133,10 +155,47 @@ native work retains all owners until whole-pair process isolation. A host See the [standalone service guide](../../npu/STANDALONE_FFN_SERVICE.md) for supported configuration, launch and shutdown behavior. +CUDA routed-experts models over `DeepEPM2NAFDConnector` have a GPU standalone +F, `--backend gpu-deepep-m2n`, under `ffn_faas/platforms/gpu/`: + +- `GPUDeepEPM2NAdapter` reuses the shared native engine CLI and F-only + launcher checks of `ffn_faas/platforms/launch.py`; with `GPUFFNConfig` it + checks on the CPU, before any device work, the M2N connector, TP 1 with EP, + the FULL_DECODE_ONLY or eager mode, the BF16 payload and lane parity with + `--enable-dbo`. The adapter resolves the model path but does not parse or + name a model family. +- `GPUFFNEnvironment` selects the GPU, preflights afdEP and DeepGEMM, builds + the F-only `VllmConfig` with `EngineArgs.from_cli_args` (`external_launcher`, + DP = F ranks, TP 1, EP), selects + the FFN-only routed-experts architecture, pins vLLM's MoE backend to Triton + before any FusedMoE is built, and creates the F-only world through vLLM's + `init_worker_distributed_environment`. +- The AFD model registry selects the FFN-only shell for the configured model. + For example, `AFDQwen3MoeRoutedExperts` keeps only `model.layers..mlp` + (native `Qwen3MoeSparseMoeBlock`) and loads through vLLM's + `DefaultModelLoader`, so the EP weight filter, FusedMoE expert loading and + strict loaded-weights check run unchanged; the embedding, final norm and LM + head are not allocated. +- `DeepEPM2NFFNBackend` constructs `GPUFFNRoutedExpertsModelRunner`, loads the + shell through that runner, starts the connector, and delegates each step to + `execute_connector_driven_step()`, the method the Worker loop also calls. + CUDA graphs are captured lazily from Attention's payloads, as on the Worker + path, so `/ready` means paired and receiving. + +`FFNBackend.execute_step()` returns `False` only when `cancel_receive()` ended +an idle receive; the host calls `cancel_receive()` after a stop request and +before joining the receiver. The GPU backend closes the M2N control plane +(aborting its own NCCL group); the NPU backend's `cancel_receive()` is a no-op. +Releasing DeepEP lanes is a barrier over every M2N rank, so a stopped GPU F +completes its release when Attention shuts down too. See the +[GPU standalone service guide](../../gpu/STANDALONE_FFN_SERVICE.md). + | Platform | Worker | Model runner | Current connectors | | --- | --- | --- | --- | | CUDA | `afd_plugin.v1.worker.AFDFFNWorker` | `GPUFFNModelRunner` | `P2pNcclAFDConnector` | | CUDA | `afd_plugin.v1.worker.AFDFFNWorker` | `GPUFFNRoutedExpertsModelRunner` | `DeepEPM2NAFDConnector` (any `AFDRoutedExpertsConnector`) | +| CUDA standalone (`ffn_faas`, no Worker) | `DeepEPM2NFFNBackend` | `GPUFFNRoutedExpertsModelRunner` | `DeepEPM2NAFDConnector` | +| NPU standalone (`ffn_faas`, no Worker) | `CAMFFNBackend` | `AsyncNPUFFNModelRunner` | `CAMAsyncAFDConnector` | | NPU | `afd_plugin.v1.worker.npu.AFDNPUFFNWorker` | `AFDNPUFFNModelRunner` (sync) or `AsyncNPUFFNModelRunner` | `CAMP2pAFDConnector`, `CAMAsyncAFDConnector` | GPU and NPU FFN runners are independent, connector-driven owners. The CUDA @@ -214,7 +273,7 @@ The V2 pair requires PP/PCP/DCP size 1, configured role ranks equal to DP x TP, static EP, and a registered AFD model. CUDA uses `P2pNcclAFDConnector` with model-specific gate placement and supports two native microbatches with Attention DP1 or DP > 1. See the -[runtime matrix](execution_platforms.md#tested-runtime-matrix) for tested +[runtime matrix](../../support.md#tested-runtime-matrix) for tested configurations and known DeepSeek-V2-Lite DBO shutdown failures. CUDA `DeepEPM2NAFDConnector` requires Attention-side gating, FFN TP 1, and the same DBO setting on both roles; see the @@ -227,21 +286,6 @@ FFN. Async FFN FULL accepts eager Attention DBO and continues to receive work items independently of Attention's threads. These limits describe the pair even though only Attention inherits a native V2 runner. -Async CAM also has an experimental, default-off `ffn_two_stream` mode for -layered W4A8/W8A8 eager/FULL. `AsyncFFNTwoStream`, owned by the Async runner, retains -two slots, compute/send streams and ready/done events. FULL owns four private -graph pools and four graphs; eager retires each slot only after its send is -observed complete. Both modes retain the bounded batch fence. Enabled eager -and FULL agree on executor, layer, capacity and core signatures before HCCL -creation. Attention must use eager MRV2 with PP=PCP=DCP=1, MoE sequence -parallelism and FlashComm1 disabled, and no speculative decoding. Attention DBO -remains supported with exactly two native microbatches under the existing -DP/TP constraints; FFN DBO stays disabled. Shutdown requires receiver -exit and device completion before release. Partial native failures retain -resources and remain visible to repeated health checks. Host tests do not -qualify native window reuse, graph capture or stream overlap; see the -[Async runtime guide](../../npu/ASYNC_FFN_RUNTIME.md#experimental-ffn-two-stream-execution). - ## Daemon step selection The worker selects one of two FFN step paths from the optional @@ -289,11 +333,18 @@ daemon thread: -> repeat ``` +On CUDA, one iteration is `GPUFFNModelRunner.execute_connector_driven_step()`, +shared by `AFDFFNWorker` and the standalone GPU FFN backend. It returns +`False` only for a receive cancelled by a requested stop; a closed control +plane without a stop request, and every other failure, raises. + ### Connector-driven loop The Ascend worker checks `connector.control_plane`. When it is `None`, as for -CAM async, the worker calls `execute_connector_driven_step()` instead of a -control-plane receive. For each layer, the runner: +CAM async, the worker calls `execute_connector_driven_step(shutdown_requested)` +instead of a control-plane receive. The runner asks the execution session once +per step whether to keep serving, so Pull adds no extra Store round trip to the +loop. For each layer, the runner: 1. receives a capacity-sized payload and its original CAM metadata through `connector.recv_attn_output()`; @@ -305,7 +356,7 @@ control-plane receive. For each layer, the runner: Async eager does not use the synchronous DP-metadata graph cache. Its FULL mode captures a complete original DR/FFN/CS transaction with device-side layer selection. The independent Async runner owns one graph and stable resources; -see [Async FFN runtime](../../npu/ASYNC_FFN_RUNTIME.md) for startup, supported +see [Async FFN runtime](../npu/async_ffn_runtime.md) for startup, supported W4A8/W8A8 configurations and the original receiver's process-cleanup boundary. DeepSeek-V4 W4A8 and W8A8 use layered GMM for both eager and FULL. Incompatible @@ -421,7 +472,7 @@ thread does not exit, native resources are not released underneath it. The connector then releases its remaining communication resources during runner shutdown. An earlier daemon failure is still surfaced after the join. Ascend uses its own connector close or asynchronous drain path before joining the -daemon; see the [Async FFN runtime](../../npu/ASYNC_FFN_RUNTIME.md). +daemon; see the [Async FFN runtime](../npu/async_ffn_runtime.md). ## Candidate invariants @@ -445,16 +496,8 @@ async model E2E coverage. ## Limitations and open issues -Current shared limits are the supported vLLM release, connector-driven FFN, -and registered role-aware model integrations. Ascend V1 native DBO accepts -exactly two ubatches. A V2-paired FFN uses the same AFD FFN runner with a -V2-compatible construction and forward-context seam. CUDA FFN consumes the -two-stage control layout and replays the matching graph when requested; -Ascend V2 consumes two-stage synchronous control metadata in eager DBO mode. -CAM async instead uses eager connector work items; native DBO runs only on -Attention and does not create a separate FFN stage scheduler. -Platform-specific limits are centralized in -[execution platforms](execution_platforms.md#tested-runtime-matrix). +Platform-specific execution limits and validation evidence are centralized in +[support matrix](../../support.md#tested-runtime-matrix). Issue [#107](https://github.com/JiusiServe/afd-plugin/issues/107) completed the optional control-plane split. Connector metadata ownership, transfer state diff --git a/docs/design/module/index.md b/docs/design/module/index.md index 0cdad9021..baba6dc12 100644 --- a/docs/design/module/index.md +++ b/docs/design/module/index.md @@ -21,10 +21,10 @@ upstream_refs: - "vLLM-Ascend target bd42eea01d8312c4df469fabaed45f77a12df4d3" verified_platform_refs: - "CUDA: tests/e2e marked e2e with AFD_E2E_BACKEND=gpu; no canonical image is recorded" - - "Ascend: current target in root installation instructions; V2-Lite and DSV4 Flash W4A8 evidence in the execution-platform matrix; older environments remain historical" + - "Ascend: current target in docs/installation.md; V2-Lite and DSV4 Flash W4A8 evidence in docs/support.md; older environments remain historical" related_issues: - "#129" -last_reviewed: 2026-10-13 +last_reviewed: 2026-11-09 --- # AFD module design @@ -35,12 +35,8 @@ verify boundaries, invariants, and validation evidence. The documents record plugin, role runtime, connector, model, platform, compatibility, and E2E testing contracts; open interfaces remain explicitly draft. -The 2026-10-13 source review used AFD `36dadb4`, vLLM `0.30.0` -(`ced6857a`), and Ascend `8d4409d6`; it is historical, not qualification of -the current pair. For the 0.32 upgrade, see the -[NPU validation record](https://github.com/vllm-project/afd-plugin/pull/507). Existing document -statuses and historical measurements retain their scope. Each module records -its remaining limitations and validation gaps. +`last_reviewed` records documentation review, not hardware qualification. +For installation, deployment and validation, use the [documentation index](../../README.md). ## Reading order @@ -62,17 +58,6 @@ Dependencies flow from role and integration modules toward shared boundaries; platform and compatibility modules adapt those boundaries to upstream runtimes. A lower-level document must not depend on a role implementation. -```text -attention_runtime ----+----> connector_contracts ----> plugin_boundary - | | -ffn_runtime ----------+ +----> execution_platforms - | -model_integration ----+----> execution_platforms - -attention_runtime ----+ -ffn_runtime ----------+----> compatibility_and_patches ----> plugin_boundary -``` - ```mermaid flowchart LR ATT["Attention runtime"] --> CONN["Connector contracts"] @@ -100,12 +85,12 @@ contract. File-level entries deliberately resolve mixed directories. | Primary document | Production paths | | --- | --- | | [Plugin boundary](plugin_boundary.md) | `afd_plugin/__init__.py`, `afd_plugin/config.py`, `afd_plugin/config_utils.py`, `afd_plugin/envs.py`, `afd_plugin/validation.py`, `afd_plugin/py.typed`, `afd_plugin/v1/__init__.py`, `afd_plugin/v1/worker/__init__.py`, `afd_plugin/v1/worker/npu/__init__.py`, `pyproject.toml` | -| [Attention runtime](attention_runtime.md) | `afd_plugin/v1/worker/attention_metadata.py`, `afd_plugin/v1/worker/attention_model_runner_v2.py`, `afd_plugin/v1/worker/attention_worker.py`, `afd_plugin/v1/worker/ubatch_wrapper.py`, `afd_plugin/v1/worker/npu/attention_model_runner_v2.py`, `afd_plugin/v1/worker/npu/attention_worker.py` | -| [FFN runtime](ffn_runtime.md) | `afd_plugin/ffn_faas/**`, `afd_plugin/v1/worker/ffn_metadata.py`, `afd_plugin/v1/worker/ffn_model_runner.py`, `afd_plugin/v1/worker/ffn_worker.py`, `afd_plugin/v1/worker/npu/ffn_model_runner.py`, `afd_plugin/v1/worker/npu/async_ffn_model_runner.py`, `afd_plugin/v1/worker/npu/ffn_worker.py` | +| [Attention runtime](attention_runtime.md) | `afd_plugin/v1/worker/attention_metadata.py`, `afd_plugin/v1/worker/attention_model_runner_v2.py`, `afd_plugin/v1/worker/attention_worker.py`, `afd_plugin/v1/worker/ubatch_wrapper.py`, `afd_plugin/v1/worker/npu/attention_model_runner_v2.py`, `afd_plugin/v1/worker/npu/attention_worker.py`, `afd_plugin/v1/worker/npu/dspark_speculator.py` | +| [FFN runtime](ffn_runtime.md) | `afd_plugin/ffn_faas/**`, `afd_plugin/v1/worker/ffn_metadata.py`, `afd_plugin/v1/worker/ffn_model_runner.py`, `afd_plugin/v1/worker/ffn_routed_experts_model_runner.py`, `afd_plugin/v1/worker/ffn_worker.py`, `afd_plugin/v1/worker/npu/ffn_model_runner.py`, `afd_plugin/v1/worker/npu/async_ffn_model_runner.py`, `afd_plugin/v1/worker/npu/ffn_worker.py` | | [Connector contracts](connector_contracts.md) | `afd_plugin/connectors/**/*.py`, `afd_plugin/connectors/npu/bin/**`, `afd_plugin/distributed/**/*.py` | | [Model integration](model_integration.md) | `afd_plugin/model_executor/**/*.py` | | [Execution platforms](execution_platforms.md) | `afd_plugin/compat/profiler.py`, `afd_plugin/compat/npu/forward_context.py`, `afd_plugin/compat/npu/ops.py`, `afd_plugin/compat/npu/profiler.py`, `afd_plugin/v1/worker/cuda_graph.py`, `afd_plugin/v1/worker/dbo.py`, `afd_plugin/v1/worker/npu/mla_graph.py`, `afd_plugin/v1/worker/npu/ubatch_runner.py`, `csrc/**`, `setup.py`, `MANIFEST.in` | -| [Compatibility and patches](compatibility_and_patches.md) | `afd_plugin/compat/__init__.py`, `afd_plugin/compat/vllm.py`, `afd_plugin/compat/npu/__init__.py`, `afd_plugin/compat/npu/feature_validation.py`, `afd_plugin/compat/npu/runtime.py`, `afd_plugin/compat/npu/runtime_config.py`, `afd_plugin/compat/patches/**/*.py` | +| [Compatibility and patches](compatibility_and_patches.md) | `afd_plugin/compat/__init__.py`, `afd_plugin/compat/vllm.py`, `afd_plugin/compat/npu/__init__.py`, `afd_plugin/compat/npu/feature_validation.py`, `afd_plugin/compat/npu/routing_ids.py`, `afd_plugin/compat/npu/runtime.py`, `afd_plugin/compat/npu/runtime_config.py`, `afd_plugin/compat/patches/**/*.py` | The routing inventory covers runtime and package code under `afd_plugin/**`, native sources under `csrc/**`, and packaging files that affect shipped @@ -123,21 +108,3 @@ The repository-level test contract is owned separately by Only identified invariant blocks in a `normative` document may use `MUST`, `MUST NOT`, or `SHOULD` as contract terms. - -## User and operational guides - -Operational guides remain separate from normative module design: - -- [NCCL P2P connector guide](../../gpu/NCCL_P2P_CONNECTOR_USER_GUIDE.md) -- [CAM P2P connector guide](../../npu/CAM_P2P_CONNECTOR_USER_GUIDE.md) -- [CAM async connector guide](../../npu/CAM_ASYNC_CONNECTOR_USER_GUIDE.md) -- [CAM routed operators](../../npu/CAM_ASYNC_ROUTED_OPS.md) -- [NPU validation](../../npu/TESTING.md) and [troubleshooting](../../npu/TROUBLESHOOTING.md) -- [Ascend NPU installation](../../../README.md#ascend-npu-installation) -- [Connector overview](../../../afd_plugin/connectors/README.md) -- [Deployment recipes](../../../recipe/README.md) - -- `afd_plugin/compat/npu/routing_ids.py` — opaque routing-id balancing boundary. - Registered during worker initialization; standalone scripts import the module - after device init. Tests: `tests/unit/compat/npu/test_routing_ids.py` and the - shared fixture in `tests/unit/fakes/routing_ids_op.py`. diff --git a/docs/design/module/model_integration.md b/docs/design/module/model_integration.md index f343b9c51..8f4ef7b23 100644 --- a/docs/design/module/model_integration.md +++ b/docs/design/module/model_integration.md @@ -31,14 +31,14 @@ verified_platform_refs: - "DeepSeek V2 Lite GPU and NPU model E2E paths" - "DeepSeek V2 Lite DeepEP M2N 2A2F E2E paths (eager, graph, graph + DBO, Attention TP2)" - "CAM async NPU model E2E path" - - "DeepSeek V4 CUDA contract tests and E2E profiles; see execution-platform matrix for hardware scope" + - "DeepSeek V4 CUDA contract tests and E2E profiles; see docs/support.md for hardware scope" - "DeepSeek V4 Ascend Hash-id boundary has focused unit coverage and a historical manual Ascend 950 run" related_issues: - "#86" - "#88" - "#105" - "#129" -last_reviewed: 2026-10-30 +last_reviewed: 2026-11-09 --- # Model integration @@ -64,10 +64,12 @@ make a backend-specific worker class the shared model API. | DeepSeek V4 CUDA role boundary | [`deepseek_v4.py`](../../../afd_plugin/model_executor/models/deepseek_v4.py) | [`test_deepseek_v4_construction.py`](../../../tests/unit/model_executor/models/test_deepseek_v4_construction.py), [`test_deepseek_v4_proxy.py`](../../../tests/unit/model_executor/models/test_deepseek_v4_proxy.py), [`test_deepseek_v4_weight_policy.py`](../../../tests/unit/model_executor/models/test_deepseek_v4_weight_policy.py) | | DeepSeek V4 Ascend Hash-id boundary | [`npu/deepseek_v4.py`](../../../afd_plugin/model_executor/models/npu/deepseek_v4.py), [`npu/deepseek_v4_attention_gate.py`](../../../afd_plugin/model_executor/models/npu/deepseek_v4_attention_gate.py) | [`test_deepseek_v4_hash_ids.py`](../../../tests/unit/model_executor/test_deepseek_v4_hash_ids.py), [NPU model role tests](../../../tests/unit/model_executor/models/npu/test_deepseek_v4.py), ids-mode cases in [`test_camp2p_token_ids.py`](../../../tests/unit/connectors/test_camp2p_token_ids.py) | | Qwen3 MoE role-aware model and weight loading | [`qwen3_moe.py`](../../../afd_plugin/model_executor/models/qwen3_moe.py) | [`test_qwen3_moe_construction.py`](../../../tests/unit/model_executor/models/test_qwen3_moe_construction.py), [`test_qwen3_moe_weight_policy.py`](../../../tests/unit/model_executor/models/test_qwen3_moe_weight_policy.py) | +| Qwen3 MoE FFN-only routed-experts shell | [`qwen3_moe.py`](../../../afd_plugin/model_executor/models/qwen3_moe.py) | [`test_qwen3_moe_construction.py`](../../../tests/unit/model_executor/models/test_qwen3_moe_construction.py) | | CUDA remote-experts boundary | [`deepseek_v2.py`](../../../afd_plugin/model_executor/models/deepseek_v2.py), [`gpu/p2p.py`](../../../afd_plugin/connectors/gpu/p2p.py) | [`test_p2p_experts_contract.py`](../../../tests/unit/connectors/test_p2p_experts_contract.py), [`test_deepseek_v2_proxy.py`](../../../tests/unit/model_executor/models/test_deepseek_v2_proxy.py) | | Forward-context adapter | [`forward_context.py`](../../../afd_plugin/model_executor/models/forward_context.py) | [`test_forward_context.py`](../../../tests/unit/model_executor/models/test_forward_context.py) | | NPU Async CAM native DBO handoff | [`npu/async_cam_dbo.py`](../../../afd_plugin/model_executor/models/npu/async_cam_dbo.py), [`npu/ubatch_runner.py`](../../../afd_plugin/v1/worker/npu/ubatch_runner.py) | [`test_async_cam_dbo.py`](../../../tests/unit/model_executor/test_async_cam_dbo.py), [`test_npu_model_runner_v2.py`](../../../tests/unit/v1/worker/test_npu_model_runner_v2.py) | | Ascend Attention-side gate | [`npu/deepseek_v2_attention_gate.py`](../../../afd_plugin/model_executor/models/npu/deepseek_v2_attention_gate.py) | Attention-gate unit cases in [`test_forward_context.py`](../../../tests/unit/model_executor/models/test_forward_context.py) | +| Ascend external-routing MoE execution | [`npu/deepseek_v4_external_routing_moe.py`](../../../afd_plugin/model_executor/models/npu/deepseek_v4_external_routing_moe.py) | [`test_deepseek_v4_external_routing_moe_cpu.py`](../../../tests/unit/model_executor/models/test_deepseek_v4_external_routing_moe_cpu.py) | | Ascend CAM dispatch layout | [`npu/async_cam_layout.py`](../../../afd_plugin/model_executor/models/npu/async_cam_layout.py) | [`test_async_cam_layout.py`](../../../tests/unit/model_executor/test_async_cam_layout.py) | | DeepSeek Ascend stage metadata ownership | [`npu/deepseek_attention_metadata.py`](../../../afd_plugin/model_executor/models/npu/deepseek_attention_metadata.py) | [`test_deepseek_attention_metadata.py`](../../../tests/unit/model_executor/test_deepseek_attention_metadata.py) | | Ascend CAM orchestration | [`npu/deepseek_v4.py`](../../../afd_plugin/model_executor/models/npu/deepseek_v4.py) | [`test_async_cam_npu.py`](../../../tests/e2e/models/deepseek_v4_flash/test_async_cam_npu.py) | @@ -84,14 +86,17 @@ registers lazy AFD wrapper paths under `AFD`-prefixed aliases. | `DeepseekV3ForCausalLM` | `AFDDeepseekV3ForCausalLM` | `AFDDeepseekV3ForCausalLM` | | `DeepseekV32ForCausalLM` | `AFDDeepseekV32ForCausalLM` | `AFDDeepseekV3ForCausalLM` | | `DeepseekV4ForCausalLM` | `AFDDeepseekV4ForCausalLM` | `AFDDeepseekV4ForCausalLM` | -| `DSparkDraftModel` | Unsupported | `AFDDSparkDeepseekV4ForCausalLM` (Async CAM: eager A, eager/FULL F) | +| `DSparkDraftModel` | Unsupported | `AFDDSparkDeepseekV4ForCausalLM` (Async CAM: eager/native decode Graph A, eager/FULL F) | | `GlmMoeDsaForCausalLM` | `AFDGlmMoeDsaForCausalLM` | `AFDGlmMoeDsaForCausalLM` | | `Qwen3MoeForCausalLM` | `AFDQwen3MoeForCausalLM` | `AFDQwen3MoeForCausalLM` | +| `Qwen3MoeRoutedExperts` | `AFDQwen3MoeRoutedExperts` | `AFDQwen3MoeRoutedExperts` (internal FFN-only alias) | | `Qwen3_5MoeForConditionalGeneration` | `AFDQwen3_5MoeForConditionalGeneration` | `AFDQwen3_5MoeForConditionalGeneration` | Only AFD workers switch their worker-local model configuration to the matching alias before constructing the AFD model runner. Non-AFD workers keep the checkpoint architecture and resolve to vLLM's native model class. +For CUDA FFN workers using a routed-experts connector, Qwen3 MoE selects the +internal FFN-only shell alias instead of the CausalLM wrapper. The DeepSeek, DeepSeek V2/V3/V3.2, and GLM aliases share the DeepSeek V2-derived implementation. DeepSeek V4, Qwen3 MoE, and Qwen3.5/3.6 each have a @@ -150,17 +155,21 @@ CUDA-only and requires synchronous `P2pNcclAFDConnector`, sequence-parallel MoE, EPLB, and the `deep_gemm_mega_moe` backend. The P2P connector validates one-dimensional `torch.int32` input IDs and preallocates their receive buffers for graph execution. See the -[runtime matrix](execution_platforms.md#tested-runtime-matrix) for tested +[runtime matrix](../../support.md#tested-runtime-matrix) for tested CUDA configurations and validation limits. ### DeepSeek V4 Ascend boundary [`npu/deepseek_v4.py`](../../../afd_plugin/model_executor/models/npu/deepseek_v4.py) wraps the native Ascend model independently of the CUDA adapter. Attention -owns attention, mHC, normalization and finalization. Synchronous CAMP2P sends -normalized hidden states and Hash-layer token IDs to FFN, which owns the -native gate and MoE. Async CAM keeps the Hash/top-k router and shared experts -on Attention and sends only routed-expert work; `dynamicQuant=1` is required. +owns attention, mHC, normalization and finalization. Synchronous CAMP2P has two +shapes, selected by the common `compute_gate_on_attention` field: with `false` +the FFN role owns the native gate and MoE, and with `true` (DeepSeek-V4, TP=1, no +DBO or ubatching) the router and shared experts run on Attention and FFN executes +only the transported routed experts. The connector-local flag of the same name is +kept for compatibility and does not select the mode. Async CAM keeps the +Hash/top-k router and shared experts on Attention and sends only routed-expert +work; `dynamicQuant=1` is required. The DSpark adapter follows the same boundary. Native checkpoint names `mtp.0..2` and KV names are retained; its CAM layer IDs start at the target's @@ -171,10 +180,14 @@ installs its own AFD metadata and opts out of cross-DP synchronization; native sampling, rejection, heads and context-KV updates are inherited. The [0.32 recipe](../../../recipe/npu/CAMAsyncAFDConnector/deepseek_v4_flash/README.md) describes A/F serving. In PD deployments, Prefill uses the upstream complete -draft model and KV connector. Attention and native draft proposal execution stay -eager. F may use the existing FULL transaction graph with one homogeneous W4A8 or W8A8 +draft model and KV connector. Attention may use native graphs with or without +PD and Attention-local FlashComm1/SP, with matching target/draft TP and DBO +disabled. The draft can remain eager or use its own native graph alongside the +target graph. Native resolution selects the effective mode for `FULL` or +`FULL_DECODE_ONLY`; device qualification uses compilation mode 0. +F may use the existing FULL transaction graph with one homogeneous W4A8 or W8A8 table covering target and draft; selection remains on device. DSpark excludes -DBO and FFN two-stream. +DBO. Shared-expert layout is model-specific. DeepSeek-V4 uses `AFDDeepseekV4SharedExperts` to preserve the native policy: @@ -188,23 +201,21 @@ separate from the routed layered executor: W4A8 applies the model's The synchronous DSV4 E2E profiles check concurrent serving but disable answer checking; the A5 profile records incorrect concurrent answers. Async DSV4 has an exact-answer smoke case; V4 true native DBO still requires model accuracy -qualification. Lite two-stage DBO has the limited result recorded in the -[0.32 validation record](https://github.com/vllm-project/afd-plugin/pull/507). See +qualification. See [E2E testing](e2e_testing.md) and the [layered guide](../../npu/CAM_ASYNC_CONNECTOR_USER_GUIDE.md#w4a8-and-w8a8-layered-gmm). ### Qwen3 MoE CUDA boundary `AFDQwen3MoeModel` uses the native `decoder_layer_type` injection hook. The -Attention role constructs native Qwen Attention and normalization modules and -uses a parameter-free `RemoteFFNProxy`; the FFN role constructs the complete -native dense MLP or `Qwen3MoeSparseMoeBlock`. Native forward order, FusedMoE, -packed parameter mapping, quantization paths, and weight loading remain owned -by vLLM. Checkpoint weights are filtered once by layer-stage path before the -native loader consumes them. +Attention role constructs native Qwen Attention and normalization modules. +With P2P, a parameter-free `RemoteFFNProxy` sends hidden states to FFN, which +runs the complete native MLP or MoE block; `compute_gate_on_attention=false` +is required. With DeepEP M2N, Attention owns gate/top-k and FFN runs routed +experts only; this requires `compute_gate_on_attention=true`, every layer to +be MoE, and no shared expert. Native weight loading follows role-based filtering. -This path is CUDA-only and requires `compute_gate_on_attention=false`. It fails -during construction for sequence-parallel MoE, EPLB, pipeline parallelism, -speculative decoding, LoRA, or NPU. +Both paths are CUDA-only and reject sequence-parallel MoE, EPLB, pipeline +parallelism, speculative decoding, and LoRA during construction. ### Qwen3.5/3.6 MoE CUDA boundary @@ -218,9 +229,7 @@ This path is CUDA-only, text-only, and requires `--language-model-only`, multimodal execution, sequence-parallel MoE, EPLB, speculative decoding, and LoRA before language-model construction; NPU model-config resolution fails before model loading. GDN DBO is unsupported under native hybrid/multimodal -DBO validation; ordinary eager and graph execution remain supported. Earlier -H20 Graph + DBO and recurrent-state/CUDA Graph parity checks describe the -removed adapter and do not qualify the current implementation. The current +DBO validation; ordinary eager and graph execution remain supported. The 2A1F suite keeps static EP and PP1. Attention-side router transport, async communication, multi-node, quantization, and performance remain outside the validated scope. @@ -234,9 +243,8 @@ ad-hoc `ForwardContext.afd_metadata` attribute. The current metadata supplies stage/request/token slicing, stage count, optional transaction ID, and a live connector reference. -Native forward-context creation can bypass an AFD model-runner call site: -Ascend V1 dummy runs do so, while V2 creates its context inside native -`execute_model()`. During those scopes, `use_afd_metadata_provider()` +Native V2 creates its context inside `execute_model()`. +During that scope, `use_afd_metadata_provider()` temporarily wraps `vllm.forward_context.create_forward_context`, lets the runner install the same `additional_kwargs` entry, and restores the original function in `finally`. This is a scoped compatibility adapter, not a permanent diff --git a/docs/design/module/plugin_boundary.md b/docs/design/module/plugin_boundary.md index 4ccd31f10..ca4ea25cc 100644 --- a/docs/design/module/plugin_boundary.md +++ b/docs/design/module/plugin_boundary.md @@ -37,7 +37,7 @@ verified_platform_refs: related_issues: - "#89" - "#129" -last_reviewed: 2026-10-13 +last_reviewed: 2026-11-09 --- # Plugin boundary @@ -129,36 +129,23 @@ presence is the activation signal: `parse_optional_afd_config()` returns raise if the required namespace is missing. Unknown top-level AFD keys are rejected and connector-specific values must be placed under `connector_extra_config`. -| Canonical key | Default | Current meaning | -| --- | --- | --- | -| `connector` | `P2pNcclAFDConnector` | Connector factory name selected from the built-in supported-name allow-list. | -| `async_dp` | `false` | Enables AFD async-DP compatibility behavior; currently valid only with `CAMAsyncAFDConnector`. | -| `role` | `attention` | Process role: `attention` or `ffn`. | -| `host`, `port` | `127.0.0.1`, `1239` | Connector rendezvous/control endpoint inputs. | -| `afd_process_group_timeout_s` | `120` | Positive AFD process-group rendezvous timeout in seconds; also bounds the NPU native DBO startup barrier. It is not a device-kernel cancellation deadline. | -| `num_attention_ranks`, `num_ffn_ranks` | `1`, `1` | AFD role-group sizes used by topology construction. | -| `compute_gate_on_attention` | `false` | Moves supported gate/MoE routing work to Attention. CUDA supports both gate placements at the remote-experts boundary; synchronous CAMP2P still requires `false`, while CAM async requires `true`. | -| `connector_extra_config` | `{}` | Envelope key parsed by the selected connector into a typed `ConnectorExtraInfo`; it is not stored on `AFDConfig`. | - -The compatibility aliases `afd_connector`, `afd_role`, `afd_port`, `afd_host`, -and `async` normalize to canonical fields. Supplying an alias and its canonical -field together is rejected. Boolean-like strings and integer-like values are -normalized; invalid types fail during parsing. +Public fields, defaults and aliases are defined in the +[configuration reference](../../configuration.md). Parsing normalizes +boolean-like strings and integer-like values; invalid types fail. Role rank is runtime state derived from the global DP rank and local PCP/TP coordinates. It is resolved once by the connector factory and is not an `AFDConfig` field. Connector construction parses `connector_extra_config` through the selected -connector class. P2P accepts only an empty mapping; CAMP2P and CAM async each -define a closed, typed schema and reject unknown fields. The detailed schemas +connector class. P2P accepts only an empty mapping; the other connectors +define closed, typed schemas and reject unknown fields. The detailed schemas belong to [connector contracts](connector_contracts.md) and the operational connector guides. -`AFDConfig.compute_hash()` currently includes connector, async mode, role, and -both role counts. This is an implementation detail used by -graph-affecting configuration paths, not a complete serialization or public -configuration identity. +`AFDConfig.compute_hash()` includes connector, async mode, role, both role +counts, and connector-specific configuration. It is used by graph-affecting +configuration paths, not as a complete serialization or public identity. ## Validation and worker selection @@ -178,7 +165,7 @@ communication resources are created. Both roles must use PP/PCP/DCP size 1, role ranks equal to DP x TP, static EP, and a registered AFD model; elastic EP and EPLB are rejected. CUDA V2 requires `P2pNcclAFDConnector`, model-specific gate placement, and exactly two microbatches with Attention DP1 or DP > 1. -See the [runtime matrix](execution_platforms.md#tested-runtime-matrix) for +See the [runtime matrix](../../support.md#tested-runtime-matrix) for tested configurations and known DeepSeek-V2-Lite DBO shutdown failures. Ascend V2 accepts synchronous `CAMP2pAFDConnector` with FFN-side gating or eager `CAMAsyncAFDConnector` with async DP and Attention-side gating. Ascend @@ -208,7 +195,7 @@ runtime integration and are not stable third-party extension interfaces. Profiling and connector diagnostics are read by their owning modules. Environment switches do not replace `additional_config["afd"]` as the activation or topology channel. See the -[runtime matrix](execution_platforms.md#tested-runtime-matrix) for hardware +[runtime matrix](../../support.md#tested-runtime-matrix) for hardware validation scope. `AFD_FORCE_SPAWN_MULTIPROCESSING=1` is an opt-in package-import adaptation for diff --git a/docs/npu/ASYNC_FFN_RUNTIME.md b/docs/design/npu/async_ffn_runtime.md similarity index 59% rename from docs/npu/ASYNC_FFN_RUNTIME.md rename to docs/design/npu/async_ffn_runtime.md index 136100682..33eb5d9bb 100644 --- a/docs/npu/ASYNC_FFN_RUNTIME.md +++ b/docs/design/npu/async_ffn_runtime.md @@ -1,25 +1,27 @@ # MRV2 Async CAM FFN eager/FULL runtime -AFD #449 captures one complete routed-expert transaction: original dispatch -receive (DR), layered W4A8/W8A8 FFN, and original combine send (CS). -Attention remains MRV2 eager. Issue #451 retires the separate MRV1 paths while -preserving this runtime and its original transport/lifecycle boundary. +FULL captures one routed-expert transaction: dispatch receive (DR), +W4A8/W8A8 FFN, and combine send (CS). Attention uses MRV2 eager or native target decode graphs. `AsyncNPUFFNModelRunner` owns Async CAM model loading, eager/FULL transactions and receive-safe resource release independently of the synchronous CAMP2P -`AFDNPUFFNModelRunner`. The FFN worker selects the owner from the connector; -neither runner inherits MRV1 or owns Attention request/KV state. +`AFDNPUFFNModelRunner`. The FFN worker selects the owner from the connector. ## Selection and startup The target is the pinned vLLM 0.32 / vLLM-Ascend stack on A3/910C, with homogeneous DeepSeek-V4 W4A8 or W8A8 layers, dynamicQuant=1, static expert placement, F TP=1 and Attention-side gating. Both roles select -`VLLM_USE_V2_MODEL_RUNNER=1`. Attention keeps `--enforce-eager`. +`VLLM_USE_V2_MODEL_RUNNER=1`. Attention normally keeps `--enforce-eager`; the +[target Attention graph](../../npu/CAM_ASYNC_CONNECTOR_USER_GUIDE.md#deepseek-v4-target-attention-graph) +profile accepts `FULL` or `FULL_DECODE_ONLY`, with native graph-mode resolution. +Attention-local FlashComm1/SP and native PD can be combined with the graphs. +DSpark draft may remain eager or use its native graph alongside target graph. +Device qualification uses compilation mode 0; other compilation modes have not +been validated. Other model-based drafts and MTP remain unsupported. Attention may enable native DBO under the existing two-microbatch constraints. Enable native DBO only on Attention; FFN consumes CAM work items without joining Attention's DBO thread schedule. -The experimental two-stream option below uses the same Attention DBO contract. The [0.32 validation record](https://github.com/vllm-project/afd-plugin/pull/507) identifies the tested W4A8 eager/FULL fixture and its limits, including graceful-exit failures. @@ -32,9 +34,9 @@ decode-only graph modes are rejected. DeepSeek-V4 DSpark uses this same F graph for target and draft routed experts. The combined layer table and startup agreement cover target0–42/draft43–45; all layers must share the same W4A8 or W8A8 layout. Native P loads the complete -draft model. Target/draft Attention and the native speculative configuration -remain eager. DSpark excludes DBO and FFN two-stream; see the -[DSpark recipe](../../recipe/npu/CAMAsyncAFDConnector/deepseek_v4_flash/README.md). +draft model. Target Attention may use the native decode graph; the draft can +remain eager or use its native query Graph. DSpark excludes DBO; see the +[DSpark recipe](../../../recipe/npu/CAMAsyncAFDConnector/deepseek_v4_flash/README.md). W8A8 reuses the loaded INT8 NZ weights, FP32 W13 scales and BF16 W2 scales. It needs no W4A8 compensation. The W2 binding requires per-layer bias entries, @@ -62,7 +64,10 @@ F HTTP API initialization is not a readiness requirement. ## Capture boundary and ownership By default, the runner owns capture/replay, weights, stable buffers, a private graph pool -and one execution stream. The connector owns CAM communication. The host-only `AsyncCAMGraphStartup` +and one execution stream. The connector owns HCCL, payload validation and +communication lifecycle coordination. Its Push session owns per-stage queues; +its Pull session owns a single window, native publish/combine, device status +and quiet/drain/stopped. The host-only `AsyncCAMGraphStartup` coordinator owns Store negotiation, stage waits and failure propagation; it owns no device resources or threads. A captured transaction receives activations, scales, expert counts and the original return metadata from DR. Device-side layer selection uses @@ -76,7 +81,8 @@ metadata goes to CS; host code does not select the layer or copy per-item inputs | CS completion tensor | Retained transfer state | | Weights and scale/compensation lists | Runner's model and executor | | Captured native workspace | Graph's private allocator pool | -| HCCL group, communication arguments and windows | Connector | +| HCCL group and final communication cleanup decision | Connector facade | +| Communication arguments/windows and pending routing references | Selected Push/Pull session | Each business batch submits as many replays as there are loaded FFN layers, then waits once for device completion. This is a bounded submission count, @@ -87,65 +93,9 @@ host metadata read or cancellation protocol. Changing weights, layout, capacity or topology requires fresh workers. There is no graph cache, online weight replacement or partial-rank restart. -## Experimental FFN two-stream execution - -Set these fields in the **FFN** `afd.connector_extra_config` to opt in: - -```json -{ - "ffn_two_stream": true, - "ffn_compute_aic_num": 12, - "ffn_compute_aiv_num": 24, - "ffn_combine_aiv_num": 24 -} -``` - -`ffn_two_stream` defaults to `false`. Both eager and FULL use exactly two -streams and two slots: receive plus the complete layered executor on the -compute stream, and FIFO combine-send on the send stream. Each send waits for -its matching compute-ready event; slot reuse waits for its previous send-done -event. The original batch metadata and CS return tensor remain owned by the -slot. The CS return tensor is an allocation anchor, not a completion flag. - -Core counts must be positive integers. Compute Cube must fit the available -AIC budget, compute Vector must be at least twice compute Cube, and the sum -of compute/send Vector limits must fit available AIV. The send stream sets -only its Vector limit. Invalid budgets, inconsistent FFN layouts and unsupported -Attention capabilities fail through pre-HCCL Store negotiation. Configuration -hashing includes connector extras. - -This option requires the layered MRV2 prerequisites above. Keep PP/PCP/DCP -at 1, SP and FlashComm1 off on both roles, and speculative decoding off. - -Attention DBO, when enabled, uses exactly two native ubatches and the existing -receive hook: the preceding raw receive is queued before the next ubatch sends -into CAM's single physical window. FFN slots remain independent work-item -buffers; they do not correspond to Attention ubatch IDs. Enable DBO only on -Attention, which remains eager, including when FFN uses FULL graph. -Streams are limited before warmup; FULL preserves -the existing matched warmup receive count. Enablement adds no startup receives. - -Eager retains each payload, intermediate and output until send completion is -observed before reallocating that slot. FULL captures **four graphs**, one -compute and one send graph per slot, each with a private pool and retained -mutable addresses. Events stay outside graph capture. Both modes preserve -the existing bounded batch size of `len(layer_ids)` and one device fence at -batch completion. Four graph objects do not mean four business transactions: -FULL submits two graph replays per transaction. - -Two-stream eager uses the same conservative shutdown rules as FULL: join the -receiver before releasing graph/tensor/model/communicator resources. Failed -capture, submission or device completion is terminal and retains resources. - -Host protocol and lifecycle tests cover this implementation. **NPU validation -is pending**: stream limits reaching custom tiling, native same-group DR/CS -window safety, split-graph capture and actual overlap/performance have not been -qualified. Keep this experimental option off in validated deployments until -the target stack passes those checks. - ## Resource lifetime -A host event stops further batch submissions. Graph reset and model/communicator +With default Push transport, a host event stops further batch submissions. Graph reset and model/communicator release are permitted only after the receive thread exits normally and device completion is confirmed. Partial capture, partial launch, execution failure or a still-live receive thread retains resources for process cleanup. A failure @@ -159,6 +109,17 @@ zero. Test reports distinguish successful inference from cleanup, and forced termination must not be described as normal shutdown. Single-stream eager and synchronous stopping policies retain their connector/process-level cleanup boundaries. +### Opt-in Pull execution session + +`connector_extra_config.ffn_mega_moe=true` selects Pull; Push remains the default. +The connector owns its execution session and all Pull status/Store handshakes. +Attention fences pending dispatch at `attention_step()` exit. FFN keeps serving +until every Attention source is quiet, drains active work, then publishes stopped; +only then may Attention and FFN release their communication resources. Failed +handshakes retain those resources. Standalone `close()` also validates shutdown +status before release: this is an additional defensive shutdown check, not a +relocated hot-path fence. + ## Computation semantics For W4A8, FULL and layered eager forward the same model `swiglu_limit` to W13. @@ -166,8 +127,20 @@ A positive `L` clamps gate above at `L` and up to `[-L, L]` before SiLU gating a quantization. Zero disables clipping. DeepSeek-V4's configured `10.0` therefore applies on both execution paths. The limit is a host scalar fixed at graph capture; layer selection remains controlled by the device `layer_index` tensor. -Layered W8A8 retains its unclipped computation and records the model limit as -metadata only. +Layered and Mega W8A8 retain unclipped computation regardless of the model +`swiglu_limit`; the executors pass zero and the native W8A8 Plan skips clamp. +Native Mega W4A8 applies the limit in its MC2 ACT epilogue and bounded-column +fallback. Native W8A8 uses an INT8 GEMM with INT32 intermediate results and a +bounded-column activation; weight scales stay FP32 for W13 and BF16 for W2. +Both plans share Pull transport and the BF16 return path. W8A8 applies weight +and token scales in Vector before activation, preserving the Layered rounding +sequence. The mode is selected once through `config[11]` (1=W4A8, 2=W8A8). + +The synthetic transport fixture accepts `--quantization W8A8 --swiglu-limit 0` +or `--swiglu-limit 10`; W8A8 ignores either model limit. Save matching Layered +eager outputs with `--save-outputs`, +then compare Mega eager and `--full-graph` using `--reference-outputs`; add +`--boundary` to exercise receive rounds, sparse experts and poisoned scratch. ## Validation @@ -178,37 +151,41 @@ A3/910C, 4A/4F, A TP4, F TP1, receive capacity 128 and DBO disabled. Select this topology still needs target-device validation. Synthetic per-group coverage does not qualify loading real per-group checkpoints. -- Host contracts: `tests/unit/connectors/test_async_cam_startup.py`, - `tests/unit/v1/worker/test_async_ffn_graph.py`, and - `tests/unit/v1/worker/test_async_ffn_runtime.py`. CPU-safe two-stream coverage - is in `tests/unit/v1/worker/test_async_ffn_two_stream.py` and - `tests/unit/test_async_cam_startup_wiring.py`; fakes verify host ordering and +- Host contracts: `tests/unit/connectors/npu/test_async_cam.py`, + `tests/unit/connectors/npu/test_async_cam_session.py`, + `tests/unit/connectors/npu/test_async_cam_startup.py`, + `tests/unit/v1/worker/test_async_ffn_graph.py`, + `tests/unit/v1/worker/test_async_ffn_runtime.py`, and + `tests/unit/connectors/npu/test_async_cam_startup_wiring.py`; fakes verify host ordering and ownership, not device execution. - Finite native comparison: `tests/npu/async_cam_runtime.py`. Run eager with `--save-outputs PREFIX`, then FULL with `--full-graph --reference-outputs PREFIX`, keeping other arguments identical. `--boundary` covers chunking and poisoned unused capacity; `--attn-ranks 2` covers different Attention sources. The test supplies finite work counts, compares outputs and checks captured addresses. + For Pull, add `--mega --idle-shutdown` to both eager and FULL runs so the + fixture exercises the real worker drain/finish handshake. Compare FULL to its + matching ON eager reference separately from the OFF-to-ON numerical check. - Real MRV2 generation: the bounded 4A/4F fixture below. Its test observers add synchronization for functional evidence, so it is not a performance test. ```bash python -m tests.e2e.models.deepseek_v4_flash.ffn_graph \ - --mode eager --model /path/to/checkpoint --afd-host "$AFD_HOST" \ - --devices 0,1,2,3,4,5,6,7 --output-dir /path/to/eager-evidence + --mode eager --mega --model /path/to/checkpoint --afd-host "$AFD_HOST" \ + --devices 0,1,2,3,4,5,6,7 --output-dir /path/to/on-eager-evidence python -m tests.e2e.models.deepseek_v4_flash.ffn_graph \ - --mode full --model /path/to/checkpoint --afd-host "$AFD_HOST" \ - --devices 0,1,2,3,4,5,6,7 --output-dir /path/to/full-evidence \ - --reference-dir /path/to/eager-evidence + --mode full --mega --model /path/to/checkpoint --afd-host "$AFD_HOST" \ + --devices 0,1,2,3,4,5,6,7 --output-dir /path/to/on-full-evidence \ + --reference-dir /path/to/on-eager-evidence +python -m tests.e2e.models.deepseek_v4_flash.ffn_graph \ + --mode eager --model /path/to/checkpoint --afd-host "$AFD_HOST" \ + --devices 0,1,2,3,4,5,6,7 --output-dir /path/to/off-eager-evidence \ + --reference-dir /path/to/on-eager-evidence ``` The fixture saves inputs, generated tokens, actual runner identities and per-rank execution evidence. FULL must capture once per F, replay during requests and -match eager tokens. These repository fixtures exercise the single-stream path. -Two-stream native and real-model qualification uses independent drivers outside -`tests/e2e` and `tests/npu`, setting the FFN connector options above and recording -four captures per F in FULL mode. Preserve driver sources, launch configuration, -source hashes, inputs and raw outputs with the evidence. +match eager tokens. Batch-end metadata samples do not prove per-layer coverage. Inference results and cleanup are recorded separately; remaining owned processes or failed requests block acceptance. Validation evidence must identify the diff --git a/docs/npu/CAM_ASYNC_ROUTED_OPS.md b/docs/design/npu/cam_async_routed_ops.md similarity index 98% rename from docs/npu/CAM_ASYNC_ROUTED_OPS.md rename to docs/design/npu/cam_async_routed_ops.md index bd17ec6b3..d51399309 100644 --- a/docs/npu/CAM_ASYNC_ROUTED_OPS.md +++ b/docs/design/npu/cam_async_routed_ops.md @@ -36,7 +36,7 @@ snapshot, under `src/comm_operator/ascend_kernels/`. The integration renames operator symbols and shared headers for AFD isolation while preserving the compact routed-only communication layout. The complete upstream MIT copyright and permission notice is retained in -[the repository license](../../LICENSE) and included in package license data. +[the repository license](../../../LICENSE) and included in package license data. ## Communication window configuration @@ -70,7 +70,7 @@ the extension CMake configuration (`AFD_SOC_VERSION`). `AFD_SKIP_ACLNN_BUILD=1` requires an existing, matching AFD vendor package with all four async ACLNN symbols. Rebuild the complete operator package after changing its source or SOC. Use the dependency stack in the -[installation guide](../../README.md#ascend-npu-installation). +[installation guide](../../installation.md#ascend-npu-installation). Load the package using the lazy loader before accessing the operators: @@ -230,7 +230,7 @@ unused cells and reuses a destination base while its TP peer remains unchanged. These changes preserve the compact protocol and do not add concurrent windows or remove empty-expert completion notifications. -Use [the standalone precision and performance harness](../../tests/npu/README.md) +Use [the standalone precision and performance harness](../../../tests/npu/README.md) to compare identical routing, TP width, window size, dtype, quantization and chunk capacity before claiming a speedup. A source optimization alone is not performance evidence; include sparse/empty peers and repeated window reuse. diff --git a/docs/gpu/DEEPEP_M2N_CONNECTOR_USER_GUIDE.md b/docs/gpu/DEEPEP_M2N_CONNECTOR_USER_GUIDE.md index f5cbd55fb..8b8e03cfc 100644 --- a/docs/gpu/DEEPEP_M2N_CONNECTOR_USER_GUIDE.md +++ b/docs/gpu/DEEPEP_M2N_CONNECTOR_USER_GUIDE.md @@ -49,6 +49,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: @@ -76,8 +86,9 @@ 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. ## Installation @@ -90,7 +101,7 @@ The connector needs three components in addition to the plugin and vLLM: ```bash python -m pip uninstall -y afdep python -m pip install --no-deps \ - "afdep-comm @ git+https://github.com/afdEP/afdEP-comm.git@" + "afdep-comm @ git+https://github.com/afdEP/afdEP-comm.git@" python -m afdep doctor --build ``` @@ -226,7 +237,63 @@ 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 [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 @@ -269,11 +336,18 @@ Send API traffic only to the Attention server. For eager, DBO, Attention TP and FP8 examples, see the [DeepEP M2N recipes](../../recipe/gpu/DeepEPM2NAFDConnector/README.md). +For BF16 Qwen3-30B-A3B, the FFN side can instead run as the standalone FFN +service (`torchrun -m afd_plugin.ffn_faas serve --backend gpu-deepep-m2n`), +which loads only the routed experts and constructs no vLLM Worker; see the +[standalone GPU FFN service guide](STANDALONE_FFN_SERVICE.md). + ## Validation - 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 diff --git a/docs/gpu/NCCL_P2P_CONNECTOR_USER_GUIDE.md b/docs/gpu/NCCL_P2P_CONNECTOR_USER_GUIDE.md index 2b598994e..bde5f1afe 100644 --- a/docs/gpu/NCCL_P2P_CONNECTOR_USER_GUIDE.md +++ b/docs/gpu/NCCL_P2P_CONNECTOR_USER_GUIDE.md @@ -80,26 +80,10 @@ AFD configuration is supplied through vLLM's `--additional-config` under the `af } ``` -### Fields - -| Field | Type | Default | Required / meaning | -| --- | --- | --- | --- | -| `role` | `"attention" \| "ffn"` | `"attention"` | Role owned by this process. Attention sends hidden states; FFN receives and returns FFN outputs. | -| `connector` | `str` | `"P2pNcclAFDConnector"` | Must be `P2pNcclAFDConnector` for this GPU path. | -| `host` | `str` | `"127.0.0.1"` | Non-empty rendezvous/control-plane host. All participating ranks must use a reachable, identical value. Host must be the first rank of FFN. | -| `port` | `int` | `1239` | Rendezvous port on `host`, valid range `1..65535`. It must be free and reachable. Subgroups share this rendezvous and need no ports of their own. | -| `num_attention_ranks` | `int` | `1` | Total number of AFD Attention ranks, including DP/TP-derived worker ranks. Must be positive. | -| `num_ffn_ranks` | `int` | `1` | Total number of AFD FFN ranks, including DP/TP-derived worker ranks. Must be positive. | -| `compute_gate_on_attention` | `bool` | `false` | When `false`, FFN owns the native gate and experts. When `true`, Attention owns the native gate and transfers router logits to the FFN external-router expert path. | -| `afd_process_group_timeout_s` | `int` | `120` | Rendezvous timeout, in seconds, for the AFD world process group join. Must be positive. The right value depends on rank count, placement, and model size; deployment tooling should compute it and pass it in. | -| `connector_extra_config` | `dict` | `{}` | Must remain empty; `P2pNcclAFDConnector` does not currently support connector-specific options. | -| `async` / `async_dp` | `bool` | `false` | Must remain `false` for `P2pNcclAFDConnector`; AFD async mode requires `CAMAsyncAFDConnector`. | - -Compatibility aliases currently accepted are `afd_role`, `afd_connector`, `afd_host`, and `afd_port`. Canonical field names should be used in new examples. - -The connector factory derives each process's role rank from its global DP rank -and local PCP/TP coordinates before connector construction. Do not configure a -role rank. +Use the [common configuration reference](../configuration.md) for defaults +and aliases. P2P requires `async_dp=false`, empty `connector_extra_config`, +and `host` on FFN rank 0. DeepSeek supports either gate placement; Qwen +requires `compute_gate_on_attention=false`. ## Topology rules @@ -131,7 +115,7 @@ export VLLM_USE_V2_MODEL_RUNNER=1 ``` The commands below illustrate a DP1/TP1 launch. See the -[runtime matrix](../design/module/execution_platforms.md#tested-runtime-matrix) +[runtime matrix](../support.md#tested-runtime-matrix) for tested configurations and known shutdown limitations. Attention side: @@ -179,11 +163,11 @@ role-rank totals equal to DP x TP, and a registered AFD architecture. Elastic EP, EPLB, sequence-parallel MoE and compile SP are rejected. Eager and `FULL_DECODE_ONLY` graph execution are accepted. V2 DBO requires -exactly two microbatches and accepts Attention DP1 or DP > 1. The generic -[GPU E2E runner](../../tests/e2e/runner.py) defines eager/graph DP2 and TP2 +exactly two microbatches and accepts Attention DP1 or DP > 1. The model-local +[GPU E2E cases](../../tests/e2e/models/deepseek_v2_lite/cases.py) define eager/graph DP2 and TP2 scenarios, including `afd-v2-eager-dbo-dp2` and `afd-v2-graph-dbo-dp2`. See the -[runtime matrix](../design/module/execution_platforms.md#tested-runtime-matrix) +[runtime matrix](../support.md#tested-runtime-matrix) for the tested scope and known DeepSeek-V2-Lite shutdown failures. A registered model does not qualify every topology. @@ -203,7 +187,7 @@ alone does not authorize replay. These behaviors are covered by the [connector tests](../../tests/unit/connectors/test_p2p_connector.py) and [FFN graph-policy tests](../../tests/unit/v1/worker/test_cuda_graph.py). -## Requirements and limitations to document in code +## Requirements and limitations - CUDA-capable PyTorch/vLLM environment with NCCL and vLLM's `PyNcclCommunicator` available. - Hidden-state tensors must reside on CUDA devices; CPU tensors are rejected. diff --git a/docs/gpu/STANDALONE_FFN_SERVICE.md b/docs/gpu/STANDALONE_FFN_SERVICE.md new file mode 100644 index 000000000..871e00c8f --- /dev/null +++ b/docs/gpu/STANDALONE_FFN_SERVICE.md @@ -0,0 +1,164 @@ +# Standalone GPU FFN service + +Run a registered CUDA routed-experts model in independent F-only processes +with `torchrun -m afd_plugin.ffn_faas serve --backend gpu-deepep-m2n`. The +service loads only the routed experts, reuses vLLM's native loader and the +existing DeepEP M2N FFN runtime, and constructs no vLLM Worker, upstream +ModelRunner, EngineCore, scheduler, KV cache or complete CausalLM. Attention +keeps its existing `vllm serve` launch, model, connector and DBO path. + +## Supported configuration + +| Supported | Rejected at startup | +| --- | --- | +| A registered CUDA routed-experts model; currently `Qwen3MoeForCausalLM`, BF16, every layer MoE, no shared experts | Unregistered models, FP8 or any `quantization_config`, dense or partial-MoE Qwen3, shared experts | +| `DeepEPM2NAFDConnector` with `payload_format: bf16` | P2P, afdEP-direct, flexible placement, FP8 payload | +| F TP=1, EP over all F ranks, linear expert placement | EPLB, elastic EP, redundant experts, PP/CP, LoRA, speculative decoding | +| Eager, FULL_DECODE_ONLY CUDA graphs, DBO with two lanes (eager and graph) | Other graph modes, more than two micro-batches | +| SM90 or newer (H100, H200, GB200) | Pre-Hopper GPUs | + +The routed experts must divide evenly across the F ranks. + +## Architecture + +```text +CPU service / host runner -> FFNBackend + DeepEPM2NFFNBackend (ffn_faas/platforms/gpu/backend.py) + GPUFFNEnvironment: GPU, F-only NCCL world and vLLM EP group + AFD routed-experts model selected by the model registry + GPUFFNRoutedExpertsModelRunner: the Worker path's FFN runner + DeepEPM2NAFDConnector: handshake, control plane, lanes, experts +``` + +Two worlds per deployment: the `torchrun` F-only world (vLLM's world, DP and +EP groups, used only to build FusedMoE and place experts) and the AFD M2N +world `[F0 .. Fn-1, A0 .. Am-1]` the connector creates at `--host:--port`. +Each F step runs `GPUFFNModelRunner.execute_connector_driven_step()`, the same +method the Worker-path `AFDFFNWorker` loop calls, so the data path, buckets, +lanes, graphs and DBO order are identical to a `vllm serve` FFN. See the +[FFN runtime design](../design/module/ffn_runtime.md#runtime-selection). + +Compared with a `vllm serve` FFN, F does not allocate the embedding, final +norm and LM head (about 1.16 GiB per rank for Qwen3-30B-A3B). + +## Launch + +Keep the Attention half of an existing recipe under +`recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/` unchanged and replace only +the FFN half. Example: 2A2F, graph and DBO on one node, Attention on GPUs 0 +and 1, F on GPUs 2 and 3: + +```bash +export VLLM_USE_V2_MODEL_RUNNER=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,"afd_process_group_timeout_s":600}}' \ + --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 & + +CUDA_VISIBLE_DEVICES=2,3 torchrun --nproc-per-node=2 \ + --master-addr=127.0.0.1 --master-port=29500 \ + -m afd_plugin.ffn_faas serve "$MODEL_PATH" --backend gpu-deepep-m2n \ + --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,"afd_process_group_timeout_s":600, + "connector_extra_config":{}}}' \ + --compilation-config '{"cudagraph_mode":"FULL_DECODE_ONLY"}' \ + --enable-dbo --status-port 29700 +``` + +Three ports: the `torchrun` master port (F-only world), the AFD rendezvous +`afd.port` (equal to Attention's) and the management port base. + +Engine options use vLLM's native engine CLI, as for the +[NPU service](../npu/STANDALONE_FFN_SERVICE.md); the AFD topology, rendezvous +and connector options go in `--additional-config`, exactly as a `vllm serve` +FFN would pass them. The executor defaults to `external_launcher`, and the GPU +backend also defaults `--skip-tokenizer-init` and `--enable-ep-weight-filter`. + +| Option | Meaning | +| --- | --- | +| `afd.host`, `afd.port` | AFD rendezvous; equal to Attention's | +| `afd.num_attention_ranks`, `afd.num_ffn_ranks` | equal to Attention's; `WORLD_SIZE`, `--data-parallel-size` and `afd.num_ffn_ranks` must agree | +| `afd.connector_extra_config` | `DeepEPM2NAFDConnector` options | +| `afd.afd_process_group_timeout_s` | M2N rendezvous timeout (default 120) | +| `--tensor-parallel-size 1 --enable-expert-parallel` | required: EP spans all F ranks | +| `--compilation-config '{"cudagraph_mode":"FULL_DECODE_ONLY"}'` or `--enforce-eager` | required: F graph or eager mode; other graph modes are rejected | +| `--enable-dbo` | two DeepEP lanes; required exactly when Attention runs `--enable-dbo` | +| `--receiver-shutdown-timeout` | host receive-loop join on stop (default 5 s) | + +F takes no Attention TP, batch size, model length or capture sizes: it adopts +Attention's TP size, dispatch capacity and bucket ladder in the M2N handshake, +and graph keys come from Attention's control payloads. + +### Configuration parity + +The handshake fails on both roles when they disagree: + +| Must agree | Attention | F | +| --- | --- | --- | +| Rank counts, AFD host/port | `afd` config | `afd` config | +| `payload_format`, `network`, `bucket_granularity`, `cpu_sync_min_tokens` | `connector_extra_config` | `connector_extra_config` | +| `num_lanes` | `--enable-dbo` | `--enable-dbo` (an explicit `num_lanes` must agree) | +| `capacity_tokens_per_rank` | derived or explicit | adopted, or explicit and equal | +| Hidden size, experts, top-k, BF16 | checkpoint | the same checkpoint files | + +`expert_backend`, `expert_num_sms`, `num_sms`, `num_qps` and the DeepEP +timeouts are F-local tuning knobs. + +### Startup ordering + +Attention joins the M2N rendezvous at the end of its model load; F joins after +its expert load. Independently started services must pair within +`afd_process_group_timeout_s` on both roles, or both fail. Raise it on both +sides to cover the expected start skew (600 s above). + +## Management and readiness + +Each local F rank serves `--status-port + LOCAL_RANK` (`0` chooses a free +port; each rank prints its address at start). + +- `/status`: phase, `loaded_experts`/`total_experts` (this rank's experts over + all layers; vLLM's loader reports no per-expert progress, so both values are + reported together after loading), `parameter_bytes`, `local_expert_first`/ + `local_expert_last`, `m2n_paired`, `attn_tp_size`, + `capacity_tokens_per_rank`, `num_buckets`, `network_hybrid`, + `compute_ep_world_size`, `num_lanes`, `graph_mode`, `dbo` and + `completed_batches` (one per control payload, including warmup, capture + and replay steps). +- `/ready`: paired with Attention and receiving. CUDA graphs are captured + later, driven by Attention's own capture; gate client traffic on + Attention's health endpoint, which comes up only after that. +- `/health`: host liveness. + +## Shutdown + +`POST /stop`, SIGINT and SIGTERM set a stop flag. An idle F cancels its +blocked control receive (the control plane aborts its own NCCL group), its +receive loop exits, and it releases CUDA graphs, DeepEP lanes, the F-only +groups and the experts, in that order. + +Releasing a DeepEP lane is a barrier over every M2N rank, so F finishes its +release only when Attention also shuts down (within the DeepEP GPU timeout, +120 s by default). Recommended order: + +1. drain client traffic; +2. `POST /stop` to every F rank; +3. stop Attention right away; +4. F exits 0 with `resources_released=true` in its status. + +A step that is in flight when the stop arrives completes first; it needs +Attention to finish the same step. If Attention dies or stops without F, +F's pending control receive cannot complete: once +`afd_process_group_timeout_s` passes, PyTorch's NCCL watchdog aborts the F +ranks (SIGABRT, so no status file is written) and `torchrun` exits non-zero. +A failure F detects itself (for example a DeepEP timeout in a step) reports +`failed` with `resources_released=false` and exits non-zero. A restarted Attention needs a +new M2N world, so restart F with it. Hot restart, reconnect and graceful +in-flight drain are not supported. diff --git a/docs/installation.md b/docs/installation.md new file mode 100644 index 000000000..94e0eba3b --- /dev/null +++ b/docs/installation.md @@ -0,0 +1,92 @@ +# Installation + +Requires Python **3.11-3.13** and [uv](https://docs.astral.sh/uv/). + +## Source setup + +```bash +git clone https://github.com/vllm-project/afd-plugin.git +cd afd-plugin +uv sync --group dev +``` + +`vllm` is an optional runtime extra so CPU-only or macOS development +environments can still run import/config tests without a CUDA wheel. + +## GPU installation + +For Linux / CUDA-capable environments only; Ascend users should skip this +command: + +```bash +uv sync --group dev --extra vllm +``` + +The optional GPU extra pins `vllm==0.32.0`. + +`DeepEPM2NAFDConnector` also needs the `afdep` extra (afdEP-comm) and +DeepGEMM, which is built from source; see +[Installation](gpu/DEEPEP_M2N_CONNECTOR_USER_GUIDE.md#installation) in +the connector guide. + +```bash +uv sync --group dev --extra vllm --extra afdep +``` + +## Ascend NPU installation + +Use the pinned source pair and tested stack below. Install a compatible driver +and firmware, then confirm device visibility with `npu-smi info`. + +| Component | Version | +| --- | --- | +| Python | Package range `3.11–3.13`; tested `3.12.13` | +| vLLM | `0.32.0`, commit `5df63a7203cdb9328d824a27b682e77c7b9928be` | +| vLLM-Ascend | commit [`bd42eea0`](https://github.com/vllm-project/vllm-ascend/commit/bd42eea01d8312c4df469fabaed45f77a12df4d3) | +| Tested CANN / torch / torch-npu | `9.2` / `2.10.0+cpu` / `2.10.0.post7` | +| Tested Triton Ascend / CAM | `3.6` / `1.3` | + +See the [support reference](support.md#validation-limits) for hardware coverage +and known shutdown limitations. + +### Environment + +The recorded A3 environment uses the pinned pair and stack above. The +repository does not claim a released container tag for the Ascend target. Use the +[installation guide at that source snapshot](https://github.com/vllm-project/vllm-ascend/blob/bd42eea01d8312c4df469fabaed45f77a12df4d3/docs/source/getting_started/installation.md) +to prepare a matching A3/openEuler environment, then install AFD from the +repository root. + +### Install AFD + +From the AFD repository root: + +```bash +AFD_BUILD_ASCEND_OPS=1 \ +SOC_VERSION=ascend910_9391 \ +python -m pip install -v --no-build-isolation --no-deps -e . +``` + +`--no-deps` preserves the matched runtime; `--no-build-isolation` uses its +CANN/torch-npu toolchain. `AFD_BUILD_ASCEND_OPS=1` forces the native build; +`0` skips it. When unset, setup detects the Ascend environment. See the +[native build guide](../csrc/npu/README.md) for build controls and artifacts. + +### Verify + +```bash +python - <<'PY' +import torch +import torch_npu +import vllm_ascend + +from afd_plugin.compat.npu import ensure_afd_ascend_ops_loaded + +assert torch.npu.is_available(), "torch-npu cannot see an Ascend device" +ensure_afd_ascend_ops_loaded() +print("AFD_OPS_OK") +PY +``` + +After `AFD_OPS_OK`, choose a [deployment recipe](../recipe/README.md). +Use the [NPU test guide](../tests/npu/README.md) to validate the runtime. diff --git a/docs/npu/CAM_ASYNC_CONNECTOR_USER_GUIDE.md b/docs/npu/CAM_ASYNC_CONNECTOR_USER_GUIDE.md index dd75f8b05..543b137a7 100644 --- a/docs/npu/CAM_ASYNC_CONNECTOR_USER_GUIDE.md +++ b/docs/npu/CAM_ASYNC_CONNECTOR_USER_GUIDE.md @@ -6,9 +6,9 @@ exchange routed-expert activations with independent FFN expert ranks through CAM async dispatch/combine operators. Use the pinned dependency pair in the -[root installation instructions](../../README.md#ascend-npu-installation). +[installation instructions](../installation.md#ascend-npu-installation). Current runtime constraints and validation evidence are summarized in the -[runtime matrix](../design/module/execution_platforms.md#tested-runtime-matrix). +[runtime matrix](../support.md#tested-runtime-matrix). > [!WARNING] > CAM async is experimental. V4 true DBO boundary qualification and performance @@ -17,11 +17,12 @@ Current runtime constraints and validation evidence are summarized in the ## When to use this connector -Use this connector for eager asynchronous Ascend inference with these requirements: +Use this connector for asynchronous Ascend inference with these requirements: - plugin-owned Ascend operators are built and installed on every node; - Attention performs MoE gating before dispatch to FFN ranks; -- Attention execution is eager and AFD async-DP is enabled; **`async_dp=true` is required**; +- Attention uses eager execution or the DeepSeek-V4 graph profile below; + **`async_dp=true` is required**; - regular prefill and autoregressive decode steps use the same CAM path; - optional ubatching uses exactly two native MRV2 DBO microbatches, including Attention DP1. @@ -32,9 +33,9 @@ unsplit locally. Both token thresholds must be at least two. Attention DP replicas make that decision independently; FFN is connector-driven. FFN also supports the experimental complete-transaction FULL mode described in -[Async FFN runtime](ASYNC_FFN_RUNTIME.md#selection-and-startup). -FULL requires MRV2 on both roles and accepts Attention-side DBO under the -existing constraints. Attention graph execution remains unsupported. +[Async FFN runtime](../design/npu/async_ffn_runtime.md#selection-and-startup). +FULL requires MRV2 on both roles and accepts eager Attention-side DBO under +the existing constraints. The target Attention graph profile is described below. ## ModelRunnerV2 deployment @@ -45,7 +46,7 @@ Attention uses native Ascend ModelRunnerV2; FFN uses the independent The V2 Async CAM contract is: - DeepSeek-V4 with homogeneous W4A8 or W8A8 routed experts on Ascend 910C/A3; -- `--enforce-eager` on Attention and `--enable-expert-parallel` on both roles, with +- eager Attention or the target graph profile below, and `--enable-expert-parallel` on both roles, with `async_dp=true` and `compute_gate_on_attention=true`; - on FFN, select eager with `--enforce-eager`; for FULL, omit `--enforce-eager` and pass `--compilation-config '{"cudagraph_mode":"FULL"}'`, following the narrow @@ -67,6 +68,41 @@ per-ubatch model forwards, and output merge. Native splitting can divide one request across microbatches. Native draft-model/MTP paths other than DeepSeek-V4 DSpark are unsupported. +### DeepSeek-V4 target Attention graph + +For a DeepSeek-V4 deployment, replace Attention's `--enforce-eager` +with native ACL decode graph selection: + +```bash +--compilation-config '{"mode":0,"cudagraph_mode":"FULL_DECODE_ONLY"}' +``` + +The profile requires PP=PCP=DCP=1, static expert placement, +and DBO/ubatching disabled. Attention-local FlashComm1/SP and native PD may remain +enabled. DSpark may stay enabled: `speculative_config.enforce_eager=true` keeps the +draft eager; `false` enables its native query Graph alongside target Graph. Target and draft use matching +TP widths. Other model-based drafts and MTP remain unsupported. PD uses the +existing upstream KV connector; the W4A8 device comparison covers +`MooncakeHybridConnector` with a native full-draft P. Draft Graph requires target Graph because +native target capture owns the draft capture lifecycle. +Attention DPs remain independent. Each Attention TP group uses the existing +CAM token split and output gather; set `attn_ranks_per_dp` to Attention TP +on both roles. + +Native Ascend owns graph inputs, batch buckets, attention-parameter updates and +replay. Ineligible batches, including prefill, retain native eager fallback. +`FULL` is also accepted and resolved by the native attention backend; the +current DSV4 backend resolves it to `FULL_DECODE_ONLY`. The device qualification +uses compilation mode 0, which disables Dynamo while retaining ACL graph +capture. Other compilation modes have not been validated by this qualification. +AFD stages local real/padded token metadata before capture/replay; Async CAM +does not send synchronous DP control messages. Existing synchronous connectors +retain their control-plane behavior. + +Keep the existing FFN execution settings. Transport warmup and F +service readiness use the existing startup protocol. This option introduces no +new graph owner, native kernel or shutdown protocol. + ## CAM async data flow CAM async removes synchronization between data-parallel replicas and avoids the @@ -139,9 +175,8 @@ from the model routed-expert count and `num_ffn_ranks`; use a model/topology in which routed experts divide evenly across FFN ranks. All roles must use the same model, routed-expert layout, quantization, HCCL address, rank counts, CAM settings. Enable native DBO only on Attention; FFN must not receive -`--enable-dbo` or native ubatching flags. Legacy connector flags may remain in -shared JSON because their normalization is Attention-only. In an async deployment, every Attention (A) -and FFN (F) process must also set `--max-num-batched-tokens` to the same value. +`--enable-dbo` or native ubatching flags. Every Attention and FFN process must +set `--max-num-batched-tokens` to the same value. ## AFD configuration @@ -169,42 +204,19 @@ There is no separate `--afd-config` option. } ``` -### Common fields - -| Field | Type | Default | Meaning and constraint | -| --- | --- | --- | --- | -| `role` | `"attention" \| "ffn"` | `"attention"` | Role owned by this process. | -| `connector` | `str` | `"P2pNcclAFDConnector"` | Must be `CAMAsyncAFDConnector`. | -| `async` / `async_dp` | `bool` | `false` | Must be `true`. `async` is the accepted compatibility alias for canonical `async_dp`. | -| `host` | `str` | `"127.0.0.1"` | **Must be the IP address of the node that owns Attention rank 0.** Every rank must use the same reachable value. | -| `port` | `int` | `1239` | HCCL rendezvous port in `1..65535`; it must be free and reachable. | -| `afd_process_group_timeout_s` | `int` | `120` | Positive AFD group rendezvous timeout in seconds; also bounds the NPU native DBO startup barrier, not device receive polling. | -| `num_attention_ranks` | `int` | `1` | Total Attention ranks across all DP and TP groups. | -| `num_ffn_ranks` | `int` | `1` | Total FFN expert ranks. | -| `compute_gate_on_attention` | `bool` | `false` | Must be `true`; CAM async runs MoE routing on Attention before dispatching to FFN ranks. | -| `connector_extra_config` | `dict` | `{}` | Connector-specific settings. Unknown top-level AFD fields are rejected. | - -Compatibility aliases `afd_role`, `afd_connector`, `afd_host`, `afd_port`, -and `async` are accepted. New configurations should use canonical names, -including `async_dp`; enable DBO through the Attention CLI flags below. +Use the [common configuration reference](../configuration.md) for defaults +and aliases. CAM async requires `async_dp=true`, +`compute_gate_on_attention=true`, and `host` on Attention rank 0. +Enable native DBO through Attention CLI flags, not connector extra config. ### CAM async `connector_extra_config` | Field | Type | Default | Meaning and constraint | | --- | --- | --- | --- | | `dynamicQuant` | `int` | `0` | Set to `1`: FFN receives dynamically quantized routed activations and scale tensors. The default `0` is not supported by the layered FFN path. | +| `ffn_mega_moe` | `bool` | `false` | Selects Mega FFN with Pull transport; requires `dynamicQuant=1` and the same setting on all A/F ranks. See the [Pull execution session](../design/npu/async_ffn_runtime.md#opt-in-pull-execution-session). | | `attn_ranks_per_dp` | `int` | `1` | Positive Attention TP rank count per DP replica. It is the CAM Attention grouping width and is independent of the FFN process's local TP size. | | `hccl_buffer_size` | `int` | unset | Positive CAM HCCL buffer size in MB. The override applies only to the `afd_async_cam` communication domain. When unset, HCCL uses `HCCL_BUFFSIZE`, then its built-in default. Configure the same value on every Attention and FFN rank in the domain. | -| `ffn_two_stream` | `bool` | `false` | Experimental FFN-only receive/compute and send streams with two slots. Requires an eligible layered MRV2 executor; eager and FULL are supported in code, with NPU qualification pending. | -| `ffn_compute_aic_num` | `int` | `12` | Positive compute Cube limit, bounded by available AIC. | -| `ffn_compute_aiv_num` | `int` | `24` | Positive compute Vector limit, at least twice the compute Cube count. | -| `ffn_combine_aiv_num` | `int` | `24` | Positive send Vector limit; compute plus send Vector must fit available AIV. | - -For `ffn_two_stream=true`, disable SP and FlashComm1 on both roles. Attention -may use native DBO with two ubatches under the constraints below; it remains -eager, and FFN must keep DBO disabled even when using FULL graph. -See [two-stream execution](ASYNC_FFN_RUNTIME.md#experimental-ffn-two-stream-execution) -for graph ownership, shutdown and the remaining hardware qualification gates. For a DP+SP deployment such as DP3TP2 Attention + DP2TP1/EP2 FFN, set `VLLM_ASCEND_ENABLE_FLASHCOMM1=1` only on Attention and explicitly set it to @@ -231,6 +243,14 @@ prefill/decode batches and must be at least the microbatch count. Native slicing including parent padding and requests crossing a microbatch boundary. The model owns SP padding and each owner's layout restoration/TP collectives. +The connector keeps configuration, topology, HCCL, payload validation and +lifecycle coordination in `async_cam.py`. Its Push session keeps a FIFO per +stage. Both sessions live in `async_cam_runtime/async_cam_session.py`. Pull owns one window +across stages, native publish/combine, device status and shutdown coordination. +The session is selected once at construction. Attention uses the existing +connector entry points; FFN uses separate receive/send for Push, or Mega's fused +transaction for Pull. Graph execution remains in the FFN ModelRunner. + CAM has one physical window per Attention rank. Before another ubatch sends, a native peer receive hook completes the preceding raw receive with its owner's explicit transfer context. The resumed owner restores layout and adds shared @@ -245,7 +265,7 @@ as DP > 1. The Ascend adapter extends native factory and threshold selection to DP1 because upstream otherwise returns early there. Async child contexts start with `dp_metadata=None`; FFN DP coordination belongs to the connector. DP1 model correctness still needs the hardware checks in the -[NPU validation guide](TESTING.md#native-dbo-contracts-and-device-execution). +[NPU validation guide](../../tests/npu/README.md#native-dbo-contracts-and-device-execution). For shape-only runtime diagnostics, set `AFD_ASYNC_MOE_LAYOUT_LOG=1` on Attention. The log reports the parent token @@ -255,7 +275,7 @@ It does not inspect tensor values or force a device synchronization. ## Requirements Use Ascend 910C, the pinned stack from the -[root installation instructions](../../README.md#ascend-npu-installation), +[installation instructions](../installation.md#ascend-npu-installation), and plugin-owned source-built operators on every node. Build the AFD operators from the repository root inside the container: @@ -280,7 +300,7 @@ rank: Async CAM host tiling reads the environment, whereas HCCL allocation can use the domain override. For the JSON example above, export `HCCL_BUFFSIZE=4096`. Changing only `hccl_buffer_size` can misalign the window capacity and tail-state address. `HCCL_BUFFSIZE` also affects unrelated groups -without explicit overrides; see the [operator window contract](CAM_ASYNC_ROUTED_OPS.md#communication-window-configuration). +without explicit overrides; see the [operator window contract](../design/npu/cam_async_routed_ops.md#communication-window-configuration). For the experimental FlashComm1/SP cases, set `VLLM_ASCEND_ENABLE_FLASHCOMM1=1` only on Attention. Set it to `0` for plain TP @@ -292,18 +312,17 @@ actual vendor library path. Missing source operators fail startup. ## Current limitations -- Attention is eager. FFN FULL targets homogeneous W4A8 or W8A8 layers on A3/910C; - see [supported configuration](ASYNC_FFN_RUNTIME.md#selection-and-startup) - and [validation commands](ASYNC_FFN_RUNTIME.md#validation). +- Attention is eager except for the DeepSeek-V4 native decode graph profile above. + FFN FULL targets homogeneous W4A8 or W8A8 layers on A3/910C; + see [supported configuration](../design/npu/async_ffn_runtime.md#selection-and-startup) + and [validation commands](../design/npu/async_ffn_runtime.md#validation). Broader models and topologies require separate validation. - Regular prefill and autoregressive decode steps are supported; context parallel execution remains outside the supported Async CAM topology. - Native DBO requires eager Attention, two ubatches, and no speculative decoding. FFN keeps its independent connector-driven execution. -- Attention DP1 reuses native DBO through scoped Ascend compatibility adapters; - the legacy request/token split selector no longer controls boundaries. -- ModelRunnerV2 rejects PP, PCP, and DCP; V1 PCP and async MoE context - parallelism remain unsupported. +- Attention DP1 reuses native DBO through scoped Ascend compatibility adapters. +- ModelRunnerV2 rejects PP, PCP, and DCP. - Routed experts should divide evenly across FFN ranks. For startup, HCCL, CAM operator, and shutdown failures, see the @@ -311,18 +330,10 @@ For startup, HCCL, CAM operator, and shutdown failures, see the ## W4A8 and W8A8 layered GMM -Async CAM FFN always uses the two layered GMM operators. FFN startup logs the -quantization, layer count and layout as `Async CAM layered GMM`. - -### Migration from the legacy FFN path - -Remove `AFD_ASYNC_CAM_LAYERED_GMM` from all launch environments; the selector -has been removed and setting it to `0` no longer selects the old path. -DeepSeek-V2 and V2-Lite, including BF16 checkpoints, no longer support Async -CAM because their FFN weights do not implement the layered contract. The -standard DeepSeek-V3.2 architecture was already rejected by model validation and -remains unsupported. Other connectors keep their own model support boundaries. -Deployments needing the old Async CAM path must remain on an earlier revision. +By default, Async CAM FFN uses two layered GMM operators. +`ffn_mega_moe=true` selects the fused Mega executor and Pull transport using +the same weight table. Startup logs its quantization, layer count and layout +as `Async CAM layered GMM`. The intended configuration is Ascend 910C / `ascend910_93`, async CAM FFN, Attention-side gate, `dynamicQuant=1`, eager execution (or the separately @@ -336,9 +347,8 @@ loader rejects positive `group_size`; the current target is limited to per-channel W4A8, or dynamic W8A8 with FP32 W13/BF16 W2 per-channel scales. W8A8 uses loaded INT8 NZ tensors and a single shared zero INT32 W2 bias; it does not require W4A8 compensation or duplicate expert weights. -DSV4 Flash W4A8 has layered off/on 300-question comparisons and -real INT8-activation/INT4-weight operator traces in the -[runtime validation record](../design/module/execution_platforms.md#tested-runtime-matrix). +Historical DSV4 Flash W4A8 layered off/on comparisons and operator traces are in the +[runtime validation record](../support.md#tested-runtime-matrix). Per-group operator coverage does not extend the checkpoint loader to per-group weights. Shared experts remain on Attention. DSV4 applies routed scaling in top-k, so @@ -349,7 +359,7 @@ For positive `L`, it computes `SiLU(min(gate, L)) * clamp(up, -L, L)` before quantization; zero keeps clipping disabled. This restores DeepSeek-V4's configured limit of `10.0`. Layered W8A8 retains its unclipped computation. FULL reuses the existing native DR and CS operators. See -[resource lifetime](ASYNC_FFN_RUNTIME.md#resource-lifetime) for completion and +[resource lifetime](../design/npu/async_ffn_runtime.md#resource-lifetime) for completion and process-cleanup requirements. Missing W4A8 `w13_scale_bias` or `w2_scale_bias`, dynamic EPLB, non-SiLU activation, incompatible parameters, other quantizations, mixed quantization, @@ -361,6 +371,7 @@ slice of the original metadata, or a device mapping, selects the layer; the original metadata goes to combine-send. Empty ranks, chunks, and ubatches use the existing communication protocol. With `AFD_CAM_OP_IO_LOG` enabled, the path logs only Tensor shape, dtype, and device without reading metadata values. +Input field names follow the registered native operator schema. The Async eager worker synchronizes after each complete FFN layer loop; FULL submits a bounded batch on one stream and waits once for that batch. These submission loops do not define request boundaries. Measure peak @@ -372,8 +383,8 @@ counts, parameter rejection, startup selection, and the absence of explicit metadata D2H in the Python path. They do not establish NPU numerical accuracy, empty-work completion, hidden synchronization, or a performance gain. Run the compiled Meta and NPU numerical tests described in the -[testing guide](TESTING.md#layered-w4a8-validation) and the -[FULL runtime checks](ASYNC_FFN_RUNTIME.md#validation). Deployments require +[testing guide](../../tests/npu/README.md#layered-operator-precision) and the +[FULL runtime checks](../design/npu/async_ffn_runtime.md#validation). Deployments require target-device validation of completion, checkpoint accuracy and performance. To switch from FULL to layered eager, pass `--enforce-eager` before restarting diff --git a/docs/npu/CAM_P2P_CONNECTOR_USER_GUIDE.md b/docs/npu/CAM_P2P_CONNECTOR_USER_GUIDE.md index c614b83b3..5fc067e5e 100644 --- a/docs/npu/CAM_P2P_CONNECTOR_USER_GUIDE.md +++ b/docs/npu/CAM_P2P_CONNECTOR_USER_GUIDE.md @@ -64,7 +64,7 @@ decoding. The pair requires static EP, PP=PCP=DCP=1, and each role's ranks equal to DP x TP. Sequence parallelism follows the model contract; synchronous DeepSeek-V2/V3 rejects effective model SP. Native MRV2 owns request slicing, capture, graph selection, replay and output merge. -AFD supplies stage-specific Ascend metadata and CAMP2P control/communication. The [NPU testing guide](TESTING.md) separates executor +AFD supplies stage-specific Ascend metadata and CAMP2P control/communication. The [NPU testing guide](../../tests/npu/README.md) separates executor checks from model correctness evidence. DBO is configured with vLLM CLI flags, not inside the `afd` object: @@ -132,27 +132,10 @@ Pass AFD configuration through vLLM's `--additional-config` option under the } ``` -### Fields - -| Field | Type | Default | Meaning and constraints | -| --- | --- | --- | --- | -| `role` | `"attention" \| "ffn"` | `"attention"` | Role owned by this process. | -| `connector` | `str` | `"P2pNcclAFDConnector"` | Set to `CAMP2pAFDConnector` for this synchronous NPU path. | -| `host` | `str` | `"127.0.0.1"` | Non-empty rendezvous host shared by every rank. It must be reachable from all participating processes. If the Attention and FFN ranks run on different machines, set this field to the FFN machine's reachable IP address. | -| `port` | `int` | `1239` | AFD rendezvous port in `1..65535`. It is separate from the vLLM HTTP service ports. | -| `afd_process_group_timeout_s` | `int` | `120` | Positive AFD group rendezvous timeout in seconds; also used by the native NPU DBO startup barrier. It does not bound a polling device kernel. | -| `num_attention_ranks` | `int` | `1` | Total number of Attention worker ranks. Must be positive. | -| `num_ffn_ranks` | `int` | `1` | Total number of FFN worker ranks. Must be positive. | -| `compute_gate_on_attention` | `bool` | `false` | Controls whether the MoE gate is computed on the Attention side or the FFN side. Currently only `false` is supported. | -| `connector_extra_config` | `dict` | `{}` | CAMP2P-specific settings such as role-specific core counts and `quant_mode`. Unknown fields are rejected. | -| `async` / `async_dp` | `bool` | `false` | Must remain `false` for the current synchronous; Ascend async mode requires `CAMAsyncAFDConnector`. | - -Compatibility aliases are accepted for `afd_role`, `afd_connector`, -`afd_host`, `afd_port`. - -The connector factory derives each process's role rank from its global DP rank -and local PCP/TP coordinates before connector construction. Do not configure a -role rank. +Use the [common configuration reference](../configuration.md) for defaults +and aliases. CAMP2P requires `async_dp=false`; `compute_gate_on_attention` may be `true` on +both roles for DeepSeek-V4 at TP=1 without DBO or ubatching. For multi-machine placement, use +the FFN machine's reachable IP as `host`. ### CAMP2P `connector_extra_config` @@ -161,7 +144,7 @@ role rank. | `core_num` | `int` | `8` | Positive default AIV core count for both roles. | | `attn_core_num` / `ffn_core_num` | `int` | unset | Positive role-specific override for `core_num`. | | `quant_mode` | `int` | `0` | CAM quantization mode. The current runtime supports only `0`. | -| `compute_gate_on_attention` | `bool` | `false` | Connector-local operator option; must remain `false`, as must the common AFD field. | +| `compute_gate_on_attention` | `bool` | `false` | Kept for backwards compatibility with older configurations and ignored for routing placement: transport derives the gate from the payload it actually carries. Omit it when deploying and set the common `afd.compute_gate_on_attention` field instead. | | `hccl_buffer_size` | `int` | unset | Positive CAMP2P HCCL buffer size in MB. The override applies to every connector-owned `afd*` communication domain and the FFN-side `afd_moe` domain; the Gloo control group is unaffected. When unset, HCCL uses `HCCL_BUFFSIZE`, then its built-in default. Use the same value on all members of each HCCL domain. | The per-domain setting avoids increasing buffers for unrelated TP, DP, or EP @@ -251,6 +234,80 @@ vllm serve /path/to/model \ }' ``` +## Single-node `4A4F` synchronous gate-on-Attention example (DeepSeek-V4, TP=1) + +This is the supported Attention-side gate deployment: synchronous CAMP2P, DeepSeek-V4, TP=1, +BF16 transport (`quant_mode: 0`), no DBO and no ubatching, four Attention ranks on devices `0-3` +and four FFN ranks on devices `4-7`. + +Set `afd.compute_gate_on_attention` to `true` on **both** roles and leave the connector-local +`compute_gate_on_attention` key out of `connector_extra_config`; that key is kept for backwards +compatibility with older configurations and does not control routing. + +The FFN service uses devices `4,5,6,7` and starts first: + +```bash +ASCEND_RT_VISIBLE_DEVICES=4,5,6,7 VLLM_PLUGINS=ascend,afd VLLM_USE_V2_MODEL_RUNNER=1 \ +vllm serve /path/to/DeepSeek-V4 \ + --host 127.0.0.1 \ + --port 18010 \ + --data-parallel-size 4 \ + --tensor-parallel-size 1 \ + --enable-expert-parallel \ + --enforce-eager \ + --max-model-len 9216 \ + --max-num-seqs 128 \ + --max-num-batched-tokens 3072 \ + --additional-config '{ + "afd": { + "role": "ffn", + "connector": "CAMP2pAFDConnector", + "host": "127.0.0.1", + "port": 6239, + "num_attention_ranks": 4, + "num_ffn_ranks": 4, + "compute_gate_on_attention": true, + "connector_extra_config": { + "quant_mode": 0 + } + } + }' +``` + +The Attention service uses devices `0,1,2,3`: + +```bash +ASCEND_RT_VISIBLE_DEVICES=0,1,2,3 VLLM_PLUGINS=ascend,afd VLLM_USE_V2_MODEL_RUNNER=1 \ +vllm serve /path/to/DeepSeek-V4 \ + --host 127.0.0.1 \ + --port 18000 \ + --data-parallel-size 4 \ + --tensor-parallel-size 1 \ + --enable-expert-parallel \ + --enforce-eager \ + --max-model-len 9216 \ + --max-num-seqs 128 \ + --max-num-batched-tokens 3072 \ + --additional-config '{ + "afd": { + "role": "attention", + "connector": "CAMP2pAFDConnector", + "host": "127.0.0.1", + "port": 6239, + "num_attention_ranks": 4, + "num_ffn_ranks": 4, + "compute_gate_on_attention": true, + "connector_extra_config": { + "quant_mode": 0 + } + } + }' +``` + +Drop `--enforce-eager` and add `--compilation-config '{"cudagraph_mode": "FULL_DECODE_ONLY", "cudagraph_capture_sizes": [8, 16, 32, 64, 128]}'` to run the same pair with ACL graphs. +The configuration is rejected when `compute_gate_on_attention=true` is combined with DBO, +ubatching, TP > 1, a connector other than `CAMP2pAFDConnector`, or a target other than DeepSeek-V4. + For model-specific launch and validation cases, see the [E2E guide](../../tests/e2e/README.md). diff --git a/docs/npu/STANDALONE_FFN_SERVICE.md b/docs/npu/STANDALONE_FFN_SERVICE.md index 4fe8ef5d2..56cbc8d09 100644 --- a/docs/npu/STANDALONE_FFN_SERVICE.md +++ b/docs/npu/STANDALONE_FFN_SERVICE.md @@ -14,7 +14,7 @@ CausalLM, vLLM Worker, upstream ModelRunner, EngineCore, scheduler or KV cache. Attention weights are not loaded by F. - The existing Attention model, router, shared experts, padding, DBO hooks and `CAMAsyncAFDConnector` are reused. F supports layered eager or FULL Graph - execution, with optional two-stream scheduling. + execution. Configuration, expert mapping, parameter loading and quantization depend on vLLM/vLLM-Ascend. Other models, quantization schemes and device backends are @@ -32,7 +32,7 @@ CPU Service / host runner -> FFNBackend ``` The backend adapts the complete CAM runtime lifecycle. Communication, compute, -Graph and two-stream scheduling remain in that runtime. See the +Graph execution remains in that runtime. See the [FFN runtime design](../design/module/ffn_runtime.md#runtime-selection) for module responsibilities and ownership. @@ -40,7 +40,7 @@ module responsibilities and ownership. Keep the existing Attention launch and its `CAMAsyncAFDConnector` configuration. See the [connector guide](CAM_ASYNC_CONNECTOR_USER_GUIDE.md) and -[eager/FULL runtime guide](ASYNC_FFN_RUNTIME.md) for its existing contract. +[eager/FULL runtime guide](../design/npu/async_ffn_runtime.md) for its existing contract. It must use eager MRV2, Attention-side gating and dynamic quantization. Native Attention DBO uses the existing integration contract; do not enable DBO on F. @@ -54,35 +54,71 @@ afd_cam_vendor=/path/to/afd_plugin/_cann_ops_custom/vendors/afd-plugin ascend_native_vendor=/path/to/vllm_ascend/_cann_ops_custom/vendors/custom_transformer export ASCEND_CUSTOM_OPP_PATH="$ascend_native_vendor:$afd_cam_vendor${ASCEND_CUSTOM_OPP_PATH:+:$ASCEND_CUSTOM_OPP_PATH}" export VLLM_USE_V2_MODEL_RUNNER=1 -# F-only setting: keep this unset or 0 in the Attention environment. -export AFD_ASYNC_CAM_LAYERED_GMM=1 export SOC_VERSION=910c export VLLM_PLUGINS=ascend,afd,ascend_model,ascend_model_loader,ascend_kv_connector # Keep BATCH_SIZE_FACTOR, HCCL/network settings and receive capacity consistent # with the existing A launch and CAM runtime guide. +export ASCEND_RT_VISIBLE_DEVICES=12,13,14,15 torchrun --nproc-per-node=4 --master-addr=127.0.0.1 --master-port=29500 \ -m afd_plugin.ffn_faas serve /path/to/model --backend ascend-cam \ - --host 127.0.0.1 --port 29600 \ - --attention-ranks 12 --attention-tp 4 --ffn-ranks 4 \ - --devices 12,13,14,15 --max-num-batched-tokens 2048 \ - --full-graph --ffn-two-stream \ - --ffn-compute-aic-num 20 --ffn-compute-aiv-num 40 \ - --ffn-combine-aiv-num 8 --status-port 29700 + --trust-remote-code --dtype bfloat16 --quantization ascend \ + --enable-expert-parallel --data-parallel-size 4 \ + --max-num-batched-tokens 2048 --max-num-seqs 256 \ + --compilation-config '{"cudagraph_mode":"FULL"}' \ + --additional-config '{ + "afd": { + "connector":"CAMAsyncAFDConnector","role":"ffn","async_dp":true, + "host":"127.0.0.1","port":29600, + "num_attention_ranks":12,"num_ffn_ranks":4, + "compute_gate_on_attention":true, + "connector_extra_config":{ + "dynamicQuant":1,"attn_ranks_per_dp":4 + } + }, + "enable_prefill_mc2":false,"enable_force_eplb":false,"enable_cpu_binding":true, + "ascend_fusion_config":{"fusion_ops_gmmswigluquant":true} + }' --status-port 29700 ``` -`--devices` indexes the process's visible logical devices, not physical IDs. -When F has a restricted visibility list, use `--devices 0,1,2,3` instead. The -external launcher's `WORLD_SIZE` is F-only. CAM creates the separate -`[A0 ... A11, F0 ... F3]` world using the same host, port, group name, topology, -quantization, layer table and startup signatures as traditional F. Use a -different rendezvous port for the F-only compute world and CAM. +Engine arguments use vLLM's `FlexibleArgumentParser`, `EngineArgs.add_cli_args` +and `EngineArgs.from_cli_args`; definitions, JSON/dotted options, defaults and +configuration validation come from the pinned upstream stack. The selected +Ascend CLI imports that layer, while importing the Service, host runner or +platform adapter module remains CPU-only. No serving engine is constructed. + +Use `--additional-config` for the existing AFD topology, connector and HCCL +options. Engine capacities and compilation are passed through +without rebuilding a second configuration. Only backend selection, management +endpoints, status files and shutdown control add standalone options. The old +`--full-graph`, `--attention-ranks`, `--attention-tp`, `--ffn-ranks`, `--devices` +and standalone stream/budget flags have been removed; update launch scripts to +the standard options above. Device selection uses the native `--device-ids` +mapping or `ASCEND_RT_VISIBLE_DEVICES` plus the launcher's `LOCAL_RANK`. + +The external launcher's `WORLD_SIZE`, `--data-parallel-size` and +`additional_config.afd.num_ffn_ranks` must agree. An omitted +`--data-parallel-rank` is taken from `RANK`; an explicit value must match it. +The executor defaults to `external_launcher` because torchrun already owns +these processes. Set `VLLM_USE_V2_MODEL_RUNNER=1` as on A. All other engine +options retain upstream defaults; the example explicitly selects the qualified +W4A8 profile. Reusing the parser does not qualify additional models, speculative +loading or serving/KV features for this service. + +The host receive-loop join uses `--receiver-shutdown-timeout` (default: five +seconds). It is separate from vLLM's `--shutdown-timeout`, whose EngineCore +shutdown semantics do not govern this standalone receiver. + +CAM creates the separate `[A0 ... A11, F0 ... F3]` world using the same host, +port, group name, topology, quantization, layer table and startup signatures as +traditional F. Use a different rendezvous port for the F-only compute world +and CAM. `--max-num-batched-tokens` must match A's dispatch envelope; it is not the -maximum request context length. F creates configuration objects for native -loading, but allocates no request KV cache. Omit `--full-graph` to use the -mature layered eager path. Two-stream execution supports both eager and FULL; -omitted core budgets use the existing CAM defaults. FULL preserves paired -warmup, private graph pools, device layer selection and FIFO slot ownership. +maximum request context length. F allocates no request KV cache. To use eager, +replace `--compilation-config` with `--enforce-eager`. FULL preserves paired +warmup, private graph pools and device layer selection. FFN two-stream +execution and its core-budget options were removed in #550; older launch +commands must omit those options. ## Management and shutdown @@ -95,7 +131,7 @@ Each local F rank exposes a management endpoint on `--status-port + LOCAL_RANK`. weights and paired CAM startup are prepared. Neither proves progress of an outstanding device receive. - `POST /stop`, SIGINT and SIGTERM request a stop between completed batches. - The host joins its receive loop for `--shutdown-timeout` seconds. + The host joins its receive loop for `--receiver-shutdown-timeout` seconds. The shutdown timeout does not cancel a CAM operator or prove device quiescence. It applies after startup; native loading and rendezvous retain their existing diff --git a/docs/npu/TESTING.md b/docs/npu/TESTING.md deleted file mode 100644 index 3eabcc06f..000000000 --- a/docs/npu/TESTING.md +++ /dev/null @@ -1,172 +0,0 @@ -# NPU validation - -Configuration and source-contract tests can run on CPU; runtime contracts -depend on vLLM/vLLM-Ascend, and device tests also need the matching Ascend -runtime and hardware. Use vLLM 0.32.0 and the vLLM-Ascend source revision in -the [installation guide](../../README.md#ascend-npu-installation). -The [0.32 validation record](https://github.com/vllm-project/afd-plugin/pull/507) distinguishes measured -coverage from deferred qualification and records the graceful-exit failures. - -From the repository root, run: - -```bash -python3 -m pytest -q -ra tests/unit -m "not gpu and not vllm_runtime" -``` - -The suite covers the NPU worker and model-runner contracts, CAM connector -behavior, and module isolation. Tests that require unavailable NPU dependencies -are skipped in CPU-only CI, so a passing CPU run does not replace this check. - -`tests/conftest.py` establishes the vLLM-Ascend import order required by the -test suite. No manual module preloading is needed. - -When diagnosing a failure, run the failing pytest node by itself first. A test -that passes alone but fails in the full suite usually indicates leaked module -or monkeypatch state. Test fixtures and mocks must also follow the concrete -vLLM types and signatures used by the pinned runtime. - -## Native DBO contracts and device execution - -Run the focused contract checks in the matching runtime environment: - -```bash -python3 -m pytest -q -ra \ - tests/unit/compat/npu/test_runtime_config.py \ - tests/unit/config/test_validation.py \ - tests/unit/model_executor/test_async_cam_dbo.py \ - tests/unit/model_executor/test_async_cam_layout.py \ - tests/unit/v1/worker/test_npu_model_runner_v2.py \ - tests/unit/v1/worker/test_npu_ubatch_outputs.py \ - tests/unit/v1/worker/test_npu_mla_graph.py -``` - -The single-device executor test uses real NPU streams/events with isolated -subprocesses for failure injection: - -```bash -ASCEND_RT_VISIBLE_DEVICES=0 python3 -m pytest -q -ra tests/npu/test_mrv2_dbo.py -``` - -It checks native child handoff, stream/context restoration and peer failure -wakeup. It does not validate CAMP2P/CAM transport, checkpoint output, or model -accuracy. The recorded Lite two-stage result does not qualify every Async CAM -DBO configuration. For broader qualification, run model comparisons at -Attention DP1 and DP > 1, plain TP and model SP, with -uneven prefill lengths and requests crossing child boundaries. Enable DBO on -Attention only. Keep two ubatches and eager execution. - -For synchronous MRV2 CAMP2P, include eager two-stage execution at Attention -DP > 1, many-to-one padding, and graph+DBO. Plain MLA graph+DBO requires -`FULL_DECODE_ONLY`. Both existing Lite graph+DBO cases must report completed -joint two-stage replays during real requests on every participating A/F rank. -Dummy/profile/capture or padded-only stages do not qualify that gate. -`tests/npu/test_mrv2_graph_inputs.py` separately checks persistent input -addresses with real NPU graph reads and modeled builders; model correctness -still belongs to E2E. Record execution, response correctness, revisions and -teardown outcomes separately. - -The original DSV4 Async concurrent case keeps native DBO thresholds. Its ten -short prompts can remain unsplit while satisfying its exact-answer, `stop`, -overlap and physical-cleanup checks; graceful exit failed in the recorded run. -DSV4 DSA two-stage qualification remains deferred. Qualify it separately -with an above-threshold workload, retaining native defaults and requiring -completed nonempty stages on every Attention rank. - -## Operator and model gates - -Use [the operator harness](../../tests/npu/README.md) for isolated routed CAM -precision/performance and [model E2E cases](../design/module/e2e_testing.md) -for serving and accuracy. DSV4 synchronous A5/A3 profiles currently check -completion/liveness with `check_answer=False`; they are not correctness gates. -Any skipped native registration or numerical case means that capability was -not validated by that run. - -## Layered W4A8 validation - -The layered routed-expert operators are `gmm_swiglu_quant_v2_layered` (W13) and -`grouped_matmul_layered` (W2). Both select their layer on device, and their -`all_*` lists hold one element per layer rather than one per expert. - -Run the CPU contract and path-selection tests first: - -```bash -pytest -q tests/unit/test_envs.py \ - tests/unit/model_executor/test_async_cam_layered.py -``` - -In an environment with matching CANN and torch-npu versions, rebuild the AFD -extension and run the device suites. All four need one device and issue no -collective: - -```bash -SOC_VERSION=910c pytest -q tests/npu/test_gmm_layered.py -SOC_VERSION=910c pytest -q tests/npu/test_gmm_swiglu_quant_v2_layered.py -AFD_RUN_ASCEND_OP_RUNTIME=1 SOC_VERSION=910c LOCAL_RANK=0 \ - pytest -o addopts= -v -ra -s \ - tests/npu/test_async_cam_layered_w4a8.py \ - tests/npu/test_async_cam_layered_w8a8.py \ - tests/npu/test_gmm_swiglu_quant_v2_layered_precision.py \ - tests/npu/test_gmm_layered_device_group_list.py -``` - -The four device suites answer different questions and none is redundant: - -| Suite | Operator | Covers | -| --- | --- | --- | -| `test_async_cam_layered_w4a8.py` | W13 then W2 | The W4A8 `AsyncCAMLayeredExecutor` chain, limits 0 and 10 | -| `test_async_cam_layered_w8a8.py` | W13 then W2 | Native W8A8 reference, device layer selection, sparse/empty/capacity rows, exact eager/FULL replay | -| `test_gmm_swiglu_quant_v2_layered_precision.py` | W13 | Default A8W8/A8W4 precision and clipped A8W4/A4W4 | -| `test_gmm_layered_device_group_list.py` | W2 | A device-produced group list, both scenarios, 4 cases | - -`test_async_cam_layered_w4a8.py` uses two layers with different nonzero INT4 -weights, nonzero compensation, per-channel and per-group scales, interleaved -nonconsecutive layer IDs, and zero, single-row, uneven, and full-capacity expert -counts. It compares valid rows with an FP32 reference and perturbs the capacity -tail. The zero-limit cases retain the random FP32 reference. Positive-limit cases -use sparse, exactly representable matmuls that cross the gate upper bound and -both up bounds, avoiding FP16 rounding followed by cancellation in the W2 -reference. Both use the same output tolerance. The W13 precision suite also -checks quantized codes, scales and dequantized outputs against an independent CPU -golden, same-shape 0/10 switching, and graph replay with device-selected layers. -These tests do not check distributed CAM completion or compare actual checkpoint -results against the non-layered path. Recheck the tolerance and supported -configurations against the project's precision standard when the operator -changes. - -The two precision suites are described under "Layered operator precision" in -[`tests/npu/README.md`](../../tests/npu/README.md), including why A8W4 cannot use -the fused built-in as its reference and why each attribute combination needs its -own activation shape. Read that section before changing a tolerance or adding a -case: the tiling cache is keyed on shape and tiling key, and the two attributes -reach the tiling only as data. - -For end-to-end validation, compare the layered path with a recorded legacy-path -baseline from an earlier revision, using the same W4A8 or W8A8 checkpoint and -deployment configuration. Record the processed parameter -shapes, dtypes, and formats; software versions; actual-path startup logs; -accuracy; peak memory; and profiler traces. Cover completion on empty ranks, -multiple chunks of one layer, interleaved DP work, and async CAM ubatching. -Check that no layer-ID or count metadata D2H occurs between receive and send, -and measure the worker group synchronization separately. Performance gains -require a controlled A/B comparison. - -For operator-level performance, use -[`tests/npu/gmm_layered_performance.py`](../../tests/npu/gmm_layered_performance.py) -rather than an end-to-end comparison. It measures both launch paths - eager and -ACL graph - and reports the fusion contribution and the graph contribution -separately, which an end-to-end run cannot separate. Anchors measured on Ascend -910C at `M = 98304`, `C = 262144`: W13 reaches about 1.80x the concatenated -baseline on the steady-state per-layer figure, `chain` about 1.46x, and W2 sits -at parity once the graph removes its launch overhead. Re-measure rather than -reuse these numbers; they are a starting point for a regression check, not a -claim about a different shape, topology, or CANN version. - -The current ACLNN entry validates `layer_index` by its one-element int64 -contract, including a slice with nonzero storage offset. The GMM path accepts -the logical 3D A8W8 weight view backed by 5D NZ storage as well as the W4A8 -layout. Rebuild the extension and vendor library together after these native -changes; Meta-only success does not exercise the ACLNN entry or kernel. -Layered GMM is allowed with MRV2 pairing, but it still runs in the independent -eager FFN runner. The current shared-expert layout also has focused coverage -in `tests/unit/model_executor/models/test_deepseek_v4_shared_experts.py` and -`test_deepseek_v2_sequence_parallel.py`. diff --git a/docs/npu/TROUBLESHOOTING.md b/docs/npu/TROUBLESHOOTING.md index 5b9c7a090..740d10154 100644 --- a/docs/npu/TROUBLESHOOTING.md +++ b/docs/npu/TROUBLESHOOTING.md @@ -169,5 +169,5 @@ additionally requires Attention DP > 1, while Async CAM permits DP1. An `NPU DBO peer failed` or `Async CAM receive failed` error means the host-side handoff was aborted. Restart the A/F worker group; waking the peer does not recover a partially used device window. Use the -[executor tests](TESTING.md#native-dbo-contracts-and-device-execution) to +[executor tests](../../tests/npu/README.md#native-dbo-contracts-and-device-execution) to separate thread/stream failures from communication and model correctness. diff --git a/docs/npu/LAYERED_W2_MDL_VALIDATION.md b/docs/reports/layered_w2_mdl_validation.md similarity index 100% rename from docs/npu/LAYERED_W2_MDL_VALIDATION.md rename to docs/reports/layered_w2_mdl_validation.md diff --git a/docs/support.md b/docs/support.md new file mode 100644 index 000000000..d45f44791 --- /dev/null +++ b/docs/support.md @@ -0,0 +1,78 @@ +# Supported configurations + +AFD is experimental. Registration and accepted configuration do not establish +hardware qualification; evidence below is scoped to its model, topology and +source revision. See [installation](installation.md) for the pinned runtime stack. + +## Models + +| Model family | Execution scope | +| --- | --- | +| DeepSeek V2 / V3 / V3.2 | CUDA P2P and synchronous Ascend paths; DeepEP M2N admits V2 only. Each role loads its required components. | +| DeepSeek V4 Flash | CUDA P2P or Ascend CAMP2P / Async CAM; backend-specific limits are in the matrix below. | +| Qwen3 MoE | CUDA; P2P gates on FFN, DeepEP M2N gates on Attention. Rejects sequence-parallel MoE, EPLB, PP, speculative decoding, LoRA and NPU. | +| Qwen3.5 / Qwen3.6 MoE | CUDA text-only with `--language-model-only`; evidence covers Qwen3.6-35B-A3B native DP4/TP1/EP4 and synchronous AFD 2A1F eager/graph. GDN DBO, Attention-side routing, PP, async and multi-node execution are unsupported; NPU/multimodal are unsupported or unverified, quantization is unverified, and no performance claim is made. | + +Architecture aliases and wrapper ownership are recorded in +[model integration](design/module/model_integration.md#model-registration). +For deployment commands, use the [recipe index](../recipe/README.md). + +## Tested runtime matrix + +This table separates accepted runtime configurations from validation evidence. +Both backends require ModelRunnerV2. A test case or a historical result does not +qualify every model/topology combination on the current target. + +| Platform/path | Execution | Ubatching | Routing/quantization limits | Evidence | +| --- | --- | --- | --- | --- | +| CUDA V2 + `P2pNcclAFDConnector` | Eager or `FULL_DECODE_ONLY` native V2 CUDA Graph | Exactly two native microbatches; Attention DP1 or DP > 1 | Model-specific gate placement; PP/CP, elastic EP, EPLB, SP MoE, and compile SP rejected; role ranks equal DP x TP | Source-specific GPU validation is linked below; coverage is limited to the recorded model/topology combinations. DeepSeek-V2-Lite has known DBO shutdown failures. Qwen3.5/3.6 GDN DBO is unsupported. | +| CUDA V2 + `DeepEPM2NAFDConnector` | Eager or `FULL_DECODE_ONLY` native V2 CUDA Graph; SM90 or newer | Exactly two native microbatches, enabled on both roles; Attention DP1 or DP > 1, TP1 or TP2 | Qwen3 MoE (all layers MoE, no shared expert) or DeepSeek-V2; `compute_gate_on_attention=true`; FFN TP1 with EP over all FFN ranks, linear placement, `allgather_reducescatter` all2all and the Triton MoE weight layout; BF16 or block-128 FP8 payload; EPLB, PP/CP and SP MoE rejected | Historical vLLM 0.30 Qwen3-30B-A3B 2A2F GSM8K E2E for eager, graph, graph + DBO, and Attention TP2 eager/graph on H100 and H200, and for eager, graph and graph + DBO on GB200; two-node 2A2F eager, graph and graph + DBO with Attention and FFN on separate nodes. DeepSeek-V2-Lite: [E2E definitions](../tests/e2e/models/deepseek_v2_lite/test_deepseek_v2_lite.py). [GPU connector device tests](../tests/gpu/README.md) | +| Ascend V2 + `CAMP2pAFDConnector` | Eager, `FULL`, or `FULL_DECODE_ONLY` native V2 ACL Graph | Eager or graph native DBO with two ubatches and Attention DP > 1; plain MLA graph DBO requires `FULL_DECODE_ONLY` | `compute_gate_on_attention` may be `false`, or `true` for DeepSeek-V4 at TP=1 without DBO/ubatching; PP/CP, elastic EP and EPLB rejected; model-specific SP policy; role ranks equal DP x TP | Historical 0.30 V2-Lite 2A2F FULL 300-question comparison and 2A1F functional checks; current 0.32 Lite evidence is in the [validation record](https://github.com/vllm-project/afd-plugin/pull/507) | +| Ascend V2 + `CAMAsyncAFDConnector` | Eager Attention; layered FFN eager or FULL | Two native DBO microbatches on Attention, including DP1 and with FFN FULL; FFN stays outside the native DBO schedule | Experimental DeepSeek-V4 W4A8/W8A8 on A3/910C; `async_dp=true`, `compute_gate_on_attention=true`, and `dynamicQuant=1`; PP/CP unsupported; plugin-owned async CAM ops required | Historical DSV4 layered comparisons and limitations are in the [validation record](https://github.com/vllm-project/afd-plugin/pull/507); they predate removal of the legacy path | + +The 300-question Ascend evidence in this table is historical. The 0.32 +[validation record](https://github.com/vllm-project/afd-plugin/pull/507) separately lists the +eight Lite 100-question cells, W4A8 fixture, and observed shutdown failures. + +The Async CAM row supports DeepSeek-V4 with static EP, PP=PCP=DCP=1, +and Attention-local SP. DeepSeek-V2/V2-Lite are no longer supported by +Async CAM; standard DeepSeek-V3.2 was already rejected by model admission. +Finegrained TP, explicit shared-expert DP configuration, and native +MTP and draft-model paths other than DeepSeek-V4 DSpark are excluded. DSpark +keeps Attention/native proposal eager and permits independent layered FFN FULL; +it does not support DBO. +FFN remains connector-driven. The row's upgrade evidence retains its original +scope; MRV2 deployment settings are documented in the +[CAM async guide](npu/CAM_ASYNC_CONNECTOR_USER_GUIDE.md#modelrunnerv2-deployment). + +Both backends target vLLM `0.32.0` (`5df63a72`); Ascend additionally targets +vLLM-Ascend `bd42eea0`. Source-specific +H20 model, topology, execution-mode and accuracy results are recorded in +the [GPU upgrade validation matrix](https://github.com/vllm-project/afd-plugin/pull/501) +and [runtime-fix validation matrix](https://github.com/vllm-project/afd-plugin/pull/502). +Those records identify the exact tested revisions. Earlier GPU results on +vLLM 0.30 do not establish vLLM 0.32 hardware qualification. + +## Validation limits + +- GPU accuracy evidence covers selected first-300 GSM8K comparisons, not + full-dataset accuracy, numerical equivalence, or performance improvement. +- DeepSeek-V2-Lite TP2 Graph + DBO and DP2 eager DBO can require force + killing during shutdown. Inference success does not imply graceful exit. +- Native V4 DBO requires a compatible upstream communication backend. The + recorded native V4 accuracy reference uses Graph without DBO. +- NPU 0.32 evidence is limited to the configurations and checks in the + [validation record](https://github.com/vllm-project/afd-plugin/pull/507); the seven AFD Lite + cases failed graceful exit. Historical NPU results do not qualify 0.32. + +Historical [NPU validation](https://github.com/vllm-project/afd-plugin/pull/425#issuecomment-6063910923) +and [DBO repair](https://github.com/vllm-project/afd-plugin/pull/430#issuecomment-6063840311) +records retain their original versions and sample counts. Synchronous DSV4 +profiles remain smoke-only, with known incorrect A5 concurrent answers; see +their [E2E limits](../tests/e2e/README.md#dsv4-flash-sync-camp2p-concurrent-requests-local-4-or-16-npus). +GPU/NPU rank topology and connector resource rules remain owned by +[connector contracts](design/module/connector_contracts.md). + +The repository does not record a canonical CUDA container or assert an +official tag for the Ascend source pin `bd42eea0`; environment evidence is +not an authoritative package tag. diff --git a/pyproject.toml b/pyproject.toml index 4f1d7518a..2687ed5de 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@v1.1.0", ] [project.urls] diff --git a/recipe/README.md b/recipe/README.md index ad6543d85..694418b0d 100644 --- a/recipe/README.md +++ b/recipe/README.md @@ -11,23 +11,6 @@ scope; the V3.2 recipes below retain their historical versions. Recipes are organized by hardware backend, connector, and model: -```text -recipe/ -├── gpu/ -│ ├── DeepEPM2NAFDConnector/ -│ │ ├── qwen3_235b_a22b_fp8/ -│ │ └── qwen3_30b_a3b/ -│ └── P2pNcclAFDConnector/ -│ └── deepseek_v2_lite/ -└── npu/ - ├── CAMAsyncAFDConnector/ - │ └── deepseek_v3_2/ - └── CAMP2pAFDConnector/ - └── deepseek_v3_2/ -``` - -Directory names follow these conventions: - - Hardware backend: `gpu` or `npu`. - Connector: the exact connector class name, such as `P2pNcclAFDConnector`, `DeepEPM2NAFDConnector`, `CAMP2pAFDConnector`, or `CAMAsyncAFDConnector`. @@ -39,11 +22,13 @@ 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+) | `DeepEPM2NAFDConnector` | Qwen3-30B-A3B | Decode | Same 2A2F layouts as the Qwen3 MoE DeepEP M2N E2E cases | [Launch examples](gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/README.md) | -| GPU (SM90+) | `DeepEPM2NAFDConnector` | Qwen3-235B-A22B-FP8 | Decode | 4A2F FP8-payload configuration run on H200 | [Launch examples](gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/README.md) | +| 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` | Qwen3-30B-A3B | Prefill / decode; decode-only benchmark | 2A2F graph, eager, graph+DBO and eager Attention-TP2 recipes; multi-pod layouts remain configurable through `LAYOUT` | [Graph, eager, DBO](gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/README.md) | +| GPU (SM90+ for `DeepEPM2NAFDConnector`) | `DeepEPM2NAFDConnector` | Qwen3-235B-A22B-FP8 | Prefill / decode; decode-only benchmark | FP8-payload 4A2F eager, graph and DBO recipes, plus two-pod 4A1F DBO; H200-class FFN GPUs | [Graph, eager, DBO](gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/README.md) | +| GPU | `P2pNcclAFDConnector` | Qwen3.5-122B-A10B-FP8 | Decode | Three-pod graph recipe (`2A0F,2A0F,0A2F`); other layouts remain configurable through `LAYOUT` | [Graph](gpu/P2pNcclAFDConnector/qwen3_5_122b_a10b_fp8/prefill_decode_colocation/multipod_2a_2a_2f_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) | +| Ascend NPU | `CAMAsyncAFDConnector` | DeepSeek-V4-Flash W4A8/W8A8 DSpark | Separate P / decode, or colocated | v0.32; F layered eager/FULL. W4A8 target/draft native decode Graph: A DP1/TP4, F EP4, with Attention-local SP and with/without PD; native full-draft P TP4 for PD. Earlier eager W4A8 PD: A DP2/TP1, F EP4. W8A8 eager A/F: A DP2/TP1, F EP8; A Graph and PD not qualified | [DSpark A/F](npu/CAMAsyncAFDConnector/deepseek_v4_flash/README.md) | Open the model-level README before running a recipe. It documents the required hardware and runtime baseline, topology, environment variables, launch order, @@ -57,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 5151e899a..7ca658db1 100644 --- a/recipe/gpu/DeepEPM2NAFDConnector/README.md +++ b/recipe/gpu/DeepEPM2NAFDConnector/README.md @@ -2,33 +2,38 @@ 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 + +Recipe names encode their layout, DP/TP shape, and mode. DeepSeek-V2-Lite and +Qwen3-30B-A3B provide `2a2f_{eager,graph,graph_dbo}_dp2tp1.sh` plus +`2a2f_eager_dp1tp2.sh`; Qwen3-235B-A22B-FP8 provides `4a2f` eager, graph and +DBO recipes plus `multipod_4a1f_4a1f_dbo.sh`. Each model README lists its +exact filenames and defaults. ```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 +├── 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 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 -- 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 DeepSeek and Qwen3-30B default `2A2F` layout uses + 4 GPUs; Qwen3-235B defaults to `4A2F` on 6 GPUs. 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 @@ -41,29 +46,100 @@ 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 -Run any script from the repository root, for example: +`LAYOUT` 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). + +## Running ```bash export MODEL_PATH=/path/model_weights/Qwen3-30B-A3B bash recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/2a2f_graph_dp2tp1.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 first start on a node compiles transport and DeepGEMM kernels, so it is -slower than later starts. +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 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. +The **`deploy-afd-k8s`** skill (`.agents/skills/deploy-afd-k8s/SKILL.md`) +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 + +```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 (`{}` 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='{}' \ + bash recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/2a2f_graph_dbo_dp2tp1.sh + +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 +``` + +Drive the Attention port (18000). The benchmark works in every layout and +mode; the connector is +[`tools/benchmarks/decode_bench.py`](../../../tools/benchmarks/decode_bench.py). + +## Configuration notes + +- Both roles set `compute_gate_on_attention: true` and run with + `--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 + 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..6f04e2dcc 100644 --- a/recipe/gpu/DeepEPM2NAFDConnector/deepseek_v2_lite/README.md +++ b/recipe/gpu/DeepEPM2NAFDConnector/deepseek_v2_lite/README.md @@ -1,46 +1,42 @@ # 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 +├── 2a2f_eager_dp2tp1.sh # eager, FFN DP2/TP1, Attention DP2/TP1 +├── 2a2f_graph_dp2tp1.sh # graph, FFN DP2/TP1, Attention DP2/TP1 +├── 2a2f_graph_dbo_dp2tp1.sh # graph + DBO, FFN DP2/TP1, Attention DP2/TP1 +└── 2a2f_eager_dp1tp2.sh # eager, FFN DP2/TP1, Attention DP1/TP2 ``` -## 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_dp1tp2.sh b/recipe/gpu/DeepEPM2NAFDConnector/deepseek_v2_lite/prefill_decode_colocation/2a2f_eager_dp1tp2.sh new file mode 100644 index 000000000..653d9b3ac --- /dev/null +++ b/recipe/gpu/DeepEPM2NAFDConnector/deepseek_v2_lite/prefill_decode_colocation/2a2f_eager_dp1tp2.sh @@ -0,0 +1,18 @@ +#!/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, 2A2F eager mode with +# Attention TP2: FFN DP2 x TP1 and Attention DP1 x TP2. +# 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/DeepSeek-V2-Lite} +ATTN_TP=${ATTN_TP:-2} FFN_TP=${FFN_TP:-1} +MODE=eager +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/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 index 3da8f033f..36b00f0bf 100644 --- 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 @@ -1,55 +1,18 @@ +#!/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, 2A2F eager mode: +# FFN DP2 x TP1 and Attention DP2 x TP1. +# Settings and running in a pod: recipe/gpu/common/README.md +# shellcheck disable=SC1090,SC2034 # the sourced launcher reads these variables -# 2A2F eager: Attention DP2 x TP1 on GPUs 0,1; FFN DP2 on GPUs 2,3. - +AFD_CONNECTOR=${AFD_CONNECTOR:-DeepEPM2NAFDConnector} +LAYOUT=${LAYOUT:-2A2F} 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 & +ATTN_TP=${ATTN_TP:-1} FFN_TP=${FFN_TP:-1} +MODE=eager +SERVE_ARGS=(--enable-ep-weight-filter --max-num-seqs 64 --max-num-batched-tokens 4096 --max-model-len 4096) -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/2a2f_graph_dbo_dp2tp1.sh b/recipe/gpu/DeepEPM2NAFDConnector/deepseek_v2_lite/prefill_decode_colocation/2a2f_graph_dbo_dp2tp1.sh index 0ea5a336e..10059bfa4 100644 --- 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 @@ -1,64 +1,19 @@ +#!/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, 2A2F FULL_DECODE_ONLY +# CUDA graphs and DBO: FFN DP2 x TP1 and Attention DP2 x TP1. +# Settings and running in a pod: recipe/gpu/common/README.md +# shellcheck disable=SC1090,SC2034 # the sourced launcher reads these variables -# 2A2F graph + DBO: Attention DP2 x TP1 on GPUs 0,1; FFN DP2 on GPUs 2,3. -# DBO must be enabled on both roles. - +AFD_CONNECTOR=${AFD_CONNECTOR:-DeepEPM2NAFDConnector} +LAYOUT=${LAYOUT:-2A2F} 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 & +ATTN_TP=${ATTN_TP:-1} FFN_TP=${FFN_TP:-1} +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) -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/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 index c7f1801ad..a68079120 100644 --- 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 @@ -1,57 +1,18 @@ +#!/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, 2A2F FULL_DECODE_ONLY +# CUDA graphs: FFN DP2 x TP1 and Attention DP2 x TP1. +# Settings and running in a pod: recipe/gpu/common/README.md +# shellcheck disable=SC1090,SC2034 # the sourced launcher reads these variables -# 2A2F FULL_DECODE_ONLY graph: Attention DP2 x TP1 on GPUs 0,1; FFN DP2 on GPUs 2,3. - +AFD_CONNECTOR=${AFD_CONNECTOR:-DeepEPM2NAFDConnector} +LAYOUT=${LAYOUT:-2A2F} 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 & +ATTN_TP=${ATTN_TP:-1} FFN_TP=${FFN_TP:-1} +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) -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/README.md b/recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/README.md index 5b5ec1e74..0ddc9ace5 100644 --- a/recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/README.md +++ b/recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/README.md @@ -1,58 +1,51 @@ # Qwen3-235B-A22B-FP8 with DeepEPM2NAFDConnector -Launch scripts for Qwen3-235B-A22B-FP8 on `DeepEPM2NAFDConnector` with vLLM -`v0.30.0`, using the FP8 dispatch payload. - -## Prerequisites - -- 6 SM90 or newer GPUs on one node with enough memory for the model. This - configuration was 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 - `/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. +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 -└── 4a2f_graph_dbo_dp4tp1.sh # 512 sequences per Attention rank, graph + DBO -``` - -## 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. -- 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. - -## 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}' +├── 4a2f_eager_dp4tp1.sh # eager, FFN DP2/TP1, Attention DP4/TP1 +├── 4a2f_graph_dp4tp1.sh # graph, FFN DP2/TP1, Attention DP4/TP1 +├── 4a2f_dbo_dp4tp1.sh # graph + DBO, FFN DP2/TP1, Attention DP4/TP1 +└── multipod_4a1f_4a1f_dbo.sh # 2 pods, graph + DBO, Attention DP8/TP1 ``` -The answer should name Paris. An FFN log line `Using TRITON Fp8 MoE backend` -confirms the expected expert-weight layout. +## Defaults + +| Setting | `4a2f_graph_dp4tp1.sh`, `4a2f_eager_dp4tp1.sh` | `4a2f_dbo_dp4tp1.sh`, `multipod_4a1f_4a1f_dbo.sh` | +| --- | --- | --- | +| `LAYOUT` | `4A2F` | `4A2F` / `4A1F,4A1F` | +| `--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. +- 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 - 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_dbo_dp4tp1.sh b/recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/4a2f_dbo_dp4tp1.sh new file mode 100644 index 000000000..1ebd1b144 --- /dev/null +++ b/recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/4a2f_dbo_dp4tp1.sh @@ -0,0 +1,19 @@ +#!/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, 4A2F FULL_DECODE_ONLY CUDA +# graphs and DBO: FFN DP2 x TP1 and Attention DP4 x TP1. +# 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:-4A2F} +MODEL_PATH=${MODEL_PATH:-/path/model_weights/Qwen3-235B-A22B-FP8} +CONNECTOR_CONFIG='"payload_format": "fp8_block128_fp32"' # FP8 payload +MODE=graph CUDAGRAPH_CAPTURE_SIZES=${CUDAGRAPH_CAPTURE_SIZES:-8,16,32,64,96,128,192,256,384,512} # ascending +SERVE_ARGS=(--enable-ep-weight-filter --max-num-seqs 512 --max-num-batched-tokens 3072 --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 "$@" diff --git a/recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/4a2f_eager_dp4tp1.sh b/recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/4a2f_eager_dp4tp1.sh new file mode 100644 index 000000000..e196f5d76 --- /dev/null +++ b/recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/4a2f_eager_dp4tp1.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-235B-A22B-FP8 on DeepEPM2NAFDConnector, 4A2F eager mode: +# FFN DP2 x TP1 and Attention DP4 x TP1. +# 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:-4A2F} +MODEL_PATH=${MODEL_PATH:-/path/model_weights/Qwen3-235B-A22B-FP8} +CONNECTOR_CONFIG='"payload_format": "fp8_block128_fp32"' # FP8 payload +MODE=eager +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/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 index 04491ad2a..c4b25f040 100644 --- 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 @@ -1,61 +1,18 @@ +#!/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, 4A2F FULL_DECODE_ONLY CUDA +# graphs: FFN DP2 x TP1 and Attention DP4 x TP1. +# Settings and running in a pod: recipe/gpu/common/README.md +# shellcheck disable=SC1090,SC2034 # the sourced launcher reads these variables -# 4A2F graph, FP8 payload: Attention DP4 x TP1 on GPUs 0-3; FFN DP2 on GPUs 4,5. - +AFD_CONNECTOR=${AFD_CONNECTOR:-DeepEPM2NAFDConnector} +LAYOUT=${LAYOUT:-4A2F} 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 256 \ - --max-num-batched-tokens 3072 \ - --max-model-len 4096 \ - --max-cudagraph-capture-size 256 \ - --compilation-config '{"cudagraph_mode": "FULL_DECODE_ONLY", "cudagraph_capture_sizes": [8, 16, 32, 64, 96, 128, 192, 256]}' \ - --cpu-distributed-timeout-seconds 3600 \ - --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 256 \ - --max-num-batched-tokens 3072 \ - --max-model-len 4096 \ - --max-cudagraph-capture-size 256 \ - --compilation-config '{"cudagraph_mode": "FULL_DECODE_ONLY", "cudagraph_capture_sizes": [8, 16, 32, 64, 96, 128, 192, 256]}' \ - --cpu-distributed-timeout-seconds 3600 \ - --host 127.0.0.1 \ - --port 18000 > attn.log 2>&1 & +CONNECTOR_CONFIG='"payload_format": "fp8_block128_fp32"' # FP8 payload +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) -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/multipod_4a1f_4a1f_dbo.sh b/recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/multipod_4a1f_4a1f_dbo.sh new file mode 100644 index 000000000..32d786a08 --- /dev/null +++ b/recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/multipod_4a1f_4a1f_dbo.sh @@ -0,0 +1,19 @@ +#!/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, two 4A1F pods with +# FULL_DECODE_ONLY CUDA graphs and DBO: Attention DP8 x TP1, FFN DP2 x TP1. +# 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:-4A1F,4A1F} +MODEL_PATH=${MODEL_PATH:-/path/model_weights/Qwen3-235B-A22B-FP8} +CONNECTOR_CONFIG='"payload_format": "fp8_block128_fp32"' # FP8 payload +MODE=graph CUDAGRAPH_CAPTURE_SIZES=${CUDAGRAPH_CAPTURE_SIZES:-8,16,32,64,96,128,192,256,384,512} # ascending +SERVE_ARGS=(--enable-ep-weight-filter --max-num-seqs 512 --max-num-batched-tokens 3072 --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 "$@" diff --git a/recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/README.md b/recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/README.md index 184806b5e..817963a4d 100644 --- a/recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/README.md +++ b/recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/README.md @@ -1,59 +1,43 @@ # Qwen3-30B-A3B with DeepEPM2NAFDConnector -Launch scripts for Qwen3-30B-A3B (BF16) on `DeepEPM2NAFDConnector` with vLLM -`v0.30.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)). - -## Prerequisites - -- 4 SM90 or newer GPUs on one node. -- The environment described in the - [connector recipes README](../README.md#prerequisites). -- 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. +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 -├── 2a2f_graph_dbo_dp2tp1.sh # Attention DP2 x TP1, graph + DBO -└── 2a2f_graph_dp1tp2.sh # Attention DP1 x TP2, FULL_DECODE_ONLY graph +├── 2a2f_eager_dp2tp1.sh # eager, FFN DP2/TP1, Attention DP2/TP1 +├── 2a2f_graph_dp2tp1.sh # graph, FFN DP2/TP1, Attention DP2/TP1 +├── 2a2f_graph_dbo_dp2tp1.sh # graph + DBO, FFN DP2/TP1, Attention DP2/TP1 +└── 2a2f_eager_dp1tp2.sh # eager, FFN DP2/TP1, Attention DP1/TP2 ``` -## Configuration 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. -- `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. -- 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. - -## 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}' -``` +## 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 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. + +## Notes -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). +- 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_dp1tp2.sh b/recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/2a2f_eager_dp1tp2.sh new file mode 100644 index 000000000..a70d42850 --- /dev/null +++ b/recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/2a2f_eager_dp1tp2.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 DeepEPM2NAFDConnector, 2A2F eager mode with +# Attention TP2: FFN DP2 x TP1 and Attention DP1 x TP2. +# 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} +ATTN_TP=${ATTN_TP:-2} FFN_TP=${FFN_TP:-1} +MODE=eager +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/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 index 18b37e29a..c563ef9da 100644 --- 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 @@ -1,53 +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 DeepEPM2NAFDConnector, 2A2F eager mode: +# FFN DP2 x TP1 and Attention DP2 x TP1. +# Settings and running in a pod: recipe/gpu/common/README.md +# shellcheck disable=SC1090,SC2034 # the sourced launcher reads these variables -# 2A2F eager: Attention DP2 x TP1 on GPUs 0,1; FFN DP2 on GPUs 2,3. - +AFD_CONNECTOR=${AFD_CONNECTOR:-DeepEPM2NAFDConnector} +LAYOUT=${LAYOUT:-2A2F} 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 & +ATTN_TP=${ATTN_TP:-1} FFN_TP=${FFN_TP:-1} +MODE=eager +SERVE_ARGS=(--enable-ep-weight-filter --max-num-seqs 64 --max-num-batched-tokens 4096 --max-model-len 4096) -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/2a2f_graph_dbo_dp2tp1.sh b/recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/2a2f_graph_dbo_dp2tp1.sh index 1736b681c..f08636036 100644 --- 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 @@ -1,62 +1,19 @@ +#!/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, 2A2F FULL_DECODE_ONLY CUDA +# graphs and DBO: FFN DP2 x TP1 and Attention DP2 x TP1. +# Settings and running in a pod: recipe/gpu/common/README.md +# shellcheck disable=SC1090,SC2034 # the sourced launcher reads these variables -# 2A2F graph + DBO: Attention DP2 x TP1 on GPUs 0,1; FFN DP2 on GPUs 2,3. -# DBO must be enabled on both roles. - +AFD_CONNECTOR=${AFD_CONNECTOR:-DeepEPM2NAFDConnector} +LAYOUT=${LAYOUT:-2A2F} 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 & +ATTN_TP=${ATTN_TP:-1} FFN_TP=${FFN_TP:-1} +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) -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/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 index 2660ec88d..85edcd752 100644 --- 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 @@ -1,55 +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 DeepEPM2NAFDConnector, 2A2F FULL_DECODE_ONLY CUDA +# graphs: FFN DP2 x TP1 and Attention DP2 x TP1. +# Settings and running in a pod: recipe/gpu/common/README.md +# shellcheck disable=SC1090,SC2034 # the sourced launcher reads these variables -# 2A2F FULL_DECODE_ONLY graph: Attention DP2 x TP1 on GPUs 0,1; FFN DP2 on GPUs 2,3. - +AFD_CONNECTOR=${AFD_CONNECTOR:-DeepEPM2NAFDConnector} +LAYOUT=${LAYOUT:-2A2F} 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 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 > attn.log 2>&1 & +ATTN_TP=${ATTN_TP:-1} FFN_TP=${FFN_TP:-1} +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) -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..8f958b62c 100644 --- a/recipe/gpu/P2pNcclAFDConnector/README.md +++ b/recipe/gpu/P2pNcclAFDConnector/README.md @@ -11,36 +11,31 @@ 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/ # explicit 2A2F, 4A4F and two-pod DBO recipes │ └── 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_5_122b_a10b_fp8/prefill_decode_colocation/ # multipod_2a_2a_2f_graph.sh ``` 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 ``` -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 +44,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. +The skill takes 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. +plus the image, model id, and PVC name, and deploys one pod per `LAYOUT` +entry (for example the DeepSeek multipod recipe's default `2A0F,0A2F`, 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 d4246f38d..accd39a42 100644 --- a/recipe/gpu/P2pNcclAFDConnector/deepseek_v2_lite/README.md +++ b/recipe/gpu/P2pNcclAFDConnector/deepseek_v2_lite/README.md @@ -3,7 +3,7 @@ End-to-end launch scripts for running DeepSeek-V2-Lite with the AFD (Attention-FFN Disaggregation) plugin targeting vLLM `v0.32.0`. These are deployment examples, not v0.32 recipe validation. See the -[runtime matrix](../../../../docs/design/module/execution_platforms.md#tested-runtime-matrix) +[runtime matrix](../../../../docs/support.md#tested-runtime-matrix) for tested E2E configurations. > [!NOTE] @@ -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,14 @@ 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 + ├── 2a2f_eager_dbo_dp2tp1.sh # 2A2F, eager + DBO, DP2/TP1 + ├── 2a2f_graph_dbo_dp2tp1.sh # 2A2F, graph + DBO, DP2/TP1 + ├── 2a2f_eager_dbo_dp1tp2.sh # 2A2F, eager + DBO, DP1/TP2 + ├── 2a2f_graph_dbo_dp1tp2.sh # 2A2F, graph + DBO, DP1/TP2 + ├── 4a4f_eager_dbo_dp2tp2.sh # 4A4F, eager + DBO, DP2/TP2 + ├── 4a4f_graph_dbo_dp2tp2.sh # 4A4F, graph + DBO, DP2/TP2 + └── multipod_2a_2f_graph_dbo_dp1tp2.sh # 2A0F,0A2F, graph + DBO, DP1/TP2 ``` ### 1. Prefill/Decode Disaggregation — `1a1f` @@ -58,31 +58,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 and +each script names its default DP/TP shape (FFN first, Attention second, and +Attention HTTP port 18000): + +| Script suffix | Attention / FFN | GPUs | +| --- | --- | --- | +| `2a2f_*_dp2tp1` | DP2 x TP1 each | 4 | +| `2a2f_*_dp1tp2` | DP1 x TP2 each | 4 | +| `4a4f_*_dp2tp2` | DP2 x TP2 each | 8 | +| `multipod_2a_2f_*_dp1tp2` | DP1 x TP2 each, one pod per role | 2 + 2 | ## Running @@ -96,7 +82,6 @@ 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 ``` @@ -122,7 +107,8 @@ 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 +116,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 +126,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 +140,7 @@ 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. +The `graph_dbo` recipes capture `FULL_DECODE_ONLY` CUDA graphs of size `64` +(`CUDAGRAPH_CAPTURE_SIZES=64`); the `eager_dbo` recipes run +`--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 index 1cac638bc..8c92d7cf4 100644 --- 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 @@ -1,55 +1,19 @@ +#!/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, 2A2F eager mode and DBO: +# FFN DP1 x TP2 and Attention DP1 x TP2. +# 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} -export VLLM_USE_V2_MODEL_RUNNER=1 +ATTN_TP=${ATTN_TP:-2} FFN_TP=${FFN_TP:-2} +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) -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 +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/2a2f_eager_dbo_dp2tp1.sh b/recipe/gpu/P2pNcclAFDConnector/deepseek_v2_lite/prefill_decode_colocation/2a2f_eager_dbo_dp2tp1.sh index edbab3812..ba600f9e1 100644 --- 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 @@ -1,55 +1,19 @@ +#!/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, 2A2F eager mode and DBO: +# FFN DP2 x TP1 and Attention DP2 x TP1. +# 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} -export VLLM_USE_V2_MODEL_RUNNER=1 +ATTN_TP=${ATTN_TP:-1} FFN_TP=${FFN_TP:-1} +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) -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 +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/2a2f_graph_dbo_dp1tp2.sh b/recipe/gpu/P2pNcclAFDConnector/deepseek_v2_lite/prefill_decode_colocation/2a2f_graph_dbo_dp1tp2.sh index e72893d22..bda209d64 100644 --- 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 @@ -1,61 +1,19 @@ +#!/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, 2A2F FULL_DECODE_ONLY CUDA +# graphs and DBO: FFN DP1 x TP2 and Attention DP1 x TP2. +# 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} -export VLLM_USE_V2_MODEL_RUNNER=1 +ATTN_TP=${ATTN_TP:-2} FFN_TP=${FFN_TP:-2} +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) -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 +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/2a2f_graph_dbo_dp2tp1.sh b/recipe/gpu/P2pNcclAFDConnector/deepseek_v2_lite/prefill_decode_colocation/2a2f_graph_dbo_dp2tp1.sh index c775780db..7baa8c391 100644 --- 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 @@ -1,61 +1,19 @@ +#!/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, 2A2F FULL_DECODE_ONLY CUDA +# graphs and DBO: FFN DP2 x TP1 and Attention DP2 x TP1. +# 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} -export VLLM_USE_V2_MODEL_RUNNER=1 +ATTN_TP=${ATTN_TP:-1} FFN_TP=${FFN_TP:-1} +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) -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 +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/4a4f_eager_dbo_dp2tp2.sh b/recipe/gpu/P2pNcclAFDConnector/deepseek_v2_lite/prefill_decode_colocation/4a4f_eager_dbo_dp2tp2.sh index e2cbf9e71..28d718f90 100644 --- 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 @@ -1,55 +1,19 @@ +#!/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, 4A4F eager mode and DBO: +# FFN DP2 x TP2 and Attention DP2 x TP2. +# 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:-4A4F} MODEL_PATH=${MODEL_PATH:-/path/model_weights/DeepSeek-V2-Lite} -export VLLM_USE_V2_MODEL_RUNNER=1 +ATTN_TP=${ATTN_TP:-2} FFN_TP=${FFN_TP:-2} +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) -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 +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/4a4f_graph_dbo_dp2tp2.sh b/recipe/gpu/P2pNcclAFDConnector/deepseek_v2_lite/prefill_decode_colocation/4a4f_graph_dbo_dp2tp2.sh index cd87ccbb4..2ecb4e93f 100644 --- 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 @@ -1,61 +1,19 @@ +#!/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, 4A4F FULL_DECODE_ONLY CUDA +# graphs and DBO: FFN DP2 x TP2 and Attention DP2 x TP2. +# 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:-4A4F} MODEL_PATH=${MODEL_PATH:-/path/model_weights/DeepSeek-V2-Lite} -export VLLM_USE_V2_MODEL_RUNNER=1 +ATTN_TP=${ATTN_TP:-2} FFN_TP=${FFN_TP:-2} +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) -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 +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 index b1a1f0b14..0a5dd95c0 100644 --- 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 @@ -1,73 +1,19 @@ +#!/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, two pods with 2A0F,0A2F, +# FULL_DECODE_ONLY CUDA graphs and DBO: each role is DP1 x TP2. +# Settings and running in a pod: recipe/gpu/common/README.md +# shellcheck disable=SC1090,SC2034 # the sourced launcher reads these variables -# 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 - +AFD_CONNECTOR=${AFD_CONNECTOR:-P2pNcclAFDConnector} +LAYOUT=${LAYOUT:-2A0F,0A2F} 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 +ATTN_TP=${ATTN_TP:-2} FFN_TP=${FFN_TP:-2} +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) -wait +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 index ec63badbd..5582c73ee 100644 --- 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 @@ -1,100 +1,21 @@ +#!/usr/bin/env bash # 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 - +# +# 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} -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 +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) -wait +source "${AFD_RECIPE_LIB:-$(dirname "${BASH_SOURCE[0]}")/../../../common/afd_colocation.sh}" +afd_colocation_serve "$@" diff --git a/recipe/gpu/common/README.md b/recipe/gpu/common/README.md new file mode 100644 index 000000000..677e95df2 --- /dev/null +++ b/recipe/gpu/common/README.md @@ -0,0 +1,161 @@ +# 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, 2A2F graph + 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} +ATTN_TP=${ATTN_TP:-1} FFN_TP=${FFN_TP:-1} +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` | recipe | 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. | + +Each personalized recipe encodes the DP/TP shape named in its filename. The +`ATTN_TP` and `FFN_TP` defaults are therefore recipe-specific rather than +always `1`; set them only when intentionally changing that shape. + +`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/2a2f_graph_dp2tp1.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/2a2f_graph_dp2tp1.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 .../2a2f_graph_dp2tp1.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 .../2a2f_graph_dp2tp1.sh +LAYOUT=2A0F,0A2F POD=AFD_1 ATTN_HEAD_HOST= FFN_HEAD_HOST= bash .../2a2f_graph_dp2tp1.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 +} diff --git a/recipe/npu/CAMAsyncAFDConnector/deepseek_v4_flash/README.md b/recipe/npu/CAMAsyncAFDConnector/deepseek_v4_flash/README.md index f295f6452..fec550967 100644 --- a/recipe/npu/CAMAsyncAFDConnector/deepseek_v4_flash/README.md +++ b/recipe/npu/CAMAsyncAFDConnector/deepseek_v4_flash/README.md @@ -34,11 +34,14 @@ and network interface. A TP1/DP1 and F EP4 use five distinct DIEs. MODEL=/absolute/path/to/dsv4-flash-dspark-w4a8 LOCAL_IP=REPLACE_WITH_COMMUNICATION_IP INTERFACE=eth0 -NUM_A=1 +A_DP=1 +A_TP=1 +NUM_A=$((A_DP * A_TP)) NUM_F=4 A_DEVICES=0 F_DEVICES=4,5,6,7 FFN_EXEC_ARGS=(--enforce-eager) +A_EXEC_ARGS=(--enforce-eager) export VLLM_USE_V2_MODEL_RUNNER=1 VLLM_ASCEND_ENABLE_FLASHCOMM1=0 export BATCH_SIZE_FACTOR=0.00390625 @@ -61,14 +64,36 @@ COMMON_ARGS=( --no-disable-hybrid-kv-cache-manager ) DSPARK_CONFIG="{\"method\":\"dspark\",\"model\":\"$MODEL\",\"num_speculative_tokens\":5,\"enforce_eager\":true}" +A_DSPARK_CONFIG="$DSPARK_CONFIG" ``` -For DP2 set `NUM_A=2` and `A_DEVICES=0,1`. The W4A8 example uses F EP4. +For DP2/TP1 set `A_DP=2`, recompute `NUM_A`, and use `A_DEVICES=0,1`. The W4A8 example uses F EP4. For a W8A8 checkpoint, size F for its larger weights; the validated W8A8 configuration uses `NUM_F=8` and `F_DEVICES=8,9,10,11,12,13,14,15`. -Target and draft Attention use eager execution. F uses layered eager or layered -FULL, with one combined target/draft weight table. Layered eager can be slower -for small decode batches, so measure the chosen deployment. +Target and draft use matching Attention TP widths; F remains TP1 with EP. +Target Attention may use eager execution or the native decode graph. +`A_DSPARK_CONFIG` selects native draft eager/Graph execution on A; F's +`DSPARK_CONFIG` and expert execution settings remain independent. To enable both +Attention graphs, set the following before launching. Keep `enforce_eager=true` +in `A_DSPARK_CONFIG` when only the target graph is wanted. + +```bash +A_EXEC_ARGS=(--compilation-config '{"mode":0,"cudagraph_mode":"FULL_DECODE_ONLY"}') +A_DSPARK_CONFIG="{\"method\":\"dspark\",\"model\":\"$MODEL\",\"num_speculative_tokens\":5,\"enforce_eager\":false}" +``` + +Attention-local FlashComm1/SP and native PD may remain enabled with both graphs. +For PD, keep the native full-draft P and KV-transfer configuration described +below. The W4A8 graph comparison covers A DP1/TP4, F EP4, and native P TP4 when +using `MooncakeHybridConnector`. Compilation mode 0 is the qualified example; +other compilation modes have not been validated here. + +Native capture buckets are normalized separately for target verification and +draft queries; choose buckets covering both query widths and `max-num-seqs`. +For A TP4, set `A_TP=4`, `A_DP=1`, recompute `NUM_A`, and use four distinct +A devices. Both roles use the same A TP width for CAM topology. F uses layered +eager or layered FULL, with one combined target/draft weight table. Layered eager can be slower for +small decode batches, so measure the chosen deployment. For F FULL, replace the F settings above before launching: @@ -78,9 +103,9 @@ FFN_EXEC_ARGS=(--compilation-config '{"cudagraph_mode":"FULL"}') F captures one DR → layered FFN → CS transaction and repeatedly replays it for both target and draft layers. The bounded replay group is independent of request -and target/draft phase boundaries. A and the native draft configuration retain -eager execution. FULL requires matching W4A8 or W8A8 layouts on all layers; unsupported -graph modes fail startup. DSpark does not support DBO or FFN two-stream. +and target/draft phase boundaries. Attention target/draft graph choices are +independent of F's execution mode. FULL requires matching W4A8 or W8A8 layouts on all layers; unsupported +graph modes fail startup. DSpark does not support DBO. Layered W4A8 forwards the model's `swiglu_limit` for both target and draft experts in eager and FULL. A positive limit clamps gate above and up on both @@ -106,17 +131,17 @@ vllm serve "$MODEL" "${COMMON_ARGS[@]}" "${FFN_EXEC_ARGS[@]}" \ \"host\":\"$LOCAL_IP\",\"port\":12499, \"num_attention_ranks\":$NUM_A,\"num_ffn_ranks\":$NUM_F, \"compute_gate_on_attention\":true, - \"connector_extra_config\":{\"dynamicQuant\":1,\"attn_ranks_per_dp\":1} + \"connector_extra_config\":{\"dynamicQuant\":1,\"attn_ranks_per_dp\":$A_TP} } }" >ffn.log 2>&1 & ``` ```bash ASCEND_RT_VISIBLE_DEVICES="$A_DEVICES" VLLM_PLUGINS=ascend,afd \ -vllm serve "$MODEL" "${COMMON_ARGS[@]}" --enforce-eager \ +vllm serve "$MODEL" "${COMMON_ARGS[@]}" "${A_EXEC_ARGS[@]}" \ --port 7100 --data-parallel-rpc-port 17100 \ - --tensor-parallel-size 1 --data-parallel-size "$NUM_A" \ - --all2all-backend allgather_reducescatter --speculative-config "$DSPARK_CONFIG" \ + --tensor-parallel-size "$A_TP" --data-parallel-size "$A_DP" \ + --all2all-backend allgather_reducescatter --speculative-config "$A_DSPARK_CONFIG" \ --additional-config "{ \"enable_cpu_binding\":false,\"enable_force_eplb\":false, \"enable_dsa_cp\":false,\"multistream_dsv4_dsa_overlap\":false, @@ -125,7 +150,7 @@ vllm serve "$MODEL" "${COMMON_ARGS[@]}" --enforce-eager \ \"host\":\"$LOCAL_IP\",\"port\":12499, \"num_attention_ranks\":$NUM_A,\"num_ffn_ranks\":$NUM_F, \"compute_gate_on_attention\":true, - \"connector_extra_config\":{\"dynamicQuant\":1,\"attn_ranks_per_dp\":1} + \"connector_extra_config\":{\"dynamicQuant\":1,\"attn_ranks_per_dp\":$A_TP} } }" >attention.log 2>&1 & ``` diff --git a/tests/conftest.py b/tests/conftest.py index 1f7a46cf4..a805bfe58 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -4,6 +4,7 @@ from __future__ import annotations +import json import os import signal import subprocess @@ -95,6 +96,18 @@ def forward_cancellation(signum, _frame) -> None: signal.signal(signum, previous_handler) +def run_case_entry( + command: list[str], output_path: Path, env: dict[str, str] | None = None +) -> None: + """Require both successful execution and canonical acceptance at pytest gates.""" + run_runner(command, env=env) + record = json.loads((output_path / "run.json").read_text()) + assert record["acceptance_passed"], ( + "task execution succeeded but the fixed scenario did not qualify: " + f"{record['configuration']}" + ) + + def download_dataset(dataset_id: str, dataset_config: str | None = None) -> None: """Download/cache a Hugging Face dataset. diff --git a/tests/e2e/README.md b/tests/e2e/README.md index de0892bcc..d651f9432 100644 --- a/tests/e2e/README.md +++ b/tests/e2e/README.md @@ -12,15 +12,18 @@ The V2 DP2 cases replace the duplicate GPU `afd-eager-2a2f`, baseline and AFD eager, graph, and graph+DBO on 2A1F. Qwen3.6 MoE retains three: baseline and AFD eager and graph on 2A1F. GDN DBO is unsupported. -The DeepSeek-V2-Lite, Qwen3 MoE, and Qwen3.6 MoE GPU model suites collect -20 cases. The seven-case ready/merge CI gate excludes local-only cases such -as `afd-v2-graph-dbo-tp2`. The separate DSV4 Flash CUDA suite is described -below. - -The current suite has 13 DeepSeek cases, four Qwen3 MoE cases, and three -Qwen3.6 MoE cases. -Qwen3.6 GDN DBO is unsupported. Workload and threshold settings are -scenario-specific; see the scenario sections below. +The single-deployment inventory contains 52 definitions: 42 CUDA cases +(18 DeepSeek-V2-Lite, 13 Qwen3 MoE, three Qwen3.6 MoE and eight DSV4 +Flash) and ten NPU cases (seven DeepSeek-V2-Lite and three DSV4 Flash). +The CUDA counts include the local DeepEP M2N and standalone FFN service +entries. CI still selects the +same 15 unique node IDs across ready/merge and weekly gates. Definitions, +collected nodes and device acceptance results are different counts. + +The ten existing GPU multi-pod entries also use model-local Cases and the +shared local deployment executor. Their existing workload/liveness/cleanup +acceptance is retained; cross-pod execution evidence is not added. The two NPU +multi-pod entries and dedicated FFN/DSpark checks retain their existing path. Ordinary cases evaluate the first 7 GSM8K samples; DBO cases have a 24-sample minimum and 12 concurrent requests. The GPU DP2 eager/eager-DBO/ @@ -59,6 +62,118 @@ few-shot examples and a 4096-token maximum model length. See the [E2E testing design](../../docs/design/module/e2e_testing.md) before adding a model or case. +## Case migration and maintenance + +All 48 single-deployment entries now use model-local Python Cases through +their existing pytest node IDs. This changes the harness, not their accepted +workloads or coverage. It does not by itself establish device qualification: +report actual CUDA/NPU results, host integration, skips and unexecuted cases +separately. No existing case is removed in this migration. + +The extension rules below apply to the migrated Case paths: all single-deployment +entries, including NPU, and the ten GPU multi-pod entries. NPU multi-pod and +dedicated FFN/DSpark executors still use their retained legacy paths; they do not +yet have the same extension contract. + +- Define deployment, workload and checks in the model's `cases.py`. + Collection reads stable IDs; checkpoint and resource binding happens in the + child process. Shared code receives the model's builder explicitly. +- Reuse `presets.py` for common synchronous eager/decode-graph deployment + construction. V2 on GPU/NPU, Qwen3 MoE and Qwen3.6 share these presets. + Models supply their DP/TP layouts, backend, connector, capture size and DBO + settings; V2 keeps its comparison-specific sequence budget explicitly. + Presets contain no model IDs, scenario tables, workloads or acceptance rules. + DSV4 keeps its specialized deployment configuration in its model directory. +- Reuse existing fields, tasks, checks and connector protocols for a new + combination. Its shared-module diff must be empty. A missing field or + capability needs an explicit capability extension and affected verification. +- Keep sample counts, thresholds, required concurrency, evidence windows, + participating ranks and cleanup requirements with the Case. Do not infer + these from its ID in shared code. +- Preserve existing acceptance before deleting an old branch or helper. + Migrate active pytest, CLI and dedicated-script callers with their behavior. + +For an existing-capability combination on a migrated path, edit the model's +`cases.py` and its stable ID list; add CI/docs selection only where needed. +The common runner, +Case schema, presets, deployment, task/check implementations and probes must have zero +diff for that addition. Implement a new task as a frozen dataclass following +`Task`: `resolve()` applies documented overrides and reports ineligibility; +`execute()` runs the workload and validates its result. Its fields are recorded +as canonical/effective configuration. Model-specific response semantics stay +in the model directory. DSV4's existing two concurrent rounds plus GSM8K are +three declared phases, not runner branches. + +The runner resets/opens/closes each observation window once. Graph, DBO and +Hash checks consume the same phase snapshot; Hash freshness references the +preceding declared window. Observation-only phases still require every rank. +Pure eager/baseline cases without required checks do not load a probe; DSV4 +AFD eager intentionally retains its required Hash probe. + +Deployment owns connector order (sync FFN then Attention; async CAM Attention +then FFN), readiness, pre/post-workload liveness and cleanup. Its explicit +`cleanup_profile` retains the scoped NPU async CAM exception; ordinary cases +use strict cleanup. `Deployment.ffn_service` runs the FFN role as the +standalone FFN service (`torchrun -m afd_plugin.ffn_faas serve`) on one host: +the runner waits for every rank's `/ready`, ends the case with `/stop` on +every rank followed by Attention shutdown, requires the launcher to exit 0 and +every rank's status file to report released resources, and sweeps the +torchrun workers by environment on cleanup. Per-role options, connector configuration, pinned Attention +address and exact compilation settings preserve the existing NPU profiles. +`Deployment.dbo` is the default for its roles. An explicit `Role.dbo` replaces +that default; `Role.dbo=None` inherits it. `DboConfig()` enables DBO with native +runtime thresholds, while an effective value of `None` leaves it disabled. +Commands, probe validation and role logging use `Deployment.dbo_for_role()`; +checks require DBO on each role they name. Async CAM can therefore enable DBO +on Attention alone without requiring it on FFN. + +GPU multi-pod entries use the model module's `multi-pod` subcommand and +`build_multi_pod_case()`. Placement and peer coordination stay in `multi_pod`; +the shared deployment executor builds and starts each local role. See the +[multi-pod execution guide](multi_pod/README.md) for the existing workload +settings, result paths and qualification limits. A pod layout does not silently +disable checks from a single-deployment Case. + +The legacy runner CLI and launch helpers remain for NPU and dedicated-script +callers. GPU-specific scenario tables, evaluation branches and Hash wrappers +have been retired from that executor. New GPU combinations extend model-local +declarations. Retire remaining legacy helpers only with their active callers. +The Case path uses the shared deployment cleanup; a legacy invocation uses its +retained cleanup. A single run does not execute both paths. + +The migrated module CLI also supports inspection and local debug execution: + +```bash +python -m tests.e2e.models.deepseek_v2_lite.cases \ + --scenario afd-v2-graph-dp2 --model /path/to/DeepSeek-V2-Lite \ + --device-backend gpu --devices 0,1,2,3 \ + --output-path /path/to/run --describe +``` + +Remove `--describe` to start the real services. Optional `--vllm-bin`, +`--api-port-base`, `--afd-port` and `--startup-timeout` bind local resources. +`AFD_E2E_API_PORT_BASE` and `AFD_E2E_AFD_PORT` provide the corresponding port +defaults. Device visibility still comes from the allocated environment. + +`run.json` records canonical/effective configuration, overrides, resource +bindings, commands, phase results and cleanup. `configuration.eligible` only +describes configuration; `execution.status` also reports actual execution. +Only `acceptance_passed=true` qualifies the fixed scenario. Lowering the +sample count or threshold with `AFD_GSM8K_LIMIT` / `AFD_GSM8K_THRESHOLD` marks +the run as debug, even if its workload succeeds. Increasing samples preserves +the existing larger-workload threshold. The module CLI allows debug execution; +the shared `run_case_entry()` pytest helper rejects an unqualified result. + +Required MRV2 checks prepare the probe before launch and store independent +per-phase evidence in `/mrv2.json`; role logs stay in the run directory. +Linux host harness checks use real child processes, sockets and signals with +stub service/evaluator CLIs. Multi-pod checks also require PyTorch's TCPStore. +They verify orchestration, not device execution: + +```bash +python -m pytest tests/e2e/test_runner.py +``` + ## Run Run from the repository root. The environment needs `vllm`, `pytest`, @@ -70,6 +185,13 @@ The selected test downloads/caches `openai/gsm8k` and its Hugging Face model when the backend model env var is unset. Point `HF_HOME` at a persistent cache if you want to reuse downloads across runs. +For an existing qualification, first locate the checkpoint from its recorded +run metadata and set the backend model path explicitly. A missing entry in one +cache directory does not establish that the model is absent or expand the +validation scope. When downloads are disallowed, use the existing model/dataset +cache with `HF_HUB_OFFLINE=1` and `HF_DATASETS_OFFLINE=1`; report missing assets +instead of preparing new downloads. + GPU: ```bash @@ -110,7 +232,7 @@ python -m pytest -q -s \ "tests/e2e/models/deepseek_v2_lite/test_deepseek_v2_lite.py::test_deepseek_v2_lite[afd-graph-2a2f]" \ "tests/e2e/models/deepseek_v2_lite/test_deepseek_v2_lite.py::test_deepseek_v2_lite[afd-graph-dbo-2a2f]" -# Qwen3 MoE (P2P cases; add the DeepEP M2N cases with -k deepep_m2n) +# Qwen3 MoE (P2P cases; run the DeepEP M2N entry separately) python -m pytest -q -s \ tests/e2e/models/qwen3_moe/test_qwen3_moe.py::test_qwen3_moe @@ -125,8 +247,10 @@ python -m pytest -q -s \ Success means seven GPU DeepSeek gate cases, four NPU DeepSeek gate cases, four GPU Qwen3 MoE P2P cases, five GPU Qwen3 MoE DeepEP M2N cases, or three GPU Qwen3.6 MoE cases, passed with zero skips. -Running the entire DeepSeek module collects 13 cases on GPU and seven on NPU, -including local-only cases. +The DeepSeek `test_deepseek_v2_lite` entry collects 13 cases on GPU and seven +on NPU, including local-only cases. The separate DeepEP M2N entry adds five +GPU-only cases. An unfiltered module collection includes those five IDs even +on NPU; select the P2P entry or the required node IDs for an NPU run. The Qwen3.5/3.6 adapter-family suite uses `Qwen/Qwen3.6-35B-A3B` and `P2pNcclAFDConnector` on CUDA. The runner now selects native ModelRunnerV2 @@ -176,9 +300,24 @@ python -m pytest -q -s \ tests/e2e/models/qwen3_moe/test_qwen3_moe.py::test_qwen3_moe_deepep_m2n ``` +`test_qwen3_moe_deepep_m2n_ffn_service` runs the eager, graph, graph + DBO +and Attention TP2 graph deployments (`afd-deepep-m2n-faas-*-2a2f`) with the +FFN side served by the +[standalone GPU FFN service](../../docs/gpu/STANDALONE_FFN_SERVICE.md) +instead of `vllm serve`. Attention launches unchanged; execution evidence is +required from Attention only, because the service has no vLLM Worker. It has +the same requirements and is not run by CI. + +```bash +python -m pytest -q -s \ + tests/e2e/models/qwen3_moe/test_qwen3_moe.py::test_qwen3_moe_deepep_m2n_ffn_service +``` + `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). @@ -237,10 +376,10 @@ and graph mode change; the graph row uses `FULL_DECODE_ONLY` with capture size 8 The GPU validator accepts exactly two microbatches at Attention DP1 or DP > 1; FFN remains connector-driven. See the -[runtime matrix](../../docs/design/module/execution_platforms.md#tested-runtime-matrix) +[runtime matrix](../../docs/support.md#tested-runtime-matrix) for tested configurations and known DeepSeek-V2-Lite shutdown failures. NPU retains the separate eager DBO and graph -constraints documented in the [execution-platform matrix](../../docs/design/module/execution_platforms.md#tested-runtime-matrix). +constraints documented in the [support matrix](../../docs/support.md#tested-runtime-matrix). The GPU harness collects bounded per-rank summaries of actual two-stage and FULL-replay execution during evaluation. Historical hardware records retain the execution assertions of the test revision that produced them. On vLLM @@ -282,7 +421,7 @@ samples, with the existing 24-sample floor for DBO. `models/deepseek_v4_flash/test_gpu.py` provides separate CUDA qualification profiles for the synchronous `P2pNcclAFDConnector`. They use CUDA launch shapes independently of the Ascend W4A8 and HCCL profiles. See the -[runtime matrix](../../docs/design/module/execution_platforms.md#tested-runtime-matrix) +[runtime matrix](../../docs/support.md#tested-runtime-matrix) for source-specific hardware validation. The basic test has four cases: @@ -400,10 +539,23 @@ SIGTERM/SIGINT cancels pending HTTP operations before service cleanup, without waiting for the request timeout. It does **not** run GSM8K or establish general model accuracy. Service liveness and owned-process cleanup are also required to pass. -This 16-NPU case allows 60 seconds for service shutdown before escalation; +NPU async CAM cases allow 60 seconds for each vLLM worker, 90 seconds for the +service, and 120 seconds for the test runner to shut down. Equivalent external +launchers should set `VLLM_WORKER_SHUTDOWN_TIMEOUT_SECONDS=60` and pass +`--shutdown-timeout 90` to both services; a service timeout alone does not +change the worker deadline. Explicit worker environment overrides are retained. +The Case declares these three budgets through `worker_shutdown_timeout_s`, +`shutdown_timeout_s` and `termination_timeout_s`. The worker environment override +is recorded in the effective configuration and overrides alongside the canonical +60-second default. The retained legacy path uses the same default constants. Attention SIGKILL escalation remains a failure. FFN uses the existing scoped async CAM cleanup exception. +vLLM 0.32 can log `AsyncLLM`'s `EngineDeadError` after an API client starts +its own shutdown. Record this as an upstream advisory only after inference +succeeds, every worker group exits gracefully, and process exit codes and +cleanup checks pass. An inference-time error or forced worker exit still fails. + ```bash export AFD_E2E_BACKEND=npu export AFD_E2E_DEVICES=0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15 diff --git a/tests/e2e/accuracy/gsm8k.py b/tests/e2e/accuracy/gsm8k.py index 39b60d4ab..99d511fd0 100644 --- a/tests/e2e/accuracy/gsm8k.py +++ b/tests/e2e/accuracy/gsm8k.py @@ -13,9 +13,12 @@ import sys import threading import time +from collections.abc import Mapping +from dataclasses import replace from pathlib import Path from types import FrameType +from tests.e2e.case import Gsm8kSettings, JsonValue from tests.e2e.process_utils import terminate_process_groups LM_EVAL_TIMEOUT_S = 7200 @@ -23,6 +26,88 @@ LM_EVAL_PROCESS_POLL_INTERVAL_S = 0.2 LM_EVAL_REAP_TIMEOUT_S = 5 LM_EVAL_READER_JOIN_TIMEOUT_S = 5 +GSM8K_FULL_SAMPLE_COUNT = 1319 +GSM8K_FULL_TIMEOUT_S = 8 * 60 * 60 +GSM8K_LIMIT_ENV = "AFD_GSM8K_LIMIT" +GSM8K_THRESHOLD_ENV = "AFD_GSM8K_THRESHOLD" + + +def resolve_settings( + canonical: Gsm8kSettings, environment: Mapping[str, str] +) -> tuple[Gsm8kSettings, tuple[str, ...]]: + """Apply documented local overrides and determine configuration eligibility.""" + limit = canonical.sample_limit + if GSM8K_LIMIT_ENV in environment: + raw = environment[GSM8K_LIMIT_ENV] + limit = None if raw == "all" else int(raw) + if limit is not None: + if limit < 1: + raise ValueError("GSM8K sample limit must be positive or 'all'") + limit = max(limit, canonical.minimum_samples) + threshold = canonical.threshold + if ( + canonical.smoke_max_samples is not None + and limit is not None + and limit <= canonical.smoke_max_samples + and canonical.smoke_threshold is not None + ): + threshold = canonical.smoke_threshold + if ( + canonical.smoke_max_samples is not None + and (limit is None or limit > canonical.smoke_max_samples) + and canonical.larger_sample_threshold is not None + ): + threshold = canonical.larger_sample_threshold + required_threshold = threshold + if GSM8K_THRESHOLD_ENV in environment: + threshold = float(environment[GSM8K_THRESHOLD_ENV]) + if not 0 <= threshold <= 1: + raise ValueError("GSM8K threshold must be a finite value between 0 and 1") + reasons = [] + canonical_count = canonical.sample_limit or GSM8K_FULL_SAMPLE_COUNT + actual_count = limit or GSM8K_FULL_SAMPLE_COUNT + if actual_count < canonical_count: + reasons.append(f"sample count reduced from {canonical_count} to {actual_count}") + if threshold < required_threshold: + reasons.append( + f"accuracy threshold reduced from {required_threshold} to {threshold}" + ) + return replace(canonical, sample_limit=limit, threshold=threshold), tuple(reasons) + + +def evaluate_settings( + settings: Gsm8kSettings, + *, + base_url: str, + model_name: str, + checkpoint: str, + output_path: Path, +) -> dict[str, JsonValue]: + """Execute the resolved Case workload and validate complete evaluator output.""" + results = _run_lm_eval( + base_url, + model_name, + output_path=str(output_path), + num_fewshot=settings.num_fewshot, + num_concurrent=settings.num_concurrent, + max_gen_toks=settings.max_gen_toks, + tokenizer=checkpoint, + limit=settings.sample_limit, + timeout_s=( + GSM8K_FULL_TIMEOUT_S if settings.sample_limit is None else LM_EVAL_TIMEOUT_S + ), + ) + count = _extract_gsm8k_sample_count(results) + expected = settings.sample_limit or GSM8K_FULL_SAMPLE_COUNT + if count != expected: + raise RuntimeError(f"GSM8K evaluated {count} samples; expected {expected}") + accuracy = _extract_gsm8k_accuracy(results) + if not accuracy >= settings.threshold: + raise RuntimeError( + f"GSM8K accuracy {accuracy:.4f} is below the required " + f"threshold {settings.threshold:.4f}" + ) + return {"sample_count": count, "accuracy": accuracy} def _finish_lm_eval_cleanup( diff --git a/tests/e2e/case.py b/tests/e2e/case.py new file mode 100644 index 000000000..b6ab5f2c1 --- /dev/null +++ b/tests/e2e/case.py @@ -0,0 +1,235 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Immutable declarations exchanged by model entries and the E2E executor. + +Importing declarations does not resolve checkpoints or import a device runtime. +New fields represent shared capabilities and require their own review. +""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass +from pathlib import Path +from typing import Protocol, TypeAlias + +JsonValue: TypeAlias = ( + str | int | float | bool | None | list["JsonValue"] | dict[str, "JsonValue"] +) + +EXECUTION_EAGER = "eager" +EXECUTION_FULL_DECODE_GRAPH = "full_decode_graph" +CLEANUP_PROFILE_STRICT = "strict" +CLEANUP_PROFILE_NPU_ASYNC_CAM_LEGACY = "npu_async_cam_legacy" +BASELINE_ROLE = "baseline" +ATTENTION_ROLE = "attention" +FFN_ROLE = "ffn" +DEFAULT_GENERATED_TOKENS = 512 +DEFAULT_FEWSHOT = 8 +DEFAULT_API_PORT = 8000 +DEFAULT_CONNECTOR_PORT = 1239 +DEFAULT_STARTUP_TIMEOUT_S = 900 +DEFAULT_SHUTDOWN_TIMEOUT_S = 10 +DEFAULT_TERMINATION_TIMEOUT_S = 20 +WORKER_SHUTDOWN_TIMEOUT_ENV = "VLLM_WORKER_SHUTDOWN_TIMEOUT_SECONDS" +NPU_ASYNC_VLLM_SHUTDOWN_TIMEOUT_S = 90 +NPU_ASYNC_WORKER_SHUTDOWN_TIMEOUT_S = 60 +NPU_ASYNC_TERMINATION_TIMEOUT_S = 120 +# Independently started services must pair within the AFD rendezvous timeout on +# both roles; the standalone FFN service loads its experts while Attention loads. +FFN_SERVICE_AFD_TIMEOUT_S = 600 +FFN_SERVICE_RECEIVER_SHUTDOWN_TIMEOUT_S = 10 +FFN_SERVICE_STOP_TIMEOUT_S = 120 + + +@dataclass(frozen=True) +class FfnService: + """The FFN role runs the standalone FFN service instead of ``vllm serve``. + + ``torchrun -m afd_plugin.ffn_faas serve --backend `` starts one + process per FFN rank on the local host. Attention launches unchanged. The + service reports per-rank readiness on its management endpoint and must + stop cleanly: ``/stop`` on every rank while idle, then Attention. + """ + + backend: str + afd_timeout_s: int = FFN_SERVICE_AFD_TIMEOUT_S + receiver_shutdown_timeout_s: int = FFN_SERVICE_RECEIVER_SHUTDOWN_TIMEOUT_S + stop_timeout_s: float = FFN_SERVICE_STOP_TIMEOUT_S + + +@dataclass(frozen=True) +class Role: + dp: int + tp: int = 1 + execution: str = EXECUTION_EAGER + dbo: DboConfig | None = None + vllm_args: tuple[str, ...] = () + additional_config: tuple[tuple[str, bool], ...] = () + + def __post_init__(self) -> None: + if self.dp < 1 or self.tp < 1: + raise ValueError("DP and TP must be positive") + if self.execution not in (EXECUTION_EAGER, EXECUTION_FULL_DECODE_GRAPH): + raise ValueError(f"unknown execution mode {self.execution!r}") + + @property + def rank_count(self) -> int: + return self.dp * self.tp + + @property + def uses_graph(self) -> bool: + return self.execution == EXECUTION_FULL_DECODE_GRAPH + + +@dataclass(frozen=True) +class DboConfig: + # None keeps the runtime default, as the existing async Attention path does. + decode_token_threshold: int | None = None + prefill_token_threshold: int | None = None + + +@dataclass(frozen=True) +class Deployment: + backend: str + attention: Role + ffn: Role | None = None + connector: str | None = None + cleanup_profile: str = CLEANUP_PROFILE_STRICT + dbo: DboConfig | None = None + cudagraph_capture_size: int | None = None + max_num_seqs: int | None = None + max_num_batched_tokens: int | None = None + compute_gate_on_attention: bool = False + connector_extra_config: tuple[tuple[str, int], ...] = () + pin_attention_address: bool = False + compilation_config: tuple[tuple[str, str | tuple[int, ...]], ...] = () + shutdown_timeout_s: int = DEFAULT_SHUTDOWN_TIMEOUT_S + # Existing worker environment overrides take precedence over this default. + worker_shutdown_timeout_s: int | None = None + termination_timeout_s: float = DEFAULT_TERMINATION_TIMEOUT_S + ffn_service: FfnService | None = None + + @property + def is_baseline(self) -> bool: + return self.ffn is None + + @property + def roles(self) -> tuple[tuple[str, Role], ...]: + if self.ffn is None: + return ((BASELINE_ROLE, self.attention),) + return ((ATTENTION_ROLE, self.attention), (FFN_ROLE, self.ffn)) + + def dbo_for_role(self, role: str) -> DboConfig | None: + """A role's DBO overrides the deployment default; None inherits it.""" + override = dict(self.roles)[role].dbo + return self.dbo if override is None else override + + +@dataclass(frozen=True) +class Gsm8kSettings: + sample_limit: int | None + threshold: float + num_fewshot: int = DEFAULT_FEWSHOT + num_concurrent: int | None = None + max_gen_toks: int = DEFAULT_GENERATED_TOKENS + minimum_samples: int = 1 + # Existing smoke profiles raise their threshold for larger workloads. + smoke_max_samples: int | None = None + smoke_threshold: float | None = None + larger_sample_threshold: float | None = None + + +@dataclass(frozen=True) +class Gsm8kTask: + settings: Gsm8kSettings + + def resolve( + self, environment: Mapping[str, str] + ) -> tuple[Gsm8kTask, tuple[str, ...]]: + from tests.e2e.accuracy.gsm8k import resolve_settings + + settings, reasons = resolve_settings(self.settings, environment) + return Gsm8kTask(settings), reasons + + def execute( + self, *, base_url: str, model_name: str, checkpoint: str, output_path: Path + ) -> dict[str, JsonValue]: + from tests.e2e.accuracy.gsm8k import evaluate_settings + + return evaluate_settings( + self.settings, + base_url=base_url, + model_name=model_name, + checkpoint=checkpoint, + output_path=output_path, + ) + + +class Task(Protocol): + """A serializable task declaration, resolved and executed in the child.""" + + def resolve( + self, environment: Mapping[str, str] + ) -> tuple[Task, tuple[str, ...]]: ... + + def execute( + self, *, base_url: str, model_name: str, checkpoint: str, output_path: Path + ) -> dict[str, JsonValue]: ... + + +@dataclass(frozen=True) +class Mrv2Check: + roles: tuple[str, ...] + graph: bool = False + dbo: bool = False + + +@dataclass(frozen=True) +class HashInputCheck: + roles: tuple[str, ...] = (ATTENTION_ROLE, FFN_ROLE) + graph: bool = False + dbo: bool = False + compare_with: str | None = None + + +@dataclass(frozen=True) +class Phase: + name: str + task: Task + checks: tuple[Mrv2Check | HashInputCheck, ...] = () + + +@dataclass(frozen=True) +class Case: + id: str + model: str + deployment: Deployment + phases: tuple[Phase, ...] + served_model_name_prefix: str + common_vllm_args: tuple[str, ...] = () + + def __post_init__(self) -> None: + if not self.id or not self.phases: + raise ValueError("Case requires an ID and at least one phase") + names = [phase.name for phase in self.phases] + if len(names) != len(set(names)): + raise ValueError("phase names must be unique within a Case") + + +@dataclass(frozen=True) +class CaseEnvironment: + model: str + vllm_bin: str + device_backend: str + devices: tuple[str, ...] + output_path: Path + api_host: str = "127.0.0.1" + api_port_base: int = DEFAULT_API_PORT + afd_host: str = "127.0.0.1" + afd_port: int = DEFAULT_CONNECTOR_PORT + startup_timeout: float = DEFAULT_STARTUP_TIMEOUT_S + + def __post_init__(self) -> None: + if len(self.devices) != len(set(self.devices)): + raise ValueError("devices must be unique") diff --git a/tests/e2e/checks/__init__.py b/tests/e2e/checks/__init__.py new file mode 100644 index 000000000..37277d264 --- /dev/null +++ b/tests/e2e/checks/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project diff --git a/tests/e2e/checks/mrv2.py b/tests/e2e/checks/mrv2.py new file mode 100644 index 000000000..7dd8b7b62 --- /dev/null +++ b/tests/e2e/checks/mrv2.py @@ -0,0 +1,372 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Host-side MRV2 probe preparation, per-phase windows and evidence checks.""" + +from __future__ import annotations + +import json +import re +import tempfile +import time +import urllib.request +from pathlib import Path + +from tests.e2e.case import Case, CaseEnvironment, HashInputCheck, Mrv2Check, Phase, Role +from tests.e2e.deployment import ServiceDeployment + +MRV2_EVIDENCE_RPC_TIMEOUT_S = 30 +MRV2_EVIDENCE_REPLAY_MIN_COUNT = 2 +PROCESS_POLL_INTERVAL_S = 0.2 +DBO_NUM_STAGES = 2 + + +class Mrv2Session: + """One shared probe per deployment, with independent immutable snapshots.""" + + def __init__(self, case: Case, environment: CaseEnvironment, run_id: str) -> None: + self.case = case + self.environment = environment + self.run_id = run_id + self.directory: tempfile.TemporaryDirectory[str] | None = None + self.start_marker: Path | None = None + self.hash_samples: dict[ + str, dict[tuple[int, int, int, int], set[tuple[int, ...]]] + ] = {} + layouts = dict(case.deployment.roles) + for phase in case.phases: + for check in phase.checks: + if not check.roles: + raise ValueError("MRV2 check requires at least one role") + for role in check.roles: + if role not in layouts: + raise ValueError(f"MRV2 requests absent role {role}") + if check.graph and not layouts[role].uses_graph: + raise ValueError( + f"graph evidence requested for eager role {role}" + ) + if check.dbo and case.deployment.dbo_for_role(role) is None: + raise ValueError( + f"DBO evidence requested for non-DBO role {role}" + ) + + def prepare(self) -> None: + if any(phase.checks for phase in self.case.phases): + self.directory = tempfile.TemporaryDirectory(prefix="afd-e2e-mrv2-") + self.start_marker = Path(self.directory.name) / "evaluation-start" + + @property + def arguments(self) -> tuple[str, ...]: + if self.start_marker is None: + return () + return ("--worker-extension-cls", "tests.e2e.mrv2_evidence.Worker") + + @property + def worker_environment(self) -> dict[str, str]: + if self.start_marker is None: + return {} + env = { + "VLLM_SERVER_DEV_MODE": "1", + "AFD_E2E_MRV2_EVIDENCE_START": str(self.start_marker), + } + if self.case.deployment.is_baseline: + env["AFD_E2E_NATIVE_BASELINE"] = "1" + hash_checks = [ + check + for phase in self.case.phases + for check in phase.checks + if isinstance(check, HashInputCheck) + ] + if hash_checks: + env["AFD_E2E_HASH_INPUT_IDS"] = "1" + env["AFD_E2E_HASH_REQUIRE_GRAPH"] = ( + "1" if any(check.graph for check in hash_checks) else "0" + ) + return env + + def rpc(self, method: str, window: str | None = None) -> None: + env = self.environment + request = urllib.request.Request( + f"http://{env.api_host}:{env.api_port_base}/collective_rpc", + data=json.dumps( + { + "method": method, + "timeout": MRV2_EVIDENCE_RPC_TIMEOUT_S, + "kwargs": {} if window is None else {"evaluation_window": window}, + } + ).encode(), + headers={"Content-Type": "application/json"}, + ) + with urllib.request.urlopen( + request, timeout=MRV2_EVIDENCE_RPC_TIMEOUT_S + ) as reply: + reply.read() + + def start(self, phase: Phase) -> None: + if not phase.checks: + return + assert self.start_marker is not None + self.start_marker.with_name("evaluation-complete").unlink(missing_ok=True) + self.rpc("afd_reset_execution_evidence", phase.name) + pending = self.start_marker.with_suffix(".pending") + pending.write_text(phase.name) + pending.replace(self.start_marker) + + def finish( + self, phase: Phase, services: ServiceDeployment + ) -> tuple[tuple[str, str], ...]: + if not phase.checks: + return () + assert self.start_marker is not None + self.rpc("afd_log_execution_evidence") + self.start_marker.with_name("evaluation-complete").write_text(phase.name) + layouts = dict(self.case.deployment.roles) + expected = { + (role, dp, tp) + for check in phase.checks + for role in check.roles + for dp in range(layouts[role].dp) + for tp in range(layouts[role].tp) + } + deadline = time.monotonic() + MRV2_EVIDENCE_RPC_TIMEOUT_S + while True: + services.ensure_alive() + snapshot = tuple( + (role, line) + for role, line in services.execution_evidence + if f"window={phase.name} " in line + ) + observed = { + (role, int(match[1]), int(match[2])) + for role, line in snapshot + if (match := re.search(r"dp_rank=(\d+) tp_rank=(\d+)", line)) + is not None + } + if expected <= observed: + break + if time.monotonic() >= deadline: + raise RuntimeError( + f"missing MRV2 ranks in {phase.name}: {sorted(expected - observed)}" + ) + time.sleep(PROCESS_POLL_INTERVAL_S) + path = self.environment.output_path / phase.name + path.mkdir(parents=True, exist_ok=True) + (path / "mrv2.json").write_text( + json.dumps( + { + "run_id": self.run_id, + "case": self.case.id, + "phase": phase.name, + "evidence": snapshot, + }, + indent=2, + ) + + "\n" + ) + for check in phase.checks: + if isinstance(check, Mrv2Check): + assert_execution_evidence( + list(snapshot), + check, + role_layouts=layouts, + backend=self.case.deployment.backend, + asynchronous=self.case.deployment.connector + == "CAMAsyncAFDConnector", + ) + else: + samples = assert_hash_inputs( + list(snapshot), + role_layouts=layouts, + graph=check.graph, + dbo=check.dbo, + ) + if check.compare_with is not None: + assert_hash_freshness( + self.hash_samples[check.compare_with], + samples, + attention_ranks=layouts["attention"].rank_count, + dbo=check.dbo, + ) + self.hash_samples[phase.name] = samples + return snapshot + + def close(self) -> None: + if self.directory is not None: + self.directory.cleanup() + + +def assert_execution_evidence( + evidence: list[tuple[str, str]], + check: Mrv2Check, + *, + role_layouts: dict[str, Role], + backend: str, + asynchronous: bool = False, +) -> None: + if not check.graph and not check.dbo: + # Observation-only windows still require every rank in finish(). + return + metric = ( + "two_stage_graph_replays" + if check.dbo and check.graph + else "two_stage_steps" + if check.dbo + else "graph_replays" + ) + minimum = ( + MRV2_EVIDENCE_REPLAY_MIN_COUNT if metric == "two_stage_graph_replays" else 1 + ) + layout_kind = "consecutive_graph" if check.graph else "two_stage" + common_layouts: set[str] | None = None + for role in check.roles: + tp_size = role_layouts[role].tp + rank_count = role_layouts[role].rank_count + expected = { + (dp, tp) for dp in range(rank_count // tp_size) for tp in range(tp_size) + } + observed = set() + layouts: dict[str, set[tuple[int, int]]] = {} + for recorded_role, line in evidence: + if recorded_role != role: + continue + match = re.search( + rf"dp_rank=(\d+) tp_rank=(\d+).*\b{metric}=(\d+)(?: |$)", line + ) + if match is not None and int(match[3]) >= minimum: + rank = (int(match[1]), int(match[2])) + observed.add(rank) + if check.dbo and not asynchronous: + payload = json.loads(line.split(" layouts=", 1)[1]) + for layout in payload[layout_kind]: + key = json.dumps(layout, separators=(",", ":")) + layouts.setdefault(key, set()).add(rank) + if missing := expected - observed: + raise RuntimeError( + f"{backend.upper()} {role} lacks live " + f"{metric}>={minimum} evidence on ranks " + f"{sorted(missing)}" + ) + if check.dbo and not asynchronous: + complete_layouts = { + layout for layout, ranks in layouts.items() if expected <= ranks + } + common_layouts = ( + complete_layouts + if common_layouts is None + else common_layouts & complete_layouts + ) + if check.dbo and not asynchronous and not common_layouts: + raise RuntimeError( + f"{backend.upper()} DBO lacks matching two-stage layouts " + "across every participating rank" + + (" with consecutive graph replays" if check.graph else "") + ) + print( + f"[execution-coverage] live {metric}>={minimum} " + "confirmed on every participating rank" + ) + + +def assert_hash_inputs( + evidence: list[tuple[str, str]], + *, + role_layouts: dict[str, Role], + graph: bool, + dbo: bool, +) -> dict[tuple[int, int, int, int], set[tuple[int, ...]]]: + """Pair bounded successful live A/F samples in original control/peer order.""" + samples: dict[str, dict[tuple[int, int, int], dict]] = { + "attention": {}, + "ffn": {}, + } + observed_ranks: dict[str, set[tuple[int, int]]] = { + "attention": set(), + "ffn": set(), + } + for role, line in evidence: + match = re.search(r"dp_rank=(\d+) tp_rank=(\d+)", line) + assert match is not None + rank = (int(match[1]), int(match[2])) + payload = json.loads(line.split(" layouts=", 1)[1]) + for sample in payload["hash_ids"]: + if not sample["completed"] or sample["control"] <= 0: + raise RuntimeError( + "Hash ID sample did not complete a live control transaction" + ) + if graph and not sample["replayed"]: + raise RuntimeError("Hash ID sample lacks completed graph replay") + key = (sample["control"], sample["stage"], sample["peer"]) + if key in samples[role]: + raise RuntimeError(f"Duplicate {role} Hash ID sample: {key}") + if role == "attention" and sample["peer"] != rank[0]: + raise RuntimeError( + f"Attention Hash ID peer does not match DP rank: {key}" + ) + samples[role][key] = sample + observed_ranks[role].add(rank) + for role, ranks in observed_ranks.items(): + rank_count = role_layouts[role].rank_count + expected = {(dp, 0) for dp in range(rank_count)} + if ranks != expected: + raise RuntimeError( + f"{role} lacks bounded Hash ID samples on ranks {expected - ranks}" + ) + if samples["attention"].keys() != samples["ffn"].keys(): + raise RuntimeError("A/F Hash ID samples do not pair by control/stage/peer") + values: dict[tuple[int, int, int, int], set[tuple[int, ...]]] = {} + covered = set() + for key, attention in samples["attention"].items(): + ffn = samples["ffn"][key] + real = attention["real_tokens"] + bucket = attention["bucket"] + if ( + ffn["bucket"] != bucket + or len(ffn["input_ids"]) != bucket + or len(attention["input_ids"]) != real + or not 0 <= real <= bucket + or attention["input_ids"] != ffn["input_ids"][:real] + ): + raise RuntimeError( + f"Native A / actual Hash-layer F input IDs differ: {key}" + ) + if real: + _, stage, peer = key + covered.add((peer, stage)) + values.setdefault((peer, stage, bucket, real), set()).add( + tuple(attention["input_ids"]) + ) + stages = range(DBO_NUM_STAGES if dbo else 1) + required = { + (peer, stage) + for peer in range(role_layouts["attention"].rank_count) + for stage in stages + } + if missing := required - covered: + raise RuntimeError( + f"Hash ID samples lack real tokens for peer/stage {sorted(missing)}" + ) + print("[hash-inputs] native A and actual Hash-layer F IDs match on every rank") + return values + + +def assert_hash_freshness( + previous: dict[tuple[int, int, int, int], set[tuple[int, ...]]], + current: dict[tuple[int, int, int, int], set[tuple[int, ...]]], + *, + attention_ranks: int, + dbo: bool, +) -> None: + """Require new real IDs across rounds at the same padded and real size.""" + changed = { + (peer, stage) + for (peer, stage, bucket, real), values in current.items() + if (peer, stage, bucket, real) in previous + and values - previous[(peer, stage, bucket, real)] + } + stages = range(DBO_NUM_STAGES if dbo else 1) + required = {(peer, stage) for peer in range(attention_ranks) for stage in stages} + if missing := required - changed: + raise RuntimeError( + "Hash input IDs did not change across concurrent rounds with the same " + f"bucket and real token count: {sorted(missing)}" + ) + print("[hash-inputs] both concurrent windows show fresh IDs at the same bucket") diff --git a/tests/e2e/deployment.py b/tests/e2e/deployment.py new file mode 100644 index 000000000..f7328a703 --- /dev/null +++ b/tests/e2e/deployment.py @@ -0,0 +1,617 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Case deployment: protocol order, role commands, readiness and owned cleanup.""" + +from __future__ import annotations + +import json +import os +import shlex +import subprocess +import sys +import threading +import time +import urllib.error +import urllib.request +from collections.abc import Callable, Mapping +from dataclasses import dataclass +from pathlib import Path + +from tests.e2e.case import ( + ATTENTION_ROLE, + CLEANUP_PROFILE_NPU_ASYNC_CAM_LEGACY, + CLEANUP_PROFILE_STRICT, + FFN_ROLE, + WORKER_SHUTDOWN_TIMEOUT_ENV, + Case, + CaseEnvironment, + DboConfig, + Deployment, + JsonValue, +) +from tests.e2e.multi_pod.layout import RoleSlot +from tests.e2e.process_utils import ( + kill_processes_matching_environment, + terminate_process_groups, +) + +REPO_ROOT = Path(__file__).resolve().parents[2] +PROCESS_POLL_INTERVAL_S = 0.2 +PROCESS_REAP_TIMEOUT_S = 5 +LOG_THREAD_JOIN_TIMEOUT_S = 2 +READY_POLL_INTERVAL_S = 2 +MRV2_EVIDENCE_PREFIX = "AFD E2E execution: " +NPU_ASYNC_PROCESS_KILL_TIMEOUT_S = 120 +DECODE_GRAPH_COMPILATION_CONFIG = {"cudagraph_mode": "FULL_DECODE_ONLY"} +# The standalone FFN service's torchrun (F-only) world listens next to the AFD +# rendezvous port; its rank r management endpoint is the FFN API port + r. +FFN_SERVICE_MASTER_PORT_OFFSET = 1 +FFN_SERVICE_STATUS_FILE = "ffn-service.json" +FFN_SERVICE_HTTP_TIMEOUT_S = 5 +FFN_SERVICE_PROCESS_KILL_TIMEOUT_S = 120 +# Starting the first role is not waiting for its connector rendezvous or API. +CONNECTOR_PROTOCOLS = ( + ("gpu", "P2pNcclAFDConnector", ("ffn", "attention")), + ("gpu", "DeepEPM2NAFDConnector", ("ffn", "attention")), + ("npu", "CAMP2pAFDConnector", ("ffn", "attention")), + ("npu", "CAMAsyncAFDConnector", ("attention", "ffn")), +) +DEPLOYMENT_ARGUMENTS = frozenset( + ( + "--config", + "--additional-config", + "--served-model-name", + "--data-parallel-size", + "--tensor-parallel-size", + "--data-parallel-size-local", + "--data-parallel-start-rank", + "--data-parallel-address", + "--data-parallel-rpc-port", + "--headless", + "--worker-extension-cls", + "--compilation-config", + "--enforce-eager", + "--max-num-seqs", + "--max-num-batched-tokens", + "--max-cudagraph-capture-size", + "--cudagraph-capture-sizes", + "--enable-dbo", + "--dbo-decode-token-threshold", + "--dbo-prefill-token-threshold", + "--shutdown-timeout", + "--host", + "--port", + ) +) + + +@dataclass(frozen=True) +class ResolvedDeployment: + case: Case + environment: CaseEnvironment + role_devices: tuple[tuple[str, tuple[str, ...]], ...] + launch_order: tuple[str, ...] + slots: tuple[RoleSlot, ...] = () + + def command(self, role: str, probe_arguments: tuple[str, ...] = ()) -> list[str]: + profile = self.case.deployment + if role == FFN_ROLE and profile.ffn_service is not None: + # The service has no vLLM Worker to host the evidence extension. + return self.ffn_service_command() + layout = dict(profile.roles)[role] + dbo = profile.dbo_for_role(role) + env = self.environment + slot = next((slot for slot in self.slots if slot.role == role), None) + cmd = [ + env.vllm_bin, + "serve", + env.model, + "--shutdown-timeout", + str(profile.shutdown_timeout_s), + "--served-model-name", + f"{self.case.served_model_name_prefix}-{role}", + "--data-parallel-size", + str(layout.dp), + "--tensor-parallel-size", + str(layout.tp), + "--enable-expert-parallel", + ] + if profile.ffn is not None: + cmd.extend( + [ + "--enable-ep-weight-filter", + "--additional-config", + json.dumps( + { + "afd": self.afd_config(role, slot), + **dict(layout.additional_config), + }, + separators=(",", ":"), + ), + ] + ) + if slot is not None and slot.spans_pods: + cmd.extend( + [ + "--data-parallel-size-local", + str(slot.dp_size_local), + "--data-parallel-address", + slot.dp_address, + "--data-parallel-rpc-port", + str(slot.dp_rpc_port), + ] + ) + # A leader must omit start-rank: vLLM otherwise selects external LB + # and rejects its headless peers. + if slot.headless: + cmd.extend( + [ + "--data-parallel-start-rank", + str(slot.dp_start_rank), + "--headless", + ] + ) + elif role == "attention" and profile.pin_attention_address: + cmd.extend( + [ + "--data-parallel-address", + env.afd_host if slot is None else slot.dp_address, + ] + ) + cmd.extend(probe_arguments) + # Native baseline emits DBO before graph flags; preserve existing CLI. + if profile.is_baseline: + cmd.extend(dbo_arguments(dbo)) + if profile.max_num_seqs is not None: + cmd.extend(["--max-num-seqs", str(profile.max_num_seqs)]) + if profile.max_num_batched_tokens is not None: + cmd.extend( + ["--max-num-batched-tokens", str(profile.max_num_batched_tokens)] + ) + if profile.compilation_config: + cmd.extend( + [ + "--compilation-config", + json.dumps(dict(profile.compilation_config), separators=(",", ":")), + ] + ) + elif layout.uses_graph: + capture_size = str(profile.cudagraph_capture_size) + cmd.extend( + [ + "--max-cudagraph-capture-size", + capture_size, + "--cudagraph-capture-sizes", + capture_size, + "--compilation-config", + json.dumps(DECODE_GRAPH_COMPILATION_CONFIG, separators=(",", ":")), + ] + ) + else: + cmd.append("--enforce-eager") + if not profile.is_baseline: + cmd.extend(dbo_arguments(dbo)) + port = env.api_port_base + (role == "ffn") + if slot is None or not slot.headless: + cmd.extend(["--host", env.api_host, "--port", str(port)]) + cmd.extend(layout.vllm_args) + cmd.extend(self.case.common_vllm_args) + return cmd + + def afd_config(self, role: str, slot: RoleSlot | None) -> dict[str, JsonValue]: + """The ``additional_config["afd"]`` one role's service starts with.""" + profile = self.case.deployment + assert profile.ffn is not None + env = self.environment + afd: dict[str, JsonValue] = { + "role": role, + "connector": profile.connector, + "host": env.afd_host if slot is None else slot.afd_host, + "port": env.afd_port, + "num_attention_ranks": profile.attention.rank_count, + "num_ffn_ranks": profile.ffn.rank_count, + } + if profile.connector == "CAMAsyncAFDConnector": + afd["async"] = True + if profile.compute_gate_on_attention: + afd["compute_gate_on_attention"] = True + if profile.connector_extra_config: + afd["connector_extra_config"] = dict(profile.connector_extra_config) + if profile.ffn_service is not None: + afd["afd_process_group_timeout_s"] = profile.ffn_service.afd_timeout_s + return afd + + def ffn_service_command(self) -> list[str]: + """``torchrun`` command of the standalone FFN service on this host. + + One process per FFN rank. The AFD configuration, execution mode and + DBO match the FFN half of the ``vllm serve`` deployment, so the + unchanged Attention command pairs with it. Request-shaping arguments + (``common_vllm_args``, batch and capture budgets) belong to Attention, + whose control payloads drive every FFN step and graph capture. + """ + profile = self.case.deployment + service, ffn = profile.ffn_service, profile.ffn + assert service is not None and ffn is not None + env = self.environment + cmd = [ + sys.executable, + "-m", + "torch.distributed.run", + f"--nproc-per-node={ffn.rank_count}", + "--nnodes=1", + f"--master-addr={env.afd_host}", + f"--master-port={env.afd_port + FFN_SERVICE_MASTER_PORT_OFFSET}", + "-m", + "afd_plugin.ffn_faas", + "serve", + env.model, + "--backend", + service.backend, + "--data-parallel-size", + str(ffn.dp), + "--tensor-parallel-size", + str(ffn.tp), + "--enable-expert-parallel", + "--additional-config", + json.dumps( + { + "afd": self.afd_config(FFN_ROLE, None), + **dict(ffn.additional_config), + }, + separators=(",", ":"), + ), + "--receiver-shutdown-timeout", + str(service.receiver_shutdown_timeout_s), + "--status-host", + env.api_host, + "--status-port", + str(self.ffn_service_status_port(0)), + "--status-file", + str(self.ffn_service_status_file), + ] + if ffn.uses_graph: + cmd.extend( + [ + "--compilation-config", + json.dumps(DECODE_GRAPH_COMPILATION_CONFIG, separators=(",", ":")), + ] + ) + else: + cmd.append("--enforce-eager") + cmd.extend(dbo_arguments(profile.dbo_for_role(FFN_ROLE))) + cmd.extend(ffn.vllm_args) + return cmd + + def ffn_service_status_port(self, rank: int) -> int: + """Management port of one FFN service rank (the FFN API port + rank).""" + return self.environment.api_port_base + 1 + rank + + @property + def ffn_service_status_file(self) -> Path: + """Base status file; rank r writes ``ffn-service-rank-.json``.""" + return self.environment.output_path / FFN_SERVICE_STATUS_FILE + + +def dbo_arguments(dbo: DboConfig | None) -> list[str]: + if dbo is None: + return [] + flags = ["--enable-dbo"] + if dbo.decode_token_threshold is not None: + flags.extend(["--dbo-decode-token-threshold", str(dbo.decode_token_threshold)]) + if dbo.prefill_token_threshold is not None: + flags.extend( + ["--dbo-prefill-token-threshold", str(dbo.prefill_token_threshold)] + ) + return flags + + +def launch_order(profile: Deployment) -> tuple[str, ...]: + if profile.is_baseline: + return ("baseline",) + for backend, connector, order in CONNECTOR_PROTOCOLS: + if (profile.backend, profile.connector) == (backend, connector): + return order + raise ValueError( + f"unsupported connector protocol: {profile.backend}/{profile.connector}" + ) + + +def resolve( + case: Case, environment: CaseEnvironment, *, slots: tuple[RoleSlot, ...] = () +) -> ResolvedDeployment: + profile = case.deployment + if environment.device_backend != profile.backend: + raise ValueError( + f"Case requires {profile.backend}, got {environment.device_backend}" + ) + if profile.cleanup_profile not in ( + CLEANUP_PROFILE_STRICT, + CLEANUP_PROFILE_NPU_ASYNC_CAM_LEGACY, + ): + raise ValueError(f"unknown cleanup profile {profile.cleanup_profile!r}") + if profile.cleanup_profile == CLEANUP_PROFILE_NPU_ASYNC_CAM_LEGACY and ( + profile.backend != "npu" or profile.connector != "CAMAsyncAFDConnector" + ): + raise ValueError("legacy NPU async cleanup requires the NPU async CAM protocol") + if profile.ffn_service is not None and (profile.ffn is None or slots): + raise ValueError("the FFN service runs a local FFN role on one host") + if any(role.uses_graph for _, role in profile.roles) and ( + profile.cudagraph_capture_size is None or profile.cudagraph_capture_size < 1 + ): + raise ValueError("graph deployment requires a positive capture size") + arguments = ( + *case.common_vllm_args, + *(arg for _, layout in profile.roles for arg in layout.vllm_args), + ) + for argument in arguments: + name = argument.split("=", 1)[0].replace("_", "-") + if name in DEPLOYMENT_ARGUMENTS: + raise ValueError(f"{name} must be declared by Deployment or Check") + required = ( + sum(len(slot.devices) for slot in slots) + if slots + else sum(role.rank_count for _, role in profile.roles) + ) + if len(environment.devices) < required: + raise ValueError( + f"Case needs {required} devices, got {len(environment.devices)}" + ) + if slots: + role_bindings = tuple( + ( + slot.role, + tuple(environment.devices[int(device)] for device in slot.devices), + ) + for slot in slots + ) + return ResolvedDeployment( + case, environment, role_bindings, launch_order(profile), slots + ) + bound = [] + offset = 0 + for role, layout in profile.roles: + bound.append((role, environment.devices[offset : offset + layout.rank_count])) + offset += layout.rank_count + return ResolvedDeployment(case, environment, tuple(bound), launch_order(profile)) + + +class ServiceDeployment: + """The single owner of all local service processes for one deployment.""" + + def __init__( + self, + resolved: ResolvedDeployment, + *, + run_id: str, + probe_arguments: tuple[str, ...] = (), + probe_environment: dict[str, str] | None = None, + extra_environment: Mapping[str, str] | None = None, + ) -> None: + self.resolved = resolved + self.run_id = run_id + self.probe_arguments = probe_arguments + self.probe_environment = probe_environment or {} + self.extra_environment = extra_environment or {} + self.processes: list[subprocess.Popen[str]] = [] + self.processes_by_role: dict[str, subprocess.Popen[str]] = {} + self.log_threads: list[threading.Thread] = [] + self.execution_evidence: list[tuple[str, str]] = [] + + def start(self) -> None: + for role in self.resolved.launch_order: + self.start_role(role) + if self.resolved.case.deployment.ffn_service is not None: + self.wait_ffn_service_ready() + self.wait_ready() + + def wait_ffn_service_ready(self) -> None: + """Wait until every FFN service rank reports ``/ready`` (paired with A). + + ``/ready`` means paired and receiving; CUDA graphs are captured later, + driven by Attention, whose API comes up only after that. + """ + profile = self.resolved.case.deployment + assert profile.ffn is not None + env = self.resolved.environment + pending = set(range(profile.ffn.rank_count)) + deadline = time.monotonic() + env.startup_timeout + last_error: BaseException | None = None + while pending and time.monotonic() < deadline: + self.ensure_alive() + for rank in sorted(pending): + try: + with urllib.request.urlopen( + self.ffn_service_url(rank, "/ready"), + timeout=FFN_SERVICE_HTTP_TIMEOUT_S, + ) as response: + if response.status == 200: + pending.discard(rank) + except (OSError, urllib.error.URLError) as exc: + last_error = exc + if pending: + time.sleep(READY_POLL_INTERVAL_S) + if pending: + raise TimeoutError( + f"FFN service ranks {sorted(pending)} were not ready: {last_error}" + ) + + def ffn_service_url(self, rank: int, path: str) -> str: + env = self.resolved.environment + port = self.resolved.ffn_service_status_port(rank) + return f"http://{env.api_host}:{port}{path}" + + def finish(self) -> None: + """Stop a deployment whose protocol defines a clean stop; else no-op. + + The standalone FFN service stops idle: ``/stop`` on every rank cancels + its idle control receive, then Attention shuts down right away, since + releasing DeepEP buffers is a barrier over every M2N rank. The service + launcher must exit 0 and every rank must report ``stopped`` with + released resources. Other deployments stop in ``close``. + """ + profile = self.resolved.case.deployment + service = profile.ffn_service + if service is None: + return + assert profile.ffn is not None + for rank in range(profile.ffn.rank_count): + request = urllib.request.Request( + self.ffn_service_url(rank, "/stop"), method="POST" + ) + with urllib.request.urlopen(request, timeout=FFN_SERVICE_HTTP_TIMEOUT_S): + pass + failures = terminate_process_groups( + [self.processes_by_role[ATTENTION_ROLE]], + termination_timeout_s=profile.termination_timeout_s, + poll_interval_s=PROCESS_POLL_INTERVAL_S, + reap_timeout_s=PROCESS_REAP_TIMEOUT_S, + ) + if failures: + raise RuntimeError("; ".join(failures)) + returncode = self.processes_by_role[FFN_ROLE].wait( + timeout=service.stop_timeout_s + ) + if returncode != 0: + raise RuntimeError(f"FFN service exited with {returncode}") + status_file = self.resolved.ffn_service_status_file + for rank in range(profile.ffn.rank_count): + rank_file = status_file.with_name(f"{status_file.stem}-rank-{rank}.json") + status = json.loads(rank_file.read_text()) + if status["phase"] != "stopped" or status["resources_released"] is not True: + raise RuntimeError( + f"FFN service rank {rank} did not stop cleanly: {status}" + ) + + def start_role(self, role: str) -> None: + """Launch one local role; multi-pod owns barriers between these calls.""" + resolved = self.resolved + env = resolved.environment + role_devices = dict(resolved.role_devices) + env.output_path.mkdir(parents=True, exist_ok=True) + command = resolved.command(role, self.probe_arguments) + visible = ",".join(role_devices[role]) + process_env = os.environ.copy() + process_env.update(self.extra_environment) + backend = resolved.case.deployment.backend + process_env[ + "ASCEND_RT_VISIBLE_DEVICES" if backend == "npu" else "CUDA_VISIBLE_DEVICES" + ] = visible + process_env["VLLM_PLUGINS"] = ( + ("ascend" if backend == "npu" else "") + if resolved.case.deployment.is_baseline + else ("ascend,afd" if backend == "npu" else "afd") + ) + process_env["VLLM_USE_V2_MODEL_RUNNER"] = "1" + process_env["AFD_E2E_RUN_ID"] = self.run_id + process_env["AFD_E2E_PROCESS_ROLE"] = role + process_env["PYTHONUNBUFFERED"] = "1" + process_env.setdefault("VLLM_ENGINE_READY_TIMEOUT_S", "18000") + worker_timeout = resolved.case.deployment.worker_shutdown_timeout_s + if worker_timeout is not None: + process_env.setdefault(WORKER_SHUTDOWN_TIMEOUT_ENV, str(worker_timeout)) + process_env.pop("AFD_PLUGIN_EARLY_ENGINE_PATCH", None) + process_env.pop("AFD_E2E_HASH_INPUT_IDS", None) + process_env.pop("AFD_E2E_HASH_REQUIRE_GRAPH", None) + process_env.pop("AFD_E2E_NATIVE_BASELINE", None) + process_env.pop("AFD_E2E_MRV2_EVIDENCE_START", None) + if backend == "npu": + process_env.pop("VLLM_ASCEND_ENABLE_FLASHCOMM1", None) + if resolved.case.deployment.dbo_for_role(role) is not None: + process_env["VLLM_LOGGING_LEVEL"] = "DEBUG" + process_env.update(self.probe_environment) + inherited = process_env.get("PYTHONPATH", "") + process_env["PYTHONPATH"] = os.pathsep.join( + path for path in (str(REPO_ROOT), inherited) if path + ) + print(f"\n=== Starting {role}: {shlex.join(command)} ===", flush=True) + process = subprocess.Popen( + command, + cwd=REPO_ROOT, + env=process_env, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + bufsize=1, + start_new_session=True, + ) + self.processes.append(process) + self.processes_by_role[role] = process + thread = threading.Thread(target=self.stream, args=(role, process), daemon=True) + self.log_threads.append(thread) + thread.start() + self.ensure_alive() + + def stream(self, role: str, process: subprocess.Popen[str]) -> None: + path = self.resolved.environment.output_path / f"{role}.log" + assert process.stdout is not None + with path.open("w") as log: + for line in process.stdout: + log.write(line) + log.flush() + print(f"[{role}] {line}", end="", flush=True) + if MRV2_EVIDENCE_PREFIX in line: + self.execution_evidence.append((role, line.strip())) + + def ensure_alive(self) -> None: + for process in self.processes: + if process.poll() is not None: + raise RuntimeError( + f"service {process.pid} exited with {process.returncode}" + ) + + def wait_ready( + self, *, path: str = "/v1/models", on_poll: Callable[[], None] | None = None + ) -> None: + env = self.resolved.environment + url = f"http://{env.api_host}:{env.api_port_base}{path}" + deadline = time.monotonic() + env.startup_timeout + last_error: BaseException | None = None + while time.monotonic() < deadline: + self.ensure_alive() + if on_poll is not None: + on_poll() + try: + with urllib.request.urlopen(url, timeout=5) as response: + if response.status == 200: + return + except (OSError, urllib.error.URLError) as exc: + last_error = exc + time.sleep(READY_POLL_INTERVAL_S) + raise TimeoutError(f"service did not become ready at {url}: {last_error}") + + def close(self) -> None: + legacy_async = ( + self.resolved.case.deployment.cleanup_profile + == CLEANUP_PROFILE_NPU_ASYNC_CAM_LEGACY + ) + ffn = self.processes_by_role.get("ffn") + failures = terminate_process_groups( + self.processes, + termination_timeout_s=self.resolved.case.deployment.termination_timeout_s, + poll_interval_s=PROCESS_POLL_INTERVAL_S, + reap_timeout_s=PROCESS_REAP_TIMEOUT_S, + deferred_sigkill_pgids=(ffn.pid,) + if legacy_async and ffn is not None + else (), + ) + # torchrun starts the FFN service workers outside the launcher's process + # group, so the service also needs the environment-matched sweep. + ffn_service = self.resolved.case.deployment.ffn_service is not None + if legacy_async or ffn_service: + failures.extend( + kill_processes_matching_environment( + {"AFD_E2E_RUN_ID": self.run_id, "AFD_E2E_PROCESS_ROLE": "ffn"}, + timeout_s=FFN_SERVICE_PROCESS_KILL_TIMEOUT_S + if ffn_service + else NPU_ASYNC_PROCESS_KILL_TIMEOUT_S, + poll_interval_s=PROCESS_POLL_INTERVAL_S, + process_name="FFN", + ) + ) + for thread in self.log_threads: + thread.join(timeout=LOG_THREAD_JOIN_TIMEOUT_S) + if thread.is_alive(): + failures.append("service log reader did not stop") + if failures: + raise RuntimeError("; ".join(failures)) diff --git a/tests/e2e/environment.py b/tests/e2e/environment.py index 359c0f577..9d270d99b 100644 --- a/tests/e2e/environment.py +++ b/tests/e2e/environment.py @@ -1,10 +1,89 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project -"""Environment parsing shared by local NPU E2E entrypoints.""" +"""Resource and environment parsing shared by E2E entrypoints.""" +import argparse import os +import sys from pathlib import Path +from tests.e2e.case import ( + DEFAULT_API_PORT, + DEFAULT_CONNECTOR_PORT, + DEFAULT_STARTUP_TIMEOUT_S, + CaseEnvironment, +) + + +def add_case_environment_arguments( + parser: argparse.ArgumentParser, + case_ids: tuple[str, ...], + *, + devices_required: bool = True, + startup_option: str = "--startup-timeout", +) -> None: + parser.add_argument("--scenario", choices=case_ids, required=True) + parser.add_argument("--model", required=True) + parser.add_argument("--vllm-bin", default="vllm") + parser.add_argument("--device-backend", choices=("gpu", "npu"), required=True) + parser.add_argument("--devices", required=devices_required, default="") + parser.add_argument("--output-path", type=Path, required=True) + parser.add_argument("--api-host", default="127.0.0.1") + parser.add_argument( + "--api-port-base", + type=int, + default=int(os.environ.get("AFD_E2E_API_PORT_BASE", str(DEFAULT_API_PORT))), + ) + parser.add_argument("--afd-host", default="127.0.0.1") + parser.add_argument( + "--afd-port", + type=int, + default=int(os.environ.get("AFD_E2E_AFD_PORT", str(DEFAULT_CONNECTOR_PORT))), + ) + parser.add_argument( + startup_option, + dest="startup_timeout", + type=float, + default=DEFAULT_STARTUP_TIMEOUT_S, + ) + parser.add_argument( + "--describe", + action="store_true", + help="Print configuration and commands without starting services", + ) + + +def parse_case_environment( + case_ids: tuple[str, ...], argv: list[str] | None = None +) -> tuple[str, CaseEnvironment, bool]: + """Read resources in the subprocess; collection never calls this function.""" + parser = argparse.ArgumentParser(description="Run a model-declared E2E Case") + add_case_environment_arguments(parser, case_ids) + args = parser.parse_args(argv) + return ( + args.scenario, + case_environment(args), + args.describe, + ) + + +def case_environment(args: argparse.Namespace) -> CaseEnvironment: + """Bind parsed resources; model and placement decisions live elsewhere.""" + return CaseEnvironment( + model=str(Path(args.model).expanduser()), + vllm_bin=args.vllm_bin, + device_backend=args.device_backend, + devices=tuple( + value.strip() for value in args.devices.split(",") if value.strip() + ), + output_path=args.output_path, + api_host=args.api_host, + api_port_base=args.api_port_base, + afd_host=args.afd_host, + afd_port=args.afd_port, + startup_timeout=args.startup_timeout, + ) + def required_env(name: str) -> str: value = os.environ.get(name) @@ -13,6 +92,35 @@ def required_env(name: str) -> str: return value +def model_case_command( + module: str, + scenario: str, + output_path: Path, + *, + backend: str, + default_devices: tuple[str, ...] = ("0", "1", "2", "3"), +) -> list[str]: + """Bind a model entry to resources without selecting its topology or checks.""" + prefix = f"AFD_{backend.upper()}_E2E" + return [ + sys.executable, + "-m", + module, + "--scenario", + scenario, + "--model", + required_env(f"{prefix}_MODEL"), + "--vllm-bin", + os.environ.get(f"{prefix}_VLLM_BIN", "vllm"), + "--device-backend", + backend, + "--devices", + os.environ.get("AFD_E2E_DEVICES", ",".join(default_devices)), + "--output-path", + str(output_path), + ] + + def devices_from_env(name: str, expected_count: int) -> list[str]: devices = [item.strip() for item in required_env(name).split(",") if item.strip()] if len(devices) != expected_count: @@ -27,3 +135,39 @@ def prepend_env_paths(env: dict[str, str], name: str, *paths: Path) -> None: env[name] = os.pathsep.join( dict.fromkeys([*(str(path) for path in paths), *existing]), ) + + +def model_multi_pod_command(module: str, scenario: str, layout: str) -> list[str]: + """Bind an existing GPU pod to a model-owned case; provisioning stays outside.""" + command = [ + sys.executable, + "-m", + module, + "multi-pod", + "--scenario", + scenario, + "--pod-layout", + layout, + "--run-id", + required_env("AFD_E2E_RUN_ID"), + "--model", + required_env("AFD_GPU_E2E_MODEL"), + "--vllm-bin", + os.environ.get("AFD_GPU_E2E_VLLM_BIN", "vllm"), + "--device-backend", + "gpu", + "--output-path", + required_env("AFD_E2E_GSM8K_OUTPUT"), + "--store-host", + required_env("AFD_E2E_STORE_HOST"), + ] + for variable, option in ( + ("AFD_E2E_STORE_PORT", "--store-port"), + ("AFD_E2E_DEVICES", "--devices"), + ): + if value := os.environ.get(variable): + command.extend([option, value]) + for value in os.environ.get("AFD_E2E_POD_ENV", "").split(";"): + if value.strip(): + command.extend(["--pod-env", value.strip()]) + return command diff --git a/tests/e2e/models/deepseek_v2_lite/cases.py b/tests/e2e/models/deepseek_v2_lite/cases.py new file mode 100644 index 000000000..ca1f97e7e --- /dev/null +++ b/tests/e2e/models/deepseek_v2_lite/cases.py @@ -0,0 +1,228 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""DeepSeek-V2-Lite IDs and model-local Case construction. + +Collection imports IDs only. build_case executes in the runner subprocess. +GPU and NPU declarations share capabilities, not scenario selection in the runner. +""" + +from __future__ import annotations + +import sys +from dataclasses import replace + +from tests.e2e.case import ( + ATTENTION_ROLE, + FFN_ROLE, + Case, + CaseEnvironment, + DboConfig, + Gsm8kSettings, + Gsm8kTask, + Mrv2Check, + Phase, +) +from tests.e2e.presets import DECODE_GRAPH, EAGER + +V2_SCENARIOS = ( + "afd-v2-eager-1a1f", + "afd-v2-eager-dp2", + "afd-v2-eager-tp2", + "afd-v2-graph-1a1f", + "afd-v2-graph-dp2", + "afd-v2-graph-tp2", + "afd-v2-eager-dbo-dp2", + "afd-v2-graph-dbo-dp2", + "afd-v2-graph-dbo-tp2", +) +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", +) +GPU_IDS = ( + "baseline-graph", + *V2_SCENARIOS, + "afd-eager-2a1f", + "afd-graph-2a1f", + "afd-graph-dbo-2a1f", +) +NPU_IDS = ( + "baseline-graph", + "afd-eager-2a2f", + "afd-graph-2a2f", + "afd-graph-dbo-2a2f", + "afd-eager-2a1f", + "afd-graph-2a1f", + "afd-graph-dbo-2a1f", +) +V2_DBO_COMPARISON_SCENARIOS = frozenset( + ("afd-v2-eager-dp2", "afd-v2-eager-dbo-dp2", "afd-v2-graph-dbo-dp2"), +) +V2_SINGLE_RANK_SCENARIOS = frozenset(("afd-v2-eager-1a1f", "afd-v2-graph-1a1f")) +MULTI_POD_IDS = ("afd-graph-2a2f", "afd-graph-dbo-2a2f") +DEFAULT_GSM8K_SAMPLE_LIMIT = 7 +V2_DBO_SAMPLE_LIMIT = 128 +SMOKE_THRESHOLD = 0.25 +ACCURACY_THRESHOLD = 0.27 +SMOKE_MAX_SAMPLES = 24 +COMPARISON_CONCURRENCY = 12 +CAPTURE_SIZE = 8 +MAX_MODEL_LEN = 4096 +DBO_DECODE_THRESHOLD = 2 +DBO_PREFILL_THRESHOLD = 8 + + +def stable_ids(backend: str) -> tuple[str, ...]: + if backend == "gpu": + return GPU_IDS + if backend == "npu": + return NPU_IDS + raise ValueError(f"unknown backend {backend!r}") + + +def build_case(case_id: str, environment: CaseEnvironment) -> Case: + backend = environment.device_backend + allowed = stable_ids(backend) + (DEEPEP_M2N_SCENARIOS if backend == "gpu" else ()) + if case_id not in allowed: + raise ValueError(f"unsupported {backend} DeepSeek-V2-Lite case {case_id!r}") + return _build_case(case_id, environment) + + +def build_multi_pod_case(case_id: str, environment: CaseEnvironment) -> Case: + if environment.device_backend != "gpu" or case_id not in MULTI_POD_IDS: + raise ValueError(f"unsupported GPU multi-pod case {case_id!r}") + case = _build_case(case_id, environment) + dbo = case.deployment.dbo is not None + # These existing distributed cases gate accuracy, liveness and cleanup. + # Their original 0.27 threshold is independent of the single-host smoke rule. + return replace( + case, + phases=( + Phase( + "evaluation", + Gsm8kTask( + Gsm8kSettings( + sample_limit=SMOKE_MAX_SAMPLES + if dbo + else DEFAULT_GSM8K_SAMPLE_LIMIT, + threshold=ACCURACY_THRESHOLD, + num_concurrent=COMPARISON_CONCURRENCY if dbo else None, + minimum_samples=SMOKE_MAX_SAMPLES if dbo else 1, + ) + ), + ), + ), + ) + + +def _build_case(case_id: str, environment: CaseEnvironment) -> Case: + # The table is model-local. No shared module selects a model or case ID. + # baseline, graph, dbo, Attention DP/TP, FFN DP/TP + profiles = { + "baseline-graph": (True, True, False, 4, 1, 0, 1), + "afd-v2-eager-1a1f": (False, False, False, 1, 1, 1, 1), + "afd-v2-eager-dp2": (False, False, False, 2, 1, 2, 1), + "afd-v2-eager-tp2": (False, False, False, 1, 2, 1, 2), + "afd-v2-graph-1a1f": (False, True, False, 1, 1, 1, 1), + "afd-v2-graph-dp2": (False, True, False, 2, 1, 2, 1), + "afd-v2-graph-tp2": (False, True, False, 1, 2, 1, 2), + "afd-v2-eager-dbo-dp2": (False, False, True, 2, 1, 2, 1), + "afd-v2-graph-dbo-dp2": (False, True, True, 2, 1, 2, 1), + "afd-v2-graph-dbo-tp2": (False, True, True, 1, 2, 1, 2), + "afd-eager-2a1f": (False, False, False, 2, 1, 1, 1), + "afd-graph-2a1f": (False, True, False, 2, 1, 1, 1), + "afd-graph-dbo-2a1f": (False, True, True, 2, 1, 1, 1), + "afd-eager-2a2f": (False, False, False, 2, 1, 2, 1), + "afd-graph-2a2f": (False, True, False, 2, 1, 2, 1), + "afd-graph-dbo-2a2f": (False, True, True, 2, 1, 2, 1), + "afd-deepep-m2n-eager-2a2f": (False, False, False, 2, 1, 2, 1), + "afd-deepep-m2n-graph-2a2f": (False, True, False, 2, 1, 2, 1), + "afd-deepep-m2n-graph-dbo-2a2f": (False, True, True, 2, 1, 2, 1), + "afd-deepep-m2n-eager-tp2-2a2f": (False, False, False, 1, 2, 2, 1), + "afd-deepep-m2n-graph-tp2-2a2f": (False, True, False, 1, 2, 2, 1), + } + backend = environment.device_backend + baseline, graph, dbo, a_dp, a_tp, f_dp, f_tp = profiles[case_id] + execution = DECODE_GRAPH if graph else EAGER + comparison = case_id in V2_DBO_COMPARISON_SCENARIOS + deepep = case_id in DEEPEP_M2N_SCENARIOS + minimum_samples = SMOKE_MAX_SAMPLES if comparison or dbo else 1 + sample_limit = ( + V2_DBO_SAMPLE_LIMIT + if comparison + else max(DEFAULT_GSM8K_SAMPLE_LIMIT, minimum_samples) + ) + # The TP2 DBO case has always retained 0.27 even for a 24-sample run. + smoke = case_id != "afd-v2-graph-dbo-tp2" + deployment = execution.deployment( + backend=backend, + connector=( + None + if baseline + else "DeepEPM2NAFDConnector" + if deepep + else "P2pNcclAFDConnector" + if backend == "gpu" + else "CAMP2pAFDConnector" + ), + attention=(a_dp, a_tp), + ffn=None if baseline else (f_dp, f_tp), + capture_size=CAPTURE_SIZE, + compute_gate_on_attention=deepep, + dbo=DboConfig(DBO_DECODE_THRESHOLD, DBO_PREFILL_THRESHOLD) if dbo else None, + ) + if comparison: + deployment = replace(deployment, max_num_seqs=CAPTURE_SIZE) + return Case( + id=case_id, + model="DeepSeek-V2-Lite", + served_model_name_prefix="deepseek-v2-lite-afd", + deployment=deployment, + common_vllm_args=(f"--max-model-len={MAX_MODEL_LEN}",), + phases=( + Phase( + name="evaluation", + task=Gsm8kTask( + Gsm8kSettings( + sample_limit=sample_limit, + threshold=ACCURACY_THRESHOLD + if comparison or not smoke + else SMOKE_THRESHOLD, + num_concurrent=COMPARISON_CONCURRENCY + if comparison or dbo + else None, + minimum_samples=minimum_samples, + smoke_max_samples=SMOKE_MAX_SAMPLES if smoke else None, + smoke_threshold=SMOKE_THRESHOLD if smoke else None, + larger_sample_threshold=ACCURACY_THRESHOLD if smoke else None, + ), + ), + checks=( + (Mrv2Check((ATTENTION_ROLE, FFN_ROLE), graph=graph, dbo=dbo),) + if not baseline and (graph or dbo) + else () + ), + ), + ), + ) + + +if __name__ == "__main__": + if sys.argv[1:2] == ["multi-pod"]: + from tests.e2e.multi_pod.runner import case_main as multi_pod_case_main + + raise SystemExit( + multi_pod_case_main(build_multi_pod_case, MULTI_POD_IDS, sys.argv[2:]) + ) + # The shared entry receives the builder, and never imports this model. + from tests.e2e.runner import case_main + + raise SystemExit( + case_main( + build_case, + tuple(dict.fromkeys((*GPU_IDS, *NPU_IDS, *DEEPEP_M2N_SCENARIOS))), + ) + ) diff --git a/tests/e2e/models/deepseek_v2_lite/test_deepseek_v2_lite.py b/tests/e2e/models/deepseek_v2_lite/test_deepseek_v2_lite.py index fa46b4118..f171dcd47 100644 --- a/tests/e2e/models/deepseek_v2_lite/test_deepseek_v2_lite.py +++ b/tests/e2e/models/deepseek_v2_lite/test_deepseek_v2_lite.py @@ -5,7 +5,6 @@ from __future__ import annotations import os -import sys from collections.abc import Iterator from pathlib import Path @@ -15,46 +14,22 @@ download_dataset, download_model, preserve_environment_variable, - run_runner, + run_case_entry, ) -from tests.e2e.runner import ( +from tests.e2e.environment import model_case_command +from tests.e2e.models.deepseek_v2_lite.cases import ( DEEPEP_M2N_SCENARIOS, - DEFAULT_GSM8K_SAMPLE_LIMIT, - GSM8K_LIMIT_ENV, - GSM8K_THRESHOLD_ENV, - V2_DBO_COMPARISON_SCENARIOS, - V2_DBO_SAMPLE_LIMIT, - V2_SCENARIOS, V2_SINGLE_RANK_SCENARIOS, + stable_ids, ) GSM8K_DATASET_ID = "openai/gsm8k" GSM8K_DATASET_CONFIG = "main" DEEPSEEK_V2_LITE_REPO_ID = "deepseek-ai/DeepSeek-V2-Lite" -DEEPSEEK_V2_LITE_MAX_MODEL_LEN = 4096 -SMOKE_MIN_ACCURACY = 0.25 -SMOKE_MAX_SAMPLES = 24 -DEFAULT_DEVICE_IDS = ("0", "1", "2", "3") -ATTENTION_DEVICE_COUNT = 2 -AFD_FFN_DEVICE_COUNT = 1 -AFD_TWO_FFN_DEVICE_COUNT = 2 -BASELINE_DEVICE_COUNT = 4 -TWO_FFN_SCENARIO_SUFFIX = "-2a2f" # CI selects exact node IDs (E2E-INV-007). GPU uses the V2 cases for the # shared 2A2F coverage; NPU retains its existing scenario IDs and now runs # every case on native Ascend MRV2 through the shared harness. -SCENARIOS = ( - "baseline-graph", - *( - ("afd-eager-2a2f", "afd-graph-2a2f", "afd-graph-dbo-2a2f") - if os.environ.get("AFD_E2E_BACKEND") == "npu" - else V2_SCENARIOS - ), - # GPU weekly CI selects Graph + DBO; eager/graph remain local-only. - "afd-eager-2a1f", - "afd-graph-2a1f", - "afd-graph-dbo-2a1f", -) +SCENARIOS = stable_ids(os.environ.get("AFD_E2E_BACKEND", "gpu")) def _required_env(name: str) -> str: @@ -89,65 +64,15 @@ def prepare_e2e_assets() -> None: def build_runner_command(scenario: str, gsm8k_output_path: Path) -> list[str]: - backend = _required_env("AFD_E2E_BACKEND") - if backend == "gpu": - model = _required_env("AFD_GPU_E2E_MODEL") - vllm_bin = os.environ.get("AFD_GPU_E2E_VLLM_BIN", "vllm") - elif backend == "npu": - model = _required_env("AFD_NPU_E2E_MODEL") - vllm_bin = os.environ.get("AFD_NPU_E2E_VLLM_BIN", "vllm") - else: - raise RuntimeError("AFD_E2E_BACKEND must be 'gpu' or 'npu'") - - env_devices = os.environ.get("AFD_E2E_DEVICES") - if env_devices: - devices = [item.strip() for item in env_devices.split(",") if item.strip()] - elif scenario in V2_SINGLE_RANK_SCENARIOS: - devices = ["0", "1"] - else: - devices = list(DEFAULT_DEVICE_IDS) - if scenario == "baseline-graph": - attention_devices = devices[:BASELINE_DEVICE_COUNT] - ffn_devices = [] - elif scenario in V2_SINGLE_RANK_SCENARIOS: - attention_devices = devices[:1] - ffn_devices = devices[1:2] - else: - attention_devices = devices[:ATTENTION_DEVICE_COUNT] - ffn_device_count = ( - AFD_TWO_FFN_DEVICE_COUNT - if scenario.endswith(TWO_FFN_SCENARIO_SUFFIX) or scenario in V2_SCENARIOS - else AFD_FFN_DEVICE_COUNT - ) - ffn_devices = devices[ - ATTENTION_DEVICE_COUNT : ATTENTION_DEVICE_COUNT + ffn_device_count - ] - - command = [ - sys.executable, - "-m", - "tests.e2e.runner", - "--model", - model, - "--vllm-bin", - vllm_bin, - "--device-backend", - backend, - f"--common-vllm-arg=--max-model-len={DEEPSEEK_V2_LITE_MAX_MODEL_LEN}", - "--attention-devices", - ",".join(attention_devices), - ] - if scenario != "baseline-graph": - command.extend(["--ffn-devices", ",".join(ffn_devices)]) - command.extend( - [ - "--scenario", - scenario, - "--gsm8k-output-path", - str(gsm8k_output_path), - ], + return model_case_command( + "tests.e2e.models.deepseek_v2_lite.cases", + scenario, + gsm8k_output_path, + backend=_required_env("AFD_E2E_BACKEND"), + default_devices=("0", "1") + if scenario in V2_SINGLE_RANK_SCENARIOS + else ("0", "1", "2", "3"), ) - return command @pytest.fixture(scope="module", autouse=True) @@ -166,29 +91,11 @@ def _prepare_e2e_assets() -> Iterator[None]: yield -def _runner_env(scenario: str) -> dict[str, str]: - env = os.environ.copy() - default_limit = ( - V2_DBO_SAMPLE_LIMIT - if scenario in V2_DBO_COMPARISON_SCENARIOS - else DEFAULT_GSM8K_SAMPLE_LIMIT - ) - sample_limit = env.get(GSM8K_LIMIT_ENV, str(default_limit)) - if ( - scenario != "afd-v2-graph-dbo-tp2" - and sample_limit != "all" - and int(sample_limit) <= SMOKE_MAX_SAMPLES - ): - # Model smoke gates require 2/7 or 6/24 correct; larger runs keep 0.27. - env.setdefault(GSM8K_THRESHOLD_ENV, str(SMOKE_MIN_ACCURACY)) - return env - - @pytest.mark.e2e @pytest.mark.parametrize("scenario", SCENARIOS, ids=SCENARIOS) def test_deepseek_v2_lite(scenario: str, tmp_path: Path) -> None: command = build_runner_command(scenario, tmp_path / scenario) - run_runner(command, env=_runner_env(scenario)) + run_case_entry(command, tmp_path / scenario) @pytest.mark.e2e @@ -198,4 +105,4 @@ def test_deepseek_v2_lite(scenario: str, tmp_path: Path) -> None: @pytest.mark.parametrize("scenario", DEEPEP_M2N_SCENARIOS, ids=DEEPEP_M2N_SCENARIOS) def test_deepseek_v2_lite_deepep_m2n(scenario: str, tmp_path: Path) -> None: command = build_runner_command(scenario, tmp_path / scenario) - run_runner(command, env=_runner_env(scenario)) + run_case_entry(command, tmp_path / scenario) diff --git a/tests/e2e/models/deepseek_v2_lite/test_deepseek_v2_lite_multi_pod.py b/tests/e2e/models/deepseek_v2_lite/test_deepseek_v2_lite_multi_pod.py index bc3d9bf87..194148ee5 100644 --- a/tests/e2e/models/deepseek_v2_lite/test_deepseek_v2_lite_multi_pod.py +++ b/tests/e2e/models/deepseek_v2_lite/test_deepseek_v2_lite_multi_pod.py @@ -14,11 +14,13 @@ from __future__ import annotations import os -import sys +from pathlib import Path import pytest -from tests.conftest import run_runner +from tests.conftest import run_case_entry +from tests.e2e.environment import model_multi_pod_command +from tests.e2e.models.deepseek_v2_lite.cases import MULTI_POD_IDS # The layout is a second axis, orthogonal to the scenario id, so accuracy # evidence stays comparable with the single-host rows. @@ -26,13 +28,9 @@ "2pod-role-split": "2A0F,0A2F", "2pod-interleaved": "1A1F,1A1F", } -MULTI_POD_CASES = [ - ("afd-graph-2a2f", "2pod-role-split"), - ("afd-graph-2a2f", "2pod-interleaved"), - ("afd-graph-dbo-2a2f", "2pod-role-split"), - ("afd-graph-dbo-2a2f", "2pod-interleaved"), -] -DEEPSEEK_V2_LITE_MAX_MODEL_LEN = 4096 +MULTI_POD_CASES = tuple( + (scenario, layout) for scenario in MULTI_POD_IDS for layout in POD_LAYOUTS +) def _required_env(name: str) -> str: @@ -58,32 +56,9 @@ def build_runner_command(scenario: str, layout_name: str) -> list[str]: invocation should still get its own value, or a leftover process from a prior run on this same pod could be mistaken for part of this run. """ - layout = POD_LAYOUTS[layout_name] - command = [ - sys.executable, - "-m", - "tests.e2e.multi_pod.runner", - "--scenario", - scenario, - "--pod-layout", - layout, - "--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={DEEPSEEK_V2_LITE_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 + return model_multi_pod_command( + "tests.e2e.models.deepseek_v2_lite.cases", scenario, POD_LAYOUTS[layout_name] + ) @pytest.mark.e2e @@ -93,4 +68,7 @@ def build_runner_command(scenario: str, layout_name: str) -> list[str]: ids=[f"{scenario}-{layout}" for scenario, layout in MULTI_POD_CASES], ) def test_multi_pod(scenario: str, layout_name: str) -> None: - run_runner(build_runner_command(scenario, layout_name)) + run_case_entry( + build_runner_command(scenario, layout_name), + Path(_required_env("AFD_E2E_GSM8K_OUTPUT")), + ) diff --git a/tests/e2e/models/deepseek_v4_flash/cases.py b/tests/e2e/models/deepseek_v4_flash/cases.py new file mode 100644 index 000000000..868928052 --- /dev/null +++ b/tests/e2e/models/deepseek_v4_flash/cases.py @@ -0,0 +1,269 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""DSV4 single-deployment declarations, separate from dedicated/multi-pod runs. + +Importing IDs is safe. Checkpoint quantization is resolved only by build_case +in the runner child. Existing launch-profile helpers remain model-owned and +are also used by the unmigrated dedicated and multi-pod executors. +""" + +from __future__ import annotations + +import argparse +import json + +from tests.e2e.case import ( + CLEANUP_PROFILE_NPU_ASYNC_CAM_LEGACY, + CLEANUP_PROFILE_STRICT, + DEFAULT_SHUTDOWN_TIMEOUT_S, + EXECUTION_EAGER, + EXECUTION_FULL_DECODE_GRAPH, + NPU_ASYNC_TERMINATION_TIMEOUT_S, + NPU_ASYNC_VLLM_SHUTDOWN_TIMEOUT_S, + NPU_ASYNC_WORKER_SHUTDOWN_TIMEOUT_S, + Case, + CaseEnvironment, + DboConfig, + Deployment, + Gsm8kSettings, + Gsm8kTask, + HashInputCheck, + Mrv2Check, + Phase, + Role, +) +from tests.e2e.models.deepseek_v4_flash import config, gpu_config +from tests.e2e.models.deepseek_v4_flash.completions import CompletionTask + +GPU_BASIC_IDS = gpu_config.DSV4_GPU_BASIC_SCENARIOS +GPU_DBO_IDS = gpu_config.DSV4_GPU_DBO_SCENARIOS +NPU_ASYNC_IDS = (config.DSV4_ASYNC_CAM_SCENARIO,) +NPU_SYNC_IDS = config.DSV4_SYNC_CAMP2P_SCENARIOS +GSM8K_SAMPLES = 7 +GSM8K_THRESHOLD = 0.27 +DBO_CONCURRENCY = 12 +DBO_MINIMUM_SAMPLES = 2 * DBO_CONCURRENCY +DBO_DECODE_THRESHOLD = 2 +DBO_PREFILL_THRESHOLD = 8 + + +def _budget(arguments: list[str], name: str) -> int | None: + """Transfer a model profile's budget to its deployment declaration.""" + if name not in arguments: + return None + index = arguments.index(name) + value = int(arguments[index + 1]) + del arguments[index : index + 2] + return value + + +def _profile_arguments( + case_id: str, environment: CaseEnvironment +) -> argparse.Namespace: + args = argparse.Namespace( + scenario=case_id, + model=environment.model, + device_backend=environment.device_backend, + afd_host=environment.afd_host, + common_vllm_arg=[], + attention_vllm_arg=[], + ffn_vllm_arg=[], + completion_output_path=environment.output_path / "completions.json", + gsm8k_output_path=environment.output_path / "gsm8k", + use_decode_bench_connector=False, + ) + if environment.device_backend == "gpu": + gpu_config.configure_scenario(args) + elif case_id in NPU_ASYNC_IDS: + config.configure_scenario(args) + else: + config.configure_sync_camp2p_scenario(args) + return args + + +def build_case(case_id: str, environment: CaseEnvironment) -> Case: + if environment.device_backend == "gpu": + return _gpu_case(case_id, environment) + if environment.device_backend == "npu" and case_id in ( + *NPU_ASYNC_IDS, + *NPU_SYNC_IDS, + ): + return _npu_case(case_id, environment) + raise ValueError(f"unsupported DSV4 {environment.device_backend} case {case_id!r}") + + +def _gpu_case(case_id: str, environment: CaseEnvironment) -> Case: + baseline, graph, dbo, attention_ranks, ffn_ranks = ( + gpu_config.DSV4_GPU_SCENARIO_SETTINGS[case_id] + ) + args = _profile_arguments(case_id, environment) + max_num_seqs = _budget(args.common_vllm_arg, "--max-num-seqs") + max_num_batched_tokens = _budget(args.common_vllm_arg, "--max-num-batched-tokens") + execution = EXECUTION_FULL_DECODE_GRAPH if graph else EXECUTION_EAGER + roles = ("baseline",) if baseline else ("attention", "ffn") + checks = ( + (Mrv2Check(roles, graph=graph, dbo=dbo),) + if not baseline or graph or dbo + else () + ) + phases = [] + for index, name in enumerate(gpu_config.DSV4_GPU_HASH_WINDOWS): + hash_checks = ( + () + if baseline + else ( + HashInputCheck( + graph=graph, + dbo=dbo, + compare_with=gpu_config.DSV4_GPU_HASH_WINDOWS[index - 1] + if index + else None, + ), + ) + ) + phases.append( + Phase( + name, + CompletionTask( + first_request_index=index * config.DSV4_CONCURRENT_REQUESTS, + answer_repetitions=gpu_config.DSV4_GPU_ANSWER_REPETITIONS, + ), + (*checks, *hash_checks), + ) + ) + phases.append( + Phase( + "gsm8k", + Gsm8kTask( + Gsm8kSettings( + sample_limit=DBO_MINIMUM_SAMPLES if dbo else GSM8K_SAMPLES, + threshold=GSM8K_THRESHOLD, + num_concurrent=DBO_CONCURRENCY if dbo else None, + minimum_samples=DBO_MINIMUM_SAMPLES if dbo else 1, + ) + ), + checks, + ) + ) + return Case( + id=case_id, + model="DeepSeek-V4-Flash", + served_model_name_prefix="dsv4-flash-gpu", + deployment=Deployment( + backend="gpu", + attention=Role(attention_ranks, execution=execution), + ffn=None if baseline else Role(ffn_ranks, execution=execution), + connector=None if baseline else "P2pNcclAFDConnector", + dbo=DboConfig(DBO_DECODE_THRESHOLD, DBO_PREFILL_THRESHOLD) if dbo else None, + cudagraph_capture_size=gpu_config.DSV4_GPU_CAPTURE_SIZE if graph else None, + max_num_seqs=max_num_seqs, + max_num_batched_tokens=max_num_batched_tokens, + ), + common_vllm_args=tuple(args.common_vllm_arg), + phases=tuple(phases), + ) + + +def _npu_case(case_id: str, environment: CaseEnvironment) -> Case: + asynchronous = case_id in NPU_ASYNC_IDS + shape = None if asynchronous else config.DSV4_SYNC_SHAPES[case_id] + args = _profile_arguments(case_id, environment) + max_num_seqs = _budget(args.common_vllm_arg, "--max-num-seqs") + max_num_batched_tokens = _budget(args.common_vllm_arg, "--max-num-batched-tokens") + graph = shape is not None and config.sync_use_graph(shape) + execution = EXECUTION_FULL_DECODE_GRAPH if graph else EXECUTION_EAGER + attention_ranks = ( + config.DSV4_ATTENTION_RANKS if shape is None else shape.attention_ranks + ) + ffn_ranks = config.DSV4_FFN_RANKS if shape is None else shape.ffn_ranks + attention_tp = ( + config.DSV4_ATTENTION_TP_SIZE if shape is None else shape.attention_tp_size + ) + ffn_tp = 1 if shape is None else shape.ffn_tp_size + if asynchronous: + # The existing async path enables native DBO only on Attention, using + # upstream defaults; it does not claim a DBO execution acceptance check. + args.attention_vllm_arg.remove("--enable-dbo") + compilation = None if shape is None else config.sync_compilation_config(shape) + return Case( + id=case_id, + model="DeepSeek-V4-Flash", + served_model_name_prefix="dsv4-flash" if asynchronous else "dsv4-flash-sync", + deployment=Deployment( + backend="npu", + attention=Role( + attention_ranks // attention_tp, + attention_tp, + execution, + dbo=DboConfig() if asynchronous else None, + vllm_args=tuple(args.attention_vllm_arg), + additional_config=tuple( + config.additional_config( + "attention", async_cam=asynchronous + ).items() + ), + ), + ffn=Role( + ffn_ranks // ffn_tp, + ffn_tp, + execution, + vllm_args=tuple(args.ffn_vllm_arg), + additional_config=tuple( + config.additional_config("ffn", async_cam=asynchronous).items() + ), + ), + connector=args.afd_connector, + compute_gate_on_attention=args.compute_gate_on_attention, + connector_extra_config=tuple( + ( + json.loads(args.afd_connector_extra_config[0]) + if args.afd_connector_extra_config + else {} + ).items() + ), + pin_attention_address=args.attention_data_parallel_address is not None, + cleanup_profile=CLEANUP_PROFILE_NPU_ASYNC_CAM_LEGACY + if asynchronous + else CLEANUP_PROFILE_STRICT, + shutdown_timeout_s=NPU_ASYNC_VLLM_SHUTDOWN_TIMEOUT_S + if asynchronous + else DEFAULT_SHUTDOWN_TIMEOUT_S, + worker_shutdown_timeout_s=NPU_ASYNC_WORKER_SHUTDOWN_TIMEOUT_S + if asynchronous + else None, + termination_timeout_s=NPU_ASYNC_TERMINATION_TIMEOUT_S + if asynchronous + else config.DSV4_PROCESS_TERMINATION_TIMEOUT_S, + cudagraph_capture_size=( + shape.cudagraph_capture_size if shape is not None and graph else None + ), + compilation_config=tuple( + (key, tuple(value) if isinstance(value, list) else value) + for key, value in compilation.items() + ) + if compilation is not None + else (), + max_num_seqs=max_num_seqs, + max_num_batched_tokens=max_num_batched_tokens, + ), + common_vllm_args=tuple(args.common_vllm_arg), + phases=( + Phase( + "evaluation", + CompletionTask( + check_answer=True if shape is None else shape.check_answer + ), + (Mrv2Check(("attention", "ffn"), graph=True),) if graph else (), + ), + ), + ) + + +if __name__ == "__main__": + from tests.e2e.runner import case_main + + raise SystemExit( + case_main( + build_case, (*GPU_BASIC_IDS, *GPU_DBO_IDS, *NPU_ASYNC_IDS, *NPU_SYNC_IDS) + ) + ) diff --git a/tests/e2e/models/deepseek_v4_flash/completions.py b/tests/e2e/models/deepseek_v4_flash/completions.py index f3302688e..e6c790047 100644 --- a/tests/e2e/models/deepseek_v4_flash/completions.py +++ b/tests/e2e/models/deepseek_v4_flash/completions.py @@ -7,10 +7,11 @@ import asyncio import json import time +from collections.abc import Mapping +from dataclasses import dataclass from pathlib import Path -import httpx - +from tests.e2e.case import JsonValue from tests.e2e.models.deepseek_v4_flash.config import ( DSV4_COMPLETION_MAX_TOKENS, DSV4_CONCURRENT_REQUESTS, @@ -28,6 +29,11 @@ def evaluate_completions( check_answer: bool = True, answer_repetitions: int = 1, first_request_index: int = 0, + num_requests: int = DSV4_CONCURRENT_REQUESTS, + max_tokens: int = DSV4_COMPLETION_MAX_TOKENS, + timeout_s: float = DSV4_REQUEST_TIMEOUT_S, + first_operand: int = DSV4_PROMPT_FIRST_OPERAND, + second_operand: int = DSV4_PROMPT_SECOND_OPERAND, ) -> None: """Run one batch of ten concurrent requests and check every response. @@ -43,12 +49,16 @@ def evaluate_completions( exercises live decode after the answer's first token, rather than only checking a one-token answer computed by the eager prefill. """ + # Case declarations are imported during collection; load the client only + # when the runner executes this task. + import httpx + results: list[dict] = [ { "index": index, "prompt": ( - f"Compute {DSV4_PROMPT_FIRST_OPERAND + index} + " - f"{DSV4_PROMPT_SECOND_OPERAND}. " + f"Compute {first_operand + index} + " + f"{second_operand}. " + ( "Reply with just the number." if answer_repetitions == 1 @@ -58,9 +68,7 @@ def evaluate_completions( ), "error": "Request cancelled before completion", } - for index in range( - first_request_index, first_request_index + DSV4_CONCURRENT_REQUESTS - ) + for index in range(first_request_index, first_request_index + num_requests) ] async def request(client: httpx.AsyncClient, item: dict) -> None: @@ -72,7 +80,7 @@ async def request(client: httpx.AsyncClient, item: dict) -> None: "model": model, "messages": [{"role": "user", "content": item["prompt"]}], "temperature": 0.0, - "max_tokens": DSV4_COMPLETION_MAX_TOKENS, + "max_tokens": max_tokens, "chat_template_kwargs": {"thinking": False}, }, ) @@ -82,9 +90,7 @@ async def request(client: httpx.AsyncClient, item: dict) -> None: response.raise_for_status() item["response"] = response.json() validate_response(item["response"], check_answer=check_answer) - expected = ( - DSV4_PROMPT_FIRST_OPERAND + item["index"] + DSV4_PROMPT_SECOND_OPERAND - ) + expected = first_operand + item["index"] + second_operand expected_content = " ".join([str(expected)] * answer_repetitions) content = item["response"]["choices"][0]["message"]["content"].strip() if check_answer and content != expected_content: @@ -101,7 +107,7 @@ async def run_requests() -> None: # All requests are scheduled before yielding. Async socket operations # cancel immediately, so SIGTERM can reach the runner's service cleanup # without waiting for a thread blocked in a 300-second HTTP read. - async with httpx.AsyncClient(timeout=DSV4_REQUEST_TIMEOUT_S) as client: + async with httpx.AsyncClient(timeout=timeout_s) as client: await asyncio.gather(*(request(client, item) for item in results)) try: @@ -122,11 +128,48 @@ async def run_requests() -> None: if max(item["started_at"] for item in results) >= min( item["finished_at"] for item in results ): - raise RuntimeError("The ten completion requests did not overlap") + raise RuntimeError("The completion requests did not overlap") for item in results: content = item["response"]["choices"][0]["message"]["content"] print(f"Request {item['index'] + 1}: {item['prompt']} -> {content}") - print(f"Concurrent completions: {len(results)}/{DSV4_CONCURRENT_REQUESTS} passed") + print(f"Concurrent completions: {len(results)}/{num_requests} passed") + + +@dataclass(frozen=True) +class CompletionTask: + """The model's concurrent answer/plumbing contract, with recorded workload.""" + + check_answer: bool = True + answer_repetitions: int = 1 + first_request_index: int = 0 + num_requests: int = DSV4_CONCURRENT_REQUESTS + max_tokens: int = DSV4_COMPLETION_MAX_TOKENS + timeout_s: float = DSV4_REQUEST_TIMEOUT_S + first_operand: int = DSV4_PROMPT_FIRST_OPERAND + second_operand: int = DSV4_PROMPT_SECOND_OPERAND + + def resolve( + self, environment: Mapping[str, str] + ) -> tuple[CompletionTask, tuple[str, ...]]: + return self, () + + def execute( + self, *, base_url: str, model_name: str, checkpoint: str, output_path: Path + ) -> dict[str, JsonValue]: + evaluate_completions( + url=f"{base_url}/v1/chat/completions", + model=model_name, + output_path=output_path / "completions.json", + check_answer=self.check_answer, + answer_repetitions=self.answer_repetitions, + first_request_index=self.first_request_index, + num_requests=self.num_requests, + max_tokens=self.max_tokens, + timeout_s=self.timeout_s, + first_operand=self.first_operand, + second_operand=self.second_operand, + ) + return {"request_count": self.num_requests, "answer_checked": self.check_answer} def validate_response(result: dict, *, check_answer: bool = True) -> None: diff --git a/tests/e2e/models/deepseek_v4_flash/ffn_graph.py b/tests/e2e/models/deepseek_v4_flash/ffn_graph.py index 4e83be3c3..71176767f 100644 --- a/tests/e2e/models/deepseek_v4_flash/ffn_graph.py +++ b/tests/e2e/models/deepseek_v4_flash/ffn_graph.py @@ -71,6 +71,7 @@ def profile(args) -> SimpleNamespace: json.dumps( { "dynamicQuant": 1, + "ffn_mega_moe": args.mega, "attn_ranks_per_dp": 4, "hccl_buffer_size": 256, } @@ -332,6 +333,7 @@ def main() -> None: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--mode", required=True, choices=("eager", "full")) parser.add_argument("--model", required=True) + parser.add_argument("--mega", action="store_true") parser.add_argument("--output-dir", type=Path, required=True) parser.add_argument("--reference-dir", type=Path) parser.add_argument("--vllm-bin", default="vllm") diff --git a/tests/e2e/models/deepseek_v4_flash/ffn_graph_evidence.py b/tests/e2e/models/deepseek_v4_flash/ffn_graph_evidence.py index 2c7f52f29..ac1c31e15 100644 --- a/tests/e2e/models/deepseek_v4_flash/ffn_graph_evidence.py +++ b/tests/e2e/models/deepseek_v4_flash/ffn_graph_evidence.py @@ -9,10 +9,12 @@ import time from functools import wraps from pathlib import Path +from types import SimpleNamespace import torch from afd_plugin.connectors.npu.async_cam import CAMAsyncAFDConnector +from afd_plugin.model_executor.npu.async_cam_mega import AsyncCAMMegaExecutor from afd_plugin.v1.worker.npu.async_ffn_model_runner import AsyncNPUFFNModelRunner from afd_plugin.v1.worker.npu.attention_model_runner_v2 import ( AFDNPUAttentionModelRunnerV2, @@ -72,6 +74,7 @@ def afd_ffn_graph_snapshot(self) -> dict: def install_observers() -> None: complete = AsyncNPUFFNModelRunner._complete_graph send_eager = CAMAsyncAFDConnector.send_ffn_output + execute_mega = AsyncCAMMegaExecutor.__call__ execute = AFDNPUAttentionModelRunnerV2.execute_model shutdown_attention = AFDNPUAttentionModelRunnerV2.shutdown shutdown_ffn = AsyncNPUFFNModelRunner.shutdown @@ -83,6 +86,8 @@ def observe(connector, context): # RPC snapshots read this host copy without touching device buffers. state = context.states header = state.token_nums_rankid_layeridx.cpu().tolist() + if header[2] < 0: + return # Pull idle returns are not layer transactions. rows = int(state.group_list.sum().item()) previous = getattr(connector, "_e2e_completion", {}) layers = dict(previous.get("layers", {})) @@ -96,11 +101,33 @@ def observe(connector, context): "zero_rows": previous.get("zero_rows", 0) + int(rows == 0), } + def mega_context(executor): + return SimpleNamespace( + states=SimpleNamespace( + token_nums_rankid_layeridx=executor.batch_info, + group_list=executor.group_list, + ) + ) + @wraps(complete) def completed(self): result = complete(self) if result: - observe(self.connector, self._graph_payload.context) + context = ( + mega_context(self._mega_executor) + if self._mega_executor is not None + else self._graph_payload.context + ) + observe(self.connector, context) + return result + + @wraps(execute_mega) + def mega_executed(self): + result = execute_mega(self) + if self.connector.graph_startup.mode == "eager": + # Never read device headers while FULL is warming up or capturing. + torch.npu.current_stream().synchronize() + observe(self.connector, mega_context(self)) return result @wraps(send_eager) @@ -168,6 +195,7 @@ def ffn_stop_requested(self): AsyncNPUFFNModelRunner._complete_graph = completed CAMAsyncAFDConnector.send_ffn_output = sent + AsyncCAMMegaExecutor.__call__ = mega_executed AFDNPUAttentionModelRunnerV2.execute_model = executed AFDNPUAttentionModelRunnerV2.shutdown = attention_closed AsyncNPUFFNModelRunner.shutdown = ffn_closed diff --git a/tests/e2e/models/deepseek_v4_flash/test_async_cam_npu.py b/tests/e2e/models/deepseek_v4_flash/test_async_cam_npu.py index 50d02550d..a40def0ce 100644 --- a/tests/e2e/models/deepseek_v4_flash/test_async_cam_npu.py +++ b/tests/e2e/models/deepseek_v4_flash/test_async_cam_npu.py @@ -5,13 +5,12 @@ from __future__ import annotations import os -import sys from pathlib import Path import pytest -from tests.conftest import run_runner -from tests.e2e.environment import devices_from_env, required_env +from tests.conftest import run_case_entry +from tests.e2e.environment import devices_from_env, model_case_command, required_env from tests.e2e.models.deepseek_v4_flash.config import ( DSV4_ASYNC_CAM_SCENARIO, DSV4_ATTENTION_RANKS, @@ -22,36 +21,26 @@ def build_runner_command(output_path: Path) -> list[str]: if required_env("AFD_E2E_BACKEND") != "npu": raise RuntimeError("DSV4 async CAM E2E requires AFD_E2E_BACKEND=npu") - devices = devices_from_env("AFD_E2E_DEVICES", DSV4_ATTENTION_RANKS + DSV4_FFN_RANKS) - return [ - sys.executable, - "-m", - "tests.e2e.runner", - "--model", - required_env("AFD_NPU_E2E_MODEL"), - "--vllm-bin", - os.environ.get("AFD_NPU_E2E_VLLM_BIN", "vllm"), - "--device-backend", - "npu", - "--attention-devices", - ",".join(devices[:DSV4_ATTENTION_RANKS]), - "--ffn-devices", - ",".join(devices[DSV4_ATTENTION_RANKS:]), - "--scenario", + devices_from_env("AFD_E2E_DEVICES", DSV4_ATTENTION_RANKS + DSV4_FFN_RANKS) + command = model_case_command( + "tests.e2e.models.deepseek_v4_flash.cases", DSV4_ASYNC_CAM_SCENARIO, - "--served-model-name-prefix", - "dsv4-flash", - "--afd-host", - required_env("HCCL_IF_IP"), - "--api-port-base", - os.environ.get("AFD_NPU_DSV4_E2E_API_PORT", "19280"), - "--afd-port", - os.environ.get("AFD_NPU_DSV4_E2E_AFD_PORT", "6455"), - "--startup-timeout", - os.environ.get("AFD_NPU_E2E_STARTUP_TIMEOUT", "1800"), - "--completion-output-path", - str(output_path), - ] + output_path, + backend="npu", + ) + command.extend( + [ + "--afd-host", + required_env("HCCL_IF_IP"), + "--api-port-base", + os.environ.get("AFD_NPU_DSV4_E2E_API_PORT", "19280"), + "--afd-port", + os.environ.get("AFD_NPU_DSV4_E2E_AFD_PORT", "6455"), + "--startup-timeout", + os.environ.get("AFD_NPU_E2E_STARTUP_TIMEOUT", "1800"), + ] + ) + return command def build_environment() -> dict[str, str]: @@ -88,7 +77,9 @@ def build_environment() -> dict[str, str]: @pytest.mark.slow @pytest.mark.parametrize("scenario", [DSV4_ASYNC_CAM_SCENARIO]) def test_deepseek_v4_flash_async_cam(scenario: str, tmp_path: Path) -> None: - run_runner( - build_runner_command(tmp_path / f"{scenario}.json"), + output_path = tmp_path / scenario + run_case_entry( + build_runner_command(output_path), + output_path, env=build_environment(), ) diff --git a/tests/e2e/models/deepseek_v4_flash/test_gpu.py b/tests/e2e/models/deepseek_v4_flash/test_gpu.py index aadaa60cd..c58b5bed6 100644 --- a/tests/e2e/models/deepseek_v4_flash/test_gpu.py +++ b/tests/e2e/models/deepseek_v4_flash/test_gpu.py @@ -5,19 +5,20 @@ from __future__ import annotations import os -import sys from pathlib import Path import pytest -from tests.conftest import download_dataset, run_runner -from tests.e2e.environment import devices_from_env, required_env +from tests.conftest import download_dataset, run_case_entry +from tests.e2e.environment import devices_from_env, model_case_command, required_env +from tests.e2e.models.deepseek_v4_flash.cases import ( + GPU_BASIC_IDS as DSV4_GPU_BASIC_SCENARIOS, +) +from tests.e2e.models.deepseek_v4_flash.cases import ( + GPU_DBO_IDS as DSV4_GPU_DBO_SCENARIOS, +) from tests.e2e.models.deepseek_v4_flash.gpu_config import ( - DSV4_GPU_BASIC_SCENARIOS, - DSV4_GPU_DBO_SCENARIOS, DSV4_GPU_DEVICE_COUNT, - DSV4_GPU_ROLE_RANKS, - DSV4_GPU_SCENARIO_SETTINGS, ) GSM8K_DATASET_ID = "openai/gsm8k" @@ -36,37 +37,19 @@ def prepare_e2e_assets() -> None: def build_runner_command(scenario: str, output_dir: Path) -> list[str]: - devices = devices_from_env("AFD_E2E_DEVICES", DSV4_GPU_DEVICE_COUNT) - baseline = DSV4_GPU_SCENARIO_SETTINGS[scenario][0] - command = [ - sys.executable, - "-m", - "tests.e2e.runner", - "--model", - required_env("AFD_GPU_E2E_MODEL"), - "--vllm-bin", - os.environ.get("AFD_GPU_E2E_VLLM_BIN", "vllm"), - "--device-backend", - "gpu", - "--attention-devices", - ",".join(devices[:DSV4_GPU_ROLE_RANKS]), - "--scenario", - scenario, - "--served-model-name-prefix", - "dsv4-flash-gpu", - "--completion-output-path", - str(output_dir / "completions.json"), - "--gsm8k-output-path", - str(output_dir / "gsm8k"), - "--startup-timeout", - os.environ.get("AFD_GPU_E2E_STARTUP_TIMEOUT", "1800"), - "--afd-port", - os.environ.get("AFD_GPU_DSV4_E2E_AFD_PORT", "6457"), - "--api-port-base", - os.environ.get("AFD_GPU_DSV4_E2E_API_PORT", "19480"), - ] - if not baseline: - command.extend(["--ffn-devices", ",".join(devices[DSV4_GPU_ROLE_RANKS:])]) + command = model_case_command( + "tests.e2e.models.deepseek_v4_flash.cases", scenario, output_dir, backend="gpu" + ) + command.extend( + [ + "--startup-timeout", + os.environ.get("AFD_GPU_E2E_STARTUP_TIMEOUT", "1800"), + "--afd-port", + os.environ.get("AFD_GPU_DSV4_E2E_AFD_PORT", "6457"), + "--api-port-base", + os.environ.get("AFD_GPU_DSV4_E2E_API_PORT", "19480"), + ] + ) return command @@ -74,7 +57,7 @@ def build_runner_command(scenario: str, output_dir: Path) -> list[str]: @pytest.mark.slow @pytest.mark.parametrize("scenario", DSV4_GPU_BASIC_SCENARIOS) def test_deepseek_v4_flash_gpu(scenario: str, tmp_path: Path) -> None: - run_runner(build_runner_command(scenario, tmp_path)) + run_case_entry(build_runner_command(scenario, tmp_path), tmp_path) @pytest.mark.e2e @@ -82,4 +65,4 @@ def test_deepseek_v4_flash_gpu(scenario: str, tmp_path: Path) -> None: @pytest.mark.parametrize("scenario", DSV4_GPU_DBO_SCENARIOS) def test_deepseek_v4_flash_gpu_dbo(scenario: str, tmp_path: Path) -> None: """Qualify DBO separately with the shared nonempty-stage execution oracle.""" - run_runner(build_runner_command(scenario, tmp_path)) + run_case_entry(build_runner_command(scenario, tmp_path), tmp_path) diff --git a/tests/e2e/models/deepseek_v4_flash/test_sync_camp2p_npu.py b/tests/e2e/models/deepseek_v4_flash/test_sync_camp2p_npu.py index 7b529df0e..0b643f4c1 100644 --- a/tests/e2e/models/deepseek_v4_flash/test_sync_camp2p_npu.py +++ b/tests/e2e/models/deepseek_v4_flash/test_sync_camp2p_npu.py @@ -14,13 +14,12 @@ from __future__ import annotations import os -import sys from pathlib import Path import pytest -from tests.conftest import run_runner -from tests.e2e.environment import devices_from_env, required_env +from tests.conftest import run_case_entry +from tests.e2e.environment import devices_from_env, model_case_command, required_env from tests.e2e.models.deepseek_v4_flash.config import ( DSV4_SYNC_CAMP2P_SCENARIOS, DSV4_SYNC_HCCL_BUFFER_SIZE_MB, @@ -46,36 +45,23 @@ def build_runner_command(scenario: str, output_path: Path) -> list[str]: shape = DSV4_SYNC_SHAPES[scenario] if required_env("AFD_E2E_BACKEND") != "npu": raise RuntimeError("DSV4 sync CAMP2P E2E requires AFD_E2E_BACKEND=npu") - devices = devices_from_env("AFD_E2E_DEVICES", shape.device_count) - return [ - sys.executable, - "-m", - "tests.e2e.runner", - "--model", - required_env("AFD_NPU_E2E_MODEL"), - "--vllm-bin", - os.environ.get("AFD_NPU_E2E_VLLM_BIN", "vllm"), - "--device-backend", - "npu", - "--attention-devices", - ",".join(devices[: shape.attention_ranks]), - "--ffn-devices", - ",".join(devices[shape.attention_ranks :]), - "--scenario", - scenario, - "--served-model-name-prefix", - "dsv4-flash-sync", - "--afd-host", - _afd_host(shape), - "--api-port-base", - os.environ.get("AFD_NPU_DSV4_SYNC_E2E_API_PORT", "19380"), - "--afd-port", - os.environ.get("AFD_NPU_DSV4_SYNC_E2E_AFD_PORT", "6456"), - "--startup-timeout", - os.environ.get("AFD_NPU_E2E_STARTUP_TIMEOUT", "1800"), - "--completion-output-path", - str(output_path), - ] + devices_from_env("AFD_E2E_DEVICES", shape.device_count) + command = model_case_command( + "tests.e2e.models.deepseek_v4_flash.cases", scenario, output_path, backend="npu" + ) + command.extend( + [ + "--afd-host", + _afd_host(shape), + "--api-port-base", + os.environ.get("AFD_NPU_DSV4_SYNC_E2E_API_PORT", "19380"), + "--afd-port", + os.environ.get("AFD_NPU_DSV4_SYNC_E2E_AFD_PORT", "6456"), + "--startup-timeout", + os.environ.get("AFD_NPU_E2E_STARTUP_TIMEOUT", "1800"), + ] + ) + return command def build_environment(scenario: str) -> dict[str, str]: @@ -139,7 +125,9 @@ def build_environment(scenario: str) -> dict[str, str]: @pytest.mark.slow @pytest.mark.parametrize("scenario", DSV4_SYNC_CAMP2P_SCENARIOS) def test_deepseek_v4_flash_sync_camp2p(scenario: str, tmp_path: Path) -> None: - run_runner( - build_runner_command(scenario, tmp_path / f"{scenario}.json"), + output_path = tmp_path / scenario + run_case_entry( + build_runner_command(scenario, output_path), + output_path, env=build_environment(scenario), ) diff --git a/tests/e2e/models/qwen3_6/cases.py b/tests/e2e/models/qwen3_6/cases.py new file mode 100644 index 000000000..1a73a9dc3 --- /dev/null +++ b/tests/e2e/models/qwen3_6/cases.py @@ -0,0 +1,74 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Model-local CUDA cases; importing IDs never resolves model assets.""" + +from __future__ import annotations + +from tests.e2e.case import ( + Case, + CaseEnvironment, + Gsm8kSettings, + Gsm8kTask, + Mrv2Check, + Phase, +) +from tests.e2e.presets import DECODE_GRAPH, EAGER + +SCENARIOS = ( + "baseline-graph", + "afd-eager-2a1f", + "afd-graph-2a1f", + # Native hybrid GDN does not support DBO. +) +GSM8K_SAMPLES = 7 +GSM8K_THRESHOLD = 0.27 +CAPTURE_SIZE = 8 +MAX_MODEL_LEN = 4096 + + +def build_case(case_id: str, environment: CaseEnvironment) -> Case: + if environment.device_backend != "gpu": + raise ValueError("Qwen3.6 MoE CUDA E2E requires the gpu backend") + profiles = { + "baseline-graph": (True, True, 4, 0), + "afd-eager-2a1f": (False, False, 2, 1), + "afd-graph-2a1f": (False, True, 2, 1), + } + baseline, graph, attention_ranks, ffn_ranks = profiles[case_id] + execution = DECODE_GRAPH if graph else EAGER + return Case( + id=case_id, + model="Qwen3.6 MoE", + served_model_name_prefix="qwen3-6-afd", + deployment=execution.deployment( + backend="gpu", + attention=(attention_ranks, 1), + ffn=None if baseline else (ffn_ranks, 1), + connector=None if baseline else "P2pNcclAFDConnector", + capture_size=CAPTURE_SIZE, + ), + common_vllm_args=( + f"--max-model-len={MAX_MODEL_LEN}", + "--language-model-only", + ), + phases=( + Phase( + "evaluation", + Gsm8kTask( + Gsm8kSettings( + sample_limit=GSM8K_SAMPLES, + threshold=GSM8K_THRESHOLD, + ) + ), + (Mrv2Check(("attention", "ffn"), graph=True),) + if not baseline and graph + else (), + ), + ), + ) + + +if __name__ == "__main__": + from tests.e2e.runner import case_main + + raise SystemExit(case_main(build_case, SCENARIOS)) diff --git a/tests/e2e/models/qwen3_6/test_qwen3_6.py b/tests/e2e/models/qwen3_6/test_qwen3_6.py index 6921e5f7f..85bc13ff4 100644 --- a/tests/e2e/models/qwen3_6/test_qwen3_6.py +++ b/tests/e2e/models/qwen3_6/test_qwen3_6.py @@ -5,7 +5,6 @@ from __future__ import annotations import os -import sys from collections.abc import Iterator from pathlib import Path @@ -15,23 +14,14 @@ download_dataset, download_model, preserve_environment_variable, - run_runner, + run_case_entry, ) +from tests.e2e.environment import model_case_command +from tests.e2e.models.qwen3_6.cases import SCENARIOS GSM8K_DATASET_ID = "openai/gsm8k" GSM8K_DATASET_CONFIG = "main" QWEN3_6_REPO_ID = "Qwen/Qwen3.6-35B-A3B" -QWEN3_6_MAX_MODEL_LEN = 4096 -DEFAULT_DEVICE_IDS = ("0", "1", "2", "3") -ATTENTION_DEVICE_COUNT = 2 -AFD_FFN_DEVICE_COUNT = 1 -BASELINE_DEVICE_COUNT = 4 -# Native MRV2 does not support DBO for Qwen3.6 hybrid GDN models. -SCENARIOS = ( - "baseline-graph", - "afd-eager-2a1f", - "afd-graph-2a1f", -) def _required_env(name: str) -> str: @@ -64,49 +54,9 @@ def build_runner_command(scenario: str, gsm8k_output_path: Path) -> list[str]: if backend != "gpu": raise RuntimeError("Qwen3.6 MoE E2E supports only the 'gpu' backend") - env_devices = os.environ.get("AFD_E2E_DEVICES") - devices = ( - [item.strip() for item in env_devices.split(",") if item.strip()] - if env_devices - else list(DEFAULT_DEVICE_IDS) - ) - if scenario == "baseline-graph": - attention_devices = devices[:BASELINE_DEVICE_COUNT] - ffn_devices = [] - else: - attention_devices = devices[:ATTENTION_DEVICE_COUNT] - ffn_devices = devices[ - ATTENTION_DEVICE_COUNT : ATTENTION_DEVICE_COUNT + AFD_FFN_DEVICE_COUNT - ] - - command = [ - sys.executable, - "-m", - "tests.e2e.runner", - "--model", - _required_env("AFD_GPU_E2E_MODEL"), - "--vllm-bin", - os.environ.get("AFD_GPU_E2E_VLLM_BIN", "vllm"), - "--device-backend", - backend, - f"--common-vllm-arg=--max-model-len={QWEN3_6_MAX_MODEL_LEN}", - "--common-vllm-arg=--language-model-only", - "--attention-devices", - ",".join(attention_devices), - ] - if scenario != "baseline-graph": - command.extend(["--ffn-devices", ",".join(ffn_devices)]) - command.extend( - [ - "--scenario", - scenario, - "--gsm8k-output-path", - str(gsm8k_output_path), - "--served-model-name-prefix", - "qwen3-6-afd", - ], + return model_case_command( + "tests.e2e.models.qwen3_6.cases", scenario, gsm8k_output_path, backend=backend ) - return command @pytest.fixture(scope="module", autouse=True) @@ -121,4 +71,4 @@ def _prepare_e2e_assets() -> Iterator[None]: @pytest.mark.parametrize("scenario", SCENARIOS, ids=SCENARIOS) def test_qwen3_6(scenario: str, tmp_path: Path) -> None: command = build_runner_command(scenario, tmp_path / scenario) - run_runner(command) + run_case_entry(command, tmp_path / scenario) diff --git a/tests/e2e/models/qwen3_moe/cases.py b/tests/e2e/models/qwen3_moe/cases.py new file mode 100644 index 000000000..46a3f1f4b --- /dev/null +++ b/tests/e2e/models/qwen3_moe/cases.py @@ -0,0 +1,166 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Model-local CUDA cases; importing IDs never resolves model assets.""" + +from __future__ import annotations + +import sys +from dataclasses import replace + +from tests.e2e.case import ( + Case, + CaseEnvironment, + DboConfig, + FfnService, + Gsm8kSettings, + Gsm8kTask, + Mrv2Check, + Phase, +) +from tests.e2e.presets import DECODE_GRAPH, EAGER + +SCENARIOS = ( + "baseline-graph", + "afd-eager-2a1f", + "afd-graph-2a1f", + "afd-graph-dbo-2a1f", +) +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", +) +# The same DeepEP M2N deployments with the FFN side served by the standalone FFN +# service (torchrun -m afd_plugin.ffn_faas) instead of `vllm serve`; Attention +# launches exactly as in the mirrored case. +DEEPEP_M2N_FFN_SERVICE_SCENARIOS = ( + "afd-deepep-m2n-faas-eager-2a2f", + "afd-deepep-m2n-faas-graph-2a2f", + "afd-deepep-m2n-faas-graph-dbo-2a2f", + "afd-deepep-m2n-faas-graph-tp2-2a2f", +) +FFN_SERVICE_BACKEND = "gpu-deepep-m2n" +MULTI_POD_IDS = ( + "afd-deepep-m2n-eager-2a2f", + "afd-deepep-m2n-graph-2a2f", + "afd-deepep-m2n-graph-dbo-2a2f", +) +MULTI_POD_GSM8K_SAMPLES = 100 + +GSM8K_SAMPLES = 7 +GSM8K_THRESHOLD = 0.27 +DBO_CONCURRENCY = 12 +DBO_MINIMUM_SAMPLES = 2 * DBO_CONCURRENCY +DBO_DECODE_THRESHOLD = 2 +DBO_PREFILL_THRESHOLD = 8 +CAPTURE_SIZE = 8 +MAX_MODEL_LEN = 4096 + + +def build_multi_pod_case(case_id: str, environment: CaseEnvironment) -> Case: + if case_id not in MULTI_POD_IDS: + raise ValueError(f"unsupported GPU multi-pod case {case_id!r}") + case = build_case(case_id, environment) + dbo = case.deployment.dbo is not None + return replace( + case, + # Preserve the served name used by the existing distributed entry. + served_model_name_prefix="deepseek-v2-lite-afd", + phases=( + Phase( + "evaluation", + Gsm8kTask( + Gsm8kSettings( + sample_limit=MULTI_POD_GSM8K_SAMPLES, + threshold=GSM8K_THRESHOLD, + num_concurrent=DBO_CONCURRENCY if dbo else None, + minimum_samples=DBO_MINIMUM_SAMPLES if dbo else 1, + ) + ), + ), + ), + ) + + +def build_case(case_id: str, environment: CaseEnvironment) -> Case: + if environment.device_backend != "gpu": + raise ValueError("Qwen3 MoE CUDA E2E requires the gpu backend") + profiles = { + "baseline-graph": (True, True, False, 4, 1, 0, 1), + "afd-eager-2a1f": (False, False, False, 2, 1, 1, 1), + "afd-graph-2a1f": (False, True, False, 2, 1, 1, 1), + "afd-graph-dbo-2a1f": (False, True, True, 2, 1, 1, 1), + "afd-deepep-m2n-eager-2a2f": (False, False, False, 2, 1, 2, 1), + "afd-deepep-m2n-graph-2a2f": (False, True, False, 2, 1, 2, 1), + "afd-deepep-m2n-graph-dbo-2a2f": (False, True, True, 2, 1, 2, 1), + "afd-deepep-m2n-eager-tp2-2a2f": (False, False, False, 1, 2, 2, 1), + "afd-deepep-m2n-graph-tp2-2a2f": (False, True, False, 1, 2, 2, 1), + "afd-deepep-m2n-faas-eager-2a2f": (False, False, False, 2, 1, 2, 1), + "afd-deepep-m2n-faas-graph-2a2f": (False, True, False, 2, 1, 2, 1), + "afd-deepep-m2n-faas-graph-dbo-2a2f": (False, True, True, 2, 1, 2, 1), + "afd-deepep-m2n-faas-graph-tp2-2a2f": (False, True, False, 1, 2, 2, 1), + } + baseline, graph, dbo, a_dp, a_tp, f_dp, f_tp = profiles[case_id] + execution = DECODE_GRAPH if graph else EAGER + ffn_service = case_id in DEEPEP_M2N_FFN_SERVICE_SCENARIOS + deepep = case_id in DEEPEP_M2N_SCENARIOS or ffn_service + deployment = execution.deployment( + backend="gpu", + attention=(a_dp, a_tp), + ffn=None if baseline else (f_dp, f_tp), + connector=None + if baseline + else "DeepEPM2NAFDConnector" + if deepep + else "P2pNcclAFDConnector", + compute_gate_on_attention=deepep, + dbo=DboConfig(DBO_DECODE_THRESHOLD, DBO_PREFILL_THRESHOLD) if dbo else None, + capture_size=CAPTURE_SIZE, + ) + if ffn_service: + deployment = replace(deployment, ffn_service=FfnService(FFN_SERVICE_BACKEND)) + # The FFN service has no vLLM Worker to host the evidence extension; + # Attention's evidence covers the paired graph/DBO path. + evidence_roles = ("attention",) if ffn_service else ("attention", "ffn") + return Case( + id=case_id, + model="Qwen3 MoE", + served_model_name_prefix="qwen3-moe-afd", + deployment=deployment, + common_vllm_args=(f"--max-model-len={MAX_MODEL_LEN}",), + phases=( + Phase( + "evaluation", + Gsm8kTask( + Gsm8kSettings( + sample_limit=DBO_MINIMUM_SAMPLES if dbo else GSM8K_SAMPLES, + threshold=GSM8K_THRESHOLD, + num_concurrent=DBO_CONCURRENCY if dbo else None, + minimum_samples=DBO_MINIMUM_SAMPLES if dbo else 1, + ) + ), + (Mrv2Check(evidence_roles, graph=graph, dbo=dbo),) + if not baseline and (graph or dbo) + else (), + ), + ), + ) + + +if __name__ == "__main__": + if sys.argv[1:2] == ["multi-pod"]: + from tests.e2e.multi_pod.runner import case_main as multi_pod_case_main + + raise SystemExit( + multi_pod_case_main(build_multi_pod_case, MULTI_POD_IDS, sys.argv[2:]) + ) + from tests.e2e.runner import case_main + + raise SystemExit( + case_main( + build_case, + (*SCENARIOS, *DEEPEP_M2N_SCENARIOS, *DEEPEP_M2N_FFN_SERVICE_SCENARIOS), + ) + ) diff --git a/tests/e2e/models/qwen3_moe/test_qwen3_moe.py b/tests/e2e/models/qwen3_moe/test_qwen3_moe.py index ec02ccab9..5ceb051fc 100644 --- a/tests/e2e/models/qwen3_moe/test_qwen3_moe.py +++ b/tests/e2e/models/qwen3_moe/test_qwen3_moe.py @@ -5,7 +5,6 @@ from __future__ import annotations import os -import sys from collections.abc import Iterator from pathlib import Path @@ -15,29 +14,18 @@ download_dataset, download_model, preserve_environment_variable, - run_runner, + run_case_entry, +) +from tests.e2e.environment import model_case_command +from tests.e2e.models.qwen3_moe.cases import ( + DEEPEP_M2N_FFN_SERVICE_SCENARIOS, + DEEPEP_M2N_SCENARIOS, + SCENARIOS, ) -from tests.e2e.runner import DEEPEP_M2N_SCENARIOS GSM8K_DATASET_ID = "openai/gsm8k" GSM8K_DATASET_CONFIG = "main" QWEN3_MOE_REPO_ID = "Qwen/Qwen3-30B-A3B" -QWEN3_MOE_MAX_MODEL_LEN = 4096 -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", - "afd-eager-2a1f", - "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: @@ -70,53 +58,9 @@ def build_runner_command(scenario: str, gsm8k_output_path: Path) -> list[str]: if backend != "gpu": raise RuntimeError("Qwen3 MoE E2E supports only the 'gpu' backend") - env_devices = os.environ.get("AFD_E2E_DEVICES") - devices = ( - [item.strip() for item in env_devices.split(",") if item.strip()] - if env_devices - else list(DEFAULT_DEVICE_IDS) + return model_case_command( + "tests.e2e.models.qwen3_moe.cases", scenario, gsm8k_output_path, backend=backend ) - if scenario == "baseline-graph": - 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 + ffn_device_count - ] - - command = [ - sys.executable, - "-m", - "tests.e2e.runner", - "--model", - _required_env("AFD_GPU_E2E_MODEL"), - "--vllm-bin", - os.environ.get("AFD_GPU_E2E_VLLM_BIN", "vllm"), - "--device-backend", - backend, - f"--common-vllm-arg=--max-model-len={QWEN3_MOE_MAX_MODEL_LEN}", - "--attention-devices", - ",".join(attention_devices), - ] - if scenario != "baseline-graph": - command.extend(["--ffn-devices", ",".join(ffn_devices)]) - command.extend( - [ - "--scenario", - scenario, - "--gsm8k-output-path", - str(gsm8k_output_path), - "--served-model-name-prefix", - "qwen3-moe-afd", - ], - ) - return command @pytest.fixture(scope="module", autouse=True) @@ -131,11 +75,22 @@ def _prepare_e2e_assets() -> Iterator[None]: @pytest.mark.parametrize("scenario", SCENARIOS, ids=SCENARIOS) def test_qwen3_moe(scenario: str, tmp_path: Path) -> None: command = build_runner_command(scenario, tmp_path / scenario) - run_runner(command) + run_case_entry(command, tmp_path / scenario) @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) + run_case_entry(command, tmp_path / scenario) + + +@pytest.mark.e2e +@pytest.mark.parametrize( + "scenario", + DEEPEP_M2N_FFN_SERVICE_SCENARIOS, + ids=DEEPEP_M2N_FFN_SERVICE_SCENARIOS, +) +def test_qwen3_moe_deepep_m2n_ffn_service(scenario: str, tmp_path: Path) -> None: + command = build_runner_command(scenario, tmp_path / scenario) + run_case_entry(command, tmp_path / scenario) 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..3461dac31 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,18 +10,19 @@ 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 import os -import sys from collections.abc import Iterator +from pathlib import Path import pytest @@ -29,26 +30,23 @@ download_dataset, download_model, preserve_environment_variable, - run_runner, + run_case_entry, ) -from tests.e2e.runner import GSM8K_LIMIT_ENV +from tests.e2e.environment import model_multi_pod_command +from tests.e2e.models.qwen3_moe.cases import MULTI_POD_IDS 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", + "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"), -] +MULTI_POD_CASES = tuple( + (scenario, layout) for layout in POD_LAYOUTS for scenario in MULTI_POD_IDS +) def _required_env(name: str) -> str: @@ -79,31 +77,9 @@ def build_runner_command(scenario: str, layout_name: str) -> list[str]: 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 + return model_multi_pod_command( + "tests.e2e.models.qwen3_moe.cases", scenario, POD_LAYOUTS[layout_name] + ) @pytest.mark.e2e @@ -113,6 +89,7 @@ def build_runner_command(scenario: str, layout_name: str) -> list[str]: 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) + run_case_entry( + build_runner_command(scenario, layout_name), + Path(_required_env("AFD_E2E_GSM8K_OUTPUT")), + ) diff --git a/tests/e2e/multi_pod/README.md b/tests/e2e/multi_pod/README.md index 07b3bcb2d..4de425428 100644 --- a/tests/e2e/multi_pod/README.md +++ b/tests/e2e/multi_pod/README.md @@ -1,373 +1,157 @@ -# `tests/e2e/multi_pod/runner.py` walkthrough +# Multi-pod E2E execution -This document explains how `runner.py` executes, using -`tests/e2e/models/deepseek_v2_lite/test_deepseek_v2_lite_multi_pod.py` as the -worked example. It is a companion to the code, not a replacement for it — -line numbers refer to `runner.py` as of this writing and may drift. +Each existing pod runs one slice of a model-declared GPU Case. The model owns +its topology and acceptance workload; `multi_pod` owns rank placement and peer +coordination; `deployment.py` owns this pod's service processes. Pods or +containers must already exist. See the +[Kubernetes walkthrough](../../../.agents/skills/run-e2e/resources/k8-multi-pod.md) +for provisioning, resources and completion handling. -## The core idea +## Entry and ownership -A "multi-pod" AFD E2E scenario splits Attention and FFN ranks across more -than one pod (or container). Something has to decide, per pod, which ranks -it should launch, how those ranks should be wired to their peers in other -pods, and when it is safe to evaluate. `runner.py` is that "something" — but -deliberately with **no master process**. The same program, with the same -argv, runs once inside *every* participating pod. Each copy: - -1. figures out its own identity from its environment (pod 0? pod 1? …), -2. derives the *entire* global plan locally, with a pure function, from - inputs every pod was given identically, -3. reads out only the slice of that plan that belongs to it, -4. coordinates with its peers through a shared rendezvous store (not through - each other directly), and -5. tears down only the children it personally started. - -This is why the module docstring says "nothing outside the pods holds test -state." A k8s Job (or a Docker container group) can restart, reschedule, or -be inspected independently, and nothing needs to be told what the plan was — -it just needs the same launch inputs again. - -## How the test reaches the runner - -`test_deepseek_v2_lite_multi_pod.py` does **not** call any function in -`runner.py` directly. Instead: - -- `MULTI_POD_CASES` pairs a scenario id (e.g. `afd-graph-2a2f`) with a layout - name (e.g. `2pod-role-split` → pod layout string `"2A0F,0A2F"`). -- `build_runner_command()` builds an argv for `python -m - tests.e2e.multi_pod.runner` — the scenario, the pod layout, a run id, the - model, output paths, and the shared rendezvous store host — pulling most of - it from required environment variables (`AFD_E2E_RUN_ID`, - `AFD_GPU_E2E_MODEL`, `AFD_E2E_GSM8K_OUTPUT`, `AFD_E2E_STORE_HOST`, …). -- `test_multi_pod()` hands that argv to `run_runner()` (`tests/conftest.py`), - which `subprocess.Popen`s it in a new process group, forwards - SIGTERM/SIGINT into that group, and raises if it exits non-zero. - -Critically, the test file's own docstring spells out the deployment -assumption: **the pods must already exist**. Something else — a human, or -`tests.e2e.multi_pod.driver.k8s` / `driver.docker` — has already created N -pods/containers and is about to (or already did) invoke this same pytest -node once inside each of them. `pytest tests/e2e/models/deepseek_v2_lite/ -test_deepseek_v2_lite_multi_pod.py::test_multi_pod[...]` running inside pod 0 -and the identical invocation running inside pod 1 together *are* one -multi-pod E2E run. `test_multi_pod()` only supplies the inputs the in-pod -runner cannot infer from its own pod's environment (scenario, layout, store -host); everything pod-specific (index, peer addresses) is resolved by the -runner itself. - -So conceptually: +The ten existing GPU entries are unchanged: four DeepSeek-V2-Lite +scenario/layout pairs and six Qwen3 MoE DeepEP M2N pairs. Their pytest +entrypoints call `model_multi_pod_command()` and `run_case_entry()`: ```text -driver (k8s Job / docker driver) - └─ pod 0: pytest ... test_multi_pod[...] → runner.py --pod-layout 2A0F,0A2F ... - └─ pod 1: pytest ... test_multi_pod[...] → runner.py --pod-layout 2A0F,0A2F ... +models//test_*_multi_pod.py + -> python -m tests.e2e.models..cases multi-pod ... + -> model-owned build_multi_pod_case() + -> multi_pod/runner.case_main() / run_case() + -> layout.py: logical topology -> this pod's RoleSlots + -> rendezvous.py: TCPStore barriers, abort and verdict + -> deployment.py: role commands, local start/readiness/cleanup + -> declared Task: workload and result validation on the evaluator ``` -The two `runner.py` invocations rendezvous with each other over a -`torch.distributed.TCPStore`, not through pytest or the driver. - -## `runner.py` code walkthrough - -### Imports and constants (lines 1–76) - -The runner deliberately imports its "pure" helpers (`identity.py`, -`layout.py`) separately from the rendezvous module, which is the only one -that pulls in `torch.distributed`. The docstring of `identity.py` calls this -out explicitly: identity resolution has to stay unit-testable without a GPU -or a cluster. - -The barrier name constants (`ADDRESS_BARRIER`, `LAUNCHED_BARRIER`, -`SERVING_BARRIER`, `VERDICT_BARRIER`, `TEARDOWN_BARRIER`) name the five -synchronization points every pod passes through, in order — they show up -again in `run_pod()`. -`ROLE_LABELS` is just for readable log prefixes (`ATTN`, `FFN`, `BASELINE`). - -### `PodProcess` (lines 79–85) - -A tiny frozen dataclass pairing a locally-launched `subprocess.Popen` with -its role and its log-line label. `processes: list[PodProcess]` is threaded -through `main()` → `run_pod()` → `launch_slot()` and is what cleanup and the -liveness checks iterate over. - -### `main()` (lines 88–156) — phases 1–5 and the top-level shape - -This is the entry point (`if __name__ == "__main__": sys.exit(main())`, -line 442). Reading it top to bottom: - -1. **Parse args, configure the scenario, parse the layout.** - `configure_scenario(args)` (imported from the single-host - `tests/e2e/runner.py`) is the same function the single-host runner uses — - it turns a scenario id like `afd-graph-2a2f` into concrete settings - (`args.num_attention_ranks`, `args.cuda_graph_full_decode_only`, - `args.afd_connector`, …). Reusing it is what keeps a multi-pod scenario's - *logical* topology (rank counts, connector, DBO, CUDA graph) identical to - its single-host counterpart — only pod placement differs. - `PodLayout.parse(args.pod_layout)` turns `"2A0F,0A2F"` into a tuple of - `PodSpec(attention=2, ffn=0)` / `PodSpec(attention=0, ffn=2)`. - `Topology.from_args(args)` captures the scenario's *logical* shape - (`RoleTopology` per role, connector, `baseline`), independent of how it is - laid out across pods. `validate_layout()` cross-checks the two: the - layout's total per-role rank counts must match the topology's, and no - role's local rank count within a pod may violate that role's TP size - (a TP group cannot span pods). - The runner then sets `args.collect_mrv2_execution_evidence = False`. For - graph and DBO scenarios the single-host runner proves the requested - execution path, not just accuracy: it loads the `tests.e2e.mrv2_evidence` - worker extension, opens the evaluation window through Attention's - `/collective_rpc` and a local marker file the FFN worker polls, and checks - every rank's summary line for graph replays or two-stage DBO steps. No pod - can drive that window for the others, so multi-pod graph and DBO scenarios - (e.g. `afd-graph-2a2f`, `afd-graph-dbo-2a2f`) launch without the extension - and pass on accuracy alone. **TODO:** share the window and the per-rank - summaries through the rendezvous store and check the merged result before - the verdict. - -2. **Resolve this pod's index.** `resolve_pod_index(layout.num_pods)` (from - `identity.py`) checks, in order: an explicit `AFD_E2E_POD_INDEX` - override, the k8s Indexed Job's `JOB_COMPLETION_INDEX`, then falls back to - the trailing integer of `HOSTNAME` (for hand-applied manifests or the - Docker driver, which names containers `..-0`, `..-1`, …). - -3. **Wait for the rendezvous store's DNS to resolve, then connect.** - `wait_for_address(args.store_host, args.store_dns_timeout)` exists - because a freshly-created pod's DNS record can lag its own start by a few - seconds; without pre-resolving, that shows up as an unexplained - `TCPStore` connection failure rather than a named "still waiting on DNS" - state. `Rendezvous(...)` then opens a `TCPStore` — pod 0 is always the - store's master (`is_master=pod_index == 0`), which is independent of - which pod *leads* a role or evaluates. - -4. **`rendezvous.verify_agreement(layout.canonical(), args.scenario)`.** - Pod 0 publishes the canonical layout string and scenario id; every other - pod reads them back and raises immediately if its own values differ. This - turns a misconfigured deployment (e.g. one pod launched with a different - `--pod-layout`) into a fast, named failure instead of a barrier timeout - thirty minutes later. - -5. **Resolve every pod's address, plan, and print.** - `resolve_addresses()` (see below) gives every pod the same ordered list - of peer addresses. `plan(topology, layout, addresses)[pod_index]` calls - the pure planning function from `layout.py` and keeps only this pod's - `PodPlan`. `print_plan()` logs it for the pod's own console/log stream. - -6. **Stale-process pre-flight.** `find_stale_run_markers(args.run_id)` scans - `/proc` for local processes whose `AFD_E2E_RUN_ID` environment variable - does not match this run's id — a leftover process from a previous run - still holding a GPU/NPU device or a port. If found, the pod publishes an - abort (so peers unwind quickly) and raises, rather than letting the - symptom surface later as an unexplained rendezvous hang. - -7. **Run the pod body under `cancellable_run(cleanup)`.** - `cancellable_run` (from `tests/e2e/cancellation.py`) is shared with the - single-host runner and gives both the same SIGTERM/SIGINT precedence: a - signal during the body unwinds immediately into cleanup; a signal that - arrives *during* cleanup is deferred until cleanup finishes; the - resulting `SystemExit` then wins over a cleanup error, which in turn wins - over a body error. `cleanup()` (defined as a closure over `processes`, - `log_threads`, and `rendezvous`) calls `terminate_processes()` on every - local child, then joins the log-streaming threads, and best-effort - publishes its own outcome to the store under `phase/{pod_index}/cleanup` - (`publish_quietly` — a store failure during teardown must never mask the - run's actual result). If the body (`run_pod`) raises, `main()` also - publishes an abort naming which pod failed, so peers still blocked on a - barrier or on the verdict key unwind within seconds instead of waiting - out their full timeout. - -8. On success, prints a `PASSED` line tagged with this pod's index, scenario, - and layout, and returns `0`. - -9. **Release the store, whatever the outcome.** Everything after connecting - runs inside a `try`/`finally` that calls `leave_rendezvous()`. Pod 0's - process *is* the store server, so it must outlive every peer's last store - call: a barrier only proves that every pod *arrived*, and a peer may still - be asleep in the barrier's poll interval when pod 0 returns from it. So - after its own cleanup each peer calls `Rendezvous.depart()` (its last - store write), and pod 0 calls `wait_for_departures()` before exiting, - bounded by `departure_timeout(args)` — the scenario's teardown budget - plus `DEPARTURE_GRACE_S`. A straggler (a peer that died or never - connected) is logged, never turned into a different result; another - arrival barrier could not replace this, since the last pod to arrive - always leaves while the others still sleep. - -### `run_pod()` (lines 159–259) — phases 6–11, the ordered protocol +For example, invoke the following in each of the two pre-existing pods, with +its own output directory and device visibility: -This is the heart of the coordination logic, and its own comments number the -phases (6 through 11) that continue from `main()`'s 1–5. Every phase either -launches something local or waits on a named rendezvous barrier; **all** -barrier waits pass `on_poll=assert_local_processes_alive`, so a barrier wait -never blocks silently past a local child dying — it notices and raises -within one poll interval. - -- **Phases 6–7 — ordered launch.** Which role starts first depends on the - connector: `uses_async_connector(args)` (true for `CAMAsyncAFDConnector`) - means Attention hosts the AFD rendezvous and must come up first; - otherwise FFN does (mirrors `RENDEZVOUS_ROLE_BY_CONNECTOR` in - `layout.py`). The runner launches this pod's slot for the *first* role - kind (`launch_slot`, if this pod holds that role), then waits on a - `f"{first_kind}-launched"` barrier before any pod starts its *second* - role. This matters because the AFD connector's own rendezvous protocol - needs the answering side present before the initiating side dials it — - getting the launch order wrong here would produce connector-level - timeouts unrelated to this runner's own barriers. - -- **Phase 8 — "launched" barrier.** Once both of this pod's local role - processes (if any) are started, it publishes `phase/{pod_index}/launched` - and waits for every pod. The code comment is explicit that "launched" - only means every pod's *children are alive*, not that they are serving — - a weaker, faster-to-reach checkpoint than readiness. - -- **Phase 9 — readiness via `/health`.** Only the pod holding the - *non-headless* Attention slot is polled for `/health` — and the comment - explains why that's sufficient: the Attention leader's API server binds - only after the AFD connector rendezvous completes, which itself requires - every FFN rank to be present. So one health check transitively proves the - whole distributed AFD world formed. An FFN engine core never builds an API - server at all (it enters a busy loop — see - `afd_plugin/compat/patches/engine_core.py:_run_ffn_busy_loop`), and a - headless DP slot starts no server by construction, so for both of those - the phase-8 liveness check was already the honest assertion. After the - optional health poll, the pod publishes `phase/{pod_index}/serving` and - waits on the `SERVING_BARRIER`. - -- **Phase 10 — exactly one evaluator.** `pod_plan.is_evaluator` is true only - for the pod holding the Attention role's leader rank (set in - `layout.plan()`). That one pod runs the accuracy/completion check against - its **own local** API through `run_scenario_evaluation`, the same - scenario-to-evaluator mapping the single-host `tests/e2e/runner.py` uses - (teardown likewise shares `process_termination_timeout`), publishes the - verdict string to - the store (`VERDICT_KEY`), and re-raises on failure after recording - `"fail: {exc}"` and publishing an abort. Every *other* pod instead calls - `rendezvous.wait_for_key(VERDICT_KEY, ...)` and blocks until that key - appears (or a peer aborts). - -- **Phase 11 — every pod holds the same verdict.** All pods (evaluator - included) pass through one more barrier (`VERDICT_BARRIER`) so that a slow - non-evaluator pod can't be torn down by its driver before it has actually - read the verdict. A final `assert_local_processes_alive()` catches a child - that died in the narrow window after the last health/liveness check, and - the pod raises if its verdict isn't `PASS_VERDICT` — this is what makes - the process's own exit code (and thus the k8s Job / Docker container exit - code, and thus `test_multi_pod()`'s subprocess result) reflect pass/fail - for the whole distributed run, not just this pod's local view. - -- **Phase 12 — teardown barrier.** A passing pod waits on - `TEARDOWN_BARRIER` before returning into cleanup, so no pod stops its - roles until every pod has passed its final liveness check. Stopping a - local role can take down its AFD peers on other pods, which a pod still - checking would misreport as a failed run. This is the only barrier - *without* `on_poll`: from here on a child exiting is expected. - -### `launch_slot()` (lines 262–290) - -Builds and starts exactly one role's process for this pod. It picks -`build_baseline_command` or `build_vllm_command(args, role=slot.role, -slot=slot)` — both imported unchanged from the single-host runner — and -passes the `RoleSlot` through so `build_vllm_command` can add the -multi-pod-only flags (`--data-parallel-size-local`, -`--data-parallel-address`, `--data-parallel-rpc-port`, plus -`--data-parallel-start-rank` and `--headless` on non-leader slots) *only* -when `slot.spans_pods` is true. The leader omits its start rank (always 0): -vLLM 0.30 reads an explicit start rank on a non-headless engine as hybrid -load balancing, which turns external when the local DP size is 1 and then -rejects the headless peers. `slot.spans_pods` is what keeps a one-pod layout -producing the exact same argv, character for character, as the single-host -runner — making a 1-pod layout a genuine control case rather than a -different code path pretending to be one. After starting the process, it is -appended to `processes`, its stdout is piped through `stream_output()` onto -a daemon thread prefixed with a pod/role label, and it is immediately polled -once to fail fast if it exited before even reaching the timeout-based checks -later. - -### `wait_for_health()` (lines 293–325) - -A small polling loop against `/health` (not `/v1/models`) on the Attention -API port, calling the caller-supplied `on_poll` liveness check on every -iteration (including the first, before the first request) so a dead child is -never masked by an HTTP-level retry. Raises `TimeoutError` naming the last -observed error if the deadline passes. - -### `deferred_sigkill_pgids()` (lines 328–335) - -Delegates to `uses_npu_async_process_cleanup(args)` (shared with the -single-host runner) to decide whether this pod's FFN process needs the NPU -async teardown allowance — FFN workers using the async CAM connector on NPU -can sit in uninterruptible HCCL teardown for tens of seconds after SIGTERM, -so their process groups get a deferred, longer-patience SIGKILL rather than -being force-killed on the same schedule as everything else. This mirrors -the equivalent logic in the single-host `tests/e2e/runner.py:main()` -verbatim, just scoped to *this pod's* FFN process instead of the whole run's. - -### `resolve_addresses()` (lines 338–363) - -Three ways to learn every pod's address, tried in order: - -1. **Explicit list** — `--pod-addresses` or the `AFD_E2E_POD_ADDRESSES` - environment variable, comma-separated. Used when the driver already - knows every pod's address up front. -2. **A template** — `--pod-address-template` with an `{index}` field (e.g. a - StatefulSet-style DNS name `afd-e2e-abc-{index}.afd-e2e-abc`). This - *skips* the address-exchange barrier entirely, since the addresses are - derivable without any coordination. -3. **Rendezvous exchange** — the fallback: this pod publishes its own - `local_address()` under `addr/{pod_index}`, waits on `ADDRESS_BARRIER` - for every pod to do the same, then reads all of them back in order. - `local_address()` prefers `HCCL_IF_IP` (an Ascend pod's HCCL interface, - which HCCL and CAM bind to), then `POD_IP`, then the hostname. - -### `publish_quietly()` and `print_plan()` (lines 366–388) - -`publish_quietly` wraps a single `rendezvous.set()` in a `try/except -Exception`, printing rather than raising — used only from `cleanup()`, -where the docstring/comment is direct: teardown's own outcome must never -fail the process, because the exit code the driver observes is what -actually reports the run's result. `print_plan` is pure logging: this pod's -address, its peer list, whether it's the evaluator, and per-slot device/DP -placement — useful for debugging a hung run from pod logs alone. - -### `parse_args()` (lines 391–439) - -Adds the multi-pod-specific CLI surface on top of -`add_scenario_arguments(parser)` (shared with the single-host runner, which -supplies `--model`, `--scenario`, connector/graph/DBO flags, etc. — but -deliberately *not* device selection, since that's derived differently by -each runner). The multi-pod-only arguments are `--pod-layout`, `--run-id`, -`--store-host` / `--store-port` / `--store-dns-timeout`, the three address -sources for `resolve_addresses()`, `--pod-env` (extra `KEY=VALUE`s merged -into every launched process's environment), and the three timeout knobs -(`--launch-timeout`, `--serving-timeout`, `--verdict-timeout`) that bound -each of the barriers above. - -## Where the shared logic actually lives - -A recurring theme above: most of what looks like "the test's logic" — -scenario configuration, vLLM command construction, GSM8K/completion -evaluation, environment building, process start/stream/terminate — is -**not** duplicated in `runner.py`. It is imported unchanged from the -single-host `tests/e2e/runner.py`. `runner.py` (this module) only adds what -is genuinely different about running as one pod among several: - -- identity resolution (`identity.py`), -- deriving a global plan and reading out one pod's slice of it - (`layout.py`), -- and cross-pod coordination through a shared store (`rendezvous.py`). - -This is also why `layout.RoleSlot` is imported directly into -`tests/e2e/runner.py` (`build_vllm_command(..., slot: RoleSlot | None = -None)`) rather than the multi-pod runner reimplementing command -construction: a `slot=None` call from the single-host runner and a -`slot=` call from here are meant to produce identical -output whenever that slot doesn't actually span pods. - -## Provisioning: who creates the pods in the first place +```bash +python -m tests.e2e.models.deepseek_v2_lite.cases multi-pod \ + --scenario afd-graph-2a2f --pod-layout 2A0F,0A2F \ + --run-id example-run --model /models/DeepSeek-V2-Lite \ + --device-backend gpu --output-path /work/results \ + --store-host afd-e2e-0.afd-e2e +``` -`runner.py` assumes N pods/containers already exist and that it is being -invoked once inside each. That provisioning step is a separate concern, -handled by `tests/e2e/multi_pod/driver/k8s.py` and `driver/docker.py`. Both -drivers describe themselves the same way: their whole job is to render and -apply the deployment (a k8s Indexed Job or a set of sibling Docker -containers), block on the platform's own completion primitive (`kubectl -wait` / a container-exit wait), collect exit codes and logs, and clean up — -they hold no test state and make no test decision. Launch order, readiness, -evaluation, and teardown are entirely owned by the pods running this module, -exactly as walked through above. +The original NPU DSV4 multi-pod entries still use the retained +`python -m tests.e2e.multi_pod.runner` CLI and NPU launch helpers. The shared +coordination sequence is the same. NPU migration is separate work. + +## Identity, placement and resources + +`resolve_pod_index()` uses `AFD_E2E_POD_INDEX`, then `JOB_COMPLETION_INDEX`, +then a trailing index in `HOSTNAME`. Pod 0 owns the TCPStore; the Attention +leader owns evaluation. These need not be the same pod. + +A layout such as `2A0F,0A2F` is independent of the Case's logical topology. +Per-role totals must match the Case, and each pod must hold a whole number of +TP groups: a TP group cannot span pods. Qwen3 DeepEP M2N retains the role-split +and interleaved layouts from main. Its transport requires equal rank counts +per NVLink domain, with FFN ranks filling the first domains or the same FFN +count in every domain. The connector's transport mapping owns these rules; +this migration adds no transport capability. +Baseline placement across multiple pods remains unsupported. + +`plan()` derives every pod's plan from the same topology, layout and addresses. +Each pod uses only its own `RoleSlot`s. A role spanning pods receives local DP +size, rendezvous address and RPC port. Only headless peers receive the DP start +rank; emitting it on the leader selects a different vLLM load-balancing mode +that rejects headless peers. Headless commands do not reserve HTTP ports. + +Slot devices are positions within the pod's allocated device list. `--devices` +or `AFD_E2E_DEVICES` can bind that list explicitly; otherwise the GPU CLI uses +inherited `CUDA_VISIBLE_DEVICES`, or the pod's local ordinal list when unset. +The service launcher partitions that list by role. + +Address resolution uses, in order: + +1. `--pod-addresses` or `AFD_E2E_POD_ADDRESSES`; +2. `--pod-address-template` with an `{index}` field, without an exchange barrier; +3. each pod's published address and an address-exchange barrier. + +The advertised local address prefers `HCCL_IF_IP`, then `POD_IP`, then hostname. +DNS resolution of the store host is bounded before TCPStore connection. +Pod 0 publishes the scenario ID and canonical layout; peers reject disagreement. +GPU execution binds this pod's allocated resources and cleans only its launched +groups. A different run marker alone does not prove stale ownership or block a +run on a shared host. The retained NPU path keeps its existing stale-run check +for its isolated pod environment; it does not clean up another run's resources. + +## Coordination sequence + +1. Read the model Case, validate placement, resolve identity and addresses. +2. Start local slots for the connector's first role; all pods pass its launched + barrier before starting the second role. GPU connector order comes from + `deployment.py`, not scenario-name branches in the coordinator. +3. Pass the launched barrier. This proves local children are alive, not ready. +4. The non-headless Attention leader waits for `/health`; other slots report + liveness. FFN busy loops and headless workers do not expose an HTTP API. + All pods then pass the serving barrier. +5. Exactly one evaluator executes the declared phases against its local API. + It publishes the workload results, effective configuration and verdict. + Peers consume those results rather than run the workload again. +6. Pass the verdict barrier and check local liveness once more. +7. Pass the teardown barrier before stopping any local role. It intentionally + omits liveness polling: a peer may now legitimately start shutting down. +8. Clean up this pod's owned process groups, reap leaders and join log readers. +9. Non-master pods publish departure after their last store access. Pod 0 waits + for departures with a bounded grace period, then releases the store. + +Barrier and verdict waits poll peer aborts and local liveness. A local failure +publishes an abort so peers waiting on coordination can unwind. Cleanup uses +`cancellable_run()`: cancellation is deferred during cleanup and takes priority +on exit, followed by cleanup errors and the original execution error. Cleanup +status publication is best effort and cannot mask those errors. + +Departure is different from arrival: an arrival barrier cannot prove that a +sleeping peer has finished its last store read. Pod 0 therefore outlives peers' +explicit departure messages; stragglers are reported after the bounded wait. +Do not automatically restart a pod into an already-advanced barrier sequence. + +## Acceptance and records + +Migration preserves existing GPU multi-pod acceptance: + +| Suite | Workload | Threshold | Execution probes | +| --- | --- | --- | --- | +| V2 graph | 7 GSM8K samples | 0.27 | None | +| V2 graph + DBO | 24 GSM8K samples, concurrency 12 | 0.27 | None | +| Qwen3 DeepEP M2N | 100 GSM8K samples; DBO concurrency 12 | 0.27 | None | + +These Cases explicitly declare no execution checks. Accuracy, local service +liveness, peer coordination and cleanup are checked; graph/DBO/Hash execution +across pods is not proved. The executor rejects a Case requesting unsupported +cross-pod execution checks instead of silently dropping them. + +Each pod writes role logs and `run.json` under its output directory. The +evaluator writes task artifacts under the phase directory (currently +`evaluation/`). `AFD_E2E_GSM8K_OUTPUT` supplies that root to the pytest entry. +Records identify the local plan/devices, actual evaluator settings and results, +coverage scope, local execution and cleanup. A weakened evaluator configuration +produces a debug verdict on every pod. The CLI may complete that workload, but +`run_case_entry()` rejects it as acceptance. + +The deployment passes only when every pod's pytest entry succeeds with +`acceptance_passed=true`. One pod's record does not prove peer cleanup. The +retained NPU CLI keeps its existing completion and scoped async-FFN cleanup +policy, including its longer process-group wait and marked-process cleanup. + +## Extending and validating + +For the migrated GPU Case path, add existing-capability combinations through +model declarations, IDs and layout selection. Do not add model/scenario branches +to shared coordination, commands, tasks or probes. Preserve the model's distributed acceptance settings explicitly; +single-deployment defaults are not automatically interchangeable. +The retained NPU executor is outside this extension contract until migrated. +GPU Case runs use deployment-owned cleanup; legacy NPU runs keep their existing +cleanup path. The coordinator does not run both for the same invocation. + +`tests/e2e/test_runner.py` exercises two real runner processes and a real +TCPStore, with stub service/evaluator CLIs. These host integrations require +PyTorch for TCPStore, but do not load models or establish device correctness. +They cover role-split/interleaved placement, the Qwen workload, debug verdict +propagation, startup/evaluator failure, cancellation and owned cleanup. +Real-device validation is reported separately, including whether placement was +same-host or cross-node. This migration adds no cross-pod execution evidence. diff --git a/tests/e2e/multi_pod/layout.py b/tests/e2e/multi_pod/layout.py index 0282ec923..0c31159b9 100644 --- a/tests/e2e/multi_pod/layout.py +++ b/tests/e2e/multi_pod/layout.py @@ -14,6 +14,8 @@ from collections.abc import Sequence from dataclasses import dataclass +from tests.e2e.case import Case + ATTENTION_ROLE = "attention" FFN_ROLE = "ffn" BASELINE_ROLE = "baseline" @@ -148,6 +150,19 @@ def rendezvous_role(self) -> str | None: ) return rendezvous_role + @classmethod + def from_case(cls, case: Case) -> Topology: + profile = case.deployment + return cls( + attention=RoleTopology(profile.attention.rank_count, profile.attention.tp), + ffn=RoleTopology( + 0 if profile.ffn is None else profile.ffn.rank_count, + 1 if profile.ffn is None else profile.ffn.tp, + ), + connector=profile.connector or "", + baseline=profile.is_baseline, + ) + @classmethod def from_args(cls, args: argparse.Namespace) -> Topology: """Build a topology from a ``configure_scenario``-populated namespace.""" diff --git a/tests/e2e/multi_pod/rendezvous.py b/tests/e2e/multi_pod/rendezvous.py index 64e34e1e6..2497a53c0 100644 --- a/tests/e2e/multi_pod/rendezvous.py +++ b/tests/e2e/multi_pod/rendezvous.py @@ -13,8 +13,6 @@ from collections.abc import Callable from datetime import timedelta -from torch.distributed import TCPStore - DEFAULT_STORE_PORT = 29500 STORE_CONNECT_TIMEOUT_S = 600 BARRIER_POLL_INTERVAL_S = 2.0 @@ -27,6 +25,7 @@ VERDICT_KEY = "verdict" ABORT_COUNT_KEY = "count/abort" PASS_VERDICT = "pass" +DEBUG_VERDICT = "debug" class PeerFailureError(RuntimeError): @@ -54,6 +53,9 @@ def __init__( connect_timeout_s: float = STORE_CONNECT_TIMEOUT_S, poll_interval_s: float = BARRIER_POLL_INTERVAL_S, ) -> None: + # Runtime-only dependency: case collection and CLI inspection need no torch. + from torch.distributed import TCPStore + if not 0 <= pod_index < num_pods: raise ValueError( f"pod index {pod_index} is outside 0..{num_pods - 1}", diff --git a/tests/e2e/multi_pod/runner.py b/tests/e2e/multi_pod/runner.py index aa3c7420f..6a8792a33 100644 --- a/tests/e2e/multi_pod/runner.py +++ b/tests/e2e/multi_pod/runner.py @@ -12,6 +12,7 @@ from __future__ import annotations import argparse +import json import os import subprocess import sys @@ -20,9 +21,12 @@ import urllib.error import urllib.request from collections.abc import Callable -from dataclasses import dataclass +from dataclasses import asdict, dataclass, replace +from tests.e2e import deployment from tests.e2e.cancellation import cancellable_run +from tests.e2e.case import Case, CaseEnvironment, JsonValue +from tests.e2e.environment import add_case_environment_arguments, case_environment from tests.e2e.multi_pod.identity import ( POD_ADDRESSES_ENV, find_stale_run_markers, @@ -42,6 +46,7 @@ validate_layout, ) from tests.e2e.multi_pod.rendezvous import ( + DEBUG_VERDICT, DEFAULT_STORE_PORT, PASS_VERDICT, VERDICT_KEY, @@ -57,6 +62,7 @@ build_env, build_vllm_command, configure_scenario, + prepare_case, print_command, process_termination_timeout, run_scenario_evaluation, @@ -166,7 +172,34 @@ def cleanup() -> None: with cancellable_run(cleanup): try: - run_pod(args, rendezvous, pod_plan, processes, log_threads) + + def check_alive() -> None: + for entry in processes: + if entry.process.poll() is not None: + raise RuntimeError( + f"{entry.label} exited with {entry.process.returncode}" + ) + + def evaluate() -> str: + run_scenario_evaluation(args) + return PASS_VERDICT + + run_pod( + args, + rendezvous, + pod_plan, + launch_order=(ATTENTION_ROLE, FFN_ROLE) + if uses_async_connector(args) + else (FFN_ROLE, ATTENTION_ROLE), + start_slot=lambda slot: launch_slot( + args, pod_plan, slot, processes, log_threads + ), + check_alive=check_alive, + wait_ready=lambda slot, check: wait_for_health( + args, slot, args.serving_timeout, on_poll=check + ), + evaluate=evaluate, + ) except BaseException as exc: rendezvous.publish_abort(f"pod {pod_index}: {exc}") raise @@ -185,34 +218,30 @@ def run_pod( args: argparse.Namespace, rendezvous: Rendezvous, pod_plan: PodPlan, - processes: list[PodProcess], - log_threads: list[threading.Thread], -) -> None: + *, + launch_order: tuple[str, ...], + start_slot: Callable[[RoleSlot], None], + check_alive: Callable[[], None], + wait_ready: Callable[[RoleSlot, Callable[[], None]], None], + evaluate: Callable[[], str], +) -> str: """Phases 6-12: ordered launch, readiness, evaluation, verdict, teardown.""" pod_index = pod_plan.index def assert_local_processes_alive() -> None: - for entry in processes: - returncode = entry.process.poll() - if returncode is not None: - message = ( - f"pod {pod_index} {entry.label} exited " - f"(rc={returncode}) during the run" - ) - rendezvous.publish_abort(message) - raise RuntimeError(message) + try: + check_alive() + except BaseException as exc: + rendezvous.publish_abort(f"pod {pod_index}: {exc}") + raise # Phases 6-7: the connector's rendezvous role launches first, and every pod # waits for it before the other role starts. - first_kind, second_kind = ( - (ATTENTION_ROLE, FFN_ROLE) - if uses_async_connector(args) - else (FFN_ROLE, ATTENTION_ROLE) - ) + first_kind, second_kind = launch_order for kind, barrier in ((first_kind, f"{first_kind}-launched"), (second_kind, None)): slot = pod_plan.slot(kind) if slot is not None: - launch_slot(args, pod_plan, slot, processes, log_threads) + start_slot(slot) if barrier is not None: rendezvous.barrier( barrier, @@ -238,12 +267,7 @@ def assert_local_processes_alive() -> None: # for both the honest assertion is liveness, published at phase 8. serving_slot = pod_plan.slot(ATTENTION_ROLE) if serving_slot is not None and not serving_slot.headless: - wait_for_health( - args, - serving_slot, - args.serving_timeout, - on_poll=assert_local_processes_alive, - ) + wait_ready(serving_slot, assert_local_processes_alive) rendezvous.set(f"phase/{pod_index}/serving", "ok") rendezvous.barrier( SERVING_BARRIER, @@ -254,14 +278,13 @@ def assert_local_processes_alive() -> None: # Phase 10: exactly one pod evaluates, against its own local API. if pod_plan.is_evaluator: try: - run_scenario_evaluation(args) + verdict = evaluate() except BaseException as exc: verdict = f"fail: {exc}" rendezvous.set(VERDICT_KEY, verdict) rendezvous.publish_abort(verdict) raise - rendezvous.set(VERDICT_KEY, PASS_VERDICT) - verdict = PASS_VERDICT + rendezvous.set(VERDICT_KEY, verdict) else: verdict = rendezvous.wait_for_key( VERDICT_KEY, @@ -277,7 +300,7 @@ def assert_local_processes_alive() -> None: ) assert_local_processes_alive() print(f"[pod-{pod_index}] verdict: {verdict}", flush=True) - if verdict != PASS_VERDICT: + if verdict not in (PASS_VERDICT, DEBUG_VERDICT): raise RuntimeError(f"pod {pod_index} run failed: {verdict}") # Phase 12: no pod tears down until every pod has passed its last liveness @@ -285,6 +308,7 @@ def assert_local_processes_alive() -> None: # still checking would report as a failed run. No on_poll here: from now on # a child exiting is expected. rendezvous.barrier(TEARDOWN_BARRIER, args.verdict_timeout) + return verdict def launch_slot( @@ -446,11 +470,189 @@ def print_plan(pod_plan: PodPlan, addresses: list[str]) -> None: ) -def parse_args() -> argparse.Namespace: +def run_case( + case: Case, environment: CaseEnvironment, args: argparse.Namespace +) -> dict: + """Coordinate an existing GPU case; no cross-pod execution probes are added.""" + if any(phase.checks for phase in case.phases): + raise ValueError( + "multi-pod execution checks require cross-pod evidence support" + ) + layout = PodLayout.parse(args.pod_layout) + topology = Topology.from_case(case) + validate_layout(topology, layout) + pod_index = resolve_pod_index(layout.num_pods) + wait_for_address(args.store_host, args.store_dns_timeout) + rendezvous = Rendezvous( + host=args.store_host, + port=args.store_port, + pod_index=pod_index, + num_pods=layout.num_pods, + ) + try: + rendezvous.verify_agreement(layout.canonical(), case.id) + addresses = resolve_addresses(args, rendezvous, layout.num_pods, pod_index) + pod_plan = plan(topology, layout, addresses)[pod_index] + print_plan(pod_plan, addresses) + environment = replace( + environment, devices=environment.devices or pod_plan.devices + ) + effective, record = prepare_case(case, environment, slots=pod_plan.slots) + resolved = deployment.resolve(effective, environment, slots=pod_plan.slots) + record["run_id"] = pod_e2e_run_id(args.run_id, pod_index) + record["pod"] = { + **asdict(pod_plan), + "layout": layout.canonical(), + "role_devices": dict(resolved.role_devices), + } + record["coverage"] = ["task_result", "local_liveness", "local_cleanup"] + services = deployment.ServiceDeployment( + resolved, + run_id=record["run_id"], + extra_environment=parse_key_values(args.pod_env, option="--pod-env"), + ) + environment.output_path.mkdir(parents=True, exist_ok=True) + + def cleanup() -> None: + try: + services.close() + except BaseException as exc: + record["cleanup"] = {"status": "failed", "error": str(exc)} + publish_quietly(rendezvous, f"phase/{pod_index}/cleanup", str(exc)) + raise + else: + record["cleanup"] = {"status": "passed"} + publish_quietly(rendezvous, f"phase/{pod_index}/cleanup", "ok") + finally: + record["processes"] = [ + {"pid": process.pid, "returncode": process.poll()} + for process in services.processes + ] + + def evaluate() -> str: + for phase in effective.phases: + phase_record: dict[str, JsonValue] = { + "name": phase.name, + "status": "running", + "checks": [], + } + record["phases"].append(phase_record) + phase_record["result"] = phase.task.execute( + base_url=f"http://{environment.api_host}:{environment.api_port_base}", + model_name=f"{case.served_model_name_prefix}-attention", + checkpoint=environment.model, + output_path=environment.output_path / phase.name, + ) + phase_record["status"] = "passed" + # This is workload/configuration result propagation, not per-rank + # graph/DBO evidence. Every pod records what the evaluator ran. + rendezvous.set( + "evaluation/result", + json.dumps( + { + "effective": record["effective"], + "configuration": record["configuration"], + "phases": record["phases"], + } + ), + ) + return ( + PASS_VERDICT if record["configuration"]["eligible"] else DEBUG_VERDICT + ) + + try: + with cancellable_run(cleanup): + record["execution"] = {"status": "running"} + try: + verdict = run_pod( + args, + rendezvous, + pod_plan, + launch_order=resolved.launch_order, + start_slot=lambda slot: services.start_role(slot.role), + check_alive=services.ensure_alive, + wait_ready=lambda slot, check: services.wait_ready( + path="/health", + on_poll=check, + ), + evaluate=evaluate, + ) + evaluation = json.loads( + rendezvous.wait_for_key( + "evaluation/result", + args.verdict_timeout, + ) + ) + record.update(evaluation) + except BaseException as exc: + rendezvous.publish_abort(f"pod {pod_index}: {exc}") + raise + record["execution"] = {"status": "passed"} + record["acceptance_passed"] = verdict == PASS_VERDICT + return record + except BaseException as exc: + status = "cancelled" if isinstance(exc, SystemExit) else "failed" + record["execution"] = {"status": status, "error": str(exc)} + for phase_record in record["phases"]: + if phase_record["status"] == "running": + phase_record.update(status=status, error=str(exc)) + raise + finally: + (environment.output_path / "run.json").write_text( + json.dumps(record, indent=2) + "\n" + ) + print( + f"[pod-{pod_index}] CASE {case.id}: {record['execution']['status']}; " + f"acceptance_passed={record['acceptance_passed']}", + flush=True, + ) + except BaseException as exc: + rendezvous.publish_abort(f"pod {pod_index}: {exc}") + raise + finally: + leave_rendezvous( + rendezvous, case.deployment.termination_timeout_s + DEPARTURE_GRACE_S + ) + + +def case_main( + build_case: Callable[[str, CaseEnvironment], Case], + case_ids: tuple[str, ...], + argv: list[str] | None = None, +) -> int: parser = argparse.ArgumentParser( - description="Run one pod's slice of a multi-pod AFD E2E scenario.", + description="Run one pod of a model-declared GPU case" ) - add_scenario_arguments(parser) + add_case_environment_arguments( + parser, + case_ids, + devices_required=False, + startup_option="--serving-timeout", + ) + add_pod_arguments(parser) + parser.set_defaults(startup_timeout=1800) + args = parser.parse_args(argv) + args.serving_timeout = args.startup_timeout + environment = case_environment(args) + if not environment.devices and os.environ.get("CUDA_VISIBLE_DEVICES"): + environment = replace( + environment, devices=tuple(os.environ["CUDA_VISIBLE_DEVICES"].split(",")) + ) + case = build_case(args.scenario, environment) + if args.describe: + print( + json.dumps( + {"canonical": asdict(case), "bindings": asdict(environment)}, + default=str, + indent=2, + ) + ) + else: + run_case(case, environment, args) + return 0 + + +def add_pod_arguments(parser: argparse.ArgumentParser) -> None: parser.add_argument( "--pod-layout", required=True, @@ -492,8 +694,16 @@ def parse_args() -> argparse.Namespace: help="KEY=VALUE added to every launched vLLM process environment.", ) parser.add_argument("--launch-timeout", type=float, default=900) - parser.add_argument("--serving-timeout", type=float, default=1800) parser.add_argument("--verdict-timeout", type=float, default=3600) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Run the retained NPU multi-pod entry.", + ) + add_scenario_arguments(parser) + add_pod_arguments(parser) + parser.add_argument("--serving-timeout", type=float, default=1800) return parser.parse_args() diff --git a/tests/e2e/operators/README.md b/tests/e2e/operators/README.md index d72cb6f80..d35a61e99 100644 --- a/tests/e2e/operators/README.md +++ b/tests/e2e/operators/README.md @@ -20,4 +20,4 @@ and repeated window use. Tolerances are fixed in the script before execution. The retired non-layered checkpoint probe remains available on earlier revisions for a before/after comparison. Validate the current layered implementation using the numerical and real-model fixtures in the -[NPU testing guide](../../../docs/npu/ASYNC_FFN_RUNTIME.md#validation). +[NPU testing guide](../../../docs/design/npu/async_ffn_runtime.md#validation). diff --git a/tests/e2e/presets.py b/tests/e2e/presets.py new file mode 100644 index 000000000..5d929d53b --- /dev/null +++ b/tests/e2e/presets.py @@ -0,0 +1,49 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Reusable execution configuration; models select cases and acceptance settings.""" + +from dataclasses import dataclass + +from tests.e2e.case import ( + EXECUTION_EAGER, + EXECUTION_FULL_DECODE_GRAPH, + DboConfig, + Deployment, + Role, +) + + +@dataclass(frozen=True) +class ExecutionPreset: + """A shared execution mode for both roles of a synchronous deployment.""" + + execution: str + + def deployment( + self, + *, + backend: str, + attention: tuple[int, int], + ffn: tuple[int, int] | None, + connector: str | None, + capture_size: int, + dbo: DboConfig | None = None, + compute_gate_on_attention: bool = False, + ) -> Deployment: + """Layouts are (DP, TP); models may replace explicit budget fields.""" + graph = self.execution == EXECUTION_FULL_DECODE_GRAPH + return Deployment( + backend=backend, + attention=Role(*attention, execution=self.execution), + ffn=Role(*ffn, execution=self.execution) if ffn is not None else None, + connector=connector, + dbo=dbo, + compute_gate_on_attention=compute_gate_on_attention, + cudagraph_capture_size=capture_size if graph else None, + max_num_seqs=capture_size if graph else None, + max_num_batched_tokens=capture_size if ffn is None else None, + ) + + +EAGER = ExecutionPreset(EXECUTION_EAGER) +DECODE_GRAPH = ExecutionPreset(EXECUTION_FULL_DECODE_GRAPH) diff --git a/tests/e2e/runner.py b/tests/e2e/runner.py index 84c4966ba..78aeedf58 100644 --- a/tests/e2e/runner.py +++ b/tests/e2e/runner.py @@ -1,7 +1,7 @@ #!/usr/bin/env python3 # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project -"""Run fixed baseline and AFD E2E scenarios on real hardware.""" +"""Execute model-declared Cases and retain the existing NPU CLI callers.""" from __future__ import annotations @@ -16,18 +16,36 @@ import time import urllib.error import urllib.request -from collections.abc import Mapping +import uuid +from collections.abc import Callable, Mapping +from dataclasses import asdict, replace from pathlib import Path from typing import Any +from tests.e2e import deployment from tests.e2e.accuracy.gsm8k import ( _extract_gsm8k_accuracy, _extract_gsm8k_sample_count, _run_lm_eval, ) from tests.e2e.cancellation import cancellable_run +from tests.e2e.case import ( + NPU_ASYNC_TERMINATION_TIMEOUT_S, + NPU_ASYNC_VLLM_SHUTDOWN_TIMEOUT_S, + NPU_ASYNC_WORKER_SHUTDOWN_TIMEOUT_S, + WORKER_SHUTDOWN_TIMEOUT_ENV, + Case, + CaseEnvironment, + JsonValue, + Mrv2Check, + Role, +) +from tests.e2e.checks.mrv2 import ( + Mrv2Session, + assert_execution_evidence, +) +from tests.e2e.environment import parse_case_environment from tests.e2e.models.deepseek_v4_flash import config as dsv4_config -from tests.e2e.models.deepseek_v4_flash import gpu_config as dsv4_gpu_config from tests.e2e.models.deepseek_v4_flash.completions import evaluate_completions from tests.e2e.models.deepseek_v4_flash.config import ( DSV4_ASYNC_CAM_SCENARIO, @@ -40,10 +58,6 @@ DSV4_SYNC_SHAPES, sync_shape, ) -from tests.e2e.models.deepseek_v4_flash.gpu_config import ( - DSV4_GPU_SCENARIO_SETTINGS, - DSV4_GPU_SCENARIOS, -) from tests.e2e.multi_pod.layout import RoleSlot from tests.e2e.process_utils import ( kill_processes_matching_environment, @@ -58,44 +72,6 @@ # Graph capture default for the scenarios that do not carry their own launch # profile; the DSV4 synchronous profiles override it. DEFAULT_CUDAGRAPH_CAPTURE_SIZE = 8 -V2_SCENARIOS = ( - "afd-v2-eager-1a1f", - "afd-v2-eager-dp2", - "afd-v2-eager-tp2", - "afd-v2-graph-1a1f", - "afd-v2-graph-dp2", - "afd-v2-graph-tp2", - "afd-v2-eager-dbo-dp2", - "afd-v2-graph-dbo-dp2", - "afd-v2-graph-dbo-tp2", -) -# 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"), -) -V2_DBO_SAMPLE_LIMIT = 128 -V2_DBO_MAX_NUM_SEQS = 8 -V2_SINGLE_RANK_SCENARIOS = frozenset( - ("afd-v2-eager-1a1f", "afd-v2-graph-1a1f"), -) -V2_TENSOR_PARALLEL_SCENARIOS = frozenset( - ("afd-v2-eager-tp2", "afd-v2-graph-tp2", "afd-v2-graph-dbo-tp2"), -) E2E_RUN_ID_ENV = "AFD_E2E_RUN_ID" E2E_PROCESS_ROLE_ENV = "AFD_E2E_PROCESS_ROLE" PROCESS_TERMINATION_TIMEOUT_S = 20 @@ -224,42 +200,15 @@ def cleanup() -> None: wait_for_openai_api(args, processes) ensure_processes_alive(processes) - if args.scenario in DSV4_GPU_SCENARIOS: - round_samples = [] - for round_index in range(dsv4_gpu_config.DSV4_GPU_CONCURRENT_ROUNDS): - window = f"concurrent-{round_index + 1}" - start_evaluation_window(args, mrv2_evidence_start, window) - run_concurrent_completion_evaluation(args, round_index=round_index) - round_samples.append( - finish_evaluation_window( - args, - mrv2_evidence_start, - window, - mrv2_execution_evidence, - processes, - ) - ) - if needs_hash_input_evidence(args): - assert_hash_input_freshness(round_samples, args) - start_evaluation_window(args, mrv2_evidence_start, "gsm8k") - run_gsm8k_evaluation(args) - finish_evaluation_window( - args, - mrv2_evidence_start, - "gsm8k", - mrv2_execution_evidence, - processes, - ) - else: - start_evaluation_window(args, mrv2_evidence_start, "evaluation") - run_scenario_evaluation(args) - finish_evaluation_window( - args, - mrv2_evidence_start, - "evaluation", - mrv2_execution_evidence, - processes, - ) + start_evaluation_window(args, mrv2_evidence_start, "evaluation") + run_scenario_evaluation(args) + finish_evaluation_window( + args, + mrv2_evidence_start, + "evaluation", + mrv2_execution_evidence, + processes, + ) ensure_processes_alive(processes) @@ -314,9 +263,6 @@ def add_scenario_arguments(parser: argparse.ArgumentParser) -> None: "afd-graph-2a2f", "afd-graph-dbo-2a2f", *DSV4_SCENARIOS, - *DSV4_GPU_SCENARIOS, - *V2_SCENARIOS, - *DEEPEP_M2N_SCENARIOS, ], required=True, help="Fixed E2E scenario to run.", @@ -409,10 +355,8 @@ def add_scenario_arguments(parser: argparse.ArgumentParser) -> None: def configure_scenario(args: argparse.Namespace) -> None: - """Set topology and features for the selected fixed scenario.""" - is_dsv4 = args.scenario in DSV4_SCENARIOS - is_dsv4_gpu = args.scenario in DSV4_GPU_SCENARIOS - scenario_settings = { + """Configure retained NPU callers; GPU cases are model-owned declarations.""" + settings = { "baseline-graph": (True, True, False, 4, 0), "afd-eager-2a1f": (False, False, False, 2, 1), "afd-graph-2a1f": (False, True, False, 2, 1), @@ -428,91 +372,51 @@ def configure_scenario(args: argparse.Namespace) -> None: DSV4_ATTENTION_RANKS, DSV4_FFN_RANKS, ), - "afd-v2-eager-1a1f": (False, False, False, 1, 1), - "afd-v2-eager-dp2": (False, False, False, 2, 2), - "afd-v2-eager-tp2": (False, False, False, 2, 2), - "afd-v2-graph-1a1f": (False, True, False, 1, 1), - "afd-v2-graph-dp2": (False, True, False, 2, 2), - "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-v2-graph-dbo-tp2": (False, True, True, 2, 2), - **DSV4_GPU_SCENARIO_SETTINGS, - "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 - # is the current suspect for the DSA operator tiling failure on that - # host. - scenario_settings[sync_scenario] = ( + for scenario, profile in DSV4_SYNC_SHAPES.items(): + # The A5 host profile keeps native DBO off: split batches are the + # current suspect for its DSA operator tiling failure. + settings[scenario] = ( False, - dsv4_config.sync_use_graph(sync_profile), + dsv4_config.sync_use_graph(profile), False, - sync_profile.attention_ranks, - sync_profile.ffn_ranks, + profile.attention_ranks, + profile.ffn_ranks, ) - baseline, use_graph, enable_dbo, attention_ranks, ffn_ranks = scenario_settings[ - args.scenario - ] + baseline, graph, dbo, attention_ranks, ffn_ranks = settings[args.scenario] args.baseline = baseline - args.cuda_graph_full_decode_only = use_graph - args.enable_dbo = enable_dbo + args.cuda_graph_full_decode_only = graph + args.enable_dbo = dbo args.num_attention_ranks = attention_ranks args.num_ffn_ranks = ffn_ranks args.attention_data_parallel_address = None - # A runner that cannot drive the evidence window turns this off. args.collect_mrv2_execution_evidence = True args.tp_size = 1 - active_sync_profile = sync_shape(args.scenario) - if args.scenario == DSV4_ASYNC_CAM_SCENARIO: - args.attention_tp_size = DSV4_ATTENTION_TP_SIZE - elif active_sync_profile is not None: - args.attention_tp_size = active_sync_profile.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: - args.ffn_tp_size = 2 - elif active_sync_profile is not None: - # A synchronous DSV4 profile sizes the FFN side exactly like its - # Attention side: A5 shards by data parallel with expert parallelism, - # A3 by tensor parallel. - args.ffn_tp_size = active_sync_profile.ffn_tp_size - else: - args.ffn_tp_size = 1 - if not is_dsv4 and args.gsm8k_output_path is None: + profile = sync_shape(args.scenario) + args.attention_tp_size = ( + DSV4_ATTENTION_TP_SIZE + if args.scenario == DSV4_ASYNC_CAM_SCENARIO + else profile.attention_tp_size + if profile is not None + else 1 + ) + args.ffn_tp_size = profile.ffn_tp_size if profile is not None else 1 + if args.scenario not in DSV4_SCENARIOS and args.gsm8k_output_path is None: raise ValueError("--gsm8k-output-path is required for GSM8K scenarios") - 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: dsv4_config.configure_sync_camp2p_scenario(args) - elif is_dsv4_gpu: - dsv4_gpu_config.configure_scenario(args) - if use_graph: + if graph: args.cudagraph_capture_size = ( - active_sync_profile.cudagraph_capture_size - if active_sync_profile is not None - else dsv4_gpu_config.DSV4_GPU_CAPTURE_SIZE - if is_dsv4_gpu + profile.cudagraph_capture_size + if profile is not None else DEFAULT_CUDAGRAPH_CAPTURE_SIZE ) if args.scenario == NPU_EAGER_DBO_SCENARIO: - # This opt-in eager cell qualifies DBO at the native 0.32 thresholds. args.dbo_decode_token_threshold = NPU_EAGER_DBO_DECODE_THRESHOLD args.dbo_prefill_token_threshold = NPU_EAGER_DBO_PREFILL_THRESHOLD - elif enable_dbo: - # vLLM validates that both DBO thresholds are at least the - # microbatch count (2), so a decode threshold of 1 is rejected. + elif dbo: args.dbo_decode_token_threshold = 2 args.dbo_prefill_token_threshold = 8 @@ -542,11 +446,7 @@ def validate_topology( f"--ffn-devices must contain exactly {args.num_ffn_ranks} device", ) if args.baseline: - baseline_ranks = ( - dsv4_gpu_config.DSV4_GPU_ROLE_RANKS - if args.scenario in DSV4_GPU_SCENARIOS - else 4 - ) + baseline_ranks = 4 if args.num_attention_ranks != baseline_ranks or args.num_ffn_ranks != 0: raise ValueError( f"baseline E2E requires {baseline_ranks} Attention ranks " @@ -555,16 +455,10 @@ def validate_topology( if role_tp_size(args, "attention") != 1: raise ValueError("baseline E2E requires Attention TP=1") return - if args.scenario in V2_SCENARIOS and args.device_backend != "gpu": - raise ValueError("ModelRunnerV2 E2E scenarios require GPU") if args.scenario == NPU_EAGER_DBO_SCENARIO and args.device_backend != "npu": raise ValueError("Lite eager DBO scenario requires NPU") - if args.scenario in DEEPEP_M2N_SCENARIOS and args.device_backend != "gpu": - raise ValueError("DeepEP M2N E2E scenarios require GPU") if args.scenario in DSV4_SCENARIOS and args.device_backend != "npu": raise ValueError("DSV4 scenarios require NPU") - if args.scenario in DSV4_GPU_SCENARIOS and args.device_backend != "gpu": - raise ValueError("CUDA DSV4 scenarios require GPU") for role, rank_count in ( ("attention", args.num_attention_ranks), ("ffn", args.num_ffn_ranks), @@ -607,15 +501,14 @@ def build_baseline_command(args: argparse.Namespace) -> list[str]: cmd.extend(dbo_vllm_arguments(args)) if args.cuda_graph_full_decode_only: capture_size = str(args.cudagraph_capture_size) - if args.scenario not in DSV4_GPU_SCENARIOS: - cmd.extend( - [ - "--max-num-seqs", - capture_size, - "--max-num-batched-tokens", - capture_size, - ], - ) + cmd.extend( + [ + "--max-num-seqs", + capture_size, + "--max-num-batched-tokens", + capture_size, + ], + ) cmd.extend( [ "--max-cudagraph-capture-size", @@ -688,7 +581,11 @@ def build_vllm_command( "serve", args.model, "--shutdown-timeout", - str(VLLM_SHUTDOWN_TIMEOUT_S), + str( + NPU_ASYNC_VLLM_SHUTDOWN_TIMEOUT_S + if is_npu and args.afd_async + else VLLM_SHUTDOWN_TIMEOUT_S + ), "--served-model-name", served_model_name(args, role), "--data-parallel-size", @@ -745,15 +642,6 @@ def build_vllm_command( ) if needs_mrv2_execution_evidence(args): cmd.extend(["--worker-extension-cls", "tests.e2e.mrv2_evidence.Worker"]) - if args.scenario in V2_DBO_COMPARISON_SCENARIOS: - # Keep request concurrency identical across the no-DBO, eager DBO - # and graph DBO comparison; retain native chunked-prefill defaults. - cmd.extend( - [ - "--max-num-seqs", - str(V2_DBO_MAX_NUM_SEQS), - ], - ) profile_compilation_config = ( None if sync_profile is None @@ -774,8 +662,7 @@ def build_vllm_command( # profiles do) keeps it: the capture size only has to cover it. max_num_seqs_args = ( [] - if args.scenario in V2_DBO_COMPARISON_SCENARIOS - or any(arg == "--max-num-seqs" for arg in args.common_vllm_arg) + if any(arg == "--max-num-seqs" for arg in args.common_vllm_arg) else ["--max-num-seqs", capture_size] ) cmd.extend( @@ -865,6 +752,8 @@ def uses_npu_async_process_cleanup(args: argparse.Namespace) -> bool: def process_termination_timeout(args: argparse.Namespace) -> float: """Return the scenario's teardown budget, shared by every runner.""" + if args.device_backend == "npu" and args.afd_async and not args.baseline: + return NPU_ASYNC_TERMINATION_TIMEOUT_S if args.scenario in DSV4_SCENARIOS: return DSV4_PROCESS_TERMINATION_TIMEOUT_S return PROCESS_TERMINATION_TIMEOUT_S @@ -909,15 +798,12 @@ def run_gsm8k_evaluation(args: argparse.Namespace) -> None: """Run the configured GSM8K workload against the scenario's public API.""" if args.gsm8k_output_path is None: raise RuntimeError("--gsm8k-output-path is required for GSM8K scenarios") - is_v2_dbo_comparison = args.scenario in V2_DBO_COMPARISON_SCENARIOS configured_limit = os.environ.get( GSM8K_LIMIT_ENV, - str( - V2_DBO_SAMPLE_LIMIT if is_v2_dbo_comparison else DEFAULT_GSM8K_SAMPLE_LIMIT - ), + str(DEFAULT_GSM8K_SAMPLE_LIMIT), ) sample_limit = None if configured_limit == "all" else int(configured_limit) - if (args.enable_dbo or is_v2_dbo_comparison) and sample_limit is not None: + if args.enable_dbo and sample_limit is not None: sample_limit = max(sample_limit, DBO_EVAL_MIN_SAMPLES) expected_sample_count = ( GSM8K_FULL_SAMPLE_COUNT if sample_limit is None else sample_limit @@ -926,7 +812,7 @@ def run_gsm8k_evaluation(args: argparse.Namespace) -> None: {"timeout_s": FULL_GSM8K_TIMEOUT_S} if sample_limit is None else {} ) scenario_options = {} - if args.enable_dbo or is_v2_dbo_comparison: + if args.enable_dbo: scenario_options["num_concurrent"] = DBO_EVAL_NUM_CONCURRENT role = "baseline" if args.baseline else "attention" results = _run_lm_eval( @@ -968,16 +854,11 @@ def execution_evidence_roles(args: argparse.Namespace) -> tuple[str, ...]: def needs_mrv2_execution_evidence(args: argparse.Namespace) -> bool: - observe_native = args.baseline and args.scenario in DSV4_GPU_SCENARIOS return ( - needs_hash_input_evidence(args) - or (not args.baseline or observe_native) + not args.baseline and (args.cuda_graph_full_decode_only or uses_dbo_execution(args)) - ) and args.collect_mrv2_execution_evidence - - -def needs_hash_input_evidence(args: argparse.Namespace) -> bool: - return args.scenario in DSV4_GPU_SCENARIOS and not args.baseline + and args.collect_mrv2_execution_evidence + ) def start_evaluation_window( @@ -1009,12 +890,7 @@ def finish_evaluation_window( scoped = [(role, line) for role, line in evidence if f"window={window} " in line] if args.cuda_graph_full_decode_only or uses_dbo_execution(args): assert_mrv2_execution_evidence(scoped, args) - return ( - assert_hash_input_evidence(scoped, args) - if needs_hash_input_evidence(args) - and window in dsv4_gpu_config.DSV4_GPU_HASH_WINDOWS - else {} - ) + return {} def mrv2_execution_evidence_rpc( @@ -1081,204 +957,35 @@ def wait_for_mrv2_execution_summaries( def assert_mrv2_execution_evidence( evidence: list[tuple[str, str]], args: argparse.Namespace ) -> None: - metric = ( - "two_stage_graph_replays" - if uses_dbo_execution(args) and args.cuda_graph_full_decode_only - else "two_stage_steps" - if uses_dbo_execution(args) - else "graph_replays" - ) - minimum = ( - MRV2_EVIDENCE_REPLAY_MIN_COUNT if metric == "two_stage_graph_replays" else 1 - ) - layout_kind = ( - "consecutive_graph" if args.cuda_graph_full_decode_only else "two_stage" - ) - common_layouts: set[str] | None = None - for role in execution_evidence_roles(args): - tp_size = role_tp_size(args, role) - rank_count = args.num_ffn_ranks if role == "ffn" else args.num_attention_ranks - expected = { - (dp, tp) for dp in range(rank_count // tp_size) for tp in range(tp_size) - } - observed = set() - layouts: dict[str, set[tuple[int, int]]] = {} - for recorded_role, line in evidence: - if recorded_role != role: - continue - match = re.search( - rf"dp_rank=(\d+) tp_rank=(\d+).*\b{metric}=(\d+)(?: |$)", line - ) - if match is not None and int(match[3]) >= minimum: - rank = (int(match[1]), int(match[2])) - observed.add(rank) - if uses_dbo_execution(args) and not uses_async_connector(args): - payload = json.loads(line.split(" layouts=", 1)[1]) - for layout in payload[layout_kind]: - key = json.dumps(layout, separators=(",", ":")) - layouts.setdefault(key, set()).add(rank) - if missing := expected - observed: - raise RuntimeError( - f"{args.device_backend.upper()} {role} lacks live " - f"{metric}>={minimum} evidence on ranks " - f"{sorted(missing)}" - ) - if uses_dbo_execution(args) and not uses_async_connector(args): - complete_layouts = { - layout for layout, ranks in layouts.items() if expected <= ranks - } - common_layouts = ( - complete_layouts - if common_layouts is None - else common_layouts & complete_layouts - ) - if ( - uses_dbo_execution(args) - and not uses_async_connector(args) - and not common_layouts - ): - raise RuntimeError( - f"{args.device_backend.upper()} DBO lacks matching two-stage layouts " - "across every participating rank" - + ( - " with consecutive graph replays" - if args.cuda_graph_full_decode_only - else "" - ) + roles = execution_evidence_roles(args) + layouts = { + role: Role( + dp=(args.num_ffn_ranks if role == "ffn" else args.num_attention_ranks) + // role_tp_size(args, role), + tp=role_tp_size(args, role), ) - print( - f"[execution-coverage] live {metric}>={minimum} " - "confirmed on every participating rank" - ) - - -def assert_hash_input_evidence( - evidence: list[tuple[str, str]], args: argparse.Namespace -) -> dict[tuple[int, int, int, int], set[tuple[int, ...]]]: - """Pair bounded successful live A/F samples in original control/peer order.""" - samples: dict[str, dict[tuple[int, int, int], dict]] = { - "attention": {}, - "ffn": {}, + for role in roles } - observed_ranks: dict[str, set[tuple[int, int]]] = { - "attention": set(), - "ffn": set(), - } - for role, line in evidence: - match = re.search(r"dp_rank=(\d+) tp_rank=(\d+)", line) - assert match is not None - rank = (int(match[1]), int(match[2])) - payload = json.loads(line.split(" layouts=", 1)[1]) - for sample in payload["hash_ids"]: - if not sample["completed"] or sample["control"] <= 0: - raise RuntimeError( - "Hash ID sample did not complete a live control transaction" - ) - if args.cuda_graph_full_decode_only and not sample["replayed"]: - raise RuntimeError("Hash ID sample lacks completed graph replay") - key = (sample["control"], sample["stage"], sample["peer"]) - if key in samples[role]: - raise RuntimeError(f"Duplicate {role} Hash ID sample: {key}") - if role == "attention" and sample["peer"] != rank[0]: - raise RuntimeError( - f"Attention Hash ID peer does not match DP rank: {key}" - ) - samples[role][key] = sample - observed_ranks[role].add(rank) - for role, ranks in observed_ranks.items(): - rank_count = args.num_ffn_ranks if role == "ffn" else args.num_attention_ranks - expected = {(dp, 0) for dp in range(rank_count)} - if ranks != expected: - raise RuntimeError( - f"{role} lacks bounded Hash ID samples on ranks {expected - ranks}" - ) - if samples["attention"].keys() != samples["ffn"].keys(): - raise RuntimeError("A/F Hash ID samples do not pair by control/stage/peer") - values: dict[tuple[int, int, int, int], set[tuple[int, ...]]] = {} - covered = set() - for key, attention in samples["attention"].items(): - ffn = samples["ffn"][key] - real = attention["real_tokens"] - bucket = attention["bucket"] - if ( - ffn["bucket"] != bucket - or len(ffn["input_ids"]) != bucket - or len(attention["input_ids"]) != real - or not 0 <= real <= bucket - or attention["input_ids"] != ffn["input_ids"][:real] - ): - raise RuntimeError( - f"Native A / actual Hash-layer F input IDs differ: {key}" - ) - if real: - _, stage, peer = key - covered.add((peer, stage)) - values.setdefault((peer, stage, bucket, real), set()).add( - tuple(attention["input_ids"]) - ) - stages = range(DBO_EVAL_NUM_UBATCHES if args.enable_dbo else 1) - required = { - (peer, stage) for peer in range(args.num_attention_ranks) for stage in stages - } - if missing := required - covered: - raise RuntimeError( - f"Hash ID samples lack real tokens for peer/stage {sorted(missing)}" - ) - print("[hash-inputs] native A and actual Hash-layer F IDs match on every rank") - return values - - -def assert_hash_input_freshness( - rounds: list[dict[tuple[int, int, int, int], set[tuple[int, ...]]]], - args: argparse.Namespace, -) -> None: - """Require new real IDs across rounds at the same padded and real size.""" - previous, current = rounds - changed = { - (peer, stage) - for (peer, stage, bucket, real), values in current.items() - if (peer, stage, bucket, real) in previous - and values - previous[(peer, stage, bucket, real)] - } - stages = range(DBO_EVAL_NUM_UBATCHES if args.enable_dbo else 1) - required = { - (peer, stage) for peer in range(args.num_attention_ranks) for stage in stages - } - if missing := required - changed: - raise RuntimeError( - "Hash input IDs did not change across concurrent rounds with the same " - f"bucket and real token count: {sorted(missing)}" - ) - print("[hash-inputs] both concurrent windows show fresh IDs at the same bucket") + assert_execution_evidence( + evidence, + Mrv2Check( + roles, graph=args.cuda_graph_full_decode_only, dbo=uses_dbo_execution(args) + ), + role_layouts=layouts, + backend=args.device_backend, + asynchronous=uses_async_connector(args), + ) -def run_concurrent_completion_evaluation( - args: argparse.Namespace, *, round_index: int | None = None -) -> None: +def run_concurrent_completion_evaluation(args: argparse.Namespace) -> None: profile = sync_shape(args.scenario) role = "baseline" if args.baseline else "attention" - output_path = Path(args.completion_output_path) - if round_index is not None: - output_path = output_path.with_stem( - f"{output_path.stem}-round-{round_index + 1}" - ) evaluate_completions( url=f"http://{args.api_host}:{attention_api_port(args)}/v1/chat/completions", model=served_model_name(args, role), - output_path=output_path, - # A synchronous profile may state that its host does not answer reliably - # yet, and then the oracle checks the concurrent plumbing only. + output_path=Path(args.completion_output_path), + # Preserve the NPU profile's answer/plumbing acceptance policy. check_answer=True if profile is None else profile.check_answer, - first_request_index=( - 0 - if round_index is None - else round_index * dsv4_config.DSV4_CONCURRENT_REQUESTS - ), - answer_repetitions=( - dsv4_gpu_config.DSV4_GPU_ANSWER_REPETITIONS - if args.scenario in DSV4_GPU_SCENARIOS - else 1 - ), ) @@ -1297,14 +1004,8 @@ def build_env( env.setdefault("VLLM_ENGINE_READY_TIMEOUT_S", "18000") env[visible_devices_env_name(args.device_backend)] = visible_devices env["VLLM_USE_V2_MODEL_RUNNER"] = "1" - if needs_hash_input_evidence(args): - env["AFD_E2E_HASH_INPUT_IDS"] = "1" - env["AFD_E2E_HASH_REQUIRE_GRAPH"] = ( - "1" if args.cuda_graph_full_decode_only else "0" - ) - else: - env.pop("AFD_E2E_HASH_INPUT_IDS", None) - env.pop("AFD_E2E_HASH_REQUIRE_GRAPH", None) + env.pop("AFD_E2E_HASH_INPUT_IDS", None) + env.pop("AFD_E2E_HASH_REQUIRE_GRAPH", None) if needs_mrv2_execution_evidence(args): env["VLLM_SERVER_DEV_MODE"] = "1" if args.baseline: @@ -1327,6 +1028,11 @@ def build_env( env[E2E_PROCESS_ROLE_ENV] = role if args.device_backend == "npu": env.pop("VLLM_ASCEND_ENABLE_FLASHCOMM1", None) + if args.afd_async and not args.baseline: + env.setdefault( + WORKER_SHUTDOWN_TIMEOUT_ENV, + str(NPU_ASYNC_WORKER_SHUTDOWN_TIMEOUT_S), + ) env.pop("AFD_PLUGIN_EARLY_ENGINE_PATCH", None) current_pythonpath = env.get("PYTHONPATH") env["PYTHONPATH"] = ( @@ -1469,5 +1175,155 @@ def shell_quote(value: str) -> str: return "'" + value.replace("'", "'\"'\"'") + "'" +def prepare_case( + case: Case, environment: CaseEnvironment, *, slots: tuple[RoleSlot, ...] = () +) -> tuple[Case, dict]: + """Resolve documented overrides without mutating the canonical declaration.""" + effective_phases = [] + reasons: list[str] = [] + for phase in case.phases: + task, phase_reasons = phase.task.resolve(os.environ) + effective_phases.append(replace(phase, task=task)) + reasons.extend(f"{phase.name}: {reason}" for reason in phase_reasons) + profile = case.deployment + override_names: tuple[str, ...] = (GSM8K_LIMIT_ENV, GSM8K_THRESHOLD_ENV) + if profile.worker_shutdown_timeout_s is not None: + override_names += (WORKER_SHUTDOWN_TIMEOUT_ENV,) + worker_timeout = os.environ.get(WORKER_SHUTDOWN_TIMEOUT_ENV) + if worker_timeout is not None: + profile = replace(profile, worker_shutdown_timeout_s=int(worker_timeout)) + effective = replace(case, deployment=profile, phases=tuple(effective_phases)) + resolved = deployment.resolve(effective, environment, slots=slots) + probe_arguments = ( + ("--worker-extension-cls", "tests.e2e.mrv2_evidence.Worker") + if any(phase.checks for phase in effective.phases) + else () + ) + record = { + "case_id": case.id, + "run_id": uuid.uuid4().hex, + "canonical": asdict(case), + "effective": asdict(effective), + "bindings": { + **asdict(environment), + "output_path": str(environment.output_path), + }, + "overrides": { + name: os.environ[name] for name in override_names if name in os.environ + }, + "configuration": { + "eligible": not reasons, + "classification": "debug" if reasons else "qualified", + "reasons": reasons, + }, + "commands": { + role: resolved.command(role, probe_arguments) + for role, _ in resolved.role_devices + }, + "execution": {"status": "not_started"}, + "cleanup": {"status": "not_started"}, + "phases": [], + "acceptance_passed": False, + } + return effective, record + + +def run_case(case: Case, environment: CaseEnvironment) -> dict: + """Execute one deployment; success follows task, evidence and cleanup.""" + effective, record = prepare_case(case, environment) + resolved = deployment.resolve(effective, environment) + probe = Mrv2Session(effective, environment, record["run_id"]) + services = deployment.ServiceDeployment(resolved, run_id=record["run_id"]) + environment.output_path.mkdir(parents=True, exist_ok=True) + + def cleanup() -> None: + failures = [] + try: + services.close() + except BaseException as exc: + failures.append(str(exc)) + finally: + try: + probe.close() + except BaseException as exc: + failures.append(str(exc)) + record["processes"] = [ + {"pid": process.pid, "returncode": process.poll()} + for process in services.processes + ] + if failures: + record["cleanup"] = {"status": "failed", "errors": failures} + raise RuntimeError("; ".join(failures)) + record["cleanup"] = {"status": "passed"} + + try: + with cancellable_run(cleanup): + probe.prepare() + services.probe_arguments = probe.arguments + services.probe_environment = probe.worker_environment + services.start() + services.ensure_alive() + record["execution"] = {"status": "running"} + for phase in effective.phases: + phase_record: dict[str, JsonValue] = { + "name": phase.name, + "status": "running", + "checks": [asdict(check) for check in phase.checks], + } + record["phases"].append(phase_record) + probe.start(phase) + role = "baseline" if effective.deployment.is_baseline else "attention" + phase_record["result"] = phase.task.execute( + base_url=f"http://{environment.api_host}:{environment.api_port_base}", + model_name=f"{case.served_model_name_prefix}-{role}", + checkpoint=environment.model, + output_path=environment.output_path / phase.name, + ) + snapshot = probe.finish(phase, services) + phase_record["evidence_records"] = len(snapshot) + phase_record["status"] = "passed" + services.ensure_alive() + services.finish() + record["execution"] = {"status": "passed"} + record["acceptance_passed"] = record["configuration"]["eligible"] + return record + except BaseException as exc: + record["execution"] = { + "status": "cancelled" if isinstance(exc, SystemExit) else "failed", + "error": str(exc), + } + for phase_record in record["phases"]: + if phase_record["status"] == "running": + phase_record["status"] = record["execution"]["status"] + phase_record["error"] = str(exc) + raise + finally: + (environment.output_path / "run.json").write_text( + json.dumps(record, indent=2) + "\n" + ) + print( + f"\nE2E CASE {case.id}: {record['execution']['status']}; " + f"configuration={record['configuration']['classification']}; " + f"acceptance_passed={record['acceptance_passed']}", + flush=True, + ) + + +def case_main( + build_case: Callable[[str, CaseEnvironment], Case], + case_ids: tuple[str, ...], + argv: list[str] | None = None, +) -> int: + """The model injects its builder; shared code never imports the model.""" + case_id, environment, describe = parse_case_environment(case_ids, argv) + case = build_case(case_id, environment) + if describe: + _, record = prepare_case(case, environment) + print(json.dumps(record, indent=2)) + else: + run_case(case, environment) + return 0 + + if __name__ == "__main__": sys.exit(main()) diff --git a/tests/e2e/test_runner.py b/tests/e2e/test_runner.py index 818e46c61..2bd27bd11 100644 --- a/tests/e2e/test_runner.py +++ b/tests/e2e/test_runner.py @@ -7,37 +7,64 @@ import json import os import signal +import socket import subprocess import sys import threading +import urllib.error +from dataclasses import replace from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer import pytest +from tests.e2e import runner +from tests.e2e.case import ( + CLEANUP_PROFILE_STRICT, + CaseEnvironment, + DboConfig, + Mrv2Check, + Phase, +) +from tests.e2e.models.deepseek_v2_lite import cases as v2_cases +from tests.e2e.models.deepseek_v4_flash import cases as dsv4_cases from tests.e2e.models.deepseek_v4_flash import test_async_cam_npu as entrypoint +from tests.e2e.models.deepseek_v4_flash.completions import CompletionTask CONCURRENT_REQUESTS = 10 CANCEL_SCRIPT = """ import json +import subprocess import sys +from dataclasses import replace from pathlib import Path from tests.e2e import runner +from tests.e2e.deployment import ServiceDeployment +from tests.e2e.models.deepseek_v4_flash.cases import build_case, NPU_ASYNC_IDS pid_path = Path(sys.argv.pop(1)) -pids = [] -start = runner.start_process - -def start_service(name, command, env): - process = start(name, [sys.executable, '-c', 'import time; time.sleep(60)'], env) - pids.append(process.pid) - pid_path.write_text(json.dumps(pids)) - return process - -runner.start_process = start_service -runner.wait_for_openai_api = lambda *args: None -sys.exit(runner.main()) + +def start_services(self): + for role in self.resolved.launch_order: + process = subprocess.Popen( + [sys.executable, '-c', 'import time; time.sleep(60)'], + start_new_session=True, + ) + self.processes.append(process) + self.processes_by_role[role] = process + pid_path.write_text(json.dumps([p.pid for p in self.processes])) + +def host_case(case_id, environment): + case = build_case(case_id, environment) + return replace(case, id="host-completion-cancellation", deployment=replace( + case.deployment, cleanup_profile="strict", termination_timeout_s=1, + )) + +# Isolate only device-service startup; execute the real task, signals and +# process cleanup. This host-only run does not qualify an NPU scenario. +ServiceDeployment.start = start_services +sys.exit(runner.case_main(host_case, NPU_ASYNC_IDS)) """ @@ -91,7 +118,11 @@ def do_POST(self): for pid in json.loads(pid_path.read_text()): with pytest.raises(ProcessLookupError): os.kill(pid, 0) - results = json.loads((tmp_path / "responses.json").read_text()) + results = json.loads( + ( + tmp_path / "responses.json" / "evaluation" / "completions.json" + ).read_text() + ) assert len(results) == CONCURRENT_REQUESTS assert all("cancelled" in row["error"] for row in results) finally: @@ -106,3 +137,654 @@ def do_POST(self): server.shutdown() server.server_close() server_thread.join(timeout=5) + + +HOST_SERVICE = r""" +import json +import os +import re +import signal +import sys +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path + +arguments = sys.argv[1:] +model = arguments[arguments.index("--served-model-name") + 1] +with open(os.environ["AFD_HOST_PIDS"], "a") as log: + log.write(str(os.getpid()) + "\n") +Path(os.environ["AFD_HOST_PIDS"]).with_name(f"{model}.json").write_text(json.dumps({ + "arguments": arguments, + "logging_level": os.environ.get("VLLM_LOGGING_LEVEL"), + "worker_shutdown_timeout": os.environ.get("VLLM_WORKER_SHUTDOWN_TIMEOUT_SECONDS"), +})) +if ( + (os.environ.get("AFD_HOST_OUTCOME") == "startup-failure" + and model.endswith("attention")) + or (os.environ.get("AFD_HOST_OUTCOME") == "peer-startup-failure" + and os.environ.get("AFD_E2E_POD_INDEX") == "1" + and model.endswith("ffn")) +): + raise SystemExit(3) +if os.environ.get("AFD_HOST_OUTCOME") == "cleanup-failure": + signal.signal(signal.SIGTERM, signal.SIG_IGN) +if "--headless" in arguments or model.endswith("-ffn"): + # Connector-driven FFN processes have no HTTP API. + signal.pause() + raise SystemExit(0) +port = int(arguments[arguments.index("--port") + 1]) + +class Handler(BaseHTTPRequestHandler): + def do_GET(self): + self.send_response(200) + self.end_headers() + self.wfile.write(json.dumps({"data": []}).encode()) + + def do_POST(self): + body = json.loads(self.rfile.read(int(self.headers["Content-Length"]))) + if self.path == "/collective_rpc": + # Host services expose no device execution evidence. + self.send_error(503, "host fixture provides no device evidence") + return + prompt = body["messages"][0]["content"] + operands = re.search(r"Compute (\d+) \+ (\d+)", prompt) + repeated = re.search(r"repeated (\d+) times", prompt) + count = int(repeated[1]) if repeated else 1 + answer = str(int(operands[1]) + int(operands[2])) + if os.environ.get("AFD_HOST_OUTCOME") == "wrong-answer": + answer = "incorrect" + choice = {"message": {"content": " ".join([answer] * count)}, + "finish_reason": "stop"} + response = {"choices": [choice]} + self.send_response(200) + self.end_headers() + self.wfile.write(json.dumps(response).encode()) + + def log_message(self, *args): + pass + +ThreadingHTTPServer(("127.0.0.1", port), Handler).serve_forever() +""" + +HOST_EVALUATOR = r""" +import argparse +import json +import os +import signal +import sys +from pathlib import Path + +parser = argparse.ArgumentParser() +parser.add_argument("--model") +parser.add_argument("--model_args") +parser.add_argument("--tasks") +parser.add_argument("--output_path") +parser.add_argument("--log_samples", action="store_true") +parser.add_argument("--num_fewshot", type=int) +parser.add_argument("--limit", type=int, default=1319) +args = parser.parse_args() +outcome = os.environ.get("AFD_HOST_OUTCOME") +if outcome == "cancel": + print("AFD_HOST_EVALUATOR_WAITING", flush=True) + signal.pause() +if outcome == "evaluator-error": + raise SystemExit(7) +output = Path(args.output_path) +output.mkdir(parents=True, exist_ok=True) +result = { + "results": {"gsm8k": {"exact_match,strict-match": 1.0}}, + "n-samples": {"gsm8k": {"effective": args.limit - (outcome == "partial-result")}}, +} +(output / "results.json").write_text(json.dumps(result)) +(output / "invocation.json").write_text(json.dumps(vars(args))) +print(json.dumps(result), flush=True) +""" + + +@pytest.fixture +def host_case_runtime(monkeypatch, tmp_path): + """Real child-process CLI boundaries; this fixture supplies no device evidence.""" + executable = tmp_path / "host-vllm" + executable.write_text(f"#!{sys.executable}\n" + HOST_SERVICE) + executable.chmod(0o700) + package = tmp_path / "stubs" / "lm_eval" + package.mkdir(parents=True) + (package / "__init__.py").write_text("") + (package / "__main__.py").write_text(HOST_EVALUATOR) + monkeypatch.setenv( + "PYTHONPATH", os.pathsep.join((str(package.parent), os.getcwd())) + ) + monkeypatch.setenv("AFD_HOST_PIDS", str(tmp_path / "service-pids")) + monkeypatch.delenv("AFD_GSM8K_LIMIT", raising=False) + monkeypatch.delenv("AFD_GSM8K_THRESHOLD", raising=False) + with socket.socket() as listener: + listener.bind(("127.0.0.1", 0)) + port = listener.getsockname()[1] + return CaseEnvironment( + model="/models/host-test", + vllm_bin=str(executable), + device_backend="gpu", + devices=("0", "1", "2", "3"), + output_path=tmp_path / "result", + api_port_base=port, + startup_timeout=10, + ) + + +@pytest.mark.parametrize( + "worker_override", + [None, "75"], + ids=["default-worker-timeout", "explicit-worker-timeout"], +) +def test_async_cam_shutdown_settings_reach_services_and_workers( + monkeypatch, host_case_runtime, worker_override +): + if worker_override is None: + monkeypatch.delenv("VLLM_WORKER_SHUTDOWN_TIMEOUT_SECONDS", raising=False) + else: + monkeypatch.setenv("VLLM_WORKER_SHUTDOWN_TIMEOUT_SECONDS", worker_override) + environment = replace( + host_case_runtime, + device_backend="npu", + devices=tuple(str(index) for index in range(16)), + ) + case = dsv4_cases.build_case(dsv4_cases.NPU_ASYNC_IDS[0], environment) + # Host stubs never daemonize. Exercise owned-group cleanup here, not the + # real CAM orphan/pidfd teardown, and keep this result out of device coverage. + case = replace( + case, + id="host-async-shutdown", + deployment=replace(case.deployment, cleanup_profile=CLEANUP_PROFILE_STRICT), + ) + record = runner.run_case(case, environment) + assert record["acceptance_passed"] is True + assert record["cleanup"]["status"] == "passed" + for role in ("attention", "ffn"): + invocation = json.loads( + ( + environment.output_path.parent + / f"{case.served_model_name_prefix}-{role}.json" + ).read_text() + ) + arguments = invocation["arguments"] + assert arguments[arguments.index("--shutdown-timeout") + 1] == "90" + assert invocation["worker_shutdown_timeout"] == (worker_override or "60") + assert record["canonical"]["deployment"]["worker_shutdown_timeout_s"] == 60 + assert record["effective"]["deployment"]["worker_shutdown_timeout_s"] == int( + worker_override or "60" + ) + assert record["effective"]["deployment"]["termination_timeout_s"] == 120 + if worker_override is not None: + assert ( + record["overrides"]["VLLM_WORKER_SHUTDOWN_TIMEOUT_SECONDS"] + == worker_override + ) + for service in record["processes"]: + with pytest.raises(ProcessLookupError): + os.kill(service["pid"], 0) + + +@pytest.mark.parametrize( + "default_dbo,attention_dbo,check_roles,expected", + [ + pytest.param( + None, + DboConfig(), + ("attention",), + {"attention": (None, None), "ffn": None}, + id="attention-only-native-defaults", + ), + pytest.param( + DboConfig(2, 8), + DboConfig(4, 16), + ("attention", "ffn"), + {"attention": (4, 16), "ffn": (2, 8)}, + id="role-override-and-inherited-default", + ), + pytest.param(None, DboConfig(), ("ffn",), None, id="reject-disabled-role"), + ], +) +def test_role_dbo_reaches_service_cli_before_requiring_device_evidence( + monkeypatch, host_case_runtime, default_dbo, attention_dbo, check_roles, expected +): # long-test: exercise configuration through launched service CLIs and the probe RPC + environment = host_case_runtime + monkeypatch.setenv("VLLM_LOGGING_LEVEL", "INFO") + source = v2_cases.build_case("afd-v2-eager-dp2", environment) + case = replace( + source, + deployment=replace( + source.deployment, + dbo=default_dbo, + attention=replace(source.deployment.attention, dbo=attention_dbo), + ), + phases=( + Phase( + "evaluation", source.phases[0].task, (Mrv2Check(check_roles, dbo=True),) + ), + ), + ) + if expected is None: + with pytest.raises(ValueError, match="DBO evidence requested"): + runner.run_case(case, environment) + assert not (environment.output_path.parent / "service-pids").exists() + return + + # The real probe RPC must be reached; no host fixture can qualify DBO. + with pytest.raises(urllib.error.HTTPError, match="503"): + runner.run_case(case, environment) + record = json.loads((environment.output_path / "run.json").read_text()) + assert record["acceptance_passed"] is False + assert record["cleanup"]["status"] == "passed" + for role, thresholds in expected.items(): + invocation = json.loads( + ( + environment.output_path.parent + / f"{case.served_model_name_prefix}-{role}.json" + ).read_text() + ) + arguments = invocation["arguments"] + assert ("--enable-dbo" in arguments) == (thresholds is not None) + if thresholds is not None: + assert ( + tuple( + int(arguments[arguments.index(flag) + 1]) + if flag in arguments + else None + for flag in ( + "--dbo-decode-token-threshold", + "--dbo-prefill-token-threshold", + ) + ) + == thresholds + ) + assert invocation["logging_level"] == ( + "DEBUG" if thresholds is not None else "INFO" + ) + for service in record["processes"]: + with pytest.raises(ProcessLookupError): + os.kill(service["pid"], 0) + + +@pytest.mark.parametrize( + "outcome,limit,threshold,error,eligible", + [ + ("pass", None, None, None, True), + ("debug-samples", "3", None, None, False), + ("debug-threshold", None, "0.1", None, False), + ("expanded", "28", None, None, True), + ("partial-result", None, None, "expected 7", True), + ("evaluator-error", None, None, "exited with code 7", True), + ("cleanup-failure", None, None, "forced SIGKILL", True), + ("startup-failure", None, None, "exited with 3", True), + ], +) +def test_case_runs_public_clis_and_records_qualification( + monkeypatch, + host_case_runtime, + outcome, + limit, + threshold, + error, + eligible, +): # long-test: integrates service launch, evaluator CLI, failure and resource cleanup + environment = host_case_runtime + monkeypatch.setenv("AFD_HOST_OUTCOME", outcome) + if limit is not None: + monkeypatch.setenv("AFD_GSM8K_LIMIT", limit) + if threshold is not None: + monkeypatch.setenv("AFD_GSM8K_THRESHOLD", threshold) + source_id = "afd-v2-eager-dp2" if outcome == "startup-failure" else "baseline-graph" + source = v2_cases.build_case(source_id, environment) + case = replace( + source, + id="host-harness", + deployment=replace(source.deployment, termination_timeout_s=0.5), + ) + if error: + with pytest.raises(RuntimeError, match=error): + runner.run_case(case, environment) + else: + runner.run_case(case, environment) + + record = json.loads((environment.output_path / "run.json").read_text()) + assert record["configuration"]["eligible"] is eligible + assert record["execution"]["status"] == ("failed" if error else "passed") + assert record["acceptance_passed"] == (eligible and error is None) + assert record["cleanup"]["status"] == ( + "failed" if outcome == "cleanup-failure" else "passed" + ) + for service in record["processes"]: + with pytest.raises(ProcessLookupError): + os.kill(service["pid"], 0) + if error is None: + invocation = json.loads( + (environment.output_path / "evaluation" / "invocation.json").read_text() + ) + settings = record["effective"]["phases"][0]["task"]["settings"] + assert invocation["limit"] == (int(limit) if limit else 7) + assert invocation["num_fewshot"] == 8 + assert "max_gen_toks=512" in invocation["model_args"] + assert settings["threshold"] == ( + float(threshold) if threshold else (0.27 if limit == "28" else 0.25) + ) + + +def test_case_comparison_debug_preserves_floor_without_qualifying( + monkeypatch, host_case_runtime +): + """Exercise the published evaluator CLI with the existing debug policy.""" + monkeypatch.setenv("AFD_GSM8K_LIMIT", "7") + environment = host_case_runtime + case = replace( + v2_cases.build_case("afd-v2-eager-dp2", environment), + id="host-comparison-debug", + ) + record = runner.run_case(case, environment) + canonical = record["canonical"]["phases"][0]["task"]["settings"] + effective = record["effective"]["phases"][0]["task"]["settings"] + assert (canonical["sample_limit"], canonical["threshold"]) == (128, 0.27) + assert (effective["sample_limit"], effective["threshold"]) == (24, 0.25) + assert record["execution"]["status"] == "passed" + assert record["configuration"]["eligible"] is False + assert record["acceptance_passed"] is False + assert record["cleanup"]["status"] == "passed" + invocation = json.loads( + (environment.output_path / "evaluation" / "invocation.json").read_text() + ) + assert invocation["limit"] == 24 + assert "num_concurrent=12" in invocation["model_args"] + for service in record["processes"]: + with pytest.raises(ProcessLookupError): + os.kill(service["pid"], 0) + + +CANCEL_CASE_SCRIPT = """ +import os +from dataclasses import replace +from pathlib import Path +from tests.e2e.case import CaseEnvironment +from tests.e2e.models.deepseek_v2_lite.cases import build_case +from tests.e2e.runner import run_case + +env = CaseEnvironment( + model="/models/host-test", vllm_bin=os.environ["AFD_HOST_VLLM"], + device_backend="gpu", devices=("0", "1", "2", "3"), + output_path=Path(os.environ["AFD_HOST_OUTPUT"]), + api_port_base=int(os.environ["AFD_HOST_PORT"]), startup_timeout=10, +) +case = replace(build_case("baseline-graph", env), id="host-harness") +run_case(case, env) +""" + + +@pytest.mark.parametrize("signum", [signal.SIGTERM, signal.SIGINT]) +def test_case_cancellation_reaps_evaluator_and_service( + monkeypatch, + host_case_runtime, + signum, +): # long-test: cancellation crosses evaluator, runner and owned service process groups + environment = host_case_runtime + monkeypatch.setenv("AFD_HOST_OUTCOME", "cancel") + monkeypatch.setenv("AFD_HOST_VLLM", environment.vllm_bin) + monkeypatch.setenv("AFD_HOST_OUTPUT", str(environment.output_path)) + monkeypatch.setenv("AFD_HOST_PORT", str(environment.api_port_base)) + waiting = threading.Event() + process = subprocess.Popen( + [sys.executable, "-c", CANCEL_CASE_SCRIPT], + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + start_new_session=True, + ) + lines = [] + + def read_output(): + assert process.stdout is not None + for line in process.stdout: + lines.append(line) + if "AFD_HOST_EVALUATOR_WAITING" in line: + waiting.set() + + thread = threading.Thread(target=read_output, daemon=True) + thread.start() + try: + assert waiting.wait(timeout=15), "".join(lines) + process.send_signal(signum) + process.wait(timeout=15) + thread.join(timeout=5) + assert process.returncode == 128 + signum, "".join(lines) + record = json.loads((environment.output_path / "run.json").read_text()) + assert record["execution"]["status"] == "cancelled" + assert record["cleanup"]["status"] == "passed" + assert record["acceptance_passed"] is False + for service in record["processes"]: + with pytest.raises(ProcessLookupError): + os.kill(service["pid"], 0) + finally: + if process.poll() is None: + os.killpg(process.pid, signal.SIGKILL) + process.wait(timeout=5) + pid_path = environment.output_path.parent / "service-pids" + if pid_path.exists(): + for pid in pid_path.read_text().splitlines(): + with contextlib.suppress(ProcessLookupError): + os.killpg(int(pid), signal.SIGKILL) + + +@pytest.mark.parametrize("outcome", ["pass", "wrong-answer"]) +def test_case_sequences_model_tasks_and_gsm8k( + monkeypatch, host_case_runtime, outcome +): # long-test: real concurrent HTTP task, evaluator subprocess and cleanup + environment = host_case_runtime + monkeypatch.setenv("AFD_HOST_OUTCOME", outcome) + source = v2_cases.build_case("baseline-graph", environment) + case = replace( + source, + id="host-model-tasks", + phases=( + Phase("first", CompletionTask(answer_repetitions=8)), + Phase( + "second", CompletionTask(first_request_index=10, answer_repetitions=8) + ), + replace(source.phases[0], name="gsm8k"), + ), + ) + if outcome == "wrong-answer": + with pytest.raises(RuntimeError, match="wrong answer"): + runner.run_case(case, environment) + else: + runner.run_case(case, environment) + record = json.loads((environment.output_path / "run.json").read_text()) + assert record["cleanup"]["status"] == "passed" + assert record["acceptance_passed"] is (outcome == "pass") + if outcome == "pass": + assert [phase["name"] for phase in record["phases"]] == [ + "first", + "second", + "gsm8k", + ] + for name, start in (("first", 0), ("second", 10)): + results = json.loads( + (environment.output_path / name / "completions.json").read_text() + ) + assert [item["index"] for item in results] == list(range(start, start + 10)) + assert all("error" not in item for item in results) + assert record["phases"][-1]["result"]["sample_count"] == 7 + else: + assert [phase["name"] for phase in record["phases"]] == ["first"] + assert not (environment.output_path / "second").exists() + assert not (environment.output_path / "gsm8k").exists() + + +@pytest.fixture +def multi_pod_runtime(host_case_runtime, tmp_path): + """Real TCPStore peers and owned service/evaluator CLIs; no device simulation.""" + + def run(model, scenario, layout, outcome): + with socket.socket() as listener: + listener.bind(("127.0.0.1", 0)) + store_port = listener.getsockname()[1] + environment = host_case_runtime + children = [] + threads = [] + lines: list[list[str]] = [[], []] + waiting = threading.Event() + other_run = None + if model == "deepseek_v2_lite" and layout == "2A0F,0A2F" and outcome == "pass": + # A different marker is not proof of a stale process or ownership + # of this deployment's devices/ports on a shared host. + other_run = subprocess.Popen( + [sys.executable, "-c", "import signal; signal.pause()"], + env={**os.environ, "AFD_E2E_RUN_ID": "host-independent-run"}, + start_new_session=True, + ) + + def read_output(index, child): + for line in child.stdout: + lines[index].append(line) + if "AFD_HOST_EVALUATOR_WAITING" in line: + waiting.set() + + try: + for index in range(2): + child_env = os.environ.copy() + child_env.update( + AFD_E2E_POD_INDEX=str(index), + AFD_HOST_OUTCOME=outcome, + AFD_HOST_PIDS=str(tmp_path / f"pod-{index}-pids"), + ) + # Only the evaluator receives a weakened override. Its debug + # verdict and effective workload must reach the other pod too. + if outcome == "debug" and index == 0: + child_env["AFD_GSM8K_LIMIT"] = "3" + command = [ + sys.executable, + "-m", + f"tests.e2e.models.{model}.cases", + "multi-pod", + "--scenario", + scenario, + "--model", + environment.model, + "--vllm-bin", + environment.vllm_bin, + "--device-backend", + "gpu", + "--devices", + "0,1", + "--output-path", + str(tmp_path / f"pod-{index}"), + "--api-port-base", + str(environment.api_port_base), + "--pod-layout", + layout, + "--run-id", + tmp_path.name, + "--store-host", + "127.0.0.1", + "--store-port", + str(store_port), + "--pod-addresses", + "127.0.0.1,127.0.0.1", + "--launch-timeout", + "15", + "--serving-timeout", + "15", + "--verdict-timeout", + "20", + ] + child = subprocess.Popen( + command, + env=child_env, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + start_new_session=True, + ) + children.append(child) + thread = threading.Thread( + target=read_output, args=(index, child), daemon=True + ) + thread.start() + threads.append(thread) + if outcome == "cancel": + assert waiting.wait(timeout=25), "".join(sum(lines, [])) + children[0].send_signal(signal.SIGTERM) + for child in children: + child.wait(timeout=45) + for thread in threads: + thread.join(timeout=5) + for index in range(2): + assert (tmp_path / f"pod-{index}" / "run.json").exists(), "".join( + sum(lines, []) + ) + records = [ + json.loads((tmp_path / f"pod-{index}" / "run.json").read_text()) + for index in range(2) + ] + # Assert before emergency fixture cleanup so it cannot hide a leak. + for record in records: + for service in record["processes"]: + with pytest.raises(ProcessLookupError): + os.kill(service["pid"], 0) + if other_run is not None: + assert other_run.poll() is None + return children, records, lines + finally: + for child in children: + if child.poll() is None: + os.killpg(child.pid, signal.SIGKILL) + child.wait(timeout=5) + if other_run is not None: + other_run.terminate() + other_run.wait(timeout=5) + for index in range(2): + pids = tmp_path / f"pod-{index}-pids" + if pids.exists(): + for pid in pids.read_text().splitlines(): + with contextlib.suppress(ProcessLookupError): + os.killpg(int(pid), signal.SIGKILL) + + return run + + +@pytest.mark.parametrize( + "model,scenario,layout,outcome,samples", + [ + ("deepseek_v2_lite", "afd-graph-2a2f", "2A0F,0A2F", "pass", 7), + ("deepseek_v2_lite", "afd-graph-2a2f", "1A1F,1A1F", "pass", 7), + ("qwen3_moe", "afd-deepep-m2n-graph-2a2f", "2A0F,0A2F", "pass", 100), + ("qwen3_moe", "afd-deepep-m2n-graph-2a2f", "1A1F,1A1F", "pass", 100), + ("deepseek_v2_lite", "afd-graph-2a2f", "2A0F,0A2F", "debug", 3), + ( + "deepseek_v2_lite", + "afd-graph-2a2f", + "1A1F,1A1F", + "peer-startup-failure", + None, + ), + ("deepseek_v2_lite", "afd-graph-2a2f", "2A0F,0A2F", "evaluator-error", None), + ("deepseek_v2_lite", "afd-graph-2a2f", "2A0F,0A2F", "cancel", None), + ], +) +def test_multi_pod_case_coordinates_workload_and_owned_cleanup( + multi_pod_runtime, model, scenario, layout, outcome, samples +): + children, records, lines = multi_pod_runtime(model, scenario, layout, outcome) + succeeded = samples is not None + for child, record, log in zip(children, records, lines, strict=True): + assert (child.returncode == 0) is succeeded, "".join(log) + assert record["cleanup"]["status"] == "passed" + assert record["acceptance_passed"] is (outcome == "pass") + assert record["coverage"] == ["task_result", "local_liveness", "local_cleanup"] + if succeeded: + assert record["execution"]["status"] == "passed" + assert record["configuration"]["eligible"] is (outcome == "pass") + assert record["phases"][0]["result"]["sample_count"] == samples + assert record["phases"][0]["checks"] == [] + if layout == "1A1F,1A1F": + headless = records[1]["commands"] + assert all("--headless" in command for command in headless.values()) + assert all("--port" not in command for command in headless.values()) + if outcome == "cancel": + assert children[0].returncode == 128 + signal.SIGTERM + assert records[0]["execution"]["status"] == "cancelled" diff --git a/tests/gpu/test_ffn_faas_gpu.py b/tests/gpu/test_ffn_faas_gpu.py new file mode 100644 index 000000000..038735929 --- /dev/null +++ b/tests/gpu/test_ffn_faas_gpu.py @@ -0,0 +1,585 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""GPU smoke for the standalone FFN service (``--backend gpu-deepep-m2n``). + +Four GPUs: two synthetic Attention ranks built with the DeepEP M2N harness +(fake vLLM config, no model) and two FFN service ranks launched by +``torchrun -m afd_plugin.ffn_faas serve`` from a synthetic Qwen3 MoE +checkpoint. The checkpoint's experts are the harness's deterministic expert +weights (Qwen3-30B-A3B sizes, two layers), so each Attention rank checks the +combined output of every layer against the harness's dense MoE reference: +that covers the FFN-only model, vLLM's loader and EP weight filter, expert +placement, the M2N handshake, eager steps and a captured CUDA graph. + +Teardown: a ``/stop`` on an idle FFN service cancels its control receive; +the service releases its DeepEP lanes together with Attention (a lane's +``destroy`` is a barrier over every M2N rank) and exits 0. When Attention +dies first, the FFN service must fail and exit non-zero. +""" + +from __future__ import annotations + +import contextlib +import json +import os +import signal +import subprocess +import sys +import time +import traceback +import urllib.request +from pathlib import Path +from typing import Any + +import pytest + +torch = pytest.importorskip("torch") +pytest.importorskip("vllm") +pytest.importorskip("safetensors") + +import torch.distributed as dist # noqa: E402 +import torch.multiprocessing as mp # noqa: E402 + +from tests.gpu.deepep_m2n_harness import ( # noqa: E402 + Checker, + available_gpus, + free_ports, + make_connector, + make_expert_weights, + make_inputs, + run_attention_step, + stage_seed, +) +from tests.unit.fakes.vllm_config import ( # noqa: E402 + M2N_HIDDEN_SIZE, + M2N_INTERMEDIATE_SIZE, + M2N_NUM_EXPERTS, + M2N_TOP_K, + M2NTopology, +) + +pytestmark = pytest.mark.gpu + +REPO_ROOT = Path(__file__).resolve().parents[2] +TOPOLOGY = M2NTopology(ffn_size=2, attn_size=2) +NUM_LAYERS = 2 +CAPACITY = 256 +BUCKET = 64 +EAGER_STEPS = 2 +GRAPH_REPLAYS = 3 +HOST = "127.0.0.1" +EXTRA_CONFIG = {"capacity_tokens_per_rank": CAPACITY} +# Short DeepEP and rendezvous timeouts for the Attention-first failure case. +PEER_EXIT_EXTRA_CONFIG = {**EXTRA_CONFIG, "cpu_timeout_s": 10, "gpu_timeout_s": 10} +PEER_EXIT_AFD_TIMEOUT_S = 30 +STARTUP_TIMEOUT_S = 900 +STOP_TIMEOUT_S = 300 +POLL_S = 0.5 +ATTENTION_EXIT_CODE = 23 +MANAGEMENT_PREFIX = '{"rank"' + + +def _skip_without_gpus() -> 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 write_qwen3_m2n_checkpoint(root: Path) -> Path: + """A BF16 Qwen3 MoE checkpoint whose experts are the harness weights. + + Every layer holds the same experts, as the harness's dense reference + assumes. HF names: ``experts..{gate,up,down}_proj``; the harness's + ``w13`` is ``[gate; up]``. The MoE gate is zero: F builds it, never runs it. + """ + from safetensors.torch import save_file + + root.mkdir(parents=True, exist_ok=True) + w13, w2 = make_expert_weights(range(M2N_NUM_EXPERTS)) + w13, w2 = w13.cpu(), w2.cpu() + tensors: dict[str, Any] = {} + for layer in range(NUM_LAYERS): + prefix = f"model.layers.{layer}.mlp" + tensors[f"{prefix}.gate.weight"] = torch.zeros( + M2N_NUM_EXPERTS, M2N_HIDDEN_SIZE, dtype=torch.bfloat16 + ) + for expert in range(M2N_NUM_EXPERTS): + gate, up = w13[expert].chunk(2, dim=0) + name = f"{prefix}.experts.{expert}" + tensors[f"{name}.gate_proj.weight"] = gate.clone() + tensors[f"{name}.up_proj.weight"] = up.clone() + tensors[f"{name}.down_proj.weight"] = w2[expert].clone() + save_file(tensors, str(root / "model.safetensors"), metadata={"format": "pt"}) + config = { + "architectures": ["Qwen3MoeForCausalLM"], + "model_type": "qwen3_moe", + "torch_dtype": "bfloat16", + "hidden_act": "silu", + "hidden_size": M2N_HIDDEN_SIZE, + "intermediate_size": 4 * M2N_HIDDEN_SIZE, + "moe_intermediate_size": M2N_INTERMEDIATE_SIZE, + "num_hidden_layers": NUM_LAYERS, + "num_attention_heads": 32, + "num_key_value_heads": 4, + "head_dim": 128, + "num_experts": M2N_NUM_EXPERTS, + "num_experts_per_tok": M2N_TOP_K, + "norm_topk_prob": True, + "decoder_sparse_step": 1, + "mlp_only_layers": [], + "vocab_size": 1024, + "max_position_embeddings": 4096, + "rms_norm_eps": 1e-6, + "rope_theta": 1000000.0, + "tie_word_embeddings": False, + } + (root / "config.json").write_text(json.dumps(config)) + return root + + +@pytest.fixture(scope="module") +def checkpoint(tmp_path_factory): + """One synthetic checkpoint (about 2.3 GB) shared by this module's tests.""" + _skip_without_gpus() + return write_qwen3_m2n_checkpoint(tmp_path_factory.mktemp("qwen3-m2n")) + + +def _send_control(connector: Any, tokens: tuple[int, ...], **flags: bool) -> None: + """Send one step's control payload and plan it, as Attention's runner does.""" + from afd_plugin.connectors import AFDControlPayload, AFDDPMetadata + + payload = AFDControlPayload( + dp_metadata_list={ + 0: AFDDPMetadata(num_tokens_across_dp_cpu=torch.tensor(tokens)) + }, + is_graph_capturing=flags.get("is_graph_capturing", False), + is_warmup=flags.get("is_warmup", False), + is_graph_replaying=flags.get("is_graph_replaying", False), + ) + connector.control_plane.send_dp_metadata_list(payload) + connector.control_plane.update_state_from_dp_metadata(payload) + + +def attention_scenario(connector: Any, args: dict[str, Any]) -> dict[str, Any]: + """Drive the FFN service like vLLM's Attention runner and check combines. + + Eager steps first; with ``graph``, then Attention's capture sequence: a + warmup payload run eagerly, a capture payload captured into a CUDA graph + on both sides, and replays with fresh inputs. Every step runs all + ``NUM_LAYERS`` MoE layers, as the FFN service runs every layer per + payload. Returns the checker's worst errors and the payload count. + """ + checker = Checker(connector) + tokens = (BUCKET,) * TOPOLOGY.attn_size + rank = connector.role_rank + payloads = 0 + + def inputs(step: int, layer: int) -> tuple[Any, Any, Any]: + return make_inputs(stage_seed(step * NUM_LAYERS + layer, 0), rank, BUCKET) + + def eager_step(step: int, **flags: bool) -> None: + _send_control(connector, tokens, **flags) + for layer in range(NUM_LAYERS): + x, ids, w = inputs(step, layer) + out = torch.empty_like(x) + run_attention_step(connector, [(x, ids, w)], [out], layer) + checker.check_combined(out, x, ids, w) + + step = 0 + for _ in range(args["eager_steps"]): + eager_step(step) + step += 1 + payloads += 1 + if args["graph"]: + eager_step(step, is_warmup=True) + step += 1 + payloads += 1 + statics = [inputs(step, layer) for layer in range(NUM_LAYERS)] + outs = [torch.empty_like(x) for x, _, _ in statics] + _send_control(connector, tokens, is_graph_capturing=True) + payloads += 1 + graph = torch.cuda.CUDAGraph() + torch.cuda.synchronize() + with torch.cuda.graph(graph): + for layer, (static, out) in enumerate(zip(statics, outs, strict=True)): + run_attention_step(connector, [static], [out], layer) + torch.cuda.synchronize() + for _ in range(GRAPH_REPLAYS): + step += 1 + for layer, static in enumerate(statics): + for static_tensor, fresh in zip( + static, inputs(step, layer), strict=True + ): + static_tensor.copy_(fresh) + _send_control(connector, tokens, is_graph_replaying=True) + graph.replay() + torch.cuda.synchronize() + payloads += 1 + for (x, ids, w), out in zip(statics, outs, strict=True): + checker.check_combined(out, x, ids, w) + return {**checker.finish(), "control_payloads": payloads} + + +def _attention_rank_main( + role_rank: int, + args: dict[str, Any], + extra_config: dict[str, Any], + ports: tuple[int, int], + out_dir: str, +) -> None: + """One synthetic Attention rank: connect, run the scenario, then release. + + The default process group spans only the Attention ranks (the FFN ranks + are another launcher's world). After the scenario the rank writes a + ``scenario`` marker, then either exits without cleanup (``exit_early``, + simulating an Attention failure) or waits for the test's ``release`` + marker and closes its connector together with the FFN service. + """ + out = Path(out_dir) + result: dict[str, Any] = {"ok": True, "error": "", "metrics": {}} + connector = None + try: + torch.cuda.set_device(role_rank) + dist.init_process_group( + "gloo", + init_method=f"tcp://{HOST}:{ports[0]}", + world_size=TOPOLOGY.attn_size, + rank=role_rank, + ) + connector = make_connector( + TOPOLOGY, + TOPOLOGY.ffn_size + role_rank, + role_rank, + extra_config, + ports[1], + HOST, + ) + connector.init_afd_connector() + result["metrics"] = attention_scenario(connector, args) + (out / f"scenario{role_rank}").write_text(json.dumps(result["metrics"])) + if args["exit_early"]: + os._exit(ATTENTION_EXIT_CODE) + deadline = time.monotonic() + STOP_TIMEOUT_S + while not (out / "release").exists() and time.monotonic() < deadline: + time.sleep(POLL_S) + except Exception: + result["ok"] = False + result["error"] = traceback.format_exc() + (out / f"scenario{role_rank}").write_text("{}") + finally: + if connector is not None: + try: + connector.close() + except Exception: + result["ok"] = False + result["error"] += traceback.format_exc() + (out / f"rank{role_rank}.json").write_text(json.dumps(result, default=str)) + if dist.is_initialized(): + dist.destroy_process_group() + + +def _start_attention( + args: dict[str, Any], + extra_config: dict[str, Any], + ports: tuple[int, int], + out: Path, +) -> list[Any]: + """Spawn the Attention ranks on the first two visible GPUs.""" + ctx = mp.get_context("spawn") + procs = [ + ctx.Process( + target=_attention_rank_main, + args=(rank, args, extra_config, ports, str(out)), + ) + for rank in range(TOPOLOGY.attn_size) + ] + for proc in procs: + proc.start() + return procs + + +def _ffn_devices() -> str: + """The visible GPUs after the Attention ranks', as a CUDA_VISIBLE_DEVICES value.""" + visible = os.environ.get("CUDA_VISIBLE_DEVICES") + ids = ( + [item.strip() for item in visible.split(",") if item.strip()] + if visible + else [str(index) for index in range(available_gpus())] + ) + return ",".join(ids[TOPOLOGY.attn_size : TOPOLOGY.world_size]) + + +def _start_ffn_service( + checkpoint: Path, + *, + afd_port: int, + master_port: int, + status_file: Path, + log_path: Path, + graph: bool, + extra_config: dict[str, Any], + afd_timeout_s: int, +) -> subprocess.Popen[bytes]: + """Launch both FFN service ranks with torchrun, as an operator would.""" + env = { + **os.environ, + "CUDA_VISIBLE_DEVICES": _ffn_devices(), + "VLLM_USE_V2_MODEL_RUNNER": "1", + "VLLM_PLUGINS": "afd", + "PYTHONPATH": os.pathsep.join( + filter(None, (str(REPO_ROOT), os.environ.get("PYTHONPATH"))) + ), + } + command = [ + sys.executable, + "-m", + "torch.distributed.run", + f"--nproc-per-node={TOPOLOGY.ffn_size}", + "--nnodes=1", + f"--master-addr={HOST}", + f"--master-port={master_port}", + "-m", + "afd_plugin.ffn_faas", + "serve", + str(checkpoint), + "--backend", + "gpu-deepep-m2n", + "--data-parallel-size", + str(TOPOLOGY.ffn_size), + "--tensor-parallel-size", + "1", + "--enable-expert-parallel", + "--additional-config", + json.dumps( + { + "afd": { + "role": "ffn", + "connector": "DeepEPM2NAFDConnector", + "host": HOST, + "port": afd_port, + "num_attention_ranks": TOPOLOGY.attn_size, + "num_ffn_ranks": TOPOLOGY.ffn_size, + "compute_gate_on_attention": True, + "afd_process_group_timeout_s": afd_timeout_s, + "connector_extra_config": extra_config, + } + } + ), + "--status-host", + HOST, + # Port 0: each rank binds a free port and prints it. + "--status-port", + "0", + "--status-file", + str(status_file), + *( + ["--compilation-config", '{"cudagraph_mode":"FULL_DECODE_ONLY"}'] + if graph + else ["--enforce-eager"] + ), + ] + with log_path.open("wb") as log: + return subprocess.Popen( + command, + cwd=REPO_ROOT, + env=env, + stdout=log, + stderr=subprocess.STDOUT, + start_new_session=True, + ) + + +def _management_ports(log_path: Path) -> dict[int, int]: + """Management port per FFN rank, from the object each rank prints at start. + + Both ranks share one log, so their announcements can land on one line. + """ + text = log_path.read_text(errors="replace") + decoder = json.JSONDecoder() + ports = {} + start = text.find(MANAGEMENT_PREFIX) + while start >= 0: + announced, end = decoder.raw_decode(text, start) + ports[int(announced["rank"])] = int(announced["management"][1]) + start = text.find(MANAGEMENT_PREFIX, end) + return ports + + +def _wait_for(predicate: Any, timeout_s: float, what: str, *procs: Any) -> None: + """Poll ``predicate`` until true; fail early if a watched process died.""" + deadline = time.monotonic() + timeout_s + while not predicate(): + for proc in procs: + code = proc.poll() if hasattr(proc, "poll") else proc.exitcode + if code is not None: + raise AssertionError(f"process exited with {code} while waiting {what}") + if time.monotonic() >= deadline: + raise AssertionError(f"timed out waiting {what}") + time.sleep(POLL_S) + + +def _status(port: int, path: str = "/status", method: str = "GET") -> dict[str, Any]: + request = urllib.request.Request(f"http://{HOST}:{port}{path}", method=method) + with urllib.request.urlopen(request, timeout=5) as reply: + return json.loads(reply.read()) + + +def _status_files(status_file: Path) -> dict[int, dict[str, Any]]: + files = {} + for rank in range(TOPOLOGY.ffn_size): + path = status_file.with_name(f"{status_file.stem}-rank-{rank}.json") + if path.exists(): + files[rank] = json.loads(path.read_text()) + return files + + +def _descendants(pid: int) -> list[int]: + """Every live descendant of ``pid``, read from ``/proc`` (Linux only). + + torchrun's workers are not in the launcher's process group, so killing + the group alone leaves them holding GPUs. + """ + children: dict[int, list[int]] = {} + for stat in Path("/proc").glob("[0-9]*/stat"): + with contextlib.suppress(OSError, ValueError, IndexError): + fields = stat.read_text().rsplit(")", 1)[1].split() + children.setdefault(int(fields[1]), []).append(int(stat.parent.name)) + found, pending = [], [pid] + while pending: + for child in children.get(pending.pop(), []): + found.append(child) + pending.append(child) + return found + + +def _kill(service: subprocess.Popen[bytes], attention: list[Any]) -> None: + """Best-effort teardown so a failed test does not leak GPU processes.""" + if service.poll() is None: + for pid in [*_descendants(service.pid), service.pid]: + with contextlib.suppress(ProcessLookupError): + os.kill(pid, signal.SIGKILL) + service.wait(timeout=30) + for proc in attention: + if proc.is_alive(): + proc.kill() + proc.join(timeout=30) + + +@pytest.mark.parametrize("graph", [False, True], ids=["eager", "graph"]) +def test_ffn_service_serves_and_stops_cleanly( # long-test: two launchers + checkpoint, tmp_path, graph +): + attention_args = {"eager_steps": EAGER_STEPS, "graph": graph, "exit_early": False} + group_port, afd_port, master_port = free_ports(3) + status_file = tmp_path / "ffn.json" + log_path = tmp_path / "ffn.log" + service = _start_ffn_service( + checkpoint, + afd_port=afd_port, + master_port=master_port, + status_file=status_file, + log_path=log_path, + graph=graph, + extra_config=EXTRA_CONFIG, + afd_timeout_s=STARTUP_TIMEOUT_S, + ) + attention = _start_attention( + attention_args, EXTRA_CONFIG, (group_port, afd_port), tmp_path + ) + try: + _wait_for( + lambda: all( + (tmp_path / f"scenario{rank}").exists() + for rank in range(TOPOLOGY.attn_size) + ), + STARTUP_TIMEOUT_S, + "for the Attention scenario", + service, + *attention, + ) + metrics = [ + json.loads((tmp_path / f"scenario{rank}").read_text()) + for rank in range(TOPOLOGY.attn_size) + ] + assert all("control_payloads" in rank_metrics for rank_metrics in metrics), ( + (tmp_path / "rank0.json").read_text() + if (tmp_path / "rank0.json").exists() + else metrics + ) + payloads = metrics[0]["control_payloads"] + ports = _management_ports(log_path) + assert sorted(ports) == list(range(TOPOLOGY.ffn_size)) + # A step is counted after its device work completes. + _wait_for( + lambda: all( + _status(port)["completed_batches"] == payloads + for port in ports.values() + ), + STOP_TIMEOUT_S, + "for the FFN service to count every step", + service, + ) + experts_per_rank = M2N_NUM_EXPERTS // TOPOLOGY.ffn_size + for rank, port in ports.items(): + status = _status(port) + assert status["phase"] == "ready" + assert status["m2n_paired"] is True + assert status["compute_ep_world_size"] == TOPOLOGY.ffn_size + assert status["capacity_tokens_per_rank"] == CAPACITY + assert status["attn_tp_size"] == 1 + assert (status["local_expert_first"], status["local_expert_last"]) == ( + rank * experts_per_rank, + (rank + 1) * experts_per_rank - 1, + ) + assert status["loaded_experts"] == status["total_experts"] + assert status["total_experts"] == NUM_LAYERS * experts_per_rank + + for port in ports.values(): + _status(port, "/stop", "POST") + (tmp_path / "release").write_text("") + assert service.wait(timeout=STOP_TIMEOUT_S) == 0, log_path.read_text()[-4000:] + for proc in attention: + proc.join(timeout=STOP_TIMEOUT_S) + results = [ + json.loads((tmp_path / f"rank{rank}.json").read_text()) + for rank in range(TOPOLOGY.attn_size) + ] + assert all(result["ok"] for result in results), results + files = _status_files(status_file) + assert sorted(files) == list(range(TOPOLOGY.ffn_size)) + for status in files.values(): + assert status["phase"] == "stopped" + assert status["resources_released"] is True + finally: + _kill(service, attention) + + +def test_ffn_service_fails_when_attention_dies_first(checkpoint, tmp_path): + attention_args = {"eager_steps": 1, "graph": False, "exit_early": True} + group_port, afd_port, master_port = free_ports(3) + status_file = tmp_path / "ffn.json" + log_path = tmp_path / "ffn.log" + service = _start_ffn_service( + checkpoint, + afd_port=afd_port, + master_port=master_port, + status_file=status_file, + log_path=log_path, + graph=False, + extra_config=PEER_EXIT_EXTRA_CONFIG, + afd_timeout_s=PEER_EXIT_AFD_TIMEOUT_S, + ) + attention = _start_attention( + attention_args, PEER_EXIT_EXTRA_CONFIG, (group_port, afd_port), tmp_path + ) + try: + for proc in attention: + proc.join(timeout=STARTUP_TIMEOUT_S) + assert proc.exitcode == ATTENTION_EXIT_CODE + code = service.wait(timeout=STOP_TIMEOUT_S) + assert code != 0, log_path.read_text()[-4000:] + for status in _status_files(status_file).values(): + assert status["resources_released"] is False + finally: + _kill(service, attention) diff --git a/tests/npu/README.md b/tests/npu/README.md index 3b93a8f84..e68210f4e 100644 --- a/tests/npu/README.md +++ b/tests/npu/README.md @@ -1,7 +1,7 @@ -# Ascend 910C operator tests +# NPU validation -Standalone harnesses, dispatcher checks, and opt-in device tests for the Ascend -operators this plugin owns. Device tests run outside CI; CPU harness checks also +Runtime checks, standalone operator harnesses, and opt-in Ascend device tests. +Device tests run outside CI; CPU harness checks also run in the PR and merge unit-test jobs. Each entry states whether it needs a device, a multi-rank launch, or only a host Python environment. @@ -36,14 +36,12 @@ connector, model weights, or shared-expert tensors. The second group covers the layered routed-expert operators, `gmm_swiglu_quant_v2_layered` (W13) and `grouped_matmul_layered` (W2), which select their layer on device from a `layer_index` input. Those are the operators -`AsyncCAMLayeredExecutor` drives, and their `all_*` lists hold one element per layer -rather than one per expert. Their validation entry points are in -[docs/npu/TESTING.md](../../docs/npu/TESTING.md), and the two precision suites -are described under [Layered operator precision](#layered-operator-precision) -below. +`AsyncCAMLayeredExecutor` drives; their `all_*` lists hold one element per layer +rather than one per expert. See +[layered operator precision](#layered-operator-precision). The async CAM harness implements the AFD schemas documented in -[CAM_ASYNC_ROUTED_OPS.md](../../docs/npu/CAM_ASYNC_ROUTED_OPS.md). The input +[routed CAM operator contract](../../docs/design/npu/cam_async_routed_ops.md). The input values, route table, quantization result, and weighted sum are all produced inside this repository from the raw inputs; the harness depends on no external reference directory at runtime. @@ -54,6 +52,66 @@ unit-test suite. The three 910C dispatcher modules share the optional `ascend_ops_runtime` fixture in `conftest.py`; requesting it loads the compiled extension, while collection and host-only tests stay independent of torch-npu. +## Runtime contract checks + +Use the pinned [Ascend stack](../../docs/installation.md#ascend-npu-installation). +From the repository root: + +```bash +python3 -m pytest -q -ra tests/unit -m "not gpu and not vllm_runtime" +``` + +Run in the matching runtime environment to exercise NPU-dependent tests; +CPU-only skips do not establish runtime correctness. `tests/conftest.py` owns +the required vLLM-Ascend import order. A test that passes alone but fails with +the full suite may indicate leaked module or monkeypatch state. + +## Native DBO contracts and device execution + +Run the focused contract checks in the matching runtime environment: + +```bash +python3 -m pytest -q -ra \ + tests/unit/compat/npu/test_runtime_config.py \ + tests/unit/config/test_validation.py \ + tests/unit/model_executor/test_async_cam_dbo.py \ + tests/unit/model_executor/test_async_cam_layout.py \ + tests/unit/v1/worker/test_npu_model_runner_v2.py \ + tests/unit/v1/worker/test_npu_ubatch_outputs.py \ + tests/unit/v1/worker/test_npu_mla_graph.py +``` + +The single-device executor test uses real NPU streams/events with isolated +subprocesses for failure injection: + +```bash +ASCEND_RT_VISIBLE_DEVICES=0 python3 -m pytest -q -ra tests/npu/test_mrv2_dbo.py +``` + +It checks native child handoff, stream/context restoration and peer failure +wakeup. It does not validate CAMP2P/CAM transport, checkpoint output, or model +accuracy. For broader qualification, run model comparisons at +Attention DP1 and DP > 1, plain TP and model SP, with +uneven prefill lengths and requests crossing child boundaries. Enable DBO on +Attention only. Keep two ubatches and eager execution. + +For synchronous MRV2 CAMP2P, include eager two-stage execution at Attention +DP > 1, many-to-one padding, and graph+DBO. Plain MLA graph+DBO requires +`FULL_DECODE_ONLY`. Both existing Lite graph+DBO cases must report completed +joint two-stage replays during real requests on every participating A/F rank. +Dummy/profile/capture or padded-only stages do not qualify that gate. +`tests/npu/test_mrv2_graph_inputs.py` separately checks persistent input +addresses with real NPU graph reads and modeled builders; model correctness +still belongs to E2E. Record execution, response correctness, revisions and +teardown outcomes separately. + +The original DSV4 Async concurrent case keeps native DBO thresholds. Its ten +short prompts can remain unsplit while satisfying its exact-answer, `stop`, +overlap and physical-cleanup checks; graceful exit failed in the recorded run. +DSV4 DSA two-stage qualification remains deferred. Qualify it separately +with an above-threshold workload, retaining native defaults and requiring +completed nonempty stages on every Attention rank. + ## Python-only validation From the repository root, no torch or Ascend installation is needed for: @@ -215,17 +273,41 @@ execution, and concurrent layers require separate hardware qualification. ### Layered operator precision -Two suites cover the layered routed-expert operators. Both are pytest modules -gated on `AFD_RUN_ASCEND_OP_RUNTIME=1`, both need one device, and neither -performs any collective. +The W4A8/W8A8 executor chains and the W13/W2 precision suites each use one +device and no collective. Run the CPU contracts first, then rebuild the matching +extension/vendor library and run the device checks: + +```bash +pytest -q tests/unit/test_envs.py tests/unit/model_executor/test_async_cam_layered.py +SOC_VERSION=910c pytest -q tests/npu/test_gmm_layered.py \ + tests/npu/test_gmm_swiglu_quant_v2_layered.py +``` ```bash SOC_VERSION=910c AFD_RUN_ASCEND_OP_RUNTIME=1 LOCAL_RANK=0 \ pytest -o addopts= -v -ra -s \ + tests/npu/test_async_cam_layered_w4a8.py \ + tests/npu/test_async_cam_layered_w8a8.py \ tests/npu/test_gmm_swiglu_quant_v2_layered_precision.py \ tests/npu/test_gmm_layered_device_group_list.py ``` +The W4A8 chain compares valid rows with an FP32 reference across distinct +nonzero weights/compensation, per-channel/per-group scales, interleaved layer +IDs, empty/uneven/full expert counts and poisoned capacity tails. Positive +clipping limits use exactly representable matmuls to avoid FP16 cancellation; +zero-limit cases retain the random reference, with the same tolerance. +The W8A8 chain compares native operators and exact eager/FULL replay. +These operator checks do not establish distributed completion or checkpoint +accuracy; use the [model E2E guide](../e2e/README.md) for those gates. + +Checkpoint comparisons must hold the checkpoint and deployment settings fixed. +Record parameter shapes/dtypes/formats, versions, startup path, accuracy, peak +memory and traces. Cover empty ranks, multi-chunk layers, interleaved DP work +and async ubatching. Check for layer-ID/count metadata D2H between receive and +send; measure group synchronization separately. Performance claims require a +controlled A/B comparison. A skipped device check establishes no qualification. + `test_gmm_swiglu_quant_v2_layered_precision.py` collects 11 cases against W13. It covers A8W8 and A8W4, and it needs two reference strategies because only one scenario has a usable built-in counterpart: @@ -436,10 +518,9 @@ Device cleanup after failure still requires validation on the target stack. ## check_routing_boundary_reuse.py Device check for the opaque routing-id boundary. Requires an Ascend host with the pinned stack -(`vllm`, `vllm-ascend`). Run it directly: +(`vllm`, `vllm-ascend`). Run from the repository root: ```bash -cd python tests/npu/check_routing_boundary_reuse.py ``` diff --git a/tests/npu/async_cam_runtime.py b/tests/npu/async_cam_runtime.py index d57e4aee9..74c7e3c59 100644 --- a/tests/npu/async_cam_runtime.py +++ b/tests/npu/async_cam_runtime.py @@ -31,6 +31,7 @@ from afd_plugin.connectors.npu.async_cam import CAMAsyncAFDConnector from afd_plugin.model_executor.npu.async_cam_layered import LayeredFFNWeights from afd_plugin.v1.worker.npu.async_ffn_model_runner import AsyncNPUFFNModelRunner +from afd_plugin.v1.worker.npu.ffn_worker import AFDNPUFFNWorker HIDDEN = 256 EXPERTS = 4 @@ -41,35 +42,106 @@ class SyntheticModel(torch.nn.Module): """Diagonal experts shared by eager/FULL runs; no checkpoint is loaded.""" - def __init__(self, device, per_channel): + def __init__( + self, + device, + per_channel, + hidden=HIDDEN, + intermediate=HIDDEN, + weight_seed=None, + routed_scale=1.0, + quantization="W4A8", + swiglu_limit=LIMIT, + experts=EXPERTS, + ): super().__init__() self.layers = [] for slot, layer_id in enumerate(LAYERS): - w13, s13, b13 = self.make_weight(device, per_channel, slot + 1, True) - w2, s2, b2 = self.make_weight(device, per_channel, 1, False) + seed = None if weight_seed is None else weight_seed + slot * 2 + w13, s13, b13 = self.make_weight( + device, + per_channel, + slot + 1, + True, + hidden, + intermediate, + seed, + quantization, + experts, + ) + w2, s2, b2 = self.make_weight( + device, + per_channel, + 1, + False, + hidden, + intermediate, + None if seed is None else seed + 1, + quantization, + experts, + ) self.layers.append( LayeredFFNWeights( - layer_id, w13, w2, s13, s2, b13, b2, per_channel, LIMIT, 1.0 + layer_id, + w13, + w2, + s13, + s2, + b13, + b2, + per_channel, + swiglu_limit, + routed_scale, + quantization=quantization, ) ) @staticmethod - def make_weight(device, per_channel, factor, gate_up): - width = HIDDEN * (2 if gate_up else 1) - values = torch.zeros(EXPERTS, HIDDEN, width, dtype=torch.int8) - diagonal = torch.arange(HIDDEN) - values[:, diagonal, diagonal] = 1 - if gate_up: - values[:, diagonal, diagonal + HIDDEN] = 1 + def make_weight( + device, + per_channel, + factor, + gate_up, + hidden=HIDDEN, + intermediate=HIDDEN, + seed=None, + quantization="W4A8", + experts=EXPERTS, + ): + k = hidden if gate_up else intermediate + width = 2 * intermediate if gate_up else hidden + if seed is None: + values = torch.zeros(experts, k, width, dtype=torch.int8) + diagonal = torch.arange(min(k, intermediate if gate_up else width)) + values[:, diagonal, diagonal] = 1 + if gate_up: + values[:, diagonal, diagonal + intermediate] = 1 + else: + generator = torch.Generator().manual_seed(seed) + values = torch.randint( + -3 if quantization == "W4A8" else -127, + 4 if quantization == "W4A8" else 128, + (experts, k, width), + dtype=torch.int8, + generator=generator, + ) groups = 1 if per_channel else 2 - scale = torch.full((EXPERTS, groups, width), float(factor)) + scale = torch.full((experts, groups, width), float(factor)) + if seed is not None: + scale *= 0.01 + 0.02 * torch.rand(scale.shape, generator=generator) if not per_channel: scale[:, 1] *= 2 - dequant = values.float() * scale.repeat_interleave(HIDDEN // groups, 1) + if quantization == "W8A8": + weight = torch_npu.npu_format_cast(values.to(device), 29) + weight_scale = scale.squeeze(1) + if not gate_up: + weight_scale = weight_scale.bfloat16() + return weight, weight_scale.to(device), None + dequant = values.float() * scale.repeat_interleave(k // groups, 1) bias = (8 * dequant.sum(1)).to(device) packed = values.reshape(-1, 2) packed = ((packed[:, 1] << 4) | (packed[:, 0] & 15)).reshape( - EXPERTS, HIDDEN, width // 2 + experts, k, width // 2 ) weight = torch_npu.npu_format_cast(packed.to(device), 29).view(torch.int32) encoded = scale.view(torch.int32).to(torch.int64) @@ -87,16 +159,46 @@ def get_async_cam_layered_weights(self): def main(): parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--per-group", action="store_true") + parser.add_argument("--mega", action="store_true") + parser.add_argument("--quantization", choices=("W4A8", "W8A8"), default="W4A8") + parser.add_argument("--swiglu-limit", type=float, default=LIMIT) + parser.add_argument( + "--experts-per-rank", type=int, choices=(4, 32, 64), default=EXPERTS + ) + parser.add_argument("--idle-probe", action="store_true") + parser.add_argument( + "--hidden-size", type=int, choices=(256, 512, 1024, 2048), default=HIDDEN + ) + parser.add_argument( + "--intermediate-size", + type=int, + choices=(256, 512, 1024, 2048, 3072, 3328, 4096, 8192), + default=HIDDEN, + ) + parser.add_argument("--weight-seed", type=int) + parser.add_argument("--routed-scale", type=float, default=1.0) parser.add_argument("--full-graph", action="store_true") - parser.add_argument("--attn-ranks", type=int, choices=(1, 2), default=1) + parser.add_argument("--attn-ranks", type=int, choices=(1, 2, 4), default=1) + parser.add_argument("--attention-tp", type=int, choices=(1, 2, 4), default=1) parser.add_argument("--boundary", action="store_true") + parser.add_argument("--receive-capacity", type=int, default=16) + parser.add_argument("--boundary-tokens", type=int, default=16) + parser.add_argument("--idle-shutdown", action="store_true") parser.add_argument("--save-outputs") parser.add_argument("--reference-outputs") args = parser.parse_args() + if not 1 <= args.receive_capacity <= 512 or not 7 <= args.boundary_tokens <= 512: + parser.error("capacity must be in [1,512], boundary tokens in [7,512]") + if args.attn_ranks % args.attention_tp: + parser.error("Attention rank count must be divisible by TP") + if (args.idle_shutdown or args.idle_probe) and not args.mega: + parser.error("idle shutdown/probe requires --mega") + if args.per_group and (args.mega or args.quantization == "W8A8"): + parser.error("native Mega requires per-channel weights") if args.full_graph and not args.reference_outputs: parser.error("FULL replay requires saved matching eager outputs") faulthandler.dump_traceback_later(60) - os.environ["BATCH_SIZE_FACTOR"] = str(16 / 262144) + os.environ["BATCH_SIZE_FACTOR"] = str(args.receive_capacity / 262144) os.environ["HCCL_BUFFSIZE"] = "64" ensure_cam_async_ops_available() rank = int(os.environ["RANK"]) @@ -129,7 +231,8 @@ def main(): "port": port[0], "connector_extra_config": { "dynamicQuant": 1, - "attn_ranks_per_dp": 1, + "ffn_mega_moe": args.mega, + "attn_ranks_per_dp": args.attention_tp, "hccl_buffer_size": 64, }, } @@ -138,21 +241,23 @@ def main(): dtype=torch.bfloat16, enforce_eager=not (args.full_graph and is_ffn), hf_config=SimpleNamespace( - hidden_size=HIDDEN, + hidden_size=args.hidden_size, num_experts_per_tok=2, - n_routed_experts=8, + n_routed_experts=2 * args.experts_per_rank, num_hidden_layers=max(LAYERS) + 1, architectures=["DeepseekV4ForCausalLM"], model_type="deepseek_v4", ), ), - scheduler_config=SimpleNamespace(max_num_batched_tokens=16), + scheduler_config=SimpleNamespace( + max_num_batched_tokens=args.boundary_tokens * args.attention_tp + ), parallel_config=SimpleNamespace( enable_dbo=False, use_ubatching=False, num_ubatches=2, - data_parallel_size=2 if is_ffn else args.attn_ranks, - tensor_parallel_size=1, + data_parallel_size=2 if is_ffn else args.attn_ranks // args.attention_tp, + tensor_parallel_size=1 if is_ffn else args.attention_tp, pipeline_parallel_size=1, prefill_context_parallel_size=1, decode_context_parallel_size=1, @@ -174,23 +279,68 @@ def main(): parse_afd_config(config), rank - args.attn_ranks if is_ffn else rank, ) - report = {"rank": rank, "per_group": args.per_group, "work": []} + report = { + "rank": rank, + "per_group": args.per_group, + "mega": args.mega, + "quantization": args.quantization, + "swiglu_limit": args.swiglu_limit, + "experts_per_rank": args.experts_per_rank, + "hidden": args.hidden_size, + "intermediate": args.intermediate_size, + "weight_seed": args.weight_seed, + "routed_scale": args.routed_scale, + "attention_tp": args.attention_tp, + "work": [], + } print(json.dumps({"rank": rank, "stage": "connector_created"}), flush=True) if is_ffn: runner = AsyncNPUFFNModelRunner(config, device, connector) runner.attach_model( - SyntheticModel(device, not args.per_group), parameter_bytes=0 + SyntheticModel( + device, + not args.per_group, + args.hidden_size, + args.intermediate_size, + None if args.weight_seed is None else args.weight_seed + rank * 10000, + args.routed_scale, + args.quantization, + args.swiglu_limit, + args.experts_per_rank, + ), + parameter_bytes=0, ) print(json.dumps({"rank": rank, "stage": "weights_ready"}), flush=True) runner.initialize_afd_connector() assert runner._layered_executor is not None + + def work_state(): + if args.mega: + mega = runner._mega_executor + return SimpleNamespace( + token_nums_rankid_layeridx=mega.batch_info, + group_list=mega.group_list, + ) + return runner._graph_payload.context.states + if args.full_graph: assert runner.effective_mode == "full" assert runner._graph_capture_count == 1 def capture_addresses(): session = runner - state = session._graph_payload.context.states + if args.mega: + mega = session._mega_executor + return tuple( + t.data_ptr() + for t in ( + mega.scratch, + mega.layer_map, + mega.anchor, + session._graph_output, + ) + ) + state = work_state() return tuple( tensor.data_ptr() for tensor in ( @@ -218,7 +368,7 @@ def forbid_python_transaction(*args, **kwargs): # Test-only per-replay observations, fenced once at the batch end. # Production batching performs no host metadata copy. session = runner - state = session._graph_payload.context.states + state = work_state() headers = torch.empty( (len(LAYERS), state.token_nums_rankid_layeridx.numel()), dtype=torch.int64, @@ -249,10 +399,12 @@ def observed_replay(): # warms the original operators concurrently before ordered DP cases. with torch.inference_mode(): if is_ffn: - for _ in range(args.attn_ranks): + for _ in range(args.attn_ranks // args.attention_tp): runner._ffn_forward_connector_driven(single=True) else: - hidden = torch.ones((1, HIDDEN), dtype=torch.bfloat16, device=device) + hidden = torch.ones( + (1, args.hidden_size), dtype=torch.bfloat16, device=device + ) context = AFDTransferContext(AFDTransferMetadata(LAYERS[0], 0, [1])) connector.send_attn_output( hidden, @@ -264,21 +416,73 @@ def observed_replay(): torch.npu.synchronize() dist.barrier() print(json.dumps({"rank": rank, "stage": "transport_ready"}), flush=True) + if args.idle_probe: + assert args.mega + if is_ffn: + pull_state = connector.comm_args.view(torch.int64) + before = int(pull_state[96].item()) + pull_state[64].fill_(8) + if args.full_graph: + runner._graph_stream.wait_stream(torch.npu.current_stream(device)) + connector.graph_startup.service_ready() + with torch.inference_mode(): + for _ in range(3): + if args.full_graph: + batch_index = 0 + runner._submit_graph(1) + runner._complete_graph() + else: + runner._ffn_forward_connector_driven(single=True) + torch.npu.synchronize() + assert int(pull_state[96].item()) == before + assert int(pull_state[32].item()) == -1 + connector.ffn_session().after_step() + pull_state[64].zero_() + torch.npu.synchronize() + report["idle_probe"] = "PASS" + dist.barrier() errors = [] + accuracy_failures: list[str] = [] + data_replays = 0 if not is_ffn: runner = thread = None cases = [(2, False, 7), (5, True, 7), (2, False, 7)] if args.boundary: - cases.extend(((5, False, 16), (2, True, 16))) + cases.extend( + ((5, False, args.boundary_tokens), (2, True, args.boundary_tokens)) + ) expected: list[tuple[int, int, int | None]] = [] for layer_id, sparse, tokens in cases: - for source in range(args.attn_ranks): - chunks = 2 if sparse and tokens == 16 and rank == args.attn_ranks else 1 - rows = 0 if sparse and rank == args.attn_ranks + 1 else None - expected.extend( - (source, layer_id, 16 if chunks == 2 else rows) for _ in range(chunks) - ) + ids = torch.arange(tokens * 2).remainder( + 2 if sparse else 2 * args.experts_per_rank + ) + local_ffn = rank - args.attn_ranks + rows = int( + ( + (ids >= local_ffn * args.experts_per_rank) + & (ids < (local_ffn + 1) * args.experts_per_rank) + ).sum() + ) + rows *= args.attention_tp + # Layered DR packs whole experts; Mega Pull slices individual rows. + chunks = [0] + for expert in range( + local_ffn * args.experts_per_rank, (local_ffn + 1) * args.experts_per_rank + ): + count = int((ids == expert).sum()) * args.attention_tp + if not args.mega: + assert count <= args.receive_capacity + if chunks[-1] + count > args.receive_capacity: + chunks.append(0) + chunks[-1] += count + for source in range(0, args.attn_ranks, args.attention_tp): + if args.mega: + # The exposed metadata is the final slice of one fused transaction. + last_rows = (rows - 1) % args.receive_capacity + 1 if rows else 0 + expected.append((source, layer_id, last_rows)) + continue + expected.extend((source, layer_id, count) for count in chunks) def observe(context): # Test-only host inspection after the matching DR/FFN/CS fence. @@ -292,7 +496,17 @@ def observe(context): } ) - if is_ffn and not args.full_graph: + if is_ffn and not args.full_graph and args.mega: + original_mega = runner._mega_executor.op + + def observe_mega(*op_args): + result = original_mega(*op_args) + torch.npu.current_stream().synchronize() + observe(SimpleNamespace(states=work_state())) + return result + + runner._mega_executor.op = observe_mega + elif is_ffn and not args.full_graph: original_send = connector.send_ffn_output def observe_eager(output, context, **kwargs): @@ -303,8 +517,26 @@ def observe_eager(output, context, **kwargs): connector.send_ffn_output = observe_eager + def finish_receiver(): + if args.idle_shutdown: + if args.full_graph: + runner._graph = graph + elif args.mega: + runner._mega_executor.op = original_mega + else: + connector.send_ffn_output = original_send + event = threading.Event() + # A local shutdown must still serve until all A ranks stop publishing. + event.set() + worker = SimpleNamespace( + _ffn_shutdown_event=event, device=device, model_runner=runner + ) + AFDNPUFFNWorker._run_ffn_server_loop(worker) + report["idle_shutdown"] = "PASS" + runner.mark_receiver_exited() + def serve(): - nonlocal batch_index + nonlocal batch_index, data_replays try: torch.npu.set_device(local_rank) if not args.full_graph: @@ -313,12 +545,14 @@ def serve(): with torch.inference_mode(): for _ in expected: runner._ffn_forward_connector_driven(single=True) + finish_receiver() return connector.graph_startup.service_ready() + replay_before = runner._graph_replay_count for start in range(0, len(expected), len(LAYERS)): batch_index = 0 # Test-only finite work count, including the final short batch. - # The production loop has no EOS or cross-role stop protocol. + # The data check uses a fixed work count before the idle-stop probe. count = min(len(LAYERS), len(expected) - start) runner._submit_graph(count) runner._complete_graph() @@ -342,10 +576,14 @@ def serve(): torch.inference_mode(), torch.npu.stream(session._graph_stream), ): - session._graph_payload.hidden_states.fill_(113) - session._graph_payload.context.states.dynamic_scales.fill_(3) + if args.mega: + session._mega_executor.scratch.fill_(113) + else: + session._graph_payload.hidden_states.fill_(113) + work_state().dynamic_scales.fill_(3) session._graph_output.fill_(123) - runner.mark_receiver_exited() + data_replays = runner._graph_replay_count - replay_before + finish_receiver() except Exception as exc: errors.append(exc) print(json.dumps({"rank": rank, "service_error": repr(exc)}), flush=True) @@ -363,19 +601,25 @@ def serve(): f"{args.reference_outputs}.rank{rank}.pt", weights_only=True ) for layer_id, sparse, tokens in cases: - for source in range(args.attn_ranks): - if rank == source: - values = torch.tensor([-12, -9, -0.5, 0.5, 9, 12, 1]).repeat(3)[:tokens] + for source in range(0, args.attn_ranks, args.attention_tp): + if source <= rank < source + args.attention_tp: + values = torch.tensor([-12, -9, -0.5, 0.5, 9, 12, 1]) + values = values.repeat((tokens + values.numel() - 1) // values.numel())[ + :tokens + ] inputs = ( (values * (rank + 1)) .bfloat16()[:, None] - .expand(-1, HIDDEN) + .expand(-1, args.hidden_size) .contiguous() ) + if args.weight_seed is not None: + feature = torch.sin(torch.arange(args.hidden_size).float() * 0.17) + inputs = (inputs.float() * feature[None, :] + 0.125).bfloat16() ids = ( torch.arange(tokens * 2) .reshape(tokens, 2) - .remainder(2 if sparse else 8) + .remainder(2 if sparse else 2 * args.experts_per_rank) ) hidden = inputs.to(device) context = AFDTransferContext(AFDTransferMetadata(layer_id, 0, [tokens])) @@ -398,9 +642,25 @@ def serve(): time.sleep(0.01) actual = output.cpu() if references is not None: - torch.testing.assert_close( - actual, references[len(outputs)], rtol=0, atol=0 - ) + try: + torch.testing.assert_close( + actual, references[len(outputs)], rtol=0, atol=0 + ) + except AssertionError as exc: + if args.save_outputs and not accuracy_failures: + torch.save( + { + "actual": actual, + "reference": references[len(outputs)], + "layer": layer_id, + "sparse": sparse, + "tokens": tokens, + }, + f"{args.save_outputs}.mismatch.rank{rank}.pt", + ) + # Finish all peer transactions before reporting numerical + # failure, so teardown cannot invalidate a pending pull. + accuracy_failures.append(str(exc)) outputs.append(actual) report["work"].append( {"layer": layer_id, "sparse": sparse, "tokens": tokens} @@ -413,8 +673,11 @@ def serve(): torch.save(outputs, f"{args.save_outputs}.rank{rank}.pt") if references is not None: assert len(references) == len(outputs) - report["eager_exact_match"] = True - assert not connector._pending_attention_payloads + report["eager_exact_match"] = not accuracy_failures + report["accuracy_failures"] = len(accuracy_failures) + assert not connector.has_pending_attention + if args.idle_shutdown: + connector.close() else: assert thread is not None thread.join(timeout=30) @@ -422,22 +685,26 @@ def serve(): if errors: raise errors[0] assert len(report["work"]) == len(expected) - for observed, (source, layer, rows) in zip( + for observed, (source, layer, expected_rows) in zip( report["work"], expected, strict=True ): assert (observed["source"], observed["layer"]) == (source, layer) - if rows is not None: - assert observed["rows"] == rows, "zero work or chunk mismatch" + if expected_rows is not None: + assert observed["rows"] == expected_rows, ( + f"zero work or chunk mismatch: {report['work']=}, {expected=}" + ) if args.full_graph: report["graph"] = { "capture_count": runner._graph_capture_count, "replay_count": runner._graph_replay_count, + "data_replay_count": data_replays, "python_transaction_disabled": True, "stable_exposed_addresses": len(captured_addresses), } - assert runner._graph_replay_count == len(expected) + assert data_replays == len(expected) # Retain the TCPStore owner until finite eager peers have consumed all work. - dist.barrier() + if not args.idle_shutdown: + dist.barrier() if is_ffn: runner.shutdown() runner.shutdown() @@ -446,12 +713,22 @@ def serve(): connector.close() assert not connector.is_initialized and connector.cam_pg is None report["cleanup_complete"] = True - reports = [None] * world_size + reports: list[dict] = [{} for _ in range(world_size)] dist.all_gather_object(reports, report) dist.destroy_process_group() faulthandler.cancel_dump_traceback_later() + failed = any(item.get("accuracy_failures", 0) for item in reports) if rank == 0: - print(json.dumps({"result": "PASS", "reports": reports}), flush=True) + print( + json.dumps({"result": "FAIL" if failed else "PASS", "reports": reports}), + flush=True, + ) + if failed: + raise AssertionError( + accuracy_failures[0] + if accuracy_failures + else "peer accuracy comparison failed" + ) if __name__ == "__main__": diff --git a/tests/npu/check_pull_layout.cpp b/tests/npu/check_pull_layout.cpp new file mode 100644 index 000000000..4ab0129cb --- /dev/null +++ b/tests/npu/check_pull_layout.cpp @@ -0,0 +1,85 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +#include +#include +#include +#include +#include +#include "utils/op_kernel/afd_pull_layout.h" +#include "utils/op_kernel/afd_native_mega_workspace.h" + +int main() { + using namespace Cam::Pull; + Layout layout{}; + assert(MakeLayout({2048, 4096, 8, 8, 8, 32, 4}, layout)); + assert(layout.payload >= layout.returned + 128ULL * 1024 * 1024); + assert(layout.bytes - layout.payload == 72ULL * 1024 * 1024); + assert(!MakeLayout({8192, 4096, 8, 8, 8, 32, 4}, layout)); + assert(!MakeLayout({2048, 4096, 8, 8, 7, 32, 4}, layout)); + assert(!MakeLayout({UINT64_MAX, 4096, 8, 8, 8, 32, 4}, layout)); + assert(TransportUbBytes({8, 256, 2, 383, 1, 256, 1}) > 192 * 1024); + assert(PublisherUbBytes({8, 256, 2, 383, 1, 256, 1}, 48) > 192 * 1024); + const Shape largePublisher{8, 7168, 8, 4, 1, 128, 1}; + assert(PublisherUbBytes(largePublisher, 48) + TransportUbBytes(largePublisher) > 192 * 1024); + assert(CombineUbBytes({8, 256, 2, 4, 1, 256, 1}, 48) > 192 * 1024); + assert(CombineUbBytes({128, 4096, 6, 4, 4, 64, 4}, 48) <= 192 * 1024); + const Cam::NativeMegaWorkspace native(512, 4096, 2048, 128, 4); + assert(native.a2 - native.a1 >= 512ULL * 4096); + assert(native.c1 - native.a2 >= 512ULL * 2048); + assert(native.c2 - native.c1 >= 8ULL * 512 * 2048); + assert(native.output - native.c2 >= 4ULL * 512 * 4096); + assert(native.bytes - native.output == 2ULL * 512 * 4096); + assert(native.output % 512 == 0); + assert(native.layer - native.tileCounts >= 8ULL * 128 * (512 / Cam::NATIVE_TILE_ROWS)); + assert(native.a1 - native.x >= 512ULL * 4096); + static_assert(Cam::NATIVE_MAX_ROWS == MAX_CAPACITY); + for (uint64_t capacity : {uint64_t(4096), uint64_t(65536), uint64_t(MAX_CAPACITY)}) { + const Cam::NativeMegaWorkspace large(capacity, 8192, 8192, 128, 4); + assert(large.a2 - large.a1 >= capacity * 8192); + assert(large.c1 - large.a2 >= capacity * 8192); + assert(large.c2 - large.c1 >= 8 * capacity * 8192); + assert(large.output - large.c2 >= 4 * capacity * 8192); + assert(large.bytes - large.output == 2 * capacity * 8192); + assert(large.output % 512 == 0); + assert(Rounds(capacity, capacity) == 1); + assert(Rounds(capacity + 1, capacity) == 2); + } + const Cam::NativeMegaWorkspace model(262144, 4096, 2048, 64, 2); + assert(model.bytes - model.output == 2ULL * 1024 * 1024 * 1024); + std::mt19937 random(20261027); + for (uint32_t trial = 0; trial < 1000; ++trial) { + const uint32_t sources = 1 + random() % 4, experts = 1 + random() % 8; + const uint64_t capacity = 1 + random() % 23; + std::vector counts(sources * experts), prefixes(counts.size()); + uint64_t rows = 0; + for (uint32_t e = 0; e < experts; ++e) { + for (uint32_t s = 0; s < sources; ++s) { + const auto i = s * experts + e; + prefixes[i] = rows; + counts[i] = random() % 29; + rows += counts[i]; + } + } + std::vector visits(rows); + for (uint64_t q = 0; q < Rounds(rows, capacity); ++q) { + uint64_t roundRows = 0; + for (uint32_t s = 0; s < sources; ++s) { + uint64_t expertPrefix = 0; + for (uint32_t e = 0; e < experts; ++e) { + const auto i = s * experts + e; + const auto slice = Intersect(prefixes[i], counts[i], q * capacity, capacity); + for (uint64_t r = 0; r < slice.count; ++r) { + ++visits[prefixes[i] + slice.skip + r]; + assert(slice.local + r < capacity); + assert(ReturnRow(0, expertPrefix, slice.skip, r) < + expertPrefix + counts[i]); + } + expertPrefix += counts[i]; + roundRows += slice.count; + } + } + assert(roundRows == std::min(capacity, rows - q * capacity)); + } + assert(std::all_of(visits.begin(), visits.end(), [](unsigned n) { return n == 1; })); + } +} diff --git a/tests/npu/standalone_cam_attention.py b/tests/npu/standalone_cam_attention.py index 41ded35a2..38bcd63e1 100644 --- a/tests/npu/standalone_cam_attention.py +++ b/tests/npu/standalone_cam_attention.py @@ -4,7 +4,7 @@ This finite functional check uses the original checkpoint on F and every main MoE layer. It has no Attention model, KV cache, HTTP inference or performance -claim. Run an eager F service first with --output, then FULL/two-stream F with +claim. Run an eager F service first with --output, then FULL F with --reference pointing to that output. Terminate the whole paired deployment after this client: the production service has no EOS or cancellable DR. """ @@ -31,15 +31,33 @@ AFDTransferMetadata, ) from afd_plugin.connectors.npu.async_cam import CAMAsyncAFDConnector # noqa: E402 -from afd_plugin.ffn_faas.platforms.npu.settings import FFNConfig # noqa: E402 -from afd_plugin.ffn_faas.platforms.npu.v4_model import ( # noqa: E402 +from afd_plugin.ffn_faas.platforms.ffn_model.deepseek_v4.npu.v4_model import ( # noqa: E402 V4CheckpointDescription, ) -def attention_config(checkpoint, settings): - additional = settings.additional_config() - additional["afd"]["role"] = "attention" +def attention_config(checkpoint, args): + additional = { + "afd": { + "connector": "CAMAsyncAFDConnector", + "role": "attention", + "async_dp": True, + "host": "127.0.0.1", + "port": args.port, + "num_attention_ranks": 1, + "num_ffn_ranks": args.ffn_ranks, + "compute_gate_on_attention": True, + "connector_extra_config": { + "dynamicQuant": 1, + "attn_ranks_per_dp": 1, + "hccl_buffer_size": 64, + }, + }, + "enable_prefill_mc2": False, + "enable_force_eplb": False, + "enable_cpu_binding": True, + "ascend_fusion_config": {"fusion_ops_gmmswigluquant": True}, + } return SimpleNamespace( additional_config=additional, model_config=SimpleNamespace( @@ -54,9 +72,7 @@ def attention_config(checkpoint, settings): model_type="deepseek_v4", ), ), - scheduler_config=SimpleNamespace( - max_num_batched_tokens=settings.max_num_batched_tokens - ), + scheduler_config=SimpleNamespace(max_num_batched_tokens=32), parallel_config=SimpleNamespace( enable_dbo=False, use_ubatching=False, @@ -89,17 +105,7 @@ def main(): ) args = parser.parse_args() checkpoint = V4CheckpointDescription.read(args.model) - settings = FFNConfig( - args.model, - attention_ranks=1, - attention_tp=1, - ffn_ranks=args.ffn_ranks, - port=args.port, - max_num_batched_tokens=32, - max_num_seqs=8, - hccl_buffer_size=64, - ) - config = attention_config(checkpoint, settings) + config = attention_config(checkpoint, args) args.output.parent.mkdir(parents=True, exist_ok=True) ensure_cam_async_ops_available() torch.npu.set_device(args.device) diff --git a/tests/npu/test_async_cam_layered_w8a8.py b/tests/npu/test_async_cam_layered_w8a8.py index 6534fe988..a5cc30805 100644 --- a/tests/npu/test_async_cam_layered_w8a8.py +++ b/tests/npu/test_async_cam_layered_w8a8.py @@ -20,8 +20,8 @@ from torch import Tensor -@pytest.fixture -def w8a8_runtime(): +@pytest.fixture(params=[0.0, 10.0]) +def w8a8_runtime(request): if os.environ.get("AFD_RUN_ASCEND_OP_RUNTIME") != "1": pytest.skip("requires opt-in 910C runtime") torch = pytest.importorskip("torch") @@ -70,16 +70,17 @@ def w8a8_runtime(): w13_bias=None, w2_bias=None, per_channel=True, - swiglu_limit=0.0, + swiglu_limit=request.param, routed_scaling_factor=1.0, quantization="W8A8", ) ) executor = AsyncCAMLayeredExecutor(layers) + assert executor.applied_swiglu_limit == 0.0 x = torch.randint( -8, 8, (capacity, hidden), dtype=torch.int8, generator=generator ).npu() - scales = torch.full((capacity,), 0.01, device="npu") + scales = torch.full((capacity,), 0.01 if request.param == 0 else 2.0, device="npu") return torch, torch_npu, layers, executor, x, scales diff --git a/tests/npu/test_cam_async_ops.py b/tests/npu/test_cam_async_ops.py index e4ef456f5..cb150892c 100644 --- a/tests/npu/test_cam_async_ops.py +++ b/tests/npu/test_cam_async_ops.py @@ -43,6 +43,8 @@ def _dispatch_recv(torch: ModuleType, *, dynamic_quant: int = 0, moe_rank_id: in "afd_async_dispatch_recv", "afd_async_combine_send", "afd_async_combine_recv", + "afd_mega_dispatch_send", + "afd_mega_combine_recv", ), ) def test_cam_dispatcher_registers_inference_and_meta_kernels( @@ -56,6 +58,137 @@ def test_cam_dispatcher_registers_inference_and_meta_kernels( assert all("shared" not in argument.name for argument in schema.arguments) +def _native_mega_inputs(torch, mode): + packed = mode == 1 + + def make(shape, dtype): + return torch.empty(shape, dtype=dtype, device="meta") + + anchor = make((1,), torch.bfloat16) + size = torch.ops.afd_ascend.afd_mega_moe_pull_workspace_size(16, 256, 256, 2, 1) + return [ + anchor, + make((4096,), torch.float16), + make((6,), torch.int64), + make((size,), torch.uint8), + [make((2, 256, 64 if packed else 512), torch.int32 if packed else torch.int8)], + [make((2, 256, 32 if packed else 256), torch.int32 if packed else torch.int8)], + [make((2, 512), torch.int64 if packed else torch.float32)], + [make((2, 256), torch.int64 if packed else torch.bfloat16)], + [make((2, 512), torch.float32)] if packed else [], + [make((2, 256), torch.float32)] if packed else [], + [8, 256, 256, 2, 2, 1, 2, 1, 3, 1, 16, mode, 1, 1], + "meta_only", + 1.0, + ] + + +@pytest.mark.parametrize("mode", [1, 2]) +def test_native_mega_meta_quantization_contract(cam_runtime, mode): + torch = cam_runtime + completion = torch.ops.afd_ascend.afd_mega_moe_pull_native( + *_native_mega_inputs(torch, mode), swiglu_limit=10.0 + ) + assert completion.shape == (1,) and completion.dtype == torch.int64 + assert completion.device.type == "meta" + + +@pytest.mark.parametrize( + "invalid", ["mode", "weight", "scale", "bias", "shape", "layers"] +) +def test_native_mega_meta_rejects_invalid_w8a8(cam_runtime, invalid): + torch = cam_runtime + inputs = _native_mega_inputs(torch, 2) + if invalid == "mode": + inputs[10][11] = 3 + elif invalid == "weight": + inputs[4][0] = inputs[4][0].to(torch.int32) + elif invalid == "scale": + inputs[7][0] = inputs[7][0].to(torch.float32) + elif invalid == "bias": + inputs[9] = [torch.empty((2, 256), dtype=torch.int32, device="meta")] + elif invalid == "shape": + inputs[4][0] = inputs[4][0].reshape(2, 512, 256) + else: + inputs[10][13] = 2 + with pytest.raises(RuntimeError, match="native Mega"): + torch.ops.afd_ascend.afd_mega_moe_pull_native(*inputs) + + +def test_mega_dispatch_meta_publishes_completion(cam_runtime): + torch = cam_runtime + state = torch.ops.afd_ascend.afd_mega_pull_init( + torch.empty(1, dtype=torch.bfloat16, device="meta"), + [8, 256, 2, 2, 1, 4, 0, 1], + [], + 1, + "meta_only", + ) + sent = torch.ops.afd_ascend.afd_mega_dispatch_send( + torch.empty((3, 256), dtype=torch.bfloat16, device="meta"), + torch.empty((3, 2), dtype=torch.int32, device="meta"), + state, + 0x41464401, + 8, + 3, + 256, + 2, + 2, + 1, + 4, + 0, + 3, + 2, + 1, + 1, + "meta_only", + ) + assert sent.shape == (1,) and sent.dtype == torch.int64 + schema = torch.ops.afd_ascend.afd_mega_dispatch_send.default._schema + comm_arg = next(arg for arg in schema.arguments if arg.name == "comm_args") + assert comm_arg.alias_info.is_write + + +def test_pull_empty_combine_preserves_logical_shape(cam_runtime): + torch = cam_runtime + anchor = torch.empty(1, dtype=torch.bfloat16, device="meta") + state = torch.ops.afd_ascend.afd_mega_pull_init( + anchor, [8, 256, 2, 2, 1, 4, 0, 1], [], 1, "meta_only" + ) + output = torch.ops.afd_ascend.afd_mega_combine_recv( + anchor, + torch.empty((0, 2), dtype=torch.int32, device="meta"), + torch.empty((0, 2), dtype=torch.float32, device="meta"), + state, + 0x41464401, + 0, + 256, + 2, + 2, + 1, + 4, + 0, + 3, + "meta_only", + 8, + 1, + ) + assert output.shape == (0, 256) and output.dtype == torch.bfloat16 + + +@pytest.mark.parametrize("layers", [(2, 2), (-1,), (65536,), (2**63 - 1,)]) +def test_pull_init_rejects_invalid_layer_map(cam_runtime, layers): + torch = cam_runtime + with pytest.raises(RuntimeError, match="layer"): + torch.ops.afd_ascend.afd_mega_pull_init( + torch.empty(1, dtype=torch.bfloat16, device="meta"), + [8, 256, 2, 2, 1, 4, 1, 1], + list(layers), + 1, + "meta_only", + ) + + @pytest.mark.parametrize("dtype_name", ["float16", "bfloat16"]) @pytest.mark.parametrize("dynamic_quant", [0, 1]) def test_cam_four_phase_meta_contracts( diff --git a/tests/unit/compat/npu/test_ascend_config_lifecycle.py b/tests/unit/compat/npu/test_ascend_config_lifecycle.py index 756eb8187..5058a58ed 100644 --- a/tests/unit/compat/npu/test_ascend_config_lifecycle.py +++ b/tests/unit/compat/npu/test_ascend_config_lifecycle.py @@ -29,7 +29,7 @@ def load_functions(path, namespace, names): - tree = ast.parse(path.read_text()) + tree = ast.parse(path.read_text(encoding="utf-8")) tree.body = [ node for node in tree.body @@ -294,7 +294,7 @@ def test_async_shared_replication_only_overrides_effective_model_sp( self.assertEqual(vars(config), before) def test_factory_binds_already_loaded_native_scheduler_callers(self): - tree = ast.parse(FACTORY_PATH.read_text()) + tree = ast.parse(FACTORY_PATH.read_text(encoding="utf-8")) aliases = next( ast.literal_eval(node.value) for node in tree.body diff --git a/tests/unit/compat/npu/test_feature_validation.py b/tests/unit/compat/npu/test_feature_validation.py index 8b4f702f2..cda824302 100644 --- a/tests/unit/compat/npu/test_feature_validation.py +++ b/tests/unit/compat/npu/test_feature_validation.py @@ -19,7 +19,10 @@ from tests.unit.fakes.vllm_config import ( # noqa: E402 async_moe_config as _async_moe_config, ) -from tests.unit.fakes.vllm_config import npu_dspark_config # noqa: E402 +from tests.unit.fakes.vllm_config import ( # noqa: E402 + npu_async_cam_target_graph_config, + npu_dspark_config, +) from tests.unit.fakes.vllm_config import ( # noqa: E402 npu_vllm_config as _vllm_config, ) @@ -35,7 +38,7 @@ [ ("attention", True, True, "deepseek_v4", None), ("ffn", True, True, "deepseek_v4", None), - ("attention", False, True, "deepseek_v4", "DSpark requires eager"), + ("attention", False, True, "deepseek_v4", "draft Graph requires target"), ( "attention", True, @@ -43,13 +46,19 @@ "deepseek_v2", "requires DeepSeek V4", ), + ("ffn", False, True, "deepseek_v4", None), ("ffn", True, False, "deepseek_v4", None), - ("attention", True, False, "deepseek_v4", "eager Attention"), + ("ffn", False, False, "deepseek_v4", None), + ("attention", True, False, "deepseek_v4", None), + ("attention", False, False, "deepseek_v4", None), ], ) def test_dspark_feature_boundaries(role, draft_eager, target_eager, model_type, error): config = npu_dspark_config(role=role, draft_eager=draft_eager) config.model_config.enforce_eager = target_eager + if role == "attention" and not target_eager: + config.compilation_config.mode = 0 + config.compilation_config.cudagraph_mode.name = "FULL_DECODE_ONLY" config.model_config.hf_text_config.model_type = model_type if error: with pytest.raises(RuntimeError, match=error): @@ -58,6 +67,16 @@ def test_dspark_feature_boundaries(role, draft_eager, target_eager, model_type, fail_if_unsupported_npu_afd_features(config) +@pytest.mark.parametrize("draft_eager", [True, False]) +def test_dspark_target_graph_reuses_tp4(monkeypatch, draft_eager): + monkeypatch.setenv("VLLM_ASCEND_ENABLE_FLASHCOMM1", "0") + config = npu_dspark_config(tp_size=4, draft_eager=draft_eager) + config.model_config.enforce_eager = False + config.compilation_config.mode = 0 + config.compilation_config.cudagraph_mode.name = "FULL_DECODE_ONLY" + fail_if_unsupported_npu_afd_features(config) + + @pytest.mark.parametrize( "mode", ["FULL_DECODE_ONLY", "PIECEWISE", "FULL_AND_PIECEWISE"] ) @@ -69,15 +88,6 @@ def test_dspark_ffn_graph_requires_full_mode(mode): fail_if_unsupported_npu_afd_features(config) -@pytest.mark.parametrize("eager", [True, False]) -def test_dspark_rejects_ffn_two_stream(eager): - config = npu_dspark_config(role="ffn") - config.model_config.enforce_eager = eager - config.additional_config["afd"]["connector_extra_config"]["ffn_two_stream"] = True - with pytest.raises(RuntimeError, match="DSpark.*two-stream"): - fail_if_unsupported_npu_afd_features(config) - - @pytest.mark.parametrize( "anchor,num_requests,capacity,allowed", [ @@ -182,7 +192,7 @@ def _mla_dbo(cudagraph_mode, **model_fields): ), pytest.param( _config_with(model={"enforce_eager": False}, **ASYNC_CAM), - "only eager", + "FULL_DECODE_ONLY", id="async-graph", ), pytest.param( @@ -286,27 +296,66 @@ def test_async_cam_rejects_models_without_layered_ffn(role, model_type): fail_if_unsupported_npu_afd_features(config) +@pytest.mark.parametrize("dp_size", [1, 2]) +@pytest.mark.parametrize("tp_size", [1, 2]) +def test_async_cam_target_graph_feature_admission(monkeypatch, dp_size, tp_size): + monkeypatch.setenv("VLLM_ASCEND_ENABLE_FLASHCOMM1", "1") + config = npu_async_cam_target_graph_config( + data_parallel_size=dp_size, + tensor_parallel_size=tp_size, + num_attention_ranks=dp_size * tp_size, + ) + config.kv_transfer_config = object() + config.parallel_config.use_sequence_parallel_moe = True + fail_if_unsupported_npu_afd_features(config) + + +@pytest.mark.parametrize( + "method,uses_draft,error", + [ + ("mtp", True, "native draft models or MTP"), + ("ngram", False, None), + ], +) +def test_async_cam_target_graph_feature_preserves_draft_model_boundary( + monkeypatch, method, uses_draft, error +): + monkeypatch.setenv("VLLM_ASCEND_ENABLE_FLASHCOMM1", "0") + config = npu_async_cam_target_graph_config() + config.speculative_config = SimpleNamespace( + method=method, + use_eagle=lambda: False, + uses_draft_model=lambda: uses_draft, + ) + with pytest.raises(RuntimeError, match=error) if error else nullcontext(): + fail_if_unsupported_npu_afd_features(config) + + def test_npu_afd_num_ubatches_reads_two_ubatch_dbo(): assert npu_afd_num_ubatches(_vllm_config(**TWO_UBATCH_DBO)) == 2 @pytest.mark.parametrize("role", ["attention", "ffn"]) @pytest.mark.parametrize( - ("factor", "max_num_batched_tokens", "rejected"), + ("factor", "max_num_batched_tokens", "mega", "rejected"), [ - ("0.015625", 8192, True), - ("0.015625", 4096, False), - ("0.5", 8192, False), - (None, 262144, False), - ("0", 16, True), - ("nan", 16, True), - ("invalid", 16, True), - ("0.01_5625", 4096, True), + ("0.015625", 8192, False, True), + ("0.015625", 8192, True, False), + ("0.015625", 4096, False, False), + ("0.5", 8192, False, False), + (None, 262144, False, False), + ("0", 16, False, True), + ("nan", 16, False, True), + ("invalid", 16, False, True), + ("0.01_5625", 4096, False, True), + ("0.00000001", 16, False, True), + ("0.00000001", 16, True, True), ], ) def test_npu_async_feature_validation_checks_receive_capacity( monkeypatch, role, + mega, factor, max_num_batched_tokens, rejected, @@ -324,7 +373,11 @@ def test_npu_async_feature_validation_checks_receive_capacity( use_v2_model_runner=True, num_attention_ranks=2, tensor_parallel_size=2 if role == "attention" else 1, - extra_config={"attn_ranks_per_dp": 2, "dynamicQuant": 1}, + extra_config={ + "attn_ranks_per_dp": 2, + "ffn_mega_moe": mega, + "dynamicQuant": 1, + }, ) config.model_config.hf_text_config.model_type = "deepseek_v4" config.scheduler_config.max_num_batched_tokens = max_num_batched_tokens diff --git a/tests/unit/compat/npu/test_sync_gate_feature_validation_cpu.py b/tests/unit/compat/npu/test_sync_gate_feature_validation_cpu.py new file mode 100644 index 000000000..aec7e0ecd --- /dev/null +++ b/tests/unit/compat/npu/test_sync_gate_feature_validation_cpu.py @@ -0,0 +1,159 @@ +# mypy: disable-error-code="attr-defined, name-defined" +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""CPU coverage for the sync gate-on-Attention feature contract. + +The public validator is the seam every AFD NPU run passes before execution, so it is +driven directly rather than through its private helpers. Only the two upstream modules +the validator reaches for -- the connector factory and the connector's extra-info type +-- are stand-ins; the validator itself is the real one. + +Keeping the assertions on the public entry matters: it fails both when the helper +starts rejecting supported shapes and when the entry stops calling the helper. +""" + +from __future__ import annotations + +import sys +import types +from types import SimpleNamespace + +import pytest + +from afd_plugin.compat.npu import feature_validation +from afd_plugin.config import CAMP2P_CONNECTOR + +_DSV4 = "deepseek_v4" + + +class _FakeExtraInfo: + """Stands in for the connector extra info the factory would build.""" + + def __init__(self) -> None: + self.validated = 0 + + def validate_supported(self) -> None: + self.validated += 1 + + +@pytest.fixture +def extra_info(monkeypatch): + info = _FakeExtraInfo() + factory = types.ModuleType("afd_plugin.connectors.factory") + factory.AFDConnectorFactory = SimpleNamespace( + parse_connector_extra_info=lambda connector, config: info, + ) + camp2p = types.ModuleType("afd_plugin.connectors.npu.camp2p") + camp2p.CAMP2PExtraInfo = _FakeExtraInfo + monkeypatch.setitem(sys.modules, "afd_plugin.connectors.factory", factory) + monkeypatch.setitem(sys.modules, "afd_plugin.connectors.npu.camp2p", camp2p) + return info + + +def _afd_config(*, role: str = "attention"): + return SimpleNamespace( + role=role, + connector=CAMP2P_CONNECTOR, + compute_gate_on_attention=True, + ) + + +def _vllm_config(*, tp: int = 1, dbo: bool = False, ubatching: bool = False): + hf_config = SimpleNamespace(model_type=_DSV4) + return SimpleNamespace( + model_config=SimpleNamespace( + hf_config=hf_config, + hf_text_config=SimpleNamespace(model_type=_DSV4), + ), + parallel_config=SimpleNamespace( + enable_dbo=dbo, + use_ubatching=ubatching, + num_ubatches=1, + tensor_parallel_size=tp, + ), + compilation_config=SimpleNamespace(cudagraph_mode=None), + speculative_config=None, + ) + + +@pytest.mark.parametrize("role", ["attention", "ffn"]) +def test_the_supported_shape_is_accepted_for_both_roles(extra_info, role): + """Gate-on-Attention at TP=1 on the sync connector must pass validation.""" + + feature_validation.fail_if_unsupported_npu_afd_features( + _vllm_config(), + afd_config=_afd_config(role=role), + ) + + assert extra_info.validated == 1 + + +@pytest.mark.parametrize( + ("kwargs", "match"), + [ + ({"dbo": True}, "DBO or ubatching"), + ({"ubatching": True}, "DBO or ubatching"), + ({"tp": 2}, "tensor-parallel size 1"), + ], +) +@pytest.mark.parametrize("role", ["attention", "ffn"]) +def test_unsupported_shapes_are_rejected(extra_info, kwargs, match, role): + with pytest.raises(RuntimeError, match=match): + feature_validation.fail_if_unsupported_npu_afd_features( + _vllm_config(**kwargs), + afd_config=_afd_config(role=role), + ) + + +def test_gate_on_ffn_keeps_its_previous_behaviour(extra_info): + """Turning the Attention-side gate off must not pick up the new restrictions.""" + + afd_config = SimpleNamespace( + role="ffn", + connector=CAMP2P_CONNECTOR, + compute_gate_on_attention=False, + ) + + feature_validation.fail_if_unsupported_npu_afd_features( + _vllm_config(tp=2), + afd_config=afd_config, + ) + + assert extra_info.validated == 1 + + +def _real_extra_info(): + """The connector's own type; importing it needs vLLM. + + ``camp2p`` pulls ``vllm.forward_context`` at module scope, so the real + class is only reachable where vLLM is installed. The validator cases + above use a stand-in and run everywhere; these two cover the class + itself and skip on a bare CPU host. + """ + + pytest.importorskip("torch") + pytest.importorskip("vllm") + from afd_plugin.connectors.npu.camp2p import CAMP2PExtraInfo + + return CAMP2PExtraInfo + + +def test_a_connector_level_gate_flag_does_not_block(): + """The switch is read from ``afd``; the connector-local copy has no say. + + Regression test for a rejection that used to live in ``validate_supported``: a + supported + shape was refused at startup whenever the flag also appeared under + ``connector_extra_config``. + """ + + _real_extra_info().from_mapping( + {"compute_gate_on_attention": True} + ).validate_supported() + + +def test_quant_mode_is_still_rejected(): + """The transport mode really is unsupported, so that check has to stay.""" + + with pytest.raises(RuntimeError, match="only quant_mode=0"): + _real_extra_info().from_mapping({"quant_mode": 4}).validate_supported() diff --git a/tests/unit/compat/patches/npu/test_mrv2_ubatching.py b/tests/unit/compat/patches/npu/test_mrv2_ubatching.py index ba812ef26..5fa429137 100644 --- a/tests/unit/compat/patches/npu/test_mrv2_ubatching.py +++ b/tests/unit/compat/patches/npu/test_mrv2_ubatching.py @@ -205,6 +205,7 @@ def dp_scope(): skip_dp.reset(token) class NativeRunner: + connector: object vllm_config: SimpleNamespace afd_config: SimpleNamespace prof: None @@ -256,6 +257,7 @@ def execute_model( future = ast.parse("from __future__ import annotations").body namespace = { "NativeRunner": NativeRunner, + "CAMAsyncAFDConnector": type("CAMAsyncAFDConnector", (), {}), "CUDAGraphMode": _GraphMode, "nullcontext": nullcontext, "step_afd_npu_profiler": lambda profiler: None, @@ -279,6 +281,7 @@ def make_runner(async_dp): Callable[[], NativeRunner], namespace["AFDNPUAttentionModelRunnerV2"] ) runner = runner_factory() + runner.connector = None runner.vllm_config = SimpleNamespace( compilation_config=SimpleNamespace(cudagraph_mode=_GraphMode.NONE) ) @@ -326,6 +329,35 @@ def test_execute_preserves_native_input_policy_and_restores_scopes( assert execute_contract.skip_dp.get() is False +@pytest.mark.parametrize("fail", [False, True]) +def test_execute_scopes_connector_step_around_native_execution(execute_contract, fail): + runner = execute_contract.make_runner(False) + events = [] + + @contextmanager + def attention_step(): + events.append("enter") + try: + yield + finally: + events.append("exit") + + runner.connector = SimpleNamespace(attention_step=attention_step) + + def install_metadata(context): + assert events == ["enter"] + events.append("native") + + runner.install_afd_metadata_on_forward_context = install_metadata + scheduler = SimpleNamespace(total_num_scheduled_tokens=7, fail=fail) + if fail: + with pytest.raises(RuntimeError, match="native execution failed"): + runner.execute_model(scheduler) + else: + runner.execute_model(scheduler) + assert events == ["enter", "native", "exit"] + + def test_execute_keeps_native_default_input_policy(execute_contract): runner = execute_contract.make_runner(False) result = runner.execute_model( diff --git a/tests/unit/compat/patches/test_signatures.py b/tests/unit/compat/patches/test_signatures.py index 0b7bd68c4..d9ebad4bb 100644 --- a/tests/unit/compat/patches/test_signatures.py +++ b/tests/unit/compat/patches/test_signatures.py @@ -237,6 +237,22 @@ def _use_installed_patch(binding, symbol, marker_name, install): "fullgraph-replay", id="fullgraph-replay", ), + pytest.param( + "vllm.v1.worker.gpu.spec_decode.dflash.cudagraph", + None, + "_prepare_dflash_inputs_to_capture", + None, + "npu-draft-capture", + id="npu-draft-capture", + ), + pytest.param( + "vllm_ascend.worker.v2.spec_decode.dflash.aclgraph", + "DFlashAclGraphManager", + "run_fullgraph", + None, + "npu-draft-replay", + id="npu-draft-replay", + ), pytest.param( "vllm_ascend.ascend_config", None, @@ -352,6 +368,26 @@ def test_patch_preserves_upstream_signature( tracker = adapter._AFDCaptureEventTracker([]) scope = adapter._use_afd_capture_input_preparation(None, tracker) + elif scope_name in {"npu-draft-capture", "npu-draft-replay"}: + from vllm.config import CUDAGraphMode + + from afd_plugin.v1.worker.npu import dspark_graph as adapter + + if scope_name == "npu-draft-capture": + runner = SimpleNamespace( + speculative_config=SimpleNamespace(enforce_eager=True), + _afd_is_graph_capturing=False, + _afd_is_graph_replaying=False, + ) + manager = SimpleNamespace( + _capture_descs={CUDAGraphMode.FULL: []}, + _max_full_descs_to_capture=None, + ) + scope = adapter.use_afd_draft_capture(runner, manager) + else: + native_class = getattr(target, owner) + binding = native_class.__new__(native_class) + scope = adapter.use_afd_draft_replay(SimpleNamespace(), binding, 1) elif scope_name == "camp2p-dispatch": from afd_plugin.v1.worker.npu import attention_model_runner_v2 as adapter diff --git a/tests/unit/config/test_config.py b/tests/unit/config/test_config.py index 263e3bca7..0a303745c 100644 --- a/tests/unit/config/test_config.py +++ b/tests/unit/config/test_config.py @@ -27,17 +27,17 @@ def test_parse_empty_additional_config_is_inactive(): def test_connector_extra_fields_participate_in_the_graph_hash(): - extra = {"ffn_two_stream": True, "ffn_compute_aic_num": 12} + extra = {"ffn_mega_moe": True, "dynamicQuant": 1} raw = { "connector": "CAMAsyncAFDConnector", "role": "ffn", "connector_extra_config": extra, } first = afd_config_from_mapping(raw) - extra["ffn_compute_aic_num"] = 10 + extra["ffn_mega_moe"] = False second = afd_config_from_mapping(raw) assert first.compute_hash() != second.compute_hash() - raw["connector_extra_config"] = {"ffn_compute_aic_num": 10, "ffn_two_stream": True} + raw["connector_extra_config"] = {"dynamicQuant": 1, "ffn_mega_moe": False} assert second.compute_hash() == afd_config_from_mapping(raw).compute_hash() diff --git a/tests/unit/config/test_validation.py b/tests/unit/config/test_validation.py index 6d8a53e0c..f61a7f7df 100644 --- a/tests/unit/config/test_validation.py +++ b/tests/unit/config/test_validation.py @@ -18,7 +18,11 @@ validate_gpu_model_runner_v2_config, validate_npu_model_runner_v2_config, ) -from tests.unit.fakes.vllm_config import mrv2_vllm_config +from tests.unit.fakes.vllm_config import ( + mrv2_vllm_config, + npu_async_cam_target_graph_config, + npu_dspark_config, +) MODEL_UTILS_MODULE = "afd_plugin.model_executor.models.model_utils" @@ -156,6 +160,97 @@ def test_gpu_validator_accepts_target_config_without_removed_enable_sp(mrv2_conf ) +@pytest.mark.parametrize("dp_size", [1, 2]) +@pytest.mark.parametrize("tp_size", [1, 2]) +def test_async_cam_target_graph_accepts_independent_attention_dps( + monkeypatch, mrv2_config, dp_size, tp_size +): + monkeypatch.setenv("VLLM_ASCEND_ENABLE_FLASHCOMM1", "0") + config = npu_async_cam_target_graph_config( + data_parallel_size=dp_size, + tensor_parallel_size=tp_size, + num_attention_ranks=dp_size * tp_size, + ) + validate_npu_model_runner_v2_config( + config, expected_role="attention", device_type="npu" + ) + + +@pytest.mark.parametrize("draft_eager", [True, False]) +def test_async_cam_target_graph_accepts_dspark_tp4( + monkeypatch, mrv2_config, draft_eager +): + monkeypatch.setenv("VLLM_ASCEND_ENABLE_FLASHCOMM1", "0") + config = npu_async_cam_target_graph_config( + tensor_parallel_size=4, num_attention_ranks=4 + ) + config.speculative_config = npu_dspark_config( + tp_size=4, draft_eager=draft_eager + ).speculative_config + validate_npu_model_runner_v2_config( + config, expected_role="attention", device_type="npu" + ) + + +@pytest.mark.parametrize( + "path,value,error", + [ + ("use_v2_model_runner", False, "ModelRunnerV2"), + ("model_config.hf_text_config.model_type", "deepseek_v2", "DeepSeek-V4"), + ("parallel_config.enable_dbo", True, "DBO/ubatching"), + ("parallel_config.use_ubatching", True, "DBO/ubatching"), + ("parallel_config.pipeline_parallel_size", 2, "PP or CP"), + ("parallel_config.prefill_context_parallel_size", 2, "PP or CP"), + ("parallel_config.decode_context_parallel_size", 2, "PP or CP"), + ("compilation_config.cudagraph_mode.name", "NONE", "supports ACL graph modes"), + ( + "compilation_config.cudagraph_mode.name", + "PIECEWISE", + "supports ACL graph modes", + ), + ], +) +def test_async_cam_target_graph_rejects_unsupported_configuration( + monkeypatch, mrv2_config, path, value, error +): + monkeypatch.setenv("VLLM_ASCEND_ENABLE_FLASHCOMM1", "0") + config = npu_async_cam_target_graph_config() + *parents, name = path.split(".") + target = config + for parent in parents: + target = getattr(target, parent) + setattr(target, name, value) + with pytest.raises(RuntimeError, match=error): + validate_npu_model_runner_v2_config( + config, expected_role="attention", device_type="npu" + ) + + +@pytest.mark.parametrize( + "path,value", + [ + ("kv_transfer_config", object()), + ("parallel_config.use_sequence_parallel_moe", True), + ("speculative_config", SimpleNamespace(method="ngram")), + ("compilation_config.cudagraph_mode.name", "FULL"), + ("compilation_config.mode", 3), + ], +) +def test_async_cam_target_graph_preserves_native_options( + monkeypatch, mrv2_config, path, value +): + monkeypatch.setenv("VLLM_ASCEND_ENABLE_FLASHCOMM1", "1") + config = npu_async_cam_target_graph_config() + *parents, name = path.split(".") + target = config + for parent in parents: + target = getattr(target, parent) + setattr(target, name, value) + validate_npu_model_runner_v2_config( + config, expected_role="attention", device_type="npu" + ) + + @pytest.mark.parametrize( "role,tp,error", [ @@ -655,8 +750,8 @@ def test_gpu_v2_validator_outcome(mrv2_config, monkeypatch, overrides, error): ), pytest.param( {**_NPU_ASYNC, "model.enforce_eager": False}, - _err("eager"), - id="reject-async-cam-graph", + _err("MRV2 DeepSeek-V4"), + id="reject-async-cam-non-dsv4-graph", ), pytest.param( {**_NPU_ASYNC, "par.pipeline_parallel_size": 2}, @@ -697,6 +792,7 @@ def test_gpu_v2_validator_outcome(mrv2_config, monkeypatch, overrides, error): ) def test_npu_v2_validator_outcome(mrv2_config, overrides, error): config, call = _validator_call(overrides) + config.model_config.hf_text_config = config.model_config.hf_config with _expect(error): validate_npu_model_runner_v2_config( config, diff --git a/tests/unit/connectors/npu/test_async_cam.py b/tests/unit/connectors/npu/test_async_cam.py new file mode 100644 index 000000000..ec7f1dbd7 --- /dev/null +++ b/tests/unit/connectors/npu/test_async_cam.py @@ -0,0 +1,317 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Native CAM message contracts and ownership through the public connector.""" + +import sys +import weakref +from types import SimpleNamespace + +import pytest + +from tests.unit.fakes.async_cam import ( + _FakeTensor, + _topk_payload, + recording_cam_connector, +) +from tests.unit.fakes.async_cam_startup import MemoryStore, startup_world +from tests.unit.fakes.async_ffn import load_async_modules + + +@pytest.fixture +def cam(monkeypatch): + _, connector, _ = load_async_modules(monkeypatch) + return SimpleNamespace( + create=lambda **kwargs: recording_cam_connector( + monkeypatch, connector, **kwargs + ), + connector_module=connector, + metadata=sys.modules["afd_plugin.connectors.metadata"], + ) + + +def test_async_connector_calls_cam_shaped_ops(cam): + h = cam.create(pcp_size=3) + connector, fake_torch = h.connector, h.torch + hidden_states = _FakeTensor((3, 16)) + metadata = cam.metadata.AFDTransferMetadata.create_attention_metadata( + layer_idx=2, + stage_idx=0, + seq_len=3, + ) + context = cam.metadata.AFDTransferContext(metadata=metadata) + + output = connector.send_attn_output( + hidden_states, + context, + **_topk_payload(3), + ) + combined = connector.recv_ffn_output( + ref_tensor=hidden_states, + ubatch_idx=0, + ) + + assert output is None + assert combined.shape == (3, 16) + assert fake_torch.ops.afd_ascend.calls[0][0] == "dispatch_send" + assert fake_torch.ops.afd_ascend.calls[1][0] == "combine_recv" + assert fake_torch.ops.afd_ascend.calls[0][1][3] == cam.connector_module.CAM_COMM_ID + assert fake_torch.ops.afd_ascend.calls[1][1][4] == cam.connector_module.CAM_COMM_ID + assert fake_torch.ops.afd_ascend.calls[0][1][5:11] == (3, 16, 2, 2, 4, 4) + assert fake_torch.ops.afd_ascend.calls[1][1][5:11] == (3, 16, 2, 2, 4, 4) + assert fake_torch.ops.afd_ascend.calls[0][1][14] == 4 + assert isinstance(context.states, cam.connector_module.AFDAsyncTransferState) + assert isinstance(context.states, cam.metadata.AFDTransferState) + + +def test_async_explicit_receive_consumes_matching_pending_dispatch(cam): + h = cam.create() + connector, fake_torch = h.connector, h.torch + hidden_states = _FakeTensor((3, 16)) + contexts = [] + payloads = [] + for layer_idx in (2, 3): + context = cam.metadata.AFDTransferContext( + metadata=cam.metadata.AFDTransferMetadata.create_attention_metadata( + layer_idx=layer_idx, stage_idx=0, seq_len=3 + ) + ) + payload = _topk_payload(3) + connector.send_attn_output(hidden_states, context, **payload) + contexts.append(context) + payloads.append(payload) + + with pytest.raises(RuntimeError, match="does not match the next pending"): + connector.recv_ffn_output(hidden_states, ubatch_idx=0, context=contexts[1]) + assert len(fake_torch.ops.afd_ascend.calls) == 2 + + connector.recv_ffn_output( + hidden_states, + ubatch_idx=0, + context=contexts[0], + **payloads[1], + ) + assert fake_torch.ops.afd_ascend.calls[-1][1][1] is payloads[0]["topk_ids"] + + connector.recv_ffn_output(hidden_states, ubatch_idx=0) + assert fake_torch.ops.afd_ascend.calls[-1][1][1] is payloads[1]["topk_ids"] + + +@pytest.mark.parametrize("mega", [False, True]) +def test_protocol_owns_stage_queue_or_exclusive_window(cam, mega): + h = cam.create(mega=mega) + hidden = _FakeTensor((3, 16)) + contexts = [ + cam.metadata.AFDTransferContext( + cam.metadata.AFDTransferMetadata.create_attention_metadata( + layer_idx=2, stage_idx=stage, seq_len=3 + ) + ) + for stage in (0, 1) + ] + h.connector.send_attn_output(hidden, contexts[0], **_topk_payload(3)) + if mega: + with pytest.raises(RuntimeError, match="Pull window is pending"): + h.connector.send_attn_output(hidden, contexts[1], **_topk_payload(3)) + assert len(h.torch.ops.afd_ascend.calls) == 1 + h.connector.recv_ffn_output(hidden, ubatch_idx=0, context=contexts[0]) + h.connector.send_attn_output(hidden, contexts[1], **_topk_payload(3)) + else: + h.connector.send_attn_output(hidden, contexts[1], **_topk_payload(3)) + h.connector.recv_ffn_output(hidden, ubatch_idx=1, context=contexts[1]) + if mega: + assert h.torch.ops.afd_ascend.calls[-1][0] == "mega_combine_recv" + assert h.torch.ops.afd_ascend.calls[-1][1][4] == 0x41464401 + assert h.torch.ops.afd_ascend.calls[-1][1][-2:] == (16, 4) + else: + assert h.connector.has_pending_attention + h.connector.recv_ffn_output(hidden, ubatch_idx=0, context=contexts[0]) + assert not h.connector.has_pending_attention + + +@pytest.mark.parametrize("mega", [False, True]) +def test_failed_combine_retains_original_routing_for_retry(cam, monkeypatch, mega): + h = cam.create(mega=mega) + hidden = _FakeTensor((3, 16)) + context = cam.metadata.AFDTransferContext( + cam.metadata.AFDTransferMetadata.create_attention_metadata( + layer_idx=2, stage_idx=0, seq_len=3 + ) + ) + payload = _topk_payload(3) + h.connector.send_attn_output(hidden, context, **payload) + operation = "afd_mega_combine_recv" if mega else "afd_async_combine_recv" + with monkeypatch.context() as patch: + + def fail(**kwargs): + raise RuntimeError("combine launch failed") + + patch.setattr(h.torch.ops.afd_ascend, operation, fail) + with pytest.raises(RuntimeError, match="combine launch failed"): + h.connector.recv_ffn_output(hidden, context=context) + assert h.connector.has_pending_attention + h.connector.recv_ffn_output(hidden, context=context) + assert h.torch.ops.afd_ascend.calls[-1][1][1] is payload["topk_ids"] + assert not h.connector.has_pending_attention + + +def test_pull_ffn_rejects_separate_receive_and_send(cam): + h = cam.create(mega=True, role="ffn") + with pytest.raises(RuntimeError, match="Mega executor"): + h.connector.recv_attn_output() + with pytest.raises(RuntimeError, match="Mega executor"): + h.connector.send_ffn_output(_FakeTensor((3, 16)), None) + assert not h.torch.ops.afd_ascend.calls + + +def test_async_ffn_side_dispatch_recv_and_combine_send(cam): + h = cam.create(role="ffn", tp_size=2, pcp_size=2) + connector, fake_torch = h.connector, h.torch + recv_output = connector.recv_attn_output(batch_size=4, layer_idx=1) + connector.send_ffn_output(recv_output.hidden_states, recv_output.context) + + states = recv_output.context.states + assert recv_output.hidden_states.shape == (connector.max_num_batched_tokens, 16) + assert states.dynamic_scales.shape == (connector.max_num_batched_tokens,) + assert states.group_list.shape == (connector.topology.expert_per_rank,) + assert fake_torch.ops.afd_ascend.calls[0][0] == "dispatch_recv" + assert fake_torch.ops.afd_ascend.calls[1][0] == "combine_send" + assert fake_torch.ops.afd_ascend.calls[0][1][11] == 2 + assert fake_torch.ops.afd_ascend.calls[1][1][12] == 2 + assert fake_torch.ops.afd_ascend.calls[1][1][2] is states.token_nums_rankid_layeridx + + +def test_async_combine_send_requires_dispatch_recv_token_metadata(cam): + h = cam.create(role="ffn") + connector = h.connector + metadata = cam.metadata.AFDTransferMetadata.create_ffn_metadata( + layer_idx=1, + stage_idx=0, + seq_lens=[4], + ) + # token_nums_rankid_layeridx is left unset (None) so combine send must fail. + states = cam.connector_module.AFDAsyncTransferState( + batch_size=4, + hidden_size=16, + topk=2, + layer_idx=1, + group_list=_FakeTensor((4,), dtype="int64"), + ) + context = cam.metadata.AFDTransferContext(metadata=metadata, states=states) + + with pytest.raises(RuntimeError, match="TokenNums_Rankid_Layeridx"): + connector.send_ffn_output(_FakeTensor((4, 16)), context) + + +@pytest.mark.parametrize("mega", [False, True]) +def test_dispatch_failure_preserves_protocol_window_ownership(cam, monkeypatch, mega): + h = cam.create(mega=mega) + connector, fake_torch = h.connector, h.torch + context = cam.metadata.AFDTransferContext( + metadata=cam.metadata.AFDTransferMetadata.create_attention_metadata( + layer_idx=2, + stage_idx=0, + seq_len=3, + ) + ) + + def fail_dispatch(*_args, **_kwargs): + raise RuntimeError("dispatch failed") + + monkeypatch.setattr( + fake_torch.ops.afd_ascend, + "afd_mega_dispatch_send" if mega else "afd_async_dispatch_send", + fail_dispatch, + ) + with pytest.raises(RuntimeError, match="dispatch failed"): + connector.send_attn_output(_FakeTensor((3, 16)), context, **_topk_payload(3)) + if mega: + with pytest.raises(RuntimeError, match="pending"): + connector.close() + else: + connector.close() + assert not connector.is_initialized + + +def test_close_releases_every_stage_routing(cam): + h = cam.create() + connector, fake_torch = h.connector, h.torch + references: list[weakref.ReferenceType] = [] + for stage in (0, 1): + context = cam.metadata.AFDTransferContext( + metadata=cam.metadata.AFDTransferMetadata.create_attention_metadata( + layer_idx=2, stage_idx=stage, seq_len=3 + ) + ) + payload = _topk_payload(3) + references.extend(weakref.ref(tensor) for tensor in payload.values()) + connector.send_attn_output(_FakeTensor((3, 16)), context, **payload) + del payload, context + fake_torch.ops.afd_ascend.calls.clear() + assert all(reference() is not None for reference in references) + connector.close() + assert all(reference() is None for reference in references) + + +@pytest.mark.parametrize( + "operation", + ["send_attn_output", "recv_ffn_output", "recv_attn_output", "send_ffn_output"], +) +def test_closed_connector_rejects_transfers_before_native_launch(cam, operation): + h = cam.create() + h.connector.close() + hidden = _FakeTensor((1, 16)) + context = cam.metadata.AFDTransferContext( + metadata=cam.metadata.AFDTransferMetadata.create_attention_metadata( + layer_idx=2, stage_idx=0, seq_len=1 + ) + ) + arguments = { + "send_attn_output": (hidden, context), + "recv_ffn_output": (hidden,), + "recv_attn_output": (), + "send_ffn_output": (hidden, context), + } + with pytest.raises(RuntimeError, match="not initialized"): + getattr(h.connector, operation)(*arguments[operation]) + assert h.torch.ops.afd_ascend.calls == [] + + +def test_reinitialization_updates_legacy_window_reader(cam): + h = cam.create() + first_window = h.connector.comm_args + h.connector.close() + assert h.connector.comm_args is None + h.connector.ffn_session().initialize("npu:1", "replacement") + replacement = h.connector.comm_args + assert replacement is not None and replacement is not first_window + assert replacement.device == "npu:1" + with pytest.raises(AttributeError): + h.connector.comm_args = first_window + + +@pytest.mark.parametrize("mega", [False, True]) +def test_window_allocation_failure_reaches_peer_before_ready(cam, monkeypatch, mega): + h = cam.create(mega=mega) + connector = h.connector + connector._initialized = False + connector.ffn_session().release() + backend = SimpleNamespace(get_hccl_comm_name=lambda rank: "window-failure") + group = SimpleNamespace(_get_backend=lambda device: backend) + monkeypatch.setattr( + cam.connector_module, "init_afd_process_group", lambda **kwargs: group + ) + store = MemoryStore() + peer = startup_world(attn=4, ffn=2, store=store)[1] + connector.graph_startup._store = store + connector.graph_startup._prefix = peer._prefix + + def fail_allocation(*args, **kwargs): + raise RuntimeError("window allocation failed") + + monkeypatch.setattr(h.torch, "empty", fail_allocation) + with pytest.raises(RuntimeError, match="window allocation failed"): + connector.init_afd_connector() + with pytest.raises(RuntimeError, match="window allocation failed"): + peer.check_failure() + assert not connector.is_initialized and connector.cam_pg is group + assert connector.comm_args is None diff --git a/tests/unit/connectors/npu/test_async_cam_session.py b/tests/unit/connectors/npu/test_async_cam_session.py new file mode 100644 index 000000000..9be1f7330 --- /dev/null +++ b/tests/unit/connectors/npu/test_async_cam_session.py @@ -0,0 +1,164 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Host session lifecycle contracts; not device execution qualification.""" + +import sys +from types import SimpleNamespace + +import pytest + +from afd_plugin.connectors.npu.async_cam_runtime import ( + async_cam_startup as startup_module, +) +from afd_plugin.connectors.npu.async_cam_runtime.async_cam_session import ( + PullCAMSession, +) +from tests.unit.fakes.async_cam import FakePullState, _FakeTorch +from tests.unit.fakes.async_cam_startup import MemoryStore, startup_world + + +@pytest.fixture +def pull_world(monkeypatch): + monkeypatch.setattr(startup_module, "GRAPH_STARTUP_TIMEOUT_SECONDS", 0) + store = MemoryStore() + startups = startup_world(store=store) + monkeypatch.setitem(sys.modules, "torch", _FakeTorch()) + states = [FakePullState() for _ in startups] + sessions = [] + for startup, state in zip(startups, states, strict=True): + session = PullCAMSession( + startup, + SimpleNamespace( + ffn_size=2, + attn_size=2, + expert_per_rank=2, + world_rank=startup.world_rank, + ), + hidden_size=16, + topk=2, + max_num_batched_tokens=8, + tp_size=1, + receive_capacity=4, + ) + session.comm_args = state + sessions.append(session) + return SimpleNamespace( + sessions=sessions, states=states, startups=startups, store=store + ) + + +@pytest.mark.parametrize("shutdown_requested", [False, True]) +@pytest.mark.parametrize("quiet_ranks", [(), (0,), (1,), (0, 1)]) +def test_keep_serving_requires_every_attention_quiet( + pull_world, quiet_ranks, shutdown_requested +): + for rank in quiet_ranks: + pull_world.store.set(f"ffn_graph/test/pull_quiet/{rank}", "ready") + for receiver in pull_world.sessions[2:]: + assert receiver.keep_serving(shutdown_requested) is (len(quiet_ranks) != 2) + + +def test_peer_failure_wins_over_quiet_shutdown(pull_world): + for rank in range(2): + pull_world.store.set(f"ffn_graph/test/pull_quiet/{rank}", "ready") + pull_world.startups[0].fail(RuntimeError("attention failed")) + with pytest.raises(RuntimeError, match="attention failed"): + pull_world.sessions[2].keep_serving(True) + + +@pytest.mark.parametrize("transactions_to_drain", [0, 1, 4, 5, None]) +def test_drain_requires_source_completion_within_capacity_bound( + pull_world, transactions_to_drain +): + session, words = pull_world.sessions[2], pull_world.states[2].words + remaining = [transactions_to_drain] + words[32] = -1 if remaining[0] == 0 else 0 + + def run_one(): + if remaining[0] is not None: + remaining[0] -= 1 + words[32] = -1 if remaining[0] == 0 else 0 + + if transactions_to_drain in (5, None): + with pytest.raises(RuntimeError, match="retain resources"): + session.drain(run_one) + if transactions_to_drain == 5: + assert remaining[0] == 1 + with pytest.raises(RuntimeError, match="not drained.*retain resources"): + session.finish_receiver() + assert not pull_world.store.check(["ffn_graph/test/pull_stopped/2"]) + else: + session.drain(run_one) + assert remaining[0] == 0 + for rank in range(2): + pull_world.store.set(f"ffn_graph/test/pull_quiet/{rank}", "ready") + session.finish_receiver() + assert pull_world.store.check(["ffn_graph/test/pull_stopped/2"]) + + +def test_failed_drain_does_not_publish_receiver_stopped(pull_world): + session, words = pull_world.sessions[2], pull_world.states[2].words + words[32] = 0 + + def fail_transaction(): + raise RuntimeError("transaction failed") + + with pytest.raises(RuntimeError, match="transaction failed"): + session.drain(fail_transaction) + with pytest.raises(RuntimeError, match="retain resources"): + session.finish_receiver() + assert not pull_world.store.check(["ffn_graph/test/pull_stopped/2"]) + + +def test_close_waits_for_all_attention_then_all_drained_receivers(pull_world): + first_a, last_a, first_f, last_f = pull_world.sessions + with pytest.raises(TimeoutError): + first_a.close_attention() + first_f.drain(lambda: pytest.fail("an idle receiver needs no transaction")) + with pytest.raises(TimeoutError): + first_f.finish_receiver() + assert not pull_world.store.check(["ffn_graph/test/pull_stopped/2"]) + with pytest.raises(TimeoutError): + last_a.close_attention() + first_f.finish_receiver() + with pytest.raises(TimeoutError): + first_a.close_attention() + last_f.drain(lambda: pytest.fail("an idle receiver needs no transaction")) + last_f.finish_receiver() + first_a.close_attention() + last_a.close_attention() + assert pull_world.store.check( + [f"ffn_graph/test/pull_stopped/{rank}" for rank in (2, 3)] + ) + pull_world.startups[3].fail(RuntimeError("drain failed")) + with pytest.raises(RuntimeError, match="drain failed"): + first_a.close_attention() + + +def test_finish_receiver_is_idempotent_after_store_release(pull_world): + session = pull_world.sessions[2] + for rank in range(2): + pull_world.store.set(f"ffn_graph/test/pull_quiet/{rank}", "ready") + session.drain(lambda: pytest.fail("an idle receiver needs no transaction")) + session.finish_receiver() + pull_world.store.closed = True + session.finish_receiver() + + +@pytest.mark.parametrize("state", ["unfinished", "active", "failed", "drained"]) +def test_shutdown_requires_finished_idle_error_free_receiver(pull_world, state): + session, words = pull_world.sessions[2], pull_world.states[2].words + if state != "unfinished": + for rank in range(2): + pull_world.store.set(f"ffn_graph/test/pull_quiet/{rank}", "ready") + session.drain(lambda: pytest.fail("an idle receiver needs no transaction")) + session.finish_receiver() + if state == "active": + words[32] = 0 + if state == "failed": + words[24] = -5 + if state == "drained": + session.validate_shutdown() + else: + with pytest.raises(RuntimeError, match="retain resources|status -5"): + session.validate_shutdown() diff --git a/tests/unit/connectors/test_async_cam_startup.py b/tests/unit/connectors/npu/test_async_cam_startup.py similarity index 62% rename from tests/unit/connectors/test_async_cam_startup.py rename to tests/unit/connectors/npu/test_async_cam_startup.py index 186019a31..0ca545f3f 100644 --- a/tests/unit/connectors/test_async_cam_startup.py +++ b/tests/unit/connectors/npu/test_async_cam_startup.py @@ -2,15 +2,19 @@ # SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project """Pure host negotiation tests; no device or model runner initialization.""" +import json +import subprocess +import sys from concurrent.futures import ThreadPoolExecutor from dataclasses import replace -from threading import Condition from types import SimpleNamespace import pytest -from afd_plugin import async_cam_startup as startup_module -from afd_plugin.async_cam_startup import ( +from afd_plugin.connectors.npu.async_cam_runtime import ( + async_cam_startup as startup_module, +) +from afd_plugin.connectors.npu.async_cam_runtime.async_cam_startup import ( AsyncCAMGraphStartup, FFNGraphPlan, resolve_ffn_mode, @@ -20,41 +24,30 @@ requested_ffn_mode, validate_graph_capacity, ) - - -class MemoryStore: - def __init__(self): - self.values = {} - self.condition = Condition() - self.closed = False - - def set(self, key, value): - assert not self.closed, "Store used after release" - with self.condition: - self.values[key] = value.encode() if isinstance(value, str) else value - self.condition.notify_all() - - def get(self, key): - assert not self.closed, "Store used after release" - with self.condition: - assert self.condition.wait_for(lambda: key in self.values, timeout=3) - return self.values[key] - - def check(self, keys): - assert not self.closed, "Store used after release" - with self.condition: - return all(key in self.values for key in keys) - - -def startup_world(attn=2, ffn=2, store=None): - store = store or MemoryStore() - owners = [] - for rank in range(attn + ffn): - owner = AsyncCAMGraphStartup(rank, attn, attn + ffn) - owner._store = store - owner._prefix = "ffn_graph/test/" - owners.append(owner) - return owners +from tests.unit.fakes.async_cam_startup import MemoryStore, startup_world + + +def test_startup_import_needs_no_torch_or_vllm_in_clean_process(): + result = subprocess.run( + [ + sys.executable, + "-c", + "import sys; " + "sys.modules.update(torch=None, vllm=None, " + "torch_npu=None, vllm_ascend=None); " + "from afd_plugin.connectors.npu.async_cam_runtime.async_cam_startup " + "import AsyncCAMGraphStartup; " + "from afd_plugin.connectors.npu.async_cam_runtime.async_cam_session " + "import PullCAMSession, PushCAMSession; " + "assert AsyncCAMGraphStartup(0, 1, 2).mode == 'eager'; " + "assert PullCAMSession.requires_drain; " + "assert not PushCAMSession.requires_drain", + ], + capture_output=True, + text=True, + check=False, + ) + assert result.returncode == 0, result.stderr def full_plan(**changes): @@ -88,6 +81,71 @@ def negotiate(owners, plans, signatures): return [future.exception(timeout=3) for future in futures] +@pytest.mark.parametrize("world", [4, 16]) +def test_shared_failure_check_uses_one_store_request(world, monkeypatch): + owners = startup_world(attn=world - 2, ffn=2) + assert not any( + negotiate( + owners, + [None] * (world - 2) + [full_plan(requested="eager")] * 2, + ["wire"] * world, + ) + ) + calls = [] + original = owners[0]._store.check + + def check(keys): + calls.append(keys) + return original(keys) + + monkeypatch.setattr(owners[0]._store, "check", check) + owners[-1].check_failure() + assert len(calls) == 1 + owners[0].fail(RuntimeError("source failed")) + for owner in owners: + with pytest.raises(RuntimeError, match="source failed"): + owner.check_failure() + + +def test_legacy_peer_keeps_per_rank_failure_detection(monkeypatch): + owners = startup_world() + original = owners[0]._publish + + def legacy_publish(stage, value="ready"): + if stage == "plan": + fields = json.loads(value) + fields.pop("shared_errors") + value = json.dumps(fields) + original(stage, value) + + monkeypatch.setattr(owners[0], "_publish", legacy_publish) + assert not any( + negotiate( + owners, [None, None] + [full_plan(requested="eager")] * 2, ["wire"] * 4 + ) + ) + owners[0]._store.set(owners[0]._keys("error", range(1))[0], "legacy failure") + for owner in owners: + with pytest.raises(RuntimeError, match="legacy failure"): + owner.check_failure() + + +def test_shared_failure_survives_compatibility_write_failure(monkeypatch): + owners = startup_world() + original = owners[0]._store.set + + def fail_rank_error(key, value): + if key.endswith("error/0"): + raise ConnectionError("connection lost") + original(key, value) + + monkeypatch.setattr(owners[0]._store, "set", fail_rank_error) + with pytest.raises(ConnectionError): + owners[0].fail(RuntimeError("original failure")) + with pytest.raises(RuntimeError, match="original failure"): + owners[-1].check_failure() + + def test_mixed_requested_modes_are_not_silently_downgraded(): with pytest.raises(ValueError, match="same execution mode"): resolve_ffn_mode([full_plan(), FFNGraphPlan()]) @@ -106,38 +164,6 @@ def test_executor_error_reaches_all_attention_and_ffn_ranks(mode): ), errors -@pytest.mark.parametrize("mode", ["eager", "full"]) -def test_two_stream_signature_and_failure_are_checked_in_both_modes(mode): - plan = full_plan(requested=mode, stream_signature="2/12C/24V/24V") - assert resolve_ffn_mode([plan, plan]) == mode - with pytest.raises(ValueError, match="two-stream execution signatures"): - resolve_ffn_mode([plan, replace(plan, stream_signature="")]) - with pytest.raises(ValueError, match="core budget"): - resolve_ffn_mode([plan, replace(plan, stream_error="core budget")]) - with pytest.raises(ValueError, match="signatures differ"): - resolve_ffn_mode([plan, replace(plan, signature="other-capacity")]) - - -@pytest.mark.parametrize("peer_failure", [False, True]) -def test_two_stream_waits_for_attention_capability_check(peer_failure): - owners = startup_world(1, 1) - - def validate(): - if peer_failure: - raise ValueError("DBO is not qualified") - - owners[0].validate_attention_streams = validate - plan = full_plan(requested="eager", stream_signature="2/12C/24V/24V") - errors = negotiate(owners, [None, plan], ["wire"] * 2) - if peer_failure: - assert all( - isinstance(error, (ValueError, RuntimeError)) and "DBO" in str(error) - for error in errors - ), errors - else: - assert errors == [None, None] - - def test_whole_world_agrees_before_service_readiness(): owners = startup_world() assert ( @@ -151,13 +177,12 @@ def test_whole_world_agrees_before_service_readiness(): owner.wait_service_ready() -@pytest.mark.parametrize("two_stream", [False, True]) @pytest.mark.parametrize("mode", ["eager", "full"]) -def test_transport_check_preserves_single_stream_eager_scope(two_stream, mode): +def test_transport_check_preserves_eager_scope(mode): owners = startup_world(1, 1) - plan = full_plan(requested=mode, stream_signature="two" if two_stream else "") + plan = full_plan(requested=mode) errors = negotiate(owners, [None, plan], ["a", "b"]) - if two_stream or mode == "full": + if mode == "full": assert all( isinstance(error, (ValueError, RuntimeError)) and "transport signatures" in str(error) @@ -211,6 +236,39 @@ def test_graph_accepts_conservative_uint16_boundary(): validate_graph_capacity(65536, 2, 65536) +def test_pull_row_slicing_accepts_capacity_below_dispatch_envelope(): + validate_graph_capacity(32, 2, 3, allow_row_slicing=True) + with pytest.raises(ValueError, match="smaller"): + validate_graph_capacity(32, 2, 3) + with pytest.raises(ValueError, match="uint16"): + validate_graph_capacity(65536, 1, 3, allow_row_slicing=True) + + +def test_pull_requires_matching_transport_even_in_eager_mode(): + owners = startup_world(1, 1) + errors = negotiate( + owners, [None, full_plan(requested="eager")], ["pull_v1:layout", "push"] + ) + assert all("transport signatures" in str(error) for error in errors) + + +@pytest.mark.parametrize( + ("bad", "message"), + [ + (FFNGraphPlan(), "signatures differ"), + (full_plan(requested="eager", layer_ids=(3,)), "layer tables"), + ], +) +def test_pull_eager_validates_executor_and_layer_tables(bad, message): + owners = startup_world(1, 2) + errors = negotiate( + owners, + [None, full_plan(requested="eager"), bad], + ["pull_v1:layout"] * 3, + ) + assert all(message in str(error) for error in errors), errors + + @pytest.mark.parametrize("stage", ["validated", "warmup", "captured"]) def test_phase_requires_all_ranks_and_failure_wins_over_ready(monkeypatch, stage): owners = startup_world() diff --git a/tests/unit/test_async_cam_startup_wiring.py b/tests/unit/connectors/npu/test_async_cam_startup_wiring.py similarity index 67% rename from tests/unit/test_async_cam_startup_wiring.py rename to tests/unit/connectors/npu/test_async_cam_startup_wiring.py index dbae8fcd4..452364201 100644 --- a/tests/unit/test_async_cam_startup_wiring.py +++ b/tests/unit/connectors/npu/test_async_cam_startup_wiring.py @@ -11,12 +11,16 @@ import pytest -from afd_plugin.async_cam_startup import AsyncCAMGraphStartup, resolve_ffn_mode from afd_plugin.compat.npu.feature_validation import ( _fail_if_unsupported_npu_afd_async_features, fail_if_unsupported_npu_afd_features, ) from afd_plugin.config import parse_afd_config +from afd_plugin.connectors.npu.async_cam_runtime.async_cam_startup import ( + AsyncCAMGraphStartup, + resolve_ffn_mode, +) +from tests.unit.fakes.async_cam_startup import startup_world from tests.unit.fakes.async_ffn import ( load_async_modules, load_host_module, @@ -132,14 +136,14 @@ def executor(layers): connector = SimpleNamespace( is_initialized=False, dynamic_quant=1, - extra_info=SimpleNamespace(ffn_two_stream=False), + extra_info=SimpleNamespace(ffn_mega_moe=False), graph_startup=AsyncCAMGraphStartup(1, 1, 2), ) connector.init_afd_connector = lambda: resolve_ffn_mode( [connector.graph_startup.plan] ) runner.connector = connector - runner._failure = runner._layered_executor = runner._two_stream = None + runner._failure = runner._layered_executor = None runner.requested_mode = mode monkeypatch.setattr( runtime, "get_ascend_device_type", lambda: runtime.AscendDeviceType.A3 @@ -173,91 +177,51 @@ def test_ffn_startup_validates_target_and_draft_weights( assert connector.graph_startup.plan.signature == "(2, 5)" -@pytest.mark.parametrize("value", [True, 0, -1, 1.5]) -def test_two_stream_rejects_invalid_core_settings(modules, value): - _, connector, _ = modules - with pytest.raises((TypeError, ValueError), match="ffn_compute_aic_num"): - connector.AFDAsyncExtraInfo.from_mapping({"ffn_compute_aic_num": value}) +@pytest.mark.parametrize("mode", ["eager", "full"]) +def test_mega_executor_failure_is_published_during_startup(modules, monkeypatch, mode): + _, _, runtime = modules + runner = runtime.AsyncNPUFFNModelRunner.__new__(runtime.AsyncNPUFFNModelRunner) + runner.model = SimpleNamespace(get_experts_layer_indices=lambda: (2,)) + runner.draft_model = None + runner._failure = runner._layered_executor = runner._mega_executor = None + runner.requested_mode = mode + connector = SimpleNamespace( + is_initialized=False, + extra_info=SimpleNamespace(ffn_mega_moe=True), + graph_startup=AsyncCAMGraphStartup(1, 1, 2), + ) + connector.init_afd_connector = lambda: resolve_ffn_mode( + [connector.graph_startup.plan] + ) + runner.connector = connector + runner._initialize_layered_executor = lambda: setattr( + runner, + "_layered_executor", + SimpleNamespace(graph_signature=lambda _: "layered"), + ) + def fail_executor(*_args): + raise ValueError("unsupported Mega weights") -def test_two_stream_config_round_trip_and_default_off(modules): - _, connector, _ = modules - default = connector.AFDAsyncExtraInfo.from_mapping({}) - assert not default.ffn_two_stream - assert "ffn_two_stream" not in default.to_mapping() - enabled = connector.AFDAsyncExtraInfo.from_mapping({"ffn_two_stream": True}) - assert enabled.ffn_compute_aic_num == 12 - assert enabled.ffn_compute_aiv_num == enabled.ffn_combine_aiv_num == 24 - assert connector.AFDAsyncExtraInfo.from_mapping(enabled.to_mapping()) == enabled - disabled = connector.AFDAsyncExtraInfo.from_mapping({"ffn_compute_aic_num": 8}) - assert connector.AFDAsyncExtraInfo.from_mapping(disabled.to_mapping()) == disabled - - -@pytest.mark.parametrize("feature", ["use_sequence_parallel_moe"]) -def test_two_stream_attention_preflight_rejects_unqualified_modes( - modules, monkeypatch, feature -): - _, module, _ = modules - connector = module.CAMAsyncAFDConnector.__new__(module.CAMAsyncAFDConnector) - connector.vllm_config = npu_vllm_config( - enable_dbo=feature == "enable_dbo", - num_ubatches=2, - use_sequence_parallel_moe=feature == "use_sequence_parallel_moe", - ) - monkeypatch.setenv( - "VLLM_ASCEND_ENABLE_FLASHCOMM1", "1" if feature == "flashcomm1" else "0" - ) - if feature in (None, "enable_dbo"): - connector._validate_two_stream_attention() - else: - with pytest.raises(ValueError, match="two-stream"): - connector._validate_two_stream_attention() + monkeypatch.setattr(runtime, "AsyncCAMMegaExecutor", fail_executor) + with pytest.raises(ValueError, match="unsupported Mega weights"): + runner.initialize_afd_connector() + assert connector.graph_startup.plan.executor_error == "unsupported Mega weights" -@pytest.mark.parametrize("role", ["attention", "ffn"]) @pytest.mark.parametrize( - "value,enabled", + "field", [ - (None, False), - ("0", False), - ("1", True), - (" TRUE ", True), - ("yes", False), + "ffn_two_stream", + "ffn_compute_aic_num", + "ffn_compute_aiv_num", + "ffn_combine_aiv_num", ], ) -def test_two_stream_flashcomm1_matches_ascend_boolean_parsing( - modules, monkeypatch, role, value, enabled -): - _, module, _ = modules - if value is None: - monkeypatch.delenv("VLLM_ASCEND_ENABLE_FLASHCOMM1", raising=False) - else: - monkeypatch.setenv("VLLM_ASCEND_ENABLE_FLASHCOMM1", value) - config = npu_vllm_config( - role=role, - connector="CAMAsyncAFDConnector", - async_dp=True, - compute_gate_on_attention=True, - use_mla=True, - ) - if role == "ffn": - - def validate(): - _fail_if_unsupported_npu_afd_async_features( - config, - parse_afd_config(config), - module.AFDAsyncExtraInfo(ffn_two_stream=True, dynamic_quant=1), - ) - - else: - connector = module.CAMAsyncAFDConnector.__new__(module.CAMAsyncAFDConnector) - connector.vllm_config = config - validate = connector._validate_two_stream_attention - if enabled: - with pytest.raises((ValueError, RuntimeError), match="two-stream"): - validate() - else: - validate() +def test_removed_two_stream_options_are_rejected(modules, field): + _, connector, _ = modules + with pytest.raises(ValueError, match="unknown AFD async connector_extra_config"): + connector.AFDAsyncExtraInfo.from_mapping({field: True}) @pytest.mark.parametrize( @@ -291,6 +255,7 @@ def create(*args, **kwargs): ) connector.__dict__.update( _initialized=False, + extra_info=connector_module.AFDAsyncExtraInfo(), afd_config=SimpleNamespace(host="localhost", port=1), topology=SimpleNamespace(world_size=2), world_rank=0, @@ -311,6 +276,7 @@ def create(*args, **kwargs): ), graph_startup=AsyncCAMGraphStartup(0, 1, 2), ) + connector._session = connector_module.PushCAMSession(connector.activation_dtype) if dbo_flag is not None: setattr(connector.vllm_config.parallel_config, dbo_flag, True) @@ -366,16 +332,21 @@ def test_graph_dbo_keeps_existing_role_and_count_validation( def test_ffn_announces_service_only_on_receiver_execution(modules): - _, _, runtime = modules + _, connector_module, runtime = modules runner = runtime.AsyncNPUFFNModelRunner.__new__(runtime.AsyncNPUFFNModelRunner) calls: list[str | tuple[str, str | int]] = [] runner.__dict__.update( prof=None, _graph=object(), _service_announced=False, + _first_step=True, + _graph_stream=MagicMock(), + device="npu:0", + session=sys.modules["tests.unit.fakes.connector"].FakeFFNSession(), layer_ids=(2, 5), connector=SimpleNamespace( - graph_startup=SimpleNamespace(service_ready=lambda: calls.append("ready")) + extra_info=connector_module.AFDAsyncExtraInfo(), + graph_startup=SimpleNamespace(service_ready=lambda: calls.append("ready")), ), _submit_graph=lambda count: calls.append(("submit", count)), _complete_graph=lambda: True, @@ -386,10 +357,33 @@ def test_ffn_announces_service_only_on_receiver_execution(modules): assert calls == ["ready", ("submit", 2), ("submit", 2)] +@pytest.mark.parametrize("session_serves", [False, True]) +def test_ffn_step_lets_session_decide_service_after_shutdown(modules, session_serves): + _, connector_module, runtime = modules + runner = runtime.AsyncNPUFFNModelRunner.__new__(runtime.AsyncNPUFFNModelRunner) + submitted: list[int] = [] + serving = [True] if session_serves else [] + runner.__dict__.update( + prof=None, + _graph=object(), + _service_announced=True, + _first_step=False, + session=sys.modules["tests.unit.fakes.connector"].FakeFFNSession( + serving=serving + ), + layer_ids=(2, 5), + connector=SimpleNamespace(extra_info=connector_module.AFDAsyncExtraInfo()), + _submit_graph=submitted.append, + _complete_graph=lambda: True, + ) + # Push stops on the shutdown request; Pull keeps serving until A is quiet. + assert runner.execute_connector_driven_step(True) is session_serves + assert submitted == ([2] if session_serves else []) + + @pytest.mark.parametrize("capture_fails", [False, True]) -@pytest.mark.parametrize("two_stream", [False, True]) def test_ffn_warmup_capture_order_and_failure_propagation( - modules, monkeypatch, capture_fails, two_stream + modules, monkeypatch, capture_fails ): # long-test: warmup/capture/readiness are one ordered startup protocol _, connector_module, runtime = modules calls: list[str | tuple[str, str | int]] = [] @@ -401,8 +395,6 @@ def initialize_connector(): calls.append("connector") assert startup.plan.signature == "weights" assert startup.plan.layer_ids == (2, 5) - assert not startup.plan.stream_error - assert bool(startup.plan.stream_signature) == two_stream startup.mode = "full" def capture(): @@ -419,10 +411,11 @@ def capture(): draft_model=None, _initialize_layered_executor=lambda: None, _layered_executor=SimpleNamespace(graph_signature=lambda capacity: "weights"), + _mega_executor=None, connector=SimpleNamespace( is_initialized=False, graph_startup=startup, - extra_info=connector_module.AFDAsyncExtraInfo(ffn_two_stream=two_stream), + extra_info=connector_module.AFDAsyncExtraInfo(), init_afd_connector=initialize_connector, attn_size=2, tp_size=1, @@ -443,7 +436,7 @@ def capture(): *(["transaction"] * 4), "device-done", "warmup", - *([] if two_stream else ["graph-resources"]), + "graph-resources", "capture", ("failure", "capture failed") if capture_fails else "captured", ] @@ -527,6 +520,7 @@ def receiver_graph_worker(ffn_owners, monkeypatch): ) connector = SimpleNamespace( is_initialized=True, + ffn_session=sys.modules["afd_plugin.connectors.base"].FFNExecutionSession, close=lambda: None, recv_attn_output=lambda: SimpleNamespace( hidden_states=object(), context=SimpleNamespace(states=state) @@ -585,20 +579,27 @@ def serve(): @pytest.mark.parametrize("state", ["live", "failed", "drained"]) -@pytest.mark.parametrize("two_stream", [False, True]) -def test_worker_retains_resources_until_clean_exit(ffn_owners, state, two_stream): +@pytest.mark.parametrize("requires_drain", [False, True]) +def test_worker_retains_resources_until_clean_exit(ffn_owners, state, requires_drain): _, runtime, module = ffn_owners calls = [] + timeouts = [] runner = runtime.AsyncNPUFFNModelRunner.__new__(runtime.AsyncNPUFFNModelRunner) - runner.effective_mode = "eager" if two_stream else "full" - runner._two_stream = object() if two_stream else None + runner.effective_mode = "full" + runner.session = SimpleNamespace(requires_drain=requires_drain) runner.connector = SimpleNamespace(close=lambda: calls.append("connector-close")) runner.shutdown = lambda: calls.append("release") worker = module.AFDNPUFFNWorker.__new__(module.AFDNPUFFNWorker) worker.model_runner = runner worker._ffn_shutdown_event = SimpleNamespace(set=lambda: calls.append("stop")) + + def join(**kwargs): + calls.append("join") + timeouts.append(kwargs["timeout"]) + worker._ffn_thread = SimpleNamespace( - join=lambda **kw: calls.append("join"), is_alive=lambda: state == "live" + join=join, + is_alive=lambda: state == "live", ) worker._ffn_loop_error = ( RuntimeError("native failure") if state == "failed" else None @@ -616,6 +617,7 @@ def test_worker_retains_resources_until_clean_exit(ffn_owners, state, two_stream worker.raise_ffn_loop_error_if_any() with pytest.raises(RuntimeError): worker.start_ffn_server_loop() + assert timeouts == [60 if requires_drain else 5] @pytest.mark.parametrize("async_connector", [False, True]) @@ -665,7 +667,10 @@ def test_ffn_runner_loads_model_and_rejects_second_load( monkeypatch.setattr(runtime, "validate_npu_model_runner_v2_config", calls.validate) monkeypatch.setattr(runtime, "fail_if_unsupported_npu_afd_features", calls.validate) monkeypatch.setattr(runtime, "create_afd_npu_profiler", lambda role: None) - connector, device, model = object(), SimpleNamespace(type="npu"), object() + connector = SimpleNamespace( + ffn_session=sys.modules["afd_plugin.connectors.base"].FFNExecutionSession + ) + device, model = SimpleNamespace(type="npu"), object() connector_name = "CAMAsyncAFDConnector" if async_connector else "CAMP2pAFDConnector" config = npu_vllm_config(role="ffn", connector=connector_name) config.model_config.enforce_eager = False @@ -715,3 +720,185 @@ def get_loader(load_config): assert runner.get_model() is model and runner.model_memory_usage == 1234 with pytest.raises(RuntimeError, match="already loaded"): runner.load_model() + + +@pytest.fixture +def connector_world(monkeypatch): + _, module, runtime = load_async_modules(monkeypatch) + owners = startup_world() + connectors = [] + for rank, startup in enumerate(owners): + role = "attention" if rank < 2 else "ffn" + config = npu_vllm_config( + role=role, + connector="CAMAsyncAFDConnector", + num_attention_ranks=2, + num_ffn_ranks=2, + extra_config={"ffn_mega_moe": True, "dynamicQuant": 1}, + ) + config.model_config.hf_config = SimpleNamespace( + hidden_size=32, num_experts_per_tok=2, n_routed_experts=4 + ) + config.model_config.dtype = "bfloat16" + connector = module.CAMAsyncAFDConnector( + rank, 0, config, parse_afd_config(config), rank % 2 + ) + connector.graph_startup = startup + connector.ffn_session().startup = startup + connector._initialized = True + connector.ffn_session().comm_args = object() + connectors.append(connector) + return SimpleNamespace(connectors=connectors, runtime=runtime) + + +def test_default_session_obeys_shutdown_without_drain(connector_world): + session = sys.modules["afd_plugin.connectors.base"].FFNExecutionSession() + assert session.keep_serving(False) + assert not session.keep_serving(True) + assert not session.requires_drain + session.before_first_step() + session.after_step() + session.drain(lambda: pytest.fail("Push must not execute a drain transaction")) + session.finish_receiver() + session.finish_receiver() + session.validate_shutdown() + + +@pytest.mark.parametrize("failure", ["unfinished", "active", "status"]) +def test_close_preserves_receiver_communication_on_unsafe_session( + connector_world, monkeypatch, failure +): + connector = connector_world.connectors[2] + resource = connector.comm_args + connector.cam_pg = group = object() + session = connector.ffn_session() + if failure != "unfinished": + monkeypatch.setattr( + connector.ffn_session(), "pull_has_active_source", lambda: False + ) + session.drain(lambda: pytest.fail("idle source must not launch")) + for attention in connector_world.connectors[:2]: + attention.graph_startup._publish("pull_quiet") + session.finish_receiver() + if failure == "active": + monkeypatch.setattr( + connector.ffn_session(), "pull_has_active_source", lambda: True + ) + elif failure == "status": + + def fail_status(): + raise RuntimeError("sticky status -5") + + monkeypatch.setattr(connector.ffn_session(), "check_pull_status", fail_status) + expected = "sticky status -5" if failure == "status" else "has not drained" + with pytest.raises(RuntimeError, match=expected): + connector.close() + assert connector.is_initialized and connector.comm_args is resource + assert connector.cam_pg is group + + +@pytest.mark.parametrize("failure", ["body", "status"]) +def test_attention_step_recovers_after_error(connector_world, monkeypatch, failure): + connector = connector_world.connectors[0] + error: list[RuntimeError | None] = [None] + + def check_status(): + if error[0] is not None: + raise error[0] + + monkeypatch.setattr(connector.ffn_session(), "check_pull_status", check_status) + expected = ValueError if failure == "body" else RuntimeError + with pytest.raises(expected, match="failed"), connector.attention_step(): + if failure == "body": + raise ValueError("model failed") + error[0] = RuntimeError("status failed") + error[0] = RuntimeError("status failed after recovery") + with ( + pytest.raises(RuntimeError, match="after recovery"), + connector.attention_step(), + ): + pass + error[0] = None + with connector.attention_step(): + pass + + +def test_nested_attention_step_defers_status_until_outer_boundary( + connector_world, monkeypatch +): + connector = connector_world.connectors[0] + + def check_status(): + raise RuntimeError("sticky status") + + monkeypatch.setattr(connector.ffn_session(), "check_pull_status", check_status) + with pytest.raises(RuntimeError, match="sticky status"), connector.attention_step(): + with connector.attention_step(): + pass + # Reaching this line proves the nested scope did not check sticky status. + completed_inner_scope = True + assert completed_inner_scope + + +@pytest.mark.parametrize("fails", [False, True]) +def test_runner_drain_returns_completed_response_or_session_error( + connector_world, fails +): + runtime = connector_world.runtime + runner = runtime.AsyncNPUFFNModelRunner.__new__(runtime.AsyncNPUFFNModelRunner) + runner.model = object() + runner.connector = connector_world.connectors[2] + runner.session = sys.modules["tests.unit.fakes.connector"].FakeFFNSession( + transactions=1, + drain_error=RuntimeError("source did not drain") if fails else None, + ) + runner._graph = None + runner.layer_ids = (2, 5) + responses = [] + runner.connector.send_ffn_output = lambda output, context: responses.append(output) + # The host executor emits a peer response; it does not simulate device compute. + runner._mega_executor = lambda: runner.connector.send_ffn_output("response", None) + if fails: + with pytest.raises(RuntimeError, match="source did not drain"): + runner.drain() + else: + runner.drain() + assert responses == ["response"] + + +@pytest.mark.parametrize("failure", ["unfinished", "active", "status"]) +def test_runner_shutdown_retains_model_and_communication_on_unsafe_session( + connector_world, monkeypatch, failure +): + runtime = connector_world.runtime + runner = runtime.AsyncNPUFFNModelRunner.__new__(runtime.AsyncNPUFFNModelRunner) + connector = connector_world.connectors[2] + runner.connector = connector + runner.session = connector.ffn_session() + runner.model = model = object() + runner._failure = None + runner._is_shutdown = False + runner._graph_receiver_exited = True + resource = connector.comm_args + if failure != "unfinished": + monkeypatch.setattr( + connector.ffn_session(), "pull_has_active_source", lambda: False + ) + runner.session.drain(lambda: pytest.fail("idle source must not launch")) + for attention in connector_world.connectors[:2]: + attention.graph_startup._publish("pull_quiet") + runner.session.finish_receiver() + if failure == "active": + monkeypatch.setattr( + connector.ffn_session(), "pull_has_active_source", lambda: True + ) + elif failure == "status": + + def fail_status(): + raise RuntimeError("status failed; retain resources") + + monkeypatch.setattr(connector.ffn_session(), "check_pull_status", fail_status) + with pytest.raises(RuntimeError, match="retain resources"): + runner.shutdown() + assert runner.get_model() is model + assert connector.is_initialized and connector.comm_args is resource diff --git a/tests/unit/connectors/test_async_cam_connector.py b/tests/unit/connectors/test_async_cam_connector.py index 46c72c541..ea8efe7f9 100644 --- a/tests/unit/connectors/test_async_cam_connector.py +++ b/tests/unit/connectors/test_async_cam_connector.py @@ -3,13 +3,11 @@ from __future__ import annotations -import weakref +import sys from types import SimpleNamespace import pytest -from afd_plugin.config import AFDConfig - pytest.importorskip("torch") pytest.importorskip("torch_npu") @@ -19,16 +17,29 @@ AFDConnectorFactory, AFDTransferContext, AFDTransferMetadata, - AFDTransferState, ) from afd_plugin.connectors.npu import async_cam as async_cam_module # noqa: E402 from afd_plugin.connectors.npu.async_cam import ( # noqa: E402 - CAM_COMM_ID, AFDAsyncExtraInfo, - AFDAsyncTransferState, CAMAsyncAFDConnector, build_async_topology, ) +from afd_plugin.connectors.npu.async_cam_runtime import ( + async_cam_session as session_module, # noqa: E402 +) +from tests.unit.fakes.async_cam import ( # noqa: E402 + _afd_config, + _dp2_afd_config, + _FakeTensor, + _FakeTorch, + _topk_payload, + cam_vllm_config, +) +from tests.unit.fakes.async_cam_startup import MemoryStore # noqa: E402 + + +def _vllm_config(**kwargs): + return cam_vllm_config(dtype=torch.bfloat16, **kwargs) class _FakeTensorLike: @@ -43,134 +54,54 @@ def __getitem__(self, item): return f"{self.name}[{start}:{stop}]" -class _FakeTensor: - def __init__(self, shape, *, dtype="bf16", device="npu:0"): - self.shape = tuple(shape) - self.dtype = dtype - self.device = device - - def to(self, *, dtype): - self.dtype = dtype - return self - - def contiguous(self): - return self - - def new_zeros(self, shape): - return _FakeTensor(shape, dtype=self.dtype, device=self.device) - - def new_empty(self, shape): - return _FakeTensor(shape, dtype=self.dtype, device=self.device) - - -class _FakeCamOps: - def __init__(self): - self.calls = [] - - def afd_async_dispatch_send(self, *args): - self.calls.append(("dispatch_send", args)) - return args[0] - - def afd_async_dispatch_recv(self, *args): - self.calls.append(("dispatch_recv", args)) - batch_size = args[3] - hidden_size = args[4] - expert_per_rank = args[8] - tp_size = args[11] - return ( - _FakeTensor((batch_size, hidden_size)), - _FakeTensor((batch_size,), dtype="fp32"), - _FakeTensor((5 + tp_size * (1 + expert_per_rank),), dtype="int64"), - _FakeTensor((expert_per_rank,), dtype="int64"), - ) - - def afd_async_combine_send(self, *args): - self.calls.append(("combine_send", args)) - return args[0] - - def afd_async_combine_recv(self, *args): - self.calls.append(("combine_recv", args)) - batch_size = args[5] - hidden_size = args[6] - return _FakeTensor((batch_size, hidden_size)) - - -class _FakeTorch: - def __init__(self): - self.bfloat16 = "bf16" - self.float16 = "fp16" - self.float32 = "fp32" - self.int32 = "int32" - self.int64 = "int64" - self.ops = SimpleNamespace(afd_ascend=_FakeCamOps()) - - def device(self, name): - return name +def test_async_extra_info_rejects_unknown_and_invalid_fields(): + with pytest.raises(ValueError, match="unknown AFD async connector_extra_config"): + AFDAsyncExtraInfo.from_mapping({"core_num": 8}) + with pytest.raises(ValueError, match="attn_ranks_per_dp must be positive"): + AFDAsyncExtraInfo.from_mapping({"attn_ranks_per_dp": 0}) - def empty(self, shape, *, dtype, device): - return _FakeTensor(shape, dtype=dtype, device=device) - def zeros(self, shape, *, dtype, device): - return _FakeTensor(shape, dtype=dtype, device=device) +def test_mega_config_round_trips_protocol_and_normalized_quantization(): + info = AFDAsyncExtraInfo.from_mapping({"ffn_mega_moe": "true", "dynamicQuant": "1"}) + assert info.ffn_mega_moe and info.dynamic_quant == 1 + assert AFDAsyncExtraInfo.from_mapping(info.to_mapping()) == info -def _vllm_config(*, tp_size: int = 1, pcp_size: int = 1, extra_config=None): - return SimpleNamespace( - additional_config={ - "afd": { - "connector_extra_config": extra_config - if extra_config is not None - else {"attn_ranks_per_dp": 4}, - }, - }, - parallel_config=SimpleNamespace( - data_parallel_size=1, - data_parallel_rank=0, - prefill_context_parallel_size=pcp_size, - tensor_parallel_size=tp_size, - ), - scheduler_config=SimpleNamespace(max_num_batched_tokens=16), - model_config=SimpleNamespace( - dtype=torch.bfloat16, - hf_config=SimpleNamespace( - hidden_size=16, - num_experts_per_tok=2, - n_routed_experts=8, - ), +@pytest.mark.parametrize("capacity", [1, 16]) +def test_mega_warmup_launches_once_per_source(monkeypatch, capacity): + monkeypatch.setenv("BATCH_SIZE_FACTOR", str(capacity / 262144)) + connector = CAMAsyncAFDConnector( + 0, + 0, + _vllm_config( + extra_config={ + "ffn_mega_moe": True, + "dynamicQuant": 1, + "attn_ranks_per_dp": 4, + } ), + _afd_config(role="ffn"), + 0, ) + assert connector.graph_startup_chunks() == 1 -def _afd_config(*, role: str): - return AFDConfig( - connector="CAMAsyncAFDConnector", - role=role, - num_attention_ranks=4, - num_ffn_ranks=2, - ) - - -def _dp2_afd_config(*, role: str): - return AFDConfig( - connector="CAMAsyncAFDConnector", - role=role, - num_attention_ranks=8, - num_ffn_ranks=8, - ) - - -def _topk_payload(batch_size: int, topk: int = 2): - return { - "topk_ids": _FakeTensor((batch_size, topk), dtype="int32"), - "topk_weights": _FakeTensor((batch_size, topk), dtype="fp32"), - } +def test_mega_requires_dynamic_quantization(): + with pytest.raises(ValueError, match="ffn_mega_moe"): + AFDAsyncExtraInfo.from_mapping({"ffn_mega_moe": True}) -def test_async_extra_info_rejects_unknown_and_invalid_fields(): +@pytest.mark.parametrize( + "field", + [ + "async_moe_ubatching", + "async_moe_num_ubatches", + "async_moe_split", + ], +) +def test_async_extra_info_rejects_removed_fields(field): with pytest.raises(ValueError, match="unknown AFD async connector_extra_config"): - AFDAsyncExtraInfo.from_mapping({"core_num": 8}) - with pytest.raises(ValueError, match="attn_ranks_per_dp must be positive"): - AFDAsyncExtraInfo.from_mapping({"attn_ranks_per_dp": 0}) + AFDAsyncExtraInfo.from_mapping({field: True}) def test_async_connector_factory_creates_import_safe_connector(): @@ -258,152 +189,6 @@ def test_async_topology_uses_one_world_for_all_dp_replicas( assert topology.expert_per_rank == 4 -def test_async_connector_calls_cam_shaped_ops(monkeypatch): - fake_torch = _FakeTorch() - monkeypatch.setattr(async_cam_module, "torch", fake_torch) - connector = CAMAsyncAFDConnector( - 0, - 0, - _vllm_config( - pcp_size=3, - extra_config={"attn_ranks_per_dp": 4}, - ), - _afd_config(role="attention"), - 0, - ) - connector._initialized = True - connector.comm_args = _FakeTensor((1,), dtype="fp16") - connector._placeholder = _FakeTensor((8, 16)) - hidden_states = _FakeTensor((3, 16)) - metadata = AFDTransferMetadata.create_attention_metadata( - layer_idx=2, - stage_idx=0, - seq_len=3, - ) - context = AFDTransferContext(metadata=metadata) - - output = connector.send_attn_output( - hidden_states, - context, - **_topk_payload(3), - ) - combined = connector.recv_ffn_output( - ref_tensor=hidden_states, - ubatch_idx=0, - ) - - assert output is None - assert combined.shape == (3, 16) - assert fake_torch.ops.afd_ascend.calls[0][0] == "dispatch_send" - assert fake_torch.ops.afd_ascend.calls[1][0] == "combine_recv" - assert fake_torch.ops.afd_ascend.calls[0][1][3] == CAM_COMM_ID - assert fake_torch.ops.afd_ascend.calls[1][1][4] == CAM_COMM_ID - assert fake_torch.ops.afd_ascend.calls[0][1][5:11] == (3, 16, 2, 2, 4, 4) - assert fake_torch.ops.afd_ascend.calls[1][1][5:11] == (3, 16, 2, 2, 4, 4) - assert fake_torch.ops.afd_ascend.calls[0][1][14] == 4 - assert isinstance(context.states, AFDAsyncTransferState) - assert isinstance(context.states, AFDTransferState) - - -def test_async_explicit_receive_consumes_matching_pending_dispatch(monkeypatch): - fake_torch = _FakeTorch() - monkeypatch.setattr(async_cam_module, "torch", fake_torch) - connector = CAMAsyncAFDConnector( - 0, 0, _vllm_config(), _afd_config(role="attention"), 0 - ) - connector._initialized = True - connector.comm_args = _FakeTensor((1,), dtype="fp16") - hidden_states = _FakeTensor((3, 16)) - contexts = [] - payloads = [] - for layer_idx in (2, 3): - context = AFDTransferContext( - metadata=AFDTransferMetadata.create_attention_metadata( - layer_idx=layer_idx, stage_idx=0, seq_len=3 - ) - ) - payload = _topk_payload(3) - connector.send_attn_output(hidden_states, context, **payload) - contexts.append(context) - payloads.append(payload) - - with pytest.raises(RuntimeError, match="does not match the next pending"): - connector.recv_ffn_output(hidden_states, ubatch_idx=0, context=contexts[1]) - assert len(fake_torch.ops.afd_ascend.calls) == 2 - - connector.recv_ffn_output( - hidden_states, - ubatch_idx=0, - context=contexts[0], - **payloads[1], - ) - assert fake_torch.ops.afd_ascend.calls[-1][1][1] is payloads[0]["topk_ids"] - - connector.recv_ffn_output(hidden_states, ubatch_idx=0) - assert fake_torch.ops.afd_ascend.calls[-1][1][1] is payloads[1]["topk_ids"] - - -def test_async_ffn_side_dispatch_recv_and_combine_send(monkeypatch): - fake_torch = _FakeTorch() - monkeypatch.setattr(async_cam_module, "torch", fake_torch) - connector = CAMAsyncAFDConnector( - 0, - 0, - _vllm_config( - pcp_size=2, - extra_config={"attn_ranks_per_dp": 2}, - ), - _afd_config(role="ffn"), - 0, - ) - connector._initialized = True - connector.comm_args = _FakeTensor((1,), dtype="fp16") - connector._placeholder = _FakeTensor((8, 16)) - - recv_output = connector.recv_attn_output(batch_size=4, layer_idx=1) - connector.send_ffn_output(recv_output.hidden_states, recv_output.context) - - states = recv_output.context.states - assert recv_output.hidden_states.shape == (connector.max_num_batched_tokens, 16) - assert states.dynamic_scales.shape == (connector.max_num_batched_tokens,) - assert states.group_list.shape == (connector.topology.expert_per_rank,) - assert fake_torch.ops.afd_ascend.calls[0][0] == "dispatch_recv" - assert fake_torch.ops.afd_ascend.calls[1][0] == "combine_send" - assert fake_torch.ops.afd_ascend.calls[0][1][11] == 2 - assert fake_torch.ops.afd_ascend.calls[1][1][12] == 2 - assert fake_torch.ops.afd_ascend.calls[1][1][2] is states.token_nums_rankid_layeridx - - -def test_async_combine_send_requires_dispatch_recv_token_metadata(monkeypatch): - fake_torch = _FakeTorch() - monkeypatch.setattr(async_cam_module, "torch", fake_torch) - connector = CAMAsyncAFDConnector( - 0, - 0, - _vllm_config(), - _afd_config(role="ffn"), - 0, - ) - connector._initialized = True - metadata = AFDTransferMetadata.create_ffn_metadata( - layer_idx=1, - stage_idx=0, - seq_lens=[4], - ) - # token_nums_rankid_layeridx is left unset (None) so combine send must fail. - states = AFDAsyncTransferState( - batch_size=4, - hidden_size=16, - topk=2, - layer_idx=1, - group_list=_FakeTensor((4,), dtype="int64"), - ) - context = AFDTransferContext(metadata=metadata, states=states) - - with pytest.raises(RuntimeError, match="TokenNums_Rankid_Layeridx"): - connector.send_ffn_output(_FakeTensor((4, 16)), context) - - def test_async_select_experts_rejects_mix_placement(): with pytest.raises(RuntimeError, match="routed-only.*mix_placement"): async_cam_module.select_cam_experts( @@ -455,29 +240,122 @@ def test_async_rejects_nondivisible_expert_placement(): build_async_topology(_afd_config(role="ffn"), 0, num_routed_experts=7) -def test_close_releases_every_stage_routing(monkeypatch): +def test_async_rejects_fused_shared_expert_ids(): + from afd_plugin.compat.npu.feature_validation import ( + _fail_if_unsupported_npu_afd_async_features, + ) + + config = _vllm_config() + config.additional_config["mix_placement"] = True + with pytest.raises(RuntimeError, match="routed-only.*mix_placement"): + _fail_if_unsupported_npu_afd_async_features( + config, + _afd_config(role="attention"), + AFDAsyncExtraInfo(), + ) + + +@pytest.fixture +def pull_attention(monkeypatch): + # CPU status words and host transfer recording, not NPU validation. fake_torch = _FakeTorch() + fake_torch.int64 = torch.int64 monkeypatch.setattr(async_cam_module, "torch", fake_torch) connector = CAMAsyncAFDConnector( - 0, 0, _vllm_config(), _afd_config(role="attention"), 0 + 0, + 0, + _vllm_config( + extra_config={ + "ffn_mega_moe": True, + "attn_ranks_per_dp": 4, + "dynamicQuant": 1, + } + ), + _afd_config(role="attention"), + 0, ) connector._initialized = True - connector.comm_args = _FakeTensor((1,), dtype="fp16") - references: list[weakref.ReferenceType] = [] - for stage in (0, 1): - context = AFDTransferContext( - metadata=AFDTransferMetadata.create_attention_metadata( - layer_idx=2, stage_idx=stage, seq_len=3 - ) + state = torch.zeros(128, dtype=torch.int64) + monkeypatch.setitem(sys.modules, "torch", fake_torch) + connector.ffn_session().comm_args = state.view(torch.float16) + store = MemoryStore() + connector.graph_startup._store = store + connector.graph_startup._prefix = "ffn_graph/test/" + for rank in (4, 5): + store.set(f"ffn_graph/test/pull_stopped/{rank}", "ready") + hidden = _FakeTensor((3, 16)) + context = AFDTransferContext( + metadata=AFDTransferMetadata.create_attention_metadata( + layer_idx=2, stage_idx=0, seq_len=3 ) - payload = _topk_payload(3) - references.extend(weakref.ref(tensor) for tensor in payload.values()) - connector.send_attn_output(_FakeTensor((3, 16)), context, **payload) - del payload, context - fake_torch.ops.afd_ascend.calls.clear() - assert all(reference() is not None for reference in references) + ) + return SimpleNamespace( + connector=connector, + state=state, + hidden=hidden, + context=context, + torch=fake_torch, + ) + + +@pytest.mark.parametrize("phase", ["launch", "status"]) +def test_pull_combine_failure_retains_window_until_completion( + pull_attention, monkeypatch, phase +): + h = pull_attention + connector = h.connector + connector.send_attn_output(h.hidden, h.context, **_topk_payload(3)) + combine = h.torch.ops.afd_ascend.afd_mega_combine_recv + + def fail(*_args, **_kwargs): + raise RuntimeError("combine failed") + + if phase == "launch": + monkeypatch.setattr(h.torch.ops.afd_ascend, "afd_mega_combine_recv", fail) + else: + h.state[session_module.PULL_FAILURE_INDEX] = -5 + with pytest.raises(RuntimeError, match="combine failed|status -5"): + connector.recv_ffn_output(h.hidden, context=h.context) + with pytest.raises(RuntimeError, match="pending"): + connector.close() + monkeypatch.setattr(h.torch.ops.afd_ascend, "afd_mega_combine_recv", combine) + h.state[session_module.PULL_FAILURE_INDEX] = 0 + connector.recv_ffn_output(h.hidden, context=h.context) connector.close() - assert all(reference() is None for reference in references) + assert not connector.is_initialized and connector.comm_args is None + + +@pytest.mark.parametrize("nested", [False, True]) +def test_attention_step_reports_status_after_queued_combine(pull_attention, nested): + h = pull_attention + connector = h.connector + with pytest.raises(RuntimeError, match="status -5"), connector.attention_step(): + connector.send_attn_output(h.hidden, h.context, **_topk_payload(3)) + h.state[session_module.PULL_FAILURE_INDEX] = -5 + if nested: + with connector.attention_step(): + output = connector.recv_ffn_output(h.hidden, context=h.context) + else: + output = connector.recv_ffn_output(h.hidden, context=h.context) + assert output.shape == h.hidden.shape + # An exception at the step boundary must restore immediate status checks. + connector.send_attn_output(h.hidden, h.context, **_topk_payload(3)) + with pytest.raises(RuntimeError, match="status -5"): + connector.recv_ffn_output(h.hidden, context=h.context) + h.state[session_module.PULL_FAILURE_INDEX] = 0 + connector.recv_ffn_output(h.hidden, context=h.context) + connector.close() + assert not connector.is_initialized + + +def test_attention_step_restores_status_checks_after_model_error(pull_attention): + h = pull_attention + with pytest.raises(ValueError, match="model failed"), h.connector.attention_step(): + raise ValueError("model failed") + h.connector.send_attn_output(h.hidden, h.context, **_topk_payload(3)) + h.state[session_module.PULL_FAILURE_INDEX] = -5 + with pytest.raises(RuntimeError, match="status -5"): + h.connector.recv_ffn_output(h.hidden, context=h.context) def test_async_routing_preserves_supplied_router_results(): diff --git a/tests/unit/connectors/test_base.py b/tests/unit/connectors/test_base.py new file mode 100644 index 000000000..d34924d8e --- /dev/null +++ b/tests/unit/connectors/test_base.py @@ -0,0 +1,19 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Host FFN lifecycle defaults shared by connector implementations.""" + +import pytest + +from afd_plugin.connectors.base import FFNExecutionSession + + +@pytest.mark.parametrize("shutdown_requested", [False, True]) +def test_default_session_stops_locally_without_drain(shutdown_requested): + session = FFNExecutionSession() + assert session.keep_serving(shutdown_requested) is not shutdown_requested + assert not session.requires_drain + session.before_first_step() + session.after_step() + session.drain(lambda: pytest.fail("Default session does not drain remote work")) + session.finish_receiver() + session.validate_shutdown() diff --git a/tests/unit/connectors/test_camp2p_connector.py b/tests/unit/connectors/test_camp2p_connector.py index 9cf18426e..d6e0e0809 100644 --- a/tests/unit/connectors/test_camp2p_connector.py +++ b/tests/unit/connectors/test_camp2p_connector.py @@ -23,13 +23,13 @@ AFDTransferState, ) from afd_plugin.connectors.metadata import AFDControlPayload, AFDDPMetadata -from afd_plugin.connectors.npu import camp2p as camp2p_module from afd_plugin.connectors.npu.camp2p import ( CAMP2pAFDConnector, CAMP2PExtraInfo, CAMP2PTransferState, build_camp2p_topology, ) +from afd_plugin.connectors.npu.camp2p import connector as camp2p_module from tests.unit.fakes.dp_metadata import DPMetadata as _FakeDPMetadata diff --git a/tests/unit/connectors/test_camp2p_token_ids.py b/tests/unit/connectors/test_camp2p_token_ids.py index adab709ce..15dd97bcc 100644 --- a/tests/unit/connectors/test_camp2p_token_ids.py +++ b/tests/unit/connectors/test_camp2p_token_ids.py @@ -25,6 +25,7 @@ import sys import types from collections.abc import Iterator +from importlib.machinery import ModuleSpec from types import SimpleNamespace from typing import Any @@ -48,6 +49,32 @@ ) +def _drop_stub_bound_modules(before: set[str], stub_objects: set[int]) -> None: + """Forget modules the stub block loaded that still hold one of its objects. + + A cached module keeps the doubles in its globals, and the parent package keeps the + child module as an attribute, so both have to be dropped for a later import to build + the real thing. + """ + + for name in [ + name + for name in sys.modules + if name not in before and name.startswith("afd_plugin") + ]: + module = sys.modules[name] + if not any( + value is not None and id(value) in stub_objects + for value in vars(module).values() + ): + continue + sys.modules.pop(name, None) + parent_name, _, leaf = name.rpartition(".") + parent = sys.modules.get(parent_name) + if parent is not None and getattr(parent, leaf, None) is module: + delattr(parent, leaf) + + @contextlib.contextmanager def _vllm_stub() -> Iterator[None]: """Expose a minimal ``vllm`` surface for the duration of one import. @@ -67,42 +94,56 @@ def _vllm_stub() -> Iterator[None]: yield return - saved = {name: sys.modules.get(name) for name in missing} - - def make(name: str, **attributes: object) -> None: + builders: dict[str, dict[str, object]] = { + "vllm": {}, + "vllm.forward_context": { + "DPMetadata": type("DPMetadata", (), {}), + "get_forward_context": lambda: None, + }, + "vllm.logger": { + "init_logger": lambda *args, **kwargs: logging.getLogger("test") + }, + "vllm.utils": {}, + "vllm.utils.torch_utils": { + "direct_register_custom_op": lambda **kwargs: None, + "is_torch_equal_or_newer": lambda *args, **kwargs: True, + }, + "vllm.distributed": {}, + "vllm.distributed.parallel_state": { + "get_pcp_group": None, + "get_tensor_model_parallel_rank": None, + }, + } + + # Only the modules that are genuinely absent are installed. Overwriting a module + # that another test already imported would outlive this context: the restore + # below only knows about the names it added. + _imported_before = set(sys.modules) + for name in missing: module = types.ModuleType(name) module.__path__ = [] # type: ignore[attr-defined] - for attribute, value in attributes.items(): + # find_spec("vllm") runs at import time; a spec-less module makes it raise. + module.__spec__ = ModuleSpec(name, None) + for attribute, value in builders[name].items(): setattr(module, attribute, value) sys.modules[name] = module - make("vllm") - make( - "vllm.forward_context", - DPMetadata=type("DPMetadata", (), {}), - get_forward_context=lambda: None, - ) - make("vllm.logger", init_logger=lambda *args, **kwargs: logging.getLogger("test")) - make("vllm.utils") - make( - "vllm.utils.torch_utils", - direct_register_custom_op=lambda **kwargs: None, - is_torch_equal_or_newer=lambda *args, **kwargs: True, - ) - make("vllm.distributed") - make( - "vllm.distributed.parallel_state", - get_pcp_group=None, - get_tensor_model_parallel_rank=None, - ) + _stub_objects = { + id(value) + for name in missing + for attribute, value in vars(sys.modules[name]).items() + if not attribute.startswith("__") and value is not None + } try: yield finally: for name in missing: sys.modules.pop(name, None) - for name, module in saved.items(): - if module is not None: - sys.modules[name] = module + # A module imported while the stub was installed keeps the doubles in its + # globals, so it must not stay cached. Only modules the stub block loaded, and + # only those that bind one of the stub objects, are dropped; evicting their + # packages would rebuild the connector factory and leave two registries. + _drop_stub_bound_modules(_imported_before, _stub_objects) with _vllm_stub(): @@ -110,12 +151,12 @@ def make(name: str, **attributes: object) -> None: AFDTransferContext, AFDTransferMetadata, ) - from afd_plugin.connectors.npu import camp2p as camp2p_module from afd_plugin.connectors.npu.camp2p import ( CAMP2pAFDConnector, prepare_token_id_transfer, received_token_ids, ) + from afd_plugin.connectors.npu.camp2p import connector as camp2p_module class _CpuTorch: @@ -326,3 +367,188 @@ def fake_a2e(*args: Any) -> tuple[Any, ...]: ] assert calls[1][-1] == 0 assert without_ids.input_ids is None + + +def _routing(connector, rows, *, columns=None): + """Routing sized the way the operator is configured.""" + + topk = connector.num_experts_per_tok if columns is None else columns + ids = torch.arange(rows * topk, dtype=torch.int64).reshape(rows, topk) + weights = torch.full((rows, topk), 0.25, dtype=torch.bfloat16) + return ids, weights + + +def test_prepare_routing_transfer_packs_int32_ids_and_float32_weights(): + """The ids channel takes int32 ids and float32 weights, not the gate dtypes.""" + + ids = torch.tensor([[0, 1], [2, 3]], dtype=torch.int64) + weights = torch.full((2, 2), 0.25, dtype=torch.bfloat16) + + wire_ids, wire_weights = camp2p_module.prepare_routing_transfer( + ids, + weights, + expected_tokens=2, + top_k=2, + ) + + assert wire_ids.dtype == torch.int32 + assert wire_ids.tolist() == [[0, 1], [2, 3]] + assert wire_weights.dtype == torch.float32 + assert wire_weights.tolist() == [[0.25, 0.25], [0.25, 0.25]] + + +def test_prepare_routing_transfer_rejects_expert_count_mismatch(): + """The wire declares ``k``, so a narrower routing tensor must not be packed.""" + + ids = torch.zeros(3, 1, dtype=torch.int32) + weights = torch.zeros(3, 1, dtype=torch.float32) + + with pytest.raises(ValueError, match="k=2"): + camp2p_module.prepare_routing_transfer( + ids, + weights, + expected_tokens=3, + top_k=2, + ) + + +def test_prepare_routing_transfer_rejects_token_count_mismatch(): + ids = torch.zeros(2, 2, dtype=torch.int32) + weights = torch.zeros(2, 2, dtype=torch.float32) + + with pytest.raises(ValueError, match="declares 3"): + camp2p_module.prepare_routing_transfer( + ids, + weights, + expected_tokens=3, + top_k=2, + ) + + +def test_received_routing_trims_operator_padded_capacity(): + """Both channel tensors are cut together, so ids and weights stay paired.""" + + ids = torch.arange(12, dtype=torch.int32).reshape(4, 3) + weights = torch.arange(12, dtype=torch.float32).reshape(4, 3) + + topk_ids, topk_weights = camp2p_module.received_routing( + ids, + weights, + expected_tokens=2, + ) + + assert topk_ids.tolist() == [[0, 1, 2], [3, 4, 5]] + assert topk_weights.tolist() == [[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]] + + +def test_received_routing_rejects_alignment_shortfall(): + ids = torch.zeros(2, 3, dtype=torch.int32) + weights = torch.zeros(2, 3, dtype=torch.float32) + + with pytest.raises(ValueError, match="routing rows"): + camp2p_module.received_routing(ids, weights, expected_tokens=3) + + +def test_send_attn_output_puts_routing_on_the_operator_ids_channel(monkeypatch): + """Attention-side routing must select the operator's ids mode and fill it.""" + + calls: list[tuple[Any, ...]] = [] + monkeypatch.setattr( + torch.ops.vllm, + "afd_camp2p_send_attn_output", + lambda *args: calls.append(args), + raising=False, + ) + monkeypatch.setattr(camp2p_module, "get_forward_context", SimpleNamespace) + connector = _connector(role="attention", rank=0) + hidden_states = torch.zeros(3, connector.hidden_size) + context = AFDTransferContext( + metadata=AFDTransferMetadata.create_attention_metadata( + layer_idx=0, + stage_idx=0, + seq_len=3, + ), + ) + ids, weights = _routing(connector, 3) + + connector.send_attn_output(hidden_states, context, routing=(ids, weights)) + + sent = calls[0] + # Trailing operator arguments: aiv_num, compute_gate, ids, scales. + assert sent[-3] == 1 + assert sent[-2].dtype == torch.int32 + assert sent[-2].tolist() == ids.tolist() + assert sent[-1].dtype == torch.float32 + assert sent[-1].tolist() == weights.to(torch.float32).tolist() + + +def test_send_attn_output_rejects_ids_and_routing_together(monkeypatch): + monkeypatch.setattr( + torch.ops.vllm, + "afd_camp2p_send_attn_output", + lambda *args: None, + raising=False, + ) + monkeypatch.setattr(camp2p_module, "get_forward_context", SimpleNamespace) + connector = _connector(role="attention", rank=0) + hidden_states = torch.zeros(3, connector.hidden_size) + context = AFDTransferContext( + metadata=AFDTransferMetadata.create_attention_metadata( + layer_idx=0, + stage_idx=0, + seq_len=3, + ), + ) + ids, weights = _routing(connector, 3) + + with pytest.raises(ValueError, match="both token ids"): + connector.send_attn_output( + hidden_states, + context, + input_ids=torch.tensor([7, 11, 13], dtype=torch.int64), + routing=(ids, weights), + ) + + +def test_recv_attn_output_reads_the_routing_slot(monkeypatch): + """The routing mode exposes the ids slot as paired top-k ids and weights.""" + + calls: list[tuple[Any, ...]] = [] + + def fake_a2e(*args: Any) -> tuple[Any, ...]: + calls.append(args) + tokens, topk = int(args[3]), int(args[5]) + ids = torch.arange(tokens * topk, dtype=torch.int32).reshape(tokens, topk) + weights = torch.full((tokens, topk), 0.5, dtype=torch.float32) + return ("hidden", ids, weights, "atten-batch", "active-mask") + + monkeypatch.setattr(torch.ops.afd_ascend, "a2e", fake_a2e, raising=False) + monkeypatch.setattr(camp2p_module, "torch", _CpuTorch()) + connector = _connector(role="ffn", rank=1) + connector.dp_metadata_list = {0: _FakeDPMetadata([2, 3, 5, 7])} + + payload = connector.recv_attn_output( + ubatch_idx=0, + layer_idx=0, + recv_routing=True, + ) + + # 10 tokens reach this FFN rank, in the same column order the gate produced. + assert calls[0][-1] == 1 + assert payload.input_ids is None + assert payload.topk_ids.shape == (10, connector.num_experts_per_tok) + assert payload.topk_ids[0].tolist() == list(range(connector.num_experts_per_tok)) + assert payload.topk_weights.shape == (10, connector.num_experts_per_tok) + assert payload.topk_weights[0].tolist() == [0.5] * connector.num_experts_per_tok + + +def test_recv_attn_output_rejects_two_wire_slot_modes(monkeypatch): + connector = _connector(role="ffn", rank=1) + + with pytest.raises(ValueError, match="exactly one"): + connector.recv_attn_output( + ubatch_idx=0, + layer_idx=0, + recv_input_ids=True, + recv_routing=True, + ) diff --git a/tests/unit/connectors/test_deepep_m2n_connector.py b/tests/unit/connectors/test_deepep_m2n_connector.py index 7e8c3cc4a..c151249aa 100644 --- a/tests/unit/connectors/test_deepep_m2n_connector.py +++ b/tests/unit/connectors/test_deepep_m2n_connector.py @@ -15,11 +15,6 @@ 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. -install_fake_afdep() - from afd_plugin.config import afd_config_from_mapping # noqa: E402 from afd_plugin.connectors import ( # noqa: E402 AFDConnectorFactory, @@ -29,23 +24,36 @@ AFDTransferMetadata, ) from afd_plugin.connectors.base import AFDControlPlaneClosedError # noqa: E402 -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, - plan_stage, - precompile_buckets, - resolve_capacity_tokens_per_rank, - resolve_num_lanes, - tokens_per_attention_rank, - validate_handshake, - validate_model_infos, +from tests.unit.fakes.afdep import ( # noqa: E402 + FakeElasticBuffer, + install_fake_afdep, + uninstall_fake_afdep, ) + +# The connector module imports afdEP's transport when it loads, so the stand-in has to +# exist for that import only. It is removed again on both paths: the bound names keep +# working, and the rest of the session sees the package as absent. +install_fake_afdep() +try: + 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, + plan_stage, + precompile_buckets, + resolve_capacity_tokens_per_rank, + resolve_num_lanes, + tokens_per_attention_rank, + validate_handshake, + validate_model_infos, + ) +finally: + uninstall_fake_afdep() from afd_plugin.connectors.gpu.deepep_m2n.experts import ( # noqa: E402 expert_row_ranges, ) @@ -53,7 +61,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, @@ -577,9 +589,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) @@ -592,18 +606,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()) @@ -827,6 +877,38 @@ def abort(group): connector.control_plane.recv_dp_metadata_list() +def test_close_destroys_the_distinct_transport_group_after_lanes(monkeypatch): + """A per-domain transport group is destroyed once, after its lanes. + + Why: afdEP-comm leaves the caller-owned Gloo group alive; keeping it + registered leaks process-group state across connector lifetimes. + """ + events = [] + + class Lane: + def destroy(self): + events.append("lane") + + monkeypatch.setattr( + connector_module, "_abort_process_group", lambda group: events.append("control") + ) + monkeypatch.setattr( + connector_module.dist, + "destroy_process_group", + lambda group: events.append("transport"), + ) + connector = _connector("ffn", 0, M2NTopology(ffn_size=2, attn_size=2)) + connector.group = SimpleNamespace() + connector.bootstrap_group = SimpleNamespace() + connector.transport_group = SimpleNamespace() + connector.lanes = [Lane()] + + connector.close() + connector.close() + + assert events == ["control", "lane", "transport"] + + @pytest.mark.parametrize( ("closes_first", "error", "expected"), [ @@ -1166,27 +1248,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/connectors/test_external_routing.py b/tests/unit/connectors/test_external_routing.py new file mode 100644 index 000000000..18dcfa66c --- /dev/null +++ b/tests/unit/connectors/test_external_routing.py @@ -0,0 +1,137 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""CPU tests for the Attention-side routing payload. + +The payload is pure tensor logic, so it is testable without an Ascend device. The +package resolves its connector lazily, so importing this module does not pull vLLM in. + +Rows carry rank- and position-specific values instead of a constant fill: a gather +or split that drops, reorders or mismatches a row has to change the observed +contents, not just the shape. +""" + +from __future__ import annotations + +import pytest + +torch = pytest.importorskip("torch") + +_TOP_K = 2 + + +@pytest.fixture(scope="module") +def routing_module(): + from afd_plugin.connectors.npu.camp2p import external_routing + + return external_routing + + +@pytest.fixture(scope="module") +def inputs_type(routing_module): + return routing_module.AFDExternalRoutingInputs + + +def _inputs(inputs_type, rows: int, rank: int = 0): + """Rows encode rank and position, so an ordering bug changes the values.""" + + base = rank * 100 + values = torch.arange(base, base + rows * _TOP_K) + return inputs_type( + topk_ids=values.to(torch.int32).reshape(rows, _TOP_K), + topk_weights=(values.to(torch.float32) / 10).reshape(rows, _TOP_K), + ) + + +_RANK_STRIDE = 100 + + +def _two_rank_gather(tensor): + """Stand-in collective whose ranks carry *different* rows. + + A real DP group gathers each rank's own tokens, so rank 1's shard is offset from + rank 0's. Repeating one shard would hide an implementation that reorders or drops + a rank. + """ + + return torch.cat([tensor, tensor + _RANK_STRIDE], dim=0) + + +def test_padding_appends_inert_rows_and_keeps_valid_ones(inputs_type): + inputs = _inputs(inputs_type, 2) + + padded = inputs.padded(2) + + assert padded.topk_ids[:2].tolist() == inputs.topk_ids.tolist() + assert padded.topk_weights[:2].tolist() == inputs.topk_weights.tolist() + assert padded.topk_ids[2:].tolist() == [[0, 0], [0, 0]] + assert padded.topk_weights[2:].tolist() == [[0.0, 0.0], [0.0, 0.0]] + assert padded.token_count == 4 + + +def test_gathering_keeps_rank_order_and_contents(inputs_type): + """Rank 0's rows come first, then rank 1's, with their own values intact.""" + + inputs = _inputs(inputs_type, 2, rank=0) + + gathered = inputs.gathered(_two_rank_gather) + + assert gathered.topk_ids.tolist() == ( + inputs.topk_ids.tolist() + (inputs.topk_ids + _RANK_STRIDE).tolist() + ) + assert gathered.topk_weights.tolist() == ( + inputs.topk_weights.tolist() + (inputs.topk_weights + _RANK_STRIDE).tolist() + ) + + +def test_split_rows_takes_one_rank_slice(inputs_type): + inputs = _inputs(inputs_type, 4) + + taken = inputs.split_rows(lambda tensor: torch.tensor_split(tensor, 2, dim=0)[1]) + + assert taken.topk_ids.tolist() == inputs.topk_ids.tolist()[2:] + assert taken.topk_weights.tolist() == inputs.topk_weights.tolist()[2:] + + +def test_non_matrix_routing_is_rejected(inputs_type): + with pytest.raises(ValueError, match="2-D"): + inputs_type( + topk_ids=torch.zeros(3, dtype=torch.int32), + topk_weights=torch.zeros(3, dtype=torch.float32), + ) + + +def test_mismatched_topk_shapes_are_rejected(inputs_type): + with pytest.raises(ValueError, match="same shape"): + inputs_type( + topk_ids=torch.zeros(2, 2, dtype=torch.int32), + topk_weights=torch.zeros(2, 3, dtype=torch.float32), + ) + + +@pytest.mark.parametrize("dtype", [torch.int32, torch.int64]) +def test_expert_indices_accept_both_index_dtypes(inputs_type, dtype): + """Both index widths are in the operator contract; neither may be rejected.""" + + inputs = inputs_type( + topk_ids=torch.zeros(2, 2, dtype=dtype), + topk_weights=torch.zeros(2, 2, dtype=torch.float32), + ) + + assert inputs.topk_ids.dtype == dtype + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bool]) +def test_expert_indices_reject_non_index_dtypes(inputs_type, dtype): + with pytest.raises(ValueError, match="expert indices"): + inputs_type( + topk_ids=torch.zeros(2, 2, dtype=dtype), + topk_weights=torch.zeros(2, 2, dtype=torch.float32), + ) + + +def test_non_float_weights_are_rejected(inputs_type): + with pytest.raises(ValueError, match="gate weights"): + inputs_type( + topk_ids=torch.zeros(2, 2, dtype=torch.int32), + topk_weights=torch.zeros(2, 2, dtype=torch.int32), + ) diff --git a/tests/unit/connectors/test_p2p_connector.py b/tests/unit/connectors/test_p2p_connector.py index dc6cb70e1..b4a3edbda 100644 --- a/tests/unit/connectors/test_p2p_connector.py +++ b/tests/unit/connectors/test_p2p_connector.py @@ -897,3 +897,50 @@ def test_p2p_subgroup_rendezvous_reuses_the_afd_world_store( def test_factory_rejects_unknown_connector(): with pytest.raises(ValueError, match="unsupported AFD connector type"): AFDConnectorFactory.get_connector_class("UnknownConnector") + + +class _LazilyExportedConnector(AFDConnectorBase): + """Stands in for a connector served only through module ``__getattr__``.""" + + def __init__(self, *args, **kwargs) -> None: # pragma: no cover - never built + raise NotImplementedError + + +@pytest.fixture +def isolated_connector_registry(monkeypatch): + """Keep a temporary registration out of the registry other tests share.""" + + monkeypatch.setattr( + AFDConnectorFactory, + "_registry", + dict(AFDConnectorFactory._registry), + ) + + +def test_factory_resolves_a_connector_served_only_by_package_getattr( + monkeypatch, isolated_connector_registry +): + """The factory must read lazily exported names out of ``__getattr__``.""" + + module = types.ModuleType("afd_lazily_exported_connector") + class_name = "_LazilyExportedConnector" + + def _resolve(name: str): + if name == class_name: + return _LazilyExportedConnector + raise AttributeError(name) + + # A module-level __getattr__ (PEP 562) is what the CAMP2P package uses. + module.__getattr__ = _resolve # type: ignore[method-assign] + monkeypatch.setitem(sys.modules, module.__name__, module) + + AFDConnectorFactory.register_connector( + "LazilyExportedConnector", + module.__name__, + class_name, + replace=True, + ) + + assert class_name not in vars(module), "the class must not be in the module dict" + resolved = AFDConnectorFactory.get_connector_class("LazilyExportedConnector") + assert resolved is _LazilyExportedConnector diff --git a/tests/unit/fakes/afdep.py b/tests/unit/fakes/afdep.py index aaf564877..f7255da8f 100644 --- a/tests/unit/fakes/afdep.py +++ b/tests/unit/fakes/afdep.py @@ -10,11 +10,12 @@ from __future__ import annotations import sys +from importlib.machinery import ModuleSpec from importlib.util import find_spec 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 v1.1.0). 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 +29,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: @@ -46,6 +53,9 @@ def __init__(self, **kwargs: object) -> None: self.kwargs = kwargs +_FAKE_MARKER = "__afd_fake_afdep__" + + 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: @@ -56,5 +66,23 @@ def install_fake_afdep() -> None: deepep.DispatchHandle = type("DispatchHandle", (), {}) # type: ignore[attr-defined] deepep.ReceiveResult = type("ReceiveResult", (), {}) # type: ignore[attr-defined] package.deepep = deepep # type: ignore[attr-defined] + # A module without a spec makes a later find_spec("afdep") raise instead of + # reporting the package as missing. + package.__spec__ = ModuleSpec("afdep", None) + deepep.__spec__ = ModuleSpec("afdep.deepep", None) + setattr(package, _FAKE_MARKER, True) + setattr(deepep, _FAKE_MARKER, True) sys.modules["afdep"] = package sys.modules["afdep.deepep"] = deepep + + +def uninstall_fake_afdep() -> None: + """Remove the stand-in again, leaving a real afdEP install untouched. + + A stand-in left in ``sys.modules`` has no ``__spec__``, so a later + ``find_spec("afdep")`` raises instead of reporting "not installed". + """ + for name in ("afdep.deepep", "afdep"): + module = sys.modules.get(name) + if module is not None and getattr(module, _FAKE_MARKER, False): + sys.modules.pop(name, None) diff --git a/tests/unit/fakes/async_cam.py b/tests/unit/fakes/async_cam.py new file mode 100644 index 000000000..9996be0a6 --- /dev/null +++ b/tests/unit/fakes/async_cam.py @@ -0,0 +1,234 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Host recording boundary for CAM native messages; no device simulation.""" + +import sys +from inspect import Parameter, Signature +from types import SimpleNamespace + +from afd_plugin.config import AFDConfig + + +class _FakeTensor: + def __init__(self, shape, *, dtype="bf16", device="npu:0"): + self.shape = (shape,) if isinstance(shape, int) else tuple(shape) + self.dtype = dtype + self.device = device + + def to(self, *, dtype): + self.dtype = dtype + return self + + def contiguous(self): + return self + + def new_zeros(self, shape): + return _FakeTensor(shape, dtype=self.dtype, device=self.device) + + def new_empty(self, shape): + return _FakeTensor(shape, dtype=self.dtype, device=self.device) + + +class FakePullState(_FakeTensor): + """Host status words at the native tensor boundary; no device execution.""" + + def __init__(self): + super().__init__((128,), dtype="fp16") + self.words = {32: -1} + + def view(self, dtype): + return self + + def __getitem__(self, index): + return SimpleNamespace( + item=lambda: self.words.get(index, 0), + fill_=lambda value: self.words.__setitem__(index, value), + ) + + +class _FakeCamOps: + def __init__(self): + self.calls = [] + + def _record(self, name, fields, args, kwargs): + # Native schema binding accepts either positional or named arguments. + signature = Signature( + [ + Parameter(field, Parameter.POSITIONAL_OR_KEYWORD) + for field in fields.split() + ] + ) + values = tuple(signature.bind(*args, **kwargs).arguments.values()) + self.calls.append((name, values)) + return values + + def afd_mega_pull_init(self, *args): + self.calls.append(("pull_init", args)) + return FakePullState() + + def afd_async_dispatch_send(self, *args, **kwargs): + values = self._record( + "dispatch_send", + "x expert_ids comm_args comm_id max_seq_len batch_size hidden_size " + "top_k moe_rank_num attn_rank_num route_expert_num_per_moe " + "attn_rank_id world_size layer_index tp_size dynamic_quant group_name", + args, + kwargs, + ) + return values[0] + + def afd_async_dispatch_recv(self, *args, **kwargs): + values = self._record( + "dispatch_recv", + "x comm_args comm_id max_seq_len hidden_size top_k moe_rank_num " + "attn_rank_num route_expert_num_per_moe moe_rank_id world_size " + "tp_size dynamic_quant group_name", + args, + kwargs, + ) + batch_size, hidden_size, expert_per_rank, tp_size = ( + values[3], + values[4], + values[8], + values[11], + ) + return ( + _FakeTensor((batch_size, hidden_size)), + _FakeTensor((batch_size,), dtype="fp32"), + _FakeTensor((5 + tp_size * (1 + expert_per_rank),), dtype="int64"), + _FakeTensor((expert_per_rank,), dtype="int64"), + ) + + def afd_async_combine_send(self, *args, **kwargs): + values = self._record( + "combine_send", + "expand_x comm_args batch_info comm_id max_seq_len hidden_size " + "top_k moe_rank_num attn_rank_num route_expert_num_per_moe " + "moe_rank_id world_size tp_size group_name", + args, + kwargs, + ) + return values[0] + + def afd_async_combine_recv(self, *args, **kwargs): + values = self._record( + "combine_recv", + "expand_x expert_ids expert_scales comm_args comm_id batch_size " + "hidden_size top_k moe_rank_num attn_rank_num route_expert_num_per_moe " + "attn_rank_id world_size group_name", + args, + kwargs, + ) + return _FakeTensor((values[5], values[6])) + + def afd_mega_dispatch_send(self, *args, **kwargs): + result = self.afd_async_dispatch_send(*args, **kwargs) + self.calls[-1] = ("mega_dispatch_send", self.calls[-1][1]) + return result + + def afd_mega_combine_recv(self, *args, max_seq_len, tp_size, **kwargs): + result = self.afd_async_combine_recv(*args, **kwargs) + self.calls[-1] = ( + "mega_combine_recv", + self.calls[-1][1] + (max_seq_len, tp_size), + ) + return result + + +class _FakeTorch: + def __init__(self): + self.bfloat16 = "bf16" + self.float16 = "fp16" + self.float32 = "fp32" + self.int32 = "int32" + self.int64 = "int64" + self.ops = SimpleNamespace(afd_ascend=_FakeCamOps()) + self.npu = SimpleNamespace(synchronize=lambda: None) + + def device(self, name): + return name + + def empty(self, shape, *, dtype, device): + return _FakeTensor(shape, dtype=dtype, device=device) + + def zeros(self, shape, *, dtype, device): + return _FakeTensor(shape, dtype=dtype, device=device) + + +def cam_vllm_config(*, tp_size=1, pcp_size=1, extra_config=None, dtype="bf16"): + return SimpleNamespace( + additional_config={ + "afd": { + "connector_extra_config": extra_config + if extra_config is not None + else {"attn_ranks_per_dp": 4}, + }, + }, + parallel_config=SimpleNamespace( + data_parallel_size=1, + data_parallel_rank=0, + prefill_context_parallel_size=pcp_size, + tensor_parallel_size=tp_size, + ), + scheduler_config=SimpleNamespace(max_num_batched_tokens=16), + model_config=SimpleNamespace( + dtype=dtype, + hf_config=SimpleNamespace( + hidden_size=16, + num_experts_per_tok=2, + n_routed_experts=8, + ), + ), + ) + + +def _afd_config(*, role: str): + return AFDConfig( + connector="CAMAsyncAFDConnector", + role=role, + num_attention_ranks=4, + num_ffn_ranks=2, + ) + + +def _dp2_afd_config(*, role: str): + return AFDConfig( + connector="CAMAsyncAFDConnector", + role=role, + num_attention_ranks=8, + num_ffn_ranks=8, + ) + + +def _topk_payload(batch_size: int, topk: int = 2): + return { + "topk_ids": _FakeTensor((batch_size, topk), dtype="int32"), + "topk_weights": _FakeTensor((batch_size, topk), dtype="fp32"), + } + + +def recording_cam_connector( + monkeypatch, module, *, role="attention", mega=False, tp_size=4, pcp_size=1 +): + fake_torch = _FakeTorch() + monkeypatch.setattr(module, "torch", fake_torch) + monkeypatch.setitem(sys.modules, "torch", fake_torch) + connector = module.CAMAsyncAFDConnector( + 0, + 0, + cam_vllm_config( + pcp_size=pcp_size, + extra_config={ + "attn_ranks_per_dp": tp_size, + "ffn_mega_moe": mega, + "dynamicQuant": 1, + }, + ), + _afd_config(role=role), + 0, + ) + monkeypatch.setattr(connector.graph_startup, "phase", lambda phase: None) + connector.ffn_session().initialize("npu:0", "recorded") + connector._initialized = True + fake_torch.ops.afd_ascend.calls.clear() + return SimpleNamespace(connector=connector, torch=fake_torch) diff --git a/tests/unit/fakes/async_cam_startup.py b/tests/unit/fakes/async_cam_startup.py new file mode 100644 index 000000000..89642b8e7 --- /dev/null +++ b/tests/unit/fakes/async_cam_startup.py @@ -0,0 +1,44 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""In-memory Store and rank world for the host startup protocol.""" + +from threading import Condition + +from afd_plugin.connectors.npu.async_cam_runtime.async_cam_startup import ( + AsyncCAMGraphStartup, +) + + +class MemoryStore: + def __init__(self): + self.values = {} + self.condition = Condition() + self.closed = False + + def set(self, key, value): + assert not self.closed, "Store used after release" + with self.condition: + self.values[key] = value.encode() if isinstance(value, str) else value + self.condition.notify_all() + + def get(self, key): + assert not self.closed, "Store used after release" + with self.condition: + assert self.condition.wait_for(lambda: key in self.values, timeout=3) + return self.values[key] + + def check(self, keys): + assert not self.closed, "Store used after release" + with self.condition: + return all(key in self.values for key in keys) + + +def startup_world(attn=2, ffn=2, store=None): + store = store or MemoryStore() + owners = [] + for rank in range(attn + ffn): + owner = AsyncCAMGraphStartup(rank, attn, attn + ffn) + owner._store = store + owner._prefix = "ffn_graph/test/" + owners.append(owner) + return owners diff --git a/tests/unit/fakes/async_ffn.py b/tests/unit/fakes/async_ffn.py index 838084860..b87b880e5 100644 --- a/tests/unit/fakes/async_ffn.py +++ b/tests/unit/fakes/async_ffn.py @@ -46,14 +46,13 @@ def stub(name, **attributes): "afd_plugin.v1.worker.ffn_metadata", ): stub(name) - stub( - "afd_plugin.connectors.base", AFDConnectorBase=object, ConnectorExtraInfo=object - ) - stub("afd_plugin.connectors.metadata", AFDTransferState=object) + stub("vllm.forward_context") + load_host_module(monkeypatch, "afd_plugin.connectors.metadata") + load_host_module(monkeypatch, "afd_plugin.connectors.base") + load_host_module(monkeypatch, "tests.unit.fakes.connector") group = load_host_module(monkeypatch, "afd_plugin.distributed.afd_process_group") stub("afd_plugin.distributed", init_afd_process_group=group.init_afd_process_group) connector = load_host_module(monkeypatch, "afd_plugin.connectors.npu.async_cam") - load_host_module(monkeypatch, "afd_plugin.v1.worker.npu.async_ffn_two_stream") runtime = load_host_module( monkeypatch, "afd_plugin.v1.worker.npu.async_ffn_model_runner" ) diff --git a/tests/unit/fakes/connector.py b/tests/unit/fakes/connector.py index 7fc57c220..7222e2270 100644 --- a/tests/unit/fakes/connector.py +++ b/tests/unit/fakes/connector.py @@ -8,6 +8,7 @@ from __future__ import annotations +from collections.abc import Callable, Iterable from types import SimpleNamespace from typing import Any @@ -17,6 +18,41 @@ AFDTransferContext, AFDTransferMetadata, ) +from afd_plugin.connectors.base import FFNExecutionSession + + +class FakeFFNSession(FFNExecutionSession): + """Configurable host lifecycle seam, with no device behavior.""" + + def __init__( + self, + *, + serving: Iterable[bool] = (), + requires_drain: bool = False, + drain_error: Exception | None = None, + transactions: int = 0, + ) -> None: + self.requires_drain = requires_drain + self.serving = iter(serving) + self.drain_error = drain_error + self.transactions = transactions + self.receiver_finished = False + + def keep_serving(self, shutdown_requested: bool) -> bool: + return next(self.serving, not shutdown_requested) + + def drain(self, run_one: Callable[[], None]) -> None: + for _ in range(self.transactions): + run_one() + if self.drain_error is not None: + raise self.drain_error + + def finish_receiver(self) -> None: + self.receiver_finished = True + + def validate_shutdown(self) -> None: + if self.requires_drain and not self.receiver_finished: + raise RuntimeError("receiver has not finished; retain resources") class FakeRoutedConnector(AFDRoutedExpertsConnector): diff --git a/tests/unit/fakes/ffn_faas.py b/tests/unit/fakes/ffn_faas.py index 861f6e8e3..de1b0d997 100644 --- a/tests/unit/fakes/ffn_faas.py +++ b/tests/unit/fakes/ffn_faas.py @@ -6,9 +6,9 @@ import atexit import json -import os import sys import threading +from argparse import ArgumentParser, Namespace from contextlib import ExitStack from dataclasses import dataclass from pathlib import Path @@ -72,18 +72,59 @@ def write_v4_checkpoint(root: Path, *, layers: int = 2, experts: int = 4) -> Pat return root +# The shape of Qwen/Qwen3-30B-A3B's config.json (no shared-expert field). +QWEN3_30B_A3B_CONFIG = { + "architectures": ["Qwen3MoeForCausalLM"], + "model_type": "qwen3_moe", + "torch_dtype": "bfloat16", + "hidden_size": 2048, + "num_hidden_layers": 48, + "num_experts": 128, + "num_experts_per_tok": 8, + "moe_intermediate_size": 768, + "decoder_sparse_step": 1, + "mlp_only_layers": [], + "norm_topk_prob": True, +} + + +def write_qwen3_checkpoint(root: Path, **overrides) -> Path: + """A Qwen3 MoE ``config.json`` only; ``None`` values remove a key.""" + root.mkdir(parents=True, exist_ok=True) + config = {**QWEN3_30B_A3B_CONFIG, **overrides} + config = {key: value for key, value in config.items() if value is not None} + (root / "config.json").write_text(json.dumps(config)) + return root + + class HostBackend: """Exercise host ownership rules, without simulating device completion.""" def __init__( - self, *, stop_during_prepare: bool = False, fail: str = "", block: bool = False + self, + *, + stop_during_prepare: bool = False, + fail: str = "", + block: bool = False, + cancellable: bool = False, + end_receive: bool = False, ) -> None: + """Configure one host-contract behavior. + + ``block`` keeps ``execute_step`` waiting like an idle receive; + ``cancellable`` lets ``cancel_receive`` end that wait the way the GPU + M2N control plane does; ``end_receive`` makes the first step report a + cancelled receive without any host stop. + """ self.stop_during_prepare = stop_during_prepare self.fail = fail self.block = block + self.cancellable = cancellable + self.end_receive = end_receive self.entered = threading.Event() self.release = threading.Event() self.finished = threading.Event() + self.cancelled = False self.released = False self.service: FFNService | None = None @@ -99,15 +140,27 @@ def start_receiver(self) -> None: if self.fail == "activate": raise RuntimeError("receiver activation failed") - def execute_step(self) -> None: + def execute_step(self) -> bool: + """Return ``False`` for a cancelled receive, else request a stop.""" assert self.service is not None self.entered.set() + if self.end_receive: + return False if self.block: self.release.wait() + if self.cancelled: + return False if self.fail == "execute": self.finished.set() raise RuntimeError("execute failed") self.service.request_stop() + return True + + def cancel_receive(self) -> None: + """End a blocked receive only when this fake models a cancellable one.""" + if self.cancellable: + self.cancelled = True + self.release.set() def receiver_exited(self) -> None: self.finished.set() @@ -123,16 +176,16 @@ def close(self) -> None: class HostPlatformAdapter: """An independent CPU host contract; it has no model/device requirements.""" - def add_arguments(self, parser) -> None: - pass + def create_parser(self) -> ArgumentParser: + return ArgumentParser() def configure(self, args) -> FFNLaunch: return FFNLaunch( rank=3, local_rank=0, - model_path=str(args.model), + model_path=str(args.model_tag), backend_name="host-contract", - shutdown_timeout=args.shutdown_timeout, + shutdown_timeout=args.receiver_shutdown_timeout, backend_factory=HostBackend, ) @@ -151,7 +204,7 @@ def run_cpu_platform(model: str, status: str) -> int: "0", "--status-file", status, - "--shutdown-timeout", + "--receiver-shutdown-timeout", "0.1", ] with ( @@ -171,24 +224,12 @@ def run_cli_host_case(model: str, marker: str, fault: str, released: bool) -> in from afd_plugin.ffn_faas import __main__ as cli atexit.register(lambda: Path(marker).write_text("ordinary shutdown")) - os.environ.update( - RANK="0", - LOCAL_RANK="0", - WORLD_SIZE="1", - LOCAL_WORLD_SIZE="1", - VLLM_USE_V2_MODEL_RUNNER="1", - AFD_ASYNC_CAM_LAYERED_GMM="1", - ) args = [ "ffn_faas", "serve", model, - "--attention-ranks", - "1", - "--attention-tp", - "1", - "--ffn-ranks", - "1", + "--backend", + "host-contract", "--status-port", "0", ] @@ -244,7 +285,39 @@ def run(self): if not released and fault != "returned_unreleased": raise RuntimeError("host-only unreleased session") + patches.enter_context( + patch.object(cli, "SUPPORTED_BACKENDS", ("host-contract",)) + ) + patches.enter_context( + patch.object(cli, "get_platform_adapter", lambda _: HostPlatformAdapter()) + ) patches.enter_context(patch.object(cli, "FFNService", HostService)) patches.enter_context(patch.object(cli, "FFNModelRunner", HostRunner)) patches.enter_context(patch.object(sys, "argv", args)) return cli.main() + + +def cam_cli_args(model: Path, **overrides) -> Namespace: + """Parsed engine arguments for CPU-only standalone launch admission tests.""" + values = { + "model": str(model), + "model_tag": None, + "data_parallel_size": 2, + "data_parallel_rank": None, + "distributed_executor_backend": "external_launcher", + "receiver_shutdown_timeout": 5.0, + "additional_config": { + "enable_cpu_binding": False, + "afd": { + "connector": "CAMAsyncAFDConnector", + "role": "ffn", + "async_dp": True, + "compute_gate_on_attention": True, + "num_attention_ranks": 12, + "num_ffn_ranks": 2, + "connector_extra_config": {"dynamicQuant": 1, "attn_ranks_per_dp": 4}, + }, + }, + } + values.update(overrides) + return Namespace(**values) diff --git a/tests/unit/fakes/npu_attention.py b/tests/unit/fakes/npu_attention.py new file mode 100644 index 000000000..80e8ae9d0 --- /dev/null +++ b/tests/unit/fakes/npu_attention.py @@ -0,0 +1,54 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Native Ascend metadata seams without device allocation.""" + +from types import SimpleNamespace + +from tests.unit.fakes.vllm_config import npu_vllm_config + + +def metadata_runner(): + from afd_plugin.v1.worker.npu.attention_model_runner_v2 import ( + AFDNPUAttentionModelRunnerV2, + ) + + runner = object.__new__(AFDNPUAttentionModelRunnerV2) + runner.vllm_config = npu_vllm_config() + runner.vllm_config.kv_transfer_config = None + runner.speculative_config = SimpleNamespace(enforce_eager=False) + runner.connector = SimpleNamespace(control_plane=None) + runner._afd_pending_metadata = None + runner._afd_suppress_metadata_send = False + runner._afd_transaction_counter = 0 + runner._is_warmup = runner._afd_is_profile = False + runner._afd_is_graph_capturing = runner._afd_is_graph_replaying = False + return runner + + +def forward_context(num_tokens): + return SimpleNamespace( + batch_descriptor=SimpleNamespace(num_tokens=num_tokens), + attn_metadata={}, + additional_kwargs={}, + ubatch_slices=None, + ) + + +class DraftGraphManager: + def __init__(self, runner, descs=(), *, fail=False): + from vllm.config import CUDAGraphMode + + self.runner = runner + self._capture_descs = {CUDAGraphMode.FULL: list(descs)} + self._max_full_descs_to_capture = None + self.fail = fail + self.emitted = [] + + def run_fullgraph(self, desc): + context = forward_context(desc.num_tokens) + self.runner.install_afd_metadata_on_forward_context(context) + metadata = context.additional_kwargs["afd_metadata"] + self.emitted.append(metadata) + if self.fail: + raise RuntimeError("native replay failed") + return metadata diff --git a/tests/unit/fakes/vllm_config.py b/tests/unit/fakes/vllm_config.py index e67a400da..c374efebd 100644 --- a/tests/unit/fakes/vllm_config.py +++ b/tests/unit/fakes/vllm_config.py @@ -45,6 +45,7 @@ def npu_parallel_config(**overrides): "pipeline_parallel_size": 1, "use_sequence_parallel_moe": False, "enable_eplb": False, + "enable_elastic_ep": False, "prefill_context_parallel_size": 1, "decode_context_parallel_size": 1, "dbo_decode_token_threshold": 1, @@ -101,10 +102,34 @@ def npu_vllm_config( fast_moe_cold_start=False, ), speculative_config=speculative_config, + kv_transfer_config=None, scheduler_config=SimpleNamespace(max_num_batched_tokens=16), ) +def npu_async_cam_target_graph_config(**parallel_overrides): + config = npu_vllm_config( + connector="CAMAsyncAFDConnector", + async_dp=True, + compute_gate_on_attention=True, + cudagraph_mode="FULL_DECODE_ONLY", + extra_config={ + "dynamicQuant": 1, + "attn_ranks_per_dp": parallel_overrides.get("tensor_parallel_size", 1), + }, + enable_expert_parallel=True, + **parallel_overrides, + ) + config.model_config.enforce_eager = False + config.model_config.hf_config = config.model_config.hf_text_config = ( + SimpleNamespace( + model_type="deepseek_v4", architectures=["DeepseekV4ForCausalLM"] + ) + ) + config.compilation_config.mode = 0 + return config + + def async_moe_config( *, role="attention", @@ -253,7 +278,7 @@ def m2n_vllm_config( ) -def npu_dspark_config(*, role="attention", draft_eager=True): +def npu_dspark_config(*, role="attention", draft_eager=True, tp_size=1): hf = SimpleNamespace( architectures=["DSparkDraftModel"], vocab_size=16, @@ -266,12 +291,14 @@ def npu_dspark_config(*, role="attention", draft_eager=True): connector="CAMAsyncAFDConnector", async_dp=True, compute_gate_on_attention=True, - extra_config={"dynamicQuant": 1}, + tensor_parallel_size=tp_size, + num_attention_ranks=tp_size, + extra_config={"dynamicQuant": 1, "attn_ranks_per_dp": tp_size}, speculative_config=SimpleNamespace( method="dspark", enforce_eager=draft_eager, num_speculative_tokens=5, - draft_parallel_config=SimpleNamespace(tensor_parallel_size=1), + draft_parallel_config=SimpleNamespace(tensor_parallel_size=tp_size), draft_model_config=SimpleNamespace(hf_config=hf, hf_text_config=hf), ), ) diff --git a/tests/unit/ffn_faas/test_v4_model.py b/tests/unit/ffn_faas/platforms/ffn_model/deepseek_v4/npu/test_v4_model.py similarity index 92% rename from tests/unit/ffn_faas/test_v4_model.py rename to tests/unit/ffn_faas/platforms/ffn_model/deepseek_v4/npu/test_v4_model.py index dd62da738..f642349cc 100644 --- a/tests/unit/ffn_faas/test_v4_model.py +++ b/tests/unit/ffn_faas/platforms/ffn_model/deepseek_v4/npu/test_v4_model.py @@ -8,15 +8,16 @@ import pytest +from afd_plugin.ffn_faas.platforms.ffn_model.deepseek_v4.npu.v4_model import ( + V4CheckpointDescription, + V4ModelAdapter, +) from afd_plugin.ffn_faas.platforms.npu import get_model from afd_plugin.ffn_faas.platforms.npu.environment import AscendFFNEnvironment from afd_plugin.ffn_faas.platforms.npu.settings import ( FFNConfig, ) -from afd_plugin.ffn_faas.platforms.npu.v4_model import ( - V4CheckpointDescription, - V4ModelAdapter, -) +from tests.unit.fakes.vllm_config import npu_vllm_config def test_main_routed_metadata_without_parameter_allocation(checkpoint_dir): @@ -91,16 +92,16 @@ def test_model_selection_rejects_other_quantization_before_loading(checkpoint_di def test_wrong_compute_world_rejected_before_native_imports( checkpoint_dir, monkeypatch ): - settings = FFNConfig(checkpoint_dir, attention_ranks=1, attention_tp=1, ffn_ranks=2) + settings = FFNConfig(checkpoint_dir, ffn_ranks=2) environment = AscendFFNEnvironment(settings) monkeypatch.setenv("RANK", "0") monkeypatch.setenv("WORLD_SIZE", "1") with pytest.raises(ValueError, match="differs from the F-only"): - environment.initialize_device() + environment.initialize_device(npu_vllm_config(role="ffn")) def test_unused_compute_owner_can_close_without_importing_runtime(checkpoint_dir): - settings = FFNConfig(checkpoint_dir, attention_ranks=1, attention_tp=1, ffn_ranks=1) + settings = FFNConfig(checkpoint_dir, ffn_ranks=1) environment = AscendFFNEnvironment(settings) environment.close() with pytest.raises(RuntimeError, match="device has not been initialized"): diff --git a/tests/unit/ffn_faas/test_gpu_adapter.py b/tests/unit/ffn_faas/test_gpu_adapter.py new file mode 100644 index 000000000..f2113e605 --- /dev/null +++ b/tests/unit/ffn_faas/test_gpu_adapter.py @@ -0,0 +1,205 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""GPU DeepEP M2N launch validation happens on the CPU, before any device.""" + +from __future__ import annotations + +import json +import os +import subprocess +import sys +from pathlib import Path + +import pytest + +from afd_plugin.ffn_faas.platforms import get_platform_adapter + +REPO_ROOT = Path(__file__).resolve().parents[3] +NATIVE_MODULES = ("torch", "vllm", "afdep", "deep_gemm") +GRAPH = ("--compilation-config", '{"cudagraph_mode":"FULL_DECODE_ONLY"}') + + +def _afd_config(**overrides) -> str: + afd = { + "role": "ffn", + "connector": "DeepEPM2NAFDConnector", + "port": 6239, + "num_attention_ranks": 2, + "num_ffn_ranks": 2, + "compute_gate_on_attention": True, + **overrides, + } + return json.dumps({"afd": afd}) + + +@pytest.fixture +def launch_env(monkeypatch): + """A two-rank torchrun environment for F rank 1.""" + for key, value in {"RANK": "1", "LOCAL_RANK": "1", "WORLD_SIZE": "2"}.items(): + monkeypatch.setenv(key, value) + + +@pytest.fixture +def configure(tmp_path): + """Parse a native ``serve`` command line and configure the GPU adapter.""" + pytest.importorskip("vllm") + from afd_plugin.ffn_faas.__main__ import create_parser + + model = tmp_path / "model" + model.mkdir() + + def run(*options, afd=None, data_parallel_size="2"): + adapter = get_platform_adapter("gpu-deepep-m2n") + args = create_parser(adapter).parse_args( + [ + "serve", + str(model), + "--backend", + "gpu-deepep-m2n", + "--data-parallel-size", + data_parallel_size, + "--enable-expert-parallel", + "--additional-config", + _afd_config(**(afd or {})), + *options, + ] + ) + return adapter.configure(args), args + + return run + + +def test_backend_selection_imports_no_native_stack(): + script = ( + "import json, sys; " + "from afd_plugin.ffn_faas.platforms import get_platform_adapter; " + "get_platform_adapter('gpu-deepep-m2n'); " + f"print(json.dumps([m for m in {NATIVE_MODULES!r} if m in sys.modules]))" + ) + result = subprocess.run( + [sys.executable, "-c", script], + cwd=REPO_ROOT, + capture_output=True, + text=True, + timeout=60, + check=True, + ) + assert json.loads(result.stdout) == [] + + +def test_valid_launch_preserves_the_model_path_and_describes_the_rank( + launch_env, configure, tmp_path +): + launch, args = configure(*GRAPH, "--enable-dbo") + + assert (launch.rank, launch.local_rank) == (1, 1) + assert launch.model_path == str((tmp_path / "model").resolve()) + assert launch.backend_name == "DeepEPM2NAFDConnector" + assert launch.initial_status == { + "payload_format": "bf16", + "num_lanes": 2, + "graph_mode": "FULL_DECODE_ONLY", + "dbo": True, + "m2n_paired": False, + } + # The engine arguments vLLM builds the F-only config from. + assert args.model == launch.model_path + assert args.data_parallel_rank == 1 + assert args.distributed_executor_backend == "external_launcher" + assert args.skip_tokenizer_init and args.enable_ep_weight_filter + + +def test_eager_launch_runs_one_lane(launch_env, configure): + launch, _ = configure("--enforce-eager") + + assert launch.initial_status["graph_mode"] == "eager" + assert launch.initial_status["num_lanes"] == 1 + assert launch.initial_status["dbo"] is False + + +@pytest.mark.parametrize( + "env, options, afd, message", + [ + ({"WORLD_SIZE": "4"}, ("--enforce-eager",), {}, "must agree"), + ({}, (), {}, "runs --enforce-eager or"), + ( + {}, + ("--compilation-config", '{"cudagraph_mode":"FULL"}'), + {}, + "runs --enforce-eager or", + ), + ( + {}, + ("--enforce-eager", "--tensor-parallel-size", "2"), + {}, + "--tensor-parallel-size 1", + ), + ( + {}, + ("--enforce-eager",), + {"connector": "P2pNcclAFDConnector"}, + "requires DeepEPM2NAFDConnector", + ), + ( + {}, + ("--enforce-eager",), + {"connector_extra_config": {"num_lanes": 2}}, + "num_lanes", + ), + ( + {}, + ("--enforce-eager",), + {"connector_extra_config": {"payload_format": "fp8_block128_fp32"}}, + "payload_format='bf16'", + ), + ], + ids=[ + "world-size", + "unset-graph-mode", + "full-graph-mode", + "tensor-parallel", + "connector", + "lanes-vs-dbo", + "fp8-payload", + ], +) +def test_invalid_launch_is_rejected( + launch_env, configure, monkeypatch, env, options, afd, message +): + for key, value in env.items(): + monkeypatch.setenv(key, value) + with pytest.raises(ValueError, match=message): + configure(*options, afd=afd) + + +def test_cli_reports_a_bad_launch_as_a_usage_error(tmp_path): + pytest.importorskip("vllm") + model = tmp_path / "model" + model.mkdir() + env = {**os.environ, "RANK": "0", "LOCAL_RANK": "0", "WORLD_SIZE": "1"} + result = subprocess.run( + [ + sys.executable, + "-m", + "afd_plugin.ffn_faas", + "serve", + str(model), + "--backend", + "gpu-deepep-m2n", + "--data-parallel-size", + "2", + "--enable-expert-parallel", + "--enforce-eager", + "--additional-config", + _afd_config(), + "--status-port", + "0", + ], + cwd=REPO_ROOT, + env=env, + capture_output=True, + text=True, + timeout=120, + ) + assert result.returncode == 2 + assert "must agree" in result.stderr diff --git a/tests/unit/ffn_faas/test_gpu_settings.py b/tests/unit/ffn_faas/test_gpu_settings.py new file mode 100644 index 000000000..3325acb72 --- /dev/null +++ b/tests/unit/ffn_faas/test_gpu_settings.py @@ -0,0 +1,77 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""GPU standalone FFN settings keep only launcher and service-fixed facts.""" + +from __future__ import annotations + +from pathlib import Path + +import pytest + +from afd_plugin.ffn_faas.platforms.gpu.settings import GPUFFNConfig + + +def _settings(**overrides): + return GPUFFNConfig( + **{ + "model": Path("/model"), + "ffn_ranks": 2, + "rank": 1, + **overrides, + } + ) + + +@pytest.mark.parametrize( + "graph, dbo, graph_mode, num_lanes", + [ + (False, False, "eager", 1), + (True, False, "FULL_DECODE_ONLY", 1), + (True, True, "FULL_DECODE_ONLY", 2), + ], + ids=["eager", "graph", "graph-dbo"], +) +def test_execution_mode_and_lanes(graph, dbo, graph_mode, num_lanes): + settings = _settings(graph=graph, dbo=dbo) + + assert settings.graph_mode == graph_mode + assert settings.num_lanes == num_lanes + + +@pytest.mark.parametrize( + "overrides, message", + [ + ({"rank": 2}, "invalid F-only rank"), + ({"ffn_ranks": 0}, "invalid F-only rank"), + ({"local_rank": -1}, "local rank must be nonnegative"), + ({"shutdown_timeout": float("inf")}, "finite and positive"), + ( + {"connector_extra_config": {"payload_format": "fp8_block128_fp32"}}, + "payload_format='bf16'", + ), + ( + {"connector_extra_config": {"num_lanes": 2}}, + "disagrees with no --enable-dbo", + ), + ( + {"dbo": True, "connector_extra_config": {"num_lanes": 1}}, + "disagrees with --enable-dbo", + ), + ], + ids=[ + "rank", + "ffn-ranks", + "local-rank", + "shutdown-timeout", + "fp8-payload", + "lanes-without-dbo", + "lanes-with-dbo", + ], +) +def test_invalid_launch_settings_are_rejected(overrides, message): + with pytest.raises(ValueError, match=message): + _settings(**overrides) + + +def test_matching_explicit_lane_count_is_accepted(): + assert _settings(dbo=True, connector_extra_config={"num_lanes": 2}).num_lanes == 2 diff --git a/tests/unit/ffn_faas/test_main.py b/tests/unit/ffn_faas/test_main.py index f6219716b..07d483b4f 100644 --- a/tests/unit/ffn_faas/test_main.py +++ b/tests/unit/ffn_faas/test_main.py @@ -175,3 +175,101 @@ def test_cpu_entry_rejects_platform_imports( assert result.returncode == 1 assert f"CPU host contract imported {package}" in result.stderr assert "platform stub ran" not in result.stderr + + +@pytest.fixture +def engine_parser(): + pytest.importorskip("vllm") + from afd_plugin.ffn_faas.__main__ import create_parser + from afd_plugin.ffn_faas.platforms.npu.adapter import AscendCAMAdapter + + return create_parser(AscendCAMAdapter()) + + +@pytest.fixture(params=[False, True], ids=["eager", "full"]) +def native_engine_case(engine_parser, request): + from vllm.engine.arg_utils import EngineArgs + + graph = request.param + config = { + "enable_cpu_binding": False, + "afd": { + "connector": "CAMAsyncAFDConnector", + "role": "ffn", + "connector_extra_config": {"hccl_buffer_size": 256}, + }, + } + mode = ( + ["--compilation-config", '{"cudagraph_mode":"FULL"}'] + if graph + else ["--enforce-eager"] + ) + parsed = engine_parser.parse_args( + [ + "serve", + "/model", + "--dtype", + "bfloat16", + "--max-num-batched-tokens", + "2K", + "--max-num-seqs", + "8", + "--data-parallel-size", + "4", + "--device-ids", + "4,5,6,7", + "--additional-config", + json.dumps(config), + "--shutdown-timeout", + "0", + "--receiver-shutdown-timeout", + "0.5", + *mode, + ] + ) + return ( + graph, + config, + EngineArgs.from_cli_args(parsed), + parsed.receiver_shutdown_timeout, + ) + + +def test_native_engine_options_preserve_configuration(native_engine_case): + graph, config, args, receiver_timeout = native_engine_case + assert args.max_num_batched_tokens == 2048 + assert args.max_num_seqs == 8 + assert args.dtype == "bfloat16" and args.data_parallel_size == 4 + assert args.device_ids == [4, 5, 6, 7] + assert args.additional_config == config + assert args.shutdown_timeout == 0 and receiver_timeout == 0.5 + assert args.enforce_eager is (not graph) + if graph: + assert args.compilation_config.cudagraph_mode.name == "FULL" + + +def test_native_dotted_json_options(engine_parser): + from vllm.engine.arg_utils import EngineArgs + + parsed = engine_parser.parse_args( + [ + "serve", + "/model", + "--additional-config.afd.role", + "ffn", + "--additional-config.afd.port", + "29600", + "--compilation-config.cudagraph_mode", + "FULL", + ] + ) + args = EngineArgs.from_cli_args(parsed) + assert args.additional_config == {"afd": {"role": "ffn", "port": 29600}} + assert args.compilation_config.cudagraph_mode.name == "FULL" + + +@pytest.mark.parametrize("flag", ["--full-graph", "--attention-ranks", "--devices"]) +def test_retired_standalone_engine_flags_are_rejected(engine_parser, flag): + with pytest.raises(SystemExit) as error: + engine_parser.parse_args(["serve", "/model", flag]) + assert error.value.code == 2 diff --git a/tests/unit/ffn_faas/test_runner.py b/tests/unit/ffn_faas/test_runner.py index a3e475f8d..7a6c99f93 100644 --- a/tests/unit/ffn_faas/test_runner.py +++ b/tests/unit/ffn_faas/test_runner.py @@ -74,6 +74,41 @@ def run(): assert backend.finished.wait(timeout=5) +def test_stop_cancels_idle_receive_and_releases(make_runner): + backend = HostBackend(block=True, cancellable=True) + runner, service = make_runner(backend) + done = threading.Event() + + def run(): + try: + runner.run() + finally: + done.set() + + thread = threading.Thread(target=run) + thread.start() + try: + assert backend.entered.wait(timeout=5) + service.request_stop() + assert done.wait(timeout=5) + assert service.snapshot()["phase"] == "stopped" + assert service.snapshot()["resources_released"] is True + assert backend.released + finally: + backend.release.set() + thread.join() + + +def test_backend_ended_receive_exits_loop_and_releases(make_runner): + backend = HostBackend(end_receive=True) + runner, service = make_runner(backend) + runner.run() + assert service.snapshot()["phase"] == "stopped" + assert service.snapshot()["resources_released"] is True + assert service.snapshot()["completed_batches"] == 0 + assert backend.released + + def test_stop_before_backend_construction_needs_no_platform(tmp_path): def unavailable(): raise RuntimeError("no device backend is available") diff --git a/tests/unit/ffn_faas/test_settings.py b/tests/unit/ffn_faas/test_settings.py index 094011399..942f7705d 100644 --- a/tests/unit/ffn_faas/test_settings.py +++ b/tests/unit/ffn_faas/test_settings.py @@ -9,73 +9,84 @@ import pytest -from afd_plugin.config import connector_extra_config_from_source, parse_afd_config +from afd_plugin.ffn_faas.platforms.npu.adapter import AscendCAMAdapter from afd_plugin.ffn_faas.platforms.npu.settings import ( FFNConfig, validate_native_expert_map, validate_static_experts, ) +from tests.unit.fakes.ffn_faas import cam_cli_args @pytest.fixture def settings(): - return FFNConfig(Path("/model"), attention_ranks=12, attention_tp=4, ffn_ranks=4) + return FFNConfig(Path("/model"), ffn_ranks=4) -def test_existing_cam_configuration_roundtrip(settings): - config = settings.additional_config() - afd = parse_afd_config(config, expected_role="ffn") - assert afd.connector == "CAMAsyncAFDConnector" - assert afd.async_dp and afd.compute_gate_on_attention - assert (afd.num_attention_ranks, afd.num_ffn_ranks) == (12, 4) - assert connector_extra_config_from_source(config) == { +@pytest.fixture +def launcher(monkeypatch): + monkeypatch.setenv("WORLD_SIZE", "2") + monkeypatch.setenv("RANK", "1") + monkeypatch.setenv("LOCAL_RANK", "1") + monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1") + + +@pytest.mark.parametrize("positional", [False, True]) +def test_engine_configuration_keeps_user_values(checkpoint_dir, launcher, positional): + args = cam_cli_args(checkpoint_dir) + if positional: + args.model = "/unused-option-model" + args.model_tag = str(checkpoint_dir) + launch = AscendCAMAdapter().configure(args) + assert (launch.rank, launch.local_rank) == (1, 1) + assert launch.model_path == str(checkpoint_dir) + assert args.model == str(checkpoint_dir) + assert args.data_parallel_rank == 1 + assert args.additional_config["enable_cpu_binding"] is False + assert args.additional_config["afd"]["connector_extra_config"] == { "dynamicQuant": 1, "attn_ranks_per_dp": 4, - "ffn_two_stream": False, } -def test_explicit_native_stream_budgets(settings): - settings = replace( - settings, - full_graph=True, - ffn_two_stream=True, - ffn_compute_aic_num=20, - ffn_compute_aiv_num=40, - ffn_combine_aiv_num=8, - hccl_buffer_size=256, - ) - assert connector_extra_config_from_source(settings.additional_config()) == { - "dynamicQuant": 1, - "attn_ranks_per_dp": 4, - "ffn_two_stream": True, - "ffn_compute_aic_num": 20, - "ffn_compute_aiv_num": 40, - "ffn_combine_aiv_num": 8, - "hccl_buffer_size": 256, - } +@pytest.mark.parametrize( + "change", + [ + {"data_parallel_size": 1}, + {"data_parallel_rank": 0}, + {"distributed_executor_backend": "mp"}, + {"receiver_shutdown_timeout": 0}, + ], +) +def test_launcher_rejects_inconsistent_engine_arguments( + checkpoint_dir, launcher, change +): + with pytest.raises(ValueError): + AscendCAMAdapter().configure(cam_cli_args(checkpoint_dir, **change)) + + +@pytest.mark.parametrize( + "field,value", [("port", 0), ("role", "attention"), ("num_ffn_ranks", 0)] +) +def test_cli_uses_existing_afd_validation(checkpoint_dir, launcher, field, value): + args = cam_cli_args(checkpoint_dir) + args.additional_config["afd"][field] = value + with pytest.raises(ValueError): + AscendCAMAdapter().configure(args) @pytest.mark.parametrize( "change", [ - {"attention_ranks": 0}, - {"attention_ranks": 11}, {"ffn_ranks": 0}, - {"attention_tp": 0}, - {"max_num_batched_tokens": 2047}, {"rank": 4}, {"local_rank": -1}, - {"device": -1}, - {"port": 0}, - {"port": 65536}, {"shutdown_timeout": 0}, {"shutdown_timeout": float("nan")}, {"shutdown_timeout": float("inf")}, - {"ffn_compute_aic_num": 0}, ], ) -def test_invalid_launch_rejected_before_native_runtime(settings, change): +def test_invalid_host_launch_rejected(settings, change): with pytest.raises(ValueError): replace(settings, **change) diff --git a/tests/unit/model_executor/models/test_deepseek_v2_construction.py b/tests/unit/model_executor/models/test_deepseek_v2_construction.py index 03756f275..ab846f7d9 100644 --- a/tests/unit/model_executor/models/test_deepseek_v2_construction.py +++ b/tests/unit/model_executor/models/test_deepseek_v2_construction.py @@ -764,7 +764,6 @@ def test_model_constructor_rejects_sequence_parallel_moe_before_allocation( def test_npu_async_attention_allows_model_owned_sp( monkeypatch, construction_env, - pp_size, ): vllm_config = _vllm_config() vllm_config.parallel_config.use_sequence_parallel_moe = True diff --git a/tests/unit/model_executor/models/test_deepseek_v4_external_routing_moe_cpu.py b/tests/unit/model_executor/models/test_deepseek_v4_external_routing_moe_cpu.py new file mode 100644 index 000000000..5cd65d9d0 --- /dev/null +++ b/tests/unit/model_executor/models/test_deepseek_v4_external_routing_moe_cpu.py @@ -0,0 +1,473 @@ +# mypy: disable-error-code="attr-defined, name-defined" +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""CPU tests for the AFD-owned external-routing MoE component. + +The component imports vLLM-Ascend types at module scope, so the tests +install stand-ins for those modules, load the real component from its file, and +drive it with fakes for the comm instance and the routed experts. + +What this covers: the routing reaches ``apply`` unchanged, only the ALLGATHER path +reduces during finalize, an unsupported comm method is rejected, comm instances are +reused across calls, and a failure in any stage leaves no per-call state behind. + +The real operators and graph execution stay with the green-zone suite. +""" + +from __future__ import annotations + +import importlib.util +import sys +import types +from dataclasses import dataclass +from importlib.machinery import ModuleSpec +from pathlib import Path + +import pytest + +torch = pytest.importorskip("torch") + +_REPO_ROOT = Path(__file__).resolve().parents[4] +_MODULE_PATH = ( + _REPO_ROOT + / "afd_plugin" + / "model_executor" + / "models" + / "npu" + / "deepseek_v4_external_routing_moe.py" +) +_ROUTING_PATH = ( + _REPO_ROOT / "afd_plugin" / "connectors" / "npu" / "camp2p" / "external_routing.py" +) + + +class _MoECommType: + ALLGATHER = "ALLGATHER" + MC2 = "MC2" + ALLTOALL = "ALLTOALL" + FUSED_MC2 = "FUSED_MC2" + + +@dataclass(frozen=True) +class _PrepareOutput: + hidden_states: torch.Tensor + router_logits: torch.Tensor | None + mc2_mask: torch.Tensor | None + padded_hidden_states_shape: torch.Size | None + pertoken_scale: torch.Tensor | None = None + + +class _PrepareFinalizeBase: + def __init__(self, moe_config): + self.moe_config = moe_config + self.tp_rank = 0 + self.tp_size = 1 + + def prepare_external(self, hidden_states, routing, **kwargs): + return ( + _PrepareOutput( + hidden_states=hidden_states, + router_logits=None, + mc2_mask=torch.ones(8, dtype=torch.bool), + padded_hidden_states_shape=None, + pertoken_scale=None, + ), + routing, + ) + + +class _CommBase: + def __init__(self, moe_config): + self.moe_config = moe_config + self.prepare_finalize = self._get_prepare_finalize() + self.finalized: list[dict] = [] + + def _get_prepare_finalize(self): + return _PrepareFinalizeBase(self.moe_config) + + def finalize(self, **kwargs): + self.finalized.append(kwargs) + return kwargs["hidden_states"] + + +class _AllGatherComm(_CommBase): + pass + + +class _MC2Comm(_CommBase): + pass + + +def _moe_config(): + return types.SimpleNamespace(dp_size=1, pcp_size=1, tp_size=1, dp_group=None) + + +_STUBBED_MODULE_NAMES = ( + "vllm", + "vllm.distributed", + "vllm.distributed.parallel_state", + "vllm_ascend", + "vllm_ascend.ops", + "vllm_ascend.ops.fused_moe", + "vllm_ascend.ops.fused_moe.dataclass", + "vllm_ascend.ops.fused_moe.dataclass.prepare_finalize", + "vllm_ascend.ops.fused_moe.prepare_finalize", + "vllm_ascend.ops.fused_moe.moe_comm_method", + "vllm_ascend.ops.fused_moe.moe_utils", + "vllm_ascend.quantization", + "vllm_ascend.quantization.quant_type", + "vllm_ascend.ascend_forward_context", + "afd_plugin", + "afd_plugin.connectors", + "afd_plugin.connectors.npu", + "afd_plugin.connectors.npu.camp2p.external_routing", +) + + +def _install_stub_modules(): + """Register the upstream modules the component imports from.""" + routing_spec = importlib.util.spec_from_file_location( + "afd_external_routing_under_test", _ROUTING_PATH + ) + assert routing_spec is not None and routing_spec.loader is not None + routing_module = importlib.util.module_from_spec(routing_spec) + routing_spec.loader.exec_module(routing_module) + + quant = types.ModuleType("vllm_ascend.quantization.quant_type") + quant.QuantType = type("QuantType", (), {"NONE": 0}) + + dataclass_module = types.ModuleType( + "vllm_ascend.ops.fused_moe.dataclass.prepare_finalize" + ) + dataclass_module.MoEPrepareOutput = _PrepareOutput + + prepare_module = types.ModuleType("vllm_ascend.ops.fused_moe.prepare_finalize") + prepare_module.PrepareAndFinalizeWithAllGather = _PrepareFinalizeBase + prepare_module.PrepareAndFinalizeWithMC2 = _PrepareFinalizeBase + + comm_module = types.ModuleType("vllm_ascend.ops.fused_moe.moe_comm_method") + comm_module.AllGatherCommImpl = _AllGatherComm + comm_module.MC2CommImpl = _MC2Comm + # The native selector is keyed by MoECommType and returns None when this + # configuration was not initialised with that method. + comm_by_type = { + _MoECommType.ALLGATHER: _AllGatherComm, + _MoECommType.MC2: _MC2Comm, + } + + def get_moe_comm_method(comm_type, moe_config=None): + impl = comm_by_type.get(comm_type) + return None if impl is None else impl(moe_config) + + comm_module.get_moe_comm_method = get_moe_comm_method + + moe_utils = types.ModuleType("vllm_ascend.ops.fused_moe.moe_utils") + moe_utils._pad_tokens_with_cat = lambda tensor, length: tensor + + parallel = types.ModuleType("vllm.distributed.parallel_state") + parallel.get_pcp_group = lambda: None + + enum_module = types.ModuleType("vllm_ascend.ascend_forward_context") + enum_module.MoECommType = _MoECommType + + context = enum_module + context._EXTRA_CTX = types.SimpleNamespace( + moe_comm_method=_AllGatherComm(_moe_config()), + moe_comm_type=_MoECommType.ALLGATHER, + mc2_mask=torch.ones(8, dtype=torch.bool), + padded_num_tokens=0, + max_tokens_across_dp=0, + max_tokens_across_pcp=0, + ) + + afd_routing = types.ModuleType("afd_plugin.connectors.npu.camp2p.external_routing") + afd_routing.AFDExternalRoutingInputs = routing_module.AFDExternalRoutingInputs + + modules = { + "vllm": types.ModuleType("vllm"), + "vllm.distributed": types.ModuleType("vllm.distributed"), + "vllm.distributed.parallel_state": parallel, + "vllm_ascend": types.ModuleType("vllm_ascend"), + "vllm_ascend.ops": types.ModuleType("vllm_ascend.ops"), + "vllm_ascend.ops.fused_moe": types.ModuleType("vllm_ascend.ops.fused_moe"), + "vllm_ascend.ops.fused_moe.dataclass": types.ModuleType( + "vllm_ascend.ops.fused_moe.dataclass" + ), + "vllm_ascend.ops.fused_moe.dataclass.prepare_finalize": dataclass_module, + "vllm_ascend.ops.fused_moe.prepare_finalize": prepare_module, + "vllm_ascend.ops.fused_moe.moe_comm_method": comm_module, + "vllm_ascend.ops.fused_moe.moe_utils": moe_utils, + "vllm_ascend.quantization": types.ModuleType("vllm_ascend.quantization"), + "vllm_ascend.quantization.quant_type": quant, + "vllm_ascend.ascend_forward_context": context, + "afd_plugin": types.ModuleType("afd_plugin"), + "afd_plugin.connectors": types.ModuleType("afd_plugin.connectors"), + "afd_plugin.connectors.npu": types.ModuleType("afd_plugin.connectors.npu"), + "afd_plugin.connectors.npu.camp2p.external_routing": afd_routing, + } + assert set(modules) == set(_STUBBED_MODULE_NAMES) + for name, module in modules.items(): + # find_spec() is called on these names at import time and raises on a + # module without a spec, so the stand-ins carry one. + module.__spec__ = ModuleSpec(name, None) + sys.modules[name] = module + + spec = importlib.util.spec_from_file_location( + "afd_deepseek_v4_external_routing_moe", _MODULE_PATH + ) + assert spec is not None and spec.loader is not None + component_module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(component_module) + return component_module, context, comm_module + + +def _restore_modules(written: dict[str, types.ModuleType | None]) -> None: + """Put back only the names this fixture replaced, and drop the ones it added.""" + + for name, original in written.items(): + if original is None: + sys.modules.pop(name, None) + else: + sys.modules[name] = original + + +@pytest.fixture(scope="module") +def loaded(): + written = {name: sys.modules.get(name) for name in _STUBBED_MODULE_NAMES} + try: + module, context, comm_module = _install_stub_modules() + except BaseException: + _restore_modules(written) + raise + try: + yield module, context, comm_module + finally: + _restore_modules(written) + + +class _Experts(torch.nn.Module): + def __init__(self, fail_at=None): + super().__init__() + self.moe_config = types.SimpleNamespace( + is_sequence_parallel=False, dp_size=1, pcp_size=1, tp_size=1, dp_group=None + ) + self.quant_type = 0 + self.ascend_pertoken_scale = "stale" + self.ascend_mc2_mask = "stale" + self.seen = [] + self.fail_at = fail_at + self.quant_method = self + + def apply( + self, *, layer, x, topk_weights, topk_ids, shared_experts, shared_experts_input + ): + if self.fail_at == "apply": + raise RuntimeError("apply failed") + self.seen.append((topk_ids.clone(), topk_weights.clone())) + return types.SimpleNamespace(routed_out=x + 1) + + +def _routing(component_module, offset, tokens=4): + ids = torch.arange(offset, offset + tokens, dtype=torch.int32).reshape(tokens, 1) + return component_module.AFDExternalRoutingInputs( + topk_ids=ids, topk_weights=torch.ones(tokens, 1) + ) + + +def test_the_routing_reaches_apply_unchanged(loaded): + module, context, _ = loaded + experts = _Experts() + component = module.AFDExternalRoutingMoE(experts) + context._EXTRA_CTX.moe_comm_method = _AllGatherComm(_moe_config()) + + hidden = torch.zeros(4, 2) + ids = torch.arange(4, dtype=torch.int32).reshape(4, 1) + weights = torch.ones(4, 1) + component.forward(hidden, topk_ids=ids, topk_weights=weights) + + assert torch.equal(experts.seen[-1][0], ids) + assert torch.equal(experts.seen[-1][1], weights) + + +@pytest.mark.parametrize( + "comm_cls,reduces", [(_AllGatherComm, True), (_MC2Comm, False)] +) +def test_only_allgather_reduces_during_finalize(loaded, comm_cls, reduces): + module, context, _ = loaded + comm = comm_cls(_moe_config()) + context._EXTRA_CTX.moe_comm_method = comm + experts = _Experts() + component = module.AFDExternalRoutingMoE(experts) + + component.forward( + torch.zeros(4, 2), + topk_ids=torch.zeros(4, 1, dtype=torch.int32), + topk_weights=torch.ones(4, 1), + ) + afd_comm = component._comm_for(context._EXTRA_CTX.moe_comm_type, comm) + assert afd_comm.finalized[-1]["reduce_results"] is reduces + + # finalize runs after the per-token state is cleared; record what it saw + seen = {} + original_finalize = afd_comm.finalize + + def recording_finalize(**kwargs): + seen["mask"] = experts.ascend_mc2_mask + seen["scale"] = experts.ascend_pertoken_scale + return original_finalize(**kwargs) + + afd_comm.finalize = recording_finalize + component.forward( + torch.zeros(4, 2), + topk_ids=torch.zeros(4, 1, dtype=torch.int32), + topk_weights=torch.ones(4, 1), + ) + assert seen["mask"] is None, "finalize ran with per-token state still set" + assert seen["scale"] is None, "finalize ran with per-token state still set" + assert context._EXTRA_CTX.moe_comm_method is comm, "the comm was not restored" + assert experts.ascend_mc2_mask is None, "per-call state leaked" + assert experts.ascend_pertoken_scale is None, "per-call state leaked" + + +def test_comm_is_reused_across_calls(loaded): + module, context, _ = loaded + native = _AllGatherComm(_moe_config()) + context._EXTRA_CTX.moe_comm_method = native + component = module.AFDExternalRoutingMoE(_Experts()) + + before = dict(component._candidates) + for _ in range(2): + component.forward( + torch.zeros(4, 2), + topk_ids=torch.zeros(4, 1, dtype=torch.int32), + topk_weights=torch.ones(4, 1), + ) + assert component._candidates == before, "a forward built a new comm" + + +def test_an_unsupported_comm_method_is_rejected(loaded): + module, context, _ = loaded + context._EXTRA_CTX.moe_comm_method = types.SimpleNamespace(moe_config=None) + component = module.AFDExternalRoutingMoE(_Experts()) + + with pytest.raises(RuntimeError, match="ALLGATHER and MC2"): + component.forward( + torch.zeros(4, 2), + topk_ids=torch.zeros(4, 1, dtype=torch.int32), + topk_weights=torch.ones(4, 1), + ) + + +def test_a_failure_in_apply_leaves_no_per_call_state(loaded): + module, context, _ = loaded + experts = _Experts(fail_at="apply") + native = _AllGatherComm(_moe_config()) + context._EXTRA_CTX.moe_comm_method = native + component = module.AFDExternalRoutingMoE(experts) + + with pytest.raises(RuntimeError, match="apply failed"): + component.forward( + torch.zeros(4, 2), + topk_ids=torch.zeros(4, 1, dtype=torch.int32), + topk_weights=torch.ones(4, 1), + ) + assert experts.ascend_pertoken_scale is None + assert experts.ascend_mc2_mask is None + assert context._EXTRA_CTX.moe_comm_method is native, "the comm was not restored" + + +class _GatherGroup: + def __init__(self, ranks): + self.ranks = ranks + + def all_gather(self, tensor, dim=0): + return torch.cat([tensor] * self.ranks, dim=dim) + + +def test_dp_gather_and_padding_align_the_routing(loaded): + """Padding then gathering must move the routing with the activation.""" + module, context, _ = loaded + group = _GatherGroup(2) + cfg = types.SimpleNamespace( + is_sequence_parallel=False, dp_size=2, pcp_size=1, tp_size=1, dp_group=group + ) + experts = _Experts() + experts.moe_config = cfg + context._EXTRA_CTX.max_tokens_across_dp = 6 + context._EXTRA_CTX.moe_comm_method = _AllGatherComm(cfg) + + component = module.AFDExternalRoutingMoE(experts) + ids = torch.arange(4, dtype=torch.int32).reshape(4, 1) + component.forward( + torch.zeros(4, 2), + topk_ids=ids, + topk_weights=torch.ones(4, 1), + ) + # padded from 4 to 6 rows, then gathered across 2 ranks + assert experts.seen[-1][0].shape[0] == 12 + + +def test_a_failure_during_prepare_clears_state_and_restores_comm(loaded): + module, context, _ = loaded + native = _AllGatherComm(_moe_config()) + context._EXTRA_CTX.moe_comm_method = native + experts = _Experts() + component = module.AFDExternalRoutingMoE(experts) + + def exploding_prepare(*args, **kwargs): + raise RuntimeError("prepare failed") + + component._comm_for( + context._EXTRA_CTX.moe_comm_type, native + ).prepare_finalize.prepare_external = exploding_prepare + with pytest.raises(RuntimeError, match="prepare failed"): + component.forward( + torch.zeros(4, 2), + topk_ids=torch.zeros(4, 1, dtype=torch.int32), + topk_weights=torch.ones(4, 1), + ) + assert experts.ascend_pertoken_scale is None + assert experts.ascend_mc2_mask is None + assert context._EXTRA_CTX.moe_comm_method is native, "the comm was not restored" + + +def test_a_missing_mc2_candidate_still_serves_allgather(loaded): + """A configuration initialised without MC2 must not need an MC2 candidate.""" + module, context, comm_module = loaded + original = comm_module.get_moe_comm_method + comm_module.get_moe_comm_method = lambda comm_type, moe_config=None: ( + None if comm_type == _MoECommType.MC2 else original(comm_type, moe_config) + ) + try: + experts = _Experts() + native = _AllGatherComm(experts.moe_config) + context._EXTRA_CTX.moe_comm_method = native + component = module.AFDExternalRoutingMoE(experts) + assert module.AFDMC2CommImpl not in component._candidates + component.forward( + torch.zeros(4, 2), + topk_ids=torch.zeros(4, 1, dtype=torch.int32), + topk_weights=torch.ones(4, 1), + ) + assert experts.seen, "the allgather candidate was not used" + finally: + comm_module.get_moe_comm_method = original + + +def test_the_comm_type_resolves_when_the_context_instance_is_unset(loaded): + """MRV2 can leave moe_comm_method unset; the comm type is always present.""" + module, context, _ = loaded + experts = _Experts() + context._EXTRA_CTX.moe_comm_type = _MoECommType.ALLGATHER + context._EXTRA_CTX.moe_comm_method = None + + component = module.AFDExternalRoutingMoE(experts) + component.forward( + torch.zeros(4, 2), + topk_ids=torch.zeros(4, 1, dtype=torch.int32), + topk_weights=torch.ones(4, 1), + ) + assert experts.seen, "expert execution did not run" + assert context._EXTRA_CTX.moe_comm_method is None, ( + "the unset instance was not restored" + ) diff --git a/tests/unit/model_executor/models/test_model_utils.py b/tests/unit/model_executor/models/test_model_utils.py index ef4f2bebb..1413b1e03 100644 --- a/tests/unit/model_executor/models/test_model_utils.py +++ b/tests/unit/model_executor/models/test_model_utils.py @@ -83,6 +83,30 @@ def test_model_config_rewrites_architecture_in_private_copy( assert hf_config.architectures == [architecture] +@pytest.mark.parametrize( + "routed_experts_connector,expected_architecture", + [ + (False, "AFDQwen3MoeForCausalLM"), + (True, "AFDQwen3MoeRoutedExperts"), + ], + ids=["full-model", "routed-experts-shell"], +) +def test_qwen3_routed_connector_selects_the_ffn_only_shell( + model_utils, routed_experts_connector, expected_architecture +): + hf_config = SimpleNamespace(architectures=["Qwen3MoeForCausalLM"]) + original = SimpleNamespace(hf_config=hf_config, hf_text_config=hf_config) + + mapped = model_utils.get_afd_model_config( + original, + device_type="cuda", + routed_experts_connector=routed_experts_connector, + ) + + assert mapped.hf_config.architectures == [expected_architecture] + assert hf_config.architectures == ["Qwen3MoeForCausalLM"] + + def test_model_config_preserves_distinct_nested_text_config(model_utils): text = SimpleNamespace(model_type="qwen3_5_moe_text") original = SimpleNamespace( 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 bdae43e44..11ac52cbd 100644 --- a/tests/unit/model_executor/models/test_qwen3_moe_construction.py +++ b/tests/unit/model_executor/models/test_qwen3_moe_construction.py @@ -235,6 +235,40 @@ def fake_native_init( assert model.afd_role == "attention" +def test_routed_experts_shell_builds_only_native_moe_blocks( + monkeypatch: pytest.MonkeyPatch, + construction_env, +) -> None: + vllm_config = _model_config(layer_count=2) + monkeypatch.setattr(adapter, "current_platform", _platform(is_cuda=True)) + monkeypatch.setattr( + adapter, + "parse_afd_config", + lambda *_args, **_kwargs: AFDConfig( + role="ffn", + connector=DEEPEP_M2N_CONNECTOR, + compute_gate_on_attention=True, + ), + ) + + model = adapter.AFDQwen3MoeRoutedExperts( + vllm_config=vllm_config, + prefix="draft", + ) + + assert construction_env["moe"] == [ + "draft.model.layers.0.mlp", + "draft.model.layers.1.mlp", + ] + assert set(dict(model.named_parameters())) == { + "model.layers.0.mlp.weight", + "model.layers.1.mlp.weight", + } + assert model.get_experts_layer_indices() == (0, 1) + assert model.is_unused_checkpoint_weight("lm_head.weight") + assert not model.is_unused_checkpoint_weight("model.layers.0.mlp.gate.weight") + + def test_attention_causal_lm_reports_no_local_moe_modules( monkeypatch: pytest.MonkeyPatch, ) -> None: diff --git a/tests/unit/model_executor/models/test_sync_attention_gate_proxy.py b/tests/unit/model_executor/models/test_sync_attention_gate_proxy.py new file mode 100644 index 000000000..486c31acc --- /dev/null +++ b/tests/unit/model_executor/models/test_sync_attention_gate_proxy.py @@ -0,0 +1,220 @@ +# mypy: disable-error-code="attr-defined, name-defined" +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project + +"""Wire protocol of the sync Attention-side gate shell. + +The shell's design claim is an ordering: send the routing, run the shared experts +while the FFN side works, then receive and merge. These tests assert that ordering +and the payload, using a fake connector. They prove the Python protocol; they do not +prove that the shared-expert work actually overlaps the remote work on a device, +which needs an NPU timeline. + +Like the other model-layer tests, this needs vLLM and is skipped on a CPU-only host. +""" + +from __future__ import annotations + +import sys +from importlib.machinery import ModuleSpec +from types import ModuleType, SimpleNamespace + +import pytest + +torch = pytest.importorskip("torch") +# The adapter module needs this specific vLLM module, and a CPU-only host can have a +# partial vLLM install where `import vllm` succeeds but this one does not. +pytest.importorskip("vllm.forward_context") +# The adapter chain imports the shared-expert module, which loads vLLM-Ascend at module +# scope; without it the collection fails instead of skipping. +pytest.importorskip("vllm_ascend") + +# afd_plugin builds its model registry from ``find_spec("torch_npu")`` when the package +# is first imported, so it is imported before the stand-in exists; otherwise a process +# without the extension would register the NPU adapter for DSV4 and keep that path after +# the stand-in is gone. +import afd_plugin # noqa: E402,F401 + +# The adapter's import chain pulls the NPU extension in at module scope, but these +# tests only drive the Python wire protocol with fakes. A temporary stand-in keeps +# the module collectable on a CPU-only host; it is removed again so other modules +# still see the extension as missing and keep skipping. +try: + import torch_npu # noqa: F401 + + _stubbed_torch_npu = False +except ModuleNotFoundError: + _torch_npu_stub = ModuleType("torch_npu") + # ``find_spec`` is called on this name at import time, and a module without a + # spec makes it raise instead of reporting the extension as missing. + _torch_npu_stub.__spec__ = ModuleSpec("torch_npu", None) + sys.modules["torch_npu"] = _torch_npu_stub + _stubbed_torch_npu = True + +# The NPU adapter; the CUDA module of the same name has no sync gate shell. +try: + from afd_plugin.model_executor.models.npu import ( + deepseek_v4 as adapter, # noqa: E402 + ) +finally: + if _stubbed_torch_npu: + del sys.modules["torch_npu"] + +TOP_K = 2 +HIDDEN = 4 +LAYER_IDX = 3 + + +class _RecordingConnector: + """Records the wire calls and the shared-expert run in one ordered list.""" + + def __init__(self, events: list[tuple]) -> None: + self.events = events + + def send_attn_output(self, hidden_states, context, **kwargs) -> None: + self.events.append(("send", hidden_states, context, kwargs)) + + def recv_ffn_output(self, *, ref_tensor, ubatch_idx): + self.events.append(("recv", ubatch_idx)) + return ref_tensor * 0.5 + + +def _fake_gate_topk(monkeypatch, events: list[tuple], topk_ids, topk_weights): + """Replace the gate helper with one that records and returns fixed routing. + + The forward does ``from package import deepseek_v4_attention_gate``. Two states + are possible: the submodule was already imported once, so the package carries a + real attribute, or it was not, so the import falls back to ``sys.modules``. + Setting both covers either one; ``raising=False`` because the attribute may not + exist yet. + """ + + import afd_plugin.model_executor.models.npu as npu_package + + helper = ModuleType( + "afd_plugin.model_executor.models.npu.deepseek_v4_attention_gate" + ) + + def compute_attention_gate_topk( + moe, hidden_states, *, input_ids=None, hidden_states_fp32=None + ): + del moe, input_ids, hidden_states_fp32 + events.append(("gate", int(hidden_states.shape[0]))) + return topk_weights, topk_ids + + helper.compute_attention_gate_topk = compute_attention_gate_topk # type: ignore[attr-defined] + monkeypatch.setitem( + sys.modules, + "afd_plugin.model_executor.models.npu.deepseek_v4_attention_gate", + helper, + ) + monkeypatch.setattr( + npu_package, + "deepseek_v4_attention_gate", + helper, + raising=False, + ) + + +def _shell(*, events, connector, shared_experts): + """Build the shell without its weight construction. + + The forward only needs the layer index and the shared-expert module, so the + linear layers the constructor would build are irrelevant here. + + ``shared_experts`` is a presence flag: True installs a recorder-backed module, + False leaves the slot empty. + """ + + del connector # the connector arrives through the forward metadata + shell = object.__new__(adapter.AFDDeepseekV4SyncAttentionGateRemoteMoE) + torch.nn.Module.__init__(shell) + shell.layer_idx = LAYER_IDX + shell.shared_experts = None + if shared_experts: + + class _Shared(torch.nn.Module): # type: ignore[name-defined] + def forward(self, hidden_states): + events.append(("shared", int(hidden_states.shape[0]))) + return torch.full_like(hidden_states, 2.0) + + shell.shared_experts = _Shared() + return shell + + +def _install_metadata(monkeypatch, connector, events): + afd_metadata = SimpleNamespace(connector=connector, stage_idx=7) + monkeypatch.setattr( + adapter, + "get_afd_metadata_from_forward_context", + lambda: afd_metadata, + ) + return afd_metadata + + +def test_forward_sends_routing_then_runs_shared_then_receives(monkeypatch): + """The shared experts must run between the send and the receive.""" + + events: list[tuple] = [] + connector = _RecordingConnector(events) + _install_metadata(monkeypatch, connector, events) + topk_ids = torch.tensor([[0, 1], [2, 3]], dtype=torch.int32) + topk_weights = torch.full((2, TOP_K), 0.5, dtype=torch.float32) + _fake_gate_topk(monkeypatch, events, topk_ids, topk_weights) + shell = _shell(events=events, connector=connector, shared_experts=True) + hidden_states = torch.full((2, HIDDEN), 8.0, dtype=torch.bfloat16) + + output = shell(hidden_states) + + assert [event[0] for event in events] == ["gate", "send", "shared", "recv"] + # The routing travels on the wire slot, not as token ids. + sent = events[1] + assert sent[3]["routing"][0] is topk_ids + assert sent[3]["routing"][1] is topk_weights + # The context identifies the layer and the stage the routing belongs to. + assert sent[2].metadata.layer_idx == LAYER_IDX + assert sent[2].metadata.seq_lens == [2] + assert events[3][1] == 7 # recv uses the metadata stage index + # routed (0.5 * input) plus shared (2.0). + assert torch.equal(output, hidden_states * 0.5 + 2.0) + + +def test_forward_without_shared_experts_returns_routed_only(monkeypatch): + """With no shared experts the output is the received routed tensor alone.""" + + events: list[tuple] = [] + connector = _RecordingConnector(events) + _install_metadata(monkeypatch, connector, events) + topk_ids = torch.zeros((2, TOP_K), dtype=torch.int32) + topk_weights = torch.ones((2, TOP_K), dtype=torch.float32) + _fake_gate_topk(monkeypatch, events, topk_ids, topk_weights) + shell = _shell(events=events, connector=connector, shared_experts=False) + hidden_states = torch.full((2, HIDDEN), 8.0, dtype=torch.bfloat16) + + output = shell(hidden_states) + + assert [event[0] for event in events] == ["gate", "send", "recv"] + assert torch.equal(output, hidden_states * 0.5) + + +def test_forward_refuses_a_context_without_afd_metadata(monkeypatch): + """A missing forward metadata is an error, not a silent activations-only send.""" + + events: list[tuple] = [] + monkeypatch.setattr( + adapter, + "get_afd_metadata_from_forward_context", + lambda: None, + ) + _fake_gate_topk( + monkeypatch, + events, + torch.zeros((2, TOP_K), dtype=torch.int32), + torch.ones((2, TOP_K), dtype=torch.float32), + ) + shell = _shell( + events=events, connector=_RecordingConnector(events), shared_experts=False + ) + + with pytest.raises(RuntimeError, match="AFD forward metadata"): + shell(torch.zeros(2, HIDDEN, dtype=torch.bfloat16)) diff --git a/tests/unit/model_executor/test_async_cam_layered.py b/tests/unit/model_executor/test_async_cam_layered.py index 32578cf72..6b6e90b1a 100644 --- a/tests/unit/model_executor/test_async_cam_layered.py +++ b/tests/unit/model_executor/test_async_cam_layered.py @@ -19,6 +19,9 @@ AsyncCAMLayeredExecutor, LayeredFFNWeights, ) +from afd_plugin.model_executor.npu.async_cam_mega import ( # noqa: E402 + AsyncCAMMegaExecutor, +) def weights(layer_idx=0, per_channel=True, quantization="W4A8"): @@ -58,6 +61,56 @@ def weights(layer_idx=0, per_channel=True, quantization="W4A8"): ) +def test_native_rejects_grouped_scale_before_receiving(operators): + executor = AsyncCAMLayeredExecutor([weights(per_channel=False)]) + with pytest.raises(ValueError, match="per-channel"): + AsyncCAMMegaExecutor(executor, None) + + +@pytest.mark.parametrize("quantization", ["W4A8", "W8A8"]) +def test_native_forwards_the_applied_swiglu_limit(monkeypatch, operators, quantization): + calls = [] + + def execute(*args): + calls.append(args) + return torch.empty(1, dtype=torch.int64) + + monkeypatch.setenv("BATCH_SIZE_FACTOR", str(16 / 262144)) + monkeypatch.setattr( + torch.ops.afd_ascend, + "afd_mega_moe_pull_workspace_size", + lambda *args: 65536, + raising=False, + ) + monkeypatch.setattr( + torch.ops.afd_ascend, "afd_mega_moe_pull_native", execute, raising=False + ) + executor = AsyncCAMMegaExecutor( + AsyncCAMLayeredExecutor( + [replace(weights(quantization=quantization), swiglu_limit=10.0)] + ), + SimpleNamespace( + max_num_batched_tokens=8, + topk=2, + ffn_size=2, + attn_size=1, + world_rank=1, + topology=SimpleNamespace(world_size=3), + tp_size=1, + comm_args=torch.empty(512, dtype=torch.float16), + group_name="test-pull", + ), + ) + executor() + assert calls[0][-1] == (10.0 if quantization == "W4A8" else 0.0) + assert calls[0][10][11] == (1 if quantization == "W4A8" else 2) + assert calls[0][4][0].dtype == ( + torch.int32 if quantization == "W4A8" else torch.int8 + ) + if quantization == "W8A8": + assert calls[0][8] == calls[0][9] == [] + + @pytest.fixture def operators(monkeypatch): calls = [] @@ -203,6 +256,84 @@ def test_repeated_interleaved_layer_ids(operators): assert [call[1][6].item() for call in operators[::2]] == [1, 0, 1, 1] +def runner_method(name, namespace): + tree = ast.parse( + Path("afd_plugin/v1/worker/npu/async_ffn_model_runner.py").read_text() + ) + cls = next( + node + for node in tree.body + if isinstance(node, ast.ClassDef) and node.name == "AsyncNPUFFNModelRunner" + ) + method = next( + node + for node in cls.body + if isinstance(node, ast.FunctionDef) and node.name == name + ) + module = ast.Module( + body=[ + ast.ImportFrom( + module="__future__", names=[ast.alias(name="annotations")], level=0 + ), + method, + ], + type_ignores=[], + ) + exec( + compile(ast.fix_missing_locations(module), "", "exec"), + namespace, + ) + return namespace[name] + + +def test_runner_preserves_raw_context_and_chunk_capacity(): + payloads = [ + SimpleNamespace( + hidden_states=torch.empty(8, 32), + context=SimpleNamespace( + states=SimpleNamespace( + dynamic_scales=torch.ones(8), + group_list=torch.tensor([0, 0]), + token_nums_rankid_layeridx=torch.tensor([99, 1, 5, 0]), + ) + ), + ) + for _ in range(2) + ] + received = iter(payloads) + sent = [] + computed = [] + connector = SimpleNamespace( + recv_attn_output=lambda: next(received), + send_ffn_output=lambda result, ctx: sent.append((result, ctx)), + ) + + def compute(*args): + computed.append(args) + return torch.empty(8, 32, dtype=torch.bfloat16) + + runner = SimpleNamespace( + model=object(), + connector=connector, + _mega_executor=None, + _layered_executor=compute, + layer_ids=(2, 5), + ) + method = runner_method( + "_ffn_forward_connector_driven", + { + "cast": cast, + "CAMAsyncAFDConnector": SimpleNamespace, + "AFDAsyncTransferState": SimpleNamespace, + }, + ) + method(runner) + assert [ctx for _, ctx in sent] == [p.context for p in payloads] + for args, payload in zip(computed, payloads, strict=True): + assert args[0] is payload.hidden_states + assert args[3] is payload.context.states.token_nums_rankid_layeridx + + def model_weights_method(monkeypatch): source = Path("afd_plugin/model_executor/models/npu/deepseek_v4.py").read_text() function = next( diff --git a/tests/unit/model_executor/test_deepseek_v4_hash_ids.py b/tests/unit/model_executor/test_deepseek_v4_hash_ids.py index f8bfc1d3f..161332c25 100644 --- a/tests/unit/model_executor/test_deepseek_v4_hash_ids.py +++ b/tests/unit/model_executor/test_deepseek_v4_hash_ids.py @@ -17,6 +17,7 @@ import sys import types from collections.abc import Iterator +from importlib.machinery import ModuleSpec from pathlib import Path from typing import Any @@ -50,34 +51,34 @@ def _vllm_stub() -> Iterator[None]: yield return - saved = {name: sys.modules.get(name) for name in missing} + builders: dict[str, dict[str, object]] = { + "vllm": {}, + "vllm.forward_context": { + "DPMetadata": type("DPMetadata", (), {}), + "ForwardContext": type("ForwardContext", (), {}), + "get_forward_context": lambda: None, + }, + "vllm.logger": { + "init_logger": lambda *args, **kwargs: logging.getLogger("test") + }, + } - def make(name: str, **attributes: object) -> None: + # Only the modules that are genuinely absent are installed: overwriting one that + # another test imported would survive this context, because the cleanup below + # only removes the names it added. + for name in missing: module = types.ModuleType(name) module.__path__ = [] # type: ignore[attr-defined] - for attribute, value in attributes.items(): + # find_spec("vllm") runs at import time; a spec-less module makes it raise. + module.__spec__ = ModuleSpec(name, None) + for attribute, value in builders[name].items(): setattr(module, attribute, value) sys.modules[name] = module - - make("vllm") - make( - "vllm.forward_context", - DPMetadata=type("DPMetadata", (), {}), - ForwardContext=type("ForwardContext", (), {}), - get_forward_context=lambda: None, - ) - make( - "vllm.logger", - init_logger=lambda *args, **kwargs: logging.getLogger("test"), - ) try: yield finally: for name in missing: sys.modules.pop(name, None) - for name, module in saved.items(): - if module is not None: - sys.modules[name] = module with _vllm_stub(): diff --git a/tests/unit/package/test_ascend_build_files.py b/tests/unit/package/test_ascend_build_files.py index 713ec23d8..3ab0b90cc 100644 --- a/tests/unit/package/test_ascend_build_files.py +++ b/tests/unit/package/test_ascend_build_files.py @@ -183,6 +183,8 @@ def _run_select_ops(*args: str) -> subprocess.CompletedProcess[str]: "afd_async_combine_recv", "grouped_matmul_swiglu_quant_v2_layered", "grouped_matmul_layered", + "afd_mega_moe_pull_native", + "afd_mega_moe_comm", ], ), ("ascend950", ["a2e", "e2a"]), diff --git a/tests/unit/package/test_cam_async_build.py b/tests/unit/package/test_cam_async_build.py index 477bd3446..c1457eecc 100644 --- a/tests/unit/package/test_cam_async_build.py +++ b/tests/unit/package/test_cam_async_build.py @@ -72,8 +72,10 @@ def test_native_sources_survive_source_distribution_manifest(monkeypatch): root = Path(__file__).resolve().parents[3] monkeypatch.chdir(root) files = FileList() + # setuptools works on manifest paths, which are always posix; native + # separators would make every required entry miss on Windows. files.allfiles = [ - str(path.relative_to(root)) + path.relative_to(root).as_posix() for path in (root / "csrc").rglob("*") if path.is_file() ] diff --git a/tests/unit/package/test_mega_activation_patch.py b/tests/unit/package/test_mega_activation_patch.py new file mode 100644 index 000000000..efaff843a --- /dev/null +++ b/tests/unit/package/test_mega_activation_patch.py @@ -0,0 +1,43 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Check that unsupported SDK quantization is rejected instead of left unpatched.""" + +import runpy +from pathlib import Path + +import pytest + +PATCH = ( + Path(__file__).resolve().parents[3] + / "csrc/npu/ascend_kernels/afd_mega_moe_pull_native/op_kernel" + / "afd_mega_moe_patch_activation.py" +) +patched_activation = runpy.run_path(str(PATCH))["patched_activation"] + + +@pytest.mark.parametrize( + ("scale", "tensor", "length"), + [ + ("127.0f / rowAbsMax", "ubGateFp32", "tileLength"), + ("127.f / GMubDequantScale", "ubOutputTmp", "ChunkTileLen"), + ], +) +def test_layered_rounding_is_applied(scale, tensor, length): + source = f""" + auto ubQuantS32 = ubGateFp32.template ReinterpretCast(); + auto scale = {scale}; + AscendC::Cast(ubQuantS32, {tensor}, AscendC::RoundMode::CAST_RINT, {length}); + AscendC::PipeBarrier(); + AscendC::SetDeqScale(static_cast(1.0)); + AscendC::Cast(ubQuantF16, ubQuantS32, AscendC::RoundMode::CAST_RINT, {length}); + """ + result = patched_activation(source) + assert "1.0f / (" in result + cast = f"Cast(ubQuantF16, {tensor}, AscendC::RoundMode::CAST_RINT, {length})" + assert cast in result + assert "ubQuantS32" not in result + + +def test_unknown_sdk_header_is_rejected(): + with pytest.raises(ValueError, match="unsupported CANN activation"): + patched_activation("unrecognized SDK implementation") diff --git a/tests/unit/package/test_package.py b/tests/unit/package/test_package.py index 8447094d6..e90309a7e 100644 --- a/tests/unit/package/test_package.py +++ b/tests/unit/package/test_package.py @@ -109,6 +109,9 @@ def test_register_afd_preserves_native_models_and_is_idempotent(registration): "AFDQwen3MoeForCausalLM": ( "afd_plugin.model_executor.models.qwen3_moe:AFDQwen3MoeForCausalLM" ), + "AFDQwen3MoeRoutedExperts": ( + "afd_plugin.model_executor.models.qwen3_moe:AFDQwen3MoeRoutedExperts" + ), "AFDQwen3_5MoeForConditionalGeneration": ( "afd_plugin.model_executor.models.qwen3_5:" "AFDQwen3_5MoeForConditionalGeneration" @@ -124,7 +127,19 @@ def test_register_afd_preserves_native_models_and_is_idempotent(registration): def test_entry_point_is_registered(): - entry_points = importlib.metadata.entry_points(group="vllm.general_plugins") - matches = [ep for ep in entry_points if ep.name == "afd"] + """The installed distribution has to expose the vLLM plugin entry point. + + Importing the package from a source tree does not register installation + metadata, so the entry point can only be checked on an installed build. + """ + try: + distribution = importlib.metadata.distribution("vllm-afd-plugin") + except importlib.metadata.PackageNotFoundError: + pytest.skip("the plugin is not installed in this environment") + matches = [ + entry_point + for entry_point in distribution.entry_points + if entry_point.group == "vllm.general_plugins" and entry_point.name == "afd" + ] assert matches assert matches[0].value == "afd_plugin:register_afd" diff --git a/tests/unit/test_sync_attention_gate_validation.py b/tests/unit/test_sync_attention_gate_validation.py new file mode 100644 index 000000000..d4680c4df --- /dev/null +++ b/tests/unit/test_sync_attention_gate_validation.py @@ -0,0 +1,96 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project + +"""Runner validation for the sync Attention-side gate prototype. + +The DBO and ubatching limits live in the NPU feature gate, which both workers run +before this validator, and are covered in +``compat/npu/test_sync_gate_feature_validation_cpu.py``. This module covers the +runner-level contract: a supported synchronous gate configuration must return +cleanly for either role. +""" + +import sys +from types import ModuleType, SimpleNamespace + +import pytest + +from afd_plugin import validation + + +def _runner_config(*, enable_dbo=False, use_ubatching=False): + """A ``VllmConfig`` surface complete enough for a clean return. + + Every field the validator reads on the accepted path is present, so the + "allowed" test can assert a real return instead of tolerating an exception. + """ + + return SimpleNamespace( + use_v2_model_runner=True, + model_config=SimpleNamespace(enforce_eager=True), + speculative_config=None, + parallel_config=SimpleNamespace( + enable_dbo=enable_dbo, + use_ubatching=use_ubatching, + num_ubatches=2, + pipeline_parallel_size=1, + prefill_context_parallel_size=1, + decode_context_parallel_size=1, + data_parallel_size=4, + tensor_parallel_size=1, + enable_expert_parallel=True, + enable_elastic_ep=False, + enable_eplb=False, + use_sequence_parallel_moe=False, + ), + compilation_config=SimpleNamespace( + pass_config=SimpleNamespace(enable_sp=False), + ), + ) + + +@pytest.fixture() +def sync_gate_afd_config(monkeypatch): + """Make the validator see the sync gate-on-Attention configuration.""" + + from afd_plugin.config import CAMP2P_CONNECTOR + + afd_config = SimpleNamespace( + connector=CAMP2P_CONNECTOR, + async_dp=False, + compute_gate_on_attention=True, + num_attention_ranks=4, + num_ffn_ranks=4, + ) + monkeypatch.setattr( + validation, + "parse_afd_config", + lambda _vllm_config, expected_role=None: afd_config, + ) + # The validator imports this lazily, and the real module path pulls vLLM in. + # Inject the leaf module so the accepted path can run on a CPU host. + for name, attributes in ( + ("afd_plugin.model_executor", {}), + ("afd_plugin.model_executor.models", {}), + ( + "afd_plugin.model_executor.models.model_utils", + {"has_afd_model_registration": lambda _model_config: True}, + ), + ): + module = ModuleType(name) + module.__path__ = [] + for attribute, value in attributes.items(): + setattr(module, attribute, value) + monkeypatch.setitem(sys.modules, name, module) + return afd_config + + +@pytest.mark.parametrize("role", ["attention", "ffn"]) +def test_validator_accepts_a_plain_sync_gate_config(sync_gate_afd_config, role): + """A supported configuration must return cleanly, not merely avoid one message.""" + + validation.validate_npu_model_runner_v2_config( + _runner_config(), + expected_role=role, + device_type="npu", + ) diff --git a/tests/unit/v1/worker/npu/test_dspark_graph.py b/tests/unit/v1/worker/npu/test_dspark_graph.py new file mode 100644 index 000000000..42884a3b0 --- /dev/null +++ b/tests/unit/v1/worker/npu/test_dspark_graph.py @@ -0,0 +1,105 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Draft graph scopes preserve independent AFD batches and native preparation.""" + +from contextlib import nullcontext + +import pytest + +pytest.importorskip("torch") +pytest.importorskip("vllm") +pytest.importorskip("vllm_ascend") + +from vllm.config import CUDAGraphMode # noqa: E402 +from vllm.v1.worker.gpu.cudagraph_utils import ( # noqa: E402 + AttentionState, + BatchExecutionDescriptor, +) + +from afd_plugin.v1.worker.npu import dspark_graph # noqa: E402 +from tests.unit.fakes.npu_attention import ( # noqa: E402 + DraftGraphManager, + forward_context, + metadata_runner, +) + + +@pytest.mark.parametrize("fail", [False, True]) +def test_draft_replay_isolates_padding_and_restores_target_context(fail): + runner = metadata_runner() + target = forward_context(12) + runner.install_afd_metadata_on_forward_context(target) + target_metadata = target.additional_kwargs["afd_metadata"] + manager = DraftGraphManager(runner, fail=fail) + desc = BatchExecutionDescriptor(CUDAGraphMode.FULL, num_tokens=15, num_reqs=3) + expected = ( + pytest.raises(RuntimeError, match="native replay failed") + if fail + else nullcontext() + ) + with ( + expected, + runner._afd_draft_context(False, False), + dspark_graph.use_afd_draft_replay(runner, manager, real_tokens=10), + ): + manager.run_fullgraph(desc) + draft = manager.emitted[0] + assert draft.tokens_lens == [15] + assert draft.tokens_unpadded_lens == [10] + assert draft.transaction_id != target_metadata.transaction_id + following = forward_context(12) + runner.install_afd_metadata_on_forward_context(following) + assert following.additional_kwargs["afd_metadata"] is target_metadata + manager.fail = False + assert manager.run_fullgraph(desc) is target_metadata + + +@pytest.mark.parametrize("fail", [False, True]) +def test_draft_capture_creates_distinct_batches_and_restores_native_prepare( + monkeypatch, fail +): + runner = metadata_runner() + desc = BatchExecutionDescriptor(CUDAGraphMode.FULL, num_tokens=10, num_reqs=2) + manager = DraftGraphManager(runner, [desc]) + native_result = AttentionState({}, {}) + monkeypatch.setattr( + dspark_graph.dflash_cudagraph, + "_prepare_dflash_inputs_to_capture", + lambda *args, **kwargs: native_result, + ) + emitted = [] + expected = ( + pytest.raises(RuntimeError, match="native capture failed") + if fail + else nullcontext() + ) + with ( + expected, + runner._afd_draft_context(True, False), + dspark_graph.use_afd_draft_capture(runner, manager), + ): + for _ in range(2): + result = dspark_graph.dflash_cudagraph._prepare_dflash_inputs_to_capture( + 2, 10, None, None, None, None, 4096, False, False + ) + assert result is native_result + context = forward_context(10) + runner.install_afd_metadata_on_forward_context(context) + emitted.append(context.additional_kwargs["afd_metadata"]) + if fail: + raise RuntimeError("native capture failed") + assert [item.tokens_lens for item in emitted] == [[10], [10]] + assert emitted[0].transaction_id != emitted[1].transaction_id + assert ( + dspark_graph.dflash_cudagraph._prepare_dflash_inputs_to_capture() + is native_result + ) + + +def test_draft_full_rejects_missing_native_capture_descriptors(): + runner = metadata_runner() + with ( + pytest.raises(RuntimeError, match="no native capture descriptors"), + dspark_graph.use_afd_draft_capture(runner, DraftGraphManager(runner)), + ): + pass diff --git a/tests/unit/v1/worker/test_async_ffn_graph.py b/tests/unit/v1/worker/test_async_ffn_graph.py index a403788e5..4454eeec0 100644 --- a/tests/unit/v1/worker/test_async_ffn_graph.py +++ b/tests/unit/v1/worker/test_async_ffn_graph.py @@ -10,8 +10,10 @@ pytest.importorskip("vllm") pytest.importorskip("torch_npu") +pytest.importorskip("vllm_ascend") import torch # noqa: E402 +from afd_plugin.connectors.base import FFNExecutionSession # noqa: E402 from afd_plugin.v1.worker.npu.async_ffn_model_runner import ( # noqa: E402 AsyncNPUFFNModelRunner, ) @@ -78,10 +80,13 @@ def send(output, context): connector.recv_attn_output = receive connector.send_ffn_output = send + connector.ffn_session = FFNExecutionSession session = AsyncNPUFFNModelRunner.__new__(AsyncNPUFFNModelRunner) session.model = torch.nn.Linear(1, 1) session._layered_executor = execute + session._mega_executor = None session.connector = connector + session.session = connector.ffn_session() session._graph = None session._graph_capture_count = session._graph_replay_count = 0 session._graph_pending = session._graph_closed = False diff --git a/tests/unit/v1/worker/test_async_ffn_runtime.py b/tests/unit/v1/worker/test_async_ffn_runtime.py index 114401afd..91c3db77a 100644 --- a/tests/unit/v1/worker/test_async_ffn_runtime.py +++ b/tests/unit/v1/worker/test_async_ffn_runtime.py @@ -4,6 +4,7 @@ from __future__ import annotations +import threading from collections import deque from types import SimpleNamespace @@ -11,6 +12,7 @@ pytest.importorskip("vllm") pytest.importorskip("torch_npu") +pytest.importorskip("vllm_ascend") import torch # noqa: E402 from afd_plugin.config import AFDConfig # noqa: E402 @@ -24,6 +26,8 @@ CAMAsyncAFDConnector, ) from afd_plugin.v1.worker.npu import async_ffn_model_runner as runtime # noqa: E402 +from afd_plugin.v1.worker.npu.ffn_worker import AFDNPUFFNWorker # noqa: E402 +from tests.unit.fakes.connector import FakeFFNSession # noqa: E402 from tests.unit.fakes.vllm_config import npu_vllm_config # noqa: E402 @@ -161,3 +165,51 @@ def execute(hidden, scales, groups, header): ): torch.testing.assert_close(sent[1], hidden.to(torch.bfloat16), rtol=0, atol=0) assert sent[2] is received_payload.context.states.token_nums_rankid_layeridx + + +@pytest.mark.parametrize("fail", [False, True]) +def test_full_receiver_exit_requires_success_even_during_shutdown(harness, fail): + worker = AFDNPUFFNWorker.__new__(AFDNPUFFNWorker) + worker.model_runner = harness.runner + harness.runner.effective_mode = "full" + harness.runner.session = FakeFFNSession(requires_drain=True) + worker._ffn_thread = None + worker._ffn_loop_error = None + worker._bind_cpus_once = lambda: None + + def serve(): + worker._ffn_shutdown_event.set() + if fail: + raise RuntimeError("partial device completion") + + worker._run_ffn_server_loop = serve + worker.start_ffn_server_loop() + assert isinstance(worker._ffn_thread, threading.Thread) + worker._ffn_thread.join(3) + assert not worker._ffn_thread.is_alive() + assert harness.runner._graph_receiver_exited is not fail + assert (worker._ffn_loop_error is not None) is fail + + +@pytest.mark.parametrize("drains", [True, False]) +def test_drain_propagates_session_result(harness, monkeypatch, drains): + h = harness + h.runner.session = FakeFFNSession( + transactions=1, + drain_error=None if drains else RuntimeError("retain resources"), + ) + h.runner._layered_executor = lambda hidden, *_args: hidden.to(torch.bfloat16) + request = payload(5, 3) + h.payloads.append(request) + monkeypatch.setattr(torch.npu, "synchronize", lambda: None) + if drains: + h.runner.drain() + else: + with pytest.raises(RuntimeError, match="retain resources"): + h.runner.drain() + returned = [ + call for call in h.calls if isinstance(call, tuple) and call[0] == "send" + ] + assert len(returned) == 1 + torch.testing.assert_close(returned[0][1], request.hidden_states.to(torch.bfloat16)) + assert returned[0][2] is request.context.states.token_nums_rankid_layeridx diff --git a/tests/unit/v1/worker/test_async_ffn_two_stream.py b/tests/unit/v1/worker/test_async_ffn_two_stream.py deleted file mode 100644 index acd2fbbfb..000000000 --- a/tests/unit/v1/worker/test_async_ffn_two_stream.py +++ /dev/null @@ -1,232 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project -"""CPU stream protocol checks; device qualification is recorded separately.""" - -import sys -import weakref -from types import SimpleNamespace -from unittest.mock import MagicMock, call - -import pytest - -from afd_plugin.async_cam_startup import AsyncCAMGraphStartup -from tests.unit.fakes.async_ffn import load_async_modules - - -@pytest.fixture -def pipeline(monkeypatch): - _, connector, runtime = load_async_modules(monkeypatch) - module = sys.modules["afd_plugin.v1.worker.npu.async_ffn_two_stream"] - npu, history, produced = module.torch.npu, MagicMock(), [] - for name in ("compute_stream", "send_stream", "ready0", "done0", "ready1", "done1"): - history.attach_mock(MagicMock(), name) - npu.NPUGraph.side_effect = lambda: MagicMock() - npu.graph_pool_handle.side_effect = object - npu.Stream.side_effect = [history.compute_stream, history.send_stream] - npu.Event.side_effect = [ - history.ready0, - history.done0, - history.ready1, - history.done1, - ] - - class Tensor: - pass - - def receive(): - hidden = Tensor() - produced.append(weakref.ref(hidden)) - state = SimpleNamespace( - dynamic_scales=1, group_list=2, token_nums_rankid_layeridx=hidden - ) - return SimpleNamespace( - hidden_states=hidden, context=SimpleNamespace(states=state) - ) - - def send(output, context): - assert output is context.states.token_nums_rankid_layeridx - - history.receive.side_effect = receive - history.compute.side_effect = lambda hidden, *a, **kw: hidden - history.send.side_effect = send - config = connector.AFDAsyncExtraInfo(ffn_two_stream=True) - owner = module.AsyncFFNTwoStream( - "npu:0", - config, - SimpleNamespace(recv_attn_output=history.receive, send_ffn_output=history.send), - history.compute, - ) - history.reset_mock() - return SimpleNamespace( - owner=owner, npu=npu, history=history, produced=produced, runtime=runtime - ) - - -@pytest.fixture -def warmup_runner(pipeline): - p = pipeline - runtime = p.runtime - runner = runtime.AsyncNPUFFNModelRunner.__new__(runtime.AsyncNPUFFNModelRunner) - runner.model = object() - runner._two_stream = p.owner - startup = AsyncCAMGraphStartup(1, 1, 2) - store = MagicMock() - startup._store = store - runner.connector = SimpleNamespace(graph_startup=startup) - p.history.done0.query.return_value = True - runner._ffn_forward_connector_driven(single=True) - p.history.reset_mock() - p.produced.clear() - return runner, store - - -@pytest.mark.parametrize("outcome", ["ready", "peer_failure", "timeout"]) -def test_warmup_observes_completion_or_failure_without_blocking_fence( - pipeline, warmup_runner, monkeypatch, outcome -): - p = pipeline - runner, store = warmup_runner - p.history.done1.query.return_value = outcome == "ready" - store.check.return_value = outcome == "peer_failure" - store.get.return_value = b"Attention rank failed" - clock = iter([0, p.runtime.GRAPH_STARTUP_TIMEOUT_SECONDS]) - monkeypatch.setattr(p.runtime.time, "monotonic", lambda: next(clock)) - p.npu.synchronize.side_effect = AssertionError("warmup must poll completion") - if outcome == "ready": - runner._ffn_forward_connector_driven(single=True) - p.owner.close(receiver_exited=True) - p.history.reset_mock() - assert all(ref() is None for ref in p.produced) - else: - error, message = ( - (RuntimeError, "Attention rank failed") - if outcome == "peer_failure" - else (TimeoutError, "warmup timed out; retain resources") - ) - with pytest.raises(error, match=message): - runner._ffn_forward_connector_driven(single=True) - with pytest.raises(RuntimeError, match="failed; retain resources"): - p.owner.close(receiver_exited=True) - p.history.reset_mock() - assert all(ref() is not None for ref in p.produced) - - -def test_eager_orders_events_and_retires_before_allocating_third_item(pipeline): - p = pipeline - p.owner.submit(3) - expected = [] - for index, slot in enumerate((0, 1, 0)): - if index == 2: - expected += ["done0.synchronize", "compute_stream.wait_event"] - expected += [ - "receive", - "compute", - f"ready{slot}.record", - "send_stream.wait_event", - "send", - f"done{slot}.record", - ] - assert [c[0] for c in p.history.mock_calls] == expected - p.history.send_stream.wait_event.assert_has_calls( - [call(p.history.ready0), call(p.history.ready1), call(p.history.ready0)] - ) - p.history.compute_stream.wait_event.assert_called_once_with(p.history.done0) - p.history.ready0.record.assert_has_calls([call(p.owner.compute)] * 2) - p.history.done0.record.assert_has_calls([call(p.owner.send)] * 2) - p.history.reset_mock() # Drop mock call arguments before checking tensor lifetime. - assert [ref() is not None for ref in p.produced] == [False, True, True] - p.owner.complete() - assert all(ref() is None for ref in p.produced) - - -def test_full_keeps_slot_addresses_and_reuses_with_device_edges(pipeline): - p = pipeline - p.owner.capture() - graphs = [ - graph - for slot in p.owner.slots - for graph in (slot.compute_graph, slot.send_graph) - ] - assert len({id(graph) for graph in graphs}) == 4 - assert len({pool for slot in p.owner.slots for pool in slot.pools}) == 4 - assert len(p.produced) == 2 and p.produced[0]() is not p.produced[1]() - p.history.reset_mock() # Call recording must not keep captured tensors alive. - p.owner.submit(5) - p.history.done0.synchronize.assert_not_called() - p.history.compute_stream.wait_event.assert_any_call(p.history.done0) - assert len(p.produced) == 2 - p.owner.complete() - assert all(ref() is not None for ref in p.produced) - with pytest.raises(RuntimeError, match="receiver exit"): - p.owner.close(receiver_exited=False) - p.owner.close(receiver_exited=True) - p.history.reset_mock() - assert all(ref() is None for ref in p.produced) - assert all(graph.reset.call_count == 1 for graph in graphs) - - -@pytest.mark.parametrize( - "failure", ["capture", "receive", "compute", "send", "retire", "replay", "fence"] -) -def test_partial_failure_retains_resources_and_disallows_retry(pipeline, failure): - p = pipeline - if failure == "replay": - p.owner.capture() - failing_call = { - "capture": p.history.send, - "receive": p.history.receive, - "compute": p.history.compute, - "send": p.history.send, - "retire": p.history.done0.synchronize, - "replay": p.owner.slots[0].compute_graph.replay - if failure == "replay" - else None, - "fence": p.npu.synchronize, - }[failure] - failing_call.side_effect = RuntimeError("native failure") - with pytest.raises(RuntimeError, match="native failure"): - if failure == "capture": - p.owner.capture() - else: - p.owner.submit(3 if failure == "retire" else 1) - p.owner.complete() - for operation in ( - p.owner.capture, - lambda: p.owner.submit(1), - lambda: p.owner.close(receiver_exited=True), - ): - with pytest.raises(RuntimeError, match="failed"): - operation() - assert all( - slot.compute_graph is None or not slot.compute_graph.reset.called - for slot in p.owner.slots - ) - p.history.reset_mock() - assert all(ref() is not None for ref in p.produced) - - -@pytest.mark.parametrize("budgets", [(25, 24, 24), (12, 23, 24), (12, 24, 25)]) -def test_invalid_budget_rejected_before_communication(monkeypatch, budgets): - _, connector, runtime = load_async_modules(monkeypatch) - config = connector.AFDAsyncExtraInfo( - ffn_two_stream=True, - ffn_compute_aic_num=budgets[0], - ffn_compute_aiv_num=budgets[1], - ffn_combine_aiv_num=budgets[2], - ) - communication = MagicMock() - with pytest.raises(ValueError, match="core limits"): - runtime.AsyncFFNTwoStream("npu:0", config, communication, MagicMock()) - assert not communication.mock_calls - - -def test_pending_work_cannot_be_resubmitted_or_closed(pipeline): - owner = pipeline.owner - owner.submit(1) - with pytest.raises(RuntimeError, match="not ready"): - owner.submit(1) - with pytest.raises(RuntimeError, match="device completion"): - owner.close(receiver_exited=True) - owner.complete() - owner.close(receiver_exited=True) - owner.close(receiver_exited=True) diff --git a/tests/unit/v1/worker/test_ffn_model_runner.py b/tests/unit/v1/worker/test_ffn_model_runner.py index 93a299f8a..ef5d65eee 100644 --- a/tests/unit/v1/worker/test_ffn_model_runner.py +++ b/tests/unit/v1/worker/test_ffn_model_runner.py @@ -50,11 +50,14 @@ def make(control_plane): worker._ffn_shutdown_event = threading.Event() worker._ffn_loop_error = None worker.device = SimpleNamespace(type="cpu") - worker.model_runner = SimpleNamespace( - connector=SimpleNamespace(control_plane=control_plane, is_initialized=True), - use_cuda_graph=False, - model="owned weights", + runner = object.__new__(GPUFFNModelRunner) + runner.connector = SimpleNamespace( + control_plane=control_plane, is_initialized=True ) + runner.use_cuda_graph = False + runner.model = "owned weights" + runner.device = SimpleNamespace(type="cpu") + worker.model_runner = runner return worker return make @@ -272,6 +275,12 @@ def _runner_with_connector_and_model(model, *, num_layers=1): runner.vllm_config = SimpleNamespace( # FFN keeps connector-driven execution within a GPU MRV2 deployment. use_v2_model_runner=True, + # A real VllmConfig always carries a model config. The forward context asks + # it for the MX scale algorithm, which reads exactly two fields: the + # architecture list and the HF text config's model type. Supplying the two + # keeps this fixture honest about the contract instead of omitting the whole + # object; both default to "not MiniMax M3", which is what these tests are. + model_config=SimpleNamespace(architectures=None, hf_text_config=None), parallel_config=SimpleNamespace( data_parallel_size=1, is_moe_model=True, @@ -603,6 +612,51 @@ def compute_experts_output( ] +def test_connector_driven_step_executes_the_received_payload(): + runner = _runner_with_connector_and_model(_FakeModel()) + runner.device = SimpleNamespace(type="cpu") + metadata = _metadata() + runner.connector.attn_outputs.append(_payload("hidden", metadata)) + runner.connector.recv_dp_metadata_list = lambda: AFDControlPayload( + dp_metadata_list={0: _FakeDPMetadata([1])}, + is_graph_capturing=False, + is_warmup=False, + ) + + assert runner.execute_connector_driven_step(cancel_requested=lambda: False) + assert runner.connector.ffn_outputs == [("ffn(hidden, layer=0)", metadata)] + + +@pytest.mark.parametrize( + "error, cancel_requested, returns_false", + [ + (AFDControlPlaneClosedError("closed"), True, True), + (AFDControlPlaneClosedError("closed"), False, False), + (RuntimeError("NCCL failure"), True, False), + ], +) +def test_connector_driven_step_reports_only_requested_cancellation( + error, cancel_requested, returns_false +): + def receive(): + raise error + + runner = _runner_with_connector_and_model(_FakeModel()) + runner.device = SimpleNamespace(type="cpu") + runner.connector.recv_dp_metadata_list = receive + + if returns_false: + assert not runner.execute_connector_driven_step( + cancel_requested=lambda: cancel_requested + ) + else: + with pytest.raises(type(error)) as caught: + runner.execute_connector_driven_step( + cancel_requested=lambda: cancel_requested + ) + assert caught.value is error + + def test_ffn_runner_requires_dp_metadata_list(): runner = object.__new__(GPUFFNModelRunner) runner.prof = None @@ -751,9 +805,8 @@ def test_ffn_worker_loop_rejects_connector_without_control_plane(): worker._ffn_shutdown_event = event worker.device = SimpleNamespace(type="cpu") - worker.model_runner = SimpleNamespace( - connector=_ConnectorDrivenFakeConnector(), - ) + worker.model_runner = object.__new__(GPUFFNModelRunner) + worker.model_runner.connector = _ConnectorDrivenFakeConnector() with pytest.raises(NotImplementedError, match="control-plane-driven"): worker._run_ffn_server_loop() diff --git a/tests/unit/v1/worker/test_model_runner_v2.py b/tests/unit/v1/worker/test_model_runner_v2.py index 8c7e35c58..2e7a55562 100644 --- a/tests/unit/v1/worker/test_model_runner_v2.py +++ b/tests/unit/v1/worker/test_model_runner_v2.py @@ -25,6 +25,7 @@ from vllm.v1.worker.gpu.ubatch_utils import UBatchState # noqa: E402 from vllm.v1.worker.ubatch_utils import UBatchSlice # noqa: E402 +from afd_plugin.config import AFDConfig # noqa: E402 from afd_plugin.model_executor.models import ( # noqa: E402 forward_context as afd_forward_context, ) @@ -1443,7 +1444,9 @@ def _prepare_gpu_worker_runtime(monkeypatch, worker_case, runner_cls): monkeypatch.setattr(native_v2, "GPUModelRunner", _NativeRunnerSentinel) monkeypatch.setattr(worker_module.Worker, "init_device", _fake_worker_init) monkeypatch.setattr( - worker_module, "assert_compatible_afd_stack", lambda *args, **kwargs: None + worker_module, + "assert_compatible_afd_stack", + lambda *args, **kwargs: AFDConfig(role=role), ) monkeypatch.setattr(worker_module.torch.accelerator, "empty_cache", lambda: None) return worker @@ -1476,7 +1479,9 @@ def unexpected_construction(*_args, **_kwargs): monkeypatch.setattr(worker_module, runner_name, unexpected_construction) monkeypatch.setattr(worker_module.Worker, "init_device", unexpected_construction) monkeypatch.setattr( - worker_module, "assert_compatible_afd_stack", lambda *args, **kwargs: None + worker_module, + "assert_compatible_afd_stack", + lambda *args, **kwargs: AFDConfig(role=role), ) with pytest.raises(RuntimeError, match="requires ModelRunnerV2"): diff --git a/tests/unit/v1/worker/test_npu_model_runner_v2.py b/tests/unit/v1/worker/test_npu_model_runner_v2.py index c35dbcafa..7cb0283c8 100644 --- a/tests/unit/v1/worker/test_npu_model_runner_v2.py +++ b/tests/unit/v1/worker/test_npu_model_runner_v2.py @@ -56,7 +56,11 @@ def _runner(*, async_dp=False): runner.afd_config = SimpleNamespace( async_dp=async_dp, num_attention_ranks=2, num_ffn_ranks=2 ) - runner.connector = SimpleNamespace(extra_info=SimpleNamespace()) + runner.connector = SimpleNamespace( + extra_info=SimpleNamespace(), + control_plane=None if async_dp else object(), + attention_step=nullcontext, + ) runner.cudagraph_manager = None runner.ubatch_runner = None runner.prof = object() @@ -198,10 +202,11 @@ def update(backend, stream, forward, tokens, config, speculative): @pytest.mark.parametrize("profile_only", [False, True]) @pytest.mark.parametrize("raise_in_native", [False, True]) -def test_capture_forwards_new_inputs_and_restores( - monkeypatch, profile_only, raise_in_native -): # long-test: capture lifecycle - runner = _runner() +@pytest.mark.parametrize("async_dp", [False, True]) +def test_capture_forwards_new_inputs_and_restores( # long-test: capture lifecycle + monkeypatch, profile_only, raise_in_native, async_dp +): + runner = _runner(async_dp=async_dp) descs = [ cudagraph_utils.BatchExecutionDescriptor( cg_mode=CUDAGraphMode.FULL, num_reqs=1, num_tokens=8 @@ -309,11 +314,72 @@ def native_capture(self, *, profile_only=False): ] if not profile_only: expected_events.extend([(32, True, False), (32, False, True)]) - assert events == expected_events + assert events == ([] if async_dp else expected_events) assert cudagraph_utils.prepare_inputs_to_capture is native_prepare assert original_prepare is not native_prepare +@pytest.mark.parametrize("async_dp", [False, True]) +@pytest.mark.parametrize("fail_replay", [False, True]) +def test_replay_real_and_padded_metadata( # long-test: native replay/context seam + monkeypatch, async_dp, fail_replay +): + runner = _runner(async_dp=async_dp) + original_metadata = runner._afd_pending_metadata + runner._afd_transaction_counter = 0 + runner.compilation_config.cudagraph_mode = CUDAGraphMode.FULL_DECODE_ONLY + runner.build_capture_dp_metadata = lambda tokens: tokens + sent = [] + runner.send_dp_metadata = lambda tokens, _slices: sent.append(tokens) + emitted = [] + desc = cudagraph_utils.BatchExecutionDescriptor( + CUDAGraphMode.FULL, num_tokens=8, num_reqs=8 + ) + + def context(): + return SimpleNamespace( + batch_descriptor=None, + attn_metadata={}, + additional_kwargs={}, + ubatch_slices=None, + ) + + class Manager: + def run_fullgraph(self, descriptor): + update_context = context() + runner.install_afd_metadata_on_forward_context(update_context) + emitted.append(update_context.additional_kwargs["afd_metadata"]) + if fail_replay: + raise RuntimeError("native replay failed") + return "native-output" + + runner.cudagraph_manager = Manager() + + def native_execute(self, *args, **kwargs): + assert should_skip_dp_coordination() is async_dp + return self.cudagraph_manager.run_fullgraph(desc) + + monkeypatch.setattr(NPUModelRunner, "execute_model", native_execute) + monkeypatch.setattr(npu_v2, "step_afd_npu_profiler", lambda _prof: None) + expected = ( + pytest.raises(RuntimeError, match="native replay failed") + if fail_replay + else nullcontext() + ) + with expected: + assert ( + runner.execute_model(SimpleNamespace(total_num_scheduled_tokens=5)) + == "native-output" + ) + assert emitted[0].tokens_lens == [8] + assert emitted[0].tokens_unpadded_lens == [5] + assert sent == ([] if async_dp else [8]) + following = context() + runner.install_afd_metadata_on_forward_context(following) + assert following.additional_kwargs["afd_metadata"] is original_metadata + assert not should_skip_dp_coordination() + + @pytest.mark.parametrize("num_tokens", [0, 1, 7]) @pytest.mark.parametrize("dp_rank", [0, 1]) def test_async_execute_uses_native_local_dp_dispatch(monkeypatch, num_tokens, dp_rank): @@ -642,7 +708,7 @@ def test_native_child_contexts_preserve_dp_and_own_ascend_kwargs(monkeypatch): assert result == contexts assert [ctx.dp_metadata for ctx in contexts] == dp assert contexts[1].attn_metadata is None - from afd_plugin.connectors.npu.camp2p import _get_group_ep + from afd_plugin.connectors.npu.camp2p.connector import _get_group_ep assert [ _get_group_ep(getattr(ctx, "ubatch_idx", 0), "stage0", "stage1", "") diff --git a/tests/unit/v1/worker/test_npu_runtime.py b/tests/unit/v1/worker/test_npu_runtime.py index 2eae24c4e..552bed742 100644 --- a/tests/unit/v1/worker/test_npu_runtime.py +++ b/tests/unit/v1/worker/test_npu_runtime.py @@ -23,6 +23,7 @@ AFDTransferMetadata, AFDTransferState, ) +from afd_plugin.connectors.base import FFNExecutionSession from tests.unit.fakes.dp_metadata import ( # noqa: E402 DPMetadataValues as _FakeDPMetadata, ) @@ -827,7 +828,9 @@ def test_npu_ffn_worker_uses_connector_driven_loop_for_async_connector(monkeypat event = threading.Event() calls: list[str | tuple[str, SimpleNamespace]] = [] - def execute_connector_driven_step(): + def execute_connector_driven_step(shutdown_requested): + if shutdown_requested: + return False calls.append("step") event.set() return True @@ -836,6 +839,7 @@ def execute_connector_driven_step(): worker.device = SimpleNamespace(type="cpu") worker.model_runner = object.__new__(ffn_worker.AsyncNPUFFNModelRunner) worker.model_runner.connector = SimpleNamespace(control_plane=None) + worker.model_runner.session = FFNExecutionSession() worker.model_runner.effective_mode = "eager" worker.model_runner.execute_connector_driven_step = execute_connector_driven_step diff --git a/tools/benchmarks/request_generator.sh b/tools/benchmarks/request_generator.sh index a60dbbf7b..3ad8e3d45 100755 --- a/tools/benchmarks/request_generator.sh +++ b/tools/benchmarks/request_generator.sh @@ -8,13 +8,13 @@ # # 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) @@ -28,7 +28,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} diff --git a/uv.lock b/uv.lock index 3c3f6e353..1f13929e7 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=v1.1.0#741b9aa5364687ac84c5049f5c6239b8cae6a7eb" } 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=v1.1.0" }, { name = "vllm", marker = "extra == 'vllm'", specifier = "==0.32.0" }, ] provides-extras = ["afdep", "vllm"]