Leoda

6. flagcx connector 升级至 vllm0.22

7 min read

1. 背景

早期在开发 flagcx connector 的时候参考 vllm 版本为0.13.0，现在的 vllm 适配国产网卡的过程中主要 vllm 版本分别为 dcu 的 0.20.0 和沐曦的 0.20.2，故需要一个 connector 能够统一这两个版本的 vllm 以及QWEN3.6(暂时)。

2. 参考与调研

由attention + Mamba 混合架构的问题内的 9 条 bug 记录，总结为需要升级增加的功能包括：

  • 支持导入新版本的 vllm 的 attention 模块 和 支持新版本的 vllm topo 库(传 list 的 backend)，支持get_attn_backend等等函数的新签名 ✅done
  • 支持HMA ✅done
  • 从 router 进来之后的一个请求获取的 block_ids变成 2 维之后的完整透传以及原来所有一维索引变成 二维来对应。✅done
  • 支持 prefill 的 tp > decode 的 tp ✅done
  • 支持混合架构模型（mamba 层 kv cach、GDN 层的 kv cache等 ），见下面 2.1 的描述 ✅done

2.1 Hybrid attention and Hybrid KV Cache Manager

参考 https://huggingface.co/Qwen/Qwen3.6-35B-A3B/blob/main/config.json 官方的 config.json，看到模型的 layers 组成为[linear_attn x 3 + full_attn] x 10来组成 40 个 layers。这需要 connector 调用Hybrid KV Cache Manager去支持这类Hybrid model。

在 vllm 0.20 内的 nixl 支持了 mamba，vllm 0.22 内的nixl 支持了 GDN(qwen3.6需要)。上周的 mooncake 的 pr 支持了 GDN。所谓 GDN 就是：

/Users/joker/Desktop/project/baai/vllm/vllm/v1/worker/gpu_model_runner.py对一分为 2， （注：gdn，sw 等等都归于 mamba 类别)。

  • AttentionSpec的话，当前就是 kv_caches[layer_name] = kv_cache.permute(*inv_order)
  • MambaSpec的话，当前就是for 循环变量后的里面装了 conv 和 ssm 的 list

image.png|562

GDN 通过当前输入进入一个 conv 卷积看到局部上下文，再通过学习到的 delta 去 ssm 全局记忆里面拿出想要的 cache 去 decode 下一个 token。

故我们的传输 kv cache 就是需要传输 conv+ssm，其中 ssm 在 nixl/mooncake 中都传输 N-2 个 token 更新后的 ssm，由decode 自己算N-1 的状态。然后 vllm 内把这两个存 cache 的 [conv,ssm]放在一个 pages 里面，从过第一个 conv 的 lens 去索引 ssm 的 base addr。ssm 的尾巴后面自动 padding 到整页末尾。因此，connector 内的 get_num_new_matched_tokens 函数内对这种情况需要对应 -1 个 token。

当然注册的时候，直接对当前这个 list 内整个pages 注册。如图

image.png

┌───────────────────────┬─────────────┬─────────────┬──────────────┬────────────────┐
│     层类型 / 场景       │ split_k_and │ blocks_firs │   拆分路径   │ 每层 region 数 │
│                       │     _v      │      t      │              │                │
├───────────────────────┼─────────────┼─────────────┼──────────────┼────────────────┤
│ attention（纯注意力模   │             │             │ 机制         │                │
│ 型, layer-first        │ True        │ False       │ A：[K,V]     │ 2              │
│ 后端）                  │             │             │ 两条         │                │
├───────────────────────┼─────────────┼─────────────┼──────────────┼────────────────┤
│ attention（GDN         │             │             │ 机制 B：登记   │                │
│ 混合模型里, 因含         │ False       │ True        │  1 条,展开劈   │ 2              │
│ mamba→blocks-first ）  │             │             │  K/V         │                │
├───────────────────────┼─────────────┼─────────────┼──────────────┼────────────────┤
│                       │             │             │ 都不拆（[con   │ 1（整页        │
│ Mamba / GDN 层         │ False       │ True        │ v] + split=F │ conv+ssm）     │
│                       │             │             │ alse）        │                │
├───────────────────────┼─────────────┼─────────────┼───────────── ─┼────────────────┤
│ MLA 层                 │ False       │ False(is_ml │ 都不拆       │ 1（整页 page_s │
│                       │             │ a)          │              │ ize_bytes）    │
└───────────────────────┴─────────────┴─────────────┴──────────────┴────────────────┘

3. 测试与验证

basic

基础安装指令
# step 0. 直接拉
docker pull vllm/vllm-openai:v0.20.2
 
docker run -d -it --network=host --gpus all --privileged --ipc=host  --entrypoint bash --ulimit memlock=-1 --ulimit stack=67108864 -v /public-flash/zhiyuan/flagcx/:/workspace/  --name 20-liuda vllm-openai:v0.20.2-liuda
 
apt update && apt install -y --no-install-recommends net-tools vim git zoxide wget
apt-get update && apt-get install -y libhwloc15
pip install -U scikit-build-core==0.11 pybind11 ninja cmake
 
 
# step 1. vllm下载与编译（pass）
# git clone --branch v0.20.2 --depth 1 https://github.com/vllm-project/vllm.git
# cd vllm
# uv venv --python 3.12 .venv
# source /workspace/liuda/vllm/.venv/bin/activate
# VLLM_USE_PRECOMPILED=1 uv pip install -e . --torch-backend=auto -v放弃了 这个源码安装慢
# uv pip install vllm==0.20.2
# deactivate
 
# step 2. 模型
modelscope download --model Qwen/Qwen3.6-35B-A3B
/workspace/liuda/models/qwen/Qwen3.6-35B-A3B
 
# step 3. 仓库, plugin内的改动放在自己的仓库内的
https://github.com/leoda1/vllm-plugin-FL.git
cd vllm-plugin-FL
git checkout new
pip install --no-build-isolation -e .
 
git clone https://github.com/flagos-ai/FlagCX.git
git submodule update --init --recursive
 
export CPATH=/usr/local/lib/python3.12/site-packages/nvidia/cu13/include:$CPATH
export LD_LIBRARY_PATH=/usr/local/cuda:$LD_LIBRARY_PATH
pip install . -v --no-build-isolation
 
git clone https://github.com/flagos-ai/FlagGems
cd FlagGems
pip install --no-build-isolation -e .

pd step（确保前面 0 error后执行这里的）

step1. 在节点 1 上起 prefill，注意这里多了一个环境变量 --no-disable-hybrid-kv-cache-manager

export NCCL_SOCKET_IFNAME=bond0
export TORCH_DEVICE_BACKEND_AUTOLOAD=0
export NCCL_IB_GID_INDEX=3
export FLAGCX_DEBUG=version
export VLLM_RPC_TIMEOUT=600000
export FLAGCX_PATH=/workspace/liuda/FlagCX
export VLLM_ENGINE_ITERATION_TIMEOUT_S=600
export PYTHONPATH=/workspace/liuda/vllm-plugin-FL:${PYTHONPATH}
export VLLM_FL_PREFER_ENABLED=false
export FLAGCX_P2P_SLICE_SIZE=65536
export VLLM_PLUGINS=fl
vllm serve /workspace/liuda/models/qwen/Qwen3.6-35B-A3B \
    --host 0.0.0.0 \
    --port 20001 \
    --tensor-parallel-size 8 \
    --seed 1024 \
    --max-model-len 40960 \
    --served-model-name base_model \
    --max-num-batched-tokens 65536 \
    --max-num-seqs 256 \
    --trust-remote-code \
    --no-disable-hybrid-kv-cache-manager \
    --kv-cache-dtype fp8 \
    --gpu-memory-utilization 0.8 \
    --kv-transfer-config \
    '{"kv_connector":"FlagCXConnector","kv_role":"kv_producer"}' > prefill-1.log 2>&1 &

step2. 在节点 2上起 decode

export NCCL_SOCKET_IFNAME=bond0
export GLOO_SOCKET_IFNAME=bond0
export NCCL_IB_GID_INDEX=3
export VLLM_RPC_TIMEOUT=600000
export FLAGCX_DEBUG=version
export FLAGCX_DEBUG_SUBSYS=
export FLAGCX_PATH=/workspace/liuda/FlagCX
export VLLM_ENGINE_ITERATION_TIMEOUT_S=600
export PYTHONPATH=/workspace/liuda/vllm-plugin-FL:${PYTHONPATH}
export VLLM_FL_PREFER_ENABLED=false
export VLLM_PLUGINS=fl
export FLAGCX_P2P_SLICE_SIZE=65536
vllm serve /workspace/liuda/models/qwen/Qwen3.6-35B-A3B \
    --host 0.0.0.0 \
    --port 20002 \
    --tensor-parallel-size 8 \
    --seed 1024 \
    --max-model-len 40960 \
    --served-model-name base_model \
    --max-num-batched-tokens 65536 \
    --max-num-seqs 256 \
    --trust-remote-code \
    --no-disable-hybrid-kv-cache-manager \
    --kv-cache-dtype fp8 \
    --gpu-memory-utilization 0.8 \
    --kv-transfer-config \
    '{"kv_connector":"FlagCXConnector","kv_role":"kv_consumer"}'> decode.log 2>&1 &

step3. 起 router，这里面的 router.py在 vllm-plugin-FL/examples/disaggregated_serving_xpyd/router.py

python3 router.py \
  --host 0.0.0.0 \
  --port 8000 \
  --prefill http://10.0.27.1:20001 8998 \
  --decode http://10.0.27.10:20002 > router.log 2>&1 &

step 4. 测试：

curl -v http://127.0.0.1:8000/v1/completions \
  -H 'Content-Type: application/json' \
  -d '{
    "model": "base_model",
    "prompt": "100 + 100 =",
    "max_tokens": 128,
    "stream": true
  }'
#!/bin/bash
SLEEP_BETWEEN=5
unset http_proxy
unset https_proxy
unset HTTP_PROXY
unset HTTPS_PROXY
export VLLM_SERVER_DEV_MODE=1
BASE_ARGS="vllm bench serve \
  --base-url http://10.0.27.10:8000 \
  --endpoint /v1/completions \
  --model /workspace/liuda/models/qwen/Qwen3.6-35B-A3B \
  --served-model-name base_model \
  --dataset-name random \
  --num-warmups 20 \
  --percentile-metrics ttft,tpot,itl,e2el \
  --metric-percentiles 50,90,99"
 
run_test() {
  local input_len=$1
  local output_len=$2
  local rps=$3
  local num_prompts=$(( rps * 5 ))
  echo "========================================"
  echo "Testing input=${input_len} output=${output_len} rps=${rps} num_prompts=${num_prompts}"
  echo "========================================"
  $BASE_ARGS \
    --random-input-len  "$input_len"  \
    --random-output-len "$output_len" \
    --request-rate      "$rps"        \
    --num-prompts       "$num_prompts"
  # echo "Sleeping ${SLEEP_BETWEEN}s before next test..."
  # curl -s -X POST http://10.0.26.138:8000/reset_prefix_cache && echo "OK"
  # sleep "$SLEEP_BETWEEN"
}
 
# input        output rps
run_test  1024   1024   75
#run_test  4096   1024   70
#run_test  8192   1024   65
#run_test  16384  1024   60
#run_test  32768  1024   45
 
echo "All tests done."

image.png|528

bug

1. flashinfer报错（放弃，直接用官方的镜像跑起来更容易）

如果报错：

WARNING 07-08 03:03:45 [gdn_linear_attn.py:733] /workspace/liuda/vllm/.venv/lib/python3.12/site-packages/flashinfer/data/cutlass/include/cutlass/cuda_host_adapter.hpp(150): error: identifier "PFN_cuTensorMapEncodeTiled" is undefined

需要去下一个正确的头文件：

 
FILE=/workspace/liuda/vllm/.venv/lib/python3.12/site-packages/flashinfer/data/cutlass/include/cutlass/cuda_host_adapter.hpp
 
cp "$FILE" "$FILE.bak"
curl -sL https://raw.githubusercontent.com/NVIDIA/cutlass/main/include/cutlass/cuda_host_adapter.hpp -o "$FILE"
rm -rf /root/.cache/flashinfer

flashinfer 的 gdn jit 零时编译的各类头文件都是 hardcode 的，没法虚拟环境里面每次报错少一个外部的头文件就 export 一个。。。

Graph View

Backlinks