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 <bool KDA, bool keep_rs_t, bool RAW, bool G_PRECOMPUTED>
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<KDA_, KEEP_, RAW_, PRE_>(q_d, k_d, v_d, g_d, b_d, rb_d, ra_d, s_d, dst_d, state_d, \
+ launch_gated_delta_net<KDA_, KEEP_, RAW_, PRE_>(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<llm_graph_input_rs> 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<int64_t>(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<std::unique_ptr<test_case>> 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