fix: stable fgb.ps dot product and N-split for 2048 harts

ggml/src/ggml-et/CMakeLists.txt

@@ -38,6 +38,16 @@ set(KERNELS
     mul_mat_Q8_0
     mul_mat_Q4_0
     mul_mat_Q4_0_matrix_engine
+    mul_mat_Q4_K
+    mul_mat_Q4_K_matrix_engine
+    mul_mat_Q6_K
+    mul_mat_Q6_K_matrix_engine
+    mul_mat_Q2_K
+    mul_mat_Q2_K_matrix_engine
+    mul_mat_Q3_K
+    mul_mat_Q3_K_matrix_engine
+    mul_mat_Q5_K
+    mul_mat_Q5_K_matrix_engine
     mul_mat_f16
     mul_mat_f16_matrix_engine
     rope_f32
@@ -103,6 +113,11 @@ set(UBERKERNEL_SUPPORTED_KERNELS
     mul_mat_f32_matrix_engine
     mul_mat_id_f32
     mul_mat_Q4_0
+    mul_mat_Q4_K
+    mul_mat_Q2_K
+    mul_mat_Q3_K
+    mul_mat_Q5_K
+    mul_mat_Q6_K
     mul_mat_Q8_0
     norm_f32
     pad_f32

ggml/src/ggml-et/et-kernels/src/flash_attn_ext_f16_me.c

@@ -358,9 +358,12 @@ int entry_point(struct ggml_et_flash_attn_ext_params * params, void * env) {
     const int is_hart1     = hart_id & 1;
     uint64_t  local_minion = (hart_id >> 1) & 0x1F;
 
-    struct ggml_tensor * q        = &params->src0;
-    struct ggml_tensor * k        = &params->src1;
-    struct ggml_tensor * v        = &params->src2;
+    struct ggml_tensor * q = &params->src0;
+    evict_region_past_l2(q->data, tensor_bytes_fa(q));
+    struct ggml_tensor * k = &params->src1;
+    evict_region_past_l2(k->data, tensor_bytes_fa(k));
+    struct ggml_tensor * v = &params->src2;
+    evict_region_past_l2(v->data, tensor_bytes_fa(v));
     struct ggml_tensor * dst      = &params->dst;
     const int32_t        has_mask = params->has_mask;
     struct ggml_tensor * mask     = has_mask ? &params->mask : (struct ggml_tensor *) 0;
@@ -371,13 +374,13 @@ int entry_point(struct ggml_et_flash_attn_ext_params * params, void * env) {
     char *       dst_data = (char *) dst->data;
 
     // et_barrier(ET_BARRIER_GLOBAL);
-    evict_region_past_l2(q->data, tensor_bytes_fa(q));
-    evict_region_past_l2(k->data, tensor_bytes_fa(k));
-    evict_region_past_l2(v->data, tensor_bytes_fa(v));
-    if (mask) {
-        evict_region_past_l2(mask->data, tensor_bytes_fa(mask));
-    }
-    et_barrier(ET_BARRIER_GLOBAL);
+    // evict_region_past_l2(q_data,   tensor_bytes_fa(q));
+    // evict_region_past_l2(k_data,   tensor_bytes_fa(k));
+    // evict_region_past_l2(v_data,   tensor_bytes_fa(v));
+    // if (mask) {
+    // evict_region_past_l2(mask->data, tensor_bytes_fa(mask));
+    // }
+    // et_barrier(ET_BARRIER_GLOBAL);
 
     const int64_t dk  = q->ne[0];
     const int64_t nq  = q->ne[1];
@@ -462,9 +465,14 @@ int entry_point(struct ggml_et_flash_attn_ext_params * params, void * env) {
     // All teams in a shire must iterate the same number of times so the
     // per-iter shire barriers stay balanced. Teams whose assigned row is
     // past total_rows still call the barriers but skip the packing work.
+
     et_barrier(ET_BARRIER_SHIRE);
     // et_barrier(ET_BARRIER_GLOBAL);
+
     if (is_hart1) {
+        // et_barrier(ET_BARRIER_GLOBAL);
+        // et_barrier(ET_BARRIER_SHIRE);
+
         uint32_t      chunk_id = 0;
         const int64_t row_base = (int64_t) shire_id + local_tile_idx * NUM_COMPUTE_SHIRES;
 
@@ -548,10 +556,10 @@ int entry_point(struct ggml_et_flash_attn_ext_params * params, void * env) {
     }
 
     // Hart 0: tensor engine compute
-#ifndef UBERKERNEL_SUPPRESS_SCP_SETUP
-    setup_cache_scp();
-#endif
-    CLEAR_TENSOR_ERROR;
+    // #ifndef UBERKERNEL_SUPPRESS_SCP_SETUP
+    //     setup_cache_scp();
+    // #endif
+    // CLEAR_TENSOR_ERROR;
 
     // Q converted to F16 (one row at a time)
     et_fp16_t q_f16[FA_DK_MAX] __attribute__((aligned(64)));

ggml/src/ggml-et/et-kernels/src/get_rows_f32.c

@@ -137,6 +137,70 @@ static void copy_q4_K_row(float * dst, const block_q4_K * src_blocks, int64_t nu
     }
 }
 
+// Copy a row of Q2_K data to F32 destination (with dequantization)
+static void copy_q2_K_row(float * dst, const block_q2_K * src_blocks, int64_t num_elements) {
+    const int64_t num_blocks = (num_elements + QK_K - 1) / QK_K;
+
+    for (int64_t block_idx = 0; block_idx < num_blocks; block_idx++) {
+        const int64_t elements_in_block = (block_idx == num_blocks - 1) ? (num_elements - block_idx * QK_K) : QK_K;
+
+        float temp_buffer[QK_K];
+        dequantize_q2_K_block(&src_blocks[block_idx], temp_buffer);
+
+        for (int64_t i = 0; i < elements_in_block; i++) {
+            dst[block_idx * QK_K + i] = temp_buffer[i];
+        }
+    }
+}
+
+// Copy a row of Q3_K data to F32 destination (with dequantization)
+static void copy_q3_K_row(float * dst, const block_q3_K * src_blocks, int64_t num_elements) {
+    const int64_t num_blocks = (num_elements + QK_K - 1) / QK_K;
+
+    for (int64_t block_idx = 0; block_idx < num_blocks; block_idx++) {
+        const int64_t elements_in_block = (block_idx == num_blocks - 1) ? (num_elements - block_idx * QK_K) : QK_K;
+
+        float temp_buffer[QK_K];
+        dequantize_q3_K_block(&src_blocks[block_idx], temp_buffer);
+
+        for (int64_t i = 0; i < elements_in_block; i++) {
+            dst[block_idx * QK_K + i] = temp_buffer[i];
+        }
+    }
+}
+
+// Copy a row of Q5_K data to F32 destination (with dequantization)
+static void copy_q5_K_row(float * dst, const block_q5_K * src_blocks, int64_t num_elements) {
+    const int64_t num_blocks = (num_elements + QK_K - 1) / QK_K;
+
+    for (int64_t block_idx = 0; block_idx < num_blocks; block_idx++) {
+        const int64_t elements_in_block = (block_idx == num_blocks - 1) ? (num_elements - block_idx * QK_K) : QK_K;
+
+        float temp_buffer[QK_K];
+        dequantize_q5_K_block(&src_blocks[block_idx], temp_buffer);
+
+        for (int64_t i = 0; i < elements_in_block; i++) {
+            dst[block_idx * QK_K + i] = temp_buffer[i];
+        }
+    }
+}
+
+// Copy a row of Q6_K data to F32 destination (with dequantization)
+static void copy_q6_K_row(float * dst, const block_q6_K * src_blocks, int64_t num_elements) {
+    const int64_t num_blocks = (num_elements + QK_K - 1) / QK_K;
+
+    for (int64_t block_idx = 0; block_idx < num_blocks; block_idx++) {
+        const int64_t elements_in_block = (block_idx == num_blocks - 1) ? (num_elements - block_idx * QK_K) : QK_K;
+
+        float temp_buffer[QK_K];
+        dequantize_q6_K_block(&src_blocks[block_idx], temp_buffer);
+
+        for (int64_t i = 0; i < elements_in_block; i++) {
+            dst[block_idx * QK_K + i] = temp_buffer[i];
+        }
+    }
+}
+
 static void dequantize_q8_0_block_cache_aligned(const block_q8_0 * block, float * dst) {
     const int8_t * qs_ptr = block->qs;
 
@@ -521,7 +585,8 @@ int entry_point(struct ggml_et_get_rows_params * params, void * env) {
     }
 
     if (src0->type != GGML_TYPE_F32 && src0->type != GGML_TYPE_F16 && src0->type != GGML_TYPE_Q8_0 &&
-        src0->type != GGML_TYPE_Q4_0 && src0->type != GGML_TYPE_Q4_K) {
+        src0->type != GGML_TYPE_Q4_0 && src0->type != GGML_TYPE_Q4_K && src0->type != GGML_TYPE_Q6_K &&
+        src0->type != GGML_TYPE_Q2_K && src0->type != GGML_TYPE_Q3_K && src0->type != GGML_TYPE_Q5_K) {
         return -1;  // Unsupported input type
     }
 
@@ -605,6 +670,31 @@ int entry_point(struct ggml_et_get_rows_params * params, void * env) {
             const block_q4_K * src_blocks       = (const block_q4_K *) src0_data + src_block_offset;
             float *            dst_row          = dst_data + dst_offset * ne00;
             copy_q4_K_row(dst_row, src_blocks, ne00);
+        } else if (src0->type == GGML_TYPE_Q6_K) {
+            // Q6_K source: dequantize while copying
+            const int64_t      blocks_per_row   = (ne00 + QK_K - 1) / QK_K;
+            const int64_t      src_block_offset = (row_index * blocks_per_row) + (batch_offset / ne00) * blocks_per_row;
+            const block_q6_K * src_blocks       = (const block_q6_K *) src0_data + src_block_offset;
+            float *            dst_row          = dst_data + dst_offset * ne00;
+            copy_q6_K_row(dst_row, src_blocks, ne00);
+        } else if (src0->type == GGML_TYPE_Q2_K) {
+            const int64_t      blocks_per_row   = (ne00 + QK_K - 1) / QK_K;
+            const int64_t      src_block_offset = (row_index * blocks_per_row) + (batch_offset / ne00) * blocks_per_row;
+            const block_q2_K * src_blocks       = (const block_q2_K *) src0_data + src_block_offset;
+            float *            dst_row          = dst_data + dst_offset * ne00;
+            copy_q2_K_row(dst_row, src_blocks, ne00);
+        } else if (src0->type == GGML_TYPE_Q3_K) {
+            const int64_t      blocks_per_row   = (ne00 + QK_K - 1) / QK_K;
+            const int64_t      src_block_offset = (row_index * blocks_per_row) + (batch_offset / ne00) * blocks_per_row;
+            const block_q3_K * src_blocks       = (const block_q3_K *) src0_data + src_block_offset;
+            float *            dst_row          = dst_data + dst_offset * ne00;
+            copy_q3_K_row(dst_row, src_blocks, ne00);
+        } else if (src0->type == GGML_TYPE_Q5_K) {
+            const int64_t      blocks_per_row   = (ne00 + QK_K - 1) / QK_K;
+            const int64_t      src_block_offset = (row_index * blocks_per_row) + (batch_offset / ne00) * blocks_per_row;
+            const block_q5_K * src_blocks       = (const block_q5_K *) src0_data + src_block_offset;
+            float *            dst_row          = dst_data + dst_offset * ne00;
+            copy_q5_K_row(dst_row, src_blocks, ne00);
         }
     }