diff --git a/ggml/src/ggml-cuda/gated_delta_net.cu b/ggml/src/ggml-cuda/gated_delta_net.cu index 6ca696fc..294d97df 100644 --- a/ggml/src/ggml-cuda/gated_delta_net.cu +++ b/ggml/src/ggml-cuda/gated_delta_net.cu @@ -60,6 +60,7 @@ gated_delta_net_cuda(const float * q, const float * raw_dt_bias, const float * raw_a, const float * curr_state, + const int32_t * state_rows, float * dst, float * state, int64_t H, @@ -95,12 +96,9 @@ gated_delta_net_cuda(const float * q, float * attn_data = dst; - // input state holds s0 only: [S_v, S_v, H, n_seqs] — seq stride is D = H * S_v * S_v. // output state layout (per-slot D * n_seqs) — same per-(seq,head) offset as before. - const int64_t state_in_offset = sequence * H * S_v * S_v + h_idx * S_v * S_v; const int64_t state_out_offset = (sequence * H + h_idx) * S_v * S_v; state += state_out_offset; - curr_state += state_in_offset; attn_data += (sequence * n_tokens * H + h_idx) * S_v; constexpr int warp_size = ggml_cuda_get_physical_warp_size() < S_v ? ggml_cuda_get_physical_warp_size() : S_v; @@ -110,6 +108,11 @@ gated_delta_net_cuda(const float * q, // state is stored transposed: M[col][i] = S[i][col], row col is contiguous ggml_cuda_pdl_sync(); + // input state: s0 of each seq, [S_v, S_v, H, n_seqs] with seq stride D = H * S_v * S_v; in rows mode the + // cache itself, seq's s0 at row state_rows[sequence]. Every thread reads its shard before it writes any + // snapshot, so a seq may read the cache rows its own snapshots overwrite. + const int64_t state_in_row = state_rows ? (int64_t) state_rows[sequence] : (int64_t) sequence; + curr_state += (state_in_row * H + h_idx) * S_v * S_v; #pragma unroll for (int c = 0; c < cols_per_warp; ++c) { #pragma unroll @@ -304,7 +307,7 @@ template static void launch_gated_delta_net( const float * q_d, const float * k_d, const float * v_d, const float * g_d, const float * b_d, const float * rb_d, const float * ra_d, const float * s_d, - float * dst_d, float * state_d, + const int32_t * rows_d, float * dst_d, float * state_d, int64_t S_v, int64_t H, int64_t n_tokens, int64_t n_seqs, int64_t sq1, int64_t sq2, int64_t sq3, int64_t sv1, int64_t sv2, int64_t sv3, @@ -327,26 +330,26 @@ static void launch_gated_delta_net( switch (S_v) { case 16: ggml_cuda_kernel_launch(gated_delta_net_cuda<16, KDA, keep_rs_t, RAW, G_PRECOMPUTED>, launch_params, - q_d, k_d, v_d, g_d, b_d, rb_d, ra_d, s_d, dst_d, state_d, H, + q_d, k_d, v_d, g_d, b_d, rb_d, ra_d, s_d, rows_d, dst_d, state_d, H, n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3, sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, state_slot_stride, K); break; case 32: ggml_cuda_kernel_launch(gated_delta_net_cuda<32, KDA, keep_rs_t, RAW, G_PRECOMPUTED>, launch_params, - q_d, k_d, v_d, g_d, b_d, rb_d, ra_d, s_d, dst_d, state_d, H, + q_d, k_d, v_d, g_d, b_d, rb_d, ra_d, s_d, rows_d, dst_d, state_d, H, n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3, sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, state_slot_stride, K); break; case 64: { ggml_cuda_kernel_launch(gated_delta_net_cuda<64, KDA, keep_rs_t, RAW, G_PRECOMPUTED>, launch_params, - q_d, k_d, v_d, g_d, b_d, rb_d, ra_d, s_d, dst_d, state_d, H, + q_d, k_d, v_d, g_d, b_d, rb_d, ra_d, s_d, rows_d, dst_d, state_d, H, n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3, sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, state_slot_stride, K); break; } case 128: { ggml_cuda_kernel_launch(gated_delta_net_cuda<128, KDA, keep_rs_t, RAW, G_PRECOMPUTED>, launch_params, - q_d, k_d, v_d, g_d, b_d, rb_d, ra_d, s_d, dst_d, state_d, H, + q_d, k_d, v_d, g_d, b_d, rb_d, ra_d, s_d, rows_d, dst_d, state_d, H, n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3, sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, state_slot_stride, K); break; @@ -395,6 +398,11 @@ static void ggml_cuda_op_gated_delta_net_impl( const float * s_d = (const float *) src_state->data; float * dst_d = (float *) dst->data; + // rows mode: src_state is the 2D cache view and src[6] holds each seq's cache row + const ggml_tensor * src_rows = dst->src[6]; + const int32_t * rows_d = src_rows ? (const int32_t *) src_rows->data : nullptr; + GGML_ASSERT(src_rows == nullptr || (src_rows->type == GGML_TYPE_I32 && src_rows->ne[0] == n_seqs)); + GGML_ASSERT(ggml_is_contiguous_rows(src_q)); GGML_ASSERT(ggml_is_contiguous_rows(src_k)); GGML_ASSERT(ggml_is_contiguous_rows(src_v)); @@ -452,7 +460,7 @@ static void ggml_cuda_op_gated_delta_net_impl( } #define GDN_LAUNCH(KDA_, KEEP_, RAW_, PRE_) \ - launch_gated_delta_net(q_d, k_d, v_d, g_d, b_d, rb_d, ra_d, s_d, dst_d, state_d, \ + launch_gated_delta_net(q_d, k_d, v_d, g_d, b_d, rb_d, ra_d, s_d, rows_d, dst_d, state_d, \ S_v, H, n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3, \ sb1, sb2, sb3, neqk1, rq3, scale, state_slot_stride, K, stream) diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 2f5ad147..d1f12fcc 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -5785,9 +5785,8 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g case GGML_OP_RWKV_WKV7: return true; case GGML_OP_GATED_DELTA_NET: - // rows-indexed state read (src[6]) not implemented on CUDA yet; - // reject so it falls back instead of silently reading src[5] as a scratch - if (op->src[6] != NULL) { + // rows-indexed state read: src[6] holds I32 cache rows + if (op->src[6] != NULL && op->src[6]->type != GGML_TYPE_I32) { return false; } //TODO: enable once MUSA compiler is solved https://github.com/ggml-org/llama.cpp/pull/19504#issuecomment-4018634327 diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 9ad7fa92..9a853e99 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -417,6 +417,8 @@ bool llm_graph_input_rs::can_reuse_rs(const llama_memory_recurrent_context * mct res &= head == mctx_cur->get_head(); res &= rs_z == mctx_cur->get_rs_z(); + res &= s_copy_main_in_place == mctx_cur->s_copy_main_in_place(params.ubatch.n_seqs); + return res; } @@ -3579,6 +3581,8 @@ static std::unique_ptr build_rs_inp_impl( inp->head = mctx_cur->get_head(); inp->rs_z = mctx_cur->get_rs_z(); + inp->s_copy_main_in_place = mctx_cur->s_copy_main_in_place(n_seqs); + return inp; } @@ -3619,16 +3623,12 @@ ggml_tensor * llm_graph_context::build_rs_cache_view( // same cache hygiene as build_rs, minus the main gather (the consumer reads // per-seq rows via inp->s_copy_main directly, inside the GDN op). // - // KNOWN LIMITATION (tracked follow-up): build_rs gathers the main rows - // BEFORE this extra relocation, so an overlapping main row is read before - // being overwritten. rows mode defers the main read into the consumer, and - // s_copy() maps a main row to an arbitrary cache slot (idx*size + src0), - // which can fall inside the extra destination [rs_head+n_seqs, rs_head+n_rs) - // during a cache reorder -- so this relocation could clobber a main row the - // consumer will later read. Not reachable on the current single-sequence - // decode path, but it is a real multi-sequence hazard; the correct fix is - // to order the relocation AFTER the GDN read (build_rs's read-before-write - // ordering), which is a graph-dependency refactor left as follow-up. + // build_rs gathers the main rows BEFORE this extra relocation; rows mode + // defers the main read into the consumer, after it. That is only safe while + // every main row lies in its own seq's cell (inp->s_copy_main_in_place): + // during a cache reorder or a fresh seq reading another cell's zeroed row, + // the relocation or another seq's snapshot write could clobber a row the + // consumer reads later, so callers keep build_rs then. ggml_tensor * state_zero = ggml_view_1d(ctx0, states, state_size*(rs_zero >= 0), rs_zero*states->nb[1]*(rs_zero >= 0)); ggml_build_forward_expand(gf, ggml_scale_inplace(ctx0, state_zero, 0)); diff --git a/src/llama-graph.h b/src/llama-graph.h index 6c2db642..2c7c98e8 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -351,6 +351,9 @@ public: // used in view offsets, need to match for valid graph reuse uint32_t head; int32_t rs_z; + + // mctx->s_copy_main_in_place(n_seqs) at build time; GDN rows mode reads states in place only then + bool s_copy_main_in_place = false; }; class llm_graph_input_cross_embd : public llm_graph_input_i { diff --git a/src/llama-memory-recurrent.cpp b/src/llama-memory-recurrent.cpp index e2990972..8e59c5d8 100644 --- a/src/llama-memory-recurrent.cpp +++ b/src/llama-memory-recurrent.cpp @@ -1270,3 +1270,16 @@ int32_t llama_memory_recurrent_context::s_copy(int i) const { } return (int32_t)(idx * mem->size) + src0; } + +bool llama_memory_recurrent_context::s_copy_main_in_place(uint32_t n_seqs) const { + if (is_full) { + return false; + } + for (uint32_t i = 0; i < n_seqs; ++i) { + const uint32_t cell_idx = i + mem->head; + if (cell_idx >= mem->size || mem->cells[cell_idx].src0 != (int32_t) cell_idx) { + return false; + } + } + return true; +} diff --git a/src/llama-memory-recurrent.h b/src/llama-memory-recurrent.h index b13b7b74..a83c766e 100644 --- a/src/llama-memory-recurrent.h +++ b/src/llama-memory-recurrent.h @@ -173,6 +173,11 @@ public: int32_t s_copy(int i) const; + // true when every main state row s_copy(i), i < n_seqs, lies in the seq's own cell head + i (any + // snapshot slot): no other seq's snapshot write and no extra-state relocation touches it, so a + // consumer may read the rows in place instead of gathering them first + bool s_copy_main_in_place(uint32_t n_seqs) const; + private: const llama_memory_status status; diff --git a/src/models/delta-net-base.cpp b/src/models/delta-net-base.cpp index ae67400c..c4ccbd51 100644 --- a/src/models/delta-net-base.cpp +++ b/src/models/delta-net-base.cpp @@ -611,7 +611,7 @@ ggml_tensor * llm_build_delta_net_base::build_recurrent_attn( // op writes the last min(n_seq_tokens, K) snapshots; trailing slots are left unwritten const int64_t n_written = std::min(n_seq_tokens, K); - if (state_rows) { + if (state_rows && gdn_state_rows_set_rows) { // rows mode: scatter the compact snapshot region [D, n_written*n_seqs] // (slot-major, starting right after the attention scores) into the // cache rows slot*mem_size + kv_head + seq -- same destinations as the diff --git a/src/models/models.h b/src/models/models.h index a8213309..c5d630ba 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -92,6 +92,10 @@ struct llm_build_delta_net_base : public llm_graph_context { ggml_tensor * gdn_raw_dt_bias = nullptr; ggml_tensor * gdn_raw_a = nullptr; + // rows mode scatters snapshots with SET_ROWS (Metal folds it into the op); false keeps the strided + // cpy, which the CUDA backend folds into the op instead + bool gdn_state_rows_set_rows = true; + ggml_tensor * build_recurrent_attn( llm_graph_input_rs * inp, ggml_tensor * ssm_states_all, @@ -2267,7 +2271,7 @@ struct llama_model_qwen35 : public llama_model_base { graph(const llama_model & model, const llm_graph_params & params); // device-dependent path choices, scanned once per graph build (not per layer) - bool gdn_state_rows_dev_ok = true; // every GPU device is Metal: fused GDN may read state rows in place + bool gdn_state_rows_dev_ok = true; // every GPU device is Metal or CUDA: fused GDN may read state rows in place bool gdn_raw_gates_dev_ok = true; // every device is CPU/Metal/CUDA/ROCm/MUSA: fused GDN takes raw gates private: ggml_tensor * build_layer_attn( diff --git a/src/models/qwen35.cpp b/src/models/qwen35.cpp index df95c5f8..6291ef28 100644 --- a/src/models/qwen35.cpp +++ b/src/models/qwen35.cpp @@ -157,7 +157,10 @@ llama_model_qwen35::graph::graph(const llama_model & model, const llm_graph_para const bool is_gpu = ggml_backend_dev_type(ldev.dev) == GGML_BACKEND_DEVICE_TYPE_GPU || ggml_backend_dev_type(ldev.dev) == GGML_BACKEND_DEVICE_TYPE_IGPU; if (is_gpu && strcmp(reg_name, "MTL") != 0) { - gdn_state_rows_dev_ok = false; + gdn_state_rows_set_rows = false; + if (strcmp(reg_name, "CUDA") != 0) { + gdn_state_rows_dev_ok = false; + } } if (strcmp(reg_name, "MTL") != 0 && strcmp(reg_name, "CUDA") != 0 && strcmp(reg_name, "ROCm") != 0 && strcmp(reg_name, "MUSA") != 0 && strcmp(reg_name, "CPU") != 0) { @@ -473,13 +476,15 @@ ggml_tensor * llama_model_qwen35::graph::build_layer_attn_linear( // ring path: read per-seq live state directly from the cache inside the // fused GDN op (rows mode) instead of a gather per layer. // GGML_GDN_STATE_GATHER=1 restores the legacy gathered path (A/B). - // rows mode (the src[6] variant) is implemented on CPU and Metal only; + // rows mode (the src[6] variant) is implemented on CPU, Metal and CUDA; // other GPU backends reject it in supports_op, which would silently move // the whole recurrent op to CPU -- keep the gathered form unless every - // GPU device in the model is Metal. + // GPU device in the model is Metal or CUDA. Rows are read in place only + // when every seq reads its own cell (see build_rs_cache_view). static const bool gdn_state_rows_env = getenv("GGML_GDN_STATE_GATHER") == nullptr; - const bool gdn_state_rows = gdn_state_rows_env && gdn_state_rows_dev_ok && cparams.n_rs_seq > 0; + const bool gdn_state_rows = gdn_state_rows_env && gdn_state_rows_dev_ok && cparams.n_rs_seq > 0 && + inp->s_copy_main_in_place; ggml_tensor * state; if (gdn_state_rows) { diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 57ead864..b41998e3 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -10438,6 +10438,12 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 4, 32, 8, 1, 1, false, false, /*K=*/3, /*rows=*/true)); // snapshot overflow (n_tokens > K) test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 4, 64, 16, 2, 1, false, false, /*K=*/4, /*rows=*/true)); test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 4, 64, 4, 2, 1, false, true, /*K=*/4, /*rows=*/true)); // KDA + // qwen35 MTP shapes: S_v 128, raw gates, K 3, decode / verify widths and a prefill tail + for (int64_t n_tokens : { 1, 2, 3, 8, 37 }) { + for (int64_t n_seqs : { 1, 2 }) { + test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 4, 128, n_tokens, n_seqs, 1, false, false, /*K=*/3, /*rows=*/true, -1, /*raw=*/true)); + } + } #if 0 // these tests are disabled to save execution time, sbut they can be handy for debugging