AI Infra 9 分钟阅读

在 8×RTX 4090 上部署 Qwen3-32B:SGLang EdgeAcc + DP Attention 从跑不起来到能服务

在 8×RTX 4090 上部署 Qwen3-32B:SGLang EdgeAcc + DP Attention 从"跑不起来"到"能服务"

一次真实的推理架构工程记录:没有 NVLink 的消费级 4090 集群,如何把 32B 稠密模型的高并发在线推理从"启动即崩溃/死锁"做到"稳定服务、吞吐翻倍"。

0. TL;DR

  • 硬件:单机 8×RTX 4090 24GB,消费级、无 NVLink,张量并行(TP)只能走 PCIe(NCCL_P2P_LEVEL=SYS)——这是所有优化的根本约束。
  • 目标:把 Qwen3-32B 做成高并发在线推理服务(并发 256、长上下文),供游戏/GUI Agent 后端调用。
  • 引擎:阿里 ENS 内部的 SGLang EdgeAcc(边缘加速版,v0.5.3-edge-acc),核心能力是 Intra-GPU Disaggregation(单卡内 prefill/decode 分离,semi-PD)
  • 难点:EdgeAcc 原生只支持 dp_size=1;要在 8 卡上扩展,需要 DP(数据并行) 乃至 DP Attention,而这条路径上埋着崩溃、死锁、跨组数据串扰等一连串坑。
  • 做法:OpenSpec 规格先行 + TDD,分两个变更集(edgeacc-dp-supportedgeacc-dp-attention)逐步打通;最后用一个幂等源码补丁 apply_fix.py 修掉 DP Attention 的空闲组死锁与消息路由错误。
  • 结果:EdgeAcc 相比开源 SGLang,同负载下总吞吐 ~2×、首字延迟(TTFT)最高 ~12× 改善

1. 背景:为什么这件事很难

32B 稠密模型单卡放不下,必须 TP=8 八卡切分。但 4090 没有 NVLink,TP 的 all-reduce 只能走 PCIe/SYS,通信开销成为吞吐天花板。单纯堆 TP 并不划算,于是引入两条正交的优化:

1.1 EdgeAcc = Intra-GPU Disaggregation(单卡内 P/D 分离)

标准 vLLM/SGLang 把 prefill(计算密集、长)和 decode(访存密集、短)混在同一个调度循环里,结果是一个长 prefill 会阻塞后面一堆 decode,拉高 token 间延迟(ITL)。

EdgeAcc 的做法:在同一张 GPU 上起两个进程——一个专做 prefill,一个专做 decode,两者通过 CUDA IPC 共享 KV cache,靠 CUDA MPS 在同卡上并发共享 SM。这样 prefill 和 decode 互不阻塞,ITL/TPOT 显著下降。这就是"semi-PD disaggregation"(半分离:物理同卡、逻辑分离)。

1.2 DP / DP Attention(跨 8 卡横向扩展)

EdgeAcc 解决"单卡内"的问题,但 8 卡怎么组织?两种思路:

  • 普通 DP:每张卡(或每组卡)独立跑一套 EdgeAcc 的 P+D,DataParallelController 在多个 DP rank 之间做负载均衡(round-robin)。进程多但简单。
  • DP Attention:把 DP 和 TP 维度合并进同一个进程组(tp_size = dp_size × attn_tp_size),attention 只在 attn_tp 子组内做 all-reduce,减少进程数、减少全局通信。更省,但通信拓扑更复杂。

示例:tp_size=8, dp_size=2, attn_tp_size=4 → 8 GPU,attn_tp 子组 {0,1,2,3}=DP0{4,5,6,7}=DP1,每组内部 4 路做 attention all-reduce。


2. 部署"跑不起来"的根因

把 EdgeAcc 扩到多卡 DP,踩到的坑逐层加深:

# 症状 根因
1 dp_size>1 直接崩溃 EdgeAcc 原生只支持 dp_size=1;dp_size>1_launch_edge_acc_subprocesses 直接回退到普通 DataParallelController,完全绕过 P/D 分裂,且 else 分支没读 DP Controller 的 ready 信号,scheduler_infos[0] 空列表越界崩溃。
2 EdgeAcc + DP Attention 启动即报错 server_args.py 里有硬性互斥 assert not (enable_intra_gpu_mode and enable_dp_attention)
3 多 DP 组下结果错乱 intra_gpu_disagg_scheduler 里 4 处用了错误的 NCCL group(全局 tp_cpu_group 而非 attn_tp_cpu_group),broadcast_pyobj 的 src rank 也算错——一个 DP 组的控制消息会覆盖另一个 DP 组的数据
4 DP Attention 下服务挂死(deadlock) DP Attention 要求每个 DP 组都参与模型前向里的集合通信 dp_gather(对整个 tp 组 all_reduce)。若某个 DP 组没有请求(idle),它会跳过 run_batch,而有请求的那个组阻塞在 dp_gather 的 all_reduce 上——死锁
5 空闲批次污染 DECODE idle 组为参与集合通信而构造的"空闲批次"(不含任何请求)被错误地送进 decode 流水线。

一句话:EdgeAcc 的 P/D 分离和 DP Attention 的集合通信,在"某些 DP 组空闲"这个边界条件上没有对齐,导致要么崩溃、要么死锁、要么串数据。


3. 解决方案与实现路径(OpenSpec 驱动)

采用规格先行(OpenSpec)+ TDD,拆成两个独立可回归的变更集,每个都保证"老模式不回归":

变更集 A:`edgeacc-dp-support`(让 EdgeAcc 支持普通 DP)

架构决策:用 DP Controller 包裹 EdgeAcc,复用成熟的负载均衡逻辑,零影响普通 DP。

TokenizerManager
    │ send_to_scheduler → s_scheduler_input_ipc_name
    ▼
DataParallelController
    │ round_robin dispatch → workers[dp_rank] (AggregatedSocket)
    ├──────────────┬──────────────┐
    ▼              ▼              ▼
  DP Rank 0      DP Rank 1      DP Rank N
  ┌────┬────┐   ┌────┬────┐
  │ P  │ D  │   │ P  │ D  │   ← 每 DP rank 内:EdgeAcc P/D 分裂
  └─┬──┴─┬─┘   └─┬──┴─┬─┘
    IPC Queue      IPC Queue    ← CUDA IPC 共享 KV cache

关键任务:

  1. Server Args 校验:enable_intra_gpu_mode + dp_attention 互斥、+ nnodes>1 互斥,提前报错而非运行时崩。
  2. 启动逻辑重构:把 engine.py_launch_{decode,prefill}_processes 抽到 intra_gpu_disagg.py,参数化 dp_rankbase_gpu_offset;dp_rank=None 时行为与 dp_size=1 完全一致(保证不回归)。
  3. DataParallelController 新增 EdgeAcc 分支:elif server_args.enable_intra_gpu_mode:launch_edge_acc_dp_schedulers;worker 用 AggregatedSocket([d_socket, p_socket])(因为 decode 也需要原始请求来构建 Req 对象)。
  4. 修复 dp_size>1 ready 信号读取,消除 scheduler_infos 空列表崩溃。

变更集 B:`edgeacc-dp-attention`(让 EdgeAcc 兼容 DP Attention)

核心洞察:DP Attention 的通信(attn_tp 子组内 all_reduce)与 EdgeAcc 的 P/D 分离(独立 NCCL 组)是正交的,架构上可以共存,不必互斥。

关键任务:

  1. 移除互斥 assert(变更 A 里加的校验,在这里放开)。
  2. 修 4 个 NCCL group/src Bug:统一用 attn_tp_cpu_group,broadcast src 统一公式 attn_dp_rank * attn_tp_size(dp_attention=False 时恒为 0,向后兼容)。
  3. DataParallelController 四路分支:在 if enable_dp_attention 之前插入 if enable_dp_attention and enable_intra_gpu_mode,control_message_step = tp_size
  4. 新增 launch_edge_acc_dp_attention_schedulers:创建 IntraGPUPortArgs,共享 tokenizer/detokenizer IPC,启动 tp_size 个 P+D 进程对,所有进程在同一个 NCCL 组内按 attn_tp 子组划分。
  5. 端到端集成测试:--enable-edge-acc --enable-dp-attention --tp 8 --dp 2 能启动、curl 正确返回、且 EdgeAcc+TP、EdgeAcc+DP 两个老模式都不回归。

4. 关键代码:`apply_fix.py`——修死锁的那一刀

最后打通 DP Attention 的,是一个幂等源码补丁脚本(今日产出),对 intra_gpu_disagg_scheduler.py 做 7 处精确替换,每处校验"恰好命中 1 次"再落盘,先备份 .bak_dp_attn。核心是三件事:

4.1 空闲 DP 组也要参与集合通信(修死锁)

Prefill 调度器的 get_next_batch_to_run 改为:先算本地批次,再无条件走 MLP 同步——即使本组空闲,也会被 prepare_mlp_sync_batch 构造成一个 idle 批次,从而加入 dp_gather 集合,避免拖死有请求的组。

def get_next_batch_to_run(self) -> Optional[ScheduleBatch]:
    # 计算本 rank 的本地 prefill 批次(空闲 DP 组可能为 None)
    ret = self._get_next_prefill_batch_local()
    # DP Attention:跨所有 attention DP 组同步。每个 rank 都必须 all-gather 自己的
    # token 数,这样没有工作的 DP 组也会构造一个 idle 批次并加入模型前向里的集合
    # dp_gather。否则空闲组会跳过 run_batch,把真正有工作的组死锁在 dp_gather 的
    # all_reduce 上(它要对整个 tp 组做 reduce)。
    if require_mlp_sync(self.server_args):
        ret = self.prepare_mlp_sync_batch(ret)
    return ret

def _get_next_prefill_batch_local(self) -> Optional[ScheduleBatch]:
    """原来的 get_next_batch_to_run 逻辑改名于此。"""
    ...

Decode 调度器同理:

def get_next_batch_to_run(self) -> Optional[ScheduleBatch]:
    ret = self.get_next_batch_to_decode()
    # DP Attention:同步所有 attention DP 组,让空闲组构造 idle 批次加入集合 dp_gather。
    if require_mlp_sync(self.server_args):
        ret = self.prepare_mlp_sync_batch(ret)
    return ret

同时删掉旧的手动 global_num_tokens 拼装(prefill/decode 各一处)——那套是 dp_size=1 时的补丁,现在由 prepare_mlp_sync_batch 统一接管。

4.2 空闲批次不许进 DECODE(修流水线污染)

def process_batch_result_prefill(self, batch, result, launch_done=None):
    # DP Attention:idle 批次只是为了让本(空闲)DP 组加入集合前向(dp_gather),
    # 它不含任何请求,绝不能被转发到 DECODE。解析掉 overlap 结果保持流水线一致后直接返回。
    if batch.forward_mode.is_idle():
        if self.enable_overlap:
            self.tp_worker.resolve_last_batch_result(launch_done)
            self.set_next_batch_sampling_info_done(batch)
        return
    next_token_logits = None
    ...

4.3 请求接收按 DP 子组路由(修跨组串扰)

Decode 调度器重写 recv_requests:只在 attn_tp 子组内广播,而不是基类那样在整个 tp 组广播(否则会覆盖其它 DP 组的 per-group 数据,破坏路由)。

def recv_requests(self) -> List[Req]:
    """EdgeAcc DECODE 调度器的请求接收。
    EdgeAcc + DP Attention 下,每条消息按 DP 组路由:decode 收到来自 DP controller 的
    原始 generate 请求 + 来自配对 prefill 的结果/元数据/清理消息(走 per-dp_rank socket)。
    这些只能在 attn_tp(DP)子组内共享。基类 Scheduler 会额外在整个 tp 组广播控制消息,
    那会覆盖其它 DP 组的数据、破坏 per-DP-group 路由,所以这里只在 attn_tp 子组内广播。
    """
    if self.attn_tp_rank == 0:
        recv_reqs = []
        while True:
            try:
                recv_req = self.recv_from_tokenizer.recv_pyobj(zmq.NOBLOCK)
            except zmq.ZMQError:
                break
            recv_reqs.append(recv_req)
        if self.recv_from_rpc is not None:
            while True:
                try:
                    recv_rpc = self.recv_from_rpc.recv_pyobj(zmq.NOBLOCK)
                except zmq.ZMQError:
                    break
                recv_reqs.append(recv_rpc)
    else:
        recv_reqs = None

    if self.attn_tp_size > 1:
        recv_reqs = broadcast_pyobj(
            recv_reqs,
            self.attn_tp_group.rank,
            self.attn_tp_cpu_group,
            src=self.attn_tp_group.ranks[0],   # 子组内 src,而非全局 rank 0
        )
    return recv_reqs

补丁自校验风格(避免"以为改了其实没命中"):

for i, (old, new) in enumerate(replacements, 1):
    count = content.count(old)
    if count != 1:
        print(f"FAIL: replacement #{i} found {count} occurrences (expected 1)")
        sys.exit(1)
    content = content.replace(old, new)
    print(f"OK: replacement #{i} applied")

5. 启动与部署实现

EdgeAcc 版是环境变量驱动的,启动脚本 sglang_launch_server.sh 负责把 env 翻成 sglang.launch_server 参数,并在开启 EdgeAcc/semi-PD 时拉起 CUDA MPS:

# CUDA ≥ 12.8 用 multiuser-server 模式,并开启 per-context SM 分区
start_mps() {
    unset CUDA_VISIBLE_DEVICES               # 让 MPS 看到所有 GPU
    export CUDA_MPS_ENABLE_PER_CTX_DEVICE_MULTIPROCESSOR_PARTITIONING=1
    if version_ge "$(get_cuda_version)" "12.8"; then
        nvidia-cuda-mps-control --multiuser-server -d
    else
        nvidia-cuda-mps-control -d
    fi
    # 恢复原 CUDA_VISIBLE_DEVICES
}

[ "$ENABLE_EDGE_ACC" == "True" ]       && { start_mps; SERVER_ARGS+=" --enable-edge-acc "; }
[ "$ENABLE_SEMI_PD_DISAGG" == "True" ] && { start_mps; SERVER_ARGS+=" --enable-semi-pd-disagg "; }

cd /sgl-workspace/sglang && python3 -m sglang.launch_server \
  --model-path ${MODEL_PATH} --served-model-name ${SERVED_MODEL_NAME} \
  --host ${SERVER_HOST} --port ${SERVER_PORT} --trust-remote-code \
  --tp ${TENSOR_PARALLEL_SIZE} ${SERVER_ARGS}

容器启动(host 网络、privileged、ipc host,挂载 /data),EdgeAcc 关键 env:

docker run -d --gpus all --privileged --network host --ipc host \
  -e SGLANG_CACHE_ROOT=/data/sglang/cache \
  -e MODEL_PATH=/data/models/Qwen3-32B -e TENSOR_PARALLEL_SIZE=8 \
  -e ENABLE_EDGE_ACC=True \
  -e TORCHINDUCTOR_COMPILE_THREADS=64 -e TORCHINDUCTOR_FX_GRAPH_CACHE=1 \
  -e SERVER_ARGS=" --enable-p2p-check --enable-torch-compile --torch-compile-max-bs 64 --disable-qps-limit " \
  alien-registry.../sglang:v0.5.3-edge-acc-...-release

EdgeAcc 相比开源版的加速点还叠加了:torch.compile / TorchInductor 图编译(算子融合 + CUDA graph)、FlashAttention-3 kernelNVSHMEM / GDRCOPY(优化无 NVLink 拓扑下的 TP 通信)、P2P 检查


6. 效果(EdgeAcc vs 开源 SGLang)

同一长上下文负载(128 并发,~2860 in / ~2180 out tokens,2560 请求,Qwen3-32B):

指标 开源 SGLang EdgeAcc 提升
压测总时长 2045 s 1014 s 2.0×
总吞吐 6289 tok/s 12702 tok/s 2.0×
输出吞吐 2732 tok/s 5487 tok/s 2.0×
平均 E2E 延迟 98.5 s 40.6 s 2.4× ↓
中位 TTFT 15028 ms 1263 ms 11.9× ↓
P99 ITL 1843 ms 111 ms 16× ↓

结论:P/D 分离把长 prefill 从 decode 路径上摘掉,首字延迟和尾延迟改善最猛;图编译 + 定制通信 kernel 把无 NVLink 下的 TP 瓶颈压下去,吞吐翻倍。


7. 经验总结

  1. 消费卡集群的推理优化,瓶颈在通信而非算力。无 NVLink 时,一切能减少/规避跨卡 all-reduce 的架构(P/D 分离、DP Attention 子组通信、NVSHMEM/GDRCOPY)都值钱。
  2. P/D 分离的收益在"解耦":prefill 计算密集、decode 访存密集,混跑互相拖累;同卡两进程 + CUDA IPC 共享 KV + MPS 共享 SM,是在单机上低成本拿到分离收益的办法。
  3. 分布式推理最隐蔽的 bug 是"边界条件下的集合通信不对齐":某个组空闲、某条消息走错 group、src rank 算错——不崩则死锁,不死锁则串数据。修法的共性是让所有 rank 在集合通信点上行为一致(空闲也要造 idle 批次入群)、把通信严格限定在正确的子组
  4. 规格先行 + 每步保证老模式不回归,是在别人的大型引擎源码上做侵入式改造时唯一能收敛的方法。幂等、可校验、可回滚的补丁脚本(count==1 断言 + .bak)比手改可靠得多。