feat: support heterogeneous qwen3.5 PD disaggregation

[Vectorwh]

中文说明

概述

本 PR 为昇腾 NPU 上的 Qwen3.5 推测解码增加异构 TP 的 PD 分离能力。当前支持 Prefill TP2、Decode TP1，并传输 Target KV Cache、MTP Draft KV Cache、Causal Conv State 和 SSM State。

Decode 收到 Prefill 生成的首 token 后，可直接接管请求并继续生成，不再重复执行 Prefill。

本 PR 还将 Decode 对两个 Prefill 分片的拉取改为并行执行。并行路径默认开启，同时保留运行时回退开关。

背景与问题

原有 PD 路径默认 Prefill 和 Decode 的 Cache 布局兼容。Qwen3.5 同时包含 Attention KV Cache、MTP Draft Cache、Causal Conv State 和 SSM Recurrent State。

当 Prefill TP 与 Decode TP 不一致时，两侧的 Cache ID、Tensor Shape、分片归属以及 Decode 最终布局均不同。

因此，只向 Decode 传递首 token 并不足以完成请求接管。Decode 要么得到不完整的状态，要么只能重新执行 Prefill。

本 PR 增加按 rank 传递的 Cache 元数据以及异构分片合并逻辑，使 Decode 能够直接使用 Prefill 产生的状态。

实现方案

拓扑检查与路由

• 扩展 PD 拓扑检查，支持 Qwen3.5 异构 TP。
• 保留原有同构场景和不支持拓扑的检查逻辑。
• 根据源端 DP/TP rank 路由 PUSH 元数据，不再假设所有源端共享同一种 Cache 布局。
• 补充拓扑检查和 PUSH 路由测试覆盖。

按 rank 传递 Cache 元数据

• 扩展 Engine、Remote Worker、Worker Service 和 Protobuf 接口，传递各 rank 的 KV Cache 和 Linear State Cache 信息。
• 将这些信息贯通普通推理、推测解码、MTP 和 DFlash Worker 路径。
• 公共 Engine 接口不直接依赖 NPU DataDist 的实现细节。

Decode Cache 分配与 Staging Cache

• 在 Prefill 完成前，提前创建 Decode 请求并分配最终 KV/Linear-State Slot。
• 为源端分片的 Target KV 和 Recurrent State 注册临时 Staging Cache。
• 将 Decode 目标 Cache ID 返回给每个 Prefill rank。
• 显式区分 Target、Draft、CONV 和 SSM Cache Role，不再按统一 Tensor 布局处理。

Cache 传输与合并

• Prefill 执行过程中，逐层将 Target KEY/VALUE 推送到 Decode Staging Cache。
• Decode 从两个 Prefill 分片恢复 Qwen3.5 CONV/SSM State，并合并为 TP1 布局。
• CONV 从各 rank 的 [Q, K, V] 重排为 Decode 侧的 [Q0, Q1, K0, K1, V0, V1]。
• SSM State 按 Tensor Parallel 分片维度拼接。
• MTP Draft KV Cache 与 Target Model Cache 分开传输。
• Decode 在首个 Decode Step 入队前完成状态接管，不重新执行 Prefill。

并行分片拉取与并发安全

• Decode 默认并行拉取两个 Prefill 源分片。
• 使用请求级互斥保护共享 Staging Tensor，避免并发 FirstGeneration RPC 相互覆盖。
• 即使请求线程抛出异常，也会等待 Worker Task 结束，避免异步任务访问已失效的栈变量。
• 可设置 XLLM_PD_PARALLEL_SHARD_PULL=0，无需重新编译即可回退到串行拉取。
• 增加 Allocation、Staging PUSH、Response RPC、Shard PULL/Merge 和 Decode FirstGeneration 的请求级耗时日志。

代码修改要点与分布

本 PR 共修改 33 个文件，代码量为 +1538/-93。改动集中在 KV 传输层，同时补齐元数据在 Distributed Runtime、Worker、Scheduler 和 Protobuf 之间的传递链路。

模块	文件数	代码量	主要职责
xllm/core/framework/kv_cache_transfer/	8	+1025/-19	异构 Cache 注册、Staging PUSH、分片 PULL/Merge、CONV/SSM 重排、拓扑检查和 PUSH 路由
xllm/core/distributed_runtime/	10	+223/-38	在 Engine、Remote Worker 和 Worker Service 之间传递各 rank 的 Cache 元数据，并管理 DataDist Cluster 连接
xllm/core/runtime/	9	+124/-2	扩展 Worker/WorkerClient 接口，并贯通普通、MTP、DFlash 和 Speculative Worker 路径
xllm/core/scheduler/	2	+111/-29	Decode 提前分配、首 token 与 Cache 元数据发送、FirstGeneration 状态接管和性能打点
xllm/models/ 与 xllm/proto/	2	+10/-0	暴露 Qwen3.5 MTP Cache 信息，并扩展 Worker RPC 消息
tests/	2	+45/-5	覆盖异构 PD 拓扑检查和按源 rank 的 PUSH 路由

关键代码入口：

• LlmDataDistTransfer::pull_and_merge_sharded_caches()：并行拉取 Prefill 分片，并按 Cache Role 合并到 Decode 最终布局。
• LlmDataDistTransfer::push_layer_registered_caches_to_staging()：等待 Prefill 各层完成后，将 Target KEY/VALUE 逐层推送到 Decode Staging Cache。
• SpecKVCacheTransfer::register_kv_cache_spec() 与 pull_kv_blocks()：注册并恢复 Target/Draft/Linear-State Cache。
• LLMEngine::pull_kv_blocks()、link_cluster() 和 unlink_cluster()：将多源 Cluster 连接和分片 Cache ID 传递到具体传输实现。
• DisaggPDScheduler::dispatch_requests()：在 Prefill 开始前创建 Decode 请求，并返回目标 Cache 元数据。
• DisaggPDScheduler::prefill_send_first_generation()：发送首 token、MTP Bootstrap 信息和源端 Cache ID。
• DisaggPDScheduler::decode_recv_first_generation()：恢复 KV/Linear State，完成首个 Decode Step 前的状态接管。

按单文件代码量看，核心实现主要位于：

• llm_data_dist_transfer.cpp：+637/-2。
• spec_kv_cache_transfer.cpp：+261/-10。
• disagg_pd_scheduler.cpp：+92/-29。
• llm_engine.cpp：+91/-30。

使能方式

部署需要先启动 etcd 和 xllm-service，随后分别启动 Prefill 和 Decode xLLM 实例。

Prefill 和 Decode 共同使用以下参数：

--enable_disagg_pd=true
--kv_cache_transfer_mode=PUSH
--num_speculative_tokens=3
--enable_mtp_draft_body_tp1=true
--enable_graph=true
--enable_graph_double_buffer=false
--enable_graph_mode_decode_no_padding=true
--enable_prefix_cache=false
--enable_chunked_prefill=true
--enable_schedule_overlap=true

角色参数如下：

Prefill: --instance_role=PREFILL，TP2 / 两个 rank
Decode:  --instance_role=DECODE， TP1 / 一个 rank

两侧必须使用相同的模型、Draft Model、Block Size、MTP 深度和 etcd 地址，并配置兼容的 DataDist 传输端口。

当前异构 Staging 布局支持 MTP3，本 PR 暂不使能 MTP5。

性能现状

测试合同：

• 模型：Qwen3.5-2B + Qwen3.5-2B-mtp。
• 固定输出 25 token，ignore_eos=true。
• 并发数：1。
• Warmup：8 个请求。
• PD 正式测试：每轮 200 个请求，共 3 轮。
• Prefill TP2 / Decode TP1 / MTP3。

并行拉取相对串行 PD Control 的收益

指标	串行 Control	并行拉取	收益
Decode KV/Linear-State 接管	17.508 ms	14.343 ms	3.165 ms / 18.1%
Linear-State 拉取与合并	13.295 ms	10.166 ms	3.130 ms
两个源分片拉取等待	7.661 ms	4.753 ms	2.908 ms
稳定 Round 2/3 E2E	191.905 ms	189.330 ms	2.575 ms
稳定 Round 2/3 TTFT	87.965 ms	87.450 ms	0.515 ms
稳定 Round 2/3 TPOT	4.330 ms/token	4.245 ms/token	0.085 ms/token

当前性能基线对比

场景	MTP	E2E mean	TTFT mean	TPOT mean	测试口径
混部 TP1	5	202.819 ms	116.805 ms	3.584 ms/token	20 请求 × 3 轮
混部 TP2	5	158.821 ms	85.536 ms	3.053 ms/token	200 × 3 复测中无外部负载的 Round 1
异构 PD P2/D1	3	189.330 ms	87.450 ms	4.245 ms/token	200 × 3 的稳定 Round 2/3

当前异构 PD 的 E2E 比干净的混部 TP2 基线慢 30.509 ms。TTFT 仅慢 1.914 ms，剩余差距主要来自 Decode 状态接管以及 TP1 Decode 的稳态生成开销。

由于混部使用 MTP5，而异构 PD 当前使用 MTP3，上表属于工程性能基线，不是相同 MTP 深度下的严格 A/B。

功能验证结论

Warning

精度尚未验证。 当前仅完成异构 KV 传输路径、输出非空和并发安全 smoke 验证。由于测试时没有空闲 NPU，pd_launch/test_accuracy.py 尚未执行，因此本 PR 当前没有可发布的精度分数，也不能声明模型精度通过。

• 正式测试 600/600 请求成功。
• 600 个响应均非空，平均输出长度严格为 25 token。
• 相对加入互斥前的候选版本，600/600 个 Prompt 一致，597 个输出逐字节一致。
• 3 个差异可在相同配置内部复现，属于配置自身的非确定性，不是 Cache 传输错误。
• Concurrency=2 补充测试 60/60 请求成功，没有 Transfer 或 Staging 错误。
• 三轮相对 Concurrency=1 参考输出的匹配数为 20/20、19/20、20/20；单个差异同样符合上述非确定性。
• 公开 commit 的 clang-format pre-commit 检查通过。

运行与性能证据采集自最终 PD 实现，随后公开分支被整理到最新 upstream main。本 PR 不修改任何第三方 submodule 指针；根据项目要求，本次发布未追加三方件编译验证。

已知限制与后续工作

• 异构 PD 当前支持 MTP3。使能 MTP5 前，需要解除 Staging Row Capacity 限制。
• Decode 仍在 FirstGeneration 接管阶段同步恢复并合并 CONV/SSM State，当前约为 10.2 ms/request，是最大的 PD 专属优化方向。
• Target KEY/VALUE 已在 Prefill 阶段推送，但 Decode 仍需执行一次本地 Staging Merge。
• 请求级 Staging Mutex 优先保证正确性，同时会串行化并发 Restore/Merge。后续可使用请求独占 Staging Row 或 Generation ID 消除该限制。
• 当前完成协议可以支持已经实现的 Target-KV Staging 路径。Recurrent State 预推送必须增加按请求、按层的完成协议，才能安全替换同步 PULL。

提交范围

公开分支只有一个 commit，并直接基于最新 upstream main。

机器专用的 Weight-Loading commit c4e88dc6f6a52f1ddfd4a6bd09cfe6725b41c3dd、本地启动脚本、Core Dump 和 Dirty Submodule 均未包含在 PR 中。

---

English Description

Summary

This PR adds heterogeneous tensor-parallel PD disaggregation for Qwen3.5 speculative decoding on NPU. It supports a Prefill instance running with TP2 and a Decode instance running with TP1, transfers the target KV cache, MTP draft KV cache, causal-convolution state, and SSM state, and lets Decode continue from the first generated token without replaying Prefill.

The implementation also makes the two Prefill-shard pulls concurrent on Decode. The parallel path is enabled by default and has a runtime fallback switch.

Motivation

The existing PD path assumes compatible cache layouts between Prefill and Decode. Qwen3.5 contains attention KV cache, MTP draft cache, causal-convolution state, and SSM recurrent state. When Prefill TP and Decode TP differ, their cache IDs, tensor shapes, shard ownership, and final Decode layouts are different.

Passing only the first token is therefore insufficient: Decode either receives incomplete state or has to execute Prefill again. This PR introduces explicit per-rank cache metadata and heterogeneous merge logic so that Decode can directly take over the request.

Implementation

Topology validation and routing

• Extend the PD topology guard for Qwen3.5 heterogeneous TP.
• Preserve existing homogeneous and unsupported-topology checks.
• Route PUSH metadata using the source DP/TP rank instead of assuming one source cache layout.
• Add topology-guard and push-route test coverage.

Per-rank cache metadata propagation

• Extend the engine, remote-worker, worker-service, and protobuf interfaces to exchange per-rank KV and linear-state cache information.
• Propagate cache information through normal, speculative, MTP, and DFlash worker paths.
• Keep the public engine path independent of NPU DataDist implementation details.

Decode cache allocation and staging

• Allocate the Decode request and its final KV/linear-state slots before Prefill execution completes.
• Register temporary staging caches for source-sharded target KV and recurrent state.
• Return the destination cache IDs to every Prefill rank.
• Keep target, draft, CONV, and SSM cache roles explicit instead of treating them as one uniform tensor layout.

Cache transfer and merge

• Prefill pushes target KEY/VALUE data layer by layer into Decode staging while Prefill is executing.
• Decode restores Qwen3.5 causal-convolution and SSM state from both Prefill shards and merges them into the TP1 layout.
• CONV data is reordered from per-rank [Q, K, V] shards to [Q0, Q1, K0, K1, V0, V1].
• SSM state is concatenated along its tensor-parallel shard dimension.
• MTP draft KV cache is transferred separately from the target-model cache.
• Decode consumes the transferred state before enqueueing the first Decode step; it does not rerun Prefill.

Parallel shard pull and safety

• Pull the two Prefill source shards concurrently on Decode by default.
• Protect shared staging tensors with a request-level mutex so concurrent first-generation RPCs cannot overwrite each other.
• Wait for the worker task even when the request thread throws, avoiding use-after-scope access.
• Provide XLLM_PD_PARALLEL_SHARD_PULL=0 as an immediate serial fallback without rebuilding.
• Add request-level timing logs for allocation, staging PUSH, response RPC, shard PULL/merge, and Decode first-generation handling.

Code change highlights and distribution

This PR changes 33 files with +1538/-93. Most of the implementation is in the KV-transfer layer, with supporting metadata plumbing across the distributed runtime, workers, scheduler, and protobuf interfaces.

Module	Files	Diff size	Responsibility
xllm/core/framework/kv_cache_transfer/	8	+1025/-19	Heterogeneous cache registration, staged PUSH, sharded PULL/merge, CONV/SSM layout conversion, topology validation, and PUSH routing
xllm/core/distributed_runtime/	10	+223/-38	Propagate per-rank cache metadata through Engine, Remote Worker, and Worker Service, and manage DataDist cluster links
xllm/core/runtime/	9	+124/-2	Extend Worker/WorkerClient interfaces across normal, MTP, DFlash, and speculative worker paths
xllm/core/scheduler/	2	+111/-29	Decode pre-allocation, first-token/cache-metadata delivery, FirstGeneration takeover, and timing instrumentation
xllm/models/ and xllm/proto/	2	+10/-0	Expose Qwen3.5 MTP cache metadata and extend worker RPC messages
tests/	2	+45/-5	Cover heterogeneous PD topology validation and source-rank-aware PUSH routing

Key code entry points:

• LlmDataDistTransfer::pull_and_merge_sharded_caches() pulls Prefill shards concurrently and merges each cache role into the final Decode layout.
• LlmDataDistTransfer::push_layer_registered_caches_to_staging() waits for each Prefill layer and pushes target KEY/VALUE into Decode staging.
• SpecKVCacheTransfer::register_kv_cache_spec() and pull_kv_blocks() register and restore target, draft, and linear-state caches.
• LLMEngine::pull_kv_blocks(), link_cluster(), and unlink_cluster() propagate multi-source cluster links and sharded cache IDs to the transfer implementation.
• DisaggPDScheduler::dispatch_requests() creates the Decode request before Prefill and returns destination cache metadata.
• DisaggPDScheduler::prefill_send_first_generation() sends the first token, MTP bootstrap data, and source cache IDs.
• DisaggPDScheduler::decode_recv_first_generation() restores KV/linear state and completes takeover before the first Decode step.

The largest implementation files are:

• llm_data_dist_transfer.cpp: +637/-2.
• spec_kv_cache_transfer.cpp: +261/-10.
• disagg_pd_scheduler.cpp: +92/-29.
• llm_engine.cpp: +91/-30.

Enablement

The deployment requires etcd and xllm-service, followed by separate Prefill and Decode xLLM instances.

Both xLLM roles use:

--enable_disagg_pd=true
--kv_cache_transfer_mode=PUSH
--num_speculative_tokens=3
--enable_mtp_draft_body_tp1=true
--enable_graph=true
--enable_graph_double_buffer=false
--enable_graph_mode_decode_no_padding=true
--enable_prefix_cache=false
--enable_chunked_prefill=true
--enable_schedule_overlap=true

Role-specific settings:

Prefill: --instance_role=PREFILL, TP2 / two ranks
Decode:  --instance_role=DECODE,  TP1 / one rank

Both roles must use the same model, draft model, block size, speculative depth, etcd endpoint, and compatible DataDist transfer ports. The current heterogeneous staging layout supports MTP depth 3. MTP5 is not enabled by this PR.

Performance status

Test contract:

• Model: Qwen3.5-2B with Qwen3.5-2B-mtp
• Output: fixed 25 tokens, ignore_eos=true
• Concurrency: 1
• Warmup: 8 requests
• Formal PD run: 200 requests x 3 rounds
• Prefill TP2 / Decode TP1 / MTP3

Parallel pull improvement over serial PD control

Metric	Serial control	Parallel pull	Improvement
Decode KV/linear-state restore	17.508 ms	14.343 ms	3.165 ms / 18.1%
Linear-state pull and merge	13.295 ms	10.166 ms	3.130 ms
Two-source-shard pull wait	7.661 ms	4.753 ms	2.908 ms
Stable Round 2/3 E2E	191.905 ms	189.330 ms	2.575 ms
Stable Round 2/3 TTFT	87.965 ms	87.450 ms	0.515 ms
Stable Round 2/3 TPOT	4.330 ms/token	4.245 ms/token	0.085 ms/token

Current baseline comparison

Scenario	MTP	E2E mean	TTFT mean	TPOT mean	Contract note
Mixed TP1	5	202.819 ms	116.805 ms	3.584 ms/token	20 requests x 3 rounds
Mixed TP2	5	158.821 ms	85.536 ms	3.053 ms/token	clean Round 1 of a 200 x 3 rerun
Heterogeneous PD P2/D1	3	189.330 ms	87.450 ms	4.245 ms/token	stable Round 2/3 of 200 x 3

The heterogeneous PD path is currently 30.509 ms slower in E2E than the clean mixed-TP2 reference. TTFT differs by only 1.914 ms; most of the remaining gap is Decode handoff plus TP1 Decode steady-state cost. These rows are engineering baselines rather than a strict same-depth A/B because mixed mode uses MTP5 while heterogeneous PD currently uses MTP3.

Functional validation

Warning

Model accuracy has not been validated yet. The current evidence only covers heterogeneous KV-transfer routing, non-empty outputs, and concurrency-safety smoke tests. pd_launch/test_accuracy.py has not been executed because no NPU was available, so this PR has no publishable accuracy score and does not claim an accuracy PASS.

• 600/600 formal requests completed successfully.
• All 600 responses were non-empty and averaged exactly 25 output tokens.
• 600/600 prompts matched the pre-mutex candidate; 597 outputs matched byte-for-byte. The three differences were reproduced as within-configuration nondeterminism rather than transfer corruption.
• A concurrency-2 safety run completed 60/60 requests with no transfer or staging errors.
• Per-round output matches against the concurrency-1 references were 20/20, 19/20, and 20/20; the one difference was consistent with the same nondeterminism above.
• clang-format pre-commit validation passed for the public commit.

The runtime and performance evidence was collected with the final PD implementation before the public branch was rebased onto the latest upstream main. This PR does not modify third-party submodule pointers; per project instruction, no additional third-party build validation was required for publication.

Known limitations and follow-up work

• Heterogeneous PD currently supports MTP3; the staging-row capacity must be generalized before MTP5 can be enabled.
• Decode still synchronously restores and merges recurrent CONV/SSM state during first-generation handoff. This path currently costs about 10.2 ms/request and is the largest remaining PD-specific optimization target.
• Target KEY/VALUE is pushed during Prefill but still requires a local staging merge on Decode.
• The request-level staging mutex favors correctness and serializes simultaneous restore/merge sections. Future work can use request-owned staging rows or generation IDs to remove this serialization.
• The current completion protocol is sufficient for the implemented target-KV staging path. Recurrent-state pre-push requires an explicit per-request/per-layer completion protocol before it can safely replace synchronous PULL.

Commit hygiene

The public branch contains one commit based directly on the latest upstream main. The local-only weight-loading commit c4e88dc6f6a52f1ddfd4a6bd09cfe6725b41c3dd, local launch scripts, core dumps, and dirty submodule worktrees are not included.