numkong 7.0.0

package/include/numkong/attention/sapphireamx.h

@@ -0,0 +1,1361 @@
+/**
+ *  @brief FlashAttention-style kernels for Intel Sapphire Rapids AMX.
+ *  @file include/numkong/attention/sapphireamx.h
+ *  @author Ash Vardanian
+ *  @date January 5, 2026
+ *
+ *  @sa include/numkong/attention.h
+ *
+ *  This file implements FlashAttention-2 style scaled dot-product attention (SDPA) optimized
+ *  for Intel AMX instructions on Sapphire Rapids CPUs. The kernel computes:
+ *
+ *      O = softmax(Q × Kᵀ / √d) × V
+ *
+ *  Key features:
+ *  - Online softmax: Mathematically exact, processes KV blocks incrementally
+ *  - Pre-packed KV cache: Amortizes packing cost for repeated inference
+ *  - GQA/MQA support: Different num_heads and num_kv_heads for grouped-query attention
+ *  - Causal masking: Optional masking for autoregressive generation
+ *
+ *  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
+ *
+ *  Performance comparison with H100 FlashAttention-2:
+ *  - H100 SXM5: ~335 TFLOPS (35% of 989 TFLOPS peak), 80GB HBM3
+ *  - 100-core SPR: ~40 TFLOPS with FlashAttention (13% of 300 TFLOPS peak)
+ *  - CPU advantage: 512GB-2TB DDR5 vs 80GB HBM → supports 10-25⨯ longer contexts
+ *
+ *  Expected performance per core:
+ *  - Decode (query_len=1, kv_len=4K): 350-450 GOPS (softmax bound)
+ *  - Prefill (query_len=64, kv_len=4K): 450-550 GOPS (better AMX utilization)
+ *  - Long context (kv_len=64K+): 250-350 GOPS (memory bandwidth bound)
+ *
+ *  Block sizes:
+ *  - Bᵣ = 16 (query block rows, matches AMX tile height)
+ *  - Bᶜ = 16 (KV block columns, fits 16×16 scores in 16 ZMM registers)
+ *
+ *  Algorithm (FlashAttention-2 style):
+ *  For each query block:
+ *    Initialize O = 0, rowsum = 0, rowmax = -∞
+ *    For each KV block:
+ *      S = Q × Kᵀ using AMX TDPBF16PS
+ *      Apply online softmax: rescale old values, accumulate new
+ *      O = rescale(O) + P × V using AMX
+ *    Finalize: normalize O by row sums
+ *
+ *  @section sapphireamx_attention_instructions Relevant Instructions
+ *
+ *      Intrinsic                   Instruction                     Sapphire
+ *      _tile_dpbf16ps              TDPBF16PS (TMM, TMM, TMM)       ~16cy (16x16x32 BF16)
+ *      _tile_dpbssd                TDPBSSD (TMM, TMM, TMM)         ~16cy (16x16x64 INT8)
+ *      _tile_loadd                 TILELOADD (TMM, MEM)            ~10cy @ p23
+ *      _tile_stored                TILESTORED (MEM, TMM)           ~10cy @ p4
+ *      _tile_zero                  TILEZERO (TMM)                  ~1cy
+ *      _mm512_fmadd_ps             VFMADD (ZMM, ZMM, ZMM)          4cy @ p05
+ *      _mm512_mul_ps               VMULPS (ZMM, ZMM, ZMM)          4cy @ p05
+ *      _mm512_max_ps               VMAXPS (ZMM, ZMM, ZMM)          4cy @ p05
+ *      _mm512_reduce_max_ps        (pseudo: VHADDPS chain)         ~8cy
+ *      _mm512_reduce_add_ps        (pseudo: VHADDPS chain)         ~8cy
+ */
+#ifndef NK_ATTENTION_SAPPHIREAMX_H
+#define NK_ATTENTION_SAPPHIREAMX_H
+
+#if NK_TARGET_X86_
+#if NK_TARGET_SAPPHIREAMX
+
+#include "numkong/types.h"
+#include "numkong/dots/sapphireamx.h"
+
+#if defined(__cplusplus)
+extern "C" {
+#endif
+
+#if defined(__clang__)
+#pragma clang attribute push(                                                                                   \
+    __attribute__((target("avx2,avx512f,avx512vl,avx512bw,avx512dq,avx512fp16,avx512bf16,f16c,fma,bmi,bmi2"))), \
+    apply_to = function)
+#elif defined(__GNUC__)
+#pragma GCC push_options
+#pragma GCC target("avx2", "avx512f", "avx512vl", "avx512bw", "avx512dq", "avx512fp16", "avx512bf16", "f16c", "fma", \
+                   "bmi", "bmi2")
+#endif
+
+/**
+ *  @brief Packed KV cache header for attention (64-byte aligned).
+ *
+ *  Layout in memory:
+ *  [header: 64 bytes][K tiles: variable][V tiles: variable]
+ *
+ *  K and V are packed in AMX tile format for efficient loading.
+ */
+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 AMX tiles
+    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 tiles from header start
+    nk_u32_t v_offset;        ///< Byte offset to V tiles from header start
+    nk_u32_t reserved[9];     ///< Pad to 64 bytes
+} nk_attention_kv_packed_header_t;
+
+/**
+ *  @brief Fast exp approximation for AVX-512.
+ *
+ *  Uses Cody-Waite range reduction + Remez minimax polynomial.
+ *  Accuracy: max error < 1 ULP for x ∈ [-87.3, 88.7] (float range).
+ *  Performance: ~15-20 cycles for 16 floats.
+ */
+
+/**
+ *  @brief Fast vectorized exp(x) approximation using AVX-512.
+ *
+ *  Algorithm:
+ *  1. Range reduction: x = n × ln(2) + r, where |r| < ln(2)/2
+ *  2. Polynomial approximation: exp(r) ≈ 1 + r + r²/2 + ... (degree 6)
+ *  3. Reconstruction: exp(x) = 2ⁿ × exp(r)
+ *
+ *  @param x Input vector (16 floats)
+ *  @return exp(x) for each element
+ */
+NK_INTERNAL __m512 nk_exp_ps_avx512_(__m512 x) {
+    // Constants for Cody-Waite range reduction
+    const __m512 log2e = _mm512_set1_ps(1.4426950408889634f);
+    const __m512 ln2_hi = _mm512_set1_ps(0.693145751953125f);
+    const __m512 ln2_lo = _mm512_set1_ps(1.42860682030941723212e-6f);
+
+    // Clamp to avoid overflow/underflow
+    const __m512 max_x = _mm512_set1_ps(88.3762626647949f);
+    const __m512 min_x = _mm512_set1_ps(-87.3365447504021f);
+    x = _mm512_max_ps(_mm512_min_ps(x, max_x), min_x);
+
+    // n = round(x / ln(2))
+    __m512 n = _mm512_roundscale_ps(_mm512_mul_ps(x, log2e), _MM_FROUND_TO_NEAREST_INT);
+
+    // r = x - n × ln(2) using Cody-Waite for precision
+    __m512 r = _mm512_fnmadd_ps(n, ln2_hi, x);
+    r = _mm512_fnmadd_ps(n, ln2_lo, r);
+
+    // Polynomial approximation for exp(r): Remez minimax degree 6
+    // Coefficients optimized for [-ln(2)/2, ln(2)/2]
+    __m512 p = _mm512_set1_ps(1.9875691500e-4f);
+    p = _mm512_fmadd_ps(p, r, _mm512_set1_ps(1.3981999507e-3f));
+    p = _mm512_fmadd_ps(p, r, _mm512_set1_ps(8.3334519073e-3f));
+    p = _mm512_fmadd_ps(p, r, _mm512_set1_ps(4.1665858030e-2f));
+    p = _mm512_fmadd_ps(p, r, _mm512_set1_ps(1.6666665459e-1f));
+    p = _mm512_fmadd_ps(p, r, _mm512_set1_ps(5.0000001201e-1f));
+    p = _mm512_fmadd_ps(p, r, _mm512_set1_ps(1.0f));
+    p = _mm512_fmadd_ps(p, r, _mm512_set1_ps(1.0f));
+
+    // Reconstruct: exp(x) = 2ⁿ × exp(r)
+    // 2ⁿ via IEEE 754 exponent manipulation
+    __m512i ni = _mm512_cvtps_epi32(n);
+    ni = _mm512_add_epi32(ni, _mm512_set1_epi32(127));
+    ni = _mm512_slli_epi32(ni, 23);
+    __m512 pow2n = _mm512_castsi512_ps(ni);
+
+    return _mm512_mul_ps(p, pow2n);
+}
+
+/**
+ *  @brief Faster exp(x) approximation using degree-4 polynomial.
+ *
+ *  Trades accuracy for speed: ~0.1% relative error (vs <0.001% for degree-6).
+ *  This is acceptable for softmax where:
+ *  - Probabilities sum to 1 (normalization absorbs errors)
+ *  - Relative ranking matters more than absolute values
+ *
+ *  Performance: ~12-15 cycles for 16 floats (vs ~18-22 for degree-6)
+ *
+ *  @param x Input vector (16 floats)
+ *  @return exp(x) approximation
+ */
+NK_INTERNAL __m512 nk_exp_ps_fast_avx512_(__m512 x) {
+    // Constants for Cody-Waite range reduction
+    const __m512 log2e = _mm512_set1_ps(1.4426950408889634f);
+    const __m512 ln2_hi = _mm512_set1_ps(0.693145751953125f);
+    const __m512 ln2_lo = _mm512_set1_ps(1.42860682030941723212e-6f);
+
+    // Clamp to avoid overflow/underflow (same as accurate version)
+    const __m512 max_x = _mm512_set1_ps(88.3762626647949f);
+    const __m512 min_x = _mm512_set1_ps(-87.3365447504021f);
+    x = _mm512_max_ps(_mm512_min_ps(x, max_x), min_x);
+
+    // n = round(x / ln(2))
+    __m512 n = _mm512_roundscale_ps(_mm512_mul_ps(x, log2e), _MM_FROUND_TO_NEAREST_INT);
+
+    // r = x - n × ln(2) using Cody-Waite for precision
+    __m512 r = _mm512_fnmadd_ps(n, ln2_hi, x);
+    r = _mm512_fnmadd_ps(n, ln2_lo, r);
+
+    // Polynomial approximation for exp(r): degree 4
+    // Optimized coefficients for [-ln(2)/2, ln(2)/2]
+    // exp(r) ≈ 1 + r + r²/2 + r³/6 + r⁴/24
+    // Using Horner form: ((c₄ × r + c₃) × r + c₂) × r + c₁) × r + c₀
+    __m512 p = _mm512_set1_ps(4.1666666667e-2f);                 // 1/24
+    p = _mm512_fmadd_ps(p, r, _mm512_set1_ps(1.6666666667e-1f)); // 1/6
+    p = _mm512_fmadd_ps(p, r, _mm512_set1_ps(5.0000000000e-1f)); // 1/2
+    p = _mm512_fmadd_ps(p, r, _mm512_set1_ps(1.0f));             // 1
+    p = _mm512_fmadd_ps(p, r, _mm512_set1_ps(1.0f));             // 1
+
+    // Reconstruct: exp(x) = 2ⁿ × exp(r)
+    __m512i ni = _mm512_cvtps_epi32(n);
+    ni = _mm512_add_epi32(ni, _mm512_set1_epi32(127));
+    ni = _mm512_slli_epi32(ni, 23);
+    __m512 pow2n = _mm512_castsi512_ps(ni);
+
+    return _mm512_mul_ps(p, pow2n);
+}
+
+/**
+ *  @brief Online softmax primitives.
+ *
+ *  These implement the online softmax algorithm from FlashAttention.
+ *  Key insight: softmax can be computed incrementally by tracking:
+ *  - m: running maximum (for numerical stability)
+ *  - l: running sum of exp(x - m)
+ *
+ *  When a new block arrives with larger values:
+ *  - Rescale old sum: l = l × exp(m_old - m_new)
+ *  - Add new contributions: l += Σ exp(x_new - m_new)
+ */
+
+/**
+ *  @brief State for online softmax computation.
+ *
+ *  Tracks per-row running maximum and sum for 16 rows.
+ */
+typedef struct {
+    __m512 row_max; ///< Running max per row (16 values)
+    __m512 row_sum; ///< Running sum of exp(x - max) per row
+} nk_attention_softmax_row_state_t;
+
+/**
+ *  @brief Update softmax state with Bᶜ=32 score block (optimized).
+ *
+ *  Computes online softmax for 16×32 score block using AVX-512.
+ *  Optimizations:
+ *  - Process 4 rows at a time for better ILP
+ *  - Keep scaled scores in registers to avoid reloading
+ *  - Vectorized row sum accumulation
+ */
+NK_INTERNAL void nk_attention_softmax_update_bc32_(nk_attention_softmax_row_state_t *state,
+                                                   nk_f32_t const *scores, // [16, 32] score block
+                                                   nk_f32_t scale,
+                                                   nk_f32_t *weights_out) { // [16, 32] output weights
+
+    __m512 scale_v = _mm512_set1_ps(scale);
+
+    // Load and scale all scores, compute per-row max
+    // Store in temporary arrays to avoid register pressure
+    __m512 s_scaled[16][2];
+    NK_ALIGN64 float row_maxes[16];
+
+    // Process 4 rows at a time for ILP
+    for (int i = 0; i < 16; i += 4) {
+        // Row i
+        s_scaled[i][0] = _mm512_mul_ps(_mm512_load_ps(scores + i * 32 + 0), scale_v);
+        s_scaled[i][1] = _mm512_mul_ps(_mm512_load_ps(scores + i * 32 + 16), scale_v);
+        __m512 m0 = _mm512_max_ps(s_scaled[i][0], s_scaled[i][1]);
+
+        // Row i+1
+        s_scaled[i + 1][0] = _mm512_mul_ps(_mm512_load_ps(scores + (i + 1) * 32 + 0), scale_v);
+        s_scaled[i + 1][1] = _mm512_mul_ps(_mm512_load_ps(scores + (i + 1) * 32 + 16), scale_v);
+        __m512 m1 = _mm512_max_ps(s_scaled[i + 1][0], s_scaled[i + 1][1]);
+
+        // Row i+2
+        s_scaled[i + 2][0] = _mm512_mul_ps(_mm512_load_ps(scores + (i + 2) * 32 + 0), scale_v);
+        s_scaled[i + 2][1] = _mm512_mul_ps(_mm512_load_ps(scores + (i + 2) * 32 + 16), scale_v);
+        __m512 m2 = _mm512_max_ps(s_scaled[i + 2][0], s_scaled[i + 2][1]);
+
+        // Row i+3
+        s_scaled[i + 3][0] = _mm512_mul_ps(_mm512_load_ps(scores + (i + 3) * 32 + 0), scale_v);
+        s_scaled[i + 3][1] = _mm512_mul_ps(_mm512_load_ps(scores + (i + 3) * 32 + 16), scale_v);
+        __m512 m3 = _mm512_max_ps(s_scaled[i + 3][0], s_scaled[i + 3][1]);
+
+        // Reduce to scalar max
+        row_maxes[i] = _mm512_reduce_max_ps(m0);
+        row_maxes[i + 1] = _mm512_reduce_max_ps(m1);
+        row_maxes[i + 2] = _mm512_reduce_max_ps(m2);
+        row_maxes[i + 3] = _mm512_reduce_max_ps(m3);
+    }
+
+    __m512 row_max_new = _mm512_load_ps(row_maxes);
+    __m512 old_max = state->row_max;
+    __m512 new_max = _mm512_max_ps(old_max, row_max_new);
+
+    // Rescale old sum
+    __m512 correction = nk_exp_ps_avx512_(_mm512_sub_ps(old_max, new_max));
+    __m512 new_sum = _mm512_mul_ps(state->row_sum, correction);
+
+    // Compute P = exp(S - new_max) and accumulate sums
+    NK_ALIGN64 float new_max_arr[16];
+    NK_ALIGN64 float row_sums[16];
+    _mm512_store_ps(new_max_arr, new_max);
+
+    // Process rows
+    for (int i = 0; i < 16; i += 2) {
+        __m512 max_i = _mm512_set1_ps(new_max_arr[i]);
+        __m512 max_i1 = _mm512_set1_ps(new_max_arr[i + 1]);
+
+        // Row i
+        __m512 p0_i = nk_exp_ps_avx512_(_mm512_sub_ps(s_scaled[i][0], max_i));
+        __m512 p1_i = nk_exp_ps_avx512_(_mm512_sub_ps(s_scaled[i][1], max_i));
+        _mm512_store_ps(weights_out + i * 32 + 0, p0_i);
+        _mm512_store_ps(weights_out + i * 32 + 16, p1_i);
+        row_sums[i] = _mm512_reduce_add_ps(p0_i) + _mm512_reduce_add_ps(p1_i);
+
+        // Row i+1
+        __m512 p0_i1 = nk_exp_ps_avx512_(_mm512_sub_ps(s_scaled[i + 1][0], max_i1));
+        __m512 p1_i1 = nk_exp_ps_avx512_(_mm512_sub_ps(s_scaled[i + 1][1], max_i1));
+        _mm512_store_ps(weights_out + (i + 1) * 32 + 0, p0_i1);
+        _mm512_store_ps(weights_out + (i + 1) * 32 + 16, p1_i1);
+        row_sums[i + 1] = _mm512_reduce_add_ps(p0_i1) + _mm512_reduce_add_ps(p1_i1);
+    }
+
+    // Add row sums to running sum vectorially
+    new_sum = _mm512_add_ps(new_sum, _mm512_load_ps(row_sums));
+
+    state->row_max = new_max;
+    state->row_sum = new_sum;
+}
+
+/**
+ *  @brief Fast softmax update using degree-4 exp polynomial.
+ *
+ *  Same algorithm as nk_attention_softmax_update_bc32_ but uses faster exp.
+ *  Trades ~0.1% accuracy for ~20% performance improvement.
+ *
+ *  Use this for inference where throughput matters more than last-bit accuracy.
+ */
+NK_INTERNAL void nk_attention_softmax_update_bc32_fast_(nk_attention_softmax_row_state_t *state,
+                                                        nk_f32_t const *scores, // [16, 32] score block
+                                                        nk_f32_t scale,
+                                                        nk_f32_t *weights_out) { // [16, 32] output weights
+
+    __m512 scale_v = _mm512_set1_ps(scale);
+
+    // Load and scale all scores, compute per-row max
+    __m512 s_scaled[16][2];
+    NK_ALIGN64 float row_maxes[16];
+
+    // Process 4 rows at a time for ILP
+    for (int i = 0; i < 16; i += 4) {
+        s_scaled[i][0] = _mm512_mul_ps(_mm512_load_ps(scores + i * 32 + 0), scale_v);
+        s_scaled[i][1] = _mm512_mul_ps(_mm512_load_ps(scores + i * 32 + 16), scale_v);
+        __m512 m0 = _mm512_max_ps(s_scaled[i][0], s_scaled[i][1]);
+
+        s_scaled[i + 1][0] = _mm512_mul_ps(_mm512_load_ps(scores + (i + 1) * 32 + 0), scale_v);
+        s_scaled[i + 1][1] = _mm512_mul_ps(_mm512_load_ps(scores + (i + 1) * 32 + 16), scale_v);
+        __m512 m1 = _mm512_max_ps(s_scaled[i + 1][0], s_scaled[i + 1][1]);
+
+        s_scaled[i + 2][0] = _mm512_mul_ps(_mm512_load_ps(scores + (i + 2) * 32 + 0), scale_v);
+        s_scaled[i + 2][1] = _mm512_mul_ps(_mm512_load_ps(scores + (i + 2) * 32 + 16), scale_v);
+        __m512 m2 = _mm512_max_ps(s_scaled[i + 2][0], s_scaled[i + 2][1]);
+
+        s_scaled[i + 3][0] = _mm512_mul_ps(_mm512_load_ps(scores + (i + 3) * 32 + 0), scale_v);
+        s_scaled[i + 3][1] = _mm512_mul_ps(_mm512_load_ps(scores + (i + 3) * 32 + 16), scale_v);
+        __m512 m3 = _mm512_max_ps(s_scaled[i + 3][0], s_scaled[i + 3][1]);
+
+        row_maxes[i] = _mm512_reduce_max_ps(m0);
+        row_maxes[i + 1] = _mm512_reduce_max_ps(m1);
+        row_maxes[i + 2] = _mm512_reduce_max_ps(m2);
+        row_maxes[i + 3] = _mm512_reduce_max_ps(m3);
+    }
+
+    __m512 row_max_new = _mm512_load_ps(row_maxes);
+    __m512 old_max = state->row_max;
+    __m512 new_max = _mm512_max_ps(old_max, row_max_new);
+
+    // Rescale old sum using fast exp
+    __m512 correction = nk_exp_ps_fast_avx512_(_mm512_sub_ps(old_max, new_max));
+    __m512 new_sum = _mm512_mul_ps(state->row_sum, correction);
+
+    // Compute P = exp(S - new_max) using fast exp
+    NK_ALIGN64 float new_max_arr[16];
+    NK_ALIGN64 float row_sums[16];
+    _mm512_store_ps(new_max_arr, new_max);
+
+    // Process rows with fast exp
+    for (int i = 0; i < 16; i += 2) {
+        __m512 max_i = _mm512_set1_ps(new_max_arr[i]);
+        __m512 max_i1 = _mm512_set1_ps(new_max_arr[i + 1]);
+
+        // Row i
+        __m512 p0_i = nk_exp_ps_fast_avx512_(_mm512_sub_ps(s_scaled[i][0], max_i));
+        __m512 p1_i = nk_exp_ps_fast_avx512_(_mm512_sub_ps(s_scaled[i][1], max_i));
+        _mm512_store_ps(weights_out + i * 32 + 0, p0_i);
+        _mm512_store_ps(weights_out + i * 32 + 16, p1_i);
+        row_sums[i] = _mm512_reduce_add_ps(p0_i) + _mm512_reduce_add_ps(p1_i);
+
+        // Row i+1
+        __m512 p0_i1 = nk_exp_ps_fast_avx512_(_mm512_sub_ps(s_scaled[i + 1][0], max_i1));
+        __m512 p1_i1 = nk_exp_ps_fast_avx512_(_mm512_sub_ps(s_scaled[i + 1][1], max_i1));
+        _mm512_store_ps(weights_out + (i + 1) * 32 + 0, p0_i1);
+        _mm512_store_ps(weights_out + (i + 1) * 32 + 16, p1_i1);
+        row_sums[i + 1] = _mm512_reduce_add_ps(p0_i1) + _mm512_reduce_add_ps(p1_i1);
+    }
+
+    new_sum = _mm512_add_ps(new_sum, _mm512_load_ps(row_sums));
+
+    state->row_max = new_max;
+    state->row_sum = new_sum;
+}
+
+/**
+ *  @brief Initialize online softmax state.
+ */
+NK_INTERNAL void nk_attention_softmax_init_(nk_attention_softmax_row_state_t *state) {
+    state->row_max = _mm512_set1_ps(NK_F32_MIN);
+    state->row_sum = _mm512_setzero_ps();
+}
+
+/**
+ *  @brief Update softmax state with new score block and compute attention weights.
+ *
+ *  For a 16×16 score block S[16][16]:
+ *  1. Compute row-wise max of S
+ *  2. Update running max: newₘₐₓ = max(oldₘₐₓ, blockₘₐₓ)
+ *  3. Rescale old sum: oldₛᵤₘ × = exp(oldₘₐₓ - newₘₐₓ)
+ *  4. Compute P = exp(S - newₘₐₓ), store for P × V
+ *  5. Update sum: newₛᵤₘ = oldₛᵤₘ + row_sum(P)
+ *
+ *  @param state        Running softmax state (updated in place)
+ *  @param scores       16×16 score block in row-major order (256 floats)
+ *  @param scale        Scaling factor (1/√head_dim)
+ *  @param weights_out  Output: 16×16 attention weights P (pre-softmax normalized)
+ */
+NK_INTERNAL void nk_attention_softmax_update_(nk_attention_softmax_row_state_t *state, nk_f32_t const *scores,
+                                              nk_f32_t scale, nk_f32_t *weights_out) {
+
+    __m512 scale_v = _mm512_set1_ps(scale);
+
+    // Load scores into 16 ZMM registers (one per row)
+    __m512 s[16];
+    for (int i = 0; i < 16; i++) { s[i] = _mm512_mul_ps(_mm512_load_ps(scores + i * 16), scale_v); }
+
+    // Per-row max (each row has 16 elements, we need max across those 16)
+    // _mm512_reduce_max_ps returns a float scalar
+    NK_ALIGN64 float row_maxes[16];
+    for (int i = 0; i < 16; i++) { row_maxes[i] = _mm512_reduce_max_ps(s[i]); }
+    __m512 row_max_new = _mm512_load_ps(row_maxes);
+
+    // Update running max
+    __m512 old_max = state->row_max;
+    __m512 new_max = _mm512_max_ps(old_max, row_max_new);
+
+    // Rescale old sum: l = l × exp(oldₘₐₓ - newₘₐₓ)
+    __m512 correction = nk_exp_ps_avx512_(_mm512_sub_ps(old_max, new_max));
+    __m512 old_sum_rescaled = _mm512_mul_ps(state->row_sum, correction);
+
+    // Compute P = exp(S - newₘₐₓ) for each row, accumulate sum
+    __m512 new_sum = old_sum_rescaled;
+    float new_max_arr[16];
+    _mm512_store_ps(new_max_arr, new_max);
+
+    for (int i = 0; i < 16; i++) {
+        __m512 max_broadcast = _mm512_set1_ps(new_max_arr[i]);
+        __m512 p = nk_exp_ps_avx512_(_mm512_sub_ps(s[i], max_broadcast));
+        _mm512_store_ps(weights_out + i * 16, p);
+
+        // Add row sum to running sum (at position i)
+        float row_sum = _mm512_reduce_add_ps(p);
+        new_sum = _mm512_mask_add_ps(new_sum, 1u << i, new_sum, _mm512_set1_ps(row_sum));
+    }
+
+    state->row_max = new_max;
+    state->row_sum = new_sum;
+}
+
+/**
+ *  @brief Rescale output accumulator when max changes.
+ *
+ *  When processing a new KV block with larger scores, previous O accumulator
+ *  needs rescaling: O = O × exp(oldₘₐₓ - newₘₐₓ)
+ *
+ *  @param output       Output accumulator [16][head_dim] in F32
+ *  @param head_dim     Head dimension
+ *  @param old_max      Previous running max per row (16 values)
+ *  @param new_max      New running max per row (16 values)
+ */
+NK_INTERNAL void nk_attention_rescale_output_(nk_f32_t *output, nk_size_t head_dim, __m512 old_max, __m512 new_max) {
+
+    __m512 correction = nk_exp_ps_avx512_(_mm512_sub_ps(old_max, new_max));
+    float corr_arr[16];
+    _mm512_store_ps(corr_arr, correction);
+
+    for (nk_size_t row = 0; row < 16; row++) {
+        __m512 corr_v = _mm512_set1_ps(corr_arr[row]);
+        for (nk_size_t col = 0; col < head_dim; col += 16) {
+            __m512 o = _mm512_load_ps(output + row * head_dim + col);
+            o = _mm512_mul_ps(o, corr_v);
+            _mm512_store_ps(output + row * head_dim + col, o);
+        }
+    }
+}
+
+NK_PUBLIC nk_size_t nk_attention_kv_packed_size_sapphireamx(nk_size_t num_kv_heads, nk_size_t head_dim,
+                                                            nk_size_t max_seq_len) {
+
+    // Pad head_dim to multiple of 32 for AMX tiles
+    nk_size_t head_dim_padded = nk_size_round_up_to_multiple_(head_dim, 32);
+
+    // Each head: seq_len × head_dim_padded BF16 values
+    // Packed in AMX tile format: 16-row tiles with pair-interleaving
+    nk_size_t tiles_per_head_col = nk_size_divide_round_up_(max_seq_len, 16);
+    nk_size_t tiles_per_head_depth = head_dim_padded / 32;
+    nk_size_t bytes_per_head = tiles_per_head_col * tiles_per_head_depth * 1024; // 1KB per tile
+
+    // K and V each have num_kv_heads heads
+    nk_size_t k_size = num_kv_heads * bytes_per_head;
+    nk_size_t v_size = num_kv_heads * bytes_per_head;
+
+    // Header + K + V, all 64-byte aligned
+    return sizeof(nk_attention_kv_packed_header_t) + k_size + v_size;
+}
+
+NK_PUBLIC void nk_attention_pack_k_sapphireamx(nk_bf16_t const *k, void *kv_packed, nk_size_t num_kv_heads,
+                                               nk_size_t seq_len, nk_size_t head_dim) {
+
+    nk_attention_kv_packed_header_t *header = (nk_attention_kv_packed_header_t *)kv_packed;
+
+    // Initialize header
+    nk_size_t head_dim_padded = nk_size_round_up_to_multiple_(head_dim, 32);
+    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_kv_packed_header_t);
+
+    nk_bf16_t *k_packed = (nk_bf16_t *)((char *)kv_packed + header->k_offset);
+
+    // For Q × Kᵀ, K acts as B matrix but transposed
+    // K[h, s, d] → Kᵀ[h, d, s]
+    // Pack Kᵀ into AMX B tile format with pair-interleaving
+
+    nk_size_t tiles_per_seq = nk_size_divide_round_up_(seq_len, 16);
+    nk_size_t tiles_per_depth = head_dim_padded / 32;
+    nk_size_t tile_size = 512; // BF16 elements per tile
+
+    for (nk_size_t h = 0; h < num_kv_heads; h++) {
+        nk_bf16_t const *k_head = k + h * seq_len * head_dim;
+        nk_bf16_t *k_head_packed = k_packed + h * tiles_per_seq * tiles_per_depth * tile_size;
+
+        // Pack tiles: iterate over seq_len tiles (columns of Kᵀ) and depth tiles
+        for (nk_size_t seq_tile = 0; seq_tile < tiles_per_seq; seq_tile++) {
+            nk_size_t seq_start = seq_tile * 16;
+            nk_size_t valid_seq = (seq_start + 16 <= seq_len) ? 16 : (seq_len - seq_start);
+
+            for (nk_size_t depth_tile = 0; depth_tile < tiles_per_depth; depth_tile++) {
+                nk_size_t depth_start = depth_tile * 32;
+                nk_size_t valid_depth = (depth_start + 32 <= head_dim) ? 32 : (head_dim - depth_start);
+
+                // Tile index in packed format
+                nk_size_t tile_idx = seq_tile * tiles_per_depth + depth_tile;
+                nk_bf16_t *tile_ptr = k_head_packed + tile_idx * tile_size;
+
+                // Pack with pair-interleaving for TDPBF16PS
+                // B tile layout: data[depth/2][col][depth%2]
+                // For Kᵀ: depth is original head_dim, col is original seq position
+                for (nk_size_t d = 0; d < 32; d += 2) {
+                    for (nk_size_t s = 0; s < 16; s++) {
+                        nk_size_t dst_idx = (d / 2) * 32 + s * 2;
+
+                        // K[h, seq_start + s, depth_start + d] and K[h, seq_start + s, depth_start + d + 1]
+                        nk_bf16_t v0 = 0, v1 = 0;
+                        if (s < valid_seq && d < valid_depth) {
+                            v0 = k_head[(seq_start + s) * head_dim + depth_start + d];
+                        }
+                        if (s < valid_seq && d + 1 < valid_depth) {
+                            v1 = k_head[(seq_start + s) * head_dim + depth_start + d + 1];
+                        }
+
+                        tile_ptr[dst_idx] = v0;
+                        tile_ptr[dst_idx + 1] = v1;
+                    }
+                }
+            }
+        }
+    }
+
+    // Calculate V offset
+    nk_size_t k_size = num_kv_heads * tiles_per_seq * tiles_per_depth * tile_size * sizeof(nk_bf16_t);
+    header->v_offset = header->k_offset + (nk_u32_t)k_size;
+}
+
+NK_PUBLIC void nk_attention_pack_v_sapphireamx(nk_bf16_t const *v, void *kv_packed, nk_size_t num_kv_heads,
+                                               nk_size_t seq_len, nk_size_t head_dim) {
+
+    nk_attention_kv_packed_header_t *header = (nk_attention_kv_packed_header_t *)kv_packed;
+    nk_size_t head_dim_padded = header->head_dim_padded;
+
+    nk_bf16_t *v_packed = (nk_bf16_t *)((char *)kv_packed + header->v_offset);
+
+    // For P @ V, P is [query_len, seq_len], V is [seq_len, head_dim]
+    // V acts as B matrix: pack with seq_len as "depth", head_dim as "columns"
+
+    nk_size_t tiles_per_seq = nk_size_divide_round_up_(seq_len, 32);          // seq_len is depth for V
+    nk_size_t tiles_per_head = nk_size_divide_round_up_(head_dim_padded, 16); // head_dim is columns
+    nk_size_t tile_size = 512;
+
+    for (nk_size_t h = 0; h < num_kv_heads; h++) {
+        nk_bf16_t const *v_head = v + h * seq_len * head_dim;
+        nk_bf16_t *v_head_packed = v_packed + h * tiles_per_seq * tiles_per_head * tile_size;
+
+        for (nk_size_t seq_tile = 0; seq_tile < tiles_per_seq; seq_tile++) {
+            nk_size_t seq_start = seq_tile * 32;
+            nk_size_t valid_seq = (seq_start + 32 <= seq_len) ? 32 : (seq_len - seq_start);
+
+            for (nk_size_t head_tile = 0; head_tile < tiles_per_head; head_tile++) {
+                nk_size_t head_start = head_tile * 16;
+                nk_size_t valid_head = (head_start + 16 <= head_dim) ? 16 : (head_dim - head_start);
+
+                nk_size_t tile_idx = seq_tile * tiles_per_head + head_tile;
+                nk_bf16_t *tile_ptr = v_head_packed + tile_idx * tile_size;
+
+                // Pack with pair-interleaving
+                // B tile: data[depth/2][col][depth%2] where depth=seq, col=head_dim
+                for (nk_size_t s = 0; s < 32; s += 2) {
+                    for (nk_size_t d = 0; d < 16; d++) {
+                        nk_size_t dst_idx = (s / 2) * 32 + d * 2;
+
+                        nk_bf16_t v0 = 0, v1 = 0;
+                        if (s < valid_seq && d < valid_head) {
+                            v0 = v_head[(seq_start + s) * head_dim + head_start + d];
+                        }
+                        if (s + 1 < valid_seq && d < valid_head) {
+                            v1 = v_head[(seq_start + s + 1) * head_dim + head_start + d];
+                        }
+
+                        tile_ptr[dst_idx] = v0;
+                        tile_ptr[dst_idx + 1] = v1;
+                    }
+                }
+            }
+        }
+    }
+}
+
+/**
+ *  @brief Extract K block from packed format: Kᵀ[head_dim, Bᶜ] for a given kv_block.
+ *
+ *  K is packed as Kᵀ for Q × Kᵀ, with pair-interleaving.
+ *  Output is in row-major F32 format: k_out[d × Bᶜ + kᵢ] = Kᵀ[d, kᵢ]
+ */
+NK_INTERNAL void nk_attention_extract_k_block_(nk_bf16_t const *k_packed, nk_f32_t *k_out, nk_size_t kv_h,
+                                               nk_size_t kv_block_start, nk_size_t valid_kv, nk_size_t head_dim,
+                                               nk_size_t kv_len) {
+
+    nk_size_t const Bc = 16;
+    nk_size_t head_dim_padded = nk_size_round_up_to_multiple_(head_dim, 32);
+    nk_size_t tiles_per_seq = nk_size_divide_round_up_(kv_len, 16);
+    nk_size_t tiles_per_depth = head_dim_padded / 32;
+    nk_size_t tile_size = 512;
+
+    nk_size_t seq_tile = kv_block_start / 16;
+    nk_size_t base_s = kv_block_start % 16;
+
+    // Get pointer to this head's K data
+    nk_bf16_t const *k_head = k_packed + kv_h * tiles_per_seq * tiles_per_depth * tile_size;
+
+    // Extract each depth tile
+    for (nk_size_t depth_tile = 0; depth_tile < tiles_per_depth; depth_tile++) {
+        nk_size_t depth_start = depth_tile * 32;
+        nk_size_t tile_idx = seq_tile * tiles_per_depth + depth_tile;
+        nk_bf16_t const *tile_ptr = k_head + tile_idx * tile_size;
+
+        // Unpack tile: pair-interleaved layout data[d/2][s][d%2]
+        for (nk_size_t d_in_tile = 0; d_in_tile < 32 && depth_start + d_in_tile < head_dim; d_in_tile++) {
+            nk_size_t d = depth_start + d_in_tile;
+            for (nk_size_t ki = 0; ki < valid_kv; ki++) {
+                nk_size_t s_in_tile = base_s + ki;
+                if (s_in_tile >= 16) continue; // Shouldn't happen if kv_block aligned
+
+                nk_size_t elem_idx = (d_in_tile / 2) * 32 + s_in_tile * 2 + (d_in_tile % 2);
+                nk_bf16_t bf16_val = tile_ptr[elem_idx];
+                nk_f32_t f32_val;
+                nk_bf16_to_f32_serial(&bf16_val, &f32_val);
+                k_out[d * Bc + ki] = f32_val;
+            }
+        }
+    }
+}
+
+/**
+ *  @brief Extract V block from packed format: V[Bᶜ, head_dim] for a given kv_block.
+ *
+ *  V is packed for P × V, with pair-interleaving.
+ *  Output is in row-major F32 format: v_out[kᵢ × head_dim + d] = V[kᵢ, d]
+ */
+NK_INTERNAL void nk_attention_extract_v_block_(nk_bf16_t const *v_packed, nk_f32_t *v_out, nk_size_t kv_h,
+                                               nk_size_t kv_block_start, nk_size_t valid_kv, nk_size_t head_dim,
+                                               nk_size_t kv_len) {
+
+    nk_size_t head_dim_padded = nk_size_round_up_to_multiple_(head_dim, 32);
+    nk_size_t tiles_per_seq = nk_size_divide_round_up_(kv_len, 32);
+    nk_size_t tiles_per_head = nk_size_divide_round_up_(head_dim_padded, 16);
+    nk_size_t tile_size = 512;
+
+    // Get pointer to this head's V data
+    nk_bf16_t const *v_head = v_packed + kv_h * tiles_per_seq * tiles_per_head * tile_size;
+
+    // For each kv position in the block
+    for (nk_size_t ki = 0; ki < valid_kv; ki++) {
+        nk_size_t kv_pos = kv_block_start + ki;
+        nk_size_t seq_tile = kv_pos / 32;
+        nk_size_t s_in_tile = kv_pos % 32;
+
+        // Extract each head_dim tile
+        for (nk_size_t head_tile = 0; head_tile < tiles_per_head; head_tile++) {
+            nk_size_t head_start = head_tile * 16;
+            nk_size_t tile_idx = seq_tile * tiles_per_head + head_tile;
+            nk_bf16_t const *tile_ptr = v_head + tile_idx * tile_size;
+
+            // Unpack: pair-interleaved layout data[s/2][d][s%2]
+            for (nk_size_t d_in_tile = 0; d_in_tile < 16 && head_start + d_in_tile < head_dim; d_in_tile++) {
+                nk_size_t d = head_start + d_in_tile;
+                nk_size_t elem_idx = (s_in_tile / 2) * 32 + d_in_tile * 2 + (s_in_tile % 2);
+                nk_bf16_t bf16_val = tile_ptr[elem_idx];
+                nk_f32_t f32_val;
+                nk_bf16_to_f32_serial(&bf16_val, &f32_val);
+                v_out[ki * head_dim + d] = f32_val;
+            }
+        }
+    }
+}
+
+NK_PUBLIC void nk_attention_bf16_sapphireamx(nk_bf16_t const *q, void const *kv_packed, nk_f32_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_kv_packed_header_t const *header = (nk_attention_kv_packed_header_t const *)kv_packed;
+    nk_size_t head_dim_padded = header->head_dim_padded;
+    nk_size_t gqa_ratio = num_heads / num_kv_heads;
+
+    // Tile sizes
+    nk_size_t const Br = 16; // Query block rows
+    nk_size_t const Bc = 16; // KV block columns
+
+    // Configure AMX tiles
+    nk_amx_tile_configure_sapphireamx_();
+
+    // Temporary buffers (aligned to 64 bytes)
+    NK_ALIGN64 nk_f32_t scores[16 * 16];  // S = Q × Kᵀ block
+    NK_ALIGN64 nk_f32_t weights[16 * 16]; // P = softmax(S)
+    NK_ALIGN64 nk_f32_t o_acc[16 * 256];  // Output accumulator (max head_dim=256)
+
+    // Packed data pointers
+    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);
+
+    // Process each head
+    for (nk_size_t h = 0; h < num_heads; h++) {
+        nk_size_t kv_h = h / gqa_ratio;
+
+        nk_bf16_t const *q_head = q + h * query_len * head_dim;
+        nk_f32_t *o_head = output + h * query_len * head_dim;
+
+        // Process query blocks
+        for (nk_size_t qb = 0; qb < query_len; qb += Br) {
+            nk_size_t valid_q = (qb + Br <= query_len) ? Br : (query_len - qb);
+
+            // Initialize softmax state and output accumulator
+            nk_attention_softmax_row_state_t softmax_state;
+            nk_attention_softmax_init_(&softmax_state);
+
+            for (nk_size_t i = 0; i < valid_q * head_dim_padded; i++) { o_acc[i] = 0.0f; }
+
+            // Temporary buffers for extracted K and V blocks
+            NK_ALIGN64 nk_f32_t k_block[16 * 256]; // Kᵀ block [head_dim, 16]
+            NK_ALIGN64 nk_f32_t v_block[16 * 256]; // V block [16, head_dim]
+            NK_ALIGN64 nk_f32_t q_block[16 * 256]; // Q block [16, head_dim]
+
+            // Pre-convert Q block to F32
+            for (nk_size_t qi = 0; qi < valid_q; qi++) {
+                for (nk_size_t d = 0; d < head_dim; d++) {
+                    nk_bf16_t q_val = q_head[(qb + qi) * head_dim + d];
+                    nk_bf16_to_f32_serial(&q_val, &q_block[qi * head_dim + d]);
+                }
+            }
+
+            // Process KV blocks
+            for (nk_size_t kvb = 0; kvb < kv_len; kvb += Bc) {
+                nk_size_t valid_kv = (kvb + Bc <= kv_len) ? Bc : (kv_len - kvb);
+
+                // Extract K block: Kᵀ[head_dim, valid_kv] using bulk extraction
+                nk_attention_extract_k_block_(k_packed, k_block, kv_h, kvb, valid_kv, head_dim, kv_len);
+
+                // Phase 1: Compute S = Q × Kᵀ using AVX-512 FMA
+                for (nk_size_t qi = 0; qi < valid_q; qi++) {
+                    for (nk_size_t ki = 0; ki < valid_kv; ki++) {
+                        __m512 sum_v = _mm512_setzero_ps();
+                        nk_size_t d = 0;
+                        // Vectorized loop over head_dim
+                        for (; d + 16 <= head_dim; d += 16) {
+                            __m512 q_v = _mm512_loadu_ps(&q_block[qi * head_dim + d]);
+                            // Kᵀ is stored as [head_dim, kv], gather is slow, use scalar for now
+                            __m512 k_v = _mm512_set_ps(
+                                k_block[(d + 15) * 16 + ki], k_block[(d + 14) * 16 + ki], k_block[(d + 13) * 16 + ki],
+                                k_block[(d + 12) * 16 + ki], k_block[(d + 11) * 16 + ki], k_block[(d + 10) * 16 + ki],
+                                k_block[(d + 9) * 16 + ki], k_block[(d + 8) * 16 + ki], k_block[(d + 7) * 16 + ki],
+                                k_block[(d + 6) * 16 + ki], k_block[(d + 5) * 16 + ki], k_block[(d + 4) * 16 + ki],
+                                k_block[(d + 3) * 16 + ki], k_block[(d + 2) * 16 + ki], k_block[(d + 1) * 16 + ki],
+                                k_block[(d + 0) * 16 + ki]);
+                            sum_v = _mm512_fmadd_ps(q_v, k_v, sum_v);
+                        }
+                        nk_f32_t sum = _mm512_reduce_add_ps(sum_v);
+                        // Scalar tail
+                        for (; d < head_dim; d++) { sum += q_block[qi * head_dim + d] * k_block[d * 16 + ki]; }
+                        scores[qi * 16 + ki] = sum;
+                    }
+                    // Zero out invalid KV positions
+                    for (nk_size_t ki = valid_kv; ki < 16; ki++) { scores[qi * 16 + ki] = NK_F32_MIN; }
+                }
+                // Zero out invalid query rows
+                for (nk_size_t qi = valid_q; qi < 16; qi++) {
+                    for (nk_size_t ki = 0; ki < 16; ki++) { scores[qi * 16 + ki] = NK_F32_MIN; }
+                }
+
+                // Phase 2: Online softmax update
+                __m512 old_max = softmax_state.row_max;
+                nk_attention_softmax_update_(&softmax_state, scores, scale, weights);
+
+                // Rescale output accumulator if max changed
+                nk_attention_rescale_output_(o_acc, head_dim_padded, old_max, softmax_state.row_max);
+
+                // Extract V block: V[valid_kv, head_dim] using bulk extraction
+                nk_attention_extract_v_block_(v_packed, v_block, kv_h, kvb, valid_kv, head_dim, kv_len);
+
+                // Phase 3: Compute O += P × V using AVX-512 FMA
+                for (nk_size_t qi = 0; qi < valid_q; qi++) {
+                    nk_size_t d = 0;
+                    // Vectorized loop over head_dim
+                    for (; d + 16 <= head_dim; d += 16) {
+                        __m512 acc_v = _mm512_loadu_ps(&o_acc[qi * head_dim_padded + d]);
+                        for (nk_size_t ki = 0; ki < valid_kv; ki++) {
+                            __m512 p_v = _mm512_set1_ps(weights[qi * 16 + ki]);
+                            __m512 v_v = _mm512_loadu_ps(&v_block[ki * head_dim + d]);
+                            acc_v = _mm512_fmadd_ps(p_v, v_v, acc_v);
+                        }
+                        _mm512_storeu_ps(&o_acc[qi * head_dim_padded + d], acc_v);
+                    }
+                    // Scalar tail
+                    for (; d < head_dim; d++) {
+                        nk_f32_t sum = o_acc[qi * head_dim_padded + d];
+                        for (nk_size_t ki = 0; ki < valid_kv; ki++) {
+                            sum += weights[qi * 16 + ki] * v_block[ki * head_dim + d];
+                        }
+                        o_acc[qi * head_dim_padded + d] = sum;
+                    }
+                }
+            }
+
+            // Finalize: normalize O by row sums
+            float row_sums[16];
+            _mm512_store_ps(row_sums, softmax_state.row_sum);
+
+            for (nk_size_t qi = 0; qi < valid_q; qi++) {
+                nk_f32_t inv_sum = 1.0f / row_sums[qi];
+                for (nk_size_t d = 0; d < head_dim; d++) {
+                    o_head[(qb + qi) * head_dim + d] = o_acc[qi * head_dim_padded + d] * inv_sum;
+                }
+            }
+        }
+    }
+}
+
+NK_PUBLIC void nk_attention_bf16_amx_bc32_sapphireamx(nk_bf16_t const *q, void const *kv_packed, nk_f32_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_kv_packed_header_t const *header = (nk_attention_kv_packed_header_t const *)kv_packed;
+    nk_size_t head_dim_padded = header->head_dim_padded;
+    nk_size_t gqa_ratio = num_heads / num_kv_heads;
+
+    // Block sizes - Bc=32 matches V tile depth granularity
+    nk_size_t const Br = 16;
+    nk_size_t const Bc = 32;
+
+    // Configure AMX tiles
+    nk_amx_tile_configure_sapphireamx_();
+
+    // Buffers
+    NK_ALIGN64 nk_f32_t scores[16 * 32];  // S [16, 32]
+    NK_ALIGN64 nk_f32_t weights[16 * 32]; // P [16, 32]
+    NK_ALIGN64 nk_f32_t o_acc[16 * 256];  // Output accumulator
+    NK_ALIGN64 nk_bf16_t q_tile[16][32];  // Q as A-tile format
+    NK_ALIGN64 nk_f32_t s_tile[16][16];   // Score tile output (for each half)
+    NK_ALIGN64 nk_bf16_t p_tile[16][32];  // P weights as A-tile format
+    NK_ALIGN64 nk_f32_t o_tile[16][16];   // Output tile from AMX
+
+    // K packing layout (16 seq positions per tile)
+    nk_size_t k_tiles_per_seq = nk_size_divide_round_up_(kv_len, 16);
+    nk_size_t tiles_per_depth = head_dim_padded / 32;
+    nk_size_t tile_size = 512; // BF16 elements per tile
+
+    // V packing layout (32 seq positions per tile)
+    nk_size_t v_tiles_per_seq = nk_size_divide_round_up_(kv_len, 32);
+    nk_size_t v_tiles_per_head = nk_size_divide_round_up_(head_dim_padded, 16);
+
+    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);
+
+    for (nk_size_t h = 0; h < num_heads; h++) {
+        nk_size_t kv_h = h / gqa_ratio;
+        nk_bf16_t const *q_head = q + h * query_len * head_dim;
+        nk_f32_t *o_head = output + h * query_len * head_dim;
+
+        // Pointer to this KV head's packed data
+        nk_bf16_t const *k_head = k_packed + kv_h * k_tiles_per_seq * tiles_per_depth * tile_size;
+        nk_bf16_t const *v_head = v_packed + kv_h * v_tiles_per_seq * v_tiles_per_head * tile_size;
+
+        for (nk_size_t qb = 0; qb < query_len; qb += Br) {
+            nk_size_t valid_q = (qb + Br <= query_len) ? Br : (query_len - qb);
+
+            nk_attention_softmax_row_state_t softmax_state;
+            nk_attention_softmax_init_(&softmax_state);
+
+            // Zero output accumulator using SIMD
+            __m512 zero = _mm512_setzero_ps();
+            for (nk_size_t i = 0; i < 16 * head_dim_padded; i += 64) {
+                _mm512_store_ps(&o_acc[i], zero);
+                _mm512_store_ps(&o_acc[i + 16], zero);
+                _mm512_store_ps(&o_acc[i + 32], zero);
+                _mm512_store_ps(&o_acc[i + 48], zero);
+            }
+
+            // Process KV blocks in chunks of 32
+            for (nk_size_t kvb = 0; kvb < kv_len; kvb += Bc) {
+                nk_size_t valid_kv = (kvb + Bc <= kv_len) ? Bc : (kv_len - kvb);
+
+                // Phase 1: S = Q × Kᵀ using AMX
+                // Need 2 K tiles per block (each K tile has 16 columns)
+                nk_size_t k_tile_idx0 = kvb / 16;        // First K tile
+                nk_size_t k_tile_idx1 = (kvb + 16) / 16; // Second K tile
+
+                // Process first half: S[0:16, 0:16]
+                _tile_zero(0); // TMM0 = score accumulator for first 16 columns
+                _tile_zero(3); // TMM3 = score accumulator for second 16 columns
+
+                for (nk_size_t dt = 0; dt < tiles_per_depth; dt++) {
+                    nk_size_t depth_start = dt * 32;
+
+                    // Load Q[qb:qb+16, depth_start:depth_start+32] into A-tile format
+                    // Use SIMD loads when possible (full 32 elements per row)
+                    if (depth_start + 32 <= head_dim) {
+                        // Full tile - use fast SIMD copy
+                        for (nk_size_t row = 0; row < valid_q; row++) {
+                            nk_bf16_t const *q_row = q_head + (qb + row) * head_dim + depth_start;
+                            // Load 32 BF16 values (64 bytes) using two 256-bit loads
+                            __m256i q0 = _mm256_loadu_si256((__m256i const *)q_row);
+                            __m256i q1 = _mm256_loadu_si256((__m256i const *)(q_row + 16));
+                            _mm256_store_si256((__m256i *)&q_tile[row][0], q0);
+                            _mm256_store_si256((__m256i *)&q_tile[row][16], q1);
+                        }
+                    }
+                    else {
+                        // Partial tile - element-by-element with padding
+                        nk_size_t valid_depth = head_dim - depth_start;
+                        for (nk_size_t row = 0; row < valid_q; row++) {
+                            nk_bf16_t const *q_row = q_head + (qb + row) * head_dim + depth_start;
+                            for (nk_size_t col = 0; col < 32; col++) {
+                                q_tile[row][col] = (col < valid_depth) ? q_row[col] : 0;
+                            }
+                        }
+                    }
+                    // Zero pad remaining rows
+                    for (nk_size_t row = valid_q; row < 16; row++) {
+                        _mm256_store_si256((__m256i *)&q_tile[row][0], _mm256_setzero_si256());
+                        _mm256_store_si256((__m256i *)&q_tile[row][16], _mm256_setzero_si256());
+                    }
+
+                    _tile_loadd(1, q_tile, 64); // A: 16×32 BF16
+
+                    // First K tile (columns 0:16)
+                    nk_bf16_t const *k_tile_ptr0 = k_head + (k_tile_idx0 * tiles_per_depth + dt) * tile_size;
+                    _tile_loadd(2, k_tile_ptr0, 64); // B: 32×16 BF16
+                    _tile_dpbf16ps(0, 1, 2);         // TMM0 += Q × K0
+
+                    // Second K tile (columns 16:32) if within bounds
+                    if (kvb + 16 < kv_len) {
+                        nk_bf16_t const *k_tile_ptr1 = k_head + (k_tile_idx1 * tiles_per_depth + dt) * tile_size;
+                        _tile_loadd(4, k_tile_ptr1, 64); // B: 32×16 BF16
+                        _tile_dpbf16ps(3, 1, 4);         // TMM3 += Q × K1
+                    }
+                }
+
+                // Store scores from TMM0 and TMM3
+                // Use SIMD for fast extraction
+                _tile_stored(0, s_tile, 64);
+
+                __m512 neg_inf = _mm512_set1_ps(NK_F32_MIN);
+
+                if (valid_q == 16 && valid_kv >= 16) {
+                    // Fast path: full first half, just copy
+                    for (nk_size_t qi = 0; qi < 16; qi++) {
+                        __m512 s0 = _mm512_load_ps(&s_tile[qi][0]);
+                        _mm512_store_ps(&scores[qi * 32], s0);
+                    }
+                }
+                else {
+                    // Partial - need masking
+                    __mmask16 kv_mask = (1u << valid_kv) - 1;
+                    for (nk_size_t qi = 0; qi < 16; qi++) {
+                        __m512 s0 = _mm512_load_ps(&s_tile[qi][0]);
+                        if (qi < valid_q) { s0 = _mm512_mask_blend_ps(kv_mask, neg_inf, s0); }
+                        else { s0 = neg_inf; }
+                        _mm512_store_ps(&scores[qi * 32], s0);
+                    }
+                }
+
+                // Second half scores (columns 16:32)
+                if (kvb + 16 < kv_len) {
+                    _tile_stored(3, s_tile, 64);
+                    nk_size_t valid_kv2 = (valid_kv > 16) ? (valid_kv - 16) : 0;
+
+                    if (valid_q == 16 && valid_kv2 >= 16) {
+                        // Fast path
+                        for (nk_size_t qi = 0; qi < 16; qi++) {
+                            __m512 s1 = _mm512_load_ps(&s_tile[qi][0]);
+                            _mm512_store_ps(&scores[qi * 32 + 16], s1);
+                        }
+                    }
+                    else {
+                        __mmask16 kv_mask2 = (valid_kv2 >= 16) ? 0xFFFF : ((1u << valid_kv2) - 1);
+                        for (nk_size_t qi = 0; qi < 16; qi++) {
+                            __m512 s1 = _mm512_load_ps(&s_tile[qi][0]);
+                            if (qi < valid_q) { s1 = _mm512_mask_blend_ps(kv_mask2, neg_inf, s1); }
+                            else { s1 = neg_inf; }
+                            _mm512_store_ps(&scores[qi * 32 + 16], s1);
+                        }
+                    }
+                }
+                else {
+                    // Mask out second half entirely
+                    for (nk_size_t qi = 0; qi < 16; qi++) { _mm512_store_ps(&scores[qi * 32 + 16], neg_inf); }
+                }
+
+                // Phase 2: online softmax (fast degree-4 exp)
+                __m512 old_max = softmax_state.row_max;
+                nk_attention_softmax_update_bc32_fast_(&softmax_state, scores, scale, weights);
+                nk_attention_rescale_output_(o_acc, head_dim_padded, old_max, softmax_state.row_max);
+
+                // Phase 3: O += P × V using AMX
+                // Convert P[16, 32] from F32 to BF16 and pack as A-tile
+                for (nk_size_t qi = 0; qi < 16; qi++) {
+                    for (nk_size_t ki = 0; ki < 32; ki += 16) {
+                        __m512 p_f32 = _mm512_loadu_ps(&weights[qi * 32 + ki]);
+                        __m256bh p_bf16 = _mm512_cvtneps_pbh(p_f32);
+                        // Store BF16 vector - cast through union or memory
+                        *(__m256bh *)&p_tile[qi][ki] = p_bf16;
+                    }
+                }
+
+                // V tile index for this block
+                nk_size_t v_seq_tile = kvb / 32;
+
+                // For each head_dim chunk of 16
+                for (nk_size_t ht = 0; ht < v_tiles_per_head; ht++) {
+                    nk_size_t head_start = ht * 16;
+
+                    // V tile is already packed: V[32, 16] in B-tile format
+                    nk_bf16_t const *v_tile_ptr = v_head + (v_seq_tile * v_tiles_per_head + ht) * tile_size;
+
+                    // Zero output tile
+                    _tile_zero(5);
+
+                    // Load P into TMM6 (A-tile: 16×32)
+                    _tile_loadd(6, p_tile, 64);
+
+                    // Load V into TMM7 (B-tile: 32×16)
+                    _tile_loadd(7, v_tile_ptr, 64);
+
+                    // O_tile = P × V
+                    _tile_dpbf16ps(5, 6, 7);
+
+                    // Store and accumulate
+                    _tile_stored(5, o_tile, 64);
+
+                    // Add to output accumulator - unrolled for all 16 rows
+                    // Even if valid_q < 16, we accumulate all (padded rows have zero weights)
+                    for (nk_size_t qi = 0; qi < 16; qi += 4) {
+                        __m512 acc0 = _mm512_load_ps(&o_acc[(qi + 0) * head_dim_padded + head_start]);
+                        __m512 acc1 = _mm512_load_ps(&o_acc[(qi + 1) * head_dim_padded + head_start]);
+                        __m512 acc2 = _mm512_load_ps(&o_acc[(qi + 2) * head_dim_padded + head_start]);
+                        __m512 acc3 = _mm512_load_ps(&o_acc[(qi + 3) * head_dim_padded + head_start]);
+
+                        acc0 = _mm512_add_ps(acc0, _mm512_load_ps(&o_tile[qi + 0][0]));
+                        acc1 = _mm512_add_ps(acc1, _mm512_load_ps(&o_tile[qi + 1][0]));
+                        acc2 = _mm512_add_ps(acc2, _mm512_load_ps(&o_tile[qi + 2][0]));
+                        acc3 = _mm512_add_ps(acc3, _mm512_load_ps(&o_tile[qi + 3][0]));
+
+                        _mm512_store_ps(&o_acc[(qi + 0) * head_dim_padded + head_start], acc0);
+                        _mm512_store_ps(&o_acc[(qi + 1) * head_dim_padded + head_start], acc1);
+                        _mm512_store_ps(&o_acc[(qi + 2) * head_dim_padded + head_start], acc2);
+                        _mm512_store_ps(&o_acc[(qi + 3) * head_dim_padded + head_start], acc3);
+                    }
+                }
+            }
+
+            // Finalize: normalize O by row sums
+            float row_sums[16];
+            _mm512_store_ps(row_sums, softmax_state.row_sum);
+            for (nk_size_t qi = 0; qi < valid_q; qi++) {
+                nk_f32_t inv_sum = 1.0f / row_sums[qi];
+                for (nk_size_t d = 0; d < head_dim; d++) {
+                    o_head[(qb + qi) * head_dim + d] = o_acc[qi * head_dim_padded + d] * inv_sum;
+                }
+            }
+        }
+    }
+}
+
+NK_PUBLIC void nk_attention_bf16_amx_optimized_sapphireamx(nk_bf16_t const *q, void const *kv_packed, nk_f32_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_kv_packed_header_t const *header = (nk_attention_kv_packed_header_t const *)kv_packed;
+    nk_size_t head_dim_padded = header->head_dim_padded;
+    nk_size_t gqa_ratio = num_heads / num_kv_heads;
+
+    nk_size_t const Br = 16;
+    nk_size_t const Bc = 32;
+
+    // Configure AMX tiles once
+    nk_amx_tile_configure_sapphireamx_();
+
+    // Tile dimensions
+    nk_size_t tiles_per_depth = head_dim_padded / 32;                           // 4 for d=128
+    nk_size_t v_tiles_per_head = nk_size_divide_round_up_(head_dim_padded, 16); // 8 for d=128
+
+    // K packing layout (16 seq positions per tile)
+    nk_size_t k_tiles_per_seq = nk_size_divide_round_up_(kv_len, 16);
+    nk_size_t tile_size = 512;
+
+    // V packing layout (32 seq positions per tile)
+    nk_size_t v_tiles_per_seq = nk_size_divide_round_up_(kv_len, 32);
+
+    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);
+
+    // Pre-allocated buffers (all L1-resident)
+    NK_ALIGN64 nk_bf16_t q_tiles[4][16][32]; // Q tiles for all depth chunks (max 4 for d=128)
+    NK_ALIGN64 nk_f32_t scores[16][32];      // Score buffer (direct tile store target)
+    NK_ALIGN64 nk_f32_t weights[16][32];     // Softmax output
+    NK_ALIGN64 nk_bf16_t p_tile[16][32];     // P weights in BF16
+    NK_ALIGN64 nk_f32_t o_tile[16][16];      // Output tile buffer
+    NK_ALIGN64 nk_f32_t o_acc[16][256];      // Output accumulator (max d=256)
+
+    __m512 neg_inf = _mm512_set1_ps(NK_F32_MIN);
+
+    for (nk_size_t h = 0; h < num_heads; h++) {
+        nk_size_t kv_h = h / gqa_ratio;
+        nk_bf16_t const *q_head = q + h * query_len * head_dim;
+        nk_f32_t *o_head = output + h * query_len * head_dim;
+
+        nk_bf16_t const *k_head = k_packed + kv_h * k_tiles_per_seq * tiles_per_depth * tile_size;
+        nk_bf16_t const *v_head = v_packed + kv_h * v_tiles_per_seq * v_tiles_per_head * tile_size;
+
+        for (nk_size_t qb = 0; qb < query_len; qb += Br) {
+            nk_size_t valid_q = (qb + Br <= query_len) ? Br : (query_len - qb);
+
+            // Pre-pack Q tiles once for all KV blocks
+            for (nk_size_t dt = 0; dt < tiles_per_depth; dt++) {
+                nk_size_t depth_start = dt * 32;
+                if (depth_start + 32 <= head_dim) {
+                    // Full tile - fast SIMD copy
+                    for (nk_size_t row = 0; row < valid_q; row++) {
+                        nk_bf16_t const *q_row = q_head + (qb + row) * head_dim + depth_start;
+                        __m256i q0 = _mm256_loadu_si256((__m256i const *)q_row);
+                        __m256i q1 = _mm256_loadu_si256((__m256i const *)(q_row + 16));
+                        _mm256_store_si256((__m256i *)&q_tiles[dt][row][0], q0);
+                        _mm256_store_si256((__m256i *)&q_tiles[dt][row][16], q1);
+                    }
+                    // Zero remaining rows
+                    for (nk_size_t row = valid_q; row < 16; row++) {
+                        _mm256_store_si256((__m256i *)&q_tiles[dt][row][0], _mm256_setzero_si256());
+                        _mm256_store_si256((__m256i *)&q_tiles[dt][row][16], _mm256_setzero_si256());
+                    }
+                }
+                else {
+                    // Partial tile with padding
+                    nk_size_t valid_depth = head_dim - depth_start;
+                    for (nk_size_t row = 0; row < 16; row++) {
+                        for (nk_size_t col = 0; col < 32; col++) {
+                            if (row < valid_q && col < valid_depth) {
+                                q_tiles[dt][row][col] = q_head[(qb + row) * head_dim + depth_start + col];
+                            }
+                            else { q_tiles[dt][row][col] = 0; }
+                        }
+                    }
+                }
+            }
+
+            // Initialize softmax state and output accumulator
+            nk_attention_softmax_row_state_t softmax_state;
+            nk_attention_softmax_init_(&softmax_state);
+
+            __m512 zero = _mm512_setzero_ps();
+            for (nk_size_t i = 0; i < 16 * head_dim_padded; i += 64) {
+                _mm512_store_ps(&o_acc[0][i], zero);
+                _mm512_store_ps(&o_acc[0][i + 16], zero);
+                _mm512_store_ps(&o_acc[0][i + 32], zero);
+                _mm512_store_ps(&o_acc[0][i + 48], zero);
+            }
+
+            // Process KV blocks
+            for (nk_size_t kvb = 0; kvb < kv_len; kvb += Bc) {
+                nk_size_t valid_kv = (kvb + Bc <= kv_len) ? Bc : (kv_len - kvb);
+                nk_size_t k_tile_idx0 = kvb / 16;
+                nk_size_t k_tile_idx1 = (kvb + 16) / 16;
+
+                // Phase 1: S = Q × Kᵀ using pre-packed Q tiles
+                _tile_zero(0); // Score cols 0:16
+                _tile_zero(3); // Score cols 16:32
+
+                for (nk_size_t dt = 0; dt < tiles_per_depth; dt++) {
+                    // Load pre-packed Q tile from L1 (not global!)
+                    _tile_loadd(1, q_tiles[dt], 64);
+
+                    // Load K tiles from global (necessary)
+                    nk_bf16_t const *k_tile_ptr0 = k_head + (k_tile_idx0 * tiles_per_depth + dt) * tile_size;
+                    _tile_loadd(2, k_tile_ptr0, 64);
+                    _tile_dpbf16ps(0, 1, 2);
+
+                    if (kvb + 16 < kv_len) {
+                        nk_bf16_t const *k_tile_ptr1 = k_head + (k_tile_idx1 * tiles_per_depth + dt) * tile_size;
+                        _tile_loadd(4, k_tile_ptr1, 64);
+                        _tile_dpbf16ps(3, 1, 4);
+                    }
+                }
+
+                // Store first 16 columns directly to scores[0:16]
+                _tile_stored(0, &scores[0][0], 128); // stride=128 bytes (32 floats)
+
+                // Store second 16 columns to scores[16:32]
+                if (kvb + 16 < kv_len) { _tile_stored(3, &scores[0][16], 128); }
+                else {
+                    // Mask out second half
+                    for (nk_size_t qi = 0; qi < 16; qi++) { _mm512_store_ps(&scores[qi][16], neg_inf); }
+                }
+
+                // Apply masking for invalid positions (only on boundaries)
+                if (valid_q < 16 || valid_kv < 32) {
+                    __mmask16 kv_mask0 = (valid_kv >= 16) ? 0xFFFF : ((1u << valid_kv) - 1);
+                    __mmask16 kv_mask1 = (valid_kv > 16) ? ((1u << (valid_kv - 16)) - 1) : 0;
+                    if (valid_kv >= 32) kv_mask1 = 0xFFFF;
+
+                    for (nk_size_t qi = 0; qi < 16; qi++) {
+                        if (qi >= valid_q) {
+                            _mm512_store_ps(&scores[qi][0], neg_inf);
+                            _mm512_store_ps(&scores[qi][16], neg_inf);
+                        }
+                        else {
+                            __m512 s0 = _mm512_load_ps(&scores[qi][0]);
+                            __m512 s1 = _mm512_load_ps(&scores[qi][16]);
+                            _mm512_store_ps(&scores[qi][0], _mm512_mask_blend_ps(kv_mask0, neg_inf, s0));
+                            _mm512_store_ps(&scores[qi][16], _mm512_mask_blend_ps(kv_mask1, neg_inf, s1));
+                        }
+                    }
+                }
+
+                // Phase 2: online softmax (fast degree-4 exp)
+                __m512 old_max = softmax_state.row_max;
+                nk_attention_softmax_update_bc32_fast_(&softmax_state, &scores[0][0], scale, &weights[0][0]);
+                nk_attention_rescale_output_(&o_acc[0][0], head_dim_padded, old_max, softmax_state.row_max);
+
+                // Phase 3: O += P × V with hoisted P tile load
+                // Convert F32 weights to BF16 P tile (once per KV block)
+                for (nk_size_t qi = 0; qi < 16; qi++) {
+                    __m512 p0 = _mm512_load_ps(&weights[qi][0]);
+                    __m512 p1 = _mm512_load_ps(&weights[qi][16]);
+                    __m256bh pb0 = _mm512_cvtneps_pbh(p0);
+                    __m256bh pb1 = _mm512_cvtneps_pbh(p1);
+                    *(__m256bh *)&p_tile[qi][0] = pb0;
+                    *(__m256bh *)&p_tile[qi][16] = pb1;
+                }
+
+                // Load P tile once, reuse for all V tiles
+                _tile_loadd(6, p_tile, 64);
+
+                nk_size_t v_seq_tile = kvb / 32;
+
+                for (nk_size_t ht = 0; ht < v_tiles_per_head; ht++) {
+                    nk_size_t head_start = ht * 16;
+
+                    // Load V tile from global
+                    nk_bf16_t const *v_tile_ptr = v_head + (v_seq_tile * v_tiles_per_head + ht) * tile_size;
+
+                    _tile_zero(5);
+                    // P already in TMM6 - no reload!
+                    _tile_loadd(7, v_tile_ptr, 64);
+                    _tile_dpbf16ps(5, 6, 7);
+
+                    // Store and accumulate
+                    _tile_stored(5, o_tile, 64);
+
+                    // Accumulate into output (unrolled)
+                    for (nk_size_t qi = 0; qi < 16; qi += 4) {
+                        __m512 acc0 = _mm512_load_ps(&o_acc[qi + 0][head_start]);
+                        __m512 acc1 = _mm512_load_ps(&o_acc[qi + 1][head_start]);
+                        __m512 acc2 = _mm512_load_ps(&o_acc[qi + 2][head_start]);
+                        __m512 acc3 = _mm512_load_ps(&o_acc[qi + 3][head_start]);
+
+                        acc0 = _mm512_add_ps(acc0, _mm512_load_ps(&o_tile[qi + 0][0]));
+                        acc1 = _mm512_add_ps(acc1, _mm512_load_ps(&o_tile[qi + 1][0]));
+                        acc2 = _mm512_add_ps(acc2, _mm512_load_ps(&o_tile[qi + 2][0]));
+                        acc3 = _mm512_add_ps(acc3, _mm512_load_ps(&o_tile[qi + 3][0]));
+
+                        _mm512_store_ps(&o_acc[qi + 0][head_start], acc0);
+                        _mm512_store_ps(&o_acc[qi + 1][head_start], acc1);
+                        _mm512_store_ps(&o_acc[qi + 2][head_start], acc2);
+                        _mm512_store_ps(&o_acc[qi + 3][head_start], acc3);
+                    }
+                }
+            }
+
+            // Finalize: normalize O by row sums
+            float row_sums[16];
+            _mm512_store_ps(row_sums, softmax_state.row_sum);
+            for (nk_size_t qi = 0; qi < valid_q; qi++) {
+                __m512 inv_sum = _mm512_set1_ps(1.0f / row_sums[qi]);
+                for (nk_size_t d = 0; d < head_dim; d += 16) {
+                    __m512 o = _mm512_load_ps(&o_acc[qi][d]);
+                    o = _mm512_mul_ps(o, inv_sum);
+                    _mm512_storeu_ps(&o_head[(qb + qi) * head_dim + d], o);
+                }
+            }
+        }
+    }
+}
+
+NK_PUBLIC void nk_attention_causal_bf16_sapphireamx(nk_bf16_t const *q, void const *kv_packed, nk_f32_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) {
+
+    // For causal attention in autoregressive decode:
+    // Query position q_pos can only attend to KV positions 0..q_pos
+    // If kv_len == query_len (prefill), we need proper masking
+    // If query_len == 1 (decode), the single query can see all KV
+
+    // Simplified: just call full attention for now
+    // TODO: Implement proper causal masking with block skipping
+    nk_attention_bf16_sapphireamx(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_SAPPHIREAMX
+#endif // NK_TARGET_X86_
+#endif // NK_ATTENTION_SAPPHIREAMX_H