[executorch][native] Keep float scales of int4 ggml embeddings exact

backends/native/runtime/engine/ggml/GgmlExecutor.cpp

@@ -590,8 +590,32 @@ class GgmlExecutable final : public EngineExecutable {
     return bytes;
   }
 
-  std::vector<uint8_t> load_q4_0(const TensorMeta& meta, ByteSpan packed)
-      const {
+  // Scales that fp16 would round, of an int4 constant only embeddings read.
+  // Q4_0 blocks hold fp16 scales, and a linear's mul_mat reads them there, so
+  // only such a weight loads as unit-scale blocks whose gathered rows are
+  // multiplied by these scales.
+  std::optional<std::vector<float>> exact_row_scales(ValueId id) const {
+    const TensorMeta& meta = method_.graph.value(id).tensor_meta();
+    const AffineGroupQuant* quant = affine_quant(meta);
+    if (quant == nullptr || quant->scale_dtype == ScalarType::Half ||
+        constant_binding(id) == nullptr ||
+        !std::ranges::all_of(
+            method_.graph.value(id).consumer_ids, [this](NodeId consumer) {
+              return op_key(method_.graph.node(consumer).target) ==
+                  "aten.embedding.default";
+            })) {
+      return std::nullopt;
+    }
+    std::vector<float> scales = load_q4_scales(meta);
+    if (std::ranges::all_of(scales, [](float scale) {
+          return ggml_fp16_to_fp32(ggml_fp32_to_fp16(scale)) == scale;
+        })) {
+      return std::nullopt;
+    }
+    return scales;
+  }
+
+  std::vector<float> load_q4_scales(const TensorMeta& meta) const {
     const AffineGroupQuant& quant = q4_affine(meta);
     const std::optional<OwnedBytes> scales =
         read_constant(quant.scale_data_key);
@@ -610,8 +634,7 @@ class GgmlExecutable final : public EngineExecutable {
             quant.zero_point_data_key + "'");
       }
     }
-    return affine_int4_to_q4_0(
-        meta, packed, to_float_vector(scales->span(), quant.scale_dtype));
+    return to_float_vector(scales->span(), quant.scale_dtype);
   }
 
   // The binding of a constant with data that no op writes.
@@ -839,14 +862,29 @@ class GgmlExecutable final : public EngineExecutable {
   }
 
   void initialize_persistent_tensors(size_t* materialized_constant_bytes) {
-    const size_t count = static_cast<size_t>(std::ranges::count_if(
-        method_.data_bindings,
-        [this](const DataBinding& binding) { return is_persistent(binding); }));
+    size_t count = 0;
+    // The float scales of int4 weights that load as unit-scale blocks.
+    std::unordered_map<ValueId, std::vector<float>> float_scales;
+    for (const DataBinding& binding : method_.data_bindings) {
+      if (!is_persistent(binding)) {
+        continue;
+      }
+      ++count;
+      const TensorMeta& meta =
+          method_.graph.value(binding.value_id).tensor_meta();
+      if (affine_quant(meta) != nullptr || packed_quant(meta) != nullptr) {
+        check_q4_consumers(binding.value_id, binding.key);
+      }
+      if (std::optional<std::vector<float>> scales =
+              exact_row_scales(binding.value_id)) {
+        float_scales.emplace(binding.value_id, std::move(*scales));
+      }
+    }
     if (count == 0) {
       return;
     }
     ggml_init_params params{
-        .mem_size = ggml_tensor_overhead() * (count + 16),
+        .mem_size = ggml_tensor_overhead() * (count + float_scales.size() + 16),
         .mem_buffer = nullptr,
         .no_alloc = true,
     };
@@ -862,13 +900,19 @@ class GgmlExecutable final : public EngineExecutable {
       }
       const TensorMeta& meta =
           method_.graph.value(binding.value_id).tensor_meta();
-      if (affine_quant(meta) != nullptr || packed_quant(meta) != nullptr) {
-        check_q4_consumers(binding.value_id, binding.key);
-      }
       ggml_tensor* tensor =
           new_tensor(persistent_context_.get(), compute_type(meta), meta.sizes);
       ggml_set_name(tensor, binding.key.c_str());
       persistent_.emplace(binding.value_id, tensor);
+      if (float_scales.contains(binding.value_id)) {
+        const AffineGroupQuant& quant = q4_affine(meta);
+        ggml_tensor* scales = new_tensor(
+            persistent_context_.get(),
+            GGML_TYPE_F32,
+            {meta.sizes[0], meta.sizes[1] / quant.group_size});
+        ggml_set_name(scales, quant.scale_data_key.c_str());
+        row_scales_.emplace(binding.value_id, scales);
+      }
     }
 
     persistent_buffer_.reset(ggml_backend_alloc_ctx_tensors(
@@ -937,7 +981,17 @@ class GgmlExecutable final : public EngineExecutable {
       ByteSpan span = bytes->span();
       std::vector<uint8_t> converted;
       if (is_q4) {
-        converted = load_q4_0(meta, span);
+        const auto entry = float_scales.find(binding.value_id);
+        if (entry == float_scales.end()) {
+          converted = affine_int4_to_q4_0(meta, span, load_q4_scales(meta));
+        } else {
+          converted = affine_int4_to_q4_0(
+              meta, span, std::vector<float>(entry->second.size(), 1.0f));
+          ggml_tensor* row_scales = row_scales_.at(binding.value_id);
+          ggml_backend_tensor_set(
+              row_scales, entry->second.data(), 0, ggml_nbytes(row_scales));
+          loaded += ggml_nbytes(row_scales);
+        }
         span = converted;
       } else if (narrow) {
         converted = narrow_to_int32(span);
@@ -1309,6 +1363,8 @@ class GgmlExecutable final : public EngineExecutable {
   ggml_backend_buffer_ptr output_staging_;
   std::vector<uint8_t> output_host_;
   std::unordered_map<ValueId, ggml_tensor*> persistent_;
+  // Float scales of the Q4_0 weights with unit block scales, by weight.
+  std::unordered_map<ValueId, ggml_tensor*> row_scales_;
   bool executed_ = false;
   std::unordered_map<ValueId, std::vector<uint8_t>> rope_freq_factors_;
   // Per-output-channel factors applied to convolution weights at load.
@@ -1333,6 +1389,7 @@ GraphInstance GgmlExecutable::build_graph(
                                inputs_,
                                outputs_,
                                persistent_,
+                               row_scales_,
                                rope_freq_factors_,
                                batch_norms_,
                                std::move(truncated),

backends/native/runtime/engine/ggml/GgmlGraphBuilder.cpp

@@ -201,6 +201,7 @@ GraphBuilder::GraphBuilder(
     const std::vector<InputSlot>& inputs,
     const std::vector<OutputSlot>& outputs,
     const std::unordered_map<ValueId, ggml_tensor*>& persistent,
+    const std::unordered_map<ValueId, ggml_tensor*>& row_scales,
     const std::unordered_map<ValueId, std::vector<uint8_t>>& rope_factor_bytes,
     const FoldedBatchNorms& batch_norms,
     std::unordered_map<ValueId, int64_t> key_extents,
@@ -209,6 +210,7 @@ GraphBuilder::GraphBuilder(
       inputs_(inputs),
       outputs_(outputs),
       persistent_(persistent),
+      row_scales_(row_scales),
       rope_factor_bytes_(rope_factor_bytes),
       batch_norms_(batch_norms),
       key_extents_(std::move(key_extents)),

backends/native/runtime/engine/ggml/GgmlGraphBuilder.h

@@ -132,6 +132,7 @@ class GraphBuilder {
       const std::vector<InputSlot>& inputs,
       const std::vector<OutputSlot>& outputs,
       const std::unordered_map<ValueId, ggml_tensor*>& persistent,
+      const std::unordered_map<ValueId, ggml_tensor*>& row_scales,
       const std::unordered_map<ValueId, std::vector<uint8_t>>&
           rope_factor_bytes,
       const FoldedBatchNorms& batch_norms,
@@ -283,6 +284,7 @@ class GraphBuilder {
   const std::vector<InputSlot>& inputs_;
   const std::vector<OutputSlot>& outputs_;
   const std::unordered_map<ValueId, ggml_tensor*>& persistent_;
+  const std::unordered_map<ValueId, ggml_tensor*>& row_scales_;
   const std::unordered_map<ValueId, std::vector<uint8_t>>& rope_factor_bytes_;
   const FoldedBatchNorms& batch_norms_;
   // Mask values (and their producers) computed over a key prefix only.

backends/native/runtime/engine/ggml/GgmlLowerLinear.cpp

@@ -48,7 +48,8 @@ void GraphBuilder::lower_linear(const Node& node) {
 }
 
 void GraphBuilder::lower_row_gather(const Node& node) {
-  const TensorRef weight = tensor_arg(node, "weight", 0);
+  const ValueId weight_id = tensor_arg_id(node, "weight", 0);
+  const TensorRef weight = tensor(weight_id);
   const Argument& arg = require_arg(node, "indices", 1).arg;
   ValueId indices_id = kInvalid;
   if (op_key(node.target) == "aten.embedding.default") {
@@ -79,8 +80,18 @@ void GraphBuilder::lower_row_gather(const Node& node) {
       reshape(ctx, indices.tensor, {ggml_nelements(indices.tensor)});
   std::vector<int64_t> sizes = indices.sizes;
   sizes.push_back(weight.sizes.at(1));
-  ggml_tensor* result =
-      reshape(ctx, ggml_get_rows(ctx, weight.tensor, flat), sizes);
+  ggml_tensor* rows = ggml_get_rows(ctx, weight.tensor, flat);
+  const auto found = row_scales_.find(weight_id);
+  if (found != row_scales_.end()) {
+    const int64_t groups = found->second->ne[0];
+    const int64_t count = ggml_nelements(flat);
+    rows = ggml_mul(
+        ctx,
+        ggml_reshape_3d(ctx, rows, weight.sizes[1] / groups, groups, count),
+        ggml_reshape_3d(
+            ctx, ggml_get_rows(ctx, found->second, flat), 1, groups, count));
+  }
+  ggml_tensor* result = reshape(ctx, rows, sizes);
   set_single_output(node, TensorRef{result, std::move(sizes)});
 }
 

backends/native/runtime/engine/ggml/GgmlTypes.cpp

@@ -105,12 +105,13 @@ const AffineGroupQuant* affine_quant(const TensorMeta& meta) {
 }
 
 // Symmetric int4 groups that are whole multiples of 32 map onto Q4_0 blocks.
-// Q4_0 stores the scale as fp16, so scales are narrowed at load.
+// Q4_0 stores the scale as fp16, so scales are narrowed at load unless the
+// executor keeps them as row scales.
 const AffineGroupQuant& q4_affine(const TensorMeta& meta) {
   const AffineGroupQuant* quant = affine_quant(meta);
   if (quant == nullptr || quant->quant_min != -8 || quant->quant_max != 7 ||
       quant->group_size <= 0 || quant->group_size % kQ40BlockElements != 0 ||
-      meta.sizes.size() != 2 || meta.sizes[0] <= 0 ||
+      meta.sizes.size() != 2 || meta.sizes[0] <= 0 || meta.sizes[1] <= 0 ||
       meta.sizes[1] % quant->group_size != 0) {
     throw std::runtime_error(
         "ggml: affine weights must be int4 [N, K] with groups divisible by 32");

backends/native/test/CMakeLists.txt

@@ -162,20 +162,13 @@ if(EXECUTORCH_BUILD_NATIVE_GGML)
   # native_ggml_sample_test_act_quant lets ggml quantize the activations of
   # weight-only quantized linears, judged against golden-int8_act.
   # native_ggml_sample_test_no_linear judges the other ops without it.
-  #
-  # TODO: Drop the int4wo f32 embedding filters once ggml keeps their float
-  # scales exact (next diff).
-  add_test(
-    NAME native_ggml_sample_test_act_quant
-    COMMAND
-      native_ggml_sample_test --packages=${_native_op_samples} --device=cpu
-      --allow_activation_quant_for_weight_only
-      --gtest_filter=-embedding_int4wo_*.*_f32
-  )
-  add_test(
-    NAME native_ggml_sample_test_no_linear
-    COMMAND native_ggml_sample_test --packages=${_native_op_samples}
-            --device=cpu --gtest_filter=-linear_*:embedding_int4wo_*.*_f32
+  add_test(NAME native_ggml_sample_test_act_quant
+           COMMAND native_ggml_sample_test --packages=${_native_op_samples}
+                   --device=cpu --allow_activation_quant_for_weight_only
+  )
+  add_test(NAME native_ggml_sample_test_no_linear
+           COMMAND native_ggml_sample_test --packages=${_native_op_samples}
+                   --device=cpu --gtest_filter=-linear_*
   )
   target_compile_options(
     native_ggml_sample_test

backends/native/test/runtime/ggml/BUCK

@@ -116,12 +116,10 @@ _VARIANTS = [
         **v.get("compat", {})
     )
     for v in _VARIANTS
-    # TODO: drop the int4wo f32 embedding filters once ggml keeps their float
-    # scales exact (next diff).
     for name, flags in [
         ("", []),
-        ("_act_quant", ["--allow_activation_quant_for_weight_only", "--gtest_filter=-embedding_int4wo_*.*_f32"]),
-        ("_no_linear", ["--gtest_filter=-linear_*:embedding_int4wo_*.*_f32"]),
+        ("_act_quant", ["--allow_activation_quant_for_weight_only"]),
+        ("_no_linear", ["--gtest_filter=-linear_*"]),
     ]
 ]