usearch 2.21.3 → 2.25.0

package/numkong/include/numkong/attention/sme.h

@@ -0,0 +1,2075 @@
+/**
+ *  @brief FlashAttention-style kernels for SME.
+ *  @file include/numkong/attention/sme.h
+ *  @author Ash Vardanian
+ *  @date January 11, 2026
+ *
+ *  @sa include/numkong/attention.h
+ *
+ *  This file implements FlashAttention-2 style scaled dot-product attention (SDPA) optimized
+ *  for ARM SME instructions on Apple M4 and similar processors. The kernel computes:
+ *
+ *      O = softmax(Q × Kᵀ / √d) × V
+ *
+ *  Key features:
+ *  - Online softmax: Mathematically exact, processes KV blocks incrementally
+ *  - Pre-packed KV cache: BFMOPA/FMOPA-interleaved format amortizes packing for repeated inference
+ *  - GQA/MQA support: Different `num_heads` and `num_kv_heads` for grouped-query attention
+ *  - Pure Streaming SVE: No NEON intrinsics for non-linear operations
+ *
+ *  Target models (2025):
+ *  - Kimi K2: `head_dim`=112, 64 heads, MHA, 128K context
+ *  - LLaMA 3.1 405B: `head_dim`=128, 128 heads, 16 KV heads (GQA 8:1), 128K context
+ *  - Qwen 2.5 72B: `head_dim`=128, 64 heads, 8 KV heads (GQA 8:1), 32K context
+ *
+ *  @section attention_sme_architecture Architecture
+ *
+ *  Both Q×Kᵀ and P×V phases use BFMOPA/FMOPA outer products on ZA tiles, eliminating
+ *  element-wise SVE loops that dominated the original implementation. The Q matrix is
+ *  pre-transposed once into a buffer matching the interleaving that ZA vertical reads
+ *  would produce, so Q×Kᵀ runs as pure memory-to-BFMOPA with no per-block ZA staging.
+ *
+ *  Block sizes:
+ *  - Bᵣ = 16 (query block rows, matches ZA32 tile height)
+ *  - Bᶜ = 32 (main prefill loop, processes two KV blocks per iteration using ZA2+ZA3)
+ *  - Bᶜ = 16 (tail loop for remaining KV positions, and decode path)
+ *
+ *  KV packing format:
+ *  - K is stored in BFMOPA-interleaved format: `K_packed[kv_block][depth_step][32]` where
+ *    `packed[2*ki + sub] = K[kv_block*16 + ki][2*depth_step + sub]`
+ *  - V is stored in BFMOPA-interleaved format: `V_packed[kv_block][dim_tile][depth_step][32]`
+ *    where `packed[2*dj + sub] = V[kv_block*16 + 2*depth_step + sub][dim_tile*16 + dj]`
+ *  - The `reserved[0]` header field stores `v_dim_tile_count` for efficient V addressing
+ *
+ *  Softmax:
+ *  - Column-wise max and exp using ZA tile vertical reads (avoids per-row horizontal extracts)
+ *  - Correction skip: when the block max does not exceed the running max, the output
+ *    accumulator rescaling is skipped entirely (common in later KV blocks)
+ *  - Degree-3 fast exp (`nk_exp_fast_f32_sve_`) saves 1 FMA per call vs degree-4
+ *  - Weights stored directly as bf16/f16 in ZA0 columns via `svzip1` (no f32 round-trip)
+ *
+ *  Decode path (query_len=1):
+ *  - Uses element-wise SVE with scalar weight broadcasts instead of BFMOPA P×V
+ *  - BFMOPA overhead too high for single-query case due to ZA setup cost
+ *
+ *  P×V prefill path:
+ *  - 4-tile BFMOPA processing: 4 dim-tiles × 8 depth steps per KV block = 32 BFMOPA ops
+ *  - ZA0-ZA3 accumulate simultaneously, read results with MOVA, add to output accumulator
+ *  - Remainder dim-tiles handled 1-at-a-time using ZA0 only
+ *
+ *  SME tile dimensions (for SVL=512, i.e., Apple M4):
+ *  - ZA32 tile: 16 × 16 `f32` elements (1KB)
+ *  - `bf16`/`f16` vectors: 32 elements per SVE vector
+ *
+ *  @section attention_sme_history Optimization History
+ *
+ *  Phase 1 (January 2026): Initial implementation using ZA staging transpose for Q×Kᵀ
+ *  and element-wise SVE for P×V. Q and K rows were loaded into ZA0/ZA1 horizontally,
+ *  read back vertically to produce interleaved vectors for BFMOPA. The P×V phase used
+ *  scalar `svmla_f32_x` loops over head_dim for each query-key pair. Softmax used
+ *  degree-4 polynomial exp with per-row horizontal max/sum. Performance: ~25-50 GFLOP/s
+ *  on Apple M4 (bf16, 8 heads, query_len=64, kv_len=4096, head_dim=128).
+ *
+ *  Phase 2 (February 2026): BFMOPA/FMOPA P×V with pre-packed V in interleaved format.
+ *  Key changes integrated:
+ *  - Q pre-transposed once into a buffer, eliminating per-block ZA staging for Q
+ *  - K pre-packed in interleaved format, enabling pure memory-to-BFMOPA Q×Kᵀ
+ *  - V pre-packed in BFMOPA-interleaved format with dim-tile blocking
+ *  - P×V uses 4-tile BFMOPA accumulation (ZA0-ZA3) with pre-extracted P columns
+ *  - Bᶜ=32 main loop for prefill (2 KV blocks per iteration via ZA2+ZA3)
+ *  - Column-wise softmax: vertical ZA reads for max/exp instead of per-row horizontal
+ *  - Correction skip when running max is unchanged
+ *  - Degree-3 fast exp (~0.5% max relative error, saves 1 FMA per call)
+ *  - Weights stored directly as bf16/f16 via `svzip1` (no f32 quantization round-trip)
+ *  Performance: ~300-400 GFLOP/s on Apple M4 (same configuration), a 6-14× improvement.
+ *
+ *  Rejected approaches:
+ *  - BFMOPA P×V for decode (query_len=1): ZA setup overhead exceeds element-wise SVE cost
+ *  - `svdot_lane` for Q×Kᵀ: lower throughput than BFMOPA on M4
+ *  - Shared ZA tiles between softmax and P×V: register pressure too high with 4-tile P×V
+ */
+#ifndef NK_ATTENTION_SME_H
+#define NK_ATTENTION_SME_H
+
+#if NK_TARGET_ARM64_
+#if NK_TARGET_SME
+
+#include "numkong/types.h"
+
+#if defined(__cplusplus)
+extern "C" {
+#endif
+
+#if defined(__clang__)
+#pragma clang attribute push(__attribute__((target("sme"))), apply_to = function)
+#elif defined(__GNUC__)
+#pragma GCC push_options
+#pragma GCC target("+sme")
+#endif
+
+/**
+ *  @brief Convert bf16 vector to f32 in registers (streaming SVE compatible).
+ *
+ *  BF16 is the upper 16 bits of F32, so we:
+ *  1. Reinterpret bf16 as u16
+ *  2. Zero-extend to u32 (unpklo for lower half)
+ *  3. Shift left by 16 to place in f32 exponent+mantissa position
+ *  4. Reinterpret as f32
+ */
+NK_INTERNAL svfloat32_t nk_bf16_to_f32_sve_(svbool_t predicate_b32x, svbfloat16_t x_bf16x) __arm_streaming {
+    svuint16_t x_u16x = svreinterpret_u16_bf16(x_bf16x);
+    svuint32_t x_u32x = svunpklo_u32(x_u16x);
+    x_u32x = svlsl_n_u32_x(predicate_b32x, x_u32x, 16);
+    return svreinterpret_f32_u32(x_u32x);
+}
+
+/**
+ *  @brief Convert f32 vector to bf16 in registers with rounding (streaming SVE compatible).
+ *
+ *  1. Reinterpret f32 as u32
+ *  2. Add rounding bias (0x8000) for round-to-nearest
+ *  3. Shift right by 16
+ *  4. Narrow to u16 and reinterpret as bf16
+ */
+NK_INTERNAL svbfloat16_t nk_f32_to_bf16_sve_(svbool_t predicate_b32x, svfloat32_t x_f32x) __arm_streaming {
+    svuint32_t x_u32x = svreinterpret_u32_f32(x_f32x);
+    x_u32x = svadd_n_u32_x(predicate_b32x, x_u32x, 0x8000); // Round to nearest
+    x_u32x = svlsr_n_u32_x(predicate_b32x, x_u32x, 16);
+    svuint16_t x_u16x = svuzp1_u16(svreinterpret_u16_u32(x_u32x), svreinterpret_u16_u32(x_u32x));
+    return svreinterpret_bf16_u16(x_u16x);
+}
+
+/**
+ *  @brief Packed KV cache header for attention (64-byte aligned).
+ *
+ *  Layout in memory:
+ *  [header: 64 bytes][K tiles: variable][V tiles: variable]
+ */
+typedef struct {
+    nk_u32_t num_kv_heads;    ///< Number of K/V heads (for GQA, may differ from Q heads)
+    nk_u32_t head_dim;        ///< Original head dimension (64, 112, 128)
+    nk_u32_t head_dim_padded; ///< Padded to multiple of 32 for SME
+    nk_u32_t seq_len;         ///< Current sequence length
+    nk_u32_t max_seq_len;     ///< Maximum sequence length (for pre-allocation)
+    nk_u32_t k_offset;        ///< Byte offset to K data from header start
+    nk_u32_t v_offset;        ///< Byte offset to V data from header start
+    nk_u32_t reserved[9];     ///< reserved[0] = v_dim_tile_count; remainder pads to 64 bytes
+} nk_attention_sme_packed_header_t;
+
+/**
+ *  @brief Fast exp approximation in Streaming SVE.
+ *
+ *  Uses Cody-Waite range reduction + Horner polynomial (degree 4).
+ *  Accuracy: ~0.1% relative error, acceptable for softmax normalization.
+ *
+ *  @param pg Active predicate
+ *  @param x Input vector
+ *  @return exp(x) approximation
+ */
+NK_INTERNAL svfloat32_t nk_exp_f32_sve_(svbool_t predicate_b32x, svfloat32_t x_f32x) __arm_streaming {
+    // Constants for Cody-Waite range reduction
+    svfloat32_t log2e_f32x = svdup_f32(1.4426950408889634f);
+    svfloat32_t ln2_high_f32x = svdup_f32(0.693145751953125f);
+    svfloat32_t ln2_low_f32x = svdup_f32(1.42860682030941723212e-6f);
+
+    // Clamp to avoid overflow/underflow
+    svfloat32_t max_x_f32x = svdup_f32(88.3762626647949f);
+    svfloat32_t min_x_f32x = svdup_f32(-87.3365447504021f);
+    x_f32x = svmax_f32_m(predicate_b32x, svmin_f32_m(predicate_b32x, x_f32x, max_x_f32x), min_x_f32x);
+
+    // n = round(x / ln(2))
+    svfloat32_t n_f32x = svrintn_f32_m(svundef_f32(), predicate_b32x, svmul_f32_m(predicate_b32x, x_f32x, log2e_f32x));
+
+    // r = x - n × ln(2) using Cody-Waite for precision
+    svfloat32_t r_f32x = svmsb_f32_m(predicate_b32x, n_f32x, ln2_high_f32x, x_f32x);
+    r_f32x = svmsb_f32_m(predicate_b32x, n_f32x, ln2_low_f32x, r_f32x);
+
+    // Polynomial approximation for exp(r): degree 4
+    // exp(r) ≈ 1 + r + r²/2 + r³/6 + r⁴/24
+    svfloat32_t p_f32x = svdup_f32(4.1666666667e-2f);                                  // 1/24
+    p_f32x = svmad_f32_m(predicate_b32x, p_f32x, r_f32x, svdup_f32(1.6666666667e-1f)); // 1/6
+    p_f32x = svmad_f32_m(predicate_b32x, p_f32x, r_f32x, svdup_f32(5.0000000000e-1f)); // 1/2
+    p_f32x = svmad_f32_m(predicate_b32x, p_f32x, r_f32x, svdup_f32(1.0f));             // 1
+    p_f32x = svmad_f32_m(predicate_b32x, p_f32x, r_f32x, svdup_f32(1.0f));             // 1
+
+    // Reconstruct: exp(x) = 2ⁿ × exp(r)
+    // 2ⁿ via IEEE 754 exponent manipulation
+    svint32_t n_i32x = svcvt_s32_f32_m(svundef_s32(), predicate_b32x, n_f32x);
+    n_i32x = svadd_s32_m(predicate_b32x, n_i32x, svdup_s32(127));
+    n_i32x = svlsl_n_s32_m(predicate_b32x, n_i32x, 23);
+    svfloat32_t pow2n_f32x = svreinterpret_f32_s32(n_i32x);
+
+    return svmul_f32_m(predicate_b32x, p_f32x, pow2n_f32x);
+}
+
+/**
+ *  @brief Degree-3 fast exp approximation. Max relative error ~0.5%.
+ *  Saves 1 FMA per call vs degree-4 nk_exp_f32_sve_.
+ */
+NK_INTERNAL svfloat32_t nk_exp_fast_f32_sve_(svbool_t predicate_b32x, svfloat32_t x_f32x) __arm_streaming {
+    svfloat32_t log2e_f32x = svdup_f32(1.4426950408889634f);
+    svfloat32_t ln2_high_f32x = svdup_f32(0.693145751953125f);
+    svfloat32_t ln2_low_f32x = svdup_f32(1.42860682030941723212e-6f);
+
+    svfloat32_t max_x_f32x = svdup_f32(88.3762626647949f);
+    svfloat32_t min_x_f32x = svdup_f32(-87.3365447504021f);
+    x_f32x = svmax_f32_m(predicate_b32x, svmin_f32_m(predicate_b32x, x_f32x, max_x_f32x), min_x_f32x);
+
+    svfloat32_t n_f32x = svrintn_f32_m(svundef_f32(), predicate_b32x, svmul_f32_m(predicate_b32x, x_f32x, log2e_f32x));
+    svfloat32_t r_f32x = svmsb_f32_m(predicate_b32x, n_f32x, ln2_high_f32x, x_f32x);
+    r_f32x = svmsb_f32_m(predicate_b32x, n_f32x, ln2_low_f32x, r_f32x);
+
+    // Degree-3: exp(r) ~ 1 + r + r^2/2 + r^3/6 (drop 1/24 term)
+    svfloat32_t p_f32x = svdup_f32(1.6666666667e-1f);                                  // 1/6
+    p_f32x = svmad_f32_m(predicate_b32x, p_f32x, r_f32x, svdup_f32(5.0000000000e-1f)); // 1/2
+    p_f32x = svmad_f32_m(predicate_b32x, p_f32x, r_f32x, svdup_f32(1.0f));             // 1
+    p_f32x = svmad_f32_m(predicate_b32x, p_f32x, r_f32x, svdup_f32(1.0f));             // 1
+
+    svint32_t n_i32x = svcvt_s32_f32_m(svundef_s32(), predicate_b32x, n_f32x);
+    n_i32x = svadd_s32_m(predicate_b32x, n_i32x, svdup_s32(127));
+    n_i32x = svlsl_n_s32_m(predicate_b32x, n_i32x, 23);
+    svfloat32_t pow2n_f32x = svreinterpret_f32_s32(n_i32x);
+
+    return svmul_f32_m(predicate_b32x, p_f32x, pow2n_f32x);
+}
+
+NK_PUBLIC nk_size_t nk_attention_packed_kv_size_bf16_sme(nk_size_t num_kv_heads, nk_size_t head_dim,
+                                                         nk_size_t max_seq_len) {
+    nk_size_t head_dim_padded = (head_dim + 31) / 32 * 32;
+    nk_size_t kv_blocks = (max_seq_len + 15) / 16;
+    nk_size_t seq_padded = kv_blocks * 16;
+    // K and V both use BFMOPA-interleaved format: [num_kv_heads, kv_blocks, depth_steps, 32]
+    nk_size_t k_size = num_kv_heads * seq_padded * head_dim_padded * sizeof(nk_bf16_t);
+    nk_size_t v_size = k_size;
+    return sizeof(nk_attention_sme_packed_header_t) + k_size + v_size;
+}
+
+NK_PUBLIC nk_size_t nk_attention_packed_kv_size_f16_sme(nk_size_t num_kv_heads, nk_size_t head_dim,
+                                                        nk_size_t max_seq_len) {
+    return nk_attention_packed_kv_size_bf16_sme(num_kv_heads, head_dim, max_seq_len);
+}
+
+static void nk_attention_pack_kv_bf16_sme_ssve_( //
+    nk_bf16_t const *k, nk_bf16_t const *v, nk_size_t num_kv_heads, nk_size_t head_dim, nk_size_t seq_len,
+    nk_size_t k_stride, nk_size_t v_stride, void *kv_packed) NK_STREAMING_ {
+
+    nk_attention_sme_packed_header_t *header = (nk_attention_sme_packed_header_t *)kv_packed;
+    nk_size_t head_dim_padded = (head_dim + 31) / 32 * 32;
+    nk_size_t dim_tile_count = (head_dim_padded + 15) / 16;
+    nk_size_t kv_block_count = (seq_len + 15) / 16;
+
+    nk_size_t k_depth_step_count = head_dim_padded / 2;
+    nk_size_t head_elems = kv_block_count * 16 * head_dim_padded;
+
+    header->num_kv_heads = (nk_u32_t)num_kv_heads;
+    header->head_dim = (nk_u32_t)head_dim;
+    header->head_dim_padded = (nk_u32_t)head_dim_padded;
+    header->seq_len = (nk_u32_t)seq_len;
+    header->k_offset = sizeof(nk_attention_sme_packed_header_t);
+    header->reserved[0] = (nk_u32_t)dim_tile_count; // v_dim_tile_count
+    header->v_offset = header->k_offset + (nk_u32_t)(num_kv_heads * head_elems * sizeof(nk_bf16_t));
+
+    nk_bf16_t *k_packed = (nk_bf16_t *)((char *)kv_packed + header->k_offset);
+    nk_bf16_t *v_packed = (nk_bf16_t *)((char *)kv_packed + header->v_offset);
+
+    for (nk_size_t h = 0; h < num_kv_heads; h++) {
+        nk_bf16_t const *k_head = k + h * k_stride;
+        nk_bf16_t const *v_head = v + h * v_stride;
+
+        // K packing: BFMOPA-interleaved format
+        // K_packed[kv_block][depth_step][32] where
+        //   packed[2*ki + sub] = K[kv_block*16 + ki][2*depth_step + sub]
+        nk_bf16_t *k_out = k_packed + h * head_elems;
+        for (nk_size_t kv_block = 0; kv_block < kv_block_count; kv_block++) {
+            for (nk_size_t depth_step = 0; depth_step < k_depth_step_count; depth_step++) {
+                nk_bf16_t *vec_out = k_out + (kv_block * k_depth_step_count + depth_step) * 32;
+                for (nk_size_t ki = 0; ki < 16; ki++) {
+                    for (nk_size_t sub = 0; sub < 2; sub++) {
+                        nk_size_t row = kv_block * 16 + ki;
+                        nk_size_t col = 2 * depth_step + sub;
+                        nk_bf16_t zero = {0};
+                        vec_out[2 * ki + sub] = (row < seq_len && col < head_dim) ? k_head[row * head_dim + col] : zero;
+                    }
+                }
+            }
+        }
+
+        // V packing: BFMOPA-interleaved format
+        nk_bf16_t *v_out = v_packed + h * head_elems;
+        for (nk_size_t kv_block = 0; kv_block < kv_block_count; kv_block++) {
+            for (nk_size_t dim_tile = 0; dim_tile < dim_tile_count; dim_tile++) {
+                for (nk_size_t depth_step = 0; depth_step < 8; depth_step++) {
+                    nk_bf16_t *vec_out = v_out + (kv_block * dim_tile_count * 8 + dim_tile * 8 + depth_step) * 32;
+                    for (nk_size_t dj = 0; dj < 16; dj++) {
+                        for (nk_size_t sub = 0; sub < 2; sub++) {
+                            nk_size_t ki = kv_block * 16 + 2 * depth_step + sub;
+                            nk_size_t d = dim_tile * 16 + dj;
+                            nk_bf16_t zero = {0};
+                            vec_out[2 * dj + sub] = (ki < seq_len && d < head_dim) ? v_head[ki * head_dim + d] : zero;
+                        }
+                    }
+                }
+            }
+        }
+    }
+}
+
+NK_PUBLIC void nk_attention_pack_kv_bf16_sme( //
+    nk_bf16_t const *k, nk_bf16_t const *v, nk_size_t num_kv_heads, nk_size_t head_dim, nk_size_t seq_len,
+    nk_size_t k_stride, nk_size_t v_stride, void *kv_packed) {
+    nk_sme_start_streaming_();
+    nk_attention_pack_kv_bf16_sme_ssve_(k, v, num_kv_heads, head_dim, seq_len, k_stride, v_stride, kv_packed);
+    nk_sme_stop_streaming_();
+}
+
+static void nk_attention_pack_kv_f16_sme_ssve_( //
+    nk_f16_t const *k, nk_f16_t const *v, nk_size_t num_kv_heads, nk_size_t head_dim, nk_size_t seq_len,
+    nk_size_t k_stride, nk_size_t v_stride, void *kv_packed) NK_STREAMING_ {
+
+    nk_attention_sme_packed_header_t *header = (nk_attention_sme_packed_header_t *)kv_packed;
+    nk_size_t head_dim_padded = (head_dim + 31) / 32 * 32;
+    nk_size_t dim_tile_count = (head_dim_padded + 15) / 16;
+    nk_size_t kv_block_count = (seq_len + 15) / 16;
+
+    nk_size_t k_depth_step_count = head_dim_padded / 2;
+    nk_size_t head_elems = kv_block_count * 16 * head_dim_padded;
+
+    header->num_kv_heads = (nk_u32_t)num_kv_heads;
+    header->head_dim = (nk_u32_t)head_dim;
+    header->head_dim_padded = (nk_u32_t)head_dim_padded;
+    header->seq_len = (nk_u32_t)seq_len;
+    header->k_offset = sizeof(nk_attention_sme_packed_header_t);
+    header->reserved[0] = (nk_u32_t)dim_tile_count; // v_dim_tile_count
+    header->v_offset = header->k_offset + (nk_u32_t)(num_kv_heads * head_elems * sizeof(nk_f16_t));
+
+    nk_f16_t *k_packed = (nk_f16_t *)((char *)kv_packed + header->k_offset);
+    nk_f16_t *v_packed = (nk_f16_t *)((char *)kv_packed + header->v_offset);
+
+    for (nk_size_t h = 0; h < num_kv_heads; h++) {
+        nk_f16_t const *k_head = k + h * k_stride;
+        nk_f16_t const *v_head = v + h * v_stride;
+
+        // K packing: FMOPA-interleaved format
+        nk_f16_t *k_out = k_packed + h * head_elems;
+        for (nk_size_t kv_block = 0; kv_block < kv_block_count; kv_block++) {
+            for (nk_size_t depth_step = 0; depth_step < k_depth_step_count; depth_step++) {
+                nk_f16_t *vec_out = k_out + (kv_block * k_depth_step_count + depth_step) * 32;
+                for (nk_size_t ki = 0; ki < 16; ki++) {
+                    for (nk_size_t sub = 0; sub < 2; sub++) {
+                        nk_size_t row = kv_block * 16 + ki;
+                        nk_size_t col = 2 * depth_step + sub;
+                        nk_f16_t zero = {0};
+                        vec_out[2 * ki + sub] = (row < seq_len && col < head_dim) ? k_head[row * head_dim + col] : zero;
+                    }
+                }
+            }
+        }
+
+        // V packing: FMOPA-interleaved format
+        nk_f16_t *v_out = v_packed + h * head_elems;
+        for (nk_size_t kv_block = 0; kv_block < kv_block_count; kv_block++) {
+            for (nk_size_t dim_tile = 0; dim_tile < dim_tile_count; dim_tile++) {
+                for (nk_size_t depth_step = 0; depth_step < 8; depth_step++) {
+                    nk_f16_t *vec_out = v_out + (kv_block * dim_tile_count * 8 + dim_tile * 8 + depth_step) * 32;
+                    for (nk_size_t dj = 0; dj < 16; dj++) {
+                        for (nk_size_t sub = 0; sub < 2; sub++) {
+                            nk_size_t ki = kv_block * 16 + 2 * depth_step + sub;
+                            nk_size_t d = dim_tile * 16 + dj;
+                            nk_f16_t zero = {0};
+                            vec_out[2 * dj + sub] = (ki < seq_len && d < head_dim) ? v_head[ki * head_dim + d] : zero;
+                        }
+                    }
+                }
+            }
+        }
+    }
+}
+
+NK_PUBLIC void nk_attention_pack_kv_f16_sme( //
+    nk_f16_t const *k, nk_f16_t const *v, nk_size_t num_kv_heads, nk_size_t head_dim, nk_size_t seq_len,
+    nk_size_t k_stride, nk_size_t v_stride, void *kv_packed) {
+    nk_sme_start_streaming_();
+    nk_attention_pack_kv_f16_sme_ssve_(k, v, num_kv_heads, head_dim, seq_len, k_stride, v_stride, kv_packed);
+    nk_sme_stop_streaming_();
+}
+
+/**
+ *  @brief Optimized bf16 attention kernel with BFMOPA P×V.
+ *
+ *  Key design choices:
+ *  - P×V uses BFMOPA with pre-packed V (4-tile accumulation) instead of element-wise SVE
+ *  - Scores read via column-wise vertical ZA reads for vectorized max/exp
+ *  - Weights stored directly as bf16 (no f32 round-trip)
+ *  - Uses degree-3 fast exp for softmax
+ *  - Correction skip when running max is unchanged
+ *  - Decode path (valid_query_count==1) remains element-wise SVE (BFMOPA overhead too high)
+ */
+__arm_new("za") static void nk_attention_bf16_sme_streaming_(
+    nk_bf16_t const *q,        // [query_len, head_dim]
+    nk_bf16_t const *k,        // [kv_len, head_dim_padded] BFMOPA-interleaved
+    nk_bf16_t const *v_packed, // BFMOPA-interleaved V for this KV head
+    nk_bf16_t *output,         // [query_len, head_dim]
+    nk_size_t query_len, nk_size_t kv_len, nk_size_t head_dim, nk_size_t head_dim_padded, nk_size_t dim_tile_count,
+    nk_f32_t scale) NK_STREAMING_ {
+
+    svbool_t const predicate_all_b32x = svptrue_b32();
+    svbool_t const predicate_all_b16x = svptrue_b16();
+    nk_size_t const valid_query_count = (query_len < 16) ? query_len : 16;
+
+    svfloat32_t row_max_f32x = svdup_f32(NK_F32_MIN);
+    svfloat32_t row_sum_f32x = svdup_f32(0.0f);
+
+    NK_ALIGN64 nk_f32_t output_accumulator[16 * 256];
+    svfloat32_t zero_f32x = svdup_f32(0.0f);
+    for (nk_size_t i = 0; i < 16 * head_dim_padded; i += svcntw()) {
+        svst1_f32(predicate_all_b32x, output_accumulator + i, zero_f32x);
+    }
+
+    nk_size_t kv_block_index = 0;
+    nk_size_t kv_start = 0;
+    svbool_t const batch_predicate_b32x = svwhilelt_b32(0u, 16u);
+
+    nk_size_t const k_depth_step_count = head_dim_padded / 2;
+
+    // Pre-transpose Q once: queries_transposed[step][16 f32 words]
+    NK_ALIGN64 nk_f32_t queries_transposed[128 * 16]; // max head_dim_padded/2 * 16 = 128 * 16
+    for (nk_size_t batch = 0; batch < head_dim_padded / 32; batch++) {
+        svzero_mask_za(nk_sme_zero_za32_tile_0_);
+        for (nk_size_t query_index = 0; query_index < valid_query_count; query_index++)
+            svld1_hor_za32(0, query_index, batch_predicate_b32x,
+                           (nk_f32_t const *)(q + query_index * head_dim + batch * 32));
+        for (nk_size_t step = 0; step < 16; step++)
+            svst1_f32(predicate_all_b32x, queries_transposed + (batch * 16 + step) * 16,
+                      svread_ver_za32_f32_m(svdup_f32(0), predicate_all_b32x, 0, step));
+    }
+
+    // Bc=32 main loop (prefill only, skipped for decode)
+    if (valid_query_count > 1) {
+        for (; kv_start + 32 <= kv_len; kv_start += 32, kv_block_index += 2) {
+            // Q×K^T: pure memory→BFMOPA, no ZA staging for Q or K
+            svzero_mask_za(nk_sme_zero_za32_tile_2_);
+            svzero_mask_za(nk_sme_zero_za32_tile_3_);
+            nk_bf16_t const *keys_block_low = k + kv_block_index * k_depth_step_count * 32;
+            nk_bf16_t const *keys_block_high = k + (kv_block_index + 1) * k_depth_step_count * 32;
+            for (nk_size_t step = 0; step < k_depth_step_count; step++) {
+                svbfloat16_t zn_bf16x = svreinterpret_bf16_f32(
+                    svld1_f32(predicate_all_b32x, queries_transposed + step * 16));
+                svbfloat16_t zm0_bf16x = svld1_bf16(predicate_all_b16x,
+                                                    (bfloat16_t const *)(keys_block_low + step * 32));
+                svbfloat16_t zm1_bf16x = svld1_bf16(predicate_all_b16x,
+                                                    (bfloat16_t const *)(keys_block_high + step * 32));
+                svmopa_za32_bf16_m(2, predicate_all_b32x, predicate_all_b32x, zn_bf16x, zm0_bf16x);
+                svmopa_za32_bf16_m(3, predicate_all_b32x, predicate_all_b32x, zn_bf16x, zm1_bf16x);
+            }
+
+            // Pass 1: Column-wise max (read ZA2/ZA3 columns vertically)
+            svfloat32_t scale_f32x = svdup_f32(scale);
+            svfloat32_t block_max_f32x = svdup_f32(NK_F32_MIN);
+            for (nk_size_t column_index = 0; column_index < 16; column_index++) {
+                svfloat32_t score_column_f32x = svmul_f32_x(
+                    predicate_all_b32x, svread_ver_za32_f32_m(svdup_f32(0), predicate_all_b32x, 2, column_index),
+                    scale_f32x);
+                block_max_f32x = svmax_f32_x(predicate_all_b32x, block_max_f32x, score_column_f32x);
+            }
+            for (nk_size_t column_index = 0; column_index < 16; column_index++) {
+                svfloat32_t score_column_f32x = svmul_f32_x(
+                    predicate_all_b32x, svread_ver_za32_f32_m(svdup_f32(0), predicate_all_b32x, 3, column_index),
+                    scale_f32x);
+                block_max_f32x = svmax_f32_x(predicate_all_b32x, block_max_f32x, score_column_f32x);
+            }
+
+            // Softmax correction (fully vectorized)
+            svfloat32_t new_max_f32x = svmax_f32_x(predicate_all_b32x, row_max_f32x, block_max_f32x);
+            svfloat32_t correction_f32x = nk_exp_fast_f32_sve_(
+                predicate_all_b32x, svsub_f32_x(predicate_all_b32x, row_max_f32x, new_max_f32x));
+            svbool_t max_changed_b32x = svcmplt_f32(predicate_all_b32x, correction_f32x, svdup_f32(1.0f));
+            nk_u32_t max_was_updated = svptest_any(predicate_all_b32x, max_changed_b32x) ? 1 : 0;
+            if (max_was_updated) row_sum_f32x = svmul_f32_x(predicate_all_b32x, row_sum_f32x, correction_f32x);
+            NK_ALIGN64 nk_f32_t corrections[16];
+            svst1_f32(predicate_all_b32x, corrections, correction_f32x);
+
+            // Pass 2: Column-wise exp + fused P write + sum
+            svfloat32_t sum_delta_f32x = svdup_f32(0.0f);
+            svzero_mask_za(nk_sme_zero_za32_tile_0_);
+            // ZA2 columns in pairs → ZA0 columns 0-7
+            for (nk_size_t column_index = 0; column_index < 16; column_index += 2) {
+                svfloat32_t score_even_f32x = svmul_f32_x(
+                    predicate_all_b32x, svread_ver_za32_f32_m(svdup_f32(0), predicate_all_b32x, 2, column_index),
+                    scale_f32x);
+                svfloat32_t score_odd_f32x = svmul_f32_x(
+                    predicate_all_b32x, svread_ver_za32_f32_m(svdup_f32(0), predicate_all_b32x, 2, column_index + 1),
+                    scale_f32x);
+                svfloat32_t weight_even_f32x = nk_exp_fast_f32_sve_(
+                    predicate_all_b32x, svsub_f32_x(predicate_all_b32x, score_even_f32x, new_max_f32x));
+                svfloat32_t weight_odd_f32x = nk_exp_fast_f32_sve_(
+                    predicate_all_b32x, svsub_f32_x(predicate_all_b32x, score_odd_f32x, new_max_f32x));
+                sum_delta_f32x = svadd_f32_x(predicate_all_b32x, sum_delta_f32x, weight_even_f32x);
+                sum_delta_f32x = svadd_f32_x(predicate_all_b32x, sum_delta_f32x, weight_odd_f32x);
+                svbfloat16_t weight_pair_bf16x = svzip1_bf16(nk_f32_to_bf16_sve_(predicate_all_b32x, weight_even_f32x),
+                                                             nk_f32_to_bf16_sve_(predicate_all_b32x, weight_odd_f32x));
+                svwrite_ver_za32_f32_m(0, column_index / 2, predicate_all_b32x,
+                                       svreinterpret_f32_bf16(weight_pair_bf16x));
+            }
+            // ZA3 columns in pairs → ZA0 columns 8-15
+            for (nk_size_t column_index = 0; column_index < 16; column_index += 2) {
+                svfloat32_t score_even_f32x = svmul_f32_x(
+                    predicate_all_b32x, svread_ver_za32_f32_m(svdup_f32(0), predicate_all_b32x, 3, column_index),
+                    scale_f32x);
+                svfloat32_t score_odd_f32x = svmul_f32_x(
+                    predicate_all_b32x, svread_ver_za32_f32_m(svdup_f32(0), predicate_all_b32x, 3, column_index + 1),
+                    scale_f32x);
+                svfloat32_t weight_even_f32x = nk_exp_fast_f32_sve_(
+                    predicate_all_b32x, svsub_f32_x(predicate_all_b32x, score_even_f32x, new_max_f32x));
+                svfloat32_t weight_odd_f32x = nk_exp_fast_f32_sve_(
+                    predicate_all_b32x, svsub_f32_x(predicate_all_b32x, score_odd_f32x, new_max_f32x));
+                sum_delta_f32x = svadd_f32_x(predicate_all_b32x, sum_delta_f32x, weight_even_f32x);
+                sum_delta_f32x = svadd_f32_x(predicate_all_b32x, sum_delta_f32x, weight_odd_f32x);
+                svbfloat16_t weight_pair_bf16x = svzip1_bf16(nk_f32_to_bf16_sve_(predicate_all_b32x, weight_even_f32x),
+                                                             nk_f32_to_bf16_sve_(predicate_all_b32x, weight_odd_f32x));
+                svwrite_ver_za32_f32_m(0, 8 + column_index / 2, predicate_all_b32x,
+                                       svreinterpret_f32_bf16(weight_pair_bf16x));
+            }
+            row_sum_f32x = svadd_f32_x(predicate_all_b32x, row_sum_f32x, sum_delta_f32x);
+            row_max_f32x = new_max_f32x;
+
+            // Extract P columns from ZA0
+            svbfloat16_t probability_column_0_f32x = svreinterpret_bf16_f32(
+                svread_ver_za32_f32_m(svdup_f32(0), predicate_all_b32x, 0, 0));
+            svbfloat16_t probability_column_1_f32x = svreinterpret_bf16_f32(
+                svread_ver_za32_f32_m(svdup_f32(0), predicate_all_b32x, 0, 1));
+            svbfloat16_t probability_column_2_f32x = svreinterpret_bf16_f32(
+                svread_ver_za32_f32_m(svdup_f32(0), predicate_all_b32x, 0, 2));
+            svbfloat16_t probability_column_3_f32x = svreinterpret_bf16_f32(
+                svread_ver_za32_f32_m(svdup_f32(0), predicate_all_b32x, 0, 3));
+            svbfloat16_t probability_column_4_f32x = svreinterpret_bf16_f32(
+                svread_ver_za32_f32_m(svdup_f32(0), predicate_all_b32x, 0, 4));
+            svbfloat16_t probability_column_5_f32x = svreinterpret_bf16_f32(
+                svread_ver_za32_f32_m(svdup_f32(0), predicate_all_b32x, 0, 5));
+            svbfloat16_t probability_column_6_f32x = svreinterpret_bf16_f32(
+                svread_ver_za32_f32_m(svdup_f32(0), predicate_all_b32x, 0, 6));
+            svbfloat16_t probability_column_7_f32x = svreinterpret_bf16_f32(
+                svread_ver_za32_f32_m(svdup_f32(0), predicate_all_b32x, 0, 7));
+            svbfloat16_t probability_column_8_f32x = svreinterpret_bf16_f32(
+                svread_ver_za32_f32_m(svdup_f32(0), predicate_all_b32x, 0, 8));
+            svbfloat16_t probability_column_9_f32x = svreinterpret_bf16_f32(
+                svread_ver_za32_f32_m(svdup_f32(0), predicate_all_b32x, 0, 9));
+            svbfloat16_t probability_column_10_f32x = svreinterpret_bf16_f32(
+                svread_ver_za32_f32_m(svdup_f32(0), predicate_all_b32x, 0, 10));
+            svbfloat16_t probability_column_11_f32x = svreinterpret_bf16_f32(
+                svread_ver_za32_f32_m(svdup_f32(0), predicate_all_b32x, 0, 11));
+            svbfloat16_t probability_column_12_f32x = svreinterpret_bf16_f32(
+                svread_ver_za32_f32_m(svdup_f32(0), predicate_all_b32x, 0, 12));
+            svbfloat16_t probability_column_13_f32x = svreinterpret_bf16_f32(
+                svread_ver_za32_f32_m(svdup_f32(0), predicate_all_b32x, 0, 13));
+            svbfloat16_t probability_column_14_f32x = svreinterpret_bf16_f32(
+                svread_ver_za32_f32_m(svdup_f32(0), predicate_all_b32x, 0, 14));
+            svbfloat16_t probability_column_15_f32x = svreinterpret_bf16_f32(
+                svread_ver_za32_f32_m(svdup_f32(0), predicate_all_b32x, 0, 15));
+
+            // Pre-apply correction once before P×V
+            svbool_t query_predicate_b16x = svwhilelt_b16_u64(0u, valid_query_count * 2);
+            nk_bf16_t const *values_block_low = v_packed + kv_block_index * dim_tile_count * 8 * 32;
+            nk_bf16_t const *values_block_high = v_packed + (kv_block_index + 1) * dim_tile_count * 8 * 32;
+
+            if (max_was_updated) {
+                for (nk_size_t query_index = 0; query_index < valid_query_count; query_index++) {
+                    svfloat32_t correction_scalar_f32x = svdup_f32(corrections[query_index]);
+                    for (nk_size_t dim_offset = 0; dim_offset < head_dim_padded; dim_offset += 16)
+                        svst1_f32(
+                            predicate_all_b32x, output_accumulator + query_index * head_dim_padded + dim_offset,
+                            svmul_f32_x(predicate_all_b32x,
+                                        svld1_f32(predicate_all_b32x,
+                                                  output_accumulator + query_index * head_dim_padded + dim_offset),
+                                        correction_scalar_f32x));
+                }
+            }
+
+            // P×V: zero → BFMOPA → read → add (no ZA writes for output_accumulator)
+            nk_size_t dim_tile = 0;
+            for (; dim_tile + 4 <= dim_tile_count; dim_tile += 4) {
+                svzero_za();
+                // Block0: 8 depth steps (KV positions 0-15)
+                svmopa_za32_bf16_m(0, query_predicate_b16x, predicate_all_b16x, probability_column_0_f32x,
+                                   svld1_bf16(predicate_all_b16x,
+                                              (bfloat16_t const *)(values_block_low + ((dim_tile + 0) * 8 + 0) * 32)));
+                svmopa_za32_bf16_m(1, query_predicate_b16x, predicate_all_b16x, probability_column_0_f32x,
+                                   svld1_bf16(predicate_all_b16x,
+                                              (bfloat16_t const *)(values_block_low + ((dim_tile + 1) * 8 + 0) * 32)));
+                svmopa_za32_bf16_m(2, query_predicate_b16x, predicate_all_b16x, probability_column_0_f32x,
+                                   svld1_bf16(predicate_all_b16x,
+                                              (bfloat16_t const *)(values_block_low + ((dim_tile + 2) * 8 + 0) * 32)));
+                svmopa_za32_bf16_m(3, query_predicate_b16x, predicate_all_b16x, probability_column_0_f32x,
+                                   svld1_bf16(predicate_all_b16x,
+                                              (bfloat16_t const *)(values_block_low + ((dim_tile + 3) * 8 + 0) * 32)));
+                svmopa_za32_bf16_m(0, query_predicate_b16x, predicate_all_b16x, probability_column_1_f32x,
+                                   svld1_bf16(predicate_all_b16x,
+                                              (bfloat16_t const *)(values_block_low + ((dim_tile + 0) * 8 + 1) * 32)));
+                svmopa_za32_bf16_m(1, query_predicate_b16x, predicate_all_b16x, probability_column_1_f32x,
+                                   svld1_bf16(predicate_all_b16x,
+                                              (bfloat16_t const *)(values_block_low + ((dim_tile + 1) * 8 + 1) * 32)));
+                svmopa_za32_bf16_m(2, query_predicate_b16x, predicate_all_b16x, probability_column_1_f32x,
+                                   svld1_bf16(predicate_all_b16x,
+                                              (bfloat16_t const *)(values_block_low + ((dim_tile + 2) * 8 + 1) * 32)));
+                svmopa_za32_bf16_m(3, query_predicate_b16x, predicate_all_b16x, probability_column_1_f32x,
+                                   svld1_bf16(predicate_all_b16x,
+                                              (bfloat16_t const *)(values_block_low + ((dim_tile + 3) * 8 + 1) * 32)));
+                svmopa_za32_bf16_m(0, query_predicate_b16x, predicate_all_b16x, probability_column_2_f32x,
+                                   svld1_bf16(predicate_all_b16x,
+                                              (bfloat16_t const *)(values_block_low + ((dim_tile + 0) * 8 + 2) * 32)));
+                svmopa_za32_bf16_m(1, query_predicate_b16x, predicate_all_b16x, probability_column_2_f32x,
+                                   svld1_bf16(predicate_all_b16x,
+                                              (bfloat16_t const *)(values_block_low + ((dim_tile + 1) * 8 + 2) * 32)));
+                svmopa_za32_bf16_m(2, query_predicate_b16x, predicate_all_b16x, probability_column_2_f32x,
+                                   svld1_bf16(predicate_all_b16x,
+                                              (bfloat16_t const *)(values_block_low + ((dim_tile + 2) * 8 + 2) * 32)));
+                svmopa_za32_bf16_m(3, query_predicate_b16x, predicate_all_b16x, probability_column_2_f32x,
+                                   svld1_bf16(predicate_all_b16x,
+                                              (bfloat16_t const *)(values_block_low + ((dim_tile + 3) * 8 + 2) * 32)));
+                svmopa_za32_bf16_m(0, query_predicate_b16x, predicate_all_b16x, probability_column_3_f32x,
+                                   svld1_bf16(predicate_all_b16x,
+                                              (bfloat16_t const *)(values_block_low + ((dim_tile + 0) * 8 + 3) * 32)));
+                svmopa_za32_bf16_m(1, query_predicate_b16x, predicate_all_b16x, probability_column_3_f32x,
+                                   svld1_bf16(predicate_all_b16x,
+                                              (bfloat16_t const *)(values_block_low + ((dim_tile + 1) * 8 + 3) * 32)));
+                svmopa_za32_bf16_m(2, query_predicate_b16x, predicate_all_b16x, probability_column_3_f32x,
+                                   svld1_bf16(predicate_all_b16x,
+                                              (bfloat16_t const *)(values_block_low + ((dim_tile + 2) * 8 + 3) * 32)));
+                svmopa_za32_bf16_m(3, query_predicate_b16x, predicate_all_b16x, probability_column_3_f32x,
+                                   svld1_bf16(predicate_all_b16x,
+                                              (bfloat16_t const *)(values_block_low + ((dim_tile + 3) * 8 + 3) * 32)));
+                svmopa_za32_bf16_m(0, query_predicate_b16x, predicate_all_b16x, probability_column_4_f32x,
+                                   svld1_bf16(predicate_all_b16x,
+                                              (bfloat16_t const *)(values_block_low + ((dim_tile + 0) * 8 + 4) * 32)));
+                svmopa_za32_bf16_m(1, query_predicate_b16x, predicate_all_b16x, probability_column_4_f32x,
+                                   svld1_bf16(predicate_all_b16x,
+                                              (bfloat16_t const *)(values_block_low + ((dim_tile + 1) * 8 + 4) * 32)));
+                svmopa_za32_bf16_m(2, query_predicate_b16x, predicate_all_b16x, probability_column_4_f32x,
+                                   svld1_bf16(predicate_all_b16x,
+                                              (bfloat16_t const *)(values_block_low + ((dim_tile + 2) * 8 + 4) * 32)));
+                svmopa_za32_bf16_m(3, query_predicate_b16x, predicate_all_b16x, probability_column_4_f32x,
+                                   svld1_bf16(predicate_all_b16x,
+                                              (bfloat16_t const *)(values_block_low + ((dim_tile + 3) * 8 + 4) * 32)));
+                svmopa_za32_bf16_m(0, query_predicate_b16x, predicate_all_b16x, probability_column_5_f32x,
+                                   svld1_bf16(predicate_all_b16x,
+                                              (bfloat16_t const *)(values_block_low + ((dim_tile + 0) * 8 + 5) * 32)));
+                svmopa_za32_bf16_m(1, query_predicate_b16x, predicate_all_b16x, probability_column_5_f32x,
+                                   svld1_bf16(predicate_all_b16x,
+                                              (bfloat16_t const *)(values_block_low + ((dim_tile + 1) * 8 + 5) * 32)));
+                svmopa_za32_bf16_m(2, query_predicate_b16x, predicate_all_b16x, probability_column_5_f32x,
+                                   svld1_bf16(predicate_all_b16x,
+                                              (bfloat16_t const *)(values_block_low + ((dim_tile + 2) * 8 + 5) * 32)));
+                svmopa_za32_bf16_m(3, query_predicate_b16x, predicate_all_b16x, probability_column_5_f32x,
+                                   svld1_bf16(predicate_all_b16x,
+                                              (bfloat16_t const *)(values_block_low + ((dim_tile + 3) * 8 + 5) * 32)));
+                svmopa_za32_bf16_m(0, query_predicate_b16x, predicate_all_b16x, probability_column_6_f32x,
+                                   svld1_bf16(predicate_all_b16x,
+                                              (bfloat16_t const *)(values_block_low + ((dim_tile + 0) * 8 + 6) * 32)));
+                svmopa_za32_bf16_m(1, query_predicate_b16x, predicate_all_b16x, probability_column_6_f32x,
+                                   svld1_bf16(predicate_all_b16x,
+                                              (bfloat16_t const *)(values_block_low + ((dim_tile + 1) * 8 + 6) * 32)));
+                svmopa_za32_bf16_m(2, query_predicate_b16x, predicate_all_b16x, probability_column_6_f32x,
+                                   svld1_bf16(predicate_all_b16x,
+                                              (bfloat16_t const *)(values_block_low + ((dim_tile + 2) * 8 + 6) * 32)));
+                svmopa_za32_bf16_m(3, query_predicate_b16x, predicate_all_b16x, probability_column_6_f32x,
+                                   svld1_bf16(predicate_all_b16x,
+                                              (bfloat16_t const *)(values_block_low + ((dim_tile + 3) * 8 + 6) * 32)));
+                svmopa_za32_bf16_m(0, query_predicate_b16x, predicate_all_b16x, probability_column_7_f32x,
+                                   svld1_bf16(predicate_all_b16x,
+                                              (bfloat16_t const *)(values_block_low + ((dim_tile + 0) * 8 + 7) * 32)));
+                svmopa_za32_bf16_m(1, query_predicate_b16x, predicate_all_b16x, probability_column_7_f32x,
+                                   svld1_bf16(predicate_all_b16x,
+                                              (bfloat16_t const *)(values_block_low + ((dim_tile + 1) * 8 + 7) * 32)));
+                svmopa_za32_bf16_m(2, query_predicate_b16x, predicate_all_b16x, probability_column_7_f32x,
+                                   svld1_bf16(predicate_all_b16x,
+                                              (bfloat16_t const *)(values_block_low + ((dim_tile + 2) * 8 + 7) * 32)));
+                svmopa_za32_bf16_m(3, query_predicate_b16x, predicate_all_b16x, probability_column_7_f32x,
+                                   svld1_bf16(predicate_all_b16x,
+                                              (bfloat16_t const *)(values_block_low + ((dim_tile + 3) * 8 + 7) * 32)));
+                // Block1: 8 depth steps (KV positions 16-31)
+                svmopa_za32_bf16_m(0, query_predicate_b16x, predicate_all_b16x, probability_column_8_f32x,
+                                   svld1_bf16(predicate_all_b16x,
+                                              (bfloat16_t const *)(values_block_high + ((dim_tile + 0) * 8 + 0) * 32)));
+                svmopa_za32_bf16_m(1, query_predicate_b16x, predicate_all_b16x, probability_column_8_f32x,
+                                   svld1_bf16(predicate_all_b16x,
+                                              (bfloat16_t const *)(values_block_high + ((dim_tile + 1) * 8 + 0) * 32)));
+                svmopa_za32_bf16_m(2, query_predicate_b16x, predicate_all_b16x, probability_column_8_f32x,
+                                   svld1_bf16(predicate_all_b16x,
+                                              (bfloat16_t const *)(values_block_high + ((dim_tile + 2) * 8 + 0) * 32)));
+                svmopa_za32_bf16_m(3, query_predicate_b16x, predicate_all_b16x, probability_column_8_f32x,
+                                   svld1_bf16(predicate_all_b16x,
+                                              (bfloat16_t const *)(values_block_high + ((dim_tile + 3) * 8 + 0) * 32)));
+                svmopa_za32_bf16_m(0, query_predicate_b16x, predicate_all_b16x, probability_column_9_f32x,
+                                   svld1_bf16(predicate_all_b16x,
+                                              (bfloat16_t const *)(values_block_high + ((dim_tile + 0) * 8 + 1) * 32)));
+                svmopa_za32_bf16_m(1, query_predicate_b16x, predicate_all_b16x, probability_column_9_f32x,
+                                   svld1_bf16(predicate_all_b16x,
+                                              (bfloat16_t const *)(values_block_high + ((dim_tile + 1) * 8 + 1) * 32)));
+                svmopa_za32_bf16_m(2, query_predicate_b16x, predicate_all_b16x, probability_column_9_f32x,
+                                   svld1_bf16(predicate_all_b16x,
+                                              (bfloat16_t const *)(values_block_high + ((dim_tile + 2) * 8 + 1) * 32)));
+                svmopa_za32_bf16_m(3, query_predicate_b16x, predicate_all_b16x, probability_column_9_f32x,
+                                   svld1_bf16(predicate_all_b16x,
+                                              (bfloat16_t const *)(values_block_high + ((dim_tile + 3) * 8 + 1) * 32)));
+                svmopa_za32_bf16_m(0, query_predicate_b16x, predicate_all_b16x, probability_column_10_f32x,
+                                   svld1_bf16(predicate_all_b16x,
+                                              (bfloat16_t const *)(values_block_high + ((dim_tile + 0) * 8 + 2) * 32)));
+                svmopa_za32_bf16_m(1, query_predicate_b16x, predicate_all_b16x, probability_column_10_f32x,
+                                   svld1_bf16(predicate_all_b16x,
+                                              (bfloat16_t const *)(values_block_high + ((dim_tile + 1) * 8 + 2) * 32)));
+                svmopa_za32_bf16_m(2, query_predicate_b16x, predicate_all_b16x, probability_column_10_f32x,
+                                   svld1_bf16(predicate_all_b16x,
+                                              (bfloat16_t const *)(values_block_high + ((dim_tile + 2) * 8 + 2) * 32)));
+                svmopa_za32_bf16_m(3, query_predicate_b16x, predicate_all_b16x, probability_column_10_f32x,
+                                   svld1_bf16(predicate_all_b16x,
+                                              (bfloat16_t const *)(values_block_high + ((dim_tile + 3) * 8 + 2) * 32)));
+                svmopa_za32_bf16_m(0, query_predicate_b16x, predicate_all_b16x, probability_column_11_f32x,
+                                   svld1_bf16(predicate_all_b16x,
+                                              (bfloat16_t const *)(values_block_high + ((dim_tile + 0) * 8 + 3) * 32)));
+                svmopa_za32_bf16_m(1, query_predicate_b16x, predicate_all_b16x, probability_column_11_f32x,
+                                   svld1_bf16(predicate_all_b16x,
+                                              (bfloat16_t const *)(values_block_high + ((dim_tile + 1) * 8 + 3) * 32)));
+                svmopa_za32_bf16_m(2, query_predicate_b16x, predicate_all_b16x, probability_column_11_f32x,
+                                   svld1_bf16(predicate_all_b16x,
+                                              (bfloat16_t const *)(values_block_high + ((dim_tile + 2) * 8 + 3) * 32)));
+                svmopa_za32_bf16_m(3, query_predicate_b16x, predicate_all_b16x, probability_column_11_f32x,
+                                   svld1_bf16(predicate_all_b16x,
+                                              (bfloat16_t const *)(values_block_high + ((dim_tile + 3) * 8 + 3) * 32)));
+                svmopa_za32_bf16_m(0, query_predicate_b16x, predicate_all_b16x, probability_column_12_f32x,
+                                   svld1_bf16(predicate_all_b16x,
+                                              (bfloat16_t const *)(values_block_high + ((dim_tile + 0) * 8 + 4) * 32)));
+                svmopa_za32_bf16_m(1, query_predicate_b16x, predicate_all_b16x, probability_column_12_f32x,
+                                   svld1_bf16(predicate_all_b16x,
+                                              (bfloat16_t const *)(values_block_high + ((dim_tile + 1) * 8 + 4) * 32)));
+                svmopa_za32_bf16_m(2, query_predicate_b16x, predicate_all_b16x, probability_column_12_f32x,
+                                   svld1_bf16(predicate_all_b16x,
+                                              (bfloat16_t const *)(values_block_high + ((dim_tile + 2) * 8 + 4) * 32)));
+                svmopa_za32_bf16_m(3, query_predicate_b16x, predicate_all_b16x, probability_column_12_f32x,
+                                   svld1_bf16(predicate_all_b16x,
+                                              (bfloat16_t const *)(values_block_high + ((dim_tile + 3) * 8 + 4) * 32)));
+                svmopa_za32_bf16_m(0, query_predicate_b16x, predicate_all_b16x, probability_column_13_f32x,
+                                   svld1_bf16(predicate_all_b16x,
+                                              (bfloat16_t const *)(values_block_high + ((dim_tile + 0) * 8 + 5) * 32)));
+                svmopa_za32_bf16_m(1, query_predicate_b16x, predicate_all_b16x, probability_column_13_f32x,
+                                   svld1_bf16(predicate_all_b16x,
+                                              (bfloat16_t const *)(values_block_high + ((dim_tile + 1) * 8 + 5) * 32)));
+                svmopa_za32_bf16_m(2, query_predicate_b16x, predicate_all_b16x, probability_column_13_f32x,
+                                   svld1_bf16(predicate_all_b16x,
+                                              (bfloat16_t const *)(values_block_high + ((dim_tile + 2) * 8 + 5) * 32)));
+                svmopa_za32_bf16_m(3, query_predicate_b16x, predicate_all_b16x, probability_column_13_f32x,
+                                   svld1_bf16(predicate_all_b16x,
+                                              (bfloat16_t const *)(values_block_high + ((dim_tile + 3) * 8 + 5) * 32)));
+                svmopa_za32_bf16_m(0, query_predicate_b16x, predicate_all_b16x, probability_column_14_f32x,
+                                   svld1_bf16(predicate_all_b16x,
+                                              (bfloat16_t const *)(values_block_high + ((dim_tile + 0) * 8 + 6) * 32)));
+                svmopa_za32_bf16_m(1, query_predicate_b16x, predicate_all_b16x, probability_column_14_f32x,
+                                   svld1_bf16(predicate_all_b16x,
+                                              (bfloat16_t const *)(values_block_high + ((dim_tile + 1) * 8 + 6) * 32)));
+                svmopa_za32_bf16_m(2, query_predicate_b16x, predicate_all_b16x, probability_column_14_f32x,
+                                   svld1_bf16(predicate_all_b16x,
+                                              (bfloat16_t const *)(values_block_high + ((dim_tile + 2) * 8 + 6) * 32)));
+                svmopa_za32_bf16_m(3, query_predicate_b16x, predicate_all_b16x, probability_column_14_f32x,
+                                   svld1_bf16(predicate_all_b16x,
+                                              (bfloat16_t const *)(values_block_high + ((dim_tile + 3) * 8 + 6) * 32)));
+                svmopa_za32_bf16_m(0, query_predicate_b16x, predicate_all_b16x, probability_column_15_f32x,
+                                   svld1_bf16(predicate_all_b16x,
+                                              (bfloat16_t const *)(values_block_high + ((dim_tile + 0) * 8 + 7) * 32)));
+                svmopa_za32_bf16_m(1, query_predicate_b16x, predicate_all_b16x, probability_column_15_f32x,
+                                   svld1_bf16(predicate_all_b16x,
+                                              (bfloat16_t const *)(values_block_high + ((dim_tile + 1) * 8 + 7) * 32)));
+                svmopa_za32_bf16_m(2, query_predicate_b16x, predicate_all_b16x, probability_column_15_f32x,
+                                   svld1_bf16(predicate_all_b16x,
+                                              (bfloat16_t const *)(values_block_high + ((dim_tile + 2) * 8 + 7) * 32)));
+                svmopa_za32_bf16_m(3, query_predicate_b16x, predicate_all_b16x, probability_column_15_f32x,
+                                   svld1_bf16(predicate_all_b16x,
+                                              (bfloat16_t const *)(values_block_high + ((dim_tile + 3) * 8 + 7) * 32)));
+                // Read BFMOPA result and ADD to output_accumulator
+                for (nk_size_t query_index = 0; query_index < valid_query_count; query_index++) {
+                    svst1_f32(
+                        predicate_all_b32x, output_accumulator + query_index * head_dim_padded + (dim_tile + 0) * 16,
+                        svadd_f32_x(predicate_all_b32x,
+                                    svld1_f32(predicate_all_b32x,
+                                              output_accumulator + query_index * head_dim_padded + (dim_tile + 0) * 16),
+                                    svread_hor_za32_f32_m(svdup_f32(0), predicate_all_b32x, 0, query_index)));
+                    svst1_f32(
+                        predicate_all_b32x, output_accumulator + query_index * head_dim_padded + (dim_tile + 1) * 16,
+                        svadd_f32_x(predicate_all_b32x,
+                                    svld1_f32(predicate_all_b32x,
+                                              output_accumulator + query_index * head_dim_padded + (dim_tile + 1) * 16),
+                                    svread_hor_za32_f32_m(svdup_f32(0), predicate_all_b32x, 1, query_index)));
+                    svst1_f32(
+                        predicate_all_b32x, output_accumulator + query_index * head_dim_padded + (dim_tile + 2) * 16,
+                        svadd_f32_x(predicate_all_b32x,
+                                    svld1_f32(predicate_all_b32x,
+                                              output_accumulator + query_index * head_dim_padded + (dim_tile + 2) * 16),
+                                    svread_hor_za32_f32_m(svdup_f32(0), predicate_all_b32x, 2, query_index)));
+                    svst1_f32(
+                        predicate_all_b32x, output_accumulator + query_index * head_dim_padded + (dim_tile + 3) * 16,
+                        svadd_f32_x(predicate_all_b32x,
+                                    svld1_f32(predicate_all_b32x,
+                                              output_accumulator + query_index * head_dim_padded + (dim_tile + 3) * 16),
+                                    svread_hor_za32_f32_m(svdup_f32(0), predicate_all_b32x, 3, query_index)));
+                }
+            }
+            // Remainder: 1 dim_tile at a time using ZA0
+            for (; dim_tile < dim_tile_count; dim_tile++) {
+                svzero_mask_za(nk_sme_zero_za32_tile_0_);
+                svmopa_za32_bf16_m(
+                    0, query_predicate_b16x, predicate_all_b16x, probability_column_0_f32x,
+                    svld1_bf16(predicate_all_b16x, (bfloat16_t const *)(values_block_low + (dim_tile * 8 + 0) * 32)));
+                svmopa_za32_bf16_m(
+                    0, query_predicate_b16x, predicate_all_b16x, probability_column_1_f32x,
+                    svld1_bf16(predicate_all_b16x, (bfloat16_t const *)(values_block_low + (dim_tile * 8 + 1) * 32)));
+                svmopa_za32_bf16_m(
+                    0, query_predicate_b16x, predicate_all_b16x, probability_column_2_f32x,
+                    svld1_bf16(predicate_all_b16x, (bfloat16_t const *)(values_block_low + (dim_tile * 8 + 2) * 32)));
+                svmopa_za32_bf16_m(
+                    0, query_predicate_b16x, predicate_all_b16x, probability_column_3_f32x,
+                    svld1_bf16(predicate_all_b16x, (bfloat16_t const *)(values_block_low + (dim_tile * 8 + 3) * 32)));
+                svmopa_za32_bf16_m(
+                    0, query_predicate_b16x, predicate_all_b16x, probability_column_4_f32x,
+                    svld1_bf16(predicate_all_b16x, (bfloat16_t const *)(values_block_low + (dim_tile * 8 + 4) * 32)));
+                svmopa_za32_bf16_m(
+                    0, query_predicate_b16x, predicate_all_b16x, probability_column_5_f32x,
+                    svld1_bf16(predicate_all_b16x, (bfloat16_t const *)(values_block_low + (dim_tile * 8 + 5) * 32)));
+                svmopa_za32_bf16_m(
+                    0, query_predicate_b16x, predicate_all_b16x, probability_column_6_f32x,
+                    svld1_bf16(predicate_all_b16x, (bfloat16_t const *)(values_block_low + (dim_tile * 8 + 6) * 32)));
+                svmopa_za32_bf16_m(
+                    0, query_predicate_b16x, predicate_all_b16x, probability_column_7_f32x,
+                    svld1_bf16(predicate_all_b16x, (bfloat16_t const *)(values_block_low + (dim_tile * 8 + 7) * 32)));
+                svmopa_za32_bf16_m(
+                    0, query_predicate_b16x, predicate_all_b16x, probability_column_8_f32x,
+                    svld1_bf16(predicate_all_b16x, (bfloat16_t const *)(values_block_high + (dim_tile * 8 + 0) * 32)));
+                svmopa_za32_bf16_m(
+                    0, query_predicate_b16x, predicate_all_b16x, probability_column_9_f32x,
+                    svld1_bf16(predicate_all_b16x, (bfloat16_t const *)(values_block_high + (dim_tile * 8 + 1) * 32)));
+                svmopa_za32_bf16_m(
+                    0, query_predicate_b16x, predicate_all_b16x, probability_column_10_f32x,
+                    svld1_bf16(predicate_all_b16x, (bfloat16_t const *)(values_block_high + (dim_tile * 8 + 2) * 32)));
+                svmopa_za32_bf16_m(
+                    0, query_predicate_b16x, predicate_all_b16x, probability_column_11_f32x,
+                    svld1_bf16(predicate_all_b16x, (bfloat16_t const *)(values_block_high + (dim_tile * 8 + 3) * 32)));
+                svmopa_za32_bf16_m(
+                    0, query_predicate_b16x, predicate_all_b16x, probability_column_12_f32x,
+                    svld1_bf16(predicate_all_b16x, (bfloat16_t const *)(values_block_high + (dim_tile * 8 + 4) * 32)));
+                svmopa_za32_bf16_m(
+                    0, query_predicate_b16x, predicate_all_b16x, probability_column_13_f32x,
+                    svld1_bf16(predicate_all_b16x, (bfloat16_t const *)(values_block_high + (dim_tile * 8 + 5) * 32)));
+                svmopa_za32_bf16_m(
+                    0, query_predicate_b16x, predicate_all_b16x, probability_column_14_f32x,
+                    svld1_bf16(predicate_all_b16x, (bfloat16_t const *)(values_block_high + (dim_tile * 8 + 6) * 32)));
+                svmopa_za32_bf16_m(
+                    0, query_predicate_b16x, predicate_all_b16x, probability_column_15_f32x,
+                    svld1_bf16(predicate_all_b16x, (bfloat16_t const *)(values_block_high + (dim_tile * 8 + 7) * 32)));
+                for (nk_size_t query_index = 0; query_index < valid_query_count; query_index++)
+                    svst1_f32(predicate_all_b32x, output_accumulator + query_index * head_dim_padded + dim_tile * 16,
+                              svadd_f32_x(predicate_all_b32x,
+                                          svld1_f32(predicate_all_b32x,
+                                                    output_accumulator + query_index * head_dim_padded + dim_tile * 16),
+                                          svread_hor_za32_f32_m(svdup_f32(0), predicate_all_b32x, 0, query_index)));
+            }
+        }
+    }
+
+    // Bc=16 tail loop (handles remaining KV positions and decode path)
+    for (; kv_start < kv_len; kv_start += 16, kv_block_index++) {
+        nk_size_t const valid_kv = ((kv_start + 16) <= kv_len) ? 16 : (kv_len - kv_start);
+
+        // Q×K^T: pure memory→BFMOPA, no ZA staging
+        svzero_mask_za(nk_sme_zero_za32_tile_2_);
+        nk_bf16_t const *k_block = k + kv_block_index * k_depth_step_count * 32;
+        for (nk_size_t step = 0; step < k_depth_step_count; step++) {
+            svbfloat16_t zn_bf16x = svreinterpret_bf16_f32(
+                svld1_f32(predicate_all_b32x, queries_transposed + step * 16));
+            svbfloat16_t zm_bf16x = svld1_bf16(predicate_all_b16x, (bfloat16_t const *)(k_block + step * 32));
+            svmopa_za32_bf16_m(2, predicate_all_b32x, predicate_all_b32x, zn_bf16x, zm_bf16x);
+        }
+
+        // Pass 1: Column-wise max (read ZA2 columns vertically)
+        svfloat32_t scale_16_f32x = svdup_f32(scale);
+        svfloat32_t block_max_16_f32x = svdup_f32(NK_F32_MIN);
+        for (nk_size_t column_index = 0; column_index < 16; column_index++) {
+            svfloat32_t score_column_f32x = svmul_f32_x(
+                predicate_all_b32x, svread_ver_za32_f32_m(svdup_f32(0), predicate_all_b32x, 2, column_index),
+                scale_16_f32x);
+            block_max_16_f32x = svmax_f32_x(predicate_all_b32x, block_max_16_f32x, score_column_f32x);
+        }
+
+        // Softmax correction (fully vectorized)
+        svfloat32_t new_max_f32x = svmax_f32_x(predicate_all_b32x, row_max_f32x, block_max_16_f32x);
+        svfloat32_t correction_f32x = nk_exp_fast_f32_sve_(predicate_all_b32x,
+                                                           svsub_f32_x(predicate_all_b32x, row_max_f32x, new_max_f32x));
+        svbool_t max_changed_16_b32x = svcmplt_f32(predicate_all_b32x, correction_f32x, svdup_f32(1.0f));
+        nk_u32_t max_was_updated_16 = svptest_any(predicate_all_b32x, max_changed_16_b32x) ? 1 : 0;
+        if (max_was_updated_16) row_sum_f32x = svmul_f32_x(predicate_all_b32x, row_sum_f32x, correction_f32x);
+        NK_ALIGN64 nk_f32_t corrections[16];
+        svst1_f32(predicate_all_b32x, corrections, correction_f32x);
+
+        // Pass 2: Column-wise exp + fused P write + sum (ZA2 → ZA0 columns 0-7)
+        svfloat32_t sum_delta_16_f32x = svdup_f32(0.0f);
+        svzero_mask_za(nk_sme_zero_za32_tile_0_);
+        for (nk_size_t column_index = 0; column_index < 16; column_index += 2) {
+            svfloat32_t score_even_f32x = svmul_f32_x(
+                predicate_all_b32x, svread_ver_za32_f32_m(svdup_f32(0), predicate_all_b32x, 2, column_index),
+                scale_16_f32x);
+            svfloat32_t score_odd_f32x = svmul_f32_x(
+                predicate_all_b32x, svread_ver_za32_f32_m(svdup_f32(0), predicate_all_b32x, 2, column_index + 1),
+                scale_16_f32x);
+            svfloat32_t weight_even_f32x = nk_exp_fast_f32_sve_(
+                predicate_all_b32x, svsub_f32_x(predicate_all_b32x, score_even_f32x, new_max_f32x));
+            svfloat32_t weight_odd_f32x = nk_exp_fast_f32_sve_(
+                predicate_all_b32x, svsub_f32_x(predicate_all_b32x, score_odd_f32x, new_max_f32x));
+            sum_delta_16_f32x = svadd_f32_x(predicate_all_b32x, sum_delta_16_f32x, weight_even_f32x);
+            sum_delta_16_f32x = svadd_f32_x(predicate_all_b32x, sum_delta_16_f32x, weight_odd_f32x);
+            svbfloat16_t weight_pair_bf16x = svzip1_bf16(nk_f32_to_bf16_sve_(predicate_all_b32x, weight_even_f32x),
+                                                         nk_f32_to_bf16_sve_(predicate_all_b32x, weight_odd_f32x));
+            svwrite_ver_za32_f32_m(0, column_index / 2, predicate_all_b32x, svreinterpret_f32_bf16(weight_pair_bf16x));
+        }
+        row_sum_f32x = svadd_f32_x(predicate_all_b32x, row_sum_f32x, sum_delta_16_f32x);
+        row_max_f32x = new_max_f32x;
+
+        if (valid_query_count == 1) {
+            // Decode path: extract f32 weights from ZA0 row 0 using SVE
+            svbfloat16_t row0_bf16x = svreinterpret_bf16_f32(
+                svread_hor_za32_f32_m(svdup_f32(0), predicate_all_b32x, 0, 0));
+            svbfloat16_t weights_even_bf16x = svuzp1_bf16(row0_bf16x, row0_bf16x);
+            svbfloat16_t weights_odd_bf16x = svuzp2_bf16(row0_bf16x, row0_bf16x);
+            NK_ALIGN64 nk_f32_t decode_weights[16];
+            svst1_f32(svwhilelt_b32(0u, 8u), decode_weights,
+                      nk_bf16_to_f32_sve_(svwhilelt_b32(0u, 8u), weights_even_bf16x));
+            svst1_f32(svwhilelt_b32(0u, 8u), decode_weights + 8,
+                      nk_bf16_to_f32_sve_(svwhilelt_b32(0u, 8u), weights_odd_bf16x));
+            NK_ALIGN64 nk_f32_t decode_weights_ordered[16];
+            for (nk_size_t i = 0; i < 8; i++) {
+                decode_weights_ordered[2 * i] = decode_weights[i];
+                decode_weights_ordered[2 * i + 1] = decode_weights[8 + i];
+            }
+            svfloat32_t corr_f32x = svdup_f32(corrections[0]);
+            for (nk_size_t d = 0; d < head_dim; d += svcntw()) {
+                svbool_t predicate_b32x = svwhilelt_b32_u64(d, head_dim);
+                svfloat32_t acc_f32x = svmul_f32_x(predicate_b32x, svld1_f32(predicate_b32x, output_accumulator + d),
+                                                   corr_f32x);
+                for (nk_size_t ki = 0; ki < valid_kv; ki++) {
+                    nk_size_t dim_tile = d / 16, depth_s = ki / 2, sub = ki % 2;
+                    nk_bf16_t const *v_vec = v_packed +
+                                             (kv_block_index * dim_tile_count * 8 + dim_tile * 8 + depth_s) * 32;
+                    svbfloat16_t packed_bf16x = svld1_bf16(predicate_all_b16x, (bfloat16_t const *)v_vec);
+                    svbfloat16_t v_selected_bf16x = (sub == 0) ? svuzp1_bf16(packed_bf16x, packed_bf16x)
+                                                               : svuzp2_bf16(packed_bf16x, packed_bf16x);
+                    acc_f32x = svmla_f32_x(predicate_b32x, acc_f32x, svdup_f32(decode_weights_ordered[ki]),
+                                           nk_bf16_to_f32_sve_(predicate_b32x, v_selected_bf16x));
+                }
+                svst1_f32(predicate_b32x, output_accumulator + d, acc_f32x);
+            }
+        }
+        else {
+            // Prefill Bc=16: extract P columns, pre-apply correction, add-after P×V
+            svbool_t query_predicate_b16x = svwhilelt_b16_u64(0u, valid_query_count * 2);
+
+            svbfloat16_t probability_column_0_f32x = svreinterpret_bf16_f32(
+                svread_ver_za32_f32_m(svdup_f32(0), predicate_all_b32x, 0, 0));
+            svbfloat16_t probability_column_1_f32x = svreinterpret_bf16_f32(
+                svread_ver_za32_f32_m(svdup_f32(0), predicate_all_b32x, 0, 1));
+            svbfloat16_t probability_column_2_f32x = svreinterpret_bf16_f32(
+                svread_ver_za32_f32_m(svdup_f32(0), predicate_all_b32x, 0, 2));
+            svbfloat16_t probability_column_3_f32x = svreinterpret_bf16_f32(
+                svread_ver_za32_f32_m(svdup_f32(0), predicate_all_b32x, 0, 3));
+            svbfloat16_t probability_column_4_f32x = svreinterpret_bf16_f32(
+                svread_ver_za32_f32_m(svdup_f32(0), predicate_all_b32x, 0, 4));
+            svbfloat16_t probability_column_5_f32x = svreinterpret_bf16_f32(
+                svread_ver_za32_f32_m(svdup_f32(0), predicate_all_b32x, 0, 5));
+            svbfloat16_t probability_column_6_f32x = svreinterpret_bf16_f32(
+                svread_ver_za32_f32_m(svdup_f32(0), predicate_all_b32x, 0, 6));
+            svbfloat16_t probability_column_7_f32x = svreinterpret_bf16_f32(
+                svread_ver_za32_f32_m(svdup_f32(0), predicate_all_b32x, 0, 7));
+
+            nk_bf16_t const *v_block = v_packed + kv_block_index * dim_tile_count * 8 * 32;
+
+            // Pre-apply correction
+            if (max_was_updated_16) {
+                for (nk_size_t query_index = 0; query_index < valid_query_count; query_index++) {
+                    svfloat32_t correction_scalar_f32x = svdup_f32(corrections[query_index]);
+                    for (nk_size_t dim_offset = 0; dim_offset < head_dim_padded; dim_offset += 16)
+                        svst1_f32(
+                            predicate_all_b32x, output_accumulator + query_index * head_dim_padded + dim_offset,
+                            svmul_f32_x(predicate_all_b32x,
+                                        svld1_f32(predicate_all_b32x,
+                                                  output_accumulator + query_index * head_dim_padded + dim_offset),
+                                        correction_scalar_f32x));
+                }
+            }
+
+            // P×V: zero → BFMOPA → read → add
+            nk_size_t dim_tile = 0;
+            for (; dim_tile + 4 <= dim_tile_count; dim_tile += 4) {
+                svzero_za();
+                svmopa_za32_bf16_m(
+                    0, query_predicate_b16x, predicate_all_b16x, probability_column_0_f32x,
+                    svld1_bf16(predicate_all_b16x, (bfloat16_t const *)(v_block + ((dim_tile + 0) * 8 + 0) * 32)));
+                svmopa_za32_bf16_m(
+                    1, query_predicate_b16x, predicate_all_b16x, probability_column_0_f32x,
+                    svld1_bf16(predicate_all_b16x, (bfloat16_t const *)(v_block + ((dim_tile + 1) * 8 + 0) * 32)));
+                svmopa_za32_bf16_m(
+                    2, query_predicate_b16x, predicate_all_b16x, probability_column_0_f32x,
+                    svld1_bf16(predicate_all_b16x, (bfloat16_t const *)(v_block + ((dim_tile + 2) * 8 + 0) * 32)));
+                svmopa_za32_bf16_m(
+                    3, query_predicate_b16x, predicate_all_b16x, probability_column_0_f32x,
+                    svld1_bf16(predicate_all_b16x, (bfloat16_t const *)(v_block + ((dim_tile + 3) * 8 + 0) * 32)));
+                svmopa_za32_bf16_m(
+                    0, query_predicate_b16x, predicate_all_b16x, probability_column_1_f32x,
+                    svld1_bf16(predicate_all_b16x, (bfloat16_t const *)(v_block + ((dim_tile + 0) * 8 + 1) * 32)));
+                svmopa_za32_bf16_m(
+                    1, query_predicate_b16x, predicate_all_b16x, probability_column_1_f32x,
+                    svld1_bf16(predicate_all_b16x, (bfloat16_t const *)(v_block + ((dim_tile + 1) * 8 + 1) * 32)));
+                svmopa_za32_bf16_m(
+                    2, query_predicate_b16x, predicate_all_b16x, probability_column_1_f32x,
+                    svld1_bf16(predicate_all_b16x, (bfloat16_t const *)(v_block + ((dim_tile + 2) * 8 + 1) * 32)));
+                svmopa_za32_bf16_m(
+                    3, query_predicate_b16x, predicate_all_b16x, probability_column_1_f32x,
+                    svld1_bf16(predicate_all_b16x, (bfloat16_t const *)(v_block + ((dim_tile + 3) * 8 + 1) * 32)));
+                svmopa_za32_bf16_m(
+                    0, query_predicate_b16x, predicate_all_b16x, probability_column_2_f32x,
+                    svld1_bf16(predicate_all_b16x, (bfloat16_t const *)(v_block + ((dim_tile + 0) * 8 + 2) * 32)));
+                svmopa_za32_bf16_m(
+                    1, query_predicate_b16x, predicate_all_b16x, probability_column_2_f32x,
+                    svld1_bf16(predicate_all_b16x, (bfloat16_t const *)(v_block + ((dim_tile + 1) * 8 + 2) * 32)));
+                svmopa_za32_bf16_m(
+                    2, query_predicate_b16x, predicate_all_b16x, probability_column_2_f32x,
+                    svld1_bf16(predicate_all_b16x, (bfloat16_t const *)(v_block + ((dim_tile + 2) * 8 + 2) * 32)));
+                svmopa_za32_bf16_m(
+                    3, query_predicate_b16x, predicate_all_b16x, probability_column_2_f32x,
+                    svld1_bf16(predicate_all_b16x, (bfloat16_t const *)(v_block + ((dim_tile + 3) * 8 + 2) * 32)));
+                svmopa_za32_bf16_m(
+                    0, query_predicate_b16x, predicate_all_b16x, probability_column_3_f32x,
+                    svld1_bf16(predicate_all_b16x, (bfloat16_t const *)(v_block + ((dim_tile + 0) * 8 + 3) * 32)));
+                svmopa_za32_bf16_m(
+                    1, query_predicate_b16x, predicate_all_b16x, probability_column_3_f32x,
+                    svld1_bf16(predicate_all_b16x, (bfloat16_t const *)(v_block + ((dim_tile + 1) * 8 + 3) * 32)));
+                svmopa_za32_bf16_m(
+                    2, query_predicate_b16x, predicate_all_b16x, probability_column_3_f32x,
+                    svld1_bf16(predicate_all_b16x, (bfloat16_t const *)(v_block + ((dim_tile + 2) * 8 + 3) * 32)));
+                svmopa_za32_bf16_m(
+                    3, query_predicate_b16x, predicate_all_b16x, probability_column_3_f32x,
+                    svld1_bf16(predicate_all_b16x, (bfloat16_t const *)(v_block + ((dim_tile + 3) * 8 + 3) * 32)));
+                svmopa_za32_bf16_m(
+                    0, query_predicate_b16x, predicate_all_b16x, probability_column_4_f32x,
+                    svld1_bf16(predicate_all_b16x, (bfloat16_t const *)(v_block + ((dim_tile + 0) * 8 + 4) * 32)));
+                svmopa_za32_bf16_m(
+                    1, query_predicate_b16x, predicate_all_b16x, probability_column_4_f32x,
+                    svld1_bf16(predicate_all_b16x, (bfloat16_t const *)(v_block + ((dim_tile + 1) * 8 + 4) * 32)));
+                svmopa_za32_bf16_m(
+                    2, query_predicate_b16x, predicate_all_b16x, probability_column_4_f32x,
+                    svld1_bf16(predicate_all_b16x, (bfloat16_t const *)(v_block + ((dim_tile + 2) * 8 + 4) * 32)));
+                svmopa_za32_bf16_m(
+                    3, query_predicate_b16x, predicate_all_b16x, probability_column_4_f32x,
+                    svld1_bf16(predicate_all_b16x, (bfloat16_t const *)(v_block + ((dim_tile + 3) * 8 + 4) * 32)));
+                svmopa_za32_bf16_m(
+                    0, query_predicate_b16x, predicate_all_b16x, probability_column_5_f32x,
+                    svld1_bf16(predicate_all_b16x, (bfloat16_t const *)(v_block + ((dim_tile + 0) * 8 + 5) * 32)));
+                svmopa_za32_bf16_m(
+                    1, query_predicate_b16x, predicate_all_b16x, probability_column_5_f32x,
+                    svld1_bf16(predicate_all_b16x, (bfloat16_t const *)(v_block + ((dim_tile + 1) * 8 + 5) * 32)));
+                svmopa_za32_bf16_m(
+                    2, query_predicate_b16x, predicate_all_b16x, probability_column_5_f32x,
+                    svld1_bf16(predicate_all_b16x, (bfloat16_t const *)(v_block + ((dim_tile + 2) * 8 + 5) * 32)));
+                svmopa_za32_bf16_m(
+                    3, query_predicate_b16x, predicate_all_b16x, probability_column_5_f32x,
+                    svld1_bf16(predicate_all_b16x, (bfloat16_t const *)(v_block + ((dim_tile + 3) * 8 + 5) * 32)));
+                svmopa_za32_bf16_m(
+                    0, query_predicate_b16x, predicate_all_b16x, probability_column_6_f32x,
+                    svld1_bf16(predicate_all_b16x, (bfloat16_t const *)(v_block + ((dim_tile + 0) * 8 + 6) * 32)));
+                svmopa_za32_bf16_m(
+                    1, query_predicate_b16x, predicate_all_b16x, probability_column_6_f32x,
+                    svld1_bf16(predicate_all_b16x, (bfloat16_t const *)(v_block + ((dim_tile + 1) * 8 + 6) * 32)));
+                svmopa_za32_bf16_m(
+                    2, query_predicate_b16x, predicate_all_b16x, probability_column_6_f32x,
+                    svld1_bf16(predicate_all_b16x, (bfloat16_t const *)(v_block + ((dim_tile + 2) * 8 + 6) * 32)));
+                svmopa_za32_bf16_m(
+                    3, query_predicate_b16x, predicate_all_b16x, probability_column_6_f32x,
+                    svld1_bf16(predicate_all_b16x, (bfloat16_t const *)(v_block + ((dim_tile + 3) * 8 + 6) * 32)));
+                svmopa_za32_bf16_m(
+                    0, query_predicate_b16x, predicate_all_b16x, probability_column_7_f32x,
+                    svld1_bf16(predicate_all_b16x, (bfloat16_t const *)(v_block + ((dim_tile + 0) * 8 + 7) * 32)));
+                svmopa_za32_bf16_m(
+                    1, query_predicate_b16x, predicate_all_b16x, probability_column_7_f32x,
+                    svld1_bf16(predicate_all_b16x, (bfloat16_t const *)(v_block + ((dim_tile + 1) * 8 + 7) * 32)));
+                svmopa_za32_bf16_m(
+                    2, query_predicate_b16x, predicate_all_b16x, probability_column_7_f32x,
+                    svld1_bf16(predicate_all_b16x, (bfloat16_t const *)(v_block + ((dim_tile + 2) * 8 + 7) * 32)));
+                svmopa_za32_bf16_m(
+                    3, query_predicate_b16x, predicate_all_b16x, probability_column_7_f32x,
+                    svld1_bf16(predicate_all_b16x, (bfloat16_t const *)(v_block + ((dim_tile + 3) * 8 + 7) * 32)));
+                for (nk_size_t query_index = 0; query_index < valid_query_count; query_index++) {
+                    svst1_f32(
+                        predicate_all_b32x, output_accumulator + query_index * head_dim_padded + (dim_tile + 0) * 16,
+                        svadd_f32_x(predicate_all_b32x,
+                                    svld1_f32(predicate_all_b32x,
+                                              output_accumulator + query_index * head_dim_padded + (dim_tile + 0) * 16),
+                                    svread_hor_za32_f32_m(svdup_f32(0), predicate_all_b32x, 0, query_index)));
+                    svst1_f32(
+                        predicate_all_b32x, output_accumulator + query_index * head_dim_padded + (dim_tile + 1) * 16,
+                        svadd_f32_x(predicate_all_b32x,
+                                    svld1_f32(predicate_all_b32x,
+                                              output_accumulator + query_index * head_dim_padded + (dim_tile + 1) * 16),
+                                    svread_hor_za32_f32_m(svdup_f32(0), predicate_all_b32x, 1, query_index)));
+                    svst1_f32(
+                        predicate_all_b32x, output_accumulator + query_index * head_dim_padded + (dim_tile + 2) * 16,
+                        svadd_f32_x(predicate_all_b32x,
+                                    svld1_f32(predicate_all_b32x,
+                                              output_accumulator + query_index * head_dim_padded + (dim_tile + 2) * 16),
+                                    svread_hor_za32_f32_m(svdup_f32(0), predicate_all_b32x, 2, query_index)));
+                    svst1_f32(
+                        predicate_all_b32x, output_accumulator + query_index * head_dim_padded + (dim_tile + 3) * 16,
+                        svadd_f32_x(predicate_all_b32x,
+                                    svld1_f32(predicate_all_b32x,
+                                              output_accumulator + query_index * head_dim_padded + (dim_tile + 3) * 16),
+                                    svread_hor_za32_f32_m(svdup_f32(0), predicate_all_b32x, 3, query_index)));
+                }
+            }
+            for (; dim_tile < dim_tile_count; dim_tile++) {
+                svzero_mask_za(nk_sme_zero_za32_tile_0_);
+                svmopa_za32_bf16_m(
+                    0, query_predicate_b16x, predicate_all_b16x, probability_column_0_f32x,
+                    svld1_bf16(predicate_all_b16x, (bfloat16_t const *)(v_block + (dim_tile * 8 + 0) * 32)));
+                svmopa_za32_bf16_m(
+                    0, query_predicate_b16x, predicate_all_b16x, probability_column_1_f32x,
+                    svld1_bf16(predicate_all_b16x, (bfloat16_t const *)(v_block + (dim_tile * 8 + 1) * 32)));
+                svmopa_za32_bf16_m(
+                    0, query_predicate_b16x, predicate_all_b16x, probability_column_2_f32x,
+                    svld1_bf16(predicate_all_b16x, (bfloat16_t const *)(v_block + (dim_tile * 8 + 2) * 32)));
+                svmopa_za32_bf16_m(
+                    0, query_predicate_b16x, predicate_all_b16x, probability_column_3_f32x,
+                    svld1_bf16(predicate_all_b16x, (bfloat16_t const *)(v_block + (dim_tile * 8 + 3) * 32)));
+                svmopa_za32_bf16_m(
+                    0, query_predicate_b16x, predicate_all_b16x, probability_column_4_f32x,
+                    svld1_bf16(predicate_all_b16x, (bfloat16_t const *)(v_block + (dim_tile * 8 + 4) * 32)));
+                svmopa_za32_bf16_m(
+                    0, query_predicate_b16x, predicate_all_b16x, probability_column_5_f32x,
+                    svld1_bf16(predicate_all_b16x, (bfloat16_t const *)(v_block + (dim_tile * 8 + 5) * 32)));
+                svmopa_za32_bf16_m(
+                    0, query_predicate_b16x, predicate_all_b16x, probability_column_6_f32x,
+                    svld1_bf16(predicate_all_b16x, (bfloat16_t const *)(v_block + (dim_tile * 8 + 6) * 32)));
+                svmopa_za32_bf16_m(
+                    0, query_predicate_b16x, predicate_all_b16x, probability_column_7_f32x,
+                    svld1_bf16(predicate_all_b16x, (bfloat16_t const *)(v_block + (dim_tile * 8 + 7) * 32)));
+                for (nk_size_t query_index = 0; query_index < valid_query_count; query_index++)
+                    svst1_f32(predicate_all_b32x, output_accumulator + query_index * head_dim_padded + dim_tile * 16,
+                              svadd_f32_x(predicate_all_b32x,
+                                          svld1_f32(predicate_all_b32x,
+                                                    output_accumulator + query_index * head_dim_padded + dim_tile * 16),
+                                          svread_hor_za32_f32_m(svdup_f32(0), predicate_all_b32x, 0, query_index)));
+            }
+        }
+    }
+
+    // Final normalization
+    NK_ALIGN64 nk_f32_t final_sums[16];
+    svst1_f32(predicate_all_b32x, final_sums, row_sum_f32x);
+
+    for (nk_size_t query_index = 0; query_index < valid_query_count; query_index++) {
+        nk_f32_t inv_sum = (final_sums[query_index] > 0.0f) ? (1.0f / final_sums[query_index]) : 0.0f;
+        svfloat32_t inv_sum_f32x = svdup_f32(inv_sum);
+
+        for (nk_size_t dim_offset = 0; dim_offset < head_dim; dim_offset += svcntw()) {
+            svbool_t predicate_b32x = svwhilelt_b32_u64(dim_offset, head_dim);
+            svfloat32_t output_f32x = svmul_f32_x(
+                predicate_b32x,
+                svld1_f32(predicate_b32x, output_accumulator + query_index * head_dim_padded + dim_offset),
+                inv_sum_f32x);
+            svbfloat16_t output_bf16x = nk_f32_to_bf16_sve_(predicate_b32x, output_f32x);
+            nk_size_t store_count = (head_dim - dim_offset) < (nk_size_t)svcntw() ? (head_dim - dim_offset)
+                                                                                  : (nk_size_t)svcntw();
+            svbool_t store_predicate_b16x = svwhilelt_b16_u64(0u, store_count);
+            svst1_bf16(store_predicate_b16x, (bfloat16_t *)(output + query_index * head_dim + dim_offset),
+                       output_bf16x);
+        }
+    }
+}
+
+NK_PUBLIC void nk_attention_bf16_sme( //
+    nk_bf16_t const *q, void const *kv_packed, nk_bf16_t *output, nk_size_t num_heads, nk_size_t num_kv_heads,
+    nk_size_t query_len, nk_size_t kv_len, nk_size_t head_dim, nk_f32_t scale) {
+
+    nk_attention_sme_packed_header_t const *header = (nk_attention_sme_packed_header_t const *)kv_packed;
+    nk_size_t head_dim_padded = header->head_dim_padded;
+    nk_size_t dim_tile_count = header->reserved[0]; // v_dim_tile_count
+    nk_size_t kv_blocks = (kv_len + 15) / 16;
+    nk_size_t kv_head_stride = kv_blocks * 16 * head_dim_padded;
+
+    nk_bf16_t const *k_packed = (nk_bf16_t const *)((char const *)kv_packed + header->k_offset);
+    nk_bf16_t const *v_packed = (nk_bf16_t const *)((char const *)kv_packed + header->v_offset);
+
+    nk_size_t group_size = (num_kv_heads > 0) ? num_heads / num_kv_heads : 1;
+
+    nk_sme_start_streaming_();
+    for (nk_size_t q_head = 0; q_head < num_heads; q_head++) {
+        nk_size_t kv_head = q_head / group_size;
+
+        nk_bf16_t const *q_ptr = q + q_head * query_len * head_dim;
+        nk_bf16_t const *k_ptr = k_packed + kv_head * kv_head_stride;
+        nk_bf16_t const *v_ptr = v_packed + kv_head * kv_head_stride;
+        nk_bf16_t *out_ptr = output + q_head * query_len * head_dim;
+
+        for (nk_size_t q_start = 0; q_start < query_len; q_start += 16) {
+            nk_size_t q_block_len = (q_start + 16 < query_len) ? 16 : (query_len - q_start);
+
+            nk_attention_bf16_sme_streaming_(q_ptr + q_start * head_dim, k_ptr, v_ptr, out_ptr + q_start * head_dim,
+                                             q_block_len, kv_len, head_dim, head_dim_padded, dim_tile_count, scale);
+        }
+    }
+    nk_sme_stop_streaming_();
+}
+
+__arm_new("za") static void nk_attention_f16_sme_streaming_( //
+    nk_f16_t const *q, nk_f16_t const *k, nk_f16_t const *v_packed, nk_f16_t *output, nk_size_t query_len,
+    nk_size_t kv_len, nk_size_t head_dim, nk_size_t head_dim_padded, nk_size_t dim_tile_count,
+    nk_f32_t scale) NK_STREAMING_ {
+
+    svbool_t const predicate_all_b32x = svptrue_b32();
+    svbool_t const predicate_all_b16x = svptrue_b16();
+    nk_size_t const valid_query_count = (query_len < 16) ? query_len : 16;
+
+    NK_ALIGN64 nk_f32_t row_max[16];
+    NK_ALIGN64 nk_f32_t row_sum[16];
+    NK_ALIGN64 nk_f32_t output_accumulator[16 * 256];
+
+    svst1_f32(predicate_all_b32x, row_max, svdup_f32(NK_F32_MIN));
+    svst1_f32(predicate_all_b32x, row_sum, svdup_f32(0.0f));
+    svfloat32_t zero_f32x = svdup_f32(0.0f);
+    for (nk_size_t i = 0; i < 16 * head_dim_padded; i += svcntw()) {
+        svst1_f32(predicate_all_b32x, output_accumulator + i, zero_f32x);
+    }
+
+    nk_size_t kv_block_index = 0;
+    nk_size_t kv_start = 0;
+    svbool_t const batch_predicate_b32x = svwhilelt_b32(0u, 16u);
+
+    nk_size_t const k_depth_step_count = head_dim_padded / 2;
+
+    // Pre-transpose Q once: queries_transposed[step][16 f32 words]
+    // queries_transposed[step] reinterpret-as-f16 = {Q[0][2s], Q[0][2s+1], Q[1][2s], Q[1][2s+1], ...}
+    // This is the same interleaving ZA0 vertical reads would produce.
+    NK_ALIGN64 nk_f32_t queries_transposed[128 * 16]; // max head_dim_padded/2 * 16 = 128 * 16
+    for (nk_size_t batch = 0; batch < head_dim_padded / 32; batch++) {
+        svzero_mask_za(nk_sme_zero_za32_tile_0_);
+        for (nk_size_t query_index = 0; query_index < valid_query_count; query_index++)
+            svld1_hor_za32(0, query_index, batch_predicate_b32x,
+                           (nk_f32_t const *)(q + query_index * head_dim + batch * 32));
+        for (nk_size_t step = 0; step < 16; step++)
+            svst1_f32(predicate_all_b32x, queries_transposed + (batch * 16 + step) * 16,
+                      svread_ver_za32_f32_m(svdup_f32(0), predicate_all_b32x, 0, step));
+    }
+
+    // === Bc=32 main loop (prefill only, skipped for decode) ===
+    if (valid_query_count > 1) {
+        for (; kv_start + 32 <= kv_len; kv_start += 32, kv_block_index += 2) {
+            // Q×K^T: pure memory→FMOPA, no ZA staging for Q or K
+            svzero_mask_za(nk_sme_zero_za32_tile_2_);
+            svzero_mask_za(nk_sme_zero_za32_tile_3_);
+            nk_f16_t const *keys_block_low = k + kv_block_index * k_depth_step_count * 32;
+            nk_f16_t const *keys_block_high = k + (kv_block_index + 1) * k_depth_step_count * 32;
+            for (nk_size_t step = 0; step < k_depth_step_count; step++) {
+                svfloat16_t zn_f16x = svreinterpret_f16_f32(
+                    svld1_f32(predicate_all_b32x, queries_transposed + step * 16));
+                svfloat16_t zm0_f16x = svld1_f16(predicate_all_b16x, (float16_t const *)(keys_block_low + step * 32));
+                svfloat16_t zm1_f16x = svld1_f16(predicate_all_b16x, (float16_t const *)(keys_block_high + step * 32));
+                svmopa_za32_f16_m(2, predicate_all_b32x, predicate_all_b32x, zn_f16x, zm0_f16x);
+                svmopa_za32_f16_m(3, predicate_all_b32x, predicate_all_b32x, zn_f16x, zm1_f16x);
+            }
+            // ZA2 = scores[query_index][0:15], ZA3 = scores[query_index][16:31]
+
+            // Pass 1: Column-wise max (read ZA2/ZA3 columns vertically)
+            svfloat32_t scale_f32x = svdup_f32(scale);
+            svfloat32_t block_max_f32x = svdup_f32(NK_F32_MIN);
+            for (nk_size_t column_index = 0; column_index < 16; column_index++) {
+                svfloat32_t score_column_f32x = svmul_f32_x(
+                    predicate_all_b32x, svread_ver_za32_f32_m(svdup_f32(0), predicate_all_b32x, 2, column_index),
+                    scale_f32x);
+                block_max_f32x = svmax_f32_x(predicate_all_b32x, block_max_f32x, score_column_f32x);
+            }
+            for (nk_size_t column_index = 0; column_index < 16; column_index++) {
+                svfloat32_t score_column_f32x = svmul_f32_x(
+                    predicate_all_b32x, svread_ver_za32_f32_m(svdup_f32(0), predicate_all_b32x, 3, column_index),
+                    scale_f32x);
+                block_max_f32x = svmax_f32_x(predicate_all_b32x, block_max_f32x, score_column_f32x);
+            }
+
+            // Softmax correction (vectorized via array load/store)
+            svfloat32_t old_max_f32x = svld1_f32(predicate_all_b32x, row_max);
+            svfloat32_t new_max_f32x = svmax_f32_x(predicate_all_b32x, old_max_f32x, block_max_f32x);
+            svfloat32_t correction_f32x = nk_exp_fast_f32_sve_(
+                predicate_all_b32x, svsub_f32_x(predicate_all_b32x, old_max_f32x, new_max_f32x));
+            svbool_t max_changed_b32x = svcmplt_f32(predicate_all_b32x, correction_f32x, svdup_f32(1.0f));
+            nk_u32_t max_was_updated = svptest_any(predicate_all_b32x, max_changed_b32x) ? 1 : 0;
+            svfloat32_t row_sum_corrected_f32x = svld1_f32(predicate_all_b32x, row_sum);
+            if (max_was_updated)
+                row_sum_corrected_f32x = svmul_f32_x(predicate_all_b32x, row_sum_corrected_f32x, correction_f32x);
+            NK_ALIGN64 nk_f32_t corrections[16];
+            svst1_f32(predicate_all_b32x, corrections, correction_f32x);
+
+            // Pass 2: Column-wise exp + fused P write + sum
+            svfloat32_t sum_delta_f32x = svdup_f32(0.0f);
+            svzero_mask_za(nk_sme_zero_za32_tile_0_);
+            // ZA2 columns in pairs -> ZA0 columns 0-7
+            for (nk_size_t column_index = 0; column_index < 16; column_index += 2) {
+                svfloat32_t score_even_f32x = svmul_f32_x(
+                    predicate_all_b32x, svread_ver_za32_f32_m(svdup_f32(0), predicate_all_b32x, 2, column_index),
+                    scale_f32x);
+                svfloat32_t score_odd_f32x = svmul_f32_x(
+                    predicate_all_b32x, svread_ver_za32_f32_m(svdup_f32(0), predicate_all_b32x, 2, column_index + 1),
+                    scale_f32x);
+                svfloat32_t weight_even_f32x = nk_exp_fast_f32_sve_(
+                    predicate_all_b32x, svsub_f32_x(predicate_all_b32x, score_even_f32x, new_max_f32x));
+                svfloat32_t weight_odd_f32x = nk_exp_fast_f32_sve_(
+                    predicate_all_b32x, svsub_f32_x(predicate_all_b32x, score_odd_f32x, new_max_f32x));
+                sum_delta_f32x = svadd_f32_x(predicate_all_b32x, sum_delta_f32x, weight_even_f32x);
+                sum_delta_f32x = svadd_f32_x(predicate_all_b32x, sum_delta_f32x, weight_odd_f32x);
+                svfloat16_t weight_pair_f16x = svzip1_f16(svcvt_f16_f32_x(predicate_all_b32x, weight_even_f32x),
+                                                          svcvt_f16_f32_x(predicate_all_b32x, weight_odd_f32x));
+                svwrite_ver_za32_f32_m(0, column_index / 2, predicate_all_b32x,
+                                       svreinterpret_f32_f16(weight_pair_f16x));
+            }
+            // ZA3 columns in pairs -> ZA0 columns 8-15
+            for (nk_size_t column_index = 0; column_index < 16; column_index += 2) {
+                svfloat32_t score_even_f32x = svmul_f32_x(
+                    predicate_all_b32x, svread_ver_za32_f32_m(svdup_f32(0), predicate_all_b32x, 3, column_index),
+                    scale_f32x);
+                svfloat32_t score_odd_f32x = svmul_f32_x(
+                    predicate_all_b32x, svread_ver_za32_f32_m(svdup_f32(0), predicate_all_b32x, 3, column_index + 1),
+                    scale_f32x);
+                svfloat32_t weight_even_f32x = nk_exp_fast_f32_sve_(
+                    predicate_all_b32x, svsub_f32_x(predicate_all_b32x, score_even_f32x, new_max_f32x));
+                svfloat32_t weight_odd_f32x = nk_exp_fast_f32_sve_(
+                    predicate_all_b32x, svsub_f32_x(predicate_all_b32x, score_odd_f32x, new_max_f32x));
+                sum_delta_f32x = svadd_f32_x(predicate_all_b32x, sum_delta_f32x, weight_even_f32x);
+                sum_delta_f32x = svadd_f32_x(predicate_all_b32x, sum_delta_f32x, weight_odd_f32x);
+                svfloat16_t weight_pair_f16x = svzip1_f16(svcvt_f16_f32_x(predicate_all_b32x, weight_even_f32x),
+                                                          svcvt_f16_f32_x(predicate_all_b32x, weight_odd_f32x));
+                svwrite_ver_za32_f32_m(0, 8 + column_index / 2, predicate_all_b32x,
+                                       svreinterpret_f32_f16(weight_pair_f16x));
+            }
+            row_sum_corrected_f32x = svadd_f32_x(predicate_all_b32x, row_sum_corrected_f32x, sum_delta_f32x);
+            svst1_f32(predicate_all_b32x, row_sum, row_sum_corrected_f32x);
+            svst1_f32(predicate_all_b32x, row_max, new_max_f32x);
+
+            // Extract P columns from ZA0
+            svfloat16_t probability_column_0_f32x = svreinterpret_f16_f32(
+                svread_ver_za32_f32_m(svdup_f32(0), predicate_all_b32x, 0, 0));
+            svfloat16_t probability_column_1_f32x = svreinterpret_f16_f32(
+                svread_ver_za32_f32_m(svdup_f32(0), predicate_all_b32x, 0, 1));
+            svfloat16_t probability_column_2_f32x = svreinterpret_f16_f32(
+                svread_ver_za32_f32_m(svdup_f32(0), predicate_all_b32x, 0, 2));
+            svfloat16_t probability_column_3_f32x = svreinterpret_f16_f32(
+                svread_ver_za32_f32_m(svdup_f32(0), predicate_all_b32x, 0, 3));
+            svfloat16_t probability_column_4_f32x = svreinterpret_f16_f32(
+                svread_ver_za32_f32_m(svdup_f32(0), predicate_all_b32x, 0, 4));
+            svfloat16_t probability_column_5_f32x = svreinterpret_f16_f32(
+                svread_ver_za32_f32_m(svdup_f32(0), predicate_all_b32x, 0, 5));
+            svfloat16_t probability_column_6_f32x = svreinterpret_f16_f32(
+                svread_ver_za32_f32_m(svdup_f32(0), predicate_all_b32x, 0, 6));
+            svfloat16_t probability_column_7_f32x = svreinterpret_f16_f32(
+                svread_ver_za32_f32_m(svdup_f32(0), predicate_all_b32x, 0, 7));
+            svfloat16_t probability_column_8_f32x = svreinterpret_f16_f32(
+                svread_ver_za32_f32_m(svdup_f32(0), predicate_all_b32x, 0, 8));
+            svfloat16_t probability_column_9_f32x = svreinterpret_f16_f32(
+                svread_ver_za32_f32_m(svdup_f32(0), predicate_all_b32x, 0, 9));
+            svfloat16_t probability_column_10_f32x = svreinterpret_f16_f32(
+                svread_ver_za32_f32_m(svdup_f32(0), predicate_all_b32x, 0, 10));
+            svfloat16_t probability_column_11_f32x = svreinterpret_f16_f32(
+                svread_ver_za32_f32_m(svdup_f32(0), predicate_all_b32x, 0, 11));
+            svfloat16_t probability_column_12_f32x = svreinterpret_f16_f32(
+                svread_ver_za32_f32_m(svdup_f32(0), predicate_all_b32x, 0, 12));
+            svfloat16_t probability_column_13_f32x = svreinterpret_f16_f32(
+                svread_ver_za32_f32_m(svdup_f32(0), predicate_all_b32x, 0, 13));
+            svfloat16_t probability_column_14_f32x = svreinterpret_f16_f32(
+                svread_ver_za32_f32_m(svdup_f32(0), predicate_all_b32x, 0, 14));
+            svfloat16_t probability_column_15_f32x = svreinterpret_f16_f32(
+                svread_ver_za32_f32_m(svdup_f32(0), predicate_all_b32x, 0, 15));
+
+            // Pre-apply correction once before P×V
+            svbool_t query_predicate_b16x = svwhilelt_b16_u64(0u, valid_query_count * 2);
+            nk_f16_t const *values_block_low = v_packed + kv_block_index * dim_tile_count * 8 * 32;
+            nk_f16_t const *values_block_high = v_packed + (kv_block_index + 1) * dim_tile_count * 8 * 32;
+
+            if (max_was_updated) {
+                for (nk_size_t query_index = 0; query_index < valid_query_count; query_index++) {
+                    svfloat32_t correction_scalar_f32x = svdup_f32(corrections[query_index]);
+                    for (nk_size_t dim_offset = 0; dim_offset < head_dim_padded; dim_offset += 16)
+                        svst1_f32(
+                            predicate_all_b32x, output_accumulator + query_index * head_dim_padded + dim_offset,
+                            svmul_f32_x(predicate_all_b32x,
+                                        svld1_f32(predicate_all_b32x,
+                                                  output_accumulator + query_index * head_dim_padded + dim_offset),
+                                        correction_scalar_f32x));
+                }
+            }
+
+            // P×V: zero -> FMOPA -> read -> add (no ZA writes for output_accumulator)
+            nk_size_t dim_tile = 0;
+            for (; dim_tile + 4 <= dim_tile_count; dim_tile += 4) {
+                svzero_za();
+                // Block0: 8 depth steps (KV positions 0-15)
+                svmopa_za32_f16_m(0, query_predicate_b16x, predicate_all_b16x, probability_column_0_f32x,
+                                  svld1_f16(predicate_all_b16x,
+                                            (float16_t const *)(values_block_low + ((dim_tile + 0) * 8 + 0) * 32)));
+                svmopa_za32_f16_m(1, query_predicate_b16x, predicate_all_b16x, probability_column_0_f32x,
+                                  svld1_f16(predicate_all_b16x,
+                                            (float16_t const *)(values_block_low + ((dim_tile + 1) * 8 + 0) * 32)));
+                svmopa_za32_f16_m(2, query_predicate_b16x, predicate_all_b16x, probability_column_0_f32x,
+                                  svld1_f16(predicate_all_b16x,
+                                            (float16_t const *)(values_block_low + ((dim_tile + 2) * 8 + 0) * 32)));
+                svmopa_za32_f16_m(3, query_predicate_b16x, predicate_all_b16x, probability_column_0_f32x,
+                                  svld1_f16(predicate_all_b16x,
+                                            (float16_t const *)(values_block_low + ((dim_tile + 3) * 8 + 0) * 32)));
+                svmopa_za32_f16_m(0, query_predicate_b16x, predicate_all_b16x, probability_column_1_f32x,
+                                  svld1_f16(predicate_all_b16x,
+                                            (float16_t const *)(values_block_low + ((dim_tile + 0) * 8 + 1) * 32)));
+                svmopa_za32_f16_m(1, query_predicate_b16x, predicate_all_b16x, probability_column_1_f32x,
+                                  svld1_f16(predicate_all_b16x,
+                                            (float16_t const *)(values_block_low + ((dim_tile + 1) * 8 + 1) * 32)));
+                svmopa_za32_f16_m(2, query_predicate_b16x, predicate_all_b16x, probability_column_1_f32x,
+                                  svld1_f16(predicate_all_b16x,
+                                            (float16_t const *)(values_block_low + ((dim_tile + 2) * 8 + 1) * 32)));
+                svmopa_za32_f16_m(3, query_predicate_b16x, predicate_all_b16x, probability_column_1_f32x,
+                                  svld1_f16(predicate_all_b16x,
+                                            (float16_t const *)(values_block_low + ((dim_tile + 3) * 8 + 1) * 32)));
+                svmopa_za32_f16_m(0, query_predicate_b16x, predicate_all_b16x, probability_column_2_f32x,
+                                  svld1_f16(predicate_all_b16x,
+                                            (float16_t const *)(values_block_low + ((dim_tile + 0) * 8 + 2) * 32)));
+                svmopa_za32_f16_m(1, query_predicate_b16x, predicate_all_b16x, probability_column_2_f32x,
+                                  svld1_f16(predicate_all_b16x,
+                                            (float16_t const *)(values_block_low + ((dim_tile + 1) * 8 + 2) * 32)));
+                svmopa_za32_f16_m(2, query_predicate_b16x, predicate_all_b16x, probability_column_2_f32x,
+                                  svld1_f16(predicate_all_b16x,
+                                            (float16_t const *)(values_block_low + ((dim_tile + 2) * 8 + 2) * 32)));
+                svmopa_za32_f16_m(3, query_predicate_b16x, predicate_all_b16x, probability_column_2_f32x,
+                                  svld1_f16(predicate_all_b16x,
+                                            (float16_t const *)(values_block_low + ((dim_tile + 3) * 8 + 2) * 32)));
+                svmopa_za32_f16_m(0, query_predicate_b16x, predicate_all_b16x, probability_column_3_f32x,
+                                  svld1_f16(predicate_all_b16x,
+                                            (float16_t const *)(values_block_low + ((dim_tile + 0) * 8 + 3) * 32)));
+                svmopa_za32_f16_m(1, query_predicate_b16x, predicate_all_b16x, probability_column_3_f32x,
+                                  svld1_f16(predicate_all_b16x,
+                                            (float16_t const *)(values_block_low + ((dim_tile + 1) * 8 + 3) * 32)));
+                svmopa_za32_f16_m(2, query_predicate_b16x, predicate_all_b16x, probability_column_3_f32x,
+                                  svld1_f16(predicate_all_b16x,
+                                            (float16_t const *)(values_block_low + ((dim_tile + 2) * 8 + 3) * 32)));
+                svmopa_za32_f16_m(3, query_predicate_b16x, predicate_all_b16x, probability_column_3_f32x,
+                                  svld1_f16(predicate_all_b16x,
+                                            (float16_t const *)(values_block_low + ((dim_tile + 3) * 8 + 3) * 32)));
+                svmopa_za32_f16_m(0, query_predicate_b16x, predicate_all_b16x, probability_column_4_f32x,
+                                  svld1_f16(predicate_all_b16x,
+                                            (float16_t const *)(values_block_low + ((dim_tile + 0) * 8 + 4) * 32)));
+                svmopa_za32_f16_m(1, query_predicate_b16x, predicate_all_b16x, probability_column_4_f32x,
+                                  svld1_f16(predicate_all_b16x,
+                                            (float16_t const *)(values_block_low + ((dim_tile + 1) * 8 + 4) * 32)));
+                svmopa_za32_f16_m(2, query_predicate_b16x, predicate_all_b16x, probability_column_4_f32x,
+                                  svld1_f16(predicate_all_b16x,
+                                            (float16_t const *)(values_block_low + ((dim_tile + 2) * 8 + 4) * 32)));
+                svmopa_za32_f16_m(3, query_predicate_b16x, predicate_all_b16x, probability_column_4_f32x,
+                                  svld1_f16(predicate_all_b16x,
+                                            (float16_t const *)(values_block_low + ((dim_tile + 3) * 8 + 4) * 32)));
+                svmopa_za32_f16_m(0, query_predicate_b16x, predicate_all_b16x, probability_column_5_f32x,
+                                  svld1_f16(predicate_all_b16x,
+                                            (float16_t const *)(values_block_low + ((dim_tile + 0) * 8 + 5) * 32)));
+                svmopa_za32_f16_m(1, query_predicate_b16x, predicate_all_b16x, probability_column_5_f32x,
+                                  svld1_f16(predicate_all_b16x,
+                                            (float16_t const *)(values_block_low + ((dim_tile + 1) * 8 + 5) * 32)));
+                svmopa_za32_f16_m(2, query_predicate_b16x, predicate_all_b16x, probability_column_5_f32x,
+                                  svld1_f16(predicate_all_b16x,
+                                            (float16_t const *)(values_block_low + ((dim_tile + 2) * 8 + 5) * 32)));
+                svmopa_za32_f16_m(3, query_predicate_b16x, predicate_all_b16x, probability_column_5_f32x,
+                                  svld1_f16(predicate_all_b16x,
+                                            (float16_t const *)(values_block_low + ((dim_tile + 3) * 8 + 5) * 32)));
+                svmopa_za32_f16_m(0, query_predicate_b16x, predicate_all_b16x, probability_column_6_f32x,
+                                  svld1_f16(predicate_all_b16x,
+                                            (float16_t const *)(values_block_low + ((dim_tile + 0) * 8 + 6) * 32)));
+                svmopa_za32_f16_m(1, query_predicate_b16x, predicate_all_b16x, probability_column_6_f32x,
+                                  svld1_f16(predicate_all_b16x,
+                                            (float16_t const *)(values_block_low + ((dim_tile + 1) * 8 + 6) * 32)));
+                svmopa_za32_f16_m(2, query_predicate_b16x, predicate_all_b16x, probability_column_6_f32x,
+                                  svld1_f16(predicate_all_b16x,
+                                            (float16_t const *)(values_block_low + ((dim_tile + 2) * 8 + 6) * 32)));
+                svmopa_za32_f16_m(3, query_predicate_b16x, predicate_all_b16x, probability_column_6_f32x,
+                                  svld1_f16(predicate_all_b16x,
+                                            (float16_t const *)(values_block_low + ((dim_tile + 3) * 8 + 6) * 32)));
+                svmopa_za32_f16_m(0, query_predicate_b16x, predicate_all_b16x, probability_column_7_f32x,
+                                  svld1_f16(predicate_all_b16x,
+                                            (float16_t const *)(values_block_low + ((dim_tile + 0) * 8 + 7) * 32)));
+                svmopa_za32_f16_m(1, query_predicate_b16x, predicate_all_b16x, probability_column_7_f32x,
+                                  svld1_f16(predicate_all_b16x,
+                                            (float16_t const *)(values_block_low + ((dim_tile + 1) * 8 + 7) * 32)));
+                svmopa_za32_f16_m(2, query_predicate_b16x, predicate_all_b16x, probability_column_7_f32x,
+                                  svld1_f16(predicate_all_b16x,
+                                            (float16_t const *)(values_block_low + ((dim_tile + 2) * 8 + 7) * 32)));
+                svmopa_za32_f16_m(3, query_predicate_b16x, predicate_all_b16x, probability_column_7_f32x,
+                                  svld1_f16(predicate_all_b16x,
+                                            (float16_t const *)(values_block_low + ((dim_tile + 3) * 8 + 7) * 32)));
+                // Block1: 8 depth steps (KV positions 16-31)
+                svmopa_za32_f16_m(0, query_predicate_b16x, predicate_all_b16x, probability_column_8_f32x,
+                                  svld1_f16(predicate_all_b16x,
+                                            (float16_t const *)(values_block_high + ((dim_tile + 0) * 8 + 0) * 32)));
+                svmopa_za32_f16_m(1, query_predicate_b16x, predicate_all_b16x, probability_column_8_f32x,
+                                  svld1_f16(predicate_all_b16x,
+                                            (float16_t const *)(values_block_high + ((dim_tile + 1) * 8 + 0) * 32)));
+                svmopa_za32_f16_m(2, query_predicate_b16x, predicate_all_b16x, probability_column_8_f32x,
+                                  svld1_f16(predicate_all_b16x,
+                                            (float16_t const *)(values_block_high + ((dim_tile + 2) * 8 + 0) * 32)));
+                svmopa_za32_f16_m(3, query_predicate_b16x, predicate_all_b16x, probability_column_8_f32x,
+                                  svld1_f16(predicate_all_b16x,
+                                            (float16_t const *)(values_block_high + ((dim_tile + 3) * 8 + 0) * 32)));
+                svmopa_za32_f16_m(0, query_predicate_b16x, predicate_all_b16x, probability_column_9_f32x,
+                                  svld1_f16(predicate_all_b16x,
+                                            (float16_t const *)(values_block_high + ((dim_tile + 0) * 8 + 1) * 32)));
+                svmopa_za32_f16_m(1, query_predicate_b16x, predicate_all_b16x, probability_column_9_f32x,
+                                  svld1_f16(predicate_all_b16x,
+                                            (float16_t const *)(values_block_high + ((dim_tile + 1) * 8 + 1) * 32)));
+                svmopa_za32_f16_m(2, query_predicate_b16x, predicate_all_b16x, probability_column_9_f32x,
+                                  svld1_f16(predicate_all_b16x,
+                                            (float16_t const *)(values_block_high + ((dim_tile + 2) * 8 + 1) * 32)));
+                svmopa_za32_f16_m(3, query_predicate_b16x, predicate_all_b16x, probability_column_9_f32x,
+                                  svld1_f16(predicate_all_b16x,
+                                            (float16_t const *)(values_block_high + ((dim_tile + 3) * 8 + 1) * 32)));
+                svmopa_za32_f16_m(0, query_predicate_b16x, predicate_all_b16x, probability_column_10_f32x,
+                                  svld1_f16(predicate_all_b16x,
+                                            (float16_t const *)(values_block_high + ((dim_tile + 0) * 8 + 2) * 32)));
+                svmopa_za32_f16_m(1, query_predicate_b16x, predicate_all_b16x, probability_column_10_f32x,
+                                  svld1_f16(predicate_all_b16x,
+                                            (float16_t const *)(values_block_high + ((dim_tile + 1) * 8 + 2) * 32)));
+                svmopa_za32_f16_m(2, query_predicate_b16x, predicate_all_b16x, probability_column_10_f32x,
+                                  svld1_f16(predicate_all_b16x,
+                                            (float16_t const *)(values_block_high + ((dim_tile + 2) * 8 + 2) * 32)));
+                svmopa_za32_f16_m(3, query_predicate_b16x, predicate_all_b16x, probability_column_10_f32x,
+                                  svld1_f16(predicate_all_b16x,
+                                            (float16_t const *)(values_block_high + ((dim_tile + 3) * 8 + 2) * 32)));
+                svmopa_za32_f16_m(0, query_predicate_b16x, predicate_all_b16x, probability_column_11_f32x,
+                                  svld1_f16(predicate_all_b16x,
+                                            (float16_t const *)(values_block_high + ((dim_tile + 0) * 8 + 3) * 32)));
+                svmopa_za32_f16_m(1, query_predicate_b16x, predicate_all_b16x, probability_column_11_f32x,
+                                  svld1_f16(predicate_all_b16x,
+                                            (float16_t const *)(values_block_high + ((dim_tile + 1) * 8 + 3) * 32)));
+                svmopa_za32_f16_m(2, query_predicate_b16x, predicate_all_b16x, probability_column_11_f32x,
+                                  svld1_f16(predicate_all_b16x,
+                                            (float16_t const *)(values_block_high + ((dim_tile + 2) * 8 + 3) * 32)));
+                svmopa_za32_f16_m(3, query_predicate_b16x, predicate_all_b16x, probability_column_11_f32x,
+                                  svld1_f16(predicate_all_b16x,
+                                            (float16_t const *)(values_block_high + ((dim_tile + 3) * 8 + 3) * 32)));
+                svmopa_za32_f16_m(0, query_predicate_b16x, predicate_all_b16x, probability_column_12_f32x,
+                                  svld1_f16(predicate_all_b16x,
+                                            (float16_t const *)(values_block_high + ((dim_tile + 0) * 8 + 4) * 32)));
+                svmopa_za32_f16_m(1, query_predicate_b16x, predicate_all_b16x, probability_column_12_f32x,
+                                  svld1_f16(predicate_all_b16x,
+                                            (float16_t const *)(values_block_high + ((dim_tile + 1) * 8 + 4) * 32)));
+                svmopa_za32_f16_m(2, query_predicate_b16x, predicate_all_b16x, probability_column_12_f32x,
+                                  svld1_f16(predicate_all_b16x,
+                                            (float16_t const *)(values_block_high + ((dim_tile + 2) * 8 + 4) * 32)));
+                svmopa_za32_f16_m(3, query_predicate_b16x, predicate_all_b16x, probability_column_12_f32x,
+                                  svld1_f16(predicate_all_b16x,
+                                            (float16_t const *)(values_block_high + ((dim_tile + 3) * 8 + 4) * 32)));
+                svmopa_za32_f16_m(0, query_predicate_b16x, predicate_all_b16x, probability_column_13_f32x,
+                                  svld1_f16(predicate_all_b16x,
+                                            (float16_t const *)(values_block_high + ((dim_tile + 0) * 8 + 5) * 32)));
+                svmopa_za32_f16_m(1, query_predicate_b16x, predicate_all_b16x, probability_column_13_f32x,
+                                  svld1_f16(predicate_all_b16x,
+                                            (float16_t const *)(values_block_high + ((dim_tile + 1) * 8 + 5) * 32)));
+                svmopa_za32_f16_m(2, query_predicate_b16x, predicate_all_b16x, probability_column_13_f32x,
+                                  svld1_f16(predicate_all_b16x,
+                                            (float16_t const *)(values_block_high + ((dim_tile + 2) * 8 + 5) * 32)));
+                svmopa_za32_f16_m(3, query_predicate_b16x, predicate_all_b16x, probability_column_13_f32x,
+                                  svld1_f16(predicate_all_b16x,
+                                            (float16_t const *)(values_block_high + ((dim_tile + 3) * 8 + 5) * 32)));
+                svmopa_za32_f16_m(0, query_predicate_b16x, predicate_all_b16x, probability_column_14_f32x,
+                                  svld1_f16(predicate_all_b16x,
+                                            (float16_t const *)(values_block_high + ((dim_tile + 0) * 8 + 6) * 32)));
+                svmopa_za32_f16_m(1, query_predicate_b16x, predicate_all_b16x, probability_column_14_f32x,
+                                  svld1_f16(predicate_all_b16x,
+                                            (float16_t const *)(values_block_high + ((dim_tile + 1) * 8 + 6) * 32)));
+                svmopa_za32_f16_m(2, query_predicate_b16x, predicate_all_b16x, probability_column_14_f32x,
+                                  svld1_f16(predicate_all_b16x,
+                                            (float16_t const *)(values_block_high + ((dim_tile + 2) * 8 + 6) * 32)));
+                svmopa_za32_f16_m(3, query_predicate_b16x, predicate_all_b16x, probability_column_14_f32x,
+                                  svld1_f16(predicate_all_b16x,
+                                            (float16_t const *)(values_block_high + ((dim_tile + 3) * 8 + 6) * 32)));
+                svmopa_za32_f16_m(0, query_predicate_b16x, predicate_all_b16x, probability_column_15_f32x,
+                                  svld1_f16(predicate_all_b16x,
+                                            (float16_t const *)(values_block_high + ((dim_tile + 0) * 8 + 7) * 32)));
+                svmopa_za32_f16_m(1, query_predicate_b16x, predicate_all_b16x, probability_column_15_f32x,
+                                  svld1_f16(predicate_all_b16x,
+                                            (float16_t const *)(values_block_high + ((dim_tile + 1) * 8 + 7) * 32)));
+                svmopa_za32_f16_m(2, query_predicate_b16x, predicate_all_b16x, probability_column_15_f32x,
+                                  svld1_f16(predicate_all_b16x,
+                                            (float16_t const *)(values_block_high + ((dim_tile + 2) * 8 + 7) * 32)));
+                svmopa_za32_f16_m(3, query_predicate_b16x, predicate_all_b16x, probability_column_15_f32x,
+                                  svld1_f16(predicate_all_b16x,
+                                            (float16_t const *)(values_block_high + ((dim_tile + 3) * 8 + 7) * 32)));
+                // Read FMOPA result and ADD to output_accumulator
+                for (nk_size_t query_index = 0; query_index < valid_query_count; query_index++) {
+                    svst1_f32(
+                        predicate_all_b32x, output_accumulator + query_index * head_dim_padded + (dim_tile + 0) * 16,
+                        svadd_f32_x(predicate_all_b32x,
+                                    svld1_f32(predicate_all_b32x,
+                                              output_accumulator + query_index * head_dim_padded + (dim_tile + 0) * 16),
+                                    svread_hor_za32_f32_m(svdup_f32(0), predicate_all_b32x, 0, query_index)));
+                    svst1_f32(
+                        predicate_all_b32x, output_accumulator + query_index * head_dim_padded + (dim_tile + 1) * 16,
+                        svadd_f32_x(predicate_all_b32x,
+                                    svld1_f32(predicate_all_b32x,
+                                              output_accumulator + query_index * head_dim_padded + (dim_tile + 1) * 16),
+                                    svread_hor_za32_f32_m(svdup_f32(0), predicate_all_b32x, 1, query_index)));
+                    svst1_f32(
+                        predicate_all_b32x, output_accumulator + query_index * head_dim_padded + (dim_tile + 2) * 16,
+                        svadd_f32_x(predicate_all_b32x,
+                                    svld1_f32(predicate_all_b32x,
+                                              output_accumulator + query_index * head_dim_padded + (dim_tile + 2) * 16),
+                                    svread_hor_za32_f32_m(svdup_f32(0), predicate_all_b32x, 2, query_index)));
+                    svst1_f32(
+                        predicate_all_b32x, output_accumulator + query_index * head_dim_padded + (dim_tile + 3) * 16,
+                        svadd_f32_x(predicate_all_b32x,
+                                    svld1_f32(predicate_all_b32x,
+                                              output_accumulator + query_index * head_dim_padded + (dim_tile + 3) * 16),
+                                    svread_hor_za32_f32_m(svdup_f32(0), predicate_all_b32x, 3, query_index)));
+                }
+            }
+            // Remainder: 1 dim_tile at a time using ZA0
+            for (; dim_tile < dim_tile_count; dim_tile++) {
+                svzero_mask_za(nk_sme_zero_za32_tile_0_);
+                svmopa_za32_f16_m(
+                    0, query_predicate_b16x, predicate_all_b16x, probability_column_0_f32x,
+                    svld1_f16(predicate_all_b16x, (float16_t const *)(values_block_low + (dim_tile * 8 + 0) * 32)));
+                svmopa_za32_f16_m(
+                    0, query_predicate_b16x, predicate_all_b16x, probability_column_1_f32x,
+                    svld1_f16(predicate_all_b16x, (float16_t const *)(values_block_low + (dim_tile * 8 + 1) * 32)));
+                svmopa_za32_f16_m(
+                    0, query_predicate_b16x, predicate_all_b16x, probability_column_2_f32x,
+                    svld1_f16(predicate_all_b16x, (float16_t const *)(values_block_low + (dim_tile * 8 + 2) * 32)));
+                svmopa_za32_f16_m(
+                    0, query_predicate_b16x, predicate_all_b16x, probability_column_3_f32x,
+                    svld1_f16(predicate_all_b16x, (float16_t const *)(values_block_low + (dim_tile * 8 + 3) * 32)));
+                svmopa_za32_f16_m(
+                    0, query_predicate_b16x, predicate_all_b16x, probability_column_4_f32x,
+                    svld1_f16(predicate_all_b16x, (float16_t const *)(values_block_low + (dim_tile * 8 + 4) * 32)));
+                svmopa_za32_f16_m(
+                    0, query_predicate_b16x, predicate_all_b16x, probability_column_5_f32x,
+                    svld1_f16(predicate_all_b16x, (float16_t const *)(values_block_low + (dim_tile * 8 + 5) * 32)));
+                svmopa_za32_f16_m(
+                    0, query_predicate_b16x, predicate_all_b16x, probability_column_6_f32x,
+                    svld1_f16(predicate_all_b16x, (float16_t const *)(values_block_low + (dim_tile * 8 + 6) * 32)));
+                svmopa_za32_f16_m(
+                    0, query_predicate_b16x, predicate_all_b16x, probability_column_7_f32x,
+                    svld1_f16(predicate_all_b16x, (float16_t const *)(values_block_low + (dim_tile * 8 + 7) * 32)));
+                svmopa_za32_f16_m(
+                    0, query_predicate_b16x, predicate_all_b16x, probability_column_8_f32x,
+                    svld1_f16(predicate_all_b16x, (float16_t const *)(values_block_high + (dim_tile * 8 + 0) * 32)));
+                svmopa_za32_f16_m(
+                    0, query_predicate_b16x, predicate_all_b16x, probability_column_9_f32x,
+                    svld1_f16(predicate_all_b16x, (float16_t const *)(values_block_high + (dim_tile * 8 + 1) * 32)));
+                svmopa_za32_f16_m(
+                    0, query_predicate_b16x, predicate_all_b16x, probability_column_10_f32x,
+                    svld1_f16(predicate_all_b16x, (float16_t const *)(values_block_high + (dim_tile * 8 + 2) * 32)));
+                svmopa_za32_f16_m(
+                    0, query_predicate_b16x, predicate_all_b16x, probability_column_11_f32x,
+                    svld1_f16(predicate_all_b16x, (float16_t const *)(values_block_high + (dim_tile * 8 + 3) * 32)));
+                svmopa_za32_f16_m(
+                    0, query_predicate_b16x, predicate_all_b16x, probability_column_12_f32x,
+                    svld1_f16(predicate_all_b16x, (float16_t const *)(values_block_high + (dim_tile * 8 + 4) * 32)));
+                svmopa_za32_f16_m(
+                    0, query_predicate_b16x, predicate_all_b16x, probability_column_13_f32x,
+                    svld1_f16(predicate_all_b16x, (float16_t const *)(values_block_high + (dim_tile * 8 + 5) * 32)));
+                svmopa_za32_f16_m(
+                    0, query_predicate_b16x, predicate_all_b16x, probability_column_14_f32x,
+                    svld1_f16(predicate_all_b16x, (float16_t const *)(values_block_high + (dim_tile * 8 + 6) * 32)));
+                svmopa_za32_f16_m(
+                    0, query_predicate_b16x, predicate_all_b16x, probability_column_15_f32x,
+                    svld1_f16(predicate_all_b16x, (float16_t const *)(values_block_high + (dim_tile * 8 + 7) * 32)));
+                for (nk_size_t query_index = 0; query_index < valid_query_count; query_index++)
+                    svst1_f32(predicate_all_b32x, output_accumulator + query_index * head_dim_padded + dim_tile * 16,
+                              svadd_f32_x(predicate_all_b32x,
+                                          svld1_f32(predicate_all_b32x,
+                                                    output_accumulator + query_index * head_dim_padded + dim_tile * 16),
+                                          svread_hor_za32_f32_m(svdup_f32(0), predicate_all_b32x, 0, query_index)));
+            }
+        }
+    }
+
+    // === Bc=16 tail loop (handles remaining KV positions and decode path) ===
+    for (; kv_start < kv_len; kv_start += 16, kv_block_index++) {
+        nk_size_t const valid_kv = ((kv_start + 16) <= kv_len) ? 16 : (kv_len - kv_start);
+
+        // Q×K^T: pure memory→FMOPA, no ZA staging
+        svzero_mask_za(nk_sme_zero_za32_tile_2_);
+        nk_f16_t const *k_block = k + kv_block_index * k_depth_step_count * 32;
+        for (nk_size_t step = 0; step < k_depth_step_count; step++) {
+            svfloat16_t zn_f16x = svreinterpret_f16_f32(svld1_f32(predicate_all_b32x, queries_transposed + step * 16));
+            svfloat16_t zm_f16x = svld1_f16(predicate_all_b16x, (float16_t const *)(k_block + step * 32));
+            svmopa_za32_f16_m(2, predicate_all_b32x, predicate_all_b32x, zn_f16x, zm_f16x);
+        }
+
+        // Pass 1: Column-wise max (read ZA2 columns vertically)
+        svfloat32_t scale_16_f32x = svdup_f32(scale);
+        svfloat32_t block_max_16_f32x = svdup_f32(NK_F32_MIN);
+        for (nk_size_t column_index = 0; column_index < 16; column_index++) {
+            svfloat32_t score_column_f32x = svmul_f32_x(
+                predicate_all_b32x, svread_ver_za32_f32_m(svdup_f32(0), predicate_all_b32x, 2, column_index),
+                scale_16_f32x);
+            block_max_16_f32x = svmax_f32_x(predicate_all_b32x, block_max_16_f32x, score_column_f32x);
+        }
+
+        svfloat32_t old_max_f32x = svld1_f32(predicate_all_b32x, row_max);
+        svfloat32_t new_max_f32x = svmax_f32_x(predicate_all_b32x, old_max_f32x, block_max_16_f32x);
+        svfloat32_t correction_f32x = nk_exp_fast_f32_sve_(predicate_all_b32x,
+                                                           svsub_f32_x(predicate_all_b32x, old_max_f32x, new_max_f32x));
+        svbool_t max_changed_16_b32x = svcmplt_f32(predicate_all_b32x, correction_f32x, svdup_f32(1.0f));
+        nk_u32_t max_was_updated_16 = svptest_any(predicate_all_b32x, max_changed_16_b32x) ? 1 : 0;
+        svfloat32_t row_sum_corrected_f32x = svld1_f32(predicate_all_b32x, row_sum);
+        if (max_was_updated_16)
+            row_sum_corrected_f32x = svmul_f32_x(predicate_all_b32x, row_sum_corrected_f32x, correction_f32x);
+        NK_ALIGN64 nk_f32_t corrections[16];
+        svst1_f32(predicate_all_b32x, corrections, correction_f32x);
+
+        // Pass 2: Column-wise exp + fused P write + sum (ZA2 → ZA0 columns 0-7)
+        svfloat32_t sum_delta_16_f32x = svdup_f32(0.0f);
+        svzero_mask_za(nk_sme_zero_za32_tile_0_);
+        for (nk_size_t column_index = 0; column_index < 16; column_index += 2) {
+            svfloat32_t score_even_f32x = svmul_f32_x(
+                predicate_all_b32x, svread_ver_za32_f32_m(svdup_f32(0), predicate_all_b32x, 2, column_index),
+                scale_16_f32x);
+            svfloat32_t score_odd_f32x = svmul_f32_x(
+                predicate_all_b32x, svread_ver_za32_f32_m(svdup_f32(0), predicate_all_b32x, 2, column_index + 1),
+                scale_16_f32x);
+            svfloat32_t weight_even_f32x = nk_exp_fast_f32_sve_(
+                predicate_all_b32x, svsub_f32_x(predicate_all_b32x, score_even_f32x, new_max_f32x));
+            svfloat32_t weight_odd_f32x = nk_exp_fast_f32_sve_(
+                predicate_all_b32x, svsub_f32_x(predicate_all_b32x, score_odd_f32x, new_max_f32x));
+            sum_delta_16_f32x = svadd_f32_x(predicate_all_b32x, sum_delta_16_f32x, weight_even_f32x);
+            sum_delta_16_f32x = svadd_f32_x(predicate_all_b32x, sum_delta_16_f32x, weight_odd_f32x);
+            svfloat16_t weight_pair_f16x = svzip1_f16(svcvt_f16_f32_x(predicate_all_b32x, weight_even_f32x),
+                                                      svcvt_f16_f32_x(predicate_all_b32x, weight_odd_f32x));
+            svwrite_ver_za32_f32_m(0, column_index / 2, predicate_all_b32x, svreinterpret_f32_f16(weight_pair_f16x));
+        }
+        row_sum_corrected_f32x = svadd_f32_x(predicate_all_b32x, row_sum_corrected_f32x, sum_delta_16_f32x);
+        svst1_f32(predicate_all_b32x, row_sum, row_sum_corrected_f32x);
+        svst1_f32(predicate_all_b32x, row_max, new_max_f32x);
+
+        if (valid_query_count == 1) {
+            // Decode path: extract f32 weights from ZA0 row 0 using SVE
+            svfloat16_t row0_f16x = svreinterpret_f16_f32(
+                svread_hor_za32_f32_m(svdup_f32(0), predicate_all_b32x, 0, 0));
+            svfloat16_t weights_even_f16x = svuzp1_f16(row0_f16x, row0_f16x);
+            svfloat16_t weights_odd_f16x = svuzp2_f16(row0_f16x, row0_f16x);
+            NK_ALIGN64 nk_f32_t decode_weights[16];
+            svst1_f32(svwhilelt_b32(0u, 8u), decode_weights, svcvt_f32_f16_x(svwhilelt_b32(0u, 8u), weights_even_f16x));
+            svst1_f32(svwhilelt_b32(0u, 8u), decode_weights + 8,
+                      svcvt_f32_f16_x(svwhilelt_b32(0u, 8u), weights_odd_f16x));
+            NK_ALIGN64 nk_f32_t decode_weights_ordered[16];
+            for (nk_size_t i = 0; i < 8; i++) {
+                decode_weights_ordered[2 * i] = decode_weights[i];
+                decode_weights_ordered[2 * i + 1] = decode_weights[8 + i];
+            }
+            svfloat32_t corr_f32x = svdup_f32(corrections[0]);
+            for (nk_size_t d = 0; d < head_dim; d += svcntw()) {
+                svbool_t predicate_b32x = svwhilelt_b32_u64(d, head_dim);
+                svfloat32_t acc_f32x = svmul_f32_x(predicate_b32x, svld1_f32(predicate_b32x, output_accumulator + d),
+                                                   corr_f32x);
+                for (nk_size_t ki = 0; ki < valid_kv; ki++) {
+                    nk_size_t dim_tile = d / 16, depth_s = ki / 2, sub = ki % 2;
+                    nk_f16_t const *v_vec = v_packed +
+                                            (kv_block_index * dim_tile_count * 8 + dim_tile * 8 + depth_s) * 32;
+                    svfloat16_t packed_f16x = svld1_f16(predicate_all_b16x, (float16_t const *)v_vec);
+                    svfloat16_t v_selected_f16x = (sub == 0) ? svuzp1_f16(packed_f16x, packed_f16x)
+                                                             : svuzp2_f16(packed_f16x, packed_f16x);
+                    acc_f32x = svmla_f32_x(predicate_b32x, acc_f32x, svdup_f32(decode_weights_ordered[ki]),
+                                           svcvt_f32_f16_x(predicate_b32x, v_selected_f16x));
+                }
+                svst1_f32(predicate_b32x, output_accumulator + d, acc_f32x);
+            }
+        }
+        else {
+            // Prefill Bc=16: extract P columns, pre-apply correction, add-after P×V
+            svbool_t query_predicate_b16x = svwhilelt_b16_u64(0u, valid_query_count * 2);
+
+            svfloat16_t probability_column_0_f32x = svreinterpret_f16_f32(
+                svread_ver_za32_f32_m(svdup_f32(0), predicate_all_b32x, 0, 0));
+            svfloat16_t probability_column_1_f32x = svreinterpret_f16_f32(
+                svread_ver_za32_f32_m(svdup_f32(0), predicate_all_b32x, 0, 1));
+            svfloat16_t probability_column_2_f32x = svreinterpret_f16_f32(
+                svread_ver_za32_f32_m(svdup_f32(0), predicate_all_b32x, 0, 2));
+            svfloat16_t probability_column_3_f32x = svreinterpret_f16_f32(
+                svread_ver_za32_f32_m(svdup_f32(0), predicate_all_b32x, 0, 3));
+            svfloat16_t probability_column_4_f32x = svreinterpret_f16_f32(
+                svread_ver_za32_f32_m(svdup_f32(0), predicate_all_b32x, 0, 4));
+            svfloat16_t probability_column_5_f32x = svreinterpret_f16_f32(
+                svread_ver_za32_f32_m(svdup_f32(0), predicate_all_b32x, 0, 5));
+            svfloat16_t probability_column_6_f32x = svreinterpret_f16_f32(
+                svread_ver_za32_f32_m(svdup_f32(0), predicate_all_b32x, 0, 6));
+            svfloat16_t probability_column_7_f32x = svreinterpret_f16_f32(
+                svread_ver_za32_f32_m(svdup_f32(0), predicate_all_b32x, 0, 7));
+
+            nk_f16_t const *v_block = v_packed + kv_block_index * dim_tile_count * 8 * 32;
+
+            if (max_was_updated_16) {
+                for (nk_size_t query_index = 0; query_index < valid_query_count; query_index++) {
+                    svfloat32_t correction_scalar_f32x = svdup_f32(corrections[query_index]);
+                    for (nk_size_t dim_offset = 0; dim_offset < head_dim_padded; dim_offset += 16)
+                        svst1_f32(
+                            predicate_all_b32x, output_accumulator + query_index * head_dim_padded + dim_offset,
+                            svmul_f32_x(predicate_all_b32x,
+                                        svld1_f32(predicate_all_b32x,
+                                                  output_accumulator + query_index * head_dim_padded + dim_offset),
+                                        correction_scalar_f32x));
+                }
+            }
+
+            nk_size_t dim_tile = 0;
+            for (; dim_tile + 4 <= dim_tile_count; dim_tile += 4) {
+                svzero_za();
+                svmopa_za32_f16_m(
+                    0, query_predicate_b16x, predicate_all_b16x, probability_column_0_f32x,
+                    svld1_f16(predicate_all_b16x, (float16_t const *)(v_block + ((dim_tile + 0) * 8 + 0) * 32)));
+                svmopa_za32_f16_m(
+                    1, query_predicate_b16x, predicate_all_b16x, probability_column_0_f32x,
+                    svld1_f16(predicate_all_b16x, (float16_t const *)(v_block + ((dim_tile + 1) * 8 + 0) * 32)));
+                svmopa_za32_f16_m(
+                    2, query_predicate_b16x, predicate_all_b16x, probability_column_0_f32x,
+                    svld1_f16(predicate_all_b16x, (float16_t const *)(v_block + ((dim_tile + 2) * 8 + 0) * 32)));
+                svmopa_za32_f16_m(
+                    3, query_predicate_b16x, predicate_all_b16x, probability_column_0_f32x,
+                    svld1_f16(predicate_all_b16x, (float16_t const *)(v_block + ((dim_tile + 3) * 8 + 0) * 32)));
+                svmopa_za32_f16_m(
+                    0, query_predicate_b16x, predicate_all_b16x, probability_column_1_f32x,
+                    svld1_f16(predicate_all_b16x, (float16_t const *)(v_block + ((dim_tile + 0) * 8 + 1) * 32)));
+                svmopa_za32_f16_m(
+                    1, query_predicate_b16x, predicate_all_b16x, probability_column_1_f32x,
+                    svld1_f16(predicate_all_b16x, (float16_t const *)(v_block + ((dim_tile + 1) * 8 + 1) * 32)));
+                svmopa_za32_f16_m(
+                    2, query_predicate_b16x, predicate_all_b16x, probability_column_1_f32x,
+                    svld1_f16(predicate_all_b16x, (float16_t const *)(v_block + ((dim_tile + 2) * 8 + 1) * 32)));
+                svmopa_za32_f16_m(
+                    3, query_predicate_b16x, predicate_all_b16x, probability_column_1_f32x,
+                    svld1_f16(predicate_all_b16x, (float16_t const *)(v_block + ((dim_tile + 3) * 8 + 1) * 32)));
+                svmopa_za32_f16_m(
+                    0, query_predicate_b16x, predicate_all_b16x, probability_column_2_f32x,
+                    svld1_f16(predicate_all_b16x, (float16_t const *)(v_block + ((dim_tile + 0) * 8 + 2) * 32)));
+                svmopa_za32_f16_m(
+                    1, query_predicate_b16x, predicate_all_b16x, probability_column_2_f32x,
+                    svld1_f16(predicate_all_b16x, (float16_t const *)(v_block + ((dim_tile + 1) * 8 + 2) * 32)));
+                svmopa_za32_f16_m(
+                    2, query_predicate_b16x, predicate_all_b16x, probability_column_2_f32x,
+                    svld1_f16(predicate_all_b16x, (float16_t const *)(v_block + ((dim_tile + 2) * 8 + 2) * 32)));
+                svmopa_za32_f16_m(
+                    3, query_predicate_b16x, predicate_all_b16x, probability_column_2_f32x,
+                    svld1_f16(predicate_all_b16x, (float16_t const *)(v_block + ((dim_tile + 3) * 8 + 2) * 32)));
+                svmopa_za32_f16_m(
+                    0, query_predicate_b16x, predicate_all_b16x, probability_column_3_f32x,
+                    svld1_f16(predicate_all_b16x, (float16_t const *)(v_block + ((dim_tile + 0) * 8 + 3) * 32)));
+                svmopa_za32_f16_m(
+                    1, query_predicate_b16x, predicate_all_b16x, probability_column_3_f32x,
+                    svld1_f16(predicate_all_b16x, (float16_t const *)(v_block + ((dim_tile + 1) * 8 + 3) * 32)));
+                svmopa_za32_f16_m(
+                    2, query_predicate_b16x, predicate_all_b16x, probability_column_3_f32x,
+                    svld1_f16(predicate_all_b16x, (float16_t const *)(v_block + ((dim_tile + 2) * 8 + 3) * 32)));
+                svmopa_za32_f16_m(
+                    3, query_predicate_b16x, predicate_all_b16x, probability_column_3_f32x,
+                    svld1_f16(predicate_all_b16x, (float16_t const *)(v_block + ((dim_tile + 3) * 8 + 3) * 32)));
+                svmopa_za32_f16_m(
+                    0, query_predicate_b16x, predicate_all_b16x, probability_column_4_f32x,
+                    svld1_f16(predicate_all_b16x, (float16_t const *)(v_block + ((dim_tile + 0) * 8 + 4) * 32)));
+                svmopa_za32_f16_m(
+                    1, query_predicate_b16x, predicate_all_b16x, probability_column_4_f32x,
+                    svld1_f16(predicate_all_b16x, (float16_t const *)(v_block + ((dim_tile + 1) * 8 + 4) * 32)));
+                svmopa_za32_f16_m(
+                    2, query_predicate_b16x, predicate_all_b16x, probability_column_4_f32x,
+                    svld1_f16(predicate_all_b16x, (float16_t const *)(v_block + ((dim_tile + 2) * 8 + 4) * 32)));
+                svmopa_za32_f16_m(
+                    3, query_predicate_b16x, predicate_all_b16x, probability_column_4_f32x,
+                    svld1_f16(predicate_all_b16x, (float16_t const *)(v_block + ((dim_tile + 3) * 8 + 4) * 32)));
+                svmopa_za32_f16_m(
+                    0, query_predicate_b16x, predicate_all_b16x, probability_column_5_f32x,
+                    svld1_f16(predicate_all_b16x, (float16_t const *)(v_block + ((dim_tile + 0) * 8 + 5) * 32)));
+                svmopa_za32_f16_m(
+                    1, query_predicate_b16x, predicate_all_b16x, probability_column_5_f32x,
+                    svld1_f16(predicate_all_b16x, (float16_t const *)(v_block + ((dim_tile + 1) * 8 + 5) * 32)));
+                svmopa_za32_f16_m(
+                    2, query_predicate_b16x, predicate_all_b16x, probability_column_5_f32x,
+                    svld1_f16(predicate_all_b16x, (float16_t const *)(v_block + ((dim_tile + 2) * 8 + 5) * 32)));
+                svmopa_za32_f16_m(
+                    3, query_predicate_b16x, predicate_all_b16x, probability_column_5_f32x,
+                    svld1_f16(predicate_all_b16x, (float16_t const *)(v_block + ((dim_tile + 3) * 8 + 5) * 32)));
+                svmopa_za32_f16_m(
+                    0, query_predicate_b16x, predicate_all_b16x, probability_column_6_f32x,
+                    svld1_f16(predicate_all_b16x, (float16_t const *)(v_block + ((dim_tile + 0) * 8 + 6) * 32)));
+                svmopa_za32_f16_m(
+                    1, query_predicate_b16x, predicate_all_b16x, probability_column_6_f32x,
+                    svld1_f16(predicate_all_b16x, (float16_t const *)(v_block + ((dim_tile + 1) * 8 + 6) * 32)));
+                svmopa_za32_f16_m(
+                    2, query_predicate_b16x, predicate_all_b16x, probability_column_6_f32x,
+                    svld1_f16(predicate_all_b16x, (float16_t const *)(v_block + ((dim_tile + 2) * 8 + 6) * 32)));
+                svmopa_za32_f16_m(
+                    3, query_predicate_b16x, predicate_all_b16x, probability_column_6_f32x,
+                    svld1_f16(predicate_all_b16x, (float16_t const *)(v_block + ((dim_tile + 3) * 8 + 6) * 32)));
+                svmopa_za32_f16_m(
+                    0, query_predicate_b16x, predicate_all_b16x, probability_column_7_f32x,
+                    svld1_f16(predicate_all_b16x, (float16_t const *)(v_block + ((dim_tile + 0) * 8 + 7) * 32)));
+                svmopa_za32_f16_m(
+                    1, query_predicate_b16x, predicate_all_b16x, probability_column_7_f32x,
+                    svld1_f16(predicate_all_b16x, (float16_t const *)(v_block + ((dim_tile + 1) * 8 + 7) * 32)));
+                svmopa_za32_f16_m(
+                    2, query_predicate_b16x, predicate_all_b16x, probability_column_7_f32x,
+                    svld1_f16(predicate_all_b16x, (float16_t const *)(v_block + ((dim_tile + 2) * 8 + 7) * 32)));
+                svmopa_za32_f16_m(
+                    3, query_predicate_b16x, predicate_all_b16x, probability_column_7_f32x,
+                    svld1_f16(predicate_all_b16x, (float16_t const *)(v_block + ((dim_tile + 3) * 8 + 7) * 32)));
+                for (nk_size_t query_index = 0; query_index < valid_query_count; query_index++) {
+                    svst1_f32(
+                        predicate_all_b32x, output_accumulator + query_index * head_dim_padded + (dim_tile + 0) * 16,
+                        svadd_f32_x(predicate_all_b32x,
+                                    svld1_f32(predicate_all_b32x,
+                                              output_accumulator + query_index * head_dim_padded + (dim_tile + 0) * 16),
+                                    svread_hor_za32_f32_m(svdup_f32(0), predicate_all_b32x, 0, query_index)));
+                    svst1_f32(
+                        predicate_all_b32x, output_accumulator + query_index * head_dim_padded + (dim_tile + 1) * 16,
+                        svadd_f32_x(predicate_all_b32x,
+                                    svld1_f32(predicate_all_b32x,
+                                              output_accumulator + query_index * head_dim_padded + (dim_tile + 1) * 16),
+                                    svread_hor_za32_f32_m(svdup_f32(0), predicate_all_b32x, 1, query_index)));
+                    svst1_f32(
+                        predicate_all_b32x, output_accumulator + query_index * head_dim_padded + (dim_tile + 2) * 16,
+                        svadd_f32_x(predicate_all_b32x,
+                                    svld1_f32(predicate_all_b32x,
+                                              output_accumulator + query_index * head_dim_padded + (dim_tile + 2) * 16),
+                                    svread_hor_za32_f32_m(svdup_f32(0), predicate_all_b32x, 2, query_index)));
+                    svst1_f32(
+                        predicate_all_b32x, output_accumulator + query_index * head_dim_padded + (dim_tile + 3) * 16,
+                        svadd_f32_x(predicate_all_b32x,
+                                    svld1_f32(predicate_all_b32x,
+                                              output_accumulator + query_index * head_dim_padded + (dim_tile + 3) * 16),
+                                    svread_hor_za32_f32_m(svdup_f32(0), predicate_all_b32x, 3, query_index)));
+                }
+            }
+            for (; dim_tile < dim_tile_count; dim_tile++) {
+                svzero_mask_za(nk_sme_zero_za32_tile_0_);
+                svmopa_za32_f16_m(
+                    0, query_predicate_b16x, predicate_all_b16x, probability_column_0_f32x,
+                    svld1_f16(predicate_all_b16x, (float16_t const *)(v_block + (dim_tile * 8 + 0) * 32)));
+                svmopa_za32_f16_m(
+                    0, query_predicate_b16x, predicate_all_b16x, probability_column_1_f32x,
+                    svld1_f16(predicate_all_b16x, (float16_t const *)(v_block + (dim_tile * 8 + 1) * 32)));
+                svmopa_za32_f16_m(
+                    0, query_predicate_b16x, predicate_all_b16x, probability_column_2_f32x,
+                    svld1_f16(predicate_all_b16x, (float16_t const *)(v_block + (dim_tile * 8 + 2) * 32)));
+                svmopa_za32_f16_m(
+                    0, query_predicate_b16x, predicate_all_b16x, probability_column_3_f32x,
+                    svld1_f16(predicate_all_b16x, (float16_t const *)(v_block + (dim_tile * 8 + 3) * 32)));
+                svmopa_za32_f16_m(
+                    0, query_predicate_b16x, predicate_all_b16x, probability_column_4_f32x,
+                    svld1_f16(predicate_all_b16x, (float16_t const *)(v_block + (dim_tile * 8 + 4) * 32)));
+                svmopa_za32_f16_m(
+                    0, query_predicate_b16x, predicate_all_b16x, probability_column_5_f32x,
+                    svld1_f16(predicate_all_b16x, (float16_t const *)(v_block + (dim_tile * 8 + 5) * 32)));
+                svmopa_za32_f16_m(
+                    0, query_predicate_b16x, predicate_all_b16x, probability_column_6_f32x,
+                    svld1_f16(predicate_all_b16x, (float16_t const *)(v_block + (dim_tile * 8 + 6) * 32)));
+                svmopa_za32_f16_m(
+                    0, query_predicate_b16x, predicate_all_b16x, probability_column_7_f32x,
+                    svld1_f16(predicate_all_b16x, (float16_t const *)(v_block + (dim_tile * 8 + 7) * 32)));
+                for (nk_size_t query_index = 0; query_index < valid_query_count; query_index++)
+                    svst1_f32(predicate_all_b32x, output_accumulator + query_index * head_dim_padded + dim_tile * 16,
+                              svadd_f32_x(predicate_all_b32x,
+                                          svld1_f32(predicate_all_b32x,
+                                                    output_accumulator + query_index * head_dim_padded + dim_tile * 16),
+                                          svread_hor_za32_f32_m(svdup_f32(0), predicate_all_b32x, 0, query_index)));
+            }
+        }
+    }
+
+    // Final normalization
+    svfloat32_t final_sum_f32x = svld1_f32(predicate_all_b32x, row_sum);
+    svfloat32_t ones_f32x = svdup_f32(1.0f);
+    svfloat32_t zeros_f32x = svdup_f32(0.0f);
+    svbool_t sum_positive_b32x = svcmpgt_f32(predicate_all_b32x, final_sum_f32x, zeros_f32x);
+    svfloat32_t inv_sum_f32x = svsel_f32(sum_positive_b32x, svdiv_f32_x(predicate_all_b32x, ones_f32x, final_sum_f32x),
+                                         zeros_f32x);
+
+    NK_ALIGN64 nk_f32_t inv_sums[16];
+    svst1_f32(predicate_all_b32x, inv_sums, inv_sum_f32x);
+
+    for (nk_size_t query_index = 0; query_index < valid_query_count; query_index++) {
+        svfloat32_t inv_sum_f32x = svdup_f32(inv_sums[query_index]);
+        for (nk_size_t dim_offset = 0; dim_offset < head_dim; dim_offset += svcntw()) {
+            svbool_t predicate_b32x = svwhilelt_b32_u64(dim_offset, head_dim);
+            svfloat32_t output_f32x = svmul_f32_x(
+                predicate_b32x,
+                svld1_f32(predicate_b32x, output_accumulator + query_index * head_dim_padded + dim_offset),
+                inv_sum_f32x);
+            svfloat16_t output_f16x = svcvt_f16_f32_x(predicate_b32x, output_f32x);
+            nk_size_t store_count = (head_dim - dim_offset) < (nk_size_t)svcntw() ? (head_dim - dim_offset)
+                                                                                  : (nk_size_t)svcntw();
+            svbool_t predicate_b16x = svwhilelt_b16_u64(0u, store_count);
+            svst1_f16(predicate_b16x, (float16_t *)(output + query_index * head_dim + dim_offset), output_f16x);
+        }
+    }
+}
+
+NK_PUBLIC void nk_attention_f16_sme( //
+    nk_f16_t const *q, void const *kv_packed, nk_f16_t *output, nk_size_t num_heads, nk_size_t num_kv_heads,
+    nk_size_t query_len, nk_size_t kv_len, nk_size_t head_dim, nk_f32_t scale) {
+
+    nk_attention_sme_packed_header_t const *header = (nk_attention_sme_packed_header_t const *)kv_packed;
+    nk_size_t head_dim_padded = header->head_dim_padded;
+    nk_size_t dim_tile_count = header->reserved[0];
+    nk_size_t kv_blocks = (kv_len + 15) / 16;
+    // K and V both use interleaved format: kv_blocks * 16 * head_dim_padded elements per head
+    nk_size_t kv_head_stride = kv_blocks * 16 * head_dim_padded;
+
+    nk_f16_t const *k_packed = (nk_f16_t const *)((char const *)kv_packed + header->k_offset);
+    nk_f16_t const *v_packed = (nk_f16_t const *)((char const *)kv_packed + header->v_offset);
+
+    nk_size_t group_size = (num_kv_heads > 0) ? num_heads / num_kv_heads : 1;
+
+    nk_sme_start_streaming_();
+    for (nk_size_t q_head = 0; q_head < num_heads; q_head++) {
+        nk_size_t kv_head = q_head / group_size;
+
+        nk_f16_t const *q_ptr = q + q_head * query_len * head_dim;
+        nk_f16_t const *k_ptr = k_packed + kv_head * kv_head_stride;
+        nk_f16_t const *v_ptr = v_packed + kv_head * kv_head_stride;
+        nk_f16_t *out_ptr = output + q_head * query_len * head_dim;
+
+        for (nk_size_t q_start = 0; q_start < query_len; q_start += 16) {
+            nk_size_t q_block_len = (q_start + 16 < query_len) ? 16 : (query_len - q_start);
+
+            nk_attention_f16_sme_streaming_(q_ptr + q_start * head_dim, k_ptr, v_ptr, out_ptr + q_start * head_dim,
+                                            q_block_len, kv_len, head_dim, head_dim_padded, dim_tile_count, scale);
+        }
+    }
+    nk_sme_stop_streaming_();
+}
+
+NK_PUBLIC void nk_attention_causal_bf16_sme(nk_bf16_t const *q, void const *kv_packed, nk_bf16_t *output,
+                                            nk_size_t num_heads, nk_size_t num_kv_heads, nk_size_t query_len,
+                                            nk_size_t kv_len, nk_size_t head_dim, nk_f32_t scale) {
+    // TODO: Implement proper causal masking with block skipping
+    // For now, delegate to full attention (correct for decode where query_len=1)
+    nk_attention_bf16_sme(q, kv_packed, output, num_heads, num_kv_heads, query_len, kv_len, head_dim, scale);
+}
+
+NK_PUBLIC void nk_attention_causal_f16_sme(nk_f16_t const *q, void const *kv_packed, nk_f16_t *output,
+                                           nk_size_t num_heads, nk_size_t num_kv_heads, nk_size_t query_len,
+                                           nk_size_t kv_len, nk_size_t head_dim, nk_f32_t scale) {
+    // TODO: Implement proper causal masking with block skipping
+    // For now, delegate to full attention (correct for decode where query_len=1)
+    nk_attention_f16_sme(q, kv_packed, output, num_heads, num_kv_heads, query_len, kv_len, head_dim, scale);
+}
+
+#if defined(__clang__)
+#pragma clang attribute pop
+#elif defined(__GNUC__)
+#pragma GCC pop_options
+#endif
+
+#if defined(__cplusplus)
+} // extern "C"
+#endif
+
+#endif // NK_TARGET_SME
+#endif // NK_TARGET_ARM64_
+#endif // NK_ATTENTION_SME_H