Prhub

#47481 [ROCm][CI] Adding nixl multiconn

原始 PR 作者 AndreasKaratzas 合并时间 2026-07-06 15:04 文件变更 4 提交数 1 评论 3 代码增减 +88 / -7

ci/build rocm v1 kv-connector

执行摘要

新增 ROCm CI 测试：Nixl 多连接器前缀缓存与精度覆盖

PR body 明确指出需要为 MI300 ROCm CI 增加 Hybrid SSM NixlConnector 前缀缓存覆盖以及 MultiConnector（Nixl+Offloading）精度和边界情况测试，以提升 KV 传输在 ROCm 平台上的可靠性验证。同时使 NIXL 集成脚本接受 ATTENTION_BACKEND 和 VLLM_SERVE_EXTRA_ARGS，便于跨后端复用。

该 PR 主要聚焦 CI 基础设施，值得关注其创建的参数化测试脚本模式，便于后续跨后端测试复用。但核心价值有限，建议在合并后考虑将关键测试步骤改为非 optional 以发挥持续验证作用。

讨论亮点
  • tjtanaa 指出新增的测试步骤均标记为 optional: true，不会在 AMD CI 中自动触发，需要手动触发并确认通过后才能合并。
  • AndreasKaratzas 随后确认所有测试已手动启用并通过，tjtanaa 批准合并。
  • 此外有 pre-commit 检查失败提示，但已在后续修复。

实现拆解

  1. 测试脚本参数化：在 run_mamba_prefix_cache_test.sh、run_multi_connector_accuracy_test.sh 和 run_multi_connector_edge_case_test.sh 中新增 ATTENTION_BACKEND 和 VLLM_SERVE_EXTRA_ARGS 环境变量读取，并将原有的硬编码 --attention-backend FLASHINFER 替换为通过 EXTRA_ARGS 动态传入，使得同一脚本可为不同注意力后端复用。
  2. 仓库根目录解析修复：在 run_multi_connector_accuracy_test.sh 和 run_multi_connector_edge_case_test.sh 中，将 GIT_ROOT=$(git rev-parse --show-toplevel) 替换为基于脚本目录的相对路径解析，避免容器内无 git 元数据时命令失败。
  3. CI 配置新增任务：在 .buildkite/test-amd.yaml 中插入 4 个新的可选测试步骤（Hybrid SSM 前缀缓存测试、MultiConnector 精度测试、MultiConnector 边界情况测试），均使用 TRITON_ATTN 后端、2-GPU MI300 池，并正确设置了 source_file_dependencies 以确保变更触发。
  4. 测试覆盖范围：新增的前缀缓存测试验证 Mamba 混合模型在 PD 分离下的前缀缓存命中率；MultiConnector 精度测试运行 gsm8k 评估；边界情况测试覆盖 block-size 边界、缓存命中/冷启动/逐出等场景，并校验 Prometheus 指标。
文件 模块 状态 重要度
tests/v1/kv_connector/nixl_integration/run_mamba_prefix_cache_test.sh 测试脚本 modified 5.07
.buildkite/test-amd.yaml CI 配置 modified 4.47
tests/v1/kv_connector/nixl_integration/run_multi_connector_accuracy_test.sh 测试脚本 modified 4.47
tests/v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh 测试脚本 modified 4.47

关键源码片段

tests/v1/kv_connector/nixl_integration/run_mamba_prefix_cache_test.sh test-coverage

最重要的测试脚本，演示了如何通过环境变量实现注意力后端动态配置

#!/bin/bash
# run_mamba_prefix_cache_test.sh – 支持动态注意力后端的 Mamba 前缀缓存测试
​
# 新增环境变量，允许外部覆盖默认后端和额外参数
ATTENTION_BACKEND=${ATTENTION_BACKEND:-FLASHINFER} # 默认保持向后兼容
VLLM_SERVE_EXTRA_ARGS=${VLLM_SERVE_EXTRA_ARGS:-}
​
# 解析额外参数：将逗号分隔的字符串转换为数组
EXTRA_ARGS=()
if [[ -n "$VLLM_SERVE_EXTRA_ARGS" ]]; then
  IFS=',' read -r -a EXTRA_ARGS <<< "$VLLM_SERVE_EXTRA_ARGS"
fi
# 如果设置了 ATTENTION_BACKEND，追加到参数数组
if [[ -n "$ATTENTION_BACKEND" ]]; then
  EXTRA_ARGS+=(--attention-backend "$ATTENTION_BACKEND")
fi
​
# Prefill 实例启动：使用参数数组代替硬编码
CUDA_VISIBLE_DEVICES=$PREFILL_GPU_ID \
VLLM_SSM_CONV_STATE_LAYOUT=DS \
VLLM_KV_CACHE_LAYOUT=HND \
VLLM_NIXL_SIDE_CHANNEL_PORT=5559 \
vllm serve $MODEL \
  --port $PREFILL_PORT \
  --enforce-eager \
  --gpu-memory-utilization $GPU_MEMORY_UTILIZATION \
  --max-model-len 16384 \
  --block-size 128 \
  --trust-remote-code \
  --enable-prefix-caching \
  --mamba-cache-mode all \
  --kv-transfer-config "$KV_CONFIG" \
  "${EXTRA_ARGS[@]}" & # 原来的 --attention-backend FLASHINFER 被移除

评论区精华

测试步骤可选标记与合并时机 other

tjtanaa 指出新增的测试都标记为 optional: true，不会在 AMD CI 中自动触发，需要手动触发后再合并。

结论：AndreasKaratzas 手动触发所有测试并确认通过，tjtanaa 批准合并。 · 已解决

风险与影响

  • CI 覆盖盲区：所有新增测试均为 optional: true，除非手动触发，否则无法自动拦截回归，可能降低 CI 有效性。
  • 后端依赖性：测试强制使用 TRITON_ATTN 后端，若该后端在特定 ROCm 版本中不可用或存在缺陷，会导致误报。
  • 资源消耗：新增 4 个每任务 180 分钟的超时测试，可能增加 CI 集群负载，尤其在手动触发时。
  • 兼容性：测试脚本改动（如 git root 解析）仅针对容器场景，若在本地有完整 git 目录的环境运行，行为不变，风险低。
  • 用户影响：无直接影响，均为 CI 和测试变更。
  • 系统影响：ROCm CI 现在涵盖更多 KV 传输场景，有助于提前发现 PD 分离下的回归。
  • 团队影响：开发者需要手动触发这些可选测试以验证相关变更，增加了少量流程成本。
  • 影响程度：中等，主要提升测试覆盖但不改变运行时逻辑。
测试标记为 optional 可能被忽略 依赖 TRITON_ATTN 后端在 ROCm 上的可用性 CI 任务超时较长可能占用集群资源

关联 Issue

未识别关联 Issue

当前没有检测到明确关联的 Issue 链接，后续同步到相关引用后会出现在这里。

完整报告

Markdown PDF

参与讨论