diff --git a/backends/native/runtime/vulkan/VulkanEngineExecutable.cpp b/backends/native/runtime/vulkan/VulkanEngineExecutable.cpp index eba42a04052..dd7af118ec0 100644 --- a/backends/native/runtime/vulkan/VulkanEngineExecutable.cpp +++ b/backends/native/runtime/vulkan/VulkanEngineExecutable.cpp @@ -51,6 +51,8 @@ #include #include +#include + #include #include #include @@ -116,8 +118,66 @@ vkapi::ScalarType to_vk_dtype(ScalarType t) { std::string("vulkan: unsupported dtype ") + scalar_type_name(t)); } +// ET-VK lacks int64 and bf16 storage, so int64 values run as int32 and bf16 +// values as f32 on the device. ScalarType device_dtype(ScalarType dtype) { - return dtype == ScalarType::Long ? ScalarType::Int : dtype; + if (dtype == ScalarType::Long) { + return ScalarType::Int; + } + if (dtype == ScalarType::BFloat16) { + return ScalarType::Float; + } + return dtype; +} + +template +T load_element(const void* data, size_t i) { + T value; + std::memcpy( + &value, static_cast(data) + i * sizeof(T), sizeof(T)); + return value; +} + +template +void store_element(void* data, size_t i, T value) { + std::memcpy(static_cast(data) + i * sizeof(T), &value, sizeof(T)); +} + +// Host copies into or out of bf16 go through f32 values rounded to bf16, so a +// bf16 tensor holds bf16 values at the host boundary whatever the host dtype. +void check_bf16_host_dtype(ScalarType dtype) { + if (dtype != ScalarType::BFloat16 && dtype != ScalarType::Float) { + throw std::runtime_error( + std::string("vulkan: bf16 copies need BFloat16 or Float host data, ") + + "not " + scalar_type_name(dtype)); + } +} + +std::vector +bf16_values_from_host(const void* data, size_t numel, ScalarType dtype) { + check_bf16_host_dtype(dtype); + std::vector values(numel); + for (size_t i = 0; i < numel; ++i) { + values[i] = dtype == ScalarType::BFloat16 + ? load_element(data, i) + : c10::BFloat16(load_element(data, i)); + } + return values; +} + +void bf16_values_to_host( + const std::vector& values, + void* data, + ScalarType dtype) { + check_bf16_host_dtype(dtype); + for (size_t i = 0; i < values.size(); ++i) { + const c10::BFloat16 value(values[i]); + if (dtype == ScalarType::BFloat16) { + store_element(data, i, value); + } else { + store_element(data, i, static_cast(value)); + } + } } size_t logical_numel(const TensorMeta& meta) { @@ -649,7 +709,8 @@ class VulkanEngineExecutable final : public EngineExecutable { std::vector zero_state_; // Bytes shared by tensor references until ET-VK consumes and releases them // during prepack: raw package bytes by canonical owner, and bytes computed - // from them (int32-narrowed Long constants, q4 repacks) by derivation key. + // from them (int32-narrowed Long constants, f32-widened bf16 constants, q4 + // repacks) by derivation key. std::unordered_map> source_constants_; std::unordered_map> @@ -792,8 +853,17 @@ class VulkanEngineExecutable final : public EngineExecutable { check_cache_position_input(data, numel, src_dtype, *bound, rows_written); pending_cache_rows_written_.at(i) = rows_written; } + const VkRef staging = input_staging_.at(i); + if (meta.dtype == ScalarType::BFloat16 || + src_dtype == ScalarType::BFloat16) { + const std::vector values = + bf16_values_from_host(data, numel, src_dtype); + graph_->maybe_cast_and_copy_into_staging( + staging, values.data(), numel, vkapi::kFloat); + return; + } graph_->maybe_cast_and_copy_into_staging( - input_staging_.at(i), data, numel, to_vk_dtype(src_dtype)); + staging, data, numel, to_vk_dtype(src_dtype)); } void execute() override { @@ -815,11 +885,17 @@ class VulkanEngineExecutable final : public EngineExecutable { if (data == nullptr) { throw std::runtime_error("vulkan: output data is null"); } + const VkRef staging = output_staging_.at(output_position); + if (output_dtype(i) == ScalarType::BFloat16 || + dst_dtype == ScalarType::BFloat16) { + std::vector values(numel); + graph_->maybe_cast_and_copy_from_staging( + staging, values.data(), numel, vkapi::kFloat); + bf16_values_to_host(values, data, dst_dtype); + return; + } graph_->maybe_cast_and_copy_from_staging( - output_staging_.at(output_position), - data, - numel, - to_vk_dtype(dst_dtype)); + staging, data, numel, to_vk_dtype(dst_dtype)); } private: @@ -857,7 +933,8 @@ class VulkanEngineExecutable final : public EngineExecutable { } const auto align4 = [](int64_t v) { return (v + 3) / 4 * 4; }; const int64_t limit = graph_->max_buffer_numel(); - const int64_t elem = static_cast(element_size(query->dtype)); + const int64_t elem = + static_cast(element_size(device_dtype(query->dtype))); int64_t weights = 0; if (__builtin_mul_overflow( query->sizes[2] * align4(query->sizes[1]), @@ -1255,8 +1332,29 @@ class VulkanEngineExecutable final : public EngineExecutable { OwnedBytes::from_vector(std::move(converted))); } storage = narrowed; + } else if (info->dtype == ScalarType::BFloat16) { + std::shared_ptr& widened = derived_constants_ + [std::to_string(info->owner.size()) + ":" + info->owner + ":f32"]; + if (!widened) { + const size_t numel = source.size() / sizeof(c10::BFloat16); + std::vector converted(numel * sizeof(float)); + for (size_t i = 0; i < numel; ++i) { + store_element( + converted.data(), + i, + static_cast( + load_element(source.data(), i))); + } + widened = std::make_shared( + OwnedBytes::from_vector(std::move(converted))); + } + storage = widened; } const ByteSpan bytes = storage->span(); + if (bytes.size() != logical_nbytes(device_meta(b.value_id))) { + throw std::runtime_error( + "vulkan: constant '" + b.key + "' does not match its device size"); + } auto* context = new std::shared_ptr(storage); return executorch::runtime::FreeableBuffer( bytes.data(), bytes.size(), release_owned_bytes, context); @@ -2171,7 +2269,7 @@ class VulkanEngineExecutable final : public EngineExecutable { } requests.push_back(AllocationRequest{ .value_id = id, - .size_bytes = logical_nbytes(value.tensor_meta()), + .size_bytes = logical_nbytes(device_meta(id)), .memory_kind = memory_kind(storage_at(id)), }); } diff --git a/backends/native/runtime/vulkan/passes/LowerRMSNorm.cpp b/backends/native/runtime/vulkan/passes/LowerRMSNorm.cpp index 1e2b1c6aab5..59ad98974cd 100644 --- a/backends/native/runtime/vulkan/passes/LowerRMSNorm.cpp +++ b/backends/native/runtime/vulkan/passes/LowerRMSNorm.cpp @@ -74,7 +74,8 @@ size_t lower_rms_norms(Graph& graph) { shape.values != std::vector{input_meta.sizes.back()} || weight_meta.sizes != shape.values || (input_meta.dtype != ScalarType::Float && - input_meta.dtype != ScalarType::Half) || + input_meta.dtype != ScalarType::Half && + input_meta.dtype != ScalarType::BFloat16) || weight_meta.dtype != input_meta.dtype || output_meta.dtype != input_meta.dtype || output_meta.sizes != input_meta.sizes) { diff --git a/backends/native/test/runtime/vulkan/test_lower_rms_norm.cpp b/backends/native/test/runtime/vulkan/test_lower_rms_norm.cpp index 6563b6041f6..73279663311 100644 --- a/backends/native/test/runtime/vulkan/test_lower_rms_norm.cpp +++ b/backends/native/test/runtime/vulkan/test_lower_rms_norm.cpp @@ -22,14 +22,14 @@ namespace ptn::vulkan { namespace { -Method make_rms_norm() { +Method make_rms_norm(ScalarType dtype = ScalarType::Float) { Method method; method.name = "forward"; Graph& graph = method.graph; graph.values = { - Value("input", ScalarType::Float, {2, 4}), - Value("weight", ScalarType::Float, {4}), - Value("output", ScalarType::Float, {2, 4}), + Value("input", dtype, {2, 4}), + Value("weight", dtype, {4}), + Value("output", dtype, {2, 4}), }; graph.value(1).role = ValueRole::Parameter; graph.nodes = { @@ -91,6 +91,19 @@ TEST(LowerRMSNormTest, ReadsUnnamedInputsByPosition) { EXPECT_EQ(lower_rms_norms(method.graph), 1); } +TEST(LowerRMSNormTest, RewritesBFloat16Norm) { + Method method = make_rms_norm(ScalarType::BFloat16); + + EXPECT_EQ(lower_rms_norms(method.graph), 1); + + const Node& norm = method.graph.node(2); + EXPECT_EQ(norm.target, "torch.ops.et_vk.rms_norm.default"); + ASSERT_EQ(norm.inputs.size(), 3); + EXPECT_EQ(norm.inputs[0].arg.as_tensor().id, 0); + EXPECT_EQ(norm.inputs[1].arg.as_tensor().id, 1); + EXPECT_NO_THROW(validate_graph(method.graph)); +} + TEST(LowerRMSNormTest, RejectsMismatchedNormalizedShape) { Method method = make_rms_norm(); method.graph.node(2).inputs[1].arg = IntListArg{{8}}; diff --git a/backends/native/test/runtime/vulkan/test_vulkan_engine.cpp b/backends/native/test/runtime/vulkan/test_vulkan_engine.cpp index 53d91ede342..5ef3fe32fb9 100644 --- a/backends/native/test/runtime/vulkan/test_vulkan_engine.cpp +++ b/backends/native/test/runtime/vulkan/test_vulkan_engine.cpp @@ -2833,6 +2833,81 @@ std::vector stored_floats(const std::vector& constants) { return tensors; } +// aten.rms_norm over the last dim of a [2, 4] input with a stored weight. +std::vector make_rms_norm_program(fbs::ScalarType dtype) { + flatbuffers::FlatBufferBuilder builder; + const auto meta = create_tensor_meta(builder, dtype, {2, 4}); + const auto weight_meta = create_tensor_meta(builder, dtype, {4}); + const std::vector> tensor_values = { + fbs::CreateTensorValueDirect(builder, "input", meta), + fbs::CreateTensorValueDirect(builder, "weight", weight_meta), + fbs::CreateTensorValueDirect(builder, "normed", meta), + }; + const std::vector> input_outputs = { + fbs::CreateOutputDirect(builder, "input")}; + const std::vector> weight_outputs = { + fbs::CreateOutputDirect(builder, "weight")}; + const std::vector> norm_outputs = { + fbs::CreateOutputDirect(builder, "normed")}; + const auto eps = fbs::CreateFloatArg(builder, 1e-5); + const std::vector> norm_inputs = { + fbs::CreateNamedArgumentDirect( + builder, "input", create_tensor_arg(builder, "input")), + fbs::CreateNamedArgumentDirect( + builder, "normalized_shape", create_int_list_arg(builder, {4})), + fbs::CreateNamedArgumentDirect( + builder, "weight", create_tensor_arg(builder, "weight")), + fbs::CreateNamedArgumentDirect( + builder, + "eps", + fbs::CreateArgument( + builder, fbs::ArgumentValue::FloatArg, eps.Union())), + }; + const std::vector> output_inputs = { + fbs::CreateNamedArgumentDirect( + builder, "", create_tensor_arg(builder, "normed"))}; + const std::vector> nodes = { + fbs::CreateNodeDirect( + builder, + "input", + fbs::OpKind::PLACEHOLDER, + "", + nullptr, + &input_outputs), + fbs::CreateNodeDirect( + builder, + "weight", + fbs::OpKind::PLACEHOLDER, + "", + nullptr, + &weight_outputs), + fbs::CreateNodeDirect( + builder, + "rms_norm", + fbs::OpKind::CALL_FUNCTION, + "torch.ops.aten.rms_norm.default", + &norm_inputs, + &norm_outputs), + fbs::CreateNodeDirect( + builder, "output", fbs::OpKind::OUTPUT, "", &output_inputs), + }; + const auto graph = fbs::CreateGraph( + builder, + builder.CreateVector(nodes), + create_strings(builder, {"input"}), + create_strings(builder, {"normed"}), + builder.CreateVector(tensor_values)); + const std::vector> constants = { + fbs::CreateNamedTensorRefDirect( + builder, "weight", "weight", weight_meta, fbs::InputKind::PARAMETER), + }; + return finish_program( + builder, + graph, + {fbs::CreateOutputSpecDirect(builder, "normed")}, + constants); +} + constexpr int64_t kHfTokens = 4; constexpr int64_t kHfHeads = 4; constexpr int64_t kHfKvHeads = 2; @@ -3439,6 +3514,79 @@ TEST(VulkanEngineTest, PreservesLogicalLongDtypeAndChecksNarrowing) { std::runtime_error); } +TEST(VulkanEngineTest, RunsBFloat16AddInFloatAndRoundsOutputToNearestEven) { + CompiledProgram compiled = + compile_program(make_add_program(fbs::ScalarType::BFLOAT16)); + EXPECT_EQ(compiled.executable->input_dtype(0), ScalarType::BFloat16); + EXPECT_EQ(compiled.executable->output_dtype(0), ScalarType::BFloat16); + + // bf16 bits: {1, 2, 3, 1} + {10, 20, 30, 2^-8}. + const std::vector left = {0x3F80, 0x4000, 0x4040, 0x3F80}; + const std::vector right = {0x4120, 0x41A0, 0x41F0, 0x3B80}; + compiled.executable->set_input( + 0, left.data(), left.size(), ScalarType::BFloat16); + compiled.executable->set_input( + 1, right.data(), right.size(), ScalarType::BFloat16); + compiled.executable->execute(); + std::vector output(4); + compiled.executable->get_output( + 0, output.data(), output.size(), ScalarType::BFloat16); + // {11, 22, 33, 1}: the f32 sum 1 + 2^-8 is a tie and rounds to even. + EXPECT_EQ(output, (std::vector{0x4130, 0x41B0, 0x4204, 0x3F80})); +} + +TEST(VulkanEngineTest, RoundsFloatHostDataOfBFloat16TensorsToBFloat16) { + CompiledProgram compiled = + compile_program(make_add_program(fbs::ScalarType::BFLOAT16)); + // 1 + 2^-9 rounds to 1 on input, and the sum 1 + 2^-8 rounds to even on + // output. + const std::vector left = {1.0f + 0x1p-9f, 1.0f, 3.0f, 1.0f}; + const std::vector right = {-1.0f, 0x1p-8f, 30.0f, 0.5f}; + compiled.executable->set_input( + 0, left.data(), left.size(), ScalarType::Float); + compiled.executable->set_input( + 1, right.data(), right.size(), ScalarType::Float); + compiled.executable->execute(); + std::vector output(4); + compiled.executable->get_output( + 0, output.data(), output.size(), ScalarType::Float); + EXPECT_EQ(output, (std::vector{0.0f, 1.0f, 33.0f, 1.5f})); + std::vector ints(4); + EXPECT_THROW( + compiled.executable->get_output( + 0, ints.data(), ints.size(), ScalarType::Int), + std::runtime_error); +} + +TEST(VulkanEngineTest, ExecutesBFloat16RMSNormWithStoredWeight) { + const std::string header = + R"({"weight":{"dtype":"BF16","shape":[4],"data_offsets":[0,8]}})"; + std::vector tensors = testing::make_safetensors(header, 8); + // bf16 bits: {0.5, 1.5, -2, 3}. + const std::vector weight = {0x3F00, 0x3FC0, 0xC000, 0x4040}; + std::memcpy( + tensors.data() + sizeof(uint64_t) + header.size(), + weight.data(), + weight.size() * sizeof(uint16_t)); + CompiledProgram compiled = compile_program( + make_rms_norm_program(fbs::ScalarType::BFLOAT16), std::move(tensors)); + // bf16 bits: rows {2, -2, 2, -2} and {4, 4, -4, 4}, each with an RMS that + // normalizes it to +-1. + const std::vector input = { + 0x4000, 0xC000, 0x4000, 0xC000, 0x4080, 0x4080, 0xC080, 0x4080}; + compiled.executable->set_input( + 0, input.data(), input.size(), ScalarType::BFloat16); + compiled.executable->execute(); + std::vector output(8); + compiled.executable->get_output( + 0, output.data(), output.size(), ScalarType::BFloat16); + // {0.5, -1.5, -2, -3} and {0.5, 1.5, 2, 3}. + EXPECT_EQ( + output, + (std::vector{ + 0x3F00, 0xBFC0, 0xC000, 0xC040, 0x3F00, 0x3FC0, 0x4000, 0x4040})); +} + TEST(VulkanEngineTest, ExecutesMetadataOnlyAliasAfterAliasLowering) { CompiledProgram compiled = compile_program(make_alias_program()); const std::vector input = {1.0f, 2.0f, 3.0f, 4.0f};