feat(server): add DFlash disk prefix cache for target layer split

[weicj]

Summary

This PR restores disk prefix-cache support for same-backend target layer split, including the DFlash draft path. Target layer split already had in-process prefix-cache support, but the restart-persistent --kv-cache-dir path was still blocked; this PR brings the split+DFlash path back to parity with the single-backend disk-backed prefix restore behavior.

Previously, dflash_server rejected --kv-cache-dir when --target-devices was enabled because the existing disk snapshot format only matched a single backend snapshot. Under target layer split, the live prefix state is sharded across multiple target shards, so a disk hit could not be safely exported, loaded after restart, and rebound back into the split backend.

This change adds a flattened layer-split disk snapshot for the same-backend path. Each shard snapshot is exported into one disk-owned CPU snapshot with shard-prefixed tensor names, then adopted back into the shard-local snapshot slots on cache lookup. For DFlash, the snapshot also persists the draft feature mirror metadata and feature rows, so a restored target prefix can continue speculative decode instead of falling back to a target-only cache state.

Changes

• Adds snapshot_ref() and snapshot_adopt() hooks to the LayerSplitAdapter / LayerSplitBackend boundary, so the server disk-cache layer can save and load snapshots through the generic backend interface.
• Adds Qwen35 target-layer-split disk snapshot export/import:
  • flattens shard-local snapshot tensors into one CPU snapshot for disk persistence;
  • prefixes shard tensors as ls<shard>_<tensor-name> so they can be rebound to the correct target shard on load;
  • stores and restores the global snap_prefill_logits tensor used by prefix restore;
  • stores and restores DFlash feature-mirror state through dflash_feature_meta and dflash_feature_data;
  • validates only the tensors required by each shard's actual cache state;
  • keeps shared adopted ctx / buf ownership from being double-freed.
• Updates disk-cache cold-start lookup so a layout learned from disk can be verified after successful backend adoption, instead of blocking the first restart-time lookup before the live layout is available.
• Updates server placement validation to allow --kv-cache-dir with same-backend --target-devices.

Notes

• Runtime validation passed on dual Pro VII / ROCm 6.3.3 with Qwen3.6-27B Q4 target and the DFlash Qwen3.6 draft, split across hip:0,hip:1. The first server process saved disk prefix cache; after restart, the second process logged [target-split] adopted disk snapshot, disk_hit=true, restore=true, and DFlash speculative decode with accepted draft tokens.
• Mixed-backend target split disk cache remains a follow-up because remote shard snapshots need IPC export/import support.

[cubic-dev-ai]

8 issues found across 44 files

Prompt for AI agents (unresolved issues)

Check if these issues are valid — if so, understand the root cause of each and fix them. If appropriate, use sub-agents to investigate and fix each issue separately.


<file name="server/src/common/dflash_draft_ipc.cpp">

<violation number="1" location="server/src/common/dflash_draft_ipc.cpp:39">
P3: The new IPC env/size helper block duplicates existing helper logic from `qwen35_target_shard_ipc.cpp`; extract/shared utility should be used to avoid behavior drift across IPC clients.</violation>
</file>

<sub>Reply with feedback, questions, or to request a fix.

Re-trigger cubic</sub>

[cubic-dev-ai]

P3: The new IPC env/size helper block duplicates existing helper logic from qwen35_target_shard_ipc.cpp; extract/shared utility should be used to avoid behavior drift across IPC clients.

Prompt for AI agents

Check if this issue is valid — if so, understand the root cause and fix it. At server/src/common/dflash_draft_ipc.cpp, line 39:

<comment>The new IPC env/size helper block duplicates existing helper logic from `qwen35_target_shard_ipc.cpp`; extract/shared utility should be used to avoid behavior drift across IPC clients.</comment>

<file context>
@@ -11,12 +11,71 @@
+    return transport;
+}
+
+bool checked_mul_size(size_t a, size_t b, size_t & out) {
+    if (a != 0 && b > std::numeric_limits<size_t>::max() / a) {
+        return false;
</file context>

[cubic-dev-ai]

2 issues found across 8 files (changes from recent commits).

Prompt for AI agents (unresolved issues)

Check if these issues are valid — if so, understand the root cause of each and fix them. If appropriate, use sub-agents to investigate and fix each issue separately.


<file name="server/src/laguna/laguna_layer_split_adapter.cpp">

<violation number="1" location="server/src/laguna/laguna_layer_split_adapter.cpp:507">
P1: `snapshot_adopt` shares `ctx/buf` across shard snapshots before validation, so failure paths can double-free the adopted snapshot memory.</violation>
</file>

<file name="server/src/common/dflash_draft_ipc_daemon.cpp">

<violation number="1">
P1: Continuing after a malformed `propose_pipe` leaves unread payload bytes in the pipe, which can desynchronize subsequent requests.</violation>
</file>

<sub>Reply with feedback, questions, or to request a fix.

Re-trigger cubic</sub>

[cubic-dev-ai]

P1: snapshot_adopt shares ctx/buf across shard snapshots before validation, so failure paths can double-free the adopted snapshot memory.

Prompt for AI agents

Check if this issue is valid — if so, understand the root cause and fix it. At server/src/laguna/laguna_layer_split_adapter.cpp, line 507:

<comment>`snapshot_adopt` shares `ctx/buf` across shard snapshots before validation, so failure paths can double-free the adopted snapshot memory.</comment>

<file context>
@@ -332,6 +367,203 @@ bool LagunaLayerSplitAdapter::snapshot_restore(int slot) {
+    for (auto & shard_snap : snap.shards) {
+        shard_snap.attn_k.assign(shards_.empty() ? 0 : shards_.front().weights.n_layer, nullptr);
+        shard_snap.attn_v.assign(shards_.empty() ? 0 : shards_.front().weights.n_layer, nullptr);
+        shard_snap.ctx = ctx;
+        shard_snap.buf = buf;
+        shard_snap.cur_pos = cur_pos;
</file context>

[cubic-dev-ai]

P1: Continuing after a malformed propose_pipe leaves unread payload bytes in the pipe, which can desynchronize subsequent requests.

Prompt for AI agents

Check if this issue is valid — if so, understand the root cause and fix it. At server/src/common/dflash_draft_ipc_daemon.cpp, line 454:

<comment>Continuing after a malformed `propose_pipe` leaves unread payload bytes in the pipe, which can desynchronize subsequent requests.</comment>

<file context>
@@ -451,7 +451,7 @@ int run_dflash_draft_ipc_daemon(const char * draft_path,
                              line.c_str());
                 stream_status(stream_fd, -1);
-                break;
+                continue;
             }
             if (!read_exact_fd(payload_fd, noise_embed.data(), bytes)) {
</file context>

Suggested change

	#include <cstddef>
	break;