afd-nano-vLLM 设计文档
输入“/”快速插入内容
afd-nano-vLLM 设计文档
afd-nano-vLLM 设计文档
用户4030用户40302月25日修改
1.
设计概述
本文的目标是在 nano-vLLM 框架上实现 Qwen-0.6B 的 Attention–FFN 解耦部署，重点关注 Attention 与 FFN 的分离部署，并实现二者之间的数据传输。其总的架构如下所示：
AFD架构总览
在 AFD 架构的初期实现中，我们采用了 Attention 与 FFN 模块的分离部署策略，将其置于不同的计算节点。架构引入了 Model Runner 与 FFN Runner 两类执行器：Model Runner 负责统筹模型推理的全流程逻辑；而 FFN Runner 则作为被动执行单元，仅循环执行“接收-计算-回传”的任务流，不参与模型的流控制。为了实现异构节点间的协同，系统引入了 AFDConnector 组件，负责在不同类型的 Rank 之间高效传输 Tensor 及相关元数据。
2.
Attention–FFN分离部署
2.1
AFD config
配置AFD的有关配置，包括通信实现类型、部署attention和FFN的GPU数量
代码块
class AFDConfig:
afd_connector: str = "dummy" # Transport backend: 'dummy', 'navie'
num_attention_servers: int = 1 # Number of attention servers
num_ffn_servers: int = 1 # Number of FFN servers
afd_server_rank: int = 0 # Server rank identifier
2.2
Model Runner
1.
ModelRunner 负责Attention部分的推理，包括schedule、prepare，相比原来的nano-vllm新增：
a.
init 函数内部初始化afd通信器(afd_connector)，用于初始化模型和传递GPU之间的通信
python
class ModelRunner:
def __init__(self, config: Config, rank: int, event: Event | list[Event]):
attn_ranks = list(range(self.afd_config.num_attention_servers))
self.attn_ranks_group = dist.new_group(ranks=attn_ranks)
self.afd_connector = None
if self.afd_config is not None:
self.afd_connector = AFDConnectorFactory.create_connector(
self.rank, self.rank, self.afd_config
)
b.
在exit()函数中增加发送信号给ffn rank，告诉其结束任务
代码块
class ModelRunner:
def exit(self):
...
#通知ffn rank结束任务
self.afd_connector.send_shutdown_signal()
...
2.3
Model FFN Runner
新增ModelFFNRunner负责FFN部分的推理，其相比Attention Model Runer，较为简单，一直等待接收来自attention rank的数据，然后将计算结果返回attention rank即可。
1.
init 初始化模型，并执行启动接收任务
代码块
class ModelFFNRunner:
def __init__(self, config: Config, rank: int, event: Event | list[Event]):
...
self.afd_connector = None
if self.afd_config is not None:
self.afd_connector = AFDConnectorFactory.create_connector(
self.rank, self.rank, self.afd_config
)
self.model = Qwen3ForCausalLM(hf_config, self.afd_connector)
load_model(self.model, config.model)
...
try:
self.ffn_loop()
except Exception as e:
print(f"[Rank {self.rank}] Fatal Error: {e}")
raise e
finally:
self.exit()