From 34a4cca9b81fc67b5b76c428b8bbc0525edddaa6 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Tue, 14 Jul 2026 14:50:05 -0700 Subject: [PATCH] Update [ghstack-poisoned] --- backends/vulkan/test/op_tests/CMakeLists.txt | 4 + .../vulkan/test/op_tests/fused_ce_test.cpp | 145 ++++++++++++++++++ backends/vulkan/test/op_tests/targets.bzl | 6 + 3 files changed, 155 insertions(+) create mode 100644 backends/vulkan/test/op_tests/fused_ce_test.cpp diff --git a/backends/vulkan/test/op_tests/CMakeLists.txt b/backends/vulkan/test/op_tests/CMakeLists.txt index 40ac8fd0a8a..2a9c625412f 100644 --- a/backends/vulkan/test/op_tests/CMakeLists.txt +++ b/backends/vulkan/test/op_tests/CMakeLists.txt @@ -124,6 +124,10 @@ if(TARGET vulkan_backend AND LIB_TORCH) linear_q4gsw_dw_test ${CMAKE_CURRENT_SOURCE_DIR}/linear_q4gsw_dw_test.cpp test_utils ) + vulkan_op_test( + fused_ce_test ${CMAKE_CURRENT_SOURCE_DIR}/fused_ce_test.cpp + test_utils + ) # Only build generated op tests if a path to tags.yaml and # native_functions.yaml is provided. These files are required for codegen. diff --git a/backends/vulkan/test/op_tests/fused_ce_test.cpp b/backends/vulkan/test/op_tests/fused_ce_test.cpp new file mode 100644 index 00000000000..88775eae574 --- /dev/null +++ b/backends/vulkan/test/op_tests/fused_ce_test.cpp @@ -0,0 +1,145 @@ +/* + * 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 + +#include + +#include +#include +#include + +#include "test_utils.h" + +#include +#include + +// +// Reference implementation: transcription of the CPU-eager fused_ce_impl +// (backends/vulkan/custom_ops_lib.py), computed with ATen (library golden). +// + +std::pair fused_ce_reference( + const at::Tensor& logits, + const at::Tensor& labels, + double n_valid) { + at::Tensor mask = labels.ge(0); + at::Tensor safe = labels.clamp_min(0).to(at::kLong); + at::Tensor lse = at::logsumexp(logits, -1); + at::Tensor picked = logits.gather(-1, safe.unsqueeze(-1)).squeeze(-1); + at::Tensor loss = + at::where(mask, (lse - picked) / n_valid, at::zeros_like(lse)).sum(); + at::Tensor softmax = at::softmax(logits, -1); + at::Tensor onehot = at::one_hot(safe, logits.size(-1)).to(logits.dtype()); + at::Tensor dlogits = at::where( + mask.unsqueeze(-1), + (softmax - onehot) / n_valid, + at::zeros_like(softmax)); + return {loss, dlogits}; +} + +void test_vulkan_fused_ce( + const int n_rows, + const int vocab, + const std::vector& labels_data, + const double n_valid) { + torch::manual_seed(0); + + at::Tensor logits = + at::rand({n_rows, vocab}, at::device(at::kCPU).dtype(at::kFloat)) * 4.0 - + 2.0; + at::Tensor labels = + at::from_blob( + const_cast(labels_data.data()), {n_rows}, at::kInt) + .clone(); + + auto ref = fused_ce_reference(logits, labels, n_valid); + at::Tensor ref_loss = ref.first; + at::Tensor ref_dlogits = ref.second; + + using namespace vkcompute; + + GraphConfig config; + ComputeGraph graph(config); + + IOValueRef r_logits = graph.add_input_tensor( + logits.sizes().vec(), vkapi::kFloat, utils::kBuffer); + IOValueRef r_labels = + graph.add_input_tensor(labels.sizes().vec(), vkapi::kInt, utils::kBuffer); + + const ValueRef r_n_valid = graph.add_scalar(n_valid); + + const ValueRef r_loss = graph.add_tensor({}, vkapi::kFloat, utils::kBuffer); + const ValueRef r_dlogits = + graph.add_tensor({n_rows, vocab}, vkapi::kFloat, utils::kBuffer); + const ValueRef r_out = graph.add_value_list({r_loss, r_dlogits}); + + VK_GET_OP_FN("et_vk.fused_ce.default") + (graph, {r_logits.value, r_labels.value, r_n_valid, r_out}); + + ValueRef staging_loss = graph.set_output_tensor(r_loss); + ValueRef staging_dlogits = graph.set_output_tensor(r_dlogits); + + graph.prepare(); + graph.prepack(); + graph.propagate_resize(); + + graph.maybe_cast_and_copy_into_staging( + r_logits.staging, logits.const_data_ptr(), logits.numel(), vkapi::kFloat); + graph.maybe_cast_and_copy_into_staging( + r_labels.staging, labels.const_data_ptr(), labels.numel(), vkapi::kInt); + + graph.execute(); + + at::Tensor vk_loss = at::zeros({}, at::device(at::kCPU).dtype(at::kFloat)); + graph.maybe_cast_and_copy_from_staging( + staging_loss, vk_loss.mutable_data_ptr(), 1, vkapi::kFloat); + + at::Tensor vk_dlogits = + at::zeros({n_rows, vocab}, at::device(at::kCPU).dtype(at::kFloat)); + graph.maybe_cast_and_copy_from_staging( + staging_dlogits, + vk_dlogits.mutable_data_ptr(), + vk_dlogits.numel(), + vkapi::kFloat); + + const double atol = 1e-4; + const double rtol = 1e-4; + + const bool loss_ok = at::allclose(ref_loss, vk_loss, rtol, atol); + const bool dlogits_ok = at::allclose(ref_dlogits, vk_dlogits, rtol, atol); + + if (!loss_ok || !dlogits_ok) { + std::cout << "fused_ce mismatch: n_rows=" << n_rows << " vocab=" << vocab + << " n_valid=" << n_valid << std::endl; + std::cout << "loss ref=" << ref_loss.item() + << " vk=" << vk_loss.item() << std::endl; + std::cout << "max dlogits diff=" + << at::max(at::abs(ref_dlogits - vk_dlogits)).item() + << std::endl; + } + ASSERT_TRUE(loss_ok); + ASSERT_TRUE(dlogits_ok); +} + +TEST(VulkanFusedCeTest, all_valid_small) { + test_vulkan_fused_ce( + /*n_rows=*/4, /*vocab=*/8, /*labels=*/{0, 3, 7, 1}, /*n_valid=*/4.0); +} + +TEST(VulkanFusedCeTest, masked_label) { + // One masked row (label < 0): contributes 0 to loss and 0 gradient. + test_vulkan_fused_ce( + /*n_rows=*/4, /*vocab=*/8, /*labels=*/{2, -1, 5, 0}, /*n_valid=*/3.0); +} + +TEST(VulkanFusedCeTest, large_vocab_strided_reduce) { + // vocab >> NWORKERS exercises the strided per-row reduction. + test_vulkan_fused_ce( + /*n_rows=*/3, /*vocab=*/200, /*labels=*/{17, -1, 199}, /*n_valid=*/2.0); +} diff --git a/backends/vulkan/test/op_tests/targets.bzl b/backends/vulkan/test/op_tests/targets.bzl index ce93d864c78..83f9b39a9a4 100644 --- a/backends/vulkan/test/op_tests/targets.bzl +++ b/backends/vulkan/test/op_tests/targets.bzl @@ -192,6 +192,12 @@ def define_common_targets(is_fbcode = False): ":test_utils", ] ) + define_test_targets( + "fused_ce_test", + extra_deps = [ + ":test_utils", + ] + ) define_test_targets( "rms_norm_test", extra_deps = [