From 81f1ef2196328541d75ddef328ff691cef6dda48 Mon Sep 17 00:00:00 2001 From: Jake Stevens Date: Thu, 30 Jul 2026 11:50:31 -0700 Subject: [PATCH 1/2] Support grouped state heads in GDN attention (#21402) Summary: Adds separate query-head and state-head configuration to the model definition as well as updates the fused channelwise_gated_delta_rule custom op to support grouped decode and prefill. Reviewed By: YIWENX14 Differential Revision: D113299906 --- extension/llm/custom_ops/custom_ops.py | 26 +- extension/llm/custom_ops/op_sdpa.cpp | 259 +++++++++++++------ extension/llm/custom_ops/op_sdpa_aot.cpp | 4 +- extension/llm/custom_ops/test_gated_delta.py | 51 +++- 4 files changed, 244 insertions(+), 96 deletions(-) diff --git a/extension/llm/custom_ops/custom_ops.py b/extension/llm/custom_ops/custom_ops.py index d09d54c6389..aaed0f387fe 100644 --- a/extension/llm/custom_ops/custom_ops.py +++ b/extension/llm/custom_ops/custom_ops.py @@ -313,26 +313,19 @@ def _validate_channelwise_gated_delta_rule_params( tensor.dtype == torch.float32 ), f"Expected {name} to be float32 but got {tensor.dtype}" - assert ( - query.shape == key.shape - ), f"Expected query and key to have matching shapes but got {query.shape} and {key.shape}" - assert ( - query.shape == decay.shape - ), f"Expected query and decay to have matching shapes but got {query.shape} and {decay.shape}" - assert ( - query.shape[:3] == value.shape[:3] - ), f"Expected query and value to match in batch/head/sequence dims but got {query.shape} and {value.shape}" - assert ( - beta.shape == query.shape[:3] - ), f"Expected beta to match query batch/head/sequence dims but got {beta.shape} and {query.shape}" + assert query.size(0) == key.size(0) and query.shape[2:] == key.shape[2:] + assert key.shape == decay.shape + assert key.shape[:3] == value.shape[:3] + assert beta.shape == key.shape[:3] + assert query.size(1) % key.size(1) == 0 assert initial_state.shape == ( - query.size(0), - query.size(1), + key.size(0), + key.size(1), query.size(3), value.size(3), ), ( "Expected initial_state to have shape " - f"{(query.size(0), query.size(1), query.size(3), value.size(3))} " + f"{(key.size(0), key.size(1), query.size(3), value.size(3))} " f"but got {initial_state.shape}" ) @@ -354,7 +347,8 @@ def channelwise_gated_delta_rule_meta( beta, initial_state, ) - return torch.empty_like(value), torch.empty_like(initial_state) + output_shape = (*query.shape[:3], value.size(3)) + return query.new_empty(output_shape), torch.empty_like(initial_state) def _validate_quantized_sdpa_params( diff --git a/extension/llm/custom_ops/op_sdpa.cpp b/extension/llm/custom_ops/op_sdpa.cpp index d6dbbb4cd39..a86c823b750 100644 --- a/extension/llm/custom_ops/op_sdpa.cpp +++ b/extension/llm/custom_ops/op_sdpa.cpp @@ -219,27 +219,30 @@ bool validate_channelwise_gated_delta_rule_args( "initial_state must be float32"); ET_CHECK_OR_RETURN_FALSE( - query.size(0) == key.size(0) && query.size(1) == key.size(1) && - query.size(2) == key.size(2) && query.size(3) == key.size(3), - "query and key must have matching shapes"); + query.size(0) == key.size(0) && query.size(2) == key.size(2) && + query.size(3) == key.size(3), + "query and key must match in batch/sequence/head-dim"); ET_CHECK_OR_RETURN_FALSE( - query.size(0) == decay.size(0) && query.size(1) == decay.size(1) && - query.size(2) == decay.size(2) && query.size(3) == decay.size(3), - "query and decay must have matching shapes"); + key.size(0) == decay.size(0) && key.size(1) == decay.size(1) && + key.size(2) == decay.size(2) && key.size(3) == decay.size(3), + "key and decay must have matching shapes"); ET_CHECK_OR_RETURN_FALSE( - query.size(0) == value.size(0) && query.size(1) == value.size(1) && - query.size(2) == value.size(2), - "query and value must match in batch/head/sequence dims"); + key.size(0) == value.size(0) && key.size(1) == value.size(1) && + key.size(2) == value.size(2), + "key and value must match in batch/head/sequence dims"); ET_CHECK_OR_RETURN_FALSE( - beta.size(0) == query.size(0) && beta.size(1) == query.size(1) && - beta.size(2) == query.size(2), - "beta must match query batch/head/sequence dims"); + beta.size(0) == key.size(0) && beta.size(1) == key.size(1) && + beta.size(2) == key.size(2), + "beta must match key batch/head/sequence dims"); ET_CHECK_OR_RETURN_FALSE( - initial_state.size(0) == query.size(0) && - initial_state.size(1) == query.size(1) && + key.size(1) > 0 && query.size(1) % key.size(1) == 0, + "query heads must be divisible by key/value heads"); + ET_CHECK_OR_RETURN_FALSE( + initial_state.size(0) == key.size(0) && + initial_state.size(1) == key.size(1) && initial_state.size(2) == query.size(3) && initial_state.size(3) == value.size(3), - "initial_state shape must match [B, H, K, V]"); + "initial_state shape must match [B, Hkv, K, V]"); for (const Tensor* tensor : {&query, &key, &value, &decay, &beta, &initial_state}) { @@ -710,23 +713,29 @@ void channelwise_gated_delta_rule_recurrence( Tensor& out, Tensor& final_state_out) { const auto batch_size = query.size(0); - const auto num_heads = query.size(1); + const auto num_query_heads = query.size(1); + const auto num_kv_heads = key.size(1); + const auto query_heads_per_kv = num_query_heads / num_kv_heads; const auto sequence_length = query.size(2); const auto k_head_dim = query.size(3); const auto v_head_dim = value.size(3); - const auto qk_batch_stride = num_heads * sequence_length * k_head_dim; + const auto query_batch_stride = + num_query_heads * sequence_length * k_head_dim; + const auto kv_batch_stride = num_kv_heads * sequence_length * k_head_dim; const auto qk_head_stride = sequence_length * k_head_dim; const auto qk_seq_stride = k_head_dim; - const auto value_batch_stride = num_heads * sequence_length * v_head_dim; + const auto value_batch_stride = num_kv_heads * sequence_length * v_head_dim; + const auto output_batch_stride = + num_query_heads * sequence_length * v_head_dim; const auto value_head_stride = sequence_length * v_head_dim; const auto value_seq_stride = v_head_dim; - const auto beta_batch_stride = num_heads * sequence_length; + const auto beta_batch_stride = num_kv_heads * sequence_length; const auto beta_head_stride = sequence_length; - const auto state_batch_stride = num_heads * k_head_dim * v_head_dim; + const auto state_batch_stride = num_kv_heads * k_head_dim * v_head_dim; const auto state_head_stride = k_head_dim * v_head_dim; const auto* query_data = query.const_data_ptr(); @@ -756,23 +765,23 @@ void channelwise_gated_delta_rule_recurrence( float* delta = scratch_data + v_head_dim; for (int64_t batch = 0; batch < batch_size; ++batch) { - for (int64_t head = 0; head < num_heads; ++head) { - const auto qk_offset = batch * qk_batch_stride + head * qk_head_stride; + for (int64_t kv_head = 0; kv_head < num_kv_heads; ++kv_head) { + const auto kv_offset = batch * kv_batch_stride + kv_head * qk_head_stride; const auto value_offset = - batch * value_batch_stride + head * value_head_stride; + batch * value_batch_stride + kv_head * value_head_stride; const auto beta_offset = - batch * beta_batch_stride + head * beta_head_stride; + batch * beta_batch_stride + kv_head * beta_head_stride; const auto state_offset = - batch * state_batch_stride + head * state_head_stride; + batch * state_batch_stride + kv_head * state_head_stride; - const auto* q_head = query_data + qk_offset; - const auto* k_head = key_data + qk_offset; - const auto* decay_head = decay_data + qk_offset; + const auto* k_head = key_data + kv_offset; + const auto* decay_head = decay_data + kv_offset; const auto* value_head = value_data + value_offset; const auto* beta_head = beta_data + beta_offset; const auto* initial_state_head = initial_state_data + state_offset; auto* state_head = state_data + state_offset; - auto* output_head = output_data + value_offset; + auto* output_group = output_data + batch * output_batch_stride + + kv_head * query_heads_per_kv * value_head_stride; // Functional: seed the running state from initial_state without mutating // the (read-only) input. @@ -781,12 +790,10 @@ void channelwise_gated_delta_rule_recurrence( } for (int64_t token = 0; token < sequence_length; ++token) { - const auto* q_t = q_head + token * qk_seq_stride; const auto* k_t = k_head + token * qk_seq_stride; const auto* decay_t = decay_head + token * qk_seq_stride; const auto* v_t = value_head + token * value_seq_stride; const float beta_t = beta_head[token]; - auto* output_t = output_head + token * value_seq_stride; // The recurrence needs only two passes over the K x V state S: // pass 1 (read-only): v_pred = (Diag(decay) S)^T k @@ -813,19 +820,50 @@ void channelwise_gated_delta_rule_recurrence( delta[v_idx] = (v_t[v_idx] - v_pred[v_idx]) * beta_t; } - // Pass 2: apply decay + rank-1 write in place, and read back the - // updated row for the output projection in the same sweep. - std::fill(output_t, output_t + v_head_dim, 0.0f); + if (query_heads_per_kv == 1) { + const auto* q_t = query_data + batch * query_batch_stride + + kv_head * qk_head_stride + token * qk_seq_stride; + auto* output_t = output_group + token * value_seq_stride; + std::fill(output_t, output_t + v_head_dim, 0.0f); + for (int64_t k_idx = 0; k_idx < k_head_dim; ++k_idx) { + const float updated_decay = decay_t[k_idx]; + const float key_value = k_t[k_idx]; + const float query_value = q_t[k_idx]; + auto* state_row = state_head + k_idx * v_head_dim; + for (int64_t v_idx = 0; v_idx < v_head_dim; ++v_idx) { + const float updated = + state_row[v_idx] * updated_decay + key_value * delta[v_idx]; + state_row[v_idx] = updated; + output_t[v_idx] += updated * query_value; + } + } + continue; + } + for (int64_t k_idx = 0; k_idx < k_head_dim; ++k_idx) { const float decay_value = decay_t[k_idx]; const float key_value = k_t[k_idx]; - const float query_value = q_t[k_idx]; auto* state_row = state_head + k_idx * v_head_dim; for (int64_t v_idx = 0; v_idx < v_head_dim; ++v_idx) { - const float updated = + state_row[v_idx] = state_row[v_idx] * decay_value + key_value * delta[v_idx]; - state_row[v_idx] = updated; - output_t[v_idx] += updated * query_value; + } + } + + for (int64_t group_head = 0; group_head < query_heads_per_kv; + ++group_head) { + const auto query_head = kv_head * query_heads_per_kv + group_head; + const auto* q_t = query_data + batch * query_batch_stride + + query_head * qk_head_stride + token * qk_seq_stride; + auto* output_t = output_group + group_head * value_head_stride + + token * value_seq_stride; + std::fill(output_t, output_t + v_head_dim, 0.0f); + for (int64_t k_idx = 0; k_idx < k_head_dim; ++k_idx) { + const float query_value = q_t[k_idx]; + const auto* state_row = state_head + k_idx * v_head_dim; + for (int64_t v_idx = 0; v_idx < v_head_dim; ++v_idx) { + output_t[v_idx] += state_row[v_idx] * query_value; + } } } } @@ -918,7 +956,9 @@ void channelwise_gated_delta_rule_chunked( Tensor& out, Tensor& final_state_out) { const auto batch_size = query.size(0); - const auto num_heads = query.size(1); + const auto num_query_heads = query.size(1); + const auto num_kv_heads = key.size(1); + const auto query_heads_per_kv = num_query_heads / num_kv_heads; const auto sequence_length = query.size(2); const auto k_head_dim = query.size(3); const auto v_head_dim = value.size(3); @@ -933,7 +973,7 @@ void channelwise_gated_delta_rule_chunked( auto* state_data = final_state_out.mutable_data_ptr(); auto* output_data = out.mutable_data_ptr(); - // Per (b, h) strides. q/k/decay are [.., T, K]; v/out are [.., T, V]. + // Per-head strides. q/k/decay are [.., T, K]; v/out are [.., T, V]. const auto qk_head_stride = sequence_length * k_head_dim; const auto v_head_stride = sequence_length * v_head_dim; const auto beta_head_stride = sequence_length; @@ -974,18 +1014,20 @@ void channelwise_gated_delta_rule_chunked( } ET_KERNEL_CHECK(ctx, arena != nullptr, MemoryAllocationFailed, ); - // Heads are independent, so parallelize over (b, h); each worker gets its own - // scratch slab (via get_thread_num()) and writes disjoint output/state. + // KV heads are independent. Each worker updates one state and writes all + // query heads assigned to it. torch::executor::parallel_for( 0, - batch_size * num_heads, + batch_size * num_kv_heads, /*grain_size=*/1, [&](int64_t bh_begin, int64_t bh_end) { for (int64_t bh = bh_begin; bh < bh_end; ++bh) { + const int64_t batch = bh / num_kv_heads; + const int64_t kv_head = bh % num_kv_heads; + const int64_t query_head_begin = kv_head * query_heads_per_kv; const int64_t tid = torch::executor::get_thread_num(); const ChunkScratch sc = ChunkScratch::view( arena + tid * size_per_thread, k_head_dim, v_head_dim); - const auto* __restrict__ q_head = query_data + bh * qk_head_stride; const auto* __restrict__ k_head = key_data + bh * qk_head_stride; const auto* __restrict__ v_head = value_data + bh * v_head_stride; const auto* __restrict__ d_head = decay_data + bh * qk_head_stride; @@ -994,7 +1036,6 @@ void channelwise_gated_delta_rule_chunked( const auto* __restrict__ init_head = initial_state_data + bh * state_head_stride; auto* __restrict__ S = state_data + bh * state_head_stride; - auto* __restrict__ o_head = output_data + bh * v_head_stride; // Seed the running state from the (read-only) initial_state. for (const auto idx : c10::irange(state_head_stride)) { @@ -1007,12 +1048,22 @@ void channelwise_gated_delta_rule_chunked( // tokens. const int64_t cur = std::min(CHUNK_SIZE, sequence_length - base); - const auto* __restrict__ q_c = q_head + base * k_head_dim; const auto* __restrict__ k_c = k_head + base * k_head_dim; const auto* __restrict__ v_c = v_head + base * v_head_dim; const auto* __restrict__ d_c = d_head + base * k_head_dim; const auto* __restrict__ beta_c = beta_head + base; - auto* __restrict__ o_c = o_head + base * v_head_dim; + const auto* __restrict__ single_q_c = query_heads_per_kv == 1 + ? query_data + + (batch * num_query_heads + query_head_begin) * + qk_head_stride + + base * k_head_dim + : nullptr; + auto* __restrict__ single_o_c = query_heads_per_kv == 1 + ? output_data + + (batch * num_query_heads + query_head_begin) * + v_head_stride + + base * v_head_dim + : nullptr; // 1. gc = cumsum_r(log decay) (per channel; resets each chunk). // Reordered so the inner loop runs contiguously over k; the running @@ -1033,8 +1084,7 @@ void channelwise_gated_delta_rule_chunked( } } - // 2. Aqk (causal, incl diag, no beta) and A = -beta_r * - // strictlower(Akk). + // 2. A = -beta_r * strictlower(Akk). // ratio = exp(gc[r,k] - gc[j,k]) = eg[r,k] * eg_inv[j,k], // factored so the exp is hoisted out of the O(BT^2) pair loop. // Safe because BT is small enough that exp(-gc) does not @@ -1046,17 +1096,18 @@ void channelwise_gated_delta_rule_chunked( for (const auto k_idx : c10::irange(k_head_dim)) { const float ratio = sc.eg[r * k_head_dim + k_idx] * sc.eg_inv[j * k_head_dim + k_idx]; - aqk += q_c[r * k_head_dim + k_idx] * - k_c[j * k_head_dim + k_idx] * ratio; + if (query_heads_per_kv == 1) { + aqk += single_q_c[r * k_head_dim + k_idx] * + k_c[j * k_head_dim + k_idx] * ratio; + } if (j < r) { akk += k_c[r * k_head_dim + k_idx] * k_c[j * k_head_dim + k_idx] * ratio; } } - // Causal: only j <= r is written, and step 5 reads only j <= r. - // The upper triangle is left uninitialized (arena is not - // zeroed). - sc.Aqk[r * CHUNK_SIZE + j] = aqk; + if (query_heads_per_kv == 1) { + sc.Aqk[r * CHUNK_SIZE + j] = aqk; + } sc.A[r * CHUNK_SIZE + j] = (j < r) ? -beta_r * akk : 0.0f; } } @@ -1106,25 +1157,30 @@ void channelwise_gated_delta_rule_chunked( } } - // 5. pv = u - w @ S_old ; o = (q ⊙ exp gc) @ S_old. Reordered so - // the inner loop runs contiguously over v; pv/o_c accumulate the - // k-sum, then pv is finalized as u - (w@S) with a single - // subtraction (bit-identical to the k-then-subtract form). + // 5. pv = u - w @ S_old. for (const auto r : c10::irange(cur)) { float* __restrict__ pv_row = sc.pv + r * v_head_dim; - float* __restrict__ o_row = o_c + r * v_head_dim; + float* __restrict__ o_row = query_heads_per_kv == 1 + ? single_o_c + r * v_head_dim + : nullptr; for (const auto v_idx : c10::irange(v_head_dim)) { pv_row[v_idx] = 0.0f; - o_row[v_idx] = 0.0f; + if (query_heads_per_kv == 1) { + o_row[v_idx] = 0.0f; + } } for (const auto k_idx : c10::irange(k_head_dim)) { const float wrk = sc.w[r * k_head_dim + k_idx]; - const float qrk = - q_c[r * k_head_dim + k_idx] * sc.eg[r * k_head_dim + k_idx]; const float* __restrict__ s_row = S + k_idx * v_head_dim; + const float qrk = query_heads_per_kv == 1 + ? single_q_c[r * k_head_dim + k_idx] * + sc.eg[r * k_head_dim + k_idx] + : 0.0f; for (const auto v_idx : c10::irange(v_head_dim)) { pv_row[v_idx] += wrk * s_row[v_idx]; - o_row[v_idx] += qrk * s_row[v_idx]; + if (query_heads_per_kv == 1) { + o_row[v_idx] += qrk * s_row[v_idx]; + } } } const float* __restrict__ u_row = sc.u + r * v_head_dim; @@ -1132,14 +1188,65 @@ void channelwise_gated_delta_rule_chunked( pv_row[v_idx] = u_row[v_idx] - pv_row[v_idx]; } } - // 5b. o += Aqk @ pv (AXPY over the contiguous v). - for (const auto r : c10::irange(cur)) { - float* __restrict__ o_row = o_c + r * v_head_dim; - for (const auto j : c10::irange(r + 1)) { - const float a = sc.Aqk[r * CHUNK_SIZE + j]; - const float* __restrict__ pv_row = sc.pv + j * v_head_dim; + + if (query_heads_per_kv == 1) { + for (const auto r : c10::irange(cur)) { + float* __restrict__ o_row = single_o_c + r * v_head_dim; + for (const auto j : c10::irange(r + 1)) { + const float a = sc.Aqk[r * CHUNK_SIZE + j]; + const float* __restrict__ pv_row = sc.pv + j * v_head_dim; + for (const auto v_idx : c10::irange(v_head_dim)) { + o_row[v_idx] += a * pv_row[v_idx]; + } + } + } + } + + // Each query head reads the shared state and pseudo-values. Aqk is + // query-specific, so its scratch is reused across the group. + for (const auto group_head : c10::irange( + query_heads_per_kv == 1 ? int64_t{0} + : query_heads_per_kv)) { + const int64_t query_head = query_head_begin + group_head; + const auto* __restrict__ q_c = query_data + + (batch * num_query_heads + query_head) * qk_head_stride + + base * k_head_dim; + auto* __restrict__ o_c = output_data + + (batch * num_query_heads + query_head) * v_head_stride + + base * v_head_dim; + + for (const auto r : c10::irange(cur)) { + for (const auto j : c10::irange(r + 1)) { + float aqk = 0.0f; + for (const auto k_idx : c10::irange(k_head_dim)) { + const float ratio = sc.eg[r * k_head_dim + k_idx] * + sc.eg_inv[j * k_head_dim + k_idx]; + aqk += q_c[r * k_head_dim + k_idx] * + k_c[j * k_head_dim + k_idx] * ratio; + } + sc.Aqk[r * CHUNK_SIZE + j] = aqk; + } + } + + for (const auto r : c10::irange(cur)) { + float* __restrict__ o_row = o_c + r * v_head_dim; for (const auto v_idx : c10::irange(v_head_dim)) { - o_row[v_idx] += a * pv_row[v_idx]; + o_row[v_idx] = 0.0f; + } + for (const auto k_idx : c10::irange(k_head_dim)) { + const float qrk = q_c[r * k_head_dim + k_idx] * + sc.eg[r * k_head_dim + k_idx]; + const float* __restrict__ s_row = S + k_idx * v_head_dim; + for (const auto v_idx : c10::irange(v_head_dim)) { + o_row[v_idx] += qrk * s_row[v_idx]; + } + } + for (const auto j : c10::irange(r + 1)) { + const float a = sc.Aqk[r * CHUNK_SIZE + j]; + const float* __restrict__ pv_row = sc.pv + j * v_head_dim; + for (const auto v_idx : c10::irange(v_head_dim)) { + o_row[v_idx] += a * pv_row[v_idx]; + } } } } @@ -1260,9 +1367,17 @@ std::tuple channelwise_gated_delta_rule_out( final_state_out.dim_order().data(), final_state_out.dim()), InvalidArgument, ret); + const ::executorch::aten::SizesType output_sizes[] = { + static_cast<::executorch::aten::SizesType>(query.size(0)), + static_cast<::executorch::aten::SizesType>(query.size(1)), + static_cast<::executorch::aten::SizesType>(query.size(2)), + static_cast<::executorch::aten::SizesType>(value.size(3))}; ET_KERNEL_CHECK_MSG( ctx, - resize_tensor(out, value.sizes()) == Error::Ok, + resize_tensor( + out, + ::executorch::aten::ArrayRef<::executorch::aten::SizesType>( + output_sizes, 4)) == Error::Ok, InvalidArgument, ret, "Failed to resize channelwise_gated_delta_rule output tensor."); diff --git a/extension/llm/custom_ops/op_sdpa_aot.cpp b/extension/llm/custom_ops/op_sdpa_aot.cpp index ce8fa88ce1b..2af2518fa26 100644 --- a/extension/llm/custom_ops/op_sdpa_aot.cpp +++ b/extension/llm/custom_ops/op_sdpa_aot.cpp @@ -438,7 +438,9 @@ std::tuple channelwise_gated_delta_rule_aten( const at::Tensor& decay, const at::Tensor& beta, const at::Tensor& initial_state) { - auto output = at::empty_like(value); + auto output = at::empty( + {query.size(0), query.size(1), query.size(2), value.size(3)}, + query.options()); auto final_state = at::empty_like(initial_state); channelwise_gated_delta_rule_out_aten( query, key, value, decay, beta, initial_state, output, final_state); diff --git a/extension/llm/custom_ops/test_gated_delta.py b/extension/llm/custom_ops/test_gated_delta.py index bac8909ee25..edd6f6e7305 100644 --- a/extension/llm/custom_ops/test_gated_delta.py +++ b/extension/llm/custom_ops/test_gated_delta.py @@ -22,14 +22,17 @@ def _make_inputs( seq_len: int = 4, k_head_dim: int = 5, v_head_dim: int = 6, + num_state_heads: int | None = None, ): + if num_state_heads is None: + num_state_heads = num_heads query = torch.randn(batch_size, num_heads, seq_len, k_head_dim) - key = torch.randn(batch_size, num_heads, seq_len, k_head_dim) - value = torch.randn(batch_size, num_heads, seq_len, v_head_dim) + key = torch.randn(batch_size, num_state_heads, seq_len, k_head_dim) + value = torch.randn(batch_size, num_state_heads, seq_len, v_head_dim) # Per-key-channel decay, passed already exponentiated (in (0, 1)). - decay = torch.rand(batch_size, num_heads, seq_len, k_head_dim) - beta = torch.sigmoid(torch.randn(batch_size, num_heads, seq_len)) - initial_state = torch.randn(batch_size, num_heads, k_head_dim, v_head_dim) + decay = torch.rand(batch_size, num_state_heads, seq_len, k_head_dim) + beta = torch.sigmoid(torch.randn(batch_size, num_state_heads, seq_len)) + initial_state = torch.randn(batch_size, num_state_heads, k_head_dim, v_head_dim) return query, key, value, decay, beta, initial_state def _reference_channelwise_gated_delta_rule( @@ -42,7 +45,10 @@ def _reference_channelwise_gated_delta_rule( initial_state: torch.Tensor, ): state = initial_state.clone() - output = torch.zeros_like(value) + output = value.new_zeros( + query.size(0), query.size(1), query.size(2), value.size(3) + ) + heads_per_state = query.size(1) // key.size(1) for token in range(query.size(2)): # Per-key-channel decay: [B, H, K, 1], already exponentiated. @@ -56,10 +62,41 @@ def _reference_channelwise_gated_delta_rule( v_pred = (state * k_t.unsqueeze(-1)).sum(dim=-2) delta = (v_t - v_pred) * beta_t state = state + k_t.unsqueeze(-1) * delta.unsqueeze(-2) - output[:, :, token] = (state * q_t.unsqueeze(-1)).sum(dim=-2) + query_state = state.repeat_interleave(heads_per_state, dim=1) + output[:, :, token] = (query_state * q_t.unsqueeze(-1)).sum(dim=-2) return output, state + def test_channelwise_gated_delta_rule_grouped_matches_reference(self): + torch.manual_seed(0) + + for seq_len in (1, 35): + with self.subTest(seq_len=seq_len): + inputs = self._make_inputs( + batch_size=2, + num_heads=4, + num_state_heads=2, + seq_len=seq_len, + ) + expected_output, expected_state = ( + self._reference_channelwise_gated_delta_rule(*inputs) + ) + actual_output, actual_state = ( + torch.ops.llama.channelwise_gated_delta_rule(*inputs) + ) + + self.assertTrue( + torch.allclose(actual_output, expected_output, atol=1e-3, rtol=1e-3) + ) + self.assertTrue( + torch.allclose(actual_state, expected_state, atol=1e-3, rtol=1e-3) + ) + + def test_channelwise_gated_delta_rule_rejects_uneven_groups(self): + inputs = self._make_inputs(num_heads=3, num_state_heads=2) + with self.assertRaises(RuntimeError): + torch.ops.llama.channelwise_gated_delta_rule(*inputs) + def test_channelwise_gated_delta_rule_matches_reference(self): torch.manual_seed(0) From e3b5b2765f4adb4736274f0dba8f93cc892a561b Mon Sep 17 00:00:00 2001 From: Jake Stevens Date: Thu, 30 Jul 2026 11:50:31 -0700 Subject: [PATCH 2/2] force neon vectorization for ios (#21479) Summary: The default implementation auto-vectorized well for Android builds, but the specific LLVM Clang 21 toolchain + compiler flags did not auto-vectorize, resulting in only scalar operations. This commit selectively forces neon vectorization for iOS Reviewed By: YIWENX14 Differential Revision: D113940046 --- extension/llm/custom_ops/op_sdpa.cpp | 188 ++++++++++++++++++++------- 1 file changed, 138 insertions(+), 50 deletions(-) diff --git a/extension/llm/custom_ops/op_sdpa.cpp b/extension/llm/custom_ops/op_sdpa.cpp index a86c823b750..79ad3388e02 100644 --- a/extension/llm/custom_ops/op_sdpa.cpp +++ b/extension/llm/custom_ops/op_sdpa.cpp @@ -18,11 +18,15 @@ #include #include #include +#include #include #include #include #include #include +#if defined(__APPLE__) && defined(__ARM_NEON) +#include +#endif // parallel_for + get_thread_num have serial fallbacks without a threadpool. #include @@ -698,6 +702,95 @@ Tensor& sdpa_with_kv_cache_out( return output; } +ET_INLINE void vec_scale(float* data, float scale, int64_t size) { + int64_t idx = 0; +#if defined(__APPLE__) && defined(__ARM_NEON) + constexpr int64_t vector_size = 4; + for (; idx + 2 * vector_size <= size; idx += 2 * vector_size) { + const float32x4_t data_vec0 = vld1q_f32(data + idx); + const float32x4_t data_vec1 = vld1q_f32(data + idx + vector_size); + vst1q_f32(data + idx, vmulq_n_f32(data_vec0, scale)); + vst1q_f32(data + idx + vector_size, vmulq_n_f32(data_vec1, scale)); + } + for (; idx + vector_size <= size; idx += vector_size) { + const float32x4_t data_vec = vld1q_f32(data + idx); + vst1q_f32(data + idx, vmulq_n_f32(data_vec, scale)); + } +#endif + for (; idx < size; ++idx) { + data[idx] *= scale; + } +} + +ET_INLINE void vec_axpy(float* y, const float* x, float alpha, int64_t size) { + int64_t idx = 0; +#if defined(__APPLE__) && defined(__ARM_NEON) + constexpr int64_t vector_size = 4; + for (; idx + 2 * vector_size <= size; idx += 2 * vector_size) { + const float32x4_t x_vec0 = vld1q_f32(x + idx); + const float32x4_t x_vec1 = vld1q_f32(x + idx + vector_size); + const float32x4_t y_vec0 = vld1q_f32(y + idx); + const float32x4_t y_vec1 = vld1q_f32(y + idx + vector_size); + vst1q_f32(y + idx, vmlaq_n_f32(y_vec0, x_vec0, alpha)); + vst1q_f32(y + idx + vector_size, vmlaq_n_f32(y_vec1, x_vec1, alpha)); + } + for (; idx + vector_size <= size; idx += vector_size) { + const float32x4_t x_vec = vld1q_f32(x + idx); + const float32x4_t y_vec = vld1q_f32(y + idx); + vst1q_f32(y + idx, vmlaq_n_f32(y_vec, x_vec, alpha)); + } +#endif + for (; idx < size; ++idx) { + y[idx] += x[idx] * alpha; + } +} + +ET_INLINE void vec_state_update_and_output( + float* state, + const float* delta, + float* output, + float decay, + float key, + float query, + int64_t size) { + int64_t idx = 0; +#if defined(__APPLE__) && defined(__ARM_NEON) + constexpr int64_t vector_size = 4; + for (; idx + 2 * vector_size <= size; idx += 2 * vector_size) { + const float32x4_t state_vec0 = vld1q_f32(state + idx); + const float32x4_t state_vec1 = vld1q_f32(state + idx + vector_size); + const float32x4_t delta_vec0 = vld1q_f32(delta + idx); + const float32x4_t delta_vec1 = vld1q_f32(delta + idx + vector_size); + const float32x4_t updated_vec0 = + vmlaq_n_f32(vmulq_n_f32(delta_vec0, key), state_vec0, decay); + const float32x4_t updated_vec1 = + vmlaq_n_f32(vmulq_n_f32(delta_vec1, key), state_vec1, decay); + const float32x4_t output_vec0 = vld1q_f32(output + idx); + const float32x4_t output_vec1 = vld1q_f32(output + idx + vector_size); + vst1q_f32(state + idx, updated_vec0); + vst1q_f32(state + idx + vector_size, updated_vec1); + vst1q_f32(output + idx, vmlaq_n_f32(output_vec0, updated_vec0, query)); + vst1q_f32( + output + idx + vector_size, + vmlaq_n_f32(output_vec1, updated_vec1, query)); + } + for (; idx + vector_size <= size; idx += vector_size) { + const float32x4_t state_vec = vld1q_f32(state + idx); + const float32x4_t delta_vec = vld1q_f32(delta + idx); + const float32x4_t updated_vec = + vmlaq_n_f32(vmulq_n_f32(delta_vec, key), state_vec, decay); + const float32x4_t output_vec = vld1q_f32(output + idx); + vst1q_f32(state + idx, updated_vec); + vst1q_f32(output + idx, vmlaq_n_f32(output_vec, updated_vec, query)); + } +#endif + for (; idx < size; ++idx) { + const float updated = state[idx] * decay + key * delta[idx]; + state[idx] = updated; + output[idx] += updated * query; + } +} + // Shared full-sequence recurrence: streams the K x V state twice per token (see // the two-pass comments inline). Used by the T==1 decode path and as the exact // prefill fallback when decay contains non-positive entries -- the chunked @@ -810,9 +903,7 @@ void channelwise_gated_delta_rule_recurrence( const float decay_value = decay_t[k_idx]; const float key_value = k_t[k_idx]; const auto* state_row = state_head + k_idx * v_head_dim; - for (int64_t v_idx = 0; v_idx < v_head_dim; ++v_idx) { - v_pred[v_idx] += (state_row[v_idx] * decay_value) * key_value; - } + vec_axpy(v_pred, state_row, decay_value * key_value, v_head_dim); } // delta = beta * (v - v_pred). @@ -826,16 +917,18 @@ void channelwise_gated_delta_rule_recurrence( auto* output_t = output_group + token * value_seq_stride; std::fill(output_t, output_t + v_head_dim, 0.0f); for (int64_t k_idx = 0; k_idx < k_head_dim; ++k_idx) { - const float updated_decay = decay_t[k_idx]; + const float decay_value = decay_t[k_idx]; const float key_value = k_t[k_idx]; const float query_value = q_t[k_idx]; auto* state_row = state_head + k_idx * v_head_dim; - for (int64_t v_idx = 0; v_idx < v_head_dim; ++v_idx) { - const float updated = - state_row[v_idx] * updated_decay + key_value * delta[v_idx]; - state_row[v_idx] = updated; - output_t[v_idx] += updated * query_value; - } + vec_state_update_and_output( + state_row, + delta, + output_t, + decay_value, + key_value, + query_value, + v_head_dim); } continue; } @@ -844,9 +937,33 @@ void channelwise_gated_delta_rule_recurrence( const float decay_value = decay_t[k_idx]; const float key_value = k_t[k_idx]; auto* state_row = state_head + k_idx * v_head_dim; - for (int64_t v_idx = 0; v_idx < v_head_dim; ++v_idx) { - state_row[v_idx] = - state_row[v_idx] * decay_value + key_value * delta[v_idx]; + int64_t idx = 0; +#if defined(__APPLE__) && defined(__ARM_NEON) + constexpr int64_t vector_size = 4; + for (; idx + 2 * vector_size <= v_head_dim; idx += 2 * vector_size) { + const float32x4_t state_vec0 = vld1q_f32(state_row + idx); + const float32x4_t state_vec1 = + vld1q_f32(state_row + idx + vector_size); + const float32x4_t delta_vec0 = vld1q_f32(delta + idx); + const float32x4_t delta_vec1 = vld1q_f32(delta + idx + vector_size); + const float32x4_t updated_vec0 = vmlaq_n_f32( + vmulq_n_f32(state_vec0, decay_value), delta_vec0, key_value); + const float32x4_t updated_vec1 = vmlaq_n_f32( + vmulq_n_f32(state_vec1, decay_value), delta_vec1, key_value); + vst1q_f32(state_row + idx, updated_vec0); + vst1q_f32(state_row + idx + vector_size, updated_vec1); + } + for (; idx + vector_size <= v_head_dim; idx += vector_size) { + const float32x4_t state_vec = vld1q_f32(state_row + idx); + const float32x4_t delta_vec = vld1q_f32(delta + idx); + const float32x4_t updated_vec = vmlaq_n_f32( + vmulq_n_f32(state_vec, decay_value), delta_vec, key_value); + vst1q_f32(state_row + idx, updated_vec); + } +#endif + for (; idx < v_head_dim; ++idx) { + state_row[idx] = + state_row[idx] * decay_value + key_value * delta[idx]; } } @@ -861,9 +978,7 @@ void channelwise_gated_delta_rule_recurrence( for (int64_t k_idx = 0; k_idx < k_head_dim; ++k_idx) { const float query_value = q_t[k_idx]; const auto* state_row = state_head + k_idx * v_head_dim; - for (int64_t v_idx = 0; v_idx < v_head_dim; ++v_idx) { - output_t[v_idx] += state_row[v_idx] * query_value; - } + vec_axpy(output_t, state_row, query_value, v_head_dim); } } } @@ -871,21 +986,6 @@ void channelwise_gated_delta_rule_recurrence( } } -// T == 1 (autoregressive decode) route. -void channelwise_gated_delta_rule_decode( - RuntimeContext& ctx, - const Tensor& query, - const Tensor& key, - const Tensor& value, - const Tensor& decay, - const Tensor& beta, - const Tensor& initial_state, - Tensor& out, - Tensor& final_state_out) { - channelwise_gated_delta_rule_recurrence( - ctx, query, key, value, decay, beta, initial_state, out, final_state_out); -} - // Chunk length for the chunkwise recurrence. Kept small so that // exp(-cumsum(log decay)) over a chunk stays finite (see eg_inv in step 1). constexpr int64_t CHUNK_SIZE = 32; @@ -1151,9 +1251,7 @@ void channelwise_gated_delta_rule_chunked( w_row[k_idx] += arj * eg_row[k_idx] * k_row[k_idx]; } const float* __restrict__ v_row = v_c + j * v_head_dim; - for (const auto v_idx : c10::irange(v_head_dim)) { - u_row[v_idx] += arj * v_row[v_idx]; - } + vec_axpy(u_row, v_row, arj, v_head_dim); } } @@ -1195,9 +1293,7 @@ void channelwise_gated_delta_rule_chunked( for (const auto j : c10::irange(r + 1)) { const float a = sc.Aqk[r * CHUNK_SIZE + j]; const float* __restrict__ pv_row = sc.pv + j * v_head_dim; - for (const auto v_idx : c10::irange(v_head_dim)) { - o_row[v_idx] += a * pv_row[v_idx]; - } + vec_axpy(o_row, pv_row, a, v_head_dim); } } } @@ -1237,16 +1333,12 @@ void channelwise_gated_delta_rule_chunked( const float qrk = q_c[r * k_head_dim + k_idx] * sc.eg[r * k_head_dim + k_idx]; const float* __restrict__ s_row = S + k_idx * v_head_dim; - for (const auto v_idx : c10::irange(v_head_dim)) { - o_row[v_idx] += qrk * s_row[v_idx]; - } + vec_axpy(o_row, s_row, qrk, v_head_dim); } for (const auto j : c10::irange(r + 1)) { const float a = sc.Aqk[r * CHUNK_SIZE + j]; const float* __restrict__ pv_row = sc.pv + j * v_head_dim; - for (const auto v_idx : c10::irange(v_head_dim)) { - o_row[v_idx] += a * pv_row[v_idx]; - } + vec_axpy(o_row, pv_row, a, v_head_dim); } } } @@ -1263,15 +1355,11 @@ void channelwise_gated_delta_rule_chunked( sc.de[r] = std::exp(gc_last - sc.gc[r * k_head_dim + k_idx]); } float* __restrict__ S_row = S + k_idx * v_head_dim; - for (const auto v_idx : c10::irange(v_head_dim)) { - S_row[v_idx] *= eg_last; - } + vec_scale(S_row, eg_last, v_head_dim); for (const auto r : c10::irange(cur)) { const float c = k_c[r * k_head_dim + k_idx] * sc.de[r]; const float* __restrict__ pv_row = sc.pv + r * v_head_dim; - for (const auto v_idx : c10::irange(v_head_dim)) { - S_row[v_idx] += c * pv_row[v_idx]; - } + vec_axpy(S_row, pv_row, c, v_head_dim); } } } // chunk @@ -1395,7 +1483,7 @@ std::tuple channelwise_gated_delta_rule_out( // -- non-positive decay (log undefined) or tiny positive decay whose // within-chunk cumulative -log(decay) would overflow exp(-gc) to inf -> NaN. if (query.size(2) == 1) { - channelwise_gated_delta_rule_decode( + channelwise_gated_delta_rule_recurrence( ctx, query, key,