NVLink fix for PR625

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

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

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<ggml_backend_meta_split_state> & 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}};

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
 

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<D, cols_per_block, 1>(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<D, cols_per_block, 1>(ctx, dst, fattn_kernel, nwarps, nbytes_shared, nthreads, need_f16_K, need_f16_V, false);
 }
 
 template <int D, ggml_type type_K, ggml_type type_V>

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><<<grid_dims, block_dims, 0, stream>>>(
                 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><<<grid_dims, block_dims, 0, stream>>>(
                 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><<<grid_dims, block_dims, 0, stream>>>(
                 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,32 +572,32 @@ static void launch_gated_delta_net(
                         gated_delta_net_cuda_grouped_cols<128, cols, width, 32, WRITE_INTER, InterT><<<grouped_grid_dims, grouped_block_dims, 0, stream>>>(
                             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));
                         dim3 grouped_block_dims(64, column_groups_per_block, 1);
                         gated_delta_net_cuda_grouped_cols<128, cols, width, 64, WRITE_INTER, InterT><<<grouped_grid_dims, grouped_block_dims, 0, stream>>>(
                             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><<<grid_dims, block_dims, 0, stream>>>(
                             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><<<grid_dims, block_dims, 0, stream>>>(
                         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><<<grid_dims, block_dims, 0, stream>>>(
                     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<true, true, true, INTER_T>(                          \
                         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<true, false, true, INTER_T>(                         \
                         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<true, false, false, INTER_T>(                        \
                         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<false, true, true, INTER_T>(                         \
                         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<false, false, true, INTER_T>(                        \
                         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<false, false, false, INTER_T>(                       \
                         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)