diff --git a/backends/cpu/runtime/providers/reference/README.md b/backends/cpu/runtime/providers/reference/README.md
index da54d68a1e4..0af63bfc4aa 100644
--- a/backends/cpu/runtime/providers/reference/README.md
+++ b/backends/cpu/runtime/providers/reference/README.md
@@ -4,7 +4,7 @@
 to registered ExecuTorch out kernels. Its implementation is named
 `Reference/registry`. Tensor dtypes map by name to ATen scalar types; quantized
 tensors are rejected because Reference routes do not interpret quantization metadata. The route table covers the
-linked portable kernel set; factory kernels stack a subset of their inputs
+linked kernel set; factory kernels stack a subset of their inputs
 because the kernels read dtype and shape from the output tensor.
 
 The provider resolves registered kernels when it is created and keeps a sorted
diff --git a/backends/cpu/runtime/providers/reference/ReferenceProvider.cpp b/backends/cpu/runtime/providers/reference/ReferenceProvider.cpp
index 73405391f46..b815aba51be 100644
--- a/backends/cpu/runtime/providers/reference/ReferenceProvider.cpp
+++ b/backends/cpu/runtime/providers/reference/ReferenceProvider.cpp
@@ -66,7 +66,7 @@ aten::ScalarType to_aten(ptn::ScalarType dtype) {
   throw std::invalid_argument("Unsupported CPU scalar type");
 }
 
-// One route per linked portable kernel. Targets pair functional overloads
+// One route per linked kernel. Targets pair functional overloads
 // with their out-variant kernels; arities come from the ATen schemas and
 // factory arg maps from the kernel implementations.
 constexpr auto kKernels = std::to_array<KernelSpec>({
@@ -87,6 +87,11 @@ constexpr auto kKernels = std::to_array<KernelSpec>({
      1,
      1,
      nullptr},
+    {"torch.ops.aten._embedding_bag.default",
+     "aten::_embedding_bag.out",
+     9,
+     4,
+     nullptr},
     {"torch.ops.aten._log_softmax.default",
      "aten::_log_softmax.out",
      3,
diff --git a/backends/cpu/test/test_runtime.py b/backends/cpu/test/test_runtime.py
index 24c72631c69..3b203962319 100644
--- a/backends/cpu/test/test_runtime.py
+++ b/backends/cpu/test/test_runtime.py
@@ -56,6 +56,64 @@ def _lower(self, model, inputs):
         # The runtime borrows the serialized program buffer.
         yield _load_for_executorch_from_buffer(program.buffer)
 
+    def test_embedding_bag_all_outputs(self):
+        class EmbeddingBag(torch.nn.Module):
+            def __init__(self, mode, weighted, last, padding):
+                super().__init__()
+                self.mode = mode
+                self.weighted = weighted
+                self.last = last
+                self.padding = padding
+
+            def forward(self, weight, indices, offsets, scales):
+                return torch.ops.aten._embedding_bag.default(
+                    weight,
+                    indices,
+                    offsets,
+                    False,
+                    self.mode,
+                    False,
+                    scales if self.weighted else None,
+                    self.last,
+                    self.padding,
+                )
+
+        weight = torch.arange(4 * 17, dtype=torch.float32).reshape(4, 17) - 20
+        scales = torch.tensor([0.5, -1.0, 2.0, 3.0])
+        i32, i64 = torch.int32, torch.int64
+        cases = (
+            # indices dtype, offsets dtype, mode, weighted, offsets, last, padding
+            (i64, i64, 0, False, [0, 2, 2], False, -1),
+            (i64, i64, 0, True, [0, 2, 2, 4], True, 1),
+            (i64, i64, 0, False, [0, 2, 2, 3], True, -1),
+            (i64, i64, 1, False, [0, 2, 2, 3], True, -1),
+            (i64, i64, 1, False, [0, 2, 2], False, 1),
+            (i64, i64, 2, False, [0, 2, 2, 4], True, 1),
+            (i32, i64, 2, False, [0, 2, 2], False, -1),
+            (i64, i32, 0, True, [0, 2, 2], False, -1),
+            (i32, i32, 1, False, [0, 2, 2, 4], True, 1),
+        )
+        for case in cases:
+            index_dtype, offset_dtype, mode, weighted, offsets, last, padding = case
+            with self.subTest(case=case):
+                model = EmbeddingBag(mode, weighted, last, padding)
+                indices = torch.tensor([1, 2, 1, 3], dtype=index_dtype)
+                offsets = torch.tensor(offsets, dtype=offset_dtype)
+                inputs = (weight, indices, offsets, scales)
+                with self._lower(model, inputs) as runtime:
+                    for current in (indices, indices.flip(0)):
+                        inputs = (weight, current, offsets, scales)
+                        actual = runtime.forward(inputs)
+                        for position, (output, expected) in enumerate(
+                            zip(actual, model(*inputs), strict=True)
+                        ):
+                            # Export types the three index outputs like
+                            # offsets; eager promotes them. The float output
+                            # (position 0) must keep its dtype.
+                            torch.testing.assert_close(
+                                output, expected, check_dtype=position == 0
+                            )
+
     def test_integer_inputs_and_outputs_preserve_precision(self):
         class Add(torch.nn.Module):
             def forward(self, x, y):
diff --git a/kernels/optimized/cpu/op_embedding_bag.cpp b/kernels/optimized/cpu/op_embedding_bag.cpp
new file mode 100644
index 00000000000..d9098fe29a9
--- /dev/null
+++ b/kernels/optimized/cpu/op_embedding_bag.cpp
@@ -0,0 +1,306 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <ATen/cpu/vec/functional.h>
+#include <ATen/cpu/vec/vec.h>
+#include <c10/util/irange.h>
+#include <executorch/kernels/portable/cpu/util/embedding_bag_util.h>
+#include <executorch/runtime/kernel/kernel_includes.h>
+#include <executorch/runtime/kernel/thread_parallel_interface.h>
+
+#include <algorithm>
+#include <tuple>
+
+namespace torch::executor::native {
+
+using Tensor = executorch::aten::Tensor;
+using ScalarType = executorch::aten::ScalarType;
+
+namespace {
+
+using Vec = at::vec::Vectorized<float>;
+
+// Export allocates offset2bag, bag_size and max_indices with the offsets
+// dtype, even though eager ATen promotes them to the wider index dtype.
+bool is_fast_path(
+    const Tensor& weight,
+    const Tensor& offsets,
+    bool include_last_offset,
+    const Tensor& out,
+    const Tensor& offset2bag,
+    const Tensor& bag_size,
+    const Tensor& max_indices) {
+  return weight.scalar_type() == ScalarType::Float &&
+      tensor_is_default_dim_order(weight) && tensor_is_default_dim_order(out) &&
+      tensor_is_default_dim_order(max_indices) &&
+      offsets.numel() > static_cast<ssize_t>(include_last_offset) &&
+      offset2bag.scalar_type() == offsets.scalar_type() &&
+      bag_size.scalar_type() == offsets.scalar_type() &&
+      max_indices.scalar_type() == offsets.scalar_type();
+}
+
+// Requires check_embedding_bag_offsets_and_indices and
+// resize_embedding_bag_outputs; every bag writes only its own output rows and
+// offset2bag range.
+template <typename CTYPE_IX, typename CTYPE_OFF>
+void embedding_bag_kernel(
+    KernelRuntimeContext& ctx,
+    const Tensor& weight,
+    const Tensor& indices,
+    const Tensor& offsets,
+    int64_t mode,
+    const std::optional<Tensor>& per_sample_weights,
+    int64_t padding_idx,
+    const EmbeddingBagLayout& layout,
+    Tensor& out,
+    Tensor& offset2bag,
+    Tensor& bag_size,
+    Tensor& max_indices) {
+  const CTYPE_IX* indices_data = indices.const_data_ptr<CTYPE_IX>();
+  const CTYPE_OFF* offsets_data = offsets.const_data_ptr<CTYPE_OFF>();
+  const int64_t num_indices = indices.numel();
+  const int64_t num_bags = out.size(0);
+  const int64_t width = weight.size(1);
+  // Partial at::vec loads and stores go through a stack buffer on NEON.
+  const int64_t vec_width = width - width % Vec::size();
+  const bool fills_offset2bag = layout.fills_offset2bag;
+  const bool truncate_last_bag = layout.truncate_last_bag;
+
+  const float* weight_data = weight.const_data_ptr<float>();
+  const float* scales = per_sample_weights.has_value()
+      ? per_sample_weights->const_data_ptr<float>()
+      : nullptr;
+  float* out_data = out.mutable_data_ptr<float>();
+  CTYPE_OFF* offset2bag_data = offset2bag.mutable_data_ptr<CTYPE_OFF>();
+  CTYPE_OFF* bag_size_data = bag_size.mutable_data_ptr<CTYPE_OFF>();
+  CTYPE_OFF* max_indices_data = max_indices.mutable_data_ptr<CTYPE_OFF>();
+
+  // Aim for about GRAIN_SIZE weight reads per task.
+  const int64_t reads_per_bag =
+      std::max<int64_t>(1, num_indices * width / num_bags);
+  const bool success = ::executorch::extension::parallel_for(
+      0,
+      num_bags,
+      std::max<int64_t>(
+          1, ::executorch::extension::internal::GRAIN_SIZE / reads_per_bag),
+      [&](const int64_t bags_begin, const int64_t bags_end) {
+        for (const auto bag : c10::irange(bags_begin, bags_end)) {
+          const int64_t begin = offsets_data[bag];
+          const int64_t end =
+              bag + 1 < num_bags ? offsets_data[bag + 1] : num_indices;
+          // truncate_last_bag implies include_last_offset, so
+          // offsets[bag + 1] exists for every bag.
+          const int64_t reduce_end =
+              truncate_last_bag ? offsets_data[bag + 1] : end;
+          float* out_row = out_data + bag * width;
+          std::fill_n(out_row, width, 0.0f);
+          CTYPE_OFF* winners = nullptr;
+          if (mode == kEmbeddingBagModeMax) {
+            winners = max_indices_data + bag * width;
+            std::fill_n(winners, width, CTYPE_OFF(0));
+          }
+          CTYPE_OFF valid = 0;
+          for (int64_t i = begin; i < end; ++i) {
+            if (fills_offset2bag) {
+              offset2bag_data[i] = static_cast<CTYPE_OFF>(bag);
+            }
+            const int64_t index = indices_data[i];
+            if (index == padding_idx) {
+              continue;
+            }
+            ++valid;
+            if (i >= reduce_end) {
+              continue;
+            }
+            const float* row = weight_data + index * width;
+            if (mode == kEmbeddingBagModeMax) {
+              for (const auto d : c10::irange(width)) {
+                if (valid == 1 || row[d] > out_row[d]) {
+                  out_row[d] = row[d];
+                  winners[d] = static_cast<CTYPE_OFF>(index);
+                }
+              }
+            } else if (scales != nullptr) {
+              const float scale = scales[i];
+              const Vec scale_vec(scale);
+              at::vec::map2<float>(
+                  [scale_vec](Vec x, Vec y) {
+                    return at::vec::fmadd(y, scale_vec, x);
+                  },
+                  out_row,
+                  out_row,
+                  row,
+                  vec_width);
+              for (const auto d : c10::irange(vec_width, width)) {
+                out_row[d] += row[d] * scale;
+              }
+            } else {
+              at::vec::map2<float>(
+                  [](Vec x, Vec y) { return x + y; },
+                  out_row,
+                  out_row,
+                  row,
+                  vec_width);
+              for (const auto d : c10::irange(vec_width, width)) {
+                out_row[d] += row[d];
+              }
+            }
+          }
+          bag_size_data[bag] =
+              mode == kEmbeddingBagModeSum ? CTYPE_OFF(0) : valid;
+          if (mode == kEmbeddingBagModeSum) {
+            max_indices_data[bag] = CTYPE_OFF(0);
+          } else if (mode == kEmbeddingBagModeMean) {
+            max_indices_data[bag] = valid;
+            if (valid != 0) {
+              const float divisor = static_cast<float>(valid);
+              const Vec divisor_vec(divisor);
+              at::vec::map<float>(
+                  [divisor_vec](Vec x) { return x / divisor_vec; },
+                  out_row,
+                  out_row,
+                  vec_width);
+              for (const auto d : c10::irange(vec_width, width)) {
+                out_row[d] /= divisor;
+              }
+            }
+          }
+        }
+      });
+  ET_KERNEL_CHECK_MSG(ctx, success, Internal, , "parallel_for failed");
+}
+
+} // namespace
+
+// Portable kernel (same-op fallback). Both libs link into the same binary.
+std::tuple<Tensor&, Tensor&, Tensor&, Tensor&> _embedding_bag_out(
+    KernelRuntimeContext& ctx,
+    const Tensor& weight,
+    const Tensor& indices,
+    const Tensor& offsets,
+    bool scale_grad_by_freq,
+    int64_t mode,
+    bool sparse,
+    const std::optional<Tensor>& per_sample_weights,
+    bool include_last_offset,
+    int64_t padding_idx,
+    Tensor& out,
+    Tensor& offset2bag,
+    Tensor& bag_size,
+    Tensor& max_indices);
+
+std::tuple<Tensor&, Tensor&, Tensor&, Tensor&> opt_embedding_bag_out(
+    KernelRuntimeContext& ctx,
+    const Tensor& weight,
+    const Tensor& indices,
+    const Tensor& offsets,
+    bool scale_grad_by_freq,
+    int64_t mode,
+    bool sparse,
+    const std::optional<Tensor>& per_sample_weights,
+    bool include_last_offset,
+    int64_t padding_idx,
+    Tensor& out,
+    Tensor& offset2bag,
+    Tensor& bag_size,
+    Tensor& max_indices) {
+  if (!is_fast_path(
+          weight,
+          offsets,
+          include_last_offset,
+          out,
+          offset2bag,
+          bag_size,
+          max_indices)) {
+    return _embedding_bag_out(
+        ctx,
+        weight,
+        indices,
+        offsets,
+        scale_grad_by_freq,
+        mode,
+        sparse,
+        per_sample_weights,
+        include_last_offset,
+        padding_idx,
+        out,
+        offset2bag,
+        bag_size,
+        max_indices);
+  }
+
+  std::tuple<Tensor&, Tensor&, Tensor&, Tensor&> ret_val(
+      out, offset2bag, bag_size, max_indices);
+
+  ET_KERNEL_CHECK(
+      ctx,
+      check_embedding_bag_args(
+          weight,
+          indices,
+          offsets,
+          mode,
+          per_sample_weights,
+          include_last_offset,
+          out,
+          offset2bag,
+          bag_size,
+          max_indices),
+      InvalidArgument,
+      ret_val);
+  ET_KERNEL_CHECK(
+      ctx,
+      check_embedding_bag_offsets_and_indices(weight, indices, offsets),
+      InvalidArgument,
+      ret_val);
+
+  EmbeddingBagLayout layout{};
+  ET_KERNEL_CHECK(
+      ctx,
+      resize_embedding_bag_outputs(
+          weight,
+          indices,
+          offsets,
+          mode,
+          per_sample_weights,
+          include_last_offset,
+          padding_idx,
+          out,
+          offset2bag,
+          bag_size,
+          max_indices,
+          layout) == Error::Ok,
+      InvalidArgument,
+      ret_val);
+
+  // @lint-ignore CLANGTIDY facebook-hte-CArray
+  static constexpr const char op_name[] = "_embedding_bag.out";
+
+  ET_SWITCH_TWO_TYPES(
+      Long, Int, indices.scalar_type(), ctx, op_name, CTYPE_IX, [&]() {
+        ET_SWITCH_TWO_TYPES(
+            Long, Int, offsets.scalar_type(), ctx, op_name, CTYPE_OFF, [&]() {
+              embedding_bag_kernel<CTYPE_IX, CTYPE_OFF>(
+                  ctx,
+                  weight,
+                  indices,
+                  offsets,
+                  mode,
+                  per_sample_weights,
+                  padding_idx,
+                  layout,
+                  out,
+                  offset2bag,
+                  bag_size,
+                  max_indices);
+            });
+      });
+
+  return ret_val;
+}
+
+} // namespace torch::executor::native
diff --git a/kernels/optimized/optimized.yaml b/kernels/optimized/optimized.yaml
index 1827411ff57..5d6d31d97a6 100644
--- a/kernels/optimized/optimized.yaml
+++ b/kernels/optimized/optimized.yaml
@@ -2,6 +2,11 @@
 #
 # This yaml file contains operators that have optimized kernels available.
 
+- op: _embedding_bag.out
+  kernels:
+    - arg_meta: null
+      kernel_name: torch::executor::opt_embedding_bag_out
+
 - op: _fft_c2r.out
   kernels:
     - arg_meta: null
diff --git a/kernels/portable/cpu/op_embedding_bag.cpp b/kernels/portable/cpu/op_embedding_bag.cpp
index 61d9f558b6b..3ab2a3512a2 100644
--- a/kernels/portable/cpu/op_embedding_bag.cpp
+++ b/kernels/portable/cpu/op_embedding_bag.cpp
@@ -7,6 +7,7 @@
  */
 
 #include <c10/util/irange.h>
+#include <executorch/kernels/portable/cpu/util/embedding_bag_util.h>
 #include <executorch/runtime/kernel/kernel_includes.h>
 
 #include <cstring>
@@ -19,19 +20,9 @@ namespace native {
 
 using Tensor = executorch::aten::Tensor;
 using ScalarType = executorch::aten::ScalarType;
-using SizesType = executorch::aten::SizesType;
 
 namespace {
 
-// aten::_embedding_bag's `mode` argument.
-constexpr int64_t kModeSum = 0;
-constexpr int64_t kModeMean = 1;
-constexpr int64_t kModeMax = 2;
-
-bool is_index_dtype(ScalarType dtype) {
-  return dtype == ScalarType::Long || dtype == ScalarType::Int;
-}
-
 // Index tensors may be Int or Long; read and write them as int64_t.
 int64_t index_at(const Tensor& t, size_t i) {
   return t.scalar_type() == ScalarType::Long
@@ -47,127 +38,6 @@ void set_index(Tensor& t, size_t i, int64_t value) {
   }
 }
 
-bool check_embedding_bag_args(
-    const Tensor& weight,
-    const Tensor& indices,
-    const Tensor& offsets,
-    int64_t mode,
-    const std::optional<Tensor>& per_sample_weights,
-    bool include_last_offset,
-    const Tensor& out,
-    const Tensor& offset2bag,
-    const Tensor& bag_size,
-    const Tensor& max_indices) {
-  ET_CHECK_OR_RETURN_FALSE(
-      weight.dim() == 2,
-      "weight must be 2-D; weight.dim() = %" ET_PRI_TENSOR_DIM,
-      weight.dim());
-  ET_CHECK_OR_RETURN_FALSE(
-      indices.dim() == 1,
-      "indices must be 1-D; indices.dim() = %" ET_PRI_TENSOR_DIM,
-      indices.dim());
-  ET_CHECK_OR_RETURN_FALSE(
-      offsets.dim() == 1,
-      "offsets must be 1-D; offsets.dim() = %" ET_PRI_TENSOR_DIM,
-      offsets.dim());
-  ET_CHECK_OR_RETURN_FALSE(
-      is_index_dtype(indices.scalar_type()),
-      "indices must be Long or Int; indices.scalar_type() = %s",
-      toString(indices.scalar_type()));
-  ET_CHECK_OR_RETURN_FALSE(
-      is_index_dtype(offsets.scalar_type()),
-      "offsets must be Long or Int; offsets.scalar_type() = %s",
-      toString(offsets.scalar_type()));
-
-  ET_LOG_AND_RETURN_IF_FALSE(tensor_is_floating_type(weight));
-  ET_LOG_AND_RETURN_IF_FALSE(tensors_have_same_dtype(weight, out));
-
-  // Eager ATen gives these the promoted indices/offsets dtype, but export
-  // plans them with meta_embedding_bag, which uses the offsets dtype.
-  for (const Tensor* t : {&offset2bag, &bag_size, &max_indices}) {
-    ET_CHECK_OR_RETURN_FALSE(
-        is_index_dtype(t->scalar_type()),
-        "index outputs must be Long or Int; got %s",
-        toString(t->scalar_type()));
-  }
-
-  ET_CHECK_OR_RETURN_FALSE(
-      mode == kModeSum || mode == kModeMean || mode == kModeMax,
-      "mode must be 0 (sum), 1 (mean) or 2 (max); mode = %" PRId64,
-      mode);
-
-  if (include_last_offset) {
-    ET_CHECK_OR_RETURN_FALSE(
-        offsets.numel() >= 1,
-        "offsets must hold at least the final offset when include_last_offset is set");
-  }
-
-  if (per_sample_weights.has_value()) {
-    const Tensor& psw = per_sample_weights.value();
-    ET_CHECK_OR_RETURN_FALSE(psw.dim() == 1, "per_sample_weights must be 1-D");
-    ET_CHECK_OR_RETURN_FALSE(
-        mode == kModeSum,
-        "per_sample_weights is only supported for mode 0 (sum); mode = %" PRId64,
-        mode);
-    ET_LOG_AND_RETURN_IF_FALSE(tensors_have_same_dtype(weight, psw));
-    ET_CHECK_OR_RETURN_FALSE(
-        psw.numel() == indices.numel(),
-        "per_sample_weights must have one entry per index; got %zd for %zd indices",
-        psw.numel(),
-        indices.numel());
-  }
-
-  return true;
-}
-
-// Runs before any weight read or output write, so every later loop can index
-// weight and outputs without bounds checks.
-bool check_offsets_and_indices(
-    const Tensor& weight,
-    const Tensor& indices,
-    const Tensor& offsets) {
-  const int64_t num_indices = indices.numel();
-  int64_t prev_offset = 0;
-  for (const auto i : c10::irange(offsets.numel())) {
-    const int64_t offset = index_at(offsets, i);
-    ET_CHECK_OR_RETURN_FALSE(
-        (i > 0 || offset == 0) && offset >= prev_offset &&
-            offset <= num_indices,
-        "offsets must start at 0 and be nondecreasing up to %" PRId64
-        "; offsets[%zu] = %" PRId64,
-        num_indices,
-        static_cast<size_t>(i),
-        offset);
-    prev_offset = offset;
-  }
-
-  // Validate every index, including those past a partial final offset.
-  const int64_t num_rows = weight.size(0);
-  for (const auto i : c10::irange(num_indices)) {
-    const int64_t row = index_at(indices, i);
-    ET_CHECK_OR_RETURN_FALSE(
-        row >= 0 && row < num_rows,
-        "index %" PRId64 " is outside weight rows [0, %" PRId64 ")",
-        row,
-        num_rows);
-  }
-  return true;
-}
-
-// Mirrors ATen's is_fast_path in EmbeddingBag.cpp.
-bool is_aten_fast_path(
-    const Tensor& weight,
-    const std::optional<Tensor>& per_sample_weights,
-    const Tensor& out,
-    int64_t padding_idx) {
-  const auto dtype = weight.scalar_type();
-  return (dtype == ScalarType::Float || dtype == ScalarType::Half ||
-          dtype == ScalarType::BFloat16) &&
-      weight.strides()[1] == 1 && out.strides()[1] == 1 && padding_idx < 0 &&
-      (!per_sample_weights.has_value() ||
-       per_sample_weights->strides()[0] == 1);
-}
-
 } // namespace
 
 std::tuple<Tensor&, Tensor&, Tensor&, Tensor&> _embedding_bag_out(
@@ -209,53 +79,32 @@ std::tuple<Tensor&, Tensor&, Tensor&, Tensor&> _embedding_bag_out(
       ret_val);
   ET_KERNEL_CHECK(
       ctx,
-      check_offsets_and_indices(weight, indices, offsets),
+      check_embedding_bag_offsets_and_indices(weight, indices, offsets),
       InvalidArgument,
       ret_val);
 
-  const int64_t num_indices = indices.numel();
-  const int64_t num_bags =
-      include_last_offset ? offsets.numel() - 1 : offsets.numel();
-  const int64_t dim = weight.size(1);
-
-  const SizesType out_sizes[2] = {
-      static_cast<SizesType>(num_bags), static_cast<SizesType>(dim)};
-  // out must be resized first: the fast-path check reads its strides.
+  EmbeddingBagLayout layout{};
   ET_KERNEL_CHECK(
       ctx,
-      resize_tensor(out, {out_sizes, 2}) == Error::Ok,
+      resize_embedding_bag_outputs(
+          weight,
+          indices,
+          offsets,
+          mode,
+          per_sample_weights,
+          include_last_offset,
+          padding_idx,
+          out,
+          offset2bag,
+          bag_size,
+          max_indices,
+          layout) == Error::Ok,
       InvalidArgument,
       ret_val);
 
-  // This kernel uses one reduction loop for both ATen paths, but follows
-  // ATen's path choice so the auxiliary outputs and the partial final bag
-  // match, including statically allocated export outputs.
-  const bool fast_path =
-      is_aten_fast_path(weight, per_sample_weights, out, padding_idx);
-  // ATen skips offset2bag only for sum on its fast path, but export sizes it
-  // with meta_embedding_bag, whose fast path differs (e.g. for BFloat16). Sum
-  // never reads it, so keep a caller-allocated num_indices layout and fill it.
-  const bool fills_offset2bag =
-      mode != kModeSum || !fast_path || offset2bag.numel() == num_indices;
-  // With include_last_offset, ATen's fast path stops the last sum/mean bag at
-  // the final offset; every other path gives it all remaining indices.
-  const bool truncate_last_bag =
-      include_last_offset && fast_path && mode != kModeMax;
-
-  // The auxiliary outputs exist for the backward pass, and the ATen CPU kernel
-  // only fills the ones each mode needs; match it so shapes agree.
-  // bag_size and the non-max max_indices are {num_bags}, a prefix of
-  // out_sizes.
-  const SizesType offset2bag_size =
-      static_cast<SizesType>(fills_offset2bag ? num_indices : 0);
-  ET_KERNEL_CHECK(
-      ctx,
-      resize_tensor(offset2bag, {&offset2bag_size, 1}) == Error::Ok &&
-          resize_tensor(bag_size, {out_sizes, 1}) == Error::Ok &&
-          resize_tensor(max_indices, {out_sizes, mode == kModeMax ? 2u : 1u}) ==
-              Error::Ok,
-      InvalidArgument,
-      ret_val);
+  const int64_t num_indices = indices.numel();
+  const int64_t num_bags = out.size(0);
+  const int64_t dim = weight.size(1);
 
   for (Tensor* t : {&bag_size, &max_indices}) {
     if (t->nbytes() > 0) {
@@ -267,13 +116,13 @@ std::tuple<Tensor&, Tensor&, Tensor&, Tensor&> _embedding_bag_out(
     return bag + 1 < num_bags ? index_at(offsets, bag + 1) : num_indices;
   };
   const auto reduction_end = [&](int64_t bag) {
-    return truncate_last_bag ? index_at(offsets, bag + 1) : bag_end(bag);
+    return layout.truncate_last_bag ? index_at(offsets, bag + 1) : bag_end(bag);
   };
 
   // The ATen CPU kernel leaves bag_size at zero for the sum mode, and mirrors
   // it into the 1-D max_indices for the mean mode, where that output carries
   // nothing else.
-  if (fills_offset2bag) {
+  if (layout.fills_offset2bag) {
     // With no bags, ATen puts every index before the first bag, in bag -1.
     if (num_bags == 0) {
       for (const auto i : c10::irange(num_indices)) {
@@ -289,10 +138,10 @@ std::tuple<Tensor&, Tensor&, Tensor&, Tensor&> _embedding_bag_out(
           ++count;
         }
       }
-      if (mode != kModeSum) {
+      if (mode != kEmbeddingBagModeSum) {
         set_index(bag_size, bag, count);
       }
-      if (mode == kModeMean) {
+      if (mode == kEmbeddingBagModeMean) {
         set_index(max_indices, bag, count);
       }
     }
@@ -325,7 +174,7 @@ std::tuple<Tensor&, Tensor&, Tensor&, Tensor&> _embedding_bag_out(
       // ATen divides by the full bag size even when the fast path truncates
       // the last bag.
       const int64_t mean_divisor =
-          mode == kModeMean ? index_at(bag_size, bag) : 0;
+          mode == kEmbeddingBagModeMean ? index_at(bag_size, bag) : 0;
       for (const auto d : c10::irange(dim)) {
         AccType acc = 0;
         bool seen = false;
@@ -336,7 +185,7 @@ std::tuple<Tensor&, Tensor&, Tensor&, Tensor&> _embedding_bag_out(
           }
           const AccType value = static_cast<AccType>(
               weight_data[row * weight_strides[0] + d * weight_strides[1]]);
-          if (mode == kModeMax) {
+          if (mode == kEmbeddingBagModeMax) {
             if (!seen || value > acc) {
               acc = value;
               set_index(
diff --git a/kernels/portable/cpu/util/embedding_bag_util.cpp b/kernels/portable/cpu/util/embedding_bag_util.cpp
new file mode 100644
index 00000000000..79908616152
--- /dev/null
+++ b/kernels/portable/cpu/util/embedding_bag_util.cpp
@@ -0,0 +1,214 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <executorch/kernels/portable/cpu/util/embedding_bag_util.h>
+
+#include <c10/util/irange.h>
+
+#include <cinttypes>
+
+namespace torch {
+namespace executor {
+
+namespace {
+
+using SizesType = executorch::aten::SizesType;
+
+bool is_index_dtype(ScalarType dtype) {
+  return dtype == ScalarType::Long || dtype == ScalarType::Int;
+}
+
+template <typename CTYPE>
+bool check_offsets(const CTYPE* offsets, int64_t count, int64_t num_indices) {
+  int64_t prev_offset = 0;
+  for (const auto i : c10::irange(count)) {
+    const int64_t offset = offsets[i];
+    ET_CHECK_OR_RETURN_FALSE(
+        (i > 0 || offset == 0) && offset >= prev_offset &&
+            offset <= num_indices,
+        "offsets must start at 0 and be nondecreasing up to %" PRId64
+        "; offsets[%zu] = %" PRId64,
+        num_indices,
+        static_cast<size_t>(i),
+        offset);
+    prev_offset = offset;
+  }
+  return true;
+}
+
+template <typename CTYPE>
+bool check_indices(const CTYPE* indices, int64_t count, int64_t num_rows) {
+  for (const auto i : c10::irange(count)) {
+    const int64_t row = indices[i];
+    ET_CHECK_OR_RETURN_FALSE(
+        row >= 0 && row < num_rows,
+        "index %" PRId64 " is outside weight rows [0, %" PRId64 ")",
+        row,
+        num_rows);
+  }
+  return true;
+}
+
+// Mirrors ATen's is_fast_path in EmbeddingBag.cpp.
+bool is_aten_fast_path(
+    const Tensor& weight,
+    const std::optional<Tensor>& per_sample_weights,
+    const Tensor& out,
+    int64_t padding_idx) {
+  const auto dtype = weight.scalar_type();
+  return (dtype == ScalarType::Float || dtype == ScalarType::Half ||
+          dtype == ScalarType::BFloat16) &&
+      weight.strides()[1] == 1 && out.strides()[1] == 1 && padding_idx < 0 &&
+      (!per_sample_weights.has_value() ||
+       per_sample_weights->strides()[0] == 1);
+}
+
+} // namespace
+
+bool check_embedding_bag_args(
+    const Tensor& weight,
+    const Tensor& indices,
+    const Tensor& offsets,
+    int64_t mode,
+    const std::optional<Tensor>& per_sample_weights,
+    bool include_last_offset,
+    const Tensor& out,
+    const Tensor& offset2bag,
+    const Tensor& bag_size,
+    const Tensor& max_indices) {
+  ET_CHECK_OR_RETURN_FALSE(
+      weight.dim() == 2,
+      "weight must be 2-D; weight.dim() = %" ET_PRI_TENSOR_DIM,
+      weight.dim());
+  ET_CHECK_OR_RETURN_FALSE(
+      indices.dim() == 1,
+      "indices must be 1-D; indices.dim() = %" ET_PRI_TENSOR_DIM,
+      indices.dim());
+  ET_CHECK_OR_RETURN_FALSE(
+      offsets.dim() == 1,
+      "offsets must be 1-D; offsets.dim() = %" ET_PRI_TENSOR_DIM,
+      offsets.dim());
+  ET_CHECK_OR_RETURN_FALSE(
+      is_index_dtype(indices.scalar_type()),
+      "indices must be Long or Int; indices.scalar_type() = %s",
+      toString(indices.scalar_type()));
+  ET_CHECK_OR_RETURN_FALSE(
+      is_index_dtype(offsets.scalar_type()),
+      "offsets must be Long or Int; offsets.scalar_type() = %s",
+      toString(offsets.scalar_type()));
+
+  ET_LOG_AND_RETURN_IF_FALSE(tensor_is_floating_type(weight));
+  ET_LOG_AND_RETURN_IF_FALSE(tensors_have_same_dtype(weight, out));
+
+  // Eager ATen gives these the promoted indices/offsets dtype, but export
+  // plans them with meta_embedding_bag, which uses the offsets dtype.
+  for (const Tensor* t : {&offset2bag, &bag_size, &max_indices}) {
+    ET_CHECK_OR_RETURN_FALSE(
+        is_index_dtype(t->scalar_type()),
+        "index outputs must be Long or Int; got %s",
+        toString(t->scalar_type()));
+  }
+
+  ET_CHECK_OR_RETURN_FALSE(
+      mode == kEmbeddingBagModeSum || mode == kEmbeddingBagModeMean ||
+          mode == kEmbeddingBagModeMax,
+      "mode must be 0 (sum), 1 (mean) or 2 (max); mode = %" PRId64,
+      mode);
+
+  if (include_last_offset) {
+    ET_CHECK_OR_RETURN_FALSE(
+        offsets.numel() >= 1,
+        "offsets must hold at least the final offset when include_last_offset is set");
+  }
+
+  if (per_sample_weights.has_value()) {
+    const Tensor& psw = per_sample_weights.value();
+    ET_CHECK_OR_RETURN_FALSE(psw.dim() == 1, "per_sample_weights must be 1-D");
+    ET_CHECK_OR_RETURN_FALSE(
+        mode == kEmbeddingBagModeSum,
+        "per_sample_weights is only supported for mode 0 (sum); mode = %" PRId64,
+        mode);
+    ET_LOG_AND_RETURN_IF_FALSE(tensors_have_same_dtype(weight, psw));
+    ET_CHECK_OR_RETURN_FALSE(
+        psw.numel() == indices.numel(),
+        "per_sample_weights must have one entry per index; got %zd for %zd indices",
+        psw.numel(),
+        indices.numel());
+  }
+
+  return true;
+}
+
+bool check_embedding_bag_offsets_and_indices(
+    const Tensor& weight,
+    const Tensor& indices,
+    const Tensor& offsets) {
+  const int64_t num_indices = indices.numel();
+  const bool offsets_ok = offsets.scalar_type() == ScalarType::Long
+      ? check_offsets(
+            offsets.const_data_ptr<int64_t>(), offsets.numel(), num_indices)
+      : check_offsets(
+            offsets.const_data_ptr<int32_t>(), offsets.numel(), num_indices);
+  if (!offsets_ok) {
+    return false;
+  }
+  // Validate every index, including those past a partial final offset.
+  const int64_t num_rows = weight.size(0);
+  return indices.scalar_type() == ScalarType::Long
+      ? check_indices(indices.const_data_ptr<int64_t>(), num_indices, num_rows)
+      : check_indices(indices.const_data_ptr<int32_t>(), num_indices, num_rows);
+}
+
+Error resize_embedding_bag_outputs(
+    const Tensor& weight,
+    const Tensor& indices,
+    const Tensor& offsets,
+    int64_t mode,
+    const std::optional<Tensor>& per_sample_weights,
+    bool include_last_offset,
+    int64_t padding_idx,
+    Tensor& out,
+    Tensor& offset2bag,
+    Tensor& bag_size,
+    Tensor& max_indices,
+    EmbeddingBagLayout& layout) {
+  const int64_t num_indices = indices.numel();
+  const SizesType out_sizes[2] = {
+      static_cast<SizesType>(
+          offsets.numel() - static_cast<int64_t>(include_last_offset)),
+      static_cast<SizesType>(weight.size(1))};
+  // out must be resized first: the fast-path check reads its strides.
+  ET_CHECK_OK_OR_RETURN_ERROR(resize_tensor(out, {out_sizes, 2}));
+
+  // Kernels may use one reduction loop for both ATen paths, but follow
+  // ATen's path choice so the auxiliary outputs and the partial final bag
+  // match, including statically allocated export outputs.
+  const bool fast_path =
+      is_aten_fast_path(weight, per_sample_weights, out, padding_idx);
+  // ATen skips offset2bag only for sum on its fast path, but export sizes it
+  // with meta_embedding_bag, whose fast path differs (e.g. for BFloat16). Sum
+  // never reads it, so keep a caller-allocated num_indices layout and fill it.
+  layout.fills_offset2bag = mode != kEmbeddingBagModeSum || !fast_path ||
+      offset2bag.numel() == num_indices;
+  layout.truncate_last_bag =
+      include_last_offset && fast_path && mode != kEmbeddingBagModeMax;
+
+  // The auxiliary outputs exist for the backward pass, and the ATen CPU kernel
+  // only fills the ones each mode needs; match it so shapes agree.
+  // bag_size and the non-max max_indices are {num_bags}, a prefix of
+  // out_sizes.
+  const SizesType offset2bag_size =
+      static_cast<SizesType>(layout.fills_offset2bag ? num_indices : 0);
+  ET_CHECK_OK_OR_RETURN_ERROR(resize_tensor(offset2bag, {&offset2bag_size, 1}));
+  ET_CHECK_OK_OR_RETURN_ERROR(resize_tensor(bag_size, {out_sizes, 1}));
+  return resize_tensor(
+      max_indices, {out_sizes, mode == kEmbeddingBagModeMax ? 2u : 1u});
+}
+
+} // namespace executor
+} // namespace torch
diff --git a/kernels/portable/cpu/util/embedding_bag_util.h b/kernels/portable/cpu/util/embedding_bag_util.h
new file mode 100644
index 00000000000..e82c77d4719
--- /dev/null
+++ b/kernels/portable/cpu/util/embedding_bag_util.h
@@ -0,0 +1,70 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#pragma once
+
+#include <executorch/runtime/kernel/kernel_includes.h>
+
+#include <optional>
+
+namespace torch {
+namespace executor {
+
+// aten::_embedding_bag's `mode` argument.
+constexpr int64_t kEmbeddingBagModeSum = 0;
+constexpr int64_t kEmbeddingBagModeMean = 1;
+constexpr int64_t kEmbeddingBagModeMax = 2;
+
+// Validates everything except the offsets and indices values.
+bool check_embedding_bag_args(
+    const Tensor& weight,
+    const Tensor& indices,
+    const Tensor& offsets,
+    int64_t mode,
+    const std::optional<Tensor>& per_sample_weights,
+    bool include_last_offset,
+    const Tensor& out,
+    const Tensor& offset2bag,
+    const Tensor& bag_size,
+    const Tensor& max_indices);
+
+// Requires check_embedding_bag_args. Once this passes, kernels can index
+// weight and every output without bounds checks.
+bool check_embedding_bag_offsets_and_indices(
+    const Tensor& weight,
+    const Tensor& indices,
+    const Tensor& offsets);
+
+// The parts of ATen's CPU path choice that change the outputs.
+struct EmbeddingBagLayout {
+  // offset2bag holds one bag id per index; otherwise it is empty.
+  bool fills_offset2bag;
+  // With include_last_offset, the last sum/mean bag stops at the final
+  // offset, while offset2bag and the mean divisor still count the indices
+  // after it. Otherwise the last bag takes every remaining index.
+  bool truncate_last_bag;
+};
+
+// Resizes the four outputs to the shapes ATen gives them and reports the
+// layout. Run both checks first, so a failing call leaves outputs untouched.
+Error resize_embedding_bag_outputs(
+    const Tensor& weight,
+    const Tensor& indices,
+    const Tensor& offsets,
+    int64_t mode,
+    const std::optional<Tensor>& per_sample_weights,
+    bool include_last_offset,
+    int64_t padding_idx,
+    Tensor& out,
+    Tensor& offset2bag,
+    Tensor& bag_size,
+    Tensor& max_indices,
+    EmbeddingBagLayout& layout);
+
+} // namespace executor
+} // namespace torch
diff --git a/kernels/portable/cpu/util/targets.bzl b/kernels/portable/cpu/util/targets.bzl
index 30315077713..b3b61e6c123 100644
--- a/kernels/portable/cpu/util/targets.bzl
+++ b/kernels/portable/cpu/util/targets.bzl
@@ -37,6 +37,7 @@ def define_common_targets():
             "//executorch/kernels/portable/cpu/util:upsample_util",
             "//executorch/kernels/portable/cpu/util:vectorized_math",
             "//executorch/kernels/portable/cpu/util:grid_sampler_2d_util",
+            "//executorch/kernels/portable/cpu/util:embedding_bag_util",
         ],
         visibility = ["PUBLIC"],
     )
@@ -415,6 +416,20 @@ def define_common_targets():
         visibility = ["//executorch/kernels/portable/cpu/...", "//executorch/kernels/optimized/cpu/..."],
     )
 
+    runtime.cxx_library(
+        name = "embedding_bag_util",
+        srcs = ["embedding_bag_util.cpp"],
+        exported_headers = ["embedding_bag_util.h"],
+        deps = [
+            "//executorch/runtime/core/exec_aten/util:scalar_type_util",
+            "//executorch/runtime/core/exec_aten/util:tensor_util",
+        ],
+        exported_deps = [
+            "//executorch/runtime/kernel:kernel_includes",
+        ],
+        visibility = ["//executorch/kernels/portable/cpu/...", "//executorch/kernels/optimized/cpu/..."],
+    )
+
     # Utility functions that can be used by operators that perform reduction
     for aten_mode in get_aten_mode_options():
         suffix = "_aten" if aten_mode else ""
diff --git a/kernels/test/CMakeLists.txt b/kernels/test/CMakeLists.txt
index 1b00ebc9fcc..a0aa2ed7871 100644
--- a/kernels/test/CMakeLists.txt
+++ b/kernels/test/CMakeLists.txt
@@ -356,6 +356,7 @@ if(EXECUTORCH_BUILD_KERNELS_OPTIMIZED)
       "op_bmm_test.cpp"
       "op_div_test.cpp"
       "op_elu_test.cpp"
+      "op_embedding_bag_test.cpp"
       "op_exp_test.cpp"
       "op_fft_c2r_test.cpp"
       "op_fft_r2c_test.cpp"
diff --git a/kernels/test/op_embedding_bag_test.cpp b/kernels/test/op_embedding_bag_test.cpp
index b5eeb0c2c38..2cdd11dd3bc 100644
--- a/kernels/test/op_embedding_bag_test.cpp
+++ b/kernels/test/op_embedding_bag_test.cpp
@@ -16,6 +16,15 @@
 
 #include <gtest/gtest.h>
 
+#ifdef ET_USE_THREADPOOL
+#include <executorch/extension/threadpool/threadpool.h>
+#endif
+
+#include <algorithm>
+#include <array>
+#include <cmath>
+#include <cstdint>
+#include <limits>
 #include <numeric>
 #include <vector>
 
@@ -25,6 +34,40 @@ using executorch::aten::Tensor;
 using std::optional;
 using torch::executor::testing::TensorFactory;
 
+namespace {
+
+// Widens the threadpool so parallel_for splits work into several tasks even on
+// single-core runners, and restores the previous size afterwards.
+class ScopedMinThreadCount {
+ public:
+  explicit ScopedMinThreadCount(uint32_t threads) {
+#ifdef ET_USE_THREADPOOL
+    auto* pool = ::executorch::extension::threadpool::get_threadpool();
+    if (pool != nullptr && pool->get_thread_count() < threads) {
+      previous_ = static_cast<uint32_t>(pool->get_thread_count());
+      pool->_unsafe_reset_threadpool(threads);
+    }
+#else
+    (void)threads;
+#endif
+  }
+  ~ScopedMinThreadCount() {
+#ifdef ET_USE_THREADPOOL
+    if (previous_ != 0) {
+      ::executorch::extension::threadpool::get_threadpool()
+          ->_unsafe_reset_threadpool(previous_);
+    }
+#endif
+  }
+  ScopedMinThreadCount(const ScopedMinThreadCount&) = delete;
+  ScopedMinThreadCount& operator=(const ScopedMinThreadCount&) = delete;
+
+ private:
+  uint32_t previous_ = 0;
+};
+
+} // namespace
+
 class OpEmbeddingBagOutTest : public OperatorTest {
  protected:
   std::tuple<Tensor&, Tensor&, Tensor&, Tensor&> op_embedding_bag_out(
@@ -134,28 +177,36 @@ class OpEmbeddingBagOutTest : public OperatorTest {
     }
   }
 
-  void expect_invalid(
+  // Index outputs use the offsets dtype, as export allocates them, so the
+  // optimized kernel takes its fast path whenever the weight allows it.
+  template <ScalarType Index, ScalarType Offset>
+  void expect_invalid_with_dtypes(
       const std::vector<int64_t>& index_values,
       const std::vector<int64_t>& offset_values,
-      int64_t mode = 0,
-      int64_t dim = 1,
-      bool include_last = false,
-      const optional<Tensor>& psw = std::nullopt) {
+      int64_t mode,
+      int64_t dim,
+      bool include_last,
+      const optional<Tensor>& psw) {
+    using IndexType = typename TensorFactory<Index>::ctype;
+    using OffsetType = typename TensorFactory<Offset>::ctype;
     TensorFactory<ScalarType::Float> tf;
-    TensorFactory<ScalarType::Long> ti;
+    TensorFactory<Index> ti;
+    TensorFactory<Offset> to;
     auto weight = tf.ones({2, static_cast<int32_t>(dim)});
-    auto indices =
-        ti.make({static_cast<int32_t>(index_values.size())}, index_values);
-    auto offsets =
-        ti.make({static_cast<int32_t>(offset_values.size())}, offset_values);
+    auto indices = ti.make(
+        {static_cast<int32_t>(index_values.size())},
+        std::vector<IndexType>(index_values.begin(), index_values.end()));
+    auto offsets = to.make(
+        {static_cast<int32_t>(offset_values.size())},
+        std::vector<OffsetType>(offset_values.begin(), offset_values.end()));
     int32_t bags =
         static_cast<int32_t>(offset_values.size()) - (include_last ? 1 : 0);
     auto out = tf.zeros({bags, static_cast<int32_t>(dim)});
     auto offset2bag =
-        ti.zeros({mode == 0 ? 0 : static_cast<int32_t>(index_values.size())});
-    auto bag_size = ti.zeros({bags});
-    auto max_indices = mode == 2 ? ti.zeros({bags, static_cast<int32_t>(dim)})
-                                 : ti.zeros({bags});
+        to.zeros({mode == 0 ? 0 : static_cast<int32_t>(index_values.size())});
+    auto bag_size = to.zeros({bags});
+    auto max_indices = mode == 2 ? to.zeros({bags, static_cast<int32_t>(dim)})
+                                 : to.zeros({bags});
     ET_EXPECT_KERNEL_FAILURE(
         context_,
         op_embedding_bag_out(
@@ -172,6 +223,39 @@ class OpEmbeddingBagOutTest : public OperatorTest {
             max_indices));
   }
 
+  void expect_invalid(
+      const std::vector<int64_t>& index_values,
+      const std::vector<int64_t>& offset_values,
+      int64_t mode = 0,
+      int64_t dim = 1,
+      bool include_last = false,
+      const optional<Tensor>& psw = std::nullopt) {
+    {
+      SCOPED_TRACE("Long indices, Long offsets");
+      expect_invalid_with_dtypes<ScalarType::Long, ScalarType::Long>(
+          index_values, offset_values, mode, dim, include_last, psw);
+    }
+    context_.fail(torch::executor::Error::Ok);
+    {
+      SCOPED_TRACE("Int indices, Int offsets");
+      expect_invalid_with_dtypes<ScalarType::Int, ScalarType::Int>(
+          index_values, offset_values, mode, dim, include_last, psw);
+    }
+    context_.fail(torch::executor::Error::Ok);
+    {
+      SCOPED_TRACE("Int indices, Long offsets");
+      expect_invalid_with_dtypes<ScalarType::Int, ScalarType::Long>(
+          index_values, offset_values, mode, dim, include_last, psw);
+    }
+    // ATen requires the promoted Long dtype for the index outputs.
+    if (!torch::executor::testing::SupportedFeatures::get()->is_aten) {
+      context_.fail(torch::executor::Error::Ok);
+      SCOPED_TRACE("Long indices, Int offsets");
+      expect_invalid_with_dtypes<ScalarType::Long, ScalarType::Int>(
+          index_values, offset_values, mode, dim, include_last, psw);
+    }
+  }
+
   template <ScalarType DTYPE>
   void test_partial_final_offset() {
     TensorFactory<DTYPE> tf;
@@ -254,6 +338,125 @@ class OpEmbeddingBagOutTest : public OperatorTest {
     EXPECT_TENSOR_EQ(max_indices, ti.zeros({2}));
   }
 
+  template <ScalarType DTYPE>
+  void test_reductions(
+      int32_t width,
+      int64_t mode,
+      bool last,
+      bool padding,
+      bool weighted = false) {
+    TensorFactory<ScalarType::Float> tf;
+    TensorFactory<DTYPE> ti;
+    const std::array<std::array<float, 3>, 4> rows = {
+        {{1, -2, 3}, {4, 5, -6}, {7, 8, 9}, {10, -11, 12}}};
+    // Exact in float, so fused and unfused multiply-adds agree.
+    const std::array<float, 4> scales = {0.5, -1, 2, 3};
+    std::vector<float> weights;
+    for (const auto& row : rows) {
+      for (int32_t d = 0; d < width; ++d) {
+        weights.push_back(row[d % 3]);
+      }
+    }
+    auto weight = tf.make({4, width}, weights);
+    auto indices = ti.make({4}, {1, 2, 1, 3});
+    auto offsets = last ? ti.make({4}, {0, 2, 2, 4}) : ti.make({3}, {0, 2, 2});
+    auto out = tf.zeros({3, width});
+    auto offset2bag = ti.zeros({mode == 0 && !padding ? 0 : 4});
+    auto bag_size = ti.zeros({3});
+    auto max_indices = mode == 2 ? ti.zeros({3, width}) : ti.zeros({3});
+    optional<Tensor> psw;
+    if (weighted) {
+      psw = tf.make({4}, std::vector<float>(scales.begin(), scales.end()));
+    }
+    op_embedding_bag_out(
+        weight,
+        indices,
+        offsets,
+        mode,
+        psw,
+        last,
+        padding ? 1 : -1,
+        out,
+        offset2bag,
+        bag_size,
+        max_indices);
+    ASSERT_EQ(context_.failure_state(), torch::executor::Error::Ok);
+    std::vector<float> expected;
+    for (int32_t bag = 0; bag < 3; ++bag) {
+      // Bags 0 and 2 hold positions {0, 1} and {2, 3}; row 1 comes first.
+      const size_t first = bag == 0 ? 0 : 2;
+      const float scale_a = weighted ? scales[first] : 1;
+      const float scale_b = weighted ? scales[first + 1] : 1;
+      for (int32_t d = 0; d < width; ++d) {
+        const float a = rows[1][d % 3];
+        const float b = rows[bag == 0 ? 2 : 3][d % 3];
+        float value = 0;
+        if (bag == 1) {
+          value = 0;
+        } else if (padding) {
+          value = b * scale_b;
+        } else if (mode == 2) {
+          value = std::max(a, b);
+        } else {
+          value = (a * scale_a + b * scale_b) / (mode == 1 ? 2 : 1);
+        }
+        expected.push_back(value);
+      }
+    }
+    EXPECT_TENSOR_EQ(out, tf.make({3, width}, expected));
+    if (offset2bag.numel() != 0) {
+      EXPECT_TENSOR_EQ(offset2bag, ti.make({4}, {0, 0, 2, 2}));
+    }
+    EXPECT_TENSOR_EQ(
+        bag_size,
+        mode == 0     ? ti.zeros({3})
+            : padding ? ti.make({3}, {1, 0, 1})
+                      : ti.make({3}, {2, 0, 2}));
+    if (mode != 2) {
+      EXPECT_TENSOR_EQ(max_indices, bag_size);
+    } else {
+      const auto* actual =
+          max_indices.template const_data_ptr<typename decltype(ti)::ctype>();
+      for (int32_t d = 0; d < width; ++d) {
+        EXPECT_EQ(actual[d], 2);
+        EXPECT_EQ(actual[width + d], 0);
+        EXPECT_EQ(actual[2 * width + d], padding || d % 3 != 1 ? 3 : 1);
+      }
+    }
+  }
+
+  template <ScalarType Index, ScalarType Offset>
+  void test_weighted_sum_mixed_index_types() {
+    TensorFactory<ScalarType::Float> tf;
+    TensorFactory<Index> ti;
+    TensorFactory<Offset> to;
+    auto weight = tf.make({3, 2}, {1, 2, 3, 4, 5, 6});
+    auto indices = ti.make({4}, {0, 1, 1, 2});
+    auto offsets = to.make({3}, {0, 2, 4});
+    auto scales = tf.make({4}, {0.5, -1, 2, 3});
+    auto out = tf.zeros({2, 2});
+    auto offset2bag = to.zeros({4});
+    auto bag_size = to.zeros({2});
+    auto max_indices = to.zeros({2});
+    op_embedding_bag_out(
+        weight,
+        indices,
+        offsets,
+        0,
+        scales,
+        true,
+        1,
+        out,
+        offset2bag,
+        bag_size,
+        max_indices);
+    ASSERT_EQ(context_.failure_state(), torch::executor::Error::Ok);
+    EXPECT_TENSOR_EQ(out, tf.make({2, 2}, {0.5, 1, 15, 18}));
+    EXPECT_TENSOR_EQ(offset2bag, to.make({4}, {0, 0, 1, 1}));
+    EXPECT_TENSOR_EQ(bag_size, to.zeros({2}));
+    EXPECT_TENSOR_EQ(max_indices, to.zeros({2}));
+  }
+
   // weight[r][c] = r * 3 + c, so a row is identifiable from its values.
   static Tensor ramp_weight(TensorFactory<ScalarType::Float>& tf) {
     std::vector<float> values(15);
@@ -637,6 +840,17 @@ TEST_F(OpEmbeddingBagOutTest, WrongPerSampleWeightsRankRejected) {
   TensorFactory<ScalarType::Float> tf;
   expect_invalid({0, 1}, {0}, 0, 1, false, tf.ones({1, 2}));
 }
+TEST_F(OpEmbeddingBagOutTest, WrongIndicesDtypeRejected) {
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen promotes Short indices to the offsets dtype");
+  expect_invalid_with_dtypes<ScalarType::Short, ScalarType::Long>(
+      {0}, {0}, 0, 1, false, std::nullopt);
+}
+TEST_F(OpEmbeddingBagOutTest, WrongOffsetsDtypeRejected) {
+  expect_invalid_with_dtypes<ScalarType::Long, ScalarType::Short>(
+      {0}, {0}, 0, 1, false, std::nullopt);
+}
 
 TEST_F(OpEmbeddingBagOutTest, WrongAuxiliaryDtypeRejected) {
   TensorFactory<ScalarType::Float> tf;
@@ -946,3 +1160,422 @@ TEST_F(OpEmbeddingBagOutTest, OffsetsDtypeAuxiliaryOutputs) {
     }
   }
 }
+
+TEST_F(OpEmbeddingBagOutTest, ReductionsAcrossVectorWidths) {
+  for (int32_t width : {0, 1, 3, 4, 7, 8, 9, 16, 17, 65}) {
+    // ATen's FBGEMM path raises a spurious index error for zero-width weights.
+    if (width == 0 &&
+        torch::executor::testing::SupportedFeatures::get()->is_aten) {
+      continue;
+    }
+    for (int64_t mode = 0; mode < 3; ++mode) {
+      for (bool last : {false, true}) {
+        for (bool padding : {false, true}) {
+          SCOPED_TRACE(
+              ::testing::Message()
+              << width << "," << mode << "," << last << "," << padding);
+          test_reductions<ScalarType::Int>(width, mode, last, padding);
+          test_reductions<ScalarType::Long>(width, mode, last, padding);
+          if (mode == 0) {
+            SCOPED_TRACE("weighted");
+            test_reductions<ScalarType::Int>(width, mode, last, padding, true);
+            test_reductions<ScalarType::Long>(width, mode, last, padding, true);
+          }
+        }
+      }
+    }
+  }
+}
+
+TEST_F(OpEmbeddingBagOutTest, WeightedSumMixedIndexTypes) {
+  test_weighted_sum_mixed_index_types<ScalarType::Int, ScalarType::Long>();
+  // ATen requires the promoted Long dtype for these outputs.
+  if (!torch::executor::testing::SupportedFeatures::get()->is_aten) {
+    test_weighted_sum_mixed_index_types<ScalarType::Long, ScalarType::Int>();
+  }
+}
+
+TEST_F(OpEmbeddingBagOutTest, PaddingOnlyAndEmptyBags) {
+  TensorFactory<ScalarType::Float> tf;
+  TensorFactory<ScalarType::Long> ti;
+  auto weight = tf.ones({2, 9});
+  for (int64_t mode = 0; mode < 3; ++mode) {
+    for (bool empty : {false, true}) {
+      SCOPED_TRACE(::testing::Message() << mode << "," << empty);
+      auto indices = empty ? ti.zeros({0}) : ti.make({2}, {1, 1});
+      auto offsets = ti.make({2}, {0, 0});
+      auto out = tf.ones({2, 9});
+      auto offset2bag = ti.zeros({empty ? 0 : 2});
+      auto bag_size = ti.ones({2});
+      auto max_indices = mode == 2 ? ti.ones({2, 9}) : ti.ones({2});
+      op_embedding_bag_out(
+          weight,
+          indices,
+          offsets,
+          mode,
+          {},
+          false,
+          1,
+          out,
+          offset2bag,
+          bag_size,
+          max_indices);
+      ASSERT_EQ(context_.failure_state(), torch::executor::Error::Ok);
+      EXPECT_TENSOR_EQ(out, tf.zeros({2, 9}));
+      EXPECT_TENSOR_EQ(bag_size, ti.zeros({2}));
+      EXPECT_TENSOR_EQ(
+          max_indices, mode == 2 ? ti.zeros({2, 9}) : ti.zeros({2}));
+    }
+  }
+}
+
+TEST_F(OpEmbeddingBagOutTest, MaxPreservesFirstTieAndNan) {
+  TensorFactory<ScalarType::Float> tf;
+  TensorFactory<ScalarType::Long> ti;
+  const float nan = std::numeric_limits<float>::quiet_NaN();
+  std::vector<float> values(18, 2);
+  values[0] = nan;
+  values[10] = nan;
+  auto weight = tf.make({2, 9}, values);
+  auto indices = ti.make({2}, {0, 1});
+  auto offsets = ti.make({1}, {0});
+  auto out = tf.zeros({1, 9});
+  auto offset2bag = ti.zeros({2});
+  auto bag_size = ti.zeros({1});
+  auto max_indices = ti.ones({1, 9});
+  op_embedding_bag_out(
+      weight,
+      indices,
+      offsets,
+      2,
+      {},
+      false,
+      -1,
+      out,
+      offset2bag,
+      bag_size,
+      max_indices);
+  ASSERT_EQ(context_.failure_state(), torch::executor::Error::Ok);
+  EXPECT_TRUE(std::isnan(out.const_data_ptr<float>()[0]));
+  for (int d = 1; d < 9; ++d) {
+    EXPECT_EQ(out.const_data_ptr<float>()[d], 2);
+  }
+  EXPECT_TENSOR_EQ(max_indices, ti.zeros({1, 9}));
+}
+
+TEST_F(OpEmbeddingBagOutTest, DynamicBoundOutputsResize) {
+  TensorFactory<ScalarType::Float> tf;
+  TensorFactory<ScalarType::Long> ti;
+  const auto bounded = torch::executor::TensorShapeDynamism::DYNAMIC_BOUND;
+  auto weight = tf.make({2, 2}, {1, 4, 3, 2});
+  auto indices = ti.make({2}, {0, 1});
+  auto offsets = ti.make({1}, {0});
+  auto out = tf.zeros({3, 3}, bounded);
+  auto offset2bag = ti.zeros({4}, bounded);
+  auto bag_size = ti.zeros({3}, bounded);
+  auto max_indices = ti.zeros({3, 3}, bounded);
+  op_embedding_bag_out(
+      weight,
+      indices,
+      offsets,
+      2,
+      {},
+      false,
+      -1,
+      out,
+      offset2bag,
+      bag_size,
+      max_indices);
+  ASSERT_EQ(context_.failure_state(), torch::executor::Error::Ok);
+  EXPECT_TENSOR_EQ(out, tf.make({1, 2}, {3, 4}));
+  EXPECT_TENSOR_EQ(offset2bag, ti.make({2}, {0, 0}));
+  EXPECT_TENSOR_EQ(bag_size, ti.make({1}, {2}));
+  EXPECT_TENSOR_EQ(max_indices, ti.make({1, 2}, {1, 0}));
+}
+
+TEST_F(OpEmbeddingBagOutTest, TransposedMaxIndices) {
+  TensorFactory<ScalarType::Float> tf;
+  TensorFactory<ScalarType::Long> ti;
+  auto weight = tf.make({3, 2}, {1, 4, 3, 2, 0, 5});
+  auto indices = ti.make({4}, {0, 1, 1, 2});
+  auto offsets = ti.make({2}, {0, 2});
+  auto out = tf.zeros({2, 2});
+  auto offset2bag = ti.zeros({4});
+  auto bag_size = ti.zeros({2});
+  // Strides (1, 2), like a transposed contiguous tensor.
+  auto max_indices = ti.make_with_dimorder({2, 2}, {0, 0, 0, 0}, {1, 0});
+  op_embedding_bag_out(
+      weight,
+      indices,
+      offsets,
+      2,
+      {},
+      false,
+      -1,
+      out,
+      offset2bag,
+      bag_size,
+      max_indices);
+  ASSERT_EQ(context_.failure_state(), torch::executor::Error::Ok);
+  EXPECT_TENSOR_EQ(out, tf.make({2, 2}, {3, 4, 3, 5}));
+  // Equivalent to tensor([[1, 0], [1, 2]]).
+  EXPECT_TENSOR_EQ(
+      max_indices, ti.make_with_dimorder({2, 2}, {1, 1, 0, 2}, {1, 0}));
+}
+
+TEST_F(OpEmbeddingBagOutTest, ManyBagsAllModes) {
+  ScopedMinThreadCount threads(4);
+  TensorFactory<ScalarType::Float> tf;
+  TensorFactory<ScalarType::Long> ti;
+  constexpr int32_t kRows = 7;
+  constexpr int32_t kWidth = 65;
+  // Enough work to split the bags across parallel tasks.
+  constexpr int32_t kBags = 512;
+  std::vector<float> weights;
+  for (int32_t row = 0; row < kRows; ++row) {
+    for (int32_t d = 0; d < kWidth; ++d) {
+      weights.push_back(static_cast<float>(row * 100 + d - 50));
+    }
+  }
+  auto weight = tf.make({kRows, kWidth}, weights);
+  // Bag b repeats row b % kRows (b % 4 + 1) times.
+  std::vector<int64_t> indices_data;
+  std::vector<int64_t> offsets_data;
+  std::vector<int64_t> offset2bag_data;
+  std::vector<int64_t> counts;
+  for (int32_t bag = 0; bag < kBags; ++bag) {
+    offsets_data.push_back(static_cast<int64_t>(indices_data.size()));
+    counts.push_back(bag % 4 + 1);
+    for (int32_t i = 0; i < counts.back(); ++i) {
+      indices_data.push_back(bag % kRows);
+      offset2bag_data.push_back(bag);
+    }
+  }
+  const auto num_indices = static_cast<int32_t>(indices_data.size());
+  auto indices = ti.make({num_indices}, indices_data);
+  auto offsets = ti.make({kBags}, offsets_data);
+  for (int64_t mode = 0; mode < 3; ++mode) {
+    SCOPED_TRACE(mode);
+    auto out = tf.zeros({kBags, kWidth});
+    auto offset2bag = ti.zeros({mode == 0 ? 0 : num_indices});
+    auto bag_size = ti.zeros({kBags});
+    auto max_indices =
+        mode == 2 ? ti.zeros({kBags, kWidth}) : ti.zeros({kBags});
+    op_embedding_bag_out(
+        weight,
+        indices,
+        offsets,
+        mode,
+        {},
+        false,
+        -1,
+        out,
+        offset2bag,
+        bag_size,
+        max_indices);
+    ASSERT_EQ(context_.failure_state(), torch::executor::Error::Ok);
+    std::vector<float> expected_out;
+    std::vector<int64_t> expected_winners;
+    for (int32_t bag = 0; bag < kBags; ++bag) {
+      for (int32_t d = 0; d < kWidth; ++d) {
+        const float value = weights[(bag % kRows) * kWidth + d];
+        expected_out.push_back(
+            mode == 0 ? value * static_cast<float>(counts[bag]) : value);
+        expected_winners.push_back(bag % kRows);
+      }
+    }
+    EXPECT_TENSOR_EQ(out, tf.make({kBags, kWidth}, expected_out));
+    if (mode != 0) {
+      EXPECT_TENSOR_EQ(offset2bag, ti.make({num_indices}, offset2bag_data));
+    }
+    EXPECT_TENSOR_EQ(
+        bag_size, mode == 0 ? ti.zeros({kBags}) : ti.make({kBags}, counts));
+    EXPECT_TENSOR_EQ(
+        max_indices,
+        mode == 0       ? ti.zeros({kBags})
+            : mode == 1 ? ti.make({kBags}, counts)
+                        : ti.make({kBags, kWidth}, expected_winners));
+  }
+}
+
+TEST_F(OpEmbeddingBagOutTest, ManyBagsPartialFinalOffset) {
+  // Four threads split the bags below into four tasks, so the last bag and
+  // its partial final offset are handled away from the first task.
+  ScopedMinThreadCount threads(4);
+  TensorFactory<ScalarType::Float> tf;
+  TensorFactory<ScalarType::Long> ti;
+  constexpr int32_t kRows = 7;
+  constexpr int32_t kWidth = 64;
+  constexpr int32_t kBags = 1024;
+  // Indices past the final offset.
+  constexpr int32_t kTail = 3;
+  constexpr int64_t kPaddingRow = 3;
+  std::vector<float> weights;
+  for (int32_t row = 0; row < kRows; ++row) {
+    for (int32_t d = 0; d < kWidth; ++d) {
+      weights.push_back(static_cast<float>(row * 100 + d - 50));
+    }
+  }
+  auto weight = tf.make({kRows, kWidth}, weights);
+  // Bag b holds rows (b + j) % kRows for j < b % 4 + 1.
+  std::vector<int64_t> indices_data;
+  std::vector<int64_t> offsets_data;
+  for (int32_t bag = 0; bag < kBags; ++bag) {
+    offsets_data.push_back(static_cast<int64_t>(indices_data.size()));
+    for (int32_t j = 0; j <= bag % 4; ++j) {
+      indices_data.push_back((bag + j) % kRows);
+    }
+  }
+  offsets_data.push_back(static_cast<int64_t>(indices_data.size()));
+  for (int32_t j = 0; j < kTail; ++j) {
+    indices_data.push_back((j * 2) % kRows);
+  }
+  const auto num_indices = static_cast<int32_t>(indices_data.size());
+  // Multiples of 0.5, so every weighted sum is exact.
+  std::vector<float> scales_data;
+  for (int32_t i = 0; i < num_indices; ++i) {
+    scales_data.push_back(static_cast<float>(i % 5 - 2) * 0.5f);
+  }
+  auto indices = ti.make({num_indices}, indices_data);
+  auto offsets = ti.make({kBags + 1}, offsets_data);
+  auto scales = tf.make({num_indices}, scales_data);
+
+  for (int64_t mode = 0; mode < 3; ++mode) {
+    for (bool padded : {false, true}) {
+      for (bool weighted : {false, true}) {
+        if (weighted && mode != 0) {
+          continue;
+        }
+        SCOPED_TRACE(
+            ::testing::Message() << mode << "," << padded << "," << weighted);
+        const int64_t padding_idx = padded ? kPaddingRow : -1;
+        // ATen's FP32 fast path (no padding) stops the last sum/mean bag at
+        // the final offset; otherwise it takes every remaining index.
+        const bool truncate = !padded && mode != 2;
+        std::vector<float> expected_out(kBags * kWidth, 0);
+        std::vector<int64_t> expected_offset2bag(num_indices);
+        std::vector<int64_t> expected_counts(kBags);
+        std::vector<int64_t> expected_winners(kBags * kWidth, 0);
+        for (int32_t bag = 0; bag < kBags; ++bag) {
+          const int64_t begin = offsets_data[bag];
+          const int64_t end =
+              bag + 1 < kBags ? offsets_data[bag + 1] : num_indices;
+          const int64_t reduce_end = truncate ? offsets_data[bag + 1] : end;
+          int64_t count = 0;
+          for (int64_t i = begin; i < end; ++i) {
+            expected_offset2bag[i] = bag;
+            const int64_t row = indices_data[i];
+            if (row == padding_idx) {
+              continue;
+            }
+            ++count;
+            if (i >= reduce_end) {
+              continue;
+            }
+            for (int32_t d = 0; d < kWidth; ++d) {
+              const float value = weights[row * kWidth + d];
+              float& acc = expected_out[bag * kWidth + d];
+              if (mode == 2) {
+                if (count == 1 || value > acc) {
+                  acc = value;
+                  expected_winners[bag * kWidth + d] = row;
+                }
+              } else {
+                acc += value * (weighted ? scales_data[i] : 1.0f);
+              }
+            }
+          }
+          expected_counts[bag] = count;
+          if (mode == 1 && count != 0) {
+            for (int32_t d = 0; d < kWidth; ++d) {
+              expected_out[bag * kWidth + d] /= static_cast<float>(count);
+            }
+          }
+        }
+
+        auto out = tf.zeros({kBags, kWidth});
+        auto offset2bag =
+            ti.full({mode == 0 && !padded ? 0 : num_indices}, -99);
+        auto bag_size = ti.full({kBags}, -99);
+        auto max_indices =
+            mode == 2 ? ti.full({kBags, kWidth}, -99) : ti.full({kBags}, -99);
+        op_embedding_bag_out(
+            weight,
+            indices,
+            offsets,
+            mode,
+            weighted ? optional<Tensor>(scales) : std::nullopt,
+            true,
+            padding_idx,
+            out,
+            offset2bag,
+            bag_size,
+            max_indices);
+        ASSERT_EQ(context_.failure_state(), torch::executor::Error::Ok);
+        EXPECT_TENSOR_CLOSE(out, tf.make({kBags, kWidth}, expected_out));
+        if (offset2bag.numel() != 0) {
+          EXPECT_TENSOR_EQ(
+              offset2bag, ti.make({num_indices}, expected_offset2bag));
+        }
+        EXPECT_TENSOR_EQ(
+            bag_size,
+            mode == 0 ? ti.zeros({kBags}) : ti.make({kBags}, expected_counts));
+        EXPECT_TENSOR_EQ(
+            max_indices,
+            mode == 0       ? ti.zeros({kBags})
+                : mode == 1 ? ti.make({kBags}, expected_counts)
+                            : ti.make({kBags, kWidth}, expected_winners));
+      }
+    }
+  }
+}
+
+TEST_F(OpEmbeddingBagOutTest, RejectedInputsLeaveDynamicOutputsUnresized) {
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen tensors have no dynamic-bound shape");
+  TensorFactory<ScalarType::Float> tf;
+  TensorFactory<ScalarType::Long> ti;
+  TensorFactory<ScalarType::Short> ts;
+  const auto bounded = torch::executor::TensorShapeDynamism::DYNAMIC_BOUND;
+  auto weight = tf.ones({2, 2});
+  struct Case {
+    const char* name;
+    Tensor indices;
+    Tensor offsets;
+  };
+  const std::vector<Case> cases = {
+      {"index past the last row", ti.make({1}, {5}), ti.make({1}, {0})},
+      {"offset past the last index",
+       ti.make({4}, {0, 1, 0, 1}),
+       ti.make({2}, {0, 5})},
+      {"Short indices", ts.make({1}, {0}), ti.make({1}, {0})},
+  };
+  for (const auto& c : cases) {
+    SCOPED_TRACE(c.name);
+    context_.fail(torch::executor::Error::Ok);
+    auto out = tf.zeros({3, 3}, bounded);
+    auto offset2bag = ti.zeros({4}, bounded);
+    auto bag_size = ti.zeros({3}, bounded);
+    auto max_indices = ti.zeros({3}, bounded);
+    ET_EXPECT_KERNEL_FAILURE(
+        context_,
+        op_embedding_bag_out(
+            weight,
+            c.indices,
+            c.offsets,
+            1,
+            {},
+            false,
+            -1,
+            out,
+            offset2bag,
+            bag_size,
+            max_indices));
+    EXPECT_EQ(out.size(0), 3);
+    EXPECT_EQ(out.size(1), 3);
+    EXPECT_EQ(offset2bag.size(0), 4);
+    EXPECT_EQ(bag_size.size(0), 3);
+    EXPECT_EQ(max_indices.size(0), 3);
+  }
+}
diff --git a/kernels/test/targets.bzl b/kernels/test/targets.bzl
index 70bce803f3e..392e773de34 100644
--- a/kernels/test/targets.bzl
+++ b/kernels/test/targets.bzl
@@ -254,7 +254,12 @@ def define_common_targets():
     _common_op_test("op_diagonal_copy_test", ["aten", "portable"])
     _common_op_test("op_div_test", ["aten", "portable", "optimized"])
     _common_op_test("op_elu_test", ["aten", "portable", "optimized"])
-    _common_op_test("op_embedding_bag_test", ["aten", "portable"])
+    _common_op_test(
+        "op_embedding_bag_test",
+        ["aten", "portable", "optimized"],
+        # Lets the test widen the pool so bags split across parallel tasks.
+        deps = ["//executorch/extension/threadpool:threadpool"],
+    )
     _common_op_test("op_embedding_test", ["aten", "portable"])
     _common_op_test("op_empty_test", ["aten", "portable"])
     _common_op_test("op_eq_test", ["aten", "portable"])
diff --git a/shim_et/xplat/executorch/build/build_variables.bzl b/shim_et/xplat/executorch/build/build_variables.bzl
index 892a9063dae..061f369fe51 100644
--- a/shim_et/xplat/executorch/build/build_variables.bzl
+++ b/shim_et/xplat/executorch/build/build_variables.bzl
@@ -251,6 +251,7 @@ KERNELS_UTIL_ALL_DEPS_SRCS = [
     "kernels/portable/cpu/util/delinearize_index.cpp",
     "kernels/portable/cpu/util/distance_util.cpp",
     "kernels/portable/cpu/util/dtype_util.cpp",
+    "kernels/portable/cpu/util/embedding_bag_util.cpp",
     "kernels/portable/cpu/util/grid_sampler_2d_util.cpp",
     "kernels/portable/cpu/util/index_util.cpp",
     "kernels/portable/cpu/util/kernel_ops_util.cpp",
@@ -271,6 +272,7 @@ OPTIMIZED_KERNELS_SRCS = [
     "kernels/optimized/cpu/op_bmm.cpp",
     "kernels/optimized/cpu/op_div.cpp",
     "kernels/optimized/cpu/op_elu.cpp",
+    "kernels/optimized/cpu/op_embedding_bag.cpp",
     "kernels/optimized/cpu/op_exp.cpp",
     "kernels/optimized/cpu/op_fft_c2r.cpp",
     "kernels/optimized/cpu/op_fft_r2c.cpp",
diff --git a/shim_et/xplat/executorch/kernels/optimized/op_registration_util.bzl b/shim_et/xplat/executorch/kernels/optimized/op_registration_util.bzl
index 8dff4325971..bb1a3be1a31 100644
--- a/shim_et/xplat/executorch/kernels/optimized/op_registration_util.bzl
+++ b/shim_et/xplat/executorch/kernels/optimized/op_registration_util.bzl
@@ -211,6 +211,15 @@ OPTIMIZED_ATEN_OPS = (
             "//executorch/runtime/core/portable_type/c10/c10:aten_headers_for_executorch",
         ],
     ),
+    op_target(
+        name = "op_embedding_bag",
+        deps = [
+            "//executorch/extension/threadpool:threadpool",
+            "//executorch/kernels/portable/cpu:op_embedding_bag",
+            "//executorch/kernels/portable/cpu/util:embedding_bag_util",
+            "//executorch/runtime/core/portable_type/c10/c10:aten_headers_for_executorch",
+        ],
+    ),
     op_target(
         name = "op_exp",
         deps = [
diff --git a/shim_et/xplat/executorch/kernels/portable/op_registration_util.bzl b/shim_et/xplat/executorch/kernels/portable/op_registration_util.bzl
index b043e995e46..be2d896a6d0 100644
--- a/shim_et/xplat/executorch/kernels/portable/op_registration_util.bzl
+++ b/shim_et/xplat/executorch/kernels/portable/op_registration_util.bzl
@@ -558,6 +558,7 @@ ATEN_OPS = (
     op_target(
         name = "op_embedding_bag",
         deps = [
+            "//executorch/kernels/portable/cpu/util:embedding_bag_util",
             "//executorch/runtime/core/exec_aten/util:scalar_type_util",
             "//executorch/runtime/core/exec_aten/util:tensor_util",
         ],