diff --git a/server/CMakeLists.txt b/server/CMakeLists.txt index 5f0cd92bb..cf1e4ee8b 100644 --- a/server/CMakeLists.txt +++ b/server/CMakeLists.txt @@ -454,6 +454,7 @@ add_library(dflash_common STATIC src/common/dynamic_backend.cpp src/common/domino_head.cpp src/common/dspark_head.cpp + src/common/dflash2_head.cpp src/common/target_shard_ipc.cpp src/common/target_shard_ipc_daemon.cpp src/common/dflash_feature_ring.cpp diff --git a/server/README.md b/server/README.md index 2ae61f481..59f10149b 100644 --- a/server/README.md +++ b/server/README.md @@ -521,6 +521,16 @@ Same DFlash + PFlash stack on AMD GPUs. PR #119 ports the Phase 2 rocWMMA flashp **RDNA4 — Radeon AI PRO R9700 (`gfx1201`, 32 GB).** First-class RDNA4 target as of this build. Qwen3.6-27B Q4_K_M + DFlash draft (`dflash-draft-3.6-q4_k_m.gguf`), `--ddtree-budget=22`: **54.65 tok/s mean DFlash decode** across the 10-prompt HumanEval suite (`bench_he.py --n-gen 256`, AL 7.14, range 36.9–93.0 tok/s) on ROCm 7.1.1. The rocWMMA Phase 2 flashprefill kernels are numerically correct on RDNA4 — ROCm 7.1's rocWMMA handles the gfx12 WMMA operand-format change internally, so no kernel changes are needed (`test_flashprefill_kernels` PASS on `gfx1201`: max diff 5e-4, e2e `flash_prefill_forward_bf16` at S=8192 in 10.7 ms/iter). Note `gfx1200` (RX 9060) and `gfx1201` (RX 9070 / R9700) are **not** code-object compatible — build for `gfx1201` explicitly for the R9700. +For Qwen3.8-27B IQ4_XS with the Q8_0 DFlash2 drafter, the drafter's metadata +block size is conservative on the R9700. `--draft-block-size 12` is the +general-purpose setting measured on `gfx1201`: 230.2 versus 159.8 aggregate +decode tok/s on the ten-prompt HumanEval benchmark (+44%), 178.5 versus 139.6 +tok/s across all 164 HumanEval+ tasks (+28%), and 145/164 versus 143/164 +pass@1. A code-heavy deployment can use `--draft-block-size 16` for 279.1 +tok/s (+75%) on the short HumanEval benchmark, at the cost of small regressions +on some low-acceptance prose prompts. Values are intentionally explicit rather +than GPU defaults because the optimum depends on the drafter and workload. + ```bash git clone --recurse-submodules https://github.com/Luce-Org/lucebox-hub && cd lucebox-hub/server diff --git a/server/deps/llama.cpp/ggml/include/ggml.h b/server/deps/llama.cpp/ggml/include/ggml.h index a83bac9fa..a4db13a1f 100644 --- a/server/deps/llama.cpp/ggml/include/ggml.h +++ b/server/deps/llama.cpp/ggml/include/ggml.h @@ -2741,6 +2741,24 @@ extern "C" { struct ggml_tensor * c, struct ggml_tensor * parent_ids); + // dflash extension: fused causal-conv step for recurrent decode/verify. + // Replaces transpose + concat(state, x) + ssm_conv + silu + state + // write-back with one kernel. + // x: [C, T, S] f32, rows contiguous (token stride may be + // larger than C, e.g. a row-slice of a stacked GEMV) + // c: [K, C] f32 depthwise conv weights + // conv_state: [K-1, C, S] f32 history; READ, then OVERWRITTEN in + // place with the last K-1 conv inputs + // conv_input_out: optional [>= K-1+T, C, S] f32; receives the full + // conv window (history rows then x rows) per channel, + // for speculative-decode rollback. May be a view. + // Returns silu(conv(x)) as [C, T, S]. CUDA/HIP only. + GGML_API struct ggml_tensor * ggml_ssm_conv_step( + struct ggml_context * ctx, + struct ggml_tensor * x, + struct ggml_tensor * c, + struct ggml_tensor * conv_state, + struct ggml_tensor * conv_input_out); // 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 @@ -2907,6 +2925,18 @@ extern "C" { struct ggml_tensor * tensor, bool skip_intermediate); + // dflash extension: let the kernel derive the gates from the raw + // projections instead of graph-side sigmoid/softplus ops: + // beta_val = sigmoid(beta_raw) + // g_val = exp(softplus(alpha_raw + dt_bias[h]) * A[h]) + // `g` then carries alpha_raw and `beta` carries beta_raw (both [1,H,T,S]); + // gate_ba is a contiguous f32 [2*H] tensor holding [dt_bias | A] + // (src[9], op_params[10] = 1). Only for the non-tree, non-KDA, + // non-SpecLA CUDA/HIP path. + GGML_API void ggml_gated_delta_net_set_raw_gates( + struct ggml_tensor * tensor, + struct ggml_tensor * gate_ba); + // dflash extension: tree-mode gated delta net for DDTree-style // speculative decoding verify. `parent_ids` is an int32 tensor of shape // [n_tokens, n_seqs] where entry [t, s] is the index within sequence s of diff --git a/server/deps/llama.cpp/ggml/src/ggml-cpu/ops.cpp b/server/deps/llama.cpp/ggml/src/ggml-cpu/ops.cpp index b10e8c75d..2c05b85f2 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cpu/ops.cpp +++ b/server/deps/llama.cpp/ggml/src/ggml-cpu/ops.cpp @@ -9332,6 +9332,8 @@ void ggml_compute_forward_flash_attn_back( static void ggml_compute_forward_ssm_conv_f32( const ggml_compute_params * params, ggml_tensor * dst) { + // dflash: the fused step mode (ggml_ssm_conv_step) is CUDA/HIP only + GGML_ASSERT(ggml_get_op_params_i32(dst, 0) == 0 && "ggml_ssm_conv_step is not supported on CPU"); const ggml_tensor * src0 = dst->src[0]; // conv_x const ggml_tensor * src1 = dst->src[1]; // conv1d.weight diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/fattn-vec.cuh b/server/deps/llama.cpp/ggml/src/ggml-cuda/fattn-vec.cuh index bcf1dd804..85a2af718 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/fattn-vec.cuh +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/fattn-vec.cuh @@ -534,7 +534,10 @@ void ggml_cuda_flash_attn_ext_vec_case_impl(ggml_backend_cuda_context & ctx, ggm const bool need_f16_K = type_K == GGML_TYPE_F16; const bool need_f16_V = type_V == GGML_TYPE_F16; constexpr size_t nbytes_shared = 0; - launch_fattn(ctx, dst, fattn_kernel, nwarps, nbytes_shared, D, need_f16_K, need_f16_V, false); + // The kernel walks the KV sequence in steps of nthreads (not D); telling + // launch_fattn so lets it split a short KV span (e.g. a 256-token window + // at head_dim 256) across two blocks per head instead of one. + launch_fattn(ctx, dst, fattn_kernel, nwarps, nbytes_shared, nthreads, need_f16_K, need_f16_V, false); } template 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 afd328070..64156d5bc 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 @@ -97,7 +97,9 @@ gated_delta_net_cuda(const float * q, int64_t sb3, const uint3 neqk1_magic, const uint3 rq3_magic, - float scale) { + float scale, + const float * gate_bias, // raw-gate mode: dt_bias[H], else nullptr + const float * gate_A) { // raw-gate mode: A[H] const uint32_t h_idx = blockIdx.x; const uint32_t sequence = blockIdx.y; // each warp owns one column, using warp-level primitives to reduce across rows @@ -196,7 +198,9 @@ gated_delta_net_cuda(const float * q, const float * beta_t = beta + gb_offset; const float * g_t = g + gb_offset * (KDA ? S_v : 1); - const float beta_val = *beta_t; + // raw-gate mode: beta = sigmoid(beta_raw); g = softplus(alpha_raw + bias) * A + const bool raw_gates = gate_bias != nullptr; + const float beta_val = raw_gates ? 1.0f / (1.0f + expf(-(*beta_t))) : *beta_t; // Cache k and q in registers float k_reg[rows_per_lane]; @@ -209,7 +213,12 @@ gated_delta_net_cuda(const float * q, } if constexpr (!KDA) { - const float g_val = expf(*g_t); + float g_log = *g_t; + if (raw_gates) { + const float a = g_log + gate_bias[h_idx]; + g_log = ((a > 20.0f) ? a : logf(1.0f + expf(a))) * gate_A[h_idx]; + } + const float g_val = expf(g_log); // kv[col] = (S^T @ k)[col] = sum_i S[i][col] * k[i] float kv_shard = 0.0f; @@ -291,7 +300,7 @@ gated_delta_net_cuda(const float * q, } } -template +template __global__ void __launch_bounds__(WARP_THREADS * 8, 2) gated_delta_net_cuda_grouped_cols(const float * q, const float * k, @@ -302,6 +311,7 @@ gated_delta_net_cuda_grouped_cols(const float * q, const int * active_slot_ids, float * dst, float * state_out, + const int * parent_ids, // TREE_MODE only; else ignored InterT * persist_inter, int64_t H, int64_t n_tokens, @@ -318,7 +328,9 @@ gated_delta_net_cuda_grouped_cols(const float * q, int64_t sb3, const uint3 neqk1_magic, const uint3 rq3_magic, - float scale) { + float scale, + const float * gate_bias, // raw-gate mode: dt_bias[H], else nullptr + const float * gate_A) { // raw-gate mode: A[H] static_assert(S_v == 128, "grouped GDN kernel is specialized for S_v=128"); static_assert(WIDTH == 16, "grouped GDN kernel expects 16-lane subgroups"); static_assert(COLS == 4, "grouped GDN kernel expects 4 columns per subgroup"); @@ -352,12 +364,16 @@ gated_delta_net_cuda_grouped_cols(const float * q, n_seqs, n_state_slots, physical_sequence, physical_state_offset); InterT * inter_states = nullptr; InterT * inter_base = nullptr; - if constexpr (WRITE_INTER) { + if constexpr (WRITE_INTER || TREE_MODE) { inter_states = persist_inter ? persist_inter : (InterT *)(dst + attn_score_elems + final_state_elems); inter_base = inter_states + (sequence * n_tokens * H + h_idx) * S_v * S_v; } + const int * parent_ids_seq = nullptr; + if constexpr (TREE_MODE) { + parent_ids_seq = parent_ids + sequence * n_tokens; + } const float * curr_state_seq = physical_sequence >= 0 ? curr_state + physical_state_offset @@ -378,6 +394,41 @@ gated_delta_net_cuda_grouped_cols(const float * q, } for (int t = 0; t < n_tokens; ++t) { + if constexpr (TREE_MODE) { + // DFS branch transition: this token continues from a state other + // than the previous token's. Reload the register shard from the + // parent's stored intermediate state (same-thread read-after-write + // on global memory, no barrier needed) or reset to the pre-block + // state for root-level siblings. + if (t > 0) { + const int parent_t = parent_ids_seq[t]; + if (parent_t == GGML_GDN_TREE_ROOT_PARENT) { +#pragma unroll + for (int c = 0; c < COLS; ++c) { + const int col = col_base + c; +#pragma unroll + for (int r = 0; r < rows_per_lane; ++r) { + const int row = r * WIDTH + lane; + state_shard[c][r] = curr_state_seq + ? curr_state_seq[col * S_v + row] + : 0.0f; + } + } + } else if (parent_t != t - 1) { + const InterT * parent_base = inter_states + + ((sequence * n_tokens + parent_t) * H + h_idx) * S_v * S_v; +#pragma unroll + for (int c = 0; c < COLS; ++c) { + const int col = col_base + c; +#pragma unroll + for (int r = 0; r < rows_per_lane; ++r) { + const int row = r * WIDTH + lane; + state_shard[c][r] = load_inter_state(parent_base, col * S_v + row); + } + } + } + } + } const float * q_t = q + iq3 * sq3 + t * sq2 + iq1 * sq1; const float * k_t = k + iq3 * sq3 + t * sq2 + iq1 * sq1; const float * v_t = v + sequence * sv3 + t * sv2 + h_idx * sv1; @@ -387,8 +438,16 @@ gated_delta_net_cuda_grouped_cols(const float * q, float g_val = 0.0f; float beta_val = 0.0f; if (threadIdx.x == 0) { - g_val = expf(g[gb_offset]); - beta_val = beta[gb_offset]; + if (gate_bias != nullptr) { + // raw-gate mode: g = exp(softplus(alpha_raw + bias) * A), beta = sigmoid(beta_raw) + const float a = g[gb_offset] + gate_bias[h_idx]; + const float sp = (a > 20.0f) ? a : logf(1.0f + expf(a)); + g_val = expf(sp * gate_A[h_idx]); + beta_val = 1.0f / (1.0f + expf(-beta[gb_offset])); + } else { + g_val = expf(g[gb_offset]); + beta_val = beta[gb_offset]; + } } g_val = __shfl_sync(0xffffffffU, g_val, 0); beta_val = __shfl_sync(0xffffffffU, beta_val, 0); @@ -455,7 +514,7 @@ gated_delta_net_cuda_grouped_cols(const float * q, } } - if constexpr (WRITE_INTER) { + if constexpr (WRITE_INTER || TREE_MODE) { #pragma unroll for (int c = 0; c < COLS; ++c) { const int col = col_base + c; @@ -497,7 +556,8 @@ static void launch_gated_delta_net( int64_t sv1, int64_t sv2, int64_t sv3, int64_t sb1, int64_t sb2, int64_t sb3, int64_t neqk1, int64_t rq3, - float scale, cudaStream_t stream) { + float scale, cudaStream_t stream, + const float * gate_bias = nullptr, const float * gate_A = nullptr) { //TODO: Add chunked kernel for even faster pre-fill const int warp_size = ggml_cuda_info().devices[ggml_cuda_get_device()].warp_size; const int num_warps = 4; @@ -521,23 +581,23 @@ static void launch_gated_delta_net( gated_delta_net_cuda<16, KDA, TREE_MODE, WRITE_INTER, InterT><<>>( q_d, k_d, v_d, g_d, b_d, s_d, active_slot_ids_d, dst_d, state_out_d, parent_ids_d, persist_inter_d, H, n_tokens, n_seqs, n_state_slots, sq1, sq2, sq3, sv1, sv2, sv3, - sb1, sb2, sb3, neqk1_magic, rq3_magic, scale); + sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, gate_bias, gate_A); break; case 32: gated_delta_net_cuda<32, KDA, TREE_MODE, WRITE_INTER, InterT><<>>( q_d, k_d, v_d, g_d, b_d, s_d, active_slot_ids_d, dst_d, state_out_d, parent_ids_d, persist_inter_d, H, n_tokens, n_seqs, n_state_slots, sq1, sq2, sq3, sv1, sv2, sv3, - sb1, sb2, sb3, neqk1_magic, rq3_magic, scale); + sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, gate_bias, gate_A); break; case 64: { gated_delta_net_cuda<64, KDA, TREE_MODE, WRITE_INTER, InterT><<>>( q_d, k_d, v_d, g_d, b_d, s_d, active_slot_ids_d, dst_d, state_out_d, parent_ids_d, persist_inter_d, H, n_tokens, n_seqs, n_state_slots, sq1, sq2, sq3, sv1, sv2, sv3, - sb1, sb2, sb3, neqk1_magic, rq3_magic, scale); + sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, gate_bias, gate_A); break; } case 128: { - if constexpr (!KDA && !TREE_MODE) { + if constexpr (!KDA) { if (use_grouped_cols && ((GGML_CUDA_CC_IS_NVIDIA(cc) && cc >= GGML_CUDA_CC_AMPERE) || GGML_CUDA_CC_IS_AMD(cc))) { @@ -549,35 +609,35 @@ static void launch_gated_delta_net( constexpr int groups_per_warp = 32 / width; dim3 grouped_grid_dims(H, n_seqs, (groups + column_groups_per_block * groups_per_warp - 1) / (column_groups_per_block * groups_per_warp)); dim3 grouped_block_dims(32, column_groups_per_block, 1); - gated_delta_net_cuda_grouped_cols<128, cols, width, 32, WRITE_INTER, InterT><<>>( - q_d, k_d, v_d, g_d, b_d, s_d, active_slot_ids_d, dst_d, state_out_d, persist_inter_d, H, + gated_delta_net_cuda_grouped_cols<128, cols, width, 32, TREE_MODE, WRITE_INTER, InterT><<>>( + q_d, k_d, v_d, g_d, b_d, s_d, active_slot_ids_d, dst_d, state_out_d, parent_ids_d, persist_inter_d, H, n_tokens, n_seqs, n_state_slots, sq1, sq2, sq3, sv1, sv2, sv3, - sb1, sb2, sb3, neqk1_magic, rq3_magic, scale); + sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, gate_bias, gate_A); } else if (warp_size == 64) { constexpr int groups_per_warp = 64 / width; dim3 grouped_grid_dims(H, n_seqs, (groups + column_groups_per_block * groups_per_warp - 1) / (column_groups_per_block * groups_per_warp)); dim3 grouped_block_dims(64, column_groups_per_block, 1); - gated_delta_net_cuda_grouped_cols<128, cols, width, 64, WRITE_INTER, InterT><<>>( - q_d, k_d, v_d, g_d, b_d, s_d, active_slot_ids_d, dst_d, state_out_d, persist_inter_d, H, + gated_delta_net_cuda_grouped_cols<128, cols, width, 64, TREE_MODE, WRITE_INTER, InterT><<>>( + q_d, k_d, v_d, g_d, b_d, s_d, active_slot_ids_d, dst_d, state_out_d, parent_ids_d, persist_inter_d, H, n_tokens, n_seqs, n_state_slots, sq1, sq2, sq3, sv1, sv2, sv3, - sb1, sb2, sb3, neqk1_magic, rq3_magic, scale); + sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, gate_bias, gate_A); } else { gated_delta_net_cuda<128, KDA, TREE_MODE, WRITE_INTER, InterT><<>>( q_d, k_d, v_d, g_d, b_d, s_d, active_slot_ids_d, dst_d, state_out_d, parent_ids_d, persist_inter_d, H, n_tokens, n_seqs, n_state_slots, sq1, sq2, sq3, sv1, sv2, sv3, - sb1, sb2, sb3, neqk1_magic, rq3_magic, scale); + sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, gate_bias, gate_A); } } else { gated_delta_net_cuda<128, KDA, TREE_MODE, WRITE_INTER, InterT><<>>( q_d, k_d, v_d, g_d, b_d, s_d, active_slot_ids_d, dst_d, state_out_d, parent_ids_d, persist_inter_d, H, n_tokens, n_seqs, n_state_slots, sq1, sq2, sq3, sv1, sv2, sv3, - sb1, sb2, sb3, neqk1_magic, rq3_magic, scale); + sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, gate_bias, gate_A); } } else { gated_delta_net_cuda<128, KDA, TREE_MODE, WRITE_INTER, InterT><<>>( q_d, k_d, v_d, g_d, b_d, s_d, active_slot_ids_d, dst_d, state_out_d, parent_ids_d, persist_inter_d, H, n_tokens, n_seqs, n_state_slots, sq1, sq2, sq3, sv1, sv2, sv3, - sb1, sb2, sb3, neqk1_magic, rq3_magic, scale); + sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, gate_bias, gate_A); } break; } @@ -912,6 +972,18 @@ void ggml_cuda_op_gated_delta_net(ggml_backend_cuda_context & ctx, ggml_tensor * const bool tree_mode = (parent_ids_d != nullptr); const bool skip_intermediate = ggml_get_op_params_i32(dst, 0) != 0; + // dflash raw-gate mode: src[9] = [dt_bias | A] (f32 [2H]); the kernel + // applies sigmoid / softplus+bias / A itself (see ggml_gated_delta_net_set_raw_gates). + const bool raw_gates = ggml_get_op_params_i32(dst, 10) != 0; + const float * gate_bias_d = nullptr; + const float * gate_A_d = nullptr; + if (raw_gates) { + GGML_ASSERT(dst->src[9] && dst->src[9]->type == GGML_TYPE_F32); + GGML_ASSERT(ggml_nelements(dst->src[9]) == 2*H); + GGML_ASSERT(!kda && !tree_mode && active_slot_ids_d == nullptr); + gate_bias_d = (const float *) dst->src[9]->data; + gate_A_d = gate_bias_d + H; + } const bool write_intermediate = tree_mode || !skip_intermediate || persist_inter_d != nullptr; // Macro to expand KDA × TREE_MODE × WRITE_INTER for a given InterT. @@ -924,34 +996,34 @@ void ggml_cuda_op_gated_delta_net(ggml_backend_cuda_context & ctx, ggml_tensor * launch_gated_delta_net( \ q_d, k_d, v_d, g_d, b_d, s_d, active_slot_ids_d, dst_d, state_out_d, parent_ids_d, persist_typed, \ S_v, H, n_tokens, n_seqs, n_state_slots, sq1, sq2, sq3, sv1, sv2, sv3, \ - sb1, sb2, sb3, neqk1, rq3, scale, stream); \ + sb1, sb2, sb3, neqk1, rq3, scale, stream, gate_bias_d, gate_A_d); \ } else if (write_intermediate) { \ launch_gated_delta_net( \ q_d, k_d, v_d, g_d, b_d, s_d, active_slot_ids_d, dst_d, state_out_d, nullptr, persist_typed, \ S_v, H, n_tokens, n_seqs, n_state_slots, sq1, sq2, sq3, sv1, sv2, sv3, \ - sb1, sb2, sb3, neqk1, rq3, scale, stream); \ + sb1, sb2, sb3, neqk1, rq3, scale, stream, gate_bias_d, gate_A_d); \ } else { \ launch_gated_delta_net( \ q_d, k_d, v_d, g_d, b_d, s_d, active_slot_ids_d, dst_d, state_out_d, nullptr, persist_typed, \ S_v, H, n_tokens, n_seqs, n_state_slots, sq1, sq2, sq3, sv1, sv2, sv3, \ - sb1, sb2, sb3, neqk1, rq3, scale, stream); \ + sb1, sb2, sb3, neqk1, rq3, scale, stream, gate_bias_d, gate_A_d); \ } \ } else { \ if (tree_mode) { \ launch_gated_delta_net( \ q_d, k_d, v_d, g_d, b_d, s_d, active_slot_ids_d, dst_d, state_out_d, parent_ids_d, persist_typed, \ S_v, H, n_tokens, n_seqs, n_state_slots, sq1, sq2, sq3, sv1, sv2, sv3, \ - sb1, sb2, sb3, neqk1, rq3, scale, stream); \ + sb1, sb2, sb3, neqk1, rq3, scale, stream, gate_bias_d, gate_A_d); \ } else if (write_intermediate) { \ launch_gated_delta_net( \ q_d, k_d, v_d, g_d, b_d, s_d, active_slot_ids_d, dst_d, state_out_d, nullptr, persist_typed, \ S_v, H, n_tokens, n_seqs, n_state_slots, sq1, sq2, sq3, sv1, sv2, sv3, \ - sb1, sb2, sb3, neqk1, rq3, scale, stream); \ + sb1, sb2, sb3, neqk1, rq3, scale, stream, gate_bias_d, gate_A_d); \ } else { \ launch_gated_delta_net( \ q_d, k_d, v_d, g_d, b_d, s_d, active_slot_ids_d, dst_d, state_out_d, nullptr, persist_typed, \ S_v, H, n_tokens, n_seqs, n_state_slots, sq1, sq2, sq3, sv1, sv2, sv3, \ - sb1, sb2, sb3, neqk1, rq3, scale, stream); \ + sb1, sb2, sb3, neqk1, rq3, scale, stream, gate_bias_d, gate_A_d); \ } \ } \ } while (0) 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 205964691..da9ee6072 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 @@ -473,7 +473,11 @@ const ggml_cuda_device_info & ggml_cuda_info() { // buffer pool for cuda (legacy) struct ggml_cuda_pool_leg : public ggml_cuda_pool { - static const int MAX_BUFFERS = 256; + // 1024 (upstream 256): LUCE_Q8_MEMO keeps one pooled q8_1 activation + // buffer per quantized matmul alive across a whole graph evaluation + // (~300 on a 64-layer hybrid), and a full pool falls back to freeing + // in-flight buffers with cudaFree. + static const int MAX_BUFFERS = 1024; int device; struct ggml_cuda_buffer { @@ -4365,6 +4369,49 @@ static bool ggml_cuda_can_fuse(const struct ggml_cgraph * cgraph, } } + // dflash: residual ADD + RMS_NORM + MUL. The add output stays live (it is + // the next residual), so this is a subgraph fusion with two outputs. + if (ops.size() == 3 && ops.begin()[0] == GGML_OP_ADD && ops.begin()[1] == GGML_OP_RMS_NORM && + ops.begin()[2] == GGML_OP_MUL) { + if (!ggml_can_fuse_subgraph(cgraph, node_idx, ops, { node_idx, node_idx + 2 })) { + return false; + } + const ggml_tensor * add = cgraph->nodes[node_idx]; + const ggml_tensor * rms = cgraph->nodes[node_idx + 1]; + const ggml_tensor * mul = cgraph->nodes[node_idx + 2]; + if (rms->src[0] != add) { + return false; + } + const ggml_tensor * w = nullptr; + if (mul->src[0] == rms) { + w = mul->src[1]; + } else if (mul->src[1] == rms) { + w = mul->src[0]; + } else { + return false; + } + const ggml_tensor * a = add->src[0]; + const ggml_tensor * b = add->src[1]; + if (a->type != GGML_TYPE_F32 || b->type != GGML_TYPE_F32 || w->type != GGML_TYPE_F32 || + add->type != GGML_TYPE_F32 || rms->type != GGML_TYPE_F32 || mul->type != GGML_TYPE_F32) { + return false; + } + if (!ggml_is_contiguous(a) || !ggml_is_contiguous(b) || !ggml_is_contiguous(w) || + !ggml_is_contiguous(add) || !ggml_is_contiguous(mul)) { + return false; + } + if (!ggml_are_same_shape(a, b) || !ggml_are_same_shape(a, add) || !ggml_are_same_shape(a, mul)) { + return false; + } + if (w->ne[0] != a->ne[0] || ggml_nelements(w) != a->ne[0]) { + return false; + } + if (ggml_backend_buft_is_cuda_split(a->buffer->buft) || ggml_backend_buft_is_cuda_split(b->buffer->buft)) { + return false; + } + return true; + } + if (!ggml_can_fuse(cgraph, node_idx, ops)) { return false; } @@ -4416,8 +4463,8 @@ 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 + if (ggml_get_op_params_i32(ssm_conv, 0) != 0) { + // the Specla (1) and dflash step (2) ssm_conv kernels apply SiLU themselves return false; } @@ -4969,6 +5016,12 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud continue; } + if (ggml_cuda_can_fuse(cgraph, i, { GGML_OP_ADD, GGML_OP_RMS_NORM, GGML_OP_MUL }, {})) { + ggml_cuda_op_add_rms_norm_mul_fused(*cuda_ctx, node, cgraph->nodes[i+1], cgraph->nodes[i+2]); + i += 2; + continue; + } + if (ggml_cuda_can_fuse(cgraph, i, { GGML_OP_RMS_NORM, GGML_OP_MUL, GGML_OP_ADD}, {})) { ggml_cuda_op_rms_norm_fused_add(*cuda_ctx, node, cgraph->nodes[i+1], cgraph->nodes[i+2]); i += 2; @@ -6216,9 +6269,10 @@ 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. + // op_params[0]: 1 = SpecLA heavy-light conv (x is [d_inner, n_tokens], + // kernel guards the final partial 128-channel block), 2 = dflash fused + // step mode (any channel count). + if (ggml_get_op_params_i32(op, 0) == 1 || ggml_get_op_params_i32(op, 0) == 2) { return true; } // assumes d_inner % threads == 0 diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/mmq.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/mmq.cu index 64a29b31f..65eeeb584 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/mmq.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/mmq.cu @@ -1,5 +1,6 @@ #include "common.cuh" #include "mmq.cuh" +#include "mmq_big.h" #include "quantize.cuh" #include "mmid.cuh" #include "rocmfp2_mix.cuh" @@ -41,12 +42,55 @@ private: } // namespace +// Big-tile dispatch (see mmq_big.h): the default instances for the dense +// hybrid types are 64x64 (GGML_CUDA_MMQ_SMALL_TILE, tuned for spec-decode +// verify widths); at prefill widths the narrow x-tile re-streams the weights, +// so wide batches take the 128x128 twin instances instead. RDNA4 only: the +// measurement is from gfx1201, and gfx1151 keeps its existing behavior. +// LUCE_MMQ_BIG_PREFILL=0 disables. +static bool lucebox_mmq_big_tile_take(const ggml_type type, const int64_t ncols_dst) { + static const bool enabled = []() { + const char * e = getenv("LUCE_MMQ_BIG_PREFILL"); + return !(e && e[0] == '0' && e[1] == '\0'); + }(); + // Measured crossover on gfx1201 (iq4_xs 17408x5120): small tile wins to + // N=64, tie at 128, big wins 16-18% at 512. Take big only where it is a + // clear win. + if (!enabled || ncols_dst < 256) { + return false; + } + const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; + if (!GGML_CUDA_CC_IS_RDNA4(cc)) { + return false; + } + switch (type) { + case GGML_TYPE_IQ4_XS: + case GGML_TYPE_Q4_K: + case GGML_TYPE_Q5_K: + case GGML_TYPE_Q6_K: + case GGML_TYPE_Q8_0: + return true; + default: + return false; + } +} + static void ggml_cuda_mul_mat_q_switch_type(ggml_backend_cuda_context & ctx, const mmq_args & args, cudaStream_t stream) { const bool is_mix_type = args.type_x == GGML_TYPE_Q2_1_ROCMFP2_MIX || args.type_x == GGML_TYPE_Q3_1_ROCMFP3_MIX; GGML_ASSERT(!is_mix_type || (args.mix_codebooks && args.mix_modes)); ++g_mmq_launch_count; + if (lucebox_mmq_big_tile_take(args.type_x, args.ncols_dst)) { + switch (args.type_x) { + case GGML_TYPE_IQ4_XS: mul_mat_q_case_big_iq4_xs(ctx, &args, stream); return; + case GGML_TYPE_Q4_K: mul_mat_q_case_big_q4_k (ctx, &args, stream); return; + case GGML_TYPE_Q5_K: mul_mat_q_case_big_q5_k (ctx, &args, stream); return; + case GGML_TYPE_Q6_K: mul_mat_q_case_big_q6_k (ctx, &args, stream); return; + case GGML_TYPE_Q8_0: mul_mat_q_case_big_q8_0 (ctx, &args, stream); return; + default: break; + } + } switch (args.type_x) { case GGML_TYPE_Q4_0: mul_mat_q_case(ctx, args, stream); diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/mmq.cuh b/server/deps/llama.cpp/ggml/src/ggml-cuda/mmq.cuh index c482ce65a..e75ac0e78 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/mmq.cuh +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/mmq.cuh @@ -112,9 +112,16 @@ struct tile_x_sizes { int sc; }; -// RDNA uses 128x128, eight-warp MMQ tiles by default. Q4_K narrows the row -// dimension to 128x64, while ROCmFPX uses 64x64 four-warp tiles. Their -// unpacking pressure makes the smaller tiles faster on gfx1151. +// RDNA uses 128x128, eight-warp MMQ tiles by default. Template instances +// compiled with GGML_CUDA_MMQ_SMALL_TILE use 64x64, four-warp tiles: +// - ROCmFPX formats: their unpacking pressure makes the smaller tile faster +// on gfx1151; +// - IQ4_XS / Q5_K / Q6_K / Q8_0 (dense hybrid targets): at spec-decode +// verify widths (N<=16) the 128-row tile leaves a 5120-row projection +// with only 40 blocks on a 64-CU gfx1201; the small tile measured +// +12-23% there (mmq_probe) at the cost of ~8% prefill throughput. +// Q4_K instead narrows only the row dimension to 128x64 (LUCEBOX_RDNA_MMQ_Y), +// the shape measured best for packed concurrent prefill on gfx1151. #ifndef LUCEBOX_RDNA_MMQ_TILE_OVERRIDE #define LUCEBOX_RDNA_MMQ_TILE_OVERRIDE 1 #endif @@ -127,7 +134,7 @@ struct tile_x_sizes { static int get_mmq_x_max_host(const int cc) { if (LUCEBOX_RDNA_TILE_HOST(cc)) { -#if defined(GGML_CUDA_ROCMFPX_MMQ_TILE) +#if defined(GGML_CUDA_MMQ_SMALL_TILE) return 64; #else return 128; @@ -144,7 +151,7 @@ static int get_mmq_x_max_host(const int cc) { static constexpr __device__ int get_mmq_x_max_device() { #if LUCEBOX_RDNA_TILE_DEVICE -#if defined(GGML_CUDA_ROCMFPX_MMQ_TILE) +#if defined(GGML_CUDA_MMQ_SMALL_TILE) return 64; #else return 128; @@ -174,7 +181,7 @@ static constexpr __device__ int get_mmq_x_max_device() { static int get_mmq_y_host(const int cc) { if (LUCEBOX_RDNA_TILE_HOST(cc)) { -#if defined(GGML_CUDA_ROCMFPX_MMQ_TILE) +#if defined(GGML_CUDA_MMQ_SMALL_TILE) return 64; #elif defined(LUCEBOX_RDNA_MMQ_Y) return LUCEBOX_RDNA_MMQ_Y; @@ -196,7 +203,7 @@ static constexpr __device__ int get_iter_k([[maybe_unused]] const ggml_type type static constexpr __device__ int get_mmq_y_device() { #if LUCEBOX_RDNA_TILE_DEVICE -#if defined(GGML_CUDA_ROCMFPX_MMQ_TILE) +#if defined(GGML_CUDA_MMQ_SMALL_TILE) return 64; #elif defined(LUCEBOX_RDNA_MMQ_Y) return LUCEBOX_RDNA_MMQ_Y; @@ -361,7 +368,7 @@ static constexpr __device__ int mmq_get_granularity_device(const int /*mmq_x*/) #if defined(GGML_USE_HIP) static int mmq_get_nwarps_host(const int cc, const int warp_size) { if (LUCEBOX_RDNA_TILE_HOST(cc)) { -#if defined(GGML_CUDA_ROCMFPX_MMQ_TILE) +#if defined(GGML_CUDA_MMQ_SMALL_TILE) return 4; #elif defined(LUCEBOX_RDNA_MMQ_Y) return 4; @@ -379,7 +386,7 @@ static int mmq_get_nwarps_host(const int /*cc*/, const int warp_size) { static constexpr __device__ int mmq_get_nwarps_device() { #if LUCEBOX_RDNA_TILE_DEVICE -#if defined(GGML_CUDA_ROCMFPX_MMQ_TILE) +#if defined(GGML_CUDA_MMQ_SMALL_TILE) return 4; #elif defined(LUCEBOX_RDNA_MMQ_Y) return 4; @@ -4333,7 +4340,7 @@ template #if defined(GGML_USE_HIP) // RDNA4 is compute-bound on MMQ (WMMA path); allow compiler to use more VGPRs // (minBlocks=1 matches NVIDIA Volta+ behavior and reduces register spilling). -#if defined(RDNA4) && !defined(GGML_CUDA_ROCMFPX_MMQ_TILE) +#if defined(RDNA4) && !defined(GGML_CUDA_MMQ_SMALL_TILE) __launch_bounds__(ggml_cuda_get_physical_warp_size()*mmq_get_nwarps_device(), 1) #elif defined(RDNA3) || defined(RDNA2) || defined(CDNA) || defined(GCN) __launch_bounds__(ggml_cuda_get_physical_warp_size()*mmq_get_nwarps_device(), 2) @@ -4895,6 +4902,14 @@ void mul_mat_q_case(ggml_backend_cuda_context & ctx, const mmq_args & args, cuda if (mmq_x % granularity != 0 || mmq_get_nbytes_shared(mmq_x, mmq_y, cc, warp_size, nwarps) > smpbo) { continue; } +#if defined(GGML_CUDA_MMQ_SMALL_TILE) + // The 64-row/4-warp tile is pathological at mmq_x == 32 on gfx1201 + // (17408x5120 IQ4_XS: N=16 443 GB/s, N=24..32 180 GB/s, N=48 315 GB/s + // in mmq_probe); a wider tile with more padding is still faster. + if (LUCEBOX_RDNA_TILE_HOST(cc) && mmq_x == 32) { + continue; + } +#endif const int ntiles_x = (args.ncols_max + mmq_x - 1) / mmq_x; diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/mmq_big.h b/server/deps/llama.cpp/ggml/src/ggml-cuda/mmq_big.h new file mode 100644 index 000000000..4c8671b4b --- /dev/null +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/mmq_big.h @@ -0,0 +1,25 @@ +// Bridges into the big-tile (128x128, 8-warp) MMQ instances that coexist with +// the GGML_CUDA_MMQ_SMALL_TILE (64x64, 4-warp) default instances on RDNA. +// +// The tile shape is baked into every device/host constexpr in mmq.cuh via +// macros, so one TU can only hold one shape. The *-big.cu template instances +// re-include mmq.cuh inside `namespace lucebox_mmq_big` with no tile macro +// defined (the upstream 128x128 RDNA default), which gives the second shape +// distinct symbols. `args` is the caller's ::mmq_args passed as void const *: +// the namespaced struct is textually identical (its layout does not depend on +// the tile macros), the bridge casts it back. +// +// Why: at spec-decode verify widths (N <= 32) the 64-row tile measured +// +12-23% (grid occupancy on a 64-CU gfx1201), but at prefill widths the +// narrow x-tile re-streams the weights (+12% MMQ time at N = 512). The +// runtime dispatch in mmq.cu picks per shape and keeps both wins. +#pragma once + +#include "common.cuh" + +// Defined in template-instances/mmq-instance--big.cu. +void mul_mat_q_case_big_iq4_xs(ggml_backend_cuda_context & ctx, const void * args, cudaStream_t stream); +void mul_mat_q_case_big_q4_k (ggml_backend_cuda_context & ctx, const void * args, cudaStream_t stream); +void mul_mat_q_case_big_q5_k (ggml_backend_cuda_context & ctx, const void * args, cudaStream_t stream); +void mul_mat_q_case_big_q6_k (ggml_backend_cuda_context & ctx, const void * args, cudaStream_t stream); +void mul_mat_q_case_big_q8_0 (ggml_backend_cuda_context & ctx, const void * args, cudaStream_t stream); diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/norm.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/norm.cu index ef98f675a..696a6f441 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/norm.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/norm.cu @@ -150,6 +150,57 @@ static __global__ void rms_norm_f32(const float * x, } } +// dflash: residual add fused into the following rms_norm * weight. +// sum = a + b (written to sum_out; it is the next residual) +// dst = rms_norm(sum) * w +// All of a, b, sum_out, dst are contiguous [ncols, R]; w is [ncols]. +template +static __global__ void add_rms_norm_mul_f32(const float * __restrict__ a, + const float * __restrict__ b, + float * __restrict__ sum_out, + float * __restrict__ dst, + const float * __restrict__ w, + const int ncols, + const float eps) { + const int64_t row = blockIdx.x; + const int tid = threadIdx.x; + + a += row * ncols; + b += row * ncols; + sum_out += row * ncols; + dst += row * ncols; + + float tmp = 0.0f; + for (int col = tid; col < ncols; col += block_size) { + const float s = a[col] + b[col]; + sum_out[col] = s; + tmp += s * s; + } + + extern __shared__ float s_sum[]; + tmp = block_reduce(tmp, s_sum); + + const float mean = tmp / ncols; + const float scale = rsqrtf(mean + eps); + + for (int col = tid; col < ncols; col += block_size) { + dst[col] = scale * sum_out[col] * w[col]; + } +} + +static void add_rms_norm_mul_f32_cuda(const float * a, const float * b, float * sum_out, float * dst, + const float * w, const int ncols, const int64_t nrows, + const float eps, cudaStream_t stream) { + const dim3 blocks_num(nrows, 1, 1); + if (ncols < 1024) { + const dim3 block_dims(256, 1, 1); + add_rms_norm_mul_f32<256><<>>(a, b, sum_out, dst, w, ncols, eps); + } else { + const dim3 block_dims(1024, 1, 1); + add_rms_norm_mul_f32<1024><<>>(a, b, sum_out, dst, w, ncols, eps); + } +} + template static __global__ void rms_norm_back_f32( const float * grad, const float * xf, float * dst, const int ncols, const float eps) { @@ -533,6 +584,36 @@ void ggml_cuda_op_rms_norm_fused(ggml_backend_cuda_context & ctx, ggml_tensor * eps, stream); } +// dflash: ADD (residual) + RMS_NORM + MUL in one launch. `add_tensor` is the +// residual add node (its output is materialized), `rms_tensor` is elided, +// `mul_tensor` receives the normalized * weight result. +void ggml_cuda_op_add_rms_norm_mul_fused(ggml_backend_cuda_context & ctx, + ggml_tensor * add_tensor, + ggml_tensor * rms_tensor, + ggml_tensor * mul_tensor) { + const ggml_tensor * a = add_tensor->src[0]; + const ggml_tensor * b = add_tensor->src[1]; + const ggml_tensor * w = (mul_tensor->src[0] == rms_tensor) ? mul_tensor->src[1] : mul_tensor->src[0]; + + float eps = 0.0f; + memcpy(&eps, rms_tensor->op_params, sizeof(float)); + GGML_ASSERT(eps >= 0.0f); + + GGML_ASSERT(a->type == GGML_TYPE_F32 && b->type == GGML_TYPE_F32 && w->type == GGML_TYPE_F32); + GGML_ASSERT(add_tensor->type == GGML_TYPE_F32 && mul_tensor->type == GGML_TYPE_F32); + GGML_ASSERT(ggml_is_contiguous(a) && ggml_is_contiguous(b) && ggml_is_contiguous(w)); + GGML_ASSERT(ggml_is_contiguous(add_tensor) && ggml_is_contiguous(mul_tensor)); + GGML_ASSERT(ggml_are_same_shape(a, b) && ggml_are_same_shape(a, add_tensor) && ggml_are_same_shape(a, mul_tensor)); + GGML_ASSERT(w->ne[0] == a->ne[0] && ggml_nelements(w) == a->ne[0]); + + const int ncols = (int) a->ne[0]; + const int64_t nrows = ggml_nrows(a); + + add_rms_norm_mul_f32_cuda((const float *) a->data, (const float *) b->data, + (float *) add_tensor->data, (float *) mul_tensor->data, + (const float *) w->data, ncols, nrows, eps, ctx.stream()); +} + void ggml_cuda_op_rms_norm_fused_add(ggml_backend_cuda_context & ctx, ggml_tensor * dst, ggml_tensor * mul_tensor, diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/norm.cuh b/server/deps/llama.cpp/ggml/src/ggml-cuda/norm.cuh index a74f63767..6313a98ce 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/norm.cuh +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/norm.cuh @@ -16,3 +16,6 @@ void ggml_cuda_op_rms_norm_fused_add(ggml_backend_cuda_context & ctx, void ggml_cuda_op_rms_norm_back(ggml_backend_cuda_context & ctx, ggml_tensor * dst); void ggml_cuda_op_l2_norm(ggml_backend_cuda_context & ctx, ggml_tensor * dst); + +// dflash: residual ADD + RMS_NORM + MUL fusion (see norm.cu) +void ggml_cuda_op_add_rms_norm_mul_fused(ggml_backend_cuda_context & ctx, ggml_tensor * add_tensor, ggml_tensor * rms_tensor, ggml_tensor * mul_tensor); diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/rope.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/rope.cu index 9cc4daf3a..e9ddffce0 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/rope.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/rope.cu @@ -23,10 +23,23 @@ struct mrope_sections { // every config with freq_scale != 1.0 or freq_factor != 1.0 (freq_factor is // applied by the callers as theta_base/freq_factor before rope_yarn()). static __device__ __forceinline__ double rope_theta_fp64(int32_t p, float theta_scale, int exp_int) { - // Dim 0: theta_scale^0 == 1 exactly. Skip pow (costly on Turing). - return (exp_int == 0) - ? (double)p - : (double)p * pow((double)theta_scale, (double)exp_int); + // Dim 0: theta_scale^0 == 1 exactly. + if (exp_int == 0) { + return (double)p; + } + // Binary exponentiation instead of pow(): the libcall dominated the whole + // rope kernel on RDNA4 (692 us vs 76 us per launch at n_tokens=512). Seven + // double multiplies keep the large-freq_base precision (the entire point + // of the fp64 path) to within 1 ulp of pow(). + double base = (double)theta_scale; + double r = 1.0; + int e = exp_int; + while (e) { + if (e & 1) { r *= base; } + base *= base; + e >>= 1; + } + return (double)p * r; } static __device__ float rope_yarn_ramp(const float low, const float high, const int i0) { 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 dc70cb9ea..9538e3414 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,6 +244,105 @@ static void ssm_conv_f32_cuda(const float * src0, const float * src1, const int } } +// dflash: fused conv step (see ggml_ssm_conv_step). One thread per channel +// walks the token loop with the K-1 history in registers, writes silu(conv), +// the optional rollback window and the new history in place. +template +static __global__ void ssm_conv_step_f32(const float * __restrict__ x, const int x_nb1, const int x_nb2, + const float * __restrict__ w, const int w_nb1, + float * state, const int st_nb1, const int st_nb2, + float * __restrict__ y, const int y_nb1, const int y_nb2, + float * ci, const int ci_nb1, const int ci_nb2, + const int C, const int T) { + const int c = blockIdx.x * blockDim.x + threadIdx.x; + const int s = blockIdx.y; + if (c >= C) return; + + const float * xs = (const float *) ((const char *) x + (size_t) s * x_nb2) + c; + float * st = (float *) ((char *) state + (size_t) s * st_nb2 + (size_t) c * st_nb1); + float * ys = (float *) ((char *) y + (size_t) s * y_nb2) + c; + float * cs = ci ? (float *) ((char *) ci + (size_t) s * ci_nb2 + (size_t) c * ci_nb1) : nullptr; + const float * wc = (const float *) ((const char *) w + (size_t) c * w_nb1); + + const int xs_stride = x_nb1 / sizeof(float); + const int ys_stride = y_nb1 / sizeof(float); + + float wt[K]; + float win[K]; // oldest first; win[K-1] is the current input +#pragma unroll + for (int k = 0; k < K; k++) { + wt[k] = wc[k]; + } +#pragma unroll + for (int j = 0; j < K - 1; j++) { + win[j] = st[j]; + if (cs) cs[j] = win[j]; + } + for (int t = 0; t < T; t++) { + const float xt = xs[(size_t) t * xs_stride]; + win[K - 1] = xt; + float acc = 0.0f; +#pragma unroll + for (int k = 0; k < K; k++) { + acc += win[k] * wt[k]; + } + ys[(size_t) t * ys_stride] = ggml_cuda_op_silu_single(acc); + if (cs) cs[K - 1 + t] = xt; +#pragma unroll + for (int j = 0; j < K - 1; j++) { + win[j] = win[j + 1]; + } + } +#pragma unroll + for (int j = 0; j < K - 1; j++) { + st[j] = win[j]; + } +} + +static void ggml_cuda_op_ssm_conv_step(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { + const ggml_tensor * x = dst->src[0]; + const ggml_tensor * w = dst->src[1]; + ggml_tensor * st = dst->src[2]; + ggml_tensor * ci = dst->src[3]; + + const int K = (int) w->ne[0]; + const int C = (int) w->ne[1]; + const int T = (int) dst->ne[1]; + const int S = (int) dst->ne[2]; + + GGML_ASSERT(x->type == GGML_TYPE_F32 && w->type == GGML_TYPE_F32 && st->type == GGML_TYPE_F32); + GGML_ASSERT(x->nb[0] == sizeof(float)); + GGML_ASSERT(w->nb[0] == sizeof(float)); + GGML_ASSERT(st->nb[0] == sizeof(float) && st->nb[1] == (size_t) (K - 1) * sizeof(float)); + GGML_ASSERT(dst->nb[0] == sizeof(float)); + if (ci) { + GGML_ASSERT(ci->type == GGML_TYPE_F32 && ci->nb[0] == sizeof(float)); + GGML_ASSERT(ci->ne[0] >= K - 1 + T); + } + + const int threads = 256; + const dim3 blocks((C + threads - 1) / threads, S, 1); + cudaStream_t stream = ctx.stream(); + + auto launch = [&](auto KK) { + constexpr int kK = decltype(KK)::value; + ssm_conv_step_f32<<>>( + (const float *) x->data, (int) x->nb[1], (int) x->nb[2], + (const float *) w->data, (int) w->nb[1], + (float *) st->data, (int) st->nb[1], (int) st->nb[2], + (float *) dst->data, (int) dst->nb[1], (int) dst->nb[2], + ci ? (float *) ci->data : nullptr, ci ? (int) ci->nb[1] : 0, ci ? (int) ci->nb[2] : 0, + C, T); + }; + switch (K) { + case 3: launch(std::integral_constant{}); break; + case 4: launch(std::integral_constant{}); break; + case 5: launch(std::integral_constant{}); break; + case 9: launch(std::integral_constant{}); break; + default: GGML_ABORT("ssm_conv_step only supports kernel sizes 3, 4, 5, 9."); + } +} + template static __global__ void ssm_conv_specla_hld_f32( const float * __restrict__ x, // [d_inner, n_t] @@ -381,6 +480,13 @@ static void ssm_conv_specla_hld_cuda(ggml_backend_cuda_context & ctx, } void ggml_cuda_op_ssm_conv(ggml_backend_cuda_context & ctx, ggml_tensor * dst, ggml_tensor * silu_dst) { + // dflash: fused step mode (silu already applied by the kernel); op_params[0] == 2 + if (ggml_get_op_params_i32(dst, 0) == 2) { + GGML_ASSERT(silu_dst == nullptr); + ggml_cuda_op_ssm_conv_step(ctx, dst); + return; + } + // SpecLA heavy-light conv; op_params[0] == 1 if (ggml_get_op_params_i32(dst, 0) == 1) { GGML_ASSERT(silu_dst == nullptr); ssm_conv_specla_hld_cuda(ctx, dst); diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/generate_cu_files.py b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/generate_cu_files.py index f87396f5a..1d8529934 100755 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/generate_cu_files.py +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/generate_cu_files.py @@ -102,12 +102,56 @@ def get_short_name(long_quant_name): "GGML_TYPE_Q2_1_ROCMFP2_MIX", "GGML_TYPE_Q3_0_ROCMFPX", "GGML_TYPE_Q3_1_ROCMFP3_MIX", + # Dense hybrid (Qwen3.5/3.8) verify widths N<=16 on gfx1201: the + # 128-row tile leaves a 5120-row projection with only 40 blocks; + # 64x64/4-warp tiles measured +12-23% on those shapes (mmq_probe). + "GGML_TYPE_IQ4_XS", + "GGML_TYPE_Q4_K", + "GGML_TYPE_Q5_K", + "GGML_TYPE_Q6_K", + "GGML_TYPE_Q8_0", }: - guard = "#define GGML_CUDA_ROCMFPX_MMQ_TILE 1\n" + guard = "#define GGML_CUDA_MMQ_SMALL_TILE 1\n" if type == "GGML_TYPE_Q4_K": guard = "#define LUCEBOX_RDNA_MMQ_Y 64\n" f.write(SOURCE_MMQ.format(type=type, guard=guard)) +BIG_TILE_TYPES = [ + "GGML_TYPE_IQ4_XS", "GGML_TYPE_Q4_K", "GGML_TYPE_Q5_K", + "GGML_TYPE_Q6_K", "GGML_TYPE_Q8_0", +] + +SOURCE_MMQ_BIG = """// This file has been autogenerated by generate_cu_files.py, do not edit manually. +// Big-tile (128x128) twin of mmq-instance-{name}.cu: no tile macro, so mmq.cuh +// compiles with the upstream RDNA default shape. Namespaced to coexist with +// the small-tile instance (see mmq_big.h). + +#include "../common.cuh" +#include "../vecdotq.cuh" +#include "../mma.cuh" + +#include +#include + +namespace lucebox_mmq_big {{ +#include "../mmq.cuh" + +DECL_MMQ_CASE({type}); +}} // namespace lucebox_mmq_big + +#include "../mmq_big.h" + +void mul_mat_q_case_big_{name}(ggml_backend_cuda_context & ctx, const void * args, cudaStream_t stream) {{ + lucebox_mmq_big::mul_mat_q_case<{type}>( + ctx, *(const lucebox_mmq_big::mmq_args *) args, stream); +}} +""" + +for type in BIG_TILE_TYPES: + name = type.replace("GGML_TYPE_", "").lower() + with open(f"mmq-instance-{name}-big.cu", "w") as f: + f.write(SOURCE_MMQ_BIG.format(type=type, name=name)) + for type in range(1, 17): with open(f"mmf-instance-ncols_{type}.cu", "w") as f: f.write(SOURCE_MMF.format(type=type)) diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-iq4_xs-big.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-iq4_xs-big.cu new file mode 100644 index 000000000..8d5c51e46 --- /dev/null +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-iq4_xs-big.cu @@ -0,0 +1,24 @@ +// This file has been autogenerated by generate_cu_files.py, do not edit manually. +// Big-tile (128x128) twin of mmq-instance-iq4_xs.cu: no tile macro, so mmq.cuh +// compiles with the upstream RDNA default shape. Namespaced to coexist with +// the small-tile instance (see mmq_big.h). + +#include "../common.cuh" +#include "../vecdotq.cuh" +#include "../mma.cuh" + +#include +#include + +namespace lucebox_mmq_big { +#include "../mmq.cuh" + +DECL_MMQ_CASE(GGML_TYPE_IQ4_XS); +} // namespace lucebox_mmq_big + +#include "../mmq_big.h" + +void mul_mat_q_case_big_iq4_xs(ggml_backend_cuda_context & ctx, const void * args, cudaStream_t stream) { + lucebox_mmq_big::mul_mat_q_case( + ctx, *(const lucebox_mmq_big::mmq_args *) args, stream); +} diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-iq4_xs.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-iq4_xs.cu index 1eb3b7430..5e2a1127a 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-iq4_xs.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-iq4_xs.cu @@ -1,5 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. +#define GGML_CUDA_MMQ_SMALL_TILE 1 #include "../mmq.cuh" DECL_MMQ_CASE(GGML_TYPE_IQ4_XS); diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q2_0_rocmfp2.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q2_0_rocmfp2.cu index 8221e1d1e..b00cd9a0c 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q2_0_rocmfp2.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q2_0_rocmfp2.cu @@ -1,6 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#define GGML_CUDA_ROCMFPX_MMQ_TILE 1 +#define GGML_CUDA_MMQ_SMALL_TILE 1 #include "../mmq.cuh" DECL_MMQ_CASE(GGML_TYPE_Q2_0_ROCMFP2); diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q2_1_rocmfp2_mix.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q2_1_rocmfp2_mix.cu index 647b4572f..f73033e33 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q2_1_rocmfp2_mix.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q2_1_rocmfp2_mix.cu @@ -1,6 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#define GGML_CUDA_ROCMFPX_MMQ_TILE 1 +#define GGML_CUDA_MMQ_SMALL_TILE 1 #include "../mmq.cuh" DECL_MMQ_CASE(GGML_TYPE_Q2_1_ROCMFP2_MIX); diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q3_0_rocmfpx.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q3_0_rocmfpx.cu index 2380af75c..486782982 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q3_0_rocmfpx.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q3_0_rocmfpx.cu @@ -1,6 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#define GGML_CUDA_ROCMFPX_MMQ_TILE 1 +#define GGML_CUDA_MMQ_SMALL_TILE 1 #include "../mmq.cuh" DECL_MMQ_CASE(GGML_TYPE_Q3_0_ROCMFPX); diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q3_1_rocmfp3_mix.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q3_1_rocmfp3_mix.cu index 1873e073f..92197f871 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q3_1_rocmfp3_mix.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q3_1_rocmfp3_mix.cu @@ -1,6 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#define GGML_CUDA_ROCMFPX_MMQ_TILE 1 +#define GGML_CUDA_MMQ_SMALL_TILE 1 #include "../mmq.cuh" DECL_MMQ_CASE(GGML_TYPE_Q3_1_ROCMFP3_MIX); diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q4_0_rocmfp4_fast.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q4_0_rocmfp4_fast.cu index 94a2bb0f5..92cb4653d 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q4_0_rocmfp4_fast.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q4_0_rocmfp4_fast.cu @@ -1,6 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#define GGML_CUDA_ROCMFPX_MMQ_TILE 1 +#define GGML_CUDA_MMQ_SMALL_TILE 1 #include "../mmq.cuh" DECL_MMQ_CASE(GGML_TYPE_Q4_0_ROCMFP4_FAST); diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q4_k-big.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q4_k-big.cu new file mode 100644 index 000000000..f750f36cb --- /dev/null +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q4_k-big.cu @@ -0,0 +1,24 @@ +// This file has been autogenerated by generate_cu_files.py, do not edit manually. +// Big-tile (128x128) twin of mmq-instance-q4_k.cu: no tile macro, so mmq.cuh +// compiles with the upstream RDNA default shape. Namespaced to coexist with +// the small-tile instance (see mmq_big.h). + +#include "../common.cuh" +#include "../vecdotq.cuh" +#include "../mma.cuh" + +#include +#include + +namespace lucebox_mmq_big { +#include "../mmq.cuh" + +DECL_MMQ_CASE(GGML_TYPE_Q4_K); +} // namespace lucebox_mmq_big + +#include "../mmq_big.h" + +void mul_mat_q_case_big_q4_k(ggml_backend_cuda_context & ctx, const void * args, cudaStream_t stream) { + lucebox_mmq_big::mul_mat_q_case( + ctx, *(const lucebox_mmq_big::mmq_args *) args, stream); +} diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q5_k-big.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q5_k-big.cu new file mode 100644 index 000000000..c414d94de --- /dev/null +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q5_k-big.cu @@ -0,0 +1,24 @@ +// This file has been autogenerated by generate_cu_files.py, do not edit manually. +// Big-tile (128x128) twin of mmq-instance-q5_k.cu: no tile macro, so mmq.cuh +// compiles with the upstream RDNA default shape. Namespaced to coexist with +// the small-tile instance (see mmq_big.h). + +#include "../common.cuh" +#include "../vecdotq.cuh" +#include "../mma.cuh" + +#include +#include + +namespace lucebox_mmq_big { +#include "../mmq.cuh" + +DECL_MMQ_CASE(GGML_TYPE_Q5_K); +} // namespace lucebox_mmq_big + +#include "../mmq_big.h" + +void mul_mat_q_case_big_q5_k(ggml_backend_cuda_context & ctx, const void * args, cudaStream_t stream) { + lucebox_mmq_big::mul_mat_q_case( + ctx, *(const lucebox_mmq_big::mmq_args *) args, stream); +} diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q5_k.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q5_k.cu index a2e90ffd5..7cf43f75e 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q5_k.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q5_k.cu @@ -1,5 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. +#define GGML_CUDA_MMQ_SMALL_TILE 1 #include "../mmq.cuh" DECL_MMQ_CASE(GGML_TYPE_Q5_K); diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q6_k-big.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q6_k-big.cu new file mode 100644 index 000000000..eb43c668b --- /dev/null +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q6_k-big.cu @@ -0,0 +1,24 @@ +// This file has been autogenerated by generate_cu_files.py, do not edit manually. +// Big-tile (128x128) twin of mmq-instance-q6_k.cu: no tile macro, so mmq.cuh +// compiles with the upstream RDNA default shape. Namespaced to coexist with +// the small-tile instance (see mmq_big.h). + +#include "../common.cuh" +#include "../vecdotq.cuh" +#include "../mma.cuh" + +#include +#include + +namespace lucebox_mmq_big { +#include "../mmq.cuh" + +DECL_MMQ_CASE(GGML_TYPE_Q6_K); +} // namespace lucebox_mmq_big + +#include "../mmq_big.h" + +void mul_mat_q_case_big_q6_k(ggml_backend_cuda_context & ctx, const void * args, cudaStream_t stream) { + lucebox_mmq_big::mul_mat_q_case( + ctx, *(const lucebox_mmq_big::mmq_args *) args, stream); +} diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q6_k.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q6_k.cu index 470938fef..8bc6b7434 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q6_k.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q6_k.cu @@ -1,5 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. +#define GGML_CUDA_MMQ_SMALL_TILE 1 #include "../mmq.cuh" DECL_MMQ_CASE(GGML_TYPE_Q6_K); diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q8_0-big.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q8_0-big.cu new file mode 100644 index 000000000..3b19f33ec --- /dev/null +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q8_0-big.cu @@ -0,0 +1,24 @@ +// This file has been autogenerated by generate_cu_files.py, do not edit manually. +// Big-tile (128x128) twin of mmq-instance-q8_0.cu: no tile macro, so mmq.cuh +// compiles with the upstream RDNA default shape. Namespaced to coexist with +// the small-tile instance (see mmq_big.h). + +#include "../common.cuh" +#include "../vecdotq.cuh" +#include "../mma.cuh" + +#include +#include + +namespace lucebox_mmq_big { +#include "../mmq.cuh" + +DECL_MMQ_CASE(GGML_TYPE_Q8_0); +} // namespace lucebox_mmq_big + +#include "../mmq_big.h" + +void mul_mat_q_case_big_q8_0(ggml_backend_cuda_context & ctx, const void * args, cudaStream_t stream) { + lucebox_mmq_big::mul_mat_q_case( + ctx, *(const lucebox_mmq_big::mmq_args *) args, stream); +} diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q8_0.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q8_0.cu index 974477bbb..fb8fcf911 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q8_0.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q8_0.cu @@ -1,5 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. +#define GGML_CUDA_MMQ_SMALL_TILE 1 #include "../mmq.cuh" DECL_MMQ_CASE(GGML_TYPE_Q8_0); diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/vecdotq.cuh b/server/deps/llama.cpp/ggml/src/ggml-cuda/vecdotq.cuh index f8dc4335d..7736d7d67 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/vecdotq.cuh +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/vecdotq.cuh @@ -29,6 +29,17 @@ static __device__ __forceinline__ int get_int_b4(const void * x, const int & i32 return ((const int *) x)[i32]; // assume at least 4 byte alignment } +// Non-temporal variant for weight streams: decode reads every weight byte +// exactly once per token, so caching them evicts the activations/KV that +// other kernels reuse. HIP lowers this to sc0/sc1 (bypass) load hints. +static __device__ __forceinline__ int get_int_b4_nt(const void * x, const int & i32) { +#if defined(GGML_USE_HIP) + return __builtin_nontemporal_load(((const int *) x) + i32); +#else + return ((const int *) x)[i32]; +#endif +} + // q4 contains 8 indices with 4 bit each. // This function selects those bytes from table that are at those indices and returns them as int2. // The first int contains the bytes with even indices in q4, the second int contains the bytes with odd indices in q4. @@ -1608,7 +1619,7 @@ static __device__ __forceinline__ float vec_dot_iq4_xs_q8_1( int sumi = 0; #pragma unroll for (int j = 0; j < 4; ++j) { - const int aux_q4 = get_int_b4(bq4->qs, iqs + j); + const int aux_q4 = get_int_b4_nt(bq4->qs, iqs + j); const int2 v = get_int_from_table_16(aux_q4, kvalues_iq4nl); const int u0 = get_int_b4(bq8_1[iqs/4].qs, j + 0); diff --git a/server/deps/llama.cpp/ggml/src/ggml.c b/server/deps/llama.cpp/ggml/src/ggml.c index 74780a434..d6c889dad 100644 --- a/server/deps/llama.cpp/ggml/src/ggml.c +++ b/server/deps/llama.cpp/ggml/src/ggml.c @@ -5891,6 +5891,53 @@ struct ggml_tensor * ggml_ssm_conv_tree( return result; } +// dflash: fused conv step. Same op id as ggml_ssm_conv; op_params[0] = 1 +// marks step mode, srcs are (x, c, conv_state, conv_input_out). +struct ggml_tensor * ggml_ssm_conv_step( + struct ggml_context * ctx, + struct ggml_tensor * x, + struct ggml_tensor * c, + struct ggml_tensor * conv_state, + struct ggml_tensor * conv_input_out) { + GGML_ASSERT(x->type == GGML_TYPE_F32); + GGML_ASSERT(c->type == GGML_TYPE_F32); + GGML_ASSERT(conv_state->type == GGML_TYPE_F32); + GGML_ASSERT(ggml_is_matrix(c)); + GGML_ASSERT(ggml_is_contiguous(c)); + GGML_ASSERT(x->nb[0] == sizeof(float)); + GGML_ASSERT(x->ne[3] == 1); + + const int64_t d_conv = c->ne[0]; + const int64_t d_inner = c->ne[1]; + const int64_t n_t = x->ne[1]; + const int64_t n_s = x->ne[2]; + + GGML_ASSERT(x->ne[0] == d_inner); + GGML_ASSERT(conv_state->ne[0] == d_conv - 1); + GGML_ASSERT(conv_state->ne[1] == d_inner); + GGML_ASSERT(conv_state->ne[2] == n_s); + GGML_ASSERT(conv_state->nb[0] == sizeof(float)); + GGML_ASSERT(conv_state->nb[1] == (size_t)(d_conv - 1) * sizeof(float)); + if (conv_input_out) { + GGML_ASSERT(conv_input_out->type == GGML_TYPE_F32); + GGML_ASSERT(conv_input_out->ne[0] >= d_conv - 1 + n_t); + GGML_ASSERT(conv_input_out->ne[1] == d_inner); + GGML_ASSERT(conv_input_out->ne[2] == n_s); + GGML_ASSERT(conv_input_out->nb[0] == sizeof(float)); + } + + struct ggml_tensor * result = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, d_inner, n_t, n_s); + ggml_set_op_params_i32(result, 0, 2); // step mode (1 = SpecLA heavy-light conv) + + result->op = GGML_OP_SSM_CONV; + result->src[0] = x; + result->src[1] = c; + result->src[2] = conv_state; + result->src[3] = conv_input_out; + + return result; +} + struct ggml_tensor * ggml_ssm_conv_specla( struct ggml_context * ctx, struct ggml_tensor * x, @@ -6764,6 +6811,28 @@ void ggml_gated_delta_net_set_skip_intermediate( tensor->nb[3] = tensor->nb[2]*tensor->ne[2]; } +// dflash: raw-gate mode (see ggml.h). [dt_bias | A] -> src[9], +// op_params[10] = 1. (src[8] / op_params[2] belong to the compact-decode and +// SpecLA variants.) +void ggml_gated_delta_net_set_raw_gates( + struct ggml_tensor * tensor, + struct ggml_tensor * gate_ba) { + GGML_ASSERT(tensor != NULL); + GGML_ASSERT(tensor->op == GGML_OP_GATED_DELTA_NET); + GGML_ASSERT(gate_ba != NULL); + GGML_ASSERT(gate_ba->type == GGML_TYPE_F32); + GGML_ASSERT(ggml_is_contiguous(gate_ba)); + const struct ggml_tensor * v = tensor->src[2]; + GGML_ASSERT(ggml_nelements(gate_ba) == 2*v->ne[1]); + // scalar gate only (no KDA), no tree mode, no SpecLA / compact decode + GGML_ASSERT(tensor->src[3]->ne[0] == 1); + GGML_ASSERT(tensor->src[6] == NULL); + GGML_ASSERT(tensor->src[8] == NULL); + GGML_ASSERT(ggml_get_op_params_i32(tensor, 2) == 0); + tensor->src[9] = gate_ba; + ggml_set_op_params_i32(tensor, 10, 1); +} + // dflash: tree-mode variant. Same op, with parent_ids plumbed into // src[6] so the CUDA kernel can branch-reload state at DFS transitions. struct ggml_tensor * ggml_gated_delta_net_tree( diff --git a/server/scripts/convert_dflash_to_gguf.py b/server/scripts/convert_dflash_to_gguf.py index b904d7ea5..3beaa3dc4 100644 --- a/server/scripts/convert_dflash_to_gguf.py +++ b/server/scripts/convert_dflash_to_gguf.py @@ -110,6 +110,37 @@ def pick(*keys): or c.get("aux_hidden_state_layer_ids")) if _tli: a["capture_layer_ids"] = [int(x) for x in _tli] + # Newer HF configs (transformers >= 5.x, e.g. the Qwen3.8 DSpark + # drafter) nest rope_theta / YaRN under rope_parameters and + # mask_token_id under dflash_config instead of top-level. + rp = c.get("rope_parameters") or c.get("rope_scaling") or {} + if isinstance(rp, dict): + if rp.get("rope_theta") is not None: + a["rope_theta"] = float(rp["rope_theta"]) + if str(rp.get("rope_type", "")).lower() == "yarn": + a["yarn_factor"] = float(rp.get("factor", 0.0)) + a["yarn_orig_ctx"] = int(rp.get("original_max_position_embeddings", 0)) + a["yarn_beta_fast"] = float(rp.get("beta_fast", 32.0)) + a["yarn_beta_slow"] = float(rp.get("beta_slow", 1.0)) + if dfc.get("mask_token_id") is not None: + a["mask_token_id"] = int(dfc["mask_token_id"]) + if dfc.get("block_size") is not None: + a["block_size"] = int(dfc["block_size"]) + # DFlash 2 (z-lab/inco): grouped dynamic convs + candidate selector. + if dfc.get("conv_kernel_size") is not None: + a["conv_kernel_size"] = int(dfc["conv_kernel_size"]) + a["conv_group_size"] = int(dfc.get("conv_group_size", 16)) + if dfc.get("selector_rank") is not None: + a["selector_rank"] = int(dfc["selector_rank"]) + a["selector_top_k"] = int(dfc.get("selector_top_k", 16)) + # Per-layer sliding-window / causal attention (Qwen3.6-style drafters + # and DFlash 2). HF: layer_types + sliding_window; a top-level + # is_causal=false (DFlash 2) makes every layer bidirectional, which is + # our default (no SWA pattern emitted). + lt = c.get("layer_types") + if lt and c.get("sliding_window") and c.get("is_causal", None) is not False: + a["swa_window"] = int(c["sliding_window"]) + a["swa_pattern"] = [str(x) == "sliding_attention" for x in lt] print(f"[info] read arch from {cfg_path}") else: print(f"[warn] no config.json next to safetensors; using 27B defaults") @@ -182,8 +213,17 @@ def map_name(name: str) -> str | None: "mlp.gate_proj.weight": f"blk.{i}.ffn_gate.weight", "mlp.up_proj.weight": f"blk.{i}.ffn_up.weight", "mlp.down_proj.weight": f"blk.{i}.ffn_down.weight", + # DFlash 2 grouped dynamic convs + "attention_conv.base_kernel": f"blk.{i}.attn_conv.base", + "attention_conv.kernel_projection.weight": f"blk.{i}.attn_conv.proj.weight", + "mlp_conv.base_kernel": f"blk.{i}.ffn_conv.base", + "mlp_conv.kernel_projection.weight": f"blk.{i}.ffn_conv.proj.weight", } return layer_map.get(rest) + # DFlash 2 candidate selector + if name == "candidate_selector.hidden_projection.weight": return "dflash.selector.hproj.weight" + if name == "candidate_selector.predecessor_codebook": return "dflash.selector.pred_cb" + if name == "candidate_selector.successor_codebook": return "dflash.selector.succ_cb" return None @@ -248,18 +288,30 @@ def bytes_to_np(raw: bytes, dtype: str, shape: list[int]) -> np.ndarray: } +# Alias sets per head tensor: SpecForge sidecar names, DS4 MTP-shard names, +# and single-file releases (e.g. RadixArk Qwen3.8-27B-DSpark) that carry the +# heads inline in the main model.safetensors. +DSPARK_MARKOV_W1_KEYS = ("dspark_markov_head.markov_w1.weight", + "mtp.2.markov_head.markov_w1.weight", + "markov_head.markov_w1.weight") +DSPARK_MARKOV_W2_KEYS = ("dspark_markov_head.markov_w2.weight", + "mtp.2.markov_head.markov_w2.weight", + "markov_head.markov_w2.weight") +DSPARK_CONF_W_KEYS = ("dspark_confidence_head.weight", + "mtp.2.confidence_head.proj.weight", + "confidence_head.proj.weight") +DSPARK_CONF_B_KEYS = ("dspark_confidence_head.bias", + "mtp.2.confidence_head.proj.bias", + "confidence_head.proj.bias") + DSPARK_TENSOR_MAP = { - ("dspark_markov_head.markov_w1.weight", - "mtp.2.markov_head.markov_w1.weight"): ("dflash.dspark.markov.w1", gguf.GGMLQuantizationType.F16), - ("dspark_markov_head.markov_w2.weight", - "mtp.2.markov_head.markov_w2.weight"): ("dflash.dspark.markov.w2", gguf.GGMLQuantizationType.F16), + DSPARK_MARKOV_W1_KEYS: ("dflash.dspark.markov.w1", gguf.GGMLQuantizationType.F16), + DSPARK_MARKOV_W2_KEYS: ("dflash.dspark.markov.w2", gguf.GGMLQuantizationType.F16), } DSPARK_CONFIDENCE_TENSOR_MAP = { - ("dspark_confidence_head.weight", - "mtp.2.confidence_head.proj.weight"): ("dflash.dspark.confidence.weight", gguf.GGMLQuantizationType.F16), - ("dspark_confidence_head.bias", - "mtp.2.confidence_head.proj.bias"): ("dflash.dspark.confidence.bias", gguf.GGMLQuantizationType.F32), + DSPARK_CONF_W_KEYS: ("dflash.dspark.confidence.weight", gguf.GGMLQuantizationType.F16), + DSPARK_CONF_B_KEYS: ("dflash.dspark.confidence.bias", gguf.GGMLQuantizationType.F32), } @@ -372,8 +424,8 @@ def add_dspark_aux_heads(writer, arch: str, aux_path: Path | None): return print(f"[info] reading DSpark aux heads from {aux_path}") - w1 = resolved[("dspark_markov_head.markov_w1.weight", "mtp.2.markov_head.markov_w1.weight")][1] - w2 = resolved[("dspark_markov_head.markov_w2.weight", "mtp.2.markov_head.markov_w2.weight")][1] + w1 = resolved[DSPARK_MARKOV_W1_KEYS][1] + w2 = resolved[DSPARK_MARKOV_W2_KEYS][1] vocab = int(w1.shape[0]) rank = int(w1.shape[1]) if tuple(w2.shape) != (vocab, rank): @@ -397,8 +449,8 @@ def add_dspark_aux_heads(writer, arch: str, aux_path: Path | None): conf_missing.append(names) continue conf_resolved[names] = (found_name, tensor, spec) - weight_names = ("dspark_confidence_head.weight", "mtp.2.confidence_head.proj.weight") - bias_names = ("dspark_confidence_head.bias", "mtp.2.confidence_head.proj.bias") + weight_names = DSPARK_CONF_W_KEYS + bias_names = DSPARK_CONF_B_KEYS if weight_names not in conf_resolved: if conf_missing: print("[warn] incomplete DSpark confidence head; Markov head will still load") @@ -470,6 +522,12 @@ def main(): writer.add_uint32(f"{ARCH}.vocab_size", a["vocab"]) writer.add_float32(f"{ARCH}.attention.layer_norm_rms_epsilon", a["rms_eps"]) writer.add_float32(f"{ARCH}.rope.freq_base", a["rope_theta"]) + if a.get("yarn_factor", 0.0) > 1.0: + writer.add_string(f"{ARCH}.rope.scaling.type", "yarn") + writer.add_float32(f"{ARCH}.rope.scaling.factor", a["yarn_factor"]) + writer.add_uint32(f"{ARCH}.rope.scaling.original_context_length", a["yarn_orig_ctx"]) + writer.add_float32(f"{ARCH}.rope.scaling.beta_fast", a["yarn_beta_fast"]) + writer.add_float32(f"{ARCH}.rope.scaling.beta_slow", a["yarn_beta_slow"]) # DFlash-specific hyperparameters writer.add_uint32(f"{ARCH}.dflash.n_target_layers", a["n_target_layers"]) @@ -484,6 +542,15 @@ def main(): elif _cap_ids: print(f"[warn] capture_layer_ids len {len(_cap_ids)} != n_target_layers " f"{a['n_target_layers']}; not embedding ids", file=sys.stderr) + if a.get("swa_pattern"): + writer.add_uint32(f"{ARCH}.attention.sliding_window", a["swa_window"]) + writer.add_array(f"{ARCH}.attention.sliding_window_pattern", [bool(x) for x in a["swa_pattern"]]) + if a.get("conv_kernel_size"): + writer.add_uint32(f"{ARCH}.dflash.dflash2.conv_kernel_size", a["conv_kernel_size"]) + writer.add_uint32(f"{ARCH}.dflash.dflash2.conv_group_size", a["conv_group_size"]) + if a.get("selector_rank"): + writer.add_uint32(f"{ARCH}.dflash.dflash2.selector_rank", a["selector_rank"]) + writer.add_uint32(f"{ARCH}.dflash.dflash2.selector_top_k", a["selector_top_k"]) # Walk + add tensors. Sort: dflash.* singletons first, then output_*, # then per-layer in numeric order — keeps the on-disk layout stable. @@ -522,7 +589,8 @@ def sort_key(t): is_norm = ( gguf_name.endswith("_norm.weight") or gguf_name == "output_norm.weight" or - gguf_name == "dflash.hidden_norm.weight" + gguf_name == "dflash.hidden_norm.weight" or + gguf_name.endswith("_conv.base") # DFlash 2 conv base kernels [2, K, hidden] ) if is_norm: arr = arr.astype(" create_backend( cfg.max_concurrency = args.max_concurrency; cfg.kv_pool_tokens = args.kv_pool_tokens; cfg.kq_stride_pad = args.kq_stride_pad; + cfg.draft_block_size = args.draft_block_size; cfg.draft_swa_window = args.draft_swa_window; cfg.draft_ctx_max = args.draft_ctx_max; cfg.fast_rollback = args.fast_rollback; diff --git a/server/src/common/dflash2_head.cpp b/server/src/common/dflash2_head.cpp new file mode 100644 index 000000000..2cbcc2eb4 --- /dev/null +++ b/server/src/common/dflash2_head.cpp @@ -0,0 +1,202 @@ +#include "dflash2_head.h" + +#include "ggml-alloc.h" + +#include +#include +#include +#include +#include + +namespace dflash::common { + +namespace { + +// Selector projection graph, built once per (drafter, backend, n_cand, K). +struct SelectorGraph { + const DraftWeights * dw = nullptr; + ggml_backend_t backend = nullptr; + int n_cand = 0; + int K = 0; + std::vector arena; + ggml_context * ctx = nullptr; + ggml_cgraph * gf = nullptr; + ggml_gallocr_t galloc = nullptr; + ggml_tensor * inp_hidden = nullptr; + ggml_tensor * inp_succ = nullptr; + ggml_tensor * inp_pred = nullptr; + ggml_tensor * hproj = nullptr; + ggml_tensor * succ = nullptr; + ggml_tensor * pred = nullptr; +}; + +SelectorGraph & selector_graph() { + static thread_local SelectorGraph g; + return g; +} + +void selector_graph_free(SelectorGraph & g) { + if (g.galloc) { ggml_gallocr_free(g.galloc); g.galloc = nullptr; } + if (g.ctx) { ggml_free(g.ctx); g.ctx = nullptr; } + g.gf = nullptr; + g.dw = nullptr; + g.n_cand = 0; + g.K = 0; +} + +} // namespace + +bool dflash2_score_candidates(const DraftWeights & dw, + ggml_backend_t backend, + DFlashTarget & target, + const float * local_hidden, + int q_len, + int32_t last_tok, + float temperature, + Dflash2TreeScores & out) { + const DraftSelectorWeights & sel = dw.selector; + if (!sel.enabled || !sel.hproj || !sel.pred_cb || !sel.succ_cb) return false; + if (q_len <= 1 || !local_hidden || !backend) return false; + const int hdim = dw.n_embd; + const int rank = sel.rank; + const int K = sel.top_k; + const int n_cand = q_len - 1; + if (hdim <= 0 || rank <= 0 || K <= 0) return false; + + // 1. Top-k candidates (log-probs) per block position through the target + // lm_head. Position 0 of local_hidden is the seed slot; candidates are + // rows 1 .. q_len-1. + if (!target.project_hidden_to_topk(local_hidden + (size_t)hdim, n_cand, K, temperature, + out.lp, out.ids)) { + return false; + } + if (out.lp.size() != (size_t)n_cand * K || out.ids.size() != (size_t)n_cand * K) return false; + + // 2. One graph on the draft backend: hproj(h) for every candidate position, + // successor rows for every candidate, predecessor rows for the seed and + // every candidate. Built once per (n_cand, K) and reused across steps. + const int n_rows_pred = 1 + n_cand * K; + SelectorGraph & g = selector_graph(); + if (!g.ctx || g.dw != &dw || g.backend != backend || g.n_cand != n_cand || g.K != K) { + selector_graph_free(g); + const size_t arena_size = ggml_tensor_overhead() * 32 + ggml_graph_overhead() + 4096; + g.arena.assign(arena_size, 0); + ggml_init_params ip{}; + ip.mem_size = g.arena.size(); + ip.mem_buffer = g.arena.data(); + ip.no_alloc = true; + g.ctx = ggml_init(ip); + if (!g.ctx) return false; + g.gf = ggml_new_graph(g.ctx); + g.inp_hidden = ggml_new_tensor_2d(g.ctx, GGML_TYPE_F32, hdim, n_cand); + g.inp_succ = ggml_new_tensor_1d(g.ctx, GGML_TYPE_I32, n_cand * K); + g.inp_pred = ggml_new_tensor_1d(g.ctx, GGML_TYPE_I32, n_rows_pred); + ggml_set_input(g.inp_hidden); + ggml_set_input(g.inp_succ); + ggml_set_input(g.inp_pred); + g.hproj = ggml_mul_mat(g.ctx, sel.hproj, g.inp_hidden); // [rank, n_cand] + g.succ = ggml_get_rows(g.ctx, sel.succ_cb, g.inp_succ); // [rank, n_cand*K] f32 + g.pred = ggml_get_rows(g.ctx, sel.pred_cb, g.inp_pred); // [rank, 1+n_cand*K] f32 + ggml_set_output(g.hproj); + ggml_set_output(g.succ); + ggml_set_output(g.pred); + ggml_build_forward_expand(g.gf, g.hproj); + ggml_build_forward_expand(g.gf, g.succ); + ggml_build_forward_expand(g.gf, g.pred); + g.galloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); + if (!g.galloc || !ggml_gallocr_alloc_graph(g.galloc, g.gf)) { + std::fprintf(stderr, "dflash2_score_candidates: gallocr_alloc_graph failed\n"); + selector_graph_free(g); + return false; + } + g.dw = &dw; g.backend = backend; g.n_cand = n_cand; g.K = K; + } + + std::vector pred_ids((size_t)n_rows_pred); + pred_ids[0] = last_tok; + std::memcpy(pred_ids.data() + 1, out.ids.data(), sizeof(int32_t) * (size_t)n_cand * K); + ggml_backend_tensor_set(g.inp_hidden, local_hidden + (size_t)hdim, 0, sizeof(float) * (size_t)hdim * n_cand); + ggml_backend_tensor_set(g.inp_succ, out.ids.data(), 0, sizeof(int32_t) * (size_t)n_cand * K); + ggml_backend_tensor_set(g.inp_pred, pred_ids.data(), 0, sizeof(int32_t) * (size_t)n_rows_pred); + if (ggml_backend_graph_compute(backend, g.gf) != GGML_STATUS_SUCCESS) { + std::fprintf(stderr, "dflash2_score_candidates: graph_compute failed\n"); + return false; + } + out.hproj.resize((size_t)rank * n_cand); + out.succ.resize((size_t)rank * n_cand * K); + out.pred.resize((size_t)rank * n_rows_pred); + ggml_backend_tensor_get_async(backend, g.hproj, out.hproj.data(), 0, sizeof(float) * out.hproj.size()); + ggml_backend_tensor_get_async(backend, g.succ, out.succ.data(), 0, sizeof(float) * out.succ.size()); + ggml_backend_tensor_get_async(backend, g.pred, out.pred.data(), 0, sizeof(float) * out.pred.size()); + ggml_backend_synchronize(backend); + out.n_cand = n_cand; out.K = K; out.rank = rank; out.seed = last_tok; + return true; +} + +bool Dflash2TreeScores::topk(const std::vector & prefix, int next_depth, + std::vector & out_lp, std::vector & out_ids) const { + const int i = next_depth - 1; // candidate position + if (i < 0 || i >= n_cand || (int)prefix.size() != i) return false; + // predecessor row: 0 = seed, 1 + (i-1)*K + j = candidate j of position i-1 + int prev_row = 0; + if (i > 0) { + const int32_t parent_tok = prefix.back(); + prev_row = -1; + for (int j = 0; j < K; ++j) { + if (ids[(size_t)(i - 1) * K + j] == parent_tok) { prev_row = 1 + (i - 1) * K + j; break; } + } + if (prev_row < 0) prev_row = 0; // unknown parent: fall back to raw log-probs via seed row? no — zero compat + } + const float * pr = pred.data() + (size_t)prev_row * rank; + const float * hp = hproj.data() + (size_t)i * rank; + std::vector> scored((size_t)K); + for (int k = 0; k < K; ++k) { + const float * sc = succ.data() + ((size_t)i * K + k) * rank; + float dot = 0.0f; + for (int r = 0; r < rank; ++r) dot += pr[r] * hp[r] * sc[r]; + scored[(size_t)k] = { lp[(size_t)i * K + k] + dot, k }; + } + std::sort(scored.begin(), scored.end(), + [](const std::pair & a, const std::pair & b) { return a.first > b.first; }); + // The compatibility dot is not on a log-prob scale; renormalize the + // adjusted scores per position (log-softmax) so the tree builder's + // cumulative best-first comparison across depths stays meaningful. + float lse = 0.0f; + const float mx = scored[0].first; + for (int k = 0; k < K; ++k) lse += std::exp(scored[(size_t)k].first - mx); + lse = mx + std::log(lse); + out_lp.resize((size_t)K); + out_ids.resize((size_t)K); + for (int k = 0; k < K; ++k) { + out_lp[(size_t)k] = scored[(size_t)k].first - lse; + out_ids[(size_t)k] = ids[(size_t)i * K + scored[(size_t)k].second]; + } + return true; +} + +bool dflash2_select_chain(const DraftWeights & dw, + ggml_backend_t backend, + DFlashTarget & target, + const float * local_hidden, + int q_len, + int32_t last_tok, + std::vector & draft_tok) { + Dflash2TreeScores sc; + if (!dflash2_score_candidates(dw, backend, target, local_hidden, q_len, last_tok, + /*temperature=*/1.0f, sc)) { + return false; + } + // Greedy path over the candidates, conditioned on the previous pick. + draft_tok.assign((size_t)q_len, last_tok); + std::vector prefix; + std::vector top_lp; + std::vector top_ids; + for (int i = 0; i < sc.n_cand; ++i) { + if (!sc.topk(prefix, i + 1, top_lp, top_ids)) return false; + draft_tok[(size_t)i + 1] = top_ids[0]; + prefix.push_back(top_ids[0]); + } + return true; +} + +} // namespace dflash::common diff --git a/server/src/common/dflash2_head.h b/server/src/common/dflash2_head.h new file mode 100644 index 000000000..f054ac4fa --- /dev/null +++ b/server/src/common/dflash2_head.h @@ -0,0 +1,56 @@ +#pragma once + +#include "dflash_target.h" +#include "internal.h" + +#include +#include + +namespace dflash::common { + +// DFlash 2 candidate selector for greedy chain drafting. +// +// For every drafted block position the target lm_head logits are reduced to +// the selector's top-k candidates (log-probs, so per-position constants do +// not matter for the argmax), then one path is traced through them: +// score(c) = logp(c) + < pred_cb[prev] * hproj(h_pos), succ_cb[c] > +// prev = argmax_c score(c) +// starting from the block seed `last_tok`. Runs the projections (hproj GEMV +// and codebook row gathers) in one small graph on `backend`, the k-way path +// search on the host. Fills draft_tok = [last_tok, tok_1 .. tok_{q_len-1}]. +bool dflash2_select_chain(const DraftWeights & dw, + ggml_backend_t backend, + DFlashTarget & target, + const float * local_hidden, + int q_len, + int32_t last_tok, + std::vector & draft_tok); + +// Selector-scored candidates for DDTree construction (DARTree-style): the +// same per-position top-k + selector projections as the chain path, kept on +// the host so the tree builder can ask for branch-conditioned scores. +struct Dflash2TreeScores { + int n_cand = 0, K = 0, rank = 0; + int32_t seed = 0; + std::vector lp; // [n_cand*K] + std::vector ids; // [n_cand*K] + std::vector hproj; // [rank*n_cand] + std::vector succ; // [rank*n_cand*K] + std::vector pred; // [rank*(1+n_cand*K)] + + // K selector-adjusted scores for position `depth-1`, conditioned on the + // prefix'/s last token. Sorted descending; false if depth out of range. + bool topk(const std::vector & prefix, int next_depth, + std::vector & out_lp, std::vector & out_ids) const; +}; + +bool dflash2_score_candidates(const DraftWeights & dw, + ggml_backend_t backend, + DFlashTarget & target, + const float * local_hidden, + int q_len, + int32_t last_tok, + float temperature, + Dflash2TreeScores & out); + +} // namespace dflash::common diff --git a/server/src/common/feature_gate.cpp b/server/src/common/feature_gate.cpp index 0bcb2f1bb..df7b03f90 100644 --- a/server/src/common/feature_gate.cpp +++ b/server/src/common/feature_gate.cpp @@ -163,6 +163,17 @@ std::string check_feature_compatibility( "' does not support PFlash compression"; } + // A block-size override changes the local draft graph itself. Remote + // drafters own that shape in the IPC process and cannot be resized here. + if (args.draft_block_size != 0) { + if (args.draft_path == nullptr) { + return "--draft-block-size requires --draft"; + } + if (args.remote_draft.enabled()) { + return "--draft-block-size requires an in-process draft"; + } + } + // ── --paged-attention × architecture, placement, and decode features // Paged decode swaps the contiguous K/V cache for a block table owned by // the monolithic qwen35 backend, so every rule below is about reaching @@ -347,6 +358,11 @@ std::vector collect_feature_warnings( arch_supports_verify_width(arch, false), split, arch, "--verify-width", "chain-spec verify width"); + warn_inert(out, args.draft_block_size != 0, + arch_supports_draft_block_size(arch, split), + arch_supports_draft_block_size(arch, false), + split, arch, "--draft-block-size", "draft block-size override"); + warn_inert(out, args.fa_window != 0, arch_supports_fa_window(arch, split), arch_supports_fa_window(arch, false), diff --git a/server/src/common/geometric_draft_topk_cuda.cu b/server/src/common/geometric_draft_topk_cuda.cu index 71086c98a..ba287656d 100644 --- a/server/src/common/geometric_draft_topk_cuda.cu +++ b/server/src/common/geometric_draft_topk_cuda.cu @@ -13,7 +13,7 @@ namespace dflash::common { namespace { -constexpr int kMaxK = 8; // ddtree_K is 8 in practice; K>kMaxK → CPU fallback +constexpr int kMaxK = 16; // ddtree_K is 8, the DFlash 2 selector uses 16; K>kMaxK → CPU fallback constexpr int kBlock = 256; // threads per block (power of two for the reduction) constexpr int kMaxSplit = 128; // max vocab splits per position (combine-block cap) @@ -380,6 +380,7 @@ bool geometric_extract_draft_topk_cuda(const void * d_logits, switch (K) { DFLASH_TOPK_CASE(1) DFLASH_TOPK_CASE(2) DFLASH_TOPK_CASE(3) DFLASH_TOPK_CASE(4) DFLASH_TOPK_CASE(5) DFLASH_TOPK_CASE(6) DFLASH_TOPK_CASE(7) DFLASH_TOPK_CASE(8) + DFLASH_TOPK_CASE(12) DFLASH_TOPK_CASE(16) default: break; } #undef DFLASH_TOPK_CASE diff --git a/server/src/common/model_capabilities.h b/server/src/common/model_capabilities.h index cf14f414b..f62087141 100644 --- a/server/src/common/model_capabilities.h +++ b/server/src/common/model_capabilities.h @@ -20,8 +20,8 @@ // // qwen35 and qwen35moe share Qwen35Config, so their rows can differ on a // column the config carries — the moe backend simply never reads the field. -// paged_attn is the one such column today; see the cross-check comment in -// backend_factory.cpp for what that costs. +// paged_attn and draft_block_size are such columns today; see the cross-check +// comment in backend_factory.cpp for what that costs. // // Note on "qwen36": it is not a dispatchable architecture. model_card.cpp's // family fallback has a branch for it, but there is no factory case, so a @@ -59,6 +59,7 @@ struct ArchCapabilities { FeatureSupport decode_draft; // --draft FeatureSupport ddtree; // --ddtree, --ddtree-budget, --ddtree-temp FeatureSupport verify_width; // --verify-width + FeatureSupport draft_block_size; // --draft-block-size FeatureSupport fa_window; // --fa-window FeatureSupport draft_swa; // --draft-swa FeatureSupport paged_attn; // --paged-attention @@ -69,13 +70,13 @@ inline constexpr FeatureSupport kMono = FeatureSupport::Monolithic; inline constexpr FeatureSupport kBoth = FeatureSupport::Both; inline constexpr ArchCapabilities kArchCapabilities[] = { -// arch split rdraft pflash offload draft ddtree vwidth fa_win dswa paged - {"qwen35", true, true, true, false, kBoth, kBoth, kNever, kBoth, kBoth, kMono}, - {"qwen35moe", false, false, false, true, kMono, kMono, kNever, kMono, kMono, kNever}, - {"laguna", true, false, false, true, kMono, kMono, kMono, kNever, kNever, kNever}, - {"qwen3", false, false, true, false, kNever, kNever, kNever, kNever, kNever, kNever}, - {"gemma4", true, false, false, false, kMono, kNever, kNever, kBoth, kNever, kNever}, - {"deepseek4", true, false, false, false, kNever, kNever, kNever, kNever, kNever, kNever}, +// arch split rdraft pflash offload draft ddtree vwidth dblock fa_win dswa paged + {"qwen35", true, true, true, false, kBoth, kBoth, kNever, kMono, kBoth, kBoth, kMono}, + {"qwen35moe", false, false, false, true, kMono, kMono, kNever, kNever, kMono, kMono, kNever}, + {"laguna", true, false, false, true, kMono, kMono, kMono, kNever, kNever, kNever, kNever}, + {"qwen3", false, false, true, false, kNever, kNever, kNever, kNever, kNever, kNever, kNever}, + {"gemma4", true, false, false, false, kMono, kNever, kNever, kNever, kBoth, kNever, kNever}, + {"deepseek4", true, false, false, false, kNever, kNever, kNever, kNever, kNever, kNever, kNever}, }; inline constexpr std::size_t kArchCount = @@ -112,6 +113,7 @@ constexpr bool row_has_both(const ArchCapabilities & c) { return c.decode_draft == FeatureSupport::Both || c.ddtree == FeatureSupport::Both || c.verify_width == FeatureSupport::Both || + c.draft_block_size == FeatureSupport::Both || c.fa_window == FeatureSupport::Both || c.draft_swa == FeatureSupport::Both || c.paged_attn == FeatureSupport::Both; @@ -233,6 +235,12 @@ inline bool arch_supports_verify_width(const std::string & arch, return detail::arch_has(arch, &ArchCapabilities::verify_width, is_layer_split); } +inline bool arch_supports_draft_block_size(const std::string & arch, + bool is_layer_split) { + return detail::arch_has( + arch, &ArchCapabilities::draft_block_size, is_layer_split); +} + inline bool arch_supports_fa_window(const std::string & arch, bool is_layer_split) { return detail::arch_has(arch, &ArchCapabilities::fa_window, is_layer_split); diff --git a/server/src/delta_net_chunked.cpp b/server/src/delta_net_chunked.cpp index c3421bf2b..d0965c7e2 100644 --- a/server/src/delta_net_chunked.cpp +++ b/server/src/delta_net_chunked.cpp @@ -68,7 +68,12 @@ DeltaNetChunkedResult build_delta_net_chunked( g = ggml_permute(ctx0, g, 0, 2, 1, 3); b = ggml_permute(ctx0, b, 0, 2, 1, 3); - const int CS = kda ? 16 : 64; // chunk size + // Chunk size. Upstream uses 64, but the [CS, CS] triangular solve with + // k = CS only takes ggml-cuda's fast warp kernel for k <= 32; at CS = 64 + // it falls into cublasStrsmBatched, which on ROCm host-loops per batch + // (measured ~365 ms per solve node on gfx1201, 35 s per 512-token + // prefill). CS = 32 keeps every op on the fast path. + const int CS = kda ? 16 : 32; // chunk size const int pad = (CS - n_tokens % CS) % CS; const int n_chunks = (int)((n_tokens + pad) / CS); @@ -193,11 +198,17 @@ DeltaNetChunkedResult build_delta_net_chunked( ggml_tensor * v_t = ggml_cont(ctx0, ggml_transpose(ctx0, v)); for (int64_t chunk = 0; chunk < n_chunks; chunk++) { - ggml_tensor * ch_k_cd = get_slice_2d(ctx0, k_cd, chunk); + // The chunk slices are strided views (chunk is dim 2 of 4D tensors), + // which pushes their matmuls into cublasGemmBatchedEx; on ROCm that + // API stages its device pointer arrays through per-call pinned host + // allocations (~1 ms of hipHostMalloc/hipFree/hipMemcpy per node, + // measured 13 s per 512-token prefill). ggml_cont restores dim-2/3 + // contiguity so every matmul takes the strided-batched fast path. + ggml_tensor * ch_k_cd = ggml_cont(ctx0, get_slice_2d(ctx0, k_cd, chunk)); ggml_tensor * ch_v_t = get_slice_2d(ctx0, v_t, chunk); - ggml_tensor * ch_kq = get_slice_2d(ctx0, kq, chunk); - ggml_tensor * ch_q_g_exp = get_slice_2d(ctx0, q_g_exp, chunk); - ggml_tensor * ch_kg_t = get_slice_2d(ctx0, kg_t, chunk); + ggml_tensor * ch_kq = ggml_cont(ctx0, get_slice_2d(ctx0, kq, chunk)); + ggml_tensor * ch_q_g_exp = ggml_cont(ctx0, get_slice_2d(ctx0, q_g_exp, chunk)); + ggml_tensor * ch_kg_t = ggml_cont(ctx0, get_slice_2d(ctx0, kg_t, chunk)); ggml_tensor * v_t_p = ggml_mul_mat(ctx0, ch_k_cd, s); diff --git a/server/src/draft/draft_gguf_loader.cpp b/server/src/draft/draft_gguf_loader.cpp index e5a04721c..58c203195 100644 --- a/server/src/draft/draft_gguf_loader.cpp +++ b/server/src/draft/draft_gguf_loader.cpp @@ -75,7 +75,7 @@ int count_attn_gate_layers(const DraftWeights & w) { bool check_shape_1d(const ggml_tensor * t, int64_t ne0, const char * name, char * buf, size_t buf_sz) { if (!t || t->ne[0] != ne0) { - std::snprintf(buf, buf_sz, "draft GGUF: Domino tensor %s shape mismatch: got [%lld], expected [%lld]", + std::snprintf(buf, buf_sz, "draft GGUF: tensor %s shape mismatch: got [%lld], expected [%lld]", name, t ? (long long)t->ne[0] : -1LL, (long long)ne0); return false; } @@ -86,7 +86,7 @@ bool check_shape_2d(const ggml_tensor * t, int64_t ne0, int64_t ne1, const char * name, char * buf, size_t buf_sz) { if (!t || t->ne[0] != ne0 || t->ne[1] != ne1) { std::snprintf(buf, buf_sz, - "draft GGUF: Domino tensor %s shape mismatch: got [%lld,%lld], expected [%lld,%lld]", + "draft GGUF: tensor %s shape mismatch: got [%lld,%lld], expected [%lld,%lld]", name, t ? (long long)t->ne[0] : -1LL, t ? (long long)t->ne[1] : -1LL, @@ -96,6 +96,21 @@ bool check_shape_2d(const ggml_tensor * t, int64_t ne0, int64_t ne1, return true; } +bool check_shape_3d(const ggml_tensor * t, int64_t ne0, int64_t ne1, int64_t ne2, + const char * name, char * buf, size_t buf_sz) { + if (!t || t->ne[0] != ne0 || t->ne[1] != ne1 || t->ne[2] != ne2) { + std::snprintf(buf, buf_sz, + "draft GGUF: tensor %s shape mismatch: got [%lld,%lld,%lld], expected [%lld,%lld,%lld]", + name, + t ? (long long)t->ne[0] : -1LL, + t ? (long long)t->ne[1] : -1LL, + t ? (long long)t->ne[2] : -1LL, + (long long)ne0, (long long)ne1, (long long)ne2); + return false; + } + return true; +} + } // namespace bool load_draft_gguf(const std::string & path, @@ -209,6 +224,14 @@ bool load_draft_gguf(const std::string & path, if (target) { out.mask_token_id = target->mask_token_id; } + // The drafter's own MASK id wins over the family default: newer drafters + // (e.g. the Qwen3.8 DSpark release) are trained with a different mask + // token than the target-side default, and drafting with the wrong mask + // embedding silently destroys acceptance. + { + const uint32_t mask_meta = read_u32("dflash.mask_token_id", 0); + if (mask_meta != 0) out.mask_token_id = (int32_t)mask_meta; + } // Upper bounds on hparams. Guards against malformed/hostile GGUFs that // would otherwise trigger huge allocations or signed-int overflow when @@ -245,6 +268,24 @@ bool load_draft_gguf(const std::string & path, if (out.rope_theta == 0.0f) { fprintf(stderr, "[draft-gguf] WARNING: rope.freq_base not found in GGUF, draft RoPE will be wrong\n"); } + // YaRN rope scaling (optional). Drafters trained with YaRN (e.g. Qwen3.8 + // DSpark: factor 32, orig ctx 8192) apply it at every position; plain + // RoPE at inference silently degrades acceptance. + { + const float yarn_factor = read_f32("rope.scaling.factor", 0.0f); + if (yarn_factor > 1.0f) { + out.rope_freq_scale = 1.0f / yarn_factor; + out.rope_ext_factor = 1.0f; + out.rope_attn_factor = read_f32("rope.scaling.attn_factor", 1.0f); + out.rope_beta_fast = read_f32("rope.scaling.beta_fast", 32.0f); + out.rope_beta_slow = read_f32("rope.scaling.beta_slow", 1.0f); + out.rope_n_ctx_orig = (int)read_u32("rope.scaling.original_context_length", 0); + fprintf(stderr, + "[draft-gguf] YaRN rope: factor=%.1f orig_ctx=%d beta=%.1f/%.1f\n", + yarn_factor, out.rope_n_ctx_orig, + out.rope_beta_fast, out.rope_beta_slow); + } + } out.layers.assign((size_t)n_layer, DraftLayer{}); auto g = [&](const char * name) -> ggml_tensor * { @@ -301,6 +342,11 @@ bool load_draft_gguf(const std::string & path, L.w_gate = fnd("ffn_gate.weight"); L.w_up = fnd("ffn_up.weight"); L.w_down = fnd("ffn_down.weight"); + // DFlash 2 grouped dynamic convs (optional) + L.attn_conv.base = fnd("attn_conv.base"); + L.attn_conv.proj = fnd("attn_conv.proj.weight"); + L.mlp_conv.base = fnd("ffn_conv.base"); + L.mlp_conv.proj = fnd("ffn_conv.proj.weight"); if (!L.attn_norm || !L.ffn_norm || !L.wq || !L.wk || !L.wv || !L.wo || !L.q_norm || !L.k_norm || !L.w_gate || !L.w_up || !L.w_down) { char b[128]; @@ -451,6 +497,64 @@ bool load_draft_gguf(const std::string & path, out.dspark.confidence_dim); } + // DFlash 2: dynamic convs in every layer + candidate selector head. + { + const int conv_k = (int)read_u32("dflash.dflash2.conv_kernel_size", 0); + int n_conv = 0; + for (const DraftLayer & L : out.layers) { + if (L.attn_conv.present() && L.mlp_conv.present()) n_conv++; + } + if (n_conv > 0 || conv_k > 0) { + if (n_conv != out.n_layer || conv_k <= 0) { + set_last_error("draft GGUF: DFlash 2 conv tensors/metadata incomplete " + "(need attn_conv/ffn_conv base+proj in every layer and " + "dflash.dflash2.conv_kernel_size)"); + ggml_free(meta_ctx); out.ctx = nullptr; gguf_free(gctx); + return false; + } + out.conv_kernel_size = conv_k; + out.conv_group_size = (int)read_u32("dflash.dflash2.conv_group_size", 16); + const DraftLayer & L0 = out.layers[0]; + const int64_t groups = out.n_embd / out.conv_group_size; + char shape_err[192]; + if (!check_shape_3d(L0.attn_conv.base, out.n_embd, conv_k, 2, "attn_conv.base", shape_err, sizeof(shape_err)) || + !check_shape_2d(L0.attn_conv.proj, out.n_embd, 2 * conv_k * groups, "attn_conv.proj", shape_err, sizeof(shape_err))) { + set_last_error(shape_err); + ggml_free(meta_ctx); out.ctx = nullptr; gguf_free(gctx); + return false; + } + std::fprintf(stderr, "[draft GGUF] DFlash 2 dynamic convs: kernel=%d group=%d\n", + out.conv_kernel_size, out.conv_group_size); + } + out.selector = DraftSelectorWeights{}; + out.selector.hproj = g("dflash.selector.hproj.weight"); + out.selector.pred_cb = g("dflash.selector.pred_cb"); + out.selector.succ_cb = g("dflash.selector.succ_cb"); + const uint32_t sel_rank = read_u32("dflash.dflash2.selector_rank", 0); + if (out.selector.hproj || out.selector.pred_cb || out.selector.succ_cb || sel_rank) { + if (!out.selector.hproj || !out.selector.pred_cb || !out.selector.succ_cb) { + set_last_error("draft GGUF: DFlash 2 selector tensors incomplete " + "(hproj.weight, pred_cb, succ_cb)"); + ggml_free(meta_ctx); out.ctx = nullptr; gguf_free(gctx); + return false; + } + out.selector.rank = sel_rank ? (int)sel_rank : (int)out.selector.hproj->ne[1]; + out.selector.top_k = (int)read_u32("dflash.dflash2.selector_top_k", 16); + char shape_err[192]; + const int64_t R = out.selector.rank; + if (!check_shape_2d(out.selector.hproj, out.n_embd, R, "selector.hproj", shape_err, sizeof(shape_err)) || + !check_shape_2d(out.selector.pred_cb, R, out.selector.pred_cb->ne[1], "selector.pred_cb", shape_err, sizeof(shape_err)) || + !check_shape_2d(out.selector.succ_cb, R, out.selector.pred_cb->ne[1], "selector.succ_cb", shape_err, sizeof(shape_err))) { + set_last_error(shape_err); + ggml_free(meta_ctx); out.ctx = nullptr; gguf_free(gctx); + return false; + } + out.selector.enabled = true; + std::fprintf(stderr, "[draft GGUF] DFlash 2 selector enabled: rank=%d top_k=%d vocab=%lld\n", + out.selector.rank, out.selector.top_k, (long long)out.selector.pred_cb->ne[1]); + } + } + // GGUF Qwen3.6 drafters carry SWA metadata emitted by the converter: // dflash-draft.attention.sliding_window = 2048 // dflash-draft.attention.sliding_window_pattern = [true,true,true,true,false] diff --git a/server/src/draft/draft_graph.cpp b/server/src/draft/draft_graph.cpp index 472c214c9..10514d324 100644 --- a/server/src/draft/draft_graph.cpp +++ b/server/src/draft/draft_graph.cpp @@ -40,6 +40,19 @@ namespace dflash::common { +// RoPE with the drafter's scaling config. YaRN-trained drafters (e.g. the +// Qwen3.8 DSpark release: factor 32, orig ctx 8192) apply the scaled rotary +// at every position, so plain-RoPE inference silently degrades acceptance. +static ggml_tensor * draft_rope(ggml_context * ctx, ggml_tensor * t, + ggml_tensor * positions, + const DraftWeights & w) { + return ggml_rope_ext(ctx, t, positions, /*freq_factors=*/nullptr, + w.head_dim, GGML_ROPE_TYPE_NEOX, w.rope_n_ctx_orig, + w.rope_theta, w.rope_freq_scale, + w.rope_ext_factor, w.rope_attn_factor, + w.rope_beta_fast, w.rope_beta_slow); +} + // Feature fusion shared by the legacy one-shot graph and the cached-KV // builders: optional per-capture RMSNorm slices, fc projection, hidden_norm. // Row-independent, so it is bit-identical whether run over the full window @@ -72,6 +85,69 @@ static ggml_tensor * draft_fuse_features( return target_feat; } +// ── DFlash 2 grouped dynamic causal conv ──────────────────────────── +// +// Two taps over the draft block (positions within the block; the block's +// first slot has no predecessor). For each tap k the coefficient is a +// per-element base kernel plus a per-group dynamic kernel projected from +// the block's normalized hidden state: +// dyn = proj @ x_norm [2*K*groups, q_len] +// coef_s_k = base[s][k] (per element) + dyn[s][k] (per group, broadcast) +// out = sum_k coef_s_k * shift_k(x) +// s = 0 ("prepare", applied to the sub-block input) or 1 ("finish", applied +// to the sub-block output); both use the dyn computed from the input. +struct DraftDynConv { + ggml_tensor * dyn = nullptr; // [2*K*groups, q_len] +}; + +static DraftDynConv draft_dyn_conv_kernel(ggml_context * ctx, + const DraftConvWeights & cw, + ggml_tensor * x_norm) { + DraftDynConv dc; + dc.dyn = ggml_mul_mat(ctx, cw.proj, x_norm); // [2*K*groups, q_len] + return dc; +} + +static ggml_tensor * draft_dyn_conv_apply(ggml_context * ctx, + const DraftWeights & w, + const DraftConvWeights & cw, + const DraftDynConv & dc, + int s, // 0 = prepare, 1 = finish + ggml_tensor * x) { // [hidden, q_len] + const int64_t hidden = x->ne[0]; + const int64_t q_len = x->ne[1]; + const int K = w.conv_kernel_size; + const int64_t gs = w.conv_group_size; + const int64_t groups = hidden / gs; + const size_t e = ggml_element_size(dc.dyn); + + ggml_tensor * out = nullptr; + for (int k = 0; k < K; ++k) { + // shift_k(x): column l takes x[:, l-k], zero for l < k + ggml_tensor * xs = x; + if (k > 0) { + if (q_len <= k) break; + ggml_tensor * head = ggml_view_2d(ctx, x, hidden, q_len - k, x->nb[1], 0); + xs = ggml_pad_ext(ctx, head, 0, 0, k, 0, 0, 0, 0, 0); // [hidden, q_len] + } + // per-group dynamic coefficient for (s, k): rows [(s*K+k)*groups, +groups) + ggml_tensor * dyn_sk = ggml_view_3d(ctx, dc.dyn, 1, groups, q_len, + e, dc.dyn->nb[1], + (size_t)((s * K + k) * groups) * e); + ggml_tensor * xs3 = ggml_reshape_3d(ctx, xs, gs, groups, q_len); + ggml_tensor * dyn3 = ggml_repeat(ctx, dyn_sk, xs3); // [gs, groups, q_len] + // per-element base coefficient base[s][k]: [hidden] at offset (s*K+k)*hidden + ggml_tensor * base_sk = ggml_view_3d(ctx, cw.base, gs, groups, 1, + cw.base->nb[0] * gs, cw.base->nb[0] * hidden, + (size_t)(s * K + k) * cw.base->nb[1]); + ggml_tensor * coef = ggml_add(ctx, dyn3, base_sk); // broadcast over q_len + ggml_tensor * term = ggml_mul(ctx, xs3, coef); + term = ggml_reshape_2d(ctx, term, hidden, q_len); + out = out ? ggml_add(ctx, out, term) : term; + } + return out; +} + DraftGraphOutputs build_draft_graph( ggml_context * ctx, const DraftWeights & w, @@ -83,7 +159,6 @@ DraftGraphOutputs build_draft_graph( const int n_kv = w.n_head_kv; const int head_dim = w.head_dim; const float eps = DFLASH27B_RMS_EPS; - const float rope_base = w.rope_theta; // ── 1. Feature fusion: target_feat = rms_norm(fc @ target_hidden_cat, hidden_norm) // fc: [5*hidden, hidden] (ggml: ne[0]=5*hidden, ne[1]=hidden) @@ -118,10 +193,18 @@ DraftGraphOutputs build_draft_graph( const int eff_total_k = eff_ctx + q_len; const int ctx_offset = use_swa ? (ctx_len - w.swa_window) : 0; + const bool dyn_conv = w.conv_kernel_size > 0 && L.attn_conv.present() && L.mlp_conv.present(); + if (!disable_attn) { // ── 2a. Attention pre-norm ggml_tensor * hn = ggml_rms_norm(ctx, h, eps); hn = ggml_mul(ctx, hn, L.attn_norm); + // DFlash 2: dynamic conv "prepare" on the attention input + DraftDynConv attn_dc; + if (dyn_conv) { + attn_dc = draft_dyn_conv_kernel(ctx, L.attn_conv, hn); + hn = draft_dyn_conv_apply(ctx, w, L.attn_conv, attn_dc, 0, hn); + } std::snprintf(probe_name, sizeof(probe_name), "draft_l%d_hn", il); ggml_set_name(hn, probe_name); @@ -185,14 +268,8 @@ DraftGraphOutputs build_draft_graph( pk = ggml_view_1d(ctx, in.positions_k, eff_total_k, ctx_offset * ggml_element_size(in.positions_k)); } - Q = ggml_rope_ext(ctx, Q, in.positions_q, /*freq_factors=*/nullptr, - head_dim, GGML_ROPE_TYPE_NEOX, /*n_ctx_orig=*/0, - rope_base, /*freq_scale=*/1.0f, - /*ext_factor=*/0.0f, /*attn_factor=*/1.0f, - /*beta_fast=*/0.0f, /*beta_slow=*/0.0f); - K = ggml_rope_ext(ctx, K, pk, nullptr, - head_dim, GGML_ROPE_TYPE_NEOX, 0, - rope_base, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f); + Q = draft_rope(ctx, Q, in.positions_q, w); + K = draft_rope(ctx, K, pk, w); // ── 2e. Permute into the layout flash_attn_ext wants // q: [n_embd_k=head_dim, n_batch=q_len, n_head, ne3] @@ -235,6 +312,9 @@ DraftGraphOutputs build_draft_graph( // ── 2g. Output projection + residual // wo: [q_dim, hidden] (ne[0]=q_dim, ne[1]=hidden) ggml_tensor * attn_out = ggml_mul_mat(ctx, L.wo, attn); // [hidden, q_len] + if (dyn_conv) { + attn_out = draft_dyn_conv_apply(ctx, w, L.attn_conv, attn_dc, 1, attn_out); + } std::snprintf(probe_name, sizeof(probe_name), "draft_l%d_attn_out", il); ggml_set_name(attn_out, probe_name); h = ggml_add(ctx, h, attn_out); @@ -246,6 +326,11 @@ DraftGraphOutputs build_draft_graph( // ── 2h. FFN pre-norm ggml_tensor * hf = ggml_rms_norm(ctx, h, eps); hf = ggml_mul(ctx, hf, L.ffn_norm); + DraftDynConv mlp_dc; + if (dyn_conv) { + mlp_dc = draft_dyn_conv_kernel(ctx, L.mlp_conv, hf); + hf = draft_dyn_conv_apply(ctx, w, L.mlp_conv, mlp_dc, 0, hf); + } // ── 2i. SwiGLU: down(silu(gate(x)) * up(x)) // w_gate, w_up: [hidden, intermediate] @@ -255,6 +340,9 @@ DraftGraphOutputs build_draft_graph( ggml_tensor * u = ggml_mul_mat(ctx, L.w_up, hf); // [inter, q_len] ggml_tensor * gu = ggml_mul(ctx, g, u); ggml_tensor * ffn_out = ggml_mul_mat(ctx, L.w_down, gu); // [hidden, q_len] + if (dyn_conv) { + ffn_out = draft_dyn_conv_apply(ctx, w, L.mlp_conv, mlp_dc, 1, ffn_out); + } h = ggml_add(ctx, h, ffn_out); std::snprintf(probe_name, sizeof(probe_name), "draft_l%d_h_after_ffn", il); @@ -309,11 +397,7 @@ static void draft_ctx_kv_rows( K = ggml_reshape_3d(ctx, K, w.head_dim, w.n_head_kv, n); K = ggml_rms_norm(ctx, K, eps); K = ggml_mul (ctx, K, L.k_norm); - K = ggml_rope_ext(ctx, K, positions, /*freq_factors=*/nullptr, - w.head_dim, GGML_ROPE_TYPE_NEOX, /*n_ctx_orig=*/0, - w.rope_theta, /*freq_scale=*/1.0f, - /*ext_factor=*/0.0f, /*attn_factor=*/1.0f, - /*beta_fast=*/0.0f, /*beta_slow=*/0.0f); + K = draft_rope(ctx, K, positions, w); // rope output is contiguous [head_dim, n_kv, n] → head-major rows view *k_rows_out = ggml_view_2d(ctx, K, (int64_t)w.head_dim * w.n_head_kv, n, K->nb[2], 0); @@ -356,7 +440,6 @@ DraftGraphOutputs build_draft_kv_step( const int n_kv = w.n_head_kv; const int head_dim = w.head_dim; const float eps = DFLASH27B_RMS_EPS; - const float rope_base = w.rope_theta; const int kv_total = cache.kv_total; static const bool disable_attn_gate = @@ -370,28 +453,30 @@ DraftGraphOutputs build_draft_kv_step( for (int il = 0; il < w.n_layer; il++) { const DraftLayer & L = w.layers[il]; const bool layer_is_swa = L.is_swa && !disable_swa; + const bool dyn_conv = w.conv_kernel_size > 0 && L.attn_conv.present() && L.mlp_conv.present(); - // ── attention pre-norm + // ── attention pre-norm (+ DFlash 2 dynamic conv "prepare") ggml_tensor * hn = ggml_rms_norm(ctx, h, eps); hn = ggml_mul(ctx, hn, L.attn_norm); + DraftDynConv attn_dc; + if (dyn_conv) { + attn_dc = draft_dyn_conv_kernel(ctx, L.attn_conv, hn); + hn = draft_dyn_conv_apply(ctx, w, L.attn_conv, attn_dc, 0, hn); + } // ── Q from noise, per-head RMSNorm, RoPE at absolute positions ggml_tensor * Q = ggml_mul_mat(ctx, L.wq, hn); Q = ggml_reshape_3d(ctx, Q, head_dim, n_head, q_len); Q = ggml_rms_norm(ctx, Q, eps); Q = ggml_mul (ctx, Q, L.q_norm); - Q = ggml_rope_ext(ctx, Q, in.positions_q, nullptr, - head_dim, GGML_ROPE_TYPE_NEOX, 0, - rope_base, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f); + Q = draft_rope(ctx, Q, in.positions_q, w); // ── noise K/V into the scratch cache slots ggml_tensor * Kn = ggml_mul_mat(ctx, L.wk, hn); Kn = ggml_reshape_3d(ctx, Kn, head_dim, n_kv, q_len); Kn = ggml_rms_norm(ctx, Kn, eps); Kn = ggml_mul (ctx, Kn, L.k_norm); - Kn = ggml_rope_ext(ctx, Kn, in.positions_q, nullptr, - head_dim, GGML_ROPE_TYPE_NEOX, 0, - rope_base, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f); + Kn = draft_rope(ctx, Kn, in.positions_q, w); ggml_tensor * Kn_rows = ggml_view_2d(ctx, Kn, (int64_t)head_dim * n_kv, q_len, Kn->nb[2], 0); ggml_tensor * Vn_rows = ggml_mul_mat(ctx, L.wv, hn); // [kv_dim, q_len] @@ -436,16 +521,27 @@ DraftGraphOutputs build_draft_kv_step( attn = ggml_reshape_2d(ctx, attn, head_dim * n_head, q_len); ggml_tensor * attn_out = ggml_mul_mat(ctx, L.wo, attn); + if (dyn_conv) { + attn_out = draft_dyn_conv_apply(ctx, w, L.attn_conv, attn_dc, 1, attn_out); + } h = ggml_add(ctx, h, attn_out); - // ── FFN + // ── FFN (+ DFlash 2 dynamic conv prepare/finish) ggml_tensor * hf = ggml_rms_norm(ctx, h, eps); hf = ggml_mul(ctx, hf, L.ffn_norm); + DraftDynConv mlp_dc; + if (dyn_conv) { + mlp_dc = draft_dyn_conv_kernel(ctx, L.mlp_conv, hf); + hf = draft_dyn_conv_apply(ctx, w, L.mlp_conv, mlp_dc, 0, hf); + } ggml_tensor * g = ggml_mul_mat(ctx, L.w_gate, hf); g = ggml_silu(ctx, g); ggml_tensor * u = ggml_mul_mat(ctx, L.w_up, hf); ggml_tensor * gu = ggml_mul(ctx, g, u); ggml_tensor * ffn_out = ggml_mul_mat(ctx, L.w_down, gu); + if (dyn_conv) { + ffn_out = draft_dyn_conv_apply(ctx, w, L.mlp_conv, mlp_dc, 1, ffn_out); + } h = ggml_add(ctx, h, ffn_out); } diff --git a/server/src/internal.h b/server/src/internal.h index e7eaaa898..f65fe7c58 100644 --- a/server/src/internal.h +++ b/server/src/internal.h @@ -76,6 +76,16 @@ struct TargetLayer { ggml_tensor * ssm_dt_bias = nullptr; // [dt_rank] per-head alpha bias ggml_tensor * ssm_norm = nullptr; // [head_v_dim] ggml_tensor * ssm_out = nullptr; // output projection after delta-net + // Zero-copy stacked projections (set by the loader when the two source + // tensors share a type and were placed back to back in the weight buffer): + // wqkv_z: rows [0, n_z) = wqkv_gate (z), rows [n_z, ...) = wqkv + // ssm_ba: rows [0, dt_rank) = ssm_beta, rows [dt_rank, ...) = ssm_alpha + // One GEMV each instead of two; nullptr when stacking was not possible. + ggml_tensor * wqkv_z = nullptr; + ggml_tensor * ssm_ba = nullptr; + // Fused raw-gate GDN kernel parameters: f32 [2 * dt_rank] = [dt_bias | A], + // one small GPU tensor per DeltaNet layer (src[9] of the GDN op). + ggml_tensor * ssm_gate_ba = nullptr; // MoE FFN (qwen35moe only; nullptr on dense qwen35) ggml_tensor * ffn_gate_inp = nullptr; // [hidden, n_expert] router @@ -147,6 +157,9 @@ struct CpuEmbedder { struct TargetWeights { ggml_context * ctx = nullptr; + ggml_context * stack_ctx = nullptr; // owns the stacked alias tensors + ggml_context * gate_ctx = nullptr; // owns the [dt_bias | A] gate tensors + ggml_backend_buffer_t gate_buf = nullptr; ggml_backend_t backend = nullptr; ggml_backend_buffer_t buf = nullptr; @@ -234,6 +247,18 @@ void free_target_weights(TargetWeights & w); // ─── Draft weights (z-lab DFlash, bf16) ─────────────────────────── +// DFlash 2 grouped dynamic causal conv (two taps over the draft block, one +// instance before/after attention and one before/after the MLP): +// dyn = proj @ x_norm [2*K*groups, q_len] +// prepare = sum_k (base[0][k] + dyn[0][k]) * shift_k(x_norm) +// finish = sum_k (base[1][k] + dyn[1][k]) * shift_k(sub_block_out) +// base is per element, dyn per group of conv_group_size elements. +struct DraftConvWeights { + ggml_tensor * base = nullptr; // [hidden, K, 2] f32 + ggml_tensor * proj = nullptr; // [hidden, 2*K*groups] + bool present() const { return base != nullptr && proj != nullptr; } +}; + struct DraftLayer { ggml_tensor * attn_norm; ggml_tensor * ffn_norm; @@ -247,6 +272,8 @@ struct DraftLayer { ggml_tensor * w_gate; ggml_tensor * w_up; ggml_tensor * w_down; + DraftConvWeights attn_conv; // optional DFlash 2 conv around attention + DraftConvWeights mlp_conv; // optional DFlash 2 conv around the MLP bool is_swa = false; // true for SWA layers (Qwen3.6 pattern) bool attn_gate_per_head = false; }; @@ -280,6 +307,18 @@ struct DraftDSparkWeights { ggml_tensor * confidence_b = nullptr; // [1] f32 }; +// DFlash 2 candidate selector: top-k candidates per block position from the +// target lm_head logits, then one path through them scored by a low-rank +// bigram form unary[c] + . +struct DraftSelectorWeights { + bool enabled = false; + int rank = 0; + int top_k = 0; + ggml_tensor * hproj = nullptr; // [hidden, rank] + ggml_tensor * pred_cb = nullptr; // [rank, vocab] predecessor codebook + ggml_tensor * succ_cb = nullptr; // [rank, vocab] successor codebook +}; + struct DraftWeights { ggml_context * ctx = nullptr; ggml_backend_t backend = nullptr; @@ -324,6 +363,12 @@ struct DraftWeights { // Optional DSpark/DeepSpec-style Markov correction head. When present, // greedy chain decode adds a low-rank previous-token bias before argmax. DraftDSparkWeights dspark; + + // Optional DFlash 2 pieces: dynamic convs live in the layers, the + // selector replaces argmax/markov projection for the drafted chain. + int conv_kernel_size = 0; // 0 = no dynamic convs + int conv_group_size = 0; + DraftSelectorWeights selector; }; bool load_draft_safetensors(const std::string & path, diff --git a/server/src/qwen35/gguf_target_loader.cpp b/server/src/qwen35/gguf_target_loader.cpp index 1f917ee69..763f3939f 100644 --- a/server/src/qwen35/gguf_target_loader.cpp +++ b/server/src/qwen35/gguf_target_loader.cpp @@ -680,16 +680,71 @@ bool load_target_gguf_partial(const std::string & path, if (!t || !should_load_target_tensor(tname, plan.layer_begin, plan.layer_end, plan.load_output, plan.skip_expert_tensors)) { continue; } - alloc_total = align_up_size(alloc_total, alignment); TargetTensorAlloc a; a.tensor = t; a.file_offset = gguf_get_data_offset(gctx) + gguf_get_tensor_offset(gctx, tid); a.file_size = gguf_get_tensor_size(gctx, tid); - a.buffer_offset = alloc_total; - alloc_total += ggml_backend_buft_get_alloc_size(buft, t); allocs.push_back(a); } + // Stacked projections: place each (first, second) pair back to back in the + // weight buffer so one alias tensor spanning both rows serves a single + // GEMV. Only for the plain single-buffer path (the TP meta allocator + // places tensors itself) and only when the pair shares type/ne0 and the + // first tensor's byte size keeps the second one aligned. + const bool can_stack = !plan.metadata_only && !ggml_backend_buft_is_meta(buft) && + std::getenv("DFLASH_QWEN35_NO_STACK") == nullptr; + if (can_stack) { + auto find_alloc = [&](const std::string & name) -> int { + for (size_t i = 0; i < allocs.size(); i++) { + if (name == allocs[i].tensor->name) return (int)i; + } + return -1; + }; + // (first, second) suffix pairs; the alias tensor stacks first's rows + // then second's, so they are emitted in that order whichever member + // the file lists first. + static const char * const kPairs[][2] = { + { ".attn_gate.weight", ".attn_qkv.weight" }, + { ".ssm_beta.weight", ".ssm_alpha.weight" }, + }; + std::vector ordered; + ordered.reserve(allocs.size()); + std::vector taken(allocs.size(), false); + for (size_t i = 0; i < allocs.size(); i++) { + if (taken[i]) continue; + const std::string name = allocs[i].tensor->name; + int first = -1, second = -1; + if (name.rfind("blk.", 0) == 0) { + for (const auto & pr : kPairs) { + for (int m = 0; m < 2; m++) { + const size_t pos = name.find(pr[m]); + if (pos == std::string::npos) continue; + const std::string prefix = name.substr(0, pos); + first = find_alloc(prefix + pr[0]); + second = find_alloc(prefix + pr[1]); + break; + } + if (first >= 0 || second >= 0) break; + } + } + if (first >= 0 && second >= 0 && !taken[(size_t)first] && !taken[(size_t)second]) { + taken[(size_t)first] = taken[(size_t)second] = true; + ordered.push_back(allocs[(size_t)first]); + ordered.push_back(allocs[(size_t)second]); + continue; + } + taken[i] = true; + ordered.push_back(allocs[i]); + } + allocs.swap(ordered); + } + for (TargetTensorAlloc & a : allocs) { + alloc_total = align_up_size(alloc_total, alignment); + a.buffer_offset = alloc_total; + alloc_total += ggml_backend_buft_get_alloc_size(buft, a.tensor); + } + // The generic meta buffer allocator must see all tensors together so it // can allocate each device from its actual slices. The legacy loader's // monolithic backing buffer would reserve alloc_total on every rank. @@ -793,6 +848,47 @@ bool load_target_gguf_partial(const std::string & path, return false; } } + if (can_stack) { + // Alias tensors over adjacent pairs. They read the same bytes as + // the two source tensors (no copy, no extra VRAM). + ggml_init_params sip{}; + sip.mem_size = (2 * n_layer + 8) * ggml_tensor_overhead(); + sip.mem_buffer = nullptr; + sip.no_alloc = true; + out.stack_ctx = ggml_init(sip); + int n_stacked = 0; + auto make_stack = [&](ggml_tensor * first, ggml_tensor * second, + const char * name) -> ggml_tensor * { + if (!first || !second || !out.stack_ctx) return nullptr; + if (first->type != second->type || first->ne[0] != second->ne[0]) return nullptr; + if (!ggml_is_contiguous(first) || !ggml_is_contiguous(second)) return nullptr; + const char * f = (const char *)first->data; + const char * sd = (const char *)second->data; + if (!f || !sd || sd != f + ggml_nbytes(first)) return nullptr; + ggml_tensor * st = ggml_new_tensor_2d(out.stack_ctx, first->type, + first->ne[0], first->ne[1] + second->ne[1]); + // The alias must not need padding the backend would want to + // clear past its end (that would scribble on the next tensor). + if (ggml_backend_buft_get_alloc_size(buft, st) != ggml_nbytes(st)) return nullptr; + ggml_set_name(st, name); + if (ggml_backend_tensor_alloc(out.buf, st, first->data) != GGML_STATUS_SUCCESS) { + return nullptr; + } + n_stacked++; + return st; + }; + for (int il = 0; il < (int)n_layer; il++) { + TargetLayer & L = out.layers[il]; + char nm[96]; + std::snprintf(nm, sizeof(nm), "blk.%d.attn_gate_qkv.stacked", il); + L.wqkv_z = make_stack(L.wqkv_gate, L.wqkv, nm); + std::snprintf(nm, sizeof(nm), "blk.%d.ssm_beta_alpha.stacked", il); + L.ssm_ba = make_stack(L.ssm_beta, L.ssm_alpha, nm); + } + if (n_stacked > 0) { + std::fprintf(stderr, "[loader] stacked %d projection pairs (zero-copy aliases)\n", n_stacked); + } + } } const size_t data_start = gguf_get_data_offset(gctx); @@ -920,6 +1016,62 @@ bool load_target_gguf_partial(const std::string & path, return false; } + // ── Fused raw-gate GDN parameters: per DeltaNet layer one f32 [2*H] + // tensor holding [dt_bias | A] so the kernel can apply + // sigmoid/softplus itself (src[9] of the GDN op). Skipped for + // metadata-only / meta (TP) loads. + if (!plan.metadata_only && !ggml_backend_buft_is_meta(buft)) { + int n_gate = 0; + for (int il = 0; il < (int)n_layer; il++) { + const TargetLayer & L = out.layers[il]; + if (L.ssm_dt_bias && L.ssm_a && L.ssm_dt_bias->data && L.ssm_a->data && + L.ssm_dt_bias->type == GGML_TYPE_F32 && L.ssm_a->type == GGML_TYPE_F32 && + ggml_nelements(L.ssm_dt_bias) == ggml_nelements(L.ssm_a)) { + n_gate++; + } + } + if (n_gate > 0) { + ggml_init_params gip{}; + gip.mem_size = (n_gate + 2) * ggml_tensor_overhead(); + gip.mem_buffer = nullptr; + gip.no_alloc = true; + out.gate_ctx = ggml_init(gip); + if (out.gate_ctx) { + for (int il = 0; il < (int)n_layer; il++) { + TargetLayer & L = out.layers[il]; + if (!(L.ssm_dt_bias && L.ssm_a && L.ssm_dt_bias->data && L.ssm_a->data && + L.ssm_dt_bias->type == GGML_TYPE_F32 && L.ssm_a->type == GGML_TYPE_F32 && + ggml_nelements(L.ssm_dt_bias) == ggml_nelements(L.ssm_a))) { + continue; + } + const int64_t h = ggml_nelements(L.ssm_a); + L.ssm_gate_ba = ggml_new_tensor_1d(out.gate_ctx, GGML_TYPE_F32, 2*h); + char nm[96]; + std::snprintf(nm, sizeof(nm), "blk.%d.ssm_gate_ba", il); + ggml_set_name(L.ssm_gate_ba, nm); + } + out.gate_buf = ggml_backend_alloc_ctx_tensors(out.gate_ctx, backend); + if (out.gate_buf) { + ggml_backend_buffer_set_usage(out.gate_buf, GGML_BACKEND_BUFFER_USAGE_WEIGHTS); + std::vector tmp; + for (int il = 0; il < (int)n_layer; il++) { + TargetLayer & L = out.layers[il]; + if (!L.ssm_gate_ba) continue; + const int64_t h = ggml_nelements(L.ssm_a); + tmp.resize((size_t)2*h); + ggml_backend_tensor_get(L.ssm_dt_bias, tmp.data(), 0, (size_t)h*sizeof(float)); + ggml_backend_tensor_get(L.ssm_a, tmp.data() + h, 0, (size_t)h*sizeof(float)); + ggml_backend_tensor_set(L.ssm_gate_ba, tmp.data(), 0, (size_t)2*h*sizeof(float)); + } + } else { + for (int il = 0; il < (int)n_layer; il++) out.layers[il].ssm_gate_ba = nullptr; + ggml_free(out.gate_ctx); + out.gate_ctx = nullptr; + } + } + } + } + if (tok_embd_off == 0 || tok_embd_type == GGML_TYPE_COUNT) { set_last_error("token_embd.weight not found or invalid type"); release_out_buffer(); @@ -958,6 +1110,9 @@ bool load_target_gguf_partial(const std::string & path, void free_target_weights(TargetWeights & w) { if (w.buf) { ggml_backend_buffer_free(w.buf); w.buf = nullptr; } if (w.ctx) { ggml_free(w.ctx); w.ctx = nullptr; } + if (w.stack_ctx) { ggml_free(w.stack_ctx); w.stack_ctx = nullptr; } + if (w.gate_buf) { ggml_backend_buffer_free(w.gate_buf); w.gate_buf = nullptr; } + if (w.gate_ctx) { ggml_free(w.gate_ctx); w.gate_ctx = nullptr; } // CpuEmbedder destructor handles the mmap automatically. w.moe_hybrid.reset(); w.layers.clear(); diff --git a/server/src/qwen35/graph_builders.cpp b/server/src/qwen35/graph_builders.cpp index ed22072c8..118963305 100644 --- a/server/src/qwen35/graph_builders.cpp +++ b/server/src/qwen35/graph_builders.cpp @@ -495,7 +495,9 @@ bool build_target_step( ggml_set_input(sg.logits_row_indices); } - sg.gf = ggml_new_graph_custom(sg.ctx, 16384, false); + // 32k nodes: the chunked delta-net prefill graph (CS = 32) reaches ~17k + // nodes at a 512-token ubatch. + sg.gf = ggml_new_graph_custom(sg.ctx, 32768, false); // Step-invariant KV write: only when topology can't vary per step. // DFLASH_QWEN35_NO_KVPAD=1 restores the legacy cpy append + exact-length diff --git a/server/src/qwen35/qwen35_backend.cpp b/server/src/qwen35/qwen35_backend.cpp index 112cf0e0a..dc8e3db79 100644 --- a/server/src/qwen35/qwen35_backend.cpp +++ b/server/src/qwen35/qwen35_backend.cpp @@ -15,6 +15,8 @@ #include "common/geometric_sampler_cuda.h" #include #endif +#include "common/dspark_head.h" +#include "common/dflash2_head.h" #include "common/io_utils.h" #include "common/restore_delta.h" #include "common/specla_mode.h" @@ -28,6 +30,7 @@ #include "flashprefill.h" #include +#include #include #include #include @@ -217,6 +220,33 @@ static bool qwen35_empty_visible_output(const std::vector & tokens, return true; } +// Drafters trained on explicit target layers (GGUF dflash.target_layer_ids) +// override the evenly-spaced derivation: capturing different layers than the +// drafter was trained on silently destroys acceptance. +static void apply_drafter_capture_layer_ids(const DraftWeights & dw, TargetWeights & w) { + if (dw.capture_layer_ids.empty()) return; + const int n = (int)dw.capture_layer_ids.size(); + bool ok = (n == w.n_capture_layers); + for (int k = 0; ok && k < n; k++) + ok = dw.capture_layer_ids[k] >= 0 && dw.capture_layer_ids[k] < w.n_layer; + if (!ok) { + std::fprintf(stderr, + "[draft] drafter target_layer_ids invalid (n=%d, slots=%d); " + "keeping derived capture layers\n", n, w.n_capture_layers); + return; + } + bool changed = false; + for (int k = 0; k < n; k++) { + changed |= w.capture_layer_ids[k] != dw.capture_layer_ids[k]; + w.capture_layer_ids[k] = dw.capture_layer_ids[k]; + } + if (changed) { + std::printf("[draft] target capture layers from drafter GGUF:"); + for (int k = 0; k < n; k++) std::printf(" %d", w.capture_layer_ids[k]); + std::printf("\n"); + } +} + // ── Construction / destruction ────────────────────────────────────────── Qwen35Backend::Qwen35Backend(const Qwen35Config & cfg) : cfg_(cfg) {} @@ -323,6 +353,7 @@ bool Qwen35Backend::init() { return false; } std::printf("[draft] loaded\n"); + apply_drafter_capture_layer_ids(dw_, w_); if (cfg_.draft_swa_window > 0) { dw_.swa_window = cfg_.draft_swa_window; @@ -331,6 +362,22 @@ bool Qwen35Backend::init() { std::printf("[draft] SWA layers: %d/%d (window=%d)\n", dw_.n_layer - 1, dw_.n_layer, dw_.swa_window); } + + // DFlash weights are sequence-length agnostic; the GGUF block size is + // the training/default verify width, not a tensor dimension. A wider + // runtime block can trade a larger target batch for fewer verification + // steps without rewriting the model file. + if (cfg_.draft_block_size != 0) { + if (cfg_.draft_block_size < 2 || cfg_.draft_block_size > 32) { + std::fprintf(stderr, + "[draft] --draft-block-size must be in [2, 32], got %d\n", + cfg_.draft_block_size); + return false; + } + std::printf("[draft] block size override: %d -> %d\n", + dw_.block_size, cfg_.draft_block_size); + dw_.block_size = cfg_.draft_block_size; + } } // Create KV cache @@ -815,6 +862,7 @@ bool Qwen35Backend::unpark(ParkTarget target) { std::fprintf(stderr, "[unpark] draft: %s\n", dflash27b_last_error()); return false; } + apply_drafter_capture_layer_ids(dw_, w_); // Re-apply rope overrides after reload. if (dw_.rope_theta != w_.rope_theta && w_.rope_theta > 0.0f) dw_.rope_theta = w_.rope_theta; @@ -1659,6 +1707,8 @@ int Qwen35Backend::do_prefill(const std::vector & tokens, // positions encode the complete context — critical for tool // definitions at prompt start to propagate into KV values that // decode-time windowed attention will later read. + static const bool prefill_timing = std::getenv("DFLASH_PREFILL_TIMING") != nullptr; + const auto t_build0 = std::chrono::steady_clock::now(); if (!build_target_step(sg_, w_, cache_, target_backend_, /*kv_start=*/kv_pos, /*n_tokens=*/n_tokens, with_mask, /*capture=*/true, @@ -1735,7 +1785,17 @@ int Qwen35Backend::do_prefill(const std::vector & tokens, } // Compute + const auto t_comp0 = std::chrono::steady_clock::now(); auto st = ggml_backend_graph_compute(target_backend_, sg_.gf); + if (prefill_timing) { + ggml_backend_synchronize(target_backend_); + const auto t_comp1 = std::chrono::steady_clock::now(); + std::fprintf(stderr, + "[prefill-timing] tokens=%d nodes=%d build+alloc=%.1fms compute=%.1fms\n", + n_tokens, ggml_graph_n_nodes(sg_.gf), + std::chrono::duration(t_comp0 - t_build0).count(), + std::chrono::duration(t_comp1 - t_comp0).count()); + } if (st != GGML_STATUS_SUCCESS) { std::fprintf(stderr, "prefill compute @%d failed\n", kv_pos); return -1; @@ -2322,6 +2382,64 @@ bool Qwen35Backend::sync_local_draft_features(int start_pos, int n_tokens) { // ── DFlash speculative decode loop ───────────────────────────────────── +static bool qwen35_dspark_enabled() { + static const bool kEnabled = []() { + const char * e = std::getenv("DFLASH_QWEN35_DSPARK"); + return e == nullptr || std::string(e) != "0"; + }(); + return kEnabled; +} + +// Confidence-gate threshold for adaptive block length (0 = gate off, verify +// the full drafted block). The drafter's AcceptRatePredictor scores each +// draft position; the chain is truncated at the first position below the +// threshold and only the confident prefix is verified. +static float qwen35_dspark_confidence_threshold() { + static const float kThreshold = []() { + const char * e = std::getenv("DFLASH_QWEN35_DSPARK_CONFIDENCE_THRESHOLD"); + if (!e) return 0.0f; + float threshold = (float)std::atof(e); + if (threshold < 0.0f) threshold = 0.0f; + if (threshold > 1.0f) threshold = 1.0f; + return threshold; + }(); + return kThreshold; +} + +// Adaptive speculation policy: a spec step (draft + heads + width-q verify) +// costs about DFLASH_QWEN35_SPEC_STEP_RATIO plain-decode steps, so it only +// pays off while the drafter gets more than (ratio - 1) of its tokens +// accepted per step. Below that (low-acceptance prose) the loop runs +// DFLASH_QWEN35_AR_BURST plain-decode steps inside the spec loop (target +// forward on the seed token only, features still captured for the drafter), +// then probes with one spec step. Set DFLASH_QWEN35_SPEC_STEP_RATIO=0 to +// disable the policy. +struct Qwen35AdaptiveSpecPolicy { + float step_ratio = 1.9f; // spec/plain step cost used until both step kinds have been timed + // (measured 54-55 vs 28.6 ms on gfx1201 for width-8 and width-16 verify) + int burst = 40; // plain-decode steps per burst (each burst ends with one spec probe step) + float ema_alpha = 0.1f; // slow EMA: ~10-step memory so bursty acceptance does not flap + bool enabled() const { return step_ratio > 1.0f && burst > 0; } + // Enter a burst only clearly below break-even (hysteresis against noise). + // `ratio` is the live spec/plain step-time ratio once measured. + float accept_threshold(float ratio) const { return 0.8f * (ratio - 1.0f); } + float accept_threshold() const { return accept_threshold(step_ratio); } +}; + +static Qwen35AdaptiveSpecPolicy qwen35_adaptive_spec_policy() { + static const Qwen35AdaptiveSpecPolicy kPolicy = []() { + Qwen35AdaptiveSpecPolicy p; + if (const char * e = std::getenv("DFLASH_QWEN35_SPEC_STEP_RATIO")) { + p.step_ratio = (float)std::atof(e); + } + if (const char * e = std::getenv("DFLASH_QWEN35_AR_BURST")) { + p.burst = std::atoi(e); + } + return p; + }(); + return kPolicy; +} + bool Qwen35Backend::do_spec_decode(int committed, int n_gen, std::vector & out_tokens, const DaemonIO & io, @@ -2520,8 +2638,36 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, auto t_dec0 = std::chrono::steady_clock::now(); + // Adaptive speculation state (see Qwen35AdaptiveSpecPolicy). + const Qwen35AdaptiveSpecPolicy adaptive = qwen35_adaptive_spec_policy(); + // Start well above the burst threshold so an unlucky opening does not + // park a predictable stream in plain decode. The probe step that ends a + // burst updates the EMA with a fast alpha (see below) so a stream that + // turned predictable leaves plain decode quickly. + float accepted_ema = 2.0f * adaptive.accept_threshold(); + int ar_burst_left = 0; + int n_ar_burst_steps = 0; + bool probe_step = false; // first spec step after a burst + // Live step-time EMAs (seconds) for the break-even ratio; 0 = not yet measured. + double t_spec_step_ema = 0.0; + double t_ar_step_ema = 0.0; + auto live_step_ratio = [&]() { + return (t_spec_step_ema > 0.0 && t_ar_step_ema > 0.0) + ? (float)(t_spec_step_ema / t_ar_step_ema) : adaptive.step_ratio; + }; + while (n_generated < n_gen) { const int need_commit_budget = n_gen - n_generated; + // Plain-decode step inside the spec loop: no drafter forward, verify + // the seed token only. Features are still captured, so the drafter + // resumes cleanly on the next probe step. + const bool ar_step = adaptive.enabled() && ar_burst_left > 0; + if (ar_step) { + ar_burst_left--; + n_ar_burst_steps++; + probe_step = (ar_burst_left == 0); + } + const auto t_step_start = std::chrono::steady_clock::now(); // Budget hook: no tail-off here. The close-token injection fires // during the emit phase (step 8) after acceptance+replay, mirroring @@ -2562,107 +2708,109 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, return false; } - // 2. Draft compute + // 2. Draft compute (skipped on plain-decode burst steps) constexpr int DRAFT_CTX_MAX_DEFAULT = 2048; - const int ring_cap = use_remote_draft ? remote_draft_.ring_cap() : feature_mirror_.cap; - const int draft_ctx = std::min(committed, - std::min(ring_cap, std::max(DRAFT_CTX_MAX_DEFAULT, cfg_.draft_ctx_max))); - const int draft_start = committed - draft_ctx; - int mirror_slot0 = 0; - const bool use_mirror_view = - !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(); - if (!remote_draft_.propose(committed, draft_ctx, noise_embed, local_hidden)) { - std::fprintf(stderr, "spec-decode: remote draft propose failed\n"); - step_graph_destroy(draft_sg); - return false; - } - } else { - // [TAG_DRAFT_KV] ring-cached drafter context KV: append newly - // committed rows instead of re-encoding the whole feature window. - static const bool draft_kv_on = []() { - const char * e = std::getenv("DFLASH_DRAFT_KV"); - return !(e && e[0] == '0' && e[1] == '\0'); - }(); - bool use_draft_kv = draft_kv_on && feature_mirror_.target_feat != nullptr; - if (use_draft_kv && draft_kv_.gf && - draft_kv_.built_for != (const void *)&dw_) { - draft_kv_free(draft_kv_); - } - if (use_draft_kv && !draft_kv_.gf) { - const int kv_cap = std::min(ring_cap, - std::max(DRAFT_CTX_MAX_DEFAULT, cfg_.draft_ctx_max)); - if (!draft_kv_init(draft_kv_, dw_, draft_backend_, kv_cap, nullptr)) { - draft_kv_free(draft_kv_); - use_draft_kv = false; - std::fprintf(stderr, - "spec-decode: draft-kv init failed; using legacy draft path\n"); - } - } - 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"); - step_graph_destroy(draft_sg); - return false; - } - ggml_backend_tensor_set(draft_kv_.inp_embed, noise_embed.data(), 0, - sizeof(float) * noise_embed.size()); - if (ggml_backend_graph_compute(draft_backend_, draft_kv_.gf) != - GGML_STATUS_SUCCESS) { - std::fprintf(stderr, "spec-decode: draft-kv compute failed\n"); + if (!ar_step) { + const int ring_cap = use_remote_draft ? remote_draft_.ring_cap() : feature_mirror_.cap; + const int draft_ctx = std::min(committed, + std::min(ring_cap, std::max(DRAFT_CTX_MAX_DEFAULT, cfg_.draft_ctx_max))); + const int draft_start = committed - draft_ctx; + int mirror_slot0 = 0; + const bool use_mirror_view = + !use_remote_draft && + draft_feature_mirror_can_view(feature_mirror_, committed, draft_ctx, mirror_slot0); + + const auto profile_draft_start = profile_start(); + if (use_remote_draft) { + local_hidden.clear(); + if (!remote_draft_.propose(committed, draft_ctx, noise_embed, local_hidden)) { + std::fprintf(stderr, "spec-decode: remote draft propose failed\n"); step_graph_destroy(draft_sg); return false; } - local_hidden.resize((size_t)hidden * q_len); - ggml_backend_tensor_get(draft_kv_.hidden_states, local_hidden.data(), 0, - sizeof(float) * local_hidden.size()); } else { - if (!build_draft_step(draft_sg, dw_, /*lm_head=*/nullptr, draft_backend_, - draft_ctx, use_mirror_view ? &feature_mirror_ : nullptr, - committed, - /*ctx_len_max=*/std::min(ring_cap, std::max(DRAFT_CTX_MAX_DEFAULT, cfg_.draft_ctx_max)))) { - std::fprintf(stderr, "spec-decode: draft build failed\n"); - step_graph_destroy(draft_sg); - return false; - } - if (!use_mirror_view && - !copy_feature_ring_range_to_tensor(feature_mirror_, draft_sg.target_hidden_cat, - draft_start, draft_ctx)) { - std::fprintf(stderr, "spec-decode: feature copy failed\n"); - step_graph_destroy(draft_sg); - return false; + // [TAG_DRAFT_KV] ring-cached drafter context KV: append newly + // committed rows instead of re-encoding the whole feature window. + static const bool draft_kv_on = []() { + const char * e = std::getenv("DFLASH_DRAFT_KV"); + return !(e && e[0] == '0' && e[1] == '\0'); + }(); + bool use_draft_kv = draft_kv_on && feature_mirror_.target_feat != nullptr; + if (use_draft_kv && draft_kv_.gf && + draft_kv_.built_for != (const void *)&dw_) { + draft_kv_free(draft_kv_); } - ggml_backend_tensor_set(draft_sg.inp_embed, noise_embed.data(), 0, - sizeof(float) * noise_embed.size()); - pos_k.resize((size_t)draft_ctx + q_len); - for (int i = 0; i < q_len; i++) pos_q[i] = draft_ctx + i; - for (int i = 0; i < draft_ctx + q_len; i++) pos_k[i] = i; - ggml_backend_tensor_set(draft_sg.positions, pos_q.data(), 0, - sizeof(int32_t) * pos_q.size()); - ggml_backend_tensor_set(draft_sg.positions_k, pos_k.data(), 0, - sizeof(int32_t) * pos_k.size()); - - auto st = ggml_backend_graph_compute(draft_backend_, draft_sg.gf); - if (st != GGML_STATUS_SUCCESS) { - std::fprintf(stderr, "spec-decode: draft compute failed\n"); - step_graph_destroy(draft_sg); - return false; + if (use_draft_kv && !draft_kv_.gf) { + const int kv_cap = std::min(ring_cap, + std::max(DRAFT_CTX_MAX_DEFAULT, cfg_.draft_ctx_max)); + if (!draft_kv_init(draft_kv_, dw_, draft_backend_, kv_cap, nullptr)) { + draft_kv_free(draft_kv_); + use_draft_kv = false; + std::fprintf(stderr, + "spec-decode: draft-kv init failed; using legacy draft path\n"); + } } + 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"); + step_graph_destroy(draft_sg); + return false; + } + ggml_backend_tensor_set(draft_kv_.inp_embed, noise_embed.data(), 0, + sizeof(float) * noise_embed.size()); + if (ggml_backend_graph_compute(draft_backend_, draft_kv_.gf) != + GGML_STATUS_SUCCESS) { + std::fprintf(stderr, "spec-decode: draft-kv compute failed\n"); + step_graph_destroy(draft_sg); + return false; + } + local_hidden.resize((size_t)hidden * q_len); + ggml_backend_tensor_get(draft_kv_.hidden_states, local_hidden.data(), 0, + sizeof(float) * local_hidden.size()); + } else { + if (!build_draft_step(draft_sg, dw_, /*lm_head=*/nullptr, draft_backend_, + draft_ctx, use_mirror_view ? &feature_mirror_ : nullptr, + committed, + /*ctx_len_max=*/std::min(ring_cap, std::max(DRAFT_CTX_MAX_DEFAULT, cfg_.draft_ctx_max)))) { + std::fprintf(stderr, "spec-decode: draft build failed\n"); + step_graph_destroy(draft_sg); + return false; + } + if (!use_mirror_view && + !copy_feature_ring_range_to_tensor(feature_mirror_, draft_sg.target_hidden_cat, + draft_start, draft_ctx)) { + std::fprintf(stderr, "spec-decode: feature copy failed\n"); + step_graph_destroy(draft_sg); + return false; + } + ggml_backend_tensor_set(draft_sg.inp_embed, noise_embed.data(), 0, + sizeof(float) * noise_embed.size()); + pos_k.resize((size_t)draft_ctx + q_len); + for (int i = 0; i < q_len; i++) pos_q[i] = draft_ctx + i; + for (int i = 0; i < draft_ctx + q_len; i++) pos_k[i] = i; + ggml_backend_tensor_set(draft_sg.positions, pos_q.data(), 0, + sizeof(int32_t) * pos_q.size()); + ggml_backend_tensor_set(draft_sg.positions_k, pos_k.data(), 0, + sizeof(int32_t) * pos_k.size()); + + auto st = ggml_backend_graph_compute(draft_backend_, draft_sg.gf); + if (st != GGML_STATUS_SUCCESS) { + std::fprintf(stderr, "spec-decode: draft compute failed\n"); + step_graph_destroy(draft_sg); + return false; + } - // Read draft hidden states to host for LM-head projection. - local_hidden.resize((size_t)hidden * q_len); - ggml_backend_tensor_get(draft_sg.hidden_states, local_hidden.data(), 0, - sizeof(float) * local_hidden.size()); + // Read draft hidden states to host for LM-head projection. + local_hidden.resize((size_t)hidden * q_len); + ggml_backend_tensor_get(draft_sg.hidden_states, local_hidden.data(), 0, + sizeof(float) * local_hidden.size()); + } } - } - profile_add(profile_draft_s, profile_draft_start); + profile_add(profile_draft_s, profile_draft_start); + } // !ar_step // ── DDTree tree-structured verify ──────────────────────────────── // When --ddtree is on and the target supports tree verify, build a @@ -2694,17 +2842,119 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, kvflash_pager_.identity_prefix_covers(committed)); const bool use_tree_verify = cfg_.ddtree_mode && target->supports_tree_verify() && kvflash_tree_ok && - !use_remote_draft && q_len > 1 && tree_special_inactive; + !use_remote_draft && q_len > 1 && tree_special_inactive && !ar_step; + // Chain-verify length for this step. The DSpark confidence gate may + // truncate the drafted block (adaptive block length); q_len stays the + // buffer-sizing upper bound. + int v_len = q_len; // DDTree consumes top-K rows directly. Avoid projecting the same // hidden block once for argmax and again for top-K on every step. - if (!use_tree_verify) { - if (!target->project_hidden_to_tokens(local_hidden.data(), q_len, draft_tok)) { - std::fprintf(stderr, "spec-decode: projection failed\n"); - step_graph_destroy(draft_sg); - return false; + if (ar_step) { + draft_tok.assign(1, last_tok); + v_len = 1; + } else if (!use_tree_verify) { + const auto profile_project_start = profile_start(); + // DFlash 2 selector (top-k candidates + low-rank path score) when + // the drafter ships it. + bool used_dspark = false; + if (dw_.selector.enabled && q_len > 1 && !sampled_verify && !use_remote_draft) { + static std::atomic s_sel_logged{false}; + if (!s_sel_logged.exchange(true)) { + std::fprintf(stderr, + "[qwen35-spec] DFlash 2 selector active for greedy chain decode " + "(rank=%d top_k=%d)\n", dw_.selector.rank, dw_.selector.top_k); + } + if (dflash2_select_chain(dw_, draft_backend_, *target, + local_hidden.data(), q_len, last_tok, draft_tok)) { + used_dspark = true; + v_len = std::max(1, (int)draft_tok.size()); + } else { + static std::atomic s_sel_warned{false}; + if (!s_sel_warned.exchange(true)) { + std::fprintf(stderr, + "[qwen35-spec] DFlash 2 selector failed; falling back to " + "base DFlash projection\n"); + } + } } - draft_tok[0] = last_tok; + // DSpark heads (markov bigram correction + optional confidence + // gate) when the drafter ships them; mirrors the laguna hook. + if (!used_dspark && qwen35_dspark_enabled() && dw_.dspark.enabled && + q_len > 1 && !sampled_verify && !use_remote_draft) { + static std::atomic s_dspark_logged{false}; + if (!s_dspark_logged.exchange(true)) { + std::fprintf(stderr, + "[qwen35-spec] DSpark Markov head active for greedy chain decode " + "(rank=%d vocab=%d confidence_dim=%d)\n", + dw_.dspark.markov_rank, dw_.dspark.vocab_size, + dw_.dspark.confidence_dim); + } + static const bool fused_dspark = []() { + const char * e = std::getenv("DFLASH_QWEN35_FUSED_DSPARK"); + return !(e && e[0] == '0' && e[1] == '\0'); + }(); + bool ds_ok = false; + const float conf_threshold = qwen35_dspark_confidence_threshold(); + if (fused_dspark) { + // One graph for every candidate: markov-corrected tokens + // plus (when gated) the confidence score per position. + std::vector conf_scores; + ds_ok = dspark_markov_correct_greedy_chain_fused( + dw_, draft_backend_, target->lm_head_tensor(), + local_hidden.data(), q_len, last_tok, draft_tok, + conf_threshold > 0.0f ? &conf_scores : nullptr); + if (ds_ok && conf_threshold > 0.0f) { + // Truncate the chain at the first low-confidence + // position: draft_tok[0] is the seed, candidate i + // scores conf_scores[i-1]. + size_t keep = 1; + while (keep < draft_tok.size() && + keep - 1 < conf_scores.size() && + conf_scores[keep - 1] >= conf_threshold) { + ++keep; + } + static const bool conf_debug = []() { + const char * e = std::getenv("DFLASH_QWEN35_DSPARK_CONF_DEBUG"); + return e && e[0] == '1'; + }(); + if (conf_debug) { + std::fprintf(stderr, "[dspark-conf] keep=%zu/%zu:", keep, draft_tok.size()); + for (float c : conf_scores) std::fprintf(stderr, " %.3f", c); + std::fprintf(stderr, "\n"); + } + draft_tok.resize(keep); + } + } + if (!ds_ok) { + ds_ok = dspark_markov_correct_greedy_chain(dw_, draft_backend_, *target, + local_hidden.data(), q_len, + last_tok, conf_threshold, + draft_tok); + } + if (ds_ok) { + used_dspark = true; + // Confidence gate truncates the drafted chain: verify + // only the confident prefix this step. + v_len = std::max(1, (int)draft_tok.size()); + } else { + static std::atomic s_dspark_warned{false}; + if (!s_dspark_warned.exchange(true)) { + std::fprintf(stderr, + "[qwen35-spec] DSpark Markov head failed; falling back to " + "base DFlash projection\n"); + } + } + } + if (!used_dspark) { + if (!target->project_hidden_to_tokens(local_hidden.data(), q_len, draft_tok)) { + std::fprintf(stderr, "spec-decode: projection failed\n"); + step_graph_destroy(draft_sg); + return false; + } + draft_tok[0] = last_tok; + } + profile_add(profile_project_s, profile_project_start); } if (use_tree_verify) { @@ -2795,10 +3045,63 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, return false; } } else { - std::vector top_lp; + // DFlash2 selector-scored tree (DARTree-style): branch scores + // are logp + selector compatibility with the branch's actual + // parent, so the tree at worst degenerates to the selector + // chain. DFLASH_QWEN35_DFLASH2_TREE=0 falls back to raw top-k. + static const bool dflash2_tree = []() { + const char * e = std::getenv("DFLASH_QWEN35_DFLASH2_TREE"); + return !(e && e[0] == '0' && e[1] == '\0'); + }(); + bool selector_tree_ok = false; + if (dflash2_tree && dw_.selector.enabled && !use_remote_draft) { + Dflash2TreeScores sc; + const auto profile_project_start = profile_start(); + if (dflash2_score_candidates(dw_, draft_backend_, *target, + local_hidden.data(), q_len, last_tok, + cfg_.ddtree_temp, sc)) { + static std::atomic s_seltree_logged{false}; + if (!s_seltree_logged.exchange(true)) { + std::fprintf(stderr, + "[qwen35-spec] DFlash2 selector scores active for DDTree candidates\n"); + } + DDTreeConditionalTopK selector_topk = + [&](const std::vector & prefix, int next_depth, + std::vector & lp, std::vector & ids) -> bool { + if (!sc.topk(prefix, next_depth, lp, ids)) return false; + if ((int)lp.size() > K) { lp.resize((size_t)K); ids.resize((size_t)K); } + return true; + }; + tree = build_ddtree_conditional( + selector_topk, L, K, cfg_.ddtree_budget, + cfg_.ddtree_chain_seed, cfg_.ddtree_tau); + selector_tree_ok = tree.n_nodes > 0; + } + profile_add(profile_project_s, profile_project_start); + } + if (!selector_tree_ok) { + 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, + static const bool dspark_tree = []() { + const char * e = std::getenv("DFLASH_QWEN35_DSPARK_TREE"); + return !(e && e[0] == '0' && e[1] == '\0'); + }(); + bool topk_ok = false; + if (dspark_tree && qwen35_dspark_enabled() && dw_.dspark.enabled) { + static std::atomic s_dstree_logged{false}; + if (!s_dstree_logged.exchange(true)) { + std::fprintf(stderr, + "[qwen35-spec] DSpark Markov head active for DDTree candidates\n"); + } + topk_ok = dspark_markov_project_topk(dw_, draft_backend_, + target->lm_head_tensor(), + local_hidden.data(), q_len, K, + cfg_.ddtree_temp, last_tok, + top_lp, top_ids); + } + if (!topk_ok && + !target->project_hidden_to_topk(local_hidden.data(), q_len, K, cfg_.ddtree_temp, top_lp, top_ids)) { std::fprintf(stderr, @@ -2812,6 +3115,7 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, 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, @@ -3083,7 +3387,7 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, int hint_fill = 0; if (hint_tokens && n_generated < (int)hint_tokens->size()) { const int hint_avail = (int)hint_tokens->size() - n_generated; - hint_fill = std::min(hint_avail, q_len - 1); + hint_fill = std::min(hint_avail, v_len - 1); for (int i = 0; i < hint_fill; i++) { draft_tok[1 + i] = (*hint_tokens)[n_generated + i]; } @@ -3094,13 +3398,18 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, io.observer("draft", draft_tok); } - // 4. Verify: snapshot KV, run target forward over draft tokens - if (!target->snapshot_kv()) { + // 4. Verify: snapshot KV, run target forward over draft tokens. + // A plain-decode step verifies only the (always accepted) seed, so + // it never rolls back: skip the snapshot copy. + const auto profile_snapshot_start = profile_start(); + if (!ar_step && !target->snapshot_kv()) { step_graph_destroy(draft_sg); return false; } + profile_add(profile_snapshot_s, profile_snapshot_start); int verify_last_tok = -1; + const auto profile_verify_start = profile_start(); if (!target->verify_batch(draft_tok, committed, verify_last_tok, &target_tok, /*capture_ssm_intermediates=*/true)) { std::fprintf(stderr, "spec-decode: verify failed\n"); @@ -3108,6 +3417,7 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, step_graph_destroy(draft_sg); return false; } + profile_add(profile_verify_s, profile_verify_start); target_forwards++; // 5. Acceptance. Greedy: longest matching prefix between draft and @@ -3118,13 +3428,13 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, int accept_n = 1; int bonus_tok = -1; if (sampled_verify) { - if (!target->read_verify_logits(q_len, verify_logits)) { + if (!target->read_verify_logits(v_len, verify_logits)) { std::fprintf(stderr, "spec-decode: verify logits read failed\n"); target->restore_kv(); step_graph_destroy(draft_sg); return false; } - const int vocab_v = (int)(verify_logits.size() / (size_t)q_len); + const int vocab_v = (int)(verify_logits.size() / (size_t)v_len); static const bool kSvDebug = []() { const char * e = std::getenv("DFLASH_SV_DEBUG"); return e != nullptr && std::string(e) == "1"; @@ -3133,7 +3443,7 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, // Row-alignment check: CPU argmax over each bulk-read row must // equal the GPU argmax (target_tok). Divergence = misaligned // or stale bulk read. - for (int i = 0; i < q_len; i++) { + for (int i = 0; i < v_len; i++) { const float * row = verify_logits.data() + (size_t)i * vocab_v; int am = 0; float best = row[0]; for (int v = 1; v < vocab_v; v++) @@ -3158,7 +3468,7 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, verify_history = out_tokens; verify_history.push_back(draft_tok[0]); bool mismatched = false; - for (int i = 0; i < q_len - 1; i++) { + for (int i = 0; i < v_len - 1; i++) { const int s = sample_logits( verify_logits.data() + (size_t)i * vocab_v, vocab_v, sampler_, verify_history, sampler_rng_); @@ -3179,11 +3489,11 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, } (void)mismatched; } else { - for (int i = 0; i < q_len - 1; i++) { + for (int i = 0; i < v_len - 1; i++) { if (draft_tok[i + 1] == target_tok[i]) accept_n++; else break; } - bonus_tok = (accept_n < q_len) ? target_tok[accept_n - 1] : -1; + bonus_tok = (accept_n < v_len) ? target_tok[accept_n - 1] : -1; } // Track hint acceptance telemetry. if (hint_fill > 0) { @@ -3208,7 +3518,14 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, int replay_last_tok = -1; bool fast_rolled_back = false; - if (use_fast_rollback) { + if (ar_step) { + // Seed-only verify: the recurrent state already sits after the + // one committed token; nothing to restore. + bonus_tok = -1; + commit_n = std::min(accept_n, need_commit_budget); + replay_last_tok = target_tok[commit_n - 1]; + fast_rolled_back = true; + } else if (use_fast_rollback) { // Fast rollback: restore SSM from captured intermediates, skip replay. // Implicit bonus: target_tok[commit_n-1] seeds next draft as draft_tok[0], // always accepted on next step. @@ -3217,7 +3534,10 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, // budget (need_commit_budget), so committing accept_n would emit // more tokens than requested. commit_n was already clamped above. commit_n = std::min(accept_n, need_commit_budget); - if (target->rollback_to(committed, commit_n)) { + const auto profile_rollback_start = profile_start(); + const bool rolled = target->rollback_to(committed, commit_n); + profile_add(profile_rollback_s, profile_rollback_start); + if (rolled) { replay_last_tok = target_tok[commit_n - 1]; fast_rolled_back = true; rollback_diag.record_fast_rollback(accept_n); @@ -3248,11 +3568,13 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, for (int i = 0; i < commit_n; i++) { replay_batch[i] = (i < accept_n) ? draft_tok[i] : bonus_tok; } + const auto profile_replay_start = profile_start(); if (!target->verify_batch(replay_batch, committed, replay_last_tok, nullptr)) { std::fprintf(stderr, "spec-decode: replay failed\n"); step_graph_destroy(draft_sg); return false; } + profile_add(profile_replay_s, profile_replay_start); target_forwards++; } @@ -3269,10 +3591,12 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, return false; } } else if (feature_mirror_.target_feat && cache_.target_feat) { + const auto profile_feature_start = profile_start(); if (!sync_local_draft_features(committed, commit_n)) { step_graph_destroy(draft_sg); return false; } + profile_add(profile_feature_s, profile_feature_start); } // 8. Emit committed tokens (stop at EOS) @@ -3416,6 +3740,27 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, n_accept_sum += std::min(accept_n, emitted); n_draft_steps++; + // Adaptive policy update on real spec steps: EMA of accepted draft + // tokens (the seed is always accepted); a low EMA schedules a burst + // of plain-decode steps, the step after the burst is a spec probe. + if (adaptive.enabled()) { + const double t_step = std::chrono::duration( + std::chrono::steady_clock::now() - t_step_start).count(); + double & t_ema = ar_step ? t_ar_step_ema : t_spec_step_ema; + t_ema = (t_ema > 0.0) ? 0.9 * t_ema + 0.1 * t_step : t_step; + } + if (adaptive.enabled() && !ar_step) { + const float accepted_drafts = (float)std::max(0, accept_n - 1); + // A probe (first spec step after a burst) weighs its result + // heavily: it is the only evidence about the current text. + const float alpha = probe_step ? 0.5f : adaptive.ema_alpha; + accepted_ema = (1.0f - alpha) * accepted_ema + alpha * accepted_drafts; + probe_step = false; + if (accepted_ema < adaptive.accept_threshold(live_step_ratio())) { + ar_burst_left = adaptive.burst; + } + } + // Notify observer with accepted tokens for this step. if (io.observer) { io.observer("verify", replay_tok); @@ -3510,6 +3855,12 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, n_generated > 0 ? n_generated / decode_s : 0.0, n_draft_steps, n_accept_sum, total_draft_pos, accept_pct, n_draft_steps > 0 ? (double)n_generated / (double)n_draft_steps : 0.0); + if (n_ar_burst_steps > 0) { + std::fprintf(stderr, "[spec-decode] adaptive: %d of %d steps ran as plain decode " + "(step ratio %.2f, accept threshold %.2f drafts/step, burst %d)\n", + n_ar_burst_steps, n_draft_steps, live_step_ratio(), + adaptive.accept_threshold(live_step_ratio()), adaptive.burst); + } if (tp_profile) { std::fprintf(stderr, "[spec-profile] draft=%.3fs project=%.3fs snapshot=%.3fs " diff --git a/server/src/qwen35/qwen35_backend.h b/server/src/qwen35/qwen35_backend.h index c2ad20d17..01bb1940d 100644 --- a/server/src/qwen35/qwen35_backend.h +++ b/server/src/qwen35/qwen35_backend.h @@ -69,6 +69,7 @@ struct Qwen35Config { int64_t kv_pool_tokens = 0; // Draft + int draft_block_size = 0; // 0 = use drafter metadata int draft_swa_window = 0; int draft_ctx_max = 4096; diff --git a/server/src/qwen35/qwen35_dflash_target.h b/server/src/qwen35/qwen35_dflash_target.h index 9c1797652..8d645628a 100644 --- a/server/src/qwen35/qwen35_dflash_target.h +++ b/server/src/qwen35/qwen35_dflash_target.h @@ -78,6 +78,7 @@ class Qwen35DFlashTarget : public DFlashTarget { int hidden_size() const override { return w_.n_embd; } int mask_token_id() const override; + ggml_tensor * lm_head_tensor() override { return w_.output; } const std::vector & capture_layer_ids() const override; // kvflash mode: verify writes are slot-mapped via the pager and the diff --git a/server/src/qwen35/qwen35_target_graph.cpp b/server/src/qwen35/qwen35_target_graph.cpp index e4615cbc2..66069e104 100644 --- a/server/src/qwen35/qwen35_target_graph.cpp +++ b/server/src/qwen35/qwen35_target_graph.cpp @@ -153,7 +153,13 @@ bool create_target_cache_partial(const TargetWeights & w, // Graph-level FWHT K-rotation (TurboQuant-style outlier spreading with // standard quant types that keep fast FA kernel paths on all arches). // Skip for TQ3_0 K cache — that type already applies WHT during quantization. - out.kv_k_rotated = (kv_k_type != GGML_TYPE_TQ3_0); + // DFLASH_KV_ROTATE=0 turns it off (two fewer launches per attention layer; + // with q8_0/f16 caches the rotation is precision-neutral). + static const bool kv_rotate_env = []() { + const char * e = std::getenv("DFLASH_KV_ROTATE"); + return !(e && e[0] == '0' && e[1] == '\0'); + }(); + out.kv_k_rotated = (kv_k_type != GGML_TYPE_TQ3_0) && kv_rotate_env; const bool needs_256_stride = kv_k_type == GGML_TYPE_TQ3_0 || kv_v_type == GGML_TYPE_TQ3_0; @@ -1045,10 +1051,19 @@ bool ensure_ssm_snapshot(TargetCache & c, ggml_backend_t backend) { static ggml_tensor * build_swiglu_ffn(ggml_context * ctx, ggml_tensor * cur, const TargetLayer & L) { - ggml_tensor * gate = apply_scale2(ctx, ggml_mul_mat(ctx, L.w_gate, cur), L.w_gate_s); // [inter, n_tokens] - gate = ggml_silu(ctx, gate); - ggml_tensor * up = apply_scale2(ctx, ggml_mul_mat(ctx, L.w_up, cur), L.w_up_s); - ggml_tensor * gu = ggml_mul(ctx, gate, up); + ggml_tensor * gate = ggml_mul_mat(ctx, L.w_gate, cur); // [inter, n_tokens] + ggml_tensor * up = ggml_mul_mat(ctx, L.w_up, cur); + ggml_tensor * gu; + if (L.w_gate_s == 1.0f && L.w_up_s == 1.0f) { + // GLU node right after the two matmuls: the CUDA/HIP backend fuses + // mul_mat(gate) + mul_mat(up) + swiglu into a single vector kernel + // for single-token decode. + gu = ggml_swiglu_split(ctx, gate, up); + } else { + gate = ggml_silu(ctx, apply_scale2(ctx, gate, L.w_gate_s)); + up = apply_scale2(ctx, up, L.w_up_s); + gu = ggml_mul(ctx, gate, up); + } return apply_scale2(ctx, ggml_mul_mat(ctx, L.w_down, gu), L.w_down_s); // [hidden, n_tokens] } @@ -1437,23 +1452,85 @@ static ggml_tensor * build_delta_net_block( GGML_ASSERT(!(use_specla_factorized || use_specla_hld) || (n_seqs == 1 && !ragged && !active_slot_ids)); + // Row-slices of a stacked projection are only contiguous for a single + // token; wider batches (verify/prefill) need a copy before reshape/unary. + auto contig = [&](ggml_tensor * t) { + return ggml_is_contiguous(t) ? t : ggml_cont(ctx, t); + }; + // ── Whole-batch projections ───────────────────────────────────── // qkv_mixed = wqkv @ cur [10240, n_tokens] - ggml_tensor * qkv_2d = apply_scale2(ctx, ggml_mul_mat(ctx, L.wqkv, cur), L.wqkv_s); - - // z = wqkv_gate @ cur [inner, n_tokens] - ggml_tensor * z = apply_scale2(ctx, ggml_mul_mat(ctx, L.wqkv_gate, cur), L.wqkv_gate_s); - - // beta = sigmoid(ssm_beta @ cur) [dt_rank, n_tokens] - ggml_tensor * beta_2d = apply_scale2(ctx, ggml_mul_mat(ctx, L.ssm_beta, cur), L.ssm_beta_s); - beta_2d = ggml_sigmoid(ctx, beta_2d); + // z = wqkv_gate @ cur [inner, n_tokens] + // One GEMV over the stacked (z | qkv) alias when the loader built it; + // qkv_2d is then a strided column view of the stacked result. + ggml_tensor * qkv_2d = nullptr; + ggml_tensor * z = nullptr; + const bool stacked_qkv_z = L.wqkv_z && L.wqkv_s == 1.0f && L.wqkv_gate_s == 1.0f; + if (stacked_qkv_z) { + const int64_t n_z = L.wqkv_gate->ne[1]; + ggml_tensor * qkvz = ggml_mul_mat(ctx, L.wqkv_z, cur); // [n_z + conv_channels, n_tokens] + const size_t e = ggml_element_size(qkvz); + z = ggml_view_2d(ctx, qkvz, n_z, n_tokens, qkvz->nb[1], 0); + qkv_2d = ggml_view_2d(ctx, qkvz, conv_channels, n_tokens, qkvz->nb[1], (size_t)n_z * e); + } else { + qkv_2d = apply_scale2(ctx, ggml_mul_mat(ctx, L.wqkv, cur), L.wqkv_s); + z = apply_scale2(ctx, ggml_mul_mat(ctx, L.wqkv_gate, cur), L.wqkv_gate_s); + } + // beta = ssm_beta @ cur [dt_rank, n_tokens] // alpha = ssm_alpha @ cur [dt_rank, n_tokens] - // g = softplus(alpha + ssm_dt_bias) * ssm_a (-A_log.exp() * softplus) - ggml_tensor * alpha = apply_scale2(ctx, ggml_mul_mat(ctx, L.ssm_alpha, cur), L.ssm_alpha_s); - alpha = ggml_add(ctx, alpha, L.ssm_dt_bias); - alpha = ggml_softplus(ctx, alpha); - ggml_tensor * g_2d = ggml_mul(ctx, alpha, L.ssm_a); + // One GEMV over the stacked (beta | alpha) alias when available. + ggml_tensor * beta_2d = nullptr; + ggml_tensor * alpha = nullptr; + const bool stacked_ba = L.ssm_ba && L.ssm_beta_s == 1.0f && L.ssm_alpha_s == 1.0f; + if (stacked_ba) { + ggml_tensor * ba = ggml_mul_mat(ctx, L.ssm_ba, cur); // [2 * dt_rank, n_tokens] + const size_t e = ggml_element_size(ba); + beta_2d = contig(ggml_view_2d(ctx, ba, num_v_heads, n_tokens, ba->nb[1], 0)); + alpha = contig(ggml_view_2d(ctx, ba, num_v_heads, n_tokens, ba->nb[1], (size_t)num_v_heads * e)); + } else { + beta_2d = apply_scale2(ctx, ggml_mul_mat(ctx, L.ssm_beta, cur), L.ssm_beta_s); + alpha = apply_scale2(ctx, ggml_mul_mat(ctx, L.ssm_alpha, cur), L.ssm_alpha_s); + } + + // Fused kernels (single-sequence chain path only): the conv step and the + // gate prep are folded into the ssm_conv_step / gated_delta_net kernels + // instead of 6-8 tiny graph ops per layer. DFLASH_QWEN35_NO_FUSED_KERNELS=1 + // keeps the op-by-op graph for A/B checks. The chunked delta-net path + // (opt-in) needs the materialized gates, so it is decided here too. + static const bool fused_kernels_env = std::getenv("DFLASH_QWEN35_NO_FUSED_KERNELS") == nullptr; + // Chunked delta-net (llama.cpp build_delta_net_chunking port, verified + // ~1e-6 vs the sequential kernel): re-expresses the recurrence as + // chunk-parallel matmuls. Prefill-shaped calls only; decode, verify + // (rollback capture), tree, ragged and SpecLA paths always keep the + // sequential fused kernel. OFF by default: on gfx1201 the sequential + // kernel wins at a 512-token ubatch (514 ms vs 667 ms per forward; the + // ~20k-node chunk graph costs more in launches than it saves in GDN + // serialization). DFLASH27B_CHUNKED=1 opts in for A/B on other + // hardware. + static const bool chunked_env_on = []() { + const char * s_env = std::getenv("DFLASH27B_CHUNKED"); + return s_env && std::atoi(s_env) == 1; + }(); + const bool chunked_call = chunked_env_on && can_skip_gdn_intermediate && !ragged && + !active_slot_ids && !use_specla_factorized && !use_specla_hld && n_tokens > 1; + const bool fused_plain = fused_kernels_env && !parent_ids && !ragged && !active_slot_ids && + !use_specla_factorized && !use_specla_hld; + const bool fused_conv = fused_plain; + const bool raw_gates = fused_plain && !chunked_call && L.ssm_gate_ba != nullptr; + + // beta = sigmoid(beta); g = softplus(alpha + ssm_dt_bias) * ssm_a + // (-A_log.exp() * softplus). In raw-gate mode the GDN kernel applies both + // itself (dt_bias / A are attached via ggml_gated_delta_net_set_raw_gates). + ggml_tensor * g_2d = nullptr; + if (raw_gates) { + g_2d = alpha; + } else { + beta_2d = ggml_sigmoid(ctx, beta_2d); + alpha = ggml_add(ctx, alpha, L.ssm_dt_bias); + alpha = ggml_softplus(ctx, alpha); + g_2d = ggml_mul(ctx, alpha, L.ssm_a); + } // ── Token-axis segments: prompt chunks first, then the decode batch ── struct DeltaSeg { @@ -1515,9 +1592,15 @@ static ggml_tensor * build_delta_net_block( (allow_inplace_state && can_skip_gdn_intermediate && !ragged && n_seq_tokens == 1); - ggml_tensor * qkv_mixed = ggml_reshape_3d(ctx, - seg_cols(qkv_2d, seg.off, seg_tokens), - conv_channels, n_seq_tokens, seg_seqs); + // qkv_2d may be a strided view of the stacked (z | qkv) projection, so + // slice it with an explicit 3D view rather than a reshape. + ggml_tensor * qkv_mixed = ggml_view_3d(ctx, qkv_2d, + conv_channels, n_seq_tokens, seg_seqs, + qkv_2d->nb[1], qkv_2d->nb[1] * n_seq_tokens, + (size_t)seg.off * qkv_2d->nb[1]); + if (use_specla_hld || use_specla_factorized) { + qkv_mixed = contig(qkv_mixed); // the SpecLA conv kernels raw-index x + } ggml_tensor * beta = ggml_reshape_4d(ctx, seg_cols(beta_2d, seg.off, seg_tokens), 1, num_v_heads, n_seq_tokens, seg_seqs); @@ -1558,6 +1641,20 @@ static ggml_tensor * build_delta_net_block( w.ssm_d_conv - 1, conv_channels, seg_seqs); } + if (fused_conv) { + // One kernel: window = [conv_state | x], silu(conv), history + // write-back, and (when capturing) the rollback window copy. + ggml_tensor * ci_dst = nullptr; + if (cap && cap->conv_input) { + const int64_t ci_len = (w.ssm_d_conv - 1) + n_tokens; + ci_dst = (ci_len == cap->conv_input->ne[0]) + ? cap->conv_input + : 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); + } + conv_out = ggml_ssm_conv_step(ctx, qkv_mixed, L.ssm_conv1d, conv_states_r, ci_dst); + } else { // qkv_mixed currently is [conv_channels, n_tokens, n_seqs]; we need // [n_tokens, conv_channels, n_seqs] to concat on dim 0. ggml_tensor * qkv_T = ggml_transpose(ctx, qkv_mixed); @@ -1636,6 +1733,7 @@ static ggml_tensor * build_delta_net_block( : 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; @@ -1664,13 +1762,31 @@ static ggml_tensor * build_delta_net_block( row_size * n_seq_tokens, v_offset * elt); - // L2 norm on Q and K - q_c = ggml_l2_norm(ctx, q_c, w.rms_eps); - k_c = ggml_l2_norm(ctx, k_c, w.rms_eps); - - // Repeat Q and K from num_k_heads to num_v_heads so they match V's layout - // (only needed if not using the fused op's broadcast support). - if (num_k_heads != num_v_heads) { + // L2 norm on Q and K: q and k heads are adjacent in conv_out, so one + // launch over the [head_k_dim, 2*num_k_heads] slab normalizes both. + { + ggml_tensor * qk_c = ggml_view_4d(ctx, conv_out, + head_k_dim, 2 * num_k_heads, n_seq_tokens, seg_seqs, + head_k_dim * elt, + row_size, + row_size * n_seq_tokens, + q_offset * elt); + ggml_tensor * qk_n = ggml_l2_norm(ctx, qk_c, w.rms_eps); // contiguous [hd, 2*Hk, T, S] + const size_t ne_ = ggml_element_size(qk_n); + q_c = ggml_view_4d(ctx, qk_n, head_k_dim, num_k_heads, n_seq_tokens, seg_seqs, + qk_n->nb[1], qk_n->nb[2], qk_n->nb[3], 0); + k_c = ggml_view_4d(ctx, qk_n, head_k_dim, num_k_heads, n_seq_tokens, seg_seqs, + qk_n->nb[1], qk_n->nb[2], qk_n->nb[3], + (size_t)num_k_heads * head_k_dim * ne_); + } + + // Repeat Q and K from num_k_heads to num_v_heads so they match V's layout. + // The fused chain/tree gated_delta_net kernels broadcast heads themselves + // (v head h reads q/k head h % num_k_heads, the same tiling ggml_repeat + // produces); the chunked, compact-decode and SpecLA paths take the + // materialized copies. + if (num_k_heads != num_v_heads && + (chunked_call || seg_active || use_specla_factorized || use_specla_hld)) { q_c = ggml_repeat_4d(ctx, q_c, head_k_dim, num_v_heads, n_seq_tokens, seg_seqs); k_c = ggml_repeat_4d(ctx, k_c, head_k_dim, num_v_heads, n_seq_tokens, seg_seqs); } @@ -1714,12 +1830,10 @@ static ggml_tensor * build_delta_net_block( // default — port produces correct shape but slightly wrong final state, // causing AL degradation and loopy output. Set DFLASH27B_CHUNKED=1 to // opt in for A/B testing while debugging. - bool use_chunked = false; - if (can_skip_gdn_intermediate && n_seq_tokens > 1) { - if (const char * s_env = std::getenv("DFLASH27B_CHUNKED")) { - use_chunked = (std::atoi(s_env) != 0); - } - } + // Chunked delta-net path (opt-in via DFLASH27B_CHUNKED, chain-only, no + // capture): decided whole-batch above; a segment only qualifies with + // more than one timestep. + const bool use_chunked = chunked_call && n_seq_tokens > 1; ggml_tensor * output = nullptr; @@ -1794,12 +1908,18 @@ static ggml_tensor * build_delta_net_block( // cache buffer — same mechanism as _tree_persist, but without tree // parent_ids. Avoids the legacy result-region cpy (and the OOB it // could cause if the result tensor has no embedded intermediate region). + // In-place final state: the kernel writes the new recurrent state + // straight into `s` (a view of the persistent ssm_state), so no + // separate 3 MB copy per layer is needed. Tree mode keeps the copy. result = inplace_state ? ggml_gated_delta_net_inplace(ctx, q_c, k_c, v_c, g_tensor, beta, s) : ggml_gated_delta_net(ctx, q_c, k_c, v_c, g_tensor, beta, s); if (persist_inter) { result->src[7] = persist_inter; } + if (raw_gates) { + ggml_gated_delta_net_set_raw_gates(result, L.ssm_gate_ba); + } } if (can_skip_gdn_intermediate) { ggml_gated_delta_net_set_skip_intermediate(result, true); @@ -1858,7 +1978,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), + contig(seg_cols(z, seg.off, seg_tokens)), head_v_dim, num_v_heads, n_seq_tokens, seg_seqs); ggml_tensor * output_n = ggml_rms_norm(ctx, rms_norm_input_f32(ctx, output), w.rms_eps); output_n = ggml_mul(ctx, output_n, L.ssm_norm); diff --git a/server/src/server/server_main.cpp b/server/src/server/server_main.cpp index 32253b823..ca46fc491 100644 --- a/server/src/server/server_main.cpp +++ b/server/src/server/server_main.cpp @@ -86,6 +86,8 @@ static void print_usage(const char * prog) { " --draft-ipc-bin Remote backend IPC daemon for mixed backends\n" " --draft-ipc-work-dir Remote draft IPC scratch directory\n" " --draft-ipc-ring-cap Remote draft feature ring capacity\n" + " --draft-block-size Dense Qwen DFlash proposal/verify width\n" + " (2..32; default: drafter metadata)\n" " --draft-swa Draft sliding-window attention size (0=off; e.g.\n" " 2048 for unsloth Qwen3.6 targets, per server/README.md.\n" " Env: DFLASH27B_DRAFT_SWA)\n" @@ -303,6 +305,18 @@ int main(int argc, char ** argv) { } } else if (std::strcmp(argv[i], "--draft-swa") == 0 && i + 1 < argc) { bargs.draft_swa_window = std::atoi(argv[++i]); + } else if (std::strcmp(argv[i], "--draft-block-size") == 0 && i + 1 < argc) { + const char * value = argv[++i]; + const char * end = value + std::strlen(value); + const auto parsed = std::from_chars( + value, end, bargs.draft_block_size); + if (parsed.ec != std::errc{} || parsed.ptr != end || + bargs.draft_block_size < 2 || bargs.draft_block_size > 32) { + std::fprintf(stderr, + "--draft-block-size expects an integer in [2, 32], got '%s'\n", + value); + return 2; + } } else if (std::strcmp(argv[i], "--draft-device") == 0 && i + 1 < argc) { if (!parse_placement_device(argv[++i], bargs.draft_device)) { std::fprintf(stderr, "[server] bad --draft-device value (expected backend:gpu)\n"); diff --git a/server/test/test_feature_gate.cpp b/server/test/test_feature_gate.cpp index 001755d36..f4ea07c7c 100644 --- a/server/test/test_feature_gate.cpp +++ b/server/test/test_feature_gate.cpp @@ -115,6 +115,21 @@ void test_feature_gate_mixed_draft_placement_requires_ipc() { args, "qwen35", PlacementBackend::Cuda).empty()); } +void test_feature_gate_draft_block_size_requires_local_draft() { + BackendArgs args; + args.model_path = "/nonexistent/model.gguf"; + args.draft_block_size = 12; + CHECK(!gate_result(args, "qwen35", PlacementBackend::Cuda).empty()); + + args.draft_path = "/nonexistent/draft.gguf"; + CHECK(gate_result(args, "qwen35", PlacementBackend::Cuda).empty()); + + args.device.backend = PlacementBackend::Cuda; + args.draft_device.backend = PlacementBackend::Hip; + args.remote_draft.ipc_bin = "/usr/bin/draft-ipc"; + CHECK(!gate_result(args, "qwen35", PlacementBackend::Cuda).empty()); +} + void test_feature_gate_pflash_requires_drafter_and_supported_arch() { BackendArgs args; args.model_path = "/nonexistent/model.gguf"; @@ -532,6 +547,7 @@ void test_feature_warnings_silent_when_supported() { args.draft_path = "/nonexistent/draft.gguf"; args.ddtree_mode = true; args.fa_window = 512; + args.draft_block_size = 12; args.draft_swa_window = 2048; // qwen35 forwards every one of these. CHECK(warn_result(args, "qwen35").empty()); @@ -569,6 +585,15 @@ void test_feature_warnings_report_inert_decode_tunables() { CHECK(!warns_about(warn_result(vw, "laguna"), "--verify-width")); CHECK(warns_about(warn_result(vw, "qwen35"), "--verify-width")); + BackendArgs db; + db.model_path = "/nonexistent/model.gguf"; + db.draft_path = "/nonexistent/draft.gguf"; + db.draft_block_size = 12; + CHECK(!warns_about(warn_result(db, "qwen35"), "--draft-block-size")); + CHECK(warns_about(warn_result(db, "qwen35moe"), "--draft-block-size")); + CHECK(parse_placement_device_list("cuda:0,cuda:1", db.device)); + CHECK(warns_about(warn_result(db, "qwen35"), "--draft-block-size")); + BackendArgs fa; fa.model_path = "/nonexistent/model.gguf"; fa.fa_window = 4096; @@ -627,6 +652,7 @@ void test_model_capability_tables() { CHECK(!arch_supports_decode_draft("qwen36", false)); CHECK(!arch_supports_ddtree("qwen36", false)); CHECK(!arch_supports_verify_width("qwen36", false)); + CHECK(!arch_supports_draft_block_size("qwen36", false)); CHECK(!arch_supports_fa_window("qwen36", false)); CHECK(!arch_supports_draft_swa("qwen36", false)); CHECK(!arch_supports_paged_attention("qwen36", false)); @@ -635,6 +661,10 @@ void test_model_capability_tables() { CHECK(arch_supports_paged_attention("qwen35", false)); CHECK(!arch_supports_paged_attention("qwen35", true)); CHECK(!arch_supports_paged_attention("qwen35moe", false)); + + CHECK(arch_supports_draft_block_size("qwen35", false)); + CHECK(!arch_supports_draft_block_size("qwen35", true)); + CHECK(!arch_supports_draft_block_size("qwen35moe", false)); } }; @@ -646,6 +676,7 @@ TEST_CASE(FeatureGateFixture, feature_gate_suite) { test_feature_gate_requires_compiled_target_backend(); test_feature_gate_ipc_options_require_ipc_binary(); test_feature_gate_mixed_draft_placement_requires_ipc(); + test_feature_gate_draft_block_size_requires_local_draft(); test_feature_gate_pflash_requires_drafter_and_supported_arch(); test_feature_gate_validates_target_split_topology(); test_feature_gate_tensor_parallel_requirements();