Skip to content

[Feature]: Move shared-expert computation to the FFN side for NPU Async AFD #331

Description

@jiangkuaixue123

Feature, motivation, and pitch

Place shared-expert computation on the FFN side for Ascend NPU Async AFD. This aligns shared and routed expert computation with the Attention/FFN split and aims to reduce compute and weight-memory pressure on the Attention side.

Proposed change

  • Load and execute shared experts on the FFN side for the targeted NPU Async model paths.
  • Carry the required inputs and metadata through the Async connector and combine shared-expert and routed-expert outputs with the model's original semantics.
  • Preserve correct stream/event ordering and tensor lifetimes across asynchronous execution.
  • Document supported models, topologies, and configuration. Audit existing model-specific handling and reuse it where applicable.

vLLM compatibility and extension points

Target the repository-pinned vLLM/vLLM-Ascend runtime pair. Implement through plugin-owned NPU model, runner, and connector extension points, reusing native operators where possible. Any necessary compatibility patch must be minimal, version-pinned, and documented with an upstream/removal plan; do not modify upstream source trees.

Alternatives considered

Keep shared-expert computation on the Attention side for paths that currently use that placement. This avoids migration work but retains the associated Attention-side compute and memory pressure.

Validation plan

  • Compare MoE outputs and end-to-end accuracy against a reference preserving the original shared-expert semantics.
  • Cover token partitioning, padding, zero-token ranks, repeated requests, and supported parallel topologies.
  • Verify Async communication ordering, output combination, and shutdown on Ascend hardware.
  • Measure Attention/FFN memory use, latency, and throughput before and after the change.

Related: #210 (zero shared-expert token handling), #287 (NPU Async prefill).

Before submitting

  • I searched existing issues and RFCs.
  • I identified whether this belongs in plugin-owned code, compat helpers, or compat patches.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions