feat: support heterogeneous qwen3.5 PD disaggregation.

[Vectorwh]

中文说明

概述

本 PR 为昇腾 NPU 上的 Qwen3.5 推测解码增加异构 Tensor Parallel 的 PD 分离能力，当前完整支持：

Prefill DP1 × TP2 → Decode DP1 × TP1

当 Prefill 与 Decode 的 TP 布局不一致时，Decode 无法直接复用原有一对一缓存传输。本 PR 引入多对一 TP 路由、异构 staging cache、按缓存角色合并以及完整的 Worker RPC 链路，使 Decode 能接管 Prefill 产生的 Target KV、CONV State、SSM State 和 MTP Draft KV，并从首 token 继续生成，无需重跑 Prefill。

Important

核心性能收益（16K+ 长输入，单轮 Discovery）

• 异构相对同构 PD： P2/D1 使用 3 张 NPU，相比 P2/D2 的 4 张卡少用 25% 资源，但仍保留 98.3%～98.7% 的系统吞吐；Req/s/NPU 提升 31.1%～31.6%，平均 E2E 仅高 1.3%～1.8%。
• 异构相对混部 TP2： 在并发 8～32、输出 32/64 token 下，P2/D1 的系统 Req/s 提升 12.5%～39.5%，平均 E2E 降低 12.2%～39.8%，TPOT 稳定在 4.62～5.95 ms/token。
• 该方案的核心价值是：用 TP1 Decode 降低 PD 的 Decode 资源成本，同时在长输入、高并发下隔离 Prefill 对 Decode 的干扰。P2/D1 为 3 卡、混部 TP2 为 2 卡，因此后一组对比只代表系统性能收益，不代表单卡效率更高。

整体框架

方案采用控制面与数据面分离，并对不同缓存使用混合 PUSH/PULL：

控制面：
Prefill Scheduler ── FirstGeneration ──→ Decode Scheduler
Decode Engine ── PullKVCache RPC ──→ Decode Worker

数据面：
Target KV：P Worker ── PUSH ──→ D staging ── merge ──→ D Target Cache
CONV/SSM：D Worker ── PULL ──→ D staging ── merge ──→ D Linear State
Draft KV：D Worker ── PULL ──→ D staging ── merge ──→ D Draft Cache

Target KV 数据量大，Prefill 每完成一层便将该层 KEY/VALUE shard PUSH 到 Decode staging，使传输与后续计算重叠。CONV/SSM 是 Prefill 结束后的最终递推状态，Draft KV 数据量较小；两者由 Decode 同步 PULL，使 shard 到达、tensor merge、设备流同步和请求入队形成明确的完成顺序。

请求时序

Decode 提前创建请求并分配目标 cache
  ↓
Prefill 执行 prompt，并按层 PUSH Target KV
  ↓
Prefill 发送 FirstGeneration、源拓扑和 cache IDs
  ↓
Decode 识别 src_tp_size != dst_tp_size
  ↓
LLMEngine 为每个 D rank 选择对应 P ranks
  ↓
恢复并合并 CONV/SSM
  ↓
合并已预推送的 Target KV
  ↓
恢复并合并 Draft KV
  ↓
同步 NPU 默认流
  ↓
全部 D Workers 成功后请求进入 Decode queue

多对一 TP 路由

一个 Decode TP rank 对应的 Prefill TP ranks 满足：

src_tp_rank = dst_tp_rank + n × dst_tp_size

等价于 dst_tp_rank = src_tp_rank % dst_tp_size。TP2→TP1 时，D-rank0 同时接收 P-rank0 和 P-rank1 的 shard。link_cluster()、unlink_cluster() 和实际恢复共用同一 owner 映射。

缓存传输与合并

缓存角色	传输方式	合并方式
Target KEY/VALUE	Prefill 分层 PUSH，Decode 本地 merge	dim 2
CONV State	Decode 同步 PULL 两个 shard	最后一维，按 Q/K/V 分段重排
SSM State	Decode 同步 PULL 两个 shard	dim 1，支持 checkpoint ID 展开
MTP Draft KV	Decode 同步 PULL	dim 2；Draft body TP1 时只拉一个完整副本

CONV 每个 Prefill rank 的布局是 [Q,K,V]。实现将 [Q0,K0,V0] + [Q1,K1,V1] 重排为 Decode TP1 所需的 [Q0,Q1,K0,K1,V0,V1]。

正确性与并发安全

• 两个 source shards 使用互不重叠的 staging 行区间；staging 按 2 MiB 对齐。
• 同一请求的两个 shard 默认并行 PULL。
• hetero_pull_mutex_ 防止并发 FirstGeneration 覆盖共享 staging。
• 请求线程异常时仍等待 shard worker task，避免 use-after-scope。
• Worker 返回成功前执行 synchronize_default_stream()。
• 任一传输、merge 或设备同步失败时释放 Decode cache，且请求不入队。
• enable_pd_parallel_shard_pull 默认开启，并行拉取两个源 TP shard；需要排查问题或回退到串行路径时，可使用 --enable_pd_parallel_shard_pull=false。

主要代码改动

• 放开非 MLA、PUSH、等 DP、Prefill TP 为 Decode TP 整数倍的异构拓扑。
• 新增 get_src_tp_ranks()，修改 Engine 建链、拆链和异构恢复路由。
• FirstGeneration 始终携带 source cluster/address、block ID 和 linear-state ID。
• 新增从 Engine、WorkerClient、RemoteWorker、CommChannel、WorkerService 到 KVCacheTransfer 的 pull_hetero_kv_blocks 调用链。
• register_hetero_staging_caches() 创建并注册 Target/Draft staging。
• push_layer_registered_caches_to_staging() 按层 PUSH Target KV。
• pull_and_merge_sharded_caches() 并行 PULL 并合并 CONV/SSM 或 Draft KV。
• merge_pre_pushed_sharded_caches() 合并已 PUSH 的 Target KV。
• SpecKVCacheTransfer::pull_hetero_kv_blocks() 编排 CONV/SSM、Target KV 和 Draft KV 的恢复顺序。
• 增加 allocation、FirstGeneration、staging PUSH、shard PULL/merge、enqueue 和 response 的 [PD-PERF] 日志。

本 PR 修改 33 个文件，代码量为 +1538/-93。

使能方式

共同参数：

--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
--enable_pd_parallel_shard_pull=true

角色参数：

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

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

16K+ 长输入性能验证（Discovery）

为验证 PD 分离在长 Prompt、高并发场景下的收益，完成了 5 种拓扑的单轮 Discovery 对比矩阵。

测试合同：

• 模型：Qwen3.5-2B + Qwen3.5-2B-mtp；MTP depth=3，Draft Body TP1。
• 数据：51 条顺序拼接请求；输入长度 16,422～20,465 token，均值 18,591.824 token。
• 输出：固定 32/64 token，ignore_eos=true；并发 8/16/32。
• Warmup：并发 8/16/32 分别执行 16/32/64 条同形状请求；每个拓扑和 shape 均独立重启服务。
• 配置：开启 graph、decode no-padding、chunked prefill 和 schedule overlap，关闭 prefix cache。
• 结果：30/30 个组合通过，1530/1530 条正式请求成功，实际输出长度全部严格等于 32 或 64 token。

异构 P2/D1 相对混部 TP2 的直接对比如下：

输出	并发	混部 TP2 Req/s	P2/D1 Req/s	吞吐提升	混部 TP2 E2E	P2/D1 E2E	E2E 降低
32	8	2.824	3.177	+12.5%	2733.7 ms	2399.2 ms	12.2%
32	16	2.321	3.175	+36.8%	6165.5 ms	4407.8 ms	28.5%
32	32	2.790	3.178	+13.9%	10026.2 ms	7259.9 ms	27.6%
64	8	2.503	3.151	+25.9%	3085.0 ms	2412.6 ms	21.8%
64	16	2.421	3.149	+30.1%	5869.8 ms	4449.1 ms	24.2%
64	32	2.258	3.150	+39.5%	12196.3 ms	7337.8 ms	39.8%

主要结论：

• 相对混部 TP2，P2/D1 在全部 6 个组合中均提高系统请求吞吐并降低平均 E2E；Decode TPOT 为 4.62～5.95 ms/token，而混部 TP2 为 19.63～185.24 ms/token。长 Prefill 与 Decode 隔离后，Decode 不再被其他请求的 Prefill 持续干扰。
• P2/D1 使用 3 张 NPU，而混部 TP2 使用 2 张，因此上述结果是系统吞吐和时延收益，不代表相对混部 TP2 的单卡效率更高。
• 与同构 P2/D2（4 卡）相比，异构 P2/D1（3 卡）只损失 1.3%～1.7% 系统 Req/s，平均 E2E 仅高 1.3%～1.8%，但少使用 25% 的 NPU；Req/s/NPU 从约 0.80 提升到约 1.05，即提高 31.1%～31.6%。这表明 Decode TP1 能以更少资源保留约 98.3%～98.7% 的 P2/D2 系统吞吐。
• PD 的平均 TTFT 仍高于混部 TP2，因为所有长 Prompt 集中在 Prefill TP2 排队；本轮收益主要体现在稳定 Decode TPOT、系统吞吐和 E2E，而不是 TTFT。
• P2×2/D2 的扩展实验显示，多 Prefill 实例可继续缓解 Prefill 排队；相对单 Prefill P2/D2，系统 Req/s 提升 81.6%～93.6%，E2E 降低 45.5%～48.5%。

完整 5 拓扑矩阵（点击展开）

输出	并发	拓扑	卡数	Req/s	Req/s/NPU	总 Token/s	Token/s/NPU	TTFT ms	TPOT ms	E2E ms
32	8	混部 TP2	2	2.824	1.412	52591.5	26295.8	1621.0	35.89	2733.7
32	8	混部 TP4	4	5.047	1.262	93995.0	23498.7	445.3	36.21	1567.9
32	8	P2/D1	3	3.177	1.059	59172.8	19724.3	2214.9	5.95	2399.2
32	8	P2/D2	4	3.221	0.805	59986.0	14996.5	2253.5	3.58	2364.4
32	8	P2×2/D2	6	5.851	0.975	108963.6	18160.6	1083.0	6.63	1288.4
32	16	混部 TP2	2	2.321	1.160	43217.4	21608.7	2407.5	121.22	6165.5
32	16	混部 TP4	4	3.643	0.911	67854.8	16963.7	858.9	79.43	3321.2
32	16	P2/D1	3	3.175	1.058	59134.5	19711.5	4226.3	5.86	4407.8
32	16	P2/D2	4	3.221	0.805	59991.7	14997.9	4234.0	3.57	4344.7
32	16	P2×2/D2	6	6.115	1.019	113889.8	18981.6	2123.8	4.18	2253.5
32	32	混部 TP2	2	2.790	1.395	51961.6	25980.8	4283.6	185.24	10026.2
32	32	混部 TP4	4	3.922	0.980	73039.8	18259.9	2067.1	136.51	6298.9
32	32	P2/D1	3	3.178	1.059	59192.3	19730.8	7083.9	5.68	7259.9
32	32	P2/D2	4	3.225	0.806	60054.0	15013.5	7024.4	3.87	7144.5
32	32	P2×2/D2	6	6.241	1.040	116235.4	19372.6	3620.4	4.44	3757.9
64	8	混部 TP2	2	2.503	1.251	46692.4	23346.2	1848.0	19.63	3085.0
64	8	混部 TP4	4	4.640	1.160	86571.4	21642.9	380.3	20.68	1683.1
64	8	P2/D1	3	3.151	1.050	58793.7	19597.9	2116.7	4.70	2412.6
64	8	P2/D2	4	3.205	0.801	59789.8	14947.5	2154.1	3.44	2371.0
64	8	P2×2/D2	6	6.153	1.026	114797.5	19132.9	975.5	3.91	1222.0
64	16	混部 TP2	2	2.421	1.210	45159.5	22579.8	2604.9	51.82	5869.8
64	16	混部 TP4	4	4.104	1.026	76564.9	19141.2	801.5	43.71	3555.4
64	16	P2/D1	3	3.149	1.050	58747.3	19582.4	4151.5	4.72	4449.1
64	16	P2/D2	4	3.194	0.798	59577.7	14894.4	4168.7	3.52	4390.2
64	16	P2×2/D2	6	6.090	1.015	113614.6	18935.8	2051.0	3.84	2293.1
64	32	混部 TP2	2	2.258	1.129	42118.8	21059.4	4826.9	116.98	12196.3
64	32	混部 TP4	4	2.806	0.702	52354.9	13088.7	2060.9	85.00	7415.9
64	32	P2/D1	3	3.150	1.050	58760.7	19586.9	7046.6	4.62	7337.8
64	32	P2/D2	4	3.191	0.798	59532.3	14883.1	6998.6	3.69	7230.9
64	32	P2×2/D2	6	6.075	1.013	113337.3	18889.5	3585.2	4.07	3841.5

Note

以上为单轮 Discovery 趋势数据，不是 3 轮 Formal 发布数据；并发 32 仅作为容量趋势。跨拓扑比较必须同时考虑实际 NPU 数。

功能验证

Warning

精度尚未验证。 当前证据覆盖异构缓存传输、非空输出和并发安全 smoke test；pd_launch/test_accuracy.py 尚未执行，因此不能声明模型精度通过。

• 正式测试 600/600 请求成功；所有响应非空，平均输出长度为 25 token。
• 相对加入互斥前的候选版本，597/600 输出逐字节一致；3 个差异可在相同配置内复现，属于配置自身非确定性。
• Concurrency=2 补充测试 60/60 成功，无 Transfer 或 Staging 错误。
• clang-format pre-commit 检查通过。

已知限制与后续工作

• 当前完整支持 Prefill TP2→Decode TP1；拓扑层更通用，但 staging/merge 仍要求每个 Decode rank 对应两个 source shards。
• 当前 staging 支持 MTP3；MTP5 需要扩展 staging row capacity。
• CONV/SSM 同步恢复约 10.2 ms/request，是最大 PD 专属优化方向。
• 请求级 staging mutex 会串行化并发 restore/merge。
• Recurrent state 改为预 PUSH 前，需要 Decode 可等待的按请求、按 shard/layer 完成协议。
• 仍需补充真实 NPU 数值一致性和正式 accuracy 测试。

---

English Description

Summary

This PR adds heterogeneous tensor-parallel PD disaggregation for Qwen3.5 speculative decoding on Ascend NPU. The complete target is Prefill DP1 × TP2 → Decode DP1 × TP1.

It reconstructs target KV, causal-convolution state, SSM recurrent state, and MTP draft KV so Decode can continue from Prefill's first generated token without replaying Prefill.

Important

Key performance benefit (16K+ input, single-round Discovery)

• Heterogeneous vs. homogeneous PD: P2/D1 uses three NPUs instead of four for P2/D2, reducing NPU usage by 25% while retaining 98.3%–98.7% of system throughput. Req/s/NPU improves by 31.1%–31.6%, with only 1.3%–1.8% higher mean E2E.
• Heterogeneous PD vs. mixed TP2: across concurrency 8–32 and 32/64-token outputs, P2/D1 improves system Req/s by 12.5%–39.5%, reduces mean E2E by 12.2%–39.8%, and keeps TPOT at 4.62–5.95 ms/token.
• The key value is lowering Decode resource cost with TP1 while isolating Decode from long-Prefill interference. P2/D1 uses three NPUs versus two for mixed TP2, so the latter comparison is a system-level gain, not a per-NPU efficiency claim.

Architecture

The implementation separates the control plane from the data plane and uses hybrid PUSH/PULL:

Control:
Prefill Scheduler ── FirstGeneration ──→ Decode Scheduler
Decode Engine ── PullKVCache RPC ──→ Decode Worker

Data:
Target KV: P Worker ── PUSH ──→ D staging ── merge ──→ Target Cache
CONV/SSM:  D Worker ── PULL ──→ D staging ── merge ──→ Linear State
Draft KV:  D Worker ── PULL ──→ D staging ── merge ──→ Draft Cache

Target KV is pushed layer by layer to overlap transfer with Prefill computation. Decode synchronously pulls CONV/SSM and Draft KV so arrival, merge, device synchronization, and enqueueing have an explicit completion order.

A Decode rank owns source ranks satisfying src_tp_rank = dst_tp_rank + n × dst_tp_size. For TP2→TP1, D-rank0 receives both Prefill shards.

Cache	Transfer	Reconstruction
Target KEY/VALUE	Layer-wise PUSH	Concatenate on dim 2
CONV	Decode PULL	Reorder Q/K/V segments on the last dimension
SSM	Decode PULL	Concatenate on dim 1, with checkpoint expansion
MTP Draft KV	Decode PULL	Concatenate on dim 2, or pull one TP1 replica

Safety

• Source shards occupy disjoint, 2 MiB-aligned staging ranges.
• Two source shards are pulled concurrently by default.
• A request-level mutex protects shared staging.
• Worker tasks are joined on exceptions.
• The NPU default stream is synchronized before success is returned.
• Decode enqueues only after every worker succeeds.
• enable_pd_parallel_shard_pull is enabled by default and pulls the two source TP shards concurrently. Use --enable_pd_parallel_shard_pull=false to fall back to serial PULL for debugging or compatibility.

Configuration

--enable_pd_parallel_shard_pull=true

This option is enabled by default. Set --enable_pd_parallel_shard_pull=false on Decode to switch heterogeneous source-shard recovery back to serial PULL.

Main changes

• Extend topology validation and add source-owner-aware TP routing.
• Propagate multi-source metadata through Engine, Worker RPC, and protobuf layers.
• Register target/draft heterogeneous staging caches.
• PUSH target KV layer by layer during Prefill.
• PULL and merge CONV/SSM and Draft KV on Decode.
• Merge pre-pushed target KV into final Decode blocks.
• Add request-level timing for allocation, FirstGeneration, transfer, merge, enqueue, and response delivery.

This PR changes 33 files with +1538/-93.

16K+ Long-Context Performance (Discovery)

A single-round Discovery matrix was added for 51 prompts of 16,422–20,465 input tokens (18,591.824 on average), fixed 32/64-token outputs with ignore_eos=true, and concurrency 8/16/32. MTP depth was 3 with Draft Body TP1. Each topology/shape was restarted independently and warmed up with 16/32/64 same-shape requests. All 30/30 combinations and 1530/1530 formal requests succeeded with exact output lengths.

Output	Concurrency	Mixed TP2 Req/s	P2/D1 Req/s	Throughput gain	Mixed TP2 E2E	P2/D1 E2E	E2E reduction
32	8	2.824	3.177	+12.5%	2733.7 ms	2399.2 ms	12.2%
32	16	2.321	3.175	+36.8%	6165.5 ms	4407.8 ms	28.5%
32	32	2.790	3.178	+13.9%	10026.2 ms	7259.9 ms	27.6%
64	8	2.503	3.151	+25.9%	3085.0 ms	2412.6 ms	21.8%
64	16	2.421	3.149	+30.1%	5869.8 ms	4449.1 ms	24.2%
64	32	2.258	3.150	+39.5%	12196.3 ms	7337.8 ms	39.8%

Across all six shapes, heterogeneous P2/D1 improves system Req/s and mean E2E over mixed TP2, with a stable 4.62–5.95 ms/token TPOT versus 19.63–185.24 ms/token for mixed TP2. This comparison is a system-level result: P2/D1 uses three NPUs while mixed TP2 uses two, so it does not claim higher per-NPU efficiency over mixed TP2.

The direct heterogeneous-value comparison is P2/D1 (3 NPUs) versus homogeneous P2/D2 (4 NPUs): P2/D1 retains 98.3%–98.7% of system Req/s with 25% fewer NPUs, while Req/s/NPU improves by 31.1%–31.6%. Mean E2E is only 1.3%–1.8% higher. P2×2/D2 additionally demonstrates scale-out headroom: versus single-Prefill P2/D2, system Req/s improves by 81.6%–93.6% and E2E falls by 45.5%–48.5%.

PD TTFT remains higher than mixed TP2 because long prompts queue on the centralized Prefill TP2 instance. The observed benefit is therefore stable Decode TPOT, system throughput, and E2E—not TTFT.

Note

These are single-round Discovery trends, not three-round Formal publishable results. Concurrency 32 is retained as a capacity trend, and cross-topology comparisons must account for NPU count. The full five-topology matrix is provided in the Chinese section above.

Validation and limitations

Warning

Model accuracy has not been validated. Current evidence covers transfer routing, non-empty outputs, and concurrency safety; formal accuracy remains pending.

• 600/600 formal requests succeeded; all responses were non-empty.
• 597/600 outputs matched byte-for-byte; the remaining differences reproduced as within-configuration nondeterminism.
• A concurrency-2 run completed 60/60 requests without transfer or staging errors.

Known limits:

• The complete path currently supports Prefill TP2→Decode TP1 and two source shards per Decode rank.
• The staging layout currently supports MTP3; MTP5 needs generalized capacity.
• Synchronous CONV/SSM restore remains the largest PD-specific latency component.
• The request-level staging mutex serializes simultaneous restore/merge sections.
• Recurrent-state pre-push requires a Decode-visible per-request/per-shard completion protocol.
• NPU numerical-consistency and formal accuracy tests remain follow-up work.

[Vectorwh]

RFC: #2060\n\nThis RFC documents the architecture, supported scope, cache resharding strategy, synchronization guarantees, and validation plan for this implementation.

[yingxudeng]

魔法数字 2 的作用是什么？

[Vectorwh]

已修改。2指的是当前只支持 Prefill TP2 到 Decode TP1。已经换成 kSupportedHeterogeneousSourceShardCount，并在 staging 注册和 pull/merge 校验里统一使用。

[yingxudeng]

仅支持 prefill TP 2 到 decode TP 1，只有这种情况嘛？其他 TP 组合是否支持。例如 qwen3.5 35b a3b 要用的话，是否可以其他 TP 组合方式

[yingxudeng]

魔法数字

[Vectorwh]

已修改

[yingxudeng]

魔法数字

[Vectorwh]

已修改

[yingxudeng]

为何 pull_and_merge_sharded_caches 内部持有锁 merge_pre_pushed_sharded_caches 无锁？

[Vectorwh]

已修改。原来锁放在 pull_and_merge_sharded_caches 里面，走 PUSH 路径时 merge_pre_pushed_sharded_caches 没有被同一把锁保护。我现在把锁提到了 pull_hetero_kv_blocks 外层，一次锁住 linear-state pull/merge、预推送 target KV merge 和 draft KV pull/merge，避免并发请求复用 staging rows 时互相覆盖。

[yingxudeng]

魔法数字

[Vectorwh]

已修改。原来的 256 是按 max_tokens_per_batch=32768、block_size=128 直接写死的，现在会根据 KVCacheConfig::block_size()、SchedulerConfig::max_tokens_per_batch() 和 max_seqs_per_batch() 算每个 shard 需要多少 staging 行，同时给每条序列可能存在的尾部未满 block 留出空间。

[yingxudeng]

魔法数字

[Vectorwh]

已修改。原来的 4 是当前 Qwen3.5 GDN linear-state 的 checkpoint stride。现在直接用同层 SSM cache 和 CONV cache 的行数比例算出 checkpoint_stride，再用它计算 staging 容量和展开 SSM id。

[yingxudeng]

pull_hetero_kv_blocks 全局持 hetero_pull_mutex_，是否会导致任意两个并发 Decode 请求都串行(即使它们的 request_id、layer、shard 完全不重叠)

[Vectorwh]

会的。目前同一个 Decode worker 上，异构 cache 的“恢复/接管”阶段还是串行的，但请求进入 Decode queue 后的生成不会被这把锁串行。原因是 staging tensor 是 worker 级共享的，不同请求会复用同一组 staging rows；如果同时 pull/merge，数据可能互相覆盖。

[yingxudeng]

原来只在pull做的事情，为何改成push也会做了

[Vectorwh]

这里没有把普通 PUSH 改成 PULL，同构 PUSH 的行为也没变。异构 TP 下走的是混合路径：Prefill 逐层把 target KEY/VALUE PUSH 到 Decode staging；Qwen3.5 的 CONV/SSM linear state 和 Draft TP1 KV 仍由 Decode 在接管时同步 PULL。所以即使 transfer mode 是 PUSH，Decode 还是要拿到源 cluster/addrs、block_ids 和 linear_state_ids，才能把剩下的状态恢复完整。

当前代码仍保留原来的 transfer mode，这些额外元数据只在异构恢复路径使用，注释里也说明了这个限制。如果不带这些信息，target KV 虽然到了，Decode 还是没法恢复 linear state 和 draft KV。

[yingxudeng]

冲突了，有空可以再rebase下

[yingxudeng]

仿照 disagg_pd_config 其他参数是不是更好？

[Vectorwh]

已修改。运行路径里不再直接读 XLLM_PD_PARALLEL_SHARD_PULL，改成正式的 disagg_pd_config 参数 enable_pd_parallel_shard_pull，默认开启。flag、JSON 配置解析和导出、支持参数列表、配置单测已补充。

[yingxudeng]

大量重复代码在 pull_and_merge_sharded_caches / merge_pre_pushed_sharded_caches / push_layer_registered_caches_to_staging 之间

[Vectorwh]

已重构。shape 判断、ID 生成、staging shard view、shard merge、回写和 2 MiB 对齐等公共逻辑都提成了 helper，三个入口里只保留各自的流程。

[yingxudeng]

为何要新增 pull_hetero_kv_blocks ？ 为何不能拓展 pull_kv_blocks

[Vectorwh]

现有 pull_kv_blocks 对应同构场景，语义是一对一 TP rank 映射和直接搬运；异构恢复需要接收多个源 TP shard，先写 staging，再按 tensor role 和分片维度重排合并，还要区分 Qwen3.5 linear state 和 Draft TP1 KV。

[yingxudeng]

llm_data_dist_transfer.cpp, transport 层只负责搬数据
但是现在 transport 现在要懂 qwen3.5 结构和细节，是否不太合理？

[Vectorwh]

已修改。tensor role、CONV Q/K/V 布局、SSM checkpoint stride、异构 staging 和 merge 这些模型/Speculative 相关逻辑，现在都放在 SpecKVCacheTransfer 里。LlmDataDistTransfer 只保留通用的 DataDist 注册、PUSH/PULL 搬运和 layer synchronization，不再理解 Qwen3.5 的 cache 结构。

[Vectorwh]

已rebase到最新

[yingxudeng]

有冲突