| 230 | } |
| 231 | |
| 232 | core::TensorValue same_self_attention( |
| 233 | core::ModuleBuildContext & ctx, |
| 234 | const core::TensorValue & input, |
| 235 | const core::TensorValue & positions, |
| 236 | const core::TensorValue * attention_mask, |
| 237 | const StableAudioSameAttentionWeights & weights, |
| 238 | const StableAudioConfig & config, |
| 239 | int64_t dim) { |
| 240 | const int64_t heads = dim / config.same_dim_heads; |
| 241 | auto qkv = modules::LinearModule({dim, dim * (config.same_differential ? 5 : 3), false, GGML_PREC_F32}) |
| 242 | .build(ctx, input, weights.to_qkv); |
| 243 | auto q = modules::SliceModule({2, 0, dim}).build(ctx, qkv); |
| 244 | auto k = modules::SliceModule({2, dim, dim}).build(ctx, qkv); |
| 245 | auto v = modules::SliceModule({2, 2 * dim, dim}).build(ctx, qkv); |
| 246 | core::TensorValue q_diff; |
| 247 | core::TensorValue k_diff; |
| 248 | if (config.same_differential) { |
| 249 | q_diff = modules::SliceModule({2, 3 * dim, dim}).build(ctx, qkv); |
| 250 | k_diff = modules::SliceModule({2, 4 * dim, dim}).build(ctx, qkv); |
| 251 | } |
| 252 | q = dynamic_tanh_norm(ctx, reshape_heads(ctx, q, heads, config.same_dim_heads), weights.q_norm); |
| 253 | k = dynamic_tanh_norm(ctx, reshape_heads(ctx, k, heads, config.same_dim_heads), weights.k_norm); |
| 254 | v = reshape_heads(ctx, v, heads, config.same_dim_heads); |
| 255 | q = modules::RoPEModule({config.same_dim_heads / 2, GGML_ROPE_TYPE_NEOX, 10000.0F}).build(ctx, q, positions); |
| 256 | k = modules::RoPEModule({config.same_dim_heads / 2, GGML_ROPE_TYPE_NEOX, 10000.0F}).build(ctx, k, positions); |
| 257 | auto q_heads = modules::TransposeModule({{0, 2, 1, 3}, q.shape.rank}).build(ctx, q); |
| 258 | auto k_heads = modules::TransposeModule({{0, 2, 1, 3}, k.shape.rank}).build(ctx, k); |
| 259 | auto v_heads = modules::TransposeModule({{0, 2, 1, 3}, v.shape.rank}).build(ctx, v); |
| 260 | auto context = same_attention(ctx, q_heads, k_heads, v_heads, attention_mask, config.same_dim_heads); |
| 261 | if (config.same_differential) { |
| 262 | q_diff = dynamic_tanh_norm(ctx, reshape_heads(ctx, q_diff, heads, config.same_dim_heads), weights.q_norm); |
| 263 | k_diff = dynamic_tanh_norm(ctx, reshape_heads(ctx, k_diff, heads, config.same_dim_heads), weights.k_norm); |
| 264 | q_diff = modules::RoPEModule({config.same_dim_heads / 2, GGML_ROPE_TYPE_NEOX, 10000.0F}).build(ctx, q_diff, positions); |
| 265 | k_diff = modules::RoPEModule({config.same_dim_heads / 2, GGML_ROPE_TYPE_NEOX, 10000.0F}).build(ctx, k_diff, positions); |
| 266 | auto qd_heads = modules::TransposeModule({{0, 2, 1, 3}, q_diff.shape.rank}).build(ctx, q_diff); |
| 267 | auto kd_heads = modules::TransposeModule({{0, 2, 1, 3}, k_diff.shape.rank}).build(ctx, k_diff); |
| 268 | auto diff_context = same_attention(ctx, qd_heads, kd_heads, v_heads, attention_mask, config.same_dim_heads); |
| 269 | context = modules::AddModule{}.build(ctx, context, scale_tensor(ctx, diff_context, -1.0F)); |
| 270 | } |
| 271 | context = core::reshape_tensor( |
| 272 | ctx, |
| 273 | ensure_contiguous(ctx, context), |
| 274 | core::TensorShape::from_dims({input.shape.dims[0], input.shape.dims[1], dim})); |
| 275 | return modules::LinearModule({dim, dim, false, GGML_PREC_F32}).build(ctx, context, weights.to_out); |
| 276 | } |
| 277 | |
| 278 | core::TensorValue same_feed_forward( |
| 279 | core::ModuleBuildContext & ctx, |
no test coverage detected