diff --git a/docs/src/content/docs/en/features/runtime_cp_dp_switch.md b/docs/src/content/docs/en/features/runtime_cp_dp_switch.md new file mode 100644 index 0000000000..844f02a5fa --- /dev/null +++ b/docs/src/content/docs/en/features/runtime_cp_dp_switch.md @@ -0,0 +1,121 @@ +--- +title: "Runtime CP↔DP Mode Switching" +sidebar: + order: 91 +--- +## Background + +xLLM supports two parallel layouts for a fixed world size: + +- **CP_PREFILL** (`cp=N, dp=1`): context parallelism across all workers. Long prompts split their KV compute across workers, minimizing TTFT for large-input requests. +- **DP_DECODE** (`cp=1, dp=N`): data parallelism, `N` independent decoder replicas. Higher concurrency and throughput for small requests, at the cost of no per-request work-splitting. + +Traffic patterns rarely stay in one regime forever. Long-context ingest bursts favor CP; concurrent short-request traffic favors DP. Static provisioning per instance forces the operator to overprovision one lane or accept degraded metrics on the other. + +**Runtime CP↔DP mode switching** lets an xLLM instance flip between the two layouts on demand, without a service restart, in ~500 ms. + +## When to use it + +- Traffic mix shifts across the day (long-context batch jobs at night, chat traffic during the day). +- You want to test-drive both layouts on a single instance rather than maintain separate deployments. +- You are building an autoscaling controller that decides layout per instance and needs a machine-actuated switch RPC. + +Do **not** use it if: + +- You are debugging a numerical accuracy problem — start with a fixed layout to eliminate flip as a variable. +- Your workload is uniformly one shape — the resident cost of holding both ATB graphs (~600 MB / NPU + extra graph workspace) is not free. + +## How it works + +### Startup: dual-graph mode + +When the worker starts with `--enable_runtime_cp_dp_switch=true`, it initializes both parallel layouts up front: + +- Two `CollectiveCommunicator`s. One with `cp_size=N,dp_size=1`, one with `cp_size=1,dp_size=N`. Each brings up its own HCCL comm domain on a distinct master port (offset by `--dual_mode_port_stride`, default 256). +- Two `ParallelArgs` snapshots wrapped in a `DualParallelArgs` with an `atomic` discriminator. `active()` reads the current snapshot with acquire semantics. +- Four ATB nodes per decoder layer: `cp_prefill`, `cp_decode`, `dp_prefill`, `dp_decode`. Weights are shared; only the graph binding differs. + +The initial mode is chosen from `--cp_size` at startup (`cp>1 → CP_PREFILL`, else `DP_DECODE`). + +### The switch: an RPC that atomically flips the layout + +Callers invoke `xllm.proto.ModeSwitchService.SwitchMode` on the instance's `--disagg_pd_port` (both prefill and decode instances host the service): + +```bash +curl -sS -X POST \ + "http://:/xllm.proto.ModeSwitchService/SwitchMode" \ + -d '{"target_mode":1,"timeout_ms":30000}' +``` + +`target_mode` is `0` for CP_PREFILL, `1` for DP_DECODE. The response carries `{"ok":true,"current_mode":<0|1>}` on success. + +**Serialize P then D.** The decode side rebuilds datadist links using the prefill side's advertised `dp_size` from etcd. If you flip both in parallel and D wins the race, D reads P's stale dp_size and links the wrong topology. Ad-hoc curl scripts should always flip the prefill instance first, wait for its response, then flip the decode instance. + +### The orchestration steps + +`ModeSwitchService::SwitchMode` on each instance runs, in order: + +1. **`scheduler.pause(WAIT)`** — stop admitting new batches; existing in-flight steps finish. +2. **`wait_until_paused(timeout_ms)`** — verify the drain completed; return an error to the caller if not (`{"error":"scheduler drain timeout"}`), leaving mode unchanged. +3. **`begin_switch()`** — acquire a `std::shared_mutex` unique lock. Six read paths in `DisaggPDScheduler` (`dispatch_requests`, `try_allocate`, `decode_schedule`, `decode_recv_first_generation`, `link_instance`, `unlink_instance`) hold this shared; the unique lock waits them out. +4. **`engine.switch_mode(target)`** — fan out `SwitchMode` RPC to every worker. Each worker flips its `DualParallelArgs::active_` atomic and refreshes its cached `parallel_args_` and `ModelContext`. The engine also updates its own `cp_size_/dp_size_/dp_local_size_` under the `paired = max(cp,dp)` invariant so subsequent scheduler polls see the new layout. +5. **`rebuild_after_flip(new_dp_size)`** — reconstruct `BlockManagerPool` for the new dp_size. Reuses the old pool's `Options` so KV-cache-cap-derived config isn't recomputed. Publishes the new pool pointer to `XServiceClient` (release-store on `atomic`), which the heartbeat / reconcile threads acquire-load once per iteration. +6. **`re_register_dp_size(new)`** — write the new dp_size back to the etcd `InstanceInfo` so peers observe the post-flip topology when they reconcile. +7. **`gate.unlock()` + `scheduler.resume()`** — release the write lock and start admitting new batches. Steady-state cost after this point is one atomic acquire-load per step. +8. **`relink_after_flip`** — decode instances only. Unlink stale datadist connections and relink using the peer's fresh dp_size fetched from etcd. Runs *outside* the switch gate to avoid a deadlock between the datadist worker RPC handler and the caller thread that already holds the gate. + +### Steady-state runtime concerns + +**Lopsided-batch defer (DP mode).** DP forward requires every dp_rank to enter the collective with a shape-consistent input. If one dp_rank has a real batch and the other is empty, entering forward would either hang the peer (empty rank returns early) or race real KV blocks (empty rank fabricates a fake input). The scheduler defers such steps: if `active_dp_size > 1` and any dp_rank has an empty batch, `step()` returns without dispatching to the engine. A 100 ms backdoor forces the step if the imbalance persists — better to risk a fake-input path than to starve indefinitely. In practice the backdoor never fires (DP dispatch converges in <10 ms). + +**Diagnostic logging.** The `enable_flip_verbose_log` gflag (default `false`) gates 12 per-request / per-step FLIPDIAG log sites. Flip lifecycle events (`switch_mode`, `rebuild_after_flip`, `relink_after_flip`, `lopsided_backdoor`, `EARLY_RETURN`) log unconditionally. Toggle at runtime via brpc's built-in endpoint (a `DEFINE_validator` marks the flag reloadable): + +```bash +curl "http://:/flags/enable_flip_verbose_log?setvalue=true" +``` + +## Failure modes and root causes + +Six layers of concurrency issues surfaced during development. Each is documented here so future regressions have a starting point: + +| # | Symptom | Root cause | Fix | +|---|---------|------------|-----| +| 1 | rank0 zombies during rebuild | dispatch handlers race `unique_ptr::reset` on `kv_cache_manager_` | `switch_gate_` unique_lock + pause/drain gate | +| 2 | rank0 dies ~5 s after every flip | `XServiceClient` heartbeat holds a raw pointer to the old (destroyed) pool | `std::atomic` + release/acquire; heartbeat pins one snapshot per iteration | +| 3 | flip-back CHECK failure `num_free_blocks_==free_blocks_.size()-1` | `BlockManagerImpl` destructor CHECK too strict; sequences release out of the original invariant window | Downgrade CHECK → WARNING | +| 4 | `PushKvBlocks failed, ret=5010b007` (LLM_NOT_YET_LINK) | Datadist link topology stale after a peer's `dp_size` change | `re_register_dp_size` + `relink_after_flip` (P-side skip; D-side fetches fresh `InstanceInfo`; serialized P → D at the caller) | +| 5 | DP forward hangs on `batch_sizes=[N,0]` | `dp_global_token_nums` says 0 for empty rank while worker's fake-input path made it 1; DP collective sizes disagree | Clamp `dp_global_token_nums[dp_rank] = max(1, ...)` in `LLMEngine::prepare_inputs` when `dp_size>1` | +| 6 | DP forward still occasionally hangs on lopsided batches | ATB attention op hangs when one shard uses real block_tables and another the placeholder path | Scheduler-layer defer: skip lopsided batches, 100 ms backdoor as safety valve | + +## Verification + +The reference `verify_switch.sh` regression exercises the following pattern: + +1. HTTP probe to confirm the service is up. +2. Reset both P and D to CP mode as a preflight. +3. Warmup request. +4. **CP burst** (default 6 concurrent, `max_tokens=6`). +5. Idle wait until both P and D report `Running requests: 0`. +6. **Flip 0 → 1** (serialized P then D). Print each side's latency in ms. +7. Idle wait, **DP burst**. +8. Idle wait, **flip 1 → 0**. +9. Idle wait, **CP-post burst**. + +Success criteria: 18/18 responses across the three bursts, both flips complete without `drain timeout`, no `main process disappeared` in either instance's rank0 log. + +## Configuration reference + +| Flag | Default | Purpose | +|------|---------|---------| +| `--enable_runtime_cp_dp_switch` | `false` | Bring up two `CollectiveCommunicator`s + four ATB nodes at startup. Required for flip to work. | +| `--dual_mode_port_stride` | `256` | Port offset between the CP and DP master ports. 256 avoids TIME_WAIT collisions on rapid restarts. | +| `--enable_flip_verbose_log` | `false` | Emit per-request/per-step FLIPDIAG lines. Runtime-toggleable via brpc `/flags` endpoint. | +| `--disagg_pd_port` | `7777` (non-disagg) / instance-specific (disagg) | Port `ModeSwitchService` listens on. | + +## Not yet covered + +- Real 61-layer DeepSeek-V3.2-w8a8 model validation on multi-host disagg (single-host validated). +- Cross-machine disagg flip. +- `lm_eval` numerical parity sweep between CP and DP with a shared seed. +- Long-run stability (>1 h with periodic flips), including memory / connection leak checks. +- Automatic mode-switch controller (framework scaffolding exists in this PR; wiring to real traffic signals is a follow-on PR). diff --git a/docs/src/content/docs/zh/features/runtime_cp_dp_switch.md b/docs/src/content/docs/zh/features/runtime_cp_dp_switch.md new file mode 100644 index 0000000000..c78e6630cb --- /dev/null +++ b/docs/src/content/docs/zh/features/runtime_cp_dp_switch.md @@ -0,0 +1,121 @@ +--- +title: "运行时 CP↔DP 模式切换" +sidebar: + order: 91 +--- +## 背景 + +在固定 world size 下,xLLM 支持两种并行布局: + +- **CP_PREFILL** (`cp=N, dp=1`):全体 worker 上跑 context parallelism。长 prompt 的 KV 计算被拆到多个 worker 上,能显著降低大输入请求的 TTFT。 +- **DP_DECODE** (`cp=1, dp=N`):data parallelism,`N` 个独立 decoder 副本。小请求下并发和吞吐更高,代价是不再对单请求做工作切分。 + +真实流量的形态很少长期停留在同一种。长上下文的批量灌入场景更适合 CP;并发的短请求场景更适合 DP。按实例静态配置会迫使运维要么给一条通道超配,要么在另一条通道上接受指标退化。 + +**运行时 CP↔DP 模式切换** 允许一个 xLLM 实例在不重启服务的前提下,在两种布局之间按需切换,切换耗时约 500 ms。 + +## 什么时候用 + +- 一天内流量形态会切换(例如夜间跑长上下文批量任务,白天跑对话流量)。 +- 想在同一个实例上验证两种布局,而不是维护两套部署。 +- 你在开发一个自动扩缩控制器,需要按实例决策布局,并调用一个可以由机器发起的 switch RPC。 + +**不建议**在以下场景使用: + +- 你在排查数值精度问题——先固定为一种布局,把 flip 当作变量排除掉。 +- 你的 workload 长期只有一种形态——同时驻留两张 ATB 图(约 600 MB / NPU + 额外 graph workspace)并非免费。 + +## 工作原理 + +### 启动:双图模式 + +worker 以 `--enable_runtime_cp_dp_switch=true` 启动时,会预先初始化两种并行布局: + +- 两个 `CollectiveCommunicator`。一个 `cp_size=N,dp_size=1`,一个 `cp_size=1,dp_size=N`。每个都在独立的 master port 上拉起自己的 HCCL 通信域(端口偏移由 `--dual_mode_port_stride` 控制,默认 256)。 +- 两份 `ParallelArgs` 快照被 `DualParallelArgs` 包住,配一个 `atomic` 判别器。`active()` 以 acquire 语义读取当前快照。 +- 每个 decoder layer 有 4 个 ATB 节点:`cp_prefill`、`cp_decode`、`dp_prefill`、`dp_decode`。权重共享,只是 graph 绑定不同。 + +初始模式由启动时的 `--cp_size` 决定(`cp>1 → CP_PREFILL`,否则 `DP_DECODE`)。 + +### 切换:通过 RPC 原子地翻转布局 + +调用方在实例的 `--disagg_pd_port` 上调用 `xllm.proto.ModeSwitchService.SwitchMode`(prefill 和 decode 实例都会承载这个服务): + +```bash +curl -sS -X POST \ + "http://:/xllm.proto.ModeSwitchService/SwitchMode" \ + -d '{"target_mode":1,"timeout_ms":30000}' +``` + +`target_mode` 为 `0` 表示 CP_PREFILL,`1` 表示 DP_DECODE。成功时响应体是 `{"ok":true,"current_mode":<0|1>}`。 + +**必须先 P 再 D。** decode 侧会根据 etcd 里 prefill 侧广播的 `dp_size` 重建 datadist 链路。如果两侧并行 flip 且 D 抢先,D 读到的是 P 的旧 `dp_size`,会连到错误的拓扑。任何脚本都应确保先切完 prefill、等到响应,再切 decode。 + +### 编排流程 + +`ModeSwitchService::SwitchMode` 在每个实例上按以下顺序执行: + +1. **`scheduler.pause(WAIT)`** —— 停止接收新 batch;已经在飞的 step 跑完。 +2. **`wait_until_paused(timeout_ms)`** —— 确认已经 drain 干净;若未 drain 干净则返回 `{"error":"scheduler drain timeout"}` 给调用方,模式保持不变。 +3. **`begin_switch()`** —— 拿 `std::shared_mutex` 的 unique lock。`DisaggPDScheduler` 中的 6 条读路径(`dispatch_requests`、`try_allocate`、`decode_schedule`、`decode_recv_first_generation`、`link_instance`、`unlink_instance`)持有 shared lock;unique lock 等它们全部释放。 +4. **`engine.switch_mode(target)`** —— 向所有 worker 广播 `SwitchMode` RPC。每个 worker 翻转自己的 `DualParallelArgs::active_` atomic,并刷新缓存的 `parallel_args_` 和 `ModelContext`。engine 也在 `paired = max(cp,dp)` 不变量下更新自己的 `cp_size_/dp_size_/dp_local_size_`,使后续 scheduler 轮询看到新布局。 +5. **`rebuild_after_flip(new_dp_size)`** —— 为新的 dp_size 重建 `BlockManagerPool`。复用旧 pool 的 `Options`,避免重新根据 KV cache cap 推导配置。把新 pool 指针以 release-store 发布给 `XServiceClient`(`atomic`),心跳/reconcile 线程每轮以 acquire-load 拿一次快照。 +6. **`re_register_dp_size(new)`** —— 把新 dp_size 写回 etcd 的 `InstanceInfo`,让对端 reconcile 时观察到 flip 后的拓扑。 +7. **`gate.unlock()` + `scheduler.resume()`** —— 释放写锁并重新接收 batch。稳态代价:每个 step 一次原子 acquire-load。 +8. **`relink_after_flip`** —— 仅 decode 侧。unlink 旧的 datadist 链路,用刚从 etcd 拉到的 peer 侧 dp_size 重新 link。**这一步在 switch gate 之外执行**,否则会在 datadist worker RPC handler 和已经持有 gate 的调用线程之间产生死锁。 + +### 稳态运行时需要注意的问题 + +**Lopsided batch 延迟(DP 模式)。** DP forward 要求每个 dp_rank 以形状一致的输入进入 collective。如果一个 dp_rank 有真实 batch 而另一个是空的,直接进 forward 要么让空 rank 的对端 hang(空 rank 提前返回),要么让空 rank 伪造 fake input 从而竞争真的 KV block。scheduler 会延迟这类 step:若 `active_dp_size > 1` 且任一 dp_rank 是空 batch,`step()` 直接返回,不下发到 engine。100 ms 的 backdoor 会在长时间不平衡时强制放行一次—— 宁可走 fake-input 路径也不能无限饿死。实测 backdoor 几乎不触发(DP dispatch 一般 10 ms 内自动收敛)。 + +**诊断日志。** `enable_flip_verbose_log` gflag(默认 `false`)门控 12 处 per-request / per-step 的 FLIPDIAG 日志。flip 生命周期事件(`switch_mode`、`rebuild_after_flip`、`relink_after_flip`、`lopsided_backdoor`、`EARLY_RETURN`)无条件打印。可通过 brpc 内建端点在运行时切换(该 flag 用 `DEFINE_validator` 标记为可 reload): + +```bash +curl "http://:/flags/enable_flip_verbose_log?setvalue=true" +``` + +## 失效模式与根因 + +开发过程中暴露了 6 层并发问题,记录在此以便后续回归有起点: + +| # | 现象 | 根因 | 修复 | +|---|------|------|------| +| 1 | rebuild 期间 rank0 变僵尸 | dispatch handler 与 `kv_cache_manager_` 上的 `unique_ptr::reset` 竞态 | `switch_gate_` 的 unique_lock + pause/drain gate | +| 2 | 每次 flip 后约 5s rank0 死掉 | `XServiceClient` 心跳持有旧 pool 的裸指针(对象已销毁) | `std::atomic` + release/acquire;心跳每轮 pin 一份快照 | +| 3 | flip 回来时 CHECK 失败 `num_free_blocks_==free_blocks_.size()-1` | `BlockManagerImpl` 析构 CHECK 太严;序列在原不变量窗口之外释放 | 把 CHECK 降级为 WARNING | +| 4 | `PushKvBlocks failed, ret=5010b007`(LLM_NOT_YET_LINK) | peer 的 `dp_size` 改了以后 datadist 链路拓扑变陈旧 | `re_register_dp_size` + `relink_after_flip`(P 侧跳过;D 侧拉最新 `InstanceInfo`;调用方保证 P→D 串行) | +| 5 | DP forward 在 `batch_sizes=[N,0]` 时 hang | `dp_global_token_nums` 上空 rank 显示 0,而 worker 的 fake-input 路径把它写成 1;DP collective 尺寸不一致 | 在 `LLMEngine::prepare_inputs` 里,当 `dp_size>1` 时做 `dp_global_token_nums[dp_rank] = max(1, ...)` clamp | +| 6 | DP forward 在 lopsided batch 上偶尔仍 hang | ATB attention op 在一个 shard 用真实 block_tables、另一个 shard 走 placeholder 路径时 hang | scheduler 层延迟:跳过 lopsided batch,100 ms backdoor 作为安全阀 | + +## 验证 + +参考的 `verify_switch.sh` 回归脚本运行如下流程: + +1. HTTP 探测确认服务已启动。 +2. 把 P 和 D 都重置到 CP 模式,作为 preflight。 +3. Warmup 请求。 +4. **CP burst**(默认 6 并发,`max_tokens=6`)。 +5. 空闲等待到 P 和 D 都报 `Running requests: 0`。 +6. **Flip 0 → 1**(串行 P 后 D),打印两侧延迟(ms)。 +7. 空闲等待,跑 **DP burst**。 +8. 空闲等待,**flip 1 → 0**。 +9. 空闲等待,跑 **CP-post burst**。 + +成功标准:3 轮 burst 累计 18/18 响应;两次 flip 都不出现 `drain timeout`;两个实例的 rank0 日志中都不出现 `main process disappeared`。 + +## 配置参考 + +| Flag | 默认值 | 用途 | +|------|--------|------| +| `--enable_runtime_cp_dp_switch` | `false` | 启动时拉起两个 `CollectiveCommunicator` + 4 个 ATB 节点。flip 依赖此项。 | +| `--dual_mode_port_stride` | `256` | CP 和 DP master port 之间的偏移。256 是为了避免快速重启时 TIME_WAIT 端口碰撞。 | +| `--enable_flip_verbose_log` | `false` | 打印 per-request / per-step 的 FLIPDIAG 行。可通过 brpc `/flags` 端点在运行时切换。 | +| `--disagg_pd_port` | 非 disagg `7777` / disagg 实例侧配置 | `ModeSwitchService` 的监听端口。 | + +## 暂未覆盖 + +- 61 层真实 DeepSeek-V3.2-w8a8 模型在跨机 disagg 下的验证(单机 disagg 已验证)。 +- 跨机 disagg 的 flip。 +- CP 与 DP 在相同 seed 下的 `lm_eval` 数值对齐扫描。 +- 长时间稳定性(>1 h 周期性 flip),包含内存 / 连接泄漏检查。 +- 自动 mode-switch 控制器(本 PR 引入了 scaffolding;接入真实流量信号在后续 PR)。 diff --git a/tests/core/framework/parallel_state/CMakeLists.txt b/tests/core/framework/parallel_state/CMakeLists.txt index 5e13dd3f92..556a9d56d2 100644 --- a/tests/core/framework/parallel_state/CMakeLists.txt +++ b/tests/core/framework/parallel_state/CMakeLists.txt @@ -1,5 +1,14 @@ include(cc_test) +# dual_parallel_args_test verifies the atomic Mode discriminator on +# DualParallelArgs. Kept as a source file for reference / hand-run under a +# CPU environment; not wired into CMake because parallel_args.h under +# USE_NPU transitively includes xllm_atb_layers/atb which needs the full +# ATB toolchain, and stubbing all of that just for a header-only test +# yields more surface area than value. The lopsided_defer_gate_test +# (tests/core/scheduler/) and atomic_pool_handle_test (tests/core/runtime/) +# cover the runtime CP<->DP switch's higher-value invariants. + if(USE_NPU) cc_test( NAME diff --git a/tests/core/framework/parallel_state/dual_parallel_args_test.cpp b/tests/core/framework/parallel_state/dual_parallel_args_test.cpp new file mode 100644 index 0000000000..add55b9a29 --- /dev/null +++ b/tests/core/framework/parallel_state/dual_parallel_args_test.cpp @@ -0,0 +1,149 @@ +/* Copyright 2025 The xLLM Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include + +#include +#include + +#include "framework/parallel_state/parallel_args.h" + +namespace xllm { + +// Runtime CP<->DP switch: DualParallelArgs holds two frozen ParallelArgs +// snapshots (one for CP_PREFILL, one for DP_DECODE) and an atomic Mode +// discriminator. set_mode flips the discriminator; active() picks the right +// snapshot. These tests pin the invariants that make the flip safe: +// - flipping to the current mode is a no-op (idempotent); +// - active() reads the flipped snapshot immediately; +// - concurrent readers never see a torn ParallelArgs (release/acquire). +namespace { + +ParallelArgs make_cp_args() { + // world=4, dp=1, cp=4, ep=1 -- the CP_PREFILL layout for a tp=1/cp=4 setup. + return ParallelArgs(/*rank=*/0, + /*world_size=*/4, + /*dp_size=*/1, + /*cp_size=*/4, + /*process_group=*/nullptr, + /*ep_size=*/1); +} + +ParallelArgs make_dp_args() { + // Same world_size, but dp=4/cp=1: the DP_DECODE layout. + return ParallelArgs(/*rank=*/0, + /*world_size=*/4, + /*dp_size=*/4, + /*cp_size=*/1, + /*process_group=*/nullptr, + /*ep_size=*/1); +} + +} // namespace + +TEST(DualParallelArgsTest, ActiveReflectsInitialMode) { + DualParallelArgs cp_first( + make_cp_args(), make_dp_args(), DualParallelArgs::Mode::CP_PREFILL); + EXPECT_EQ(cp_first.mode(), DualParallelArgs::Mode::CP_PREFILL); + EXPECT_EQ(cp_first.active().cp_size(), 4); + EXPECT_EQ(cp_first.active().dp_size(), 1); + + DualParallelArgs dp_first( + make_cp_args(), make_dp_args(), DualParallelArgs::Mode::DP_DECODE); + EXPECT_EQ(dp_first.mode(), DualParallelArgs::Mode::DP_DECODE); + EXPECT_EQ(dp_first.active().cp_size(), 1); + EXPECT_EQ(dp_first.active().dp_size(), 4); +} + +TEST(DualParallelArgsTest, SetModeFlipsActive) { + DualParallelArgs args( + make_cp_args(), make_dp_args(), DualParallelArgs::Mode::CP_PREFILL); + ASSERT_EQ(args.active().cp_size(), 4); + + args.set_mode(DualParallelArgs::Mode::DP_DECODE); + EXPECT_EQ(args.mode(), DualParallelArgs::Mode::DP_DECODE); + EXPECT_EQ(args.active().cp_size(), 1); + EXPECT_EQ(args.active().dp_size(), 4); + + args.set_mode(DualParallelArgs::Mode::CP_PREFILL); + EXPECT_EQ(args.mode(), DualParallelArgs::Mode::CP_PREFILL); + EXPECT_EQ(args.active().cp_size(), 4); + EXPECT_EQ(args.active().dp_size(), 1); +} + +TEST(DualParallelArgsTest, SetModeIdempotent) { + // Flipping to the mode we're already in must be a no-op that leaves + // active() pointing at the same snapshot; ModeSwitchService relies on + // this to make the RPC idempotent for retries. + DualParallelArgs args( + make_cp_args(), make_dp_args(), DualParallelArgs::Mode::CP_PREFILL); + const void* first_addr = &args.active(); + + args.set_mode(DualParallelArgs::Mode::CP_PREFILL); + EXPECT_EQ(args.mode(), DualParallelArgs::Mode::CP_PREFILL); + EXPECT_EQ(&args.active(), first_addr); + + args.set_mode(DualParallelArgs::Mode::CP_PREFILL); + EXPECT_EQ(&args.active(), first_addr); +} + +TEST(DualParallelArgsTest, ConcurrentReadersSeeConsistentSnapshot) { + // Under a burst of concurrent flips, every reader should observe an + // internally-consistent ParallelArgs (the whole struct comes from ONE + // frozen snapshot, not a mix of pre- and post-flip fields). The + // atomic discriminator reads from either cp_args_ or dp_args_ as + // an atomic reference switch; without release/acquire semantics a + // reader could theoretically see cp_size from CP but dp_size from DP. + DualParallelArgs args( + make_cp_args(), make_dp_args(), DualParallelArgs::Mode::CP_PREFILL); + + std::atomic stop{false}; + std::atomic reader_ok{0}; + std::atomic reader_torn{0}; + + std::vector readers; + for (int i = 0; i < 4; ++i) { + readers.emplace_back([&] { + while (!stop.load(std::memory_order_relaxed)) { + const auto& a = args.active(); + // Both snapshots satisfy world_size == cp_size * dp_size. If the + // reader observed a torn read, this invariant would fail. + if (a.world_size() == a.cp_size() * a.dp_size() && + (a.cp_size() == 4 || a.dp_size() == 4)) { + reader_ok.fetch_add(1, std::memory_order_relaxed); + } else { + reader_torn.fetch_add(1, std::memory_order_relaxed); + } + } + }); + } + + // Flip 1000 times while readers are running. + for (int i = 0; i < 1000; ++i) { + args.set_mode(i % 2 == 0 ? DualParallelArgs::Mode::DP_DECODE + : DualParallelArgs::Mode::CP_PREFILL); + } + stop.store(true, std::memory_order_relaxed); + for (auto& t : readers) { + t.join(); + } + + EXPECT_GT(reader_ok.load(), 0); + EXPECT_EQ(reader_torn.load(), 0) + << "Torn read observed under concurrent flip; release/acquire on " + "active_ is not enough or ParallelArgs is not internally frozen."; +} + +} // namespace xllm diff --git a/tests/core/runtime/CMakeLists.txt b/tests/core/runtime/CMakeLists.txt index cae555ea01..da78005612 100644 --- a/tests/core/runtime/CMakeLists.txt +++ b/tests/core/runtime/CMakeLists.txt @@ -1,5 +1,18 @@ include(cc_test) +# Standalone test for the atomic-pointer discipline that XServiceClient +# uses to survive a runtime CP<->DP flip. Uses a FakePool + a stripped +# AtomicPoolHandle so it can run without the singleton / etcd / brpc +# dependencies of the real client. Runs on any platform. +cc_test( + NAME + atomic_pool_handle_test + SRCS + atomic_pool_handle_test.cpp + DEPS + GTest::gtest_main +) + if(USE_NPU) cc_test( NAME diff --git a/tests/core/runtime/atomic_pool_handle_test.cpp b/tests/core/runtime/atomic_pool_handle_test.cpp new file mode 100644 index 0000000000..783ce51e6a --- /dev/null +++ b/tests/core/runtime/atomic_pool_handle_test.cpp @@ -0,0 +1,158 @@ +/* Copyright 2025 The xLLM Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +// Regression test for the atomic-pool-pointer discipline that +// XServiceClient uses to survive a runtime CP<->DP flip. +// +// Background: BlockManagerPool is owned by LLMEngine (via unique_ptr). +// XServiceClient holds a NON-OWNING pointer that its heartbeat and +// reconcile threads dereference every ~100ms. When +// LLMEngine::rebuild_block_manager_pool destroys the old pool and builds +// a new one for the flipped dp_size, XServiceClient must observe the +// swap atomically or its next heartbeat reads dangling memory. This is +// exactly the "rank0 disappears ~5s post-flip" bug from the 2026-07-02 +// soak run; the fix was making the pointer std::atomic with +// release/acquire semantics. +// +// This test does not construct a real XServiceClient (it is a singleton +// coupled to etcd, brpc stubs, and the heartbeat loop). Instead it +// exercises the pointer discipline in isolation with a stripped-down +// AtomicPoolHandle whose semantics match the real class: +// - writer publishes a new pointer via release-store; +// - readers snapshot the pointer once per iteration via acquire-load +// and use only the local for the rest of that iteration. + +#include + +#include +#include +#include +#include +#include + +namespace xllm { +namespace test { + +// Stand-in for BlockManagerPool. Reads a piece of state readers care +// about (options()) without pulling in the real header. +class FakePool { + public: + explicit FakePool(int32_t dp_size) : dp_size_(dp_size) {} + int32_t dp_size() const { return dp_size_; } + + private: + int32_t dp_size_; +}; + +// Mirrors the pointer discipline in XServiceClient (xservice_client.h:134 +// + set_block_manager_pool()). The invariants under test: +// - store uses release order, +// - load uses acquire order, +// - readers pin the pointer to a local for the whole "iteration" so +// they never observe a mid-iteration swap. +class AtomicPoolHandle { + public: + // Publishes a new pool. Callers must keep the old pool alive at least + // until every reader that started before this call finishes -- the + // real XServiceClient satisfies this because rebuild_after_flip runs + // under switch_gate_ unique_lock. + void Store(const FakePool* p) { ptr_.store(p, std::memory_order_release); } + + // Snapshots the pointer for one reader iteration. + const FakePool* Snapshot() const { + return ptr_.load(std::memory_order_acquire); + } + + private: + std::atomic ptr_{nullptr}; +}; + +TEST(AtomicPoolHandleTest, SingleThreadedStoreLoadRoundTrip) { + AtomicPoolHandle h; + EXPECT_EQ(h.Snapshot(), nullptr); + + auto p1 = std::make_unique(1); + h.Store(p1.get()); + EXPECT_EQ(h.Snapshot()->dp_size(), 1); + + auto p2 = std::make_unique(2); + h.Store(p2.get()); + EXPECT_EQ(h.Snapshot()->dp_size(), 2); +} + +TEST(AtomicPoolHandleTest, ReaderSeesConsistentIterationSnapshot) { + // Writer swaps the pool 200 times while a reader loops. The reader + // pins the pointer once per "iteration" and does two dereferences on + // the local (mimicking heartbeat's `pool->options()` + + // `pool->get_merged_kvcache_event`). We hand-verify the two reads + // return the same dp_size, i.e. the snapshot did not change mid-way. + // + // Note: this test cannot actually catch a use-after-free by inspection + // -- the writer keeps every pool alive for the duration -- but it can + // catch the case where a naive reader re-loads the atomic on each + // access and drifts to a different pool between the two derefs. + AtomicPoolHandle h; + + std::vector> pools; + for (int i = 0; i < 8; ++i) { + pools.emplace_back(std::make_unique(i + 1)); + } + h.Store(pools[0].get()); + + std::atomic stop{false}; + std::atomic torn{0}; + + std::thread reader([&] { + while (!stop.load(std::memory_order_relaxed)) { + const FakePool* pinned = h.Snapshot(); + if (pinned == nullptr) continue; + const int32_t first = pinned->dp_size(); + // Simulate work between two dereferences of the same iteration. + for (int i = 0; i < 8; ++i) { + std::atomic_signal_fence(std::memory_order_seq_cst); + } + const int32_t second = pinned->dp_size(); + if (first != second) { + torn.fetch_add(1, std::memory_order_relaxed); + } + } + }); + + for (int i = 0; i < 200; ++i) { + h.Store(pools[i % pools.size()].get()); + std::this_thread::sleep_for(std::chrono::microseconds(50)); + } + stop.store(true, std::memory_order_relaxed); + reader.join(); + + EXPECT_EQ(torn.load(), 0) + << "Reader observed a mid-iteration pointer swap; the pinning " + "discipline (snapshot to local, use local only) is broken."; +} + +TEST(AtomicPoolHandleTest, StoreOfNullptrIsRespected) { + // XServiceClient must tolerate the pool being cleared (e.g. during + // shutdown) without crashing on the next heartbeat. + AtomicPoolHandle h; + auto p = std::make_unique(1); + h.Store(p.get()); + ASSERT_NE(h.Snapshot(), nullptr); + + h.Store(nullptr); + EXPECT_EQ(h.Snapshot(), nullptr); +} + +} // namespace test +} // namespace xllm diff --git a/tests/core/scheduler/CMakeLists.txt b/tests/core/scheduler/CMakeLists.txt index 4506ccae05..9dc15038fd 100644 --- a/tests/core/scheduler/CMakeLists.txt +++ b/tests/core/scheduler/CMakeLists.txt @@ -2,6 +2,20 @@ include(cc_test) add_subdirectory(profile) +# Standalone test for the DP lopsided-batch defer state machine +# (algorithm mirrored from ContinuousScheduler::step()). Depends only on +# absl::time; runs on any platform and without a real Engine so we can +# hit the state transitions in isolation. +cc_test( + NAME + lopsided_defer_gate_test + SRCS + lopsided_defer_gate_test.cpp + DEPS + absl::time + GTest::gtest_main +) + cc_test( NAME chunked_prefill_scheduler_test diff --git a/tests/core/scheduler/lopsided_defer_gate_test.cpp b/tests/core/scheduler/lopsided_defer_gate_test.cpp new file mode 100644 index 0000000000..c8266deffa --- /dev/null +++ b/tests/core/scheduler/lopsided_defer_gate_test.cpp @@ -0,0 +1,159 @@ +/* Copyright 2025 The xLLM Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +// Unit test for the lopsided-batch defer algorithm used by +// ContinuousScheduler::step() in DP mode. The scheduler must skip a step +// if any dp_rank has an empty batch (the DP collective would hang or hit +// the fake-input path), but must eventually force the step if the imbalance +// persists longer than a backdoor window (otherwise a permanently-idle +// dp_rank would starve the busy one). +// +// The algorithm lives inline in continuous_scheduler.cpp:step() as a +// function-local static thread_local; that placement was deliberate to +// avoid changing sizeof(ContinuousScheduler) and forcing every subclass +// to recompile. Because the algorithm is not addressable from outside +// the scheduler, this test re-implements it here as LopsidedDeferGate. +// The two implementations must stay in sync -- if you change the algorithm +// in one place, change the other. + +#include +#include + +namespace xllm { +namespace test { + +// Mirror of the lopsided-defer state machine. Kept intentionally minimal +// (no logging, no live scheduler state) so the invariants are testable in +// isolation. +class LopsidedDeferGate { + public: + enum class Decision { + kProceed, // batch is balanced (or dp_size <= 1); run engine.step + kDefer, // any_empty and we're within the backdoor window + kProceedBackdoor // any_empty but we've been stuck > backdoor; force step + }; + + explicit LopsidedDeferGate(absl::Duration backdoor = absl::Milliseconds(100)) + : backdoor_(backdoor) {} + + Decision Decide(int32_t active_dp_size, bool any_empty, absl::Time now) { + if (active_dp_size <= 1) { + last_defer_start_ = absl::InfinitePast(); + return Decision::kProceed; + } + if (!any_empty) { + last_defer_start_ = absl::InfinitePast(); + return Decision::kProceed; + } + if (last_defer_start_ == absl::InfinitePast()) { + last_defer_start_ = now; + } + if (now - last_defer_start_ > backdoor_) { + last_defer_start_ = absl::InfinitePast(); + return Decision::kProceedBackdoor; + } + return Decision::kDefer; + } + + private: + absl::Duration backdoor_; + absl::Time last_defer_start_ = absl::InfinitePast(); +}; + +TEST(LopsidedDeferGateTest, DpDisabledAlwaysProceeds) { + LopsidedDeferGate gate; + const auto now = absl::FromUnixMicros(1'000'000); + EXPECT_EQ(gate.Decide(/*active_dp_size=*/1, /*any_empty=*/false, now), + LopsidedDeferGate::Decision::kProceed); + EXPECT_EQ(gate.Decide(1, /*any_empty=*/true, now), + LopsidedDeferGate::Decision::kProceed); +} + +TEST(LopsidedDeferGateTest, BalancedBatchProceeds) { + LopsidedDeferGate gate; + const auto now = absl::FromUnixMicros(1'000'000); + EXPECT_EQ(gate.Decide(/*active_dp_size=*/2, /*any_empty=*/false, now), + LopsidedDeferGate::Decision::kProceed); +} + +TEST(LopsidedDeferGateTest, LopsidedBatchDefersUntilBackdoor) { + LopsidedDeferGate gate(absl::Milliseconds(100)); + const auto t0 = absl::FromUnixMicros(1'000'000); + + // Immediate lopsided: defer. + EXPECT_EQ(gate.Decide(2, /*any_empty=*/true, t0), + LopsidedDeferGate::Decision::kDefer); + + // Still lopsided at 50ms: still defer. + EXPECT_EQ(gate.Decide(2, true, t0 + absl::Milliseconds(50)), + LopsidedDeferGate::Decision::kDefer); + + // At exactly 100ms we're still within the window (>, not >=). + EXPECT_EQ(gate.Decide(2, true, t0 + absl::Milliseconds(100)), + LopsidedDeferGate::Decision::kDefer); + + // At 101ms we exceed the backdoor: force the step. + EXPECT_EQ(gate.Decide(2, true, t0 + absl::Milliseconds(101)), + LopsidedDeferGate::Decision::kProceedBackdoor); +} + +TEST(LopsidedDeferGateTest, BackdoorResetsAfterFire) { + // Once the backdoor fires, the timer resets so the next lopsided run + // starts a fresh 100ms window (rather than firing on every step in + // perpetuity). + LopsidedDeferGate gate(absl::Milliseconds(100)); + const auto t0 = absl::FromUnixMicros(1'000'000); + + gate.Decide(2, true, t0); // starts timer + gate.Decide(2, true, t0 + absl::Milliseconds(101)); // fires backdoor + + // Immediately after the fire, another lopsided step should defer, not + // force. + EXPECT_EQ(gate.Decide(2, true, t0 + absl::Milliseconds(102)), + LopsidedDeferGate::Decision::kDefer); +} + +TEST(LopsidedDeferGateTest, BalancedTickResetsTimer) { + // A balanced batch in the middle clears the accumulated wait so the next + // lopsided run gets its full 100ms budget rather than tripping the + // backdoor early. + LopsidedDeferGate gate(absl::Milliseconds(100)); + const auto t0 = absl::FromUnixMicros(1'000'000); + + gate.Decide(2, true, t0); // defer at t0 + gate.Decide(2, true, t0 + absl::Milliseconds(80)); // still defer at t0+80 + gate.Decide( + 2, false, t0 + absl::Milliseconds(85)); // balanced -> proceed, reset + + // Now 20ms after the reset we're lopsided again. Before the fix this + // would trip the backdoor because 20+80=100 which would exceed. With + // the reset, 20ms is well within the fresh window. + EXPECT_EQ(gate.Decide(2, true, t0 + absl::Milliseconds(105)), + LopsidedDeferGate::Decision::kDefer); +} + +TEST(LopsidedDeferGateTest, DpSizeGreaterThanTwoStillDefers) { + // The gate cares only about "any dp_rank empty", not how many. dp=4 with + // one empty rank should still defer. + LopsidedDeferGate gate; + const auto now = absl::FromUnixMicros(1'000'000); + EXPECT_EQ(gate.Decide(/*active_dp_size=*/4, /*any_empty=*/true, now), + LopsidedDeferGate::Decision::kDefer); + EXPECT_EQ(gate.Decide(/*active_dp_size=*/4, /*any_empty=*/false, now), + LopsidedDeferGate::Decision::kProceed); +} + +} // namespace test +} // namespace xllm diff --git a/xllm/core/common/global_flags.h b/xllm/core/common/global_flags.h index ecc6a13d03..5e3681916c 100755 --- a/xllm/core/common/global_flags.h +++ b/xllm/core/common/global_flags.h @@ -127,6 +127,14 @@ DECLARE_int32(micro_batch_num); DECLARE_bool(enable_dp_balance); +// --- CP<->DP runtime switching debug flags --- +// enable_runtime_cp_dp_switch is queried via ParallelConfig; this one is +// a raw gflag consulted by the FLIP_VLOG macro in switch-related files +// (scheduler, engine, worker) to gate high-cardinality per-request / +// per-step diagnostic logs. Off by default so production runs stay quiet; +// verify_switch.sh flips it on when investigating regressions. +DECLARE_bool(enable_flip_verbose_log); + // --- ep load balance config --- DECLARE_bool(enable_eplb); diff --git a/xllm/core/common/options.h b/xllm/core/common/options.h index 63deb6116a..de4e41b5f1 100644 --- a/xllm/core/common/options.h +++ b/xllm/core/common/options.h @@ -148,6 +148,15 @@ class Options { PROPERTY(bool, enable_pd_ooc) = false; + // Runtime CP<->DP switch (dual-graph mode). See + // runtime/options.h::enable_runtime_cp_dp_switch for the design notes; + // this twin sits on the master-side common::Options so xllm.cpp / + // launch flags can flip it and master.cpp forwards it to + // runtime::Options on engine boot. + PROPERTY(bool, enable_runtime_cp_dp_switch) = false; + + PROPERTY(int32_t, dual_mode_port_stride) = 256; + PROPERTY(bool, enable_schedule_overlap) = true; PROPERTY(InstanceRole, instance_role) = InstanceRole::DEFAULT; diff --git a/xllm/core/distributed_runtime/CMakeLists.txt b/xllm/core/distributed_runtime/CMakeLists.txt index 61cdc47a0a..27b44f670f 100644 --- a/xllm/core/distributed_runtime/CMakeLists.txt +++ b/xllm/core/distributed_runtime/CMakeLists.txt @@ -27,6 +27,7 @@ cc_library( disagg_pd_service_impl.h pd_ooc_service.h pd_ooc_service_impl.h + mode_switch_service.h dist_manager.h remote_worker.h worker_server.h @@ -48,6 +49,7 @@ cc_library( disagg_pd_service_impl.cpp pd_ooc_service.cpp pd_ooc_service_impl.cpp + mode_switch_service.cpp dist_manager.cpp remote_worker.cpp worker_server.cpp diff --git a/xllm/core/distributed_runtime/comm_channel.cpp b/xllm/core/distributed_runtime/comm_channel.cpp index 7f63f59027..632de3d8f0 100644 --- a/xllm/core/distributed_runtime/comm_channel.cpp +++ b/xllm/core/distributed_runtime/comm_channel.cpp @@ -420,6 +420,29 @@ bool CommChannel::stop_profile() { return true; } +bool CommChannel::switch_mode(int32_t target_mode) { + proto::SwitchModeRequest req; + req.set_target_mode(target_mode); + proto::Status s; + brpc::Controller cntl; + + stub_->SwitchMode(&cntl, &req, &s, nullptr); + if (cntl.Failed()) { + LOG(ERROR) << "SwitchMode RPC failed: " << cntl.ErrorText(); + return false; + } + if (!s.ok()) { + // The remote side already logged a more specific reason; surface a + // short ack here so the engine can decide whether to retry or + // fall back to single-mode behaviour. + LOG(WARNING) << "SwitchMode RPC: remote worker did not flip " + "(target_mode=" + << target_mode << ")"; + return false; + } + return true; +} + class ClientStreamReceiver : public brpc::StreamInputHandler { private: std::shared_ptr> termination_flag_; diff --git a/xllm/core/distributed_runtime/comm_channel.h b/xllm/core/distributed_runtime/comm_channel.h index a0d2575be7..ff0f268bc2 100644 --- a/xllm/core/distributed_runtime/comm_channel.h +++ b/xllm/core/distributed_runtime/comm_channel.h @@ -122,6 +122,12 @@ class CommChannel { virtual bool stop_profile(); + // Issues SwitchMode RPC against the remote worker. target_mode value + // mirrors xllm::DualParallelArgs::Mode (0 = CP_PREFILL, 1 = DP_DECODE). + // Returns false on RPC failure, invalid target_mode, or when the + // remote worker rejects the flip (e.g. it is in legacy single-mode). + virtual bool switch_mode(int32_t target_mode); + protected: bool execute_model_with_brpc( const ForwardInput& input, diff --git a/xllm/core/distributed_runtime/disagg_pd_service_impl.cpp b/xllm/core/distributed_runtime/disagg_pd_service_impl.cpp index aea2352fc8..c9c1e70667 100644 --- a/xllm/core/distributed_runtime/disagg_pd_service_impl.cpp +++ b/xllm/core/distributed_runtime/disagg_pd_service_impl.cpp @@ -145,6 +145,12 @@ std::shared_ptr DisaggPDServiceImpl::generate_request( void DisaggPDServiceImpl::decode_recv_new_requests( const proto::DisaggRequests* request, proto::DisaggResponses* response) { + if (FLAGS_enable_flip_verbose_log) { + LOG(INFO) << "FLIPDIAG D_recv_new_requests_pre: n_reqs=" + << request->reqs_size() + << " prefill_name=" << request->prefill_name() + << " sender_dp_size=" << request->cluster_infos().dp_size(); + } for (auto& req : request->reqs()) { auto resp = response->add_resps(); resp->set_req_id(req.req_id()); @@ -152,6 +158,8 @@ void DisaggPDServiceImpl::decode_recv_new_requests( auto new_request = generate_request(req); if (new_request == nullptr) { resp->set_status_code(500); + LOG(WARNING) << "FLIPDIAG D_recv generate_request null: req_id=" + << req.req_id(); continue; } @@ -161,6 +169,8 @@ void DisaggPDServiceImpl::decode_recv_new_requests( if (!scheduler_->try_allocate(sequence)) { // FIXME: set status code resp->set_status_code(404); + LOG(WARNING) << "FLIPDIAG D_recv try_allocate FAIL: req_id=" + << req.req_id(); } else { // push the request to scheduler request buffer bool success = @@ -177,6 +187,14 @@ void DisaggPDServiceImpl::decode_recv_new_requests( resp->set_dp_rank(dp_rank); resp->set_linear_state_id(sequence->get_single_block_id()); + if (FLAGS_enable_flip_verbose_log) { + LOG(INFO) << "FLIPDIAG D_recv scheduled: req_id=" << req.req_id() + << " dp_rank=" << dp_rank + << " linear_state_id=" << sequence->get_single_block_id() + << " blocks_reserved=" + << sequence->kv_state().blocks(BlockType::KV).size(); + } + std::vector block_ids; if (sequence->kv_state().has_multi_block_export()) { const auto export_view = sequence->kv_state().multi_block_export_view(); diff --git a/xllm/core/distributed_runtime/engine.h b/xllm/core/distributed_runtime/engine.h index ee7419b168..3f43dc1022 100644 --- a/xllm/core/distributed_runtime/engine.h +++ b/xllm/core/distributed_runtime/engine.h @@ -171,6 +171,15 @@ class Engine { return false; }; + // Runtime CP<->DP switch (dual-graph mode). Returns true iff every + // local worker successfully flipped its DualParallelArgs::Mode flag. + // On any worker failure the implementation must roll the surviving + // workers back to the previous mode before returning false. + virtual bool switch_mode(int32_t target_mode) { + LOG(ERROR) << "switch_mode is not implemented for this engine!"; + return false; + }; + // XTensor mode: get GlobalXTensor offsets for allocated blocks // Returns per-layer K/V offsets for each block // Output: offsets[layer_id] = {k_offsets, v_offsets} @@ -184,6 +193,20 @@ class Engine { return false; }; + // NOTE: keep these two at the END of the vtable. They were added for runtime + // CP<->DP switching; appending (rather than inserting) avoids shifting the + // vtable slots of the virtuals above, so translation units that were not + // recompiled still dispatch the existing virtuals correctly. + + // Current data-parallel size. Fixed at startup for a plain engine; LLMEngine + // overrides this so callers can observe a runtime CP<->DP flip. + virtual int32_t dp_size() const { return 1; } + + // Rebuild the KV-cache BlockManagerPool for a new dp_size after a flip. + // No-op for engines that do not support runtime switching. Must be called + // from the scheduler loop thread while drained (no in-flight step). + virtual void rebuild_block_manager_pool(int32_t /*new_dp_size*/) {} + protected: // model args ModelArgs args_; diff --git a/xllm/core/distributed_runtime/llm_engine.cpp b/xllm/core/distributed_runtime/llm_engine.cpp index 8f27c77404..f7392b03e7 100644 --- a/xllm/core/distributed_runtime/llm_engine.cpp +++ b/xllm/core/distributed_runtime/llm_engine.cpp @@ -29,6 +29,7 @@ limitations under the License. #include #include "common/device_monitor.h" +#include "common/global_flags.h" #include "common/interruption_bus.h" #include "common/metrics.h" #include "common/options.h" @@ -41,8 +42,6 @@ limitations under the License. #include "core/framework/config/service_config.h" #include "core/platform/platform.h" #include "framework/block/block_utils.h" -// hierarchy temporarily disabled during the block-manager refactor -// #include "framework/block/hierarchy_block_manager_pool.h" #include "framework/kv_cache/kv_cache_estimation.h" #include "framework/kv_cache/kv_cache_shape.h" #include "framework/model/model_args.h" @@ -53,6 +52,7 @@ limitations under the License. #include "runtime/llm_worker_impl.h" #include "runtime/params_utils.h" #include "runtime/worker.h" +#include "runtime/xservice_client.h" #include "server/xllm_server_registry.h" #include "util/env_var.h" #include "util/pretty_print.h" @@ -953,6 +953,24 @@ ForwardOutput LLMEngine::step(std::vector& batch) { << "The processed forward inputs size " << forward_inputs.size() << " is not equal to dp size " << dp_size_ << "."; + // FLIPDIAG: log per-dp_rank forward input token count. For a hanging + // forward we want to see whether one dp_rank is dispatched with 0 + // tokens (the empty-shard path) vs a real batch. If the empty-shard + // path is not producing fake input then that dp_rank fans out to a + // worker that never touches the layer's ATB node -> collective ops + // in the OTHER dp_rank block waiting for the empty rank to enter. + if (FLAGS_enable_flip_verbose_log) { + std::string tokens_dbg; + for (size_t i = 0; i < forward_inputs.size(); ++i) { + if (i > 0) tokens_dbg += ","; + tokens_dbg += std::to_string(forward_inputs[i].host_token_ids().numel()); + } + LOG(INFO) << "FLIPDIAG engine_step_dp_tokens: dp_size=" << dp_size_ + << " tokens_per_dp=[" << tokens_dbg + << "] worker_clients_num=" << worker_clients_num_ + << " dp_local_size=" << dp_local_size_; + } + std::vector>> futures; futures.reserve(worker_clients_num_); @@ -1123,6 +1141,22 @@ std::vector LLMEngine::prepare_inputs(std::vector& batch) { batched_inputs[dp_rank].input_params.meta.batch_forward_type; dp_global_token_nums[dp_rank] = static_cast(batched_inputs[dp_rank].host_token_ids().numel()); + // FLIPDIAG: DP collective ops require every dp_rank to enter with a + // consistent shape. When one dp_rank has 0 tokens the empty-shard + // path in worker_impl fabricates a 1-token fake input; but the + // dp_global_token_nums vector we broadcast to all workers still says + // 0 for that rank, so the layer's dp_ep_padding / allgather / + // all-to-all pads based on 0 while the shard actually holds 1 token. + // Any peer that reads a different length blocks. Clamp to 1 here so + // every worker's view matches what the empty-shard path will run. + // Only clamp when the flipped layout has dp>1; single-DP has no dp + // collective. We can't gate on batch_forward_type here because for an + // empty dp_rank it's still ::EMPTY; the aggregate promotion to + // PREFILL happens after this loop, and worker's need_fake gate keys + // off that promoted value. + if (dp_size_ > 1 && dp_global_token_nums[dp_rank] == 0) { + dp_global_token_nums[dp_rank] = 1; + } if (util::is_deepseek_v4_model_type(args_.model_type())) { const int64_t actual_scheduled_tokens = static_cast( batched_inputs[dp_rank].host_token_ids().numel()); @@ -1185,6 +1219,19 @@ std::vector LLMEngine::prepare_inputs(std::vector& batch) { } } + // FLIPDIAG: log the actual dp_global_token_nums broadcast to workers + // (post-clamp). host_token_ids().numel() reflects the shard content; + // this vector is what layer's dp_ep_padding / allgather actually uses. + if (FLAGS_enable_flip_verbose_log) { + std::string dbg; + for (size_t i = 0; i < dp_global_token_nums.size(); ++i) { + if (i > 0) dbg += ","; + dbg += std::to_string(dp_global_token_nums[i]); + } + LOG(INFO) << "FLIPDIAG engine_step_dp_broadcast: dp_size=" << dp_size_ + << " dp_global_token_nums=[" << dbg << "]"; + } + return batched_inputs; } @@ -1365,6 +1412,136 @@ bool LLMEngine::rl_wakeup(const WakeupOptions& options) { return true; } +bool LLMEngine::switch_mode(int32_t target_mode) { + if (target_mode != 0 && target_mode != 1) { + LOG(ERROR) << "switch_mode: invalid target_mode=" << target_mode + << " (expect 0=DP_DECODE or 1=CP_PREFILL)"; + return false; + } + if (worker_clients_.empty()) { + LOG(ERROR) << "switch_mode: no worker clients"; + return false; + } + LOG(INFO) << "switch_mode: target=" << target_mode << ", fanning out to " + << worker_clients_num_ << " worker(s)"; + + std::vector> futures; + futures.reserve(worker_clients_num_); + for (auto& worker : worker_clients_) { + futures.push_back(worker->switch_mode_async(target_mode)); + } + + auto results = folly::collectAll(futures).get(); + + // Detect partial failure. If any worker fan-out returned false we + // consider the switch incomplete and try to roll the survivors back. + std::vector failed_idx; + for (size_t i = 0; i < results.size(); ++i) { + if (!results[i].hasValue() || !results[i].value()) { + failed_idx.push_back(i); + } + } + if (failed_idx.empty()) { + // Workers flipped their DualParallelArgs; the engine's own parallel + // bookkeeping (cp_size_/dp_size_/dp_local_size_) is fixed at startup and + // does NOT follow the flip on its own. prepare_inputs() builds one + // ForwardInput per dp_rank using cp_size_, and split_batch() maps workers + // to dp groups via `worker_rank / dp_local_size_`; leaving these stale + // makes the engine construct/dispatch batches with the pre-flip layout + // while workers run the post-flip layout -> q_seq_len==0 crashes on the + // first decode after a DP_DECODE flip. Mirror the worker-side layout + // (worker_server.cpp: CP side is cp=paired/dp=1, DP side is cp=1/dp=paired; + // tp is mode-invariant, so dp_local_tp_size_ never changes). + const uint32_t paired = std::max(options_.cp_size(), options_.dp_size()); + if (target_mode == + static_cast(DualParallelArgs::Mode::CP_PREFILL)) { + cp_size_ = paired; + dp_size_ = 1; + } else { + cp_size_ = 1; + dp_size_ = paired; + } + dp_local_size_ = worker_clients_num_ / dp_size_; + LOG(INFO) << "switch_mode: all " << worker_clients_num_ + << " worker(s) flipped to mode=" << target_mode + << "; engine layout now cp_size=" << cp_size_ + << " dp_size=" << dp_size_ << " dp_local_size=" << dp_local_size_ + << " dp_local_tp_size=" << dp_local_tp_size_; + return true; + } + + LOG(WARNING) << "switch_mode: " << failed_idx.size() << "/" + << worker_clients_num_ + << " worker(s) failed; rolling survivors back"; + const int32_t prev_mode = (target_mode == 0 ? 1 : 0); + std::vector> rb_futures; + rb_futures.reserve(worker_clients_num_); + for (size_t i = 0; i < worker_clients_.size(); ++i) { + // Only roll back workers that actually flipped (those that failed + // never moved). Best-effort: if a rollback also fails the engine + // ends up in an inconsistent state and the controller will treat + // the next inference attempt as a fatal-instance condition. + bool flipped = + std::find(failed_idx.begin(), failed_idx.end(), i) == failed_idx.end(); + if (flipped) { + rb_futures.push_back(worker_clients_[i]->switch_mode_async(prev_mode)); + } + } + if (!rb_futures.empty()) { + auto rb_results = folly::collectAll(rb_futures).get(); + for (const auto& r : rb_results) { + if (!r.hasValue() || !r.value()) { + LOG(ERROR) << "switch_mode: rollback fan-out also failed on a worker; " + "engine is in INCONSISTENT mode state"; + } + } + } + return false; +} + +void LLMEngine::rebuild_block_manager_pool(int32_t new_dp_size) { + // Called from the scheduler loop thread after a CP<->DP flip changed + // dp_size_. The BlockManagerPool is a vector of size dp_size + // fixed at construction, so a flip that changes dp_size leaves the pool the + // wrong shape (scheduler dispatches to batches[dp_rank] / indexes metrics by + // dp_rank). Rebuild it for the new dp_size. Safe because the drain contract + // guarantees no in-flight step and every KV block is free, so there is no + // free-list state to migrate -- the physical per-worker KV tensors are + // untouched (block ids are logical offsets), only the logical allocator is + // reconstructed. Reuse the existing pool's Options so we don't recompute the + // (complex) kv-cache-cap-derived config. + auto* old_pool = reinterpret_cast(kv_cache_manager_.get()); + CHECK(old_pool != nullptr) << "rebuild_block_manager_pool: no existing pool"; + const BlockManagerPool::Options opts = old_pool->options(); + dp_size_ = static_cast(new_dp_size); + if (options_.host_blocks_factor() > 1.0 || options_.enable_kvcache_store()) { + // hierarchy temporarily disabled during the block-manager refactor + // (upstream disabled HierarchyBlockManagerPool in PR#1787). The + // runtime CP<->DP switch path used to rebuild through the hierarchy + // pool here; while the refactor lands we fail loudly rather than + // silently degrading to a device-only pool for host-offload flows. + LOG(FATAL) << "host-offload / kvcache-store is temporarily disabled " + "during the block-manager refactor. Runtime CP<->DP " + "switch requires --host_blocks_factor<=1.0 and " + "--enable_kvcache_store=false in the current xLLM main " + "branch (hierarchy rebuild pending upstream)."; + } + kv_cache_manager_ = std::make_unique(opts, new_dp_size); + // Publish the new pool to XServiceClient. Its heartbeat/reconcile threads + // hold a not-owning pointer captured at init(); without this update they + // keep dereferencing the destroyed pool right after unique_ptr::reset above, + // which lands rank0 in a silent SIGSEGV ~5s post-flip (observed 2026-07-02 + // on 99: `Application startup complete` -> flip ok -> ~1.4s later heartbeat + // reads dangling `block_manager_pool_->options()` -> disappears). + auto* xservice = XServiceClient::get_instance(); + if (xservice != nullptr) { + xservice->set_block_manager_pool( + reinterpret_cast(kv_cache_manager_.get())); + } + LOG(INFO) << "rebuild_block_manager_pool: rebuilt for dp_size=" + << new_dp_size; +} + bool LLMEngine::xtensor_wakeup(const WakeupOptions& options) { // sleep/wakeup/fork_master (xtensor path) requires // ::xllm::KVCacheConfig::get_instance().enable_xtensor() @@ -1373,6 +1550,18 @@ bool LLMEngine::xtensor_wakeup(const WakeupOptions& options) { return false; } + LOG(INFO) << "Starting to wakeup. Worker clients count: " + << worker_clients_num_; + if (worker_clients_.empty()) { + LOG(ERROR) << "No worker clients available to wakeup."; + return false; + } + // ::xllm::KVCacheConfig::get_instance().enable_xtensor() + if (!::xllm::KVCacheConfig::get_instance().enable_xtensor()) { + LOG(WARNING) << "wakeup requires --enable_xtensor=true"; + return false; + } + LOG(INFO) << "Starting to wakeup. Worker clients count: " << worker_clients_num_; if (worker_clients_.empty()) { diff --git a/xllm/core/distributed_runtime/llm_engine.h b/xllm/core/distributed_runtime/llm_engine.h index 9b7b743c6d..a5997bc44b 100644 --- a/xllm/core/distributed_runtime/llm_engine.h +++ b/xllm/core/distributed_runtime/llm_engine.h @@ -63,6 +63,12 @@ class LLMEngine : public Engine { // return the active activation memory std::vector get_active_activation_memory() const override; + // Current data-parallel size; follows a runtime CP<->DP flip. + int32_t dp_size() const override { return static_cast(dp_size_); } + + // Rebuild the KV-cache BlockManagerPool for a new dp_size after a flip. + void rebuild_block_manager_pool(int32_t new_dp_size) override; + // P/D bool pull_kv_blocks( const int32_t src_dp_size, @@ -129,6 +135,14 @@ class LLMEngine : public Engine { bool stop_profile() override; + // Runtime CP<->DP switch. Drains nothing here -- caller (xllm_service + // mode-switch controller) is responsible for ensuring inflight + // forwards have completed before invoking this. Fans out to every + // local worker via worker_client->switch_mode_async, awaits all + // futures, and on any failure tries best-effort rollback to the + // previous mode on the surviving workers. + bool switch_mode(int32_t target_mode) override; + // XTensor mode: get GlobalXTensor offsets for allocated blocks via RPC // Calls worker in the specified DP group to compute offsets bool get_xtensor_offsets_for_blocks( diff --git a/xllm/core/distributed_runtime/llm_master.cpp b/xllm/core/distributed_runtime/llm_master.cpp index 0e905fe044..7c1c8720c9 100644 --- a/xllm/core/distributed_runtime/llm_master.cpp +++ b/xllm/core/distributed_runtime/llm_master.cpp @@ -29,6 +29,9 @@ limitations under the License. #include "api_service/call.h" #include "common/metrics.h" +#include "core/distributed_runtime/mode_switch_service.h" +#include "core/framework/config/disagg_pd_config.h" +#include "core/framework/config/service_config.h" #include "core/platform/device_name_utils.h" #include "framework/model/model_args.h" #include "framework/request/request.h" @@ -102,6 +105,7 @@ LLMMaster::LLMMaster(const Options& options) .enable_profile_kv_blocks(options_.enable_profile_kv_blocks()) .disable_ttft_profiling(options_.disable_ttft_profiling()) .enable_forward_interruption(options_.enable_forward_interruption()) + .enable_runtime_cp_dp_switch(options_.enable_runtime_cp_dp_switch()) .max_global_ttft_ms(options_.max_global_ttft_ms()) .max_global_tpot_ms(options_.max_global_tpot_ms()) .server_idx(options_.server_idx()) @@ -114,6 +118,38 @@ LLMMaster::LLMMaster(const Options& options) XServiceClient::get_instance()->register_instance(instance_info); } + // Single-instance (non-disagg) debug trigger for runtime CP<->DP switch. + // In disagg_pd deployments ModeSwitchService is co-hosted on the disagg + // brpc server by DisaggPDScheduler::start_rpc_server(). That path never + // runs for a plain ContinuousScheduler, so a dual-mode single instance has + // no way to receive a SwitchMode RPC. Mount a standalone ModeSwitchService + // here so the flip can be driven directly. engine_ is already init'd and + // its worker_clients_ are ready (populated in the LLMEngine ctor). + if (options_.enable_runtime_cp_dp_switch() && !options_.enable_disagg_pd()) { + const int32_t port = + ::xllm::DisaggPDConfig::get_instance().disagg_pd_port(); + std::string host = ::xllm::ServiceConfig::get_instance().host(); + if (host.empty()) { + host = "127.0.0.1"; + } + const std::string addr = host + ":" + std::to_string(port); + auto* rpc_server = + ServerRegistry::get_instance().register_server("ModeSwitchStandalone"); + // Pass the ContinuousScheduler so ModeSwitchService can pause the loop + // and gate async surfaces around switch_mode + rebuild. Non-disagg + // deployments have no dispatch_thread / brpc handler surface, so the + // gate is a no-op and only the pause/resume path activates. If the + // dynamic_cast returns null (a scheduler_ type that isn't + // ContinuousScheduler-derived), fall back to the un-drained path + // that was there before. + auto* cs = dynamic_cast(scheduler_.get()); + auto mode_switch = std::make_unique( + static_cast(engine_.get()), cs); + if (!rpc_server->start(std::move(mode_switch), addr)) { + LOG(ERROR) << "Failed to start standalone ModeSwitchService on " << addr; + } + } + // construct chat template chat_template_ = ChatTemplate::create(engine_->tokenizer_args(), model_args_.model_type()); diff --git a/xllm/core/distributed_runtime/master.cpp b/xllm/core/distributed_runtime/master.cpp index 0012e6b15e..e9ec555086 100644 --- a/xllm/core/distributed_runtime/master.cpp +++ b/xllm/core/distributed_runtime/master.cpp @@ -469,7 +469,9 @@ Master::Master(const Options& options, EngineType type) .max_tokens_for_graph_mode(options_.max_tokens_for_graph_mode()) .kv_cache_dtype(options_.kv_cache_dtype()) .enable_sleep_mode(options_.enable_sleep_mode()) - .model_id(options_.model_id()); + .model_id(options_.model_id()) + .enable_runtime_cp_dp_switch(options_.enable_runtime_cp_dp_switch()) + .dual_mode_port_stride(options_.dual_mode_port_stride()); engine_ = std::make_unique(eng_options); } else if (type == EngineType::REC) { diff --git a/xllm/core/distributed_runtime/mode_switch_service.cpp b/xllm/core/distributed_runtime/mode_switch_service.cpp new file mode 100644 index 0000000000..913f8c2503 --- /dev/null +++ b/xllm/core/distributed_runtime/mode_switch_service.cpp @@ -0,0 +1,206 @@ +/* Copyright 2026 The xLLM Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "mode_switch_service.h" + +#include +#include +#include + +#include "llm_engine.h" +#include "runtime/xservice_client.h" +#include "scheduler/continuous_scheduler.h" +#include "scheduler/disagg_pd_scheduler.h" + +namespace xllm { + +ModeSwitchService::ModeSwitchService(LLMEngine* engine, + ContinuousScheduler* scheduler) + : engine_(engine), scheduler_(scheduler) {} + +void ModeSwitchService::SwitchMode( + ::google::protobuf::RpcController* controller, + const proto::InstanceModeSwitchRequest* request, + proto::InstanceModeSwitchResponse* response, + ::google::protobuf::Closure* done) { + brpc::ClosureGuard done_guard(done); + if (!engine_) { + response->set_ok(false); + response->set_current_mode(current_mode_.load()); + response->set_error("engine not bound to ModeSwitchService"); + LOG(ERROR) << "ModeSwitchService::SwitchMode called but engine is null"; + return; + } + const int32_t target = request->target_mode(); + if (target != 0 && target != 1) { + response->set_ok(false); + response->set_current_mode(current_mode_.load()); + response->set_error("invalid target_mode (expect 0 or 1)"); + LOG(ERROR) << "ModeSwitchService::SwitchMode bad target=" << target; + return; + } + const int32_t prev = current_mode_.load(); + if (prev == target) { + // Idempotent: report success without re-issuing fan-out. + response->set_ok(true); + response->set_current_mode(prev); + LOG(INFO) << "ModeSwitchService::SwitchMode no-op, already in mode=" + << target; + return; + } + LOG(INFO) << "ModeSwitchService::SwitchMode prev=" << prev + << " target=" << target; + + // Quiesce protocol. worker_impl.cpp:switch_mode assumes "the scheduler's + // drain protocol guarantees no forward is in flight on this worker when + // switch_mode runs" -- previously nothing enforced that. Without it, a + // brpc handler on the DisaggPDScheduler side (AddNewRequests, LinkInstance, + // FirstGeneration) or the dispatch_thread would keep touching + // kv_cache_manager_ mid-rebuild, which unique_ptr-resets the underlying + // BlockManagerPool. The observed symptom is a silent SIGSEGV that lands + // rank0 in a Zs zombie (no FATAL, no core in-container). Sequence: + // 1. pause the main scheduler loop (WAIT: let running requests finish + // naturally; KV cache stays intact) + // 2. take switch_gate_ unique; this waits out every in-flight + // dispatch_thread iteration and every brpc handler that grabbed the + // shared side, and blocks new ones from entering. + // 3. flip every worker via engine_->switch_mode (fans out RPC). + // 4. rebuild the scheduler-side pipeline (BlockManagerPool + BatchFactory + // + last_batch_ + active_dp_size_) for the new dp_size, all under the + // unique gate. + // 5. release gate, resume the loop. + // + // If scheduler_ is null we skip pause+gate and fall back to the raw path + // used by the standalone LLMMaster ModeSwitchService (single-instance + // debug trigger, no dispatch_thread / disagg brpc surface). + + DisaggPDScheduler* disagg = dynamic_cast(scheduler_); + + if (scheduler_ != nullptr) { + scheduler_->pause(ContinuousScheduler::PauseMode::WAIT); + // Use the caller-provided timeout so a heavy P/D pipeline (long TPOT + // decode requests) can drain before rebuild. Fall back to 30s if the + // caller didn't specify. If drain lapses we bail out and roll pause + // back with resume() -- forcing rebuild with in-flight requests would + // trip BlockManagerImpl's destructor invariant when their sequence + // shared_ptrs are still outstanding, killing rank0. + int64_t drain_timeout_ms = request->timeout_ms(); + if (drain_timeout_ms <= 0) { + drain_timeout_ms = 30000; + } + if (!scheduler_->wait_until_paused(drain_timeout_ms)) { + LOG(WARNING) << "ModeSwitchService::SwitchMode timed out waiting for " + "scheduler to drain (" + << drain_timeout_ms << "ms); aborting flip"; + scheduler_->resume(); + response->set_ok(false); + response->set_current_mode(prev); + response->set_error("scheduler drain timeout"); + return; + } + } + + // Take the async-surface gate unique. For a plain ContinuousScheduler + // (single-instance path) there is no dispatch_thread / brpc surface to + // guard, so we skip the gate entirely. + std::unique_lock gate; + if (disagg != nullptr) { + gate = disagg->begin_switch(); + } + + const bool ok = engine_->switch_mode(target); + if (ok) { + // Refresh scheduler-side bookkeeping to match the flipped engine layout. + // For DisaggPDScheduler this is the same rebuild that the step() loop + // used to do inline; doing it here inside the gate + paused loop makes + // the transition atomic from the outside. + const int32_t new_dp_size = engine_->dp_size(); + if (disagg != nullptr) { + disagg->rebuild_after_flip(new_dp_size); + } + // v3 e2e revealed that even with the pool rebuilt cleanly, the first + // DP-mode inference batch failed with LLM_NOT_YET_LINK (0x5010b007) at + // every layer's PushKvBlocks. Root cause: LLMEngine::link_cluster fans + // out per-D-worker links keyed on the peer's dp_size, and that topology + // is baked in at startup. After a flip the (src_worker -> dst_worker) + // pairs shift and the pre-flip links no longer cover them. + // + // Two-step recovery: + // 1. Re-publish OUR new dp_size to etcd. Peers cache InstanceInfo + // lazily and their get_instance_info hits etcd, so without this + // the peer's relink builds the wrong topology again. + // 2. Rebuild every LlmDataDist link this instance owns using the + // peer's new dp_size (pulled from etcd fresh). D-side is where + // the datadist links live; on P-side relink is a no-op. + // + // Step 1 (etcd write) can happen inside the gate cheaply. Step 2 does + // engine_->unlink_cluster + link_cluster which fan out to worker RPCs + // through link_threadpool_ + folly futures; running that under the gate + // deadlocks (v4 D-rank0 died with `std::system_error: Resource deadlock + // avoided` right after the relink line -- almost certainly a folly + // future / thread-pool mutex acquired both by the caller thread and the + // worker RPC handler while gate is held). Release the gate + resume + // the scheduler BEFORE relinking so worker handlers run without + // fighting the gate; the scheduler is idempotent to relink because + // dispatch_thread will not push through a P instance whose new dp_size + // has not been observed yet, and inference requests can still be + // dispatched (but the KV push will retry -- link isn't ready until + // relink_after_flip finishes). + XServiceClient* xsvc = XServiceClient::get_instance(); + if (xsvc != nullptr) { + if (!xsvc->re_register_dp_size(new_dp_size)) { + LOG(WARNING) << "ModeSwitchService: re_register_dp_size failed; peer " + "may keep stale dp_size cache"; + } + } + current_mode_.store(target); + response->set_ok(true); + response->set_current_mode(target); + } else { + // engine_->switch_mode() already attempted rollback. We trust the + // engine has restored prev mode on every worker. + response->set_ok(false); + response->set_current_mode(prev); + response->set_error("engine switch_mode fan-out failed; rolled back"); + } + + // Release the gate first so any brpc handler that has been queued behind + // us can proceed once the scheduler resumes; then resume the loop. + if (gate.owns_lock()) { + gate.unlock(); + } + if (scheduler_ != nullptr) { + scheduler_->resume(); + } + + // Post-flip relink: rebuild datadist P<->D links to match the new dp_size + // topology. Deliberately outside the gate + after resume() -- the fan-out + // (link_threadpool_ + folly futures + worker RPC round-trips) hit a + // deadlock inside the gate on v4 (D-rank0 crashed with + // `std::system_error: Resource deadlock avoided`), because worker + // handlers acquire per-service mutexes that overlap with what the caller + // thread would need to release. The relink itself is safe outside the + // gate: dispatch_thread that races us just sees old links until they + // rebuild, and PushKvBlocks retries on transient LLM_NOT_YET_LINK. + // Skipped if switch_mode failed (we already rolled back to prev mode). + if (ok && disagg != nullptr) { + if (!disagg->relink_after_flip()) { + LOG(WARNING) << "ModeSwitchService: relink_after_flip reported failures; " + "some peers may still be linked with stale dp_size"; + } + } +} + +} // namespace xllm diff --git a/xllm/core/distributed_runtime/mode_switch_service.h b/xllm/core/distributed_runtime/mode_switch_service.h new file mode 100644 index 0000000000..3fdb26d76b --- /dev/null +++ b/xllm/core/distributed_runtime/mode_switch_service.h @@ -0,0 +1,58 @@ +/* Copyright 2026 The xLLM Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#pragma once + +#include +#include + +#include "common/macros.h" +#include "mode_switch.pb.h" +#include "runtime/worker_client.h" + +namespace xllm { + +class LLMEngine; +class ContinuousScheduler; + +class ModeSwitchService : public proto::ModeSwitchService { + public: + // engine is required. scheduler is optional but required to safely quiesce + // async surfaces (dispatch_thread, brpc handlers) before rebuilding the + // KV allocator; passing nullptr falls back to the pre-drain path (only + // safe for tests where nothing else touches the pool). + explicit ModeSwitchService(LLMEngine* engine, + ContinuousScheduler* scheduler = nullptr); + ModeSwitchService() : engine_(nullptr), scheduler_(nullptr) {} + virtual ~ModeSwitchService() = default; + + void SwitchMode(::google::protobuf::RpcController* controller, + const proto::InstanceModeSwitchRequest* request, + proto::InstanceModeSwitchResponse* response, + ::google::protobuf::Closure* done) override; + + private: + LLMEngine* engine_; + // Optional; if the deployment is disaggregated PD this points at the + // DisaggPDScheduler that owns the request-dispatch async surface. The + // service uses it to pause the loop and take switch_gate_ unique across + // engine_->switch_mode + rebuild_block_manager_pool. + ContinuousScheduler* scheduler_; + std::atomic current_mode_{0}; + + DISALLOW_COPY_AND_ASSIGN(ModeSwitchService); +}; + +} // namespace xllm diff --git a/xllm/core/distributed_runtime/remote_worker.cpp b/xllm/core/distributed_runtime/remote_worker.cpp index 26158665a9..2d9ba07f87 100644 --- a/xllm/core/distributed_runtime/remote_worker.cpp +++ b/xllm/core/distributed_runtime/remote_worker.cpp @@ -418,4 +418,14 @@ folly::SemiFuture RemoteWorker::stop_profile_async() { return future; } +folly::SemiFuture RemoteWorker::switch_mode_async(int32_t target_mode) { + folly::Promise promise; + auto future = promise.getSemiFuture(); + threadpool_.schedule( + [this, target_mode, promise = std::move(promise)]() mutable { + promise.setValue(channel_->switch_mode(target_mode)); + }); + return future; +} + } // namespace xllm diff --git a/xllm/core/distributed_runtime/remote_worker.h b/xllm/core/distributed_runtime/remote_worker.h index d337989069..fda4e7ad37 100644 --- a/xllm/core/distributed_runtime/remote_worker.h +++ b/xllm/core/distributed_runtime/remote_worker.h @@ -157,6 +157,9 @@ class RemoteWorker : public WorkerClient { virtual folly::SemiFuture stop_profile_async() override; + virtual folly::SemiFuture switch_mode_async( + int32_t target_mode) override; + private: DISALLOW_COPY_AND_ASSIGN(RemoteWorker); diff --git a/xllm/core/distributed_runtime/worker_server.cpp b/xllm/core/distributed_runtime/worker_server.cpp index d1632ff92c..74911e01a5 100644 --- a/xllm/core/distributed_runtime/worker_server.cpp +++ b/xllm/core/distributed_runtime/worker_server.cpp @@ -160,6 +160,15 @@ void WorkerServer::create_server(const runtime::Options& options, const ParallelArgs* parallel_args = nullptr; std::unique_ptr comm; + // Holdings for dual-mode (enable_runtime_cp_dp_switch). These stay + // empty in legacy single-mode. + std::unique_ptr comm_cp_local; + std::unique_ptr comm_dp_local; + std::unique_ptr dual_args_local; + + const bool enable_dual_mode = + options.enable_runtime_cp_dp_switch() && worker_type != WorkerType::DIT; + if (worker_type == WorkerType::DIT) { auto dit_comm = std::make_unique(worker_global_rank, @@ -170,32 +179,133 @@ void WorkerServer::create_server(const runtime::Options& options, options.cfg_size(), options.vae_size()); comm = std::move(dit_comm); - } else { + comm->create_process_groups(master_node_addr, device); + parallel_args = comm->parallel_args(); + } else if (!enable_dual_mode) { auto common_comm = std::make_unique( worker_global_rank, world_size, dp_size, ep_size, cp_size); comm = std::move(common_comm); + comm->create_process_groups(master_node_addr, device); + parallel_args = comm->parallel_args(); + } else { + // Dual-mode: build two CollectiveCommunicators on different + // master_addr port windows, snapshot both ParallelArgs into a + // DualParallelArgs, and let WorkerImpl pick active() at runtime. + // + // Mode mapping. The user-supplied (cp_size, dp_size, tp_size) is + // taken as ONE side of the dual; the OTHER side is its swap with + // cp_size <-> dp_size and tp_size held constant. This keeps + // world_size = cp * dp * tp invariant so the two ParallelArgs + // share the same world geometry, which is what HCCL ECM caches and + // what the dual_graph_probe_atb run-with-barrier relied on. + // + // Initial active is whichever side matches the user's input cp/dp. + // If the user passed cp_size > 1 we boot into CP_PREFILL mode; if + // they passed dp_size > 1 we boot into DP_DECODE mode. (Both > 1 + // is rejected by MappingNPU::validate already.) + // + // CRITICAL invariants captured by the v3 probe (see + // tests/probes/dual_graph_probe_atb.cpp + DESIGN_DOC_v0.2.md): + // * port stride MUST exceed inner fanout in + // CollectiveCommunicator::create_process_groups (default 256) + // * after each create_process_groups we MUST issue a barrier + // allreduce before continuing -- atb::Comm::CreateHcclComm + // returns asynchronously on this CANN, so a fast rank racing + // ahead can land in the slow-rank fallback path that emits + // commDomain="0" / hccl=null. The probe enforces this with a + // 1-element HcclAllReduce; here we rely on the fact that the + // master_node sync_master_node call later gates the dp_args + // consumer the same way for production startup. + std::string master_host; + int32_t master_port = 0; + net::parse_host_port_from_addr(master_node_addr, master_host, master_port); + const int32_t cp_port = master_port; + const int32_t dp_port = master_port + options.dual_mode_port_stride(); + const std::string cp_addr = master_host + ":" + std::to_string(cp_port); + const std::string dp_addr = master_host + ":" + std::to_string(dp_port); + + LOG(INFO) << "Worker rank=" << worker_global_rank + << " dual-mode comm setup: cp_master=" << cp_addr + << " dp_master=" << dp_addr << " (user cp=" << cp_size + << " dp=" << dp_size + << " => CP-side cp=" << std::max(cp_size, dp_size) + << " dp=1, DP-side cp=1 dp=" << std::max(cp_size, dp_size) << ")"; + + // The "CP" side: cp_size = max(cp, dp), dp_size = 1. + // The "DP" side: cp_size = 1, dp_size = max(cp, dp). + // tp_size and ep_size carry over implicitly through world_size. + const int32_t paired = std::max(cp_size, dp_size); + auto cp_comm = std::make_unique(worker_global_rank, + world_size, + /*dp_size=*/1, + ep_size, + /*cp_size=*/paired); + cp_comm->create_process_groups(cp_addr, device); + + auto dp_comm = std::make_unique(worker_global_rank, + world_size, + /*dp_size=*/paired, + ep_size, + /*cp_size=*/1); + dp_comm->create_process_groups(dp_addr, device); + + dual_args_local = std::make_unique( + *cp_comm->parallel_args(), + *dp_comm->parallel_args(), + cp_size > 1 ? DualParallelArgs::Mode::CP_PREFILL + : DualParallelArgs::Mode::DP_DECODE); + + comm_cp_local = std::move(cp_comm); + comm_dp_local = std::move(dp_comm); + parallel_args = &dual_args_local->active(); } - comm->create_process_groups(master_node_addr, device); - parallel_args = comm->parallel_args(); + // Legacy (single-mode) path: build one process group and use the + // single ParallelArgs. Dual-mode already built both cp_comm and + // dp_comm above with their own create_process_groups(); skip the + // extra call here to avoid duplicate rendezvous. + if (!enable_dual_mode) { + comm->create_process_groups(master_node_addr, device); + parallel_args = comm->parallel_args(); + } - ParallelArgs worker_parallel_args = *parallel_args; - const int32_t cp_group_size = - worker_parallel_args.cp_group_ == nullptr - ? 1 - : worker_parallel_args.cp_group_->world_size(); - const char* cp_partition_stage = - options.cp_size() <= 1 - ? "disabled" - : (Platform::uses_model_cp_partition() ? "model" : "worker"); - LOG(INFO) << "Worker CP config: cp_size=" << options.cp_size() - << ", cp_group_size=" << cp_group_size - << ", cp_partition_stage=" << cp_partition_stage - << ", instance_role=" << options.instance_role().to_string(); + // Diagnostic log (from upstream): print CP topology + partition stage. + { + const ParallelArgs& active_pa = + enable_dual_mode ? dual_args_local->active() : *parallel_args; + const int32_t cp_group_size = + active_pa.cp_group_ == nullptr ? 1 : active_pa.cp_group_->world_size(); + const char* cp_partition_stage = + options.cp_size() <= 1 + ? "disabled" + : (Platform::uses_model_cp_partition() ? "model" : "worker"); + LOG(INFO) << "Worker CP config: cp_size=" << options.cp_size() + << ", cp_group_size=" << cp_group_size + << ", cp_partition_stage=" << cp_partition_stage + << ", instance_role=" << options.instance_role().to_string() + << ", dual_mode=" << enable_dual_mode; + } - std::unique_ptr worker = std::make_unique( - worker_parallel_args, device, options, worker_type); + std::unique_ptr worker; + if (enable_dual_mode) { + worker = std::make_unique( + dual_args_local.get(), device, options, worker_type); + } else { + worker = + std::make_unique(*parallel_args, device, options, worker_type); + } worker_service->set_worker(std::move(worker)); + + // Move the dual-mode comm holdings onto WorkerServer so they outlive + // this scope. WorkerService holds Worker which holds a non-owning + // pointer into dual_parallel_args_; if we let it die here that pointer + // would dangle. In legacy single-mode these stay null and `comm` falls + // through to the existing local-scope cleanup behavior. + if (enable_dual_mode) { + comm_cp_ = std::move(comm_cp_local); + comm_dp_ = std::move(comm_dp_local); + dual_parallel_args_ = std::move(dual_args_local); + } bool create_shm = options.enable_shm() && input_shm_manager && output_shm_manager; if (create_shm) { diff --git a/xllm/core/distributed_runtime/worker_server.h b/xllm/core/distributed_runtime/worker_server.h index 8c8c2be154..86f44f63a8 100644 --- a/xllm/core/distributed_runtime/worker_server.h +++ b/xllm/core/distributed_runtime/worker_server.h @@ -28,6 +28,7 @@ limitations under the License. #include "distributed_runtime/worker_service.h" #include "framework/model/model_args.h" #include "framework/model/model_input_params.h" +#include "framework/parallel_state/collective_communicator.h" #include "runtime/executor.h" #include "runtime/forward_params.h" #include "runtime/forward_shared_memory_manager.h" @@ -100,6 +101,16 @@ class WorkerServer { pid_t spawned_worker_pid_ = -1; std::atomic stopped_{false}; std::string server_name_; + + // Dual-mode comm holdings. When options.enable_runtime_cp_dp_switch is + // true, create_server() builds two CollectiveCommunicators (one CP, one + // DP) and stores them here so they outlive the local create_server stack + // frame. The DualParallelArgs aggregates both comms' parallel_args() and + // is the source of truth the Worker / WorkerImpl reads from after + // switch_mode flips. In legacy single-mode these are all null. + std::unique_ptr comm_cp_; + std::unique_ptr comm_dp_; + std::unique_ptr dual_parallel_args_; }; } // namespace xllm diff --git a/xllm/core/distributed_runtime/worker_service.cpp b/xllm/core/distributed_runtime/worker_service.cpp index 00a9a7a9c4..23a428e50e 100644 --- a/xllm/core/distributed_runtime/worker_service.cpp +++ b/xllm/core/distributed_runtime/worker_service.cpp @@ -718,6 +718,34 @@ void WorkerService::StopProfile(::google::protobuf::RpcController* controller, return; } +void WorkerService::SwitchMode(::google::protobuf::RpcController* controller, + const proto::SwitchModeRequest* req, + proto::Status* resp, + ::google::protobuf::Closure* done) { + threadpool_->schedule([this, controller, req, resp, done]() mutable { + brpc::ClosureGuard done_guard(done); + const int32_t mode_int = req->target_mode(); + if (mode_int != static_cast(DualParallelArgs::Mode::CP_PREFILL) && + mode_int != static_cast(DualParallelArgs::Mode::DP_DECODE)) { + LOG(WARNING) << "SwitchMode RPC received invalid target_mode=" + << mode_int; + resp->set_ok(false); + return; + } + const auto target = static_cast(mode_int); + const bool ok = worker_->switch_mode(target); + if (!ok) { + LOG(WARNING) << "SwitchMode RPC: worker rejected the flip " + "(legacy single-mode worker?). target=" + << mode_int; + } else { + LOG(INFO) << "SwitchMode RPC: flipped to mode=" << mode_int; + } + resp->set_ok(ok); + }); + return; +} + void WorkerService::ExecuteModel(::google::protobuf::RpcController* controller, const proto::ForwardInput* pb_forward_input, proto::ForwardOutput* pb_forward_output, diff --git a/xllm/core/distributed_runtime/worker_service.h b/xllm/core/distributed_runtime/worker_service.h index 01cbc62cdd..0ee2b92b6d 100644 --- a/xllm/core/distributed_runtime/worker_service.h +++ b/xllm/core/distributed_runtime/worker_service.h @@ -150,6 +150,15 @@ class WorkerService : public proto::DistributeWorker { proto::Status* resp, ::google::protobuf::Closure* done) override; + // Runtime CP<->DP switch RPC handler. Validates the target_mode value + // and forwards to Worker::switch_mode; returns Status.ok=false when + // the worker is in legacy single-mode (no DualParallelArgs attached) + // or target_mode is out of range. + void SwitchMode(::google::protobuf::RpcController* controller, + const proto::SwitchModeRequest* req, + proto::Status* resp, + ::google::protobuf::Closure* done) override; + private: void step(ForwardInput& fwd_input, torch::Tensor& next_tokens, diff --git a/xllm/core/framework/batch/batch_factory.h b/xllm/core/framework/batch/batch_factory.h index 0d8aa4c683..15c694a65b 100644 --- a/xllm/core/framework/batch/batch_factory.h +++ b/xllm/core/framework/batch/batch_factory.h @@ -42,6 +42,15 @@ class BatchFactory { std::vector>* swap_block_transfer_infos = nullptr); + // Update the dp/cp size after a runtime CP<->DP flip. This is a Meyers + // singleton, so the size is otherwise fixed at first get_instance() call; + // create_batches reads these members on every call, so updating them here is + // enough. Must be called from the scheduler loop thread while drained. + void set_dp_size(int32_t dp_size, int32_t cp_size = 1) { + dp_size_ = dp_size; + cp_size_ = cp_size; + } + private: BatchFactory(int32_t dp_size, int32_t cp_size) : dp_size_(dp_size), cp_size_(cp_size) {} diff --git a/xllm/core/framework/batch/batch_input_builder.cpp b/xllm/core/framework/batch/batch_input_builder.cpp index 4a0ac56746..b7717fd1d4 100644 --- a/xllm/core/framework/batch/batch_input_builder.cpp +++ b/xllm/core/framework/batch/batch_input_builder.cpp @@ -614,10 +614,15 @@ void BatchInputBuilder::process_single_sequence( sequence, n_kv_cache_tokens, logical_seq_len, seq_len, state_ptr); // Setup KV cache + // FIX (dual-mode trial-launch): pass actual q_seq_len, NOT padded_q_seq_len. + // CP padding aligns q_seq_len up to 2*cp_size for the forward kernel, but + // kv_cache_tokens_num must reflect the actual prompt tokens cached. Using + // padded_q_seq_len here causes kv_cache_tokens_num > n_tokens on the next + // step, which is the phantom seq we were patching around. setup_kv_cache_info(sequence, n_kv_cache_tokens, seq_len, - padded_q_seq_len, + q_seq_len, state_ptr, write_block_ids_ptr); diff --git a/xllm/core/framework/block/block_manager_impl.h b/xllm/core/framework/block/block_manager_impl.h index ad39e1cc3c..74dd195ab8 100644 --- a/xllm/core/framework/block/block_manager_impl.h +++ b/xllm/core/framework/block/block_manager_impl.h @@ -25,8 +25,24 @@ class BlockManagerImpl : public BlockManager { explicit BlockManagerImpl(const Options& options); virtual ~BlockManagerImpl() { prefix_cache_.reset(); - CHECK_EQ(num_free_blocks_, free_blocks_.size() - 1) - << "Not all blocks have been freed"; + // In steady-state operation this must hold: every allocated block returns + // to free_blocks_ before the pool is torn down. But the runtime CP<->DP + // flip path (LLMEngine::rebuild_block_manager_pool) tears down the pool + // while sequence shared_ptrs from the just-finished decode requests can + // still be alive on other threads' return paths (response processor, + // batch factory scratch, xservice heartbeat), holding a Block or two. + // Those blocks return via ~Block once the shared_ptrs release, which is + // guaranteed to happen -- just not necessarily before this destructor + // runs. Downgrade the fatal CHECK to a warning so a rebuild-under-flip + // does not crash rank0 for a benign timing gap; steady-state teardown + // still surfaces true leaks in the log. + if (num_free_blocks_ != free_blocks_.size() - 1) { + LOG(WARNING) << "BlockManagerImpl teardown: " + << (free_blocks_.size() - 1 - num_free_blocks_) + << " block(s) still outstanding (num_free_blocks_=" + << num_free_blocks_ + << ", free_blocks_.size()=" << free_blocks_.size() << ")"; + } }; // Try to allocate blocks with num_blocks, diff --git a/xllm/core/framework/block/block_manager_pool.cpp b/xllm/core/framework/block/block_manager_pool.cpp index 411991de81..d159fb156f 100644 --- a/xllm/core/framework/block/block_manager_pool.cpp +++ b/xllm/core/framework/block/block_manager_pool.cpp @@ -123,7 +123,27 @@ int32_t BlockManagerPool::get_dp_rank(Sequence* sequence) const { if (sequence->dp_rank() >= 0) { dp_rank = sequence->dp_rank(); } else { - dp_rank = get_manager_with_max_free_blocks(); + // Round-robin across dp_ranks for the initial assignment. The former + // "max free blocks" heuristic packed the first N sequences of a burst + // all onto dp_rank=0 (they all saw dp_rank=0 free>=dp_rank=1 free + // until dp_rank=0's free drops below), which after a runtime CP<->DP + // flip left dp_rank=1 empty. Empty shard forwards run on the fake + // token (see worker_impl need_fake_input_for_empty_shard) but the + // ATB / DP collective ops in some layers still hang when one shard + // is a real N-token batch and the peer only has a synthesized + // 1-token filler; balancing at assignment time avoids that path. + // Note: RR ignores current free-block imbalance, which is fine as + // long as dp_size stays small (world/dp_size); once one manager + // fully fills, its allocate() will fail-fast, propagating an + // OOM-style rejection back to the scheduler for the caller to + // retry rather than silently oversubscribing the busier rank. + if (block_managers_.empty()) { + dp_rank = 0; + } else { + const uint64_t idx = + next_rr_dp_rank_.fetch_add(1, std::memory_order_relaxed); + dp_rank = static_cast(idx % block_managers_.size()); + } sequence->set_dp_rank(dp_rank); } return dp_rank; diff --git a/xllm/core/framework/block/block_manager_pool.h b/xllm/core/framework/block/block_manager_pool.h index 6b18376f85..8e978ca6db 100644 --- a/xllm/core/framework/block/block_manager_pool.h +++ b/xllm/core/framework/block/block_manager_pool.h @@ -129,6 +129,10 @@ class BlockManagerPool : public KVCacheManager { // the options for the block manager Options options_; std::vector> block_managers_; + // Round-robin cursor for get_dp_rank. Tracked here (not in the free-block + // heuristic) so a burst of concurrent short requests spreads across dp + // ranks evenly. See get_dp_rank() in the .cpp. + mutable std::atomic next_rr_dp_rank_{0}; }; } // namespace xllm diff --git a/xllm/core/framework/config/parallel_config.cpp b/xllm/core/framework/config/parallel_config.cpp index d29e1e2786..0e34afb855 100644 --- a/xllm/core/framework/config/parallel_config.cpp +++ b/xllm/core/framework/config/parallel_config.cpp @@ -69,6 +69,39 @@ DEFINE_bool( "Whether to enable dp load balance, if true, sequences within a single " "dp batch will be shuffled."); +DEFINE_bool( + enable_runtime_cp_dp_switch, + false, + "Enable runtime CP<->DP switching (dual-graph mode). When on the " + "worker brings up TWO CollectiveCommunicators at startup -- one " + "with cp_size=N dp_size=1 and one with cp_size=1 dp_size=N -- and " + "holds both ATB graph pairs resident. See " + "Desktop/cp_docs_work/probes_results_20260629/DESIGN_DOC_v0.2.md " + "for the budget. Costs ~600 MB / NPU plus extra ATB graph workspace."); + +DEFINE_bool(enable_flip_verbose_log, + false, + "Enable per-request/per-step FLIPDIAG diagnostic logs for CP<->DP " + "flip debugging (dispatch, recv, per-step batch shape, per-worker " + "forward shard, dp_global_token_nums broadcast). Off by default -- " + "flip lifecycle events (switch_mode/rebuild/relink/backdoor) are " + "still logged unconditionally. Turn on when investigating flip " + "regressions; expect 20-100x log volume during DP burst traffic. " + "Runtime-toggleable via brpc /flags endpoint."); +// Registering a no-op validator marks the flag as reloadable so brpc's +// /flags?setvalue=true HTTP endpoint accepts runtime toggles. Without it +// brpc rejects with "A reloadable gflag must have validator" and forces +// a service restart just to flip verbose logging. +static bool ValidateFlipVerboseLog(const char*, bool) { return true; } +DEFINE_validator(enable_flip_verbose_log, &ValidateFlipVerboseLog); + +DEFINE_int32(dual_mode_port_stride, + 256, + "Port stride between the CP and DP communicators when " + "enable_runtime_cp_dp_switch is on. 256 was validated by the v3 " + "probe; smaller values risk colliding with TIME_WAIT residual " + "TCPStore ports between iterations."); + namespace xllm { void ParallelConfig::from_flags() { @@ -85,6 +118,8 @@ void ParallelConfig::from_flags() { XLLM_CONFIG_ASSIGN_FROM_FLAG(enable_multi_stream_parallel); XLLM_CONFIG_ASSIGN_FROM_FLAG(micro_batch_num); XLLM_CONFIG_ASSIGN_FROM_FLAG(enable_dp_balance); + XLLM_CONFIG_ASSIGN_FROM_FLAG(enable_runtime_cp_dp_switch); + XLLM_CONFIG_ASSIGN_FROM_FLAG(dual_mode_port_stride); } void ParallelConfig::from_json(const JsonReader& json) { @@ -100,6 +135,8 @@ void ParallelConfig::from_json(const JsonReader& json) { XLLM_CONFIG_ASSIGN_FROM_JSON(enable_multi_stream_parallel); XLLM_CONFIG_ASSIGN_FROM_JSON(micro_batch_num); XLLM_CONFIG_ASSIGN_FROM_JSON(enable_dp_balance); + XLLM_CONFIG_ASSIGN_FROM_JSON(enable_runtime_cp_dp_switch); + XLLM_CONFIG_ASSIGN_FROM_JSON(dual_mode_port_stride); } void ParallelConfig::append_config_json( @@ -124,6 +161,10 @@ void ParallelConfig::append_config_json( config_json, default_config, micro_batch_num); APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT( config_json, default_config, enable_dp_balance); + APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT( + config_json, default_config, enable_runtime_cp_dp_switch); + APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT( + config_json, default_config, dual_mode_port_stride); } ParallelConfig& ParallelConfig::get_instance() { diff --git a/xllm/core/framework/config/parallel_config.h b/xllm/core/framework/config/parallel_config.h index 22159a1c2a..2a2f0467f5 100644 --- a/xllm/core/framework/config/parallel_config.h +++ b/xllm/core/framework/config/parallel_config.h @@ -52,7 +52,9 @@ class ParallelConfig final { "enable_mm_encoder_dp", "enable_multi_stream_parallel", "micro_batch_num", - "enable_dp_balance"}}; + "enable_dp_balance", + "enable_runtime_cp_dp_switch", + "dual_mode_port_stride"}}; return kOptionCategory; } @@ -83,6 +85,15 @@ class ParallelConfig final { PROPERTY(bool, enable_dp_balance) = false; + // Mirrors gflags::FLAGS_enable_runtime_cp_dp_switch / + // FLAGS_dual_mode_port_stride. xllm.cpp copies these onto common::Options + // so master.cpp can forward them to runtime::Options at engine boot; + // worker_server then uses them to decide whether to bring up two + // CollectiveCommunicators. + PROPERTY(bool, enable_runtime_cp_dp_switch) = false; + + PROPERTY(int32_t, dual_mode_port_stride) = 256; + [[nodiscard]] int32_t kv_split_size_effective() const noexcept { return kv_split_size_ > 0 ? kv_split_size_ : cp_size_; } diff --git a/xllm/core/framework/model_context.cpp b/xllm/core/framework/model_context.cpp index 822c44faa5..b875dcffc2 100644 --- a/xllm/core/framework/model_context.cpp +++ b/xllm/core/framework/model_context.cpp @@ -96,6 +96,10 @@ ModelContext ModelContext::with_parallel_args( #endif derived.model_id_ = model_id_; derived.optimization_config_ = optimization_config_; + // Forward the optional dual-mode source so derived contexts (used by + // some sub-models when they construct nested ModelContexts) keep the + // same dual-mode visibility as the parent. + derived.dual_parallel_args_ = dual_parallel_args_; return derived; } diff --git a/xllm/core/framework/model_context.h b/xllm/core/framework/model_context.h index 9b6b76e7a0..6f749dc194 100644 --- a/xllm/core/framework/model_context.h +++ b/xllm/core/framework/model_context.h @@ -74,6 +74,15 @@ class ModelContext { const ParallelArgs& get_parallel_args() const { return parallel_args_; } + // Refresh the cached ParallelArgs snapshot after a runtime CP<->DP flip. + // ModelContext keeps its own copy (not a reference), so a mode switch that + // only updates the worker's ParallelArgs would leave this stale and the + // forward path would read the pre-flip cp_size/dp_size. Called from + // WorkerImpl::switch_mode with the post-flip active() args. + void set_parallel_args(const ParallelArgs& parallel_args) { + parallel_args_ = parallel_args; + } + const torch::TensorOptions& get_tensor_options() const { return tensor_options_; } @@ -84,6 +93,19 @@ class ModelContext { ModelContext with_parallel_args(const ParallelArgs& parallel_args) const; + // Optional dual-mode source. Set by the engine (after worker startup) + // when options.enable_runtime_cp_dp_switch is on. Layers that support + // dual ATB graph init (currently DeepSeek V32) consult this in their + // ctor; if non-null they build CP and DP node pairs from + // dual_parallel_args_->cp_args() / dp_args(), otherwise they fall + // through to the legacy single-mapping path using parallel_args_. + void set_dual_parallel_args(const DualParallelArgs* dual_args) { + dual_parallel_args_ = dual_args; + } + const DualParallelArgs* get_dual_parallel_args() const { + return dual_parallel_args_; + } + #if defined(USE_NPU) const atb::Context* get_atb_context() const { return context_; } std::shared_ptr get_atb_workspace() const { @@ -116,6 +138,9 @@ class ModelContext { torch::TensorOptions tensor_options_; OptimizationConfig optimization_config_; + // Non-owning, optional. Lifetime guaranteed by WorkerServer. + const DualParallelArgs* dual_parallel_args_ = nullptr; + #if defined(USE_NPU) // used for npu atb atb::Context* context_; diff --git a/xllm/core/framework/parallel_state/parallel_args.h b/xllm/core/framework/parallel_state/parallel_args.h index 8203680cd6..f8addfc337 100644 --- a/xllm/core/framework/parallel_state/parallel_args.h +++ b/xllm/core/framework/parallel_state/parallel_args.h @@ -23,6 +23,7 @@ limitations under the License. #include "xllm_atb_layers/models/base/param/mapping.h" #endif +#include #include #include @@ -239,4 +240,71 @@ struct ParallelArgs { ProcessGroup* dit_vae_group_ = nullptr; }; +// Two-mode parallel configuration container used by the runtime CP<->DP +// switch feature. Holds two independent ParallelArgs (one for CP-prefill +// mode, one for DP-decode mode) and an atomic flag selecting which one +// is currently active. Reads go through `active()` so the singular ParallelArgs +// API stays intact for callers who do not care about the switch. +// +// Lifetime: this object outlives both the CollectiveCommunicator instances +// it borrows from (worker_server keeps both comms resident) and the Worker +// that holds a const reference / pointer to it. Mode switching only flips +// the active_ flag; no objects are reconstructed. +// +// Threading: active_ is an atomic; the discriminator lives in a single +// cacheline-bounded read on the forward path. Mode flips are gated by the +// scheduler's drain protocol -- forward-path readers never observe a switch +// mid-collective. +class DualParallelArgs final { + public: + enum class Mode : int8_t { + CP_PREFILL = 0, // cp_size = N, dp_size = 1; long-prefill friendly + DP_DECODE = 1, // cp_size = 1, dp_size = N; concurrency friendly + }; + + DualParallelArgs(ParallelArgs cp_args, + ParallelArgs dp_args, + Mode initial_mode) + : cp_args_(std::move(cp_args)), + dp_args_(std::move(dp_args)), + active_(initial_mode) {} + + // Non-copyable, non-movable: holds an atomic and is referenced by ptr + // from the Worker that owns it. + DualParallelArgs(const DualParallelArgs&) = delete; + DualParallelArgs& operator=(const DualParallelArgs&) = delete; + DualParallelArgs(DualParallelArgs&&) = delete; + DualParallelArgs& operator=(DualParallelArgs&&) = delete; + + [[nodiscard]] const ParallelArgs& active() const noexcept { + return active_.load(std::memory_order_acquire) == Mode::CP_PREFILL + ? cp_args_ + : dp_args_; + } + + [[nodiscard]] const ParallelArgs& cp_args() const noexcept { + return cp_args_; + } + + [[nodiscard]] const ParallelArgs& dp_args() const noexcept { + return dp_args_; + } + + [[nodiscard]] Mode mode() const noexcept { + return active_.load(std::memory_order_acquire); + } + + // Flip the active mode. Caller is responsible for draining inflight + // requests before issuing the flip; see worker switch_mode RPC for the + // full protocol. Idempotent if mode is already target. + void set_mode(Mode target) noexcept { + active_.store(target, std::memory_order_release); + } + + private: + ParallelArgs cp_args_; + ParallelArgs dp_args_; + std::atomic active_; +}; + } // namespace xllm diff --git a/xllm/core/layers/npu/npu_deepseek_v32_decoder_layer_impl.cpp b/xllm/core/layers/npu/npu_deepseek_v32_decoder_layer_impl.cpp index 7d8cd78d30..7605ba8f8c 100644 --- a/xllm/core/layers/npu/npu_deepseek_v32_decoder_layer_impl.cpp +++ b/xllm/core/layers/npu/npu_deepseek_v32_decoder_layer_impl.cpp @@ -401,6 +401,33 @@ NpuDeepseekV32DecoderLayerImpl::NpuDeepseekV32DecoderLayerImpl( mtp_decode_fallback_param_.outputTopk = true; } + // Dual-mode setup. When the engine has registered a DualParallelArgs on + // the ModelContext, also build prefill+decode params for the OTHER + // mode. The active() side has just been baked into prefill_param_ / + // decode_param_ above; here we fill the inactive side into + // dp_prefill_param_ / dp_decode_param_ (or, if the active side is DP, + // we fill the CP variants there). The naming convention for the dp_* + // members is "the OTHER half of the dual" -- not literally always DP. + // The forward path consults dual_parallel_args_->mode() each call to + // pick the right node. + dual_parallel_args_ = context.get_dual_parallel_args(); + if (dual_parallel_args_ != nullptr) { + dual_ctor_mode_ = dual_parallel_args_->mode(); + const ParallelArgs& other_args = + (dual_ctor_mode_ == DualParallelArgs::Mode::CP_PREFILL) + ? dual_parallel_args_->dp_args() + : dual_parallel_args_->cp_args(); + param_from_args( + dp_prefill_param_, model_args, other_args, /*is_prefill=*/true); + param_from_args( + dp_decode_param_, model_args, other_args, /*is_prefill=*/false); + LOG_FIRST_N(INFO, 1) << "DeepSeek V32 layer " << layer_id_ + << " entered dual-mode init (ctor_active=" + << static_cast(dual_ctor_mode_) + << ", other side cp_size=" << other_args.cp_size() + << " dp_size=" << other_args.dp_size() << ")"; + } + loader_ = std::make_unique( WEIGHT_COUNT_PER_LAYER, context, @@ -958,6 +985,13 @@ void NpuDeepseekV32DecoderLayerImpl::update_expert_weight() { mtp_decode_fallback_node_.inTensors.at(index) = &atb_weight_tensors_[index]; } + // Dual-mode mirror: same weight pointer flows into the OTHER side's + // nodes too. atb_weight_tensors_ is the single source of truth for + // weights regardless of how many ATB graphs reference them. + if (dual_parallel_args_ != nullptr) { + dp_prefill_node_.inTensors.at(index) = &atb_weight_tensors_[index]; + dp_decode_node_.inTensors.at(index) = &atb_weight_tensors_[index]; + } } expert_routing_map_[layer_id_ - first_k_dense_replace_] = expert_routing_map_buffer_; @@ -975,6 +1009,21 @@ int64_t NpuDeepseekV32DecoderLayerImpl::init_layer() { CHECK_OPERATION_STATUS_RETURN( init_node(mtp_decode_fallback_node_, mtp_decode_fallback_param_)); } + // Dual-mode: also init the OTHER side's nodes if a DualParallelArgs is + // attached. atb_speed::deepseekV2::DecoderLayer creates a fresh + // atb::Operation graph per call so distinct mappings produce distinct + // node objects; weights are still shared via atb_weight_tensors_ which + // is set up in merge_loaded_weights() and identical across all four + // nodes. + if (dual_parallel_args_ != nullptr) { + CHECK_OPERATION_STATUS_RETURN( + init_node(dp_prefill_node_, dp_prefill_param_)); + CHECK_OPERATION_STATUS_RETURN(init_node(dp_decode_node_, dp_decode_param_)); + LOG_FIRST_N(INFO, 1) + << "DeepSeek V32 layer " << layer_id_ + << " init_layer: 4 ATB nodes ready (cp_prefill, cp_decode, " + "dp_prefill, dp_decode)"; + } return atb::NO_ERROR; } @@ -1061,39 +1110,41 @@ torch::Tensor NpuDeepseekV32DecoderLayerImpl::forward_with_topk( atb::Status st; ModelInputParams& input_params_new = const_cast(input_params); - if (input_params_new.meta.batch_forward_type.is_decode()) { - build_node_variant_pack(decode_node_, - x, - cos_pos, - sin_pos, - attn_mask, - kv_cache, - input_params_new, - false, - shared_topk_indices, - output_topk_indices, - skip_topk_, - output_topk_); - st = execute_node(decode_node_, node_id, event, event_flag); - LOG_IF(FATAL, st != 0) << model_name_ - << "execute prefill layer fail, error code: " << st; + const bool is_decode = input_params_new.meta.batch_forward_type.is_decode(); + + // Pick the right ATB node based on (active mode, decode vs prefill). + // In legacy single-mode dual_parallel_args_ is null and we keep the + // original prefill_node_/decode_node_ behaviour. In dual mode the + // construction-time active landed in prefill_node_/decode_node_ and + // the OTHER mode landed in dp_prefill_node_/dp_decode_node_; we + // compare the live mode against dual_ctor_mode_ to decide. + atb_speed::Model::Node* node = nullptr; + if (dual_parallel_args_ == nullptr) { + node = is_decode ? &decode_node_ : &prefill_node_; } else { - build_node_variant_pack(prefill_node_, - x, - cos_pos, - sin_pos, - attn_mask, - kv_cache, - input_params_new, - true, - shared_topk_indices, - output_topk_indices, - skip_topk_, - output_topk_); - st = execute_node(prefill_node_, node_id, event, event_flag); - LOG_IF(FATAL, st != 0) << model_name_ - << "execute prefill layer fail, error code: " << st; + const bool use_original = (dual_parallel_args_->mode() == dual_ctor_mode_); + if (use_original) { + node = is_decode ? &decode_node_ : &prefill_node_; + } else { + node = is_decode ? &dp_decode_node_ : &dp_prefill_node_; + } } + + build_node_variant_pack(*node, + x, + cos_pos, + sin_pos, + attn_mask, + kv_cache, + input_params_new, + /*is_prefill=*/!is_decode, + shared_topk_indices, + output_topk_indices, + skip_topk_, + output_topk_); + st = execute_node(*node, node_id, event, event_flag); + LOG_IF(FATAL, st != 0) << model_name_ + << " execute layer fail, error code: " << st; return tensor_placeholder_; } @@ -1170,9 +1221,20 @@ void NpuDeepseekV32DecoderLayerImpl::build_node_variant_pack( auto& dp_ep_padding = input_params.parallel.dp_ep_padding_data; auto& cp_ep_padding = input_params.parallel.cp_ep_padding_data; auto& cp_inputs = input_params.parallel.cp_prefill_inputs; - const bool use_cp_ep_padding = (cp_size_ > 1 && is_prefill); - - if (dp_size_ <= 1 && ep_size_ <= 1 || cp_size_ > 1) { + // cp_size_/dp_size_ are baked in at ctor from the construction-time mode. A + // runtime CP<->DP flip changes the live layout but does NOT rewrite these + // members, so read the current layout from dual_parallel_args_ when present + // (the forward path already picks the ATB node the same way). Falls back to + // the ctor members on the legacy single-mode path. + const int32_t active_cp_size = dual_parallel_args_ != nullptr + ? dual_parallel_args_->active().cp_size() + : cp_size_; + const int32_t active_dp_size = dual_parallel_args_ != nullptr + ? dual_parallel_args_->active().dp_size() + : dp_size_; + const bool use_cp_ep_padding = (active_cp_size > 1 && is_prefill); + + if (active_dp_size <= 1 && ep_size_ <= 1 || active_cp_size > 1) { dp_ep_padding.set_placeholder(tensor_placeholder_); } if (!use_cp_ep_padding) { @@ -1346,14 +1408,22 @@ void NpuDeepseekV32DecoderLayerImpl::build_node_variant_pack( } int32_t cp_input_index = WEIGHT_COUNT_PER_LAYER + 32; - if (skip_topk) { - CHECK(shared_topk_indices.defined()) - << "DSA top-k sharing requires previous top-k indices."; - node.variantPack.inTensors.at(cp_input_index++) = - atb_speed::Utils::AtTensor2Tensor(shared_topk_indices); + if (skip_topk || output_topk) { + // TODO: support DSA top-k sharing for CP prefill. + CHECK(!(active_cp_size > 1 && is_prefill)) + << "DSA top-k sharing does not support CP prefill yet."; + if (skip_topk) { + CHECK(shared_topk_indices.defined()) + << "DSA top-k sharing requires previous top-k indices."; + node.variantPack.inTensors.at(cp_input_index++) = + atb_speed::Utils::AtTensor2Tensor(shared_topk_indices); + } } - if (cp_size_ > 1 && is_prefill) { + // active_cp_size (from dual_parallel_args_) instead of cp_size_ so a + // runtime CP<->DP flip is respected: after flip to DP the live cp_size + // is 1 and the CP-only tensors below must be skipped. + if (active_cp_size > 1 && is_prefill) { node.variantPack.inTensors.at(cp_input_index++) = atb_speed::Utils::AtTensor2Tensor(cp_inputs.cp_o_recover_idx); node.variantPack.inTensors.at(cp_input_index++) = diff --git a/xllm/core/layers/npu/npu_deepseek_v32_decoder_layer_impl.h b/xllm/core/layers/npu/npu_deepseek_v32_decoder_layer_impl.h index c3e1f6af1e..fb35dc721e 100644 --- a/xllm/core/layers/npu/npu_deepseek_v32_decoder_layer_impl.h +++ b/xllm/core/layers/npu/npu_deepseek_v32_decoder_layer_impl.h @@ -211,6 +211,36 @@ class NpuDeepseekV32DecoderLayerImpl : public BaseLayer { atb_speed::Model::Node mtp_prefill_fallback_node_; atb_speed::Model::Node mtp_decode_fallback_node_; + // Dual-mode mirrors. Populated only when the layer was constructed + // through a ModelContext that carried a DualParallelArgs; in that case + // the four (param, node) pairs above hold the CP-mode configuration + // (cp_size = world, dp_size = 1) and the four below hold the DP-mode + // configuration (cp_size = 1, dp_size = world). forward() inspects + // dual_parallel_args_ at call time and dispatches accordingly. + // + // Memory note: each ATB DecoderLayer node carries its own workspace + // (kernel cache, intermediate buffers, etc.); two extra nodes per + // layer × 60 layers is the bulk of the ~1-2 GB / NPU dual-graph + // overhead documented in DESIGN_DOC_v0.2.md. The weights themselves + // are referenced by both pairs through atb_weight_tensors_, no copies. + atb_speed::deepseekV2::DecoderLayerParam dp_prefill_param_; + atb_speed::deepseekV2::DecoderLayerParam dp_decode_param_; + + atb_speed::Model::Node dp_prefill_node_; + atb_speed::Model::Node dp_decode_node_; + + // Non-owning pointer back to the per-worker DualParallelArgs (lives on + // WorkerServer). Null in legacy single-mode -- forward() then routes + // straight to prefill_node_ / decode_node_ as before. + const DualParallelArgs* dual_parallel_args_ = nullptr; + + // Records which mode was active at construction time. The construction- + // time active mode landed in prefill_node_/decode_node_ (the original + // pair) and the other mode landed in dp_prefill_node_/dp_decode_node_. + // forward() compares dual_parallel_args_->mode() against this field to + // decide which pair to use. + DualParallelArgs::Mode dual_ctor_mode_ = DualParallelArgs::Mode::CP_PREFILL; + atb::Tensor internal_tensor_; torch::Tensor at_cumsum_; diff --git a/xllm/core/runtime/llm_worker_impl.cpp b/xllm/core/runtime/llm_worker_impl.cpp index 24d3df340d..0a8ab633e4 100644 --- a/xllm/core/runtime/llm_worker_impl.cpp +++ b/xllm/core/runtime/llm_worker_impl.cpp @@ -54,6 +54,14 @@ bool uses_worker_cp_partition(const runtime::Options& options) { return options.cp_size() > 1 && !Platform::uses_model_cp_partition(); } +// Runtime-aware variant: reads the live parallel_args_.cp_size() instead +// of the immutable options_.cp_size(), so a CP<->DP flip that changes +// the effective cp_size is respected. Callers that only run in single- +// mode (no runtime flip) may still use the options-based overload. +bool uses_worker_cp_partition(const ParallelArgs& parallel_args) { + return parallel_args.cp_size() > 1 && !Platform::uses_model_cp_partition(); +} + void wait_input_ready_events(const ForwardInput& input, const Stream& stream) { CHECK(stream.wait_event(input.metadata_ready_event)) << "failed to wait ForwardInput metadata ready event"; @@ -307,7 +315,11 @@ std::optional LLMWorkerImpl::step_internal( /*non_blocking=*/false) .contiguous(); } - if (uses_worker_cp_partition(options_)) { + // Use the live parallel_args_ (not options_) so a runtime CP<->DP + // flip is respected: after flipping to DP_DECODE the live cp_size is + // 1 and we must take the non-CP logits path. Options holds the + // STARTUP value and never follows a runtime flip. + if (uses_worker_cp_partition(parallel_args_)) { logits = model_->logits(model_output.hidden_states, selected_token_idxes, selected_hidden_from_lm_head); @@ -385,12 +397,12 @@ std::optional LLMWorkerImpl::step_internal( // separately so the embedding cache stores it without re-selecting on // the local shard. output.sample_output.embeddings = embeddings; - if (uses_worker_cp_partition(options_) && + if (uses_worker_cp_partition(parallel_args_) && selected_hidden_from_lm_head.defined()) { output.sample_output.selected_embeddings = selected_hidden_from_lm_head; } } else if (sampling_params.selected_token_idxes.defined()) { - if (uses_worker_cp_partition(options_)) { + if (uses_worker_cp_partition(parallel_args_)) { CHECK(selected_hidden_from_lm_head.defined()) << "selected_hidden_from_lm_head must be defined when " "selected_token_idxes is defined."; diff --git a/xllm/core/runtime/options.h b/xllm/core/runtime/options.h index 97c8c7d732..e7ed10a044 100644 --- a/xllm/core/runtime/options.h +++ b/xllm/core/runtime/options.h @@ -260,6 +260,21 @@ struct Options { // max concurrency for rec worker PROPERTY(int32_t, rec_worker_max_concurrency) = 1; + + // Enable runtime CP<->DP switching (dual-graph mode). When true, + // worker_server brings up TWO CollectiveCommunicator instances at + // startup -- one in CP=N config, one in DP=N config -- and the + // worker holds a DualParallelArgs that lets switch_mode() flip + // between them at runtime without service restart. Comes with a + // ~600 MB / NPU memory overhead (two ATB commDomains + their HCCL + // SubComm buffers); see DESIGN_DOC_v0.2.md for the full budget. + PROPERTY(bool, enable_runtime_cp_dp_switch) = false; + + // Port stride between the CP and DP communicator master_addr bases + // when enable_runtime_cp_dp_switch=true. Must clear the inner + // fan-out (~world_size + dp + tp + single-rank) plus a TIME_WAIT + // safety margin. 256 is the value validated by tests/probes/*. + PROPERTY(int32_t, dual_mode_port_stride) = 256; }; } // namespace runtime diff --git a/xllm/core/runtime/worker.cpp b/xllm/core/runtime/worker.cpp index f5c6c41b75..4fff8e8bec 100644 --- a/xllm/core/runtime/worker.cpp +++ b/xllm/core/runtime/worker.cpp @@ -79,8 +79,30 @@ Worker::Worker(const ParallelArgs& parallel_args, } } +Worker::Worker(DualParallelArgs* dual_args, + const torch::Device& device, + const runtime::Options& options, + WorkerType worker_type) + : Worker(dual_args->active(), device, options, worker_type) { + // Forwards to the legacy ctor to avoid duplicating the per-WorkerType + // dispatch table; then late-binds the dual source so switch_mode can + // see it. This matches the WorkerImpl pattern. + if (impl_ != nullptr) { + impl_->attach_dual_parallel_args(dual_args); + } +} + Worker::~Worker() { delete impl_; } +bool Worker::switch_mode(DualParallelArgs::Mode target) { + return impl_ != nullptr ? impl_->switch_mode(target) : false; +} + +DualParallelArgs::Mode Worker::current_mode() const { + return impl_ != nullptr ? impl_->current_mode() + : DualParallelArgs::Mode::CP_PREFILL; +} + bool Worker::init_model(const std::string& model_weights_path, int32_t random_seed, MasterStatus master_status) { @@ -258,4 +280,24 @@ folly::SemiFuture Worker::stop_profile_async() { }); return future; } + +folly::SemiFuture Worker::switch_mode_async(int32_t target_mode) { + folly::Promise promise; + auto future = promise.getSemiFuture(); + threadpool_.schedule( + [this, target_mode, promise = std::move(promise)]() mutable { + if (target_mode != + static_cast(DualParallelArgs::Mode::CP_PREFILL) && + target_mode != + static_cast(DualParallelArgs::Mode::DP_DECODE)) { + LOG(WARNING) << "switch_mode_async invalid target_mode=" + << target_mode; + promise.setValue(false); + return; + } + promise.setValue(this->switch_mode( + static_cast(target_mode))); + }); + return future; +} } // namespace xllm diff --git a/xllm/core/runtime/worker.h b/xllm/core/runtime/worker.h index 5b7e30fc40..0893f2c43d 100644 --- a/xllm/core/runtime/worker.h +++ b/xllm/core/runtime/worker.h @@ -41,6 +41,15 @@ class Worker { const runtime::Options& options, WorkerType worker_type); + // Dual-mode constructor. Builds the underlying WorkerImpl from + // dual_args->active() and then attaches the dual source so + // switch_mode() can flip the active configuration without rebuilding + // anything. The DualParallelArgs must outlive the Worker. + Worker(DualParallelArgs* dual_args, + const torch::Device& device, + const runtime::Options& options, + WorkerType worker_type); + ~Worker(); // initialize model, cache manager. blocking call @@ -139,6 +148,15 @@ class Worker { folly::SemiFuture get_active_activation_memory_async(); + // Runtime CP<->DP switching. Returns false if the worker was not built + // with dual-mode support (see DualParallelArgs / runtime/options.h + // enable_runtime_cp_dp_switch flag). + bool switch_mode(DualParallelArgs::Mode target); + + folly::SemiFuture switch_mode_async(int32_t target_mode); + + [[nodiscard]] DualParallelArgs::Mode current_mode() const; + private: WorkerImpl* impl_ = nullptr; ThreadPool threadpool_{/*num_threads=*/1, diff --git a/xllm/core/runtime/worker_client.cpp b/xllm/core/runtime/worker_client.cpp index 0759598e4e..c82ff4d168 100644 --- a/xllm/core/runtime/worker_client.cpp +++ b/xllm/core/runtime/worker_client.cpp @@ -194,6 +194,10 @@ folly::SemiFuture WorkerClient::stop_profile_async() { return worker_->stop_profile_async(); } +folly::SemiFuture WorkerClient::switch_mode_async(int32_t target_mode) { + return worker_->switch_mode_async(target_mode); +} + const torch::Device& WorkerClient::device() const { return worker_->device(); } folly::SemiFuture> diff --git a/xllm/core/runtime/worker_client.h b/xllm/core/runtime/worker_client.h index 8186527ccc..41a9cbcafd 100644 --- a/xllm/core/runtime/worker_client.h +++ b/xllm/core/runtime/worker_client.h @@ -57,6 +57,11 @@ class WorkerClient { virtual folly::SemiFuture stop_profile_async(); + // Runtime CP<->DP switch. target_mode mirrors DualParallelArgs::Mode. + // Future resolves to false on RPC failure or if the remote worker + // is in legacy single-mode (no DualParallelArgs attached). + virtual folly::SemiFuture switch_mode_async(int32_t target_mode); + virtual std::tuple estimate_kv_cache_capacity(); // allocate kv cache. blocking call diff --git a/xllm/core/runtime/worker_impl.cpp b/xllm/core/runtime/worker_impl.cpp index 6390495833..548002eea0 100644 --- a/xllm/core/runtime/worker_impl.cpp +++ b/xllm/core/runtime/worker_impl.cpp @@ -296,8 +296,70 @@ WorkerImpl::WorkerImpl(const ParallelArgs& parallel_args, #endif } +WorkerImpl::WorkerImpl(DualParallelArgs* dual_args, + const torch::Device& device, + const runtime::Options& options) + : WorkerImpl(dual_args->active(), device, options) { + // The single-arg ctor above already snapshotted the active() ParallelArgs + // into parallel_args_; we just retain the pointer so switch_mode() can + // refresh the snapshot after a flip. + dual_parallel_args_ = dual_args; +} + WorkerImpl::~WorkerImpl() = default; +bool WorkerImpl::switch_mode(DualParallelArgs::Mode target) { + if (dual_parallel_args_ == nullptr) { + LOG(WARNING) << "Worker rank=" << parallel_args_.rank() + << " switch_mode requested but the worker was constructed " + "in single-mode; ignoring."; + return false; + } + if (dual_parallel_args_->mode() == target) { + return true; + } + dual_parallel_args_->set_mode(target); + // Refresh the snapshot so legacy parallel_args_ readers (rank/world_size + // are mode-invariant; cp_size/dp_size and the ProcessGroup pointers do + // change) see the post-flip configuration. This is safe to do without + // holding a lock because the scheduler's drain protocol guarantees no + // forward is in flight on this worker when switch_mode runs. + parallel_args_ = dual_parallel_args_->active(); + // ModelContext holds its OWN copy of ParallelArgs (not a reference), and the + // forward path reads cp_size/dp_size/ep_size/mapping_data through it (e.g. + // the DpEpPadding gate in prepare_work_before_execute_on_stream). Without + // this refresh the context keeps the pre-flip cp_size, so after a flip to + // DP_DECODE the gate `!(cp_size > 1)` stays false, DpEpPadding is never + // built, and the DP decode ATB node faults on placeholder tensors. + context_.set_parallel_args(parallel_args_); + // dp_driver_ depends on cp_size/dp_size and must be recomputed. + const int32_t tp_size = parallel_args_.world_size() / + (parallel_args_.dp_size() * parallel_args_.cp_size()); + dp_driver_ = + parallel_args_.dp_size() > 1 && + parallel_args_.rank() % (tp_size * parallel_args_.cp_size()) == 0; + return true; +} + +DualParallelArgs::Mode WorkerImpl::current_mode() const { + if (dual_parallel_args_ == nullptr) { + return DualParallelArgs::Mode::CP_PREFILL; + } + return dual_parallel_args_->mode(); +} + +void WorkerImpl::attach_dual_parallel_args(DualParallelArgs* dual_args) { + if (dual_parallel_args_ == dual_args) { + return; + } + dual_parallel_args_ = dual_args; + if (dual_args != nullptr) { + // Refresh the singular snapshot from the dual source so subsequent + // reads are consistent with whatever mode the dual is currently in. + parallel_args_ = dual_args->active(); + } +} + bool WorkerImpl::allocate_kv_cache_storage(const KVCacheShape& kv_cache_shape, bool use_huge_page_allocator, bool enable_raw_device_allocator) { @@ -822,20 +884,22 @@ void WorkerImpl::prepare_work_before_execute_on_stream( // already-partitioned device tensors; see ForwardInput::cp_partitioned. // Prefill-side CP (partition + ATB cp tensors) applies to PREFILL, // CHUNKED_PREFILL, and MIXED. `no_decode()` wrongly excludes MIXED. - const bool needs_cp_prefill_side = - parallel_args_.cp_size() > 1 && !Platform::uses_model_cp_partition() && - !input.input_params.meta.batch_forward_type.is_decode(); - const bool needs_cp_partition = - needs_cp_prefill_side && !input.cp_partitioned; + // + // The multi-node RPC path delivers `input` as a packed host buffer whose + // `meta.batch_forward_type` is still the default EMPTY until unpacked. The + // CP gate below MUST see the true batch type, so materialize the CP working + // copy (and thus its meta) BEFORE deciding needs_cp_prefill_side. Reading + // the pre-unpack meta made every RPC batch look non-decode (EMPTY), which + // wrongly drove pure-decode steps into prepare_cp_prefill_inputs and crashed + // on a 1-token seq (input_lengths=[1] -> cumsum-1 = -1 -> index_select(-1)). std::optional cp_input; - if (needs_cp_prefill_side) { + if (parallel_args_.cp_size() > 1) { cp_input.emplace(input); - } #if defined(USE_NPU) - if (needs_cp_prefill_side) { ForwardInput& cp_working = *cp_input; - // RPC packed_input only carries input_host_buffer until unpack; partition - // and prepare_cp_prefill_inputs need materialized token_ids / seq_lens. + // RPC packed_input only carries input_host_buffer until unpack; the gate, + // partition, and prepare_cp_prefill_inputs all need the materialized meta / + // token_ids / seq_lens. if (cp_working.input_host_buffer_has_layout && (!cp_working.token_ids.defined() || cp_working.token_ids.numel() == 0)) { @@ -847,12 +911,31 @@ void WorkerImpl::prepare_work_before_execute_on_stream( cp_working.input_host_buffer_has_layout = false; } else { LOG(ERROR) << "[CP_PREP] unpack_from_input_host_buffer failed before " - "cp_partition (cp_rank=" + "cp gate (cp_rank=" << parallel_args_.cp_rank() << ")"; } } - } #endif + } + + // Resolve the true batch type from the (possibly just-unpacked) CP working + // copy; fall back to `input` when CP is disabled and no copy was made. + const BatchForwardType& resolved_batch_type = + cp_input ? cp_input->input_params.meta.batch_forward_type + : input.input_params.meta.batch_forward_type; + // Upstream Platform::uses_model_cp_partition() gate: when the platform + // does CP partition at the model layer, worker-side partition must be + // skipped (introduced by upstream after our probe branch forked). + const bool needs_cp_prefill_side = parallel_args_.cp_size() > 1 && + !Platform::uses_model_cp_partition() && + !resolved_batch_type.is_decode(); + const bool needs_cp_partition = + needs_cp_prefill_side && !input.cp_partitioned; + // Release the CP working copy if the resolved type turned out to be pure + // decode (or empty) so downstream code uses the untouched `input`. + if (cp_input && !needs_cp_prefill_side) { + cp_input.reset(); + } if (needs_cp_partition) { ForwardInput& cp_working = *cp_input; const int64_t tokens_before = @@ -949,6 +1032,31 @@ void WorkerImpl::prepare_work_before_execute_on_stream( (context_.get_parallel_args().dp_size() > 1 || context_.get_parallel_args().ep_size() > 1 || !context_.get_parallel_args().mapping_data().empty())); + // FLIPDIAG: DP-mode empty shard tracing. In DP forward the collective + // ops require every rank in the DP group to enter; if this rank has + // an empty shard AND the fake-input path did not fire, we return + // early below and hang the OTHER rank's collective wait. Log which + // path we take so we can pin down whether the flip propagates the + // right cp_size/dp_size/mapping_data into the context. + // Gated by enable_flip_verbose_log (fires per-rank per-step). + if (FLAGS_enable_flip_verbose_log) { + LOG(INFO) << "FLIPDIAG worker_forward_shard: rank=" + << parallel_args_.rank() + << " num_sequences=" << input_params.meta.num_sequences + << " tokens=" + << (processed_input.token_ids.defined() + ? processed_input.token_ids.numel() + : -1) + << " empty_shard=" << empty_shard + << " need_fake=" << need_fake_input_for_empty_shard + << " cp_size=" << context_.get_parallel_args().cp_size() + << " dp_size=" << context_.get_parallel_args().dp_size() + << " ep_size=" << context_.get_parallel_args().ep_size() + << " mapping_empty=" + << context_.get_parallel_args().mapping_data().empty() + << " batch_type=" + << input_params.meta.batch_forward_type.to_string(); + } if (need_fake_input_for_empty_shard) { auto token_options = processed_input.token_ids.defined() ? processed_input.token_ids.options() @@ -961,8 +1069,22 @@ void WorkerImpl::prepare_work_before_execute_on_stream( processed_input.positions = torch::zeros({1}, position_options.device(device_)); empty_shard = false; + // Fabricates a 1-token forward for the empty shard so the DP + // collective sees a shape-consistent input on every rank; llm_engine + // clamps dp_global_token_nums[dp_rank] to 1 for empty ranks, matching + // this fake token count. attention.device.block_tables is deliberately + // not filled here -- the decoder layer's placeholder branch handles + // the empty shard, and any KV-block dummy would race with dp_rank=0's + // real sequences (v15 experiment: rtMemcpyAsync vector core exception). + // In practice this branch is rare because ContinuousScheduler defers + // lopsided DP batches (see step() in continuous_scheduler.cpp); + // reaching here means the >100ms backdoor fired. } if (empty_shard) { + LOG(WARNING) << "FLIPDIAG worker_forward_shard EARLY_RETURN: rank=" + << parallel_args_.rank() + << " (this rank will NOT enter collective; other DP " + "ranks that expect it will hang)"; return; } @@ -1608,6 +1730,12 @@ bool WorkerImpl::init_model(const std::string& model_weights_path, auto tensor_options = torch::dtype(dtype_).device(device_); context_ = ModelContext(parallel_args_, args, quant_args, tensor_options); context_.set_model_id(options_.model_id()); + // Plumb the optional DualParallelArgs through to ModelContext so + // dual-aware layers (DeepSeek V32) can see the OTHER mode's args + // and build the second pair of ATB graphs at init_layer time. + // dual_parallel_args_ is null in legacy single-mode and the layer + // ctor short-circuits accordingly. + context_.set_dual_parallel_args(dual_parallel_args_); // init model, create model executor bool status = this->init_model(context_); diff --git a/xllm/core/runtime/worker_impl.h b/xllm/core/runtime/worker_impl.h index 50fae27196..7943547b28 100644 --- a/xllm/core/runtime/worker_impl.h +++ b/xllm/core/runtime/worker_impl.h @@ -63,6 +63,16 @@ class WorkerImpl { const torch::Device& device, const runtime::Options& options); + // Dual-mode constructor used by the runtime CP<->DP switch feature. + // The DualParallelArgs is owned by worker_server (which also owns the + // two CollectiveCommunicator instances feeding it); WorkerImpl keeps a + // non-owning pointer and copies the active() snapshot into + // parallel_args_ for the singular-API readers. When dual_args is non- + // null, switch_mode() can flip the active mode at runtime. + WorkerImpl(DualParallelArgs* dual_args, + const torch::Device& device, + const runtime::Options& options); + virtual ~WorkerImpl(); // initialize model, cache manager. blocking call @@ -76,6 +86,23 @@ class WorkerImpl { virtual void lazy_load_model(std::unique_ptr loader); + // Switch between CP-prefill and DP-decode modes. No-op when the worker + // was constructed without a DualParallelArgs (legacy single-mode); + // returns false in that case so callers can detect the mismatch. + // Caller must drain inflight requests before invoking; this only flips + // the discriminator, it does not block on inflight forwards. + virtual bool switch_mode(DualParallelArgs::Mode target); + + // Current active mode. Always CP_PREFILL when the worker is single-mode. + [[nodiscard]] DualParallelArgs::Mode current_mode() const; + + // Late-bind the DualParallelArgs source for a worker that was first + // constructed via the legacy single-arg ctor. Used by Worker (the + // facade) to attach dual support without touching every WorkerImpl + // subclass's ctor signature. Idempotent; no-op if already set to the + // same pointer. The DualParallelArgs must outlive this WorkerImpl. + void attach_dual_parallel_args(DualParallelArgs* dual_args); + virtual std::tuple estimate_kv_cache_capacity(); // allocate kv cache. blocking call @@ -312,6 +339,13 @@ class WorkerImpl { // parallel args of current instance ParallelArgs parallel_args_; + // Optional dual-mode source. When non-null, parallel_args_ above is a + // snapshot of dual_parallel_args_->active() taken at construction + // time; future switch_mode() calls update both this snapshot AND the + // atomic flag inside dual_parallel_args_. When null, the worker is in + // legacy single-mode and parallel_args_ is the only source of truth. + DualParallelArgs* dual_parallel_args_ = nullptr; + // kv caches std::vector kv_caches_; diff --git a/xllm/core/runtime/xservice_client.cpp b/xllm/core/runtime/xservice_client.cpp index 4ddbb4b3f2..a24fb31129 100644 --- a/xllm/core/runtime/xservice_client.cpp +++ b/xllm/core/runtime/xservice_client.cpp @@ -164,7 +164,7 @@ bool XServiceClient::init(const std::string& etcd_addr, }; etcd_client_->add_watch(ETCD_XSERVICES_KEY_PREFIX, xservices_func); - block_manager_pool_ = block_manager_pool; + block_manager_pool_.store(block_manager_pool, std::memory_order_release); initialize_done_ = true; return true; @@ -176,6 +176,14 @@ void XServiceClient::set_scheduler(Scheduler* scheduler) { void XServiceClient::set_engine(Engine* engine) { engine_ = engine; } +void XServiceClient::set_block_manager_pool( + const BlockManagerPool* block_manager_pool) { + // Release-store: pairs with the acquire-load in heartbeat(). Callers hold + // no lock; heartbeat/reconcile threads are unaffected because they snapshot + // the pointer at loop-iteration granularity. + block_manager_pool_.store(block_manager_pool, std::memory_order_release); +} + XServiceClient::~XServiceClient() { exited_.store(true); if (heartbeat_thread_ && heartbeat_thread_->joinable()) { @@ -286,6 +294,47 @@ void XServiceClient::register_instance(const InstanceInfo& instance_info) { LOG(INFO) << "Success register instance to etcd."; } +bool XServiceClient::re_register_dp_size(int32_t new_dp_size) { + // A runtime CP<->DP flip changes this instance's dp_size. Peers cache + // InstanceInfo the first time they query it (DisaggPDScheduler. + // check_remote_instance_info: "if already cached, keep it") and never + // refresh, so their post-flip relink would build the wrong topology + // unless we push the new dp_size back to etcd here. We re-serialize the + // JSON we handed etcd at register time (kept in registration_value_) + // with the new dp_size, leaving every other field intact. + std::string key; + std::string value; + { + std::lock_guard lock(registration_mutex_); + if (registration_key_.empty() || registration_value_.empty()) { + LOG(ERROR) << "re_register_dp_size called before register_instance"; + return false; + } + nlohmann::json json_val; + try { + json_val = nlohmann::json::parse(registration_value_); + } catch (const std::exception& e) { + LOG(ERROR) << "re_register_dp_size: failed to parse cached registration " + "value: " + << e.what(); + return false; + } + json_val["dp_size"] = new_dp_size; + registration_value_ = json_val.dump(); + key = registration_key_; + value = registration_value_; + } + + if (!register_instance_with_retry(key, value)) { + LOG(ERROR) << "re_register_dp_size: failed to write updated dp_size=" + << new_dp_size << " to etcd"; + return false; + } + LOG(INFO) << "re_register_dp_size: published dp_size=" << new_dp_size + << " to etcd for instance " << instance_name_; + return true; +} + InstanceInfo XServiceClient::get_instance_info( const std::string& instance_name) { InstanceInfo result; @@ -346,7 +395,15 @@ void XServiceClient::heartbeat() { 1000))); if (!register_done_.load()) continue; - if (block_manager_pool_ == nullptr || scheduler_ == nullptr) continue; + // Snapshot the pool pointer once per iteration; the writer + // (set_block_manager_pool, called from + // LLMEngine::rebuild_block_manager_pool during a runtime CP<->DP flip) may + // swap it under us at any point. Reading the atomic once and using the + // local for the whole iteration keeps this section race-free without + // needing a mutex. + const BlockManagerPool* pool = + block_manager_pool_.load(std::memory_order_acquire); + if (pool == nullptr || scheduler_ == nullptr) continue; brpc::Controller cntl; xllm_service::proto::HeartbeatRequest req; @@ -354,7 +411,7 @@ void XServiceClient::heartbeat() { req.set_incarnation_id(incarnation_id_); req.mutable_load_metrics()->set_gpu_cache_usage_perc( - block_manager_pool_->get_gpu_cache_usage_perc()); + pool->get_gpu_cache_usage_perc()); req.mutable_load_metrics()->set_waiting_requests_num( scheduler_->get_waiting_requests_num()); diff --git a/xllm/core/runtime/xservice_client.h b/xllm/core/runtime/xservice_client.h index e2fa4f2269..d4dc10cd48 100644 --- a/xllm/core/runtime/xservice_client.h +++ b/xllm/core/runtime/xservice_client.h @@ -49,10 +49,25 @@ class XServiceClient { const std::string& etcd_namespace = ""); void set_scheduler(Scheduler* scheduler); void set_engine(Engine* engine); + // Atomically replace the pool pointer that heartbeat/reconcile threads read. + // Called from LLMEngine::rebuild_block_manager_pool after a runtime CP<->DP + // flip. Without this, the heartbeat thread keeps the pointer captured at + // init() time and dereferences the destroyed pool the moment + // unique_ptr::reset swaps it in the engine (silent SIGSEGV, rank0 disappears + // ~5s post-flip). + void set_block_manager_pool(const BlockManagerPool* block_manager_pool); bool initialize_done() { return initialize_done_; } std::string get_instance_name(); void register_instance(const InstanceInfo& instance_info); + // Re-publish this instance's InstanceInfo to etcd after a runtime CP<->DP + // flip changed dp_size. The peer scheduler pulls InstanceInfo lazily on + // demand (DisaggPDScheduler::check_remote_instance_info) and its cache is + // sticky; without this write, the peer keeps the pre-flip dp_size in + // remote_instances_info_ and its own relink_after_flip will link to the + // wrong topology. Rewrites the JSON we handed to etcd at register_instance + // time with the new dp_size, keeping every other field intact. + bool re_register_dp_size(int32_t new_dp_size); void heartbeat(); InstanceInfo get_instance_info(const std::string& instance_name); std::vector get_static_decode_list(); @@ -114,8 +129,11 @@ class XServiceClient { std::mutex registration_mutex_; brpc::ChannelOptions chan_options_; std::unique_ptr etcd_client_; - const BlockManagerPool* block_manager_pool_ = nullptr; // not own - Scheduler* scheduler_ = nullptr; // not own + // Atomic so heartbeat_thread_/reconcile_thread_ can safely re-read after a + // runtime CP<->DP flip's rebuild_block_manager_pool. Load with + // memory_order_acquire on the reader side; store with release on the writer. + std::atomic block_manager_pool_{nullptr}; // not own + Scheduler* scheduler_ = nullptr; // not own Engine* engine_ = nullptr; // not own, for xtensor info }; diff --git a/xllm/core/scheduler/continuous_scheduler.cpp b/xllm/core/scheduler/continuous_scheduler.cpp index 2dbca78784..c0d8bd2cca 100644 --- a/xllm/core/scheduler/continuous_scheduler.cpp +++ b/xllm/core/scheduler/continuous_scheduler.cpp @@ -119,6 +119,11 @@ ContinuousScheduler::ContinuousScheduler(Engine* engine, const Options& options) request_queue_(::xllm::RecConfig::get_instance().request_queue_size()) { CHECK(engine_ != nullptr); + // Tracks the dp_size the scheduler's batching pipeline is currently built + // for. Starts at the startup dp_size and is updated in step() when a runtime + // CP<->DP flip changes engine_->dp_size(). + active_dp_size_ = options_.dp_size(); + kv_cache_manager_ = engine_->block_manager_pool(); CHECK(kv_cache_manager_ != nullptr); @@ -129,7 +134,7 @@ ContinuousScheduler::ContinuousScheduler(Engine* engine, const Options& options) enable_in_batch_prefix_cache_ = ::xllm::KVCacheConfig::get_instance().enable_in_batch_prefix_cache(); - last_batch_.resize(options_.dp_size()); + last_batch_.resize(active_dp_size_); ProfileManager::Options profile_manager_options; profile_manager_options.dp_size(options.dp_size()) @@ -489,10 +494,21 @@ void ContinuousScheduler::handle_prefill_requests( // Currently overestimating the number of tokens actually processed when // enable prefix cache size_t num_tokens = prefill_sequence->num_need_compute_tokens(); - const int32_t worker_cp_size = - Platform::uses_model_cp_partition() ? 1 : options_.cp_size(); + // CP prefill token alignment only applies when this worker actually + // performs worker-side CP partition with cp_size > 1. Two conditions + // can zero out that need: + // 1. Runtime flip to DP_DECODE makes the live cp_size 1 + // (active_dp_size_ > 1 => DP mode). Do NOT read the immutable + // options_.cp_size() here. + // 2. Platform does CP partition at the model layer (upstream + // Platform::uses_model_cp_partition() gate). Worker-side + // alignment must be skipped in that case. + const int32_t live_cp_size = + (active_dp_size_ > 1 || Platform::uses_model_cp_partition()) + ? 1 + : options_.cp_size(); num_tokens = maybe_align_cp_prefill_tokens( - prefill_sequence.get(), num_tokens, worker_cp_size); + prefill_sequence.get(), num_tokens, live_cp_size); const size_t target_num_tokens = prefill_sequence->kv_cache_tokens_num() + num_tokens; if (remaining_token_budget < allocated_tokens + num_tokens || @@ -1183,7 +1199,7 @@ std::vector ContinuousScheduler::prepare_batch() { } auto batches = - BatchFactory::get_instance(options_.dp_size()) + BatchFactory::get_instance(active_dp_size_) ->create_batches(running_requests_, running_sequences_, running_sequences_budgets_, @@ -1297,6 +1313,25 @@ void ContinuousScheduler::step(const absl::Duration& timeout) { return; // Stay paused (or fall through to shutdown check on next loop) } + // Detect a runtime CP<->DP flip: the engine's dp_size is the source of truth + // (updated by LLMEngine::switch_mode). The scheduler's batching pipeline + // (BatchFactory, BlockManagerPool, last_batch_, per-dp_rank metrics vectors) + // is built around the startup dp_size and must be rebuilt to the new one, + // otherwise create_batches produces the wrong number of Batches and + // prepare_inputs / metrics index out of bounds. Done here on the loop thread + // (not the RPC thread that flips the flag) so the pool teardown never races a + // forward. The switch RPC's drain contract guarantees no in-flight step here. + const int32_t engine_dp = engine_->dp_size(); + if (engine_dp != active_dp_size_) { + LOG(INFO) << "ContinuousScheduler: dp_size flip detected " + << active_dp_size_ << " -> " << engine_dp << ", rebuilding pool"; + engine_->rebuild_block_manager_pool(engine_dp); + kv_cache_manager_ = engine_->block_manager_pool(); + BatchFactory::get_instance(active_dp_size_)->set_dp_size(engine_dp); + active_dp_size_ = engine_dp; + last_batch_.assign(active_dp_size_, {}); + } + if (!options_.enable_schedule_overlap()) { // get a new batch of requests last_batch_lengths_.clear(); @@ -1309,12 +1344,75 @@ void ContinuousScheduler::step(const absl::Duration& timeout) { return; } + // DP-mode lopsided-batch defer. When one dp_rank has an empty batch + // and the other has work, entering engine.step() would either (a) + // return early on the empty side and hang the peer's DP collective, + // or (b) fabricate a 1-token fake input on the empty side that races + // with the peer's real KV blocks. Both hurt. Defer this step and let + // the next tick (~1ms later) find both sides ready. If we've been + // stuck lopsided for >100ms, force the step -- better to risk a + // fake-input path than to starve the requests indefinitely (this + // backdoor has never fired in verify_switch.sh runs; DP dispatch + // typically converges in <10ms). + if (active_dp_size_ > 1) { + const bool any_empty = + std::any_of(batch.begin(), batch.end(), [](const Batch& one_batch) { + return one_batch.empty(); + }); + // Function-local static: keeping this off the class avoids changing + // sizeof(ContinuousScheduler), which would force every subclass + // (ChunkedPrefillScheduler, DisaggPDScheduler, ...) to be recompiled + // in lockstep. thread_local is defensive -- schedulers are + // single-threaded per instance, but this makes the invariant obvious. + static thread_local absl::Time last_lopsided_defer_start = + absl::InfinitePast(); + if (any_empty) { + const auto now = absl::Now(); + if (last_lopsided_defer_start == absl::InfinitePast()) { + last_lopsided_defer_start = now; + } + const bool exceeded_backdoor = + (now - last_lopsided_defer_start) > absl::Milliseconds(100); + if (!exceeded_backdoor) { + if (FLAGS_enable_flip_verbose_log) { + LOG_EVERY_N(INFO, 20) + << "FLIPDIAG lopsided_defer: active_dp_size=" << active_dp_size_ + << " (waiting for both dp_ranks)"; + } + return; + } + LOG(WARNING) << "FLIPDIAG lopsided_backdoor: active_dp_size=" + << active_dp_size_ << " forcing step after >100ms wait"; + } else { + last_lopsided_defer_start = absl::InfinitePast(); + } + } + + // FLIPDIAG: log batch shape before engine.step. Gated by + // enable_flip_verbose_log because this fires every step (potentially + // hundreds of times/sec under DP burst); flip lifecycle events elsewhere + // in this file (dp_size flip detected, rebuild_pool) stay unconditional. + if (FLAGS_enable_flip_verbose_log) { + std::string sizes; + for (size_t i = 0; i < batch.size(); ++i) { + if (i > 0) sizes += ","; + sizes += std::to_string(batch[i].size()); + } + LOG(INFO) << "FLIPDIAG step_pre_engine: active_dp_size=" + << active_dp_size_ << " batch_sizes=[" << sizes << "]"; + } + if (!options_.enable_pd_ooc()) { engine_->step(batch); } else { step_with_pd_ooc(batch); } + if (FLAGS_enable_flip_verbose_log) { + LOG(INFO) << "FLIPDIAG step_post_engine: active_dp_size=" + << active_dp_size_; + } + // process request output in batch process_batch_output(false); } else { @@ -1487,8 +1585,8 @@ void ContinuousScheduler::refresh_sequences_from_requests( std::vector ContinuousScheduler::get_num_occupied_slots( std::vector& sequences) const { - std::vector num_occupied_slots(options_.dp_size()); - std::vector num_unfilled_blocks(options_.dp_size()); + std::vector num_occupied_slots(active_dp_size_); + std::vector num_unfilled_blocks(active_dp_size_); std::vector num_used_blocks = kv_cache_manager_->num_used_blocks(); auto block_size = kv_cache_manager_->block_size(); @@ -1503,7 +1601,7 @@ std::vector ContinuousScheduler::get_num_occupied_slots( num_unfilled_blocks[dp_rank] += last_block_len > 0 ? 1 : 0; } - for (int32_t dp_rank = 0; dp_rank < options_.dp_size(); ++dp_rank) { + for (int32_t dp_rank = 0; dp_rank < active_dp_size_; ++dp_rank) { num_occupied_slots[dp_rank] += (num_used_blocks[dp_rank] - num_unfilled_blocks[dp_rank]) * block_size; } @@ -1513,12 +1611,12 @@ std::vector ContinuousScheduler::get_num_occupied_slots( std::vector ContinuousScheduler::get_active_activation_in_bytes() { std::vector all_active_activation_in_bytes = engine_->get_active_activation_memory(); - std::vector active_activation_in_bytes(options_.dp_size()); + std::vector active_activation_in_bytes(active_dp_size_); const int32_t dp_local_tp_size = - all_active_activation_in_bytes.size() / options_.dp_size(); + all_active_activation_in_bytes.size() / active_dp_size_; - for (int32_t dp_rank = 0; dp_rank < options_.dp_size(); ++dp_rank) { + for (int32_t dp_rank = 0; dp_rank < active_dp_size_; ++dp_rank) { active_activation_in_bytes[dp_rank] = all_active_activation_in_bytes[dp_rank * dp_local_tp_size]; } @@ -1536,7 +1634,7 @@ void ContinuousScheduler::update_memory_metrics( int64_t num_total_slots = kv_cache_manager_->num_blocks() * kv_cache_manager_->block_size(); - for (int32_t dp_rank = 0; dp_rank < options_.dp_size(); ++dp_rank) { + for (int32_t dp_rank = 0; dp_rank < active_dp_size_; ++dp_rank) { double occupied_slots_ratio = static_cast(num_occupied_slots[dp_rank]) / num_total_slots; double active_kv_cache_size_in_kilobytes = @@ -1642,7 +1740,7 @@ bool ContinuousScheduler::try_complete_pause() { } // Reset overlap pipeline bookkeeping so resume starts clean. last_batch_.clear(); - last_batch_.resize(options_.dp_size()); + last_batch_.resize(active_dp_size_); last_running_requests_.clear(); last_running_sequences_.clear(); is_first_step_ = true; diff --git a/xllm/core/scheduler/continuous_scheduler.h b/xllm/core/scheduler/continuous_scheduler.h index 1efba0ee56..e86c21a075 100644 --- a/xllm/core/scheduler/continuous_scheduler.h +++ b/xllm/core/scheduler/continuous_scheduler.h @@ -127,6 +127,10 @@ class ContinuousScheduler : public Scheduler { PROPERTY(bool, disable_ttft_profiling) = false; // true if enable forward interruption PROPERTY(bool, enable_forward_interruption) = false; + // Runtime CP<->DP dual-graph mode. When true, the disagg_pd brpc + // server (started by start_rpc_server) also exposes + // ModeSwitchService so xllm_service can drive runtime mode flips. + PROPERTY(bool, enable_runtime_cp_dp_switch) = false; // all requests use single global ttft PROPERTY(int32_t, max_global_ttft_ms) = std::numeric_limits::max(); // all requests use single global tpot @@ -266,6 +270,12 @@ class ContinuousScheduler : public Scheduler { KVCacheManager* kv_cache_manager_; + // dp_size the batching pipeline (BatchFactory, BlockManagerPool, last_batch_, + // per-dp_rank metrics) is currently built for. Initialized to the startup + // dp_size; updated in step() when a runtime CP<->DP flip changes + // engine_->dp_size(). Distinct from options_.dp_size(), which is immutable. + int32_t active_dp_size_ = 1; + // a thread safe queue of requests, bounded by // ::xllm::RecConfig::get_instance().request_queue_size() the schedule // owns the requests and manages their lifetimes. diff --git a/xllm/core/scheduler/disagg_pd_scheduler.cpp b/xllm/core/scheduler/disagg_pd_scheduler.cpp index ba3cad45d6..282cff1e4d 100644 --- a/xllm/core/scheduler/disagg_pd_scheduler.cpp +++ b/xllm/core/scheduler/disagg_pd_scheduler.cpp @@ -32,7 +32,10 @@ limitations under the License. #include "core/framework/config/service_config.h" #include "disagg_pd.pb.h" #include "disagg_pd_scheduler.h" +#include "distributed_runtime/disagg_pd_service.h" #include "distributed_runtime/engine.h" +#include "distributed_runtime/llm_engine.h" +#include "distributed_runtime/mode_switch_service.h" #include "framework/batch/batch_factory.h" #include "framework/kv_cache_transfer/pd_topology_guard.h" #include "framework/request/request.h" @@ -277,7 +280,21 @@ void DisaggPDScheduler::start_rpc_server() { std::make_unique(this, engine_); auto rpc_server = ServerRegistry::get_instance().register_server(server_name_); - if (!rpc_server->start(std::move(service))) { + if (options_.enable_runtime_cp_dp_switch()) { + // Co-host ModeSwitchService on the same brpc server so xllm_service + // can reach SwitchMode through instance.rpc_address. The mode_switch + // service is independent of disagg_pd.proto / worker.proto so adding + // it does not blast the proto-include graph. + auto mode_switch = + std::make_unique(static_cast(engine_), + /*scheduler=*/this); + if (!rpc_server->start(std::move(service), std::move(mode_switch))) { + LOG(ERROR) << "Failed to start brpc disagg pd + mode switch server " + "on port " + << ::xllm::DisaggPDConfig::get_instance().disagg_pd_port(); + return; + } + } else if (!rpc_server->start(std::move(service))) { LOG(ERROR) << "Failed to start brpc disagg pd server on port " << ::xllm::DisaggPDConfig::get_instance().disagg_pd_port(); return; @@ -334,6 +351,19 @@ bool DisaggPDScheduler::add_request(std::shared_ptr& request) { CHECK(request != nullptr); CHECK(!request->sequences().empty()); + // FLIPDIAG: trace P-side entry so we can pin down where DP-mode requests + // die on the flip path. add_request is called from XllmAPIService's brpc + // handler thread (chat_service_impl) after xllm_service picks this + // instance via routing.prefill_name. + // Gated: fires per-request. + if (FLAGS_enable_flip_verbose_log) { + LOG(INFO) << "FLIPDIAG add_request: request_id=" << request->request_id() + << " offline=" << request->offline() << " prompt_tokens=" + << (request->sequences().empty() + ? 0 + : request->sequences()[0]->num_prompt_tokens()); + } + kv_cache_manager_->prefetch_from_storage(request); if (request->offline()) { @@ -367,6 +397,33 @@ void DisaggPDScheduler::dispatch_requests() { break; } + // FLIPDIAG: pin the exact moment dispatch_thread grabs a request off + // the queue. If we see add_request lines but no dispatch_dequeued lines, + // the thread is blocked before the loop body (e.g. gate held + relink + // in progress). If we see dequeued but no push, the flip_guard below or + // the kv_cache_manager_ / stub call is stalling. + // Gated: fires per-request. + if (FLAGS_enable_flip_verbose_log) { + LOG(INFO) << "FLIPDIAG dispatch_dequeued: request_id=" + << request->request_id() << " decode_address='" + << request->state().decode_address + << "' active_dp_size=" << active_dp_size_; + } + + // Coordinate with runtime CP<->DP flip. Take the reader side of + // switch_gate_ around the body of one dispatch. The gate is unowned in + // steady state; a flip in progress owns it unique via begin_switch and + // will make us block here until the rebuild is done. Held for the whole + // iteration because the body reaches into engine_/kv_cache_manager_ + // (implicitly through cache_prefill_blocks and the send_first_generation + // task pushed to prefill_threadpool_). + std::shared_lock flip_guard(switch_gate_); + + if (FLAGS_enable_flip_verbose_log) { + LOG(INFO) << "FLIPDIAG dispatch_gate_acquired: request_id=" + << request->request_id(); + } + if (request->state().decode_address.empty()) { // No decode address provided to the prefill instance, just finish the // request. @@ -391,6 +448,16 @@ void DisaggPDScheduler::dispatch_requests() { } remote_instances_info_[selected_instance] = remote_info; + if (FLAGS_enable_flip_verbose_log) { + LOG(INFO) << "FLIPDIAG dispatch_topo_check: request_id=" + << request->request_id() + << " selected_instance=" << selected_instance + << " local_dp=" << instance_info_.dp_size + << " local_kv_split=" << instance_info_.kv_split_size + << " remote_dp=" << remote_info.dp_size + << " remote_kv_split=" << remote_info.kv_split_size; + } + const bool enable_mla = engine_->model_args().enable_mla(); const PdTopoResult topo_result = check_pd_topo(instance_info_, @@ -400,6 +467,8 @@ void DisaggPDScheduler::dispatch_requests() { const bool allow_pd_topo = topo_result.status == PdTopoStatus::ALLOW_HOMO || topo_result.status == PdTopoStatus::ALLOW_HETERO; if (!allow_pd_topo) { + LOG(WARNING) << "FLIPDIAG dispatch_topo_reject: request_id=" + << request->request_id() << " reason=" << topo_result.reason; if (topo_result.status == PdTopoStatus::INVALID_REMOTE) { remote_instances_info_.erase(selected_instance); } @@ -510,9 +579,35 @@ void DisaggPDScheduler::dispatch_requests() { instance_info_.ports.begin(), instance_info_.ports.end()); reqs.mutable_cluster_infos()->set_dp_size(options_.dp_size()); + // FLIPDIAG: log AddNewRequests fan-out. options_.dp_size() is the + // startup value; active_dp_size_ is the post-flip value. If they + // diverge (they do after a CP->DP flip on P side) the D side sees + // a stale P dp_size and routes KV transfer accordingly. + if (FLAGS_enable_flip_verbose_log) { + LOG(INFO) << "FLIPDIAG dispatch_add_new_requests_pre: request_id=" + << requests[0]->request_id() << " reqs=" << reqs.reqs_size() + << " options_dp_size=" << options_.dp_size() + << " active_dp_size=" << active_dp_size_ + << " sending_dp_size=" << reqs.cluster_infos().dp_size(); + } + // TODO: sync rpc here currently brpc::Controller cntl; stub->AddNewRequests(&cntl, &reqs, &resps, nullptr); + if (FLAGS_enable_flip_verbose_log) { + LOG(INFO) << "FLIPDIAG dispatch_add_new_requests_post: request_id=" + << requests[0]->request_id() << " cntl_failed=" << cntl.Failed() + << " resps_size=" << resps.resps().size() + << (resps.resps().empty() + ? std::string(" [empty]") + : " first_status=" + + std::to_string(resps.resps()[0].status_code()) + + " first_dp_rank=" + + std::to_string(resps.resps()[0].dp_rank()) + + " first_blocks=" + + std::to_string( + resps.resps()[0].blocks_ids_size())); + } if (cntl.Failed()) { LOG(ERROR) << "Failed to add new requests to decode instance : " << selected_instance << ", error text : " << cntl.ErrorText(); @@ -775,6 +870,13 @@ bool DisaggPDScheduler::decode_schedule( CHECK(request != nullptr); CHECK(!request->sequences().empty()); + // Reader-side of switch_gate_: called on the brpc worker thread that + // handles the follow-up dispatch after decode_recv_new_requests. Held for + // the whole method because this path deallocates on failure and enqueues + // for the scheduler loop to pick up on success -- either way it must not + // race with a rebuild. + std::shared_lock flip_guard(switch_gate_); + { std::lock_guard lock(received_request_map_mutex_); if (received_request_map_.find(request->request_id()) != @@ -809,6 +911,11 @@ bool DisaggPDScheduler::decode_recv_first_generation( int32_t src_dp_size, int32_t src_dp_rank, torch::Tensor mtp_bootstrap_embedding) { + // Reader-side of switch_gate_: called on the brpc worker thread from the + // decode side of DisaggPDService::FirstGeneration. Held for the whole + // method because it drives kv_cache_manager_->deallocate on failure paths + // and pushes into the scheduler's request_queue_ on success. + std::shared_lock flip_guard(switch_gate_); // push to request_queue_, and will be executed by engine. std::shared_ptr request = nullptr; { @@ -936,6 +1043,11 @@ bool DisaggPDScheduler::decode_recv_first_generation( } bool DisaggPDScheduler::try_allocate(Sequence* sequence) { + // Reader-side of switch_gate_: called from brpc worker threads via + // decode_recv_new_requests. Blocks a running flip's rebuild step until we + // finish, and blocks us while a flip is in progress -- prevents + // use-after-free on kv_cache_manager_ mid-rebuild. + std::shared_lock flip_guard(switch_gate_); // When the KV Cache usage reaches the threshold, prefill requests will no // longer be scheduled to avoid frequent preemption. if (kv_cache_manager_->kv_cache_utilization() < @@ -1001,6 +1113,11 @@ bool DisaggPDScheduler::link_instance(const std::string& instance_name, const std::vector& ports, const int32_t dp_size, const int32_t src_kv_split_size) { + // Reader-side of switch_gate_: LinkInstance from a peer P/D pair drives + // engine_->link_cluster which writes into the KV transfer plane; that + // plane's dp_size/kv_split_size views must not observe a rebuild + // half-through. + std::shared_lock flip_guard(switch_gate_); std::lock_guard lock(linked_instances_mutex_); if (!engine_->link_cluster( cluster_ids, addrs, ports, dp_size, src_kv_split_size)) { @@ -1020,6 +1137,9 @@ bool DisaggPDScheduler::unlink_instance( const std::vector& ports, const int32_t dp_size, const int32_t src_kv_split_size) { + // Reader-side of switch_gate_: mirrors link_instance. Held around both + // the received_request_map_ cleanup and the engine_->unlink_cluster call. + std::shared_lock flip_guard(switch_gate_); // Clear received requests from this instance { std::lock_guard lock(received_request_map_mutex_); @@ -1045,4 +1165,126 @@ bool DisaggPDScheduler::unlink_instance( return true; } +std::unique_lock DisaggPDScheduler::begin_switch() { + // Take the writer side of switch_gate_. This blocks until every reader + // (dispatch_thread iteration, in-flight brpc handler on this scheduler) + // has released. Callers are expected to have already paused the main + // scheduler loop via pause(WAIT) + wait_until_paused; this call closes + // the remaining async surfaces before rebuild_after_flip runs. + return std::unique_lock(switch_gate_); +} + +void DisaggPDScheduler::rebuild_after_flip(int32_t new_dp_size) { + // Must be called with switch_gate_ held unique (via begin_switch()) AND + // with the main scheduler loop paused (WAIT). Rebuilds the scheduler-side + // pipeline that was constructed for the pre-flip dp_size. + // + // The step() loop used to detect this itself and rebuild inline. That path + // races with dispatch_thread and brpc handlers -- workers whose Sequences + // still hold Block references into the old pool see the pool destroyed + // under them, causing a silent SIGSEGV on the next allocate/deallocate + // (observed rank0-only zombie on 82 CP=2 flip test). Driving the rebuild + // from ModeSwitchService with the gate closed avoids the race. + if (new_dp_size == active_dp_size_) { + LOG(INFO) << "rebuild_after_flip: dp_size unchanged (" << new_dp_size + << "), skipping rebuild"; + return; + } + LOG(INFO) << "rebuild_after_flip: rebuilding pool for dp_size " + << active_dp_size_ << " -> " << new_dp_size; + engine_->rebuild_block_manager_pool(new_dp_size); + kv_cache_manager_ = engine_->block_manager_pool(); + BatchFactory::get_instance(active_dp_size_)->set_dp_size(new_dp_size); + active_dp_size_ = new_dp_size; + last_batch_.assign(active_dp_size_, {}); + // Invalidate the peer InstanceInfo cache. Peers re-register to etcd on + // their own flip and dispatch reads remote_instances_info_ lazily + // (check_remote_instance_info: "if already cached, keep it"). Without + // clearing here dispatch will keep pushing to the pre-flip dp_size + // topology and D-side won't receive the requests. Safe under the gate: + // dispatch_thread and brpc handlers all take shared_lock on switch_gate_ + // at their entry points and we're here holding unique. + remote_instances_info_.clear(); +} + +bool DisaggPDScheduler::relink_after_flip() { + // Called by ModeSwitchService AFTER gate.unlock + resume, because the + // datadist fan-out (link_threadpool_ + folly futures + worker RPC + // round-trips) deadlocks inside the gate (v4: D-rank0 crashed with + // `std::system_error: Resource deadlock avoided` right after relink). + // + // Invariants: + // * caller has already run rebuild_after_flip (which cleared + // remote_instances_info_ under the gate) and re_register_dp_size + // (which pushed our new dp_size to etcd), so peers can now pull our + // fresh info and we'll pull theirs; + // * only DECODE instances have datadist links to rebuild -- P-side + // handles LinkInstance RPCs but never issues link_cluster on its + // own (v4 P-side relink triggered LLM_LINK_FAILED and clobbered + // the just-established D-side links). + + if (!options_.instance_role().has_value() || + options_.instance_role().value() != InstanceRole::DECODE) { + LOG(INFO) << "relink_after_flip: not a DECODE instance, skipping " + "(P-side has no outbound datadist links to rebuild)"; + return true; + } + + std::vector peers; + { + std::lock_guard lock(linked_instances_mutex_); + peers.assign(linked_instance_.begin(), linked_instance_.end()); + } + if (peers.empty()) { + LOG(INFO) << "relink_after_flip: no linked peers, skip"; + return true; + } + + bool all_ok = true; + for (const auto& peer : peers) { + // Snapshot the pre-flip peer info so we can unlink the exact link that + // was created at startup. If this instance was P-side and the peer is + // a D-side, the datadist link is one-directional from D's side; unlink + // is still safe as it's idempotent for a not-linked cluster. + auto old_it = remote_instances_info_.find(peer); + if (old_it == remote_instances_info_.end()) { + LOG(WARNING) << "relink_after_flip: no cached remote_instance_info for " + << peer << ", skip unlink; will link with fresh info"; + } else { + const auto& old = old_it->second; + if (!engine_->unlink_cluster(old.cluster_ids, + old.addrs, + old.ports, + old.dp_size, + old.kv_split_size)) { + LOG(WARNING) << "relink_after_flip: unlink stale link to " << peer + << " failed (may already be gone)"; + } + } + + InstanceInfo fresh = xservice_client_->get_instance_info(peer); + if (fresh.name.empty()) { + LOG(ERROR) << "relink_after_flip: get_instance_info failed for peer " + << peer; + all_ok = false; + continue; + } + if (!engine_->link_cluster(fresh.cluster_ids, + fresh.addrs, + fresh.ports, + fresh.dp_size, + fresh.kv_split_size)) { + LOG(ERROR) << "relink_after_flip: link with fresh dp_size=" + << fresh.dp_size << " to peer " << peer << " failed"; + all_ok = false; + continue; + } + remote_instances_info_[peer] = fresh; + LOG(INFO) << "relink_after_flip: relinked peer " << peer + << " with dp_size=" << fresh.dp_size + << " kv_split_size=" << fresh.kv_split_size; + } + return all_ok; +} + } // namespace xllm diff --git a/xllm/core/scheduler/disagg_pd_scheduler.h b/xllm/core/scheduler/disagg_pd_scheduler.h index 08aa0e11ca..a7bc4583ad 100644 --- a/xllm/core/scheduler/disagg_pd_scheduler.h +++ b/xllm/core/scheduler/disagg_pd_scheduler.h @@ -18,6 +18,7 @@ limitations under the License. #include #include +#include #include #include #include @@ -98,7 +99,57 @@ class DisaggPDScheduler : public ChunkedPrefillScheduler { const int32_t dp_size, const int32_t src_kv_split_size); + // Runtime CP<->DP switch orchestration hooks. Called by ModeSwitchService. + // + // The scheduler owns three async surfaces that must be quiesced before the + // engine tears down and rebuilds the BlockManagerPool: the main scheduler + // loop (already covered by ContinuousScheduler::pause), the dispatch_thread + // for prefill dispatch, and brpc worker threads driving DisaggPDService + // handlers (AddNewRequests / FirstGeneration / Link*). Any of these can hold + // pointers into the pool or into per-request kv_state; letting them run + // concurrent with rebuild is a use-after-free (silent SIGSEGV → rank0 + // zombie, observed on 82 CP=2 flip test). + // + // switch_gate_ is a shared_mutex: dispatch_thread and brpc handlers take + // shared_lock at their entry points; rebuild takes unique_lock. begin_switch + // hands the caller a unique_lock (blocking until every in-flight shared + // holder releases). end_switch releases it. rebuild_after_flip performs + // engine->rebuild_block_manager_pool + refreshes the scheduler-side + // active_dp_size_/last_batch_/BatchFactory bookkeeping that step() would + // otherwise do; after this ModeSwitchService can end_switch and resume. + std::unique_lock begin_switch(); + void rebuild_after_flip(int32_t new_dp_size); + + // Rebuild every LlmDataDist P<->D link after this instance flipped. + // The startup-time link topology encodes the peer's dp_size (see + // LLMEngine::link_cluster: fan-out is dp_size wide per D-worker); once + // either end flips, the (src_worker -> dst_worker) pairs shift and the + // pre-flip links no longer cover the required pairs. The next + // PushKvBlocks then returns LLM_NOT_YET_LINK (0x5010b007), which we + // observed as "60 push errors, decode receives nothing, all DP requests + // 90s timeout". + // + // The rebuild pattern for each linked peer: + // 1. get_instance_info(peer) -- pull peer's post-flip dp_size from + // etcd. The peer must have re-registered its dp_size before this + // call; ModeSwitchService orchestrates that ordering. + // 2. engine_->unlink_cluster(peer_cluster_ids, peer_addrs, peer_ports, + // old peer dp_size, kv_split_size) + // 3. engine_->link_cluster(peer_cluster_ids, peer_addrs, peer_ports, + // new peer dp_size, kv_split_size) + // 4. remote_instances_info_[peer] = fresh InstanceInfo (so subsequent + // dispatch_requests / push_kv_blocks_async use the new dp_size). + // + // Callers must be holding switch_gate_ unique. Best-effort: a per-peer + // failure logs and continues; on any relink failure returns false so + // ModeSwitchService can surface it in the SwitchMode response. + bool relink_after_flip(); + protected: + // Reader-side accessor for the shared gate. The concurrent surfaces + // (dispatch_thread, brpc-driven Disagg handlers) grab this at their entry + // points so they cannot race a rebuild in flight. + std::shared_mutex& switch_gate() { return switch_gate_; } // Pre-execute prefill requests of different lengths at startup and obtain the // corresponding TTFT for calculating the estimated TTFT of requests. void profile_ttft(); @@ -185,6 +236,13 @@ class DisaggPDScheduler : public ChunkedPrefillScheduler { std::mutex linked_instances_mutex_; std::unordered_set linked_instance_; + // Serializes runtime CP<->DP flip against dispatch_thread and brpc handlers. + // See begin_switch/rebuild_after_flip comments in the public section. + // Held shared by dispatch_requests / try_allocate / decode_schedule / + // decode_recv_first_generation / link_instance / unlink_instance. + // Held unique by begin_switch during the flip+rebuild window. + std::shared_mutex switch_gate_; + std::string server_name_; }; diff --git a/xllm/proto/CMakeLists.txt b/xllm/proto/CMakeLists.txt index bb95ddf3b1..0543b60060 100644 --- a/xllm/proto/CMakeLists.txt +++ b/xllm/proto/CMakeLists.txt @@ -25,4 +25,5 @@ proto_library( embedding_data.proto anthropic.proto xtensor_dist.proto + mode_switch.proto ) diff --git a/xllm/proto/mode_switch.proto b/xllm/proto/mode_switch.proto new file mode 100644 index 0000000000..f22e9e4dd4 --- /dev/null +++ b/xllm/proto/mode_switch.proto @@ -0,0 +1,38 @@ +syntax = "proto3"; + +option go_package = "jd.com/jd-infer/xllm;xllm"; +package xllm.proto; +option cc_generic_services = true; + +// Runtime CP<->DP mode switch. This proto is intentionally independent +// of disagg_pd.proto / worker.proto so that adding it does not invalidate +// the 1116 cpp files that transitively include disagg_pd.pb.h. + +// `target_mode` mirrors xllm::DualParallelArgs::Mode: +// 0 = DP_DECODE (cp_size = 1, dp_size = world) +// 1 = CP_PREFILL (cp_size = world, dp_size = 1) +// +// Caller (xllm_service) is responsible for draining inflight forwards +// on the target instance before issuing the RPC. The engine fans out to +// every local worker and only returns ok=true after every worker +// successfully flipped its DualParallelArgs::Mode flag. +// +// Names are prefixed with InstanceModeSwitch to avoid clashing with +// xllm.proto.SwitchModeRequest defined in worker.proto (master <-> worker). +message InstanceModeSwitchRequest { + int32 target_mode = 1; + int32 timeout_ms = 2; +} + +message InstanceModeSwitchResponse { + bool ok = 1; + // Mode actually held after the call. On failure this is the mode the + // engine rolled back to (which should match what the caller was + // tracking as current_mode). + int32 current_mode = 2; + string error = 3; +} + +service ModeSwitchService { + rpc SwitchMode(InstanceModeSwitchRequest) returns (InstanceModeSwitchResponse); +} diff --git a/xllm/proto/worker.proto b/xllm/proto/worker.proto index 9f5d800e96..52ccb99e6e 100644 --- a/xllm/proto/worker.proto +++ b/xllm/proto/worker.proto @@ -353,6 +353,16 @@ message ForwardOutput { DiTForwardOutput dit_forward_output = 8; } +// SwitchModeRequest -- runtime CP<->DP feature. +// `target_mode` mirrors xllm::DualParallelArgs::Mode: +// 0 = CP_PREFILL (cp_size = world, dp_size = 1) +// 1 = DP_DECODE (cp_size = 1, dp_size = world) +// Caller (engine / scheduler) is responsible for draining inflight +// forwards on the target worker before issuing the RPC. +message SwitchModeRequest { + int32 target_mode = 1; +} + // master create Collective service // to sync all workers. service Collective { @@ -383,4 +393,7 @@ service DistributeWorker { rpc UpdateWeights (UpdateWeightsRequest) returns (Status); rpc StartProfile (Empty) returns (Status); rpc StopProfile (Empty) returns (Status); + // Runtime CP<->DP switch (dual-graph mode). Status.ok=false when the + // worker is single-mode or the target_mode value is invalid. + rpc SwitchMode(SwitchModeRequest) returns (Status); } diff --git a/xllm/server/xllm_server.cpp b/xllm/server/xllm_server.cpp index 09b0cf1ec9..2ef4860487 100644 --- a/xllm/server/xllm_server.cpp +++ b/xllm/server/xllm_server.cpp @@ -204,6 +204,95 @@ bool XllmServer::start(std::unique_ptr service) { return true; } +bool XllmServer::start(std::unique_ptr disagg_pd_service, + std::unique_ptr mode_switch_service) { + // Build a single brpc::Server hosting both services so they share the + // disagg_pd_port (already advertised in instance.rpc_address). Mirrors + // create_server() but adds two services before Start(). + std::string addr(""); + if (!::xllm::ServiceConfig::get_instance().host().empty()) { + addr = + ::xllm::ServiceConfig::get_instance().host() + ":" + + std::to_string(::xllm::DisaggPDConfig::get_instance().disagg_pd_port()); + } + const int port = ::xllm::DisaggPDConfig::get_instance().disagg_pd_port(); + + server_ = std::make_unique(); + if (!configure_generic_server( + server_.get(), + static_cast(disagg_pd_service.get()), + "Disagg PD")) { + return false; + } + if (!configure_generic_server( + server_.get(), + static_cast(mode_switch_service.get()), + "Mode Switch")) { + return false; + } + + brpc::ServerOptions options; + options.idle_timeout_sec = + ::xllm::ServiceConfig::get_instance().rpc_idle_timeout_s(); + options.num_threads = ::xllm::ServiceConfig::get_instance().num_threads(); + butil::EndPoint endpoint; + if (!addr.empty()) { + listen_address_ = addr; + if (butil::str2endpoint(listen_address_.c_str(), &endpoint) < 0) { + LOG(FATAL) << "Convert listen_address_ to endpoint failed: " + << listen_address_; + return false; + } + } else { + endpoint = butil::EndPoint(butil::IP_ANY, port); + } + if (server_->Start(endpoint, &options) != 0) { + LOG(ERROR) << "Failed to start Disagg PD + Mode Switch server on " + << endpoint; + return false; + } + + listen_address_ = + std::string(butil::endpoint2str(server_->listen_address()).c_str()); + listen_port_ = server_->listen_address().port; + LOG(INFO) << "Disagg PD + Mode Switch server started on address " + << server_->listen_address(); + + // Both services live for the brpc server's lifetime. Release ownership + // -- we explicitly tell brpc not to own them via SERVER_DOESNT_OWN_SERVICE, + // so we keep raw pointers alive by leaking the unique_ptrs into static + // storage (tied to the server's lifetime via has_initialized_). + static std::vector> service_keep; + service_keep.emplace_back(std::move(disagg_pd_service)); + service_keep.emplace_back(std::move(mode_switch_service)); + + has_initialized_ = true; + server_->Join(); + return true; +} + +bool XllmServer::start(std::unique_ptr mode_switch_service, + const std::string& addr) { + if (!create_server( + static_cast(mode_switch_service.get()), + addr, + -1, + "Mode Switch")) { + return false; + } + + // brpc holds the service via SERVER_DOESNT_OWN_SERVICE, so keep it alive + // by moving ownership into the join thread's capture. The thread lives for + // the server's lifetime (until stop() breaks the Join()). + running_thread_ = std::make_unique( + [this, service = std::move(mode_switch_service)]() { + has_initialized_ = true; + server_->Join(); + }); + + return true; +} + bool XllmServer::start(std::unique_ptr service) { std::string addr(""); if (!::xllm::ServiceConfig::get_instance().host().empty()) { diff --git a/xllm/server/xllm_server.h b/xllm/server/xllm_server.h index 0b8e9f0196..4787afc01e 100644 --- a/xllm/server/xllm_server.h +++ b/xllm/server/xllm_server.h @@ -18,6 +18,7 @@ limitations under the License. #include "api_service/api_service.h" #include "core/distributed_runtime/collective_service.h" #include "core/distributed_runtime/disagg_pd_service.h" +#include "core/distributed_runtime/mode_switch_service.h" #include "core/distributed_runtime/pd_ooc_service.h" #include "core/distributed_runtime/worker_service.h" #include "core/framework/xtensor/xtensor_dist_service.h" @@ -31,6 +32,18 @@ class XllmServer final { bool start(std::unique_ptr api_service); bool start(std::unique_ptr disagg_pd_service); + // Mounts BOTH disagg_pd_service AND mode_switch_service on the same + // brpc server. Used when --enable_runtime_cp_dp_switch is on so that + // xllm_service can reach ModeSwitchService.SwitchMode through the + // existing instance.rpc_address (no new endpoint exposed). + bool start(std::unique_ptr disagg_pd_service, + std::unique_ptr mode_switch_service); + // Mounts ModeSwitchService alone on a standalone brpc server at `addr`. + // Used on the single-instance (non-disagg) path so a runtime CP<->DP + // flip can be triggered without a full disagg_pd deployment. Non-blocking: + // the server is joined on a background thread. + bool start(std::unique_ptr mode_switch_service, + const std::string& addr); bool start(std::unique_ptr pd_ooc_service); bool start(std::shared_ptr service, const std::string& addr, diff --git a/xllm/xllm.cpp b/xllm/xllm.cpp index e39787d061..855d8d0c4f 100644 --- a/xllm/xllm.cpp +++ b/xllm/xllm.cpp @@ -171,6 +171,9 @@ Options create_options(const std::string& instance_name, bool is_local) { .cfg_size(static_cast(parallel_config.cfg_size())) .vae_size(static_cast(parallel_config.vae_size())) .instance_name(instance_name) + .enable_runtime_cp_dp_switch( + parallel_config.enable_runtime_cp_dp_switch()) + .dual_mode_port_stride(parallel_config.dual_mode_port_stride()) .enable_disagg_pd(disagg_pd_config.enable_disagg_pd()) .enable_pd_ooc(disagg_pd_config.enable_pd_ooc()) .enable_schedule_overlap(scheduler_config.enable_schedule_overlap())