Adds Intel AMX TMUL hardware acceleration for FlashAttention using INT8

CMakeLists.txt

@@ -62,7 +62,7 @@ if(GEMMA_ONEDNN_BRGEMM OR GEMMA_ONEDNN_MATMUL)
 endif()
 include(${CMAKE_CURRENT_LIST_DIR}/cmake/GemmaFetch.cmake)
 
-gemma_fetch(highway GIT_REPOSITORY https://github.com/google/highway.git GIT_TAG 97a5dd1af1a43a8b2ccbd31e556b4139bffbdafd EXCLUDE_FROM_ALL)
+gemma_fetch(highway GIT_REPOSITORY https://github.com/google/highway.git GIT_TAG 353597727402dfc7b28e5b1474766e366ceafd24 EXCLUDE_FROM_ALL)
 
 if(CMAKE_CXX_COMPILER_ID STREQUAL "GNU" AND
    CMAKE_CXX_COMPILER_VERSION VERSION_GREATER_EQUAL 15)

MODULE.bazel

@@ -19,7 +19,7 @@ bazel_dep(name = "google_benchmark", version = "1.8.5")
 # Require a more recent version.
 git_override(
     module_name = "highway",
-    commit = "97a5dd1af1a43a8b2ccbd31e556b4139bffbdafd",
+    commit = "353597727402dfc7b28e5b1474766e366ceafd24",
     remote = "https://github.com/google/highway",
 )
 

examples/hello_world/CMakeLists.txt

@@ -35,7 +35,7 @@ else()
   endfunction()
 endif()
 
-gemma_fetch(highway GIT_REPOSITORY https://github.com/google/highway.git GIT_TAG 97a5dd1af1a43a8b2ccbd31e556b4139bffbdafd)
+gemma_fetch(highway GIT_REPOSITORY https://github.com/google/highway.git GIT_TAG 353597727402dfc7b28e5b1474766e366ceafd24)
 if(CMAKE_CXX_COMPILER_ID STREQUAL "GNU" AND
    CMAKE_CXX_COMPILER_VERSION VERSION_GREATER_EQUAL 15)
   target_compile_definitions(hwy PUBLIC HWY_DISABLED_TARGETS=HWY_AVX10_2)

examples/simplified_gemma/CMakeLists.txt

@@ -35,7 +35,7 @@ else()
   endfunction()
 endif()
 
-gemma_fetch(highway GIT_REPOSITORY https://github.com/google/highway.git GIT_TAG 97a5dd1af1a43a8b2ccbd31e556b4139bffbdafd)
+gemma_fetch(highway GIT_REPOSITORY https://github.com/google/highway.git GIT_TAG 353597727402dfc7b28e5b1474766e366ceafd24)
 if(CMAKE_CXX_COMPILER_ID STREQUAL "GNU" AND
    CMAKE_CXX_COMPILER_VERSION VERSION_GREATER_EQUAL 15)
   target_compile_definitions(hwy PUBLIC HWY_DISABLED_TARGETS=HWY_AVX10_2)

gemma/activations.h

@@ -131,7 +131,19 @@ struct AttentionActivations {
         inv_timescale_global(CreateInvTimescale(
             allocator, max_qkv_dim,
             layer_config.post_qk == PostQKType::HalfRope,
-            config.global_rope_theta, config.partial_rotary_factor)) {
+            config.global_rope_theta, config.partial_rotary_factor)),
+        s_att_q(config.is_encoder_decoder ? config.decoder_num_layers
+                                          : config.num_layers,
+                max_workers),
+        s_att_k(config.is_encoder_decoder ? config.decoder_num_layers
+                                          : config.num_layers,
+                max_workers),
+        s_att_v(config.is_encoder_decoder ? config.decoder_num_layers
+                                          : config.num_layers,
+                max_workers),
+        s_att_out(config.is_encoder_decoder ? config.decoder_num_layers
+                                            : config.num_layers,
+                  max_workers) {
     // Batch size can be 0 in experimental code so do not assert.
     if (batch_size == 0) {
       static std::atomic_flag warned = ATOMIC_FLAG_INIT;
@@ -246,6 +258,13 @@ struct AttentionActivations {
   // Rope
   MatStorageT<float> inv_timescale;
   MatStorageT<float> inv_timescale_global;
+
+  // Only active when GCPP_TENSOR_STATS.
+  TensorStats s_att_q;
+  TensorStats s_att_k;
+  TensorStats s_att_v;
+  TensorStats s_att_out;
+
   // Replication factor to help evenly share work over threads.
   static constexpr size_t kThreadReplicationFactor = 4;
 };
@@ -301,6 +320,10 @@ struct AttentionActivationsPtrs {
     int8_queries = &activations.int8_queries;
     float_queries = &activations.float_queries;
     q_scales = &activations.q_scales;
+    s_att_q = &activations.s_att_q;
+    s_att_k = &activations.s_att_k;
+    s_att_v = &activations.s_att_v;
+    s_att_out = &activations.s_att_out;
   }
 
   void SetBatchSize(size_t batch_size) {
@@ -383,6 +406,12 @@ struct AttentionActivationsPtrs {
   hwy::Divisor div_heads;
   // Query scaling factor for attention computation.
   float query_scale;
+
+  // Only active when GCPP_TENSOR_STATS.
+  TensorStats* s_att_q = nullptr;
+  TensorStats* s_att_k = nullptr;
+  TensorStats* s_att_v = nullptr;
+  TensorStats* s_att_out = nullptr;
 };
 
 static inline size_t MoEBatchSize(const LayerConfig& layer_config,
@@ -646,6 +675,10 @@ struct Activations {
   }
 
   ~Activations() {
+    attention_storage.s_att_q.ReduceAndPrint("att_q");
+    attention_storage.s_att_k.ReduceAndPrint("att_k");
+    attention_storage.s_att_v.ReduceAndPrint("att_v");
+    attention_storage.s_att_out.ReduceAndPrint("att_out");
     s_ffw_in.ReduceAndPrint("ffw_in");
     s_ffw_hidden.ReduceAndPrint("ffw_hidden");
     s_ffw_out.ReduceAndPrint("ffw_out");

gemma/configs.cc

@@ -1304,6 +1304,7 @@ constexpr std::pair<const char*, AttentionImpl> kAttentionImplNameToEnum[] = {
     {"flash_matrix_accumulation", AttentionImpl::kFlashMatrixAccumulation},
     {"int8_matrix_accumulation", AttentionImpl::kInt8MatrixAccumulation},
     {"flash_amx", AttentionImpl::kFlashAMX},
+    {"flash_amx_int8", AttentionImpl::kFlashAMXInt8},
 };
 
 std::string GetAttentionImplName(AttentionImpl impl) {

gemma/configs.h

@@ -151,6 +151,7 @@ enum class AttentionImpl {
   kInt8MatrixAccumulation,
   kFlashTransposedQsInt8,
   kFlashAMX,
+  kFlashAMXInt8,
   kSentinel,
 };
 
@@ -164,14 +165,23 @@ static inline bool IsTiledAttention(AttentionImpl impl) {
          impl == AttentionImpl::kFlashTransposedQsInt8 ||
          impl == AttentionImpl::kInt8MatrixAccumulation ||
          impl == AttentionImpl::kFlashMatrixAccumulation ||
-         impl == AttentionImpl::kFlashAMX;
+         impl == AttentionImpl::kFlashAMX ||
+         impl == AttentionImpl::kFlashAMXInt8;
 }
 
 static inline bool IsBF16TransposedQsAttention(AttentionImpl impl) {
   return impl == AttentionImpl::kFlashTransposedQsBF16 ||
          impl == AttentionImpl::kFlashAMX;
 }
 
+// Int8 implementations sharing the VNNI KV tile layout (K as
+// [qkv_dim/4][kTileSize][4], V as [kTileSize/4][qkv_dim][4], BF16 K/V scales,
+// then int32 K sums) and int8 queries with per-query scales.
+static inline bool IsInt8VNNIAttention(AttentionImpl impl) {
+  return impl == AttentionImpl::kFlashTransposedQsInt8 ||
+         impl == AttentionImpl::kFlashAMXInt8;
+}
+
 // Post attention and ffw normalization type.
 enum class PostNormType {
   None,

gemma/flash_attention.cc

@@ -1882,6 +1882,27 @@ void DispatchTileFlashAttentionReturnExpSumsAndMaxLogitsAMX(
       att_out, exp_denominator_sums, max_logits);
 }
 
+void DispatchTileFlashAttentionReturnExpSumsAndMaxLogitsAMXInt8(
+    hwy::Span<const MatPtr> kvs, size_t q_count,
+    const int8_t* HWY_RESTRICT q_base, const hwy::Span<const float> q_scales,
+    hwy::Span<const size_t> start_pos_per_query,
+    hwy::Span<const size_t> last_pos_per_query, const float att_cap,
+    MatPtrT<float>& att_out, float* HWY_RESTRICT exp_denominator_sums,
+    float* HWY_RESTRICT max_logits,
+    hwy::AlignedVector<uint8_t>* worker_workspace) {
+  TileFlashAttentionReturnExpSumsAndMaxLogitsAMXInt8(
+      kvs, q_count, q_base, q_scales, start_pos_per_query, last_pos_per_query,
+      att_cap, att_out, exp_denominator_sums, max_logits, worker_workspace);
+}
+
+bool AmxInt8Available() {
+#if GEMMA_HAVE_AMX_INT8
+  return hwy::HaveTile64BMatMulI8();
+#else
+  return false;
+#endif
+}
+
 void DispatchTileFlashAttentionReturnExpSumsAndMaxLogitsInt16(
     hwy::Span<const MatPtr> kvs, size_t q_count,
     const int16_t* HWY_RESTRICT q_base, const hwy::Span<const float> q_scales,

gemma/flash_attention.h

@@ -77,6 +77,18 @@ namespace gcpp {
       MatPtrT<float>& att_out, float* HWY_RESTRICT exp_denominator_sums,      \
       float* HWY_RESTRICT max_logits);                                        \
                                                                               \
+  void DispatchTileFlashAttentionReturnExpSumsAndMaxLogitsAMXInt8(            \
+      hwy::Span<const MatPtr> kvs, size_t q_count,                            \
+      const int8_t* HWY_RESTRICT q_base, hwy::Span<const float> q_scales,     \
+      hwy::Span<const size_t> start_pos_per_query,                            \
+      hwy::Span<const size_t> last_pos_per_query, const float att_cap,        \
+      MatPtrT<float>& att_out, float* HWY_RESTRICT exp_denominator_sums,      \
+      float* HWY_RESTRICT max_logits,                                         \
+      hwy::AlignedVector<uint8_t>* worker_workspace);                         \
+                                                                              \
+  /* True if the AMX-INT8 kernel is compiled in and usable on this CPU. */    \
+  bool AmxInt8Available();                                                    \
+                                                                              \
   void DispatchTileFlashAttentionReturnExpSumsAndMaxLogitsInt16(              \
       hwy::Span<const MatPtr> kvs, size_t q_count,                            \
       const int16_t* HWY_RESTRICT q_base, hwy::Span<const float> q_scales,    \