diff --git a/server/CMakeLists.txt b/server/CMakeLists.txt index 46ecfda81..fe2e06301 100644 --- a/server/CMakeLists.txt +++ b/server/CMakeLists.txt @@ -428,6 +428,7 @@ add_library(dflash_common STATIC src/kv_cache.cpp src/kv_quant.cpp src/delta_net_chunked.cpp + src/delta_net_specla.cpp # Laguna-XS.2 (Poolside) target arch src/laguna/laguna_target_loader.cpp src/laguna/laguna_target_graph.cpp @@ -592,6 +593,10 @@ if(DFLASH27B_GPU_BACKEND STREQUAL "hip") target_sources(dflash_common PRIVATE src/common/geometric_draft_topk_cuda.cu) set_source_files_properties(src/common/geometric_draft_topk_cuda.cu PROPERTIES LANGUAGE HIP) + # Fused SpecLA accepted-state commit (one launch for all delta layers). + target_sources(dflash_common PRIVATE src/common/specla_commit_cuda.cu) + set_source_files_properties(src/common/specla_commit_cuda.cu + PROPERTIES LANGUAGE HIP) # PUBLIC so test consumers (test_dflash / test_draft_topk_cuda) also take the # GPU draft top-K path instead of the CPU fallback. target_compile_definitions(dflash_common PUBLIC DFLASH27B_HAVE_DRAFT_TOPK=1) @@ -626,7 +631,8 @@ elseif(DFLASH27B_GPU_BACKEND STREQUAL "cuda") target_sources(dflash_common PRIVATE src/flashprefill_select.cpp src/flashprefill.cpp - src/common/geometric_draft_topk_cuda.cu) + src/common/geometric_draft_topk_cuda.cu + src/common/specla_commit_cuda.cu) # PUBLIC so consumers (e.g. the test_dflash executable) also see the macro # and take the GPU draft top-K path instead of the CPU fallback. Same macro # name as the HIP branch above (backend-neutral). @@ -1518,6 +1524,18 @@ if(DFLASH27B_TESTS) list(APPEND _raw_unit_test_targets test_recurrent_snapshot) endif() + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_delta_net_specla.cpp") + # GPU parity test: SpecLA topology-masked verify + factor-based state + # reconstruction vs the fused sequential gated-delta-net kernel. + # Exits 77 (ctest SKIP) when no GPU is present. + add_executable(test_delta_net_specla test/test_delta_net_specla.cpp) + target_include_directories(test_delta_net_specla PRIVATE + ${DFLASH27B_SRC_INCLUDE_DIRS}) + target_link_libraries(test_delta_net_specla PRIVATE + dflash_common ggml ${DFLASH27B_GGML_BACKEND_TARGET}) + list(APPEND _raw_unit_test_targets test_delta_net_specla) + endif() + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_server_unit.cpp") set(_server_unit_sources test/test_unit_main.cpp @@ -1530,6 +1548,7 @@ if(DFLASH27B_TESTS) test/test_admission.cpp test/test_restore_delta.cpp test/test_chain_rollback_policy.cpp + test/test_ddtree_tau.cpp test/test_anchor_transitive.cpp test/test_drafter_early_exit_score_range.cpp test/test_drafter_tail_capture_guard.cpp diff --git a/server/README.md b/server/README.md index 1e96b697d..2ae61f481 100644 --- a/server/README.md +++ b/server/README.md @@ -170,6 +170,12 @@ Run it directly: --model-name luce-dflash ``` +Use `--specla` to enable speculative linear-attention verification when the +target supports it. The runtime chooses a compatible proposal adapter; the +current monolithic Qwen3.5/Qwen3.6 path uses DDTree with tested defaults. +Hardware- or checkpoint-specific overrides remain available through +`--ddtree-budget`, `--ddtree-tau`, `--specla-top-k`, and `--draft-swa`. + ### Compression proxy mode `dflash_server` can run as a **PFlash compression proxy** in front of any diff --git a/server/deps/llama.cpp/ggml/include/ggml.h b/server/deps/llama.cpp/ggml/include/ggml.h index acac40c1f..a83bac9fa 100644 --- a/server/deps/llama.cpp/ggml/include/ggml.h +++ b/server/deps/llama.cpp/ggml/include/ggml.h @@ -2741,6 +2741,26 @@ extern "C" { struct ggml_tensor * c, struct ggml_tensor * parent_ids); + // SpecLA heavy-light convolution. Applies compact accepted inputs to the + // durable conv state, then verifies the current tree without committing + // speculative inputs. Current input factors are written directly to the + // persistent double-buffer selected by factor_ptrs/layer/bank; the result + // packs [conv output | boundary windows] and already includes SiLU. + GGML_API struct ggml_tensor * ggml_ssm_conv_specla( + struct ggml_context * ctx, + struct ggml_tensor * x, + struct ggml_tensor * c, + struct ggml_tensor * state, + struct ggml_tensor * hld, + struct ggml_tensor * factor_ptrs, + int n_layers, + int layer, + int pending_bank, + int n_boundaries, + int n_chains, + int n_waves, + int max_parallel_chains); + GGML_API struct ggml_tensor * ggml_ssm_scan( struct ggml_context * ctx, struct ggml_tensor * s, @@ -2924,6 +2944,28 @@ extern "C" { struct ggml_tensor * parent_ids, struct ggml_tensor * persist_inter); + // SpecLA state-resident heavy-light verify. The kernel applies the compact + // factors accepted in the preceding step, writes only that committed base + // state, then verifies the current HLD chains while writing raw + // (k, delta, log-gate) factors directly to the opposite persistent bank. + GGML_API struct ggml_tensor * ggml_gated_delta_net_specla( + struct ggml_context * ctx, + struct ggml_tensor * q, + struct ggml_tensor * k, + struct ggml_tensor * v, + struct ggml_tensor * g, + struct ggml_tensor * beta, + struct ggml_tensor * state, + struct ggml_tensor * hld, + struct ggml_tensor * factor_ptrs, + int n_layers, + int layer, + int pending_bank, + int n_boundaries, + int n_chains, + int n_waves, + int max_parallel_chains); + // custom operators typedef void (*ggml_custom1_op_t)(struct ggml_tensor * dst , const struct ggml_tensor * a, int ith, int nth, void * userdata); diff --git a/server/deps/llama.cpp/ggml/src/ggml-cpu/ggml-cpu.cpp b/server/deps/llama.cpp/ggml/src/ggml-cpu/ggml-cpu.cpp index d31c7e985..5d6e93a68 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cpu/ggml-cpu.cpp +++ b/server/deps/llama.cpp/ggml/src/ggml-cpu/ggml-cpu.cpp @@ -472,6 +472,15 @@ static bool ggml_backend_cpu_device_supports_op(ggml_backend_dev_t dev, const st return !ggml_flash_attn_ext_is_ds4(op); case GGML_OP_PAGED_ATTN: return false; + case GGML_OP_SSM_CONV: + // The Specla layout (op param 0 == 1) needs the packed HLD state and + // is only supported by the CUDA kernel; the generic CPU kernel would + // silently compute garbage. + return ggml_get_op_params_i32(op, 0) != 1; + case GGML_OP_GATED_DELTA_NET: + // The Specla GDN variant (op param 2 == 1) is stateful via HLD and is + // only supported by the CUDA kernel. + return ggml_get_op_params_i32(op, 2) != 1; case GGML_OP_OUT_PROD: return (src0->type == GGML_TYPE_F32 || (ggml_is_quantized(src0->type) && src0->ne[2] == src1->ne[2] && src0->ne[3] == src1->ne[3])) && src1->type == GGML_TYPE_F32 && op->type == GGML_TYPE_F32; diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/gated_delta_net.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/gated_delta_net.cu index 76f2de9da..afd328070 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/gated_delta_net.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/gated_delta_net.cu @@ -587,7 +587,226 @@ static void launch_gated_delta_net( } } +template +__global__ void gated_delta_net_specla_hld_cuda( + const float * __restrict__ q, + const float * __restrict__ k, + const float * __restrict__ v, + const float * __restrict__ g, + const float * __restrict__ beta, + float * __restrict__ durable_state, + const int * __restrict__ meta, + const int64_t * __restrict__ factor_ptrs, + float * __restrict__ packed, + int64_t H, + int64_t n_tokens, + int n_layers, + int layer, + int pending_bank, + int64_t sq1, + int64_t sq2, + int64_t sv1, + int64_t sv2, + int64_t sb1, + int64_t sb2, + int n_chains, + int wave, + float scale) { + const int h_idx = blockIdx.x; + const int wave_chain = blockIdx.y; + constexpr int warp_size = + ggml_cuda_get_physical_warp_size() < S_v ? + ggml_cuda_get_physical_warp_size() : S_v; + constexpr int rows_per_lane = (S_v + warp_size - 1) / warp_size; + const int lane = threadIdx.x; + const int col = blockIdx.z * blockDim.y + threadIdx.y; + if (col >= S_v) return; + + const int order_off = meta[6]; + const int offsets_off = meta[7]; + const int parent_off = meta[8]; + const int boundary_off = meta[9]; + const int wave_off = meta[10]; + int chain = 0; + while (chain < n_chains && meta[wave_off + chain] < wave) ++chain; + chain += wave_chain; + if (chain >= n_chains || meta[wave_off + chain] != wave) return; + + const int64_t plane_offset = (int64_t)h_idx*S_v*S_v; + float * state_plane = durable_state + plane_offset; + const int pbase = pending_bank*4; + const int cbase = (1 - pending_bank)*4; + const float * pending_k = (const float *)(uintptr_t)factor_ptrs[pbase + 0]; + const float * pending_v = (const float *)(uintptr_t)factor_ptrs[pbase + 1]; + const float * pending_g = (const float *)(uintptr_t)factor_ptrs[pbase + 2]; + float * current_k = (float *)(uintptr_t)factor_ptrs[cbase + 0]; + float * current_v = (float *)(uintptr_t)factor_ptrs[cbase + 1]; + float * current_g = (float *)(uintptr_t)factor_ptrs[cbase + 2]; + float state_shard[rows_per_lane]; + const int parent_boundary = meta[parent_off + chain]; + const int64_t attn_elems = (int64_t)S_v*H*n_tokens; + const int64_t boundary_base = attn_elems; + + if (parent_boundary < 0) { +#pragma unroll + for (int r = 0; r < rows_per_lane; ++r) { + const int row = r*warp_size + lane; + state_shard[r] = state_plane[(int64_t)col*S_v + row]; + } + + // Delayed commit of the preceding accepted path. Factors have already + // been compacted into path order, so this is the exact serial + // recurrence and does not touch any rejected branch. + const int pending_count = meta[5]; + for (int t = 0; t < pending_count; ++t) { + const int64_t th = ((int64_t)t*n_layers + layer)*H + h_idx; + const float g_val = expf(pending_g[th]); + const float delta = pending_v[th*S_v + col]; +#pragma unroll + for (int r = 0; r < rows_per_lane; ++r) { + const int row = r*warp_size + lane; + const float k_val = pending_k[th*S_v + row]; + state_shard[r] = fmaf(k_val, delta, g_val*state_shard[r]); + } + } +#pragma unroll + for (int r = 0; r < rows_per_lane; ++r) { + const int row = r*warp_size + lane; + state_plane[(int64_t)col*S_v + row] = state_shard[r]; + } + } else { + const float * boundary = packed + boundary_base + + ((int64_t)parent_boundary*H + h_idx)*S_v*S_v + + (int64_t)col*S_v; +#pragma unroll + for (int r = 0; r < rows_per_lane; ++r) { + const int row = r*warp_size + lane; + state_shard[r] = boundary[row]; + } + } + + const int begin = meta[offsets_off + chain]; + const int end = meta[offsets_off + chain + 1]; + for (int p = begin; p < end; ++p) { + const int node = meta[order_off + p]; + const float * q_t = q + (int64_t)node*sq2 + (int64_t)h_idx*sq1; + const float * k_t = k + (int64_t)node*sq2 + (int64_t)h_idx*sq1; + const float * v_t = v + (int64_t)node*sv2 + (int64_t)h_idx*sv1; + const int64_t gb = (int64_t)node*sb2 + (int64_t)h_idx*sb1; + const float g_log = g[gb]; + const float g_val = expf(g_log); + const float beta_val = beta[gb]; + + float kv_partial = 0.0f; + float k_reg[rows_per_lane]; + float q_reg[rows_per_lane]; +#pragma unroll + for (int r = 0; r < rows_per_lane; ++r) { + const int row = r*warp_size + lane; + k_reg[r] = k_t[row]; + q_reg[r] = q_t[row]; + kv_partial += state_shard[r]*k_reg[r]; + } + const float kv_col = warp_reduce_sum(kv_partial); + const float delta = (v_t[col] - g_val*kv_col)*beta_val; + + float attn_partial = 0.0f; +#pragma unroll + for (int r = 0; r < rows_per_lane; ++r) { + state_shard[r] = fmaf(k_reg[r], delta, g_val*state_shard[r]); + attn_partial += state_shard[r]*q_reg[r]; + } + const float attn_col = warp_reduce_sum(attn_partial); + if (lane == 0) { + packed[((int64_t)node*H + h_idx)*S_v + col] = + attn_col*scale; + const int64_t nh = ((int64_t)node*n_layers + layer)*H + h_idx; + current_v[nh*S_v + col] = delta; + if (col == 0) current_g[nh] = g_log; + } + if (col == 0) { + const int64_t nh = ((int64_t)node*n_layers + layer)*H + h_idx; +#pragma unroll + for (int r = 0; r < rows_per_lane; ++r) { + const int row = r*warp_size + lane; + current_k[nh*S_v + row] = k_reg[r]; + } + } + + const int boundary_slot = meta[boundary_off + node]; + if (boundary_slot >= 0) { + float * boundary = packed + boundary_base + + ((int64_t)boundary_slot*H + h_idx)*S_v*S_v + + (int64_t)col*S_v; +#pragma unroll + for (int r = 0; r < rows_per_lane; ++r) { + const int row = r*warp_size + lane; + boundary[row] = state_shard[r]; + } + } + } +} + +static void launch_gated_delta_net_specla( + ggml_backend_cuda_context & ctx, + ggml_tensor * dst) { + ggml_tensor * q = dst->src[0]; + ggml_tensor * k = dst->src[1]; + ggml_tensor * v = dst->src[2]; + ggml_tensor * g = dst->src[3]; + ggml_tensor * beta = dst->src[4]; + ggml_tensor * state = dst->src[5]; + ggml_tensor * hld = dst->src[6]; + ggml_tensor * factor_ptrs = dst->src[7]; + const int S_v = (int)v->ne[0]; + const int H = (int)v->ne[1]; + const int n_tokens = (int)v->ne[2]; + const int n_chains = ggml_get_op_params_i32(dst, 4); + const int n_waves = ggml_get_op_params_i32(dst, 5); + const int n_layers = ggml_get_op_params_i32(dst, 6); + const int layer = ggml_get_op_params_i32(dst, 7); + const int pending_bank = ggml_get_op_params_i32(dst, 8); + const int max_parallel_chains = ggml_get_op_params_i32(dst, 9); + GGML_ASSERT(v->ne[3] == 1 && g->ne[0] == 1); + GGML_ASSERT(hld->type == GGML_TYPE_I32 && n_chains > 0 && n_waves > 0); + GGML_ASSERT(ggml_is_contiguous(state)); + + const int warp_size = ggml_cuda_info().devices[ggml_cuda_get_device()].warp_size; + constexpr int num_warps = 4; + const dim3 block(warp_size <= S_v ? warp_size : S_v, num_warps, 1); + const dim3 grid((unsigned)H, (unsigned)max_parallel_chains, + (unsigned)((S_v + num_warps - 1)/num_warps)); + const float scale = 1.0f/sqrtf((float)S_v); + auto launch = [&](auto SV, int wave) { + constexpr int kSV = decltype(SV)::value; + gated_delta_net_specla_hld_cuda<<>>( + (const float *)q->data, (const float *)k->data, + (const float *)v->data, (const float *)g->data, + (const float *)beta->data, (float *)state->data, + (const int *)hld->data, + (const int64_t *)factor_ptrs->data, + (float *)dst->data, H, n_tokens, n_layers, layer, pending_bank, + q->nb[1]/sizeof(float), q->nb[2]/sizeof(float), + v->nb[1]/sizeof(float), v->nb[2]/sizeof(float), + beta->nb[1]/sizeof(float), beta->nb[2]/sizeof(float), + n_chains, wave, scale); + }; + for (int wave = 0; wave < n_waves; ++wave) { + switch (S_v) { + case 16: launch(std::integral_constant{}, wave); break; + case 32: launch(std::integral_constant{}, wave); break; + case 64: launch(std::integral_constant{}, wave); break; + case 128: launch(std::integral_constant{}, wave); break; + default: GGML_ABORT("Unsupported SpecLA GDN state size"); + } + } +} + void ggml_cuda_op_gated_delta_net(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { + if (ggml_get_op_params_i32(dst, 2) == 1) { + launch_gated_delta_net_specla(ctx, dst); + return; + } ggml_tensor * src_q = dst->src[0]; ggml_tensor * src_k = dst->src[1]; ggml_tensor * src_v = dst->src[2]; diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu index 1c55d51d1..da0740634 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu @@ -4362,6 +4362,11 @@ static bool ggml_cuda_can_fuse(const struct ggml_cgraph * cgraph, const ggml_tensor * ssm_conv = cgraph->nodes[node_idx]; const ggml_tensor * silu = cgraph->nodes[node_idx+1]; + if (ggml_get_op_params_i32(ssm_conv, 0) == 1) { + // the Specla ssm_conv kernel applies SiLU itself + return false; + } + if (ssm_conv->type != GGML_TYPE_F32 || silu->type != GGML_TYPE_F32) { return false; } @@ -6157,6 +6162,11 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g } } case GGML_OP_SSM_CONV: { + if (ggml_get_op_params_i32(op, 0) == 1) { + // Specla layout: x is [d_inner, n_tokens], so d_inner = ne[0] + // and the kernel guards the final partial 128-channel block. + return true; + } // assumes d_inner % threads == 0 return op->src[0]->ne[1] % 128 == 0; } diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/ssm-conv.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/ssm-conv.cu index e6ce26f72..dc70cb9ea 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/ssm-conv.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/ssm-conv.cu @@ -244,7 +244,148 @@ static void ssm_conv_f32_cuda(const float * src0, const float * src1, const int } } +template +static __global__ void ssm_conv_specla_hld_f32( + const float * __restrict__ x, // [d_inner, n_t] + const float * __restrict__ weight, // [d_conv, d_inner] + float * __restrict__ state, // [d_conv-1, d_inner] + const int * __restrict__ meta, + const int64_t * __restrict__ factor_ptrs, + float * __restrict__ packed, + int d_inner, + int n_t, + int n_layers, + int layer, + int pending_bank, + int n_chains, + int wave) { + const int wave_chain = blockIdx.x; + const int channel = blockIdx.y * blockDim.x + threadIdx.x; + if (channel >= d_inner) return; + + const int order_off = meta[6]; + const int offsets_off = meta[7]; + const int parent_off = meta[8]; + const int boundary_off = meta[9]; + const int wave_off = meta[10]; + int chain = 0; + while (chain < n_chains && meta[wave_off + chain] < wave) ++chain; + chain += wave_chain; + if (chain >= n_chains || meta[wave_off + chain] != wave) return; + + float window[d_conv - 1]; + const float * pending_x = (const float *)(uintptr_t) + factor_ptrs[pending_bank*4 + 3]; + float * current_x = (float *)(uintptr_t) + factor_ptrs[(1 - pending_bank)*4 + 3]; + const int parent_boundary = meta[parent_off + chain]; + if (parent_boundary < 0) { +#pragma unroll + for (int j = 0; j < d_conv - 1; ++j) { + window[j] = state[(size_t)channel * (d_conv - 1) + j]; + } + // Delayed commit: compact accepted inputs from the preceding verify + // are consumed before the current root chain. Only this committed + // window is written to durable state. + const int pending_count = meta[5]; + for (int t = 0; t < pending_count; ++t) { +#pragma unroll + for (int j = 0; j < d_conv - 2; ++j) window[j] = window[j + 1]; + window[d_conv - 2] = pending_x[ + (size_t)channel + (size_t)d_inner*(layer + (size_t)n_layers*t)]; + } +#pragma unroll + for (int j = 0; j < d_conv - 1; ++j) { + state[(size_t)channel * (d_conv - 1) + j] = window[j]; + } + } else { + const size_t boundary_base = + ((size_t)n_t + (size_t)parent_boundary*(d_conv - 1))*d_inner; +#pragma unroll + for (int j = 0; j < d_conv - 1; ++j) { + window[j] = packed[boundary_base + (size_t)j*d_inner + channel]; + } + } + + const int begin = meta[offsets_off + chain]; + const int end = meta[offsets_off + chain + 1]; + for (int p = begin; p < end; ++p) { + const int node = meta[order_off + p]; + const float x_val = x[(size_t)node*d_inner + channel]; + float sum = 0.0f; +#pragma unroll + for (int j = 0; j < d_conv - 1; ++j) { + sum += window[j] * weight[(size_t)channel*d_conv + j]; + } + sum += x_val * weight[(size_t)channel*d_conv + d_conv - 1]; + packed[(size_t)node*d_inner + channel] = + ggml_cuda_op_silu_single(sum); + current_x[(size_t)channel + + (size_t)d_inner*(layer + (size_t)n_layers*node)] = x_val; + +#pragma unroll + for (int j = 0; j < d_conv - 2; ++j) window[j] = window[j + 1]; + window[d_conv - 2] = x_val; + const int boundary = meta[boundary_off + node]; + if (boundary >= 0) { + const size_t boundary_base = + ((size_t)n_t + (size_t)boundary*(d_conv - 1))*d_inner; +#pragma unroll + for (int j = 0; j < d_conv - 1; ++j) { + packed[boundary_base + (size_t)j*d_inner + channel] = window[j]; + } + } + } +} + +static void ssm_conv_specla_hld_cuda(ggml_backend_cuda_context & ctx, + ggml_tensor * dst) { + ggml_tensor * x = dst->src[0]; + ggml_tensor * weight = dst->src[1]; + ggml_tensor * state = dst->src[2]; + ggml_tensor * hld = dst->src[3]; + ggml_tensor * factor_ptrs = dst->src[4]; + const int d_conv = (int)weight->ne[0]; + const int d_inner = (int)x->ne[0]; + const int n_t = (int)x->ne[1]; + const int n_chains = ggml_get_op_params_i32(dst, 2); + const int n_waves = ggml_get_op_params_i32(dst, 3); + const int n_layers = ggml_get_op_params_i32(dst, 4); + const int layer = ggml_get_op_params_i32(dst, 5); + const int pending_bank = ggml_get_op_params_i32(dst, 6); + const int max_parallel_chains = ggml_get_op_params_i32(dst, 7); + GGML_ASSERT(x->type == GGML_TYPE_F32 && weight->type == GGML_TYPE_F32); + GGML_ASSERT(state->type == GGML_TYPE_F32 && factor_ptrs->type == GGML_TYPE_I64); + GGML_ASSERT(hld->type == GGML_TYPE_I32 && n_chains > 0 && n_waves > 0); + + const dim3 block(128); + const dim3 grid((unsigned)max_parallel_chains, + (unsigned)((d_inner + 127)/128), 1); + auto launch = [&](auto DC, int wave) { + constexpr int kDC = decltype(DC)::value; + ssm_conv_specla_hld_f32<<>>( + (const float *)x->data, (const float *)weight->data, + (float *)state->data, (const int *)hld->data, + (const int64_t *)factor_ptrs->data, (float *)dst->data, + d_inner, n_t, n_layers, layer, pending_bank, n_chains, wave); + }; + for (int wave = 0; wave < n_waves; ++wave) { + switch (d_conv) { + case 3: launch(std::integral_constant{}, wave); break; + case 4: launch(std::integral_constant{}, wave); break; + case 5: launch(std::integral_constant{}, wave); break; + case 9: launch(std::integral_constant{}, wave); break; + default: GGML_ABORT("SpecLA ssm_conv supports kernel sizes 3, 4, 5, 9."); + } + } +} + void ggml_cuda_op_ssm_conv(ggml_backend_cuda_context & ctx, ggml_tensor * dst, ggml_tensor * silu_dst) { + if (ggml_get_op_params_i32(dst, 0) == 1) { + GGML_ASSERT(silu_dst == nullptr); + ssm_conv_specla_hld_cuda(ctx, dst); + return; + } const struct ggml_tensor * src0 = dst->src[0]; // conv_x const struct ggml_tensor * src1 = dst->src[1]; // conv1d.weight // dflash27b_ggml: optional src[2] = parent_ids (i32) enables tree mode diff --git a/server/deps/llama.cpp/ggml/src/ggml.c b/server/deps/llama.cpp/ggml/src/ggml.c index 86e4520d6..74780a434 100644 --- a/server/deps/llama.cpp/ggml/src/ggml.c +++ b/server/deps/llama.cpp/ggml/src/ggml.c @@ -5891,6 +5891,66 @@ struct ggml_tensor * ggml_ssm_conv_tree( return result; } +struct ggml_tensor * ggml_ssm_conv_specla( + struct ggml_context * ctx, + struct ggml_tensor * x, + struct ggml_tensor * c, + struct ggml_tensor * state, + struct ggml_tensor * hld, + struct ggml_tensor * factor_ptrs, + int n_layers, + int layer, + int pending_bank, + int n_boundaries, + int n_chains, + int n_waves, + int max_parallel_chains) { + GGML_ASSERT(ggml_is_3d(x)); + GGML_ASSERT(ggml_is_matrix(c)); + GGML_ASSERT(ggml_is_matrix(state)); + GGML_ASSERT(hld->type == GGML_TYPE_I32 && ggml_is_contiguous(hld)); + GGML_ASSERT(factor_ptrs->type == GGML_TYPE_I64 && + ggml_nelements(factor_ptrs) == 8 && + ggml_is_contiguous(factor_ptrs)); + GGML_ASSERT(x->type == GGML_TYPE_F32 && c->type == GGML_TYPE_F32 && + state->type == GGML_TYPE_F32); + // The CUDA kernel raw-indexes these tensors and does not consume strides. + GGML_ASSERT(ggml_is_contiguous(x)); + GGML_ASSERT(ggml_is_contiguous(c)); + GGML_ASSERT(ggml_is_contiguous(state)); + GGML_ASSERT(x->ne[0] == c->ne[1]); + GGML_ASSERT(state->ne[0] == c->ne[0] - 1 && state->ne[1] == c->ne[1]); + GGML_ASSERT(x->ne[2] == 1 && state->ne[2] == 1); + + const int64_t d_inner = x->ne[0]; + const int64_t n_t = x->ne[1]; + GGML_ASSERT(c->ne[0] == 3 || c->ne[0] == 4 || + c->ne[0] == 5 || c->ne[0] == 9); + GGML_ASSERT(n_layers > 0 && layer >= 0 && layer < n_layers); + GGML_ASSERT(pending_bank == 0 || pending_bank == 1); + GGML_ASSERT(n_boundaries >= 0 && n_chains > 0 && n_waves > 0); + GGML_ASSERT(max_parallel_chains > 0 && max_parallel_chains <= n_chains); + // n_boundaries is used only for output sizing below; the runtime boundary + // layout is packed in the HLD meta tensor. + const int64_t packed_rows = n_t + (c->ne[0] - 1)*n_boundaries; + struct ggml_tensor * result = + ggml_new_tensor_2d(ctx, GGML_TYPE_F32, d_inner, packed_rows); + ggml_set_op_params_i32(result, 0, 1); + ggml_set_op_params_i32(result, 2, n_chains); + ggml_set_op_params_i32(result, 3, n_waves); + ggml_set_op_params_i32(result, 4, n_layers); + ggml_set_op_params_i32(result, 5, layer); + ggml_set_op_params_i32(result, 6, pending_bank); + ggml_set_op_params_i32(result, 7, max_parallel_chains); + result->op = GGML_OP_SSM_CONV; + result->src[0] = x; + result->src[1] = c; + result->src[2] = state; + result->src[3] = hld; + result->src[4] = factor_ptrs; + return result; +} + // ggml_ssm_scan struct ggml_tensor * ggml_ssm_scan( @@ -6763,6 +6823,79 @@ struct ggml_tensor * ggml_gated_delta_net_tree_persist( return result; } +struct ggml_tensor * ggml_gated_delta_net_specla( + struct ggml_context * ctx, + struct ggml_tensor * q, + struct ggml_tensor * k, + struct ggml_tensor * v, + struct ggml_tensor * g, + struct ggml_tensor * beta, + struct ggml_tensor * state, + struct ggml_tensor * hld, + struct ggml_tensor * factor_ptrs, + int n_layers, + int layer, + int pending_bank, + int n_boundaries, + int n_chains, + int n_waves, + int max_parallel_chains) { + GGML_ASSERT(q->type == GGML_TYPE_F32 && k->type == GGML_TYPE_F32); + GGML_ASSERT(v->type == GGML_TYPE_F32 && g->type == GGML_TYPE_F32); + GGML_ASSERT(beta->type == GGML_TYPE_F32 && state->type == GGML_TYPE_F32); + GGML_ASSERT(hld->type == GGML_TYPE_I32 && ggml_is_contiguous(hld)); + GGML_ASSERT(factor_ptrs->type == GGML_TYPE_I64 && + ggml_nelements(factor_ptrs) == 8 && + ggml_is_contiguous(factor_ptrs)); + GGML_ASSERT(ggml_is_contiguous_rows(q)); + GGML_ASSERT(ggml_is_contiguous_rows(k)); + GGML_ASSERT(ggml_is_contiguous_rows(v)); + GGML_ASSERT(ggml_are_same_shape(q, k) && ggml_are_same_shape(q, v)); + // The CUDA launcher passes q strides for both q and k. + GGML_ASSERT(ggml_are_same_stride(q, k)); + GGML_ASSERT(ggml_is_contiguous(g) && ggml_is_contiguous(beta)); + GGML_ASSERT(ggml_is_contiguous(state)); + + const int64_t S = v->ne[0]; + const int64_t H = v->ne[1]; + const int64_t T = v->ne[2]; + GGML_ASSERT(S == 16 || S == 32 || S == 64 || S == 128); + GGML_ASSERT(v->ne[3] == 1); + GGML_ASSERT(ggml_are_same_shape(g, beta)); + GGML_ASSERT(g->ne[0] == 1 && g->ne[1] == H && + g->ne[2] == T && g->ne[3] == 1); + GGML_ASSERT(state->ne[0] == S && state->ne[1] == S && + state->ne[2] == H && state->ne[3] == 1); + GGML_ASSERT(ggml_nelements(state) == S*S*H); + GGML_ASSERT(n_layers > 0 && layer >= 0 && layer < n_layers); + GGML_ASSERT(pending_bank == 0 || pending_bank == 1); + GGML_ASSERT(n_boundaries >= 0 && n_chains > 0 && n_waves > 0); + GGML_ASSERT(max_parallel_chains > 0 && max_parallel_chains <= n_chains); + // n_boundaries is used only for output sizing below; the runtime boundary + // layout is packed in the HLD meta tensor. + const int64_t packed = S*H*T + (int64_t)n_boundaries*S*S*H; + struct ggml_tensor * result = + ggml_new_tensor_1d(ctx, GGML_TYPE_F32, packed); + ggml_set_op_params_i32(result, 1, 1); + ggml_set_op_params_i32(result, 2, 1); + ggml_set_op_params_i32(result, 4, n_chains); + ggml_set_op_params_i32(result, 5, n_waves); + ggml_set_op_params_i32(result, 6, n_layers); + ggml_set_op_params_i32(result, 7, layer); + ggml_set_op_params_i32(result, 8, pending_bank); + ggml_set_op_params_i32(result, 9, max_parallel_chains); + result->op = GGML_OP_GATED_DELTA_NET; + result->src[0] = q; + result->src[1] = k; + result->src[2] = v; + result->src[3] = g; + result->src[4] = beta; + result->src[5] = state; + result->src[6] = hld; + result->src[7] = factor_ptrs; + return result; +} + //////////////////////////////////////////////////////////////////////////////// struct ggml_hash_set ggml_hash_set_new(size_t size) { diff --git a/server/docs/SPECLA.md b/server/docs/SPECLA.md new file mode 100644 index 000000000..9f9c1c596 --- /dev/null +++ b/server/docs/SPECLA.md @@ -0,0 +1,203 @@ +# SpecLA for the Qwen3.6 Gated-DeltaNet target + +This is the Qwen runtime implementation of *SpecLA: Efficient Speculative +Decoding for Linear-Attention Models* (arXiv:2607.16673). It is enabled only +for the single-device Qwen target with exact fast rollback; tensor-parallel, +layer-split, and KVFlash targets keep their established state paths. + +The important distinction from the old experimental path is that the normal +route is now the paper's state-resident, chain-decomposed verifier with delayed +raw factors. The UT-factorized verifier and immediate DeltaConstruct commit +remain only as a compatibility/reference route. + +## Paper-to-code map + +### §4.1: state-resident serial verification + +`ggml_gated_delta_net_specla` processes every adjacent node in a chain inside +one kernel while its recurrent-state tile stays in registers. The kernel tiles +the value dimension and keeps the full key dimension inside a warp, avoiding a +cross-block reduction for `S^T q`. It reads durable state once and does not +write it between adjacent candidate tokens. + +`ggml_ssm_conv_specla` does the same for the causal depthwise-convolution +window. This matters for Qwen: correct GDN recurrence alone is insufficient if +siblings accidentally inherit convolution history from DFS neighbours. + +Chain verification is represented by one chain and one dependency wave, so it +uses these same kernels without tree-specific overhead. + +### §4.2: fully factorized tree verification + +`build_delta_net_specla` implements the topology-masked UT transform. Host +inputs provide strict-ancestor, inclusive-ancestor, and identity masks; the +builder returns candidate outputs plus corrected values and cumulative gates. +It is retained as a numerical reference and fallback for callers without an +HLD schedule. It is not the production Qwen route because its short-window +setup and reduction cost lose to the serial kernel at this shape. + +### §4.3: chain-decomposed hybrid verification + +`make_specla_hld_schedule` performs deterministic heavy-light decomposition: + +- the largest child subtree is the heavy continuation; +- heavy edges stay in one state-resident chain; +- light edges create chain boundaries; +- chains are grouped into dependency waves; and +- all ready chains in a wave launch in parallel. + +Only states at light-edge boundaries are materialized. A chain loads either +durable root state or its parent boundary, executes serially, emits per-node +outputs/factors, and writes only boundary states needed by later waves. Both +GDN and convolution consume the same packed parent topology. + +### §5.1: accepted-factor buffering + +Dense per-candidate recurrent checkpoints are replaced by two consolidated +FP32 factor banks. The HLD kernels write directly into the current bank: + +- normalized key `k`; +- the already-computed Delta-rule residual/update vector; +- log decay `g`; and +- raw convolution input. + +This is sufficient to replay the exact serial recurrence for an accepted path +without regenerating projections or storing a dense state per node. A chain +acceptance rotates banks with no device copy. A branch acceptance gathers its +arbitrary DFS indices into path order with one compaction kernel. + +### §5.2: delayed fused update and verify + +After selection, accepted factors remain pending. At the next verification, +each state-resident kernel: + +1. loads its durable state tile; +2. applies the preceding accepted GDN/conv factors; +3. writes that committed state once; +4. immediately verifies the current chain from the same live tile; and +5. records current candidates in the other bank. + +There is no standalone commit kernel or recurrent-state snapshot in the normal +loop. If generation ends, switches to autoregressive decode, or is cancelled +before another verification, `finish_speculative_state()` materializes the one +remaining pending path exactly once. + +The older immediate `specla_commit_accepted` implementation remains for the +factorized fallback and tests; it is not used by the normal HLD route. + +### §6.1: confidence-guided pruning + +`build_ddtree` and `build_ddtree_conditional` score a node by cumulative path +log probability and keep + +``` +q(v) >= q* - tau_tree +``` + +before applying the node budget. Expansion is best-first and the retained set +is ancestor-closed. A finite `--ddtree-tau` also shrinks the actual target +batch; pruned nodes are not replaced by fake padding in SpecLA mode. + +### §6.2: target-aligned drafting + +The Qwen DFlash draft consumes the target's recurrent execution features, so +the feature-alignment interface exists. Its released checkpoint is a five-layer +block-diffusion drafter, however, not the paper's specially trained one-layer +EAGLE drafter. + +The runtime includes exact prefix-conditioned tree construction. Set +`DFLASH_SPECLA_CONDITIONAL_DRAFT=1` to rerun the draft for every expanded +prefix. This is the faithful algorithmic experiment, but it is intentionally +off by default: per-node reruns of this five-layer draft cost more than a 27B +target verification. Realizing the paper's §6.2 speedup requires training the +small recurrent-feature EAGLE checkpoint; a runtime patch cannot synthesize +that model artifact. + +SpecLA defaults to the paper's top-k=4 tree width. Use +`--specla-top-k ` when a checkpoint or workload benefits from a different +width. `DFLASH_SPECLA_TOPK=` remains available for non-CLI harnesses. + +## Running + +```sh +# Enable SpecLA and let the runtime select a compatible proposal adapter. +# Qwen3.6 currently selects DDTree with budget 22, tau 6, and top-k 4. +# Older drafts without embedded sliding-window metadata may additionally need +# --draft-swa=2048. +build/test_dflash TARGET.gguf DRAFT.gguf prompt.bin 128 out.bin --specla + +# Expensive exact branch-conditioned drafting experiment. +DFLASH_SPECLA=1 DFLASH_SPECLA_CONDITIONAL_DRAFT=1 build/dflash_server ... +``` + +Tau 6 was best in the current ten-prompt probe and is the `--specla` default. +Use `--ddtree-tau` to tune it for the intended checkpoint and workload; a +tighter margin changed batching at numerically sensitive logits and was not +consistently faster. `--ddtree-budget`, `--specla-top-k`, and `--draft-swa` +remain explicit for the same setup-dependent reason. + +`DFLASH_SPECLA=1` remains a compatibility switch for non-CLI integrations. +`DFLASH_SPECLA_CONDITIONAL_DRAFT=1` and `DFLASH_SPECLA_FUSED_COMMIT=0` are +advanced algorithm/debug controls, not required for normal use. + +KVFlash uses a pager-backed attention cache that cannot migrate SpecLA factor +state. Combining `--specla` with `--kvflash ` therefore prints a +warning and falls back to ordinary DDTree verification; startup reports SpecLA +as off. + +SpecLA does not intrinsically require DDTree. It consumes speculative +candidates plus their parent topology; a chain is a degenerate tree. DDTree is +the proposal adapter currently connected for Qwen3.6. A future DSpark adapter +can feed the same user-facing `--specla` mode when its target recurrence and +factor-capture path are supported. + +## Correctness coverage + +- CPU tests cover cumulative-score pruning, ancestor closure, budget caps, + exact-prefix conditional queries, and deterministic HLD schedules. +- GPU tests compare HLD chains and trees with the sequential GDN reference at + the Qwen state shape, including non-empty pending factors and sibling + boundaries. +- Convolution tests cover tree ancestry, delayed pending windows, and direct + factor capture. +- Factorized UT output and DeltaConstruct remain cross-checked against the + serial recurrence. +- Runtime exit, tree sibling, greedy chain/tree, and sampled bonus paths all + use the same double-bank lifecycle; the final flush is GPU-tested directly. + +On the current gfx1151 probe, 16-node and 22-node HLD trees produced identical +128-token hashes even though their accepted boundaries differed. The legacy +F16 checkpoint rollback did not match that stream. Dynamic tree sizes can +still perturb near-tie logits in Qwen's full-attention layers; that is a +batched floating-point issue outside the GDN recurrence and is why confidence +margin changes require output checks. + +## Current Qwen3.6-27B measurements + +Target: Qwen3.6-27B Q4_K_M. Draft: five-layer DFlash Q8_0 with SWA 2048. +Workload: ten cached HumanEval-style prompts, 128 generated tokens, gfx1151. + +| route | tree | mean accepted/step | mean decode | +|---|---:|---:|---:| +| SpecLA off reference | 22, top-8 | 5.64 | 25.14 tok/s | +| completed HLD, no pruning | 22, top-4 | 5.58 | 26.05 tok/s | +| completed HLD, `tau=6` | 22, top-4 | 5.95 | **27.19 tok/s** | +| completed HLD, `tau=6` | 22, top-8 | 5.90 | 26.92 tok/s | + +The recommended route is 8.2% faster than the SpecLA-off reference, but that +reference runs the same 22-node tree at **top-8** while the recommended route +runs at **top-4**, so the headline number mixes the tree-width change into the +SpecLA/pruning effect. The comparable pairs are: at top-8, completed HLD with +`tau=6` reaches 26.92 tok/s versus 25.14 tok/s for the off reference (+7.1%); +within completed HLD, `tau=6` raises top-4 throughput from 26.05 to 27.19 +tok/s (+4.4%); and with HLD plus `tau=6` held fixed, top-4 is 1.0% faster than +top-8 (27.19 versus 26.92 tok/s). The table does not include a pair that +isolates HLD alone. On one fixed 22-node step, HLD reduced target verification +from 191.53 ms to 184.11 ms; the full-model gain is smaller than the paper's +GDN-1.3B result because Qwen has full-attention layers and large +projections/FFNs that HLD does not accelerate, and this draft's acceptance is +below the paper's best workloads. + +The paper reports 1.42x mixed, 1.70x GSM8K, and 1.06x HumanEval end-to-end +speedups on an H100 with a pure GDN-1.3B target and its trained EAGLE-style +drafter. Those figures are not directly transferable to this 27B hybrid model. diff --git a/server/scripts/bench_he.py b/server/scripts/bench_he.py index a4da24f6d..664c3964e 100644 --- a/server/scripts/bench_he.py +++ b/server/scripts/bench_he.py @@ -236,8 +236,10 @@ def tokenize_prompt(prompt: str, out_path: Path, tokenizer) -> int: def run_test_dflash(prompt_path: Path, n_gen: int, fast_rollback: bool, + specla: bool = False, ddtree_budget: int | None = None, ddtree_temp: float | None = None, + ddtree_tau: float | None = None, ddtree_no_chain_seed: bool = False, extra_args: list[str] | None = None, extra_env: dict[str, str] | None = None) -> dict: @@ -247,11 +249,15 @@ def run_test_dflash(prompt_path: Path, n_gen: int, fast_rollback: bool, ] if fast_rollback: cmd.append("--fast-rollback") + if specla: + cmd.append("--specla") if ddtree_budget is not None: cmd.append("--ddtree") cmd.append(f"--ddtree-budget={ddtree_budget}") if ddtree_temp is not None: cmd.append(f"--ddtree-temp={ddtree_temp}") + if ddtree_tau is not None: + cmd.append(f"--ddtree-tau={ddtree_tau}") if ddtree_no_chain_seed: cmd.append("--ddtree-no-chain-seed") if extra_args: @@ -315,8 +321,12 @@ def main(): ap.add_argument("--skip-tokenize", action="store_true") ap.add_argument("--ddtree-budget", type=int, default=None, help="Enable DDTree mode with this node budget (e.g. 15, 32, 64)") + ap.add_argument("--specla", action="store_true", + help="Enable SpecLA with its tested defaults") ap.add_argument("--ddtree-temp", type=float, default=None, help="Sharpen draft logits with this temperature (T<1 widens top-1/top-2 gap)") + ap.add_argument("--ddtree-tau", type=float, default=None, + help="SpecLA cumulative path-log-probability pruning margin") ap.add_argument("--ddtree-no-chain-seed", action="store_true", help="Use paper's pure best-first (no chain pre-seed)") ap.add_argument("--draft-feature-mirror", action="store_true", @@ -423,8 +433,10 @@ def main(): try: r = run_test_dflash(path, args.n_gen, fast_rollback=(args.mode == "fast" and not args.target_split_dflash), + specla=args.specla, ddtree_budget=args.ddtree_budget, ddtree_temp=args.ddtree_temp, + ddtree_tau=args.ddtree_tau, ddtree_no_chain_seed=args.ddtree_no_chain_seed, extra_args=extra_args, extra_env=extra_env) diff --git a/server/src/common/backend_args.h b/server/src/common/backend_args.h index ee7b55d33..d69749a9b 100644 --- a/server/src/common/backend_args.h +++ b/server/src/common/backend_args.h @@ -5,6 +5,8 @@ #pragma once +#include + #include "placement/placement_config.h" #include "placement/remote_draft_config.h" #include "placement/remote_target_shard_config.h" @@ -75,10 +77,12 @@ struct BackendArgs { int draft_ctx_max = 4096; bool fast_rollback = true; bool seq_verify = false; + bool specla_mode = false; bool ddtree_mode = false; int ddtree_budget = 22; float ddtree_temp = 1.0f; bool ddtree_chain_seed = true; + float ddtree_tau = std::numeric_limits::infinity(); int verify_width = 0; // chain spec verify width; 0 = adaptive bool use_feature_mirror = false; }; diff --git a/server/src/common/backend_factory.cpp b/server/src/common/backend_factory.cpp index 0d1ccc61b..813798a2d 100644 --- a/server/src/common/backend_factory.cpp +++ b/server/src/common/backend_factory.cpp @@ -282,6 +282,7 @@ std::unique_ptr create_backend( cfg.ddtree_budget = args.ddtree_budget; cfg.ddtree_temp = args.ddtree_temp; cfg.ddtree_chain_seed = args.ddtree_chain_seed; + cfg.ddtree_tau = args.ddtree_tau; cfg.use_feature_mirror = args.use_feature_mirror; auto backend = std::make_unique(cfg); @@ -308,6 +309,7 @@ std::unique_ptr create_backend( cfg.ddtree_budget = args.ddtree_budget; cfg.ddtree_temp = args.ddtree_temp; cfg.ddtree_chain_seed = args.ddtree_chain_seed; + cfg.ddtree_tau = args.ddtree_tau; cfg.use_feature_mirror = args.use_feature_mirror; auto backend = std::make_unique(cfg); diff --git a/server/src/common/chain_rollback_policy.h b/server/src/common/chain_rollback_policy.h index 7ad2b0d44..38874e87b 100644 --- a/server/src/common/chain_rollback_policy.h +++ b/server/src/common/chain_rollback_policy.h @@ -26,7 +26,8 @@ inline bool split_chain_fast_rollback_enabled() { } inline ChainRollbackPolicy resolve_chain_rollback_policy( - bool tensor_parallel = false) { + bool tensor_parallel = false, + bool exact_fast_rollback = false) { ChainRollbackPolicy policy; policy.checkpoint_f32 = env_flag_enabled("DFLASH_SINGLE_CHAIN_CHECKPOINT_F32"); policy.diagnostics = env_flag_enabled("DFLASH_SINGLE_CHAIN_ROLLBACK_DIAG"); @@ -48,6 +49,13 @@ inline ChainRollbackPolicy resolve_chain_rollback_policy( if (tensor_parallel) { policy.fast_rollback_threshold = 1; } + // Exact device-side rollback (SpecLA factor commit, or an equivalent + // implementation) is profitable from the first accepted token. This is + // an actual target capability, not an environment flag: the requested + // mode may be unavailable for the active cache/backend. + if (exact_fast_rollback) { + policy.fast_rollback_threshold = 1; + } return policy; } diff --git a/server/src/common/ddtree.cpp b/server/src/common/ddtree.cpp index 08ca33464..9c083996c 100644 --- a/server/src/common/ddtree.cpp +++ b/server/src/common/ddtree.cpp @@ -61,145 +61,298 @@ void extract_draft_topk(const float * logits, } } -DDTree build_ddtree(const float * top_log_probs, - const int32_t * top_token_ids, - int L, int K, int budget, - bool chain_seed) { +namespace { + +// Fill the ancestor-only visibility mask: slot v can see every slot on the +// root-to-v path (including itself), nothing else. visibility is row-major +// (1 + n_nodes)^2, iteration order exploits parents[i] < i (DFS/level order). +void build_visibility(DDTree & tree) { + const int N = 1 + tree.n_nodes; + tree.visibility.assign((size_t)N * N, 0); + tree.visibility[0] = 1; // root sees itself + for (int i = 1; i < N; i++) { + const int p = tree.parents[(size_t)i]; + for (int j = 0; j < i; j++) { + tree.visibility[(size_t)i * N + j] = + tree.visibility[(size_t)p * N + j]; + } + tree.visibility[(size_t)i * N + i] = 1; + } +} + +// Shared best-first DDTree construction. +// +// `topk(prefix, depth, log_probs, token_ids)` returns the depth-th (1-based) +// position's sorted top-K distribution conditioned on `prefix` (the token ids +// chosen so far). It fills both out vectors with exactly K entries on success +// and returns false on failure (e.g. an exact prefix that cannot be scored). +// `depth` is 1-based and `prefix` contains the previously chosen token ids. +// +// Candidates are kept in a max-heap keyed by cumulative path log-probability +// q(v); popping in descending q makes the SpecLA confidence window +// (keep q(v) >= q* - tau_tree) a single early-stop comparison. +// +// `eager_siblings` selects between the two historical expansion schedules: +// the conditional builder pushes every sibling/child rank eagerly, while the +// precomputed builder pushes only the next sibling and rank-0 child lazily. +// The schedules produce the same candidate set, but `std::priority_queue` +// tie-breaks differently on equal cumulative scores, so each public builder +// keeps its original schedule. +DDTree build_ddtree_impl(const DDTreeConditionalTopK & topk, + int L, int K, int budget, + bool chain_seed, + bool eager_siblings, + float tau_tree, + const float * precomputed_log_probs, + const int32_t * precomputed_token_ids) { DDTree tree; - if (budget <= 0 || L <= 0) { - tree.parents.push_back(-1); - tree.child_maps.emplace_back(); + tree.parents.push_back(-1); + tree.child_maps.emplace_back(); + + // A negative tau would reject even q* itself and leave a degenerate + // proposal (the conditional Qwen path treats that as a failure). Clamp it + // to zero: q* still passes, and positive margins prune normally. + if (tau_tree < 0.0f) tau_tree = 0.0f; + + if (budget <= 0 || L <= 0 || K <= 0) { + tree.visibility.assign(1, 1); + return tree; + } + + // Precomputed rows are prefix-independent, so the lazy builder can index + // them directly without allocating/copying two K-element vectors for + // every sibling and child expansion. + const bool precomputed = !eager_siblings && precomputed_log_probs && + precomputed_token_ids; + + // Fetch the depth-1 distribution. Its rank-0 score is q*, the best + // candidate score in the whole tree. + std::vector lp; + std::vector ids; + if (!precomputed && + (!topk({}, 1, lp, ids) || + (int)lp.size() < K || (int)ids.size() < K)) { tree.visibility.assign(1, 1); return tree; } + const float q_star = precomputed ? precomputed_log_probs[0] : lp[0]; - struct HeapEntry { - float neg_logw; - std::vector ranks; - int parent_index; + struct Candidate { + float logw; + int parent; int depth; int rank; - float logw; + int32_t token; + std::vector path; // token ids root->this node }; - struct HeapCmp { - bool operator()(const HeapEntry & a, const HeapEntry & b) const { - return a.neg_logw > b.neg_logw; + struct Worse { + bool operator()(const Candidate & a, const Candidate & b) const { + return a.logw < b.logw; } }; - std::priority_queue, HeapCmp> heap; + std::priority_queue, Worse> heap; - tree.token_ids.reserve(budget); - tree.depths.reserve(budget); - tree.parents.reserve(budget + 1); - tree.parents.push_back(-1); - tree.child_maps.emplace_back(); + auto push_scored_candidate = [&](int parent, int depth, float logw, + const std::vector & prefix, + int rank, int32_t token) { + Candidate c; + c.logw = logw; + c.parent = parent; + c.depth = depth; + c.rank = rank; + c.token = token; + if (!precomputed) { + c.path = prefix; + c.path.push_back(c.token); + } + heap.push(std::move(c)); + }; + + // Push one rank of one depth's distribution. The candidate's path is + // `prefix` + the candidate's own token. + auto push_candidate = [&](int parent, int depth, float parent_logw, + const std::vector & prefix, + int rank, + const std::vector & probs, + const std::vector & toks) { + push_scored_candidate(parent, depth, + parent_logw + probs[(size_t)rank], + prefix, rank, toks[(size_t)rank]); + }; + + // Push ranks [first_rank, K) of one depth's distribution as siblings of + // `parent`. + auto push_children = [&](int parent, int depth, float parent_logw, + const std::vector & prefix, + const std::vector & probs, + const std::vector & toks, + int first_rank) { + for (int rank = first_rank; rank < K; ++rank) { + push_candidate(parent, depth, parent_logw, prefix, rank, + probs, toks); + } + }; if (chain_seed) { + // Defensively pre-seed the top-1 chain (up to the budget), preserving + // ancestor closure: stop at the first depth whose cumulative top-1 + // score leaves the confidence window (every deeper descendant would + // be out of window too). + std::vector prefix; + float cumulative = 0.0f; + int parent = 0; + std::vector cur_lp = lp; + std::vector cur_ids = ids; const int chain_depth = std::min(L, budget); - float cum_logw = 0.0f; - int prev_idx = 0; - for (int d = 1; d <= chain_depth; d++) { - const int32_t tok_id = top_token_ids[(size_t)(d - 1) * K + 0]; - cum_logw += top_log_probs[(size_t)(d - 1) * K + 0]; - - const int cur_idx = tree.n_nodes + 1; - tree.token_ids.push_back(tok_id); - tree.depths.push_back(d); - tree.parents.push_back(prev_idx); + for (int depth = 1; depth <= chain_depth; ++depth) { + const int row = (depth - 1)*K; + const float rank0_logp = precomputed + ? precomputed_log_probs[row] : cur_lp[0]; + const int32_t rank0_token = precomputed + ? precomputed_token_ids[row] : cur_ids[0]; + const float next_logw = cumulative + rank0_logp; + if (q_star - next_logw > tau_tree) break; + + // Sibling candidates below the top-1 at this depth. + if (eager_siblings) { + push_children(parent, depth, cumulative, prefix, + cur_lp, cur_ids, 1); + } else if (K > 1) { + if (precomputed) { + push_scored_candidate( + parent, depth, + cumulative + precomputed_log_probs[row + 1], + prefix, 1, precomputed_token_ids[row + 1]); + } else { + push_candidate(parent, depth, cumulative, prefix, 1, + cur_lp, cur_ids); + } + } + + const int node = tree.n_nodes + 1; + const int32_t token = rank0_token; + tree.token_ids.push_back(token); + tree.depths.push_back(depth); + tree.parents.push_back(parent); tree.child_maps.emplace_back(); - tree.child_maps[prev_idx][tok_id] = cur_idx; + tree.child_maps[(size_t)parent][token] = node; tree.n_nodes++; - if (K > 1) { - const float sibling_logw = cum_logw - - top_log_probs[(size_t)(d - 1) * K + 0] - + top_log_probs[(size_t)(d - 1) * K + 1]; - heap.push({ - -sibling_logw, - {1}, - prev_idx, - d, - 1, - sibling_logw, - }); + if (!precomputed) prefix.push_back(token); + cumulative = next_logw; + parent = node; + + if (depth == L) break; + if (!precomputed && + (!topk(prefix, depth + 1, cur_lp, cur_ids) || + (int)cur_lp.size() < K || (int)cur_ids.size() < K)) { + break; } - prev_idx = cur_idx; } } else { - const float root_logw = top_log_probs[0 * K + 0]; - heap.push({ - -root_logw, - {0}, - 0, - 1, - 0, - root_logw, - }); + // Pure best-first: every depth-1 candidate is a root child. The lazy + // schedule starts with rank 0 alone and discovers its siblings as it + // pops, matching the precomputed builder's original heap layout. + if (eager_siblings) { + push_children(0, 1, 0.0f, {}, lp, ids, 0); + } else if (precomputed) { + push_scored_candidate(0, 1, precomputed_log_probs[0], {}, 0, + precomputed_token_ids[0]); + } else { + push_candidate(0, 1, 0.0f, {}, 0, lp, ids); + } } while (!heap.empty() && tree.n_nodes < budget) { - HeapEntry top = heap.top(); + Candidate c = heap.top(); heap.pop(); + // Best-first pops in descending q(v): the first out-of-window + // candidate proves every remaining one is out of the window too. + if (q_star - c.logw > tau_tree) break; - const int depth_minus_1 = top.depth - 1; - const int rank = top.rank; - const int32_t token_id = top_token_ids[(size_t)depth_minus_1 * K + rank]; - - const int current_index = tree.n_nodes + 1; - tree.token_ids.push_back(token_id); - tree.depths.push_back(top.depth); - tree.parents.push_back(top.parent_index); + const int node = tree.n_nodes + 1; + tree.token_ids.push_back(c.token); + tree.depths.push_back(c.depth); + tree.parents.push_back(c.parent); tree.child_maps.emplace_back(); - tree.child_maps[top.parent_index][token_id] = current_index; + tree.child_maps[(size_t)c.parent][c.token] = node; tree.n_nodes++; - if (rank + 1 < K) { - const float sibling_logw = top.logw - - top_log_probs[(size_t)depth_minus_1 * K + rank] - + top_log_probs[(size_t)depth_minus_1 * K + rank + 1]; - std::vector sibling_ranks = top.ranks; - sibling_ranks.back() = rank + 1; - heap.push({ - -sibling_logw, - std::move(sibling_ranks), - top.parent_index, - top.depth, - rank + 1, - sibling_logw, - }); - } - - if (top.depth < L) { - const float child_logw = top.logw - + top_log_probs[(size_t)top.depth * K + 0]; - std::vector child_ranks = top.ranks; - child_ranks.push_back(0); - heap.push({ - -child_logw, - std::move(child_ranks), - current_index, - top.depth + 1, - 0, - child_logw, - }); - } - } - - // Build ancestor-only visibility mask. - const int N = 1 + tree.n_nodes; - tree.visibility.assign((size_t)N * N, 0); - tree.visibility[0 * N + 0] = 1; - for (int i = 1; i < N; i++) { - const int p = tree.parents[i]; - for (int j = 0; j < i; j++) { - tree.visibility[(size_t)i * N + j] = tree.visibility[(size_t)p * N + j]; + if (eager_siblings) { + if (c.depth < L && + topk(c.path, c.depth + 1, lp, ids) && + (int)lp.size() >= K && (int)ids.size() >= K) { + push_children(node, c.depth + 1, c.logw, c.path, lp, ids, 0); + } + } else { + // Lazy schedule: expose the next sibling and the rank-0 child one + // at a time (the original precomputed-array expansion order). + if (precomputed) { + const int row = (c.depth - 1)*K; + if (c.rank + 1 < K) { + const float parent_logw = + c.logw - precomputed_log_probs[row + c.rank]; + push_scored_candidate( + c.parent, c.depth, + parent_logw + precomputed_log_probs[row + c.rank + 1], + {}, c.rank + 1, + precomputed_token_ids[row + c.rank + 1]); + } + if (c.depth < L) { + const int child_row = c.depth*K; + push_scored_candidate( + node, c.depth + 1, + c.logw + precomputed_log_probs[child_row], + {}, 0, precomputed_token_ids[child_row]); + } + } else if (c.rank + 1 < K) { + std::vector parent_lp; + std::vector parent_ids; + std::vector parent_prefix = c.path; + parent_prefix.pop_back(); + if (topk(parent_prefix, c.depth, parent_lp, parent_ids) && + (int)parent_lp.size() >= K && + (int)parent_ids.size() >= K) { + push_candidate(c.parent, c.depth, + c.logw - parent_lp[(size_t)c.rank], + parent_prefix, c.rank + 1, + parent_lp, parent_ids); + } + } + if (!precomputed && c.depth < L && + topk(c.path, c.depth + 1, lp, ids) && + (int)lp.size() >= K && (int)ids.size() >= K) { + push_candidate(node, c.depth + 1, c.logw, c.path, 0, lp, ids); + } } - tree.visibility[(size_t)i * N + i] = 1; } + build_visibility(tree); return tree; } +} // namespace + +DDTree build_ddtree(const float * top_log_probs, + const int32_t * top_token_ids, + int L, int K, int budget, + bool chain_seed, + float tau_tree) { + return build_ddtree_impl(DDTreeConditionalTopK{}, L, K, budget, + chain_seed, /*eager_siblings=*/false, tau_tree, + top_log_probs, top_token_ids); +} + +DDTree build_ddtree_conditional(const DDTreeConditionalTopK & next_topk, + int L, int K, int budget, + bool chain_seed, + float tau_tree) { + return build_ddtree_impl(next_topk, L, K, budget, chain_seed, + /*eager_siblings=*/true, tau_tree, + /*precomputed_log_probs=*/nullptr, + /*precomputed_token_ids=*/nullptr); +} + std::vector follow_verified_tree(const DDTree & tree, const int32_t * posterior, int & out_next_token, diff --git a/server/src/common/ddtree.h b/server/src/common/ddtree.h index afe22f226..026dc4ca4 100644 --- a/server/src/common/ddtree.h +++ b/server/src/common/ddtree.h @@ -9,6 +9,8 @@ #pragma once #include +#include +#include #include #include #include @@ -48,10 +50,35 @@ void extract_draft_topk(const float * logits, // K: top-K per position // budget: maximum number of non-root tree nodes // chain_seed: pre-seed full top-1 chain (defensive) vs pure best-first +// tau_tree: SpecLA confidence margin (arXiv:2607.16673 §6.1): only +// candidates with cumulative path log-probability +// q(v) >= q* - tau_tree are expanded, where q* is the best +// candidate score. Because best-first pops in descending q, +// one comparison prunes width and depth jointly, and the +// retained set is ancestor-closed by construction. The node +// budget still applies on top. Non-finite (default) = off. +// The chain seed is pruned at the first out-of-window depth, +// preserving ancestor closure. DDTree build_ddtree(const float * top_log_probs, const int32_t * top_token_ids, int L, int K, int budget, - bool chain_seed = true); + bool chain_seed = true, + float tau_tree = std::numeric_limits::infinity()); + +// Branch-conditioned variant used by SpecLA. The callback receives the token +// prefix of a tree node and the 1-based depth to predict next, and returns K +// sorted log-probabilities/token ids for that exact prefix. +using DDTreeConditionalTopK = std::function & prefix, + int next_depth, + std::vector & log_probs, + std::vector & token_ids)>; + +DDTree build_ddtree_conditional(const DDTreeConditionalTopK & next_topk, + int L, int K, int budget, + bool chain_seed = true, + float tau_tree = + std::numeric_limits::infinity()); // Walk the verified tree following the target's argmax (posterior) at each // node. Returns the list of flat-tree indices that make up the accepted path diff --git a/server/src/common/dflash_spec_decode.cpp b/server/src/common/dflash_spec_decode.cpp index f43d71d97..9bf903359 100644 --- a/server/src/common/dflash_spec_decode.cpp +++ b/server/src/common/dflash_spec_decode.cpp @@ -86,7 +86,8 @@ bool run_dflash_spec_decode( int n_accept_sum = 0; int n_hint_proposed = 0; int n_hint_accepted = 0; - const ChainRollbackPolicy rollback_policy = resolve_chain_rollback_policy(); + const ChainRollbackPolicy rollback_policy = + resolve_chain_rollback_policy(false, target.exact_fast_rollback()); RollbackDiag rollback_diag; auto t_dec0 = std::chrono::steady_clock::now(); @@ -243,9 +244,13 @@ bool run_dflash_spec_decode( fast_rolled_back = true; rollback_diag.record_fast_rollback(accept_n); } else { - // Rollback failed (e.g. CUDA error / unsupported state type). + if (!target.rollback_failure_is_recoverable()) { + std::fprintf(stderr, "dflash-spec rollback_to failed after " + "an in-place commit attempt; aborting\n"); + return false; + } // The pre-verify snapshot is still valid, so degrade to the - // legacy restore+replay path below instead of aborting. + // legacy restore+replay path below. std::fprintf(stderr, "dflash-spec rollback_to failed; " "falling back to restore+replay\n"); rollback_diag.record_failed_fallback(); @@ -290,6 +295,10 @@ bool run_dflash_spec_decode( if (io.is_cancelled()) break; if (hit_eos) break; } + if (!target.finish_speculative_state()) { + std::fprintf(stderr, "dflash-spec final recurrent-state flush failed\n"); + return false; + } if (!use_remote_draft && draft_backend) ggml_backend_synchronize(draft_backend); auto t_dec1 = std::chrono::steady_clock::now(); const double decode_s = std::chrono::duration(t_dec1 - t_dec0).count(); diff --git a/server/src/common/dflash_target.h b/server/src/common/dflash_target.h index af7cd756f..e369bfd8f 100644 --- a/server/src/common/dflash_target.h +++ b/server/src/common/dflash_target.h @@ -62,6 +62,14 @@ struct DFlashTarget { // When true, verify_batch captures intermediates and rollback_to() works. virtual bool supports_fast_rollback() const { return false; } + // Whether fast rollback is an exact, low-overhead device-side commit for + // which the replay breakeven threshold does not apply. + virtual bool exact_fast_rollback() const { return false; } + + // Whether restore+replay remains safe after rollback_to() returns false. + // In-place commit implementations override this while active. + virtual bool rollback_failure_is_recoverable() const { return true; } + // Roll back recurrent state to position `commit_n` within the last // verify batch (0-indexed). Uses SSM intermediate states captured during // verify. Also truncates KV to `base_pos + commit_n`. No replay needed. @@ -70,6 +78,11 @@ struct DFlashTarget { (void)base_pos; (void)commit_n; return false; } + // Flush any accepted recurrent factors intentionally left pending for + // fusion with the next verify. Called once before a successful decode + // returns (EOS, cancellation, or token budget). + virtual bool finish_speculative_state() { return true; } + // ── DDTree tree-structured verify ─────────────────────────────── // Whether this target can verify a draft tree (ancestor-masked batched // forward over DFS-ordered tree nodes). When false, callers fall back to diff --git a/server/src/common/kvflash_pager.h b/server/src/common/kvflash_pager.h index c9af59a38..d9113dbca 100644 --- a/server/src/common/kvflash_pager.h +++ b/server/src/common/kvflash_pager.h @@ -599,13 +599,19 @@ struct KvFlashAutoBudget { int speed_cap_tokens = 16384; }; +// Whether the operator requested any KVFlash pool. "auto" is a real request +// even though its final size is not known until the backend measures VRAM. +inline bool kvflash_pool_requested(const char * value) { + return value != nullptr && + (std::strcmp(value, "auto") == 0 || std::atoi(value) > 0); +} + // The compatibility gate can reject a fixed KVFlash pool before model setup. // "auto" is deliberately excluded: only the backend's VRAM-aware sizing can // determine whether an automatic pool will actually be active. inline bool kvflash_fixed_pool_requested(const char * value) { - return value != nullptr && - std::strcmp(value, "auto") != 0 && - std::atoi(value) > 0; + return kvflash_pool_requested(value) && + std::strcmp(value, "auto") != 0; } // Pool size from DFLASH_KVFLASH for a backend with `cfg` protections: diff --git a/server/src/common/specla_commit_cuda.cu b/server/src/common/specla_commit_cuda.cu new file mode 100644 index 000000000..9dfb5b3de --- /dev/null +++ b/server/src/common/specla_commit_cuda.cu @@ -0,0 +1,293 @@ +// Fused SpecLA DeltaConstruct commit — see specla_commit_cuda.h. + +#include "specla_commit_cuda.h" + +#include +#include + +namespace dflash::common { + +namespace { + +constexpr int kBlock = 256; +constexpr int kMaxAccept = 64; // >= every verify window / ddtree budget in use + +// Grid: x over the S_k*S_v state elements of one (head, layer) plane, +// y over head-layer planes (hl = h + H*l). +// Factor element (·, h, l, t) lives at block index fo = h + H*(l + L*t) — +// identical for fk (stride S_k), fv (stride S_v), and fg (stride 1). +__global__ void specla_commit_kernel(float * const * ssm_ptrs, + const float * __restrict__ fk, + const float * __restrict__ fv, + const float * __restrict__ fg, + const int32_t * __restrict__ idx, + int A, int S_k, int S_v, int H, int L) { + const int hl = blockIdx.y; + const int l = hl / H; + const int h = hl % H; + + // Per-plane scalars shared by every thread: the accepted-end gate and the + // per-token decay weights along the accepted path. + __shared__ float s_w[kMaxAccept]; + __shared__ float s_gA_exp; + __shared__ int s_tok[kMaxAccept]; + if (threadIdx.x < A) { + s_tok[threadIdx.x] = idx[threadIdx.x]; + } + __syncthreads(); + if (threadIdx.x == 0) { + const int tA = s_tok[A - 1]; + const float gA = fg[(size_t)h + (size_t)H * (l + (size_t)L * tA)]; + s_gA_exp = expf(gA); + for (int t = 0; t < A; t++) { + const size_t fo = (size_t)h + (size_t)H * (l + (size_t)L * s_tok[t]); + s_w[t] = expf(gA - fg[fo]); + } + } + __syncthreads(); + + const int e = blockIdx.x * blockDim.x + threadIdx.x; + if (e >= S_k * S_v) return; + const int i = e % S_k; // state row (k-dim, ne0) + const int c = e / S_k; // state column + + float * s_plane = ssm_ptrs[l] + (size_t)h * S_k * S_v; + float acc = s_gA_exp * s_plane[(size_t)c * S_k + i]; + for (int t = 0; t < A; t++) { + const size_t fo = (size_t)h + (size_t)H * (l + (size_t)L * s_tok[t]); + acc += s_w[t] * fk[(size_t)fo * S_k + i] * fv[(size_t)fo * S_v + c]; + } + s_plane[(size_t)c * S_k + i] = acc; +} + +__global__ void specla_compact_kernel( + const float * __restrict__ src_k, + const float * __restrict__ src_v, + const float * __restrict__ src_g, + const float * __restrict__ src_conv, + float * __restrict__ dst_k, + float * __restrict__ dst_v, + float * __restrict__ dst_g, + float * __restrict__ dst_conv, + const int32_t * __restrict__ idx, + int A, int S_k, int S_v, int H, int L, int C) { + const size_t e = (size_t)blockIdx.x*blockDim.x + threadIdx.x; + const size_t k_n = (size_t)A*L*H*S_k; + const size_t v_n = (size_t)A*L*H*S_v; + const size_t g_n = (size_t)A*L*H; + const size_t c_n = (size_t)A*L*C; + if (e < k_n) { + const int t = e/(L*H*S_k); + const size_t inner = e%(L*H*S_k); + dst_k[e] = src_k[(size_t)idx[t]*L*H*S_k + inner]; + } + if (e < v_n) { + const int t = e/(L*H*S_v); + const size_t inner = e%(L*H*S_v); + dst_v[e] = src_v[(size_t)idx[t]*L*H*S_v + inner]; + } + if (e < g_n) { + const int t = e/(L*H); + const size_t inner = e%(L*H); + dst_g[e] = src_g[(size_t)idx[t]*L*H + inner]; + } + if (e < c_n) { + const int t = e/(L*C); + const size_t inner = e%(L*C); + dst_conv[e] = src_conv[(size_t)idx[t]*L*C + inner]; + } +} + +__global__ void specla_flush_state_raw_kernel( + float * const * ssm_ptrs, + const float * __restrict__ fk, + const float * __restrict__ fv, + const float * __restrict__ fg, + int A, int S_k, int S_v, int H, int L) { + const int hl = blockIdx.y; + const int l = hl/H; + const int h = hl%H; + const int e = blockIdx.x*blockDim.x + threadIdx.x; + if (e >= S_k*S_v) return; + const int row = e%S_k; + const int col = e/S_k; + float * plane = ssm_ptrs[l] + (size_t)h*S_k*S_v; + float acc = plane[(size_t)col*S_k + row]; + for (int t = 0; t < A; ++t) { + const size_t hl_t = (size_t)h + (size_t)H*(l + (size_t)L*t); + acc = fmaf(fk[hl_t*S_k + row], fv[hl_t*S_v + col], + expf(fg[hl_t])*acc); + } + plane[(size_t)col*S_k + row] = acc; +} + +__global__ void specla_flush_conv_raw_kernel( + float * const * conv_ptrs, + const float * __restrict__ conv, + int A, int L, int C, int d_conv) { + const int e = blockIdx.x*blockDim.x + threadIdx.x; + if (e >= L*C) return; + const int l = e/C; + const int channel = e%C; + float * state = conv_ptrs[l] + (size_t)channel*(d_conv - 1); + for (int t = 0; t < A; ++t) { + for (int j = 0; j < d_conv - 2; ++j) state[j] = state[j + 1]; + state[d_conv - 2] = + conv[(size_t)channel + (size_t)C*(l + (size_t)L*t)]; + } +} + +} // namespace + +bool specla_commit_fused(float * const * ssm_ptrs_dev, + const float * fk, + const float * fv, + const float * fg, + const int32_t * idx_dev, + int A, int S_k, int S_v, int H, int n_delta, + void * stream, + bool * launched) { + if (launched) *launched = false; + if (!ssm_ptrs_dev || !fk || !fv || !fg || !idx_dev) return false; + if (A <= 0 || A > kMaxAccept || S_k <= 0 || S_v <= 0 || H <= 0 || n_delta <= 0) { + return false; + } + const int planes = H * n_delta; + if (planes > 65535) return false; // grid.y limit + + dim3 block(kBlock); + dim3 grid(((unsigned)(S_k * S_v) + kBlock - 1) / kBlock, (unsigned)planes); + (void)cudaGetLastError(); // discard any unrelated prior launch status + specla_commit_kernel<<>>( + ssm_ptrs_dev, fk, fv, fg, idx_dev, A, S_k, S_v, H, n_delta); + if (cudaGetLastError() != cudaSuccess) return false; + if (launched) *launched = true; + return cudaStreamSynchronize((cudaStream_t)stream) == cudaSuccess; +} + +bool specla_compact_fused( + const float * src_k, const float * src_v, const float * src_g, + const float * src_conv, float * dst_k, float * dst_v, float * dst_g, + float * dst_conv, const int32_t * idx_dev, int A, int S_k, int S_v, + int H, int n_delta, int conv_channels, void * stream) { + if (!src_k || !src_v || !src_g || !src_conv || !dst_k || !dst_v || + !dst_g || !dst_conv || !idx_dev || A <= 0 || A > kMaxAccept) { + return false; + } + const size_t n = std::max( + std::max((size_t)A*n_delta*H*S_k, (size_t)A*n_delta*H*S_v), + std::max((size_t)A*n_delta*H, (size_t)A*n_delta*conv_channels)); + const dim3 block(kBlock); + const dim3 grid((unsigned)((n + kBlock - 1)/kBlock)); + (void)cudaGetLastError(); + specla_compact_kernel<<>>( + src_k, src_v, src_g, src_conv, dst_k, dst_v, dst_g, dst_conv, + idx_dev, A, S_k, S_v, H, n_delta, conv_channels); + if (cudaGetLastError() != cudaSuccess) return false; + return cudaStreamSynchronize((cudaStream_t)stream) == cudaSuccess; +} + +bool specla_flush_raw_fused( + float * const * ssm_ptrs_dev, float * const * conv_ptrs_dev, + const float * fk, const float * fv, const float * fg, + const float * conv, int A, int S_k, int S_v, int H, int n_delta, + int conv_channels, int d_conv, void * stream) { + if (!ssm_ptrs_dev || !conv_ptrs_dev || !fk || !fv || !fg || !conv || + A <= 0 || A > kMaxAccept || d_conv < 2) return false; + const dim3 block(kBlock); + const dim3 state_grid( + ((unsigned)(S_k*S_v) + kBlock - 1)/kBlock, + (unsigned)(H*n_delta)); + (void)cudaGetLastError(); + specla_flush_state_raw_kernel<<>>( + ssm_ptrs_dev, fk, fv, fg, A, S_k, S_v, H, n_delta); + const dim3 conv_grid( + ((unsigned)(n_delta*conv_channels) + kBlock - 1)/kBlock); + specla_flush_conv_raw_kernel<<>>( + conv_ptrs_dev, conv, A, n_delta, conv_channels, d_conv); + if (cudaGetLastError() != cudaSuccess) return false; + return cudaStreamSynchronize((cudaStream_t)stream) == cudaSuccess; +} + +bool specla_commit_conv_raw_fused(float * const * conv_ptrs_dev, + const float * conv, + int A, int n_tokens, + int n_delta, int conv_channels, + int d_conv, void * stream) { + if (!conv_ptrs_dev || !conv || A <= 0 || A > n_tokens || + n_tokens <= 0 || + n_delta <= 0 || conv_channels <= 0 || d_conv < 2) { + return false; + } + const dim3 block(kBlock); + const dim3 grid( + ((unsigned)(n_delta*conv_channels) + kBlock - 1)/kBlock); + (void)cudaGetLastError(); + specla_flush_conv_raw_kernel<<>>( + conv_ptrs_dev, conv, A, n_delta, conv_channels, d_conv); + if (cudaGetLastError() != cudaSuccess) return false; + return cudaStreamSynchronize((cudaStream_t)stream) == cudaSuccess; +} + +bool specla_rotate_pending_factors(const SpeclaFactorBanks & banks, + const int32_t * idx_dev, + int pending_bank, + bool walked_sibling, + int commit_n, + int S_k, int S_v, int H, + int n_delta, int conv_channels, + void * stream, + int * out_pending_bank) { + if (!out_pending_bank || pending_bank < 0 || pending_bank > 1 || + commit_n <= 0) { + return false; + } + for (int b = 0; b < 2; ++b) { + if (!banks.k[b] || !banks.v[b] || !banks.g[b] || !banks.conv[b]) { + return false; + } + } + + const int current_bank = 1 - pending_bank; + if (!walked_sibling) { + // A chain acceptance is already in path order: only the host-side + // bank role changes. + *out_pending_bank = current_bank; + return true; + } + + // A sibling walk scatters the accepted path across the current bank. + // Compact it into the old pending bank, which then becomes pending. + if (!idx_dev) return false; + const bool ok = specla_compact_fused( + banks.k[current_bank], banks.v[current_bank], + banks.g[current_bank], banks.conv[current_bank], + banks.k[pending_bank], banks.v[pending_bank], + banks.g[pending_bank], banks.conv[pending_bank], + idx_dev, commit_n, S_k, S_v, H, n_delta, conv_channels, stream); + if (ok) *out_pending_bank = pending_bank; + return ok; +} + +bool specla_flush_pending_factors(const SpeclaFactorBanks & banks, + float * const * ssm_ptrs_dev, + float * const * conv_ptrs_dev, + int pending_bank, + int pending_count, + int S_k, int S_v, int H, + int n_delta, int conv_channels, int d_conv, + void * stream) { + if (pending_count <= 0) return true; + if (pending_bank < 0 || pending_bank > 1 || !ssm_ptrs_dev || + !conv_ptrs_dev || !banks.k[pending_bank] || !banks.v[pending_bank] || + !banks.g[pending_bank] || !banks.conv[pending_bank]) { + return false; + } + return specla_flush_raw_fused( + ssm_ptrs_dev, conv_ptrs_dev, + banks.k[pending_bank], banks.v[pending_bank], + banks.g[pending_bank], banks.conv[pending_bank], + pending_count, S_k, S_v, H, n_delta, conv_channels, d_conv, stream); +} + +} // namespace dflash::common diff --git a/server/src/common/specla_commit_cuda.h b/server/src/common/specla_commit_cuda.h new file mode 100644 index 000000000..e0a300efc --- /dev/null +++ b/server/src/common/specla_commit_cuda.h @@ -0,0 +1,121 @@ +// SpecLA factor lifecycle helpers (docs/SPECLA.md). +// +// The first entry is the immediate DeltaConstruct compatibility path used by +// the fully factorized verifier: +// +// S_l ← exp(g⁺_A) · S_l + Σ_{t ∈ path} exp(g⁺_A − g⁺_t) · k_t ⊗ ṽ_t +// +// The ggml-graph implementation of the same math needs ~6 small ops per layer +// (48 layers ⇒ hundreds of kernel launches ⇒ ~5 ms of launch overhead per +// commit on ROCm without HIP graphs). The production HLD path uses the compact +// and final-flush helpers below; its normal accepted update is delayed into +// the next verification kernel. +// +// Compiled for CUDA and (via the hip_compat shim with +// LANGUAGE HIP) for ROCm — same shared-.cu pattern as +// geometric_draft_topk_cuda.cu. + +#pragma once + +#include +#include + +namespace dflash::common { + +// Consolidated double-bank factor pointers. Bank 0 and bank 1 alternate +// between "pending" (consumed by the next verify / final flush) and "current" +// (receiving the just-run verify's factors). +struct SpeclaFactorBanks { + float * k[2] = {nullptr, nullptr}; + float * v[2] = {nullptr, nullptr}; + float * g[2] = {nullptr, nullptr}; + float * conv[2] = {nullptr, nullptr}; +}; + +// ssm_ptrs_dev: DEVICE array of n_delta pointers, one per delta layer's +// [S_k, S_v, H] f32 state tensor (ne0 = k-dim rows). +// fk/fv/fg: consolidated factor buffers, f32, token axis outermost: +// fk [S_k, H, n_delta, T], fv [S_v, H, n_delta, T], +// fg [H, n_delta, T]. +// idx_dev: DEVICE array of A accepted token indices, deepest last. +// stream: cudaStream_t / hipStream_t (nullptr = default stream). +// launched: set once the kernel has been accepted by the runtime. A false +// return with launched=true is not safe to retry in place. +bool specla_commit_fused(float * const * ssm_ptrs_dev, + const float * fk, + const float * fv, + const float * fg, + const int32_t * idx_dev, + int A, int S_k, int S_v, int H, int n_delta, + void * stream, + bool * launched = nullptr); + +// Compact an accepted tree path from arbitrary flat-node indices into +// contiguous token slots in the alternate bank. All delta layers and both +// recurrent factor families are copied by one kernel launch. +bool specla_compact_fused(const float * src_k, + const float * src_v, + const float * src_g, + const float * src_conv, + float * dst_k, + float * dst_v, + float * dst_g, + float * dst_conv, + const int32_t * idx_dev, + int A, int S_k, int S_v, int H, + int n_delta, int conv_channels, + void * stream); + +// Materialize a final delayed path when generation ends before another verify +// can consume it. Factors are raw per-token recurrence terms in path order. +bool specla_flush_raw_fused(float * const * ssm_ptrs_dev, + float * const * conv_ptrs_dev, + const float * fk, + const float * fv, + const float * fg, + const float * conv, + int A, int S_k, int S_v, int H, + int n_delta, int conv_channels, int d_conv, + void * stream); + +// Apply an accepted prefix from a raw token-major convolution factor bank to +// every durable per-layer convolution window. `conv` is [C, L, T]; each +// state pointer addresses a contiguous [d_conv-1, C] tensor. +bool specla_commit_conv_raw_fused(float * const * conv_ptrs_dev, + const float * conv, + int A, int n_tokens, + int n_delta, int conv_channels, + int d_conv, void * stream); + +// Shared bank-lifecycle helpers used by both the production Qwen35 target and +// the test_dflash harness. `pending_bank` is the bank the just-run verify +// consumed; the opposite bank received that verify's factors. +// +// Rotate the just-verified factors into the pending role. A pure chain +// acceptance is already contiguous and only switches banks. A tree acceptance +// that walked a sibling is compacted from the current bank into the old +// pending bank in accepted-path order via specla_compact_fused; `idx_dev` +// must already hold the `commit_n` accepted DFS indices. On success the new +// pending bank is written to `out_pending_bank`. +bool specla_rotate_pending_factors(const SpeclaFactorBanks & banks, + const int32_t * idx_dev, + int pending_bank, + bool walked_sibling, + int commit_n, + int S_k, int S_v, int H, + int n_delta, int conv_channels, + void * stream, + int * out_pending_bank); + +// Apply the pending bank's raw factors to the durable SSM/conv states. Used at +// generation boundaries where no next verify will consume the pending path. +bool specla_flush_pending_factors(const SpeclaFactorBanks & banks, + float * const * ssm_ptrs_dev, + float * const * conv_ptrs_dev, + int pending_bank, + int pending_count, + int S_k, int S_v, int H, + int n_delta, int conv_channels, int d_conv, + void * stream); + +} // namespace dflash::common diff --git a/server/src/common/specla_mode.h b/server/src/common/specla_mode.h new file mode 100644 index 000000000..8448ec6c0 --- /dev/null +++ b/server/src/common/specla_mode.h @@ -0,0 +1,56 @@ +// SpecLA runtime mode (docs/SPECLA.md, arXiv:2607.16673). +// +// DFLASH_SPECLA=1 switches the single-target qwen35 speculative verifier to +// the paper's state-resident path. A heavy-light schedule runs adjacent tree +// nodes while the GDN and convolution state tiles are live, records raw +// per-node factors, and applies the previously accepted factor buffer at the +// beginning of the next verification. Prefill and ordinary AR decode retain +// their normal state writebacks. +#pragma once + +#include +#include +#include +#include + +namespace dflash::common { + +inline bool specla_enabled() { + static const bool on = []() { + const char * v = std::getenv("DFLASH_SPECLA"); + return v != nullptr && v[0] != '\0' && std::strcmp(v, "0") != 0; + }(); + return on; +} + +// Section 6.2 assumes a tiny, trained EAGLE-style draft layer that can roll a +// beam out prefix by prefix. Qwen's current DFlash draft is a substantially +// larger block-diffusion model; rerunning all five layers for every expanded +// node is useful for algorithm experiments but normally costs more than it +// saves. Keep the exact branch-conditioned builder available without making +// it the production default for this model. +inline bool specla_conditional_draft_enabled() { + static const bool on = []() { + const char * v = std::getenv("DFLASH_SPECLA_CONDITIONAL_DRAFT"); + return v != nullptr && v[0] != '\0' && std::strcmp(v, "0") != 0; + }(); + return on; +} + +inline int specla_tree_topk() { + static const int topk = []() { + const char * v = std::getenv("DFLASH_SPECLA_TOPK"); + if (!v || v[0] == '\0') return 4; // paper's end-to-end setting + char * end = nullptr; + errno = 0; + const long parsed = std::strtol(v, &end, 10); + if (errno == ERANGE || end == v || *end != '\0' || + parsed <= 0 || parsed > INT_MAX) { + return 4; // malformed/out-of-range → documented default + } + return static_cast(parsed); + }(); + return topk; +} + +} // namespace dflash::common diff --git a/server/src/common/step_graph.h b/server/src/common/step_graph.h index cdd4d0210..900a7305e 100644 --- a/server/src/common/step_graph.h +++ b/server/src/common/step_graph.h @@ -36,6 +36,12 @@ struct StepGraph { ggml_tensor * positions = nullptr; ggml_tensor * attn_mask = nullptr; // may be null ggml_tensor * parent_ids = nullptr; // DDTree tree-mode; null for chain mode + // SpecLA topology masks ([n_tokens, n_tokens] f32, host-filled; see + // delta_net_specla.h). Created only when DFLASH_SPECLA capture is active. + ggml_tensor * specla_m_strict = nullptr; + ggml_tensor * specla_m_incl = nullptr; + ggml_tensor * specla_m_eye = nullptr; + ggml_tensor * specla_hld = nullptr; ggml_tensor * target_hidden_cat = nullptr; // draft only ggml_tensor * positions_k = nullptr; // draft only ggml_tensor * pad_mask_full = nullptr; // draft only; padded-ctx mask @@ -92,6 +98,8 @@ inline void step_graph_free(StepGraph & sg) { sg.built_view = false; sg.hidden_input = nullptr; sg.parent_ids = nullptr; + sg.specla_m_strict = sg.specla_m_incl = sg.specla_m_eye = nullptr; + sg.specla_hld = nullptr; sg.kv_write_rows = nullptr; sg.active_slot_ids = nullptr; sg.state_slot_ids = nullptr; diff --git a/server/src/delta_net_specla.cpp b/server/src/delta_net_specla.cpp new file mode 100644 index 000000000..a106cb221 --- /dev/null +++ b/server/src/delta_net_specla.cpp @@ -0,0 +1,275 @@ +// SpecLA topology-masked delta-net verification — see delta_net_specla.h. +// +// Derived from build_delta_net_chunked (itself a port of llama.cpp's +// build_delta_net_chunking), specialized to a single speculative window: +// - one chunk, no padding, no cross-chunk state loop; +// - ggml_tri causal masks replaced by host-filled topology masks so tree +// drafts verify in the same factorized form (paper §4.2: +// A_tree = M_strict ⊙ (K K_bᵀ) ⊙ exp(g⁺_t − g⁺_u), T = (I + A)⁻¹); +// - the recurrent state is read-only; instead of a new_state the builder +// returns the per-token factors (ṽ, g⁺) for accepted-state +// reconstruction (paper §5.1). +// +// Orientation conventions below follow ggml_mul_mat(a, b): with +// a = [K, M, b2, b3] and b = [K, N, b2', b3'], the result is +// C[M, N] = Σ_K a[K, M]·b[K, N], batched over dims 2..3 with broadcast. + +#include "delta_net_specla.h" + +#include +#include +#include + +namespace dflash::common { + +SpecLAHLDSchedule make_specla_hld_schedule(const int32_t * parents, + int n, + int pending_count) { + SpecLAHLDSchedule out; + if (!parents || n <= 0 || parents[0] >= 0 || pending_count < 0) { + return out; + } + + std::vector> children((size_t)n); + for (int node = 1; node < n; ++node) { + const int parent = parents[node]; + if (parent < 0 || parent >= node) return {}; + children[(size_t)parent].push_back(node); + } + + // Pick the largest subtree as the heavy continuation. Ties retain the + // topological child order, making schedules deterministic for tests and + // CUDA graph reuse. + std::vector subtree((size_t)n, 1); + std::vector heavy((size_t)n, -1); + for (int node = n - 1; node >= 0; --node) { + int best_size = 0; + for (int child : children[(size_t)node]) { + subtree[(size_t)node] += subtree[(size_t)child]; + if (subtree[(size_t)child] > best_size) { + best_size = subtree[(size_t)child]; + heavy[(size_t)node] = child; + } + } + } + + struct Chain { + std::vector nodes; + int parent_node = -1; + int wave = 0; + }; + std::vector chains; + std::vector node_chain((size_t)n, -1); + for (int head = 0; head < n; ++head) { + const int parent = parents[head]; + if (head != 0 && heavy[(size_t)parent] == head) continue; + Chain chain; + chain.parent_node = parent; + for (int node = head; node >= 0; node = heavy[(size_t)node]) { + node_chain[(size_t)node] = (int)chains.size(); + chain.nodes.push_back(node); + } + chains.push_back(std::move(chain)); + } + for (size_t c = 0; c < chains.size(); ++c) { + const int parent = chains[c].parent_node; + chains[c].wave = parent < 0 ? 0 : chains[(size_t)node_chain[(size_t)parent]].wave + 1; + } + std::stable_sort(chains.begin(), chains.end(), + [](const Chain & a, const Chain & b) { return a.wave < b.wave; }); + + std::vector order; + std::vector chain_offsets; + std::vector chain_parent_boundary; + std::vector chain_wave; + std::vector node_boundary((size_t)n, -1); + order.reserve((size_t)n); + chain_offsets.reserve(chains.size() + 1); + chain_parent_boundary.reserve(chains.size()); + chain_wave.reserve(chains.size()); + chain_offsets.push_back(0); + + int n_boundaries = 0; + int n_waves = 0; + int max_parallel_chains = 0; + int previous_wave = -1; + int chains_in_wave = 0; + for (const Chain & chain : chains) { + int boundary = -1; + if (chain.parent_node >= 0) { + int32_t & slot = node_boundary[(size_t)chain.parent_node]; + if (slot < 0) slot = n_boundaries++; + boundary = slot; + } + chain_parent_boundary.push_back(boundary); + chain_wave.push_back(chain.wave); + n_waves = std::max(n_waves, chain.wave + 1); + if (chain.wave != previous_wave) { + previous_wave = chain.wave; + chains_in_wave = 0; + } + max_parallel_chains = std::max(max_parallel_chains, ++chains_in_wave); + for (int node : chain.nodes) order.push_back(node); + chain_offsets.push_back((int32_t)order.size()); + } + + // Packed metadata ABI shared with gated_delta_net.cu and ssm-conv.cu. + // Header: magic, N, C, W, B, pending_count, five section offsets, and + // the maximum number of chains in one dependency wave. + constexpr int kHeader = 12; + const int order_off = kHeader; + const int offsets_off = order_off + n; + const int parent_off = offsets_off + (int)chains.size() + 1; + const int boundary_off = parent_off + (int)chains.size(); + const int wave_off = boundary_off + n; + out.packed.assign((size_t)wave_off + chains.size(), 0); + out.packed[0] = 0x534c4148; // "SLAH" + out.packed[1] = n; + out.packed[2] = (int32_t)chains.size(); + out.packed[3] = n_waves; + out.packed[4] = n_boundaries; + out.packed[5] = pending_count; + out.packed[6] = order_off; + out.packed[7] = offsets_off; + out.packed[8] = parent_off; + out.packed[9] = boundary_off; + out.packed[10] = wave_off; + out.packed[11] = max_parallel_chains; + std::copy(order.begin(), order.end(), out.packed.begin() + order_off); + std::copy(chain_offsets.begin(), chain_offsets.end(), out.packed.begin() + offsets_off); + std::copy(chain_parent_boundary.begin(), chain_parent_boundary.end(), out.packed.begin() + parent_off); + std::copy(node_boundary.begin(), node_boundary.end(), out.packed.begin() + boundary_off); + std::copy(chain_wave.begin(), chain_wave.end(), out.packed.begin() + wave_off); + + out.n_nodes = n; + out.n_chains = (int)chains.size(); + out.n_waves = n_waves; + out.n_boundaries = n_boundaries; + out.max_parallel_chains = max_parallel_chains; + return out; +} + +DeltaNetSpecLAResult build_delta_net_specla( + ggml_context * ctx0, + ggml_tensor * q, + ggml_tensor * k, + ggml_tensor * v, + ggml_tensor * g, + ggml_tensor * b, + ggml_tensor * s, + ggml_tensor * m_strict, + ggml_tensor * m_incl, + ggml_tensor * m_eye) { + const int64_t S_k = q->ne[0]; + const int64_t H_v = q->ne[1]; + const int64_t n = q->ne[2]; + + const int64_t S_v = v->ne[0]; + + // Same layer family as the chunked GDA path: scalar gate per value head, + // q/k already repeated to H_v heads by the caller. + GGML_ASSERT(S_k == S_v); + GGML_ASSERT(k->ne[0] == S_k && k->ne[1] == H_v && k->ne[2] == n); + GGML_ASSERT(v->ne[1] == H_v && v->ne[2] == n); + GGML_ASSERT(g->ne[0] == 1 && g->ne[1] == H_v && g->ne[2] == n); + GGML_ASSERT(b->ne[0] == 1 && b->ne[1] == H_v && b->ne[2] == n); + GGML_ASSERT(s->ne[0] == S_v && s->ne[1] == S_v && s->ne[2] == H_v); + GGML_ASSERT(q->ne[3] == 1 && s->ne[3] == 1); // spec-decode verify is single-seq + GGML_ASSERT(m_strict->ne[0] == n && m_strict->ne[1] == n); + GGML_ASSERT(m_incl->ne[0] == n && m_incl->ne[1] == n); + GGML_ASSERT(m_eye->ne[0] == n && m_eye->ne[1] == n); + + const float scale = 1.0f / sqrtf((float)S_k); + q = ggml_scale(ctx0, q, scale); + + // [S, H, n, 1] → [S, n, H, 1] → [S, n, 1, H]; gates [1, H, n, 1] → [1, n, 1, H] + q = ggml_reshape_4d(ctx0, ggml_cont(ctx0, ggml_permute(ctx0, q, 0, 2, 1, 3)), S_k, n, 1, H_v); + k = ggml_reshape_4d(ctx0, ggml_cont(ctx0, ggml_permute(ctx0, k, 0, 2, 1, 3)), S_k, n, 1, H_v); + v = ggml_reshape_4d(ctx0, ggml_cont(ctx0, ggml_permute(ctx0, v, 0, 2, 1, 3)), S_v, n, 1, H_v); + g = ggml_reshape_4d(ctx0, ggml_cont(ctx0, ggml_permute(ctx0, g, 0, 2, 1, 3)), 1, n, 1, H_v); + b = ggml_reshape_4d(ctx0, ggml_cont(ctx0, ggml_permute(ctx0, b, 0, 2, 1, 3)), 1, n, 1, H_v); + + ggml_tensor * v_b = ggml_mul(ctx0, v, b); + ggml_tensor * k_b = ggml_mul(ctx0, k, b); + + // Path-cumulative gate g⁺: for chains this is the plain cumsum; in + // general g⁺_t = Σ_{u ∈ anc(t) ∪ t} g_u = m_inclᵀ · g. + ggml_tensor * g_col = ggml_cont(ctx0, ggml_transpose(ctx0, g)); // [n, 1, 1, H] + ggml_tensor * g_ps = ggml_mul_mat(ctx0, m_incl, g_col); // [n, 1, 1, H] + + // Pairwise decay D[u, t] = exp(g⁺_t − g⁺_u) on ancestor-or-self pairs. + // Masking the exponent (not the exponential) makes out-of-window entries + // exp(0)=1, mirroring the chunked path's tri-then-exp; they are zeroed + // by the m_strict/m_incl products on kb/kq below. + ggml_tensor * g_ps_row = ggml_reshape_4d(ctx0, g_ps, 1, n, 1, H_v); + g_ps_row = ggml_repeat_4d(ctx0, g_ps_row, n, n, 1, H_v); // [u, t, 1, H] = g⁺_t + ggml_tensor * decay = ggml_sub(ctx0, g_ps_row, g_ps); // g⁺_t − g⁺_u + decay = ggml_exp(ctx0, ggml_mul(ctx0, decay, m_incl)); + + // A[u, t] = M_strict[u, t] · (k_u · β_t k_t) · D[u, t] + ggml_tensor * kb = ggml_mul_mat(ctx0, k, k_b); // [u, t, 1, H] + kb = ggml_mul(ctx0, kb, decay); + ggml_tensor * attn = ggml_mul(ctx0, kb, m_strict); + + // kq[u, t] = M_incl[u, t] · (k_u · q_t) · D[u, t] — the output-side + // attention onto corrected values. + ggml_tensor * kq = ggml_mul_mat(ctx0, k, q); + kq = ggml_mul(ctx0, kq, decay); + kq = ggml_mul(ctx0, kq, m_incl); + + // T = (I + A)⁻¹ via a lower-triangular solve: X = solve(I+A, −A), T = X + I. + // Topological node order keeps A strictly lower-triangular under any tree. + ggml_tensor * lhs = ggml_add(ctx0, attn, m_eye); + ggml_tensor * lin_solve = ggml_solve_tri(ctx0, lhs, ggml_neg(ctx0, attn), true, true, false); + ggml_tensor * t_mat = ggml_add(ctx0, lin_solve, m_eye); // [i, j, 1, H], T[j, i] entries + + // ṽ = T (β v − exp(g⁺) β k S₀): the corrected candidate updates. + ggml_tensor * v_corr = ggml_mul_mat(ctx0, + ggml_cont(ctx0, ggml_transpose(ctx0, v_b)), t_mat); // [S_v, t, 1, H] = (T v_b)_t + ggml_tensor * g_exp = ggml_exp(ctx0, g_ps); // [n, 1, 1, H] + ggml_tensor * kbg = ggml_mul(ctx0, + ggml_cont(ctx0, ggml_transpose(ctx0, k_b)), g_exp); // [t, S_k, 1, H] = e^{g⁺_t} β_t k_t + ggml_tensor * k_cd = ggml_mul_mat(ctx0, kbg, t_mat); // [S_k, t, 1, H] = (T e^{g⁺} k_b)_t + + ggml_tensor * s_r = ggml_reshape_4d(ctx0, s, S_v, S_v, 1, H_v); + ggml_tensor * v_state = ggml_mul_mat(ctx0, k_cd, s_r); // [t, S_v, 1, H] = S₀ᵀ (T e^{g⁺} k_b)_t + ggml_tensor * v_new = ggml_sub(ctx0, + ggml_cont(ctx0, ggml_transpose(ctx0, v_corr)), v_state); // [t, S_v, 1, H] = ṽ_t + + // o_t = e^{g⁺_t} S₀ᵀ q_t + Σ_u kq[u, t] ṽ_u + ggml_tensor * v_attn = ggml_mul_mat(ctx0, v_new, kq); // [S_v, t, 1, H] + ggml_tensor * g_exp_row = ggml_cont(ctx0, ggml_transpose(ctx0, g_exp)); // [1, n, 1, H] + ggml_tensor * q_g = ggml_mul(ctx0, q, g_exp_row); // [S_k, n, 1, H] + ggml_tensor * attn_inter = ggml_mul_mat(ctx0, s_r, q_g); // [S_v, t, 1, H] + ggml_tensor * o = ggml_add(ctx0, attn_inter, v_attn); // [S_v, n, 1, H] + + DeltaNetSpecLAResult r; + // [S_v, n, 1, H] → [S_v, H, n, 1] to match the fused op's output layout. + r.output = ggml_cont(ctx0, ggml_permute(ctx0, o, 0, 2, 3, 1)); + r.v_new = v_new; + r.g_ps = g_ps; + return r; +} + +void fill_specla_masks(const int32_t * parents, int n, + float * m_strict, float * m_incl, float * m_eye) { + // Masks are [ne0 = u (ancestor), ne1 = t (node)]: element t*n + u. + std::memset(m_strict, 0, sizeof(float) * (size_t)n * n); + std::memset(m_incl, 0, sizeof(float) * (size_t)n * n); + std::memset(m_eye, 0, sizeof(float) * (size_t)n * n); + for (int t = 0; t < n; t++) { + m_eye[(size_t)t * n + t] = 1.0f; + m_incl[(size_t)t * n + t] = 1.0f; + // Ancestor closure: t inherits its parent's ancestor-or-self row. + const int p = parents[t]; + if (p < 0) continue; + GGML_ASSERT(p < t); // DFS/topological order + for (int u = 0; u <= p; u++) { + const float a = m_incl[(size_t)p * n + u]; + m_incl[(size_t)t * n + u] = a; + m_strict[(size_t)t * n + u] = a; + } + } +} + +} // namespace dflash::common diff --git a/server/src/delta_net_specla.h b/server/src/delta_net_specla.h new file mode 100644 index 000000000..93b5cb329 --- /dev/null +++ b/server/src/delta_net_specla.h @@ -0,0 +1,75 @@ +// SpecLA topology-masked delta-net verification (arXiv:2607.16673 §4.2, §5.1). +// +// Single-window UT-transform verify for speculative windows (n <= 64 tokens), +// derived from build_delta_net_chunked but: +// - masks are host-filled ancestor-topology inputs instead of ggml_tri, so +// one builder serves both chain drafts (lower-triangular masks) and +// DFS-ordered tree drafts (ancestor masks); +// - the recurrent state is NOT advanced: verification reads the committed +// state S0 only, and besides per-node outputs it exposes the compact +// factors (corrected values v_new and path-cumulative gates g_ps) that +// DeltaConstruct needs to advance the state along the accepted path: +// S_A = exp(g_ps_A) * S0 + sum_{t in path(A)} exp(g_ps_A - g_ps_t) k_t (x) v_new_t +// +// See server/docs/SPECLA.md for the derivation and the commit-side consumer. +#pragma once + +#include "ggml.h" +#include +#include + +namespace dflash::common { + +struct DeltaNetSpecLAResult { + ggml_tensor * output; // [S_v, H_v, n_tokens, 1] — same layout as the fused op's output slice + ggml_tensor * v_new; // corrected values ṽ: [n_tokens, S_v, 1, H_v] + ggml_tensor * g_ps; // path-cumulative gate: [n_tokens, 1, 1, H_v] +}; + +// Heavy-light schedule consumed by the fused SpecLA recurrent kernels. Chains +// are grouped by dependency wave; every chain in a wave can execute in +// parallel because its parent boundary was produced by an earlier wave. +// The packed vector is a compact int32 ABI shared with the CUDA/HIP kernels. +struct SpecLAHLDSchedule { + std::vector packed; + int n_nodes = 0; + int n_chains = 0; + int n_waves = 0; + int n_boundaries = 0; + int max_parallel_chains = 0; +}; + +SpecLAHLDSchedule make_specla_hld_schedule(const int32_t * parents, + int n, + int pending_count); + +// q,k,v,g,b,s use the exact shapes build_delta_net_block passes to +// build_delta_net_chunked: q/k [S_k, H_v, n, 1] (post l2-norm, post repeat), +// v [S_v, H_v, n, 1], g/b [1, H_v, n, 1], s [S_v, S_v, H_v, 1]. +// +// m_strict / m_incl / m_eye are F32 [n, n] input tensors filled host-side +// with the draft topology over DFS-ordered nodes: element (ne0=u, ne1=t) is +// 1.0f when u is a strict ancestor of t (m_strict) or an ancestor-or-self of +// t (m_incl), else 0.0f; m_eye is the identity. For a chain draft these are +// simply the strict lower triangle and the lower triangle with diagonal. +// m_eye is a host input (not ggml_fill on a view, whose output would alias +// another node's data) so no graph node aliases live intermediate storage. +DeltaNetSpecLAResult build_delta_net_specla( + ggml_context * ctx0, + ggml_tensor * q, + ggml_tensor * k, + ggml_tensor * v, + ggml_tensor * g, + ggml_tensor * b, + ggml_tensor * s, + ggml_tensor * m_strict, + ggml_tensor * m_incl, + ggml_tensor * m_eye); + +// Host-side helpers: fill the three topology masks from a parent-pointer +// array over DFS/topologically ordered nodes (parents[t] < t, root = -1). +// `dst` buffers hold n*n floats each. For chains pass parents[t] = t-1. +void fill_specla_masks(const int32_t * parents, int n, + float * m_strict, float * m_incl, float * m_eye); + +} // namespace dflash::common diff --git a/server/src/internal.h b/server/src/internal.h index fffda41ac..e7eaaa898 100644 --- a/server/src/internal.h +++ b/server/src/internal.h @@ -406,12 +406,51 @@ struct TargetCache { // 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 - // delta layer. Holds the full concat(old_conv_state, qkv_new_tokens) - // that was fed to ggml_ssm_conv. Spec decode slices - // [commit_n..commit_n+kernel-2] along dim 0 for conv state rollback. + // conv_input_cache: normally [(kernel-1) + max_q_len, conv_channels] + // f32, one per delta layer, holding the full concat fed to + // ggml_ssm_conv. In SpecLA mode these are [conv_channels, max_q_len] + // strided views into conv_factor_all and hold raw current inputs. std::vector ssm_intermediate; // size = n_delta (48) std::vector conv_input_cache; // size = n_delta (48) + std::vector conv_input_cache_alt;// SpecLA factor bank 1 + ggml_tensor * conv_factor_all = nullptr; + ggml_tensor * conv_factor_all_alt = nullptr; + + // SpecLA factor buffers (allocated instead of ssm_intermediate when + // DFLASH_SPECLA=1 on the single-target path). Two token-major banks let a + // verify consume the preceding accepted path while writing its own raw + // factors without aliasing: + // factor_k_all: [S_k, H_v, n_delta, max_q_len] f32 + // factor_v_new_all: [S_v, H_v, n_delta, max_q_len] f32 + // factor_g_ps_all: [H_v, n_delta, max_q_len] f32 + // conv_factor_all: [conv_channels, n_delta, max_q_len] f32 + // The per-layer vectors below are persistent VIEWS into these (created + // after buffer allocation), shaped like independent per-layer buffers so + // the capture path treats them exactly like other cache tensors. + ggml_tensor * factor_k_all = nullptr; + ggml_tensor * factor_v_new_all = nullptr; + ggml_tensor * factor_g_ps_all = nullptr; + ggml_tensor * factor_k_all_alt = nullptr; + ggml_tensor * factor_v_new_all_alt = nullptr; + ggml_tensor * factor_g_ps_all_alt = nullptr; + std::vector factor_k; // view [S_k, H_v, max_q_len] + std::vector factor_v_new; // view [S_v, H_v, max_q_len] + std::vector factor_g_ps; // view [H_v, max_q_len] + std::vector factor_k_alt; + std::vector factor_v_new_alt; + std::vector factor_g_ps_alt; + // Host-side bank state. pending_bank is read by the next verify; the + // opposite bank receives that verify's factors. + int specla_pending_bank = 0; + int specla_pending_count = 0; + // Device-side accepted-index scratch and uploaded pointer tables. + ggml_tensor * specla_idx = nullptr; // i32 [max_q_len] + ggml_tensor * specla_state_ptrs = nullptr; // i64 [n_delta] + ggml_tensor * specla_conv_state_ptrs = nullptr; // i64 [n_delta] + // i64 [8]: base pointers for bank0 {k,v,g,conv}, then bank1. HLD kernels + // write their compact factors here directly, avoiding four cpy nodes per + // delta layer. Consolidated layout is token-major [t, layer, head, dim]. + ggml_tensor * specla_factor_ptrs = nullptr; // Rolling target layer features captured during target forward passes. // Shape [5 * hidden, target_feat_cap] bf16. target_feat_cap is typically @@ -611,7 +650,19 @@ bool migrate_prefill_cache(const TargetWeights & w, int max_ctx, int max_verify_tokens, ggml_backend_t backend, - TargetCache & cache); + TargetCache & cache, + bool enable_specla = true); + +// Compatibility commit for the fully factorized §4.2 fallback. The production +// HLD route instead keeps raw accepted factors pending and consumes them in +// the next state-resident verify (§5.2). Commits the bank the just-run verify +// wrote (1 - specla_pending_bank) into durable SSM and conv state, so it must +// be called before the host-side bank rotation. accepted_idx is in path order. +bool specla_commit_accepted(TargetCache & cache, + ggml_backend_t backend, + const int32_t * accepted_idx, + int A); + // ─── Target forward graph ───────────────────────────────────────── @@ -627,12 +678,32 @@ bool migrate_prefill_cache(const TargetWeights & w, // ssm_intermediate_states: [S_v, S_v, H_v, q_len] f32 // Element t on axis 3 holds the DeltaNet state after processing verify // token t. Rollback reads offset (commit_n-1) * S_v*S_v*H*elt. -// conv_input: [(kernel-1) + q_len, conv_channels, 1] f32 -// Full concat(old_conv_state, qkv_new_tokens) fed to ggml_ssm_conv. -// Rollback reads slice [commit_n..commit_n+kernel-2] along dim 0. +// conv_input: normally [(kernel-1) + q_len, conv_channels, 1] f32 with the +// full concat fed to ggml_ssm_conv. In SpecLA mode it is a strided +// [conv_channels, q_len, 1] view receiving raw current inputs. struct DeltaNetCapture { ggml_tensor * ssm_intermediate_states = nullptr; ggml_tensor * conv_input = nullptr; + + // SpecLA factor capture (DFLASH_SPECLA=1, docs/SPECLA.md). Persistent F32 + // aliases into the bank written by this verify. In the HLD path the + // historical field names hold raw serial-recurrence terms: + // factor_k: [S_k, H_v, max_verify_tokens] — post-l2norm keys + // factor_v_new: [S_v, H_v, max_verify_tokens] — raw Delta-rule delta + // factor_g_ps: [H_v, max_verify_tokens] — raw log-decay g + // The factorized compatibility route uses corrected ṽ and cumulative g⁺ + // in the same slots. ssm_intermediate_states stays null in SpecLA mode. + ggml_tensor * factor_k = nullptr; + ggml_tensor * factor_v_new = nullptr; + ggml_tensor * factor_g_ps = nullptr; + ggml_tensor * pending_factor_k = nullptr; + ggml_tensor * pending_factor_v_new = nullptr; + ggml_tensor * pending_factor_g = nullptr; + ggml_tensor * pending_conv_input = nullptr; + ggml_tensor * factor_ptrs = nullptr; + int factor_n_layers = 0; + int factor_layer = -1; + int pending_bank = 0; }; // One contiguous prompt chunk on the flattened token axis of a concurrent @@ -721,6 +792,17 @@ struct QwenGraphInputs { // layer into cache.q_cap (KVFlash target-QK scorer). Step-invariant: // node properties depend only on n_tokens and the layer index. bool q_capture = false; + + // SpecLA topology masks for the fully factorized compatibility route. + ggml_tensor * specla_m_strict = nullptr; + ggml_tensor * specla_m_incl = nullptr; + ggml_tensor * specla_m_eye = nullptr; + // Packed heavy-light schedule for the state-resident SpecLA kernels. + ggml_tensor * specla_hld = nullptr; + int specla_n_chains = 0; + int specla_n_waves = 0; + int specla_n_boundaries = 0; + int specla_max_parallel_chains = 0; }; struct QwenGraphOutputs { diff --git a/server/src/qwen35/graph_builders.cpp b/server/src/qwen35/graph_builders.cpp index ad3d56b58..ed22072c8 100644 --- a/server/src/qwen35/graph_builders.cpp +++ b/server/src/qwen35/graph_builders.cpp @@ -1,10 +1,14 @@ #include "graph_builders.h" +#include "common/specla_mode.h" +#include "delta_net_specla.h" + #include "ggml-alloc.h" #include #include #include +#include #include namespace dflash::common { @@ -515,6 +519,19 @@ bool build_target_step( ggml_set_input(sg.kv_write_rows); } + SpecLAHLDSchedule hld_schedule; + if (capture_delta_intermediate && specla_enabled() && !cache.factor_k.empty()) { + std::vector parents((size_t)n_tokens); + for (int t = 0; t < n_tokens; ++t) parents[(size_t)t] = t - 1; + hld_schedule = make_specla_hld_schedule( + parents.data(), n_tokens, cache.specla_pending_count); + if (hld_schedule.packed.empty()) return false; + sg.specla_hld = ggml_new_tensor_1d( + sg.ctx, GGML_TYPE_I32, hld_schedule.packed.size()); + ggml_set_name(sg.specla_hld, "specla_hld"); + ggml_set_input(sg.specla_hld); + } + QwenGraphInputs gi{}; gi.inp_embed = sg.inp_embed; gi.positions = sg.positions; @@ -541,6 +558,14 @@ bool build_target_step( gi.logits_row_indices = sg.logits_row_indices; gi.prefill_segments = prefill_segments; gi.n_prefill_segments = n_prefill_segments; + gi.specla_m_strict = sg.specla_m_strict; + gi.specla_m_incl = sg.specla_m_incl; + gi.specla_m_eye = sg.specla_m_eye; + gi.specla_hld = sg.specla_hld; + gi.specla_n_chains = hld_schedule.n_chains; + gi.specla_n_waves = hld_schedule.n_waves; + gi.specla_n_boundaries = hld_schedule.n_boundaries; + gi.specla_max_parallel_chains = hld_schedule.max_parallel_chains; QwenGraphOutputs go = build_qwen35_graph(sg.ctx, sg.gf, w, cache, gi); if (!go.logits) return false; @@ -557,7 +582,12 @@ bool build_target_step( if (!sg.alloc) { sg.alloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); } - return ggml_gallocr_alloc_graph(sg.alloc, sg.gf); + if (!ggml_gallocr_alloc_graph(sg.alloc, sg.gf)) return false; + if (sg.specla_hld) { + ggml_backend_tensor_set(sg.specla_hld, hld_schedule.packed.data(), 0, + hld_schedule.packed.size()*sizeof(int32_t)); + } + return true; } // ── build_target_step_tree ────────────────────────────────────── @@ -570,7 +600,8 @@ bool build_target_step_tree( int kv_start, int n_tokens, int fa_window, - int kq_stride_pad) { + int kq_stride_pad, + const SpecLAHLDSchedule * hld_schedule) { step_graph_free(sg); ggml_init_params ip{}; @@ -600,6 +631,29 @@ bool build_target_step_tree( ggml_set_name(sg.parent_ids, "parent_ids"); ggml_set_input(sg.parent_ids); + // SpecLA tree verify: ancestor masks over the DFS-ordered nodes replace + // the sequential kernel's parent_ids state fanout (parent_ids still + // steers the tree conv). Host-filled by verify_tree from tree.parents. + if (specla_enabled() && !cache.factor_k.empty() && hld_schedule) { + if (hld_schedule->n_nodes != n_tokens || hld_schedule->packed.empty()) { + return false; + } + sg.specla_hld = ggml_new_tensor_1d( + sg.ctx, GGML_TYPE_I32, hld_schedule->packed.size()); + ggml_set_name(sg.specla_hld, "specla_hld"); + ggml_set_input(sg.specla_hld); + } else if (specla_enabled() && !cache.factor_k.empty()) { + sg.specla_m_strict = ggml_new_tensor_2d(sg.ctx, GGML_TYPE_F32, n_tokens, n_tokens); + sg.specla_m_incl = ggml_new_tensor_2d(sg.ctx, GGML_TYPE_F32, n_tokens, n_tokens); + sg.specla_m_eye = ggml_new_tensor_2d(sg.ctx, GGML_TYPE_F32, n_tokens, n_tokens); + ggml_set_name(sg.specla_m_strict, "specla_m_strict"); + ggml_set_name(sg.specla_m_incl, "specla_m_incl"); + ggml_set_name(sg.specla_m_eye, "specla_m_eye"); + ggml_set_input(sg.specla_m_strict); + ggml_set_input(sg.specla_m_incl); + ggml_set_input(sg.specla_m_eye); + } + sg.gf = ggml_new_graph_custom(sg.ctx, 16384, false); QwenGraphInputs gi{}; @@ -612,6 +666,14 @@ bool build_target_step_tree( gi.capture_layers = true; gi.capture_delta_intermediate = true; gi.parent_ids = sg.parent_ids; + gi.specla_m_strict = sg.specla_m_strict; + gi.specla_m_incl = sg.specla_m_incl; + gi.specla_m_eye = sg.specla_m_eye; + gi.specla_hld = sg.specla_hld; + gi.specla_n_chains = hld_schedule ? hld_schedule->n_chains : 0; + gi.specla_n_waves = hld_schedule ? hld_schedule->n_waves : 0; + gi.specla_n_boundaries = hld_schedule ? hld_schedule->n_boundaries : 0; + gi.specla_max_parallel_chains = hld_schedule ? hld_schedule->max_parallel_chains : 0; QwenGraphOutputs go = build_qwen35_graph(sg.ctx, sg.gf, w, cache, gi); if (!go.logits) return false; @@ -627,7 +689,12 @@ bool build_target_step_tree( if (!sg.alloc) { sg.alloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); } - return ggml_gallocr_alloc_graph(sg.alloc, sg.gf); + if (!ggml_gallocr_alloc_graph(sg.alloc, sg.gf)) return false; + if (sg.specla_hld) { + ggml_backend_tensor_set(sg.specla_hld, hld_schedule->packed.data(), 0, + hld_schedule->packed.size()*sizeof(int32_t)); + } + return true; } diff --git a/server/src/qwen35/graph_builders.h b/server/src/qwen35/graph_builders.h index cdbaf75ed..a56a0f79b 100644 --- a/server/src/qwen35/graph_builders.h +++ b/server/src/qwen35/graph_builders.h @@ -17,6 +17,7 @@ #include "step_graph.h" #include "attn_masks.h" // align_up, KQ_MASK_PAD #include "internal.h" // TargetWeights, TargetCache +#include "delta_net_specla.h" #include "ggml.h" #include "ggml-backend.h" @@ -144,7 +145,8 @@ bool build_target_step_tree( int kv_start, int n_tokens, int fa_window = 0, - int kq_stride_pad = KQ_MASK_PAD); + int kq_stride_pad = KQ_MASK_PAD, + const SpecLAHLDSchedule * specla_hld = nullptr); // LM-head projection: project draft hidden states through the target output matrix. bool build_lm_head_projection_step( diff --git a/server/src/qwen35/qwen35_backend.cpp b/server/src/qwen35/qwen35_backend.cpp index 490f86d3b..112cf0e0a 100644 --- a/server/src/qwen35/qwen35_backend.cpp +++ b/server/src/qwen35/qwen35_backend.cpp @@ -17,6 +17,7 @@ #endif #include "common/io_utils.h" #include "common/restore_delta.h" +#include "common/specla_mode.h" #include "qwen35_tensor_parallel.h" #include "qwen3/qwen3_drafter.h" #include "qwen3/qwen3_kvflash_scorer.h" @@ -28,6 +29,7 @@ #include #include +#include #include #include #include @@ -1587,9 +1589,11 @@ int Qwen35Backend::do_prefill(const std::vector & tokens, const int max_verify_tokens = cfg_.ddtree_mode ? std::max(dw_.block_size, cfg_.ddtree_budget + 1) : dw_.block_size; + const bool enable_specla = cfg_.fast_rollback && + !cfg_.device.is_tensor_parallel() && !kvflash_active(); if (!migrate_prefill_cache(w_, cfg_.device.max_ctx, max_verify_tokens, - target_backend_, cache_)) { + target_backend_, cache_, enable_specla)) { std::fprintf(stderr, "prefill: rollback cache migration failed: %s\n", dflash27b_last_error()); return -1; @@ -2428,15 +2432,22 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, // ── DFlash spec-decode: draft → verify → accept → replay ────────── DFlashTarget * target = dflash_target(); + auto finish_speculative_state = [&]() { + if (target->finish_speculative_state()) return true; + std::fprintf(stderr, "spec-decode: final SpecLA state flush failed\n"); + return false; + }; const bool use_remote_draft = cfg_.remote_draft.enabled() && remote_draft_.active(); const int q_len = dw_.block_size > 0 ? dw_.block_size : DFLASH27B_DRAFT_BLOCK_SIZE; const int max_verify_tokens = cfg_.ddtree_mode ? std::max(dw_.block_size, cfg_.ddtree_budget + 1) : dw_.block_size; if ((cfg_.fast_rollback || cfg_.ddtree_mode) && !cache_.rollback_ctx) { + const bool enable_specla = cfg_.fast_rollback && + !cfg_.device.is_tensor_parallel() && !kvflash_active(); if (!migrate_prefill_cache(w_, cfg_.device.max_ctx, max_verify_tokens, - target_backend_, cache_)) { + target_backend_, cache_, enable_specla)) { std::fprintf(stderr, "spec-decode: rollback cache migration failed: %s\n", dflash27b_last_error()); return false; @@ -2462,7 +2473,8 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, int n_hint_accepted = 0; int target_forwards = 0; const ChainRollbackPolicy rollback_policy = - resolve_chain_rollback_policy(cfg_.device.is_tensor_parallel()); + resolve_chain_rollback_policy(cfg_.device.is_tensor_parallel(), + target->exact_fast_rollback()); const int fast_rollback_threshold = rollback_policy.fast_rollback_threshold; RollbackDiag rollback_diag; @@ -2525,10 +2537,12 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, cache_.last_tok = out_tokens.back(); const int ar_n_gen = n_gen - n_generated; if (ar_n_gen <= 0) { + if (!finish_speculative_state()) return false; log_target_forward_stats(); io.emit(-1); return true; } + if (!finish_speculative_state()) return false; BudgetHook tail_hook = budget_hook ? *budget_hook : BudgetHook{}; bool ok = do_ar_decode(committed, ar_n_gen, out_tokens, io, tail_hook, forced_close_out, @@ -2559,6 +2573,7 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, !use_remote_draft && draft_feature_mirror_can_view(feature_mirror_, committed, draft_ctx, mirror_slot0); + bool used_draft_kv = false; const auto profile_draft_start = profile_start(); if (use_remote_draft) { local_hidden.clear(); @@ -2590,6 +2605,7 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, } } if (use_draft_kv) { + used_draft_kv = true; if (!draft_kv_begin_step(draft_kv_, dw_, draft_backend_, feature_mirror_, committed)) { std::fprintf(stderr, "spec-decode: draft-kv step prep failed\n"); @@ -2693,24 +2709,118 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, if (use_tree_verify) { const int L = q_len - 1; - const int K = (cfg_.ddtree_budget > L) ? 8 : 1; - std::vector top_lp; - std::vector top_ids; - const auto profile_project_start = profile_start(); - if (!target->project_hidden_to_topk(local_hidden.data(), q_len, K, - cfg_.ddtree_temp, top_lp, top_ids)) { - std::fprintf(stderr, "spec-decode: ddtree topk projection failed\n"); - step_graph_destroy(draft_sg); - return false; - } - profile_add(profile_project_s, profile_project_start); - // Tree depth L draws from draft rows 1..q_len-1 (row 0 = the seed). - // Known limitation: branch descendants beyond depth 1 still come - // from one spine-conditioned block-draft forward, so a confident - // draft may not beat the chain. - DDTree tree = build_ddtree(top_lp.data() + (size_t)K, top_ids.data() + (size_t)K, - L, K, cfg_.ddtree_budget, cfg_.ddtree_chain_seed); - const int N = cfg_.ddtree_budget + 1; // fixed alloc width + // The paper's end-to-end tree route uses top-k=4. Wider top-k + // spends projection and scheduling work on low-ranked siblings; + // keep the legacy width only outside SpecLA. + const int K = (cfg_.ddtree_budget > L) + ? (target->exact_fast_rollback() + ? std::min(specla_tree_topk(), w_.n_vocab) + : 8) + : 1; + DDTree tree; + if (target->exact_fast_rollback() && K > 1 && + specla_conditional_draft_enabled()) { + // SpecLA §6.2: every expanded branch is proposed from that + // branch's exact token prefix, avoiding the old one-spine + // approximation. This is experimental for the current + // five-layer drafter because every prefix requires a rerun. + bool conditional_ok = true; + std::vector row_hidden((size_t)hidden); + DDTreeConditionalTopK conditional_topk = + [&](const std::vector & prefix, int next_depth, + std::vector & top_lp, + std::vector & top_ids) -> bool { + if (!conditional_ok || next_depth < 1 || next_depth > L || + (int)prefix.size() != next_depth - 1) { + conditional_ok = false; + return false; + } + + if (prefix.empty()) { + std::memcpy(row_hidden.data(), + local_hidden.data() + (size_t)hidden, + sizeof(float) * (size_t)hidden); + } else { + noise_ids[0] = last_tok; + for (int i = 1; i < q_len; ++i) { + noise_ids[(size_t)i] = i <= (int)prefix.size() + ? prefix[(size_t)i - 1] + : target->mask_token_id(); + } + if (!target->embed_tokens(noise_ids.data(), q_len, + noise_embed.data())) { + conditional_ok = false; + return false; + } + const auto branch_draft_start = profile_start(); + ggml_tensor * branch_input = used_draft_kv + ? draft_kv_.inp_embed : draft_sg.inp_embed; + ggml_cgraph * branch_graph = used_draft_kv + ? draft_kv_.gf : draft_sg.gf; + ggml_tensor * branch_hidden = used_draft_kv + ? draft_kv_.hidden_states : draft_sg.hidden_states; + if (!branch_input || !branch_graph || !branch_hidden) { + conditional_ok = false; + return false; + } + ggml_backend_tensor_set(branch_input, noise_embed.data(), 0, + sizeof(float) * noise_embed.size()); + if (ggml_backend_graph_compute(draft_backend_, branch_graph) != + GGML_STATUS_SUCCESS) { + conditional_ok = false; + return false; + } + ggml_backend_tensor_get( + branch_hidden, row_hidden.data(), + (size_t)next_depth * hidden * sizeof(float), + sizeof(float) * (size_t)hidden); + profile_add(profile_draft_s, branch_draft_start); + } + + const auto branch_project_start = profile_start(); + const bool ok = target->project_hidden_to_topk( + row_hidden.data(), 1, K, cfg_.ddtree_temp, + top_lp, top_ids); + profile_add(profile_project_s, branch_project_start); + conditional_ok = conditional_ok && ok; + return ok; + }; + tree = build_ddtree_conditional( + conditional_topk, L, K, cfg_.ddtree_budget, + cfg_.ddtree_chain_seed, cfg_.ddtree_tau); + if (!conditional_ok || tree.n_nodes == 0) { + std::fprintf(stderr, + "spec-decode: conditioned ddtree proposal failed\n"); + step_graph_destroy(draft_sg); + return false; + } + } else { + std::vector top_lp; + std::vector top_ids; + const auto profile_project_start = profile_start(); + if (!target->project_hidden_to_topk(local_hidden.data(), q_len, K, + cfg_.ddtree_temp, + top_lp, top_ids)) { + std::fprintf(stderr, + "spec-decode: ddtree topk projection failed\n"); + step_graph_destroy(draft_sg); + return false; + } + profile_add(profile_project_s, profile_project_start); + tree = build_ddtree( + top_lp.data() + (size_t)K, + top_ids.data() + (size_t)K, + L, K, cfg_.ddtree_budget, cfg_.ddtree_chain_seed, + cfg_.ddtree_tau); + } + // SpecLA schedules the retained topology directly. Never execute + // fake padding nodes: confidence pruning must reduce target work, + // and this loop already rebuilds the topology-dependent graph. + const int N = target->exact_fast_rollback() + ? 1 + tree.n_nodes + : (std::isfinite(cfg_.ddtree_tau) + ? 1 + tree.n_nodes + : cfg_.ddtree_budget + 1); std::vector flat_tokens((size_t)N, 0); flat_tokens[0] = last_tok; @@ -2905,7 +3015,13 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, if (can_commit_bonus) { const int bonus_pos = committed + total_emitted; std::vector bonus_vec(1, next_token); - if (!target->verify_batch(bonus_vec, bonus_pos, bonus_last_tok, nullptr)) { + // A delayed-state target must consume the accepted tree path + // inside this verify and capture the committed bonus as the + // next pending factor. A normal target has already materialized + // the tree rollback and may keep the ordinary writeback path. + const bool delayed_bonus = target->exact_fast_rollback(); + if (!target->verify_batch(bonus_vec, bonus_pos, bonus_last_tok, + nullptr, delayed_bonus)) { std::fprintf(stderr, "spec-decode: tree bonus replay failed\n"); step_graph_destroy(draft_sg); return false; @@ -2921,6 +3037,12 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, step_graph_destroy(draft_sg); return false; } + if (delayed_bonus && !target->rollback_to(bonus_pos, 1)) { + std::fprintf(stderr, + "spec-decode: tree bonus factor commit failed\n"); + step_graph_destroy(draft_sg); + return false; + } out_tokens.push_back(next_token); io.emit(next_token); @@ -3100,9 +3222,14 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, fast_rolled_back = true; rollback_diag.record_fast_rollback(accept_n); } else { - // Rollback failed (CUDA error / unsupported state type). The - // pre-verify snapshot is still valid, so degrade to the legacy - // restore+replay path below instead of aborting the request. + if (!target->rollback_failure_is_recoverable()) { + std::fprintf(stderr, "spec-decode: rollback_to failed after " + "an in-place commit attempt; aborting\n"); + step_graph_destroy(draft_sg); + return false; + } + // The pre-verify snapshot is still valid, so degrade to the + // legacy restore+replay path below. std::fprintf(stderr, "spec-decode: rollback_to failed; " "falling back to restore+replay\n"); rollback_diag.record_failed_fallback(); @@ -3303,10 +3430,12 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, (float)((double)n_accept_sum / (double)total_draft_pos); const int ar_n_gen = n_gen - n_generated; if (ar_n_gen <= 0) { + if (!finish_speculative_state()) return false; log_target_forward_stats(); io.emit(-1); return true; } + if (!finish_speculative_state()) return false; BudgetHook tail_hook = budget_hook ? *budget_hook : BudgetHook{}; bool ok = do_ar_decode(committed, ar_n_gen, out_tokens, io, tail_hook, forced_close_out, @@ -3346,10 +3475,12 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, (float)((double)n_accept_sum / (double)total_draft_pos); const int ar_n_gen = n_gen - n_generated; if (ar_n_gen <= 0) { + if (!finish_speculative_state()) return false; log_target_forward_stats(); io.emit(-1); return true; } + if (!finish_speculative_state()) return false; BudgetHook tail_hook = budget_hook ? *budget_hook : BudgetHook{}; tail_hook.close_token_ids.clear(); bool ok = do_ar_decode(committed, ar_n_gen, out_tokens, io, @@ -3362,6 +3493,10 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, if (hit_eos) break; } + if (!finish_speculative_state()) { + step_graph_destroy(draft_sg); + return false; + } step_graph_destroy(draft_sg); auto t_dec1 = std::chrono::steady_clock::now(); diff --git a/server/src/qwen35/qwen35_backend.h b/server/src/qwen35/qwen35_backend.h index a6a508f73..c2ad20d17 100644 --- a/server/src/qwen35/qwen35_backend.h +++ b/server/src/qwen35/qwen35_backend.h @@ -18,6 +18,8 @@ #include "placement/remote_draft_config.h" #include "step_graph.h" #include "ddtree.h" + +#include #include "dflash_feature_ring.h" #include "common/dflash_draft_kv.h" #include "common/concurrency/paged_kv_pool.h" @@ -83,6 +85,8 @@ struct Qwen35Config { int ddtree_budget = 22; float ddtree_temp = 1.0f; bool ddtree_chain_seed = true; + // SpecLA confidence margin on cumulative path log-prob (off when inf). + float ddtree_tau = std::numeric_limits::infinity(); bool use_feature_mirror = false; }; diff --git a/server/src/qwen35/qwen35_dflash_target.cpp b/server/src/qwen35/qwen35_dflash_target.cpp index e6849e918..bf6e76e5b 100644 --- a/server/src/qwen35/qwen35_dflash_target.cpp +++ b/server/src/qwen35/qwen35_dflash_target.cpp @@ -1,11 +1,13 @@ // Qwen35DFlashTarget — DFlashTarget adapter for qwen35 hybrid models. #include "qwen35_dflash_target.h" +#include "delta_net_specla.h" #include "graph_builders.h" #include "step_graph.h" #include "attn_masks.h" #include "prefill_helpers.h" #include "common/geometric_draft_topk_cuda.h" +#include "common/specla_commit_cuda.h" #include "ggml-backend-impl.h" // gpu_runtime_compat.h maps the raw cudaStream_t / cudaMemcpy* symbols used // below (rollback_to / rollback_to_tree) onto their HIP equivalents. Without @@ -414,11 +416,26 @@ bool Qwen35DFlashTarget::verify_tree( } const int hidden = w_.n_embd; + // Root-inclusive topology. Padding nodes are independent root children; + // their outputs are masked, but scheduling them keeps every downstream + // activation initialized without contaminating a real branch. + std::vector parent_ids(N, 0); + parent_ids[0] = -1; + for (int i = 1; i < N_actual; i++) parent_ids[i] = (int32_t)tree.parents[i]; + SpecLAHLDSchedule hld; + const SpecLAHLDSchedule * hld_ptr = nullptr; + if (specla_active()) { + hld = make_specla_hld_schedule( + parent_ids.data(), N, cache_.specla_pending_count); + if (hld.packed.empty()) return false; + hld_ptr = &hld; + } + // Tree-verify graph: ancestor-masked batched forward over DFS-ordered nodes. // Capture per-node SSM intermediates so rollback_to_tree() can restore. if (!build_target_step_tree(sg_, w_, cache_, backend_, /*kv_start=*/committed, /*n_tokens=*/N, - fa_window_, kq_stride_pad_)) { + fa_window_, kq_stride_pad_, hld_ptr)) { std::fprintf(stderr, "verify_tree: build_target_step_tree failed\n"); return false; } @@ -470,10 +487,28 @@ bool Qwen35DFlashTarget::verify_tree( std::fprintf(stderr, "verify_tree: step graph has no parent_ids tensor\n"); return false; } - std::vector parent_ids(N, 0); - parent_ids[0] = -1; - for (int i = 1; i < N_actual; i++) parent_ids[i] = (int32_t)tree.parents[i]; - ggml_backend_tensor_set(sg_.parent_ids, parent_ids.data(), 0, sizeof(int32_t) * N); + // HLD carries the topology for both DeltaNet and convolution, so the + // legacy parent tensor is intentionally disconnected and unallocated. + if (sg_.parent_ids->buffer) { + ggml_backend_tensor_set(sg_.parent_ids, parent_ids.data(), 0, + sizeof(int32_t) * N); + } else if (!sg_.specla_hld) { + std::fprintf(stderr, "verify_tree: parent_ids tensor is unallocated\n"); + return false; + } + + // SpecLA tree verify: ancestor masks over the same root-inclusive flat + // node order. Padding nodes hang off the root; their outputs/factors are + // never read. + if (sg_.specla_m_strict) { + std::vector ms((size_t)N * N); + std::vector mi((size_t)N * N); + std::vector me((size_t)N * N); + fill_specla_masks(parent_ids.data(), N, ms.data(), mi.data(), me.data()); + ggml_backend_tensor_set(sg_.specla_m_strict, ms.data(), 0, sizeof(float) * ms.size()); + ggml_backend_tensor_set(sg_.specla_m_incl, mi.data(), 0, sizeof(float) * mi.size()); + ggml_backend_tensor_set(sg_.specla_m_eye, me.data(), 0, sizeof(float) * me.size()); + } auto st = ggml_backend_graph_compute(backend_, sg_.gf); if (st != GGML_STATUS_SUCCESS) { @@ -544,19 +579,82 @@ bool Qwen35DFlashTarget::rollback_to_tree( const int n_delta = (int)sg_.delta_captures.size(); GGML_ASSERT(!cache_.ssm_state.empty()); const bool meta_backend = is_meta_tensor(cache_.ssm_state.front()); + const bool specla = specla_active(); + if (specla && meta_backend) return false; // TP meta is outside SpecLA scope + if (specla && !sg_.specla_hld) { + // The fully factorized tree fallback would need per-ancestry conv + // commit from the accepted DFS path; production tree verify always + // builds an HLD schedule, so fail closed instead of committing a + // chain-window conv state. + std::fprintf(stderr, + "rollback_to_tree: factorized SpecLA tree fallback is unsupported\n"); + return false; + } + + if (specla) { + if (!cache_.factor_k_all || !cache_.factor_v_new_all || + !cache_.factor_g_ps_all || !cache_.conv_factor_all || + !cache_.factor_k_all_alt || !cache_.factor_v_new_all_alt || + !cache_.factor_g_ps_all_alt || !cache_.conv_factor_all_alt) { + return false; + } + for (int il = 0; il < n_delta; il++) { + const DeltaNetCapture & cap = sg_.delta_captures[il]; + if (!cap.factor_k || !cap.factor_v_new || !cap.factor_g_ps || + !cap.conv_input || rollback_dfs >= cap.factor_k->ne[2]) return false; + } + + SpeclaFactorBanks banks; + banks.k[0] = (float *)cache_.factor_k_all->data; + banks.v[0] = (float *)cache_.factor_v_new_all->data; + banks.g[0] = (float *)cache_.factor_g_ps_all->data; + banks.conv[0] = (float *)cache_.conv_factor_all->data; + banks.k[1] = (float *)cache_.factor_k_all_alt->data; + banks.v[1] = (float *)cache_.factor_v_new_all_alt->data; + banks.g[1] = (float *)cache_.factor_g_ps_all_alt->data; + banks.conv[1] = (float *)cache_.conv_factor_all_alt->data; + + const int old_pending_bank = cache_.specla_pending_bank; + if (walked_sibling) { + if (!cache_.specla_idx || !cache_.specla_idx->data) return false; + std::vector idx(accepted_dfs.begin(), accepted_dfs.end()); + ggml_backend_tensor_set(cache_.specla_idx, idx.data(), 0, + idx.size()*sizeof(int32_t)); + } + int new_pending_bank = old_pending_bank; + if (!specla_rotate_pending_factors( + banks, + walked_sibling ? (const int32_t *)cache_.specla_idx->data : nullptr, + old_pending_bank, walked_sibling, commit_n, + (int)cache_.factor_k_all->ne[0], + (int)cache_.factor_v_new_all->ne[0], + (int)cache_.factor_k_all->ne[1], + n_delta, (int)cache_.conv_factor_all->ne[0], + /*stream=*/nullptr, &new_pending_bank)) { + std::fprintf(stderr, "rollback_to_tree: SpecLA factor rotation failed\n"); + return false; + } + cache_.specla_pending_bank = new_pending_bank; + cache_.specla_pending_count = commit_n; + } + cudaStream_t stream = nullptr; for (int il = 0; il < n_delta; il++) { const DeltaNetCapture & cap = sg_.delta_captures[il]; - if (!cap.ssm_intermediate_states || !cap.conv_input) { + if ((!specla && !cap.ssm_intermediate_states) || !cap.conv_input) { std::fprintf(stderr, "rollback_to_tree: missing capture at layer %d\n", il); return false; } + if (specla) { + continue; + } if (rollback_dfs >= (int)cap.ssm_intermediate_states->ne[3]) { std::fprintf(stderr, "rollback_to_tree: rollback_dfs %d >= captured slots %d (layer %d)\n", rollback_dfs, (int)cap.ssm_intermediate_states->ne[3], il); return false; } // SSM state ← intermediate[rollback_dfs] (dequantize Q8_0/F16 → f32). + { const size_t ssm_elems = (size_t)cache_.ssm_state[il]->ne[0] * (size_t)cache_.ssm_state[il]->ne[1] * @@ -598,6 +696,7 @@ bool Qwen35DFlashTarget::rollback_to_tree( (int64_t)ssm_elems, stream); } } + } // end non-SpecLA SSM restore // Conv state ← the K-1 most recent inputs along rollback_dfs's ancestry. const int K_conv = 4; @@ -752,6 +851,10 @@ bool Qwen35DFlashTarget::rollback_to_tree( } bool Qwen35DFlashTarget::snapshot_kv() { + // SpecLA applies only the already-committed pending path to durable state + // during verify; current candidates remain in the factor bank. There is + // therefore no speculative durable mutation to snapshot or undo. + if (specla_active()) return true; if (!cache_.ssm_state.empty() && is_meta_tensor(cache_.ssm_state.front())) { return copy_meta_recurrent_state( cache_.ssm_state, cache_.conv_state, @@ -761,6 +864,16 @@ bool Qwen35DFlashTarget::snapshot_kv() { } bool Qwen35DFlashTarget::restore_kv() { + if (specla_active()) { + // A successful SpecLA verify has already folded the *previous* + // accepted factors into the durable state, while the factors produced + // by that verify are still only candidates. Restoring therefore means + // dropping the pending candidate selection, not copying a full state + // snapshot. Leaving the old count live would apply it a second time + // when a replay graph starts. + cache_.specla_pending_count = 0; + return true; + } if (!cache_.ssm_state.empty() && is_meta_tensor(cache_.ssm_state.front())) { return copy_meta_recurrent_state( cache_.ssm_state_snap, cache_.conv_state_snap, @@ -770,11 +883,11 @@ bool Qwen35DFlashTarget::restore_kv() { } bool Qwen35DFlashTarget::supports_fast_rollback() const { - // Pure capability. Fast-rollback only restores recurrent SSM/conv state and - // defers the bonus token, so it is pager-safe even while paging: committed - // KV rows are written slot-mapped by verify_batch(), and the deferred bonus - // is re-fed at the next committed position on the following step. - return fast_rollback_; + // KVFlash requires the set-rows write path, which is mutually exclusive + // with recurrent capture. Report the runtime capability, not just the + // requested mode, so callers snapshot and replay instead of attempting a + // rollback from missing/stale captures. + return fast_rollback_ && pager_ == nullptr; } bool Qwen35DFlashTarget::rollback_to(int base_pos, int commit_n) { @@ -811,6 +924,12 @@ bool Qwen35DFlashTarget::rollback_to(int base_pos, int commit_n) { return false; } + // SpecLA kept the current candidates out of durable state, so acceptance + // must rotate their factor bank even when the whole window matched. + if (specla_active()) { + return rollback_to_specla(base_pos, commit_n); + } + // If all tokens accepted, the SSM state after processing all q_len tokens // is exactly what we want — no rollback needed, just fix cur_pos. const int q_len = cache_.cur_pos - base_pos; @@ -956,6 +1075,117 @@ bool Qwen35DFlashTarget::rollback_to(int base_pos, int commit_n) { return true; } +bool Qwen35DFlashTarget::rollback_to_specla(int base_pos, int commit_n) { + const int n_delta = (int)sg_.delta_captures.size(); + const int q_len = cache_.cur_pos - base_pos; + if (n_delta == 0 || commit_n <= 0 || commit_n > q_len) return false; + if (!cache_.ssm_state.empty() && is_meta_tensor(cache_.ssm_state.front())) { + // Tensor-parallel meta backends are outside SpecLA's scope; fail so + // the caller degrades to restore+replay (which stays correct: the + // capture verify did not mutate state and replay runs writeback-on). + return false; + } + + // The just-computed verify wrote the bank opposite the one it consumed. + // A chain acceptance is already compact and in path order, so rollback is + // only a host-side bank rotation. The next verify applies these factors + // inside its state-resident kernels. + if (!cache_.factor_k_all || !cache_.factor_v_new_all || + !cache_.factor_g_ps_all || !cache_.conv_factor_all || + !cache_.factor_k_all_alt || !cache_.factor_v_new_all_alt || + !cache_.factor_g_ps_all_alt || !cache_.conv_factor_all_alt) { + return false; + } + for (int il = 0; il < n_delta; il++) { + const DeltaNetCapture & cap = sg_.delta_captures[il]; + if (!cap.factor_k || !cap.factor_v_new || !cap.factor_g_ps || + !cap.conv_input || commit_n > cap.factor_k->ne[2] || + commit_n > cap.conv_input->ne[1]) { + std::fprintf(stderr, "rollback_to_specla: factor capture bad layer=%d\n", il); + return false; + } + } + + if (!sg_.specla_hld) { + // Fully factorized fallback: no next HLD kernel will consume the + // pending factors, so commit the just-verified bank to durable state + // (and its accepted conv window) immediately. + std::vector idx((size_t)commit_n); + for (int i = 0; i < commit_n; i++) idx[(size_t)i] = i; + if (!specla_commit_accepted(cache_, backend_, idx.data(), commit_n)) { + std::fprintf(stderr, "rollback_to_specla: factorized commit failed\n"); + return false; + } + cache_.cur_pos = base_pos + commit_n; + return true; + } + + SpeclaFactorBanks banks; + banks.k[0] = (float *)cache_.factor_k_all->data; + banks.v[0] = (float *)cache_.factor_v_new_all->data; + banks.g[0] = (float *)cache_.factor_g_ps_all->data; + banks.conv[0] = (float *)cache_.conv_factor_all->data; + banks.k[1] = (float *)cache_.factor_k_all_alt->data; + banks.v[1] = (float *)cache_.factor_v_new_all_alt->data; + banks.g[1] = (float *)cache_.factor_g_ps_all_alt->data; + banks.conv[1] = (float *)cache_.conv_factor_all_alt->data; + + int new_pending_bank = cache_.specla_pending_bank; + if (!specla_rotate_pending_factors( + banks, /*idx_dev=*/nullptr, cache_.specla_pending_bank, + /*walked_sibling=*/false, commit_n, + (int)cache_.factor_k_all->ne[0], + (int)cache_.factor_v_new_all->ne[0], + (int)cache_.factor_k_all->ne[1], + n_delta, (int)cache_.conv_factor_all->ne[0], + /*stream=*/nullptr, &new_pending_bank)) { + std::fprintf(stderr, "rollback_to_specla: factor bank rotation failed\n"); + return false; + } + cache_.specla_pending_bank = new_pending_bank; + cache_.specla_pending_count = commit_n; + cache_.cur_pos = base_pos + commit_n; + return true; +} + +bool Qwen35DFlashTarget::finish_speculative_state() { + if (!specla_active() || cache_.specla_pending_count == 0) return true; + if (!cache_.factor_k_all || !cache_.factor_v_new_all || + !cache_.factor_g_ps_all || !cache_.conv_factor_all || + !cache_.factor_k_all_alt || !cache_.factor_v_new_all_alt || + !cache_.factor_g_ps_all_alt || !cache_.conv_factor_all_alt || + !cache_.specla_state_ptrs || !cache_.specla_conv_state_ptrs) { + return false; + } + + SpeclaFactorBanks banks; + banks.k[0] = (float *)cache_.factor_k_all->data; + banks.v[0] = (float *)cache_.factor_v_new_all->data; + banks.g[0] = (float *)cache_.factor_g_ps_all->data; + banks.conv[0] = (float *)cache_.conv_factor_all->data; + banks.k[1] = (float *)cache_.factor_k_all_alt->data; + banks.v[1] = (float *)cache_.factor_v_new_all_alt->data; + banks.g[1] = (float *)cache_.factor_g_ps_all_alt->data; + banks.conv[1] = (float *)cache_.conv_factor_all_alt->data; + + const int n_delta = (int)cache_.ssm_state.size(); + const int conv_channels = (int)cache_.conv_factor_all->ne[0]; + if (!specla_flush_pending_factors( + banks, + (float * const *)cache_.specla_state_ptrs->data, + (float * const *)cache_.specla_conv_state_ptrs->data, + cache_.specla_pending_bank, cache_.specla_pending_count, + (int)cache_.factor_k_all->ne[0], + (int)cache_.factor_v_new_all->ne[0], + (int)cache_.factor_k_all->ne[1], + n_delta, conv_channels, w_.ssm_d_conv, + /*stream=*/nullptr)) { + return false; + } + cache_.specla_pending_count = 0; + return true; +} + bool Qwen35DFlashTarget::is_eos(int token) const { return is_eos_tok(token, w_); } diff --git a/server/src/qwen35/qwen35_dflash_target.h b/server/src/qwen35/qwen35_dflash_target.h index 3c8864b6b..9c1797652 100644 --- a/server/src/qwen35/qwen35_dflash_target.h +++ b/server/src/qwen35/qwen35_dflash_target.h @@ -44,7 +44,10 @@ class Qwen35DFlashTarget : public DFlashTarget { bool snapshot_kv() override; bool restore_kv() override; bool supports_fast_rollback() const override; + bool exact_fast_rollback() const override { return specla_active(); } + bool rollback_failure_is_recoverable() const override { return !specla_active(); } bool rollback_to(int base_pos, int commit_n) override; + bool finish_speculative_state() override; bool supports_tree_verify() const override; bool verify_tree(int committed, @@ -99,6 +102,19 @@ class Qwen35DFlashTarget : public DFlashTarget { KvFlashPager * pager_ = nullptr; bool fast_rollback_ = false; + // SpecLA (DFLASH_SPECLA=1, docs/SPECLA.md): true when the cache was + // migrated with factor buffers. Capture-verify then runs the + // topology-masked factor path, never mutates durable SSM/conv state + // (snapshot/restore become no-ops), and rollback commits via + // DeltaConstruct instead of dense checkpoint copies. + bool specla_active() const { + return fast_rollback_ && pager_ == nullptr && !cache_.factor_k.empty(); + } + + // SpecLA chain commit: DeltaConstruct over the accepted prefix plus a + // fused shift/append of raw convolution factors. + bool rollback_to_specla(int base_pos, int commit_n); + // Cached vector form of capture layer IDs (built once in constructor). std::vector capture_ids_; diff --git a/server/src/qwen35/qwen35_layer_split_dflash_target.h b/server/src/qwen35/qwen35_layer_split_dflash_target.h index b12a1dfed..15ebd6b2c 100644 --- a/server/src/qwen35/qwen35_layer_split_dflash_target.h +++ b/server/src/qwen35/qwen35_layer_split_dflash_target.h @@ -61,6 +61,11 @@ class Qwen35LayerSplitDFlashTarget : public DFlashTarget { bool rollback_to_tree(int committed, const DDTree & tree, const std::vector & accepted_dfs) override; + bool rollback_failure_is_recoverable() const override { + // A context-fatal CUDA error poisons the context: restore+replay is + // not safe to attempt, so tell the spec-decode loop to abort. + return !last_rollback_context_fatal_; + } bool last_rollback_context_fatal() const { return last_rollback_context_fatal_; } bool is_eos(int token) const override; diff --git a/server/src/qwen35/qwen35_target_graph.cpp b/server/src/qwen35/qwen35_target_graph.cpp index 1082df6f5..e4615cbc2 100644 --- a/server/src/qwen35/qwen35_target_graph.cpp +++ b/server/src/qwen35/qwen35_target_graph.cpp @@ -32,10 +32,15 @@ #include "internal.h" #include "delta_net_chunked.h" +#include "delta_net_specla.h" #include "kv_quant.h" #include "qwen35_ops.h" #include "qwen35moe_ffn.h" #include "common/chain_rollback_policy.h" +#include "common/specla_commit_cuda.h" +#include "common/specla_mode.h" + +#include "ggml-alloc.h" #include #include @@ -371,6 +376,47 @@ bool create_target_cache_partial(const TargetWeights & w, return true; } +static void refresh_specla_state_ptrs(TargetCache & c) { + if (!c.specla_state_ptrs) return; + const int n_delta = (int)c.ssm_state.size(); + if (c.specla_state_ptrs->ne[0] != n_delta) return; + std::vector ptrs((size_t)n_delta, 0); + for (int dn = 0; dn < n_delta; dn++) { + if (!c.ssm_state[dn]) return; + ptrs[(size_t)dn] = (int64_t)(intptr_t)c.ssm_state[dn]->data; + } + ggml_backend_tensor_set(c.specla_state_ptrs, ptrs.data(), 0, + sizeof(int64_t) * ptrs.size()); + if (c.specla_conv_state_ptrs && + c.specla_conv_state_ptrs->ne[0] == n_delta && + (int)c.conv_state.size() == n_delta) { + for (int dn = 0; dn < n_delta; ++dn) { + if (!c.conv_state[dn]) return; + ptrs[(size_t)dn] = (int64_t)(intptr_t)c.conv_state[dn]->data; + } + ggml_backend_tensor_set(c.specla_conv_state_ptrs, ptrs.data(), 0, + sizeof(int64_t) * ptrs.size()); + } + if (c.specla_factor_ptrs && c.specla_factor_ptrs->ne[0] == 8 && + c.factor_k_all && c.factor_v_new_all && c.factor_g_ps_all && + c.conv_factor_all && c.factor_k_all_alt && + c.factor_v_new_all_alt && c.factor_g_ps_all_alt && + c.conv_factor_all_alt) { + const int64_t factor_ptrs[8] = { + (int64_t)(intptr_t)c.factor_k_all->data, + (int64_t)(intptr_t)c.factor_v_new_all->data, + (int64_t)(intptr_t)c.factor_g_ps_all->data, + (int64_t)(intptr_t)c.conv_factor_all->data, + (int64_t)(intptr_t)c.factor_k_all_alt->data, + (int64_t)(intptr_t)c.factor_v_new_all_alt->data, + (int64_t)(intptr_t)c.factor_g_ps_all_alt->data, + (int64_t)(intptr_t)c.conv_factor_all_alt->data, + }; + ggml_backend_tensor_set(c.specla_factor_ptrs, factor_ptrs, 0, + sizeof(factor_ptrs)); + } +} + void free_target_cache(TargetCache & c) { if (c.base_buf) { ggml_backend_buffer_free(c.base_buf); c.base_buf = nullptr; } if (c.base_ctx) { ggml_free(c.base_ctx); c.base_ctx = nullptr; } @@ -384,6 +430,27 @@ void free_target_cache(TargetCache & c) { c.conv_state_snap.clear(); c.ssm_intermediate.clear(); c.conv_input_cache.clear(); + c.conv_input_cache_alt.clear(); + c.factor_k.clear(); + c.factor_v_new.clear(); + c.factor_g_ps.clear(); + c.factor_k_alt.clear(); + c.factor_v_new_alt.clear(); + c.factor_g_ps_alt.clear(); + c.factor_k_all = nullptr; + c.factor_v_new_all = nullptr; + c.factor_g_ps_all = nullptr; + c.factor_k_all_alt = nullptr; + c.factor_v_new_all_alt = nullptr; + c.factor_g_ps_all_alt = nullptr; + c.conv_factor_all = nullptr; + c.conv_factor_all_alt = nullptr; + c.specla_idx = nullptr; + c.specla_state_ptrs = nullptr; + c.specla_conv_state_ptrs = nullptr; + c.specla_factor_ptrs = nullptr; + c.specla_pending_bank = 0; + c.specla_pending_count = 0; c.target_feat = nullptr; c.q_cap = nullptr; c.cur_pos = 0; @@ -391,10 +458,13 @@ void free_target_cache(TargetCache & c) { void reset_target_cache(TargetCache & c) { c.cur_pos = 0; + c.specla_pending_bank = 0; + c.specla_pending_count = 0; if (c.backend && ggml_backend_buft_is_meta( ggml_backend_get_default_buffer_type(c.backend))) { if (c.base_buf) ggml_backend_buffer_clear(c.base_buf, 0); if (c.rollback_buf) ggml_backend_buffer_clear(c.rollback_buf, 0); + refresh_specla_state_ptrs(c); return; } std::vector zeros(1 * 1024 * 1024, 0); @@ -413,9 +483,14 @@ void reset_target_cache(TargetCache & c) { } } } + // reset_target_cache clears the whole rollback buffer, including this + // persistent device pointer table. Re-upload it for daemon request 2+. + refresh_specla_state_ptrs(c); } void reset_recurrent_state(TargetCache & c) { + c.specla_pending_bank = 0; + c.specla_pending_count = 0; // Device-side clear of the whole base buffer (KV + SSM + conv): with the // step-invariant decode the FA span is 256-padded and mask-less, so stale // K/V rows from the PREVIOUS request inside the padded tail would be @@ -459,7 +534,8 @@ bool migrate_prefill_cache(const TargetWeights & w, int max_ctx, int max_verify_tokens, ggml_backend_t backend, - TargetCache & cache) { + TargetCache & cache, + bool enable_specla) { // Already migrated (e.g. daemon mode second+ request after reset_target_cache). if (cache.rollback_ctx) return true; @@ -474,8 +550,22 @@ bool migrate_prefill_cache(const TargetWeights & w, cache.conv_state_snap.assign(n_delta, nullptr); cache.ssm_intermediate.assign(n_delta, nullptr); cache.conv_input_cache.assign(n_delta, nullptr); - - const int rb_tensors = 4 * n_delta; + cache.conv_input_cache_alt.assign(n_delta, nullptr); + + // SpecLA replaces the dense per-token state checkpoints with compact + // per-token factor buffers (~(S_k+S_v+1)·H_v·max_q·4B per layer vs + // state_size·max_q per layer) — the paper's §5.1 memory trade. + const bool specla = enable_specla && specla_enabled(); + if (specla) { + cache.factor_k.assign(n_delta, nullptr); + cache.factor_v_new.assign(n_delta, nullptr); + cache.factor_g_ps.assign(n_delta, nullptr); + cache.factor_k_alt.assign(n_delta, nullptr); + cache.factor_v_new_alt.assign(n_delta, nullptr); + cache.factor_g_ps_alt.assign(n_delta, nullptr); + } + + const int rb_tensors = (specla ? 12 : 4) * n_delta + (specla ? 1 : 0); ggml_init_params ip{}; ip.mem_size = (size_t)(rb_tensors + 16) * ggml_tensor_overhead(); ip.mem_buffer = nullptr; @@ -489,6 +579,47 @@ bool migrate_prefill_cache(const TargetWeights & w, const ggml_type checkpoint_type = rollback_policy.checkpoint_f32 ? GGML_TYPE_F32 : GGML_TYPE_F16; + const int head_k_dim = w.ssm_d_state; + if (specla) { + // Consolidated double-buffered factors. Token-major layout lets one + // compaction kernel gather an arbitrary accepted tree path. + cache.factor_k_all = ggml_new_tensor_4d(cache.rollback_ctx, GGML_TYPE_F32, + head_k_dim, w.ssm_dt_rank, n_delta, max_verify_tokens); + cache.factor_v_new_all = ggml_new_tensor_4d(cache.rollback_ctx, GGML_TYPE_F32, + head_v_dim, w.ssm_dt_rank, n_delta, max_verify_tokens); + cache.factor_g_ps_all = ggml_new_tensor_3d(cache.rollback_ctx, GGML_TYPE_F32, + w.ssm_dt_rank, n_delta, max_verify_tokens); + cache.factor_k_all_alt = ggml_new_tensor_4d(cache.rollback_ctx, GGML_TYPE_F32, + head_k_dim, w.ssm_dt_rank, n_delta, max_verify_tokens); + cache.factor_v_new_all_alt = ggml_new_tensor_4d(cache.rollback_ctx, GGML_TYPE_F32, + head_v_dim, w.ssm_dt_rank, n_delta, max_verify_tokens); + cache.factor_g_ps_all_alt = ggml_new_tensor_3d(cache.rollback_ctx, GGML_TYPE_F32, + w.ssm_dt_rank, n_delta, max_verify_tokens); + cache.conv_factor_all = ggml_new_tensor_3d(cache.rollback_ctx, GGML_TYPE_F32, + conv_ch, n_delta, max_verify_tokens); + cache.conv_factor_all_alt = ggml_new_tensor_3d(cache.rollback_ctx, GGML_TYPE_F32, + conv_ch, n_delta, max_verify_tokens); + ggml_set_name(cache.factor_k_all, "specla_factor_k_all"); + ggml_set_name(cache.factor_v_new_all, "specla_factor_v_all"); + ggml_set_name(cache.factor_g_ps_all, "specla_factor_g_all"); + ggml_set_name(cache.factor_k_all_alt, "specla_factor_k_all_alt"); + ggml_set_name(cache.factor_v_new_all_alt, "specla_factor_v_all_alt"); + ggml_set_name(cache.factor_g_ps_all_alt, "specla_factor_g_all_alt"); + ggml_set_name(cache.conv_factor_all, "specla_conv_factor_all"); + ggml_set_name(cache.conv_factor_all_alt, "specla_conv_factor_all_alt"); + cache.specla_idx = ggml_new_tensor_1d(cache.rollback_ctx, GGML_TYPE_I32, + max_verify_tokens); + cache.specla_state_ptrs = ggml_new_tensor_1d(cache.rollback_ctx, GGML_TYPE_I64, + n_delta); + cache.specla_conv_state_ptrs = ggml_new_tensor_1d( + cache.rollback_ctx, GGML_TYPE_I64, n_delta); + cache.specla_factor_ptrs = ggml_new_tensor_1d( + cache.rollback_ctx, GGML_TYPE_I64, 8); + ggml_set_name(cache.specla_idx, "specla_idx"); + ggml_set_name(cache.specla_state_ptrs, "specla_state_ptrs"); + ggml_set_name(cache.specla_conv_state_ptrs, "specla_conv_state_ptrs"); + ggml_set_name(cache.specla_factor_ptrs, "specla_factor_ptrs"); + } int dn_idx = 0; for (int il = 0; il < w.n_layer; il++) { if (((il + 1) % w.full_attention_interval) != 0) { @@ -496,21 +627,33 @@ bool migrate_prefill_cache(const TargetWeights & w, head_v_dim, head_v_dim, w.ssm_dt_rank); ggml_tensor * Cn = ggml_new_tensor_2d(cache.rollback_ctx, GGML_TYPE_F32, w.ssm_d_conv - 1, conv_ch); - ggml_tensor * Si = ggml_new_tensor_4d(cache.rollback_ctx, checkpoint_type, - head_v_dim, head_v_dim, - w.ssm_dt_rank, max_verify_tokens); - ggml_tensor * Ci = ggml_new_tensor_3d(cache.rollback_ctx, GGML_TYPE_F32, - (w.ssm_d_conv - 1) + max_verify_tokens, - conv_ch, 1); + ggml_tensor * Ci = specla ? nullptr + : ggml_new_tensor_3d(cache.rollback_ctx, GGML_TYPE_F32, + (w.ssm_d_conv - 1) + max_verify_tokens, + conv_ch, 1); + ggml_tensor * Ci_alt = nullptr; char name[64]; std::snprintf(name, sizeof(name), "ssm_state_snap_%d", il); ggml_set_name(Sn, name); std::snprintf(name, sizeof(name), "conv_state_snap_%d", il); ggml_set_name(Cn, name); - std::snprintf(name, sizeof(name), "ssm_intermediate_%d", il); ggml_set_name(Si, name); - std::snprintf(name, sizeof(name), "conv_input_cache_%d", il); ggml_set_name(Ci, name); + if (Ci) { + std::snprintf(name, sizeof(name), "conv_input_cache_%d", il); + ggml_set_name(Ci, name); + } cache.ssm_state_snap[dn_idx] = Sn; cache.conv_state_snap[dn_idx] = Cn; - cache.ssm_intermediate[dn_idx] = Si; cache.conv_input_cache[dn_idx] = Ci; + cache.conv_input_cache_alt[dn_idx] = Ci_alt; + if (Ci_alt) { + std::snprintf(name, sizeof(name), "conv_input_cache_alt_%d", il); + ggml_set_name(Ci_alt, name); + } + if (!specla) { + ggml_tensor * Si = ggml_new_tensor_4d(cache.rollback_ctx, checkpoint_type, + head_v_dim, head_v_dim, + w.ssm_dt_rank, max_verify_tokens); + std::snprintf(name, sizeof(name), "ssm_intermediate_%d", il); ggml_set_name(Si, name); + cache.ssm_intermediate[dn_idx] = Si; + } dn_idx++; } } @@ -553,6 +696,235 @@ bool migrate_prefill_cache(const TargetWeights & w, } } + // SpecLA: per-layer factor views into the consolidated buffers, created + // after allocation so they carry live data/buffer pointers. Shaped like + // stand-alone per-layer tensors ([.., max_q] with a cross-layer token + // stride) so the capture path treats them like any other cache tensor. + if (specla) { + ggml_tensor * Fk = cache.factor_k_all; + ggml_tensor * Fv = cache.factor_v_new_all; + ggml_tensor * Fg = cache.factor_g_ps_all; + ggml_tensor * Fk_alt = cache.factor_k_all_alt; + ggml_tensor * Fv_alt = cache.factor_v_new_all_alt; + ggml_tensor * Fg_alt = cache.factor_g_ps_all_alt; + for (int dn = 0; dn < n_delta; dn++) { + cache.factor_k[dn] = ggml_view_3d(cache.rollback_ctx, Fk, + Fk->ne[0], Fk->ne[1], max_verify_tokens, + Fk->nb[1], Fk->nb[3], (size_t)dn * Fk->nb[2]); + cache.factor_v_new[dn] = ggml_view_3d(cache.rollback_ctx, Fv, + Fv->ne[0], Fv->ne[1], max_verify_tokens, + Fv->nb[1], Fv->nb[3], (size_t)dn * Fv->nb[2]); + cache.factor_g_ps[dn] = ggml_view_2d(cache.rollback_ctx, Fg, + Fg->ne[0], max_verify_tokens, + Fg->nb[2], (size_t)dn * Fg->nb[1]); + cache.factor_k_alt[dn] = ggml_view_3d(cache.rollback_ctx, Fk_alt, + Fk_alt->ne[0], Fk_alt->ne[1], max_verify_tokens, + Fk_alt->nb[1], Fk_alt->nb[3], (size_t)dn * Fk_alt->nb[2]); + cache.factor_v_new_alt[dn] = ggml_view_3d(cache.rollback_ctx, Fv_alt, + Fv_alt->ne[0], Fv_alt->ne[1], max_verify_tokens, + Fv_alt->nb[1], Fv_alt->nb[3], (size_t)dn * Fv_alt->nb[2]); + cache.factor_g_ps_alt[dn] = ggml_view_2d(cache.rollback_ctx, Fg_alt, + Fg_alt->ne[0], max_verify_tokens, + Fg_alt->nb[2], (size_t)dn * Fg_alt->nb[1]); + cache.conv_input_cache[dn] = ggml_view_3d(cache.rollback_ctx, + cache.conv_factor_all, cache.conv_factor_all->ne[0], + max_verify_tokens, 1, cache.conv_factor_all->nb[2], + cache.conv_factor_all->nb[2]*max_verify_tokens, + (size_t)dn*cache.conv_factor_all->nb[1]); + cache.conv_input_cache_alt[dn] = ggml_view_3d(cache.rollback_ctx, + cache.conv_factor_all_alt, cache.conv_factor_all_alt->ne[0], + max_verify_tokens, 1, cache.conv_factor_all_alt->nb[2], + cache.conv_factor_all_alt->nb[2]*max_verify_tokens, + (size_t)dn*cache.conv_factor_all_alt->nb[1]); + } + // State addresses are stable for the cache's lifetime. The same helper + // is also called after reset_target_cache clears rollback scratch. + refresh_specla_state_ptrs(cache); + } + + return true; +} + +// SpecLA DeltaConstruct commit (docs/SPECLA.md): one small graph advancing all +// delta-net layers' durable SSM states along the accepted path, +// S_A = exp(g⁺_A) S0 + Σ_{t∈path} exp(g⁺_A − g⁺_t) k_t ⊗ ṽ_t, +// from the factor buffers the last verify captured. The verify graph never +// wrote the speculative state back, so cache.ssm_state still holds S0 here. +bool specla_commit_accepted(TargetCache & cache, + ggml_backend_t backend, + const int32_t * accepted_idx, + int A) { + const int n_delta = (int)cache.factor_k.size(); + if (n_delta == 0 || A <= 0 || !accepted_idx || !backend) return false; + + // The just-run factorized verify wrote the bank opposite the one its + // capture setup treated as pending (the capture views point at + // factor_k_alt when specla_pending_bank == 0). This function is called + // BEFORE the host-side bank rotation, so commit that opposite bank; + // committing bank 0 unconditionally replays stale factors on the first + // verify. + const int selected_bank = 1 - cache.specla_pending_bank; + const bool alt = selected_bank != 0; + ggml_tensor * Fk = alt ? cache.factor_k_all_alt : cache.factor_k_all; + ggml_tensor * Fv = alt ? cache.factor_v_new_all_alt : cache.factor_v_new_all; + ggml_tensor * Fg = alt ? cache.factor_g_ps_all_alt : cache.factor_g_ps_all; + ggml_tensor * Fc = alt ? cache.conv_factor_all_alt : cache.conv_factor_all; + if (!Fk || !Fv || !Fg || !Fc || !cache.specla_conv_state_ptrs || + !cache.specla_conv_state_ptrs->data) { + return false; + } + + for (int i = 0; i < A; i++) { + if (accepted_idx[i] < 0 || accepted_idx[i] >= Fk->ne[3]) return false; + // The factorized fallback's convolution capture is in flat-token + // order. Tree factorized verify is built with an HLD schedule in + // production; fail closed rather than commit a scattered DFS path. + if (accepted_idx[i] != i) return false; + } + + // Pre-validate the convolution commit targets before mutating SSM state. + // The selected consolidated bank is [channels, layers, tokens]. Commit + // shifts each durable K-1 window and appends raw inputs [0, A). + if (Fc->type != GGML_TYPE_F32 || !ggml_is_contiguous(Fc) || + Fc->ne[1] != n_delta || A > Fc->ne[2]) { + return false; + } + int d_conv = 0; + for (int il = 0; il < n_delta; il++) { + ggml_tensor * dst = (il < (int)cache.conv_state.size()) + ? cache.conv_state[il] : nullptr; + if (!dst || dst->type != GGML_TYPE_F32 || !ggml_is_contiguous(dst) || + dst->ne[1] != Fc->ne[0]) { + return false; + } + const int layer_d_conv = (int)dst->ne[0] + 1; + if (layer_d_conv < 2 || (d_conv != 0 && d_conv != layer_d_conv)) { + return false; + } + d_conv = layer_d_conv; + } + + const int64_t S_k = Fk->ne[0]; + const int64_t H = Fk->ne[1]; + const int64_t S_v = Fv->ne[0]; + const int64_t HL = H * n_delta; + + // Fast path: one fused kernel updates every layer's state in place. + // Escape hatch DFLASH_SPECLA_FUSED_COMMIT=0 falls back to the ggml-graph + // implementation below (also the fallback on any launch failure). + static const bool kFusedCommit = []() { + const char * v = std::getenv("DFLASH_SPECLA_FUSED_COMMIT"); + return v == nullptr || v[0] != '0'; + }(); + bool ok = false; + if (kFusedCommit && cache.specla_idx && cache.specla_state_ptrs) { + ggml_backend_tensor_set(cache.specla_idx, accepted_idx, 0, + sizeof(int32_t) * (size_t)A); + bool launched = false; + if (specla_commit_fused( + (float * const *)cache.specla_state_ptrs->data, + (const float *)Fk->data, + (const float *)Fv->data, + (const float *)Fg->data, + (const int32_t *)cache.specla_idx->data, + A, (int)S_k, (int)S_v, (int)H, n_delta, + /*stream=*/nullptr, &launched)) { + ok = true; + } else if (launched) { + std::fprintf(stderr, + "specla_commit_accepted: fused kernel execution failed\n"); + return false; + } else { + std::fprintf(stderr, + "specla_commit_accepted: fused launch rejected; using graph path\n"); + } + } + + if (!ok) { + // Persistent metadata arena + allocator, reused across steps like the + // verify step graph — avoids per-commit gallocr churn. + static thread_local std::vector s_arena; + static thread_local ggml_gallocr_t s_galloc = nullptr; + const size_t graph_nodes = (size_t)n_delta * 8 + 64; + ggml_init_params ip{}; + ip.mem_size = graph_nodes * 4 * ggml_tensor_overhead() + + ggml_graph_overhead_custom(graph_nodes, false); + if (s_arena.size() < ip.mem_size) s_arena.resize(ip.mem_size); + ip.mem_buffer = s_arena.data(); + ip.no_alloc = true; + ggml_context * ctx = ggml_init(ip); + if (!ctx) return false; + ggml_cgraph * gf = ggml_new_graph_custom(ctx, graph_nodes, false); + + ggml_tensor * idx = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, A); + ggml_set_input(idx); + + // Gather the accepted token slices across ALL layers at once (the token + // axis is outermost in the consolidated buffers). + ggml_tensor * k_sel = ggml_get_rows(ctx, + ggml_reshape_2d(ctx, Fk, S_k * HL, Fk->ne[3]), idx); // [S_k*HL, A] + ggml_tensor * v_sel = ggml_get_rows(ctx, + ggml_reshape_2d(ctx, Fv, S_v * HL, Fv->ne[3]), idx); // [S_v*HL, A] + ggml_tensor * g_sel = ggml_get_rows(ctx, + ggml_reshape_2d(ctx, Fg, HL, Fg->ne[2]), idx); // [HL, A] + k_sel = ggml_reshape_3d(ctx, k_sel, S_k, HL, A); + v_sel = ggml_reshape_3d(ctx, v_sel, S_v, HL, A); + + // w[hl, t] = exp(g⁺_A − g⁺_t); the deepest accepted node is last. + ggml_tensor * gA = ggml_view_2d(ctx, g_sel, HL, 1, g_sel->nb[1], + (size_t)(A - 1) * g_sel->nb[1]); + ggml_tensor * w_dec = ggml_exp(ctx, ggml_neg(ctx, ggml_sub(ctx, g_sel, gA))); + + // Σ_t (k_t · w_t) ⊗ ṽ_t for every (layer, head) in one batched matmul. + ggml_tensor * kg = ggml_mul(ctx, k_sel, ggml_reshape_3d(ctx, w_dec, 1, HL, A)); + ggml_tensor * kg_t = ggml_cont(ctx, ggml_permute(ctx, kg, 1, 2, 0, 3)); // [A, S_k, HL] + ggml_tensor * v_t = ggml_cont(ctx, ggml_permute(ctx, v_sel, 1, 2, 0, 3)); // [A, S_v, HL] + ggml_tensor * upd = ggml_mul_mat(ctx, kg_t, v_t); // [S_k, S_v, HL] + ggml_tensor * gA_exp = ggml_exp(ctx, ggml_cont(ctx, gA)); // [HL, 1] + + // Per-layer tail: S ← exp(g⁺_A)·S + upd (states are separate tensors). + for (int il = 0; il < n_delta; il++) { + ggml_tensor * S = cache.ssm_state[il]; + if (!S) { ggml_free(ctx); return false; } + ggml_tensor * upd_l = ggml_view_3d(ctx, upd, S_k, S_v, H, + upd->nb[1], upd->nb[2], (size_t)il * H * upd->nb[2]); + ggml_tensor * gA_l = ggml_reshape_3d(ctx, + ggml_cont(ctx, ggml_view_1d(ctx, gA_exp, H, (size_t)il * H * gA_exp->nb[0])), + 1, 1, H); + ggml_tensor * s_new = ggml_add(ctx, ggml_mul(ctx, S, gA_l), upd_l); + ggml_build_forward_expand(gf, ggml_cpy(ctx, s_new, S)); + } + + if (!s_galloc) { + s_galloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); + } + ok = s_galloc != nullptr && ggml_gallocr_alloc_graph(s_galloc, gf); + if (ok) { + ggml_backend_tensor_set(idx, accepted_idx, 0, sizeof(int32_t) * A); + ok = ggml_backend_graph_compute(backend, gf) == GGML_STATUS_SUCCESS; + } + ggml_free(ctx); + if (ok) ggml_backend_synchronize(backend); + } + if (!ok) return false; + + // Apply the accepted raw convolution factors across all layers. The + // token-major kernel handles partial and full acceptance without needing + // a synthetic K-1 prefix in the factor bank. + if (!specla_commit_conv_raw_fused( + (float * const *)cache.specla_conv_state_ptrs->data, + (const float *)Fc->data, A, (int)Fc->ne[2], n_delta, + (int)Fc->ne[0], d_conv, /*stream=*/nullptr)) { + std::fprintf(stderr, + "specla_commit_accepted: raw conv factor commit failed\n"); + return false; + } + + // The selected bank has been consumed into durable state. Mark it as the + // pending role with nothing outstanding so the next verify writes the + // opposite bank and the final flush is a no-op. + cache.specla_pending_bank = selected_bank; + cache.specla_pending_count = 0; return true; } @@ -1015,7 +1387,19 @@ static ggml_tensor * build_delta_net_block( int n_prefill_segments = 0, ggml_tensor * active_slot_ids = nullptr, ggml_tensor * state_slot_ids = nullptr, - bool allow_inplace_state = false + bool allow_inplace_state = false, + // SpecLA topology masks (all three non-null together): route the + // recurrence through the topology-masked factor-capture verify. + // parent_ids then only steers the tree conv; the recurrence gets its + // topology from the masks. See docs/SPECLA.md. + ggml_tensor * specla_m_strict = nullptr, + ggml_tensor * specla_m_incl = nullptr, + ggml_tensor * specla_m_eye = nullptr, + ggml_tensor * specla_hld = nullptr, + int specla_n_boundaries = 0, + int specla_n_chains = 0, + int specla_n_waves = 0, + int specla_max_parallel_chains = 0 ) { const int head_k_dim = w.ssm_d_state; const int num_k_heads = w.ssm_n_group; @@ -1040,6 +1424,18 @@ static ggml_tensor * build_delta_net_block( } GGML_ASSERT(!ragged || (!cap && !parent_ids)); const bool can_skip_gdn_intermediate = skip_gdn_intermediate && !parent_ids && !cap; + // Fully factorized SpecLA fallback: the current candidates do not mutate + // durable state. The HLD route below may materialize the *previously* + // accepted pending path while keeping current candidates speculative. + const bool use_specla_factorized = cap && cap->factor_k && cap->factor_v_new && + cap->factor_g_ps && specla_m_strict && specla_m_incl && specla_m_eye; + const bool use_specla_hld = cap && cap->factor_ptrs && + cap->factor_n_layers > 0 && cap->factor_layer >= 0 && + cap->factor_layer < cap->factor_n_layers && specla_hld && + specla_n_chains > 0 && specla_n_waves > 0 && + specla_max_parallel_chains > 0; + GGML_ASSERT(!(use_specla_factorized || use_specla_hld) || + (n_seqs == 1 && !ragged && !active_slot_ids)); // ── Whole-batch projections ───────────────────────────────────── // qkv_mixed = wqkv @ cur [10240, n_tokens] @@ -1129,6 +1525,22 @@ static ggml_tensor * build_delta_net_block( seg_cols(g_2d, seg.off, seg_tokens), 1, num_v_heads, n_seq_tokens, seg_seqs); + ggml_tensor * conv_out = nullptr; + if (use_specla_hld) { + // Delayed convolution commit + HLD verify. The result is already + // SiLU'd; raw inputs are written directly to the persistent bank. + ggml_tensor * packed = ggml_ssm_conv_specla( + ctx, qkv_mixed, L.ssm_conv1d, seg.conv_st, + specla_hld, cap->factor_ptrs, cap->factor_n_layers, + cap->factor_layer, cap->pending_bank, + specla_n_boundaries, specla_n_chains, specla_n_waves, + specla_max_parallel_chains); + const size_t f32 = sizeof(float); + conv_out = ggml_view_3d(ctx, packed, + conv_channels, n_seq_tokens, 1, + (size_t)conv_channels*f32, + (size_t)conv_channels*n_seq_tokens*f32, 0); + } else { // ── Fetch conv state [kernel-1, conv_channels] and prepend to qkv_mixed // along the token axis to form the convolution input. ggml_tensor * conv_states_r = nullptr; @@ -1162,39 +1574,56 @@ static ggml_tensor * build_delta_net_block( // past graph_compute). After graph_compute, the cache buffer's data is // always valid; the rollback code slices it at commit_n. if (cap && cap->conv_input) { - // conv_input may be shorter than the pre-allocated cache - // (e.g. during prefill when n_tokens < max_verify_tokens). - // Copy into a matching-sized view of the cache destination. - const int64_t ci_len = conv_input->ne[0]; - ggml_tensor * dst; - if (ci_len == cap->conv_input->ne[0]) { - dst = cap->conv_input; - } else { - dst = ggml_view_3d(ctx, cap->conv_input, - ci_len, cap->conv_input->ne[1], cap->conv_input->ne[2], + if (use_specla_factorized) { + // The consolidated SpecLA bank is [channels, layers, tokens]. + // Capture only the raw current inputs; the compatibility commit + // shifts the durable K-1 window and appends accepted tokens. + GGML_ASSERT(qkv_mixed->ne[0] == cap->conv_input->ne[0]); + GGML_ASSERT(n_seq_tokens <= cap->conv_input->ne[1]); + ggml_tensor * dst = ggml_view_3d(ctx, cap->conv_input, + cap->conv_input->ne[0], n_seq_tokens, 1, cap->conv_input->nb[1], cap->conv_input->nb[2], 0); + GGML_ASSERT(ggml_nelements(qkv_mixed) == ggml_nelements(dst)); + ggml_build_forward_expand(gf, ggml_cpy(ctx, qkv_mixed, dst)); + } else { + // conv_input may be shorter than the pre-allocated cache + // (e.g. during prefill when n_tokens < max_verify_tokens). + // Copy into a matching-sized view of the cache destination. + const int64_t ci_len = conv_input->ne[0]; + ggml_tensor * dst; + if (ci_len == cap->conv_input->ne[0]) { + dst = cap->conv_input; + } else { + dst = ggml_view_3d(ctx, cap->conv_input, + ci_len, cap->conv_input->ne[1], cap->conv_input->ne[2], + cap->conv_input->nb[1], cap->conv_input->nb[2], 0); + } + GGML_ASSERT(ggml_nelements(conv_input) == ggml_nelements(dst)); + ggml_build_forward_expand(gf, ggml_cpy(ctx, conv_input, dst)); } - GGML_ASSERT(ggml_nelements(conv_input) == ggml_nelements(dst)); - ggml_build_forward_expand(gf, ggml_cpy(ctx, conv_input, dst)); } // ── Save the last (kernel-1) steps back to the conv state - ggml_tensor * last_conv = ggml_view_3d(ctx, conv_input, - w.ssm_d_conv - 1, conv_channels, seg_seqs, - conv_input->nb[1], conv_input->nb[2], - (conv_input->ne[0] - (w.ssm_d_conv - 1)) * ggml_element_size(conv_input)); - if (seg_active) { - const int64_t slab = - (int64_t)(w.ssm_d_conv - 1) * conv_channels; - ggml_tensor * compact_last = ggml_reshape_2d( - ctx, ggml_cont(ctx, last_conv), slab, seg_seqs); - ggml_tensor * all_conv = ggml_reshape_2d( - ctx, seg.conv_st, slab, seg.conv_st->ne[2]); - ggml_build_forward_expand( - gf, ggml_set_rows_masked( - ctx, all_conv, compact_last, active_slot_ids)); - } else { - ggml_build_forward_expand(gf, ggml_cpy(ctx, last_conv, seg.conv_st)); + // SpecLA: skipped — the window is speculative; the commit path + // shifts conv_state and appends the accepted raw inputs at commit time. + if (!use_specla_factorized) { + ggml_tensor * last_conv = ggml_view_3d(ctx, conv_input, + w.ssm_d_conv - 1, conv_channels, seg_seqs, + conv_input->nb[1], conv_input->nb[2], + (conv_input->ne[0] - (w.ssm_d_conv - 1)) * ggml_element_size(conv_input)); + if (seg_active) { + const int64_t slab = + (int64_t)(w.ssm_d_conv - 1) * conv_channels; + ggml_tensor * compact_last = ggml_reshape_2d( + ctx, ggml_cont(ctx, last_conv), slab, seg_seqs); + ggml_tensor * all_conv = ggml_reshape_2d( + ctx, seg.conv_st, slab, seg.conv_st->ne[2]); + ggml_build_forward_expand( + gf, ggml_set_rows_masked( + ctx, all_conv, compact_last, active_slot_ids)); + } else { + ggml_build_forward_expand(gf, ggml_cpy(ctx, last_conv, seg.conv_st)); + } } // ── 1D conv + silu @@ -1202,10 +1631,11 @@ static ggml_tensor * build_delta_net_block( // their conv window from their actual tree parent instead of the DFS // predecessor. Without this, siblings get garbage logits (the conv // output would mix unrelated branches). - ggml_tensor * conv_out = parent_ids + conv_out = parent_ids ? ggml_ssm_conv_tree(ctx, conv_input, L.ssm_conv1d, parent_ids) : ggml_ssm_conv (ctx, conv_input, L.ssm_conv1d); conv_out = ggml_silu(ctx, conv_out); + } // conv_out: [conv_channels, n_tokens, n_seqs] const int64_t q_offset = 0; @@ -1293,6 +1723,54 @@ static ggml_tensor * build_delta_net_block( ggml_tensor * output = nullptr; + if (use_specla_hld) { + ggml_tensor * result = ggml_gated_delta_net_specla( + ctx, q_c, k_c, v_c, g_tensor, beta, s, specla_hld, + cap->factor_ptrs, cap->factor_n_layers, cap->factor_layer, + cap->pending_bank, specla_n_boundaries, + specla_n_chains, specla_n_waves, specla_max_parallel_chains); + const int64_t S = head_v_dim; + const int64_t H = num_v_heads; + const int64_t T = n_seq_tokens; + const size_t f32 = sizeof(float); + const size_t factor_bytes = (size_t)S*H*T*f32; + output = ggml_view_4d(ctx, result, S, H, T, 1, + S*f32, (size_t)S*H*f32, factor_bytes, 0); + goto after_delta_net; + } + + if (use_specla_factorized) { + auto r = build_delta_net_specla(ctx, q_c, k_c, v_c, g_tensor, beta, s, + specla_m_strict, specla_m_incl, specla_m_eye); + output = r.output; + + // Factor capture: prefix views of the persistent per-layer buffers. + // k as fed to the recurrence (post l2-norm, post head repeat). + { + ggml_tensor * dst_k = ggml_view_4d(ctx, cap->factor_k, + cap->factor_k->ne[0], cap->factor_k->ne[1], n_seq_tokens, 1, + cap->factor_k->nb[1], cap->factor_k->nb[2], cap->factor_k->nb[2] * n_seq_tokens, 0); + ggml_build_forward_expand(gf, ggml_cpy(ctx, k_c, dst_k)); + + // ṽ [n, S_v, 1, H] → logical [S_v, H, n, 1] to match the buffer. + ggml_tensor * src_v = ggml_cont(ctx, ggml_permute(ctx, r.v_new, 2, 0, 3, 1)); + ggml_tensor * dst_v = ggml_view_4d(ctx, cap->factor_v_new, + cap->factor_v_new->ne[0], cap->factor_v_new->ne[1], n_seq_tokens, 1, + cap->factor_v_new->nb[1], cap->factor_v_new->nb[2], cap->factor_v_new->nb[2] * n_seq_tokens, 0); + ggml_build_forward_expand(gf, ggml_cpy(ctx, src_v, dst_v)); + + // g⁺ [n, 1, 1, H] → logical [H, n]. + ggml_tensor * src_g = ggml_cont(ctx, ggml_permute(ctx, r.g_ps, 1, 2, 3, 0)); + ggml_tensor * dst_g = ggml_view_2d(ctx, cap->factor_g_ps, + cap->factor_g_ps->ne[0], n_seq_tokens, cap->factor_g_ps->nb[1], 0); + ggml_build_forward_expand(gf, ggml_cpy(ctx, src_g, dst_g)); + } + // Compatibility fallback for callers without an HLD schedule. No + // new_state and no state writeback: verification is read-only on + // the durable state; specla_commit_accepted() advances it. + goto after_delta_net; + } + if (use_chunked) { auto r = build_delta_net_chunked(ctx, q_c, k_c, v_c, g_tensor, beta, s); output = r.output; @@ -1377,6 +1855,7 @@ static ggml_tensor * build_delta_net_block( } } +after_delta_net: // ── Gated output norm: rms_norm(output) * silu(z_4d) ggml_tensor * z_4d = ggml_reshape_4d(ctx, seg_cols(z, seg.off, seg_tokens), @@ -1612,6 +2091,29 @@ QwenGraphOutputs build_qwen35_graph( // valid because they're cache-resident, not gallocr-managed. cap_ptr->ssm_intermediate_states = cache.ssm_intermediate[dn_idx]; cap_ptr->conv_input = cache.conv_input_cache[dn_idx]; + if (!cache.factor_k.empty()) { + const bool pending_alt = cache.specla_pending_bank != 0; + cap_ptr->pending_factor_k = pending_alt + ? cache.factor_k_alt[dn_idx] : cache.factor_k[dn_idx]; + cap_ptr->pending_factor_v_new = pending_alt + ? cache.factor_v_new_alt[dn_idx] : cache.factor_v_new[dn_idx]; + cap_ptr->pending_factor_g = pending_alt + ? cache.factor_g_ps_alt[dn_idx] : cache.factor_g_ps[dn_idx]; + cap_ptr->pending_conv_input = pending_alt + ? cache.conv_input_cache_alt[dn_idx] : cache.conv_input_cache[dn_idx]; + cap_ptr->factor_k = pending_alt + ? cache.factor_k[dn_idx] : cache.factor_k_alt[dn_idx]; + cap_ptr->factor_v_new = pending_alt + ? cache.factor_v_new[dn_idx] : cache.factor_v_new_alt[dn_idx]; + cap_ptr->factor_g_ps = pending_alt + ? cache.factor_g_ps[dn_idx] : cache.factor_g_ps_alt[dn_idx]; + cap_ptr->conv_input = pending_alt + ? cache.conv_input_cache[dn_idx] : cache.conv_input_cache_alt[dn_idx]; + cap_ptr->factor_ptrs = cache.specla_factor_ptrs; + cap_ptr->factor_n_layers = (int)cache.factor_k.size(); + cap_ptr->factor_layer = dn_idx; + cap_ptr->pending_bank = cache.specla_pending_bank; + } } ggml_tensor * conv_st = cache.conv_state[dn_idx]; ggml_tensor * ssm_st = cache.ssm_state[dn_idx]; @@ -1643,7 +2145,13 @@ QwenGraphOutputs build_qwen35_graph( in.active_slot_ids, in.state_slot_ids, /*allow_inplace_state=*/ - in.n_prefill_tokens == 0); + in.n_prefill_tokens == 0, + in.specla_m_strict, in.specla_m_incl, + in.specla_m_eye, in.specla_hld, + in.specla_n_boundaries, + in.specla_n_chains, + in.specla_n_waves, + in.specla_max_parallel_chains); dn_idx++; } diff --git a/server/src/server/server_main.cpp b/server/src/server/server_main.cpp index 7373aace4..3f79799aa 100644 --- a/server/src/server/server_main.cpp +++ b/server/src/server/server_main.cpp @@ -22,12 +22,15 @@ #include "common/moe_hybrid_routing_stats.h" #include "common/platform_env.h" #include "common/peer_access.h" +#include "common/specla_mode.h" #include "placement/pflash_placement.h" #include "placement/draft_residency.h" #include "kvflash_pager.h" #include +#include #include +#include #include #include #include @@ -120,8 +123,13 @@ static void print_usage(const char * prog) { " --prefill-cache-slots Full prompt/prefill cache slots (default: 0)\n" " --fast-rollback Enable speculative fast rollback (default: on)\n" " --no-fast-rollback Disable speculative fast rollback, even with --ddtree\n" + " --specla Enable speculative linear-attention verification\n" + " when supported (Qwen3.6 uses DDTree automatically)\n" + " --specla-top-k SpecLA draft-tree width (default: 4)\n" " --ddtree Enable DDTree speculative decode\n" " --ddtree-budget DDTree budget (default: 22)\n" + " --ddtree-tau Confidence margin on cumulative log-prob\n" + " (default: 6 with --specla; otherwise off)\n" " --verify-width laguna chain spec verify width (default: base 8,\n" " trimmed per step by drafter confidence; N = fixed base)\n" " --adaptive-experts [tau] MoE expert-count gating on verify batches\n" @@ -236,6 +244,9 @@ int main(int argc, char ** argv) { bool fast_rollback_forced_off = false; bool target_split_fast_rollback_cli = false; bool adaptive_experts_set = false; // --adaptive-experts (MoE architectures only) + bool ddtree_tau_set = false; + bool specla_top_k_set = false; + int specla_top_k = 4; // Track which thinking-budget tunables the operator set via CLI. // Those values win over the model card (spec §3.1: "Explicit CLI @@ -407,11 +418,38 @@ int main(int argc, char ** argv) { sconfig.prefill_cache_cap = std::atoi(argv[++i]); } else if (std::strcmp(argv[i], "--fast-rollback") == 0) { bargs.fast_rollback = true; + } else if (std::strcmp(argv[i], "--specla") == 0) { + bargs.specla_mode = true; + bargs.fast_rollback = true; + } else if (std::strcmp(argv[i], "--specla-top-k") == 0 && i + 1 < argc) { + const char * value = argv[++i]; + const char * end = value + std::strlen(value); + const auto parsed = std::from_chars(value, end, specla_top_k); + if (parsed.ec != std::errc{} || parsed.ptr != end || specla_top_k <= 0) { + std::fprintf(stderr, + "--specla-top-k expects a positive integer, got '%s'\n", value); + return 2; + } + specla_top_k_set = true; } else if (std::strcmp(argv[i], "--ddtree") == 0) { bargs.ddtree_mode = true; bargs.fast_rollback = true; } else if (std::strcmp(argv[i], "--ddtree-budget") == 0 && i + 1 < argc) { bargs.ddtree_budget = std::atoi(argv[++i]); + } else if (std::strcmp(argv[i], "--ddtree-tau") == 0 && i + 1 < argc) { + const char * value = argv[++i]; + char * end = nullptr; + errno = 0; + const float tau = std::strtof(value, &end); + if (errno == ERANGE || end == value || *end != '\0' || + !std::isfinite(tau) || tau < 0.0f) { + std::fprintf(stderr, + "--ddtree-tau expects a non-negative finite number, got '%s'\n", + value); + return 2; + } + bargs.ddtree_tau = tau; + ddtree_tau_set = true; } else if (std::strcmp(argv[i], "--adaptive-experts") == 0) { const char * tau = "0.80"; if (i + 1 < argc && argv[i + 1][0] != '-') { @@ -616,6 +654,19 @@ int main(int argc, char ** argv) { return 2; } } + if (specla_top_k_set && !bargs.specla_mode) { + std::fprintf(stderr, "[server] --specla-top-k requires --specla\n"); + return 2; + } + if (bargs.specla_mode && fast_rollback_forced_off) { + std::fprintf(stderr, + "[server] --specla is incompatible with --no-fast-rollback\n"); + return 2; + } + if (bargs.specla_mode && !ddtree_tau_set) { + bargs.ddtree_tau = 6.0f; + } + if (fast_rollback_forced_off) { bargs.fast_rollback = false; target_split_fast_rollback_cli = false; @@ -688,6 +739,8 @@ int main(int argc, char ** argv) { } const ResolvedBackendPlan & backend_plan = backend_preparation.plan; const std::string & arch = backend_plan.arch(); + const bool kvflash_requested = + kvflash_pool_requested(std::getenv("DFLASH_KVFLASH")); if (target_split_fast_rollback_cli && arch != "qwen35") { std::fprintf(stderr, "[server] --target-split-fast-rollback is only supported for " @@ -695,6 +748,49 @@ int main(int argc, char ** argv) { return 2; } + // SpecLA is the verification mode, not a proposal algorithm. Select the + // proposal adapter supported by this model. Qwen3.6 currently has a + // DDTree adapter; future model families may select DSpark here instead. + if (bargs.specla_mode) { + const bool supported = arch == "qwen35" && !bargs.device.is_multi_device(); + if (supported) { + if (!bargs.draft_path) { + std::fprintf(stderr, + "[server] Qwen3.6 SpecLA requires --draft \n"); + return 2; + } + bargs.ddtree_mode = true; + if (kvflash_requested) { + // KVFlash installs a paged attention-KV layout and therefore + // disables the factor-cache migration required by SpecLA. + // Keep the compatible proposal adapter, but report the + // effective verification mode accurately. + std::fprintf(stderr, + "[server] warning: --specla is unavailable with KVFlash; " + "using ordinary DDTree verification\n"); + bargs.specla_mode = false; + unset_environment_variable("DFLASH_SPECLA"); + } else { + set_environment_variable("DFLASH_SPECLA", "1", true); + if (specla_top_k_set) { + set_environment_variable( + "DFLASH_SPECLA_TOPK", std::to_string(specla_top_k).c_str(), true); + } else { + specla_top_k = specla_tree_topk(); + } + } + } else { + std::fprintf(stderr, + "[server] warning: --specla is unavailable for architecture '%s' " + "with placement %s; using the architecture's normal decode path\n", + arch.c_str(), placement_device_name(bargs.device).c_str()); + bargs.specla_mode = false; + if (!ddtree_tau_set) { + bargs.ddtree_tau = std::numeric_limits::infinity(); + } + } + } + // Paged decode owns its K/V through a block table that the snapshot format // cannot describe yet, so the caches it would restore into are turned off. // This rewrites ServerConfig rather than rejecting the launch, which is why @@ -1108,6 +1204,11 @@ int main(int argc, char ** argv) { "[server] │ Use --fa-window 0 for tool-call workloads.\n"); } std::fprintf(stderr, "[server] │ ddtree = %s\n", bargs.ddtree_mode ? "ON" : "off"); + std::fprintf(stderr, "[server] │ specla = %s\n", bargs.specla_mode ? "ON" : "off"); + if (bargs.specla_mode) { + std::fprintf(stderr, "[server] │ specla_top_k = %d\n", specla_top_k); + std::fprintf(stderr, "[server] │ ddtree_tau = %.3g\n", bargs.ddtree_tau); + } 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", diff --git a/server/test/test_chain_rollback_policy.cpp b/server/test/test_chain_rollback_policy.cpp index 13780b89a..874b25355 100644 --- a/server/test/test_chain_rollback_policy.cpp +++ b/server/test/test_chain_rollback_policy.cpp @@ -53,6 +53,9 @@ TEST_CASE(ChainRollbackPolicyFixture, policy_defaults_and_env_parsing) { policy = resolve_chain_rollback_policy(true); CHECK(!policy.checkpoint_f32); CHECK(policy.fast_rollback_threshold == 1); + policy = resolve_chain_rollback_policy(false, true); + CHECK(!policy.checkpoint_f32); + CHECK(policy.fast_rollback_threshold == 1); setenv("DFLASH_SINGLE_CHAIN_CHECKPOINT_F32", "1", 1); // Boolean flags follow the project's non-empty, non-"0" convention. diff --git a/server/test/test_ddtree_tau.cpp b/server/test/test_ddtree_tau.cpp new file mode 100644 index 000000000..c2413c8e4 --- /dev/null +++ b/server/test/test_ddtree_tau.cpp @@ -0,0 +1,154 @@ +// SpecLA confidence-guided draft-tree pruning (arXiv:2607.16673 §6.1). +// +// build_ddtree's best-first expansion pops candidates in descending +// cumulative path log-probability q(v), so the tau_tree window +// (keep q(v) >= q* - tau) is a single early-stop comparison. These tests pin +// the contract: the margin prunes exactly the out-of-window candidates, the +// retained set stays ancestor-closed, the budget still caps the tree, and the +// default (infinite tau) reproduces the unpruned tree. + +#include "CppUnitTestFramework.hpp" +#include "ddtree.h" + +#include +#include +#include + +using dflash::common::build_ddtree; +using dflash::common::build_ddtree_conditional; +using dflash::common::DDTree; + +namespace { +struct DdtreeTauFixture {}; + +// L=3 positions, K=2 ranks. Rank-0 chain is strong; rank-1 siblings weak. +// depth 1: {-0.1, -3.0} +// depth 2: {-0.2, -3.5} +// depth 3: {-0.3, -4.0} +// Cumulative scores: top chain -0.1/-0.3/-0.6; the best sibling is -3.0. +const float kLp[6] = { -0.1f, -3.0f, -0.2f, -3.5f, -0.3f, -4.0f }; +const int32_t kIds[6] = { 10, 11, 20, 21, 30, 31 }; + +bool ancestor_closed(const DDTree & t) { + for (int i = 1; i <= t.n_nodes; i++) { + const int p = t.parents[i]; + if (p < 0 || p > t.n_nodes || p >= i) return false; + } + return true; +} +} // namespace + +TEST_CASE(DdtreeTauFixture, infinite_tau_matches_unpruned_tree) { + DDTree base = build_ddtree(kLp, kIds, 3, 2, 8, /*chain_seed=*/false); + DDTree inf = build_ddtree(kLp, kIds, 3, 2, 8, /*chain_seed=*/false, + std::numeric_limits::infinity()); + CHECK(base.n_nodes == inf.n_nodes); + CHECK(base.token_ids == inf.token_ids); + CHECK(base.parents == inf.parents); +} + +TEST_CASE(DdtreeTauFixture, margin_prunes_low_confidence_branches) { + // Window of 1.0 below q* = -0.1 keeps the top chain (-0.1, -0.3, -0.6) + // and prunes every sibling (best sibling -3.0). + DDTree t = build_ddtree(kLp, kIds, 3, 2, 8, /*chain_seed=*/false, 1.0f); + CHECK(t.n_nodes == 3); + CHECK(t.token_ids == (std::vector{10, 20, 30})); + CHECK(ancestor_closed(t)); + + // A wide window admits the siblings again and expansion runs to the + // budget cap (the candidate space under the siblings exceeds it). + DDTree wide = build_ddtree(kLp, kIds, 3, 2, 8, /*chain_seed=*/false, 10.0f); + CHECK(wide.n_nodes == 8); + CHECK(ancestor_closed(wide)); +} + +TEST_CASE(DdtreeTauFixture, budget_still_caps_within_window) { + DDTree t = build_ddtree(kLp, kIds, 3, 2, 2, /*chain_seed=*/false, 10.0f); + CHECK(t.n_nodes == 2); + CHECK(ancestor_closed(t)); +} + +TEST_CASE(DdtreeTauFixture, precomputed_fast_path_preserves_legacy_order) { + const std::vector expected_tokens = + {10, 20, 30, 11, 20, 30, 21, 30}; + const std::vector expected_depths = + {1, 2, 3, 1, 2, 3, 2, 3}; + const std::vector expected_parents = + {-1, 0, 1, 2, 0, 4, 5, 1, 7}; + const std::vector expected_visibility = { + 1,0,0,0,0,0,0,0,0, + 1,1,0,0,0,0,0,0,0, + 1,1,1,0,0,0,0,0,0, + 1,1,1,1,0,0,0,0,0, + 1,0,0,0,1,0,0,0,0, + 1,0,0,0,1,1,0,0,0, + 1,0,0,0,1,1,1,0,0, + 1,1,0,0,0,0,0,1,0, + 1,1,0,0,0,0,0,1,1, + }; + for (bool chain_seed : {false, true}) { + DDTree t = build_ddtree(kLp, kIds, 3, 2, 8, chain_seed); + CHECK(t.token_ids == expected_tokens); + CHECK(t.depths == expected_depths); + CHECK(t.parents == expected_parents); + CHECK(t.visibility == expected_visibility); + } +} + +TEST_CASE(DdtreeTauFixture, chain_seed_respects_margin) { + // The root child is q* and remains. Deeper top-1 nodes fall outside the + // tiny window, so the seed stops while preserving ancestor closure. + DDTree t = build_ddtree(kLp, kIds, 3, 2, 8, /*chain_seed=*/true, 0.05f); + CHECK(t.n_nodes == 1); + CHECK(t.token_ids == (std::vector{10})); + CHECK(ancestor_closed(t)); +} + +TEST_CASE(DdtreeTauFixture, branch_descendants_use_their_exact_prefix) { + std::vector> queried; + auto next_topk = [&](const std::vector & prefix, int depth, + std::vector & lp, + std::vector & ids) { + CHECK(depth == (int)prefix.size() + 1); + queried.push_back(prefix); + lp = {-0.10f, -0.20f}; + if (prefix.empty()) ids = {10, 11}; + else if (prefix[0] == 10) ids = {20, 21}; + else ids = {30, 31}; + return true; + }; + + DDTree t = build_ddtree_conditional( + next_topk, /*L=*/2, /*K=*/2, /*budget=*/6, + /*chain_seed=*/false, /*tau_tree=*/10.0f); + CHECK(t.n_nodes == 6); + CHECK(ancestor_closed(t)); + + bool saw_ten = false, saw_eleven = false; + for (const auto & prefix : queried) { + saw_ten = saw_ten || prefix == std::vector{10}; + saw_eleven = saw_eleven || prefix == std::vector{11}; + } + CHECK(saw_ten); + CHECK(saw_eleven); + + const int node10 = t.child_maps[0].at(10); + const int node11 = t.child_maps[0].at(11); + CHECK(t.child_maps[node10].count(20) == 1); + CHECK(t.child_maps[node11].count(30) == 1); +} + +TEST_CASE(DdtreeTauFixture, conditioned_tau_prunes_and_stays_ancestor_closed) { + auto next_topk = [](const std::vector & prefix, int, + std::vector & lp, + std::vector & ids) { + lp = {-0.10f, -4.0f}; + ids = {100 + (int32_t)prefix.size(), 200 + (int32_t)prefix.size()}; + return true; + }; + DDTree t = build_ddtree_conditional( + next_topk, /*L=*/4, /*K=*/2, /*budget=*/12, + /*chain_seed=*/false, /*tau_tree=*/0.6f); + CHECK(t.n_nodes == 4); + CHECK(ancestor_closed(t)); +} diff --git a/server/test/test_delta_net_specla.cpp b/server/test/test_delta_net_specla.cpp new file mode 100644 index 000000000..19692f619 --- /dev/null +++ b/server/test/test_delta_net_specla.cpp @@ -0,0 +1,877 @@ +// GPU parity test for the SpecLA topology-masked delta-net verify builder +// (src/delta_net_specla.cpp) against the fused sequential ggml_gated_delta_net +// kernel, which advances the recurrent state token by token and is the +// numerical ground truth for the recurrence. +// +// Checks, per shape/topology case: +// 1. per-node outputs of build_delta_net_specla match the fused kernel +// (chain cases use ggml_gated_delta_net, tree cases the _tree variant); +// 2. host-side DeltaConstruct over the captured factors — +// S_A = exp(g⁺_A) S0 + Σ_{u ∈ path(A)} exp(g⁺_A − g⁺_u) k_u ⊗ ṽ_u +// — matches the fused kernel's per-token intermediate state at EVERY +// possible accepted endpoint A (every prefix of a chain, every node of a +// tree). This is the correctness contract the factor-based accepted-state +// commit (SPECLA.md §1-§3) relies on; +// 3. g⁺ equals the host-computed ancestor path sum of g. + +#include "delta_net_specla.h" +#include "specla_commit_cuda.h" + +#include "ggml.h" +#include "ggml-alloc.h" +#include "ggml-backend.h" +#include "ggml-cuda.h" // ggml_backend_cuda_init; maps to HIP under GGML_USE_HIP + +#include +#include +#include +#include +#include +#include +#include +#include + +using dflash::common::build_delta_net_specla; +using dflash::common::fill_specla_masks; +using dflash::common::make_specla_hld_schedule; + +static int failures = 0; + +#define CHECK_MSG(cond, ...) do { \ + if (!(cond)) { \ + std::fprintf(stderr, "FAIL %s:%d: ", __FILE__, __LINE__); \ + std::fprintf(stderr, __VA_ARGS__); \ + std::fprintf(stderr, "\n"); \ + failures++; \ + } \ +} while (0) + +namespace { + +struct CaseInputs { + int S = 0, H = 0, n = 0; + std::vector parents; // parents[0] = -1, parents[t] < t + std::vector q, k, v, g, b, s0; // layouts as fed to the builders +}; + +CaseInputs make_inputs(int S, int H, int n, const std::vector & parents, + unsigned seed) { + CaseInputs in; + in.S = S; in.H = H; in.n = n; in.parents = parents; + std::mt19937 rng(seed); + std::normal_distribution nd(0.0f, 1.0f); + std::uniform_real_distribution ud(0.0f, 1.0f); + + auto fill_unit_heads = [&](std::vector & dst) { + // [S, H, n]: one l2-normalized S-vector per (head, token), like the + // post-l2_norm q/k the real block feeds the recurrence. + dst.resize((size_t)S * H * n); + for (int t = 0; t < n; t++) { + for (int h = 0; h < H; h++) { + float norm2 = 0.0f; + float * vec = dst.data() + (size_t)t * S * H + (size_t)h * S; + for (int s = 0; s < S; s++) { vec[s] = nd(rng); norm2 += vec[s] * vec[s]; } + const float inv = 1.0f / std::sqrt(norm2 + 1e-6f); + for (int s = 0; s < S; s++) vec[s] *= inv; + } + } + }; + fill_unit_heads(in.q); + fill_unit_heads(in.k); + + in.v.resize((size_t)S * H * n); + for (auto & x : in.v) x = 0.5f * nd(rng); + in.g.resize((size_t)H * n); // [1, H, n] — log-decay, negative + for (auto & x : in.g) x = -(0.05f + 1.5f * ud(rng)); + in.b.resize((size_t)H * n); // [1, H, n] — sigmoid-like in (0,1) + for (auto & x : in.b) x = 0.1f + 0.8f * ud(rng); + in.s0.resize((size_t)S * S * H); + for (auto & x : in.s0) x = 0.1f * nd(rng); + return in; +} + +struct RefOutputs { + std::vector attn; // [S*H per token][n] token-major + std::vector inter; // [S*S*H per token][n] state after node t +}; + +struct SpecLAOutputs { + std::vector out; // [S, H, n] — same layout as RefOutputs::attn + std::vector v_new; // [n, S_v, 1, H] + std::vector g_ps; // [n, 1, 1, H] +}; + +struct PendingFactors { + int count = 0; + std::vector k; + std::vector delta; + std::vector g; + std::vector state_after; +}; + +PendingFactors make_pending_factors(const CaseInputs & in) { + PendingFactors out; + out.count = in.n; + out.k = in.k; + out.g = in.g; + out.delta.resize(in.v.size()); + out.state_after = in.s0; + for (int t = 0; t < in.n; ++t) { + for (int h = 0; h < in.H; ++h) { + const size_t th = (size_t)t*in.H + h; + const float decay = std::exp(in.g[th]); + const float beta = in.b[th]; + const float * k = in.k.data() + th*in.S; + const float * v = in.v.data() + th*in.S; + float * delta = out.delta.data() + th*in.S; + float * state = out.state_after.data() + (size_t)h*in.S*in.S; + for (int col = 0; col < in.S; ++col) { + float kv = 0.0f; + for (int row = 0; row < in.S; ++row) { + kv += state[(size_t)col*in.S + row]*k[row]; + } + delta[col] = (v[col] - decay*kv)*beta; + } + for (int col = 0; col < in.S; ++col) { + for (int row = 0; row < in.S; ++row) { + float & cell = state[(size_t)col*in.S + row]; + cell = std::fma(k[row], delta[col], decay*cell); + } + } + } + } + return out; +} + +float max_abs_diff(const std::vector & a, + const std::vector & b); +std::vector chain_parents(int n); +std::vector random_tree_parents(int n, unsigned seed); + +struct GraphEnv { + ggml_context * ctx = nullptr; + ggml_cgraph * gf = nullptr; + ggml_gallocr_t galloc = nullptr; + + explicit GraphEnv(size_t n_tensors = 512) { + ggml_init_params ip{}; + ip.mem_size = n_tensors * ggml_tensor_overhead() + ggml_graph_overhead(); + ip.no_alloc = true; + ctx = ggml_init(ip); + gf = ggml_new_graph_custom(ctx, 2048, false); + } + ~GraphEnv() { + if (galloc) ggml_gallocr_free(galloc); + if (ctx) ggml_free(ctx); + } + bool alloc_and_run(ggml_backend_t backend) { + galloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); + if (!ggml_gallocr_alloc_graph(galloc, gf)) return false; + return true; + } +}; + +void set_f32(ggml_tensor * t, const std::vector & host) { + GGML_ASSERT((size_t)ggml_nelements(t) == host.size()); + ggml_backend_tensor_set(t, host.data(), 0, host.size() * sizeof(float)); +} + +std::vector get_f32(const ggml_tensor * t, size_t off_elems, size_t n_elems) { + std::vector out(n_elems); + ggml_backend_tensor_get(t, out.data(), off_elems * sizeof(float), + n_elems * sizeof(float)); + return out; +} + +// Reference pass: fused sequential kernel, chain (plain op) or tree variant. +bool run_reference(ggml_backend_t backend, const CaseInputs & in, bool tree_op, + RefOutputs & ref) { + const int S = in.S, H = in.H, n = in.n; + GraphEnv env; + ggml_context * ctx = env.ctx; + + ggml_tensor * q = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, S, H, n, 1); + ggml_tensor * k = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, S, H, n, 1); + ggml_tensor * v = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, S, H, n, 1); + ggml_tensor * g = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 1, H, n, 1); + ggml_tensor * b = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 1, H, n, 1); + ggml_tensor * s = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, S, S, H, 1); + for (ggml_tensor * t : {q, k, v, g, b, s}) ggml_set_input(t); + + ggml_tensor * parent_ids = nullptr; + ggml_tensor * result = nullptr; + if (tree_op) { + parent_ids = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, n); + ggml_set_input(parent_ids); + result = ggml_gated_delta_net_tree(ctx, q, k, v, g, b, s, parent_ids); + } else { + result = ggml_gated_delta_net(ctx, q, k, v, g, b, s); + } + // Intermediates deliberately kept (no set_skip_intermediate): they are the + // per-token ground-truth states DeltaConstruct is validated against. + ggml_set_output(result); + ggml_build_forward_expand(env.gf, result); + + if (!env.alloc_and_run(backend)) return false; + set_f32(q, in.q); set_f32(k, in.k); set_f32(v, in.v); + set_f32(g, in.g); set_f32(b, in.b); set_f32(s, in.s0); + if (parent_ids) { + ggml_backend_tensor_set(parent_ids, in.parents.data(), 0, + sizeof(int32_t) * n); + } + if (ggml_backend_graph_compute(backend, env.gf) != GGML_STATUS_SUCCESS) return false; + + // Packed result: [ attn: S*H*n | final_state: S*S*H | inter: S*S*H*n ] + ref.attn = get_f32(result, 0, (size_t)S * H * n); + ref.inter = get_f32(result, (size_t)S * H * n + (size_t)S * S * H, + (size_t)S * S * H * n); + return true; +} + +bool run_specla(ggml_backend_t backend, const CaseInputs & in, SpecLAOutputs & out) { + const int S = in.S, H = in.H, n = in.n; + GraphEnv env; + ggml_context * ctx = env.ctx; + + ggml_tensor * q = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, S, H, n, 1); + ggml_tensor * k = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, S, H, n, 1); + ggml_tensor * v = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, S, H, n, 1); + ggml_tensor * g = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 1, H, n, 1); + ggml_tensor * b = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 1, H, n, 1); + ggml_tensor * s = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, S, S, H, 1); + ggml_tensor * m_strict = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n, n); + ggml_tensor * m_incl = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n, n); + ggml_tensor * m_eye = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n, n); + for (ggml_tensor * t : {q, k, v, g, b, s, m_strict, m_incl, m_eye}) ggml_set_input(t); + + auto r = build_delta_net_specla(ctx, q, k, v, g, b, s, m_strict, m_incl, m_eye); + for (ggml_tensor * t : {r.output, r.v_new, r.g_ps}) { + ggml_set_output(t); + ggml_build_forward_expand(env.gf, t); + } + + if (!env.alloc_and_run(backend)) return false; + set_f32(q, in.q); set_f32(k, in.k); set_f32(v, in.v); + set_f32(g, in.g); set_f32(b, in.b); set_f32(s, in.s0); + std::vector ms((size_t)n * n), mi((size_t)n * n), me((size_t)n * n); + fill_specla_masks(in.parents.data(), n, ms.data(), mi.data(), me.data()); + set_f32(m_strict, ms); set_f32(m_incl, mi); set_f32(m_eye, me); + if (ggml_backend_graph_compute(backend, env.gf) != GGML_STATUS_SUCCESS) return false; + + out.out = get_f32(r.output, 0, (size_t)S * H * n); + out.v_new = get_f32(r.v_new, 0, (size_t)n * S * H); + out.g_ps = get_f32(r.g_ps, 0, (size_t)n * H); + return true; +} + +bool run_hld_specla(ggml_backend_t backend, const CaseInputs & in, + SpecLAOutputs & out, float & durable_diff, + const PendingFactors * pending = nullptr) { + const int S = in.S, H = in.H, n = in.n; + const auto schedule = make_specla_hld_schedule( + in.parents.data(), n, pending ? pending->count : 0); + if (schedule.packed.empty()) return false; + + GraphEnv env; + ggml_context * ctx = env.ctx; + ggml_init_params fp_ip{}; + fp_ip.mem_size = 16*ggml_tensor_overhead(); + fp_ip.no_alloc = true; + ggml_context * fp_ctx = ggml_init(fp_ip); + ggml_tensor * banks[8] = { + ggml_new_tensor_4d(fp_ctx, GGML_TYPE_F32, S, H, 1, n), + ggml_new_tensor_4d(fp_ctx, GGML_TYPE_F32, S, H, 1, n), + ggml_new_tensor_3d(fp_ctx, GGML_TYPE_F32, H, 1, n), + ggml_new_tensor_3d(fp_ctx, GGML_TYPE_F32, 1, 1, n), + ggml_new_tensor_4d(fp_ctx, GGML_TYPE_F32, S, H, 1, n), + ggml_new_tensor_4d(fp_ctx, GGML_TYPE_F32, S, H, 1, n), + ggml_new_tensor_3d(fp_ctx, GGML_TYPE_F32, H, 1, n), + ggml_new_tensor_3d(fp_ctx, GGML_TYPE_F32, 1, 1, n), + }; + ggml_backend_buffer_t fp_buf = ggml_backend_alloc_ctx_tensors(fp_ctx, backend); + if (!fp_buf) { ggml_free(fp_ctx); return false; } + auto free_fp = [&]() { + ggml_backend_buffer_free(fp_buf); + ggml_free(fp_ctx); + }; + ggml_tensor * q = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, S, H, n, 1); + ggml_tensor * k = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, S, H, n, 1); + ggml_tensor * v = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, S, H, n, 1); + ggml_tensor * g = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 1, H, n, 1); + ggml_tensor * b = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 1, H, n, 1); + ggml_tensor * s = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, S, S, H, 1); + ggml_tensor * meta = ggml_new_tensor_1d( + ctx, GGML_TYPE_I32, (int64_t)schedule.packed.size()); + ggml_tensor * factor_ptrs = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, 8); + for (ggml_tensor * t : {q, k, v, g, b, s, meta, factor_ptrs}) { + ggml_set_input(t); + } + ggml_tensor * result = ggml_gated_delta_net_specla( + ctx, q, k, v, g, b, s, meta, factor_ptrs, + /*n_layers=*/1, /*layer=*/0, /*pending_bank=*/0, + schedule.n_boundaries, schedule.n_chains, schedule.n_waves, + schedule.max_parallel_chains); + ggml_set_output(result); + ggml_build_forward_expand(env.gf, result); + + if (!env.alloc_and_run(backend)) { + free_fp(); + return false; + } + set_f32(q, in.q); set_f32(k, in.k); set_f32(v, in.v); + set_f32(g, in.g); set_f32(b, in.b); set_f32(s, in.s0); + ggml_backend_tensor_set(meta, schedule.packed.data(), 0, + schedule.packed.size()*sizeof(int32_t)); + int64_t ptrs[8]; + for (int i = 0; i < 8; ++i) ptrs[i] = (int64_t)(intptr_t)banks[i]->data; + ggml_backend_tensor_set(factor_ptrs, ptrs, 0, sizeof(ptrs)); + if (pending) { + GGML_ASSERT(pending->count <= n); + ggml_backend_tensor_set(banks[0], pending->k.data(), 0, + pending->k.size()*sizeof(float)); + ggml_backend_tensor_set(banks[1], pending->delta.data(), 0, + pending->delta.size()*sizeof(float)); + ggml_backend_tensor_set(banks[2], pending->g.data(), 0, + pending->g.size()*sizeof(float)); + } + if (ggml_backend_graph_compute(backend, env.gf) != GGML_STATUS_SUCCESS) { + free_fp(); + return false; + } + + const size_t factors = (size_t)S*H*n; + out.out = get_f32(result, 0, factors); + out.v_new = get_f32(banks[5], 0, factors); + out.g_ps = get_f32(banks[6], 0, (size_t)H*n); + durable_diff = max_abs_diff( + pending ? pending->state_after : in.s0, + get_f32(s, 0, in.s0.size())); + + const std::vector captured_k = get_f32(banks[4], 0, factors); + CHECK_MSG(max_abs_diff(in.k, captured_k) == 0.0f, + "HLD raw k capture differs from input"); + CHECK_MSG(max_abs_diff(in.g, out.g_ps) == 0.0f, + "HLD raw gate capture differs from input"); + free_fp(); + return true; +} + +void run_hld_delayed_case(ggml_backend_t backend) { + constexpr int S = 64; + constexpr int H = 4; + const CaseInputs pending_input = make_inputs( + S, H, 3, chain_parents(3), 201); + const PendingFactors pending = make_pending_factors(pending_input); + CaseInputs current = make_inputs( + S, H, 13, random_tree_parents(13, 202), 202); + current.s0 = pending.state_after; + CaseInputs kernel_current = current; + kernel_current.s0 = pending_input.s0; + + RefOutputs reference; + SpecLAOutputs hld; + float durable_diff = INFINITY; + const bool ok = run_reference(backend, current, true, reference) && + run_hld_specla(backend, kernel_current, hld, durable_diff, &pending); + CHECK_MSG(ok, "HLD delayed-update/reference compute failed"); + if (!ok) return; + const float out_diff = max_abs_diff(reference.attn, hld.out); + CHECK_MSG(out_diff <= 5e-4f, + "HLD delayed-update output diff %.3e", out_diff); + CHECK_MSG(durable_diff <= 5e-6f, + "HLD delayed-update durable-state diff %.3e", durable_diff); + std::printf("%-28s HLD out=%.3e durable=%.3e\n", + "hld-delayed-tree", out_diff, durable_diff); +} + +void run_hld_conv_delayed_case(ggml_backend_t backend) { + constexpr int C = 96; + constexpr int K = 4; + constexpr int N = 13; + constexpr int P = 3; + const std::vector parents = random_tree_parents(N, 302); + const auto schedule = make_specla_hld_schedule(parents.data(), N, P); + std::mt19937 rng(301); + std::normal_distribution nd(0.0f, 0.2f); + std::vector x((size_t)C*N), weight((size_t)K*C); + std::vector state((size_t)(K - 1)*C); + std::vector pending((size_t)C*P); + for (auto * vec : {&x, &weight, &state, &pending}) { + for (float & value : *vec) value = nd(rng); + } + + std::vector durable = state; + for (int t = 0; t < P; ++t) { + for (int c = 0; c < C; ++c) { + float * window = durable.data() + (size_t)c*(K - 1); + for (int j = 0; j < K - 2; ++j) window[j] = window[j + 1]; + window[K - 2] = pending[(size_t)t*C + c]; + } + } + std::vector reference((size_t)C*N); + std::vector node_states((size_t)N*C*(K - 1)); + for (int node = 0; node < N; ++node) { + const int parent = parents[(size_t)node]; + for (int c = 0; c < C; ++c) { + float window[K - 1]; + const float * source = parent < 0 + ? durable.data() + (size_t)c*(K - 1) + : node_states.data() + ((size_t)parent*C + c)*(K - 1); + for (int j = 0; j < K - 1; ++j) window[j] = source[j]; + float sum = x[(size_t)node*C + c]*weight[(size_t)c*K + K - 1]; + for (int j = 0; j < K - 1; ++j) { + sum += window[j]*weight[(size_t)c*K + j]; + } + reference[(size_t)node*C + c] = sum/(1.0f + std::exp(-sum)); + float * endpoint = node_states.data() + ((size_t)node*C + c)*(K - 1); + for (int j = 0; j < K - 2; ++j) endpoint[j] = window[j + 1]; + endpoint[K - 2] = x[(size_t)node*C + c]; + } + } + + GraphEnv env; + ggml_context * ctx = env.ctx; + ggml_init_params fp_ip{}; + fp_ip.mem_size = 4*ggml_tensor_overhead(); + fp_ip.no_alloc = true; + ggml_context * fp_ctx = ggml_init(fp_ip); + ggml_tensor * pending_bank = ggml_new_tensor_3d( + fp_ctx, GGML_TYPE_F32, C, 1, N); + ggml_tensor * current_bank = ggml_new_tensor_3d( + fp_ctx, GGML_TYPE_F32, C, 1, N); + ggml_backend_buffer_t fp_buf = ggml_backend_alloc_ctx_tensors(fp_ctx, backend); + CHECK_MSG(fp_buf != nullptr, "conv factor bank allocation failed"); + if (!fp_buf) { ggml_free(fp_ctx); return; } + auto free_fp = [&]() { + ggml_backend_buffer_free(fp_buf); + ggml_free(fp_ctx); + }; + + ggml_tensor * tx = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, C, N, 1); + ggml_tensor * tw = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, K, C); + ggml_tensor * ts = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, K - 1, C); + ggml_tensor * meta = ggml_new_tensor_1d( + ctx, GGML_TYPE_I32, (int64_t)schedule.packed.size()); + ggml_tensor * ptr_table = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, 8); + for (ggml_tensor * tensor : {tx, tw, ts, meta, ptr_table}) ggml_set_input(tensor); + ggml_tensor * result = ggml_ssm_conv_specla( + ctx, tx, tw, ts, meta, ptr_table, + /*n_layers=*/1, /*layer=*/0, /*pending_bank=*/0, + schedule.n_boundaries, schedule.n_chains, schedule.n_waves, + schedule.max_parallel_chains); + const bool supported = ggml_backend_supports_op(backend, result); + CHECK_MSG(supported, + "CUDA rejected partial-block SpecLA conv channels=%d", C); + if (!supported) { + free_fp(); + return; + } + ggml_set_output(result); + ggml_build_forward_expand(env.gf, result); + const bool allocated = env.alloc_and_run(backend); + CHECK_MSG(allocated, "conv HLD graph allocation failed"); + if (!allocated) { + free_fp(); + return; + } + set_f32(tx, x); set_f32(tw, weight); set_f32(ts, state); + ggml_backend_tensor_set(meta, schedule.packed.data(), 0, + schedule.packed.size()*sizeof(int32_t)); + ggml_backend_tensor_set(pending_bank, pending.data(), 0, + pending.size()*sizeof(float)); + int64_t ptrs[8]{}; + ptrs[3] = (int64_t)(intptr_t)pending_bank->data; + ptrs[7] = (int64_t)(intptr_t)current_bank->data; + ggml_backend_tensor_set(ptr_table, ptrs, 0, sizeof(ptrs)); + const bool computed = ggml_backend_graph_compute(backend, env.gf) == GGML_STATUS_SUCCESS; + CHECK_MSG(computed, "conv HLD graph compute failed"); + if (computed) { + const float out_diff = max_abs_diff( + reference, get_f32(result, 0, reference.size())); + const float state_diff = max_abs_diff( + durable, get_f32(ts, 0, durable.size())); + const float factor_diff = max_abs_diff( + x, get_f32(current_bank, 0, x.size())); + CHECK_MSG(out_diff <= 2e-6f, "conv HLD output diff %.3e", out_diff); + CHECK_MSG(state_diff == 0.0f, "conv delayed-state diff %.3e", state_diff); + CHECK_MSG(factor_diff == 0.0f, "conv factor capture diff %.3e", factor_diff); + std::printf("%-28s HLD out=%.3e durable=%.3e\n", + "conv-hld-delayed-tree", out_diff, state_diff); + } + free_fp(); +} + +void run_factorized_conv_commit_case(ggml_backend_t backend) { + constexpr int C = 7; + constexpr int L = 2; + constexpr int T = 4; + constexpr int K = 4; + constexpr int W = K - 1; + + ggml_init_params ip{}; + ip.mem_size = 8*ggml_tensor_overhead(); + ip.no_alloc = true; + ggml_context * ctx = ggml_init(ip); + ggml_tensor * bank = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, C, L, T); + ggml_tensor * state0 = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, W, C); + ggml_tensor * state1 = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, W, C); + ggml_tensor * ptr_table = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, L); + ggml_backend_buffer_t buffer = ggml_backend_alloc_ctx_tensors(ctx, backend); + CHECK_MSG(buffer != nullptr, "factorized conv commit allocation failed"); + if (!buffer) { + ggml_free(ctx); + return; + } + + std::vector factors((size_t)C*L*T); + for (int t = 0; t < T; ++t) { + for (int l = 0; l < L; ++l) { + for (int c = 0; c < C; ++c) { + factors[(size_t)c + (size_t)C*(l + L*t)] = + 100.0f*t + 10.0f*l + c; + } + } + } + std::vector> initial((size_t)L, + std::vector((size_t)W*C)); + for (int l = 0; l < L; ++l) { + for (int c = 0; c < C; ++c) { + for (int j = 0; j < W; ++j) { + initial[(size_t)l][(size_t)j + (size_t)W*c] = + -100.0f*l - 10.0f*c - j; + } + } + } + set_f32(bank, factors); + int64_t ptrs[L] = { + (int64_t)(intptr_t)state0->data, + (int64_t)(intptr_t)state1->data, + }; + ggml_backend_tensor_set(ptr_table, ptrs, 0, sizeof(ptrs)); + ggml_tensor * states[L] = {state0, state1}; + + for (int accepted : {1, T}) { + for (int l = 0; l < L; ++l) set_f32(states[l], initial[(size_t)l]); + std::vector> expected = initial; + for (int t = 0; t < accepted; ++t) { + for (int l = 0; l < L; ++l) { + for (int c = 0; c < C; ++c) { + float * window = expected[(size_t)l].data() + (size_t)W*c; + for (int j = 0; j < W - 1; ++j) window[j] = window[j + 1]; + window[W - 1] = + factors[(size_t)c + (size_t)C*(l + L*t)]; + } + } + } + const bool committed = dflash::common::specla_commit_conv_raw_fused( + (float * const *)ptr_table->data, (const float *)bank->data, + accepted, T, L, C, K, /*stream=*/nullptr); + CHECK_MSG(committed, "factorized conv commit failed A=%d", accepted); + if (committed) { + for (int l = 0; l < L; ++l) { + const float diff = max_abs_diff( + expected[(size_t)l], get_f32(states[l], 0, (size_t)W*C)); + CHECK_MSG(diff == 0.0f, + "factorized conv commit A=%d layer=%d diff %.3e", + accepted, l, diff); + } + } + } + CHECK_MSG(!dflash::common::specla_commit_conv_raw_fused( + (float * const *)ptr_table->data, (const float *)bank->data, + T + 1, T, L, C, K, /*stream=*/nullptr), + "factorized conv commit accepted an out-of-bounds window"); + + ggml_backend_buffer_free(buffer); + ggml_free(ctx); +} + +float max_abs_diff(const std::vector & a, const std::vector & b) { + if (a.size() != b.size()) { + CHECK_MSG(false, "max_abs_diff size mismatch: %zu != %zu", a.size(), b.size()); + return std::numeric_limits::infinity(); + } + float m = 0.0f; + for (size_t i = 0; i < a.size(); i++) m = std::max(m, std::fabs(a[i] - b[i])); + return m; +} + +void run_case(ggml_backend_t backend, const char * name, const CaseInputs & in, + bool tree_op, float out_tol, float state_tol) { + const int S = in.S, H = in.H, n = in.n; + + RefOutputs ref; + SpecLAOutputs sp; + if (!run_reference(backend, in, tree_op, ref)) { + CHECK_MSG(false, "%s: reference compute failed", name); + return; + } + if (!run_specla(backend, in, sp)) { + CHECK_MSG(false, "%s: specla compute failed", name); + return; + } + + // 1. Per-node outputs. + const float out_diff = max_abs_diff(ref.attn, sp.out); + CHECK_MSG(out_diff <= out_tol, "%s: output max diff %.3e > %.3e", name, out_diff, out_tol); + + // 3. g⁺ vs host ancestor path sums (g host layout [H, n] per token block). + float gps_diff = 0.0f; + std::vector gps_host((size_t)n * H); + for (int t = 0; t < n; t++) { + for (int h = 0; h < H; h++) { + float acc = 0.0f; + for (int u = t; u >= 0; u = in.parents[u]) acc += in.g[(size_t)u * H + h]; + gps_host[(size_t)t * H + h] = acc; + // sp.g_ps layout [n, 1, 1, H]: (t, h) at t + h*n + gps_diff = std::max(gps_diff, + std::fabs(acc - sp.g_ps[(size_t)h * n + t])); + } + } + CHECK_MSG(gps_diff <= 1e-5f, "%s: g_ps max diff %.3e", name, gps_diff); + + // 2. DeltaConstruct at every accepted endpoint t: reconstruct S_t from + // {k, ṽ, g⁺} along root→t and compare with the kernel's state after t. + float state_diff = 0.0f; + std::vector s_rec((size_t)S * S * H); + for (int t = 0; t < n; t++) { + // path root→t + std::vector path; + for (int u = t; u >= 0; u = in.parents[u]) path.push_back(u); + for (int h = 0; h < H; h++) { + const float gA = gps_host[(size_t)t * H + h]; + const float decay0 = std::exp(gA); + for (int c = 0; c < S; c++) { + for (int sk = 0; sk < S; sk++) { + s_rec[(size_t)h * S * S + (size_t)c * S + sk] = + decay0 * in.s0[(size_t)h * S * S + (size_t)c * S + sk]; + } + } + for (int u : path) { + const float w = std::exp(gA - gps_host[(size_t)u * H + h]); + // k host layout [S, H, n]; ṽ layout [n, S_v, 1, H] + const float * ku = in.k.data() + (size_t)u * S * H + (size_t)h * S; + for (int c = 0; c < S; c++) { + const float wv = w * sp.v_new[(size_t)h * n * S + (size_t)c * n + u]; + float * dst = s_rec.data() + (size_t)h * S * S + (size_t)c * S; + for (int sk = 0; sk < S; sk++) dst[sk] += wv * ku[sk]; + } + } + } + const std::vector s_ref(ref.inter.begin() + (size_t)t * S * S * H, + ref.inter.begin() + (size_t)(t + 1) * S * S * H); + state_diff = std::max(state_diff, max_abs_diff(s_rec, s_ref)); + } + CHECK_MSG(state_diff <= state_tol, "%s: DeltaConstruct state max diff %.3e > %.3e", + name, state_diff, state_tol); + + std::printf("%-28s S=%-3d H=%-2d n=%-3d out=%.3e g+=%.3e state=%.3e\n", + name, S, H, n, out_diff, gps_diff, state_diff); +} + +void run_hld_case(ggml_backend_t backend, const char * name, + const CaseInputs & in, bool tree_op, float tolerance) { + RefOutputs ref; + SpecLAOutputs hld; + float durable_diff = INFINITY; + if (!run_reference(backend, in, tree_op, ref) || + !run_hld_specla(backend, in, hld, durable_diff)) { + CHECK_MSG(false, "%s: HLD/reference compute failed", name); + return; + } + const float out_diff = max_abs_diff(ref.attn, hld.out); + CHECK_MSG(out_diff <= tolerance, "%s: HLD output diff %.3e > %.3e", + name, out_diff, tolerance); + CHECK_MSG(durable_diff == 0.0f, + "%s: zero-pending verify mutated durable state %.3e", + name, durable_diff); + std::printf("%-28s HLD out=%.3e durable=%.3e\n", + name, out_diff, durable_diff); +} + +void test_hld_schedule() { + const std::vector parents = {-1, 0, 1, 1, 3, 0, 5, 5}; + const auto hld = make_specla_hld_schedule( + parents.data(), (int)parents.size(), /*pending_count=*/3); + CHECK_MSG(!hld.packed.empty(), "HLD schedule unexpectedly empty"); + if (hld.packed.empty()) return; + CHECK_MSG(hld.n_nodes == (int)parents.size(), "HLD node count mismatch"); + CHECK_MSG(hld.n_chains >= 2 && hld.n_waves >= 2, + "HLD tree was not decomposed into dependent waves"); + CHECK_MSG(hld.packed[0] == 0x534c4148 && hld.packed[5] == 3, + "HLD ABI header mismatch"); + const int order_off = hld.packed[6]; + std::vector order( + hld.packed.begin() + order_off, + hld.packed.begin() + order_off + parents.size()); + std::sort(order.begin(), order.end()); + for (int i = 0; i < (int)order.size(); ++i) { + CHECK_MSG(order[(size_t)i] == i, "HLD schedule lost/duplicated node %d", i); + } +} + +std::vector chain_parents(int n) { + std::vector p(n); + for (int t = 0; t < n; t++) p[t] = t - 1; + return p; +} + +std::vector random_tree_parents(int n, unsigned seed) { + std::mt19937 rng(seed); + std::vector p(n); + p[0] = -1; + for (int t = 1; t < n; t++) { + // Bias toward recent nodes so trees have realistic depth. + std::uniform_int_distribution d(std::max(0, t - 4), t - 1); + p[t] = d(rng); + } + return p; +} + +void test_production_commit_kernel(ggml_backend_t backend) { + constexpr int S_k = 5, S_v = 7, H = 3, L = 4, T = 6; + const std::vector accepted = {0, 2, 5}; + + ggml_init_params ip{}; + ip.mem_size = 64 * ggml_tensor_overhead(); + ip.no_alloc = true; + ggml_context * ctx = ggml_init(ip); + std::vector states((size_t)L); + for (int l = 0; l < L; l++) { + states[(size_t)l] = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, S_k, S_v, H); + } + ggml_tensor * fk = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, S_k, H, L, T); + ggml_tensor * fv = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, S_v, H, L, T); + ggml_tensor * fg = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, H, L, T); + ggml_tensor * idx = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, T); + ggml_tensor * state_ptrs = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, L); + + ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors(ctx, backend); + CHECK_MSG(buf != nullptr, "production commit: allocation failed"); + if (!buf) { ggml_free(ctx); return; } + + auto pattern = [](size_t i, float scale) { + return scale * (float)((int)(i % 17) - 8); + }; + std::vector fk_h((size_t)S_k * H * L * T); + std::vector fv_h((size_t)S_v * H * L * T); + std::vector fg_h((size_t)H * L * T); + for (size_t i = 0; i < fk_h.size(); i++) fk_h[i] = pattern(i, 0.013f); + for (size_t i = 0; i < fv_h.size(); i++) fv_h[i] = pattern(i + 3, 0.017f); + for (size_t i = 0; i < fg_h.size(); i++) fg_h[i] = -0.02f * (float)(1 + i % 9); + set_f32(fk, fk_h); set_f32(fv, fv_h); set_f32(fg, fg_h); + ggml_backend_tensor_set(idx, accepted.data(), 0, accepted.size() * sizeof(int32_t)); + + std::vector> state_h((size_t)L); + std::vector state_ptr_h((size_t)L); + for (int l = 0; l < L; l++) { + auto & s = state_h[(size_t)l]; + s.resize((size_t)S_k * S_v * H); + for (size_t i = 0; i < s.size(); i++) s[i] = pattern(i + (size_t)l, 0.01f); + set_f32(states[(size_t)l], s); + state_ptr_h[(size_t)l] = (int64_t)(intptr_t)states[(size_t)l]->data; + } + ggml_backend_tensor_set(state_ptrs, state_ptr_h.data(), 0, + state_ptr_h.size() * sizeof(int64_t)); + + bool launched = false; + bool ok = dflash::common::specla_commit_fused( + (float * const *)state_ptrs->data, (const float *)fk->data, + (const float *)fv->data, (const float *)fg->data, + (const int32_t *)idx->data, (int)accepted.size(), + S_k, S_v, H, L, nullptr, &launched); + CHECK_MSG(ok && launched, "production SSM commit failed (launched=%d)", (int)launched); + if (ok) { + float max_diff = 0.0f; + for (int l = 0; l < L; l++) { + std::vector expected = state_h[(size_t)l]; + for (int h = 0; h < H; h++) { + const int tA = accepted.back(); + const size_t ga_off = (size_t)h + (size_t)H * (l + L * tA); + const float gA = fg_h[ga_off]; + for (int c = 0; c < S_v; c++) { + for (int i = 0; i < S_k; i++) { + const size_t se = (size_t)h * S_k * S_v + (size_t)c * S_k + i; + float value = std::exp(gA) * expected[se]; + for (int t : accepted) { + const size_t fo = (size_t)h + (size_t)H * (l + L * t); + value += std::exp(gA - fg_h[fo]) * + fk_h[fo * S_k + i] * fv_h[fo * S_v + c]; + } + expected[se] = value; + } + } + } + max_diff = std::max(max_diff, + max_abs_diff(expected, get_f32(states[(size_t)l], 0, expected.size()))); + } + CHECK_MSG(max_diff <= 2e-6f, "production SSM commit max diff %.3e", max_diff); + } + + ggml_backend_buffer_free(buf); + ggml_free(ctx); +} + +} // namespace + +int main() { + ggml_backend_t backend = ggml_backend_cuda_init(0); + if (!backend) { + std::fprintf(stderr, "test_delta_net_specla: no GPU backend available\n"); + return 77; // ctest SKIP + } + + const float kOutTol = 5e-4f; + const float kStateTol = 5e-4f; + + test_hld_schedule(); + test_production_commit_kernel(backend); + + run_hld_case(backend, "hld-chain-model-shape", + make_inputs(128, 8, 16, chain_parents(16), 101), + false, kOutTol); + run_hld_case(backend, "hld-tree-model-shape", + make_inputs(128, 8, 24, random_tree_parents(24, 102), 102), + true, kOutTol); + run_hld_delayed_case(backend); + run_hld_conv_delayed_case(backend); + run_factorized_conv_commit_case(backend); + + // Chain drafts vs the plain fused op. + run_case(backend, "chain-small", + make_inputs(64, 4, 8, chain_parents(8), 1), false, kOutTol, kStateTol); + run_case(backend, "chain-model-shape", + make_inputs(128, 8, 16, chain_parents(16), 2), false, kOutTol, kStateTol); + run_case(backend, "chain-n1", + make_inputs(64, 2, 1, chain_parents(1), 3), false, kOutTol, kStateTol); + run_case(backend, "chain-odd", + make_inputs(64, 2, 33, chain_parents(33), 4), false, kOutTol, kStateTol); + + // Tree drafts vs the fused tree op. + run_case(backend, "tree-chain-shaped", + make_inputs(64, 4, 8, chain_parents(8), 5), true, kOutTol, kStateTol); + { + std::vector star(9, 0); + star[0] = -1; + run_case(backend, "tree-star", + make_inputs(64, 4, 9, star, 6), true, kOutTol, kStateTol); + } + run_case(backend, "tree-random-small", + make_inputs(64, 4, 15, random_tree_parents(15, 42), 7), true, kOutTol, kStateTol); + run_case(backend, "tree-model-shape", + make_inputs(128, 8, 31, random_tree_parents(31, 43), 8), true, kOutTol, kStateTol); + + // Full qwen35-27B delta-net shape: S=128, H_v=48, 16-token verify window. + run_case(backend, "chain-qwen35-27b", + make_inputs(128, 48, 16, chain_parents(16), 9), false, kOutTol, kStateTol); + run_case(backend, "tree-qwen35-27b", + make_inputs(128, 48, 24, random_tree_parents(24, 44), 10), true, kOutTol, kStateTol); + + ggml_backend_free(backend); + if (failures) { + std::fprintf(stderr, "test_delta_net_specla: %d failure(s)\n", failures); + return 1; + } + std::printf("test_delta_net_specla: all cases passed\n"); + return 0; +} diff --git a/server/test/test_dflash.cpp b/server/test/test_dflash.cpp index d2827db2b..91e96604e 100644 --- a/server/test/test_dflash.cpp +++ b/server/test/test_dflash.cpp @@ -19,11 +19,16 @@ // #include "dflash27b.h" +#include #include "internal.h" +#include "delta_net_specla.h" +#include "specla_commit_cuda.h" +#include "specla_mode.h" #include "draft_graph.h" #include "qwen3_drafter.h" #include "gpu_runtime_compat.h" #include "chain_rollback_policy.h" +#include "platform_env.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. @@ -59,10 +64,12 @@ using to_fp32_cuda_t = void (*)(const void *, float *, int64_t, cudaStream_t); to_fp32_cuda_t ggml_get_to_fp32_cuda(ggml_type type); #include +#include #include #include #include #include +#include #ifdef _WIN32 #define setenv(name, value, overwrite) _putenv_s(name, value) @@ -753,6 +760,11 @@ int main(int argc, char ** argv) { int ddtree_budget = 64; float ddtree_temp = 1.0f; // softmax temperature for top-K extract bool ddtree_chain_seed = true; // pre-seed full chain (vs paper's pure best-first) + float ddtree_tau = std::numeric_limits::infinity(); // SpecLA confidence margin + bool ddtree_tau_set = false; + bool specla_mode = false; + bool specla_top_k_set = false; + int specla_top_k = 4; bool profile_scaling = false; // microbench: time target forward at varying N bool time_breakdown = false; // one-token time breakdown: prefill/decode/verify × ctx size bool hybrid_bench_only = false; // skip monolithic scenarios, run only hybrid/pipelined @@ -804,11 +816,43 @@ int main(int argc, char ** argv) { if (std::strcmp(argv[i], "--daemon") == 0) daemon_mode = true; else if (std::strcmp(argv[i], "--seq-verify") == 0) seq_verify = true; else if (std::strcmp(argv[i], "--fast-rollback") == 0) fast_rollback = true; + else if (std::strcmp(argv[i], "--specla") == 0) { + specla_mode = true; + ddtree_mode = true; + fast_rollback = true; + } + else if (std::strncmp(argv[i], "--specla-top-k=", 15) == 0) { + const char * value = argv[i] + 15; + char * end = nullptr; + const long parsed = std::strtol(value, &end, 10); + if (end == value || *end != '\0' || parsed <= 0 || parsed > INT_MAX) { + std::fprintf(stderr, "bad --specla-top-k value: %s\n", value); + return 2; + } + specla_top_k = (int)parsed; + specla_top_k_set = true; + } else if (std::strcmp(argv[i], "--ddtree") == 0) { ddtree_mode = true; fast_rollback = true; } else if (std::strncmp(argv[i], "--ddtree-budget=", 16) == 0) { ddtree_budget = std::atoi(argv[i] + 16); if (ddtree_budget <= 0) ddtree_budget = 64; } + else if (std::strncmp(argv[i], "--ddtree-tau=", 13) == 0) { + const char * tau_str = argv[i] + 13; + char * end = nullptr; + errno = 0; + const float tau = std::strtof(tau_str, &end); + // Mirror --ddtree-temp/--ddtree-budget: invalid values fall back + // to the documented default instead of silently enabling the + // pruned-tree path (atof("garbage") would read as 0.0). + if (end == tau_str || *end != '\0' || + !std::isfinite(tau) || tau <= 0.0f) { + ddtree_tau = std::numeric_limits::infinity(); + } else { + ddtree_tau = tau; + ddtree_tau_set = true; + } + } else if (std::strncmp(argv[i], "--ddtree-temp=", 14) == 0) { ddtree_temp = (float)std::atof(argv[i] + 14); if (ddtree_temp <= 0.0f) ddtree_temp = 1.0f; @@ -1045,10 +1089,25 @@ int main(int argc, char ** argv) { (void)n; #endif }; + if (specla_mode && seq_verify) { + std::fprintf(stderr, "--specla and --seq-verify are mutually exclusive\n"); + return 2; + } if (fast_rollback && seq_verify && !ddtree_mode) { std::fprintf(stderr, "--fast-rollback and --seq-verify are mutually exclusive\n"); return 2; } + if (specla_mode) { + if (!ddtree_tau_set) ddtree_tau = 6.0f; + set_environment_variable("DFLASH_SPECLA", "1", true); + if (specla_top_k_set) { + set_environment_variable( + "DFLASH_SPECLA_TOPK", std::to_string(specla_top_k).c_str(), true); + } + } else if (specla_top_k_set) { + std::fprintf(stderr, "--specla-top-k requires --specla\n"); + return 2; + } if (target_split_dflash) target_split_load_draft = true; if (target_gpus.empty()) target_gpus.push_back(target_gpu); if (target_gpus.size() == 1) target_gpu = target_gpus[0]; @@ -3256,7 +3315,13 @@ int main(int argc, char ** argv) { // DDTree top-K: use GPU argmax for draft_tok; full logits transfer // only when DDTree needs top-K (K>1) for sibling expansion. - const int ddtree_K = (ddtree_budget > q_len - 1) ? 8 : 1; + // Match the SpecLA paper's tree route (top-k=4) when factor-buffered + // state is active; retain the historical top-8 baseline otherwise. + const int ddtree_K = (ddtree_budget > q_len - 1) + ? (!cache.factor_k.empty() + ? std::min(dflash::common::specla_tree_topk(), vocab) + : 8) + : 1; if (draft_hidden_bridge) { for (int i = 0; i < q_len; i++) { @@ -3385,14 +3450,38 @@ int main(int argc, char ** argv) { ddtree_top_log_probs.data(), ddtree_top_token_ids.data(), L, ddtree_K, ddtree_budget, - ddtree_chain_seed); + ddtree_chain_seed, ddtree_tau); const int N_actual = 1 + tree.n_nodes; // actual tree size - const int N = ddtree_budget + 1; // fixed allocation size for gallocr reuse + // The HLD graph is topology-specific and rebuilt for this round, + // so padding pruned trees only wastes target work. + const int N = !cache.factor_k.empty() + ? N_actual + : (std::isfinite(ddtree_tau) ? N_actual : ddtree_budget + 1); + + // Root-inclusive parent topology, including harmless padding + // children. SpecLA consumes this schedule during graph build so + // every delta layer executes state-resident HLD chains. + std::vector parent_ids(N, 0); + parent_ids[0] = -1; + for (int i = 1; i < N_actual; i++) { + parent_ids[i] = (int32_t)tree.parents[i]; + } + SpecLAHLDSchedule hld; + const SpecLAHLDSchedule * hld_ptr = nullptr; + if (!cache.factor_k.empty()) { + hld = make_specla_hld_schedule( + parent_ids.data(), N, cache.specla_pending_count); + if (hld.packed.empty()) { + std::fprintf(stderr, "ddtree HLD schedule failed\n"); + return 1; + } + hld_ptr = &hld; + } if (!build_target_step_tree(sg, w, cache, backend, /*kv_start=*/committed, /*n_tokens=*/N, - g_fa_window, g_kq_stride_pad)) { + g_fa_window, g_kq_stride_pad, hld_ptr)) { std::fprintf(stderr, "ddtree verify build failed\n"); return 1; } T_verify_build = sync_us(); @@ -3450,14 +3539,23 @@ int main(int argc, char ** argv) { ggml_backend_tensor_set(sg.attn_mask, mask_buf.data(), 0, sizeof(uint16_t) * mask_buf.size()); - // parent_ids: actual tree nodes, then padding → point to root (slot 0) - std::vector parent_ids(N, 0); - parent_ids[0] = -1; - for (int i = 1; i < N_actual; i++) parent_ids[i] = (int32_t)tree.parents[i]; - // Padding slots: parent=0 (root). DeltaNet kernel processes them - // but their outputs are never used (masked out in attention). - ggml_backend_tensor_set(sg.parent_ids, parent_ids.data(), 0, - sizeof(int32_t) * N); + // Padding slots remain root children. Their outputs are ignored. + if (sg.parent_ids->buffer) { + ggml_backend_tensor_set(sg.parent_ids, parent_ids.data(), 0, + sizeof(int32_t) * N); + } + + // SpecLA: ancestor masks over the same root-inclusive node order. + if (sg.specla_m_strict) { + std::vector sp_ms((size_t)N * N); + std::vector sp_mi((size_t)N * N); + std::vector sp_me((size_t)N * N); + fill_specla_masks(parent_ids.data(), N, + sp_ms.data(), sp_mi.data(), sp_me.data()); + ggml_backend_tensor_set(sg.specla_m_strict, sp_ms.data(), 0, sizeof(float) * sp_ms.size()); + ggml_backend_tensor_set(sg.specla_m_incl, sp_mi.data(), 0, sizeof(float) * sp_mi.size()); + ggml_backend_tensor_set(sg.specla_m_eye, sp_me.data(), 0, sizeof(float) * sp_me.size()); + } T_verify_set = sync_us(); tt_verify_set += std::chrono::duration(T_verify_set - T_verify_build).count(); @@ -3587,13 +3685,8 @@ int main(int argc, char ** argv) { // the next iteration feeds it to w.embedder.embed(), that fails, // and the decode loop returns 1 without writing the output file // or printing the summary line (issue #191). - if (hit_eos || last_tok < 0 || IS_EOS_TOK(last_tok, w)) { - committed += commit_n; - n_generated += commit_n; - n_accept_sum += commit_n; - n_draft_steps++; - break; - } + const bool stop_after_tree_commit = + hit_eos || last_tok < 0 || IS_EOS_TOK(last_tok, w); // Rollback: per-layer DeltaNet SSM and conv state + KV compaction // for full-attention layers. @@ -3617,12 +3710,65 @@ int main(int argc, char ** argv) { { const int n_delta = (int)sg.delta_captures.size(); cudaStream_t stream = nullptr; + // SpecLA: keep the accepted raw factors pending. A pure spine + // is already contiguous and only rotates banks; a sibling + // walk is compacted into path order in one kernel. The next + // HLD verify applies both GDN and conv factors while their + // state tiles are resident. + const bool specla_commit = !cache.factor_k.empty(); + if (specla_commit) { + if (!cache.factor_k_all || !cache.factor_v_new_all || + !cache.factor_g_ps_all || !cache.conv_factor_all || + !cache.factor_k_all_alt || !cache.factor_v_new_all_alt || + !cache.factor_g_ps_all_alt || !cache.conv_factor_all_alt) { + std::fprintf(stderr, "ddtree SpecLA factor banks missing\n"); + return 1; + } + SpeclaFactorBanks banks; + banks.k[0] = (float *)cache.factor_k_all->data; + banks.v[0] = (float *)cache.factor_v_new_all->data; + banks.g[0] = (float *)cache.factor_g_ps_all->data; + banks.conv[0] = (float *)cache.conv_factor_all->data; + banks.k[1] = (float *)cache.factor_k_all_alt->data; + banks.v[1] = (float *)cache.factor_v_new_all_alt->data; + banks.g[1] = (float *)cache.factor_g_ps_all_alt->data; + banks.conv[1] = (float *)cache.conv_factor_all_alt->data; + + const int old_pending_bank = cache.specla_pending_bank; + if (walked_sibling_for_rollback) { + if (!cache.specla_idx || !cache.specla_idx->data) { + std::fprintf(stderr, "ddtree SpecLA index buffer missing\n"); + return 1; + } + std::vector acc_idx( + accepted.begin(), accepted.begin() + commit_n); + ggml_backend_tensor_set(cache.specla_idx, acc_idx.data(), 0, + acc_idx.size()*sizeof(int32_t)); + } + int new_pending_bank = old_pending_bank; + if (!specla_rotate_pending_factors( + banks, + walked_sibling_for_rollback + ? (const int32_t *)cache.specla_idx->data : nullptr, + old_pending_bank, walked_sibling_for_rollback, commit_n, + (int)cache.factor_k_all->ne[0], + (int)cache.factor_v_new_all->ne[0], + (int)cache.factor_k_all->ne[1], + n_delta, (int)cache.conv_factor_all->ne[0], + /*stream=*/nullptr, &new_pending_bank)) { + std::fprintf(stderr, "ddtree SpecLA factor rotation failed\n"); + return 1; + } + cache.specla_pending_bank = new_pending_bank; + cache.specla_pending_count = commit_n; + } for (int il = 0; il < n_delta; il++) { const DeltaNetCapture & cap = sg.delta_captures[il]; - if (!cap.ssm_intermediate_states || !cap.conv_input) { + if ((!specla_commit && !cap.ssm_intermediate_states) || !cap.conv_input) { std::fprintf(stderr, "ddtree rollback: missing capture layer %d\n", il); return 1; } + if (specla_commit) continue; // SSM state rollback: source is cache.ssm_intermediate_states // ([S_v, S_v, H_v, max_verify_tokens]) at slot rollback_dfs. // Destination is cache.ssm_state[il] (f32). Use ggml's @@ -3781,6 +3927,7 @@ int main(int argc, char ** argv) { n_generated += commit_n; n_accept_sum += commit_n; // for stats n_draft_steps++; + if (stop_after_tree_commit) break; continue; // skip the rest of the verify/commit logic for this iter } @@ -3803,6 +3950,21 @@ int main(int argc, char ** argv) { ggml_backend_tensor_set(sg.inp_embed, verify_embed.data(), 0, sizeof(float) * verify_embed.size()); + // SpecLA: chain topology masks (parents[t] = t-1), host-filled + // like the attention mask below. + if (sg.specla_m_strict) { + std::vector sp_parents(q_len); + for (int t = 0; t < q_len; t++) sp_parents[t] = t - 1; + std::vector sp_ms((size_t)q_len * q_len); + std::vector sp_mi((size_t)q_len * q_len); + std::vector sp_me((size_t)q_len * q_len); + fill_specla_masks(sp_parents.data(), q_len, + sp_ms.data(), sp_mi.data(), sp_me.data()); + ggml_backend_tensor_set(sg.specla_m_strict, sp_ms.data(), 0, sizeof(float) * sp_ms.size()); + ggml_backend_tensor_set(sg.specla_m_incl, sp_mi.data(), 0, sizeof(float) * sp_mi.size()); + ggml_backend_tensor_set(sg.specla_m_eye, sp_me.data(), 0, sizeof(float) * sp_me.size()); + } + // M-RoPE axis-major layout: [axis0_tok0..axis0_tokN-1, axis1_..., axis2_..., axis3_...]. // First 3 axes hold the token position; axis 3 is always 0 for text. for (int i = 0; i < q_len; i++) { @@ -3946,7 +4108,47 @@ int main(int argc, char ** argv) { // Rollback SSM + conv state unless we fully accepted (in which case // state after processing all q_len tokens is exactly what we want). - if (commit_n < q_len) { + // + // SpecLA (DFLASH_SPECLA=1): current candidates remain outside the + // durable state, so their bank is rotated even on full acceptance. + const bool specla_commit = !cache.factor_k.empty(); + if (specla_commit) { + if (!cache.factor_k_all || !cache.factor_v_new_all || + !cache.factor_g_ps_all || !cache.conv_factor_all || + !cache.factor_k_all_alt || !cache.factor_v_new_all_alt || + !cache.factor_g_ps_all_alt || !cache.conv_factor_all_alt) { + std::fprintf(stderr, "SpecLA factor banks missing\n"); + return 1; + } + SpeclaFactorBanks banks; + banks.k[0] = (float *)cache.factor_k_all->data; + banks.v[0] = (float *)cache.factor_v_new_all->data; + banks.g[0] = (float *)cache.factor_g_ps_all->data; + banks.conv[0] = (float *)cache.conv_factor_all->data; + banks.k[1] = (float *)cache.factor_k_all_alt->data; + banks.v[1] = (float *)cache.factor_v_new_all_alt->data; + banks.g[1] = (float *)cache.factor_g_ps_all_alt->data; + banks.conv[1] = (float *)cache.conv_factor_all_alt->data; + + // The accepted prefix is already contiguous in the bank the + // HLD verify just produced. Rotate it into the pending role; + // the next verify fuses its recurrent update with state load. + int new_pending_bank = cache.specla_pending_bank; + if (!specla_rotate_pending_factors( + banks, /*idx_dev=*/nullptr, cache.specla_pending_bank, + /*walked_sibling=*/false, commit_n, + (int)cache.factor_k_all->ne[0], + (int)cache.factor_v_new_all->ne[0], + (int)cache.factor_k_all->ne[1], + (int)cache.factor_k.size(), + (int)cache.conv_factor_all->ne[0], + /*stream=*/nullptr, &new_pending_bank)) { + std::fprintf(stderr, "SpecLA factor bank rotation failed\n"); + return 1; + } + cache.specla_pending_bank = new_pending_bank; + cache.specla_pending_count = commit_n; + } else if (commit_n < q_len) { const int rollback_idx = commit_n - 1; // index into per-step intermediates // Temporary ctx for view tensors (no data alloc — views inherit // data pointers from their already-live sources). @@ -4136,6 +4338,46 @@ int main(int argc, char ** argv) { n_draft_steps++; } + // A pending path is normally consumed by the following HLD verify. At a + // generation boundary there is no following verify, so materialize it + // once to leave reusable cache/snapshot state exact. + if (!cache.factor_k.empty() && cache.specla_pending_count > 0) { + if (!cache.factor_k_all || !cache.factor_v_new_all || + !cache.factor_g_ps_all || !cache.conv_factor_all || + !cache.factor_k_all_alt || !cache.factor_v_new_all_alt || + !cache.factor_g_ps_all_alt || !cache.conv_factor_all_alt || + !cache.specla_state_ptrs || !cache.specla_conv_state_ptrs) { + std::fprintf(stderr, "final SpecLA state flush buffers missing\n"); + return 1; + } + SpeclaFactorBanks banks; + banks.k[0] = (float *)cache.factor_k_all->data; + banks.v[0] = (float *)cache.factor_v_new_all->data; + banks.g[0] = (float *)cache.factor_g_ps_all->data; + banks.conv[0] = (float *)cache.conv_factor_all->data; + banks.k[1] = (float *)cache.factor_k_all_alt->data; + banks.v[1] = (float *)cache.factor_v_new_all_alt->data; + banks.g[1] = (float *)cache.factor_g_ps_all_alt->data; + banks.conv[1] = (float *)cache.conv_factor_all_alt->data; + + if (!specla_flush_pending_factors( + banks, + (float * const *)cache.specla_state_ptrs->data, + (float * const *)cache.specla_conv_state_ptrs->data, + cache.specla_pending_bank, cache.specla_pending_count, + (int)cache.factor_k_all->ne[0], + (int)cache.factor_v_new_all->ne[0], + (int)cache.factor_k_all->ne[1], + (int)cache.ssm_state.size(), + (int)cache.conv_factor_all->ne[0], w.ssm_d_conv, + /*stream=*/nullptr)) { + std::fprintf(stderr, "final SpecLA state flush failed\n"); + return 1; + } + cache.specla_pending_count = 0; + ggml_backend_synchronize(target_backend); + } + auto t_gen1 = std::chrono::steady_clock::now(); double gen_s = std::chrono::duration(t_gen1 - t_gen0).count(); double tps = n_generated / std::max(1e-9, gen_s); diff --git a/server/test/test_kvflash_pool_sizing.cpp b/server/test/test_kvflash_pool_sizing.cpp index 07091be97..7bdea044c 100644 --- a/server/test/test_kvflash_pool_sizing.cpp +++ b/server/test/test_kvflash_pool_sizing.cpp @@ -20,6 +20,11 @@ struct KvflashPoolSizingFixture {}; } TEST_CASE(KvflashPoolSizingFixture, kvflash_pool_sizing_suite) { + REQUIRE(!kvflash_pool_requested(nullptr)); + REQUIRE(!kvflash_pool_requested("0")); + REQUIRE(kvflash_pool_requested("auto")); + REQUIRE(kvflash_pool_requested("4096")); + REQUIRE(!kvflash_fixed_pool_requested(nullptr)); REQUIRE(!kvflash_fixed_pool_requested("0")); REQUIRE(!kvflash_fixed_pool_requested("auto"));