diff --git a/cmake/onnxruntime_mlas.cmake b/cmake/onnxruntime_mlas.cmake index 46f9ca2c38e95..a6dee4fcc0a66 100644 --- a/cmake/onnxruntime_mlas.cmake +++ b/cmake/onnxruntime_mlas.cmake @@ -59,6 +59,8 @@ onnxruntime_add_static_library(onnxruntime_mlas ${MLAS_SRC_DIR}/flashattn_qkv.cpp ${MLAS_SRC_DIR}/flashattn_gqa.cpp ${MLAS_SRC_DIR}/qkv_quant.cpp + ${MLAS_SRC_DIR}/linear_attention.h + ${MLAS_SRC_DIR}/linear_attention.cpp ${MLAS_SRC_DIR}/cast.cpp ${MLAS_SRC_DIR}/layernorm.cpp ${MLAS_SRC_DIR}/rotary_embedding.h @@ -245,6 +247,7 @@ function(setup_mlas_source_for_windows) ${MLAS_SRC_DIR}/intrinsics/avx512/silu_avx512f.cpp ${MLAS_SRC_DIR}/intrinsics/avx512/quantize_avx512f.cpp ${MLAS_SRC_DIR}/intrinsics/avx512/sconv_nchw_depthwise_multiplier_greater_than_1_avx512f.cpp + ${MLAS_SRC_DIR}/linear_attention_kernel_avx512f.cpp ) set_source_files_properties(${mlas_platform_srcs_avx512} PROPERTIES COMPILE_FLAGS "/arch:AVX512") @@ -928,6 +931,7 @@ else() ${MLAS_SRC_DIR}/intrinsics/avx512/silu_avx512f.cpp ${MLAS_SRC_DIR}/intrinsics/avx512/quantize_avx512f.cpp ${MLAS_SRC_DIR}/intrinsics/avx512/sconv_nchw_depthwise_multiplier_greater_than_1_avx512f.cpp + ${MLAS_SRC_DIR}/linear_attention_kernel_avx512f.cpp ) set_source_files_properties(${mlas_platform_srcs_avx512f} PROPERTIES COMPILE_FLAGS "-mavx512f") diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md index 96cfc7863f731..9d9e174d39452 100644 --- a/docs/ContribOperators.md +++ b/docs/ContribOperators.md @@ -2895,7 +2895,7 @@ This version of the operator has been available since version 1 of the 'com.micr
output : T
-
Attention output with 3D packed shape (B, T, H_q * d_v).
+
Attention output with 3D packed shape (B, T, max(H_q, H_kv) * d_v). Standard GQA emits one output per query head; inverse GQA, where H_kv exceeds H_q, emits one per KV head.
present_state : S
Updated recurrent state with shape (B, H_kv, d_k, d_v), or (W, B, H_kv, d_k, d_v) when state_window = W > 0. Slot W-1 is the state after the last token; slot j is the state after token (T - W + j).
diff --git a/onnxruntime/contrib_ops/cpu/bert/linear_attention.cc b/onnxruntime/contrib_ops/cpu/bert/linear_attention.cc index b08f7b17d4d15..a385eb1f95b96 100644 --- a/onnxruntime/contrib_ops/cpu/bert/linear_attention.cc +++ b/onnxruntime/contrib_ops/cpu/bert/linear_attention.cc @@ -4,13 +4,15 @@ #include "contrib_ops/cpu/bert/linear_attention.h" #include "contrib_ops/cpu/bert/linear_attention_helper.h" +#include "core/framework/allocator.h" #include "core/framework/tensorprotoutils.h" #include "core/common/safeint.h" #include "core/mlas/inc/mlas.h" #include "core/platform/threadpool.h" +#include #include -#include +#include using onnxruntime::concurrency::ThreadPool; @@ -23,9 +25,11 @@ namespace contrib { // MLFloat16 CPU support would require input/output conversion buffers // (MlasConvertHalfToFloatBuffer / MlasConvertFloatToHalfBuffer). // -// MLAS usage: MlasGemm is used for retrieval (S^T @ k), state update (k ⊗ delta), -// and query readout (S^T @ q) when d_k * d_v >= 4096. Smaller dimensions use -// scalar loops to avoid MLAS overhead. +// The recurrence itself lives in MLAS (MlasLinearAttention, see +// core/mlas/lib/linear_attention.cpp), which owns the threading, the state +// layout and the ISA dispatch. This kernel is plumbing: it validates inputs, +// derives the head geometry, allocates the output/state/scratch buffers, and +// hands MLAS a descriptor. #define REGISTER_KERNEL_TYPED(T) \ ONNX_OPERATOR_TYPED_KERNEL_EX( \ LinearAttention, \ @@ -56,6 +60,16 @@ LinearAttention::LinearAttention(const OpKernelInfo& info) : OpKernel(info) { update_rule_ == "delta" || update_rule_ == "gated_delta", "update_rule must be one of: linear, gated, delta, gated_delta"); + if (update_rule_ == "linear") { + rule_ = MlasLinearAttentionRuleLinear; + } else if (update_rule_ == "gated") { + rule_ = MlasLinearAttentionRuleGated; + } else if (update_rule_ == "delta") { + rule_ = MlasLinearAttentionRuleDelta; + } else { + rule_ = MlasLinearAttentionRuleGatedDelta; + } + scale_ = info.GetAttrOrDefault("scale", 0.0f); int64_t chunk_size = info.GetAttrOrDefault("chunk_size", 64); @@ -67,253 +81,6 @@ LinearAttention::LinearAttention(const OpKernelInfo& info) : OpKernel(info) { state_window_ = 0; } -namespace { - -// Process a single (batch, kv_head) pair across all timesteps. -// This is the hot inner loop — called once per (b, h_kv) combination, -// potentially from different threads. -void ProcessHead( - float* S, // State matrix [d_k, d_v], in-place updated - const float* q_data, // Packed Q: (B, T, H_q*d_k) - const float* k_data, // Packed K: (B, T, n_k*d_k) - const float* v_data, // Packed V: (B, T, H_kv*d_v) - const float* decay_data, // Decay gates (may be nullptr) - const float* beta_data, // Beta rates (may be nullptr) - float* output_data, // Output: (B, T, H_q*d_v) - float* retrieved_buf, // Pre-allocated scratch buffer [d_v] - int64_t batch_idx, - int h_kv, - int h_k, // Key head index (may differ from h_kv when n_k != kv_num_heads) - int64_t seq_len, - int64_t d_k, - int64_t d_v, - int q_num_heads, - int kv_num_heads, - int n_k_heads, - int heads_per_group, - int64_t output_hidden, - float scale, - bool needs_decay, - bool decay_per_key_dim, - bool needs_beta, - bool beta_per_head, - bool needs_retrieval) { - const size_t dk = static_cast(d_k); - const size_t dv = static_cast(d_v); - const bool use_mlas = (d_k * d_v >= 4096); - - for (int64_t t = 0; t < seq_len; ++t) { - // Pointers into packed 3D tensors at position [batch_idx, t, head*d] - const float* kt = k_data + (batch_idx * seq_len + t) * (n_k_heads * d_k) + h_k * d_k; - const float* vt = v_data + (batch_idx * seq_len + t) * (kv_num_heads * d_v) + h_kv * d_v; - - // ---- Step 1: Apply decay S *= exp(g_t) ---- - if (needs_decay) { - if (decay_per_key_dim) { - const float* gt = decay_data + (batch_idx * seq_len + t) * (kv_num_heads * d_k) + h_kv * d_k; - for (int64_t i = 0; i < d_k; ++i) { - float exp_g = std::exp(gt[i]); - // Scale row i of S by exp_g - for (int64_t j = 0; j < d_v; ++j) { - S[i * d_v + j] *= exp_g; - } - } - } else { - const float* gt = decay_data + (batch_idx * seq_len + t) * kv_num_heads + h_kv; - float exp_g = std::exp(gt[0]); - for (int64_t i = 0; i < d_k * d_v; ++i) { - S[i] *= exp_g; - } - } - } - - // ---- Step 2: Retrieval = S^T @ k_t ---- - if (needs_retrieval) { - if (use_mlas) { - MlasGemm( - CblasNoTrans, - CblasNoTrans, - 1, - dv, - dk, - 1.0f, - kt, - dk, - S, - dv, - 0.0f, - retrieved_buf, - dv, - nullptr, - nullptr); - } else { - for (int64_t j = 0; j < d_v; ++j) { - float acc = 0.0f; - for (int64_t i = 0; i < d_k; ++i) { - acc += S[i * d_v + j] * kt[i]; - } - retrieved_buf[static_cast(j)] = acc; - } - } - } - - // ---- Step 3: State update ---- - if (needs_beta) { - float bt; - if (beta_per_head) { - bt = beta_data[(batch_idx * seq_len + t) * kv_num_heads + h_kv]; - } else { - bt = beta_data[(batch_idx * seq_len + t) * 1]; - } - // Compute delta = beta * (v_t - retrieved) in-place into retrieved_buf - for (size_t j = 0; j < dv; ++j) { - retrieved_buf[j] = bt * (vt[j] - retrieved_buf[j]); - } - // S += k_t outer delta - if (use_mlas) { - MlasGemm( - CblasNoTrans, - CblasNoTrans, - dk, - dv, - 1, - 1.0f, - kt, - 1, - retrieved_buf, - dv, - 1.0f, - S, - dv, - nullptr, - nullptr); - } else { - for (int64_t i = 0; i < d_k; ++i) { - float* s_row = S + i * d_v; - const float ki = kt[i]; - for (int64_t j = 0; j < d_v; ++j) { - s_row[j] += ki * retrieved_buf[static_cast(j)]; - } - } - } - } else { - // linear/gated: S += k_t outer v_t - if (use_mlas) { - MlasGemm( - CblasNoTrans, - CblasNoTrans, - dk, - dv, - 1, - 1.0f, - kt, - 1, - vt, - dv, - 1.0f, - S, - dv, - nullptr, - nullptr); - } else { - for (int64_t i = 0; i < d_k; ++i) { - float* s_row = S + i * d_v; - const float ki = kt[i]; - for (int64_t j = 0; j < d_v; ++j) { - s_row[j] += ki * vt[j]; - } - } - } - } - - // ---- Step 4: Query readout for each q head in this kv group ---- - // o_t = scale * q_t^T @ S -> [1, d_v] - // - // Standard GQA (heads_per_group > 0): multiple Q heads share this KV state. - // Inverse GQA (heads_per_group == 0): multiple KV states share one Q head. - if (heads_per_group > 0) { - for (int g = 0; g < heads_per_group; ++g) { - int h_q = h_kv * heads_per_group + g; - const float* qt = q_data + (batch_idx * seq_len + t) * (q_num_heads * d_k) + h_q * d_k; - float* ot = output_data + (batch_idx * seq_len + t) * output_hidden + h_q * d_v; - - if (use_mlas) { - // Use alpha=1.0 to hit the MLAS M=1 gemv fast path, then scale output. - MlasGemm( - CblasNoTrans, - CblasNoTrans, - 1, - dv, - dk, - 1.0f, - qt, - dk, - S, - dv, - 0.0f, - ot, - dv, - nullptr, - nullptr); - if (scale != 1.0f) { - for (size_t j = 0; j < dv; ++j) { - ot[j] *= scale; - } - } - } else { - for (int64_t j = 0; j < d_v; ++j) { - float acc = 0.0f; - for (int64_t i = 0; i < d_k; ++i) { - acc += qt[i] * S[i * d_v + j]; - } - ot[j] = scale * acc; - } - } - } - } else { - // Inverse GQA: this KV head's Q is determined by h_kv * q_num / kv_num - int h_q = h_kv * q_num_heads / kv_num_heads; - const float* qt = q_data + (batch_idx * seq_len + t) * (q_num_heads * d_k) + h_q * d_k; - float* ot = output_data + (batch_idx * seq_len + t) * output_hidden + h_kv * d_v; - - if (use_mlas) { - // Use alpha=1.0 to hit the MLAS M=1 gemv fast path, then scale output. - MlasGemm( - CblasNoTrans, - CblasNoTrans, - 1, - dv, - dk, - 1.0f, - qt, - dk, - S, - dv, - 0.0f, - ot, - dv, - nullptr, - nullptr); - if (scale != 1.0f) { - for (size_t j = 0; j < dv; ++j) { - ot[j] *= scale; - } - } - } else { - for (int64_t j = 0; j < d_v; ++j) { - float acc = 0.0f; - for (int64_t i = 0; i < d_k; ++i) { - acc += qt[i] * S[i * d_v + j]; - } - ot[j] = scale * acc; - } - } - } - } -} - -} // anonymous namespace - template Status LinearAttention::Compute(OpKernelContext* context) const { // ==== Input Retrieval ==== @@ -376,9 +143,10 @@ Status LinearAttention::Compute(OpKernelContext* context) const { } // ==== Validate optional inputs based on update_rule ==== - bool needs_decay = (update_rule_ == "gated" || update_rule_ == "gated_delta"); - bool needs_beta = (update_rule_ == "delta" || update_rule_ == "gated_delta"); - bool needs_retrieval = (update_rule_ == "delta" || update_rule_ == "gated_delta"); + bool needs_decay = (rule_ == MlasLinearAttentionRuleGated || + rule_ == MlasLinearAttentionRuleGatedDelta); + bool needs_beta = (rule_ == MlasLinearAttentionRuleDelta || + rule_ == MlasLinearAttentionRuleGatedDelta); ORT_RETURN_IF_NOT(!needs_decay || decay_tensor != nullptr, "decay input is required for update_rule=", update_rule_); @@ -388,7 +156,7 @@ Status LinearAttention::Compute(OpKernelContext* context) const { const float* decay_data = decay_tensor ? decay_tensor->Data() : nullptr; const float* beta_data = beta_tensor ? beta_tensor->Data() : nullptr; - bool decay_per_key_dim = false; + MLAS_LINEAR_ATTENTION_DECAY_LAYOUT decay_layout = MlasLinearAttentionDecayNone; if (decay_tensor != nullptr) { const auto& decay_shape = decay_tensor->Shape(); ORT_RETURN_IF_NOT(decay_shape.NumDimensions() == 3, @@ -398,14 +166,15 @@ Status LinearAttention::Compute(OpKernelContext* context) const { "), got (", decay_shape[0], ", ", decay_shape[1], ")"); int64_t decay_last = decay_shape[2]; if (decay_last == kv_num_heads_ * d_k) { - decay_per_key_dim = true; + decay_layout = MlasLinearAttentionDecayPerKeyDim; } else { ORT_RETURN_IF_NOT(decay_last == kv_num_heads_, "decay last dim must be H_kv or H_kv*d_k"); + decay_layout = MlasLinearAttentionDecayPerHead; } } - bool beta_per_head = false; + MLAS_LINEAR_ATTENTION_BETA_LAYOUT beta_layout = MlasLinearAttentionBetaNone; if (beta_tensor != nullptr) { const auto& beta_shape = beta_tensor->Shape(); ORT_RETURN_IF_NOT(beta_shape.NumDimensions() == 3, @@ -415,14 +184,28 @@ Status LinearAttention::Compute(OpKernelContext* context) const { "), got (", beta_shape[0], ", ", beta_shape[1], ")"); int64_t beta_last = beta_shape[2]; if (beta_last == kv_num_heads_) { - beta_per_head = true; + beta_layout = MlasLinearAttentionBetaPerHead; } else { ORT_RETURN_IF_NOT(beta_last == 1, "beta last dim must be H_kv or 1"); + beta_layout = MlasLinearAttentionBetaShared; } } + // MLAS only consumes decay/beta that the rule actually uses; a tensor + // supplied for a rule that ignores it stays ignored. + if (!needs_decay) { + decay_data = nullptr; + decay_layout = MlasLinearAttentionDecayNone; + } + if (!needs_beta) { + beta_data = nullptr; + beta_layout = MlasLinearAttentionBetaNone; + } + // ==== Initialize state: write directly into output present_state ==== - // state_window_ is always 0 on CPU, so this is the legacy (B, H_kv, d_k, d_v) shape. + // state_window_ is always 0 on CPU, so this is the legacy (B, H_kv, d_k, d_v) + // shape. MLAS updates this buffer in place and never initializes it, so the + // seeding below is mandatory. TensorShape state_shape; ORT_RETURN_IF_ERROR(linear_attention_helper::CheckInputs( state_window_, static_cast(batch_size), kv_num_heads_, @@ -434,7 +217,9 @@ Status LinearAttention::Compute(OpKernelContext* context) const { if (past_state_tensor != nullptr) { const float* ps_data = past_state_tensor->Data(); - std::memcpy(state_data, ps_data, static_cast(total_state) * sizeof(float)); + if (ps_data != state_data) { + std::memcpy(state_data, ps_data, static_cast(total_state) * sizeof(float)); + } } else { std::memset(state_data, 0, static_cast(total_state) * sizeof(float)); } @@ -442,66 +227,68 @@ Status LinearAttention::Compute(OpKernelContext* context) const { // ==== Allocate output ==== // Output hidden dim: max(q_num_heads, kv_num_heads) * d_v // Standard GQA: q_num_heads * d_v; Inverse GQA: kv_num_heads * d_v - int64_t output_hidden = std::max(q_num_heads_, kv_num_heads_) * d_v; + int64_t output_hidden = static_cast( + MlasLinearAttentionOutputHiddenSize(q_num_heads_, kv_num_heads_, static_cast(d_v))); TensorShape output_shape({batch_size, seq_len, output_hidden}); Tensor* output_tensor = context->Output(0, output_shape); float* output_data = output_tensor->MutableData(); - // ==== GQA head mapping ==== + // ==== Validate the GQA head mapping ==== // Standard GQA: q_num_heads >= kv_num_heads, multiple Q heads per KV group. // Inverse GQA: q_num_heads < kv_num_heads (e.g., Qwen3.5 9B: n_k=16, n_kv=32). // Also n_k_heads may differ from both (K has its own head count). - int heads_per_group; // Q heads per KV group (0 if inverse GQA) + // MLAS derives the actual mappings; the checks here exist to report a Status. if (q_num_heads_ >= kv_num_heads_) { ORT_RETURN_IF_NOT(q_num_heads_ % kv_num_heads_ == 0, "q_num_heads must be divisible by kv_num_heads"); - heads_per_group = q_num_heads_ / kv_num_heads_; } else { ORT_RETURN_IF_NOT(kv_num_heads_ % q_num_heads_ == 0, "kv_num_heads must be divisible by q_num_heads (inverse GQA)"); - heads_per_group = 0; // signals inverse GQA to ProcessHead } // K-to-KV head mapping: when n_k < kv_num_heads, multiple KV heads share one K head ORT_RETURN_IF_NOT(kv_num_heads_ % n_k_heads == 0, "kv_num_heads must be divisible by n_k_heads"); - int kv_per_k_head = kv_num_heads_ / n_k_heads; - - // ==== Thread-parallel over (batch, kv_head) pairs ==== - // Each (b, h_kv) pair is fully independent — the state matrix for each - // head is disjoint, and the sequential token dependency is within a - // single head only. This gives us batch_size * kv_num_heads parallel tasks. - int64_t total_tasks = batch_size * kv_num_heads_; - - // Cost estimate: per task processes seq_len tokens, each doing ~3*d_k*d_v FLOPs - double cost_per_task = static_cast(seq_len) * static_cast(d_k * d_v) * 3.0; + // ==== Hand off to MLAS ==== + // Work is partitioned over (batch, kv_head) pairs: each pair owns a disjoint + // state matrix, and the sequential token dependency lives entirely within a + // single pair. auto* tp = context->GetOperatorThreadPool(); - ThreadPool::TryParallelFor( - tp, - static_cast(total_tasks), - cost_per_task, - [&](std::ptrdiff_t first, std::ptrdiff_t last) { - // Pre-allocate scratch buffer per thread-batch to avoid malloc in hot loop - std::vector retrieved_buf(static_cast(d_v)); - - for (std::ptrdiff_t task = first; task < last; ++task) { - int64_t b = task / kv_num_heads_; - int h_kv = static_cast(task % kv_num_heads_); - int h_k = h_kv / kv_per_k_head; // map KV head to K head - - float* S = state_data + (b * kv_num_heads_ + h_kv) * state_per_head; - - ProcessHead( - S, q_data, k_data, v_data, decay_data, beta_data, output_data, - retrieved_buf.data(), - b, h_kv, h_k, seq_len, d_k, d_v, - q_num_heads_, kv_num_heads_, n_k_heads, heads_per_group, output_hidden, - s, needs_decay, decay_per_key_dim, needs_beta, beta_per_head, - needs_retrieval); - } - }); + MlasLinearAttentionArgs args; + args.batch_size = static_cast(batch_size); + args.sequence_length = static_cast(seq_len); + args.q_num_heads = q_num_heads_; + args.kv_num_heads = kv_num_heads_; + args.k_num_heads = n_k_heads; + args.k_head_size = static_cast(d_k); + args.v_head_size = static_cast(d_v); + args.rule = rule_; + args.decay_layout = decay_layout; + args.beta_layout = beta_layout; + args.scale = s; + + const int thread_count = std::max(1, ThreadPool::DegreeOfParallelism(tp)); + args.thread_count = thread_count; + args.buffer_size_per_thread = + MlasLinearAttentionBufferSizePerThread(args.k_head_size, args.v_head_size); + + AllocatorPtr allocator; + ORT_RETURN_IF_ERROR(context->GetTempSpaceAllocator(&allocator)); + const size_t buffer_bytes = SafeInt(args.buffer_size_per_thread) * thread_count; + IAllocatorUniquePtr buffer = IAllocator::MakeUniquePtr(allocator, buffer_bytes); + args.buffer = reinterpret_cast(buffer.get()); + + args.query = q_data; + args.key = k_data; + args.value = v_data; + args.decay = decay_data; + args.beta = beta_data; + args.state = state_data; + args.output = output_data; + + MlasLinearAttention(&args, tp); return Status::OK(); } diff --git a/onnxruntime/contrib_ops/cpu/bert/linear_attention.h b/onnxruntime/contrib_ops/cpu/bert/linear_attention.h index 1fa57028bb95b..c956a8d8094fd 100644 --- a/onnxruntime/contrib_ops/cpu/bert/linear_attention.h +++ b/onnxruntime/contrib_ops/cpu/bert/linear_attention.h @@ -5,6 +5,7 @@ #include "core/common/common.h" #include "core/framework/op_kernel.h" +#include "core/mlas/inc/mlas.h" #include @@ -21,6 +22,7 @@ class LinearAttention final : public OpKernel { int q_num_heads_; int kv_num_heads_; std::string update_rule_; + MLAS_LINEAR_ATTENTION_RULE rule_; // update_rule_, resolved once in the constructor float scale_; int chunk_size_; // Always 0 on CPU (a state window is CUDA-only), but kept so the shared shape helper in diff --git a/onnxruntime/core/graph/contrib_ops/bert_defs.cc b/onnxruntime/core/graph/contrib_ops/bert_defs.cc index cb67692929414..7f2de758fa9d3 100644 --- a/onnxruntime/core/graph/contrib_ops/bert_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/bert_defs.cc @@ -2658,7 +2658,11 @@ ONNX_MS_OPERATOR_SET_SCHEMA( OpSchema::Optional) .Output(0, "output", - "Attention output with 3D packed shape (B, T, H_q * d_v).", + // Kept free of angle brackets: gen_contrib_doc.py interpolates this + // description straight into an HTML
element without escaping. + "Attention output with 3D packed shape (B, T, max(H_q, H_kv) * d_v). " + "Standard GQA emits one output per query head; inverse GQA, where " + "H_kv exceeds H_q, emits one per KV head.", "T") .Output(1, "present_state", @@ -2682,7 +2686,7 @@ ONNX_MS_OPERATOR_SET_SCHEMA( int64_t q_num_heads = (q_num_heads_attr && q_num_heads_attr->has_i()) ? q_num_heads_attr->i() : 0; int64_t kv_num_heads = (kv_num_heads_attr && kv_num_heads_attr->has_i()) ? kv_num_heads_attr->i() : 0; - // Output 0: (B, T, H_q * d_v) — 3D packed + // Output 0: (B, T, max(H_q, H_kv) * d_v) — 3D packed if (hasInputShape(ctx, 0) && hasInputShape(ctx, 2) && q_num_heads > 0 && kv_num_heads > 0) { auto& query_shape = getInputShape(ctx, 0); auto& value_shape = getInputShape(ctx, 2); diff --git a/onnxruntime/core/mlas/inc/mlas.h b/onnxruntime/core/mlas/inc/mlas.h index 0f9cbd10aa3cd..501f2f6a2ff71 100644 --- a/onnxruntime/core/mlas/inc/mlas.h +++ b/onnxruntime/core/mlas/inc/mlas.h @@ -2482,6 +2482,154 @@ MlasFlashAttentionGQA( MLAS_THREADPOOL* ThreadPool ); +// +// Linear (recurrent) attention. +// +// Computes, independently for every (batch, kv head) pair, a token-sequential +// recurrence over a state matrix S of shape [k_head_size, v_head_size]: +// +// for t in [0, sequence_length): +// (1) decay : S *= exp(g_t) (gated / gated_delta only) +// (2) retrieval : r = S^T k_t (delta / gated_delta only) +// (3) update : S += k_t (x) d_t d_t = v_t (linear / gated) +// d_t = beta_t * (v_t - r) (delta / gated_delta) +// (4) readout : o_{t,h} = scale * q_{t,h}^T S for every query head h +// mapped to this kv head +// +// All tensors are FP32, packed row-major with no padding; strides are implied +// by the head counts and head sizes documented on each field. +// +// Head mapping: +// Key heads - kv_num_heads % k_num_heads == 0; consecutive kv heads share +// a key head, h_k = h_kv / (kv_num_heads / k_num_heads). +// Standard GQA - q_num_heads >= kv_num_heads and divisible by it; each kv +// head serves q_num_heads / kv_num_heads consecutive query +// heads and writes the matching output heads. +// Inverse GQA - q_num_heads < kv_num_heads and divides it; each kv head +// reads the single shared query head +// h_q = h_kv * q_num_heads / kv_num_heads and writes output +// head h_kv. +// + +typedef enum { + MlasLinearAttentionRuleLinear = 0, // no decay, no beta + MlasLinearAttentionRuleGated = 1, // decay, no beta + MlasLinearAttentionRuleDelta = 2, // no decay, beta + retrieval + MlasLinearAttentionRuleGatedDelta = 3, // decay, beta + retrieval +} MLAS_LINEAR_ATTENTION_RULE; + +typedef enum { + MlasLinearAttentionDecayNone = 0, // rule has no decay term + MlasLinearAttentionDecayPerHead = 1, // decay is [B, T, kv_num_heads] + MlasLinearAttentionDecayPerKeyDim = 2, // decay is [B, T, kv_num_heads * k_head_size] +} MLAS_LINEAR_ATTENTION_DECAY_LAYOUT; + +typedef enum { + MlasLinearAttentionBetaNone = 0, // rule has no beta term + MlasLinearAttentionBetaPerHead = 1, // beta is [B, T, kv_num_heads] + MlasLinearAttentionBetaShared = 2, // beta is [B, T, 1] (broadcast over heads) +} MLAS_LINEAR_ATTENTION_BETA_LAYOUT; + +struct MlasLinearAttentionArgs { + // + // Problem shape. + // + // + // Preconditions the caller is responsible for, in the manner of the other + // MLAS attention entry points: these are assumed, not validated. The CPU EP + // enforces them and reports a Status (see contrib_ops/cpu/bert/linear_attention.cc). + // + // k_num_heads >= 1, and kv_num_heads % k_num_heads == 0 + // q_num_heads >= 1, kv_num_heads >= 1, and whichever of q_num_heads / + // kv_num_heads is larger must be a multiple of the smaller + // k_head_size >= 1, v_head_size >= 1 + // + int batch_size; // B + int sequence_length; // T + int q_num_heads; // H_q + int kv_num_heads; // H_kv - one independent recurrent state per (batch, kv head) + int k_num_heads; // H_k - key heads; H_kv % H_k == 0 + int k_head_size; // d_k + int v_head_size; // d_v + + // + // Behaviour. + // + MLAS_LINEAR_ATTENTION_RULE rule; + + // Must be *None if and only if the rule carries no decay / beta term. + MLAS_LINEAR_ATTENTION_DECAY_LAYOUT decay_layout; + MLAS_LINEAR_ATTENTION_BETA_LAYOUT beta_layout; + + // Query readout scale. The caller resolves any default (e.g. 1/sqrt(d_k)). + float scale; + + // + // Threading and scratch. + // + int thread_count; // number of partitions; >= 1 + float* buffer; // thread_count * buffer_size_per_thread bytes + size_t buffer_size_per_thread; // BYTES; see MlasLinearAttentionBufferSizePerThread + + // + // Tensors. All FP32, packed row-major, no padding. + // + const float* query; // [B, T, H_q * d_k] + const float* key; // [B, T, H_k * d_k] + const float* value; // [B, T, H_kv * d_v] + // Must be non-null whenever the rule carries the corresponding term, i.e. + // decay for Gated / GatedDelta and beta for Delta / GatedDelta. + const float* decay; // nullptr, or [B, T, H_kv] / [B, T, H_kv * d_k] per decay_layout + const float* beta; // nullptr, or [B, T, H_kv] / [B, T, 1] per beta_layout + + // IN/OUT. [B, H_kv, d_k, d_v] row-major; the leading dimension of each + // d_k x d_v slice is d_v. MLAS does NOT initialize this buffer - the caller + // must zero it or copy the past state in beforehand. On return it holds the + // post-sequence state. + float* state; + + // OUT. [B, T, MlasLinearAttentionOutputHiddenSize(...)]. Fully overwritten + // when sequence_length > 0; no accumulation is performed. + float* output; +}; + +/** + * @brief Hidden size of the LinearAttention output tensor. + * + * Standard GQA (q_num_heads >= kv_num_heads): q_num_heads * v_head_size. + * Inverse GQA (q_num_heads < kv_num_heads): kv_num_heads * v_head_size. + */ +size_t +MLASCALL +MlasLinearAttentionOutputHiddenSize( + int q_num_heads, + int kv_num_heads, + int v_head_size +); + +/** + * @brief Per-thread scratch size, in bytes, for MlasLinearAttention. Multiply + * by MlasLinearAttentionArgs::thread_count to size the buffer. + */ +size_t +MLASCALL +MlasLinearAttentionBufferSizePerThread( + int k_head_size, + int v_head_size +); + +/** + * @brief FP32 linear / gated / delta / gated-delta recurrent attention. + * @param args Arguments; args->state is updated in place. + * @param ThreadPool Thread pool + */ +void +MLASCALL +MlasLinearAttention( + MlasLinearAttentionArgs* args, + MLAS_THREADPOOL* ThreadPool +); + /** * @brief Enumeration of supported GELU algorithm variants. * diff --git a/onnxruntime/core/mlas/lib/linear_attention.cpp b/onnxruntime/core/mlas/lib/linear_attention.cpp new file mode 100644 index 0000000000000..37d378c15ea48 --- /dev/null +++ b/onnxruntime/core/mlas/lib/linear_attention.cpp @@ -0,0 +1,419 @@ +/*++ + +Copyright (c) Microsoft Corporation. All rights reserved. + +Licensed under the MIT License. + +Module Name: + + linear_attention.cpp + +Abstract: + + This module implements the FP32 linear (recurrent) attention kernels used + by the LinearAttention contrib operator. + + Unlike softmax attention, the recurrence is strictly sequential in the + token dimension, so parallelism comes entirely from the (batch, kv head) + dimensions - each pair owns a private state matrix and is fully + independent. + +--*/ + +#include "linear_attention.h" + +#include +#include +#include + +// +// Below this state size the SGEMM call overhead (argument marshalling, packing +// decisions, kernel selection) dominates the ~2 * d_k * d_v flops of the M=1 / +// K=1 shapes involved, so scalar loops win. +// +constexpr size_t MlasLinearAttentionSgemmThreshold = 4096; + +size_t +MLASCALL +MlasLinearAttentionOutputHiddenSize( + int q_num_heads, + int kv_num_heads, + int v_head_size +) +{ + return static_cast(std::max(q_num_heads, kv_num_heads)) * + static_cast(v_head_size); +} + +size_t +MLASCALL +MlasLinearAttentionBufferSizePerThread( + int k_head_size, + int v_head_size +) +{ + // + // v_head_size floats hold the retrieval / delta vector. k_head_size floats + // are reserved so a fused kernel can materialize exp(g_t) without growing + // the contract. Round up so each thread's slice starts on the preferred + // buffer alignment, which is queried rather than assumed to be 64. + // + const size_t Bytes = (static_cast(k_head_size) + + static_cast(v_head_size)) * sizeof(float); + + const size_t BufferAlignment = MlasGetPreferredBufferAlignment(); + + return (Bytes + BufferAlignment - 1) & ~(BufferAlignment - 1); +} + +void +MlasLinearAttentionProcessHead( + const MLAS_LINEAR_ATTENTION_WORK* Work +) +{ + const size_t dk = Work->KHeadSize; + const size_t dv = Work->VHeadSize; + const size_t heads_per_group = Work->HeadsPerGroup; + const float scale = Work->Scale; + + const bool needs_decay = (Work->Rule == MlasLinearAttentionRuleGated || + Work->Rule == MlasLinearAttentionRuleGatedDelta); + const bool needs_beta = (Work->Rule == MlasLinearAttentionRuleDelta || + Work->Rule == MlasLinearAttentionRuleGatedDelta); + const bool needs_retrieval = needs_beta; + const bool decay_per_key_dim = (Work->DecayLayout == MlasLinearAttentionDecayPerKeyDim); + + const bool use_sgemm = (dk * dv >= MlasLinearAttentionSgemmThreshold); + + float* S = Work->State; + float* retrieved_buf = Work->Scratch; + + for (size_t t = 0; t < Work->SequenceLength; ++t) { + const float* kt = Work->Key + t * Work->KeyTokenStride; + const float* vt = Work->Value + t * Work->ValueTokenStride; + + // ---- Step 1: Apply decay S *= exp(g_t) ---- + if (needs_decay) { + const float* gt = Work->Decay + t * Work->DecayTokenStride; + if (decay_per_key_dim) { + for (size_t i = 0; i < dk; ++i) { + float exp_g = std::exp(gt[i]); + // Scale row i of S by exp_g + for (size_t j = 0; j < dv; ++j) { + S[i * dv + j] *= exp_g; + } + } + } else { + float exp_g = std::exp(gt[0]); + for (size_t i = 0; i < dk * dv; ++i) { + S[i] *= exp_g; + } + } + } + + // ---- Step 2: Retrieval = S^T @ k_t ---- + if (needs_retrieval) { + if (use_sgemm) { + MlasSgemmOperation(CblasNoTrans, CblasNoTrans, 1, dv, dk, + 1.0f, kt, dk, S, dv, 0.0f, retrieved_buf, dv); + } else { + for (size_t j = 0; j < dv; ++j) { + float acc = 0.0f; + for (size_t i = 0; i < dk; ++i) { + acc += S[i * dv + j] * kt[i]; + } + retrieved_buf[j] = acc; + } + } + } + + // ---- Step 3: State update ---- + if (needs_beta) { + const float bt = Work->Beta[t * Work->BetaTokenStride]; + // Compute delta = beta * (v_t - retrieved) in-place into retrieved_buf + for (size_t j = 0; j < dv; ++j) { + retrieved_buf[j] = bt * (vt[j] - retrieved_buf[j]); + } + // S += k_t outer delta + if (use_sgemm) { + MlasSgemmOperation(CblasNoTrans, CblasNoTrans, dk, dv, 1, + 1.0f, kt, 1, retrieved_buf, dv, 1.0f, S, dv); + } else { + for (size_t i = 0; i < dk; ++i) { + float* s_row = S + i * dv; + const float ki = kt[i]; + for (size_t j = 0; j < dv; ++j) { + s_row[j] += ki * retrieved_buf[j]; + } + } + } + } else { + // linear/gated: S += k_t outer v_t + if (use_sgemm) { + MlasSgemmOperation(CblasNoTrans, CblasNoTrans, dk, dv, 1, + 1.0f, kt, 1, vt, dv, 1.0f, S, dv); + } else { + for (size_t i = 0; i < dk; ++i) { + float* s_row = S + i * dv; + const float ki = kt[i]; + for (size_t j = 0; j < dv; ++j) { + s_row[j] += ki * vt[j]; + } + } + } + } + + // ---- Step 4: Query readout, o_t = scale * q_t^T @ S -> [1, d_v] ---- + // + // Standard GQA: HeadsPerGroup consecutive query heads share this state. + // Inverse GQA: HeadsPerGroup is 1 and Query already points at the + // shared query head for this kv head. + // + const float* qt_base = Work->Query + t * Work->QueryTokenStride; + float* ot_base = Work->Output + t * Work->OutputTokenStride; + + for (size_t g = 0; g < heads_per_group; ++g) { + const float* qt = qt_base + g * dk; + float* ot = ot_base + g * dv; + + if (use_sgemm) { + // Use alpha=1.0 to hit the MLAS M=1 gemv fast path, then scale output. + MlasSgemmOperation(CblasNoTrans, CblasNoTrans, 1, dv, dk, + 1.0f, qt, dk, S, dv, 0.0f, ot, dv); + if (scale != 1.0f) { + for (size_t j = 0; j < dv; ++j) { + ot[j] *= scale; + } + } + } else { + for (size_t j = 0; j < dv; ++j) { + float acc = 0.0f; + for (size_t i = 0; i < dk; ++i) { + acc += qt[i] * S[i * dv + j]; + } + ot[j] = scale * acc; + } + } + } + } +} + +const MLAS_LINEAR_ATTENTION_DISPATCH MlasLinearAttentionDispatchDefault = { + MlasLinearAttentionProcessHead +}; + +// +// Number of independent work items: one per (batch, kv head) pair. +// +static +std::ptrdiff_t +MlasLinearAttentionTaskCount( + const MlasLinearAttentionArgs* args +) +{ + return static_cast(args->batch_size) * + static_cast(args->kv_num_heads); +} + +static +std::ptrdiff_t +MlasLinearAttentionPartitionCount( + const MlasLinearAttentionArgs* args +) +{ + const std::ptrdiff_t total_task_count = MlasLinearAttentionTaskCount(args); + + std::ptrdiff_t thread_count = static_cast(args->thread_count); + if (thread_count > total_task_count) { + thread_count = total_task_count; + } + if (thread_count < 1) { + thread_count = 1; + } + + return thread_count; +} + +void +MlasLinearAttentionThreaded( + void* argptr, + std::ptrdiff_t thread_id +) +{ + const MlasLinearAttentionArgs* args = reinterpret_cast(argptr); + + const std::ptrdiff_t total_task_count = MlasLinearAttentionTaskCount(args); + const std::ptrdiff_t thread_count = MlasLinearAttentionPartitionCount(args); + + if (thread_id >= thread_count) { + return; + } + + std::ptrdiff_t quotient = total_task_count / thread_count; + std::ptrdiff_t remainder = total_task_count % thread_count; + std::ptrdiff_t task_start; + std::ptrdiff_t task_end; + if (thread_id < remainder) { + task_start = (quotient + 1) * thread_id; + task_end = task_start + quotient + 1; + } else { + task_start = quotient * thread_id + remainder; + task_end = task_start + quotient; + } + + // + // Loop invariants. All offsets are computed in size_t: a packed query + // tensor at B=8, T=8192, H_q=32, d_k=128 already exceeds INT32_MAX. + // + const std::ptrdiff_t kv_num_heads = static_cast(args->kv_num_heads); + const size_t dk = static_cast(args->k_head_size); + const size_t dv = static_cast(args->v_head_size); + const size_t seq_len = static_cast(args->sequence_length); + const size_t state_per_head = dk * dv; + + const size_t q_token_stride = static_cast(args->q_num_heads) * dk; + const size_t k_token_stride = static_cast(args->k_num_heads) * dk; + const size_t v_token_stride = static_cast(args->kv_num_heads) * dv; + const size_t o_token_stride = MlasLinearAttentionOutputHiddenSize( + args->q_num_heads, args->kv_num_heads, args->v_head_size); + + const std::ptrdiff_t kv_per_k_head = kv_num_heads / static_cast(args->k_num_heads); + + // + // Standard GQA (H_q >= H_kv) gives each kv head H_q / H_kv consecutive + // query heads; inverse GQA (H_q < H_kv) gives it a single shared one. Both + // reduce to "HeadsPerGroup consecutive query heads from h_q0, written to + // HeadsPerGroup consecutive output heads from h_out0". + // + const bool inverse_gqa = (args->q_num_heads < args->kv_num_heads); + const std::ptrdiff_t heads_per_group = + inverse_gqa ? 1 : (static_cast(args->q_num_heads) / kv_num_heads); + + size_t decay_token_stride = 0; + if (args->decay_layout == MlasLinearAttentionDecayPerKeyDim) { + decay_token_stride = static_cast(args->kv_num_heads) * dk; + } else if (args->decay_layout == MlasLinearAttentionDecayPerHead) { + decay_token_stride = static_cast(args->kv_num_heads); + } + + size_t beta_token_stride = 0; + if (args->beta_layout == MlasLinearAttentionBetaPerHead) { + beta_token_stride = static_cast(args->kv_num_heads); + } else if (args->beta_layout == MlasLinearAttentionBetaShared) { + beta_token_stride = 1; + } + + float* scratch = reinterpret_cast( + reinterpret_cast(args->buffer) + thread_id * args->buffer_size_per_thread); + + const MLAS_LINEAR_ATTENTION_DISPATCH* dispatch = GetMlasPlatform().LinearAttentionDispatch; + + for (std::ptrdiff_t task_index = task_start; task_index < task_end; ++task_index) { + const std::ptrdiff_t b = task_index / kv_num_heads; + const std::ptrdiff_t h_kv = task_index % kv_num_heads; + const std::ptrdiff_t h_k = h_kv / kv_per_k_head; + + const std::ptrdiff_t h_q0 = + inverse_gqa ? (h_kv * static_cast(args->q_num_heads) / kv_num_heads) + : (h_kv * heads_per_group); + const std::ptrdiff_t h_out0 = inverse_gqa ? h_kv : (h_kv * heads_per_group); + + // Token 0 of this batch item. + const size_t seq_base = static_cast(b) * seq_len; + + MLAS_LINEAR_ATTENTION_WORK work; + + work.State = args->state + + (static_cast(b) * static_cast(args->kv_num_heads) + + static_cast(h_kv)) * state_per_head; + work.Query = args->query + seq_base * q_token_stride + static_cast(h_q0) * dk; + work.Key = args->key + seq_base * k_token_stride + static_cast(h_k) * dk; + work.Value = args->value + seq_base * v_token_stride + static_cast(h_kv) * dv; + work.Output = args->output + seq_base * o_token_stride + static_cast(h_out0) * dv; + + if (args->decay != nullptr && args->decay_layout != MlasLinearAttentionDecayNone) { + const size_t head_offset = + (args->decay_layout == MlasLinearAttentionDecayPerKeyDim) + ? static_cast(h_kv) * dk + : static_cast(h_kv); + work.Decay = args->decay + seq_base * decay_token_stride + head_offset; + } else { + work.Decay = nullptr; + } + + if (args->beta != nullptr && args->beta_layout != MlasLinearAttentionBetaNone) { + const size_t head_offset = (args->beta_layout == MlasLinearAttentionBetaPerHead) + ? static_cast(h_kv) + : 0; + work.Beta = args->beta + seq_base * beta_token_stride + head_offset; + } else { + work.Beta = nullptr; + } + + work.QueryTokenStride = q_token_stride; + work.KeyTokenStride = k_token_stride; + work.ValueTokenStride = v_token_stride; + work.DecayTokenStride = (work.Decay != nullptr) ? decay_token_stride : 0; + work.BetaTokenStride = (work.Beta != nullptr) ? beta_token_stride : 0; + work.OutputTokenStride = o_token_stride; + + work.SequenceLength = seq_len; + work.KHeadSize = dk; + work.VHeadSize = dv; + work.HeadsPerGroup = static_cast(heads_per_group); + work.Scale = args->scale; + work.Rule = args->rule; + work.DecayLayout = args->decay_layout; + work.Scratch = scratch; + + if (dispatch != nullptr && dispatch->ProcessHead != nullptr) { + dispatch->ProcessHead(&work); + } else { + MlasLinearAttentionProcessHead(&work); + } + } +} + +void +MLASCALL +MlasLinearAttention( + MlasLinearAttentionArgs* args, + MLAS_THREADPOOL* ThreadPool +) +{ + // + // Preconditions, documented on MlasLinearAttentionArgs. Like the other MLAS + // attention entry points this trusts its caller in release builds - the CPU + // EP validates and reports a Status - but the invariants are asserted here + // so a mis-filled args struct fails loudly in a debug build rather than + // dividing by zero or dereferencing null deep inside a kernel. + // + assert(args->k_num_heads > 0); + assert(args->q_num_heads > 0 && args->kv_num_heads > 0); + assert(args->kv_num_heads % args->k_num_heads == 0); + assert(args->q_num_heads >= args->kv_num_heads + ? args->q_num_heads % args->kv_num_heads == 0 + : args->kv_num_heads % args->q_num_heads == 0); + assert(args->k_head_size > 0 && args->v_head_size > 0); + assert((args->rule != MlasLinearAttentionRuleGated && + args->rule != MlasLinearAttentionRuleGatedDelta) || + (args->decay != nullptr && args->decay_layout != MlasLinearAttentionDecayNone)); + assert((args->rule != MlasLinearAttentionRuleDelta && + args->rule != MlasLinearAttentionRuleGatedDelta) || + (args->beta != nullptr && args->beta_layout != MlasLinearAttentionBetaNone)); + + if (args->sequence_length <= 0 || MlasLinearAttentionTaskCount(args) <= 0) { + return; + } + + // + // Clamping the partition count to the task count keeps threads from being + // scheduled with nothing to do, and lets the common B=1, H_kv=1 decode + // case take the inline path in MlasExecuteThreaded. + // + MlasExecuteThreaded( + MlasLinearAttentionThreaded, + static_cast(args), + MlasLinearAttentionPartitionCount(args), + ThreadPool); +} diff --git a/onnxruntime/core/mlas/lib/linear_attention.h b/onnxruntime/core/mlas/lib/linear_attention.h new file mode 100644 index 0000000000000..aeb92cdc70a0b --- /dev/null +++ b/onnxruntime/core/mlas/lib/linear_attention.h @@ -0,0 +1,105 @@ +/*++ + +Copyright (c) Microsoft Corporation. All rights reserved. + +Licensed under the MIT License. + +Module Name: + + linear_attention.h + +Abstract: + + This module includes the work-item description and kernel dispatch + structure for the linear (recurrent) attention kernels. + +--*/ + +#pragma once + +#include "mlasi.h" + +// +// One unit of linear-attention work: the complete token loop for a single +// (batch, kv head) pair, i.e. a single state matrix S. +// +// The threaded routine resolves every base pointer and stride before invoking +// a kernel, so a kernel performs no tensor index arithmetic and knows nothing +// about batching, GQA head mapping or key-head sharing. +// +struct MLAS_LINEAR_ATTENTION_WORK { + // + // State matrix S, [KHeadSize, VHeadSize] row-major with leading dimension + // VHeadSize. Read and updated in place; never initialized by the kernel. + // + float* State; + + // + // Token 0 of this (batch, head). Advance by the matching token stride to + // reach token t. Query points at the first of HeadsPerGroup consecutive + // query heads, Output at the first of HeadsPerGroup consecutive output + // heads. + // + const float* Query; + const float* Key; + const float* Value; + const float* Decay; // nullptr when the rule carries no decay term + const float* Beta; // nullptr when the rule carries no beta term + float* Output; + + size_t QueryTokenStride; + size_t KeyTokenStride; + size_t ValueTokenStride; + size_t DecayTokenStride; // 0 when Decay is nullptr + size_t BetaTokenStride; // 0 when Beta is nullptr + size_t OutputTokenStride; + + size_t SequenceLength; + size_t KHeadSize; // d_k + size_t VHeadSize; // d_v + size_t HeadsPerGroup; // query heads read / output heads written per token + + float Scale; + + MLAS_LINEAR_ATTENTION_RULE Rule; + + // + // Distinguishes a scalar broadcast (PerHead) from a per-row vector + // (PerKeyDim). The beta layout needs no equivalent field: the base pointer + // and BetaTokenStride already encode both cases. + // + MLAS_LINEAR_ATTENTION_DECAY_LAYOUT DecayLayout; + + // + // At least MlasLinearAttentionBufferSizePerThread(KHeadSize, VHeadSize) + // bytes of thread-private scratch. + // + float* Scratch; +}; + +/** + * @brief Dispatch structure for linear-attention recurrence kernels. + * + * The dispatch granularity is one complete (batch, kv head) sequence: a kernel + * owns the loop over all timesteps for a single state matrix. This is + * deliberate - the state matrix is d_k * d_v floats (8-64 KB, L1/L2 resident), + * so a vectorized kernel must be free to fuse decay, retrieval, update and + * readout into a single streaming pass over it. Dispatching finer-grained + * primitives would force one traversal of that working set per primitive. + */ +struct MLAS_LINEAR_ATTENTION_DISPATCH { + typedef void(ProcessHead_Fn)( + const MLAS_LINEAR_ATTENTION_WORK* Work + ); + + ProcessHead_Fn* ProcessHead = nullptr; +}; + +// +// Portable reference kernel. Registered as the baseline dispatch entry for +// every target, and used directly if no dispatch entry is available. +// +void +MlasLinearAttentionProcessHead( + const MLAS_LINEAR_ATTENTION_WORK* Work +); diff --git a/onnxruntime/core/mlas/lib/linear_attention_kernel_avx512f.cpp b/onnxruntime/core/mlas/lib/linear_attention_kernel_avx512f.cpp new file mode 100644 index 0000000000000..dd2fa8719a301 --- /dev/null +++ b/onnxruntime/core/mlas/lib/linear_attention_kernel_avx512f.cpp @@ -0,0 +1,451 @@ +/*++ + +Copyright (c) Microsoft Corporation. All rights reserved. + +Licensed under the MIT License. + +Module Name: + + linear_attention_kernel_avx512f.cpp + +Abstract: + + This module implements the AVX-512 linear (recurrent) attention kernel. + + The portable kernel walks the state matrix S once per step of the + recurrence - decay, retrieval, rank-1 update and readout - which is roughly + 4 reads and 2 writes of S per token. This kernel fuses all four into two + passes (one read, one write) using the identity below, and processes S one + column panel at a time so the second pass re-reads the panel from L1. + + With dec[i] = exp(g_t[i]) (or 1 when the rule has no decay), and the + pre-weighted vectors k'[i] = dec[i]*k[i] and q'_g[i] = dec[i]*q_g[i]: + + retrieved[j] = sum_i (k'[i] * S_old[i,j]) + upd[j] = beta * (v[j] - retrieved[j]) (= v[j] with no beta) + S_new[i,j] = dec[i] * S_old[i,j] + k[i] * upd[j] + o_g[j] = scale * ( sum_i (q'_g[i] * S_old[i,j]) + (q_g . k) * upd[j] ) + + The last line is the load-bearing one: the readout is expressible from + S_old plus a rank-1 correction, so it never needs the written-back S_new. + Note the rank-1 term sits outside the sum. Pass 1 accumulates `retrieved` + and every head's a_g = sum_i (q'_g[i] * S_old[i,:]) from a single read of + S_old; pass 2 writes S_new in place. + + Both kernels are templated on HAS_DECAY. Without decay the pre-weighting + disappears (k' = k, q' = q, so the kernels read the input vectors directly + and copy nothing), and pass 2 collapses from a multiply plus an FMA to a + single FMA. That matters: for the `linear` rule, which is the one the + portable path handles in the fewest passes, the redundant multiply by 1 was + enough to lose to a tuned SGEMM. + + This reassociates the floating-point sums relative to the portable kernel, + so results agree to tolerance rather than bit-exactly. + +--*/ + +#include "linear_attention.h" + +#include +#include + +// +// With decay, the kernels hold the pre-weighted vectors on the stack, which is +// what bounds KHeadSize. NOUT <= 8 puts the worst case at 9 * 256 * 4 = 9 KB. +// +constexpr size_t MlasLinearAttentionAvx512MaxKHeadSize = 256; + +namespace { + +// +// Single readout head (MHA, or inverse GQA where one query head is shared). +// +// Requires d_k % 16 == 0, d_k <= 256, d_v % 32 == 0. Accumulators stay in ZMM +// registers across both passes; nothing here relies on auto-vectorization, +// which previously spilled the pass-1 accumulators and emitted a serial scalar +// horizontal reduction for q.k. +// +template +MLAS_FORCEINLINE +void +FusedTokenN1Avx512( + float* __restrict S, + size_t d_k, + size_t d_v, + const float* __restrict decvec, + const float* __restrict kt, + const float* __restrict vt, + const float* __restrict q0, + float* __restrict o0, + float scale, + float beta_val, + bool needs_beta, + bool needs_retrieval +) +{ + // + // Staging for the decay-weighted vectors. Without decay the kernel reads q0 + // and kt directly, so the buffers are never written and are sized away. + // + constexpr size_t StageK = HAS_DECAY ? MlasLinearAttentionAvx512MaxKHeadSize : 1; + float wkv_buf[StageK]; + float wqv_buf[StageK]; + + const float* __restrict wqv = q0; + const float* __restrict wkv = kt; + + if constexpr (HAS_DECAY) { + for (size_t i = 0; i < d_k; ++i) { + wqv_buf[i] = decvec[i] * q0[i]; + } + wqv = wqv_buf; + if (needs_retrieval) { + for (size_t i = 0; i < d_k; ++i) { + wkv_buf[i] = decvec[i] * kt[i]; + } + wkv = wkv_buf; + } + } + + // + // qk = q0 . k, accumulated in a vector with a single horizontal reduce. + // + __m512 qkacc = _mm512_setzero_ps(); + for (size_t i = 0; i < d_k; i += 16) { + qkacc = _mm512_fmadd_ps(_mm512_loadu_ps(q0 + i), _mm512_loadu_ps(kt + i), qkacc); + } + + const __m512 sc = _mm512_set1_ps(scale); + const __m512 qkv = _mm512_set1_ps(_mm512_reduce_add_ps(qkacc)); + const __m512 b = _mm512_set1_ps(beta_val); + + // + // One 32-column panel of S at a time: pass 1 -> upd -> output -> pass 2, + // all while the panel (d_k x 32) is hot in L1, so pass 2 re-reads it from + // L1 rather than L2 and the retrieved / a / upd intermediates never leave + // ZMM registers. + // + for (size_t j0 = 0; j0 < d_v; j0 += 32) { + __m512 r0 = _mm512_setzero_ps(); + __m512 r1 = _mm512_setzero_ps(); + __m512 a0 = _mm512_setzero_ps(); + __m512 a1 = _mm512_setzero_ps(); + + if (needs_retrieval) { + for (size_t i = 0; i < d_k; ++i) { + const __m512 wk = _mm512_set1_ps(wkv[i]); + const __m512 wq = _mm512_set1_ps(wqv[i]); + const float* __restrict Si = S + i * d_v + j0; + const __m512 s0 = _mm512_loadu_ps(Si); + const __m512 s1 = _mm512_loadu_ps(Si + 16); + r0 = _mm512_fmadd_ps(wk, s0, r0); + r1 = _mm512_fmadd_ps(wk, s1, r1); + a0 = _mm512_fmadd_ps(wq, s0, a0); + a1 = _mm512_fmadd_ps(wq, s1, a1); + } + } else { + for (size_t i = 0; i < d_k; ++i) { + const __m512 wq = _mm512_set1_ps(wqv[i]); + const float* __restrict Si = S + i * d_v + j0; + a0 = _mm512_fmadd_ps(wq, _mm512_loadu_ps(Si), a0); + a1 = _mm512_fmadd_ps(wq, _mm512_loadu_ps(Si + 16), a1); + } + } + + __m512 u0; + __m512 u1; + if (needs_beta) { + u0 = _mm512_mul_ps(b, _mm512_sub_ps(_mm512_loadu_ps(vt + j0), r0)); + u1 = _mm512_mul_ps(b, _mm512_sub_ps(_mm512_loadu_ps(vt + j0 + 16), r1)); + } else { + u0 = _mm512_loadu_ps(vt + j0); + u1 = _mm512_loadu_ps(vt + j0 + 16); + } + + _mm512_storeu_ps(o0 + j0, _mm512_mul_ps(sc, _mm512_fmadd_ps(qkv, u0, a0))); + _mm512_storeu_ps(o0 + j0 + 16, _mm512_mul_ps(sc, _mm512_fmadd_ps(qkv, u1, a1))); + + for (size_t i = 0; i < d_k; ++i) { + const __m512 kk = _mm512_set1_ps(kt[i]); + float* __restrict Si = S + i * d_v + j0; + if constexpr (HAS_DECAY) { + const __m512 dc = _mm512_set1_ps(decvec[i]); + _mm512_storeu_ps(Si, + _mm512_fmadd_ps(dc, _mm512_loadu_ps(Si), _mm512_mul_ps(kk, u0))); + _mm512_storeu_ps(Si + 16, + _mm512_fmadd_ps(dc, _mm512_loadu_ps(Si + 16), _mm512_mul_ps(kk, u1))); + } else { + _mm512_storeu_ps(Si, _mm512_fmadd_ps(kk, u0, _mm512_loadu_ps(Si))); + _mm512_storeu_ps(Si + 16, _mm512_fmadd_ps(kk, u1, _mm512_loadu_ps(Si + 16))); + } + } + } +} + +// +// Standard GQA: NOUT query heads share this state. Direct generalization of +// FusedTokenN1Avx512 - one read of each S_old panel feeds `retrieved` and all +// NOUT heads' a_g accumulators (1 + NOUT FMAs per S load), then a shared upd, +// per-head output, and the shared in-place pass 2. +// +// Templated on NOUT so the per-head accumulators have compile-time indices and +// stay in registers; a runtime loop would spill them. NLANE is the panel width +// in ZMM registers. Live registers are NLANE * (1 + NOUT) accumulators plus +// NLANE loads - AVX-512 embedded broadcast makes the wkv/wqv multipliers memory +// operands, so the weights cost no registers. +// +// Requires d_k % 16 == 0, d_k <= 256, d_v % (NLANE * 16) == 0. +// +template +MLAS_FORCEINLINE +void +FusedTokenGQAAvx512( + float* __restrict S, + size_t d_k, + size_t d_v, + const float* __restrict decvec, + const float* __restrict kt, + const float* __restrict vt, + const float* __restrict q_base, + float* __restrict o_base, + float scale, + float beta_val, + bool needs_beta, + bool needs_retrieval +) +{ + constexpr size_t MaxK = MlasLinearAttentionAvx512MaxKHeadSize; + constexpr size_t PW = NLANE * 16; // panel width in floats + + // + // Staging for the decay-weighted vectors, sized away without decay: the + // kernel then reads the query rows and kt directly and copies nothing. + // + float wkv_buf[HAS_DECAY ? MaxK : 1]; + float wqv_buf[HAS_DECAY ? NOUT * MaxK : 1]; + + // + // Without decay the per-head weights are just the query rows themselves, + // so nothing is copied. NOUT is a compile-time bound, so this array of + // pointers stays in registers. + // + const float* wqv[NOUT]; + const float* __restrict wkv = kt; + + for (int g = 0; g < NOUT; ++g) { + const float* __restrict qg = q_base + static_cast(g) * d_k; + if constexpr (HAS_DECAY) { + float* __restrict wq = wqv_buf + static_cast(g) * MaxK; + for (size_t i = 0; i < d_k; ++i) { + wq[i] = decvec[i] * qg[i]; + } + wqv[g] = wq; + } else { + wqv[g] = qg; + } + } + + if constexpr (HAS_DECAY) { + if (needs_retrieval) { + for (size_t i = 0; i < d_k; ++i) { + wkv_buf[i] = decvec[i] * kt[i]; + } + wkv = wkv_buf; + } + } + + float qk[NOUT]; + for (int g = 0; g < NOUT; ++g) { + const float* __restrict qg = q_base + static_cast(g) * d_k; + __m512 acc = _mm512_setzero_ps(); + for (size_t i = 0; i < d_k; i += 16) { + acc = _mm512_fmadd_ps(_mm512_loadu_ps(qg + i), _mm512_loadu_ps(kt + i), acc); + } + qk[g] = _mm512_reduce_add_ps(acc); + } + + const __m512 sc = _mm512_set1_ps(scale); + const __m512 bb = _mm512_set1_ps(beta_val); + + for (size_t j0 = 0; j0 < d_v; j0 += PW) { + __m512 r[NLANE]; + __m512 a[NOUT][NLANE]; + for (int l = 0; l < NLANE; ++l) { + r[l] = _mm512_setzero_ps(); + } + for (int g = 0; g < NOUT; ++g) { + for (int l = 0; l < NLANE; ++l) { + a[g][l] = _mm512_setzero_ps(); + } + } + + for (size_t i = 0; i < d_k; ++i) { + const float* __restrict Si = S + i * d_v + j0; + __m512 s[NLANE]; + for (int l = 0; l < NLANE; ++l) { + s[l] = _mm512_loadu_ps(Si + l * 16); + } + if (needs_retrieval) { + const __m512 wk = _mm512_set1_ps(wkv[i]); + for (int l = 0; l < NLANE; ++l) { + r[l] = _mm512_fmadd_ps(wk, s[l], r[l]); + } + } + for (int g = 0; g < NOUT; ++g) { + const __m512 wq = _mm512_set1_ps(wqv[g][i]); + for (int l = 0; l < NLANE; ++l) { + a[g][l] = _mm512_fmadd_ps(wq, s[l], a[g][l]); + } + } + } + + __m512 u[NLANE]; + for (int l = 0; l < NLANE; ++l) { + const __m512 vv = _mm512_loadu_ps(vt + j0 + l * 16); + u[l] = needs_beta ? _mm512_mul_ps(bb, _mm512_sub_ps(vv, r[l])) : vv; + } + + for (int g = 0; g < NOUT; ++g) { + const __m512 qkg = _mm512_set1_ps(qk[g]); + float* __restrict og = o_base + static_cast(g) * d_v + j0; + for (int l = 0; l < NLANE; ++l) { + _mm512_storeu_ps(og + l * 16, + _mm512_mul_ps(sc, _mm512_fmadd_ps(qkg, u[l], a[g][l]))); + } + } + + for (size_t i = 0; i < d_k; ++i) { + const __m512 kk = _mm512_set1_ps(kt[i]); + float* __restrict Si = S + i * d_v + j0; + if constexpr (HAS_DECAY) { + const __m512 dc = _mm512_set1_ps(decvec[i]); + for (int l = 0; l < NLANE; ++l) { + _mm512_storeu_ps(Si + l * 16, + _mm512_fmadd_ps(dc, _mm512_loadu_ps(Si + l * 16), + _mm512_mul_ps(kk, u[l]))); + } + } else { + for (int l = 0; l < NLANE; ++l) { + _mm512_storeu_ps(Si + l * 16, + _mm512_fmadd_ps(kk, u[l], _mm512_loadu_ps(Si + l * 16))); + } + } + } + } +} + +template +MLAS_FORCEINLINE +void +FusedTokenAvx512( + size_t n_out, + float* __restrict S, + size_t d_k, + size_t d_v, + const float* __restrict decvec, + const float* __restrict kt, + const float* __restrict vt, + const float* __restrict q_base, + float* __restrict o_base, + float scale, + float beta_val, + bool needs_beta, + bool needs_retrieval +) +{ + switch (n_out) { + case 1: + FusedTokenN1Avx512(S, d_k, d_v, decvec, kt, vt, q_base, o_base, + scale, beta_val, needs_beta, needs_retrieval); + break; + case 2: + FusedTokenGQAAvx512<2, 2, HAS_DECAY>(S, d_k, d_v, decvec, kt, vt, q_base, o_base, + scale, beta_val, needs_beta, needs_retrieval); + break; + case 4: + FusedTokenGQAAvx512<4, 2, HAS_DECAY>(S, d_k, d_v, decvec, kt, vt, q_base, o_base, + scale, beta_val, needs_beta, needs_retrieval); + break; + default: // 8 + FusedTokenGQAAvx512<8, 2, HAS_DECAY>(S, d_k, d_v, decvec, kt, vt, q_base, o_base, + scale, beta_val, needs_beta, needs_retrieval); + break; + } +} + +} // namespace + +void +MlasLinearAttentionProcessHeadAvx512F( + const MLAS_LINEAR_ATTENTION_WORK* Work +) +{ + const size_t d_k = Work->KHeadSize; + const size_t d_v = Work->VHeadSize; + const size_t n_out = Work->HeadsPerGroup; + + // + // The panel loops are unmasked, and the pre-weight buffers are fixed-size. + // Anything outside that envelope goes to the portable kernel, which has no + // shape restrictions. n_out == 16 is deliberately excluded: benchmarking + // showed the 16-column single-lane variant it would need cannot match a + // batched GEMM, unlike n_out in {1, 2, 4, 8}. + // + const bool shape_ok = (d_k % 16 == 0) && + (d_k <= MlasLinearAttentionAvx512MaxKHeadSize) && + (d_v % 32 == 0); + const bool n_out_ok = (n_out == 1 || n_out == 2 || n_out == 4 || n_out == 8); + + if (!shape_ok || !n_out_ok) { + MlasLinearAttentionProcessHead(Work); + return; + } + + const bool needs_decay = (Work->Rule == MlasLinearAttentionRuleGated || + Work->Rule == MlasLinearAttentionRuleGatedDelta); + const bool needs_beta = (Work->Rule == MlasLinearAttentionRuleDelta || + Work->Rule == MlasLinearAttentionRuleGatedDelta); + const bool needs_retrieval = needs_beta; + const bool decay_per_key_dim = (Work->DecayLayout == MlasLinearAttentionDecayPerKeyDim); + + float* __restrict S = Work->State; + float* __restrict decvec = Work->Scratch; // d_k floats + + for (size_t t = 0; t < Work->SequenceLength; ++t) { + const float* __restrict kt = Work->Key + t * Work->KeyTokenStride; + const float* __restrict vt = Work->Value + t * Work->ValueTokenStride; + + const float beta_val = needs_beta ? Work->Beta[t * Work->BetaTokenStride] : 0.0f; + + const float* __restrict q_base = Work->Query + t * Work->QueryTokenStride; + float* __restrict o_base = Work->Output + t * Work->OutputTokenStride; + + if (needs_decay) { + // + // Materialize the decay vector once per token, so the two passes + // and the pre-weighting all read it instead of recomputing exp. + // + if (decay_per_key_dim) { + const float* gt = Work->Decay + t * Work->DecayTokenStride; + for (size_t i = 0; i < d_k; ++i) { + decvec[i] = std::exp(gt[i]); + } + } else { + const float exp_g = std::exp(Work->Decay[t * Work->DecayTokenStride]); + for (size_t i = 0; i < d_k; ++i) { + decvec[i] = exp_g; + } + } + + FusedTokenAvx512(n_out, S, d_k, d_v, decvec, kt, vt, q_base, o_base, + Work->Scale, beta_val, needs_beta, needs_retrieval); + } else { + FusedTokenAvx512(n_out, S, d_k, d_v, nullptr, kt, vt, q_base, o_base, + Work->Scale, beta_val, needs_beta, needs_retrieval); + } + } +} + +const MLAS_LINEAR_ATTENTION_DISPATCH MlasLinearAttentionDispatchAvx512F = { + MlasLinearAttentionProcessHeadAvx512F +}; diff --git a/onnxruntime/core/mlas/lib/mlasi.h b/onnxruntime/core/mlas/lib/mlasi.h index b8c0ea1690294..6ceeff96755f9 100644 --- a/onnxruntime/core/mlas/lib/mlasi.h +++ b/onnxruntime/core/mlas/lib/mlasi.h @@ -1597,6 +1597,13 @@ extern const MLAS_KV_QUANT_GEMM_DISPATCH MlasKVQuantGemmDispatchAvx2; extern const MLAS_KV_QUANT_GEMM_DISPATCH MlasKVQuantGemmDispatchAvx512Vnni; extern const MLAS_KV_QUANT_GEMM_DISPATCH MlasKVQuantGemmDispatchNeon; +// +// Linear (recurrent) attention dispatch structure. +// +struct MLAS_LINEAR_ATTENTION_DISPATCH; +extern const MLAS_LINEAR_ATTENTION_DISPATCH MlasLinearAttentionDispatchDefault; +extern const MLAS_LINEAR_ATTENTION_DISPATCH MlasLinearAttentionDispatchAvx512F; + // // Quantized depthwise convolution kernels. // @@ -1845,6 +1852,7 @@ MLAS_COMPUTE_TANH_FP16_KERNEL* TanhFP16KernelRoutine = nullptr; const MLAS_SOFTMAX_DISPATCH* SoftmaxDispatch{nullptr}; const MLAS_ELTWISE_DISPATCH* EltwiseDispatch{nullptr}; const MLAS_KV_QUANT_GEMM_DISPATCH* KVQuantGemmDispatch{nullptr}; + const MLAS_LINEAR_ATTENTION_DISPATCH* LinearAttentionDispatch{nullptr}; }; inline diff --git a/onnxruntime/core/mlas/lib/platform.cpp b/onnxruntime/core/mlas/lib/platform.cpp index 68519a5987fd0..50c3764595378 100644 --- a/onnxruntime/core/mlas/lib/platform.cpp +++ b/onnxruntime/core/mlas/lib/platform.cpp @@ -277,6 +277,12 @@ Return Value: this->CastF16ToF32Kernel = nullptr; this->CastF32ToF16Kernel = nullptr; + // + // Linear attention: portable baseline for every target. ISA-specific + // implementations override this in the target-specific blocks below. + // + this->LinearAttentionDispatch = &MlasLinearAttentionDispatchDefault; + #if defined(MLAS_TARGET_RISCV64) this->GemmFloatKernel = nullptr; this->GemmU8S8Dispatch = &MlasGemmQuantDispatchDefault; @@ -562,6 +568,7 @@ Return Value: this->ComputeExpF32Kernel = MlasComputeExpF32KernelAvx512F; this->ComputeSumExpF32Kernel = MlasComputeSumExpF32KernelAvx512F; this->ReduceMaximumF32Kernel = MlasReduceMaximumF32KernelAvx512F; + this->LinearAttentionDispatch = &MlasLinearAttentionDispatchAvx512F; this->QuantizeLinearS8Kernel = MlasQuantizeLinearS8KernelAvx512F; this->QuantizeLinearU8Kernel = MlasQuantizeLinearU8KernelAvx512F; this->NchwcBlockSize = 16; diff --git a/onnxruntime/test/contrib_ops/linear_attention_op_test.cc b/onnxruntime/test/contrib_ops/linear_attention_op_test.cc index 3c3a8e71e105f..365119a8479dd 100644 --- a/onnxruntime/test/contrib_ops/linear_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/linear_attention_op_test.cc @@ -145,7 +145,7 @@ void LinearAttentionGQAReference( const std::vector* initial_state, // (B, kv_num_heads, dk, dv) const std::vector* decay, // (B, kv_num_heads, T[, dk]) const std::vector* beta, // (B, kv_num_heads, T) - std::vector& output, // (B, kv_num_heads, T, dv) + std::vector& output, // (B, max(q_num_heads, kv_num_heads), T, dv) std::vector& final_state) { // (B, kv_num_heads, dk, dv) int bht_kv = batch_size * kv_num_heads * seq_length; bool decay_broadcast_dk = (decay != nullptr && static_cast(decay->size()) == bht_kv); @@ -1429,10 +1429,21 @@ TEST(ContribOpLinearAttentionTest, GatedDeltaRule_StandardGQA_N16) { RunStandardGQA("gated_delta", 16, 1, 32, 64); } +// All four update rules at a fixed n_out=4 ratio: "linear"/"gated" take the no-retrieval +// (no-beta) branch, "delta"/"gated_delta" the retrieval branch, and "gated"/"gated_delta" +// additionally apply the decay gate. TEST(ContribOpLinearAttentionTest, LinearRule_StandardGQA_N4) { RunStandardGQA("linear", 8, 2, 32, 64); } +TEST(ContribOpLinearAttentionTest, GatedRule_StandardGQA_N4) { + RunStandardGQA("gated", 8, 2, 32, 64); +} + +TEST(ContribOpLinearAttentionTest, DeltaRule_StandardGQA_N4) { + RunStandardGQA("delta", 8, 2, 32, 64); +} + TEST(ContribOpLinearAttentionTest, GatedDeltaRule_StandardGQA_N4_Dim128) { RunStandardGQA("gated_delta", 8, 2, 128, 128); } diff --git a/onnxruntime/test/mlas/bench/bench_linear_attention.cpp b/onnxruntime/test/mlas/bench/bench_linear_attention.cpp new file mode 100644 index 0000000000000..b4fce247e3561 --- /dev/null +++ b/onnxruntime/test/mlas/bench/bench_linear_attention.cpp @@ -0,0 +1,245 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +// +// Benchmark for the linear (recurrent) attention kernel, MlasLinearAttention. +// +// BM_LinearAttentionScalar is a self-contained scalar recurrence in this +// translation unit - no SGEMM, no dispatch - so it stays a fixed reference +// point as ISA-specific kernels are added behind the dispatch. +// + +#include "mlas.h" +#include "core/util/thread_utils.h" +#include "benchmark/benchmark.h" +#include "bench_util.h" + +#include +#include +#include +#include + +namespace { + +MLAS_THREADPOOL* GetBenchThreadPool() { + static OrtThreadPoolParams tpo = []() { + OrtThreadPoolParams params{}; + params.thread_pool_size = 8; + params.auto_set_affinity = true; + return params; + }(); + static std::unique_ptr tp( + onnxruntime::concurrency::CreateThreadPool(&onnxruntime::Env::Default(), + tpo, onnxruntime::concurrency::ThreadPoolType::INTRA_OP)); + return tp.get(); +} + +struct LinearAttentionProblem { + int batch_size, seq_len, q_num_heads, kv_num_heads, k_num_heads, k_head_size, v_head_size; + MLAS_LINEAR_ATTENTION_RULE rule; + + std::vector query, key, value, decay, beta, state, output; + + explicit LinearAttentionProblem(benchmark::State& state_arg) { + batch_size = static_cast(state_arg.range(0)); + seq_len = static_cast(state_arg.range(1)); + q_num_heads = static_cast(state_arg.range(2)); + kv_num_heads = static_cast(state_arg.range(3)); + k_num_heads = static_cast(state_arg.range(4)); + k_head_size = static_cast(state_arg.range(5)); + v_head_size = static_cast(state_arg.range(6)); + rule = static_cast(state_arg.range(7)); + + const size_t tokens = static_cast(batch_size) * seq_len; + query = RandomVectorUniform(tokens * q_num_heads * k_head_size, -1.0f, 1.0f); + key = RandomVectorUniform(tokens * k_num_heads * k_head_size, -1.0f, 1.0f); + value = RandomVectorUniform(tokens * kv_num_heads * v_head_size, -1.0f, 1.0f); + // exp(g) <= 1 keeps the recurrence bounded over long sequences. + decay = RandomVectorUniform(tokens * kv_num_heads * k_head_size, -1.0f, 0.0f); + beta = RandomVectorUniform(tokens * kv_num_heads, 0.0f, 1.0f); + state.assign(static_cast(batch_size) * kv_num_heads * k_head_size * v_head_size, 0.0f); + output.assign(tokens * MlasLinearAttentionOutputHiddenSize(q_num_heads, kv_num_heads, v_head_size), + 0.0f); + } + + bool HasDecay() const { + return rule == MlasLinearAttentionRuleGated || rule == MlasLinearAttentionRuleGatedDelta; + } + + bool HasBeta() const { + return rule == MlasLinearAttentionRuleDelta || rule == MlasLinearAttentionRuleGatedDelta; + } + + // Per-invocation flop estimate, used for the throughput counter. + double Flops() const { + const double sd = static_cast(k_head_size) * v_head_size; + const double group = (q_num_heads >= kv_num_heads) + ? static_cast(q_num_heads) / kv_num_heads + : 1.0; + double per_token = 2.0 * sd; // rank-1 state update + per_token += 2.0 * sd * group; // query readout + if (HasDecay()) per_token += sd; // decay + if (HasBeta()) per_token += 2.0 * sd; // retrieval + return per_token * batch_size * seq_len * kv_num_heads; + } +}; + +// Scalar recurrence, no SGEMM and no dispatch: the "no optimization" baseline. +void ScalarLinearAttention(LinearAttentionProblem& p) { + const size_t dk = static_cast(p.k_head_size); + const size_t dv = static_cast(p.v_head_size); + const size_t T = static_cast(p.seq_len); + const size_t H_q = static_cast(p.q_num_heads); + const size_t H_kv = static_cast(p.kv_num_heads); + const size_t H_k = static_cast(p.k_num_heads); + const size_t q_stride = H_q * dk; + const size_t k_stride = H_k * dk; + const size_t v_stride = H_kv * dv; + const size_t o_stride = MlasLinearAttentionOutputHiddenSize(p.q_num_heads, p.kv_num_heads, + p.v_head_size); + const size_t decay_stride = H_kv * dk; + const size_t beta_stride = H_kv; + const float scale = 1.0f / std::sqrt(static_cast(dk)); + const bool has_decay = p.HasDecay(); + const bool has_beta = p.HasBeta(); + const size_t group = (H_q >= H_kv) ? H_q / H_kv : 1; + + std::vector retrieved(dv); + + for (size_t b = 0; b < static_cast(p.batch_size); ++b) { + for (size_t h_kv = 0; h_kv < H_kv; ++h_kv) { + const size_t h_k = h_kv / (H_kv / H_k); + float* S = p.state.data() + (b * H_kv + h_kv) * dk * dv; + + for (size_t t = 0; t < T; ++t) { + const float* kt = p.key.data() + (b * T + t) * k_stride + h_k * dk; + const float* vt = p.value.data() + (b * T + t) * v_stride + h_kv * dv; + + if (has_decay) { + const float* gt = p.decay.data() + (b * T + t) * decay_stride + h_kv * dk; + for (size_t i = 0; i < dk; ++i) { + const float g = std::exp(gt[i]); + for (size_t j = 0; j < dv; ++j) { + S[i * dv + j] *= g; + } + } + } + + if (has_beta) { + for (size_t j = 0; j < dv; ++j) { + float acc = 0.0f; + for (size_t i = 0; i < dk; ++i) { + acc += S[i * dv + j] * kt[i]; + } + retrieved[j] = acc; + } + const float bt = p.beta[(b * T + t) * beta_stride + h_kv]; + for (size_t j = 0; j < dv; ++j) { + retrieved[j] = bt * (vt[j] - retrieved[j]); + } + } else { + std::copy_n(vt, dv, retrieved.begin()); + } + + for (size_t i = 0; i < dk; ++i) { + const float ki = kt[i]; + for (size_t j = 0; j < dv; ++j) { + S[i * dv + j] += ki * retrieved[j]; + } + } + + for (size_t g = 0; g < group; ++g) { + const size_t h_q = (H_q >= H_kv) ? (h_kv * group + g) : (h_kv * H_q / H_kv); + const size_t h_out = (H_q >= H_kv) ? (h_kv * group + g) : h_kv; + const float* qt = p.query.data() + (b * T + t) * q_stride + h_q * dk; + float* ot = p.output.data() + (b * T + t) * o_stride + h_out * dv; + for (size_t j = 0; j < dv; ++j) { + float acc = 0.0f; + for (size_t i = 0; i < dk; ++i) { + acc += qt[i] * S[i * dv + j]; + } + ot[j] = scale * acc; + } + } + } + } + } +} + +void LinearAttentionArgs(benchmark::internal::Benchmark* b) { + b->ArgNames({"B", "T", "Hq", "Hkv", "Hk", "dk", "dv", "rule"}); + for (int rule : {static_cast(MlasLinearAttentionRuleLinear), + static_cast(MlasLinearAttentionRuleGated), + static_cast(MlasLinearAttentionRuleDelta), + static_cast(MlasLinearAttentionRuleGatedDelta)}) { + // Below the SGEMM threshold (32*64 = 2048). + for (int t : {1, 128, 1024}) { + b->Args({1, t, 16, 16, 16, 32, 64, rule}); + } + // At and above the threshold: production gated-delta-net shape. + for (int t : {1, 128, 1024, 4096}) { + b->Args({1, t, 16, 16, 16, 128, 128, rule}); + } + b->Args({1, 1024, 32, 8, 8, 128, 128, rule}); // standard GQA, 4 q heads per kv head + b->Args({1, 1024, 16, 32, 16, 128, 128, rule}); // inverse GQA (Qwen3.5 9B style) + b->Args({8, 512, 16, 16, 16, 128, 128, rule}); // batched prefill: multi-thread partition + b->Args({1, 4096, 1, 1, 1, 128, 128, rule}); // B*Hkv == 1: inline fast path + } +} + +} // namespace + +static void BM_LinearAttentionScalar(benchmark::State& state) { + LinearAttentionProblem p(state); + + ScalarLinearAttention(p); // warm up + + for (auto _ : state) { + std::fill(p.state.begin(), p.state.end(), 0.0f); + ScalarLinearAttention(p); + benchmark::DoNotOptimize(p.output.data()); + } + state.counters["flops"] = benchmark::Counter(p.Flops(), benchmark::Counter::kIsIterationInvariantRate); +} + +static void BM_LinearAttention(benchmark::State& state) { + LinearAttentionProblem p(state); + auto* tp = GetBenchThreadPool(); + + MlasLinearAttentionArgs args; + args.batch_size = p.batch_size; + args.sequence_length = p.seq_len; + args.q_num_heads = p.q_num_heads; + args.kv_num_heads = p.kv_num_heads; + args.k_num_heads = p.k_num_heads; + args.k_head_size = p.k_head_size; + args.v_head_size = p.v_head_size; + args.rule = p.rule; + args.decay_layout = p.HasDecay() ? MlasLinearAttentionDecayPerKeyDim : MlasLinearAttentionDecayNone; + args.beta_layout = p.HasBeta() ? MlasLinearAttentionBetaPerHead : MlasLinearAttentionBetaNone; + args.scale = 1.0f / std::sqrt(static_cast(p.k_head_size)); + args.thread_count = onnxruntime::concurrency::ThreadPool::DegreeOfParallelism(tp); + args.buffer_size_per_thread = MlasLinearAttentionBufferSizePerThread(p.k_head_size, p.v_head_size); + + std::vector scratch((args.buffer_size_per_thread / sizeof(float)) * args.thread_count); + args.buffer = scratch.data(); + args.query = p.query.data(); + args.key = p.key.data(); + args.value = p.value.data(); + args.decay = p.HasDecay() ? p.decay.data() : nullptr; + args.beta = p.HasBeta() ? p.beta.data() : nullptr; + args.state = p.state.data(); + args.output = p.output.data(); + + MlasLinearAttention(&args, tp); // warm up + + for (auto _ : state) { + std::fill(p.state.begin(), p.state.end(), 0.0f); + MlasLinearAttention(&args, tp); + benchmark::DoNotOptimize(p.output.data()); + } + state.counters["flops"] = benchmark::Counter(p.Flops(), benchmark::Counter::kIsIterationInvariantRate); +} + +BENCHMARK(BM_LinearAttentionScalar)->Apply(LinearAttentionArgs)->UseRealTime(); +BENCHMARK(BM_LinearAttention)->Apply(LinearAttentionArgs)->UseRealTime(); diff --git a/onnxruntime/test/mlas/unittest/test_linear_attention.cpp b/onnxruntime/test/mlas/unittest/test_linear_attention.cpp new file mode 100644 index 0000000000000..f3dd2d76c65a5 --- /dev/null +++ b/onnxruntime/test/mlas/unittest/test_linear_attention.cpp @@ -0,0 +1,367 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "test_util.h" + +#include +#include +#include +#include + +// +// Tests MlasLinearAttention against an independent scalar oracle written +// straight from the recurrence definition. +// +class MlasLinearAttentionTest : public MlasTestBase { + private: + MatrixGuardBuffer BufferQ; + MatrixGuardBuffer BufferK; + MatrixGuardBuffer BufferV; + MatrixGuardBuffer BufferDecay; + MatrixGuardBuffer BufferBeta; + MatrixGuardBuffer BufferInitialState; + MatrixGuardBuffer BufferState; + MatrixGuardBuffer BufferStateRef; + MatrixGuardBuffer BufferOutput; + MatrixGuardBuffer BufferOutputRef; + MatrixGuardBuffer BufferScratch; + + static void Fill(float* data, size_t count, std::mt19937& rng, float lo, float hi) { + std::uniform_real_distribution dist(lo, hi); + for (size_t i = 0; i < count; ++i) { + data[i] = dist(rng); + } + } + + static bool RuleHasDecay(MLAS_LINEAR_ATTENTION_RULE rule) { + return rule == MlasLinearAttentionRuleGated || rule == MlasLinearAttentionRuleGatedDelta; + } + + static bool RuleHasBeta(MLAS_LINEAR_ATTENTION_RULE rule) { + return rule == MlasLinearAttentionRuleDelta || rule == MlasLinearAttentionRuleGatedDelta; + } + + // + // Independent reference. Deliberately written from the recurrence rather than + // adapted from the MLAS kernel, and it recomputes every index mapping + // (key-head sharing, standard/inverse GQA) from first principles. + // + static void Reference(const MlasLinearAttentionArgs& args, + const float* initial_state, + float* state, + float* output) { + const size_t dk = static_cast(args.k_head_size); + const size_t dv = static_cast(args.v_head_size); + const size_t T = static_cast(args.sequence_length); + const size_t H_q = static_cast(args.q_num_heads); + const size_t H_kv = static_cast(args.kv_num_heads); + const size_t H_k = static_cast(args.k_num_heads); + + const size_t q_stride = H_q * dk; + const size_t k_stride = H_k * dk; + const size_t v_stride = H_kv * dv; + const size_t o_stride = MlasLinearAttentionOutputHiddenSize( + args.q_num_heads, args.kv_num_heads, args.v_head_size); + + const bool has_decay = RuleHasDecay(args.rule); + const bool has_beta = RuleHasBeta(args.rule); + const bool decay_per_key_dim = (args.decay_layout == MlasLinearAttentionDecayPerKeyDim); + const size_t decay_stride = decay_per_key_dim ? H_kv * dk : H_kv; + const size_t beta_stride = (args.beta_layout == MlasLinearAttentionBetaPerHead) ? H_kv : 1; + + const size_t state_per_head = dk * dv; + std::vector retrieved(dv); + + for (size_t b = 0; b < static_cast(args.batch_size); ++b) { + for (size_t h_kv = 0; h_kv < H_kv; ++h_kv) { + // Consecutive kv heads share a key head (floor division, not modulo). + const size_t h_k = h_kv / (H_kv / H_k); + + float* S = state + (b * H_kv + h_kv) * state_per_head; + std::copy_n(initial_state + (b * H_kv + h_kv) * state_per_head, state_per_head, S); + + for (size_t t = 0; t < T; ++t) { + const float* kt = args.key + (b * T + t) * k_stride + h_k * dk; + const float* vt = args.value + (b * T + t) * v_stride + h_kv * dv; + + if (has_decay) { + const float* gt = args.decay + (b * T + t) * decay_stride + + (decay_per_key_dim ? h_kv * dk : h_kv); + for (size_t i = 0; i < dk; ++i) { + const float g = std::exp(decay_per_key_dim ? gt[i] : gt[0]); + for (size_t j = 0; j < dv; ++j) { + S[i * dv + j] *= g; + } + } + } + + if (has_beta) { + for (size_t j = 0; j < dv; ++j) { + double acc = 0.0; + for (size_t i = 0; i < dk; ++i) { + acc += static_cast(S[i * dv + j]) * kt[i]; + } + retrieved[j] = static_cast(acc); + } + const float bt = args.beta[(b * T + t) * beta_stride + + ((args.beta_layout == MlasLinearAttentionBetaPerHead) ? h_kv : 0)]; + for (size_t j = 0; j < dv; ++j) { + retrieved[j] = bt * (vt[j] - retrieved[j]); + } + } else { + std::copy_n(vt, dv, retrieved.begin()); + } + + for (size_t i = 0; i < dk; ++i) { + for (size_t j = 0; j < dv; ++j) { + S[i * dv + j] += kt[i] * retrieved[j]; + } + } + + // Readout. Standard GQA gives each kv head H_q/H_kv consecutive query + // heads; inverse GQA gives it one shared query head. + if (H_q >= H_kv) { + const size_t group = H_q / H_kv; + for (size_t g = 0; g < group; ++g) { + const size_t h_q = h_kv * group + g; + const float* qt = args.query + (b * T + t) * q_stride + h_q * dk; + float* ot = output + (b * T + t) * o_stride + h_q * dv; + for (size_t j = 0; j < dv; ++j) { + double acc = 0.0; + for (size_t i = 0; i < dk; ++i) { + acc += static_cast(qt[i]) * S[i * dv + j]; + } + ot[j] = args.scale * static_cast(acc); + } + } + } else { + const size_t h_q = h_kv * H_q / H_kv; + const float* qt = args.query + (b * T + t) * q_stride + h_q * dk; + float* ot = output + (b * T + t) * o_stride + h_kv * dv; + for (size_t j = 0; j < dv; ++j) { + double acc = 0.0; + for (size_t i = 0; i < dk; ++i) { + acc += static_cast(qt[i]) * S[i * dv + j]; + } + ot[j] = args.scale * static_cast(acc); + } + } + } + } + } + } + + void Test(int B, int T, int H_q, int H_kv, int H_k, int dk, int dv, + MLAS_LINEAR_ATTENTION_RULE rule, + MLAS_LINEAR_ATTENTION_DECAY_LAYOUT decay_layout, + MLAS_LINEAR_ATTENTION_BETA_LAYOUT beta_layout, + float scale, int thread_count, bool nonzero_initial_state) { + const bool has_decay = RuleHasDecay(rule); + const bool has_beta = RuleHasBeta(rule); + if (!has_decay) { + decay_layout = MlasLinearAttentionDecayNone; + } + if (!has_beta) { + beta_layout = MlasLinearAttentionBetaNone; + } + + const size_t tokens = static_cast(B) * T; + const size_t q_elems = tokens * H_q * dk; + const size_t k_elems = tokens * H_k * dk; + const size_t v_elems = tokens * H_kv * dv; + const size_t o_elems = tokens * MlasLinearAttentionOutputHiddenSize(H_q, H_kv, dv); + const size_t s_elems = static_cast(B) * H_kv * dk * dv; + const size_t decay_elems = + has_decay ? tokens * ((decay_layout == MlasLinearAttentionDecayPerKeyDim) + ? static_cast(H_kv) * dk + : static_cast(H_kv)) + : 1; + const size_t beta_elems = + has_beta ? tokens * ((beta_layout == MlasLinearAttentionBetaPerHead) + ? static_cast(H_kv) + : size_t(1)) + : 1; + + // Seed from the shape so failures are reproducible. + std::mt19937 rng(static_cast(B * 1000003 + T * 10007 + H_q * 1009 + + H_kv * 101 + H_k * 31 + dk * 7 + dv + + static_cast(rule) * 3 + thread_count)); + + float* q = BufferQ.GetBuffer(q_elems); + float* k = BufferK.GetBuffer(k_elems); + float* v = BufferV.GetBuffer(v_elems); + float* decay = BufferDecay.GetBuffer(decay_elems); + float* beta = BufferBeta.GetBuffer(beta_elems); + float* initial_state = BufferInitialState.GetBuffer(s_elems); + float* state = BufferState.GetBuffer(s_elems); + float* state_ref = BufferStateRef.GetBuffer(s_elems); + float* output = BufferOutput.GetBuffer(o_elems); + float* output_ref = BufferOutputRef.GetBuffer(o_elems); + + Fill(q, q_elems, rng, -1.0f, 1.0f); + Fill(k, k_elems, rng, -1.0f, 1.0f); + Fill(v, v_elems, rng, -1.0f, 1.0f); + // Keep exp(g) <= 1 so the recurrence stays bounded over long sequences. + Fill(decay, decay_elems, rng, -1.0f, 0.0f); + Fill(beta, beta_elems, rng, 0.0f, 1.0f); + if (nonzero_initial_state) { + Fill(initial_state, s_elems, rng, -0.5f, 0.5f); + } else { + std::fill_n(initial_state, s_elems, 0.0f); + } + + // MLAS never initializes the state; the caller seeds it. + std::copy_n(initial_state, s_elems, state); + + MlasLinearAttentionArgs args; + args.batch_size = B; + args.sequence_length = T; + args.q_num_heads = H_q; + args.kv_num_heads = H_kv; + args.k_num_heads = H_k; + args.k_head_size = dk; + args.v_head_size = dv; + args.rule = rule; + args.decay_layout = decay_layout; + args.beta_layout = beta_layout; + args.scale = scale; + args.thread_count = thread_count; + args.buffer_size_per_thread = MlasLinearAttentionBufferSizePerThread(dk, dv); + args.buffer = BufferScratch.GetBuffer( + (args.buffer_size_per_thread / sizeof(float)) * thread_count); + args.query = q; + args.key = k; + args.value = v; + args.decay = has_decay ? decay : nullptr; + args.beta = has_beta ? beta : nullptr; + args.state = state; + args.output = output; + + // Run the oracle from a *copy* of the initial state - args.state is updated + // in place, so sharing the buffer would compare a result against itself. + Reference(args, initial_state, state_ref, output_ref); + + MlasLinearAttention(&args, thread_count > 1 ? GetMlasThreadPool() : nullptr); + + std::ostringstream tag; + tag << "B=" << B << " T=" << T << " Hq=" << H_q << " Hkv=" << H_kv << " Hk=" << H_k + << " dk=" << dk << " dv=" << dv << " rule=" << static_cast(rule) + << " decay=" << static_cast(decay_layout) << " beta=" << static_cast(beta_layout) + << " scale=" << scale << " threads=" << thread_count + << " init_state=" << nonzero_initial_state; + const std::string context = tag.str(); + + for (size_t i = 0; i < o_elems; ++i) { + const float tol = 1e-4f * std::max(1.0f, std::abs(output_ref[i])); + ASSERT_NEAR(output[i], output_ref[i], tol) << "output[" << i << "] " << context; + } + // The state is an operator output too: a bad update can be invisible in + // `output` at short sequence lengths. + for (size_t i = 0; i < s_elems; ++i) { + const float tol = 1e-4f * std::max(1.0f, std::abs(state_ref[i])); + ASSERT_NEAR(state[i], state_ref[i], tol) << "state[" << i << "] " << context; + } + } + + public: + static const char* GetTestSuiteName() { + return "LinearAttention"; + } + + void ExecuteShort(void) override { + // (d_k, d_v) pairs chosen to straddle two independent boundaries: + // - the d_k*d_v >= 4096 SGEMM threshold in the portable kernel (4096 is + // the inclusive boundary); + // - the AVX-512 kernel's eligibility envelope, d_k % 16 == 0 && + // d_k <= 256 && d_v % 32 == 0. {32,64}, {64,64}, {128,128} are + // eligible; the rest exercise the fallback to the portable kernel, + // including {272,32} which fails only the d_k <= 256 bound. + static const int kShapes[][2] = { + {8, 8}, + {16, 16}, + {32, 64}, + {24, 40}, + {64, 64}, + {48, 96}, + {128, 128}, + {272, 32}, + }; + static const MLAS_LINEAR_ATTENTION_RULE kRules[] = { + MlasLinearAttentionRuleLinear, + MlasLinearAttentionRuleGated, + MlasLinearAttentionRuleDelta, + MlasLinearAttentionRuleGatedDelta, + }; + // {H_q, H_kv, H_k}: MHA, standard GQA, standard GQA + key sharing, key + // sharing alone, inverse GQA, inverse GQA + key sharing, large group. + // The implied heads-per-group covers every value the AVX-512 kernel + // specializes on (1, 2, 4, 8) plus 16, which must fall back. + static const int kHeads[][3] = { + {1, 1, 1}, // n_out = 1 + {2, 2, 2}, // n_out = 1 + {4, 2, 2}, // n_out = 2 + {4, 2, 1}, // n_out = 2, key sharing + {8, 2, 2}, // n_out = 4 + {8, 2, 1}, // n_out = 4, key sharing + {16, 2, 1}, // n_out = 8, key sharing + {16, 1, 1}, // n_out = 16 - outside the kernel's specializations + {4, 4, 2}, // n_out = 1, key sharing + {2, 4, 4}, // inverse GQA + {2, 4, 2}, // inverse GQA + key sharing + }; + + // All rules x all shapes x both layouts, at a fixed head config. + for (auto rule : kRules) { + for (const auto& shape : kShapes) { + Test(1, 5, 2, 2, 2, shape[0], shape[1], rule, + MlasLinearAttentionDecayPerHead, MlasLinearAttentionBetaPerHead, + 0.125f, 1, false); + Test(1, 5, 2, 2, 2, shape[0], shape[1], rule, + MlasLinearAttentionDecayPerKeyDim, MlasLinearAttentionBetaShared, + 0.125f, 1, true); + } + } + + // All head configs, on both sides of the threshold, for the two rules that + // exercise every step of the recurrence and the simplest one. + for (const auto& heads : kHeads) { + for (auto rule : {MlasLinearAttentionRuleLinear, MlasLinearAttentionRuleGatedDelta}) { + Test(2, 7, heads[0], heads[1], heads[2], 32, 64, rule, + MlasLinearAttentionDecayPerKeyDim, MlasLinearAttentionBetaPerHead, + 0.125f, 1, true); + Test(1, 3, heads[0], heads[1], heads[2], 64, 64, rule, + MlasLinearAttentionDecayPerHead, MlasLinearAttentionBetaShared, + 1.0f, 1, false); + } + } + + // Sequence lengths, including the T=1 decode case, and scale == 1.0 which + // skips the post-SGEMM scaling pass. + for (int T : {1, 3, 8, 17}) { + Test(1, T, 4, 2, 2, 64, 64, MlasLinearAttentionRuleGatedDelta, + MlasLinearAttentionDecayPerKeyDim, MlasLinearAttentionBetaPerHead, + 1.0f, 1, true); + Test(2, T, 2, 4, 2, 32, 64, MlasLinearAttentionRuleDelta, + MlasLinearAttentionDecayNone, MlasLinearAttentionBetaShared, + 1.0f, 1, true); + } + + // Thread partitioning: 8 threads against B*H_kv == 1 exercises the clamp to + // the task count; 3 threads against 8 tasks exercises the remainder split. + for (int threads : {1, 2, 3, 8}) { + Test(1, 9, 1, 1, 1, 64, 64, MlasLinearAttentionRuleGatedDelta, + MlasLinearAttentionDecayPerKeyDim, MlasLinearAttentionBetaPerHead, + 0.125f, threads, true); + Test(2, 9, 8, 4, 2, 32, 64, MlasLinearAttentionRuleGatedDelta, + MlasLinearAttentionDecayPerHead, MlasLinearAttentionBetaPerHead, + 0.125f, threads, true); + } + } +}; + +static UNUSED_VARIABLE bool added_to_main = AddTestRegister([](bool is_short_execute) { + if (is_short_execute) { + return MlasDirectShortExecuteTests::RegisterShortExecute(); + } + return size_t(0); +});