diff --git a/backends/native/runtime/engine/ggml/GgmlLowerLinear.cpp b/backends/native/runtime/engine/ggml/GgmlLowerLinear.cpp index 01625b42a09..d2ce76d172c 100644 --- a/backends/native/runtime/engine/ggml/GgmlLowerLinear.cpp +++ b/backends/native/runtime/engine/ggml/GgmlLowerLinear.cpp @@ -48,19 +48,31 @@ void GraphBuilder::lower_linear(const Node& node) { } void GraphBuilder::lower_row_gather(const Node& node) { - const TensorRef weight = tensor_arg(node, "weight", 0); + TensorRef weight = tensor_arg(node, "weight", 0); const Argument& arg = require_arg(node, "indices", 1).arg; ValueId indices_id = kInvalid; + // x[:, rows] of a [1, T, D] x, as Hugging Face selects logits_to_keep rows, + // gathers rows of the [T, D] matrix and keeps the leading 1. + bool leading_one = false; if (op_key(node.target) == "aten.embedding.default") { indices_id = arg.as_tensor().id; } else { const std::vector ids = arg.kind() == ArgKind::TensorList ? arg.as_tensor_list().ids : arg.as_optional_tensor_list().ids; - if (ids.size() != 1 || !valid(ids.front())) { - throw std::runtime_error("ggml: index supports one row-index tensor"); + leading_one = ids.size() == 2 && !valid(ids.front()) && + weight.sizes.size() == 3 && weight.sizes.front() == 1; + if (leading_one) { + weight.sizes.erase(weight.sizes.begin()); + } else if (ids.size() != 1) { + throw std::runtime_error( + "ggml: index supports one row-index tensor, after at most one " + "leading size-1 dimension"); + } + indices_id = ids.back(); + if (!valid(indices_id)) { + throw std::runtime_error("ggml: index needs a row-index tensor"); } - indices_id = ids.front(); } const ScalarType dtype = method_.graph.value(indices_id).tensor_meta().dtype; if (weight.sizes.size() != 2 || @@ -79,6 +91,9 @@ void GraphBuilder::lower_row_gather(const Node& node) { reshape(ctx, indices.tensor, {ggml_nelements(indices.tensor)}); std::vector sizes = indices.sizes; sizes.push_back(weight.sizes.at(1)); + if (leading_one) { + sizes.insert(sizes.begin(), 1); + } ggml_tensor* result = reshape(ctx, ggml_get_rows(ctx, weight.tensor, flat), sizes); set_single_output(node, TensorRef{result, std::move(sizes)}); diff --git a/backends/native/test/runtime/ggml/test_ggml_engine.cpp b/backends/native/test/runtime/ggml/test_ggml_engine.cpp index 38be8cf5d7e..0b33143b00a 100644 --- a/backends/native/test/runtime/ggml/test_ggml_engine.cpp +++ b/backends/native/test/runtime/ggml/test_ggml_engine.cpp @@ -282,9 +282,18 @@ std::unique_ptr make_context( host, program_bytes, kquant_block(kQ4KBlockBytes), options); } -// A tensor (by name), an int, a float, an int list, or None. -using TestArg = std:: - variant, std::monostate>; +// An optional tensor list: a tensor name per entry, or nullopt for None. +using TestOptionalTensors = std::vector>; + +// A tensor (by name), an int, a float, an int list, None, or an optional +// tensor list. +using TestArg = std::variant< + std::string, + int64_t, + double, + std::vector, + std::monostate, + TestOptionalTensors>; // A one-method program assembled value by value; every node's output value // carries the node's name. @@ -451,6 +460,22 @@ class TestProgram { fbs::ArgumentValue::IntListArg, fbs::CreateIntListArgDirect(builder_, list).Union()); } + if (const auto* tensors = std::get_if(&arg)) { + std::vector> names; + std::vector has_value; + for (const auto& name : *tensors) { + names.push_back(builder_.CreateString(name.value_or(""))); + has_value.push_back(name.has_value()); + } + return fbs::CreateArgument( + builder_, + fbs::ArgumentValue::OptionalTensorListArg, + fbs::CreateOptionalTensorListArg( + builder_, + builder_.CreateVector(names), + builder_.CreateVector(has_value)) + .Union()); + } return fbs::CreateArgument( builder_, fbs::ArgumentValue::FloatArg, @@ -2291,6 +2316,39 @@ TEST(GgmlEngineTest, OffGraphAttentionNeedsAMatchingCache) { std::runtime_error); } +// x[:, rows] of a [1, T, D] x, as Hugging Face selects logits_to_keep rows. +TEST(GgmlEngineTest, GathersRowsAfterALeadingOne) { + TestProgram program; + program.input("x", fbs::ScalarType::FLOAT, {1, 6, 8}); + program.input("rows", fbs::ScalarType::LONG, {2}); + program.call( + "selected", + "torch.ops.aten.index.Tensor", + {{"self", std::string("x")}, + {"indices", TestOptionalTensors{std::nullopt, "rows"}}}, + fbs::ScalarType::FLOAT, + {1, 2, 8}); + auto host = GgmlEngineHost::create(); + auto context = make_context(*host, program.finish("selected")); + auto executable = context->compile("forward"); + std::vector x(48); + std::iota(x.begin(), x.end(), 0.0f); + const std::vector rows = {4, 1}; + executable->set_input(0, x.data(), x.size(), ScalarType::Float); + executable->set_input(1, rows.data(), rows.size(), ScalarType::Long); + executable->execute(); + EXPECT_EQ(executable->output_sizes(0), (std::vector{1, 2, 8})); + std::vector out(16); + executable->get_output(0, out.data(), out.size(), ScalarType::Float); + for (int64_t i = 0; i < 8; ++i) { + EXPECT_EQ(out[i], x[4 * 8 + i]); + EXPECT_EQ(out[8 + i], x[1 * 8 + i]); + } + const std::vector outside = {6, 0}; + executable->set_input(1, outside.data(), outside.size(), ScalarType::Long); + EXPECT_ANY_THROW(executable->execute()); +} + TEST(GgmlEngineTest, RejectsQuantizedGraphInputsAndOutputs) { TestProgram as_input; as_input.input("x", fbs::ScalarType::FLOAT, {1, 256});