From 0ec3e739be7b14da2337cdf947d4cbcb022d6069 Mon Sep 17 00:00:00 2001 From: Kawrakow Date: Wed, 25 Feb 2026 05:27:55 +0000 Subject: [PATCH] Don't re-apply L2 norm - it has already been done --- ggml/src/ggml-cuda/delta-net.cu | 62 ++++++++++----------------------- ggml/src/iqk/iqk_mul_mat.cpp | 30 ++++------------ 2 files changed, 25 insertions(+), 67 deletions(-) diff --git a/ggml/src/ggml-cuda/delta-net.cu b/ggml/src/ggml-cuda/delta-net.cu index 367ae67d..54ad3c85 100644 --- a/ggml/src/ggml-cuda/delta-net.cu +++ b/ggml/src/ggml-cuda/delta-net.cu @@ -84,55 +84,34 @@ __global__ void delta_net_recurrent_f32( float * sQ = smem; // HEAD_DIM float * sK = sQ + HEAD_DIM; // HEAD_DIM float * sV = sK + HEAD_DIM; // HEAD_DIM - float * sKBeta = sV + HEAD_DIM; // HEAD_DIM (plain k for state update) - float * sVBeta = sKBeta + HEAD_DIM; // HEAD_DIM (v * sigmoid(beta)) - float * sOut = sVBeta + HEAD_DIM; // HEAD_DIM - float * sKCumdecay = sOut + HEAD_DIM; // HEAD_DIM (k * sigmoid(beta) * exp(g)) - float * sVNew = sKCumdecay + HEAD_DIM; // HEAD_DIM (v_beta - v_prime) + float * sVNew = sV + HEAD_DIM; // HEAD_DIM const float scale = rsqrtf((float)HEAD_DIM); __shared__ float sum_helper[block_size/WARP_SIZE]; // Copy initial state to output buffer (will be updated in place) - for (int i = tid; i < HEAD_DIM * HEAD_DIM; i += blockDim.x) { + for (int i = tid; i < HEAD_DIM * HEAD_DIM; i += block_size) { state_dst[i] = state_src[i]; } - __syncthreads(); // Process each token sequentially for (int64_t t = 0; t < n_tokens; t++) { - float q_sq = 0.0f; - float k_sq = 0.0f; - for (int i = tid; i < HEAD_DIM; i += blockDim.x) { + float sum_kq = 0.0f; + for (int i = tid; i < HEAD_DIM; i += block_size) { sQ[i] = q_ptr[t * qkv_stride_token + i]; sK[i] = k_ptr[t * qkv_stride_token + i]; sV[i] = v_ptr[t * qkv_stride_token + i]; - q_sq += sQ[i] * sQ[i]; - k_sq += sK[i] * sK[i]; + sum_kq += sK[i] * sQ[i]; } - q_sq = reduce_sum(q_sq, sum_helper); - k_sq = reduce_sum(k_sq, sum_helper); - - float q_norm = rsqrtf(q_sq + eps); - float k_norm = rsqrtf(k_sq + eps); + sum_kq = reduce_sum(sum_kq, sum_helper); float beta_val = sigmoid_f(beta_ptr[t]); float decay = expf(fminf(g_ptr[t], 50.0f)); - float sum = 0; - for (int i = tid; i < HEAD_DIM; i += blockDim.x) { - sQ[i] = sQ[i] * q_norm * scale; - sK[i] = sK[i] * k_norm; - sKBeta[i] = sK[i]; - sVBeta[i] = sV[i] * beta_val; - sKCumdecay[i] = sK[i] * beta_val * decay; - sum += sK[i] * sQ[i]; - } - float attn_score = reduce_sum(sum, sum_helper); - //__syncthreads(); + float attn_score = sum_kq * scale; for (int row_out = warp_id; row_out < HEAD_DIM; row_out += NUM_WARPS) { float sum1 = 0.0f; @@ -140,16 +119,15 @@ __global__ void delta_net_recurrent_f32( #pragma unroll for (int col = lane_id; col < HEAD_DIM; col += WARP_SIZE) { float sval = state_dst[row_out + col * HEAD_DIM]; - sum1 += sval * sKCumdecay[col]; + sum1 += sval * sK[col]; sum2 += sval * sQ[col]; } - sum1 = warp_reduce_sum(sum1); - sum2 = warp_reduce_sum(sum2); + sum1 = warp_reduce_sum(sum1) * beta_val * decay; + sum2 = warp_reduce_sum(sum2) * scale * decay; if (lane_id == 0) { - sVNew[row_out] = sVBeta[row_out] - sum1; + sVNew[row_out] = sV[row_out] * beta_val - sum1; float v_attn = sVNew[row_out] * attn_score; - //sOut[row_out] = sum2 * decay + v_attn; - out_base[t * out_token_stride + row_out] = sum2 * decay + v_attn; + out_base[t * out_token_stride + row_out] = sum2 + v_attn; } } __syncthreads(); @@ -158,18 +136,14 @@ __global__ void delta_net_recurrent_f32( #pragma unroll for (int row = lane_id; row < HEAD_DIM; row += WARP_SIZE) { float state_val = state_dst[row + out_dim * HEAD_DIM]; - float safe_decay = decay; - if (isnan(safe_decay) || isinf(safe_decay)) { - safe_decay = 1.0f; - } - float new_state_val = safe_decay * state_val + sVNew[row] * sKBeta[out_dim]; + float new_state_val = decay * state_val + sVNew[row] * sK[out_dim]; new_state_val = fminf(fmaxf(new_state_val, -1e6f), 1e6f); state_dst[row + out_dim * HEAD_DIM] = new_state_val; } } - if (t < n_tokens - 1) { - __syncthreads(); - } + //if (t < n_tokens - 1) { + // __syncthreads(); + //} } } @@ -410,7 +384,9 @@ static void delta_net_f32_cuda( // Shared memory: 9 * head_dim (for Q, K, V, KBeta, VBeta, Out, KCumdecay, VPrime, VNew) // Plus 6 floats for Norm[2], g_val, beta_val, decay, attn_score - const size_t smem_size = (9 * head_dim + 6) * sizeof(float); + //const size_t smem_size = (9 * head_dim + 6) * sizeof(float); + //const size_t smem_size = (4 * head_dim + 2 * n_tokens) * sizeof(float); + const size_t smem_size = 4 * head_dim * sizeof(float); // Use templated kernel for common head dimensions, generic for others if (head_dim == 64) { diff --git a/ggml/src/iqk/iqk_mul_mat.cpp b/ggml/src/iqk/iqk_mul_mat.cpp index 956b33e0..9a6ecd90 100644 --- a/ggml/src/iqk/iqk_mul_mat.cpp +++ b/ggml/src/iqk/iqk_mul_mat.cpp @@ -1397,7 +1397,6 @@ void iqk_fused_delta_net_impl(int n_heads, int n_tokens, int n_seqs, static_assert(head_dim % 8 == 0); #endif - const float eps = 1e-6f; const float scale = 1.0f / sqrtf((float) head_dim); float v_new_buf[head_dim]; @@ -1428,42 +1427,25 @@ void iqk_fused_delta_net_impl(int n_heads, int n_tokens, int n_seqs, const float g_val = g_data[g_head_offset + t]; const float beta_raw = beta_data[g_head_offset + t]; - float q_norm_sq = 0.0f; - float k_norm_sq = 0.0f; float kq_sum = 0.0f; #ifdef __AVX2__ - auto vqsum = _mm256_setzero_ps(); - auto vksum = _mm256_setzero_ps(); auto vqksum = _mm256_setzero_ps(); for (int i = 0; i < head_dim; i += 8) { auto vq = _mm256_loadu_ps(q_t + i); auto vk = _mm256_loadu_ps(k_t + i); - vqsum = _mm256_fmadd_ps(vq, vq, vqsum); - vksum = _mm256_fmadd_ps(vk, vk, vksum); vqksum = _mm256_fmadd_ps(vk, vq, vqksum); } - q_norm_sq = hsum_float_8(vqsum); - k_norm_sq = hsum_float_8(vksum); kq_sum = hsum_float_8(vqksum); #else for (int i = 0; i < head_dim; ++i) { - q_norm_sq += q_t[i] * q_t[i]; - k_norm_sq += k_t[i] * k_t[i]; - kq_sum += k_t[i] * q_t[i]; + kq_sum += k_t[i] * q_t[i]; } #endif - const float q_norm_inv = 1.0f / sqrtf(q_norm_sq + eps); - const float k_norm_inv = 1.0f / sqrtf(k_norm_sq + eps); const float beta_val = 1.0f / (1.0f + expf(-beta_raw)); const float decay = expf(fminf(g_val, 50.0f)); - float attn_score = kq_sum * k_norm_inv * q_norm_inv * scale; - - //float attn_score = 0.0f; - //for (int i = 0; i < head_dim; ++i) { - // attn_score += (k_t[i] * k_norm_inv) * (q_t[i] * q_norm_inv * scale); - //} + float attn_score = kq_sum * scale; float * out_t = out_data + out_head_offset + t * out_token_stride; @@ -1479,9 +1461,9 @@ void iqk_fused_delta_net_impl(int n_heads, int n_tokens, int n_seqs, } } for (int row = 0; row < head_dim; ++row) { - const float v_new = v_t[row] * beta_val - v_prime[row] * beta_val * decay * k_norm_inv; + const float v_new = v_t[row] * beta_val - v_prime[row] * beta_val * decay; v_new_buf[row] = v_new; - out_t[row] = out_val[row] * decay * q_norm_inv * scale + v_new * attn_score; + out_t[row] = out_val[row] * decay * scale + v_new * attn_score; } #ifdef __AVX2__ @@ -1489,7 +1471,7 @@ void iqk_fused_delta_net_impl(int n_heads, int n_tokens, int n_seqs, auto vmin = _mm256_set1_ps(-1e6f); auto vmax = _mm256_set1_ps( 1e6f); for (int col = 0; col < head_dim; ++col) { - auto vk = _mm256_set1_ps(k_t[col] * k_norm_inv); + auto vk = _mm256_set1_ps(k_t[col]); for (int row = 0; row < head_dim; row += 8) { auto vs = _mm256_loadu_ps(state + col * head_dim + row); auto vn = _mm256_loadu_ps(v_new_buf + row); @@ -1503,7 +1485,7 @@ void iqk_fused_delta_net_impl(int n_heads, int n_tokens, int n_seqs, } #else for (int col = 0; col < head_dim; ++col) { - const float k_col = k_t[col] * k_norm_inv; + const float k_col = k_t[col]; for (int row = 0; row < head_dim; ++row) { float s = state[row + col * head_dim]; s = decay * s + v_new_buf[row] * k_col;