numkong 7.5.0 → 7.6.0

package/include/numkong/dot/skylake.h

@@ -156,6 +156,36 @@ NK_INTERNAL void nk_dot_through_f32_update_skylake_(nk_dot_through_f32_state_sky
     state->sum_f32x16 = _mm512_fmadd_ps(a.zmm_ps, b.zmm_ps, state->sum_f32x16);
 }
 
+/**
+ *  @brief E5M2 byte-batched update: consumes 64 raw E5M2 bytes per call and widens inline.
+ *  Two independent FMA chains (each 2-deep) merge into the single state accumulator at exit.
+ *  Keeps register pressure at one __m512 across calls while breaking the FMA dep chain.
+ */
+NK_INTERNAL void nk_dot_e5m2x64_update_skylake_(nk_dot_through_f32_state_skylake_t_ *state, nk_b512_vec_t a_bytes,
+                                                nk_b512_vec_t b_bytes, nk_size_t depth_offset,
+                                                nk_size_t active_dimensions) {
+    nk_unused_(depth_offset);
+    nk_unused_(active_dimensions);
+    __m512i const zero_u8x64 = _mm512_setzero_si512();
+    __m512i a_even_f16x32 = _mm512_unpacklo_epi8(zero_u8x64, a_bytes.zmm);
+    __m512i a_odd_f16x32 = _mm512_unpackhi_epi8(zero_u8x64, a_bytes.zmm);
+    __m512i b_even_f16x32 = _mm512_unpacklo_epi8(zero_u8x64, b_bytes.zmm);
+    __m512i b_odd_f16x32 = _mm512_unpackhi_epi8(zero_u8x64, b_bytes.zmm);
+    __m512 a_first_f32x16 = _mm512_cvtph_ps(_mm512_castsi512_si256(a_even_f16x32));
+    __m512 a_second_f32x16 = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(a_even_f16x32, 1));
+    __m512 a_third_f32x16 = _mm512_cvtph_ps(_mm512_castsi512_si256(a_odd_f16x32));
+    __m512 a_fourth_f32x16 = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(a_odd_f16x32, 1));
+    __m512 b_first_f32x16 = _mm512_cvtph_ps(_mm512_castsi512_si256(b_even_f16x32));
+    __m512 b_second_f32x16 = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(b_even_f16x32, 1));
+    __m512 b_third_f32x16 = _mm512_cvtph_ps(_mm512_castsi512_si256(b_odd_f16x32));
+    __m512 b_fourth_f32x16 = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(b_odd_f16x32, 1));
+    __m512 first_chain_f32x16 = _mm512_mul_ps(a_first_f32x16, b_first_f32x16);
+    __m512 second_chain_f32x16 = _mm512_mul_ps(a_second_f32x16, b_second_f32x16);
+    first_chain_f32x16 = _mm512_fmadd_ps(a_third_f32x16, b_third_f32x16, first_chain_f32x16);
+    second_chain_f32x16 = _mm512_fmadd_ps(a_fourth_f32x16, b_fourth_f32x16, second_chain_f32x16);
+    state->sum_f32x16 = _mm512_add_ps(state->sum_f32x16, _mm512_add_ps(first_chain_f32x16, second_chain_f32x16));
+}
+
 /**
  *  @brief Finalizes 4x low-precision dot-products placing them into 4x consecutive 32-bit slots.
  *  @sa nk_dot_f16x16_udpate_skylake, nk_dot_bf16x16_udpate_skylake
@@ -543,7 +573,7 @@ NK_PUBLIC void nk_dot_e4m3_skylake(nk_e4m3_t const *a_scalars, nk_e4m3_t const *
 
 nk_dot_e4m3_skylake_cycle:
     if (count_scalars < 16) {
-        __mmask16 mask = (__mmask16)_bzhi_u32(0xFFFF, count_scalars);
+        __mmask16 mask = (__mmask16)_bzhi_u32(0xFFFF, (unsigned int)count_scalars);
         a_e4m3_u8x16 = _mm_maskz_loadu_epi8(mask, a_scalars);
         b_e4m3_u8x16 = _mm_maskz_loadu_epi8(mask, b_scalars);
         count_scalars = 0;
@@ -563,27 +593,44 @@ nk_dot_e4m3_skylake_cycle:
 
 NK_PUBLIC void nk_dot_e5m2_skylake(nk_e5m2_t const *a_scalars, nk_e5m2_t const *b_scalars, nk_size_t count_scalars,
                                    nk_f32_t *result) {
-    __m128i a_e5m2_u8x16, b_e5m2_u8x16;
-    __m512 sum_f32x16 = _mm512_setzero_ps();
+    // E5M2 shares F16 bias (15): vpunpck*bw against zero places the byte as F16 encoding,
+    // so we inline the widen rather than calling the helper 4× — same ops, cleaner code.
+    __m512 first_chain_f32x16 = _mm512_setzero_ps();
+    __m512 second_chain_f32x16 = _mm512_setzero_ps();
+    __m512i const zero_u8x64 = _mm512_setzero_si512();
+    __m512i a_u8x64, b_u8x64;
 
 nk_dot_e5m2_skylake_cycle:
-    if (count_scalars < 16) {
-        __mmask16 mask = (__mmask16)_bzhi_u32(0xFFFF, count_scalars);
-        a_e5m2_u8x16 = _mm_maskz_loadu_epi8(mask, a_scalars);
-        b_e5m2_u8x16 = _mm_maskz_loadu_epi8(mask, b_scalars);
+    if (count_scalars < 64) {
+        __mmask64 mask = _bzhi_u64(0xFFFFFFFFFFFFFFFFULL, (unsigned int)count_scalars);
+        a_u8x64 = _mm512_maskz_loadu_epi8(mask, a_scalars);
+        b_u8x64 = _mm512_maskz_loadu_epi8(mask, b_scalars);
         count_scalars = 0;
     }
     else {
-        a_e5m2_u8x16 = _mm_loadu_si128((__m128i const *)a_scalars);
-        b_e5m2_u8x16 = _mm_loadu_si128((__m128i const *)b_scalars);
-        a_scalars += 16, b_scalars += 16, count_scalars -= 16;
+        a_u8x64 = _mm512_loadu_si512((__m512i const *)a_scalars);
+        b_u8x64 = _mm512_loadu_si512((__m512i const *)b_scalars);
+        a_scalars += 64, b_scalars += 64, count_scalars -= 64;
     }
-    __m512 a_f32x16 = nk_e5m2x16_to_f32x16_skylake_(a_e5m2_u8x16);
-    __m512 b_f32x16 = nk_e5m2x16_to_f32x16_skylake_(b_e5m2_u8x16);
-    sum_f32x16 = _mm512_fmadd_ps(a_f32x16, b_f32x16, sum_f32x16);
+    __m512i a_even_f16x32 = _mm512_unpacklo_epi8(zero_u8x64, a_u8x64);
+    __m512i a_odd_f16x32 = _mm512_unpackhi_epi8(zero_u8x64, a_u8x64);
+    __m512i b_even_f16x32 = _mm512_unpacklo_epi8(zero_u8x64, b_u8x64);
+    __m512i b_odd_f16x32 = _mm512_unpackhi_epi8(zero_u8x64, b_u8x64);
+    __m512 a_first_f32x16 = _mm512_cvtph_ps(_mm512_castsi512_si256(a_even_f16x32));
+    __m512 a_second_f32x16 = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(a_even_f16x32, 1));
+    __m512 a_third_f32x16 = _mm512_cvtph_ps(_mm512_castsi512_si256(a_odd_f16x32));
+    __m512 a_fourth_f32x16 = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(a_odd_f16x32, 1));
+    __m512 b_first_f32x16 = _mm512_cvtph_ps(_mm512_castsi512_si256(b_even_f16x32));
+    __m512 b_second_f32x16 = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(b_even_f16x32, 1));
+    __m512 b_third_f32x16 = _mm512_cvtph_ps(_mm512_castsi512_si256(b_odd_f16x32));
+    __m512 b_fourth_f32x16 = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(b_odd_f16x32, 1));
+    first_chain_f32x16 = _mm512_fmadd_ps(a_first_f32x16, b_first_f32x16, first_chain_f32x16);
+    second_chain_f32x16 = _mm512_fmadd_ps(a_second_f32x16, b_second_f32x16, second_chain_f32x16);
+    first_chain_f32x16 = _mm512_fmadd_ps(a_third_f32x16, b_third_f32x16, first_chain_f32x16);
+    second_chain_f32x16 = _mm512_fmadd_ps(a_fourth_f32x16, b_fourth_f32x16, second_chain_f32x16);
     if (count_scalars) goto nk_dot_e5m2_skylake_cycle;
 
-    *result = nk_reduce_add_f32x16_skylake_(sum_f32x16);
+    *result = nk_reduce_add_f32x16_skylake_(_mm512_add_ps(first_chain_f32x16, second_chain_f32x16));
 }
 
 NK_PUBLIC void nk_dot_e2m3_skylake(nk_e2m3_t const *a_scalars, nk_e2m3_t const *b_scalars, nk_size_t count_scalars,

package/include/numkong/dots/README.md

@@ -72,6 +72,8 @@ Int8 data is quad-interleaved: [a₀, a₁, a₂, a₃, a₀, a₁, a₂, a₃,
 Tile configuration via `LDTILECFG` sets row counts and column byte-widths per tile — allows undersized tiles at matrix edges without masking.
 Morton Z-curve ordering for tile traversal improves cache reuse when both A and B exceed L2.
 This eliminates the explicit M×N×K loop nesting and register file pressure of vector ISAs — the entire dot-product reduction happens inside the tile instruction.
+FP8 inputs on Sapphire AMX go through an on-the-fly E4M3/E5M2 → BF16 pack via the Ice Lake `VPERMI2W` LUT helpers — port-5-bound but the simplest correct route to feed `TDPBF16PS` tiles.
+Granite Rapids adds `TDPFP16PS` (same tile shape, FP16 operands); the E5M2 variant widens inputs with a single `VPUNPCK*BW` against zero into FP16 tiles at pack time and then reuses the native FP16 compute loop — keeps the intermediate at FP16 precision instead of truncating to BF16 like the Sapphire path.
 
 ### SME Outer-Product Streaming
 

package/include/numkong/dots/graniteamx.h

@@ -718,6 +718,440 @@ NK_PUBLIC void nk_dots_symmetric_f16_graniteamx(
 
 #pragma endregion F16 Native
 
+#pragma region E5M2 Source (widened to FP16 tiles)
+
+/*  E5M2 Granite AMX kernels: same F16 tile shapes, same TDPFP16PS compute body as the F16 path.
+ *  The only difference is byte-to-word widening during A-load and B-pack: E5M2 shares F16's
+ *  exponent bias (15), so `(byte << 8)` is the exact F16 bit pattern for every E5M2 value,
+ *  including zero/subnormals/Inf/NaN. Tile buffers hold F16 after widen — "e5m2" in the
+ *  typedefs refers to the source dtype, not the on-tile representation.
+ *
+ *  The tile types below alias the F16 types (identical memory layout) so we can reuse the
+ *  F16 init/store/output2x2/update internal helpers by pointer cast at the boundary. Only the
+ *  public entry points and the load-A widen helper are new.
+ */
+
+typedef nk_dots_f16_a16x32_graniteamx_t nk_dots_e5m2_a16x32_graniteamx_t;
+typedef nk_dots_f16_b32x16_graniteamx_t nk_dots_e5m2_b32x16_graniteamx_t;
+typedef nk_dots_f16_state_graniteamx_t nk_dots_e5m2_state_graniteamx_t;
+typedef nk_dots_f16_state2x2_graniteamx_t nk_dots_e5m2_state2x2_graniteamx_t;
+
+/* Load A tile from E5M2 row-major source, widen to F16 via `(byte << 8)` into the F16 tile buffer. */
+NK_INTERNAL void nk_dots_e5m2_load_a_graniteamx_(        //
+    nk_dots_e5m2_a16x32_graniteamx_t *a_tile,            //
+    nk_e5m2_t const *src, nk_size_t src_stride_elements, //
+    nk_size_t valid_rows, nk_size_t valid_cols) {
+
+    __mmask32 column_mask = (valid_cols >= 32) ? 0xFFFFFFFF : ((__mmask32)1 << valid_cols) - 1;
+    __m512i zero_i16x32 = _mm512_setzero_si512();
+
+    for (nk_size_t row_idx = 0; row_idx < 16; row_idx++) {
+        if (row_idx < valid_rows) {
+            __m256i e5m2_u8x32 = _mm256_maskz_loadu_epi8(column_mask, src + row_idx * src_stride_elements);
+            __m512i word_u16x32 = _mm512_cvtepu8_epi16(e5m2_u8x32);
+            __m512i f16_bits_i16x32 = _mm512_slli_epi16(word_u16x32, 8);
+            _mm512_store_si512((__m512i *)a_tile->data[row_idx], f16_bits_i16x32);
+        }
+        else { _mm512_store_si512((__m512i *)a_tile->data[row_idx], zero_i16x32); }
+    }
+    nk_compiler_barrier_sapphireamx_();
+}
+
+NK_PUBLIC nk_size_t nk_dots_packed_size_e5m2_graniteamx(nk_size_t column_count, nk_size_t depth) {
+    nk_size_t const tmm_rows = 16;
+    nk_size_t const tmm_cols = 32;
+    nk_size_t const tile_bytes = 512 * sizeof(nk_f16_t); // Tiles hold F16 after widen: same 1KB as F16.
+
+    nk_size_t const full_column_tiles = column_count / tmm_rows;
+    nk_size_t const tiles_along_depth = nk_size_divide_round_up_(depth, tmm_cols);
+    nk_size_t const column_remainder_count = column_count - full_column_tiles * tmm_rows;
+
+    nk_size_t size = sizeof(nk_dots_amx_packed_header_t);
+    size += full_column_tiles * tiles_along_depth * tile_bytes;
+    if (column_remainder_count > 0) size += column_remainder_count * depth * sizeof(nk_f16_t);
+    size += column_count * sizeof(nk_f32_t);
+
+    return size;
+}
+
+NK_PUBLIC void nk_dots_pack_e5m2_graniteamx(                     //
+    nk_e5m2_t const *b, nk_size_t column_count, nk_size_t depth, //
+    nk_size_t b_stride_in_bytes, void *b_packed) {
+
+    nk_size_t const tmm_rows = 16;
+    nk_size_t const tmm_cols = 32;
+    nk_size_t const tile_elements = 512;
+    nk_size_t const tile_bytes = tile_elements * sizeof(nk_f16_t);
+    nk_size_t const b_stride_elements = b_stride_in_bytes; // E5M2: 1 byte per element
+
+    nk_size_t const column_tiles_count = column_count / tmm_rows;
+    nk_size_t const depth_tiles_count = nk_size_divide_round_up_(depth, tmm_cols);
+    nk_size_t const column_remainder_count = column_count - column_tiles_count * tmm_rows;
+    nk_size_t const total_tiles = column_tiles_count * depth_tiles_count;
+
+    nk_dots_amx_packed_header_t *header = (nk_dots_amx_packed_header_t *)b_packed;
+    header->full_column_tiles = (nk_u32_t)column_tiles_count;
+    header->full_depth_tiles = (nk_u32_t)depth_tiles_count;
+    header->column_remainder_count = (nk_u32_t)column_remainder_count;
+
+    nk_size_t const tiles_offset = sizeof(nk_dots_amx_packed_header_t);
+    nk_size_t const column_edge_offset = tiles_offset + total_tiles * tile_bytes;
+    header->column_edge_offset = (nk_u32_t)column_edge_offset;
+
+    nk_f16_t *tiles_ptr = (nk_f16_t *)((char *)b_packed + tiles_offset);
+    nk_f16_t *column_edge_ptr = (nk_f16_t *)((char *)b_packed + column_edge_offset);
+
+    // Pack tiles: widen-and-gather 16 strided E5M2 rows into an F16-aligned source tile,
+    // then transpose via the shared BF16/F16 transposer (same 16-bit lane layout).
+    for (nk_size_t column_tile_idx = 0; column_tile_idx < column_tiles_count; column_tile_idx++) {
+        for (nk_size_t depth_tile_idx = 0; depth_tile_idx < depth_tiles_count; depth_tile_idx++) {
+
+            nk_size_t const tile_index = column_tile_idx * depth_tiles_count + depth_tile_idx;
+            nk_f16_t *tile_output = tiles_ptr + tile_index * tile_elements;
+
+            nk_size_t const src_row_start = column_tile_idx * tmm_rows;
+            nk_size_t const src_column_start = depth_tile_idx * tmm_cols;
+            nk_size_t const columns_to_pack = (src_column_start + tmm_cols <= depth) ? tmm_cols
+                                                                                     : (depth - src_column_start);
+
+            nk_dots_bf16_a16x32_sapphireamx_t source_tile;
+            __mmask32 depth_mask = (__mmask32)((columns_to_pack < 32) ? ((1U << columns_to_pack) - 1) : ~0U);
+            for (nk_size_t row_idx = 0; row_idx < tmm_rows; row_idx++) {
+                nk_e5m2_t const *source_row = b + (src_row_start + row_idx) * b_stride_elements + src_column_start;
+                __m256i e5m2_u8x32 = _mm256_maskz_loadu_epi8(depth_mask, source_row);
+                __m512i word_u16x32 = _mm512_cvtepu8_epi16(e5m2_u8x32);
+                __m512i f16_bits_i16x32 = _mm512_slli_epi16(word_u16x32, 8);
+                _mm512_store_si512(&source_tile.data[row_idx][0], f16_bits_i16x32);
+            }
+
+            nk_dots_bf16_b32x16_sapphireamx_t transposed_tile;
+            nk_dots_pack_bf16_transposed_sapphireamx_(&source_tile, &transposed_tile);
+            for (nk_size_t i = 0; i < tile_bytes; i += 64)
+                _mm512_storeu_si512((char *)tile_output + i, _mm512_load_si512((char const *)&transposed_tile + i));
+        }
+    }
+
+    // Column-remainder: widen and store row-major (depth contiguous per column)
+    if (column_remainder_count > 0) {
+        nk_size_t const remainder_start_row = column_tiles_count * tmm_rows;
+        for (nk_size_t row_idx = 0; row_idx < column_remainder_count; row_idx++) {
+            nk_e5m2_t const *src_row = b + (remainder_start_row + row_idx) * b_stride_elements;
+            nk_f16_t *dst_row = column_edge_ptr + row_idx * depth;
+            nk_size_t column_idx = 0;
+            for (; column_idx + 32 <= depth; column_idx += 32) {
+                __m256i e5m2_u8x32 = _mm256_loadu_si256((__m256i const *)(src_row + column_idx));
+                __m512i word_u16x32 = _mm512_cvtepu8_epi16(e5m2_u8x32);
+                __m512i f16_bits_i16x32 = _mm512_slli_epi16(word_u16x32, 8);
+                _mm512_storeu_si512(dst_row + column_idx, f16_bits_i16x32);
+            }
+            if (column_idx < depth) {
+                __mmask32 tail_mask = (__mmask32)((1U << (depth - column_idx)) - 1);
+                __m256i e5m2_u8x32 = _mm256_maskz_loadu_epi8(tail_mask, src_row + column_idx);
+                __m512i word_u16x32 = _mm512_cvtepu8_epi16(e5m2_u8x32);
+                __m512i f16_bits_i16x32 = _mm512_slli_epi16(word_u16x32, 8);
+                _mm512_mask_storeu_epi16(dst_row + column_idx, tail_mask, f16_bits_i16x32);
+            }
+        }
+    }
+
+    nk_size_t norms_offset = column_edge_offset +
+                             (column_remainder_count > 0 ? column_remainder_count * depth * sizeof(nk_f16_t) : 0);
+    header->norms_byte_offset = (nk_u32_t)norms_offset;
+    nk_f32_t *norms = (nk_f32_t *)((char *)b_packed + norms_offset);
+    for (nk_size_t col = 0; col < column_count; col++)
+        norms[col] = nk_dots_reduce_sumsq_e5m2_(b + col * b_stride_elements, depth);
+}
+
+NK_PUBLIC void nk_dots_packed_e5m2_graniteamx(             //
+    nk_e5m2_t const *a, void const *b_packed, nk_f32_t *c, //
+    nk_size_t rows_count, nk_size_t cols_count, nk_size_t depth, nk_size_t a_stride_bytes, nk_size_t c_stride_bytes) {
+    nk_unused_(cols_count);
+
+    nk_dots_amx_packed_header_t const *header = (nk_dots_amx_packed_header_t const *)b_packed;
+    nk_size_t const column_tiles_count = header->full_column_tiles;
+    nk_size_t const depth_tiles_count = header->full_depth_tiles;
+    nk_size_t const column_remainder_count = header->column_remainder_count;
+
+    nk_f16_t const *b_tiles_base = (nk_f16_t const *)((char const *)b_packed + sizeof(nk_dots_amx_packed_header_t));
+    nk_f16_t const *col_edge_ptr = (nk_f16_t const *)((char const *)b_packed + header->column_edge_offset);
+
+    nk_size_t const a_stride_elements = a_stride_bytes; // E5M2: 1 byte per element
+    nk_size_t const c_stride_elements = c_stride_bytes / sizeof(nk_f32_t);
+
+    nk_size_t const tile_depth = 32;
+    nk_size_t const tile_size = 512;
+    nk_size_t const full_cols = column_tiles_count * 16;
+
+    nk_size_t const row_blocks_count = nk_size_divide_round_up_(rows_count, 32);
+    nk_size_t const col_blocks_count = column_tiles_count / 2;
+
+    if (depth_tiles_count == 0) return;
+
+    nk_dots_e5m2_a16x32_graniteamx_t a_tile_top, a_tile_bottom;
+    nk_dots_e5m2_state2x2_graniteamx_t c_accum_buffer;
+
+    nk_size_t const full_depth_tiles_count = depth / tile_depth;
+    nk_size_t const depth_remainder = depth % tile_depth;
+
+    nk_amx_tile_configure_sapphireamx_();
+
+    for (nk_size_t row_block_idx = 0; row_block_idx < row_blocks_count; row_block_idx++) {
+        nk_size_t const row_block_start = row_block_idx * 32;
+        nk_size_t const valid_rows_count = (row_block_start + 32 <= rows_count) ? 32 : (rows_count - row_block_start);
+        nk_size_t const is_full_row_block = (valid_rows_count == 32);
+
+        for (nk_size_t column_block_idx = 0; column_block_idx < col_blocks_count; column_block_idx++) {
+            nk_size_t const col_block_start = column_block_idx * 32;
+            nk_size_t const b_column_left_base = (column_block_idx * 2) * depth_tiles_count;
+            nk_size_t const b_column_right_base = (column_block_idx * 2 + 1) * depth_tiles_count;
+
+            _tile_zero(4);
+            _tile_zero(5);
+            _tile_zero(6);
+            _tile_zero(7);
+
+            // E5M2 A needs widen-then-tile-load every time (no fast in-place path), so we always
+            // route through the widen buffer a_tile_{top,bottom} rather than `_tile_loadd` from A.
+            for (nk_size_t depth_tile_idx = 0; depth_tile_idx < depth_tiles_count; depth_tile_idx++) {
+                nk_size_t const depth_offset = depth_tile_idx * tile_depth;
+                nk_size_t const valid_depth = (depth_tile_idx < full_depth_tiles_count) ? tile_depth : depth_remainder;
+                nk_size_t const rows_in_high_tile = is_full_row_block
+                                                        ? 16
+                                                        : ((valid_rows_count > 16) ? 16 : valid_rows_count);
+                nk_size_t const rows_in_low_tile = is_full_row_block
+                                                       ? 16
+                                                       : ((valid_rows_count > 16) ? valid_rows_count - 16 : 0);
+
+                nk_dots_e5m2_load_a_graniteamx_(&a_tile_top, a + row_block_start * a_stride_elements + depth_offset,
+                                                a_stride_elements, rows_in_high_tile, valid_depth);
+                if (rows_in_low_tile > 0) {
+                    nk_dots_e5m2_load_a_graniteamx_(&a_tile_bottom,
+                                                    a + (row_block_start + 16) * a_stride_elements + depth_offset,
+                                                    a_stride_elements, rows_in_low_tile, valid_depth);
+                }
+
+                nk_dots_f16_b32x16_graniteamx_t const *b_tile_left = //
+                    (nk_dots_f16_b32x16_graniteamx_t const *)(b_tiles_base +
+                                                              (b_column_left_base + depth_tile_idx) * tile_size);
+                nk_dots_f16_b32x16_graniteamx_t const *b_tile_right = //
+                    (nk_dots_f16_b32x16_graniteamx_t const *)(b_tiles_base +
+                                                              (b_column_right_base + depth_tile_idx) * tile_size);
+
+                _tile_loadd(0, a_tile_top.data, 64);
+                if (rows_in_low_tile > 0) _tile_loadd(1, a_tile_bottom.data, 64);
+                _tile_loadd(2, b_tile_left->data, 64);
+                _tile_loadd(3, b_tile_right->data, 64);
+
+                _tile_dpfp16ps(4, 0, 2);
+                _tile_dpfp16ps(5, 0, 3);
+                if (rows_in_low_tile > 0) {
+                    _tile_dpfp16ps(6, 1, 2);
+                    _tile_dpfp16ps(7, 1, 3);
+                }
+            }
+
+            if (is_full_row_block) {
+                nk_f32_t *c_block = c + row_block_start * c_stride_elements + col_block_start;
+                _tile_stored(4, c_block, c_stride_bytes);
+                _tile_stored(5, c_block + 16, c_stride_bytes);
+                _tile_stored(6, (nk_f32_t *)((char *)c_block + 16 * c_stride_bytes), c_stride_bytes);
+                _tile_stored(7, (nk_f32_t *)((char *)c_block + 16 * c_stride_bytes) + 16, c_stride_bytes);
+            }
+            else {
+                _tile_stored(4, c_accum_buffer.c[0][0].data, 64);
+                _tile_stored(5, c_accum_buffer.c[0][1].data, 64);
+                _tile_stored(6, c_accum_buffer.c[1][0].data, 64);
+                _tile_stored(7, c_accum_buffer.c[1][1].data, 64);
+                nk_dots_f16_output2x2_graniteamx_(&c_accum_buffer,
+                                                  c + row_block_start * c_stride_elements + col_block_start,
+                                                  c_stride_elements, valid_rows_count, 32);
+            }
+        }
+    }
+
+    // Odd column tile
+    if (column_tiles_count % 2 == 1) {
+        nk_size_t const column_tile_idx = column_tiles_count - 1;
+        nk_size_t const col_start = column_tile_idx * 16;
+        nk_size_t const b_column_base = column_tile_idx * depth_tiles_count;
+
+        for (nk_size_t row_block_idx = 0; row_block_idx < row_blocks_count; row_block_idx++) {
+            nk_size_t const row_block_start = row_block_idx * 32;
+            nk_size_t const valid_rows_count = (row_block_start + 32 <= rows_count) ? 32
+                                                                                    : (rows_count - row_block_start);
+            nk_size_t const rows_in_high_tile = (valid_rows_count > 16) ? 16 : valid_rows_count;
+            nk_size_t const rows_in_low_tile = (valid_rows_count > 16) ? valid_rows_count - 16 : 0;
+
+            nk_dots_f16_state_graniteamx_t c_high_state, c_low_state;
+
+            _tile_zero(4);
+            _tile_zero(6);
+
+            for (nk_size_t depth_tile_idx = 0; depth_tile_idx < depth_tiles_count; depth_tile_idx++) {
+                nk_size_t const depth_offset = depth_tile_idx * tile_depth;
+                nk_size_t const valid_depth = (depth_tile_idx < full_depth_tiles_count) ? tile_depth : depth_remainder;
+
+                nk_dots_e5m2_load_a_graniteamx_(&a_tile_top, a + row_block_start * a_stride_elements + depth_offset,
+                                                a_stride_elements, rows_in_high_tile, valid_depth);
+                if (rows_in_low_tile > 0) {
+                    nk_dots_e5m2_load_a_graniteamx_(&a_tile_bottom,
+                                                    a + (row_block_start + 16) * a_stride_elements + depth_offset,
+                                                    a_stride_elements, rows_in_low_tile, valid_depth);
+                }
+
+                nk_dots_f16_b32x16_graniteamx_t const *b_tile = //
+                    (nk_dots_f16_b32x16_graniteamx_t const *)(b_tiles_base +
+                                                              (b_column_base + depth_tile_idx) * tile_size);
+
+                _tile_loadd(0, a_tile_top.data, 64);
+                if (rows_in_low_tile > 0) _tile_loadd(1, a_tile_bottom.data, 64);
+                _tile_loadd(2, b_tile->data, 64);
+
+                _tile_dpfp16ps(4, 0, 2);
+                if (rows_in_low_tile > 0) _tile_dpfp16ps(6, 1, 2);
+            }
+
+            _tile_stored(4, c_high_state.data, 64);
+            if (rows_in_low_tile > 0) _tile_stored(6, c_low_state.data, 64);
+
+            nk_dots_f16_store_graniteamx_(&c_high_state, c + row_block_start * c_stride_elements + col_start,
+                                          c_stride_elements, rows_in_high_tile, 16);
+            if (rows_in_low_tile > 0) {
+                nk_dots_f16_store_graniteamx_(&c_low_state, c + (row_block_start + 16) * c_stride_elements + col_start,
+                                              c_stride_elements, rows_in_low_tile, 16);
+            }
+        }
+    }
+
+    // Column edge (fewer than 16 extra columns, stored row-major as F16 after pack)
+    if (column_remainder_count > 0) {
+        for (nk_size_t row_block_idx = 0; row_block_idx < row_blocks_count; row_block_idx++) {
+            nk_size_t const row_block_start = row_block_idx * 32;
+            nk_size_t const valid_rows_count = (row_block_start + 32 <= rows_count) ? 32
+                                                                                    : (rows_count - row_block_start);
+            nk_size_t const rows_in_high_tile = (valid_rows_count > 16) ? 16 : valid_rows_count;
+            nk_size_t const rows_in_low_tile = (valid_rows_count > 16) ? valid_rows_count - 16 : 0;
+
+            nk_dots_f16_state_graniteamx_t c_high_state, c_low_state;
+            nk_dots_bf16_a16x32_sapphireamx_t b_as_a;
+            nk_dots_bf16_b32x16_sapphireamx_t b_tile;
+
+            _tile_zero(4);
+            _tile_zero(6);
+
+            for (nk_size_t depth_tile_idx = 0; depth_tile_idx < depth_tiles_count; depth_tile_idx++) {
+                nk_size_t const depth_offset = depth_tile_idx * tile_depth;
+                nk_size_t const valid_depth = (depth_tile_idx < full_depth_tiles_count) ? tile_depth : depth_remainder;
+
+                nk_dots_e5m2_load_a_graniteamx_(&a_tile_top, a + row_block_start * a_stride_elements + depth_offset,
+                                                a_stride_elements, rows_in_high_tile, valid_depth);
+                if (rows_in_low_tile > 0) {
+                    nk_dots_e5m2_load_a_graniteamx_(&a_tile_bottom,
+                                                    a + (row_block_start + 16) * a_stride_elements + depth_offset,
+                                                    a_stride_elements, rows_in_low_tile, valid_depth);
+                }
+
+                nk_dots_bf16_load_a_sapphireamx_(&b_as_a, (nk_bf16_t const *)(col_edge_ptr + depth_offset), depth,
+                                                 column_remainder_count, valid_depth);
+                nk_dots_pack_bf16_transposed_sapphireamx_(&b_as_a, &b_tile);
+
+                _tile_loadd(0, a_tile_top.data, 64);
+                if (rows_in_low_tile > 0) _tile_loadd(1, a_tile_bottom.data, 64);
+                _tile_loadd(2, b_tile.data, 64);
+
+                _tile_dpfp16ps(4, 0, 2);
+                if (rows_in_low_tile > 0) _tile_dpfp16ps(6, 1, 2);
+            }
+
+            _tile_stored(4, c_high_state.data, 64);
+            if (rows_in_low_tile > 0) _tile_stored(6, c_low_state.data, 64);
+
+            nk_dots_f16_store_graniteamx_(&c_high_state, c + row_block_start * c_stride_elements + full_cols,
+                                          c_stride_elements, rows_in_high_tile, column_remainder_count);
+            if (rows_in_low_tile > 0) {
+                nk_dots_f16_store_graniteamx_(&c_low_state, c + (row_block_start + 16) * c_stride_elements + full_cols,
+                                              c_stride_elements, rows_in_low_tile, column_remainder_count);
+            }
+        }
+    }
+
+    _tile_release();
+}
+
+NK_PUBLIC void nk_dots_symmetric_e5m2_graniteamx(                                  //
+    nk_e5m2_t const *vectors, nk_size_t vectors_count, nk_size_t depth,            //
+    nk_size_t stride_in_bytes, nk_f32_t *result, nk_size_t result_stride_in_bytes, //
+    nk_size_t row_start, nk_size_t row_count) {
+
+    nk_size_t const stride_elements = stride_in_bytes; // E5M2: 1 byte per element
+    nk_size_t const result_stride_elements = result_stride_in_bytes / sizeof(nk_f32_t);
+
+    nk_size_t const row_end = (row_count == 0)
+                                  ? vectors_count
+                                  : (row_start + row_count < vectors_count ? row_start + row_count : vectors_count);
+
+    nk_size_t const depth_tiles = nk_size_divide_round_up_(depth, 32);
+    nk_size_t const depth_tile_groups = nk_size_divide_round_up_(depth_tiles, 3);
+
+    nk_dots_e5m2_a16x32_graniteamx_t a_tiles[3];
+    nk_dots_e5m2_a16x32_graniteamx_t b_src_tiles[3];
+    nk_dots_f16_b32x16_graniteamx_t b_tiles[3];
+    nk_dots_f16_state_graniteamx_t state;
+
+    nk_amx_tile_configure_sapphireamx_();
+
+    for (nk_size_t row_tile = row_start; row_tile < row_end; row_tile += 16) {
+        nk_size_t const valid_rows = (row_tile + 16 <= row_end) ? 16 : (row_end - row_tile);
+
+        for (nk_size_t col_tile = 0; col_tile < vectors_count; col_tile += 16) {
+            nk_size_t const valid_cols = (col_tile + 16 <= vectors_count) ? 16 : (vectors_count - col_tile);
+
+            nk_dots_f16_init_graniteamx_(&state);
+
+            for (nk_size_t depth_group_idx = 0; depth_group_idx < depth_tile_groups; depth_group_idx++) {
+                nk_size_t const depth_base = depth_group_idx * 96;
+
+                for (int tile_idx = 0; tile_idx < 3; tile_idx++) {
+                    nk_size_t const depth_start = depth_base + tile_idx * 32;
+                    nk_size_t const valid_depth = (depth_start + 32 <= depth)
+                                                      ? 32
+                                                      : (depth > depth_start ? depth - depth_start : 0);
+
+                    nk_dots_e5m2_load_a_graniteamx_(                        //
+                        &a_tiles[tile_idx],                                 //
+                        vectors + row_tile * stride_elements + depth_start, //
+                        stride_elements, valid_rows, valid_depth);
+
+                    if (row_tile == col_tile) {
+                        nk_dots_pack_bf16_transposed_sapphireamx_(
+                            (nk_dots_bf16_a16x32_sapphireamx_t const *)&a_tiles[tile_idx],
+                            (nk_dots_bf16_b32x16_sapphireamx_t *)&b_tiles[tile_idx]);
+                    }
+                    else {
+                        nk_dots_e5m2_load_a_graniteamx_(                        //
+                            &b_src_tiles[tile_idx],                             //
+                            vectors + col_tile * stride_elements + depth_start, //
+                            stride_elements, valid_cols, valid_depth);
+                        nk_dots_pack_bf16_transposed_sapphireamx_(
+                            (nk_dots_bf16_a16x32_sapphireamx_t const *)&b_src_tiles[tile_idx],
+                            (nk_dots_bf16_b32x16_sapphireamx_t *)&b_tiles[tile_idx]);
+                    }
+                }
+
+                nk_dots_f16_update_graniteamx_(                    //
+                    &state, &a_tiles[0], &a_tiles[1], &a_tiles[2], //
+                    &b_tiles[0], &b_tiles[1], &b_tiles[2]);
+            }
+
+            nk_dots_f16_store_graniteamx_(                                     //
+                &state, result + row_tile * result_stride_elements + col_tile, //
+                result_stride_elements, valid_rows, valid_cols);
+        }
+    }
+}
+
+#pragma endregion E5M2 Source
+
 #if defined(__clang__)
 #pragma clang attribute pop
 #elif defined(__GNUC__)

package/include/numkong/dots/haswell.h

@@ -115,45 +115,45 @@ nk_define_cross_packed_(dots, bf16, haswell, bf16, bf16, f32, nk_b256_vec_t, nk_
                         nk_partial_store_b32x4_haswell_,
                         /*depth_simd_dimensions=*/16, /*dimensions_per_value=*/1)
 
-/* E4M3 GEMM: depth_simd_dimensions=8 (8 e4m3s = 8 bytes) → upcasted to 8×f32 (256-bit) */
-nk_define_cross_pack_size_(dots, e4m3, haswell, e4m3, f32, /*norm_value_type=*/f32, /*depth_simd_dimensions=*/8,
+/* E4M3 GEMM: depth_simd_dimensions=32 (byte-level batch; widen inside the update helper) */
+nk_define_cross_pack_size_(dots, e4m3, haswell, e4m3, f32, /*norm_value_type=*/f32, /*depth_simd_dimensions=*/32,
                            /*dimensions_per_value=*/1)
-nk_define_cross_pack_(dots, e4m3, haswell, e4m3, f32, nk_b256_vec_t, nk_load_e4m3x8_to_f32x8_haswell_,
-                      nk_partial_load_e4m3x8_to_f32x8_haswell_, nk_store_b256_haswell_, nk_partial_store_b32x8_serial_,
-                      /*simd_width=*/8, /*norm_value_type=*/f32, nk_dots_reduce_sumsq_e4m3_,
-                      /*depth_simd_dimensions=*/8, /*dimensions_per_value=*/1)
+nk_define_cross_pack_(dots, e4m3, haswell, e4m3, f32, nk_b256_vec_t, nk_load_b256_haswell_,
+                      nk_partial_load_b8x32_serial_, nk_store_b256_haswell_, nk_partial_store_b8x32_serial_,
+                      /*simd_width=*/32, /*norm_value_type=*/f32, nk_dots_reduce_sumsq_e4m3_,
+                      /*depth_simd_dimensions=*/32, /*dimensions_per_value=*/1)
 nk_define_cross_symmetric_(dots, e4m3, haswell, e4m3, f32, nk_b256_vec_t, nk_dot_through_f32_state_haswell_t_,
-                           nk_b128_vec_t, nk_dot_through_f32_init_haswell_, nk_load_e4m3x8_to_f32x8_haswell_,
-                           nk_partial_load_e4m3x8_to_f32x8_haswell_, nk_dot_through_f32_update_haswell_,
+                           nk_b128_vec_t, nk_dot_through_f32_init_haswell_, nk_load_b256_haswell_,
+                           nk_partial_load_b8x32_serial_, nk_dot_e4m3x32_update_haswell_,
                            nk_dot_through_f32_finalize_haswell_, nk_store_b128_haswell_,
                            nk_partial_store_b32x4_haswell_,
-                           /*depth_simd_dimensions=*/8, /*dimensions_per_value=*/1)
+                           /*depth_simd_dimensions=*/32, /*dimensions_per_value=*/1)
 nk_define_cross_packed_(dots, e4m3, haswell, e4m3, f32, f32, nk_b256_vec_t, nk_dot_through_f32_state_haswell_t_,
-                        nk_b128_vec_t, nk_dot_through_f32_init_haswell_, nk_load_e4m3x8_to_f32x8_haswell_,
-                        nk_partial_load_e4m3x8_to_f32x8_haswell_, nk_load_b256_haswell_, nk_partial_load_b32x8_serial_,
-                        nk_dot_through_f32_update_haswell_, nk_dot_through_f32_finalize_haswell_,
-                        nk_store_b128_haswell_, nk_partial_store_b32x4_haswell_,
-                        /*depth_simd_dimensions=*/8, /*dimensions_per_value=*/1)
+                        nk_b128_vec_t, nk_dot_through_f32_init_haswell_, nk_load_b256_haswell_,
+                        nk_partial_load_b8x32_serial_, nk_load_b256_haswell_, nk_partial_load_b8x32_serial_,
+                        nk_dot_e4m3x32_update_haswell_, nk_dot_through_f32_finalize_haswell_, nk_store_b128_haswell_,
+                        nk_partial_store_b32x4_haswell_,
+                        /*depth_simd_dimensions=*/32, /*dimensions_per_value=*/1)
 
-/* E5M2 GEMM: depth_simd_dimensions=8 (8 e5m2s = 8 bytes) → upcasted to 8×f32 (256-bit) */
-nk_define_cross_pack_size_(dots, e5m2, haswell, e5m2, f32, /*norm_value_type=*/f32, /*depth_simd_dimensions=*/8,
+/* E5M2 GEMM: depth_simd_dimensions=32 (byte-level batch; widen inside the update helper) */
+nk_define_cross_pack_size_(dots, e5m2, haswell, e5m2, f32, /*norm_value_type=*/f32, /*depth_simd_dimensions=*/32,
                            /*dimensions_per_value=*/1)
-nk_define_cross_pack_(dots, e5m2, haswell, e5m2, f32, nk_b256_vec_t, nk_load_e5m2x8_to_f32x8_haswell_,
-                      nk_partial_load_e5m2x8_to_f32x8_haswell_, nk_store_b256_haswell_, nk_partial_store_b32x8_serial_,
-                      /*simd_width=*/8, /*norm_value_type=*/f32, nk_dots_reduce_sumsq_e5m2_,
-                      /*depth_simd_dimensions=*/8, /*dimensions_per_value=*/1)
+nk_define_cross_pack_(dots, e5m2, haswell, e5m2, f32, nk_b256_vec_t, nk_load_b256_haswell_,
+                      nk_partial_load_b8x32_serial_, nk_store_b256_haswell_, nk_partial_store_b8x32_serial_,
+                      /*simd_width=*/32, /*norm_value_type=*/f32, nk_dots_reduce_sumsq_e5m2_,
+                      /*depth_simd_dimensions=*/32, /*dimensions_per_value=*/1)
 nk_define_cross_symmetric_(dots, e5m2, haswell, e5m2, f32, nk_b256_vec_t, nk_dot_through_f32_state_haswell_t_,
-                           nk_b128_vec_t, nk_dot_through_f32_init_haswell_, nk_load_e5m2x8_to_f32x8_haswell_,
-                           nk_partial_load_e5m2x8_to_f32x8_haswell_, nk_dot_through_f32_update_haswell_,
+                           nk_b128_vec_t, nk_dot_through_f32_init_haswell_, nk_load_b256_haswell_,
+                           nk_partial_load_b8x32_serial_, nk_dot_e5m2x32_update_haswell_,
                            nk_dot_through_f32_finalize_haswell_, nk_store_b128_haswell_,
                            nk_partial_store_b32x4_haswell_,
-                           /*depth_simd_dimensions=*/8, /*dimensions_per_value=*/1)
+                           /*depth_simd_dimensions=*/32, /*dimensions_per_value=*/1)
 nk_define_cross_packed_(dots, e5m2, haswell, e5m2, f32, f32, nk_b256_vec_t, nk_dot_through_f32_state_haswell_t_,
-                        nk_b128_vec_t, nk_dot_through_f32_init_haswell_, nk_load_e5m2x8_to_f32x8_haswell_,
-                        nk_partial_load_e5m2x8_to_f32x8_haswell_, nk_load_b256_haswell_, nk_partial_load_b32x8_serial_,
-                        nk_dot_through_f32_update_haswell_, nk_dot_through_f32_finalize_haswell_,
-                        nk_store_b128_haswell_, nk_partial_store_b32x4_haswell_,
-                        /*depth_simd_dimensions=*/8, /*dimensions_per_value=*/1)
+                        nk_b128_vec_t, nk_dot_through_f32_init_haswell_, nk_load_b256_haswell_,
+                        nk_partial_load_b8x32_serial_, nk_load_b256_haswell_, nk_partial_load_b8x32_serial_,
+                        nk_dot_e5m2x32_update_haswell_, nk_dot_through_f32_finalize_haswell_, nk_store_b128_haswell_,
+                        nk_partial_store_b32x4_haswell_,
+                        /*depth_simd_dimensions=*/32, /*dimensions_per_value=*/1)
 
 /* E2M3 GEMM: integer LUT path, depth_simd_dimensions=32 (32 e2m3s = 32 bytes = AVX2 register width) */
 nk_define_cross_pack_size_(dots, e2m3, haswell, e2m3, e2m3, /*norm_value_type=*/f32, /*depth_simd_dimensions=*/32,

package/include/numkong/dots/sapphireamx.h

@@ -73,7 +73,7 @@
 #if NK_TARGET_SAPPHIREAMX
 
 #include "numkong/cast/icelake.h" // For FP8 ↔ BF16 conversions
-#include "numkong/dots/serial.h"  // For nk_dots_reduce_sumsq_bf16_
+#include "numkong/dots/serial.h"  // `nk_dots_reduce_sumsq_bf16_`
 
 #if defined(__cplusplus)
 extern "C" {