[RFC]: ATTN-FFN Disaggregation for MoE Models

[chopper0126]

Motivation.

In large-scale MoE inference, the benefits of sparsity in the decoding phase drive continuous expansion of expert parallelism (EP). Prior designs (e.g., DeepEP[1]) improve throughput at scale but have limited scalability because they place EP shards across data-parallel (DP) ranks. Recently, Attention–FFN disaggregation (AFD) has been proposed by ByteDance[2], StepFun[3], and Huawei[4]. The motivation is straightforward: the attention phase is memory-bound, whereas the FFN/expert phase is compute-bound, so a single homogeneous deployment cannot optimize both simultaneously. Module-wise heterogeneous placement is beneficial as EP continues to scale.
Based on these production insights, we propose introducing AFD in vLLM: decouple Attention and FFN/experts at resource levels so they can scale independently, and overlap communication in a “perfectly balanced” pipeline to improve throughput. The initial revision targets an eager-mode MVP without custom kernels: it runs multiple DP replicas on the Attention side and EP replicas on the MoE side, enabling cross-node M2N routing with correctness guarantees. Subsequently, we will provide stable AFD interfaces for communication, load balancing, and elasticity.

[1] https://github.com/deepseek-ai/DeepEP
[2] Zhu R, Jiang Z, Jin C, et al. MegaScale-Infer: Serving Mixture-of-Experts at Scale with Disaggregated Expert Parallelism[J]. arXiv preprint arXiv:2504.02263, 2025.
[3] Wang B, Wang B, Wan C, et al. Step-3 is Large yet Affordable: Model-system Co-design for Cost-effective Decoding[J]. arXiv preprint arXiv:2507.19427, 2025.
[4] Xiao A, He B, Zhang B, et al. xDeepServe: Model-as-a-Service on Huawei CloudMatrix384[J]. arXiv preprint arXiv:2508.02520, 2025.

Proposed Change.

We propose to start with two core changes:

1. AFD on vLLM: Provide initial AFD capabilities for DeepSeek-v3 and other MoE models, under eager mode, with torch communication APIs for M2N; design doc

2. AFConnector interface for XPUs: to support custom AF communication backends, e.g., stepmesh. design doc #issues

Roadmaps

• AFD on vLLM (MVP) @chopper0126 @xiaoshudian555 @wabluy

  • Disaggregated MoE deployment:
    • Model side: redefine model structure, weight loading, and module graph;
    • Service side: define a global topo to launch a standalone MoE service with torchrun, enable communication with Attention nodes (temporarily via P2P);
  • Accuracy evaluation: validate the correctness of outputs;
  • Pipeline support: configurable micro-batch size and asynchronous M2N communications;
  • Performance optimizations: integrate features such as MTP and support custom load-balancing strategies.
  • Extensibility: support custom MoE models and refine the codebase for AFConnector.

• AFConnector interface @jiangkuaixue123

  • collect IPs and ranks of attention&ffn instances and initialize the necessary communication resources
  • tensor level async send/recv and all-gather/reduce-scatter APIs to transfer tensors between attention and ffn instances
  • hardware agnostic interface to support various hardware backends, communication libraries, and NICs
  • plugin mechanism by Python setuptools entrypoints

After completing the two core changes, we also propose the following features:

• Graph mode support: to accelerate AFD on XPUs;

• Full disaggregation (AF + PD): Introduce PD disaggregation on top of AFD; @hsliuustc0106

• Unified benchmarking framework: Standardize throughput/latency evaluation and dataset baselines;

• Elastic and fault-tolerant AFD: Support online scale-out/in of A/F workers while preserving KV consistency and correct routing, improving fault tolerance; @leonkang1 @fangyuchu @wabluy

• Service management: Provide a unified interface to integrate with solutions like AIBrix and enable reliable scaling management. @chopper0126

We welcome contributors in the community to discuss and co-work on the above topics and brainstorm more AFD optimizations after the first two changes :)

Feedback Period.

No response

CC List.

@chopper0126 @xiaoshudian555 @wabluy @leonkang1 @fangyuchu @jianzs @hsliuustc0106 @jiangkuaixue123

Any Other Things.

No response

Before submitting a new issue...

• Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.

[Just-it]

When can we get a demo?

[david6666666]

Looking forward to follow-up updates, especially demo.

[yiakwy-xpu-ml-framework-team]

The best way and mature method to implement the method, better to be called Decoding Stage FFN-Attention Model Sharding proposed in step3, is to impelment dynamo graph transforming (ffn_attn_sharding_transform.cpp / ffn_attn_sharding_transform.py) : sharding FFN-Attention to different sub graphs place them on different devices, then create communication ops between FFN and attention if their devices are different.

Once we sharded the graph, then we can implement so call multi-stages (device 0 (stage 0)-> device 1 (stage 1) -> device0 (stage 2) -> device 1 (stage 3) ....) ping-pong pipeline.

This communication is not trival, since the communication ops is dynamically, so better to be inserted in computing transforming stage. Hence I recommend to do it with dynamo.

The first challenging is how to optimize these ping-pong pipeline. I recommend to start off from distributed gemm, nvshmem shared memmory based communication instead of traditional nccl ops , which requires a lot of work to optimize in-comming ringbuffer, output-comming ringbuffer processing overlap when mutiple micro batches in a device.

The second challenge is that the model does not garantee to work well with computing bound MLA and memory bound Qwen/GPT-OSS GQA. So sopporting step3 may be first step followed by this optimization.

[Oliver-ss]

Hi, we’re currently working on bringing AFD to the open-source vLLM (built on our stepmesh communication library and step3 model). We’ve got an initial design draft ready and are now cleaning up the code. We expect to open a preliminary PR in a week or two, and would love for everyone to join the discussion when it’s out.

[wabluy]

The best way and mature method to implement the method, better to be called Decoding Stage FFN-Attention Model Sharding proposed in step3, is to impelment dynamo graph transforming (ffn_attn_sharding_transform.cpp / ffn_attn_sharding_transform.py) : sharding FFN-Attention to different sub graphs place them on different devices, then create communication ops between FFN and attention if their devices are different.

Once we sharded the graph, then we can implement so call multi-stages (device 0 (stage 0)-> device 1 (stage 1) -> device0 (stage 2) -> device 1 (stage 3) ....) ping-pong pipeline.

This communication is not trival, since the communication ops is dynamically, so better to be inserted in computing transforming stage. Hence I recommend to do it with dynamo.

The first challenging is how to optimize these ping-pong pipeline. I recomned to start off from distributed gemm, nvshmem shared memmory based communication instead of traditional nccl ops , which requires a lot of work to optimize in-comming ringbuffer, output-comming ringbuffer processing overlap when mutiple micro batches in a device.

The second challenge is the model does not garantee to work well with computing bound MLA and Qwen/GPT-OSS GQA. So sopporting step3 may be first step followed by this optimization.

Hi @yiakwy-xpu-ml-framework-team , thanks for the detailed proposal. We agree that graph mode (via Dynamo) is a good path for performance on GPU: sharding FFN and Attention into separate subgraphs placed on different devices, with communication ops inserted during the transform pass.

Given we plan to support heterogeneous hardware and multiple backends, we’d like a unified, framework-level approach in vLLM:

• A portable communication-op abstraction in vLLM with various backends(NCCL/HCCL/NVSHMEM/Stepmesh, etc.).

We will then focus on communication and scheduling, reducing bubbles under multiple micro-batches with different high-performance communication backends.

For AFD, I agree that the potential risks are with compute-bound MLA and GQA variants. We propose to first land vanilla AFD, then iterate with model-specific deployment optimizations.

We’ll share a more detailed design and initial PR in the coming days. Feedback on the comm-op abstraction and scheduler interface would be especially helpful. Thanks again—looking forward to collaborating.

[wabluy]

Hi, we’re currently working on bringing AFD to the open-source vLLM (built on our stepmesh communication library and step3 model). We’ve got an initial design draft ready and are now cleaning up the code. We expect to open a preliminary PR in a week or two, and would love for everyone to join the discussion when it’s out.

Hi Stepfun team @Oliver-ss —thank you for the impactful contributions, and great to see AFD coming to vLLM. We’re pursuing a unified design in this area and would like to align to ensure API-level interoperability. Would you be open to a brief sync to coordinate? We’ll share our draft and invite co-review on the initial PR. Looking forward to more collaborations on open-source optimizations with you.

[yiakwy-xpu-ml-framework-team]

@wabluy Great let 's work on it. I am especailly interested in nvshem based communication banckend for NV/AMD GPU. And I will work POC and MVP on SGLang in more details.

@Oliver-ss as far as I understand, step3 currently only supports RoCE mode NIC, in order to achieve performance lead in a MVP solution, IB based solution may have higher priority and step3 communication libray can serve as one of the backends.

[youkaichao]

I created a slack channel #feat-afd in vllm slack, feel free to join the discussion.

[chopper0126]

I created a slack channel #feat-afd in vllm slack, feel free to join the discussion.

Hi, thanks for creating a discussion on #feat-afd, it will help us make this feature better.

[guanyu1998]

1. sglang ：[Feature] Support AF Disaggregation
   [Feature] Support AF Disaggregation sgl-project/sglang#9401
2. sglang ：[RFC] Support FFN-Attention Disaggregation Strategy for scientific research and low capable NIC devices
   [RFC] Support FFN-Attention Disaggregation Strategy for scientific research and low capable NIC devices sgl-project/sglang#9347
   3.support Multi-Matrix Factorization Attention and Attention-FFN Disaggregation
   [Roadmap] support Multi-Matrix Factorization Attention and Attention-FFN Disaggregation sgl-project/sglang#8360

[Dong-Jiahuan]

Hi, we’re currently working on bringing AFD to the open-source vLLM (built on our stepmesh communication library and step3 model). We’ve got an initial design draft ready and are now cleaning up the code. We expect to open a preliminary PR in a week or two, and would love for everyone to join the discussion when it’s out.

Hi there! Just wanted to follow up on the AFD integration for vLLM you mentioned a couple of weeks ago. How’s everything going? Looking forward to the PR and the discussion once it’s ready!