[shlomitk1]

Purpose

Enable FFN-side EPLB (Expert Parallel Load Balancing) for CUDA AFD deployments on the pinned vLLM 0.30.0.

The FFN role now owns vLLM's native EPLBController and follows the upstream GPUModelRunner wiring: the loaded model registers its native DeepseekV2MoE layers (via the inherited MixtureOfExperts protocol), and every connector-driven step prepares the router's unpadded token count and steps the EPLB state. Rearrangement, load-window all-reduce, in-place map
commit (CUDA-graph compatible), and expert-weight transfer (nixl/torch_gloo auto-selected) all run in upstream EplbState machinery — no plugin duplication.

Two deployment recipes enable EPLB on DeepSeek-V2-Lite (2A2F colocation, single-node and multipod) with 32 redundant experts (96 physical, 48 per FFN rank at EP size 2).

Design notes: docs/design/module/ffn_runtime.md ("FFN-side EPLB" section).

Issue

• Related issue(s): #
• Closing keyword, only if fully resolved: Closes #

Scope

• In scope:
  • GPUFFNModelRunner EPLB wiring: controller construction, model registration on load, per-step prepare_forward + step, dummy stepping for warmup/graph-capture paths
  • Unpadded token accounting for the router's map-and-record kernel (real_attention_token_total, placeholder rows excluded)
  • Validation: EPLB allowed on the FFN role only; rejected on Attention and with DBO/ubatching
  • Unit tests for all of the above; two GPU recipes; module doc update
• Out of scope:
  • NPU FFN EPLB (NPU validator still rejects enable_eplb)
  • EPLB + DBO/ubatching on the FFN role (rejected — see Implementation Notes)
  • An E2E scenario under tests/e2e (recipes are the E2E vehicle for now; needs a GPU run first)
  • Elastic EP integration (setup_eplb_from_mapping not wired)

Implementation Notes

• Why FFN-only: with compute_gate_on_attention=false, all MoE routing, expert weights, and EP-group membership live on the FFN ranks; the Attention role ships hidden states and its CUDA remote-expert shells reject EPLB (AFDDeepseekV2DeepseekV2DecoderLayer construction check retained as defense in depth). EPLB flags therefore go on the FFN serve command only
  (as in the recipes).
• Extension point used: the runner reuses vLLM 0.30's EPLBController (vllm/v1/worker/gpu/eplb_utils.py) rather than reimplementing stepping. The connector daemon loop is the "once per engine step" cadence that upstream gets from sample_tokens/pool/_dummy_run decorators; the new eplb_step seam is documented with a patch-reason block referencing the pinned
  commit (ced6857afa0ea7b2e3f0846a62e1394e90f15607).
• Step semantics: eager forwards and graph replays are real steps (the record kernel runs inside captured graphs and reads device-side should_record/map tensors, which are committed in place after rearrangement). Warmup and graph-capture steps are dummy steps that keep every FFN rank's rearrangement counter in lockstep — required or the rearrangement
  collectives hang.
• Token counting: the unpadded count is the step's total real Attention tokens (TP-expanded, idle-peer placeholder rows excluded). Each FFN rank's router view covers those tokens either through TP-replicated delivery (FFN DP=1) or the MoE-level all-gather (FFN DP>1), matching what the map-and-record kernel processes. The shared TP-expansion logic was
  extracted to avoid a third copy of the connector's expansion rules.
• DBO restriction: FFN EPLB rejects enable_dbo/use_ubatching. The record kernel selects its per-ubatch token tensor via dbo_current_ubatch_id(), and the FFN daemon loop never installs vLLM's DBO thread context, so multi-stage steps would read the wrong count. Enforced in validation and with a single-stage runtime check.
• No monkey patches added; the eplb process group and override_envs_for_eplb come free from the inherited native Worker init path.

Test Plan

CPU-only (this PR):

• uv run pytest tests/unit — full unit suite
• uv run ruff check / ruff format --check on touched files
• uv run mypy on touched plugin modules
• bash -n on both recipes

GPU-gated (follow-up, not run here): serve both recipes' commands against DeepSeek-V2-Lite, confirm the FFN log shows EPLB is enabled for MoE part of model, a Rearranging experts line after step_interval - step_interval/4 steps, and GSM8K smoke accuracy; then wire an E2E scenario.

Test Result

• uv run pytest tests/unit -p no:warnings → 1255 passed, 114 skipped, 55 subtests passed (exit 0)
• New coverage: accept-ffn-eplb, reject-attention-eplb, reject-ffn-eplb-dbo, reject-ffn-eplb-ubatching (validator); 8 runner tests (load/registration, real vs dummy stepping, replay stepping, TP-expanded count computation, disabled no-op, multi-stage rejection, capture-path stepping); 7 parametrized real_attention_token_total cases
• ruff check → All checks passed; ruff format --check → clean
• mypy (3 files) → Success: no issues found
• GPU E2E: not run — no GPU host in this session; skipped per plan above
• bash -n on both recipes → syntax OK

Docs Impact

• Files updated: docs/design/module/ffn_runtime.md (new "FFN-side EPLB" section; amended the V2-pair static-EP sentence to note the FFN-only EPLB exception), plus the two recipes under recipe/gpu/P2pNcclAFDConnector/deepseek_v2_lite/prefill_decode_colocation/
• If none, reason: N/A

────────────────────────────────────────────────────────────────────────────────

Essential PR Checklist

• Purpose is clear and linked to public context when possible.
• Scope is bounded.
• Compatibility with vLLM v0.30.0 is considered.
• No changes are made to the vLLM source checkout.
• Plugin-owned classes or explicit dotted class paths are preferred over monkey patches.
• Any compat shim or monkey patch is isolated, idempotent, version-guarded, documented, and tested. (no new patches; the one new seam, eplb_step, is documented with a patch-reason block and unit-tested)
• Imports remain CPU-safe; CUDA-heavy work is delayed or GPU-gated. (new import vllm.v1.worker.gpu.eplb_utils is no heavier than the module's existing vLLM imports; tests/unit still collects and passes without CUDA)
• Validation evidence is included, including skipped GPU tests when applicable.
• Documentation impact is stated.

────────────────────────────────────────────────────────────────────────────────