Objective
Extend AIConfigurator (AIC) to model Attention-FFN Disaggregation (AFD), where Attention runs on A-Workers and FFN/MoE on F-Workers within each decode step. This enables independent scaling of the two dominant workload types in Transformer decoding and supports GPU+LPU heterogeneous deployment analysis.
Goals
- AFD Serving Mode: Add AFD as a first-class mode alongside aggregated and P/D-disaggregated, supporting Dense and MoE architectures with search over A-Worker / F-Worker configs.
- AFD Cost Modeling: Model per-layer A↔F communication, pipeline overlap, and latency balancing between heterogeneous worker pools.
- Side-by-Side Comparison: Enable AFD vs. aggregated vs. P/D-disaggregated comparison in a single AIC run.
- Hardware Abstraction: Design device-agnostic hardware specs to lay the groundwork for GPU+LPU mixed deployments.
Non-Goals
- Runtime AFD serving implementation (offline analysis only).
- LPU profiling data collection (currently uses analytical estimation).
- Changes to existing aggregated or P/D-disaggregated modes.
Background
AIC currently supports aggregated and P/D-disaggregated modes, both of which colocate Attention and FFN/MoE on the same device during Decode. Under typical batch sizes, both Attention (KV-cache bound) and FFN (weight bound) are memory-bandwidth-bound, yet neither can be independently scaled. AFD separates them: Attention stays on A-Workers while FFN concentrates on F-Workers, increasing effective batch size to push FFN toward compute-bound.
Unlike P/D disaggregation — which transfers KV-cache once per request — AFD exchanges activations twice per layer, making communication modeling accuracy critical. AFD is orthogonal to P/D and can be used together. AIC's existing pipeline (search-space enumeration, profiling, rate-matching, Pareto filtering) provides the structural template for this extension.
Roadmap
Implementation is split into three phases, starting with CLI estimate mode to deliver a functional MVP, followed by default mode integration and accuracy verification.
Phase 1: CLI Estimate Mode (~5.30)
| # |
Task |
Key Points |
| 1 |
Add AFD mode to estimate CLI |
- Add AFD mode and its related parameters to input
- Orthogonal to P/D disaggregation mode: can be applied to P or D independently, or used in combination
- Output performance metrics follow the existing format |
| 2 |
AFDInferenceSession |
- run_afd performance simulation, including micro batch overlap
- Can be used in combination with P/D disaggregation mode
- GPU Memory (HBM) bound check |
| 3 |
Model partitioning |
- Generic build_afd_ops_partition()
- Support partitioning for both prefill and decode phases
- Support both MoE and dense models, implement one of each |
| 4 |
AFDTransfer operation |
- Simulate communication patterns between A-Workers and F-Workers, assuming uniform probability for MoE
- Support Attn DP/TP × FFN TP/EP combinations
- Compatible with multiple backends; implement one variant: P2P + AG/RS |
Timeline
| Deadline |
Task |
| 4.23 – 5.7 |
- [ ] Add AFD mode to CLI, submit PR (1)
- [ ] Create PR (3), draft implementation
- [ ] Detailed design for key modules (if needed) |
| ~5.21 |
- [ ] Create PRs (2, 4), draft implementation |
| ~6.4 (complete) |
- [ ] Add necessary unit tests and integration tests
- [ ] Complete PR revisions and merge |
Phase 2: CLI Default Mode
| # |
Task |
Key Points |
| 1 |
Add AFD mode to default CLI |
- CLI input and columns output
- Integrate upper-layer CLI args
- Align output format and schema |
| 2 |
Search space generation |
- TaskConfig / TaskRunner integration (_disagg_defaults_layer → build_afd_parallel_lists)
- AFD config layer
- Generate candidate configuration set
- Hard constraints |
| 3 |
Pareto analysis integration |
- Implement afd_pareto → find_best_afd_result_under_constraints
- Search for optimal solution within candidate set
- Soft constraints (pruning) |
Phase 3: Accuracy Verification and Feature Extension
TBD
Objective
Extend AIConfigurator (AIC) to model Attention-FFN Disaggregation (AFD), where Attention runs on A-Workers and FFN/MoE on F-Workers within each decode step. This enables independent scaling of the two dominant workload types in Transformer decoding and supports GPU+LPU heterogeneous deployment analysis.
Goals
Non-Goals
Background
AIC currently supports aggregated and P/D-disaggregated modes, both of which colocate Attention and FFN/MoE on the same device during Decode. Under typical batch sizes, both Attention (KV-cache bound) and FFN (weight bound) are memory-bandwidth-bound, yet neither can be independently scaled. AFD separates them: Attention stays on A-Workers while FFN concentrates on F-Workers, increasing effective batch size to push FFN toward compute-bound.
Unlike P/D disaggregation — which transfers KV-cache once per request — AFD exchanges activations twice per layer, making communication modeling accuracy critical. AFD is orthogonal to P/D and can be used together. AIC's existing pipeline (search-space enumeration, profiling, rate-matching, Pareto filtering) provides the structural template for this extension.
Roadmap
Implementation is split into three phases, starting with CLI estimate mode to deliver a functional MVP, followed by default mode integration and accuracy verification.
Phase 1: CLI Estimate Mode (~5.30)
- Orthogonal to P/D disaggregation mode: can be applied to P or D independently, or used in combination
- Output performance metrics follow the existing format
run_afdperformance simulation, including micro batch overlap- Can be used in combination with P/D disaggregation mode
- GPU Memory (HBM) bound check
build_afd_ops_partition()- Support partitioning for both prefill and decode phases
- Support both MoE and dense models, implement one of each
- Support Attn DP/TP × FFN TP/EP combinations
- Compatible with multiple backends; implement one variant: P2P + AG/RS
Timeline
- [ ] Create PR (3), draft implementation
- [ ] Detailed design for key modules (if needed)
- [ ] Complete PR revisions and merge
Phase 2: CLI Default Mode
- Integrate upper-layer CLI args
- Align output format and schema
_disagg_defaults_layer→build_afd_parallel_lists)- AFD config layer
- Generate candidate configuration set
- Hard constraints
afd_pareto→find_best_afd_result_under_constraints- Search for optimal solution within candidate set
- Soft constraints (pruning)
Phase 3: Accuracy Verification and Feature Extension
TBD