bugfix: prevent SHM H2D races during graph replay.

[pjgao]

Description

问题

Qwen3.5 在 TP2、MTP、draft body TP1、ACL Graph、schedule overlap、SHM 和 HCCL AIV 同时开启时，可能出现两个 TP rank 进入不一致的等待顺序，表现为请求长期不返回等异常情况（大部分情况没问题，少数异常场景有问题，可复现场景见下文脚本）。

该问题在 MTP5 压力配置下可以稳定观察到；MTP 会提高每轮输入准备和 speculative state 更新频率，因此更容易放大原有的跨线程时序问题。

问题来源

PR #1964 为消除 Qwen3 Gated DeltaNet / CausalConv1d 前的 H2D 空泡，在 deserialize_forward_input_payload() 中提前创建 linear_state_indices：

if (!input_params.embedding.linear_state_ids.empty()) {
  input_params.embedding.linear_state_indices =
      torch::tensor(input_params.embedding.linear_state_ids, torch::kInt)
          .to(device, /*non_blocking=*/true);
}

这段代码的数据和值没有问题，问题在于执行位置发生了变化：

• 修改前，缺失的 linear_state_indices 在 worker 的 ModelInputParams::to(device) 阶段按需创建，受 worker prepare / graph task 的既有顺序约束；
• 修改后，在 SHM 模式下，这次 non-blocking H2D 会从 SHM polling thread 发起，不再受 worker ordered prepare 的顺序约束；
• polling-thread H2D 与计算线程的 graph task update/replay 并发后，两个 TP rank 可能进入不同的等待顺序，最终导致请求挂起。

根因

根因是 NPU Graph + schedule overlap + SHM 路径把 device materialization 提前到了 SHM polling thread。

POSIX SHM payload 可被 producer 复用，而异步 H2D 又要求源 host 数据在传输完成前保持有效；同时，graph replay 和输入准备需要遵循 worker 的有序执行关系。原实现既没有把这次 H2D 纳入 worker ordered prepare，也没有为异步传输建立独立、稳定且由 ForwardInput 持有的 host snapshot，因此破坏了该路径原有的生命周期和执行顺序约束。

修改

本 PR 对这条路径做了整体调整：

• 在 NPU enable_graph && enable_schedule_overlap 场景中，SHM polling thread 只读取 payload，不再发起 device materialization；
• 把可复用的 POSIX SHM payload 复制到由 ForwardInput 持有的真实 pinned host buffer，确保异步 H2D 期间源数据生命周期稳定；
• 在 worker ordered prepare 中重新解包 owned host snapshot，并完成 device materialization，使 H2D、graph task update 和 replay 回到同一有序路径；
• 所有从 host payload 构造的 tensor view 都以最终 owned pinned buffer 为基址，保证 view 与 owner 始终对应；
• linear_state_indices 仍在模型执行前完成准备，但不再由 polling thread 提前发起 H2D。

最终产品代码修改 3 个文件，共 +44/-13。

原来：Polling thread 提前 H2D，与 Graph replay 并发，造成 TP rank 时序分叉和请求挂起。
改后：保存 owned pinned snapshot，H2D 回到 Worker Prepare Stream，恢复有序执行。

影响范围与验证边界

• NPU Graph + schedule overlap + SHM：device materialization 从 polling thread 移到 worker ordered prepare

最小复现

下面的脚本使用Qwen3.5 2B的TP2复现卡住问题，需要配置 target/draft 模型路径和 xLLM binary。模型路径、设备号和端口均为占位符，不包含机器或用户信息。客户端使用 EvalScope 1.9.1+ 的固定 prompt，不依赖外部数据集。

1. 启动服务

#!/usr/bin/env bash
set -euo pipefail

BIN=${BIN:?set BIN to the xLLM executable}
MODEL_PATH=${MODEL_PATH:?set MODEL_PATH}
DRAFT_MODEL_PATH=${DRAFT_MODEL_PATH:?set DRAFT_MODEL_PATH}
DEVICE_0=${DEVICE_0:-0}
DEVICE_1=${DEVICE_1:-1}
API_PORT=${API_PORT:-18000}
RANK1_PORT=${RANK1_PORT:-18001}
MASTER_PORT=${MASTER_PORT:-28000}

source /usr/local/Ascend/cann-9.0.0/set_env.sh
source /usr/local/Ascend/nnal/atb/9.0.0/set_env.sh
export ASCEND_CUSTOM_OPP_PATH=/usr/local/Ascend/cann-9.0.0/opp/vendors/custom_xllm_math:
export XLLM_HCCL_TP_AIV=1
export XLLM_NPU_REDUCE_VIA_ALLGATHER_MAX_NUMEL=0
export TASK_QUEUE_ENABLE=1
export HCCL_BLOCKING_WAIT=0

args=(
  --model "$MODEL_PATH"
  --nnodes=2
  --master_node_addr="127.0.0.1:$MASTER_PORT"
  --max_memory_utilization=0.80
  --max_tokens_per_batch=32768
  --max_seqs_per_batch=32
  --max_concurrent_requests=32
  --block_size=128
  --communication_backend=hccl
  --enable_prefix_cache=false
  --enable_chunked_prefill=true
  --enable_graph=true
  --enable_schedule_overlap=true
  --enable_shm=true
  --draft_model "$DRAFT_MODEL_PATH"
  --num_speculative_tokens=5
  --enable_mtp_draft_body_tp1=true
  --enable_opt_validate_probs=true
  --enable_graph_double_buffer=false
  --enable_graph_mode_decode_no_padding=true
  --max_linear_state_cache_slots=0
)

ASCEND_RT_VISIBLE_DEVICES="$DEVICE_0" \
  "$BIN" "${args[@]}" --node_rank=0 --port="$API_PORT" &
pid0=$!
ASCEND_RT_VISIBLE_DEVICES="$DEVICE_1" \
  "$BIN" "${args[@]}" --node_rank=1 --port="$RANK1_PORT" &
pid1=$!
trap 'kill "$pid0" "$pid1" 2>/dev/null || true' EXIT INT TERM
wait

2. 发送请求

服务 ready 后在另一个终端执行：

evalscope perf \
  --model Qwen35-27B \
  --api openai \
  --url http://127.0.0.1:18000/v1/chat/completions \
  --prompt '请分析异步内存生命周期。' \
  --apply-chat-template \
  --parallel 16 \
  --number 64 \
  --max-tokens 512 \
  --temperature 1.0 \
  --no-stream \
  --read-timeout 900 \
  --total-timeout 900

判定标准：

• main/RED：请求长期不返回，两个 rank 仍存活但服务不再推进；
• PR/GREEN：完成 64/64，两个 rank 无目标错误，并在压测后继续通过 /v1/models 和一次真实生成请求。

（speculative token broadcast 会改变 rank 时序和复现概率）

基线验证

测试锁定基线为 origin/main@6bc3fd249c81651cc79e7e38dd39a2e36c52231b；该提交只用于保留当次 RED/GREEN 数字，不代表当前最新 main。

版本	Warmup	正式请求	Server E2E (ms)	Server TTFT (ms)	Server TPOT (ms)	接收率	结论
origin/main@6bc3fd24 + 相同机器 patch	8/8	0（首条卡住）	N/A	N/A	N/A	N/A	HANG_REPRODUCED
已验证 PR 快照 4fd3343b	8/8	60/60	137.440	72.000	2.727	63.018%	PASS

main baseline 在 8 条 warmup 后，首个正式请求超过 120 秒没有完成；已验证 PR 快照稳定完成 60/60，并在压测后保持服务可用。

二进制 SHA256：

• main baseline：9160ce348094da4b197a982e3e50f618247e5724b0d7974d4d056a6d43f757f2
• PR head：2e37c4f230d426600c99390efcd60eed8ed8917a37f39de1d902020a03f8d6ff

性能验证

性能验证使用 Qwen3.5-2B BF16、target TP2、draft body TP1、MTP5、HCCL AIV、Graph、schedule overlap、SHM、optimized validate probs、graph decode no padding；before/after 使用同一对 NPU、相同 CPU affinity、权重、机器 patch 和构建参数；log_request 数据集，temperature 使用 API 默认值 0.0，warmup 8，正式 3×20，请求输出 25 tokens。判退阈值为 3%。

完整 PR 相对锁定 main 的结果见上表：main 在首个正式请求挂起，没有可比较的性能数字；PR head 完成 60/60，Server E2E 137.440 ms、TTFT 72.000 ms、TPOT 2.727 ms，未观察到逐轮退化。由于 baseline 没有完成正式请求，本结果只能证明修复后的绝对性能与稳定性，不能计算 before/after 性能增益。
从代码白盒分析，该修改对性能无影响。

Related Issues

Related to #1965.

Change Type

• Bug fix
• New feature
• Performance improvement
• Refactor
• Documentation
• Test
• Build or CI

Pull Request Checklist

PR Title and Commit Messages

• The PR title and each commit message follow the xLLM commit format: : .

Pre-commit Checks

• I have installed pre-commit by running pip install pre-commit or an equivalent command.
• I have installed the hooks with pre-commit install.
• I have run pre-commit run --all-files and fixed any reported issues.

Self Review

• I have self-reviewed the code according to .agents/skills/code-review/references/custom-code-style.md, especially code written or assisted by AI.
• I have rebased this PR onto the latest main branch.

Build and Test Coverage

• Tests have been added or updated as needed.
• CUDA: python setup.py build test has passed on a CUDA machine.
• NPU: python setup.py build test has passed on an NPU machine.
• MLU: python setup.py build test has passed on an MLU machine.

Reviewer Notes

PR #1971 修复随机采样下的跨 TP rank speculative token 分叉；本 PR 修复 SHM host ownership 与异步 H2D ordering 路径。两者可能表现为相似的 MTP 挂起，但根因不同。

The three development commits were rebased onto main@49651e8f and squashed into 1a0c2cd. The final three-file diff remains +44/-13, and its stable patch-id is unchanged (d7e4ac7d1ab659e8ae8db84c77438296a53c2202). The validated runtime binaries above correspond to the pre-squash snapshot with the same product diff; the rebased head has not been rebuilt, so the NPU build/test checkbox remains unchecked.

[yingxudeng]

!XXXX 这种语义是不是不太直观，bool defer_graph_overlap_h2d 换成类似 如下可读性会不会好一些（类似原来的 ForwardSyncPolicy、RegisterLengthPolicy）。具体叫什么名字您可以自己再看看。 enum class InputDeviceCopyPolicy : int8_t {
COPY_TO_DEVICE_ON_READ = 0,
DEFER_COPY_TO_PREPARE_STREAM = 1,
};

[yingxudeng]

这个分支是不是会影响所有模型？ deepseek v4 是否还能正常推理？