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/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-backend-meta.cpp b/server/deps/llama.cpp/ggml/src/ggml-backend-meta.cpp index 9e68aa8c4..c0d3fc175 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-backend-meta.cpp +++ b/server/deps/llama.cpp/ggml/src/ggml-backend-meta.cpp @@ -825,6 +825,18 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( }; auto handle_ssm_conv = [&](const std::vector & src_ss) -> ggml_backend_meta_split_state { + // Step / SpecLA variants (ggml_ssm_conv_step, ggml_ssm_conv_specla): + // x [C,T,S] -> out [C,T,S]; the channel axis stays axis 0, while the + // weight [K,C] and conv_state [K-1,C,S] carry the same channel + // partition on axis 1. The axis layout (not the op_params flag, whose + // encoding differs between the step and SpecLA variants) determines + // the split. + if (tensor->src[2] != nullptr && + src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_0 && + src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_1 && + src_ss[2].axis == GGML_BACKEND_SPLIT_AXIS_1) { + return {GGML_BACKEND_SPLIT_AXIS_0, {0}, 1, {1}}; + } if (src_ss[0].axis == src_ss[1].axis) { if (src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_0) { return {GGML_BACKEND_SPLIT_AXIS_1, {0}, 1, {1}}; 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..4671b51e9 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; @@ -318,7 +327,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"); @@ -387,8 +398,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); @@ -497,7 +516,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,19 +541,19 @@ 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: { @@ -552,7 +572,7 @@ static void launch_gated_delta_net( 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, 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)); @@ -560,24 +580,24 @@ static void launch_gated_delta_net( 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, 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 +932,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 +956,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.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/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/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..5dbaf6c5b 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,8 +102,16 @@ 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)) 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-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.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.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.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(" Q4_0_ROCMFP4_FAST + Q8_0 2D weights (attn_*, ssm_*) -> Q8_0_ROCMFPX + Q6_K 2D weights (attn_output, output) -> Q6_0_ROCMFPX + token_embd, norms, 1D tensors, conv -> unchanged + +The from_float quantizers live in libggml (built with the rocmfpx types), so +this script requires --libggml (or auto-discovery under server/build-hip*). + +Usage: + python requant_target_rocmfp.py in.gguf out.gguf [--libggml path] +""" +import argparse +import concurrent.futures +import ctypes +import glob +import math +import os +import sys +from pathlib import Path + +import numpy as np + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "deps" / "llama.cpp" / "gguf-py")) + +import gguf # noqa: E402 +from gguf import GGUFReader, GGUFWriter, GGMLQuantizationType # noqa: E402 +from gguf.quants import dequantize # noqa: E402 + +ROCMFP_TYPE_IDS = { + "Q4_0_ROCMFP4": 100, + "Q4_0_ROCMFP4_FAST": 101, + "Q6_0_ROCMFPX": 102, + "Q8_0_ROCMFPX": 103, +} + +# source ggml type -> destination rocmfp type name +REQUANT_POLICY = { + GGMLQuantizationType.Q4_K: "Q4_0_ROCMFP4_FAST", + GGMLQuantizationType.Q8_0: "Q8_0_ROCMFPX", + GGMLQuantizationType.Q6_K: "Q6_0_ROCMFPX", +} + +KEEP_NAMES = {"token_embd.weight"} + + +class GgmlTypeTraits(ctypes.Structure): + _fields_ = [ + ("type_name", ctypes.c_char_p), + ("blck_size", ctypes.c_int64), + ("blck_size_interleave", ctypes.c_int64), + ("type_size", ctypes.c_size_t), + ("is_quantized", ctypes.c_bool), + ("to_float", ctypes.c_void_p), + ("from_float_ref", ctypes.c_void_p), + ] + + +_GGML_FROM_FLOAT_T = ctypes.CFUNCTYPE(None, ctypes.POINTER(ctypes.c_float), + ctypes.c_void_p, ctypes.c_int64) + + +class GgmlLib: + def __init__(self, path: str): + self.path = path + self.lib = ctypes.CDLL(path) + self.lib.ggml_get_type_traits.restype = ctypes.POINTER(GgmlTypeTraits) + self.lib.ggml_get_type_traits.argtypes = [ctypes.c_int] + self.lib.ggml_quantize_init.restype = None + self.lib.ggml_quantize_init.argtypes = [ctypes.c_int] + self.lib.ggml_row_size.restype = ctypes.c_size_t + self.lib.ggml_row_size.argtypes = [ctypes.c_int, ctypes.c_int64] + self.lib.ggml_blck_size.restype = ctypes.c_int64 + self.lib.ggml_blck_size.argtypes = [ctypes.c_int] + self.lib.ggml_type_size.restype = ctypes.c_size_t + self.lib.ggml_type_size.argtypes = [ctypes.c_int] + self._from_float_cache: dict[int, object] = {} + self._workers = max(1, int(os.environ.get("CONV_QUANT_THREADS", + os.cpu_count() or 8))) + + def blck_size(self, type_id: int) -> int: + return int(self.lib.ggml_blck_size(type_id)) + + def type_size(self, type_id: int) -> int: + return int(self.lib.ggml_type_size(type_id)) + + def row_size(self, type_id: int, n_per_row: int) -> int: + return int(self.lib.ggml_row_size(type_id, n_per_row)) + + def _from_float(self, type_id: int): + fn = self._from_float_cache.get(type_id) + if fn is None: + self.lib.ggml_quantize_init(type_id) + traits = self.lib.ggml_get_type_traits(type_id).contents + if not traits.from_float_ref: + raise RuntimeError(f"type {type_id} has no from_float_ref quantizer") + fn = ctypes.cast(traits.from_float_ref, _GGML_FROM_FLOAT_T) + self._from_float_cache[type_id] = fn + return fn + + def quantize(self, type_id: int, arr_f32: np.ndarray) -> np.ndarray: + arr = np.ascontiguousarray(arr_f32, dtype=np.float32) + n_per_row = arr.shape[-1] + nrows = arr.size // n_per_row + blck = self.blck_size(type_id) + if n_per_row % blck != 0: + raise RuntimeError(f"n_per_row {n_per_row} not a multiple of blck_size " + f"{blck} for type {type_id}") + row_bytes = self.row_size(type_id, n_per_row) + total = row_bytes * nrows + dst = (ctypes.c_char * total)() + dst_addr = ctypes.addressof(dst) + src_addr = arr.ctypes.data_as(ctypes.c_void_p).value + fn = self._from_float(type_id) + ELEM = 4 + workers = min(self._workers, nrows) + + def _quant_rows(r0: int): + r1 = min(r0 + chunk_rows, nrows) + nr = r1 - r0 + s = ctypes.cast(src_addr + r0 * n_per_row * ELEM, ctypes.POINTER(ctypes.c_float)) + d = ctypes.cast(dst_addr + r0 * row_bytes, ctypes.c_void_p) + fn(s, d, ctypes.c_int64(nr * n_per_row)) + + if workers <= 1: + chunk_rows = nrows + _quant_rows(0) + else: + chunk_rows = max(1, math.ceil(nrows / workers)) + with concurrent.futures.ThreadPoolExecutor(max_workers=workers) as ex: + list(ex.map(_quant_rows, range(0, nrows, chunk_rows))) + buf = np.frombuffer(bytes(dst), dtype=np.uint8).copy() + return buf.reshape((*arr.shape[:-1], row_bytes)) + + +def find_libggml(explicit: str | None) -> str | None: + if explicit: + return explicit if os.path.exists(explicit) else None + root = Path(__file__).resolve().parent.parent + for pat in ("build-hip*/**/libggml-base*.so", "build-hip*/**/libggml.so", + "build*/**/libggml-base*.so"): + hits = sorted(glob.glob(str(root / pat), recursive=True)) + if hits: + return hits[0] + return None + + +def register_rocmfp_type(type_id: int, lib: GgmlLib) -> None: + bs = lib.blck_size(type_id) + ts = lib.type_size(type_id) + gguf.constants.GGML_QUANT_SIZES[type_id] = (bs, ts) + try: + gguf.quants.GGML_QUANT_SIZES[type_id] = (bs, ts) + except Exception: + pass + + +def copy_metadata(r: GGUFReader, w: GGUFWriter) -> None: + skip = {"GGUF.version", "GGUF.tensor_count", "GGUF.kv_count", "general.architecture"} + T = gguf.GGUFValueType + for f in r.fields.values(): + if f.name in skip: + continue + ftype = f.types[0] + val = f.parts[f.data[0]] + if ftype == T.STRING: + w.add_string(f.name, bytes(val).decode()) + elif ftype == T.ARRAY: + sub = f.types[1] + vals = [f.parts[i] for i in f.data] + if sub == T.STRING: + w.add_array(f.name, [bytes(v).decode() for v in vals]) + else: + w.add_array(f.name, [np.asarray(v)[0].item() for v in vals]) + elif ftype == T.BOOL: + w.add_bool(f.name, bool(val[0])) + elif ftype == T.FLOAT32: + w.add_float32(f.name, float(val[0])) + elif ftype == T.FLOAT64: + w.add_float64(f.name, float(val[0])) + else: + fn = {T.UINT32: w.add_uint32, T.INT32: w.add_int32, + T.UINT64: w.add_uint64, T.INT64: w.add_int64, + T.UINT8: w.add_uint8, T.INT8: w.add_int8, + T.UINT16: w.add_uint16, T.INT16: w.add_int16}[ftype] + fn(f.name, val[0].item()) + + +def main() -> int: + ap = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument("input") + ap.add_argument("output") + ap.add_argument("--libggml", default=None) + ap.add_argument("--skip-output", action="store_true", + help="keep output.weight (lm_head) at its source type") + ap.add_argument("--only-q4k", action="store_true", + help="requantize only Q4_K tensors (FFN); keep Q8_0/Q6_K native") + args = ap.parse_args() + + lib_path = find_libggml(args.libggml) + if not lib_path: + print("error: libggml not found; pass --libggml", file=sys.stderr) + return 1 + lib = GgmlLib(lib_path) + print(f"[info] libggml: {lib_path}") + for tid in ROCMFP_TYPE_IDS.values(): + register_rocmfp_type(tid, lib) + + r = GGUFReader(args.input) + arch = None + for f in r.fields.values(): + if f.name == "general.architecture": + arch = bytes(f.parts[f.data[0]]).decode() + if not arch: + print("error: no general.architecture in input", file=sys.stderr) + return 1 + + w = GGUFWriter(args.output, arch) + copy_metadata(r, w) + + n_q = n_keep = 0 + for t in r.tensors: + shape = [int(x) for x in t.shape] # ggml ne order + dst_name = REQUANT_POLICY.get(t.tensor_type) + keep = ( + dst_name is None or len(shape) != 2 or t.name in KEEP_NAMES or + "norm" in t.name or shape[0] % 256 != 0 or + (args.skip_output and t.name == "output.weight") or + (args.only_q4k and t.tensor_type != GGMLQuantizationType.Q4_K) + ) + if keep: + w.add_tensor(t.name, np.array(t.data), raw_dtype=t.tensor_type) + n_keep += 1 + continue + type_id = ROCMFP_TYPE_IDS[dst_name] + f32 = dequantize(t.data, t.tensor_type).reshape(shape[::-1]) + buf = lib.quantize(type_id, f32) + w.add_tensor(t.name, buf, raw_dtype=type_id) + n_q += 1 + print(f"[requant] {t.name:36s} {t.tensor_type.name:5s} -> {dst_name} {tuple(shape)}") + + print(f"[info] writing {args.output} (requantized {n_q}, kept {n_keep})") + w.write_header_to_file() + w.write_kv_data_to_file() + w.write_tensors_to_file() + w.close() + print("[done]") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/server/src/common/dflash2_head.cpp b/server/src/common/dflash2_head.cpp new file mode 100644 index 000000000..e4d6325ac --- /dev/null +++ b/server/src/common/dflash2_head.cpp @@ -0,0 +1,156 @@ +#include "dflash2_head.h" + +#include "ggml-alloc.h" + +#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_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) { + 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. + std::vector cand_lp; + std::vector cand_ids; + if (!target.project_hidden_to_topk(local_hidden + (size_t)hdim, n_cand, K, /*temperature=*/1.0f, + cand_lp, cand_ids)) { + return false; + } + if (cand_lp.size() != (size_t)n_cand * K || cand_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 (the path picks its predecessor among them). The + // graph shape only depends on (n_cand, K), so it is built once 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_select_chain: 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, cand_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, cand_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_select_chain: graph_compute failed\n"); + return false; + } + std::vector h_hproj((size_t)rank * n_cand); + std::vector h_succ((size_t)rank * n_cand * K); + std::vector h_pred((size_t)rank * n_rows_pred); + ggml_backend_tensor_get_async(backend, g.hproj, h_hproj.data(), 0, sizeof(float) * h_hproj.size()); + ggml_backend_tensor_get_async(backend, g.succ, h_succ.data(), 0, sizeof(float) * h_succ.size()); + ggml_backend_tensor_get_async(backend, g.pred, h_pred.data(), 0, sizeof(float) * h_pred.size()); + ggml_backend_synchronize(backend); + + // 3. Path search: greedy over the candidates, conditioned on the previous pick. + draft_tok.assign((size_t)q_len, last_tok); + int prev_row = 0; // row in h_pred: 0 = seed, 1 + i*K + k = candidate k of position i + for (int i = 0; i < n_cand; ++i) { + const float * pr = h_pred.data() + (size_t)prev_row * rank; + const float * hp = h_hproj.data() + (size_t)i * rank; + float best = -INFINITY; + int best_k = 0; + for (int k = 0; k < K; ++k) { + const float * sc = h_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]; + const float score = cand_lp[(size_t)i * K + k] + dot; + if (score > best) { best = score; best_k = k; } + } + draft_tok[(size_t)i + 1] = cand_ids[(size_t)i * K + best_k]; + prev_row = 1 + i * K + best_k; + } + 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..446a8646c --- /dev/null +++ b/server/src/common/dflash2_head.h @@ -0,0 +1,29 @@ +#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); + +} // namespace dflash::common 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/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/qwen35_backend.cpp b/server/src/qwen35/qwen35_backend.cpp index 112cf0e0a..31bb757f4 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" @@ -23,11 +25,13 @@ #include "qwen3/qwen3_kvflash_scorer.h" #include "ggml-cuda.h" +#include "ggml-backend-impl.h" #include "common/snapshot_backend.h" #include "pflash_ggml_adapter.h" #include "flashprefill.h" #include +#include #include #include #include @@ -217,6 +221,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 +354,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; @@ -815,6 +847,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; @@ -2322,6 +2355,91 @@ 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; +} + +// Resolve the target lm_head to a tensor whose data is readable on the +// draft backend. In tensor-parallel mode the lm_head is a meta tensor whose +// data pointer is a placeholder; the real data is mirrored (a full copy) on +// every rank. Return the simple tensor of the rank that matches the draft +// GPU (a local read on the draft backend), else nullptr so the caller falls +// back to the host chain. +static ggml_tensor * resolve_fused_lm_head(ggml_tensor * lm_head, + const DevicePlacement & device, + int draft_gpu) { + if (!lm_head) return nullptr; + if (!lm_head->buffer || + !ggml_backend_buft_is_meta(ggml_backend_buffer_get_type(lm_head->buffer))) { + return lm_head; // simple tensor: data pointer is real, use as-is + } + // Meta tensor: the lm_head is mirrored (full copy) on every rank. Read + // it from the rank that matches the draft GPU so the fused graph does a + // local read on the draft backend. + for (size_t j = 0; j < device.layer_split_gpus.size(); ++j) { + if (device.layer_split_gpus[j] == draft_gpu) { + ggml_tensor * simple = ggml_backend_meta_simple_tensor(lm_head, j); + return (simple && simple->data) ? simple : nullptr; + } + } + // Draft GPU is not a target rank; fall back to the host chain. + return nullptr; +} + +// 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"); + 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; } - 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; + // [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,128 @@ 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. + // The lm_head must be readable on the draft backend; in + // tensor-parallel mode that is the simple tensor of the + // rank matching the draft GPU (lm_head is mirrored on + // every rank). + std::vector conf_scores; + ggml_tensor * fused_lm_head = + resolve_fused_lm_head(target->lm_head_tensor(), + cfg_.device, cfg_.draft_gpu); + if (fused_lm_head) { + ds_ok = dspark_markov_correct_greedy_chain_fused( + dw_, draft_backend_, fused_lm_head, + 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 +3054,33 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, return false; } } else { - std::vector top_lp; + 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"); + } + ggml_tensor * fused_lm_head = + resolve_fused_lm_head(target->lm_head_tensor(), + cfg_.device, cfg_.draft_gpu); + if (fused_lm_head) { + topk_ok = dspark_markov_project_topk(dw_, draft_backend_, + fused_lm_head, + 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, @@ -3083,7 +3365,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 +3376,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 +3395,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 +3406,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 +3421,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 +3446,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 +3467,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 +3496,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 +3512,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 +3546,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 +3569,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 +3718,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 +3833,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_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..fe2d6404d 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,76 @@ 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; + bool chunked_env = false; + if (can_skip_gdn_intermediate && n_tokens > 1) { + if (const char * s_env = std::getenv("DFLASH27B_CHUNKED")) { + chunked_env = (std::atoi(s_env) != 0); + } + } + 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_env && 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 +1583,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 +1632,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 +1724,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 +1753,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_env || 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 +1821,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_env && n_seq_tokens > 1; ggml_tensor * output = nullptr; @@ -1794,12 +1899,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 +1969,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);