diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 3de636f2d..3cfe3e446 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -279,7 +279,7 @@ jobs: -o "$RUNNER_TEMP/hip_smoke" .github/ci/hip_smoke.cpp "$RUNNER_TEMP/hip_smoke" - - name: Build + test ROCmFP formats + - name: Build + test ROCm formats and inference core run: | cmake -S server -B "$RUNNER_TEMP/rocmfp-build" \ -DDFLASH27B_GPU_BACKEND=hip \ @@ -290,11 +290,11 @@ jobs: -DCMAKE_HIP_FLAGS=-DDFLASH_WAVE_SIZE=32 cmake --build "$RUNNER_TEMP/rocmfp-build" \ --target test_rocmfp4 test_rocmfpx test_rocmfp4_hip_tail test_rocmfpx_mmq \ - test_deepseek4_mmid_grouped_cuda \ + test_deepseek4_mmid_grouped_cuda test_recurrent_snapshot test_server_unit \ --parallel 8 ctest --test-dir "$RUNNER_TEMP/rocmfp-build" \ --output-on-failure \ - -R 'rocmfp4_reference|rocmfpx_reference|rocmfp4_hip_tail|rocmfpx_mmq|deepseek4_mmid_grouped_cuda' + -R 'rocmfp4_reference|rocmfpx_reference|rocmfp4_hip_tail|rocmfpx_mmq|deepseek4_mmid_grouped_cuda|recurrent_snapshot|ChainRollbackPolicy' build-windows: name: Build Windows (MSVC + CUDA, library + server targets) diff --git a/README.md b/README.md index bc524efde..ce72c5918 100644 --- a/README.md +++ b/README.md @@ -346,6 +346,8 @@ When compression is on, the request path picks one of three modes automatically, | `DFLASH27B_KV_TQ3=1` | (default) | Preset TQ3_0 K+V (3.5 bpv, fits 256K @ 24 GB) | | `DFLASH27B_KV_Q4=1` | off | Q4_0 K+V (4.5 bpv, legacy, ~128K ceiling) | | `--prefix-cache-slots N` | — | Live prefix-cache slot count | +| `DFLASH_PREFIX_CACHE_SLOTS=N` | `32` | Container-entrypoint equivalent of `--prefix-cache-slots`; the native binary itself uses the CLI flag. | +| `DFLASH_PREFILL_CACHE_SLOTS=N` | `0` | Container-entrypoint equivalent of `--prefill-cache-slots`; the native binary itself uses the CLI flag. | | `--kv-cache-dir ` | — | Persist prefix cache to disk | | `--kv-cache-budget N` | — | On-disk cache size cap | @@ -379,6 +381,8 @@ Pages the attention KV cache through a fixed pool of GPU slots; cold 64-token ch | `--target-gpu N` | `0` | Target GPU index | | `--draft-gpu N` | same as target | Draft GPU index; offload draft to a second GPU | | `--target-devices ` / `--target-layer-split` | single GPU | Layer-split target across GPUs | +| `--target-split-fast-rollback` | off | Qwen35 local layer-split only: enable exact F32 per-token checkpoints and skip accepted-token replay. Adds checkpoint VRAM (~1.65 GiB for the measured Qwen3.6-27B q=16 split). | +| `DFLASH_SPLIT_FAST_ROLLBACK=1` | off | Environment equivalent of `--target-split-fast-rollback`. | | `--draft-ipc-bin ` | — | Out-of-process draft binary (mixed CUDA/HIP) | | `--peer-access` | off | Enable P2P between target GPUs | | `--chunk N` | backend default | Prefill ubatch size | @@ -394,6 +398,7 @@ For MoE targets (`laguna`, `qwen35`/`qwen36`) whose experts don't fit in VRAM. ` | Flag / env | Default | Effect | |---|---|---| | `--spark` | off | One-flag autotune: enable the bounded expert cache, size it from the VRAM target, auto-load and keep persisting a placement profile (`.gguf.spark.csv`). | +| `--spark-slots ` | auto | Explicit expert-cache slots per layer; overrides Spark auto-sizing. | | `--spark-vram ` | whole card | Total VRAM Spark may use; it sizes the hot tier + cache + KV under this cap. | | `DFLASH_SPARK=1` | off | Env equivalent of `--spark`. | | `DFLASH_SPARK_VRAM_MB=N` | — | Env equivalent of `--spark-vram` (in MB). | diff --git a/server/CMakeLists.txt b/server/CMakeLists.txt index 9512598da..1abc4723f 100644 --- a/server/CMakeLists.txt +++ b/server/CMakeLists.txt @@ -1026,6 +1026,22 @@ if(DFLASH27B_TESTS) "${CMAKE_CURRENT_SOURCE_DIR}/scripts/entrypoint.sh") endif() + # CPU-only contract test for the fail-closed layer-split tree boundary. + add_executable(test_qwen35_split_tree_guard + test/test_qwen35_split_tree_guard.cpp) + target_include_directories(test_qwen35_split_tree_guard PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/src/qwen35) + list(APPEND _raw_unit_test_targets test_qwen35_split_tree_guard) + + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_recurrent_snapshot.cpp") + add_executable(test_recurrent_snapshot test/test_recurrent_snapshot.cpp) + target_include_directories(test_recurrent_snapshot PRIVATE + ${DFLASH27B_SRC_INCLUDE_DIRS}) + target_link_libraries(test_recurrent_snapshot PRIVATE + dflash_common ggml ${DFLASH27B_GGML_BACKEND_TARGET}) + list(APPEND _raw_unit_test_targets test_recurrent_snapshot) + endif() + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_server_unit.cpp") set(_server_unit_sources test/test_unit_main.cpp @@ -1170,6 +1186,10 @@ if(DFLASH27B_TESTS) set(_unit_ctest_name draft_topk_cuda) elseif(_unit_target STREQUAL "test_deepseek4_unit") set(_unit_ctest_name deepseek4_unit) + elseif(_unit_target STREQUAL "test_qwen35_split_tree_guard") + set(_unit_ctest_name qwen35_split_tree_guard) + elseif(_unit_target STREQUAL "test_recurrent_snapshot") + set(_unit_ctest_name recurrent_snapshot) endif() add_test(NAME "${_unit_ctest_name}" COMMAND ${_unit_target}) if(_unit_target STREQUAL "test_deepseek4_mmid_grouped_cuda") diff --git a/server/docs/ENVIRONMENT.md b/server/docs/ENVIRONMENT.md index 59de2d292..a85e1fa65 100644 --- a/server/docs/ENVIRONMENT.md +++ b/server/docs/ENVIRONMENT.md @@ -33,6 +33,12 @@ consolidation of this list into CLI flags is tracked as follow-up work. | `DFLASH_MOE_TP_*` / `DFLASH_MOE_HYBRID_PREFILL_EAGER` | unset | BURN-IN: model-neutral names for common heterogeneous-MoE scheduling and kernel policy. Existing `DFLASH_DS4_*` names remain compatibility aliases. | | `DFLASH_MMID_TELEMETRY` | unset | DEBUG: report MUL_MAT_ID dispatch, MMVQ variant, and per-node graph compatibility. | | `DFLASH_KVFLASH` | unset | Prefer the CLI: `--kvflash` (token count or `auto`). | +| `DFLASH_PREFIX_CACHE_SLOTS` | 32 | Container-entrypoint equivalent of `--prefix-cache-slots`; not read directly by the native binary. | +| `DFLASH_PREFILL_CACHE_SLOTS` | 0 | Container-entrypoint equivalent of `--prefill-cache-slots`; not read directly by the native binary. | +| `DFLASH_SPLIT_FAST_ROLLBACK` | unset | OPT-IN: exact F32 checkpoints and replay-free rollback for local qwen35 target layer splits. Prefer `--target-split-fast-rollback`; adds checkpoint VRAM (~1.65 GiB for the measured Qwen3.6-27B q=16 split). | +| `DFLASH_STALL_TOOL_PREFIX` | unset | OPT-IN: recover a stalled tool call by injecting the prepared tool prefix when generation stops after an action suffix. | +| `DFLASH_DS4_SPEC` / `DFLASH_DS4_DRAFT` / `DFLASH_DS4_DRAFT_GPU` | unset | OPT-IN: enable DeepSeek4 DSpark, select its draft GGUF, and optionally select the local drafter GPU. See `DS4.md`. | +| `DFLASH_DS4_CUDA_LAYERS` | auto | Override the DeepSeek4 heterogeneous layer-split heuristic. See `DS4.md`. | ## Full inventory (generated) @@ -53,6 +59,8 @@ consolidation of this list into CLI flags is tracked as follow-up work. - `DFLASH_ADAPTIVE_WIDTH_MIN` - adaptive_verify_width.h - `DFLASH_ADAPTIVE_WIDTH_THETA` - adaptive_verify_width.h - `DFLASH_COLD_THREADS` - moe_expert_compute_cpu.cpp +- `DFLASH_CUDA_MMVQ_MOE_ALIGN_SHARED_IDS` - moe_hybrid_ffn_eval.cpp +- `DFLASH_CUDA_MMVQ_MOE_KERNEL` - moe_hybrid_ffn_eval.cpp - `DFLASH_DISABLE_DRAFT_ATTN` - draft_graph.cpp - `DFLASH_DISABLE_DRAFT_ATTN_GATE` - draft_graph.cpp - `DFLASH_DISABLE_DRAFT_AUX_NORMS` - draft_graph.cpp @@ -64,14 +72,29 @@ consolidation of this list into CLI flags is tracked as follow-up work. - `DFLASH_DRAFT_KV` - laguna_backend.cpp, qwen35_backend.cpp - `DFLASH_DRAFT_PERSIST` - laguna_backend.cpp - `DFLASH_DROP_COLD` - qwen35moe_backend.cpp, qwen35moe_pipelined_decode.cpp +- `DFLASH_DS4_ADAPTIVE_WIDTH` - deepseek4_dspark_spec.cpp +- `DFLASH_DS4_CUDA_LAYERS` - deepseek4_layer_split_adapter.cpp +- `DFLASH_DS4_DENSE_TP_MASK` - deepseek4_loader.cpp +- `DFLASH_DS4_DENSE_TP_STRIX_FRACTION` - deepseek4_loader.cpp +- `DFLASH_DS4_DRAFT` - deepseek4_backend.cpp - `DFLASH_DS4_DRAFT_GPU` - deepseek4_backend.cpp +- `DFLASH_DS4_DSPARK_DEBUG` - deepseek4_graph.cpp +- `DFLASH_DS4_FUSED_VERIFY` - deepseek4_dspark_spec.cpp, deepseek4_loader.cpp - `DFLASH_DS4_HOTNESS_CSV` - deepseek4_backend.cpp - `DFLASH_DS4_MOE_TP` - deepseek4_backend.cpp - `DFLASH_DS4_MOE_TP_GPU` - deepseek4_backend.cpp - `DFLASH_DS4_MOE_TP_INPROC` - deepseek4_backend.cpp +- `DFLASH_DS4_ROUTING_STATS_OUT` - deepseek4_backend.cpp +- `DFLASH_DS4_SEQ_VERIFY` - deepseek4_dspark_spec.cpp +- `DFLASH_DS4_SPEC` - deepseek4_backend.cpp +- `DFLASH_DS4_SPEC_Q` - deepseek4_dspark_spec.cpp - `DFLASH_DS4_TIMING` - deepseek4_backend.cpp, deepseek4_target_shard_ipc_daemon.cpp +- `DFLASH_DS4_TP_CAPTURE_CACHE_SLOTS` - deepseek4_fused_verify.inc +- `DFLASH_DS4_TP_FUSED_CACHE_SLOTS` - deepseek4_fused_verify.inc +- `DFLASH_DS4_TOPK` - deepseek4_graph.cpp - `DFLASH_EXPERT_BUDGET_MB` - deepseek4_backend.cpp, laguna_backend.cpp, qwen35moe_backend.cpp - `DFLASH_EXPERT_BUDGET_PCT` - laguna_backend.cpp +- `DFLASH_FAST_ROLLBACK_THRESHOLD` - chain_rollback_policy.h - `DFLASH_FEATURE_DTYPE` - dflash_feature_ring.cpp - `DFLASH_FP_ALPHA` - http_server.cpp, qwen3_graph.cpp, server_main.cpp - `DFLASH_FP_CHUNK_S` - qwen3_graph.cpp @@ -133,22 +156,34 @@ consolidation of this list into CLI flags is tracked as follow-up work. - `DFLASH_MMQ_SUB_BATCH` - moe_hybrid_ffn_eval.cpp - `DFLASH_MODEL_CARDS_DIR` - model_card.cpp - `DFLASH_MOE_COLD_BACKEND` - deepseek4_loader.cpp +- `DFLASH_MOE_COMPACT_MATERIALIZED` - moe_hybrid_ffn_eval.cpp +- `DFLASH_MOE_DUPLICATE_HOT_ON_COLD` - moe_hybrid_storage.cpp +- `DFLASH_MOE_EXPERT_COMPUTE_DAEMON_TOKEN_LOOP` - moe_expert_compute_ipc.cpp - `DFLASH_MOE_EXPERT_COMPUTE_IPC_BATCH_CAPACITY` - moe_expert_compute_ipc.cpp - `DFLASH_MOE_EXPERT_COMPUTE_IPC_DTYPE` - moe_expert_compute_ipc.cpp +- `DFLASH_MOE_EXPERT_COMPUTE_IPC_GPU` - deepseek4_backend.cpp - `DFLASH_MOE_EXPERT_COMPUTE_IPC_MODE` - moe_hybrid_ffn_eval.cpp - `DFLASH_MOE_EXPERT_COMPUTE_IPC_PROFILE` - moe_expert_compute_ipc.cpp - `DFLASH_MOE_EXPERT_COMPUTE_IPC_SHARED_BYTES` - moe_expert_compute_ipc.cpp - `DFLASH_MOE_EXPERT_COMPUTE_IPC_TRANSPORT` - moe_expert_compute_ipc.cpp - `DFLASH_MOE_EXPERT_COMPUTE_THREADS` - moe_expert_compute_cpu.cpp +- `DFLASH_MOE_EXPERT_MAJOR_GPU_REDUCE` - moe_hybrid_ffn_eval.cpp +- `DFLASH_MOE_EXPERT_MAJOR_PREFILL` - moe_hybrid_ffn_eval.cpp - `DFLASH_MOE_FIXED_SLOT_GRAPHS` - moe_hybrid_ffn_eval.cpp - `DFLASH_MOE_FIXED_SLOT_MAX` - moe_hybrid_ffn_eval.cpp +- `DFLASH_MOE_FULL_COLD_PARALLEL` - moe_hybrid_ffn_eval.cpp +- `DFLASH_MOE_FUSED_COMBINE` - moe_hybrid_ffn_eval.cpp +- `DFLASH_MOE_PREFILL_DEVICE_INPUT` - deepseek4_graph.cpp - `DFLASH_MOE_PREFILL_HOT_SUB_BATCH` - moe_hybrid_ffn_eval.cpp +- `DFLASH_MOE_PREFILL_MASKED_COLD` - moe_hybrid_ffn_eval.cpp - `DFLASH_MOE_PREFILL_PERSISTENT_OWNER_ALLOC` - deepseek4_graph.cpp - `DFLASH_NO_MASK` - laguna_backend.cpp - `DFLASH_NO_MOE_ROUTER_FUSE` - qwen35moe_ffn.cpp - `DFLASH_NO_MOE_SWIGLU_FUSE` - qwen35moe_ffn.cpp - `DFLASH_NO_PREAD` - deepseek4_loader.cpp - `DFLASH_PROF` - prof_env.h +- `DFLASH_PREFILL_CACHE_SLOTS` - scripts/entrypoint.sh (maps to `--prefill-cache-slots`) +- `DFLASH_PREFIX_CACHE_SLOTS` - scripts/entrypoint.sh (maps to `--prefix-cache-slots`) - `DFLASH_QWEN35MOE_CACHE_SLOTS` - qwen35moe_backend.cpp - `DFLASH_QWEN35MOE_HOTNESS` - qwen35moe_backend.cpp - `DFLASH_QWEN35MOE_NEXT_PLACEMENT_OUT` - qwen35moe_backend.cpp @@ -161,8 +196,14 @@ consolidation of this list into CLI flags is tracked as follow-up work. - `DFLASH_QWEN35_NO_KVPAD` - graph_builders.cpp - `DFLASH_SAMPLED_VERIFY` - laguna_backend.cpp, qwen35_backend.cpp - `DFLASH_SHARE_DIR` - http_server.cpp +- `DFLASH_SINGLE_CHAIN_CHECKPOINT_F32` - chain_rollback_policy.h +- `DFLASH_SINGLE_CHAIN_ROLLBACK_DIAG` - chain_rollback_policy.h - `DFLASH_SPARK` - laguna_backend.cpp, qwen35moe_backend.cpp - `DFLASH_SPARK_VRAM_MB` - laguna_backend.cpp, qwen35moe_backend.cpp +- `DFLASH_SPLIT_CAPTURE_SELFTEST` - qwen35_layer_split_dflash_target.cpp +- `DFLASH_SPLIT_CHAIN_ROLLBACK_DIAG` - qwen35_layer_split_dflash_target.cpp, qwen35_target_graph.cpp +- `DFLASH_SPLIT_FAST_ROLLBACK` - chain_rollback_policy.h +- `DFLASH_STALL_TOOL_PREFIX` - http_server.cpp - `DFLASH_SV_DEBUG` - qwen35_backend.cpp - `DFLASH_TARGET_SHARD_IPC_SHARED_BYTES` - target_shard_ipc.cpp - `DFLASH_TARGET_SHARD_IPC_TRANSPORT` - target_shard_ipc.cpp @@ -171,6 +212,8 @@ consolidation of this list into CLI flags is tracked as follow-up work. - `DFLASH_VERIFY_WIDTH` - qwen35moe_backend.cpp - `FAST_ROLLBACK_DIAG` - qwen35_dflash_target.cpp - `HOME` - spark_corpus.cpp +- `LUCE_CUDA_I32_REPEAT` - moe_hybrid_ffn_eval.cpp +- `LUCE_MMVQ_MAX_NCOLS` - deepseek4_backend.cpp - `LUCE_QK_FUSE_LAYERS` - laguna_target_graph.cpp - `LUCE_QK_FUSE_MODE` - laguna_target_graph.cpp - `PFLASH_DRAFTER_EARLY_EXIT_N` - qwen3_graph.cpp diff --git a/server/docs/TOOL_PREFIX_CACHE.md b/server/docs/TOOL_PREFIX_CACHE.md index ad9538a43..dc1e31740 100644 --- a/server/docs/TOOL_PREFIX_CACHE.md +++ b/server/docs/TOOL_PREFIX_CACHE.md @@ -40,8 +40,9 @@ never restores past the stable prefix. No tool-specific flag is required. The native server default enables the in-memory prefix cache with 32 slots. Direct container launches inherit that -default; set `DFLASH_PREFIX_CACHE_SLOTS=N` to override it or explicitly set it -to `0` to disable prefix reuse. +default. Pass `--prefix-cache-slots N` to the native binary, or set +`DFLASH_PREFIX_CACHE_SLOTS=N` through `server/scripts/entrypoint.sh`; use `0` +to disable prefix reuse. ## Reproducible benchmark diff --git a/server/hip_compat/cuda_runtime.h b/server/hip_compat/cuda_runtime.h index 4ca749899..b534d179c 100644 --- a/server/hip_compat/cuda_runtime.h +++ b/server/hip_compat/cuda_runtime.h @@ -29,6 +29,9 @@ using cudaDeviceProp = hipDeviceProp_t; // Error codes #define cudaSuccess hipSuccess #define cudaErrorInvalidValue hipErrorInvalidValue +#define cudaErrorIllegalAddress hipErrorIllegalAddress +#define cudaErrorAssert hipErrorAssert +#define cudaErrorLaunchFailure hipErrorLaunchFailure // Memory functions #define cudaMalloc hipMalloc @@ -49,6 +52,7 @@ using cudaDeviceProp = hipDeviceProp_t; #define cudaStreamSynchronize hipStreamSynchronize #define cudaStreamDefault hipStreamDefault #define cudaStreamNonBlocking hipStreamNonBlocking +#define cudaStreamPerThread hipStreamPerThread // Device functions #define cudaGetDevice hipGetDevice @@ -73,6 +77,8 @@ using cudaDeviceProp = hipDeviceProp_t; // Error checking #define cudaGetLastError hipGetLastError #define cudaGetErrorString hipGetErrorString +#define cudaGetErrorName hipGetErrorName +#define cudaPeekAtLastError hipPeekAtLastError // Launch bounds #define __launch_bounds__ __launch_bounds__ diff --git a/server/src/common/chain_rollback_policy.h b/server/src/common/chain_rollback_policy.h index 57eb1318c..fa94cddc1 100644 --- a/server/src/common/chain_rollback_policy.h +++ b/server/src/common/chain_rollback_policy.h @@ -18,6 +18,13 @@ inline bool env_flag_enabled(const char * name) { return value != nullptr && value[0] != '\0' && std::strcmp(value, "0") != 0; } +// Layer-split fast rollback needs exact F32 per-token checkpoints. Keep the +// feature, its capture work, and its additional memory behind one shared +// opt-in so allocation and runtime dispatch cannot drift apart. +inline bool split_chain_fast_rollback_enabled() { + return env_flag_enabled("DFLASH_SPLIT_FAST_ROLLBACK"); +} + inline ChainRollbackPolicy resolve_chain_rollback_policy() { ChainRollbackPolicy policy; policy.checkpoint_f32 = env_flag_enabled("DFLASH_SINGLE_CHAIN_CHECKPOINT_F32"); diff --git a/server/src/common/gpu_runtime_compat.h b/server/src/common/gpu_runtime_compat.h index 19a7927b0..0cfc9d76b 100644 --- a/server/src/common/gpu_runtime_compat.h +++ b/server/src/common/gpu_runtime_compat.h @@ -14,11 +14,16 @@ #define cudaDeviceSynchronize hipDeviceSynchronize #define cudaErrorPeerAccessAlreadyEnabled hipErrorPeerAccessAlreadyEnabled #define cudaErrorPeerAccessNotEnabled hipErrorPeerAccessNotEnabled +#define cudaErrorIllegalAddress hipErrorIllegalAddress +#define cudaErrorAssert hipErrorAssert +#define cudaErrorLaunchFailure hipErrorLaunchFailure #define cudaError_t hipError_t #define cudaFree hipFree #define cudaGetDeviceCount hipGetDeviceCount #define cudaGetErrorString hipGetErrorString +#define cudaGetErrorName hipGetErrorName #define cudaGetLastError hipGetLastError +#define cudaPeekAtLastError hipPeekAtLastError #define cudaMalloc hipMalloc #define cudaMemcpy2D hipMemcpy2D #define cudaMemcpy2DAsync hipMemcpy2DAsync @@ -33,6 +38,7 @@ #define cudaSetDevice hipSetDevice #define cudaStreamSynchronize hipStreamSynchronize #define cudaStream_t hipStream_t +#define cudaStreamPerThread hipStreamPerThread #define cudaSuccess hipSuccess #define cudaDeviceProp hipDeviceProp_t #define cudaDeviceReset hipDeviceReset diff --git a/server/src/internal.h b/server/src/internal.h index 6f5a37405..ceadfc18b 100644 --- a/server/src/internal.h +++ b/server/src/internal.h @@ -392,7 +392,9 @@ struct TargetCache { // persistent cache memory (not tracked by the per-call gallocr), matching // SGLang's mamba_caches.intermediate_ssm / intermediate_conv_window pattern. // - // ssm_intermediate: [S_v, S_v, H_v, max_q_len] f32, one per delta layer. + // ssm_intermediate: [S_v, S_v, H_v, max_q_len], checkpoint dtype + // (Q8_0 for direct caches, F16 for migrated single-target caches, or + // F32 for opt-in exact rollback), one per delta layer. // Element t on axis 3 holds the DeltaNet recurrent state after // processing verify token t. Spec decode commits t = commit_n - 1. // conv_input_cache: [(kernel-1) + max_q_len, conv_channels] f32, one per @@ -420,9 +422,9 @@ struct TargetCache { }; // Snapshot the current SSM+conv state into TargetCache::*_snap tensors. -void snapshot_ssm_state(TargetCache & c); +bool snapshot_ssm_state(TargetCache & c, ggml_backend_t backend); // Restore the SSM+conv state from the snapshot. -void restore_ssm_state(TargetCache & c); +bool restore_ssm_state(TargetCache & c, ggml_backend_t backend); // Allocate rollback snapshot tensors mirroring live ssm/conv state (MoE path). bool ensure_ssm_snapshot(TargetCache & c, ggml_backend_t backend); @@ -536,6 +538,9 @@ bool create_target_cache(const TargetWeights & w, bool prefill_only = false, int ctx_alloc = 0); +// `f32_ssm_intermediates` enables exact per-token checkpoints for the opt-in +// layer-split fast rollback path. The default preserves the established Q8_0 +// allocation and avoids its ~1.65 GiB incremental memory cost. bool create_target_cache_partial(const TargetWeights & w, int max_ctx, int max_verify_tokens, @@ -545,7 +550,8 @@ bool create_target_cache_partial(const TargetWeights & w, int layer_begin, int layer_end, bool allocate_target_feat, - int ctx_alloc = 0); + int ctx_alloc = 0, + bool f32_ssm_intermediates = false); void free_target_cache(TargetCache & c); diff --git a/server/src/qwen35/layer_split_daemon_loop.cpp b/server/src/qwen35/layer_split_daemon_loop.cpp index fe2862756..488fd28e0 100644 --- a/server/src/qwen35/layer_split_daemon_loop.cpp +++ b/server/src/qwen35/layer_split_daemon_loop.cpp @@ -6,6 +6,7 @@ #include "layer_split_forward.h" // free_qwen35_layer_split_shards #include "dflash_feature_ring.h" #include "common/io_utils.h" +#include "common/chain_rollback_policy.h" #include "common/sampler.h" #include "common/layer_split_utils.h" #include "common/gguf_inspect.h" @@ -55,7 +56,11 @@ int run_layer_split_daemon(const LayerSplitDaemonConfig & cfg) { shard.backend, shard.cache, /*prefill_only=*/!cfg.run_dflash, shard.layer_begin, shard.layer_end, - /*allocate_target_feat=*/false)) { + /*allocate_target_feat=*/false, + /*ctx_alloc=*/0, + /*f32_ssm_intermediates=*/ + cfg.run_dflash && + split_chain_fast_rollback_enabled())) { std::fprintf(stderr, "target-split load/cache gpu=%d: %s\n", shard.gpu, dflash27b_last_error()); free_qwen35_layer_split_shards(shards); diff --git a/server/src/qwen35/layer_split_forward.cpp b/server/src/qwen35/layer_split_forward.cpp index 55440d603..a66383cc9 100644 --- a/server/src/qwen35/layer_split_forward.cpp +++ b/server/src/qwen35/layer_split_forward.cpp @@ -15,12 +15,35 @@ #include #include +#include #include namespace dflash::common { namespace { +int qwen35_delta_index_for_layer(const TargetWeights & w, int layer_idx) { + int delta_idx = 0; + for (int il = 0; il < layer_idx; ++il) { + if (((il + 1) % w.full_attention_interval) != 0) ++delta_idx; + } + return delta_idx; +} + +bool qwen35_split_ssm_rollback_storage_present( + const Qwen35LayerSplitShard & shard, int layer_idx) { + const bool is_attn = (((layer_idx + 1) % shard.weights.full_attention_interval) == 0); + if (is_attn) return false; + const int delta_idx = qwen35_delta_index_for_layer(shard.weights, layer_idx); + return delta_idx >= 0 && + delta_idx < (int)shard.cache.ssm_intermediate.size() && + delta_idx < (int)shard.cache.conv_input_cache.size() && + shard.cache.ssm_intermediate[(size_t)delta_idx] && + shard.cache.conv_input_cache[(size_t)delta_idx] && + shard.cache.ssm_intermediate[(size_t)delta_idx]->type == GGML_TYPE_F32 && + shard.cache.conv_input_cache[(size_t)delta_idx]->type == GGML_TYPE_F32; +} + bool fill_qwen35_kvflash_inputs( StepGraph & sg, const TargetWeights & w, @@ -165,7 +188,9 @@ bool run_qwen35_layer_split_forward( std::vector * logits_out, DFlashDraftIpcClient * remote_draft, ggml_type activation_type, - KvFlashPager * kvflash) { + KvFlashPager * kvflash, + bool capture_ssm_intermediates, + Qwen35SplitCaptureStats * capture_stats) { if (shards.empty() || tokens.empty()) return false; const int hidden = shards.front().weights.n_embd; const int vocab = shards.front().weights.n_vocab; @@ -243,8 +268,30 @@ bool run_qwen35_layer_split_forward( } const bool is_attn = (((il + 1) % embed_source.full_attention_interval) == 0); + const bool owns_ssm_rollback_layer = !is_attn; const int capture_idx = target_capture_index(embed_source.capture_layer_ids, embed_source.n_capture_layers, il); + const bool owns_feature_tap = capture_idx >= 0; + const bool has_ssm_rollback_storage = owns_ssm_rollback_layer && + qwen35_split_ssm_rollback_storage_present(*shard, il); + const bool capture_this_layer = capture_ssm_intermediates && has_ssm_rollback_storage; + if (capture_stats) { + const size_t shard_idx = (size_t)(shard - shards.data()); + if (owns_ssm_rollback_layer && shard_idx < capture_stats->layers_owned_per_shard.size()) { + capture_stats->layers_owned_per_shard[shard_idx]++; + } + if (owns_feature_tap && shard_idx < capture_stats->feature_taps_owned_per_shard.size()) { + capture_stats->feature_taps_owned_per_shard[shard_idx]++; + } + } + if (capture_ssm_intermediates && owns_ssm_rollback_layer && !has_ssm_rollback_storage) { + if (capture_stats) capture_stats->missing_owner_count++; + std::fprintf(stderr, + "[target-split][capture] missing owner ssm rollback storage layer=%d gpu=%d\n", + il, shard->gpu); + activation_pair_free(acts); + return false; + } for (int start = 0; start < n_tokens_total; start += ubatch) { const int n = std::min(ubatch, n_tokens_total - start); const int kv_start = base_pos + start; @@ -258,8 +305,8 @@ bool run_qwen35_layer_split_forward( if (!build_layer_step(shard->layer_graph, shard->weights, shard->cache, shard->backend, il, act_in, act_out, start, n, kv_start, with_mask, - /*capture=*/false, fa_window, kq_stride_pad, - kvflash != nullptr)) { + /*capture=*/capture_this_layer, fa_window, kq_stride_pad, + kvflash != nullptr, /*tree_mode=*/false)) { std::fprintf(stderr, "target-split build layer=%d @%d gpu=%d\n", il, start, shard->gpu); activation_pair_free(acts); @@ -300,6 +347,13 @@ bool run_qwen35_layer_split_forward( activation_pair_free(acts); return false; } + if (capture_this_layer && capture_stats) { + const size_t shard_idx = (size_t)(shard - shards.data()); + capture_stats->enabled++; + if (shard_idx < capture_stats->slots_written_per_shard.size()) { + capture_stats->slots_written_per_shard[shard_idx] += (uint64_t)n; + } + } if ((feature_ring || remote_draft) && capture_idx >= 0) { if (feature_ring && !copy_capture_slice_to_draft_ring(*feature_ring, capture_idx, @@ -321,6 +375,12 @@ bool run_qwen35_layer_split_forward( activation_pair_free(acts); return false; } + if (capture_stats) { + const size_t shard_idx = (size_t)(shard - shards.data()); + if (shard_idx < capture_stats->feature_slots_written_per_shard.size()) { + capture_stats->feature_slots_written_per_shard[shard_idx] += (uint64_t)n; + } + } } } std::swap(act_in, act_out); @@ -362,7 +422,10 @@ bool run_qwen35_layer_split_layers_from_activation( DraftFeatureMirror * feature_ring, DFlashDraftIpcClient * remote_draft, KvFlashPager * kvflash, - bool kvflash_preallocated = false) { + bool capture_ssm_intermediates, + Qwen35SplitCaptureStats * capture_stats, + bool kvflash_preallocated = false, + const Qwen35SplitTreeInputs * tree_inputs = nullptr) { if (shards.empty() || !acts.a || !acts.b || n_tokens_total <= 0) return false; if (kvflash && fa_window > 0) { std::fprintf(stderr, @@ -402,8 +465,29 @@ bool run_qwen35_layer_split_layers_from_activation( } const bool is_attn = (((il + 1) % shard->weights.full_attention_interval) == 0); + const bool owns_ssm_rollback_layer = !is_attn; const int capture_idx = target_capture_index(shard->weights.capture_layer_ids, shard->weights.n_capture_layers, il); + const bool owns_feature_tap = capture_idx >= 0; + const bool has_ssm_rollback_storage = owns_ssm_rollback_layer && + qwen35_split_ssm_rollback_storage_present(*shard, il); + const bool capture_this_layer = capture_ssm_intermediates && has_ssm_rollback_storage; + if (capture_stats) { + const size_t shard_idx = (size_t)(shard - shards.data()); + if (owns_ssm_rollback_layer && shard_idx < capture_stats->layers_owned_per_shard.size()) { + capture_stats->layers_owned_per_shard[shard_idx]++; + } + if (owns_feature_tap && shard_idx < capture_stats->feature_taps_owned_per_shard.size()) { + capture_stats->feature_taps_owned_per_shard[shard_idx]++; + } + } + if (capture_ssm_intermediates && owns_ssm_rollback_layer && !has_ssm_rollback_storage) { + if (capture_stats) capture_stats->missing_owner_count++; + std::fprintf(stderr, + "[target-split][capture] missing owner ssm rollback storage layer=%d gpu=%d\n", + il, shard->gpu); + return false; + } for (int start = 0; start < n_tokens_total; start += ubatch) { const int n = std::min(ubatch, n_tokens_total - start); const int kv_start = base_pos + start; @@ -417,8 +501,9 @@ bool run_qwen35_layer_split_layers_from_activation( if (!build_layer_step(shard->layer_graph, shard->weights, shard->cache, shard->backend, il, act_in, act_out, start, n, kv_start, with_mask, - /*capture=*/false, fa_window, kq_stride_pad, - kvflash != nullptr)) { + /*capture=*/capture_this_layer, fa_window, kq_stride_pad, + kvflash != nullptr, + /*tree_mode=*/tree_inputs != nullptr)) { std::fprintf(stderr, "target-split build layer=%d @%d gpu=%d\n", il, start, shard->gpu); return false; @@ -451,12 +536,30 @@ bool run_qwen35_layer_split_layers_from_activation( ggml_backend_tensor_set(shard->layer_graph.attn_mask, mask_buf.data(), 0, sizeof(uint16_t) * mask_buf.size()); } + if (tree_inputs && shard->layer_graph.parent_ids) { + if (!tree_inputs->parent_ids || tree_inputs->n_actual != n_tokens_total) { + std::fprintf(stderr, + "target-split tree verify missing parent ids n_actual=%d expected=%d\n", + tree_inputs ? tree_inputs->n_actual : 0, n_tokens_total); + return false; + } + ggml_backend_tensor_set(shard->layer_graph.parent_ids, + tree_inputs->parent_ids + start, 0, + sizeof(int32_t) * (size_t)n); + } auto st = ggml_backend_graph_compute(shard->backend, shard->layer_graph.gf); if (st != GGML_STATUS_SUCCESS) { std::fprintf(stderr, "target-split compute layer=%d @%d gpu=%d status=%d\n", il, start, shard->gpu, (int)st); return false; } + if (capture_this_layer && capture_stats) { + const size_t shard_idx = (size_t)(shard - shards.data()); + capture_stats->enabled++; + if (shard_idx < capture_stats->slots_written_per_shard.size()) { + capture_stats->slots_written_per_shard[shard_idx] += (uint64_t)n; + } + } if ((captures_out || feature_ring || remote_draft) && capture_idx >= 0) { if (captures_out) { Qwen35TargetCaptureSlice capture; @@ -490,6 +593,12 @@ bool run_qwen35_layer_split_layers_from_activation( il, capture_idx, shard->gpu); return false; } + if (capture_stats) { + const size_t shard_idx = (size_t)(shard - shards.data()); + if (shard_idx < capture_stats->feature_slots_written_per_shard.size()) { + capture_stats->feature_slots_written_per_shard[shard_idx] += (uint64_t)n; + } + } } } std::swap(act_in, act_out); @@ -516,11 +625,14 @@ bool run_qwen35_layer_split_forward_from_activation( std::vector * logits_out, std::vector * captures_out, KvFlashPager * kvflash, - bool kvflash_preallocated) { + bool kvflash_preallocated, + const Qwen35SplitTreeInputs * tree_inputs, + bool capture_ssm_intermediates) { if (!run_qwen35_layer_split_layers_from_activation( shards, acts, base_pos, n_tokens_total, ubatch, kq_stride_pad, fa_window, captures_out, nullptr, nullptr, kvflash, - kvflash_preallocated)) { + capture_ssm_intermediates, + /*capture_stats=*/nullptr, kvflash_preallocated, tree_inputs)) { return false; } @@ -547,6 +659,40 @@ bool run_qwen35_layer_split_forward_from_activation( return true; } +bool run_qwen35_layer_split_tree_verify_from_activation( + std::vector & shards, + ActivationPair & acts, + int base_pos, + int n_tokens_total, + int ubatch, + int & last_tok, + int kq_stride_pad, + int fa_window, + const Qwen35SplitTreeInputs & tree_inputs, + std::vector * argmax_out, + std::vector * logits_out, + std::vector * captures_out, + KvFlashPager * kvflash, + bool kvflash_preallocated) { + if (tree_inputs.n_actual != n_tokens_total || !tree_inputs.parent_ids) { + std::fprintf(stderr, + "target-split tree verify invalid inputs n_actual=%d expected=%d parent_ids=%p\n", + tree_inputs.n_actual, n_tokens_total, + static_cast(tree_inputs.parent_ids)); + return false; + } + if (tree_inputs.visibility) { + std::fprintf(stderr, + "target-split tree verify visibility mask is not yet wired in split prewindow patch\n"); + return false; + } + return run_qwen35_layer_split_forward_from_activation( + shards, acts, base_pos, n_tokens_total, ubatch, last_tok, + kq_stride_pad, fa_window, argmax_out, logits_out, captures_out, + kvflash, kvflash_preallocated, + &tree_inputs, /*capture_ssm_intermediates=*/true); +} + bool run_qwen35_mixed_layer_split_forward( std::vector & local_shards, Qwen35TargetShardIpcClient & remote_shard, @@ -561,7 +707,9 @@ bool run_qwen35_mixed_layer_split_forward( std::vector * logits_out, DraftFeatureMirror * feature_ring, DFlashDraftIpcClient * remote_draft, - KvFlashPager * kvflash) { + KvFlashPager * kvflash, + bool capture_ssm_intermediates, + Qwen35SplitCaptureStats * capture_stats) { if (!remote_shard.active() || tokens.empty() || local_shards.empty() || local_shards.front().layer_begin != 0 || local_shards.back().layer_end <= 0) { @@ -611,7 +759,8 @@ bool run_qwen35_mixed_layer_split_forward( if (!run_qwen35_layer_split_layers_from_activation( local_shards, acts, base_pos, n_tokens_total, ubatch, kq_stride_pad, fa_window, nullptr, feature_ring, remote_draft, - kvflash, kvflash != nullptr)) { + kvflash, capture_ssm_intermediates, capture_stats, + kvflash != nullptr)) { activation_pair_free(acts); return false; } diff --git a/server/src/qwen35/layer_split_forward.h b/server/src/qwen35/layer_split_forward.h index 8d436b1b9..c9ddb3b9a 100644 --- a/server/src/qwen35/layer_split_forward.h +++ b/server/src/qwen35/layer_split_forward.h @@ -64,7 +64,9 @@ bool run_qwen35_layer_split_forward( std::vector * logits_out = nullptr, DFlashDraftIpcClient * remote_draft = nullptr, ggml_type activation_type = GGML_TYPE_F32, - KvFlashPager * kvflash = nullptr); + KvFlashPager * kvflash = nullptr, + bool capture_ssm_intermediates = false, + Qwen35SplitCaptureStats * capture_stats = nullptr); bool run_qwen35_layer_split_forward_from_activation( std::vector & shards, @@ -79,6 +81,24 @@ bool run_qwen35_layer_split_forward_from_activation( std::vector * logits_out = nullptr, std::vector * captures_out = nullptr, KvFlashPager * kvflash = nullptr, + bool kvflash_preallocated = false, + const Qwen35SplitTreeInputs * tree_inputs = nullptr, + bool capture_ssm_intermediates = false); + +bool run_qwen35_layer_split_tree_verify_from_activation( + std::vector & shards, + ActivationPair & acts, + int base_pos, + int n_tokens_total, + int ubatch, + int & last_tok, + int kq_stride_pad, + int fa_window, + const Qwen35SplitTreeInputs & tree_inputs, + std::vector * argmax_out = nullptr, + std::vector * logits_out = nullptr, + std::vector * captures_out = nullptr, + KvFlashPager * kvflash = nullptr, bool kvflash_preallocated = false); bool run_qwen35_mixed_layer_split_forward( @@ -95,7 +115,9 @@ bool run_qwen35_mixed_layer_split_forward( std::vector * logits_out = nullptr, DraftFeatureMirror * feature_ring = nullptr, DFlashDraftIpcClient * remote_draft = nullptr, - KvFlashPager * kvflash = nullptr); + KvFlashPager * kvflash = nullptr, + bool capture_ssm_intermediates = false, + Qwen35SplitCaptureStats * capture_stats = nullptr); // Free all shards (weights, cache, backend). void free_qwen35_layer_split_shards(std::vector & shards); diff --git a/server/src/qwen35/layer_split_types.h b/server/src/qwen35/layer_split_types.h index 6e222d767..997cce244 100644 --- a/server/src/qwen35/layer_split_types.h +++ b/server/src/qwen35/layer_split_types.h @@ -16,6 +16,8 @@ #include "ggml-backend.h" #include +#include +#include namespace dflash::common { @@ -27,4 +29,31 @@ struct Qwen35LayerSplitShard : LayerSplitShardMeta { StepGraph layer_graph; }; +struct Qwen35SplitTreeInputs { + const int32_t * parent_ids = nullptr; + const uint8_t * visibility = nullptr; + int n_actual = 0; + int committed = 0; +}; + +struct Qwen35SplitCaptureStats { + uint64_t requested = 0; + uint64_t enabled = 0; + uint64_t missing_owner_count = 0; + std::vector layers_owned_per_shard; + std::vector slots_written_per_shard; + std::vector feature_taps_owned_per_shard; + std::vector feature_slots_written_per_shard; + + void reset(size_t n_shards) { + requested = 0; + enabled = 0; + missing_owner_count = 0; + layers_owned_per_shard.assign(n_shards, 0); + slots_written_per_shard.assign(n_shards, 0); + feature_taps_owned_per_shard.assign(n_shards, 0); + feature_slots_written_per_shard.assign(n_shards, 0); + } +}; + } // namespace dflash::common diff --git a/server/src/qwen35/qwen35_dflash_target.cpp b/server/src/qwen35/qwen35_dflash_target.cpp index 19be74052..18fe94b46 100644 --- a/server/src/qwen35/qwen35_dflash_target.cpp +++ b/server/src/qwen35/qwen35_dflash_target.cpp @@ -504,13 +504,11 @@ bool Qwen35DFlashTarget::rollback_to_tree( } bool Qwen35DFlashTarget::snapshot_kv() { - snapshot_ssm_state(cache_); - return true; + return snapshot_ssm_state(cache_, backend_); } bool Qwen35DFlashTarget::restore_kv() { - restore_ssm_state(cache_); - return true; + return restore_ssm_state(cache_, backend_); } bool Qwen35DFlashTarget::supports_fast_rollback() const { diff --git a/server/src/qwen35/qwen35_layer_split_adapter.cpp b/server/src/qwen35/qwen35_layer_split_adapter.cpp index fa918e7ef..a323e6a93 100644 --- a/server/src/qwen35/qwen35_layer_split_adapter.cpp +++ b/server/src/qwen35/qwen35_layer_split_adapter.cpp @@ -3,6 +3,7 @@ #include "qwen35_layer_split_adapter.h" #include "common/backend_precision.h" +#include "common/chain_rollback_policy.h" #include "common/dflash_spec_decode.h" #include "common/gguf_inspect.h" #include "common/layer_split_utils.h" @@ -90,7 +91,10 @@ bool Qwen35LayerSplitAdapter::init() { /*prefill_only=*/!cfg_.run_dflash, shard.layer_begin, shard.layer_end, /*allocate_target_feat=*/false, - kvflash_tokens_)) { + kvflash_tokens_, + /*f32_ssm_intermediates=*/ + cfg_.run_dflash && + split_chain_fast_rollback_enabled())) { std::fprintf(stderr, "[target-split] cache gpu=%d: %s\n", shard.gpu, dflash27b_last_error()); return false; @@ -543,7 +547,8 @@ bool Qwen35LayerSplitAdapter::prefill(const std::vector & prompt, &prefill_last_logits_, (cfg_.run_dflash && !remote_draft_.active()) ? &feature_ring_ : nullptr, remote_draft_.active() ? &remote_draft_ : nullptr, - kvflash_active() ? &kvflash_pager_ : nullptr); + kvflash_active() ? &kvflash_pager_ : nullptr, + /*capture_ssm_intermediates=*/false, /*capture_stats=*/nullptr); if (ok && kvflash_active()) { kvflash_sync_history(prompt, base_pos); kvflash_pager_.zero_free_blocks(); @@ -558,7 +563,8 @@ bool Qwen35LayerSplitAdapter::prefill(const std::vector & prompt, &prefill_last_logits_, cfg_.run_dflash ? &remote_draft_ : nullptr, activation_type_, - kvflash_active() ? &kvflash_pager_ : nullptr); + kvflash_active() ? &kvflash_pager_ : nullptr, + /*capture_ssm_intermediates=*/false, /*capture_stats=*/nullptr); if (ok && kvflash_active()) { kvflash_sync_history(prompt, base_pos); kvflash_pager_.zero_free_blocks(); @@ -1290,7 +1296,8 @@ bool Qwen35LayerSplitAdapter::decode_ar( logits_out, (cfg_.run_dflash && !remote_draft_.active()) ? &feature_ring_ : nullptr, remote_draft_.active() ? &remote_draft_ : nullptr, - kvflash_active() ? &kvflash_pager_ : nullptr); + kvflash_active() ? &kvflash_pager_ : nullptr, + /*capture_ssm_intermediates=*/false, /*capture_stats=*/nullptr); } return run_qwen35_layer_split_forward( shards_, shards_.front().weights, one, pos, 1, next_tok, @@ -1300,7 +1307,8 @@ bool Qwen35LayerSplitAdapter::decode_ar( logits_out, cfg_.run_dflash ? &remote_draft_ : nullptr, activation_type_, - kvflash_active() ? &kvflash_pager_ : nullptr); + kvflash_active() ? &kvflash_pager_ : nullptr, + /*capture_ssm_intermediates=*/false, /*capture_stats=*/nullptr); }, [&](int tok) { return is_eos_tok(tok, w); }, out_tokens, io); diff --git a/server/src/qwen35/qwen35_layer_split_dflash_target.cpp b/server/src/qwen35/qwen35_layer_split_dflash_target.cpp index 9228fa99c..7b0b8a4e9 100644 --- a/server/src/qwen35/qwen35_layer_split_dflash_target.cpp +++ b/server/src/qwen35/qwen35_layer_split_dflash_target.cpp @@ -1,14 +1,75 @@ // Qwen35LayerSplitDFlashTarget — DFlashTarget adapter for qwen35 layer-split. #include "qwen35_layer_split_dflash_target.h" +#include "qwen35_layer_split_tree_guard.h" #include "internal.h" #include "graph_builders.h" #include "step_graph.h" #include "common/kvflash_pager.h" +#include "common/chain_rollback_policy.h" +#include "common/gpu_runtime_compat.h" + +#include +#include +#include +#include +#include + +using to_fp32_cuda_t = void (*)(const void *, float *, int64_t, cudaStream_t); +extern "C++" to_fp32_cuda_t ggml_get_to_fp32_cuda(ggml_type type); namespace dflash::common { +namespace { + +static bool split_rollback_context_fatal(cudaError_t err) { + return err == cudaErrorIllegalAddress || + err == cudaErrorAssert || + err == cudaErrorLaunchFailure; +} + +static bool split_fast_rollback_storage_ready( + const std::vector & shards) { + bool found_owned_delta_layer = false; + for (const auto & shard : shards) { + const auto & cache = shard.cache; + const auto & w = shard.weights; + int dn_idx = 0; + for (int il = 0; il < w.n_layer; ++il) { + const bool is_attn = (((il + 1) % w.full_attention_interval) == 0); + if (is_attn) continue; + const bool owns_layer = il >= shard.layer_begin && il < shard.layer_end; + if (!owns_layer) { + ++dn_idx; + continue; + } + found_owned_delta_layer = true; + if (dn_idx >= (int)cache.ssm_state.size() || + dn_idx >= (int)cache.conv_state.size() || + dn_idx >= (int)cache.ssm_intermediate.size() || + dn_idx >= (int)cache.conv_input_cache.size()) { + return false; + } + const ggml_tensor * ssm_state = cache.ssm_state[(size_t)dn_idx]; + const ggml_tensor * conv_state = cache.conv_state[(size_t)dn_idx]; + const ggml_tensor * ssm_inter = cache.ssm_intermediate[(size_t)dn_idx]; + const ggml_tensor * conv_input = cache.conv_input_cache[(size_t)dn_idx]; + if (!ssm_state || !conv_state || !ssm_inter || !conv_input || + ssm_state->type != GGML_TYPE_F32 || + conv_state->type != GGML_TYPE_F32 || + ssm_inter->type != GGML_TYPE_F32 || + conv_input->type != GGML_TYPE_F32) { + return false; + } + ++dn_idx; + } + } + return found_owned_delta_layer; +} + +} // namespace + Qwen35LayerSplitDFlashTarget::~Qwen35LayerSplitDFlashTarget() { step_graph_destroy(proj_sg_); } @@ -42,23 +103,107 @@ bool Qwen35LayerSplitDFlashTarget::verify_batch( std::vector * all_argmax, bool capture_ssm_intermediates) { if (shards_.empty()) return false; + if (rollback_poisoned_) { + if (std::getenv("DFLASH_SPLIT_CHAIN_ROLLBACK_DIAG") != nullptr) { + std::fprintf(stderr, + "[target-split][rollback-poison] verify_batch_refused_while_poisoned=1 base_pos=%d n_tokens=%zu\n", + base_pos, tokens.size()); + } + return false; + } + Qwen35SplitCaptureStats capture_stats; + const bool capture_selftest = std::getenv("DFLASH_SPLIT_CAPTURE_SELFTEST") != nullptr; + const bool capture_diag = std::getenv("DFLASH_SPLIT_CHAIN_ROLLBACK_DIAG") != nullptr; + const bool local_only = !(remote_target_shard_ && remote_target_shard_->active()); + const bool storage_ready = split_fast_rollback_storage_ready(shards_); + const bool capture_enabled = capture_ssm_intermediates && + split_chain_fast_rollback_enabled() && local_only && storage_ready; + capture_stats.reset(shards_.size()); + if (capture_ssm_intermediates) capture_stats.requested++; + + bool ok = false; if (remote_target_shard_ && remote_target_shard_->active()) { - return run_qwen35_mixed_layer_split_forward( + ok = run_qwen35_mixed_layer_split_forward( shards_, *remote_target_shard_, shards_.front().weights, tokens, base_pos, (int)tokens.size(), last_tok, kq_stride_pad_, fa_window_, all_argmax, /*logits_out=*/nullptr, feature_ring_, remote_draft_, - kvflash_); + kvflash_, capture_enabled, &capture_stats); + } else { + ok = run_qwen35_layer_split_forward( + shards_, shards_.front().weights, tokens, base_pos, (int)tokens.size(), + last_tok, kq_stride_pad_, fa_window_, + feature_ring_, + all_argmax, /*logits_out=*/nullptr, remote_draft_, + /*activation_type=*/GGML_TYPE_F32, kvflash_, + capture_enabled, &capture_stats); } - return run_qwen35_layer_split_forward( - shards_, shards_.front().weights, tokens, base_pos, (int)tokens.size(), - last_tok, kq_stride_pad_, fa_window_, - feature_ring_, - all_argmax, /*logits_out=*/nullptr, remote_draft_, - /*activation_type=*/GGML_TYPE_F32, kvflash_); + + if (capture_selftest || capture_diag) { + std::fprintf(stderr, + "[target-split][capture] split_capture_requested=%llu split_capture_enabled=%llu split_capture_missing_owner_count=%llu ok=%d\n", + (unsigned long long)capture_stats.requested, + (unsigned long long)capture_stats.enabled, + (unsigned long long)capture_stats.missing_owner_count, ok ? 1 : 0); + for (size_t i = 0; i < capture_stats.layers_owned_per_shard.size(); ++i) { + const int gpu = i < shards_.size() ? shards_[i].gpu : -1; + std::fprintf(stderr, + "[target-split][capture] shard=%zu gpu=%d split_capture_layers_owned_per_shard=%llu split_capture_slots_written_per_shard=%llu\n", + i, gpu, + (unsigned long long)capture_stats.layers_owned_per_shard[i], + (unsigned long long)capture_stats.slots_written_per_shard[i]); + const unsigned long long feature_owned = + i < capture_stats.feature_taps_owned_per_shard.size() + ? (unsigned long long)capture_stats.feature_taps_owned_per_shard[i] : 0ULL; + const unsigned long long feature_written = + i < capture_stats.feature_slots_written_per_shard.size() + ? (unsigned long long)capture_stats.feature_slots_written_per_shard[i] : 0ULL; + std::fprintf(stderr, + "[target-split][capture] shard=%zu gpu=%d split_feature_taps_owned_per_shard=%llu split_feature_slots_written_per_shard=%llu target_feat_absent_by_config=true\n", + i, gpu, feature_owned, feature_written); + } + } + bool capture_gate_ok = false; + bool feature_gate_ok = true; + bool shard_gate_ok = !capture_stats.layers_owned_per_shard.empty(); + if (capture_enabled) { + capture_gate_ok = ok && capture_stats.enabled > 0 && + capture_stats.missing_owner_count == 0; + for (size_t i = 0; i < capture_stats.layers_owned_per_shard.size(); ++i) { + if (capture_stats.layers_owned_per_shard[i] > 0 && + capture_stats.slots_written_per_shard[i] == 0) { + shard_gate_ok = false; + } + if (i < capture_stats.feature_taps_owned_per_shard.size() && + capture_stats.feature_taps_owned_per_shard[i] > 0 && + capture_stats.feature_slots_written_per_shard[i] == 0) { + feature_gate_ok = false; + } + } + // Stage 3 tree verification depends on capture storage/config validation, + // not on whether the most recent forward pass requested a fresh capture. + // Once the current target instance proves its capture configuration, keep + // that gate sticky across no-capture verification/replay passes. A future + // storage/config rebuild creates a new target instance; do not clear this + // on ordinary restore_kv() or no-capture passes. + split_capture_validated_ = split_capture_validated_ || + (capture_gate_ok && shard_gate_ok && feature_gate_ok); + } + if (capture_selftest && capture_ssm_intermediates) { + if (!split_capture_validated_) { + std::fprintf(stderr, + "[target-split][capture] self-test failed closed env_enabled=%d local_only=%d storage_ready=%d\n", + split_chain_fast_rollback_enabled() ? 1 : 0, + local_only ? 1 : 0, storage_ready ? 1 : 0); + return false; + } + } + return ok; } bool Qwen35LayerSplitDFlashTarget::snapshot_kv() { - for (auto & shard : shards_) snapshot_ssm_state(shard.cache); + for (auto & shard : shards_) { + if (!snapshot_ssm_state(shard.cache, shard.backend)) return false; + } if (remote_target_shard_ && remote_target_shard_->active()) { return remote_target_shard_->snapshot_kv(); } @@ -71,7 +216,595 @@ bool Qwen35LayerSplitDFlashTarget::restore_kv() { return false; } } - for (auto & shard : shards_) restore_ssm_state(shard.cache); + for (auto & shard : shards_) { + if (!restore_ssm_state(shard.cache, shard.backend)) return false; + } + rollback_poisoned_ = false; + return true; +} + +bool Qwen35LayerSplitDFlashTarget::supports_fast_rollback() const { + const bool env_enabled = split_chain_fast_rollback_enabled(); + const bool local_only = !(remote_target_shard_ && remote_target_shard_->active()); + const bool storage_ready = split_fast_rollback_storage_ready(shards_); + const bool supported = env_enabled && local_only && storage_ready && + split_capture_validated_; + if (std::getenv("DFLASH_SPLIT_CHAIN_ROLLBACK_DIAG") != nullptr) { + std::fprintf(stderr, + "[target-split][chain-rollback] split_chain_fast_rollback_supported=%d env_enabled=%d capture_validated=%d local_only=%d storage_ready=%d split_tree_verify_supported=0 split_rollback_to_tree_supported=0\n", + supported ? 1 : 0, env_enabled ? 1 : 0, + split_capture_validated_ ? 1 : 0, local_only ? 1 : 0, + storage_ready ? 1 : 0); + } + return supported; +} + +bool Qwen35LayerSplitDFlashTarget::rollback_to(int base_pos, int commit_n) { + const auto t0 = std::chrono::steady_clock::now(); + const bool diag = std::getenv("DFLASH_SPLIT_CHAIN_ROLLBACK_DIAG") != nullptr; + last_rollback_context_fatal_ = false; + auto fail = [&](const char * why, cudaError_t err = cudaSuccess) { + const bool has_cuda_error = err != cudaSuccess; + const bool context_fatal = has_cuda_error && split_rollback_context_fatal(err); + last_rollback_context_fatal_ = context_fatal; + if (diag) { + std::fprintf(stderr, + "[target-split][chain-rollback] split_chain_fast_rollback_fail=1 reason=%s split_chain_fast_rollback_fallback_restore_replay=%d split_chain_rollback_cuda_error_name=%s split_chain_rollback_failure_class=%s split_chain_rollback_context_fatal=%d\n", + why, + context_fatal ? 0 : 1, + has_cuda_error ? cudaGetErrorName(err) : "cudaSuccess", + context_fatal ? "context-fatal" : "recoverable-return", + context_fatal ? 1 : 0); + } + return false; + }; + + if (!supports_fast_rollback()) return fail("unsupported_or_not_validated"); + if (remote_target_shard_ && remote_target_shard_->active()) return fail("remote_target_shard_active"); + if (shards_.empty()) return fail("no_shards"); + if (commit_n <= 0) return fail("commit_n_nonpositive"); + + int q_len = -1; + for (const auto & shard : shards_) { + const int shard_q = shard.cache.cur_pos - base_pos; + if (shard_q < 0) return fail("negative_q_len"); + if (q_len < 0) q_len = shard_q; + if (q_len != shard_q) return fail("cross_shard_cur_pos_mismatch"); + } + + if (commit_n >= q_len) { + for (auto & shard : shards_) shard.cache.cur_pos = base_pos + commit_n; + if (diag) { + std::fprintf(stderr, + "[target-split][chain-rollback] split_chain_fast_rollback_attempts=1 split_chain_fast_rollback_success=1 split_chain_rollback_commit_n=%d split_chain_rollback_slot_idx=%d split_chain_rollback_context_device_per_shard=none_no_restore split_chain_rollback_stream_source_per_shard=none_no_restore split_chain_feature_ring_action=none_required_position_aligned split_tree_verify_supported=0 split_rollback_to_tree_supported=0\n", + commit_n, commit_n - 1); + } + return true; + } + + const int rollback_idx = commit_n - 1; + std::vector restored_per_shard(shards_.size(), 0); + int prior_device = -1; + (void)cudaGetDevice(&prior_device); + + for (size_t si = 0; si < shards_.size(); ++si) { + auto & shard = shards_[si]; + auto & cache = shard.cache; + const auto & w = shard.weights; + cudaError_t ce = cudaSetDevice(shard.gpu); + if (ce != cudaSuccess) return fail("cuda_set_device_failed", ce); + const cudaStream_t shard_stream = cudaStreamPerThread; + if (diag) { + int current_device = -1; + cudaGetDevice(¤t_device); + std::fprintf(stderr, + "[target-split][chain-rollback] shard=%zu gpu=%d split_chain_rollback_context_device_per_shard=%d split_chain_rollback_stream_source_per_shard=cudaStreamPerThread\n", + si, shard.gpu, current_device); + } + int dn_idx = 0; + for (int il = 0; il < w.n_layer; ++il) { + const bool is_attn = (((il + 1) % w.full_attention_interval) == 0); + if (is_attn) continue; + const bool owns_layer = il >= shard.layer_begin && il < shard.layer_end; + if (!owns_layer) { dn_idx++; continue; } + if (dn_idx >= (int)cache.ssm_state.size() || dn_idx >= (int)cache.conv_state.size() || + dn_idx >= (int)cache.ssm_intermediate.size() || dn_idx >= (int)cache.conv_input_cache.size()) { + return fail("capture_index_oob"); + } + ggml_tensor * ssm_state = cache.ssm_state[dn_idx]; + ggml_tensor * conv_state = cache.conv_state[dn_idx]; + ggml_tensor * ssm_inter = cache.ssm_intermediate[dn_idx]; + ggml_tensor * conv_input = cache.conv_input_cache[dn_idx]; + if (!ssm_state || !conv_state || !ssm_inter || !conv_input) return fail("missing_capture_storage"); + if (rollback_idx >= (int)ssm_inter->ne[3]) return fail("rollback_idx_oob"); + + const size_t ssm_elems = (size_t)ssm_state->ne[0] * + (size_t)ssm_state->ne[1] * (size_t)ssm_state->ne[2]; + const void * ssm_src = (const char *)ssm_inter->data + + (size_t)rollback_idx * ssm_inter->nb[3]; + if (ssm_inter->type == GGML_TYPE_F32) { + ce = cudaMemcpyAsync(ssm_state->data, ssm_src, + ssm_elems * sizeof(float), + cudaMemcpyDeviceToDevice, shard_stream); + if (ce != cudaSuccess) return fail("ssm_f32_copy_failed", ce); + } else { + const auto to_fp32 = ggml_get_to_fp32_cuda(ssm_inter->type); + if (!to_fp32) return fail("missing_ssm_converter"); + to_fp32(ssm_src, (float *)ssm_state->data, (int64_t)ssm_elems, shard_stream); + ce = cudaPeekAtLastError(); + if (ce != cudaSuccess) return fail("ssm_convert_launch_failed", ce); + } + + const int K_conv = w.ssm_d_conv; + if (commit_n + K_conv - 1 > (int)conv_input->ne[0]) return fail("conv_input_oob"); + const int row_cnt = (int)conv_input->ne[1]; + const size_t elt = ggml_element_size(conv_input); + const size_t dpitch = (size_t)(K_conv - 1) * elt; + const size_t spitch = conv_input->nb[1]; + const size_t width = (size_t)(K_conv - 1) * elt; + const void * conv_src = (const char *)conv_input->data + (size_t)commit_n * elt; + ce = cudaMemcpy2DAsync(conv_state->data, dpitch, + conv_src, spitch, + width, row_cnt, + cudaMemcpyDeviceToDevice, shard_stream); + if (ce != cudaSuccess) return fail("conv_copy_failed", ce); + restored_per_shard[si]++; + dn_idx++; + } + cudaError_t sync = cudaStreamSynchronize(shard_stream); + if (sync != cudaSuccess) return fail("stream_sync_failed", sync); + } + if (prior_device >= 0) (void)cudaSetDevice(prior_device); + + for (auto & shard : shards_) shard.cache.cur_pos = base_pos + commit_n; + const auto t1 = std::chrono::steady_clock::now(); + const auto latency_us = std::chrono::duration_cast(t1 - t0).count(); + if (diag) { + std::fprintf(stderr, + "[target-split][chain-rollback] split_chain_fast_rollback_attempts=1 split_chain_fast_rollback_success=1 split_chain_fast_rollback_fail=0 split_chain_fast_rollback_fallback_restore_replay=0 split_chain_rollback_commit_n=%d split_chain_rollback_slot_idx=%d split_chain_rollback_latency_us=%lld split_chain_rollback_cuda_error_name=cudaSuccess split_chain_rollback_failure_class=recoverable-return split_chain_rollback_context_fatal=0 split_chain_feature_ring_action=none_required_position_aligned split_tree_verify_supported=0 split_rollback_to_tree_supported=0\n", + commit_n, rollback_idx, (long long)latency_us); + for (size_t si = 0; si < restored_per_shard.size(); ++si) { + std::fprintf(stderr, + "[target-split][chain-rollback] shard=%zu gpu=%d split_chain_rollback_layers_restored_per_shard=%d\n", + si, shards_[si].gpu, restored_per_shard[si]); + } + } + return true; +} + +bool Qwen35LayerSplitDFlashTarget::supports_tree_verify() const { + // Production fail-closed boundary: sibling visibility and depth positions + // are not wired for layer-split execution. Evidence overlays may expose a + // separate test-only seam, but ordinary runtime must never advertise this. + return false; +} + +bool Qwen35LayerSplitDFlashTarget::verify_tree( + int committed, + const DDTree & tree, + const std::vector & flat_tokens, + int n_alloc, + std::vector & posterior_out, + std::vector * logits_out) { + auto fail_precheck = [&](const char * reason) { + std::fprintf(stderr, + "[target-split][pure-chain-guard] verify_precheck_fail=%s committed=%d n_alloc=%d n_actual=%d n_nodes=%d parents_size=%zu flat_tokens_size=%zu\n", + reason, committed, n_alloc, 1 + tree.n_nodes, tree.n_nodes, + tree.parents.size(), flat_tokens.size()); + return false; + }; + const int n_actual = 1 + tree.n_nodes; + if (n_actual <= 0) return fail_precheck("n_actual_nonpositive"); + if ((int)tree.parents.size() < n_actual) return fail_precheck("parents_size_lt_n_actual"); + if (!qwen35_split_run_if_root_inclusive_pure_chain( + tree.parents.data(), tree.parents.size(), (size_t)n_actual)) { + return fail_precheck("not_root_inclusive_pure_chain"); + } + if (rollback_poisoned_) return fail_precheck("rollback_poisoned"); + if (!supports_tree_verify()) return fail_precheck("production_capability_disabled"); + if (n_alloc < n_actual) return fail_precheck("n_alloc_lt_n_actual"); + if ((int)flat_tokens.size() < n_actual) return fail_precheck("flat_tokens_size_lt_n_actual"); + if ((int)tree.token_ids.size() < tree.n_nodes) return fail_precheck("token_ids_size_lt_n_nodes"); + if ((int)tree.depths.size() < tree.n_nodes) return fail_precheck("depths_size_lt_n_nodes"); + if (tree.visibility.size() < (size_t)n_actual * (size_t)n_actual) { + return fail_precheck("visibility_size_lt_n_actual_sq"); + } + if (committed < 0) return fail_precheck("negative_committed"); + required_tree_slots_ = n_alloc; + + // Padded graph reuse is retained only behind the disabled production seam. + // Padding is embedded as zeros and cannot enter posterior decisions. + std::vector padded_tokens((size_t)n_alloc, 0); + std::copy(flat_tokens.begin(), flat_tokens.begin() + n_actual, padded_tokens.begin()); + + std::vector parent_ids((size_t)n_alloc, 0); + parent_ids[0] = -1; + for (int s = 1; s < n_actual; ++s) { + const int pflat = tree.parents[(size_t)s]; + if (pflat < 0 || pflat >= s) return fail_precheck("parent_value_oob"); + parent_ids[(size_t)s] = (int32_t)pflat; + } + std::vector host_act((size_t)n_alloc * (size_t)hidden_size(), 0.0f); + if (!embed_tokens(padded_tokens.data(), n_actual, host_act.data())) return fail_precheck("embed_tokens_failed"); + + ActivationPair acts; + if (!activation_pair_init(acts, shards_.front().backend, hidden_size(), n_alloc)) { + return fail_precheck("activation_pair_init_failed"); + } + ggml_backend_tensor_set(acts.a, host_act.data(), 0, + sizeof(float) * host_act.size()); + + Qwen35SplitTreeInputs tree_inputs; + tree_inputs.parent_ids = parent_ids.data(); + tree_inputs.visibility = nullptr; + tree_inputs.n_actual = n_alloc; + tree_inputs.committed = committed; + + int last_tok = -1; + std::vector all_argmax; + std::vector tree_captures; + const bool ok = run_qwen35_layer_split_tree_verify_from_activation( + shards_, acts, committed, n_alloc, n_alloc, last_tok, + kq_stride_pad_, fa_window_, tree_inputs, &all_argmax, logits_out, + &tree_captures, kvflash_, /*kvflash_preallocated=*/false); + activation_pair_free(acts); + if (!ok) return fail_precheck("tree_verify_from_activation_failed"); + if ((int)all_argmax.size() < n_actual) return fail_precheck("all_argmax_size_lt_n_actual"); + + posterior_out.assign(all_argmax.begin(), all_argmax.begin() + n_actual); + if (logits_out && !logits_out->empty()) { + const int vocab = shards_.back().weights.n_vocab; + if ((int)logits_out->size() >= n_alloc * vocab) { + logits_out->resize((size_t)n_actual * (size_t)vocab); + } + } + return true; +} + +bool Qwen35LayerSplitDFlashTarget::rollback_to_tree( + int committed, + const DDTree & tree, + const std::vector & accepted_dfs) { + const int n_actual_ri = 1 + tree.n_nodes; + if (n_actual_ri <= 0 || + !qwen35_split_run_if_root_inclusive_pure_chain( + tree.parents.data(), tree.parents.size(), (size_t)n_actual_ri)) { + std::fprintf(stderr, + "[target-split][pure-chain-guard] rollback_precheck_fail=not_root_inclusive_pure_chain n_actual=%d parents_size=%zu\n", + n_actual_ri, tree.parents.size()); + return false; + } + const auto t0 = std::chrono::steady_clock::now(); + const bool diag = std::getenv("DFLASH_SPLIT_CHAIN_ROLLBACK_DIAG") != nullptr; + last_rollback_context_fatal_ = false; + enum class Phase { AValidate, BRestore, CBarrier1, DCompact, EBarrier2, FCommit }; + auto phase_name = [](Phase p) -> const char * { + switch (p) { + case Phase::AValidate: return "A_validate_all"; + case Phase::BRestore: return "B_restore_recurrent_all"; + case Phase::CBarrier1: return "C_barrier_1"; + case Phase::DCompact: return "D_compact_kv_and_feature_all"; + case Phase::EBarrier2: return "E_barrier_2"; + case Phase::FCommit: return "F_commit_advance"; + } + return "unknown"; + }; + Phase phase = Phase::AValidate; + auto fail = [&](const char * why, cudaError_t err = cudaSuccess, + int shard_idx = -1, int layer_id = -1, int dn_idx = -1) { + const bool has_cuda_error = err != cudaSuccess; + const bool context_fatal = has_cuda_error && split_rollback_context_fatal(err); + last_rollback_context_fatal_ = context_fatal || phase != Phase::AValidate; + if (phase != Phase::AValidate) rollback_poisoned_ = true; + if (diag) { + std::fprintf(stderr, + "[target-split][pure-chain-rollback] split_rollback_to_tree_fail=1 phase=%s reason=%s shard=%d layer=%d dn_idx=%d cuda_error_name=%s split_rollback_to_tree_context_fatal=%d split_rollback_to_tree_poisoned=%d\n", + phase_name(phase), why, shard_idx, layer_id, dn_idx, + has_cuda_error ? cudaGetErrorName(err) : "cudaSuccess", + context_fatal ? 1 : 0, last_rollback_context_fatal_ ? 1 : 0); + } + return false; + }; + + if (rollback_poisoned_) return fail("rollback_poisoned"); + if (!supports_tree_verify()) return fail("unsupported_or_not_validated"); + if (remote_target_shard_ && remote_target_shard_->active()) return fail("remote_target_shard_active"); + if (shards_.empty()) return fail("no_shards"); + if (committed < 0) return fail("negative_committed"); + const int commit_n = (int)accepted_dfs.size(); + if (commit_n <= 0) return fail("commit_n_nonpositive"); + if ((int)tree.parents.size() < n_actual_ri) return fail("parents_size_lt_n_actual"); + if (tree.parents[0] != -1) return fail("parents_root_not_minus1"); + // Parents use the root-inclusive flat-slot contract: slot 0 is the + // synthetic root and each later value names its parent flat slot. + auto parent_flat_slot = [&](int flat_slot) -> int { + if (flat_slot == 0) return -1; + if (flat_slot < 0 || flat_slot > tree.n_nodes) return -2; + const int pflat = tree.parents[(size_t)flat_slot]; + if (pflat < 0 || pflat >= flat_slot) return -2; + return pflat; + }; + if (accepted_dfs[0] != 0) return fail("accepted_dfs_not_starting_at_root"); + for (int d = 0; d < commit_n; ++d) { + const int dfs_idx = accepted_dfs[(size_t)d]; + if (dfs_idx < 0 || dfs_idx > tree.n_nodes) return fail("accepted_dfs_oob"); + if (d > 0) { + if (accepted_dfs[(size_t)d] <= accepted_dfs[(size_t)d - 1]) { + return fail("accepted_dfs_not_strictly_increasing"); + } + int cur = dfs_idx; + const int prev_flat = accepted_dfs[(size_t)d - 1]; + bool parent_found = false; + for (int guard = 0; guard <= tree.n_nodes && cur >= 0; ++guard) { + if (cur == prev_flat) { parent_found = true; break; } + cur = parent_flat_slot(cur); + } + if (!parent_found) return fail("accepted_path_not_parent_chain"); + } + } + // accepted_dfs is root-inclusive, so the deepest accepted flat slot is + // used directly with no index conversion. + const int rollback_dfs = accepted_dfs.back(); + const bool walked_sibling = [&]() { + for (int i = 0; i < commit_n; ++i) if (accepted_dfs[(size_t)i] != i) return true; + return false; + }(); + + + struct OwnedDelta { size_t si; int layer_id; int dn_idx; }; + std::vector owned; + owned.reserve(64); + for (size_t si = 0; si < shards_.size(); ++si) { + const auto & shard = shards_[si]; + const auto & cache = shard.cache; + const auto & w = shard.weights; + int dn_idx = 0; + for (int il = 0; il < w.n_layer; ++il) { + const bool is_attn = (((il + 1) % w.full_attention_interval) == 0); + if (is_attn) continue; + const bool owns_layer = il >= shard.layer_begin && il < shard.layer_end; + if (!owns_layer) { dn_idx++; continue; } + if (dn_idx >= (int)cache.ssm_state.size() || dn_idx >= (int)cache.conv_state.size() || + dn_idx >= (int)cache.ssm_intermediate.size() || dn_idx >= (int)cache.conv_input_cache.size()) { + return fail("capture_index_oob", cudaSuccess, (int)si, il, dn_idx); + } + ggml_tensor * ssm_state = cache.ssm_state[(size_t)dn_idx]; + ggml_tensor * conv_state = cache.conv_state[(size_t)dn_idx]; + ggml_tensor * ssm_inter = cache.ssm_intermediate[(size_t)dn_idx]; + ggml_tensor * conv_input = cache.conv_input_cache[(size_t)dn_idx]; + if (!ssm_state || !conv_state || !ssm_inter || !conv_input) { + return fail("missing_capture_storage", cudaSuccess, (int)si, il, dn_idx); + } + if (ssm_inter->type != GGML_TYPE_F32 || conv_input->type != GGML_TYPE_F32 || + ssm_state->type != GGML_TYPE_F32 || conv_state->type != GGML_TYPE_F32) { + return fail("non_f32_recurrent_storage", cudaSuccess, (int)si, il, dn_idx); + } + // Both operands are root-inclusive flat-slot domain: ne[3] was + // filled by tree token t directly into persistent slot t. + if (rollback_dfs >= (int)ssm_inter->ne[3]) { + return fail("rollback_dfs_oob", cudaSuccess, (int)si, il, dn_idx); + } + const int K_conv = w.ssm_d_conv; + if (K_conv <= 1 || conv_input->ne[0] < K_conv || + conv_state->ne[0] < K_conv - 1 || + conv_input->ne[1] != conv_state->ne[1]) { + return fail("conv_shape_invalid", cudaSuccess, (int)si, il, dn_idx); + } + if (!walked_sibling) { + if (rollback_dfs + K_conv > (int)conv_input->ne[0]) { + return fail("conv_contiguous_oob", cudaSuccess, (int)si, il, dn_idx); + } + } else { + int virt = rollback_dfs; + for (int k = K_conv - 2; k >= 0; --k) { + const int sx_slot = (K_conv - 1) + virt; + if (sx_slot < 0 || sx_slot >= (int)conv_input->ne[0]) { + return fail("conv_ancestry_oob", cudaSuccess, (int)si, il, dn_idx); + } + virt = (virt >= 0) ? parent_flat_slot(virt) : (virt - 1); + } + } + owned.push_back({si, il, dn_idx}); + dn_idx++; + } + } + if (owned.empty()) return fail("no_owned_delta_layers"); + + if (feature_ring_) { + if (!feature_ring_->target_feat || feature_ring_->cap <= 0 || + commit_n > feature_ring_->cap || committed < 0) { + return fail("feature_ring_precondition_failed"); + } + } + + int prior_device = -1; + (void)cudaGetDevice(&prior_device); + std::vector restored_per_shard(shards_.size(), 0); + + phase = Phase::BRestore; + for (const auto & od : owned) { + auto & shard = shards_[od.si]; + auto & cache = shard.cache; + const auto & w = shard.weights; + cudaError_t ce = cudaSetDevice(shard.gpu); + if (ce != cudaSuccess) return fail("cuda_set_device_failed", ce, (int)od.si, od.layer_id, od.dn_idx); + const cudaStream_t stream = cudaStreamPerThread; + ggml_tensor * ssm_state = cache.ssm_state[(size_t)od.dn_idx]; + ggml_tensor * conv_state = cache.conv_state[(size_t)od.dn_idx]; + ggml_tensor * ssm_inter = cache.ssm_intermediate[(size_t)od.dn_idx]; + ggml_tensor * conv_input = cache.conv_input_cache[(size_t)od.dn_idx]; + const size_t ssm_elems = (size_t)ssm_state->ne[0] * + (size_t)ssm_state->ne[1] * (size_t)ssm_state->ne[2]; + // Exact-domain restore: root-inclusive flat slot -> root-inclusive + // ssm_intermediate ne[3] slot. I2 probes neighbors empirically. + const void * ssm_src = (const char *)ssm_inter->data + + (size_t)rollback_dfs * ssm_inter->nb[3]; + ce = cudaMemcpyAsync(ssm_state->data, ssm_src, ssm_elems * sizeof(float), + cudaMemcpyDeviceToDevice, stream); + if (ce != cudaSuccess) return fail("ssm_copy_failed", ce, (int)od.si, od.layer_id, od.dn_idx); + + const int K_conv = w.ssm_d_conv; + const int row_cnt = (int)conv_input->ne[1]; + const size_t elt = ggml_element_size(conv_input); + const size_t dpitch = (size_t)(K_conv - 1) * elt; + const size_t spitch = conv_input->nb[1]; + if (!walked_sibling) { + // conv_input row domain is [K_conv-1 prefix | root-inclusive + // verify rows]: flat slot t is physical row (K_conv-1)+t, so the + // K_conv-1 history rows ending at accepted flat slot rollback_dfs + // are physical rows rollback_dfs+1 .. rollback_dfs+K_conv-1 + // (directive 168 §1B; not an isolated conv-row change — it is + // re-derived together with the corrected rollback_dfs domain). + const void * conv_src = (const char *)conv_input->data + (size_t)(rollback_dfs + 1) * elt; + ce = cudaMemcpy2DAsync(conv_state->data, dpitch, conv_src, spitch, + (size_t)(K_conv - 1) * elt, row_cnt, + cudaMemcpyDeviceToDevice, stream); + if (ce != cudaSuccess) return fail("conv_contiguous_copy_failed", ce, (int)od.si, od.layer_id, od.dn_idx); + } else { + std::vector virt((size_t)(K_conv - 1)); + int cur = rollback_dfs; + for (int k = K_conv - 2; k >= 0; --k) { + virt[(size_t)k] = cur; + cur = (cur >= 0) ? parent_flat_slot(cur) : (cur - 1); + } + for (int k = 0; k < K_conv - 1; ++k) { + // virt is root-inclusive flat-slot domain; K_conv-1 converts + // it to the corresponding row in conv_input's concatenation. + const int sx_slot = (K_conv - 1) + virt[(size_t)k]; + const void * src_col = (const char *)conv_input->data + (size_t)sx_slot * elt; + char * dst_col = (char *)conv_state->data + (size_t)k * elt; + ce = cudaMemcpy2DAsync(dst_col, dpitch, src_col, spitch, elt, row_cnt, + cudaMemcpyDeviceToDevice, stream); + if (ce != cudaSuccess) return fail("conv_ancestry_copy_failed", ce, (int)od.si, od.layer_id, od.dn_idx); + } + } + restored_per_shard[od.si]++; + } + + phase = Phase::CBarrier1; + for (size_t si = 0; si < shards_.size(); ++si) { + cudaError_t ce = cudaSetDevice(shards_[si].gpu); + if (ce != cudaSuccess) return fail("cuda_set_device_failed", ce, (int)si); + ce = cudaStreamSynchronize(cudaStreamPerThread); + if (ce != cudaSuccess) return fail("stream_sync_failed", ce, (int)si); + } + + phase = Phase::DCompact; + for (size_t si = 0; si < shards_.size(); ++si) { + auto & shard = shards_[si]; + auto & cache = shard.cache; + cudaError_t ce = cudaSetDevice(shard.gpu); + if (ce != cudaSuccess) return fail("cuda_set_device_failed", ce, (int)si); + const cudaStream_t stream = cudaStreamPerThread; + for (int d = 0; d < commit_n; ++d) { + // Root-inclusive accepted flat slot s was written at committed+s + // during the tree pass (root row at committed+0); no +1 + // conversion (directive 168 §1B). + const int src_dfs = accepted_dfs[(size_t)d]; + for (size_t l = 0; l < cache.attn_k.size(); ++l) { + ggml_tensor * ck = cache.attn_k[l]; + ggml_tensor * cv = l < cache.attn_v.size() ? cache.attn_v[l] : nullptr; + if (!ck || !cv) continue; + const int src_pos = committed + src_dfs; + const int dst_pos = committed + d; + if (src_pos < 0 || dst_pos < 0 || + src_pos >= (int)ck->ne[1] || dst_pos >= (int)ck->ne[1] || + src_pos >= (int)cv->ne[1] || dst_pos >= (int)cv->ne[1]) { + return fail("kv_compact_oob", cudaSuccess, (int)si, -1, (int)l); + } + const int n_kv = (int)ck->ne[2]; + for (int h = 0; h < n_kv; ++h) { + const size_t k_bytes = ck->nb[1]; + const size_t v_bytes = cv->nb[1]; + const size_t k_src = (size_t)src_pos * ck->nb[1] + (size_t)h * ck->nb[2]; + const size_t k_dst = (size_t)dst_pos * ck->nb[1] + (size_t)h * ck->nb[2]; + const size_t v_src = (size_t)src_pos * cv->nb[1] + (size_t)h * cv->nb[2]; + const size_t v_dst = (size_t)dst_pos * cv->nb[1] + (size_t)h * cv->nb[2]; + ce = cudaMemcpyAsync((char *)ck->data + k_dst, + (const char *)ck->data + k_src, + k_bytes, cudaMemcpyDeviceToDevice, stream); + if (ce != cudaSuccess) return fail("kv_k_copy_failed", ce, (int)si, -1, (int)l); + ce = cudaMemcpyAsync((char *)cv->data + v_dst, + (const char *)cv->data + v_src, + v_bytes, cudaMemcpyDeviceToDevice, stream); + if (ce != cudaSuccess) return fail("kv_v_copy_failed", ce, (int)si, -1, (int)l); + } + } + } + } + + long long feature_realign_latency_us = 0; + if (feature_ring_) { + const auto feature_t0 = std::chrono::steady_clock::now(); + // One N+1 staging buffer: accepted-DFS rows plus one guard row, then + // read back in committed-spine order. + // The helper reads and writes by logical position, matching + // draft_feature_mirror_sync_range(). This prewindow path uses host F32 + // staging for inspectable exactness; a device-staging fast path is an + // explicit follow-up optimization, not claimed here. + const int fc_in = feature_ring_->n_target_layers * feature_ring_->hidden_size; + if (fc_in <= 0) return fail("feature_ring_width_invalid"); + const int stage_rows = commit_n + 1; + if (stage_rows > feature_ring_->cap) return fail("feature_stage_rows_exceed_cap"); + std::vector source_stage((size_t)stage_rows * (size_t)fc_in, 0.0f); + std::vector one_row; + for (int d = 0; d < commit_n; ++d) { + // Root-inclusive accepted flat slot: feature row for accepted + // element d lives at committed+accepted_dfs[d]; no +1 (168 §1B). + if (!copy_feature_ring_range_to_host_f32(*feature_ring_, committed + accepted_dfs[(size_t)d], 1, one_row)) { + return fail("feature_source_stage_copy_failed"); + } + if ((int)one_row.size() != fc_in) return fail("feature_source_stage_width_mismatch"); + std::copy(one_row.begin(), one_row.end(), + source_stage.begin() + (size_t)d * (size_t)fc_in); + } + if (!copy_host_f32_to_feature_ring_range(*feature_ring_, committed, commit_n, source_stage)) { + return fail("feature_destination_backfill_failed"); + } + const auto feature_t1 = std::chrono::steady_clock::now(); + feature_realign_latency_us = (long long)std::chrono::duration_cast( + feature_t1 - feature_t0).count(); + } + + phase = Phase::EBarrier2; + for (size_t si = 0; si < shards_.size(); ++si) { + cudaError_t ce = cudaSetDevice(shards_[si].gpu); + if (ce != cudaSuccess) return fail("cuda_set_device_failed", ce, (int)si); + ce = cudaStreamSynchronize(cudaStreamPerThread); + if (ce != cudaSuccess) return fail("stream_sync_failed", ce, (int)si); + } + if (feature_ring_) { + cudaError_t ce = cudaSetDevice(feature_ring_->device); + if (ce != cudaSuccess) return fail("feature_cuda_set_device_failed", ce); + ce = cudaDeviceSynchronize(); + if (ce != cudaSuccess) return fail("feature_device_sync_failed", ce); + } + + phase = Phase::FCommit; + if (kvflash_ && !kvflash_->alloc_span(committed, commit_n)) { + return fail("kvflash_alloc_span_failed"); + } + for (auto & shard : shards_) { + shard.cache.cur_pos = committed + commit_n; + shard.cache.last_tok = -1; + } + if (prior_device >= 0) (void)cudaSetDevice(prior_device); + const auto t1 = std::chrono::steady_clock::now(); + const auto latency_us = std::chrono::duration_cast(t1 - t0).count(); + if (diag) { + std::fprintf(stderr, + "[target-split][pure-chain-rollback] split_rollback_to_tree_success=1 split_rollback_to_tree_commit_n=%d split_rollback_to_tree_rollback_dfs=%d split_rollback_to_tree_walked_sibling=%d split_rollback_to_tree_latency_us=%lld split_rollback_to_tree_feature_realign=%d split_rollback_to_tree_post_verify_required=1\n", + commit_n, rollback_dfs, walked_sibling ? 1 : 0, + (long long)latency_us, feature_ring_ ? 1 : 0); + std::fprintf(stderr, + "[target-split][pure-chain-rollback] split_rollback_to_tree_feature_realign_latency_us=%lld split_rollback_to_tree_feature_staging=device_followup_host_f32_prewindow split_rollback_to_tree_feature_stage_rows=%d\n", + feature_realign_latency_us, feature_ring_ ? (commit_n + 1) : 0); + for (size_t si = 0; si < restored_per_shard.size(); ++si) { + std::fprintf(stderr, + "[target-split][pure-chain-rollback] shard=%zu gpu=%d split_rollback_to_tree_layers_restored_per_shard=%d\n", + si, shards_[si].gpu, restored_per_shard[si]); + } + } return true; } diff --git a/server/src/qwen35/qwen35_layer_split_dflash_target.h b/server/src/qwen35/qwen35_layer_split_dflash_target.h index 274a4a2db..b12a1dfed 100644 --- a/server/src/qwen35/qwen35_layer_split_dflash_target.h +++ b/server/src/qwen35/qwen35_layer_split_dflash_target.h @@ -20,6 +20,7 @@ #include "qwen35_target_shard_ipc.h" #include "step_graph.h" +#include #include namespace dflash::common { @@ -48,6 +49,20 @@ class Qwen35LayerSplitDFlashTarget : public DFlashTarget { bool snapshot_kv() override; bool restore_kv() override; + bool supports_fast_rollback() const override; + bool rollback_to(int base_pos, int commit_n) override; + bool supports_tree_verify() const override; + bool verify_tree(int committed, + const DDTree & tree, + const std::vector & flat_tokens, + int n_alloc, + std::vector & posterior_out, + std::vector * logits_out = nullptr) override; + bool rollback_to_tree(int committed, + const DDTree & tree, + const std::vector & accepted_dfs) override; + bool last_rollback_context_fatal() const { return last_rollback_context_fatal_; } + bool is_eos(int token) const override; bool embed_tokens(const int32_t * tokens, int n, @@ -72,6 +87,10 @@ class Qwen35LayerSplitDFlashTarget : public DFlashTarget { std::vector capture_ids_; StepGraph proj_sg_; + bool split_capture_validated_ = false; + bool last_rollback_context_fatal_ = false; + bool rollback_poisoned_ = false; + mutable int required_tree_slots_ = 1; }; } // namespace dflash::common diff --git a/server/src/qwen35/qwen35_layer_split_tree_guard.h b/server/src/qwen35/qwen35_layer_split_tree_guard.h new file mode 100644 index 000000000..a5f245609 --- /dev/null +++ b/server/src/qwen35/qwen35_layer_split_tree_guard.h @@ -0,0 +1,33 @@ +#pragma once + +#include +#include + +namespace dflash::common { + +using Qwen35SplitPureChainHook = bool (*)(void * context); + +// The only retained layer-split tree seam is the validated root-inclusive +// pure chain: slot 0 is the synthetic root and every later slot follows its +// immediate predecessor. The hook makes ordering observable in CPU-only tests. +inline bool qwen35_split_run_if_root_inclusive_pure_chain( + const int32_t * parents, + std::size_t parent_count, + std::size_t n_actual, + Qwen35SplitPureChainHook hook = nullptr, + void * hook_context = nullptr) { + if (parents == nullptr || n_actual == 0 || parent_count < n_actual) { + return false; + } + if (parents[0] != -1) { + return false; + } + for (std::size_t slot = 1; slot < n_actual; ++slot) { + if (parents[slot] != static_cast(slot - 1)) { + return false; + } + } + return hook == nullptr || hook(hook_context); +} + +} // namespace dflash::common diff --git a/server/src/qwen35/qwen35_target_graph.cpp b/server/src/qwen35/qwen35_target_graph.cpp index 24808177a..6d93caaa2 100644 --- a/server/src/qwen35/qwen35_target_graph.cpp +++ b/server/src/qwen35/qwen35_target_graph.cpp @@ -93,7 +93,8 @@ bool create_target_cache_partial(const TargetWeights & w, int layer_begin, int layer_end, bool allocate_target_feat, - int ctx_alloc) { + int ctx_alloc, + bool f32_ssm_intermediates) { if (layer_begin < 0) layer_begin = 0; if (layer_end < 0 || layer_end > w.n_layer) layer_end = w.n_layer; if (layer_begin > layer_end) { @@ -235,9 +236,16 @@ bool create_target_cache_partial(const TargetWeights & w, head_v_dim, head_v_dim, w.ssm_dt_rank); ggml_tensor * Cn = ggml_new_tensor_2d(out.rollback_ctx, GGML_TYPE_F32, w.ssm_d_conv - 1, conv_ch); - ggml_tensor * Si = ggml_new_tensor_4d(out.rollback_ctx, GGML_TYPE_Q8_0, + // I0 domain: ne[3] is the root-inclusive flat verify-token + // domain. Tree capture writes t=0 synthetic root through the + // final/padded flat slot directly into slot t. + const ggml_type ssm_intermediate_type = f32_ssm_intermediates + ? GGML_TYPE_F32 : GGML_TYPE_Q8_0; + ggml_tensor * Si = ggml_new_tensor_4d(out.rollback_ctx, ssm_intermediate_type, head_v_dim, head_v_dim, w.ssm_dt_rank, max_verify_tokens); + // I0 domain: ne[0] is [K_conv-1 prefix rows | + // root-inclusive verify rows]. ggml_tensor * Ci = ggml_new_tensor_3d(out.rollback_ctx, GGML_TYPE_F32, (w.ssm_d_conv - 1) + max_verify_tokens, conv_ch, 1); @@ -255,6 +263,27 @@ bool create_target_cache_partial(const TargetWeights & w, } out.rollback_buf = ggml_backend_alloc_ctx_tensors(out.rollback_ctx, backend); + if (std::getenv("DFLASH_SPLIT_CHAIN_ROLLBACK_DIAG")) { + int owned_delta_layers = 0; + for (int il = 0; il < w.n_layer; ++il) { + if (((il + 1) % w.full_attention_interval) != 0 && il >= layer_begin && il < layer_end) { + owned_delta_layers++; + } + } + const size_t elems_per_slot_per_layer = (size_t)head_v_dim * (size_t)head_v_dim * (size_t)w.ssm_dt_rank; + const size_t f32_bytes_per_slot_per_layer = elems_per_slot_per_layer * sizeof(float); + const size_t q8_bytes_per_slot_per_layer = ((elems_per_slot_per_layer + 31) / 32) * 34; + const size_t f32_total = f32_bytes_per_slot_per_layer * (size_t)max_verify_tokens * (size_t)owned_delta_layers; + const size_t q8_total = q8_bytes_per_slot_per_layer * (size_t)max_verify_tokens * (size_t)owned_delta_layers; + std::fprintf(stderr, + "[target-split][chain-rollback] split_ssm_intermediate_dtype=%s split_ssm_intermediate_persist_dtype_dst=%s split_ssm_intermediate_persist_quantized=%d layer_begin=%d layer_end=%d owned_delta_layers=%d max_verify_tokens=%d split_ssm_intermediate_f32_bytes=%zu split_ssm_intermediate_incremental_bytes_over_q8=%zu\n", + f32_ssm_intermediates ? "F32" : "Q8_0", + f32_ssm_intermediates ? "F32" : "Q8_0", + f32_ssm_intermediates ? 0 : 1, + layer_begin, layer_end, owned_delta_layers, max_verify_tokens, f32_total, + f32_ssm_intermediates && f32_total > q8_total + ? f32_total - q8_total : 0); + } if (!out.rollback_buf) { set_last_error("ggml_backend_alloc_ctx_tensors failed for rollback cache"); ggml_free(out.rollback_ctx); @@ -446,24 +475,56 @@ bool migrate_prefill_cache(const TargetWeights & w, return true; } -// Snapshot/restore SSM+conv state for speculative rollback. Uses device-side -// tensor copy (ggml_backend_tensor_copy). Called outside of any compute graph. -void snapshot_ssm_state(TargetCache & c) { +// Snapshot/restore SSM+conv state for speculative rollback. Queue all device +// copies on one backend stream, then synchronize once for the complete snapshot. +static bool recurrent_snapshot_layout_valid(const TargetCache & c) { + const size_t n = c.ssm_state.size(); + if (c.ssm_state_snap.size() != n || c.conv_state.size() != n || + c.conv_state_snap.size() != n) { + return false; + } + for (size_t i = 0; i < n; i++) { + const bool owns_state = c.ssm_state[i] || c.ssm_state_snap[i] || + c.conv_state[i] || c.conv_state_snap[i]; + if (!owns_state) continue; + if (!c.ssm_state[i] || !c.ssm_state_snap[i] || + !c.conv_state[i] || !c.conv_state_snap[i] || + c.ssm_state[i]->type != c.ssm_state_snap[i]->type || + !ggml_are_same_shape(c.ssm_state[i], c.ssm_state_snap[i]) || + !ggml_are_same_stride(c.ssm_state[i], c.ssm_state_snap[i]) || + c.conv_state[i]->type != c.conv_state_snap[i]->type || + !ggml_are_same_shape(c.conv_state[i], c.conv_state_snap[i]) || + !ggml_are_same_stride(c.conv_state[i], c.conv_state_snap[i])) { + return false; + } + } + return true; +} + +bool snapshot_ssm_state(TargetCache & c, ggml_backend_t backend) { + if (!backend || !recurrent_snapshot_layout_valid(c)) return false; for (size_t i = 0; i < c.ssm_state.size(); i++) { - if (!c.ssm_state[i] || !c.ssm_state_snap[i]) continue; - ggml_backend_tensor_copy(c.ssm_state[i], c.ssm_state_snap[i]); - if (!c.conv_state[i] || !c.conv_state_snap[i]) continue; - ggml_backend_tensor_copy(c.conv_state[i], c.conv_state_snap[i]); + if (!c.ssm_state[i]) continue; + ggml_backend_tensor_copy_async( + backend, backend, c.ssm_state[i], c.ssm_state_snap[i]); + ggml_backend_tensor_copy_async( + backend, backend, c.conv_state[i], c.conv_state_snap[i]); } + ggml_backend_synchronize(backend); + return true; } -void restore_ssm_state(TargetCache & c) { +bool restore_ssm_state(TargetCache & c, ggml_backend_t backend) { + if (!backend || !recurrent_snapshot_layout_valid(c)) return false; for (size_t i = 0; i < c.ssm_state.size(); i++) { - if (!c.ssm_state_snap[i] || !c.ssm_state[i]) continue; - ggml_backend_tensor_copy(c.ssm_state_snap[i], c.ssm_state[i]); - if (!c.conv_state_snap[i] || !c.conv_state[i]) continue; - ggml_backend_tensor_copy(c.conv_state_snap[i], c.conv_state[i]); + if (!c.ssm_state[i]) continue; + ggml_backend_tensor_copy_async( + backend, backend, c.ssm_state_snap[i], c.ssm_state[i]); + ggml_backend_tensor_copy_async( + backend, backend, c.conv_state_snap[i], c.conv_state[i]); } + ggml_backend_synchronize(backend); + return true; } // Allocate SSM/conv rollback snapshot tensors by mirroring the live recurrent @@ -910,13 +971,12 @@ static ggml_tensor * build_delta_net_block( // path which handles F32→Q8_0 quantization automatically. // persist_inter: when capture is requested, route the kernel's per-token // intermediate-state writes DIRECTLY into the persistent cache buffer via - // src[7], avoiding the legacy result-region cpy. Works for BOTH tree and - // non-tree (chain-verify) capture — the kernel checks src[7] regardless of - // tree mode, and write_inter is forced true whenever src[7] is non-null. - // This also keeps non-tree capture safe if the result tensor is compacted - // and no longer embeds per-token intermediate states. - // Q8_0 intermediates fall through (persist requires F32/F16); the legacy - // cpy path below handles F32→Q8_0 quantization for that case (guarded). + // src[7], avoiding the legacy result-region cpy. This works for both tree + // and non-tree (chain-verify) capture and preserves upstream #469 semantics. + // Stage 2 split-chain rollback allocates F32 intermediates, so its checkpoint + // path is never quantized. In tree mode, n_seq_tokens is root-inclusive and + // flat slot t is persisted directly at ne[3] slot t. + // Q8_0 intermediates fall through to the guarded legacy copy path below. ggml_tensor * persist_inter = (cap && cap->ssm_intermediate_states && (cap->ssm_intermediate_states->type == GGML_TYPE_F32 || cap->ssm_intermediate_states->type == GGML_TYPE_F16)) diff --git a/server/src/qwen35/qwen35_target_shard_ipc_daemon.cpp b/server/src/qwen35/qwen35_target_shard_ipc_daemon.cpp index 7f9162191..1d7f6064d 100644 --- a/server/src/qwen35/qwen35_target_shard_ipc_daemon.cpp +++ b/server/src/qwen35/qwen35_target_shard_ipc_daemon.cpp @@ -551,8 +551,11 @@ int run_qwen35_target_shard_ipc_daemon(const char * target_path, if (!enable_dflash) { stream_status(stream_fd, -1); } else { - for (auto & shard : shards) snapshot_ssm_state(shard.cache); - stream_status(stream_fd, 0); + bool ok = true; + for (auto & shard : shards) { + if (!snapshot_ssm_state(shard.cache, shard.backend)) ok = false; + } + stream_status(stream_fd, ok ? 0 : -1); } continue; } @@ -560,8 +563,11 @@ int run_qwen35_target_shard_ipc_daemon(const char * target_path, if (!enable_dflash) { stream_status(stream_fd, -1); } else { - for (auto & shard : shards) restore_ssm_state(shard.cache); - stream_status(stream_fd, 0); + bool ok = true; + for (auto & shard : shards) { + if (!restore_ssm_state(shard.cache, shard.backend)) ok = false; + } + stream_status(stream_fd, ok ? 0 : -1); } continue; } diff --git a/server/src/qwen35moe/qwen35moe_backend.cpp b/server/src/qwen35moe/qwen35moe_backend.cpp index d7e118934..78052a757 100644 --- a/server/src/qwen35moe/qwen35moe_backend.cpp +++ b/server/src/qwen35moe/qwen35moe_backend.cpp @@ -2078,7 +2078,11 @@ bool Qwen35MoeBackend::do_hybrid_spec_decode(int committed, int n_gen, draft_tok[0] = last_tok; // 4. Verify: snapshot recurrent state, then run ALL draft tokens batched - snapshot_ssm_state(target_cache()); + if (!snapshot_ssm_state(target_cache(), target_backend())) { + std::fprintf(stderr, "[hybrid-spec] recurrent-state snapshot failed\n"); + step_graph_destroy(draft_sg); + return false; + } target_tok.resize(verify_width); bool verify_ok = hybrid_forward_batch( @@ -2086,7 +2090,9 @@ bool Qwen35MoeBackend::do_hybrid_spec_decode(int committed, int n_gen, act_cur, target_tok, /*capture_features=*/false); if (!verify_ok) { std::fprintf(stderr, "[hybrid-spec] verify failed\n"); - restore_ssm_state(target_cache()); + if (!restore_ssm_state(target_cache(), target_backend())) { + std::fprintf(stderr, "[hybrid-spec] recurrent-state restore failed\n"); + } step_graph_destroy(draft_sg); return false; } @@ -2106,7 +2112,11 @@ bool Qwen35MoeBackend::do_hybrid_spec_decode(int committed, int n_gen, } // 6. Restore and replay accepted tokens - restore_ssm_state(target_cache()); + if (!restore_ssm_state(target_cache(), target_backend())) { + std::fprintf(stderr, "[hybrid-spec] recurrent-state restore failed\n"); + step_graph_destroy(draft_sg); + return false; + } std::vector replay_tok((size_t)commit_n); for (int i = 0; i < commit_n; i++) { diff --git a/server/src/server/server_main.cpp b/server/src/server/server_main.cpp index d5ecf71c9..6d654848d 100644 --- a/server/src/server/server_main.cpp +++ b/server/src/server/server_main.cpp @@ -15,6 +15,7 @@ #include "chat_template.h" #include "model_card.h" #include "common/backend_factory.h" +#include "common/chain_rollback_policy.h" #include "common/layer_split_utils.h" #include "common/spark_corpus.h" #include "common/moe_routing_collector.h" @@ -87,6 +88,9 @@ static void print_usage(const char * prog) { " --target-shard-ipc-work-dir Remote target shard IPC scratch directory\n" " --target-devices Reserved layer-split devices, e.g. cuda:0,cuda:1\n" " --target-layer-split Reserved layer-split weights\n" + " --target-split-fast-rollback Opt in to exact F32 checkpoints for local\n" + " qwen35 layer splits (extra VRAM; env:\n" + " DFLASH_SPLIT_FAST_ROLLBACK=1)\n" " --peer-access Enable peer access for multi-GPU placement\n" " --chunk Chunked-prefill chunk size (default: 512)\n" " --ds4-fused-decode Enable DeepSeek4 single-graph GPU decode\n" @@ -140,6 +144,9 @@ static void print_usage(const char * prog) { #else " Default: per model family (laguna q8_0, else q4_0)\n" #endif + " --kvflash Enable bounded KV residency\n" + " --kvflash-policy drafter, lru, or qk (default: drafter)\n" + " --kvflash-tau Drafter-policy reselect interval (default: 64)\n" "\n" "PFlash (speculative prefill compression):\n" " --prefill-compression off|auto|always (default: off)\n" @@ -175,7 +182,7 @@ static void print_usage(const char * prog) { " prefix; auto:N uses the last N requests.\n" " A plain N caches the first N prompt tokens.\n" " --disk-prefix-cache-compress Enable FlowKV aged-history compression composed\n" - " with the disk cache. Requires --pflash-drafter.\n" + " with the disk cache. Requires --prefill-drafter.\n" " compress=false default is byte-identical to base.\n" "\n" "Chat template (optional, e.g. froggeric Qwen3.6 template for tool-using\n" @@ -185,6 +192,11 @@ static void print_usage(const char * prog) { " renderer. Empty or missing falls back\n" " to the hardcoded template.\n" "\n" + "MoE expert placement:\n" + " --spark Enable self-tuning hot/cold expert placement\n" + " --spark-slots Explicit expert-cache slots per layer\n" + " --spark-vram Total VRAM target (default: whole card)\n" + "\n" "Expert routing analysis:\n" " --freq Enable expert frequency tracking + print analysis at shutdown\n" " --collect-routing Log binary routing data (hidden states + expert IDs)\n" @@ -210,6 +222,7 @@ int main(int argc, char ** argv) { bool target_device_seen = false; bool target_devices_seen = false; bool fast_rollback_forced_off = false; + bool target_split_fast_rollback_cli = false; bool adaptive_experts_set = false; // --adaptive-experts (MoE architectures only) // Track which thinking-budget tunables the operator set via CLI. @@ -299,6 +312,8 @@ int main(int argc, char ** argv) { std::fprintf(stderr, "[server] bad --target-layer-split value\n"); return 2; } + } else if (std::strcmp(argv[i], "--target-split-fast-rollback") == 0) { + target_split_fast_rollback_cli = true; } else if (std::strcmp(argv[i], "--peer-access") == 0) { bargs.device.peer_access = true; } else if (std::strcmp(argv[i], "--chunk") == 0 && i + 1 < argc) { @@ -544,7 +559,28 @@ int main(int argc, char ** argv) { return 2; } } - if (fast_rollback_forced_off) bargs.fast_rollback = false; + if (fast_rollback_forced_off) { + bargs.fast_rollback = false; + target_split_fast_rollback_cli = false; + // This is the global rollback kill switch, including an externally + // supplied layer-split opt-in. + unset_environment_variable("DFLASH_SPLIT_FAST_ROLLBACK"); + } else if (target_split_fast_rollback_cli) { + if (!bargs.device.is_layer_split()) { + std::fprintf(stderr, + "[server] --target-split-fast-rollback requires " + "--target-devices with at least two local devices\n"); + return 2; + } + if (bargs.device.is_mixed_layer_split() || + bargs.remote_target_shard.enabled()) { + std::fprintf(stderr, + "[server] --target-split-fast-rollback supports only local " + "same-backend target splits\n"); + return 2; + } + set_environment_variable("DFLASH_SPLIT_FAST_ROLLBACK", "1", true); + } // Resolve documented environment defaults before factory preparation so // compatibility warnings describe the effective backend configuration. @@ -583,6 +619,12 @@ int main(int argc, char ** argv) { } const ResolvedBackendPlan & backend_plan = backend_preparation.plan; const std::string & arch = backend_plan.arch(); + if (target_split_fast_rollback_cli && arch != "qwen35") { + std::fprintf(stderr, + "[server] --target-split-fast-rollback is only supported for " + "qwen35 targets (detected '%s')\n", arch.c_str()); + return 2; + } // Sync max_ctx: if --max-ctx was not provided, use the backend's default. // This prevents the HTTP server from accepting prompts larger than the @@ -980,6 +1022,10 @@ int main(int argc, char ** argv) { } std::fprintf(stderr, "[server] │ ddtree = %s\n", bargs.ddtree_mode ? "ON" : "off"); std::fprintf(stderr, "[server] │ fast_rollback = %s\n", bargs.fast_rollback ? "ON" : "off"); + if (bargs.device.is_layer_split()) { + std::fprintf(stderr, "[server] │ split_rollback = %s\n", + split_chain_fast_rollback_enabled() ? "ON" : "off"); + } std::fprintf(stderr, "[server] │ ddtree_budget = %d\n", bargs.ddtree_budget); std::fprintf(stderr, "[server] │ prefix_cache = %d slots\n", sconfig.prefix_cache_cap); std::fprintf(stderr, "[server] │ prefill_cache = %d slots\n", sconfig.prefill_cache_cap); diff --git a/server/test/test_chain_rollback_policy.cpp b/server/test/test_chain_rollback_policy.cpp index 0ad661dca..ea7c87a67 100644 --- a/server/test/test_chain_rollback_policy.cpp +++ b/server/test/test_chain_rollback_policy.cpp @@ -1,6 +1,9 @@ #include "CppUnitTestFramework.hpp" #include "scoped_env.h" #include "chain_rollback_policy.h" +#include "internal.h" + +#include "ggml-cpu.h" #include #include @@ -9,6 +12,7 @@ using dflash::common::resolve_chain_rollback_policy; using dflash::common::RollbackDiag; +using dflash::common::split_chain_fast_rollback_enabled; namespace { struct ChainRollbackPolicyFixture {}; @@ -62,6 +66,71 @@ TEST_CASE(ChainRollbackPolicyFixture, policy_defaults_and_env_parsing) { clear_policy_env(); } +TEST_CASE(ChainRollbackPolicyFixture, split_fast_rollback_is_explicitly_opt_in) { + const luce_test::ScopedEnvVar split_fast("DFLASH_SPLIT_FAST_ROLLBACK", nullptr); + unsetenv("DFLASH_SPLIT_FAST_ROLLBACK"); + CHECK(!split_chain_fast_rollback_enabled()); + + setenv("DFLASH_SPLIT_FAST_ROLLBACK", "1", 1); + CHECK(split_chain_fast_rollback_enabled()); + setenv("DFLASH_SPLIT_FAST_ROLLBACK", "true", 1); + CHECK(split_chain_fast_rollback_enabled()); + setenv("DFLASH_SPLIT_FAST_ROLLBACK", "0", 1); + CHECK(!split_chain_fast_rollback_enabled()); + setenv("DFLASH_SPLIT_FAST_ROLLBACK", "", 1); + CHECK(!split_chain_fast_rollback_enabled()); + unsetenv("DFLASH_SPLIT_FAST_ROLLBACK"); +} + +TEST_CASE(ChainRollbackPolicyFixture, split_checkpoint_dtype_is_gated_at_allocation) { + const luce_test::ScopedEnvVar kv_f16("DFLASH27B_KV_F16", nullptr); + const luce_test::ScopedEnvVar kv_q4("DFLASH27B_KV_Q4", nullptr); + const luce_test::ScopedEnvVar kv_tq3("DFLASH27B_KV_TQ3", nullptr); + const luce_test::ScopedEnvVar kv_k("DFLASH27B_KV_K", nullptr); + const luce_test::ScopedEnvVar kv_v("DFLASH27B_KV_V", nullptr); + + ggml_backend_t backend = ggml_backend_cpu_init(); + CHECK(backend != nullptr); + if (!backend) return; + + dflash::common::TargetWeights weights; + weights.n_layer = 2; + weights.full_attention_interval = 2; + weights.n_embd_head_k = 32; + weights.n_embd_head_v = 32; + weights.n_head = 1; + weights.n_head_kv = 1; + weights.n_embd = 32; + weights.n_capture_layers = 0; + weights.ssm_d_inner = 32; + weights.ssm_d_state = 1; + weights.ssm_dt_rank = 1; + weights.ssm_n_group = 1; + weights.ssm_d_conv = 2; + + auto check_type = [&](bool f32_checkpoints, ggml_type expected) { + dflash::common::TargetCache cache; + const bool ok = dflash::common::create_target_cache_partial( + weights, /*max_ctx=*/1, /*max_verify_tokens=*/2, backend, cache, + /*prefill_only=*/false, /*layer_begin=*/0, /*layer_end=*/2, + /*allocate_target_feat=*/false, /*ctx_alloc=*/0, + /*f32_ssm_intermediates=*/f32_checkpoints); + CHECK(ok); + if (ok) { + CHECK(cache.ssm_intermediate.size() == 1); + CHECK(cache.ssm_intermediate[0] != nullptr); + if (cache.ssm_intermediate[0]) { + CHECK(cache.ssm_intermediate[0]->type == expected); + } + } + dflash::common::free_target_cache(cache); + }; + + check_type(false, GGML_TYPE_Q8_0); + check_type(true, GGML_TYPE_F32); + ggml_backend_free(backend); +} + TEST_CASE(ChainRollbackPolicyFixture, diagnostics_accumulator_and_print_contract) { const luce_test::ScopedEnvVar checkpoint("DFLASH_SINGLE_CHAIN_CHECKPOINT_F32", nullptr); const luce_test::ScopedEnvVar threshold("DFLASH_FAST_ROLLBACK_THRESHOLD", nullptr); diff --git a/server/test/test_dflash.cpp b/server/test/test_dflash.cpp index fc6a341f4..ef6ba86a8 100644 --- a/server/test/test_dflash.cpp +++ b/server/test/test_dflash.cpp @@ -23,6 +23,7 @@ #include "draft_graph.h" #include "qwen3_drafter.h" #include "gpu_runtime_compat.h" +#include "chain_rollback_policy.h" #include "laguna_daemon.h" // arch dispatch - laguna targets are served by // dflash::common::run_laguna_daemon() instead of the // qwen35 + DFlash + DDTree pipeline below. @@ -369,7 +370,11 @@ static int run_target_layer_split_harness( shard.backend, shard.cache, /*prefill_only=*/!run_dflash, shard.layer_begin, shard.layer_end, - allocate_target_feat)) { + allocate_target_feat, + /*ctx_alloc=*/0, + /*f32_ssm_intermediates=*/ + run_dflash && + split_chain_fast_rollback_enabled())) { std::fprintf(stderr, "target-split cache gpu=%d: %s\n", shard.gpu, dflash27b_last_error()); free_qwen35_layer_split_shards(shards); @@ -2139,11 +2144,13 @@ int main(int argc, char ** argv) { psg2 = StepGraph{}; migrate_prefill_cache(w, max_ctx, max_verify_tokens, target_backend, cache); - snapshot_ssm_state(cache); + check(snapshot_ssm_state(cache, target_backend), + "snapshot recurrent state succeeded"); std::vector logits_full(vocab_t), logits_win(vocab_t); bool ok = decode_one(psg2, 512, lt2, 512, 0, logits_full.data()); check(ok, "decode full-attention succeeded"); - restore_ssm_state(cache); + check(restore_ssm_state(cache, target_backend), + "restore recurrent state succeeded"); ok = decode_one(psg2, 512, lt2, 512, 2048, logits_win.data()); check(ok, "decode window=2048 succeeded"); @@ -2175,11 +2182,13 @@ int main(int argc, char ** argv) { psg3 = StepGraph{}; migrate_prefill_cache(w, max_ctx, max_verify_tokens, target_backend, cache); - snapshot_ssm_state(cache); + check(snapshot_ssm_state(cache, target_backend), + "snapshot recurrent state succeeded"); std::vector logits_full(vocab_t), logits_win(vocab_t); bool ok = decode_one(psg3, 4096, lt3, 4096, 0, logits_full.data()); check(ok, "decode full-attention succeeded"); - restore_ssm_state(cache); + check(restore_ssm_state(cache, target_backend), + "restore recurrent state succeeded"); ok = decode_one(psg3, 4096, lt3, 4096, 1024, logits_win.data()); check(ok, "decode window=1024 succeeded"); @@ -3321,7 +3330,10 @@ int main(int argc, char ** argv) { // gated_delta_net kernel captures per-step intermediate states, so // we don't need a pre-verify snapshot to restore from). if (!fast_rollback) { - snapshot_ssm_state(cache); + if (!snapshot_ssm_state(cache, target_backend)) { + std::fprintf(stderr, "snapshot recurrent state failed\n"); + return 1; + } } auto T_snap = sync_us(); tt_snap += std::chrono::duration(T_snap - T_draft_logits).count(); @@ -4033,7 +4045,10 @@ int main(int argc, char ** argv) { if (hit_eos) break; } else { // ── Legacy replay path ── - restore_ssm_state(cache); + if (!restore_ssm_state(cache, target_backend)) { + std::fprintf(stderr, "restore recurrent state failed\n"); + return 1; + } auto T_restore = sync_us(); tt_restore += std::chrono::duration(T_restore - T_accept).count(); std::vector replay_tok(commit_n); diff --git a/server/test/test_qwen35_split_tree_guard.cpp b/server/test/test_qwen35_split_tree_guard.cpp new file mode 100644 index 000000000..f348a7de3 --- /dev/null +++ b/server/test/test_qwen35_split_tree_guard.cpp @@ -0,0 +1,51 @@ +#include "qwen35_layer_split_tree_guard.h" + +#include +#include +#include + +using dflash::common::qwen35_split_run_if_root_inclusive_pure_chain; + +namespace { + +bool mark_executed(void * context) { + *static_cast(context) = true; + return true; +} + +bool expect(const char * name, const std::vector & parents, + std::size_t n_actual, bool expected) { + bool executed = false; + const bool result = qwen35_split_run_if_root_inclusive_pure_chain( + parents.data(), parents.size(), n_actual, mark_executed, &executed); + const bool ok = result == expected && executed == expected; + std::fprintf(stderr, + "qwen35_split_tree_guard case=%s result=%d sentinel_executed=%d expected=%d pass=%d\n", + name, result ? 1 : 0, executed ? 1 : 0, expected ? 1 : 0, ok ? 1 : 0); + return ok; +} + +} // namespace + +int main() { + bool ok = true; + ok &= expect("single_root", {-1}, 1, true); + ok &= expect("pure_chain", {-1, 0, 1, 2, 3, 4}, 6, true); + ok &= expect("malformed_root", {0, 0, 1}, 3, false); + ok &= expect("skipped_parent", {-1, 0, 0}, 3, false); + ok &= expect("sibling", {-1, 0, 0, 2}, 4, false); + ok &= expect("cycle_or_forward_parent", {-1, 2, 1}, 3, false); + ok &= expect("out_of_range", {-1, 0, 7}, 3, false); + ok &= expect("truncated", {-1, 0}, 3, false); + + bool null_executed = false; + const bool null_ok = !qwen35_split_run_if_root_inclusive_pure_chain( + nullptr, 0, 1, mark_executed, &null_executed) && !null_executed; + std::fprintf(stderr, + "qwen35_split_tree_guard case=null result=%d sentinel_executed=%d pass=%d\n", + 0, null_executed ? 1 : 0, null_ok ? 1 : 0); + ok &= null_ok; + + std::fprintf(stderr, "qwen35_split_tree_guard overall=%s\n", ok ? "PASS" : "FAIL"); + return ok ? 0 : 1; +} diff --git a/server/test/test_recurrent_snapshot.cpp b/server/test/test_recurrent_snapshot.cpp new file mode 100644 index 000000000..a9b247caf --- /dev/null +++ b/server/test/test_recurrent_snapshot.cpp @@ -0,0 +1,109 @@ +#include "internal.h" + +#include "ggml-backend.h" +#include "ggml-cpu.h" +#include "ggml.h" + +#include +#include + +using dflash::common::TargetCache; +using dflash::common::restore_ssm_state; +using dflash::common::snapshot_ssm_state; + +static int failures = 0; + +#define CHECK(expr) do { \ + if (!(expr)) { \ + std::fprintf(stderr, "FAIL %s:%d: %s\n", __FILE__, __LINE__, #expr); \ + failures++; \ + } \ +} while (0) + +static void set_tensor(ggml_tensor * tensor, const std::vector & values) { + CHECK(ggml_nelements(tensor) == (int64_t)values.size()); + ggml_backend_tensor_set(tensor, values.data(), 0, + values.size() * sizeof(float)); +} + +static std::vector get_tensor(const ggml_tensor * tensor) { + std::vector values((size_t)ggml_nelements(tensor)); + ggml_backend_tensor_get(tensor, values.data(), 0, + values.size() * sizeof(float)); + return values; +} + +int main() { + ggml_backend_t backend = ggml_backend_cpu_init(); + CHECK(backend != nullptr); + if (!backend) return 1; + + ggml_init_params params{}; + params.mem_size = 8 * ggml_tensor_overhead(); + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + CHECK(ctx != nullptr); + if (!ctx) { + ggml_backend_free(backend); + return 1; + } + + ggml_tensor * ssm = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 4, 3); + ggml_tensor * ssm_snap = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 4, 3); + ggml_tensor * conv = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 5, 2); + ggml_tensor * conv_snap = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 5, 2); + ggml_backend_buffer_t buffer = ggml_backend_alloc_ctx_tensors(ctx, backend); + CHECK(buffer != nullptr); + if (!buffer) { + ggml_free(ctx); + ggml_backend_free(backend); + return 1; + } + + TargetCache cache; + cache.ssm_state = {ssm, nullptr}; + cache.ssm_state_snap = {ssm_snap, nullptr}; + cache.conv_state = {conv, nullptr}; + cache.conv_state_snap = {conv_snap, nullptr}; + + const std::vector ssm_original = { + 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, + }; + const std::vector conv_original = { + 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, + }; + const std::vector ssm_mutated(ssm_original.size(), -1.0f); + const std::vector conv_mutated(conv_original.size(), -2.0f); + + set_tensor(ssm, ssm_original); + set_tensor(conv, conv_original); + CHECK(snapshot_ssm_state(cache, backend)); + set_tensor(ssm, ssm_mutated); + set_tensor(conv, conv_mutated); + CHECK(restore_ssm_state(cache, backend)); + CHECK(get_tensor(ssm) == ssm_original); + CHECK(get_tensor(conv) == conv_original); + + // A partial shard may omit a complete recurrent-state quartet, but an + // asymmetric quartet must fail validation before any copy is queued. + set_tensor(ssm, ssm_mutated); + cache.conv_state_snap[0] = nullptr; + CHECK(!snapshot_ssm_state(cache, backend)); + CHECK(get_tensor(ssm_snap) == ssm_original); + cache.conv_state_snap[0] = conv_snap; + + CHECK(!snapshot_ssm_state(cache, nullptr)); + cache.ssm_state_snap.pop_back(); + CHECK(!restore_ssm_state(cache, backend)); + + ggml_backend_buffer_free(buffer); + ggml_free(ctx); + ggml_backend_free(backend); + + if (failures != 0) { + std::fprintf(stderr, "%d recurrent snapshot test(s) failed\n", failures); + return 1; + } + std::printf("recurrent snapshot tests passed\n"); + return 0; +}